I am done

This commit is contained in:
2024-10-30 22:14:35 +01:00
parent 720dc28c09
commit 40e2a747cf
36901 changed files with 5011519 additions and 0 deletions

View File

@ -0,0 +1,271 @@
r"""
N-dim array module for SymPy.
Four classes are provided to handle N-dim arrays, given by the combinations
dense/sparse (i.e. whether to store all elements or only the non-zero ones in
memory) and mutable/immutable (immutable classes are SymPy objects, but cannot
change after they have been created).
Examples
========
The following examples show the usage of ``Array``. This is an abbreviation for
``ImmutableDenseNDimArray``, that is an immutable and dense N-dim array, the
other classes are analogous. For mutable classes it is also possible to change
element values after the object has been constructed.
Array construction can detect the shape of nested lists and tuples:
>>> from sympy import Array
>>> a1 = Array([[1, 2], [3, 4], [5, 6]])
>>> a1
[[1, 2], [3, 4], [5, 6]]
>>> a1.shape
(3, 2)
>>> a1.rank()
2
>>> from sympy.abc import x, y, z
>>> a2 = Array([[[x, y], [z, x*z]], [[1, x*y], [1/x, x/y]]])
>>> a2
[[[x, y], [z, x*z]], [[1, x*y], [1/x, x/y]]]
>>> a2.shape
(2, 2, 2)
>>> a2.rank()
3
Otherwise one could pass a 1-dim array followed by a shape tuple:
>>> m1 = Array(range(12), (3, 4))
>>> m1
[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]
>>> m2 = Array(range(12), (3, 2, 2))
>>> m2
[[[0, 1], [2, 3]], [[4, 5], [6, 7]], [[8, 9], [10, 11]]]
>>> m2[1,1,1]
7
>>> m2.reshape(4, 3)
[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]]
Slice support:
>>> m2[:, 1, 1]
[3, 7, 11]
Elementwise derivative:
>>> from sympy.abc import x, y, z
>>> m3 = Array([x**3, x*y, z])
>>> m3.diff(x)
[3*x**2, y, 0]
>>> m3.diff(z)
[0, 0, 1]
Multiplication with other SymPy expressions is applied elementwisely:
>>> (1+x)*m3
[x**3*(x + 1), x*y*(x + 1), z*(x + 1)]
To apply a function to each element of the N-dim array, use ``applyfunc``:
>>> m3.applyfunc(lambda x: x/2)
[x**3/2, x*y/2, z/2]
N-dim arrays can be converted to nested lists by the ``tolist()`` method:
>>> m2.tolist()
[[[0, 1], [2, 3]], [[4, 5], [6, 7]], [[8, 9], [10, 11]]]
>>> isinstance(m2.tolist(), list)
True
If the rank is 2, it is possible to convert them to matrices with ``tomatrix()``:
>>> m1.tomatrix()
Matrix([
[0, 1, 2, 3],
[4, 5, 6, 7],
[8, 9, 10, 11]])
Products and contractions
-------------------------
Tensor product between arrays `A_{i_1,\ldots,i_n}` and `B_{j_1,\ldots,j_m}`
creates the combined array `P = A \otimes B` defined as
`P_{i_1,\ldots,i_n,j_1,\ldots,j_m} := A_{i_1,\ldots,i_n}\cdot B_{j_1,\ldots,j_m}.`
It is available through ``tensorproduct(...)``:
>>> from sympy import Array, tensorproduct
>>> from sympy.abc import x,y,z,t
>>> A = Array([x, y, z, t])
>>> B = Array([1, 2, 3, 4])
>>> tensorproduct(A, B)
[[x, 2*x, 3*x, 4*x], [y, 2*y, 3*y, 4*y], [z, 2*z, 3*z, 4*z], [t, 2*t, 3*t, 4*t]]
In case you don't want to evaluate the tensor product immediately, you can use
``ArrayTensorProduct``, which creates an unevaluated tensor product expression:
>>> from sympy.tensor.array.expressions import ArrayTensorProduct
>>> ArrayTensorProduct(A, B)
ArrayTensorProduct([x, y, z, t], [1, 2, 3, 4])
Calling ``.as_explicit()`` on ``ArrayTensorProduct`` is equivalent to just calling
``tensorproduct(...)``:
>>> ArrayTensorProduct(A, B).as_explicit()
[[x, 2*x, 3*x, 4*x], [y, 2*y, 3*y, 4*y], [z, 2*z, 3*z, 4*z], [t, 2*t, 3*t, 4*t]]
Tensor product between a rank-1 array and a matrix creates a rank-3 array:
>>> from sympy import eye
>>> p1 = tensorproduct(A, eye(4))
>>> p1
[[[x, 0, 0, 0], [0, x, 0, 0], [0, 0, x, 0], [0, 0, 0, x]], [[y, 0, 0, 0], [0, y, 0, 0], [0, 0, y, 0], [0, 0, 0, y]], [[z, 0, 0, 0], [0, z, 0, 0], [0, 0, z, 0], [0, 0, 0, z]], [[t, 0, 0, 0], [0, t, 0, 0], [0, 0, t, 0], [0, 0, 0, t]]]
Now, to get back `A_0 \otimes \mathbf{1}` one can access `p_{0,m,n}` by slicing:
>>> p1[0,:,:]
[[x, 0, 0, 0], [0, x, 0, 0], [0, 0, x, 0], [0, 0, 0, x]]
Tensor contraction sums over the specified axes, for example contracting
positions `a` and `b` means
`A_{i_1,\ldots,i_a,\ldots,i_b,\ldots,i_n} \implies \sum_k A_{i_1,\ldots,k,\ldots,k,\ldots,i_n}`
Remember that Python indexing is zero starting, to contract the a-th and b-th
axes it is therefore necessary to specify `a-1` and `b-1`
>>> from sympy import tensorcontraction
>>> C = Array([[x, y], [z, t]])
The matrix trace is equivalent to the contraction of a rank-2 array:
`A_{m,n} \implies \sum_k A_{k,k}`
>>> tensorcontraction(C, (0, 1))
t + x
To create an expression representing a tensor contraction that does not get
evaluated immediately, use ``ArrayContraction``, which is equivalent to
``tensorcontraction(...)`` if it is followed by ``.as_explicit()``:
>>> from sympy.tensor.array.expressions import ArrayContraction
>>> ArrayContraction(C, (0, 1))
ArrayContraction([[x, y], [z, t]], (0, 1))
>>> ArrayContraction(C, (0, 1)).as_explicit()
t + x
Matrix product is equivalent to a tensor product of two rank-2 arrays, followed
by a contraction of the 2nd and 3rd axes (in Python indexing axes number 1, 2).
`A_{m,n}\cdot B_{i,j} \implies \sum_k A_{m, k}\cdot B_{k, j}`
>>> D = Array([[2, 1], [0, -1]])
>>> tensorcontraction(tensorproduct(C, D), (1, 2))
[[2*x, x - y], [2*z, -t + z]]
One may verify that the matrix product is equivalent:
>>> from sympy import Matrix
>>> Matrix([[x, y], [z, t]])*Matrix([[2, 1], [0, -1]])
Matrix([
[2*x, x - y],
[2*z, -t + z]])
or equivalently
>>> C.tomatrix()*D.tomatrix()
Matrix([
[2*x, x - y],
[2*z, -t + z]])
Diagonal operator
-----------------
The ``tensordiagonal`` function acts in a similar manner as ``tensorcontraction``,
but the joined indices are not summed over, for example diagonalizing
positions `a` and `b` means
`A_{i_1,\ldots,i_a,\ldots,i_b,\ldots,i_n} \implies A_{i_1,\ldots,k,\ldots,k,\ldots,i_n}
\implies \tilde{A}_{i_1,\ldots,i_{a-1},i_{a+1},\ldots,i_{b-1},i_{b+1},\ldots,i_n,k}`
where `\tilde{A}` is the array equivalent to the diagonal of `A` at positions
`a` and `b` moved to the last index slot.
Compare the difference between contraction and diagonal operators:
>>> from sympy import tensordiagonal
>>> from sympy.abc import a, b, c, d
>>> m = Matrix([[a, b], [c, d]])
>>> tensorcontraction(m, [0, 1])
a + d
>>> tensordiagonal(m, [0, 1])
[a, d]
In short, no summation occurs with ``tensordiagonal``.
Derivatives by array
--------------------
The usual derivative operation may be extended to support derivation with
respect to arrays, provided that all elements in the that array are symbols or
expressions suitable for derivations.
The definition of a derivative by an array is as follows: given the array
`A_{i_1, \ldots, i_N}` and the array `X_{j_1, \ldots, j_M}`
the derivative of arrays will return a new array `B` defined by
`B_{j_1,\ldots,j_M,i_1,\ldots,i_N} := \frac{\partial A_{i_1,\ldots,i_N}}{\partial X_{j_1,\ldots,j_M}}`
The function ``derive_by_array`` performs such an operation:
>>> from sympy import derive_by_array
>>> from sympy.abc import x, y, z, t
>>> from sympy import sin, exp
With scalars, it behaves exactly as the ordinary derivative:
>>> derive_by_array(sin(x*y), x)
y*cos(x*y)
Scalar derived by an array basis:
>>> derive_by_array(sin(x*y), [x, y, z])
[y*cos(x*y), x*cos(x*y), 0]
Deriving array by an array basis: `B^{nm} := \frac{\partial A^m}{\partial x^n}`
>>> basis = [x, y, z]
>>> ax = derive_by_array([exp(x), sin(y*z), t], basis)
>>> ax
[[exp(x), 0, 0], [0, z*cos(y*z), 0], [0, y*cos(y*z), 0]]
Contraction of the resulting array: `\sum_m \frac{\partial A^m}{\partial x^m}`
>>> tensorcontraction(ax, (0, 1))
z*cos(y*z) + exp(x)
"""
from .dense_ndim_array import MutableDenseNDimArray, ImmutableDenseNDimArray, DenseNDimArray
from .sparse_ndim_array import MutableSparseNDimArray, ImmutableSparseNDimArray, SparseNDimArray
from .ndim_array import NDimArray, ArrayKind
from .arrayop import tensorproduct, tensorcontraction, tensordiagonal, derive_by_array, permutedims
from .array_comprehension import ArrayComprehension, ArrayComprehensionMap
Array = ImmutableDenseNDimArray
__all__ = [
'MutableDenseNDimArray', 'ImmutableDenseNDimArray', 'DenseNDimArray',
'MutableSparseNDimArray', 'ImmutableSparseNDimArray', 'SparseNDimArray',
'NDimArray', 'ArrayKind',
'tensorproduct', 'tensorcontraction', 'tensordiagonal', 'derive_by_array',
'permutedims', 'ArrayComprehension', 'ArrayComprehensionMap',
'Array',
]

View File

@ -0,0 +1,399 @@
import functools, itertools
from sympy.core.sympify import _sympify, sympify
from sympy.core.expr import Expr
from sympy.core import Basic, Tuple
from sympy.tensor.array import ImmutableDenseNDimArray
from sympy.core.symbol import Symbol
from sympy.core.numbers import Integer
class ArrayComprehension(Basic):
"""
Generate a list comprehension.
Explanation
===========
If there is a symbolic dimension, for example, say [i for i in range(1, N)] where
N is a Symbol, then the expression will not be expanded to an array. Otherwise,
calling the doit() function will launch the expansion.
Examples
========
>>> from sympy.tensor.array import ArrayComprehension
>>> from sympy import symbols
>>> i, j, k = symbols('i j k')
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
>>> a
ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
>>> a.doit()
[[11, 12, 13], [21, 22, 23], [31, 32, 33], [41, 42, 43]]
>>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k))
>>> b.doit()
ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k))
"""
def __new__(cls, function, *symbols, **assumptions):
if any(len(l) != 3 or None for l in symbols):
raise ValueError('ArrayComprehension requires values lower and upper bound'
' for the expression')
arglist = [sympify(function)]
arglist.extend(cls._check_limits_validity(function, symbols))
obj = Basic.__new__(cls, *arglist, **assumptions)
obj._limits = obj._args[1:]
obj._shape = cls._calculate_shape_from_limits(obj._limits)
obj._rank = len(obj._shape)
obj._loop_size = cls._calculate_loop_size(obj._shape)
return obj
@property
def function(self):
"""The function applied across limits.
Examples
========
>>> from sympy.tensor.array import ArrayComprehension
>>> from sympy import symbols
>>> i, j = symbols('i j')
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
>>> a.function
10*i + j
"""
return self._args[0]
@property
def limits(self):
"""
The list of limits that will be applied while expanding the array.
Examples
========
>>> from sympy.tensor.array import ArrayComprehension
>>> from sympy import symbols
>>> i, j = symbols('i j')
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
>>> a.limits
((i, 1, 4), (j, 1, 3))
"""
return self._limits
@property
def free_symbols(self):
"""
The set of the free_symbols in the array.
Variables appeared in the bounds are supposed to be excluded
from the free symbol set.
Examples
========
>>> from sympy.tensor.array import ArrayComprehension
>>> from sympy import symbols
>>> i, j, k = symbols('i j k')
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
>>> a.free_symbols
set()
>>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k+3))
>>> b.free_symbols
{k}
"""
expr_free_sym = self.function.free_symbols
for var, inf, sup in self._limits:
expr_free_sym.discard(var)
curr_free_syms = inf.free_symbols.union(sup.free_symbols)
expr_free_sym = expr_free_sym.union(curr_free_syms)
return expr_free_sym
@property
def variables(self):
"""The tuples of the variables in the limits.
Examples
========
>>> from sympy.tensor.array import ArrayComprehension
>>> from sympy import symbols
>>> i, j, k = symbols('i j k')
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
>>> a.variables
[i, j]
"""
return [l[0] for l in self._limits]
@property
def bound_symbols(self):
"""The list of dummy variables.
Note
====
Note that all variables are dummy variables since a limit without
lower bound or upper bound is not accepted.
"""
return [l[0] for l in self._limits if len(l) != 1]
@property
def shape(self):
"""
The shape of the expanded array, which may have symbols.
Note
====
Both the lower and the upper bounds are included while
calculating the shape.
Examples
========
>>> from sympy.tensor.array import ArrayComprehension
>>> from sympy import symbols
>>> i, j, k = symbols('i j k')
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
>>> a.shape
(4, 3)
>>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k+3))
>>> b.shape
(4, k + 3)
"""
return self._shape
@property
def is_shape_numeric(self):
"""
Test if the array is shape-numeric which means there is no symbolic
dimension.
Examples
========
>>> from sympy.tensor.array import ArrayComprehension
>>> from sympy import symbols
>>> i, j, k = symbols('i j k')
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
>>> a.is_shape_numeric
True
>>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k+3))
>>> b.is_shape_numeric
False
"""
for _, inf, sup in self._limits:
if Basic(inf, sup).atoms(Symbol):
return False
return True
def rank(self):
"""The rank of the expanded array.
Examples
========
>>> from sympy.tensor.array import ArrayComprehension
>>> from sympy import symbols
>>> i, j, k = symbols('i j k')
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
>>> a.rank()
2
"""
return self._rank
def __len__(self):
"""
The length of the expanded array which means the number
of elements in the array.
Raises
======
ValueError : When the length of the array is symbolic
Examples
========
>>> from sympy.tensor.array import ArrayComprehension
>>> from sympy import symbols
>>> i, j = symbols('i j')
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
>>> len(a)
12
"""
if self._loop_size.free_symbols:
raise ValueError('Symbolic length is not supported')
return self._loop_size
@classmethod
def _check_limits_validity(cls, function, limits):
#limits = sympify(limits)
new_limits = []
for var, inf, sup in limits:
var = _sympify(var)
inf = _sympify(inf)
#since this is stored as an argument, it should be
#a Tuple
if isinstance(sup, list):
sup = Tuple(*sup)
else:
sup = _sympify(sup)
new_limits.append(Tuple(var, inf, sup))
if any((not isinstance(i, Expr)) or i.atoms(Symbol, Integer) != i.atoms()
for i in [inf, sup]):
raise TypeError('Bounds should be an Expression(combination of Integer and Symbol)')
if (inf > sup) == True:
raise ValueError('Lower bound should be inferior to upper bound')
if var in inf.free_symbols or var in sup.free_symbols:
raise ValueError('Variable should not be part of its bounds')
return new_limits
@classmethod
def _calculate_shape_from_limits(cls, limits):
return tuple([sup - inf + 1 for _, inf, sup in limits])
@classmethod
def _calculate_loop_size(cls, shape):
if not shape:
return 0
loop_size = 1
for l in shape:
loop_size = loop_size * l
return loop_size
def doit(self, **hints):
if not self.is_shape_numeric:
return self
return self._expand_array()
def _expand_array(self):
res = []
for values in itertools.product(*[range(inf, sup+1)
for var, inf, sup
in self._limits]):
res.append(self._get_element(values))
return ImmutableDenseNDimArray(res, self.shape)
def _get_element(self, values):
temp = self.function
for var, val in zip(self.variables, values):
temp = temp.subs(var, val)
return temp
def tolist(self):
"""Transform the expanded array to a list.
Raises
======
ValueError : When there is a symbolic dimension
Examples
========
>>> from sympy.tensor.array import ArrayComprehension
>>> from sympy import symbols
>>> i, j = symbols('i j')
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
>>> a.tolist()
[[11, 12, 13], [21, 22, 23], [31, 32, 33], [41, 42, 43]]
"""
if self.is_shape_numeric:
return self._expand_array().tolist()
raise ValueError("A symbolic array cannot be expanded to a list")
def tomatrix(self):
"""Transform the expanded array to a matrix.
Raises
======
ValueError : When there is a symbolic dimension
ValueError : When the rank of the expanded array is not equal to 2
Examples
========
>>> from sympy.tensor.array import ArrayComprehension
>>> from sympy import symbols
>>> i, j = symbols('i j')
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
>>> a.tomatrix()
Matrix([
[11, 12, 13],
[21, 22, 23],
[31, 32, 33],
[41, 42, 43]])
"""
from sympy.matrices import Matrix
if not self.is_shape_numeric:
raise ValueError("A symbolic array cannot be expanded to a matrix")
if self._rank != 2:
raise ValueError('Dimensions must be of size of 2')
return Matrix(self._expand_array().tomatrix())
def isLambda(v):
LAMBDA = lambda: 0
return isinstance(v, type(LAMBDA)) and v.__name__ == LAMBDA.__name__
class ArrayComprehensionMap(ArrayComprehension):
'''
A subclass of ArrayComprehension dedicated to map external function lambda.
Notes
=====
Only the lambda function is considered.
At most one argument in lambda function is accepted in order to avoid ambiguity
in value assignment.
Examples
========
>>> from sympy.tensor.array import ArrayComprehensionMap
>>> from sympy import symbols
>>> i, j, k = symbols('i j k')
>>> a = ArrayComprehensionMap(lambda: 1, (i, 1, 4))
>>> a.doit()
[1, 1, 1, 1]
>>> b = ArrayComprehensionMap(lambda a: a+1, (j, 1, 4))
>>> b.doit()
[2, 3, 4, 5]
'''
def __new__(cls, function, *symbols, **assumptions):
if any(len(l) != 3 or None for l in symbols):
raise ValueError('ArrayComprehension requires values lower and upper bound'
' for the expression')
if not isLambda(function):
raise ValueError('Data type not supported')
arglist = cls._check_limits_validity(function, symbols)
obj = Basic.__new__(cls, *arglist, **assumptions)
obj._limits = obj._args
obj._shape = cls._calculate_shape_from_limits(obj._limits)
obj._rank = len(obj._shape)
obj._loop_size = cls._calculate_loop_size(obj._shape)
obj._lambda = function
return obj
@property
def func(self):
class _(ArrayComprehensionMap):
def __new__(cls, *args, **kwargs):
return ArrayComprehensionMap(self._lambda, *args, **kwargs)
return _
def _get_element(self, values):
temp = self._lambda
if self._lambda.__code__.co_argcount == 0:
temp = temp()
elif self._lambda.__code__.co_argcount == 1:
temp = temp(functools.reduce(lambda a, b: a*b, values))
return temp

