def inner(left, right):
"""Return inner product of two symbols."""
left, right = preprocess_binary(left, right)
# simplify multiply by scalar zero, being careful about shape
if pybamm.is_scalar_zero(left):
return pybamm.zeros_like(right)
if pybamm.is_scalar_zero(right):
return pybamm.zeros_like(left)
# if one of the children is a zero matrix, we have to be careful about shapes
if pybamm.is_matrix_zero(left) or pybamm.is_matrix_zero(right):
return pybamm.zeros_like(pybamm.Inner(left, right))
# anything multiplied by a scalar one returns itself
if pybamm.is_scalar_one(left):
return right
if pybamm.is_scalar_one(right):
return left
return pybamm.simplify_if_constant(pybamm.Inner(left, right))
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 inner(left, right):
"""
Return inner product of two symbols.
"""
return pybamm.Inner(left, right)
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))