From 88a3a2cc3a1cf92efd4f7e0e70e4cde432afac23 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Tue, 28 Nov 2023 13:40:42 +0000
Subject: [PATCH 01/13] Split Problem into Base, Fitting and Design

---
 examples/notebooks/spm_nlopt.ipynb   |   2 +-
 examples/scripts/spm_CMAES.py        |   2 +-
 examples/scripts/spm_IRPropMin.py    |   2 +-
 examples/scripts/spm_SNES.py         |   2 +-
 examples/scripts/spm_XNES.py         |   2 +-
 examples/scripts/spm_adam.py         |   2 +-
 examples/scripts/spm_descent.py      |   2 +-
 examples/scripts/spm_nlopt.py        |   4 +-
 examples/scripts/spm_pso.py          |   2 +-
 pybop/__init__.py                    |   2 +-
 pybop/_problem.py                    | 133 ++++++++++++++++++++++-----
 pybop/models/base_model.py           |   8 +-
 tests/unit/test_cost.py              |   4 +-
 tests/unit/test_optimisation.py      |   7 +-
 tests/unit/test_parameterisations.py |  12 +--
 tests/unit/test_problem.py           |   6 +-
 16 files changed, 146 insertions(+), 46 deletions(-)

diff --git a/examples/notebooks/spm_nlopt.ipynb b/examples/notebooks/spm_nlopt.ipynb
index 811575160..50f0877d3 100644
--- a/examples/notebooks/spm_nlopt.ipynb
+++ b/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)"
       ]
     },
diff --git a/examples/scripts/spm_CMAES.py b/examples/scripts/spm_CMAES.py
index 7f044409e..6ac42df13 100644
--- a/examples/scripts/spm_CMAES.py
+++ b/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(parameters, dataset, model=model)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.CMAES)
 optim.set_max_iterations(100)
diff --git a/examples/scripts/spm_IRPropMin.py b/examples/scripts/spm_IRPropMin.py
index 2d4dd2ec4..d3aee1152 100644
--- a/examples/scripts/spm_IRPropMin.py
+++ b/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(parameters, dataset, model=model)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.IRPropMin)
 optim.set_max_iterations(100)
diff --git a/examples/scripts/spm_SNES.py b/examples/scripts/spm_SNES.py
index f5db3c9b9..7416ef966 100644
--- a/examples/scripts/spm_SNES.py
+++ b/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(parameters, dataset, model=model)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.SNES)
 optim.set_max_iterations(100)
diff --git a/examples/scripts/spm_XNES.py b/examples/scripts/spm_XNES.py
index 37939245f..e6f35bdd5 100644
--- a/examples/scripts/spm_XNES.py
+++ b/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(parameters, dataset, model=model)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.XNES)
 optim.set_max_iterations(100)
diff --git a/examples/scripts/spm_adam.py b/examples/scripts/spm_adam.py
index 27949e9ac..347a4f86f 100644
--- a/examples/scripts/spm_adam.py
+++ b/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(parameters, dataset, model=model)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.Adam)
 optim.set_max_iterations(100)
diff --git a/examples/scripts/spm_descent.py b/examples/scripts/spm_descent.py
index 85f77f262..508c5c0d7 100644
--- a/examples/scripts/spm_descent.py
+++ b/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(parameters, dataset, model=model)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.GradientDescent)
 optim.optimiser.set_learning_rate(0.025)
diff --git a/examples/scripts/spm_nlopt.py b/examples/scripts/spm_nlopt.py
index 19401ed45..bc6f59307 100644
--- a/examples/scripts/spm_nlopt.py
+++ b/examples/scripts/spm_nlopt.py
@@ -32,7 +32,9 @@
 
 # Define the cost to optimise
 signal = "Terminal voltage [V]"
-problem = pybop.Problem(model, parameters, dataset, signal=signal, init_soc=0.98)
+problem = pybop.FittingProblem(
+    parameters, dataset, model=model, signal=signal, init_soc=0.98
+)
 cost = pybop.RootMeanSquaredError(problem)
 
 # Build the optimisation problem
diff --git a/examples/scripts/spm_pso.py b/examples/scripts/spm_pso.py
index 9a9cb5aab..8757299ab 100644
--- a/examples/scripts/spm_pso.py
+++ b/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(parameters, dataset, model=model)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.PSO)
 optim.set_max_iterations(100)
diff --git a/pybop/__init__.py b/pybop/__init__.py
index 0e933b8e2..ed8adfe18 100644
--- a/pybop/__init__.py
+++ b/pybop/__init__.py
@@ -70,7 +70,7 @@
 #
 # Problem class
 #
-from ._problem import Problem
+from ._problem import FittingProblem, DesignProblem
 
 #
 # Plotting class
diff --git a/pybop/_problem.py b/pybop/_problem.py
index 469b65047..3d12ff9f3 100644
--- a/pybop/_problem.py
+++ b/pybop/_problem.py
@@ -1,30 +1,82 @@
 import numpy as np
 
 
