-- SPDX-License-Identifier: MPL-2.0

{- | Bounded Double forecasting models. AR fits use checked QR least squares.
ARMA/SARIMA minimize conditional sum of squares, with presample observations
and innovations set to zero, using analytic innovation derivatives and damped
Levenberg--Marquardt. These are local conditional fits, not exact Gaussian
maximum likelihood; stationarity/invertibility and prediction intervals are
not inferred. Differencing is integrated explicitly when forecasting.
-}
module HMath.TimeSeries.Models (
    Smoothing (..),
    smoothForecast,
    Decomposition (..),
    decompose,
    STLConfig (..),
    stl,
    defaultSTLConfig,
    ARModel,
    fitAR,
    arCoefficients,
    forecastAR,
    SARIMAOrder (..),
    ARIMAModel,
    fitARIMA,
    fitMA,
    fitARMA,
    fitSARIMA,
    arimaParameters,
    arimaResiduals,
    arimaInnovationVariance,
    forecastARIMA,
    VARModel,
    fitVAR,
    varCoefficients,
    forecastVAR,
    kalmanSeries,
) where

import Control.Monad (foldM)
import Data.List (sort, transpose)
import HMath.Error
import qualified HMath.Matrix as M
import qualified HMath.Matrix.QR as QR
import qualified HMath.Optimization.Common as O
import qualified HMath.Optimization.Local as O
import qualified HMath.Signal as S
import HMath.TimeSeries.Core

{- | Exponential smoothing: Simple alpha; Holt alpha beta; HoltWinters period alpha beta
gamma. Rates lie in [0,1] and seasonal fitting requires complete initial cycles.
-}
data Smoothing = Simple Double | Holt Double Double | HoltWinters Int Double Double Double
    deriving (Smoothing -> Smoothing -> Bool
(Smoothing -> Smoothing -> Bool)
-> (Smoothing -> Smoothing -> Bool) -> Eq Smoothing
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Smoothing -> Smoothing -> Bool
== :: Smoothing -> Smoothing -> Bool
$c/= :: Smoothing -> Smoothing -> Bool
/= :: Smoothing -> Smoothing -> Bool
Eq, Int -> Smoothing -> ShowS
[Smoothing] -> ShowS
Smoothing -> String
(Int -> Smoothing -> ShowS)
-> (Smoothing -> String)
-> ([Smoothing] -> ShowS)
-> Show Smoothing
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Smoothing -> ShowS
showsPrec :: Int -> Smoothing -> ShowS
$cshow :: Smoothing -> String
show :: Smoothing -> String
$cshowList :: [Smoothing] -> ShowS
showList :: [Smoothing] -> ShowS
Show)

{- | Fitted one-step means and future means. Simple/Holt initialize from the
first sample/first difference. Additive Holt--Winters initializes the level
and seasonal indices from a complete first cycle, and the trend from the
difference between the first two cycle means. Initial fitted cycle is the
initializer, not an out-of-sample prediction.
-}
smoothForecast :: Smoothing -> Int -> [Double] -> Either MathError ([Double], [Double])
smoothForecast :: Smoothing
-> Int -> [Double] -> Either MathError ([Double], [Double])
smoothForecast Smoothing
method Int
horizon [Double]
xs = do
    Int -> [Double] -> Either MathError ()
validate Int
1 [Double]
xs
    Bool -> String -> Either MathError ()
require (Int
horizon Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
horizon Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
16384) String
"invalid forecast horizon"
    let rate :: Double -> Bool
rate Double
x = Double -> Bool
finite Double
x Bool -> Bool -> Bool
&& Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 Bool -> Bool -> Bool
&& Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
1
    (fitted, future) <- case Smoothing
method of
        Simple Double
alpha -> do
            Bool -> String -> Either MathError ()
require (Double -> Bool
rate Double
alpha) String
"smoothing rate outside [0,1]"
            let states :: [Double]
states = (Double -> Double -> Double) -> Double -> [Double] -> [Double]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl (\Double
l Double
x -> Double
alpha Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
alpha) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
l) (([Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) [Double]
xs) ((Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1) [Double]
xs)
            ([Double], [Double]) -> Either MathError ([Double], [Double])
forall a b. b -> Either a b
Right (([Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) [Double]
xs Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double] -> [Double]
forall a. HasCallStack => [a] -> [a]
init [Double]
states, Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
horizon ([Double] -> Double
forall a. HasCallStack => [a] -> a
last [Double]
states))
        Holt Double
alpha Double
beta -> do
            Bool -> String -> Either MathError ()
require ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
rate [Double
alpha, Double
beta]) String
"invalid Holt inputs"
            let step :: (Double, Double, [Double]) -> Double -> (Double, Double, [Double])
step (Double
l, Double
b, [Double]
fs) Double
x =
                    let predicted :: Double
predicted = Double
l Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b
                        next :: Double
next = Double
alpha Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
alpha) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
predicted
                        trend :: Double
trend = Double
beta Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
next Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
l) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
beta) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b
                     in (Double
next, Double
trend, [Double]
fs [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double
predicted])
                (Double
finalLevel, Double
finalTrend, [Double]
finalFitted) = ((Double, Double, [Double])
 -> Double -> (Double, Double, [Double]))
-> (Double, Double, [Double])
-> [Double]
-> (Double, Double, [Double])
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (Double, Double, [Double]) -> Double -> (Double, Double, [Double])
step (([Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) [Double]
xs, [Double]
xs [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- ([Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) [Double]
xs, [([Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) [Double]
xs]) ((Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1) [Double]
xs)
            ([Double], [Double]) -> Either MathError ([Double], [Double])
forall a b. b -> Either a b
Right ([Double]
finalFitted, [Double
finalLevel Double -> Double -> Double
forall a. Num a => a -> a -> a
+ 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
finalTrend | Int
k <- [Int
1 .. Int
horizon]])
        HoltWinters Int
period Double
alpha Double
beta Double
gamma -> do
            Bool -> String -> Either MathError ()
require
                (Int
period Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 Bool -> Bool -> Bool
&& Int
period Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1024 Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
period Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
rate [Double
alpha, Double
beta, Double
gamma])
                String
"invalid Holt-Winters inputs"
            let level :: Double
level = [Double] -> Double
average (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
period [Double]
xs)
                trend :: Double
trend = ([Double] -> Double
average (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
period (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
period [Double]
xs)) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
level) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
period
                seasons :: [Double]
seasons = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Num a => a -> a -> a
subtract Double
level) (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
period [Double]
xs)
                step :: (Double, Double, [Double], [Double])
-> (Int, Double) -> (Double, Double, [Double], [Double])
step (Double
l, Double
b, [Double]
ss, [Double]
fs) (Int
t, Double
x) =
                    let j :: Int
j = Int
t Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
period
                        s :: Double
s = [Double]
ss [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j
                        prediction :: Double
prediction = Double
l Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
s
                        next :: Double
next = Double
alpha Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
s) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
alpha) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
l Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b)
                        bt :: Double
bt = Double
beta Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
next Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
l) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
beta) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b
                        st :: Double
st = Double
gamma Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
next) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
gamma) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
s
                     in (Double
next, Double
bt, Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
j [Double]
ss [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double
st] [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Double]
ss, [Double]
fs [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double
prediction])
                (Double
finalLevel, Double
finalTrend, [Double]
finalSeasonal, [Double]
finalFitted) = ((Double, Double, [Double], [Double])
 -> (Int, Double) -> (Double, Double, [Double], [Double]))
