In the context of a stochastic ray tracer, I'd like to decouple the MC integration (path tracing, bidirectional path tracing) from sample generation (uniform random, stratified, poisson, metropolis, ...). Most of this is already implemented, but it's tedious to use. So I ditched that and try build something nicer, by splitting sampled computations in two phases: In SampleGen you are allowed to request a random value using the mk1d and mk2d functions, which are then supplied with actual Floats by the sampling algorithm. Those values can be examined in SampleRun to do the actual computation. Here's some code with the interesting bits of a stratified sampler and it's use:
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
import Control.Applicative
import Control.Monad.State.Strict
import Control.Monad.Primitive
import System.Random.MWC as MWC
-- allows to construct sampled computations
newtype SampleGen s m a = SampleGen (StateT s m a)
deriving ( Functor, Applicative, Monad
, MonadState s, MonadTrans )
-- allows to evaluate sampled computations constructed in SampleGen
newtype SampleRun s m a = SampleRun (StateT s m a)
deriving ( Functor, Applicative, Monad
, MonadState s )
-- a sampled computation, parametrized over the generator's state g,
-- the evaluator's state r, the underlying monad m and the result
-- type a
type Sampled g r m a = SampleGen g m (SampleRun r m a)
----------------------
-- Stratified Sampling
----------------------
-- | we just count the number of requested 1D samples
type StratGen = Int
-- | the pre-computed values and a RNG for additional ones
type StratRun m = ([Float], Gen (PrimState m))
-- | specialization of Sampled for stratified sampling
type Stratified m a = Sampled StratGen (StratRun m) m a
-- | gives a sampled value in [0..1), this is kind
-- of the "prime" value, upon which all computations
-- are built
mk1d :: PrimMonad m => Stratified m Float
mk1d = do
n1d <- get
put $ n1d + 1
return $ SampleRun $ do
fs <- gets fst
if length fs > n1d
then return (fs !! n1d)
else gets snd >>= lift . MWC.uniform
-- | gives a pair of stratified values, should really also
-- be a "prime" value, but here we just construct them
-- from two 1D samples for fun
mk2d :: (Functor m, PrimMonad m) => Stratified m (Float, Float)
mk2d = mk1d >>= \f1 -> mk1d >>= \f2 ->
return $ (,) <$> f1 <*> f2
-- | evaluates a stratified computation
runStratified
:: (PrimMonad m)
=> Int -- ^ number of samples
-> Stratified m a -- ^ computation to evaluate
-> m [a] -- ^ the values produced, a list of nsamples values
runStratified nsamples (SampleGen c) = do
(SampleRun x, n1d) <- runStateT c 0
-- let's just pretend I'd use n1d to actually
-- compute stratified samples
gen <- MWC.create
replicateM nsamples $ evalStateT x ([{- samples would go here #-}], gen)
-- estimate Pi by Monte Carlo sampling
-- mcPi :: (Functor m, PrimMonad m) => Sampled g r m Float
mcPi :: (Functor m, PrimMonad m) => Stratified m Float
mcPi = do
v <- mk2d
return $ v >>= \(x, y) -> return $ if x * x + y * y < 1 then 4 else 0
main :: IO ()
main = do
vs <- runStratified 10000 mcPi :: IO [Float]
print $ sum vs / fromIntegral (length vs)
The missing part here is that in it's current form, the mcPi function has the type
mcPi :: (Functor m, PrimMonad m) => Stratified m Float
while it should really be something like
mcPi :: (Functor m, PrimMonad m) => Sampled g r m Float
Admitted, the four type parameters on Sampled aren't exactly beautiful, but at least something like this would be useful. In summary, I'm looking for something allowing to express computations like mcPi independent of the sampling algorithm, e.g.:
- a uniform random sampler does not need to maintain any state in the
SampleGenphase, and needs only a RNG in theSampleRunphase - both, the stratified and the poisson disk sampler (and probably others) keep track of the number of 1D and 2D samples needed and precompute them into a vector, and they would be allowed to share a
SampleGenandSampleRunimplementation, to differ only in what happens inbetweenSampleGenandSampleRun(how the vector is actually filled) - a metropolis sampler would use a lazy sample generation technique in it's
SampleRunphase
I'd like to compile it using GHC, so extensions like MultiParamTypeClasses and TypeFamilies are ok to me, but I did not come up with anything remotely usable.
PS: As motivation, some pretty pictures. And the code in it's current form is on GitHub
