{-# 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 :: forall (v :: * -> *).
(Vector v Double, Vector v Int) =>
v Double -> v Double
autocovariance v Double
a = forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map Int -> Double
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a. (Vector v a, Enum a) => a -> a -> v a
G.enumFromTo Int
0 forall a b. (a -> b) -> a -> b
$ Int
lforall a. Num a => a -> a -> a
-Int
2
  where
    f :: Int -> Double
f Int
k = forall (v :: * -> *). Vector v Double => v Double -> Double
sum (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 forall a. Num a => a -> a -> a
(*) (forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
G.take (Int
lforall a. Num a => a -> a -> a
-Int
k) v Double
c) (forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
Int -> Int -> v a -> v a
G.slice Int
k (Int
lforall a. Num a => a -> a -> a
-Int
k) v Double
c))
          forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
l
    c :: v Double
c   = forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map (forall a. Num a => a -> a -> a
subtract (forall (v :: * -> *). Vector v Double => v Double -> Double
mean v Double
a)) v Double
a
    l :: Int
l   = 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 :: forall (v :: * -> *).
(Vector v Double, Vector v Int) =>
v Double -> (v Double, v Double, v Double)
autocorrelation v Double
a = (v Double
r, forall {a}. (Num a, Vector v a) => (Double -> Double -> a) -> v a
ci (-), forall {a}. (Num a, Vector v a) => (Double -> Double -> a) -> v a
ci forall a. Num a => a -> a -> a
(+))
  where
    r :: v Double
r           = forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map (forall a. Fractional a => a -> a -> a
/ forall (v :: * -> *) a. Vector v a => v a -> a
G.head v Double
c) v Double
c
      where c :: v Double
c   = forall (v :: * -> *).
(Vector v Double, Vector v Int) =>
v Double -> v Double
autocovariance v Double
a
    dllse :: v Double
dllse       = forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map Double -> Double
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a. Vector v a => (a -> a -> a) -> v a -> v a
G.scanl1 forall a. Num a => a -> a -> a
(+) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map Double -> Double
square forall a b. (a -> b) -> a -> b
$ v Double
r
      where f :: Double -> Double
f Double
v = Double
1.96 forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
sqrt ((Double
v forall a. Num a => a -> a -> a
* Double
2 forall a. Num a => a -> a -> a
+ Double
1) forall a. Fractional a => a -> a -> a
/ Double
l)
    l :: Double
l           = forall a b. (Integral a, Num b) => a -> b
fromIntegral (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        = forall (v :: * -> *) a. Vector v a => a -> v a -> v a
G.cons a
1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a. Vector v a => v a -> v a
G.tail forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map (Double -> Double -> a
f (-Double
1forall a. Fractional a => a -> a -> a
/Double
l)) forall a b. (a -> b) -> a -> b
$ v Double
dllse