Merge pull request #1959 from pybamm-team/i1863-casadi-idklu

I1898 make evaluator functions consistent across casadi, python and jax

Author: martinjrobins

pybamm/expression_tree/operations/evaluate_python.py

@@ -465,7 +465,7 @@ def __init__(self, symbol):
         # add function def to first line
         python_str = (
             "def evaluate(constants, t=None, y=None, "
-            "y_dot=None, inputs=None, known_evals=None):\n" + python_str
+            "inputs=None):\n" + python_str
         )
 
         # calculate the final variable that will output the result of calling `evaluate`
@@ -491,21 +491,17 @@ def __init__(self, symbol):
         compiled_function = compile(python_str, result_var, "exec")
         exec(compiled_function)
 
-    def evaluate(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None):
+    def __call__(self, t=None, y=None, inputs=None):
         """
-        Acts as a drop-in replacement for :func:`pybamm.Symbol.evaluate`
+        evaluate function
         """
         # generated code assumes y is a column vector
         if y is not None and y.ndim == 1:
             y = y.reshape(-1, 1)
 
-        result = self._evaluate(self._constants, t, y, y_dot, inputs, known_evals)
+        result = self._evaluate(self._constants, t, y, inputs)
 
-        # don't need known_evals, but need to reproduce Symbol.evaluate signature
-        if known_evals is not None:
-            return result, known_evals
-        else:
-            return result
+        return result
 
     def __getstate__(self):
         # Control the state of instances of EvaluatorPython
@@ -581,7 +577,7 @@ def __init__(self, symbol):
         python_str = python_str.replace("\n", "\n   ")
 
         # add function def to first line
-        args = "t=None, y=None, y_dot=None, inputs=None, known_evals=None"
+        args = "t=None, y=None, inputs=None"
         if self._arg_list:
             args = ",".join(self._arg_list) + ", " + args
         python_str = "def evaluate_jax({}):\n".format(args) + python_str
@@ -628,23 +624,23 @@ def get_jacobian(self):
     def get_sensitivities(self):
         n = len(self._arg_list)
 
-        # forward mode autodiff wrt inputs, which is argument 3 after arg_list
-        jacobian_evaluate = jax.jacfwd(self._evaluate_jax, argnums=3 + n)
+        # forward mode autodiff wrt inputs, which is argument 2 after arg_list
+        jacobian_evaluate = jax.jacfwd(self._evaluate_jax, argnums=2 + n)
 
         self._sens_evaluate = jax.jit(
             jacobian_evaluate, static_argnums=self._static_argnums
         )
 
         return EvaluatorJaxSensitivities(self._sens_evaluate, self._constants)
 
-    def debug(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None):
+    def debug(self, t=None, y=None, inputs=None):
         # generated code assumes y is a column vector
         if y is not None and y.ndim == 1:
             y = y.reshape(-1, 1)
 
         # execute code
         jaxpr = jax.make_jaxpr(self._evaluate_jax)(
-            *self._constants, t, y, y_dot, inputs, known_evals
+            *self._constants, t, y, inputs
         ).jaxpr
         print("invars:", jaxpr.invars)
         print("outvars:", jaxpr.outvars)
@@ -654,52 +650,45 @@ def debug(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None):
         print()
         print("jaxpr:", jaxpr)
 
-    def evaluate(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None):
+    def __call__(self, t=None, y=None, inputs=None):
         """
-        Acts as a drop-in replacement for :func:`pybamm.Symbol.evaluate`
+        evaluate function
         """
         # generated code assumes y is a column vector
         if y is not None and y.ndim == 1:
             y = y.reshape(-1, 1)
 
-        result = self._jit_evaluate(*self._constants, t, y, y_dot, inputs, known_evals)
+        result = self._jit_evaluate(*self._constants, t, y, inputs)
 
-        # don't need known_evals, but need to reproduce Symbol.evaluate signature
-        if known_evals is not None:
-            return result, known_evals
-        else:
-            return result
+        return result
 
 
 class EvaluatorJaxJacobian:
     def __init__(self, jac_evaluate, constants):
         self._jac_evaluate = jac_evaluate
         self._constants = constants
 
-    def evaluate(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None):
+    def __call__(self, t=None, y=None, inputs=None):
         """
-        Acts as a drop-in replacement for :func:`pybamm.Symbol.evaluate`
+        evaluate function
         """
         # generated code assumes y is a column vector
         if y is not None and y.ndim == 1:
             y = y.reshape(-1, 1)
 
         # execute code
-        result = self._jac_evaluate(*self._constants, t, y, y_dot, inputs, known_evals)
+        result = self._jac_evaluate(*self._constants, t, y, inputs)
         result = result.reshape(result.shape[0], -1)
 
-        if known_evals is not None:
-            return result, known_evals
-        else:
-            return result
+        return result
 
 
 class EvaluatorJaxSensitivities:
     def __init__(self, jac_evaluate, constants):
         self._jac_evaluate = jac_evaluate
         self._constants = constants
 
-    def evaluate(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None):
+    def __call__(self, t=None, y=None, inputs=None):
         """
         Acts as a drop-in replacement for :func:`pybamm.Symbol.evaluate`
         """
@@ -708,9 +697,6 @@ def evaluate(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None):
             y = y.reshape(-1, 1)
 
         # execute code
-        result = self._jac_evaluate(*self._constants, t, y, y_dot, inputs, known_evals)
+        result = self._jac_evaluate(*self._constants, t, y, inputs)
 
-        if known_evals is not None:
-            return result, known_evals
-        else:
-            return result
+        return result

pybamm/solvers/algebraic_solver.py

@@ -67,7 +67,8 @@ def _integrate(self, model, t_eval, inputs_dict=None):
 
         y0 = model.y0
         if isinstance(y0, casadi.DM):
-            y0 = y0.full().flatten()
+            y0 = y0.full()
+        y0 = y0.flatten()
 
         # The casadi algebraic solver can read rhs equations, but leaves them unchanged
         # i.e. the part of the solution vector that corresponds to the differential
@@ -80,7 +81,16 @@ def _integrate(self, model, t_eval, inputs_dict=None):
             len_rhs = model.rhs_eval(t_eval[0], y0, inputs).shape[0]
         y0_diff, y0_alg = np.split(y0, [len_rhs])
 
-        algebraic = model.algebraic_eval
+        test_result = model.algebraic_eval(0, y0, inputs)
+
+        if isinstance(test_result, casadi.DM):
+            def algebraic(t, y):
+                result = model.algebraic_eval(t, y, inputs)
+                return result.full().flatten()
+        else:
+            def algebraic(t, y):
+                result = model.algebraic_eval(t, y, inputs)
+                return result.flatten()
 
         y_alg = np.empty((len(y0_alg), len(t_eval)))
 
@@ -91,7 +101,7 @@ def _integrate(self, model, t_eval, inputs_dict=None):
             def root_fun(y_alg):
                 "Evaluates algebraic using y"
                 y = np.concatenate([y0_diff, y_alg])
-                out = algebraic(t, y, inputs)
+                out = algebraic(t, y)
                 pybamm.logger.debug(
                     "Evaluating algebraic equations at t={}, L2-norm is {}".format(
                         t * model.timescale_eval, np.linalg.norm(out)