-> (Double, Double, [Double], [Double])
-> [(Int, Double)]
-> (Double, Double, [Double], [Double])
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (Double, Double, [Double], [Double])
-> (Int, Double) -> (Double, Double, [Double], [Double])
step (Double
level, Double
trend, [Double]
seasons, Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
period [Double]
xs) ([Int] -> [Double] -> [(Int, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
period ..] (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
period [Double]
xs))
            ([Double], [Double]) -> Either MathError ([Double], [Double])
forall a b. b -> Either a b
Right
                ( [Double]
finalFitted
                , [ Double
finalLevel Double -> Double -> Double
forall a. Num a => a -> a -> a
+ 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
finalTrend Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double]
finalSeasonal [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
period)
                  | Int
k <- [Int
1 .. Int
horizon]
                  ]
                )
    _ <- checked (fitted ++ future)
    Right (fitted, future)

-- | Additive reconstruction is trend+seasonal+remainder at every input index.
data Decomposition = Decomposition
    { Decomposition -> [Double]
trendComponent :: [Double]
    -- ^ Estimated trend values aligned with original observations.
    , Decomposition -> [Double]
seasonalComponent :: [Double]
    -- ^ Estimated seasonal values aligned with original observations.
    , Decomposition -> [Double]
remainderComponent :: [Double]
    -- ^ Residual values after subtracting trend and seasonality.
    }
    deriving (Decomposition -> Decomposition -> Bool
(Decomposition -> Decomposition -> Bool)
-> (Decomposition -> Decomposition -> Bool) -> Eq Decomposition
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Decomposition -> Decomposition -> Bool
== :: Decomposition -> Decomposition -> Bool
$c/= :: Decomposition -> Decomposition -> Bool
/= :: Decomposition -> Decomposition -> Bool
Eq, Int -> Decomposition -> ShowS
[Decomposition] -> ShowS
Decomposition -> String
(Int -> Decomposition -> ShowS)
-> (Decomposition -> String)
-> ([Decomposition] -> ShowS)
-> Show Decomposition
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Decomposition -> ShowS
showsPrec :: Int -> Decomposition -> ShowS
$cshow :: Decomposition -> String
show :: Decomposition -> String
$cshowList :: [Decomposition] -> ShowS
showList :: [Decomposition] -> ShowS
Show)

{- | Classical additive decomposition with local linear trend smoothing and
centered phase means. Requires at least two complete periods.
-}
decompose :: Int -> [Double] -> Either MathError Decomposition
decompose :: Int -> [Double] -> Either MathError Decomposition
decompose Int
period [Double]
xs = do
    Int -> [Double] -> Either MathError ()
validate Int
1 [Double]
xs
    Bool -> String -> Either MathError ()
require
        (Int
period Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 Bool -> Bool -> Bool
&& Int
period Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1024 Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
period Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
2048)
        String
"invalid decomposition period/sample count"
    let trend :: [Double]
trend = Int -> [Double] -> [Double] -> [Double]
loess (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
period Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs) Double
1) [Double]
xs
        detrended :: [Double]
detrended = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) [Double]
xs [Double]
trend
        phases :: [Double]
