Approx.hs

-- | Approximate credal inference via one-pass interval WMC.  Bounds are
--   sound outer approximations, but are not tight in general.
--
--   Throws an error if the evidence is unsatisfiable, i.e. the credal set is empty.
module Imp.Inference.Approx
  ( marginalApprox
  , intervalProbabilityApprox
  , intervalExpectationApprox
  ) where

import Control.Monad.State.Strict (runState)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

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

-- | Flips get a point interval, Knights the full unit interval.
intervalWeights :: Weight -> (IntervalS, IntervalS)
intervalWeights w = case w of
  Prob p   -> (IntervalS (1 - p) (1 - p), IntervalS p p)
  Knight _ -> (IntervalS 0 1, IntervalS 0 1)

-- | Group the return values by @key@, returning the total-mass interval
--   and one interval per group
compileApprox :: (Ord a, Ord k) => Imp g a -> (a -> k) -> (IntervalS, Map k IntervalS)
compileApprox prog key =
  let (mgr, weights, _, worlds) = compile prog
      blocks = Map.fromListWith (++) [ (key v, [g]) | (v, g) <- Map.toList worlds ]
      (bdds, mgr')      = runState (traverse bddAny blocks) mgr
      (evidence, mgr'') = runState (bddAny (Map.elems bdds)) mgr'
  in ( wmc intervalWeights mgr'' weights evidence
     , wmcBatch intervalWeights mgr'' weights bdds )

-- | Interval division where the quotient stays in @[0, 1]@.
intervalDiv :: IntervalS -> IntervalS -> IntervalS
intervalDiv (IntervalS nLo nHi) (IntervalS dLo dHi) =
  let div' a b = if b > 0 then a / b else 0
  in  IntervalS (div' nLo dHi) (div' nHi (max nHi dLo))

-- | Divide a count box by the total box, tightened by the complementary count.
condition :: IntervalS -> IntervalS -> IntervalS -> (Double, Double)
condition total@(IntervalS _ tHi) count remainder =
  let IntervalS lo hi   = intervalDiv count total
      IntervalS rLo rHi = intervalDiv remainder total
      lo' = max lo (1 - rHi)
      hi' = min hi (1 - rLo)
  in if tHi <= 0 then error "No feasible probabilities"
                 -- clamp to min in case of rounding errors
                 else (min lo' hi', hi')

-- | Approximate per-value marginal bounds via interval WMC.
marginalApprox :: Ord a => Imp g a -> Map a (Double, Double)
marginalApprox prog =
  let (total, counts) = compileApprox prog id
      remainder v = sumS (Map.elems (Map.delete v counts))
  in Map.mapWithKey (\v c -> condition total c (remainder v)) counts

-- | Approximate interval probability via interval WMC.
intervalProbabilityApprox :: Ord a => Imp g a -> (a -> Bool) -> (Double, Double)
intervalProbabilityApprox prog predicate =
  let (total, counts) = compileApprox prog predicate
      getInterval b = Map.findWithDefault zero b counts
  in condition total (getInterval True) (getInterval False)

-- | Approximate lower and upper expectations via interval WMC, clamped to
--   the range of scores the posterior can still reach.
intervalExpectationApprox :: Ord a => Imp g a -> (a -> Double) -> (Double, Double)
intervalExpectationApprox prog score =
  let reachable = [ (score v, b) | (v, b) <- Map.toList (marginalApprox prog), snd b > 0 ]
      scores    = map fst reachable
      (los, his) = unzip [ if s >= 0 then (s * lo, s * hi) else (s * hi, s * lo)
                         | (s, (lo, hi)) <- reachable ]
  in (max (minimum scores) (sum los), min (maximum scores) (sum his))