improve floating point random generation and include tree for stdgen

1 files changed, 26 insertions, 7 deletions

diff --git a/src/PPL/Sampling.hs b/src/PPL/Sampling.hs
index 574025a..71a61fd 100644
--- a/src/PPL/Sampling.hs
+++ b/src/PPL/Sampling.hs
@@ -18,11 +18,30 @@ import Data.IORef
 import Data.List (foldl')
 import Data.Vector qualified as V
 import Data.Vector.Hashtables qualified as H
-import Numeric.Log
-import PPL.Internal
+import Numeric.Log hiding (sum)
+import PPL.Internal hiding (newTree, split)
 import Streaming.Prelude (Of, Stream, yield)
-import System.Random (StdGen)
-import System.Random qualified as R
+import System.IO.Unsafe
+import System.Random (StdGen, split)
+
+{-# INLINE newTree #-}
+newTree :: IORef (HashMap Integer Double, StdGen) -> Tree
+newTree s = go 1
+  where
+    go :: Integer -> Tree
+    go id =
+      Tree
+        ( unsafePerformIO $ do
+            (m, g) <- readIORef s
+            H.lookup m id >>= \case
+              Nothing -> do
+                let (x, g') = random g
+                H.insert m id x
+                writeIORef s (m, g')
+                pure x
+              Just x -> pure x
+        )
+        (go (2 * id), go (2 * id + 1))
 
 data MinHeap a k = Node a k (MinHeap a k) (MinHeap a k) | Empty
   deriving (Show)
@@ -41,14 +60,14 @@ toList (Node k v l r) = v : toList (merge l r)
 
 mh :: (MonadIO m) => StdGen -> Double -> Meas a -> Stream (Of (a, Log Double)) m ()
 mh g p m = do
-  let (g0, g1) = R.split g
+  let (g0, g1) = split g
   hm <- liftIO $ H.initialize 0
   omega <- liftIO $ newIORef (hm, g0)
   let (x, w) = head $ samples m $ newTree omega
   step g1 omega x w
   where
     step !g0 !omega !x !w = do
-      let (Exp . log -> r, R.split -> (g1, g2)) = R.random g0
+      let (Exp . log -> r, split -> (g1, g2)) = random g0
       omega' <- mutate g1 omega
       let (!x', !w') = head $ samples m $ newTree omega'
           ratio = w' / w
@@ -64,7 +83,7 @@ mh g p m = do
       (m, g0) <- readIORef omega
       m' <- H.clone m
       ks <- H.keys m
-      let (rs :: [Double], qs :: [Double]) = (R.randoms *** R.randoms) (R.split g)
+      let (rs :: [Double], qs :: [Double]) = (randoms *** randoms) (split g)
           ks' = toList $ foldl' (\m -> uncurry (insert m)) Empty $ zip rs $ V.toList ks
           n = fromIntegral (V.length ks)
       void $ zipWithM_ (\k q -> H.insert m' k q) (take (1 + floor (p * fromIntegral n)) ks') qs