DSL.hs

-- | Tests the graded monad primitives and the grade algebra.  The type
--   signatures below only compile if 'Merge', 'Union' and 'TagAll' agree.
{-# LANGUAGE QualifiedDo #-}
module Test.DSL (tests) where

import Prelude hiding (return, (>>=), (>>), flip)
import Test.Tasty
import Test.Tasty.HUnit
import Test.Util (assertApprox, assertBounds, assertDist, assertMap, knightNames)

import qualified Imp.DSL as Imp
import Imp.DSL (Imp, flip, ifThenElse, interval, knight, observe, tag)
import Imp.Inference

-- | 'Merge' sorts, so binding out of alphabetical order still gives a sorted grade.
sorted :: Imp '["a", "b", "c"] Bool
sorted = Imp.do
  c <- knight @"c"
  a <- knight @"a"
  b <- knight @"b"
  Imp.return (a && b && c)

-- | The grade of a branch is the 'Union' of its arms.  Explicit because
--   this module does not enable @RebindableSyntax@.
branched :: Imp '["x", "y"] Bool
branched = Imp.do
  c <- flip 0.5
  ifThenElse c
    (Imp.do { a <- knight @"x"; Imp.return a })
    (Imp.do { b <- knight @"y"; Imp.return (not b) })

-- | 'TagAll' prefixes every name in the subprogram.
tagged :: Imp '["t.k"] Bool
tagged = tag @"t" (knight @"k")

-- | Nested tags compose left to right.
nested :: Imp '["a.b.k"] Bool
nested = tag @"a" (tag @"b" (knight @"k"))

-- | An 'interval' keeps its bounds under a tag.
taggedInterval :: Imp '["m.i"] Bool
taggedInterval = tag @"m" (interval @"i" 0.25 0.75)

-- | The README's conditioning example.
conditioned :: Imp '["bias"] Bool
conditioned = Imp.do
  biased <- interval @"bias" 0.3 0.7
  observe biased
  Imp.return biased

tests :: TestTree
tests = testGroup "DSL"
  [ testGroup "Primitives"
    [ testCase "flip p is a Bernoulli" $
        assertDist "flip 0.7" [(False, 0.3), (True, 0.7)] (preciseMarginal (flip 0.7))

    , testCase "knight is completely unconstrained" $
        assertMap "knight" [(False, (0, 1)), (True, (0, 1))] (marginal (knight @"k"))

    , testCase "interval lo hi bounds both outcomes" $ do
        assertMap "interval" [(False, (0.25, 0.75)), (True, (0.25, 0.75))]
          (marginal (interval @"i" 0.25 0.75))
        length (credalVertices (interval @"i" 0.25 0.75)) @?= 2

    , testCase "observe renormalises away the rejected world" $ do
        let prog = Imp.do { h <- flip 0.5; observe h; Imp.return h }
        assertDist "observed" [(True, 1.0)] (preciseMarginal prog)

    , testCase "(>>) discards the first result but keeps its evidence" $
        assertDist "sequenced" [(False, 0.75), (True, 0.25)]
          (preciseMarginal (observe True Imp.>> flip 0.25))

    , testCase "graded fmap maps the value, not the weight" $ do
        assertDist "fmap" [(False, 0.7), (True, 0.3)]
          (preciseMarginal (Imp.fmap not (flip 0.7)))
        assertDist "<$>" [(False, 0.7), (True, 0.3)]
          (preciseMarginal (not Imp.<$> flip 0.7))
    ]

  , testGroup "Grades"
    [ testCase "Merge sorts names regardless of bind order" $
        knightNames sorted @?= ["a", "b", "c"]

    , testCase "Union: a probabilistic condition reaches both branches" $ do
        knightNames branched @?= ["x", "y"]
        length (credalVertices branched) @?= 4
        assertMap "branched" [(False, (0.0, 1.0)), (True, (0.0, 1.0))] (marginal branched)

    , testCase "ifThenElse on non-Imp branches is the ordinary conditional" $ do
        ifThenElse True "yes" "no" @?= "yes"
        ifThenElse False "yes" "no" @?= "no"

    , testCase "tag prefixes the Knightian name" $
        knightNames tagged @?= ["t.k"]

    , testCase "nested tags join with dots, outermost first" $
        knightNames nested @?= ["a.b.k"]

    , testCase "the empty tag leaves the name alone" $
        knightNames (tag @"t" (tag @"" (knight @"k")) :: Imp '["t.k"] Bool) @?= ["t.k"]

    , testCase "tag preserves the interval bounds it wraps" $ do
        knightNames taggedInterval @?= ["m.i"]
        assertMap "tagged interval" [(False, (0.25, 0.75)), (True, (0.25, 0.75))]
          (marginal taggedInterval)
    ]

  , testGroup "Conditioning"
    [ testCase "conditioned: every backend gives P(True) = 1" $ do
        assertMap "marginal" [(False, (0, 0)), (True, (1, 1))] (marginal conditioned)
        assertMap "marginalSymbolic" [(False, (0, 0)), (True, (1, 1))]
          (marginalSymbolic conditioned)
        assertMap "marginalApprox" [(False, (0, 0)), (True, (1, 1))]
          (marginalApprox conditioned)
        assertBounds "intervalProbability" (1, 1) (intervalProbability conditioned id)
        assertBounds "intervalProbabilityApprox" (1, 1)
          (intervalProbabilityApprox conditioned id)
        assertBounds "intervalProbabilitySymbolic" (1, 1)
          (intervalProbabilitySymbolic conditioned id)

    , testCase "conditioned: optimization cannot leave the credal set" $ do
        assertApprox "ascent" 1.0 (snd (optimizeProbability conditioned id 200 0.1))
        assertApprox "descent" 1.0 (snd (optimizeProbability conditioned id 200 (-0.1)))
    ]
  ]