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 | ||| Kernel functions, particularly for use in Gaussian processes.
17 | module Spidr.Model.Kernel
18 |
19 | import Data.Nat
20 |
21 | import Spidr.Tensor
22 |
23 | ||| A `Kernel` function maps pairs of points in a feature space to the covariance between those two
24 | ||| points in some target space.
25 | |||
26 | ||| @features The shape of the feature domain.
27 | public export 0
28 | Kernel : (0 features : Shape) -> Type
29 | Kernel features =
30 |   {sk, sk' : _} ->
31 |   Tensor (sk :: features) F64 ->
32 |   Tensor (sk' :: features) F64 ->
33 |   Tag $ Tensor [sk, sk'] F64
34 |
35 | scaledL2Norm :
36 |   Tensor [] F64 ->
37 |   {d, n, n' : _} ->
38 |   Tensor [n, S d] F64 ->
39 |   Tensor [n', S d] F64 ->
40 |   Tag $ Tensor [n, n'] F64
41 | scaledL2Norm len x x' =
42 |   let xs = broadcast {to = [n, n', S d]} $ expand 1 x
43 |    in reduce @{Sum} [2] $ ((xs - broadcast (expand 0 x')) / broadcast len) ^ fill 2.0
44 |
45 | ||| The radial basis function, or squared exponential kernel. This is a stationary kernel with form
46 | |||
47 | ||| (\mathbf x_i, \mathbf x_j) \mapsto \exp \left(- \frac{r^2}{2l^2} \right)
48 | |||
49 | ||| where `r^2 = (\mathbf x_i - \mathbf x_j)^ \intercal (\mathbf x_i - \mathbf x_j)` and the
50 | ||| length scale `l > 0`.
51 | |||
52 | ||| Two points that are close in feature space will be more tightly correlated than points that
53 | ||| are further apart. The distance over which the correlation reduces is given by the length
54 | ||| scale `l`. Smaller length scales result in faster-varying target values.
55 | |||
56 | ||| @lengthScale The length scale `l`.
57 | export
58 | rbf : (lengthScale : Tensor [] F64) -> {d : _} -> Kernel [S d]
59 | rbf lengthScale x x' = pure $ exp (- !(scaledL2Norm lengthScale x x') / fill 2.0)
60 |
61 | ||| The Matern kernel for parameter 5/2. This is a stationary kernel with form
62 | |||
63 | ||| (\mathbf x_i, \mathbf x_j) \mapsto \sigma^2 \left(
64 | |||   1 + \frac{\sqrt{5}r}{l} + \frac{5 r^2}{3 l^2}
65 | ||| \right) \exp \left( -\frac{\sqrt{5}r}{l} \right)
66 | |||
67 | ||| where `r^2 = (\mathbf x_i - \mathbf x_j)^ \intercal (\mathbf x_i - \mathbf x_j)` and the
68 | ||| length scale `l > 0`.
69 | |||
70 | ||| @amplitude The amplitude `\sigma`.
71 | ||| @length_scale The length scale `l`.
72 | export
73 | matern52 :
74 |   (amplitude : Tensor [] F64) -> (length_scale : Tensor [] F64) -> {d : _} -> Kernel [S d]
75 | matern52 amp len x x' = do
76 |   d2 <- tag $ fill 5.0 * !(scaledL2Norm len x x')
77 |   d <- tag $ d2 ^ fill 0.5
78 |   pure $ (broadcast $ amp ^ 2.0) * (d2 / fill 3.0 + d + fill 1.0) * exp (- d)
79 |