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,23 @@
"""A module to manipulate symbolic objects with indices including tensors
"""
from .indexed import IndexedBase, Idx, Indexed
from .index_methods import get_contraction_structure, get_indices
from .functions import shape
from .array import (MutableDenseNDimArray, ImmutableDenseNDimArray,
MutableSparseNDimArray, ImmutableSparseNDimArray, NDimArray, tensorproduct,
tensorcontraction, tensordiagonal, derive_by_array, permutedims, Array,
DenseNDimArray, SparseNDimArray,)
__all__ = [
'IndexedBase', 'Idx', 'Indexed',
'get_contraction_structure', 'get_indices',
'shape',
'MutableDenseNDimArray', 'ImmutableDenseNDimArray',
'MutableSparseNDimArray', 'ImmutableSparseNDimArray', 'NDimArray',
'tensorproduct', 'tensorcontraction', 'tensordiagonal', 'derive_by_array', 'permutedims',
'Array', 'DenseNDimArray', 'SparseNDimArray',
]

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

View File

@ -0,0 +1,154 @@
from collections.abc import Iterable
from functools import singledispatch
from sympy.core.expr import Expr
from sympy.core.mul import Mul
from sympy.core.singleton import S
from sympy.core.sympify import sympify
from sympy.core.parameters import global_parameters
class TensorProduct(Expr):
"""
Generic class for tensor products.
"""
is_number = False
def __new__(cls, *args, **kwargs):
from sympy.tensor.array import NDimArray, tensorproduct, Array
from sympy.matrices.expressions.matexpr import MatrixExpr
from sympy.matrices.matrixbase import MatrixBase
from sympy.strategies import flatten
args = [sympify(arg) for arg in args]
evaluate = kwargs.get("evaluate", global_parameters.evaluate)
if not evaluate:
obj = Expr.__new__(cls, *args)
return obj
arrays = []
other = []
scalar = S.One
for arg in args:
if isinstance(arg, (Iterable, MatrixBase, NDimArray)):
arrays.append(Array(arg))
elif isinstance(arg, (MatrixExpr,)):
other.append(arg)
else:
scalar *= arg
coeff = scalar*tensorproduct(*arrays)
if len(other) == 0:
return coeff
if coeff != 1:
newargs = [coeff] + other
else:
newargs = other
obj = Expr.__new__(cls, *newargs, **kwargs)
return flatten(obj)
def rank(self):
return len(self.shape)
def _get_args_shapes(self):
from sympy.tensor.array import Array
return [i.shape if hasattr(i, "shape") else Array(i).shape for i in self.args]
@property
def shape(self):
shape_list = self._get_args_shapes()
return sum(shape_list, ())
def __getitem__(self, index):
index = iter(index)
return Mul.fromiter(
arg.__getitem__(tuple(next(index) for i in shp))
for arg, shp in zip(self.args, self._get_args_shapes())
)
@singledispatch
def shape(expr):
"""
Return the shape of the *expr* as a tuple. *expr* should represent
suitable object such as matrix or array.
Parameters
==========
expr : SymPy object having ``MatrixKind`` or ``ArrayKind``.
Raises
======
NoShapeError : Raised when object with wrong kind is passed.
Examples
========
This function returns the shape of any object representing matrix or array.
>>> from sympy import shape, Array, ImmutableDenseMatrix, Integral
>>> from sympy.abc import x
>>> A = Array([1, 2])
>>> shape(A)
(2,)
>>> shape(Integral(A, x))
(2,)
>>> M = ImmutableDenseMatrix([1, 2])
>>> shape(M)
(2, 1)
>>> shape(Integral(M, x))
(2, 1)
You can support new type by dispatching.
>>> from sympy import Expr
>>> class NewExpr(Expr):
... pass
>>> @shape.register(NewExpr)
... def _(expr):
... return shape(expr.args[0])
>>> shape(NewExpr(M))
(2, 1)
If unsuitable expression is passed, ``NoShapeError()`` will be raised.
>>> shape(Integral(x, x))
Traceback (most recent call last):
...
sympy.tensor.functions.NoShapeError: shape() called on non-array object: Integral(x, x)
Notes
=====
Array-like classes (such as ``Matrix`` or ``NDimArray``) has ``shape``
property which returns its shape, but it cannot be used for non-array
classes containing array. This function returns the shape of any
registered object representing array.
"""
if hasattr(expr, "shape"):
return expr.shape
raise NoShapeError(
"%s does not have shape, or its type is not registered to shape()." % expr)
class NoShapeError(Exception):
"""
Raised when ``shape()`` is called on non-array object.
This error can be imported from ``sympy.tensor.functions``.
Examples
========
>>> from sympy import shape
>>> from sympy.abc import x
>>> shape(x)
Traceback (most recent call last):
...
sympy.tensor.functions.NoShapeError: shape() called on non-array object: x
"""
pass

View File

