Refactors BaseCost.evaluate/S1 into BaseCost.compute

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Features
 
+- [#436](https://github.com/pybop-team/PyBOP/pull/436) - Refactors `BaseCost.evaluate/S1` & `BaseCost._evaluate/S1` into `BaseCost.__call__` & `BaseCost.compute`. `BaseCost.compute` directly acts on the predictions, while `BaseCost.__call__` calls `BaseProblem.evaluate/S1` before `BaseCost.compute`. Fully separates `__call__` and `compute` for FittingCosts, Likelihoods. Bugfixes to GaussianLogLikelihood calculation.
 - [#450](https://github.com/pybop-team/PyBOP/pull/450) - Adds support for IDAKLU with output variables, and corresponding examples, tests.
 - [#364](https://github.com/pybop-team/PyBOP/pull/364) - Adds the MultiFittingProblem class and the multi_fitting example script.
 - [#444](https://github.com/pybop-team/PyBOP/issues/444) - Merge `BaseModel` `build()` and `rebuild()` functionality.
@@ -20,6 +21,7 @@
 
 ## Breaking Changes
 
+- [#436](https://github.com/pybop-team/PyBOP/pull/436) - **API Change:** The functionality from `BaseCost.evaluate/S1` & `BaseCost._evaluate/S1` is represented in `BaseCost.__call__` & `BaseCost.compute`. `BaseCost.compute` directly acts on the predictions, while `BaseCost.__call__` calls `BaseProblem.evaluate/S1` before `BaseCost.compute`. `compute` has optional args for gradient cost calculations.
 - [#424](https://github.com/pybop-team/PyBOP/issues/424) - Replaces the `init_soc` input to `FittingProblem` with the option to pass an initial OCV value, updates `BaseModel` and fixes `multi_model_identification.ipynb` and `spm_electrode_design.ipynb`.
 
 # [v24.6.1](https://github.com/pybop-team/PyBOP/tree/v24.6.1) - 2024-07-31

examples/scripts/spm_MLE.py

@@ -2,8 +2,14 @@
 
 import pybop
 
-# Define model
+# Define model and set initial parameter values
 parameter_set = pybop.ParameterSet.pybamm("Chen2020")
+parameter_set.update(
+    {
+        "Negative electrode active material volume fraction": 0.63,
+        "Positive electrode active material volume fraction": 0.51,
+    }
+)
 model = pybop.lithium_ion.SPM(parameter_set=parameter_set)
 
 # Fitting parameters
@@ -19,31 +25,39 @@
     ),
 )
 
-# Set initial parameter values
-parameter_set.update(
-    {
-        "Negative electrode active material volume fraction": 0.63,
-        "Positive electrode active material volume fraction": 0.51,
-    }
-)
 # Generate data
-sigma = 0.005
-t_eval = np.arange(0, 900, 3)
-values = model.predict(t_eval=t_eval)
-corrupt_values = values["Voltage [V]"].data + np.random.normal(0, sigma, len(t_eval))
+sigma = 0.002
+experiment = pybop.Experiment(
+    [
+        (
+            "Discharge at 0.5C for 3 minutes (3 second period)",
+            "Charge at 0.5C for 3 minutes (3 second period)",
+        ),
+    ]
+)
+values = model.predict(initial_state={"Initial SoC": 0.5}, experiment=experiment)
+
+
+def noise(sigma):
+    return np.random.normal(0, sigma, len(values["Voltage [V]"].data))
+
 
 # Form dataset
 dataset = pybop.Dataset(
     {
-        "Time [s]": t_eval,
+        "Time [s]": values["Time [s]"].data,
         "Current function [A]": values["Current [A]"].data,
-        "Voltage [V]": corrupt_values,
+        "Voltage [V]": values["Voltage [V]"].data + noise(sigma),
+        "Bulk open-circuit voltage [V]": values["Bulk open-circuit voltage [V]"].data
+        + noise(sigma),
     }
 )
 
+
+signal = ["Voltage [V]", "Bulk open-circuit voltage [V]"]
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(model, parameters, dataset)
-likelihood = pybop.GaussianLogLikelihoodKnownSigma(problem, sigma0=sigma)
+problem = pybop.FittingProblem(model, parameters, dataset, signal=signal)
+likelihood = pybop.GaussianLogLikelihood(problem, sigma0=sigma * 4)
 optim = pybop.IRPropMin(
     likelihood,
     max_unchanged_iterations=20,

examples/standalone/cost.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 import pybop
 
 
@@ -43,7 +45,9 @@ def __init__(self, problem=None):
         )
         self.x0 = self.parameters.initial_value()
 
-    def compute(self, inputs):
+    def compute(
+        self, y: dict, dy: np.ndarray = None, inputs=None, calculate_grad: bool = False
+    ):
         """
         Compute the cost for a given parameter value.
 

pybop/costs/_likelihoods.py

@@ -39,46 +39,34 @@ def __init__(self, problem: BaseProblem, sigma0: Union[list[float], float]):
         self._offset = -0.5 * self.n_time_data * np.log(2 * np.pi * self.sigma2)
         self._multip = -1 / (2.0 * self.sigma2)
 
-    def compute(self, inputs: Inputs) -> float:
+    def compute(
+        self,
+        y: dict,
+        dy: np.ndarray = None,
+        inputs: Inputs = None,
+        calculate_grad: bool = False,
+    ) -> Union[float, tuple[float, np.ndarray]]:
         """
         Compute the Gaussian log-likelihood for the given parameters with known sigma.
 
         This method only computes the likelihood, without calling the problem.evaluateS1.
         """
+        # Verify we have dy if calculate_grad is True
+        self.verify_args(dy, calculate_grad)
 
-        if not self.verify_prediction(self.y):
-            return -np.inf
-
-        e = np.asarray(
-            [
-                np.sum(
-                    self._offset
-                    + self._multip
-                    * np.sum((self._target[signal] - self.y[signal]) ** 2.0)
-                )
-                for signal in self.signal
-            ]
-        )
-
-        return e.item() if self.n_outputs == 1 else np.sum(e)
+        # Early return if the prediction is not verified
+        if not self.verify_prediction(y):
+            return (-np.inf, -self.grad_fail) if calculate_grad else -np.inf
 
-    def computeS1(self, inputs: Inputs) -> tuple[float, np.ndarray]:
-        """
-        Compute the Gaussian log-likelihood and it's gradient for the given parameters with known sigma.
+        # Calculate residuals and error
+        r = np.asarray([self._target[signal] - y[signal] for signal in self.signal])
+        e = np.sum(self._offset + self._multip * np.sum(r**2.0))
 
-        This method only computes the likelihood, without calling the problem.evaluateS1.
-        """
-        if not self.verify_prediction(self.y):
-            return -np.inf, -self._de * np.ones(self.n_parameters)
+        if calculate_grad:
+            dl = np.sum((np.sum((r * dy.T), axis=2) / self.sigma2), axis=1)
+            return e, dl
 
-        likelihood = self.compute(inputs)
-
-        r = np.asarray(
-            [self._target[signal] - self.y[signal] for signal in self.signal]
-        )
-        dl = np.sum((np.sum((r * self.dy.T), axis=2) / self.sigma2), axis=1)
-
-        return likelihood, dl
+        return e
 
     def check_sigma0(self, sigma0: Union[np.ndarray, float]):
         """
@@ -121,7 +109,6 @@ def __init__(
         super().__init__(problem)
         self._dsigma_scale = dsigma_scale
         self._logpi = -0.5 * self.n_time_data * np.log(2 * np.pi)
-        self._has_separable_problem = False
 
         self.sigma = Parameters()
         self._add_sigma_parameters(sigma0)
@@ -152,6 +139,7 @@ def _add_single_sigma(self, index, value):
                     f"Sigma for output {index+1}",
                     initial_value=value,
                     prior=Uniform(0.5 * value, 1.5 * value),
+                    bounds=[1e-8, 3 * value],
                 )
             )
         else:
@@ -173,7 +161,13 @@ def dsigma_scale(self, new_value):
             raise ValueError("dsigma_scale must be non-negative")
         self._dsigma_scale = new_value
 
-    def compute(self, inputs: Inputs) -> float:
+    def compute(
+        self,
+        y: dict,
+        dy: np.ndarray = None,
+        inputs: Inputs = None,
+        calculate_grad: bool = False,
+    ) -> Union[float, tuple[float, np.ndarray]]:
         """
         Compute the Gaussian log-likelihood for the given parameters.
 
@@ -190,68 +184,30 @@ def compute(self, inputs: Inputs) -> float:
         float
             The log-likelihood value, or -inf if the standard deviations are non-positive.
         """
-        self.parameters.update(values=list(inputs.values()))
-
-        sigma = self.sigma.current_value()
-        if np.any(sigma <= 0):
-            return -np.inf
-
-        self.y = self.problem.evaluate(self.problem.parameters.as_dict())
-        if not self.verify_prediction(self.y):
-            return -np.inf
-
-        e = np.asarray(
-            [
-                np.sum(
-                    self._logpi
-                    - self.n_time_data * np.log(sigma)
-                    - np.sum((self._target[signal] - self.y[signal]) ** 2.0)
-                    / (2.0 * sigma**2.0)
-                )
-                for signal in self.signal
-            ]
-        )
-
-        return e.item() if self.n_outputs == 1 else np.sum(e)
-
-    def computeS1(self, inputs: Inputs) -> tuple[float, np.ndarray]:
-        """
-        Compute the Gaussian log-likelihood and it's gradient for the given parameters.
-
-        This method only computes the likelihood, without calling the problem.evaluateS1.
-
-        Parameters
-        ----------
-        inputs : Inputs
-            The parameters for which to evaluate the log-likelihood.
-
-        Returns
-        -------
-        Tuple[float, np.ndarray]
-            The log-likelihood and its gradient.
-        """
-        self.parameters.update(values=list(inputs.values()))
-
+        # Verify we have dy if calculate_grad is True
+        self.verify_args(dy, calculate_grad)
         sigma = self.sigma.current_value()
-        if np.any(sigma <= 0):
-            return -np.inf, -self._de * np.ones(self.n_parameters)
-
-        self.y, self.dy = self.problem.evaluateS1(self.problem.parameters.as_dict())
-        if not self.verify_prediction(self.y):
-            return -np.inf, -self._de * np.ones(self.n_parameters)
 
-        likelihood = self.compute(inputs)
+        if not self.verify_prediction(y):
+            return (-np.inf, -self.grad_fail) if calculate_grad else -np.inf
 
-        r = np.asarray(
-            [self._target[signal] - self.y[signal] for signal in self.signal]
+        # Calculate residuals and error
+        r = np.asarray([self._target[signal] - y[signal] for signal in self.signal])
+        e = np.sum(
+            self._logpi
+            - self.n_time_data * np.log(sigma)
+            - np.sum(r**2.0, axis=1) / (2.0 * sigma**2.0)
         )
-        dl = np.sum((np.sum((r * self.dy.T), axis=2) / (sigma**2.0)), axis=1)
-        dsigma = (
-            -self.n_time_data / sigma + np.sum(r**2.0, axis=1) / (sigma**3.0)
-        ) / self._dsigma_scale
-        dl = np.concatenate((dl.flatten(), dsigma))
 
-        return likelihood, dl
+        if calculate_grad:
+            dl = np.sum((np.sum((r * dy.T), axis=2) / (sigma**2.0)), axis=1)
+            dsigma = (
+                -self.n_time_data / sigma + np.sum(r**2.0, axis=1) / (sigma**3.0)
+            ) / self._dsigma_scale
+            dl = np.concatenate((dl.flatten(), dsigma))
+            return e, dl
+
+        return e
 
 
 class MAP(BaseLikelihood):
@@ -290,7 +246,13 @@ def __init__(self, problem, likelihood, sigma0=None, gradient_step=1e-3):
         self.parameters = self.likelihood.parameters
         self._has_separable_problem = self.likelihood.has_separable_problem
 
-    def compute(self, inputs: Inputs) -> float:
+    def compute(
+        self,
+        y: dict,
+        dy: np.ndarray = None,
+        inputs: Inputs = None,
+        calculate_grad: bool = False,
+    ) -> Union[float, tuple[float, np.ndarray]]:
         """
         Compute the Maximum a Posteriori for the given parameters.
 
@@ -308,70 +270,41 @@ def compute(self, inputs: Inputs) -> float:
         float
             The maximum a posteriori cost.
         """
+        # Verify we have dy if calculate_grad is True
+        self.verify_args(dy, calculate_grad)
+
         log_prior = sum(
             self.parameters[key].prior.logpdf(value) for key, value in inputs.items()
         )
 
         if not np.isfinite(log_prior).any():
-            return -np.inf
-
-        if self._has_separable_problem:
-            self.likelihood.y = self.y
-        log_likelihood = self.likelihood.evaluate(inputs)
+            return (-np.inf, -self.grad_fail) if calculate_grad else -np.inf
 
-        posterior = log_likelihood + log_prior
-        return posterior
-
-    def computeS1(self, inputs: Inputs) -> tuple[float, np.ndarray]:
-        """
-        Compute the Maximum a Posteriori, and it's gradient for the given parameters.
-
-        If self._has_separable_problem is True, then this method only computes the
-        likelihood, without calling the problem.evaluateS1(). Else, problem.evaluateS1()
-        is called before computing the likelihood.
-
-        Parameters
-        ----------
-        inputs : Inputs
-            The parameters for which to compute the cost and gradient.
+        if calculate_grad:
+            log_likelihood, dl = self.likelihood.compute(
+                y, dy, inputs, calculate_grad=True
+            )
 
-        Returns
-        -------
-        tuple
-            A tuple containing the cost and the gradient. The cost is a float,
-            and the gradient is an array-like of the same length as `x`.
-
-        Raises
-        ------
-        ValueError
-            If an error occurs during the calculation of the cost or gradient.
-        """
-        log_prior = sum(
-            self.parameters[key].prior.logpdf(value) for key, value in inputs.items()
-        )
-        if not np.isfinite(log_prior).any():
-            return -np.inf, -self._de * np.ones(self.n_parameters)
+            # Compute a finite difference approximation of the gradient of the log prior
+            delta = self.parameters.initial_value() * self.gradient_step
+            prior_gradient = []
 
-        if self._has_separable_problem:
-            self.likelihood.y, self.likelihood.dy = (self.y, self.dy)
-        log_likelihood, dl = self.likelihood.evaluateS1(inputs)
+            for parameter, step_size in zip(self.problem.parameters, delta):
+                param_value = inputs[parameter.name]
 
-        # Compute a finite difference approximation of the gradient of the log prior
-        delta = self.parameters.initial_value() * self.gradient_step
-        prior_gradient = []
+                log_prior_upper = parameter.prior.logpdf(param_value * (1 + step_size))
+                log_prior_lower = parameter.prior.logpdf(param_value * (1 - step_size))
 
-        for parameter, step_size in zip(self.problem.parameters, delta):
-            param_value = inputs[parameter.name]
+                gradient = (log_prior_upper - log_prior_lower) / (
+                    2 * step_size * param_value + np.finfo(float).eps
+                )
+                prior_gradient.append(gradient)
 
-            log_prior_upper = parameter.prior.logpdf(param_value * (1 + step_size))
-            log_prior_lower = parameter.prior.logpdf(param_value * (1 - step_size))
+            posterior = log_likelihood + log_prior
+            total_gradient = dl + prior_gradient
 
-            gradient = (log_prior_upper - log_prior_lower) / (
-                2 * step_size * param_value + np.finfo(float).eps
-            )
-            prior_gradient.append(gradient)
+            return posterior, total_gradient
 
+        log_likelihood = self.likelihood.compute(y, inputs)
         posterior = log_likelihood + log_prior
-        total_gradient = dl + prior_gradient
-
-        return posterior, total_gradient
+        return posterior