@@ -1,3 +1,4 @@
1+{-# LANGUAGE GADTs #-}
12{-# LANGUAGE AllowAmbiguousTypes #-}
23{-# LANGUAGE TypeApplications #-}
34{-# LANGUAGE KindSignatures #-}
@@ -10,6 +11,15 @@ module Dist.Internal
1011 (
1112 Dist(..),
1213 Prob,
14+
15+ fmapD,
16+ unitD,
17+ multD,
18+ selectD,
19+ returnD,
20+ bindD,
21+ (??),
22+
1323 choose,
1424 shape,
1525 linear,
@@ -33,14 +43,78 @@ type Prob = Double
3343
3444-- | Type synonym for column vector matrices. This represents a probability
3545-- distribution.
36−type Dist a = (FromLists Prob () (Normalize a)) => Matrix Prob () a
46+newtype Dist a = D (Matrix Prob () a)
47+
48+-- | Functor instance
49+fmapD ::
50+ ( Bounded a,
51+ Bounded b,
52+ Enum a,
53+ Enum b,
54+ Eq b,
55+ KnownNat (Count (Normalize a)),
56+ KnownNat (Count (Normalize b)),
57+ FromLists Prob (Normalize b) (Normalize a)
58+ )
59+ =>
60+ (a -> b) -> Dist a -> Dist b
61+fmapD f (D m) = D (fromF' f `comp` m)
62+
63+-- | Applicative/Monoidal instance 'unit' function
64+unitD :: Dist ()
65+unitD = D (one 1)
66+
67+-- | Applicative/Monoidal instance 'mult' function
68+multD ::
69+ ( KnownNat (Count (Normalize a)),
70+ KnownNat (Count (Normalize b)),
71+ KnownNat (Count (Normalize (a, b))),
72+ FromLists Prob (Normalize (a, b)) (Normalize a),
73+ FromLists Prob (Normalize (a, b)) (Normalize b),
74+ Normalize (a, b) ~ Normalize (Normalize a, Normalize b)
75+ ) => Dist a -> Dist b -> Dist (a, b)
76+multD (D a) (D b) = D (khatri a b)
77+
78+-- | Selective instance function
79+selectD ::
80+ ( Normalize (Either a b) ~ Either (Normalize a) (Normalize b),
81+ FromLists Prob (Normalize b) (Normalize b),
82+ KnownNat (Count (Normalize b))
83+ ) => Dist (Either a b) -> Matrix Prob a b -> Dist b
84+selectD (D d) m = D (junc m identity `comp` d)
85+
86+
87+-- | Monad instance 'return' function
88+returnD :: forall a . (Enum a, FromLists Prob () (Normalize a), KnownNat (Count a)) => a -> Dist a
89+returnD a = D (col l)
90+ where
91+ i = fromInteger $ natVal (Proxy :: Proxy (Count a))
92+ x = fromEnum a
93+ l = take x [0,0..] ++ [1] ++ take (i - x - 1) [0,0..]
94+
95+-- | Monad instance '(>>=)' function
96+bindD :: Dist a -> Matrix Prob a b -> Dist b
97+bindD (D d) m = D (m `comp` d)
98+
99+-- | Extract probabilities given an Event.
100+(??) ::
101+ ( Enum a,
102+ KnownNat (Count a),
103+ FromLists Prob () (Normalize a)
104+ ) => (a -> Bool) -> Dist a -> Prob
105+(??) p d =
106+ let l = toValues d
107+ x = filter (p . fst) l
108+ in sum . map snd $ x
109+
110+-- Distribution Construction
37111
38112-- | Constructs a Bernoulli distribution
39−choose :: Prob -> Dist a
40−choose prob = col [prob, 1 - prob]
113+choose :: (FromLists Prob () (Normalize a)) => Prob -> Dist a
114+choose prob = D (col [prob, 1 - prob])
41115
42116-- | Creates a distribution given a shape function
43−shape :: (Prob -> Prob) -> [a] -> Dist a
117+shape :: (FromLists Prob () (Normalize a)) => (Prob -> Prob) -> [a] -> Dist a
44118shape _ [] = error "Probability.shape: empty list"
45119shape f xs =
46120 let incr = 1 / fromIntegral (length xs - 1)
@@ -48,24 +122,24 @@ shape f xs =
48122 in fromFreqs (zip xs ps)
49123
50124-- | Constructs a Linear distribution
51−linear :: [a] -> Dist a
125+linear :: (FromLists Prob () (Normalize a)) => [a] -> Dist a
52126linear = shape id
53127
54128-- | Constructs an Uniform distribution
55−uniform :: [a] -> Dist a
129+uniform :: (FromLists Prob () (Normalize a)) => [a] -> Dist a
56130uniform = shape (const 1)
57131
58132-- | Constructs an Negative Exponential distribution
59−negExp :: [a] -> Dist a
133+negExp :: (FromLists Prob () (Normalize a)) => [a] -> Dist a
60134negExp = shape (\x -> exp (-x))
61135
62136-- | Constructs an Normal distribution
63−normal :: [a] -> Dist a
137+normal :: (FromLists Prob () (Normalize a)) => [a] -> Dist a
64138normal = shape (normalCurve 0.5 0.5)
65139
66140-- | Transforms a 'Dist' into a list of pairs.
67141toValues :: forall a . (Enum a, KnownNat (Count a), FromLists Prob () (Normalize a)) => Dist a -> [(a, Prob)]
68−toValues d =
142+toValues (D d) =
69143 let rows = fromInteger (natVal (Proxy :: Proxy (Count a)))
70144 probs = toList d
71145 res = zip (map toEnum [0..rows]) probs
@@ -89,8 +163,8 @@ prettyPrintDist = putStrLn . prettyDist @a
89163
90164-- Auxiliary functions
91165
92−fromFreqs :: [(a,Prob)] -> Dist a
93−fromFreqs xs = col (map (\(x,p) -> p/q) xs)
166+fromFreqs :: (FromLists Prob () (Normalize a)) => [(a,Prob)] -> Dist a
167+fromFreqs xs = D (col (map (\(x,p) -> p/q) xs))
94168 where q = sum $ map snd xs
95169
96170normalCurve :: Prob -> Prob -> Prob -> Prob