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