Allow nonlinear constraints #353

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Features
 
+- [#353](https://github.com/pybop-team/PyBOP/issues/353) - Allow user-defined check_params functions to enforce nonlinear constraints, and enable SciPy constrained optimisation methods
 - [#222](https://github.com/pybop-team/PyBOP/issues/222) - Adds an example for performing and electrode balancing.
 - [#441](https://github.com/pybop-team/PyBOP/issues/441) - Adds an example for estimating constants within a `pybamm.FunctionalParameter`.
 - [#405](https://github.com/pybop-team/PyBOP/pull/405) - Adds frequency-domain based EIS prediction methods via `model.simulateEIS` and updates to `problem.evaluate` with examples and tests.

examples/scripts/ecm_tau_constraints.py

@@ -0,0 +1,192 @@
+import numpy as np
+
+import pybop
+
+"""
+When fitting empirical models, the parameters we are able to identify
+will be constrained from the data that's available. For example, it's
+no good trying to fit an RC timescale of 0.1 s from data sampled at
+1 Hz! Likewise, an RC timescale of 100 s cannot be meaningfully fitted
+to just 10 s of data. To ensure the optimiser doesn't propose
+excessively long or short timescales - beyond what can reasonably be
+inferred from the data - it is common to apply nonlinear constraints
+on the parameter space. This script fits an RC pair with the
+constraint 0.5 <= R1 * C1 <= 1, to highlight a possible method for
+applying constraints on the timescales.
+
+An alternative approach is given i the ecm_trust-constr notebook,
+which can lead to better results and higher optimisation efficiency
+when good timescale guesses are available.
+"""
+
+# Import the ECM parameter set from JSON
+# parameter_set = pybop.ParameterSet(
+#     json_path="examples/scripts/parameters/initial_ecm_parameters.json"
+# )
+# parameter_set.import_parameters()
+
+
+# Alternatively, define the initial parameter set with a dictionary
+# Add definitions for R's, C's, and initial overpotentials for any additional RC elements
+parameter_set = {
+    "chemistry": "ecm",
+    "Initial SoC": 0.5,
+    "Initial temperature [K]": 25 + 273.15,
+    "Cell capacity [A.h]": 5,
+    "Nominal cell capacity [A.h]": 5,
+    "Ambient temperature [K]": 25 + 273.15,
+    "Current function [A]": 5,
+    "Upper voltage cut-off [V]": 4.2,
+    "Lower voltage cut-off [V]": 3.0,
+    "Cell thermal mass [J/K]": 1000,
+    "Cell-jig heat transfer coefficient [W/K]": 10,
+    "Jig thermal mass [J/K]": 500,
+    "Jig-air heat transfer coefficient [W/K]": 10,
+    "Open-circuit voltage [V]": pybop.empirical.Thevenin().default_parameter_values[
+        "Open-circuit voltage [V]"
+    ],
+    "R0 [Ohm]": 0.001,
+    "Element-1 initial overpotential [V]": 0,
+    "Element-2 initial overpotential [V]": 0,
+    "R1 [Ohm]": 0.0002,
+    "R2 [Ohm]": 0.0003,
+    "C1 [F]": 10000,
+    "C2 [F]": 5000,
+    "Entropic change [V/K]": 0.0004,
+}
+
+
+def get_parameter_checker(
+    tau_mins: float | list[float],
+    tau_maxs: float | list[float],
+    fitted_rc_pair_indices: int | list[int],
+):
+    """Returns a function to check parameters against given tau bounds.
+    The resulting check_params function will be sent off to PyBOP; the
+    rest of the code does some light checking of the constraints.
+
+    Parameters
+    ----------
+    tau_mins: float | list[float]
+        Lower bounds on timescale tau_i = Ri * Ci
+    tau_maxs: float | list[float]
+        Upper bounds on timescale tau_i = Ri * Ci
+    fitted_rc_pair_indices: int | list[float]
+        The index of each RC pair whose parameters are to be fitted.
+        Eg. [1, 2] means fitting R1, R2, C1, C2. The timescale of RC
+        pair fitted_rc_pair_indices[j] is constrained to be in the
+        range tau_mins[j] <= R * C <= tau_maxs[j]
+
+    Returns
+    -------
+    check_params
+        Function to check the proposed parameter values match the
+        requested constraints
+
+    """
+
+    # Ensure inputs are lists
+    tau_mins = [tau_mins] if not isinstance(tau_mins, list) else tau_mins
+    tau_maxs = [tau_maxs] if not isinstance(tau_maxs, list) else tau_maxs
+    fitted_rc_pair_indices = (
+        [fitted_rc_pair_indices]
+        if not isinstance(fitted_rc_pair_indices, list)
+        else fitted_rc_pair_indices
+    )
+
+    # Validate input lengths
+    if len(tau_mins) != len(fitted_rc_pair_indices) or len(tau_maxs) != len(
+        fitted_rc_pair_indices
+    ):
+        raise ValueError(
+            "tau_mins and tau_maxs must have the same length as fitted_rc_pair_indices"
+        )
+
+    def check_params(
+        inputs: dict[str, float] = None,
+        parameter_set=None,
+        allow_infeasible_solutions: bool = False,
+    ) -> bool:
+        """Checks if the given inputs are within the tau bounds."""
+        # Allow simulation to run if inputs are None
+        if inputs is None or inputs == {}:
+            return True
+
+        # Check every respective R*C against tau bounds
+        print(inputs)
+        for i, tau_min, tau_max in zip(fitted_rc_pair_indices, tau_mins, tau_maxs):
+            tau = inputs[f"R{i} [Ohm]"] * inputs[f"C{i} [F]"]
+            if not tau_min <= tau <= tau_max:
+                return False
+        return True
+
+    return check_params
+
+
+# Define the model
+params = pybop.ParameterSet(params_dict=parameter_set)
+model = pybop.empirical.Thevenin(
+    parameter_set=params,
+    check_params=get_parameter_checker(
+        0, 0.5, 1
+    ),  # Set the model up to automatically check parameters
+    options={"number of rc elements": 2},
+)
+
+# Fitting parameters
+parameters = pybop.Parameters(
+    pybop.Parameter(
+        "R0 [Ohm]",
+        prior=pybop.Gaussian(0.0002, 0.0001),
+        bounds=[1e-4, 1e-2],
+    ),
+    pybop.Parameter(
+        "R1 [Ohm]",
+        prior=pybop.Gaussian(0.0001, 0.0001),
+        bounds=[1e-5, 1e-2],
+    ),
+    pybop.Parameter(
+        "C1 [F]",
+        prior=pybop.Gaussian(10000, 2500),
+        bounds=[2500, 5e4],
+    ),
+)
+
+sigma = 0.001
+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))
+
+# Form dataset
+dataset = pybop.Dataset(
+    {
+        "Time [s]": t_eval,
+        "Current function [A]": values["Current [A]"].data,
+        "Voltage [V]": corrupt_values,
+    }
+)
+
+# Generate problem, cost function, and optimisation class
+problem = pybop.FittingProblem(model, parameters, dataset)
+cost = pybop.SumSquaredError(problem)
+optim = pybop.XNES(
+    cost,
+    allow_infeasible_solutions=False,
+    max_iterations=100,
+)
+
+x, final_cost = optim.run()
+print("Estimated parameters:", x)
+
+
+# Plot the time series
+pybop.plot_dataset(dataset)
+
+# Plot the timeseries output
+pybop.quick_plot(problem, problem_inputs=x, title="Optimised Comparison")
+
+# Plot convergence
+pybop.plot_convergence(optim)
+
+# Plot the parameter traces
+pybop.plot_parameters(optim)

pybop/models/base_model.py

@@ -1,6 +1,6 @@
 import copy
 from dataclasses import dataclass
-from typing import Optional, Union
+from typing import Callable, Optional, Union
 
 import casadi
 import numpy as np
@@ -70,6 +70,7 @@ def __init__(
         self,
         name: str = "Base Model",
         parameter_set: Optional[ParameterSet] = None,
+        check_params: Callable = None,
         eis=False,
     ):
         """
@@ -82,6 +83,15 @@ def __init__(
         parameter_set : pybop.ParameterSet, optional
             A PyBOP ParameterSet, PyBaMM ParameterValues object or a dictionary containing the
             parameter values.
+        check_params : Callable, optional
+            A compatibility check for the model parameters. Function, with
+            signature
+                check_params(
+                    inputs: dict,
+                    allow_infeasible_solutions: bool, optional
+                )
+            Returns true if parameters are valid, False otherwise. Can be
+            used to impose constraints on valid parameters.
 
         Additional Attributes
         ---------------------
@@ -105,6 +115,7 @@ def __init__(
             self._parameter_set = parameter_set.copy()
         else:  # a pybop parameter set
             self._parameter_set = pybamm.ParameterValues(parameter_set.params).copy()
+        self.param_checker = check_params
 
         self.pybamm_model = None
         self.parameters = Parameters()
@@ -798,6 +809,8 @@ def _check_params(
         bool
             A boolean which signifies whether the parameters are compatible.
         """
+        if self.param_checker:
+            return self.param_checker(inputs, allow_infeasible_solutions)
         return True
 
     def copy(self):

pybop/models/empirical/base_ecm.py

@@ -45,6 +45,7 @@ def __init__(
         var_pts=None,
         spatial_methods=None,
         solver=None,
+        check_params=None,
         eis=False,
         **model_kwargs,
     ):
@@ -65,7 +66,9 @@ def __init__(
                 print("Setting open-circuit voltage to default function")
                 parameter_set["Open-circuit voltage [V]"] = default_ocp
 
-        super().__init__(name=name, parameter_set=parameter_set, eis=eis)
+        super().__init__(
+            name=name, parameter_set=parameter_set, check_params=check_params, eis=eis
+        )
         self.pybamm_model = pybamm_model
         self._unprocessed_model = self.pybamm_model
 
@@ -114,6 +117,8 @@ def _check_params(
         bool
             A boolean which signifies whether the parameters are compatible.
         """
+        if self.param_checker:
+            return self.param_checker(inputs, allow_infeasible_solutions)
         return True
 
     def get_initial_state(

pybop/optimisers/scipy_optimisers.py

@@ -1,5 +1,8 @@
+import warnings
+from typing import Union
+
 import numpy as np
-from scipy.optimize import OptimizeResult, differential_evolution, minimize
+from scipy.optimize import Bounds, OptimizeResult, differential_evolution, minimize
 
 from pybop import BaseOptimiser, Result
 
@@ -52,12 +55,18 @@ def _sanitise_inputs(self):
 
         # Convert bounds to SciPy format
         if isinstance(self.bounds, dict):
-            self._scipy_bounds = [
-                (lower, upper)
-                for lower, upper in zip(self.bounds["lower"], self.bounds["upper"])
-            ]
-        else:
+            self._scipy_bounds = Bounds(
+                self.bounds["lower"], self.bounds["upper"], True
+            )
+        elif isinstance(self.bounds, list):
+            lb, ub = zip(*self.bounds)
+            self._scipy_bounds = Bounds(lb, ub, True)
+        elif isinstance(self.bounds, Bounds) or self.bounds is None:
             self._scipy_bounds = self.bounds
+        else:
+            raise TypeError(
+                "Bounds provided must be either type dict, list or SciPy.optimize.bounds object."
+            )
 
     def _run(self):
         """
@@ -70,10 +79,15 @@ def _run(self):
         """
         result = self._run_optimiser()
 
+        try:
+            nit = result.nit
+        except AttributeError:
+            nit = -1
+
         return Result(
             x=result.x,
             final_cost=self.cost(result.x),
-            n_iterations=result.nit,
+            n_iterations=nit,
             scipy_result=result,
         )
 
@@ -107,6 +121,7 @@ def __init__(self, cost, **optimiser_kwargs):
         optimiser_options = dict(method="Nelder-Mead", jac=False)
         optimiser_options.update(**optimiser_kwargs)
         super().__init__(cost, **optimiser_options)
+        self._cost0 = 1.0
 
     def _set_up_optimiser(self):
         """
@@ -169,17 +184,45 @@ def _run_optimiser(self):
         self.inf_count = 0
 
         # Add callback storing history of parameter values
-        def callback(intermediate_result: OptimizeResult):
-            self.log["x_best"].append(intermediate_result.x)
+
+        def base_callback(intermediate_result: Union[OptimizeResult, np.ndarray]):
+            """
+            Log intermediate optimisation solutions. Depending on the
+            optimisation algorithm, intermediate_result may be either
+            a OptimizeResult or an array of parameter values, with a
+            try/except ensuring both cases are handled correctly.
+            """
+            if isinstance(intermediate_result, OptimizeResult):
+                x_best = intermediate_result.x
+                cost = intermediate_result.fun
+            else:
+                x_best = intermediate_result
+                cost = self.cost(x_best)
+
+            self.log["x_best"].append(x_best)
             self.log["cost"].append(
-                intermediate_result.fun if self.minimising else -intermediate_result.fun
+                (-1 if not self.minimising else 1) * cost * self._cost0
             )
 
+        callback = (
+            base_callback
+            if self._options["method"] != "trust-constr"
+            else lambda x, intermediate_result: base_callback(intermediate_result)
+        )
+
         # Compute the absolute initial cost and resample if required
         self._cost0 = np.abs(self.cost(self.x0))
         if np.isinf(self._cost0):
             for _i in range(1, self.num_resamples):
-                self.x0 = self.parameters.rvs()
+                try:
+                    self.x0 = self.parameters.rvs()
+                except AttributeError:
+                    warnings.warn(
+                        "Parameter does not have a prior distribution. Stopping resampling.",
+                        UserWarning,
+                        stacklevel=2,
+                    )
+                    break
                 self._cost0 = np.abs(self.cost(self.x0))
                 if not np.isinf(self._cost0):
                     break
@@ -249,9 +292,8 @@ def _set_up_optimiser(self):
         if self._scipy_bounds is None:
             raise ValueError("Bounds must be specified for differential_evolution.")
         else:
-            if not all(
-                np.isfinite(value) for pair in self._scipy_bounds for value in pair
-            ):
+            bnds = self._scipy_bounds
+            if not (np.isfinite(bnds.lb).all() and np.isfinite(bnds.ub).all()):
                 raise ValueError("Bounds must be specified for differential_evolution.")
 
         # Apply default maxiter and tolerance