Repository: pybamm-team/PyBaMM

Issue:
Theoretical energy integral is very slow
This for loop is killing us and should be vectorized: https://github.com/pybamm-team/PyBaMM/blob/33d21874fbc39fd81ff89a8cc5638f6a81f8c31d/pybamm/models/full_battery_models/lithium_ion/electrode_soh.py#L1007-L1011

Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b02df8ed4..c1cf4b91e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## Bug fixes
 
+- Fixed bug in calculation of theoretical energy that made it very slow ([#3506](https://github.com/pybamm-team/PyBaMM/pull/3506))
 - Fixed a bug where the JaxSolver would fails when using GPU support with no input parameters ([#3423](https://github.com/pybamm-team/PyBaMM/pull/3423))
 
 # [v23.9rc0](https://github.com/pybamm-team/PyBaMM/tree/v23.9rc0) - 2023-10-31
diff --git a/pybamm/expression_tree/binary_operators.py b/pybamm/expression_tree/binary_operators.py
index 749384e9b..bde9a1727 100644
--- a/pybamm/expression_tree/binary_operators.py
+++ b/pybamm/expression_tree/binary_operators.py
@@ -14,12 +14,12 @@ import pybamm
 def _preprocess_binary(left, right):
     if isinstance(left, numbers.Number):
         left = pybamm.Scalar(left)
-    if isinstance(right, numbers.Number):
-        right = pybamm.Scalar(right)
     elif isinstance(left, np.ndarray):
         if left.ndim > 1:
             raise ValueError("left must be a 1D array")
         left = pybamm.Vector(left)
+    if isinstance(right, numbers.Number):
+        right = pybamm.Scalar(right)
     elif isinstance(right, np.ndarray):
         if right.ndim > 1:
             raise ValueError("right must be a 1D array")
diff --git a/pybamm/expression_tree/broadcasts.py b/pybamm/expression_tree/broadcasts.py
index 32cf2c002..d30762ad7 100644
--- a/pybamm/expression_tree/broadcasts.py
+++ b/pybamm/expression_tree/broadcasts.py
@@ -546,8 +546,10 @@ def full_like(symbols, fill_value):
         return array_type(entries, domains=sum_symbol.domains)
 
     except NotImplementedError:
-        if sum_symbol.shape_for_testing == (1, 1) or sum_symbol.shape_for_testing == (
-            1,
+        if (
+            sum_symbol.shape_for_testing == (1, 1)
+            or sum_symbol.shape_for_testing == (1,)
+            or sum_symbol.domain == []
         ):
             return pybamm.Scalar(fill_value)
         if sum_symbol.evaluates_on_edges("primary"):
diff --git a/pybamm/models/full_battery_models/lithium_ion/electrode_soh.py b/pybamm/models/full_battery_models/lithium_ion/electrode_soh.py
index c6a445f31..d975de859 100644
--- a/pybamm/models/full_battery_models/lithium_ion/electrode_soh.py
+++ b/pybamm/models/full_battery_models/lithium_ion/electrode_soh.py
@@ -410,10 +410,7 @@ class ElectrodeSOHSolver:
         # Calculate theoretical energy
         # TODO: energy calc for MSMR
         if self.options["open-circuit potential"] != "MSMR":
-            energy = pybamm.lithium_ion.electrode_soh.theoretical_energy_integral(
-                self.parameter_values,
-                sol_dict,
-            )
+            energy = self.theoretical_energy_integral(sol_dict)
             sol_dict.update({"Maximum theoretical energy [W.h]": energy})
         return sol_dict
 
@@ -829,6 +826,27 @@ class ElectrodeSOHSolver:
         sol = self.solve(inputs)
         return [sol["Un(x_0)"], sol["Un(x_100)"], sol["Up(y_100)"], sol["Up(y_0)"]]
 
+    def theoretical_energy_integral(self, inputs, points=1000):
+        x_0 = inputs["x_0"]
+        y_0 = inputs["y_0"]
+        x_100 = inputs["x_100"]
+        y_100 = inputs["y_100"]
+        Q_p = inputs["Q_p"]
+        x_vals = np.linspace(x_100, x_0, num=points)
+        y_vals = np.linspace(y_100, y_0, num=points)
+        # Calculate OCV at each stoichiometry
+        param = self.param
+        T = param.T_amb_av(0)
+        Vs = self.parameter_values.evaluate(
+            param.p.prim.U(y_vals, T) - param.n.prim.U(x_vals, T)
+        ).flatten()
+        # Calculate dQ
+        Q = Q_p * (y_0 - y_100)
+        dQ = Q / (points - 1)
+        # Integrate and convert to W-h
+        E = np.trapz(Vs, dx=dQ)
+        return E
+
 
 def get_initial_stoichiometries(
     initial_value,
@@ -972,7 +990,7 @@ def get_min_max_ocps(
     return esoh_solver.get_min_max_ocps()
 
 
-def theoretical_energy_integral(parameter_values, inputs, points=100):
+def theoretical_energy_integral(parameter_values, param, inputs, points=100):
     """
     Calculate maximum energy possible from a cell given OCV, initial soc, and final soc
     given voltage limits, open-circuit potentials, etc defined by parameter_values
@@ -991,30 +1009,8 @@ def theoretical_energy_integral(parameter_values, inputs, points=100):
     E
         The total energy of the cell in Wh
     """
-    x_0 = inputs["x_0"]
-    y_0 = inputs["y_0"]
-    x_100 = inputs["x_100"]
-    y_100 = inputs["y_100"]
-    Q_p = inputs["Q_p"]
-    x_vals = np.linspace(x_100, x_0, num=points)
-    y_vals = np.linspace(y_100, y_0, num=points)
-    # Calculate OCV at each stoichiometry
-    param = pybamm.LithiumIonParameters()
-    y = pybamm.standard_spatial_vars.y
-    z = pybamm.standard_spatial_vars.z
-    T = pybamm.yz_average(param.T_amb(y, z, 0))
-    Vs = np.empty(x_vals.shape)
-    for i in range(x_vals.size):
-        Vs[i] = (
-            parameter_values.evaluate(param.p.prim.U(y_vals[i], T)).item()
-            - parameter_values.evaluate(param.n.prim.U(x_vals[i], T)).item()
-        )
-    # Calculate dQ
-    Q = Q_p * (y_0 - y_100)
-    dQ = Q / (points - 1)
-    # Integrate and convert to W-h
-    E = np.trapz(Vs, dx=dQ)
-    return E
+    esoh_solver = ElectrodeSOHSolver(parameter_values, param)
+    return esoh_solver.theoretical_energy_integral(inputs, points=points)
 
 
 def calculate_theoretical_energy(
@@ -1045,6 +1041,7 @@ def calculate_theoretical_energy(
     Q_p = parameter_values.evaluate(pybamm.LithiumIonParameters().p.prim.Q_init)
     E = theoretical_energy_integral(
         parameter_values,
+        pybamm.LithiumIonParameters(),
         {"x_100": x_100, "x_0": x_0, "y_100": y_100, "y_0": y_0, "Q_p": Q_p},
         points=points,
     )
diff --git a/pybamm/parameters/lithium_ion_parameters.py b/pybamm/parameters/lithium_ion_parameters.py
index 726e876aa..c459a4ef1 100644
--- a/pybamm/parameters/lithium_ion_parameters.py
+++ b/pybamm/parameters/lithium_ion_parameters.py
@@ -50,6 +50,7 @@ class LithiumIonParameters(BaseParameters):
         self.T_ref = self.therm.T_ref
         self.T_init = self.therm.T_init
         self.T_amb = self.therm.T_amb
+        self.T_amb_av = self.therm.T_amb_av
         self.h_edge = self.therm.h_edge
         self.h_total = self.therm.h_total
         self.rho_c_p_eff = self.therm.rho_c_p_eff
diff --git a/pybamm/parameters/thermal_parameters.py b/pybamm/parameters/thermal_parameters.py
index ea1dd1206..8e92ff8d3 100644
--- a/pybamm/parameters/thermal_parameters.py
+++ b/pybamm/parameters/thermal_parameters.py
@@ -51,6 +51,12 @@ class ThermalParameters(BaseParameters):
             },
         )
 
+    def T_amb_av(self, t):
+        """YZ-averaged ambient temperature [K]"""
+        y = pybamm.standard_spatial_vars.y
+        z = pybamm.standard_spatial_vars.z
+        return pybamm.yz_average(self.T_amb(y, z, t))
+
     def h_edge(self, y, z):
         """Cell edge heat transfer coefficient [W.m-2.K-1]"""
         inputs = {


Test patch:
diff --git a/tests/unit/test_models/test_full_battery_models/test_lithium_ion/test_electrode_soh.py b/tests/unit/test_models/test_full_battery_models/test_lithium_ion/test_electrode_soh.py
index 628017d5d..e5e79a6ae 100644
--- a/tests/unit/test_models/test_full_battery_models/test_lithium_ion/test_electrode_soh.py
+++ b/tests/unit/test_models/test_full_battery_models/test_lithium_ion/test_electrode_soh.py
@@ -40,9 +40,7 @@ class TestElectrodeSOH(TestCase):
                     k: sol_split[k].data[0]
                     for k in ["x_0", "y_0", "x_100", "y_100", "Q_p"]
                 }
-                energy = pybamm.lithium_ion.electrode_soh.theoretical_energy_integral(
-                    parameter_values, inputs
-                )
+                energy = esoh_solver.theoretical_energy_integral(inputs)
                 self.assertAlmostEqual(sol[key], energy, places=5)
 
         # should still work with old inputs