{-# LANGUAGE GADTs #-}

-- |
-- Module      :  Diagrams.CubicSpline.Internal
-- Copyright   :  (c) 2011 diagrams-lib team (see LICENSE)
-- License     :  BSD-style (see LICENSE)
-- Maintainer  :  diagrams-discuss@googlegroups.com
--
-- A /cubic spline/ is a smooth, connected sequence of cubic curves
-- passing through a given sequence of points.  This module implements
-- a straightforward spline generation algorithm based on solving
-- tridiagonal systems of linear equations.
module Diagrams.CubicSpline.Internal (
  -- * Solving for spline coefficents
  solveCubicSplineDerivatives,
  solveCubicSplineDerivativesClosed,
  solveCubicSplineCoefficients,
) where

import Data.List
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty.Compat as NE
import Diagrams.CubicSpline.NonSingleton (NonSingleton (..))
import qualified Diagrams.CubicSpline.NonSingleton as NS
import Diagrams.Solve.Tridiagonal

-- | Use the tri-diagonal solver with the appropriate parameters for an open cubic spline.
--
--   See e.g. https://observablehq.com/@jrus/cubic-spline
solveCubicSplineDerivatives :: Fractional a => NonSingleton a -> NonEmpty a
solveCubicSplineDerivatives :: forall a. Fractional a => NonSingleton a -> NonEmpty a
solveCubicSplineDerivatives NonSingleton a
xs = NonEmpty a -> NonEmpty a -> NonEmpty a -> NonEmpty a -> NonEmpty a
forall a.
Fractional a =>
NonEmpty a -> NonEmpty a -> NonEmpty a -> NonEmpty a -> NonEmpty a
solveTriDiagonal NonEmpty a
as NonEmpty a
bs NonEmpty a
as (NonSingleton a -> NonEmpty a
forall a. NonSingleton a -> NonEmpty a
NS.toNonEmpty NonSingleton a
ds)
 where
  as :: NonEmpty a
as = (a -> a) -> NonEmpty a -> NonEmpty a
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NE.map (a -> a -> a
forall a b. a -> b -> a
const a
1) (NonSingleton a -> NonEmpty a
forall a. NonSingleton a -> NonEmpty a
NS.tail NonSingleton a
ds)
  bs :: NonEmpty a
bs = a
2 a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a -> a -> a
forall a b. a -> b -> a
const a
4) (NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NE.tail (NonSingleton a -> NonEmpty a
forall a. NonSingleton a -> NonEmpty a
NS.tail NonSingleton a
ds)) [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
2]
  -- zipWith f [x1, x2, x3, ..., xn, xn] [x0, x0, x1, ... x{n-2}, x{n-1}]
  ds :: NonSingleton a
ds = (a -> a -> a) -> NonSingleton a -> NonSingleton a -> NonSingleton a
forall a b c.
(a -> b -> c) -> NonSingleton a -> NonSingleton b -> NonSingleton c
NS.zipWith a -> a -> a
forall {a}. Num a => a -> a -> a
f (NonEmpty a -> NonSingleton a
forall a. NonEmpty a -> NonSingleton a
NS.dupLast (NonSingleton a -> NonEmpty a
forall a. NonSingleton a -> NonEmpty a
NS.tail NonSingleton a
xs)) (NonSingleton a -> NonSingleton a
forall a. NonSingleton a -> NonSingleton a
NS.dupFirst NonSingleton a
xs)
  f :: a -> a -> a
f a
a a
b = a
3 a -> a -> a
forall {a}. Num a => a -> a -> a
* (a
a a -> a -> a
forall {a}. Num a => a -> a -> a
- a
b)

-- | Use the cyclic-tri-diagonal solver with the appropriate parameters for a closed cubic spline.
solveCubicSplineDerivativesClosed :: Fractional a => NonSingleton a -> NonEmpty a
solveCubicSplineDerivativesClosed :: forall a. Fractional a => NonSingleton a -> NonEmpty a
solveCubicSplineDerivativesClosed NonSingleton a
xs = NonEmpty a
-> NonEmpty a -> NonEmpty a -> NonEmpty a -> a -> a -> NonEmpty a
forall a.
Fractional a =>
NonEmpty a
-> NonEmpty a -> NonEmpty a -> NonEmpty a -> a -> a -> NonEmpty a
solveCyclicTriDiagonal NonEmpty a
as NonEmpty a
bs NonEmpty a
as NonEmpty a
ds a
1 a
1
 where
  as :: NonEmpty a
as = (a -> a) -> NonEmpty a -> NonEmpty a
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a -> a -> a
forall a b. a -> b -> a
const a
1) (NonSingleton a -> NonEmpty a
forall a. NonSingleton a -> NonEmpty a
NS.tail NonSingleton a
xs)
  bs :: NonEmpty a