@ -0,0 +1,469 @@
"""Module with functions operating on IndexedBase, Indexed and Idx objects
- Check shape conformance
- Determine indices in resulting expression
etc.
Methods in this module could be implemented by calling methods on Expr
objects instead. When things stabilize this could be a useful
refactoring.
"""
from functools import reduce
from sympy.core.function import Function
from sympy.functions import exp, Piecewise
from sympy.tensor.indexed import Idx, Indexed
from sympy.utilities import sift
from collections import OrderedDict
class IndexConformanceException(Exception):
pass
def _unique_and_repeated(inds):
"""
Returns the unique and repeated indices. Also note, from the examples given below
that the order of indices is maintained as given in the input.
Examples
========
>>> from sympy.tensor.index_methods import _unique_and_repeated
>>> _unique_and_repeated([2, 3, 1, 3, 0, 4, 0])
([2, 1, 4], [3, 0])
"""
uniq = OrderedDict()
for i in inds:
if i in uniq:
uniq[i] = 0
else:
uniq[i] = 1
return sift(uniq, lambda x: uniq[x], binary=True)
def _remove_repeated(inds):
"""
Removes repeated objects from sequences
Returns a set of the unique objects and a tuple of all that have been
removed.
Examples
========
>>> from sympy.tensor.index_methods import _remove_repeated
>>> l1 = [1, 2, 3, 2]
>>> _remove_repeated(l1)
({1, 3}, (2,))
"""
u, r = _unique_and_repeated(inds)
return set(u), tuple(r)
def _get_indices_Mul(expr, return_dummies=False):
"""Determine the outer indices of a Mul object.
Examples
========
>>> from sympy.tensor.index_methods import _get_indices_Mul
>>> from sympy.tensor.indexed import IndexedBase, Idx
>>> i, j, k = map(Idx, ['i', 'j', 'k'])
>>> x = IndexedBase('x')
>>> y = IndexedBase('y')
>>> _get_indices_Mul(x[i, k]*y[j, k])
({i, j}, {})
>>> _get_indices_Mul(x[i, k]*y[j, k], return_dummies=True)
({i, j}, {}, (k,))
"""
inds = list(map(get_indices, expr.args))
inds, syms = list(zip(*inds))
inds = list(map(list, inds))
inds = list(reduce(lambda x, y: x + y, inds))
inds, dummies = _remove_repeated(inds)
symmetry = {}
for s in syms:
for pair in s:
if pair in symmetry:
symmetry[pair] *= s[pair]
else:
symmetry[pair] = s[pair]
if return_dummies:
return inds, symmetry, dummies
else:
return inds, symmetry
def _get_indices_Pow(expr):
"""Determine outer indices of a power or an exponential.
A power is considered a universal function, so that the indices of a Pow is
just the collection of indices present in the expression. This may be
viewed as a bit inconsistent in the special case:
x[i]**2 = x[i]*x[i] (1)
The above expression could have been interpreted as the contraction of x[i]
with itself, but we choose instead to interpret it as a function
lambda y: y**2
applied to each element of x (a universal function in numpy terms). In
order to allow an interpretation of (1) as a contraction, we need
contravariant and covariant Idx subclasses. (FIXME: this is not yet
implemented)
Expressions in the base or exponent are subject to contraction as usual,
but an index that is present in the exponent, will not be considered
contractable with its own base. Note however, that indices in the same
exponent can be contracted with each other.
Examples
========
>>> from sympy.tensor.index_methods import _get_indices_Pow
>>> from sympy import Pow, exp, IndexedBase, Idx
>>> A = IndexedBase('A')
>>> x = IndexedBase('x')
>>> i, j, k = map(Idx, ['i', 'j', 'k'])
>>> _get_indices_Pow(exp(A[i, j]*x[j]))
({i}, {})
>>> _get_indices_Pow(Pow(x[i], x[i]))
({i}, {})
>>> _get_indices_Pow(Pow(A[i, j]*x[j], x[i]))
({i}, {})
"""
base, exp = expr.as_base_exp()
binds, bsyms = get_indices(base)
einds, esyms = get_indices(exp)
inds = binds | einds
# FIXME: symmetries from power needs to check special cases, else nothing
symmetries = {}
return inds, symmetries
def _get_indices_Add(expr):
"""Determine outer indices of an Add object.
In a sum, each term must have the same set of outer indices. A valid
expression could be
x(i)*y(j) - x(j)*y(i)
But we do not allow expressions like:
x(i)*y(j) - z(j)*z(j)
FIXME: Add support for Numpy broadcasting
Examples
========
>>> from sympy.tensor.index_methods import _get_indices_Add
>>> from sympy.tensor.indexed import IndexedBase, Idx
>>> i, j, k = map(Idx, ['i', 'j', 'k'])
>>> x = IndexedBase('x')
>>> y = IndexedBase('y')
>>> _get_indices_Add(x[i] + x[k]*y[i, k])
({i}, {})
"""
inds = list(map(get_indices, expr.args))
inds, syms = list(zip(*inds))
# allow broadcast of scalars
non_scalars = [x for x in inds if x != set()]
if not non_scalars:
return set(), {}
if not all(x == non_scalars[0] for x in non_scalars[1:]):
raise IndexConformanceException("Indices are not consistent: %s" % expr)
if not reduce(lambda x, y: x != y or y, syms):
symmetries = syms[0]
else:
# FIXME: search for symmetries
symmetries = {}
return non_scalars[0], symmetries
def get_indices(expr):
"""Determine the outer indices of expression ``expr``
By *outer* we mean indices that are not summation indices. Returns a set
and a dict. The set contains outer indices and the dict contains
information about index symmetries.
Examples
========
>>> from sympy.tensor.index_methods import get_indices
>>> from sympy import symbols
>>> from sympy.tensor import IndexedBase
>>> x, y, A = map(IndexedBase, ['x', 'y', 'A'])
>>> i, j, a, z = symbols('i j a z', integer=True)
The indices of the total expression is determined, Repeated indices imply a
summation, for instance the trace of a matrix A:
>>> get_indices(A[i, i])
(set(), {})
In the case of many terms, the terms are required to have identical
outer indices. Else an IndexConformanceException is raised.
>>> get_indices(x[i] + A[i, j]*y[j])
({i}, {})
:Exceptions:
An IndexConformanceException means that the terms ar not compatible, e.g.
>>> get_indices(x[i] + y[j]) #doctest: +SKIP
(...)
IndexConformanceException: Indices are not consistent: x(i) + y(j)
.. warning::
The concept of *outer* indices applies recursively, starting on the deepest
level. This implies that dummies inside parenthesis are assumed to be
summed first, so that the following expression is handled gracefully:
>>> get_indices((x[i] + A[i, j]*y[j])*x[j])
({i, j}, {})
This is correct and may appear convenient, but you need to be careful
with this as SymPy will happily .expand() the product, if requested. The
resulting expression would mix the outer ``j`` with the dummies inside
the parenthesis, which makes it a different expression. To be on the
safe side, it is best to avoid such ambiguities by using unique indices
for all contractions that should be held separate.
"""
# We call ourself recursively to determine indices of sub expressions.
# break recursion
if isinstance(expr, Indexed):
c = expr.indices
inds, dummies = _remove_repeated(c)
return inds, {}
elif expr is None:
return set(), {}
elif isinstance(expr, Idx):
return {expr}, {}
elif expr.is_Atom:
return set(), {}
# recurse via specialized functions
else:
if expr.is_Mul:
return _get_indices_Mul(expr)
elif expr.is_Add:
return _get_indices_Add(expr)
elif expr.is_Pow or isinstance(expr, exp):
return _get_indices_Pow(expr)
elif isinstance(expr, Piecewise):
# FIXME: No support for Piecewise yet
return set(), {}
elif isinstance(expr, Function):
# Support ufunc like behaviour by returning indices from arguments.
# Functions do not interpret repeated indices across arguments
# as summation
ind0 = set()
for arg in expr.args:
ind, sym = get_indices(arg)
ind0 |= ind
return ind0, sym
# this test is expensive, so it should be at the end
elif not expr.has(Indexed):
return set(), {}
raise NotImplementedError(
"FIXME: No specialized handling of type %s" % type(expr))
def get_contraction_structure(expr):
"""Determine dummy indices of ``expr`` and describe its structure
By *dummy* we mean indices that are summation indices.
The structure of the expression is determined and described as follows:
1) A conforming summation of Indexed objects is described with a dict where
the keys are summation indices and the corresponding values are sets
containing all terms for which the summation applies. All Add objects
in the SymPy expression tree are described like this.
2) For all nodes in the SymPy expression tree that are *not* of type Add, the
following applies:
If a node discovers contractions in one of its arguments, the node
itself will be stored as a key in the dict. For that key, the
corresponding value is a list of dicts, each of which is the result of a
recursive call to get_contraction_structure(). The list contains only
dicts for the non-trivial deeper contractions, omitting dicts with None
as the one and only key.
.. Note:: The presence of expressions among the dictionary keys indicates
multiple levels of index contractions. A nested dict displays nested
contractions and may itself contain dicts from a deeper level. In
practical calculations the summation in the deepest nested level must be
calculated first so that the outer expression can access the resulting
indexed object.
Examples
========
>>> from sympy.tensor.index_methods import get_contraction_structure
>>> from sympy import default_sort_key
>>> from sympy.tensor import IndexedBase, Idx
>>> x, y, A = map(IndexedBase, ['x', 'y', 'A'])
>>> i, j, k, l = map(Idx, ['i', 'j', 'k', 'l'])
>>> get_contraction_structure(x[i]*y[i] + A[j, j])
{(i,): {x[i]*y[i]}, (j,): {A[j, j]}}
>>> get_contraction_structure(x[i]*y[j])
{None: {x[i]*y[j]}}
A multiplication of contracted factors results in nested dicts representing
the internal contractions.
>>> d = get_contraction_structure(x[i, i]*y[j, j])
>>> sorted(d.keys(), key=default_sort_key)
[None, x[i, i]*y[j, j]]
In this case, the product has no contractions:
>>> d[None]
{x[i, i]*y[j, j]}
Factors are contracted "first":
>>> sorted(d[x[i, i]*y[j, j]], key=default_sort_key)
[{(i,): {x[i, i]}}, {(j,): {y[j, j]}}]
A parenthesized Add object is also returned as a nested dictionary. The
term containing the parenthesis is a Mul with a contraction among the
arguments, so it will be found as a key in the result. It stores the
dictionary resulting from a recursive call on the Add expression.
>>> d = get_contraction_structure(x[i]*(y[i] + A[i, j]*x[j]))
>>> sorted(d.keys(), key=default_sort_key)
[(A[i, j]*x[j] + y[i])*x[i], (i,)]
>>> d[(i,)]
{(A[i, j]*x[j] + y[i])*x[i]}
>>> d[x[i]*(A[i, j]*x[j] + y[i])]
[{None: {y[i]}, (j,): {A[i, j]*x[j]}}]
Powers with contractions in either base or exponent will also be found as
keys in the dictionary, mapping to a list of results from recursive calls:
>>> d = get_contraction_structure(A[j, j]**A[i, i])
>>> d[None]
{A[j, j]**A[i, i]}
>>> nested_contractions = d[A[j, j]**A[i, i]]
>>> nested_contractions[0]
{(j,): {A[j, j]}}
>>> nested_contractions[1]
{(i,): {A[i, i]}}
The description of the contraction structure may appear complicated when
represented with a string in the above examples, but it is easy to iterate
over:
>>> from sympy import Expr
>>> for key in d:
... if isinstance(key, Expr):
... continue
... for term in d[key]:
... if term in d:
... # treat deepest contraction first
... pass
... # treat outermost contactions here
"""
# We call ourself recursively to inspect sub expressions.
if isinstance(expr, Indexed):
junk, key = _remove_repeated(expr.indices)
return {key or None: {expr}}
elif expr.is_Atom:
return {None: {expr}}
elif expr.is_Mul:
junk, junk, key = _get_indices_Mul(expr, return_dummies=True)
result = {key or None: {expr}}
# recurse on every factor
nested = []
for fac in expr.args:
facd = get_contraction_structure(fac)
if not (None in facd and len(facd) == 1):
nested.append(facd)
if nested:
result[expr] = nested
return result
elif expr.is_Pow or isinstance(expr, exp):
# recurse in base and exp separately. If either has internal
# contractions we must include ourselves as a key in the returned dict
b, e = expr.as_base_exp()
dbase = get_contraction_structure(b)
dexp = get_contraction_structure(e)
dicts = []
for d in dbase, dexp:
if not (None in d and len(d) == 1):
dicts.append(d)
result = {None: {expr}}
if dicts:
result[expr] = dicts
return result
elif expr.is_Add:
# Note: we just collect all terms with identical summation indices, We
# do nothing to identify equivalent terms here, as this would require
# substitutions or pattern matching in expressions of unknown
# complexity.
result = {}
for term in expr.args:
# recurse on every term
d = get_contraction_structure(term)
for key in d:
if key in result:
result[key] |= d[key]
else:
result[key] = d[key]
return result
elif isinstance(expr, Piecewise):
# FIXME: No support for Piecewise yet
return {None: expr}
elif isinstance(expr, Function):
# Collect non-trivial contraction structures in each argument
# We do not report repeated indices in separate arguments as a
# contraction
deeplist = []
for arg in expr.args:
deep = get_contraction_structure(arg)
if not (None in deep and len(deep) == 1):
deeplist.append(deep)
d = {None: {expr}}
if deeplist:
d[expr] = deeplist
return d
# this test is expensive, so it should be at the end
elif not expr.has(Indexed):
return {None: {expr}}
raise NotImplementedError(
"FIXME: No specialized handling of type %s" % type(expr))

