def __init__(self, mesh=None, spatial_methods=None):
self._mesh = mesh
if mesh is None:
self._spatial_methods = {}
else:
# Unpack macroscale to the constituent subdomains
if "macroscale" in spatial_methods.keys():
method = spatial_methods["macroscale"]
spatial_methods["negative electrode"] = method
spatial_methods["separator"] = method
spatial_methods["positive electrode"] = method
self._spatial_methods = spatial_methods
for domain, method in self._spatial_methods.items():
method.build(mesh)
# Check zero-dimensional methods are only applied to zero-dimensional
# meshes
if isinstance(method, pybamm.ZeroDimensionalSpatialMethod):
if not isinstance(mesh[domain], pybamm.SubMesh0D):
raise pybamm.DiscretisationError(
"Zero-dimensional spatial method for the "
"{} domain requires a zero-dimensional submesh".
format(domain))
self.bcs = {}
self.y_slices = {}
self._discretised_symbols = {}
self.external_variables = {}
def check_discretised_or_discretise_inplace_if_0D(self):
"""
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 self.is_discretised is False:
try:
disc = pybamm.Discretisation()
disc.process_model(self)
except pybamm.DiscretisationError as e:
raise pybamm.DiscretisationError(
"Cannot automatically discretise model, model should be "
"discretised before exporting casadi functions ({})".
format(e))
def __init__(self, geometry, submesh_types, var_pts):
super().__init__()
# convert var_pts to an id dict
var_id_pts = {var.id: pts for var, pts in var_pts.items()}
# create submesh_pts from var_pts
submesh_pts = {}
for domain in geometry:
# create mesh generator if just class is passed (will throw an error
# later if the mesh needed parameters)
if not isinstance(
submesh_types[domain], pybamm.MeshGenerator
) and issubclass(submesh_types[domain], pybamm.SubMesh):
submesh_types[domain] = pybamm.MeshGenerator(submesh_types[domain])
# Zero dimensional submesh case (only one point)
if issubclass(submesh_types[domain].submesh_type, pybamm.SubMesh0D):
submesh_pts[domain] = 1
# other cases
else:
submesh_pts[domain] = {}
if len(list(geometry[domain].keys())) > 3:
raise pybamm.GeometryError("Too many keys provided")
for var in list(geometry[domain].keys()):
if var in ["primary", "secondary"]:
raise pybamm.GeometryError(
"Geometry should no longer be given keys 'primary' or "
"'secondary'. See pybamm.battery_geometry() for example"
)
# skip over tabs key
if var != "tabs":
# Raise error if the number of points for a particular
# variable haven't been provided, unless that variable
# doesn't appear in the geometry
if (
var.id not in var_id_pts.keys()
and var.domain[0] in geometry.keys()
):
raise KeyError(
"Points not given for a variable in domain {}".format(
domain
)
)
# Otherwise add to the dictionary of submesh points
submesh_pts[domain][var.id] = var_id_pts[var.id]
self.submesh_pts = submesh_pts
# Input domain order manually
self.domain_order = []
# First the macroscale domains, whose order we care about
for domain in ["negative electrode", "separator", "positive electrode"]:
if domain in geometry:
self.domain_order.append(domain)
# Then the remaining domains
for domain in geometry:
if domain not in ["negative electrode", "separator", "positive electrode"]:
self.domain_order.append(domain)
# evaluate any expressions in geometry
for domain in geometry:
for spatial_variable, spatial_limits in geometry[domain].items():
# process tab information if using 1 or 2D current collectors
if spatial_variable == "tabs":
for tab, position_size in spatial_limits.items():
for position_size, sym in position_size.items():
if isinstance(sym, pybamm.Symbol):
sym_eval = sym.evaluate()
geometry[domain]["tabs"][tab][position_size] = sym_eval
else:
for lim, sym in spatial_limits.items():
if isinstance(sym, pybamm.Symbol):
try:
sym_eval = sym.evaluate()
except NotImplementedError as error:
if sym.has_symbol_of_classes(pybamm.Parameter):
raise pybamm.DiscretisationError(
"Parameter values have not yet been set for "
"geometry. Make sure that something like "
"`param.process_geometry(geometry)` has been "
"run."
)
else:
raise error
elif isinstance(sym, numbers.Number):
sym_eval = sym
geometry[domain][spatial_variable][lim] = sym_eval
# Create submeshes
for domain in geometry:
self[domain] = submesh_types[domain](geometry[domain], submesh_pts[domain])
# add ghost meshes
self.add_ghost_meshes()
def process_model(self, model, inplace=True, check_model=True):
"""Discretise a model.
Currently inplace, could be changed to return a new model.
Parameters
----------
model : :class:`pybamm.BaseModel`
Model to dicretise. Must have attributes rhs, initial_conditions and
boundary_conditions (all dicts of {variable: equation})
inplace : bool, optional
If True, discretise the model in place. Otherwise, return a new
discretised model. Default is True.
check_model : bool, optional
If True, model checks are performed after discretisation. For large
systems these checks can be slow, so can be skipped by setting this
option to False. When developing, testing or debugging it is recommened
to leave this option as True as it may help to identify any errors.
Default is True.
Returns
-------
model_disc : :class:`pybamm.BaseModel`
The discretised model. Note that if ``inplace`` is True, model will
have also been discretised in place so model == model_disc. If
``inplace`` is False, model != model_disc
Raises
------
:class:`pybamm.ModelError`
If an empty model is passed (`model.rhs = {}` and `model.algebraic = {}` and
`model.variables = {}`)
"""
if model.is_discretised is True:
raise pybamm.ModelError(
"Cannot re-discretise a model. "
"Set 'inplace=False' when first discretising a model to then be able "
"to discretise it more times (e.g. for convergence studies).")
pybamm.logger.info("Start discretising {}".format(model.name))
# Make sure model isn't empty
if (len(model.rhs) == 0 and len(model.algebraic) == 0
and len(model.variables) == 0):
raise pybamm.ModelError("Cannot discretise empty model")
# Check well-posedness to avoid obscure errors
model.check_well_posedness()
# Prepare discretisation
# set variables (we require the full variable not just id)
variables = list(model.rhs.keys()) + list(model.algebraic.keys())
if self.spatial_methods == {} and any(var.domain != []
for var in variables):
for var in variables:
if var.domain != []:
raise pybamm.DiscretisationError(
"Spatial method has not been given "
"for variable {} with domain {}".format(
var.name, var.domain))
# Set the y split for variables
pybamm.logger.info("Set variable slices for {}".format(model.name))
self.set_variable_slices(variables)
# Keep a record of y_slices in the model
model.y_slices = self.y_slices_explicit
# now add extrapolated external variables to the boundary conditions
# if required by the spatial method
self._preprocess_external_variables(model)
self.set_external_variables(model)
# set boundary conditions (only need key ids for boundary_conditions)
pybamm.logger.info("Discretise boundary conditions for {}".format(
model.name))
self.bcs = self.process_boundary_conditions(model)
pybamm.logger.info("Set internal boundary conditions for {}".format(
model.name))
self.set_internal_boundary_conditions(model)
# set up inplace vs not inplace
if inplace:
# any changes to model_disc attributes will change model attributes
# since they point to the same object
model_disc = model
else:
# create an empty copy of the original model
model_disc = model.new_copy()
model_disc.bcs = self.bcs
pybamm.logger.info("Discretise initial conditions for {}".format(
model.name))
ics, concat_ics = self.process_initial_conditions(model)
model_disc.initial_conditions = ics
model_disc.concatenated_initial_conditions = concat_ics
# Discretise variables (applying boundary conditions)
# Note that we **do not** discretise the keys of model.rhs,
# model.initial_conditions and model.boundary_conditions
pybamm.logger.info("Discretise variables for {}".format(model.name))
model_disc.variables = self.process_dict(model.variables)
# Process parabolic and elliptic equations
pybamm.logger.info("Discretise model equations for {}".format(
model.name))
rhs, concat_rhs, alg, concat_alg = self.process_rhs_and_algebraic(
model)
model_disc.rhs, model_disc.concatenated_rhs = rhs, concat_rhs
model_disc.algebraic, model_disc.concatenated_algebraic = alg, concat_alg
# Process events
processed_events = []
pybamm.logger.info("Discretise events for {}".format(model.name))
for event in model.events:
pybamm.logger.debug("Discretise event '{}'".format(event.name))
processed_event = pybamm.Event(
event.name, self.process_symbol(event.expression),
event.event_type)
processed_events.append(processed_event)
model_disc.events = processed_events
# Create mass matrix
pybamm.logger.info("Create mass matrix for {}".format(model.name))
model_disc.mass_matrix, model_disc.mass_matrix_inv = self.create_mass_matrix(
model_disc)
# Check that resulting model makes sense
if check_model:
pybamm.logger.info("Performing model checks for {}".format(
model.name))
self.check_model(model_disc)
pybamm.logger.info("Finish discretising {}".format(model.name))
# Record that the model has been discretised
model_disc.is_discretised = True
return model_disc
def export_casadi_objects(self, variable_names, input_parameter_order=None):
"""
Export the constituent parts of the model (rhs, algebraic, initial conditions,
etc) as casadi objects.
Parameters
----------
variable_names : list
Variables to be exported alongside the model structure
input_parameter_order : list, optional
Order in which the input parameters should be stacked. If None, the order
returned by :meth:`BaseModel.input_parameters` is used
Returns
-------
casadi_dict : dict
Dictionary of {str: casadi object} pairs representing the model in casadi
format
"""
# 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 self.is_discretised is False:
try:
disc = pybamm.Discretisation()
disc.process_model(self)
except pybamm.DiscretisationError as e:
raise pybamm.DiscretisationError(
"Cannot automatically discretise model, model should be "
"discretised before exporting casadi functions ({})".format(e)
)
# Create casadi functions for the model
t_casadi = casadi.MX.sym("t")
y_diff = casadi.MX.sym("y_diff", self.concatenated_rhs.size)
y_alg = casadi.MX.sym("y_alg", self.concatenated_algebraic.size)
y_casadi = casadi.vertcat(y_diff, y_alg)
# Read inputs
inputs_wrong_order = {}
for input_param in self.input_parameters:
name = input_param.name
inputs_wrong_order[name] = casadi.MX.sym(name, input_param._expected_size)
# Read external variables
external_casadi = {}
for external_varaiable in self.external_variables:
name = external_varaiable.name
ev_size = external_varaiable._evaluate_for_shape().shape[0]
external_casadi[name] = casadi.MX.sym(name, ev_size)
# Sort according to input_parameter_order
if input_parameter_order is None:
inputs = inputs_wrong_order
else:
inputs = {name: inputs_wrong_order[name] for name in input_parameter_order}
# Set up external variables and inputs
# Put external variables first like the integrator expects
ext_and_in = {**external_casadi, **inputs}
inputs_stacked = casadi.vertcat(*[p for p in ext_and_in.values()])
# Convert initial conditions to casadi form
y0 = self.concatenated_initial_conditions.to_casadi(
t_casadi, y_casadi, inputs=inputs
)
x0 = y0[: self.concatenated_rhs.size]
z0 = y0[self.concatenated_rhs.size :]
# Convert rhs and algebraic to casadi form and calculate jacobians
rhs = self.concatenated_rhs.to_casadi(t_casadi, y_casadi, inputs=ext_and_in)
jac_rhs = casadi.jacobian(rhs, y_casadi)
algebraic = self.concatenated_algebraic.to_casadi(
t_casadi, y_casadi, inputs=inputs
)
jac_algebraic = casadi.jacobian(algebraic, y_casadi)
# For specified variables, convert to casadi
variables = OrderedDict()
for name in variable_names:
var = self.variables[name]
variables[name] = var.to_casadi(t_casadi, y_casadi, inputs=ext_and_in)
casadi_dict = {
"t": t_casadi,
"x": y_diff,
"z": y_alg,
"inputs": inputs_stacked,
"rhs": rhs,
"algebraic": algebraic,
"jac_rhs": jac_rhs,
"jac_algebraic": jac_algebraic,
"variables": variables,
"x0": x0,
"z0": z0,
}
return casadi_dict
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")