phases =
            [[Double] -> Double
average [Double
v | (Int
i, Double
v) <- [Int] -> [Double] -> [(Int, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [Double]
detrended, Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
period Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j] | Int
j <- [Int
0 .. Int
period Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
        seasonal :: [Double]
seasonal = [[Double]
phases [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
period) Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double] -> Double
average [Double]
phases | Int
i <- [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    [Double] -> [Double] -> [Double] -> Either MathError Decomposition
decomposition [Double]
xs [Double]
trend [Double]
seasonal

{- | Seasonal period,odd LOESS windows,and inner/robust iteration counts for STL
decomposition. Use defaultSTLConfig then checked record updates.
-}
data STLConfig = STLConfig
    { STLConfig -> Int
stlPeriod :: Int
    -- ^ Positive seasonal cycle length in observations.
    , STLConfig -> Int
stlSeasonalWindow :: Int
    -- ^ Odd LOESS window for seasonal subseries.
    , STLConfig -> Int
stlTrendWindow :: Int
    -- ^ Odd LOESS window for trend smoothing.
    , STLConfig -> Int
stlLowPassWindow :: Int
    -- ^ Odd low-pass LOESS smoothing window.
    , STLConfig -> Int
stlInnerIterations :: Int
    -- ^ Number of STL seasonal/trend backfitting passes.
    , STLConfig -> Int
stlRobustIterations :: Int
    -- ^ Number of robust outer reweighting passes; zero disables robust reweighting.
    }
    deriving (STLConfig -> STLConfig -> Bool
(STLConfig -> STLConfig -> Bool)
-> (STLConfig -> STLConfig -> Bool) -> Eq STLConfig
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: STLConfig -> STLConfig -> Bool
== :: STLConfig -> STLConfig -> Bool
$c/= :: STLConfig -> STLConfig -> Bool
/= :: STLConfig -> STLConfig -> Bool
Eq, Int -> STLConfig -> ShowS
[STLConfig] -> ShowS
STLConfig -> String
(Int -> STLConfig -> ShowS)
-> (STLConfig -> String)
-> ([STLConfig] -> ShowS)
-> Show STLConfig
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> STLConfig -> ShowS
showsPrec :: Int -> STLConfig -> ShowS
$cshow :: STLConfig -> String
show :: STLConfig -> String
$cshowList :: [STLConfig] -> ShowS
showList :: [STLConfig] -> ShowS
Show)

{- | Given seasonal period, construct STL options with seasonal window 7, odd trend window
above 2*period, odd low-pass window at/above period, two inner iterations and no robust
outer iterations. stl validates the resulting configuration.
-}
defaultSTLConfig :: Int -> STLConfig
defaultSTLConfig :: Int -> STLConfig
defaultSTLConfig Int
period = Int -> Int -> Int -> Int -> Int -> Int -> STLConfig
STLConfig Int
period Int
7 (Int -> Int
forall {a}. Integral a => a -> a
oddAbove (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
period Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) (Int -> Int
forall {a}. Integral a => a -> a
oddAbove Int
period) Int
2 Int
0
  where
    oddAbove :: a -> a
oddAbove a
x = if a -> Bool
forall a. Integral a => a -> Bool
odd a
x then a
x else a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a
1

{- | Cleveland-style STL reference backfitting: smooth seasonal subseries by
degree-one tricube LOESS with one-cycle extrapolation; remove a period,
period, three-point moving-average low-pass followed by LOESS, then smooth
the deseasonalized trend. Optional bisquare robust outer iterations use six
times median absolute residual. Every index is evaluated (no jump shortcut).
See https://search.r-project.org/R/refmans/stats/html/stl.html .
-}
stl :: STLConfig -> [Double] -> Either MathError Decomposition
stl :: STLConfig -> [Double] -> Either MathError Decomposition
stl STLConfig
config [Double]
xs = do
    Int -> [Double] -> Either MathError ()
validate Int
1 [Double]
xs
    let p :: Int
p = STLConfig -> Int
stlPeriod STLConfig
config
        n :: Int
n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs
        windows :: [Int]
windows = [STLConfig -> Int
stlSeasonalWindow STLConfig
config, STLConfig -> Int
stlTrendWindow STLConfig
config, STLConfig -> Int
stlLowPassWindow STLConfig
config]
        ni :: Int
ni = STLConfig -> Int
stlInnerIterations STLConfig
config
        no :: Int
no = STLConfig -> Int
stlRobustIterations STLConfig
config
    Bool -> String -> Either MathError ()
require
        ( Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2
            Bool -> Bool -> Bool
&& Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
512
            Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
p
            Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1024
            Bool -> Bool -> Bool
&& (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Int
w -> Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
3 Bool -> Bool -> Bool
&& Int -> Bool
forall a. Integral a => a -> Bool
odd Int
w Bool -> Bool -> Bool
&& Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
4095) [Int]
windows
            Bool -> Bool -> Bool
&& Int
ni Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1
            Bool -> Bool -> Bool
&& Int
ni Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
10
            Bool -> Bool -> Bool
&& Int
no Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
            Bool -> Bool -> Bool
&& Int
no Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
10
        )
        String
"invalid STL configuration/sample count"
    Bool -> String -> Either MathError ()
require (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
ni Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
no Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
20000000) String
"STL work limit"
    let outer :: Int -> [Double] -> [Double] -> ([Double], [Double])
outer Int
k [Double]
weights [Double]
trend =
            let ([Double]
tr, [Double]
se) = (([Double], [Double]) -> ([Double], [Double]))
-> ([Double], [Double]) -> [([Double], [Double])]
forall a. (a -> a) -> a -> [a]
iterate ([Double] -> ([Double], [Double]) -> ([Double], [Double])
forall {b}. [Double] -> ([Double], b) -> ([Double], [Double])
inner [Double]
weights) ([Double]
trend, Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
n Double
0) [([Double], [Double])] -> Int -> ([Double], [Double])
forall a. HasCallStack => [a] -> Int -> a
!! Int
ni
                residual :: [Double]
residual = (Double -> Double -> Double -> Double)
-> [Double] -> [Double] -> [Double] -> [Double]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\Double
x Double
t Double
s -> Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
s) [Double]
xs [Double]
tr [Double]
se
                scale :: Double
scale = Double
6 Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double] -> Double
forall {a}. (Fractional a, Ord a) => [a] -> a
median ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Num a => a -> a
abs [Double]
residual)
                nextWeights :: [Double]
nextWeights =
                    if Double
scale Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
                        then (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Double
r -> if Double
r Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
1 else Double
0) [Double]
residual
                        else [let u :: Double
u = Double -> Double
forall a. Num a => a -> a
abs Double
r Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
scale in if Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
1 then Double
0 else (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
u) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) | Double
r <- [Double]
residual]
             in if Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
no then ([Double]
tr, [Double]
se) else Int -> [Double] -> [Double] -> ([Double], [Double])
outer (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Double]
nextWeights [Double]
tr
        inner :: [Double] -> ([Double], b) -> ([Double], [Double])
inner [Double]
weights ([Double]
trend, b
_) =
            let detrended :: [Double]
detrended = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) [Double]
xs [Double]
trend
                sub :: Int -> [(a, Double, Double)]
sub Int
j =
                    [ (Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k, Double
v, [Double]
weights [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
                    | (Int
k, Int
i) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [Int
j, Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    , let v :: Double
v = [Double]
detrended [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i
                    ]
                phaseAt :: Int -> Double
phaseAt Int
i =
                    let j :: Int
j = Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
p; t :: Double
t = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
j) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
p)
                     in Int -> [(Double, Double, Double)] -> Double -> Double
loessAt (STLConfig -> Int
stlSeasonalWindow STLConfig
config) (Int -> [(Double, Double, Double)]
forall {a}. Num a => Int -> [(a, Double, Double)]
sub Int
j) Double
t
                extended :: [Double]
extended = [Int -> Double
phaseAt Int
i | Int
i <- [-Int
p .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                low0 :: [Double]
low0 = Int -> [Double] -> [Double]
moving Int
p [Double]
extended
                low1 :: [Double]
low1 = Int -> [Double] -> [Double]
moving Int
p [Double]
low0
                low2 :: [Double]
low2 = Int -> [Double] -> [Double]
moving Int
3 [Double]
low1
                low :: [Double]
low = Int -> [Double] -> [Double] -> [Double]
loess (STLConfig -> Int
stlLowPassWindow STLConfig
config) (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
n Double
1) [Double]
low2
                seasonal :: [Double]
seasonal = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
n (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
p [Double]
extended)) [Double]
low
                trend' :: [Double]
trend' = Int -> [Double] -> [Double] -> [Double]
loess (STLConfig -> Int
stlTrendWindow STLConfig
config) [Double]
weights ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) [Double]
xs [Double]
seasonal)
             in ([Double]
trend', [Double]
seasonal)
        ([Double]
finalTrend, [Double]
finalSeasonal) = Int -> [Double] -> [Double] -> ([Double], [Double])
outer Int
0 (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
n Double
1) (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
n Double
0)
    [Double] -> [Double] -> [Double] -> Either MathError Decomposition
decomposition [Double]
xs [Double]
finalTrend [Double]
finalSeasonal
  where
    moving :: Int -> [Double] -> [Double]
moving Int
k [Double]
ys = [[Double] -> Double
average (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
k (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
i [Double]
ys)) | Int
i <- [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
ys Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k]]
    median :: [a] -> a
median [a]
ys =
        let zs :: [a]
zs = [a] -> [a]
forall a. Ord a => [a] -> [a]
sort [a]
ys; n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
zs
         in if Int -> Bool
forall a. Integral a => a -> Bool
odd Int
n then [a]
zs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) else ([a]
zs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) a -> a -> a
forall a. Num a => a -> a -> a
+ [a]
zs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)) a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
2

loess :: Int -> [Double] -> [Double] -> [Double]
loess :: Int -> [Double] -> [Double] -> [Double]
loess Int
window [Double]
weights [Double]
xs = [Int -> [(Double, Double, Double)] -> Double -> Double
loessAt Int
window [(Double, Double, Double)]
sample (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i) | Int
i <- [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
  where
    sample :: [(Double, Double, Double)]
sample = (Int -> Double -> Double -> (Double, Double, Double))
-> [Int] -> [Double] -> [Double] -> [(Double, Double, Double)]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\Int
i Double
x Double
w -> (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i, Double
x, Double
w)) [Int
0 :: Int ..] [Double]
xs [Double]
weights

-- Local linear weighted least squares centered at the query. Select nearest
-- window points with deterministic ties; a widened radius keeps all endpoints
-- from receiving zero weight, including the two-point/extrapolation cases.
loessAt :: Int -> [(Double, Double, Double)] -> Double -> Double
loessAt :: Int -> [(Double, Double, Double)] -> Double -> Double
loessAt Int
window [(Double, Double, Double)]
sample Double
query =
    let nearest :: [(Double, Double, Double)]
nearest = Int -> [(Double, Double, Double)] -> [(Double, Double, Double)]
forall a. Int -> [a] -> [a]
take Int
window ([(Double, Double, Double)] -> [(Double, Double, Double)]
sortOnDistance [(Double, Double, Double)]
sample)
        radius :: Double
radius = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
1e-12 ([Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (((Double, Double, Double) -> Double)
-> [(Double, Double, Double)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\(Double
t, Double
_, Double
_) -> Double -> Double
forall a. Num a => a -> a
abs (Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
query)) [(Double, Double, Double)]
nearest) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
1.000001)
        rows :: [(Double, Double, Double)]
rows =
            [ (Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
query, Double
x, Double
w Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- (Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
1 (Double -> Double
forall a. Num a => a -> a
abs (Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
query) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
radius)) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
3 :: Int)) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
3 :: Int))
            | (Double
t, Double
x, Double
w) <- [(Double, Double, Double)]
nearest
            ]
        sw :: Double
