Repository: pybamm-team/PyBaMM

Issue:
refactor `process` function in BaseSolver to produce "raw" functions with a common signature
**Summary**
The evaluate signature for expression trees is:
`evaluate(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None) -> result, known_evals`

The call signature used for calls by the solvers are:
`__call__(t, y, inputs) -> result`

in `BaseSolver`, the `process` function is used to (optionally) convert an pybamm expression tree into one of pybamm.EvaluatorJax, pybamm.EvaluatorPython or Casadi functions, and then wrap them in a `SolverCallable` so that the call signature for the solvers above are satisfied. It also generates "raw" functions (the 1st return value), which can be:
1. the `evaluate` function of a pybamm expression tree
2. the `evaluate` function of a EvaluatorJax or EvaluatorPython class
3. a casadi function

The main issue here is that the signature of 3 is the same as the call signature above, while 1 and 2 use the `evaluate` signature. This is mainly why `SolverCallable` is needed to wrap them in a common interface.

I propose that we:
1. remove the 1st "raw" option above, so that these "raw" functions are either one of the evaluator functions, or a casadi function
2. ensure that the `evaluate` function of the Evaluator classes use the call signature used by the solvers
3. remove the `SolverCallable` class entirely, and have the solvers all use the raw functions directly (which are either the evaluate methods of one of the Evaluator classes, or a casadi function)

 **Motivation**

The motivation is to:
- speed up the calls by the solvers. The `SolverCallable` needs to be generic across all formats, but the raw casadi and evaluator functions have the potential to be faster. Note that this is currently done by the casadi solver anyway, which uses raw casadi functions rather than the `SolverCallable`
- simplify the BaseSolver code: it currently needs to keep track of raw functions (generally casadi function) and `SolverCallable` functions, this way it only keeps track of the raw functions
- allow all solvers access to the raw functions. At the moment only the casadi solver uses the raw casadi functions, but other solvers might need these as well (see #1863)



Discussion:

Patch:
diff --git a/pybamm/expression_tree/operations/evaluate_python.py b/pybamm/expression_tree/operations/evaluate_python.py
index 6949ccb3c..d0fb7ccd9 100644
--- a/pybamm/expression_tree/operations/evaluate_python.py
+++ b/pybamm/expression_tree/operations/evaluate_python.py
@@ -465,7 +465,7 @@ class EvaluatorPython:
         # 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 @@ class EvaluatorPython:
         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 @@ class EvaluatorJax:
         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,8 +624,8 @@ class EvaluatorJax:
     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
@@ -637,14 +633,14 @@ class EvaluatorJax:
 
         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,21 +650,17 @@ class EvaluatorJax:
         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:
@@ -676,22 +668,19 @@ class EvaluatorJaxJacobian:
         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:
@@ -699,7 +688,7 @@ class EvaluatorJaxSensitivities:
         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 @@ class EvaluatorJaxSensitivities:
             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
diff --git a/pybamm/solvers/algebraic_solver.py b/pybamm/solvers/algebraic_solver.py
index 10671e241..81f0348c6 100644
--- a/pybamm/solvers/algebraic_solver.py
+++ b/pybamm/solvers/algebraic_solver.py
@@ -67,7 +67,8 @@ class AlgebraicSolver(pybamm.BaseSolver):
 
         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 @@ class AlgebraicSolver(pybamm.BaseSolver):
             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 @@ class AlgebraicSolver(pybamm.BaseSolver):
             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)
diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py
index d83219856..6bc05e129 100644
--- a/pybamm/solvers/base_solver.py
+++ b/pybamm/solvers/base_solver.py
@@ -108,7 +108,11 @@ class BaseSolver(object):
         t_eval : numeric type, optional
             The times (in seconds) at which to compute the solution
         """
-        pybamm.logger.info("Start solver set-up")
+
+        if ics_only:
+            pybamm.logger.info("Start solver set-up, initial_conditions only")
+        else:
+            pybamm.logger.info("Start solver set-up")
 
         # Check model.algebraic for ode solvers
         if self.ode_solver is True and len(model.algebraic) > 0:
@@ -208,6 +212,9 @@ class BaseSolver(object):
                     num_parameters += len(inputs[name])
             model.len_rhs_sens = model.len_rhs * num_parameters
             model.len_alg_sens = model.len_alg * num_parameters
+        else:
+            model.len_rhs_sens = 0
+            model.len_alg_sens = 0
 
         if model.convert_to_format != "casadi":
             # Create Jacobian from concatenated rhs and algebraic
@@ -240,7 +247,59 @@ class BaseSolver(object):
             else:
                 y_and_S = y_casadi
 
-        def process(func, name, use_jacobian=None):
+        # if we will change the equations to include the explicit sensitivity
+        # equations, then we also need to update the mass matrix and bounds
+        if calculate_sensitivities_explicit:
+            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
+            if model.bounds[0].shape[0] == model.len_rhs_and_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
+            ):
+
+                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"
+                    )
+                )
+        else:
+            # take care if calculate_sensitivites used then not used
+            if model.bounds[0].shape[0] > model.len_rhs_and_alg:
+                model.bounds = (
+                    model.bounds[0][: model.len_rhs_and_alg],
+                    model.bounds[1][: model.len_rhs_and_alg],
+                )
+            if (
+                model.mass_matrix is not None
+                and model.mass_matrix.shape[0] > model.len_rhs_and_alg
+            ):
+                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
+                    ]
+                )
+
+        def process(symbol, name, use_jacobian=None):
             def report(string):
                 # don't log event conversion
                 if "event" not in string:
@@ -251,7 +310,7 @@ class BaseSolver(object):
 
             if model.convert_to_format == "jax":
                 report(f"Converting {name} to jax")
-                func = pybamm.EvaluatorJax(func)
+                func = pybamm.EvaluatorJax(symbol)
                 jacp = None
                 if model.calculate_sensitivities:
                     report(
@@ -261,18 +320,14 @@ class BaseSolver(object):
                         )
                     )
                     jacp = func.get_sensitivities()
-                    jacp = jacp.evaluate
                 if use_jacobian:
                     report(f"Calculating jacobian for {name} using jax")
                     jac = func.get_jacobian()
-                    jac = jac.evaluate
                 else:
                     jac = None
 
-                func = func.evaluate
-
             elif model.convert_to_format != "casadi":
-                # Process with pybamm functions, optionally converting
+                # Process with pybamm functions, converting
                 # to python evaluator
                 if model.calculate_sensitivities:
                     report(
@@ -282,20 +337,20 @@ class BaseSolver(object):
                         )
                     )
                     jacp_dict = {
-                        p: func.diff(pybamm.InputParameter(p))
+                        p: symbol.diff(pybamm.InputParameter(p))
                         for p in model.calculate_sensitivities
                     }
-                    if model.convert_to_format == "python":
-                        report(f"Converting sensitivities for {name} to python")
-                        jacp_dict = {
-                            p: pybamm.EvaluatorPython(jacp)
-                            for p, jacp in jacp_dict.items()
-                        }
+
+                    report(f"Converting sensitivities for {name} to python")
+                    jacp_dict = {
+                        p: pybamm.EvaluatorPython(jacp)
+                        for p, jacp in jacp_dict.items()
+                    }
 
                     # jacp should be a function that returns a dict of sensitivities
                     def jacp(*args, **kwargs):
                         return {
-                            k: v.evaluate(*args, **kwargs) for k, v in jacp_dict.items()
+                            k: v(*args, **kwargs) for k, v in jacp_dict.items()
                         }
 
                 else:
@@ -303,24 +358,21 @@ class BaseSolver(object):
 
                 if use_jacobian:
                     report(f"Calculating jacobian for {name}")
-                    jac = jacobian.jac(func, y)
+                    jac = jacobian.jac(symbol, y)
                     if model.convert_to_format == "python":
                         report(f"Converting jacobian for {name} to python")
                         jac = pybamm.EvaluatorPython(jac)
-                    jac = jac.evaluate
                 else:
                     jac = None
 
-                if model.convert_to_format == "python":
-                    report(f"Converting {name} to python")
-                    func = pybamm.EvaluatorPython(func)
-
-                func = func.evaluate
+                report(f"Converting {name} to python")
+                func = pybamm.EvaluatorPython(symbol)
 
             else:
                 # Process with CasADi
                 report(f"Converting {name} to CasADi")
-                func = func.to_casadi(t_casadi, y_casadi, inputs=p_casadi)
+                casadi_expression = symbol.to_casadi(
+                    t_casadi, y_casadi, inputs=p_casadi)
                 # Add sensitivity vectors to the rhs and algebraic equations
                 jacp = None
                 if calculate_sensitivities_explicit:
@@ -334,59 +386,59 @@ class BaseSolver(object):
                             "Creating explicit forward sensitivity equations "
                             "for rhs using CasADi"
                         )
-                        df_dx = casadi.jacobian(func, y_diff)
-                        df_dp = casadi.jacobian(func, pS_casadi_stacked)
+                        df_dx = casadi.jacobian(casadi_expression, y_diff)
+                        df_dp = casadi.jacobian(casadi_expression, pS_casadi_stacked)
                         S_x_mat = S_x.reshape(
                             (model.len_rhs, pS_casadi_stacked.shape[0])
                         )
                         if model.len_alg == 0:
                             S_rhs = (df_dx @ S_x_mat + df_dp).reshape((-1, 1))
                         else:
-                            df_dz = casadi.jacobian(func, y_alg)
+                            df_dz = casadi.jacobian(casadi_expression, y_alg)
                             S_z_mat = S_z.reshape(
                                 (model.len_alg, pS_casadi_stacked.shape[0])
                             )
                             S_rhs = (df_dx @ S_x_mat + df_dz @ S_z_mat + df_dp).reshape(
                                 (-1, 1)
                             )
-                        func = casadi.vertcat(func, S_rhs)
+                        casadi_expression = casadi.vertcat(casadi_expression, S_rhs)
                     if name == "algebraic" and model.len_alg > 0:
                         report(
                             "Creating explicit forward sensitivity equations "
                             "for algebraic using CasADi"
                         )
-                        dg_dz = casadi.jacobian(func, y_alg)
-                        dg_dp = casadi.jacobian(func, pS_casadi_stacked)
+                        dg_dz = casadi.jacobian(casadi_expression, y_alg)
+                        dg_dp = casadi.jacobian(casadi_expression, pS_casadi_stacked)
                         S_z_mat = S_z.reshape(
                             (model.len_alg, pS_casadi_stacked.shape[0])
                         )
                         if model.len_rhs == 0:
                             S_alg = (dg_dz @ S_z_mat + dg_dp).reshape((-1, 1))
                         else:
-                            dg_dx = casadi.jacobian(func, y_diff)
+                            dg_dx = casadi.jacobian(casadi_expression, y_diff)
                             S_x_mat = S_x.reshape(
                                 (model.len_rhs, pS_casadi_stacked.shape[0])
                             )
                             S_alg = (dg_dx @ S_x_mat + dg_dz @ S_z_mat + dg_dp).reshape(
                                 (-1, 1)
                             )
-                        func = casadi.vertcat(func, S_alg)
+                        casadi_expression = casadi.vertcat(casadi_expression, S_alg)
                     if name == "initial_conditions":
                         if model.len_rhs == 0 or model.len_alg == 0:
-                            S_0 = casadi.jacobian(func, pS_casadi_stacked).reshape(
-                                (-1, 1)
-                            )
-                            func = casadi.vertcat(func, S_0)
+                            S_0 = casadi.jacobian(
+                                casadi_expression, pS_casadi_stacked
+                            ).reshape((-1, 1))
+                            casadi_expression = casadi.vertcat(casadi_expression, S_0)
                         else:
-                            x0 = func[: model.len_rhs]
-                            z0 = func[model.len_rhs :]
+                            x0 = casadi_expression[: model.len_rhs]
+                            z0 = casadi_expression[model.len_rhs :]
                             Sx_0 = casadi.jacobian(x0, pS_casadi_stacked).reshape(
                                 (-1, 1)
                             )
                             Sz_0 = casadi.jacobian(z0, pS_casadi_stacked).reshape(
                                 (-1, 1)
                             )
-                            func = casadi.vertcat(x0, Sx_0, z0, Sz_0)
+                            casadi_expression = casadi.vertcat(x0, Sx_0, z0, Sz_0)
                 elif model.calculate_sensitivities:
                     report(
                         (
@@ -397,18 +449,19 @@ class BaseSolver(object):
                     )
                     jacp_dict = {}
                     for pname in model.calculate_sensitivities:
-                        p_diff = casadi.jacobian(func, p_casadi[pname])
+                        p_diff = casadi.jacobian(casadi_expression, p_casadi[pname])
                         jacp_dict[pname] = casadi.Function(
-                            name, [t_casadi, y_casadi, p_casadi_stacked], [p_diff]
+                            name, [t_casadi, y_and_S, p_casadi_stacked], [p_diff]
                         )
-                    # jacp should be a function that returns a dict of sensitivities
 
+                    # jacp should be a casadi_expressiontion that returns
+                    # a dict of sensitivities
                     def jacp(*args, **kwargs):
                         return {k: v(*args, **kwargs) for k, v in jacp_dict.items()}
 
                 if use_jacobian:
                     report(f"Calculating jacobian for {name} using CasADi")
-                    jac_casadi = casadi.jacobian(func, y_and_S)
+                    jac_casadi = casadi.jacobian(casadi_expression, y_and_S)
                     jac = casadi.Function(
                         name, [t_casadi, y_and_S, p_casadi_stacked], [jac_casadi]
                     )
@@ -416,26 +469,10 @@ class BaseSolver(object):
                     jac = None
 
                 func = casadi.Function(
-                    name, [t_casadi, y_and_S, p_casadi_stacked], [func]
+                    name, [t_casadi, y_and_S, p_casadi_stacked], [casadi_expression]
                 )
-            if name == "residuals":
-                func_call = Residuals(func, name, model)
-            else:
-                func_call = SolverCallable(func, name, model)
-            if jac is not None:
-                jac_call = SolverCallable(jac, name + "_jac", model)
-            else:
-                jac_call = None
-            if jacp is not None:
-                jacp_call = SensitivityCallable(
-                    jacp,
-                    name + "_sensitivity_wrt_inputs",
-                    model,
-                    model.convert_to_format,
-                )
-            else:
-                jacp_call = None
-            return func, func_call, jac_call, jacp_call
+
+            return func, jac, jacp
 
         # Process initial conditions
         initial_conditions = process(
@@ -443,241 +480,194 @@ class BaseSolver(object):
             "initial_conditions",
             use_jacobian=False,
         )[0]
-        init_eval = InitialConditions(initial_conditions, model)
+        model.initial_conditions_eval = initial_conditions
 
-        if calculate_sensitivities_explicit:
-            y0_total_size = (
-                model.len_rhs + model.len_rhs_sens + model.len_alg + model.len_alg_sens
-            )
-            init_eval.y_dummy = np.zeros((y0_total_size, 1))
+        # evaluate initial condition
+        y0_total_size = (
+            model.len_rhs + model.len_rhs_sens + model.len_alg + model.len_alg_sens
+        )
+        y_zero = np.zeros((y0_total_size, 1))
+        if model.convert_to_format == "casadi":
+            # stack inputs
+            inputs_casadi = casadi.vertcat(*[x for x in inputs.values()])
+            model.y0 = initial_conditions(0, y_zero, inputs_casadi)
         else:
-            init_eval.y_dummy = np.zeros((model.len_rhs_and_alg, 1))
-
-        # Calculate initial conditions
-        model.y0 = init_eval(inputs)
-        model.init_eval = init_eval
-
-        if not ics_only:
-            # Check for heaviside and modulo functions in rhs and algebraic and add
-            # discontinuity events if these exist.
-            # Note: only checks for the case of t < X, t <= X, X < t, or X <= t,
-            # but also accounts for the fact that t might be dimensional
-            # Only do this for DAE models as ODE models can deal with discontinuities
-            # fine
-            if len(model.algebraic) > 0:
-                for symbol in itertools.chain(
-                    model.concatenated_rhs.pre_order(),
-                    model.concatenated_algebraic.pre_order(),
-                ):
-                    if isinstance(symbol, _Heaviside):
-                        found_t = False
-                        # Dimensionless
-                        if symbol.right.id == pybamm.t.id:
-                            expr = symbol.left
-                            found_t = True
-                        elif symbol.left.id == pybamm.t.id:
-                            expr = symbol.right
-                            found_t = True
-                        # Dimensional
-                        elif symbol.right.id == (pybamm.t * model.timescale_eval).id:
-                            expr = symbol.left / symbol.right.right
-                            found_t = True
-                        elif symbol.left.id == (pybamm.t * model.timescale_eval).id:
-                            expr = symbol.right / symbol.left.right
-                            found_t = True
-
-                        # Update the events if the heaviside function depended on t
-                        if found_t:
+            model.y0 = initial_conditions(0, y_zero, inputs)
+
+        if ics_only:
+            pybamm.logger.info("Finish solver set-up")
+            return
+
+        # Check for heaviside and modulo functions in rhs and algebraic and add
+        # discontinuity events if these exist.
+        # Note: only checks for the case of t < X, t <= X, X < t, or X <= t,
+        # but also accounts for the fact that t might be dimensional
+        # Only do this for DAE models as ODE models can deal with discontinuities
+        # fine
+        if len(model.algebraic) > 0:
+            for symbol in itertools.chain(
+                model.concatenated_rhs.pre_order(),
+                model.concatenated_algebraic.pre_order(),
+            ):
+                if isinstance(symbol, _Heaviside):
+                    found_t = False
+                    # Dimensionless
+                    if symbol.right.id == pybamm.t.id:
+                        expr = symbol.left
+                        found_t = True
+                    elif symbol.left.id == pybamm.t.id:
+                        expr = symbol.right
+                        found_t = True
+                    # Dimensional
+                    elif symbol.right.id == (pybamm.t * model.timescale_eval).id:
+                        expr = symbol.left / symbol.right.right
+                        found_t = True
+                    elif symbol.left.id == (pybamm.t * model.timescale_eval).id:
+                        expr = symbol.right / symbol.left.right
+                        found_t = True
+
+                    # Update the events if the heaviside function depended on t
+                    if found_t:
+                        model.events.append(
+                            pybamm.Event(
+                                str(symbol), expr, pybamm.EventType.DISCONTINUITY,
+                            )
+                        )
+                elif isinstance(symbol, pybamm.Modulo):
+                    found_t = False
+                    # Dimensionless
+                    if symbol.left.id == pybamm.t.id:
+                        expr = symbol.right
+                        found_t = True
+                    # Dimensional
+                    elif symbol.left.id == (pybamm.t * model.timescale_eval).id:
+                        expr = symbol.right / symbol.left.right
+                        found_t = True
+
+                    # Update the events if the modulo function depended on t
+                    if found_t:
+                        if t_eval is None:
+                            N_events = 200
+                        else:
+                            N_events = t_eval[-1] // expr.value
+
+                        for i in np.arange(N_events):
                             model.events.append(
                                 pybamm.Event(
-                                    str(symbol), expr, pybamm.EventType.DISCONTINUITY
+                                    str(symbol), expr * pybamm.Scalar(i + 1),
+                                    pybamm.EventType.DISCONTINUITY
                                 )
                             )
-                    elif isinstance(symbol, pybamm.Modulo):
-                        found_t = False
-                        # Dimensionless
-                        if symbol.left.id == pybamm.t.id:
-                            expr = symbol.right
-                            found_t = True
-                        # Dimensional
-                        elif symbol.left.id == (pybamm.t * model.timescale_eval).id:
-                            expr = symbol.right / symbol.left.right
-                            found_t = True
-
-                        # Update the events if the modulo function depended on t
-                        if found_t:
-                            if t_eval is None:
-                                N_events = 200
-                            else:
-                                N_events = t_eval[-1] // expr.value
-
-                            for i in np.arange(N_events):
-                                model.events.append(
-                                    pybamm.Event(
-                                        str(symbol),
-                                        expr * pybamm.Scalar(i + 1),
-                                        pybamm.EventType.DISCONTINUITY,
-                                    )
-                                )
 
-            # Process rhs, algebraic and event expressions
-            rhs, rhs_eval, jac_rhs, jacp_rhs = process(model.concatenated_rhs, "RHS")
-            algebraic, algebraic_eval, jac_algebraic, jacp_algebraic = process(
-                model.concatenated_algebraic, "algebraic"
-            )
-            casadi_terminate_events = []
-            terminate_events_eval = []
-            interpolant_extrapolation_events_eval = []
-            discontinuity_events_eval = []
-            for n, event in enumerate(model.events):
-                if event.event_type == pybamm.EventType.DISCONTINUITY:
-                    # discontinuity events are evaluated before the solver is called,
-                    # so don't need to process them
-                    discontinuity_events_eval.append(event)
-                elif event.event_type == pybamm.EventType.SWITCH:
-                    if (
-                        isinstance(self, pybamm.CasadiSolver)
-                        and self.mode == "fast with events"
-                        and model.algebraic != {}
-                    ):
-                        # Save some events to casadi_terminate_events for the 'fast with
-                        # events' mode of the casadi solver
-                        # We only need to do this if the model is a DAE model
-                        # see #1082
-                        k = 20
-                        init_sign = float(
-                            np.sign(event.evaluate(0, model.y0.full(), inputs=inputs))
-                        )
-                        # We create a sigmoid for each event which will multiply the
-                        # rhs. Doing * 2 - 1 ensures that when the event is crossed,
-                        # the sigmoid is zero. Hence the rhs is zero and the solution
-                        # stays constant for the rest of the simulation period
-                        # We can then cut off the part after the event was crossed
-                        event_sigmoid = (
-                            pybamm.sigmoid(0, init_sign * event.expression, k) * 2 - 1
-                        )
-                        event_casadi = process(
-                            event_sigmoid, f"event_{n}", use_jacobian=False
-                        )[0]
-                        # use the actual casadi object as this will go into the rhs
-                        casadi_terminate_events.append(event_casadi)
-                else:
-                    # use the function call
-                    event_eval = process(
-                        event.expression, f"event_{n}", use_jacobian=False
-                    )[1]
-                    if event.event_type == pybamm.EventType.TERMINATION:
-                        terminate_events_eval.append(event_eval)
-                    elif event.event_type == pybamm.EventType.INTERPOLANT_EXTRAPOLATION:
-                        interpolant_extrapolation_events_eval.append(event_eval)
-
-            # Add the solver attributes
-            model.rhs_eval = rhs_eval
-            model.algebraic_eval = algebraic_eval
-            model.jac_algebraic_eval = jac_algebraic
-            model.casadi_terminate_events = casadi_terminate_events
-            model.terminate_events_eval = terminate_events_eval
-            model.discontinuity_events_eval = discontinuity_events_eval
-            model.interpolant_extrapolation_events_eval = (
-                interpolant_extrapolation_events_eval
-            )
+        # Process rhs, algebraic, residual and event expressions
+        # and wrap in callables
+        rhs, jac_rhs, jacp_rhs = process(model.concatenated_rhs, "RHS")
 
-            # if we have changed the equations to include the explicit sensitivity
-            # equations, then we also need to update the mass matrix and bounds
-            if calculate_sensitivities_explicit:
-                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
-                if model.bounds[0].shape[0] == model.len_rhs_and_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
-                ):
+        algebraic, jac_algebraic, jacp_algebraic = process(
+            model.concatenated_algebraic, "algebraic"
+        )
 
-                    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"
-                        )
-                    )
-            else:
-                # take care if calculate_sensitivites used then not used
-                if model.bounds[0].shape[0] > model.len_rhs_and_alg:
-                    model.bounds = (
-                        model.bounds[0][: model.len_rhs_and_alg],
-                        model.bounds[1][: model.len_rhs_and_alg],
-                    )
+        # combine rhs and algebraic functions
+        if len(model.rhs) == 0:
+            rhs_algebraic = model.concatenated_algebraic
+        elif len(model.algebraic) == 0:
+            rhs_algebraic = model.concatenated_rhs
+        else:
+            rhs_algebraic = pybamm.NumpyConcatenation(
+                model.concatenated_rhs, model.concatenated_algebraic
+            )
+        rhs_algebraic, jac_rhs_algebraic, jacp_rhs_algebraic = process(
+            rhs_algebraic, "rhs_algebraic"
+        )
+
+        casadi_terminate_events = []
+        terminate_events = []
+        interpolant_extrapolation_events = []
+        discontinuity_events = []
+        for n, event in enumerate(model.events):
+            if event.event_type == pybamm.EventType.DISCONTINUITY:
+                # discontinuity events are evaluated before the solver is called,
+                # so don't need to process them
+                discontinuity_events.append(event)
+            elif event.event_type == pybamm.EventType.SWITCH:
                 if (
-                    model.mass_matrix is not None
-                    and model.mass_matrix.shape[0] > model.len_rhs_and_alg
+                    isinstance(self, pybamm.CasadiSolver)
+                    and self.mode == "fast with events"
+                    and model.algebraic != {}
                 ):
-                    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
-                        ]
-                    )
-
-            # Save CasADi functions for the CasADi solver
-            # 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.root_method, pybamm.CasadiAlgebraicSolver) or isinstance(
-                self, (pybamm.CasadiSolver, pybamm.CasadiAlgebraicSolver)
-            ):
-                # can use DAE solver to solve model with algebraic equations only
-                if len(model.rhs) > 0:
-                    mass_matrix_inv = casadi.MX(model.mass_matrix_inv.entries)
-                    explicit_rhs = mass_matrix_inv @ rhs(
-                        t_casadi, y_and_S, p_casadi_stacked
+                    # Save some events to casadi_terminate_events for the 'fast with
+                    # events' mode of the casadi solver
+                    # We only need to do this if the model is a DAE model
+                    # see #1082
+                    k = 20
+                    init_sign = float(
+                        np.sign(event.evaluate(0, model.y0.full(), inputs=inputs))
                     )
-                    model.casadi_rhs = casadi.Function(
-                        "rhs", [t_casadi, y_and_S, p_casadi_stacked], [explicit_rhs]
+                    # We create a sigmoid for each event which will multiply the
+                    # rhs. Doing * 2 - 1 ensures that when the event is crossed,
+                    # the sigmoid is zero. Hence the rhs is zero and the solution
+                    # stays constant for the rest of the simulation period
+                    # We can then cut off the part after the event was crossed
+                    event_sigmoid = (
+                        pybamm.sigmoid(0, init_sign * event.expression, k) * 2 - 1
                     )
-                model.casadi_algebraic = algebraic
-                model.casadi_sensitivities_rhs = jacp_rhs
-                model.casadi_sensitivities_algebraic = jacp_algebraic
-            if len(model.rhs) == 0:
-                # No rhs equations: residuals is algebraic only
-                model.residuals_eval = Residuals(algebraic, "residuals", model)
-                model.jacobian_eval = jac_algebraic
-                model.sensitivities_eval = jacp_algebraic
-            elif len(model.algebraic) == 0:
-                # No algebraic equations: residuals is rhs only
-                model.residuals_eval = Residuals(rhs, "residuals", model)
-                model.jacobian_eval = jac_rhs
-                model.sensitivities_eval = jacp_rhs
-            # Calculate consistent initial conditions for the algebraic equations
+                    event_casadi = process(
+                        event_sigmoid, f"event_{n}", use_jacobian=False
+                    )[0]
+                    # use the actual casadi object as this will go into the rhs
+                    casadi_terminate_events.append(event_casadi)
             else:
-                all_states = pybamm.NumpyConcatenation(
-                    model.concatenated_rhs, model.concatenated_algebraic
+                # use the function call
+                event_call = process(
+                    event.expression, f"event_{n}", use_jacobian=False
+                )[0]
+                if event.event_type == pybamm.EventType.TERMINATION:
+                    terminate_events.append(event_call)
+                elif event.event_type == pybamm.EventType.INTERPOLANT_EXTRAPOLATION:
+                    interpolant_extrapolation_events.append(event_call)
+
+        # Add the solver attributes
+        model.rhs_eval = rhs
+        model.algebraic_eval = algebraic
+        model.rhs_algebraic_eval = rhs_algebraic
+
+        model.terminate_events_eval = terminate_events
+        model.discontinuity_events_eval = discontinuity_events
+        model.interpolant_extrapolation_events_eval = interpolant_extrapolation_events
+
+        model.jac_rhs_eval = jac_rhs
+        model.jacp_rhs_eval = jacp_rhs
+
+        model.jac_algebraic_eval = jac_algebraic
+        model.jacp_algebraic_eval = jacp_algebraic
+
+        model.jac_rhs_algebraic_eval = jac_rhs_algebraic
+        model.jacp_rhs_algebraic_eval = jacp_rhs_algebraic
+
+        # Save CasADi functions for the CasADi solver
+        # Save CasADi functions for solvers that use CasADi
+        # Note: when we pass to casadi the ode part of the problem must be in
+        if isinstance(self.root_method, pybamm.CasadiAlgebraicSolver) or isinstance(
+            self, (pybamm.CasadiSolver, pybamm.CasadiAlgebraicSolver)
+        ):
+            # can use DAE solver to solve model with algebraic equations only
+            if len(model.rhs) > 0:
+                mass_matrix_inv = casadi.MX(model.mass_matrix_inv.entries)
+                explicit_rhs = mass_matrix_inv @ rhs(
+                    t_casadi, y_and_S, p_casadi_stacked
+                )
+                model.casadi_rhs = casadi.Function(
+                    "rhs", [t_casadi, y_and_S, p_casadi_stacked], [explicit_rhs]
                 )
-                # Process again, uses caching so should be quick
-                residuals_eval, jacobian_eval, jacobian_wrtp_eval = process(
-                    all_states, "residuals"
-                )[1:]
-                model.residuals_eval = residuals_eval
-                model.jacobian_eval = jacobian_eval
-                model.sensitivities_eval = jacobian_wrtp_eval
+            model.casadi_terminate_events = casadi_terminate_events
+            model.casadi_algebraic = algebraic
+            model.casadi_sensitivities = jacp_rhs_algebraic
+            model.casadi_sensitivities_rhs = jacp_rhs
+            model.casadi_sensitivities_algebraic = jacp_algebraic
 
         pybamm.logger.info("Finish solver set-up")
 
-    def _set_initial_conditions(self, model, inputs, update_rhs):
+    def _set_initial_conditions(self, model, inputs_dict, update_rhs):
         """
         Set initial conditions for the model. This is skipped if the solver is an
         algebraic solver (since this would make the algebraic solver redundant), and if
