commit 620d56e17b410aa20251e2be592a938686434262
parent c33b02c0089e2b2a2c2bfddf249ec1d386c69b17
Author: Jared Tobin <jared@jtobin.ca>
Date: Tue, 22 Oct 2013 10:37:28 +1300
Add Chinese Restaurant Process example.
Diffstat:
3 files changed, 117 insertions(+), 1 deletion(-)
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1 @@
+*swp
diff --git a/src/Examples.hs b/src/Examples.hs
@@ -1,11 +1,18 @@
+{-# LANGUAGE TemplateHaskell #-}
+
import Control.Applicative
+import Control.Arrow
import Control.Error
+import Control.Lens hiding (to)
import Control.Monad
import Control.Monad.Primitive
import Control.Monad.Trans
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
+import Data.IntMap.Strict (IntMap)
+import qualified Data.IntMap.Strict as IntMap
import Data.Vector (singleton)
+import qualified Data.Traversable as Traversable
import Measurable.Generic
import Numeric.SpecFunctions
import Statistics.Distribution
@@ -167,7 +174,7 @@ gaussianMixtureModel n observed g = do
fromObservations samples
--- | A bizarre measure.
+-- | A bizarre random measure.
weirdMeasure
:: Fractional r
=> [Group]
@@ -226,3 +233,11 @@ main = do
print mixtureModelMean
+ weirdProbabilityOfTrue <- expectation id $
+ containing [True] <$> weirdMeasure [B, B, A, C, A, C, C, A, B, A] []
+
+ print weirdProbabilityOfTrue
+
+
+
+
diff --git a/src/Examples/ChineseRestaurantProcess.hs b/src/Examples/ChineseRestaurantProcess.hs
@@ -0,0 +1,100 @@
+{-# LANGUAGE TemplateHaskell #-}
+
+import Control.Applicative
+import Control.Lens
+import Control.Monad.Trans
+import Data.IntMap.Strict (IntMap)
+import qualified Data.IntMap.Strict as IntMap
+import Measurable.Generic
+
+data Table = Table {
+ _number :: {-# UNPACK #-} !Int
+ , _people :: {-# UNPACK #-} !Int
+ } deriving (Eq, Show)
+
+instance Ord Table where
+ t1 < t2 = _people t1 < _people t2
+
+$(makeLenses ''Table)
+
+-- | Mass function for a given table. It's dependent on the state of the
+-- restaurant via 'n' and 'newestTable'.
+tableMass :: (Fractional a, Integral b) => b -> a -> Table -> Table -> a
+tableMass n a newestTable table
+ | table^.number == newestTable^.number = a / (fromIntegral n + a)
+ | otherwise = fromIntegral (table^.people) / (fromIntegral n + a)
+
+-- | A dependent measure over tables.
+tableMeasure
+ :: (Fractional r, Integral b, Applicative m, Monad m, Traversable t)
+ => b
+ -> r
+ -> Table
+ -> t Table
+ -> MeasureT r m Table
+tableMeasure n a newestTable =
+ fromMassFunction (return . tableMass n a newestTable)
+
+-- | A probability measure over restaurants, represented by IntMaps.
+restaurantMeasure
+ :: (Fractional r, Monad m, Applicative m)
+ => r
+ -> IntMap Table
+ -> MeasureT r m (IntMap Table)
+restaurantMeasure a restaurant = do
+ let numberOfCustomers = sumOf (traverse.people) restaurant
+ numberOfTables = lengthOf traverse restaurant
+ nextTableNum = succ numberOfTables
+ possibleTable = Table nextTableNum 1
+ possibleRestaurant = IntMap.insert nextTableNum possibleTable restaurant
+
+ table <- tableMeasure numberOfCustomers a possibleTable possibleRestaurant
+
+ let newTable | table^.number == possibleTable^.number = table
+ | otherwise = table&people %~ succ
+
+ return $ IntMap.insert (newTable^.number) newTable restaurant
+
+-- | The Chinese Restaurant process.
+--
+-- This implementation is dismally inefficient as-is, but appears to be
+-- correct. I think I need to look at doing memoization under the hood.
+chineseRestaurantProcess
+ :: (Enum a, Eq a, Fractional r, Monad m, Applicative m, Num a)
+ => a
+ -> r
+ -> MeasureT r m (IntMap Table)
+chineseRestaurantProcess n a = go n IntMap.empty
+ where go 0 restaurant = return restaurant
+ go j restaurant = restaurantMeasure a restaurant >>= go (pred j)
+
+main :: IO ()
+main = do
+ -- Easily verified by hand.
+ meanTinyRestaurant <- expectation (fromIntegral . lengthOf traverse)
+ (chineseRestaurantProcess 2 1)
+
+ -- Easily verified by hand.
+ meanBiggerRestaurant <- expectation (fromIntegral . lengthOf traverse)
+ (chineseRestaurantProcess 3 1)
+
+ meanBigRestaurant <- expectation (fromIntegral . lengthOf traverse)
+ (chineseRestaurantProcess 9 1)
+
+ meanBigRestaurantAntisocial <- expectation (fromIntegral . lengthOf traverse)
+ (chineseRestaurantProcess 9 3)
+
+ -- We can answer other questions by changing our transformation function.
+ -- Trivial example: the expected number of customers for a CRP observed for n
+ -- epochs is always n.
+
+ differentQuestionSameMeasure <-
+ expectation (fromIntegral . sumOf (traverse.people))
+ (chineseRestaurantProcess 9 3)
+
+ print meanTinyRestaurant
+ print meanBiggerRestaurant
+ print meanBigRestaurant
+ print meanBigRestaurantAntisocial
+ print differentQuestionSameMeasure
+
