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

-- dunai
import qualified Control.Monad.Trans.MSF.Reader as DunaiReader

-- dunai-bayes
import qualified Data.MonadicStreamFunction.Bayes as DunaiBayes

-- rhine
import FRP.Rhine

-- * Inference methods

-- | Run the Sequential Monte Carlo algorithm continuously on a 'ClSF'.
runPopulationCl ::
  forall m cl a b.
  (Monad m) =>
  -- | Number of particles
  Int ->
  -- | Resampler (see 'Control.Monad.Bayes.Population' for some standard choices)
  (forall x. Population m x -> Population 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 (Population m) cl a b ->
  ClSF m cl a [(b, Log Double)]
runPopulationCl :: forall (m :: * -> *) cl a b.
Monad m =>
Int
-> (forall x. Population m x -> Population m x)
-> ClSF (Population m) cl a b
-> ClSF m cl a [(b, Log Double)]
runPopulationCl Int
nParticles forall x. Population m x -> Population m x
resampler = forall (m :: * -> *) r a b.
Monad m =>
MSF m (r, a) b -> MSF (ReaderT r m) a b
DunaiReader.readerS forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a b.
Monad m =>
Int
-> (forall x. Population m x -> Population m x)
-> MSF (Population m) a b
-> MSF m a [(b, Log Double)]
DunaiBayes.runPopulationS Int
nParticles forall x. Population m x -> Population m x
resampler forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) r a b.
Monad m =>
MSF (ReaderT r m) a b -> MSF m (r, a) b
DunaiReader.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 = forall (m :: * -> *) b cl a. Monad m => m b -> ClSF m cl a b
constMCl forall a b. (a -> b) -> a -> b
$ 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. StochasticProcessF td Double Double
whiteNoiseVarying = forall (m :: * -> *) a b cl. Monad m => (a -> m b) -> ClSF m cl a b
arrMCl forall a b. (a -> b) -> a -> b
$ 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 = forall (m :: * -> *) cl a. Monad m => ClSF m cl a (Diff (Time cl))
sinceLastS forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> forall (m :: * -> *) a b cl. Monad m => (a -> m b) -> ClSF m cl a b
arrMCl Diff td -> m v
incrementor forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> forall v s (m :: * -> *). (VectorSpace v s, Monad m) => MSF 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 = forall (m :: * -> *) v td.
(MonadDistribution m, VectorSpace v (Diff td)) =>
(Diff td -> m v) -> Behaviour m td v
levy forall a b. (a -> b) -> a -> b
$ \Diff td
diffTime -> forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> m Double
normal Double
0 forall a b. (a -> b) -> a -> b
$ forall a. Floating a => a -> a
sqrt forall a b. (a -> b) -> a -> b
$ Diff td
diffTime forall a. Fractional a => a -> a -> a
/ Diff td
timescale
brownianMotion :: forall (m :: * -> *) td.
(MonadDistribution m, Diff td ~ Double) =>
Diff td -> Behaviour m td Double
brownianMotion = 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 <- forall (m :: * -> *) cl a. Monad m => ClSF m cl a (Diff (Time cl))
sinceLastS -< ()
  let stdDev :: Double
stdDev = forall a. Floating a => a -> a
sqrt forall a b. (a -> b) -> a -> b
$ Double
diffTime forall a. Fractional a => a -> a -> a
/ Diff td
timeScale
  Double
increment <-
    if Double
stdDev forall a. Ord a => a -> a -> Bool
> Double
0
      then forall (m :: * -> *) a b. Monad m => (a -> m b) -> MSF m a b
arrM forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> m Double
normal Double
0 -< Double
stdDev
      else forall (a :: * -> * -> *) b. Arrow a => a b b
returnA -< Double
0
  forall v s (m :: * -> *). (VectorSpace v s, Monad m) => MSF m v v
sumS -< Double
increment
brownianMotionVarying :: forall td.
(Diff td ~ Double) =>
StochasticProcessF td (Diff td) Double
brownianMotionVarying = forall 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 = forall (m :: * -> *) td.
(MonadDistribution m, Diff td ~ Double) =>
Diff td -> Behaviour m td Double
wiener Diff td
timescale forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr 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 = forall td.
(Diff td ~ Double) =>
StochasticProcessF td (Diff td) Double
wienerVarying forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr 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 = forall (m :: * -> *) a b. Monad m => (a -> m b) -> MSF m a b
arrM forall a b. (a -> b) -> a -> b
$ \Diff td
rate -> forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT forall a b. (a -> b) -> a -> b
$ \TimeInfo cl
timeInfo -> forall (m :: * -> *). MonadDistribution m => Double -> m Int
poisson forall a b. (a -> b) -> a -> b
$ forall a b. (Real a, Fractional b) => a -> b
realToFrac forall a b. (a -> b) -> a -> b
$ forall cl. TimeInfo cl -> Diff (Time cl)
sinceLast TimeInfo cl
timeInfo 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 = forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (forall a b. a -> b -> a
const Diff td
rate) forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> 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 <- forall (m :: * -> *) cl a. Monad m => ClSF m cl a (Diff (Time cl))
sinceInitS -< ()
  forall (m :: * -> *) a s.
Monad m =>
(a -> s -> s) -> s -> MSF m a s
accumulateWith forall a. Num a => a -> a -> a
(+) Int
0 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
<<< forall (m :: * -> *) td.
(MonadDistribution m, Real (Diff td), Fractional (Diff td)) =>
BehaviourF m td (Diff td) Int
poissonInhomogeneous -< Diff td
gamma forall a. Fractional a => a -> a -> a
/ Diff td
t forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
exp (-Diff td
t 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 = forall (m :: * -> *) a b cl. Monad m => (a -> m b) -> ClSF m cl a b
arrMCl forall (m :: * -> *). MonadDistribution m => Double -> m Bool
bernoulli