@@ -688,38 +678,53 @@ class BaseSolver(object):
         ----------
         model : :class:`pybamm.BaseModel`
             The model for which to calculate initial conditions.
-        inputs : dict
+        inputs_dict : dict
             Any input parameters to pass to the model when solving
         update_rhs : bool
             Whether to update the rhs. True for 'solve', False for 'step'.
 
         """
 
+        y0_total_size = (
+            model.len_rhs + model.len_rhs_sens + model.len_alg + model.len_alg_sens
+        )
+        y_zero = np.zeros((y0_total_size, 1))
+
+        if model.convert_to_format == "casadi":
+            # stack inputs
+            inputs = casadi.vertcat(*[x for x in inputs_dict.values()])
+        else:
+            inputs = inputs_dict
+
         if self.algebraic_solver is True:
             # Don't update model.y0
             return None
         elif len(model.algebraic) == 0:
             if update_rhs is True:
                 # Recalculate initial conditions for the rhs equations
-                y0 = model.init_eval(inputs)
+                y0 = model.initial_conditions_eval(0, y_zero, inputs)
             else:
                 # Don't update model.y0
                 return None
         else:
             if update_rhs is True:
                 # Recalculate initial conditions for the rhs equations
-                y0_from_inputs = model.init_eval(inputs)
+                y0_from_inputs = model.initial_conditions_eval(0, y_zero, inputs)
                 # Reuse old solution for algebraic equations
                 y0_from_model = model.y0
                 len_rhs = model.len_rhs
                 # update model.y0, which is used for initialising the algebraic solver
                 if len_rhs == 0:
                     model.y0 = y0_from_model
-                else:
+                elif isinstance(y0_from_inputs, casadi.DM):
                     model.y0 = casadi.vertcat(
                         y0_from_inputs[:len_rhs], y0_from_model[len_rhs:]
                     )
-            y0 = self.calculate_consistent_state(model, 0, inputs)
+                else:
+                    model.y0 = np.vstack(
+                        (y0_from_inputs[:len_rhs], y0_from_model[len_rhs:])
+                    )
+            y0 = self.calculate_consistent_state(model, 0, inputs_dict)
         # Make y0 a function of inputs if doing symbolic with casadi
         model.y0 = y0
 
@@ -734,7 +739,7 @@ class BaseSolver(object):
             The model for which to calculate initial conditions.
         time : float
             The time at which to calculate the states
-        inputs_dict : dict, optional
+        inputs: dict, optional
             Any input parameters to pass to the model when solving
 
         Returns
@@ -756,8 +761,6 @@ class BaseSolver(object):
         pybamm.logger.debug("Found consistent states")
 
         y0 = root_sol.all_ys[0]
-        if isinstance(y0, np.ndarray):
-            y0 = y0.flatten()
         return y0
 
     def solve(
@@ -1234,7 +1237,6 @@ class BaseSolver(object):
                 model.y0 = (
                     model.set_initial_conditions_from(old_solution)
                     .concatenated_initial_conditions.evaluate(0, inputs=ext_and_inputs)
-                    .flatten()
                 )
         set_up_time = timer.time()
 
@@ -1438,96 +1440,3 @@ class BaseSolver(object):
 
         ext_and_inputs = {**external_variables, **ordered_inputs}
         return ext_and_inputs
-
-
-class SolverCallable:
-    """A class that will be called by the solver when integrating"""
-
-    def __init__(self, function, name, model):
-        self._function = function
-        if isinstance(function, casadi.Function):
-            self.form = "casadi"
-        else:
-            self.form = "python"
-        self.name = name
-        self.model = model
-        self.timescale = self.model.timescale_eval
-
-    def __call__(self, t, y, inputs):
-        pybamm.logger.debug(
-            "Evaluating {} for {} at t={}".format(
-                self.name, self.model.name, t * self.timescale
-            )
-        )
-        if self.name in ["RHS", "algebraic", "residuals", "event"]:
-
-            return self.function(t, y, inputs).flatten()
-        else:
-            return self.function(t, y, inputs)
-
-    def function(self, t, y, inputs):
-        if self.form == "casadi":
-            states_eval = self._function(t, y, inputs)
-            if self.name in ["RHS", "algebraic", "residuals", "event"]:
-                return states_eval.full()
-            else:
-                # keep jacobians sparse
-                return states_eval
-        else:
-            return self._function(t, y, inputs=inputs, known_evals={})[0]
-
-
-class SensitivityCallable:
-    """A class that will be called by the solver when integrating"""
-
-    def __init__(self, function, name, model, form):
-        self._function = function
-        self.form = form
-        self.name = name
-        self.model = model
-        self.timescale = self.model.timescale_eval
-
-    def __call__(self, t, y, inputs):
-        pybamm.logger.debug(
-            "Evaluating sensitivities of {} for {} at t={}".format(
-                self.name, self.model.name, t * self.timescale
-            )
-        )
-        return self.function(t, y, inputs)
-
-    def function(self, t, y, inputs):
-        if self.form == "casadi":
-            return self._function(t, y, inputs)
-        elif self.form == "jax":
-            return self._function(t, y, inputs=inputs)
-        else:
-            ret_with_known_evals = self._function(t, y, inputs=inputs, known_evals={})
-            return {k: v[0] for k, v in ret_with_known_evals.items()}
-
-
-class Residuals(SolverCallable):
-    """Returns information about residuals at time t and state y"""
-
-    def __init__(self, function, name, model):
-        super().__init__(function, name, model)
-        if model.mass_matrix is not None:
-            self.mass_matrix = model.mass_matrix.entries
-
-    def __call__(self, t, y, ydot, inputs):
-        states_eval = super().__call__(t, y, inputs)
-        return states_eval - self.mass_matrix @ ydot
-
-
-class InitialConditions(SolverCallable):
-    """Returns initial conditions given inputs"""
-
-    def __init__(self, function, model):
-        super().__init__(function, "initial conditions", model)
-
-    def __call__(self, inputs):
-        if self.form == "casadi":
-            if isinstance(inputs, dict):
-                inputs = casadi.vertcat(*[x for x in inputs.values()])
-            return self._function(0, self.y_dummy, inputs)
-        else:
-            return self._function(0, self.y_dummy, inputs=inputs).flatten()
diff --git a/pybamm/solvers/casadi_solver.py b/pybamm/solvers/casadi_solver.py
index 0c8b3c747..edd9da1a0 100644
--- a/pybamm/solvers/casadi_solver.py
+++ b/pybamm/solvers/casadi_solver.py
@@ -482,7 +482,7 @@ class CasadiSolver(pybamm.BaseSolver):
         if model.interpolant_extrapolation_events_eval:
             inputs = casadi.vertcat(*[x for x in solution.all_inputs[-1].values()])
             extrap_event = [
-                event(solution.t[-1], solution.y[:, -1], inputs=inputs)
+                event(solution.t[-1], solution.y[:, -1], inputs)
                 for event in model.interpolant_extrapolation_events_eval
             ]
 
diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py
index e7c204175..3a0753deb 100644
--- a/pybamm/solvers/idaklu_solver.py
+++ b/pybamm/solvers/idaklu_solver.py
@@ -161,7 +161,7 @@ class IDAKLUSolver(pybamm.BaseSolver):
             Any external variables or input parameters to pass to the model when solving
         """
         inputs_dict = inputs_dict or {}
