{-# LANGUAGE ParallelListComp, BangPatterns #-}
module Math.Polynomial.Chebyshev where

import Math.Polynomial
import Data.List

-- |The Chebyshev polynomials of the first kind with 'Integer' coefficients.
ts :: [Poly Integer]
ts :: [Poly Integer]
ts = Endianness -> [Integer] -> Poly Integer
forall a. (Num a, Eq a) => Endianness -> [a] -> Poly a
poly Endianness
LE [1] Poly Integer -> [Poly Integer] -> [Poly Integer]
forall a. a -> [a] -> [a]
: 
    [ Poly Integer -> Poly Integer -> Poly Integer
forall a. (Num a, Eq a) => Poly a -> Poly a -> Poly a
addPoly (Poly Integer
forall a. (Num a, Eq a) => Poly a
x                Poly Integer -> Poly Integer -> Poly Integer
forall a. (Num a, Eq a) => Poly a -> Poly a -> Poly a
`multPoly` Poly Integer
t_n)
              (Endianness -> [Integer] -> Poly Integer
forall a. (Num a, Eq a) => Endianness -> [a] -> Poly a
poly Endianness
LE [-1,0,1] Poly Integer -> Poly Integer -> Poly Integer
forall a. (Num a, Eq a) => Poly a -> Poly a -> Poly a
`multPoly` Poly Integer
u_n)
    | Poly Integer
t_n <- [Poly Integer]
ts
    | Poly Integer
u_n <- Endianness -> [Integer] -> Poly Integer
forall a. (Num a, Eq a) => Endianness -> [a] -> Poly a
poly Endianness
LE [0] Poly Integer -> [Poly Integer] -> [Poly Integer]
forall a. a -> [a] -> [a]
: [Poly Integer]
us
    ]

-- The Chebyshev polynomials of the second kind with 'Integer' coefficients.
us :: [Poly Integer]
us :: [Poly Integer]
us = 
    [ Poly Integer -> Poly Integer -> Poly Integer
forall a. (Num a, Eq a) => Poly a -> Poly a -> Poly a
addPoly Poly Integer
t_n (Poly Integer -> Poly Integer -> Poly Integer
forall a. (Num a, Eq a) => Poly a -> Poly a -> Poly a
multPoly Poly Integer
forall a. (Num a, Eq a) => Poly a
x Poly Integer
u_n)
    | Poly Integer
t_n <- [Poly Integer]
ts
    | Poly Integer
u_n <- Endianness -> [Integer] -> Poly Integer
forall a. (Num a, Eq a) => Endianness -> [a] -> Poly a
poly Endianness
LE [0] Poly Integer -> [Poly Integer] -> [Poly Integer]
forall a. a -> [a] -> [a]
: [Poly Integer]
us
    ]

-- |Compute the coefficients of the n'th Chebyshev polynomial of the first kind.
t :: (Num a, Eq a) => Int -> Poly a
t :: Int -> Poly a
t n :: Int
n | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= 0    = Endianness -> [a] -> Poly a
forall a. (Num a, Eq a) => Endianness -> [a] -> Poly a
poly Endianness
LE ([a] -> Poly a) -> (Poly Integer -> [a]) -> Poly Integer -> Poly a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> a) -> [Integer] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> a
forall a. Num a => Integer -> a
fromInteger ([Integer] -> [a])
-> (Poly Integer -> [Integer]) -> Poly Integer -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Endianness -> Poly Integer -> [Integer]
forall a. (Num a, Eq a) => Endianness -> Poly a -> [a]
polyCoeffs Endianness
LE (Poly Integer -> Poly a) -> Poly Integer -> Poly a
forall a b. (a -> b) -> a -> b
$ [Poly Integer]
ts [Poly Integer] -> Int -> Poly Integer
forall a. [a] -> Int -> a
!! Int
n
    | Bool
otherwise = [Char] -> Poly a
forall a. HasCallStack => [Char] -> a
error "t: negative index"

