{-# LANGUAGE BangPatterns              #-}
{-# LANGUAGE ExistentialQuantification #-}
-- |
-- Module    : Numeric.Series
-- Copyright : (c) 2016 Alexey Khudyakov
-- License   : BSD3
--
-- Maintainer  : alexey.skladnoy@gmail.com, bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Functions for working with infinite sequences. In particular
-- summation of series and evaluation of continued fractions.
module Numeric.Series (
    -- * Data type for infinite sequences.
    Sequence(..)
    -- * Constructors
  , enumSequenceFrom
  , enumSequenceFromStep
  , scanSequence
    -- * Summation of series
  , sumSeries
  , sumPowerSeries
  , sequenceToList
    -- * Evaluation of continued fractions
  , evalContFractionB
  ) where

import Control.Applicative
import Data.List (unfoldr)

import Numeric.MathFunctions.Constants (m_epsilon)


----------------------------------------------------------------

-- | Infinite series. It's represented as opaque state and step
--   function.
data Sequence a = forall s. Sequence s (s -> (a,s))

instance Functor Sequence where
  fmap :: forall a b. (a -> b) -> Sequence a -> Sequence b
fmap a -> b
f (Sequence s
s0 s -> (a, s)
step) = s -> (s -> (b, s)) -> Sequence b
forall a s. s -> (s -> (a, s)) -> Sequence a
Sequence s
s0 (\s
s -> let (a
a,s
s') = s -> (a, s)
step s
s in (a -> b
f a
a, s
s'))
  {-# INLINE fmap #-}

instance Applicative Sequence where
  pure :: forall a. a -> Sequence a
pure a
a = () -> (() -> (a, ())) -> Sequence a
forall a s. s -> (s -> (a, s)) -> Sequence a
Sequence () (\() -> (a
a,()))
  Sequence s
sA s -> (a -> b, s)
fA <*> :: forall a b. Sequence (a -> b) -> Sequence a -> Sequence b
<*> Sequence s
sB s -> (a, s)
fB = (s, s) -> ((s, s) -> (b, (s, s))) -> Sequence b
forall a s. s -> (s -> (a, s)) -> Sequence a
Sequence (s
sA,s
sB) (((s, s) -> (b, (s, s))) -> Sequence b)
-> ((s, s) -> (b, (s, s))) -> Sequence b
forall a b. (a -> b) -> a -> b
$ \(!s
sa,!s
sb) ->
    let (a -> b
a,s
sa') = s -> (a -> b, s)
fA s
sa
        (a
b,s
sb') = s -> (a, s)
fB s
sb
    in (a -> b
a a
b, (s
sa',s
sb'))
  {-# INLINE pure  #-}
  {-# INLINE (<*>) #-}

-- | Elementwise operations with sequences
instance Num a => Num (Sequence a) where
  + :: Sequence a -> Sequence a -> Sequence a
(+) = (a -> a -> a) -> Sequence a -> Sequence a -> Sequence a
forall a b c.
(a -> b -> c) -> Sequence a -> Sequence b -> Sequence c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> a -> a
forall a. Num a => a -> a -> a
(+)
  * :: Sequence a -> Sequence a -> Sequence a
(*) = (a -> a -> a) -> Sequence a -> Sequence a -> Sequence a
forall a b c.
(a -> b -> c) -> Sequence a -> Sequence b -> Sequence c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> a -> a
forall a. Num a => a -> a -> a
(*)
  (-) = (a -> a -> a) -> Sequence a -> Sequence a -> Sequence a
forall a b c.
(a -> b -> c) -> Sequence a -> Sequence b -> Sequence c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (-)
  {-# INLINE (+) #-}
  {-# INLINE (*) #-}
  {-# INLINE (-) #-}
  abs :: Sequence a -> Sequence a
abs         = (a -> a) -> Sequence a -> Sequence a
forall a b. (a -> b) -> Sequence a -> Sequence b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
forall a. Num a => a -> a
abs
  signum :: Sequence a -> Sequence a
signum      = (a -> a) -> Sequence a -> Sequence a
forall a b. (a -> b) -> Sequence a -> Sequence b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
forall a. Num a => a -> a
signum
  fromInteger :: Integer -> Sequence a
fromInteger = a -> Sequence a
forall a. a -> Sequence a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Sequence a) -> (Integer -> a) -> Integer -> Sequence a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> a
forall a. Num a => Integer -> a
fromInteger
  {-# INLINE abs         #-}
  {-# INLINE signum      #-}
  {-# INLINE fromInteger #-}

-- | Elementwise operations with sequences
instance Fractional a => Fractional (Sequence a) where
  / :: Sequence a -> Sequence a -> Sequence a
(/)          = (a -> a -> a) -> Sequence a -> Sequence a -> Sequence a
forall a b c.
(a -> b -> c) -> Sequence a -> Sequence b -> Sequence c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> a -> a
forall a. Fractional a => a -> a -> a
(/)
  recip :: Sequence a -> Sequence a
recip        = (a -> a) -> Sequence a -> Sequence a
forall a b. (a -> b) -> Sequence a -> Sequence b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
forall a. Fractional a => a -> a
recip
  fromRational :: Rational -> Sequence a
fromRational = a -> Sequence a
forall a. a -> Sequence a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Sequence a) -> (Rational -> a) -> Rational -> Sequence a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rational -> a
forall a. Fractional a => Rational -> a
fromRational
  {-# INLINE (/)          #-}
  {-# INLINE recip        #-}
  {-# INLINE fromRational #-}



----------------------------------------------------------------
-- Constructors
----------------------------------------------------------------

-- | @enumSequenceFrom x@ generate sequence:
--
-- \[ a_n = x + n \]
enumSequenceFrom :: Num a => a -> Sequence a
enumSequenceFrom :: forall a. Num a => a -> Sequence a
enumSequenceFrom a
i = a -> (a -> (a, a)) -> Sequence a
forall a s. s -> (s -> (a, s)) -> Sequence a
Sequence a
i (\a
n -> (a
n,a
na -> a -> a
forall a. Num a => a -> a -> a
+a
1))
{-# INLINE enumSequenceFrom #-}

-- | @enumSequenceFromStep x d@ generate sequence:
--
-- \[ a_n = x + nd \]
enumSequenceFromStep :: Num a => a -> a -> Sequence a
enumSequenceFromStep :: forall a. Num a => a -> a -> Sequence a
enumSequenceFromStep a
n a
d = a -> (a -> (a, a)) -> Sequence a
forall a s. s -> (s -> (a, s)) -> Sequence a
Sequence a
n (\a
i -> (a
i,a
ia -> a -> a
forall a. Num a => a -> a -> a
+a
d))
{-# INLINE enumSequenceFromStep #-}

-- | Analog of 'scanl' for sequence.
scanSequence :: (b -> a -> b) -> b -> Sequence a -> Sequence b
{-# INLINE scanSequence #-}
scanSequence :: forall b a. (b -> a -> b) -> b -> Sequence a -> Sequence b
scanSequence b -> a -> b
f b
b0 (Sequence s
s0 s -> (a, s)
step) = (b, s) -> ((b, s) -> (b, (b, s))) -> Sequence b
forall a s. s -> (s -> (a, s)) -> Sequence a
Sequence (b
b0,s
s0) (((b, s) -> (b, (b, s))) -> Sequence b)
-> ((b, s) -> (b, (b, s))) -> Sequence b
forall a b. (a -> b) -> a -> b
$ \(b
b,s
s) ->
  let (a
a,s
s') = s -> (a, s)
step s
s
      b' :: b
b'     = b -> a -> b
f b
b a
a
  in (b
b,(b
b',s
s'))


----------------------------------------------------------------
-- Evaluation of series
----------------------------------------------------------------

-- | Calculate sum of series
--
-- \[ \sum_{i=0}^\infty a_i \]
--
-- Summation is stopped when
--
-- \[ a_{n+1} < \varepsilon\sum_{i=0}^n a_i \]
--
-- where ε is machine precision ('m_epsilon')
sumSeries :: Sequence Double -> Double
{-# INLINE sumSeries #-}
sumSeries :: Sequence Double -> Double
sumSeries (Sequence s
sInit s -> (Double, s)
step)
  = Double -> s -> Double
go Double
x0 s
s0
  where 
    (Double
x0,s
s0) = s -> (Double, s)
step s
sInit
    go :: Double -> s -> Double
go Double
x s
s | Double -> Double
forall a. Num a => a -> a
abs (Double
dDouble -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
x) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
m_epsilon = Double -> s -> Double
go Double
x' s
s'
           | Bool
otherwise              = Double
x'
      where (Double
d,s
s') = s -> (Double, s)
step s
s
            x' :: Double
x'     = Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
d

-- | Calculate sum of series
--
-- \[ \sum_{i=0}^\infty x^ia_i \]
--
-- Calculation is stopped when next value in series is less than
-- ε·sum.
sumPowerSeries :: Double -> Sequence Double -> Double
sumPowerSeries :: Double -> Sequence Double -> Double
sumPowerSeries Double
x Sequence Double
ser = Sequence Double -> Double
sumSeries (Sequence Double -> Double) -> Sequence Double -> Double
forall a b. (a -> b) -> a -> b
$ (Double -> Double -> Double)
-> Sequence Double -> Sequence Double -> Sequence Double
forall a b c.
(a -> b -> c) -> Sequence a -> Sequence b -> Sequence c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) ((Double -> Double -> Double)
-> Double -> Sequence Double -> Sequence Double
forall b a. (b -> a -> b) -> b -> Sequence a -> Sequence b
scanSequence Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) Double
1 (Double -> Sequence Double
forall a. a -> Sequence a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
x)) Sequence Double
ser
{-# INLINE sumPowerSeries #-}

-- | Convert series to infinite list
sequenceToList :: Sequence a -> [a]
sequenceToList :: forall a. Sequence a -> [a]
sequenceToList (Sequence s
s s -> (a, s)
f) = (s -> Maybe (a, s)) -> s -> [a]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr ((a, s) -> Maybe (a, s)
forall a. a -> Maybe a
Just ((a, s) -> Maybe (a, s)) -> (s -> (a, s)) -> s -> Maybe (a, s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> (a, s)
f) s
s



----------------------------------------------------------------
-- Evaluation of continued fractions
----------------------------------------------------------------

-- |
-- Evaluate continued fraction using modified Lentz algorithm.
-- Sequence contain pairs (a[i],b[i]) which form following expression:
--
-- \[
-- b_0 + \frac{a_1}{b_1+\frac{a_2}{b_2+\frac{a_3}{b_3 + \cdots}}}
-- \]
--
-- Modified Lentz algorithm is described in Numerical recipes 5.2
-- "Evaluation of Continued Fractions"
evalContFractionB :: Sequence (Double,Double) -> Double
{-# INLINE evalContFractionB #-}
evalContFractionB :: Sequence (Double, Double) -> Double
evalContFractionB (Sequence s
sInit s -> ((Double, Double), s)
step)
  = let ((Double
_,Double
b0),s
s0) = s -> ((Double, Double), s)
step s
sInit
        f0 :: Double
f0          = Double -> Double
maskZero Double
b0
    in  Double -> Double -> Double -> s -> Double
go Double
f0 Double
f0 Double
0 s
s0
  where
    tiny :: Double
tiny = Double
1e-60
    maskZero :: Double -> Double
maskZero Double
0 = Double
tiny
    maskZero Double
x = Double
x
    
    go :: Double -> Double -> Double -> s -> Double
go Double
f Double
c Double
d s
s
      | Double -> Double
forall a. Num a => a -> a
abs (Double
delta Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
m_epsilon = Double -> Double -> Double -> s -> Double
go Double
f' Double
c' Double
d' s
s'
      | Bool
otherwise                    = Double
f'
      where
          ((Double
a,Double
b),s
s') = s -> ((Double, Double), s)
step s
s
          d' :: Double
d'    = Double -> Double
forall a. Fractional a => a -> a
recip (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double -> Double
maskZero (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
aDouble -> Double -> Double
forall a. Num a => a -> a -> a
*Double
d
          c' :: Double
c'    = Double -> Double
maskZero (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
aDouble -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
c 
          delta :: Double
delta = Double
c'Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
d'
          f' :: Double
f'    = Double
fDouble -> Double -> Double
forall a. Num a => a -> a -> a
*Double
delta