-        if model.rhs_eval.form == "casadi":
+        if model.convert_to_format == "casadi":
             # stack inputs
             inputs = casadi.vertcat(*[x for x in inputs_dict.values()])
             # raise warning about casadi format being slow
@@ -174,7 +174,7 @@ class IDAKLUSolver(pybamm.BaseSolver):
         else:
             inputs = inputs_dict
 
-        if model.jacobian_eval is None:
+        if model.jac_rhs_algebraic_eval is None:
             raise pybamm.SolverError("KLU requires the Jacobian to be provided")
 
         try:
@@ -184,7 +184,8 @@ class IDAKLUSolver(pybamm.BaseSolver):
 
         y0 = model.y0
         if isinstance(y0, casadi.DM):
-            y0 = y0.full().flatten()
+            y0 = y0.full()
+        y0 = y0.flatten()
 
         rtol = self.rtol
         atol = self._check_atol_type(atol, y0.size)
@@ -195,22 +196,28 @@ class IDAKLUSolver(pybamm.BaseSolver):
             mass_matrix = model.mass_matrix.entries
 
         # construct residuals function by binding inputs
-        def resfn(t, y, ydot):
-            return model.residuals_eval(t, y, ydot, inputs)
-
-        if model.jacobian_eval:
-            jac_y0_t0 = model.jacobian_eval(t_eval[0], y0, inputs)
-            if sparse.issparse(jac_y0_t0):
-
-                def jacfn(t, y, cj):
-                    j = model.jacobian_eval(t, y, inputs) - cj * mass_matrix
-                    return j
-
-            else:
+        if model.convert_to_format == "casadi":
+            def resfn(t, y, ydot):
+                return (
+                    model.rhs_algebraic_eval(t, y, inputs).full().flatten()
+                    - mass_matrix @ ydot
+                )
+        else:
+            def resfn(t, y, ydot):
+                return (
+                    model.rhs_algebraic_eval(t, y, inputs).flatten()
+                    - mass_matrix @ ydot
+                )
 