View File

@ -0,0 +1,793 @@
r"""Module that defines indexed objects.
The classes ``IndexedBase``, ``Indexed``, and ``Idx`` represent a
matrix element ``M[i, j]`` as in the following diagram::
1) The Indexed class represents the entire indexed object.
|
___|___
' '
M[i, j]
/ \__\______
| |
| |
| 2) The Idx class represents indices; each Idx can
| optionally contain information about its range.
|
3) IndexedBase represents the 'stem' of an indexed object, here `M`.
The stem used by itself is usually taken to represent the entire
array.
There can be any number of indices on an Indexed object. No
transformation properties are implemented in these Base objects, but
implicit contraction of repeated indices is supported.
Note that the support for complicated (i.e. non-atomic) integer
expressions as indices is limited. (This should be improved in
future releases.)
Examples
========
To express the above matrix element example you would write:
>>> from sympy import symbols, IndexedBase, Idx
>>> M = IndexedBase('M')
>>> i, j = symbols('i j', cls=Idx)
>>> M[i, j]
M[i, j]
Repeated indices in a product implies a summation, so to express a
matrix-vector product in terms of Indexed objects:
>>> x = IndexedBase('x')
>>> M[i, j]*x[j]
M[i, j]*x[j]
If the indexed objects will be converted to component based arrays, e.g.
with the code printers or the autowrap framework, you also need to provide
(symbolic or numerical) dimensions. This can be done by passing an
optional shape parameter to IndexedBase upon construction:
>>> dim1, dim2 = symbols('dim1 dim2', integer=True)
>>> A = IndexedBase('A', shape=(dim1, 2*dim1, dim2))
>>> A.shape
(dim1, 2*dim1, dim2)
>>> A[i, j, 3].shape
(dim1, 2*dim1, dim2)
If an IndexedBase object has no shape information, it is assumed that the
array is as large as the ranges of its indices:
>>> n, m = symbols('n m', integer=True)
>>> i = Idx('i', m)
>>> j = Idx('j', n)
>>> M[i, j].shape
(m, n)
>>> M[i, j].ranges
[(0, m - 1), (0, n - 1)]
The above can be compared with the following:
>>> A[i, 2, j].shape
(dim1, 2*dim1, dim2)
>>> A[i, 2, j].ranges
[(0, m - 1), None, (0, n - 1)]
To analyze the structure of indexed expressions, you can use the methods
get_indices() and get_contraction_structure():
>>> from sympy.tensor import get_indices, get_contraction_structure
>>> get_indices(A[i, j, j])
({i}, {})
>>> get_contraction_structure(A[i, j, j])
{(j,): {A[i, j, j]}}
See the appropriate docstrings for a detailed explanation of the output.
"""
# TODO: (some ideas for improvement)
#
# o test and guarantee numpy compatibility
# - implement full support for broadcasting
# - strided arrays
#
# o more functions to analyze indexed expressions
# - identify standard constructs, e.g matrix-vector product in a subexpression
#
# o functions to generate component based arrays (numpy and sympy.Matrix)
# - generate a single array directly from Indexed
# - convert simple sub-expressions
#
# o sophisticated indexing (possibly in subclasses to preserve simplicity)
# - Idx with range smaller than dimension of Indexed
# - Idx with stepsize != 1
# - Idx with step determined by function call
from collections.abc import Iterable
from sympy.core.numbers import Number
from sympy.core.assumptions import StdFactKB
from sympy.core import Expr, Tuple, sympify, S
from sympy.core.symbol import _filter_assumptions, Symbol
from sympy.core.logic import fuzzy_bool, fuzzy_not
from sympy.core.sympify import _sympify
from sympy.functions.special.tensor_functions import KroneckerDelta
from sympy.multipledispatch import dispatch
from sympy.utilities.iterables import is_sequence, NotIterable
from sympy.utilities.misc import filldedent
class IndexException(Exception):
pass
class Indexed(Expr):
"""Represents a mathematical object with indices.
>>> from sympy import Indexed, IndexedBase, Idx, symbols
>>> i, j = symbols('i j', cls=Idx)
>>> Indexed('A', i, j)
A[i, j]
It is recommended that ``Indexed`` objects be created by indexing ``IndexedBase``:
``IndexedBase('A')[i, j]`` instead of ``Indexed(IndexedBase('A'), i, j)``.
>>> A = IndexedBase('A')
>>> a_ij = A[i, j] # Prefer this,
>>> b_ij = Indexed(A, i, j) # over this.
>>> a_ij == b_ij
True
"""
is_Indexed = True
is_symbol = True
is_Atom = True
def __new__(cls, base, *args, **kw_args):
from sympy.tensor.array.ndim_array import NDimArray
from sympy.matrices.matrixbase import MatrixBase
if not args:
raise IndexException("Indexed needs at least one index.")
if isinstance(base, (str, Symbol)):
base = IndexedBase(base)
elif not hasattr(base, '__getitem__') and not isinstance(base, IndexedBase):
raise TypeError(filldedent("""
The base can only be replaced with a string, Symbol,
IndexedBase or an object with a method for getting
items (i.e. an object with a `__getitem__` method).
"""))
args = list(map(sympify, args))
if isinstance(base, (NDimArray, Iterable, Tuple, MatrixBase)) and all(i.is_number for i in args):
if len(args) == 1:
return base[args[0]]
else:
return base[args]
base = _sympify(base)
obj = Expr.__new__(cls, base, *args, **kw_args)
IndexedBase._set_assumptions(obj, base.assumptions0)
return obj
def _hashable_content(self):
return super()._hashable_content() + tuple(sorted(self.assumptions0.items()))
@property
def name(self):
return str(self)
@property
def _diff_wrt(self):
"""Allow derivatives with respect to an ``Indexed`` object."""
return True
def _eval_derivative(self, wrt):
from sympy.tensor.array.ndim_array import NDimArray
if isinstance(wrt, Indexed) and wrt.base == self.base:
if len(self.indices) != len(wrt.indices):
msg = "Different # of indices: d({!s})/d({!s})".format(self,
wrt)
raise IndexException(msg)
result = S.One
for index1, index2 in zip(self.indices, wrt.indices):
result *= KroneckerDelta(index1, index2)
return result
elif isinstance(self.base, NDimArray):
from sympy.tensor.array import derive_by_array
return Indexed(derive_by_array(self.base, wrt), *self.args[1:])
else:
if Tuple(self.indices).has(wrt):
return S.NaN
return S.Zero
@property
def assumptions0(self):
return {k: v for k, v in self._assumptions.items() if v is not None}
@property
def base(self):
"""Returns the ``IndexedBase`` of the ``Indexed`` object.
Examples
========
>>> from sympy import Indexed, IndexedBase, Idx, symbols
>>> i, j = symbols('i j', cls=Idx)
>>> Indexed('A', i, j).base
A
>>> B = IndexedBase('B')
>>> B == B[i, j].base
True
"""
return self.args[0]
@property
def indices(self):
"""
Returns the indices of the ``Indexed`` object.
Examples
========
>>> from sympy import Indexed, Idx, symbols
>>> i, j = symbols('i j', cls=Idx)
>>> Indexed('A', i, j).indices
(i, j)
"""
return self.args[1:]
@property
def rank(self):
"""
Returns the rank of the ``Indexed`` object.
Examples
========
>>> from sympy import Indexed, Idx, symbols
>>> i, j, k, l, m = symbols('i:m', cls=Idx)
>>> Indexed('A', i, j).rank
2
>>> q = Indexed('A', i, j, k, l, m)
>>> q.rank
5
>>> q.rank == len(q.indices)
True
"""
return len(self.args) - 1
@property
def shape(self):
"""Returns a list with dimensions of each index.
Dimensions is a property of the array, not of the indices. Still, if
the ``IndexedBase`` does not define a shape attribute, it is assumed
that the ranges of the indices correspond to the shape of the array.
>>> from sympy import IndexedBase, Idx, symbols
>>> n, m = symbols('n m', integer=True)
>>> i = Idx('i', m)
>>> j = Idx('j', m)
>>> A = IndexedBase('A', shape=(n, n))
>>> B = IndexedBase('B')
>>> A[i, j].shape
(n, n)
>>> B[i, j].shape
(m, m)
"""
if self.base.shape:
return self.base.shape
sizes = []
for i in self.indices:
upper = getattr(i, 'upper', None)
lower = getattr(i, 'lower', None)
if None in (upper, lower):
raise IndexException(filldedent("""
Range is not defined for all indices in: %s""" % self))
try:
size = upper - lower + 1
except TypeError:
raise IndexException(filldedent("""
Shape cannot be inferred from Idx with
undefined range: %s""" % self))
sizes.append(size)
return Tuple(*sizes)
@property
def ranges(self):
"""Returns a list of tuples with lower and upper range of each index.
If an index does not define the data members upper and lower, the
corresponding slot in the list contains ``None`` instead of a tuple.
Examples
========
>>> from sympy import Indexed,Idx, symbols
>>> Indexed('A', Idx('i', 2), Idx('j', 4), Idx('k', 8)).ranges
[(0, 1), (0, 3), (0, 7)]
>>> Indexed('A', Idx('i', 3), Idx('j', 3), Idx('k', 3)).ranges
[(0, 2), (0, 2), (0, 2)]
>>> x, y, z = symbols('x y z', integer=True)
>>> Indexed('A', x, y, z).ranges
[None, None, None]
"""
ranges = []
sentinel = object()
for i in self.indices:
upper = getattr(i, 'upper', sentinel)
lower = getattr(i, 'lower', sentinel)
if sentinel not in (upper, lower):
ranges.append((lower, upper))
else:
ranges.append(None)
return ranges
def _sympystr(self, p):
indices = list(map(p.doprint, self.indices))
return "%s[%s]" % (p.doprint(self.base), ", ".join(indices))
@property
def free_symbols(self):
base_free_symbols = self.base.free_symbols
indices_free_symbols = {
fs for i in self.indices for fs in i.free_symbols}
if base_free_symbols:
return {self} | base_free_symbols | indices_free_symbols
else:
return indices_free_symbols
@property
def expr_free_symbols(self):
from sympy.utilities.exceptions import sympy_deprecation_warning
sympy_deprecation_warning("""
The expr_free_symbols property is deprecated. Use free_symbols to get
the free symbols of an expression.
""",
deprecated_since_version="1.9",
active_deprecations_target="deprecated-expr-free-symbols")
return {self}
class IndexedBase(Expr, NotIterable):
"""Represent the base or stem of an indexed object
The IndexedBase class represent an array that contains elements. The main purpose
of this class is to allow the convenient creation of objects of the Indexed
class. The __getitem__ method of IndexedBase returns an instance of
Indexed. Alone, without indices, the IndexedBase class can be used as a
notation for e.g. matrix equations, resembling what you could do with the
Symbol class. But, the IndexedBase class adds functionality that is not
available for Symbol instances:
- An IndexedBase object can optionally store shape information. This can
be used in to check array conformance and conditions for numpy
broadcasting. (TODO)
- An IndexedBase object implements syntactic sugar that allows easy symbolic
representation of array operations, using implicit summation of
repeated indices.
- The IndexedBase object symbolizes a mathematical structure equivalent
to arrays, and is recognized as such for code generation and automatic
compilation and wrapping.
>>> from sympy.tensor import IndexedBase, Idx
>>> from sympy import symbols
>>> A = IndexedBase('A'); A
A
>>> type(A)
<class 'sympy.tensor.indexed.IndexedBase'>
When an IndexedBase object receives indices, it returns an array with named
axes, represented by an Indexed object:
>>> i, j = symbols('i j', integer=True)
>>> A[i, j, 2]
A[i, j, 2]
>>> type(A[i, j, 2])
<class 'sympy.tensor.indexed.Indexed'>
The IndexedBase constructor takes an optional shape argument. If given,
it overrides any shape information in the indices. (But not the index
ranges!)
>>> m, n, o, p = symbols('m n o p', integer=True)
>>> i = Idx('i', m)
>>> j = Idx('j', n)
>>> A[i, j].shape
(m, n)
>>> B = IndexedBase('B', shape=(o, p))
>>> B[i, j].shape
(o, p)
Assumptions can be specified with keyword arguments the same way as for Symbol:
>>> A_real = IndexedBase('A', real=True)
>>> A_real.is_real
True
>>> A != A_real
True
Assumptions can also be inherited if a Symbol is used to initialize the IndexedBase:
>>> I = symbols('I', integer=True)
>>> C_inherit = IndexedBase(I)
>>> C_explicit = IndexedBase('I', integer=True)
>>> C_inherit == C_explicit
True
"""
is_symbol = True
is_Atom = True
@staticmethod
def _set_assumptions(obj, assumptions):
"""Set assumptions on obj, making sure to apply consistent values."""
tmp_asm_copy = assumptions.copy()
is_commutative = fuzzy_bool(assumptions.get('commutative', True))
assumptions['commutative'] = is_commutative
obj._assumptions = StdFactKB(assumptions)
obj._assumptions._generator = tmp_asm_copy # Issue #8873
def __new__(cls, label, shape=None, *, offset=S.Zero, strides=None, **kw_args):
from sympy.matrices.matrixbase import MatrixBase
from sympy.tensor.array.ndim_array import NDimArray
assumptions, kw_args = _filter_assumptions(kw_args)
if isinstance(label, str):
label = Symbol(label, **assumptions)
elif isinstance(label, Symbol):
assumptions = label._merge(assumptions)
elif isinstance(label, (MatrixBase, NDimArray)):
return label
elif isinstance(label, Iterable):
return _sympify(label)
else:
label = _sympify(label)
if is_sequence(shape):
shape = Tuple(*shape)
elif shape is not None:
shape = Tuple(shape)
if shape is not None:
obj = Expr.__new__(cls, label, shape)
else:
obj = Expr.__new__(cls, label)
obj._shape = shape
obj._offset = offset
obj._strides = strides
obj._name = str(label)
IndexedBase._set_assumptions(obj, assumptions)
return obj
@property
def name(self):
return self._name
def _hashable_content(self):
return super()._hashable_content() + tuple(sorted(self.assumptions0.items()))
@property
def assumptions0(self):
return {k: v for k, v in self._assumptions.items() if v is not None}
def __getitem__(self, indices, **kw_args):
if is_sequence(indices):
# Special case needed because M[*my_tuple] is a syntax error.
if self.shape and len(self.shape) != len(indices):
raise IndexException("Rank mismatch.")
return Indexed(self, *indices, **kw_args)
else:
if self.shape and len(self.shape) != 1:
raise IndexException("Rank mismatch.")
return Indexed(self, indices, **kw_args)
@property
def shape(self):
"""Returns the shape of the ``IndexedBase`` object.
Examples
========
>>> from sympy import IndexedBase, Idx
>>> from sympy.abc import x, y
>>> IndexedBase('A', shape=(x, y)).shape
(x, y)
Note: If the shape of the ``IndexedBase`` is specified, it will override
any shape information given by the indices.
>>> A = IndexedBase('A', shape=(x, y))
>>> B = IndexedBase('B')
>>> i = Idx('i', 2)
>>> j = Idx('j', 1)
>>> A[i, j].shape
(x, y)
>>> B[i, j].shape
(2, 1)
"""
return self._shape
@property
def strides(self):
"""Returns the strided scheme for the ``IndexedBase`` object.
Normally this is a tuple denoting the number of
steps to take in the respective dimension when traversing
an array. For code generation purposes strides='C' and
strides='F' can also be used.
strides='C' would mean that code printer would unroll
in row-major order and 'F' means unroll in column major
order.
"""
return self._strides
@property
def offset(self):
"""Returns the offset for the ``IndexedBase`` object.
This is the value added to the resulting index when the
2D Indexed object is unrolled to a 1D form. Used in code
generation.
Examples
==========
>>> from sympy.printing import ccode
>>> from sympy.tensor import IndexedBase, Idx
>>> from sympy import symbols
>>> l, m, n, o = symbols('l m n o', integer=True)
>>> A = IndexedBase('A', strides=(l, m, n), offset=o)
>>> i, j, k = map(Idx, 'ijk')
>>> ccode(A[i, j, k])
'A[l*i + m*j + n*k + o]'
"""
return self._offset
@property
def label(self):
"""Returns the label of the ``IndexedBase`` object.
Examples
========
>>> from sympy import IndexedBase
>>> from sympy.abc import x, y
>>> IndexedBase('A', shape=(x, y)).label
A
"""
return self.args[0]
def _sympystr(self, p):
return p.doprint(self.label)
class Idx(Expr):
"""Represents an integer index as an ``Integer`` or integer expression.
There are a number of ways to create an ``Idx`` object. The constructor
takes two arguments:
``label``
An integer or a symbol that labels the index.
``range``
Optionally you can specify a range as either
* ``Symbol`` or integer: This is interpreted as a dimension. Lower and
upper bounds are set to ``0`` and ``range - 1``, respectively.
* ``tuple``: The two elements are interpreted as the lower and upper
bounds of the range, respectively.
Note: bounds of the range are assumed to be either integer or infinite (oo
and -oo are allowed to specify an unbounded range). If ``n`` is given as a
bound, then ``n.is_integer`` must not return false.
For convenience, if the label is given as a string it is automatically
converted to an integer symbol. (Note: this conversion is not done for
range or dimension arguments.)
Examples
========
>>> from sympy import Idx, symbols, oo
>>> n, i, L, U = symbols('n i L U', integer=True)
If a string is given for the label an integer ``Symbol`` is created and the
bounds are both ``None``:
>>> idx = Idx('qwerty'); idx
qwerty
>>> idx.lower, idx.upper
(None, None)
Both upper and lower bounds can be specified:
>>> idx = Idx(i, (L, U)); idx
i
>>> idx.lower, idx.upper
(L, U)
When only a single bound is given it is interpreted as the dimension
and the lower bound defaults to 0:
>>> idx = Idx(i, n); idx.lower, idx.upper
(0, n - 1)
>>> idx = Idx(i, 4); idx.lower, idx.upper
(0, 3)
>>> idx = Idx(i, oo); idx.lower, idx.upper
(0, oo)
"""
is_integer = True
is_finite = True
is_real = True
is_symbol = True
is_Atom = True
_diff_wrt = True
def __new__(cls, label, range=None, **kw_args):
if isinstance(label, str):
label = Symbol(label, integer=True)
label, range = list(map(sympify, (label, range)))
if label.is_Number:
if not label.is_integer:
raise TypeError("Index is not an integer number.")
return label
if not label.is_integer:
raise TypeError("Idx object requires an integer label.")
elif is_sequence(range):
if len(range) != 2:
raise ValueError(filldedent("""
Idx range tuple must have length 2, but got %s""" % len(range)))
for bound in range:
if (bound.is_integer is False and bound is not S.Infinity
and bound is not S.NegativeInfinity):
raise TypeError("Idx object requires integer bounds.")
args = label, Tuple(*range)
elif isinstance(range, Expr):
if range is not S.Infinity and fuzzy_not(range.is_integer):
raise TypeError("Idx object requires an integer dimension.")
args = label, Tuple(0, range - 1)
elif range:
raise TypeError(filldedent("""
The range must be an ordered iterable or
integer SymPy expression."""))
else:
args = label,
obj = Expr.__new__(cls, *args, **kw_args)
obj._assumptions["finite"] = True
obj._assumptions["real"] = True
return obj
@property
def label(self):
"""Returns the label (Integer or integer expression) of the Idx object.
Examples
========
>>> from sympy import Idx, Symbol
>>> x = Symbol('x', integer=True)
>>> Idx(x).label
x
>>> j = Symbol('j', integer=True)
>>> Idx(j).label
j
>>> Idx(j + 1).label
j + 1
"""
return self.args[0]
@property
def lower(self):
"""Returns the lower bound of the ``Idx``.
Examples
========
>>> from sympy import Idx
>>> Idx('j', 2).lower
0
>>> Idx('j', 5).lower
0
>>> Idx('j').lower is None
True
"""
try:
return self.args[1][0]
except IndexError:
return
@property
def upper(self):
"""Returns the upper bound of the ``Idx``.
Examples
========
>>> from sympy import Idx
>>> Idx('j', 2).upper
1
>>> Idx('j', 5).upper
4
>>> Idx('j').upper is None
True
"""
try:
return self.args[1][1]
except IndexError:
return
def _sympystr(self, p):
return p.doprint(self.label)
@property
def name(self):
return self.label.name if self.label.is_Symbol else str(self.label)
@property
def free_symbols(self):
return {self}
@dispatch(Idx, Idx)
def _eval_is_ge(lhs, rhs): # noqa:F811
other_upper = rhs if rhs.upper is None else rhs.upper
other_lower = rhs if rhs.lower is None else rhs.lower
if lhs.lower is not None and (lhs.lower >= other_upper) == True:
return True
if lhs.upper is not None and (lhs.upper < other_lower) == True:
return False
return None
@dispatch(Idx, Number) # type:ignore
def _eval_is_ge(lhs, rhs): # noqa:F811
other_upper = rhs
other_lower = rhs
if lhs.lower is not None and (lhs.lower >= other_upper) == True:
return True
if lhs.upper is not None and (lhs.upper < other_lower) == True:
return False
return None
@dispatch(Number, Idx) # type:ignore
def _eval_is_ge(lhs, rhs): # noqa:F811
other_upper = lhs
other_lower = lhs
if rhs.upper is not None and (rhs.upper <= other_lower) == True:
return True
if rhs.lower is not None and (rhs.lower > other_upper) == True:
return False
return None

