#1477 put back algebraic solver sens tests

Author: martinjrobins

tests/unit/test_solvers/test_casadi_algebraic_solver.py

@@ -5,6 +5,8 @@
 import pybamm
 import unittest
 import numpy as np
+from scipy.optimize import least_squares
+import tests
 
 
 class TestCasadiAlgebraicSolver(unittest.TestCase):
@@ -172,6 +174,157 @@ def test_solve_with_input(self):
         np.testing.assert_array_equal(solution.y, -7)
 
 
+class TestCasadiAlgebraicSolverSensitivity(unittest.TestCase):
+    def test_solve_with_symbolic_input(self):
+        # Simple system: a single algebraic equation
+        var = pybamm.Variable("var")
+        model = pybamm.BaseModel()
+        model.algebraic = {var: var + pybamm.InputParameter("param")}
+        model.initial_conditions = {var: 2}
+        model.variables = {"var": var}
+
+        # create discretisation
+        disc = pybamm.Discretisation()
+        disc.process_model(model)
+
+        # Solve
+        solver = pybamm.CasadiAlgebraicSolver()
+        solution = solver.solve(model, [0])
+        np.testing.assert_array_equal(solution["var"].value({"param": 7}), -7)
+        np.testing.assert_array_equal(solution["var"].value({"param": 3}), -3)
+        np.testing.assert_array_equal(solution["var"].sensitivity({"param": 3}), -1)
+
+    def test_least_squares_fit(self):
+        # Simple system: a single algebraic equation
+        var = pybamm.Variable("var", domain="negative electrode")
+        model = pybamm.BaseModel()
+        p = pybamm.InputParameter("p")
+        q = pybamm.InputParameter("q")
+        model.algebraic = {var: (var - p)}
+        model.initial_conditions = {var: 3}
+        model.variables = {"objective": (var - q) ** 2 + (p - 3) ** 2}
+
+        # create discretisation
+        disc = tests.get_discretisation_for_testing()
+        disc.process_model(model)
+
+        # Solve
+        solver = pybamm.CasadiAlgebraicSolver()
+        solution = solver.solve(model, [0])
+        sol_var = solution["objective"]
+
+        def objective(x):
+            return sol_var.value({"p": x[0], "q": x[1]}).full().flatten()
+
+        # without jacobian
+        lsq_sol = least_squares(objective, [2, 2], method="lm")
+        np.testing.assert_array_almost_equal(lsq_sol.x, [3, 3], decimal=3)
+
+        def jac(x):
+            return sol_var.sensitivity({"p": x[0], "q": x[1]})
+
+        # with jacobian
+        lsq_sol = least_squares(objective, [2, 2], jac=jac, method="lm")
+        np.testing.assert_array_almost_equal(lsq_sol.x, [3, 3], decimal=3)
+
+    def test_solve_with_symbolic_input_1D_scalar_input(self):
+        var = pybamm.Variable("var", "negative electrode")
+        model = pybamm.BaseModel()
+        param = pybamm.InputParameter("param")
+        model.algebraic = {var: var + param}
+        model.initial_conditions = {var: 2}
+        model.variables = {"var": var}
+
+        # create discretisation
+        disc = tests.get_discretisation_for_testing()
+        disc.process_model(model)
+
+        # Solve - scalar input
+        solver = pybamm.CasadiAlgebraicSolver()
+        solution = solver.solve(model, [0])
+        np.testing.assert_array_equal(solution["var"].value({"param": 7}), -7)
+        np.testing.assert_array_equal(solution["var"].value({"param": 3}), -3)
+        np.testing.assert_array_equal(solution["var"].sensitivity({"param": 3}), -1)
+
+    def test_solve_with_symbolic_input_1D_vector_input(self):
+        var = pybamm.Variable("var", "negative electrode")
+        model = pybamm.BaseModel()
+        param = pybamm.InputParameter("param", "negative electrode")
+        model.algebraic = {var: var + param}
+        model.initial_conditions = {var: 2}
+        model.variables = {"var": var}
+
+        # create discretisation
+        disc = tests.get_discretisation_for_testing()
+        disc.process_model(model)
+
+        # Solve - scalar input
+        solver = pybamm.CasadiAlgebraicSolver()
+        solution = solver.solve(model, [0])
+        n = disc.mesh["negative electrode"].npts
+
+        solver = pybamm.CasadiAlgebraicSolver()
+        solution = solver.solve(model, [0])
+        p = np.linspace(0, 1, n)[:, np.newaxis]
+        np.testing.assert_array_almost_equal(
+            solution["var"].value({"param": 3 * np.ones(n)}), -3
+        )
+        np.testing.assert_array_almost_equal(
+            solution["var"].value({"param": 2 * p}), -2 * p
+        )
+        np.testing.assert_array_almost_equal(
+            solution["var"].sensitivity({"param": 3 * np.ones(n)}), -np.eye(40)
+        )
+        np.testing.assert_array_almost_equal(
+            solution["var"].sensitivity({"param": p}), -np.eye(40)
+        )
+
+    def test_solve_with_symbolic_input_in_initial_conditions(self):
+        # Simple system: a single algebraic equation
+        var = pybamm.Variable("var")
+        model = pybamm.BaseModel()
+        model.algebraic = {var: var + 2}
+        model.initial_conditions = {var: pybamm.InputParameter("param")}
+        model.variables = {"var": var}
+
+        # create discretisation
+        disc = pybamm.Discretisation()
+        disc.process_model(model)
+
+        # Solve
+        solver = pybamm.CasadiAlgebraicSolver()
+        solution = solver.solve(model, [0])
+        np.testing.assert_array_equal(solution["var"].value({"param": 7}), -2)
+        np.testing.assert_array_equal(solution["var"].value({"param": 3}), -2)
+        np.testing.assert_array_equal(solution["var"].sensitivity({"param": 3}), 0)
+
+    def test_least_squares_fit_input_in_initial_conditions(self):
+        # Simple system: a single algebraic equation
+        var = pybamm.Variable("var", domain="negative electrode")
+        model = pybamm.BaseModel()
+        p = pybamm.InputParameter("p")
+        q = pybamm.InputParameter("q")
+        model.algebraic = {var: (var - p)}
+        model.initial_conditions = {var: p}
+        model.variables = {"objective": (var - q) ** 2 + (p - 3) ** 2}
+
+        # create discretisation
+        disc = tests.get_discretisation_for_testing()
+        disc.process_model(model)
+
+        # Solve
+        solver = pybamm.CasadiAlgebraicSolver()
+        solution = solver.solve(model, [0])
+        sol_var = solution["objective"]
+
+        def objective(x):
+            return sol_var.value({"p": x[0], "q": x[1]}).full().flatten()
+
+        # without jacobian
+        lsq_sol = least_squares(objective, [2, 2], method="lm")
+        np.testing.assert_array_almost_equal(lsq_sol.x, [3, 3], decimal=3)
+
+
 if __name__ == "__main__":
     print("Add -v for more debug output")
     import sys