{-# LANGUAGE FlexibleContexts #-}
-- |
-- Module    : Statistics.Test.KolmogorovSmirnov
-- Copyright : (c) 2011 Aleksey Khudyakov
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Kolmogov-Smirnov tests are non-parametric tests for assessing
-- whether given sample could be described by distribution or whether
-- two samples have the same distribution. It's only applicable to
-- continuous distributions.
module Statistics.Test.KolmogorovSmirnov (
    -- * Kolmogorov-Smirnov test
    kolmogorovSmirnovTest
  , kolmogorovSmirnovTestCdf
  , kolmogorovSmirnovTest2
    -- * Evaluate statistics
  , kolmogorovSmirnovCdfD
  , kolmogorovSmirnovD
  , kolmogorovSmirnov2D
    -- * Probabilities
  , kolmogorovSmirnovProbability
    -- * References
    -- $references
  , module Statistics.Test.Types
  ) where

import Control.Monad (when)
import Prelude hiding (exponent, sum)
import Statistics.Distribution (Distribution(..))
import Statistics.Function (gsort, unsafeModify)
import Statistics.Matrix (center, for, fromVector)
import qualified Statistics.Matrix as Mat
import Statistics.Test.Types
import Statistics.Types (mkPValue)
import qualified Data.Vector          as V
import qualified Data.Vector.Storable as S
import qualified Data.Vector.Unboxed  as U
import qualified Data.Vector.Generic  as G
import           Data.Vector.Generic    ((!))
import qualified Data.Vector.Unboxed.Mutable as M


----------------------------------------------------------------
-- Test
----------------------------------------------------------------

-- | Check that sample could be described by distribution. Returns
--   @Nothing@ is sample is empty
--
--   This test uses Marsaglia-Tsang-Wang exact algorithm for
--   calculation of p-value.
kolmogorovSmirnovTest :: (Distribution d, G.Vector v Double)
                      => d        -- ^ Distribution
                      -> v Double -- ^ Data sample
                      -> Maybe (Test ())
{-# INLINE kolmogorovSmirnovTest #-}
kolmogorovSmirnovTest :: forall d (v :: * -> *).
(Distribution d, Vector v Double) =>
d -> v Double -> Maybe (Test ())
kolmogorovSmirnovTest d
d
  = (Double -> Double) -> v Double -> Maybe (Test ())
forall (v :: * -> *).
Vector v Double =>
(Double -> Double) -> v Double -> Maybe (Test ())
kolmogorovSmirnovTestCdf (d -> Double -> Double
forall d. Distribution d => d -> Double -> Double
cumulative d
d)


-- | Variant of 'kolmogorovSmirnovTest' which uses CDF in form of
--   function.
kolmogorovSmirnovTestCdf :: (G.Vector v Double)
                         => (Double -> Double) -- ^ CDF of distribution
                         -> v Double           -- ^ Data sample
                         -> Maybe (Test ())
{-# INLINE kolmogorovSmirnovTestCdf #-}
kolmogorovSmirnovTestCdf :: forall (v :: * -> *).
Vector v Double =>
(Double -> Double) -> v Double -> Maybe (Test ())
kolmogorovSmirnovTestCdf Double -> Double
cdf v Double
sample
  | v Double -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v Double
sample = Maybe (Test ())
forall a. Maybe a
Nothing
  | Bool
otherwise     = Test () -> Maybe (Test ())
forall a. a -> Maybe a
Just Test
      { testSignificance :: PValue Double
testSignificance = Double -> PValue Double
forall a. (Ord a, Num a) => a -> PValue a
mkPValue (Double -> PValue Double) -> Double -> PValue Double
forall a b. (a -> b) -> a -> b
$ Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
prob
      , testStatistics :: Double
testStatistics   = Double
d
      , testDistribution :: ()
testDistribution = ()
      }
  where
    d :: Double
d    = (Double -> Double) -> v Double -> Double
forall (v :: * -> *).
Vector v Double =>
(Double -> Double) -> v Double -> Double
kolmogorovSmirnovCdfD Double -> Double
cdf v Double
sample
    prob :: Double
prob = Int -> Double -> Double
kolmogorovSmirnovProbability (v Double -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v Double
sample) Double
d


-- | Two sample Kolmogorov-Smirnov test. It tests whether two data
--   samples could be described by the same distribution without
--   making any assumptions about it. If either of samples is empty
--   returns Nothing.
--
--   This test uses approximate formula for computing p-value.
kolmogorovSmirnovTest2 :: (G.Vector v Double)
                       => v Double -- ^ Sample 1
                       -> v Double -- ^ Sample 2
                       -> Maybe (Test ())
kolmogorovSmirnovTest2 :: forall (v :: * -> *).
Vector v Double =>
v Double -> v Double -> Maybe (Test ())
kolmogorovSmirnovTest2 v Double
xs1 v Double
xs2
  | v Double -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v Double
xs1 Bool -> Bool -> Bool
|| v Double -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v Double
xs2 = Maybe (Test ())
forall a. Maybe a
Nothing
  | Bool
otherwise                = Test () -> Maybe (Test ())
forall a. a -> Maybe a
Just Test
      { testSignificance :: PValue Double
testSignificance = Double -> PValue Double
forall a. (Ord a, Num a) => a -> PValue a
mkPValue (Double -> PValue Double) -> Double -> PValue Double
forall a b. (a -> b) -> a -> b
$ Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall {a}. (Ord a, Floating a) => a -> a
prob Double
d
      , testStatistics :: Double
testStatistics   = Double
d
      , testDistribution :: ()
testDistribution = ()
      }
  where
    d :: Double
d    = v Double -> v Double -> Double
forall (v :: * -> *).
Vector v Double =>
v Double -> v Double -> Double
kolmogorovSmirnov2D v Double
xs1 v Double
xs2
         Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
en Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.12 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.11Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
en)
    -- Effective number of data points
    n1 :: Double
n1   = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (v Double -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v Double
xs1)
    n2 :: Double
n2   = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (v Double -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v Double
xs2)
    en :: Double
en   = Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double
n1 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
n2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
n1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
n2)
    --
    prob :: a -> a
prob a
z
      | a
z a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<  a
0    = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"kolmogorovSmirnov2D: internal error"
      | a
z a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0    = a
0
      | a
z a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<  a
1.18 = let y :: a
y = a -> a
forall a. Floating a => a -> a
exp( -a
1.23370055013616983 a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
za -> a -> a
forall a. Num a => a -> a -> a
*a
z) )
                    in  a