File diff suppressed because it is too large Load Diff

View File

@ -0,0 +1,57 @@
from sympy.tensor.functions import TensorProduct
from sympy.matrices.dense import Matrix
from sympy.matrices.expressions.matexpr import MatrixSymbol
from sympy.tensor.array import Array
from sympy.abc import x, y, z
from sympy.abc import i, j, k, l
A = MatrixSymbol("A", 3, 3)
B = MatrixSymbol("B", 3, 3)
C = MatrixSymbol("C", 3, 3)
def test_TensorProduct_construction():
assert TensorProduct(3, 4) == 12
assert isinstance(TensorProduct(A, A), TensorProduct)
expr = TensorProduct(TensorProduct(x, y), z)
assert expr == x*y*z
expr = TensorProduct(TensorProduct(A, B), C)
assert expr == TensorProduct(A, B, C)
expr = TensorProduct(Matrix.eye(2), Array([[0, -1], [1, 0]]))
assert expr == Array([
[
[[0, -1], [1, 0]],
[[0, 0], [0, 0]]
],
[
[[0, 0], [0, 0]],
[[0, -1], [1, 0]]
]
])
def test_TensorProduct_shape():
expr = TensorProduct(3, 4, evaluate=False)
assert expr.shape == ()
assert expr.rank() == 0
expr = TensorProduct(Array([1, 2]), Array([x, y]), evaluate=False)
assert expr.shape == (2, 2)
assert expr.rank() == 2
expr = TensorProduct(expr, expr, evaluate=False)
assert expr.shape == (2, 2, 2, 2)
assert expr.rank() == 4
expr = TensorProduct(Matrix.eye(2), Array([[0, -1], [1, 0]]), evaluate=False)
assert expr.shape == (2, 2, 2, 2)
assert expr.rank() == 4
def test_TensorProduct_getitem():
expr = TensorProduct(A, B)
assert expr[i, j, k, l] == A[i, j]*B[k, l]