-- |Compute the coefficients of the n'th Chebyshev polynomial of the second kind.
u :: (Num a, Eq a) => Int -> Poly a
u :: Int -> Poly a
u n :: Int
n | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= 0    = Endianness -> [a] -> Poly a
forall a. (Num a, Eq a) => Endianness -> [a] -> Poly a
poly Endianness
LE ([a] -> Poly a) -> (Poly Integer -> [a]) -> Poly Integer -> Poly a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> a) -> [Integer] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> a
forall a. Num a => Integer -> a
fromInteger ([Integer] -> [a])
-> (Poly Integer -> [Integer]) -> Poly Integer -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Endianness -> Poly Integer -> [Integer]
forall a. (Num a, Eq a) => Endianness -> Poly a -> [a]
polyCoeffs Endianness
LE (Poly Integer -> Poly a) -> Poly Integer -> Poly a
forall a b. (a -> b) -> a -> b
$ [Poly Integer]
us [Poly Integer] -> Int -> Poly Integer
forall a. [a] -> Int -> a
!! Int
n
    | Bool
otherwise = [Char] -> Poly a
forall a. HasCallStack => [Char] -> a
error "u: negative index"

-- |Evaluate the n'th Chebyshev polynomial of the first kind at a point X.  
-- Both more efficient and more numerically stable than computing the 
-- coefficients and evaluating the polynomial.
evalT :: Num a => Int -> a -> a
evalT :: Int -> a -> a
evalT n :: Int
n x :: a
x = (a, a) -> a
forall a b. (a, b) -> a
fst (Int -> a -> (a, a)
forall a. Num a => Int -> a -> (a, a)
evalTU Int
n a
x)

-- |Evaluate all the Chebyshev polynomials of the first kind at a point X.
evalTs :: Num a => a -> [a]
evalTs :: a -> [a]
evalTs = ([a], [a]) -> [a]
forall a b. (a, b) -> a
fst (([a], [a]) -> [a]) -> (a -> ([a], [a])) -> a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ([a], [a])
forall a. Num a => a -> ([a], [a])
evalTsUs

-- |Evaluate the n'th Chebyshev polynomial of the second kind at a point X.  
-- Both more efficient and more numerically stable than computing the 
-- coefficients and evaluating the polynomial.
evalU :: Num a => Int -> a -> a
evalU :: Int -> a -> a
evalU n :: Int
n x :: a
x = (a, a) -> a
forall a b. (a, b) -> b
snd (Int -> a -> (a, a)
forall a. Num a => Int -> a -> (a, a)
evalTU Int
n a
x)

-- |Evaluate all the Chebyshev polynomials of the second kind at a point X.
evalUs :: Num a => a -> [a]
evalUs :: a -> [a]
evalUs = ([a], [a]) -> [a]
forall a b. (a, b) -> b
snd (([a], [a]) -> [a]) -> (a -> ([a], [a])) -> a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ([a], [a])
forall a. Num a => a -> ([a], [a])
evalTsUs

-- |Evaluate the n'th Chebyshev polynomials of both kinds at a point X.
evalTU :: Num a => Int -> a -> (a,a)
evalTU :: Int -> a -> (a, a)
evalTU n :: Int
n x :: a
x = Int -> a -> a -> (a, a)
forall t. (Eq t, Num t) => t -> a -> a -> (a, a)
go Int
n 1 0
    where
        go :: t -> a -> a -> (a, a)
go !t
0 !a
t_n !a
u_n = (a
t_n, a
u_n)
        go !t
n !a
t_n !a
u_n = t -> a -> a -> (a, a)
go (t
nt -> t -> t
forall a. Num a => a -> a -> a
-1) a
t_np1 a
u_np1
            where
                t_np1 :: a
t_np1 = a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
t_n a -> a -> a
forall a. Num a => a -> a -> a
- (1a -> a -> a
forall a. Num a => a -> a -> a
-a
xa -> a -> a
forall a. Num a => a -> a -> a
*a
x)a -> a -> a
forall a. Num a => a -> a -> a
*a
u_n
                u_np1 :: a
u_np1 = a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
u_n a -> a -> a
forall a. Num a => a -> a -> a
+ a
t_n

-- |Evaluate all the Chebyshev polynomials of both kinds at a point X.
evalTsUs :: Num a => a -> ([a], [a])
evalTsUs :: a -> ([a], [a])
evalTsUs x :: a
x = ([a]
ts, [a] -> [a]
forall a. [a] -> [a]
tail [a]
us)
    where
        ts :: [a]
ts = 1 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
t_n a -> a -> a
forall a. Num a => a -> a -> a
- (1a -> a -> a
forall a. Num a => a -> a -> a
-a
xa -> a -> a
forall a. Num a => a -> a -> a
*a
x)a -> a -> a
forall a. Num a => a -> a -> a
*a
u_n  | a
t_n <- [a]
ts | a
u_n <- [a]
us]
        us :: [a]
