refactor the API for random generation again. Now PGF contains probabilities in the abstract syntax

1 files changed, 29 insertions, 9 deletions

diff --git a/src/runtime/haskell/PGF/Probabilistic.hs b/src/runtime/haskell/PGF/Probabilistic.hs
index a256983c9..873f17be4 100644
--- a/src/runtime/haskell/PGF/Probabilistic.hs
+++ b/src/runtime/haskell/PGF/Probabilistic.hs
@@ -2,6 +2,8 @@ module PGF.Probabilistic
          ( Probabilities(..)
          , mkProbabilities                 -- :: PGF -> M.Map CId Double -> Probabilities
          , defaultProbabilities            -- :: PGF -> Probabilities
+         , getProbabilities
+         , setProbabilities
          , showProbabilities               -- :: Probabilities -> String
          , readProbabilitiesFromFile       -- :: FilePath -> PGF -> IO Probabilities
 
@@ -15,7 +17,7 @@ import PGF.Macros
 
 import qualified Data.Map as Map
 import Data.List (sortBy,partition)
-import Data.Maybe (fromMaybe)
+import Data.Maybe (fromMaybe, fromJust)
 
 -- | An abstract data structure which represents
 -- the probabilities for the different functions in a grammar.
@@ -51,7 +53,7 @@ mkProbabilities pgf probs =
       cats1 = Map.map (\(_,fs) -> fill fs) (cats (abstract pgf))
   in Probs funs1 cats1
   where
-    fill fs = pad [(Map.lookup f probs,f) | f <- fs]
+    fill fs = pad [(Map.lookup f probs,f) | (_,f) <- fs]
       where
         pad :: [(Maybe Double,a)] -> [(Double,a)]
         pad pfs = [(fromMaybe deflt mb_p,f) | (mb_p,f) <- pfs]
@@ -64,16 +66,34 @@ mkProbabilities pgf probs =
 defaultProbabilities :: PGF -> Probabilities
 defaultProbabilities pgf = mkProbabilities pgf Map.empty
 
+getProbabilities :: PGF -> Probabilities
+getProbabilities pgf = Probs {
+  funProbs = Map.map (\(_,_,_,p) -> p) (funs (abstract pgf)),
+  catProbs = Map.map (\(_,fns) -> fns) (cats (abstract pgf))
+  }
+
+setProbabilities :: Probabilities -> PGF -> PGF
+setProbabilities probs pgf = pgf {
+  abstract = (abstract pgf) {
+    funs = mapUnionWith (\(ty,a,df,_) p -> (ty,a,df,p)) (funs (abstract pgf)) (funProbs probs),
+    cats = mapUnionWith (\(hypos,_) fns -> (hypos,fns)) (cats (abstract pgf)) (catProbs probs)
+  }}
+  where
+    mapUnionWith f map1 map2 = 
+      Map.mapWithKey (\k v -> f v (fromJust (Map.lookup k map2))) map1
+
 -- | compute the probability of a given tree
-probTree :: Probabilities -> Expr -> Double
-probTree probs t = case t of
-  EApp f e -> probTree probs f * probTree probs e
-  EFun f   -> maybe 1 id $ Map.lookup f (funProbs probs)
+probTree :: PGF -> Expr -> Double
+probTree pgf t = case t of
+  EApp f e -> probTree pgf f * probTree pgf e
+  EFun f   -> case Map.lookup f (funs (abstract pgf)) of
+                Just (_,_,_,p) -> p
+                Nothing        -> 1
   _ -> 1
 
 -- | rank from highest to lowest probability
-rankTreesByProbs :: Probabilities -> [Expr] -> [(Expr,Double)]
-rankTreesByProbs probs ts = sortBy (\ (_,p) (_,q) -> compare q p) 
-  [(t, probTree probs t) | t <- ts]
+rankTreesByProbs :: PGF -> [Expr] -> [(Expr,Double)]
+rankTreesByProbs pgf ts = sortBy (\ (_,p) (_,q) -> compare q p) 
+  [(t, probTree pgf t) | t <- ts]