View File

@ -0,0 +1,227 @@
from sympy.core import symbols, S, Pow, Function
from sympy.functions import exp
from sympy.testing.pytest import raises
from sympy.tensor.indexed import Idx, IndexedBase
from sympy.tensor.index_methods import IndexConformanceException
from sympy.tensor.index_methods import (get_contraction_structure, get_indices)
def test_trivial_indices():
x, y = symbols('x y')
assert get_indices(x) == (set(), {})
assert get_indices(x*y) == (set(), {})
assert get_indices(x + y) == (set(), {})
assert get_indices(x**y) == (set(), {})
def test_get_indices_Indexed():
x = IndexedBase('x')
i, j = Idx('i'), Idx('j')
assert get_indices(x[i, j]) == ({i, j}, {})
assert get_indices(x[j, i]) == ({j, i}, {})
def test_get_indices_Idx():
f = Function('f')
i, j = Idx('i'), Idx('j')
assert get_indices(f(i)*j) == ({i, j}, {})
assert get_indices(f(j, i)) == ({j, i}, {})
assert get_indices(f(i)*i) == (set(), {})
def test_get_indices_mul():
x = IndexedBase('x')
y = IndexedBase('y')
i, j = Idx('i'), Idx('j')
assert get_indices(x[j]*y[i]) == ({i, j}, {})
assert get_indices(x[i]*y[j]) == ({i, j}, {})
def test_get_indices_exceptions():
x = IndexedBase('x')
y = IndexedBase('y')
i, j = Idx('i'), Idx('j')
raises(IndexConformanceException, lambda: get_indices(x[i] + y[j]))
def test_scalar_broadcast():
x = IndexedBase('x')
y = IndexedBase('y')
i, j = Idx('i'), Idx('j')
assert get_indices(x[i] + y[i, i]) == ({i}, {})
assert get_indices(x[i] + y[j, j]) == ({i}, {})
def test_get_indices_add():
x = IndexedBase('x')
y = IndexedBase('y')
A = IndexedBase('A')
i, j, k = Idx('i'), Idx('j'), Idx('k')
assert get_indices(x[i] + 2*y[i]) == ({i}, {})
assert get_indices(y[i] + 2*A[i, j]*x[j]) == ({i}, {})
assert get_indices(y[i] + 2*(x[i] + A[i, j]*x[j])) == ({i}, {})
assert get_indices(y[i] + x[i]*(A[j, j] + 1)) == ({i}, {})
assert get_indices(
y[i] + x[i]*x[j]*(y[j] + A[j, k]*x[k])) == ({i}, {})
def test_get_indices_Pow():
x = IndexedBase('x')
y = IndexedBase('y')
A = IndexedBase('A')
i, j, k = Idx('i'), Idx('j'), Idx('k')
assert get_indices(Pow(x[i], y[j])) == ({i, j}, {})
assert get_indices(Pow(x[i, k], y[j, k])) == ({i, j, k}, {})
assert get_indices(Pow(A[i, k], y[k] + A[k, j]*x[j])) == ({i, k}, {})
assert get_indices(Pow(2, x[i])) == get_indices(exp(x[i]))
# test of a design decision, this may change:
assert get_indices(Pow(x[i], 2)) == ({i}, {})
def test_get_contraction_structure_basic():
x = IndexedBase('x')
y = IndexedBase('y')
i, j = Idx('i'), Idx('j')
assert get_contraction_structure(x[i]*y[j]) == {None: {x[i]*y[j]}}
assert get_contraction_structure(x[i] + y[j]) == {None: {x[i], y[j]}}
assert get_contraction_structure(x[i]*y[i]) == {(i,): {x[i]*y[i]}}
assert get_contraction_structure(
1 + x[i]*y[i]) == {None: {S.One}, (i,): {x[i]*y[i]}}
assert get_contraction_structure(x[i]**y[i]) == {None: {x[i]**y[i]}}
def test_get_contraction_structure_complex():
x = IndexedBase('x')
y = IndexedBase('y')
A = IndexedBase('A')
i, j, k = Idx('i'), Idx('j'), Idx('k')
expr1 = y[i] + A[i, j]*x[j]
d1 = {None: {y[i]}, (j,): {A[i, j]*x[j]}}
assert get_contraction_structure(expr1) == d1
expr2 = expr1*A[k, i] + x[k]
d2 = {None: {x[k]}, (i,): {expr1*A[k, i]}, expr1*A[k, i]: [d1]}
assert get_contraction_structure(expr2) == d2
def test_contraction_structure_simple_Pow():
x = IndexedBase('x')
y = IndexedBase('y')
i, j, k = Idx('i'), Idx('j'), Idx('k')
ii_jj = x[i, i]**y[j, j]
assert get_contraction_structure(ii_jj) == {
None: {ii_jj},
ii_jj: [
{(i,): {x[i, i]}},
{(j,): {y[j, j]}}
]
}
ii_jk = x[i, i]**y[j, k]
assert get_contraction_structure(ii_jk) == {
None: {x[i, i]**y[j, k]},
x[i, i]**y[j, k]: [
{(i,): {x[i, i]}}
]
}
def test_contraction_structure_Mul_and_Pow():
x = IndexedBase('x')
y = IndexedBase('y')
i, j, k = Idx('i'), Idx('j'), Idx('k')
i_ji = x[i]**(y[j]*x[i])
assert get_contraction_structure(i_ji) == {None: {i_ji}}
ij_i = (x[i]*y[j])**(y[i])
assert get_contraction_structure(ij_i) == {None: {ij_i}}
j_ij_i = x[j]*(x[i]*y[j])**(y[i])
assert get_contraction_structure(j_ij_i) == {(j,): {j_ij_i}}
j_i_ji = x[j]*x[i]**(y[j]*x[i])
assert get_contraction_structure(j_i_ji) == {(j,): {j_i_ji}}
ij_exp_kki = x[i]*y[j]*exp(y[i]*y[k, k])
result = get_contraction_structure(ij_exp_kki)
expected = {
(i,): {ij_exp_kki},
ij_exp_kki: [{
None: {exp(y[i]*y[k, k])},
exp(y[i]*y[k, k]): [{
None: {y[i]*y[k, k]},
y[i]*y[k, k]: [{(k,): {y[k, k]}}]
}]}
]
}
assert result == expected
def test_contraction_structure_Add_in_Pow():
x = IndexedBase('x')
y = IndexedBase('y')
i, j, k = Idx('i'), Idx('j'), Idx('k')
s_ii_jj_s = (1 + x[i, i])**(1 + y[j, j])
expected = {
None: {s_ii_jj_s},
s_ii_jj_s: [
{None: {S.One}, (i,): {x[i, i]}},
{None: {S.One}, (j,): {y[j, j]}}
]
}
result = get_contraction_structure(s_ii_jj_s)
assert result == expected
s_ii_jk_s = (1 + x[i, i]) ** (1 + y[j, k])
expected_2 = {
None: {(x[i, i] + 1)**(y[j, k] + 1)},
s_ii_jk_s: [
{None: {S.One}, (i,): {x[i, i]}}
]
}
result_2 = get_contraction_structure(s_ii_jk_s)
assert result_2 == expected_2
def test_contraction_structure_Pow_in_Pow():
x = IndexedBase('x')
y = IndexedBase('y')
z = IndexedBase('z')
i, j, k = Idx('i'), Idx('j'), Idx('k')
ii_jj_kk = x[i, i]**y[j, j]**z[k, k]
expected = {
None: {ii_jj_kk},
ii_jj_kk: [
{(i,): {x[i, i]}},
{
None: {y[j, j]**z[k, k]},
y[j, j]**z[k, k]: [
{(j,): {y[j, j]}},
{(k,): {z[k, k]}}
]
}
]
}
assert get_contraction_structure(ii_jj_kk) == expected
def test_ufunc_support():
f = Function('f')
g = Function('g')
x = IndexedBase('x')
y = IndexedBase('y')
i, j = Idx('i'), Idx('j')
a = symbols('a')
assert get_indices(f(x[i])) == ({i}, {})
assert get_indices(f(x[i], y[j])) == ({i, j}, {})
assert get_indices(f(y[i])*g(x[i])) == (set(), {})
assert get_indices(f(a, x[i])) == ({i}, {})
assert get_indices(f(a, y[i], x[j])*g(x[i])) == ({j}, {})
assert get_indices(g(f(x[i]))) == ({i}, {})
assert get_contraction_structure(f(x[i])) == {None: {f(x[i])}}
assert get_contraction_structure(
f(y[i])*g(x[i])) == {(i,): {f(y[i])*g(x[i])}}
assert get_contraction_structure(
f(y[i])*g(f(x[i]))) == {(i,): {f(y[i])*g(f(x[i]))}}
assert get_contraction_structure(
f(x[j], y[i])*g(x[i])) == {(i,): {f(x[j], y[i])*g(x[i])}}

