import casadi
import pybamm
import numpy as np
[docs]
class CasadiAlgebraicSolver(pybamm.BaseSolver):
"""Solve a discretised model which contains only (time independent) algebraic
equations using CasADi's root finding algorithm.
Note: this solver could be extended for quasi-static models, or models in
which the time derivative is manually discretised and results in a (possibly
nonlinear) algebaric system at each time level.
Parameters
----------
tol : float, optional
The tolerance for the solver (default is 1e-6).
extra_options : dict, optional
Any options to pass to the CasADi rootfinder.
Please consult `CasADi documentation <https://web.casadi.org/python-api/#rootfinding>`_ for
details.
"""
def __init__(self, tol=1e-6, extra_options=None):
super().__init__()
self.tol = tol
self.name = "CasADi algebraic solver"
self.algebraic_solver = True
self.extra_options = extra_options or {}
pybamm.citations.register("Andersson2019")
@property
def tol(self):
return self._tol
@tol.setter
def tol(self, value):
self._tol = value
def _integrate(self, model, t_eval, inputs_dict=None):
"""
Calculate the solution of the algebraic equations through root-finding
Parameters
----------
model : :class:`pybamm.BaseModel`
The model whose solution to calculate.
t_eval : :class:`numpy.array`, size (k,)
The times at which to compute the solution
inputs_dict : dict, optional
Any input parameters to pass to the model when solving.
"""
# Record whether there are any symbolic inputs
inputs_dict = inputs_dict or {}
# Create casadi objects for the root-finder
inputs = casadi.vertcat(*[v for v in inputs_dict.values()])
y0 = model.y0
# 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
# equations will be equal to the initial condition provided. This allows this
# solver to be used for initialising the DAE solvers
if model.rhs == {}:
len_rhs = 0
y0_diff = casadi.DM()
y0_alg = y0
else:
# Check y0 to see if it includes sensitivities
if model.len_rhs_and_alg == y0.shape[0]:
len_rhs = model.len_rhs
else:
len_rhs = model.len_rhs + model.len_rhs_sens
y0_diff = y0[:len_rhs]
y0_alg = y0[len_rhs:]
y_alg = None
# Set up
t_sym = casadi.MX.sym("t")
y_alg_sym = casadi.MX.sym("y_alg", y0_alg.shape[0])
y_sym = casadi.vertcat(y0_diff, y_alg_sym)
alg = model.casadi_algebraic(t_sym, y_sym, inputs)
# Set constraints vector in the casadi format
# Constrain the unknowns. 0 (default): no constraint on ui, 1: ui >= 0.0,
# -1: ui <= 0.0, 2: ui > 0.0, -2: ui < 0.0.
constraints = np.zeros_like(model.bounds[0], dtype=int)
# If the lower bound is positive then the variable must always be positive
constraints[model.bounds[0] >= 0] = 1
# If the upper bound is negative then the variable must always be negative
constraints[model.bounds[1] <= 0] = -1
# Set up rootfinder
roots = casadi.rootfinder(
"roots",
"newton",
dict(x=y_alg_sym, p=t_sym, g=alg),
{
**self.extra_options,
"abstol": self.tol,
"constraints": list(constraints[len_rhs:]),
},
)
timer = pybamm.Timer()
integration_time = 0
for _, t in enumerate(t_eval):
# Solve
try:
timer.reset()
y_alg_sol = roots(y0_alg, t)
integration_time += timer.time()
success = True
message = None
# Check final output
y_sol = casadi.vertcat(y0_diff, y_alg_sol)
fun = model.casadi_algebraic(t, y_sol, inputs)
except RuntimeError as err:
success = False
message = err.args[0]
fun = None
# If there are no symbolic inputs, check the function is below the tol
# Skip this check if there are symbolic inputs
if success and (
not any(np.isnan(fun)) and np.all(casadi.fabs(fun) < self.tol)
):
# update initial guess for the next iteration
y0_alg = y_alg_sol
y0 = casadi.vertcat(y0_diff, y0_alg)
# update solution array
if y_alg is None:
y_alg = y_alg_sol
else:
y_alg = casadi.horzcat(y_alg, y_alg_sol)
elif not success:
raise pybamm.SolverError(
f"Could not find acceptable solution: {message}"
)
elif any(np.isnan(fun)):
raise pybamm.SolverError(
"Could not find acceptable solution: solver returned NaNs"
)
else:
raise pybamm.SolverError(
f"""
Could not find acceptable solution: solver terminated
successfully, but maximum solution error ({casadi.mmax(casadi.fabs(fun))})
above tolerance ({self.tol})
"""
)
# Concatenate differential part
y_diff = casadi.horzcat(*[y0_diff] * len(t_eval))
y_sol = casadi.vertcat(y_diff, y_alg)
# Return solution object (no events, so pass None to t_event, y_event)
try:
explicit_sensitivities = bool(model.calculate_sensitivities)
except AttributeError:
explicit_sensitivities = False
sol = pybamm.Solution(
[t_eval],
y_sol,
model,
inputs_dict,
termination="final time",
sensitivities=explicit_sensitivities,
)
sol.integration_time = integration_time
return sol
