chebyshev and other orthogonal polynomial bases (Spread polynomials borked)

Chebyshev.hs

{-# LANGUAGE MultiParamTypeClasses, FunctionalDependencies, FlexibleInstances, FlexibleContexts, UndecidableInstances #-}
module Numeric.Polynomial.Chebyshev where

import qualified Data.Map as Map
import Data.List (foldl')

zipWithT :: (a -> b -> c) -> ([a] -> [c]) -> ([b] -> [c]) -> [a] -> [b] -> [c]
zipWithT f fa fb = go where
  go (a:as) (b:bs) = f a b : go as bs
  go [] bs = fb bs
  go as [] = fa as

-- class InnerProductSpace p a | p -> a where
--  prod :: p -> p -> a

sum' :: Num a => [a] -> a
sum' = foldl' (+) 0

-- evaluating a polynomial from another representation:

class Eval p a | p -> a where
  (.*) :: a -> p -> p
  (./) :: p -> a -> p
  at :: p -> a -> a
  toPower :: p -> Power a
  fromPower :: Power a -> p
  mulx :: p -> p

-- The standard (power/MacLaurin) basis
newtype Power a = Power { runPower :: [a] } deriving (Eq,Show)

instance Fractional a => Eval (Power a) a where
  Power coefs `at` x = sum' $ zipWith (*) coefs (iterate (*x) 1)
  toPower = id
  fromPower = id
  a .* Power as = Power (map (a*) as)
  Power as ./ a = Power (map (/a) as)
  mulx (Power as) = Power (0:as)

instance Fractional a => Num (Power a) where
  Power a + Power b = Power (zipWithT (+) id id a b)
  Power a - Power b = Power (zipWithT (-) id (map negate) a b)
  Power (a:as) * bbs@(Power (b:bs)) = Power (a*b : runPower (Power (map (a *) bs) + Power as * bbs))
  Power [] * bs = Power []
  as * Power [] = Power []
  fromInteger n = Power [fromInteger n]
  negate (Power as) = Power (map negate as)
  abs = error "borked numerical tower"
  signum = error "borked numerical tower"

instance Fractional a => Fractional (Power a) where
  a / b = a * recip b
  recip (Power []) = Power [recip 0]
  recip (Power (a:as)) = Power (ooa : go [ooa]) where
    ooa = 1/a
    go acc = b : go (b:acc) where 
      b = - sum' (zipWith (*) as acc) / a
  fromRational a = Power [fromRational a]
    
-- The factorial basis x^n/n! used by exponential generating functions
newtype Factorial a = Factorial { runFactorial :: [a] } deriving (Show,Eq)

instance Fractional a => Num (Factorial a) where
  Factorial a + Factorial b = Factorial (zipWithT (+) id id a b)
  Factorial a - Factorial b = Factorial (zipWithT (-) id (map negate) a b)
  Factorial (a:as) * bbs@(Factorial (b:bs)) = Factorial (a*b : runFactorial (Factorial (map (a *) bs) + Factorial as * bbs))
  as * bs = fromPower (toPower as * toPower bs)
  -- Factorial [] * bs = Factorial []
  -- as * Factorial [] = Factorial []
  -- fromInteger n = Factorial [fromInteger n]
  negate (Factorial as) = Factorial (map negate as)
  fromInteger n = Factorial [fromInteger n]
  abs = error "borked numerical tower"
  signum = error "borked numerical tower"

facts :: Num a => [a]
facts = go 1 1 where
  go n m | n `seq` m `seq` False = undefined
  go n m = n : go (n*m) (m + 1)
  
factPowers :: Fractional a => a -> [a]
factPowers x = go 1 1 where
  go xn m | xn `seq` m `seq` False = undefined
  go xn m = xn : go (xn*x/m) (m + 1)

instance Fractional a => Eval (Factorial a) a where
  -- Factorial coefs `at` x = sum' $ zipWith3 (\coef xn f -> coef * xn * f) coefs (iterate (*x) 1) facts
  Factorial coefs `at` x = sum' $ zipWith (*) coefs (factPowers x)
  toPower (Factorial coefs) = Power (zipWith (/) coefs facts)
  fromPower (Power coefs) = Factorial (zipWith (*) coefs facts)
  a .* Factorial as = Factorial (map (a*) as)
  Factorial as ./ a = Factorial (map (/a) as)
  mulx (Factorial as) = Factorial $ 0 : zipWith (*) (iterate (+1) 1) as

-- Chebyshev polynomials of the first kind
newtype T a = T [a] deriving (Eq, Show)

t :: Num a => a -> [a]
t x = ts where
  ts = 1 : x : zipWith (\tnm1x tnx -> 2*x*tnx - tnm1x) ts (tail ts)

-- other recurrences
-- t_{2n}(x)   = t_2(t_n(x)) = 2 * (t_n(x)^2) - 1
-- t_{2n+1}(x) = (T_{n+1}(x) + T_n(x))^2/(x + 1) - 1

instance Fractional a => Eval (T a) a where
  T as `at` a = sum' $ zipWith (*) as (t a)
  toPower (T as) = T (map (Power . return) as) `at` Power [0,1]
  fromPower (Power as) = sum' $ zipWith (.*) as (iterate mulx 1) 
  a .* T as = T (map (a*) as)
  T as ./ a = T (map (/a) as)
  mulx (T []) = T []
  mulx (T [0]) = T []
  mulx (T (a:as)) = T tmp + T (a*0 : a : tmp) where tmp = map (/2) as 

instance Fractional a => Num (T a) where
  T a + T b = T (zipWithT (+) id id a b)
  T a - T b = T (zipWithT (-) id (map negate) a b)
  -- TODO: improve this
  T as * T bs = T ds where
    ds = flip fill 0 $ 
         Map.toList $ 
         Map.fromListWith (+) 
         [ (f i j, a*b) 
         | (i, a) <- zip [0..] as
         , (j, b) <- zip [0..] bs
         , f <- [(+), \ i' j' -> abs (i' - j')] 
         ] where
    fill ((k,c):cs) n | n == k = c : (fill cs $! n + 1)
    fill [] _ = []
    fill cs n = 0 : (fill cs $! n + 1)
  -- as * T [] = T []
  -- T [] * bs = T []
  -- as * bs = fromPower (toPower as * toPower bs)
  fromInteger n = T [fromInteger n]
  negate (T as) = T (map negate as)
  abs = error "borked numerical tower"
  signum = error "borked numerical tower"

-- Chebyshev polynomials of the second kind
newtype U a = U [a] deriving (Show,Eq)

u :: Num a => a -> [a]
u x = us where
  us = 1 : 2*x : zipWith (\unm1x unx -> 2*x*unx - unm1x) us (tail us)

instance Fractional a => Eval (U a) a where
  U as `at` a = sum' $ zipWith (*) as (u a)
  toPower (U as) = U (map (Power . return) as) `at` Power [0,1]
  fromPower (Power as) = sum' $ zipWith (.*) as (iterate mulx 1)
  a .* U as = U (map (a*) as)
  U as ./ a = U (map (/a) as)
  mulx (U []) = U []
  mulx (U [0]) = U []
  mulx (U (a:as)) = U tmp + U (a*0 : a/2 : tmp) where tmp = map (/2) as 
  -- x * u n x = (u (n + 1) x + u (n - 1) x) / 2

instance Fractional a => Num (U a) where
  U as + U bs = U (zipWithT (+) id id as bs)
  U as - U bs = U (zipWithT (-) id (map negate) as bs)
  as * bs = fromPower (toPower as * toPower bs)
  fromInteger n = U [fromInteger n]
  negate (U as) = U (map negate as)
  abs = error "borked numerical tower"
  signum = error "borked numerical tower"

newtype Spread a = Spread [a] deriving (Eq,Show)

spread :: Num a => a -> [a]
spread x = 1 : ss where twox = 2*x; ss = x : zipWith (\snm1x snx -> 2*(1 - twox)*snx - snm1x + twox) (0:ss) ss

instance Fractional a => Eval (Spread a) a where
  Spread as `at` a = sum' $ zipWith (*) as (spread a)
  toPower (Spread as) = Spread (map (Power . return) as) `at` Power [0,1]
  fromPower (Power as) = sum' $ zipWith (.*) as (iterate mulx 1)
  a .* Spread as  = Spread $ map (a*) as
  Spread as ./ a = Spread (map (/a) as)
  mulx (Spread [])  = Spread []
  mulx (Spread [0]) = Spread []
  mulx (Spread [a]) = Spread [0,a]
  mulx (Spread (a:as)) = (Spread (0:as) ./ 2) - Spread t' - Spread (a*0: a-0.5: t) where 
    t = map (/4) as
    t' = case t of
      [] -> []
      (t0:ts) -> 0:ts
-- mulx (Spread (a:as)) = (Spread (0:as) ./ 2) + 2*s1 snm1 - snp1 where 
--    snm1 = Spread $ case as of 
--      (x:xs) -> 0:xs
--      [] -> []
--    snp1 = Spread $ a*0:a:as
-- mulx (Spread [a,b]) = Spread $ error "TODO"
-- s S_n(s) = (2 S_n(s) + S_1(s) - S_(n-1)(s) - S_(n+1)(s)) / 4
-- mulx (Spread 
-- S_{n-1}(x) - 2*x*S_{n-1}(x) = (S_{n}(x) + S_{n-2} - 2x)/2

instance Fractional a => Num (Spread a) where
  Spread as + Spread bs = Spread (zipWithT (+) id id as bs)
  Spread as - Spread bs = Spread (zipWithT (-) id (map negate) as bs)
  as * bs = fromPower (toPower as * toPower bs)
  fromInteger n = Spread [fromInteger n]
  negate (Spread as) = Spread (map negate as)
  abs = error "borked numerical tower"
  signum = error "borked numerical tower"