Update acquisition cost implementation

CHANGELOG.rst

@@ -1,6 +1,7 @@
 Changelog
 =========
 
+- Add option to use additive or multiplicative adjustment in any acquisition method
 - 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
 - Drop official Python support for 3.7 and 3.8 as GPy is not officially supported for these versions

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,16 +223,14 @@ class LCBSC(AcquisitionBase):
 
     """
 
-    def __init__(self, *args, delta=None, additive_cost=None, **kwargs):
+    def __init__(self, *args, delta=None, **kwargs):
         """Initialize LCBSC.
 
         Parameters
         ----------
         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
 
 

elfi/methods/bo/utils.py

@@ -111,28 +111,28 @@ 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
         self.gradient = gradient
         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
+    -------
+    AdditiveAcquisition
+
+    """
+    class AdditiveAcquisition(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 AdditiveAcquisition
+
+
+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
+    -------
+    MultiplicativeAcquisition
+
+    """
+    class MultiplicativeAcquisition(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 MultiplicativeAcquisition

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.')

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