module FRP.Rhine.Bayes where

-- transformers
import Control.Monad.Trans.Reader (ReaderT (..))

-- log-domain
import Numeric.Log hiding (sum)

-- monad-bayes
import Control.Monad.Bayes.Class
import Control.Monad.Bayes.Population

-- automaton
import qualified Data.Automaton.Trans.Reader as AutomatonReader

-- rhine-bayes
import qualified Data.Automaton.Bayes as AutomatonBayes

-- rhine
import FRP.Rhine

-- * Inference methods

-- | Run the Sequential Monte Carlo algorithm continuously on a 'ClSF'.
runPopulationCl ::
  forall m cl a b.
  (Monad m, MonadDistribution m) =>
  -- | Number of particles
  Int ->
  -- | Resampler (see 'Control.Monad.Bayes.PopulationT' for some standard choices)
  (forall x m. (MonadDistribution m) => PopulationT m x -> PopulationT m x) ->
  -- | A signal function modelling the stochastic process on which to perform inference.
  --   @a@ represents observations upon which the model should condition, using e.g. 'score'.
  --   It can also additionally contain hyperparameters.
  --   @b@ is the type of estimated current state.
  ClSF (PopulationT m) cl a b ->
  ClSF m cl a [(b, Log Double)]
runPopulationCl :: forall (m :: * -> *) cl a b.
(Monad m, MonadDistribution m) =>
Int
-> (forall x (m :: * -> *).
    MonadDistribution m =>
    PopulationT m x -> PopulationT m x)
-> ClSF (PopulationT m) cl a b
-> ClSF m cl a [(b, Log Double)]
runPopulationCl Int
nParticles forall x (m :: * -> *).
MonadDistribution m =>
PopulationT m x -> PopulationT m x
resampler = Automaton m (TimeInfo cl, a) [(b, Log Double)]
-> Automaton (ReaderT (TimeInfo cl) m) a [(b, Log Double)]
forall (m :: * -> *) r a b.
Monad m =>
Automaton m (r, a) b -> Automaton (ReaderT r m) a b
AutomatonReader.readerS (Automaton m (TimeInfo cl, a) [(b, Log Double)]
 -> Automaton (ReaderT (TimeInfo cl) m) a [(b, Log Double)])
-> (ClSF (PopulationT m) cl a b
    -> Automaton m (TimeInfo cl, a) [(b, Log Double)])
