Merge pull request #224 from pybop-team/177c-plotting-capabilities

Datatype restructure for optimisation objects

examples/scripts/exp_UKF.py

@@ -42,15 +42,16 @@
 simulator = pybop.Observer(parameters, model, signal=["2y"], x0=x0)
 simulator._time_data = t_eval
 measurements = simulator.evaluate(x0)
-measurements = measurements[:, 0]
 
 # Verification step: Compare by plotting
 go = pybop.PlotlyManager().go
 line1 = go.Scatter(x=t_eval, y=corrupt_values, name="Corrupt values", mode="markers")
 line2 = go.Scatter(
     x=t_eval, y=expected_values, name="Expected trajectory", mode="lines"
 )
-line3 = go.Scatter(x=t_eval, y=measurements, name="Observed values", mode="markers")
+line3 = go.Scatter(
+    x=t_eval, y=measurements["2y"], name="Observed values", mode="markers"
+)
 fig = go.Figure(data=[line1, line2, line3])
 
 # Form dataset
@@ -85,10 +86,11 @@
 
 # Verification step: Find the maximum likelihood estimate given the true parameters
 estimation = observer.evaluate(x0)
-estimation = estimation[:, 0]
 
 # Verification step: Add the estimate to the plot
-line4 = go.Scatter(x=t_eval, y=estimation, name="Estimated trajectory", mode="lines")
+line4 = go.Scatter(
+    x=t_eval, y=estimation["2y"], name="Estimated trajectory", mode="lines"
+)
 fig.add_trace(line4)
 fig.show()
 

examples/scripts/spm_CMAES.py

@@ -33,11 +33,13 @@
         "Time [s]": t_eval,
         "Current function [A]": values["Current [A]"].data,
         "Voltage [V]": corrupt_values,
+        "Bulk open-circuit voltage [V]": values["Bulk open-circuit voltage [V]"].data,
     }
 )
 
+signal = ["Voltage [V]", "Bulk open-circuit voltage [V]"]
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(model, parameters, dataset)
+problem = pybop.FittingProblem(model, parameters, dataset, signal=signal)
 cost = pybop.SumSquaredError(problem)
 optim = pybop.Optimisation(cost, optimiser=pybop.CMAES)
 optim.set_max_iterations(100)

examples/scripts/spm_adam.py

@@ -20,25 +20,50 @@
 ]
 
 # Generate data
-sigma = 0.001
-t_eval = np.arange(0, 900, 2)
-values = model.predict(t_eval=t_eval)
-corrupt_values = values["Voltage [V]"].data + np.random.normal(0, sigma, len(t_eval))
+init_soc = 0.5
+sigma = 0.003
+experiment = pybop.Experiment(
+    [
+        (
+            "Discharge at 0.5C for 3 minutes (1 second period)",
+            "Charge at 0.5C for 3 minutes (1 second period)",
+        ),
+    ]
+    * 2
+)
+values = model.predict(init_soc=init_soc, experiment=experiment)
+
+
+def noise(sigma):
+    return np.random.normal(0, sigma, len(values["Voltage [V]"].data))
+
 
 # Form dataset
 dataset = pybop.Dataset(
     {
-        "Time [s]": t_eval,
+        "Time [s]": values["Time [s]"].data,
         "Current function [A]": values["Current [A]"].data,
-        "Voltage [V]": corrupt_values,
+        "Voltage [V]": values["Voltage [V]"].data + noise(sigma),
+        "Bulk open-circuit voltage [V]": values["Bulk open-circuit voltage [V]"].data
+        + noise(sigma),
     }
 )
 
+signal = ["Voltage [V]", "Bulk open-circuit voltage [V]"]
 # Generate problem, cost function, and optimisation class
-problem = pybop.FittingProblem(model, parameters, dataset)
-cost = pybop.SumSquaredError(problem)
-optim = pybop.Optimisation(cost, optimiser=pybop.Adam, verbose=True)
+problem = pybop.FittingProblem(
+    model, parameters, dataset, signal=signal, init_soc=init_soc
+)
+cost = pybop.RootMeanSquaredError(problem)
+optim = pybop.Optimisation(
+    cost,
+    optimiser=pybop.Adam,
+    verbose=True,
+    allow_infeasible_solutions=True,
+    sigma0=0.05,
+)
 optim.set_max_iterations(100)
+optim.set_max_unchanged_iterations(45)
 
 # Run optimisation
 x, final_cost = optim.run()
@@ -55,3 +80,6 @@
 
 # Plot the cost landscape with optimisation path
 pybop.plot_optim2d(optim, steps=15)
+
+# Plot the cost and gradient landscapes
+pybop.plot_cost2d(cost, gradient=True, steps=3)

examples/standalone/problem.py

@@ -14,10 +14,13 @@ def __init__(
         model=None,
         check_model=True,
         signal=None,
+        default_variables=None,
         init_soc=None,
         x0=None,
     ):
