Repository: pybamm-team/PyBaMM
Pull request #4879 (2025-02-26 20:17:16)

Fix IDAKLU bug with non-identity inverse mass matrices
# Description

For more info, see: https://github.com/pybamm-team/PyBaMM/pull/4878#discussion_r1972113603

This is a minor bug that practically only affects FEM simulations using the `IDAKLUSolver`

I will make a followup PR in `pybammsolvers` to address a related issue there: https://github.com/pybamm-team/pybammsolvers/issues/31

## Type of change

Please add a line in the relevant section of [CHANGELOG.md](https://github.com/pybamm-team/PyBaMM/blob/develop/CHANGELOG.md) to document the change (include PR #)

# Important checks:

Please confirm the following before marking the PR as ready for review:
- No style issues: `nox -s pre-commit`
- All tests pass: `nox -s tests`
- The documentation builds: `nox -s doctests`
- Code is commented for hard-to-understand areas
- Tests added that prove fix is effective or that feature works


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pybamm-team/PyBaMM/pull/4879?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pybamm-team) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 98.71%. Comparing base [(`6f20992`)](https://app.codecov.io/gh/pybamm-team/PyBaMM/commit/6f20992513b219e5753d62d690c38a45029fd537?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pybamm-team) to head [(`6951398`)](https://app.codecov.io/gh/pybamm-team/PyBaMM/commit/695139880c2ff6a13a8e4d37bbb11250e103d6e0?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pybamm-team).
> Report is 1 commits behind head on develop.

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff            @@
##           develop    #4879   +/-   ##
========================================
  Coverage    98.70%   98.71%           
========================================
  Files          304      304           
  Lines        23483    23495   +12     
========================================
+ Hits         23180    23192   +12     
  Misses         303      303           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pybamm-team/PyBaMM/pull/4879?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pybamm-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pybamm-team).


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3987ec38f..0ae99152c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,6 +16,7 @@
 
 ## Bug fixes
 
