<commit_before>{-# LANGUAGE BangPatterns #-}

module NestedSampling.Sampler where

import System.IO
import Control.Monad (replicateM, when)
import Control.Monad.Primitive (RealWorld)
import Data.List
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import System.Random.MWC (Gen)
import qualified System.Random.MWC as MWC

import NestedSampling.SpikeSlab
import NestedSampling.RNG
import NestedSampling.Utils

-- A sampler
data Sampler = Sampler
               {
                   numParticles      :: {-# UNPACK #-} !Int,
                   mcmcSteps         :: {-# UNPACK #-} !Int,
                   theParticles      :: !(V.Vector (U.Vector Double)),
                   theLogLikelihoods :: !(U.Vector Double),
                   iteration         :: {-# UNPACK #-} !Int,
                   logZ              :: {-# UNPACK #-} !Double,
                   information       :: {-# UNPACK #-} !Double
               } deriving Show

-- | Choose a particle to copy, that isn't number k.
chooseCopy :: Int -> Int -> Gen RealWorld -> IO Int
chooseCopy k n = loop where
  loop prng = do
    index <- MWC.uniformR (0, n - 1) prng
    if   index == k
    then loop prng
    else return index

-- Generate a sampler with n particles and m mcmc steps
generateSampler :: Int -> Int -> Gen RealWorld -> IO Sampler
generateSampler n m gen
    | n <= 1    = undefined
    | m <= 0    = undefined
    | otherwise = do
        putStr ("Generating " ++ (show n) ++ " particles\
                        \ from the prior...")
        hFlush stdout
        theParticles <- V.replicateM n (fromPrior gen)
        let lls = V.map logLikelihood theParticles
        putStrLn "done."
        return Sampler {numParticles=n, mcmcSteps=m,
                        theParticles=theParticles,
                        theLogLikelihoods=U.convert lls, iteration=1,
                        logZ = -1E300, information=0.0}

-- Find the index and the log likelihood value of the worst particle
--
-- NB (jtobin):
--   It may be better to use something like a heap so that we can access the
--   min element in O(1).
findWorstParticle :: Sampler -> (Int, Double)
findWorstParticle sampler = (idx, U.unsafeIndex lls idx)
    where
        idx = U.minIndex lls
        lls = theLogLikelihoods sampler

-- Function to do a single metropolis update
-- Input: A vector of parameters and a loglikelihood threshold
-- Output: An IO action which would return a new vector of parameters
-- and its log likelihood
metropolisUpdate :: Double -> ((U.Vector Double), Double) -> Gen RealWorld
                                -> IO ((U.Vector Double), Double)
metropolisUpdate threshold (x, logL) gen = do
    (proposal, logH) <- perturb x gen
    let a = exp logH
    uu <- rand gen
    let llProposal = logLikelihood proposal
    let accept = (uu < a) && (llProposal > threshold)
    return $
      if   accept
      then (proposal, llProposal)
      else (x, logL)

-- Function to do many metropolis updates
metropolisUpdates :: Int -> Double -> ((U.Vector Double), Double) -> Gen RealWorld
                                -> IO ((U.Vector Double), Double)
metropolisUpdates = loop where
  loop n threshold (x, logL) gen
    | n < 0     = undefined
    | n == 0    = return (x, logL)
    | otherwise = do
                    next  <- metropolisUpdate threshold (x, logL) gen
                    loop (n-1) threshold next gen
{-# INLINE metropolisUpdates #-}

-- Do many NestedSampling iterations
nestedSamplingIterations :: Int -> Sampler -> Gen RealWorld -> IO Sampler
nestedSamplingIterations = loop where
  loop n sampler gen
    | n < 0     = undefined
    | n == 0    = return sampler
    | otherwise = do
                    next <- nestedSamplingIteration sampler gen
                    loop (n-1) next gen
{-# INLINE nestedSamplingIterations #-}

-- Save a particle to disk
writeToFile :: Bool -> (Double, Double) -> U.Vector Double -> IO ()
writeToFile append (logw, logl) particle = do
    -- Open the file
    sampleInfo <- openFile "sample_info.txt"
                        (if append then AppendMode else WriteMode)

    hPutStrLn sampleInfo $ (show logw) ++ " " ++ (show logl)
    hClose sampleInfo

    -- Open the file
    sample <- openFile "sample.txt"
                        (if append then AppendMode else WriteMode)

    let particle' = U.toList particle
    hPutStrLn sample $ foldl' (++) [] $ [show x ++ " " | x <- particle']        
    hClose sample

-- Do a single NestedSampling iteration
nestedSamplingIteration :: Sampler -> Gen RealWorld -> IO Sampler
nestedSamplingIteration sampler gen = do
    let it = iteration sampler
    let logX = negate $ (fromIntegral it) / (fromIntegral (numParticles sampler)) :: Double
    let worst = findWorstParticle sampler

    let iWorst = fst worst
    let n = numParticles sampler
    copy <- chooseCopy iWorst n gen

    -- Approximate log prior weight of the worst particle
    let logPriorWeight = -(k/n) + log (exp (1.0/n) - 1.0) where
          k = fromIntegral it
          n = fromIntegral (numParticles sampler)
    let logPost = logPriorWeight + logLike where
          logLike = snd worst
    let logZ' = logsumexp (logZ sampler) logPost 


    let information' = information sampler
--    let information = exp 

    --H = exp(Obj[worst].logWt - logZnew) * Obj[worst].logL
    -- + exp(logZ - logZnew) * (H + logZ) - logZnew;
    --logZ = logZnew;

    -- Print some stuff from time to time
    let display = it `mod` (numParticles sampler) == 0 :: Bool
    if display then do
        putStr $ "Iteration " ++ (show it) ++ ". "
        putStr $ "log(X) = " ++ (show logX) ++ ". "
        putStrLn $ "log(L) = " ++ (show $ snd worst) ++ ". "
        putStrLn $ "log(Z) = " ++ (show logZ') ++ ".\n"

    else return ()
--        putStr $ "log(L) = " ++ (show $ snd worst) ++ ". " }
--        when (it == 1)


    -- Write to file
    writeToFile (it /= 1) (logPriorWeight, snd worst)
                        $ (theParticles sampler) V.! iWorst


    -- Copy a surviving particle
    let particle = (V.unsafeIndex (theParticles sampler) copy,
                        U.unsafeIndex (theLogLikelihoods sampler) copy)

    -- Do Metropolis
    let update = metropolisUpdates (mcmcSteps sampler) (snd worst)<commit_msg>Calculate and print the information<commit_after>{-# LANGUAGE BangPatterns #-}

module NestedSampling.Sampler where

import System.IO
import Control.Monad (replicateM, when)
import Control.Monad.Primitive (RealWorld)
import Data.List
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import System.Random.MWC (Gen)
import qualified System.Random.MWC as MWC

import NestedSampling.SpikeSlab
import NestedSampling.RNG
import NestedSampling.Utils

-- A sampler
data Sampler = Sampler
               {
                   numParticles      :: {-# UNPACK #-} !Int,
                   mcmcSteps         :: {-# UNPACK #-} !Int,
                   theParticles      :: !(V.Vector (U.Vector Double)),
                   theLogLikelihoods :: !(U.Vector Double),
                   iteration         :: {-# UNPACK #-} !Int,
                   logZ              :: {-# UNPACK #-} !Double,
                   information       :: {-# UNPACK #-} !Double
               } deriving Show

-- | Choose a particle to copy, that isn't number k.
chooseCopy :: Int -> Int -> Gen RealWorld -> IO Int
chooseCopy k n = loop where
  loop prng = do
    index <- MWC.uniformR (0, n - 1) prng
    if   index == k
    then loop prng
    else return index

-- Generate a sampler with n particles and m mcmc steps
generateSampler :: Int -> Int -> Gen RealWorld -> IO Sampler
generateSampler n m gen
    | n <= 1    = undefined
    | m <= 0    = undefined
    | otherwise = do
        putStr ("Generating " ++ (show n) ++ " particles\
                        \ from the prior...")
        hFlush stdout
        theParticles <- V.replicateM n (fromPrior gen)
        let lls = V.map logLikelihood theParticles
        putStrLn "done."
        return Sampler {numParticles=n, mcmcSteps=m,
                        theParticles=theParticles,
                        theLogLikelihoods=U.convert lls, iteration=1,
                        logZ = -1E300, information=0.0}

-- Find the index and the log likelihood value of the worst particle
--
-- NB (jtobin):
--   It may be better to use something like a heap so that we can access the
--   min element in O(1).
findWorstParticle :: Sampler -> (Int, Double)
findWorstParticle sampler = (idx, U.unsafeIndex lls idx)
    where
        idx = U.minIndex lls
        lls = theLogLikelihoods sampler

-- Function to do a single metropolis update
-- Input: A vector of parameters and a loglikelihood threshold
-- Output: An IO action which would return a new vector of parameters
-- and its log likelihood
metropolisUpdate :: Double -> ((U.Vector Double), Double) -> Gen RealWorld
                                -> IO ((U.Vector Double), Double)
metropolisUpdate threshold (x, logL) gen = do
    (proposal, logH) <- perturb x gen
    let a = exp logH
    uu <- rand gen
    let llProposal = logLikelihood proposal
    let accept = (uu < a) && (llProposal > threshold)
    return $
      if   accept
      then (proposal, llProposal)
      else (x, logL)

-- Function to do many metropolis updates
metropolisUpdates :: Int -> Double -> ((U.Vector Double), Double) -> Gen RealWorld
                                -> IO ((U.Vector Double), Double)
metropolisUpdates = loop where
  loop n threshold (x, logL) gen
    | n < 0     = undefined
    | n == 0    = return (x, logL)
    | otherwise = do
                    next  <- metropolisUpdate threshold (x, logL) gen
                    loop (n-1) threshold next gen
{-# INLINE metropolisUpdates #-}

-- Do many NestedSampling iterations
nestedSamplingIterations :: Int -> Sampler -> Gen RealWorld -> IO Sampler
nestedSamplingIterations = loop where
  loop n sampler gen
    | n < 0     = undefined
    | n == 0    = return sampler
    | otherwise = do
                    next <- nestedSamplingIteration sampler gen
                    loop (n-1) next gen
{-# INLINE nestedSamplingIterations #-}

-- Save a particle to disk
writeToFile :: Bool -> (Double, Double) -> U.Vector Double -> IO ()
writeToFile append (logw, logl) particle = do
    -- Open the file
    sampleInfo <- openFile "sample_info.txt"
                        (if append then AppendMode else WriteMode)

    hPutStrLn sampleInfo $ (show logw) ++ " " ++ (show logl)
    hClose sampleInfo

    -- Open the file
    sample <- openFile "sample.txt"
                        (if append then AppendMode else WriteMode)

    let particle' = U.toList particle
    hPutStrLn sample $ foldl' (++) [] $ [show x ++ " " | x <- particle']        
    hClose sample

-- Do a single NestedSampling iteration
nestedSamplingIteration :: Sampler -> Gen RealWorld -> IO Sampler
nestedSamplingIteration sampler gen = do
    let it = iteration sampler
    let logX = negate $ (fromIntegral it) / (fromIntegral (numParticles sampler)) :: Double
    let worst = findWorstParticle sampler

    let iWorst = fst worst
    let n = numParticles sampler
    copy <- chooseCopy iWorst n gen

    -- Approximate log prior weight of the worst particle
    let logPriorWeight = -(k/n) + log (exp (1.0/n) - 1.0) where
          k = fromIntegral it
          n = fromIntegral (numParticles sampler)
    let logPost = logPriorWeight + logLike where
          logLike = snd worst
    let logZ' = logsumexp (logZ sampler) logPost 

    let information' = exp (logPost - logZ') * (snd worst)
                       + exp (logZ sampler - logZ') *
                          (information sampler + logZ sampler) - logZ'

    --H = exp(Obj[worst].logWt - logZnew) * Obj[worst].logL
    -- + exp(logZ - logZnew) * (H + logZ) - logZnew;
    --logZ = logZnew;

    -- Print some stuff from time to time
    let display = it `mod` (numParticles sampler) == 0 :: Bool
    if display then do
        putStr   $ "Iteration " ++ (show it) ++ ". "
        putStr   $ "ln(X) = " ++ (show logX) ++ ". "
        putStrLn $ "ln(L) = " ++ (show $ snd worst) ++ "."
        putStr   $ "ln(Z) = " ++ (show logZ') ++ ", "
        putStrLn $ "H = " ++ (show information') ++ " nats.\n"

    else return ()
--        putStr $ "log(L) = " ++ (show $ snd worst) ++ ". " }
--        when (it == 1)


    -- Write to file
    writeToFile (it /= 1) (logPriorWeight, snd worst)
                        $ (theParticles sampler) V.! iWorst


    -- Copy a surviving particle
    let particle = (V.unsafeIndex (theParticles sampler) copy,
                        U.unsafeIndex (theLogLikelihoods sampler) copy)

    -- Do Metropolis
    let update = metropolisUpdates (mcmcSteps sampler) (snd worst)
{'fraction_non_alphanumeric': 0.10456758636679805, 'fraction_numerical': 0.0038642862663266094, 'mean_word_length': 2.1008866522885215, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 49, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '335346', 'n_tokens_mistral': 3999, 'n_tokens_neox': 3813, 'n_words': 1306}