-        super().__init__(parameters, model, check_model, signal, init_soc, x0)
+        super().__init__(
+            parameters, model, check_model, signal, default_variables, init_soc, x0
+        )
         self._dataset = dataset.data
 
         # Check that the dataset contains time and current
@@ -37,8 +40,7 @@ def __init__(
                 raise ValueError(
                     f"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
+        self._target = {signal: self._dataset[signal] for signal in self.signal}
 
     def evaluate(self, x):
         """
@@ -55,7 +57,7 @@ def evaluate(self, x):
             The model output y(t) simulated with inputs x.
         """
 
-        return x[0] * self._time_data + x[1]
+        return {signal: x[0] * self._time_data + x[1] for signal in self.signal}
 
     def evaluateS1(self, x):
         """
@@ -73,8 +75,8 @@ def evaluateS1(self, x):
             with given inputs x.
         """
 
-        y = x[0] * self._time_data + x[1]
+        y = {signal: x[0] * self._time_data + x[1] for signal in self.signal}
 
-        dy = np.dstack([self._time_data, np.zeros(self._time_data.shape)])
+        dy = [self._time_data, np.zeros(self._time_data.shape)]
 
-        return (np.asarray(y), np.asarray(dy))
+        return (y, np.asarray(dy))

noxfile.py

@@ -33,7 +33,7 @@ def coverage(session):
         "--cov-report=xml",
     )
     session.run(
-        "pytest", "--plots", "--cov", "--cov-append", "--cov-report=xml", "-n", "1"
+        "pytest", "--plots", "--cov", "--cov-append", "--cov-report=xml", "-n", "0"
     )
 
 

pybop/_optimisation.py

@@ -175,15 +175,16 @@ def _run_pybop(self):
         final_cost : float
             The final cost associated with the best parameters.
         """
-        x, final_cost = self.optimiser.optimise(
+        result = self.optimiser.optimise(
             cost_function=self.cost,
             x0=self.x0,
             bounds=self.bounds,
             maxiter=self._max_iterations,
         )
         self.log = self.optimiser.log
+        self._iterations = result.nit
 
-        return x, final_cost
+        return result.x, self.cost(result.x)
 
     def _run_pints(self):
         """

pybop/_problem.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pybop
 
 
 class BaseProblem:
@@ -15,6 +16,8 @@ class BaseProblem:
         Flag to indicate if the model should be checked (default: True).
     signal: List[str]
       The signal to observe.
+    additional_variables : List[str], optional
+        Additional variables to observe and store in the solution (default: []).
     init_soc : float, optional
         Initial state of charge (default: None).
     x0 : np.ndarray, optional
@@ -27,6 +30,7 @@ def __init__(
         model=None,
         check_model=True,
         signal=["Voltage [V]"],
+        additional_variables=[],
         init_soc=None,
         x0=None,
     ):
@@ -45,6 +49,11 @@ def __init__(
         self._time_data = None
         self._target = None
 
+        if isinstance(model, (pybop.BaseModel, pybop.lithium_ion.EChemBaseModel)):
+            self.additional_variables = additional_variables
+        else:
+            self.additional_variables = []
+
         # Set bounds
         self.bounds = dict(
             lower=[param.bounds[0] for param in self.parameters],
@@ -138,7 +147,13 @@ class FittingProblem(BaseProblem):
     dataset : Dataset
         Dataset object containing the data to fit the model to.
     signal : str, optional
-        The signal to fit (default: "Voltage [V]").
+        The variable used for fitting (default: "Voltage [V]").
+    additional_variables : List[str], optional
+        Additional variables to observe and store in the solution (default additions are: ["Time [s]"]).
+    init_soc : float, optional
+        Initial state of charge (default: None).
+    x0 : np.ndarray, optional
+        Initial parameter values (default: None).
     """
 
     def __init__(
@@ -148,10 +163,14 @@ def __init__(
         dataset,
         check_model=True,
         signal=["Voltage [V]"],
+        additional_variables=[],
         init_soc=None,
         x0=None,
     ):
-        super().__init__(parameters, model, check_model, signal, init_soc, x0)
+        additional_variables += ["Time [s]"]
+        super().__init__(
+            parameters, model, check_model, signal, additional_variables, init_soc, x0
+        )
         self._dataset = dataset.data
         self.x = self.x0
 
@@ -161,12 +180,13 @@ def __init__(
         # Unpack time and target data
         self._time_data = self._dataset["Time [s]"]
         self.n_time_data = len(self._time_data)
-        target = [self._dataset[signal] for signal in self.signal]
-        self._target = np.vstack(target).T
+        self._target = {signal: self._dataset[signal] for signal in self.signal}
 
         # Add useful parameters to model
         if model is not None:
             self._model.signal = self.signal
+            self._model.additional_variables = self.additional_variables
+            self._model.n_parameters = self.n_parameters
             self._model.n_outputs = self.n_outputs
             self._model.n_time_data = self.n_time_data
 
@@ -193,14 +213,14 @@ def evaluate(self, x):
         y : np.ndarray
             The model output y(t) simulated with inputs x.
         """
-        if (x != self.x).any() and self._model.matched_parameters:
+        if np.any(x != self.x) and self._model.matched_parameters:
             for i, param in enumerate(self.parameters):
                 param.update(value=x[i])
 
             self._model.rebuild(parameters=self.parameters)
             self.x = x
 
-        y = np.asarray(self._model.simulate(inputs=x, t_eval=self._time_data))
+        y = self._model.simulate(inputs=x, t_eval=self._time_data)
 
         return y
 
@@ -229,7 +249,7 @@ def evaluateS1(self, x):
             t_eval=self._time_data,
         )
 
-        return (np.asarray(y), np.asarray(dy))
+        return (y, np.asarray(dy))
 
 
 class DesignProblem(BaseProblem):
@@ -246,6 +266,16 @@ class DesignProblem(BaseProblem):
         List of parameters for the problem.
     experiment : object
         The experimental setup to apply the model to.
+    check_model : bool, optional
+        Flag to indicate if the model parameters should be checked for feasibility each iteration (default: True).
+    signal : str, optional
+        The signal to fit (default: "Voltage [V]").
+    additional_variables : List[str], optional
+        Additional variables to observe and store in the solution (default additions are: ["Time [s]", "Current [A]"]).
+    init_soc : float, optional
+        Initial state of charge (default: None).
+    x0 : np.ndarray, optional
+        Initial parameter values (default: None).
     """
 
     def __init__(
@@ -255,10 +285,14 @@ def __init__(
         experiment,
         check_model=True,
         signal=["Voltage [V]"],
+        additional_variables=[],
         init_soc=None,
         x0=None,
     ):
-        super().__init__(parameters, model, check_model, signal, init_soc, x0)
+        additional_variables += ["Time [s]", "Current [A]"]
+        super().__init__(
+            parameters, model, check_model, signal, additional_variables, init_soc, x0
+        )
         self.experiment = experiment
 
         # Build the model if required
@@ -278,8 +312,8 @@ def __init__(
 
         # Add an example dataset for plotting comparison
         sol = self.evaluate(self.x0)
-        self._time_data = sol[:, -1]
-        self._target = sol[:, 0:-1]
+        self._time_data = sol["Time [s]"]
+        self._target = {key: sol[key] for key in self.signal}
         self._dataset = None
 
     def evaluate(self, x):
@@ -307,6 +341,8 @@ def evaluate(self, x):
             return sol
 
         else:
-            predictions = [sol[signal].data for signal in self.signal + ["Time [s]"]]
+            predictions = {}
+            for signal in self.signal + self.additional_variables:
+                predictions[signal] = sol[signal].data
 
-            return np.vstack(predictions).T
+        return predictions

pybop/costs/base_cost.py

@@ -33,6 +33,7 @@ def __init__(self, problem):
             self.bounds = problem.bounds
             self.n_parameters = problem.n_parameters
             self.n_outputs = problem.n_outputs
+            self.signal = problem.signal
 
     def __call__(self, x, grad=None):
         """

pybop/costs/design_costs.py

@@ -57,9 +57,12 @@ def update_simulation_data(self, initial_conditions):
         if self.update_capacity:
             self.problem.model.approximate_capacity(self.problem.x0)
         solution = self.problem.evaluate(initial_conditions)
-        self.problem._time_data = solution[:, -1]
-        self.problem._target = solution[:, 0:-1]
-        self.dt = solution[1, -1] - solution[0, -1]
+
+        if "Time [s]" not in solution:
+            raise ValueError("The solution does not contain time data.")
+        self.problem._time_data = solution["Time [s]"]
+        self.problem._target = {key: solution[key] for key in self.problem.signal}
+        self.dt = solution["Time [s]"][1] - solution["Time [s]"][0]
 
     def _evaluate(self, x, grad=None):
         """
@@ -123,7 +126,7 @@ def _evaluate(self, x, grad=None):
                     self.problem.model.approximate_capacity(x)
                 solution = self.problem.evaluate(x)
 
-                voltage, current = solution[:, 0], solution[:, 1]
+                voltage, current = solution["Voltage [V]"], solution["Current [A]"]
                 negative_energy_density = -np.trapz(voltage * current, dx=self.dt) / (
                     3600 * self.problem.model.cell_mass(self.parameter_set)
                 )
@@ -181,7 +184,7 @@ def _evaluate(self, x, grad=None):
                     self.problem.model.approximate_capacity(x)
                 solution = self.problem.evaluate(x)
 
-                voltage, current = solution[:, 0], solution[:, 1]
+                voltage, current = solution["Voltage [V]"], solution["Current [A]"]
                 negative_energy_density = -np.trapz(voltage * current, dx=self.dt) / (
                     3600 * self.problem.model.cell_volume(self.parameter_set)
                 )