From 07d2328aeff7bb21a3d7a44d976b46e2bea3a6be Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Wed, 13 Dec 2023 19:19:21 +0000
Subject: [PATCH 01/17] Make Dataset contain a dictionary

---
 examples/notebooks/spm_nlopt.ipynb   | 12 ++++----
 examples/scripts/ecm_CMAES.py        | 13 +++++----
 examples/scripts/spm_CMAES.py        | 14 ++++++----
 examples/scripts/spm_IRPropMin.py    | 13 +++++----
 examples/scripts/spm_SNES.py         | 12 ++++----
 examples/scripts/spm_XNES.py         | 13 +++++----
 examples/scripts/spm_adam.py         | 14 ++++++----
 examples/scripts/spm_descent.py      | 14 ++++++----
 examples/scripts/spm_nlopt.py        | 12 ++++----
 examples/scripts/spm_pso.py          | 13 +++++----
 examples/scripts/spm_scipymin.py     | 12 ++++----
 pybop/_problem.py                    |  6 ++--
 pybop/datasets/base_dataset.py       | 10 ++++---
 pybop/models/base_model.py           |  4 +--
 tests/unit/test_cost.py              | 12 ++++----
 tests/unit/test_optimisation.py      | 12 ++++----
 tests/unit/test_parameterisations.py | 42 ++++++++++++++--------------
 tests/unit/test_problem.py           | 12 ++++----
 18 files changed, 137 insertions(+), 103 deletions(-)

diff --git a/examples/notebooks/spm_nlopt.ipynb b/examples/notebooks/spm_nlopt.ipynb
index b64a2a9ab..baf89804c 100644
--- a/examples/notebooks/spm_nlopt.ipynb
+++ b/examples/notebooks/spm_nlopt.ipynb
@@ -277,11 +277,13 @@
       "outputs": [],
       "source": [
         "pyb_model = pybop.lithium_ion.SPM()\n",
-        "dataset = [\n",
-        "    pybop.Dataset(\"Time [s]\", synthetic_sol[\"Time [s]\"].data),\n",
-        "    pybop.Dataset(\"Current function [A]\", synthetic_sol[\"Current [A]\"].data),\n",
-        "    pybop.Dataset(\"Terminal voltage [V]\", corrupt_V),\n",
-        "]"
+        "dataset = pybop.Dataset(\n",
+        "    {\n",
+        "        \"Time [s]\": synthetic_sol[\"Time [s]\"].data,\n",
+        "        \"Current function [A]\": synthetic_sol[\"Current [A]\"].data,\n",
+        "        \"Terminal voltage [V]\": corrupt_V,\n",
+        "    }\n",
+        ")"
       ]
     },
     {
diff --git a/examples/scripts/ecm_CMAES.py b/examples/scripts/ecm_CMAES.py
index 07f575fcc..d60a8a8c8 100644
--- a/examples/scripts/ecm_CMAES.py
+++ b/examples/scripts/ecm_CMAES.py
@@ -61,11 +61,14 @@
 values = model.predict(t_eval=t_eval)
 corrupt_values = values["Voltage [V]"].data + np.random.normal(0, sigma, len(t_eval))
 
-dataset = [
-    pybop.Dataset("Time [s]", t_eval),
-    pybop.Dataset("Current function [A]", values["Current [A]"].data),
-    pybop.Dataset("Voltage [V]", corrupt_values),
-]
+# 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)
diff --git a/examples/scripts/spm_CMAES.py b/examples/scripts/spm_CMAES.py
index 6e4527b6c..63aa0cc65 100644
--- a/examples/scripts/spm_CMAES.py
+++ b/examples/scripts/spm_CMAES.py
@@ -25,12 +25,14 @@
 values = model.predict(t_eval=t_eval)
 corrupt_values = values["Voltage [V]"].data + np.random.normal(0, sigma, len(t_eval))
 
-# Form dataset for optimisation
-dataset = [
-    pybop.Dataset("Time [s]", t_eval),
-    pybop.Dataset("Current function [A]", values["Current [A]"].data),
-    pybop.Dataset("Voltage [V]", corrupt_values),
-]
+# 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)
diff --git a/examples/scripts/spm_IRPropMin.py b/examples/scripts/spm_IRPropMin.py
index 794bc3c34..b764184d3 100644
--- a/examples/scripts/spm_IRPropMin.py
+++ b/examples/scripts/spm_IRPropMin.py
@@ -24,11 +24,14 @@
 values = model.predict(t_eval=t_eval)
 corrupt_values = values["Voltage [V]"].data + np.random.normal(0, sigma, len(t_eval))
 
-dataset = [
-    pybop.Dataset("Time [s]", t_eval),
-    pybop.Dataset("Current function [A]", values["Current [A]"].data),
-    pybop.Dataset("Voltage [V]", corrupt_values),
-]
+# 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)
diff --git a/examples/scripts/spm_SNES.py b/examples/scripts/spm_SNES.py
index ac24c9494..cc5995a99 100644
--- a/examples/scripts/spm_SNES.py
+++ b/examples/scripts/spm_SNES.py
@@ -24,11 +24,13 @@
 values = model.predict(t_eval=t_eval)
 corrupt_values = values["Voltage [V]"].data + np.random.normal(0, sigma, len(t_eval))
 
-dataset = [
-    pybop.Dataset("Time [s]", t_eval),
-    pybop.Dataset("Current function [A]", values["Current [A]"].data),
-    pybop.Dataset("Voltage [V]", corrupt_values),
-]
+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)
diff --git a/examples/scripts/spm_XNES.py b/examples/scripts/spm_XNES.py
index 81d8c26c7..7f46a352e 100644
--- a/examples/scripts/spm_XNES.py
+++ b/examples/scripts/spm_XNES.py
@@ -24,11 +24,14 @@
 values = model.predict(t_eval=t_eval)
 corrupt_values = values["Voltage [V]"].data + np.random.normal(0, sigma, len(t_eval))
 
-dataset = [
-    pybop.Dataset("Time [s]", t_eval),
-    pybop.Dataset("Current function [A]", values["Current [A]"].data),
-    pybop.Dataset("Voltage [V]", corrupt_values),
-]
+# 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)
diff --git a/examples/scripts/spm_adam.py b/examples/scripts/spm_adam.py
index b8654f2a4..fba51adb3 100644
--- a/examples/scripts/spm_adam.py
+++ b/examples/scripts/spm_adam.py
@@ -25,12 +25,14 @@
 values = model.predict(t_eval=t_eval)
 corrupt_values = values["Voltage [V]"].data + np.random.normal(0, sigma, len(t_eval))
 
