<commit_before>{-# LANGUAGE RankNTypes, BangPatterns, ScopedTypeVariables, FlexibleInstances #-}

module Math.Probably.HamMC where

import Math.Probably.MCMC
import qualified Math.Probably.PDF as PDF
import Math.Probably.RandIO
import Math.Probably.FoldingStats
import Math.Probably.Sampler
import Control.Applicative

import Numeric.LinearAlgebra
import Numeric.AD
import Text.Printf
import System.IO
import Data.Maybe

import Statistics.Test.KolmogorovSmirnov
import Statistics.Test.MannWhitneyU
import qualified Data.Vector.Unboxed as U
import qualified Control.Monad.State.Strict as S

import Debug.Trace
import Data.IORef
import Control.Spoon


data HMCPar = HMCPar { hpXi :: !(Vector Double),
                       hpPi :: !Double,
                       hpL :: !Int,
                       hpEpsilon :: !Double,
                       hpCount :: !Int,
                       hpAccept :: !Int }

hmc1 :: (Vector Double -> (Double,Vector Double)) 
         -> HMCPar
         -> Sampler HMCPar
hmc1 postGrad (HMCPar current_q current_U l eps count accept) = do
   let dims = dim current_q
       grad_u = snd . postGrad
       u = fst . postGrad
   current_p <- fmap fromList $ gaussManyUnitD dims
   
   let step :: Int -> Vector Double -> Vector Double -> (Vector Double, Vector Double)
       step n p q  -- note that my l is Radford Neal's L+1
        | n == 0 = let qfinal = q + scale eps p in -- only q not p update in last loop
                   (negate $ p - scale (eps/2) (grad_u qfinal), -- after Neal's loop
                    qfinal)
        | otherwise =
            let q1 = q + scale eps p
                p1 = p - scale eps (grad_u q)
            in step (n-1) p1 q1

   let (propose_p, propose_q) = step l (current_p - scale (eps/2) (grad_u current_q)) current_q
       current_K =  (current_p `dot` current_p) / 2
       propose_U = u propose_q
       propose_K =  (propose_p `dot` propose_p) / 2
       ratio = exp $ current_U - propose_U + current_K - propose_K
   u <- unitSample -- 0 to 1
   return $ if u < ratio
               then HMCPar propose_q propose_U l eps (count+1) (accept+1)
               else HMCPar current_q current_U l eps (count+1) (accept)

runHMC postgrad nsam hmc0 = go nsam hmc0 [] where
  go 0 hmcp xs = return (hmcp, xs)
  go n hmcp xs = do hmc1 <- sample $ hmc1 postgrad hmcp<commit_msg>hmc: fix bugs<commit_after>{-# LANGUAGE RankNTypes, BangPatterns, ScopedTypeVariables, FlexibleInstances #-}

module Math.Probably.HamMC where

import Math.Probably.MCMC
import qualified Math.Probably.PDF as PDF
import Math.Probably.RandIO
import Math.Probably.FoldingStats
import Math.Probably.Sampler
import Control.Applicative

import Numeric.LinearAlgebra
import Numeric.AD
import Text.Printf
import System.IO
import Data.Maybe

import Statistics.Test.KolmogorovSmirnov
import Statistics.Test.MannWhitneyU
import qualified Data.Vector.Unboxed as U
import qualified Control.Monad.State.Strict as S

import Debug.Trace
import Data.IORef
import Control.Spoon


data HMCPar = HMCPar { hpXi :: !(Vector Double),
                       hpPi :: !Double,
                       hpL :: !Int,
                       hpEpsilon :: !Double,
                       hpCount :: !Int,
                       hpAccept :: !Int }
    deriving Show

hmc1 :: (Vector Double -> (Double,Vector Double)) 
         -> HMCPar
         -> Sampler HMCPar
hmc1 postGrad (HMCPar current_q current_U l eps count accept) = do
   let dims = dim current_q
       grad_u = negate . snd . postGrad
       u = negate . fst . postGrad
   current_p <- fmap fromList $ gaussManyUnitD dims
   
   let step :: Int -> Vector Double -> Vector Double -> (Vector Double, Vector Double)
       step n p q  -- note that my l is Radford Neal's L+1
        | n == 0 = let qfinal = q + scale eps p in -- only q not p update in last loop
                   (negate $ p - scale (eps/2) (grad_u qfinal), -- after Neal's loop
                    qfinal)
        | otherwise =
            let q1 = q + scale eps p
                p1 = p - scale eps (grad_u q)
            in step (n-1) p1 q1

   let (propose_p, propose_q) = step l (current_p - scale (eps/2) (grad_u current_q)) current_q
       current_K =  (current_p `dot` current_p) / 2
       propose_U = u propose_q
       propose_K =  (propose_p `dot` propose_p) / 2
       ratio = exp $ current_U - propose_U + current_K - propose_K
   u <- unitSample -- 0 to 1
   return $ if u < ratio
               then HMCPar propose_q propose_U l eps (count+1) (accept+1)
               else HMCPar current_q current_U l eps (count+1) (accept)

runHMC postgrad nsam hmc0 = go nsam hmc0 [] where
  go 0 hmcp xs = do io $ putStrLn $ "done with "++show hmcp
                    return (hmcp, xs)
  go n hmcp xs = do hmc1 <- sample $ hmc1 postgrad hmcp
{'fraction_non_alphanumeric': 0.08505747126436781, 'fraction_numerical': 0.00961337513061651, 'mean_word_length': 2.2163978494623655, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 15, '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': '7210723', 'n_tokens_mistral': 1604, 'n_tokens_neox': 1516, 'n_words': 514}