View File

@ -0,0 +1,129 @@
from __future__ import annotations
from sympy.core.expr import Expr
from sympy.core.function import Derivative
from sympy.core.numbers import Integer
from sympy.matrices.matrixbase import MatrixBase
from .ndim_array import NDimArray
from .arrayop import derive_by_array
from sympy.matrices.expressions.matexpr import MatrixExpr
from sympy.matrices.expressions.special import ZeroMatrix
from sympy.matrices.expressions.matexpr import _matrix_derivative
class ArrayDerivative(Derivative):
is_scalar = False
def __new__(cls, expr, *variables, **kwargs):
obj = super().__new__(cls, expr, *variables, **kwargs)
if isinstance(obj, ArrayDerivative):
obj._shape = obj._get_shape()
return obj
def _get_shape(self):
shape = ()
for v, count in self.variable_count:
if hasattr(v, "shape"):
for i in range(count):
shape += v.shape
if hasattr(self.expr, "shape"):
shape += self.expr.shape
return shape
@property
def shape(self):
return self._shape
@classmethod
def _get_zero_with_shape_like(cls, expr):
if isinstance(expr, (MatrixBase, NDimArray)):
return expr.zeros(*expr.shape)
elif isinstance(expr, MatrixExpr):
return ZeroMatrix(*expr.shape)
else:
raise RuntimeError("Unable to determine shape of array-derivative.")
@staticmethod
def _call_derive_scalar_by_matrix(expr: Expr, v: MatrixBase) -> Expr:
return v.applyfunc(lambda x: expr.diff(x))
@staticmethod
def _call_derive_scalar_by_matexpr(expr: Expr, v: MatrixExpr) -> Expr:
if expr.has(v):
return _matrix_derivative(expr, v)
else:
return ZeroMatrix(*v.shape)
@staticmethod
def _call_derive_scalar_by_array(expr: Expr, v: NDimArray) -> Expr:
return v.applyfunc(lambda x: expr.diff(x))
@staticmethod
def _call_derive_matrix_by_scalar(expr: MatrixBase, v: Expr) -> Expr:
return _matrix_derivative(expr, v)
@staticmethod
def _call_derive_matexpr_by_scalar(expr: MatrixExpr, v: Expr) -> Expr:
return expr._eval_derivative(v)
@staticmethod
def _call_derive_array_by_scalar(expr: NDimArray, v: Expr) -> Expr:
return expr.applyfunc(lambda x: x.diff(v))
@staticmethod
def _call_derive_default(expr: Expr, v: Expr) -> Expr | None:
if expr.has(v):
return _matrix_derivative(expr, v)
else:
return None
@classmethod
def _dispatch_eval_derivative_n_times(cls, expr, v, count):
# Evaluate the derivative `n` times. If
# `_eval_derivative_n_times` is not overridden by the current
# object, the default in `Basic` will call a loop over
# `_eval_derivative`:
if not isinstance(count, (int, Integer)) or ((count <= 0) == True):
return None
# TODO: this could be done with multiple-dispatching:
if expr.is_scalar:
if isinstance(v, MatrixBase):
result = cls._call_derive_scalar_by_matrix(expr, v)
elif isinstance(v, MatrixExpr):
result = cls._call_derive_scalar_by_matexpr(expr, v)
elif isinstance(v, NDimArray):
result = cls._call_derive_scalar_by_array(expr, v)
elif v.is_scalar:
# scalar by scalar has a special
return super()._dispatch_eval_derivative_n_times(expr, v, count)
else:
return None
elif v.is_scalar:
if isinstance(expr, MatrixBase):
result = cls._call_derive_matrix_by_scalar(expr, v)
elif isinstance(expr, MatrixExpr):
result = cls._call_derive_matexpr_by_scalar(expr, v)
elif isinstance(expr, NDimArray):
result = cls._call_derive_array_by_scalar(expr, v)
else:
return None
else:
# Both `expr` and `v` are some array/matrix type:
if isinstance(expr, MatrixBase) or isinstance(v, MatrixBase):
result = derive_by_array(expr, v)
elif isinstance(expr, MatrixExpr) and isinstance(v, MatrixExpr):
result = cls._call_derive_default(expr, v)
elif isinstance(expr, MatrixExpr) or isinstance(v, MatrixExpr):
# if one expression is a symbolic matrix expression while the other isn't, don't evaluate:
return None
else:
result = derive_by_array(expr, v)
if result is None:
return None
if count == 1:
return result
else:
return cls._dispatch_eval_derivative_n_times(result, v, count - 1)

View File

@ -0,0 +1,528 @@
import itertools
from collections.abc import Iterable
from sympy.core._print_helpers import Printable
from sympy.core.containers import Tuple
from sympy.core.function import diff
from sympy.core.singleton import S
from sympy.core.sympify import _sympify
from sympy.tensor.array.ndim_array import NDimArray
from sympy.tensor.array.dense_ndim_array import DenseNDimArray, ImmutableDenseNDimArray
from sympy.tensor.array.sparse_ndim_array import SparseNDimArray
def _arrayfy(a):
from sympy.matrices import MatrixBase
if isinstance(a, NDimArray):
return a
if isinstance(a, (MatrixBase, list, tuple, Tuple)):
return ImmutableDenseNDimArray(a)
return a
def tensorproduct(*args):
"""
Tensor product among scalars or array-like objects.
The equivalent operator for array expressions is ``ArrayTensorProduct``,
which can be used to keep the expression unevaluated.
Examples
========
>>> from sympy.tensor.array import tensorproduct, Array
>>> from sympy.abc import x, y, z, t
>>> A = Array([[1, 2], [3, 4]])
>>> B = Array([x, y])
>>> tensorproduct(A, B)
[[[x, y], [2*x, 2*y]], [[3*x, 3*y], [4*x, 4*y]]]
>>> tensorproduct(A, x)
[[x, 2*x], [3*x, 4*x]]
>>> tensorproduct(A, B, B)
[[[[x**2, x*y], [x*y, y**2]], [[2*x**2, 2*x*y], [2*x*y, 2*y**2]]], [[[3*x**2, 3*x*y], [3*x*y, 3*y**2]], [[4*x**2, 4*x*y], [4*x*y, 4*y**2]]]]
Applying this function on two matrices will result in a rank 4 array.
>>> from sympy import Matrix, eye
>>> m = Matrix([[x, y], [z, t]])
>>> p = tensorproduct(eye(3), m)
>>> p
[[[[x, y], [z, t]], [[0, 0], [0, 0]], [[0, 0], [0, 0]]], [[[0, 0], [0, 0]], [[x, y], [z, t]], [[0, 0], [0, 0]]], [[[0, 0], [0, 0]], [[0, 0], [0, 0]], [[x, y], [z, t]]]]
See Also
========
sympy.tensor.array.expressions.array_expressions.ArrayTensorProduct
"""
from sympy.tensor.array import SparseNDimArray, ImmutableSparseNDimArray
if len(args) == 0:
return S.One
if len(args) == 1:
return _arrayfy(args[0])
from sympy.tensor.array.expressions.array_expressions import _CodegenArrayAbstract
from sympy.tensor.array.expressions.array_expressions import ArrayTensorProduct
from sympy.tensor.array.expressions.array_expressions import _ArrayExpr
from sympy.matrices.expressions.matexpr import MatrixSymbol
if any(isinstance(arg, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)) for arg in args):
return ArrayTensorProduct(*args)
if len(args) > 2:
return tensorproduct(tensorproduct(args[0], args[1]), *args[2:])
# length of args is 2:
a, b = map(_arrayfy, args)
if not isinstance(a, NDimArray) or not isinstance(b, NDimArray):
return a*b
if isinstance(a, SparseNDimArray) and isinstance(b, SparseNDimArray):
lp = len(b)
new_array = {k1*lp + k2: v1*v2 for k1, v1 in a._sparse_array.items() for k2, v2 in b._sparse_array.items()}
return ImmutableSparseNDimArray(new_array, a.shape + b.shape)
product_list = [i*j for i in Flatten(a) for j in Flatten(b)]
return ImmutableDenseNDimArray(product_list, a.shape + b.shape)
def _util_contraction_diagonal(array, *contraction_or_diagonal_axes):
array = _arrayfy(array)
# Verify contraction_axes:
taken_dims = set()
for axes_group in contraction_or_diagonal_axes:
if not isinstance(axes_group, Iterable):
raise ValueError("collections of contraction/diagonal axes expected")
dim = array.shape[axes_group[0]]
for d in axes_group:
if d in taken_dims:
raise ValueError("dimension specified more than once")
if dim != array.shape[d]:
raise ValueError("cannot contract or diagonalize between axes of different dimension")
taken_dims.add(d)
rank = array.rank()
remaining_shape = [dim for i, dim in enumerate(array.shape) if i not in taken_dims]
cum_shape = [0]*rank
_cumul = 1
for i in range(rank):
cum_shape[rank - i - 1] = _cumul
_cumul *= int(array.shape[rank - i - 1])
# DEFINITION: by absolute position it is meant the position along the one
# dimensional array containing all the tensor components.
# Possible future work on this module: move computation of absolute
# positions to a class method.
# Determine absolute positions of the uncontracted indices:
remaining_indices = [[cum_shape[i]*j for j in range(array.shape[i])]
for i in range(rank) if i not in taken_dims]
# Determine absolute positions of the contracted indices:
summed_deltas = []
for axes_group in contraction_or_diagonal_axes:
lidx = []
for js in range(array.shape[axes_group[0]]):
lidx.append(sum(cum_shape[ig] * js for ig in axes_group))
summed_deltas.append(lidx)
return array, remaining_indices, remaining_shape, summed_deltas
def tensorcontraction(array, *contraction_axes):
"""
Contraction of an array-like object on the specified axes.
The equivalent operator for array expressions is ``ArrayContraction``,
which can be used to keep the expression unevaluated.
Examples
========
>>> from sympy import Array, tensorcontraction
>>> from sympy import Matrix, eye
>>> tensorcontraction(eye(3), (0, 1))
3
>>> A = Array(range(18), (3, 2, 3))
>>> A
[[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]], [[12, 13, 14], [15, 16, 17]]]
>>> tensorcontraction(A, (0, 2))
[21, 30]
Matrix multiplication may be emulated with a proper combination of
``tensorcontraction`` and ``tensorproduct``
>>> from sympy import tensorproduct
>>> from sympy.abc import a,b,c,d,e,f,g,h
>>> m1 = Matrix([[a, b], [c, d]])
>>> m2 = Matrix([[e, f], [g, h]])
>>> p = tensorproduct(m1, m2)
>>> p
[[[[a*e, a*f], [a*g, a*h]], [[b*e, b*f], [b*g, b*h]]], [[[c*e, c*f], [c*g, c*h]], [[d*e, d*f], [d*g, d*h]]]]
>>> tensorcontraction(p, (1, 2))
[[a*e + b*g, a*f + b*h], [c*e + d*g, c*f + d*h]]
>>> m1*m2
Matrix([
[a*e + b*g, a*f + b*h],
[c*e + d*g, c*f + d*h]])
See Also
========
sympy.tensor.array.expressions.array_expressions.ArrayContraction
"""
from sympy.tensor.array.expressions.array_expressions import _array_contraction
from sympy.tensor.array.expressions.array_expressions import _CodegenArrayAbstract
from sympy.tensor.array.expressions.array_expressions import _ArrayExpr
from sympy.matrices.expressions.matexpr import MatrixSymbol
if isinstance(array, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)):
return _array_contraction(array, *contraction_axes)
array, remaining_indices, remaining_shape, summed_deltas = _util_contraction_diagonal(array, *contraction_axes)
# Compute the contracted array:
#
# 1. external for loops on all uncontracted indices.
# Uncontracted indices are determined by the combinatorial product of
# the absolute positions of the remaining indices.
# 2. internal loop on all contracted indices.
# It sums the values of the absolute contracted index and the absolute
# uncontracted index for the external loop.
contracted_array = []
for icontrib in itertools.product(*remaining_indices):
index_base_position = sum(icontrib)
isum = S.Zero
for sum_to_index in itertools.product(*summed_deltas):
idx = array._get_tuple_index(index_base_position + sum(sum_to_index))
isum += array[idx]
contracted_array.append(isum)
if len(remaining_indices) == 0:
assert len(contracted_array) == 1
return contracted_array[0]
return type(array)(contracted_array, remaining_shape)
def tensordiagonal(array, *diagonal_axes):
"""
Diagonalization of an array-like object on the specified axes.
This is equivalent to multiplying the expression by Kronecker deltas
uniting the axes.
The diagonal indices are put at the end of the axes.
The equivalent operator for array expressions is ``ArrayDiagonal``, which
can be used to keep the expression unevaluated.
Examples
========
``tensordiagonal`` acting on a 2-dimensional array by axes 0 and 1 is
equivalent to the diagonal of the matrix:
>>> from sympy import Array, tensordiagonal
>>> from sympy import Matrix, eye
>>> tensordiagonal(eye(3), (0, 1))
[1, 1, 1]
>>> from sympy.abc import a,b,c,d
>>> m1 = Matrix([[a, b], [c, d]])
>>> tensordiagonal(m1, [0, 1])
[a, d]
In case of higher dimensional arrays, the diagonalized out dimensions
are appended removed and appended as a single dimension at the end:
>>> A = Array(range(18), (3, 2, 3))
>>> A
[[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]], [[12, 13, 14], [15, 16, 17]]]
>>> tensordiagonal(A, (0, 2))
[[0, 7, 14], [3, 10, 17]]
>>> from sympy import permutedims
>>> tensordiagonal(A, (0, 2)) == permutedims(Array([A[0, :, 0], A[1, :, 1], A[2, :, 2]]), [1, 0])
True
See Also
========
sympy.tensor.array.expressions.array_expressions.ArrayDiagonal
"""
if any(len(i) <= 1 for i in diagonal_axes):
raise ValueError("need at least two axes to diagonalize")
from sympy.tensor.array.expressions.array_expressions import _ArrayExpr
from sympy.tensor.array.expressions.array_expressions import _CodegenArrayAbstract
from sympy.tensor.array.expressions.array_expressions import ArrayDiagonal, _array_diagonal
from sympy.matrices.expressions.matexpr import MatrixSymbol
if isinstance(array, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)):
return _array_diagonal(array, *diagonal_axes)
ArrayDiagonal._validate(array, *diagonal_axes)
array, remaining_indices, remaining_shape, diagonal_deltas = _util_contraction_diagonal(array, *diagonal_axes)
# Compute the diagonalized array:
#
# 1. external for loops on all undiagonalized indices.
# Undiagonalized indices are determined by the combinatorial product of
# the absolute positions of the remaining indices.
# 2. internal loop on all diagonal indices.
# It appends the values of the absolute diagonalized index and the absolute
# undiagonalized index for the external loop.
diagonalized_array = []
diagonal_shape = [len(i) for i in diagonal_deltas]
for icontrib in itertools.product(*remaining_indices):
index_base_position = sum(icontrib)
isum = []
for sum_to_index in itertools.product(*diagonal_deltas):
idx = array._get_tuple_index(index_base_position + sum(sum_to_index))
isum.append(array[idx])
isum = type(array)(isum).reshape(*diagonal_shape)
diagonalized_array.append(isum)
return type(array)(diagonalized_array, remaining_shape + diagonal_shape)
def derive_by_array(expr, dx):
r"""
Derivative by arrays. Supports both arrays and scalars.
The equivalent operator for array expressions is ``array_derive``.
Explanation
===========
Given the array `A_{i_1, \ldots, i_N}` and the array `X_{j_1, \ldots, j_M}`
this function will return a new array `B` defined by
`B_{j_1,\ldots,j_M,i_1,\ldots,i_N} := \frac{\partial A_{i_1,\ldots,i_N}}{\partial X_{j_1,\ldots,j_M}}`
Examples
========
>>> from sympy import derive_by_array
>>> from sympy.abc import x, y, z, t
>>> from sympy import cos
>>> derive_by_array(cos(x*t), x)
-t*sin(t*x)
>>> derive_by_array(cos(x*t), [x, y, z, t])
[-t*sin(t*x), 0, 0, -x*sin(t*x)]
>>> derive_by_array([x, y**2*z], [[x, y], [z, t]])
[[[1, 0], [0, 2*y*z]], [[0, y**2], [0, 0]]]
"""
from sympy.matrices import MatrixBase
from sympy.tensor.array import SparseNDimArray
array_types = (Iterable, MatrixBase, NDimArray)
if isinstance(dx, array_types):
dx = ImmutableDenseNDimArray(dx)
for i in dx:
if not i._diff_wrt:
raise ValueError("cannot derive by this array")
if isinstance(expr, array_types):
if isinstance(expr, NDimArray):
expr = expr.as_immutable()
else:
expr = ImmutableDenseNDimArray(expr)
if isinstance(dx, array_types):
if isinstance(expr, SparseNDimArray):
lp = len(expr)
new_array = {k + i*lp: v
for i, x in enumerate(Flatten(dx))
for k, v in expr.diff(x)._sparse_array.items()}
else:
new_array = [[y.diff(x) for y in Flatten(expr)] for x in Flatten(dx)]
return type(expr)(new_array, dx.shape + expr.shape)
else:
return expr.diff(dx)
else:
expr = _sympify(expr)
if isinstance(dx, array_types):
return ImmutableDenseNDimArray([expr.diff(i) for i in Flatten(dx)], dx.shape)
else:
dx = _sympify(dx)
return diff(expr, dx)
def permutedims(expr, perm=None, index_order_old=None, index_order_new=None):
"""
Permutes the indices of an array.
Parameter specifies the permutation of the indices.
The equivalent operator for array expressions is ``PermuteDims``, which can
be used to keep the expression unevaluated.
Examples
========
>>> from sympy.abc import x, y, z, t
>>> from sympy import sin
>>> from sympy import Array, permutedims
>>> a = Array([[x, y, z], [t, sin(x), 0]])
>>> a
[[x, y, z], [t, sin(x), 0]]
>>> permutedims(a, (1, 0))
[[x, t], [y, sin(x)], [z, 0]]
If the array is of second order, ``transpose`` can be used:
>>> from sympy import transpose
>>> transpose(a)
[[x, t], [y, sin(x)], [z, 0]]
Examples on higher dimensions:
>>> b = Array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
>>> permutedims(b, (2, 1, 0))
[[[1, 5], [3, 7]], [[2, 6], [4, 8]]]
>>> permutedims(b, (1, 2, 0))
[[[1, 5], [2, 6]], [[3, 7], [4, 8]]]
An alternative way to specify the same permutations as in the previous
lines involves passing the *old* and *new* indices, either as a list or as
a string:
>>> permutedims(b, index_order_old="cba", index_order_new="abc")
[[[1, 5], [3, 7]], [[2, 6], [4, 8]]]
>>> permutedims(b, index_order_old="cab", index_order_new="abc")
[[[1, 5], [2, 6]], [[3, 7], [4, 8]]]
``Permutation`` objects are also allowed:
>>> from sympy.combinatorics import Permutation
>>> permutedims(b, Permutation([1, 2, 0]))
[[[1, 5], [2, 6]], [[3, 7], [4, 8]]]
See Also
========
sympy.tensor.array.expressions.array_expressions.PermuteDims
"""
from sympy.tensor.array import SparseNDimArray
from sympy.tensor.array.expressions.array_expressions import _ArrayExpr
from sympy.tensor.array.expressions.array_expressions import _CodegenArrayAbstract
from sympy.tensor.array.expressions.array_expressions import _permute_dims
from sympy.matrices.expressions.matexpr import MatrixSymbol
from sympy.tensor.array.expressions import PermuteDims
from sympy.tensor.array.expressions.array_expressions import get_rank
perm = PermuteDims._get_permutation_from_arguments(perm, index_order_old, index_order_new, get_rank(expr))
if isinstance(expr, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)):
return _permute_dims(expr, perm)
if not isinstance(expr, NDimArray):
expr = ImmutableDenseNDimArray(expr)
from sympy.combinatorics import Permutation
if not isinstance(perm, Permutation):
perm = Permutation(list(perm))
if perm.size != expr.rank():
raise ValueError("wrong permutation size")
# Get the inverse permutation:
iperm = ~perm
new_shape = perm(expr.shape)
if isinstance(expr, SparseNDimArray):
return type(expr)({tuple(perm(expr._get_tuple_index(k))): v
for k, v in expr._sparse_array.items()}, new_shape)
indices_span = perm([range(i) for i in expr.shape])
new_array = [None]*len(expr)
for i, idx in enumerate(itertools.product(*indices_span)):
t = iperm(idx)
new_array[i] = expr[t]
return type(expr)(new_array, new_shape)
class Flatten(Printable):
"""
Flatten an iterable object to a list in a lazy-evaluation way.
Notes
=====
This class is an iterator with which the memory cost can be economised.
Optimisation has been considered to ameliorate the performance for some
specific data types like DenseNDimArray and SparseNDimArray.
Examples
========
>>> from sympy.tensor.array.arrayop import Flatten
>>> from sympy.tensor.array import Array
>>> A = Array(range(6)).reshape(2, 3)
>>> Flatten(A)
Flatten([[0, 1, 2], [3, 4, 5]])
>>> [i for i in Flatten(A)]
[0, 1, 2, 3, 4, 5]
"""
def __init__(self, iterable):
from sympy.matrices.matrixbase import MatrixBase
from sympy.tensor.array import NDimArray
if not isinstance(iterable, (Iterable, MatrixBase)):
raise NotImplementedError("Data type not yet supported")
if isinstance(iterable, list):
iterable = NDimArray(iterable)
self._iter = iterable
self._idx = 0
def __iter__(self):
return self
def __next__(self):
from sympy.matrices.matrixbase import MatrixBase
if len(self._iter) > self._idx:
if isinstance(self._iter, DenseNDimArray):
result = self._iter._array[self._idx]
elif isinstance(self._iter, SparseNDimArray):
if self._idx in self._iter._sparse_array:
result = self._iter._sparse_array[self._idx]
else:
result = 0
elif isinstance(self._iter, MatrixBase):
result = self._iter[self._idx]
elif hasattr(self._iter, '__next__'):
result = next(self._iter)
else:
result = self._iter[self._idx]
else:
raise StopIteration
self._idx += 1
return result
def next(self):
return self.__next__()
def _sympystr(self, printer):
return type(self).__name__ + '(' + printer._print(self._iter) + ')'