-# Dataset definition
-dataset = [
-    pybop.Dataset("Time [s]", t_eval),
-    pybop.Dataset("Current function [A]", values["Current [A]"].data),
-    pybop.Dataset("Voltage [V]", corrupt_values),
-]
+# 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)
diff --git a/examples/scripts/spm_descent.py b/examples/scripts/spm_descent.py
index a1936a953..4f7ddf572 100644
--- a/examples/scripts/spm_descent.py
+++ b/examples/scripts/spm_descent.py
@@ -25,12 +25,14 @@
 values = model.predict(t_eval=t_eval)
 corrupt_values = values["Voltage [V]"].data + np.random.normal(0, sigma, len(t_eval))
 
-# Dataset definition
-dataset = [
-    pybop.Dataset("Time [s]", t_eval),
-    pybop.Dataset("Current function [A]", values["Current [A]"].data),
-    pybop.Dataset("Voltage [V]", corrupt_values),
-]
+# 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)
diff --git a/examples/scripts/spm_nlopt.py b/examples/scripts/spm_nlopt.py
index 258a5256c..31f55bae2 100644
--- a/examples/scripts/spm_nlopt.py
+++ b/examples/scripts/spm_nlopt.py
@@ -3,11 +3,13 @@
 
 # Form dataset
 Measurements = pd.read_csv("examples/scripts/Chen_example.csv", comment="#").to_numpy()
-dataset = [
-    pybop.Dataset("Time [s]", Measurements[:, 0]),
-    pybop.Dataset("Current function [A]", Measurements[:, 1]),
-    pybop.Dataset("Voltage [V]", Measurements[:, 2]),
-]
+dataset = pybop.Dataset(
+    {
+        "Time [s]": Measurements[:, 0],
+        "Current function [A]": Measurements[:, 1],
+        "Voltage [V]": Measurements[:, 2],
+    }
+)
 
 # Define model
 parameter_set = pybop.ParameterSet.pybamm("Chen2020")
diff --git a/examples/scripts/spm_pso.py b/examples/scripts/spm_pso.py
index 07bd9669c..fd55552d1 100644
--- a/examples/scripts/spm_pso.py
+++ b/examples/scripts/spm_pso.py
@@ -24,11 +24,14 @@
 values = model.predict(t_eval=t_eval)
 corrupt_values = values["Voltage [V]"].data + np.random.normal(0, sigma, len(t_eval))
 
-dataset = [
-    pybop.Dataset("Time [s]", t_eval),
-    pybop.Dataset("Current function [A]", values["Current [A]"].data),
-    pybop.Dataset("Voltage [V]", corrupt_values),
-]
+# 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)
diff --git a/examples/scripts/spm_scipymin.py b/examples/scripts/spm_scipymin.py
index ec9173781..797c640a4 100644
--- a/examples/scripts/spm_scipymin.py
+++ b/examples/scripts/spm_scipymin.py
@@ -3,11 +3,13 @@
 
 # Form dataset
 Measurements = pd.read_csv("examples/scripts/Chen_example.csv", comment="#").to_numpy()
-dataset = [
-    pybop.Dataset("Time [s]", Measurements[:, 0]),
-    pybop.Dataset("Current function [A]", Measurements[:, 1]),
-    pybop.Dataset("Voltage [V]", Measurements[:, 2]),
-]
+dataset = pybop.Dataset(
+    {
+        "Time [s]": Measurements[:, 0],
+        "Current function [A]": Measurements[:, 1],
+        "Voltage [V]": Measurements[:, 2],
+    }
+)
 
 # Define model
 parameter_set = pybop.ParameterSet.pybamm("Chen2020")
