module Bayes.VariableElimination.Buckets(
Buckets(..)
, EliminationOrder(..)
, IsBucketItem(..)
, createBuckets
, getBucket
, updateBucket
, addBucket
, removeFromBucket
, marginalizeOneVariable
) where
import Bayes.PrivateTypes
import qualified Data.Map as M
import Data.List(partition,minimumBy,(\\),find,foldl')
import Data.Maybe(fromJust)
type EliminationOrder dv = [dv]
data Buckets f = Buckets !(EliminationOrder DV) !(M.Map DV [f])
instance Show f => Show (Buckets f) where
show (Buckets v m) = "BUCKET\n" ++ show v ++ "\n" ++ concatMap disp (M.toList m)
where
disp (v,f) = "Bucket for " ++ show v ++ "\n" ++ concatMap dispElem f ++ "\n----\n"
dispElem f = show f ++ "\n"
class IsBucketItem f where
scalarItem :: f -> Bool
itemProduct :: [f] -> f
itemProjectOut :: DV -> f -> f
itemContainsVariable :: f -> DV -> Bool
addDVToBucket :: IsBucketItem f => ([f],M.Map DV [f]) -> DV -> ([f],M.Map DV [f])
addDVToBucket (rf, m) dv =
let (fk,remaining) = partition (flip itemContainsVariable dv) rf
in
(remaining, M.insert dv fk m)
createBuckets :: (IsBucketItem f)
=> [f]
-> EliminationOrder DV
-> EliminationOrder DV
-> Buckets f
createBuckets s e r =
let
theOrder = e ++ r
(_,b) = foldl' addDVToBucket (s,M.empty) theOrder
in
Buckets theOrder b
getBucket :: DV
-> Buckets f
-> [f]
getBucket dv (Buckets _ m) = fromJust $ M.lookup dv m
updateBucket :: IsBucketItem f
=> DV
-> f
-> Buckets f
-> Buckets f
updateBucket dv f b@(Buckets e m) =
if scalarItem f
then
Buckets (remainingVarsToProcess e) (M.insert dv [f] m)
else
let b' = removeFromBucket dv b
in
addBucket b' f
where
remainingVarsToProcess [] = []
remainingVarsToProcess l = tail l
addBucket :: IsBucketItem f => Buckets f -> f -> Buckets f
addBucket (Buckets e b) f =
let inBucket = find (f `itemContainsVariable`) e
in
case inBucket of
Nothing -> Buckets e b
Just bucket -> Buckets e (M.insertWith' (++) bucket [f] b)
removeFromBucket :: DV -> Buckets f -> Buckets f
removeFromBucket dv (Buckets [] m) = Buckets [] (M.delete dv m)
removeFromBucket dv (Buckets e m) = Buckets (tail e) (M.delete dv m)
marginalizeOneVariable :: IsBucketItem f => Buckets f -> DV -> Buckets f
marginalizeOneVariable currentBucket dv =
let fk = getBucket dv currentBucket
p = itemProduct fk
f' = itemProjectOut dv p
in
updateBucket dv f' currentBucket