sw = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double
w | (Double
_, Double
_, Double
w) <- [(Double, Double, Double)]
rows]
        sx :: Double
sx = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double
w Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
t | (Double
t, Double
_, Double
w) <- [(Double, Double, Double)]
rows]
        sxx :: Double
sxx = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double
w Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
t | (Double
t, Double
_, Double
w) <- [(Double, Double, Double)]
rows]
        sy :: Double
sy = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double
w Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x | (Double
_, Double
x, Double
w) <- [(Double, Double, Double)]
rows]
        sxy :: Double
sxy = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double
w Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x | (Double
t, Double
x, Double
w) <- [(Double, Double, Double)]
rows]
        den :: Double
den = Double
sw Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sxx Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
sx Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sx
     in if Double
sw Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
            then [Double] -> Double
average [Double
x | (Double
_, Double
x, Double
_) <- [(Double, Double, Double)]
nearest]
            else if Double -> Double
forall a. Num a => a -> a
abs Double
den Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
1e-14 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
1 (Double
sw Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sxx) then Double
sy Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sw else (Double
sxx Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sy Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
sx Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sxy) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
den
  where
    sortOnDistance :: [(Double, Double, Double)] -> [(Double, Double, Double)]
sortOnDistance = (((Double, Double), (Double, Double, Double))
 -> (Double, Double, Double))
-> [((Double, Double), (Double, Double, Double))]
-> [(Double, Double, Double)]
forall a b. (a -> b) -> [a] -> [b]
map ((Double, Double), (Double, Double, Double))
-> (Double, Double, Double)
forall a b. (a, b) -> b
snd ([((Double, Double), (Double, Double, Double))]
 -> [(Double, Double, Double)])
-> ([(Double, Double, Double)]
    -> [((Double, Double), (Double, Double, Double))])
-> [(Double, Double, Double)]
-> [(Double, Double, Double)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [((Double, Double), (Double, Double, Double))]
-> [((Double, Double), (Double, Double, Double))]
forall a. Ord a => [a] -> [a]
sort ([((Double, Double), (Double, Double, Double))]
 -> [((Double, Double), (Double, Double, Double))])
-> ([(Double, Double, Double)]
    -> [((Double, Double), (Double, Double, Double))])
-> [(Double, Double, Double)]
-> [((Double, Double), (Double, Double, Double))]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Double, Double, Double)
 -> ((Double, Double), (Double, Double, Double)))
-> [(Double, Double, Double)]
-> [((Double, Double), (Double, Double, Double))]
forall a b. (a -> b) -> [a] -> [b]
map (\r :: (Double, Double, Double)
r@(Double
t, Double
_, Double
_) -> ((Double -> Double
forall a. Num a => a -> a
abs (Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
query), Double
t), (Double, Double, Double)
r))

decomposition :: [Double] -> [Double] -> [Double] -> Either MathError Decomposition
decomposition :: [Double] -> [Double] -> [Double] -> Either MathError Decomposition
decomposition [Double]
xs [Double]
trend [Double]
seasonal = do
    let residual :: [Double]
residual = (Double -> Double -> Double -> Double)
-> [Double] -> [Double] -> [Double] -> [Double]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\Double
x Double
t Double
s -> Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
s) [Double]
xs [Double]
trend [Double]
seasonal
    _ <- [Double] -> Either MathError [Double]
