<commit_before>module NestedSampling.SpikeSlab where

import Control.Monad.Primitive (RealWorld)
import Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import NestedSampling.Utils
import System.Random.MWC (Gen, uniform, uniformR)
import System.Random.MWC.Distributions (standard)

-- The SpikeSlab model --

-- Log likelihood function
logLikelihood :: U.Vector Double -> Double
logLikelihood params = logsumexp (logl1 + log 100.0) logl2
    where logl1 = (fromIntegral n)*(c - log u)
                        - 0.5*(U.sum $ U.map (\x -> ((x - shift)/u)**2) params)
          logl2 = (fromIntegral n)*(c - log v)
                        - 0.5*(U.sum $ U.map (\x -> (x/v)**2) params)
          c = -0.5*log(2*pi)    :: Double
          u = 0.01              :: Double
          v = 0.1               :: Double
          shift = 0.031         :: Double
          n = U.length params   :: Int

-- fromPrior is an IO action that returns a vector of doubles
-- representing a point in the parameter space
fromPrior :: Gen RealWorld -> IO (U.Vector Double)
fromPrior gen = do
    x <- U.replicateM 20 (uniform gen)
    return $ U.map (\a -> a - 0.5) x

-- Perturb takes a list of doubles as input
-- and returns an IO action that returns the
-- perturbed particle and the logH value.
perturb :: U.Vector Double -> Gen RealWorld -> IO (U.Vector Double, Double)
perturb params gen = do
    -- Choose a parameter to perturb
    k <- uniformR (0, U.length params - 1) gen

    -- Draw from randh
    rh <- randh gen

    -- NB (jtobin):
    --   note that we can't use unsafeThaw here as the params vector could
    --   still be used elsewhere (i.e. in the non-accepting branch of
    --   metropolisUpdate).
    params' <- do
      mvec <- U.thaw params
      UM.unsafeModify mvec (`perturbSingle` rh) k
      U.unsafeFreeze mvec

    return (params', 0.0)

-- Perturb a single double using the provided randh and
-- wrap back into [-0.5, 0.5]
perturbSingle :: Double -> Double -> Double
perturbSingle x rh = (`wrap` (-0.5, 0.5)) $ x + rh

-- My favourite heavy tailed distribution
randh :: Gen RealWorld -> IO Double
randh gen = do
            a <- standard gen
            b <- uniform gen
            n <- standard gen
            return $! transform a b n

-- Function that transforms (a, b, n) -> x
-- for randh
transform :: Double -> Double -> Double -> Double
transform a b n =
  let t = a/sqrt (- (log b))
  in  10.0**(1.5 - 3.0*(abs t))*n
<commit_msg>Make 'transform' local to randh.<commit_after>module NestedSampling.SpikeSlab where

import Control.Monad.Primitive (RealWorld)
import Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import NestedSampling.Utils
import System.Random.MWC (Gen, uniform, uniformR)
import System.Random.MWC.Distributions (standard)

-- The SpikeSlab model --

-- Log likelihood function
logLikelihood :: U.Vector Double -> Double
logLikelihood params = logsumexp (logl1 + log 100.0) logl2
    where logl1 = (fromIntegral n)*(c - log u)
                        - 0.5*(U.sum $ U.map (\x -> ((x - shift)/u)**2) params)
          logl2 = (fromIntegral n)*(c - log v)
                        - 0.5*(U.sum $ U.map (\x -> (x/v)**2) params)
          c = -0.5*log(2*pi)    :: Double
          u = 0.01              :: Double
          v = 0.1               :: Double
          shift = 0.031         :: Double
          n = U.length params   :: Int

-- fromPrior is an IO action that returns a vector of doubles
-- representing a point in the parameter space
fromPrior :: Gen RealWorld -> IO (U.Vector Double)
fromPrior gen = do
    x <- U.replicateM 20 (uniform gen)
    return $ U.map (\a -> a - 0.5) x

-- Perturb takes a list of doubles as input
-- and returns an IO action that returns the
-- perturbed particle and the logH value.
perturb :: U.Vector Double -> Gen RealWorld -> IO (U.Vector Double, Double)
perturb params gen = do
    -- Choose a parameter to perturb
    k  <- uniformR (0, U.length params - 1) gen
    rh <- randh gen

    -- NB (jtobin):
    --   note that we can't use unsafeThaw here as the params vector could
    --   still be used elsewhere (i.e. in the non-accepting branch of
    --   metropolisUpdate).
    params' <- do
      mvec <- U.thaw params
      UM.unsafeModify mvec (`perturbSingle` rh) k
      U.unsafeFreeze mvec

    return (params', 0.0)

-- Perturb a single double using the provided randh and
-- wrap back into [-0.5, 0.5]
perturbSingle :: Double -> Double -> Double
perturbSingle x rh = (`wrap` (-0.5, 0.5)) $ x + rh

-- My favourite heavy tailed distribution
randh :: Gen RealWorld -> IO Double
randh gen = do
    a <- standard gen
    b <- uniform gen
    n <- standard gen
    return $! transform a b n
  where
    transform a b n =
      let t = a/sqrt (- (log b))
      in  10.0**(1.5 - 3.0*(abs t))*n

{'fraction_non_alphanumeric': 0.1181460790399338, 'fraction_numerical': 0.020277260500724186, 'mean_word_length': 2.165684348395547, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 27, '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': '4988043', 'n_tokens_mistral': 1612, 'n_tokens_neox': 1495, 'n_words': 570}