Optimize.hs

-- | Approximate credal inference by gradient ascent over the Knightian
--   weights.  Values are inner approximations: every iterate is a point
--   inside the credal set, so they never overshoot the exact bounds.
--
--   Throws an error if the evidence is unsatisfiable, i.e. the credal set is empty.
module Imp.Inference.Optimize
  ( optimizeExpectation
  , optimizeProbability
  ) where

import Control.Applicative ((<|>))
import Control.Monad (guard)
import Data.Bool (bool)
import Data.Maybe (fromMaybe)
import Data.Ord (clamp)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import qualified Data.Vector as V

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

-- | A constant dual number.
dualConst :: Double -> DualS
dualConst x = DualS x V.empty

-- | Dual number division; undefined for a non-positive denominator.
dualDiv :: DualS -> DualS -> Maybe DualS
dualDiv (DualS a da) (DualS b db)
  | b <= 0    = Nothing
  | otherwise = Just (DualS (a / b) (V.zipWith (\a' b' -> (a' * b - a * b') / (b * b)) da db))

-- | The dual-number reading of the variable weights at the given parameters.
dualWeights :: V.Vector Double -> Weight -> (DualS, DualS)
dualWeights params w = case w of
  Prob p   -> (dualConst (1 - p), dualConst p)
  Knight i -> let p = params V.! i
                  k = length params
              in ( DualS (1 - p) (V.generate k (\j -> if j == i then -1 else 0))
                 , DualS p       (V.generate k (\j -> if j == i then  1 else 0)) )

-- | Gradient ascent over the Knightian variables to maximize expected score over
--   the credal set. Ascent starts deterministically at weight @0.5@.
--
--   Returns @(weights, expectation)@ where @weights@ maps each Knightian
--   variable name to its optimized weight.
--
--   Stops early when the gradient magnitude falls below @1e-6@.
--
--   Use a negative learning rate for gradient /descent/ (minimization).
optimizeExpectation :: Ord a
                    => Imp g a
                    -> (a -> Double)
                    -> Int              -- ^ maximum steps
                    -> Double           -- ^ learning rate
                    -> (Map String Double, Double)
optimizeExpectation prog score steps lr =
  let (mgr, weights, knights, worlds) = compile prog
      params0 = V.replicate (Map.size knights) 0.5
      go n params = do
        let counts = wmcBatch (dualWeights params) mgr weights worlds
            aggr   = sumS [ dualConst (score v) .*. count | (v, count) <- Map.toList counts ]
        DualS e de <- dualDiv aggr (sumS (Map.elems counts))
        let params'  = V.zipWith (\p dp -> clamp (0, 1) (p + lr * dp)) params de
            continue = n > 0 && V.sum (V.map (\x -> x * x) de) > 1e-12
        (guard continue >> go (n - 1) params') <|> Just ((params V.!) <$> knights, e)
  in fromMaybe (error "No feasible probabilities") (go steps params0)

-- | Gradient ascent over the Knightian variables to maximize P(event).
optimizeProbability :: Ord a
                    => Imp g a
                    -> (a -> Bool)
                    -> Int              -- ^ maximum steps
                    -> Double           -- ^ learning rate
                    -> (Map String Double, Double)
optimizeProbability prog predicate steps lr =
  optimizeExpectation prog (bool 0 1 . predicate) steps lr