Skip to content

Commit

Permalink
#1477 unit tests pass
Browse files Browse the repository at this point in the history
  • Loading branch information
martinjrobins committed Jul 16, 2021
1 parent d9ff546 commit 6e91335
Show file tree
Hide file tree
Showing 14 changed files with 127 additions and 272 deletions.
1 change: 1 addition & 0 deletions pybamm/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,7 @@ def version(formatted=False):
#
from .solvers.solution import Solution
from .solvers.processed_variable import ProcessedVariable
from .solvers.processed_symbolic_variable import ProcessedSymbolicVariable
from .solvers.base_solver import BaseSolver
from .solvers.dummy_solver import DummySolver
from .solvers.algebraic_solver import AlgebraicSolver
Expand Down
29 changes: 20 additions & 9 deletions pybamm/solvers/base_solver.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
from scipy.sparse import block_diag
import multiprocessing as mp
import warnings
import numbers


class BaseSolver(object):
Expand Down Expand Up @@ -228,6 +227,13 @@ def set_up(self, model, inputs=None, t_eval=None,
if calculate_sensitivites and not isinstance(self, pybamm.IDAKLUSolver):
calculate_sensitivites_explicit = True

if calculate_sensitivites_explicit and model.convert_to_format != 'casadi':
raise NotImplementedError(
"Sensitivities only supported for:\n"
" - model.convert_to_format = 'casadi'\n"
" - IDAKLUSolver (any convert_to_format)"
)

# save sensitivity parameters so we can identify them later on
# (FYI: this is used in the Solution class)
model.calculate_sensitivities = calculate_sensitivites
Expand Down Expand Up @@ -288,8 +294,8 @@ def report(string):
jacp = None
if calculate_sensitivites_explicit:
raise NotImplementedError(
"sensitivities using convert_to_format = 'jax' "
"only implemented for IDAKLUSolver"
"explicit sensitivity equations not supported for "
"convert_to_format='jax'"
)
elif calculate_sensitivites:
report((
Expand All @@ -310,11 +316,12 @@ def report(string):
elif model.convert_to_format != "casadi":
# Process with pybamm functions, optionally converting
# to python evaluator
if calculate_sensitivites:
if calculate_sensitivites_explicit:
raise NotImplementedError(
"sensitivities only implemented with "
"convert_to_format = 'casadi' or convert_to_format = 'jax'"
"explicit sensitivity equations not supported for "
"convert_to_format='{}'".format(model.convert_to_format)
)
elif calculate_sensitivites:
report((
f"Calculating sensitivities for {name} with respect "
f"to parameters {calculate_sensitivites}"
Expand Down Expand Up @@ -367,7 +374,9 @@ def jacp(*args, **kwargs):
# for details
if name == "RHS" and model.len_rhs > 0:
report(
"Creating explicit forward sensitivity equations for rhs using CasADi")
"Creating explicit forward sensitivity equations "
"for rhs using CasADi"
)
df_dx = casadi.jacobian(func, y_diff)
df_dp = casadi.jacobian(func, pS_casadi_stacked)
S_x_mat = S_x.reshape(
Expand All @@ -386,7 +395,8 @@ def jacp(*args, **kwargs):
func = casadi.vertcat(func, S_rhs)
if name == "algebraic" and model.len_alg > 0:
report(
"Creating explicit forward sensitivity equations for algebraic using CasADi"
"Creating explicit forward sensitivity equations "
"for algebraic using CasADi"
)
dg_dz = casadi.jacobian(func, y_alg)
dg_dp = casadi.jacobian(func, pS_casadi_stacked)
Expand Down Expand Up @@ -812,6 +822,7 @@ def solve(
"""
pybamm.logger.info("Start solving {} with {}".format(model.name, self.name))
self.calculate_sensitivites = calculate_sensitivities

# Make sure model isn't empty
if len(model.rhs) == 0 and len(model.algebraic) == 0:
Expand Down Expand Up @@ -1401,7 +1412,7 @@ def _set_up_ext_and_inputs(self, model, external_variables, inputs):
name = input_param.name
if name not in inputs:
# Don't allow symbolic inputs if using `sensitivity`
if self.sensitivity == "explicit forward":
if self.calculate_sensitivites:
raise pybamm.SolverError(
"Cannot have symbolic inputs if explicitly solving forward"
"sensitivity equations"
Expand Down
73 changes: 33 additions & 40 deletions pybamm/solvers/casadi_algebraic_solver.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,9 +32,6 @@ def __init__(self, tol=1e-6, extra_options=None):
self.extra_options = extra_options or {}
pybamm.citations.register("Andersson2019")

self.rootfinders = {}
self.y_sols = {}

@property
def tol(self):
return self._tol
Expand Down Expand Up @@ -102,6 +99,14 @@ def _integrate(self, model, t_eval, inputs_dict=None):

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)

t_and_inputs_sym = casadi.vertcat(t_sym, symbolic_inputs)
alg = model.casadi_algebraic(t_sym, y_sym, inputs)

# Check interpolant extrapolation
if model.interpolant_extrapolation_events_eval:
extrap_event = [
Expand All @@ -116,7 +121,9 @@ def _integrate(self, model, t_eval, inputs_dict=None):
event.event_type
== pybamm.EventType.INTERPOLANT_EXTRAPOLATION
and (
event.expression.evaluate(0, y0.full(), inputs=inputs)
event.expression.evaluate(
0, y0.full(), inputs=inputs_dict
)
< self.extrap_tol
)
):
Expand All @@ -129,40 +136,26 @@ def _integrate(self, model, t_eval, inputs_dict=None):
"outside these bounds.".format(extrap_event_names)
)

if model in self.rootfinders:
roots = self.rootfinders[model]
else:
# Set up
t_sym = casadi.MX.sym("t")
y0_diff_sym = casadi.MX.sym("y0_diff", y0_diff.shape[0])
y_alg_sym = casadi.MX.sym("y_alg", y0_alg.shape[0])
y_sym = casadi.vertcat(y0_diff_sym, y_alg_sym)

t_y0diff_inputs_sym = casadi.vertcat(t_sym, y0_diff_sym, symbolic_inputs)
alg = model.casadi_algebraic(t_sym, y_sym, symbolic_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_y0diff_inputs_sym, g=alg),
{
**self.extra_options,
"abstol": self.tol,
"constraints": list(constraints[len_rhs:]),
},
)

self.rootfinders[model] = roots
# 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_and_inputs_sym, g=alg),
{
**self.extra_options,
"abstol": self.tol,
"constraints": list(constraints[len_rhs:]),
},
)

timer = pybamm.Timer()
integration_time = 0
Expand All @@ -182,11 +175,11 @@ def _integrate(self, model, t_eval, inputs_dict=None):
y_alg = casadi.horzcat(y_alg, y0_alg)
# Otherwise calculate new y_sol
else:
t_y0_diff_inputs = casadi.vertcat(t, y0_diff, symbolic_inputs)
t_eval_inputs_sym = casadi.vertcat(t, symbolic_inputs)
# Solve
try:
timer.reset()
y_alg_sol = roots(y0_alg, t_y0_diff_inputs)
y_alg_sol = roots(y0_alg, t_eval_inputs_sym)
integration_time += timer.time()
success = True
message = None
Expand Down
3 changes: 0 additions & 3 deletions pybamm/solvers/casadi_solver.py
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,6 @@ def _integrate(self, model, t_eval, inputs_dict=None):
Any external variables or input parameters to pass to the model when solving
"""


# are we solving explicit forward equations?
explicit_sensitivities = bool(self.calculate_sensitivites)

Expand Down Expand Up @@ -612,8 +611,6 @@ def _run_integrator(self, model, y0, inputs_dict, inputs, t_eval, use_grid=True)
else:
integrator = self.integrators[model]["no grid"]

symbolic_inputs = casadi.MX.sym("inputs", inputs.shape[0])

len_rhs = model.concatenated_rhs.size

# Check y0 to see if it includes sensitivities
Expand Down
17 changes: 1 addition & 16 deletions pybamm/solvers/idaklu_solver.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,14 +38,6 @@ class IDAKLUSolver(pybamm.BaseSolver):
The tolerance for the initial-condition solver (default is 1e-6).
extrap_tol : float, optional
The tolerance to assert whether extrapolation occurs or not (default is 0).
sensitivity : str, optional
Whether (and how) to calculate sensitivities when solving. Options are:
- "explicit forward": explicitly formulate the sensitivity equations. \
The formulation is as per "Park, S., Kato, D., Gima, Z., \
Klein, R., & Moura, S. (2018). Optimal experimental design for parameterization\
of an electrochemical lithium-ion battery model. Journal of The Electrochemical\
Society, 165(7), A1309.". See #1100 for details \
"""

def __init__(
Expand All @@ -56,17 +48,11 @@ def __init__(
root_tol=1e-6,
extrap_tol=0,
max_steps="deprecated",
sensitivity="idas"
):

if idaklu_spec is None:
raise ImportError("KLU is not installed")

if sensitivity == "explicit forward":
raise NotImplementedError(
"Cannot use explicit forward equations with IDAKLUSolver"
)

super().__init__(
"ida",
rtol,
Expand All @@ -75,7 +61,6 @@ def __init__(
root_tol,
extrap_tol,
max_steps,
sensitivity=sensitivity,
)
self.name = "IDA KLU solver"

Expand Down Expand Up @@ -339,7 +324,7 @@ def sensfn(resvalS, t, y, yp, yS, ypS):
name: sol.yS[i].transpose() for i, name in enumerate(sens0.keys())
}
else:
yS_out = None
yS_out = False
if sol.flag in [0, 2]:
# 0 = solved for all t_eval
if sol.flag == 0:
Expand Down
4 changes: 3 additions & 1 deletion pybamm/solvers/processed_variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -524,7 +524,9 @@ def initialise_sensitivity_explicit_forward(self):
p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()])

# Convert variable to casadi format for differentiating
var_casadi = self.base_variables[0].to_casadi(t_casadi, y_casadi, inputs=p_casadi)
var_casadi = self.base_variables[0].to_casadi(
t_casadi, y_casadi, inputs=p_casadi
)
dvar_dy = casadi.jacobian(var_casadi, y_casadi)
dvar_dp = casadi.jacobian(var_casadi, p_casadi_stacked)

Expand Down
11 changes: 2 additions & 9 deletions pybamm/solvers/scipy_solver.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,12 +25,6 @@ class ScipySolver(pybamm.BaseSolver):
Any options to pass to the solver.
Please consult `SciPy documentation <https://tinyurl.com/yafgqg9y>`_ for
details.
sensitivity : str, optional
Whether (and how) to calculate sensitivities when solving. Options are:
- None: no sensitivities
- "explicit forward": explicitly formulate the sensitivity equations. \
See :class:`pybamm.BaseSolver`
"""

def __init__(
Expand All @@ -40,14 +34,12 @@ def __init__(
atol=1e-6,
extrap_tol=0,
extra_options=None,
sensitivity=None,
):
super().__init__(
method=method,
rtol=rtol,
atol=atol,
extrap_tol=extrap_tol,
sensitivity=sensitivity,
)
self.ode_solver = True
self.extra_options = extra_options or {}
Expand Down Expand Up @@ -136,7 +128,8 @@ def event_fn(t, y):
t_event = None
y_event = np.array(None)
sol = pybamm.Solution(
sol.t, sol.y, model, inputs_dict, t_event, y_event, termination
sol.t, sol.y, model, inputs_dict, t_event, y_event, termination,
sensitivities=bool(self.calculate_sensitivites)
)
sol.integration_time = integration_time
return sol
Expand Down
4 changes: 2 additions & 2 deletions pybamm/solvers/solution.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ def __init__(
self._sensitivities = {}
# if solution consists of explicit sensitivity equations, extract them
if (
sensitivities == True
sensitivities is True
and all_models[0] is not None
and not isinstance(all_ys[0], casadi.Function)
and all_models[0].len_rhs_and_alg != all_ys[0].shape[0]
Expand All @@ -97,7 +97,7 @@ def __init__(
self._all_ys[0], self._sensitivities = \
self._extract_explicit_sensitivities(
all_models[0], all_ys[0], all_ts[0], self.all_inputs[0]
)
)
elif isinstance(sensitivities, dict):
self._sensitivities = sensitivities
else:
Expand Down
4 changes: 2 additions & 2 deletions tests/unit/test_solvers/test_base_solver.py
Original file line number Diff line number Diff line change
Expand Up @@ -163,12 +163,12 @@ def __init__(self):
)
self.convert_to_format = "casadi"
self.bounds = (-np.inf * np.ones(4), np.inf * np.ones(4))
<<<<<<< HEAD
self.len_rhs = 1
self.len_rhs_and_alg = 4
self.interpolant_extrapolation_events_eval = []
>>>>>>> develop

def rhs_eval(self, t, y, inputs):
return y[0:1]

def algebraic_eval(self, t, y, inputs):
return (y[1:] - vec[1:]) ** 2
Expand Down
Loading

0 comments on commit 6e91335

Please sign in to comment.