Split Problem into Base, Fitting and Design

CHANGELOG.md

@@ -1,9 +1,10 @@
 # [Unreleased](https://github.com/pybop-team/PyBOP)
 
 - [#116](https://github.com/pybop-team/PyBOP/issues/116) - Adds PSO, SNES, XNES, ADAM, and IPropMin optimisers to PintsOptimisers() class
+- [#38](https://github.com/pybop-team/PyBOP/issues/38) - Restructures the Problem classes ahead of adding a design optimisation example
 
 # [v23.11](https://github.com/pybop-team/PyBOP/releases/tag/v23.11)
 - Initial release
 - Adds Pints, NLOpt, and SciPy optimisers
 - Adds SumofSquareError and RootMeanSquareError cost functions
-- Adds Parameter and dataset classes
+- Adds Parameter and Dataset classes

examples/notebooks/spm_nlopt.ipynb

@@ -334,7 +334,7 @@
       "source": [
         "# Define the cost to optimise\n",
         "signal = \"Terminal voltage [V]\"\n",
-        "problem = pybop.Problem(pyb_model, parameters, dataset, signal=signal)\n",
+        "problem = pybop.Problem(parameters, dataset, model=pyb_model, signal=signal)\n",
         "cost = pybop.RootMeanSquaredError(problem)"
       ]
     },

examples/scripts/spm_CMAES.py

@@ -33,7 +33,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.Problem(model, parameters, dataset)
+problem = pybop.FittingProblem(model, parameters, dataset)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.CMAES)
 optim.set_max_iterations(100)

examples/scripts/spm_IRPropMin.py

@@ -33,7 +33,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.Problem(model, parameters, dataset)
+problem = pybop.FittingProblem(model, parameters, dataset)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.IRPropMin)
 optim.set_max_iterations(100)

examples/scripts/spm_SNES.py

@@ -33,7 +33,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.Problem(model, parameters, dataset)
+problem = pybop.FittingProblem(model, parameters, dataset)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.SNES)
 optim.set_max_iterations(100)

examples/scripts/spm_XNES.py

@@ -33,7 +33,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.Problem(model, parameters, dataset)
+problem = pybop.FittingProblem(model, parameters, dataset)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.XNES)
 optim.set_max_iterations(100)

examples/scripts/spm_adam.py

@@ -36,7 +36,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.Problem(model, parameters, dataset)
+problem = pybop.FittingProblem(model, parameters, dataset)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.Adam)
 optim.set_max_iterations(100)

examples/scripts/spm_descent.py

@@ -36,7 +36,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.Problem(model, parameters, dataset)
+problem = pybop.FittingProblem(model, parameters, dataset)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.GradientDescent)
 optim.optimiser.set_learning_rate(0.025)

examples/scripts/spm_nlopt.py

@@ -32,7 +32,7 @@
 
 # Define the cost to optimise
 signal = "Terminal voltage [V]"
-problem = pybop.Problem(model, parameters, dataset, signal=signal, init_soc=0.98)
+problem = pybop.FittingProblem(model, parameters, dataset, signal=signal, init_soc=0.98)
 cost = pybop.RootMeanSquaredError(problem)
 
 # Build the optimisation problem

examples/scripts/spm_pso.py

@@ -33,7 +33,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.Problem(model, parameters, dataset)
+problem = pybop.FittingProblem(model, parameters, dataset)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.PSO)
 optim.set_max_iterations(100)

pybop/__init__.py

@@ -70,7 +70,7 @@
 #
 # Problem class
 #
-from ._problem import Problem
+from ._problem import FittingProblem, DesignProblem
 
 #
 # Plotting class

pybop/_problem.py

@@ -1,30 +1,78 @@
 import numpy as np
 
 
-class Problem:
+class BaseProblem:
     """
-    Defines a PyBOP single output problem, follows the PINTS interface.
+    Defines the PyBOP base problem, following the PINTS interface.
     """
 
     def __init__(
         self,
-        model,
         parameters,
-        dataset,
-        signal="Terminal voltage [V]",
+        model=None,
         check_model=True,
         init_soc=None,
         x0=None,
     ):
         self._model = model
-        self.parameters = parameters
-        self.signal = signal
-        self._model.signal = self.signal
-        self._dataset = {o.name: o for o in dataset}
         self.check_model = check_model
+        self.parameters = parameters
         self.init_soc = init_soc
         self.x0 = x0
         self.n_parameters = len(self.parameters)
