def test_sigmoid(self):
a = pybamm.Scalar(1)
b = pybamm.StateVector(slice(0, 1))
sigm = pybamm.sigmoid(a, b, 10)
self.assertAlmostEqual(sigm.evaluate(y=np.array([2]))[0, 0], 1)
self.assertEqual(sigm.evaluate(y=np.array([1])), 0.5)
self.assertAlmostEqual(sigm.evaluate(y=np.array([0]))[0, 0], 0)
self.assertEqual(str(sigm), "(1.0 + tanh(10.0 * (y[0:1] - 1.0))) / 2.0")
sigm = pybamm.sigmoid(b, a, 10)
self.assertAlmostEqual(sigm.evaluate(y=np.array([2]))[0, 0], 0)
self.assertEqual(sigm.evaluate(y=np.array([1])), 0.5)
self.assertAlmostEqual(sigm.evaluate(y=np.array([0]))[0, 0], 1)
self.assertEqual(str(sigm), "(1.0 + tanh(10.0 * (1.0 - y[0:1]))) / 2.0")
def __ge__(self, other):
"""return a :class:`EqualHeaviside` object, or a smooth approximation."""
k = pybamm.settings.heaviside_smoothing
# Return exact approximation if that is the setting or the outcome is a constant
# (i.e. no need for smoothing)
if k == "exact" or (is_constant(self) and is_constant(other)):
out = pybamm.EqualHeaviside(other, self)
else:
out = pybamm.sigmoid(other, self, k)
return pybamm.simplify_if_constant(out)
def test_sigmoid(self):
# Test that smooth heaviside is used when the setting is changed
a = pybamm.Symbol("a")
b = pybamm.Symbol("b")
pybamm.settings.heaviside_smoothing = 10
self.assertEqual(str(a < b), str(pybamm.sigmoid(a, b, 10)))
self.assertEqual(str(a <= b), str(pybamm.sigmoid(a, b, 10)))
self.assertEqual(str(a > b), str(pybamm.sigmoid(b, a, 10)))
self.assertEqual(str(a >= b), str(pybamm.sigmoid(b, a, 10)))
# But exact heavisides should still be used if both variables are constant
a = pybamm.Scalar(1)
b = pybamm.Scalar(2)
self.assertEqual(str(a < b), str(pybamm.Scalar(1)))
self.assertEqual(str(a <= b), str(pybamm.Scalar(1)))
self.assertEqual(str(a > b), str(pybamm.Scalar(0)))
self.assertEqual(str(a >= b), str(pybamm.Scalar(0)))
# Change setting back for other tests
pybamm.settings.heaviside_smoothing = "exact"
def set_up(self, model, inputs=None, t_eval=None):
"""Unpack model, perform checks, 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
"""
pybamm.logger.info("Start solver set-up")
# 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":
# 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.verbose(string)
if use_jacobian is None:
use_jacobian = model.use_jacobian
if model.convert_to_format != "casadi":
# Process with pybamm functions
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.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_eval).id:
expr = symbol.left.new_copy() / symbol.right.right.new_copy()
found_t = True
elif symbol.left.id == (pybamm.t * model.timescale_eval).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_eval).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"
)
# Calculate initial conditions
model.y0 = init_eval(inputs)
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.init_eval = init_eval
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
)
# 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")