{- | Gaussian processes in the `LazyPPL` library.

Gaussian processes are random functions. We provide the Wiener process (`wiener`) and a generic Gaussian process (`gp`). 

Although the implementation uses hidden state, it is still safe, i.e. statistically commutative and discardable. 

For illustrations, see the [Gaussian process regression](https://lazyppl-team.github.io/GaussianProcessDemo.html) 
or [Wiener process regression](https://lazyppl-team.github.io/WienerDemo.html). The latter also illustrates composing the Wiener process function to get jump diffusion.
-}


module LazyPPL.Distributions.GP (wiener, gp) where

import LazyPPL (Prob,Tree(Tree),Prob(Prob),runProb)
import LazyPPL.Distributions (normal,iid)
import Numeric.LinearAlgebra hiding (step,size,find)
import Numeric.LinearAlgebra.Data hiding (step,size,find)
import Data.List (find)
import Data.Map (empty,lookup,insert,size,keys,elems,Map)
import Data.IORef
import System.IO.Unsafe
import Prelude hiding ((<>))

{-| [Wiener process](https://en.wikipedia.org/wiki/Wiener_process) (Brownian motion).

This is a random function. 

The Wiener process is implemented by using a Brownian bridge and a hidden memo table. -}
wiener :: Prob (Double -> Double)
wiener :: Prob (Double -> Double)
wiener = (Tree -> Double -> Double) -> Prob (Double -> Double)
forall a. (Tree -> a) -> Prob a
Prob ((Tree -> Double -> Double) -> Prob (Double -> Double))
-> (Tree -> Double -> Double) -> Prob (Double -> Double)
forall a b. (a -> b) -> a -> b
$ \(Tree Double
r [Tree]
gs) ->
                   IO (Double -> Double) -> Double -> Double
forall a. IO a -> a
unsafePerformIO (IO (Double -> Double) -> Double -> Double)
-> IO (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ do
                                 IORef (Map Double Double)
ref <- Map Double Double -> IO (IORef (Map Double Double))
forall a. a -> IO (IORef a)
newIORef Map Double Double
forall k a. Map k a
Data.Map.empty
                                 IORef (Map Double Double)
-> (Map Double Double -> Map Double Double) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef (Map Double Double)
ref (Double -> Double -> Map Double Double -> Map Double Double
forall k a. Ord k => k -> a -> Map k a -> Map k a
Data.Map.insert Double
0 Double
0)
                                 (Double -> Double) -> IO (Double -> Double)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Double -> Double) -> IO (Double -> Double))
-> (Double -> Double) -> IO (Double -> Double)
forall a b. (a -> b) -> a -> b
$ \Double
x -> IO Double -> Double
forall a. IO a -> a
unsafePerformIO (IO Double -> Double) -> IO Double -> Double
forall a b. (a -> b) -> a -> b
$ do
                                        Map Double Double
table <- IORef (Map Double Double) -> IO (Map Double Double)
forall a. IORef a -> IO a
readIORef IORef (Map Double Double)
ref
                                        case Double -> Map Double Double -> Maybe Double
forall k a. Ord k => k -> Map k a -> Maybe a
Data.Map.lookup Double
x Map Double Double
table of
                                             Just Double
y -> do {Double -> IO Double
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Double
y}
                                             Maybe Double
Nothing -> do let lower :: Maybe (Double, Double)
lower = do {Double
l <- Double -> [Double] -> Maybe Double
findMaxLower Double
x (Map Double Double -> [Double]
forall k a. Map k a -> [k]
keys Map Double Double
table) ;
                                                                           Double
v <- Double -> Map Double Double -> Maybe Double
forall k a. Ord k => k -> Map k a -> Maybe a
Data.Map.lookup Double
l Map Double Double
table ; (Double, Double) -> Maybe (Double, Double)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double
l,Double
v) }
                                                           let upper :: Maybe (Double, Double)
upper = do {Double
u <- (Double -> Bool) -> [Double] -> Maybe Double
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
x) (Map Double Double -> [Double]
forall k a. Map k a -> [k]
keys Map Double Double
table) ;
                                                                           Double
v <- Double -> Map Double Double -> Maybe Double
forall k a. Ord k => k -> Map k a -> Maybe a
Data.Map.lookup Double
u Map Double Double
table ; (Double, Double) -> Maybe (Double, Double)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double
u,Double
v) }
                                                           let m :: Prob Double
m = Maybe (Double, Double)
-> Double -> Maybe (Double, Double) -> Prob Double
bridge Maybe (Double, Double)
lower Double
x Maybe (Double, Double)
upper
                                                           let y :: Double
y = Prob Double -> Tree -> Double
forall a. Prob a -> Tree -> a
runProb Prob Double
m ([Tree]
gs [Tree] -> Int -> Tree
forall a. HasCallStack => [a] -> Int -> a
!! (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Map Double Double -> Int
forall k a. Map k a -> Int
size Map Double Double
table))
                                                           IORef (Map Double Double)
