<commit_before>{-# LANGUAGE
  FlexibleContexts
 #-}

module Inference where

import Control.Arrow (first,second)
import Data.Either
import Data.Number.LogFloat
import Data.Typeable
import Control.Monad.Trans.Maybe
import Control.Monad.State.Lazy
import Control.Monad.Writer.Lazy

import Base
import Sampler
import Rejection
import Weighted
import Particle
import Empirical
import Dist

-- | Rejection sampling.
rejection :: RejectionT Sampler a -> Sampler a
rejection d = do
  m <- runRejectionT d
  case m of Just x  -> return x
            Nothing -> rejection d

-- | Simple importance sampling from the prior.
importance :: WeightedT Sampler a -> Sampler (a,LogFloat)
importance = runWeightedT

-- | Multiple importance samples with post-processing.
importance' :: (Ord a, Typeable a) => Int -> EmpiricalT Sampler a -> Sampler [(a,Double)]
importance' n d = fmap (enumerate . categorical) $ runEmpiricalT $ population n >> d

-- | Sequential Monte Carlo from the prior.
smc :: Int -> ParticleT (EmpiricalT Sampler) a -> EmpiricalT Sampler a
smc n d = fmap (\(Left x) -> x) $ run start where
    start = runParticleT $ lift (population n) >> d
    step particles = resample $ particles >>= advance
    run particles = do
      finished <- lift $ Empirical.all isLeft particles
      if finished then particles else run (step particles)

-- | `smc` with post-processing.
smc' :: (Ord a, Typeable a) => Int -> ParticleT (EmpiricalT Sampler) a -> Sampler [(a,Double)]
smc' n d = fmap (enumerate . categorical) $ runEmpiricalT $ smc n d



-- | Metropolis-Hastings kernel. Generates a new value and the MH ratio.
newtype MHKernel m a = MHKernel {runMHKernel :: a -> m (a,LogFloat)}

-- | Metropolis-Hastings algorithm. The idea is that the kernel handles the
-- part of ratio that results from MonadDist effects and transition function,
-- while state carries the factor associated with MonadBayes effects.<commit_msg>Made types of sampling inference algorithms polymorphic<commit_after>{-# LANGUAGE
  FlexibleContexts
 #-}

module Inference where

import Control.Arrow (first,second)
import Data.Either
import Data.Number.LogFloat
import Data.Typeable
import Control.Monad.Trans.Maybe
import Control.Monad.State.Lazy
import Control.Monad.Writer.Lazy

import Base
import Sampler
import Rejection
import Weighted
import Particle
import Empirical
import Dist

-- | Rejection sampling.
rejection :: MonadDist m => RejectionT m a -> m a
rejection d = do
  m <- runRejectionT d
  case m of Just x  -> return x
            Nothing -> rejection d

-- | Simple importance sampling from the prior.
importance :: MonadDist m => WeightedT m a -> m (a,LogFloat)
importance = runWeightedT

-- | Multiple importance samples with post-processing.
importance' :: (Ord a, Typeable a, MonadDist m) =>
               Int -> EmpiricalT m a -> m [(a,Double)]
importance' n d = fmap (enumerate . categorical) $ runEmpiricalT $ population n >> d

-- | Sequential Monte Carlo from the prior.
smc :: MonadDist m => Int -> ParticleT (EmpiricalT m) a -> EmpiricalT m a
smc n d = fmap (\(Left x) -> x) $ run start where
    start = runParticleT $ lift (population n) >> d
    step particles = resample $ particles >>= advance
    run particles = do
      finished <- lift $ Empirical.all isLeft particles
      if finished then particles else run (step particles)

-- | `smc` with post-processing.
smc' :: (Ord a, Typeable a, MonadDist m) => Int ->
        ParticleT (EmpiricalT m) a -> m [(a,Double)]
smc' n d = fmap (enumerate . categorical) $ runEmpiricalT $ smc n d



-- | Metropolis-Hastings kernel. Generates a new value and the MH ratio.
newtype MHKernel m a = MHKernel {runMHKernel :: a -> m (a,LogFloat)}

-- | Metropolis-Hastings algorithm. The idea is that the kernel handles the
-- part of ratio that results from MonadDist effects and transition function,
-- while state carries the factor associated with MonadBayes effects.
{'fraction_non_alphanumeric': 0.09565217391304348, 'fraction_numerical': 0.0, 'mean_word_length': 3.5742690058479534, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 46, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3778518', 'n_tokens_mistral': 1190, 'n_tokens_neox': 1150, 'n_words': 465}