diff --git a/pybop/_problem.py b/pybop/_problem.py
index 8cf37ec34..1fdaebd98 100644
--- a/pybop/_problem.py
+++ b/pybop/_problem.py
@@ -72,7 +72,7 @@ def __init__(
         if model is not None:
             self._model.signal = signal
         self.signal = signal
-        self._dataset = {o.name: o for o in dataset}
+        self._dataset = dataset.data
         self.n_outputs = len([self.signal])
 
         # Check that the dataset contains time and current
@@ -80,9 +80,9 @@ def __init__(
             if name not in self._dataset:
                 raise ValueError(f"expected {name} in list of dataset")
 
-        self._time_data = self._dataset["Time [s]"].data
+        self._time_data = self._dataset["Time [s]"]
         self.n_time_data = len(self._time_data)
-        self._target = self._dataset[signal].data
+        self._target = self._dataset[signal]
 
         if np.any(self._time_data < 0):
             raise ValueError("Times can not be negative.")
diff --git a/pybop/datasets/base_dataset.py b/pybop/datasets/base_dataset.py
index ed194ae48..7ac2c8689 100644
--- a/pybop/datasets/base_dataset.py
+++ b/pybop/datasets/base_dataset.py
@@ -6,12 +6,14 @@ class Dataset:
     Class for experimental observations.
     """
 
-    def __init__(self, name, data):
-        self.name = name
-        self.data = data
+    def __init__(self, data_dictionary):
+        if not isinstance(data_dictionary, dict):
+            raise ValueError("The input to pybop.Dataset must be a dictionary.")
+        self.data = data_dictionary
+        self.names = self.data.keys()
 
     def __repr__(self):
-        return f"Dataset: {self.name} \n Data: {self.data}"
+        return f"Dataset: {type(self.data)} \n Contains: {self.names}"
 
     def Interpolant(self):
         if self.variable == "time":
diff --git a/pybop/models/base_model.py b/pybop/models/base_model.py
index 848ab029b..637cf4937 100644
--- a/pybop/models/base_model.py
+++ b/pybop/models/base_model.py
@@ -86,8 +86,8 @@ def set_params(self):
         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,
-                    self.dataset["Current function [A]"].data,
+                    self.dataset["Time [s]"],
+                    self.dataset["Current function [A]"],
                     pybamm.t,
                 )
                 # Set t_eval
diff --git a/tests/unit/test_cost.py b/tests/unit/test_cost.py
index ed14655ef..26516685f 100644
--- a/tests/unit/test_cost.py
+++ b/tests/unit/test_cost.py
@@ -25,11 +25,13 @@ def test_costs(self, cut_off):
         # Form dataset
         x0 = np.array([0.52])
         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),
-        ]
+        dataset = pybop.Dataset(
+            {
+                "Time [s]": solution["Time [s]"].data,
+                "Current function [A]": solution["Current [A]"].data,
+                "Voltage [V]": solution["Terminal voltage [V]"].data,
+            }
+        )
 
         # Construct Problem
         signal = "Voltage [V]"
diff --git a/tests/unit/test_optimisation.py b/tests/unit/test_optimisation.py
index dc96367cf..16941fda7 100644
--- a/tests/unit/test_optimisation.py
+++ b/tests/unit/test_optimisation.py
@@ -11,11 +11,13 @@ class TestOptimisation:
 
     @pytest.fixture
     def dataset(self):
-        return [
-            pybop.Dataset("Time [s]", np.linspace(0, 360, 10)),
-            pybop.Dataset("Current function [A]", np.zeros(10)),
-            pybop.Dataset("Terminal voltage [V]", np.ones(10)),
-        ]
+        return pybop.Dataset(
+            {
+                "Time [s]": np.linspace(0, 360, 10),
+                "Current function [A]": np.zeros(10),
+                "Terminal voltage [V]": np.ones(10),
+            }
+        )
 
     @pytest.fixture
     def parameters(self):
diff --git a/tests/unit/test_parameterisations.py b/tests/unit/test_parameterisations.py
index af1a4d573..db80d601c 100644
--- a/tests/unit/test_parameterisations.py
+++ b/tests/unit/test_parameterisations.py
@@ -38,13 +38,13 @@ def x0(self):
     def test_spm(self, parameters, model, x0, init_soc):
         # Form dataset
         solution = self.getdata(model, x0, init_soc)
-        dataset = [
-            pybop.Dataset("Time [s]", solution["Time [s]"].data),
-            pybop.Dataset("Current function [A]", solution["Current [A]"].data),
-            pybop.Dataset(
-                "Terminal voltage [V]", solution["Terminal voltage [V]"].data
-            ),
-        ]
+        dataset = pybop.Dataset(
+            {
+                "Time [s]": solution["Time [s]"].data,
+                "Current function [A]": solution["Current [A]"].data,
+                "Terminal voltage [V]": solution["Terminal voltage [V]"].data,
+            }
+        )
 
         # Define the cost to optimise
         signal = "Terminal voltage [V]"
@@ -71,13 +71,13 @@ def spm_cost(self, parameters, model, x0):
         # Form dataset
         init_soc = 0.5
         solution = self.getdata(model, x0, init_soc)
-        dataset = [
-            pybop.Dataset("Time [s]", solution["Time [s]"].data),
-            pybop.Dataset("Current function [A]", solution["Current [A]"].data),
-            pybop.Dataset(
-                "Terminal voltage [V]", solution["Terminal voltage [V]"].data
-            ),
-        ]
+        dataset = pybop.Dataset(
+            {
+                "Time [s]": solution["Time [s]"].data,
+                "Current function [A]": solution["Current [A]"].data,
+                "Terminal voltage [V]": solution["Terminal voltage [V]"].data,
+            }
+        )
 
         # Define the cost to optimise
         signal = "Terminal voltage [V]"
@@ -151,13 +151,13 @@ def test_model_misparameterisation(self, parameters, model, x0, init_soc):
 
         # Form dataset
         solution = self.getdata(second_model, x0, init_soc)
-        dataset = [
-            pybop.Dataset("Time [s]", solution["Time [s]"].data),
-            pybop.Dataset("Current function [A]", solution["Current [A]"].data),
-            pybop.Dataset(
-                "Terminal voltage [V]", solution["Terminal voltage [V]"].data
-            ),
-        ]
+        dataset = pybop.Dataset(
+            {
+                "Time [s]": solution["Time [s]"].data,
+                "Current function [A]": solution["Current [A]"].data,
+                "Terminal voltage [V]": solution["Terminal voltage [V]"].data,
+            }
+        )
 
         # Define the cost to optimise
         signal = "Terminal voltage [V]"
diff --git a/tests/unit/test_problem.py b/tests/unit/test_problem.py
index 85d246df3..db2c3e278 100644
--- a/tests/unit/test_problem.py
+++ b/tests/unit/test_problem.py
@@ -53,11 +53,13 @@ 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),
-        ]
+        return pybop.Dataset(
+            {
+                "Time [s]": solution["Time [s]"].data,
+                "Current function [A]": solution["Current [A]"].data,
+                "Voltage [V]": solution["Terminal voltage [V]"].data,
+            }
+        )
 
     @pytest.fixture
     def signal(self):

From ef367646bad3184dbbb21325ecedaf20183b527a Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 14 Dec 2023 09:43:24 +0000
Subject: [PATCH 02/17] Extend signal to a list of strings

---
 pybop/_problem.py            | 26 ++++++++++++++++++--------
 pybop/models/base_model.py   | 27 +++++++++++++++------------
 pybop/plotting/quick_plot.py | 33 +++++++++++++++++----------------
 3 files changed, 50 insertions(+), 36 deletions(-)

diff --git a/pybop/_problem.py b/pybop/_problem.py
index 1fdaebd98..c5024a60f 100644
--- a/pybop/_problem.py
+++ b/pybop/_problem.py
@@ -69,28 +69,38 @@ def __init__(
         x0=None,
     ):
         super().__init__(parameters, model, check_model, init_soc, x0)
-        if model is not None:
-            self._model.signal = signal
+        if isinstance(signal, str):
+            signal = [signal]
         self.signal = signal
+        self.n_outputs = len(self.signal)
         self._dataset = dataset.data
-        self.n_outputs = len([self.signal])
 
         # Check that the dataset contains time and current
-        for name in ["Time [s]", "Current function [A]", signal]:
+        for name in ["Time [s]", "Current function [A]"] + signal:
             if name not in self._dataset:
                 raise ValueError(f"expected {name} in list of dataset")
 
         self._time_data = self._dataset["Time [s]"]
         self.n_time_data = len(self._time_data)
-        self._target = self._dataset[signal]
-
         if np.any(self._time_data < 0):
             raise ValueError("Times can not be negative.")
         if np.any(self._time_data[:-1] >= self._time_data[1:]):
             raise ValueError("Times must be increasing.")
 
-        if len(self._target) != len(self._time_data):
-            raise ValueError("Time data and signal data must be the same length.")
+        target = [self._dataset[signal] for signal in self.signal]
+        self._target = np.vstack(target).T
+        if self.n_outputs == 1:
+            if len(self._target) != self.n_time_data:
+                raise ValueError("Time data and target data must be the same length.")
+        else:
+            if self._target.shape != (self.n_time_data, self.n_outputs):
+                raise ValueError("Time data and target data must be the same shape.")
+
+        # Add useful parameters to model
+        if model is not None:
+            self._model.signal = self.signal
+            self._model.n_outputs = self.n_outputs
+            self._model.n_time_data = self.n_time_data
 
         # Build the model
         if self._model._built_model is None:
diff --git a/pybop/models/base_model.py b/pybop/models/base_model.py
index 637cf4937..0f4a39623 100644
--- a/pybop/models/base_model.py
+++ b/pybop/models/base_model.py
@@ -111,9 +111,11 @@ def simulate(self, inputs, t_eval):
             if not isinstance(inputs, dict):
                 inputs = {key: inputs[i] for i, key in enumerate(self.fit_keys)}
 
-            return self.solver.solve(self.built_model, inputs=inputs, t_eval=t_eval)[
-                self.signal
-            ].data
+            sol = self.solver.solve(self.built_model, inputs=inputs, t_eval=t_eval)
+
+            predictions = [sol[signal].data for signal in self.signal]
+
+            return np.vstack(predictions).T
 
     def simulateS1(self, inputs, t_eval):
         """
@@ -134,15 +136,16 @@ def simulateS1(self, inputs, t_eval):
                 calculate_sensitivities=True,
             )
 
-            return (
-                sol[self.signal].data,
-                np.asarray(
-                    [
-                        sol[self.signal].sensitivities[key].toarray()
-                        for key in self.fit_keys
-                    ]
-                ).T,
-            )
+            predictions = [sol[signal].data for signal in self.signal]
+
+            sensitivities = [
+                np.array(
+                    [[sol[signal].sensitivities[key]] for signal in self.signal]
+                ).reshape(len(sol[self.signal[0]].data), self.n_outputs)
+                for key in self.fit_keys
+            ]
+
+            return (np.vstack(predictions).T, np.dstack(sensitivities))
 
     def predict(
         self,
diff --git a/pybop/plotting/quick_plot.py b/pybop/plotting/quick_plot.py
index d6628760b..b6f26c755 100644
--- a/pybop/plotting/quick_plot.py
+++ b/pybop/plotting/quick_plot.py
@@ -199,21 +199,22 @@ def quick_plot(params, cost, title="Scatter Plot", width=1024, height=576):
     model_output = cost.problem.evaluate(params)
     target_output = cost.problem.target()
 
-    # Create the figure using the StandardPlot class
-    fig = pybop.StandardPlot(
-        x=time_data,
-        y=model_output,
-        cost=cost,
-        y2=target_output,
-        xaxis_title="Time [s]",
-        yaxis_title=cost.problem.signal,
-        title=title,
-        trace_name="Model",
-        width=width,
-        height=height,
-    )()
-
-    # Display the figure
-    fig.show()
+    for i in range(0, cost.problem.n_outputs):
+        # Create the figure using the StandardPlot class
+        fig = pybop.StandardPlot(
+            x=time_data,
+            y=model_output[:, i],
+            cost=cost,
+            y2=target_output[:, i],
+            xaxis_title="Time [s]",
+            yaxis_title=cost.problem.signal[i],
+            title=title,
+            trace_name="Model",
+            width=width,
+            height=height,
+        )()
+
+        # Display the figure
+        fig.show()
 
     return fig

From 61fd9b8f7b402557c8e983f4a57313b3544498f5 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 14 Dec 2023 10:50:42 +0000
Subject: [PATCH 03/17] Align with set_max_iterations

---
 pybop/optimisers/nlopt_optimize.py   |  6 +++---
 pybop/optimisers/scipy_optimisers.py | 17 ++++++++++-------
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/pybop/optimisers/nlopt_optimize.py b/pybop/optimisers/nlopt_optimize.py
index 7d78a699f..93bb59e87 100644
--- a/pybop/optimisers/nlopt_optimize.py
+++ b/pybop/optimisers/nlopt_optimize.py
@@ -11,7 +11,7 @@ class NLoptOptimize(BaseOptimiser):
     def __init__(self, n_param, xtol=None, method=None, maxiter=None):
         super().__init__()
         self.n_param = n_param
-        self.maxiter = maxiter
+        self._max_iterations = maxiter
 
         if method is not None:
             self.optim = nlopt.opt(method, self.n_param)
@@ -48,8 +48,8 @@ def cost_wrapper(x, grad):
         self.optim.set_upper_bounds(bounds["upper"])
 
         # Set max iterations
-        if self.maxiter is not None:
-            self.optim.set_maxeval(self.maxiter)
+        if self._max_iterations is not None:
+            self.optim.set_maxeval(self._max_iterations)
 
         # Run the optimser
         x = self.optim.optimize(x0)
diff --git a/pybop/optimisers/scipy_optimisers.py b/pybop/optimisers/scipy_optimisers.py
index 59f9e6388..842fc05d9 100644
--- a/pybop/optimisers/scipy_optimisers.py
+++ b/pybop/optimisers/scipy_optimisers.py
@@ -11,11 +11,8 @@ def __init__(self, method=None, bounds=None, maxiter=None):
         super().__init__()
         self.method = method
         self.bounds = bounds
-        self.maxiter = maxiter
-        if self.maxiter is not None:
-            self.options = {"maxiter": self.maxiter}
-        else:
-            self.options = {}
+        self.options = {}
+        self._max_iterations = maxiter
 
         if self.method is None:
             self.method = "COBYLA"  # "L-BFGS-B"
@@ -44,6 +41,12 @@ def callback(x):
                 (lower, upper) for lower, upper in zip(bounds["lower"], bounds["upper"])
             )
 
+        # Set max iterations
+        if self._max_iterations is not None:
+            self.options = {"maxiter": self._max_iterations}
+        else:
+            self.options.pop("maxiter", None)
+
         output = minimize(
             cost_function,
             x0,
@@ -81,7 +84,7 @@ def __init__(self, bounds=None, strategy="best1bin", maxiter=1000, popsize=15):
         super().__init__()
         self.bounds = bounds
         self.strategy = strategy
-        self.maxiter = maxiter
+        self._max_iterations = maxiter
         self.popsize = popsize
 
     def _runoptimise(self, cost_function, x0=None, bounds=None):
@@ -124,7 +127,7 @@ def callback(x, convergence):
             cost_function,
             bounds,
             strategy=self.strategy,
-            maxiter=self.maxiter,
+            maxiter=self._max_iterations,
             popsize=self.popsize,
             callback=callback,
         )

From dc5ec7e8fc52bce63780d6c73f218fe14baeea78 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 14 Dec 2023 10:51:54 +0000
Subject: [PATCH 04/17] Turn signal into list

---
 examples/scripts/spm_nlopt.py    | 2 +-
 examples/scripts/spm_scipymin.py | 2 +-
 pybop/_problem.py                | 2 +-
 tests/unit/test_cost.py          | 2 +-
 tests/unit/test_optimisation.py  | 2 +-
 5 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/examples/scripts/spm_nlopt.py b/examples/scripts/spm_nlopt.py
index 31f55bae2..7ba0acd8b 100644
--- a/examples/scripts/spm_nlopt.py
+++ b/examples/scripts/spm_nlopt.py
@@ -32,7 +32,7 @@
 ]
 
 # Define the cost to optimise
