Additional tests for idaklu solver with output_variables

pybamm/solvers/idaklu_solver.py

@@ -678,7 +678,7 @@ def _integrate(self, model, t_eval, inputs_dict=None):
                     if isinstance(
                         model.variables_and_events[var], pybamm.ExplicitTimeIntegral
                     ):
-                        continue
+                        continue  # pragma: no cover
                     len_of_var = (
                         self._setup["var_casadi_fcns"][var](0, 0, 0).sparsity().nnz()
                     )

pybamm/solvers/processed_variable_var.py

@@ -116,22 +116,21 @@ def _unroll_nnz(self, realdata=None):
         # unroll in nnz != numel, otherwise copy
         if realdata is None:
             realdata = self.base_variables_data
-        # sp = self.base_variables_casadi[0](0, 0, 0).sparsity()
-        # if sp.nnz() != sp.numel():
-        #     data = [None] * len(realdata)
-        #     for datak in range(len(realdata)):
-        #         data[datak] = np.zeros(self.base_eval_shape[0] * len(self.t_pts))
-        #         var_data = realdata[0].flatten()
-        #         k = 0
-        #         for t_i in range(len(self.t_pts)):
-        #             base = t_i * sp.numel()
-        #             for r in sp.row():
-        #                 data[datak][base + r] = var_data[k]
-        #                 k = k + 1
-        # else:
-        #     data = realdata
-        # return data
-        return realdata
+        sp = self.base_variables_casadi[0](0, 0, 0).sparsity()
+        if sp.nnz() != sp.numel():
+            data = [None] * len(realdata)
+            for datak in range(len(realdata)):
+                data[datak] = np.zeros(self.base_eval_shape[0] * len(self.t_pts))
+                var_data = realdata[0].flatten()
+                k = 0
+                for t_i in range(len(self.t_pts)):
+                    base = t_i * sp.numel()
+                    for r in sp.row():
+                        data[datak][base + r] = var_data[k]
+                        k = k + 1
+        else:
+            data = realdata
+        return data
 
     def unroll_0D(self, realdata=None):
         if realdata is None:

tests/integration/test_solvers/test_idaklu.py

@@ -90,6 +90,78 @@ def test_changing_grid(self):
             # solve
             solver.solve(model_disc, t_eval)
 
+    def test_with_output_variables(self):
+        # Construct a model and solve for all vairables, then test
+        # the 'output_variables' option for each variable in turn, confirming
+        # equivalence
+
+        # construct model
+        model = pybamm.lithium_ion.DFN()
+        geometry = model.default_geometry
+        param = model.default_parameter_values
+        input_parameters = {}
+        param.update({key: "[input]" for key in input_parameters})
+        param.process_model(model)
+        param.process_geometry(geometry)
+        var_pts = {"x_n": 50, "x_s": 50, "x_p": 50, "r_n": 5, "r_p": 5}
+        mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts)
+        disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
+        disc.process_model(model)
+        t_eval = np.linspace(0, 3600, 100)
+
+        options = {
+            'linear_solver': 'SUNLinSol_KLU',
+            'jacobian': 'sparse',
+            'num_threads': 4,
+        }
+
+        # Select all output_variables (NB: This can be very slow to run all solves)
+        # output_variables = [m for m, (k, v) in
+        #                     zip(model.variable_names(), model.variables.items())
+        #                     if not isinstance(v, pybamm.ExplicitTimeIntegral)]
+
+        # Use a selection of variables (of different types)
+        output_variables = [
+            "Voltage [V]",
+            "Time [min]",
+            "Current [A]",
+            "r_n [m]",
+            "x [m]",
+            "Gradient of negative electrolyte potential [V.m-1]",
+            "Negative particle flux [mol.m-2.s-1]",
+            "Discharge capacity [A.h]",
+            "Throughput capacity [A.h]",
+        ]
+
+        solver_all = pybamm.IDAKLUSolver(
+            atol=1e-8, rtol=1e-8,
+            options=options,
+        )
+        sol_all = solver_all.solve(
+            model,
+            t_eval,
+            inputs=input_parameters,
+            calculate_sensitivities=True,
+        )
+
+        for varname in output_variables:
+            print(varname)
+            solver = pybamm.IDAKLUSolver(
+                atol=1e-8, rtol=1e-8,
+                options=options,
+                output_variables=[varname],
+            )
+
+            sol = solver.solve(
+                model,
+                t_eval,
+                inputs=input_parameters,
+                calculate_sensitivities=True,
+            )
+
+            # Compare output to sol_all
+            self.assertTrue(np.allclose(sol[varname].data, sol_all[varname].data))
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")