-                def jacfn(t, y, cj):
-                    jac_eval = model.jacobian_eval(t, y, inputs) - cj * mass_matrix
-                    return sparse.csr_matrix(jac_eval)
+        jac_y0_t0 = model.jac_rhs_algebraic_eval(t_eval[0], y0, inputs)
+        if sparse.issparse(jac_y0_t0):
+            def jacfn(t, y, cj):
+                j = model.jac_rhs_algebraic_eval(t, y, inputs) - cj * mass_matrix
+                return j
+        else:
+            def jacfn(t, y, cj):
+                jac_eval = model.jac_rhs_algebraic_eval(t, y, inputs) - cj * mass_matrix
+                return sparse.csr_matrix(jac_eval)
 
         class SundialsJacobian:
             def __init__(self):
@@ -257,8 +264,8 @@ class IDAKLUSolver(pybamm.BaseSolver):
         ids = np.concatenate((rhs_ids, alg_ids))
 
         number_of_sensitivity_parameters = 0
-        if model.sensitivities_eval is not None:
-            sens0 = model.sensitivities_eval(t=0, y=y0, inputs=inputs)
+        if model.jacp_rhs_algebraic_eval is not None:
+            sens0 = model.jacp_rhs_algebraic_eval(0, y0, inputs)
             number_of_sensitivity_parameters = len(sens0.keys())
 
         def sensfn(resvalS, t, y, yp, yS, ypS):
@@ -288,9 +295,9 @@ class IDAKLUSolver(pybamm.BaseSolver):
 
             """
 
-            dFdy = model.jacobian_eval(t, y, inputs)
+            dFdy = model.jac_rhs_algebraic_eval(t, y, inputs)
             dFdyd = mass_matrix
-            dFdp = model.sensitivities_eval(t, y, inputs)
+            dFdp = model.jacp_rhs_algebraic_eval(t, y, inputs)
 
             for i, dFdp_i in enumerate(dFdp.values()):
                 resvalS[i][:] = dFdy @ yS[i] - dFdyd @ ypS[i] + dFdp_i
diff --git a/pybamm/solvers/scikits_dae_solver.py b/pybamm/solvers/scikits_dae_solver.py
index eff037cc7..cdc86d81e 100644
--- a/pybamm/solvers/scikits_dae_solver.py
+++ b/pybamm/solvers/scikits_dae_solver.py
@@ -94,15 +94,29 @@ class ScikitsDaeSolver(pybamm.BaseSolver):
 
         y0 = model.y0
         if isinstance(y0, casadi.DM):
-            y0 = y0.full().flatten()
+            y0 = y0.full()
+        y0 = y0.flatten()
 
-        residuals = model.residuals_eval
+        rhs_algebraic_eval = model.rhs_algebraic_eval
         events = model.terminate_events_eval
-        jacobian = model.jacobian_eval
-        mass_matrix = model.mass_matrix.entries
+        jacobian = model.jac_rhs_algebraic_eval
+        if model.convert_to_format == "jax":
+            mass_matrix = model.mass_matrix.entries.toarray()
+        else:
+            mass_matrix = model.mass_matrix.entries
 
-        def eqsres(t, y, ydot, return_residuals):
-            return_residuals[:] = residuals(t, y, ydot, inputs)
+        if model.convert_to_format == "casadi":
+            def eqsres(t, y, ydot, return_residuals):
+                return_residuals[:] = (
+                    rhs_algebraic_eval(t, y, inputs).full().flatten()
+                    - mass_matrix @ ydot
+                )
+        else:
+            def eqsres(t, y, ydot, return_residuals):
+                return_residuals[:] = (
+                    rhs_algebraic_eval(t, y, inputs).flatten()
+                    - mass_matrix @ ydot
+                )
 
         def rootfn(t, y, ydot, return_root):
             return_root[:] = [event(t, y, inputs) for event in events]
@@ -117,13 +131,10 @@ class ScikitsDaeSolver(pybamm.BaseSolver):
         if jacobian:
             jac_y0_t0 = jacobian(t_eval[0], y0, inputs)
             if sparse.issparse(jac_y0_t0):
-
                 def jacfn(t, y, ydot, residuals, cj, J):
                     jac_eval = jacobian(t, y, inputs) - cj * mass_matrix
                     J[:][:] = jac_eval.toarray()
-
             else:
-
                 def jacfn(t, y, ydot, residuals, cj, J):
                     jac_eval = jacobian(t, y, inputs) - cj * mass_matrix
                     J[:][:] = jac_eval
diff --git a/pybamm/solvers/scikits_ode_solver.py b/pybamm/solvers/scikits_ode_solver.py
index 0c8913b9c..addeb728f 100644
--- a/pybamm/solvers/scikits_ode_solver.py
+++ b/pybamm/solvers/scikits_ode_solver.py
@@ -84,21 +84,26 @@ class ScikitsOdeSolver(pybamm.BaseSolver):
 
         """
         inputs_dict = inputs_dict or {}
