refactor: integration of pbeis functionality

examples/scripts/eis_fitting.py

@@ -0,0 +1,68 @@
+import numpy as np
+
+import pybop
+
+# Define model
+parameter_set = pybop.ParameterSet.pybamm("Chen2020")
+model = pybop.lithium_ion.DFN(
+    parameter_set=parameter_set, options={"surface form": "differential"}
+)
+
+# Fitting parameters
+parameters = pybop.Parameters(
+    pybop.Parameter(
+        "Negative electrode active material volume fraction",
+        prior=pybop.Gaussian(0.6, 0.05),
+    ),
+    pybop.Parameter(
+        "Positive electrode active material volume fraction",
+        prior=pybop.Gaussian(0.48, 0.05),
+    ),
+)
+
+# Generate data
+sigma = 0.001
+t_eval = np.arange(0, 900, 3)
+values = model.predict(t_eval=t_eval)
+corrupt_values = values["Voltage [V]"].data + np.random.normal(0, sigma, len(t_eval))
+
+# Form dataset
+dataset = pybop.Dataset(
+    {
+        "Frequency [Hz]": np.logspace(-4, 5, 300),
+        "Current function [A]": np.ones(300) * 0.0,
+        "Impedance": np.ones(300),
+    }
+)
+
+signal = ["Impedance"]
+# Generate problem, cost function, and optimisation class
+problem = pybop.EISProblem(model, parameters, dataset, signal=signal)
+prediction = problem.evaluate(np.array([0.75, 0.665]))
+# fig = px.scatter(x=prediction["Impedance"].real, y=-prediction["Impedance"].imag)
+# fig.show()
+# cost = pybop.SumSquaredError(problem)
+# optim = pybop.CMAES(cost, max_iterations=100)
+
+# # Run the optimisation
+# x, final_cost = optim.run()
+# print("True parameters:", parameters.true_value())
+# print("Estimated parameters:", x)
+
+# # Plot the time series
+# pybop.plot_dataset(dataset)
+
+# # Plot the timeseries output
+# pybop.quick_plot(problem, problem_inputs=x, title="Optimised Comparison")
+
+# # Plot convergence
+# pybop.plot_convergence(optim)
+
+# # Plot the parameter traces
+# pybop.plot_parameters(optim)
+
+# # Plot the cost landscape
+# pybop.plot2d(cost, steps=15)
+
+# # Plot the cost landscape with optimisation path
+# pybop.plot2d(optim, steps=15)

pybop/__init__.py

@@ -77,6 +77,7 @@
 from .problems.base_problem import BaseProblem
 from .problems.fitting_problem import FittingProblem
 from .problems.design_problem import DesignProblem
+from .problems.eis_problem import EISProblem
 
 #
 # Cost function class

pybop/_utils.py

@@ -47,11 +47,7 @@ def process_model(self, unprocessed_model, inplace=True):
             new model with parameter values set. Default is True.
         """
 
-        if inplace:
-            model = unprocessed_model
-        else:
-            model = unprocessed_model.new_copy()
-
+        model = unprocessed_model if inplace else unprocessed_model.new_copy()
         for variable, equation in unprocessed_model.rhs.items():
             pybamm.logger.verbose(f"Replacing symbols in {variable!r} (rhs)")
             model.rhs[self.process_symbol(variable)] = self.process_symbol(equation)
@@ -89,11 +85,6 @@ def process_model(self, unprocessed_model, inplace=True):
             )
         model.events = new_events
 
-        # Set external variables
-        # model.external_variables = [
-        #     self.process_symbol(var) for var in unprocessed_model.external_variables
-        # ]
-
         pybamm.logger.info(f"Finish replacing symbols in {model.name}")
 
         return model

pybop/models/base_model.py

@@ -102,6 +102,7 @@ def build(
             The initial state of charge to be used in simulations.
         """
         self.dataset = dataset
+        self.eis = eis
         if parameters is not None:
             self.parameters = parameters
             self.classify_and_update_parameters(self.parameters)
