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