tests/unit/test_solvers/test_idaklu_solver.py

@@ -76,22 +76,6 @@ def test_model_events(self):
             # create discretisation
             disc = pybamm.Discretisation()
             model_disc = disc.process_model(model, inplace=False)
-            # Invalid atol (dict) raises error, valid options are float or ndarray
-            self.assertRaises(
-                pybamm.SolverError,
-                pybamm.IDAKLUSolver(
-                    rtol=1e-8, atol={'key': 'value'},
-                    root_method=root_method)
-            )
-            # output_variables only valid with convert_to_format=='casadi'
-            if form == "python" or form == "jax":
-                self.assertRaises(
-                    pybamm.SolverError,
-                    pybamm.IDAKLUSolver(
-                        rtol=1e-8, atol=1e-8,
-                        output_variables=['var'],
-                        root_method=root_method)
-                )
             # Solve
             solver = pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-8, root_method=root_method)
             t_eval = np.linspace(0, 1, 100)
@@ -101,6 +85,10 @@ def test_model_events(self):
                 solution.y[0], np.exp(0.1 * solution.t), decimal=5
             )
 
+            # Check invalid atol type raises an error
+            with self.assertRaises(pybamm.SolverError):
+                solver._check_atol_type({'key': 'value'}, [])
+
             # enforce events that won't be triggered
             model.events = [pybamm.Event("an event", var + 1)]
             model_disc = disc.process_model(model, inplace=False)
@@ -211,7 +199,6 @@ def test_sensitivites_initial_condition(self):
 
             disc = pybamm.Discretisation()
             disc.process_model(model)
-
             solver = pybamm.IDAKLUSolver(output_variables=output_variables)
 
             t_eval = np.linspace(0, 3, 100)

tests/unit/test_solvers/test_processed_variable_var.py

@@ -32,7 +32,7 @@ def to_casadi(var_pybamm, y, inputs=None):
 
 
 def process_and_check_2D_variable(
-    var, first_spatial_var, second_spatial_var, disc=None
+    var, first_spatial_var, second_spatial_var, disc=None, geometry_options={}
 ):
     # first_spatial_var should be on the "smaller" domain, i.e "r" for an "r-x" variable
     if disc is None:
@@ -49,20 +49,20 @@ def process_and_check_2D_variable(
     y_sol = np.ones(len(second_sol) * len(first_sol))[:, np.newaxis] * np.linspace(0, 5)
 
     var_casadi = to_casadi(var_sol, y_sol)
-    processed_var = pybamm.ProcessedVariableVar(
+    model = tests.get_base_model_with_battery_geometry(**geometry_options)
+    pybamm.ProcessedVariableVar(
         [var_sol],
         [var_casadi],
         [y_sol],
-        pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}),
+        pybamm.Solution(t_sol, y_sol, model, {}),
         warn=False,
     )
-    # Ordering from idaklu with output_variables set is different to
-    # the full solver
-    y_sol = y_sol.reshape((y_sol.shape[1], y_sol.shape[0])).transpose()
-    np.testing.assert_array_equal(
-        processed_var.entries,
-        np.reshape(y_sol, [len(first_sol), len(second_sol), len(t_sol)]),
-    )
+    # NB: ProcessedVariableVar does not interpret y in the same way as
+    #  ProcessedVariable; a better test of equivalence is to check that the
+    #  results are the same between IDAKLUSolver with (and without)
+    #  output_variables. This is implemented in the integration test:
+    #    tests/integration/test_solvers/test_idaklu_solver.py
+    #    ::test_output_variables
     return y_sol, first_sol, second_sol, t_sol
 
 
@@ -175,18 +175,17 @@ def test_processed_variable_1D(self):
         processed_var.mesh.nodes, processed_var.mesh.edges = \
             processed_var.mesh.edges, processed_var.mesh.nodes
 
-        # Check no errors with domain-specific attributes
-        # (see ProcessedVariableVar.initialise_2D() for details)
+        # Check that there are no errors with domain-specific attributes
+        #  (see ProcessedVariableVar.initialise_1D() for details)
         domain_list = [
-            ["particle", "electrode"],
-            ["separator", "current collector"],
-            ["particle", "particle size"],
-            ["particle size", "electrode"],
-            ["particle size", "current collector"]
+            "particle",
+            "separator",
+            "current collector",
+            "particle size",
+            "random-non-specific-domain",
         ]
-        for domain, secondary in domain_list:
+        for domain in domain_list:
             processed_var.domain[0] = domain
-            processed_var.domains["secondary"] = [secondary]
             processed_var.initialise_1D()
 
     def test_processed_variable_1D_unknown_domain(self):