bs = (a -> a) -> NonEmpty a -> NonEmpty a
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a -> a -> a
forall a b. a -> b -> a
const a
4) (NonSingleton a -> NonEmpty a
forall a. NonSingleton a -> NonEmpty a
NS.toNonEmpty NonSingleton a
xs)

  -- zipWith f [x1, x2, x3, ..., xn, x0] [xn, x0, x1, ..., x{n-1}]
  -- ds = take l $ zipWith f (drop 1 xs') (drop (l - 1) xs') where l = length xs
  xsNE :: NonEmpty a
xsNE = NonSingleton a -> NonEmpty a
forall a. NonSingleton a -> NonEmpty a
NS.toNonEmpty NonSingleton a
xs
  ds :: NonEmpty a
ds = (a -> a -> a) -> NonEmpty a -> NonEmpty a -> NonEmpty a
forall a b c.
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
NE.zipWith a -> a -> a
forall {a}. Num a => a -> a -> a
f (NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NE.tail NonEmpty a
xsNE [a] -> a -> NonEmpty a
forall a. [a] -> a -> NonEmpty a
|: NonEmpty a -> a
forall a. NonEmpty a -> a
NE.head NonEmpty a
xsNE) (a -> NonEmpty a -> NonEmpty a
forall a. a -> NonEmpty a -> NonEmpty a
NE.cons (NonEmpty a -> a
forall a. NonEmpty a -> a
NE.last NonEmpty a
xsNE) NonEmpty a
xsNE)
  f :: a -> a -> a
f a
a a
b = a
3 a -> a -> a
forall {a}. Num a => a -> a -> a
* (a
a a -> a -> a
forall {a}. Num a => a -> a -> a
- a
b)

  (|:) :: [a] -> a -> NonEmpty a
  |: :: forall a. [a] -> a -> NonEmpty a
(|:) [a]
ys a
y = (a -> NonEmpty a -> NonEmpty a) -> NonEmpty a -> [a] -> NonEmpty a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> NonEmpty a -> NonEmpty a
forall a. a -> NonEmpty a -> NonEmpty a
NE.cons (a -> NonEmpty a
forall a. a -> NonEmpty a
NE.singleton a
y) [a]
ys

-- | Use the cyclic-tri-diagonal solver with the appropriate parameters for a closed cubic spline.
solveCubicSplineCoefficients :: Fractional a => Bool -> NonSingleton a -> [[a]]
solveCubicSplineCoefficients :: forall a. Fractional a => Bool -> NonSingleton a -> [[a]]
solveCubicSplineCoefficients Bool
closed NonSingleton a
xs =
  [ [a
x, a
d, a
3 a -> a -> a
forall {a}. Num a => a -> a -> a
* (a
x1 a -> a -> a
forall {a}. Num a => a -> a -> a
- a
x) a -> a -> a
forall {a}. Num a => a -> a -> a
- a
2 a -> a -> a
forall {a}. Num a => a -> a -> a
* a
d a -> a -> a
forall {a}. Num a => a -> a -> a
- a
d1, a
2 a -> a -> a
forall {a}. Num a => a -> a -> a
* (a
x a -> a -> a
forall {a}. Num a => a -> a -> a
- a
x1) a -> a -> a
forall {a}. Num a => a -> a -> a
+ a
d a -> a -> a
forall {a}. Num a => a -> a -> a
+ a
d1]
  | (a
x, a
x1, a
d, a
d1) <- [a] -> [a] -> [a] -> [a] -> [(a, a, a, a)]
forall a b c d. [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]
zip4 (NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty a
xs') (NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NE.tail NonEmpty a
xs') (NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty a
ds') (NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NE.tail NonEmpty a
ds')
  ]
 where
  ds :: NonEmpty a
ds
    | Bool
closed = NonSingleton a -> NonEmpty a
forall a. Fractional a => NonSingleton a -> NonEmpty a
solveCubicSplineDerivativesClosed NonSingleton a
xs
    | Bool
otherwise = NonSingleton a -> NonEmpty a
forall a. Fractional a => NonSingleton a -> NonEmpty a
solveCubicSplineDerivatives NonSingleton a
xs
  close :: NonEmpty a -> NonEmpty a
close NonEmpty a
as
    | Bool
closed = NonEmpty a
as NonEmpty a -> NonEmpty a -> NonEmpty a
forall a. Semigroup a => a -> a -> a
<> a -> NonEmpty a
forall a. a -> NonEmpty a
NE.singleton (NonEmpty a -> a
forall a. NonEmpty a -> a
NE.head NonEmpty a
as)
    | Bool
otherwise = NonEmpty a
as
  xs' :: NonEmpty a
xs' = NonEmpty a -> NonEmpty a
close (NonSingleton a -> NonEmpty a
forall a. NonSingleton a -> NonEmpty a
NS.toNonEmpty NonSingleton a
xs)
  ds' :: NonEmpty a
ds' = NonEmpty a -> NonEmpty a
close NonEmpty a
ds