View File

@ -0,0 +1,206 @@
import functools
from typing import List
from sympy.core.basic import Basic
from sympy.core.containers import Tuple
from sympy.core.singleton import S
from sympy.core.sympify import _sympify
from sympy.tensor.array.mutable_ndim_array import MutableNDimArray
from sympy.tensor.array.ndim_array import NDimArray, ImmutableNDimArray, ArrayKind
from sympy.utilities.iterables import flatten
class DenseNDimArray(NDimArray):
_array: List[Basic]
def __new__(self, *args, **kwargs):
return ImmutableDenseNDimArray(*args, **kwargs)
@property
def kind(self) -> ArrayKind:
return ArrayKind._union(self._array)
def __getitem__(self, index):
"""
Allows to get items from N-dim array.
Examples
========
>>> from sympy import MutableDenseNDimArray
>>> a = MutableDenseNDimArray([0, 1, 2, 3], (2, 2))
>>> a
[[0, 1], [2, 3]]
>>> a[0, 0]
0
>>> a[1, 1]
3
>>> a[0]
[0, 1]
>>> a[1]
[2, 3]
Symbolic index:
>>> from sympy.abc import i, j
>>> a[i, j]
[[0, 1], [2, 3]][i, j]
Replace `i` and `j` to get element `(1, 1)`:
>>> a[i, j].subs({i: 1, j: 1})
3
"""
syindex = self._check_symbolic_index(index)
if syindex is not None:
return syindex
index = self._check_index_for_getitem(index)
if isinstance(index, tuple) and any(isinstance(i, slice) for i in index):
sl_factors, eindices = self._get_slice_data_for_array_access(index)
array = [self._array[self._parse_index(i)] for i in eindices]
nshape = [len(el) for i, el in enumerate(sl_factors) if isinstance(index[i], slice)]
return type(self)(array, nshape)
else:
index = self._parse_index(index)
return self._array[index]
@classmethod
def zeros(cls, *shape):
list_length = functools.reduce(lambda x, y: x*y, shape, S.One)
return cls._new(([0]*list_length,), shape)
def tomatrix(self):
"""
Converts MutableDenseNDimArray to Matrix. Can convert only 2-dim array, else will raise error.
Examples
========
>>> from sympy import MutableDenseNDimArray
>>> a = MutableDenseNDimArray([1 for i in range(9)], (3, 3))
>>> b = a.tomatrix()
>>> b
Matrix([
[1, 1, 1],
[1, 1, 1],
[1, 1, 1]])
"""
from sympy.matrices import Matrix
if self.rank() != 2:
raise ValueError('Dimensions must be of size of 2')
return Matrix(self.shape[0], self.shape[1], self._array)
def reshape(self, *newshape):
"""
Returns MutableDenseNDimArray instance with new shape. Elements number
must be suitable to new shape. The only argument of method sets
new shape.
Examples
========
>>> from sympy import MutableDenseNDimArray
>>> a = MutableDenseNDimArray([1, 2, 3, 4, 5, 6], (2, 3))
>>> a.shape
(2, 3)
>>> a
[[1, 2, 3], [4, 5, 6]]
>>> b = a.reshape(3, 2)
>>> b.shape
(3, 2)
>>> b
[[1, 2], [3, 4], [5, 6]]
"""
new_total_size = functools.reduce(lambda x,y: x*y, newshape)
if new_total_size != self._loop_size:
raise ValueError('Expecting reshape size to %d but got prod(%s) = %d' % (
self._loop_size, str(newshape), new_total_size))
# there is no `.func` as this class does not subtype `Basic`:
return type(self)(self._array, newshape)
class ImmutableDenseNDimArray(DenseNDimArray, ImmutableNDimArray): # type: ignore
def __new__(cls, iterable, shape=None, **kwargs):
return cls._new(iterable, shape, **kwargs)
@classmethod
def _new(cls, iterable, shape, **kwargs):
shape, flat_list = cls._handle_ndarray_creation_inputs(iterable, shape, **kwargs)
shape = Tuple(*map(_sympify, shape))
cls._check_special_bounds(flat_list, shape)
flat_list = flatten(flat_list)
flat_list = Tuple(*flat_list)
self = Basic.__new__(cls, flat_list, shape, **kwargs)
self._shape = shape
self._array = list(flat_list)
self._rank = len(shape)
self._loop_size = functools.reduce(lambda x,y: x*y, shape, 1)
return self
def __setitem__(self, index, value):
raise TypeError('immutable N-dim array')
def as_mutable(self):
return MutableDenseNDimArray(self)
def _eval_simplify(self, **kwargs):
from sympy.simplify.simplify import simplify
return self.applyfunc(simplify)
class MutableDenseNDimArray(DenseNDimArray, MutableNDimArray):
def __new__(cls, iterable=None, shape=None, **kwargs):
return cls._new(iterable, shape, **kwargs)
@classmethod
def _new(cls, iterable, shape, **kwargs):
shape, flat_list = cls._handle_ndarray_creation_inputs(iterable, shape, **kwargs)
flat_list = flatten(flat_list)
self = object.__new__(cls)
self._shape = shape
self._array = list(flat_list)
self._rank = len(shape)
self._loop_size = functools.reduce(lambda x,y: x*y, shape) if shape else len(flat_list)
return self
def __setitem__(self, index, value):
"""Allows to set items to MutableDenseNDimArray.
Examples
========
>>> from sympy import MutableDenseNDimArray
>>> a = MutableDenseNDimArray.zeros(2, 2)
>>> a[0,0] = 1
>>> a[1,1] = 1
>>> a
[[1, 0], [0, 1]]
"""
if isinstance(index, tuple) and any(isinstance(i, slice) for i in index):
value, eindices, slice_offsets = self._get_slice_data_for_array_assignment(index, value)
for i in eindices:
other_i = [ind - j for ind, j in zip(i, slice_offsets) if j is not None]
self._array[self._parse_index(i)] = value[other_i]
else:
index = self._parse_index(index)
self._setter_iterable_check(value)
value = _sympify(value)
self._array[index] = value
def as_immutable(self):
return ImmutableDenseNDimArray(self)
@property
def free_symbols(self):
return {i for j in self._array for i in j.free_symbols}

View File

@ -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

File diff suppressed because it is too large Load Diff

View File

@ -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()

View File

@ -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)

View File

@ -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)

View File

@ -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)

View File

@ -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)

View File

@ -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

View File

@ -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, ()

View File

@ -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

View File

@ -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)

View File

@ -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))

View File

@ -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()

View File

@ -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]

View File

@ -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])

View File

@ -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]))

View File

@ -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))

View File

@ -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)

View File

@ -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

View File

@ -0,0 +1,13 @@
from sympy.tensor.array.ndim_array import NDimArray
class MutableNDimArray(NDimArray):
def as_immutable(self):
raise NotImplementedError("abstract method")
def as_mutable(self):
return self
def _sympy_(self):
return self.as_immutable()

View File