@@ -111,6 +112,21 @@ def build(
 
         if eis:
             self.set_up_for_eis(self.pybamm_model)
+            self.parameter_set["Current function [A]"] = 0
+            sim = pybamm.Simulation(
+                self.pybamm_model,
+                geometry=self.geometry,
+                parameter_values=self.parameter_set,
+                submesh_types=self.submesh_types,
+                var_pts=self.var_pts,
+                spatial_methods=self.spatial_methods,
+            )
+            sim.build()
+            self._built_model = sim.built_model
+
+            V_scale = getattr(self.pybamm_model.variables["Voltage [V]"], "scale", 1)
+            I_scale = getattr(self.pybamm_model.variables["Current [A]"], "scale", 1)
+            self.z_scale = self._parameter_set.evaluate(V_scale / I_scale)
 
         if self._built_model:
             return
@@ -120,7 +136,7 @@ def build(
         else:
             if not self.pybamm_model._built:
                 self.pybamm_model.build_model()
-            self.set_params()
+            self.set_params(eis=self.eis)
 
             self._mesh = pybamm.Mesh(self.geometry, self.submesh_types, self.var_pts)
             self._disc = pybamm.Discretisation(
@@ -134,7 +150,6 @@ def build(
 
             # Clear solver and setup model
             self._solver._model_set_up = {}
-            # self._solver.set_up(self._built_model, inputs={})
 
         self.n_states = self._built_model.len_rhs_and_alg  # len_rhs + len_alg
 
@@ -163,7 +178,7 @@ def set_init_soc(self, init_soc: float):
         # Save solved initial SOC in case we need to rebuild the model
         self._built_initial_soc = init_soc
 
-    def set_params(self, rebuild=False):
+    def set_params(self, eis=False, rebuild=False):
         """
         Assign the parameters to the model.
 
@@ -177,7 +192,11 @@ def set_params(self, rebuild=False):
         for key in self.standard_parameters.keys():
             self._parameter_set[key] = "[input]"
 
-        if self.dataset is not None and (not self.rebuild_parameters or not rebuild):
+        if (
+            not eis
+            and self.dataset is not None
+            and (not self.rebuild_parameters or not rebuild)
+        ):
             if (
                 self.parameters is None
                 or "Current function [A]" not in self.parameters.keys()
@@ -441,11 +460,9 @@ def simulate(
 
             return y
 
-    def simulateEIS(
-        self, inputs: Inputs, t_eval, frequencies: list
-    ) -> dict[str, np.ndarray]:
+    def simulateEIS(self, inputs: Inputs, f_eval: list) -> dict[str, np.ndarray]:
         """
-        Execute the forward model simulation with electrochemical impedence spectroscopy
+        Compute the forward model simulation with electrochemical impedence spectroscopy
         and return the result.
 
         Parameters
@@ -466,63 +483,60 @@ def simulateEIS(
         ValueError
             If the model has not been built before simulation.
         """
+        inputs = self.parameters.verify(inputs)
+
         if self._built_model is None:
             raise ValueError("Model must be built before calling simulate")
         else:
-            # if self.matched_parameters and not self.non_matched_parameters:
-            #     sol = self.solver.solve(self.built_model, t_eval=t_eval)
-
-            # else:
-            if not isinstance(inputs, dict):
-                inputs = {key: inputs[i] for i, key in enumerate(self.fit_keys)}
-
-            if self.check_params(
-                inputs=inputs,
-                allow_infeasible_solutions=self.allow_infeasible_solutions,
-            ):
-                model = self.built_model
-                inputs = (
-                    casadi.vertcat(*[x for x in inputs.values()])
-                    if model.convert_to_format == "casadi"
-                    else inputs
-                )
-
-                # Set up the solver for new inputs
-                self._solver.set_up(self._built_model, inputs=inputs)
+            if self.rebuild_parameters and not self.standard_parameters:
+                raise NotImplementedError
+            else:
+                if self.check_params(
+                    inputs=inputs,
+                    allow_infeasible_solutions=self.allow_infeasible_solutions,
+                ):
+                    model = self._built_model
+                    self.inputs = (
+                        casadi.vertcat(*[x for x in inputs.values()])
+                        if model.convert_to_format == "casadi"
+                        else inputs
+                    )
 
-                # Extract necessary attributes from the model
-                self.M = self._built_model.mass_matrix.entries
-                self.y0 = self._built_model.concatenated_initial_conditions.entries
-                self.J = self._built_model.jac_rhs_algebraic_eval(
-                    0, self.y0, []
-                ).sparse()
+                    # Set up the solver for new inputs
+                    self._solver.set_up(self._built_model, inputs=inputs)
 
-                # Convert to Compressed Sparse Column format
-                self.M = csc_matrix(self.M)
-                self.J = csc_matrix(self.J)
+                    # Extract necessary attributes from the model
+                    self.M = self._built_model.mass_matrix.entries
+                    self.y0 = self._built_model.concatenated_initial_conditions.entries
+                    self.J = self._built_model.jac_rhs_algebraic_eval(
+                        0, self.y0, []
+                    ).sparse()
 
-                # Add forcing to the RHS on the current density
-                self.b = np.zeros_like(self.y0)
-                self.b[-1] = -1
+                    # Convert to Compressed Sparse Column format
+                    self.M = csc_matrix(self.M)
+                    self.J = csc_matrix(self.J)
 
-                zs = []
-                for frequency in frequencies:
-                    # Compute the system matrix
-                    A = 1.0j * 2 * np.pi * frequency * self.M - self.J
+                    # Add forcing to the RHS on the current density
+                    self.b = np.zeros_like(self.y0)
+                    self.b[-1] = -1
 
-                    # Factorize the matrix
-                    lu = splu(A)
+                    zs = []
+                    for frequency in f_eval:
+                        # Compute the system matrix
+                        A = 1.0j * 2 * np.pi * frequency * self.M - self.J
 
-                    # Solve the system
-                    x = lu.solve(self.b)
+                        # Factorize the matrix
+                        lu = splu(A)
 
-                    # Calculate and store the impedance
-                    z = -x[-2][0] / x[-1][0]
-                    zs.append(z)
-            else:
-                return {"Impedence": [np.inf]}
+                        # Solve the system
+                        x = lu.solve(self.b)
 
-            return {"Impedence": zs}
+                        # Calculate and store the impedance
+                        z = -x[-2][0] / x[-1][0]
+                        zs.append(z)
+                    return {"Impedance": np.asarray(zs) * self.z_scale}
+                else:
+                    return {signal: [np.inf] for signal in self.signal}
 
     def simulateS1(self, inputs: Inputs, t_eval: np.array):
         """