def test_numpy_concatenation_vector_scalar(self):
# with entries
y = np.linspace(0, 1, 10)[:, np.newaxis]
a = pybamm.Vector(y)
b = pybamm.Scalar(16)
c = pybamm.Scalar(3)
conc = pybamm.NumpyConcatenation(a, b, c)
np.testing.assert_array_equal(
conc.evaluate(y=y),
np.concatenate([y, np.array([[16]]),
np.array([[3]])]))
# with y_slice
a = pybamm.StateVector(slice(0, 10))
conc = pybamm.NumpyConcatenation(a, b, c)
np.testing.assert_array_equal(
conc.evaluate(y=y),
np.concatenate([y, np.array([[16]]),
np.array([[3]])]))
# with time
b = pybamm.t
conc = pybamm.NumpyConcatenation(a, b, c)
np.testing.assert_array_equal(
conc.evaluate(16, y),
np.concatenate([y, np.array([[16]]),
np.array([[3]])]))
def test_jac_of_numpy_concatenation(self):
u = pybamm.StateVector(slice(0, 2))
y0 = np.ones(2)
# Multiple children
func = pybamm.NumpyConcatenation(u, u)
jacobian = np.array([[1, 0], [0, 1], [1, 0], [0, 1]])
dfunc_dy = func.jac(u).evaluate(y=y0)
np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
# One child
self.assertEqual(u.jac(u).id, pybamm.NumpyConcatenation(u).jac(u).id)
def test_numpy_concatenation_vectors(self):
# with entries
y = np.linspace(0, 1, 15)[:, np.newaxis]
a = pybamm.Vector(y[:5])
b = pybamm.Vector(y[5:9])
c = pybamm.Vector(y[9:])
conc = pybamm.NumpyConcatenation(a, b, c)
np.testing.assert_array_equal(conc.evaluate(None, y), y)
# with y_slice
a = pybamm.StateVector(slice(0, 10))
b = pybamm.StateVector(slice(10, 15))
c = pybamm.StateVector(slice(15, 23))
conc = pybamm.NumpyConcatenation(a, b, c)
y = np.linspace(0, 1, 23)[:, np.newaxis]
np.testing.assert_array_equal(conc.evaluate(None, y), y)
def test_concatenations(self):
y = np.linspace(0, 1, 10)[:, np.newaxis]
a = pybamm.Vector(y)
b = pybamm.Scalar(16)
c = pybamm.Scalar(3)
conc = pybamm.NumpyConcatenation(a, b, c)
self.assert_casadi_equal(conc.to_casadi(),
casadi.MX(conc.evaluate()),
evalf=True)
# Domain concatenation
mesh = get_mesh_for_testing()
a_dom = ["negative electrode"]
b_dom = ["separator"]
a = 2 * pybamm.Vector(np.ones_like(mesh[a_dom[0]].nodes), domain=a_dom)
b = pybamm.Vector(np.ones_like(mesh[b_dom[0]].nodes), domain=b_dom)
conc = pybamm.DomainConcatenation([b, a], mesh)
self.assert_casadi_equal(conc.to_casadi(),
casadi.MX(conc.evaluate()),
evalf=True)
# 2d
disc = get_1p1d_discretisation_for_testing()
a = pybamm.Variable("a", domain=a_dom)
b = pybamm.Variable("b", domain=b_dom)
conc = pybamm.Concatenation(a, b)
disc.set_variable_slices([conc])
expr = disc.process_symbol(conc)
y = casadi.SX.sym("y", expr.size)
x = expr.to_casadi(None, y)
f = casadi.Function("f", [x], [x])
y_eval = np.linspace(0, 1, expr.size)
self.assert_casadi_equal(f(y_eval), casadi.SX(expr.evaluate(y=y_eval)))
def test_numpy_concatenation(self):
a = pybamm.Variable("a")
b = pybamm.Variable("b")
c = pybamm.Variable("c")
self.assertEqual(
pybamm.numpy_concatenation(pybamm.numpy_concatenation(a, b), c).id,
pybamm.NumpyConcatenation(a, b, c).id,
)
def test_symbol_new_copy(self):
a = pybamm.Parameter("a")
b = pybamm.Parameter("b")
v_n = pybamm.Variable("v", "negative electrode")
x_n = pybamm.standard_spatial_vars.x_n
v_s = pybamm.Variable("v", "separator")
vec = pybamm.Vector([1, 2, 3, 4, 5])
mat = pybamm.Matrix([[1, 2], [3, 4]])
mesh = get_mesh_for_testing()
for symbol in [
a + b,
a - b,
a * b,
a / b,
a**b,
-a,
abs(a),
pybamm.Function(np.sin, a),
pybamm.FunctionParameter("function", {"a": a}),
pybamm.grad(v_n),
pybamm.div(pybamm.grad(v_n)),
pybamm.upwind(v_n),
pybamm.IndefiniteIntegral(v_n, x_n),
pybamm.BackwardIndefiniteIntegral(v_n, x_n),
pybamm.BoundaryValue(v_n, "right"),
pybamm.BoundaryGradient(v_n, "right"),
pybamm.PrimaryBroadcast(a, "domain"),
pybamm.SecondaryBroadcast(v_n, "current collector"),
pybamm.FullBroadcast(a, "domain",
{"secondary": "other domain"}),
pybamm.concatenation(v_n, v_s),
pybamm.NumpyConcatenation(a, b, v_s),
pybamm.DomainConcatenation([v_n, v_s], mesh),
pybamm.Parameter("param"),
pybamm.InputParameter("param"),
pybamm.StateVector(slice(0, 56)),
pybamm.Matrix(np.ones((50, 40))),
pybamm.SpatialVariable("x", ["negative electrode"]),
pybamm.t,
pybamm.Index(vec, 1),
pybamm.NotConstant(a),
pybamm.ExternalVariable(
"external variable",
20,
domain="test",
auxiliary_domains={"secondary": "test2"},
),
pybamm.minimum(a, b),
pybamm.maximum(a, b),
pybamm.SparseStack(mat, mat),
]:
self.assertEqual(symbol.id, symbol.new_copy().id)
def test_symbol_new_copy(self):
a = pybamm.Scalar(0)
b = pybamm.Scalar(1)
v_n = pybamm.Variable("v", "negative electrode")
x_n = pybamm.standard_spatial_vars.x_n
v_s = pybamm.Variable("v", "separator")
vec = pybamm.Vector(np.array([1, 2, 3, 4, 5]))
mesh = get_mesh_for_testing()
for symbol in [
a + b,
a - b,
a * b,
a / b,
a**b,
-a,
abs(a),
pybamm.Function(np.sin, a),
pybamm.FunctionParameter("function", {"a": a}),
pybamm.grad(v_n),
pybamm.div(pybamm.grad(v_n)),
pybamm.Integral(a, pybamm.t),
pybamm.IndefiniteIntegral(v_n, x_n),
pybamm.BackwardIndefiniteIntegral(v_n, x_n),
pybamm.BoundaryValue(v_n, "right"),
pybamm.BoundaryGradient(v_n, "right"),
pybamm.PrimaryBroadcast(a, "domain"),
pybamm.SecondaryBroadcast(v_n, "current collector"),
pybamm.FullBroadcast(a, "domain",
{"secondary": "other domain"}),
pybamm.Concatenation(v_n, v_s),
pybamm.NumpyConcatenation(a, b, v_s),
pybamm.DomainConcatenation([v_n, v_s], mesh),
pybamm.Parameter("param"),
pybamm.InputParameter("param"),
pybamm.StateVector(slice(0, 56)),
pybamm.Matrix(np.ones((50, 40))),
pybamm.SpatialVariable("x", ["negative electrode"]),
pybamm.t,
pybamm.Index(vec, 1),
]:
self.assertEqual(symbol.id, symbol.new_copy().id)
def test_numpy_concatenation_simplify(self):
a = pybamm.Variable("a")
b = pybamm.Variable("b")
c = pybamm.Variable("c")
# simplifying flattens the concatenations into a single concatenation
self.assertEqual(
pybamm.NumpyConcatenation(pybamm.NumpyConcatenation(a, b),
c).simplify().id,
pybamm.NumpyConcatenation(a, b, c).id,
)
self.assertEqual(
pybamm.NumpyConcatenation(a, pybamm.NumpyConcatenation(
b, c)).simplify().id,
pybamm.NumpyConcatenation(a, b, c).id,
)
def test_model_solver_manually_update_initial_conditions(self):
# Create model
model = pybamm.BaseModel()
var1 = pybamm.Variable("var1")
model.rhs = {var1: -var1}
model.initial_conditions = {var1: 1}
# Solve
solver = pybamm.ScipySolver(rtol=1e-8, atol=1e-8)
t_eval = np.linspace(0, 5, 100)
solution = solver.solve(model, t_eval)
np.testing.assert_array_almost_equal(solution.y[0],
1 * np.exp(-solution.t),
decimal=5)
# Change initial conditions and solve again
model.concatenated_initial_conditions = pybamm.NumpyConcatenation(
pybamm.Vector([[2]]))
solution = solver.solve(model, t_eval)
np.testing.assert_array_almost_equal(solution.y[0],
2 * np.exp(-solution.t),
decimal=5)
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_find_symbols(self):
a = pybamm.StateVector(slice(0, 1))
b = pybamm.StateVector(slice(1, 2))
# test a + b
constant_symbols = OrderedDict()
variable_symbols = OrderedDict()
expr = a + b
pybamm.find_symbols(expr, constant_symbols, variable_symbols)
self.assertEqual(len(constant_symbols), 0)
# test keys of known_symbols
self.assertEqual(list(variable_symbols.keys())[0], a.id)
self.assertEqual(list(variable_symbols.keys())[1], b.id)
self.assertEqual(list(variable_symbols.keys())[2], expr.id)
# test values of variable_symbols
self.assertEqual(list(variable_symbols.values())[0], "y[:1][[True]]")
self.assertEqual(
list(variable_symbols.values())[1], "y[:2][[False, True]]")
var_a = pybamm.id_to_python_variable(a.id)
var_b = pybamm.id_to_python_variable(b.id)
self.assertEqual(
list(variable_symbols.values())[2], "{} + {}".format(var_a, var_b))
# test identical subtree
constant_symbols = OrderedDict()
variable_symbols = OrderedDict()
expr = a + b + b
pybamm.find_symbols(expr, constant_symbols, variable_symbols)
self.assertEqual(len(constant_symbols), 0)
# test keys of variable_symbols
self.assertEqual(list(variable_symbols.keys())[0], a.id)
self.assertEqual(list(variable_symbols.keys())[1], b.id)
self.assertEqual(list(variable_symbols.keys())[2], expr.children[0].id)
self.assertEqual(list(variable_symbols.keys())[3], expr.id)
# test values of variable_symbols
self.assertEqual(list(variable_symbols.values())[0], "y[:1][[True]]")
self.assertEqual(
list(variable_symbols.values())[1], "y[:2][[False, True]]")
self.assertEqual(
list(variable_symbols.values())[2], "{} + {}".format(var_a, var_b))
var_child = pybamm.id_to_python_variable(expr.children[0].id)
self.assertEqual(
list(variable_symbols.values())[3],
"{} + {}".format(var_child, var_b))
# test unary op
constant_symbols = OrderedDict()
variable_symbols = OrderedDict()
expr = a + (-b)
pybamm.find_symbols(expr, constant_symbols, variable_symbols)
self.assertEqual(len(constant_symbols), 0)
# test keys of variable_symbols
self.assertEqual(list(variable_symbols.keys())[0], a.id)
self.assertEqual(list(variable_symbols.keys())[1], b.id)
self.assertEqual(list(variable_symbols.keys())[2], expr.children[1].id)
self.assertEqual(list(variable_symbols.keys())[3], expr.id)
# test values of variable_symbols
self.assertEqual(list(variable_symbols.values())[0], "y[:1][[True]]")
self.assertEqual(
list(variable_symbols.values())[1], "y[:2][[False, True]]")
self.assertEqual(
list(variable_symbols.values())[2], "-{}".format(var_b))
var_child = pybamm.id_to_python_variable(expr.children[1].id)
self.assertEqual(
list(variable_symbols.values())[3],
"{} + {}".format(var_a, var_child))
# test function
constant_symbols = OrderedDict()
variable_symbols = OrderedDict()
expr = pybamm.Function(test_function, a)
pybamm.find_symbols(expr, constant_symbols, variable_symbols)
self.assertEqual(list(constant_symbols.keys())[0], expr.id)
self.assertEqual(list(constant_symbols.values())[0], test_function)
self.assertEqual(list(variable_symbols.keys())[0], a.id)
self.assertEqual(list(variable_symbols.keys())[1], expr.id)
self.assertEqual(list(variable_symbols.values())[0], "y[:1][[True]]")
var_funct = pybamm.id_to_python_variable(expr.id, True)
self.assertEqual(
list(variable_symbols.values())[1],
"{}({})".format(var_funct, var_a))
# test matrix
constant_symbols = OrderedDict()
variable_symbols = OrderedDict()
A = pybamm.Matrix(np.array([[1, 2], [3, 4]]))
pybamm.find_symbols(A, constant_symbols, variable_symbols)
self.assertEqual(len(variable_symbols), 0)
self.assertEqual(list(constant_symbols.keys())[0], A.id)
np.testing.assert_allclose(
list(constant_symbols.values())[0], np.array([[1, 2], [3, 4]]))
# test sparse matrix
constant_symbols = OrderedDict()
variable_symbols = OrderedDict()
A = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[0, 2], [0, 4]])))
pybamm.find_symbols(A, constant_symbols, variable_symbols)
self.assertEqual(len(variable_symbols), 0)
self.assertEqual(list(constant_symbols.keys())[0], A.id)
np.testing.assert_allclose(
list(constant_symbols.values())[0].toarray(), A.entries.toarray())
# test numpy concatentate
constant_symbols = OrderedDict()
variable_symbols = OrderedDict()
expr = pybamm.NumpyConcatenation(a, b)
pybamm.find_symbols(expr, constant_symbols, variable_symbols)
self.assertEqual(len(constant_symbols), 0)
self.assertEqual(list(variable_symbols.keys())[0], a.id)
self.assertEqual(list(variable_symbols.keys())[1], b.id)
self.assertEqual(list(variable_symbols.keys())[2], expr.id)
self.assertEqual(
list(variable_symbols.values())[2],
"np.concatenate(({},{}))".format(var_a, var_b),
)
# test domain concatentate
constant_symbols = OrderedDict()
variable_symbols = OrderedDict()
expr = pybamm.NumpyConcatenation(a, b)
pybamm.find_symbols(expr, constant_symbols, variable_symbols)
self.assertEqual(len(constant_symbols), 0)
self.assertEqual(list(variable_symbols.keys())[0], a.id)
self.assertEqual(list(variable_symbols.keys())[1], b.id)
self.assertEqual(list(variable_symbols.keys())[2], expr.id)
self.assertEqual(
list(variable_symbols.values())[2],
"np.concatenate(({},{}))".format(var_a, var_b),
)
# test that Concatentation throws
expr = pybamm.Concatenation(a, b)
with self.assertRaises(NotImplementedError):
pybamm.find_symbols(expr, constant_symbols, variable_symbols)
# test that these nodes throw
for expr in (pybamm.Variable("a"), pybamm.Parameter("a")):
with self.assertRaises(NotImplementedError):
pybamm.find_symbols(expr, constant_symbols, variable_symbols)
def set_initial_conditions_from(self, solution, inplace=True):
"""
Update initial conditions with the final states from a Solution object or from
a dictionary.
This assumes that, for each variable in self.initial_conditions, there is a
corresponding variable in the solution with the same name and size.
Parameters
----------
solution : :class:`pybamm.Solution`, or dict
The solution to use to initialize the model
inplace : bool
Whether to modify the model inplace or create a new model
"""
if inplace is True:
model = self
else:
model = self.new_copy()
if isinstance(solution, pybamm.Solution):
solution = solution.last_state
for var, equation in model.initial_conditions.items():
if isinstance(var, pybamm.Variable):
try:
final_state = solution[var.name]
except KeyError as e:
raise pybamm.ModelError(
"To update a model from a solution, each variable in "
"model.initial_conditions must appear in the solution with "
"the same key as the variable name. In the solution provided, "
f"{e.args[0]}")
if isinstance(solution, pybamm.Solution):
final_state = final_state.data
if final_state.ndim == 1:
final_state_eval = final_state[-1:]
elif final_state.ndim == 2:
final_state_eval = final_state[:, -1]
elif final_state.ndim == 3:
final_state_eval = final_state[:, :, -1].flatten(order="F")
else:
raise NotImplementedError("Variable must be 0D, 1D, or 2D")
model.initial_conditions[var] = pybamm.Vector(final_state_eval)
elif isinstance(var, pybamm.Concatenation):
children = []
for child in var.orphans:
try:
final_state = solution[child.name]
except KeyError as e:
raise pybamm.ModelError(
"To update a model from a solution, each variable in "
"model.initial_conditions must appear in the solution with "
"the same key as the variable name. In the solution "
f"provided, {e.args[0]}")
if isinstance(solution, pybamm.Solution):
final_state = final_state.data
if final_state.ndim == 2:
final_state_eval = final_state[:, -1]
else:
raise NotImplementedError(
"Variable in concatenation must be 1D")
children.append(final_state_eval)
model.initial_conditions[var] = pybamm.Vector(
np.concatenate(children))
else:
raise NotImplementedError(
"Variable must have type 'Variable' or 'Concatenation'")
# Also update the concatenated initial conditions if the model is already
# discretised
if model.is_discretised:
# Unpack slices for sorting
y_slices = {var.id: slce for var, slce in model.y_slices.items()}
slices = []
for symbol in model.initial_conditions.keys():
if isinstance(symbol, pybamm.Concatenation):
# must append the slice for the whole concatenation, so that
# equations get sorted correctly
slices.append(
slice(
y_slices[symbol.children[0].id][0].start,
y_slices[symbol.children[-1].id][0].stop,
))
else:
slices.append(y_slices[symbol.id][0])
equations = list(model.initial_conditions.values())
# sort equations according to slices
sorted_equations = [eq for _, eq in sorted(zip(slices, equations))]
model.concatenated_initial_conditions = pybamm.NumpyConcatenation(
*sorted_equations)
return model
def concatenate(self, *symbols, sparse=False):
if sparse:
return pybamm.SparseStack(*symbols)
else:
return pybamm.NumpyConcatenation(*symbols)
def set_up(self, model, inputs=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
"""
inputs = inputs or {}
y0 = model.concatenated_initial_conditions.evaluate(0, None, inputs)
# Set model timescale
model.timescale_eval = model.timescale.evaluate(u=inputs)
# 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))
if self.ode_solver is True:
self.root_method = None
if (isinstance(self, pybamm.CasadiSolver) or self.root_method
== "casadi") and model.convert_to_format != "casadi":
pybamm.logger.warning(
f"Converting {model.name} to CasADi for solving with CasADi solver"
)
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, np.size(y0)))
# 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", len(model.concatenated_rhs.evaluate(0, y0, inputs)))
y_alg = casadi.MX.sym(
"y_alg",
len(model.concatenated_algebraic.evaluate(0, y0, inputs)))
y_casadi = casadi.vertcat(y_diff, y_alg)
u_casadi = {}
for name, value in inputs.items():
if isinstance(value, numbers.Number):
u_casadi[name] = casadi.MX.sym(name)
else:
u_casadi[name] = casadi.MX.sym(name, value.shape[0])
u_casadi_stacked = casadi.vertcat(*[u for u in u_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 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
# Check for heaviside 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):
# Dimensionless
if symbol.right.id == pybamm.t.id:
expr = symbol.left
elif symbol.left.id == pybamm.t.id:
expr = symbol.right
# Dimensional
elif symbol.right.id == (pybamm.t * model.timescale).id:
expr = symbol.left.new_copy(
) / symbol.right.right.new_copy()
elif symbol.left.id == (pybamm.t * model.timescale).id:
expr = symbol.right.new_copy(
) / symbol.left.right.new_copy()
model.events.append(
pybamm.Event(str(symbol), expr.new_copy(),
pybamm.EventType.DISCONTINUITY))
# 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.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
# 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 self.root_method == "casadi" or isinstance(self,
pybamm.CasadiSolver):
mass_matrix_inv = casadi.MX(model.mass_matrix_inv.entries)
explicit_rhs = mass_matrix_inv @ rhs(t_casadi, y_casadi,
u_casadi_stacked)
model.casadi_rhs = casadi.Function(
"rhs", [t_casadi, y_casadi, u_casadi_stacked], [explicit_rhs])
model.casadi_algebraic = algebraic
# Calculate consistent initial conditions for the algebraic equations
if len(model.algebraic) > 0:
all_states = pybamm.NumpyConcatenation(
model.concatenated_rhs, model.concatenated_algebraic)
# Process again, uses caching so should be quick
residuals, residuals_eval, jacobian_eval = process(
all_states, "residuals")
model.residuals_eval = residuals_eval
model.jacobian_eval = jacobian_eval
y0_guess = y0.flatten()
model.y0 = self.calculate_consistent_state(model, 0, y0_guess,
inputs)
else:
# can use DAE solver to solve ODE model
model.residuals_eval = Residuals(rhs, "residuals", model)
model.jacobian_eval = jac_rhs
model.y0 = y0.flatten()
pybamm.logger.info("Finish solver set-up")
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")
def test_evaluator_jax(self):
a = pybamm.StateVector(slice(0, 1))
b = pybamm.StateVector(slice(1, 2))
y_tests = [
np.array([[2.0], [3.0]]),
np.array([[1.0], [3.0]]),
np.array([1.0, 3.0]),
]
t_tests = [1.0, 2.0]
# test a * b
expr = a * b
evaluator = pybamm.EvaluatorJax(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.EvaluatorJax(expr)
result = evaluator.evaluate(t=None, y=np.array([[2], [3]]))
self.assertEqual(result, 12)
# test exp
expr = pybamm.exp(a * b)
evaluator = pybamm.EvaluatorJax(expr)
result = evaluator.evaluate(t=None, y=np.array([[2], [3]]))
self.assertEqual(result, np.exp(6))
# test a constant expression
expr = pybamm.Scalar(2) * pybamm.Scalar(3)
evaluator = pybamm.EvaluatorJax(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.EvaluatorJax(expr)
for y in y_tests:
result = evaluator.evaluate(t=None, y=y)
np.testing.assert_allclose(result, expr.evaluate(t=None, y=y))
# test something with time
expr = a * pybamm.t
evaluator = pybamm.EvaluatorJax(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.EvaluatorJax(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.EvaluatorJax(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.EvaluatorJax(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.EvaluatorJax(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.EvaluatorJax(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.EvaluatorJax(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-vector 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.EvaluatorJax(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 the sparse-scalar multiplication
A = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1, 0], [0, 4]])))
for expr in [
A * pybamm.t @ pybamm.StateVector(slice(0, 2)),
pybamm.t * A @ pybamm.StateVector(slice(0, 2)),
]:
evaluator = pybamm.EvaluatorJax(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 the sparse-scalar division
A = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1, 0], [0, 4]])))
expr = A / (1.0 + pybamm.t) @ pybamm.StateVector(slice(0, 2))
evaluator = pybamm.EvaluatorJax(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]])))
a = pybamm.StateVector(slice(0, 1))
expr = pybamm.SparseStack(A, a * B)
with self.assertRaises(NotImplementedError):
evaluator = pybamm.EvaluatorJax(expr)
# test sparse mat-mat mult
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]])))
a = pybamm.StateVector(slice(0, 1))
expr = A @ (a * B)
with self.assertRaises(NotImplementedError):
evaluator = pybamm.EvaluatorJax(expr)
# test numpy concatenation
a = pybamm.Vector(np.array([[1], [2]]))
b = pybamm.Vector(np.array([[3]]))
expr = pybamm.NumpyConcatenation(a, b)
evaluator = pybamm.EvaluatorJax(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 Inner
A = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1]])))
v = pybamm.StateVector(slice(0, 1))
for expr in [
pybamm.Inner(A, v) @ v,
pybamm.Inner(v, A) @ v,
pybamm.Inner(v, v) @ v
]:
evaluator = pybamm.EvaluatorJax(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))
def add_ghost_nodes(self, symbol, discretised_symbol, bcs):
"""
Add ghost nodes to a symbol.
For Dirichlet bcs, for a boundary condition "y = a at the left-hand boundary",
we concatenate a ghost node to the start of the vector y with value "2*a - y1"
where y1 is the value of the first node.
Similarly for the right-hand boundary condition.
For Dirichlet bcs, for a boundary condition "y = a at the left-hand boundary",
we concatenate a ghost node to the start of the vector y with value "2*a - y1"
where y1 is the value of the first node.
Similarly for the right-hand boundary condition.
For Neumann bcs, for a boundary condition "dy/dx = b at the left-hand boundary",
we concatenate a ghost node to the start of the vector y with value "b*h + y1"
where y1 is the value of the first node and h is the mesh size.
Similarly for the right-hand boundary condition.
Parameters
----------
domain : list of strings
The domain of the symbol for which to add ghost nodes
bcs : dict of tuples (:class:`pybamm.Scalar`, str)
Dictionary (with keys "left" and "right") of boundary conditions. Each
boundary condition consists of a value and a flag indicating its type
(e.g. "Dirichlet")
Returns
-------
:class:`pybamm.Symbol` (shape (n+2, n))
`Matrix @ discretised_symbol + bcs_vector`. When evaluated, this gives the
discretised_symbol, with appropriate ghost nodes concatenated at each end.
"""
# get relevant grid points
submesh_list = self.mesh.combine_submeshes(*symbol.domain)
# Prepare sizes and empty bcs_vector
n = submesh_list[0].npts
sec_pts = len(submesh_list)
bcs_vector = pybamm.Vector(np.array([])) # starts empty
lbc_value, lbc_type = bcs["left"]
rbc_value, rbc_type = bcs["right"]
for i in range(sec_pts):
if lbc_value.evaluates_to_number():
lbc_i = lbc_value
else:
lbc_i = lbc_value[i]
if rbc_value.evaluates_to_number():
rbc_i = rbc_value
else:
rbc_i = rbc_value[i]
if lbc_type == "Dirichlet":
left_ghost_constant = 2 * lbc_i
elif lbc_type == "Neumann":
dx = 2 * (submesh_list[0].nodes[0] - submesh_list[0].edges[0])
left_ghost_constant = -dx * lbc_i
else:
raise ValueError(
"boundary condition must be Dirichlet or Neumann, not '{}'".format(
lbc_type
)
)
if rbc_type == "Dirichlet":
right_ghost_constant = 2 * rbc_i
elif rbc_type == "Neumann":
dx = 2 * (submesh_list[0].edges[-1] - submesh_list[0].nodes[-1])
right_ghost_constant = dx * rbc_i
else:
raise ValueError(
"boundary condition must be Dirichlet or Neumann, not '{}'".format(
rbc_type
)
)
# concatenate
bcs_vector = pybamm.NumpyConcatenation(
bcs_vector,
left_ghost_constant,
pybamm.Vector(np.zeros(n)),
right_ghost_constant,
)
# Make matrix to calculate ghost nodes
bc_factors = {"Dirichlet": -1, "Neumann": 1}
left_factor = bc_factors[lbc_type]
right_factor = bc_factors[rbc_type]
# coo_matrix takes inputs (data, (row, col)) and puts data[i] at the point
# (row[i], col[i]) for each index of data.
left_ghost_vector = coo_matrix(([left_factor], ([0], [0])), shape=(1, n))
right_ghost_vector = coo_matrix(([right_factor], ([0], [n - 1])), shape=(1, n))
sub_matrix = vstack([left_ghost_vector, eye(n), right_ghost_vector])
# repeat matrix for secondary dimensions
# Convert to csr_matrix so that we can take the index (row-slicing), which is
# not supported by the default kron format
# Note that this makes column-slicing inefficient, but this should not be an
# issue
matrix = csr_matrix(kron(eye(sec_pts), sub_matrix))
return pybamm.Matrix(matrix) @ discretised_symbol + bcs_vector
def concatenate(self, *symbols):
return pybamm.NumpyConcatenation(*symbols)
