#1477 improve coverage

Author: martinjrobins

pybamm/solvers/base_solver.py

@@ -607,17 +607,26 @@ def jacp(*args, **kwargs):
             )
 
             # if we have changed the equations to include the explicit sensitivity
-            # equations, then we also need to update the mass matrix
+            # equations, then we also need to update the mass matrix and bounds
             if calculate_sensitivities_explicit:
-                n_inputs = model.len_rhs_sens // model.len_rhs
+                if model.len_rhs != 0:
+                    n_inputs = model.len_rhs_sens // model.len_rhs
+                elif model.len_alg != 0:
+                    n_inputs = model.len_alg_sens // model.len_alg
+                model.bounds = (
+                    np.repeat(model.bounds[0], n_inputs + 1),
+                    np.repeat(model.bounds[1], n_inputs + 1),
+                )
                 if (model.mass_matrix is not None
                         and model.mass_matrix.shape[0] == model.len_rhs_and_alg):
-                    model.mass_matrix_inv = pybamm.Matrix(
-                        block_diag(
-                            [model.mass_matrix_inv.entries] * (n_inputs + 1),
-                            format="csr"
+
+                    if model.mass_matrix_inv is not None:
+                        model.mass_matrix_inv = pybamm.Matrix(
+                            block_diag(
+                                [model.mass_matrix_inv.entries] * (n_inputs + 1),
+                                format="csr"
+                            )
                         )
-                    )
                     model.mass_matrix = pybamm.Matrix(
                         block_diag(
                             [model.mass_matrix.entries] * (n_inputs + 1), format="csr"
@@ -627,10 +636,11 @@ def jacp(*args, **kwargs):
                 # take care if calculate_sensitivites used then not used
                 if (model.mass_matrix is not None and
                         model.mass_matrix.shape[0] > model.len_rhs_and_alg):
-                    model.mass_matrix_inv = pybamm.Matrix(
-                        model.mass_matrix_inv.entries[:model.len_rhs,
-                                                      :model.len_rhs]
-                    )
+                    if model.mass_matrix_inv is not None:
+                        model.mass_matrix_inv = pybamm.Matrix(
+                            model.mass_matrix_inv.entries[:model.len_rhs,
+                                                          :model.len_rhs]
+                        )
                     model.mass_matrix = pybamm.Matrix(
                         model.mass_matrix.entries[:model.len_rhs_and_alg,
                                                   :model.len_rhs_and_alg]
@@ -1428,12 +1438,6 @@ def _set_up_ext_and_inputs(self, model, external_variables, inputs,
         for input_param in model.input_parameters:
             name = input_param.name
             if name not in inputs:
-                # Don't allow symbolic inputs if using `sensitivity`
-                if calculate_sensitivities:
-                    raise pybamm.SolverError(
-                        "Cannot have symbolic inputs if explicitly solving forward"
-                        "sensitivity equations"
-                    )
                 # Only allow symbolic inputs for CasadiSolver and CasadiAlgebraicSolver
                 if not isinstance(
                     self, (pybamm.CasadiSolver, pybamm.CasadiAlgebraicSolver)

pybamm/solvers/casadi_algebraic_solver.py

@@ -229,8 +229,11 @@ def _integrate(self, model, t_eval, inputs_dict=None):
         y_sol = casadi.vertcat(y_diff, y_alg)
 
         # Return solution object (no events, so pass None to t_event, y_event)
+
+        explicit_sensitivities = bool(model.calculate_sensitivities)
         sol = pybamm.Solution(
-            [t_eval], y_sol, model, inputs_dict, termination="success"
+            [t_eval], y_sol, model, inputs_dict, termination="success",
+            sensitivities=explicit_sensitivities
         )
         sol.integration_time = integration_time
         return sol

pybamm/solvers/casadi_solver.py

@@ -123,9 +123,6 @@ def _integrate(self, model, t_eval, inputs_dict=None):
             Any external variables or input parameters to pass to the model when solving
         """
 
-        # are we solving explicit forward equations?
-        explicit_sensitivities = bool(model.calculate_sensitivities)
-
         # Record whether there are any symbolic inputs
         inputs_dict = inputs_dict or {}
         has_symbolic_inputs = any(
@@ -275,10 +272,6 @@ def _integrate(self, model, t_eval, inputs_dict=None):
                     # update y0
                     y0 = solution.all_ys[-1][:, -1]
 
-            # now we extract sensitivities from the solution
-            if (explicit_sensitivities):
-                solution.extract_explicit_sensitivities()
-
             return solution
 
     def _solve_for_event(self, coarse_solution, init_event_signs):
@@ -544,11 +537,7 @@ def create_integrator(self, model, inputs, t_eval=None, use_event_switch=False):
             # set up and solve
             t = casadi.MX.sym("t")
             p = casadi.MX.sym("p", inputs.shape[0])
-            # If the initial conditions depend on inputs, evaluate the function
-            if isinstance(model.y0, casadi.Function):
-                y0 = model.y0(p)
-            else:
-                y0 = model.y0
+            y0 = model.y0
 
             y_diff = casadi.MX.sym("y_diff", rhs(0, y0, p).shape[0])
             y_alg = casadi.MX.sym("y_alg", algebraic(0, y0, p).shape[0])

pybamm/solvers/solution.py

@@ -182,7 +182,10 @@ def _extract_explicit_sensitivities(self, model, y, t_eval, inputs):
         n_rhs = model.len_rhs
         n_alg = model.len_alg
         # Get the point where the algebraic equations start
-        n_p = model.len_rhs_sens // model.len_rhs
+        if model.len_rhs != 0:
+            n_p = model.len_rhs_sens // model.len_rhs
+        else:
+            n_p = model.len_alg_sens // model.len_alg
         len_rhs_and_sens = model.len_rhs + model.len_rhs_sens
 
         n_t = len(t_eval)

tests/unit/test_solvers/test_casadi_solver.py

@@ -817,6 +817,38 @@ def test_solve_sensitivity_scalar_var_vector_input(self):
             np.vstack([-2 * t * np.exp(-p_eval * t) * l_n / n for t in t_eval]),
         )
 
+    def test_solve_sensitivity_then_no_sensitivity(self):
+        # Create model
+        model = pybamm.BaseModel()
+        var = pybamm.Variable("var")
+        p = pybamm.InputParameter("p")
+        model.rhs = {var: p * var}
+        model.initial_conditions = {var: 1}
+        model.variables = {"var squared": var ** 2}
+
+        # Solve
+        # Make sure that passing in extra options works
+        solver = pybamm.CasadiSolver(
+            mode="fast", rtol=1e-10, atol=1e-10
+        )
+        t_eval = np.linspace(0, 1, 80)
+        solution = solver.solve(model, t_eval, inputs={"p": 0.1},
+                                calculate_sensitivities=True)
+
+        # check sensitivities
+        np.testing.assert_allclose(
+            solution.sensitivities["p"],
+            (solution.t * np.exp(0.1 * solution.t))[:, np.newaxis],
+        )
+
+        solution = solver.solve(model, t_eval, inputs={"p": 0.1})
+
+        np.testing.assert_array_equal(solution.t, t_eval)
+        np.testing.assert_allclose(solution.y, np.exp(0.1 * solution.t).reshape(1, -1))
+        np.testing.assert_allclose(
+            solution["var squared"].data, np.exp(0.1 * solution.t) ** 2
+        )
+
 
 class TestCasadiSolverDAEsWithForwardSensitivityEquations(unittest.TestCase):
     def test_solve_sensitivity_scalar_var_scalar_input(self):
@@ -858,180 +890,33 @@ def test_solve_sensitivity_scalar_var_scalar_input(self):
             atol=1e-7
         )
 
-    def test_solve_sensitivity_vector_var_scalar_input(self):
-        var = pybamm.Variable("var", "negative electrode")
-        model = pybamm.BaseModel()
-        # Set length scales to avoid warning
-        model.length_scales = {"negative electrode": 1}
-        param = pybamm.InputParameter("param")
-        model.rhs = {var: -param * var}
-        model.initial_conditions = {var: 2}
-        model.variables = {"var": var}
-
-        # create discretisation
-        disc = get_discretisation_for_testing()
-        disc.process_model(model)
-        n = disc.mesh["negative electrode"].npts
-
-        # Solve - scalar input
-        solver = pybamm.CasadiSolver()
-        t_eval = np.linspace(0, 1)
-        solution = solver.solve(model, t_eval, inputs={"param": 7},
-                                calculate_sensitivities=["param"])
-        np.testing.assert_array_almost_equal(
-            solution["var"].data, np.tile(2 * np.exp(-7 * t_eval), (n, 1)), decimal=4,
-        )
-        np.testing.assert_array_almost_equal(
-            solution["var"].sensitivities["param"],
-            np.repeat(-2 * t_eval * np.exp(-7 * t_eval), n)[:, np.newaxis],
-            decimal=4,
-        )
-
-        # More complicated model
+    def test_solve_sensitivity_algebraic(self):
         # Create model
         model = pybamm.BaseModel()
-        # Set length scales to avoid warning
-        model.length_scales = {"negative electrode": 1}
-        var = pybamm.Variable("var", "negative electrode")
+        var = pybamm.Variable("var")
         p = pybamm.InputParameter("p")
-        q = pybamm.InputParameter("q")
-        r = pybamm.InputParameter("r")
-        s = pybamm.InputParameter("s")
-        model.rhs = {var: p * q}
-        model.initial_conditions = {var: r}
-        model.variables = {"var times s": var * s}
-
-        # Discretise
-        disc.process_model(model)
+        model.algebraic = {var: var - p * pybamm.t}
+        model.initial_conditions = {var: 0}
+        model.variables = {"var squared": var ** 2}
 
         # Solve
         # Make sure that passing in extra options works
-        solver = pybamm.CasadiSolver(
-            rtol=1e-10, atol=1e-10,
-        )
+        solver = pybamm.CasadiAlgebraicSolver(tol=1e-10)
         t_eval = np.linspace(0, 1, 80)
-        solution = solver.solve(
-            model, t_eval, inputs={"p": 0.1, "q": 2, "r": -1, "s": 0.5},
-            calculate_sensitivities=True,
-        )
-        np.testing.assert_allclose(solution.y, np.tile(-1 + 0.2 * solution.t, (n, 1)))
-        np.testing.assert_allclose(
-            solution.sensitivities["p"], np.repeat(2 * solution.t, n)[:, np.newaxis],
-        )
-        np.testing.assert_allclose(
-            solution.sensitivities["q"], np.repeat(0.1 * solution.t, n)[:, np.newaxis],
-        )
-        np.testing.assert_allclose(solution.sensitivities["r"], 1)
-        np.testing.assert_allclose(solution.sensitivities["s"], 0)
-        np.testing.assert_allclose(
-            solution.sensitivities["all"],
-            np.hstack(
-                [
-                    solution.sensitivities["p"],
-                    solution.sensitivities["q"],
-                    solution.sensitivities["r"],
-                    solution.sensitivities["s"],
-                ]
-            ),
-        )
+        solution = solver.solve(model, t_eval, inputs={"p": 0.1},
+                                calculate_sensitivities=True)
+        np.testing.assert_array_equal(solution.t, t_eval)
+        np.testing.assert_allclose(solution.y[0], 0.1 * solution.t)
         np.testing.assert_allclose(
-            solution["var times s"].data, np.tile(0.5 * (-1 + 0.2 * solution.t), (n, 1))
+            solution.sensitivities["p"], solution.t.reshape(-1, 1), atol=1e-7
         )
         np.testing.assert_allclose(
-            solution["var times s"].sensitivities["p"],
-            np.repeat(0.5 * (2 * solution.t), n)[:, np.newaxis],
+            solution["var squared"].data, (0.1 * solution.t) ** 2
         )
         np.testing.assert_allclose(
-            solution["var times s"].sensitivities["q"],
-            np.repeat(0.5 * (0.1 * solution.t), n)[:, np.newaxis],
-        )
-        np.testing.assert_allclose(solution["var times s"].sensitivities["r"], 0.5)
-        np.testing.assert_allclose(
-            solution["var times s"].sensitivities["s"],
-            np.repeat(-1 + 0.2 * solution.t, n)[:, np.newaxis],
-        )
-        np.testing.assert_allclose(
-            solution["var times s"].sensitivities["all"],
-            np.hstack(
-                [
-                    solution["var times s"].sensitivities["p"],
-                    solution["var times s"].sensitivities["q"],
-                    solution["var times s"].sensitivities["r"],
-                    solution["var times s"].sensitivities["s"],
-                ]
-            ),
-        )
-
-    def test_solve_sensitivity_scalar_var_vector_input(self):
-        var = pybamm.Variable("var", "negative electrode")
-        model = pybamm.BaseModel()
-        # Set length scales to avoid warning
-        model.length_scales = {"negative electrode": 1}
-
-        param = pybamm.InputParameter("param", "negative electrode")
-        model.rhs = {var: -param * var}
-        model.initial_conditions = {var: 2}
-        model.variables = {
-            "var": var,
-            "integral of var": pybamm.Integral(var, pybamm.standard_spatial_vars.x_n),
-        }
-
-        # create discretisation
-        mesh = get_mesh_for_testing(xpts=5)
-        spatial_methods = {"macroscale": pybamm.FiniteVolume()}
-        disc = pybamm.Discretisation(mesh, spatial_methods)
-        disc.process_model(model)
-        n = disc.mesh["negative electrode"].npts
-
-        # Solve - constant input
-        solver = pybamm.CasadiSolver(
-            mode="fast", rtol=1e-10, atol=1e-10
-        )
-        t_eval = np.linspace(0, 1)
-        solution = solver.solve(model, t_eval, inputs={"param": 7 * np.ones(n)},
-                                calculate_sensitivities=True)
-        l_n = mesh["negative electrode"].edges[-1]
-        np.testing.assert_array_almost_equal(
-            solution["var"].data, np.tile(2 * np.exp(-7 * t_eval), (n, 1)), decimal=4,
-        )
-
-        np.testing.assert_array_almost_equal(
-            solution["var"].sensitivities["param"],
-            np.vstack([np.eye(n) * -2 * t * np.exp(-7 * t) for t in t_eval]),
-        )
-        np.testing.assert_array_almost_equal(
-            solution["integral of var"].data, 2 * np.exp(-7 * t_eval) * l_n, decimal=4,
-        )
-        np.testing.assert_array_almost_equal(
-            solution["integral of var"].sensitivities["param"],
-            np.tile(-2 * t_eval * np.exp(-7 * t_eval) * l_n / n, (n, 1)).T,
-        )
-
-        # Solve - linspace input
-        p_eval = np.linspace(1, 2, n)
-        solution = solver.solve(model, t_eval, inputs={"param": p_eval},
-                                calculate_sensitivities=True)
-        l_n = mesh["negative electrode"].edges[-1]
-        np.testing.assert_array_almost_equal(
-            solution["var"].data, 2 * np.exp(-p_eval[:, np.newaxis] * t_eval), decimal=4
-        )
-        np.testing.assert_array_almost_equal(
-            solution["var"].sensitivities["param"],
-            np.vstack([np.diag(-2 * t * np.exp(-p_eval * t)) for t in t_eval]),
-        )
-
-        np.testing.assert_array_almost_equal(
-            solution["integral of var"].data,
-            np.sum(
-                2
-                * np.exp(-p_eval[:, np.newaxis] * t_eval)
-                * mesh["negative electrode"].d_edges[:, np.newaxis],
-                axis=0,
-            ),
-        )
-        np.testing.assert_array_almost_equal(
-            solution["integral of var"].sensitivities["param"],
-            np.vstack([-2 * t * np.exp(-p_eval * t) * l_n / n for t in t_eval]),
+            solution["var squared"].sensitivities["p"],
+            (2 * 0.1 * solution.t ** 2).reshape(-1, 1),
+            atol=1e-7
         )