checked ([Double]
trend [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
seasonal [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
residual)
    Right (Decomposition trend seasonal residual)

-- | Intercept followed by lag coefficients; observations stored chronologically.
data ARModel = ARModel [Double] [Double] deriving (ARModel -> ARModel -> Bool
(ARModel -> ARModel -> Bool)
-> (ARModel -> ARModel -> Bool) -> Eq ARModel
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ARModel -> ARModel -> Bool
== :: ARModel -> ARModel -> Bool
$c/= :: ARModel -> ARModel -> Bool
/= :: ARModel -> ARModel -> Bool
Eq, Int -> ARModel -> ShowS
[ARModel] -> ShowS
ARModel -> String
(Int -> ARModel -> ShowS)
-> (ARModel -> String) -> ([ARModel] -> ShowS) -> Show ARModel
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ARModel -> ShowS
showsPrec :: Int -> ARModel -> ShowS
$cshow :: ARModel -> String
show :: ARModel -> String
$cshowList :: [ARModel] -> ShowS
showList :: [ARModel] -> ShowS
Show)

{- | Return AR intercept followed by lag-1 through lag-p coefficients. The model retains
observations chronologically for forecastAR.
-}
arCoefficients :: ARModel -> [Double]
arCoefficients :: ARModel -> [Double]
arCoefficients (ARModel [Double]
b [Double]
_) = [Double]
b

{- | Given AR order then finite chronological samples, fit intercept and lag coefficients by
checked OLS. Require order in [0,64], \<=4096 samples and more than 2*order+1
observations; invalid/ill-conditioned fits fail.
-}
fitAR :: Int -> [Double] -> Either MathError ARModel
fitAR :: Int -> [Double] -> Either MathError ARModel
fitAR Int
order [Double]
xs = do
    Int -> [Double] -> Either MathError ()
validate Int
2 [Double]
xs
    Bool -> String -> Either MathError ()
require
        (Int
order Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
order Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
64 Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
order Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
4096)
        String
"invalid AR order/sample count"
    b <-
        [[Double]] -> [Double] -> Either MathError [Double]
ols [[Double
1] [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [[Double]
xs [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
j) | Int
j <- [Int
1 .. Int
order]] | Int
t <- [Int
order .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
order [Double]
xs)
    Right (ARModel b xs)

{- | Given horizon in [0,16384] then fitted AR model, recursively forecast future samples
from latest observations. Zero horizon returns []; invalid horizon or non-finite
predictions fail.
-}
forecastAR :: Int -> ARModel -> Either MathError [Double]
forecastAR :: Int -> ARModel -> Either MathError [Double]
forecastAR Int
horizon (ARModel [Double]
b [Double]
xs) = do
    Bool -> String -> Either MathError ()
require (Int
horizon Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
horizon Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
16384) String
"invalid AR forecast horizon"
    let p :: Int
p = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
b Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        go :: t -> [Double] -> [Double] -> Either MathError [Double]
go t
0 [Double]
_ [Double]
acc = [Double] -> Either MathError [Double]
forall a b. b -> Either a b
Right ([Double] -> [Double]
forall a. [a] -> [a]
reverse [Double]
acc)
        go t
k [Double]
past [Double]
acc = do
            let value :: Double
value = ([Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) [Double]
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) ((Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1) [Double]
b) (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
p [Double]
past))
            _ <- [Double] -> Either MathError [Double]
checked [Double
value]
            go (k - 1) (value : past) (value : acc)
    Int -> [Double] -> [Double] -> Either MathError [Double]
forall {t}.
(Eq t, Num t) =>
t -> [Double] -> [Double] -> Either MathError [Double]
go Int
horizon ([Double] -> [Double]
forall a. [a] -> [a]
reverse [Double]
xs) []

{- | Ordinary (p,d,q),seasonal (P,D,Q),and cycle period s. fitSARIMA validates supported
order/sample budgets.
-}
data SARIMAOrder = SARIMAOrder
    { SARIMAOrder -> Int
arOrder :: Int
    -- ^ Ordinary autoregressive order p.
    , SARIMAOrder -> Int
differenceOrder :: Int
    -- ^ Ordinary differencing order d.
    , SARIMAOrder -> Int
maOrder :: Int
    -- ^ Ordinary moving-average order q.
    , SARIMAOrder -> Int
seasonalAROrder :: Int
    -- ^ Seasonal autoregressive order P.
    , SARIMAOrder -> Int
seasonalDifferenceOrder :: Int
    -- ^ Seasonal differencing order D.
    , SARIMAOrder -> Int
seasonalMAOrder :: Int
    -- ^ Seasonal moving-average order Q.
    , SARIMAOrder -> Int
seasonalPeriod :: Int
    -- ^ Seasonal cycle length s used to place seasonal polynomial lags.
    }
    deriving (SARIMAOrder -> SARIMAOrder -> Bool
(SARIMAOrder -> SARIMAOrder -> Bool)
-> (SARIMAOrder -> SARIMAOrder -> Bool) -> Eq SARIMAOrder
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SARIMAOrder -> SARIMAOrder -> Bool
== :: SARIMAOrder -> SARIMAOrder -> Bool
$c/= :: SARIMAOrder -> SARIMAOrder -> Bool
/= :: SARIMAOrder -> SARIMAOrder -> Bool
Eq, Int -> SARIMAOrder -> ShowS
[SARIMAOrder] -> ShowS
SARIMAOrder -> String
(Int -> SARIMAOrder -> ShowS)
-> (SARIMAOrder -> String)
-> ([SARIMAOrder] -> ShowS)
-> Show SARIMAOrder
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SARIMAOrder -> ShowS
showsPrec :: Int -> SARIMAOrder -> ShowS
$cshow :: SARIMAOrder -> String
show :: SARIMAOrder -> String
$cshowList :: [SARIMAOrder] -> ShowS
showList :: [SARIMAOrder] -> ShowS
Show)

{- | Opaque fitted ARIMA/SARIMA state with differencing,parameters,innovation residuals,and
original history for forecasting.
-}
data ARIMAModel = ARIMAModel SARIMAOrder [Double] [Double] [Double] [Double] Double
    deriving (ARIMAModel -> ARIMAModel -> Bool
(ARIMAModel -> ARIMAModel -> Bool)
-> (ARIMAModel -> ARIMAModel -> Bool) -> Eq ARIMAModel
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ARIMAModel -> ARIMAModel -> Bool
== :: ARIMAModel -> ARIMAModel -> Bool
$c/= :: ARIMAModel -> ARIMAModel -> Bool
/= :: ARIMAModel -> ARIMAModel -> Bool
Eq, Int -> ARIMAModel -> ShowS
[ARIMAModel] -> ShowS
ARIMAModel -> String
(Int -> ARIMAModel -> ShowS)
-> (ARIMAModel -> String)
-> ([ARIMAModel] -> ShowS)
-> Show ARIMAModel
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ARIMAModel -> ShowS
showsPrec :: Int -> ARIMAModel -> ShowS
$cshow :: ARIMAModel -> String
show :: ARIMAModel -> String
$cshowList :: [ARIMAModel] -> ShowS
showList :: [ARIMAModel] -> ShowS
Show)

{- | Return fitted intercept,ordinary AR,ordinary MA,seasonal AR,seasonal MA parameters in
that order. Intercept belongs to the differenced process.
-}
arimaParameters :: ARIMAModel -> [Double]
arimaParameters :: ARIMAModel -> [Double]
arimaParameters (ARIMAModel SARIMAOrder
_ [Double]
b [Double]
_ [Double]
_ [Double]
_ Double
_) = [Double]
b

{- | Return fitted innovation residuals in chronological order, on the model's differenced
process.
-}
arimaResiduals :: ARIMAModel -> [Double]
arimaResiduals :: ARIMAModel -> [Double]
arimaResiduals (ARIMAModel SARIMAOrder
_ [Double]
_ [Double]
_ [Double]
_ [Double]
es Double
_) = [Double]
es

-- | Return fitted innovation-variance estimate stored in an ARIMA/SARIMA model.
arimaInnovationVariance :: ARIMAModel -> Double
arimaInnovationVariance :: ARIMAModel -> Double
arimaInnovationVariance (ARIMAModel SARIMAOrder
_ [Double]
_ [Double]
_ [Double]
_ [Double]
_ Double
v) = Double
v

{- | Given optimization options,p,d,q,then finite chronological series, fit nonseasonal
ARIMA through fitSARIMA with seasonal orders zero. Invalid order/sample bounds,
unstable or nonconvergent/ill-conditioned fit returns Left MathError.
-}
fitARIMA :: O.OptimizeOptions -> Int -> Int -> Int -> [Double] -> Either MathError ARIMAModel
fitARIMA :: OptimizeOptions
-> Int -> Int -> Int -> [Double] -> Either MathError ARIMAModel
fitARIMA OptimizeOptions
options Int
p Int
d Int
q = OptimizeOptions
-> SARIMAOrder -> [Double] -> Either MathError ARIMAModel
fitSARIMA OptimizeOptions
options (Int -> Int -> Int -> Int -> Int -> Int -> Int -> SARIMAOrder
SARIMAOrder Int
p Int
d Int
q Int
0 Int
0 Int
0 Int
1)

{- | Given optimization options,MA order q,then samples, fit ARIMA(0,0,q).
Order/sample/optimization errors propagate from fitARIMA.
-}
fitMA :: O.OptimizeOptions -> Int -> [Double] -> Either MathError ARIMAModel
fitMA :: OptimizeOptions -> Int -> [Double] -> Either MathError ARIMAModel
fitMA OptimizeOptions
options Int
q = OptimizeOptions
-> Int -> Int -> Int -> [Double] -> Either MathError ARIMAModel
fitARIMA OptimizeOptions
options Int
0 Int
0 Int
q

{- | Given optimization options,AR order p,MA order q,then samples, fit ARIMA(p,0,q).
Invalid or nonconvergent fits return Left MathError.
-}
fitARMA :: O.OptimizeOptions -> Int -> Int -> [Double] -> Either MathError ARIMAModel
fitARMA :: OptimizeOptions
-> Int -> Int -> [Double] -> Either MathError ARIMAModel
fitARMA OptimizeOptions
options Int
p Int
q = OptimizeOptions
-> Int -> Int -> Int -> [Double] -> Either MathError ARIMAModel
fitARIMA OptimizeOptions
options Int
p Int
0 Int
q

{- | Parameters are intercept, ordinary AR, ordinary MA, seasonal AR, seasonal
MA. Multiplicative seasonal polynomials include cross-lag interactions.
Intercept always belongs to the differenced process. Caller supplies fit
budgets; exhausted or ill-conditioned fits return MathError.
-}
fitSARIMA :: O.OptimizeOptions -> SARIMAOrder -> [Double] -> Either MathError ARIMAModel
fitSARIMA :: OptimizeOptions
-> SARIMAOrder -> [Double] -> Either MathError ARIMAModel
fitSARIMA OptimizeOptions
options SARIMAOrder
order [Double]
xs = do
    Int -> [Double] -> Either MathError ()
validate Int
3 [Double]
xs
    OptimizeOptions -> Either MathError ()
O.checkedOptions OptimizeOptions
options
    let p :: Int
p = SARIMAOrder -> Int
arOrder SARIMAOrder
order
        q :: Int
q = SARIMAOrder -> Int
maOrder SARIMAOrder
order
        bigP :: Int
bigP = SARIMAOrder -> Int
seasonalAROrder SARIMAOrder
order
        bigQ :: Int
bigQ = SARIMAOrder -> Int
seasonalMAOrder SARIMAOrder
order
        s :: Int
s = SARIMAOrder -> Int
seasonalPeriod SARIMAOrder
order
        d :: Int
d = SARIMAOrder -> Int
differenceOrder SARIMAOrder
order
        bigD :: Int
bigD = SARIMAOrder -> Int
seasonalDifferenceOrder SARIMAOrder
order
        parameters :: Int
parameters = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bigP Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bigQ
        burn :: Int
burn = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bigP Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
s) (Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bigQ Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
s)
    Bool -> String -> Either MathError ()
require
        ( (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Int
k -> Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
8) [Int
p, Int
q, Int
bigP, Int
bigQ]
            Bool -> Bool -> Bool
&& Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
            Bool -> Bool -> Bool
&& Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
2
            Bool -> Bool -> Bool
&& Int
bigD Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
            Bool -> Bool -> Bool
&& Int
bigD Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
2
            Bool -> Bool -> Bool
&& Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1
            Bool -> Bool -> Bool
&& Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
128
            Bool -> Bool -> Bool
&& Int
parameters Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
13
            Bool -> Bool -> Bool
&& Int
burn Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
256
            Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
2048
        )
        String
"SARIMA order/sample limit"
    zs <- Int -> Int -> [Double] -> Either MathError [Double]
difference Int
d Int
1 [Double]
xs Either MathError [Double]
-> ([Double] -> Either MathError [Double])
-> Either MathError [Double]
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Int -> [Double] -> Either MathError [Double]
difference Int
bigD Int
s
    require (length zs > burn + parameters) "insufficient SARIMA observations"
    require
        ( toInteger (O.optMaxEvaluations options)
            * toInteger (length zs)
            * toInteger parameters
            * toInteger (max 1 burn)
            <= 100000000
        )
        "SARIMA fit work limit"
    let rows [Double]
b = SARIMAOrder
-> [Double] -> [Double] -> Either MathError ([Double], [[Double]])
innovations SARIMAOrder
order [Double]
b [Double]
zs
        scale = Double -> Double
forall a. Floating a => a -> a
sqrt (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
zs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
burn))
        residual [Double]
b = do (es, _) <- [Double] -> Either MathError ([Double], [[Double]])
rows [Double]
b; checked (map (/ scale) (drop burn es))
        jacobian [Double]
b = do
            (_, js) <- [Double] -> Either MathError ([Double], [[Double]])
rows [Double]
b
            let out = ([Double] -> [Double]) -> [[Double]] -> [[Double]]
forall a b. (a -> b) -> [a] -> [b]
map ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
scale)) (Int -> [[Double]] -> [[Double]]
forall a. Int -> [a] -> [a]
drop Int
burn [[Double]]
js)
            _ <- checked (concat out)
            Right out
        initial = [Double] -> Double
