Include PyBaMM solution in OptimizationResult

pybop/optimisers/base_optimiser.py

@@ -302,6 +302,8 @@ class OptimisationResult:
         Number of iterations performed by the optimiser.
     scipy_result : scipy.optimize.OptimizeResult, optional
         The result obtained from a SciPy optimiser.
+    pybamm_solution: pybamm.Solution or list[pybamm.Solution], optional
+        The final solution object(s) obtained from the optimisation.
     """
 
     def __init__(
@@ -313,6 +315,7 @@ def __init__(
         optim: Optional[BaseOptimiser] = None,
         time: Optional[float] = None,
         scipy_result=None,
+        pybamm_solution=None,
     ):
         self.x = x
         self.cost = cost
@@ -328,6 +331,8 @@ def __init__(
         else:
             self.x0 = None
 
+        self.pybamm_solution = pybamm_solution
+
         # Check that the parameters produce finite cost, and are physically viable
         self._validate_parameters()
         self.check_physical_viability(self.x)
@@ -403,5 +408,6 @@ def __str__(self) -> str:
             f"  Final cost: {self.final_cost}\n"
             f"  Optimisation time: {self.time} seconds\n"
             f"  Number of iterations: {self.n_iterations}\n"
-            f"  SciPy result available: {'Yes' if self.scipy_result else 'No'}"
+            f"  SciPy result available: {'Yes' if self.scipy_result else 'No'}\n"
+            f"  PyBaMM Solution available: {'Yes' if self.pybamm_solution else 'No'}"
         )

pybop/optimisers/base_pints_optimiser.py

@@ -359,6 +359,7 @@ def fun(x):
             n_iterations=self._iterations,
             optim=self,
             time=total_time,
+            pybamm_solution=self.cost.problem.solution,
         )
 
     def f_guessed_tracking(self):

pybop/optimisers/scipy_optimisers.py

@@ -88,15 +88,18 @@ def _run(self):
             nit = -1
 
         return OptimisationResult(
-            x=self._transformation.to_model(result.x)
-            if self._transformation
-            else result.x,
+            x=(
+                self._transformation.to_model(result.x)
+                if self._transformation
+                else result.x
+            ),
             cost=self.cost,
             final_cost=self.cost(result.x, apply_transform=True),
             n_iterations=nit,
             scipy_result=result,
             optim=self,
             time=total_time,
+            pybamm_solution=self.cost.problem.solution,
         )
 
 

pybop/problems/base_problem.py

@@ -96,6 +96,9 @@ def __init__(
                 output_variables=self.output_vars,
             )
 
+        # to store pybamm solution objects
+        self._solution = None
+
     def set_initial_state(self, initial_state: Optional[dict] = None):
         """
         Set the initial state to be applied to evaluations of the problem.
@@ -195,3 +198,7 @@ def domain_data(self, domain_data):
     @property
     def dataset(self):
         return self._dataset
+
+    @property
+    def solution(self):
+        return self._solution

pybop/problems/design_problem.py

@@ -148,6 +148,8 @@ def evaluate(self, inputs: Inputs):
                 for signal in [*self.signal, *self.additional_variables]
             }
 
+        self._solution = sol
+
         return {
             signal: sol[signal].data
             for signal in self.signal + self.additional_variables

pybop/problems/fitting_problem.py

@@ -130,6 +130,7 @@ def evaluate(self, inputs: Inputs) -> dict[str, np.ndarray[np.float64]]:
                     print(f"Simulation error: {e}")
                 return {signal: self.failure_output for signal in self.signal}
 
+            self._solution = sol
             return {
                 signal: sol[signal].data
                 for signal in (self.signal + self.additional_variables)

pybop/problems/multi_fitting_problem.py

@@ -94,14 +94,17 @@ def evaluate(self, inputs: Inputs):
         self.parameters.update(values=list(inputs.values()))
 
         combined_signal = []
+        combined_solutions = []
 
         for problem in self.problems:
             problem_inputs = problem.parameters.as_dict()
             signal_values = problem.evaluate(problem_inputs)
+            combined_solutions.append(problem.solution)
 
             # Collect signals
             for signal in problem.signal:
                 combined_signal.extend(signal_values[signal])
+        self._solution = combined_solutions
 
         return {"Combined signal": np.asarray(combined_signal)}
 

tests/unit/test_optimisation.py

@@ -507,7 +507,8 @@ def test_halting(self, cost):
             f"  Final cost: {results.final_cost}\n"
             f"  Optimisation time: {results.time} seconds\n"
             f"  Number of iterations: {results.n_iterations}\n"
-            f"  SciPy result available: No"
+            f"  SciPy result available: No\n"
+            f"  PyBaMM Solution available: No"
         )
 
         # Test guessed values
@@ -550,7 +551,8 @@ def test_halting(self, cost):
             f"  Final cost: {results.final_cost}\n"
             f"  Optimisation time: {results.time} seconds\n"
             f"  Number of iterations: {results.n_iterations}\n"
-            f"  SciPy result available: No"
+            f"  SciPy result available: No\n"
+            f"  PyBaMM Solution available: Yes"
         )
 
         optim.set_max_unchanged_iterations()