Inference.hs

-- | Behaviour of the inference backends on minimal inline programs, plus
--   cross-backend agreement.  Example-program facts live in Test.Examples.
{-# LANGUAGE QualifiedDo #-}
module Test.Inference (tests) where

import Prelude hiding (return, (>>=), (>>), flip)
import qualified Data.Map.Strict as Map
import Test.Tasty
import Test.Tasty.HUnit
import Test.Util (assertApprox, assertBounds, assertContains, assertDist, assertMap,
                  assertNoFeasible)
import qualified Imp.DSL as Imp
import Imp.DSL (Imp, flip, knight, observe)
import Imp.DSL.Combinators (GenNames, knightN)
import Imp.Inference
import Imp.Examples.Ellsberg as E
import Imp.Examples.IMDP as IMDP
import Imp.Examples.Knightian as K
import Imp.Examples.MontyHall as MH
import Imp.Examples.Polytope as P
import Imp.Examples.TwoChild as TC

andCoins :: Imp '[] Bool
andCoins = Imp.do
  a <- flip 0.5
  b <- flip 0.5
  Imp.return (a && b)

-- | A flip widened by a Knightian choice.
orKnight :: Imp '["k"] Bool
orKnight = Imp.do
  a <- flip 0.5
  k <- knight @"k"
  Imp.return (a || k)

-- | Conditioning on a Knightian choice.
observedKnight :: Imp '["k"] Bool
observedKnight = Imp.do
  k <- knight @"k"
  observe k
  Imp.return k

-- | Conditioning and widening by a Knightian choice.
observedOr :: Imp '["k"] Bool
observedOr = Imp.do
  k <- knight @"k"
  h <- flip 0.5
  observe (k || h)
  Imp.return h

-- | Evidence satisfiable only via a rare fault.  A feasibility threshold
--   scaled too loosely discards a genuinely feasible corner.
rareFault :: Imp (GenNames 4 "s") Bool
rareFault = Imp.do
  ss <- knightN @4 @"s"
  f  <- flip 1e-18
  observe (not (or ss) || f)
  Imp.return (or ss)

-- | Empty credal set.
infeasible :: Imp '[] Bool
infeasible = Imp.do
  h <- flip 0.5
  observe (h && not h)
  Imp.return h

-- | Mixed-sign reward.
reward :: Position -> Double
reward v = case v of IMDP.P2 -> 10; IMDP.P1 -> -1; _ -> 0

-- | Mixed-sign score over the three-way polytope.
mixedScore :: P.Three -> Double
mixedScore v = case v of P.Red -> 2.0; P.Green -> -1.0; P.Blue -> 0.5

-- | Enumerate and Symbolic must agree on both marginals and expectations.
agree :: (Ord a, Show a) => String -> Imp g a -> (a -> Double) -> Assertion
agree label prog score = do
  let enum = marginal prog
      sym  = marginalSymbolic prog
  Map.keys sym @?= Map.keys enum
  mapM_ (\(v, (e, s)) -> assertBounds (label ++ " marginal " ++ show v) e s)
        (Map.toList (Map.intersectionWith (,) enum sym))
  assertBounds (label ++ " expectation")
    (intervalExpectation prog score) (intervalExpectationSymbolic prog score)

tests :: TestTree
tests = testGroup "Inference"
  [ testGroup "Precise"
    [ testCase "fair coin" $
        assertDist "flip 0.5" [(False, 0.5), (True, 0.5)] (preciseMarginal (flip 0.5))

    , testCase "biased coin" $
        assertDist "flip 0.7" [(False, 0.3), (True, 0.7)] (preciseMarginal (flip 0.7))

    , testCase "AND of two fair coins" $
        assertDist "and" [(False, 0.75), (True, 0.25)] (preciseMarginal andCoins)
    ]

  , testGroup "Exact enumeration"
    [ testCase "no Knightian vars: bounds collapse to a point" $ do
        assertBounds "P(flip 0.5)" (0.5, 0.5) (intervalProbability (flip 0.5) id)
        assertBounds "P(flip 0.7)" (0.7, 0.7) (intervalProbability (flip 0.7) id)
        assertBounds "P(and)" (0.25, 0.25) (intervalProbability andCoins id)

    , testCase "flip OR knight: P(True) = [0.5, 1]" $
        assertBounds "P" (0.5, 1.0) (intervalProbability orKnight id)

    , testCase "expectation of a scaled indicator" $
        assertBounds "E" (6.0, 6.0)
          (intervalExpectation (flip 0.6) (\b -> if b then 10.0 else 0.0))

    , testCase "conditioning can make a Knightian valuation infeasible" $ do
        assertMap "marginal" [(False, (0, 0)), (True, (1, 1))] (marginal observedKnight)
        length (credalVertices observedKnight) @?= 1

    , testCase "observedOr: P(True) = [0.5, 1]" $
        assertBounds "P" (0.5, 1.0) (intervalProbability observedOr id)
    ]

  , testGroup "Interval approximation"
    [ testCase "exact when each Knightian var is read once" $
        assertMap "marginalApprox" [(False, (1/3, 1/3)), (True, (2/3, 2/3))]
          (marginalApprox MH.montyHall)

    , testCase "exact when the Knightian name is shared" $
        assertMap "marginalApprox"
          [ (K.Red, (0.0, 1.0)), (K.Green, (0.0, 0.5)), (K.Blue, (0.0, 0.5)) ]
          (marginalApprox K.dependent)

    , testCase "loose but sound on simpleRobot" $ do
        assertMap "marginalApprox"
          [ (IMDP.P0, (0.0, 0.25)), (IMDP.P1, (0.0, 1.0)), (IMDP.P2, (0.0, 1.0)) ]
          (marginalApprox IMDP.simpleRobot)
        containsExact "simpleRobot" (marginal IMDP.simpleRobot)
                      (marginalApprox IMDP.simpleRobot)

    , testCase "contains the exact marginal" $ do
        containsExact "complexRobot" (marginal IMDP.complexRobot)
                      (marginalApprox IMDP.complexRobot)
        containsExact "polytope2" (marginal P.polytope2) (marginalApprox P.polytope2)
        containsExact "twoChild" (marginal TC.twoChild) (marginalApprox TC.twoChild)

    , testCase "contains the exact probability under conditioning" $ do
        assertContains "observedKnight" (intervalProbability observedKnight id)
                       (intervalProbabilityApprox observedKnight id)
        assertContains "observedOr" (intervalProbability observedOr id)
                       (intervalProbabilityApprox observedOr id)
        assertContains "twoChild" (intervalProbability TC.twoChild id)
                       (intervalProbabilityApprox TC.twoChild id)

    , testCase "contains the exact expectation" $ do
        assertContains "complexRobot" (intervalExpectation IMDP.complexRobot reward)
                       (intervalExpectationApprox IMDP.complexRobot reward)
        assertContains "observedOr"
          (intervalExpectation observedOr (\b -> if b then 3.0 else -1.0))
          (intervalExpectationApprox observedOr (\b -> if b then 3.0 else -1.0))

    , testCase "a value with no remaining mass drops out of the score range" $ do
        assertBounds "positive on True" (3.0, 3.0)
          (intervalExpectationApprox observedKnight (\b -> if b then 3.0 else -1.0))
        assertBounds "negative on True" (-1.0, -1.0)
          (intervalExpectationApprox observedKnight (\b -> if b then -1.0 else 3.0))

    , testCase "the score range clamps a sum of loose per-value boxes" $ do
        -- The P0/P1 boxes sum past the extreme score; the clamp pulls it back.
        assertBounds "two negatives" (-1.0, 0.0)
          (intervalExpectationApprox IMDP.simpleRobot
             (\v -> case v of IMDP.P2 -> 0; _ -> -1))
        assertBounds "two positives" (0.0, 1.0)
          (intervalExpectationApprox IMDP.simpleRobot
             (\v -> case v of IMDP.P2 -> 0; _ -> 1))

    , testCase "no Knightian vars: expectation is exact" $
        assertBounds "E" (6.0, 6.0)
          (intervalExpectationApprox (flip 0.6) (\b -> if b then 10.0 else 0.0))

    , testCase "a predicate no world satisfies" $ do
        assertBounds "const True" (1.0, 1.0)
          (intervalProbabilityApprox (flip 0.5) (const True))
        assertBounds "const False" (0.0, 0.0)
          (intervalProbabilityApprox (flip 0.5) (const False))
    ]

  , testGroup "Symbolic"
    [ testCase "no Knightian: P(True) = [0.7, 0.7]" $
        assertBounds "P" (0.7, 0.7) (intervalProbabilitySymbolic (flip 0.7) id)

    , testCase "dependent P(Green) = [0, 0.5]" $
        assertBounds "P" (0.0, 0.5) (intervalProbabilitySymbolic K.dependent (== K.Green))

    , testCase "ellsberg P(Red) = [1/3, 1/3], P(Black) = [0, 2/3]" $ do
        assertBounds "Red" (1/3, 1/3) (intervalProbabilitySymbolic E.ellsberg (== E.Red))
        assertBounds "Black" (0, 2/3) (intervalProbabilitySymbolic E.ellsberg (== E.Black))

    , testCase "montyHall: P(switch wins) = [2/3, 2/3]" $
        assertBounds "P" (2/3, 2/3) (intervalProbabilitySymbolic MH.montyHall id)

    , testCase "polytope P(R) = [0.2, 0.8]" $
        assertBounds "P" (0.2, 0.8) (intervalProbabilitySymbolic P.polytope (== P.Red))

    , testCase "expectation: ellsberg E[bet Red] = [1/3, 1/3]" $
        assertBounds "E" (1/3, 1/3)
          (intervalExpectationSymbolic E.ellsberg (\b -> if b == E.Red then 1.0 else 0.0))

    , testCase "expectation: no Knightian is exact" $
        assertBounds "E" (6.0, 6.0)
          (intervalExpectationSymbolic (flip 0.6) (\b -> if b then 10.0 else 0.0))

    , testCase "at least as tight as the interval approximation" $ do
        assertContains "ellsberg Black" (intervalProbabilitySymbolic E.ellsberg (== E.Black))
                       (intervalProbabilityApprox E.ellsberg (== E.Black))
        assertContains "polytope2 Red" (intervalProbabilitySymbolic P.polytope2 (== P.Red))
                       (intervalProbabilityApprox P.polytope2 (== P.Red))
    ]

  , testGroup "Enumeration and Symbolic agree"
    [ testCase "dependent" $ agree "dependent" K.dependent
        (\v -> case v of K.Red -> 1.0; K.Green -> 0.5; K.Blue -> 0.0)
    , testCase "complexRobot (mixed-sign reward)" $
        agree "complexRobot" IMDP.complexRobot reward
    , testCase "polytope2 (mixed-sign score)" $
        agree "polytope2" P.polytope2 mixedScore
    , testCase "montyHall" $
        agree "montyHall" MH.montyHall (\b -> if b then 1.0 else 0.0)
    , testCase "twoChild (observe)" $
        agree "twoChild" TC.twoChild (\b -> if b then 3.0 else -1.0)
    , testCase "observedKnight (an infeasible corner)" $
        agree "observedKnight" observedKnight (\b -> if b then 1.0 else 0.0)
    , testCase "observedOr (non-constant denominator)" $
        agree "observedOr" observedOr (\b -> if b then 3.0 else -1.0)
    , testCase "rareFault (a feasible corner with tiny evidence)" $
        agree "rareFault" rareFault (\b -> if b then 1.0 else 0.0)
    ]

  , testGroup "Gradient optimization"
    [ testCase "complexRobot: ascent on P(P2) reaches the exact upper bound" $ do
        let (weights, prob) = optimizeProbability IMDP.complexRobot (== IMDP.P2) 200 0.1
        -- move*.b has zero gradient once move*.f is 1, so it keeps its 0.5 init.
        weights @?= Map.fromList
          [("move1.b", 0.5), ("move1.f", 1.0), ("move2.b", 0.5), ("move2.f", 1.0)]
        assertApprox "P(P2)" 0.64 prob
        assertApprox "= exact upper" prob
          (snd (intervalProbability IMDP.complexRobot (== IMDP.P2)))

    , testCase "complexRobot: ascent on the reward prefers backing off P1" $ do
        let (weights, val) = optimizeExpectation IMDP.complexRobot reward 200 0.1
        weights @?= Map.fromList
          [("move1.b", 0.5), ("move1.f", 1.0), ("move2.b", 1.0), ("move2.f", 1.0)]
        assertApprox "E[reward]" 6.112 val
        assertApprox "= exact upper" (snd (intervalExpectation IMDP.complexRobot reward)) val

    , testCase "complexRobot: descent is not a mirror of ascent" $ do
        let (weights, val) = optimizeExpectation IMDP.complexRobot reward 200 (-0.1)
        weights @?= Map.fromList
          [("move1.b", 0.5), ("move1.f", 0.0), ("move2.b", 0.0), ("move2.f", 0.0)]
        assertApprox "E[reward]" 2.0 val
        assertApprox "= exact lower" (fst (intervalExpectation IMDP.complexRobot reward)) val

    , testCase "ellsberg: P(Red) is precise, so ascent cannot move it" $ do
        let (weights, prob) = optimizeProbability E.ellsberg (== E.Red) 200 0.1
        weights @?= Map.fromList [("split", 0.5)]
        assertApprox "P(Red)" (1/3) prob

    , testCase "ellsberg: E[bet Red] is precise too" $
        assertApprox "E" (1/3)
          (snd (optimizeExpectation E.ellsberg
                  (\b -> if b == E.Red then 1 else 0) 200 0.1))

    , testCase "conditioning: iterates stay inside the credal set" $ do
        assertApprox "ascent" 1.0 (snd (optimizeProbability observedOr id 200 0.1))
        assertApprox "descent" 0.5 (snd (optimizeProbability observedOr id 200 (-0.1)))
        let (lo, hi) = intervalProbability observedOr id
        assertBounds "exact" (0.5, 1.0) (lo, hi)

    , testCase "one ascent step moves by the exact gradient" $ do
        -- P(h | k or h) = 1/(1 + p_k), so the derivative at 0.5 is -4/9.
        let (weights, prob) = optimizeProbability observedOr id 1 1.0
        assertApprox "weight" (0.5 - 4/9) (weights Map.! "k")
        assertApprox "P" (18/19) prob

    , testCase "no Knightian vars: empty weights and the exact value" $ do
        let (weights, prob) = optimizeProbability (flip 0.6) id 100 0.1
        weights @?= Map.empty
        assertApprox "prob" 0.6 prob
        let (weights', val) = optimizeExpectation (flip 0.6) (\b -> if b then 10 else 0) 100 0.1
        weights' @?= Map.empty
        assertApprox "E[f]" 6.0 val
    ]

  , testGroup "Empty credal set"
    [ testCase "credalVertices is empty rather than an error" $
        credalVertices infeasible @?= []

    , testCase "every bound-producing entry point throws" $ do
        assertNoFeasible "preciseMarginal" (preciseMarginal infeasible)
        assertNoFeasible "marginal" (marginal infeasible)
        assertNoFeasible "intervalProbability" (intervalProbability infeasible id)
        assertNoFeasible "intervalExpectation"
          (intervalExpectation infeasible (\b -> if b then 1 else 0))
        assertNoFeasible "marginalApprox" (marginalApprox infeasible)
        assertNoFeasible "intervalProbabilityApprox" (intervalProbabilityApprox infeasible id)
        assertNoFeasible "intervalExpectationApprox"
          (intervalExpectationApprox infeasible (\b -> if b then 1 else 0))
        assertNoFeasible "marginalSymbolic" (marginalSymbolic infeasible)
        assertNoFeasible "intervalProbabilitySymbolic"
          (intervalProbabilitySymbolic infeasible id)
        assertNoFeasible "intervalExpectationSymbolic"
          (intervalExpectationSymbolic infeasible (\b -> if b then 1 else 0))
        assertNoFeasible "optimizeProbability" (optimizeProbability infeasible id 100 0.1)
        assertNoFeasible "optimizeExpectation"
          (optimizeExpectation infeasible (\b -> if b then 1 else 0) 100 0.1)
    ]
  ]

-- | Assert that the approximate bounds contain the exact ones.
containsExact :: (Ord a, Show a)
              => String -> Map.Map a (Double, Double) -> Map.Map a (Double, Double)
              -> Assertion
containsExact label exact approx = do
  Map.keys approx @?= Map.keys exact
  mapM_ (\(v, (e, a)) -> assertContains (label ++ " " ++ show v) e a)
        (Map.toList (Map.intersectionWith (,) exact approx))