Issue 762 casadi mass

pybamm/discretisations/discretisation.py

@@ -4,7 +4,8 @@
 import pybamm
 import numpy as np
 from collections import defaultdict, OrderedDict
-from scipy.sparse import block_diag, csr_matrix
+from scipy.sparse import block_diag, csc_matrix, csr_matrix
+from scipy.sparse.linalg import inv
 
 
 def has_bc_of_form(symbol, side, bcs, form):
@@ -206,7 +207,9 @@ def process_model(self, model, inplace=True, check_model=True):
 
         # Create mass matrix
         pybamm.logger.info("Create mass matrix for {}".format(model.name))
-        model_disc.mass_matrix = self.create_mass_matrix(model_disc)
+        model_disc.mass_matrix, model_disc.mass_matrix_inv = self.create_mass_matrix(
+            model_disc
+        )
 
         # Check that resulting model makes sense
         if check_model:
@@ -532,10 +535,14 @@ def create_mass_matrix(self, model):
         -------
         :class:`pybamm.Matrix`
             The mass matrix
+        :class:`pybamm.Matrix`
+            The inverse of the ode part of the mass matrix (required by solvers
+            which only accept the ODEs in explicit form)
         """
         # Create list of mass matrices for each equation to be put into block
         # diagonal mass matrix for the model
         mass_list = []
+        mass_inv_list = []
 
         # get a list of model rhs variables that are sorted according to
         # where they are in the state vector
@@ -560,12 +567,26 @@ def create_mass_matrix(self, model):
             if var.domain == []:
                 # If variable domain empty then mass matrix is just 1
                 mass_list.append(1.0)
+                mass_inv_list.append(1.0)
             else:
-                mass_list.append(
+                mass = (
                     self.spatial_methods[var.domain[0]]
                     .mass_matrix(var, self.bcs)
                     .entries
                 )
+                mass_list.append(mass)
+                if isinstance(
+                    self.spatial_methods[var.domain[0]],
+                    (pybamm.ZeroDimensionalMethod, pybamm.FiniteVolume),
+                ):
+                    # for 0D methods the mass matrix is just a scalar 1 and for
+                    # finite volumes the mass matrix is identity, so no need to
+                    # compute the inverse
+                    mass_inv_list.append(mass)
+                else:
+                    # inverse is more efficient in csc format
+                    mass_inv = inv(csc_matrix(mass))
+                    mass_inv_list.append(mass_inv)
 
         # Create lumped mass matrix (of zeros) of the correct shape for the
         # discretised algebraic equations
@@ -574,10 +595,14 @@ def create_mass_matrix(self, model):
             mass_algebraic = csr_matrix((mass_algebraic_size, mass_algebraic_size))
             mass_list.append(mass_algebraic)
 
-        # Create block diagonal (sparse) mass matrix
-        mass_matrix = block_diag(mass_list, format="csr")
+        # Create block diagonal (sparse) mass matrix and inverse (if model has odes)
+        mass_matrix = pybamm.Matrix(block_diag(mass_list, format="csr"))
+        if model.rhs.keys():
+            mass_matrix_inv = pybamm.Matrix(block_diag(mass_inv_list, format="csr"))
+        else:
+            mass_matrix_inv = None
 
-        return pybamm.Matrix(mass_matrix)
+        return mass_matrix, mass_matrix_inv
 
     def create_jacobian(self, model):
         """Creates Jacobian of the discretised model.

pybamm/models/base_model.py

@@ -60,6 +60,9 @@ class BaseModel(object):
     mass_matrix : :class:`pybamm.Matrix`
         After discretisation, contains the mass matrix for the model. This is computed
         automatically
+    mass_matrix_inv : :class:`pybamm.Matrix`
+        After discretisation, contains the inverse mass matrix for the differential
+        (rhs) part of model. This is computed automatically
     jacobian : :class:`pybamm.Concatenation`
         Contains the Jacobian for the model. If model.use_jacobian is True, the
         Jacobian is computed automatically during solver set up
@@ -88,7 +91,7 @@ class BaseModel(object):
         - "casadi": convert into CasADi expression tree, which then uses CasADi's \
         algorithm to calculate the Jacobian.
 