-        if model.rhs_eval.form == "casadi":
+        if model.convert_to_format == "casadi":
             inputs = casadi.vertcat(*[x for x in inputs_dict.values()])
         else:
             inputs = inputs_dict
 
         y0 = model.y0
         if isinstance(y0, casadi.DM):
-            y0 = y0.full().flatten()
+            y0 = y0.full()
+        y0 = y0.flatten()
 
         derivs = model.rhs_eval
         events = model.terminate_events_eval
-        jacobian = model.jacobian_eval
+        jacobian = model.jac_rhs_eval
 
-        def eqsydot(t, y, return_ydot):
-            return_ydot[:] = derivs(t, y, inputs)
+        if model.convert_to_format == "casadi":
+            def eqsydot(t, y, return_ydot):
+                return_ydot[:] = derivs(t, y, inputs).full().flatten()
+        else:
+            def eqsydot(t, y, return_ydot):
+                return_ydot[:] = derivs(t, y, inputs).flatten()
 
         def rootfn(t, y, return_root):
             return_root[:] = [event(t, y, inputs) for event in events]
diff --git a/pybamm/solvers/scipy_solver.py b/pybamm/solvers/scipy_solver.py
index 38b30103a..b9a41f2e5 100644
--- a/pybamm/solvers/scipy_solver.py
+++ b/pybamm/solvers/scipy_solver.py
@@ -78,19 +78,29 @@ class ScipySolver(pybamm.BaseSolver):
         # Initial conditions
         y0 = model.y0
         if isinstance(y0, casadi.DM):
-            y0 = y0.full().flatten()
+            y0 = y0.full()
+        y0 = y0.flatten()
 
         # check for user-supplied Jacobian
         implicit_methods = ["Radau", "BDF", "LSODA"]
         if np.any([self.method in implicit_methods]):
-            if model.jacobian_eval:
+            if model.jac_rhs_eval:
+                def jacobian(t, y):
+                    return model.jac_rhs_eval(t, y, inputs)
                 extra_options.update(
-                    {"jac": lambda t, y: model.jacobian_eval(t, y, inputs)}
+                    {"jac": jacobian}
                 )
 
+        # rhs equation
+        if model.convert_to_format == 'casadi':
+            def rhs(t, y):
+                return model.rhs_eval(t, y, inputs).full().reshape(-1)
+        else:
+            def rhs(t, y):
+                return model.rhs_eval(t, y, inputs).reshape(-1)
+
         # make events terminal so that the solver stops when they are reached
         if model.terminate_events_eval:
-
             def event_wrapper(event):
                 def event_fn(t, y):
                     return event(t, y, inputs)
@@ -103,7 +113,7 @@ class ScipySolver(pybamm.BaseSolver):
 
         timer = pybamm.Timer()
         sol = it.solve_ivp(
-            lambda t, y: model.rhs_eval(t, y, inputs),
+            rhs,
             (t_eval[0], t_eval[-1]),
             y0,
             t_eval=t_eval,


Test patch:
diff --git a/tests/integration/test_models/standard_model_tests.py b/tests/integration/test_models/standard_model_tests.py
index bac6f9c86..66ceb58d7 100644
--- a/tests/integration/test_models/standard_model_tests.py
+++ b/tests/integration/test_models/standard_model_tests.py
@@ -191,10 +191,10 @@ class OptimisationsTest(object):
                 eqn_eval, known_evals = eqn.evaluate(0, y, known_evals={})
             elif to_python:
                 evaluator = pybamm.EvaluatorPython(eqn)
-                eqn_eval = evaluator.evaluate(0, y)
+                eqn_eval = evaluator(0, y)
             elif to_jax:
                 evaluator = pybamm.EvaluatorJax(eqn)
-                eqn_eval = evaluator.evaluate(0, y)
+                eqn_eval = evaluator(0, y)
             else:
                 eqn_eval = eqn.evaluate(0, y)
 
diff --git a/tests/unit/test_expression_tree/test_operations/test_evaluate_python.py b/tests/unit/test_expression_tree/test_operations/test_evaluate_python.py
index 9d30bf046..004f51898 100644
--- a/tests/unit/test_expression_tree/test_operations/test_evaluate_python.py
+++ b/tests/unit/test_expression_tree/test_operations/test_evaluate_python.py
@@ -206,13 +206,13 @@ class TestEvaluate(unittest.TestCase):
         )
 
         evaluator = pybamm.EvaluatorPython(expr)
-        result = evaluator.evaluate(y=y)
+        result = evaluator(y=y)
         np.testing.assert_allclose(result, expr.evaluate(y=y))
 
         # check that concatenating a single domain is consistent
         expr = pybamm.DomainConcatenation([a], mesh)
         evaluator = pybamm.EvaluatorPython(expr)
-        result = evaluator.evaluate(y=y)
+        result = evaluator(y=y)
         np.testing.assert_allclose(result, expr.evaluate(y=y))
 
         # check the reordering in case a child vector has to be split up
@@ -247,7 +247,7 @@ class TestEvaluate(unittest.TestCase):
         )
 
         evaluator = pybamm.EvaluatorPython(expr)
-        result = evaluator.evaluate(y=y)
+        result = evaluator(y=y)
         np.testing.assert_allclose(result, expr.evaluate(y=y))
 
     def test_domain_concatenation_2D(self):
@@ -273,13 +273,13 @@ class TestEvaluate(unittest.TestCase):
         self.assertEqual(len(constant_symbols), 0)
 
         evaluator = pybamm.EvaluatorPython(expr)
-        result = evaluator.evaluate(y=y)
+        result = evaluator(y=y)
         np.testing.assert_allclose(result, expr.evaluate(y=y))
 
         # check that concatenating a single domain is consistent
         expr = disc.process_symbol(pybamm.concatenation(a))
         evaluator = pybamm.EvaluatorPython(expr)
-        result = evaluator.evaluate(y=y)
+        result = evaluator(y=y)
         np.testing.assert_allclose(result, expr.evaluate(y=y))
 
     def test_to_python(self):
@@ -307,40 +307,40 @@ class TestEvaluate(unittest.TestCase):
         # test a * b
         expr = a * b
         evaluator = pybamm.EvaluatorPython(expr)
-        result = evaluator.evaluate(t=None, y=np.array([[2], [3]]))
+        result = evaluator(t=None, y=np.array([[2], [3]]))
         self.assertEqual(result, 6)
-        result = evaluator.evaluate(t=None, y=np.array([[1], [3]]))
+        result = evaluator(t=None, y=np.array([[1], [3]]))
         self.assertEqual(result, 3)
 
         # test function(a*b)
         expr = pybamm.Function(test_function, a * b)
         evaluator = pybamm.EvaluatorPython(expr)
-        result = evaluator.evaluate(t=None, y=np.array([[2], [3]]))
+        result = evaluator(t=None, y=np.array([[2], [3]]))
         self.assertEqual(result, 12)
 
         expr = pybamm.Function(test_function2, a, b)
         evaluator = pybamm.EvaluatorPython(expr)
-        result = evaluator.evaluate(t=None, y=np.array([[2], [3]]))
+        result = evaluator(t=None, y=np.array([[2], [3]]))
         self.assertEqual(result, 5)
 
         # test a constant expression
         expr = pybamm.Scalar(2) * pybamm.Scalar(3)
         evaluator = pybamm.EvaluatorPython(expr)
-        result = evaluator.evaluate()
+        result = evaluator()
         self.assertEqual(result, 6)
 
         # test a larger expression
         expr = a * b + b + a ** 2 / b + 2 * a + b / 2 + 4
         evaluator = pybamm.EvaluatorPython(expr)
         for y in y_tests:
