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 | ||| Acquisition functionality for Bayesian optimization. Acquisition functions quantify the
 17 | ||| usefulness of points in a search domain, towards the task of finding a function optimum.
 18 | module Spidr.BayesianOptimization.Acquisition
 19 |
 20 | import Control.Monad.Reader
 21 | import Control.Monad.Identity
 22 | import Data.Nat
 23 |
 24 | import Spidr.Distribution
 25 | import Spidr.Tensor
 26 | import Spidr.Model.Supervised
 27 | import Spidr.Model
 28 |
 29 | %prefix_record_projections off
 30 |
 31 | ||| A `DataModel` packages data with a model over that data.
 32 | public export
 33 | record DataModel modelType {auto probabilisticModel : ProbabilisticModel f t marginal modelType} where
 34 |   constructor MkDataModel
 35 |
 36 |   ||| A probabilistic model
 37 |   model : modelType
 38 |
 39 |   ||| The data the model is trained on
 40 |   dataset : Dataset f t
 41 |
 42 | %prefix_record_projections on
 43 |
 44 | ||| An `Acquisition` function quantifies how useful it would be to query the objective at a given
 45 | ||| set of points, towards the goal of optimizing the objective.
 46 | |||
 47 | ||| @batchSize The number of points in the feature domain that the `Acquisition` evaluates
 48 | |||   at once.
 49 | ||| @features The shape of the feature domain.
 50 | public export 0
 51 | Acquisition : (0 batchSize : Nat) -> {auto 0 _ : GT batchSize 0} -> (0 features : Shape) -> Type
 52 | Acquisition batchSize features = Tensor (batchSize :: features) F64 -> Tag $ Tensor [] F64
 53 |
 54 | ||| Construct the acquisition function that estimates the absolute improvement in the best
 55 | ||| observation if we were to evaluate the objective at a given point.
 56 | |||
 57 | ||| @model The model over the historic data.
 58 | ||| @best The current best observation.
 59 | export
 60 | expectedImprovement :
 61 |   ProbabilisticModel features [1] Gaussian m =>
 62 |   (model : m) ->
 63 |   (best : Tensor [] F64) ->
 64 |   Acquisition 1 features
 65 | expectedImprovement model best at = do
 66 |   best <- tag best
 67 |   marginal <- tag =<< marginalise model at
 68 |   let best' = broadcast {to = [_, 1]} best
 69 |   pdf <- tag =<< pdf marginal best'
 70 |   cdf <- tag =<< cdf marginal best'
 71 |   let mean = squeeze !(mean {event = [1]} {dim = 1} marginal)
 72 |       variance = squeeze !(variance {event = [1]} marginal)
 73 |   pure $ (best - mean) * cdf + variance * pdf
 74 |
 75 | ||| Build an acquisition function that returns the absolute improvement, expected by the model, in
 76 | ||| the observation value at each point.
 77 | export
 78 | expectedImprovementByModel :
 79 |   ProbabilisticModel features [1] Gaussian modelType =>
 80 |   ReaderT (DataModel modelType) Tag $ Acquisition 1 features
 81 | expectedImprovementByModel = MkReaderT $ \env => do
 82 |   marginal <- marginalise env.model env.dataset.features
 83 |   best <- tag $ squeeze !(reduce @{Min} [0] !(mean {event = [1]} marginal))
 84 |   pure $ expectedImprovement env.model best
 85 |
 86 | ||| Build an acquisition function that returns the probability that any given point will take a
 87 | ||| value less than the specified `limit`.
 88 | export
 89 | probabilityOfFeasibility :
 90 |   (limit : Tensor [] F64) ->
 91 |   ClosedFormDistribution [1] dist =>
 92 |   ProbabilisticModel features [1] dist modelType =>
 93 |   ReaderT (DataModel modelType) Tag $ Acquisition 1 features
 94 | probabilityOfFeasibility limit =
 95 |   asks $ \env, at => do cdf !(marginalise env.model at) (broadcast {to = [_, 1]} limit)
 96 |
 97 | ||| Build an acquisition function that returns the negative of the lower confidence bound of the
 98 | ||| probabilistic model. The variance contribution is weighted by a factor `beta`.
 99 | |||
100 | ||| @beta The weighting given to the variance contribution.
101 | export
102 | negativeLowerConfidenceBound :
103 |   (beta : Double) ->
104 |   {auto 0 betaNonNegative : beta >= 0 = True} ->
105 |   ProbabilisticModel features [1] Gaussian modelType =>
106 |   ReaderT (DataModel modelType) Tag $ Acquisition 1 features
107 | negativeLowerConfidenceBound beta = asks $ \env, at => do
108 |   marginal <- tag =<< marginalise env.model at
109 |   pure $ squeeze $
110 |     !(mean {event = [1]} marginal) - fill beta * !(variance {event = [1]} marginal)
111 |
112 | ||| Build the expected improvement acquisition function in the context of a constraint on the input
113 | ||| domain, where points that do not satisfy the constraint do not offer an improvement. The
114 | ||| complete acquisition function is built from a constraint acquisition function, which quantifies
115 | ||| whether specified points in the input space satisfy the constraint.
116 | |||
117 | ||| **NOTE** This function is not yet implemented.
118 | export
119 | expectedConstrainedImprovement :
120 |   (limit : Tensor [] F64) ->
121 |   ProbabilisticModel features [1] Gaussian modelType =>
122 |   ReaderT (DataModel modelType) Tag (Acquisition 1 features -> Acquisition 1 features)
123 |