Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,178 @@
|
||||
r"""
|
||||
Array expressions are expressions representing N-dimensional arrays, without
|
||||
evaluating them. These expressions represent in a certain way abstract syntax
|
||||
trees of operations on N-dimensional arrays.
|
||||
|
||||
Every N-dimensional array operator has a corresponding array expression object.
|
||||
|
||||
Table of correspondences:
|
||||
|
||||
=============================== =============================
|
||||
Array operator Array expression operator
|
||||
=============================== =============================
|
||||
tensorproduct ArrayTensorProduct
|
||||
tensorcontraction ArrayContraction
|
||||
tensordiagonal ArrayDiagonal
|
||||
permutedims PermuteDims
|
||||
=============================== =============================
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
``ArraySymbol`` objects are the N-dimensional equivalent of ``MatrixSymbol``
|
||||
objects in the matrix module:
|
||||
|
||||
>>> from sympy.tensor.array.expressions import ArraySymbol
|
||||
>>> from sympy.abc import i, j, k
|
||||
>>> A = ArraySymbol("A", (3, 2, 4))
|
||||
>>> A.shape
|
||||
(3, 2, 4)
|
||||
>>> A[i, j, k]
|
||||
A[i, j, k]
|
||||
>>> A.as_explicit()
|
||||
[[[A[0, 0, 0], A[0, 0, 1], A[0, 0, 2], A[0, 0, 3]],
|
||||
[A[0, 1, 0], A[0, 1, 1], A[0, 1, 2], A[0, 1, 3]]],
|
||||
[[A[1, 0, 0], A[1, 0, 1], A[1, 0, 2], A[1, 0, 3]],
|
||||
[A[1, 1, 0], A[1, 1, 1], A[1, 1, 2], A[1, 1, 3]]],
|
||||
[[A[2, 0, 0], A[2, 0, 1], A[2, 0, 2], A[2, 0, 3]],
|
||||
[A[2, 1, 0], A[2, 1, 1], A[2, 1, 2], A[2, 1, 3]]]]
|
||||
|
||||
Component-explicit arrays can be added inside array expressions:
|
||||
|
||||
>>> from sympy import Array
|
||||
>>> from sympy import tensorproduct
|
||||
>>> from sympy.tensor.array.expressions import ArrayTensorProduct
|
||||
>>> a = Array([1, 2, 3])
|
||||
>>> b = Array([i, j, k])
|
||||
>>> expr = ArrayTensorProduct(a, b, b)
|
||||
>>> expr
|
||||
ArrayTensorProduct([1, 2, 3], [i, j, k], [i, j, k])
|
||||
>>> expr.as_explicit() == tensorproduct(a, b, b)
|
||||
True
|
||||
|
||||
Constructing array expressions from index-explicit forms
|
||||
--------------------------------------------------------
|
||||
|
||||
Array expressions are index-implicit. This means they do not use any indices to
|
||||
represent array operations. The function ``convert_indexed_to_array( ... )``
|
||||
may be used to convert index-explicit expressions to array expressions.
|
||||
It takes as input two parameters: the index-explicit expression and the order
|
||||
of the indices:
|
||||
|
||||
>>> from sympy.tensor.array.expressions import convert_indexed_to_array
|
||||
>>> from sympy import Sum
|
||||
>>> A = ArraySymbol("A", (3, 3))
|
||||
>>> B = ArraySymbol("B", (3, 3))
|
||||
>>> convert_indexed_to_array(A[i, j], [i, j])
|
||||
A
|
||||
>>> convert_indexed_to_array(A[i, j], [j, i])
|
||||
PermuteDims(A, (0 1))
|
||||
>>> convert_indexed_to_array(A[i, j] + B[j, i], [i, j])
|
||||
ArrayAdd(A, PermuteDims(B, (0 1)))
|
||||
>>> convert_indexed_to_array(Sum(A[i, j]*B[j, k], (j, 0, 2)), [i, k])
|
||||
ArrayContraction(ArrayTensorProduct(A, B), (1, 2))
|
||||
|
||||
The diagonal of a matrix in the array expression form:
|
||||
|
||||
>>> convert_indexed_to_array(A[i, i], [i])
|
||||
ArrayDiagonal(A, (0, 1))
|
||||
|
||||
The trace of a matrix in the array expression form:
|
||||
|
||||
>>> convert_indexed_to_array(Sum(A[i, i], (i, 0, 2)), [i])
|
||||
ArrayContraction(A, (0, 1))
|
||||
|
||||
Compatibility with matrices
|
||||
---------------------------
|
||||
|
||||
Array expressions can be mixed with objects from the matrix module:
|
||||
|
||||
>>> from sympy import MatrixSymbol
|
||||
>>> from sympy.tensor.array.expressions import ArrayContraction
|
||||
>>> M = MatrixSymbol("M", 3, 3)
|
||||
>>> N = MatrixSymbol("N", 3, 3)
|
||||
|
||||
Express the matrix product in the array expression form:
|
||||
|
||||
>>> from sympy.tensor.array.expressions import convert_matrix_to_array
|
||||
>>> expr = convert_matrix_to_array(M*N)
|
||||
>>> expr
|
||||
ArrayContraction(ArrayTensorProduct(M, N), (1, 2))
|
||||
|
||||
The expression can be converted back to matrix form:
|
||||
|
||||
>>> from sympy.tensor.array.expressions import convert_array_to_matrix
|
||||
>>> convert_array_to_matrix(expr)
|
||||
M*N
|
||||
|
||||
Add a second contraction on the remaining axes in order to get the trace of `M \cdot N`:
|
||||
|
||||
>>> expr_tr = ArrayContraction(expr, (0, 1))
|
||||
>>> expr_tr
|
||||
ArrayContraction(ArrayContraction(ArrayTensorProduct(M, N), (1, 2)), (0, 1))
|
||||
|
||||
Flatten the expression by calling ``.doit()`` and remove the nested array contraction operations:
|
||||
|
||||
>>> expr_tr.doit()
|
||||
ArrayContraction(ArrayTensorProduct(M, N), (0, 3), (1, 2))
|
||||
|
||||
Get the explicit form of the array expression:
|
||||
|
||||
>>> expr.as_explicit()
|
||||
[[M[0, 0]*N[0, 0] + M[0, 1]*N[1, 0] + M[0, 2]*N[2, 0], M[0, 0]*N[0, 1] + M[0, 1]*N[1, 1] + M[0, 2]*N[2, 1], M[0, 0]*N[0, 2] + M[0, 1]*N[1, 2] + M[0, 2]*N[2, 2]],
|
||||
[M[1, 0]*N[0, 0] + M[1, 1]*N[1, 0] + M[1, 2]*N[2, 0], M[1, 0]*N[0, 1] + M[1, 1]*N[1, 1] + M[1, 2]*N[2, 1], M[1, 0]*N[0, 2] + M[1, 1]*N[1, 2] + M[1, 2]*N[2, 2]],
|
||||
[M[2, 0]*N[0, 0] + M[2, 1]*N[1, 0] + M[2, 2]*N[2, 0], M[2, 0]*N[0, 1] + M[2, 1]*N[1, 1] + M[2, 2]*N[2, 1], M[2, 0]*N[0, 2] + M[2, 1]*N[1, 2] + M[2, 2]*N[2, 2]]]
|
||||
|
||||
Express the trace of a matrix:
|
||||
|
||||
>>> from sympy import Trace
|
||||
>>> convert_matrix_to_array(Trace(M))
|
||||
ArrayContraction(M, (0, 1))
|
||||
>>> convert_matrix_to_array(Trace(M*N))
|
||||
ArrayContraction(ArrayTensorProduct(M, N), (0, 3), (1, 2))
|
||||
|
||||
Express the transposition of a matrix (will be expressed as a permutation of the axes:
|
||||
|
||||
>>> convert_matrix_to_array(M.T)
|
||||
PermuteDims(M, (0 1))
|
||||
|
||||
Compute the derivative array expressions:
|
||||
|
||||
>>> from sympy.tensor.array.expressions import array_derive
|
||||
>>> d = array_derive(M, M)
|
||||
>>> d
|
||||
PermuteDims(ArrayTensorProduct(I, I), (3)(1 2))
|
||||
|
||||
Verify that the derivative corresponds to the form computed with explicit matrices:
|
||||
|
||||
>>> d.as_explicit()
|
||||
[[[[1, 0, 0], [0, 0, 0], [0, 0, 0]], [[0, 1, 0], [0, 0, 0], [0, 0, 0]], [[0, 0, 1], [0, 0, 0], [0, 0, 0]]], [[[0, 0, 0], [1, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 1, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 1], [0, 0, 0]]], [[[0, 0, 0], [0, 0, 0], [1, 0, 0]], [[0, 0, 0], [0, 0, 0], [0, 1, 0]], [[0, 0, 0], [0, 0, 0], [0, 0, 1]]]]
|
||||
>>> Me = M.as_explicit()
|
||||
>>> Me.diff(Me)
|
||||
[[[[1, 0, 0], [0, 0, 0], [0, 0, 0]], [[0, 1, 0], [0, 0, 0], [0, 0, 0]], [[0, 0, 1], [0, 0, 0], [0, 0, 0]]], [[[0, 0, 0], [1, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 1, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 1], [0, 0, 0]]], [[[0, 0, 0], [0, 0, 0], [1, 0, 0]], [[0, 0, 0], [0, 0, 0], [0, 1, 0]], [[0, 0, 0], [0, 0, 0], [0, 0, 1]]]]
|
||||
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
"ArraySymbol", "ArrayElement", "ZeroArray", "OneArray",
|
||||
"ArrayTensorProduct",
|
||||
"ArrayContraction",
|
||||
"ArrayDiagonal",
|
||||
"PermuteDims",
|
||||
"ArrayAdd",
|
||||
"ArrayElementwiseApplyFunc",
|
||||
"Reshape",
|
||||
"convert_array_to_matrix",
|
||||
"convert_matrix_to_array",
|
||||
"convert_array_to_indexed",
|
||||
"convert_indexed_to_array",
|
||||
"array_derive",
|
||||
]
|
||||
|
||||
from sympy.tensor.array.expressions.array_expressions import ArrayTensorProduct, ArrayAdd, PermuteDims, ArrayDiagonal, \
|
||||
ArrayContraction, Reshape, ArraySymbol, ArrayElement, ZeroArray, OneArray, ArrayElementwiseApplyFunc
|
||||
from sympy.tensor.array.expressions.arrayexpr_derivatives import array_derive
|
||||
from sympy.tensor.array.expressions.from_array_to_indexed import convert_array_to_indexed
|
||||
from sympy.tensor.array.expressions.from_array_to_matrix import convert_array_to_matrix
|
||||
from sympy.tensor.array.expressions.from_indexed_to_array import convert_indexed_to_array
|
||||
from sympy.tensor.array.expressions.from_matrix_to_array import convert_matrix_to_array
|
||||
+1969
File diff suppressed because it is too large
Load Diff
+194
@@ -0,0 +1,194 @@
|
||||
import operator
|
||||
from functools import reduce, singledispatch
|
||||
|
||||
from sympy.core.expr import Expr
|
||||
from sympy.core.singleton import S
|
||||
from sympy.matrices.expressions.hadamard import HadamardProduct
|
||||
from sympy.matrices.expressions.inverse import Inverse
|
||||
from sympy.matrices.expressions.matexpr import (MatrixExpr, MatrixSymbol)
|
||||
from sympy.matrices.expressions.special import Identity, OneMatrix
|
||||
from sympy.matrices.expressions.transpose import Transpose
|
||||
from sympy.combinatorics.permutations import _af_invert
|
||||
from sympy.matrices.expressions.applyfunc import ElementwiseApplyFunction
|
||||
from sympy.tensor.array.expressions.array_expressions import (
|
||||
_ArrayExpr, ZeroArray, ArraySymbol, ArrayTensorProduct, ArrayAdd,
|
||||
PermuteDims, ArrayDiagonal, ArrayElementwiseApplyFunc, get_rank,
|
||||
get_shape, ArrayContraction, _array_tensor_product, _array_contraction,
|
||||
_array_diagonal, _array_add, _permute_dims, Reshape)
|
||||
from sympy.tensor.array.expressions.from_matrix_to_array import convert_matrix_to_array
|
||||
|
||||
|
||||
@singledispatch
|
||||
def array_derive(expr, x):
|
||||
"""
|
||||
Derivatives (gradients) for array expressions.
|
||||
"""
|
||||
raise NotImplementedError(f"not implemented for type {type(expr)}")
|
||||
|
||||
|
||||
@array_derive.register(Expr)
|
||||
def _(expr: Expr, x: _ArrayExpr):
|
||||
return ZeroArray(*x.shape)
|
||||
|
||||
|
||||
@array_derive.register(ArrayTensorProduct)
|
||||
def _(expr: ArrayTensorProduct, x: Expr):
|
||||
args = expr.args
|
||||
addend_list = []
|
||||
for i, arg in enumerate(expr.args):
|
||||
darg = array_derive(arg, x)
|
||||
if darg == 0:
|
||||
continue
|
||||
args_prev = args[:i]
|
||||
args_succ = args[i+1:]
|
||||
shape_prev = reduce(operator.add, map(get_shape, args_prev), ())
|
||||
shape_succ = reduce(operator.add, map(get_shape, args_succ), ())
|
||||
addend = _array_tensor_product(*args_prev, darg, *args_succ)
|
||||
tot1 = len(get_shape(x))
|
||||
tot2 = tot1 + len(shape_prev)
|
||||
tot3 = tot2 + len(get_shape(arg))
|
||||
tot4 = tot3 + len(shape_succ)
|
||||
perm = list(range(tot1, tot2)) + \
|
||||
list(range(tot1)) + list(range(tot2, tot3)) + \
|
||||
list(range(tot3, tot4))
|
||||
addend = _permute_dims(addend, _af_invert(perm))
|
||||
addend_list.append(addend)
|
||||
if len(addend_list) == 1:
|
||||
return addend_list[0]
|
||||
elif len(addend_list) == 0:
|
||||
return S.Zero
|
||||
else:
|
||||
return _array_add(*addend_list)
|
||||
|
||||
|
||||
@array_derive.register(ArraySymbol)
|
||||
def _(expr: ArraySymbol, x: _ArrayExpr):
|
||||
if expr == x:
|
||||
return _permute_dims(
|
||||
ArrayTensorProduct.fromiter(Identity(i) for i in expr.shape),
|
||||
[2*i for i in range(len(expr.shape))] + [2*i+1 for i in range(len(expr.shape))]
|
||||
)
|
||||
return ZeroArray(*(x.shape + expr.shape))
|
||||
|
||||
|
||||
@array_derive.register(MatrixSymbol)
|
||||
def _(expr: MatrixSymbol, x: _ArrayExpr):
|
||||
m, n = expr.shape
|
||||
if expr == x:
|
||||
return _permute_dims(
|
||||
_array_tensor_product(Identity(m), Identity(n)),
|
||||
[0, 2, 1, 3]
|
||||
)
|
||||
return ZeroArray(*(x.shape + expr.shape))
|
||||
|
||||
|
||||
@array_derive.register(Identity)
|
||||
def _(expr: Identity, x: _ArrayExpr):
|
||||
return ZeroArray(*(x.shape + expr.shape))
|
||||
|
||||
|
||||
@array_derive.register(OneMatrix)
|
||||
def _(expr: OneMatrix, x: _ArrayExpr):
|
||||
return ZeroArray(*(x.shape + expr.shape))
|
||||
|
||||
|
||||
@array_derive.register(Transpose)
|
||||
def _(expr: Transpose, x: Expr):
|
||||
# D(A.T, A) ==> (m,n,i,j) ==> D(A_ji, A_mn) = d_mj d_ni
|
||||
# D(B.T, A) ==> (m,n,i,j) ==> D(B_ji, A_mn)
|
||||
fd = array_derive(expr.arg, x)
|
||||
return _permute_dims(fd, [0, 1, 3, 2])
|
||||
|
||||
|
||||
@array_derive.register(Inverse)
|
||||
def _(expr: Inverse, x: Expr):
|
||||
mat = expr.I
|
||||
dexpr = array_derive(mat, x)
|
||||
tp = _array_tensor_product(-expr, dexpr, expr)
|
||||
mp = _array_contraction(tp, (1, 4), (5, 6))
|
||||
pp = _permute_dims(mp, [1, 2, 0, 3])
|
||||
return pp
|
||||
|
||||
|
||||
@array_derive.register(ElementwiseApplyFunction)
|
||||
def _(expr: ElementwiseApplyFunction, x: Expr):
|
||||
assert get_rank(expr) == 2
|
||||
assert get_rank(x) == 2
|
||||
fdiff = expr._get_function_fdiff()
|
||||
dexpr = array_derive(expr.expr, x)
|
||||
tp = _array_tensor_product(
|
||||
ElementwiseApplyFunction(fdiff, expr.expr),
|
||||
dexpr
|
||||
)
|
||||
td = _array_diagonal(
|
||||
tp, (0, 4), (1, 5)
|
||||
)
|
||||
return td
|
||||
|
||||
|
||||
@array_derive.register(ArrayElementwiseApplyFunc)
|
||||
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)
|
||||
|
||||
|
||||
@array_derive.register(MatrixExpr)
|
||||
def _(expr: MatrixExpr, x: Expr):
|
||||
cg = convert_matrix_to_array(expr)
|
||||
return array_derive(cg, x)
|
||||
|
||||
|
||||
@array_derive.register(HadamardProduct)
|
||||
def _(expr: HadamardProduct, x: Expr):
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
@array_derive.register(ArrayContraction)
|
||||
def _(expr: ArrayContraction, x: Expr):
|
||||
fd = array_derive(expr.expr, x)
|
||||
rank_x = len(get_shape(x))
|
||||
contraction_indices = expr.contraction_indices
|
||||
new_contraction_indices = [tuple(j + rank_x for j in i) for i in contraction_indices]
|
||||
return _array_contraction(fd, *new_contraction_indices)
|
||||
|
||||
|
||||
@array_derive.register(ArrayDiagonal)
|
||||
def _(expr: ArrayDiagonal, x: Expr):
|
||||
dsubexpr = array_derive(expr.expr, x)
|
||||
rank_x = len(get_shape(x))
|
||||
diag_indices = [[j + rank_x for j in i] for i in expr.diagonal_indices]
|
||||
return _array_diagonal(dsubexpr, *diag_indices)
|
||||
|
||||
|
||||
@array_derive.register(ArrayAdd)
|
||||
def _(expr: ArrayAdd, x: Expr):
|
||||
return _array_add(*[array_derive(arg, x) for arg in expr.args])
|
||||
|
||||
|
||||
@array_derive.register(PermuteDims)
|
||||
def _(expr: PermuteDims, x: Expr):
|
||||
de = array_derive(expr.expr, x)
|
||||
perm = [0, 1] + [i + 2 for i in expr.permutation.array_form]
|
||||
return _permute_dims(de, perm)
|
||||
|
||||
|
||||
@array_derive.register(Reshape)
|
||||
def _(expr: Reshape, x: Expr):
|
||||
de = array_derive(expr.expr, x)
|
||||
return Reshape(de, get_shape(x) + expr.shape)
|
||||
|
||||
|
||||
def matrix_derive(expr, x):
|
||||
from sympy.tensor.array.expressions.from_array_to_matrix import convert_array_to_matrix
|
||||
ce = convert_matrix_to_array(expr)
|
||||
dce = array_derive(ce, x)
|
||||
return convert_array_to_matrix(dce).doit()
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
from sympy.tensor.array.expressions import from_array_to_indexed
|
||||
from sympy.utilities.decorator import deprecated
|
||||
|
||||
|
||||
_conv_to_from_decorator = deprecated(
|
||||
"module has been renamed by replacing 'conv_' with 'from_' in its name",
|
||||
deprecated_since_version="1.11",
|
||||
active_deprecations_target="deprecated-conv-array-expr-module-names",
|
||||
)
|
||||
|
||||
|
||||
convert_array_to_indexed = _conv_to_from_decorator(from_array_to_indexed.convert_array_to_indexed)
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
from sympy.tensor.array.expressions import from_array_to_matrix
|
||||
from sympy.tensor.array.expressions.conv_array_to_indexed import _conv_to_from_decorator
|
||||
|
||||
convert_array_to_matrix = _conv_to_from_decorator(from_array_to_matrix.convert_array_to_matrix)
|
||||
_array2matrix = _conv_to_from_decorator(from_array_to_matrix._array2matrix)
|
||||
_remove_trivial_dims = _conv_to_from_decorator(from_array_to_matrix._remove_trivial_dims)
|
||||
+4
@@ -0,0 +1,4 @@
|
||||
from sympy.tensor.array.expressions import from_indexed_to_array
|
||||
from sympy.tensor.array.expressions.conv_array_to_indexed import _conv_to_from_decorator
|
||||
|
||||
convert_indexed_to_array = _conv_to_from_decorator(from_indexed_to_array.convert_indexed_to_array)
|
||||
+4
@@ -0,0 +1,4 @@
|
||||
from sympy.tensor.array.expressions import from_matrix_to_array
|
||||
from sympy.tensor.array.expressions.conv_array_to_indexed import _conv_to_from_decorator
|
||||
|
||||
convert_matrix_to_array = _conv_to_from_decorator(from_matrix_to_array.convert_matrix_to_array)
|
||||
+84
@@ -0,0 +1,84 @@
|
||||
import collections.abc
|
||||
import operator
|
||||
from itertools import accumulate
|
||||
|
||||
from sympy import Mul, Sum, Dummy, Add
|
||||
from sympy.tensor.array.expressions import PermuteDims, ArrayAdd, ArrayElementwiseApplyFunc, Reshape
|
||||
from sympy.tensor.array.expressions.array_expressions import ArrayTensorProduct, get_rank, ArrayContraction, \
|
||||
ArrayDiagonal, get_shape, _get_array_element_or_slice, _ArrayExpr
|
||||
from sympy.tensor.array.expressions.utils import _apply_permutation_to_list
|
||||
|
||||
|
||||
def convert_array_to_indexed(expr, indices):
|
||||
return _ConvertArrayToIndexed().do_convert(expr, indices)
|
||||
|
||||
|
||||
class _ConvertArrayToIndexed:
|
||||
|
||||
def __init__(self):
|
||||
self.count_dummies = 0
|
||||
|
||||
def do_convert(self, expr, indices):
|
||||
if isinstance(expr, ArrayTensorProduct):
|
||||
cumul = list(accumulate([0] + [get_rank(arg) for arg in expr.args]))
|
||||
indices_grp = [indices[cumul[i]:cumul[i+1]] for i in range(len(expr.args))]
|
||||
return Mul.fromiter(self.do_convert(arg, ind) for arg, ind in zip(expr.args, indices_grp))
|
||||
if isinstance(expr, ArrayContraction):
|
||||
new_indices = [None for i in range(get_rank(expr.expr))]
|
||||
limits = []
|
||||
bottom_shape = get_shape(expr.expr)
|
||||
for contraction_index_grp in expr.contraction_indices:
|
||||
d = Dummy(f"d{self.count_dummies}")
|
||||
self.count_dummies += 1
|
||||
dim = bottom_shape[contraction_index_grp[0]]
|
||||
limits.append((d, 0, dim-1))
|
||||
for i in contraction_index_grp:
|
||||
new_indices[i] = d
|
||||
j = 0
|
||||
for i in range(len(new_indices)):
|
||||
if new_indices[i] is None:
|
||||
new_indices[i] = indices[j]
|
||||
j += 1
|
||||
newexpr = self.do_convert(expr.expr, new_indices)
|
||||
return Sum(newexpr, *limits)
|
||||
if isinstance(expr, ArrayDiagonal):
|
||||
new_indices = [None for i in range(get_rank(expr.expr))]
|
||||
ind_pos = expr._push_indices_down(expr.diagonal_indices, list(range(len(indices))), get_rank(expr))
|
||||
for i, index in zip(ind_pos, indices):
|
||||
if isinstance(i, collections.abc.Iterable):
|
||||
for j in i:
|
||||
new_indices[j] = index
|
||||
else:
|
||||
new_indices[i] = index
|
||||
newexpr = self.do_convert(expr.expr, new_indices)
|
||||
return newexpr
|
||||
if isinstance(expr, PermuteDims):
|
||||
permuted_indices = _apply_permutation_to_list(expr.permutation, indices)
|
||||
return self.do_convert(expr.expr, permuted_indices)
|
||||
if isinstance(expr, ArrayAdd):
|
||||
return Add.fromiter(self.do_convert(arg, indices) for arg in expr.args)
|
||||
if isinstance(expr, _ArrayExpr):
|
||||
return expr.__getitem__(tuple(indices))
|
||||
if isinstance(expr, ArrayElementwiseApplyFunc):
|
||||
return expr.function(self.do_convert(expr.expr, indices))
|
||||
if isinstance(expr, Reshape):
|
||||
shape_up = expr.shape
|
||||
shape_down = get_shape(expr.expr)
|
||||
cumul = list(accumulate([1] + list(reversed(shape_up)), operator.mul))
|
||||
one_index = Add.fromiter(i*s for i, s in zip(reversed(indices), cumul))
|
||||
dest_indices = [None for _ in shape_down]
|
||||
c = 1
|
||||
for i, e in enumerate(reversed(shape_down)):
|
||||
if c == 1:
|
||||
if i == len(shape_down) - 1:
|
||||
dest_indices[i] = one_index
|
||||
else:
|
||||
dest_indices[i] = one_index % e
|
||||
elif i == len(shape_down) - 1:
|
||||
dest_indices[i] = one_index // c
|
||||
else:
|
||||
dest_indices[i] = one_index // c % e
|
||||
c *= e
|
||||
dest_indices.reverse()
|
||||
return self.do_convert(expr.expr, dest_indices)
|
||||
return _get_array_element_or_slice(expr, indices)
|
||||
+1004
File diff suppressed because it is too large
Load Diff
+257
@@ -0,0 +1,257 @@
|
||||
from collections import defaultdict
|
||||
|
||||
from sympy import Function
|
||||
from sympy.combinatorics.permutations import _af_invert
|
||||
from sympy.concrete.summations import Sum
|
||||
from sympy.core.add import Add
|
||||
from sympy.core.mul import Mul
|
||||
from sympy.core.numbers import Integer
|
||||
from sympy.core.power import Pow
|
||||
from sympy.core.sorting import default_sort_key
|
||||
from sympy.functions.special.tensor_functions import KroneckerDelta
|
||||
from sympy.tensor.array.expressions import ArrayElementwiseApplyFunc
|
||||
from sympy.tensor.indexed import (Indexed, IndexedBase)
|
||||
from sympy.combinatorics import Permutation
|
||||
from sympy.matrices.expressions.matexpr import MatrixElement
|
||||
from sympy.tensor.array.expressions.array_expressions import ArrayDiagonal, \
|
||||
get_shape, ArrayElement, _array_tensor_product, _array_diagonal, _array_contraction, _array_add, \
|
||||
_permute_dims, OneArray, ArrayAdd
|
||||
from sympy.tensor.array.expressions.utils import _get_argindex, _get_diagonal_indices
|
||||
|
||||
|
||||
def convert_indexed_to_array(expr, first_indices=None):
|
||||
r"""
|
||||
Parse indexed expression into a form useful for code generation.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy.tensor.array.expressions.from_indexed_to_array import convert_indexed_to_array
|
||||
>>> from sympy import MatrixSymbol, Sum, symbols
|
||||
|
||||
>>> i, j, k, d = symbols("i j k d")
|
||||
>>> M = MatrixSymbol("M", d, d)
|
||||
>>> N = MatrixSymbol("N", d, d)
|
||||
|
||||
Recognize the trace in summation form:
|
||||
|
||||
>>> expr = Sum(M[i, i], (i, 0, d-1))
|
||||
>>> convert_indexed_to_array(expr)
|
||||
ArrayContraction(M, (0, 1))
|
||||
|
||||
Recognize the extraction of the diagonal by using the same index `i` on
|
||||
both axes of the matrix:
|
||||
|
||||
>>> expr = M[i, i]
|
||||
>>> convert_indexed_to_array(expr)
|
||||
ArrayDiagonal(M, (0, 1))
|
||||
|
||||
This function can help perform the transformation expressed in two
|
||||
different mathematical notations as:
|
||||
|
||||
`\sum_{j=0}^{N-1} A_{i,j} B_{j,k} \Longrightarrow \mathbf{A}\cdot \mathbf{B}`
|
||||
|
||||
Recognize the matrix multiplication in summation form:
|
||||
|
||||
>>> expr = Sum(M[i, j]*N[j, k], (j, 0, d-1))
|
||||
>>> convert_indexed_to_array(expr)
|
||||
ArrayContraction(ArrayTensorProduct(M, N), (1, 2))
|
||||
|
||||
Specify that ``k`` has to be the starting index:
|
||||
|
||||
>>> convert_indexed_to_array(expr, first_indices=[k])
|
||||
ArrayContraction(ArrayTensorProduct(N, M), (0, 3))
|
||||
"""
|
||||
|
||||
result, indices = _convert_indexed_to_array(expr)
|
||||
|
||||
if any(isinstance(i, (int, Integer)) for i in indices):
|
||||
result = ArrayElement(result, indices)
|
||||
indices = []
|
||||
|
||||
if not first_indices:
|
||||
return result
|
||||
|
||||
def _check_is_in(elem, indices):
|
||||
if elem in indices:
|
||||
return True
|
||||
if any(elem in i for i in indices if isinstance(i, frozenset)):
|
||||
return True
|
||||
return False
|
||||
|
||||
repl = {j: i for i in indices if isinstance(i, frozenset) for j in i}
|
||||
first_indices = [repl.get(i, i) for i in first_indices]
|
||||
for i in first_indices:
|
||||
if not _check_is_in(i, indices):
|
||||
first_indices.remove(i)
|
||||
first_indices.extend([i for i in indices if not _check_is_in(i, first_indices)])
|
||||
|
||||
def _get_pos(elem, indices):
|
||||
if elem in indices:
|
||||
return indices.index(elem)
|
||||
for i, e in enumerate(indices):
|
||||
if not isinstance(e, frozenset):
|
||||
continue
|
||||
if elem in e:
|
||||
return i
|
||||
raise ValueError("not found")
|
||||
|
||||
permutation = _af_invert([_get_pos(i, first_indices) for i in indices])
|
||||
if isinstance(result, ArrayAdd):
|
||||
return _array_add(*[_permute_dims(arg, permutation) for arg in result.args])
|
||||
else:
|
||||
return _permute_dims(result, permutation)
|
||||
|
||||
|
||||
def _convert_indexed_to_array(expr):
|
||||
if isinstance(expr, Sum):
|
||||
function = expr.function
|
||||
summation_indices = expr.variables
|
||||
subexpr, subindices = _convert_indexed_to_array(function)
|
||||
subindicessets = {j: i for i in subindices if isinstance(i, frozenset) for j in i}
|
||||
summation_indices = sorted({subindicessets.get(i, i) for i in summation_indices}, key=default_sort_key)
|
||||
# TODO: check that Kronecker delta is only contracted to one other element:
|
||||
kronecker_indices = set()
|
||||
if isinstance(function, Mul):
|
||||
for arg in function.args:
|
||||
if not isinstance(arg, KroneckerDelta):
|
||||
continue
|
||||
arg_indices = sorted(set(arg.indices), key=default_sort_key)
|
||||
if len(arg_indices) == 2:
|
||||
kronecker_indices.update(arg_indices)
|
||||
kronecker_indices = sorted(kronecker_indices, key=default_sort_key)
|
||||
# Check dimensional consistency:
|
||||
shape = get_shape(subexpr)
|
||||
if shape:
|
||||
for ind, istart, iend in expr.limits:
|
||||
i = _get_argindex(subindices, ind)
|
||||
if istart != 0 or iend+1 != shape[i]:
|
||||
raise ValueError("summation index and array dimension mismatch: %s" % ind)
|
||||
contraction_indices = []
|
||||
subindices = list(subindices)
|
||||
if isinstance(subexpr, ArrayDiagonal):
|
||||
diagonal_indices = list(subexpr.diagonal_indices)
|
||||
dindices = subindices[-len(diagonal_indices):]
|
||||
subindices = subindices[:-len(diagonal_indices)]
|
||||
for index in summation_indices:
|
||||
if index in dindices:
|
||||
position = dindices.index(index)
|
||||
contraction_indices.append(diagonal_indices[position])
|
||||
diagonal_indices[position] = None
|
||||
diagonal_indices = [i for i in diagonal_indices if i is not None]
|
||||
for i, ind in enumerate(subindices):
|
||||
if ind in summation_indices:
|
||||
pass
|
||||
if diagonal_indices:
|
||||
subexpr = _array_diagonal(subexpr.expr, *diagonal_indices)
|
||||
else:
|
||||
subexpr = subexpr.expr
|
||||
|
||||
axes_contraction = defaultdict(list)
|
||||
for i, ind in enumerate(subindices):
|
||||
include = all(j not in kronecker_indices for j in ind) if isinstance(ind, frozenset) else ind not in kronecker_indices
|
||||
if ind in summation_indices and include:
|
||||
axes_contraction[ind].append(i)
|
||||
subindices[i] = None
|
||||
for k, v in axes_contraction.items():
|
||||
if any(i in kronecker_indices for i in k) if isinstance(k, frozenset) else k in kronecker_indices:
|
||||
continue
|
||||
contraction_indices.append(tuple(v))
|
||||
free_indices = [i for i in subindices if i is not None]
|
||||
indices_ret = list(free_indices)
|
||||
indices_ret.sort(key=lambda x: free_indices.index(x))
|
||||
return _array_contraction(
|
||||
subexpr,
|
||||
*contraction_indices,
|
||||
free_indices=free_indices
|
||||
), tuple(indices_ret)
|
||||
if isinstance(expr, Mul):
|
||||
args, indices = zip(*[_convert_indexed_to_array(arg) for arg in expr.args])
|
||||
# Check if there are KroneckerDelta objects:
|
||||
kronecker_delta_repl = {}
|
||||
for arg in args:
|
||||
if not isinstance(arg, KroneckerDelta):
|
||||
continue
|
||||
# Diagonalize two indices:
|
||||
i, j = arg.indices
|
||||
kindices = set(arg.indices)
|
||||
if i in kronecker_delta_repl:
|
||||
kindices.update(kronecker_delta_repl[i])
|
||||
if j in kronecker_delta_repl:
|
||||
kindices.update(kronecker_delta_repl[j])
|
||||
kindices = frozenset(kindices)
|
||||
for index in kindices:
|
||||
kronecker_delta_repl[index] = kindices
|
||||
# Remove KroneckerDelta objects, their relations should be handled by
|
||||
# ArrayDiagonal:
|
||||
newargs = []
|
||||
newindices = []
|
||||
for arg, loc_indices in zip(args, indices):
|
||||
if isinstance(arg, KroneckerDelta):
|
||||
continue
|
||||
newargs.append(arg)
|
||||
newindices.append(loc_indices)
|
||||
flattened_indices = [kronecker_delta_repl.get(j, j) for i in newindices for j in i]
|
||||
diagonal_indices, ret_indices = _get_diagonal_indices(flattened_indices)
|
||||
tp = _array_tensor_product(*newargs)
|
||||
if diagonal_indices:
|
||||
return _array_diagonal(tp, *diagonal_indices), ret_indices
|
||||
else:
|
||||
return tp, ret_indices
|
||||
if isinstance(expr, MatrixElement):
|
||||
indices = expr.args[1:]
|
||||
diagonal_indices, ret_indices = _get_diagonal_indices(indices)
|
||||
if diagonal_indices:
|
||||
return _array_diagonal(expr.args[0], *diagonal_indices), ret_indices
|
||||
else:
|
||||
return expr.args[0], ret_indices
|
||||
if isinstance(expr, ArrayElement):
|
||||
indices = expr.indices
|
||||
diagonal_indices, ret_indices = _get_diagonal_indices(indices)
|
||||
if diagonal_indices:
|
||||
return _array_diagonal(expr.name, *diagonal_indices), ret_indices
|
||||
else:
|
||||
return expr.name, ret_indices
|
||||
if isinstance(expr, Indexed):
|
||||
indices = expr.indices
|
||||
diagonal_indices, ret_indices = _get_diagonal_indices(indices)
|
||||
if diagonal_indices:
|
||||
return _array_diagonal(expr.base, *diagonal_indices), ret_indices
|
||||
else:
|
||||
return expr.args[0], ret_indices
|
||||
if isinstance(expr, IndexedBase):
|
||||
raise NotImplementedError
|
||||
if isinstance(expr, KroneckerDelta):
|
||||
return expr, expr.indices
|
||||
if isinstance(expr, Add):
|
||||
args, indices = zip(*[_convert_indexed_to_array(arg) for arg in expr.args])
|
||||
args = list(args)
|
||||
# Check if all indices are compatible. Otherwise expand the dimensions:
|
||||
index0 = []
|
||||
shape0 = []
|
||||
for arg, arg_indices in zip(args, indices):
|
||||
arg_indices_set = set(arg_indices)
|
||||
arg_indices_missing = arg_indices_set.difference(index0)
|
||||
index0.extend([i for i in arg_indices if i in arg_indices_missing])
|
||||
arg_shape = get_shape(arg)
|
||||
shape0.extend([arg_shape[i] for i, e in enumerate(arg_indices) if e in arg_indices_missing])
|
||||
for i, (arg, arg_indices) in enumerate(zip(args, indices)):
|
||||
if len(arg_indices) < len(index0):
|
||||
missing_indices_pos = [i for i, e in enumerate(index0) if e not in arg_indices]
|
||||
missing_shape = [shape0[i] for i in missing_indices_pos]
|
||||
arg_indices = tuple(index0[j] for j in missing_indices_pos) + arg_indices
|
||||
args[i] = _array_tensor_product(OneArray(*missing_shape), args[i])
|
||||
permutation = Permutation([arg_indices.index(j) for j in index0])
|
||||
# Perform index permutations:
|
||||
args[i] = _permute_dims(args[i], permutation)
|
||||
return _array_add(*args), tuple(index0)
|
||||
if isinstance(expr, Pow):
|
||||
subexpr, subindices = _convert_indexed_to_array(expr.base)
|
||||
if isinstance(expr.exp, (int, Integer)):
|
||||
diags = zip(*[(2*i, 2*i + 1) for i in range(expr.exp)])
|
||||
arr = _array_diagonal(_array_tensor_product(*[subexpr for i in range(expr.exp)]), *diags)
|
||||
return arr, subindices
|
||||
if isinstance(expr, Function):
|
||||
subexpr, subindices = _convert_indexed_to_array(expr.args[0])
|
||||
return ArrayElementwiseApplyFunc(type(expr), subexpr), subindices
|
||||
return expr, ()
|
||||
+87
@@ -0,0 +1,87 @@
|
||||
from sympy import KroneckerProduct
|
||||
from sympy.core.basic import Basic
|
||||
from sympy.core.function import Lambda
|
||||
from sympy.core.mul import Mul
|
||||
from sympy.core.numbers import Integer
|
||||
from sympy.core.power import Pow
|
||||
from sympy.core.singleton import S
|
||||
from sympy.core.symbol import (Dummy, symbols)
|
||||
from sympy.matrices.expressions.hadamard import (HadamardPower, HadamardProduct)
|
||||
from sympy.matrices.expressions.matadd import MatAdd
|
||||
from sympy.matrices.expressions.matmul import MatMul
|
||||
from sympy.matrices.expressions.matpow import MatPow
|
||||
from sympy.matrices.expressions.trace import Trace
|
||||
from sympy.matrices.expressions.transpose import Transpose
|
||||
from sympy.matrices.expressions.matexpr import MatrixExpr
|
||||
from sympy.tensor.array.expressions.array_expressions import \
|
||||
ArrayElementwiseApplyFunc, _array_tensor_product, _array_contraction, \
|
||||
_array_diagonal, _array_add, _permute_dims, Reshape
|
||||
|
||||
|
||||
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
|
||||
+808
@@ -0,0 +1,808 @@
|
||||
import random
|
||||
|
||||
from sympy import tensordiagonal, eye, KroneckerDelta, Array
|
||||
from sympy.core.symbol import symbols
|
||||
from sympy.functions.elementary.trigonometric import (cos, sin)
|
||||
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, Reshape
|
||||
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_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
|
||||
|
||||
|
||||
def test_zero_array():
|
||||
assert ZeroArray() == 0
|
||||
assert ZeroArray().is_Integer
|
||||
|
||||
za = ZeroArray(3, 2, 4)
|
||||
assert za.shape == (3, 2, 4)
|
||||
za_e = za.as_explicit()
|
||||
assert za_e.shape == (3, 2, 4)
|
||||
|
||||
m, n, k = symbols("m n k")
|
||||
za = ZeroArray(m, n, k, 2)
|
||||
assert za.shape == (m, n, k, 2)
|
||||
raises(ValueError, lambda: za.as_explicit())
|
||||
|
||||
|
||||
def test_one_array():
|
||||
assert OneArray() == 1
|
||||
assert OneArray().is_Integer
|
||||
|
||||
oa = OneArray(3, 2, 4)
|
||||
assert oa.shape == (3, 2, 4)
|
||||
oa_e = oa.as_explicit()
|
||||
assert oa_e.shape == (3, 2, 4)
|
||||
|
||||
m, n, k = symbols("m n k")
|
||||
oa = OneArray(m, n, k, 2)
|
||||
assert oa.shape == (m, n, k, 2)
|
||||
raises(ValueError, lambda: oa.as_explicit())
|
||||
|
||||
|
||||
def test_arrayexpr_contraction_construction():
|
||||
|
||||
cg = _array_contraction(A)
|
||||
assert cg == A
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(A, B), (1, 0))
|
||||
assert cg == _array_contraction(_array_tensor_product(A, B), (0, 1))
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(M, N), (0, 1))
|
||||
indtup = cg._get_contraction_tuples()
|
||||
assert indtup == [[(0, 0), (0, 1)]]
|
||||
assert cg._contraction_tuples_to_contraction_indices(cg.expr, indtup) == [(0, 1)]
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(M, N), (1, 2))
|
||||
indtup = cg._get_contraction_tuples()
|
||||
assert indtup == [[(0, 1), (1, 0)]]
|
||||
assert cg._contraction_tuples_to_contraction_indices(cg.expr, indtup) == [(1, 2)]
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(M, M, N), (1, 4), (2, 5))
|
||||
indtup = cg._get_contraction_tuples()
|
||||
assert indtup == [[(0, 0), (1, 1)], [(0, 1), (2, 0)]]
|
||||
assert cg._contraction_tuples_to_contraction_indices(cg.expr, indtup) == [(0, 3), (1, 4)]
|
||||
|
||||
# Test removal of trivial contraction:
|
||||
assert _array_contraction(a, (1,)) == a
|
||||
assert _array_contraction(
|
||||
_array_tensor_product(a, b), (0, 2), (1,), (3,)) == _array_contraction(
|
||||
_array_tensor_product(a, b), (0, 2))
|
||||
|
||||
|
||||
def test_arrayexpr_array_flatten():
|
||||
|
||||
# Flatten nested ArrayTensorProduct objects:
|
||||
expr1 = _array_tensor_product(M, N)
|
||||
expr2 = _array_tensor_product(P, Q)
|
||||
expr = _array_tensor_product(expr1, expr2)
|
||||
assert expr == _array_tensor_product(M, N, P, Q)
|
||||
assert expr.args == (M, N, P, Q)
|
||||
|
||||
# Flatten mixed ArrayTensorProduct and ArrayContraction objects:
|
||||
cg1 = _array_contraction(expr1, (1, 2))
|
||||
cg2 = _array_contraction(expr2, (0, 3))
|
||||
|
||||
expr = _array_tensor_product(cg1, cg2)
|
||||
assert expr == _array_contraction(_array_tensor_product(M, N, P, Q), (1, 2), (4, 7))
|
||||
|
||||
expr = _array_tensor_product(M, cg1)
|
||||
assert expr == _array_contraction(_array_tensor_product(M, M, N), (3, 4))
|
||||
|
||||
# Flatten nested ArrayContraction objects:
|
||||
cgnested = _array_contraction(cg1, (0, 1))
|
||||
assert cgnested == _array_contraction(_array_tensor_product(M, N), (0, 3), (1, 2))
|
||||
|
||||
cgnested = _array_contraction(_array_tensor_product(cg1, cg2), (0, 3))
|
||||
assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 6), (1, 2), (4, 7))
|
||||
|
||||
cg3 = _array_contraction(_array_tensor_product(M, N, P, Q), (1, 3), (2, 4))
|
||||
cgnested = _array_contraction(cg3, (0, 1))
|
||||
assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 5), (1, 3), (2, 4))
|
||||
|
||||
cgnested = _array_contraction(cg3, (0, 3), (1, 2))
|
||||
assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 7), (1, 3), (2, 4), (5, 6))
|
||||
|
||||
cg4 = _array_contraction(_array_tensor_product(M, N, P, Q), (1, 5), (3, 7))
|
||||
cgnested = _array_contraction(cg4, (0, 1))
|
||||
assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 2), (1, 5), (3, 7))
|
||||
|
||||
cgnested = _array_contraction(cg4, (0, 1), (2, 3))
|
||||
assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 2), (1, 5), (3, 7), (4, 6))
|
||||
|
||||
cg = _array_diagonal(cg4)
|
||||
assert cg == cg4
|
||||
assert isinstance(cg, type(cg4))
|
||||
|
||||
# Flatten nested ArrayDiagonal objects:
|
||||
cg1 = _array_diagonal(expr1, (1, 2))
|
||||
cg2 = _array_diagonal(expr2, (0, 3))
|
||||
cg3 = _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 3), (2, 4))
|
||||
cg4 = _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 5), (3, 7))
|
||||
|
||||
cgnested = _array_diagonal(cg1, (0, 1))
|
||||
assert cgnested == _array_diagonal(_array_tensor_product(M, N), (1, 2), (0, 3))
|
||||
|
||||
cgnested = _array_diagonal(cg3, (1, 2))
|
||||
assert cgnested == _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 3), (2, 4), (5, 6))
|
||||
|
||||
cgnested = _array_diagonal(cg4, (1, 2))
|
||||
assert cgnested == _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 5), (3, 7), (2, 4))
|
||||
|
||||
cg = _array_add(M, N)
|
||||
cg2 = _array_add(cg, P)
|
||||
assert isinstance(cg2, ArrayAdd)
|
||||
assert cg2.args == (M, N, P)
|
||||
assert cg2.shape == (k, k)
|
||||
|
||||
expr = _array_tensor_product(_array_diagonal(X, (0, 1)), _array_diagonal(A, (0, 1)))
|
||||
assert expr == _array_diagonal(_array_tensor_product(X, A), (0, 1), (2, 3))
|
||||
|
||||
expr1 = _array_diagonal(_array_tensor_product(X, A), (1, 2))
|
||||
expr2 = _array_tensor_product(expr1, a)
|
||||
assert expr2 == _permute_dims(_array_diagonal(_array_tensor_product(X, A, a), (1, 2)), [0, 1, 4, 2, 3])
|
||||
|
||||
expr1 = _array_contraction(_array_tensor_product(X, A), (1, 2))
|
||||
expr2 = _array_tensor_product(expr1, a)
|
||||
assert isinstance(expr2, ArrayContraction)
|
||||
assert isinstance(expr2.expr, ArrayTensorProduct)
|
||||
|
||||
cg = _array_tensor_product(_array_diagonal(_array_tensor_product(A, X, Y), (0, 3), (1, 5)), a, b)
|
||||
assert cg == _permute_dims(_array_diagonal(_array_tensor_product(A, X, Y, a, b), (0, 3), (1, 5)), [0, 1, 6, 7, 2, 3, 4, 5])
|
||||
|
||||
|
||||
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_arrayexpr_array_shape():
|
||||
expr = _array_tensor_product(M, N, P, Q)
|
||||
assert expr.shape == (k, k, k, k, k, k, k, k)
|
||||
Z = MatrixSymbol("Z", m, n)
|
||||
expr = _array_tensor_product(M, Z)
|
||||
assert expr.shape == (k, k, m, n)
|
||||
expr2 = _array_contraction(expr, (0, 1))
|
||||
assert expr2.shape == (m, n)
|
||||
expr2 = _array_diagonal(expr, (0, 1))
|
||||
assert expr2.shape == (m, n, k)
|
||||
exprp = _permute_dims(expr, [2, 1, 3, 0])
|
||||
assert exprp.shape == (m, k, n, k)
|
||||
expr3 = _array_tensor_product(N, Z)
|
||||
expr2 = _array_add(expr, expr3)
|
||||
assert expr2.shape == (k, k, m, n)
|
||||
|
||||
# Contraction along axes with discordant dimensions:
|
||||
raises(ValueError, lambda: _array_contraction(expr, (1, 2)))
|
||||
# Also diagonal needs the same dimensions:
|
||||
raises(ValueError, lambda: _array_diagonal(expr, (1, 2)))
|
||||
# Diagonal requires at least to axes to compute the diagonal:
|
||||
raises(ValueError, lambda: _array_diagonal(expr, (1,)))
|
||||
|
||||
|
||||
def test_arrayexpr_permutedims_sink():
|
||||
|
||||
cg = _permute_dims(_array_tensor_product(M, N), [0, 1, 3, 2], nest_permutation=False)
|
||||
sunk = nest_permutation(cg)
|
||||
assert sunk == _array_tensor_product(M, _permute_dims(N, [1, 0]))
|
||||
|
||||
cg = _permute_dims(_array_tensor_product(M, N), [1, 0, 3, 2], nest_permutation=False)
|
||||
sunk = nest_permutation(cg)
|
||||
assert sunk == _array_tensor_product(_permute_dims(M, [1, 0]), _permute_dims(N, [1, 0]))
|
||||
|
||||
cg = _permute_dims(_array_tensor_product(M, N), [3, 2, 1, 0], nest_permutation=False)
|
||||
sunk = nest_permutation(cg)
|
||||
assert sunk == _array_tensor_product(_permute_dims(N, [1, 0]), _permute_dims(M, [1, 0]))
|
||||
|
||||
cg = _permute_dims(_array_contraction(_array_tensor_product(M, N), (1, 2)), [1, 0], nest_permutation=False)
|
||||
sunk = nest_permutation(cg)
|
||||
assert sunk == _array_contraction(_permute_dims(_array_tensor_product(M, N), [[0, 3]]), (1, 2))
|
||||
|
||||
cg = _permute_dims(_array_tensor_product(M, N), [1, 0, 3, 2], nest_permutation=False)
|
||||
sunk = nest_permutation(cg)
|
||||
assert sunk == _array_tensor_product(_permute_dims(M, [1, 0]), _permute_dims(N, [1, 0]))
|
||||
|
||||
cg = _permute_dims(_array_contraction(_array_tensor_product(M, N, P), (1, 2), (3, 4)), [1, 0], nest_permutation=False)
|
||||
sunk = nest_permutation(cg)
|
||||
assert sunk == _array_contraction(_permute_dims(_array_tensor_product(M, N, P), [[0, 5]]), (1, 2), (3, 4))
|
||||
|
||||
|
||||
def test_arrayexpr_push_indices_up_and_down():
|
||||
|
||||
indices = list(range(12))
|
||||
|
||||
contr_diag_indices = [(0, 6), (2, 8)]
|
||||
assert ArrayContraction._push_indices_down(contr_diag_indices, indices) == (1, 3, 4, 5, 7, 9, 10, 11, 12, 13, 14, 15)
|
||||
assert ArrayContraction._push_indices_up(contr_diag_indices, indices) == (None, 0, None, 1, 2, 3, None, 4, None, 5, 6, 7)
|
||||
|
||||
assert ArrayDiagonal._push_indices_down(contr_diag_indices, indices, 10) == (1, 3, 4, 5, 7, 9, (0, 6), (2, 8), None, None, None, None)
|
||||
assert ArrayDiagonal._push_indices_up(contr_diag_indices, indices, 10) == (6, 0, 7, 1, 2, 3, 6, 4, 7, 5, None, None)
|
||||
|
||||
contr_diag_indices = [(1, 2), (7, 8)]
|
||||
assert ArrayContraction._push_indices_down(contr_diag_indices, indices) == (0, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15)
|
||||
assert ArrayContraction._push_indices_up(contr_diag_indices, indices) == (0, None, None, 1, 2, 3, 4, None, None, 5, 6, 7)
|
||||
|
||||
assert ArrayDiagonal._push_indices_down(contr_diag_indices, indices, 10) == (0, 3, 4, 5, 6, 9, (1, 2), (7, 8), None, None, None, None)
|
||||
assert ArrayDiagonal._push_indices_up(contr_diag_indices, indices, 10) == (0, 6, 6, 1, 2, 3, 4, 7, 7, 5, None, None)
|
||||
|
||||
|
||||
def test_arrayexpr_split_multiple_contractions():
|
||||
a = MatrixSymbol("a", k, 1)
|
||||
b = MatrixSymbol("b", k, 1)
|
||||
A = MatrixSymbol("A", k, k)
|
||||
B = MatrixSymbol("B", k, k)
|
||||
C = MatrixSymbol("C", k, k)
|
||||
X = MatrixSymbol("X", k, k)
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(A.T, a, b, b.T, (A*X*b).applyfunc(cos)), (1, 2, 8), (5, 6, 9))
|
||||
expected = _array_contraction(_array_tensor_product(A.T, DiagMatrix(a), OneArray(1), b, b.T, (A*X*b).applyfunc(cos)), (1, 3), (2, 9), (6, 7, 10))
|
||||
assert cg.split_multiple_contractions().dummy_eq(expected)
|
||||
|
||||
# Check no overlap of lines:
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(A, a, C, a, B), (1, 2, 4), (5, 6, 8), (3, 7))
|
||||
assert cg.split_multiple_contractions() == cg
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(a, b, A), (0, 2, 4), (1, 3))
|
||||
assert cg.split_multiple_contractions() == cg
|
||||
|
||||
|
||||
def test_arrayexpr_nested_permutations():
|
||||
|
||||
cg = _permute_dims(_permute_dims(M, (1, 0)), (1, 0))
|
||||
assert cg == M
|
||||
|
||||
times = 3
|
||||
plist1 = [list(range(6)) for i in range(times)]
|
||||
plist2 = [list(range(6)) for i in range(times)]
|
||||
|
||||
for i in range(times):
|
||||
random.shuffle(plist1[i])
|
||||
random.shuffle(plist2[i])
|
||||
|
||||
plist1.append([2, 5, 4, 1, 0, 3])
|
||||
plist2.append([3, 5, 0, 4, 1, 2])
|
||||
|
||||
plist1.append([2, 5, 4, 0, 3, 1])
|
||||
plist2.append([3, 0, 5, 1, 2, 4])
|
||||
|
||||
plist1.append([5, 4, 2, 0, 3, 1])
|
||||
plist2.append([4, 5, 0, 2, 3, 1])
|
||||
|
||||
Me = M.subs(k, 3).as_explicit()
|
||||
Ne = N.subs(k, 3).as_explicit()
|
||||
Pe = P.subs(k, 3).as_explicit()
|
||||
cge = tensorproduct(Me, Ne, Pe)
|
||||
|
||||
for permutation_array1, permutation_array2 in zip(plist1, plist2):
|
||||
p1 = Permutation(permutation_array1)
|
||||
p2 = Permutation(permutation_array2)
|
||||
|
||||
cg = _permute_dims(
|
||||
_permute_dims(
|
||||
_array_tensor_product(M, N, P),
|
||||
p1),
|
||||
p2
|
||||
)
|
||||
result = _permute_dims(
|
||||
_array_tensor_product(M, N, P),
|
||||
p2*p1
|
||||
)
|
||||
assert cg == result
|
||||
|
||||
# Check that `permutedims` behaves the same way with explicit-component arrays:
|
||||
result1 = _permute_dims(_permute_dims(cge, p1), p2)
|
||||
result2 = _permute_dims(cge, p2*p1)
|
||||
assert result1 == result2
|
||||
|
||||
|
||||
def test_arrayexpr_contraction_permutation_mix():
|
||||
|
||||
Me = M.subs(k, 3).as_explicit()
|
||||
Ne = N.subs(k, 3).as_explicit()
|
||||
|
||||
cg1 = _array_contraction(PermuteDims(_array_tensor_product(M, N), Permutation([0, 2, 1, 3])), (2, 3))
|
||||
cg2 = _array_contraction(_array_tensor_product(M, N), (1, 3))
|
||||
assert cg1 == cg2
|
||||
cge1 = tensorcontraction(permutedims(tensorproduct(Me, Ne), Permutation([0, 2, 1, 3])), (2, 3))
|
||||
cge2 = tensorcontraction(tensorproduct(Me, Ne), (1, 3))
|
||||
assert cge1 == cge2
|
||||
|
||||
cg1 = _permute_dims(_array_tensor_product(M, N), Permutation([0, 1, 3, 2]))
|
||||
cg2 = _array_tensor_product(M, _permute_dims(N, Permutation([1, 0])))
|
||||
assert cg1 == cg2
|
||||
|
||||
cg1 = _array_contraction(
|
||||
_permute_dims(
|
||||
_array_tensor_product(M, N, P, Q), Permutation([0, 2, 3, 1, 4, 5, 7, 6])),
|
||||
(1, 2), (3, 5)
|
||||
)
|
||||
cg2 = _array_contraction(
|
||||
_array_tensor_product(M, N, P, _permute_dims(Q, Permutation([1, 0]))),
|
||||
(1, 5), (2, 3)
|
||||
)
|
||||
assert cg1 == cg2
|
||||
|
||||
cg1 = _array_contraction(
|
||||
_permute_dims(
|
||||
_array_tensor_product(M, N, P, Q), Permutation([1, 0, 4, 6, 2, 7, 5, 3])),
|
||||
(0, 1), (2, 6), (3, 7)
|
||||
)
|
||||
cg2 = _permute_dims(
|
||||
_array_contraction(
|
||||
_array_tensor_product(M, P, Q, N),
|
||||
(0, 1), (2, 3), (4, 7)),
|
||||
[1, 0]
|
||||
)
|
||||
assert cg1 == cg2
|
||||
|
||||
cg1 = _array_contraction(
|
||||
_permute_dims(
|
||||
_array_tensor_product(M, N, P, Q), Permutation([1, 0, 4, 6, 7, 2, 5, 3])),
|
||||
(0, 1), (2, 6), (3, 7)
|
||||
)
|
||||
cg2 = _permute_dims(
|
||||
_array_contraction(
|
||||
_array_tensor_product(_permute_dims(M, [1, 0]), N, P, Q),
|
||||
(0, 1), (3, 6), (4, 5)
|
||||
),
|
||||
Permutation([1, 0])
|
||||
)
|
||||
assert cg1 == cg2
|
||||
|
||||
|
||||
def test_arrayexpr_permute_tensor_product():
|
||||
cg1 = _permute_dims(_array_tensor_product(M, N, P, Q), Permutation([2, 3, 1, 0, 5, 4, 6, 7]))
|
||||
cg2 = _array_tensor_product(N, _permute_dims(M, [1, 0]),
|
||||
_permute_dims(P, [1, 0]), Q)
|
||||
assert cg1 == cg2
|
||||
|
||||
# TODO: reverse operation starting with `PermuteDims` and getting down to `bb`...
|
||||
cg1 = _permute_dims(_array_tensor_product(M, N, P, Q), Permutation([2, 3, 4, 5, 0, 1, 6, 7]))
|
||||
cg2 = _array_tensor_product(N, P, M, Q)
|
||||
assert cg1 == cg2
|
||||
|
||||
cg1 = _permute_dims(_array_tensor_product(M, N, P, Q), Permutation([2, 3, 4, 6, 5, 7, 0, 1]))
|
||||
assert cg1.expr == _array_tensor_product(N, P, Q, M)
|
||||
assert cg1.permutation == Permutation([0, 1, 2, 4, 3, 5, 6, 7])
|
||||
|
||||
cg1 = _array_contraction(
|
||||
_permute_dims(
|
||||
_array_tensor_product(N, Q, Q, M),
|
||||
[2, 1, 5, 4, 0, 3, 6, 7]),
|
||||
[1, 2, 6])
|
||||
cg2 = _permute_dims(_array_contraction(_array_tensor_product(Q, Q, N, M), (3, 5, 6)), [0, 2, 3, 1, 4])
|
||||
assert cg1 == cg2
|
||||
|
||||
cg1 = _array_contraction(
|
||||
_array_contraction(
|
||||
_array_contraction(
|
||||
_array_contraction(
|
||||
_permute_dims(
|
||||
_array_tensor_product(N, Q, Q, M),
|
||||
[2, 1, 5, 4, 0, 3, 6, 7]),
|
||||
[1, 2, 6]),
|
||||
[1, 3, 4]),
|
||||
[1]),
|
||||
[0])
|
||||
cg2 = _array_contraction(_array_tensor_product(M, N, Q, Q), (0, 3, 5), (1, 4, 7), (2,), (6,))
|
||||
assert cg1 == cg2
|
||||
|
||||
|
||||
def test_arrayexpr_canonicalize_diagonal__permute_dims():
|
||||
tp = _array_tensor_product(M, Q, N, P)
|
||||
expr = _array_diagonal(
|
||||
_permute_dims(tp, [0, 1, 2, 4, 7, 6, 3, 5]), (2, 4, 5), (6, 7),
|
||||
(0, 3))
|
||||
result = _array_diagonal(tp, (2, 6, 7), (3, 5), (0, 4))
|
||||
assert expr == result
|
||||
|
||||
tp = _array_tensor_product(M, N, P, Q)
|
||||
expr = _array_diagonal(_permute_dims(tp, [0, 5, 2, 4, 1, 6, 3, 7]), (1, 2, 6), (3, 4))
|
||||
result = _array_diagonal(_array_tensor_product(M, P, N, Q), (3, 4, 5), (1, 2))
|
||||
assert expr == result
|
||||
|
||||
|
||||
def test_arrayexpr_canonicalize_diagonal_contraction():
|
||||
tp = _array_tensor_product(M, N, P, Q)
|
||||
expr = _array_contraction(_array_diagonal(tp, (1, 3, 4)), (0, 3))
|
||||
result = _array_diagonal(_array_contraction(_array_tensor_product(M, N, P, Q), (0, 6)), (0, 2, 3))
|
||||
assert expr == result
|
||||
|
||||
expr = _array_contraction(_array_diagonal(tp, (0, 1, 2, 3, 7)), (1, 2, 3))
|
||||
result = _array_contraction(_array_tensor_product(M, N, P, Q), (0, 1, 2, 3, 5, 6, 7))
|
||||
assert expr == result
|
||||
|
||||
expr = _array_contraction(_array_diagonal(tp, (0, 2, 6, 7)), (1, 2, 3))
|
||||
result = _array_diagonal(_array_contraction(tp, (3, 4, 5)), (0, 2, 3, 4))
|
||||
assert expr == result
|
||||
|
||||
td = _array_diagonal(_array_tensor_product(M, N, P, Q), (0, 3))
|
||||
expr = _array_contraction(td, (2, 1), (0, 4, 6, 5, 3))
|
||||
result = _array_contraction(_array_tensor_product(M, N, P, Q), (0, 1, 3, 5, 6, 7), (2, 4))
|
||||
assert expr == result
|
||||
|
||||
|
||||
def test_arrayexpr_array_wrong_permutation_size():
|
||||
cg = _array_tensor_product(M, N)
|
||||
raises(ValueError, lambda: _permute_dims(cg, [1, 0]))
|
||||
raises(ValueError, lambda: _permute_dims(cg, [1, 0, 2, 3, 5, 4]))
|
||||
|
||||
|
||||
def test_arrayexpr_nested_array_elementwise_add():
|
||||
cg = _array_contraction(_array_add(
|
||||
_array_tensor_product(M, N),
|
||||
_array_tensor_product(N, M)
|
||||
), (1, 2))
|
||||
result = _array_add(
|
||||
_array_contraction(_array_tensor_product(M, N), (1, 2)),
|
||||
_array_contraction(_array_tensor_product(N, M), (1, 2))
|
||||
)
|
||||
assert cg == result
|
||||
|
||||
cg = _array_diagonal(_array_add(
|
||||
_array_tensor_product(M, N),
|
||||
_array_tensor_product(N, M)
|
||||
), (1, 2))
|
||||
result = _array_add(
|
||||
_array_diagonal(_array_tensor_product(M, N), (1, 2)),
|
||||
_array_diagonal(_array_tensor_product(N, M), (1, 2))
|
||||
)
|
||||
assert cg == result
|
||||
|
||||
|
||||
def test_arrayexpr_array_expr_zero_array():
|
||||
za1 = ZeroArray(k, l, m, n)
|
||||
zm1 = ZeroMatrix(m, n)
|
||||
|
||||
za2 = ZeroArray(k, m, m, n)
|
||||
zm2 = ZeroMatrix(m, m)
|
||||
zm3 = ZeroMatrix(k, k)
|
||||
|
||||
assert _array_tensor_product(M, N, za1) == ZeroArray(k, k, k, k, k, l, m, n)
|
||||
assert _array_tensor_product(M, N, zm1) == ZeroArray(k, k, k, k, m, n)
|
||||
|
||||
assert _array_contraction(za1, (3,)) == ZeroArray(k, l, m)
|
||||
assert _array_contraction(zm1, (1,)) == ZeroArray(m)
|
||||
assert _array_contraction(za2, (1, 2)) == ZeroArray(k, n)
|
||||
assert _array_contraction(zm2, (0, 1)) == 0
|
||||
|
||||
assert _array_diagonal(za2, (1, 2)) == ZeroArray(k, n, m)
|
||||
assert _array_diagonal(zm2, (0, 1)) == ZeroArray(m)
|
||||
|
||||
assert _permute_dims(za1, [2, 1, 3, 0]) == ZeroArray(m, l, n, k)
|
||||
assert _permute_dims(zm1, [1, 0]) == ZeroArray(n, m)
|
||||
|
||||
assert _array_add(za1) == za1
|
||||
assert _array_add(zm1) == ZeroArray(m, n)
|
||||
tp1 = _array_tensor_product(MatrixSymbol("A", k, l), MatrixSymbol("B", m, n))
|
||||
assert _array_add(tp1, za1) == tp1
|
||||
tp2 = _array_tensor_product(MatrixSymbol("C", k, l), MatrixSymbol("D", m, n))
|
||||
assert _array_add(tp1, za1, tp2) == _array_add(tp1, tp2)
|
||||
assert _array_add(M, zm3) == M
|
||||
assert _array_add(M, N, zm3) == _array_add(M, N)
|
||||
|
||||
|
||||
def test_arrayexpr_array_expr_applyfunc():
|
||||
|
||||
A = ArraySymbol("A", (3, k, 2))
|
||||
aaf = ArrayElementwiseApplyFunc(sin, A)
|
||||
assert aaf.shape == (3, k, 2)
|
||||
|
||||
|
||||
def test_edit_array_contraction():
|
||||
cg = _array_contraction(_array_tensor_product(A, B, C, D), (1, 2, 5))
|
||||
ecg = _EditArrayContraction(cg)
|
||||
assert ecg.to_array_contraction() == cg
|
||||
|
||||
ecg.args_with_ind[1], ecg.args_with_ind[2] = ecg.args_with_ind[2], ecg.args_with_ind[1]
|
||||
assert ecg.to_array_contraction() == _array_contraction(_array_tensor_product(A, C, B, D), (1, 3, 4))
|
||||
|
||||
ci = ecg.get_new_contraction_index()
|
||||
new_arg = _ArgE(X)
|
||||
new_arg.indices = [ci, ci]
|
||||
ecg.args_with_ind.insert(2, new_arg)
|
||||
assert ecg.to_array_contraction() == _array_contraction(_array_tensor_product(A, C, X, B, D), (1, 3, 6), (4, 5))
|
||||
|
||||
assert ecg.get_contraction_indices() == [[1, 3, 6], [4, 5]]
|
||||
assert [[tuple(j) for j in i] for i in ecg.get_contraction_indices_to_ind_rel_pos()] == [[(0, 1), (1, 1), (3, 0)], [(2, 0), (2, 1)]]
|
||||
assert [list(i) for i in ecg.get_mapping_for_index(0)] == [[0, 1], [1, 1], [3, 0]]
|
||||
assert [list(i) for i in ecg.get_mapping_for_index(1)] == [[2, 0], [2, 1]]
|
||||
raises(ValueError, lambda: ecg.get_mapping_for_index(2))
|
||||
|
||||
ecg.args_with_ind.pop(1)
|
||||
assert ecg.to_array_contraction() == _array_contraction(_array_tensor_product(A, X, B, D), (1, 4), (2, 3))
|
||||
|
||||
ecg.args_with_ind[0].indices[1] = ecg.args_with_ind[1].indices[0]
|
||||
ecg.args_with_ind[1].indices[1] = ecg.args_with_ind[2].indices[0]
|
||||
assert ecg.to_array_contraction() == _array_contraction(_array_tensor_product(A, X, B, D), (1, 2), (3, 4))
|
||||
|
||||
ecg.insert_after(ecg.args_with_ind[1], _ArgE(C))
|
||||
assert ecg.to_array_contraction() == _array_contraction(_array_tensor_product(A, X, C, B, D), (1, 2), (3, 6))
|
||||
|
||||
|
||||
def test_array_expressions_no_canonicalization():
|
||||
|
||||
tp = _array_tensor_product(M, N, P)
|
||||
|
||||
# ArrayTensorProduct:
|
||||
|
||||
expr = ArrayTensorProduct(tp, N)
|
||||
assert str(expr) == "ArrayTensorProduct(ArrayTensorProduct(M, N, P), N)"
|
||||
assert expr.doit() == ArrayTensorProduct(M, N, P, N)
|
||||
|
||||
expr = ArrayTensorProduct(ArrayContraction(M, (0, 1)), N)
|
||||
assert str(expr) == "ArrayTensorProduct(ArrayContraction(M, (0, 1)), N)"
|
||||
assert expr.doit() == ArrayContraction(ArrayTensorProduct(M, N), (0, 1))
|
||||
|
||||
expr = ArrayTensorProduct(ArrayDiagonal(M, (0, 1)), N)
|
||||
assert str(expr) == "ArrayTensorProduct(ArrayDiagonal(M, (0, 1)), N)"
|
||||
assert expr.doit() == PermuteDims(ArrayDiagonal(ArrayTensorProduct(M, N), (0, 1)), [2, 0, 1])
|
||||
|
||||
expr = ArrayTensorProduct(PermuteDims(M, [1, 0]), N)
|
||||
assert str(expr) == "ArrayTensorProduct(PermuteDims(M, (0 1)), N)"
|
||||
assert expr.doit() == PermuteDims(ArrayTensorProduct(M, N), [1, 0, 2, 3])
|
||||
|
||||
# ArrayContraction:
|
||||
|
||||
expr = ArrayContraction(_array_contraction(tp, (0, 2)), (0, 1))
|
||||
assert isinstance(expr, ArrayContraction)
|
||||
assert isinstance(expr.expr, ArrayContraction)
|
||||
assert str(expr) == "ArrayContraction(ArrayContraction(ArrayTensorProduct(M, N, P), (0, 2)), (0, 1))"
|
||||
assert expr.doit() == ArrayContraction(tp, (0, 2), (1, 3))
|
||||
|
||||
expr = ArrayContraction(ArrayContraction(ArrayContraction(tp, (0, 1)), (0, 1)), (0, 1))
|
||||
assert expr.doit() == ArrayContraction(tp, (0, 1), (2, 3), (4, 5))
|
||||
# assert expr._canonicalize() == ArrayContraction(ArrayContraction(tp, (0, 1)), (0, 1), (2, 3))
|
||||
|
||||
expr = ArrayContraction(ArrayDiagonal(tp, (0, 1)), (0, 1))
|
||||
assert str(expr) == "ArrayContraction(ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 1)), (0, 1))"
|
||||
assert expr.doit() == ArrayDiagonal(ArrayContraction(ArrayTensorProduct(N, M, P), (0, 1)), (0, 1))
|
||||
|
||||
expr = ArrayContraction(PermuteDims(M, [1, 0]), (0, 1))
|
||||
assert str(expr) == "ArrayContraction(PermuteDims(M, (0 1)), (0, 1))"
|
||||
assert expr.doit() == ArrayContraction(M, (0, 1))
|
||||
|
||||
# ArrayDiagonal:
|
||||
|
||||
expr = ArrayDiagonal(ArrayDiagonal(tp, (0, 2)), (0, 1))
|
||||
assert str(expr) == "ArrayDiagonal(ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 2)), (0, 1))"
|
||||
assert expr.doit() == ArrayDiagonal(tp, (0, 2), (1, 3))
|
||||
|
||||
expr = ArrayDiagonal(ArrayDiagonal(ArrayDiagonal(tp, (0, 1)), (0, 1)), (0, 1))
|
||||
assert expr.doit() == ArrayDiagonal(tp, (0, 1), (2, 3), (4, 5))
|
||||
assert expr._canonicalize() == expr.doit()
|
||||
|
||||
expr = ArrayDiagonal(ArrayContraction(tp, (0, 1)), (0, 1))
|
||||
assert str(expr) == "ArrayDiagonal(ArrayContraction(ArrayTensorProduct(M, N, P), (0, 1)), (0, 1))"
|
||||
assert expr.doit() == expr
|
||||
|
||||
expr = ArrayDiagonal(PermuteDims(M, [1, 0]), (0, 1))
|
||||
assert str(expr) == "ArrayDiagonal(PermuteDims(M, (0 1)), (0, 1))"
|
||||
assert expr.doit() == ArrayDiagonal(M, (0, 1))
|
||||
|
||||
# ArrayAdd:
|
||||
|
||||
expr = ArrayAdd(M)
|
||||
assert isinstance(expr, ArrayAdd)
|
||||
assert expr.doit() == M
|
||||
|
||||
expr = ArrayAdd(ArrayAdd(M, N), P)
|
||||
assert str(expr) == "ArrayAdd(ArrayAdd(M, N), P)"
|
||||
assert expr.doit() == ArrayAdd(M, N, P)
|
||||
|
||||
expr = ArrayAdd(M, ArrayAdd(N, ArrayAdd(P, M)))
|
||||
assert expr.doit() == ArrayAdd(M, N, P, M)
|
||||
assert expr._canonicalize() == ArrayAdd(M, N, ArrayAdd(P, M))
|
||||
|
||||
expr = ArrayAdd(M, ZeroArray(k, k), N)
|
||||
assert str(expr) == "ArrayAdd(M, ZeroArray(k, k), N)"
|
||||
assert expr.doit() == ArrayAdd(M, N)
|
||||
|
||||
# PermuteDims:
|
||||
|
||||
expr = PermuteDims(PermuteDims(M, [1, 0]), [1, 0])
|
||||
assert str(expr) == "PermuteDims(PermuteDims(M, (0 1)), (0 1))"
|
||||
assert expr.doit() == M
|
||||
|
||||
expr = PermuteDims(PermuteDims(PermuteDims(M, [1, 0]), [1, 0]), [1, 0])
|
||||
assert expr.doit() == PermuteDims(M, [1, 0])
|
||||
assert expr._canonicalize() == expr.doit()
|
||||
|
||||
# Reshape
|
||||
|
||||
expr = Reshape(A, (k**2,))
|
||||
assert expr.shape == (k**2,)
|
||||
assert isinstance(expr, Reshape)
|
||||
|
||||
|
||||
def test_array_expr_construction_with_functions():
|
||||
|
||||
tp = tensorproduct(M, N)
|
||||
assert tp == ArrayTensorProduct(M, N)
|
||||
|
||||
expr = tensorproduct(A, eye(2))
|
||||
assert expr == ArrayTensorProduct(A, eye(2))
|
||||
|
||||
# Contraction:
|
||||
|
||||
expr = tensorcontraction(M, (0, 1))
|
||||
assert expr == ArrayContraction(M, (0, 1))
|
||||
|
||||
expr = tensorcontraction(tp, (1, 2))
|
||||
assert expr == ArrayContraction(tp, (1, 2))
|
||||
|
||||
expr = tensorcontraction(tensorcontraction(tp, (1, 2)), (0, 1))
|
||||
assert expr == ArrayContraction(tp, (0, 3), (1, 2))
|
||||
|
||||
# Diagonalization:
|
||||
|
||||
expr = tensordiagonal(M, (0, 1))
|
||||
assert expr == ArrayDiagonal(M, (0, 1))
|
||||
|
||||
expr = tensordiagonal(tensordiagonal(tp, (0, 1)), (0, 1))
|
||||
assert expr == ArrayDiagonal(tp, (0, 1), (2, 3))
|
||||
|
||||
# Permutation of dimensions:
|
||||
|
||||
expr = permutedims(M, [1, 0])
|
||||
assert expr == PermuteDims(M, [1, 0])
|
||||
|
||||
expr = permutedims(PermuteDims(tp, [1, 0, 2, 3]), [0, 1, 3, 2])
|
||||
assert expr == PermuteDims(tp, [1, 0, 3, 2])
|
||||
|
||||
expr = PermuteDims(tp, index_order_new=["a", "b", "c", "d"], index_order_old=["d", "c", "b", "a"])
|
||||
assert expr == PermuteDims(tp, [3, 2, 1, 0])
|
||||
|
||||
arr = Array(range(32)).reshape(2, 2, 2, 2, 2)
|
||||
expr = PermuteDims(arr, index_order_new=["a", "b", "c", "d", "e"], index_order_old=['b', 'e', 'a', 'd', 'c'])
|
||||
assert expr == PermuteDims(arr, [2, 0, 4, 3, 1])
|
||||
assert expr.as_explicit() == permutedims(arr, index_order_new=["a", "b", "c", "d", "e"], index_order_old=['b', 'e', 'a', 'd', 'c'])
|
||||
|
||||
|
||||
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_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_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)
|
||||
+78
@@ -0,0 +1,78 @@
|
||||
from sympy.core.symbol import symbols
|
||||
from sympy.functions.elementary.trigonometric import (cos, sin)
|
||||
from sympy.matrices.expressions.matexpr import MatrixSymbol
|
||||
from sympy.matrices.expressions.special import Identity
|
||||
from sympy.matrices.expressions.applyfunc import ElementwiseApplyFunction
|
||||
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))
|
||||
|
||||
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)
|
||||
))
|
||||
|
||||
cg = Reshape(A, (k**2,))
|
||||
res = array_derive(cg, A)
|
||||
assert res == Reshape(PermuteDims(ArrayTensorProduct(I, I), [0, 2, 1, 3]), (k, k, k**2))
|
||||
+63
@@ -0,0 +1,63 @@
|
||||
from sympy.core.symbol import Symbol
|
||||
from sympy.matrices.expressions.matexpr import MatrixSymbol
|
||||
from sympy.tensor.array.arrayop import (permutedims, tensorcontraction, tensordiagonal, tensorproduct)
|
||||
from sympy.tensor.array.dense_ndim_array import ImmutableDenseNDimArray
|
||||
from sympy.tensor.array.expressions.array_expressions import ZeroArray, OneArray, ArraySymbol, \
|
||||
ArrayTensorProduct, PermuteDims, ArrayDiagonal, ArrayContraction, ArrayAdd
|
||||
from sympy.testing.pytest import raises
|
||||
|
||||
|
||||
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_array_as_explicit_matrix_symbol():
|
||||
|
||||
A = MatrixSymbol("A", 3, 3)
|
||||
B = MatrixSymbol("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])
|
||||
|
||||
expr = ArrayAdd(ArrayTensorProduct(A, B), ArrayTensorProduct(B, A))
|
||||
assert expr.as_explicit() == expr.args[0].as_explicit() + expr.args[1].as_explicit()
|
||||
+61
@@ -0,0 +1,61 @@
|
||||
from sympy import Sum, Dummy, sin
|
||||
from sympy.tensor.array.expressions import ArraySymbol, ArrayTensorProduct, ArrayContraction, PermuteDims, \
|
||||
ArrayDiagonal, ArrayAdd, OneArray, ZeroArray, convert_indexed_to_array, ArrayElementwiseApplyFunc, Reshape
|
||||
from sympy.tensor.array.expressions.from_array_to_indexed import convert_array_to_indexed
|
||||
|
||||
from sympy.abc import i, j, k, l, m, n, o
|
||||
|
||||
|
||||
def test_convert_array_to_indexed_main():
|
||||
A = ArraySymbol("A", (3, 3, 3))
|
||||
B = ArraySymbol("B", (3, 3))
|
||||
C = ArraySymbol("C", (3, 3))
|
||||
|
||||
d_ = Dummy("d_")
|
||||
|
||||
assert convert_array_to_indexed(A, [i, j, k]) == A[i, j, k]
|
||||
|
||||
expr = ArrayTensorProduct(A, B, C)
|
||||
conv = convert_array_to_indexed(expr, [i,j,k,l,m,n,o])
|
||||
assert conv == A[i,j,k]*B[l,m]*C[n,o]
|
||||
assert convert_indexed_to_array(conv, [i,j,k,l,m,n,o]) == expr
|
||||
|
||||
expr = ArrayContraction(A, (0, 2))
|
||||
assert convert_array_to_indexed(expr, [i]).dummy_eq(Sum(A[d_, i, d_], (d_, 0, 2)))
|
||||
|
||||
expr = ArrayDiagonal(A, (0, 2))
|
||||
assert convert_array_to_indexed(expr, [i, j]) == A[j, i, j]
|
||||
|
||||
expr = PermuteDims(A, [1, 2, 0])
|
||||
conv = convert_array_to_indexed(expr, [i, j, k])
|
||||
assert conv == A[k, i, j]
|
||||
assert convert_indexed_to_array(conv, [i, j, k]) == expr
|
||||
|
||||
expr = ArrayAdd(B, C, PermuteDims(C, [1, 0]))
|
||||
conv = convert_array_to_indexed(expr, [i, j])
|
||||
assert conv == B[i, j] + C[i, j] + C[j, i]
|
||||
assert convert_indexed_to_array(conv, [i, j]) == expr
|
||||
|
||||
expr = ArrayElementwiseApplyFunc(sin, A)
|
||||
conv = convert_array_to_indexed(expr, [i, j, k])
|
||||
assert conv == sin(A[i, j, k])
|
||||
assert convert_indexed_to_array(conv, [i, j, k]).dummy_eq(expr)
|
||||
|
||||
assert convert_array_to_indexed(OneArray(3, 3), [i, j]) == 1
|
||||
assert convert_array_to_indexed(ZeroArray(3, 3), [i, j]) == 0
|
||||
|
||||
expr = Reshape(A, (27,))
|
||||
assert convert_array_to_indexed(expr, [i]) == A[i // 9, i // 3 % 3, i % 3]
|
||||
|
||||
X = ArraySymbol("X", (2, 3, 4, 5, 6))
|
||||
expr = Reshape(X, (2*3*4*5*6,))
|
||||
assert convert_array_to_indexed(expr, [i]) == X[i // 360, i // 120 % 3, i // 30 % 4, i // 6 % 5, i % 6]
|
||||
|
||||
expr = Reshape(X, (4, 9, 2, 2, 5))
|
||||
one_index = 180*i + 20*j + 10*k + 5*l + m
|
||||
expected = X[one_index // (3*4*5*6), one_index // (4*5*6) % 3, one_index // (5*6) % 4, one_index // 6 % 5, one_index % 6]
|
||||
assert convert_array_to_indexed(expr, [i, j, k, l, m]) == expected
|
||||
|
||||
X = ArraySymbol("X", (2*3*5,))
|
||||
expr = Reshape(X, (2, 3, 5))
|
||||
assert convert_array_to_indexed(expr, [i, j, k]) == X[15*i + 5*j + k]
|
||||
+689
@@ -0,0 +1,689 @@
|
||||
from sympy import Lambda, S, Dummy, KroneckerProduct
|
||||
from sympy.core.symbol import symbols
|
||||
from sympy.functions.elementary.miscellaneous import sqrt
|
||||
from sympy.functions.elementary.trigonometric import cos, sin
|
||||
from sympy.matrices.expressions.hadamard import HadamardProduct, HadamardPower
|
||||
from sympy.matrices.expressions.special import (Identity, OneMatrix, ZeroMatrix)
|
||||
from sympy.matrices.expressions.matexpr import MatrixElement
|
||||
from sympy.tensor.array.expressions.from_matrix_to_array import convert_matrix_to_array
|
||||
from sympy.tensor.array.expressions.from_array_to_matrix import _support_function_tp1_recognize, \
|
||||
_array_diag2contr_diagmatrix, convert_array_to_matrix, _remove_trivial_dims, _array2matrix, \
|
||||
_combine_removed, identify_removable_identity_matrices, _array_contraction_to_diagonal_multiple_identity
|
||||
from sympy.matrices.expressions.matexpr import MatrixSymbol
|
||||
from sympy.combinatorics import Permutation
|
||||
from sympy.matrices.expressions.diagonal import DiagMatrix, DiagonalMatrix
|
||||
from sympy.matrices import Trace, MatMul, Transpose
|
||||
from sympy.tensor.array.expressions.array_expressions import ZeroArray, OneArray, \
|
||||
ArrayElement, ArraySymbol, ArrayElementwiseApplyFunc, _array_tensor_product, _array_contraction, \
|
||||
_array_diagonal, _permute_dims, PermuteDims, ArrayAdd, ArrayDiagonal, ArrayContraction, ArrayTensorProduct
|
||||
from sympy.testing.pytest import raises
|
||||
|
||||
|
||||
i, j, k, l, m, n = symbols("i j k l m n")
|
||||
|
||||
I = Identity(k)
|
||||
I1 = Identity(1)
|
||||
|
||||
M = MatrixSymbol("M", k, k)
|
||||
N = MatrixSymbol("N", k, k)
|
||||
P = MatrixSymbol("P", k, k)
|
||||
Q = MatrixSymbol("Q", k, k)
|
||||
|
||||
A = MatrixSymbol("A", k, k)
|
||||
B = MatrixSymbol("B", k, k)
|
||||
C = MatrixSymbol("C", k, k)
|
||||
D = MatrixSymbol("D", k, k)
|
||||
|
||||
X = MatrixSymbol("X", k, k)
|
||||
Y = MatrixSymbol("Y", k, k)
|
||||
|
||||
a = MatrixSymbol("a", k, 1)
|
||||
b = MatrixSymbol("b", k, 1)
|
||||
c = MatrixSymbol("c", k, 1)
|
||||
d = MatrixSymbol("d", k, 1)
|
||||
|
||||
x = MatrixSymbol("x", k, 1)
|
||||
y = MatrixSymbol("y", k, 1)
|
||||
|
||||
|
||||
def test_arrayexpr_convert_array_to_matrix():
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(M), (0, 1))
|
||||
assert convert_array_to_matrix(cg) == Trace(M)
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(M, N), (0, 1), (2, 3))
|
||||
assert convert_array_to_matrix(cg) == Trace(M) * Trace(N)
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(M, N), (0, 3), (1, 2))
|
||||
assert convert_array_to_matrix(cg) == Trace(M * N)
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(M, N), (0, 2), (1, 3))
|
||||
assert convert_array_to_matrix(cg) == Trace(M * N.T)
|
||||
|
||||
cg = convert_matrix_to_array(M * N * P)
|
||||
assert convert_array_to_matrix(cg) == M * N * P
|
||||
|
||||
cg = convert_matrix_to_array(M * N.T * P)
|
||||
assert convert_array_to_matrix(cg) == M * N.T * P
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(M,N,P,Q), (1, 2), (5, 6))
|
||||
assert convert_array_to_matrix(cg) == _array_tensor_product(M * N, P * Q)
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(-2, M, N), (1, 2))
|
||||
assert convert_array_to_matrix(cg) == -2 * M * N
|
||||
|
||||
a = MatrixSymbol("a", k, 1)
|
||||
b = MatrixSymbol("b", k, 1)
|
||||
c = MatrixSymbol("c", k, 1)
|
||||
cg = PermuteDims(
|
||||
_array_contraction(
|
||||
_array_tensor_product(
|
||||
a,
|
||||
ArrayAdd(
|
||||
_array_tensor_product(b, c),
|
||||
_array_tensor_product(c, b),
|
||||
)
|
||||
), (2, 4)), [0, 1, 3, 2])
|
||||
assert convert_array_to_matrix(cg) == a * (b.T * c + c.T * b)
|
||||
|
||||
za = ZeroArray(m, n)
|
||||
assert convert_array_to_matrix(za) == ZeroMatrix(m, n)
|
||||
|
||||
cg = _array_tensor_product(3, M)
|
||||
assert convert_array_to_matrix(cg) == 3 * M
|
||||
|
||||
# Partial conversion to matrix multiplication:
|
||||
expr = _array_contraction(_array_tensor_product(M, N, P, Q), (0, 2), (1, 4, 6))
|
||||
assert convert_array_to_matrix(expr) == _array_contraction(_array_tensor_product(M.T*N, P, Q), (0, 2, 4))
|
||||
|
||||
x = MatrixSymbol("x", k, 1)
|
||||
cg = PermuteDims(
|
||||
_array_contraction(_array_tensor_product(OneArray(1), x, OneArray(1), DiagMatrix(Identity(1))),
|
||||
(0, 5)), Permutation(1, 2, 3))
|
||||
assert convert_array_to_matrix(cg) == x
|
||||
|
||||
expr = ArrayAdd(M, PermuteDims(M, [1, 0]))
|
||||
assert convert_array_to_matrix(expr) == M + Transpose(M)
|
||||
|
||||
|
||||
def test_arrayexpr_convert_array_to_matrix2():
|
||||
cg = _array_contraction(_array_tensor_product(M, N), (1, 3))
|
||||
assert convert_array_to_matrix(cg) == M * N.T
|
||||
|
||||
cg = PermuteDims(_array_tensor_product(M, N), Permutation([0, 1, 3, 2]))
|
||||
assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T)
|
||||
|
||||
cg = _array_tensor_product(M, PermuteDims(N, Permutation([1, 0])))
|
||||
assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T)
|
||||
|
||||
cg = _array_contraction(
|
||||
PermuteDims(
|
||||
_array_tensor_product(M, N, P, Q), Permutation([0, 2, 3, 1, 4, 5, 7, 6])),
|
||||
(1, 2), (3, 5)
|
||||
)
|
||||
assert convert_array_to_matrix(cg) == _array_tensor_product(M * P.T * Trace(N), Q.T)
|
||||
|
||||
cg = _array_contraction(
|
||||
_array_tensor_product(M, N, P, PermuteDims(Q, Permutation([1, 0]))),
|
||||
(1, 5), (2, 3)
|
||||
)
|
||||
assert convert_array_to_matrix(cg) == _array_tensor_product(M * P.T * Trace(N), Q.T)
|
||||
|
||||
cg = _array_tensor_product(M, PermuteDims(N, [1, 0]))
|
||||
assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T)
|
||||
|
||||
cg = _array_tensor_product(PermuteDims(M, [1, 0]), PermuteDims(N, [1, 0]))
|
||||
assert convert_array_to_matrix(cg) == _array_tensor_product(M.T, N.T)
|
||||
|
||||
cg = _array_tensor_product(PermuteDims(N, [1, 0]), PermuteDims(M, [1, 0]))
|
||||
assert convert_array_to_matrix(cg) == _array_tensor_product(N.T, M.T)
|
||||
|
||||
cg = _array_contraction(M, (0,), (1,))
|
||||
assert convert_array_to_matrix(cg) == OneMatrix(1, k)*M*OneMatrix(k, 1)
|
||||
|
||||
cg = _array_contraction(x, (0,), (1,))
|
||||
assert convert_array_to_matrix(cg) == OneMatrix(1, k)*x
|
||||
|
||||
Xm = MatrixSymbol("Xm", m, n)
|
||||
cg = _array_contraction(Xm, (0,), (1,))
|
||||
assert convert_array_to_matrix(cg) == OneMatrix(1, m)*Xm*OneMatrix(n, 1)
|
||||
|
||||
|
||||
def test_arrayexpr_convert_array_to_diagonalized_vector():
|
||||
|
||||
# Check matrix recognition over trivial dimensions:
|
||||
|
||||
cg = _array_tensor_product(a, b)
|
||||
assert convert_array_to_matrix(cg) == a * b.T
|
||||
|
||||
cg = _array_tensor_product(I1, a, b)
|
||||
assert convert_array_to_matrix(cg) == a * b.T
|
||||
|
||||
# Recognize trace inside a tensor product:
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(A, B, C), (0, 3), (1, 2))
|
||||
assert convert_array_to_matrix(cg) == Trace(A * B) * C
|
||||
|
||||
# Transform diagonal operator to contraction:
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(A, a), (1, 2))
|
||||
assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(A, OneArray(1), DiagMatrix(a)), (1, 3))
|
||||
assert convert_array_to_matrix(cg) == A * DiagMatrix(a)
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(a, b), (0, 2))
|
||||
assert _array_diag2contr_diagmatrix(cg) == _permute_dims(
|
||||
_array_contraction(_array_tensor_product(DiagMatrix(a), OneArray(1), b), (0, 3)), [1, 2, 0]
|
||||
)
|
||||
assert convert_array_to_matrix(cg) == b.T * DiagMatrix(a)
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(A, a), (0, 2))
|
||||
assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(A, OneArray(1), DiagMatrix(a)), (0, 3))
|
||||
assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a)
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(I, x, I1), (0, 2), (3, 5))
|
||||
assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(I, OneArray(1), I1, DiagMatrix(x)), (0, 5))
|
||||
assert convert_array_to_matrix(cg) == DiagMatrix(x)
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(I, x, A, B), (1, 2), (5, 6))
|
||||
assert _array_diag2contr_diagmatrix(cg) == _array_diagonal(_array_contraction(_array_tensor_product(I, OneArray(1), A, B, DiagMatrix(x)), (1, 7)), (5, 6))
|
||||
# TODO: this is returning a wrong result:
|
||||
# convert_array_to_matrix(cg)
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(I1, a, b), (1, 3, 5))
|
||||
assert convert_array_to_matrix(cg) == a*b.T
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(I1, a, b), (1, 3))
|
||||
assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(OneArray(1), a, b, I1), (2, 6))
|
||||
assert convert_array_to_matrix(cg) == a*b.T
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(x, I1), (1, 2))
|
||||
assert isinstance(cg, ArrayDiagonal)
|
||||
assert cg.diagonal_indices == ((1, 2),)
|
||||
assert convert_array_to_matrix(cg) == x
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(x, I), (0, 2))
|
||||
assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(OneArray(1), I, DiagMatrix(x)), (1, 3))
|
||||
assert convert_array_to_matrix(cg).doit() == DiagMatrix(x)
|
||||
|
||||
raises(ValueError, lambda: _array_diagonal(x, (1,)))
|
||||
|
||||
# Ignore identity matrices with contractions:
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(I, A, I, I), (0, 2), (1, 3), (5, 7))
|
||||
assert cg.split_multiple_contractions() == cg
|
||||
assert convert_array_to_matrix(cg) == Trace(A) * I
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(Trace(A) * I, I, I), (1, 5), (3, 4))
|
||||
assert cg.split_multiple_contractions() == cg
|
||||
assert convert_array_to_matrix(cg).doit() == Trace(A) * I
|
||||
|
||||
# Add DiagMatrix when required:
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(A, a), (1, 2))
|
||||
assert cg.split_multiple_contractions() == cg
|
||||
assert convert_array_to_matrix(cg) == A * a
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(A, a, B), (1, 2, 4))
|
||||
assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1), B), (1, 2), (3, 5))
|
||||
assert convert_array_to_matrix(cg) == A * DiagMatrix(a) * B
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(A, a, B), (0, 2, 4))
|
||||
assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1), B), (0, 2), (3, 5))
|
||||
assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a) * B
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(A, a, b, a.T, B), (0, 2, 4, 7, 9))
|
||||
assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1),
|
||||
DiagMatrix(b), OneArray(1), DiagMatrix(a), OneArray(1), B),
|
||||
(0, 2), (3, 5), (6, 9), (8, 12))
|
||||
assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a) * DiagMatrix(b) * DiagMatrix(a) * B.T
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(I1, I1, I1), (1, 2, 4))
|
||||
assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(I1, I1, OneArray(1), I1), (1, 2), (3, 5))
|
||||
assert convert_array_to_matrix(cg) == 1
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(I, I, I, I, A), (1, 2, 8), (5, 6, 9))
|
||||
assert convert_array_to_matrix(cg.split_multiple_contractions()).doit() == A
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(A, a, C, a, B), (1, 2, 4), (5, 6, 8))
|
||||
expected = _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1), C, DiagMatrix(a), OneArray(1), B), (1, 3), (2, 5), (6, 7), (8, 10))
|
||||
assert cg.split_multiple_contractions() == expected
|
||||
assert convert_array_to_matrix(cg) == A * DiagMatrix(a) * C * DiagMatrix(a) * B
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(a, I1, b, I1, (a.T*b).applyfunc(cos)), (1, 2, 8), (5, 6, 9))
|
||||
expected = _array_contraction(_array_tensor_product(a, I1, OneArray(1), b, I1, OneArray(1), (a.T*b).applyfunc(cos)),
|
||||
(1, 3), (2, 10), (6, 8), (7, 11))
|
||||
assert cg.split_multiple_contractions().dummy_eq(expected)
|
||||
assert convert_array_to_matrix(cg).doit().dummy_eq(MatMul(a, (a.T * b).applyfunc(cos), b.T))
|
||||
|
||||
|
||||
def test_arrayexpr_convert_array_contraction_tp_additions():
|
||||
a = ArrayAdd(
|
||||
_array_tensor_product(M, N),
|
||||
_array_tensor_product(N, M)
|
||||
)
|
||||
tp = _array_tensor_product(P, a, Q)
|
||||
expr = _array_contraction(tp, (3, 4))
|
||||
expected = _array_tensor_product(
|
||||
P,
|
||||
ArrayAdd(
|
||||
_array_contraction(_array_tensor_product(M, N), (1, 2)),
|
||||
_array_contraction(_array_tensor_product(N, M), (1, 2)),
|
||||
),
|
||||
Q
|
||||
)
|
||||
assert expr == expected
|
||||
assert convert_array_to_matrix(expr) == _array_tensor_product(P, M * N + N * M, Q)
|
||||
|
||||
expr = _array_contraction(tp, (1, 2), (3, 4), (5, 6))
|
||||
result = _array_contraction(
|
||||
_array_tensor_product(
|
||||
P,
|
||||
ArrayAdd(
|
||||
_array_contraction(_array_tensor_product(M, N), (1, 2)),
|
||||
_array_contraction(_array_tensor_product(N, M), (1, 2)),
|
||||
),
|
||||
Q
|
||||
), (1, 2), (3, 4))
|
||||
assert expr == result
|
||||
assert convert_array_to_matrix(expr) == P * (M * N + N * M) * Q
|
||||
|
||||
|
||||
def test_arrayexpr_convert_array_to_implicit_matmul():
|
||||
# Trivial dimensions are suppressed, so the result can be expressed in matrix form:
|
||||
|
||||
cg = _array_tensor_product(a, b)
|
||||
assert convert_array_to_matrix(cg) == a * b.T
|
||||
|
||||
cg = _array_tensor_product(a, b, I)
|
||||
assert convert_array_to_matrix(cg) == _array_tensor_product(a*b.T, I)
|
||||
|
||||
cg = _array_tensor_product(I, a, b)
|
||||
assert convert_array_to_matrix(cg) == _array_tensor_product(I, a*b.T)
|
||||
|
||||
cg = _array_tensor_product(a, I, b)
|
||||
assert convert_array_to_matrix(cg) == _array_tensor_product(a, I, b)
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(I, I), (1, 2))
|
||||
assert convert_array_to_matrix(cg) == I
|
||||
|
||||
cg = PermuteDims(_array_tensor_product(I, Identity(1)), [0, 2, 1, 3])
|
||||
assert convert_array_to_matrix(cg) == I
|
||||
|
||||
|
||||
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_arrayexpr_convert_array_to_matrix_diag2contraction_diagmatrix():
|
||||
cg = _array_diagonal(_array_tensor_product(M, a), (1, 2))
|
||||
res = _array_diag2contr_diagmatrix(cg)
|
||||
assert res.shape == cg.shape
|
||||
assert res == _array_contraction(_array_tensor_product(M, OneArray(1), DiagMatrix(a)), (1, 3))
|
||||
|
||||
raises(ValueError, lambda: _array_diagonal(_array_tensor_product(a, M), (1, 2)))
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(a.T, M), (1, 2))
|
||||
res = _array_diag2contr_diagmatrix(cg)
|
||||
assert res.shape == cg.shape
|
||||
assert res == _array_contraction(_array_tensor_product(OneArray(1), M, DiagMatrix(a.T)), (1, 4))
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(a.T, M, N, b.T), (1, 2), (4, 7))
|
||||
res = _array_diag2contr_diagmatrix(cg)
|
||||
assert res.shape == cg.shape
|
||||
assert res == _array_contraction(
|
||||
_array_tensor_product(OneArray(1), M, N, OneArray(1), DiagMatrix(a.T), DiagMatrix(b.T)), (1, 7), (3, 9))
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(a, M, N, b.T), (0, 2), (4, 7))
|
||||
res = _array_diag2contr_diagmatrix(cg)
|
||||
assert res.shape == cg.shape
|
||||
assert res == _array_contraction(
|
||||
_array_tensor_product(OneArray(1), M, N, OneArray(1), DiagMatrix(a), DiagMatrix(b.T)), (1, 6), (3, 9))
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(a, M, N, b.T), (0, 4), (3, 7))
|
||||
res = _array_diag2contr_diagmatrix(cg)
|
||||
assert res.shape == cg.shape
|
||||
assert res == _array_contraction(
|
||||
_array_tensor_product(OneArray(1), M, N, OneArray(1), DiagMatrix(a), DiagMatrix(b.T)), (3, 6), (2, 9))
|
||||
|
||||
I1 = Identity(1)
|
||||
x = MatrixSymbol("x", k, 1)
|
||||
A = MatrixSymbol("A", k, k)
|
||||
cg = _array_diagonal(_array_tensor_product(x, A.T, I1), (0, 2))
|
||||
assert _array_diag2contr_diagmatrix(cg).shape == cg.shape
|
||||
assert _array2matrix(cg).shape == cg.shape
|
||||
|
||||
|
||||
def test_arrayexpr_convert_array_to_matrix_support_function():
|
||||
|
||||
assert _support_function_tp1_recognize([], [2 * k]) == 2 * k
|
||||
|
||||
assert _support_function_tp1_recognize([(1, 2)], [A, 2 * k, B, 3]) == 6 * k * A * B
|
||||
|
||||
assert _support_function_tp1_recognize([(0, 3), (1, 2)], [A, B]) == Trace(A * B)
|
||||
|
||||
assert _support_function_tp1_recognize([(1, 2)], [A, B]) == A * B
|
||||
assert _support_function_tp1_recognize([(0, 2)], [A, B]) == A.T * B
|
||||
assert _support_function_tp1_recognize([(1, 3)], [A, B]) == A * B.T
|
||||
assert _support_function_tp1_recognize([(0, 3)], [A, B]) == A.T * B.T
|
||||
|
||||
assert _support_function_tp1_recognize([(1, 2), (5, 6)], [A, B, C, D]) == _array_tensor_product(A * B, C * D)
|
||||
assert _support_function_tp1_recognize([(1, 4), (3, 6)], [A, B, C, D]) == PermuteDims(
|
||||
_array_tensor_product(A * C, B * D), [0, 2, 1, 3])
|
||||
|
||||
assert _support_function_tp1_recognize([(0, 3), (1, 4)], [A, B, C]) == B * A * C
|
||||
|
||||
assert _support_function_tp1_recognize([(9, 10), (1, 2), (5, 6), (3, 4), (7, 8)],
|
||||
[X, Y, A, B, C, D]) == X * Y * A * B * C * D
|
||||
|
||||
assert _support_function_tp1_recognize([(9, 10), (1, 2), (5, 6), (3, 4)],
|
||||
[X, Y, A, B, C, D]) == _array_tensor_product(X * Y * A * B, C * D)
|
||||
|
||||
assert _support_function_tp1_recognize([(1, 7), (3, 8), (4, 11)], [X, Y, A, B, C, D]) == PermuteDims(
|
||||
_array_tensor_product(X * B.T, Y * C, A.T * D.T), [0, 2, 4, 1, 3, 5]
|
||||
)
|
||||
|
||||
assert _support_function_tp1_recognize([(0, 1), (3, 6), (5, 8)], [X, A, B, C, D]) == PermuteDims(
|
||||
_array_tensor_product(Trace(X) * A * C, B * D), [0, 2, 1, 3])
|
||||
|
||||
assert _support_function_tp1_recognize([(1, 2), (3, 4), (5, 6), (7, 8)], [A, A, B, C, D]) == A ** 2 * B * C * D
|
||||
assert _support_function_tp1_recognize([(1, 2), (3, 4), (5, 6), (7, 8)], [X, A, B, C, D]) == X * A * B * C * D
|
||||
|
||||
assert _support_function_tp1_recognize([(1, 6), (3, 8), (5, 10)], [X, Y, A, B, C, D]) == PermuteDims(
|
||||
_array_tensor_product(X * B, Y * C, A * D), [0, 2, 4, 1, 3, 5]
|
||||
)
|
||||
|
||||
assert _support_function_tp1_recognize([(1, 4), (3, 6)], [A, B, C, D]) == PermuteDims(
|
||||
_array_tensor_product(A * C, B * D), [0, 2, 1, 3])
|
||||
|
||||
assert _support_function_tp1_recognize([(0, 4), (1, 7), (2, 5), (3, 8)], [X, A, B, C, D]) == C*X.T*B*A*D
|
||||
|
||||
assert _support_function_tp1_recognize([(0, 4), (1, 7), (2, 5), (3, 8)], [X, A, B, C, D]) == C*X.T*B*A*D
|
||||
|
||||
|
||||
def test_convert_array_to_hadamard_products():
|
||||
|
||||
expr = HadamardProduct(M, N)
|
||||
cg = convert_matrix_to_array(expr)
|
||||
ret = convert_array_to_matrix(cg)
|
||||
assert ret == expr
|
||||
|
||||
expr = HadamardProduct(M, N)*P
|
||||
cg = convert_matrix_to_array(expr)
|
||||
ret = convert_array_to_matrix(cg)
|
||||
assert ret == expr
|
||||
|
||||
expr = Q*HadamardProduct(M, N)*P
|
||||
cg = convert_matrix_to_array(expr)
|
||||
ret = convert_array_to_matrix(cg)
|
||||
assert ret == expr
|
||||
|
||||
expr = Q*HadamardProduct(M, N.T)*P
|
||||
cg = convert_matrix_to_array(expr)
|
||||
ret = convert_array_to_matrix(cg)
|
||||
assert ret == expr
|
||||
|
||||
expr = HadamardProduct(M, N)*HadamardProduct(Q, P)
|
||||
cg = convert_matrix_to_array(expr)
|
||||
ret = convert_array_to_matrix(cg)
|
||||
assert expr == ret
|
||||
|
||||
expr = P.T*HadamardProduct(M, N)*HadamardProduct(Q, P)
|
||||
cg = convert_matrix_to_array(expr)
|
||||
ret = convert_array_to_matrix(cg)
|
||||
assert expr == ret
|
||||
|
||||
# ArrayDiagonal should be converted
|
||||
cg = _array_diagonal(_array_tensor_product(M, N, Q), (1, 3), (0, 2, 4))
|
||||
ret = convert_array_to_matrix(cg)
|
||||
expected = PermuteDims(_array_diagonal(_array_tensor_product(HadamardProduct(M.T, N.T), Q), (1, 2)), [1, 0, 2])
|
||||
assert expected == ret
|
||||
|
||||
# Special case that should return the same expression:
|
||||
cg = _array_diagonal(_array_tensor_product(HadamardProduct(M, N), Q), (0, 2))
|
||||
ret = convert_array_to_matrix(cg)
|
||||
assert ret == cg
|
||||
|
||||
# Hadamard products with traces:
|
||||
|
||||
expr = Trace(HadamardProduct(M, N))
|
||||
cg = convert_matrix_to_array(expr)
|
||||
ret = convert_array_to_matrix(cg)
|
||||
assert ret == Trace(HadamardProduct(M.T, N.T))
|
||||
|
||||
expr = Trace(A*HadamardProduct(M, N))
|
||||
cg = convert_matrix_to_array(expr)
|
||||
ret = convert_array_to_matrix(cg)
|
||||
assert ret == Trace(HadamardProduct(M, N)*A)
|
||||
|
||||
expr = Trace(HadamardProduct(A, M)*N)
|
||||
cg = convert_matrix_to_array(expr)
|
||||
ret = convert_array_to_matrix(cg)
|
||||
assert ret == Trace(HadamardProduct(M.T, N)*A)
|
||||
|
||||
# These should not be converted into Hadamard products:
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(M, N), (0, 1, 2, 3))
|
||||
ret = convert_array_to_matrix(cg)
|
||||
assert ret == cg
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(A), (0, 1))
|
||||
ret = convert_array_to_matrix(cg)
|
||||
assert ret == cg
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(M, N, P), (0, 2, 4), (1, 3, 5))
|
||||
assert convert_array_to_matrix(cg) == HadamardProduct(M, N, P)
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(M, N, P), (0, 3, 4), (1, 2, 5))
|
||||
assert convert_array_to_matrix(cg) == HadamardProduct(M, P, N.T)
|
||||
|
||||
cg = _array_diagonal(_array_tensor_product(I, I1, x), (1, 4), (3, 5))
|
||||
assert convert_array_to_matrix(cg) == DiagMatrix(x)
|
||||
|
||||
|
||||
def test_identify_removable_identity_matrices():
|
||||
|
||||
D = DiagonalMatrix(MatrixSymbol("D", k, k))
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(A, B, I), (1, 2, 4, 5))
|
||||
expected = _array_contraction(_array_tensor_product(A, B), (1, 2))
|
||||
assert identify_removable_identity_matrices(cg) == expected
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(A, B, C, I), (1, 3, 5, 6, 7))
|
||||
expected = _array_contraction(_array_tensor_product(A, B, C), (1, 3, 5))
|
||||
assert identify_removable_identity_matrices(cg) == expected
|
||||
|
||||
# Tests with diagonal matrices:
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(A, B, D), (1, 2, 4, 5))
|
||||
ret = identify_removable_identity_matrices(cg)
|
||||
expected = _array_contraction(_array_tensor_product(A, B, D), (1, 4), (2, 5))
|
||||
assert ret == expected
|
||||
|
||||
cg = _array_contraction(_array_tensor_product(A, B, D, M, N), (1, 2, 4, 5, 6, 8))
|
||||
ret = identify_removable_identity_matrices(cg)
|
||||
assert ret == cg
|
||||
|
||||
|
||||
def test_combine_removed():
|
||||
|
||||
assert _combine_removed(6, [0, 1, 2], [0, 1, 2]) == [0, 1, 2, 3, 4, 5]
|
||||
assert _combine_removed(8, [2, 5], [1, 3, 4]) == [1, 2, 4, 5, 6]
|
||||
assert _combine_removed(8, [7], []) == [7]
|
||||
|
||||
|
||||
def test_array_contraction_to_diagonal_multiple_identities():
|
||||
|
||||
expr = _array_contraction(_array_tensor_product(A, B, I, C), (1, 2, 4), (5, 6))
|
||||
assert _array_contraction_to_diagonal_multiple_identity(expr) == (expr, [])
|
||||
assert convert_array_to_matrix(expr) == _array_contraction(_array_tensor_product(A, B, C), (1, 2, 4))
|
||||
|
||||
expr = _array_contraction(_array_tensor_product(A, I, I), (1, 2, 4))
|
||||
assert _array_contraction_to_diagonal_multiple_identity(expr) == (A, [2])
|
||||
assert convert_array_to_matrix(expr) == A
|
||||
|
||||
expr = _array_contraction(_array_tensor_product(A, I, I, B), (1, 2, 4), (3, 6))
|
||||
assert _array_contraction_to_diagonal_multiple_identity(expr) == (expr, [])
|
||||
|
||||
expr = _array_contraction(_array_tensor_product(A, I, I, B), (1, 2, 3, 4, 6))
|
||||
assert _array_contraction_to_diagonal_multiple_identity(expr) == (expr, [])
|
||||
|
||||
|
||||
def test_convert_array_element_to_matrix():
|
||||
|
||||
expr = ArrayElement(M, (i, j))
|
||||
assert convert_array_to_matrix(expr) == MatrixElement(M, i, j)
|
||||
|
||||
expr = ArrayElement(_array_contraction(_array_tensor_product(M, N), (1, 3)), (i, j))
|
||||
assert convert_array_to_matrix(expr) == MatrixElement(M*N.T, i, j)
|
||||
|
||||
expr = ArrayElement(_array_tensor_product(M, N), (i, j, m, n))
|
||||
assert convert_array_to_matrix(expr) == expr
|
||||
|
||||
|
||||
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 test_array2matrix():
|
||||
# See issue https://github.com/sympy/sympy/pull/22877
|
||||
expr = PermuteDims(ArrayContraction(ArrayTensorProduct(x, I, I1, x), (0, 3), (1, 7)), Permutation(2, 3))
|
||||
expected = PermuteDims(ArrayTensorProduct(x*x.T, I1), Permutation(3)(1, 2))
|
||||
assert _array2matrix(expr) == expected
|
||||
|
||||
|
||||
def test_recognize_broadcasting():
|
||||
expr = ArrayTensorProduct(x.T*x, A)
|
||||
assert _remove_trivial_dims(expr) == (KroneckerProduct(x.T*x, A), [0, 1])
|
||||
|
||||
expr = ArrayTensorProduct(A, x.T*x)
|
||||
assert _remove_trivial_dims(expr) == (KroneckerProduct(A, x.T*x), [2, 3])
|
||||
|
||||
expr = ArrayTensorProduct(A, B, x.T*x, C)
|
||||
assert _remove_trivial_dims(expr) == (ArrayTensorProduct(A, KroneckerProduct(B, x.T*x), C), [4, 5])
|
||||
|
||||
# Always prefer matrix multiplication to Kronecker product, if possible:
|
||||
expr = ArrayTensorProduct(a, b, x.T*x)
|
||||
assert _remove_trivial_dims(expr) == (a*x.T*x*b.T, [1, 3, 4, 5])
|
||||
+205
@@ -0,0 +1,205 @@
|
||||
from sympy import tanh
|
||||
from sympy.concrete.summations import Sum
|
||||
from sympy.core.symbol import symbols
|
||||
from sympy.functions.special.tensor_functions import KroneckerDelta
|
||||
from sympy.matrices.expressions.matexpr import MatrixSymbol
|
||||
from sympy.matrices.expressions.special import Identity
|
||||
from sympy.tensor.array.expressions import ArrayElementwiseApplyFunc
|
||||
from sympy.tensor.indexed import IndexedBase
|
||||
from sympy.combinatorics import Permutation
|
||||
from sympy.tensor.array.expressions.array_expressions import ArrayContraction, ArrayTensorProduct, \
|
||||
ArrayDiagonal, ArrayAdd, PermuteDims, ArrayElement, _array_tensor_product, _array_contraction, _array_diagonal, \
|
||||
_array_add, _permute_dims, ArraySymbol, OneArray
|
||||
from sympy.tensor.array.expressions.from_array_to_matrix import convert_array_to_matrix
|
||||
from sympy.tensor.array.expressions.from_indexed_to_array import convert_indexed_to_array, _convert_indexed_to_array
|
||||
from sympy.testing.pytest import raises
|
||||
|
||||
|
||||
A, B = symbols("A B", cls=IndexedBase)
|
||||
i, j, k, l, m, n = symbols("i j k l m n")
|
||||
d0, d1, d2, d3 = symbols("d0:4")
|
||||
|
||||
I = Identity(k)
|
||||
|
||||
M = MatrixSymbol("M", k, k)
|
||||
N = MatrixSymbol("N", k, k)
|
||||
P = MatrixSymbol("P", k, k)
|
||||
Q = MatrixSymbol("Q", k, k)
|
||||
|
||||
a = MatrixSymbol("a", k, 1)
|
||||
b = MatrixSymbol("b", k, 1)
|
||||
c = MatrixSymbol("c", k, 1)
|
||||
d = MatrixSymbol("d", k, 1)
|
||||
|
||||
|
||||
def test_arrayexpr_convert_index_to_array_support_function():
|
||||
expr = M[i, j]
|
||||
assert _convert_indexed_to_array(expr) == (M, (i, j))
|
||||
expr = M[i, j]*N[k, l]
|
||||
assert _convert_indexed_to_array(expr) == (ArrayTensorProduct(M, N), (i, j, k, l))
|
||||
expr = M[i, j]*N[j, k]
|
||||
assert _convert_indexed_to_array(expr) == (ArrayDiagonal(ArrayTensorProduct(M, N), (1, 2)), (i, k, j))
|
||||
expr = Sum(M[i, j]*N[j, k], (j, 0, k-1))
|
||||
assert _convert_indexed_to_array(expr) == (ArrayContraction(ArrayTensorProduct(M, N), (1, 2)), (i, k))
|
||||
expr = M[i, j] + N[i, j]
|
||||
assert _convert_indexed_to_array(expr) == (ArrayAdd(M, N), (i, j))
|
||||
expr = M[i, j] + N[j, i]
|
||||
assert _convert_indexed_to_array(expr) == (ArrayAdd(M, PermuteDims(N, Permutation([1, 0]))), (i, j))
|
||||
expr = M[i, j] + M[j, i]
|
||||
assert _convert_indexed_to_array(expr) == (ArrayAdd(M, PermuteDims(M, Permutation([1, 0]))), (i, j))
|
||||
expr = (M*N*P)[i, j]
|
||||
assert _convert_indexed_to_array(expr) == (_array_contraction(ArrayTensorProduct(M, N, P), (1, 2), (3, 4)), (i, j))
|
||||
expr = expr.function # Disregard summation in previous expression
|
||||
ret1, ret2 = _convert_indexed_to_array(expr)
|
||||
assert ret1 == ArrayDiagonal(ArrayTensorProduct(M, N, P), (1, 2), (3, 4))
|
||||
assert str(ret2) == "(i, j, _i_1, _i_2)"
|
||||
expr = KroneckerDelta(i, j)*M[i, k]
|
||||
assert _convert_indexed_to_array(expr) == (M, ({i, j}, k))
|
||||
expr = KroneckerDelta(i, j)*KroneckerDelta(j, k)*M[i, l]
|
||||
assert _convert_indexed_to_array(expr) == (M, ({i, j, k}, l))
|
||||
expr = KroneckerDelta(j, k)*(M[i, j]*N[k, l] + N[i, j]*M[k, l])
|
||||
assert _convert_indexed_to_array(expr) == (_array_diagonal(_array_add(
|
||||
ArrayTensorProduct(M, N),
|
||||
_permute_dims(ArrayTensorProduct(M, N), Permutation(0, 2)(1, 3))
|
||||
), (1, 2)), (i, l, frozenset({j, k})))
|
||||
expr = KroneckerDelta(j, m)*KroneckerDelta(m, k)*(M[i, j]*N[k, l] + N[i, j]*M[k, l])
|
||||
assert _convert_indexed_to_array(expr) == (_array_diagonal(_array_add(
|
||||
ArrayTensorProduct(M, N),
|
||||
_permute_dims(ArrayTensorProduct(M, N), Permutation(0, 2)(1, 3))
|
||||
), (1, 2)), (i, l, frozenset({j, m, k})))
|
||||
expr = KroneckerDelta(i, j)*KroneckerDelta(j, k)*KroneckerDelta(k,m)*M[i, 0]*KroneckerDelta(m, n)
|
||||
assert _convert_indexed_to_array(expr) == (M, ({i, j, k, m, n}, 0))
|
||||
expr = M[i, i]
|
||||
assert _convert_indexed_to_array(expr) == (ArrayDiagonal(M, (0, 1)), (i,))
|
||||
|
||||
|
||||
def test_arrayexpr_convert_indexed_to_array_expression():
|
||||
|
||||
s = Sum(A[i]*B[i], (i, 0, 3))
|
||||
cg = convert_indexed_to_array(s)
|
||||
assert cg == ArrayContraction(ArrayTensorProduct(A, B), (0, 1))
|
||||
|
||||
expr = M*N
|
||||
result = ArrayContraction(ArrayTensorProduct(M, N), (1, 2))
|
||||
elem = expr[i, j]
|
||||
assert convert_indexed_to_array(elem) == result
|
||||
|
||||
expr = M*N*M
|
||||
elem = expr[i, j]
|
||||
result = _array_contraction(_array_tensor_product(M, M, N), (1, 4), (2, 5))
|
||||
cg = convert_indexed_to_array(elem)
|
||||
assert cg == result
|
||||
|
||||
cg = convert_indexed_to_array((M * N * P)[i, j])
|
||||
assert cg == _array_contraction(ArrayTensorProduct(M, N, P), (1, 2), (3, 4))
|
||||
|
||||
cg = convert_indexed_to_array((M * N.T * P)[i, j])
|
||||
assert cg == _array_contraction(ArrayTensorProduct(M, N, P), (1, 3), (2, 4))
|
||||
|
||||
expr = -2*M*N
|
||||
elem = expr[i, j]
|
||||
cg = convert_indexed_to_array(elem)
|
||||
assert cg == ArrayContraction(ArrayTensorProduct(-2, M, N), (1, 2))
|
||||
|
||||
|
||||
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)
|
||||
|
||||
s = tanh(A[i]*B[j])
|
||||
cg = convert_indexed_to_array(s, [i, j])
|
||||
assert cg.dummy_eq(ArrayElementwiseApplyFunc(tanh, ArrayTensorProduct(A, B)))
|
||||
|
||||
|
||||
def test_arrayexpr_convert_indexed_to_array_and_back_to_matrix():
|
||||
|
||||
expr = a.T*b
|
||||
elem = expr[0, 0]
|
||||
cg = convert_indexed_to_array(elem)
|
||||
assert cg == ArrayElement(ArrayContraction(ArrayTensorProduct(a, b), (0, 2)), [0, 0])
|
||||
|
||||
expr = M[i,j] + N[i,j]
|
||||
p1, p2 = _convert_indexed_to_array(expr)
|
||||
assert convert_array_to_matrix(p1) == M + N
|
||||
|
||||
expr = M[i,j] + N[j,i]
|
||||
p1, p2 = _convert_indexed_to_array(expr)
|
||||
assert convert_array_to_matrix(p1) == M + N.T
|
||||
|
||||
expr = M[i,j]*N[k,l] + N[i,j]*M[k,l]
|
||||
p1, p2 = _convert_indexed_to_array(expr)
|
||||
assert convert_array_to_matrix(p1) == ArrayAdd(
|
||||
ArrayTensorProduct(M, N),
|
||||
ArrayTensorProduct(N, M))
|
||||
|
||||
expr = (M*N*P)[i, j]
|
||||
p1, p2 = _convert_indexed_to_array(expr)
|
||||
assert convert_array_to_matrix(p1) == M * N * P
|
||||
|
||||
expr = Sum(M[i,j]*(N*P)[j,m], (j, 0, k-1))
|
||||
p1, p2 = _convert_indexed_to_array(expr)
|
||||
assert convert_array_to_matrix(p1) == M * N * P
|
||||
|
||||
expr = Sum((P[j, m] + P[m, j])*(M[i,j]*N[m,n] + N[i,j]*M[m,n]), (j, 0, k-1), (m, 0, k-1))
|
||||
p1, p2 = _convert_indexed_to_array(expr)
|
||||
assert convert_array_to_matrix(p1) == M * P * N + M * P.T * N + N * P * M + N * P.T * M
|
||||
|
||||
|
||||
def test_arrayexpr_convert_indexed_to_array_out_of_bounds():
|
||||
|
||||
expr = Sum(M[i, i], (i, 0, 4))
|
||||
raises(ValueError, lambda: convert_indexed_to_array(expr))
|
||||
expr = Sum(M[i, i], (i, 0, k))
|
||||
raises(ValueError, lambda: convert_indexed_to_array(expr))
|
||||
expr = Sum(M[i, i], (i, 1, k-1))
|
||||
raises(ValueError, lambda: convert_indexed_to_array(expr))
|
||||
|
||||
expr = Sum(M[i, j]*N[j,m], (j, 0, 4))
|
||||
raises(ValueError, lambda: convert_indexed_to_array(expr))
|
||||
expr = Sum(M[i, j]*N[j,m], (j, 0, k))
|
||||
raises(ValueError, lambda: convert_indexed_to_array(expr))
|
||||
expr = Sum(M[i, j]*N[j,m], (j, 1, k-1))
|
||||
raises(ValueError, lambda: convert_indexed_to_array(expr))
|
||||
|
||||
|
||||
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]))
|
||||
+128
@@ -0,0 +1,128 @@
|
||||
from sympy import Lambda, KroneckerProduct
|
||||
from sympy.core.symbol import symbols, Dummy
|
||||
from sympy.matrices.expressions.hadamard import (HadamardPower, HadamardProduct)
|
||||
from sympy.matrices.expressions.inverse import Inverse
|
||||
from sympy.matrices.expressions.matexpr import MatrixSymbol
|
||||
from sympy.matrices.expressions.matpow import MatPow
|
||||
from sympy.matrices.expressions.special import Identity
|
||||
from sympy.matrices.expressions.trace import Trace
|
||||
from sympy.matrices.expressions.transpose import Transpose
|
||||
from sympy.tensor.array.expressions.array_expressions import ArrayTensorProduct, ArrayContraction, \
|
||||
PermuteDims, ArrayDiagonal, ArrayElementwiseApplyFunc, _array_contraction, _array_tensor_product, Reshape
|
||||
from sympy.tensor.array.expressions.from_array_to_matrix import convert_array_to_matrix
|
||||
from sympy.tensor.array.expressions.from_matrix_to_array import convert_matrix_to_array
|
||||
|
||||
i, j, k, l, m, n = symbols("i j k l m n")
|
||||
|
||||
I = Identity(k)
|
||||
|
||||
M = MatrixSymbol("M", k, k)
|
||||
N = MatrixSymbol("N", k, k)
|
||||
P = MatrixSymbol("P", k, k)
|
||||
Q = MatrixSymbol("Q", k, k)
|
||||
|
||||
A = MatrixSymbol("A", k, k)
|
||||
B = MatrixSymbol("B", k, k)
|
||||
C = MatrixSymbol("C", k, k)
|
||||
D = MatrixSymbol("D", k, k)
|
||||
|
||||
X = MatrixSymbol("X", k, k)
|
||||
Y = MatrixSymbol("Y", k, k)
|
||||
|
||||
a = MatrixSymbol("a", k, 1)
|
||||
b = MatrixSymbol("b", k, 1)
|
||||
c = MatrixSymbol("c", k, 1)
|
||||
d = MatrixSymbol("d", k, 1)
|
||||
|
||||
|
||||
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))
|
||||
|
||||
expr = KroneckerProduct(A, B)
|
||||
cg = convert_matrix_to_array(expr)
|
||||
assert cg == Reshape(PermuteDims(ArrayTensorProduct(A, B), [0, 2, 1, 3]), (k**2, k**2))
|
||||
|
||||
expr = KroneckerProduct(A, B, C, D)
|
||||
cg = convert_matrix_to_array(expr)
|
||||
assert cg == Reshape(PermuteDims(ArrayTensorProduct(A, B, C, D), [0, 2, 4, 6, 1, 3, 5, 7]), (k**4, k**4))
|
||||
+22
@@ -0,0 +1,22 @@
|
||||
from sympy import MatrixSymbol, symbols, Sum
|
||||
from sympy.tensor.array.expressions import conv_array_to_indexed, from_array_to_indexed, ArrayTensorProduct, \
|
||||
ArrayContraction, conv_array_to_matrix, from_array_to_matrix, conv_matrix_to_array, from_matrix_to_array, \
|
||||
conv_indexed_to_array, from_indexed_to_array
|
||||
from sympy.testing.pytest import warns
|
||||
from sympy.utilities.exceptions import SymPyDeprecationWarning
|
||||
|
||||
|
||||
def test_deprecated_conv_module_results():
|
||||
|
||||
M = MatrixSymbol("M", 3, 3)
|
||||
N = MatrixSymbol("N", 3, 3)
|
||||
i, j, d = symbols("i j d")
|
||||
|
||||
x = ArrayContraction(ArrayTensorProduct(M, N), (1, 2))
|
||||
y = Sum(M[i, d]*N[d, j], (d, 0, 2))
|
||||
|
||||
with warns(SymPyDeprecationWarning, test_stacklevel=False):
|
||||
assert conv_array_to_indexed.convert_array_to_indexed(x, [i, j]).dummy_eq(from_array_to_indexed.convert_array_to_indexed(x, [i, j]))
|
||||
assert conv_array_to_matrix.convert_array_to_matrix(x) == from_array_to_matrix.convert_array_to_matrix(x)
|
||||
assert conv_matrix_to_array.convert_matrix_to_array(M*N) == from_matrix_to_array.convert_matrix_to_array(M*N)
|
||||
assert conv_indexed_to_array.convert_indexed_to_array(y) == from_indexed_to_array.convert_indexed_to_array(y)
|
||||
@@ -0,0 +1,123 @@
|
||||
import bisect
|
||||
from collections import defaultdict
|
||||
|
||||
from sympy.combinatorics import Permutation
|
||||
from sympy.core.containers import Tuple
|
||||
from sympy.core.numbers import Integer
|
||||
|
||||
|
||||
def _get_mapping_from_subranks(subranks):
|
||||
mapping = {}
|
||||
counter = 0
|
||||
for i, rank in enumerate(subranks):
|
||||
for j in range(rank):
|
||||
mapping[counter] = (i, j)
|
||||
counter += 1
|
||||
return mapping
|
||||
|
||||
|
||||
def _get_contraction_links(args, subranks, *contraction_indices):
|
||||
mapping = _get_mapping_from_subranks(subranks)
|
||||
contraction_tuples = [[mapping[j] for j in i] for i in contraction_indices]
|
||||
dlinks = defaultdict(dict)
|
||||
for links in contraction_tuples:
|
||||
if len(links) == 2:
|
||||
(arg1, pos1), (arg2, pos2) = links
|
||||
dlinks[arg1][pos1] = (arg2, pos2)
|
||||
dlinks[arg2][pos2] = (arg1, pos1)
|
||||
continue
|
||||
|
||||
return args, dict(dlinks)
|
||||
|
||||
|
||||
def _sort_contraction_indices(pairing_indices):
|
||||
pairing_indices = [Tuple(*sorted(i)) for i in pairing_indices]
|
||||
pairing_indices.sort(key=lambda x: min(x))
|
||||
return pairing_indices
|
||||
|
||||
|
||||
def _get_diagonal_indices(flattened_indices):
|
||||
axes_contraction = defaultdict(list)
|
||||
for i, ind in enumerate(flattened_indices):
|
||||
if isinstance(ind, (int, Integer)):
|
||||
# If the indices is a number, there can be no diagonal operation:
|
||||
continue
|
||||
axes_contraction[ind].append(i)
|
||||
axes_contraction = {k: v for k, v in axes_contraction.items() if len(v) > 1}
|
||||
# Put the diagonalized indices at the end:
|
||||
ret_indices = [i for i in flattened_indices if i not in axes_contraction]
|
||||
diag_indices = list(axes_contraction)
|
||||
diag_indices.sort(key=lambda x: flattened_indices.index(x))
|
||||
diagonal_indices = [tuple(axes_contraction[i]) for i in diag_indices]
|
||||
ret_indices += diag_indices
|
||||
ret_indices = tuple(ret_indices)
|
||||
return diagonal_indices, ret_indices
|
||||
|
||||
|
||||
def _get_argindex(subindices, ind):
|
||||
for i, sind in enumerate(subindices):
|
||||
if ind == sind:
|
||||
return i
|
||||
if isinstance(sind, (set, frozenset)) and ind in sind:
|
||||
return i
|
||||
raise IndexError("%s not found in %s" % (ind, subindices))
|
||||
|
||||
|
||||
def _apply_recursively_over_nested_lists(func, arr):
|
||||
if isinstance(arr, (tuple, list, Tuple)):
|
||||
return tuple(_apply_recursively_over_nested_lists(func, i) for i in arr)
|
||||
elif isinstance(arr, Tuple):
|
||||
return Tuple.fromiter(_apply_recursively_over_nested_lists(func, i) for i in arr)
|
||||
else:
|
||||
return func(arr)
|
||||
|
||||
|
||||
def _build_push_indices_up_func_transformation(flattened_contraction_indices):
|
||||
shifts = {0: 0}
|
||||
i = 0
|
||||
cumulative = 0
|
||||
while i < len(flattened_contraction_indices):
|
||||
j = 1
|
||||
while i+j < len(flattened_contraction_indices):
|
||||
if flattened_contraction_indices[i] + j != flattened_contraction_indices[i+j]:
|
||||
break
|
||||
j += 1
|
||||
cumulative += j
|
||||
shifts[flattened_contraction_indices[i]] = cumulative
|
||||
i += j
|
||||
shift_keys = sorted(shifts.keys())
|
||||
|
||||
def func(idx):
|
||||
return shifts[shift_keys[bisect.bisect_right(shift_keys, idx)-1]]
|
||||
|
||||
def transform(j):
|
||||
if j in flattened_contraction_indices:
|
||||
return None
|
||||
else:
|
||||
return j - func(j)
|
||||
|
||||
return transform
|
||||
|
||||
|
||||
def _build_push_indices_down_func_transformation(flattened_contraction_indices):
|
||||
N = flattened_contraction_indices[-1]+2
|
||||
|
||||
shifts = [i for i in range(N) if i not in flattened_contraction_indices]
|
||||
|
||||
def transform(j):
|
||||
if j < len(shifts):
|
||||
return shifts[j]
|
||||
else:
|
||||
return j + shifts[-1] - len(shifts) + 1
|
||||
|
||||
return transform
|
||||
|
||||
|
||||
def _apply_permutation_to_list(perm: Permutation, target_list: list):
|
||||
"""
|
||||
Permute a list according to the given permutation.
|
||||
"""
|
||||
new_list = [None for i in range(perm.size)]
|
||||
for i, e in enumerate(target_list):
|
||||
new_list[perm(i)] = e
|
||||
return new_list
|
||||
Reference in New Issue
Block a user