Generalise to assume multiple models and datasets

pybop/problems/fitting_problem.py

@@ -1,9 +1,18 @@
 import numpy as np
 
-from pybop import BaseProblem
+import pybop
 
 
-class FittingProblem(BaseProblem):
+"""
+TODO:
+1. Allow multiple datasets, one model
+    Include weights function: func[list[errors]]->scalar error
+2. Allow multiple models
+3. Allow serial, parallel dispatch
+"""
+
+
+class FittingProblem(pybop.BaseProblem):
     """
     Problem class for fitting (parameter estimation) problems.
 
@@ -42,41 +51,51 @@ def __init__(
         additional_variables.extend(["Time [s]"])
         additional_variables = list(set(additional_variables))
 
+        if not hasattr(model, "__len__"):
+            model = [model] * dataset.n_datasets
+        # TODO check valid (1,1), (1,n), (n,1), (n,n) model/dataset pairings
+
         super().__init__(
             parameters, model, check_model, signal, additional_variables, init_soc, x0
         )
-        self._dataset = dataset.data
+        self._dataset = dataset
         self.x = self.x0
 
         # Check that the dataset contains time and current
         dataset.check(self.signal + ["Current function [A]"])
 
         # Unpack time and target data
-        self._time_data = self._dataset["Time [s]"]
-        self.n_time_data = len(self._time_data)
         self.set_target(dataset)
+        try:
+            self._time_data = [data["Time [s]"] for data in self._dataset]
+        except TypeError:
+            self._time_data = [self._dataset["Time [s]"]]
 
         # Add useful parameters to model
         if model is not None:
-            self._model.signal = self.signal
-            self._model.additional_variables = self.additional_variables
-            self._model.n_parameters = self.n_parameters
-            self._model.n_outputs = self.n_outputs
-            self._model.n_time_data = self.n_time_data
-
-            # Build the model from scratch
-            if self._model._built_model is not None:
-                self._model._model_with_set_params = None
-                self._model._built_model = None
-                self._model._built_initial_soc = None
-                self._model._mesh = None
-                self._model._disc = None
-            self._model.build(
-                dataset=self._dataset,
-                parameters=self.parameters,
-                check_model=self.check_model,
-                init_soc=self.init_soc,
-            )
+            for thismodel, thisdataset in zip(self._model, self._dataset):
+                thismodel.signal = self.signal
+                thismodel.additional_variables = (
+                    self.additional_variables
+                )
+                # TODO generalise to allow different numbers of parameters, outputs
+                thismodel.n_parameters = self.n_parameters
+                thismodel.n_outputs = self.n_outputs
+
+                # Build the model from scratch
+                if thismodel._built_model is not None:
+                    thismodel._model_with_set_params = None
+                    thismodel._built_model = None
+                    thismodel._built_initial_soc = None
+                    thismodel._mesh = None
+                    thismodel._disc = None
+                # TODO generalise to allow different initial SOCs
+                thismodel.build(
+                    dataset=thisdataset,
+                    parameters=self.parameters,
+                    check_model=self.check_model,
+                    init_soc=self.init_soc,
+                )
 
     def evaluate(self, x):
         """
@@ -89,17 +108,20 @@ def evaluate(self, x):
 
         Returns
         -------
-        y : np.ndarray
-            The model output y(t) simulated with inputs x.
+        y : list[dict[str, np.ndarray[np.float64]]
+            List of model outputs y(t) simulated with inputs x, with
+            one dict for each dataset
         """
-        if np.any(x != self.x) and self._model.matched_parameters:
-            for i, param in enumerate(self.parameters):
-                param.update(value=x[i])
+        y = []
+        for model, ts in zip(self._model, self._time_data):
+            if np.any(x != self.x) and model.matched_parameters:
+                for i, param in enumerate(self.parameters):
+                    param.update(value=x[i])
 
-            self._model.rebuild(parameters=self.parameters)
-            self.x = x
+                model.rebuild(parameters=self.parameters)
+                self.x = x
 
-        y = self._model.simulate(inputs=x, t_eval=self._time_data)
+            y.append(model.simulate(inputs=x, t_eval=ts))
 
         return y
 
@@ -116,16 +138,41 @@ def evaluateS1(self, x):
         -------
         tuple
             A tuple containing the simulation result y(t) and the sensitivities dy/dx(t) evaluated
-            with given inputs x.
+            with given inputs x, for each dataset.
         """
         if self._model.matched_parameters:
             raise RuntimeError(
                 "Gradient not available when using geometric parameters."
             )
 
-        y, dy = self._model.simulateS1(
-            inputs=x,
-            t_eval=self._time_data,
-        )
+        sims = [
+            model.simulateS1(inputs=x, t_eval=ts)
+            for ts, model in zip(self._time_data, self._model)
+        ]
+        y = [s[0] for s in sims]
+        dy = [s[1] for s in sims]
 
         return (y, np.asarray(dy))
+
+    def set_target(self, dataset):
+        """
+        Set the target dataset.
+
+        Parameters
+        ----------
+        target : np.ndarray
+            The target dataset array.
+        """
+        if self.signal is None:
+            raise ValueError("Signal must be defined to set target.")
+        if not isinstance(dataset, (pybop.Dataset, pybop.DatasetCollection)):
+            raise ValueError("Dataset must be a pybop Dataset object.")
+
+        # List, with one signal/target dict per dataset
+        try:
+            self._target = [
+                {signal: dataset[signal][i] for signal in self.signal}
+                for i in range(dataset.n_datasets)
+            ]
+        except AttributeError:
+            self._target = [{signal: dataset[signal] for signal in self.signal}]