+- Fixed a bug in simulating FEM models with the `IDAKLUSolver`. ([#4879](https://github.com/pybamm-team/PyBaMM/pull/4879))
 - Moved concentration inside x-averaged when calculating LLI due to LAM variables ([#4858](https://github.com/pybamm-team/PyBaMM/pull/4858))
 - Fixed a bug that caused the variable `"Loss of lithium due to {domain} lithium plating"`to have the domain `"current collector"` (should not have any domain at all) if the `"x-average side reactions"` option was set to `"true"`. ([#4844](https://github.com/pybamm-team/PyBaMM/pull/4844))
 - Fixed interpolation bug in `pybamm.QuickPlot` with spatial variables. ([#4841](https://github.com/pybamm-team/PyBaMM/pull/4841))
diff --git a/src/pybamm/models/base_model.py b/src/pybamm/models/base_model.py
index a4e58600b..b5670320b 100644
--- a/src/pybamm/models/base_model.py
+++ b/src/pybamm/models/base_model.py
@@ -69,6 +69,7 @@ class BaseModel:
         self._parameters = None
         self._input_parameters = None
         self._parameter_info = None
+        self._is_standard_form_dae = None
         self._variables_casadi = {}
         self._geometry = pybamm.Geometry({})
 
@@ -429,6 +430,16 @@ class BaseModel:
             self._input_parameters = self._find_symbols(pybamm.InputParameter)
         return self._input_parameters
 
+    @property
+    def is_standard_form_dae(self):
+        """
+        Check if the model is a DAE in standard form with a mass matrix that is all
+        zeros except for along the diagonal, which is either ones or zeros.
+        """
+        if self._is_standard_form_dae is None:
+            self._is_standard_form_dae = self._check_standard_form_dae()
+        return self._is_standard_form_dae
+
     @property
     def calc_esoh(self):
         """Whether to include eSOH variables in the summary variables."""
@@ -549,6 +560,30 @@ class BaseModel:
 
         return max_name_length, max_type_length
 
+    def _check_standard_form_dae(self):
+        """
+        Check if the model is a DAE in standard form with a mass matrix that is all
+        zeros except for along the diagonal, which is either ones or zeros.
+
+        For example, the following is standard form:
+        M*y' = f(y, y', t)
+
+        M = [I 0
+             0 0]
+
+        The following explicit ODE is also a standard form DAE:
+        M = I
+
+        The following is not standard form:
+        M = [2I 0
+             0 0]
+        """
+        if self.mass_matrix is None or self.mass_matrix_inv is None:
+            return False
+        # Check that the mass matrix inverse is an identity matrix
+        mass_matrix_inv = self.mass_matrix_inv.entries.toarray()
+        return np.allclose(mass_matrix_inv, np.eye(mass_matrix_inv.shape[0]))
+
     def _format_table_row(
         self, param_name, param_type, max_name_length, max_type_length
     ):
diff --git a/src/pybamm/solvers/idaklu_solver.py b/src/pybamm/solvers/idaklu_solver.py
index 30149ea0a..4150818ac 100644
--- a/src/pybamm/solvers/idaklu_solver.py
+++ b/src/pybamm/solvers/idaklu_solver.py
@@ -633,6 +633,7 @@ class IDAKLUSolver(pybamm.BaseSolver):
             "ids": ids,  # array
             "sensitivity_names": sensitivity_names,
             "number_of_sensitivity_parameters": number_of_sensitivity_parameters,
+            "standard_form_dae": model.is_standard_form_dae,  # bool
             "output_variables": self.output_variables,
             "var_fcns": self.computed_var_fcns,
             "var_idaklu_fcns": self.var_idaklu_fcns,
@@ -990,8 +991,12 @@ class IDAKLUSolver(pybamm.BaseSolver):
         rhs0 = rhs0.flatten()
 
         # for the differential terms, ydot = M^-1 * (rhs)
-        # M^-1 is the identity matrix, so we can just use rhs
-        ydot0[: model.len_rhs] = rhs0
+        if model.is_standard_form_dae:
+            # M^-1 is the identity matrix, so we can just use rhs
+            ydot0[: model.len_rhs] = rhs0
+        else:
+            # M^-1 is not the identity matrix, so we need to use the mass matrix
+            ydot0[: model.len_rhs] = model.mass_matrix_inv.entries @ rhs0
 
         return ydot0
 


Test patch:
diff --git a/tests/unit/test_discretisations/test_discretisation.py b/tests/unit/test_discretisations/test_discretisation.py
index 270a2eaa3..776827638 100644
--- a/tests/unit/test_discretisations/test_discretisation.py
+++ b/tests/unit/test_discretisations/test_discretisation.py
@@ -470,6 +470,8 @@ class TestDiscretise:
             np.eye(submesh.nodes.shape[0]), model.mass_matrix.entries.toarray()
         )
 
+        assert model.is_standard_form_dae
+
         # Create StateVector to differentiate model with respect to
         y = pybamm.StateVector(slice(0, submesh.npts))
 
@@ -527,6 +529,8 @@ class TestDiscretise:
             np.eye(np.size(y0)), model.mass_matrix.entries.toarray()
         )
 
+        assert model.is_standard_form_dae
+
         # Create StateVector to differentiate model with respect to
         y = pybamm.StateVector(slice(0, np.size(y0)))
 
@@ -640,6 +644,7 @@ class TestDiscretise:
         np.testing.assert_array_equal(
             mass.toarray(), model.mass_matrix.entries.toarray()
         )
+        assert model.is_standard_form_dae
 
         # jacobian
         y = pybamm.StateVector(slice(0, np.size(y0)))
@@ -723,6 +728,8 @@ class TestDiscretise:
         mass = np.zeros((np.size(submesh.nodes), np.size(submesh.nodes)))
         np.testing.assert_array_equal(mass, model.mass_matrix.entries.toarray())
 
+        assert not model.is_standard_form_dae
+
         # jacobian
         y = pybamm.StateVector(slice(0, np.size(y0)))
         jacobian = model.concatenated_algebraic.jac(y).evaluate(0, y0)
diff --git a/tests/unit/test_solvers/test_casadi_solver.py b/tests/unit/test_solvers/test_casadi_solver.py
index d80af4944..79b562e2e 100644
--- a/tests/unit/test_solvers/test_casadi_solver.py
+++ b/tests/unit/test_solvers/test_casadi_solver.py
@@ -476,6 +476,8 @@ class TestCasadiSolver:
         mass_matrix_inv = 0.1 * eye(int(mass_matrix.shape[0] / 2))
         model.mass_matrix_inv = pybamm.Matrix(mass_matrix_inv)
 
+        assert not model.is_standard_form_dae
+
         # Solve
         solver = pybamm.CasadiSolver(rtol=1e-8, atol=1e-8)
         t_eval = np.linspace(0, 1, 100)
diff --git a/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py
index fd421e5a7..05ed69217 100644
--- a/tests/unit/test_solvers/test_idaklu_solver.py
+++ b/tests/unit/test_solvers/test_idaklu_solver.py
@@ -4,6 +4,7 @@ import pytest
 import numpy as np
 
 import pybamm
+from scipy.sparse import eye
 from tests import get_discretisation_for_testing
 
 
@@ -1265,3 +1266,35 @@ class TestIDAKLUSolver:
 
         with pytest.warns(pybamm.SolverWarning, match="extrapolation occurred for"):
             solver.solve(model, t_eval=[0, 1])
+
+    def test_model_solver_with_non_identity_mass(self):
+        model = pybamm.BaseModel()
+        var1 = pybamm.Variable("var1", domain="negative electrode")
+        var2 = pybamm.Variable("var2", domain="negative electrode")
+        model.rhs = {var1: var1}
+        model.algebraic = {var2: 2 * var1 - var2}
+        model.initial_conditions = {var1: 1, var2: 2}
+        disc = get_discretisation_for_testing()
+        disc.process_model(model)
+
+        # FV discretisation has identity mass. Manually set the mass matrix to
+        # be a diag of 10s here for testing. Note that the algebraic part is all
+        # zeros
+        mass_matrix = 10 * model.mass_matrix.entries
+        model.mass_matrix = pybamm.Matrix(mass_matrix)
+
+        # Note that mass_matrix_inv is just the inverse of the ode block of the
+        # mass matrix
+        mass_matrix_inv = 0.1 * eye(mass_matrix.shape[0] // 2)
+        model.mass_matrix_inv = pybamm.Matrix(mass_matrix_inv)
+
+        assert not model.is_standard_form_dae
+
+        # Solve
+        solver = pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-8)
+        t_interp = np.linspace(0, 1, 100)
+        t_eval = [t_interp[0], t_interp[-1]]
+        solution = solver.solve(model, t_eval, t_interp=t_interp)
+        np.testing.assert_array_equal(solution.t, t_interp)
+        np.testing.assert_allclose(solution.y[0], np.exp(0.1 * solution.t))
+        np.testing.assert_allclose(solution.y[-1], 2 * np.exp(0.1 * solution.t))