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