-        Default is "python".
+        Default is "casadi".
 
     """
 
@@ -107,6 +110,7 @@ def __init__(self, name="Unnamed model"):
         self._concatenated_algebraic = None
         self._concatenated_initial_conditions = None
         self._mass_matrix = None
+        self._mass_matrix_inv = None
         self._jacobian = None
         self._jacobian_algebraic = None
         self.external_variables = []
@@ -249,6 +253,14 @@ def mass_matrix(self):
     def mass_matrix(self, mass_matrix):
         self._mass_matrix = mass_matrix
 
+    @property
+    def mass_matrix_inv(self):
+        return self._mass_matrix_inv
+
+    @mass_matrix_inv.setter
+    def mass_matrix_inv(self, mass_matrix_inv):
+        self._mass_matrix_inv = mass_matrix_inv
+
     @property
     def jacobian(self):
         return self._jacobian

pybamm/solvers/casadi_solver.py

@@ -93,11 +93,11 @@ def solve(self, model, t_eval, inputs=None):
         """
         if self.mode == "fast":
             # Solve model normally by calling the solve method from parent class
-            return super().solve(model, t_eval)
+            return super().solve(model, t_eval, inputs=inputs)
         elif model.events == {}:
             pybamm.logger.info("No events found, running fast mode")
             # Solve model normally by calling the solve method from parent class
-            return super().solve(model, t_eval)
+            return super().solve(model, t_eval, inputs=inputs)
         elif self.mode == "safe":
             # Step-and-check
             timer = pybamm.Timer()
@@ -122,7 +122,7 @@ def solve(self, model, t_eval, inputs=None):
                     # different to t_eval, but shouldn't matter too much as it should
                     # only happen near events.
                     try:
-                        current_step_sol = self.step(model, dt)
+                        current_step_sol = self.step(model, dt, inputs=inputs)
                         solved = True
                     except pybamm.SolverError:
                         dt /= 2
@@ -199,12 +199,7 @@ def compute_solution(self, model, t_eval, inputs=None):
         solve_start_time = timer.time()
         pybamm.logger.debug("Calling DAE solver")
         solution = self.integrate_casadi(
-            self.casadi_rhs,
-            self.casadi_algebraic,
-            self.y0,
-            t_eval,
-            inputs,
-            mass_matrix=model.mass_matrix.entries,
+            self.casadi_rhs, self.casadi_algebraic, self.y0, t_eval, inputs
         )
         solve_time = timer.time() - solve_start_time
 
@@ -213,9 +208,7 @@ def compute_solution(self, model, t_eval, inputs=None):
 
         return solution, solve_time, termination
 
