Initialisation du repository de Beta
This commit is contained in:
commit
14985f6dbb
9469 changed files with 1903273 additions and 0 deletions
|
|
@ -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
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
File diff suppressed because it is too large
Load diff
|
|
@ -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()
|
||||
|
|
@ -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)
|
||||
|
|
@ -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)
|
||||
|
|
@ -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)
|
||||
|
|
@ -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)
|
||||
|
|
@ -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)
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -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, ()
|
||||
|
|
@ -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
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -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)
|
||||
|
|
@ -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))
|
||||
|
|
@ -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()
|
||||
|
|
@ -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]
|
||||
|
|
@ -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])
|
||||
|
|
@ -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]))
|
||||
|
|
@ -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))
|
||||
|
|
@ -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
|
||||
Loading…
Add table
Add a link
Reference in a new issue