-> ClSF (PopulationT m) cl a b
-> Automaton (ReaderT (TimeInfo cl) m) a [(b, Log Double)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int
-> (forall x. PopulationT m x -> PopulationT m x)
-> Automaton (PopulationT m) (TimeInfo cl, a) b
-> Automaton m (TimeInfo cl, a) [(b, Log Double)]
forall (m :: * -> *) a b.
Monad m =>
Int
-> (forall x. PopulationT m x -> PopulationT m x)
-> Automaton (PopulationT m) a b
-> Automaton m a [(b, Log Double)]
AutomatonBayes.runPopulationS Int
nParticles PopulationT m x -> PopulationT m x
forall x. PopulationT m x -> PopulationT m x
forall x (m :: * -> *).
MonadDistribution m =>
PopulationT m x -> PopulationT m x
resampler (Automaton (PopulationT m) (TimeInfo cl, a) b
 -> Automaton m (TimeInfo cl, a) [(b, Log Double)])
-> (ClSF (PopulationT m) cl a b
    -> Automaton (PopulationT m) (TimeInfo cl, a) b)
-> ClSF (PopulationT m) cl a b
-> Automaton m (TimeInfo cl, a) [(b, Log Double)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ClSF (PopulationT m) cl a b
-> Automaton (PopulationT m) (TimeInfo cl, a) b
forall (m :: * -> *) r a b.
Monad m =>
Automaton (ReaderT r m) a b -> Automaton m (r, a) b
AutomatonReader.runReaderS

-- * Short standard library of stochastic processes

-- | A stochastic process is a behaviour that uses, as only effect, random sampling.
type StochasticProcess time a = forall m. (MonadDistribution m) => Behaviour m time a

-- | Like 'StochasticProcess', but with a live input.
type StochasticProcessF time a b = forall m. (MonadDistribution m) => BehaviourF m time a b

-- | White noise, that is, an independent normal distribution at every time step.
whiteNoise :: Double -> StochasticProcess td Double
whiteNoise :: forall td. Double -> StochasticProcess td Double
whiteNoise Double
sigma = m Double -> ClSF m cl arbitrary Double
forall (m :: * -> *) b cl a. Monad m => m b -> ClSF m cl a b
constMCl (m Double -> ClSF m cl arbitrary Double)
-> m Double -> ClSF m cl arbitrary Double
forall a b. (a -> b) -> a -> b
$ Double -> Double -> m Double
forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> m Double
normal Double
0 Double
sigma

-- | Like 'whiteNoise', that is, an independent normal distribution at every time step.
whiteNoiseVarying :: StochasticProcessF td Double Double
whiteNoiseVarying :: forall td (m :: * -> *).
MonadDistribution m =>
BehaviourF m td Double Double
whiteNoiseVarying = (Double -> m Double) -> ClSF m cl Double Double
forall (m :: * -> *) a b cl. Monad m => (a -> m b) -> ClSF m cl a b
arrMCl ((Double -> m Double) -> ClSF m cl Double Double)
-> (Double -> m Double) -> ClSF m cl Double Double
forall a b. (a -> b) -> a -> b
$ Double -> Double -> m Double
forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> m Double
normal Double
0

-- | Construct a Lévy process from the increment between time steps.
levy ::
  (MonadDistribution m, VectorSpace v (Diff td)) =>
  -- | The increment function at every time step. The argument is the difference between times.
  (Diff td -> m v) ->
  Behaviour m td v
levy :: forall (m :: * -> *) v td.
(MonadDistribution m, VectorSpace v (Diff td)) =>
(Diff td -> m v) -> Behaviour m td v
levy Diff td -> m v
incrementor = Automaton (ReaderT (TimeInfo cl) m) arbitrary (Diff td)
ClSF m cl arbitrary (Diff (Time cl))
forall (m :: * -> *) cl a. Monad m => ClSF m cl a (Diff (Time cl))
sinceLastS Automaton (ReaderT (TimeInfo cl) m) arbitrary (Diff td)
-> Automaton (ReaderT (TimeInfo cl) m) (Diff td) v
-> Automaton (ReaderT (TimeInfo cl) m) arbitrary v
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> (Diff td -> m v) -> Automaton (ReaderT (TimeInfo cl) m) (Diff td) v
forall (m :: * -> *) a b cl. Monad m => (a -> m b) -> ClSF m cl a b
arrMCl Diff td -> m v
incrementor Automaton (ReaderT (TimeInfo cl) m) (Diff td) v
-> Automaton (ReaderT (TimeInfo cl) m) v v
-> Automaton (ReaderT (TimeInfo cl) m) (Diff td) v
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Automaton (ReaderT (TimeInfo cl) m) v v
forall (m :: * -> *) v s.
(Monad m, VectorSpace v s) =>
Automaton m v v
sumS

-- | The Wiener process, also known as Brownian motion.
wiener
  , brownianMotion ::
    (MonadDistribution m, Diff td ~ Double) =>
    -- | Time scale of variance.
    Diff td ->
    Behaviour m td Double
wiener :: forall (m :: * -> *) td.
(MonadDistribution m, Diff td ~ Double) =>
Diff td -> Behaviour m td Double
wiener Diff td
timescale = (Diff td -> m Double) -> Behaviour m td Double
forall (m :: * -> *) v td.
(MonadDistribution m, VectorSpace v (Diff td)) =>
(Diff td -> m v) -> Behaviour m td v
levy ((Diff td -> m Double) -> Behaviour m td Double)
-> (Diff td -> m Double) -> Behaviour m td Double
forall a b. (a -> b) -> a -> b
$ \Diff td
diffTime -> Double -> Double -> m Double
forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> m Double
normal Double
0 (Double -> m Double) -> Double -> m Double
forall a b. (a -> b) -> a -> b
$ Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double
Diff td
diffTime Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
Diff td
timescale
brownianMotion :: forall (m :: * -> *) td.
(MonadDistribution m, Diff td ~ Double) =>
Diff td -> Behaviour m td Double
brownianMotion = Diff td -> ClSF m cl arbitrary Double
Diff td -> Behaviour m td Double
forall (m :: * -> *) td.
(MonadDistribution m, Diff td ~ Double) =>
Diff td -> Behaviour m td Double
wiener

-- | The Wiener process, also known as Brownian motion, with varying variance parameter.
wienerVarying
  , brownianMotionVarying ::
    (Diff td ~ Double) =>
    StochasticProcessF td (Diff td) Double
wienerVarying :: forall td.
(Diff td ~ Double) =>
StochasticProcessF td (Diff td) Double
wienerVarying = proc Diff td
timeScale -> do
  Double
diffTime <- Automaton (ReaderT (TimeInfo cl) m) () Double
ClSF m cl () (Diff (Time cl))
forall (m :: * -> *) cl a. Monad m => ClSF m cl a (Diff (Time cl))
sinceLastS -< ()
  let stdDev :: Double
stdDev = Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double
diffTime Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
Diff td
timeScale
  Double
increment <-
    if Double
stdDev Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0
      then (Double -> ReaderT (TimeInfo cl) m Double)
-> Automaton (ReaderT (TimeInfo cl) m) Double Double
forall (m :: * -> *) a b.
Functor m =>
(a -> m b) -> Automaton m a b
arrM ((Double -> ReaderT (TimeInfo cl) m Double)
 -> Automaton (ReaderT (TimeInfo cl) m) Double Double)
-> (Double -> ReaderT (TimeInfo cl) m Double)
-> Automaton (ReaderT (TimeInfo cl) m) Double Double
forall a b. (a -> b) -> a -> b
$ Double -> Double -> ReaderT (TimeInfo cl) m Double
forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> m Double
normal Double
0 -< Double
stdDev
      else Automaton (ReaderT (TimeInfo cl) m) Double Double
forall (a :: * -> * -> *) b. Arrow a => a b b
returnA -< Double
0
  Automaton (ReaderT (TimeInfo cl) m) Double Double
forall (m :: * -> *) v s.
(Monad m, VectorSpace v s) =>
Automaton m v v
sumS -< Double
increment
brownianMotionVarying :: forall td.
(Diff td ~ Double) =>
StochasticProcessF td (Diff td) Double
brownianMotionVarying = ClSF m cl (Diff td) Double
BehaviourF m td (Diff td) Double
forall td.
(Diff td ~ Double) =>
StochasticProcessF td (Diff td) Double
StochasticProcessF td (Diff td) Double
wienerVarying

-- | The 'wiener' process transformed to the Log domain, also called the geometric Wiener process.
wienerLogDomain ::
  (Diff td ~ Double) =>
  -- | Time scale of variance
  Diff td ->
  StochasticProcess td (Log Double)
wienerLogDomain :: forall td.
(Diff td ~ Double) =>
Diff td -> StochasticProcess td (Log Double)
wienerLogDomain Diff td
timescale = Diff td -> Behaviour m td Double
forall (m :: * -> *) td.
(MonadDistribution m, Diff td ~ Double) =>
Diff td -> Behaviour m td Double
wiener Diff td
timescale ClSF m cl arbitrary Double
-> Automaton (ReaderT (TimeInfo cl) m) Double (Log Double)
-> Automaton (ReaderT (TimeInfo cl) m) arbitrary (Log Double)
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> (Double -> Log Double)
-> Automaton (ReaderT (TimeInfo cl) m) Double (Log Double)
forall b c. (b -> c) -> Automaton (ReaderT (TimeInfo cl) m) b c
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr Double -> Log Double
forall a. a -> Log a
Exp

-- | See 'wienerLogDomain' and 'wienerVarying'.
wienerVaryingLogDomain ::
  (Diff td ~ Double) =>
  StochasticProcessF td (Diff td) (Log Double)
wienerVaryingLogDomain :: forall td.
(Diff td ~ Double) =>
StochasticProcessF td (Diff td) (Log Double)
wienerVaryingLogDomain = Automaton (ReaderT (TimeInfo cl) m) Double Double
ClSF m cl (Diff td) Double
BehaviourF m td (Diff td) Double
forall td.
(Diff td ~ Double) =>
StochasticProcessF td (Diff td) Double
StochasticProcessF td (Diff td) Double
wienerVarying Automaton (ReaderT (TimeInfo cl) m) Double Double
-> Automaton (ReaderT (TimeInfo cl) m) Double (Log Double)
-> Automaton (ReaderT (TimeInfo cl) m) Double (Log Double)
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> (Double -> Log Double)
-> Automaton (ReaderT (TimeInfo cl) m) Double (Log Double)
forall b c. (b -> c) -> Automaton (ReaderT (TimeInfo cl) m) b c
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr Double -> Log Double
forall a. a -> Log a
Exp

{- | Inhomogeneous Poisson point process, as described in:
  https://en.wikipedia.org/wiki/Poisson_point_process#Inhomogeneous_Poisson_point_process

  * The input is the inverse of the current rate or intensity.
    It corresponds to the average duration between two events.
  * The output is the number of events since the last tick.
-}
poissonInhomogeneous ::
  (MonadDistribution m, Real (Diff td), Fractional (Diff td)) =>
  BehaviourF m td (Diff td) Int
poissonInhomogeneous :: forall (m :: * -> *) td.
(MonadDistribution m, Real (Diff td), Fractional (Diff td)) =>
BehaviourF m td (Diff td) Int
poissonInhomogeneous = (Diff td -> ReaderT (TimeInfo cl) m Int)
-> Automaton (ReaderT (TimeInfo cl) m) (Diff td) Int
forall (m :: * -> *) a b.
Functor m =>
(a -> m b) -> Automaton m a b
arrM ((Diff td -> ReaderT (TimeInfo cl) m Int)
 -> Automaton (ReaderT (TimeInfo cl) m) (Diff td) Int)
-> (Diff td -> ReaderT (TimeInfo cl) m Int)
-> Automaton (ReaderT (TimeInfo cl) m) (Diff td) Int
forall a b. (a -> b) -> a -> b
$ \Diff td
rate -> (TimeInfo cl -> m Int) -> ReaderT (TimeInfo cl) m Int
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((TimeInfo cl -> m Int) -> ReaderT (TimeInfo cl) m Int)
-> (TimeInfo cl -> m Int) -> ReaderT (TimeInfo cl) m Int
forall a b. (a -> b) -> a -> b
$ \TimeInfo cl
timeInfo -> Double -> m Int
forall (m :: * -> *). MonadDistribution m => Double -> m Int
poisson (Double -> m Int) -> Double -> m Int
forall a b. (a -> b) -> a -> b
$ Diff td -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac (Diff td -> Double) -> Diff td -> Double
forall a b. (a -> b) -> a -> b
$ TimeInfo cl -> Diff (Time cl)
forall cl. TimeInfo cl -> Diff (Time cl)
sinceLast TimeInfo cl
timeInfo Diff td -> Diff td -> Diff td
forall a. Fractional a => a -> a -> a
/ Diff td
rate

-- | Like 'poissonInhomogeneous', but the rate is constant.
poissonHomogeneous ::
  (MonadDistribution m, Real (Diff td), Fractional (Diff td)) =>
  -- | The (constant) rate of the process
  Diff td ->
  BehaviourF m td () Int
poissonHomogeneous :: forall (m :: * -> *) td.
(MonadDistribution m, Real (Diff td), Fractional (Diff td)) =>
Diff td -> BehaviourF m td () Int
poissonHomogeneous Diff td
rate = (() -> Diff td) -> Automaton (ReaderT (TimeInfo cl) m) () (Diff td)
forall b c. (b -> c) -> Automaton (ReaderT (TimeInfo cl) m) b c
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (Diff td -> () -> Diff td
forall a b. a -> b -> a
const Diff td
rate) Automaton (ReaderT (TimeInfo cl) m) () (Diff td)
-> Automaton (ReaderT (TimeInfo cl) m) (Diff td) Int
-> Automaton (ReaderT (TimeInfo cl) m) () Int
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Automaton (ReaderT (TimeInfo cl) m) (Diff td) Int
BehaviourF m td (Diff td) Int
forall (m :: * -> *) td.
(MonadDistribution m, Real (Diff td), Fractional (Diff td)) =>
BehaviourF m td (Diff td) Int
poissonInhomogeneous

{- | The Gamma process, https://en.wikipedia.org/wiki/Gamma_process.

  The live input corresponds to inverse shape parameter, which is variance over mean.
-}
gammaInhomogeneous ::
  (MonadDistribution m, Real (Diff td), Fractional (Diff td), Floating (Diff td)) =>
  -- | The scale parameter
  Diff td ->
  BehaviourF m td (Diff td) Int
gammaInhomogeneous :: forall (m :: * -> *) td.
(MonadDistribution m, Real (Diff td), Fractional (Diff td),
 Floating (Diff td)) =>
Diff td -> BehaviourF m td (Diff td) Int
gammaInhomogeneous Diff td
gamma = proc Diff td
rate -> do
  Diff td
t <- Automaton (ReaderT (TimeInfo cl) m) () (Diff td)
ClSF m cl () (Diff (Time cl))
forall (m :: * -> *) cl a. Monad m => ClSF m cl a (Diff (Time cl))
sinceInitS -< ()
  (Int -> Int -> Int)
-> Int -> Automaton (ReaderT (TimeInfo cl) m) Int Int
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Automaton m a b
accumulateWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 Automaton (ReaderT (TimeInfo cl) m) Int Int
-> Automaton (ReaderT (TimeInfo cl) m) (Diff td) Int
-> Automaton (ReaderT (TimeInfo cl) m) (Diff td) Int
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
<<< Automaton (ReaderT (TimeInfo cl) m) (Diff td) Int
BehaviourF m td (Diff td) Int
forall (m :: * -> *) td.
(MonadDistribution m, Real (Diff td), Fractional (Diff td)) =>
BehaviourF m td (Diff td) Int
poissonInhomogeneous -< Diff td
gamma Diff td -> Diff td -> Diff td
forall a. Fractional a => a -> a -> a
/ Diff td
t Diff td -> Diff td -> Diff td
forall a. Num a => a -> a -> a
* Diff td -> Diff td
forall a. Floating a => a -> a
exp (-Diff td
t Diff td -> Diff td -> Diff td
forall a. Fractional a => a -> a -> a
/ Diff td
rate)

{- | The inhomogeneous Bernoulli process, https://en.wikipedia.org/wiki/Bernoulli_process

  Throws a coin to a given probability at each tick.
  The live input is the probability.
-}
bernoulliInhomogeneous :: (MonadDistribution m) => BehaviourF m td Double Bool
bernoulliInhomogeneous :: forall (m :: * -> *) td.
MonadDistribution m =>
BehaviourF m td Double Bool
bernoulliInhomogeneous = (Double -> m Bool) -> ClSF m cl Double Bool
forall (m :: * -> *) a b cl. Monad m => (a -> m b) -> ClSF m cl a b
arrMCl Double -> m Bool
forall (m :: * -> *). MonadDistribution m => Double -> m Bool
bernoulli