2.25675833419102515 a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
sqrt( -a -> a
forall a. Floating a => a -> a
log a
y ) a -> a -> a
forall a. Num a => a -> a -> a
* (a
y a -> a -> a
forall a. Num a => a -> a -> a
+ a
ya -> a -> a
forall a. Floating a => a -> a -> a
**a
9 a -> a -> a
forall a. Num a => a -> a -> a
+ a
ya -> a -> a
forall a. Floating a => a -> a -> a
**a
25 a -> a -> a
forall a. Num a => a -> a -> a
+ a
ya -> a -> a
forall a. Floating a => a -> a -> a
**a
49)
      | Bool
otherwise = let x :: a
x = a -> a
forall a. Floating a => a -> a
exp(-a
2 a -> a -> a
forall a. Num a => a -> a -> a
* a
z a -> a -> a
forall a. Num a => a -> a -> a
* a
z)
                    in  a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
2a -> a -> a
forall a. Num a => a -> a -> a
*(a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
xa -> a -> a
forall a. Floating a => a -> a -> a
**a
4 a -> a -> a
forall a. Num a => a -> a -> a
+ a
xa -> a -> a
forall a. Floating a => a -> a -> a
**a
9)
{-# INLINABLE  kolmogorovSmirnovTest2 #-}
{-# SPECIALIZE kolmogorovSmirnovTest2 :: U.Vector Double -> U.Vector Double -> Maybe (Test ()) #-}
{-# SPECIALIZE kolmogorovSmirnovTest2 :: V.Vector Double -> V.Vector Double -> Maybe (Test ()) #-}
{-# SPECIALIZE kolmogorovSmirnovTest2 :: S.Vector Double -> S.Vector Double -> Maybe (Test ()) #-}
-- FIXME: Find source for approximation for D



----------------------------------------------------------------
-- Kolmogorov's statistic
----------------------------------------------------------------

-- | Calculate Kolmogorov's statistic /D/ for given cumulative
--   distribution function (CDF) and data sample. If sample is empty
--   returns 0.
kolmogorovSmirnovCdfD :: G.Vector v Double
                      => (Double -> Double) -- ^ CDF function
                      -> v Double           -- ^ Sample
                      -> Double
kolmogorovSmirnovCdfD :: forall (v :: * -> *).
Vector v Double =>
(Double -> Double) -> v Double -> Double
kolmogorovSmirnovCdfD Double -> Double
cdf v Double
sample
  | v Double -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v Double
sample = Double
0
  | Bool
otherwise     = v Double -> Double
forall (v :: * -> *) a. (Vector v a, Ord a) => v a -> a
G.maximum
                  (v Double -> Double) -> v Double -> Double
forall a b. (a -> b) -> a -> b
$ (Double -> Double -> Double -> Double)
-> v Double -> v Double -> v Double -> v Double
forall (v :: * -> *) a b c d.
(Vector v a, Vector v b, Vector v c, Vector v d) =>
(a -> b -> c -> d) -> v a -> v b -> v c -> v d
G.zipWith3 (\Double
p Double
a Double
b -> Double -> Double
forall a. Num a => a -> a
abs (Double
pDouble -> Double -> Double
forall a. Num a => a -> a -> a
-Double
a) Double -> Double -> Double
forall a. Ord a => a -> a -> a
`max` Double -> Double
forall a. Num a => a -> a
abs (Double
pDouble -> Double -> Double
forall a. Num a => a -> a -> a
-Double
b))
                    v Double
ps v Double
steps (v Double -> v Double
forall (v :: * -> *) a. Vector v a => v a -> v a
G.tail v Double
steps)
  where
    xs :: v Double
xs = v Double -> v Double
forall e (v :: * -> *). (Ord e, Vector v e) => v e -> v e
gsort v Double
sample
    n :: Int
n  = v Double -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v Double
xs
    --
    ps :: v Double
ps    = (Double -> Double) -> v Double -> v Double
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map Double -> Double
cdf v Double
xs
    steps :: v Double
steps = (Double -> Double) -> v Double -> v Double
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)
          (v Double -> v Double) -> v Double -> v Double
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> Double) -> v Double
forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
G.generate (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINABLE  kolmogorovSmirnovCdfD #-}
{-# SPECIALIZE kolmogorovSmirnovCdfD :: (Double -> Double) -> U.Vector Double -> Double #-}
{-# SPECIALIZE kolmogorovSmirnovCdfD :: (Double -> Double) -> V.Vector Double -> Double #-}
{-# SPECIALIZE kolmogorovSmirnovCdfD :: (Double -> Double) -> S.Vector Double -> Double #-}


-- | Calculate Kolmogorov's statistic /D/ for given cumulative
--   distribution function (CDF) and data sample. If sample is empty
--   returns 0.
kolmogorovSmirnovD :: (Distribution d, G.Vector v Double)
                   => d         -- ^ Distribution
                   -> v Double  -- ^ Sample
                   -> Double
kolmogorovSmirnovD :: forall d (v :: * -> *).
(Distribution d, Vector v Double) =>
d -> v Double -> Double
kolmogorovSmirnovD d
d = (Double -> Double) -> v Double -> Double
forall (v :: * -> *).
Vector v Double =>
(Double -> Double) -> v Double -> Double
kolmogorovSmirnovCdfD (d -> Double -> Double
forall d. Distribution d => d -> Double -> Double
cumulative d
d)
{-# INLINE kolmogorovSmirnovD #-}


-- | Calculate Kolmogorov's statistic /D/ for two data samples. If
--   either of samples is empty returns 0.
kolmogorovSmirnov2D :: (G.Vector v Double)
                    => v Double   -- ^ First sample
                    -> v Double   -- ^ Second sample
                    -> Double
kolmogorovSmirnov2D :: forall (v :: * -> *).
Vector v Double =>
v Double -> v Double -> Double
kolmogorovSmirnov2D v Double
sample1 v Double
sample2
  | v Double -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v Double
sample1 Bool -> Bool -> Bool
|| v Double -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v Double
sample2 = Double
0
  | Bool
otherwise                        = Double -> Int -> Int -> Double
worker Double
0 Int
0 Int
0
  where
    xs1 :: v Double
xs1 = v Double -> v Double
forall e (v :: * -> *). (Ord e, Vector v e) => v e -> v e
gsort v Double
sample1
    xs2 :: v Double
xs2 = v Double -> v Double
forall e (v :: * -> *). (Ord e, Vector v e) => v e -> v e
gsort v Double
sample2
    n1 :: Int
n1  = v Double -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v Double
xs1
    n2 :: Int
n2  = v Double -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v Double
xs2
    en1 :: Double
en1 = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n1
    en2 :: Double
en2 = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n2
    -- Find new index
    skip :: a -> Int -> v a -> Int
skip a
x Int
i v a
xs = Int -> Int
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
      where go :: Int -> Int
go Int
n | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
xs = Int
n
                 | v a
xs v a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
! Int
n a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x      = Int -> Int
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                 | Bool
otherwise        = Int
n
    -- Main loop
    worker :: Double -> Int -> Int -> Double
worker Double
d Int
i1 Int
i2
      | Int
i1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n1 Bool -> Bool -> Bool
|| Int
i2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n2 = Double
d
      | Bool
otherwise            = Double -> Int -> Int -> Double
worker Double
d' Int
i1' Int
i2'
      where
        d1 :: Double
d1  = v Double
xs1 v Double -> Int -> Double
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
! Int
i1
        d2 :: Double
d2  = v Double
xs2 v Double -> Int -> Double
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
! Int
i2
        i1' :: Int
i1' | Double
d1 Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
d2  = Double -> Int -> v Double -> Int
forall {v :: * -> *} {a}.
(Vector v a, Eq a) =>
a -> Int -> v a -> Int
skip Double
d1 Int
i1 v Double
xs1
            | Bool
otherwise = Int
i1
        i2' :: Int
i2' | Double
d2 Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
d1  = Double -> Int -> v Double -> Int
forall {v :: * -> *} {a}.
(Vector v a, Eq a) =>
a -> Int -> v a -> Int
skip Double
d2 Int
i2 v Double
xs2
            | Bool
otherwise = Int
i2
        d' :: Double
d'  = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
d (Double -> Double
forall a. Num a => a -> a
abs (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i1' Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
en1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i2' Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
en2)
{-# INLINABLE  kolmogorovSmirnov2D #-}
{-# SPECIALIZE kolmogorovSmirnov2D :: U.Vector Double -> U.Vector Double -> Double #-}
{-# SPECIALIZE kolmogorovSmirnov2D :: V.Vector Double -> V.Vector Double -> Double #-}
{-# SPECIALIZE kolmogorovSmirnov2D :: S.Vector Double -> S.Vector Double -> Double #-}



-- | Calculate cumulative probability function for Kolmogorov's
--   distribution with /n/ parameters or probability of getting value
--   smaller than /d/ with n-elements sample.
--
--   It uses algorithm by Marsgalia et. al. and provide at least
--   7-digit accuracy.
kolmogorovSmirnovProbability :: Int    -- ^ Size of the sample
                             -> Double -- ^ D value
                             -> Double
kolmogorovSmirnovProbability :: Int -> Double -> Double
kolmogorovSmirnovProbability Int
n Double
d
  -- Avoid potentially lengthy calculations for large N and D > 0.999
  | Double
s Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
7.24 Bool -> Bool -> Bool
|| (Double
s Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
3.76 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
99) = Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
exp( -(Double
2.000071 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.331 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt Double
n' Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1.409 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
n') Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
s)
  -- Exact computation
  | Bool
otherwise = KSMatrix -> Double
fini (KSMatrix -> Double) -> KSMatrix -> Double
forall a b. (a -> b) -> a -> b
$ Int -> Matrix -> KSMatrix
KSMatrix Int
0 Matrix
matrix KSMatrix -> Int -> KSMatrix
`power` Int
n
  where
    s :: Double
s  = Double
n' Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
d Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
d
    n' :: Double
n' = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n

    size :: Int
size = Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    k :: Int
k    = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double
n' Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
d) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    h :: Double
h    = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
n' Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
d
    -- Calculate initial matrix
    matrix :: Matrix
matrix =
      let m :: Vector Double
m = (forall s. ST s (MVector s Double)) -> Vector Double
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
U.create ((forall s. ST s (MVector s Double)) -> Vector Double)
-> (forall s. ST s (MVector s Double)) -> Vector Double
forall a b. (a -> b) -> a -> b
$ do
            MVector s Double
mat <- Int -> ST s (MVector (PrimState (ST s)) Double)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
M.new (Int
sizeInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
size)
            -- Fill matrix with 0 and 1s
            Int -> Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *).
Monad m =>
Int -> Int -> (Int -> m ()) -> m ()
for Int
0 Int
size ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
row ->
              Int -> Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *).
Monad m =>
Int -> Int -> (Int -> m ()) -> m ()
for Int
0 Int
size ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
col -> do
                let val :: Double
val | Int
row Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
col = Double
1
                        | Bool
otherwise      = Double
0 :: Double
                MVector (PrimState (ST s)) Double -> Int -> Double -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
M.write MVector s Double
MVector (PrimState (ST s)) Double
mat (Int
row Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
col) Double
val
            -- Correct left column/bottom row
            Int -> Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *).
Monad m =>
Int -> Int -> (Int -> m ()) -> m ()
for Int
0 Int
size ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
              let delta :: Double
delta = Double
h Double -> Int -> Double
forall a b. (Fractional a, Integral b) => a -> b -> a
^^ (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
              MVector s Double -> Int -> (Double -> Double) -> ST s ()
forall s. MVector s Double -> Int -> (Double -> Double) -> ST s ()
unsafeModify MVector s Double
mat (Int
i    Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
size)         (Double -> Double -> Double
forall a. Num a => a -> a -> a
subtract Double
delta)
              MVector s Double -> Int -> (Double -> Double) -> ST s ()
forall s. MVector s Double -> Int -> (Double -> Double) -> ST s ()
unsafeModify MVector s Double
mat (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) (Double -> Double -> Double
forall a. Num a => a -> a -> a
subtract Double
delta)
            -- Correct corner element if needed
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Double
2Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
h Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
              MVector s Double -> Int -> (Double -> Double) -> ST s ()
forall s. MVector s Double -> Int -> (Double -> Double) -> ST s ()
unsafeModify MVector s Double
mat ((Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
size) (Double -> Double -> Double
forall a. Num a => a -> a -> a
+ ((Double
2Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
h Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
size))
            -- Divide diagonals by factorial
            let divide :: Double -> Int -> ST s ()
divide Double
g Int
num
                  | Int
num Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
size = () -> ST s ()
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
                  | Bool
otherwise   = do Int -> Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *).
Monad m =>
Int -> Int -> (Int -> m ()) -> m ()
for Int
num Int
size ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                                       MVector s Double -> Int -> (Double -> Double) -> ST s ()
forall s. MVector s Double -> Int -> (Double -> Double) -> ST s ()
unsafeModify MVector s Double
mat (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
num) (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
g)
                                     Double -> Int -> ST s ()
divide (Double
g Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
numInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
2)) (Int
numInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
            Double -> Int -> ST s ()
divide Double
2 Int
1
            MVector s Double -> ST s (MVector s Double)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return MVector s Double
mat
      in Int -> Int -> Vector Double -> Matrix
fromVector Int
size Int
size Vector Double
m
    -- Last calculation
    fini :: KSMatrix -> Double
fini (KSMatrix Int
e Matrix
m) = Int -> Double -> Int -> Double
forall {t} {t}.
(Integral t, Ord t, Fractional t) =>
Int -> t -> t -> t
loop Int
1 (Matrix -> Double
center Matrix
m) Int
e
      where
        loop :: Int -> t -> t -> t
loop Int
i t
ss t
eQ
          | Int
i  Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n       = t
ss t -> t -> t
forall a. Num a => a -> a -> a
* t
10 t -> t -> t
forall a b. (Fractional a, Integral b) => a -> b -> a
^^ t
eQ
          | t
ss' t -> t -> Bool
forall a. Ord a => a -> a -> Bool
< t
1e-140 = Int -> t -> t -> t
loop (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (t
ss' t -> t -> t
forall a. Num a => a -> a -> a
* t
1e140) (t
eQ t -> t -> t
forall a. Num a => a -> a -> a
- t
140)
          | Bool
otherwise    = Int -> t -> t -> t
loop (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)  t
ss'           t
eQ
          where ss' :: t
ss' = t
ss t -> t -> t
forall a. Num a => a -> a -> a
* Int -> t
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i t -> t -> t
forall a. Fractional a => a -> a -> a
/ Int -> t
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n

data KSMatrix = KSMatrix Int Mat.Matrix


multiply :: KSMatrix -> KSMatrix -> KSMatrix
multiply :: KSMatrix -> KSMatrix -> KSMatrix
multiply (KSMatrix Int
e1 Matrix
m1) (KSMatrix Int
e2 Matrix
m2) = Int -> Matrix -> KSMatrix
KSMatrix (Int
e1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
e2) (Matrix -> Matrix -> Matrix
Mat.multiply Matrix
m1 Matrix
m2)

power :: KSMatrix -> Int -> KSMatrix
power :: KSMatrix -> Int -> KSMatrix
power KSMatrix
mat Int
1 = KSMatrix
mat
power KSMatrix
mat Int
n = KSMatrix -> KSMatrix
avoidOverflow KSMatrix
res
  where
    mat2 :: KSMatrix
mat2 = KSMatrix -> Int -> KSMatrix
power KSMatrix
mat (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
2)
    pow :: KSMatrix
pow  = KSMatrix -> KSMatrix -> KSMatrix
multiply KSMatrix
mat2 KSMatrix
mat2
    res :: KSMatrix
res | Int -> Bool
forall a. Integral a => a -> Bool
odd Int
n     = KSMatrix -> KSMatrix -> KSMatrix
multiply KSMatrix
pow KSMatrix
mat
        | Bool
otherwise = KSMatrix
pow

avoidOverflow :: KSMatrix -> KSMatrix
avoidOverflow :: KSMatrix -> KSMatrix
avoidOverflow ksm :: KSMatrix
ksm@(KSMatrix Int
e Matrix
m)
  | Matrix -> Double
center Matrix
m Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1e140 = Int -> Matrix -> KSMatrix
KSMatrix (Int
e Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
140) ((Double -> Double) -> Matrix -> Matrix
Mat.map (Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
1e-140) Matrix
m)
  | Bool
otherwise        = KSMatrix
ksm


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

-- $references
--
-- * G. Marsaglia, W. W. Tsang, J. Wang (2003) Evaluating Kolmogorov's
--   distribution, Journal of Statistical Software, American
--   Statistical Association, vol. 8(i18).