+
+        # Set bounds
+        self.bounds = dict(
+            lower=[param.bounds[0] for param in self.parameters],
+            upper=[param.bounds[1] for param in self.parameters],
+        )
+
+        # Sample from prior for x0
+        if x0 is None:
+            self.x0 = np.zeros(self.n_parameters)
+            for i, param in enumerate(self.parameters):
+                self.x0[i] = param.rvs(1)
+        elif len(x0) != self.n_parameters:
+            raise ValueError("x0 dimensions do not match number of parameters")
+
+        # Add the initial values to the parameter definitions
+        for i, param in enumerate(self.parameters):
+            param.update(value=self.x0[i])
+
+    def evaluate(self, parameters):
+        """
+        Evaluate the model with the given parameters and return the signal.
+        """
+        raise NotImplementedError
+
+    def evaluateS1(self, parameters):
+        """
+        Evaluate the model with the given parameters and return the signal and
+        its derivatives.
+        """
+        raise NotImplementedError
+
+
+class FittingProblem(BaseProblem):
+    """
+    Defines the problem class for a fitting (parameter estimation) problem.
+    """
+
+    def __init__(
+        self,
+        model,
+        parameters,
+        dataset,
+        signal="Terminal voltage [V]",
+        check_model=True,
+        init_soc=None,
+        x0=None,
+    ):
+        super().__init__(parameters, model, check_model, init_soc, x0)
+        if model is not None:
+            self._model.signal = signal
+        self.signal = signal
+        self._dataset = {o.name: o for o in dataset}
         self.n_outputs = len([self.signal])
 
         # Check that the dataset contains time and current