@@ -217,6 +216,126 @@ def test_processed_variable_1D_unknown_domain(self):
         c_casadi = to_casadi(c, y_sol)
         pybamm.ProcessedVariableVar([c], [c_casadi], [y_sol], solution, warn=False)
 
+    def test_processed_variable_2D_x_r(self):
+        var = pybamm.Variable(
+            "var",
+            domain=["negative particle"],
+            auxiliary_domains={"secondary": ["negative electrode"]},
+        )
+        x = pybamm.SpatialVariable("x", domain=["negative electrode"])
+        r = pybamm.SpatialVariable(
+            "r",
+            domain=["negative particle"],
+            auxiliary_domains={"secondary": ["negative electrode"]},
+        )
+
+        disc = tests.get_p2d_discretisation_for_testing()
+        process_and_check_2D_variable(var, r, x, disc=disc)
+
+    def test_processed_variable_2D_R_x(self):
+        var = pybamm.Variable(
+            "var",
+            domain=["negative particle size"],
+            auxiliary_domains={"secondary": ["negative electrode"]},
+        )
+        R = pybamm.SpatialVariable(
+            "R",
+            domain=["negative particle size"],
+            auxiliary_domains={"secondary": ["negative electrode"]},
+        )
+        x = pybamm.SpatialVariable("x", domain=["negative electrode"])
+
+        disc = tests.get_size_distribution_disc_for_testing()
+        process_and_check_2D_variable(
+            var,
+            R,
+            x,
+            disc=disc,
+            geometry_options={"options": {"particle size": "distribution"}},
+        )
+
+    def test_processed_variable_2D_R_z(self):
+        var = pybamm.Variable(
+            "var",
+            domain=["negative particle size"],
+            auxiliary_domains={"secondary": ["current collector"]},
+        )
+        R = pybamm.SpatialVariable(
+            "R",
+            domain=["negative particle size"],
+            auxiliary_domains={"secondary": ["current collector"]},
+        )
+        z = pybamm.SpatialVariable("z", domain=["current collector"])
+
+        disc = tests.get_size_distribution_disc_for_testing()
+        process_and_check_2D_variable(
+            var,
+            R,
+            z,
+            disc=disc,
+            geometry_options={"options": {"particle size": "distribution"}},
+        )
+
+    def test_processed_variable_2D_r_R(self):
+        var = pybamm.Variable(
+            "var",
+            domain=["negative particle"],
+            auxiliary_domains={"secondary": ["negative particle size"]},
+        )
+        r = pybamm.SpatialVariable(
+            "r",
+            domain=["negative particle"],
+            auxiliary_domains={"secondary": ["negative particle size"]},
+        )
+        R = pybamm.SpatialVariable("R", domain=["negative particle size"])
+
+        disc = tests.get_size_distribution_disc_for_testing()
+        process_and_check_2D_variable(
+            var,
+            r,
+            R,
+            disc=disc,
+            geometry_options={"options": {"particle size": "distribution"}},
+        )
+
+    def test_processed_variable_2D_x_z(self):
+        var = pybamm.Variable(
+            "var",
+            domain=["negative electrode", "separator"],
+            auxiliary_domains={"secondary": "current collector"},
+        )
+        x = pybamm.SpatialVariable(
+            "x",
+            domain=["negative electrode", "separator"],
+            auxiliary_domains={"secondary": "current collector"},
+        )
+        z = pybamm.SpatialVariable("z", domain=["current collector"])
+
+        disc = tests.get_1p1d_discretisation_for_testing()
+        y_sol, x_sol, z_sol, t_sol = process_and_check_2D_variable(var, x, z, disc=disc)
+        del x_sol
+
+        # On edges
+        x_s_edge = pybamm.Matrix(
+            np.tile(disc.mesh["separator"].edges, len(z_sol)),
+            domain="separator",
+            auxiliary_domains={"secondary": "current collector"},
+        )
+        x_s_edge.mesh = disc.mesh["separator"]
+        x_s_edge.secondary_mesh = disc.mesh["current collector"]
+        x_s_casadi = to_casadi(x_s_edge, y_sol)
+        processed_x_s_edge = pybamm.ProcessedVariable(
+            [x_s_edge],
+            [x_s_casadi],
+            pybamm.Solution(
+                t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {}
+            ),
+            warn=False,
+        )
+        np.testing.assert_array_equal(
+            x_s_edge.entries.flatten(), processed_x_s_edge.entries[:, :, 0].T.flatten()
+        )
+
     def test_processed_variable_2D_space_only(self):
         var = pybamm.Variable(
             "var",