-signal = "Voltage [V]"
+signal = ["Voltage [V]"]
 problem = pybop.FittingProblem(model, parameters, dataset, signal=signal, init_soc=0.98)
 cost = pybop.RootMeanSquaredError(problem)
 
diff --git a/examples/scripts/spm_scipymin.py b/examples/scripts/spm_scipymin.py
index 797c640a4..af81b6c2f 100644
--- a/examples/scripts/spm_scipymin.py
+++ b/examples/scripts/spm_scipymin.py
@@ -32,7 +32,7 @@
 ]
 
 # Define the cost to optimise
-signal = "Voltage [V]"
+signal = ["Voltage [V]"]
 problem = pybop.FittingProblem(model, parameters, dataset, signal=signal, init_soc=0.98)
 cost = pybop.RootMeanSquaredError(problem)
 
diff --git a/pybop/_problem.py b/pybop/_problem.py
index c5024a60f..60d0f0361 100644
--- a/pybop/_problem.py
+++ b/pybop/_problem.py
@@ -63,7 +63,7 @@ def __init__(
         model,
         parameters,
         dataset,
-        signal="Voltage [V]",
+        signal=["Voltage [V]"],
         check_model=True,
         init_soc=None,
         x0=None,
diff --git a/tests/unit/test_cost.py b/tests/unit/test_cost.py
index 26516685f..362499482 100644
--- a/tests/unit/test_cost.py
+++ b/tests/unit/test_cost.py
@@ -34,7 +34,7 @@ def test_costs(self, cut_off):
         )
 
         # Construct Problem