@@ -44,30 +92,71 @@ def __init__(
         if len(self._target) != len(self._time_data):
             raise ValueError("Time data and signal data must be the same length.")
 
-        # Set bounds
-        self.bounds = dict(
-            lower=[param.bounds[0] for param in self.parameters],
-            upper=[param.bounds[1] for param in self.parameters],
+        # Build the model
+        if self._model._built_model is None:
+            self._model.build(
+                dataset=self._dataset,
+                parameters={o.name: o.value for o in self.parameters},
+                check_model=self.check_model,
+                init_soc=self.init_soc,
+            )
+
+    def evaluate(self, parameters):
+        """
+        Evaluate the model with the given parameters and return the signal.
+        """
+
+        y = np.asarray(self._model.simulate(inputs=parameters, t_eval=self._time_data))
+
+        return y
+
+    def evaluateS1(self, parameters):
+        """
+        Evaluate the model with the given parameters and return the signal and
+        its derivatives.
+        """
+
+        y, dy = self._model.simulateS1(
+            inputs=parameters,
+            t_eval=self._time_data,
         )
 
-        # Sample from prior for x0
-        if x0 is None:
-            self.x0 = np.zeros(self.n_parameters)
-            for i, param in enumerate(self.parameters):
-                self.x0[i] = param.rvs(1)
-        elif len(x0) != self.n_parameters:
-            raise ValueError("x0 dimensions do not match number of parameters")
+        return (np.asarray(y), np.asarray(dy))
 
-        # Add the initial values to the parameter definitions
-        for i, param in enumerate(self.parameters):
-            param.update(value=self.x0[i])
+    def target(self):
+        """
+        Returns the target dataset.
+        """
+        return self._target
 
-        # Set the fitting parameters and build the model
-        self.fit_parameters = {o.name: o.value for o in parameters}
-        if self._model._built_model is None:
+
+class DesignProblem(BaseProblem):
+    """
+    Defines the problem class for a design optimiation problem.
+    """
+
+    def __init__(
+        self,
+        model,
+        parameters,
+        experiment,
+        check_model=True,
+        init_soc=None,
+        x0=None,
+    ):
+        super().__init__(parameters, model, check_model, init_soc, x0)
+        self.experiment = experiment
+        self._target = None
+
+        # Build the model if required
+        if experiment is not None:
+            # Leave the build until later to apply the experiment
+            self._model.parameters = {o.name: o.value for o in self.parameters}
+
+        elif self._model._built_model is None:
             self._model.build(
-                dataset=self._dataset,
-                fit_parameters=self.fit_parameters,
+                experiment=self.experiment,
+                parameters={o.name: o.value for o in self.parameters},
                 check_model=self.check_model,
                 init_soc=self.init_soc,
             )
@@ -86,11 +175,9 @@ def evaluateS1(self, parameters):
         Evaluate the model with the given parameters and return the signal and
         its derivatives.
         """
-        for i, key in enumerate(self.fit_parameters):
-            self.fit_parameters[key] = parameters[i]
 
         y, dy = self._model.simulateS1(
-            inputs=self.fit_parameters,
+            inputs=parameters,
             t_eval=self._time_data,
         )
 

pybop/models/base_model.py

@@ -10,14 +10,14 @@ class BaseModel:
     def __init__(self, name="Base Model"):
         self.name = name
         self.pybamm_model = None
-        self.fit_parameters = None
+        self.parameters = None
         self.dataset = None
         self.signal = None
 
     def build(
         self,
         dataset=None,
-        fit_parameters=None,
+        parameters=None,
         check_model=True,
         init_soc=None,
     ):
@@ -26,10 +26,10 @@ def build(
         For PyBaMM forward models, this method follows a
         similar process to pybamm.Simulation.build().
         """
-        self.fit_parameters = fit_parameters
+        self.parameters = parameters
         self.dataset = dataset
-        if self.fit_parameters is not None:
-            self.fit_keys = list(self.fit_parameters.keys())
+        if self.parameters is not None:
+            self.fit_keys = list(self.parameters.keys())
 
         if init_soc is not None:
             self.set_init_soc(init_soc)
@@ -78,12 +78,12 @@ def set_params(self):
         if self.model_with_set_params:
             return
 
-        if self.fit_parameters is not None:
+        if self.parameters is not None:
             # set input parameters in parameter set from fitting parameters
-            for i in self.fit_parameters.keys():
+            for i in self.parameters.keys():
                 self._parameter_set[i] = "[input]"
 
-        if self.dataset is not None and self.fit_parameters is not None:
+        if self.dataset is not None and self.parameters is not None:
             if "Current function [A]" not in self.fit_keys:
                 self.parameter_set["Current function [A]"] = pybamm.Interpolant(
                     self.dataset["Time [s]"].data,
@@ -109,9 +109,7 @@ def simulate(self, inputs, t_eval):
             raise ValueError("Model must be built before calling simulate")
         else:
             if not isinstance(inputs, dict):
-                inputs_dict = {
-                    key: inputs[i] for i, key in enumerate(self.fit_parameters)
-                }
+                inputs_dict = {key: inputs[i] for i, key in enumerate(self.parameters)}
                 return self.solver.solve(
                     self.built_model, inputs=inputs_dict, t_eval=t_eval
                 )[self.signal].data
@@ -130,9 +128,7 @@ def simulateS1(self, inputs, t_eval):
             raise ValueError("Model must be built before calling simulate")
         else:
             if not isinstance(inputs, dict):
-                inputs_dict = {
-                    key: inputs[i] for i, key in enumerate(self.fit_parameters)
-                }
+                inputs_dict = {key: inputs[i] for i, key in enumerate(self.parameters)}
 
                 sol = self.solver.solve(
                     self.built_model,
@@ -171,6 +167,8 @@ def predict(
         """
         parameter_set = parameter_set or self._parameter_set
         if inputs is not None:
+            if not isinstance(inputs, dict):
+                inputs = {key: inputs[i] for i, key in enumerate(self.parameters)}
             parameter_set.update(inputs)
         if self._unprocessed_model is not None:
             if experiment is None:

tests/unit/test_cost.py

@@ -34,7 +34,7 @@ def test_costs(self, cut_off):
         # Construct Problem
         signal = "Voltage [V]"
         model.parameter_set.update({"Lower voltage cut-off [V]": cut_off})
-        problem = pybop.Problem(model, parameters, dataset, signal=signal, x0=x0)
+        problem = pybop.FittingProblem(model, parameters, dataset, signal=signal, x0=x0)
 
         # Base Cost
         base_cost = pybop.BaseCost(problem)

tests/unit/test_optimisation.py

@@ -29,8 +29,12 @@ def parameters(self):
 
     @pytest.fixture
     def problem(self, parameters, dataset):
-        return pybop.Problem(
-            pybop.lithium_ion.SPM(), parameters, dataset, signal="Terminal voltage [V]"
+        model = pybop.lithium_ion.SPM()
+        return pybop.FittingProblem(
+            model,
+            parameters,
+            dataset,
+            signal="Terminal voltage [V]",
         )
 
     @pytest.fixture