added support for binomial distributions

src/Data/PWhile/Expression.hs

@@ -19,7 +19,7 @@ module Data.PWhile.Expression
     -- distributions
   -- , Frac (..)
   , Dist (..)
-  , discrete, dirac, rand, bernoulli, unif
+  , discrete, dirac, rand, bernoulli, binomial, unif
   , prettyFrac
   )
 where
@@ -127,6 +127,14 @@ dirac e = Discrete [(1,e)]
 bernoulli :: (Exp, Exp) -> e -> e -> Dist e
 bernoulli (a,b) e1 e2 = Discrete [(a, e1), (b, e2)]
 
+binomial :: (Int, Int) -> Int -> Dist Exp
+binomial (pa,pb) n = Discrete [(fromIntegral $ binom n k * pa^k * (pb - pa)^k, fromIntegral k)
+                              | k <- [0..n]]
+  where
+    binom _ 0 = 1
+    binom 0 _ = 0
+    binom m k = binom (m-1) (k-1) * m `div` k
+
 -- pretty
 ----------------------------------------------------------------------
 

src/PWhile/Parser.hs

@@ -59,7 +59,7 @@ keywords :: [String]
 keywords =
   [ "def", "while","skip","if","then","else","do"
   ,"and","or","not","abort","tick","var","prob","assume", "consume"
-  ,"unif","rand","nondet","ber"]
+  ,"unif","rand","nondet","ber","bin"]
 
 keyword :: String -> Parser ()
 keyword w = lexeme $ try $ string w *> notFollowedBy alphaNumChar
@@ -178,7 +178,7 @@ pAssignment :: Parser Program
 pAssignment = do
   v <- pVar
   void $ symbol "="
-  pDirac v <|> pRand v <|> pUnif v <|> pBer v
+  pDirac v <|> pRand v <|> pUnif v <|> pBer v <|> pBin v
 
 
 
@@ -200,17 +200,27 @@ pBer v = do
   (a,b) <- symbol "ber" *> pTuple integer
   return $ (v .~) $ bernoulli (1,1) (E.constant a) (E.constant b)
 
+pBin :: Exp -> Parser Program
+pBin v = do
+  (a,b,c) <- symbol "bin" *> pTriple integer
+  return $ (v .~) $ binomial (b,c) a
+
 pTuple :: Parser a -> Parser (a,a)
 pTuple p = parens $ p >>= \a -> comma >> p >>= \b -> return (a,b)
 
+pTriple :: Parser a -> Parser (a,a,a)
+pTriple p = parens $ p >>= \a -> comma >> p >>= \b -> comma >> p >>= \ c -> return (a,b,c)
+
 pDef :: Parser Program
 pDef =
   snd <$> withBlock' (keyword "def" *> identifier *> symbol "():" *> whitespaceN *> pVars)
 
 pAssume :: Parser Program
 pAssume = do
-  (g,tt) <- withBlock' (keyword "assume" *> pGuard)
-  return $ assume g tt
+  g <- keyword "assume" *> pBExp
+  return (ite' g skip abort)
+  -- (g,tt) <- withBlock' (keyword "assume" *> pGuard)
+  -- return $ assume g tt
 
 pElse :: Parser Program
 pElse = do