{-# LANGUAGE ParallelListComp #-}
module Math.Gamma.Incomplete
    ( lowerGammaCF, pCF
    , lowerGammaHypGeom, lnLowerGammaHypGeom, pHypGeom
    , upperGammaCF, lnUpperGammaConvergents, qCF
    ) where

import {-# SOURCE #-}  Math.Gamma
import Math.ContinuedFraction
import Math.Sequence.Converge

-- |Continued fraction representation of the lower incomplete gamma function.
lowerGammaCF :: (Floating a, Ord a) => a -> a -> Math.ContinuedFraction.CF a
lowerGammaCF :: a -> a -> CF a
lowerGammaCF a
s a
z = a -> [(a, a)] -> CF a
forall a. a -> [(a, a)] -> CF a
gcf a
0
    [ (a
p,a
q)
    | a
p <- a -> a -> a
forall a. (Floating a, Ord a) => a -> a -> a
pow_x_s_div_exp_x a
s a
z
        a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
interleave
            [a -> a
forall a. Num a => a -> a
negate a
spn a -> a -> a
forall a. Num a => a -> a -> a
* a
z | a
spn <- (a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate (a
1a -> a -> a
forall a. Num a => a -> a -> a
+) a
s]
            [a
n a -> a -> a
forall a. Num a => a -> a -> a
* a
z          | a
n   <- (a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate (a
1a -> a -> a
forall a. Num a => a -> a -> a
+) a
1]
    | a
q <- (a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate (a
1a -> a -> a
forall a. Num a => a -> a -> a
+) a
s
    ]

-- |Lower incomplete gamma function, computed using Kummer's confluent
-- hypergeometric function M(a;b;x).  Specifically, this uses the identity:
-- 
-- gamma(s,x) = x**s * exp (-x) / s * M(1; 1+s; x)
-- 
-- From Abramowitz & Stegun (6.5.12).
--
-- Recommended for use when x < s+1
lowerGammaHypGeom :: (Eq b, Floating b) => b -> b -> b
lowerGammaHypGeom :: b -> b -> b
lowerGammaHypGeom b
0 b
0 = b
0b -> b -> b
forall a. Fractional a => a -> a -> a
/b
0
lowerGammaHypGeom b
s b
x = b
x b -> b -> b
forall a. Floating a => a -> a -> a
** b
s b -> b -> b
forall a. Num a => a -> a -> a
* b -> b
forall a. Floating a => a -> a
exp (b -> b
forall a. Num a => a -> a
negate b
x) b -> b -> b
forall a. Fractional a => a -> a -> a
/ b
s b -> b -> b
forall a. Num a => a -> a -> a
* b -> b -> b
forall a. (Eq a, Fractional a) => a -> a -> a
m_1_sp1 b
s b
x

-- |Natural logarithm of lower gamma function, based on the same identity as
-- 'lowerGammaHypGeom' and evaluated carefully to avoid overflow and underflow.
-- Recommended for use when x < s+1
lnLowerGammaHypGeom :: (Eq a, Floating a) => a -> a -> a
lnLowerGammaHypGeom :: a -> a -> a
lnLowerGammaHypGeom a
0 a
0 = a
0a -> a -> a
forall a. Fractional a => a -> a -> a
/a
0
lnLowerGammaHypGeom a
s a
x 
    = a -> a
forall a. Floating a => a -> a
log ((a -> a
forall a. Num a => a -> a
signum a
x)a -> a -> a
forall a. Floating a => a -> a -> a
**a
s a -> a -> a
forall a. Num a => a -> a -> a
* a
sign_m a -> a -> a
forall a. Fractional a => a -> a -> a
/ a -> a
forall a. Num a => a -> a
signum a
s)
    a -> a -> a
forall a. Num a => a -> a -> a
+ a
sa -> a -> a
forall a. Num a => a -> a -> a
*a -> a
forall a. Floating a => a -> a
log (a -> a
forall a. Num a => a -> a
abs a
x) a -> a -> a
forall a. Num a => a -> a -> a
- a
x a -> a -> a
forall a. Num a => a -> a -> a
- a -> a
forall a. Floating a => a -> a
log (a -> a
forall a. Num a => a -> a
abs a
s) a -> a -> a
forall a. Num a => a -> a -> a
+ a
log_m
    where
        (a
sign_m, a
log_m) = a -> a -> (a, a)
forall a. (Eq a, Floating a) => a -> a -> (a, a)
log_m_1_sp1 a
s a
x

-- |Continued fraction representation of the regularized lower incomplete gamma function.
pCF :: (Gamma a, Ord a, Enum a) => a -> a -> CF a
pCF :: a -> a -> CF a
pCF a
s a
x = a -> [(a, a)] -> CF a
forall a. a -> [(a, a)] -> CF a
gcf a
0
    [ (a
p,a
q)
    | a
p <- a -> a -> a
forall a. (Gamma a, Ord a) => a -> a -> a
pow_x_s_div_gamma_s_div_exp_x a
s a
x
        a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
interleave
            [a -> a
forall a. Num a => a -> a
negate a
spn a -> a -> a
forall a. Num a => a -> a -> a
* a
x | a
spn <- [a
s..]]
            [a
n a -> a -> a
forall a. Num a => a -> a -> a
* a
x          | a
n   <- [a
1..]]
    | a
q <- [a
s..]
    ]

-- |Regularized lower incomplete gamma function, computed using Kummer's
-- confluent hypergeometric function.  Uses same identity as 'lowerGammaHypGeom'.
-- 
-- Recommended for use when x < s+1
pHypGeom :: (Gamma a, Ord a) => a -> a -> a
pHypGeom :: a -> a -> a
pHypGeom a
0 a
0 = a
0a -> a -> a
forall a. Fractional a => a -> a -> a
/a
0
pHypGeom a
s a
x
    | a
s a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0
    = a -> a
forall a. Num a => a -> a
signum a
x a -> a -> a
forall a. Floating a => a -> a -> a
** a
s a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
sin (a
forall a. Floating a => a
pia -> a -> a
forall a. Num a => a -> a -> a
*a
s) a -> a -> a
forall a. Fractional a => a -> a -> a
/ (-a
forall a. Floating a => a
pi)
    a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
exp (a
s a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
log (a -> a
forall a. Num a => a -> a
abs a
x) a -> a -> a
forall a. Num a => a -> a -> a
- a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a -> a
forall a. Gamma a => a -> a
lnGamma  (-a
s)) a -> a -> a
forall a. Num a => a -> a -> a
* a -> a -> a
forall a. (Eq a, Fractional a) => a -> a -> a
m_1_sp1 a
s a
x

    | a
s a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 Bool -> Bool -> Bool
|| a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0
    = a
0

    | Bool
otherwise
    = a -> a
forall a. Num a => a -> a
signum a
x a -> a -> a
forall a. Floating a => a -> a -> a
** a
s a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
exp (a
s a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
log (a -> a
forall a. Num a => a -> a
abs a
x) a -> a -> a
forall a. Num a => a -> a -> a
- a
x a -> a -> a
forall a. Num a => a -> a -> a
- a -> a
forall a. Gamma a => a -> a
lnGamma (a
sa -> a -> a
forall a. Num a => a -> a -> a
+a
1)) a -> a -> a
forall a. Num a => a -> a -> a
* a -> a -> a
forall a. (Eq a, Fractional a) => a -> a -> a
m_1_sp1 a
s a
x

-- |Continued fraction representation of the regularized upper incomplete gamma function.
-- Recommended for use when x >= s+1
qCF :: (Gamma a, Ord a, Enum a) => a -> a -> CF a
qCF :: a -> a -> CF a
qCF a
s a
x = a -> [(a, a)] -> CF a
forall a. a -> [(a, a)] -> CF a
gcf a
0
    [ (a
p,a
q)
    | a
p <- a -> a -> a
forall a. (Gamma a, Ord a) => a -> a -> a
pow_x_s_div_gamma_s_div_exp_x a
s a
x
        a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (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
1..] ((a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate (a -> a -> a
forall a. Num a => a -> a -> a
subtract a
1) (a
sa -> a -> a
forall a. Num a => a -> a -> a
-a
1))
    | a
q <- [a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
s | a
n <- [a
1,a
3..]]
    ]

-- |Continued fraction representation of the upper incomplete gamma function.
-- Recommended for use when x >= s+1
upperGammaCF :: (Floating a, Ord a) => a -> a -> CF a
upperGammaCF :: a -> a -> CF a
upperGammaCF a
s a
z = a -> [(a, a)] -> CF a
forall a. a -> [(a, a)] -> CF a
gcf a
0
    [ (a
p,a
q)
    | a
p <- a -> a -> a
forall a. (Floating a, Ord a) => a -> a -> a
pow_x_s_div_exp_x a
s a
z
        a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (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 -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate (a
1a -> a -> a
forall a. Num a => a -> a -> a
+) a
1) ((a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate (a -> a -> a
forall a. Num a => a -> a -> a
subtract a
1) (a
sa -> a -> a
forall a. Num a => a -> a -> a
-a
1))
    | a
q <- [a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
z a -> a -> a
forall a. Num a => a -> a -> a
- a
s | a
n <- (a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate (a
2a -> a -> a
forall a. Num a => a -> a -> a
+) a
1]
    ]

-- |Natural logarithms of the convergents of the upper gamma function, 
-- evaluated carefully to avoid overflow and underflow.
-- Recommended for use when x >= s+1
lnUpperGammaConvergents :: (Eq a, Floating a) => a -> a -> [a]
lnUpperGammaConvergents :: a -> a -> [a]
lnUpperGammaConvergents a
s a
x = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a
a a -> a -> a
forall a. Num a => a -> a -> a
-) ([[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (CF a -> [[a]]
eval CF a
theCF)) 
    where 
        eval :: CF a -> [[a]]
eval = ([(a, a)] -> [a]) -> [[(a, a)]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (((a, a) -> a) -> [(a, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, a) -> a
forall a. Floating a => (a, a) -> a
evalSign) ([[(a, a)]] -> [[a]]) -> (CF a -> [[(a, a)]]) -> CF a -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> (a, a))
-> ((a, a) -> (a, a) -> (a, a))
-> ((a, a) -> (a, a))
-> a
-> CF a
-> [[(a, a)]]
forall a b.
(Fractional a, Eq a) =>
(a -> b) -> (b -> b -> b) -> (b -> b) -> a -> CF a -> [[b]]
modifiedLentzWith a -> (a, a)
forall a. Floating a => a -> (a, a)
signLog (a, a) -> (a, a) -> (a, a)
forall a b. (Num a, Num b) => (a, b) -> (a, b) -> (a, b)
addSignLog (a, a) -> (a, a)
forall b a. Num b => (a, b) -> (a, b)
negateSignLog a
1e-30
        
        a :: a
a = a
s a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
log a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
x
        theCF :: CF a
theCF = a -> [(a, a)] -> CF a
forall a. a -> [(a, a)] -> CF a
gcf (a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
s)
            [ (a
p,a
q)
            | a
p <- (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 -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate (a
1a -> a -> a
forall a. Num a => a -> a -> a
+) a
1) ((a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate (a -> a -> a
forall a. Num a => a -> a -> a
subtract a
1) (a
sa -> a -> a
forall a. Num a => a -> a -> a
-a
1))
            | a
q <- [a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
s | a
n <- (a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate (a
2a -> a -> a
forall a. Num a => a -> a -> a
+) a
3]
            ]

---- various utility functions ----

evalSign :: Floating a => (a,a) -> a
evalSign :: (a, a) -> a
evalSign (a
s,a
x) = a -> a
forall a. Floating a => a -> a
log a
s a -> a -> a
forall a. Num a => a -> a -> a
+ a
x

signLog :: Floating a => a -> (a,a)
signLog :: a -> (a, a)
signLog a
x = (a -> a
forall a. Num a => a -> a
signum a
x, a -> a
forall a. Floating a => a -> a
log (a -> a
forall a. Num a => a -> a
abs a
x))

addSignLog :: (Num a, Num b) => (a,b) -> (a,b) -> (a,b)
addSignLog :: (a, b) -> (a, b) -> (a, b)
addSignLog (a
xS,b
xL) (a
yS,b
yL) = (a
xSa -> a -> a
forall a. Num a => a -> a -> a
*a
yS, b
xLb -> b -> b
forall a. Num a => a -> a -> a
+b
yL)

negateSignLog :: (Num b) => (a,b) -> (a,b)
negateSignLog :: (a, b) -> (a, b)
negateSignLog (a
s,b
l) = (a
s, b -> b
forall a. Num a => a -> a
negate b
l)

-- |Special case of Kummer's confluent hypergeometric function, used
-- in lower gamma functions.
-- 
-- m_1_sp1 s z = M(1;s+1;z)
-- 
m_1_sp1 :: (Eq a, Fractional a) => a -> a -> a
m_1_sp1 :: a -> a -> a
m_1_sp1 a
s a
z = [a] -> a
forall a. Eq a => [a] -> a
converge ([a] -> a) -> ([a] -> [a]) -> [a] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> a) -> a -> [a] -> [a]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl a -> a -> a
forall a. Num a => a -> a -> a
(+) a
0 ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> a) -> a -> [a] -> [a]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl a -> a -> a
forall a. Num a => a -> a -> a
(*) a
1 ([a] -> a) -> [a] -> a
forall a b. (a -> b) -> a -> b
$
    [a
z a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
x | a
x <- (a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate (a
1a -> a -> a
forall a. Num a => a -> a -> a
+) (a
sa -> a -> a
forall a. Num a => a -> a -> a
+a
1)]

log_m_1_sp1 :: (Eq a, Floating a) => a -> a -> (a,a)
log_m_1_sp1 :: a -> a -> (a, a)
log_m_1_sp1 a
s a
z = [(a, a)] -> (a, a)
forall a. Eq a => [a] -> a
converge ([[(a, a)]] -> [(a, a)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (a -> a -> [[(a, a)]]
forall a. (Eq a, Floating a) => a -> a -> [[(a, a)]]
log_m_1_sp1_convergents a
s a
z))

log_m_1_sp1_convergents :: (Eq a, Floating a) => a -> a -> [[(a,a)]]
log_m_1_sp1_convergents :: a -> a -> [[(a, a)]]
log_m_1_sp1_convergents a
s a
z
    = (a -> (a, a))
-> ((a, a) -> (a, a) -> (a, a))
-> ((a, a) -> (a, a))
-> a
-> CF a
-> [[(a, a)]]
forall a b.
(Fractional a, Eq a) =>
(a -> b) -> (b -> b -> b) -> (b -> b) -> a -> CF a -> [[b]]
modifiedLentzWith a -> (a, a)
forall a. Floating a => a -> (a, a)
signLog (a, a) -> (a, a) -> (a, a)
forall a b. (Num a, Num b) => (a, b) -> (a, b) -> (a, b)
addSignLog (a, a) -> (a, a)
forall b a. Num b => (a, b) -> (a, b)
negateSignLog a
1e-30
    (CF a -> [[(a, a)]]) -> CF a -> [[(a, a)]]
forall a b. (a -> b) -> a -> b
$ [a] -> CF a
forall a. Num a => [a] -> CF a
sumPartialProducts (a
1a -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a
z a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
x | a
x <- (a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate (a
1a -> a -> a
forall a. Num a => a -> a -> a
+) (a
sa -> a -> a
forall a. Num a => a -> a -> a
+a
1)])

interleave :: [a] -> [a] -> [a]
interleave :: [a] -> [a] -> [a]
interleave [] [a]
_ = []
interleave [a]
_ [] = []
interleave (a
x:[a]
xs) [a]
ys = a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
interleave [a]
ys [a]
xs

-- A common subexpression appearing in both 'pCF' and 'qCF'.
pow_x_s_div_gamma_s_div_exp_x :: (Gamma a, Ord a) => a -> a -> a
pow_x_s_div_gamma_s_div_exp_x :: a -> a -> a
pow_x_s_div_gamma_s_div_exp_x a
s a
x 
    | a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0     = a -> a
forall a. Floating a => a -> a
exp (a -> a
forall a. Floating a => a -> a
log a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
s a -> a -> a
forall a. Num a => a -> a -> a
- a
x a -> a -> a
forall a. Num a => a -> a -> a
- a -> a
forall a. Gamma a => a -> a
lnGamma a
s)
    | Bool
otherwise = a
x a -> a -> a
forall a. Floating a => a -> a -> a
** a
s a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a -> a
forall a. Floating a => a -> a
exp a
x a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Gamma a => a -> a
gamma a
s)

-- The corresponding subexpression from 'lowerGammaCF' and 'upperGammaCF'
pow_x_s_div_exp_x :: (Floating a, Ord a) => a -> a -> a
pow_x_s_div_exp_x :: a -> a -> a
pow_x_s_div_exp_x a
s a
x 
    | a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0     = a -> a
forall a. Floating a => a -> a
exp (a -> a
forall a. Floating a => a -> a
log a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
s a -> a -> a
forall a. Num a => a -> a -> a
- a
x)
    | Bool
otherwise = a
x a -> a -> a
forall a. Floating a => a -> a -> a
** a
s a -> a -> a
forall a. Fractional a => a -> a -> a
/ a -> a
forall a. Floating a => a -> a
exp a
x