-            result = evaluator.evaluate(t=None, y=y)
+            result = evaluator(t=None, y=y)
             self.assertEqual(result, expr.evaluate(t=None, y=y))
 
         # test something with time
         expr = a * pybamm.t
         evaluator = pybamm.EvaluatorPython(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             self.assertEqual(result, expr.evaluate(t=t, y=y))
 
         # test something with a matrix multiplication
@@ -348,7 +348,7 @@ class TestEvaluate(unittest.TestCase):
         expr = A @ pybamm.StateVector(slice(0, 2))
         evaluator = pybamm.EvaluatorPython(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         # test something with a heaviside
@@ -356,13 +356,13 @@ class TestEvaluate(unittest.TestCase):
         expr = a <= pybamm.StateVector(slice(0, 2))
         evaluator = pybamm.EvaluatorPython(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         expr = a > pybamm.StateVector(slice(0, 2))
         evaluator = pybamm.EvaluatorPython(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         # test something with a minimum or maximum
@@ -370,20 +370,20 @@ class TestEvaluate(unittest.TestCase):
         expr = pybamm.minimum(a, pybamm.StateVector(slice(0, 2)))
         evaluator = pybamm.EvaluatorPython(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         expr = pybamm.maximum(a, pybamm.StateVector(slice(0, 2)))
         evaluator = pybamm.EvaluatorPython(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         # test something with an index
         expr = pybamm.Index(A @ pybamm.StateVector(slice(0, 2)), 0)
         evaluator = pybamm.EvaluatorPython(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             self.assertEqual(result, expr.evaluate(t=t, y=y))
 
         # test something with a sparse matrix multiplication
@@ -393,13 +393,13 @@ class TestEvaluate(unittest.TestCase):
         expr = A @ B @ C @ pybamm.StateVector(slice(0, 2))
         evaluator = pybamm.EvaluatorPython(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         expr = B @ pybamm.StateVector(slice(0, 2))
         evaluator = pybamm.EvaluatorPython(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         # test numpy concatenation
@@ -412,12 +412,12 @@ class TestEvaluate(unittest.TestCase):
         expr = pybamm.NumpyConcatenation(a, b)
         evaluator = pybamm.EvaluatorPython(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
         expr = pybamm.NumpyConcatenation(a, c)
         evaluator = pybamm.EvaluatorPython(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         # test sparse stack
@@ -427,18 +427,18 @@ class TestEvaluate(unittest.TestCase):
         expr = pybamm.SparseStack(A, a * B)
         evaluator = pybamm.EvaluatorPython(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y).toarray()
+            result = evaluator(t=t, y=y).toarray()
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y).toarray())
         expr = pybamm.SparseStack(A)
         evaluator = pybamm.EvaluatorPython(expr)
-        result = evaluator.evaluate().toarray()
+        result = evaluator().toarray()
         np.testing.assert_allclose(result, expr.evaluate().toarray())
 
         # test Inner
         expr = pybamm.Inner(a, b)
         evaluator = pybamm.EvaluatorPython(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         v = pybamm.StateVector(slice(0, 2))
@@ -446,7 +446,7 @@ class TestEvaluate(unittest.TestCase):
         for expr in [pybamm.Inner(A, v), pybamm.Inner(v, A)]:
             evaluator = pybamm.EvaluatorPython(expr)
             for t, y in zip(t_tests, y_tests):
-                result = evaluator.evaluate(t=t, y=y).toarray()
+                result = evaluator(t=t, y=y).toarray()
                 np.testing.assert_allclose(result, expr.evaluate(t=t, y=y).toarray())
 
         y_tests = [np.array([[2], [3], [4], [5]]), np.array([[1], [3], [2], [1]])]
@@ -456,7 +456,7 @@ class TestEvaluate(unittest.TestCase):
         expr = a * b
         evaluator = pybamm.EvaluatorPython(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
     @unittest.skipIf(not pybamm.have_jax(), "jax or jaxlib is not installed")
@@ -487,41 +487,41 @@ class TestEvaluate(unittest.TestCase):
         # test a * b
         expr = a * b
         evaluator = pybamm.EvaluatorJax(expr)
-        result = evaluator.evaluate(t=None, y=np.array([[2], [3]]))
+        result = evaluator(t=None, y=np.array([[2], [3]]))
         self.assertEqual(result, 6)
-        result = evaluator.evaluate(t=None, y=np.array([[1], [3]]))
+        result = evaluator(t=None, y=np.array([[1], [3]]))
         self.assertEqual(result, 3)
 
         # test function(a*b)
         expr = pybamm.Function(test_function, a * b)
         evaluator = pybamm.EvaluatorJax(expr)
-        result = evaluator.evaluate(t=None, y=np.array([[2], [3]]))
+        result = evaluator(t=None, y=np.array([[2], [3]]))
         self.assertEqual(result, 12)
 
         # test exp
         expr = pybamm.exp(a * b)
         evaluator = pybamm.EvaluatorJax(expr)
-        result = evaluator.evaluate(t=None, y=np.array([[2], [3]]))
+        result = evaluator(t=None, y=np.array([[2], [3]]))
         self.assertEqual(result, np.exp(6))
 
         # test a constant expression
         expr = pybamm.Scalar(2) * pybamm.Scalar(3)
         evaluator = pybamm.EvaluatorJax(expr)
-        result = evaluator.evaluate()
+        result = evaluator()
         self.assertEqual(result, 6)
 
         # test a larger expression
         expr = a * b + b + a ** 2 / b + 2 * a + b / 2 + 4
         evaluator = pybamm.EvaluatorJax(expr)
         for y in y_tests:
-            result = evaluator.evaluate(t=None, y=y)
+            result = evaluator(t=None, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=None, y=y))
 
         # test something with time
         expr = a * pybamm.t
         evaluator = pybamm.EvaluatorJax(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             self.assertEqual(result, expr.evaluate(t=t, y=y))
 
         # test something with a matrix multiplication
@@ -529,7 +529,7 @@ class TestEvaluate(unittest.TestCase):
         expr = A @ pybamm.StateVector(slice(0, 2))
         evaluator = pybamm.EvaluatorJax(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         # test something with a heaviside
@@ -537,13 +537,13 @@ class TestEvaluate(unittest.TestCase):
         expr = a <= pybamm.StateVector(slice(0, 2))
         evaluator = pybamm.EvaluatorJax(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         expr = a > pybamm.StateVector(slice(0, 2))
         evaluator = pybamm.EvaluatorJax(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         # test something with a minimum or maximum
@@ -551,20 +551,20 @@ class TestEvaluate(unittest.TestCase):
         expr = pybamm.minimum(a, pybamm.StateVector(slice(0, 2)))
         evaluator = pybamm.EvaluatorJax(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         expr = pybamm.maximum(a, pybamm.StateVector(slice(0, 2)))
         evaluator = pybamm.EvaluatorJax(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         # test something with an index
         expr = pybamm.Index(A @ pybamm.StateVector(slice(0, 2)), 0)
         evaluator = pybamm.EvaluatorJax(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             self.assertEqual(result, expr.evaluate(t=t, y=y))
 
         # test something with a sparse matrix-vector multiplication
@@ -574,7 +574,7 @@ class TestEvaluate(unittest.TestCase):
         expr = A @ B @ C @ pybamm.StateVector(slice(0, 2))
         evaluator = pybamm.EvaluatorJax(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         # test the sparse-scalar multiplication
@@ -585,7 +585,7 @@ class TestEvaluate(unittest.TestCase):
         ]:
             evaluator = pybamm.EvaluatorJax(expr)
             for t, y in zip(t_tests, y_tests):
-                result = evaluator.evaluate(t=t, y=y)
+                result = evaluator(t=t, y=y)
                 np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         # test the sparse-scalar division
@@ -593,7 +593,7 @@ class TestEvaluate(unittest.TestCase):
         expr = A / (1.0 + pybamm.t) @ pybamm.StateVector(slice(0, 2))
         evaluator = pybamm.EvaluatorJax(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         # test sparse stack
@@ -618,7 +618,7 @@ class TestEvaluate(unittest.TestCase):
         expr = pybamm.NumpyConcatenation(a, b)
         evaluator = pybamm.EvaluatorJax(expr)
         for t, y in zip(t_tests, y_tests):
-            result = evaluator.evaluate(t=t, y=y)
+            result = evaluator(t=t, y=y)
             np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
         # test Inner
@@ -631,7 +631,7 @@ class TestEvaluate(unittest.TestCase):
         ]:
             evaluator = pybamm.EvaluatorJax(expr)
             for t, y in zip(t_tests, y_tests):
-                result = evaluator.evaluate(t=t, y=y)
+                result = evaluator(t=t, y=y)
                 np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
 
     @unittest.skipIf(not pybamm.have_jax(), "jax or jaxlib is not installed")
@@ -645,15 +645,15 @@ class TestEvaluate(unittest.TestCase):
         evaluator_jac_test = evaluator.get_jacobian()
         evaluator_jac = pybamm.EvaluatorJax(expr_jac)
         for y in y_tests:
-            result_test = evaluator_jac_test.evaluate(t=None, y=y)
-            result_true = evaluator_jac.evaluate(t=None, y=y)
+            result_test = evaluator_jac_test(t=None, y=y)
+            result_true = evaluator_jac(t=None, y=y)
             np.testing.assert_allclose(result_test, result_true)
 
     @unittest.skipIf(not pybamm.have_jax(), "jax or jaxlib is not installed")
     def test_evaluator_jax_debug(self):
         a = pybamm.StateVector(slice(0, 1))
         expr = a ** 2
-        y_test = np.array([[2.0], [3.0]])
+        y_test = np.array([2.0, 3.0])
         evaluator = pybamm.EvaluatorJax(expr)
         evaluator.debug(y=y_test)
 
@@ -662,7 +662,7 @@ class TestEvaluate(unittest.TestCase):
         a = pybamm.InputParameter("a")
         expr = a ** 2
         evaluator = pybamm.EvaluatorJax(expr)
-        result = evaluator.evaluate(inputs={"a": 2})
+        result = evaluator(inputs={"a": 2})
         self.assertEqual(result, 4)
 
     @unittest.skipIf(not pybamm.have_jax(), "jax or jaxlib is not installed")
diff --git a/tests/unit/test_solvers/test_algebraic_solver.py b/tests/unit/test_solvers/test_algebraic_solver.py
index 15e8a7bbd..74c8fe2eb 100644
--- a/tests/unit/test_solvers/test_algebraic_solver.py
+++ b/tests/unit/test_solvers/test_algebraic_solver.py
@@ -44,9 +44,12 @@ class TestAlgebraicSolver(unittest.TestCase):
             timescale_eval = 1
             length_scales = {}
             jac_algebraic_eval = None
-            convert_to_format = "python"
             len_rhs_and_alg = 1
 
+            def __init__(self):
+                super().__init__()
+                self.convert_to_format = "python"
+
             def algebraic_eval(self, t, y, inputs):
                 return y + 2
 
@@ -68,9 +71,12 @@ class TestAlgebraicSolver(unittest.TestCase):
             timescale_eval = 1
             length_scales = {}
             jac_algebraic_eval = None
-            convert_to_format = "casadi"
             len_rhs_and_alg = 1
 
+            def __init__(self):
+                super().__init__()
+                self.convert_to_format = "python"
+
             def algebraic_eval(self, t, y, inputs):
                 # algebraic equation has no real root
                 return y ** 2 + 1
@@ -99,9 +105,12 @@ class TestAlgebraicSolver(unittest.TestCase):
             rhs = {}
             timescale_eval = 1
             length_scales = {}
-            convert_to_format = "python"
             len_rhs_and_alg = 2
 
+            def __init__(self):
+                super().__init__()
+                self.convert_to_format = "python"
+
             def algebraic_eval(self, t, y, inputs):
                 return A @ y - b
 
diff --git a/tests/unit/test_solvers/test_base_solver.py b/tests/unit/test_solvers/test_base_solver.py
index 9d5bcd9bc..a96adab59 100644
--- a/tests/unit/test_solvers/test_base_solver.py
+++ b/tests/unit/test_solvers/test_base_solver.py
@@ -175,7 +175,7 @@ class TestBaseSolver(unittest.TestCase):
 
         model = VectorModel()
         init_cond = solver.calculate_consistent_state(model)
-        np.testing.assert_array_almost_equal(init_cond, vec)
+        np.testing.assert_array_almost_equal(init_cond.flatten(), vec)
         # with casadi
         init_cond = solver_with_casadi.calculate_consistent_state(model)
         np.testing.assert_array_almost_equal(init_cond.full().flatten(), vec)
@@ -186,7 +186,7 @@ class TestBaseSolver(unittest.TestCase):
 
         model.jac_algebraic_eval = jac_dense
         init_cond = solver.calculate_consistent_state(model)
-        np.testing.assert_array_almost_equal(init_cond, vec)
+        np.testing.assert_array_almost_equal(init_cond.flatten(), vec)
 
         # With sparse jacobian
         def jac_sparse(t, y, inputs):
@@ -196,7 +196,7 @@ class TestBaseSolver(unittest.TestCase):
 
         model.jac_algebraic_eval = jac_sparse
         init_cond = solver.calculate_consistent_state(model)
-        np.testing.assert_array_almost_equal(init_cond, vec)
+        np.testing.assert_array_almost_equal(init_cond.flatten(), vec)
 
     def test_fail_consistent_initial_conditions(self):
         class Model:
@@ -375,7 +375,8 @@ class TestBaseSolver(unittest.TestCase):
                 else:
                     use_inputs = inputs
 
-                sens = model.sensitivities_eval(t, y, use_inputs)
+                sens = model.jacp_rhs_algebraic_eval(t, y, use_inputs)
+
                 np.testing.assert_allclose(
                     sens["a"], exact_diff_a(y, inputs["a"], inputs["b"])
                 )
diff --git a/tests/unit/test_solvers/test_jax_bdf_solver.py b/tests/unit/test_solvers/test_jax_bdf_solver.py
index 8dcdb04a6..1a61df1e1 100644
--- a/tests/unit/test_solvers/test_jax_bdf_solver.py
+++ b/tests/unit/test_solvers/test_jax_bdf_solver.py
@@ -33,7 +33,7 @@ class TestJaxBDFSolver(unittest.TestCase):
         rhs = pybamm.EvaluatorJax(model.concatenated_rhs)
 
         def fun(y, t):
-            return rhs.evaluate(t=t, y=y).reshape(-1)
+            return rhs(t=t, y=y).reshape(-1)
 
         t0 = time.perf_counter()
         y = pybamm.jax_bdf_integrate(fun, y0, t_eval, rtol=1e-8, atol=1e-8)
@@ -105,7 +105,7 @@ class TestJaxBDFSolver(unittest.TestCase):
         rhs = pybamm.EvaluatorJax(model.concatenated_rhs)
 
         def fun(y, t, inputs):
-            return rhs.evaluate(t=t, y=y, inputs=inputs).reshape(-1)
+            return rhs(t=t, y=y, inputs=inputs).reshape(-1)
 
         h = 0.0001
         rate = 0.1
@@ -183,7 +183,7 @@ class TestJaxBDFSolver(unittest.TestCase):
         rhs = pybamm.EvaluatorJax(model.concatenated_rhs)
 
         def fun(y, t, inputs):
-            return rhs.evaluate(t=t, y=y, inputs=inputs).reshape(-1)
+            return rhs(t=t, y=y, inputs=inputs).reshape(-1)
 
         y = pybamm.jax_bdf_integrate(
             fun, y0, t_eval, {"rate": 0.1}, rtol=1e-9, atol=1e-9
diff --git a/tests/unit/test_solvers/test_scikits_solvers.py b/tests/unit/test_solvers/test_scikits_solvers.py
index de4f496cc..d6f6e561d 100644
--- a/tests/unit/test_solvers/test_scikits_solvers.py
+++ b/tests/unit/test_solvers/test_scikits_solvers.py
@@ -49,10 +49,10 @@ class TestScikitsSolvers(unittest.TestCase):
             length_scales = {}
             convert_to_format = "python"
 
-            def residuals_eval(self, t, y, ydot, inputs):
-                return np.array([0.5 * np.ones_like(y[0]) - ydot[0], 2 * y[0] - y[1]])
+            def rhs_algebraic_eval(self, t, y, inputs):
+                return np.array([0.5 * np.ones_like(y[0]), 2 * y[0] - y[1]])
 
-            def jacobian_eval(self, t, y, inputs):
+            def jac_rhs_algebraic_eval(self, t, y, inputs):
                 return np.array([[0.0, 0.0], [2.0, -1.0]])
 
         model = Model()
@@ -101,12 +101,12 @@ class TestScikitsSolvers(unittest.TestCase):
             convert_to_format = "python"
             len_rhs_and_alg = 2
 
-            def residuals_eval(self, t, y, ydot, inputs):
+            def rhs_algebraic_eval(self, t, y, inputs):
                 return np.array(
-                    [0.5 * np.ones_like(y[0]) - 4 * ydot[0], 2.0 * y[0] - y[1]]
+                    [0.5 * np.ones_like(y[0]), 2.0 * y[0] - y[1]]
                 )
 
-            def jacobian_eval(self, t, y, inputs):
+            def jac_rhs_algebraic_eval(self, t, y, inputs):
                 return np.array([[0.0, 0.0], [2.0, -1.0]])
 
         model = Model()
@@ -214,6 +214,28 @@ class TestScikitsSolvers(unittest.TestCase):
         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))
 
+    @unittest.skipIf(not pybamm.have_jax(), "jax or jaxlib is not installed")
+    def test_model_solver_dae_jax(self):
+        model = pybamm.BaseModel()
+        model.convert_to_format = "jax"
+        whole_cell = ["negative electrode", "separator", "positive electrode"]
+        var1 = pybamm.Variable("var1", domain=whole_cell)
+        var2 = pybamm.Variable("var2", domain=whole_cell)
+        model.rhs = {var1: 0.1 * var1}
+        model.algebraic = {var2: 2 * var1 - var2}
+        model.initial_conditions = {var1: 1, var2: 2}
+        model.use_jacobian = False
+        disc = get_discretisation_for_testing()
+        disc.process_model(model)
+
+        # Solve
+        solver = pybamm.ScikitsDaeSolver(rtol=1e-8, atol=1e-8, root_method="lm")
+        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))
+
     def test_model_solver_dae_bad_ics_python(self):
         model = pybamm.BaseModel()
         model.convert_to_format = "python"
@@ -760,9 +782,11 @@ class TestScikitsSolvers(unittest.TestCase):
     def test_model_step_nonsmooth_events(self):
         # Create model
         model = pybamm.BaseModel()
-        model.timescale = pybamm.Scalar(1)
         var1 = pybamm.Variable("var1")
         var2 = pybamm.Variable("var2")
+
+        # if this is 1 it gets simplified out
+        model.timescale = pybamm.Scalar(1.000001)
         a = 0.6
         discontinuities = (np.arange(3) + 1) * a
 
@@ -787,7 +811,9 @@ class TestScikitsSolvers(unittest.TestCase):
         end_time = 3
         step_solution = None
         while time < end_time:
-            step_solution = step_solver.step(step_solution, model, dt=dt, npts=10)
+            step_solution = step_solver.step(
+                step_solution, model, dt=dt, npts=10
+            )
             time += dt
         np.testing.assert_array_less(step_solution.y[0, :-1], 0.55)
         np.testing.assert_array_less(step_solution.y[-1, :-1], 1.2)