Probability Trees

{-# language DeriveFunctor, DeriveFoldable #-}
{-# language PatternSynonyms, ViewPatterns #-}

module ProbTree where

import Data.Monoid
import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)
import Control.Arrow
import Data.Ratio
import Data.Foldable

Previously, I tried to figure out how to make the probability monad more “listy”. I read a little more about the topic (especially Erwig and Kollmansberger 2006; and Kidd 2007).

I then thought about what a probability monad would look like if it was based on other data structures. I feel like the standard version really wants to be:

newtype ProperProb a = ProperProb
  { yes :: Map a (Product Rational) }

But of course a monad instance isn’t allowed.

Similar to a map, though, is a binary tree:

data BinaryTree a = Leaf
                  | Node (BinaryTree a) a (BinaryTree a)

And it feels better for probability - flatter, somehow. Transmuting it into a probability-thing:

data Odds a = Certain a
            | Choice (Odds a) Rational (Odds a)
            deriving (Eq, Functor, Foldable, Show)

That looks good to me. A choice between two different branches feels more natural than a choice between a head and a tail.

The fold is similar to before, with an unfold for good measure:

foldOdds :: (b -> Rational -> b -> b) -> (a -> b) -> Odds a -> b
foldOdds f b = r where
  r (Certain x) = b x
  r (Choice xs p ys) = f (r xs) p (r ys)
  
unfoldOdds :: (b -> Either a (b,Rational,b)) -> b -> Odds a
unfoldOdds f = r where
  r b = case f b of
    Left a -> Certain a
    Right (x,p,y) -> Choice (r x) p (r y)
  
fi :: Bool -> a -> a -> a
fi True  t _ = t
fi False _ f = f

I changed the pattern synonym a little:

unRatio :: Num a => Rational -> (a,a)
unRatio = numerator   &&& denominator 
      >>> fromInteger *** fromInteger

pattern n :% d <- (unRatio -> (n,d))

Then, the probOf function:

probOf :: Eq a => a -> Odds a -> Rational
probOf e = foldOdds f b where
  b x = fi (e == x) 1 0
  f x (n:%d) y = (x * n + y * d) / (n + d)

This version doesn’t have the option for short-circuiting on the first value it finds.

For generating from lists, you can try to evenly divide the list among each branch.

fromListOdds :: (([b], Int) -> Integer) -> (b -> a) -> [b] -> Maybe (Odds a)
fromListOdds fr e = r where
  r [] = Nothing
  r xs = Just (unfoldOdds f (xs, length xs))
  f ([x],_) = Left (e x)
  f (xs ,n) = Right ((ys,l), fr (ys,l) % fr (zs,r), (zs,r)) where
    l = n `div` 2
    r = n - l
    (ys,zs) = splitAt l xs

equalOdds :: [a] -> Maybe (Odds a)
equalOdds = fromListOdds (fromIntegral . snd) id

fromDistrib :: [(a,Integer)] -> Maybe (Odds a)
fromDistrib = fromListOdds (sum . map snd . fst) fst

What’s really nice about this version is the fact that the old append is just the Choice constructor, leaving the instances to be really nice:

flatten :: Odds (Odds a) -> Odds a
flatten = foldOdds Choice id

instance Applicative Odds where
  pure = Certain
  fs <*> xs = flatten (fmap (<$> xs) fs)
  
instance Monad Odds where
  x >>= f = flatten (f <$> x)

Finally, as a bonus, to remove duplicates:

lcd :: Foldable f => f Rational -> Integer
lcd = foldl' (\a e -> lcm a (denominator e)) 1

toDistrib :: Odds a -> [(a,Integer)]
toDistrib = factorOut . foldOdds f b where
  b x = [(x,1)]
  f l p r = (map.fmap) (n%t*) l ++ (map.fmap) (d%t*) r where
    n = numerator p
    d = denominator p
    t = n + d
  factorOut xs = (map.fmap) (numerator . (lcd'*)) xs where
    lcd' = fromIntegral . lcd . map snd $ xs

counts :: (Ord a, Num n) => [(a,n)] -> [(a,n)]
counts = 
  Map.assocs . 
  Map.fromListWith (+)
      
compress :: Ord a => Odds a -> Odds a
compress xs = let Just ys = (fromDistrib . counts . toDistrib) xs in ys

After reading yet more on this, I found that the main issue with the monad is its performance. Two articles in particular: Larsen (2011), and Ścibior, Ghahramani, and Gordon (2015), refer to a GADT implementation of the monad which maximises laziness.

References