Add saving/loading example

examples/scripts/getting_started/linked_parameters.py

@@ -62,18 +62,10 @@
 parameter_values.update(
     {
         "Positive electrode thickness [m]": pybop.Parameter(
-            distribution=pybop.Gaussian(
-                7.56e-05,
-                0.1e-05,
-                truncated_at=[65e-06, 10e-05],
-            ),
+            distribution=pybop.Gaussian(75.6e-6, 1e-6, truncated_at=[65e-6, 100e-6]),
         ),
         "Positive electrode active material volume fraction": pybop.Parameter(
-            distribution=pybop.Gaussian(
-                0.6,
-                0.15,
-                truncated_at=[0.1, 0.9],
-            ),
+            distribution=pybop.Gaussian(0.6, 0.15, truncated_at=[0.1, 0.9]),
         ),
     }
 )

examples/scripts/getting_started/maximum_likelihood.py

@@ -35,11 +35,7 @@
 parameter_values.update(
     {
         "Negative electrode active material volume fraction": pybop.Parameter(
-            distribution=pybop.Gaussian(
-                0.6,
-                0.05,
-                truncated_at=[0.5, 0.8],
-            )
+            distribution=pybop.Gaussian(0.6, 0.05, truncated_at=[0.5, 0.8])
         ),
         "Positive electrode active material volume fraction": pybop.Parameter(
             distribution=pybop.Gaussian(0.48, 0.05),

examples/scripts/getting_started/optimising_with_adamw.py

@@ -31,11 +31,7 @@
 parameter_values.update(
     {
         "Negative electrode active material volume fraction": pybop.Parameter(
-            distribution=pybop.Gaussian(
-                0.68,
-                0.05,
-                truncated_at=[0.4, 0.9],
-            ),
+            distribution=pybop.Gaussian(0.68, 0.05, truncated_at=[0.4, 0.9]),
             initial_value=0.45,
         ),
         "Positive electrode active material volume fraction": pybop.Parameter(

examples/scripts/getting_started/optimising_with_scipy_minimize.py

@@ -41,18 +41,10 @@
 parameter_values.update(
     {
         "Negative electrode active material volume fraction": pybop.Parameter(
-            distribution=pybop.Gaussian(
-                0.6,
-                0.05,
-                truncated_at=[0.5, 0.8],
-            ),
+            distribution=pybop.Gaussian(0.6, 0.05, truncated_at=[0.5, 0.8]),
         ),
         "Positive electrode active material volume fraction": pybop.Parameter(
-            distribution=pybop.Gaussian(
-                0.48,
-                0.05,
-                truncated_at=[0.4, 0.7],
-            ),
+            distribution=pybop.Gaussian(0.48, 0.05, truncated_at=[0.4, 0.7]),
         ),
     }
 )

examples/scripts/getting_started/optimising_with_simulated_annealing.py

@@ -33,18 +33,10 @@
 parameter_values.update(
     {
         "Negative electrode active material volume fraction": pybop.Parameter(
-            distribution=pybop.Gaussian(
-                0.6,
-                0.1,
-                truncated_at=[0.4, 0.85],
-            ),
+            distribution=pybop.Gaussian(0.6, 0.1, truncated_at=[0.4, 0.85]),
         ),
         "Positive electrode active material volume fraction": pybop.Parameter(
-            distribution=pybop.Gaussian(
-                0.6,
-                0.1,
-                truncated_at=[0.4, 0.85],
-            ),
+            distribution=pybop.Gaussian(0.6, 0.1, truncated_at=[0.4, 0.85]),
         ),
     }
 )

examples/scripts/getting_started/saving_and_loading_results.py

@@ -0,0 +1,80 @@
+import os
+
+import numpy as np
+import pybamm
+
+import pybop
+
+"""
+This example shows how to save and load an optimisation (or sampling) result.
+First we run an example optimisation to generate a result.
+"""
+
+# Define model and parameter values
+model = pybamm.lithium_ion.SPM()
+parameter_values = pybamm.ParameterValues("Chen2020")
+
+# Generate a synthetic dataset
+sigma = 5e-3
+t_eval = np.linspace(0, 500, 240)
+solution = pybamm.Simulation(model, parameter_values=parameter_values).solve(
+    t_eval=t_eval
+)
+dataset = pybop.Dataset(
+    {
+        "Time [s]": t_eval,
+        "Current [A]": solution["Current [A]"](t_eval),
+        "Voltage [V]": pybop.add_noise(solution["Voltage [V]"](t_eval), sigma),
+    }
+)
+
+# Fitting parameters
+parameter_values.update(
+    {
+        "Negative electrode active material volume fraction": pybop.Parameter(
+            distribution=pybop.Gaussian(0.68, 0.05, truncated_at=[0.4, 0.9]),
+            initial_value=0.45,
+        ),
+        "Positive electrode active material volume fraction": pybop.Parameter(
+            distribution=pybop.Gaussian(0.58, 0.05, truncated_at=[0.4, 0.9]),
+            initial_value=0.45,
+        ),
+    }
+)
+
+# Build the problem
+simulator = pybop.pybamm.Simulator(
+    model, parameter_values=parameter_values, protocol=dataset
+)
+cost = pybop.MeanAbsoluteError(dataset)
+problem = pybop.Problem(simulator, cost)
+
+# Set up the optimiser
+options = pybop.PintsOptions(max_iterations=150, verbose=True)
+optim = pybop.PSO(problem, options=options)
+
+# Run the optimisation
+result = optim.run()
+
+# Save the result: either pickle the whole result or save the data in
+# one of these formats: "pickle", "json", "matlab"
+save_path = "examples/results/"
+os.makedirs(os.path.dirname(save_path), exist_ok=True)
+result.save(save_path + "saved_result_object.pkl")
+result.save_data(save_path + "saved_result_data.json", to_format="json")
+
+# Load the result
+result_from_pkl = pybop.Result.load(save_path + "saved_result_object.pkl")
+result_from_json = pybop.Result.load_data(
+    save_path + "saved_result_data.json", file_format="json"
+)
+
+# Plot the optimisation result from .pkl
+result_from_pkl.plot_convergence()
+result_from_pkl.plot_parameters()
+result_from_pkl.plot_surface(bounds=problem.parameters.get_bounds_array())
+
+# Plot the optimisation result from .json (it is the same)
+result_from_json.plot_convergence()
+result_from_json.plot_parameters()
+result_from_json.plot_surface(bounds=problem.parameters.get_bounds_array())

pybop/_result.py

@@ -347,6 +347,7 @@ def data_dict(self) -> dict:
 
         return {
             "minimising": self._minimising,
+            "parameter_names": self.problem.parameters.names,
             "method_name": self.method_name,
             "n_runs": self.n_runs,
             "best_run": self._best_run,
@@ -425,6 +426,7 @@ def load_data(filename: str, file_format: str = "pickle") -> dict:
             file_format=file_format,
             data_keys_0d=["_minimising", "n_runs", "best_run"],
             data_keys_1d=[
+                "parameter_names",
                 "method_name",
                 "best_cost",
                 "initial_cost",
@@ -439,6 +441,8 @@ def load_data(filename: str, file_format: str = "pickle") -> dict:
         # Create a dummy problem
         problem = types.SimpleNamespace()
         problem.minimising = data["minimising"]
+        problem.parameters = types.SimpleNamespace()
+        problem.parameters.names = data["parameter_names"]
 
         # Create one logging result for each run
         n_runs = data["n_runs"]

pybop/_utils.py

@@ -117,11 +117,13 @@ def save_data_dict(
         df = pd.DataFrame(data_dict_copy)
         return df.to_csv(filename, index=False)
     elif to_format == "json":
+        # Format with indentation
+        json_clean = json.dumps(data_dict, cls=NumpyEncoder, indent=4) + "\n"
         if filename is None:
-            return json.dumps(data_dict, cls=NumpyEncoder)
+            return json_clean
         else:
             with open(filename, "w") as outfile:
-                json.dump(data_dict, outfile, cls=NumpyEncoder)
+                outfile.write(json_clean)
     else:
         raise ValueError(f"format '{to_format}' is not supported")
 

pybop/plot/voronoi.py

@@ -259,7 +259,7 @@ def surface(
     points = result.x_model
     parameters = result.problem.parameters
 
-    if points[0].shape[0] != 2:
+    if np.shape(points[0])[0] != 2:
         raise ValueError("This plot method requires two parameters.")
 
     x_optim, y_optim = map(list, zip(*points, strict=False))

tests/unit/test_optimisation.py

@@ -754,7 +754,7 @@ def compare_result_data(self, result1, result2):
         np.testing.assert_array_equal(result1._time, result2._time)
 
     @pytest.mark.parametrize("to_format", ["json", "matlab", "pickle"])
-    def test_save_result_data(self, result, problem, to_format, tmp_path):
+    def test_save_result_data(self, result, to_format, tmp_path):
         test_stub = tmp_path / "test"
 
         if to_format == "matlab":
@@ -766,14 +766,14 @@ def test_save_result_data(self, result, problem, to_format, tmp_path):
         # Test save result
         result.save_data(filename, to_format=to_format)
 
-        result_load = OptimisationResult.load_data(filename, file_format=to_format)
+        result_load = pybop.Result.load_data(filename, file_format=to_format)
         self.compare_result_data(result, result_load)
 
         # Test save combined result
-        result_combined = OptimisationResult.combine([result, result])
+        result_combined = pybop.Result.combine([result, result])
         result_combined.save_data(filename, to_format=to_format)
 
-        result_load = OptimisationResult.load_data(filename, file_format=to_format)
+        result_load = pybop.Result.load_data(filename, file_format=to_format)
         self.compare_result_data(result_combined, result_load)
 
     def test_save_result(self, result, tmp_path):
@@ -782,6 +782,6 @@ def test_save_result(self, result, tmp_path):
         # test save whole result
         filename = f"{test_stub}.pickle"
         result.save(filename)
-        result_load = OptimisationResult.load(filename)
+        result_load = pybop.Result.load(filename)
         self.compare_result_data(result, result_load)
         assert result.problem.parameters.names == result_load.problem.parameters.names

tests/unit/test_sampling.py

@@ -351,14 +351,14 @@ def test_save(self, log_pdf, n_chains, MCMC, tmp_path):
             result.save_data(filename, to_format=to_format)
 
             # load result
-            result_load = SamplingResult.load_data(filename, file_format=to_format)
+            result_load = pybop.Result.load_data(filename, file_format=to_format)
             self.compare_result_data(result, result_load)
             assert sampler2.logger is None
 
         # test save whole result
         filename = f"{test_stub}.pickle"
         result.save(filename)
-        result_load = SamplingResult.load(filename)
+        result_load = pybop.Result.load(filename)
         self.compare_result_data(result, result_load)
         assert result.problem.parameters.names == result_load.problem.parameters.names
         np.testing.assert_array_equal(result.chains, result_load.chains)