-        signal = "Voltage [V]"
+        signal = ["Voltage [V]"]
         model.parameter_set.update({"Lower voltage cut-off [V]": cut_off})
         problem = pybop.FittingProblem(model, parameters, dataset, signal=signal, x0=x0)
 
diff --git a/tests/unit/test_optimisation.py b/tests/unit/test_optimisation.py
index 16941fda7..8e7edc347 100644
--- a/tests/unit/test_optimisation.py
+++ b/tests/unit/test_optimisation.py
@@ -36,7 +36,7 @@ def problem(self, parameters, dataset):
             model,
             parameters,
             dataset,
-            signal="Terminal voltage [V]",
+            signal=["Terminal voltage [V]"],
         )
 
     @pytest.fixture

From 7b0949af4fb844ca9a340a086cb187cba64ae775 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 14 Dec 2023 10:53:43 +0000
Subject: [PATCH 05/17] Parameterise and test multiple signals

---
 tests/unit/test_parameterisations.py | 129 +++++++++++++++++----------
 1 file changed, 81 insertions(+), 48 deletions(-)

diff --git a/tests/unit/test_parameterisations.py b/tests/unit/test_parameterisations.py
index db80d601c..aa2c637e0 100644
--- a/tests/unit/test_parameterisations.py
+++ b/tests/unit/test_parameterisations.py
@@ -47,7 +47,7 @@ def test_spm(self, parameters, model, x0, init_soc):
         )
 
         # Define the cost to optimise
-        signal = "Terminal voltage [V]"
+        signal = ["Terminal voltage [V]"]
         problem = pybop.FittingProblem(
             model, parameters, dataset, signal=signal, init_soc=init_soc
         )
@@ -80,66 +80,99 @@ def spm_cost(self, parameters, model, x0):
         )
 
         # Define the cost to optimise
-        signal = "Terminal voltage [V]"
+        signal = ["Terminal voltage [V]"]
         problem = pybop.FittingProblem(
             model, parameters, dataset, signal=signal, init_soc=init_soc
         )
         return pybop.SumSquaredError(problem)
 
-    @pytest.mark.unit
-    def test_spm_optimisers(self, spm_cost, x0):
-        # Select optimisers
-        optimisers = [
+    @pytest.mark.parametrize(
+        "optimiser",
+        [
             pybop.NLoptOptimize,
             pybop.SciPyMinimize,
             pybop.SciPyDifferentialEvolution,
-            pybop.CMAES,
             pybop.Adam,
+            pybop.CMAES,
             pybop.GradientDescent,
+            pybop.IRPropMin,
             pybop.PSO,
-            pybop.XNES,
             pybop.SNES,
-            pybop.IRPropMin,
-        ]
+            pybop.XNES,
+        ],
+    )
+    @pytest.mark.unit
+    def test_spm_optimisers(self, optimiser, spm_cost, x0):
+        # Test each optimiser
+        parameterisation = pybop.Optimisation(cost=spm_cost, optimiser=optimiser)
+        parameterisation.set_max_unchanged_iterations(iterations=15, threshold=5e-4)
+
+        if optimiser in [pybop.CMAES]:
+            parameterisation.set_f_guessed_tracking(True)
+            assert parameterisation._use_f_guessed is True
+            parameterisation.set_max_iterations(1)
+            x, final_cost = parameterisation.run()
+
+            parameterisation.set_f_guessed_tracking(False)
+            parameterisation.set_max_iterations(100)
+
+            x, final_cost = parameterisation.run()
+            assert parameterisation._max_iterations == 100
+
+        elif optimiser in [pybop.GradientDescent]:
+            parameterisation.optimiser.set_learning_rate(0.025)
+            parameterisation.set_max_iterations(100)
+            x, final_cost = parameterisation.run()
+
+        else:
+            parameterisation.set_max_iterations(100)
+            x, final_cost = parameterisation.run()
+
+        # Assertions
+        np.testing.assert_allclose(final_cost, 0, atol=1e-2)
+        np.testing.assert_allclose(x, x0, atol=1e-1)
+
+    @pytest.fixture
+    def spm_cost_2(self, parameters, model, x0):
+        # Form dataset
+        init_soc = 0.5
+        solution = self.getdata(model, x0, init_soc)
+        dataset = pybop.Dataset(
+            {
+                "Time [s]": solution["Time [s]"].data,
+                "Current function [A]": solution["Current [A]"].data,
+                "Terminal voltage [V]": solution["Terminal voltage [V]"].data,
+            }
+        )
 
+        # Define the cost to optimise
+        signal = ["Terminal voltage [V]", "Time [s]"]
+        problem = pybop.FittingProblem(
+            model, parameters, dataset, signal=signal, init_soc=init_soc
+        )
+        return pybop.SumSquaredError(problem)
+
+    @pytest.mark.parametrize(
+        "optimiser",
+        [
+            pybop.NLoptOptimize,
+            pybop.SciPyMinimize,
+            pybop.Adam,
+            pybop.CMAES,
+        ],
+    )
+    @pytest.mark.unit
+    def test_multiple_signals(self, optimiser, spm_cost_2, x0):
         # Test each optimiser
-        for optimiser in optimisers:
-            parameterisation = pybop.Optimisation(cost=spm_cost, optimiser=optimiser)
-            parameterisation.set_max_unchanged_iterations(iterations=15, threshold=5e-4)
-
-            if optimiser in [pybop.CMAES]:
-                parameterisation.set_f_guessed_tracking(True)
-                assert parameterisation._use_f_guessed is True
-                parameterisation.set_max_iterations(1)
-                x, final_cost = parameterisation.run()
-
-                parameterisation.set_f_guessed_tracking(False)
-                parameterisation.set_max_iterations(100)
-
-                x, final_cost = parameterisation.run()
-                assert parameterisation._max_iterations == 100
-
-            elif optimiser in [pybop.GradientDescent]:
-                parameterisation.optimiser.set_learning_rate(0.025)
-                parameterisation.set_max_iterations(100)
-                x, final_cost = parameterisation.run()
-
-            elif optimiser in [
-                pybop.PSO,
-                pybop.XNES,
-                pybop.SNES,
-                pybop.Adam,
-                pybop.IRPropMin,
-            ]:
-                parameterisation.set_max_iterations(100)
-                x, final_cost = parameterisation.run()
-
-            else:
-                x, final_cost = parameterisation.run()
-
-            # Assertions
-            np.testing.assert_allclose(final_cost, 0, atol=1e-2)
-            np.testing.assert_allclose(x, x0, atol=1e-1)
+        parameterisation = pybop.Optimisation(cost=spm_cost_2, optimiser=optimiser)
+        parameterisation.set_max_unchanged_iterations(iterations=15, threshold=5e-4)
+        parameterisation.set_max_iterations(100)
+
+        x, final_cost = parameterisation.run()
+
+        # Assertions
+        np.testing.assert_allclose(final_cost, 0, atol=1e-2)
+        np.testing.assert_allclose(x, x0, atol=1e-1)
 
     @pytest.mark.parametrize("init_soc", [0.3, 0.7])
     @pytest.mark.unit