@ -0,0 +1,600 @@
from sympy.core.basic import Basic
from sympy.core.containers import (Dict, Tuple)
from sympy.core.expr import Expr
from sympy.core.kind import Kind, NumberKind, UndefinedKind
from sympy.core.numbers import Integer
from sympy.core.singleton import S
from sympy.core.sympify import sympify
from sympy.external.gmpy import SYMPY_INTS
from sympy.printing.defaults import Printable
import itertools
from collections.abc import Iterable
class ArrayKind(Kind):
"""
Kind for N-dimensional array in SymPy.
This kind represents the multidimensional array that algebraic
operations are defined. Basic class for this kind is ``NDimArray``,
but any expression representing the array can have this.
Parameters
==========
element_kind : Kind
Kind of the element. Default is :obj:NumberKind `<sympy.core.kind.NumberKind>`,
which means that the array contains only numbers.
Examples
========
Any instance of array class has ``ArrayKind``.
>>> from sympy import NDimArray
>>> NDimArray([1,2,3]).kind
ArrayKind(NumberKind)
Although expressions representing an array may be not instance of
array class, it will have ``ArrayKind`` as well.
>>> from sympy import Integral
>>> from sympy.tensor.array import NDimArray
>>> from sympy.abc import x
>>> intA = Integral(NDimArray([1,2,3]), x)
>>> isinstance(intA, NDimArray)
False
>>> intA.kind
ArrayKind(NumberKind)
Use ``isinstance()`` to check for ``ArrayKind` without specifying
the element kind. Use ``is`` with specifying the element kind.
>>> from sympy.tensor.array import ArrayKind
>>> from sympy.core import NumberKind
>>> boolA = NDimArray([True, False])
>>> isinstance(boolA.kind, ArrayKind)
True
>>> boolA.kind is ArrayKind(NumberKind)
False
See Also
========
shape : Function to return the shape of objects with ``MatrixKind``.
"""
def __new__(cls, element_kind=NumberKind):
obj = super().__new__(cls, element_kind)
obj.element_kind = element_kind
return obj
def __repr__(self):
return "ArrayKind(%s)" % self.element_kind
@classmethod
def _union(cls, kinds) -> 'ArrayKind':
elem_kinds = {e.kind for e in kinds}
if len(elem_kinds) == 1:
elemkind, = elem_kinds
else:
elemkind = UndefinedKind
return ArrayKind(elemkind)
class NDimArray(Printable):
"""N-dimensional array.
Examples
========
Create an N-dim array of zeros:
>>> from sympy import MutableDenseNDimArray
>>> a = MutableDenseNDimArray.zeros(2, 3, 4)
>>> a
[[[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]], [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]]
Create an N-dim array from a list;
>>> a = MutableDenseNDimArray([[2, 3], [4, 5]])
>>> a
[[2, 3], [4, 5]]
>>> b = MutableDenseNDimArray([[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]])
>>> b
[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]]
Create an N-dim array from a flat list with dimension shape:
>>> a = MutableDenseNDimArray([1, 2, 3, 4, 5, 6], (2, 3))
>>> a
[[1, 2, 3], [4, 5, 6]]
Create an N-dim array from a matrix:
>>> from sympy import Matrix
>>> a = Matrix([[1,2],[3,4]])
>>> a
Matrix([
[1, 2],
[3, 4]])
>>> b = MutableDenseNDimArray(a)
>>> b
[[1, 2], [3, 4]]
Arithmetic operations on N-dim arrays
>>> a = MutableDenseNDimArray([1, 1, 1, 1], (2, 2))
>>> b = MutableDenseNDimArray([4, 4, 4, 4], (2, 2))
>>> c = a + b
>>> c
[[5, 5], [5, 5]]
>>> a - b
[[-3, -3], [-3, -3]]
"""
_diff_wrt = True
is_scalar = False
def __new__(cls, iterable, shape=None, **kwargs):
from sympy.tensor.array import ImmutableDenseNDimArray
return ImmutableDenseNDimArray(iterable, shape, **kwargs)
def __getitem__(self, index):
raise NotImplementedError("A subclass of NDimArray should implement __getitem__")
def _parse_index(self, index):
if isinstance(index, (SYMPY_INTS, Integer)):
if index >= self._loop_size:
raise ValueError("Only a tuple index is accepted")
return index
if self._loop_size == 0:
raise ValueError("Index not valid with an empty array")
if len(index) != self._rank:
raise ValueError('Wrong number of array axes')
real_index = 0
# check if input index can exist in current indexing
for i in range(self._rank):
if (index[i] >= self.shape[i]) or (index[i] < -self.shape[i]):
raise ValueError('Index ' + str(index) + ' out of border')
if index[i] < 0:
real_index += 1
real_index = real_index*self.shape[i] + index[i]
return real_index
def _get_tuple_index(self, integer_index):
index = []
for sh in reversed(self.shape):
index.append(integer_index % sh)
integer_index //= sh
index.reverse()
return tuple(index)
def _check_symbolic_index(self, index):
# Check if any index is symbolic:
tuple_index = (index if isinstance(index, tuple) else (index,))
if any((isinstance(i, Expr) and (not i.is_number)) for i in tuple_index):
for i, nth_dim in zip(tuple_index, self.shape):
if ((i < 0) == True) or ((i >= nth_dim) == True):
raise ValueError("index out of range")
from sympy.tensor import Indexed
return Indexed(self, *tuple_index)
return None
def _setter_iterable_check(self, value):
from sympy.matrices.matrixbase import MatrixBase
if isinstance(value, (Iterable, MatrixBase, NDimArray)):
raise NotImplementedError
@classmethod
def _scan_iterable_shape(cls, iterable):
def f(pointer):
if not isinstance(pointer, Iterable):
return [pointer], ()
if len(pointer) == 0:
return [], (0,)
result = []
elems, shapes = zip(*[f(i) for i in pointer])
if len(set(shapes)) != 1:
raise ValueError("could not determine shape unambiguously")
for i in elems:
result.extend(i)
return result, (len(shapes),)+shapes[0]
return f(iterable)
@classmethod
def _handle_ndarray_creation_inputs(cls, iterable=None, shape=None, **kwargs):
from sympy.matrices.matrixbase import MatrixBase
from sympy.tensor.array import SparseNDimArray
if shape is None:
if iterable is None:
shape = ()
iterable = ()
# Construction of a sparse array from a sparse array
elif isinstance(iterable, SparseNDimArray):
return iterable._shape, iterable._sparse_array
# Construct N-dim array from another N-dim array:
elif isinstance(iterable, NDimArray):
shape = iterable.shape
# Construct N-dim array from an iterable (numpy arrays included):
elif isinstance(iterable, Iterable):
iterable, shape = cls._scan_iterable_shape(iterable)
# Construct N-dim array from a Matrix:
elif isinstance(iterable, MatrixBase):
shape = iterable.shape
else:
shape = ()
iterable = (iterable,)
if isinstance(iterable, (Dict, dict)) and shape is not None:
new_dict = iterable.copy()
for k in new_dict:
if isinstance(k, (tuple, Tuple)):
new_key = 0
for i, idx in enumerate(k):
new_key = new_key * shape[i] + idx
iterable[new_key] = iterable[k]
del iterable[k]
if isinstance(shape, (SYMPY_INTS, Integer)):
shape = (shape,)
if not all(isinstance(dim, (SYMPY_INTS, Integer)) for dim in shape):
raise TypeError("Shape should contain integers only.")
return tuple(shape), iterable
def __len__(self):
"""Overload common function len(). Returns number of elements in array.
Examples
========
>>> from sympy import MutableDenseNDimArray
>>> a = MutableDenseNDimArray.zeros(3, 3)
>>> a
[[0, 0, 0], [0, 0, 0], [0, 0, 0]]
>>> len(a)
9
"""
return self._loop_size
@property
def shape(self):
"""
Returns array shape (dimension).
Examples
========
>>> from sympy import MutableDenseNDimArray
>>> a = MutableDenseNDimArray.zeros(3, 3)
>>> a.shape
(3, 3)
"""
return self._shape
def rank(self):
"""
Returns rank of array.
Examples
========
>>> from sympy import MutableDenseNDimArray
>>> a = MutableDenseNDimArray.zeros(3,4,5,6,3)
>>> a.rank()
5
"""
return self._rank
def diff(self, *args, **kwargs):
"""
Calculate the derivative of each element in the array.
Examples
========
>>> from sympy import ImmutableDenseNDimArray
>>> from sympy.abc import x, y
>>> M = ImmutableDenseNDimArray([[x, y], [1, x*y]])
>>> M.diff(x)
[[1, 0], [0, y]]
"""
from sympy.tensor.array.array_derivatives import ArrayDerivative
kwargs.setdefault('evaluate', True)
return ArrayDerivative(self.as_immutable(), *args, **kwargs)
def _eval_derivative(self, base):
# Types are (base: scalar, self: array)
return self.applyfunc(lambda x: base.diff(x))
def _eval_derivative_n_times(self, s, n):
return Basic._eval_derivative_n_times(self, s, n)
def applyfunc(self, f):
"""Apply a function to each element of the N-dim array.
Examples
========
>>> from sympy import ImmutableDenseNDimArray
>>> m = ImmutableDenseNDimArray([i*2+j for i in range(2) for j in range(2)], (2, 2))
>>> m
[[0, 1], [2, 3]]
>>> m.applyfunc(lambda i: 2*i)
[[0, 2], [4, 6]]
"""
from sympy.tensor.array import SparseNDimArray
from sympy.tensor.array.arrayop import Flatten
if isinstance(self, SparseNDimArray) and f(S.Zero) == 0:
return type(self)({k: f(v) for k, v in self._sparse_array.items() if f(v) != 0}, self.shape)
return type(self)(map(f, Flatten(self)), self.shape)
def _sympystr(self, printer):
def f(sh, shape_left, i, j):
if len(shape_left) == 1:
return "["+", ".join([printer._print(self[self._get_tuple_index(e)]) for e in range(i, j)])+"]"
sh //= shape_left[0]
return "[" + ", ".join([f(sh, shape_left[1:], i+e*sh, i+(e+1)*sh) for e in range(shape_left[0])]) + "]" # + "\n"*len(shape_left)
if self.rank() == 0:
return printer._print(self[()])
return f(self._loop_size, self.shape, 0, self._loop_size)
def tolist(self):
"""
Converting MutableDenseNDimArray to one-dim list
Examples
========
>>> from sympy import MutableDenseNDimArray
>>> a = MutableDenseNDimArray([1, 2, 3, 4], (2, 2))
>>> a
[[1, 2], [3, 4]]
>>> b = a.tolist()
>>> b
[[1, 2], [3, 4]]
"""
def f(sh, shape_left, i, j):
if len(shape_left) == 1:
return [self[self._get_tuple_index(e)] for e in range(i, j)]
result = []
sh //= shape_left[0]
for e in range(shape_left[0]):
result.append(f(sh, shape_left[1:], i+e*sh, i+(e+1)*sh))
return result
return f(self._loop_size, self.shape, 0, self._loop_size)
def __add__(self, other):
from sympy.tensor.array.arrayop import Flatten
if not isinstance(other, NDimArray):
return NotImplemented
if self.shape != other.shape:
raise ValueError("array shape mismatch")
result_list = [i+j for i,j in zip(Flatten(self), Flatten(other))]
return type(self)(result_list, self.shape)
def __sub__(self, other):
from sympy.tensor.array.arrayop import Flatten
if not isinstance(other, NDimArray):
return NotImplemented
if self.shape != other.shape:
raise ValueError("array shape mismatch")
result_list = [i-j for i,j in zip(Flatten(self), Flatten(other))]
return type(self)(result_list, self.shape)
def __mul__(self, other):
from sympy.matrices.matrixbase import MatrixBase
from sympy.tensor.array import SparseNDimArray
from sympy.tensor.array.arrayop import Flatten
if isinstance(other, (Iterable, NDimArray, MatrixBase)):
raise ValueError("scalar expected, use tensorproduct(...) for tensorial product")
other = sympify(other)
if isinstance(self, SparseNDimArray):
if other.is_zero:
return type(self)({}, self.shape)
return type(self)({k: other*v for (k, v) in self._sparse_array.items()}, self.shape)
result_list = [i*other for i in Flatten(self)]
return type(self)(result_list, self.shape)
def __rmul__(self, other):
from sympy.matrices.matrixbase import MatrixBase
from sympy.tensor.array import SparseNDimArray
from sympy.tensor.array.arrayop import Flatten
if isinstance(other, (Iterable, NDimArray, MatrixBase)):
raise ValueError("scalar expected, use tensorproduct(...) for tensorial product")
other = sympify(other)
if isinstance(self, SparseNDimArray):
if other.is_zero:
return type(self)({}, self.shape)
return type(self)({k: other*v for (k, v) in self._sparse_array.items()}, self.shape)
result_list = [other*i for i in Flatten(self)]
return type(self)(result_list, self.shape)
def __truediv__(self, other):
from sympy.matrices.matrixbase import MatrixBase
from sympy.tensor.array import SparseNDimArray
from sympy.tensor.array.arrayop import Flatten
if isinstance(other, (Iterable, NDimArray, MatrixBase)):
raise ValueError("scalar expected")
other = sympify(other)
if isinstance(self, SparseNDimArray) and other != S.Zero:
return type(self)({k: v/other for (k, v) in self._sparse_array.items()}, self.shape)
result_list = [i/other for i in Flatten(self)]
return type(self)(result_list, self.shape)
def __rtruediv__(self, other):
raise NotImplementedError('unsupported operation on NDimArray')
def __neg__(self):
from sympy.tensor.array import SparseNDimArray
from sympy.tensor.array.arrayop import Flatten
if isinstance(self, SparseNDimArray):
return type(self)({k: -v for (k, v) in self._sparse_array.items()}, self.shape)
result_list = [-i for i in Flatten(self)]
return type(self)(result_list, self.shape)
def __iter__(self):
def iterator():
if self._shape:
for i in range(self._shape[0]):
yield self[i]
else:
yield self[()]
return iterator()
def __eq__(self, other):
"""
NDimArray instances can be compared to each other.
Instances equal if they have same shape and data.
Examples
========
>>> from sympy import MutableDenseNDimArray
>>> a = MutableDenseNDimArray.zeros(2, 3)
>>> b = MutableDenseNDimArray.zeros(2, 3)
>>> a == b
True
>>> c = a.reshape(3, 2)
>>> c == b
False
>>> a[0,0] = 1
>>> b[0,0] = 2
>>> a == b
False
"""
from sympy.tensor.array import SparseNDimArray
if not isinstance(other, NDimArray):
return False
if not self.shape == other.shape:
return False
if isinstance(self, SparseNDimArray) and isinstance(other, SparseNDimArray):
return dict(self._sparse_array) == dict(other._sparse_array)
return list(self) == list(other)
def __ne__(self, other):
return not self == other
def _eval_transpose(self):
if self.rank() != 2:
raise ValueError("array rank not 2")
from .arrayop import permutedims
return permutedims(self, (1, 0))
def transpose(self):
return self._eval_transpose()
def _eval_conjugate(self):
from sympy.tensor.array.arrayop import Flatten
return self.func([i.conjugate() for i in Flatten(self)], self.shape)
def conjugate(self):
return self._eval_conjugate()
def _eval_adjoint(self):
return self.transpose().conjugate()
def adjoint(self):
return self._eval_adjoint()
def _slice_expand(self, s, dim):
if not isinstance(s, slice):
return (s,)
start, stop, step = s.indices(dim)
return [start + i*step for i in range((stop-start)//step)]
def _get_slice_data_for_array_access(self, index):
sl_factors = [self._slice_expand(i, dim) for (i, dim) in zip(index, self.shape)]
eindices = itertools.product(*sl_factors)
return sl_factors, eindices
def _get_slice_data_for_array_assignment(self, index, value):
if not isinstance(value, NDimArray):
value = type(self)(value)
sl_factors, eindices = self._get_slice_data_for_array_access(index)
slice_offsets = [min(i) if isinstance(i, list) else None for i in sl_factors]
# TODO: add checks for dimensions for `value`?
return value, eindices, slice_offsets
@classmethod
def _check_special_bounds(cls, flat_list, shape):
if shape == () and len(flat_list) != 1:
raise ValueError("arrays without shape need one scalar value")
if shape == (0,) and len(flat_list) > 0:
raise ValueError("if array shape is (0,) there cannot be elements")
def _check_index_for_getitem(self, index):
if isinstance(index, (SYMPY_INTS, Integer, slice)):
index = (index,)
if len(index) < self.rank():
index = tuple(index) + \
tuple(slice(None) for i in range(len(index), self.rank()))
if len(index) > self.rank():
raise ValueError('Dimension of index greater than rank of array')
return index
class ImmutableNDimArray(NDimArray, Basic):
_op_priority = 11.0
def __hash__(self):
return Basic.__hash__(self)
def as_immutable(self):
return self
def as_mutable(self):
raise NotImplementedError("abstract method")

View File

@ -0,0 +1,196 @@
from sympy.core.basic import Basic
from sympy.core.containers import (Dict, Tuple)
from sympy.core.singleton import S
from sympy.core.sympify import _sympify
from sympy.tensor.array.mutable_ndim_array import MutableNDimArray
from sympy.tensor.array.ndim_array import NDimArray, ImmutableNDimArray
from sympy.utilities.iterables import flatten
import functools
class SparseNDimArray(NDimArray):
def __new__(self, *args, **kwargs):
return ImmutableSparseNDimArray(*args, **kwargs)
def __getitem__(self, index):
"""
Get an element from a sparse N-dim array.
Examples
========
>>> from sympy import MutableSparseNDimArray
>>> a = MutableSparseNDimArray(range(4), (2, 2))
>>> a
[[0, 1], [2, 3]]
>>> a[0, 0]
0
>>> a[1, 1]
3
>>> a[0]
[0, 1]
>>> a[1]
[2, 3]
Symbolic indexing:
>>> from sympy.abc import i, j
>>> a[i, j]
[[0, 1], [2, 3]][i, j]
Replace `i` and `j` to get element `(0, 0)`:
>>> a[i, j].subs({i: 0, j: 0})
0
"""
syindex = self._check_symbolic_index(index)
if syindex is not None:
return syindex
index = self._check_index_for_getitem(index)
# `index` is a tuple with one or more slices:
if isinstance(index, tuple) and any(isinstance(i, slice) for i in index):
sl_factors, eindices = self._get_slice_data_for_array_access(index)
array = [self._sparse_array.get(self._parse_index(i), S.Zero) for i in eindices]
nshape = [len(el) for i, el in enumerate(sl_factors) if isinstance(index[i], slice)]
return type(self)(array, nshape)
else:
index = self._parse_index(index)
return self._sparse_array.get(index, S.Zero)
@classmethod
def zeros(cls, *shape):
"""
Return a sparse N-dim array of zeros.
"""
return cls({}, shape)
def tomatrix(self):
"""
Converts MutableDenseNDimArray to Matrix. Can convert only 2-dim array, else will raise error.
Examples
========
>>> from sympy import MutableSparseNDimArray
>>> a = MutableSparseNDimArray([1 for i in range(9)], (3, 3))
>>> b = a.tomatrix()
>>> b
Matrix([
[1, 1, 1],
[1, 1, 1],
[1, 1, 1]])
"""
from sympy.matrices import SparseMatrix
if self.rank() != 2:
raise ValueError('Dimensions must be of size of 2')
mat_sparse = {}
for key, value in self._sparse_array.items():
mat_sparse[self._get_tuple_index(key)] = value
return SparseMatrix(self.shape[0], self.shape[1], mat_sparse)
def reshape(self, *newshape):
new_total_size = functools.reduce(lambda x,y: x*y, newshape)
if new_total_size != self._loop_size:
raise ValueError("Invalid reshape parameters " + newshape)
return type(self)(self._sparse_array, newshape)
class ImmutableSparseNDimArray(SparseNDimArray, ImmutableNDimArray): # type: ignore
def __new__(cls, iterable=None, shape=None, **kwargs):
shape, flat_list = cls._handle_ndarray_creation_inputs(iterable, shape, **kwargs)
shape = Tuple(*map(_sympify, shape))
cls._check_special_bounds(flat_list, shape)
loop_size = functools.reduce(lambda x,y: x*y, shape) if shape else len(flat_list)
# Sparse array:
if isinstance(flat_list, (dict, Dict)):
sparse_array = Dict(flat_list)
else:
sparse_array = {}
for i, el in enumerate(flatten(flat_list)):
if el != 0:
sparse_array[i] = _sympify(el)
sparse_array = Dict(sparse_array)
self = Basic.__new__(cls, sparse_array, shape, **kwargs)
self._shape = shape
self._rank = len(shape)
self._loop_size = loop_size
self._sparse_array = sparse_array
return self
def __setitem__(self, index, value):
raise TypeError("immutable N-dim array")
def as_mutable(self):
return MutableSparseNDimArray(self)
class MutableSparseNDimArray(MutableNDimArray, SparseNDimArray):
def __new__(cls, iterable=None, shape=None, **kwargs):
shape, flat_list = cls._handle_ndarray_creation_inputs(iterable, shape, **kwargs)
self = object.__new__(cls)
self._shape = shape
self._rank = len(shape)
self._loop_size = functools.reduce(lambda x,y: x*y, shape) if shape else len(flat_list)
# Sparse array:
if isinstance(flat_list, (dict, Dict)):
self._sparse_array = dict(flat_list)
return self
self._sparse_array = {}
for i, el in enumerate(flatten(flat_list)):
if el != 0:
self._sparse_array[i] = _sympify(el)
return self
def __setitem__(self, index, value):
"""Allows to set items to MutableDenseNDimArray.
Examples
========
>>> from sympy import MutableSparseNDimArray
>>> a = MutableSparseNDimArray.zeros(2, 2)
>>> a[0, 0] = 1
>>> a[1, 1] = 1
>>> a
[[1, 0], [0, 1]]
"""
if isinstance(index, tuple) and any(isinstance(i, slice) for i in index):
value, eindices, slice_offsets = self._get_slice_data_for_array_assignment(index, value)
for i in eindices:
other_i = [ind - j for ind, j in zip(i, slice_offsets) if j is not None]
other_value = value[other_i]
complete_index = self._parse_index(i)
if other_value != 0:
self._sparse_array[complete_index] = other_value
elif complete_index in self._sparse_array:
self._sparse_array.pop(complete_index)
else:
index = self._parse_index(index)
value = _sympify(value)
if value == 0 and index in self._sparse_array:
self._sparse_array.pop(index)
else:
self._sparse_array[index] = value
def as_immutable(self):
return ImmutableSparseNDimArray(self)
@property
def free_symbols(self):
return {i for j in self._sparse_array.values() for i in j.free_symbols}

View File

@ -0,0 +1,78 @@
from sympy.tensor.array.array_comprehension import ArrayComprehension, ArrayComprehensionMap
from sympy.tensor.array import ImmutableDenseNDimArray
from sympy.abc import i, j, k, l
from sympy.testing.pytest import raises
from sympy.matrices import Matrix
def test_array_comprehension():
a = ArrayComprehension(i*j, (i, 1, 3), (j, 2, 4))
b = ArrayComprehension(i, (i, 1, j+1))
c = ArrayComprehension(i+j+k+l, (i, 1, 2), (j, 1, 3), (k, 1, 4), (l, 1, 5))
d = ArrayComprehension(k, (i, 1, 5))
e = ArrayComprehension(i, (j, k+1, k+5))
assert a.doit().tolist() == [[2, 3, 4], [4, 6, 8], [6, 9, 12]]
assert a.shape == (3, 3)
assert a.is_shape_numeric == True
assert a.tolist() == [[2, 3, 4], [4, 6, 8], [6, 9, 12]]
assert a.tomatrix() == Matrix([
[2, 3, 4],
[4, 6, 8],
[6, 9, 12]])
assert len(a) == 9
assert isinstance(b.doit(), ArrayComprehension)
assert isinstance(a.doit(), ImmutableDenseNDimArray)
assert b.subs(j, 3) == ArrayComprehension(i, (i, 1, 4))
assert b.free_symbols == {j}
assert b.shape == (j + 1,)
assert b.rank() == 1
assert b.is_shape_numeric == False
assert c.free_symbols == set()
assert c.function == i + j + k + l
assert c.limits == ((i, 1, 2), (j, 1, 3), (k, 1, 4), (l, 1, 5))
assert c.doit().tolist() == [[[[4, 5, 6, 7, 8], [5, 6, 7, 8, 9], [6, 7, 8, 9, 10], [7, 8, 9, 10, 11]],
[[5, 6, 7, 8, 9], [6, 7, 8, 9, 10], [7, 8, 9, 10, 11], [8, 9, 10, 11, 12]],
[[6, 7, 8, 9, 10], [7, 8, 9, 10, 11], [8, 9, 10, 11, 12], [9, 10, 11, 12, 13]]],
[[[5, 6, 7, 8, 9], [6, 7, 8, 9, 10], [7, 8, 9, 10, 11], [8, 9, 10, 11, 12]],
[[6, 7, 8, 9, 10], [7, 8, 9, 10, 11], [8, 9, 10, 11, 12], [9, 10, 11, 12, 13]],
[[7, 8, 9, 10, 11], [8, 9, 10, 11, 12], [9, 10, 11, 12, 13], [10, 11, 12, 13, 14]]]]
assert c.free_symbols == set()
assert c.variables == [i, j, k, l]
assert c.bound_symbols == [i, j, k, l]
assert d.doit().tolist() == [k, k, k, k, k]
assert len(e) == 5
raises(TypeError, lambda: ArrayComprehension(i*j, (i, 1, 3), (j, 2, [1, 3, 2])))
raises(ValueError, lambda: ArrayComprehension(i*j, (i, 1, 3), (j, 2, 1)))
raises(ValueError, lambda: ArrayComprehension(i*j, (i, 1, 3), (j, 2, j+1)))
raises(ValueError, lambda: len(ArrayComprehension(i*j, (i, 1, 3), (j, 2, j+4))))
raises(TypeError, lambda: ArrayComprehension(i*j, (i, 0, i + 1.5), (j, 0, 2)))
raises(ValueError, lambda: b.tolist())
raises(ValueError, lambda: b.tomatrix())
raises(ValueError, lambda: c.tomatrix())
def test_arraycomprehensionmap():
a = ArrayComprehensionMap(lambda i: i+1, (i, 1, 5))
assert a.doit().tolist() == [2, 3, 4, 5, 6]
assert a.shape == (5,)
assert a.is_shape_numeric
assert a.tolist() == [2, 3, 4, 5, 6]
assert len(a) == 5
assert isinstance(a.doit(), ImmutableDenseNDimArray)
expr = ArrayComprehensionMap(lambda i: i+1, (i, 1, k))
assert expr.doit() == expr
assert expr.subs(k, 4) == ArrayComprehensionMap(lambda i: i+1, (i, 1, 4))
assert expr.subs(k, 4).doit() == ImmutableDenseNDimArray([2, 3, 4, 5])
b = ArrayComprehensionMap(lambda i: i+1, (i, 1, 2), (i, 1, 3), (i, 1, 4), (i, 1, 5))
assert b.doit().tolist() == [[[[2, 3, 4, 5, 6], [3, 5, 7, 9, 11], [4, 7, 10, 13, 16], [5, 9, 13, 17, 21]],
[[3, 5, 7, 9, 11], [5, 9, 13, 17, 21], [7, 13, 19, 25, 31], [9, 17, 25, 33, 41]],
[[4, 7, 10, 13, 16], [7, 13, 19, 25, 31], [10, 19, 28, 37, 46], [13, 25, 37, 49, 61]]],
[[[3, 5, 7, 9, 11], [5, 9, 13, 17, 21], [7, 13, 19, 25, 31], [9, 17, 25, 33, 41]],
[[5, 9, 13, 17, 21], [9, 17, 25, 33, 41], [13, 25, 37, 49, 61], [17, 33, 49, 65, 81]],
[[7, 13, 19, 25, 31], [13, 25, 37, 49, 61], [19, 37, 55, 73, 91], [25, 49, 73, 97, 121]]]]
# tests about lambda expression
assert ArrayComprehensionMap(lambda: 3, (i, 1, 5)).doit().tolist() == [3, 3, 3, 3, 3]
assert ArrayComprehensionMap(lambda i: i+1, (i, 1, 5)).doit().tolist() == [2, 3, 4, 5, 6]
raises(ValueError, lambda: ArrayComprehensionMap(i*j, (i, 1, 3), (j, 2, 4)))
a = ArrayComprehensionMap(lambda i, j: i+j, (i, 1, 5))
raises(ValueError, lambda: a.doit())

View File

@ -0,0 +1,52 @@
from sympy.core.symbol import symbols
from sympy.matrices.dense import Matrix
from sympy.matrices.expressions.matexpr import MatrixSymbol
from sympy.tensor.array.ndim_array import NDimArray
from sympy.matrices.matrixbase import MatrixBase
from sympy.tensor.array.array_derivatives import ArrayDerivative
x, y, z, t = symbols("x y z t")
m = Matrix([[x, y], [z, t]])
M = MatrixSymbol("M", 3, 2)
N = MatrixSymbol("N", 4, 3)
def test_array_derivative_construction():
d = ArrayDerivative(x, m, evaluate=False)
assert d.shape == (2, 2)
expr = d.doit()
assert isinstance(expr, MatrixBase)
assert expr.shape == (2, 2)
d = ArrayDerivative(m, m, evaluate=False)
assert d.shape == (2, 2, 2, 2)
expr = d.doit()
assert isinstance(expr, NDimArray)
assert expr.shape == (2, 2, 2, 2)
d = ArrayDerivative(m, x, evaluate=False)
assert d.shape == (2, 2)
expr = d.doit()
assert isinstance(expr, MatrixBase)
assert expr.shape == (2, 2)
d = ArrayDerivative(M, N, evaluate=False)
assert d.shape == (4, 3, 3, 2)
expr = d.doit()
assert isinstance(expr, ArrayDerivative)
assert expr.shape == (4, 3, 3, 2)
d = ArrayDerivative(M, (N, 2), evaluate=False)
assert d.shape == (4, 3, 4, 3, 3, 2)
expr = d.doit()
assert isinstance(expr, ArrayDerivative)
assert expr.shape == (4, 3, 4, 3, 3, 2)
d = ArrayDerivative(M.as_explicit(), (N.as_explicit(), 2), evaluate=False)
assert d.doit().shape == (4, 3, 4, 3, 3, 2)
expr = d.doit()
assert isinstance(expr, NDimArray)
assert expr.shape == (4, 3, 4, 3, 3, 2)

View File

@ -0,0 +1,361 @@
import itertools
import random
from sympy.combinatorics import Permutation
from sympy.combinatorics.permutations import _af_invert
from sympy.testing.pytest import raises
from sympy.core.function import diff
from sympy.core.symbol import symbols
from sympy.functions.elementary.complexes import (adjoint, conjugate, transpose)
from sympy.functions.elementary.exponential import (exp, log)
from sympy.functions.elementary.trigonometric import (cos, sin)
from sympy.tensor.array import Array, ImmutableDenseNDimArray, ImmutableSparseNDimArray, MutableSparseNDimArray
from sympy.tensor.array.arrayop import tensorproduct, tensorcontraction, derive_by_array, permutedims, Flatten, \
tensordiagonal
def test_import_NDimArray():
from sympy.tensor.array import NDimArray
del NDimArray
def test_tensorproduct():
x,y,z,t = symbols('x y z t')
from sympy.abc import a,b,c,d
assert tensorproduct() == 1
assert tensorproduct([x]) == Array([x])
assert tensorproduct([x], [y]) == Array([[x*y]])
assert tensorproduct([x], [y], [z]) == Array([[[x*y*z]]])
assert tensorproduct([x], [y], [z], [t]) == Array([[[[x*y*z*t]]]])
assert tensorproduct(x) == x
assert tensorproduct(x, y) == x*y
assert tensorproduct(x, y, z) == x*y*z
assert tensorproduct(x, y, z, t) == x*y*z*t
for ArrayType in [ImmutableDenseNDimArray, ImmutableSparseNDimArray]:
A = ArrayType([x, y])
B = ArrayType([1, 2, 3])
C = ArrayType([a, b, c, d])
assert tensorproduct(A, B, C) == ArrayType([[[a*x, b*x, c*x, d*x], [2*a*x, 2*b*x, 2*c*x, 2*d*x], [3*a*x, 3*b*x, 3*c*x, 3*d*x]],
[[a*y, b*y, c*y, d*y], [2*a*y, 2*b*y, 2*c*y, 2*d*y], [3*a*y, 3*b*y, 3*c*y, 3*d*y]]])
assert tensorproduct([x, y], [1, 2, 3]) == tensorproduct(A, B)
assert tensorproduct(A, 2) == ArrayType([2*x, 2*y])
assert tensorproduct(A, [2]) == ArrayType([[2*x], [2*y]])
assert tensorproduct([2], A) == ArrayType([[2*x, 2*y]])
assert tensorproduct(a, A) == ArrayType([a*x, a*y])
assert tensorproduct(a, A, B) == ArrayType([[a*x, 2*a*x, 3*a*x], [a*y, 2*a*y, 3*a*y]])
assert tensorproduct(A, B, a) == ArrayType([[a*x, 2*a*x, 3*a*x], [a*y, 2*a*y, 3*a*y]])
assert tensorproduct(B, a, A) == ArrayType([[a*x, a*y], [2*a*x, 2*a*y], [3*a*x, 3*a*y]])
# tests for large scale sparse array
for SparseArrayType in [ImmutableSparseNDimArray, MutableSparseNDimArray]:
a = SparseArrayType({1:2, 3:4},(1000, 2000))
b = SparseArrayType({1:2, 3:4},(1000, 2000))
assert tensorproduct(a, b) == ImmutableSparseNDimArray({2000001: 4, 2000003: 8, 6000001: 8, 6000003: 16}, (1000, 2000, 1000, 2000))
def test_tensorcontraction():
from sympy.abc import a,b,c,d,e,f,g,h,i,j,k,l,m,n,o,p,q,r,s,t,u,v,w,x
B = Array(range(18), (2, 3, 3))
assert tensorcontraction(B, (1, 2)) == Array([12, 39])
C1 = Array([a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s, t, u, v, w, x], (2, 3, 2, 2))
assert tensorcontraction(C1, (0, 2)) == Array([[a + o, b + p], [e + s, f + t], [i + w, j + x]])
assert tensorcontraction(C1, (0, 2, 3)) == Array([a + p, e + t, i + x])
assert tensorcontraction(C1, (2, 3)) == Array([[a + d, e + h, i + l], [m + p, q + t, u + x]])
def test_derivative_by_array():
from sympy.abc import i, j, t, x, y, z
bexpr = x*y**2*exp(z)*log(t)
sexpr = sin(bexpr)
cexpr = cos(bexpr)
a = Array([sexpr])
assert derive_by_array(sexpr, t) == x*y**2*exp(z)*cos(x*y**2*exp(z)*log(t))/t
assert derive_by_array(sexpr, [x, y, z]) == Array([bexpr/x*cexpr, 2*y*bexpr/y**2*cexpr, bexpr*cexpr])
assert derive_by_array(a, [x, y, z]) == Array([[bexpr/x*cexpr], [2*y*bexpr/y**2*cexpr], [bexpr*cexpr]])
assert derive_by_array(sexpr, [[x, y], [z, t]]) == Array([[bexpr/x*cexpr, 2*y*bexpr/y**2*cexpr], [bexpr*cexpr, bexpr/log(t)/t*cexpr]])
assert derive_by_array(a, [[x, y], [z, t]]) == Array([[[bexpr/x*cexpr], [2*y*bexpr/y**2*cexpr]], [[bexpr*cexpr], [bexpr/log(t)/t*cexpr]]])
assert derive_by_array([[x, y], [z, t]], [x, y]) == Array([[[1, 0], [0, 0]], [[0, 1], [0, 0]]])
assert derive_by_array([[x, y], [z, t]], [[x, y], [z, t]]) == Array([[[[1, 0], [0, 0]], [[0, 1], [0, 0]]],
[[[0, 0], [1, 0]], [[0, 0], [0, 1]]]])
assert diff(sexpr, t) == x*y**2*exp(z)*cos(x*y**2*exp(z)*log(t))/t
assert diff(sexpr, Array([x, y, z])) == Array([bexpr/x*cexpr, 2*y*bexpr/y**2*cexpr, bexpr*cexpr])
assert diff(a, Array([x, y, z])) == Array([[bexpr/x*cexpr], [2*y*bexpr/y**2*cexpr], [bexpr*cexpr]])
assert diff(sexpr, Array([[x, y], [z, t]])) == Array([[bexpr/x*cexpr, 2*y*bexpr/y**2*cexpr], [bexpr*cexpr, bexpr/log(t)/t*cexpr]])
assert diff(a, Array([[x, y], [z, t]])) == Array([[[bexpr/x*cexpr], [2*y*bexpr/y**2*cexpr]], [[bexpr*cexpr], [bexpr/log(t)/t*cexpr]]])
assert diff(Array([[x, y], [z, t]]), Array([x, y])) == Array([[[1, 0], [0, 0]], [[0, 1], [0, 0]]])
assert diff(Array([[x, y], [z, t]]), Array([[x, y], [z, t]])) == Array([[[[1, 0], [0, 0]], [[0, 1], [0, 0]]],
[[[0, 0], [1, 0]], [[0, 0], [0, 1]]]])
# test for large scale sparse array
for SparseArrayType in [ImmutableSparseNDimArray, MutableSparseNDimArray]:
b = MutableSparseNDimArray({0:i, 1:j}, (10000, 20000))
assert derive_by_array(b, i) == ImmutableSparseNDimArray({0: 1}, (10000, 20000))
assert derive_by_array(b, (i, j)) == ImmutableSparseNDimArray({0: 1, 200000001: 1}, (2, 10000, 20000))
#https://github.com/sympy/sympy/issues/20655
U = Array([x, y, z])
E = 2
assert derive_by_array(E, U) == ImmutableDenseNDimArray([0, 0, 0])
def test_issue_emerged_while_discussing_10972():
ua = Array([-1,0])
Fa = Array([[0, 1], [-1, 0]])
po = tensorproduct(Fa, ua, Fa, ua)
assert tensorcontraction(po, (1, 2), (4, 5)) == Array([[0, 0], [0, 1]])
sa = symbols('a0:144')
po = Array(sa, [2, 2, 3, 3, 2, 2])
assert tensorcontraction(po, (0, 1), (2, 3), (4, 5)) == sa[0] + sa[108] + sa[111] + sa[124] + sa[127] + sa[140] + sa[143] + sa[16] + sa[19] + sa[3] + sa[32] + sa[35]
assert tensorcontraction(po, (0, 1, 4, 5), (2, 3)) == sa[0] + sa[111] + sa[127] + sa[143] + sa[16] + sa[32]
assert tensorcontraction(po, (0, 1), (4, 5)) == Array([[sa[0] + sa[108] + sa[111] + sa[3], sa[112] + sa[115] + sa[4] + sa[7],
sa[11] + sa[116] + sa[119] + sa[8]], [sa[12] + sa[120] + sa[123] + sa[15],
sa[124] + sa[127] + sa[16] + sa[19], sa[128] + sa[131] + sa[20] + sa[23]],
[sa[132] + sa[135] + sa[24] + sa[27], sa[136] + sa[139] + sa[28] + sa[31],
sa[140] + sa[143] + sa[32] + sa[35]]])
assert tensorcontraction(po, (0, 1), (2, 3)) == Array([[sa[0] + sa[108] + sa[124] + sa[140] + sa[16] + sa[32], sa[1] + sa[109] + sa[125] + sa[141] + sa[17] + sa[33]],
[sa[110] + sa[126] + sa[142] + sa[18] + sa[2] + sa[34], sa[111] + sa[127] + sa[143] + sa[19] + sa[3] + sa[35]]])
def test_array_permutedims():
sa = symbols('a0:144')
for ArrayType in [ImmutableDenseNDimArray, ImmutableSparseNDimArray]:
m1 = ArrayType(sa[:6], (2, 3))
assert permutedims(m1, (1, 0)) == transpose(m1)
assert m1.tomatrix().T == permutedims(m1, (1, 0)).tomatrix()
assert m1.tomatrix().T == transpose(m1).tomatrix()
assert m1.tomatrix().C == conjugate(m1).tomatrix()
assert m1.tomatrix().H == adjoint(m1).tomatrix()
assert m1.tomatrix().T == m1.transpose().tomatrix()
assert m1.tomatrix().C == m1.conjugate().tomatrix()
assert m1.tomatrix().H == m1.adjoint().tomatrix()
raises(ValueError, lambda: permutedims(m1, (0,)))
raises(ValueError, lambda: permutedims(m1, (0, 0)))
raises(ValueError, lambda: permutedims(m1, (1, 2, 0)))
# Some tests with random arrays:
dims = 6
shape = [random.randint(1,5) for i in range(dims)]
elems = [random.random() for i in range(tensorproduct(*shape))]
ra = ArrayType(elems, shape)
perm = list(range(dims))
# Randomize the permutation:
random.shuffle(perm)
# Test inverse permutation:
assert permutedims(permutedims(ra, perm), _af_invert(perm)) == ra
# Test that permuted shape corresponds to action by `Permutation`:
assert permutedims(ra, perm).shape == tuple(Permutation(perm)(shape))
z = ArrayType.zeros(4,5,6,7)
assert permutedims(z, (2, 3, 1, 0)).shape == (6, 7, 5, 4)
assert permutedims(z, [2, 3, 1, 0]).shape == (6, 7, 5, 4)
assert permutedims(z, Permutation([2, 3, 1, 0])).shape == (6, 7, 5, 4)
po = ArrayType(sa, [2, 2, 3, 3, 2, 2])
raises(ValueError, lambda: permutedims(po, (1, 1)))
raises(ValueError, lambda: po.transpose())
raises(ValueError, lambda: po.adjoint())
assert permutedims(po, reversed(range(po.rank()))) == ArrayType(
[[[[[[sa[0], sa[72]], [sa[36], sa[108]]], [[sa[12], sa[84]], [sa[48], sa[120]]], [[sa[24],
sa[96]], [sa[60], sa[132]]]],
[[[sa[4], sa[76]], [sa[40], sa[112]]], [[sa[16],
sa[88]], [sa[52], sa[124]]],
[[sa[28], sa[100]], [sa[64], sa[136]]]],
[[[sa[8],
sa[80]], [sa[44], sa[116]]], [[sa[20], sa[92]], [sa[56], sa[128]]], [[sa[32],
sa[104]], [sa[68], sa[140]]]]],
[[[[sa[2], sa[74]], [sa[38], sa[110]]], [[sa[14],
sa[86]], [sa[50], sa[122]]], [[sa[26], sa[98]], [sa[62], sa[134]]]],
[[[sa[6],
sa[78]], [sa[42], sa[114]]], [[sa[18], sa[90]], [sa[54], sa[126]]], [[sa[30],
sa[102]], [sa[66], sa[138]]]],
[[[sa[10], sa[82]], [sa[46], sa[118]]], [[sa[22],
sa[94]], [sa[58], sa[130]]],
[[sa[34], sa[106]], [sa[70], sa[142]]]]]],
[[[[[sa[1],
sa[73]], [sa[37], sa[109]]], [[sa[13], sa[85]], [sa[49], sa[121]]], [[sa[25],
sa[97]], [sa[61], sa[133]]]],
[[[sa[5], sa[77]], [sa[41], sa[113]]], [[sa[17],
sa[89]], [sa[53], sa[125]]],
[[sa[29], sa[101]], [sa[65], sa[137]]]],
[[[sa[9],
sa[81]], [sa[45], sa[117]]], [[sa[21], sa[93]], [sa[57], sa[129]]], [[sa[33],
sa[105]], [sa[69], sa[141]]]]],
[[[[sa[3], sa[75]], [sa[39], sa[111]]], [[sa[15],
sa[87]], [sa[51], sa[123]]], [[sa[27], sa[99]], [sa[63], sa[135]]]],
[[[sa[7],
sa[79]], [sa[43], sa[115]]], [[sa[19], sa[91]], [sa[55], sa[127]]], [[sa[31],
sa[103]], [sa[67], sa[139]]]],
[[[sa[11], sa[83]], [sa[47], sa[119]]], [[sa[23],
sa[95]], [sa[59], sa[131]]],
[[sa[35], sa[107]], [sa[71], sa[143]]]]]]])
assert permutedims(po, (1, 0, 2, 3, 4, 5)) == ArrayType(
[[[[[[sa[0], sa[1]], [sa[2], sa[3]]], [[sa[4], sa[5]], [sa[6], sa[7]]], [[sa[8], sa[9]], [sa[10],
sa[11]]]],
[[[sa[12], sa[13]], [sa[14], sa[15]]], [[sa[16], sa[17]], [sa[18],
sa[19]]], [[sa[20], sa[21]], [sa[22], sa[23]]]],
[[[sa[24], sa[25]], [sa[26],
sa[27]]], [[sa[28], sa[29]], [sa[30], sa[31]]], [[sa[32], sa[33]], [sa[34],
sa[35]]]]],
[[[[sa[72], sa[73]], [sa[74], sa[75]]], [[sa[76], sa[77]], [sa[78],
sa[79]]], [[sa[80], sa[81]], [sa[82], sa[83]]]],
[[[sa[84], sa[85]], [sa[86],
sa[87]]], [[sa[88], sa[89]], [sa[90], sa[91]]], [[sa[92], sa[93]], [sa[94],
sa[95]]]],
[[[sa[96], sa[97]], [sa[98], sa[99]]], [[sa[100], sa[101]], [sa[102],
sa[103]]],
[[sa[104], sa[105]], [sa[106], sa[107]]]]]], [[[[[sa[36], sa[37]], [sa[38],
sa[39]]],
[[sa[40], sa[41]], [sa[42], sa[43]]],
[[sa[44], sa[45]], [sa[46],
sa[47]]]],
[[[sa[48], sa[49]], [sa[50], sa[51]]],
[[sa[52], sa[53]], [sa[54],
sa[55]]],
[[sa[56], sa[57]], [sa[58], sa[59]]]],
[[[sa[60], sa[61]], [sa[62],
sa[63]]],
[[sa[64], sa[65]], [sa[66], sa[67]]],
[[sa[68], sa[69]], [sa[70],
sa[71]]]]], [
[[[sa[108], sa[109]], [sa[110], sa[111]]],
[[sa[112], sa[113]], [sa[114],
sa[115]]],
[[sa[116], sa[117]], [sa[118], sa[119]]]],
[[[sa[120], sa[121]], [sa[122],
sa[123]]],
[[sa[124], sa[125]], [sa[126], sa[127]]],
[[sa[128], sa[129]], [sa[130],
sa[131]]]],
[[[sa[132], sa[133]], [sa[134], sa[135]]],
[[sa[136], sa[137]], [sa[138],
sa[139]]],
[[sa[140], sa[141]], [sa[142], sa[143]]]]]]])
assert permutedims(po, (0, 2, 1, 4, 3, 5)) == ArrayType(
[[[[[[sa[0], sa[1]], [sa[4], sa[5]], [sa[8], sa[9]]], [[sa[2], sa[3]], [sa[6], sa[7]], [sa[10],
sa[11]]]],
[[[sa[36], sa[37]], [sa[40], sa[41]], [sa[44], sa[45]]], [[sa[38],
sa[39]], [sa[42], sa[43]], [sa[46], sa[47]]]]],
[[[[sa[12], sa[13]], [sa[16],
sa[17]], [sa[20], sa[21]]], [[sa[14], sa[15]], [sa[18], sa[19]], [sa[22],
sa[23]]]],
[[[sa[48], sa[49]], [sa[52], sa[53]], [sa[56], sa[57]]], [[sa[50],
sa[51]], [sa[54], sa[55]], [sa[58], sa[59]]]]],
[[[[sa[24], sa[25]], [sa[28],
sa[29]], [sa[32], sa[33]]], [[sa[26], sa[27]], [sa[30], sa[31]], [sa[34],
sa[35]]]],
[[[sa[60], sa[61]], [sa[64], sa[65]], [sa[68], sa[69]]], [[sa[62],
sa[63]], [sa[66], sa[67]], [sa[70], sa[71]]]]]],
[[[[[sa[72], sa[73]], [sa[76],
sa[77]], [sa[80], sa[81]]], [[sa[74], sa[75]], [sa[78], sa[79]], [sa[82],
sa[83]]]],
[[[sa[108], sa[109]], [sa[112], sa[113]], [sa[116], sa[117]]], [[sa[110],
sa[111]], [sa[114], sa[115]],
[sa[118], sa[119]]]]],
[[[[sa[84], sa[85]], [sa[88],
sa[89]], [sa[92], sa[93]]], [[sa[86], sa[87]], [sa[90], sa[91]], [sa[94],
sa[95]]]],
[[[sa[120], sa[121]], [sa[124], sa[125]], [sa[128], sa[129]]], [[sa[122],
sa[123]], [sa[126], sa[127]],
[sa[130], sa[131]]]]],
[[[[sa[96], sa[97]], [sa[100],
sa[101]], [sa[104], sa[105]]], [[sa[98], sa[99]], [sa[102], sa[103]], [sa[106],
sa[107]]]],
[[[sa[132], sa[133]], [sa[136], sa[137]], [sa[140], sa[141]]], [[sa[134],
sa[135]], [sa[138], sa[139]],
[sa[142], sa[143]]]]]]])
po2 = po.reshape(4, 9, 2, 2)
assert po2 == ArrayType([[[[sa[0], sa[1]], [sa[2], sa[3]]], [[sa[4], sa[5]], [sa[6], sa[7]]], [[sa[8], sa[9]], [sa[10], sa[11]]], [[sa[12], sa[13]], [sa[14], sa[15]]], [[sa[16], sa[17]], [sa[18], sa[19]]], [[sa[20], sa[21]], [sa[22], sa[23]]], [[sa[24], sa[25]], [sa[26], sa[27]]], [[sa[28], sa[29]], [sa[30], sa[31]]], [[sa[32], sa[33]], [sa[34], sa[35]]]], [[[sa[36], sa[37]], [sa[38], sa[39]]], [[sa[40], sa[41]], [sa[42], sa[43]]], [[sa[44], sa[45]], [sa[46], sa[47]]], [[sa[48], sa[49]], [sa[50], sa[51]]], [[sa[52], sa[53]], [sa[54], sa[55]]], [[sa[56], sa[57]], [sa[58], sa[59]]], [[sa[60], sa[61]], [sa[62], sa[63]]], [[sa[64], sa[65]], [sa[66], sa[67]]], [[sa[68], sa[69]], [sa[70], sa[71]]]], [[[sa[72], sa[73]], [sa[74], sa[75]]], [[sa[76], sa[77]], [sa[78], sa[79]]], [[sa[80], sa[81]], [sa[82], sa[83]]], [[sa[84], sa[85]], [sa[86], sa[87]]], [[sa[88], sa[89]], [sa[90], sa[91]]], [[sa[92], sa[93]], [sa[94], sa[95]]], [[sa[96], sa[97]], [sa[98], sa[99]]], [[sa[100], sa[101]], [sa[102], sa[103]]], [[sa[104], sa[105]], [sa[106], sa[107]]]], [[[sa[108], sa[109]], [sa[110], sa[111]]], [[sa[112], sa[113]], [sa[114], sa[115]]], [[sa[116], sa[117]], [sa[118], sa[119]]], [[sa[120], sa[121]], [sa[122], sa[123]]], [[sa[124], sa[125]], [sa[126], sa[127]]], [[sa[128], sa[129]], [sa[130], sa[131]]], [[sa[132], sa[133]], [sa[134], sa[135]]], [[sa[136], sa[137]], [sa[138], sa[139]]], [[sa[140], sa[141]], [sa[142], sa[143]]]]])
assert permutedims(po2, (3, 2, 0, 1)) == ArrayType([[[[sa[0], sa[4], sa[8], sa[12], sa[16], sa[20], sa[24], sa[28], sa[32]], [sa[36], sa[40], sa[44], sa[48], sa[52], sa[56], sa[60], sa[64], sa[68]], [sa[72], sa[76], sa[80], sa[84], sa[88], sa[92], sa[96], sa[100], sa[104]], [sa[108], sa[112], sa[116], sa[120], sa[124], sa[128], sa[132], sa[136], sa[140]]], [[sa[2], sa[6], sa[10], sa[14], sa[18], sa[22], sa[26], sa[30], sa[34]], [sa[38], sa[42], sa[46], sa[50], sa[54], sa[58], sa[62], sa[66], sa[70]], [sa[74], sa[78], sa[82], sa[86], sa[90], sa[94], sa[98], sa[102], sa[106]], [sa[110], sa[114], sa[118], sa[122], sa[126], sa[130], sa[134], sa[138], sa[142]]]], [[[sa[1], sa[5], sa[9], sa[13], sa[17], sa[21], sa[25], sa[29], sa[33]], [sa[37], sa[41], sa[45], sa[49], sa[53], sa[57], sa[61], sa[65], sa[69]], [sa[73], sa[77], sa[81], sa[85], sa[89], sa[93], sa[97], sa[101], sa[105]], [sa[109], sa[113], sa[117], sa[121], sa[125], sa[129], sa[133], sa[137], sa[141]]], [[sa[3], sa[7], sa[11], sa[15], sa[19], sa[23], sa[27], sa[31], sa[35]], [sa[39], sa[43], sa[47], sa[51], sa[55], sa[59], sa[63], sa[67], sa[71]], [sa[75], sa[79], sa[83], sa[87], sa[91], sa[95], sa[99], sa[103], sa[107]], [sa[111], sa[115], sa[119], sa[123], sa[127], sa[131], sa[135], sa[139], sa[143]]]]])
# test for large scale sparse array
for SparseArrayType in [ImmutableSparseNDimArray, MutableSparseNDimArray]:
A = SparseArrayType({1:1, 10000:2}, (10000, 20000, 10000))
assert permutedims(A, (0, 1, 2)) == A
assert permutedims(A, (1, 0, 2)) == SparseArrayType({1: 1, 100000000: 2}, (20000, 10000, 10000))
B = SparseArrayType({1:1, 20000:2}, (10000, 20000))
assert B.transpose() == SparseArrayType({10000: 1, 1: 2}, (20000, 10000))
def test_permutedims_with_indices():
A = Array(range(32)).reshape(2, 2, 2, 2, 2)
indices_new = list("abcde")
indices_old = list("ebdac")
new_A = permutedims(A, index_order_new=indices_new, index_order_old=indices_old)
for a, b, c, d, e in itertools.product(range(2), range(2), range(2), range(2), range(2)):
assert new_A[a, b, c, d, e] == A[e, b, d, a, c]
indices_old = list("cabed")
new_A = permutedims(A, index_order_new=indices_new, index_order_old=indices_old)
for a, b, c, d, e in itertools.product(range(2), range(2), range(2), range(2), range(2)):
assert new_A[a, b, c, d, e] == A[c, a, b, e, d]
raises(ValueError, lambda: permutedims(A, index_order_old=list("aacde"), index_order_new=list("abcde")))
raises(ValueError, lambda: permutedims(A, index_order_old=list("abcde"), index_order_new=list("abcce")))
raises(ValueError, lambda: permutedims(A, index_order_old=list("abcde"), index_order_new=list("abce")))
raises(ValueError, lambda: permutedims(A, index_order_old=list("abce"), index_order_new=list("abce")))
raises(ValueError, lambda: permutedims(A, [2, 1, 0, 3, 4], index_order_old=list("abcde")))
raises(ValueError, lambda: permutedims(A, [2, 1, 0, 3, 4], index_order_new=list("abcde")))
def test_flatten():
from sympy.matrices.dense import Matrix
for ArrayType in [ImmutableDenseNDimArray, ImmutableSparseNDimArray, Matrix]:
A = ArrayType(range(24)).reshape(4, 6)
assert list(Flatten(A)) == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]
for i, v in enumerate(Flatten(A)):
assert i == v
def test_tensordiagonal():
from sympy.matrices.dense import eye
expr = Array(range(9)).reshape(3, 3)
raises(ValueError, lambda: tensordiagonal(expr, [0], [1]))
raises(ValueError, lambda: tensordiagonal(expr, [0, 0]))
assert tensordiagonal(eye(3), [0, 1]) == Array([1, 1, 1])
assert tensordiagonal(expr, [0, 1]) == Array([0, 4, 8])
x, y, z = symbols("x y z")
expr2 = tensorproduct([x, y, z], expr)
assert tensordiagonal(expr2, [1, 2]) == Array([[0, 4*x, 8*x], [0, 4*y, 8*y], [0, 4*z, 8*z]])
assert tensordiagonal(expr2, [0, 1]) == Array([[0, 3*y, 6*z], [x, 4*y, 7*z], [2*x, 5*y, 8*z]])
assert tensordiagonal(expr2, [0, 1, 2]) == Array([0, 4*y, 8*z])
# assert tensordiagonal(expr2, [0]) == permutedims(expr2, [1, 2, 0])
# assert tensordiagonal(expr2, [1]) == permutedims(expr2, [0, 2, 1])
# assert tensordiagonal(expr2, [2]) == expr2
# assert tensordiagonal(expr2, [1], [2]) == expr2
# assert tensordiagonal(expr2, [0], [1]) == permutedims(expr2, [2, 0, 1])
a, b, c, X, Y, Z = symbols("a b c X Y Z")
expr3 = tensorproduct([x, y, z], [1, 2, 3], [a, b, c], [X, Y, Z])
assert tensordiagonal(expr3, [0, 1, 2, 3]) == Array([x*a*X, 2*y*b*Y, 3*z*c*Z])
assert tensordiagonal(expr3, [0, 1], [2, 3]) == tensorproduct([x, 2*y, 3*z], [a*X, b*Y, c*Z])
# assert tensordiagonal(expr3, [0], [1, 2], [3]) == tensorproduct([x, y, z], [a, 2*b, 3*c], [X, Y, Z])
assert tensordiagonal(tensordiagonal(expr3, [2, 3]), [0, 1]) == tensorproduct([a*X, b*Y, c*Z], [x, 2*y, 3*z])
raises(ValueError, lambda: tensordiagonal([[1, 2, 3], [4, 5, 6]], [0, 1]))
raises(ValueError, lambda: tensordiagonal(expr3.reshape(3, 3, 9), [1, 2]))

View File

@ -0,0 +1,452 @@
from copy import copy
from sympy.tensor.array.dense_ndim_array import ImmutableDenseNDimArray
from sympy.core.containers import Dict
from sympy.core.function import diff
from sympy.core.numbers import Rational
from sympy.core.singleton import S
from sympy.core.symbol import (Symbol, symbols)
from sympy.matrices import SparseMatrix
from sympy.tensor.indexed import (Indexed, IndexedBase)
from sympy.matrices import Matrix
from sympy.tensor.array.sparse_ndim_array import ImmutableSparseNDimArray
from sympy.testing.pytest import raises
def test_ndim_array_initiation():
arr_with_no_elements = ImmutableDenseNDimArray([], shape=(0,))
assert len(arr_with_no_elements) == 0
assert arr_with_no_elements.rank() == 1
raises(ValueError, lambda: ImmutableDenseNDimArray([0], shape=(0,)))
raises(ValueError, lambda: ImmutableDenseNDimArray([1, 2, 3], shape=(0,)))
raises(ValueError, lambda: ImmutableDenseNDimArray([], shape=()))
raises(ValueError, lambda: ImmutableSparseNDimArray([0], shape=(0,)))
raises(ValueError, lambda: ImmutableSparseNDimArray([1, 2, 3], shape=(0,)))
raises(ValueError, lambda: ImmutableSparseNDimArray([], shape=()))
arr_with_one_element = ImmutableDenseNDimArray([23])
assert len(arr_with_one_element) == 1
assert arr_with_one_element[0] == 23
assert arr_with_one_element[:] == ImmutableDenseNDimArray([23])
assert arr_with_one_element.rank() == 1
arr_with_symbol_element = ImmutableDenseNDimArray([Symbol('x')])
assert len(arr_with_symbol_element) == 1
assert arr_with_symbol_element[0] == Symbol('x')
assert arr_with_symbol_element[:] == ImmutableDenseNDimArray([Symbol('x')])
assert arr_with_symbol_element.rank() == 1
number5 = 5
vector = ImmutableDenseNDimArray.zeros(number5)
assert len(vector) == number5
assert vector.shape == (number5,)
assert vector.rank() == 1
vector = ImmutableSparseNDimArray.zeros(number5)
assert len(vector) == number5
assert vector.shape == (number5,)
assert vector._sparse_array == Dict()
assert vector.rank() == 1
n_dim_array = ImmutableDenseNDimArray(range(3**4), (3, 3, 3, 3,))
assert len(n_dim_array) == 3 * 3 * 3 * 3
assert n_dim_array.shape == (3, 3, 3, 3)
assert n_dim_array.rank() == 4
array_shape = (3, 3, 3, 3)
sparse_array = ImmutableSparseNDimArray.zeros(*array_shape)
assert len(sparse_array._sparse_array) == 0
assert len(sparse_array) == 3 * 3 * 3 * 3
assert n_dim_array.shape == array_shape
assert n_dim_array.rank() == 4
one_dim_array = ImmutableDenseNDimArray([2, 3, 1])
assert len(one_dim_array) == 3
assert one_dim_array.shape == (3,)
assert one_dim_array.rank() == 1
assert one_dim_array.tolist() == [2, 3, 1]
shape = (3, 3)
array_with_many_args = ImmutableSparseNDimArray.zeros(*shape)
assert len(array_with_many_args) == 3 * 3
assert array_with_many_args.shape == shape
assert array_with_many_args[0, 0] == 0
assert array_with_many_args.rank() == 2
shape = (int(3), int(3))
array_with_long_shape = ImmutableSparseNDimArray.zeros(*shape)
assert len(array_with_long_shape) == 3 * 3
assert array_with_long_shape.shape == shape
assert array_with_long_shape[int(0), int(0)] == 0
assert array_with_long_shape.rank() == 2
vector_with_long_shape = ImmutableDenseNDimArray(range(5), int(5))
assert len(vector_with_long_shape) == 5
assert vector_with_long_shape.shape == (int(5),)
assert vector_with_long_shape.rank() == 1
raises(ValueError, lambda: vector_with_long_shape[int(5)])
from sympy.abc import x
for ArrayType in [ImmutableDenseNDimArray, ImmutableSparseNDimArray]:
rank_zero_array = ArrayType(x)
assert len(rank_zero_array) == 1
assert rank_zero_array.shape == ()
assert rank_zero_array.rank() == 0
assert rank_zero_array[()] == x
raises(ValueError, lambda: rank_zero_array[0])
def test_reshape():
array = ImmutableDenseNDimArray(range(50), 50)
assert array.shape == (50,)
assert array.rank() == 1
array = array.reshape(5, 5, 2)
assert array.shape == (5, 5, 2)
assert array.rank() == 3
assert len(array) == 50
def test_getitem():
for ArrayType in [ImmutableDenseNDimArray, ImmutableSparseNDimArray]:
array = ArrayType(range(24)).reshape(2, 3, 4)
assert array.tolist() == [[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]]
assert array[0] == ArrayType([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]])
assert array[0, 0] == ArrayType([0, 1, 2, 3])
value = 0
for i in range(2):
for j in range(3):
for k in range(4):
assert array[i, j, k] == value
value += 1
raises(ValueError, lambda: array[3, 4, 5])
raises(ValueError, lambda: array[3, 4, 5, 6])
raises(ValueError, lambda: array[3, 4, 5, 3:4])
def test_iterator():
array = ImmutableDenseNDimArray(range(4), (2, 2))
assert array[0] == ImmutableDenseNDimArray([0, 1])
assert array[1] == ImmutableDenseNDimArray([2, 3])
array = array.reshape(4)
j = 0
for i in array:
assert i == j
j += 1
def test_sparse():
sparse_array = ImmutableSparseNDimArray([0, 0, 0, 1], (2, 2))
assert len(sparse_array) == 2 * 2
# dictionary where all data is, only non-zero entries are actually stored:
assert len(sparse_array._sparse_array) == 1
assert sparse_array.tolist() == [[0, 0], [0, 1]]
for i, j in zip(sparse_array, [[0, 0], [0, 1]]):
assert i == ImmutableSparseNDimArray(j)
def sparse_assignment():
sparse_array[0, 0] = 123
assert len(sparse_array._sparse_array) == 1
raises(TypeError, sparse_assignment)
assert len(sparse_array._sparse_array) == 1
assert sparse_array[0, 0] == 0
assert sparse_array/0 == ImmutableSparseNDimArray([[S.NaN, S.NaN], [S.NaN, S.ComplexInfinity]], (2, 2))
# test for large scale sparse array
# equality test
assert ImmutableSparseNDimArray.zeros(100000, 200000) == ImmutableSparseNDimArray.zeros(100000, 200000)
# __mul__ and __rmul__
a = ImmutableSparseNDimArray({200001: 1}, (100000, 200000))
assert a * 3 == ImmutableSparseNDimArray({200001: 3}, (100000, 200000))
assert 3 * a == ImmutableSparseNDimArray({200001: 3}, (100000, 200000))
assert a * 0 == ImmutableSparseNDimArray({}, (100000, 200000))
assert 0 * a == ImmutableSparseNDimArray({}, (100000, 200000))
# __truediv__
assert a/3 == ImmutableSparseNDimArray({200001: Rational(1, 3)}, (100000, 200000))
# __neg__
assert -a == ImmutableSparseNDimArray({200001: -1}, (100000, 200000))
def test_calculation():
a = ImmutableDenseNDimArray([1]*9, (3, 3))
b = ImmutableDenseNDimArray([9]*9, (3, 3))
c = a + b
for i in c:
assert i == ImmutableDenseNDimArray([10, 10, 10])
assert c == ImmutableDenseNDimArray([10]*9, (3, 3))
assert c == ImmutableSparseNDimArray([10]*9, (3, 3))
c = b - a
for i in c:
assert i == ImmutableDenseNDimArray([8, 8, 8])
assert c == ImmutableDenseNDimArray([8]*9, (3, 3))
assert c == ImmutableSparseNDimArray([8]*9, (3, 3))
def test_ndim_array_converting():
dense_array = ImmutableDenseNDimArray([1, 2, 3, 4], (2, 2))
alist = dense_array.tolist()
assert alist == [[1, 2], [3, 4]]
matrix = dense_array.tomatrix()
assert (isinstance(matrix, Matrix))
for i in range(len(dense_array)):
assert dense_array[dense_array._get_tuple_index(i)] == matrix[i]
assert matrix.shape == dense_array.shape
assert ImmutableDenseNDimArray(matrix) == dense_array
assert ImmutableDenseNDimArray(matrix.as_immutable()) == dense_array
assert ImmutableDenseNDimArray(matrix.as_mutable()) == dense_array
sparse_array = ImmutableSparseNDimArray([1, 2, 3, 4], (2, 2))
alist = sparse_array.tolist()
assert alist == [[1, 2], [3, 4]]
matrix = sparse_array.tomatrix()
assert(isinstance(matrix, SparseMatrix))
for i in range(len(sparse_array)):
assert sparse_array[sparse_array._get_tuple_index(i)] == matrix[i]
assert matrix.shape == sparse_array.shape
assert ImmutableSparseNDimArray(matrix) == sparse_array
assert ImmutableSparseNDimArray(matrix.as_immutable()) == sparse_array
assert ImmutableSparseNDimArray(matrix.as_mutable()) == sparse_array
def test_converting_functions():
arr_list = [1, 2, 3, 4]
arr_matrix = Matrix(((1, 2), (3, 4)))
# list
arr_ndim_array = ImmutableDenseNDimArray(arr_list, (2, 2))
assert (isinstance(arr_ndim_array, ImmutableDenseNDimArray))
assert arr_matrix.tolist() == arr_ndim_array.tolist()
# Matrix
arr_ndim_array = ImmutableDenseNDimArray(arr_matrix)
assert (isinstance(arr_ndim_array, ImmutableDenseNDimArray))
assert arr_matrix.tolist() == arr_ndim_array.tolist()
assert arr_matrix.shape == arr_ndim_array.shape
def test_equality():
first_list = [1, 2, 3, 4]
second_list = [1, 2, 3, 4]
third_list = [4, 3, 2, 1]
assert first_list == second_list
assert first_list != third_list
first_ndim_array = ImmutableDenseNDimArray(first_list, (2, 2))
second_ndim_array = ImmutableDenseNDimArray(second_list, (2, 2))
fourth_ndim_array = ImmutableDenseNDimArray(first_list, (2, 2))
assert first_ndim_array == second_ndim_array
def assignment_attempt(a):
a[0, 0] = 0
raises(TypeError, lambda: assignment_attempt(second_ndim_array))
assert first_ndim_array == second_ndim_array
assert first_ndim_array == fourth_ndim_array
def test_arithmetic():
a = ImmutableDenseNDimArray([3 for i in range(9)], (3, 3))
b = ImmutableDenseNDimArray([7 for i in range(9)], (3, 3))
c1 = a + b
c2 = b + a
assert c1 == c2
d1 = a - b
d2 = b - a
assert d1 == d2 * (-1)
e1 = a * 5
e2 = 5 * a
e3 = copy(a)
e3 *= 5
assert e1 == e2 == e3
f1 = a / 5
f2 = copy(a)
f2 /= 5
assert f1 == f2
assert f1[0, 0] == f1[0, 1] == f1[0, 2] == f1[1, 0] == f1[1, 1] == \
f1[1, 2] == f1[2, 0] == f1[2, 1] == f1[2, 2] == Rational(3, 5)
assert type(a) == type(b) == type(c1) == type(c2) == type(d1) == type(d2) \
== type(e1) == type(e2) == type(e3) == type(f1)
z0 = -a
assert z0 == ImmutableDenseNDimArray([-3 for i in range(9)], (3, 3))
def test_higher_dimenions():
m3 = ImmutableDenseNDimArray(range(10, 34), (2, 3, 4))
assert m3.tolist() == [[[10, 11, 12, 13],
[14, 15, 16, 17],
[18, 19, 20, 21]],
[[22, 23, 24, 25],
[26, 27, 28, 29],
[30, 31, 32, 33]]]
assert m3._get_tuple_index(0) == (0, 0, 0)
assert m3._get_tuple_index(1) == (0, 0, 1)
assert m3._get_tuple_index(4) == (0, 1, 0)
assert m3._get_tuple_index(12) == (1, 0, 0)
assert str(m3) == '[[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]]'
m3_rebuilt = ImmutableDenseNDimArray([[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]])
assert m3 == m3_rebuilt
m3_other = ImmutableDenseNDimArray([[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]], (2, 3, 4))
assert m3 == m3_other
def test_rebuild_immutable_arrays():
sparr = ImmutableSparseNDimArray(range(10, 34), (2, 3, 4))
densarr = ImmutableDenseNDimArray(range(10, 34), (2, 3, 4))
assert sparr == sparr.func(*sparr.args)
assert densarr == densarr.func(*densarr.args)
def test_slices():
md = ImmutableDenseNDimArray(range(10, 34), (2, 3, 4))
assert md[:] == ImmutableDenseNDimArray(range(10, 34), (2, 3, 4))
assert md[:, :, 0].tomatrix() == Matrix([[10, 14, 18], [22, 26, 30]])
assert md[0, 1:2, :].tomatrix() == Matrix([[14, 15, 16, 17]])
assert md[0, 1:3, :].tomatrix() == Matrix([[14, 15, 16, 17], [18, 19, 20, 21]])
assert md[:, :, :] == md
sd = ImmutableSparseNDimArray(range(10, 34), (2, 3, 4))
assert sd == ImmutableSparseNDimArray(md)
assert sd[:] == ImmutableSparseNDimArray(range(10, 34), (2, 3, 4))
assert sd[:, :, 0].tomatrix() == Matrix([[10, 14, 18], [22, 26, 30]])
assert sd[0, 1:2, :].tomatrix() == Matrix([[14, 15, 16, 17]])
assert sd[0, 1:3, :].tomatrix() == Matrix([[14, 15, 16, 17], [18, 19, 20, 21]])
assert sd[:, :, :] == sd
def test_diff_and_applyfunc():
from sympy.abc import x, y, z
md = ImmutableDenseNDimArray([[x, y], [x*z, x*y*z]])
assert md.diff(x) == ImmutableDenseNDimArray([[1, 0], [z, y*z]])
assert diff(md, x) == ImmutableDenseNDimArray([[1, 0], [z, y*z]])
sd = ImmutableSparseNDimArray(md)
assert sd == ImmutableSparseNDimArray([x, y, x*z, x*y*z], (2, 2))
assert sd.diff(x) == ImmutableSparseNDimArray([[1, 0], [z, y*z]])
assert diff(sd, x) == ImmutableSparseNDimArray([[1, 0], [z, y*z]])
mdn = md.applyfunc(lambda x: x*3)
assert mdn == ImmutableDenseNDimArray([[3*x, 3*y], [3*x*z, 3*x*y*z]])
assert md != mdn
sdn = sd.applyfunc(lambda x: x/2)
assert sdn == ImmutableSparseNDimArray([[x/2, y/2], [x*z/2, x*y*z/2]])
assert sd != sdn
sdp = sd.applyfunc(lambda x: x+1)
assert sdp == ImmutableSparseNDimArray([[x + 1, y + 1], [x*z + 1, x*y*z + 1]])
assert sd != sdp
def test_op_priority():
from sympy.abc import x
md = ImmutableDenseNDimArray([1, 2, 3])
e1 = (1+x)*md
e2 = md*(1+x)
assert e1 == ImmutableDenseNDimArray([1+x, 2+2*x, 3+3*x])
assert e1 == e2
sd = ImmutableSparseNDimArray([1, 2, 3])
e3 = (1+x)*sd
e4 = sd*(1+x)
assert e3 == ImmutableDenseNDimArray([1+x, 2+2*x, 3+3*x])
assert e3 == e4
def test_symbolic_indexing():
x, y, z, w = symbols("x y z w")
M = ImmutableDenseNDimArray([[x, y], [z, w]])
i, j = symbols("i, j")
Mij = M[i, j]
assert isinstance(Mij, Indexed)
Ms = ImmutableSparseNDimArray([[2, 3*x], [4, 5]])
msij = Ms[i, j]
assert isinstance(msij, Indexed)
for oi, oj in [(0, 0), (0, 1), (1, 0), (1, 1)]:
assert Mij.subs({i: oi, j: oj}) == M[oi, oj]
assert msij.subs({i: oi, j: oj}) == Ms[oi, oj]
A = IndexedBase("A", (0, 2))
assert A[0, 0].subs(A, M) == x
assert A[i, j].subs(A, M) == M[i, j]
assert M[i, j].subs(M, A) == A[i, j]
assert isinstance(M[3 * i - 2, j], Indexed)
assert M[3 * i - 2, j].subs({i: 1, j: 0}) == M[1, 0]
assert isinstance(M[i, 0], Indexed)
assert M[i, 0].subs(i, 0) == M[0, 0]
assert M[0, i].subs(i, 1) == M[0, 1]
assert M[i, j].diff(x) == ImmutableDenseNDimArray([[1, 0], [0, 0]])[i, j]
assert Ms[i, j].diff(x) == ImmutableSparseNDimArray([[0, 3], [0, 0]])[i, j]
Mo = ImmutableDenseNDimArray([1, 2, 3])
assert Mo[i].subs(i, 1) == 2
Mos = ImmutableSparseNDimArray([1, 2, 3])
assert Mos[i].subs(i, 1) == 2
raises(ValueError, lambda: M[i, 2])
raises(ValueError, lambda: M[i, -1])
raises(ValueError, lambda: M[2, i])
raises(ValueError, lambda: M[-1, i])
raises(ValueError, lambda: Ms[i, 2])
raises(ValueError, lambda: Ms[i, -1])
raises(ValueError, lambda: Ms[2, i])
raises(ValueError, lambda: Ms[-1, i])
def test_issue_12665():
# Testing Python 3 hash of immutable arrays:
arr = ImmutableDenseNDimArray([1, 2, 3])
# This should NOT raise an exception:
hash(arr)
def test_zeros_without_shape():
arr = ImmutableDenseNDimArray.zeros()
assert arr == ImmutableDenseNDimArray(0)
def test_issue_21870():
a0 = ImmutableDenseNDimArray(0)
assert a0.rank() == 0
a1 = ImmutableDenseNDimArray(a0)
assert a1.rank() == 0

