switching to high quality piper tts and added label translations
This commit is contained in:
Matthias Hinrichs
@@ -0,0 +1,23 @@
|
||||
"""A module to manipulate symbolic objects with indices including tensors
|
||||
|
||||
"""
|
||||
from .indexed import IndexedBase, Idx, Indexed
|
||||
from .index_methods import get_contraction_structure, get_indices
|
||||
from .functions import shape
|
||||
from .array import (MutableDenseNDimArray, ImmutableDenseNDimArray,
|
||||
MutableSparseNDimArray, ImmutableSparseNDimArray, NDimArray, tensorproduct,
|
||||
tensorcontraction, tensordiagonal, derive_by_array, permutedims, Array,
|
||||
DenseNDimArray, SparseNDimArray,)
|
||||
|
||||
__all__ = [
|
||||
'IndexedBase', 'Idx', 'Indexed',
|
||||
|
||||
'get_contraction_structure', 'get_indices',
|
||||
|
||||
'shape',
|
||||
|
||||
'MutableDenseNDimArray', 'ImmutableDenseNDimArray',
|
||||
'MutableSparseNDimArray', 'ImmutableSparseNDimArray', 'NDimArray',
|
||||
'tensorproduct', 'tensorcontraction', 'tensordiagonal', 'derive_by_array', 'permutedims',
|
||||
'Array', 'DenseNDimArray', 'SparseNDimArray',
|
||||
]
|
||||
@@ -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
|
||||
@@ -0,0 +1,154 @@
|
||||
from collections.abc import Iterable
|
||||
from functools import singledispatch
|
||||
|
||||
from sympy.core.expr import Expr
|
||||
from sympy.core.mul import Mul
|
||||
from sympy.core.singleton import S
|
||||
from sympy.core.sympify import sympify
|
||||
from sympy.core.parameters import global_parameters
|
||||
|
||||
|
||||
class TensorProduct(Expr):
|
||||
"""
|
||||
Generic class for tensor products.
|
||||
"""
|
||||
is_number = False
|
||||
|
||||
def __new__(cls, *args, **kwargs):
|
||||
from sympy.tensor.array import NDimArray, tensorproduct, Array
|
||||
from sympy.matrices.expressions.matexpr import MatrixExpr
|
||||
from sympy.matrices.matrixbase import MatrixBase
|
||||
from sympy.strategies import flatten
|
||||
|
||||
args = [sympify(arg) for arg in args]
|
||||
evaluate = kwargs.get("evaluate", global_parameters.evaluate)
|
||||
|
||||
if not evaluate:
|
||||
obj = Expr.__new__(cls, *args)
|
||||
return obj
|
||||
|
||||
arrays = []
|
||||
other = []
|
||||
scalar = S.One
|
||||
for arg in args:
|
||||
if isinstance(arg, (Iterable, MatrixBase, NDimArray)):
|
||||
arrays.append(Array(arg))
|
||||
elif isinstance(arg, (MatrixExpr,)):
|
||||
other.append(arg)
|
||||
else:
|
||||
scalar *= arg
|
||||
|
||||
coeff = scalar*tensorproduct(*arrays)
|
||||
if len(other) == 0:
|
||||
return coeff
|
||||
if coeff != 1:
|
||||
newargs = [coeff] + other
|
||||
else:
|
||||
newargs = other
|
||||
obj = Expr.__new__(cls, *newargs, **kwargs)
|
||||
return flatten(obj)
|
||||
|
||||
def rank(self):
|
||||
return len(self.shape)
|
||||
|
||||
def _get_args_shapes(self):
|
||||
from sympy.tensor.array import Array
|
||||
return [i.shape if hasattr(i, "shape") else Array(i).shape for i in self.args]
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
shape_list = self._get_args_shapes()
|
||||
return sum(shape_list, ())
|
||||
|
||||
def __getitem__(self, index):
|
||||
index = iter(index)
|
||||
return Mul.fromiter(
|
||||
arg.__getitem__(tuple(next(index) for i in shp))
|
||||
for arg, shp in zip(self.args, self._get_args_shapes())
|
||||
)
|
||||
|
||||
|
||||
@singledispatch
|
||||
def shape(expr):
|
||||
"""
|
||||
Return the shape of the *expr* as a tuple. *expr* should represent
|
||||
suitable object such as matrix or array.
|
||||
|
||||
Parameters
|
||||
==========
|
||||
|
||||
expr : SymPy object having ``MatrixKind`` or ``ArrayKind``.
|
||||
|
||||
Raises
|
||||
======
|
||||
|
||||
NoShapeError : Raised when object with wrong kind is passed.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
This function returns the shape of any object representing matrix or array.
|
||||
|
||||
>>> from sympy import shape, Array, ImmutableDenseMatrix, Integral
|
||||
>>> from sympy.abc import x
|
||||
>>> A = Array([1, 2])
|
||||
>>> shape(A)
|
||||
(2,)
|
||||
>>> shape(Integral(A, x))
|
||||
(2,)
|
||||
>>> M = ImmutableDenseMatrix([1, 2])
|
||||
>>> shape(M)
|
||||
(2, 1)
|
||||
>>> shape(Integral(M, x))
|
||||
(2, 1)
|
||||
|
||||
You can support new type by dispatching.
|
||||
|
||||
>>> from sympy import Expr
|
||||
>>> class NewExpr(Expr):
|
||||
... pass
|
||||
>>> @shape.register(NewExpr)
|
||||
... def _(expr):
|
||||
... return shape(expr.args[0])
|
||||
>>> shape(NewExpr(M))
|
||||
(2, 1)
|
||||
|
||||
If unsuitable expression is passed, ``NoShapeError()`` will be raised.
|
||||
|
||||
>>> shape(Integral(x, x))
|
||||
Traceback (most recent call last):
|
||||
...
|
||||
sympy.tensor.functions.NoShapeError: shape() called on non-array object: Integral(x, x)
|
||||
|
||||
Notes
|
||||
=====
|
||||
|
||||
Array-like classes (such as ``Matrix`` or ``NDimArray``) has ``shape``
|
||||
property which returns its shape, but it cannot be used for non-array
|
||||
classes containing array. This function returns the shape of any
|
||||
registered object representing array.
|
||||
|
||||
"""
|
||||
if hasattr(expr, "shape"):
|
||||
return expr.shape
|
||||
raise NoShapeError(
|
||||
"%s does not have shape, or its type is not registered to shape()." % expr)
|
||||
|
||||
|
||||
class NoShapeError(Exception):
|
||||
"""
|
||||
Raised when ``shape()`` is called on non-array object.
|
||||
|
||||
This error can be imported from ``sympy.tensor.functions``.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy import shape
|
||||
>>> from sympy.abc import x
|
||||
>>> shape(x)
|
||||
Traceback (most recent call last):
|
||||
...
|
||||
sympy.tensor.functions.NoShapeError: shape() called on non-array object: x
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,469 @@
|
||||
"""Module with functions operating on IndexedBase, Indexed and Idx objects
|
||||
|
||||
- Check shape conformance
|
||||
- Determine indices in resulting expression
|
||||
|
||||
etc.
|
||||
|
||||
Methods in this module could be implemented by calling methods on Expr
|
||||
objects instead. When things stabilize this could be a useful
|
||||
refactoring.
|
||||
"""
|
||||
|
||||
from functools import reduce
|
||||
|
||||
from sympy.core.function import Function
|
||||
from sympy.functions import exp, Piecewise
|
||||
from sympy.tensor.indexed import Idx, Indexed
|
||||
from sympy.utilities import sift
|
||||
|
||||
from collections import OrderedDict
|
||||
|
||||
class IndexConformanceException(Exception):
|
||||
pass
|
||||
|
||||
def _unique_and_repeated(inds):
|
||||
"""
|
||||
Returns the unique and repeated indices. Also note, from the examples given below
|
||||
that the order of indices is maintained as given in the input.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy.tensor.index_methods import _unique_and_repeated
|
||||
>>> _unique_and_repeated([2, 3, 1, 3, 0, 4, 0])
|
||||
([2, 1, 4], [3, 0])
|
||||
"""
|
||||
uniq = OrderedDict()
|
||||
for i in inds:
|
||||
if i in uniq:
|
||||
uniq[i] = 0
|
||||
else:
|
||||
uniq[i] = 1
|
||||
return sift(uniq, lambda x: uniq[x], binary=True)
|
||||
|
||||
def _remove_repeated(inds):
|
||||
"""
|
||||
Removes repeated objects from sequences
|
||||
|
||||
Returns a set of the unique objects and a tuple of all that have been
|
||||
removed.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy.tensor.index_methods import _remove_repeated
|
||||
>>> l1 = [1, 2, 3, 2]
|
||||
>>> _remove_repeated(l1)
|
||||
({1, 3}, (2,))
|
||||
|
||||
"""
|
||||
u, r = _unique_and_repeated(inds)
|
||||
return set(u), tuple(r)
|
||||
|
||||
|
||||
def _get_indices_Mul(expr, return_dummies=False):
|
||||
"""Determine the outer indices of a Mul object.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy.tensor.index_methods import _get_indices_Mul
|
||||
>>> from sympy.tensor.indexed import IndexedBase, Idx
|
||||
>>> i, j, k = map(Idx, ['i', 'j', 'k'])
|
||||
>>> x = IndexedBase('x')
|
||||
>>> y = IndexedBase('y')
|
||||
>>> _get_indices_Mul(x[i, k]*y[j, k])
|
||||
({i, j}, {})
|
||||
>>> _get_indices_Mul(x[i, k]*y[j, k], return_dummies=True)
|
||||
({i, j}, {}, (k,))
|
||||
|
||||
"""
|
||||
|
||||
inds = list(map(get_indices, expr.args))
|
||||
inds, syms = list(zip(*inds))
|
||||
|
||||
inds = list(map(list, inds))
|
||||
inds = list(reduce(lambda x, y: x + y, inds))
|
||||
inds, dummies = _remove_repeated(inds)
|
||||
|
||||
symmetry = {}
|
||||
for s in syms:
|
||||
for pair in s:
|
||||
if pair in symmetry:
|
||||
symmetry[pair] *= s[pair]
|
||||
else:
|
||||
symmetry[pair] = s[pair]
|
||||
|
||||
if return_dummies:
|
||||
return inds, symmetry, dummies
|
||||
else:
|
||||
return inds, symmetry
|
||||
|
||||
|
||||
def _get_indices_Pow(expr):
|
||||
"""Determine outer indices of a power or an exponential.
|
||||
|
||||
A power is considered a universal function, so that the indices of a Pow is
|
||||
just the collection of indices present in the expression. This may be
|
||||
viewed as a bit inconsistent in the special case:
|
||||
|
||||
x[i]**2 = x[i]*x[i] (1)
|
||||
|
||||
The above expression could have been interpreted as the contraction of x[i]
|
||||
with itself, but we choose instead to interpret it as a function
|
||||
|
||||
lambda y: y**2
|
||||
|
||||
applied to each element of x (a universal function in numpy terms). In
|
||||
order to allow an interpretation of (1) as a contraction, we need
|
||||
contravariant and covariant Idx subclasses. (FIXME: this is not yet
|
||||
implemented)
|
||||
|
||||
Expressions in the base or exponent are subject to contraction as usual,
|
||||
but an index that is present in the exponent, will not be considered
|
||||
contractable with its own base. Note however, that indices in the same
|
||||
exponent can be contracted with each other.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy.tensor.index_methods import _get_indices_Pow
|
||||
>>> from sympy import Pow, exp, IndexedBase, Idx
|
||||
>>> A = IndexedBase('A')
|
||||
>>> x = IndexedBase('x')
|
||||
>>> i, j, k = map(Idx, ['i', 'j', 'k'])
|
||||
>>> _get_indices_Pow(exp(A[i, j]*x[j]))
|
||||
({i}, {})
|
||||
>>> _get_indices_Pow(Pow(x[i], x[i]))
|
||||
({i}, {})
|
||||
>>> _get_indices_Pow(Pow(A[i, j]*x[j], x[i]))
|
||||
({i}, {})
|
||||
|
||||
"""
|
||||
base, exp = expr.as_base_exp()
|
||||
binds, bsyms = get_indices(base)
|
||||
einds, esyms = get_indices(exp)
|
||||
|
||||
inds = binds | einds
|
||||
|
||||
# FIXME: symmetries from power needs to check special cases, else nothing
|
||||
symmetries = {}
|
||||
|
||||
return inds, symmetries
|
||||
|
||||
|
||||
def _get_indices_Add(expr):
|
||||
"""Determine outer indices of an Add object.
|
||||
|
||||
In a sum, each term must have the same set of outer indices. A valid
|
||||
expression could be
|
||||
|
||||
x(i)*y(j) - x(j)*y(i)
|
||||
|
||||
But we do not allow expressions like:
|
||||
|
||||
x(i)*y(j) - z(j)*z(j)
|
||||
|
||||
FIXME: Add support for Numpy broadcasting
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy.tensor.index_methods import _get_indices_Add
|
||||
>>> from sympy.tensor.indexed import IndexedBase, Idx
|
||||
>>> i, j, k = map(Idx, ['i', 'j', 'k'])
|
||||
>>> x = IndexedBase('x')
|
||||
>>> y = IndexedBase('y')
|
||||
>>> _get_indices_Add(x[i] + x[k]*y[i, k])
|
||||
({i}, {})
|
||||
|
||||
"""
|
||||
|
||||
inds = list(map(get_indices, expr.args))
|
||||
inds, syms = list(zip(*inds))
|
||||
|
||||
# allow broadcast of scalars
|
||||
non_scalars = [x for x in inds if x != set()]
|
||||
if not non_scalars:
|
||||
return set(), {}
|
||||
|
||||
if not all(x == non_scalars[0] for x in non_scalars[1:]):
|
||||
raise IndexConformanceException("Indices are not consistent: %s" % expr)
|
||||
if not reduce(lambda x, y: x != y or y, syms):
|
||||
symmetries = syms[0]
|
||||
else:
|
||||
# FIXME: search for symmetries
|
||||
symmetries = {}
|
||||
|
||||
return non_scalars[0], symmetries
|
||||
|
||||
|
||||
def get_indices(expr):
|
||||
"""Determine the outer indices of expression ``expr``
|
||||
|
||||
By *outer* we mean indices that are not summation indices. Returns a set
|
||||
and a dict. The set contains outer indices and the dict contains
|
||||
information about index symmetries.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy.tensor.index_methods import get_indices
|
||||
>>> from sympy import symbols
|
||||
>>> from sympy.tensor import IndexedBase
|
||||
>>> x, y, A = map(IndexedBase, ['x', 'y', 'A'])
|
||||
>>> i, j, a, z = symbols('i j a z', integer=True)
|
||||
|
||||
The indices of the total expression is determined, Repeated indices imply a
|
||||
summation, for instance the trace of a matrix A:
|
||||
|
||||
>>> get_indices(A[i, i])
|
||||
(set(), {})
|
||||
|
||||
In the case of many terms, the terms are required to have identical
|
||||
outer indices. Else an IndexConformanceException is raised.
|
||||
|
||||
>>> get_indices(x[i] + A[i, j]*y[j])
|
||||
({i}, {})
|
||||
|
||||
:Exceptions:
|
||||
|
||||
An IndexConformanceException means that the terms ar not compatible, e.g.
|
||||
|
||||
>>> get_indices(x[i] + y[j]) #doctest: +SKIP
|
||||
(...)
|
||||
IndexConformanceException: Indices are not consistent: x(i) + y(j)
|
||||
|
||||
.. warning::
|
||||
The concept of *outer* indices applies recursively, starting on the deepest
|
||||
level. This implies that dummies inside parenthesis are assumed to be
|
||||
summed first, so that the following expression is handled gracefully:
|
||||
|
||||
>>> get_indices((x[i] + A[i, j]*y[j])*x[j])
|
||||
({i, j}, {})
|
||||
|
||||
This is correct and may appear convenient, but you need to be careful
|
||||
with this as SymPy will happily .expand() the product, if requested. The
|
||||
resulting expression would mix the outer ``j`` with the dummies inside
|
||||
the parenthesis, which makes it a different expression. To be on the
|
||||
safe side, it is best to avoid such ambiguities by using unique indices
|
||||
for all contractions that should be held separate.
|
||||
|
||||
"""
|
||||
# We call ourself recursively to determine indices of sub expressions.
|
||||
|
||||
# break recursion
|
||||
if isinstance(expr, Indexed):
|
||||
c = expr.indices
|
||||
inds, dummies = _remove_repeated(c)
|
||||
return inds, {}
|
||||
elif expr is None:
|
||||
return set(), {}
|
||||
elif isinstance(expr, Idx):
|
||||
return {expr}, {}
|
||||
elif expr.is_Atom:
|
||||
return set(), {}
|
||||
|
||||
|
||||
# recurse via specialized functions
|
||||
else:
|
||||
if expr.is_Mul:
|
||||
return _get_indices_Mul(expr)
|
||||
elif expr.is_Add:
|
||||
return _get_indices_Add(expr)
|
||||
elif expr.is_Pow or isinstance(expr, exp):
|
||||
return _get_indices_Pow(expr)
|
||||
|
||||
elif isinstance(expr, Piecewise):
|
||||
# FIXME: No support for Piecewise yet
|
||||
return set(), {}
|
||||
elif isinstance(expr, Function):
|
||||
# Support ufunc like behaviour by returning indices from arguments.
|
||||
# Functions do not interpret repeated indices across arguments
|
||||
# as summation
|
||||
ind0 = set()
|
||||
for arg in expr.args:
|
||||
ind, sym = get_indices(arg)
|
||||
ind0 |= ind
|
||||
return ind0, sym
|
||||
|
||||
# this test is expensive, so it should be at the end
|
||||
elif not expr.has(Indexed):
|
||||
return set(), {}
|
||||
raise NotImplementedError(
|
||||
"FIXME: No specialized handling of type %s" % type(expr))
|
||||
|
||||
|
||||
def get_contraction_structure(expr):
|
||||
"""Determine dummy indices of ``expr`` and describe its structure
|
||||
|
||||
By *dummy* we mean indices that are summation indices.
|
||||
|
||||
The structure of the expression is determined and described as follows:
|
||||
|
||||
1) A conforming summation of Indexed objects is described with a dict where
|
||||
the keys are summation indices and the corresponding values are sets
|
||||
containing all terms for which the summation applies. All Add objects
|
||||
in the SymPy expression tree are described like this.
|
||||
|
||||
2) For all nodes in the SymPy expression tree that are *not* of type Add, the
|
||||
following applies:
|
||||
|
||||
If a node discovers contractions in one of its arguments, the node
|
||||
itself will be stored as a key in the dict. For that key, the
|
||||
corresponding value is a list of dicts, each of which is the result of a
|
||||
recursive call to get_contraction_structure(). The list contains only
|
||||
dicts for the non-trivial deeper contractions, omitting dicts with None
|
||||
as the one and only key.
|
||||
|
||||
.. Note:: The presence of expressions among the dictionary keys indicates
|
||||
multiple levels of index contractions. A nested dict displays nested
|
||||
contractions and may itself contain dicts from a deeper level. In
|
||||
practical calculations the summation in the deepest nested level must be
|
||||
calculated first so that the outer expression can access the resulting
|
||||
indexed object.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy.tensor.index_methods import get_contraction_structure
|
||||
>>> from sympy import default_sort_key
|
||||
>>> from sympy.tensor import IndexedBase, Idx
|
||||
>>> x, y, A = map(IndexedBase, ['x', 'y', 'A'])
|
||||
>>> i, j, k, l = map(Idx, ['i', 'j', 'k', 'l'])
|
||||
>>> get_contraction_structure(x[i]*y[i] + A[j, j])
|
||||
{(i,): {x[i]*y[i]}, (j,): {A[j, j]}}
|
||||
>>> get_contraction_structure(x[i]*y[j])
|
||||
{None: {x[i]*y[j]}}
|
||||
|
||||
A multiplication of contracted factors results in nested dicts representing
|
||||
the internal contractions.
|
||||
|
||||
>>> d = get_contraction_structure(x[i, i]*y[j, j])
|
||||
>>> sorted(d.keys(), key=default_sort_key)
|
||||
[None, x[i, i]*y[j, j]]
|
||||
|
||||
In this case, the product has no contractions:
|
||||
|
||||
>>> d[None]
|
||||
{x[i, i]*y[j, j]}
|
||||
|
||||
Factors are contracted "first":
|
||||
|
||||
>>> sorted(d[x[i, i]*y[j, j]], key=default_sort_key)
|
||||
[{(i,): {x[i, i]}}, {(j,): {y[j, j]}}]
|
||||
|
||||
A parenthesized Add object is also returned as a nested dictionary. The
|
||||
term containing the parenthesis is a Mul with a contraction among the
|
||||
arguments, so it will be found as a key in the result. It stores the
|
||||
dictionary resulting from a recursive call on the Add expression.
|
||||
|
||||
>>> d = get_contraction_structure(x[i]*(y[i] + A[i, j]*x[j]))
|
||||
>>> sorted(d.keys(), key=default_sort_key)
|
||||
[(A[i, j]*x[j] + y[i])*x[i], (i,)]
|
||||
>>> d[(i,)]
|
||||
{(A[i, j]*x[j] + y[i])*x[i]}
|
||||
>>> d[x[i]*(A[i, j]*x[j] + y[i])]
|
||||
[{None: {y[i]}, (j,): {A[i, j]*x[j]}}]
|
||||
|
||||
Powers with contractions in either base or exponent will also be found as
|
||||
keys in the dictionary, mapping to a list of results from recursive calls:
|
||||
|
||||
>>> d = get_contraction_structure(A[j, j]**A[i, i])
|
||||
>>> d[None]
|
||||
{A[j, j]**A[i, i]}
|
||||
>>> nested_contractions = d[A[j, j]**A[i, i]]
|
||||
>>> nested_contractions[0]
|
||||
{(j,): {A[j, j]}}
|
||||
>>> nested_contractions[1]
|
||||
{(i,): {A[i, i]}}
|
||||
|
||||
The description of the contraction structure may appear complicated when
|
||||
represented with a string in the above examples, but it is easy to iterate
|
||||
over:
|
||||
|
||||
>>> from sympy import Expr
|
||||
>>> for key in d:
|
||||
... if isinstance(key, Expr):
|
||||
... continue
|
||||
... for term in d[key]:
|
||||
... if term in d:
|
||||
... # treat deepest contraction first
|
||||
... pass
|
||||
... # treat outermost contactions here
|
||||
|
||||
"""
|
||||
|
||||
# We call ourself recursively to inspect sub expressions.
|
||||
|
||||
if isinstance(expr, Indexed):
|
||||
junk, key = _remove_repeated(expr.indices)
|
||||
return {key or None: {expr}}
|
||||
elif expr.is_Atom:
|
||||
return {None: {expr}}
|
||||
elif expr.is_Mul:
|
||||
junk, junk, key = _get_indices_Mul(expr, return_dummies=True)
|
||||
result = {key or None: {expr}}
|
||||
# recurse on every factor
|
||||
nested = []
|
||||
for fac in expr.args:
|
||||
facd = get_contraction_structure(fac)
|
||||
if not (None in facd and len(facd) == 1):
|
||||
nested.append(facd)
|
||||
if nested:
|
||||
result[expr] = nested
|
||||
return result
|
||||
elif expr.is_Pow or isinstance(expr, exp):
|
||||
# recurse in base and exp separately. If either has internal
|
||||
# contractions we must include ourselves as a key in the returned dict
|
||||
b, e = expr.as_base_exp()
|
||||
dbase = get_contraction_structure(b)
|
||||
dexp = get_contraction_structure(e)
|
||||
|
||||
dicts = []
|
||||
for d in dbase, dexp:
|
||||
if not (None in d and len(d) == 1):
|
||||
dicts.append(d)
|
||||
result = {None: {expr}}
|
||||
if dicts:
|
||||
result[expr] = dicts
|
||||
return result
|
||||
elif expr.is_Add:
|
||||
# Note: we just collect all terms with identical summation indices, We
|
||||
# do nothing to identify equivalent terms here, as this would require
|
||||
# substitutions or pattern matching in expressions of unknown
|
||||
# complexity.
|
||||
result = {}
|
||||
for term in expr.args:
|
||||
# recurse on every term
|
||||
d = get_contraction_structure(term)
|
||||
for key in d:
|
||||
if key in result:
|
||||
result[key] |= d[key]
|
||||
else:
|
||||
result[key] = d[key]
|
||||
return result
|
||||
|
||||
elif isinstance(expr, Piecewise):
|
||||
# FIXME: No support for Piecewise yet
|
||||
return {None: expr}
|
||||
elif isinstance(expr, Function):
|
||||
# Collect non-trivial contraction structures in each argument
|
||||
# We do not report repeated indices in separate arguments as a
|
||||
# contraction
|
||||
deeplist = []
|
||||
for arg in expr.args:
|
||||
deep = get_contraction_structure(arg)
|
||||
if not (None in deep and len(deep) == 1):
|
||||
deeplist.append(deep)
|
||||
d = {None: {expr}}
|
||||
if deeplist:
|
||||
d[expr] = deeplist
|
||||
return d
|
||||
|
||||
# this test is expensive, so it should be at the end
|
||||
elif not expr.has(Indexed):
|
||||
return {None: {expr}}
|
||||
raise NotImplementedError(
|
||||
"FIXME: No specialized handling of type %s" % type(expr))
|
||||
@@ -0,0 +1,793 @@
|
||||
r"""Module that defines indexed objects.
|
||||
|
||||
The classes ``IndexedBase``, ``Indexed``, and ``Idx`` represent a
|
||||
matrix element ``M[i, j]`` as in the following diagram::
|
||||
|
||||
1) The Indexed class represents the entire indexed object.
|
||||
|
|
||||
___|___
|
||||
' '
|
||||
M[i, j]
|
||||
/ \__\______
|
||||
| |
|
||||
| |
|
||||
| 2) The Idx class represents indices; each Idx can
|
||||
| optionally contain information about its range.
|
||||
|
|
||||
3) IndexedBase represents the 'stem' of an indexed object, here `M`.
|
||||
The stem used by itself is usually taken to represent the entire
|
||||
array.
|
||||
|
||||
There can be any number of indices on an Indexed object. No
|
||||
transformation properties are implemented in these Base objects, but
|
||||
implicit contraction of repeated indices is supported.
|
||||
|
||||
Note that the support for complicated (i.e. non-atomic) integer
|
||||
expressions as indices is limited. (This should be improved in
|
||||
future releases.)
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
To express the above matrix element example you would write:
|
||||
|
||||
>>> from sympy import symbols, IndexedBase, Idx
|
||||
>>> M = IndexedBase('M')
|
||||
>>> i, j = symbols('i j', cls=Idx)
|
||||
>>> M[i, j]
|
||||
M[i, j]
|
||||
|
||||
Repeated indices in a product implies a summation, so to express a
|
||||
matrix-vector product in terms of Indexed objects:
|
||||
|
||||
>>> x = IndexedBase('x')
|
||||
>>> M[i, j]*x[j]
|
||||
M[i, j]*x[j]
|
||||
|
||||
If the indexed objects will be converted to component based arrays, e.g.
|
||||
with the code printers or the autowrap framework, you also need to provide
|
||||
(symbolic or numerical) dimensions. This can be done by passing an
|
||||
optional shape parameter to IndexedBase upon construction:
|
||||
|
||||
>>> dim1, dim2 = symbols('dim1 dim2', integer=True)
|
||||
>>> A = IndexedBase('A', shape=(dim1, 2*dim1, dim2))
|
||||
>>> A.shape
|
||||
(dim1, 2*dim1, dim2)
|
||||
>>> A[i, j, 3].shape
|
||||
(dim1, 2*dim1, dim2)
|
||||
|
||||
If an IndexedBase object has no shape information, it is assumed that the
|
||||
array is as large as the ranges of its indices:
|
||||
|
||||
>>> n, m = symbols('n m', integer=True)
|
||||
>>> i = Idx('i', m)
|
||||
>>> j = Idx('j', n)
|
||||
>>> M[i, j].shape
|
||||
(m, n)
|
||||
>>> M[i, j].ranges
|
||||
[(0, m - 1), (0, n - 1)]
|
||||
|
||||
The above can be compared with the following:
|
||||
|
||||
>>> A[i, 2, j].shape
|
||||
(dim1, 2*dim1, dim2)
|
||||
>>> A[i, 2, j].ranges
|
||||
[(0, m - 1), None, (0, n - 1)]
|
||||
|
||||
To analyze the structure of indexed expressions, you can use the methods
|
||||
get_indices() and get_contraction_structure():
|
||||
|
||||
>>> from sympy.tensor import get_indices, get_contraction_structure
|
||||
>>> get_indices(A[i, j, j])
|
||||
({i}, {})
|
||||
>>> get_contraction_structure(A[i, j, j])
|
||||
{(j,): {A[i, j, j]}}
|
||||
|
||||
See the appropriate docstrings for a detailed explanation of the output.
|
||||
"""
|
||||
|
||||
# TODO: (some ideas for improvement)
|
||||
#
|
||||
# o test and guarantee numpy compatibility
|
||||
# - implement full support for broadcasting
|
||||
# - strided arrays
|
||||
#
|
||||
# o more functions to analyze indexed expressions
|
||||
# - identify standard constructs, e.g matrix-vector product in a subexpression
|
||||
#
|
||||
# o functions to generate component based arrays (numpy and sympy.Matrix)
|
||||
# - generate a single array directly from Indexed
|
||||
# - convert simple sub-expressions
|
||||
#
|
||||
# o sophisticated indexing (possibly in subclasses to preserve simplicity)
|
||||
# - Idx with range smaller than dimension of Indexed
|
||||
# - Idx with stepsize != 1
|
||||
# - Idx with step determined by function call
|
||||
from collections.abc import Iterable
|
||||
|
||||
from sympy.core.numbers import Number
|
||||
from sympy.core.assumptions import StdFactKB
|
||||
from sympy.core import Expr, Tuple, sympify, S
|
||||
from sympy.core.symbol import _filter_assumptions, Symbol
|
||||
from sympy.core.logic import fuzzy_bool, fuzzy_not
|
||||
from sympy.core.sympify import _sympify
|
||||
from sympy.functions.special.tensor_functions import KroneckerDelta
|
||||
from sympy.multipledispatch import dispatch
|
||||
from sympy.utilities.iterables import is_sequence, NotIterable
|
||||
from sympy.utilities.misc import filldedent
|
||||
|
||||
|
||||
class IndexException(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class Indexed(Expr):
|
||||
"""Represents a mathematical object with indices.
|
||||
|
||||
>>> from sympy import Indexed, IndexedBase, Idx, symbols
|
||||
>>> i, j = symbols('i j', cls=Idx)
|
||||
>>> Indexed('A', i, j)
|
||||
A[i, j]
|
||||
|
||||
It is recommended that ``Indexed`` objects be created by indexing ``IndexedBase``:
|
||||
``IndexedBase('A')[i, j]`` instead of ``Indexed(IndexedBase('A'), i, j)``.
|
||||
|
||||
>>> A = IndexedBase('A')
|
||||
>>> a_ij = A[i, j] # Prefer this,
|
||||
>>> b_ij = Indexed(A, i, j) # over this.
|
||||
>>> a_ij == b_ij
|
||||
True
|
||||
|
||||
"""
|
||||
is_Indexed = True
|
||||
is_symbol = True
|
||||
is_Atom = True
|
||||
|
||||
def __new__(cls, base, *args, **kw_args):
|
||||
from sympy.tensor.array.ndim_array import NDimArray
|
||||
from sympy.matrices.matrixbase import MatrixBase
|
||||
|
||||
if not args:
|
||||
raise IndexException("Indexed needs at least one index.")
|
||||
if isinstance(base, (str, Symbol)):
|
||||
base = IndexedBase(base)
|
||||
elif not hasattr(base, '__getitem__') and not isinstance(base, IndexedBase):
|
||||
raise TypeError(filldedent("""
|
||||
The base can only be replaced with a string, Symbol,
|
||||
IndexedBase or an object with a method for getting
|
||||
items (i.e. an object with a `__getitem__` method).
|
||||
"""))
|
||||
args = list(map(sympify, args))
|
||||
if isinstance(base, (NDimArray, Iterable, Tuple, MatrixBase)) and all(i.is_number for i in args):
|
||||
if len(args) == 1:
|
||||
return base[args[0]]
|
||||
else:
|
||||
return base[args]
|
||||
|
||||
base = _sympify(base)
|
||||
|
||||
obj = Expr.__new__(cls, base, *args, **kw_args)
|
||||
|
||||
IndexedBase._set_assumptions(obj, base.assumptions0)
|
||||
|
||||
return obj
|
||||
|
||||
def _hashable_content(self):
|
||||
return super()._hashable_content() + tuple(sorted(self.assumptions0.items()))
|
||||
|
||||
@property
|
||||
def name(self):
|
||||
return str(self)
|
||||
|
||||
@property
|
||||
def _diff_wrt(self):
|
||||
"""Allow derivatives with respect to an ``Indexed`` object."""
|
||||
return True
|
||||
|
||||
def _eval_derivative(self, wrt):
|
||||
from sympy.tensor.array.ndim_array import NDimArray
|
||||
|
||||
if isinstance(wrt, Indexed) and wrt.base == self.base:
|
||||
if len(self.indices) != len(wrt.indices):
|
||||
msg = "Different # of indices: d({!s})/d({!s})".format(self,
|
||||
wrt)
|
||||
raise IndexException(msg)
|
||||
result = S.One
|
||||
for index1, index2 in zip(self.indices, wrt.indices):
|
||||
result *= KroneckerDelta(index1, index2)
|
||||
return result
|
||||
elif isinstance(self.base, NDimArray):
|
||||
from sympy.tensor.array import derive_by_array
|
||||
return Indexed(derive_by_array(self.base, wrt), *self.args[1:])
|
||||
else:
|
||||
if Tuple(self.indices).has(wrt):
|
||||
return S.NaN
|
||||
return S.Zero
|
||||
|
||||
@property
|
||||
def assumptions0(self):
|
||||
return {k: v for k, v in self._assumptions.items() if v is not None}
|
||||
|
||||
@property
|
||||
def base(self):
|
||||
"""Returns the ``IndexedBase`` of the ``Indexed`` object.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy import Indexed, IndexedBase, Idx, symbols
|
||||
>>> i, j = symbols('i j', cls=Idx)
|
||||
>>> Indexed('A', i, j).base
|
||||
A
|
||||
>>> B = IndexedBase('B')
|
||||
>>> B == B[i, j].base
|
||||
True
|
||||
|
||||
"""
|
||||
return self.args[0]
|
||||
|
||||
@property
|
||||
def indices(self):
|
||||
"""
|
||||
Returns the indices of the ``Indexed`` object.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy import Indexed, Idx, symbols
|
||||
>>> i, j = symbols('i j', cls=Idx)
|
||||
>>> Indexed('A', i, j).indices
|
||||
(i, j)
|
||||
|
||||
"""
|
||||
return self.args[1:]
|
||||
|
||||
@property
|
||||
def rank(self):
|
||||
"""
|
||||
Returns the rank of the ``Indexed`` object.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy import Indexed, Idx, symbols
|
||||
>>> i, j, k, l, m = symbols('i:m', cls=Idx)
|
||||
>>> Indexed('A', i, j).rank
|
||||
2
|
||||
>>> q = Indexed('A', i, j, k, l, m)
|
||||
>>> q.rank
|
||||
5
|
||||
>>> q.rank == len(q.indices)
|
||||
True
|
||||
|
||||
"""
|
||||
return len(self.args) - 1
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
"""Returns a list with dimensions of each index.
|
||||
|
||||
Dimensions is a property of the array, not of the indices. Still, if
|
||||
the ``IndexedBase`` does not define a shape attribute, it is assumed
|
||||
that the ranges of the indices correspond to the shape of the array.
|
||||
|
||||
>>> from sympy import IndexedBase, Idx, symbols
|
||||
>>> n, m = symbols('n m', integer=True)
|
||||
>>> i = Idx('i', m)
|
||||
>>> j = Idx('j', m)
|
||||
>>> A = IndexedBase('A', shape=(n, n))
|
||||
>>> B = IndexedBase('B')
|
||||
>>> A[i, j].shape
|
||||
(n, n)
|
||||
>>> B[i, j].shape
|
||||
(m, m)
|
||||
"""
|
||||
|
||||
if self.base.shape:
|
||||
return self.base.shape
|
||||
sizes = []
|
||||
for i in self.indices:
|
||||
upper = getattr(i, 'upper', None)
|
||||
lower = getattr(i, 'lower', None)
|
||||
if None in (upper, lower):
|
||||
raise IndexException(filldedent("""
|
||||
Range is not defined for all indices in: %s""" % self))
|
||||
try:
|
||||
size = upper - lower + 1
|
||||
except TypeError:
|
||||
raise IndexException(filldedent("""
|
||||
Shape cannot be inferred from Idx with
|
||||
undefined range: %s""" % self))
|
||||
sizes.append(size)
|
||||
return Tuple(*sizes)
|
||||
|
||||
@property
|
||||
def ranges(self):
|
||||
"""Returns a list of tuples with lower and upper range of each index.
|
||||
|
||||
If an index does not define the data members upper and lower, the
|
||||
corresponding slot in the list contains ``None`` instead of a tuple.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy import Indexed,Idx, symbols
|
||||
>>> Indexed('A', Idx('i', 2), Idx('j', 4), Idx('k', 8)).ranges
|
||||
[(0, 1), (0, 3), (0, 7)]
|
||||
>>> Indexed('A', Idx('i', 3), Idx('j', 3), Idx('k', 3)).ranges
|
||||
[(0, 2), (0, 2), (0, 2)]
|
||||
>>> x, y, z = symbols('x y z', integer=True)
|
||||
>>> Indexed('A', x, y, z).ranges
|
||||
[None, None, None]
|
||||
|
||||
"""
|
||||
ranges = []
|
||||
sentinel = object()
|
||||
for i in self.indices:
|
||||
upper = getattr(i, 'upper', sentinel)
|
||||
lower = getattr(i, 'lower', sentinel)
|
||||
if sentinel not in (upper, lower):
|
||||
ranges.append((lower, upper))
|
||||
else:
|
||||
ranges.append(None)
|
||||
return ranges
|
||||
|
||||
def _sympystr(self, p):
|
||||
indices = list(map(p.doprint, self.indices))
|
||||
return "%s[%s]" % (p.doprint(self.base), ", ".join(indices))
|
||||
|
||||
@property
|
||||
def free_symbols(self):
|
||||
base_free_symbols = self.base.free_symbols
|
||||
indices_free_symbols = {
|
||||
fs for i in self.indices for fs in i.free_symbols}
|
||||
if base_free_symbols:
|
||||
return {self} | base_free_symbols | indices_free_symbols
|
||||
else:
|
||||
return indices_free_symbols
|
||||
|
||||
@property
|
||||
def expr_free_symbols(self):
|
||||
from sympy.utilities.exceptions import sympy_deprecation_warning
|
||||
sympy_deprecation_warning("""
|
||||
The expr_free_symbols property is deprecated. Use free_symbols to get
|
||||
the free symbols of an expression.
|
||||
""",
|
||||
deprecated_since_version="1.9",
|
||||
active_deprecations_target="deprecated-expr-free-symbols")
|
||||
|
||||
return {self}
|
||||
|
||||
|
||||
class IndexedBase(Expr, NotIterable):
|
||||
"""Represent the base or stem of an indexed object
|
||||
|
||||
The IndexedBase class represent an array that contains elements. The main purpose
|
||||
of this class is to allow the convenient creation of objects of the Indexed
|
||||
class. The __getitem__ method of IndexedBase returns an instance of
|
||||
Indexed. Alone, without indices, the IndexedBase class can be used as a
|
||||
notation for e.g. matrix equations, resembling what you could do with the
|
||||
Symbol class. But, the IndexedBase class adds functionality that is not
|
||||
available for Symbol instances:
|
||||
|
||||
- An IndexedBase object can optionally store shape information. This can
|
||||
be used in to check array conformance and conditions for numpy
|
||||
broadcasting. (TODO)
|
||||
- An IndexedBase object implements syntactic sugar that allows easy symbolic
|
||||
representation of array operations, using implicit summation of
|
||||
repeated indices.
|
||||
- The IndexedBase object symbolizes a mathematical structure equivalent
|
||||
to arrays, and is recognized as such for code generation and automatic
|
||||
compilation and wrapping.
|
||||
|
||||
>>> from sympy.tensor import IndexedBase, Idx
|
||||
>>> from sympy import symbols
|
||||
>>> A = IndexedBase('A'); A
|
||||
A
|
||||
>>> type(A)
|
||||
<class 'sympy.tensor.indexed.IndexedBase'>
|
||||
|
||||
When an IndexedBase object receives indices, it returns an array with named
|
||||
axes, represented by an Indexed object:
|
||||
|
||||
>>> i, j = symbols('i j', integer=True)
|
||||
>>> A[i, j, 2]
|
||||
A[i, j, 2]
|
||||
>>> type(A[i, j, 2])
|
||||
<class 'sympy.tensor.indexed.Indexed'>
|
||||
|
||||
The IndexedBase constructor takes an optional shape argument. If given,
|
||||
it overrides any shape information in the indices. (But not the index
|
||||
ranges!)
|
||||
|
||||
>>> m, n, o, p = symbols('m n o p', integer=True)
|
||||
>>> i = Idx('i', m)
|
||||
>>> j = Idx('j', n)
|
||||
>>> A[i, j].shape
|
||||
(m, n)
|
||||
>>> B = IndexedBase('B', shape=(o, p))
|
||||
>>> B[i, j].shape
|
||||
(o, p)
|
||||
|
||||
Assumptions can be specified with keyword arguments the same way as for Symbol:
|
||||
|
||||
>>> A_real = IndexedBase('A', real=True)
|
||||
>>> A_real.is_real
|
||||
True
|
||||
>>> A != A_real
|
||||
True
|
||||
|
||||
Assumptions can also be inherited if a Symbol is used to initialize the IndexedBase:
|
||||
|
||||
>>> I = symbols('I', integer=True)
|
||||
>>> C_inherit = IndexedBase(I)
|
||||
>>> C_explicit = IndexedBase('I', integer=True)
|
||||
>>> C_inherit == C_explicit
|
||||
True
|
||||
"""
|
||||
is_symbol = True
|
||||
is_Atom = True
|
||||
|
||||
@staticmethod
|
||||
def _set_assumptions(obj, assumptions):
|
||||
"""Set assumptions on obj, making sure to apply consistent values."""
|
||||
tmp_asm_copy = assumptions.copy()
|
||||
is_commutative = fuzzy_bool(assumptions.get('commutative', True))
|
||||
assumptions['commutative'] = is_commutative
|
||||
obj._assumptions = StdFactKB(assumptions)
|
||||
obj._assumptions._generator = tmp_asm_copy # Issue #8873
|
||||
|
||||
def __new__(cls, label, shape=None, *, offset=S.Zero, strides=None, **kw_args):
|
||||
from sympy.matrices.matrixbase import MatrixBase
|
||||
from sympy.tensor.array.ndim_array import NDimArray
|
||||
|
||||
assumptions, kw_args = _filter_assumptions(kw_args)
|
||||
if isinstance(label, str):
|
||||
label = Symbol(label, **assumptions)
|
||||
elif isinstance(label, Symbol):
|
||||
assumptions = label._merge(assumptions)
|
||||
elif isinstance(label, (MatrixBase, NDimArray)):
|
||||
return label
|
||||
elif isinstance(label, Iterable):
|
||||
return _sympify(label)
|
||||
else:
|
||||
label = _sympify(label)
|
||||
|
||||
if is_sequence(shape):
|
||||
shape = Tuple(*shape)
|
||||
elif shape is not None:
|
||||
shape = Tuple(shape)
|
||||
|
||||
if shape is not None:
|
||||
obj = Expr.__new__(cls, label, shape)
|
||||
else:
|
||||
obj = Expr.__new__(cls, label)
|
||||
obj._shape = shape
|
||||
obj._offset = offset
|
||||
obj._strides = strides
|
||||
obj._name = str(label)
|
||||
|
||||
IndexedBase._set_assumptions(obj, assumptions)
|
||||
return obj
|
||||
|
||||
@property
|
||||
def name(self):
|
||||
return self._name
|
||||
|
||||
def _hashable_content(self):
|
||||
return super()._hashable_content() + tuple(sorted(self.assumptions0.items()))
|
||||
|
||||
@property
|
||||
def assumptions0(self):
|
||||
return {k: v for k, v in self._assumptions.items() if v is not None}
|
||||
|
||||
def __getitem__(self, indices, **kw_args):
|
||||
if is_sequence(indices):
|
||||
# Special case needed because M[*my_tuple] is a syntax error.
|
||||
if self.shape and len(self.shape) != len(indices):
|
||||
raise IndexException("Rank mismatch.")
|
||||
return Indexed(self, *indices, **kw_args)
|
||||
else:
|
||||
if self.shape and len(self.shape) != 1:
|
||||
raise IndexException("Rank mismatch.")
|
||||
return Indexed(self, indices, **kw_args)
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
"""Returns the shape of the ``IndexedBase`` object.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy import IndexedBase, Idx
|
||||
>>> from sympy.abc import x, y
|
||||
>>> IndexedBase('A', shape=(x, y)).shape
|
||||
(x, y)
|
||||
|
||||
Note: If the shape of the ``IndexedBase`` is specified, it will override
|
||||
any shape information given by the indices.
|
||||
|
||||
>>> A = IndexedBase('A', shape=(x, y))
|
||||
>>> B = IndexedBase('B')
|
||||
>>> i = Idx('i', 2)
|
||||
>>> j = Idx('j', 1)
|
||||
>>> A[i, j].shape
|
||||
(x, y)
|
||||
>>> B[i, j].shape
|
||||
(2, 1)
|
||||
|
||||
"""
|
||||
return self._shape
|
||||
|
||||
@property
|
||||
def strides(self):
|
||||
"""Returns the strided scheme for the ``IndexedBase`` object.
|
||||
|
||||
Normally this is a tuple denoting the number of
|
||||
steps to take in the respective dimension when traversing
|
||||
an array. For code generation purposes strides='C' and
|
||||
strides='F' can also be used.
|
||||
|
||||
strides='C' would mean that code printer would unroll
|
||||
in row-major order and 'F' means unroll in column major
|
||||
order.
|
||||
|
||||
"""
|
||||
|
||||
return self._strides
|
||||
|
||||
@property
|
||||
def offset(self):
|
||||
"""Returns the offset for the ``IndexedBase`` object.
|
||||
|
||||
This is the value added to the resulting index when the
|
||||
2D Indexed object is unrolled to a 1D form. Used in code
|
||||
generation.
|
||||
|
||||
Examples
|
||||
==========
|
||||
>>> from sympy.printing import ccode
|
||||
>>> from sympy.tensor import IndexedBase, Idx
|
||||
>>> from sympy import symbols
|
||||
>>> l, m, n, o = symbols('l m n o', integer=True)
|
||||
>>> A = IndexedBase('A', strides=(l, m, n), offset=o)
|
||||
>>> i, j, k = map(Idx, 'ijk')
|
||||
>>> ccode(A[i, j, k])
|
||||
'A[l*i + m*j + n*k + o]'
|
||||
|
||||
"""
|
||||
return self._offset
|
||||
|
||||
@property
|
||||
def label(self):
|
||||
"""Returns the label of the ``IndexedBase`` object.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy import IndexedBase
|
||||
>>> from sympy.abc import x, y
|
||||
>>> IndexedBase('A', shape=(x, y)).label
|
||||
A
|
||||
|
||||
"""
|
||||
return self.args[0]
|
||||
|
||||
def _sympystr(self, p):
|
||||
return p.doprint(self.label)
|
||||
|
||||
|
||||
class Idx(Expr):
|
||||
"""Represents an integer index as an ``Integer`` or integer expression.
|
||||
|
||||
There are a number of ways to create an ``Idx`` object. The constructor
|
||||
takes two arguments:
|
||||
|
||||
``label``
|
||||
An integer or a symbol that labels the index.
|
||||
``range``
|
||||
Optionally you can specify a range as either
|
||||
|
||||
* ``Symbol`` or integer: This is interpreted as a dimension. Lower and
|
||||
upper bounds are set to ``0`` and ``range - 1``, respectively.
|
||||
* ``tuple``: The two elements are interpreted as the lower and upper
|
||||
bounds of the range, respectively.
|
||||
|
||||
Note: bounds of the range are assumed to be either integer or infinite (oo
|
||||
and -oo are allowed to specify an unbounded range). If ``n`` is given as a
|
||||
bound, then ``n.is_integer`` must not return false.
|
||||
|
||||
For convenience, if the label is given as a string it is automatically
|
||||
converted to an integer symbol. (Note: this conversion is not done for
|
||||
range or dimension arguments.)
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy import Idx, symbols, oo
|
||||
>>> n, i, L, U = symbols('n i L U', integer=True)
|
||||
|
||||
If a string is given for the label an integer ``Symbol`` is created and the
|
||||
bounds are both ``None``:
|
||||
|
||||
>>> idx = Idx('qwerty'); idx
|
||||
qwerty
|
||||
>>> idx.lower, idx.upper
|
||||
(None, None)
|
||||
|
||||
Both upper and lower bounds can be specified:
|
||||
|
||||
>>> idx = Idx(i, (L, U)); idx
|
||||
i
|
||||
>>> idx.lower, idx.upper
|
||||
(L, U)
|
||||
|
||||
When only a single bound is given it is interpreted as the dimension
|
||||
and the lower bound defaults to 0:
|
||||
|
||||
>>> idx = Idx(i, n); idx.lower, idx.upper
|
||||
(0, n - 1)
|
||||
>>> idx = Idx(i, 4); idx.lower, idx.upper
|
||||
(0, 3)
|
||||
>>> idx = Idx(i, oo); idx.lower, idx.upper
|
||||
(0, oo)
|
||||
|
||||
"""
|
||||
|
||||
is_integer = True
|
||||
is_finite = True
|
||||
is_real = True
|
||||
is_symbol = True
|
||||
is_Atom = True
|
||||
_diff_wrt = True
|
||||
|
||||
def __new__(cls, label, range=None, **kw_args):
|
||||
|
||||
if isinstance(label, str):
|
||||
label = Symbol(label, integer=True)
|
||||
label, range = list(map(sympify, (label, range)))
|
||||
|
||||
if label.is_Number:
|
||||
if not label.is_integer:
|
||||
raise TypeError("Index is not an integer number.")
|
||||
return label
|
||||
|
||||
if not label.is_integer:
|
||||
raise TypeError("Idx object requires an integer label.")
|
||||
|
||||
elif is_sequence(range):
|
||||
if len(range) != 2:
|
||||
raise ValueError(filldedent("""
|
||||
Idx range tuple must have length 2, but got %s""" % len(range)))
|
||||
for bound in range:
|
||||
if (bound.is_integer is False and bound is not S.Infinity
|
||||
and bound is not S.NegativeInfinity):
|
||||
raise TypeError("Idx object requires integer bounds.")
|
||||
args = label, Tuple(*range)
|
||||
elif isinstance(range, Expr):
|
||||
if range is not S.Infinity and fuzzy_not(range.is_integer):
|
||||
raise TypeError("Idx object requires an integer dimension.")
|
||||
args = label, Tuple(0, range - 1)
|
||||
elif range:
|
||||
raise TypeError(filldedent("""
|
||||
The range must be an ordered iterable or
|
||||
integer SymPy expression."""))
|
||||
else:
|
||||
args = label,
|
||||
|
||||
obj = Expr.__new__(cls, *args, **kw_args)
|
||||
obj._assumptions["finite"] = True
|
||||
obj._assumptions["real"] = True
|
||||
return obj
|
||||
|
||||
@property
|
||||
def label(self):
|
||||
"""Returns the label (Integer or integer expression) of the Idx object.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy import Idx, Symbol
|
||||
>>> x = Symbol('x', integer=True)
|
||||
>>> Idx(x).label
|
||||
x
|
||||
>>> j = Symbol('j', integer=True)
|
||||
>>> Idx(j).label
|
||||
j
|
||||
>>> Idx(j + 1).label
|
||||
j + 1
|
||||
|
||||
"""
|
||||
return self.args[0]
|
||||
|
||||
@property
|
||||
def lower(self):
|
||||
"""Returns the lower bound of the ``Idx``.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy import Idx
|
||||
>>> Idx('j', 2).lower
|
||||
0
|
||||
>>> Idx('j', 5).lower
|
||||
0
|
||||
>>> Idx('j').lower is None
|
||||
True
|
||||
|
||||
"""
|
||||
try:
|
||||
return self.args[1][0]
|
||||
except IndexError:
|
||||
return
|
||||
|
||||
@property
|
||||
def upper(self):
|
||||
"""Returns the upper bound of the ``Idx``.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy import Idx
|
||||
>>> Idx('j', 2).upper
|
||||
1
|
||||
>>> Idx('j', 5).upper
|
||||
4
|
||||
>>> Idx('j').upper is None
|
||||
True
|
||||
|
||||
"""
|
||||
try:
|
||||
return self.args[1][1]
|
||||
except IndexError:
|
||||
return
|
||||
|
||||
def _sympystr(self, p):
|
||||
return p.doprint(self.label)
|
||||
|
||||
@property
|
||||
def name(self):
|
||||
return self.label.name if self.label.is_Symbol else str(self.label)
|
||||
|
||||
@property
|
||||
def free_symbols(self):
|
||||
return {self}
|
||||
|
||||
|
||||
@dispatch(Idx, Idx)
|
||||
def _eval_is_ge(lhs, rhs): # noqa:F811
|
||||
|
||||
other_upper = rhs if rhs.upper is None else rhs.upper
|
||||
other_lower = rhs if rhs.lower is None else rhs.lower
|
||||
|
||||
if lhs.lower is not None and (lhs.lower >= other_upper) == True:
|
||||
return True
|
||||
if lhs.upper is not None and (lhs.upper < other_lower) == True:
|
||||
return False
|
||||
return None
|
||||
|
||||
|
||||
@dispatch(Idx, Number) # type:ignore
|
||||
def _eval_is_ge(lhs, rhs): # noqa:F811
|
||||
|
||||
other_upper = rhs
|
||||
other_lower = rhs
|
||||
|
||||
if lhs.lower is not None and (lhs.lower >= other_upper) == True:
|
||||
return True
|
||||
if lhs.upper is not None and (lhs.upper < other_lower) == True:
|
||||
return False
|
||||
return None
|
||||
|
||||
|
||||
@dispatch(Number, Idx) # type:ignore
|
||||
def _eval_is_ge(lhs, rhs): # noqa:F811
|
||||
|
||||
other_upper = lhs
|
||||
other_lower = lhs
|
||||
|
||||
if rhs.upper is not None and (rhs.upper <= other_lower) == True:
|
||||
return True
|
||||
if rhs.lower is not None and (rhs.lower > other_upper) == True:
|
||||
return False
|
||||
return None
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,57 @@
|
||||
from sympy.tensor.functions import TensorProduct
|
||||
from sympy.matrices.dense import Matrix
|
||||
from sympy.matrices.expressions.matexpr import MatrixSymbol
|
||||
from sympy.tensor.array import Array
|
||||
from sympy.abc import x, y, z
|
||||
from sympy.abc import i, j, k, l
|
||||
|
||||
|
||||
A = MatrixSymbol("A", 3, 3)
|
||||
B = MatrixSymbol("B", 3, 3)
|
||||
C = MatrixSymbol("C", 3, 3)
|
||||
|
||||
|
||||
def test_TensorProduct_construction():
|
||||
assert TensorProduct(3, 4) == 12
|
||||
assert isinstance(TensorProduct(A, A), TensorProduct)
|
||||
|
||||
expr = TensorProduct(TensorProduct(x, y), z)
|
||||
assert expr == x*y*z
|
||||
|
||||
expr = TensorProduct(TensorProduct(A, B), C)
|
||||
assert expr == TensorProduct(A, B, C)
|
||||
|
||||
expr = TensorProduct(Matrix.eye(2), Array([[0, -1], [1, 0]]))
|
||||
assert expr == Array([
|
||||
[
|
||||
[[0, -1], [1, 0]],
|
||||
[[0, 0], [0, 0]]
|
||||
],
|
||||
[
|
||||
[[0, 0], [0, 0]],
|
||||
[[0, -1], [1, 0]]
|
||||
]
|
||||
])
|
||||
|
||||
|
||||
def test_TensorProduct_shape():
|
||||
|
||||
expr = TensorProduct(3, 4, evaluate=False)
|
||||
assert expr.shape == ()
|
||||
assert expr.rank() == 0
|
||||
|
||||
expr = TensorProduct(Array([1, 2]), Array([x, y]), evaluate=False)
|
||||
assert expr.shape == (2, 2)
|
||||
assert expr.rank() == 2
|
||||
expr = TensorProduct(expr, expr, evaluate=False)
|
||||
assert expr.shape == (2, 2, 2, 2)
|
||||
assert expr.rank() == 4
|
||||
|
||||
expr = TensorProduct(Matrix.eye(2), Array([[0, -1], [1, 0]]), evaluate=False)
|
||||
assert expr.shape == (2, 2, 2, 2)
|
||||
assert expr.rank() == 4
|
||||
|
||||
|
||||
def test_TensorProduct_getitem():
|
||||
expr = TensorProduct(A, B)
|
||||
assert expr[i, j, k, l] == A[i, j]*B[k, l]
|
||||
@@ -0,0 +1,227 @@
|
||||
from sympy.core import symbols, S, Pow, Function
|
||||
from sympy.functions import exp
|
||||
from sympy.testing.pytest import raises
|
||||
from sympy.tensor.indexed import Idx, IndexedBase
|
||||
from sympy.tensor.index_methods import IndexConformanceException
|
||||
|
||||
from sympy.tensor.index_methods import (get_contraction_structure, get_indices)
|
||||
|
||||
|
||||
def test_trivial_indices():
|
||||
x, y = symbols('x y')
|
||||
assert get_indices(x) == (set(), {})
|
||||
assert get_indices(x*y) == (set(), {})
|
||||
assert get_indices(x + y) == (set(), {})
|
||||
assert get_indices(x**y) == (set(), {})
|
||||
|
||||
|
||||
def test_get_indices_Indexed():
|
||||
x = IndexedBase('x')
|
||||
i, j = Idx('i'), Idx('j')
|
||||
assert get_indices(x[i, j]) == ({i, j}, {})
|
||||
assert get_indices(x[j, i]) == ({j, i}, {})
|
||||
|
||||
|
||||
def test_get_indices_Idx():
|
||||
f = Function('f')
|
||||
i, j = Idx('i'), Idx('j')
|
||||
assert get_indices(f(i)*j) == ({i, j}, {})
|
||||
assert get_indices(f(j, i)) == ({j, i}, {})
|
||||
assert get_indices(f(i)*i) == (set(), {})
|
||||
|
||||
|
||||
def test_get_indices_mul():
|
||||
x = IndexedBase('x')
|
||||
y = IndexedBase('y')
|
||||
i, j = Idx('i'), Idx('j')
|
||||
assert get_indices(x[j]*y[i]) == ({i, j}, {})
|
||||
assert get_indices(x[i]*y[j]) == ({i, j}, {})
|
||||
|
||||
|
||||
def test_get_indices_exceptions():
|
||||
x = IndexedBase('x')
|
||||
y = IndexedBase('y')
|
||||
i, j = Idx('i'), Idx('j')
|
||||
raises(IndexConformanceException, lambda: get_indices(x[i] + y[j]))
|
||||
|
||||
|
||||
def test_scalar_broadcast():
|
||||
x = IndexedBase('x')
|
||||
y = IndexedBase('y')
|
||||
i, j = Idx('i'), Idx('j')
|
||||
assert get_indices(x[i] + y[i, i]) == ({i}, {})
|
||||
assert get_indices(x[i] + y[j, j]) == ({i}, {})
|
||||
|
||||
|
||||
def test_get_indices_add():
|
||||
x = IndexedBase('x')
|
||||
y = IndexedBase('y')
|
||||
A = IndexedBase('A')
|
||||
i, j, k = Idx('i'), Idx('j'), Idx('k')
|
||||
assert get_indices(x[i] + 2*y[i]) == ({i}, {})
|
||||
assert get_indices(y[i] + 2*A[i, j]*x[j]) == ({i}, {})
|
||||
assert get_indices(y[i] + 2*(x[i] + A[i, j]*x[j])) == ({i}, {})
|
||||
assert get_indices(y[i] + x[i]*(A[j, j] + 1)) == ({i}, {})
|
||||
assert get_indices(
|
||||
y[i] + x[i]*x[j]*(y[j] + A[j, k]*x[k])) == ({i}, {})
|
||||
|
||||
|
||||
def test_get_indices_Pow():
|
||||
x = IndexedBase('x')
|
||||
y = IndexedBase('y')
|
||||
A = IndexedBase('A')
|
||||
i, j, k = Idx('i'), Idx('j'), Idx('k')
|
||||
assert get_indices(Pow(x[i], y[j])) == ({i, j}, {})
|
||||
assert get_indices(Pow(x[i, k], y[j, k])) == ({i, j, k}, {})
|
||||
assert get_indices(Pow(A[i, k], y[k] + A[k, j]*x[j])) == ({i, k}, {})
|
||||
assert get_indices(Pow(2, x[i])) == get_indices(exp(x[i]))
|
||||
|
||||
# test of a design decision, this may change:
|
||||
assert get_indices(Pow(x[i], 2)) == ({i}, {})
|
||||
|
||||
|
||||
def test_get_contraction_structure_basic():
|
||||
x = IndexedBase('x')
|
||||
y = IndexedBase('y')
|
||||
i, j = Idx('i'), Idx('j')
|
||||
assert get_contraction_structure(x[i]*y[j]) == {None: {x[i]*y[j]}}
|
||||
assert get_contraction_structure(x[i] + y[j]) == {None: {x[i], y[j]}}
|
||||
assert get_contraction_structure(x[i]*y[i]) == {(i,): {x[i]*y[i]}}
|
||||
assert get_contraction_structure(
|
||||
1 + x[i]*y[i]) == {None: {S.One}, (i,): {x[i]*y[i]}}
|
||||
assert get_contraction_structure(x[i]**y[i]) == {None: {x[i]**y[i]}}
|
||||
|
||||
|
||||
def test_get_contraction_structure_complex():
|
||||
x = IndexedBase('x')
|
||||
y = IndexedBase('y')
|
||||
A = IndexedBase('A')
|
||||
i, j, k = Idx('i'), Idx('j'), Idx('k')
|
||||
expr1 = y[i] + A[i, j]*x[j]
|
||||
d1 = {None: {y[i]}, (j,): {A[i, j]*x[j]}}
|
||||
assert get_contraction_structure(expr1) == d1
|
||||
expr2 = expr1*A[k, i] + x[k]
|
||||
d2 = {None: {x[k]}, (i,): {expr1*A[k, i]}, expr1*A[k, i]: [d1]}
|
||||
assert get_contraction_structure(expr2) == d2
|
||||
|
||||
|
||||
def test_contraction_structure_simple_Pow():
|
||||
x = IndexedBase('x')
|
||||
y = IndexedBase('y')
|
||||
i, j, k = Idx('i'), Idx('j'), Idx('k')
|
||||
ii_jj = x[i, i]**y[j, j]
|
||||
assert get_contraction_structure(ii_jj) == {
|
||||
None: {ii_jj},
|
||||
ii_jj: [
|
||||
{(i,): {x[i, i]}},
|
||||
{(j,): {y[j, j]}}
|
||||
]
|
||||
}
|
||||
|
||||
ii_jk = x[i, i]**y[j, k]
|
||||
assert get_contraction_structure(ii_jk) == {
|
||||
None: {x[i, i]**y[j, k]},
|
||||
x[i, i]**y[j, k]: [
|
||||
{(i,): {x[i, i]}}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def test_contraction_structure_Mul_and_Pow():
|
||||
x = IndexedBase('x')
|
||||
y = IndexedBase('y')
|
||||
i, j, k = Idx('i'), Idx('j'), Idx('k')
|
||||
|
||||
i_ji = x[i]**(y[j]*x[i])
|
||||
assert get_contraction_structure(i_ji) == {None: {i_ji}}
|
||||
ij_i = (x[i]*y[j])**(y[i])
|
||||
assert get_contraction_structure(ij_i) == {None: {ij_i}}
|
||||
j_ij_i = x[j]*(x[i]*y[j])**(y[i])
|
||||
assert get_contraction_structure(j_ij_i) == {(j,): {j_ij_i}}
|
||||
j_i_ji = x[j]*x[i]**(y[j]*x[i])
|
||||
assert get_contraction_structure(j_i_ji) == {(j,): {j_i_ji}}
|
||||
ij_exp_kki = x[i]*y[j]*exp(y[i]*y[k, k])
|
||||
result = get_contraction_structure(ij_exp_kki)
|
||||
expected = {
|
||||
(i,): {ij_exp_kki},
|
||||
ij_exp_kki: [{
|
||||
None: {exp(y[i]*y[k, k])},
|
||||
exp(y[i]*y[k, k]): [{
|
||||
None: {y[i]*y[k, k]},
|
||||
y[i]*y[k, k]: [{(k,): {y[k, k]}}]
|
||||
}]}
|
||||
]
|
||||
}
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_contraction_structure_Add_in_Pow():
|
||||
x = IndexedBase('x')
|
||||
y = IndexedBase('y')
|
||||
i, j, k = Idx('i'), Idx('j'), Idx('k')
|
||||
s_ii_jj_s = (1 + x[i, i])**(1 + y[j, j])
|
||||
expected = {
|
||||
None: {s_ii_jj_s},
|
||||
s_ii_jj_s: [
|
||||
{None: {S.One}, (i,): {x[i, i]}},
|
||||
{None: {S.One}, (j,): {y[j, j]}}
|
||||
]
|
||||
}
|
||||
result = get_contraction_structure(s_ii_jj_s)
|
||||
assert result == expected
|
||||
|
||||
s_ii_jk_s = (1 + x[i, i]) ** (1 + y[j, k])
|
||||
expected_2 = {
|
||||
None: {(x[i, i] + 1)**(y[j, k] + 1)},
|
||||
s_ii_jk_s: [
|
||||
{None: {S.One}, (i,): {x[i, i]}}
|
||||
]
|
||||
}
|
||||
result_2 = get_contraction_structure(s_ii_jk_s)
|
||||
assert result_2 == expected_2
|
||||
|
||||
|
||||
def test_contraction_structure_Pow_in_Pow():
|
||||
x = IndexedBase('x')
|
||||
y = IndexedBase('y')
|
||||
z = IndexedBase('z')
|
||||
i, j, k = Idx('i'), Idx('j'), Idx('k')
|
||||
ii_jj_kk = x[i, i]**y[j, j]**z[k, k]
|
||||
expected = {
|
||||
None: {ii_jj_kk},
|
||||
ii_jj_kk: [
|
||||
{(i,): {x[i, i]}},
|
||||
{
|
||||
None: {y[j, j]**z[k, k]},
|
||||
y[j, j]**z[k, k]: [
|
||||
{(j,): {y[j, j]}},
|
||||
{(k,): {z[k, k]}}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
assert get_contraction_structure(ii_jj_kk) == expected
|
||||
|
||||
|
||||
def test_ufunc_support():
|
||||
f = Function('f')
|
||||
g = Function('g')
|
||||
x = IndexedBase('x')
|
||||
y = IndexedBase('y')
|
||||
i, j = Idx('i'), Idx('j')
|
||||
a = symbols('a')
|
||||
|
||||
assert get_indices(f(x[i])) == ({i}, {})
|
||||
assert get_indices(f(x[i], y[j])) == ({i, j}, {})
|
||||
assert get_indices(f(y[i])*g(x[i])) == (set(), {})
|
||||
assert get_indices(f(a, x[i])) == ({i}, {})
|
||||
assert get_indices(f(a, y[i], x[j])*g(x[i])) == ({j}, {})
|
||||
assert get_indices(g(f(x[i]))) == ({i}, {})
|
||||
|
||||
assert get_contraction_structure(f(x[i])) == {None: {f(x[i])}}
|
||||
assert get_contraction_structure(
|
||||
f(y[i])*g(x[i])) == {(i,): {f(y[i])*g(x[i])}}
|
||||
assert get_contraction_structure(
|
||||
f(y[i])*g(f(x[i]))) == {(i,): {f(y[i])*g(f(x[i]))}}
|
||||
assert get_contraction_structure(
|
||||
f(x[j], y[i])*g(x[i])) == {(i,): {f(x[j], y[i])*g(x[i])}}
|
||||
@@ -0,0 +1,511 @@
|
||||
from sympy.core import symbols, Symbol, Tuple, oo, Dummy
|
||||
from sympy.tensor.indexed import IndexException
|
||||
from sympy.testing.pytest import raises
|
||||
from sympy.utilities.iterables import iterable
|
||||
|
||||
# import test:
|
||||
from sympy.concrete.summations import Sum
|
||||
from sympy.core.function import Function, Subs, Derivative
|
||||
from sympy.core.relational import (StrictLessThan, GreaterThan,
|
||||
StrictGreaterThan, LessThan)
|
||||
from sympy.core.singleton import S
|
||||
from sympy.functions.elementary.exponential import exp, log
|
||||
from sympy.functions.elementary.trigonometric import cos, sin
|
||||
from sympy.functions.special.tensor_functions import KroneckerDelta
|
||||
from sympy.series.order import Order
|
||||
from sympy.sets.fancysets import Range
|
||||
from sympy.tensor.indexed import IndexedBase, Idx, Indexed
|
||||
|
||||
|
||||
def test_Idx_construction():
|
||||
i, a, b = symbols('i a b', integer=True)
|
||||
assert Idx(i) != Idx(i, 1)
|
||||
assert Idx(i, a) == Idx(i, (0, a - 1))
|
||||
assert Idx(i, oo) == Idx(i, (0, oo))
|
||||
|
||||
x = symbols('x', integer=False)
|
||||
raises(TypeError, lambda: Idx(x))
|
||||
raises(TypeError, lambda: Idx(0.5))
|
||||
raises(TypeError, lambda: Idx(i, x))
|
||||
raises(TypeError, lambda: Idx(i, 0.5))
|
||||
raises(TypeError, lambda: Idx(i, (x, 5)))
|
||||
raises(TypeError, lambda: Idx(i, (2, x)))
|
||||
raises(TypeError, lambda: Idx(i, (2, 3.5)))
|
||||
|
||||
|
||||
def test_Idx_properties():
|
||||
i, a, b = symbols('i a b', integer=True)
|
||||
assert Idx(i).is_integer
|
||||
assert Idx(i).name == 'i'
|
||||
assert Idx(i + 2).name == 'i + 2'
|
||||
assert Idx('foo').name == 'foo'
|
||||
|
||||
|
||||
def test_Idx_bounds():
|
||||
i, a, b = symbols('i a b', integer=True)
|
||||
assert Idx(i).lower is None
|
||||
assert Idx(i).upper is None
|
||||
assert Idx(i, a).lower == 0
|
||||
assert Idx(i, a).upper == a - 1
|
||||
assert Idx(i, 5).lower == 0
|
||||
assert Idx(i, 5).upper == 4
|
||||
assert Idx(i, oo).lower == 0
|
||||
assert Idx(i, oo).upper is oo
|
||||
assert Idx(i, (a, b)).lower == a
|
||||
assert Idx(i, (a, b)).upper == b
|
||||
assert Idx(i, (1, 5)).lower == 1
|
||||
assert Idx(i, (1, 5)).upper == 5
|
||||
assert Idx(i, (-oo, oo)).lower is -oo
|
||||
assert Idx(i, (-oo, oo)).upper is oo
|
||||
|
||||
|
||||
def test_Idx_fixed_bounds():
|
||||
i, a, b, x = symbols('i a b x', integer=True)
|
||||
assert Idx(x).lower is None
|
||||
assert Idx(x).upper is None
|
||||
assert Idx(x, a).lower == 0
|
||||
assert Idx(x, a).upper == a - 1
|
||||
assert Idx(x, 5).lower == 0
|
||||
assert Idx(x, 5).upper == 4
|
||||
assert Idx(x, oo).lower == 0
|
||||
assert Idx(x, oo).upper is oo
|
||||
assert Idx(x, (a, b)).lower == a
|
||||
assert Idx(x, (a, b)).upper == b
|
||||
assert Idx(x, (1, 5)).lower == 1
|
||||
assert Idx(x, (1, 5)).upper == 5
|
||||
assert Idx(x, (-oo, oo)).lower is -oo
|
||||
assert Idx(x, (-oo, oo)).upper is oo
|
||||
|
||||
|
||||
def test_Idx_inequalities():
|
||||
i14 = Idx("i14", (1, 4))
|
||||
i79 = Idx("i79", (7, 9))
|
||||
i46 = Idx("i46", (4, 6))
|
||||
i35 = Idx("i35", (3, 5))
|
||||
|
||||
assert i14 <= 5
|
||||
assert i14 < 5
|
||||
assert not (i14 >= 5)
|
||||
assert not (i14 > 5)
|
||||
|
||||
assert 5 >= i14
|
||||
assert 5 > i14
|
||||
assert not (5 <= i14)
|
||||
assert not (5 < i14)
|
||||
|
||||
assert LessThan(i14, 5)
|
||||
assert StrictLessThan(i14, 5)
|
||||
assert not GreaterThan(i14, 5)
|
||||
assert not StrictGreaterThan(i14, 5)
|
||||
|
||||
assert i14 <= 4
|
||||
assert isinstance(i14 < 4, StrictLessThan)
|
||||
assert isinstance(i14 >= 4, GreaterThan)
|
||||
assert not (i14 > 4)
|
||||
|
||||
assert isinstance(i14 <= 1, LessThan)
|
||||
assert not (i14 < 1)
|
||||
assert i14 >= 1
|
||||
assert isinstance(i14 > 1, StrictGreaterThan)
|
||||
|
||||
assert not (i14 <= 0)
|
||||
assert not (i14 < 0)
|
||||
assert i14 >= 0
|
||||
assert i14 > 0
|
||||
|
||||
from sympy.abc import x
|
||||
|
||||
assert isinstance(i14 < x, StrictLessThan)
|
||||
assert isinstance(i14 > x, StrictGreaterThan)
|
||||
assert isinstance(i14 <= x, LessThan)
|
||||
assert isinstance(i14 >= x, GreaterThan)
|
||||
|
||||
assert i14 < i79
|
||||
assert i14 <= i79
|
||||
assert not (i14 > i79)
|
||||
assert not (i14 >= i79)
|
||||
|
||||
assert i14 <= i46
|
||||
assert isinstance(i14 < i46, StrictLessThan)
|
||||
assert isinstance(i14 >= i46, GreaterThan)
|
||||
assert not (i14 > i46)
|
||||
|
||||
assert isinstance(i14 < i35, StrictLessThan)
|
||||
assert isinstance(i14 > i35, StrictGreaterThan)
|
||||
assert isinstance(i14 <= i35, LessThan)
|
||||
assert isinstance(i14 >= i35, GreaterThan)
|
||||
|
||||
iNone1 = Idx("iNone1")
|
||||
iNone2 = Idx("iNone2")
|
||||
|
||||
assert isinstance(iNone1 < iNone2, StrictLessThan)
|
||||
assert isinstance(iNone1 > iNone2, StrictGreaterThan)
|
||||
assert isinstance(iNone1 <= iNone2, LessThan)
|
||||
assert isinstance(iNone1 >= iNone2, GreaterThan)
|
||||
|
||||
|
||||
def test_Idx_inequalities_current_fails():
|
||||
i14 = Idx("i14", (1, 4))
|
||||
|
||||
assert S(5) >= i14
|
||||
assert S(5) > i14
|
||||
assert not (S(5) <= i14)
|
||||
assert not (S(5) < i14)
|
||||
|
||||
|
||||
def test_Idx_func_args():
|
||||
i, a, b = symbols('i a b', integer=True)
|
||||
ii = Idx(i)
|
||||
assert ii.func(*ii.args) == ii
|
||||
ii = Idx(i, a)
|
||||
assert ii.func(*ii.args) == ii
|
||||
ii = Idx(i, (a, b))
|
||||
assert ii.func(*ii.args) == ii
|
||||
|
||||
|
||||
def test_Idx_subs():
|
||||
i, a, b = symbols('i a b', integer=True)
|
||||
assert Idx(i, a).subs(a, b) == Idx(i, b)
|
||||
assert Idx(i, a).subs(i, b) == Idx(b, a)
|
||||
|
||||
assert Idx(i).subs(i, 2) == Idx(2)
|
||||
assert Idx(i, a).subs(a, 2) == Idx(i, 2)
|
||||
assert Idx(i, (a, b)).subs(i, 2) == Idx(2, (a, b))
|
||||
|
||||
|
||||
def test_IndexedBase_sugar():
|
||||
i, j = symbols('i j', integer=True)
|
||||
a = symbols('a')
|
||||
A1 = Indexed(a, i, j)
|
||||
A2 = IndexedBase(a)
|
||||
assert A1 == A2[i, j]
|
||||
assert A1 == A2[(i, j)]
|
||||
assert A1 == A2[[i, j]]
|
||||
assert A1 == A2[Tuple(i, j)]
|
||||
assert all(a.is_Integer for a in A2[1, 0].args[1:])
|
||||
|
||||
|
||||
def test_IndexedBase_subs():
|
||||
i = symbols('i', integer=True)
|
||||
a, b = symbols('a b')
|
||||
A = IndexedBase(a)
|
||||
B = IndexedBase(b)
|
||||
assert A[i] == B[i].subs(b, a)
|
||||
C = {1: 2}
|
||||
assert C[1] == A[1].subs(A, C)
|
||||
|
||||
|
||||
def test_IndexedBase_shape():
|
||||
i, j, m, n = symbols('i j m n', integer=True)
|
||||
a = IndexedBase('a', shape=(m, m))
|
||||
b = IndexedBase('a', shape=(m, n))
|
||||
assert b.shape == Tuple(m, n)
|
||||
assert a[i, j] != b[i, j]
|
||||
assert a[i, j] == b[i, j].subs(n, m)
|
||||
assert b.func(*b.args) == b
|
||||
assert b[i, j].func(*b[i, j].args) == b[i, j]
|
||||
raises(IndexException, lambda: b[i])
|
||||
raises(IndexException, lambda: b[i, i, j])
|
||||
F = IndexedBase("F", shape=m)
|
||||
assert F.shape == Tuple(m)
|
||||
assert F[i].subs(i, j) == F[j]
|
||||
raises(IndexException, lambda: F[i, j])
|
||||
|
||||
|
||||
def test_IndexedBase_assumptions():
|
||||
i = Symbol('i', integer=True)
|
||||
a = Symbol('a')
|
||||
A = IndexedBase(a, positive=True)
|
||||
for c in (A, A[i]):
|
||||
assert c.is_real
|
||||
assert c.is_complex
|
||||
assert not c.is_imaginary
|
||||
assert c.is_nonnegative
|
||||
assert c.is_nonzero
|
||||
assert c.is_commutative
|
||||
assert log(exp(c)) == c
|
||||
|
||||
assert A != IndexedBase(a)
|
||||
assert A == IndexedBase(a, positive=True, real=True)
|
||||
assert A[i] != Indexed(a, i)
|
||||
|
||||
|
||||
def test_IndexedBase_assumptions_inheritance():
|
||||
I = Symbol('I', integer=True)
|
||||
I_inherit = IndexedBase(I)
|
||||
I_explicit = IndexedBase('I', integer=True)
|
||||
|
||||
assert I_inherit.is_integer
|
||||
assert I_explicit.is_integer
|
||||
assert I_inherit.label.is_integer
|
||||
assert I_explicit.label.is_integer
|
||||
assert I_inherit == I_explicit
|
||||
|
||||
|
||||
def test_issue_17652():
|
||||
"""Regression test issue #17652.
|
||||
|
||||
IndexedBase.label should not upcast subclasses of Symbol
|
||||
"""
|
||||
class SubClass(Symbol):
|
||||
pass
|
||||
|
||||
x = SubClass('X')
|
||||
assert type(x) == SubClass
|
||||
base = IndexedBase(x)
|
||||
assert type(x) == SubClass
|
||||
assert type(base.label) == SubClass
|
||||
|
||||
|
||||
def test_Indexed_constructor():
|
||||
i, j = symbols('i j', integer=True)
|
||||
A = Indexed('A', i, j)
|
||||
assert A == Indexed(Symbol('A'), i, j)
|
||||
assert A == Indexed(IndexedBase('A'), i, j)
|
||||
raises(TypeError, lambda: Indexed(A, i, j))
|
||||
raises(IndexException, lambda: Indexed("A"))
|
||||
assert A.free_symbols == {A, A.base.label, i, j}
|
||||
|
||||
|
||||
def test_Indexed_func_args():
|
||||
i, j = symbols('i j', integer=True)
|
||||
a = symbols('a')
|
||||
A = Indexed(a, i, j)
|
||||
assert A == A.func(*A.args)
|
||||
|
||||
|
||||
def test_Indexed_subs():
|
||||
i, j, k = symbols('i j k', integer=True)
|
||||
a, b = symbols('a b')
|
||||
A = IndexedBase(a)
|
||||
B = IndexedBase(b)
|
||||
assert A[i, j] == B[i, j].subs(b, a)
|
||||
assert A[i, j] == A[i, k].subs(k, j)
|
||||
|
||||
|
||||
def test_Indexed_properties():
|
||||
i, j = symbols('i j', integer=True)
|
||||
A = Indexed('A', i, j)
|
||||
assert A.name == 'A[i, j]'
|
||||
assert A.rank == 2
|
||||
assert A.indices == (i, j)
|
||||
assert A.base == IndexedBase('A')
|
||||
assert A.ranges == [None, None]
|
||||
raises(IndexException, lambda: A.shape)
|
||||
|
||||
n, m = symbols('n m', integer=True)
|
||||
assert Indexed('A', Idx(
|
||||
i, m), Idx(j, n)).ranges == [Tuple(0, m - 1), Tuple(0, n - 1)]
|
||||
assert Indexed('A', Idx(i, m), Idx(j, n)).shape == Tuple(m, n)
|
||||
raises(IndexException, lambda: Indexed("A", Idx(i, m), Idx(j)).shape)
|
||||
|
||||
|
||||
def test_Indexed_shape_precedence():
|
||||
i, j = symbols('i j', integer=True)
|
||||
o, p = symbols('o p', integer=True)
|
||||
n, m = symbols('n m', integer=True)
|
||||
a = IndexedBase('a', shape=(o, p))
|
||||
assert a.shape == Tuple(o, p)
|
||||
assert Indexed(
|
||||
a, Idx(i, m), Idx(j, n)).ranges == [Tuple(0, m - 1), Tuple(0, n - 1)]
|
||||
assert Indexed(a, Idx(i, m), Idx(j, n)).shape == Tuple(o, p)
|
||||
assert Indexed(
|
||||
a, Idx(i, m), Idx(j)).ranges == [Tuple(0, m - 1), (None, None)]
|
||||
assert Indexed(a, Idx(i, m), Idx(j)).shape == Tuple(o, p)
|
||||
|
||||
|
||||
def test_complex_indices():
|
||||
i, j = symbols('i j', integer=True)
|
||||
A = Indexed('A', i, i + j)
|
||||
assert A.rank == 2
|
||||
assert A.indices == (i, i + j)
|
||||
|
||||
|
||||
def test_not_interable():
|
||||
i, j = symbols('i j', integer=True)
|
||||
A = Indexed('A', i, i + j)
|
||||
assert not iterable(A)
|
||||
|
||||
|
||||
def test_Indexed_coeff():
|
||||
N = Symbol('N', integer=True)
|
||||
len_y = N
|
||||
i = Idx('i', len_y-1)
|
||||
y = IndexedBase('y', shape=(len_y,))
|
||||
a = (1/y[i+1]*y[i]).coeff(y[i])
|
||||
b = (y[i]/y[i+1]).coeff(y[i])
|
||||
assert a == b
|
||||
|
||||
|
||||
def test_differentiation():
|
||||
from sympy.functions.special.tensor_functions import KroneckerDelta
|
||||
i, j, k, l = symbols('i j k l', cls=Idx)
|
||||
a = symbols('a')
|
||||
m, n = symbols("m, n", integer=True, finite=True)
|
||||
assert m.is_real
|
||||
h, L = symbols('h L', cls=IndexedBase)
|
||||
hi, hj = h[i], h[j]
|
||||
|
||||
expr = hi
|
||||
assert expr.diff(hj) == KroneckerDelta(i, j)
|
||||
assert expr.diff(hi) == KroneckerDelta(i, i)
|
||||
|
||||
expr = S(2) * hi
|
||||
assert expr.diff(hj) == S(2) * KroneckerDelta(i, j)
|
||||
assert expr.diff(hi) == S(2) * KroneckerDelta(i, i)
|
||||
assert expr.diff(a) is S.Zero
|
||||
|
||||
assert Sum(expr, (i, -oo, oo)).diff(hj) == Sum(2*KroneckerDelta(i, j), (i, -oo, oo))
|
||||
assert Sum(expr.diff(hj), (i, -oo, oo)) == Sum(2*KroneckerDelta(i, j), (i, -oo, oo))
|
||||
assert Sum(expr, (i, -oo, oo)).diff(hj).doit() == 2
|
||||
|
||||
assert Sum(expr.diff(hi), (i, -oo, oo)).doit() == Sum(2, (i, -oo, oo)).doit()
|
||||
assert Sum(expr, (i, -oo, oo)).diff(hi).doit() is oo
|
||||
|
||||
expr = a * hj * hj / S(2)
|
||||
assert expr.diff(hi) == a * h[j] * KroneckerDelta(i, j)
|
||||
assert expr.diff(a) == hj * hj / S(2)
|
||||
assert expr.diff(a, 2) is S.Zero
|
||||
|
||||
assert Sum(expr, (i, -oo, oo)).diff(hi) == Sum(a*KroneckerDelta(i, j)*h[j], (i, -oo, oo))
|
||||
assert Sum(expr.diff(hi), (i, -oo, oo)) == Sum(a*KroneckerDelta(i, j)*h[j], (i, -oo, oo))
|
||||
assert Sum(expr, (i, -oo, oo)).diff(hi).doit() == a*h[j]
|
||||
|
||||
assert Sum(expr, (j, -oo, oo)).diff(hi) == Sum(a*KroneckerDelta(i, j)*h[j], (j, -oo, oo))
|
||||
assert Sum(expr.diff(hi), (j, -oo, oo)) == Sum(a*KroneckerDelta(i, j)*h[j], (j, -oo, oo))
|
||||
assert Sum(expr, (j, -oo, oo)).diff(hi).doit() == a*h[i]
|
||||
|
||||
expr = a * sin(hj * hj)
|
||||
assert expr.diff(hi) == 2*a*cos(hj * hj) * hj * KroneckerDelta(i, j)
|
||||
assert expr.diff(hj) == 2*a*cos(hj * hj) * hj
|
||||
|
||||
expr = a * L[i, j] * h[j]
|
||||
assert expr.diff(hi) == a*L[i, j]*KroneckerDelta(i, j)
|
||||
assert expr.diff(hj) == a*L[i, j]
|
||||
assert expr.diff(L[i, j]) == a*h[j]
|
||||
assert expr.diff(L[k, l]) == a*KroneckerDelta(i, k)*KroneckerDelta(j, l)*h[j]
|
||||
assert expr.diff(L[i, l]) == a*KroneckerDelta(j, l)*h[j]
|
||||
|
||||
assert Sum(expr, (j, -oo, oo)).diff(L[k, l]) == Sum(a * KroneckerDelta(i, k) * KroneckerDelta(j, l) * h[j], (j, -oo, oo))
|
||||
assert Sum(expr, (j, -oo, oo)).diff(L[k, l]).doit() == a * KroneckerDelta(i, k) * h[l]
|
||||
|
||||
assert h[m].diff(h[m]) == 1
|
||||
assert h[m].diff(h[n]) == KroneckerDelta(m, n)
|
||||
assert Sum(a*h[m], (m, -oo, oo)).diff(h[n]) == Sum(a*KroneckerDelta(m, n), (m, -oo, oo))
|
||||
assert Sum(a*h[m], (m, -oo, oo)).diff(h[n]).doit() == a
|
||||
assert Sum(a*h[m], (n, -oo, oo)).diff(h[n]) == Sum(a*KroneckerDelta(m, n), (n, -oo, oo))
|
||||
assert Sum(a*h[m], (m, -oo, oo)).diff(h[m]).doit() == oo*a
|
||||
|
||||
|
||||
def test_indexed_series():
|
||||
A = IndexedBase("A")
|
||||
i = symbols("i", integer=True)
|
||||
assert sin(A[i]).series(A[i]) == A[i] - A[i]**3/6 + A[i]**5/120 + Order(A[i]**6, A[i])
|
||||
|
||||
|
||||
def test_indexed_is_constant():
|
||||
A = IndexedBase("A")
|
||||
i, j, k = symbols("i,j,k")
|
||||
assert not A[i].is_constant()
|
||||
assert A[i].is_constant(j)
|
||||
assert not A[1+2*i, k].is_constant()
|
||||
assert not A[1+2*i, k].is_constant(i)
|
||||
assert A[1+2*i, k].is_constant(j)
|
||||
assert not A[1+2*i, k].is_constant(k)
|
||||
|
||||
|
||||
def test_issue_12533():
|
||||
d = IndexedBase('d')
|
||||
assert IndexedBase(range(5)) == Range(0, 5, 1)
|
||||
assert d[0].subs(Symbol("d"), range(5)) == 0
|
||||
assert d[0].subs(d, range(5)) == 0
|
||||
assert d[1].subs(d, range(5)) == 1
|
||||
assert Indexed(Range(5), 2) == 2
|
||||
|
||||
|
||||
def test_issue_12780():
|
||||
n = symbols("n")
|
||||
i = Idx("i", (0, n))
|
||||
raises(TypeError, lambda: i.subs(n, 1.5))
|
||||
|
||||
|
||||
def test_issue_18604():
|
||||
m = symbols("m")
|
||||
assert Idx("i", m).name == 'i'
|
||||
assert Idx("i", m).lower == 0
|
||||
assert Idx("i", m).upper == m - 1
|
||||
m = symbols("m", real=False)
|
||||
raises(TypeError, lambda: Idx("i", m))
|
||||
|
||||
def test_Subs_with_Indexed():
|
||||
A = IndexedBase("A")
|
||||
i, j, k = symbols("i,j,k")
|
||||
x, y, z = symbols("x,y,z")
|
||||
f = Function("f")
|
||||
|
||||
assert Subs(A[i], A[i], A[j]).diff(A[j]) == 1
|
||||
assert Subs(A[i], A[i], x).diff(A[i]) == 0
|
||||
assert Subs(A[i], A[i], x).diff(A[j]) == 0
|
||||
assert Subs(A[i], A[i], x).diff(x) == 1
|
||||
assert Subs(A[i], A[i], x).diff(y) == 0
|
||||
assert Subs(A[i], A[i], A[j]).diff(A[k]) == KroneckerDelta(j, k)
|
||||
assert Subs(x, x, A[i]).diff(A[j]) == KroneckerDelta(i, j)
|
||||
assert Subs(f(A[i]), A[i], x).diff(A[j]) == 0
|
||||
assert Subs(f(A[i]), A[i], A[k]).diff(A[j]) == Derivative(f(A[k]), A[k])*KroneckerDelta(j, k)
|
||||
assert Subs(x, x, A[i]**2).diff(A[j]) == 2*KroneckerDelta(i, j)*A[i]
|
||||
assert Subs(A[i], A[i], A[j]**2).diff(A[k]) == 2*KroneckerDelta(j, k)*A[j]
|
||||
|
||||
assert Subs(A[i]*x, x, A[i]).diff(A[i]) == 2*A[i]
|
||||
assert Subs(A[i]*x, x, A[i]).diff(A[j]) == 2*A[i]*KroneckerDelta(i, j)
|
||||
assert Subs(A[i]*x, x, A[j]).diff(A[i]) == A[j] + A[i]*KroneckerDelta(i, j)
|
||||
assert Subs(A[i]*x, x, A[j]).diff(A[j]) == A[i] + A[j]*KroneckerDelta(i, j)
|
||||
assert Subs(A[i]*x, x, A[i]).diff(A[k]) == 2*A[i]*KroneckerDelta(i, k)
|
||||
assert Subs(A[i]*x, x, A[j]).diff(A[k]) == KroneckerDelta(i, k)*A[j] + KroneckerDelta(j, k)*A[i]
|
||||
|
||||
assert Subs(A[i]*x, A[i], x).diff(A[i]) == 0
|
||||
assert Subs(A[i]*x, A[i], x).diff(A[j]) == 0
|
||||
assert Subs(A[i]*x, A[j], x).diff(A[i]) == x
|
||||
assert Subs(A[i]*x, A[j], x).diff(A[j]) == x*KroneckerDelta(i, j)
|
||||
assert Subs(A[i]*x, A[i], x).diff(A[k]) == 0
|
||||
assert Subs(A[i]*x, A[j], x).diff(A[k]) == x*KroneckerDelta(i, k)
|
||||
|
||||
|
||||
def test_complicated_derivative_with_Indexed():
|
||||
x, y = symbols("x,y", cls=IndexedBase)
|
||||
sigma = symbols("sigma")
|
||||
i, j, k = symbols("i,j,k")
|
||||
m0,m1,m2,m3,m4,m5 = symbols("m0:6")
|
||||
f = Function("f")
|
||||
|
||||
expr = f((x[i] - y[i])**2/sigma)
|
||||
_xi_1 = symbols("xi_1", cls=Dummy)
|
||||
assert expr.diff(x[m0]).dummy_eq(
|
||||
(x[i] - y[i])*KroneckerDelta(i, m0)*\
|
||||
2*Subs(
|
||||
Derivative(f(_xi_1), _xi_1),
|
||||
(_xi_1,),
|
||||
((x[i] - y[i])**2/sigma,)
|
||||
)/sigma
|
||||
)
|
||||
assert expr.diff(x[m0]).diff(x[m1]).dummy_eq(
|
||||
2*KroneckerDelta(i, m0)*\
|
||||
KroneckerDelta(i, m1)*Subs(
|
||||
Derivative(f(_xi_1), _xi_1),
|
||||
(_xi_1,),
|
||||
((x[i] - y[i])**2/sigma,)
|
||||
)/sigma + \
|
||||
4*(x[i] - y[i])**2*KroneckerDelta(i, m0)*KroneckerDelta(i, m1)*\
|
||||
Subs(
|
||||
Derivative(f(_xi_1), _xi_1, _xi_1),
|
||||
(_xi_1,),
|
||||
((x[i] - y[i])**2/sigma,)
|
||||
)/sigma**2
|
||||
)
|
||||
|
||||
|
||||
def test_IndexedBase_commutative():
|
||||
t = IndexedBase('t', commutative=False)
|
||||
u = IndexedBase('u', commutative=False)
|
||||
v = IndexedBase('v')
|
||||
assert t[0]*v[0] == v[0]*t[0]
|
||||
assert t[0]*u[0] != u[0]*t[0]
|
||||
@@ -0,0 +1,13 @@
|
||||
from sympy.tensor.tensor import TensorIndexType, tensor_indices, TensorHead
|
||||
from sympy import I
|
||||
|
||||
def test_printing_TensMul():
|
||||
R3 = TensorIndexType('R3', dim=3)
|
||||
p, q = tensor_indices("p q", R3)
|
||||
K = TensorHead("K", [R3])
|
||||
|
||||
assert repr(2*K(p)) == "2*K(p)"
|
||||
assert repr(-K(p)) == "-K(p)"
|
||||
assert repr(-2*K(p)*K(q)) == "-2*K(p)*K(q)"
|
||||
assert repr(-I*K(p)) == "-I*K(p)"
|
||||
assert repr(I*K(p)) == "I*K(p)"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,31 @@
|
||||
from sympy.tensor.tensor import (Tensor, TensorIndexType, TensorSymmetry,
|
||||
tensor_indices, TensorHead, TensorElement)
|
||||
from sympy.tensor import Array
|
||||
from sympy.core.symbol import Symbol
|
||||
|
||||
|
||||
def test_tensor_element():
|
||||
L = TensorIndexType("L")
|
||||
i, j, k, l, m, n = tensor_indices("i j k l m n", L)
|
||||
|
||||
A = TensorHead("A", [L, L], TensorSymmetry.no_symmetry(2))
|
||||
|
||||
a = A(i, j)
|
||||
|
||||
assert isinstance(TensorElement(a, {}), Tensor)
|
||||
assert isinstance(TensorElement(a, {k: 1}), Tensor)
|
||||
|
||||
te1 = TensorElement(a, {Symbol("i"): 1})
|
||||
assert te1.free == [(j, 0)]
|
||||
assert te1.get_free_indices() == [j]
|
||||
assert te1.dum == []
|
||||
|
||||
te2 = TensorElement(a, {i: 1})
|
||||
assert te2.free == [(j, 0)]
|
||||
assert te2.get_free_indices() == [j]
|
||||
assert te2.dum == []
|
||||
|
||||
assert te1 == te2
|
||||
|
||||
array = Array([[1, 2], [3, 4]])
|
||||
assert te1.replace_with_arrays({A(i, j): array}, [j]) == array[1, :]
|
||||
@@ -0,0 +1,464 @@
|
||||
from sympy import sin, cos
|
||||
from sympy.testing.pytest import raises
|
||||
|
||||
from sympy.tensor.toperators import PartialDerivative
|
||||
from sympy.tensor.tensor import (TensorIndexType,
|
||||
tensor_indices,
|
||||
TensorHead, tensor_heads)
|
||||
from sympy.core.numbers import Rational
|
||||
from sympy.core.symbol import symbols
|
||||
from sympy.matrices.dense import diag
|
||||
from sympy.tensor.array import Array
|
||||
|
||||
from sympy.core.random import randint
|
||||
|
||||
|
||||
L = TensorIndexType("L")
|
||||
i, j, k, m, m1, m2, m3, m4 = tensor_indices("i j k m m1 m2 m3 m4", L)
|
||||
i0 = tensor_indices("i0", L)
|
||||
L_0, L_1 = tensor_indices("L_0 L_1", L)
|
||||
|
||||
A, B, C, D = tensor_heads("A B C D", [L])
|
||||
|
||||
H = TensorHead("H", [L, L])
|
||||
|
||||
|
||||
def test_invalid_partial_derivative_valence():
|
||||
raises(ValueError, lambda: PartialDerivative(C(j), D(-j)))
|
||||
raises(ValueError, lambda: PartialDerivative(C(-j), D(j)))
|
||||
|
||||
|
||||
def test_tensor_partial_deriv():
|
||||
# Test flatten:
|
||||
expr = PartialDerivative(PartialDerivative(A(i), A(j)), A(k))
|
||||
assert expr == PartialDerivative(A(i), A(j), A(k))
|
||||
assert expr.expr == A(i)
|
||||
assert expr.variables == (A(j), A(k))
|
||||
assert expr.get_indices() == [i, -j, -k]
|
||||
assert expr.get_free_indices() == [i, -j, -k]
|
||||
|
||||
expr = PartialDerivative(PartialDerivative(A(i), A(j)), A(i))
|
||||
assert expr.expr == A(L_0)
|
||||
assert expr.variables == (A(j), A(L_0))
|
||||
|
||||
expr1 = PartialDerivative(A(i), A(j))
|
||||
assert expr1.expr == A(i)
|
||||
assert expr1.variables == (A(j),)
|
||||
|
||||
expr2 = A(i)*PartialDerivative(H(k, -i), A(j))
|
||||
assert expr2.get_indices() == [L_0, k, -L_0, -j]
|
||||
|
||||
expr2b = A(i)*PartialDerivative(H(k, -i), A(-j))
|
||||
assert expr2b.get_indices() == [L_0, k, -L_0, j]
|
||||
|
||||
expr3 = A(i)*PartialDerivative(B(k)*C(-i) + 3*H(k, -i), A(j))
|
||||
assert expr3.get_indices() == [L_0, k, -L_0, -j]
|
||||
|
||||
expr4 = (A(i) + B(i))*PartialDerivative(C(j), D(j))
|
||||
assert expr4.get_indices() == [i, L_0, -L_0]
|
||||
|
||||
expr4b = (A(i) + B(i))*PartialDerivative(C(-j), D(-j))
|
||||
assert expr4b.get_indices() == [i, -L_0, L_0]
|
||||
|
||||
expr5 = (A(i) + B(i))*PartialDerivative(C(-i), D(j))
|
||||
assert expr5.get_indices() == [L_0, -L_0, -j]
|
||||
|
||||
|
||||
def test_replace_arrays_partial_derivative():
|
||||
|
||||
x, y, z, t = symbols("x y z t")
|
||||
|
||||
expr = PartialDerivative(A(i), B(j))
|
||||
repl = expr.replace_with_arrays({A(i): [sin(x)*cos(y), x**3*y**2], B(i): [x, y]})
|
||||
assert repl == Array([[cos(x)*cos(y), -sin(x)*sin(y)], [3*x**2*y**2, 2*x**3*y]])
|
||||
repl = expr.replace_with_arrays({A(i): [sin(x)*cos(y), x**3*y**2], B(i): [x, y]}, [-j, i])
|
||||
assert repl == Array([[cos(x)*cos(y), 3*x**2*y**2], [-sin(x)*sin(y), 2*x**3*y]])
|
||||
|
||||
# d(A^i)/d(A_j) = d(g^ik A_k)/d(A_j) = g^ik delta_jk
|
||||
expr = PartialDerivative(A(i), A(-j))
|
||||
assert expr.get_free_indices() == [i, j]
|
||||
assert expr.get_indices() == [i, j]
|
||||
assert expr.replace_with_arrays({A(i): [x, y], L: diag(1, 1)}, [i, j]) == Array([[1, 0], [0, 1]])
|
||||
assert expr.replace_with_arrays({A(i): [x, y], L: diag(1, -1)}, [i, j]) == Array([[1, 0], [0, -1]])
|
||||
assert expr.replace_with_arrays({A(-i): [x, y], L: diag(1, 1)}, [i, j]) == Array([[1, 0], [0, 1]])
|
||||
assert expr.replace_with_arrays({A(-i): [x, y], L: diag(1, -1)}, [i, j]) == Array([[1, 0], [0, -1]])
|
||||
|
||||
expr = PartialDerivative(A(i), A(j))
|
||||
assert expr.get_free_indices() == [i, -j]
|
||||
assert expr.get_indices() == [i, -j]
|
||||
assert expr.replace_with_arrays({A(i): [x, y]}, [i, -j]) == Array([[1, 0], [0, 1]])
|
||||
assert expr.replace_with_arrays({A(i): [x, y], L: diag(1, 1)}, [i, -j]) == Array([[1, 0], [0, 1]])
|
||||
assert expr.replace_with_arrays({A(i): [x, y], L: diag(1, -1)}, [i, -j]) == Array([[1, 0], [0, 1]])
|
||||
assert expr.replace_with_arrays({A(-i): [x, y], L: diag(1, 1)}, [i, -j]) == Array([[1, 0], [0, 1]])
|
||||
assert expr.replace_with_arrays({A(-i): [x, y], L: diag(1, -1)}, [i, -j]) == Array([[1, 0], [0, 1]])
|
||||
|
||||
expr = PartialDerivative(A(-i), A(-j))
|
||||
assert expr.get_free_indices() == [-i, j]
|
||||
assert expr.get_indices() == [-i, j]
|
||||
assert expr.replace_with_arrays({A(-i): [x, y]}, [-i, j]) == Array([[1, 0], [0, 1]])
|
||||
assert expr.replace_with_arrays({A(-i): [x, y], L: diag(1, 1)}, [-i, j]) == Array([[1, 0], [0, 1]])
|
||||
assert expr.replace_with_arrays({A(-i): [x, y], L: diag(1, -1)}, [-i, j]) == Array([[1, 0], [0, 1]])
|
||||
assert expr.replace_with_arrays({A(i): [x, y], L: diag(1, 1)}, [-i, j]) == Array([[1, 0], [0, 1]])
|
||||
assert expr.replace_with_arrays({A(i): [x, y], L: diag(1, -1)}, [-i, j]) == Array([[1, 0], [0, 1]])
|
||||
|
||||
expr = PartialDerivative(A(i), A(i))
|
||||
assert expr.get_free_indices() == []
|
||||
assert expr.get_indices() == [L_0, -L_0]
|
||||
assert expr.replace_with_arrays({A(i): [x, y], L: diag(1, 1)}, []) == 2
|
||||
assert expr.replace_with_arrays({A(i): [x, y], L: diag(1, -1)}, []) == 2
|
||||
|
||||
expr = PartialDerivative(A(-i), A(-i))
|
||||
assert expr.get_free_indices() == []
|
||||
assert expr.get_indices() == [-L_0, L_0]
|
||||
assert expr.replace_with_arrays({A(i): [x, y], L: diag(1, 1)}, []) == 2
|
||||
assert expr.replace_with_arrays({A(i): [x, y], L: diag(1, -1)}, []) == 2
|
||||
|
||||
expr = PartialDerivative(H(i, j) + H(j, i), A(i))
|
||||
assert expr.get_indices() == [L_0, j, -L_0]
|
||||
assert expr.get_free_indices() == [j]
|
||||
|
||||
expr = PartialDerivative(H(i, j) + H(j, i), A(k))*B(-i)
|
||||
assert expr.get_indices() == [L_0, j, -k, -L_0]
|
||||
assert expr.get_free_indices() == [j, -k]
|
||||
|
||||
expr = PartialDerivative(A(i)*(H(-i, j) + H(j, -i)), A(j))
|
||||
assert expr.get_indices() == [L_0, -L_0, L_1, -L_1]
|
||||
assert expr.get_free_indices() == []
|
||||
|
||||
expr = A(j)*A(-j) + expr
|
||||
assert expr.get_indices() == [L_0, -L_0, L_1, -L_1]
|
||||
assert expr.get_free_indices() == []
|
||||
|
||||
expr = A(i)*(B(j)*PartialDerivative(C(-j), D(i)) + C(j)*PartialDerivative(D(-j), B(i)))
|
||||
assert expr.get_indices() == [L_0, L_1, -L_1, -L_0]
|
||||
assert expr.get_free_indices() == []
|
||||
|
||||
expr = A(i)*PartialDerivative(C(-j), D(i))
|
||||
assert expr.get_indices() == [L_0, -j, -L_0]
|
||||
assert expr.get_free_indices() == [-j]
|
||||
|
||||
|
||||
def test_expand_partial_derivative_sum_rule():
|
||||
tau = symbols("tau")
|
||||
|
||||
# check sum rule for D(tensor, symbol)
|
||||
expr1aa = PartialDerivative(A(i), tau)
|
||||
|
||||
assert expr1aa._expand_partial_derivative() == PartialDerivative(A(i), tau)
|
||||
|
||||
expr1ab = PartialDerivative(A(i) + B(i), tau)
|
||||
|
||||
assert (expr1ab._expand_partial_derivative() ==
|
||||
PartialDerivative(A(i), tau) +
|
||||
PartialDerivative(B(i), tau))
|
||||
|
||||
expr1ac = PartialDerivative(A(i) + B(i) + C(i), tau)
|
||||
|
||||
assert (expr1ac._expand_partial_derivative() ==
|
||||
PartialDerivative(A(i), tau) +
|
||||
PartialDerivative(B(i), tau) +
|
||||
PartialDerivative(C(i), tau))
|
||||
|
||||
# check sum rule for D(tensor, D(j))
|
||||
expr1ba = PartialDerivative(A(i), D(j))
|
||||
|
||||
assert expr1ba._expand_partial_derivative() ==\
|
||||
PartialDerivative(A(i), D(j))
|
||||
expr1bb = PartialDerivative(A(i) + B(i), D(j))
|
||||
|
||||
assert (expr1bb._expand_partial_derivative() ==
|
||||
PartialDerivative(A(i), D(j)) +
|
||||
PartialDerivative(B(i), D(j)))
|
||||
|
||||
expr1bc = PartialDerivative(A(i) + B(i) + C(i), D(j))
|
||||
assert expr1bc._expand_partial_derivative() ==\
|
||||
PartialDerivative(A(i), D(j))\
|
||||
+ PartialDerivative(B(i), D(j))\
|
||||
+ PartialDerivative(C(i), D(j))
|
||||
|
||||
# check sum rule for D(tensor, H(j, k))
|
||||
expr1ca = PartialDerivative(A(i), H(j, k))
|
||||
assert expr1ca._expand_partial_derivative() ==\
|
||||
PartialDerivative(A(i), H(j, k))
|
||||
expr1cb = PartialDerivative(A(i) + B(i), H(j, k))
|
||||
assert (expr1cb._expand_partial_derivative() ==
|
||||
PartialDerivative(A(i), H(j, k))
|
||||
+ PartialDerivative(B(i), H(j, k)))
|
||||
expr1cc = PartialDerivative(A(i) + B(i) + C(i), H(j, k))
|
||||
assert (expr1cc._expand_partial_derivative() ==
|
||||
PartialDerivative(A(i), H(j, k))
|
||||
+ PartialDerivative(B(i), H(j, k))
|
||||
+ PartialDerivative(C(i), H(j, k)))
|
||||
|
||||
# check sum rule for D(D(tensor, D(j)), H(k, m))
|
||||
expr1da = PartialDerivative(A(i), (D(j), H(k, m)))
|
||||
assert expr1da._expand_partial_derivative() ==\
|
||||
PartialDerivative(A(i), (D(j), H(k, m)))
|
||||
expr1db = PartialDerivative(A(i) + B(i), (D(j), H(k, m)))
|
||||
assert expr1db._expand_partial_derivative() ==\
|
||||
PartialDerivative(A(i), (D(j), H(k, m)))\
|
||||
+ PartialDerivative(B(i), (D(j), H(k, m)))
|
||||
expr1dc = PartialDerivative(A(i) + B(i) + C(i), (D(j), H(k, m)))
|
||||
assert expr1dc._expand_partial_derivative() ==\
|
||||
PartialDerivative(A(i), (D(j), H(k, m)))\
|
||||
+ PartialDerivative(B(i), (D(j), H(k, m)))\
|
||||
+ PartialDerivative(C(i), (D(j), H(k, m)))
|
||||
|
||||
|
||||
def test_expand_partial_derivative_constant_factor_rule():
|
||||
nneg = randint(0, 1000)
|
||||
pos = randint(1, 1000)
|
||||
neg = -randint(1, 1000)
|
||||
|
||||
c1 = Rational(nneg, pos)
|
||||
c2 = Rational(neg, pos)
|
||||
c3 = Rational(nneg, neg)
|
||||
|
||||
expr2a = PartialDerivative(nneg*A(i), D(j))
|
||||
assert expr2a._expand_partial_derivative() ==\
|
||||
nneg*PartialDerivative(A(i), D(j))
|
||||
|
||||
expr2b = PartialDerivative(neg*A(i), D(j))
|
||||
assert expr2b._expand_partial_derivative() ==\
|
||||
neg*PartialDerivative(A(i), D(j))
|
||||
|
||||
expr2ca = PartialDerivative(c1*A(i), D(j))
|
||||
assert expr2ca._expand_partial_derivative() ==\
|
||||
c1*PartialDerivative(A(i), D(j))
|
||||
|
||||
expr2cb = PartialDerivative(c2*A(i), D(j))
|
||||
assert expr2cb._expand_partial_derivative() ==\
|
||||
c2*PartialDerivative(A(i), D(j))
|
||||
|
||||
expr2cc = PartialDerivative(c3*A(i), D(j))
|
||||
assert expr2cc._expand_partial_derivative() ==\
|
||||
c3*PartialDerivative(A(i), D(j))
|
||||
|
||||
|
||||
def test_expand_partial_derivative_full_linearity():
|
||||
nneg = randint(0, 1000)
|
||||
pos = randint(1, 1000)
|
||||
neg = -randint(1, 1000)
|
||||
|
||||
c1 = Rational(nneg, pos)
|
||||
c2 = Rational(neg, pos)
|
||||
c3 = Rational(nneg, neg)
|
||||
|
||||
# check full linearity
|
||||
p = PartialDerivative(42, D(j))
|
||||
assert p and not p._expand_partial_derivative()
|
||||
|
||||
expr3a = PartialDerivative(nneg*A(i) + pos*B(i), D(j))
|
||||
assert expr3a._expand_partial_derivative() ==\
|
||||
nneg*PartialDerivative(A(i), D(j))\
|
||||
+ pos*PartialDerivative(B(i), D(j))
|
||||
|
||||
expr3b = PartialDerivative(nneg*A(i) + neg*B(i), D(j))
|
||||
assert expr3b._expand_partial_derivative() ==\
|
||||
nneg*PartialDerivative(A(i), D(j))\
|
||||
+ neg*PartialDerivative(B(i), D(j))
|
||||
|
||||
expr3c = PartialDerivative(neg*A(i) + pos*B(i), D(j))
|
||||
assert expr3c._expand_partial_derivative() ==\
|
||||
neg*PartialDerivative(A(i), D(j))\
|
||||
+ pos*PartialDerivative(B(i), D(j))
|
||||
|
||||
expr3d = PartialDerivative(c1*A(i) + c2*B(i), D(j))
|
||||
assert expr3d._expand_partial_derivative() ==\
|
||||
c1*PartialDerivative(A(i), D(j))\
|
||||
+ c2*PartialDerivative(B(i), D(j))
|
||||
|
||||
expr3e = PartialDerivative(c2*A(i) + c1*B(i), D(j))
|
||||
assert expr3e._expand_partial_derivative() ==\
|
||||
c2*PartialDerivative(A(i), D(j))\
|
||||
+ c1*PartialDerivative(B(i), D(j))
|
||||
|
||||
expr3f = PartialDerivative(c2*A(i) + c3*B(i), D(j))
|
||||
assert expr3f._expand_partial_derivative() ==\
|
||||
c2*PartialDerivative(A(i), D(j))\
|
||||
+ c3*PartialDerivative(B(i), D(j))
|
||||
|
||||
expr3g = PartialDerivative(c3*A(i) + c2*B(i), D(j))
|
||||
assert expr3g._expand_partial_derivative() ==\
|
||||
c3*PartialDerivative(A(i), D(j))\
|
||||
+ c2*PartialDerivative(B(i), D(j))
|
||||
|
||||
expr3h = PartialDerivative(c3*A(i) + c1*B(i), D(j))
|
||||
assert expr3h._expand_partial_derivative() ==\
|
||||
c3*PartialDerivative(A(i), D(j))\
|
||||
+ c1*PartialDerivative(B(i), D(j))
|
||||
|
||||
expr3i = PartialDerivative(c1*A(i) + c3*B(i), D(j))
|
||||
assert expr3i._expand_partial_derivative() ==\
|
||||
c1*PartialDerivative(A(i), D(j))\
|
||||
+ c3*PartialDerivative(B(i), D(j))
|
||||
|
||||
|
||||
def test_expand_partial_derivative_product_rule():
|
||||
# check product rule
|
||||
expr4a = PartialDerivative(A(i)*B(j), D(k))
|
||||
|
||||
assert expr4a._expand_partial_derivative() == \
|
||||
PartialDerivative(A(i), D(k))*B(j)\
|
||||
+ A(i)*PartialDerivative(B(j), D(k))
|
||||
|
||||
expr4b = PartialDerivative(A(i)*B(j)*C(k), D(m))
|
||||
assert expr4b._expand_partial_derivative() ==\
|
||||
PartialDerivative(A(i), D(m))*B(j)*C(k)\
|
||||
+ A(i)*PartialDerivative(B(j), D(m))*C(k)\
|
||||
+ A(i)*B(j)*PartialDerivative(C(k), D(m))
|
||||
|
||||
expr4c = PartialDerivative(A(i)*B(j), C(k), D(m))
|
||||
assert expr4c._expand_partial_derivative() ==\
|
||||
PartialDerivative(A(i), C(k), D(m))*B(j) \
|
||||
+ PartialDerivative(A(i), C(k))*PartialDerivative(B(j), D(m))\
|
||||
+ PartialDerivative(A(i), D(m))*PartialDerivative(B(j), C(k))\
|
||||
+ A(i)*PartialDerivative(B(j), C(k), D(m))
|
||||
|
||||
|
||||
def test_eval_partial_derivative_expr_by_symbol():
|
||||
|
||||
tau, alpha = symbols("tau alpha")
|
||||
|
||||
expr1 = PartialDerivative(tau**alpha, tau)
|
||||
assert expr1._perform_derivative() == alpha * 1 / tau * tau ** alpha
|
||||
|
||||
expr2 = PartialDerivative(2*tau + 3*tau**4, tau)
|
||||
assert expr2._perform_derivative() == 2 + 12 * tau ** 3
|
||||
|
||||
expr3 = PartialDerivative(2*tau + 3*tau**4, alpha)
|
||||
assert expr3._perform_derivative() == 0
|
||||
|
||||
|
||||
def test_eval_partial_derivative_single_tensors_by_scalar():
|
||||
|
||||
tau, mu = symbols("tau mu")
|
||||
|
||||
expr = PartialDerivative(tau**mu, tau)
|
||||
assert expr._perform_derivative() == mu*tau**mu/tau
|
||||
|
||||
expr1a = PartialDerivative(A(i), tau)
|
||||
assert expr1a._perform_derivative() == 0
|
||||
|
||||
expr1b = PartialDerivative(A(-i), tau)
|
||||
assert expr1b._perform_derivative() == 0
|
||||
|
||||
expr2a = PartialDerivative(H(i, j), tau)
|
||||
assert expr2a._perform_derivative() == 0
|
||||
|
||||
expr2b = PartialDerivative(H(i, -j), tau)
|
||||
assert expr2b._perform_derivative() == 0
|
||||
|
||||
expr2c = PartialDerivative(H(-i, j), tau)
|
||||
assert expr2c._perform_derivative() == 0
|
||||
|
||||
expr2d = PartialDerivative(H(-i, -j), tau)
|
||||
assert expr2d._perform_derivative() == 0
|
||||
|
||||
|
||||
def test_eval_partial_derivative_single_1st_rank_tensors_by_tensor():
|
||||
|
||||
expr1 = PartialDerivative(A(i), A(j))
|
||||
assert expr1._perform_derivative() - L.delta(i, -j) == 0
|
||||
|
||||
expr2 = PartialDerivative(A(i), A(-j))
|
||||
assert expr2._perform_derivative() - L.metric(i, L_0) * L.delta(-L_0, j) == 0
|
||||
|
||||
expr3 = PartialDerivative(A(-i), A(-j))
|
||||
assert expr3._perform_derivative() - L.delta(-i, j) == 0
|
||||
|
||||
expr4 = PartialDerivative(A(-i), A(j))
|
||||
assert expr4._perform_derivative() - L.metric(-i, -L_0) * L.delta(L_0, -j) == 0
|
||||
|
||||
expr5 = PartialDerivative(A(i), B(j))
|
||||
expr6 = PartialDerivative(A(i), C(j))
|
||||
expr7 = PartialDerivative(A(i), D(j))
|
||||
expr8 = PartialDerivative(A(i), H(j, k))
|
||||
assert expr5._perform_derivative() == 0
|
||||
assert expr6._perform_derivative() == 0
|
||||
assert expr7._perform_derivative() == 0
|
||||
assert expr8._perform_derivative() == 0
|
||||
|
||||
expr9 = PartialDerivative(A(i), A(i))
|
||||
assert expr9._perform_derivative() - L.delta(L_0, -L_0) == 0
|
||||
|
||||
expr10 = PartialDerivative(A(-i), A(-i))
|
||||
assert expr10._perform_derivative() - L.delta(-L_0, L_0) == 0
|
||||
|
||||
|
||||
def test_eval_partial_derivative_single_2nd_rank_tensors_by_tensor():
|
||||
|
||||
expr1 = PartialDerivative(H(i, j), H(m, m1))
|
||||
assert expr1._perform_derivative() - L.delta(i, -m) * L.delta(j, -m1) == 0
|
||||
|
||||
expr2 = PartialDerivative(H(i, j), H(-m, m1))
|
||||
assert expr2._perform_derivative() - L.metric(i, L_0) * L.delta(-L_0, m) * L.delta(j, -m1) == 0
|
||||
|
||||
expr3 = PartialDerivative(H(i, j), H(m, -m1))
|
||||
assert expr3._perform_derivative() - L.delta(i, -m) * L.metric(j, L_0) * L.delta(-L_0, m1) == 0
|
||||
|
||||
expr4 = PartialDerivative(H(i, j), H(-m, -m1))
|
||||
assert expr4._perform_derivative() - L.metric(i, L_0) * L.delta(-L_0, m) * L.metric(j, L_1) * L.delta(-L_1, m1) == 0
|
||||
|
||||
def test_eval_partial_derivative_divergence_type():
|
||||
expr1a = PartialDerivative(A(i), A(i))
|
||||
expr1b = PartialDerivative(A(i), A(k))
|
||||
expr1c = PartialDerivative(L.delta(-i, k) * A(i), A(k))
|
||||
|
||||
assert (expr1a._perform_derivative()
|
||||
- (L.delta(-i, k) * expr1b._perform_derivative())).contract_delta(L.delta) == 0
|
||||
|
||||
assert (expr1a._perform_derivative()
|
||||
- expr1c._perform_derivative()).contract_delta(L.delta) == 0
|
||||
|
||||
expr2a = PartialDerivative(H(i, j), H(i, j))
|
||||
expr2b = PartialDerivative(H(i, j), H(k, m))
|
||||
expr2c = PartialDerivative(L.delta(-i, k) * L.delta(-j, m) * H(i, j), H(k, m))
|
||||
|
||||
assert (expr2a._perform_derivative()
|
||||
- (L.delta(-i, k) * L.delta(-j, m) * expr2b._perform_derivative())).contract_delta(L.delta) == 0
|
||||
|
||||
assert (expr2a._perform_derivative()
|
||||
- expr2c._perform_derivative()).contract_delta(L.delta) == 0
|
||||
|
||||
|
||||
def test_eval_partial_derivative_expr1():
|
||||
|
||||
tau, alpha = symbols("tau alpha")
|
||||
|
||||
# this is only some special expression
|
||||
# tested: vector derivative
|
||||
# tested: scalar derivative
|
||||
# tested: tensor derivative
|
||||
base_expr1 = A(i)*H(-i, j) + A(i)*A(-i)*A(j) + tau**alpha*A(j)
|
||||
|
||||
tensor_derivative = PartialDerivative(base_expr1, H(k, m))._perform_derivative()
|
||||
vector_derivative = PartialDerivative(base_expr1, A(k))._perform_derivative()
|
||||
scalar_derivative = PartialDerivative(base_expr1, tau)._perform_derivative()
|
||||
|
||||
assert (tensor_derivative - A(L_0)*L.metric(-L_0, -L_1)*L.delta(L_1, -k)*L.delta(j, -m)) == 0
|
||||
|
||||
assert (vector_derivative - (tau**alpha*L.delta(j, -k) +
|
||||
L.delta(L_0, -k)*A(-L_0)*A(j) +
|
||||
A(L_0)*L.metric(-L_0, -L_1)*L.delta(L_1, -k)*A(j) +
|
||||
A(L_0)*A(-L_0)*L.delta(j, -k) +
|
||||
L.delta(L_0, -k)*H(-L_0, j))).expand().doit() == 0
|
||||
|
||||
assert (vector_derivative.contract_metric(L.metric).contract_delta(L.delta) -
|
||||
(tau**alpha*L.delta(j, -k) + A(L_0)*A(-L_0)*L.delta(j, -k) + H(-k, j) + 2*A(j)*A(-k))).expand().doit() == 0
|
||||
|
||||
assert scalar_derivative - alpha*1/tau*tau**alpha*A(j) == 0
|
||||
|
||||
|
||||
def test_eval_partial_derivative_mixed_scalar_tensor_expr2():
|
||||
|
||||
tau, alpha = symbols("tau alpha")
|
||||
|
||||
base_expr2 = A(i)*A(-i) + tau**2
|
||||
|
||||
vector_expression = PartialDerivative(base_expr2, A(k))._perform_derivative()
|
||||
assert (vector_expression -
|
||||
(L.delta(L_0, -k)*A(-L_0) + A(L_0)*L.metric(-L_0, -L_1)*L.delta(L_1, -k))).expand().doit() == 0
|
||||
|
||||
scalar_expression = PartialDerivative(base_expr2, tau)._perform_derivative()
|
||||
assert scalar_expression == 2*tau
|
||||
@@ -0,0 +1,256 @@
|
||||
from sympy import permutedims
|
||||
from sympy.core.numbers import Number
|
||||
from sympy.core.singleton import S
|
||||
from sympy.core.symbol import Symbol
|
||||
from sympy.core.sympify import sympify
|
||||
from sympy.tensor.tensor import Tensor, TensExpr, TensAdd, TensMul
|
||||
|
||||
|
||||
class PartialDerivative(TensExpr):
|
||||
"""
|
||||
Partial derivative for tensor expressions.
|
||||
|
||||
Examples
|
||||
========
|
||||
|
||||
>>> from sympy.tensor.tensor import TensorIndexType, TensorHead
|
||||
>>> from sympy.tensor.toperators import PartialDerivative
|
||||
>>> from sympy import symbols
|
||||
>>> L = TensorIndexType("L")
|
||||
>>> A = TensorHead("A", [L])
|
||||
>>> B = TensorHead("B", [L])
|
||||
>>> i, j, k = symbols("i j k")
|
||||
|
||||
>>> expr = PartialDerivative(A(i), A(j))
|
||||
>>> expr
|
||||
PartialDerivative(A(i), A(j))
|
||||
|
||||
The ``PartialDerivative`` object behaves like a tensorial expression:
|
||||
|
||||
>>> expr.get_indices()
|
||||
[i, -j]
|
||||
|
||||
Notice that the deriving variables have opposite valence than the
|
||||
printed one: ``A(j)`` is printed as covariant, but the index of the
|
||||
derivative is actually contravariant, i.e. ``-j``.
|
||||
|
||||
Indices can be contracted:
|
||||
|
||||
>>> expr = PartialDerivative(A(i), A(i))
|
||||
>>> expr
|
||||
PartialDerivative(A(L_0), A(L_0))
|
||||
>>> expr.get_indices()
|
||||
[L_0, -L_0]
|
||||
|
||||
The method ``.get_indices()`` always returns all indices (even the
|
||||
contracted ones). If only uncontracted indices are needed, call
|
||||
``.get_free_indices()``:
|
||||
|
||||
>>> expr.get_free_indices()
|
||||
[]
|
||||
|
||||
Nested partial derivatives are flattened:
|
||||
|
||||
>>> expr = PartialDerivative(PartialDerivative(A(i), A(j)), A(k))
|
||||
>>> expr
|
||||
PartialDerivative(A(i), A(j), A(k))
|
||||
>>> expr.get_indices()
|
||||
[i, -j, -k]
|
||||
|
||||
Replace a derivative with array values:
|
||||
|
||||
>>> from sympy.abc import x, y
|
||||
>>> from sympy import sin, log
|
||||
>>> compA = [sin(x), log(x)*y**3]
|
||||
>>> compB = [x, y]
|
||||
>>> expr = PartialDerivative(A(i), B(j))
|
||||
>>> expr.replace_with_arrays({A(i): compA, B(i): compB})
|
||||
[[cos(x), 0], [y**3/x, 3*y**2*log(x)]]
|
||||
|
||||
The returned array is indexed by `(i, -j)`.
|
||||
|
||||
Be careful that other SymPy modules put the indices of the deriving
|
||||
variables before the indices of the derivand in the derivative result.
|
||||
For example:
|
||||
|
||||
>>> expr.get_free_indices()
|
||||
[i, -j]
|
||||
|
||||
>>> from sympy import Matrix, Array
|
||||
>>> Matrix(compA).diff(Matrix(compB)).reshape(2, 2)
|
||||
[[cos(x), y**3/x], [0, 3*y**2*log(x)]]
|
||||
>>> Array(compA).diff(Array(compB))
|
||||
[[cos(x), y**3/x], [0, 3*y**2*log(x)]]
|
||||
|
||||
These are the transpose of the result of ``PartialDerivative``,
|
||||
as the matrix and the array modules put the index `-j` before `i` in the
|
||||
derivative result. An array read with index order `(-j, i)` is indeed the
|
||||
transpose of the same array read with index order `(i, -j)`. By specifying
|
||||
the index order to ``.replace_with_arrays`` one can get a compatible
|
||||
expression:
|
||||
|
||||
>>> expr.replace_with_arrays({A(i): compA, B(i): compB}, [-j, i])
|
||||
[[cos(x), y**3/x], [0, 3*y**2*log(x)]]
|
||||
"""
|
||||
|
||||
def __new__(cls, expr, *variables):
|
||||
|
||||
# Flatten:
|
||||
if isinstance(expr, PartialDerivative):
|
||||
variables = expr.variables + variables
|
||||
expr = expr.expr
|
||||
|
||||
args, indices, free, dum = cls._contract_indices_for_derivative(
|
||||
S(expr), variables)
|
||||
|
||||
obj = TensExpr.__new__(cls, *args)
|
||||
|
||||
obj._indices = indices
|
||||
obj._free = free
|
||||
obj._dum = dum
|
||||
return obj
|
||||
|
||||
@property
|
||||
def coeff(self):
|
||||
return S.One
|
||||
|
||||
@property
|
||||
def nocoeff(self):
|
||||
return self
|
||||
|
||||
@classmethod
|
||||
def _contract_indices_for_derivative(cls, expr, variables):
|
||||
variables_opposite_valence = []
|
||||
|
||||
for i in variables:
|
||||
if isinstance(i, Tensor):
|
||||
i_free_indices = i.get_free_indices()
|
||||
variables_opposite_valence.append(
|
||||
i.xreplace({k: -k for k in i_free_indices}))
|
||||
elif isinstance(i, Symbol):
|
||||
variables_opposite_valence.append(i)
|
||||
|
||||
args, indices, free, dum = TensMul._tensMul_contract_indices(
|
||||
[expr] + variables_opposite_valence, replace_indices=True)
|
||||
|
||||
for i in range(1, len(args)):
|
||||
args_i = args[i]
|
||||
if isinstance(args_i, Tensor):
|
||||
i_indices = args[i].get_free_indices()
|
||||
args[i] = args[i].xreplace({k: -k for k in i_indices})
|
||||
|
||||
return args, indices, free, dum
|
||||
|
||||
def doit(self, **hints):
|
||||
args, indices, free, dum = self._contract_indices_for_derivative(self.expr, self.variables)
|
||||
|
||||
obj = self.func(*args)
|
||||
obj._indices = indices
|
||||
obj._free = free
|
||||
obj._dum = dum
|
||||
|
||||
return obj
|
||||
|
||||
def _expand_partial_derivative(self):
|
||||
args, indices, free, dum = self._contract_indices_for_derivative(self.expr, self.variables)
|
||||
|
||||
obj = self.func(*args)
|
||||
obj._indices = indices
|
||||
obj._free = free
|
||||
obj._dum = dum
|
||||
|
||||
result = obj
|
||||
|
||||
if not args[0].free_symbols:
|
||||
return S.Zero
|
||||
elif isinstance(obj.expr, TensAdd):
|
||||
# take care of sums of multi PDs
|
||||
result = obj.expr.func(*[
|
||||
self.func(a, *obj.variables)._expand_partial_derivative()
|
||||
for a in result.expr.args])
|
||||
elif isinstance(obj.expr, TensMul):
|
||||
# take care of products of multi PDs
|
||||
if len(obj.variables) == 1:
|
||||
# derivative with respect to single variable
|
||||
terms = []
|
||||
mulargs = list(obj.expr.args)
|
||||
for ind in range(len(mulargs)):
|
||||
if not isinstance(sympify(mulargs[ind]), Number):
|
||||
# a number coefficient is not considered for
|
||||
# expansion of PartialDerivative
|
||||
d = self.func(mulargs[ind], *obj.variables)._expand_partial_derivative()
|
||||
terms.append(TensMul(*(mulargs[:ind]
|
||||
+ [d]
|
||||
+ mulargs[(ind + 1):])))
|
||||
result = TensAdd.fromiter(terms)
|
||||
else:
|
||||
# derivative with respect to multiple variables
|
||||
# decompose:
|
||||
# partial(expr, (u, v))
|
||||
# = partial(partial(expr, u).doit(), v).doit()
|
||||
result = obj.expr # init with expr
|
||||
for v in obj.variables:
|
||||
result = self.func(result, v)._expand_partial_derivative()
|
||||
# then throw PD on it
|
||||
|
||||
return result
|
||||
|
||||
def _perform_derivative(self):
|
||||
result = self.expr
|
||||
for v in self.variables:
|
||||
if isinstance(result, TensExpr):
|
||||
result = result._eval_partial_derivative(v)
|
||||
else:
|
||||
if v._diff_wrt:
|
||||
result = result._eval_derivative(v)
|
||||
else:
|
||||
result = S.Zero
|
||||
return result
|
||||
|
||||
def get_indices(self):
|
||||
return self._indices
|
||||
|
||||
def get_free_indices(self):
|
||||
free = sorted(self._free, key=lambda x: x[1])
|
||||
return [i[0] for i in free]
|
||||
|
||||
def _replace_indices(self, repl):
|
||||
expr = self.expr.xreplace(repl)
|
||||
mirrored = {-k: -v for k, v in repl.items()}
|
||||
variables = [i.xreplace(mirrored) for i in self.variables]
|
||||
return self.func(expr, *variables)
|
||||
|
||||
@property
|
||||
def expr(self):
|
||||
return self.args[0]
|
||||
|
||||
@property
|
||||
def variables(self):
|
||||
return self.args[1:]
|
||||
|
||||
def _extract_data(self, replacement_dict):
|
||||
from .array import derive_by_array, tensorcontraction
|
||||
indices, array = self.expr._extract_data(replacement_dict)
|
||||
for variable in self.variables:
|
||||
var_indices, var_array = variable._extract_data(replacement_dict)
|
||||
var_indices = [-i for i in var_indices]
|
||||
coeff_array, var_array = zip(*[i.as_coeff_Mul() for i in var_array])
|
||||
dim_before = len(array.shape)
|
||||
array = derive_by_array(array, var_array)
|
||||
dim_after = len(array.shape)
|
||||
dim_increase = dim_after - dim_before
|
||||
array = permutedims(array, [i + dim_increase for i in range(dim_before)] + list(range(dim_increase)))
|
||||
array = array.as_mutable()
|
||||
varindex = var_indices[0]
|
||||
# Remove coefficients of base vector:
|
||||
coeff_index = [0] + [slice(None) for i in range(len(indices))]
|
||||
for i, coeff in enumerate(coeff_array):
|
||||
coeff_index[0] = i
|
||||
array[tuple(coeff_index)] /= coeff
|
||||
if -varindex in indices:
|
||||
pos = indices.index(-varindex)
|
||||
array = tensorcontraction(array, (0, pos+1))
|
||||
indices.pop(pos)
|
||||
else:
|
||||
indices.append(varindex)
|
||||
return indices, array
|
||||
Reference in New Issue
Block a user