us = 0 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
u_n a -> a -> a
forall a. Num a => a -> a -> a
+ a
t_n          | a
t_n <- [a]
ts | a
u_n <- [a]
us]

-- |Compute the roots of the n'th Chebyshev polynomial of the first kind.
tRoots :: Floating a => Int -> [a]
tRoots :: Int -> [a]
tRoots   n :: Int
n = [a -> a
forall a. Floating a => a -> a
cos (a
forall a. Floating a => a
pi a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n a -> a -> a
forall a. Num a => a -> a -> a
* (Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k a -> a -> a
forall a. Num a => a -> a -> a
+ 0.5)) | Int
k <- [0..Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1]]

-- |Compute the extreme points of the n'th Chebyshev polynomial of the first kind.
tExtrema :: Floating a => Int -> [a]
tExtrema :: Int -> [a]
tExtrema n :: Int
n = [a -> a
forall a. Floating a => a -> a
cos (a
forall a. Floating a => a
pi a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n a -> a -> a
forall a. Num a => a -> a -> a
*  Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k       ) | Int
k <- [0..Int
n]]

-- |@chebyshevFit n f@ returns a list of N coefficients @cs@ such that 
-- @f x@ ~= @sum (zipWith (*) cs (evalTs x))@ on the interval -1 < x < 1.
-- 
-- The N roots of the N'th Chebyshev polynomial are the fitting points at 
-- which the function will be evaluated and at which the approximation will be
-- exact.  These points always lie within the interval -1 < x < 1.  Outside
-- this interval, the approximation will diverge quickly.
--
-- This function deviates from most chebyshev-fit implementations in that it
-- returns the first coefficient pre-scaled so that the series evaluation 
-- operation is a simple inner product, since in most other algorithms
-- operating on chebyshev series, that factor is almost always a nuissance.
chebyshevFit :: Floating a => Int -> (a -> a) -> [a]
chebyshevFit :: Int -> (a -> a) -> [a]
chebyshevFit n :: Int
n f :: a -> a
f = 
    [ a
oneOrTwo a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n 
    a -> a -> a
forall a. Num a => a -> a -> a
* [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> a -> a
forall a. Num a => a -> a -> a
(*) [a]
ts [a]
fxs)
    | [a]
ts <- [[a]] -> [[a]]
forall a. [[a]] -> [[a]]
transpose [[a]]
txs
    | a
oneOrTwo <- 1 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: a -> [a]
forall a. a -> [a]
repeat 2
    ]
    where
        txs :: [[a]]
txs = (a -> [a]) -> [a] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
n ([a] -> [a]) -> (a -> [a]) -> a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> [a]
forall a. Num a => a -> [a]
evalTs) [a]
xs
        fxs :: [a]
fxs = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map a -> a
f [a]
xs
        xs :: [a]
xs = Int -> [a]
forall a. Floating a => Int -> [a]
tRoots Int
n

-- |Evaluate a Chebyshev series expansion with a finite number of terms.
-- 
-- Note that this function expects the first coefficient to be pre-scaled
-- by 1/2, which is what is produced by 'chebyshevFit'.  Thus, this computes
-- a simple inner product of the given list with a matching-length sequence of
-- chebyshev polynomials.
evalChebyshevSeries :: Num a => [a] -> a -> a
evalChebyshevSeries :: [a] -> a -> a
evalChebyshevSeries     []  _ = 0
evalChebyshevSeries (c0 :: a
c0:cs :: [a]
cs) x :: a
x = 
        let b1 :: a
b1:b2 :: a
b2:_ = [a] -> [a]
forall a. [a] -> [a]
reverse [a]
bs
         in a
xa -> a -> a
forall a. Num a => a -> a -> a
*a
b1 a -> a -> a
forall a. Num a => a -> a -> a
- a
b2 a -> a -> a
forall a. Num a => a -> a -> a
+ a
c0
    where
        -- Clenshaw's recurrence formula
        bs :: [a]
bs = 0 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: 0 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [2a -> a -> a
forall a. Num a => a -> a -> a
*a
xa -> a -> a
forall a. Num a => a -> a -> a
*a
b1 a -> a -> a
forall a. Num a => a -> a -> a
- a
b2 a -> a -> a
forall a. Num a => a -> a -> a
+ a
c | b2 :: a
b2:b1 :: a
b1:_ <- [a] -> [[a]]
forall a. [a] -> [[a]]
tails [a]
bs | a
c <- [a] -> [a]
forall a. [a] -> [a]
reverse [a]
cs]