Implement distribution functor

ma.cabal

@@ -35,6 +35,7 @@ library
                      , MA.Functors.Bag
                      , MA.Functors.FixedProduct
                      , MA.Functors.MonoidValued
+                     , MA.Functors.Distribution
                      , MA.Functors.SomeFunctor
                      , MA.Parser
                      , MA.Parser.Lexer
@@ -107,6 +108,7 @@ test-suite spec
                      , MA.Functors.PowersetSpec
                      , MA.Functors.MonoidValuedSpec
                      , MA.Functors.BagSpec
+                     , MA.Functors.DistributionSpec
                      , MA.FunctorExpression.ParserSpec
                      , MA.FunctorExpression.PrettySpec
                      , MA.FunctorExpression.SortsSpec

src/MA/Functors.hs

@@ -8,12 +8,13 @@ import           MA.Functors.FixedProduct (fixedproduct)
 import           MA.Functors.MonoidValued (intValued, realValued)
 import           MA.Functors.Powerset (powerset)
 import           MA.Functors.Bag (bag)
+import           MA.Functors.Distribution (distribution)
 import           MA.Functors.SomeFunctor
 import           MA.FunctorExpression.Parser
 
 registeredFunctors :: [[FunctorParser SomeFunctor]]
 registeredFunctors =
   [ [someFunctor intValued, someFunctor realValued]
-  , [someFunctor powerset, someFunctor bag]
+  , [someFunctor powerset, someFunctor bag, someFunctor distribution]
   , [someFunctor fixedproduct]
   ]

src/MA/Functors/Distribution.hs

+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE TemplateHaskell #-}
+module MA.Functors.Distribution
+  ( distribution
+  , Distribution(..)
+  ) where
+
+import           Prelude hiding (init)
+import           Control.Monad (when)
+
+import           Text.Megaparsec
+import           Data.Eq.Deriving (deriveEq1)
+import           Text.Show.Deriving (deriveShow1)
+
+import           MA.RefinementInterface
+import qualified MA.Parser.Lexer as L
+import           MA.FunctorExpression.Parser
+import           MA.Coalgebra.RefinementTypes
+import           MA.Coalgebra.Parser
+import           MA.Functors.MonoidValued
+
+
+newtype Distribution x = Distribution x
+  deriving (Functor, Foldable, Traversable)
+
+distribution :: FunctorParser Distribution
+distribution = Prefix ((L.symbol "D" <|> L.symbol "Ɗ") >> pure Distribution)
+
+type instance Label Distribution = Label (MonoidValued Double)
+type instance Weight Distribution = Weight (MonoidValued Double)
+type instance H1 Distribution = H1 (MonoidValued Double)
+type instance H3 Distribution = H3 (MonoidValued Double)
+
+$(deriveEq1 ''Distribution)
+$(deriveShow1 ''Distribution)
+
+deriving instance Show (Distribution ())
+
+instance ParseMorphism Distribution where
+  parseMorphismPoint (Distribution inner) = do
+    (h1, succs) <- parseMorphismPoint (RealValued inner)
+
+    when (h1 /= 1) $
+      fail "distribution: Sum of outgoing labels is not 1"
+
+    return (h1, succs)
+
+instance RefinementInterface Distribution where
+  init _ _ = (0, 1)
+  update weightsToS (toRest, toC) =
+    let
+      toS = sum weightsToS
+      toCwithoutS = toC - toS
+      isOk x = x >= 0 && x <= 1
+     in
+      if isOk toRest && isOk toCwithoutS && isOk toS
+        then mkRes (toRest, toCwithoutS, toS)
+        else mkRes (0, 0, 1)
+    where
+      mkRes h3@(a, b, c) = ((a + b, c), h3, (a + c, b))

src/MA/Functors/MonoidValued.hs

 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE FlexibleInstances #-}
 
+-- | The monoid valued functor @M^X@.
+--
+-- The 'RefinementInterface' for this functor is implemented in an abstract
+-- fashion over any type that satisfies the 'Num' constraint, but concrete
+-- parsers only exist for integers and reals.
 module MA.Functors.MonoidValued
   ( intValued
   , realValued
@@ -20,7 +25,8 @@ import qualified MA.Parser.Lexer as L
 import           MA.Parser.Types
 
 data MonoidValued m a where
-  RealValued :: a -> MonoidValued Double a
+  RealValued :: a -> MonoidValued Double a  -- FIXME: Handle rounding and
+                                            -- comparison of doubles sensibly
   IntValued  :: a -> MonoidValued Int a
 
 deriving instance Show (MonoidValued m ())

tests/MA/Functors/DistributionSpec.hs

+module MA.Functors.DistributionSpec (spec) where
+
+import           Test.Hspec
+import           Test.Hspec.Megaparsec
+
+import           Data.Proxy
+import           Control.Monad.ST
+import           Data.Either
+
+import qualified Data.Vector as V
+
+import           Data.MorphismEncoding (Encoding)
+import qualified Data.MorphismEncoding as Encoding
+import qualified Data.Partition as Part
+import           MA.Algorithm
+import           MA.Coalgebra.Parser
+import           MA.FunctorExpression.Desorting
+import           MA.FunctorExpression.Parser
+import           MA.FunctorExpression.Type
+import           MA.Functors.Distribution
+
+
+spec :: Spec
+spec = do
+  functorExpressionSpec
+  parseMorphismPointSpec
+  refineSpec
+
+functorExpressionSpec :: Spec
+functorExpressionSpec = describe "functorExpression" $ do
+  it "parses a normal D correctly" $
+    parseFunctorExpression [[distribution]] "" "DX" `shouldParse` (Functor 1 (Distribution Variable))
+
+  it "parses a unicode D (Ɗ) correctly" $
+    parseFunctorExpression [[distribution]] "" "ƊX" `shouldParse` (Functor 1 (Distribution Variable))
+
+parseMorphismPointSpec :: Spec
+parseMorphismPointSpec = describe "parseMorphismPoint" $ do
+  it "works for a simple example" $
+    (snd <$>
+     parseMorphisms
+       (Functor 1 (Distribution Variable))
+       ""
+       "x: {x: 0.5, y: 0.5}\ny: {x: 1.0}") `shouldParse`
+    encoding [(1, 1), (1, 1)] [(0, (1, 0.5), 0), (0, (1, 0.5), 1), (1, (1, 1), 0)]
+
+  it "errors if esge weight sum isn't 1" $
+    parseMorphisms (Functor 1 (Distribution Variable)) "" `shouldFailOn`
+      "x: {x: 0.5}"
+
+refineSpec :: Spec
+refineSpec = describe "refining" $ do
+  let f = Functor 1 (Distribution Variable)
+
+  it "handles states with incoming edges greater than 1" $ do
+    let x = parseMorphisms f "" "x: {z: 1.0}\ny: {z: 1.0}\nz: {y: 1.0}"
+    x `shouldSatisfy` isRight
+    let Right (_, enc) = x
+    stToIO (refine (Proxy @(Desorted Distribution)) enc) `shouldReturn` Part.fromBlocks [[0, 1, 2]]
+
+-- FIXME: Remove duplicate definition of this function
+encoding :: [h1] -> [(Int, l, Int)] -> Encoding l h1
+encoding h1 es = Encoding.new (V.fromList h1) (V.fromList (map toEdge es))
+  where
+    toEdge (from, lab, to) = Encoding.Edge from lab to