average [Double]
zs Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate (Int
parameters Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Double
0
    b <-
        if parameters == 1
            then Right initial
            else O.optPoint <$> O.leastSquares O.SquaredLoss 0.01 options residual jacobian initial
    (es, _) <- rows b
    let variance = [Double] -> Double
average [Double
e Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
e | Double
e <- Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
burn [Double]
es]
    _ <- checked [variance]
    Right (ARIMAModel order b xs zs es variance)

polynomialProduct :: [Double] -> [Double] -> [Double]
polynomialProduct :: [Double] -> [Double] -> [Double]
polynomialProduct [Double]
a [Double]
b =
    [ [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[Double]
a [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double]
b [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) | Int
i <- [Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
b Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) .. Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
k ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
a Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)]]
    | Int
k <- [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
b Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2]
    ]
seasonalPolynomial :: Int -> [Double] -> [Double]
seasonalPolynomial :: Int -> [Double] -> [Double]
seasonalPolynomial Int
s [Double]
values = [if Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then [Double]
values [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
s) else Double
0 | Int
i <- [Int
0 .. Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
* ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
values Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)]]

modelPolynomials :: SARIMAOrder -> [Double] -> ([Double], [Double])
modelPolynomials :: SARIMAOrder -> [Double] -> ([Double], [Double])
modelPolynomials SARIMAOrder
order [Double]
b =
    let p :: Int
p = SARIMAOrder -> Int
arOrder SARIMAOrder
order
        q :: Int
q = SARIMAOrder -> Int
maOrder SARIMAOrder
order
        bigP :: Int
bigP = SARIMAOrder -> Int
seasonalAROrder SARIMAOrder
order
        bigQ :: Int
bigQ = SARIMAOrder -> Int
seasonalMAOrder SARIMAOrder
order
        s :: Int
s = SARIMAOrder -> Int
seasonalPeriod SARIMAOrder
order
        phi :: [Double]
phi = Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
p (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1 [Double]
b)
        theta :: [Double]
theta = Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
q (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p) [Double]
b)
        bigPhi :: [Double]
bigPhi = Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
bigP (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
q) [Double]
b)
        bigTheta :: [Double]
bigTheta = Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
bigQ (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bigP) [Double]
b)
     in ( (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map
            Double -> Double
forall a. Num a => a -> a
negate
            ((Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1) ([Double] -> [Double] -> [Double]
polynomialProduct (Double
1 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Num a => a -> a
negate [Double]
phi) (Int -> [Double] -> [Double]
seasonalPolynomial Int
s (Double
1 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Num a => a -> a
negate [Double]
bigPhi))))
        , (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1) ([Double] -> [Double] -> [Double]
polynomialProduct (Double
1 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double]
theta) (Int -> [Double] -> [Double]
seasonalPolynomial Int
s (Double
1 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double]
bigTheta)))
        )

innovations :: SARIMAOrder -> [Double] -> [Double] -> Either MathError ([Double], [[Double]])
innovations :: SARIMAOrder
-> [Double] -> [Double] -> Either MathError ([Double], [[Double]])
innovations SARIMAOrder
order [Double]
b [Double]
zs = do
    _ <- [Double] -> Either MathError [Double]
checked [Double]
b
    let (ar, ma) = modelPolynomials order b
        derivatives = SARIMAOrder -> [Double] -> [([Double], [Double])]
coefficientPartials SARIMAOrder
order [Double]
b
        past [a]
values Int
k = if Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
values then [a]
values [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
k else a
0
        step ([Double]
es, [[Double]]
js) (Int
t, Double
z) =
            let ypast :: [Double]
ypast = [Double] -> [Double]
forall a. [a] -> [a]
reverse (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
t [Double]
zs)
                e :: Double
e = Double
z Double -> Double -> Double
forall a. Num a => a -> a -> a
- ([Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) [Double]
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
ar [Double]
ypast) Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
ma [Double]
es)
                deriv :: Int -> ([Double], [Double]) -> Double
deriv Int
j ([Double]
da, [Double]
dm) =
                    -(if Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Double
1 else Double
0)
                        Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
da [Double]
ypast)
                        Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
dm [Double]
es)
                        Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double
m Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double] -> Int -> Double
forall {a}. Num a => [a] -> Int -> a
past (([Double] -> Double) -> [[Double]] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map ([Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j) [[Double]]
js) Int
k | (Int
k, Double
m) <- [Int] -> [Double] -> [(Int, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Double]
ma]
                row :: [Double]
row = (Int -> ([Double], [Double]) -> Double)
-> [Int] -> [([Double], [Double])] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> ([Double], [Double]) -> Double
deriv [Int
0 ..] [([Double], [Double])]
derivatives
             in (Double
e Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double]
es, [Double]
row [Double] -> [[Double]] -> [[Double]]
forall a. a -> [a] -> [a]
: [[Double]]
js)
        (er, jr) = foldl step ([], []) (zip [0 ..] zs)
    _ <- checked (er ++ concat jr)
    Right (reverse er, reverse jr)

-- Differentiate polynomial products directly. Subtracting f(b+1)-f(b)
-- would lose a unit perturbation entirely for coefficients above 2^53.
coefficientPartials :: SARIMAOrder -> [Double] -> [([Double], [Double])]
coefficientPartials :: SARIMAOrder -> [Double] -> [([Double], [Double])]
coefficientPartials SARIMAOrder
order [Double]
b = [Int -> ([Double], [Double])
partial Int
j | Int
j <- [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
b Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
  where
    p :: Int
p = SARIMAOrder -> Int
arOrder SARIMAOrder
order
    q :: Int
q = SARIMAOrder -> Int
maOrder SARIMAOrder
order
    bigP :: Int
bigP = SARIMAOrder -> Int
seasonalAROrder SARIMAOrder
order
    bigQ :: Int
bigQ = SARIMAOrder -> Int
seasonalMAOrder SARIMAOrder
order
    s :: Int
s = SARIMAOrder -> Int
seasonalPeriod SARIMAOrder
order
    ar :: [Double]
ar = Double
1 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Num a => a -> a
negate (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
p (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1 [Double]
b))
    ma :: [Double]
ma = Double
1 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
q (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p) [Double]
b)
    seasonalAR :: [Double]
seasonalAR = Int -> [Double] -> [Double]
seasonalPolynomial Int
s (Double
1 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Num a => a -> a
negate (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
bigP (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
q) [Double]
b)))
    seasonalMA :: [Double]
seasonalMA = Int -> [Double] -> [Double]
seasonalPolynomial Int
s (Double
1 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
bigQ (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bigP) [Double]
b))
    basis :: a -> a -> a -> [a]
basis a
n a
index a
sign = [if a
k a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
index then a
sign else a
0 | a
k <- [a
0 .. a
n]]
    partial :: Int -> ([Double], [Double])
partial Int
j
        | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1 Bool -> Bool -> Bool
&& Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
p =
            ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Num a => a -> a
negate (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1 ([Double] -> [Double] -> [Double]
polynomialProduct (Int -> Int -> Double -> [Double]
forall {a} {a}. (Enum a, Eq a, Num a, Num a) => a -> a -> a -> [a]
basis Int
p Int
j (-Double
1)) [Double]
seasonalAR)), Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate (Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bigQ Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
s) Double
0)
        | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
p Bool -> Bool -> Bool
&& Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
q =
            (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bigP Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
s) Double
0, Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1 ([Double] -> [Double] -> [Double]
polynomialProduct (Int -> Int -> Double -> [Double]
forall {a} {a}. (Enum a, Eq a, Num a, Num a) => a -> a -> a -> [a]
basis Int
q (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
p) Double
1) [Double]
seasonalMA))
        | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
q Bool -> Bool -> Bool
&& Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bigP =
            ( (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Num a => a -> a
negate (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1 ([Double] -> [Double] -> [Double]
polynomialProduct [Double]
ar (Int -> [Double] -> [Double]
seasonalPolynomial Int
s (Int -> Int -> Double -> [Double]
forall {a} {a}. (Enum a, Eq a, Num a, Num a) => a -> a -> a -> [a]
basis Int
bigP (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
q) (-Double
1)))))
            , Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate (Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bigQ Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
s) Double
0
            )
        | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bigP =
            ( Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bigP Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
s) Double
0
            , Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1 ([Double] -> [Double] -> [Double]
polynomialProduct [Double]
ma (Int -> [Double] -> [Double]
seasonalPolynomial Int
s (Int -> Int -> Double -> [Double]
forall {a} {a}. (Enum a, Eq a, Num a, Num a) => a -> a -> a -> [a]
basis Int
bigQ (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bigP) Double
1)))
            )
        | Bool
otherwise = (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bigP Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
s) Double
0, Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate (Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bigQ Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
s) Double
0)

{- | Given horizon in [0,16384] then fitted ARIMA/SARIMA model, forecast recursively with
future innovations set to zero and undo ordinary/seasonal differencing. Invalid horizon
or non-finite predictions fail.
-}
forecastARIMA :: Int -> ARIMAModel -> Either MathError [Double]
forecastARIMA :: Int -> ARIMAModel -> Either MathError [Double]
forecastARIMA Int
horizon (ARIMAModel SARIMAOrder
order [Double]
b [Double]
xs [Double]
zs [Double]
es Double
_) = do
    Bool -> String -> Either MathError ()
require (Int
horizon Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
horizon Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
16384) String
"invalid SARIMA forecast horizon"
    let ([Double]
ar, [Double]
ma) = SARIMAOrder -> [Double] -> ([Double], [Double])
modelPolynomials SARIMAOrder
order [Double]
b
        diff :: [Double]
diff =
            ([Double] -> [Double] -> [Double])
-> [Double] -> [[Double]] -> [Double]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
                [Double] -> [Double] -> [Double]
polynomialProduct
                [Double
1]
                ( Int -> [Double] -> [[Double]]
forall a. Int -> a -> [a]
replicate (SARIMAOrder -> Int
differenceOrder SARIMAOrder
order) [Double
1, -Double
1]
                    [[Double]] -> [[Double]] -> [[Double]]
forall a. [a] -> [a] -> [a]
++ Int -> [Double] -> [[Double]]
forall a. Int -> a -> [a]
replicate (SARIMAOrder -> Int
seasonalDifferenceOrder SARIMAOrder
order) (Int -> [Double] -> [Double]
seasonalPolynomial (SARIMAOrder -> Int
seasonalPeriod SARIMAOrder
order) [Double
1, -Double
1])
                )
        loop :: t
-> [Double]
-> [Double]
-> [Double]
-> [Double]
-> Either MathError [Double]
loop t
0 [Double]
_ [Double]
_ [Double]
_ [Double]
acc = [Double] -> Either MathError [Double]
forall a b. b -> Either a b
Right ([Double] -> [Double]
forall a. [a] -> [a]
reverse [Double]
acc)
        loop t
k [Double]
original [Double]
transformed [Double]
errors [Double]
acc = do
            let z :: Double
z = ([Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) [Double]
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
ar [Double]
transformed) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
ma [Double]
errors)
                value :: Double
value = Double
z Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) ((Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1) [Double]
diff) [Double]
original)
            _ <- [Double] -> Either MathError [Double]
checked [Double
z, Double
value]
            loop (k - 1) (value : original) (z : transformed) (0 : errors) (value : acc)
    Int
