0 | {--
 1 | Copyright (C) 2021  Joel Berkeley
 2 |
 3 | This program is free software: you can redistribute it and/or modify
 4 | it under the terms of the GNU Affero General Public License as published
 5 | by the Free Software Foundation, either version 3 of the License, or
 6 | (at your option) any later version.
 7 |
 8 | This program is distributed in the hope that it will be useful,
 9 | but WITHOUT ANY WARRANTY; without even the implied warranty of
10 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11 | GNU Affero General Public License for more details.
12 |
13 | You should have received a copy of the GNU Affero General Public License
14 | along with this program.  If not, see <https://www.gnu.org/licenses/>.
15 | --}
16 | ||| Probability distributions.
17 | module Spidr.Distribution
18 |
19 | import Data.Nat
20 |
21 | import Spidr.Tensor
22 | import Spidr.Constants
23 |
24 | ||| A joint, or multivariate distribution over a tensor of floating point values, where the first
25 | ||| two central moments (mean and covariance) are known. Every sub-event is assumed to have the
26 | ||| same shape.
27 | |||
28 | ||| @dist Constructs the distribution from the shape of each sub-event and the number of events in
29 | |||   the distribution.
30 | public export
31 | interface Distribution (0 dist : (0 event : Shape) -> (0 dim : Nat) -> Type) where
32 |   ||| The mean of the distribution.
33 |   mean : dist event dim -> Tag $ Tensor (dim :: event) F64
34 |
35 |   ||| The covariance, or correlation, between sub-events.
36 |   cov : dist event dim -> Tag $ Tensor (dim :: dim :: event) F64
37 |
38 | ||| The variance of a single random variable.
39 | export
40 | variance : {event : _} -> Distribution dist => dist event 1 -> Tag $ Tensor (1 :: event) F64
41 | variance dist = squeeze {from = (1 :: 1 :: event)} <$> cov dist
42 |
43 | ||| A joint, or multivariate distribution over a tensor of floating point values, where the density
44 | ||| function and corresponding cumulative density function are known (either analytically or via
45 | ||| approximation). Every sub-event is assumed to have the same shape.
46 | |||
47 | ||| @event The shape of each sub-event.
48 | ||| @dist Constructs the distribution from the shape of each sub-event and the number of events in
49 | |||   the distribution.
50 | public export
51 | interface Distribution dist  =>
52 |   ClosedFormDistribution (0 event : Shape)
53 |     (0 dist : (0 event : Shape) -> (0 dim : Nat) -> Type) where
54 |       ||| The probability density function of the distribution at the specified point.
55 |       pdf : dist event (S d) -> Tensor (S d :: event) F64 -> Tag $ Tensor [] F64
56 |
57 |       ||| The cumulative distribution function of the distribution at the specified point (that is,
58 |       ||| the probability the random variable takes a value less than or equal to the given point).
59 |       cdf : dist event (S d) -> Tensor (S d :: event) F64 -> Tag $ Tensor [] F64
60 |
61 | ||| A joint Gaussian distribution.
62 | |||
63 | ||| @event The shape of each sub-event.
64 | ||| @dim The number of sub-events.
65 | public export
66 | data Gaussian : (0 event : Shape) -> (0 dim : Nat) -> Type where
67 |   ||| @mean The mean of the events.
68 |   ||| @cov The covariance between events.
69 |   MkGaussian : {d : Nat} -> (mean : Tensor (S d :: event) F64) ->
70 |                (cov : Tensor (S d :: S d :: event) F64) ->
71 |                Gaussian event (S d)
72 |
73 | export
74 | Taggable (Gaussian event dim) where
75 |   tag (MkGaussian mean cov) = [| MkGaussian (tag mean) (tag cov) |]
76 |
77 | export
78 | Distribution Gaussian where
79 |   mean (MkGaussian mean' _) = pure mean'
80 |   cov (MkGaussian _ cov') = pure cov'
81 |
82 | ||| **NOTE** `cdf` is implemented only for univariate `Gaussian`.
83 | export
84 | ClosedFormDistribution [1] Gaussian where
85 |   pdf (MkGaussian {d} mean cov) x = do
86 |     cholCov <- tag !(cholesky $ squeeze {to = [S d, S d]} cov)
87 |     tri <- tag $ cholCov |\ squeeze (x - mean)
88 |     let exponent = - tri @@ tri / 2.0
89 |     covSqrtDet <- reduce @{Prod} [0] (diag cholCov)
90 |     let denominator = fromDouble (pow (2.0 * pi) (cast (S d) / 2.0)) * covSqrtDet
91 |     pure (exp exponent / denominator)
92 |
93 |   cdf (MkGaussian {d = S _} _ _) _ =
94 |     assert_total $ idris_crash "CDF not implemented for multivariate Gaussian"
95 |   cdf (MkGaussian {d = 0} mean cov) x =
96 |     pure $ (1.0 + erf (squeeze (x - mean) / (sqrt (squeeze cov * 2.0)))) / 2.0
97 |