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