-class Problem:
+class BaseProblem:
     """
     Defines a PyBOP single output problem, follows 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
+        if model is not None:
+            self.check_model = check_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.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])
+        self.fit_parameters = {o.name: o.value for o in parameters}
+
+        # Then build the model
+
+    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 a PyBOP single output problem, follows the PINTS interface.
+    """
+
+    def __init__(
+        self,
+        parameters,
+        dataset,
+        model=None,
+        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,29 +96,62 @@ 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,
+                fit_parameters=self.fit_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.
+        """
+        for i, key in enumerate(self.fit_parameters):
+            self.fit_parameters[key] = parameters[i]
+
+        y, dy = self._model.simulateS1(
+            inputs=self.fit_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])
+
+class DesignProblem(BaseProblem):
+    """
+    Defines a PyBOP single output problem, follows the PINTS interface.
+    """
+
+    def __init__(
+        self,
+        parameters,
+        experiment,
+        model=None,
+        check_model=True,
+        init_soc=None,
+        x0=None,
+    ):
+        super().__init__(parameters, model, check_model, init_soc, x0)
+        self.experiment = experiment
+        self._target = None
 
         # 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:
             self._model.build(
-                dataset=self._dataset,
+                experiment=self.experiment,
                 fit_parameters=self.fit_parameters,
                 check_model=self.check_model,
                 init_soc=self.init_soc,
diff --git a/pybop/models/base_model.py b/pybop/models/base_model.py
index ced38437e..ca780ee2c 100644
--- a/pybop/models/base_model.py
+++ b/pybop/models/base_model.py
@@ -17,6 +17,7 @@ def __init__(self, name="Base Model"):
     def build(
         self,
         dataset=None,
+        experiment=None,
         fit_parameters=None,
         check_model=True,
         init_soc=None,
@@ -28,13 +29,18 @@ def build(
         """
         self.fit_parameters = fit_parameters
         self.dataset = dataset
+        self.experiment = experiment
         if self.fit_parameters is not None:
             self.fit_keys = list(self.fit_parameters.keys())
 
         if init_soc is not None:
             self.set_init_soc(init_soc)
 
-        if self._built_model:
+        if experiment is not None:
+            # Leave the build until later to apply the experiment
+            return
+
+        elif self._built_model:
             return
 
         elif self.pybamm_model.is_discretised:
diff --git a/tests/unit/test_cost.py b/tests/unit/test_cost.py
index 0c7e329f3..8408c992e 100644
--- a/tests/unit/test_cost.py
+++ b/tests/unit/test_cost.py
@@ -34,7 +34,9 @@ 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(
+            parameters, dataset, model=model, signal=signal, x0=x0
+        )
 
         # Base Cost
         base_cost = pybop.BaseCost(problem)
diff --git a/tests/unit/test_optimisation.py b/tests/unit/test_optimisation.py
index 21ecf16ac..19c899841 100644
--- a/tests/unit/test_optimisation.py
+++ b/tests/unit/test_optimisation.py
@@ -29,8 +29,11 @@ def parameters(self):
 
     @pytest.fixture
     def problem(self, parameters, dataset):
-        return pybop.Problem(
-            pybop.lithium_ion.SPM(), parameters, dataset, signal="Terminal voltage [V]"
+        return pybop.FittingProblem(
+            parameters,
+            dataset,
+            model=pybop.lithium_ion.SPM(),
+            signal="Terminal voltage [V]",
         )
 
     @pytest.fixture
diff --git a/tests/unit/test_parameterisations.py b/tests/unit/test_parameterisations.py
index c18100638..1ca77e83b 100644
--- a/tests/unit/test_parameterisations.py
+++ b/tests/unit/test_parameterisations.py
@@ -44,8 +44,8 @@ def test_spm(self, init_soc):
 
         # Define the cost to optimise
         signal = "Terminal voltage [V]"
-        problem = pybop.Problem(
-            model, parameters, dataset, signal=signal, init_soc=init_soc
+        problem = pybop.FittingProblem(
+            parameters, dataset, model=model, signal=signal, init_soc=init_soc
         )
         cost = pybop.RootMeanSquaredError(problem)
 
@@ -97,8 +97,8 @@ def test_spm_optimisers(self, init_soc):
 
         # Define the cost to optimise
         signal = "Terminal voltage [V]"
-        problem = pybop.Problem(
-            model, parameters, dataset, signal=signal, init_soc=init_soc
+        problem = pybop.FittingProblem(
+            parameters, dataset, model=model, signal=signal, init_soc=init_soc
         )
         cost = pybop.SumSquaredError(problem)
 
@@ -193,8 +193,8 @@ def test_model_misparameterisation(self, init_soc):
 
         # Define the cost to optimise
         signal = "Terminal voltage [V]"
-        problem = pybop.Problem(
-            model, parameters, dataset, signal=signal, init_soc=init_soc
+        problem = pybop.FittingProblem(
+            parameters, dataset, model=model, signal=signal, init_soc=init_soc
         )
         cost = pybop.RootMeanSquaredError(problem)
 
diff --git a/tests/unit/test_problem.py b/tests/unit/test_problem.py
index aa470d9b4..da53f88c2 100644
--- a/tests/unit/test_problem.py
+++ b/tests/unit/test_problem.py
@@ -39,10 +39,12 @@ def test_problem(self):
 
         # Test incorrect number of initial parameter values
         with pytest.raises(ValueError):
-            pybop.Problem(model, parameters, dataset, signal=signal, x0=np.array([]))
+            pybop.FittingProblem(
+                parameters, dataset, model=model, signal=signal, x0=np.array([])
+            )
 
         # Construct Problem
-        problem = pybop.Problem(model, parameters, dataset, signal=signal)
+        problem = pybop.FittingProblem(parameters, dataset, model=model, signal=signal)
 
         assert problem._model == model
         assert problem._model._built_model is not None

From 46eff73203e48671625cf7138970b9d2151bb91a Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Tue, 28 Nov 2023 14:35:34 +0000
Subject: [PATCH 02/13] Add test for DesignProblem

---
 pybop/models/base_model.py |  2 ++
 tests/unit/test_problem.py | 53 +++++++++++++++++++++++++++++++++++++-
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/pybop/models/base_model.py b/pybop/models/base_model.py
index ca780ee2c..34307fa70 100644
--- a/pybop/models/base_model.py
+++ b/pybop/models/base_model.py
@@ -177,6 +177,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.fit_parameters)}
             parameter_set.update(inputs)
         if self._unprocessed_model is not None:
             if experiment is None:
diff --git a/tests/unit/test_problem.py b/tests/unit/test_problem.py
index da53f88c2..102684b20 100644
--- a/tests/unit/test_problem.py
+++ b/tests/unit/test_problem.py
@@ -10,7 +10,7 @@ class TestProblem:
     """
 
     @pytest.mark.unit
-    def test_problem(self):
+    def test_fitting_problem(self):
         # Define model
         model = pybop.lithium_ion.SPM()
         parameters = [
@@ -74,3 +74,54 @@ def getdata(self, model, x0):
         )
         sim = model.predict(experiment=experiment)
         return sim
+
+    def test_design_problem(self):
+        # Define model
+        model = pybop.lithium_ion.SPM()
+        model.parameter_set = model.pybamm_model.default_parameter_values
+        x0 = np.array([0.52, 0.63])
+        model.parameter_set.update(
+            {
+                "Negative electrode active material volume fraction": x0[0],
+                "Positive electrode active material volume fraction": x0[1],
+            }
+        )
+
+        # Define target experiment
+        experiment = pybamm.Experiment(
+            [
+                (
+                    "Discharge at 1C for 5 minutes (1 second period)",
+                    "Rest for 2 minutes (1 second period)",
+                    "Charge at 1C for 5 minutes (1 second period)",
+                    "Rest for 2 minutes (1 second period)",
+                ),
+            ]
+            * 2
+        )
+
+        parameters = [
+            pybop.Parameter(
+                "Negative electrode active material volume fraction",
+                prior=pybop.Gaussian(0.5, 0.02),
+                bounds=[0.375, 0.625],
+            ),
+            pybop.Parameter(
+                "Positive electrode active material volume fraction",
+                prior=pybop.Gaussian(0.65, 0.02),
+                bounds=[0.525, 0.75],
+            ),
+        ]
+
+        # Test incorrect number of initial parameter values
+        with pytest.raises(ValueError):
+            pybop.DesignProblem(parameters, experiment, model=model, x0=np.array([]))
+
+        # Construct Problem
+        problem = pybop.DesignProblem(parameters, experiment, model=model)
+
+        assert problem._model == model
+        # problem._model._built_model is None (with an experiment)
+
+        # Test model.predict
+        model.predict(inputs=[0.5, 0.5], experiment=experiment)

From 68311621768f599ffcb025db14ffda47d7ab74aa Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Tue, 28 Nov 2023 15:32:22 +0000
Subject: [PATCH 03/13] Add test for BaseProblem

---
 tests/unit/test_problem.py | 40 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/tests/unit/test_problem.py b/tests/unit/test_problem.py
index 102684b20..b1ee55509 100644
--- a/tests/unit/test_problem.py
+++ b/tests/unit/test_problem.py
@@ -9,6 +9,46 @@ class TestProblem:
     A class to test the problem class.
     """
 
+    @pytest.mark.unit
+    def test_base_problem(self):
+        # Define model
+        model = pybop.lithium_ion.SPM()
+        model.parameter_set = model.pybamm_model.default_parameter_values
+        x0 = np.array([0.52, 0.63])
+        model.parameter_set.update(
+            {
+                "Negative electrode active material volume fraction": x0[0],
+                "Positive electrode active material volume fraction": x0[1],
+            }
+        )
+
+        parameters = [
+            pybop.Parameter(
+                "Negative electrode active material volume fraction",
+                prior=pybop.Gaussian(0.5, 0.02),
+                bounds=[0.375, 0.625],
+            ),
+            pybop.Parameter(
+                "Positive electrode active material volume fraction",
+                prior=pybop.Gaussian(0.65, 0.02),
+                bounds=[0.525, 0.75],
+            ),
+        ]
+
+        # Test incorrect number of initial parameter values
+        with pytest.raises(ValueError):
+            pybop._problem.BaseProblem(parameters, model=model, x0=np.array([]))
+
+        # Construct Problem
+        problem = pybop._problem.BaseProblem(parameters, model=model)
+
+        assert problem._model == model
+
+        with pytest.raises(NotImplementedError):
+            problem.evaluate([0.5, 0.5])
+        with pytest.raises(NotImplementedError):
+            problem.evaluateS1([0.5, 0.5])
+
     @pytest.mark.unit
     def test_fitting_problem(self):
         # Define model

From 00942203a9d32820d7b68ec5368897f27cbd97ce Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Tue, 28 Nov 2023 18:31:57 +0000
Subject: [PATCH 04/13] Rename fit_parameters to parameters

---
 pybop/_problem.py          | 26 ++++++++++++--------------
 pybop/models/base_model.py | 26 +++++++++++---------------
 2 files changed, 23 insertions(+), 29 deletions(-)

diff --git a/pybop/_problem.py b/pybop/_problem.py
index 3d12ff9f3..1f79afb2a 100644
--- a/pybop/_problem.py
+++ b/pybop/_problem.py
@@ -3,7 +3,7 @@
 
 class BaseProblem:
     """
-    Defines a PyBOP single output problem, follows the PINTS interface.
+    Defines the PyBOP base problem, following the PINTS interface.
     """
 
     def __init__(
@@ -39,9 +39,7 @@ def __init__(
         # Add the initial values to the parameter definitions
         for i, param in enumerate(self.parameters):
             param.update(value=self.x0[i])
-        self.fit_parameters = {o.name: o.value for o in parameters}
-
-        # Then build the model
+        self.parameters = {o.name: o.value for o in parameters}
 
     def evaluate(self, parameters):
         """
@@ -59,7 +57,7 @@ def evaluateS1(self, parameters):
 
 class FittingProblem(BaseProblem):
     """
-    Defines a PyBOP single output problem, follows the PINTS interface.
+    Defines the problem class for a fitting (parameter estimation) problem.
     """
 
     def __init__(
@@ -100,7 +98,7 @@ def __init__(
         if self._model._built_model is None:
             self._model.build(
                 dataset=self._dataset,
-                fit_parameters=self.fit_parameters,
+                parameters=self.parameters,
                 check_model=self.check_model,
                 init_soc=self.init_soc,
             )
@@ -119,11 +117,11 @@ 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]
+        for i, key in enumerate(self.parameters):
+            self.parameters[key] = parameters[i]
 
         y, dy = self._model.simulateS1(
-            inputs=self.fit_parameters,
+            inputs=self.parameters,
             t_eval=self._time_data,
         )
 
@@ -132,7 +130,7 @@ def evaluateS1(self, parameters):
 
 class DesignProblem(BaseProblem):
     """
-    Defines a PyBOP single output problem, follows the PINTS interface.
+    Defines the problem class for a design optimiation problem.
     """
 
     def __init__(
@@ -152,7 +150,7 @@ def __init__(
         if self._model._built_model is None:
             self._model.build(
                 experiment=self.experiment,
-                fit_parameters=self.fit_parameters,
+                parameters=self.parameters,
                 check_model=self.check_model,
                 init_soc=self.init_soc,
             )
@@ -171,11 +169,11 @@ 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]
+        for i, key in enumerate(self.parameters):
+            self.parameters[key] = parameters[i]
 
         y, dy = self._model.simulateS1(
-            inputs=self.fit_parameters,
+            inputs=self.parameters,
             t_eval=self._time_data,
         )
 
diff --git a/pybop/models/base_model.py b/pybop/models/base_model.py
index 34307fa70..d7ac1a6ca 100644
--- a/pybop/models/base_model.py
+++ b/pybop/models/base_model.py
@@ -10,7 +10,7 @@ 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
 
@@ -18,7 +18,7 @@ def build(
         self,
         dataset=None,
         experiment=None,
-        fit_parameters=None,
+        parameters=None,
         check_model=True,
         init_soc=None,
     ):
@@ -27,11 +27,11 @@ 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
         self.experiment = experiment
-        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)
@@ -84,12 +84,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,
@@ -115,9 +115,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
@@ -136,9 +134,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,
@@ -178,7 +174,7 @@ 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.fit_parameters)}
+                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:

From 08f435c49b5a1d9a740a6a7f03a1d955c98801eb Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 30 Nov 2023 13:22:38 +0000
Subject: [PATCH 05/13] Make model required except for BaseProblem

---
 pybop/_problem.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/pybop/_problem.py b/pybop/_problem.py
index 1f79afb2a..f08608c58 100644
--- a/pybop/_problem.py
+++ b/pybop/_problem.py
@@ -15,8 +15,7 @@ def __init__(
         x0=None,
     ):
         self._model = model
-        if model is not None:
-            self.check_model = check_model
+        self.check_model = check_model
         self.parameters = parameters
         self.init_soc = init_soc
         self.x0 = x0
@@ -64,7 +63,7 @@ def __init__(
         self,
         parameters,
         dataset,
-        model=None,
+        model,
         signal="Terminal voltage [V]",
         check_model=True,
         init_soc=None,
@@ -137,7 +136,7 @@ def __init__(
         self,
         parameters,
         experiment,
-        model=None,
+        model,
         check_model=True,
         init_soc=None,
         x0=None,
@@ -146,7 +145,7 @@ def __init__(
         self.experiment = experiment
         self._target = None
 
-        # Set the fitting parameters and build the model
+        # Set the input parameters and build the model
         if self._model._built_model is None:
             self._model.build(
                 experiment=self.experiment,

From 5b11cf3f8647bcee48f2d941e7c27ef18d90508c Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 30 Nov 2023 13:29:39 +0000
Subject: [PATCH 06/13] Check for experiment prior to build

---
 pybop/_problem.py          | 8 ++++++--
 pybop/models/base_model.py | 6 +-----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/pybop/_problem.py b/pybop/_problem.py
index f08608c58..07df3a866 100644
--- a/pybop/_problem.py
+++ b/pybop/_problem.py
@@ -145,8 +145,12 @@ def __init__(
         self.experiment = experiment
         self._target = None
 
-        # Set the input parameters and build the model
-        if self._model._built_model is None:
+        # Build the model if required
+        if experiment is not None:
+            # Leave the build until later to apply the experiment
+            return
+
+        elif self._model._built_model is None:
             self._model.build(
                 experiment=self.experiment,
                 parameters=self.parameters,
diff --git a/pybop/models/base_model.py b/pybop/models/base_model.py
index d7ac1a6ca..4ff35685a 100644
--- a/pybop/models/base_model.py
+++ b/pybop/models/base_model.py
@@ -36,11 +36,7 @@ def build(
         if init_soc is not None:
             self.set_init_soc(init_soc)
 
-        if experiment is not None:
-            # Leave the build until later to apply the experiment
-            return
-
-        elif self._built_model:
+        if self._built_model:
             return
 
         elif self.pybamm_model.is_discretised:

From 7cecdc73f983c173761e7b9b98f167d5977784a1 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 30 Nov 2023 13:58:05 +0000
Subject: [PATCH 07/13] Reset model as required input

---
 examples/scripts/spm_CMAES.py     | 2 +-
 examples/scripts/spm_IRPropMin.py | 2 +-
 examples/scripts/spm_SNES.py      | 2 +-
 examples/scripts/spm_XNES.py      | 2 +-
 examples/scripts/spm_adam.py      | 2 +-
 examples/scripts/spm_descent.py   | 2 +-
 examples/scripts/spm_nlopt.py     | 4 +---
 examples/scripts/spm_pso.py       | 2 +-
 8 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/examples/scripts/spm_CMAES.py b/examples/scripts/spm_CMAES.py
index 6ac42df13..0404c55e2 100644
--- a/examples/scripts/spm_CMAES.py
+++ b/examples/scripts/spm_CMAES.py
@@ -33,7 +33,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(parameters, dataset, model=model)
+problem = pybop.FittingProblem(parameters, dataset, model)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.CMAES)
 optim.set_max_iterations(100)
diff --git a/examples/scripts/spm_IRPropMin.py b/examples/scripts/spm_IRPropMin.py
index d3aee1152..8df43dd7d 100644
--- a/examples/scripts/spm_IRPropMin.py
+++ b/examples/scripts/spm_IRPropMin.py
@@ -33,7 +33,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(parameters, dataset, model=model)
+problem = pybop.FittingProblem(parameters, dataset, model)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.IRPropMin)
 optim.set_max_iterations(100)
diff --git a/examples/scripts/spm_SNES.py b/examples/scripts/spm_SNES.py
index 7416ef966..9d3d0f857 100644
--- a/examples/scripts/spm_SNES.py
+++ b/examples/scripts/spm_SNES.py
@@ -33,7 +33,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(parameters, dataset, model=model)
+problem = pybop.FittingProblem(parameters, dataset, model)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.SNES)
 optim.set_max_iterations(100)
diff --git a/examples/scripts/spm_XNES.py b/examples/scripts/spm_XNES.py
index e6f35bdd5..f7b9e6873 100644
--- a/examples/scripts/spm_XNES.py
+++ b/examples/scripts/spm_XNES.py
@@ -33,7 +33,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(parameters, dataset, model=model)
+problem = pybop.FittingProblem(parameters, dataset, model)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.XNES)
 optim.set_max_iterations(100)
diff --git a/examples/scripts/spm_adam.py b/examples/scripts/spm_adam.py
index 347a4f86f..56795ebe0 100644
--- a/examples/scripts/spm_adam.py
+++ b/examples/scripts/spm_adam.py
@@ -36,7 +36,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(parameters, dataset, model=model)
+problem = pybop.FittingProblem(parameters, dataset, model)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.Adam)
 optim.set_max_iterations(100)
diff --git a/examples/scripts/spm_descent.py b/examples/scripts/spm_descent.py
index 508c5c0d7..0d4a16596 100644
--- a/examples/scripts/spm_descent.py
+++ b/examples/scripts/spm_descent.py
@@ -36,7 +36,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(parameters, dataset, model=model)
+problem = pybop.FittingProblem(parameters, dataset, model)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.GradientDescent)
 optim.optimiser.set_learning_rate(0.025)
diff --git a/examples/scripts/spm_nlopt.py b/examples/scripts/spm_nlopt.py
index bc6f59307..d46b93aa3 100644
--- a/examples/scripts/spm_nlopt.py
+++ b/examples/scripts/spm_nlopt.py
@@ -32,9 +32,7 @@
 
 # Define the cost to optimise
 signal = "Terminal voltage [V]"
-problem = pybop.FittingProblem(
-    parameters, dataset, model=model, signal=signal, init_soc=0.98
-)
+problem = pybop.FittingProblem(parameters, dataset, model, signal=signal, init_soc=0.98)
 cost = pybop.RootMeanSquaredError(problem)
 
 # Build the optimisation problem
diff --git a/examples/scripts/spm_pso.py b/examples/scripts/spm_pso.py
index 8757299ab..8313e2f38 100644
--- a/examples/scripts/spm_pso.py
+++ b/examples/scripts/spm_pso.py
@@ -33,7 +33,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(parameters, dataset, model=model)
+problem = pybop.FittingProblem(parameters, dataset, model)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.PSO)
 optim.set_max_iterations(100)

From cf09bbffb9f81ced610962747d7da7222fce9d05 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 30 Nov 2023 13:58:47 +0000
Subject: [PATCH 08/13] Define parameters when not building

---
 pybop/_problem.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pybop/_problem.py b/pybop/_problem.py
index 07df3a866..82e43c6ab 100644
--- a/pybop/_problem.py
+++ b/pybop/_problem.py
@@ -148,7 +148,7 @@ def __init__(
         # Build the model if required
         if experiment is not None:
             # Leave the build until later to apply the experiment
-            return
+            self._model.parameters = self.parameters
 
         elif self._model._built_model is None:
             self._model.build(

From edf4c824c483288704c262f930330802a78db49c Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Fri, 1 Dec 2023 12:01:26 +0000
Subject: [PATCH 09/13] Refactor test_problem and update Changelog

---
 CHANGELOG.md               |   3 +-
 tests/unit/test_problem.py | 147 +++++++++++++------------------------
 2 files changed, 51 insertions(+), 99 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c5780d436..df0d53f8c 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/tests/unit/test_problem.py b/tests/unit/test_problem.py
index b1ee55509..6de7771e6 100644
--- a/tests/unit/test_problem.py
+++ b/tests/unit/test_problem.py
@@ -9,20 +9,25 @@ class TestProblem:
     A class to test the problem class.
     """
 
-    @pytest.mark.unit
-    def test_base_problem(self):
-        # Define model
+    @pytest.fixture
+    def x0(self):
+        return np.array([0.52, 0.63])
+
+    @pytest.fixture
+    def model(self, x0):
         model = pybop.lithium_ion.SPM()
         model.parameter_set = model.pybamm_model.default_parameter_values
-        x0 = np.array([0.52, 0.63])
         model.parameter_set.update(
             {
                 "Negative electrode active material volume fraction": x0[0],
                 "Positive electrode active material volume fraction": x0[1],
             }
         )
+        return model
 
-        parameters = [
+    @pytest.fixture
+    def parameters(self):
+        return [
             pybop.Parameter(
                 "Negative electrode active material volume fraction",
                 prior=pybop.Gaussian(0.5, 0.02),
@@ -35,6 +40,35 @@ def test_base_problem(self):
             ),
         ]
 
+    @pytest.fixture
+    def experiment(self):
+        return pybamm.Experiment(
+            [
+                (
+                    "Discharge at 1C for 5 minutes (1 second period)",
+                    "Rest for 2 minutes (1 second period)",
+                    "Charge at 1C for 5 minutes (1 second period)",
+                    "Rest for 2 minutes (1 second period)",
+                ),
+            ]
+            * 2
+        )
+
+    @pytest.fixture
+    def dataset(self, model, experiment):
+        solution = model.predict(experiment=experiment)
+        return [
+            pybop.Dataset("Time [s]", solution["Time [s]"].data),
+            pybop.Dataset("Current function [A]", solution["Current [A]"].data),
+            pybop.Dataset("Voltage [V]", solution["Terminal voltage [V]"].data),
+        ]
+
+    @pytest.fixture
+    def signal(self):
+        return "Voltage [V]"
+
+    @pytest.mark.unit
+    def test_base_problem(self, parameters, model):
         # Test incorrect number of initial parameter values
         with pytest.raises(ValueError):
             pybop._problem.BaseProblem(parameters, model=model, x0=np.array([]))
@@ -50,41 +84,15 @@ def test_base_problem(self):
             problem.evaluateS1([0.5, 0.5])
 
     @pytest.mark.unit
-    def test_fitting_problem(self):
-        # Define model
-        model = pybop.lithium_ion.SPM()
-        parameters = [
-            pybop.Parameter(
-                "Negative electrode active material volume fraction",
-                prior=pybop.Gaussian(0.5, 0.02),
-                bounds=[0.375, 0.625],
-            ),
-            pybop.Parameter(
-                "Positive electrode active material volume fraction",
-                prior=pybop.Gaussian(0.65, 0.02),
-                bounds=[0.525, 0.75],
-            ),
-        ]
-        signal = "Voltage [V]"
-
-        # Form dataset
-        x0 = np.array([0.52, 0.63])
-        solution = self.getdata(model, x0)
-
-        dataset = [
-            pybop.Dataset("Time [s]", solution["Time [s]"].data),
-            pybop.Dataset("Current function [A]", solution["Current [A]"].data),
-            pybop.Dataset("Voltage [V]", solution["Terminal voltage [V]"].data),
-        ]
-
+    def test_fitting_problem(self, parameters, dataset, model, signal):
         # Test incorrect number of initial parameter values
         with pytest.raises(ValueError):
             pybop.FittingProblem(
-                parameters, dataset, model=model, signal=signal, x0=np.array([])
+                parameters, dataset, model, signal=signal, x0=np.array([])
             )
 
         # Construct Problem
-        problem = pybop.FittingProblem(parameters, dataset, model=model, signal=signal)
+        problem = pybop.FittingProblem(parameters, dataset, model, signal=signal)
 
         assert problem._model == model
         assert problem._model._built_model is not None
@@ -92,76 +100,19 @@ def test_fitting_problem(self):
         # Test model.simulate
         model.simulate(inputs=[0.5, 0.5], t_eval=np.linspace(0, 10, 100))
 
-    def getdata(self, model, x0):
-        model.parameter_set = model.pybamm_model.default_parameter_values
-
-        model.parameter_set.update(
-            {
-                "Negative electrode active material volume fraction": x0[0],
-                "Positive electrode active material volume fraction": x0[1],
-            }
-        )
-        experiment = pybamm.Experiment(
-            [
-                (
-                    "Discharge at 1C for 5 minutes (1 second period)",
-                    "Rest for 2 minutes (1 second period)",
-                    "Charge at 1C for 5 minutes (1 second period)",
-                    "Rest for 2 minutes (1 second period)",
-                ),
-            ]
-            * 2
-        )
-        sim = model.predict(experiment=experiment)
-        return sim
-
-    def test_design_problem(self):
-        # Define model
-        model = pybop.lithium_ion.SPM()
-        model.parameter_set = model.pybamm_model.default_parameter_values
-        x0 = np.array([0.52, 0.63])
-        model.parameter_set.update(
-            {
-                "Negative electrode active material volume fraction": x0[0],
-                "Positive electrode active material volume fraction": x0[1],
-            }
-        )
-
-        # Define target experiment
-        experiment = pybamm.Experiment(
-            [
-                (
-                    "Discharge at 1C for 5 minutes (1 second period)",
-                    "Rest for 2 minutes (1 second period)",
-                    "Charge at 1C for 5 minutes (1 second period)",
-                    "Rest for 2 minutes (1 second period)",
-                ),
-            ]
-            * 2
-        )
-
-        parameters = [
-            pybop.Parameter(
-                "Negative electrode active material volume fraction",
-                prior=pybop.Gaussian(0.5, 0.02),
-                bounds=[0.375, 0.625],
-            ),
-            pybop.Parameter(
-                "Positive electrode active material volume fraction",
-                prior=pybop.Gaussian(0.65, 0.02),
-                bounds=[0.525, 0.75],
-            ),
-        ]
-
+    @pytest.mark.unit
+    def test_design_problem(self, parameters, experiment, model):
         # Test incorrect number of initial parameter values
         with pytest.raises(ValueError):
-            pybop.DesignProblem(parameters, experiment, model=model, x0=np.array([]))
+            pybop.DesignProblem(parameters, experiment, model, x0=np.array([]))
 
         # Construct Problem
-        problem = pybop.DesignProblem(parameters, experiment, model=model)
+        problem = pybop.DesignProblem(parameters, experiment, model)
 
         assert problem._model == model
-        # problem._model._built_model is None (with an experiment)
+        assert (
+            problem._model._built_model is None
+        )  # building postponed with input experiment
 
         # Test model.predict
         model.predict(inputs=[0.5, 0.5], experiment=experiment)

From c2c161baa787f409dd217e698d3255b86308b142 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Fri, 1 Dec 2023 12:16:43 +0000
Subject: [PATCH 10/13] Move setting of x0

---
 tests/unit/test_problem.py | 24 ++++++++++--------------
 1 file changed, 10 insertions(+), 14 deletions(-)

diff --git a/tests/unit/test_problem.py b/tests/unit/test_problem.py
index 6de7771e6..0908e8186 100644
--- a/tests/unit/test_problem.py
+++ b/tests/unit/test_problem.py
@@ -10,20 +10,8 @@ class TestProblem:
     """
 
     @pytest.fixture
-    def x0(self):
-        return np.array([0.52, 0.63])
-
-    @pytest.fixture
-    def model(self, x0):
-        model = pybop.lithium_ion.SPM()
-        model.parameter_set = model.pybamm_model.default_parameter_values
-        model.parameter_set.update(
-            {
-                "Negative electrode active material volume fraction": x0[0],
-                "Positive electrode active material volume fraction": x0[1],
-            }
-        )
-        return model
+    def model(self):
+        return pybop.lithium_ion.SPM()
 
     @pytest.fixture
     def parameters(self):
@@ -56,6 +44,14 @@ def experiment(self):
 
     @pytest.fixture
     def dataset(self, model, experiment):
+        model.parameter_set = model.pybamm_model.default_parameter_values
+        x0 = np.array([0.52, 0.63])
+        model.parameter_set.update(
+            {
+                "Negative electrode active material volume fraction": x0[0],
+                "Positive electrode active material volume fraction": x0[1],
+            }
+        )
         solution = model.predict(experiment=experiment)
         return [
             pybop.Dataset("Time [s]", solution["Time [s]"].data),

From 88d8131f717eebb71187e969de049896dc552fde Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Fri, 1 Dec 2023 16:36:34 +0000
Subject: [PATCH 11/13] Reset order of inputs

---
 examples/scripts/spm_CMAES.py        | 2 +-
 examples/scripts/spm_IRPropMin.py    | 2 +-
 examples/scripts/spm_SNES.py         | 2 +-
 examples/scripts/spm_XNES.py         | 2 +-
 examples/scripts/spm_adam.py         | 2 +-
 examples/scripts/spm_descent.py      | 2 +-
 examples/scripts/spm_nlopt.py        | 2 +-
 examples/scripts/spm_pso.py          | 2 +-
 pybop/_problem.py                    | 4 ++--
 pybop/models/base_model.py           | 2 --
 tests/unit/test_cost.py              | 4 +---
 tests/unit/test_optimisation.py      | 3 ++-
 tests/unit/test_parameterisations.py | 6 +++---
 tests/unit/test_problem.py           | 8 ++++----
 14 files changed, 20 insertions(+), 23 deletions(-)

diff --git a/examples/scripts/spm_CMAES.py b/examples/scripts/spm_CMAES.py
index 0404c55e2..56a2aca7c 100644
--- a/examples/scripts/spm_CMAES.py
+++ b/examples/scripts/spm_CMAES.py
@@ -33,7 +33,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(parameters, dataset, model)
+problem = pybop.FittingProblem(model, parameters, dataset)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.CMAES)
 optim.set_max_iterations(100)
diff --git a/examples/scripts/spm_IRPropMin.py b/examples/scripts/spm_IRPropMin.py
index 8df43dd7d..f620ff7e7 100644
--- a/examples/scripts/spm_IRPropMin.py
+++ b/examples/scripts/spm_IRPropMin.py
@@ -33,7 +33,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(parameters, dataset, model)
+problem = pybop.FittingProblem(model, parameters, dataset)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.IRPropMin)
 optim.set_max_iterations(100)
diff --git a/examples/scripts/spm_SNES.py b/examples/scripts/spm_SNES.py
index 9d3d0f857..c45362cf0 100644
--- a/examples/scripts/spm_SNES.py
+++ b/examples/scripts/spm_SNES.py
@@ -33,7 +33,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(parameters, dataset, model)
+problem = pybop.FittingProblem(model, parameters, dataset)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.SNES)
 optim.set_max_iterations(100)
diff --git a/examples/scripts/spm_XNES.py b/examples/scripts/spm_XNES.py
index f7b9e6873..3f3960a34 100644
--- a/examples/scripts/spm_XNES.py
+++ b/examples/scripts/spm_XNES.py
@@ -33,7 +33,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(parameters, dataset, model)
+problem = pybop.FittingProblem(model, parameters, dataset)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.XNES)
 optim.set_max_iterations(100)
diff --git a/examples/scripts/spm_adam.py b/examples/scripts/spm_adam.py
index 56795ebe0..0560a09f4 100644
--- a/examples/scripts/spm_adam.py
+++ b/examples/scripts/spm_adam.py
@@ -36,7 +36,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(parameters, dataset, model)
+problem = pybop.FittingProblem(model, parameters, dataset)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.Adam)
 optim.set_max_iterations(100)
diff --git a/examples/scripts/spm_descent.py b/examples/scripts/spm_descent.py
index 0d4a16596..a5ec52342 100644
--- a/examples/scripts/spm_descent.py
+++ b/examples/scripts/spm_descent.py
@@ -36,7 +36,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(parameters, dataset, model)
+problem = pybop.FittingProblem(model, parameters, dataset)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.GradientDescent)
 optim.optimiser.set_learning_rate(0.025)
diff --git a/examples/scripts/spm_nlopt.py b/examples/scripts/spm_nlopt.py
index d46b93aa3..733aea719 100644
--- a/examples/scripts/spm_nlopt.py
+++ b/examples/scripts/spm_nlopt.py
@@ -32,7 +32,7 @@
 
 # Define the cost to optimise
 signal = "Terminal voltage [V]"
-problem = pybop.FittingProblem(parameters, dataset, model, 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
diff --git a/examples/scripts/spm_pso.py b/examples/scripts/spm_pso.py
index 8313e2f38..f94e06940 100644
--- a/examples/scripts/spm_pso.py
+++ b/examples/scripts/spm_pso.py
@@ -33,7 +33,7 @@
 ]
 
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(parameters, dataset, model)
+problem = pybop.FittingProblem(model, parameters, dataset)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.PSO)
 optim.set_max_iterations(100)
diff --git a/pybop/_problem.py b/pybop/_problem.py
index 82e43c6ab..5d7050757 100644
--- a/pybop/_problem.py
+++ b/pybop/_problem.py
@@ -61,9 +61,9 @@ class FittingProblem(BaseProblem):
 
     def __init__(
         self,
+        model,
         parameters,
         dataset,
-        model,
         signal="Terminal voltage [V]",
         check_model=True,
         init_soc=None,
@@ -134,9 +134,9 @@ class DesignProblem(BaseProblem):
 
     def __init__(
         self,
+        model,
         parameters,
         experiment,
-        model,
         check_model=True,
         init_soc=None,
         x0=None,
diff --git a/pybop/models/base_model.py b/pybop/models/base_model.py
index 4ff35685a..20100ab64 100644
--- a/pybop/models/base_model.py
+++ b/pybop/models/base_model.py
@@ -17,7 +17,6 @@ def __init__(self, name="Base Model"):
     def build(
         self,
         dataset=None,
-        experiment=None,
         parameters=None,
         check_model=True,
         init_soc=None,
@@ -29,7 +28,6 @@ def build(
         """
         self.parameters = parameters
         self.dataset = dataset
-        self.experiment = experiment
         if self.parameters is not None:
             self.fit_keys = list(self.parameters.keys())
 
diff --git a/tests/unit/test_cost.py b/tests/unit/test_cost.py
index 8408c992e..ed14655ef 100644
--- a/tests/unit/test_cost.py
+++ b/tests/unit/test_cost.py
@@ -34,9 +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.FittingProblem(
-            parameters, dataset, model=model, signal=signal, x0=x0
-        )
+        problem = pybop.FittingProblem(model, parameters, dataset, signal=signal, x0=x0)
 
         # Base Cost
         base_cost = pybop.BaseCost(problem)
diff --git a/tests/unit/test_optimisation.py b/tests/unit/test_optimisation.py
index 19c899841..269d20b84 100644
--- a/tests/unit/test_optimisation.py
+++ b/tests/unit/test_optimisation.py
@@ -29,10 +29,11 @@ def parameters(self):
 
     @pytest.fixture
     def problem(self, parameters, dataset):
+        model = pybop.lithium_ion.SPM()
         return pybop.FittingProblem(
+            model,
             parameters,
             dataset,
-            model=pybop.lithium_ion.SPM(),
             signal="Terminal voltage [V]",
         )
 
diff --git a/tests/unit/test_parameterisations.py b/tests/unit/test_parameterisations.py
index 1ca77e83b..398f137a3 100644
--- a/tests/unit/test_parameterisations.py
+++ b/tests/unit/test_parameterisations.py
@@ -45,7 +45,7 @@ def test_spm(self, init_soc):
         # Define the cost to optimise
         signal = "Terminal voltage [V]"
         problem = pybop.FittingProblem(
-            parameters, dataset, model=model, signal=signal, init_soc=init_soc
+            model, parameters, dataset, signal=signal, init_soc=init_soc
         )
         cost = pybop.RootMeanSquaredError(problem)
 
@@ -98,7 +98,7 @@ def test_spm_optimisers(self, init_soc):
         # Define the cost to optimise
         signal = "Terminal voltage [V]"
         problem = pybop.FittingProblem(
-            parameters, dataset, model=model, signal=signal, init_soc=init_soc
+            model, parameters, dataset, signal=signal, init_soc=init_soc
         )
         cost = pybop.SumSquaredError(problem)
 
@@ -194,7 +194,7 @@ def test_model_misparameterisation(self, init_soc):
         # Define the cost to optimise
         signal = "Terminal voltage [V]"
         problem = pybop.FittingProblem(
-            parameters, dataset, model=model, signal=signal, init_soc=init_soc
+            model, parameters, dataset, signal=signal, init_soc=init_soc
         )
         cost = pybop.RootMeanSquaredError(problem)
 
diff --git a/tests/unit/test_problem.py b/tests/unit/test_problem.py
index 0908e8186..85d246df3 100644
--- a/tests/unit/test_problem.py
+++ b/tests/unit/test_problem.py
@@ -84,11 +84,11 @@ def test_fitting_problem(self, parameters, dataset, model, signal):
         # Test incorrect number of initial parameter values
         with pytest.raises(ValueError):
             pybop.FittingProblem(
-                parameters, dataset, model, signal=signal, x0=np.array([])
+                model, parameters, dataset, signal=signal, x0=np.array([])
             )
 
         # Construct Problem
-        problem = pybop.FittingProblem(parameters, dataset, model, signal=signal)
+        problem = pybop.FittingProblem(model, parameters, dataset, signal=signal)
 
         assert problem._model == model
         assert problem._model._built_model is not None
@@ -100,10 +100,10 @@ def test_fitting_problem(self, parameters, dataset, model, signal):
     def test_design_problem(self, parameters, experiment, model):
         # Test incorrect number of initial parameter values
         with pytest.raises(ValueError):
-            pybop.DesignProblem(parameters, experiment, model, x0=np.array([]))
+            pybop.DesignProblem(model, parameters, experiment, x0=np.array([]))
 
         # Construct Problem
-        problem = pybop.DesignProblem(parameters, experiment, model)
+        problem = pybop.DesignProblem(model, parameters, experiment)
 
         assert problem._model == model
         assert (

From df71e016750370901c7a93c05da7cadcf12436d6 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Sat, 2 Dec 2023 05:28:09 +0000
Subject: [PATCH 12/13] Add target definition

---
 pybop/_problem.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/pybop/_problem.py b/pybop/_problem.py
index 5d7050757..d6fd455c6 100644
--- a/pybop/_problem.py
+++ b/pybop/_problem.py
@@ -126,6 +126,12 @@ def evaluateS1(self, parameters):
 
         return (np.asarray(y), np.asarray(dy))
 
+    def target(self):
+        """
+        Returns the target dataset.
+        """
+        return self._target
+
 
 class DesignProblem(BaseProblem):
     """

From f46fa6f2fb9964801d123a903d46264595216f62 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Mon, 4 Dec 2023 11:51:46 +0000
Subject: [PATCH 13/13] Change saved parameters type

---
 pybop/_problem.py | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/pybop/_problem.py b/pybop/_problem.py
index d6fd455c6..201164bcd 100644
--- a/pybop/_problem.py
+++ b/pybop/_problem.py
@@ -38,7 +38,6 @@ def __init__(
         # Add the initial values to the parameter definitions
         for i, param in enumerate(self.parameters):
             param.update(value=self.x0[i])
-        self.parameters = {o.name: o.value for o in parameters}
 
     def evaluate(self, parameters):
         """
@@ -97,7 +96,7 @@ def __init__(
         if self._model._built_model is None:
             self._model.build(
                 dataset=self._dataset,
-                parameters=self.parameters,
+                parameters={o.name: o.value for o in self.parameters},
                 check_model=self.check_model,
                 init_soc=self.init_soc,
             )
@@ -116,11 +115,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.parameters):
-            self.parameters[key] = parameters[i]
 
         y, dy = self._model.simulateS1(
-            inputs=self.parameters,
+            inputs=parameters,
             t_eval=self._time_data,
         )
 
@@ -154,12 +151,12 @@ def __init__(
         # Build the model if required
         if experiment is not None:
             # Leave the build until later to apply the experiment
-            self._model.parameters = self.parameters
+            self._model.parameters = {o.name: o.value for o in self.parameters}
 
         elif self._model._built_model is None:
             self._model.build(
                 experiment=self.experiment,
-                parameters=self.parameters,
+                parameters={o.name: o.value for o in self.parameters},
                 check_model=self.check_model,
                 init_soc=self.init_soc,
             )
@@ -178,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.parameters):
-            self.parameters[key] = parameters[i]
 
         y, dy = self._model.simulateS1(
-            inputs=self.parameters,
+            inputs=parameters,
             t_eval=self._time_data,
         )