Semiring.hs

-- | The 'Semiring' class and the four WMC semirings.
module Imp.Semiring
  ( Semiring(..)
  , ProbS(..)
  , DualS(..)
  , IntervalS(..)
  , PolyS(..)
  , sumS
  ) where

import Data.Bits ((.|.))
import qualified Data.Map.Strict as Map
import qualified Data.Vector as V

-- | The algebra WMC is parameterised over.
class Semiring a where
  zero :: a
  one  :: a
  (.+.) :: a -> a -> a
  (.*.) :: a -> a -> a

infixl 6 .+.
infixl 7 .*.

-- | Sum a list in the semiring.
sumS :: Semiring s => [s] -> s
sumS = foldl' (.+.) zero

-- | Probability semiring for precise WMC.
newtype ProbS = ProbS { unProb :: Double }
  deriving newtype (Show, Eq)

instance Semiring ProbS where
  zero = ProbS 0
  one  = ProbS 1
  ProbS a .+. ProbS b = ProbS (a + b)
  ProbS a .*. ProbS b = ProbS (a * b)

-- | Dual number semiring for forward-mode AD through WMC.
data DualS = DualS !Double !(V.Vector Double)
  deriving stock (Show, Eq)

instance Semiring DualS where
  zero = DualS 0 V.empty
  one  = DualS 1 V.empty
  DualS a da .+. DualS b db = DualS (a + b) (vplus da db)
  DualS a da .*. DualS b db =
    DualS (a * b) (vplus (vscale a db) (vscale b da))

vplus :: V.Vector Double -> V.Vector Double -> V.Vector Double
vplus a b
  | V.null a  = b
  | V.null b  = a
  | otherwise = V.zipWith (+) a b

vscale :: Double -> V.Vector Double -> V.Vector Double
vscale s v = V.map (* s) v

-- | Interval semiring for one-pass, sound WMC.
data IntervalS = IntervalS !Double !Double
  deriving stock (Show, Eq)

instance Semiring IntervalS where
  zero = IntervalS 0 0
  one  = IntervalS 1 1
  IntervalS a b .+. IntervalS c d = IntervalS (a + c) (b + d)
  IntervalS a b .*. IntervalS c d = IntervalS (a * c) (b * d)

-- | Polynomial semiring for symbolic WMC.
newtype PolyS = PolyS { unPoly :: Map.Map Integer Double }
  deriving stock (Show, Eq)

instance Semiring PolyS where
  zero = PolyS Map.empty
  one  = PolyS (Map.singleton 0 1)
  PolyS a .+. PolyS b = PolyS $ Map.filter (/= 0) $ Map.unionWith (+) a b
  PolyS a .*. PolyS b = PolyS $ Map.filter (/= 0) $ Map.fromListWith (+)
    [ (ma .|. mb, ca * cb)
    | (ma, ca) <- Map.toList a
    , (mb, cb) <- Map.toList b
    ]