View File

@ -0,0 +1,511 @@
from sympy.core import symbols, Symbol, Tuple, oo, Dummy
from sympy.tensor.indexed import IndexException
from sympy.testing.pytest import raises
from sympy.utilities.iterables import iterable
# import test:
from sympy.concrete.summations import Sum
from sympy.core.function import Function, Subs, Derivative
from sympy.core.relational import (StrictLessThan, GreaterThan,
StrictGreaterThan, LessThan)
from sympy.core.singleton import S
from sympy.functions.elementary.exponential import exp, log
from sympy.functions.elementary.trigonometric import cos, sin
from sympy.functions.special.tensor_functions import KroneckerDelta
from sympy.series.order import Order
from sympy.sets.fancysets import Range
from sympy.tensor.indexed import IndexedBase, Idx, Indexed
def test_Idx_construction():
i, a, b = symbols('i a b', integer=True)
assert Idx(i) != Idx(i, 1)
assert Idx(i, a) == Idx(i, (0, a - 1))
assert Idx(i, oo) == Idx(i, (0, oo))
x = symbols('x', integer=False)
raises(TypeError, lambda: Idx(x))
raises(TypeError, lambda: Idx(0.5))
raises(TypeError, lambda: Idx(i, x))
raises(TypeError, lambda: Idx(i, 0.5))
raises(TypeError, lambda: Idx(i, (x, 5)))
raises(TypeError, lambda: Idx(i, (2, x)))
raises(TypeError, lambda: Idx(i, (2, 3.5)))
def test_Idx_properties():
i, a, b = symbols('i a b', integer=True)
assert Idx(i).is_integer
assert Idx(i).name == 'i'
assert Idx(i + 2).name == 'i + 2'
assert Idx('foo').name == 'foo'
def test_Idx_bounds():
i, a, b = symbols('i a b', integer=True)
assert Idx(i).lower is None
assert Idx(i).upper is None
assert Idx(i, a).lower == 0
assert Idx(i, a).upper == a - 1
assert Idx(i, 5).lower == 0
assert Idx(i, 5).upper == 4
assert Idx(i, oo).lower == 0
assert Idx(i, oo).upper is oo
assert Idx(i, (a, b)).lower == a
assert Idx(i, (a, b)).upper == b
assert Idx(i, (1, 5)).lower == 1
assert Idx(i, (1, 5)).upper == 5
assert Idx(i, (-oo, oo)).lower is -oo
assert Idx(i, (-oo, oo)).upper is oo
def test_Idx_fixed_bounds():
i, a, b, x = symbols('i a b x', integer=True)
assert Idx(x).lower is None
assert Idx(x).upper is None
assert Idx(x, a).lower == 0
assert Idx(x, a).upper == a - 1
assert Idx(x, 5).lower == 0
assert Idx(x, 5).upper == 4
assert Idx(x, oo).lower == 0
assert Idx(x, oo).upper is oo
assert Idx(x, (a, b)).lower == a
assert Idx(x, (a, b)).upper == b
assert Idx(x, (1, 5)).lower == 1
assert Idx(x, (1, 5)).upper == 5
assert Idx(x, (-oo, oo)).lower is -oo
assert Idx(x, (-oo, oo)).upper is oo
def test_Idx_inequalities():
i14 = Idx("i14", (1, 4))
i79 = Idx("i79", (7, 9))
i46 = Idx("i46", (4, 6))
i35 = Idx("i35", (3, 5))
assert i14 <= 5
assert i14 < 5
assert not (i14 >= 5)
assert not (i14 > 5)
assert 5 >= i14
assert 5 > i14
assert not (5 <= i14)
assert not (5 < i14)
assert LessThan(i14, 5)
assert StrictLessThan(i14, 5)
assert not GreaterThan(i14, 5)
assert not StrictGreaterThan(i14, 5)
assert i14 <= 4
assert isinstance(i14 < 4, StrictLessThan)
assert isinstance(i14 >= 4, GreaterThan)
assert not (i14 > 4)
assert isinstance(i14 <= 1, LessThan)
assert not (i14 < 1)
assert i14 >= 1
assert isinstance(i14 > 1, StrictGreaterThan)
assert not (i14 <= 0)
assert not (i14 < 0)
assert i14 >= 0
assert i14 > 0
from sympy.abc import x
assert isinstance(i14 < x, StrictLessThan)
assert isinstance(i14 > x, StrictGreaterThan)
assert isinstance(i14 <= x, LessThan)
assert isinstance(i14 >= x, GreaterThan)
assert i14 < i79
assert i14 <= i79
assert not (i14 > i79)
assert not (i14 >= i79)
assert i14 <= i46
assert isinstance(i14 < i46, StrictLessThan)
assert isinstance(i14 >= i46, GreaterThan)
assert not (i14 > i46)
assert isinstance(i14 < i35, StrictLessThan)
assert isinstance(i14 > i35, StrictGreaterThan)
assert isinstance(i14 <= i35, LessThan)
assert isinstance(i14 >= i35, GreaterThan)
iNone1 = Idx("iNone1")
iNone2 = Idx("iNone2")
assert isinstance(iNone1 < iNone2, StrictLessThan)
assert isinstance(iNone1 > iNone2, StrictGreaterThan)
assert isinstance(iNone1 <= iNone2, LessThan)
assert isinstance(iNone1 >= iNone2, GreaterThan)
def test_Idx_inequalities_current_fails():
i14 = Idx("i14", (1, 4))
assert S(5) >= i14
assert S(5) > i14
assert not (S(5) <= i14)
assert not (S(5) < i14)
def test_Idx_func_args():
i, a, b = symbols('i a b', integer=True)
ii = Idx(i)
assert ii.func(*ii.args) == ii
ii = Idx(i, a)
assert ii.func(*ii.args) == ii
ii = Idx(i, (a, b))
assert ii.func(*ii.args) == ii
def test_Idx_subs():
i, a, b = symbols('i a b', integer=True)
assert Idx(i, a).subs(a, b) == Idx(i, b)
assert Idx(i, a).subs(i, b) == Idx(b, a)
assert Idx(i).subs(i, 2) == Idx(2)
assert Idx(i, a).subs(a, 2) == Idx(i, 2)
assert Idx(i, (a, b)).subs(i, 2) == Idx(2, (a, b))
def test_IndexedBase_sugar():
i, j = symbols('i j', integer=True)
a = symbols('a')
A1 = Indexed(a, i, j)
A2 = IndexedBase(a)
assert A1 == A2[i, j]
assert A1 == A2[(i, j)]
assert A1 == A2[[i, j]]
assert A1 == A2[Tuple(i, j)]
assert all(a.is_Integer for a in A2[1, 0].args[1:])
def test_IndexedBase_subs():
i = symbols('i', integer=True)
a, b = symbols('a b')
A = IndexedBase(a)
B = IndexedBase(b)
assert A[i] == B[i].subs(b, a)
C = {1: 2}
assert C[1] == A[1].subs(A, C)
def test_IndexedBase_shape():
i, j, m, n = symbols('i j m n', integer=True)
a = IndexedBase('a', shape=(m, m))
b = IndexedBase('a', shape=(m, n))
assert b.shape == Tuple(m, n)
assert a[i, j] != b[i, j]
assert a[i, j] == b[i, j].subs(n, m)
assert b.func(*b.args) == b
assert b[i, j].func(*b[i, j].args) == b[i, j]
raises(IndexException, lambda: b[i])
raises(IndexException, lambda: b[i, i, j])
F = IndexedBase("F", shape=m)
assert F.shape == Tuple(m)
assert F[i].subs(i, j) == F[j]
raises(IndexException, lambda: F[i, j])
def test_IndexedBase_assumptions():
i = Symbol('i', integer=True)
a = Symbol('a')
A = IndexedBase(a, positive=True)
for c in (A, A[i]):
assert c.is_real
assert c.is_complex
assert not c.is_imaginary
assert c.is_nonnegative
assert c.is_nonzero
assert c.is_commutative
assert log(exp(c)) == c
assert A != IndexedBase(a)
assert A == IndexedBase(a, positive=True, real=True)
assert A[i] != Indexed(a, i)
def test_IndexedBase_assumptions_inheritance():
I = Symbol('I', integer=True)
I_inherit = IndexedBase(I)
I_explicit = IndexedBase('I', integer=True)
assert I_inherit.is_integer
assert I_explicit.is_integer
assert I_inherit.label.is_integer
assert I_explicit.label.is_integer
assert I_inherit == I_explicit
def test_issue_17652():
"""Regression test issue #17652.
IndexedBase.label should not upcast subclasses of Symbol
"""
class SubClass(Symbol):
pass
x = SubClass('X')
assert type(x) == SubClass
base = IndexedBase(x)
assert type(x) == SubClass
assert type(base.label) == SubClass
def test_Indexed_constructor():
i, j = symbols('i j', integer=True)
A = Indexed('A', i, j)
assert A == Indexed(Symbol('A'), i, j)
assert A == Indexed(IndexedBase('A'), i, j)
raises(TypeError, lambda: Indexed(A, i, j))
raises(IndexException, lambda: Indexed("A"))
assert A.free_symbols == {A, A.base.label, i, j}
def test_Indexed_func_args():
i, j = symbols('i j', integer=True)
a = symbols('a')
A = Indexed(a, i, j)
assert A == A.func(*A.args)
def test_Indexed_subs():
i, j, k = symbols('i j k', integer=True)
a, b = symbols('a b')
A = IndexedBase(a)
B = IndexedBase(b)
assert A[i, j] == B[i, j].subs(b, a)
assert A[i, j] == A[i, k].subs(k, j)
def test_Indexed_properties():
i, j = symbols('i j', integer=True)
A = Indexed('A', i, j)
assert A.name == 'A[i, j]'
assert A.rank == 2
assert A.indices == (i, j)
assert A.base == IndexedBase('A')
assert A.ranges == [None, None]
raises(IndexException, lambda: A.shape)
n, m = symbols('n m', integer=True)
assert Indexed('A', Idx(
i, m), Idx(j, n)).ranges == [Tuple(0, m - 1), Tuple(0, n - 1)]
assert Indexed('A', Idx(i, m), Idx(j, n)).shape == Tuple(m, n)
raises(IndexException, lambda: Indexed("A", Idx(i, m), Idx(j)).shape)
def test_Indexed_shape_precedence():
i, j = symbols('i j', integer=True)
o, p = symbols('o p', integer=True)
n, m = symbols('n m', integer=True)
a = IndexedBase('a', shape=(o, p))
assert a.shape == Tuple(o, p)
assert Indexed(
a, Idx(i, m), Idx(j, n)).ranges == [Tuple(0, m - 1), Tuple(0, n - 1)]
assert Indexed(a, Idx(i, m), Idx(j, n)).shape == Tuple(o, p)
assert Indexed(
a, Idx(i, m), Idx(j)).ranges == [Tuple(0, m - 1), (None, None)]
assert Indexed(a, Idx(i, m), Idx(j)).shape == Tuple(o, p)
def test_complex_indices():
i, j = symbols('i j', integer=True)
A = Indexed('A', i, i + j)
assert A.rank == 2
assert A.indices == (i, i + j)
def test_not_interable():
i, j = symbols('i j', integer=True)
A = Indexed('A', i, i + j)
assert not iterable(A)
def test_Indexed_coeff():
N = Symbol('N', integer=True)
len_y = N
i = Idx('i', len_y-1)
y = IndexedBase('y', shape=(len_y,))
a = (1/y[i+1]*y[i]).coeff(y[i])
b = (y[i]/y[i+1]).coeff(y[i])
assert a == b
def test_differentiation():
from sympy.functions.special.tensor_functions import KroneckerDelta
i, j, k, l = symbols('i j k l', cls=Idx)
a = symbols('a')
m, n = symbols("m, n", integer=True, finite=True)
assert m.is_real
h, L = symbols('h L', cls=IndexedBase)
hi, hj = h[i], h[j]
expr = hi
assert expr.diff(hj) == KroneckerDelta(i, j)
assert expr.diff(hi) == KroneckerDelta(i, i)
expr = S(2) * hi
assert expr.diff(hj) == S(2) * KroneckerDelta(i, j)
assert expr.diff(hi) == S(2) * KroneckerDelta(i, i)
assert expr.diff(a) is S.Zero
assert Sum(expr, (i, -oo, oo)).diff(hj) == Sum(2*KroneckerDelta(i, j), (i, -oo, oo))
assert Sum(expr.diff(hj), (i, -oo, oo)) == Sum(2*KroneckerDelta(i, j), (i, -oo, oo))
assert Sum(expr, (i, -oo, oo)).diff(hj).doit() == 2
assert Sum(expr.diff(hi), (i, -oo, oo)).doit() == Sum(2, (i, -oo, oo)).doit()
assert Sum(expr, (i, -oo, oo)).diff(hi).doit() is oo
expr = a * hj * hj / S(2)
assert expr.diff(hi) == a * h[j] * KroneckerDelta(i, j)
assert expr.diff(a) == hj * hj / S(2)
assert expr.diff(a, 2) is S.Zero
assert Sum(expr, (i, -oo, oo)).diff(hi) == Sum(a*KroneckerDelta(i, j)*h[j], (i, -oo, oo))
assert Sum(expr.diff(hi), (i, -oo, oo)) == Sum(a*KroneckerDelta(i, j)*h[j], (i, -oo, oo))
assert Sum(expr, (i, -oo, oo)).diff(hi).doit() == a*h[j]
assert Sum(expr, (j, -oo, oo)).diff(hi) == Sum(a*KroneckerDelta(i, j)*h[j], (j, -oo, oo))
assert Sum(expr.diff(hi), (j, -oo, oo)) == Sum(a*KroneckerDelta(i, j)*h[j], (j, -oo, oo))
assert Sum(expr, (j, -oo, oo)).diff(hi).doit() == a*h[i]
expr = a * sin(hj * hj)
assert expr.diff(hi) == 2*a*cos(hj * hj) * hj * KroneckerDelta(i, j)
assert expr.diff(hj) == 2*a*cos(hj * hj) * hj
expr = a * L[i, j] * h[j]
assert expr.diff(hi) == a*L[i, j]*KroneckerDelta(i, j)
assert expr.diff(hj) == a*L[i, j]
assert expr.diff(L[i, j]) == a*h[j]
assert expr.diff(L[k, l]) == a*KroneckerDelta(i, k)*KroneckerDelta(j, l)*h[j]
assert expr.diff(L[i, l]) == a*KroneckerDelta(j, l)*h[j]
assert Sum(expr, (j, -oo, oo)).diff(L[k, l]) == Sum(a * KroneckerDelta(i, k) * KroneckerDelta(j, l) * h[j], (j, -oo, oo))
assert Sum(expr, (j, -oo, oo)).diff(L[k, l]).doit() == a * KroneckerDelta(i, k) * h[l]
assert h[m].diff(h[m]) == 1
assert h[m].diff(h[n]) == KroneckerDelta(m, n)
assert Sum(a*h[m], (m, -oo, oo)).diff(h[n]) == Sum(a*KroneckerDelta(m, n), (m, -oo, oo))
assert Sum(a*h[m], (m, -oo, oo)).diff(h[n]).doit() == a
assert Sum(a*h[m], (n, -oo, oo)).diff(h[n]) == Sum(a*KroneckerDelta(m, n), (n, -oo, oo))
assert Sum(a*h[m], (m, -oo, oo)).diff(h[m]).doit() == oo*a
def test_indexed_series():
A = IndexedBase("A")
i = symbols("i", integer=True)
assert sin(A[i]).series(A[i]) == A[i] - A[i]**3/6 + A[i]**5/120 + Order(A[i]**6, A[i])
def test_indexed_is_constant():
A = IndexedBase("A")
i, j, k = symbols("i,j,k")
assert not A[i].is_constant()
assert A[i].is_constant(j)
assert not A[1+2*i, k].is_constant()
assert not A[1+2*i, k].is_constant(i)
assert A[1+2*i, k].is_constant(j)
assert not A[1+2*i, k].is_constant(k)
def test_issue_12533():
d = IndexedBase('d')
assert IndexedBase(range(5)) == Range(0, 5, 1)
assert d[0].subs(Symbol("d"), range(5)) == 0
assert d[0].subs(d, range(5)) == 0
assert d[1].subs(d, range(5)) == 1
assert Indexed(Range(5), 2) == 2
def test_issue_12780():
n = symbols("n")
i = Idx("i", (0, n))
raises(TypeError, lambda: i.subs(n, 1.5))
def test_issue_18604():
m = symbols("m")
assert Idx("i", m).name == 'i'
assert Idx("i", m).lower == 0
assert Idx("i", m).upper == m - 1
m = symbols("m", real=False)
raises(TypeError, lambda: Idx("i", m))
def test_Subs_with_Indexed():
A = IndexedBase("A")
i, j, k = symbols("i,j,k")
x, y, z = symbols("x,y,z")
f = Function("f")
assert Subs(A[i], A[i], A[j]).diff(A[j]) == 1
assert Subs(A[i], A[i], x).diff(A[i]) == 0
assert Subs(A[i], A[i], x).diff(A[j]) == 0
assert Subs(A[i], A[i], x).diff(x) == 1
assert Subs(A[i], A[i], x).diff(y) == 0
assert Subs(A[i], A[i], A[j]).diff(A[k]) == KroneckerDelta(j, k)
assert Subs(x, x, A[i]).diff(A[j]) == KroneckerDelta(i, j)
assert Subs(f(A[i]), A[i], x).diff(A[j]) == 0
assert Subs(f(A[i]), A[i], A[k]).diff(A[j]) == Derivative(f(A[k]), A[k])*KroneckerDelta(j, k)
assert Subs(x, x, A[i]**2).diff(A[j]) == 2*KroneckerDelta(i, j)*A[i]
assert Subs(A[i], A[i], A[j]**2).diff(A[k]) == 2*KroneckerDelta(j, k)*A[j]
assert Subs(A[i]*x, x, A[i]).diff(A[i]) == 2*A[i]
assert Subs(A[i]*x, x, A[i]).diff(A[j]) == 2*A[i]*KroneckerDelta(i, j)
assert Subs(A[i]*x, x, A[j]).diff(A[i]) == A[j] + A[i]*KroneckerDelta(i, j)
assert Subs(A[i]*x, x, A[j]).diff(A[j]) == A[i] + A[j]*KroneckerDelta(i, j)
assert Subs(A[i]*x, x, A[i]).diff(A[k]) == 2*A[i]*KroneckerDelta(i, k)
assert Subs(A[i]*x, x, A[j]).diff(A[k]) == KroneckerDelta(i, k)*A[j] + KroneckerDelta(j, k)*A[i]
assert Subs(A[i]*x, A[i], x).diff(A[i]) == 0
assert Subs(A[i]*x, A[i], x).diff(A[j]) == 0
assert Subs(A[i]*x, A[j], x).diff(A[i]) == x
assert Subs(A[i]*x, A[j], x).diff(A[j]) == x*KroneckerDelta(i, j)
assert Subs(A[i]*x, A[i], x).diff(A[k]) == 0
assert Subs(A[i]*x, A[j], x).diff(A[k]) == x*KroneckerDelta(i, k)
def test_complicated_derivative_with_Indexed():
x, y = symbols("x,y", cls=IndexedBase)
sigma = symbols("sigma")
i, j, k = symbols("i,j,k")
m0,m1,m2,m3,m4,m5 = symbols("m0:6")
f = Function("f")
expr = f((x[i] - y[i])**2/sigma)
_xi_1 = symbols("xi_1", cls=Dummy)
assert expr.diff(x[m0]).dummy_eq(
(x[i] - y[i])*KroneckerDelta(i, m0)*\
2*Subs(
Derivative(f(_xi_1), _xi_1),
(_xi_1,),
((x[i] - y[i])**2/sigma,)
)/sigma
)
assert expr.diff(x[m0]).diff(x[m1]).dummy_eq(
2*KroneckerDelta(i, m0)*\
KroneckerDelta(i, m1)*Subs(
Derivative(f(_xi_1), _xi_1),
(_xi_1,),
((x[i] - y[i])**2/sigma,)
)/sigma + \
4*(x[i] - y[i])**2*KroneckerDelta(i, m0)*KroneckerDelta(i, m1)*\
Subs(
Derivative(f(_xi_1), _xi_1, _xi_1),
(_xi_1,),
((x[i] - y[i])**2/sigma,)
)/sigma**2
)
def test_IndexedBase_commutative():
t = IndexedBase('t', commutative=False)
u = IndexedBase('u', commutative=False)
v = IndexedBase('v')
assert t[0]*v[0] == v[0]*t[0]
assert t[0]*u[0] != u[0]*t[0]

