Fix integration test logic, add gradient landscape plots, pin pytest …

…version due to breaking change in 8.1.0

examples/scripts/exp_UKF.py

@@ -86,10 +86,11 @@
 
 # Verification step: Find the maximum likelihood estimate given the true parameters
 estimation = observer.evaluate(x0)
-estimation = estimation["2y"]
 
 # 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_adam.py

@@ -20,34 +20,50 @@
 ]
 
 # Generate data
-sigma = 0.01
-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,
-        "Bulk open-circuit voltage [V]": values["Bulk open-circuit voltage [V]"].data,
+        "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, signal=signal)
+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=sigma,
+    sigma0=0.05,
 )
 optim.set_max_iterations(100)
-optim.set_max_unchanged_iterations(20)
+optim.set_max_unchanged_iterations(45)
 
 # Run optimisation
 x, final_cost = optim.run()
@@ -64,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)

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/costs/fitting_costs.py

@@ -82,9 +82,9 @@
         for key in y:
             if key not in ["Time [s]", "Discharge capacity [A.h]"]:
                 if len(y.get(key, [])) != len(self._target.get(key, [])):
                     e = np.float64(np.inf)
                     de = self._de * np.ones(self.n_parameters)
                     return e, de
 
         r = np.array(
             [
@@ -98,16 +98,17 @@
             r = r.reshape(self.problem.n_time_data)
             dy = dy.reshape(self.n_parameters, self.problem.n_time_data)
             e = np.sqrt(np.mean(r**2))
-            de = np.mean((r * dy), axis=1) / np.sqrt(
-                np.mean((r * dy) ** 2, axis=1) + np.finfo(float).eps
+            de = np.mean((r * dy), axis=1) / (
+                np.sqrt(np.mean((r * dy) ** 2, axis=1) + np.finfo(float).eps)
             )
             return e.item(), de.flatten()
 
         else:
             r = r.reshape(self.n_outputs, self.problem.n_time_data)
             e = np.sqrt(np.mean(r**2, axis=1))
-            de = np.mean((r[:, :, np.newaxis] * dy), axis=1) / np.sqrt(
-                np.mean((r[:, :, np.newaxis] * dy) ** 2, axis=1) + np.finfo(float).eps
+            de = np.mean((r[:, :, np.newaxis] * dy), axis=1) / (
+                np.sqrt(np.mean((r[:, :, np.newaxis] * dy) ** 2, axis=1))
+                + np.finfo(float).eps
             )
             return np.sum(e), np.sum(de, axis=1)
 
@@ -162,7 +163,7 @@
 
         e = np.array(
             [
-                np.sum(((prediction[signal] - self._target[signal]) ** 2), axis=0)
+                np.sum(((prediction[signal] - self._target[signal]) ** 2))
                 for signal in prediction
                 if signal not in ["Time [s]", "Discharge capacity [A.h]"]
             ]
@@ -196,9 +197,9 @@
         for key in y:
             if key not in ["Time [s]", "Discharge capacity [A.h]"]:
                 if len(y.get(key, [])) != len(self._target.get(key, [])):
                     e = np.float64(np.inf)
                     de = self._de * np.ones(self.n_parameters)
                     return e, de
 
         r = np.array(
             [
@@ -217,7 +218,7 @@
 
         else:
             r = r.reshape(self.n_outputs, self.problem.n_time_data)
-            e = np.sum(r**2, axis=0)
+            e = np.sum(r**2, axis=1)
             de = 2 * np.sum((r[:, :, np.newaxis] * dy), axis=1)
             return np.sum(e), np.sum(de, axis=1)
 

pybop/plotting/plot_cost2d.py

@@ -3,7 +3,9 @@
 import numpy as np
 
 
-def plot_cost2d(cost, bounds=None, steps=10, show=True, **layout_kwargs):
+def plot_cost2d(
+    cost, gradient=False, bounds=None, steps=10, show=True, **layout_kwargs
+):
     """
     Plot a 2D visualisation of a cost landscape using Plotly.
 
@@ -54,6 +56,23 @@
         for j, yj in enumerate(y):
             costs[j, i] = cost(np.array([xi, yj]))
 
+    if gradient:
+        grad_parameter_costs = []
+
+        # Determine the number of gradient outputs from cost.evaluateS1
+        num_gradients = len(cost.evaluateS1(np.array([x[0], y[0]]))[1])
+
+        # Create an array to hold each gradient output & populate
+        grads = [np.zeros((len(y), len(x))) for _ in range(num_gradients)]
+        for i, xi in enumerate(x):
+            for j, yj in enumerate(y):
+                (*current_grads,) = cost.evaluateS1(np.array([xi, yj]))[1]
+                for k, grad_output in enumerate(current_grads):
+                    grads[k][j, i] = grad_output
+
+        # Append the arrays to the grad_parameter_costs list
+        grad_parameter_costs.extend(grads)
+
     # Import plotly only when needed
     go = pybop.PlotlyManager().go
 
@@ -80,6 +99,32 @@
     elif show:
         fig.show()
 
+    if gradient:
+        grad_figs = []
+        for i, grad_costs in enumerate(grad_parameter_costs):
+            # Update title for gradient plots
+            updated_layout_options = layout_options.copy()
+            updated_layout_options["title"] = f"Gradient for Parameter: {i+1}"
+
+            # Create contour plot with updated layout options
+            grad_layout = go.Layout(updated_layout_options)
+
+            # Create fig
+            grad_fig = go.Figure(
+                data=[go.Contour(x=x, y=y, z=grad_costs)], layout=grad_layout
+            )
+            grad_fig.update_layout(**layout_kwargs)
+
+            if "ipykernel" in sys.modules and show:
+                grad_fig.show("svg")
+            elif show:
+                grad_fig.show()
+
+            # append grad_fig to list
+            grad_figs.append(grad_fig)
+
+        return fig, grad_figs
+
     return fig
 
 

pyproject.toml

@@ -41,7 +41,7 @@ dev = [
     "nox",
     "nbmake",
     "pre-commit",
-    "pytest>=6",
+    "pytest<=8",
     "pytest-cov",
     "pytest-mock",
     "pytest-xdist",