@@ -160,7 +193,7 @@ def test_model_misparameterisation(self, parameters, model, x0, init_soc):
         )
 
         # Define the cost to optimise
-        signal = "Terminal voltage [V]"
+        signal = ["Terminal voltage [V]"]
         problem = pybop.FittingProblem(
             model, parameters, dataset, signal=signal, init_soc=init_soc
         )

From 2e9bc8f821178931b507751043c314a6be90297a Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 14 Dec 2023 11:03:27 +0000
Subject: [PATCH 06/17] Update CHANGELOG.md

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 164050d07..2da53a9e0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@
 - [#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
 - [#120](https://github.com/pybop-team/PyBOP/issues/120) - Updates the parameterisation test settings including the number of iterations
+- [#145](https://github.com/pybop-team/PyBOP/issues/145) - Reformats Dataset to contain a dictionary and signal into a list of strings
 
 ## Bug Fixes
 

From 5f9a90244c14ad2e1dfcdf2549209e64528b204c Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 14 Dec 2023 11:14:09 +0000
Subject: [PATCH 07/17] Relax test_multiple_signals assertion

---
 tests/unit/test_parameterisations.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unit/test_parameterisations.py b/tests/unit/test_parameterisations.py
index aa2c637e0..b83537116 100644
--- a/tests/unit/test_parameterisations.py
+++ b/tests/unit/test_parameterisations.py
@@ -171,7 +171,7 @@ def test_multiple_signals(self, optimiser, spm_cost_2, x0):
         x, final_cost = parameterisation.run()
 
         # Assertions
-        np.testing.assert_allclose(final_cost, 0, atol=1e-2)
+        np.testing.assert_allclose(final_cost, 0, atol=2e-2)
         np.testing.assert_allclose(x, x0, atol=1e-1)
 
     @pytest.mark.parametrize("init_soc", [0.3, 0.7])

From 69563bc80266df99a6b002167e0cbb37b903b02f Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 14 Dec 2023 11:45:55 +0000
Subject: [PATCH 08/17] Create test_dataset.py

---
 tests/unit/test_dataset.py | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)
 create mode 100644 tests/unit/test_dataset.py

diff --git a/tests/unit/test_dataset.py b/tests/unit/test_dataset.py
new file mode 100644
index 000000000..fcb00c220
--- /dev/null
+++ b/tests/unit/test_dataset.py
@@ -0,0 +1,36 @@
+import pytest
+import pybop
+import numpy as np
+
+
+class TestDataset:
+    """
+    Class to test dataset construction.
+    """
+
+    @pytest.mark.unit
+    def test_dataset(self):
+        # Construct and simulate model
+        model = pybop.lithium_ion.SPM()
+        model.parameter_set = model.pybamm_model.default_parameter_values
+        solution = model.predict(t_eval=np.linspace(0, 10, 100))
+
+        # Form dataset
+        data_dictionary = {
+            "Time [s]": solution["Time [s]"].data,
+            "Current function [A]": solution["Current [A]"].data,
+            "Voltage [V]": solution["Terminal voltage [V]"].data,
+        }
+        dataset = pybop.Dataset(data_dictionary)         
+
+        # Test repr
+        print(dataset)
+
+        # Test data structure
+        assert dataset.data == data_dictionary
+        
+        # Test exception for non-dictionary inputs
+        with pytest.raises(ValueError):
+            pybop.Dataset(["StringInputShouldNotWork"])
+        with pytest.raises(ValueError):
+            pybop.Dataset(solution["Time [s]"].data)

From 4695d3bd6e639ac1f96518a130067eb012fff51a Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 14 Dec 2023 11:46:18 +0000
Subject: [PATCH 09/17] style: pre-commit fixes

---
 tests/unit/test_dataset.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/unit/test_dataset.py b/tests/unit/test_dataset.py
index fcb00c220..4715b6c36 100644
--- a/tests/unit/test_dataset.py
+++ b/tests/unit/test_dataset.py
@@ -21,14 +21,14 @@ def test_dataset(self):
             "Current function [A]": solution["Current [A]"].data,
             "Voltage [V]": solution["Terminal voltage [V]"].data,
         }
-        dataset = pybop.Dataset(data_dictionary)         
+        dataset = pybop.Dataset(data_dictionary)
 
         # Test repr
         print(dataset)
 
         # Test data structure
         assert dataset.data == data_dictionary
-        
+
         # Test exception for non-dictionary inputs
         with pytest.raises(ValueError):
             pybop.Dataset(["StringInputShouldNotWork"])

From fb17d00cbeabee6f6b48973708c9468258e1b1ad Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 14 Dec 2023 12:06:32 +0000
Subject: [PATCH 10/17] Update Dataset to accept PyBaMM Solution

---
 pybop/datasets/base_dataset.py | 2 ++
 tests/unit/test_dataset.py     | 8 ++++++--
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/pybop/datasets/base_dataset.py b/pybop/datasets/base_dataset.py
index 7ac2c8689..e56c88716 100644
--- a/pybop/datasets/base_dataset.py
+++ b/pybop/datasets/base_dataset.py
@@ -7,6 +7,8 @@ class Dataset:
     """
 
     def __init__(self, data_dictionary):
+        if isinstance(data_dictionary, pybamm.solvers.solution.Solution):
+            data_dictionary = data_dictionary.get_data_dict()
         if not isinstance(data_dictionary, dict):
             raise ValueError("The input to pybop.Dataset must be a dictionary.")
         self.data = data_dictionary
diff --git a/tests/unit/test_dataset.py b/tests/unit/test_dataset.py
index 4715b6c36..b5cfa96dc 100644
--- a/tests/unit/test_dataset.py
+++ b/tests/unit/test_dataset.py
@@ -18,8 +18,8 @@ def test_dataset(self):
         # Form dataset
         data_dictionary = {
             "Time [s]": solution["Time [s]"].data,
-            "Current function [A]": solution["Current [A]"].data,
-            "Voltage [V]": solution["Terminal voltage [V]"].data,
+            "Current [A]": solution["Current [A]"].data,
+            "Terminal voltage [V]": solution["Terminal voltage [V]"].data,
         }
         dataset = pybop.Dataset(data_dictionary)
 
@@ -34,3 +34,7 @@ def test_dataset(self):
             pybop.Dataset(["StringInputShouldNotWork"])
         with pytest.raises(ValueError):
             pybop.Dataset(solution["Time [s]"].data)
+
+        # Test conversion of pybamm solution into dictionary
+        assert dataset.data == pybop.Dataset(solution).data
+        assert dataset.names == pybop.Dataset(solution).names

From 6184b2afe9023e723b16adafcca90bfc0fbd31c6 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 14 Dec 2023 12:31:04 +0000
Subject: [PATCH 11/17] Update test_problem.py

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

diff --git a/tests/unit/test_problem.py b/tests/unit/test_problem.py
index db2c3e278..b1ba7579c 100644
--- a/tests/unit/test_problem.py
+++ b/tests/unit/test_problem.py
@@ -98,6 +98,22 @@ def test_fitting_problem(self, parameters, dataset, model, signal):
         # Test model.simulate
         model.simulate(inputs=[0.5, 0.5], t_eval=np.linspace(0, 10, 100))
 
+        # Test problem construction errors
+        for bad_dataset in [
+            pybop.Dataset({"Time [s]": np.array([0])}),
+            pybop.Dataset(
+                {"Time [s]": np.array([-1]), "Current function [A]": np.array([0])}
+            ),
+            pybop.Dataset(
+                {"Time [s]": np.array([1, 0]), "Current function [A]": np.array([0, 0])}
+            ),
+            pybop.Dataset(
+                {"Time [s]": np.array([0]), "Current function [A]": np.array([0, 0])}
+            ),
+        ]:
+            with pytest.raises(ValueError):
+                pybop.FittingProblem(model, parameters, bad_dataset, signal=signal)
+
     @pytest.mark.unit
     def test_design_problem(self, parameters, experiment, model):
         # Test incorrect number of initial parameter values

From c5fc469cf32e7f69e50406036450f4cdb2b95fab Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 14 Dec 2023 13:02:04 +0000
Subject: [PATCH 12/17] Update test_problem.py

---
 tests/unit/test_problem.py | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/tests/unit/test_problem.py b/tests/unit/test_problem.py
index b1ba7579c..98be1f602 100644
--- a/tests/unit/test_problem.py
+++ b/tests/unit/test_problem.py
@@ -102,13 +102,25 @@ def test_fitting_problem(self, parameters, dataset, model, signal):
         for bad_dataset in [
             pybop.Dataset({"Time [s]": np.array([0])}),
             pybop.Dataset(
-                {"Time [s]": np.array([-1]), "Current function [A]": np.array([0])}
+                {
+                    "Time [s]": np.array([-1]),
+                    "Current function [A]": np.array([0]),
+                    "Voltage [V]": np.array([0]),
+                }
             ),
             pybop.Dataset(
-                {"Time [s]": np.array([1, 0]), "Current function [A]": np.array([0, 0])}
+                {
+                    "Time [s]": np.array([1, 0]),
+                    "Current function [A]": np.array([0, 0]),
+                    "Voltage [V]": np.array([0, 0]),
+                }
             ),
             pybop.Dataset(
-                {"Time [s]": np.array([0]), "Current function [A]": np.array([0, 0])}
+                {
+                    "Time [s]": np.array([0]),
+                    "Current function [A]": np.array([0, 0]),
+                    "Voltage [V]": np.array([0, 0]),
+                }
             ),
         ]:
             with pytest.raises(ValueError):

From d00b4c637cc0d20d781cc76187c6f1ced550ab7f Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 14 Dec 2023 15:07:14 +0000
Subject: [PATCH 13/17] Merge changes to Dataset

---
 pybop/_dataset.py              | 18 ++++++++++--------
 pybop/datasets/base_dataset.py | 24 ------------------------
 2 files changed, 10 insertions(+), 32 deletions(-)
 delete mode 100644 pybop/datasets/base_dataset.py

diff --git a/pybop/_dataset.py b/pybop/_dataset.py
index 9a5f66506..1263ace3e 100644
--- a/pybop/_dataset.py
+++ b/pybop/_dataset.py
@@ -18,20 +18,22 @@ class Dataset:
 
     """
 
-    def __init__(self, name, data):
+    def __init__(self, data_dictionary):
         """
         Initialize a Dataset instance with a name and data.
 
         Parameters
         ----------
-        name : str
-            The name for the dataset.
-        data : array-like
+        data_dictionary : dict or instance of pybamm.solvers.solution.Solution
             The experimental data to store within the dataset.
         """
 
-        self.name = name
-        self.data = data
+        if isinstance(data_dictionary, pybamm.solvers.solution.Solution):
+            data_dictionary = data_dictionary.get_data_dict()
+        if not isinstance(data_dictionary, dict):
+            raise ValueError("The input to pybop.Dataset must be a dictionary.")
+        self.data = data_dictionary
+        self.names = self.data.keys()
 
     def __repr__(self):
         """
@@ -40,9 +42,9 @@ def __repr__(self):
         Returns
         -------
         str
-            A string that includes the name and data of the dataset.
+            A string that includes the type and contents of the dataset.
         """
-        return f"Dataset: {self.name} \n Data: {self.data}"
+        return f"Dataset: {type(self.data)} \n Contains: {self.names}"
 
     def Interpolant(self):
         """
diff --git a/pybop/datasets/base_dataset.py b/pybop/datasets/base_dataset.py
deleted file mode 100644
index e56c88716..000000000
--- a/pybop/datasets/base_dataset.py
+++ /dev/null
@@ -1,24 +0,0 @@
-import pybamm
-
-
-class Dataset:
-    """
-    Class for experimental observations.
-    """
-
-    def __init__(self, data_dictionary):
-        if isinstance(data_dictionary, pybamm.solvers.solution.Solution):
-            data_dictionary = data_dictionary.get_data_dict()
-        if not isinstance(data_dictionary, dict):
-            raise ValueError("The input to pybop.Dataset must be a dictionary.")
-        self.data = data_dictionary
-        self.names = self.data.keys()
-
-    def __repr__(self):
-        return f"Dataset: {type(self.data)} \n Contains: {self.names}"
-
-    def Interpolant(self):
-        if self.variable == "time":
-            self.Interpolant = pybamm.Interpolant(self.x, self.y, pybamm.t)
-        else:
-            NotImplementedError("Only time interpolation is supported")

From fbdbae3b47da32a9eec8c56950c0abb7701cf209 Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 14 Dec 2023 16:19:21 +0000
Subject: [PATCH 14/17] Move signal to BaseProblem

---
 pybop/_problem.py | 66 +++++++++++++++++++++++++----------------------
 1 file changed, 35 insertions(+), 31 deletions(-)

diff --git a/pybop/_problem.py b/pybop/_problem.py
index d49ba8294..df9da0b4f 100644
--- a/pybop/_problem.py
+++ b/pybop/_problem.py
@@ -24,15 +24,22 @@ def __init__(
         parameters,
         model=None,
         check_model=True,
+        signal=["Voltage [V]"],
         init_soc=None,
         x0=None,
     ):
+        self.parameters = parameters
         self._model = model
         self.check_model = check_model
-        self.parameters = parameters
+        if isinstance(signal, str):
+            signal = [signal]
+        self.signal = signal
         self.init_soc = init_soc
         self.x0 = x0
         self.n_parameters = len(self.parameters)
+        self.n_outputs = len(self.signal)
+        self._time_data = None
+        self._target = None
 
         # Set bounds
         self.bounds = dict(
@@ -84,6 +91,28 @@ def evaluateS1(self, x):
         """
         raise NotImplementedError
 
+    def time_data(self):
+        """
+        Returns the time data.
+
+        Returns
+        -------
+        np.ndarray
+            The time array.
+        """
+        return self._time_data
+
+    def target(self):
+        """
+        Return the target dataset.
+
+        Returns
+        -------
+        np.ndarray
+            The target dataset array.
+        """
+        return self._target
+
 
 class FittingProblem(BaseProblem):
     """
@@ -108,20 +137,16 @@ def __init__(
         model,
         parameters,
         dataset,
-        signal=["Voltage [V]"],
         check_model=True,
+        signal=["Voltage [V]"],
         init_soc=None,
         x0=None,
     ):
-        super().__init__(parameters, model, check_model, init_soc, x0)
-        if isinstance(signal, str):
-            signal = [signal]
-        self.signal = signal
-        self.n_outputs = len(self.signal)
+        super().__init__(parameters, model, check_model, signal, init_soc, x0)
         self._dataset = dataset.data
 
         # Check that the dataset contains time and current
-        for name in ["Time [s]", "Current function [A]"] + signal:
+        for name in ["Time [s]", "Current function [A]"] + self.signal:
             if name not in self._dataset:
                 raise ValueError(f"expected {name} in list of dataset")
 
@@ -187,17 +212,6 @@ def evaluateS1(self, x):
 
         return (np.asarray(y), np.asarray(dy))
 
-    def target(self):
-        """
-        Return the target dataset.
-
-        Returns
-        -------
-        np.ndarray
-            The target dataset array.
-        """
-        return self._target
-
 
 class DesignProblem(BaseProblem):
     """
@@ -221,12 +235,12 @@ def __init__(
         parameters,
         experiment,
         check_model=True,
+        signal=["Voltage [V]"],
         init_soc=None,
         x0=None,
     ):
-        super().__init__(parameters, model, check_model, init_soc, x0)
+        super().__init__(parameters, model, check_model, signal, init_soc, x0)
         self.experiment = experiment
-        self._target = None
 
         # Build the model if required
         if experiment is not None:
@@ -273,13 +287,3 @@ def evaluateS1(self, x):
         )
 
         return (np.asarray(y), np.asarray(dy))
