{-# LANGUAGE FlexibleContexts #-}
-- |
-- Module    : Statistics.Autocorrelation
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Functions for computing autocovariance and autocorrelation of a
-- sample.

module Statistics.Autocorrelation
    (
      autocovariance
    , autocorrelation
    ) where

import Prelude hiding (sum)
import Statistics.Function (square)
import Statistics.Sample (mean)
import Statistics.Sample.Internal (sum)
import qualified Data.Vector.Generic as G

-- | Compute the autocovariance of a sample, i.e. the covariance of
-- the sample against a shifted version of itself.
autocovariance :: (G.Vector v Double, G.Vector v Int) => v Double -> v Double
autocovariance :: v Double -> v Double
autocovariance v Double
a = (Int -> Double) -> v Int -> v Double
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map Int -> Double
f (v Int -> v Double) -> (Int -> v Int) -> Int -> v Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> v Int
forall (v :: * -> *) a. (Vector v a, Enum a) => a -> a -> v a
G.enumFromTo Int
0 (Int -> v Double) -> Int -> v Double
forall a b. (a -> b) -> a -> b
$ Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
2
  where
    f :: Int -> Double
f Int
k = v Double -> Double
forall (v :: * -> *). Vector v Double => v Double -> Double
sum ((Double -> Double -> Double) -> v Double -> v Double -> v Double
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> v a -> v b -> v c
G.zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) (Int -> v Double -> v Double
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
G.take (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
k) v Double
c) (Int -> Int -> v Double -> v Double
forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> v a
G.slice Int
k (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
k) v Double
c))
          Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
l
    c :: v Double
c   = (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. Num a => a -> a -> a
subtract (v Double -> Double
forall (v :: * -> *). Vector v Double => v Double -> Double
mean v Double
a)) v Double
a
    l :: Int
l   = v Double -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v Double
a

-- | Compute the autocorrelation function of a sample, and the upper
-- and lower bounds of confidence intervals for each element.
--
-- /Note/: The calculation of the 95% confidence interval assumes a
-- stationary Gaussian process.
autocorrelation :: (G.Vector v Double, G.Vector v Int) => v Double -> (v Double, v Double, v Double)
autocorrelation :: v Double -> (v Double, v Double, v Double)
autocorrelation v Double
a = (v Double
r, (Double -> Double -> Double) -> v Double
forall a. (Vector v a, Num a) => (Double -> Double -> a) -> v a
ci (-), (Double -> Double -> Double) -> v Double
forall a. (Vector v a, Num a) => (Double -> Double -> a) -> v a
ci Double -> Double -> Double
forall a. Num a => a -> a -> a
(+))
  where
    r :: v Double
r           = (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
/ v Double -> Double
forall (v :: * -> *) a. Vector v a => v a -> a
G.head v Double
c) v Double
c
      where c :: v Double
c   = v Double -> v Double
forall (v :: * -> *).
(Vector v Double, Vector v Int) =>
v Double -> v Double
autocovariance v Double
a
    dllse :: v Double
dllse       = (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
f (v Double -> v Double)
-> (v Double -> v Double) -> v Double -> v Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> Double -> Double) -> v Double -> v Double
forall (v :: * -> *) a. Vector v a => (a -> a -> a) -> v a -> v a
G.scanl1 Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) (v Double -> v Double)
-> (v Double -> v Double) -> v Double -> v Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (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
square (v Double -> v Double) -> v Double -> v Double
forall a b. (a -> b) -> a -> b
$ v Double
r
      where f :: Double -> Double
f Double
v = Double
1.96 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sqrt ((Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
l)
    l :: Double
l           = 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
a)
    ci :: (Double -> Double -> a) -> v a
ci Double -> Double -> a
f        = a -> v a -> v a
forall (v :: * -> *) a. Vector v a => a -> v a -> v a
G.cons a
1 (v a -> v a) -> (v Double -> v a) -> v Double -> v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v a -> v a
forall (v :: * -> *) a. Vector v a => v a -> v a
G.tail (v a -> v a) -> (v Double -> v a) -> v Double -> v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> a) -> v Double -> v a
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map (Double -> Double -> a
f (-Double
1Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
l)) (v Double -> v a) -> v Double -> v a
forall a b. (a -> b) -> a -> b
$ v Double
dllse