<commit_before>{-# LANGUAGE TypeFamilies, RankNTypes, FlexibleContexts #-}
module Language.Hakaru.Inference where

import Prelude hiding (Real)

import Language.Hakaru.Lazy
import Language.Hakaru.Compose
import Language.Hakaru.Syntax
import Language.Hakaru.Expect (Expect(..), Expect', normalize)

priorAsProposal :: Mochastic repr =>
                   repr (Measure (a,b)) -> repr (a,b) -> repr (Measure (a,b))
priorAsProposal p x = bern (1/2) `bind` \c ->
                      p `bind` \x' ->
                      dirac (if_ c
                             (pair (fst_ x ) (snd_ x'))
                             (pair (fst_ x') (snd_ x )))

mh :: (Mochastic repr, Integrate repr, Lambda repr,
       a ~ Expect' a, Order_ a, Backward a a) =>
      (forall repr'. (Mochastic repr') => repr' a -> repr' (Measure a)) ->
      (forall repr'. (Mochastic repr') => repr' (Measure a)) ->
      repr (a -> Measure (a, Prob))
mh proposal target =
  let_ (lam (d unit)) $ \mu ->
  lam $ \old ->
    proposal old `bind` \new ->
    dirac (pair new (mu `app` pair new old / mu `app` pair old new))
  where d:_ = density (\dummy -> ununit dummy $ bindx target proposal)

mcmc :: (Mochastic repr, Integrate repr, Lambda repr,
         a ~ Expect' a, Order_ a, Backward a a) =>
        (forall repr'. (Mochastic repr') => repr' a -> repr' (Measure a)) ->
        (forall repr'. (Mochastic repr') => repr' (Measure a)) ->
        repr (a -> Measure a)
mcmc proposal target =
  let_ (mh proposal target) $ \f ->
  lam $ \old ->
    app f old `bind` \new_ratio ->
    unpair new_ratio $ \new ratio ->
    bern (min_ 1 ratio) `bind` \accept ->
    dirac (if_ accept new old)

gibbsProposal :: (Expect' a ~ a, Expect' b ~ b,
                  Backward a a, Order_ a, 
                  Mochastic repr, Integrate repr, Lambda repr) =>
                 Cond (Expect repr) () (Measure (a,b)) ->
                 repr (a, b) -> repr (Measure (a, b))
gibbsProposal p x = q (fst_ x) `bind` \x' -> dirac (pair (fst_ x) x')
  where d:_ = runDisintegrate p
        q y = normalize (app (app d unit) (Expect y))              

-- Slice sampling can be thought of:
--
-- slice target x = do
--      u  <- uniform(0, density(target, x))  
--      x' <- lebesgue
--      condition (density(target, x') >= u) true
--      return x'

-- un-normalized slice
slice :: (Mochastic repr, Integrate repr, Lambda repr) =>
         (forall s t. Lazy s (Compose [] t (Expect repr))
                             (Measure Real)) ->
         repr (Real -> Measure Real)
slice target = lam $ \x ->
               uniform 0 (densAt x) `bind` \u ->
               lebesgue `bind` \x' ->
               if_ (u `less_` densAt x')
                 (dirac x')
                 (superpose [])
  where d:_ = density (\dummy -> ununit dummy target)
        densAt x = fromProb $ d unit x

sliceX :: (Mochastic repr, Integrate repr, Lambda repr,
          a ~ Expect' a, Order_ a, Backward a a) =>
          (forall repr' . (Mochastic repr') => repr' (Measure a)) ->
          repr (Measure (a, Real))
sliceX target = target `bindx` \x ->
                uniform 0 (densAt x)
  where d:_ = density (\dummy -> ununit dummy target)
        densAt x = fromProb $ d unit x

incompleteBeta :: Integrate repr => repr Prob -> repr Prob -> repr Prob -> repr Prob
incompleteBeta x a b = integrate 0 (fromProb x)
                       (\t -> pow_ (unsafeProb t    ) (fromProb a-1) *
                              pow_ (unsafeProb $ 1-t) (fromProb b-1))
<commit_msg>slice is normalized now, code should be nicer<commit_after>{-# LANGUAGE TypeFamilies, RankNTypes, FlexibleContexts #-}
module Language.Hakaru.Inference where

import Prelude hiding (Real)

import Language.Hakaru.Lazy
import Language.Hakaru.Compose
import Language.Hakaru.Syntax
import Language.Hakaru.Expect (Expect(..), Expect', normalize)

priorAsProposal :: Mochastic repr =>
                   repr (Measure (a,b)) -> repr (a,b) -> repr (Measure (a,b))
priorAsProposal p x = bern (1/2) `bind` \c ->
                      p `bind` \x' ->
                      dirac (if_ c
                             (pair (fst_ x ) (snd_ x'))
                             (pair (fst_ x') (snd_ x )))

mh :: (Mochastic repr, Integrate repr, Lambda repr,
       a ~ Expect' a, Order_ a, Backward a a) =>
      (forall repr'. (Mochastic repr') => repr' a -> repr' (Measure a)) ->
      (forall repr'. (Mochastic repr') => repr' (Measure a)) ->
      repr (a -> Measure (a, Prob))
mh proposal target =
  let_ (lam (d unit)) $ \mu ->
  lam $ \old ->
    proposal old `bind` \new ->
    dirac (pair new (mu `app` pair new old / mu `app` pair old new))
  where d:_ = density (\dummy -> ununit dummy $ bindx target proposal)

mcmc :: (Mochastic repr, Integrate repr, Lambda repr,
         a ~ Expect' a, Order_ a, Backward a a) =>
        (forall repr'. (Mochastic repr') => repr' a -> repr' (Measure a)) ->
        (forall repr'. (Mochastic repr') => repr' (Measure a)) ->
        repr (a -> Measure a)
mcmc proposal target =
  let_ (mh proposal target) $ \f ->
  lam $ \old ->
    app f old `bind` \new_ratio ->
    unpair new_ratio $ \new ratio ->
    bern (min_ 1 ratio) `bind` \accept ->
    dirac (if_ accept new old)

gibbsProposal :: (Expect' a ~ a, Expect' b ~ b,
                  Backward a a, Order_ a, 
                  Mochastic repr, Integrate repr, Lambda repr) =>
                 Cond (Expect repr) () (Measure (a,b)) ->
                 repr (a, b) -> repr (Measure (a, b))
gibbsProposal p x = q (fst_ x) `bind` \x' -> dirac (pair (fst_ x) x')
  where d:_ = runDisintegrate p
        q y = normalize (app (app d unit) (Expect y))              

-- Slice sampling can be thought of:
--
-- slice target x = do
--      u  <- uniform(0, density(target, x))  
--      x' <- lebesgue
--      condition (density(target, x') >= u) true
--      return x'

slice :: (Mochastic repr, Integrate repr, Lambda repr) =>
         (forall repr' . (Mochastic repr') => repr' (Measure Real)) ->
         repr (Real -> Measure Real)
slice target = lam $ \x ->
               uniform 0 (densAt x) `bind` \u ->
               slice' target u
  where d:_ = density (\dummy -> ununit dummy target)
        densAt x = fromProb $ d unit x

slice' :: (Mochastic repr, Integrate repr, Lambda repr) =>
       (forall repr' . (Mochastic repr') => repr' (Measure Real)) ->
       repr Real -> repr (Measure Real)
slice' target u = normalize $
        lebesgue `bind` \x' ->
        if_ (Expect u `less_` densAt x')
            (dirac x')
            (superpose [])
  where d:_ = density (\dummy -> ununit dummy target)
        densAt x = fromProb $ d unit x


sliceX :: (Mochastic repr, Integrate repr, Lambda repr,
          a ~ Expect' a, Order_ a, Backward a a) =>
          (forall repr' . (Mochastic repr') => repr' (Measure a)) ->
          repr (Measure (a, Real))
sliceX target = target `bindx` \x ->
                uniform 0 (densAt x)
  where d:_ = density (\dummy -> ununit dummy target)
        densAt x = fromProb $ d unit x

incompleteBeta :: Integrate repr => repr Prob -> repr Prob -> repr Prob -> repr Prob
incompleteBeta x a b = integrate 0 (fromProb x)
                       (\t -> pow_ (unsafeProb t    ) (fromProb a-1) *
                              pow_ (unsafeProb $ 1-t) (fromProb b-1))

{'fraction_non_alphanumeric': 0.1398380677919583, 'fraction_numerical': 0.0027446136956223413, 'mean_word_length': 1.969845150774246, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 109, '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': '6931589', 'n_tokens_mistral': 2515, 'n_tokens_neox': 2302, 'n_words': 859}