-
-    def target(self):
-        """
-        Return the target dataset (not applicable for design problems).
-
-        Returns
-        -------
-        None
-        """
-        return self._target

From a03396a36d4d4e2f36e079941b175d0a81a010bf Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 14 Dec 2023 16:42:15 +0000
Subject: [PATCH 15/17] Update to spm_two_signal_cost

---
 tests/unit/test_parameterisations.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tests/unit/test_parameterisations.py b/tests/unit/test_parameterisations.py
index b83537116..8f3ffde9e 100644
--- a/tests/unit/test_parameterisations.py
+++ b/tests/unit/test_parameterisations.py
@@ -133,7 +133,7 @@ def test_spm_optimisers(self, optimiser, spm_cost, x0):
         np.testing.assert_allclose(x, x0, atol=1e-1)
 
     @pytest.fixture
-    def spm_cost_2(self, parameters, model, x0):
+    def spm_two_signal_cost(self, parameters, model, x0):
         # Form dataset
         init_soc = 0.5
         solution = self.getdata(model, x0, init_soc)
@@ -162,9 +162,11 @@ def spm_cost_2(self, parameters, model, x0):
         ],
     )
     @pytest.mark.unit
-    def test_multiple_signals(self, optimiser, spm_cost_2, x0):
+    def test_multiple_signals(self, optimiser, spm_two_signal_cost, x0):
         # Test each optimiser
