def _(expr: ArrayElementwiseApplyFunc, x: Expr):
fdiff = expr._get_function_fdiff()
subexpr = expr.expr
dsubexpr = array_derive(subexpr, x)
tp = _array_tensor_product(dsubexpr,
ArrayElementwiseApplyFunc(fdiff, subexpr))
b = get_rank(x)
c = get_rank(expr)
diag_indices = [(b + i, b + c + i) for i in range(c)]
return _array_diagonal(tp, *diag_indices)
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)
assert res == ArrayContraction(ArrayTensorProduct(I, I), (1, 3))
cg = ArrayDiagonal(X, (0, 1))
res = array_derive(cg, X)
assert res == ArrayDiagonal(ArrayTensorProduct(I, I), (1, 3))
cg = ElementwiseApplyFunction(sin, X)
res = array_derive(cg, X)
assert res.dummy_eq(ArrayDiagonal(
ArrayTensorProduct(
ElementwiseApplyFunction(cos, X),
I,
I
), (0, 3), (1, 5)))
cg = ArrayElementwiseApplyFunc(sin, X)
res = array_derive(cg, X)
assert res.dummy_eq(ArrayDiagonal(
ArrayTensorProduct(
I,
I,
ArrayElementwiseApplyFunc(cos, X)
), (1, 4), (3, 5)))
res = array_derive(A1, A1)
assert res == PermuteDims(
ArrayTensorProduct(Identity(3), Identity(2), Identity(k)),
[0, 2, 4, 1, 3, 5]
)
cg = ArrayElementwiseApplyFunc(sin, A1)
res = array_derive(cg, A1)
assert res.dummy_eq(ArrayDiagonal(
ArrayTensorProduct(
Identity(3), Identity(2), Identity(k),
ArrayElementwiseApplyFunc(cos, A1)
), (1, 6), (3, 7), (5, 8)
))
def test_convert_array_elementwise_function_to_matrix():
d = Dummy("d")
expr = ArrayElementwiseApplyFunc(Lambda(d, sin(d)), x.T*y)
assert convert_array_to_matrix(expr) == sin(x.T*y)
expr = ArrayElementwiseApplyFunc(Lambda(d, d**2), x.T*y)
assert convert_array_to_matrix(expr) == (x.T*y)**2
expr = ArrayElementwiseApplyFunc(Lambda(d, sin(d)), x)
assert convert_array_to_matrix(expr).dummy_eq(x.applyfunc(sin))
expr = ArrayElementwiseApplyFunc(Lambda(d, 1 / (2 * sqrt(d))), x)
assert convert_array_to_matrix(expr) == S.Half * HadamardPower(x, -S.Half)
def convert_matrix_to_array(expr: MatrixExpr) -> Basic:
if isinstance(expr, MatMul):
args_nonmat = []
args = []
for arg in expr.args:
if isinstance(arg, MatrixExpr):
args.append(arg)
else:
args_nonmat.append(convert_matrix_to_array(arg))
contractions = [(2 * i + 1, 2 * i + 2) for i in range(len(args) - 1)]
scalar = ArrayTensorProduct.fromiter(
args_nonmat) if args_nonmat else S.One
if scalar == 1:
tprod = ArrayTensorProduct(
*[convert_matrix_to_array(arg) for arg in args])
else:
tprod = ArrayTensorProduct(
scalar, *[convert_matrix_to_array(arg) for arg in args])
return ArrayContraction(tprod, *contractions)
elif isinstance(expr, MatAdd):
return ArrayAdd(*[convert_matrix_to_array(arg) for arg in expr.args])
elif isinstance(expr, Transpose):
return PermuteDims(convert_matrix_to_array(expr.args[0]), [1, 0])
elif isinstance(expr, Trace):
inner_expr = convert_matrix_to_array(expr.arg)
return ArrayContraction(inner_expr, (0, len(inner_expr.shape) - 1))
elif isinstance(expr, Mul):
return ArrayTensorProduct.fromiter(
convert_matrix_to_array(i) for i in expr.args)
elif isinstance(expr, Pow):
base = convert_matrix_to_array(expr.base)
if (expr.exp > 0) == True:
return ArrayTensorProduct.fromiter(base for i in range(expr.exp))
else:
return expr
elif isinstance(expr, MatPow):
base = convert_matrix_to_array(expr.base)
if expr.exp.is_Integer != True:
b = symbols("b", cls=Dummy)
return ArrayElementwiseApplyFunc(Lambda(b, b**expr.exp),
convert_matrix_to_array(base))
elif (expr.exp > 0) == True:
return convert_matrix_to_array(
MatMul.fromiter(base for i in range(expr.exp)))
else:
return expr
elif isinstance(expr, HadamardProduct):
tp = ArrayTensorProduct.fromiter(expr.args)
diag = [[2 * i for i in range(len(expr.args))],
[2 * i + 1 for i in range(len(expr.args))]]
return ArrayDiagonal(tp, *diag)
elif isinstance(expr, HadamardPower):
base, exp = expr.args
return convert_matrix_to_array(
HadamardProduct.fromiter(base for i in range(exp)))
else:
return expr
def test_array_expr_as_explicit_with_explicit_component_arrays():
# Test if .as_explicit() works with explicit-component arrays
# nested in array expressions:
from sympy.abc import x, y, z, t
A = Array([[x, y], [z, t]])
assert ArrayTensorProduct(A, A).as_explicit() == tensorproduct(A, A)
assert ArrayDiagonal(A, (0, 1)).as_explicit() == tensordiagonal(A, (0, 1))
assert ArrayContraction(A, (0, 1)).as_explicit() == tensorcontraction(A, (0, 1))
assert ArrayAdd(A, A).as_explicit() == A + A
assert ArrayElementwiseApplyFunc(sin, A).as_explicit() == A.applyfunc(sin)
assert PermuteDims(A, [1, 0]).as_explicit() == permutedims(A, [1, 0])
assert Reshape(A, [4]).as_explicit() == A.reshape(4)
def _(expr: ArrayElementwiseApplyFunc):
subexpr = _array2matrix(expr.expr)
if isinstance(subexpr, MatrixExpr):
if subexpr.shape != (1, 1):
d = expr.function.bound_symbols[0]
w = Wild("w", exclude=[d])
p = Wild("p", exclude=[d])
m = expr.function.expr.match(w * d**p)
if m is not None:
return m[w] * HadamardPower(subexpr, m[p])
return ElementwiseApplyFunction(expr.function, subexpr)
else:
return ArrayElementwiseApplyFunc(expr.function, subexpr)
def test_arrayexpr_array_expr_applyfunc():
A = ArraySymbol("A", (3, k, 2))
aaf = ArrayElementwiseApplyFunc(sin, A)
assert aaf.shape == (3, k, 2)
def _(expr: ArrayElementwiseApplyFunc):
subexpr, removed = _remove_trivial_dims(expr.expr)
return ArrayElementwiseApplyFunc(expr.function, subexpr), removed
def _(expr: ArrayElementwiseApplyFunc):
subexpr = _array2matrix(expr.expr)
if isinstance(subexpr, MatrixExpr):
return ElementwiseApplyFunction(expr.function, subexpr)
else:
return ArrayElementwiseApplyFunc(expr.function, subexpr)
def convert_matrix_to_array(expr: Basic) -> Basic:
if isinstance(expr, MatMul):
args_nonmat = []
args = []
for arg in expr.args:
if isinstance(arg, MatrixExpr):
args.append(arg)
else:
args_nonmat.append(convert_matrix_to_array(arg))
contractions = [(2*i+1, 2*i+2) for i in range(len(args)-1)]
scalar = _array_tensor_product(*args_nonmat) if args_nonmat else S.One
if scalar == 1:
tprod = _array_tensor_product(
*[convert_matrix_to_array(arg) for arg in args])
else:
tprod = _array_tensor_product(
scalar,
*[convert_matrix_to_array(arg) for arg in args])
return _array_contraction(
tprod,
*contractions
)
elif isinstance(expr, MatAdd):
return _array_add(
*[convert_matrix_to_array(arg) for arg in expr.args]
)
elif isinstance(expr, Transpose):
return _permute_dims(
convert_matrix_to_array(expr.args[0]), [1, 0]
)
elif isinstance(expr, Trace):
inner_expr: MatrixExpr = convert_matrix_to_array(expr.arg) # type: ignore
return _array_contraction(inner_expr, (0, len(inner_expr.shape) - 1))
elif isinstance(expr, Mul):
return _array_tensor_product(*[convert_matrix_to_array(i) for i in expr.args])
elif isinstance(expr, Pow):
base = convert_matrix_to_array(expr.base)
if (expr.exp > 0) == True:
return _array_tensor_product(*[base for i in range(expr.exp)])
else:
return expr
elif isinstance(expr, MatPow):
base = convert_matrix_to_array(expr.base)
if expr.exp.is_Integer != True:
b = symbols("b", cls=Dummy)
return ArrayElementwiseApplyFunc(Lambda(b, b**expr.exp), convert_matrix_to_array(base))
elif (expr.exp > 0) == True:
return convert_matrix_to_array(MatMul.fromiter(base for i in range(expr.exp)))
else:
return expr
elif isinstance(expr, HadamardProduct):
tp = _array_tensor_product(*[convert_matrix_to_array(arg) for arg in expr.args])
diag = [[2*i for i in range(len(expr.args))], [2*i+1 for i in range(len(expr.args))]]
return _array_diagonal(tp, *diag)
elif isinstance(expr, HadamardPower):
base, exp = expr.args
if isinstance(exp, Integer) and exp > 0:
return convert_matrix_to_array(HadamardProduct.fromiter(base for i in range(exp)))
else:
d = Dummy("d")
return ArrayElementwiseApplyFunc(Lambda(d, d**exp), base)
elif isinstance(expr, KroneckerProduct):
kp_args = [convert_matrix_to_array(arg) for arg in expr.args]
permutation = [2*i for i in range(len(kp_args))] + [2*i + 1 for i in range(len(kp_args))]
return Reshape(_permute_dims(_array_tensor_product(*kp_args), permutation), expr.shape)
else:
return expr
def test_arrayexpr_convert_matrix_to_array():
expr = M * N
result = ArrayContraction(ArrayTensorProduct(M, N), (1, 2))
assert convert_matrix_to_array(expr) == result
expr = M * N * M
result = _array_contraction(ArrayTensorProduct(M, N, M), (1, 2), (3, 4))
assert convert_matrix_to_array(expr) == result
expr = Transpose(M)
assert convert_matrix_to_array(expr) == PermuteDims(M, [1, 0])
expr = M * Transpose(N)
assert convert_matrix_to_array(expr) == _array_contraction(
_array_tensor_product(M, PermuteDims(N, [1, 0])), (1, 2))
expr = 3 * M * N
res = convert_matrix_to_array(expr)
rexpr = convert_array_to_matrix(res)
assert expr == rexpr
expr = 3 * M + N * M.T * M + 4 * k * N
res = convert_matrix_to_array(expr)
rexpr = convert_array_to_matrix(res)
assert expr == rexpr
expr = Inverse(M) * N
rexpr = convert_array_to_matrix(convert_matrix_to_array(expr))
assert expr == rexpr
expr = M**2
rexpr = convert_array_to_matrix(convert_matrix_to_array(expr))
assert expr == rexpr
expr = M * (2 * N + 3 * M)
res = convert_matrix_to_array(expr)
rexpr = convert_array_to_matrix(res)
assert expr == rexpr
expr = Trace(M)
result = ArrayContraction(M, (0, 1))
assert convert_matrix_to_array(expr) == result
expr = 3 * Trace(M)
result = ArrayContraction(ArrayTensorProduct(3, M), (0, 1))
assert convert_matrix_to_array(expr) == result
expr = 3 * Trace(Trace(M) * M)
result = ArrayContraction(ArrayTensorProduct(3, M, M), (0, 1), (2, 3))
assert convert_matrix_to_array(expr) == result
expr = 3 * Trace(M)**2
result = ArrayContraction(ArrayTensorProduct(3, M, M), (0, 1), (2, 3))
assert convert_matrix_to_array(expr) == result
expr = HadamardProduct(M, N)
result = ArrayDiagonal(ArrayTensorProduct(M, N), (0, 2), (1, 3))
assert convert_matrix_to_array(expr) == result
expr = HadamardProduct(M * N, N * M)
result = ArrayDiagonal(
ArrayContraction(ArrayTensorProduct(M, N, N, M), (1, 2), (5, 6)),
(0, 2), (1, 3))
assert convert_matrix_to_array(expr) == result
expr = HadamardPower(M, 2)
result = ArrayDiagonal(ArrayTensorProduct(M, M), (0, 2), (1, 3))
assert convert_matrix_to_array(expr) == result
expr = HadamardPower(M * N, 2)
result = ArrayDiagonal(
ArrayContraction(ArrayTensorProduct(M, N, M, N), (1, 2), (5, 6)),
(0, 2), (1, 3))
assert convert_matrix_to_array(expr) == result
expr = HadamardPower(M, n)
d0 = Dummy("d0")
result = ArrayElementwiseApplyFunc(Lambda(d0, d0**n), M)
assert convert_matrix_to_array(expr).dummy_eq(result)
expr = M**2
assert isinstance(expr, MatPow)
assert convert_matrix_to_array(expr) == ArrayContraction(
ArrayTensorProduct(M, M), (1, 2))
expr = a.T * b
cg = convert_matrix_to_array(expr)
assert cg == ArrayContraction(ArrayTensorProduct(a, b), (0, 2))
