def test_discretise_spatial_variable(self):
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {
"macroscale": pybamm.FiniteVolume(),
"negative particle": pybamm.FiniteVolume(),
"positive particle": pybamm.FiniteVolume(),
}
disc = pybamm.Discretisation(mesh, spatial_methods)
# space
x1 = pybamm.SpatialVariable("x", ["negative electrode"])
x1_disc = disc.process_symbol(x1)
self.assertIsInstance(x1_disc, pybamm.Vector)
np.testing.assert_array_equal(
x1_disc.evaluate(),
disc.mesh["negative electrode"][0].nodes[:, np.newaxis])
x2 = pybamm.SpatialVariable("x", ["negative electrode", "separator"])
x2_disc = disc.process_symbol(x2)
self.assertIsInstance(x2_disc, pybamm.Vector)
np.testing.assert_array_equal(
x2_disc.evaluate(),
disc.mesh.combine_submeshes("negative electrode",
"separator")[0].nodes[:, np.newaxis],
)
r = 3 * pybamm.SpatialVariable("r", ["negative particle"])
r_disc = disc.process_symbol(r)
self.assertIsInstance(r_disc, pybamm.Vector)
np.testing.assert_array_equal(
r_disc.evaluate(),
3 * disc.mesh["negative particle"][0].nodes[:, np.newaxis],
)
def test_model_solver_multiple_inputs_jax_format_error(self):
# Create model
model = pybamm.BaseModel()
model.convert_to_format = "jax"
domain = ["negative electrode", "separator", "positive electrode"]
var = pybamm.Variable("var", domain=domain)
model.rhs = {var: -pybamm.InputParameter("rate") * var}
model.initial_conditions = {var: 2 * pybamm.InputParameter("rate")}
# No need to set parameters; can use base discretisation (no spatial
# operators)
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)
solver = pybamm.ScipySolver(rtol=1e-8, atol=1e-8, method="RK45")
t_eval = np.linspace(0, 10, 100)
ninputs = 8
inputs_list = [{"rate": 0.01 * (i + 1)} for i in range(ninputs)]
with self.assertRaisesRegex(
pybamm.SolverError,
(
"Cannot solve list of inputs with multiprocessing "
'when model in format "jax".'
),
):
solver.solve(model, t_eval, inputs=inputs_list, nproc=2)
def test_model_solver_with_external(self):
# Create model
model = pybamm.BaseModel()
domain = ["negative electrode", "separator", "positive electrode"]
var1 = pybamm.Variable("var1", domain=domain)
var2 = pybamm.Variable("var2", domain=domain)
model.rhs = {var1: -var2}
model.initial_conditions = {var1: 1}
model.external_variables = [var2]
model.variables = {"var1": var1, "var2": var2}
# No need to set parameters; can use base discretisation (no spatial
# operators)
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)
# Solve
solver = pybamm.CasadiSolver(rtol=1e-8, atol=1e-8)
t_eval = np.linspace(0, 10, 100)
solution = solver.solve(model,
t_eval,
external_variables={"var2": 0.5})
np.testing.assert_allclose(solution.y.full()[0],
1 - 0.5 * solution.t,
rtol=1e-06)
def test_model_solver_multiple_inputs_happy_path(self):
for convert_to_format in ["python", "casadi"]:
# Create model
model = pybamm.BaseModel()
model.convert_to_format = convert_to_format
domain = ["negative electrode", "separator", "positive electrode"]
var = pybamm.Variable("var", domain=domain)
model.rhs = {var: -pybamm.InputParameter("rate") * var}
model.initial_conditions = {var: 1}
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)
solver = pybamm.ScipySolver(rtol=1e-8, atol=1e-8, method="RK45")
t_eval = np.linspace(0, 10, 100)
ninputs = 8
inputs_list = [{"rate": 0.01 * (i + 1)} for i in range(ninputs)]
solutions = solver.solve(model, t_eval, inputs=inputs_list, nproc=2)
for i in range(ninputs):
with self.subTest(i=i):
solution = solutions[i]
np.testing.assert_array_equal(solution.t, t_eval)
np.testing.assert_allclose(
solution.y[0], np.exp(-0.01 * (i + 1) * solution.t)
)
def test_model_solver_multiple_inputs_discontinuity_error(self):
# Create model
model = pybamm.BaseModel()
model.convert_to_format = "casadi"
domain = ["negative electrode", "separator", "positive electrode"]
var = pybamm.Variable("var", domain=domain)
model.rhs = {var: -pybamm.InputParameter("rate") * var}
model.initial_conditions = {var: 1}
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)
solver = pybamm.ScipySolver(rtol=1e-8, atol=1e-8, method="RK45")
t_eval = np.linspace(0, 10, 100)
ninputs = 8
inputs_list = [{"rate": 0.01 * (i + 1)} for i in range(ninputs)]
model.events = [
pybamm.Event(
"discontinuity",
pybamm.Scalar(t_eval[-1] / 2),
event_type=pybamm.EventType.DISCONTINUITY,
)
]
with self.assertRaisesRegex(
pybamm.SolverError,
(
"Cannot solve for a list of input parameters"
" sets with discontinuities"
),
):
solver.solve(model, t_eval, inputs=inputs_list, nproc=2)
def test_solver_with_inputs(self):
# Create model
model = pybamm.BaseModel()
model.convert_to_format = "jax"
domain = ["negative electrode", "separator", "positive electrode"]
var = pybamm.Variable("var", domain=domain)
model.rhs = {var: -pybamm.InputParameter("rate") * var}
model.initial_conditions = {var: 1}
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)
# Solve
t_eval = np.linspace(0, 10, 80)
y0 = model.concatenated_initial_conditions.evaluate().reshape(-1)
rhs = pybamm.EvaluatorJax(model.concatenated_rhs)
def fun(y, t, inputs):
return rhs.evaluate(t=t, y=y, inputs=inputs).reshape(-1)
y = pybamm.jax_bdf_integrate(
fun, y0, t_eval, {"rate": 0.1}, rtol=1e-9, atol=1e-9
)
np.testing.assert_allclose(y[:, 0].reshape(-1), np.exp(-0.1 * t_eval))
def test_model_solver_with_inputs(self):
# Create model
model = pybamm.BaseModel()
domain = ["negative electrode", "separator", "positive electrode"]
var = pybamm.Variable("var", domain=domain)
model.rhs = {var: -pybamm.InputParameter("rate") * var}
model.initial_conditions = {var: 1}
model.events = [pybamm.Event("var=0.5", pybamm.min(var - 0.5))]
# No need to set parameters; can use base discretisation (no spatial
# operators)
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)
# Solve
solver = pybamm.CasadiSolver(rtol=1e-8, atol=1e-8)
t_eval = np.linspace(0, 10, 100)
solution = solver.solve(model, t_eval, inputs={"rate": 0.1})
self.assertLess(len(solution.t), len(t_eval))
np.testing.assert_allclose(solution.y[0], np.exp(-0.1 * solution.t), rtol=1e-04)
# Without grid
solver = pybamm.CasadiSolver(mode="safe without grid", rtol=1e-8, atol=1e-8)
t_eval = np.linspace(0, 10, 100)
solution = solver.solve(model, t_eval, inputs={"rate": 0.1})
self.assertLess(len(solution.t), len(t_eval))
np.testing.assert_allclose(solution.y[0], np.exp(-0.1 * solution.t), rtol=1e-04)
solution = solver.solve(model, t_eval, inputs={"rate": 1.1})
self.assertLess(len(solution.t), len(t_eval))
np.testing.assert_allclose(solution.y[0], np.exp(-1.1 * solution.t), rtol=1e-04)
def test_solver_doesnt_support_events(self):
# Create model
model = pybamm.BaseModel()
model.convert_to_format = "jax"
domain = ["negative electrode", "separator", "positive electrode"]
var = pybamm.Variable("var", domain=domain)
model.rhs = {var: -0.1 * var}
model.initial_conditions = {var: 1}
# needs to work with multiple events (to avoid bug where only last event is
# used)
model.events = [
pybamm.Event("var=0.5", pybamm.min(var - 0.5)),
pybamm.Event("var=-0.5", pybamm.min(var + 0.5)),
]
# No need to set parameters; can use base discretisation (no spatial operators)
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)
# Solve
solver = pybamm.JaxSolver()
t_eval = np.linspace(0, 10, 100)
with self.assertRaisesRegex(RuntimeError,
"Terminate events not supported"):
solver.solve(model, t_eval)
def test_model_solver_dae_with_jacobian(self):
# Create simple test model
model = pybamm.BaseModel()
whole_cell = ["negative electrode", "separator", "positive electrode"]
var1 = pybamm.Variable("var1", domain=whole_cell)
var2 = pybamm.Variable("var2", domain=whole_cell)
model.rhs = {var1: 0.1 * var1}
model.algebraic = {var2: 2 * var1 - var2}
model.initial_conditions = {var1: 1.0, var2: 2.0}
model.initial_conditions_ydot = {var1: 0.1, var2: 0.2}
disc = get_discretisation_for_testing()
disc.process_model(model)
# Add user-supplied Jacobian to model
mesh = get_mesh_for_testing()
combined_submesh = mesh.combine_submeshes("negative electrode",
"separator",
"positive electrode")
N = combined_submesh[0].npts
def jacobian(t, y):
return np.block([
[0.1 * np.eye(N), np.zeros((N, N))],
[2.0 * np.eye(N), -1.0 * np.eye(N)],
])
# Solve
solver = pybamm.ScikitsDaeSolver(rtol=1e-8, atol=1e-8)
t_eval = np.linspace(0, 1, 100)
solution = solver.solve(model, t_eval)
np.testing.assert_array_equal(solution.t, t_eval)
np.testing.assert_allclose(solution.y[0], np.exp(0.1 * solution.t))
np.testing.assert_allclose(solution.y[-1],
2 * np.exp(0.1 * solution.t))
def test_model_solver_with_inputs(self):
# Create model
model = pybamm.BaseModel()
model.convert_to_format = "jax"
domain = ["negative electrode", "separator", "positive electrode"]
var = pybamm.Variable("var", domain=domain)
model.rhs = {var: -pybamm.InputParameter("rate") * var}
model.initial_conditions = {var: 1}
# No need to set parameters; can use base discretisation (no spatial
# operators)
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)
# Solve
solver = pybamm.JaxSolver(rtol=1e-8, atol=1e-8)
t_eval = np.linspace(0, 5, 80)
t0 = time.perf_counter()
solution = solver.solve(model, t_eval, inputs={"rate": 0.1})
t_first_solve = time.perf_counter() - t0
np.testing.assert_allclose(solution.y[0], np.exp(-0.1 * solution.t),
rtol=1e-6, atol=1e-6)
t0 = time.perf_counter()
solution = solver.solve(model, t_eval, inputs={"rate": 0.2})
t_second_solve = time.perf_counter() - t0
np.testing.assert_allclose(solution.y[0], np.exp(-0.2 * solution.t),
rtol=1e-6, atol=1e-6)
self.assertLess(t_second_solve, t_first_solve)
def test_discretise_spatial_variable(self):
# create discretisation
mesh = get_mesh_for_testing()
spatial_method = pybamm.SpatialMethod()
spatial_method.build(mesh)
# centre
x1 = pybamm.SpatialVariable("x", ["negative electrode"])
x2 = pybamm.SpatialVariable("x", ["negative electrode", "separator"])
r = pybamm.SpatialVariable("r", ["negative particle"])
for var in [x1, x2, r]:
var_disc = spatial_method.spatial_variable(var)
self.assertIsInstance(var_disc, pybamm.Vector)
np.testing.assert_array_equal(
var_disc.evaluate()[:, 0], mesh.combine_submeshes(*var.domain)[0].nodes
)
# edges
x1_edge = pybamm.SpatialVariableEdge("x", ["negative electrode"])
x2_edge = pybamm.SpatialVariableEdge("x", ["negative electrode", "separator"])
r_edge = pybamm.SpatialVariableEdge("r", ["negative particle"])
for var in [x1_edge, x2_edge, r_edge]:
var_disc = spatial_method.spatial_variable(var)
self.assertIsInstance(var_disc, pybamm.Vector)
np.testing.assert_array_equal(
var_disc.evaluate()[:, 0], mesh.combine_submeshes(*var.domain)[0].edges
)
def test_model_solver_ode_jacobian_python(self):
model = pybamm.BaseModel()
model.convert_to_format = "python"
whole_cell = ["negative electrode", "separator", "positive electrode"]
var1 = pybamm.Variable("var1", domain=whole_cell)
var2 = pybamm.Variable("var2", domain=whole_cell)
model.rhs = {var1: var1, var2: 1 - var1}
model.initial_conditions = {var1: 1.0, var2: -1.0}
model.variables = {"var1": var1, "var2": var2}
disc = get_discretisation_for_testing()
disc.process_model(model)
# Add user-supplied Jacobian to model
mesh = get_mesh_for_testing()
combined_submesh = mesh.combine_submeshes(
"negative electrode", "separator", "positive electrode"
)
N = combined_submesh[0].npts
# Solve
solver = pybamm.ScikitsOdeSolver(rtol=1e-9, atol=1e-9)
t_eval = np.linspace(0, 1, 100)
solution = solver.solve(model, t_eval)
np.testing.assert_array_equal(solution.t, t_eval)
T, Y = solution.t, solution.y
np.testing.assert_array_almost_equal(
model.variables["var1"].evaluate(T, Y),
np.ones((N, T.size)) * np.exp(T[np.newaxis, :]),
)
np.testing.assert_array_almost_equal(
model.variables["var2"].evaluate(T, Y),
np.ones((N, T.size)) * (T[np.newaxis, :] - np.exp(T[np.newaxis, :])),
)
def test_model_step(self):
# Create model
model = pybamm.BaseModel()
domain = ["negative electrode", "separator", "positive electrode"]
var = pybamm.Variable("var", domain=domain)
model.rhs = {var: 0.1 * var}
model.initial_conditions = {var: 1}
# No need to set parameters; can use base discretisation (no spatial operators)
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)
solver = pybamm.CasadiSolver(rtol=1e-8, atol=1e-8, method="idas")
# Step once
dt = 0.1
step_sol = solver.step(model, dt)
np.testing.assert_array_equal(step_sol.t, [0, dt])
np.testing.assert_allclose(step_sol.y[0], np.exp(0.1 * step_sol.t))
# Step again (return 5 points)
step_sol_2 = solver.step(model, dt, npts=5)
np.testing.assert_array_equal(step_sol_2.t, np.linspace(dt, 2 * dt, 5))
np.testing.assert_allclose(step_sol_2.y[0], np.exp(0.1 * step_sol_2.t))
# append solutions
step_sol.append(step_sol_2)
# Check steps give same solution as solve
t_eval = step_sol.t
solution = solver.solve(model, t_eval)
np.testing.assert_allclose(solution.y[0], step_sol.y[0])
def test_mass_matrix_shape(self):
"""
Test mass matrix shape
"""
# one equation
whole_cell = ["negative electrode", "separator", "positive electrode"]
c = pybamm.Variable("c", domain=whole_cell)
N = pybamm.grad(c)
model = pybamm.BaseModel()
model.rhs = {c: pybamm.div(N)}
model.initial_conditions = {c: pybamm.Scalar(0)}
model.boundary_conditions = {
c: {
"left": (0, "Dirichlet"),
"right": (0, "Dirichlet")
}
}
model.variables = {"c": c, "N": N}
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
combined_submesh = mesh.combine_submeshes(*whole_cell)
disc.process_model(model)
# mass matrix
mass = np.eye(combined_submesh[0].npts)
np.testing.assert_array_equal(mass,
model.mass_matrix.entries.toarray())
def test_model_solver_with_event_python(self):
# Create model
model = pybamm.BaseModel()
model.convert_to_format = "python"
domain = ["negative electrode", "separator", "positive electrode"]
var = pybamm.Variable("var", domain=domain)
model.rhs = {var: -0.1 * var}
model.initial_conditions = {var: 1}
# needs to work with multiple events (to avoid bug where only last event is
# used)
model.events = [
pybamm.Event("var=0.5", pybamm.min(var - 0.5)),
pybamm.Event("var=-0.5", pybamm.min(var + 0.5)),
]
# No need to set parameters; can use base discretisation (no spatial operators)
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)
# Solve
solver = pybamm.ScipySolver(rtol=1e-8, atol=1e-8, method="RK45")
t_eval = np.linspace(0, 10, 100)
solution = solver.solve(model, t_eval)
self.assertLess(len(solution.t), len(t_eval))
np.testing.assert_array_equal(solution.t, t_eval[:len(solution.t)])
np.testing.assert_allclose(solution.y[0], np.exp(-0.1 * solution.t))
def test_model_solver_with_event_with_casadi(self):
# Create model
model = pybamm.BaseModel()
for use_jacobian in [True, False]:
model.use_jacobian = use_jacobian
model.convert_to_format = "casadi"
domain = ["negative electrode", "separator", "positive electrode"]
var = pybamm.Variable("var", domain=domain)
model.rhs = {var: -0.1 * var}
model.initial_conditions = {var: 1}
model.events = {"var=0.5": pybamm.min(var - 0.5)}
# No need to set parameters; can use base discretisation (no spatial
# operators)
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)
# Solve
solver = pybamm.ScipySolver(rtol=1e-8, atol=1e-8, method="RK45")
t_eval = np.linspace(0, 10, 100)
solution = solver.solve(model, t_eval)
self.assertLess(len(solution.t), len(t_eval))
np.testing.assert_array_equal(solution.t, t_eval[:len(solution.t)])
np.testing.assert_allclose(solution.y[0],
np.exp(-0.1 * solution.t))
def test_model_solver_python(self):
# Create model
model = pybamm.BaseModel()
model.convert_to_format = "python"
domain = ["negative electrode", "separator", "positive electrode"]
var = pybamm.Variable("var", domain=domain)
model.rhs = {var: 0.1 * var}
model.initial_conditions = {var: 1}
# No need to set parameters; can use base discretisation (no spatial operators)
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)
# Solve
# Make sure that passing in extra options works
solver = pybamm.ScipySolver(rtol=1e-8,
atol=1e-8,
method="RK45",
extra_options={"first_step": 1e-4})
t_eval = np.linspace(0, 1, 80)
solution = solver.solve(model, t_eval)
np.testing.assert_array_equal(solution.t, t_eval)
np.testing.assert_allclose(solution.y[0], np.exp(0.1 * solution.t))
# Test time
self.assertEqual(solution.total_time,
solution.solve_time + solution.set_up_time)
self.assertEqual(solution.termination, "final time")
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_grad_div_shapes_mixed_domain(self):
"""
Test grad and div with Dirichlet boundary conditions (applied by grad on var)
"""
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
# grad
var = pybamm.Variable("var",
domain=["negative electrode", "separator"])
grad_eqn = pybamm.grad(var)
boundary_conditions = {
var.id: {
"left": (pybamm.Scalar(1), "Dirichlet"),
"right": (pybamm.Scalar(1), "Dirichlet"),
}
}
disc.bcs = boundary_conditions
disc.set_variable_slices([var])
grad_eqn_disc = disc.process_symbol(grad_eqn)
combined_submesh = mesh.combine_submeshes("negative electrode",
"separator")
constant_y = np.ones_like(combined_submesh[0].nodes[:, np.newaxis])
np.testing.assert_array_equal(
grad_eqn_disc.evaluate(None, constant_y),
np.zeros_like(combined_submesh[0].edges[:, np.newaxis]),
)
# div: test on linear y (should have laplacian zero) so change bcs
linear_y = combined_submesh[0].nodes
N = pybamm.grad(var)
div_eqn = pybamm.div(N)
boundary_conditions = {
var.id: {
"left": (pybamm.Scalar(0), "Dirichlet"),
"right":
(pybamm.Scalar(combined_submesh[0].edges[-1]), "Dirichlet"),
}
}
disc.bcs = boundary_conditions
grad_eqn_disc = disc.process_symbol(grad_eqn)
np.testing.assert_array_almost_equal(
grad_eqn_disc.evaluate(None, linear_y),
np.ones_like(combined_submesh[0].edges[:, np.newaxis]),
)
div_eqn_disc = disc.process_symbol(div_eqn)
np.testing.assert_array_almost_equal(
div_eqn_disc.evaluate(None, linear_y),
np.zeros_like(combined_submesh[0].nodes[:, np.newaxis]),
)
def test_numpy_domain_concatenation(self):
# create mesh
mesh = get_mesh_for_testing()
a_dom = ["negative electrode"]
b_dom = ["positive electrode"]
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)
# concatenate them the "wrong" way round to check they get reordered correctly
conc = pybamm.DomainConcatenation([b, a], mesh)
np.testing.assert_array_equal(
conc.evaluate(),
np.concatenate([
np.full(mesh[a_dom[0]].npts, 2),
np.full(mesh[b_dom[0]].npts, 1)
])[:, np.newaxis],
)
# test size and shape
self.assertEqual(conc.size, mesh[a_dom[0]].npts + mesh[b_dom[0]].npts)
self.assertEqual(conc.shape,
(mesh[a_dom[0]].npts + mesh[b_dom[0]].npts, 1))
# check the reordering in case a child vector has to be split up
a_dom = ["separator"]
b_dom = ["negative electrode", "positive electrode"]
a = 2 * pybamm.Vector(np.ones_like(mesh[a_dom[0]].nodes), domain=a_dom)
b = pybamm.Vector(
np.concatenate([
np.full(mesh[b_dom[0]].npts, 1),
np.full(mesh[b_dom[1]].npts, 3)
])[:, np.newaxis],
domain=b_dom,
)
conc = pybamm.DomainConcatenation([a, b], mesh)
np.testing.assert_array_equal(
conc.evaluate(),
np.concatenate([
np.full(mesh[b_dom[0]].npts, 1),
np.full(mesh[a_dom[0]].npts, 2),
np.full(mesh[b_dom[1]].npts, 3),
])[:, np.newaxis],
)
# test size and shape
self.assertEqual(
conc.size,
mesh[b_dom[0]].npts + mesh[a_dom[0]].npts + mesh[b_dom[1]].npts,
)
self.assertEqual(
conc.shape,
(
mesh[b_dom[0]].npts + mesh[a_dom[0]].npts +
mesh[b_dom[1]].npts,
1,
),
)
def test_model_solver_ode_with_jacobian(self):
# Create model
model = pybamm.BaseModel()
whole_cell = ["negative electrode", "separator", "positive electrode"]
var1 = pybamm.Variable("var1", domain=whole_cell)
var2 = pybamm.Variable("var2", domain=whole_cell)
model.rhs = {var1: var1, var2: 1 - var1}
model.initial_conditions = {var1: 1.0, var2: -1.0}
model.variables = {"var1": var1, "var2": var2}
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)
# Add user-supplied Jacobian to model
combined_submesh = mesh.combine_submeshes("negative electrode",
"separator",
"positive electrode")
N = combined_submesh[0].npts
# construct jacobian in order of model.rhs
J = []
for var in model.rhs.keys():
if var.id == var1.id:
J.append([np.eye(N), np.zeros((N, N))])
else:
J.append([-1.0 * np.eye(N), np.zeros((N, N))])
J = np.block(J)
def jacobian(t, y):
return J
model.jacobian = jacobian
# Solve
solver = pybamm.ScipySolver(rtol=1e-9, atol=1e-9)
t_eval = np.linspace(0, 1, 100)
solution = solver.solve(model, t_eval)
np.testing.assert_array_equal(solution.t, t_eval)
T, Y = solution.t, solution.y
np.testing.assert_array_almost_equal(
model.variables["var1"].evaluate(T, Y),
np.ones((N, T.size)) * np.exp(T[np.newaxis, :]),
)
np.testing.assert_array_almost_equal(
model.variables["var2"].evaluate(T, Y),
np.ones(
(N, T.size)) * (T[np.newaxis, :] - np.exp(T[np.newaxis, :])),
)
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_add_ghost_nodes(self):
# Set up
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
# Add ghost nodes
whole_cell = ["negative electrode", "separator", "positive electrode"]
var = pybamm.Variable("var", domain=whole_cell)
disc.set_variable_slices([var])
discretised_symbol = pybamm.StateVector(*disc.y_slices[var.id])
bcs = {
"left": (pybamm.Scalar(0), "Dirichlet"),
"right": (pybamm.Scalar(3), "Dirichlet"),
}
# Test
sp_meth = pybamm.FiniteVolume()
sp_meth.build(mesh)
sym_ghost, _ = sp_meth.add_ghost_nodes(var, discretised_symbol, bcs)
combined_submesh = mesh.combine_submeshes(*whole_cell)
y_test = np.linspace(0, 1, combined_submesh[0].npts)
np.testing.assert_array_equal(
sym_ghost.evaluate(y=y_test)[1:-1],
discretised_symbol.evaluate(y=y_test))
self.assertEqual(
(sym_ghost.evaluate(y=y_test)[0] + sym_ghost.evaluate(y=y_test)[1])
/ 2, 0)
self.assertEqual((sym_ghost.evaluate(y=y_test)[-2] +
sym_ghost.evaluate(y=y_test)[-1]) / 2, 3)
# test errors
bcs = {
"left": (pybamm.Scalar(0), "x"),
"right": (pybamm.Scalar(3), "Neumann")
}
with self.assertRaisesRegex(ValueError, "boundary condition must be"):
sp_meth.add_ghost_nodes(var, discretised_symbol, bcs)
with self.assertRaisesRegex(ValueError, "boundary condition must be"):
sp_meth.add_neumann_values(var, discretised_symbol, bcs,
var.domain)
bcs = {
"left": (pybamm.Scalar(0), "Neumann"),
"right": (pybamm.Scalar(3), "x")
}
with self.assertRaisesRegex(ValueError, "boundary condition must be"):
sp_meth.add_ghost_nodes(var, discretised_symbol, bcs)
with self.assertRaisesRegex(ValueError, "boundary condition must be"):
sp_meth.add_neumann_values(var, discretised_symbol, bcs,
var.domain)
def test_add_ghost_nodes_concatenation(self):
# Set up
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
# Add ghost nodes
whole_cell = ["negative electrode", "separator", "positive electrode"]
var_n = pybamm.Variable("var", domain=["negative electrode"])
var_s = pybamm.Variable("var", domain=["separator"])
var_p = pybamm.Variable("var", domain=["positive electrode"])
var = pybamm.Concatenation(var_n, var_s, var_p)
disc.set_variable_slices([var])
discretised_symbol = disc.process_symbol(var)
bcs = {
"left": (pybamm.Scalar(0), "Dirichlet"),
"right": (pybamm.Scalar(3), "Dirichlet"),
}
# Test
combined_submesh = mesh.combine_submeshes(*whole_cell)
y_test = np.ones_like(combined_submesh[0].nodes[:, np.newaxis])
# both
sp_meth = pybamm.FiniteVolume()
sp_meth.build(mesh)
symbol_plus_ghost_both, _ = sp_meth.add_ghost_nodes(
var, discretised_symbol, bcs
)
np.testing.assert_array_equal(
symbol_plus_ghost_both.evaluate(None, y_test)[1:-1],
discretised_symbol.evaluate(None, y_test),
)
self.assertEqual(
(
symbol_plus_ghost_both.evaluate(None, y_test)[0]
+ symbol_plus_ghost_both.evaluate(None, y_test)[1]
)
/ 2,
0,
)
self.assertEqual(
(
symbol_plus_ghost_both.evaluate(None, y_test)[-2]
+ symbol_plus_ghost_both.evaluate(None, y_test)[-1]
)
/ 2,
3,
)
def test_broadcast_checks(self):
child = pybamm.Symbol("sym", domain=["negative electrode"])
symbol = pybamm.BoundaryGradient(child, "left")
mesh = get_mesh_for_testing()
spatial_method = pybamm.SpatialMethod()
spatial_method.build(mesh)
with self.assertRaisesRegex(TypeError, "Cannot process BoundaryGradient"):
spatial_method.boundary_value_or_flux(symbol, child)
mesh = get_1p1d_mesh_for_testing()
spatial_method = pybamm.SpatialMethod()
spatial_method.build(mesh)
with self.assertRaisesRegex(NotImplementedError, "Cannot process 2D symbol"):
spatial_method.boundary_value_or_flux(symbol, child)
def test_jac_of_domain_concatenation(self):
# create mesh
mesh = get_mesh_for_testing()
y = pybamm.StateVector(slice(0, 100))
# Jacobian of a DomainConcatenation of constants is a zero matrix of the
# appropriate size
a_dom = ["negative electrode"]
b_dom = ["separator"]
c_dom = ["positive electrode"]
a_npts = mesh[a_dom[0]].npts
b_npts = mesh[b_dom[0]].npts
c_npts = mesh[c_dom[0]].npts
a = 2 * pybamm.Vector(np.ones(a_npts), domain=a_dom)
b = pybamm.Vector(np.ones(b_npts), domain=b_dom)
c = 3 * pybamm.Vector(np.ones(c_npts), domain=c_dom)
conc = pybamm.DomainConcatenation([a, b, c], mesh)
jac = conc.jac(y).evaluate().toarray()
np.testing.assert_array_equal(jac, np.zeros((100, 100)))
# Jacobian of a DomainConcatenation of StateVectors
a = 2 * pybamm.StateVector(slice(0, a_npts), domain=a_dom)
b = pybamm.StateVector(slice(a_npts, a_npts + b_npts), domain=b_dom)
c = 3 * pybamm.StateVector(
slice(a_npts + b_npts, a_npts + b_npts + c_npts), domain=c_dom
)
conc = pybamm.DomainConcatenation([a, b, c], mesh)
y0 = np.ones(100)
jac = conc.jac(y).evaluate(y=y0).toarray()
np.testing.assert_array_equal(
jac,
np.diag(
np.concatenate(
[2 * np.ones(a_npts), np.ones(b_npts), 3 * np.ones(c_npts)]
)
),
)
# multi=domain case not implemented
a = 2 * pybamm.StateVector(slice(0, a_npts), domain=a_dom)
b = pybamm.StateVector(
slice(a_npts, a_npts + b_npts + c_npts), domain=b_dom + c_dom
)
conc = pybamm.DomainConcatenation([a, b], mesh)
with self.assertRaisesRegex(
NotImplementedError, "jacobian only implemented for when each child has"
):
conc.jac(y)
def test_domain_concatenation_domains(self):
mesh = get_mesh_for_testing()
# ensure concatenated domains are sorted correctly
a = pybamm.Symbol("a", domain=["negative electrode"])
b = pybamm.Symbol("b", domain=["separator", "positive electrode"])
c = pybamm.Symbol("c", domain=["negative particle"])
conc = pybamm.DomainConcatenation([c, a, b], mesh)
self.assertEqual(
conc.domain,
[
"negative electrode",
"separator",
"positive electrode",
"negative particle",
],
)
def test_grad_div_broadcast(self):
# create mesh and discretisation
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
mesh = get_mesh_for_testing()
disc = pybamm.Discretisation(mesh, spatial_methods)
a = pybamm.PrimaryBroadcast(1, "negative electrode")
grad_a = disc.process_symbol(pybamm.grad(a))
np.testing.assert_array_equal(grad_a.evaluate(), 0)
a_edge = pybamm.PrimaryBroadcastToEdges(1, "negative electrode")
div_a = disc.process_symbol(pybamm.div(a_edge))
np.testing.assert_array_equal(div_a.evaluate(), 0)
div_grad_a = disc.process_symbol(pybamm.div(pybamm.grad(a)))
np.testing.assert_array_equal(div_grad_a.evaluate(), 0)
def test_semi_explicit_model(self):
# Create model
model = pybamm.BaseModel()
model.convert_to_format = "jax"
domain = ["negative electrode", "separator", "positive electrode"]
var = pybamm.Variable("var", domain=domain)
var2 = pybamm.Variable("var2", domain=domain)
model.rhs = {var: 0.1 * var}
model.algebraic = {var2: var2 - 2.0 * var}
# give inconsistent initial conditions, should calculate correct ones
model.initial_conditions = {var: 1.0, var2: 1.0}
# No need to set parameters; can use base discretisation (no spatial operators)
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)
# Solve
solver = pybamm.JaxSolver(
method='BDF', rtol=1e-8, atol=1e-8
)
t_eval = np.linspace(0, 1, 80)
t0 = time.perf_counter()
solution = solver.solve(model, t_eval)
t_first_solve = time.perf_counter() - t0
np.testing.assert_array_equal(solution.t, t_eval)
soln = np.exp(0.1 * solution.t)
np.testing.assert_allclose(solution.y[0], soln,
rtol=1e-7, atol=1e-7)
np.testing.assert_allclose(solution.y[-1], 2 * soln,
rtol=1e-7, atol=1e-7)
# Test time
self.assertEqual(
solution.total_time, solution.solve_time + solution.set_up_time
)
self.assertEqual(solution.termination, "final time")
t0 = time.perf_counter()
second_solution = solver.solve(model, t_eval)
t_second_solve = time.perf_counter() - t0
self.assertLess(t_second_solve, t_first_solve)
np.testing.assert_array_equal(second_solution.y, solution.y)
def test_solver(self):
# Create model
model = pybamm.BaseModel()
model.convert_to_format = "jax"
domain = ["negative electrode", "separator", "positive electrode"]
var = pybamm.Variable("var", domain=domain)
model.rhs = {var: 0.1 * var}
model.initial_conditions = {var: 1}
# No need to set parameters; can use base discretisation (no spatial operators)
# create discretisation
mesh = get_mesh_for_testing()
spatial_methods = {"macroscale": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)
# Solve
t_eval = np.linspace(0.0, 1.0, 80)
y0 = model.concatenated_initial_conditions.evaluate().reshape(-1)
rhs = pybamm.EvaluatorJax(model.concatenated_rhs)
def fun(y, t):
return rhs.evaluate(t=t, y=y).reshape(-1)
t0 = time.perf_counter()
y = pybamm.jax_bdf_integrate(fun, y0, t_eval, rtol=1e-8, atol=1e-8)
t1 = time.perf_counter() - t0
# test accuracy
np.testing.assert_allclose(y[:, 0],
np.exp(0.1 * t_eval),
rtol=1e-6,
atol=1e-6)
t0 = time.perf_counter()
y = pybamm.jax_bdf_integrate(fun, y0, t_eval, rtol=1e-8, atol=1e-8)
t2 = time.perf_counter() - t0
# second run should be much quicker
self.assertLess(t2, t1)
# test second run is accurate
np.testing.assert_allclose(y[:, 0],
np.exp(0.1 * t_eval),
rtol=1e-6,
atol=1e-6)
