{-# LANGUAGE ParallelListComp #-}
module Math.Polynomial.Lagrange
    ( lagrangeBasis
    , lagrange
    , lagrangeWeights
    ) where

import Math.Polynomial

-- given a list, return one list containing each element of the original list
-- paired with all the other elements of the list.
select :: [a] -> [(a,[a])]
select :: [a] -> [(a, [a])]
select [] = []
select (x :: a
x:xs :: [a]
xs) = (a
x, [a]
xs) (a, [a]) -> [(a, [a])] -> [(a, [a])]
forall a. a -> [a] -> [a]
: [(a
y, a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys) | (y :: a
y, ys :: [a]
ys) <- [a] -> [(a, [a])]
forall a. [a] -> [(a, [a])]
select [a]
xs]

-- |Returns the Lagrange basis set of polynomials associated with a set of 
-- points. This is the set of polynomials each of which is @1@ at its 
-- corresponding point in the input list and @0@ at all others.
--
-- These polynomials are especially convenient, mathematically, for 
-- interpolation.  The interpolating polynomial for a set of points  @(x,y)@ 
-- is given by using the @y@s as coefficients for the basis given by 
-- @lagrangeBasis xs@.  Computationally, this is not an especially stable 
-- procedure though.  'Math.Polynomial.Interpolation.lagrangePolyFit'
-- implements a slightly better algorithm based on the same idea.  
-- 
-- Generally it is better to not compute the coefficients at all.  
-- 'Math.Polynomial.Interpolation.polyInterp' evaluates the interpolating
-- polynomial directly, and is both quicker and more stable than any method
-- I know of that computes the coefficients.
lagrangeBasis :: (Fractional a, Eq a) => [a] -> [Poly a]
lagrangeBasis :: [a] -> [Poly a]
lagrangeBasis xs :: [a]
xs =
    [ (Poly a -> Poly a -> Poly a) -> [Poly a] -> Poly a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 Poly a -> Poly a -> Poly a
forall a. (Num a, Eq a) => Poly a -> Poly a -> Poly a
multPoly
        [ if a
q a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= 0
            then Endianness -> [a] -> Poly a
forall a. (Num a, Eq a) => Endianness -> [a] -> Poly a
poly Endianness
LE [a -> a
forall a. Num a => a -> a
negate a
x_ja -> a -> a
forall a. Fractional a => a -> a -> a
/a
q, 1a -> a -> a
forall a. Fractional a => a -> a -> a
/a
q]
            else [Char] -> Poly a
forall a. HasCallStack => [Char] -> a
error ("lagrangeBasis: duplicate root")
        | a
x_j <- [a]
otherXs
        , let q :: a
q = a
x_i a -> a -> a
forall a. Num a => a -> a -> a
- a
x_j
        ]
    | (x_i :: a
x_i, otherXs :: [a]
otherXs) <- [a] -> [(a, [a])]
forall a. [a] -> [(a, [a])]
select [a]
xs
    ]

-- |Construct the Lagrange "master polynomial" for the Lagrange barycentric form:
-- That is, the monic polynomial with a root at each point in the input list.
lagrange :: (Num a, Eq a) => [a] -> Poly a
lagrange :: [a] -> Poly a
lagrange [] = Poly a
forall a. (Num a, Eq a) => Poly a
one
lagrange xs :: [a]
xs = (Poly a -> Poly a -> Poly a) -> [Poly a] -> Poly a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 Poly a -> Poly a -> Poly a
forall a. (Num a, Eq a) => Poly a -> Poly a -> Poly a
multPoly
    [ Endianness -> [a] -> Poly a
forall a. (Num a, Eq a) => Endianness -> [a] -> Poly a
poly Endianness
LE [a -> a
forall a. Num a => a -> a
negate a
x_i, 1]
    | a
x_i <- [a]
xs
    ]

-- |Compute the weights associated with each abscissa in the Lagrange
-- barycentric form.
lagrangeWeights :: Fractional a => [a] -> [a]
lagrangeWeights :: [a] -> [a]
lagrangeWeights xs :: [a]
xs = 
    [ a -> a
forall a. Fractional a => a -> a
recip (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$ [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product
        [ a
x_i a -> a -> a
forall a. Num a => a -> a -> a
- a
x_j
        | a
x_j <- [a]
otherXs
        ]
    | (x_i :: a
x_i, otherXs :: [a]
otherXs) <- [a] -> [(a, [a])]
forall a. [a] -> [(a, [a])]
select [a]
xs
    ]