def jac(self, variable, known_jacs=None, clear_domain=True): """ Differentiate a symbol with respect to a (slice of) a State Vector. See :class:`pybamm.Jacobian`. """ jac = pybamm.Jacobian(known_jacs, clear_domain=clear_domain) return jac.jac(self, variable)
def create_jacobian(self, model): """Creates Jacobian of the discretised model. Note that the model is assumed to be of the form M*y_dot = f(t,y), where M is the (possibly singular) mass matrix. The Jacobian is df/dy. Note: At present, calculation of the Jacobian is deferred until after simplification, since it is much faster to compute the Jacobian of the simplified model. However, in some use cases (e.g. running the same model multiple times but with different parameters) it may be more efficient to compute the Jacobian once, before simplification, so that parameters in the Jacobian can be updated (see PR #670). Parameters ---------- model : :class:`pybamm.BaseModel` Discretised model. Must have attributes rhs, initial_conditions and boundary_conditions (all dicts of {variable: equation}) Returns ------- :class:`pybamm.Concatenation` The expression trees corresponding to the Jacobian of the model """ # create state vector to differentiate with respect to y = pybamm.StateVector( slice(0, np.size(model.concatenated_initial_conditions))) # set up Jacobian object, for re-use of dict jacobian = pybamm.Jacobian() # calculate Jacobian of rhs by equation jac_rhs_eqn_dict = {} for eqn_key, eqn in model.rhs.items(): pybamm.logger.debug( "Calculating block of Jacobian for {!r}".format(eqn_key.name)) jac_rhs_eqn_dict[eqn_key] = jacobian.jac(eqn, y) jac_rhs = self._concatenate_in_order(jac_rhs_eqn_dict, sparse=True) # calculate Jacobian of algebraic by equation jac_algebraic_eqn_dict = {} for eqn_key, eqn in model.algebraic.items(): pybamm.logger.debug( "Calculating block of Jacobian for {!r}".format(eqn_key.name)) jac_algebraic_eqn_dict[eqn_key] = jacobian.jac(eqn, y) jac_algebraic = self._concatenate_in_order(jac_algebraic_eqn_dict, sparse=True) # full Jacobian if model.rhs.keys() and model.algebraic.keys(): jac = pybamm.SparseStack(jac_rhs, jac_algebraic) elif not model.algebraic.keys(): jac = jac_rhs else: jac = jac_algebraic return jac, jac_rhs, jac_algebraic
def jac(self, variable, known_jacs=None, clear_domain=True): """ Differentiate a symbol with respect to a (slice of) a StateVector or StateVectorDot. See :class:`pybamm.Jacobian`. """ jac = pybamm.Jacobian(known_jacs, clear_domain=clear_domain) if not isinstance(variable, (pybamm.StateVector, pybamm.StateVectorDot)): raise TypeError( "Jacobian can only be taken with respect to a 'StateVector' " "or 'StateVectorDot', but {} is a {}".format(variable, type(variable)) ) return jac.jac(self, variable)
def jac(self, variable, known_jacs=None): """ Differentiate a symbol with respect to a (slice of) a State Vector. See :class:`pybamm.Jacobian`. """ return pybamm.Jacobian(known_jacs).jac(self, variable)
def set_up(self, model): """Unpack model, perform checks, simplify and calculate jacobian. Parameters ---------- model : :class:`pybamm.BaseModel` The model whose solution to calculate. Must have attributes rhs and initial_conditions Raises ------ :class:`pybamm.SolverError` If the model contains any algebraic equations (in which case a DAE solver should be used instead) """ # Check for algebraic equations if len(model.algebraic) > 0: raise pybamm.SolverError( """Cannot use ODE solver to solve model with DAEs""") # create simplified rhs and event expressions concatenated_rhs = model.concatenated_rhs events = model.events if model.use_simplify: # set up simplification object, for re-use of dict simp = pybamm.Simplification() # create simplified rhs and event expressions pybamm.logger.info("Simplifying RHS") concatenated_rhs = simp.simplify(concatenated_rhs) pybamm.logger.info("Simplifying events") events = { name: simp.simplify(event) for name, event in events.items() } y0 = model.concatenated_initial_conditions[:, 0] if model.use_jacobian: # Create Jacobian from concatenated rhs y = pybamm.StateVector(slice(0, np.size(y0))) # set up Jacobian object, for re-use of dict jacobian = pybamm.Jacobian() pybamm.logger.info("Calculating jacobian") jac_rhs = jacobian.jac(concatenated_rhs, y) model.jacobian = jac_rhs model.jacobian_rhs = jac_rhs if model.use_simplify: pybamm.logger.info("Simplifying jacobian") jac_rhs = simp.simplify(jac_rhs) if model.convert_to_format == "python": pybamm.logger.info("Converting jacobian to python") jac_rhs = pybamm.EvaluatorPython(jac_rhs) else: jac_rhs = None if model.convert_to_format == "python": pybamm.logger.info("Converting RHS to python") concatenated_rhs = pybamm.EvaluatorPython(concatenated_rhs) pybamm.logger.info("Converting events to python") events = { name: pybamm.EvaluatorPython(event) for name, event in events.items() } # Create function to evaluate rhs def dydt(t, y): pybamm.logger.debug("Evaluating RHS for {} at t={}".format( model.name, t)) y = y[:, np.newaxis] dy = concatenated_rhs.evaluate(t, y, known_evals={})[0] return dy[:, 0] # Create event-dependent function to evaluate events def event_fun(event): def eval_event(t, y): return event.evaluate(t, y) return eval_event event_funs = [event_fun(event) for event in events.values()] # Create function to evaluate jacobian if jac_rhs is not None: def jacobian(t, y): return jac_rhs.evaluate(t, y, known_evals={})[0] else: jacobian = None # Add the solver attributes # Note: these are the (possibly) converted to python version rhs, algebraic # etc. The expression tree versions of these are attributes of the model self.y0 = y0 self.dydt = dydt self.events = events self.event_funs = event_funs self.jacobian = jacobian
def set_up(self, model, inputs=None): """Unpack model, perform checks, simplify and calculate jacobian. Parameters ---------- model : :class:`pybamm.BaseModel` The model whose solution to calculate. Must have attributes rhs and initial_conditions inputs : dict, optional Any input parameters to pass to the model when solving """ inputs = inputs or {} y0 = model.concatenated_initial_conditions.evaluate(0, None, inputs) # Set model timescale model.timescale_eval = model.timescale.evaluate(u=inputs) # Check model.algebraic for ode solvers if self.ode_solver is True and len(model.algebraic) > 0: raise pybamm.SolverError( "Cannot use ODE solver '{}' to solve DAE model".format( self.name)) if self.ode_solver is True: self.root_method = None if (isinstance(self, pybamm.CasadiSolver) or self.root_method == "casadi") and model.convert_to_format != "casadi": pybamm.logger.warning( f"Converting {model.name} to CasADi for solving with CasADi solver" ) model.convert_to_format = "casadi" if model.convert_to_format != "casadi": simp = pybamm.Simplification() # Create Jacobian from concatenated rhs and algebraic y = pybamm.StateVector(slice(0, np.size(y0))) # set up Jacobian object, for re-use of dict jacobian = pybamm.Jacobian() else: # Convert model attributes to casadi t_casadi = casadi.MX.sym("t") y_diff = casadi.MX.sym( "y_diff", len(model.concatenated_rhs.evaluate(0, y0, inputs))) y_alg = casadi.MX.sym( "y_alg", len(model.concatenated_algebraic.evaluate(0, y0, inputs))) y_casadi = casadi.vertcat(y_diff, y_alg) u_casadi = {} for name, value in inputs.items(): if isinstance(value, numbers.Number): u_casadi[name] = casadi.MX.sym(name) else: u_casadi[name] = casadi.MX.sym(name, value.shape[0]) u_casadi_stacked = casadi.vertcat(*[u for u in u_casadi.values()]) def process(func, name, use_jacobian=None): def report(string): # don't log event conversion if "event" not in string: pybamm.logger.info(string) if use_jacobian is None: use_jacobian = model.use_jacobian if model.convert_to_format != "casadi": # Process with pybamm functions if model.use_simplify: report(f"Simplifying {name}") func = simp.simplify(func) if use_jacobian: report(f"Calculating jacobian for {name}") jac = jacobian.jac(func, y) if model.use_simplify: report(f"Simplifying jacobian for {name}") jac = simp.simplify(jac) 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 else: # Process with CasADi report(f"Converting {name} to CasADi") func = func.to_casadi(t_casadi, y_casadi, u_casadi) if use_jacobian: report(f"Calculating jacobian for {name} using CasADi") jac_casadi = casadi.jacobian(func, y_casadi) jac = casadi.Function( name, [t_casadi, y_casadi, u_casadi_stacked], [jac_casadi]) else: jac = None func = casadi.Function(name, [t_casadi, y_casadi, u_casadi_stacked], [func]) if name == "residuals": func_call = Residuals(func, name, model) else: func_call = SolverCallable(func, name, model) func_call.set_inputs(inputs) if jac is not None: jac_call = SolverCallable(jac, name + "_jac", model) jac_call.set_inputs(inputs) else: jac_call = None return func, func_call, jac_call # Check for heaviside 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, pybamm.Heaviside): # Dimensionless if symbol.right.id == pybamm.t.id: expr = symbol.left elif symbol.left.id == pybamm.t.id: expr = symbol.right # Dimensional elif symbol.right.id == (pybamm.t * model.timescale).id: expr = symbol.left.new_copy( ) / symbol.right.right.new_copy() elif symbol.left.id == (pybamm.t * model.timescale).id: expr = symbol.right.new_copy( ) / symbol.left.right.new_copy() model.events.append( pybamm.Event(str(symbol), expr.new_copy(), pybamm.EventType.DISCONTINUITY)) # Process rhs, algebraic and event expressions rhs, rhs_eval, jac_rhs = process(model.concatenated_rhs, "RHS") algebraic, algebraic_eval, jac_algebraic = process( model.concatenated_algebraic, "algebraic") terminate_events_eval = [ process(event.expression, "event", use_jacobian=False)[1] for event in model.events if event.event_type == pybamm.EventType.TERMINATION ] # discontinuity events are evaluated before the solver is called, so don't need # to process them discontinuity_events_eval = [ event for event in model.events if event.event_type == pybamm.EventType.DISCONTINUITY ] # Add the solver attributes model.rhs_eval = rhs_eval model.algebraic_eval = algebraic_eval model.jac_algebraic_eval = jac_algebraic model.terminate_events_eval = terminate_events_eval model.discontinuity_events_eval = discontinuity_events_eval # 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 self.root_method == "casadi" or isinstance(self, pybamm.CasadiSolver): mass_matrix_inv = casadi.MX(model.mass_matrix_inv.entries) explicit_rhs = mass_matrix_inv @ rhs(t_casadi, y_casadi, u_casadi_stacked) model.casadi_rhs = casadi.Function( "rhs", [t_casadi, y_casadi, u_casadi_stacked], [explicit_rhs]) model.casadi_algebraic = algebraic # Calculate consistent initial conditions for the algebraic equations if len(model.algebraic) > 0: all_states = pybamm.NumpyConcatenation( model.concatenated_rhs, model.concatenated_algebraic) # Process again, uses caching so should be quick residuals, residuals_eval, jacobian_eval = process( all_states, "residuals") model.residuals_eval = residuals_eval model.jacobian_eval = jacobian_eval y0_guess = y0.flatten() model.y0 = self.calculate_consistent_state(model, 0, y0_guess, inputs) else: # can use DAE solver to solve ODE model model.residuals_eval = Residuals(rhs, "residuals", model) model.jacobian_eval = jac_rhs model.y0 = y0.flatten() pybamm.logger.info("Finish solver set-up")
def set_up(self, model): """Unpack model, perform checks, simplify and calculate jacobian. Parameters ---------- model : :class:`pybamm.BaseModel` The model whose solution to calculate. Must have attributes rhs and initial_conditions Returns ------- concatenated_algebraic : :class:`pybamm.Concatenation` Algebraic equations, which should evaluate to zero jac : :class:`pybamm.SparseStack` Jacobian matrix for the differential and algebraic equations Raises ------ :class:`pybamm.SolverError` If the model contains any time derivatives, i.e. rhs equations (in which case an ODE or DAE solver should be used instead) """ if len(model.rhs) > 0: raise pybamm.SolverError( """Cannot use algebraic solver to solve model with time derivatives""" ) # create simplified algebraic expressions concatenated_algebraic = model.concatenated_algebraic if model.use_simplify: # set up simplification object, for re-use of dict simp = pybamm.Simplification() pybamm.logger.info("Simplifying algebraic") concatenated_algebraic = simp.simplify(concatenated_algebraic) if model.use_jacobian: # Create Jacobian from concatenated algebraic y = pybamm.StateVector( slice(0, np.size(model.concatenated_initial_conditions))) # set up Jacobian object, for re-use of dict jacobian = pybamm.Jacobian() pybamm.logger.info("Calculating jacobian") jac = jacobian.jac(concatenated_algebraic, y) model.jacobian = jac model.jacobian_algebraic = jac if model.use_simplify: pybamm.logger.info("Simplifying jacobian") jac = simp.simplify(jac) if model.convert_to_format == "python": pybamm.logger.info("Converting jacobian to python") jac = pybamm.EvaluatorPython(jac) else: jac = None if model.convert_to_format == "python": pybamm.logger.info("Converting algebraic to python") concatenated_algebraic = pybamm.EvaluatorPython( concatenated_algebraic) return concatenated_algebraic, jac
def set_up(self, model, inputs=None, t_eval=None): """Unpack model, perform checks, simplify and calculate jacobian. Parameters ---------- model : :class:`pybamm.BaseModel` The model whose solution to calculate. Must have attributes rhs and initial_conditions inputs : dict, optional Any input parameters to pass to the model when solving t_eval : numeric type, optional The times (in seconds) at which to compute the solution """ # Check model.algebraic for ode solvers if self.ode_solver is True and len(model.algebraic) > 0: raise pybamm.SolverError( "Cannot use ODE solver '{}' to solve DAE model".format( self.name)) # Check model.rhs for algebraic solvers if self.algebraic_solver is True and len(model.rhs) > 0: raise pybamm.SolverError( """Cannot use algebraic solver to solve model with time derivatives""" ) # casadi solver won't allow solving algebraic model so we have to raise an # error here if isinstance(self, pybamm.CasadiSolver) and len(model.rhs) == 0: raise pybamm.SolverError( "Cannot use CasadiSolver to solve algebraic model, " "use CasadiAlgebraicSolver instead") # Discretise model if it isn't already discretised # This only works with purely 0D models, as otherwise the mesh and spatial # method should be specified by the user if model.is_discretised is False: try: disc = pybamm.Discretisation() disc.process_model(model) except pybamm.DiscretisationError as e: raise pybamm.DiscretisationError( "Cannot automatically discretise model, " "model should be discretised before solving ({})".format( e)) inputs = inputs or {} # Set model timescale model.timescale_eval = model.timescale.evaluate(inputs=inputs) # Set model lengthscales model.length_scales_eval = { domain: scale.evaluate(inputs=inputs) for domain, scale in model.length_scales.items() } if (isinstance(self, (pybamm.CasadiSolver, pybamm.CasadiAlgebraicSolver)) ) and model.convert_to_format != "casadi": pybamm.logger.warning( "Converting {} to CasADi for solving with CasADi solver". format(model.name)) model.convert_to_format = "casadi" if (isinstance(self.root_method, pybamm.CasadiAlgebraicSolver) and model.convert_to_format != "casadi"): pybamm.logger.warning( "Converting {} to CasADi for calculating ICs with CasADi". format(model.name)) model.convert_to_format = "casadi" if model.convert_to_format != "casadi": simp = pybamm.Simplification() # Create Jacobian from concatenated rhs and algebraic y = pybamm.StateVector( slice(0, model.concatenated_initial_conditions.size)) # set up Jacobian object, for re-use of dict jacobian = pybamm.Jacobian() else: # Convert model attributes to casadi t_casadi = casadi.MX.sym("t") y_diff = casadi.MX.sym("y_diff", model.concatenated_rhs.size) y_alg = casadi.MX.sym("y_alg", model.concatenated_algebraic.size) y_casadi = casadi.vertcat(y_diff, y_alg) p_casadi = {} for name, value in inputs.items(): if isinstance(value, numbers.Number): p_casadi[name] = casadi.MX.sym(name) else: p_casadi[name] = casadi.MX.sym(name, value.shape[0]) p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()]) def process(func, name, use_jacobian=None): def report(string): # don't log event conversion if "event" not in string: pybamm.logger.info(string) if use_jacobian is None: use_jacobian = model.use_jacobian if model.convert_to_format != "casadi": # Process with pybamm functions if model.use_simplify: report(f"Simplifying {name}") func = simp.simplify(func) if model.convert_to_format == "jax": report(f"Converting {name} to jax") jax_func = pybamm.EvaluatorJax(func) if use_jacobian: report(f"Calculating jacobian for {name}") jac = jacobian.jac(func, y) if model.use_simplify: report(f"Simplifying jacobian for {name}") jac = simp.simplify(jac) if model.convert_to_format == "python": report(f"Converting jacobian for {name} to python") jac = pybamm.EvaluatorPython(jac) elif model.convert_to_format == "jax": report(f"Converting jacobian for {name} to jax") jac = jax_func.get_jacobian() jac = jac.evaluate else: jac = None if model.convert_to_format == "python": report(f"Converting {name} to python") func = pybamm.EvaluatorPython(func) if model.convert_to_format == "jax": report(f"Converting {name} to jax") func = jax_func func = func.evaluate else: # Process with CasADi report(f"Converting {name} to CasADi") func = func.to_casadi(t_casadi, y_casadi, inputs=p_casadi) if use_jacobian: report(f"Calculating jacobian for {name} using CasADi") jac_casadi = casadi.jacobian(func, y_casadi) jac = casadi.Function( name, [t_casadi, y_casadi, p_casadi_stacked], [jac_casadi]) else: jac = None func = casadi.Function(name, [t_casadi, y_casadi, p_casadi_stacked], [func]) 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 return func, func_call, jac_call # 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, pybamm.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).id: expr = symbol.left.new_copy( ) / symbol.right.right.new_copy() found_t = True elif symbol.left.id == (pybamm.t * model.timescale).id: expr = symbol.right.new_copy( ) / symbol.left.right.new_copy() found_t = True # Update the events if the heaviside function depended on t if found_t: model.events.append( pybamm.Event( str(symbol), expr.new_copy(), 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).id: expr = symbol.right.new_copy( ) / symbol.left.right.new_copy() 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.new_copy() * pybamm.Scalar(i + 1), pybamm.EventType.DISCONTINUITY, )) # Process initial conditions initial_conditions = process( model.concatenated_initial_conditions, "initial_conditions", use_jacobian=False, )[0] init_eval = InitialConditions(initial_conditions, model) # Process rhs, algebraic and event expressions rhs, rhs_eval, jac_rhs = process(model.concatenated_rhs, "RHS") algebraic, algebraic_eval, jac_algebraic = process( model.concatenated_algebraic, "algebraic") terminate_events_eval = [ process(event.expression, "event", use_jacobian=False)[1] for event in model.events if event.event_type == pybamm.EventType.TERMINATION ] # discontinuity events are evaluated before the solver is called, so don't need # to process them discontinuity_events_eval = [ event for event in model.events if event.event_type == pybamm.EventType.DISCONTINUITY ] # Add the solver attributes model.init_eval = init_eval model.rhs_eval = rhs_eval model.algebraic_eval = algebraic_eval model.jac_algebraic_eval = jac_algebraic model.terminate_events_eval = terminate_events_eval model.discontinuity_events_eval = discontinuity_events_eval # Calculate initial conditions model.y0 = init_eval(inputs) # 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_casadi, p_casadi_stacked) model.casadi_rhs = casadi.Function( "rhs", [t_casadi, y_casadi, p_casadi_stacked], [explicit_rhs]) model.casadi_algebraic = 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 elif len(model.algebraic) == 0: # No algebraic equations: residuals is rhs only model.residuals_eval = Residuals(rhs, "residuals", model) model.jacobian_eval = jac_rhs # Calculate consistent initial conditions for the algebraic equations else: all_states = pybamm.NumpyConcatenation( model.concatenated_rhs, model.concatenated_algebraic) # Process again, uses caching so should be quick residuals_eval, jacobian_eval = process(all_states, "residuals")[1:] model.residuals_eval = residuals_eval model.jacobian_eval = jacobian_eval pybamm.logger.info("Finish solver set-up")
def set_up(self, model): """Unpack model, perform checks, simplify and calculate jacobian. Parameters ---------- model : :class:`pybamm.BaseModel` The model whose solution to calculate. Must have attributes rhs and initial_conditions Raises ------ :class:`pybamm.SolverError` If the model contains any algebraic equations (in which case a DAE solver should be used instead) """ # create simplified rhs, algebraic and event expressions concatenated_rhs = model.concatenated_rhs concatenated_algebraic = model.concatenated_algebraic events = model.events if model.use_simplify: # set up simplification object, for re-use of dict simp = pybamm.Simplification() pybamm.logger.info("Simplifying RHS") concatenated_rhs = simp.simplify(concatenated_rhs) pybamm.logger.info("Simplifying algebraic") concatenated_algebraic = simp.simplify(concatenated_algebraic) pybamm.logger.info("Simplifying events") events = { name: simp.simplify(event) for name, event in events.items() } if model.use_jacobian: # Create Jacobian from concatenated rhs and algebraic y = pybamm.StateVector( slice(0, np.size(model.concatenated_initial_conditions))) # set up Jacobian object, for re-use of dict jacobian = pybamm.Jacobian() pybamm.logger.info("Calculating jacobian") jac_rhs = jacobian.jac(concatenated_rhs, y) jac_algebraic = jacobian.jac(concatenated_algebraic, y) jac = pybamm.SparseStack(jac_rhs, jac_algebraic) model.jacobian = jac model.jacobian_rhs = jac_rhs model.jacobian_algebraic = jac_algebraic if model.use_simplify: pybamm.logger.info("Simplifying jacobian") jac_algebraic = simp.simplify(jac_algebraic) jac = simp.simplify(jac) if model.use_to_python: pybamm.logger.info("Converting jacobian to python") jac_algebraic = pybamm.EvaluatorPython(jac_algebraic) jac = pybamm.EvaluatorPython(jac) def jac_alg_fn(t, y): return jac_algebraic.evaluate(t, y) else: jac = None jac_alg_fn = None if model.use_to_python: pybamm.logger.info("Converting RHS to python") concatenated_rhs = pybamm.EvaluatorPython(concatenated_rhs) pybamm.logger.info("Converting algebraic to python") concatenated_algebraic = pybamm.EvaluatorPython( concatenated_algebraic) pybamm.logger.info("Converting events to python") events = { name: pybamm.EvaluatorPython(event) for name, event in events.items() } # Calculate consistent initial conditions for the algebraic equations def rhs(t, y): return concatenated_rhs.evaluate(t, y, known_evals={})[0][:, 0] def algebraic(t, y): return concatenated_algebraic.evaluate(t, y, known_evals={})[0][:, 0] if len(model.algebraic) > 0: y0 = self.calculate_consistent_initial_conditions( rhs, algebraic, model.concatenated_initial_conditions[:, 0], jac_alg_fn) else: # can use DAE solver to solve ODE model y0 = model.concatenated_initial_conditions[:, 0] # Create functions to evaluate residuals def residuals(t, y, ydot): pybamm.logger.debug("Evaluating residuals for {} at t={}".format( model.name, t)) y = y[:, np.newaxis] rhs_eval, known_evals = concatenated_rhs.evaluate(t, y, known_evals={}) # reuse known_evals alg_eval = concatenated_algebraic.evaluate( t, y, known_evals=known_evals)[0] # turn into 1D arrays rhs_eval = rhs_eval[:, 0] alg_eval = alg_eval[:, 0] return (np.concatenate( (rhs_eval, alg_eval)) - model.mass_matrix.entries @ ydot) # Create event-dependent function to evaluate events def event_fun(event): def eval_event(t, y): return event.evaluate(t, y) return eval_event event_funs = [event_fun(event) for event in events.values()] # Create function to evaluate jacobian if jac is not None: def jacobian(t, y): return jac.evaluate(t, y, known_evals={})[0] else: jacobian = None # Add the solver attributes # Note: these are the (possibly) converted to python version rhs, algebraic # etc. The expression tree versions of these are attributes of the model self.y0 = y0 self.rhs = rhs self.algebraic = algebraic self.residuals = residuals self.events = events self.event_funs = event_funs self.jacobian = jacobian