def test_domain_concatenation_2D(self):
disc = get_1p1d_discretisation_for_testing()
a_dom = ["negative electrode"]
b_dom = ["separator"]
a = pybamm.Variable("a", domain=a_dom)
b = pybamm.Variable("b", domain=b_dom)
conc = pybamm.concatenation(2 * a, 3 * b)
disc.set_variable_slices([a, b])
expr = disc.process_symbol(conc)
self.assertIsInstance(expr, pybamm.DomainConcatenation)
y = np.empty((expr._size, 1))
for i in range(len(y)):
y[i] = i
constant_symbols = OrderedDict()
variable_symbols = OrderedDict()
pybamm.find_symbols(expr, constant_symbols, variable_symbols)
self.assertEqual(len(constant_symbols), 0)
evaluator = pybamm.EvaluatorPython(expr)
result = evaluator.evaluate(y=y)
np.testing.assert_allclose(result, expr.evaluate(y=y))
# check that concatenating a single domain is consistent
expr = disc.process_symbol(pybamm.Concatenation(a))
evaluator = pybamm.EvaluatorPython(expr)
result = evaluator.evaluate(y=y)
np.testing.assert_allclose(result, expr.evaluate(y=y))
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:
y = pybamm.StateVector(
slice(0, np.size(model.concatenated_initial_conditions)))
pybamm.logger.info("Calculating jacobian")
jac = concatenated_algebraic.jac(y)
if model.use_simplify:
pybamm.logger.info("Simplifying jacobian")
jac = jac.simplify()
if model.use_to_python:
pybamm.logger.info("Converting jacobian to python")
jac = pybamm.EvaluatorPython(jac)
else:
jac = None
if model.use_to_python:
pybamm.logger.info("Converting algebraic to python")
concatenated_algebraic = pybamm.EvaluatorPython(
concatenated_algebraic)
return concatenated_algebraic, jac
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
def evaluate_model(self,
simplify=False,
use_known_evals=False,
to_python=False):
result = np.empty((0, 1))
for eqn in [
self.model.concatenated_rhs, self.model.concatenated_algebraic
]:
if simplify:
eqn = eqn.simplify()
y = self.model.concatenated_initial_conditions.evaluate(t=0)
if use_known_evals:
eqn_eval, known_evals = eqn.evaluate(0, y, known_evals={})
elif to_python:
evaluator = pybamm.EvaluatorPython(eqn)
eqn_eval = evaluator.evaluate(0, y)
else:
eqn_eval = eqn.evaluate(0, y)
if eqn_eval.shape == (0, ):
eqn_eval = eqn_eval[:, np.newaxis]
result = np.concatenate([result, eqn_eval])
return result
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 simplified rhs
y = pybamm.StateVector(
slice(0, np.size(model.concatenated_initial_conditions)))
pybamm.logger.info("Calculating jacobian")
jac_rhs = concatenated_rhs.jac(y)
jac_algebraic = concatenated_algebraic.jac(y)
jac = pybamm.SparseStack(jac_rhs, jac_algebraic)
model.jacobian = jac
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
self.y0 = y0
self.rhs = rhs
self.algebraic = algebraic
self.residuals = residuals
self.events = events
self.event_funs = event_funs
self.jacobian = jacobian
def test_evaluator_python(self):
a = pybamm.StateVector(slice(0, 1))
b = pybamm.StateVector(slice(1, 2))
y_tests = [np.array([[2], [3]]), np.array([[1], [3]])]
t_tests = [1, 2]
# test a * b
expr = a * b
evaluator = pybamm.EvaluatorPython(expr)
result = evaluator.evaluate(t=None, y=np.array([[2], [3]]))
self.assertEqual(result, 6)
result = evaluator.evaluate(t=None, y=np.array([[1], [3]]))
self.assertEqual(result, 3)
# test function(a*b)
expr = pybamm.Function(test_function, a * b)
evaluator = pybamm.EvaluatorPython(expr)
result = evaluator.evaluate(t=None, y=np.array([[2], [3]]))
self.assertEqual(result, 12)
# test a constant expression
expr = pybamm.Scalar(2) * pybamm.Scalar(3)
evaluator = pybamm.EvaluatorPython(expr)
result = evaluator.evaluate()
self.assertEqual(result, 6)
# test a larger expression
expr = a * b + b + a**2 / b + 2 * a + b / 2 + 4
evaluator = pybamm.EvaluatorPython(expr)
for y in y_tests:
result = evaluator.evaluate(t=None, y=y)
self.assertEqual(result, expr.evaluate(t=None, y=y))
# test something with time
expr = a * pybamm.t
evaluator = pybamm.EvaluatorPython(expr)
for t, y in zip(t_tests, y_tests):
result = evaluator.evaluate(t=t, y=y)
self.assertEqual(result, expr.evaluate(t=t, y=y))
# test something with a matrix multiplication
A = pybamm.Matrix(np.array([[1, 2], [3, 4]]))
expr = A @ pybamm.StateVector(slice(0, 2))
evaluator = pybamm.EvaluatorPython(expr)
for t, y in zip(t_tests, y_tests):
result = evaluator.evaluate(t=t, y=y)
np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
# test something with a heaviside
a = pybamm.Vector(np.array([1, 2]))
expr = a <= pybamm.StateVector(slice(0, 2))
evaluator = pybamm.EvaluatorPython(expr)
for t, y in zip(t_tests, y_tests):
result = evaluator.evaluate(t=t, y=y)
np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
expr = a > pybamm.StateVector(slice(0, 2))
evaluator = pybamm.EvaluatorPython(expr)
for t, y in zip(t_tests, y_tests):
result = evaluator.evaluate(t=t, y=y)
np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
# test something with a minimum or maximum
a = pybamm.Vector(np.array([1, 2]))
expr = pybamm.minimum(a, pybamm.StateVector(slice(0, 2)))
evaluator = pybamm.EvaluatorPython(expr)
for t, y in zip(t_tests, y_tests):
result = evaluator.evaluate(t=t, y=y)
np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
expr = pybamm.maximum(a, pybamm.StateVector(slice(0, 2)))
evaluator = pybamm.EvaluatorPython(expr)
for t, y in zip(t_tests, y_tests):
result = evaluator.evaluate(t=t, y=y)
np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
# test something with an index
expr = pybamm.Index(A @ pybamm.StateVector(slice(0, 2)), 0)
evaluator = pybamm.EvaluatorPython(expr)
for t, y in zip(t_tests, y_tests):
result = evaluator.evaluate(t=t, y=y)
self.assertEqual(result, expr.evaluate(t=t, y=y))
# test something with a sparse matrix multiplication
A = pybamm.Matrix(np.array([[1, 2], [3, 4]]))
B = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1, 0], [0, 4]])))
C = pybamm.Matrix(scipy.sparse.coo_matrix(np.array([[1, 0], [0, 4]])))
expr = A @ B @ C @ pybamm.StateVector(slice(0, 2))
evaluator = pybamm.EvaluatorPython(expr)
for t, y in zip(t_tests, y_tests):
result = evaluator.evaluate(t=t, y=y)
np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
# test numpy concatenation
a = pybamm.Vector(np.array([[1], [2]]))
b = pybamm.Vector(np.array([[3]]))
expr = pybamm.NumpyConcatenation(a, b)
evaluator = pybamm.EvaluatorPython(expr)
for t, y in zip(t_tests, y_tests):
result = evaluator.evaluate(t=t, y=y)
np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
# test sparse stack
A = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1, 0], [0, 4]])))
B = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[2, 0], [5, 0]])))
expr = pybamm.SparseStack(A, B)
evaluator = pybamm.EvaluatorPython(expr)
for t, y in zip(t_tests, y_tests):
result = evaluator.evaluate(t=t, y=y).toarray()
np.testing.assert_allclose(result,
expr.evaluate(t=t, y=y).toarray())
# test Inner
v = pybamm.Vector(np.ones(5), domain="test")
w = pybamm.Vector(2 * np.ones(5), domain="test")
expr = pybamm.Inner(v, w)
evaluator = pybamm.EvaluatorPython(expr)
result = evaluator.evaluate()
np.testing.assert_allclose(result, expr.evaluate())
def test_domain_concatenation(self):
disc = get_discretisation_for_testing()
mesh = disc.mesh
a_dom = ["negative electrode"]
b_dom = ["positive electrode"]
a_pts = mesh[a_dom[0]].npts
b_pts = mesh[b_dom[0]].npts
a = pybamm.StateVector(slice(0, a_pts), domain=a_dom)
b = pybamm.StateVector(slice(a_pts, a_pts + b_pts), domain=b_dom)
y = np.empty((a_pts + b_pts, 1))
for i in range(len(y)):
y[i] = i
# concatenate them the "wrong" way round to check they get reordered correctly
expr = pybamm.DomainConcatenation([b, a], mesh)
constant_symbols = OrderedDict()
variable_symbols = OrderedDict()
pybamm.find_symbols(expr, constant_symbols, variable_symbols)
self.assertEqual(list(variable_symbols.keys())[0], b.id)
self.assertEqual(list(variable_symbols.keys())[1], a.id)
self.assertEqual(list(variable_symbols.keys())[2], expr.id)
var_a = pybamm.id_to_python_variable(a.id)
var_b = pybamm.id_to_python_variable(b.id)
self.assertEqual(len(constant_symbols), 0)
self.assertEqual(
list(variable_symbols.values())[2],
"np.concatenate(({}[0:{}],{}[0:{}]))".format(
var_a, a_pts, var_b, b_pts),
)
evaluator = pybamm.EvaluatorPython(expr)
result = evaluator.evaluate(y=y)
np.testing.assert_allclose(result, expr.evaluate(y=y))
# check that concatenating a single domain is consistent
expr = pybamm.DomainConcatenation([a], mesh)
evaluator = pybamm.EvaluatorPython(expr)
result = evaluator.evaluate(y=y)
np.testing.assert_allclose(result, expr.evaluate(y=y))
# check the reordering in case a child vector has to be split up
a_dom = ["separator"]
b_dom = ["negative electrode", "positive electrode"]
b0_pts = mesh[b_dom[0]].npts
a0_pts = mesh[a_dom[0]].npts
b1_pts = mesh[b_dom[1]].npts
a = pybamm.StateVector(slice(0, a0_pts), domain=a_dom)
b = pybamm.StateVector(slice(a0_pts, a0_pts + b0_pts + b1_pts),
domain=b_dom)
y = np.empty((a0_pts + b0_pts + b1_pts, 1))
for i in range(len(y)):
y[i] = i
var_a = pybamm.id_to_python_variable(a.id)
var_b = pybamm.id_to_python_variable(b.id)
expr = pybamm.DomainConcatenation([a, b], mesh)
constant_symbols = OrderedDict()
variable_symbols = OrderedDict()
pybamm.find_symbols(expr, constant_symbols, variable_symbols)
b0_str = "{}[0:{}]".format(var_b, b0_pts)
a0_str = "{}[0:{}]".format(var_a, a0_pts)
b1_str = "{}[{}:{}]".format(var_b, b0_pts, b0_pts + b1_pts)
self.assertEqual(len(constant_symbols), 0)
self.assertEqual(
list(variable_symbols.values())[2],
"np.concatenate(({},{},{}))".format(b0_str, a0_str, b1_str),
)
evaluator = pybamm.EvaluatorPython(expr)
result = evaluator.evaluate(y=y)
np.testing.assert_allclose(result, expr.evaluate(y=y))
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