-> (Map Double Double -> Map Double Double) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef (Map Double Double)
ref (Double -> Double -> Map Double Double -> Map Double Double
forall k a. Ord k => k -> a -> Map k a -> Map k a
Data.Map.insert Double
x Double
y)
                                                           Double -> IO Double
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Double
y

bridge :: Maybe (Double,Double) -> Double -> Maybe (Double,Double) -> Prob Double
-- not needed since the table is always initialized with (0, 0)
-- bridge Nothing y Nothing = if y==0 then return 0 else normal 0 (sqrt y) 
bridge :: Maybe (Double, Double)
-> Double -> Maybe (Double, Double) -> Prob Double
bridge (Just (Double
x,Double
x')) Double
y Maybe (Double, Double)
Nothing = Double -> Double -> Prob Double
normal Double
x' (Double -> Double
forall a. Floating a => a -> a
sqrt (Double
yDouble -> Double -> Double
forall a. Num a => a -> a -> a
-Double
x))
bridge Maybe (Double, Double)
Nothing Double
y (Just (Double
z,Double
z')) = Double -> Double -> Prob Double
normal Double
z' (Double -> Double
forall a. Floating a => a -> a
sqrt (Double
zDouble -> Double -> Double
forall a. Num a => a -> a -> a
-Double
y))
bridge (Just (Double
x,Double
x')) Double
y (Just (Double
z,Double
z')) = Double -> Double -> Prob Double
normal (Double
x' Double -> Double -> Double
forall a. Num a => a -> a -> a
+ ((Double
yDouble -> Double -> Double
forall a. Num a => a -> a -> a
-Double
x)Double -> Double -> Double
forall a. Num a => a -> a -> a
*(Double
z'Double -> Double -> Double
forall a. Num a => a -> a -> a
-Double
x')Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/(Double
zDouble -> Double -> Double
forall a. Num a => a -> a -> a
-Double
x))) (Double -> Double
forall a. Floating a => a -> a
sqrt ((Double
zDouble -> Double -> Double
forall a. Num a => a -> a -> a
-Double
y)Double -> Double -> Double
forall a. Num a => a -> a -> a
*(Double
yDouble -> Double -> Double
forall a. Num a => a -> a -> a
-Double
x)Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/(Double
zDouble -> Double -> Double
forall a. Num a => a -> a -> a
-Double
x)))

findMaxLower :: Double -> [Double] -> Maybe Double 
findMaxLower :: Double -> [Double] -> Maybe Double
findMaxLower Double
d [] = Maybe Double
forall a. Maybe a
Nothing
findMaxLower Double
d (Double
x:[Double]
xs) = let y :: Maybe Double
y = Double -> [Double] -> Maybe Double
findMaxLower Double
d [Double]
xs in
                       case Maybe Double
y of 
                           Maybe Double
Nothing -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
d then Double -> Maybe Double
forall a. a -> Maybe a
Just Double
x else Maybe Double
forall a. Maybe a
Nothing 
                           Just Double
m -> do 
                                          if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
m Bool -> Bool -> Bool
&& Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
d then Double -> Maybe Double
forall a. a -> Maybe a
Just Double
x else Double -> Maybe Double
forall a. a -> Maybe a
Just Double
m 




{-| [Gaussian process](https://en.wikipedia.org/wiki/Gaussian_process) with given covariance function.

This is a random function.

The function is defined by using linear algebra and a hidden memo table.
Although it uses hidden state, it is still safe, i.e. statistically commutative and discardable. --}
gp :: (Double -> Double -> Double) -- ^ Covariance function
      -> Prob (Double -> Double) -- ^ Returns a random function
gp :: (Double -> Double -> Double) -> Prob (Double -> Double)
gp Double -> Double -> Double
cov = do [Double]
ns <- Prob Double -> Prob [Double]
forall a. Prob a -> Prob [a]
iid (Prob Double -> Prob [Double]) -> Prob Double -> Prob [Double]
forall a b. (a -> b) -> a -> b
$ Double -> Double -> Prob Double
normal Double
0 Double
1
            (Double -> Double) -> Prob (Double -> Double)
forall a. a -> Prob a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Double -> Double) -> Prob (Double -> Double))
-> (Double -> Double) -> Prob (Double -> Double)
forall a b. (a -> b) -> a -> b
$ IO (Double -> Double) -> Double -> Double
forall a. IO a -> a
unsafePerformIO (IO (Double -> Double) -> Double -> Double)
-> IO (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ do
                                 IORef (Map Double Double)
ref <- Map Double Double -> IO (IORef (Map Double Double))
forall a. a -> IO (IORef a)
newIORef Map Double Double
forall k a. Map k a
Data.Map.empty
                                 IORef (Map Double Double)
-> (Map Double Double -> Map Double Double) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef (Map Double Double)
ref (Double -> Double -> Map Double Double -> Map Double Double
forall k a. Ord k => k -> a -> Map k a -> Map k a
Data.Map.insert Double
0 Double
0)
                                 (Double -> Double) -> IO (Double -> Double)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Double -> Double) -> IO (Double -> Double))
-> (Double -> Double) -> IO (Double -> Double)
forall a b. (a -> b) -> a -> b
$ \Double
x -> IO Double -> Double
forall a. IO a -> a
unsafePerformIO (IO Double -> Double) -> IO Double -> Double
forall a b. (a -> b) -> a -> b
$ do
                                        Map Double Double
table <- IORef (Map Double Double) -> IO (Map Double Double)
forall a. IORef a -> IO a
readIORef IORef (Map Double Double)
ref
                                        case Double -> Map Double Double -> Maybe Double
forall k a. Ord k => k -> Map k a -> Maybe a
Data.Map.lookup Double
x Map Double Double
table of
                                             Just Double
y -> do {Double -> IO Double
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Double
y}
                                             Maybe Double
Nothing -> do let y :: Double
y = (Double -> Double -> Double)
-> Map Double Double -> Double -> Double -> Double
step Double -> Double -> Double
cov Map Double Double
table Double
x (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ [Double]
ns [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Map Double Double -> Int
forall k a. Map k a -> Int
size Map Double Double
table)
                                                           IORef (Map Double Double)
-> (Map Double Double -> Map Double Double) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef (Map Double Double)
ref (Double -> Double -> Map Double Double -> Map Double Double
forall k a. Ord k => k -> a -> Map k a -> Map k a
Data.Map.insert Double
x Double
y)
                                                           Double -> IO Double
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Double
y

step :: (Double -> Double -> Double) -> Data.Map.Map Double Double -> Double -> Double -> Double
step :: (Double -> Double -> Double)
-> Map Double Double -> Double -> Double -> Double
step Double -> Double -> Double
cov Map Double Double
table Double
x Double
seed =
       let sig11 :: Double
sig11 = Double -> Double -> Double
cov Double
x Double
x 
           sig12 :: Matrix Double
sig12 = Int -> [Double] -> Matrix Double
matrix (Map Double Double -> Int
forall k a. Map k a -> Int
size Map Double Double
table) [Double -> Double -> Double
cov Double
x Double
b | Double
b <-Map Double Double -> [Double]
forall k a. Map k a -> [k]
keys Map Double Double
table] 
           sig21 :: Matrix Double
sig21 = Int -> [Double] -> Matrix Double
matrix Int
1 [Double -> Double -> Double
cov Double
a Double
x | Double
a <-Map Double Double -> [Double]
forall k a. Map k a -> [k]
keys Map Double Double
table] 
           sig22 :: Matrix Double
sig22 = Int -> [Double] -> Matrix Double
matrix (Map Double Double -> Int
forall k a. Map k a -> Int
size Map Double Double
table) [Double -> Double -> Double
cov Double
a Double
b | Double
a <- Map Double Double -> [Double]
forall k a. Map k a -> [k]
keys Map Double Double
table , Double
b <-Map Double Double -> [Double]
forall k a. Map k a -> [k]
keys Map Double Double
table] 
           regCoeff :: Matrix Double
regCoeff = Matrix Double
sig12 Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> (Double -> Matrix Double -> Matrix Double
forall t. Field t => Double -> Matrix t -> Matrix t
pinvTol Double
1E8 Matrix Double
sig22)
           mu :: Double
mu = (Matrix Double
regCoeff Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> (Int -> [Double] -> Matrix Double
matrix Int
1 ([Double] -> Matrix Double) -> [Double] -> Matrix Double
forall a b. (a -> b) -> a -> b
$ Map Double Double -> [Double]
forall k a. Map k a -> [a]
elems Map Double Double
table)) Matrix Double -> IndexOf Matrix -> Double
forall (c :: * -> *) e. Container c e => c e -> IndexOf c -> e
`atIndex` (Int
0,Int
0)
           var :: Double
var = Double
sig11 Double -> Double -> Double
forall a. Num a => a -> a -> a
- ((Matrix Double
regCoeff Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double
sig21) Matrix Double -> IndexOf Matrix -> Double
forall (c :: * -> *) e. Container c e => c e -> IndexOf c -> e
`atIndex` (Int
0,Int
0)) in
       Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
seed Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ if Double
var Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> -Double
0.01 then (Double -> Double
forall a. Num a => a -> a
abs Double
var) else [Char] -> Double
forall a. HasCallStack => [Char] -> a
error (Double -> [Char]
forall a. Show a => a -> [Char]
show Double
var))