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 | ||| Function optimizers.
 17 | module Spidr.Optimize
 18 |
 19 | import Data.Vect
 20 | import Data.DPair
 21 | import Data.List.HasLength
 22 |
 23 | import Spidr.Data.Nat
 24 | import Spidr.Tensor
 25 |
 26 | ||| An `Optimizer` finds the value, in a `Tensor`-valued feature space, which (typically
 27 | ||| approximately) optimizes a scalar-valued function over that space.
 28 | |||
 29 | ||| @domain The type of the domain over which to find the optimal value.
 30 | public export 0
 31 | Optimizer : (0 domain : Type) -> Type
 32 | Optimizer a = (a -> Tag $ Tensor [] F64) -> Tag a
 33 |
 34 | stack : {ds : _} -> {dtype : _} -> Vect d (Tensor ds dtype) -> Tensor (d :: ds) dtype
 35 | stack [] = tensor []
 36 | stack (x :: xs) = concat 0 (expand 0 x) (stack xs)
 37 |
 38 | unstack : {d : _} -> Tensor (d :: ds) dtype -> Tag $ Vect d (Tensor ds dtype)
 39 | unstack x = tag x <&> \x => Subset.tabulate $ \(Element i lt) => slice [at i] x
 40 |
 41 | vect : (xs : List a) -> {auto len : HasLength p xs} -> Vect p a
 42 | vect [] @{Z} = []
 43 | vect (x :: xs) @{S len} = x :: vect xs
 44 |
 45 | ||| For each of `d` axes, return a grid of values that scales linearly between `bound` and `bound'`.
 46 | ||| Bounds need not be ordered. In the n-th grid, values vary along the n-th axis, and are constant
 47 | ||| across other axes. The number of values along the axis is the n-th `density`.
 48 | |||
 49 | ||| For example, for
 50 | ||| ```
 51 | ||| lo, hi : Tensor [2] F64
 52 | ||| lo = tensor [0.0, -0.5]
 53 | ||| hi = tensor [1.5, 1.0]
 54 | ||| ```
 55 | ||| `meshGrid [3, 4] lo hi` yields `[x, y]` where
 56 | ||| ```
 57 | ||| x, y : Tensor [3, 4] F64
 58 | ||| x = tensor [
 59 | |||     [0.0 , 0.0 , 0.0 , 0.0 ]
 60 | |||   , [0.75, 0.75, 0.75, 0.75]
 61 | |||   , [1.5 , 1.5 , 1.5 , 1.5 ]
 62 | |||   ]
 63 | ||| y = tensor [
 64 | |||     [-0.5, 0.0, 0.5, 1.0]
 65 | |||   , [-0.5, 0.0, 0.5, 1.0]
 66 | |||   , [-0.5, 0.0, 0.5, 1.0]
 67 | |||   ]
 68 | ||| ```
 69 | |||
 70 | ||| @density The number of values for each axis.
 71 | ||| @bound One bound of the values for each axis.
 72 | ||| @bound' The other bounds.
 73 | export
 74 | meshGrid :
 75 |   {d : _} ->
 76 |   (density : List Nat) ->
 77 |   {auto hl : HasLength d density} ->
 78 |   {auto nonEmpty : All (`GT` 1) density} ->
 79 |   (bound, bound' : Tensor [d] F64) ->
 80 |   Tag $ Vect d $ Tensor density F64
 81 | meshGrid density bound bound' = do
 82 |   let iotas = Subset.tabulate $ \(Element dim lt) => iota {inBounds = ltLengthIsInBound lt hl} dim
 83 |       density' = F64.fromDouble . cast . pred <$> vect density
 84 |   bound <- traverse tag =<< unstack bound
 85 |   bound' <- traverse tag =<< unstack bound'
 86 |   pure [| rescale iotas density' bound bound' |]
 87 |
 88 |   where
 89 |   rescale : Tensor density F64 -> (dens, b, b' : Tensor [] F64) -> Tensor density F64
 90 |   rescale is dens b b' = broadcast (b' - b) * (is / broadcast dens) + broadcast b
 91 |
 92 |   ltLengthIsInBound : {m : _} -> {xs : _} -> LT m n -> HasLength n xs -> InBounds m xs
 93 |   ltLengthIsInBound _ Z impossible
 94 |   ltLengthIsInBound {m = 0} {xs = _ :: _} _ _ = InFirst
 95 |   ltLengthIsInBound {m = S _} (LTESucc lt) (S hl) = InLater (ltLengthIsInBound lt hl)
 96 |
 97 | ||| Grid search of a scalar-valued function. Grid search approximates the optimum by evaluating the
 98 | ||| objective over a bounded, evenly-spaced grid. Bounds need not be ordered.
 99 | |||
100 | ||| @density The density of the grid.
101 | ||| @bound One bound of the grid.
102 | ||| @bound' The other bound of the grid.
103 | export
104 | gridSearch :
105 |   {d : _} ->
106 |   (density : List Nat) ->
107 |   {auto hl : HasLength d density} ->
108 |   {auto nonEmpty : All (`GT` 1) density} ->
109 |   (bound, bound' : Tensor [d] F64) ->
110 |   Optimizer $ Tensor [d] F64
111 | gridSearch {hl = Z} [] _ _ _ = pure (tensor [])
112 | gridSearch {hl = S _} density@(dn :: dns) bound bound' f = do
113 |   grid <- meshGrid density bound bound'
114 |   grid <- tag (stack $ reshape {sizesEqual = flattenable density} <$> grid).T
115 |   [out] <- !(vmap {shapes = [_], dtypes = [_]} $ \x => do pure [!(f x)]) grid
116 |   idx <- argminAS out $ sizeSuccIsSucc (dn :: dns)
117 |   pure $ slice [at idx] grid
118 |
119 |   where
120 |   argminAS : Tensor [s] F64 -> (0 _ : IsSucc s) -> Tag $ Tensor [] U64
121 |   argminAS x ItIsSucc = argmin x
122 |
123 |   0 sizeSuccIsSucc : (xs : List Nat) -> {auto iss : All (`GT` 1) xs} -> IsSucc $ product xs
124 |   sizeSuccIsSucc [] = ItIsSucc
125 |   sizeSuccIsSucc (x :: xs) {iss = is :: _} = multSuccIsSucc (gtIsSucc is) (sizeSuccIsSucc xs)
126 |
127 | ||| The limited-memory BFGS (L-BFGS) optimization tactic, see
128 | |||
129 | ||| Nocedal, Jorge, Updating quasi-Newton matrices with limited storage.
130 | ||| Math. Comp. 35 (1980), no. 151, 773–782.
131 | |||
132 | ||| available at
133 | |||
134 | ||| https://www.ams.org/journals/mcom/1980-35-151/S0025-5718-1980-0572855-7/
135 | |||
136 | ||| **NOTE** This function is not yet implemented.
137 | |||
138 | ||| @initialPoints The points from which to start optimization.
139 | export
140 | lbfgs : (initialPoints : Tensor [n] F64) -> Optimizer $ Tensor [n] F64
141 |