Symbolic.hs

-- | Exact credal inference via a symbolic (multilinear-polynomial) semiring.
--
--   Bounds are extracted by optimizing the objective over the parameter box corners.
--
--   Throws an error if the evidence is unsatisfiable, i.e. the credal set is empty.
module Imp.Inference.Symbolic
  ( marginalSymbolic
  , intervalProbabilitySymbolic
  , intervalExpectationSymbolic
  ) where

import Data.Bits (bit, testBit, clearBit, (.|.))
import Data.Bool (bool)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

import Imp.BDD.Compile (compile)
import Imp.BDD.WMC (Weight(..), wmcBatch)
import Imp.DSL (Imp)
import Imp.Semiring

-- | A constant polynomial.
polyConst :: Double -> PolyS
polyConst p = PolyS (if p == 0 then Map.empty else Map.singleton 0 p)

-- | The polynomial reading of the variable weights.
polyWeights :: Weight -> (PolyS, PolyS)
polyWeights w = case w of
  Prob p   -> (polyConst (1 - p), polyConst p)
  -- Two indicators per Knightian variable, one for @0@ and one for @1@.
  Knight i -> ( PolyS (Map.singleton (bit (2 * i))     1)
              , PolyS (Map.singleton (bit (2 * i + 1)) 1) )

-- | Each value's count polynomial together with the total-mass polynomial.
compilePolys :: Ord a => Imp g a -> (Map a PolyS, PolyS)
compilePolys prog =
  let (mgr, weights, _, worlds) = compile prog
      counts = wmcBatch polyWeights mgr weights worlds
  in (counts, sumS (Map.elems counts))

-- | Exact per-value marginal bounds, computed symbolically.
marginalSymbolic :: Ord a => Imp g a -> Map a (Double, Double)
marginalSymbolic prog =
  let (counts, total) = compilePolys prog
  in Map.map (optimizeRatio total) counts

-- | Exact lower and upper probability of an event, computed symbolically.
intervalProbabilitySymbolic :: Ord a => Imp g a -> (a -> Bool) -> (Double, Double)
intervalProbabilitySymbolic prog predicate =
  intervalExpectationSymbolic prog (bool 0 1 . predicate)

-- | Exact lower and upper expectation of a real-valued function, computed symbolically.
intervalExpectationSymbolic :: Ord a => Imp g a -> (a -> Double) -> (Double, Double)
intervalExpectationSymbolic prog score =
  let (counts, total) = compilePolys prog
      aggr = sumS [polyConst (score v) .*. count | (v, count) <- Map.toList counts]
  in optimizeRatio total aggr

-- | Bounds over the feasible corners of the parameter box.
optimizeRatio :: PolyS -> PolyS -> (Double, Double)
optimizeRatio den num =
  let mask = foldl' (.|.) 0 (Map.keys (unPoly num) ++ Map.keys (unPoly den))
      free = [ i | i <- takeWhile (\i -> bit (2 * i) <= mask) [0 ..]
                 , testBit mask (2 * i) || testBit mask (2 * i + 1) ]
  in case cornerRatios free den num of
       []     -> error "No feasible probabilities"
       ratios -> (minimum ratios, maximum ratios)

-- | Ratio at every feasible corner reachable by fixing the Knightian variables to @0@ and @1@.
cornerRatios :: [Int] -> PolyS -> PolyS -> [Double]
cornerRatios [] den num =
  let constDen = constTerm den in [ constTerm num / constDen | constDen > 0 ]
cornerRatios (i : rest) den num =
     cornerRatios rest (fixKnight i False den) (fixKnight i False num)
  ++ cornerRatios rest (fixKnight i True  den) (fixKnight i True  num)

-- | Fix Knightian parameter @i@ to @0@ or @1@.
fixKnight :: Int -> Bool -> PolyS -> PolyS
fixKnight i b (PolyS m) =
  PolyS $ Map.filter (/= 0) $ Map.fromListWith (+)
    [ (clearBit mask thisBit, c) | (mask, c) <- Map.toList m, not (testBit mask thatBit) ]
  where
    thisBit = if b then 2 * i + 1 else 2 * i
    thatBit = if b then 2 * i     else 2 * i + 1

-- | The constant term (value once every parameter is fixed).
constTerm :: PolyS -> Double
constTerm (PolyS m) = Map.findWithDefault 0 0 m