Inference over a 2D physics model

{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell       #-}


module Physics where
import Apecs.Physics
import Apecs.Physics.Gloss 

import Graphics.Gloss.Export (exportPicturesToGif,GifLooping(LoopingForever))

import Control.Monad (replicateM, when) 

import LazyPPL
import LazyPPL.Distributions (uniformbounded, normalPdf)

import Numeric.Log(Log(Exp),ln)
import System.Random (setStdGen, mkStdGen)
import Graphics.Matplotlib hiding (density)

import System.IO.Unsafe (unsafePerformIO)
import Data.List (findIndex)
import Data.Monoid (Product(Product))

exportGif filename pics = do
   putStrLn $ "Plotting GIF to " ++ (show filename) ++ "..."
   exportPicturesToGif 2 LoopingForever (288,288) (makeColor 1 1 1 0) filename (\t -> pics !! (floor t)) [0..(fromIntegral $ length pics - 1)]
   putStrLn $ "Done."

plotCoords :: String -> [[(Double,Double)]] -> Double -> Double -> Double -> IO ()
plotCoords filename  xyss ymin ymax alpha = 
    do  putStrLn $ "Plotting " ++ filename ++ "..."
        file filename $ mp # "plot.axis('off')" % mp # "colors = plot.cm.gnuplot2(np.linspace(0,1," # (floor ((fromIntegral $ length xyss) * 1.5) :: Integer) # ")) " % foldl (\a i -> a % plot (map fst (xyss !! i)) (map snd (xyss !! i)) @@ [o1 "o-", o2 "color" (lit ("colors[" ++ (show (length xyss - i)) ++ "]")), o2 "linewidth" (1.0 :: Double), o2 "alpha" alpha, o2 "ms" (0 :: Int)]) mp [0..(length xyss -1)] % (plot ([4,4,5,5] :: [Double]) ([-4,-5,-5,-4] :: [Double]) @@ [o1 "o-", o2 "c" "black", o2 "ms" (0 :: Int)] % xlim (-5 :: Int) (5.1 :: Double) % ylim ymin ymax) 
        putStrLn "Done."
        return ()

getTrajectory :: [(Double,Double,Double)] -> IO [(Double,Double)]
getTrajectory bumpers =
  do w <- initWorld
     runWith w $ do
       b <- initialize bumpers
       xys <- run b 500 
       return xys
  where
    -- run produces the trajectory until the endCriterion
    -- or until the fuel runs out
    -- (in case the ball somehow gets stuck)
    run b fuel = do
      stepPhysics (1/60)
      (Position (V2 x y)) <- get b
      if fuel < 0 then return [(x,y)] else
        if endCriterion (x,y) then return [(x,y)] else
        do
           rest <- run b (fuel - 1)
           return $ (x,y) : rest

getPics :: [(Double,Double,Double)] -> IO [Picture]
getPics bumpers =
  do w <- initWorld
     runWith w $ do
       b <- initialize bumpers
       xypics <- mapM (\_ -> do
                       stepPhysics (1/60)
                       (Position (V2 x y)) <- get b
                       pic <- draw 
                       return (x,y,pic)
                   ) [1..1000]
       let (Just n) = findIndex (\(x,y,_) -> endCriterion (x,y)) xypics
       return $ map (\(_,_,pic) -> pic) $ take (n + 50) xypics
  where
    draw = do
      -- draw the current scene
      pic <- foldDrawM drawBody
      let cam = Camera (V2 (8) (-8)) 9
      return $ cameraTransform cam pic 

runModelMH n trajfile anifile = do
  -- Use MH with p=0.15 (0.15 ~= 1/6, and approx six samples needed)
  samples <- mh 0.15 model
  -- The ball enters the cup when the likelihood is > 0.0876
  let p = Product $ Exp $ ln $ 0.0876415
  let (Just i) = findIndex (\(_,w) -> w > p) samples
  putStrLn $ "[MH] Ball entered the cup after " ++ show i ++ " samples."
  -- Output illustrations to files:
  when (trajfile /= "") $ do
     -- Extract the trajectories of the first n samples
     xyss <- mapM getTrajectory (map fst $ take n samples)
     -- Plot the trajectories together
     plotCoords trajfile xyss (-5) 5 0.01
  when (anifile /= "") $ do
    -- Plot the last trajectory
    pics <- getPics $ fst $ samples !! n
    exportGif anifile pics

runModelBruteForce n trajfile = do
  samples <- weightedsamples model
  -- The ball enters the cup when the likelihood is > 0.0876
  let p = Exp $ ln $ 0.0876415
  let (Just i) = findIndex (\(_,w) -> w > p) samples
  putStrLn $ "[BF] Ball entered the cup after " ++ show i ++ " samples."
  -- Extract the trajectories of the first n samples
  xyss <- mapM getTrajectory (map fst $ take n samples)
  -- Plot the trajectories together
  plotCoords trajfile xyss (-5) 5 0.1

main :: IO ()
main = do
     -- The distribution is multimodal,
     -- so let's restart four times for four examples.
     -- (Statistically better to use mhirreducible.)
     -- We plot the trajectories of the last run.
     setStdGen (mkStdGen 1)
     runModelMH 1000 "" "images/physics1.gif"
     setStdGen (mkStdGen 42)
     runModelMH 1000 "" "images/physics2.gif"
     setStdGen (mkStdGen 47)
     runModelMH 1000 "" "images/physics3.gif"
     setStdGen (mkStdGen 57)
     runModelMH 1000 "images/physics4.svg" "images/physics4.gif"
     -- We also plot what brute force looks like with 4000 trials
     runModelBruteForce 4000 "images/physics5.svg"