@@ -1,3 +1,4 @@
1+{-# LANGUAGE ConstraintKinds #-}
12{-# LANGUAGE GADTs #-}
23{-# LANGUAGE AllowAmbiguousTypes #-}
34{-# LANGUAGE TypeApplications #-}
@@ -8,10 +9,18 @@
89{-# LANGUAGE RankNTypes #-}
910
1011module LAoP.Dist.Internal
11− (
12+ (
1213 Dist(..),
1314 Prob,
1415
16+ Countable,
17+ CountableN,
18+ CountableDimensionsN,
19+ FromListsN,
20+ Liftable,
21+ TrivialP,
22+ TrivialE,
23+
1524 fmapD,
1625 unitD,
1726 multD,
@@ -19,7 +28,7 @@ module LAoP.Dist.Internal
1928 returnD,
2029 bindD,
2130 (??),
22−
31+
2332 choose,
2433 shape,
2534 linear,
@@ -32,7 +41,8 @@ module LAoP.Dist.Internal
3241 )
3342 where
3443
35−import LAoP.Matrix.Type
44+import LAoP.Matrix.Type hiding (TrivialP, TrivialE, Countable, CountableDimensions, CountableN, CountableDimensionsN, Liftable, FromListsN)
45+import qualified LAoP.Matrix.Internal as I
3646import LAoP.Utils
3747import GHC.TypeLits
3848import Data.Proxy
@@ -45,16 +55,20 @@ type Prob = Double
4555-- distribution.
4656newtype Dist a = D (Matrix Prob () a)
4757
58+-- | Constraint type synonyms to keep the type signatures less convoluted
59+type Countable a = KnownNat (I.Count a)
60+type CountableN a = KnownNat (I.Count (I.Normalize a))
61+type CountableDimensionsN a b = (CountableN a, CountableN b)
62+type FromListsN a b = I.FromLists Prob (I.Normalize a) (I.Normalize b)
63+type Liftable a b = (Bounded a, Bounded b, Enum a, Enum b, Eq b, Num Prob, Ord Prob)
64+type TrivialP a b = Normalize (a, b) ~ Normalize (Normalize a, Normalize b)
65+type TrivialE a b = Normalize (Either a b) ~ Either (Normalize a) (Normalize b)
66+
4867-- | Functor instance
4968fmapD ::
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)
69+ ( Liftable a b,
70+ CountableDimensionsN a b,
71+ FromListsN b a
5872 )
5973 =>
6074 (a -> b) -> Dist a -> Dist b
@@ -66,26 +80,25 @@ unitD = D (one 1)
6680
6781-- | Applicative/Monoidal instance 'mult' function
6882multD ::
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)
83+ ( CountableDimensionsN a b,
84+ CountableN (a, b),
85+ FromListsN (a, b) a,
86+ FromListsN (a, b) b,
87+ TrivialP a b
7588 ) => Dist a -> Dist b -> Dist (a, b)
7689multD (D a) (D b) = D (khatri a b)
7790
7891-- | Selective instance function
7992selectD ::
80− ( Normalize (Either a b) ~ Either (Normalize a) (Normalize b),
81− FromLists Prob (Normalize b) (Normalize b),
82− KnownNat (Count (Normalize b))
93+ ( TrivialE a b,
94+ FromListsN b b,
95+ CountableN b
8396 ) => Dist (Either a b) -> Matrix Prob a b -> Dist b
8497selectD (D d) m = D (junc m identity `comp` d)
8598
8699
87100-- | Monad instance 'return' function
88−returnD :: forall a . (Enum a, FromLists Prob () (Normalize a), KnownNat (Count a)) => a -> Dist a
101+returnD :: forall a . (Enum a, FromListsN () a, Countable a) => a -> Dist a
89102returnD a = D (col l)
90103 where
91104 i = fromInteger $ natVal (Proxy :: Proxy (Count a))
@@ -99,8 +112,8 @@ bindD (D d) m = D (m `comp` d)
99112-- | Extract probabilities given an Event.
100113(??) ::
101114 ( Enum a,
102− KnownNat (Count a),
103− FromLists Prob () (Normalize a)
115+ Countable a,
116+ FromListsN () a
104117 ) => (a -> Bool) -> Dist a -> Prob
105118(??) p d =
106119 let l = toValues d
@@ -110,11 +123,11 @@ bindD (D d) m = D (m `comp` d)
110123-- Distribution Construction
111124
112125-- | Constructs a Bernoulli distribution
113−choose :: (FromLists Prob () (Normalize a)) => Prob -> Dist a
126+choose :: (FromListsN () a) => Prob -> Dist a
114127choose prob = D (col [prob, 1 - prob])
115128
116129-- | Creates a distribution given a shape function
117−shape :: (FromLists Prob () (Normalize a)) => (Prob -> Prob) -> [a] -> Dist a
130+shape :: (FromListsN () a) => (Prob -> Prob) -> [a] -> Dist a
118131shape _ [] = error "Probability.shape: empty list"
119132shape f xs =
120133 let incr = 1 / fromIntegral (length xs - 1)
@@ -122,23 +135,23 @@ shape f xs =
122135 in fromFreqs (zip xs ps)
123136
124137-- | Constructs a Linear distribution
125−linear :: (FromLists Prob () (Normalize a)) => [a] -> Dist a
138+linear :: (FromListsN () a) => [a] -> Dist a
126139linear = shape id
127140
128141-- | Constructs an Uniform distribution
129−uniform :: (FromLists Prob () (Normalize a)) => [a] -> Dist a
142+uniform :: (FromListsN () a) => [a] -> Dist a
130143uniform = shape (const 1)
131144
132145-- | Constructs an Negative Exponential distribution
133−negExp :: (FromLists Prob () (Normalize a)) => [a] -> Dist a
146+negExp :: (FromListsN () a) => [a] -> Dist a
134147negExp = shape (\x -> exp (-x))
135148
136149-- | Constructs an Normal distribution
137−normal :: (FromLists Prob () (Normalize a)) => [a] -> Dist a
150+normal :: (FromListsN () a) => [a] -> Dist a
138151normal = shape (normalCurve 0.5 0.5)
139152
140153-- | Transforms a 'Dist' into a list of pairs.
141−toValues :: forall a . (Enum a, KnownNat (Count a), FromLists Prob () (Normalize a)) => Dist a -> [(a, Prob)]
154+toValues :: forall a . (Enum a, Countable a, FromListsN () a) => Dist a -> [(a, Prob)]
142155toValues (D d) =
143156 let rows = fromInteger (natVal (Proxy :: Proxy (Count a)))
144157 probs = toList d
@@ -146,7 +159,7 @@ toValues (D d) =
146159 in res
147160
148161-- | Pretty a distribution
149−prettyDist :: forall a. (Show a, Enum a, KnownNat (Count a), FromLists Prob () (Normalize a)) => Dist a -> String
162+prettyDist :: forall a. (Show a, Enum a, Countable a, FromListsN () a) => Dist a -> String
150163prettyDist d =
151164 let values = sortBy (\(a, p1) (b, p2) -> compare p2 p1) (toValues @a d)
152165 w = maximum (map (length . show . fst) values)
@@ -158,12 +171,12 @@ prettyDist d =
158171 showR n x = show x ++ " "
159172
160173-- | Pretty Print a distribution
161−prettyPrintDist :: forall a . (Show a, Enum a, KnownNat (Count a), FromLists Prob () (Normalize a)) => Dist a -> IO ()
174+prettyPrintDist :: forall a . (Show a, Enum a, Countable a, FromListsN () a) => Dist a -> IO ()
162175prettyPrintDist = putStrLn . prettyDist @a
163176
164177-- Auxiliary functions
165178
166−fromFreqs :: (FromLists Prob () (Normalize a)) => [(a,Prob)] -> Dist a
179+fromFreqs :: (FromListsN () a) => [(a,Prob)] -> Dist a
167180fromFreqs xs = D (col (map (\(x,p) -> p/q) xs))
168181 where q = sum $ map snd xs
169182