View File

@ -0,0 +1,374 @@
from copy import copy
from sympy.tensor.array.dense_ndim_array import MutableDenseNDimArray
from sympy.core.function import diff
from sympy.core.numbers import Rational
from sympy.core.singleton import S
from sympy.core.symbol import Symbol
from sympy.core.sympify import sympify
from sympy.matrices import SparseMatrix
from sympy.matrices import Matrix
from sympy.tensor.array.sparse_ndim_array import MutableSparseNDimArray
from sympy.testing.pytest import raises
def test_ndim_array_initiation():
arr_with_one_element = MutableDenseNDimArray([23])
assert len(arr_with_one_element) == 1
assert arr_with_one_element[0] == 23
assert arr_with_one_element.rank() == 1
raises(ValueError, lambda: arr_with_one_element[1])
arr_with_symbol_element = MutableDenseNDimArray([Symbol('x')])
assert len(arr_with_symbol_element) == 1
assert arr_with_symbol_element[0] == Symbol('x')
assert arr_with_symbol_element.rank() == 1
number5 = 5
vector = MutableDenseNDimArray.zeros(number5)
assert len(vector) == number5
assert vector.shape == (number5,)
assert vector.rank() == 1
raises(ValueError, lambda: arr_with_one_element[5])
vector = MutableSparseNDimArray.zeros(number5)
assert len(vector) == number5
assert vector.shape == (number5,)
assert vector._sparse_array == {}
assert vector.rank() == 1
n_dim_array = MutableDenseNDimArray(range(3**4), (3, 3, 3, 3,))
assert len(n_dim_array) == 3 * 3 * 3 * 3
assert n_dim_array.shape == (3, 3, 3, 3)
assert n_dim_array.rank() == 4
raises(ValueError, lambda: n_dim_array[0, 0, 0, 3])
raises(ValueError, lambda: n_dim_array[3, 0, 0, 0])
raises(ValueError, lambda: n_dim_array[3**4])
array_shape = (3, 3, 3, 3)
sparse_array = MutableSparseNDimArray.zeros(*array_shape)
assert len(sparse_array._sparse_array) == 0
assert len(sparse_array) == 3 * 3 * 3 * 3
assert n_dim_array.shape == array_shape
assert n_dim_array.rank() == 4
one_dim_array = MutableDenseNDimArray([2, 3, 1])
assert len(one_dim_array) == 3
assert one_dim_array.shape == (3,)
assert one_dim_array.rank() == 1
assert one_dim_array.tolist() == [2, 3, 1]
shape = (3, 3)
array_with_many_args = MutableSparseNDimArray.zeros(*shape)
assert len(array_with_many_args) == 3 * 3
assert array_with_many_args.shape == shape
assert array_with_many_args[0, 0] == 0
assert array_with_many_args.rank() == 2
shape = (int(3), int(3))
array_with_long_shape = MutableSparseNDimArray.zeros(*shape)
assert len(array_with_long_shape) == 3 * 3
assert array_with_long_shape.shape == shape
assert array_with_long_shape[int(0), int(0)] == 0
assert array_with_long_shape.rank() == 2
vector_with_long_shape = MutableDenseNDimArray(range(5), int(5))
assert len(vector_with_long_shape) == 5
assert vector_with_long_shape.shape == (int(5),)
assert vector_with_long_shape.rank() == 1
raises(ValueError, lambda: vector_with_long_shape[int(5)])
from sympy.abc import x
for ArrayType in [MutableDenseNDimArray, MutableSparseNDimArray]:
rank_zero_array = ArrayType(x)
assert len(rank_zero_array) == 1
assert rank_zero_array.shape == ()
assert rank_zero_array.rank() == 0
assert rank_zero_array[()] == x
raises(ValueError, lambda: rank_zero_array[0])
def test_sympify():
from sympy.abc import x, y, z, t
arr = MutableDenseNDimArray([[x, y], [1, z*t]])
arr_other = sympify(arr)
assert arr_other.shape == (2, 2)
assert arr_other == arr
def test_reshape():
array = MutableDenseNDimArray(range(50), 50)
assert array.shape == (50,)
assert array.rank() == 1
array = array.reshape(5, 5, 2)
assert array.shape == (5, 5, 2)
assert array.rank() == 3
assert len(array) == 50
def test_iterator():
array = MutableDenseNDimArray(range(4), (2, 2))
assert array[0] == MutableDenseNDimArray([0, 1])
assert array[1] == MutableDenseNDimArray([2, 3])
array = array.reshape(4)
j = 0
for i in array:
assert i == j
j += 1
def test_getitem():
for ArrayType in [MutableDenseNDimArray, MutableSparseNDimArray]:
array = ArrayType(range(24)).reshape(2, 3, 4)
assert array.tolist() == [[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]]
assert array[0] == ArrayType([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]])
assert array[0, 0] == ArrayType([0, 1, 2, 3])
value = 0
for i in range(2):
for j in range(3):
for k in range(4):
assert array[i, j, k] == value
value += 1
raises(ValueError, lambda: array[3, 4, 5])
raises(ValueError, lambda: array[3, 4, 5, 6])
raises(ValueError, lambda: array[3, 4, 5, 3:4])
def test_sparse():
sparse_array = MutableSparseNDimArray([0, 0, 0, 1], (2, 2))
assert len(sparse_array) == 2 * 2
# dictionary where all data is, only non-zero entries are actually stored:
assert len(sparse_array._sparse_array) == 1
assert sparse_array.tolist() == [[0, 0], [0, 1]]
for i, j in zip(sparse_array, [[0, 0], [0, 1]]):
assert i == MutableSparseNDimArray(j)
sparse_array[0, 0] = 123
assert len(sparse_array._sparse_array) == 2
assert sparse_array[0, 0] == 123
assert sparse_array/0 == MutableSparseNDimArray([[S.ComplexInfinity, S.NaN], [S.NaN, S.ComplexInfinity]], (2, 2))
# when element in sparse array become zero it will disappear from
# dictionary
sparse_array[0, 0] = 0
assert len(sparse_array._sparse_array) == 1
sparse_array[1, 1] = 0
assert len(sparse_array._sparse_array) == 0
assert sparse_array[0, 0] == 0
# test for large scale sparse array
# equality test
a = MutableSparseNDimArray.zeros(100000, 200000)
b = MutableSparseNDimArray.zeros(100000, 200000)
assert a == b
a[1, 1] = 1
b[1, 1] = 2
assert a != b
# __mul__ and __rmul__
assert a * 3 == MutableSparseNDimArray({200001: 3}, (100000, 200000))
assert 3 * a == MutableSparseNDimArray({200001: 3}, (100000, 200000))
assert a * 0 == MutableSparseNDimArray({}, (100000, 200000))
assert 0 * a == MutableSparseNDimArray({}, (100000, 200000))
# __truediv__
assert a/3 == MutableSparseNDimArray({200001: Rational(1, 3)}, (100000, 200000))
# __neg__
assert -a == MutableSparseNDimArray({200001: -1}, (100000, 200000))
def test_calculation():
a = MutableDenseNDimArray([1]*9, (3, 3))
b = MutableDenseNDimArray([9]*9, (3, 3))
c = a + b
for i in c:
assert i == MutableDenseNDimArray([10, 10, 10])
assert c == MutableDenseNDimArray([10]*9, (3, 3))
assert c == MutableSparseNDimArray([10]*9, (3, 3))
c = b - a
for i in c:
assert i == MutableSparseNDimArray([8, 8, 8])
assert c == MutableDenseNDimArray([8]*9, (3, 3))
assert c == MutableSparseNDimArray([8]*9, (3, 3))
def test_ndim_array_converting():
dense_array = MutableDenseNDimArray([1, 2, 3, 4], (2, 2))
alist = dense_array.tolist()
assert alist == [[1, 2], [3, 4]]
matrix = dense_array.tomatrix()
assert (isinstance(matrix, Matrix))
for i in range(len(dense_array)):
assert dense_array[dense_array._get_tuple_index(i)] == matrix[i]
assert matrix.shape == dense_array.shape
assert MutableDenseNDimArray(matrix) == dense_array
assert MutableDenseNDimArray(matrix.as_immutable()) == dense_array
assert MutableDenseNDimArray(matrix.as_mutable()) == dense_array
sparse_array = MutableSparseNDimArray([1, 2, 3, 4], (2, 2))
alist = sparse_array.tolist()
assert alist == [[1, 2], [3, 4]]
matrix = sparse_array.tomatrix()
assert(isinstance(matrix, SparseMatrix))
for i in range(len(sparse_array)):
assert sparse_array[sparse_array._get_tuple_index(i)] == matrix[i]
assert matrix.shape == sparse_array.shape
assert MutableSparseNDimArray(matrix) == sparse_array
assert MutableSparseNDimArray(matrix.as_immutable()) == sparse_array
assert MutableSparseNDimArray(matrix.as_mutable()) == sparse_array
def test_converting_functions():
arr_list = [1, 2, 3, 4]
arr_matrix = Matrix(((1, 2), (3, 4)))
# list
arr_ndim_array = MutableDenseNDimArray(arr_list, (2, 2))
assert (isinstance(arr_ndim_array, MutableDenseNDimArray))
assert arr_matrix.tolist() == arr_ndim_array.tolist()
# Matrix
arr_ndim_array = MutableDenseNDimArray(arr_matrix)
assert (isinstance(arr_ndim_array, MutableDenseNDimArray))
assert arr_matrix.tolist() == arr_ndim_array.tolist()
assert arr_matrix.shape == arr_ndim_array.shape
def test_equality():
first_list = [1, 2, 3, 4]
second_list = [1, 2, 3, 4]
third_list = [4, 3, 2, 1]
assert first_list == second_list
assert first_list != third_list
first_ndim_array = MutableDenseNDimArray(first_list, (2, 2))
second_ndim_array = MutableDenseNDimArray(second_list, (2, 2))
third_ndim_array = MutableDenseNDimArray(third_list, (2, 2))
fourth_ndim_array = MutableDenseNDimArray(first_list, (2, 2))
assert first_ndim_array == second_ndim_array
second_ndim_array[0, 0] = 0
assert first_ndim_array != second_ndim_array
assert first_ndim_array != third_ndim_array
assert first_ndim_array == fourth_ndim_array
def test_arithmetic():
a = MutableDenseNDimArray([3 for i in range(9)], (3, 3))
b = MutableDenseNDimArray([7 for i in range(9)], (3, 3))
c1 = a + b
c2 = b + a
assert c1 == c2
d1 = a - b
d2 = b - a
assert d1 == d2 * (-1)
e1 = a * 5
e2 = 5 * a
e3 = copy(a)
e3 *= 5
assert e1 == e2 == e3
f1 = a / 5
f2 = copy(a)
f2 /= 5
assert f1 == f2
assert f1[0, 0] == f1[0, 1] == f1[0, 2] == f1[1, 0] == f1[1, 1] == \
f1[1, 2] == f1[2, 0] == f1[2, 1] == f1[2, 2] == Rational(3, 5)
assert type(a) == type(b) == type(c1) == type(c2) == type(d1) == type(d2) \
== type(e1) == type(e2) == type(e3) == type(f1)
z0 = -a
assert z0 == MutableDenseNDimArray([-3 for i in range(9)], (3, 3))
def test_higher_dimenions():
m3 = MutableDenseNDimArray(range(10, 34), (2, 3, 4))
assert m3.tolist() == [[[10, 11, 12, 13],
[14, 15, 16, 17],
[18, 19, 20, 21]],
[[22, 23, 24, 25],
[26, 27, 28, 29],
[30, 31, 32, 33]]]
assert m3._get_tuple_index(0) == (0, 0, 0)
assert m3._get_tuple_index(1) == (0, 0, 1)
assert m3._get_tuple_index(4) == (0, 1, 0)
assert m3._get_tuple_index(12) == (1, 0, 0)
assert str(m3) == '[[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]]'
m3_rebuilt = MutableDenseNDimArray([[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]])
assert m3 == m3_rebuilt
m3_other = MutableDenseNDimArray([[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]], (2, 3, 4))
assert m3 == m3_other
def test_slices():
md = MutableDenseNDimArray(range(10, 34), (2, 3, 4))
assert md[:] == MutableDenseNDimArray(range(10, 34), (2, 3, 4))
assert md[:, :, 0].tomatrix() == Matrix([[10, 14, 18], [22, 26, 30]])
assert md[0, 1:2, :].tomatrix() == Matrix([[14, 15, 16, 17]])
assert md[0, 1:3, :].tomatrix() == Matrix([[14, 15, 16, 17], [18, 19, 20, 21]])
assert md[:, :, :] == md
sd = MutableSparseNDimArray(range(10, 34), (2, 3, 4))
assert sd == MutableSparseNDimArray(md)
assert sd[:] == MutableSparseNDimArray(range(10, 34), (2, 3, 4))
assert sd[:, :, 0].tomatrix() == Matrix([[10, 14, 18], [22, 26, 30]])
assert sd[0, 1:2, :].tomatrix() == Matrix([[14, 15, 16, 17]])
assert sd[0, 1:3, :].tomatrix() == Matrix([[14, 15, 16, 17], [18, 19, 20, 21]])
assert sd[:, :, :] == sd
def test_slices_assign():
a = MutableDenseNDimArray(range(12), shape=(4, 3))
b = MutableSparseNDimArray(range(12), shape=(4, 3))
for i in [a, b]:
assert i.tolist() == [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]]
i[0, :] = [2, 2, 2]
assert i.tolist() == [[2, 2, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]]
i[0, 1:] = [8, 8]
assert i.tolist() == [[2, 8, 8], [3, 4, 5], [6, 7, 8], [9, 10, 11]]
i[1:3, 1] = [20, 44]
assert i.tolist() == [[2, 8, 8], [3, 20, 5], [6, 44, 8], [9, 10, 11]]
def test_diff():
from sympy.abc import x, y, z
md = MutableDenseNDimArray([[x, y], [x*z, x*y*z]])
assert md.diff(x) == MutableDenseNDimArray([[1, 0], [z, y*z]])
assert diff(md, x) == MutableDenseNDimArray([[1, 0], [z, y*z]])
sd = MutableSparseNDimArray(md)
assert sd == MutableSparseNDimArray([x, y, x*z, x*y*z], (2, 2))
assert sd.diff(x) == MutableSparseNDimArray([[1, 0], [z, y*z]])
assert diff(sd, x) == MutableSparseNDimArray([[1, 0], [z, y*z]])

