Hypergeometric.hs

-- | See Shaw, Ewart [\"Hypergeometric Functions and CDFs in J\"](https://www.jsoftware.com/papers/jhyper.pdf).
module Math.Hypergeometric ( hypergeometric
                           , euler
                           , erf
                           , ncdf
                           ) where

import           Data.Functor ((<$>))

risingFactorial :: Num a => a -> Int -> a
risingFactorial _ 0 = 1
risingFactorial a n = (a + fromIntegral n - 1) * risingFactorial a (n-1)

factorial :: Num a => Int -> a
factorial n = product (fromIntegral <$> [1..n])

{-# SPECIALIZE ncdf :: Double -> Double #-}
{-# SPECIALIZE ncdf :: Float -> Float #-}
-- | CDF of the standard normal \( N(0,1) \)
ncdf :: (Ord a, Floating a) => a -> a
ncdf z = (1/2) * (1 + erf (z / sqrt 2))

{-# SPECIALIZE erf :: Double -> Double #-}
{-# SPECIALIZE erf :: Float -> Float #-}
-- | [erf](https://mathworld.wolfram.com/Erf.html)
erf :: (Ord a, Floating a) => a -> a
erf z = (2 * z * exp (-(z^(2::Int))) / sqrt pi) * hypergeometric [1] [3/2] (z^(2::Int))

{-# SPECIALIZE euler :: Double -> Double -> Double -> Double -> Double #-}
{-# SPECIALIZE euler :: Float -> Float -> Float -> Float -> Float #-}
-- | Euler's transform:
--
-- \( \displaystyle _2F_1(a,b;c;z) = (1-z)^{-a} {}_2F_1\left(a,c-b;c;\frac{z}{z-1}\right) \)
--
-- Koekoek, Roelef and Swarttouw, René F. [The Askey-scheme of hypergeometric orthogonal polynomials and its q-analogue](https://arxiv.org/abs/math/9602214).
--
-- @since 0.1.5.0
euler :: (Ord a, Floating a)
      => a -- ^ \(a\)
      -> a -- ^ \(b\)
      -> a -- ^ \(c\)
      -> a -- ^ \(z\)
      -> a
euler a b c z = hypergeometric [a, c-b] [c] (z/z-1)

{-# SPECIALIZE hypergeometric :: [Double] -> [Double] -> Double -> Double #-}
{-# SPECIALIZE hypergeometric :: [Float] -> [Float] -> Float -> Float #-}
-- | \( _pF_q(a_1,\ldots,a_p;b_1,\ldots,b_q;z) = \displaystyle\sum_{n=0}^\infty\frac{(a_1)_n\cdots(a_p)_n}{(b_1)_b\cdots(b_q)_n}\frac{z^n}{n!} \)
--
-- The radius of convergence is
--
-- \( \rho = \begin{cases} \infty & \text{if} & p<q+1 \\ 1 & \text{if} & p=q+1 \\ 0 & \text{if} & p>q+1 \\ \end{cases} \)
--
-- This iterates until the result stabilizes.
hypergeometric :: (Ord a, Fractional a)
               => [a] -- ^ \( a_1,\ldots,a_p \)
               -> [a] -- ^ \( b_1,\ldots,b_q \)
               -> a -- ^ \( z \)
               -> a
hypergeometric as bs z = sumUntilEq
    [ (product (fmap (`risingFactorial` n) as) / product (fmap (`risingFactorial` n) bs)) * (zϵ ^ n) / factorial n | n <- [0..] ]
  where
    zϵ = 𝜌 z
    𝜌 = case compare (length as) (length bs+1) of
        LT -> id
        EQ -> if z > 1 then error "Outside the radius of convergence." else id
        GT -> if z /= 0 then error "Outside the radius of convergence." else id

sumUntilEq :: (Eq a, Num a) => [a] -> a
sumUntilEq = sumUntilEqLoop 0

sumUntilEqLoop :: (Eq a, Num a) => a -> [a] -> a
sumUntilEqLoop acc (x:y:xs) =
    if step0 == step1
        then step0
        else sumUntilEqLoop step1 xs
    where step0 = acc + x
          step1 = acc + x + y