-        parameterisation = pybop.Optimisation(cost=spm_cost_2, optimiser=optimiser)
+        parameterisation = pybop.Optimisation(
+            cost=spm_two_signal_cost, optimiser=optimiser
+        )
         parameterisation.set_max_unchanged_iterations(iterations=15, threshold=5e-4)
         parameterisation.set_max_iterations(100)
 

From ee3848d7527f884ec4a091e8f0cc78e20926694b Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 14 Dec 2023 17:02:26 +0000
Subject: [PATCH 16/17] Add test on signal type

---
 pybop/_problem.py          | 2 ++
 tests/unit/test_problem.py | 3 +++
 2 files changed, 5 insertions(+)

diff --git a/pybop/_problem.py b/pybop/_problem.py
index df9da0b4f..ca5b7dd61 100644
--- a/pybop/_problem.py
+++ b/pybop/_problem.py
@@ -33,6 +33,8 @@ def __init__(
         self.check_model = check_model
         if isinstance(signal, str):
             signal = [signal]
+        elif not all(isinstance(item, str) for item in signal):
+            raise ValueError("Signal should be either a string or list of strings.")
         self.signal = signal
         self.init_soc = init_soc
         self.x0 = x0
diff --git a/tests/unit/test_problem.py b/tests/unit/test_problem.py
index 98be1f602..976388e94 100644
--- a/tests/unit/test_problem.py
+++ b/tests/unit/test_problem.py
@@ -81,6 +81,9 @@ def test_base_problem(self, parameters, model):
         with pytest.raises(NotImplementedError):
             problem.evaluateS1([0.5, 0.5])
 
+        with pytest.raises(ValueError):
+            pybop._problem.BaseProblem(parameters, model=model, signal=[0.5, 0.5])
+
     @pytest.mark.unit
     def test_fitting_problem(self, parameters, dataset, model, signal):
         # Test incorrect number of initial parameter values

From ca63085060d7cb5d9b4d78664adcab99478d9b5c Mon Sep 17 00:00:00 2001
From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com>
Date: Thu, 14 Dec 2023 17:19:03 +0000
Subject: [PATCH 17/17] Remove brackets

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

diff --git a/pybop/models/base_model.py b/pybop/models/base_model.py
index d8d2296b6..6e679cb28 100644
--- a/pybop/models/base_model.py
+++ b/pybop/models/base_model.py
@@ -215,7 +215,7 @@ def simulateS1(self, inputs, t_eval):
                 for key in self.fit_keys
             ]
 
-            return (np.vstack(predictions).T, np.dstack(sensitivities))
+            return np.vstack(predictions).T, np.dstack(sensitivities)
 
     def predict(
         self,