def test_array_as_explicit_call():
assert ZeroArray(3, 2, 4).as_explicit() == ImmutableDenseNDimArray.zeros(
3, 2, 4)
assert OneArray(3, 2, 4).as_explicit() == ImmutableDenseNDimArray(
[1 for i in range(3 * 2 * 4)]).reshape(3, 2, 4)
k = Symbol("k")
X = ArraySymbol("X", k, 3, 2)
raises(ValueError, lambda: X.as_explicit())
raises(ValueError, lambda: ZeroArray(k, 2, 3).as_explicit())
raises(ValueError, lambda: OneArray(2, k, 2).as_explicit())
A = ArraySymbol("A", 3, 3)
B = ArraySymbol("B", 3, 3)
texpr = tensorproduct(A, B)
assert isinstance(texpr, ArrayTensorProduct)
assert texpr.as_explicit() == tensorproduct(A.as_explicit(),
B.as_explicit())
texpr = tensorcontraction(A, (0, 1))
assert isinstance(texpr, ArrayContraction)
assert texpr.as_explicit() == A[0, 0] + A[1, 1] + A[2, 2]
texpr = tensordiagonal(A, (0, 1))
assert isinstance(texpr, ArrayDiagonal)
assert texpr.as_explicit() == ImmutableDenseNDimArray(
[A[0, 0], A[1, 1], A[2, 2]])
texpr = permutedims(A, [1, 0])
assert isinstance(texpr, PermuteDims)
assert texpr.as_explicit() == permutedims(A.as_explicit(), [1, 0])
def test_arrayexpr_convert_indexed_to_array_broadcast():
A = ArraySymbol("A", (3, 3))
B = ArraySymbol("B", (3, 3))
expr = A[i, j] + B[k, l]
O2 = OneArray(3, 3)
expected = ArrayAdd(ArrayTensorProduct(A, O2), ArrayTensorProduct(O2, B))
assert convert_indexed_to_array(expr) == expected
assert convert_indexed_to_array(expr, [i, j, k, l]) == expected
assert convert_indexed_to_array(expr, [l, k, i, j]) == ArrayAdd(
PermuteDims(ArrayTensorProduct(O2, A), [1, 0, 2, 3]),
PermuteDims(ArrayTensorProduct(B, O2), [1, 0, 2, 3]))
expr = A[i, j] + B[j, k]
O1 = OneArray(3)
assert convert_indexed_to_array(expr, [i, j, k]) == ArrayAdd(
ArrayTensorProduct(A, O1), ArrayTensorProduct(O1, B))
C = ArraySymbol("C", (d0, d1))
D = ArraySymbol("D", (d3, d1))
expr = C[i, j] + D[k, j]
assert convert_indexed_to_array(expr, [i, j, k]) == ArrayAdd(
ArrayTensorProduct(C, OneArray(d3)),
PermuteDims(ArrayTensorProduct(OneArray(d0), D), [0, 2, 1]))
X = ArraySymbol("X", (5, 3))
expr = X[i, n] - X[j, n]
assert convert_indexed_to_array(expr, [i, j, n]) == ArrayAdd(
ArrayTensorProduct(-1, OneArray(5), X),
PermuteDims(ArrayTensorProduct(X, OneArray(5)), [0, 2, 1]))
raises(ValueError, lambda: convert_indexed_to_array(C[i, j] + D[i, j]))
def test_array_expr_reshape():
A = MatrixSymbol("A", 2, 2)
B = ArraySymbol("B", (2, 2, 2))
C = Array([1, 2, 3, 4])
expr = Reshape(A, (4,))
assert expr.expr == A
assert expr.shape == (4,)
assert expr.as_explicit() == Array([A[0, 0], A[0, 1], A[1, 0], A[1, 1]])
expr = Reshape(B, (2, 4))
assert expr.expr == B
assert expr.shape == (2, 4)
ee = expr.as_explicit()
assert isinstance(ee, ImmutableDenseNDimArray)
assert ee.shape == (2, 4)
assert ee == Array([[B[0, 0, 0], B[0, 0, 1], B[0, 1, 0], B[0, 1, 1]], [B[1, 0, 0], B[1, 0, 1], B[1, 1, 0], B[1, 1, 1]]])
expr = Reshape(A, (k, 2))
assert expr.shape == (k, 2)
raises(ValueError, lambda: Reshape(A, (2, 3)))
raises(ValueError, lambda: Reshape(A, (3,)))
expr = Reshape(C, (2, 2))
assert expr.expr == C
assert expr.shape == (2, 2)
assert expr.doit() == Array([[1, 2], [3, 4]])
def test_array_printer():
A = ArraySymbol('A', (4, 4, 6, 6, 6))
I = IndexedBase('I')
i, j, k = Idx('i', (0, 1)), Idx('j', (2, 3)), Idx('k', (4, 5))
prntr = NumPyPrinter()
assert prntr.doprint(ZeroArray(5)) == 'numpy.zeros((5,))'
assert prntr.doprint(OneArray(5)) == 'numpy.ones((5,))'
assert prntr.doprint(ArrayContraction(
A, [2, 3])) == 'numpy.einsum("abccd->abd", A)'
assert prntr.doprint(I) == 'I'
assert prntr.doprint(ArrayDiagonal(
A, [2, 3, 4])) == 'numpy.einsum("abccc->abc", A)'
assert prntr.doprint(ArrayDiagonal(
A, [0, 1], [2, 3])) == 'numpy.einsum("aabbc->cab", A)'
assert prntr.doprint(ArrayContraction(
A, [2], [3])) == 'numpy.einsum("abcde->abe", A)'
assert prntr.doprint(Assignment(I[i, j, k], I[i, j, k])) == 'I = I'
prntr = TensorflowPrinter()
assert prntr.doprint(ZeroArray(5)) == 'tensorflow.zeros((5,))'
assert prntr.doprint(OneArray(5)) == 'tensorflow.ones((5,))'
assert prntr.doprint(ArrayContraction(
A, [2, 3])) == 'tensorflow.linalg.einsum("abccd->abd", A)'
assert prntr.doprint(I) == 'I'
assert prntr.doprint(ArrayDiagonal(
A, [2, 3, 4])) == 'tensorflow.linalg.einsum("abccc->abc", A)'
assert prntr.doprint(ArrayDiagonal(
A, [0, 1], [2, 3])) == 'tensorflow.linalg.einsum("aabbc->cab", A)'
assert prntr.doprint(ArrayContraction(
A, [2], [3])) == 'tensorflow.linalg.einsum("abcde->abe", A)'
assert prntr.doprint(Assignment(I[i, j, k], I[i, j, k])) == 'I = I'
def test_arrayexpr_convert_array_element_to_array_expression():
A = ArraySymbol("A", (k, ))
B = ArraySymbol("B", (k, ))
s = Sum(A[i] * B[i], (i, 0, k - 1))
cg = convert_indexed_to_array(s)
assert cg == ArrayContraction(ArrayTensorProduct(A, B), (0, 1))
s = A[i] * B[i]
cg = convert_indexed_to_array(s)
assert cg == ArrayDiagonal(ArrayTensorProduct(A, B), (0, 1))
s = A[i] * B[j]
cg = convert_indexed_to_array(s, [i, j])
assert cg == ArrayTensorProduct(A, B)
cg = convert_indexed_to_array(s, [j, i])
assert cg == ArrayTensorProduct(B, A)
def test_array_symbol_and_element():
A = ArraySymbol("A", (2, ))
A0 = ArrayElement(A, (0, ))
A1 = ArrayElement(A, (1, ))
assert A[0] == A0
assert A[1] != A0
assert A.as_explicit() == ImmutableDenseNDimArray([A0, A1])
A2 = tensorproduct(A, A)
assert A2.shape == (2, 2)
# TODO: not yet supported:
# assert A2.as_explicit() == Array([[A[0]*A[0], A[1]*A[0]], [A[0]*A[1], A[1]*A[1]]])
A3 = tensorcontraction(A2, (0, 1))
assert A3.shape == ()
# TODO: not yet supported:
# assert A3.as_explicit() == Array([])
A = ArraySymbol("A", (2, 3, 4))
Ae = A.as_explicit()
assert Ae == ImmutableDenseNDimArray(
[[[ArrayElement(A, (i, j, k)) for k in range(4)] for j in range(3)]
for i in range(2)])
p = _permute_dims(A, Permutation(0, 2, 1))
assert isinstance(p, PermuteDims)
def test_array_element_expressions():
# Check commutative property:
assert M[0, 0]*N[0, 0] == N[0, 0]*M[0, 0]
# Check derivatives:
assert M[0, 0].diff(M[0, 0]) == 1
assert M[0, 0].diff(M[1, 0]) == 0
assert M[0, 0].diff(N[0, 0]) == 0
assert M[0, 1].diff(M[i, j]) == KroneckerDelta(i, 0)*KroneckerDelta(j, 1)
assert M[0, 1].diff(N[i, j]) == 0
K4 = ArraySymbol("K4", shape=(k, k, k, k))
assert K4[i, j, k, l].diff(K4[1, 2, 3, 4]) == (
KroneckerDelta(i, 1)*KroneckerDelta(j, 2)*KroneckerDelta(k, 3)*KroneckerDelta(l, 4)
)
def test_arrayexpr_array_diagonal():
cg = _array_diagonal(M, (1, 0))
assert cg == _array_diagonal(M, (0, 1))
cg = _array_diagonal(_array_tensor_product(M, N, P), (4, 1), (2, 0))
assert cg == _array_diagonal(_array_tensor_product(M, N, P), (1, 4), (0, 2))
cg = _array_diagonal(_array_tensor_product(M, N), (1, 2), (3,), allow_trivial_diags=True)
assert cg == _permute_dims(_array_diagonal(_array_tensor_product(M, N), (1, 2)), [0, 2, 1])
Ax = ArraySymbol("Ax", shape=(1, 2, 3, 4, 3, 5, 6, 2, 7))
cg = _array_diagonal(Ax, (1, 7), (3,), (2, 4), (6,), allow_trivial_diags=True)
assert cg == _permute_dims(_array_diagonal(Ax, (1, 7), (2, 4)), [0, 2, 4, 5, 1, 6, 3])
cg = _array_diagonal(M, (0,), allow_trivial_diags=True)
assert cg == _permute_dims(M, [1, 0])
raises(ValueError, lambda: _array_diagonal(M, (0, 0)))
def test_array_symbol_and_element():
A = ArraySymbol("A", (2,))
A0 = ArrayElement(A, (0,))
A1 = ArrayElement(A, (1,))
assert A[0] == A0
assert A[1] != A0
assert A.as_explicit() == ImmutableDenseNDimArray([A0, A1])
A2 = tensorproduct(A, A)
assert A2.shape == (2, 2)
# TODO: not yet supported:
# assert A2.as_explicit() == Array([[A[0]*A[0], A[1]*A[0]], [A[0]*A[1], A[1]*A[1]]])
A3 = tensorcontraction(A2, (0, 1))
assert A3.shape == ()
# TODO: not yet supported:
# assert A3.as_explicit() == Array([])
A = ArraySymbol("A", (2, 3, 4))
Ae = A.as_explicit()
assert Ae == ImmutableDenseNDimArray(
[[[ArrayElement(A, (i, j, k)) for k in range(4)] for j in range(3)] for i in range(2)])
p = _permute_dims(A, Permutation(0, 2, 1))
assert isinstance(p, PermuteDims)
A = ArraySymbol("A", (2,))
raises(IndexError, lambda: A[()])
raises(IndexError, lambda: A[0, 1])
raises(ValueError, lambda: A[-1])
raises(ValueError, lambda: A[2])
O = OneArray(3, 4)
Z = ZeroArray(m, n)
raises(IndexError, lambda: O[()])
raises(IndexError, lambda: O[1, 2, 3])
raises(ValueError, lambda: O[3, 0])
raises(ValueError, lambda: O[0, 4])
assert O[1, 2] == 1
assert Z[1, 2] == 0
from sympy.tensor.array.expressions.array_expressions import ArraySymbol, ArrayTensorProduct, \
PermuteDims, ArrayDiagonal, ArrayElementwiseApplyFunc, ArrayContraction
from sympy.tensor.array.expressions.arrayexpr_derivatives import array_derive
k = symbols("k")
I = Identity(k)
X = MatrixSymbol("X", k, k)
x = MatrixSymbol("x", k, 1)
A = MatrixSymbol("A", k, k)
B = MatrixSymbol("B", k, k)
C = MatrixSymbol("C", k, k)
D = MatrixSymbol("D", k, k)
A1 = ArraySymbol("A", 3, 2, k)
def test_arrayexpr_derivatives1():
res = array_derive(X, X)
assert res == PermuteDims(ArrayTensorProduct(I, I), [0, 2, 1, 3])
cg = ArrayTensorProduct(A, X, B)
res = array_derive(cg, X)
assert res == PermuteDims(
ArrayTensorProduct(I, A, I, B),
[0, 4, 2, 3, 1, 5, 6, 7])
cg = ArrayContraction(X, (0, 1))
res = array_derive(cg, X)
from sympy.tensor.array.expressions.array_expressions import ArraySymbol, ArrayTensorProduct, \
PermuteDims, ArrayDiagonal, ArrayElementwiseApplyFunc, ArrayContraction, _permute_dims, Reshape
from sympy.tensor.array.expressions.arrayexpr_derivatives import array_derive
k = symbols("k")
I = Identity(k)
X = MatrixSymbol("X", k, k)
x = MatrixSymbol("x", k, 1)
A = MatrixSymbol("A", k, k)
B = MatrixSymbol("B", k, k)
C = MatrixSymbol("C", k, k)
D = MatrixSymbol("D", k, k)
A1 = ArraySymbol("A", (3, 2, k))
def test_arrayexpr_derivatives1():
res = array_derive(X, X)
assert res == PermuteDims(ArrayTensorProduct(I, I), [0, 2, 1, 3])
cg = ArrayTensorProduct(A, X, B)
res = array_derive(cg, X)
assert res == _permute_dims(ArrayTensorProduct(I, A, I, B),
[0, 4, 2, 3, 1, 5, 6, 7])
cg = ArrayContraction(X, (0, 1))
res = array_derive(cg, X)
assert res == ArrayContraction(ArrayTensorProduct(I, I), (1, 3))
def test_arrayexpr_array_expr_applyfunc():
A = ArraySymbol("A", (3, k, 2))
aaf = ArrayElementwiseApplyFunc(sin, A)
assert aaf.shape == (3, k, 2)
from sympy.matrices.expressions.diagonal import DiagMatrix
from sympy.matrices.expressions.matexpr import MatrixSymbol
from sympy.matrices.expressions.special import ZeroMatrix
from sympy.tensor.array.arrayop import (permutedims, tensorcontraction, tensorproduct)
from sympy.tensor.array.dense_ndim_array import ImmutableDenseNDimArray
from sympy.combinatorics import Permutation
from sympy.tensor.array.expressions.array_expressions import ZeroArray, OneArray, ArraySymbol, ArrayElement, \
PermuteDims, ArrayContraction, ArrayTensorProduct, ArrayDiagonal, \
ArrayAdd, nest_permutation, ArrayElementwiseApplyFunc, _EditArrayContraction, _ArgE, _array_tensor_product, \
_array_contraction, _array_diagonal, _array_add, _permute_dims
from sympy.testing.pytest import raises
i, j, k, l, m, n = symbols("i j k l m n")
M = ArraySymbol("M", (k, k))
N = ArraySymbol("N", (k, k))
P = ArraySymbol("P", (k, k))
Q = ArraySymbol("Q", (k, k))
A = ArraySymbol("A", (k, k))
B = ArraySymbol("B", (k, k))
C = ArraySymbol("C", (k, k))
D = ArraySymbol("D", (k, k))
X = ArraySymbol("X", (k, k))
Y = ArraySymbol("Y", (k, k))
a = ArraySymbol("a", (k, 1))
b = ArraySymbol("b", (k, 1))
c = ArraySymbol("c", (k, 1))
import random
from sympy import symbols, ImmutableDenseNDimArray, tensorproduct, tensorcontraction, permutedims, MatrixSymbol, \
ZeroMatrix, sin, cos, DiagMatrix
from sympy.combinatorics import Permutation
from sympy.tensor.array.expressions.array_expressions import ZeroArray, OneArray, ArraySymbol, ArrayElement, \
PermuteDims, ArrayContraction, ArrayTensorProduct, ArrayDiagonal, \
ArrayAdd, nest_permutation, ArrayElementwiseApplyFunc, _EditArrayContraction, _ArgE
from sympy.testing.pytest import raises
i, j, k, l, m, n = symbols("i j k l m n")
M = ArraySymbol("M", k, k)
N = ArraySymbol("N", k, k)
P = ArraySymbol("P", k, k)
Q = ArraySymbol("Q", k, k)
A = ArraySymbol("A", k, k)
B = ArraySymbol("B", k, k)
C = ArraySymbol("C", k, k)
D = ArraySymbol("D", k, k)
X = ArraySymbol("X", k, k)
Y = ArraySymbol("Y", k, k)
a = ArraySymbol("a", k, 1)
b = ArraySymbol("b", k, 1)
c = ArraySymbol("c", k, 1)
d = ArraySymbol("d", k, 1)
def test_printing_str_array_expressions():
assert sstr(ArraySymbol("A", 2, 3, 4)) == "A"
assert sstr(ArrayElement("A", (2, 1/(1-x), 0))) == "A[2, 1/(1 - x), 0]"
def test_arrayexpr_convert_array_to_matrix_remove_trivial_dims():
# Tensor Product:
assert _remove_trivial_dims(_array_tensor_product(a, b)) == (a * b.T, [1, 3])
assert _remove_trivial_dims(_array_tensor_product(a.T, b)) == (a * b.T, [0, 3])
assert _remove_trivial_dims(_array_tensor_product(a, b.T)) == (a * b.T, [1, 2])
assert _remove_trivial_dims(_array_tensor_product(a.T, b.T)) == (a * b.T, [0, 2])
assert _remove_trivial_dims(_array_tensor_product(I, a.T, b.T)) == (_array_tensor_product(I, a * b.T), [2, 4])
assert _remove_trivial_dims(_array_tensor_product(a.T, I, b.T)) == (_array_tensor_product(a.T, I, b.T), [])
assert _remove_trivial_dims(_array_tensor_product(a, I)) == (_array_tensor_product(a, I), [])
assert _remove_trivial_dims(_array_tensor_product(I, a)) == (_array_tensor_product(I, a), [])
assert _remove_trivial_dims(_array_tensor_product(a.T, b.T, c, d)) == (
_array_tensor_product(a * b.T, c * d.T), [0, 2, 5, 7])
assert _remove_trivial_dims(_array_tensor_product(a.T, I, b.T, c, d, I)) == (
_array_tensor_product(a.T, I, b*c.T, d, I), [4, 7])
# Addition:
cg = ArrayAdd(_array_tensor_product(a, b), _array_tensor_product(c, d))
assert _remove_trivial_dims(cg) == (a * b.T + c * d.T, [1, 3])
# Permute Dims:
cg = PermuteDims(_array_tensor_product(a, b), Permutation(3)(1, 2))
assert _remove_trivial_dims(cg) == (a * b.T, [2, 3])
cg = PermuteDims(_array_tensor_product(a, I, b), Permutation(5)(1, 2, 3, 4))
assert _remove_trivial_dims(cg) == (cg, [])
cg = PermuteDims(_array_tensor_product(I, b, a), Permutation(5)(1, 2, 4, 5, 3))
assert _remove_trivial_dims(cg) == (PermuteDims(_array_tensor_product(I, b * a.T), [0, 2, 3, 1]), [4, 5])
# Diagonal:
cg = _array_diagonal(_array_tensor_product(M, a), (1, 2))
assert _remove_trivial_dims(cg) == (cg, [])
# Contraction:
cg = _array_contraction(_array_tensor_product(M, a), (1, 2))
assert _remove_trivial_dims(cg) == (cg, [])
# A few more cases to test the removal and shift of nested removed axes
# with array contractions and array diagonals:
tp = _array_tensor_product(
OneMatrix(1, 1),
M,
x,
OneMatrix(1, 1),
Identity(1),
)
expr = _array_contraction(tp, (1, 8))
rexpr, removed = _remove_trivial_dims(expr)
assert removed == [0, 5, 6, 7]
expr = _array_contraction(tp, (1, 8), (3, 4))
rexpr, removed = _remove_trivial_dims(expr)
assert removed == [0, 3, 4, 5]
expr = _array_diagonal(tp, (1, 8))
rexpr, removed = _remove_trivial_dims(expr)
assert removed == [0, 5, 6, 7, 8]
expr = _array_diagonal(tp, (1, 8), (3, 4))
rexpr, removed = _remove_trivial_dims(expr)
assert removed == [0, 3, 4, 5, 6]
expr = _array_diagonal(_array_contraction(_array_tensor_product(A, x, I, I1), (1, 2, 5)), (1, 4))
rexpr, removed = _remove_trivial_dims(expr)
assert removed == [2, 3]
cg = _array_diagonal(_array_tensor_product(PermuteDims(_array_tensor_product(x, I1), Permutation(1, 2, 3)), (x.T*x).applyfunc(sqrt)), (2, 4), (3, 5))
rexpr, removed = _remove_trivial_dims(cg)
assert removed == [1, 2]
# Contractions with identity matrices need to be followed by a permutation
# in order
cg = _array_contraction(_array_tensor_product(A, B, C, M, I), (1, 8))
ret, removed = _remove_trivial_dims(cg)
assert ret == PermuteDims(_array_tensor_product(A, B, C, M), [0, 2, 3, 4, 5, 6, 7, 1])
assert removed == []
cg = _array_contraction(_array_tensor_product(A, B, C, M, I), (1, 8), (3, 4))
ret, removed = _remove_trivial_dims(cg)
assert ret == PermuteDims(_array_contraction(_array_tensor_product(A, B, C, M), (3, 4)), [0, 2, 3, 4, 5, 1])
assert removed == []
# Trivial matrices are sometimes inserted into MatMul expressions:
cg = _array_tensor_product(b*b.T, a.T*a)
ret, removed = _remove_trivial_dims(cg)
assert ret == b*a.T*a*b.T
assert removed == [2, 3]
Xs = ArraySymbol("X", (3, 2, k))
cg = _array_tensor_product(M, Xs, b.T*c, a*a.T, b*b.T, c.T*d)
ret, removed = _remove_trivial_dims(cg)
assert ret == _array_tensor_product(M, Xs, a*b.T*c*c.T*d*a.T, b*b.T)
assert removed == [5, 6, 11, 12]
cg = _array_diagonal(_array_tensor_product(I, I1, x), (1, 4), (3, 5))
assert _remove_trivial_dims(cg) == (PermuteDims(_array_diagonal(_array_tensor_product(I, x), (1, 2)), Permutation(1, 2)), [1])
expr = _array_diagonal(_array_tensor_product(x, I, y), (0, 2))
assert _remove_trivial_dims(expr) == (PermuteDims(_array_tensor_product(DiagMatrix(x), y), [1, 2, 3, 0]), [0])
expr = _array_diagonal(_array_tensor_product(x, I, y), (0, 2), (3, 4))
assert _remove_trivial_dims(expr) == (expr, [])
def test_printing_str_array_expressions():
assert sstr(ArraySymbol("A", (2, 3, 4))) == "A"
assert sstr(ArrayElement("A", (2, 1/(1-x), 0))) == "A[2, 1/(1 - x), 0]"
M = MatrixSymbol("M", 3, 3)
N = MatrixSymbol("N", 3, 3)
assert sstr(ArrayElement(M*N, [x, 0])) == "(M*N)[x, 0]"