def _binary_simplify(self, left, right):
""" See :meth:`pybamm.BinaryOperator.simplify()`. """
# anything multiplied by a scalar zero returns a scalar zero
if is_scalar_zero(left):
if isinstance(right, pybamm.Array):
return pybamm.Array(np.zeros(right.shape))
else:
return pybamm.Scalar(0)
if is_scalar_zero(right):
if isinstance(left, pybamm.Array):
return pybamm.Array(np.zeros(left.shape))
else:
return pybamm.Scalar(0)
# if one of the children is a zero matrix, we have to be careful about shapes
if is_matrix_zero(left) or is_matrix_zero(right):
shape = (left * right).shape
if len(shape) == 1 or shape[1] == 1:
return pybamm.Vector(np.zeros(shape))
else:
return pybamm.Matrix(csr_matrix(shape))
# anything multiplied by a scalar one returns itself
if is_one(left):
return right
if is_one(right):
return left
return pybamm.simplify_multiplication_division(self.__class__, left,
right)
def _binary_simplify(self, left, right):
"""
See :meth:`pybamm.BinaryOperator._binary_simplify()`.
Note
----
We check for scalars first, then matrices. This is because
(Zero Matrix) - (Zero Scalar)
should return (Zero Matrix), not -(Zero Scalar).
"""
# anything added by a scalar zero returns the other child
if is_scalar_zero(left):
return -right
if is_scalar_zero(right):
return left
# Check matrices after checking scalars
if is_matrix_zero(left):
if isinstance(right, pybamm.Scalar):
return pybamm.Array(-right.value * np.ones(left.shape_for_testing))
else:
return -right
if is_matrix_zero(right):
if isinstance(left, pybamm.Scalar):
return pybamm.Array(left.value * np.ones(right.shape_for_testing))
else:
return left
return pybamm.simplify_addition_subtraction(self.__class__, left, right)
def _binary_simplify(self, left, right):
""" See :meth:`pybamm.BinaryOperator.simplify()`. """
# zero divided by zero returns nan scalar
if is_scalar_zero(left) and is_scalar_zero(right):
return pybamm.Scalar(np.nan)
# zero divided by anything returns zero
if is_scalar_zero(left):
if right.shape_for_testing == ():
return pybamm.Scalar(0)
else:
return pybamm.Array(np.zeros(right.shape_for_testing))
# anything divided by zero returns inf
if is_scalar_zero(right):
if left.shape_for_testing == ():
return pybamm.Scalar(np.inf)
else:
return pybamm.Array(np.inf * np.ones(left.shape_for_testing))
# anything divided by one is itself
if is_one(right):
return left
return pybamm.simplify_multiplication_division(self.__class__, left,
right)
def test_special_functions(self):
a = pybamm.Array(np.array([1, 2, 3, 4, 5]))
self.assert_casadi_equal(pybamm.max(a).to_casadi(),
casadi.MX(5),
evalf=True)
self.assert_casadi_equal(pybamm.min(a).to_casadi(),
casadi.MX(1),
evalf=True)
b = pybamm.Array(np.array([-2]))
c = pybamm.Array(np.array([3]))
self.assert_casadi_equal(pybamm.Function(np.abs, b).to_casadi(),
casadi.MX(2),
evalf=True)
self.assert_casadi_equal(pybamm.Function(np.abs, c).to_casadi(),
casadi.MX(3),
evalf=True)
for np_fun in [
np.sqrt,
np.tanh,
np.cosh,
np.sinh,
np.exp,
np.log,
np.sign,
np.sin,
np.cos,
np.arccosh,
np.arcsinh,
]:
self.assert_casadi_equal(pybamm.Function(np_fun, c).to_casadi(),
casadi.MX(np_fun(3)),
evalf=True)
def test_special_functions(self):
a = pybamm.Array(np.array([1, 2, 3, 4, 5]))
self.assertEqual(pybamm.max(a).to_casadi(), casadi.SX(5))
self.assertEqual(pybamm.min(a).to_casadi(), casadi.SX(1))
b = pybamm.Array(np.array([-2]))
c = pybamm.Array(np.array([3]))
self.assertEqual(pybamm.Function(np.abs, b).to_casadi(), casadi.SX(2))
self.assertEqual(pybamm.Function(np.abs, c).to_casadi(), casadi.SX(3))
def test_plot(self):
x = pybamm.Array(np.array([0, 3, 10]))
y = pybamm.Array(np.array([6, 16, 78]))
pybamm.plot(x, y, testing=True)
_, ax = plt.subplots()
ax_out = pybamm.plot(x, y, ax=ax, testing=True)
self.assertEqual(ax_out, ax)
def meshgrid(x, y, **kwargs):
"""
Return coordinate matrices as from coordinate vectors by calling
`numpy.meshgrid` with keyword arguments 'kwargs'. For a list of 'kwargs'
see the `numpy meshgrid documentation <https://tinyurl.com/y8azewrj>`_
"""
[X, Y] = np.meshgrid(x.entries, y.entries)
X = pybamm.Array(X)
Y = pybamm.Array(Y)
return X, Y
def test_plot2D(self):
x = pybamm.Array(np.array([0, 3, 10]))
y = pybamm.Array(np.array([6, 16, 78]))
X, Y = pybamm.meshgrid(x, y)
# plot with array directly
pybamm.plot2D(x, y, Y, testing=True)
# plot with meshgrid
pybamm.plot2D(X, Y, Y, testing=True)
_, ax = plt.subplots()
ax_out = pybamm.plot2D(X, Y, Y, ax=ax, testing=True)
self.assertEqual(ax_out, ax)
def test_convert_array_symbols(self):
# Arrays
a = np.array([1, 2, 3, 4, 5])
pybamm_a = pybamm.Array(a)
self.assert_casadi_equal(pybamm_a.to_casadi(), casadi.MX(a))
casadi_t = casadi.MX.sym("t")
casadi_y = casadi.MX.sym("y", 10)
casadi_y_dot = casadi.MX.sym("y_dot", 10)
pybamm_t = pybamm.Time()
pybamm_y = pybamm.StateVector(slice(0, 10))
pybamm_y_dot = pybamm.StateVectorDot(slice(0, 10))
# Time
self.assertEqual(pybamm_t.to_casadi(casadi_t, casadi_y), casadi_t)
# State Vector
self.assert_casadi_equal(pybamm_y.to_casadi(casadi_t, casadi_y),
casadi_y)
# State Vector Dot
self.assert_casadi_equal(
pybamm_y_dot.to_casadi(casadi_t, casadi_y, casadi_y_dot),
casadi_y_dot)
def _binary_simplify(self, left, right):
""" See :meth:`pybamm.BinaryOperator._binary_simplify()`. """
# zero divided by zero returns nan scalar
if is_scalar_zero(left) and is_scalar_zero(right):
return pybamm.Scalar(np.nan)
# zero divided by anything returns zero (being careful about shape)
if is_scalar_zero(left):
return zeros_of_shape(right.shape_for_testing)
# matrix zero divided by anything returns matrix zero (i.e. itself)
if is_matrix_zero(left):
return left
# anything divided by zero returns inf
if is_scalar_zero(right):
if left.shape_for_testing == ():
return pybamm.Scalar(np.inf)
else:
return pybamm.Array(np.inf * np.ones(left.shape_for_testing))
# anything divided by one is itself
if is_scalar_one(right):
return left
return pybamm.simplify_multiplication_division(self.__class__, left, right)
def linspace(start, stop, num=50, **kwargs):
"""
Creates a linearly spaced array by calling `numpy.linspace` with keyword
arguments 'kwargs'. For a list of 'kwargs' see the
`numpy linspace documentation <https://tinyurl.com/yc4ne47x>`_
"""
return pybamm.Array(np.linspace(start, stop, num, **kwargs))
def test_special_functions(self):
a = pybamm.Array(np.array([1, 2, 3, 4, 5]))
self.assert_casadi_equal(pybamm.max(a).to_casadi(),
casadi.MX(5),
evalf=True)
self.assert_casadi_equal(pybamm.min(a).to_casadi(),
casadi.MX(1),
evalf=True)
b = pybamm.Array(np.array([-2]))
c = pybamm.Array(np.array([3]))
self.assert_casadi_equal(pybamm.Function(np.abs, b).to_casadi(),
casadi.MX(2),
evalf=True)
self.assert_casadi_equal(pybamm.Function(np.abs, c).to_casadi(),
casadi.MX(3),
evalf=True)
def test_plot2D_fail(self):
x = pybamm.Array(np.array([0]))
with self.assertRaisesRegex(TypeError, "x must be 'pybamm.Array'"):
pybamm.plot2D("bad", x, x)
with self.assertRaisesRegex(TypeError, "y must be 'pybamm.Array'"):
pybamm.plot2D(x, "bad", x)
with self.assertRaisesRegex(TypeError, "z must be 'pybamm.Array'"):
pybamm.plot2D(x, x, "bad")
def test_plot2D(self):
x = pybamm.Array(np.array([0, 3, 10]))
y = pybamm.Array(np.array([6, 16, 78]))
X, Y = pybamm.meshgrid(x, y)
# plot with array directly
pybamm.plot2D(x,
y,
Y,
xlabel="x",
ylabel="y",
title="title",
testing=True)
# plot with meshgrid
pybamm.plot2D(X,
Y,
Y,
xlabel="x",
ylabel="y",
title="title",
testing=True)
def test_evaluate(self):
parameter_values = pybamm.ParameterValues({"a": 1, "b": 2, "c": 3})
a = pybamm.Parameter("a")
b = pybamm.Parameter("b")
c = pybamm.Parameter("c")
self.assertEqual(parameter_values.evaluate(a), 1)
self.assertEqual(parameter_values.evaluate(a + (b * c)), 7)
y = pybamm.StateVector(slice(0, 1))
with self.assertRaises(ValueError):
parameter_values.evaluate(y)
array = pybamm.Array(np.array([1, 2, 3]))
with self.assertRaises(ValueError):
parameter_values.evaluate(array)
def test_to_equation(self):
# Test print_name
pybamm.Addition.print_name = "test"
self.assertEqual(pybamm.Addition(1, 2).to_equation(), sympy.symbols("test"))
# Test Power
self.assertEqual(pybamm.Power(7, 2).to_equation(), 49)
# Test Division
self.assertEqual(pybamm.Division(10, 2).to_equation(), 5)
# Test Matrix Multiplication
arr1 = pybamm.Array([[1, 0], [0, 1]])
arr2 = pybamm.Array([[4, 1], [2, 2]])
self.assertEqual(
pybamm.MatrixMultiplication(arr1, arr2).to_equation(),
sympy.Matrix([[4.0, 1.0], [2.0, 2.0]]),
)
# Test EqualHeaviside
self.assertEqual(pybamm.EqualHeaviside(1, 0).to_equation(), False)
# Test NotEqualHeaviside
self.assertEqual(pybamm.NotEqualHeaviside(2, 4).to_equation(), True)
def test_evaluate(self):
parameter_values = pybamm.ParameterValues({"a": 1, "b": 2, "c": 3})
a = pybamm.Parameter("a")
b = pybamm.Parameter("b")
c = pybamm.Parameter("c")
self.assertEqual(parameter_values.evaluate(a), 1)
self.assertEqual(parameter_values.evaluate(a + (b * c)), 7)
d = pybamm.Parameter("a") + pybamm.Parameter("b") * pybamm.Array([4, 5])
np.testing.assert_array_equal(
parameter_values.evaluate(d), np.array([9, 11])[:, np.newaxis]
)
y = pybamm.StateVector(slice(0, 1))
with self.assertRaises(ValueError):
parameter_values.evaluate(y)
def test_convert_array_symbols(self):
# Arrays
a = np.array([1, 2, 3, 4, 5])
pybamm_a = pybamm.Array(a)
self.assertTrue(casadi.is_equal(pybamm_a.to_casadi(), casadi.SX(a)))
casadi_t = casadi.SX.sym("t")
casadi_y = casadi.SX.sym("y", 10)
pybamm_t = pybamm.Time()
pybamm_y = pybamm.StateVector(slice(0, 10))
# Time
self.assertEqual(pybamm_t.to_casadi(casadi_t, casadi_y), casadi_t)
# State Vector
self.assertTrue(
casadi.is_equal(pybamm_y.to_casadi(casadi_t, casadi_y), casadi_y))
# outer product
outer = pybamm.Outer(pybamm_a, pybamm_a)
self.assertTrue(
casadi.is_equal(outer.to_casadi(), casadi.SX(outer.evaluate())))
def find_symbols(symbol, constant_symbols, variable_symbols, output_jax=False):
"""
This function converts an expression tree to a dictionary of node id's and strings
specifying valid python code to calculate that nodes value, given y and t.
The function distinguishes between nodes that represent constant nodes in the tree
(e.g. a pybamm.Matrix), and those that are variable (e.g. subtrees that contain
pybamm.StateVector). The former are put in `constant_symbols`, the latter in
`variable_symbols`
Note that it is important that the arguments `constant_symbols` and
`variable_symbols` be an *ordered* dict, since the final ordering of the code lines
are important for the calculations. A dict is specified rather than a list so that
identical subtrees (which give identical id's) are not recalculated in the code
Parameters
----------
symbol : :class:`pybamm.Symbol`
The symbol or expression tree to convert
constant_symbol: collections.OrderedDict
The output dictionary of constant symbol ids to lines of code
variable_symbol: collections.OrderedDict
The output dictionary of variable (with y or t) symbol ids to lines of code
output_jax: bool
If True, only numpy and jax operations will be used in the generated code,
raises NotImplNotImplementedError if any SparseStack or Mat-Mat multiply
operations are used
"""
# constant symbols that are not numbers are stored in a list of constants, which are
# passed into the generated function constant symbols that are numbers are written
# directly into the code
if symbol.is_constant():
value = symbol.evaluate()
if not isinstance(value, numbers.Number):
if output_jax and scipy.sparse.issparse(value):
# convert any remaining sparse matrices to our custom coo matrix
constant_symbols[symbol.id] = create_jax_coo_matrix(value)
else:
constant_symbols[symbol.id] = value
return
# process children recursively
for child in symbol.children:
find_symbols(child, constant_symbols, variable_symbols, output_jax)
# calculate the variable names that will hold the result of calculating the
# children variables
children_vars = []
for child in symbol.children:
if child.is_constant():
child_eval = child.evaluate()
if isinstance(child_eval, numbers.Number):
children_vars.append(str(child_eval))
else:
children_vars.append(id_to_python_variable(child.id, True))
else:
children_vars.append(id_to_python_variable(child.id, False))
if isinstance(symbol, pybamm.BinaryOperator):
# Multiplication and Division need special handling for scipy sparse matrices
# TODO: we can pass through a dummy y and t to get the type and then hardcode
# the right line, avoiding these checks
if isinstance(symbol, pybamm.Multiplication):
dummy_eval_left = symbol.children[0].evaluate_for_shape()
dummy_eval_right = symbol.children[1].evaluate_for_shape()
if scipy.sparse.issparse(dummy_eval_left):
if output_jax and is_scalar(dummy_eval_right):
symbol_str = "{0}.scalar_multiply({1})".format(
children_vars[0], children_vars[1]
)
else:
symbol_str = "{0}.multiply({1})".format(
children_vars[0], children_vars[1]
)
elif scipy.sparse.issparse(dummy_eval_right):
if output_jax and is_scalar(dummy_eval_left):
symbol_str = "{1}.scalar_multiply({0})".format(
children_vars[0], children_vars[1]
)
else:
symbol_str = "{1}.multiply({0})".format(
children_vars[0], children_vars[1]
)
else:
symbol_str = "{0} * {1}".format(children_vars[0], children_vars[1])
elif isinstance(symbol, pybamm.Division):
dummy_eval_left = symbol.children[0].evaluate_for_shape()
dummy_eval_right = symbol.children[1].evaluate_for_shape()
if scipy.sparse.issparse(dummy_eval_left):
if output_jax and is_scalar(dummy_eval_right):
symbol_str = "{0}.scalar_multiply(1/{1})".format(
children_vars[0], children_vars[1]
)
else:
symbol_str = "{0}.multiply(1/{1})".format(
children_vars[0], children_vars[1]
)
else:
symbol_str = "{0} / {1}".format(children_vars[0], children_vars[1])
elif isinstance(symbol, pybamm.Inner):
dummy_eval_left = symbol.children[0].evaluate_for_shape()
dummy_eval_right = symbol.children[1].evaluate_for_shape()
if scipy.sparse.issparse(dummy_eval_left):
if output_jax and is_scalar(dummy_eval_right):
symbol_str = "{0}.scalar_multiply({1})".format(
children_vars[0], children_vars[1]
)
else:
symbol_str = "{0}.multiply({1})".format(
children_vars[0], children_vars[1]
)
elif scipy.sparse.issparse(dummy_eval_right):
if output_jax and is_scalar(dummy_eval_left):
symbol_str = "{1}.scalar_multiply({0})".format(
children_vars[0], children_vars[1]
)
else:
symbol_str = "{1}.multiply({0})".format(
children_vars[0], children_vars[1]
)
else:
symbol_str = "{0} * {1}".format(children_vars[0], children_vars[1])
elif isinstance(symbol, pybamm.Minimum):
symbol_str = "np.minimum({},{})".format(children_vars[0], children_vars[1])
elif isinstance(symbol, pybamm.Maximum):
symbol_str = "np.maximum({},{})".format(children_vars[0], children_vars[1])
elif isinstance(symbol, pybamm.MatrixMultiplication):
dummy_eval_left = symbol.children[0].evaluate_for_shape()
dummy_eval_right = symbol.children[1].evaluate_for_shape()
if output_jax and (
scipy.sparse.issparse(dummy_eval_left)
and scipy.sparse.issparse(dummy_eval_right)
):
raise NotImplementedError(
"sparse mat-mat multiplication not supported "
"for output_jax == True"
)
else:
symbol_str = (
children_vars[0] + " " + symbol.name + " " + children_vars[1]
)
else:
symbol_str = children_vars[0] + " " + symbol.name + " " + children_vars[1]
elif isinstance(symbol, pybamm.UnaryOperator):
# Index has a different syntax than other univariate operations
if isinstance(symbol, pybamm.Index):
symbol_str = "{}[{}:{}]".format(
children_vars[0], symbol.slice.start, symbol.slice.stop
)
else:
symbol_str = symbol.name + children_vars[0]
elif isinstance(symbol, pybamm.Function):
children_str = ""
for child_var in children_vars:
if children_str == "":
children_str = child_var
else:
children_str += ", " + child_var
if isinstance(symbol.function, np.ufunc):
# write any numpy functions directly
symbol_str = "np.{}({})".format(symbol.function.__name__, children_str)
else:
# unknown function, store it as a constant and call this in the
# generated code
constant_symbols[symbol.id] = symbol.function
funct_var = id_to_python_variable(symbol.id, True)
symbol_str = "{}({})".format(funct_var, children_str)
elif isinstance(symbol, pybamm.Concatenation):
# no need to concatenate if there is only a single child
if isinstance(symbol, pybamm.NumpyConcatenation):
if len(children_vars) == 1:
symbol_str = children_vars[0]
else:
symbol_str = "np.concatenate(({}))".format(",".join(children_vars))
elif isinstance(symbol, pybamm.SparseStack):
if len(children_vars) == 1:
symbol_str = children_vars[0]
else:
if output_jax:
raise NotImplementedError
else:
symbol_str = "scipy.sparse.vstack(({}))".format(
",".join(children_vars)
)
# DomainConcatenation specifies a particular ordering for the concatenation,
# which we must follow
elif isinstance(symbol, pybamm.DomainConcatenation):
slice_starts = []
all_child_vectors = []
for i in range(symbol.secondary_dimensions_npts):
child_vectors = []
for child_var, slices in zip(children_vars, symbol._children_slices):
for child_dom, child_slice in slices.items():
slice_starts.append(symbol._slices[child_dom][i].start)
child_vectors.append(
"{}[{}:{}]".format(
child_var, child_slice[i].start, child_slice[i].stop
)
)
all_child_vectors.extend(
[v for _, v in sorted(zip(slice_starts, child_vectors))]
)
if len(children_vars) > 1 or symbol.secondary_dimensions_npts > 1:
symbol_str = "np.concatenate(({}))".format(",".join(all_child_vectors))
else:
symbol_str = "{}".format(",".join(children_vars))
else:
raise NotImplementedError
# Note: we assume that y is being passed as a column vector
elif isinstance(symbol, pybamm.StateVector):
indices = np.argwhere(symbol.evaluation_array).reshape(-1).astype(np.int32)
consecutive = np.all(indices[1:] - indices[:-1] == 1)
if len(indices) == 1 or consecutive:
symbol_str = "y[{}:{}]".format(indices[0], indices[-1] + 1)
else:
indices_array = pybamm.Array(indices)
constant_symbols[indices_array.id] = indices
index_name = id_to_python_variable(indices_array.id, True)
symbol_str = "y[{}]".format(index_name)
elif isinstance(symbol, pybamm.Time):
symbol_str = "t"
elif isinstance(symbol, pybamm.InputParameter):
symbol_str = "inputs['{}']".format(symbol.name)
else:
raise NotImplementedError(
"Conversion to python not implemented for a symbol of type '{}'".format(
type(symbol)
)
)
variable_symbols[symbol.id] = symbol_str
def test_plot(self):
x = pybamm.Array(np.array([0, 3, 10]))
y = pybamm.Array(np.array([6, 16, 78]))
pybamm.plot(x, y, xlabel="x", ylabel="y", title="title", testing=True)
def test_name(self):
arr = pybamm.Array(np.array([1, 2, 3]))
self.assertEqual(arr.name, "Array of shape (3, 1)")
def test_to_equation(self):
self.assertEqual(
pybamm.Array([1, 2]).to_equation(), sympy.Array([[1.0], [2.0]]))
def test_list_entries(self):
vect = pybamm.Array([1, 2, 3])
np.testing.assert_array_equal(vect.entries, np.array([[1], [2], [3]]))
vect = pybamm.Array([[1], [2], [3]])
np.testing.assert_array_equal(vect.entries, np.array([[1], [2], [3]]))