View File

@ -0,0 +1,73 @@
from sympy.testing.pytest import raises
from sympy.functions.elementary.trigonometric import sin, cos
from sympy.matrices.dense import Matrix
from sympy.simplify import simplify
from sympy.tensor.array import Array
from sympy.tensor.array.dense_ndim_array import (
ImmutableDenseNDimArray, MutableDenseNDimArray)
from sympy.tensor.array.sparse_ndim_array import (
ImmutableSparseNDimArray, MutableSparseNDimArray)
from sympy.abc import x, y
mutable_array_types = [
MutableDenseNDimArray,
MutableSparseNDimArray
]
array_types = [
ImmutableDenseNDimArray,
ImmutableSparseNDimArray,
MutableDenseNDimArray,
MutableSparseNDimArray
]
def test_array_negative_indices():
for ArrayType in array_types:
test_array = ArrayType([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]])
assert test_array[:, -1] == Array([5, 10])
assert test_array[:, -2] == Array([4, 9])
assert test_array[:, -3] == Array([3, 8])
assert test_array[:, -4] == Array([2, 7])
assert test_array[:, -5] == Array([1, 6])
assert test_array[:, 0] == Array([1, 6])
assert test_array[:, 1] == Array([2, 7])
assert test_array[:, 2] == Array([3, 8])
assert test_array[:, 3] == Array([4, 9])
assert test_array[:, 4] == Array([5, 10])
raises(ValueError, lambda: test_array[:, -6])
raises(ValueError, lambda: test_array[-3, :])
assert test_array[-1, -1] == 10
def test_issue_18361():
A = Array([sin(2 * x) - 2 * sin(x) * cos(x)])
B = Array([sin(x)**2 + cos(x)**2, 0])
C = Array([(x + x**2)/(x*sin(y)**2 + x*cos(y)**2), 2*sin(x)*cos(x)])
assert simplify(A) == Array([0])
assert simplify(B) == Array([1, 0])
assert simplify(C) == Array([x + 1, sin(2*x)])
def test_issue_20222():
A = Array([[1, 2], [3, 4]])
B = Matrix([[1,2],[3,4]])
raises(TypeError, lambda: A - B)
def test_issue_17851():
for array_type in array_types:
A = array_type([])
assert isinstance(A, array_type)
assert A.shape == (0,)
assert list(A) == []
def test_issue_and_18715():
for array_type in mutable_array_types:
A = array_type([0, 1, 2])
A[0] += 5
assert A[0] == 5

View File

@ -0,0 +1,22 @@
from sympy.tensor.array import (ImmutableDenseNDimArray,
ImmutableSparseNDimArray, MutableDenseNDimArray, MutableSparseNDimArray)
from sympy.abc import x, y, z
def test_NDim_array_conv():
MD = MutableDenseNDimArray([x, y, z])
MS = MutableSparseNDimArray([x, y, z])
ID = ImmutableDenseNDimArray([x, y, z])
IS = ImmutableSparseNDimArray([x, y, z])
assert MD.as_immutable() == ID
assert MD.as_mutable() == MD
assert MS.as_immutable() == IS
assert MS.as_mutable() == MS
assert ID.as_immutable() == ID
assert ID.as_mutable() == MD
assert IS.as_immutable() == IS
assert IS.as_mutable() == MS