diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 325e32da..df8a6d85 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -1,6 +1,7 @@ Changelog ========= +- Add option to use additive or multiplicative adjustment in any acquisition method - Add `arziv`-mocking to rtd-setup - Add convenience method for obtaining elfi samples as `InferenceData`` to be used with `arviz` - Improve `randmaxvar` batch acquisitions and initialisation by enabling sampling from prior diff --git a/elfi/methods/bo/acquisition.py b/elfi/methods/bo/acquisition.py index fa94bba7..0f44c656 100644 --- a/elfi/methods/bo/acquisition.py +++ b/elfi/methods/bo/acquisition.py @@ -7,7 +7,7 @@ import scipy.stats as ss import elfi.methods.mcmc as mcmc -from elfi.methods.bo.utils import CostFunction, minimize +from elfi.methods.bo.utils import minimize from elfi.methods.utils import resolve_sigmas logger = logging.getLogger(__name__) @@ -223,7 +223,7 @@ class LCBSC(AcquisitionBase): """ - def __init__(self, *args, delta=None, additive_cost=None, **kwargs): + def __init__(self, *args, delta=None, **kwargs): """Initialize LCBSC. Parameters @@ -231,8 +231,6 @@ def __init__(self, *args, delta=None, additive_cost=None, **kwargs): delta: float, optional In between (0, 1). Default is 1/exploration_rate. If given, overrides the exploration_rate. - additive_cost: CostFunction, optional - Cost function output is added to the base acquisition value. """ if delta is not None: @@ -244,10 +242,6 @@ def __init__(self, *args, delta=None, additive_cost=None, **kwargs): self.name = 'lcbsc' self.label_fn = 'Confidence Bound' - if additive_cost is not None and not isinstance(additive_cost, CostFunction): - raise TypeError("Additive cost must be type CostFunction.") - self.additive_cost = additive_cost - @property def delta(self): """Return the inverse of exploration rate.""" @@ -275,8 +269,6 @@ def evaluate(self, x, t=None): """ mean, var = self.model.predict(x, noiseless=True) value = mean - np.sqrt(self._beta(t) * var) - if self.additive_cost is not None: - value += self.additive_cost.evaluate(x) return value def evaluate_gradient(self, x, t=None): @@ -296,8 +288,6 @@ def evaluate_gradient(self, x, t=None): mean, var = self.model.predict(x, noiseless=True) grad_mean, grad_var = self.model.predictive_gradients(x) value = grad_mean - 0.5 * grad_var * np.sqrt(self._beta(t) / var) - if self.additive_cost is not None: - value += self.additive_cost.evaluate_gradient(x) return value diff --git a/elfi/methods/bo/utils.py b/elfi/methods/bo/utils.py index 45ec826b..3e94043e 100644 --- a/elfi/methods/bo/utils.py +++ b/elfi/methods/bo/utils.py @@ -111,20 +111,20 @@ def minimize(fun, return locs[ind_min], vals[ind_min] -class CostFunction: - """Convenience class for modelling acquisition costs.""" +class AdjustmentFunction: + """Convenience class for modelling acquisition function adjustments.""" def __init__(self, function, gradient, scale=1): - """Initialise CostFunction. + """Initialise AdjustmentFunction. Parameters ---------- function : callable - Function that returns cost function value. + Function that returns adjustment function value. gradient : callable - Function that returns cost function gradient. + Function that returns adjustment function gradient. scale : float, optional - Cost function is multiplied with scale. + Adjustment function is multiplied with scale. """ self.function = function @@ -132,7 +132,7 @@ def __init__(self, function, gradient, scale=1): self.scale = scale def evaluate(self, x): - """Return cost function value evaluated at x. + """Return adjustment function value evaluated at x. Parameters ---------- @@ -148,7 +148,7 @@ def evaluate(self, x): return self.scale * self.function(x).reshape(n, 1) def evaluate_gradient(self, x): - """Return cost function gradient evaluated at x. + """Return adjustment function gradient evaluated at x. Parameters ---------- @@ -162,3 +162,67 @@ def evaluate_gradient(self, x): x = np.atleast_2d(x) n, input_dim = x.shape return self.scale * self.gradient(x).reshape(n, input_dim) + + +def make_additive_acq(acquisition_class, function): + """Make acquisition function adjusted with an additive term. + + Parameters + ---------- + acquisition_class : Type[elfi.methods.bo.acquisition.AcquisitionBase] + Acquisition function to be adjusted. + function : AdjustmentFunction + Function added to the base acquisition function. + + Returns + ------- + Type[AdjustedAcquisition] + + """ + class AdjustedAcquisition(acquisition_class): + + def __init__(self, model, **kwargs): + super().__init__(model=model, **kwargs) + self._func = function + + def evaluate(self, theta_new, t=None): + return super().evaluate(theta_new, t=t) + self._func.evaluate(theta_new) + + def evaluate_gradient(self, theta_new, t=None): + t1 = super().evaluate_gradient(theta_new, t=t) + t2 = self._func.evaluate_gradient(theta_new) + return t1 + t2 + + return AdjustedAcquisition + + +def make_multiplicative_acq(acquisition_class, function): + """Make acquisition function adjusted with a multiplictive term. + + Parameters + ---------- + acquisition_class : Type[elfi.methods.bo.acquisition.AcquisitionBase] + Acquisition function to be adjusted. + function : AdjustmentFunction + Function that multiplies the base acquisition function. + + Returns + ------- + Type[AdjustedAcquisition] + + """ + class AdjustedAcquisition(acquisition_class): + + def __init__(self, model, **kwargs): + super().__init__(model=model, **kwargs) + self._func = function + + def evaluate(self, theta_new, t=None): + return super().evaluate(theta_new, t=t) * self._func.evaluate(theta_new) + + def evaluate_gradient(self, theta_new, t=None): + t1 = super().evaluate_gradient(theta_new, t=t) * self._func.evaluate(theta_new) + t2 = super().evaluate(theta_new, t=t) * self._func.evaluate_gradient(theta_new) + return t1 + t2 + + return AdjustedAcquisition diff --git a/elfi/methods/inference/bolfire.py b/elfi/methods/inference/bolfire.py index cef3d612..4ab51528 100644 --- a/elfi/methods/inference/bolfire.py +++ b/elfi/methods/inference/bolfire.py @@ -10,7 +10,7 @@ from elfi.loader import get_sub_seed from elfi.methods.bo.acquisition import LCBSC, AcquisitionBase from elfi.methods.bo.gpy_regression import GPyRegression -from elfi.methods.bo.utils import CostFunction +from elfi.methods.bo.utils import AdjustmentFunction, make_additive_acq from elfi.methods.classifier import Classifier, LogisticRegression from elfi.methods.inference.parameter_inference import ModelBased from elfi.methods.posteriors import BOLFIREPosterior @@ -333,14 +333,15 @@ def _resolve_target_model(self, target_model): def _resolve_acquisition_method(self, acquisition_method): """Resolve acquisition method.""" if acquisition_method is None: - # Model prior log-probabilities as an additive cost - cost = CostFunction(self.prior.logpdf, self.prior.gradient_logpdf, scale=-1) - return LCBSC(model=self.target_model, - prior=self.prior, - noise_var=self.acq_noise_var, - exploration_rate=self.exploration_rate, - seed=self.seed, - additive_cost=cost) + # LCBSC with prior log-probabilities as an additive term + prior_term = AdjustmentFunction(self.prior.logpdf, self.prior.gradient_logpdf, + scale=-1) + acq_method = make_additive_acq(LCBSC, prior_term) + return acq_method(model=self.target_model, + prior=self.prior, + noise_var=self.acq_noise_var, + exploration_rate=self.exploration_rate, + seed=self.seed) if isinstance(acquisition_method, AcquisitionBase): return acquisition_method raise TypeError('acquisition_method must be an instance of AcquisitionBase.') diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index ef8b9fd2..565bf8bf 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -6,7 +6,8 @@ import elfi from elfi.examples.ma2 import get_model -from elfi.methods.bo.utils import CostFunction, minimize, stochastic_optimization +from elfi.methods.bo.utils import (AdjustmentFunction, minimize, stochastic_optimization, + make_additive_acq, make_multiplicative_acq) from elfi.methods.density_ratio_estimation import DensityRatioEstimation from elfi.methods.utils import (GMDistribution, normalize_weights, numgrad, numpy_to_python_type, sample_object_to_dict, weighted_sample_quantile, weighted_var) @@ -283,20 +284,99 @@ def test_ratio_estimation(self): assert np.max(np.abs(test_w - test_w_estim)) < 0.1 assert np.abs(np.max(test_w) - densratio.max_ratio()) < 0.1 -class TestCostFunction: +class TestAdjustmentFunction: def test_evaluate(self): def fun(x): return x[0]**2 + (x[1] - 1)**4 - cost = CostFunction(elfi.tools.vectorize(fun), None, scale=10) + adjustment = AdjustmentFunction(elfi.tools.vectorize(fun), None, scale=10) x = np.array([0.5, 0.5]) - assert np.isclose(10 * fun(x), cost.evaluate(x)) + assert np.isclose(10 * fun(x), adjustment.evaluate(x)) def test_evaluate_gradient(self): def grad(x): return np.array([2 * x[0], 4 * (x[1] - 1)**3]) - cost = CostFunction(None, elfi.tools.vectorize(grad), scale=10) + adjustment = AdjustmentFunction(None, elfi.tools.vectorize(grad), scale=10) x = np.array([0.5, 0.5]) - assert np.allclose(10 * grad(x), cost.evaluate_gradient(x)) + assert np.allclose(10 * grad(x), adjustment.evaluate_gradient(x)) + + +class TestAdjustedAcquisition: + class CustomAcquisition(elfi.methods.bo.acquisition.AcquisitionBase): + def __init__(self, model, scale=1): + self.model = model + self.scale = scale + + def evaluate(self, x, t=None): + return self.scale * x * np.sin(x) + + def evaluate_gradient(self, x, t=None): + return self.scale * (np.sin(x) + x * np.cos(x)) + + def minimize(self): + return minimize(self.evaluate, ((0, 10), ), grad=self.evaluate_gradient) + + def get_adjustment(self, scale): + def func(x): + return scale * x + + def grad(x): + return scale * np.ones_like(x) + + return AdjustmentFunction(func, grad) + + def test_evaluate_additive(self): + acq = self.CustomAcquisition(None, scale=-1) + adj = self.get_adjustment(0.5) + adjusted_acq = make_additive_acq(self.CustomAcquisition, adj)(None, scale=-1) + x = np.arange(10).reshape(10, 1) + y = acq.evaluate(x) + adj.evaluate(x) + test_y = adjusted_acq.evaluate(x) + assert np.allclose(test_y, y) + + def test_evaluate_gradient_additive(self): + acq = self.CustomAcquisition(None, scale=-1) + adj = self.get_adjustment(0.5) + adjusted_acq = make_additive_acq(self.CustomAcquisition, adj)(None, scale=-1) + x = np.arange(10).reshape(10, 1) + y = acq.evaluate_gradient(x) + adj.evaluate_gradient(x) + test_y = adjusted_acq.evaluate_gradient(x) + assert np.allclose(test_y, y) + + def test_minimize_additive(self): + acq = self.CustomAcquisition(None, scale=-1) + loc1, val1 = acq.minimize() + adj = self.get_adjustment(0.5) + adjusted_acq = make_additive_acq(self.CustomAcquisition, adj)(None, scale=-1) + loc2, val2 = adjusted_acq.minimize() + assert loc1 > loc2 + assert val1 < val2 + + def test_evaluate_multiplicative(self): + acq = self.CustomAcquisition(None, scale=-1) + adj = self.get_adjustment(0.1) + adjusted_acq = make_multiplicative_acq(self.CustomAcquisition, adj)(None, scale=-1) + x = np.arange(10).reshape(10, 1) + y = acq.evaluate(x) * adj.evaluate(x) + test_y = adjusted_acq.evaluate(x) + assert np.allclose(test_y, y) + + def test_evaluate_gradient_multiplicative(self): + acq = self.CustomAcquisition(None, scale=-1) + adj = self.get_adjustment(0.1) + adjusted_acq = make_multiplicative_acq(self.CustomAcquisition, adj)(None, scale=-1) + x = np.arange(10).reshape(10, 1) + y = acq.evaluate(x) * adj.evaluate_gradient(x) + acq.evaluate_gradient(x) * adj.evaluate(x) + test_y = adjusted_acq.evaluate_gradient(x) + assert np.allclose(test_y, y) + + def test_minimize_multiplicative(self): + acq = self.CustomAcquisition(None, scale=-1) + loc1, val1 = acq.minimize() + adj = self.get_adjustment(0.1) + adjusted_acq = make_multiplicative_acq(self.CustomAcquisition, adj)(None, scale=-1) + loc2, val2 = adjusted_acq.minimize() + assert loc1 < loc2 + assert val1 < val2