-> [Double]
-> [Double]
-> [Double]
-> [Double]
-> Either MathError [Double]
forall {t}.
(Eq t, Num t) =>
t
-> [Double]
-> [Double]
-> [Double]
-> [Double]
-> Either MathError [Double]
loop Int
horizon ([Double] -> [Double]
forall a. [a] -> [a]
reverse [Double]
xs) ([Double] -> [Double]
forall a. [a] -> [a]
reverse [Double]
zs) ([Double] -> [Double]
forall a. [a] -> [a]
reverse [Double]
es) []

{- | Opaque fitted vector autoregression with intercept/lag coefficients and chronological
multivariate history.
-}
data VARModel = VARModel Int [[Double]] [[Double]] deriving (VARModel -> VARModel -> Bool
(VARModel -> VARModel -> Bool)
-> (VARModel -> VARModel -> Bool) -> Eq VARModel
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: VARModel -> VARModel -> Bool
== :: VARModel -> VARModel -> Bool
$c/= :: VARModel -> VARModel -> Bool
/= :: VARModel -> VARModel -> Bool
Eq, Int -> VARModel -> ShowS
[VARModel] -> ShowS
VARModel -> String
(Int -> VARModel -> ShowS)
-> (VARModel -> String) -> ([VARModel] -> ShowS) -> Show VARModel
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> VARModel -> ShowS
showsPrec :: Int -> VARModel -> ShowS
$cshow :: VARModel -> String
show :: VARModel -> String
$cshowList :: [VARModel] -> ShowS
showList :: [VARModel] -> ShowS
Show)

-- | Rows are response components, columns are intercept then lag-major predictors.
varCoefficients :: VARModel -> [[Double]]
varCoefficients :: VARModel -> [[Double]]
varCoefficients (VARModel Int
_ [[Double]]
b [[Double]]
_) = [[Double]]
b

{- | Given lag order then chronological multivariate rows, fit response-wise OLS with
intercept and lag-major features. Require finite rectangular data, dimensions\<=16,
rows\<=2048, order in [1,16], sufficient samples and \<=20000000 fitting work.
Invalid/ill-conditioned fits fail.
-}
fitVAR :: Int -> [[Double]] -> Either MathError VARModel
fitVAR :: Int -> [[Double]] -> Either MathError VARModel
fitVAR Int
order [[Double]]
xs = do
    Bool -> String -> Either MathError ()
require (Bool -> Bool
not ([[Double]] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Double]]
xs) Bool -> Bool -> Bool
&& Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (([[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) [[Double]]
xs))) String
"empty VAR data"
    let d :: Int
d = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (([[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) [[Double]]
xs); n :: Int
n = [[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
xs; width :: Int
width = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
order Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
d
    Bool -> String -> Either MathError ()
require
        ( Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
16
            Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
2048
            Bool -> Bool -> Bool
&& Int
order Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1
            Bool -> Bool -> Bool
&& Int
order Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
16
            Bool -> Bool -> Bool
&& ([Double] -> Bool) -> [[Double]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
d) (Int -> Bool) -> ([Double] -> Int) -> [Double] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) [[Double]]
xs
            Bool -> Bool -> Bool
&& Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
order Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
width
            Bool -> Bool -> Bool
&& Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
width Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
width Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
20000000
        )
        String
"VAR shape/order/work limit"
    _ <- [Double] -> Either MathError [Double]
checked ([[Double]] -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Double]]
xs)
    let design = [Double
1 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [[Double]] -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[[Double]]
xs [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! (Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
j) | Int
j <- [Int
1 .. Int
order]] | Int
t <- [Int
order .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    b <- mapM (ols design) (transpose (drop order xs))
    Right (VARModel order b xs)

{- | Given horizon then VAR model, recursively return future multivariate observation rows.
Require horizon in [0,16384] and horizon*components\<=262144; invalid bounds or
non-finite predictions fail.
-}
forecastVAR :: Int -> VARModel -> Either MathError [[Double]]
forecastVAR :: Int -> VARModel -> Either MathError [[Double]]
forecastVAR Int
horizon (VARModel Int
order [[Double]]
b [[Double]]
xs) = do
    Bool -> String -> Either MathError ()
require
        (Int
horizon Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
horizon Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
16384 Bool -> Bool -> Bool
&& Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
horizon Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
b) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
262144)
        String
"VAR forecast limit"
    let loop :: t -> [[Double]] -> [[Double]] -> Either MathError [[Double]]
loop t
0 [[Double]]
_ [[Double]]
acc = [[Double]] -> Either MathError [[Double]]
forall a b. b -> Either a b
Right ([[Double]] -> [[Double]]
forall a. [a] -> [a]
reverse [[Double]]
acc)
        loop t
k [[Double]]
past [[Double]]
acc = do
            let features :: [Double]
features = Double
1 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [[Double]] -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Int -> [[Double]] -> [[Double]]
forall a. Int -> [a] -> [a]
take Int
order [[Double]]
past); next :: [Double]
next = [[Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
row [Double]
features) | [Double]
row <- [[Double]]
b]
            _ <- [Double] -> Either MathError [Double]
checked [Double]
next
            loop (k - 1) (next : past) (next : acc)
    Int -> [[Double]] -> [[Double]] -> Either MathError [[Double]]
forall {t}.
(Eq t, Num t) =>
t -> [[Double]] -> [[Double]] -> Either MathError [[Double]]
loop Int
horizon ([[Double]] -> [[Double]]
forall a. [a] -> [a]
reverse [[Double]]
xs) []

{- | Reuse the checked Joseph-form linear Kalman update for each complete
observation. Returns posterior states, excluding the caller's initial state.
-}
kalmanSeries ::
    [[Double]] ->
    [[Double]] ->
    [[Double]] ->
    [[Double]] ->
    S.GaussianState ->
    [[Double]] ->
    Either MathError [S.GaussianState]
kalmanSeries :: [[Double]]
-> [[Double]]
-> [[Double]]
-> [[Double]]
-> GaussianState
-> [[Double]]
-> Either MathError [GaussianState]
kalmanSeries [[Double]]
transition [[Double]]
observation [[Double]]
q [[Double]]
r GaussianState
initial [[Double]]
measured = do
    Bool -> String -> Either MathError ()
require ([[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
measured Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
16384) String
"Kalman series limit"
    (_, states) <-
        ((GaussianState, [GaussianState])
 -> [Double] -> Either MathError (GaussianState, [GaussianState]))
-> (GaussianState, [GaussianState])
-> [[Double]]
-> Either MathError (GaussianState, [GaussianState])
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
            (\(GaussianState
old, [GaussianState]
acc) [Double]
y -> do next <- [[Double]]
-> [[Double]]
-> [[Double]]
-> [[Double]]
-> [Double]
-> GaussianState
-> Either MathError GaussianState
S.kalman [[Double]]
transition [[Double]]
observation [[Double]]
q [[Double]]
r [Double]
y GaussianState
old; Right (next, next : acc))
            (GaussianState
initial, [])
            [[Double]]
measured
    Right (reverse states)

ols :: [[Double]] -> [Double] -> Either MathError [Double]
ols :: [[Double]] -> [Double] -> Either MathError [Double]
ols [[Double]]
rows [Double]
target = do
    let n :: Int
n = [[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
rows; p :: Int
p = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (([[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) [[Double]]
rows)
    matrix <- Int -> Int -> [Double] -> Either MathError (Matrix Double)
forall a. Int -> Int -> [a] -> Either MathError (Matrix a)
M.fromListChecked Int
n Int
p ([[Double]] -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Double]]
rows)
    QR.leastSquares matrix target >>= checked