-    def integrate_casadi(
-        self, rhs, algebraic, y0, t_eval, inputs=None, mass_matrix=None
-    ):
+    def integrate_casadi(self, rhs, algebraic, y0, t_eval, inputs=None):
         """
         Solve a DAE model defined by residuals with initial conditions y0.
 
@@ -230,10 +223,6 @@ def integrate_casadi(
             The times at which to compute the solution
         inputs : dict, optional
             Any input parameters to pass to the model when solving
-        mass_matrix : array_like, optional
-            The (sparse) mass matrix for the chosen spatial method. This is only passed
-            to check that the mass matrix is diagonal with 1s for the odes and 0s for
-            the algebraic equations, as CasADi does not allow to pass mass matrices.
         """
         inputs = inputs or {}
         options = {

pybamm/solvers/dae_solver.py

@@ -212,6 +212,8 @@ def get_event_class(event):
         self.event_funs = [get_event_class(event) for event in events.values()]
         self.jacobian = jacobian
 
+        pybamm.logger.info("Finish solver set-up")
+
     def set_up_casadi(self, model, inputs=None):
         """Convert model to casadi format and use their inbuilt functionalities.
 
@@ -336,8 +338,17 @@ def get_event_class(event):
         self.jacobian = jacobian
 
         # Save CasADi functions for the CasADi solver
-        self.casadi_rhs = concatenated_rhs_fn
-        self.casadi_algebraic = concatenated_algebraic_fn
+        # Note: when we pass to casadi the ode part of the problem must be in explicit
+        # form so we pre-multiply by the inverse of the mass matrix
+        if isinstance(self, pybamm.CasadiSolver):
+            mass_matrix_inv = casadi.MX(model.mass_matrix_inv.entries)
+            explicit_rhs = mass_matrix_inv @ concatenated_rhs
+            self.casadi_rhs = casadi.Function(
+                "rhs", [t_casadi, y_casadi_w_ext, u_casadi_stacked], [explicit_rhs]
+            )
+            self.casadi_algebraic = concatenated_algebraic_fn
+
+        pybamm.logger.info("Finish solver set-up")
 
     def set_inputs_and_external(self, inputs):
         """

pybamm/solvers/ode_solver.py

@@ -151,6 +151,8 @@ def get_event_class(event):
         self.event_funs = [get_event_class(event) for event in events.values()]
         self.jacobian = jacobian
 
+        pybamm.logger.info("Finish solver set-up")
+
     def set_up_casadi(self, model, inputs=None):
         """Convert model to casadi format and use their inbuilt functionalities.
 

tests/unit/test_discretisations/test_discretisation.py

@@ -9,10 +9,12 @@
     get_mesh_for_testing,
     get_discretisation_for_testing,
     get_1p1d_discretisation_for_testing,
+    get_2p1d_mesh_for_testing,
 )
 from tests.shared import SpatialMethodForTesting
 
-from scipy.sparse import block_diag
+from scipy.sparse import block_diag, csc_matrix
+from scipy.sparse.linalg import inv
 
 
 class TestDiscretise(unittest.TestCase):
@@ -934,6 +936,37 @@ def test_process_with_no_check(self):
         disc = get_discretisation_for_testing()
         disc.process_model(model, check_model=False)
 
+    def test_mass_matirx_inverse(self):
+        # get mesh
+        mesh = get_2p1d_mesh_for_testing(ypts=5, zpts=5)
+        spatial_methods = {
+            "macroscale": pybamm.FiniteVolume(),
+            "current collector": pybamm.ScikitFiniteElement(),
+        }
+        # create model
+        a = pybamm.Variable("a", domain="negative electrode")
+        b = pybamm.Variable("b", domain="current collector")
+        model = pybamm.BaseModel()
+        model.rhs = {a: pybamm.Laplacian(a), b: 4 * pybamm.Laplacian(b)}
+        model.initial_conditions = {a: pybamm.Scalar(3), b: pybamm.Scalar(10)}
+        model.boundary_conditions = {
+            a: {"left": (0, "Neumann"), "right": (0, "Neumann")},
+            b: {"negative tab": (0, "Neumann"), "positive tab": (0, "Neumann")},
+        }
+        model.variables = {"a": a, "b": b}
+
+        # create discretisation
+        disc = pybamm.Discretisation(mesh, spatial_methods)
+        disc.process_model(model)
+
+        # test that computing mass matrix block-by-block (as is done during
+        # discretisation) gives the correct result
+        # Note: inverse is more efficient in csc format
+        mass_inv = inv(csc_matrix(model.mass_matrix.entries))
+        np.testing.assert_equal(
+            model.mass_matrix_inv.entries.toarray(), mass_inv.toarray()
+        )
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")

tests/unit/test_solvers/test_casadi_solver.py

@@ -6,6 +6,7 @@
 import unittest
 import numpy as np
 from tests import get_mesh_for_testing, get_discretisation_for_testing
+from scipy.sparse import eye
 
 
 class TestCasadiSolver(unittest.TestCase):
@@ -206,12 +207,41 @@ def test_model_solver_with_inputs(self):
         disc = pybamm.Discretisation(mesh, spatial_methods)
         disc.process_model(model)
         # Solve
-        solver = pybamm.ScipySolver(rtol=1e-8, atol=1e-8, method="RK45")
+        solver = pybamm.CasadiSolver(rtol=1e-8, atol=1e-8)
         t_eval = np.linspace(0, 10, 100)
         solution = solver.solve(model, t_eval, inputs={"rate": 0.1})
         self.assertLess(len(solution.t), len(t_eval))
         np.testing.assert_array_equal(solution.t, t_eval[: len(solution.t)])
-        np.testing.assert_allclose(solution.y[0], np.exp(-0.1 * solution.t))
+        np.testing.assert_allclose(solution.y[0], np.exp(-0.1 * solution.t), rtol=1e-06)
+
+    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(int(mass_matrix.shape[0] / 2))
+        model.mass_matrix_inv = pybamm.Matrix(mass_matrix_inv)
+
+        # Solve
+        solver = pybamm.CasadiSolver(rtol=1e-8, atol=1e-8, method="idas")
+        t_eval = np.linspace(0, 1, 100)
+        solution = solver.solve(model, t_eval)
+        np.testing.assert_array_equal(solution.t, t_eval)
+        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))
 
 
 if __name__ == "__main__":