View File

@ -0,0 +1,13 @@
from sympy.tensor.tensor import TensorIndexType, tensor_indices, TensorHead
from sympy import I
def test_printing_TensMul():
R3 = TensorIndexType('R3', dim=3)
p, q = tensor_indices("p q", R3)
K = TensorHead("K", [R3])
assert repr(2*K(p)) == "2*K(p)"
assert repr(-K(p)) == "-K(p)"
assert repr(-2*K(p)*K(q)) == "-2*K(p)*K(q)"
assert repr(-I*K(p)) == "-I*K(p)"
assert repr(I*K(p)) == "I*K(p)"

File diff suppressed because it is too large Load Diff

View File

@ -0,0 +1,31 @@
from sympy.tensor.tensor import (Tensor, TensorIndexType, TensorSymmetry,
tensor_indices, TensorHead, TensorElement)
from sympy.tensor import Array
from sympy.core.symbol import Symbol
def test_tensor_element():
L = TensorIndexType("L")
i, j, k, l, m, n = tensor_indices("i j k l m n", L)
A = TensorHead("A", [L, L], TensorSymmetry.no_symmetry(2))
a = A(i, j)
assert isinstance(TensorElement(a, {}), Tensor)
assert isinstance(TensorElement(a, {k: 1}), Tensor)
te1 = TensorElement(a, {Symbol("i"): 1})
assert te1.free == [(j, 0)]
assert te1.get_free_indices() == [j]
assert te1.dum == []
te2 = TensorElement(a, {i: 1})
assert te2.free == [(j, 0)]
assert te2.get_free_indices() == [j]
assert te2.dum == []
assert te1 == te2
array = Array([[1, 2], [3, 4]])
assert te1.replace_with_arrays({A(i, j): array}, [j]) == array[1, :]

Some files were not shown because too many files have changed in this diff Show More