Working with probabilities in Haskell

probability.hs

{-# LANGUAGE TypeFamilies,
             GeneralizedNewtypeDeriving #-}

-- Example based on http://www.randomhacks.net/2007/02/21/randomly-sampled-distributions/

module NewEx where

import Control.Monad
import Control.Monad.Trans
import Data.List

data Human = Man | Woman | SuperHuman deriving (Show, Eq, Ord)
type Friends = [Human]

newtype PerhapsT m a = PerhapsT { runPerhapsT :: m (Perhaps a) }

type Dist = PerhapsT []

newtype Prob = P Float deriving (Eq, Ord, Show, Num)

data Perhaps a = Perhaps a Prob

never :: Perhaps a
never = Perhaps undefined 0

neverHappens :: Perhaps a -> Bool
neverHappens never = True
neverHappens _ = False

instance Functor Perhaps where
  fmap f (Perhaps x p) = Perhaps (f x) p

instance Applicative Perhaps where
  pure x = Perhaps x 1
  (Perhaps f _) <*> (Perhaps x p) = Perhaps (f x) p

instance Monad Perhaps where
  return x = Perhaps x 1
  ph >>= f  | neverHappens ph  = never
            | otherwise        = Perhaps x (p1 * p2)
    where (Perhaps (Perhaps x p1) p2) = fmap f ph

instance MonadTrans PerhapsT where
  lift x = PerhapsT (liftM return x)

instance Monad m => Functor (PerhapsT m) where
  fmap = liftM

instance (Monad m) => Applicative (PerhapsT m) where
  pure = return
  (<*>) = ap

instance Monad m => Monad (PerhapsT m) where
  return = lift . return
  m >>= f = PerhapsT bound
    where bound = do
            ph <- runPerhapsT m
            case ph of
              (Perhaps x1 p1)  | p1 == 0    -> return never
                               | otherwise  -> do
                (Perhaps x2 p2) <- runPerhapsT (f x1)
                return (Perhaps x2 (p1 * p2))

uniform :: [a] -> Dist a
uniform = weighted . map (\x -> (x, 1))

weighted :: [(a, Float)] -> Dist a
weighted [] =
  error "Empty probability distributuion"
weighted xws = PerhapsT (map weight xws)
  where weight (x,w) = Perhaps x (P (w / sum))
        sum = foldl (\w1 (_,w2) -> w1+w2) 0 xws

enum :: [a] -> [Float] -> Dist a
enum as ps = weighted $ zip as ps

prollyFriends = enum [Man, Woman, SuperHuman] [0.488, 0.511, 0.001]

generateFriends :: Int -> [Perhaps Friends]
generateFriends n = runPerhapsT $ replicateM n prollyFriends

instance (Show a) => Show (Perhaps a) where
  show (Perhaps h (P p)) = show h ++ " " ++ (show prob) ++ "%"
    where prob = (myround 3 p * 100)

myround n f = (fromInteger $ round $ f * (10^n)) / (10.0^^n)

numSuperHeroes n = merge . group' . sort' $ map (fmap $ elem SuperHuman) $ generateFriends n

groupFn (Perhaps a _) (Perhaps b _) = a == b

merge = map (foldl mergeProb (Perhaps True 0))

group' = groupBy groupFn

sort' = sortBy sortFn

sortFn :: (Ord a) => Perhaps a -> Perhaps a -> Ordering
sortFn (Perhaps a _) (Perhaps b _)
  | a > b = GT
  | a < b = LT
  | otherwise = EQ

mergeProb :: Perhaps Bool -> Perhaps Bool -> Perhaps Bool
mergeProb (Perhaps a (P p)) (Perhaps b (P p')) = Perhaps (a && b) (P $ p + p')