diff --git a/.gitattributes b/.gitattributes index d77be6c455d46c96e4f44c3df49ba9f1a06640f5..36f069f53f906c6a627f0df413b33636be714433 100644 --- a/.gitattributes +++ b/.gitattributes @@ -437,3 +437,5 @@ tuning-competition-baseline/.venv/lib/python3.11/site-packages/nvidia/cudnn/lib/ .venv/lib/python3.11/site-packages/transformers/models/wav2vec2/__pycache__/modeling_tf_wav2vec2.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text .venv/lib/python3.11/site-packages/transformers/models/wav2vec2_conformer/__pycache__/modeling_wav2vec2_conformer.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text .venv/lib/python3.11/site-packages/transformers/models/whisper/__pycache__/modeling_whisper.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text +.venv/lib/python3.11/site-packages/transformers/utils/__pycache__/dummy_tf_objects.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text +.venv/lib/python3.11/site-packages/transformers/utils/__pycache__/dummy_pt_objects.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text diff --git a/.venv/lib/python3.11/site-packages/sympy/stats/sampling/tests/__pycache__/__init__.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/stats/sampling/tests/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..72ad19c7864226176e4877b7a11e5d084ae2c726 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/stats/sampling/tests/__pycache__/__init__.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/stats/sampling/tests/__pycache__/test_sample_continuous_rv.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/stats/sampling/tests/__pycache__/test_sample_continuous_rv.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..418e16fbd25e4847d04343869ddb43037e070905 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/stats/sampling/tests/__pycache__/test_sample_continuous_rv.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/__init__.py b/.venv/lib/python3.11/site-packages/sympy/tensor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a832614b1d48e26bf01e16f040f34dd412e8e32b --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/__init__.py @@ -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', +] diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/__init__.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..eca2eb4c6c58cb113517b6e41737e9d97abbb84e --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__init__.py @@ -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', +] diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/__init__.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..19193da35b54c76c18c3a25c2878857dee44501d Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/__init__.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/array_comprehension.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/array_comprehension.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e0fec100583d31a63621d0bd2154793d7261010d Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/array_comprehension.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/array_derivatives.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/array_derivatives.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..adf850dce7ea1e6f3dab0444c4726e43bd77c379 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/array_derivatives.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/arrayop.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/arrayop.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..62c8f518d438da704d19869e91524d648900fa3c Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/arrayop.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/dense_ndim_array.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/dense_ndim_array.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f860724cb3307be4cbc75b68b29f96065e243d3 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/dense_ndim_array.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/mutable_ndim_array.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/mutable_ndim_array.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..22f8215ed21524889a20192d6514c5ad99a83ea5 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/mutable_ndim_array.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/ndim_array.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/ndim_array.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46cd9a423cc38f238db7aa851a34d0d348380eb8 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/ndim_array.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/sparse_ndim_array.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/sparse_ndim_array.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fb098b429556cdba4d737a1774aa2428cba779be Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/__pycache__/sparse_ndim_array.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/array_comprehension.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/array_comprehension.py new file mode 100644 index 0000000000000000000000000000000000000000..95702f499f3e40597fd0144929138ac1329962ee --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/array_comprehension.py @@ -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 diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/array_derivatives.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/array_derivatives.py new file mode 100644 index 0000000000000000000000000000000000000000..a38db6caefe256a8c7e1f3415b78351b3787fee9 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/array_derivatives.py @@ -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) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/arrayop.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/arrayop.py new file mode 100644 index 0000000000000000000000000000000000000000..a81e6b381a8a93f0cd585278a4be0259b06406dd --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/arrayop.py @@ -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) + ')' diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/dense_ndim_array.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/dense_ndim_array.py new file mode 100644 index 0000000000000000000000000000000000000000..576e452c55d8d374ca1f72c553f3a64de7227d43 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/dense_ndim_array.py @@ -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} diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__init__.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f1658241782cdf0e38a30c43a6d67f9811297f4c --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__init__.py @@ -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 diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/__init__.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6880c12f2cdcac4990379485f5fbfadf6b529ca4 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/__init__.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/arrayexpr_derivatives.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/arrayexpr_derivatives.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9515ee2b3435736ed916ec6c21b4bfe890b1b520 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/arrayexpr_derivatives.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/conv_array_to_indexed.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/conv_array_to_indexed.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9ce73d695ec602cd2a56ade4a2db1c1a49b44f34 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/conv_array_to_indexed.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/conv_array_to_matrix.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/conv_array_to_matrix.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..61281a8dc3c2cd06743639df6d06d6012710fb9e Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/conv_array_to_matrix.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/conv_indexed_to_array.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/conv_indexed_to_array.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..093b4af9edb6d69a304e938499ce2a8a4f7ffb78 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/conv_indexed_to_array.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/conv_matrix_to_array.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/conv_matrix_to_array.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e813409b01aa121879c0dc74af797ea88fbdbe56 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/conv_matrix_to_array.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/from_array_to_indexed.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/from_array_to_indexed.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a0877c86a43d70c5f1a733dd2c9f397d9b226442 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/from_array_to_indexed.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/from_array_to_matrix.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/from_array_to_matrix.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c0ff581d892cbba6658fd39ebf4e0a3297e0a47 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/from_array_to_matrix.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/from_indexed_to_array.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/from_indexed_to_array.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fd51bd8da101825c33dfd778f4262245a35bf2dd Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/from_indexed_to_array.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/from_matrix_to_array.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/from_matrix_to_array.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..201c23871a7c89af7208134c197f86fce85f0aa6 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/from_matrix_to_array.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/utils.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3594ea2eb046ef3548c8cc5b698c7dd8ac5d136d Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/__pycache__/utils.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/array_expressions.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/array_expressions.py new file mode 100644 index 0000000000000000000000000000000000000000..60912c0a4e23f0146869e6c57b8f7eb12d3fee16 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/array_expressions.py @@ -0,0 +1,1967 @@ +import collections.abc +import operator +from collections import defaultdict, Counter +from functools import reduce +import itertools +from itertools import accumulate +from typing import Optional, List, Tuple as tTuple + +import typing + +from sympy.core.numbers import Integer +from sympy.core.relational import Equality +from sympy.functions.special.tensor_functions import KroneckerDelta +from sympy.core.basic import Basic +from sympy.core.containers import Tuple +from sympy.core.expr import Expr +from sympy.core.function import (Function, Lambda) +from sympy.core.mul import Mul +from sympy.core.singleton import S +from sympy.core.sorting import default_sort_key +from sympy.core.symbol import (Dummy, Symbol) +from sympy.matrices.matrixbase import MatrixBase +from sympy.matrices.expressions.diagonal import diagonalize_vector +from sympy.matrices.expressions.matexpr import MatrixExpr +from sympy.matrices.expressions.special import ZeroMatrix +from sympy.tensor.array.arrayop import (permutedims, tensorcontraction, tensordiagonal, tensorproduct) +from sympy.tensor.array.dense_ndim_array import ImmutableDenseNDimArray +from sympy.tensor.array.ndim_array import NDimArray +from sympy.tensor.indexed import (Indexed, IndexedBase) +from sympy.matrices.expressions.matexpr import MatrixElement +from sympy.tensor.array.expressions.utils import _apply_recursively_over_nested_lists, _sort_contraction_indices, \ + _get_mapping_from_subranks, _build_push_indices_up_func_transformation, _get_contraction_links, \ + _build_push_indices_down_func_transformation +from sympy.combinatorics import Permutation +from sympy.combinatorics.permutations import _af_invert +from sympy.core.sympify import _sympify + + +class _ArrayExpr(Expr): + shape: tTuple[Expr, ...] + + def __getitem__(self, item): + if not isinstance(item, collections.abc.Iterable): + item = (item,) + ArrayElement._check_shape(self, item) + return self._get(item) + + def _get(self, item): + return _get_array_element_or_slice(self, item) + + +class ArraySymbol(_ArrayExpr): + """ + Symbol representing an array expression + """ + + def __new__(cls, symbol, shape: typing.Iterable) -> "ArraySymbol": + if isinstance(symbol, str): + symbol = Symbol(symbol) + # symbol = _sympify(symbol) + shape = Tuple(*map(_sympify, shape)) + obj = Expr.__new__(cls, symbol, shape) + return obj + + @property + def name(self): + return self._args[0] + + @property + def shape(self): + return self._args[1] + + def as_explicit(self): + if not all(i.is_Integer for i in self.shape): + raise ValueError("cannot express explicit array with symbolic shape") + data = [self[i] for i in itertools.product(*[range(j) for j in self.shape])] + return ImmutableDenseNDimArray(data).reshape(*self.shape) + + +class ArrayElement(Expr): + """ + An element of an array. + """ + + _diff_wrt = True + is_symbol = True + is_commutative = True + + def __new__(cls, name, indices): + if isinstance(name, str): + name = Symbol(name) + name = _sympify(name) + if not isinstance(indices, collections.abc.Iterable): + indices = (indices,) + indices = _sympify(tuple(indices)) + cls._check_shape(name, indices) + obj = Expr.__new__(cls, name, indices) + return obj + + @classmethod + def _check_shape(cls, name, indices): + indices = tuple(indices) + if hasattr(name, "shape"): + index_error = IndexError("number of indices does not match shape of the array") + if len(indices) != len(name.shape): + raise index_error + if any((i >= s) == True for i, s in zip(indices, name.shape)): + raise ValueError("shape is out of bounds") + if any((i < 0) == True for i in indices): + raise ValueError("shape contains negative values") + + @property + def name(self): + return self._args[0] + + @property + def indices(self): + return self._args[1] + + def _eval_derivative(self, s): + if not isinstance(s, ArrayElement): + return S.Zero + + if s == self: + return S.One + + if s.name != self.name: + return S.Zero + + return Mul.fromiter(KroneckerDelta(i, j) for i, j in zip(self.indices, s.indices)) + + +class ZeroArray(_ArrayExpr): + """ + Symbolic array of zeros. Equivalent to ``ZeroMatrix`` for matrices. + """ + + def __new__(cls, *shape): + if len(shape) == 0: + return S.Zero + shape = map(_sympify, shape) + obj = Expr.__new__(cls, *shape) + return obj + + @property + def shape(self): + return self._args + + def as_explicit(self): + if not all(i.is_Integer for i in self.shape): + raise ValueError("Cannot return explicit form for symbolic shape.") + return ImmutableDenseNDimArray.zeros(*self.shape) + + def _get(self, item): + return S.Zero + + +class OneArray(_ArrayExpr): + """ + Symbolic array of ones. + """ + + def __new__(cls, *shape): + if len(shape) == 0: + return S.One + shape = map(_sympify, shape) + obj = Expr.__new__(cls, *shape) + return obj + + @property + def shape(self): + return self._args + + def as_explicit(self): + if not all(i.is_Integer for i in self.shape): + raise ValueError("Cannot return explicit form for symbolic shape.") + return ImmutableDenseNDimArray([S.One for i in range(reduce(operator.mul, self.shape))]).reshape(*self.shape) + + def _get(self, item): + return S.One + + +class _CodegenArrayAbstract(Basic): + + @property + def subranks(self): + """ + Returns the ranks of the objects in the uppermost tensor product inside + the current object. In case no tensor products are contained, return + the atomic ranks. + + Examples + ======== + + >>> from sympy.tensor.array import tensorproduct, tensorcontraction + >>> from sympy import MatrixSymbol + >>> M = MatrixSymbol("M", 3, 3) + >>> N = MatrixSymbol("N", 3, 3) + >>> P = MatrixSymbol("P", 3, 3) + + Important: do not confuse the rank of the matrix with the rank of an array. + + >>> tp = tensorproduct(M, N, P) + >>> tp.subranks + [2, 2, 2] + + >>> co = tensorcontraction(tp, (1, 2), (3, 4)) + >>> co.subranks + [2, 2, 2] + """ + return self._subranks[:] + + def subrank(self): + """ + The sum of ``subranks``. + """ + return sum(self.subranks) + + @property + def shape(self): + return self._shape + + def doit(self, **hints): + deep = hints.get("deep", True) + if deep: + return self.func(*[arg.doit(**hints) for arg in self.args])._canonicalize() + else: + return self._canonicalize() + +class ArrayTensorProduct(_CodegenArrayAbstract): + r""" + Class to represent the tensor product of array-like objects. + """ + + def __new__(cls, *args, **kwargs): + args = [_sympify(arg) for arg in args] + + canonicalize = kwargs.pop("canonicalize", False) + + ranks = [get_rank(arg) for arg in args] + + obj = Basic.__new__(cls, *args) + obj._subranks = ranks + shapes = [get_shape(i) for i in args] + + if any(i is None for i in shapes): + obj._shape = None + else: + obj._shape = tuple(j for i in shapes for j in i) + if canonicalize: + return obj._canonicalize() + return obj + + def _canonicalize(self): + args = self.args + args = self._flatten(args) + + ranks = [get_rank(arg) for arg in args] + + # Check if there are nested permutation and lift them up: + permutation_cycles = [] + for i, arg in enumerate(args): + if not isinstance(arg, PermuteDims): + continue + permutation_cycles.extend([[k + sum(ranks[:i]) for k in j] for j in arg.permutation.cyclic_form]) + args[i] = arg.expr + if permutation_cycles: + return _permute_dims(_array_tensor_product(*args), Permutation(sum(ranks)-1)*Permutation(permutation_cycles)) + + if len(args) == 1: + return args[0] + + # If any object is a ZeroArray, return a ZeroArray: + if any(isinstance(arg, (ZeroArray, ZeroMatrix)) for arg in args): + shapes = reduce(operator.add, [get_shape(i) for i in args], ()) + return ZeroArray(*shapes) + + # If there are contraction objects inside, transform the whole + # expression into `ArrayContraction`: + contractions = {i: arg for i, arg in enumerate(args) if isinstance(arg, ArrayContraction)} + if contractions: + ranks = [_get_subrank(arg) if isinstance(arg, ArrayContraction) else get_rank(arg) for arg in args] + cumulative_ranks = list(accumulate([0] + ranks))[:-1] + tp = _array_tensor_product(*[arg.expr if isinstance(arg, ArrayContraction) else arg for arg in args]) + contraction_indices = [tuple(cumulative_ranks[i] + k for k in j) for i, arg in contractions.items() for j in arg.contraction_indices] + return _array_contraction(tp, *contraction_indices) + + diagonals = {i: arg for i, arg in enumerate(args) if isinstance(arg, ArrayDiagonal)} + if diagonals: + inverse_permutation = [] + last_perm = [] + ranks = [get_rank(arg) for arg in args] + cumulative_ranks = list(accumulate([0] + ranks))[:-1] + for i, arg in enumerate(args): + if isinstance(arg, ArrayDiagonal): + i1 = get_rank(arg) - len(arg.diagonal_indices) + i2 = len(arg.diagonal_indices) + inverse_permutation.extend([cumulative_ranks[i] + j for j in range(i1)]) + last_perm.extend([cumulative_ranks[i] + j for j in range(i1, i1 + i2)]) + else: + inverse_permutation.extend([cumulative_ranks[i] + j for j in range(get_rank(arg))]) + inverse_permutation.extend(last_perm) + tp = _array_tensor_product(*[arg.expr if isinstance(arg, ArrayDiagonal) else arg for arg in args]) + ranks2 = [_get_subrank(arg) if isinstance(arg, ArrayDiagonal) else get_rank(arg) for arg in args] + cumulative_ranks2 = list(accumulate([0] + ranks2))[:-1] + diagonal_indices = [tuple(cumulative_ranks2[i] + k for k in j) for i, arg in diagonals.items() for j in arg.diagonal_indices] + return _permute_dims(_array_diagonal(tp, *diagonal_indices), _af_invert(inverse_permutation)) + + return self.func(*args, canonicalize=False) + + @classmethod + def _flatten(cls, args): + args = [i for arg in args for i in (arg.args if isinstance(arg, cls) else [arg])] + return args + + def as_explicit(self): + return tensorproduct(*[arg.as_explicit() if hasattr(arg, "as_explicit") else arg for arg in self.args]) + + +class ArrayAdd(_CodegenArrayAbstract): + r""" + Class for elementwise array additions. + """ + + def __new__(cls, *args, **kwargs): + args = [_sympify(arg) for arg in args] + ranks = [get_rank(arg) for arg in args] + ranks = list(set(ranks)) + if len(ranks) != 1: + raise ValueError("summing arrays of different ranks") + shapes = [arg.shape for arg in args] + if len({i for i in shapes if i is not None}) > 1: + raise ValueError("mismatching shapes in addition") + + canonicalize = kwargs.pop("canonicalize", False) + + obj = Basic.__new__(cls, *args) + obj._subranks = ranks + if any(i is None for i in shapes): + obj._shape = None + else: + obj._shape = shapes[0] + if canonicalize: + return obj._canonicalize() + return obj + + def _canonicalize(self): + args = self.args + + # Flatten: + args = self._flatten_args(args) + + shapes = [get_shape(arg) for arg in args] + args = [arg for arg in args if not isinstance(arg, (ZeroArray, ZeroMatrix))] + if len(args) == 0: + if any(i for i in shapes if i is None): + raise NotImplementedError("cannot handle addition of ZeroMatrix/ZeroArray and undefined shape object") + return ZeroArray(*shapes[0]) + elif len(args) == 1: + return args[0] + return self.func(*args, canonicalize=False) + + @classmethod + def _flatten_args(cls, args): + new_args = [] + for arg in args: + if isinstance(arg, ArrayAdd): + new_args.extend(arg.args) + else: + new_args.append(arg) + return new_args + + def as_explicit(self): + return reduce( + operator.add, + [arg.as_explicit() if hasattr(arg, "as_explicit") else arg for arg in self.args]) + + +class PermuteDims(_CodegenArrayAbstract): + r""" + Class to represent permutation of axes of arrays. + + Examples + ======== + + >>> from sympy.tensor.array import permutedims + >>> from sympy import MatrixSymbol + >>> M = MatrixSymbol("M", 3, 3) + >>> cg = permutedims(M, [1, 0]) + + The object ``cg`` represents the transposition of ``M``, as the permutation + ``[1, 0]`` will act on its indices by switching them: + + `M_{ij} \Rightarrow M_{ji}` + + This is evident when transforming back to matrix form: + + >>> from sympy.tensor.array.expressions.from_array_to_matrix import convert_array_to_matrix + >>> convert_array_to_matrix(cg) + M.T + + >>> N = MatrixSymbol("N", 3, 2) + >>> cg = permutedims(N, [1, 0]) + >>> cg.shape + (2, 3) + + There are optional parameters that can be used as alternative to the permutation: + + >>> from sympy.tensor.array.expressions import ArraySymbol, PermuteDims + >>> M = ArraySymbol("M", (1, 2, 3, 4, 5)) + >>> expr = PermuteDims(M, index_order_old="ijklm", index_order_new="kijml") + >>> expr + PermuteDims(M, (0 2 1)(3 4)) + >>> expr.shape + (3, 1, 2, 5, 4) + + Permutations of tensor products are simplified in order to achieve a + standard form: + + >>> from sympy.tensor.array import tensorproduct + >>> M = MatrixSymbol("M", 4, 5) + >>> tp = tensorproduct(M, N) + >>> tp.shape + (4, 5, 3, 2) + >>> perm1 = permutedims(tp, [2, 3, 1, 0]) + + The args ``(M, N)`` have been sorted and the permutation has been + simplified, the expression is equivalent: + + >>> perm1.expr.args + (N, M) + >>> perm1.shape + (3, 2, 5, 4) + >>> perm1.permutation + (2 3) + + The permutation in its array form has been simplified from + ``[2, 3, 1, 0]`` to ``[0, 1, 3, 2]``, as the arguments of the tensor + product `M` and `N` have been switched: + + >>> perm1.permutation.array_form + [0, 1, 3, 2] + + We can nest a second permutation: + + >>> perm2 = permutedims(perm1, [1, 0, 2, 3]) + >>> perm2.shape + (2, 3, 5, 4) + >>> perm2.permutation.array_form + [1, 0, 3, 2] + """ + + def __new__(cls, expr, permutation=None, index_order_old=None, index_order_new=None, **kwargs): + from sympy.combinatorics import Permutation + expr = _sympify(expr) + expr_rank = get_rank(expr) + permutation = cls._get_permutation_from_arguments(permutation, index_order_old, index_order_new, expr_rank) + permutation = Permutation(permutation) + permutation_size = permutation.size + if permutation_size != expr_rank: + raise ValueError("Permutation size must be the length of the shape of expr") + + canonicalize = kwargs.pop("canonicalize", False) + + obj = Basic.__new__(cls, expr, permutation) + obj._subranks = [get_rank(expr)] + shape = get_shape(expr) + if shape is None: + obj._shape = None + else: + obj._shape = tuple(shape[permutation(i)] for i in range(len(shape))) + if canonicalize: + return obj._canonicalize() + return obj + + def _canonicalize(self): + expr = self.expr + permutation = self.permutation + if isinstance(expr, PermuteDims): + subexpr = expr.expr + subperm = expr.permutation + permutation = permutation * subperm + expr = subexpr + if isinstance(expr, ArrayContraction): + expr, permutation = self._PermuteDims_denestarg_ArrayContraction(expr, permutation) + if isinstance(expr, ArrayTensorProduct): + expr, permutation = self._PermuteDims_denestarg_ArrayTensorProduct(expr, permutation) + if isinstance(expr, (ZeroArray, ZeroMatrix)): + return ZeroArray(*[expr.shape[i] for i in permutation.array_form]) + plist = permutation.array_form + if plist == sorted(plist): + return expr + return self.func(expr, permutation, canonicalize=False) + + @property + def expr(self): + return self.args[0] + + @property + def permutation(self): + return self.args[1] + + @classmethod + def _PermuteDims_denestarg_ArrayTensorProduct(cls, expr, permutation): + # Get the permutation in its image-form: + perm_image_form = _af_invert(permutation.array_form) + args = list(expr.args) + # Starting index global position for every arg: + cumul = list(accumulate([0] + expr.subranks)) + # Split `perm_image_form` into a list of list corresponding to the indices + # of every argument: + perm_image_form_in_components = [perm_image_form[cumul[i]:cumul[i+1]] for i in range(len(args))] + # Create an index, target-position-key array: + ps = [(i, sorted(comp)) for i, comp in enumerate(perm_image_form_in_components)] + # Sort the array according to the target-position-key: + # In this way, we define a canonical way to sort the arguments according + # to the permutation. + ps.sort(key=lambda x: x[1]) + # Read the inverse-permutation (i.e. image-form) of the args: + perm_args_image_form = [i[0] for i in ps] + # Apply the args-permutation to the `args`: + args_sorted = [args[i] for i in perm_args_image_form] + # Apply the args-permutation to the array-form of the permutation of the axes (of `expr`): + perm_image_form_sorted_args = [perm_image_form_in_components[i] for i in perm_args_image_form] + new_permutation = Permutation(_af_invert([j for i in perm_image_form_sorted_args for j in i])) + return _array_tensor_product(*args_sorted), new_permutation + + @classmethod + def _PermuteDims_denestarg_ArrayContraction(cls, expr, permutation): + if not isinstance(expr, ArrayContraction): + return expr, permutation + if not isinstance(expr.expr, ArrayTensorProduct): + return expr, permutation + args = expr.expr.args + subranks = [get_rank(arg) for arg in expr.expr.args] + + contraction_indices = expr.contraction_indices + contraction_indices_flat = [j for i in contraction_indices for j in i] + cumul = list(accumulate([0] + subranks)) + + # Spread the permutation in its array form across the args in the corresponding + # tensor-product arguments with free indices: + permutation_array_blocks_up = [] + image_form = _af_invert(permutation.array_form) + counter = 0 + for i, e in enumerate(subranks): + current = [] + for j in range(cumul[i], cumul[i+1]): + if j in contraction_indices_flat: + continue + current.append(image_form[counter]) + counter += 1 + permutation_array_blocks_up.append(current) + + # Get the map of axis repositioning for every argument of tensor-product: + index_blocks = [list(range(cumul[i], cumul[i+1])) for i, e in enumerate(expr.subranks)] + index_blocks_up = expr._push_indices_up(expr.contraction_indices, index_blocks) + inverse_permutation = permutation**(-1) + index_blocks_up_permuted = [[inverse_permutation(j) for j in i if j is not None] for i in index_blocks_up] + + # Sorting key is a list of tuple, first element is the index of `args`, second element of + # the tuple is the sorting key to sort `args` of the tensor product: + sorting_keys = list(enumerate(index_blocks_up_permuted)) + sorting_keys.sort(key=lambda x: x[1]) + + # Now we can get the permutation acting on the args in its image-form: + new_perm_image_form = [i[0] for i in sorting_keys] + # Apply the args-level permutation to various elements: + new_index_blocks = [index_blocks[i] for i in new_perm_image_form] + new_index_perm_array_form = _af_invert([j for i in new_index_blocks for j in i]) + new_args = [args[i] for i in new_perm_image_form] + new_contraction_indices = [tuple(new_index_perm_array_form[j] for j in i) for i in contraction_indices] + new_expr = _array_contraction(_array_tensor_product(*new_args), *new_contraction_indices) + new_permutation = Permutation(_af_invert([j for i in [permutation_array_blocks_up[k] for k in new_perm_image_form] for j in i])) + return new_expr, new_permutation + + @classmethod + def _check_permutation_mapping(cls, expr, permutation): + subranks = expr.subranks + index2arg = [i for i, arg in enumerate(expr.args) for j in range(expr.subranks[i])] + permuted_indices = [permutation(i) for i in range(expr.subrank())] + new_args = list(expr.args) + arg_candidate_index = index2arg[permuted_indices[0]] + current_indices = [] + new_permutation = [] + inserted_arg_cand_indices = set() + for i, idx in enumerate(permuted_indices): + if index2arg[idx] != arg_candidate_index: + new_permutation.extend(current_indices) + current_indices = [] + arg_candidate_index = index2arg[idx] + current_indices.append(idx) + arg_candidate_rank = subranks[arg_candidate_index] + if len(current_indices) == arg_candidate_rank: + new_permutation.extend(sorted(current_indices)) + local_current_indices = [j - min(current_indices) for j in current_indices] + i1 = index2arg[i] + new_args[i1] = _permute_dims(new_args[i1], Permutation(local_current_indices)) + inserted_arg_cand_indices.add(arg_candidate_index) + current_indices = [] + new_permutation.extend(current_indices) + + # TODO: swap args positions in order to simplify the expression: + # TODO: this should be in a function + args_positions = list(range(len(new_args))) + # Get possible shifts: + maps = {} + cumulative_subranks = [0] + list(accumulate(subranks)) + for i in range(len(subranks)): + s = {index2arg[new_permutation[j]] for j in range(cumulative_subranks[i], cumulative_subranks[i+1])} + if len(s) != 1: + continue + elem = next(iter(s)) + if i != elem: + maps[i] = elem + + # Find cycles in the map: + lines = [] + current_line = [] + while maps: + if len(current_line) == 0: + k, v = maps.popitem() + current_line.append(k) + else: + k = current_line[-1] + if k not in maps: + current_line = [] + continue + v = maps.pop(k) + if v in current_line: + lines.append(current_line) + current_line = [] + continue + current_line.append(v) + for line in lines: + for i, e in enumerate(line): + args_positions[line[(i + 1) % len(line)]] = e + + # TODO: function in order to permute the args: + permutation_blocks = [[new_permutation[cumulative_subranks[i] + j] for j in range(e)] for i, e in enumerate(subranks)] + new_args = [new_args[i] for i in args_positions] + new_permutation_blocks = [permutation_blocks[i] for i in args_positions] + new_permutation2 = [j for i in new_permutation_blocks for j in i] + return _array_tensor_product(*new_args), Permutation(new_permutation2) # **(-1) + + @classmethod + def _check_if_there_are_closed_cycles(cls, expr, permutation): + args = list(expr.args) + subranks = expr.subranks + cyclic_form = permutation.cyclic_form + cumulative_subranks = [0] + list(accumulate(subranks)) + cyclic_min = [min(i) for i in cyclic_form] + cyclic_max = [max(i) for i in cyclic_form] + cyclic_keep = [] + for i, cycle in enumerate(cyclic_form): + flag = True + for j in range(len(cumulative_subranks) - 1): + if cyclic_min[i] >= cumulative_subranks[j] and cyclic_max[i] < cumulative_subranks[j+1]: + # Found a sinkable cycle. + args[j] = _permute_dims(args[j], Permutation([[k - cumulative_subranks[j] for k in cyclic_form[i]]])) + flag = False + break + if flag: + cyclic_keep.append(cyclic_form[i]) + return _array_tensor_product(*args), Permutation(cyclic_keep, size=permutation.size) + + def nest_permutation(self): + r""" + DEPRECATED. + """ + ret = self._nest_permutation(self.expr, self.permutation) + if ret is None: + return self + return ret + + @classmethod + def _nest_permutation(cls, expr, permutation): + if isinstance(expr, ArrayTensorProduct): + return _permute_dims(*cls._check_if_there_are_closed_cycles(expr, permutation)) + elif isinstance(expr, ArrayContraction): + # Invert tree hierarchy: put the contraction above. + cycles = permutation.cyclic_form + newcycles = ArrayContraction._convert_outer_indices_to_inner_indices(expr, *cycles) + newpermutation = Permutation(newcycles) + new_contr_indices = [tuple(newpermutation(j) for j in i) for i in expr.contraction_indices] + return _array_contraction(PermuteDims(expr.expr, newpermutation), *new_contr_indices) + elif isinstance(expr, ArrayAdd): + return _array_add(*[PermuteDims(arg, permutation) for arg in expr.args]) + return None + + def as_explicit(self): + expr = self.expr + if hasattr(expr, "as_explicit"): + expr = expr.as_explicit() + return permutedims(expr, self.permutation) + + @classmethod + def _get_permutation_from_arguments(cls, permutation, index_order_old, index_order_new, dim): + if permutation is None: + if index_order_new is None or index_order_old is None: + raise ValueError("Permutation not defined") + return PermuteDims._get_permutation_from_index_orders(index_order_old, index_order_new, dim) + else: + if index_order_new is not None: + raise ValueError("index_order_new cannot be defined with permutation") + if index_order_old is not None: + raise ValueError("index_order_old cannot be defined with permutation") + return permutation + + @classmethod + def _get_permutation_from_index_orders(cls, index_order_old, index_order_new, dim): + if len(set(index_order_new)) != dim: + raise ValueError("wrong number of indices in index_order_new") + if len(set(index_order_old)) != dim: + raise ValueError("wrong number of indices in index_order_old") + if len(set.symmetric_difference(set(index_order_new), set(index_order_old))) > 0: + raise ValueError("index_order_new and index_order_old must have the same indices") + permutation = [index_order_old.index(i) for i in index_order_new] + return permutation + + +class ArrayDiagonal(_CodegenArrayAbstract): + r""" + Class to represent the diagonal operator. + + Explanation + =========== + + In a 2-dimensional array it returns the diagonal, this looks like the + operation: + + `A_{ij} \rightarrow A_{ii}` + + The diagonal over axes 1 and 2 (the second and third) of the tensor product + of two 2-dimensional arrays `A \otimes B` is + + `\Big[ A_{ab} B_{cd} \Big]_{abcd} \rightarrow \Big[ A_{ai} B_{id} \Big]_{adi}` + + In this last example the array expression has been reduced from + 4-dimensional to 3-dimensional. Notice that no contraction has occurred, + rather there is a new index `i` for the diagonal, contraction would have + reduced the array to 2 dimensions. + + Notice that the diagonalized out dimensions are added as new dimensions at + the end of the indices. + """ + + def __new__(cls, expr, *diagonal_indices, **kwargs): + expr = _sympify(expr) + diagonal_indices = [Tuple(*sorted(i)) for i in diagonal_indices] + canonicalize = kwargs.get("canonicalize", False) + + shape = get_shape(expr) + if shape is not None: + cls._validate(expr, *diagonal_indices, **kwargs) + # Get new shape: + positions, shape = cls._get_positions_shape(shape, diagonal_indices) + else: + positions = None + if len(diagonal_indices) == 0: + return expr + obj = Basic.__new__(cls, expr, *diagonal_indices) + obj._positions = positions + obj._subranks = _get_subranks(expr) + obj._shape = shape + if canonicalize: + return obj._canonicalize() + return obj + + def _canonicalize(self): + expr = self.expr + diagonal_indices = self.diagonal_indices + trivial_diags = [i for i in diagonal_indices if len(i) == 1] + if len(trivial_diags) > 0: + trivial_pos = {e[0]: i for i, e in enumerate(diagonal_indices) if len(e) == 1} + diag_pos = {e: i for i, e in enumerate(diagonal_indices) if len(e) > 1} + diagonal_indices_short = [i for i in diagonal_indices if len(i) > 1] + rank1 = get_rank(self) + rank2 = len(diagonal_indices) + rank3 = rank1 - rank2 + inv_permutation = [] + counter1 = 0 + indices_down = ArrayDiagonal._push_indices_down(diagonal_indices_short, list(range(rank1)), get_rank(expr)) + for i in indices_down: + if i in trivial_pos: + inv_permutation.append(rank3 + trivial_pos[i]) + elif isinstance(i, (Integer, int)): + inv_permutation.append(counter1) + counter1 += 1 + else: + inv_permutation.append(rank3 + diag_pos[i]) + permutation = _af_invert(inv_permutation) + if len(diagonal_indices_short) > 0: + return _permute_dims(_array_diagonal(expr, *diagonal_indices_short), permutation) + else: + return _permute_dims(expr, permutation) + if isinstance(expr, ArrayAdd): + return self._ArrayDiagonal_denest_ArrayAdd(expr, *diagonal_indices) + if isinstance(expr, ArrayDiagonal): + return self._ArrayDiagonal_denest_ArrayDiagonal(expr, *diagonal_indices) + if isinstance(expr, PermuteDims): + return self._ArrayDiagonal_denest_PermuteDims(expr, *diagonal_indices) + if isinstance(expr, (ZeroArray, ZeroMatrix)): + positions, shape = self._get_positions_shape(expr.shape, diagonal_indices) + return ZeroArray(*shape) + return self.func(expr, *diagonal_indices, canonicalize=False) + + @staticmethod + def _validate(expr, *diagonal_indices, **kwargs): + # Check that no diagonalization happens on indices with mismatched + # dimensions: + shape = get_shape(expr) + for i in diagonal_indices: + if any(j >= len(shape) for j in i): + raise ValueError("index is larger than expression shape") + if len({shape[j] for j in i}) != 1: + raise ValueError("diagonalizing indices of different dimensions") + if not kwargs.get("allow_trivial_diags", False) and len(i) <= 1: + raise ValueError("need at least two axes to diagonalize") + if len(set(i)) != len(i): + raise ValueError("axis index cannot be repeated") + + @staticmethod + def _remove_trivial_dimensions(shape, *diagonal_indices): + return [tuple(j for j in i) for i in diagonal_indices if shape[i[0]] != 1] + + @property + def expr(self): + return self.args[0] + + @property + def diagonal_indices(self): + return self.args[1:] + + @staticmethod + def _flatten(expr, *outer_diagonal_indices): + inner_diagonal_indices = expr.diagonal_indices + all_inner = [j for i in inner_diagonal_indices for j in i] + all_inner.sort() + # TODO: add API for total rank and cumulative rank: + total_rank = _get_subrank(expr) + inner_rank = len(all_inner) + outer_rank = total_rank - inner_rank + shifts = [0 for i in range(outer_rank)] + counter = 0 + pointer = 0 + for i in range(outer_rank): + while pointer < inner_rank and counter >= all_inner[pointer]: + counter += 1 + pointer += 1 + shifts[i] += pointer + counter += 1 + outer_diagonal_indices = tuple(tuple(shifts[j] + j for j in i) for i in outer_diagonal_indices) + diagonal_indices = inner_diagonal_indices + outer_diagonal_indices + return _array_diagonal(expr.expr, *diagonal_indices) + + @classmethod + def _ArrayDiagonal_denest_ArrayAdd(cls, expr, *diagonal_indices): + return _array_add(*[_array_diagonal(arg, *diagonal_indices) for arg in expr.args]) + + @classmethod + def _ArrayDiagonal_denest_ArrayDiagonal(cls, expr, *diagonal_indices): + return cls._flatten(expr, *diagonal_indices) + + @classmethod + def _ArrayDiagonal_denest_PermuteDims(cls, expr: PermuteDims, *diagonal_indices): + back_diagonal_indices = [[expr.permutation(j) for j in i] for i in diagonal_indices] + nondiag = [i for i in range(get_rank(expr)) if not any(i in j for j in diagonal_indices)] + back_nondiag = [expr.permutation(i) for i in nondiag] + remap = {e: i for i, e in enumerate(sorted(back_nondiag))} + new_permutation1 = [remap[i] for i in back_nondiag] + shift = len(new_permutation1) + diag_block_perm = [i + shift for i in range(len(back_diagonal_indices))] + new_permutation = new_permutation1 + diag_block_perm + return _permute_dims( + _array_diagonal( + expr.expr, + *back_diagonal_indices + ), + new_permutation + ) + + def _push_indices_down_nonstatic(self, indices): + transform = lambda x: self._positions[x] if x < len(self._positions) else None + return _apply_recursively_over_nested_lists(transform, indices) + + def _push_indices_up_nonstatic(self, indices): + + def transform(x): + for i, e in enumerate(self._positions): + if (isinstance(e, int) and x == e) or (isinstance(e, tuple) and x in e): + return i + + return _apply_recursively_over_nested_lists(transform, indices) + + @classmethod + def _push_indices_down(cls, diagonal_indices, indices, rank): + positions, shape = cls._get_positions_shape(range(rank), diagonal_indices) + transform = lambda x: positions[x] if x < len(positions) else None + return _apply_recursively_over_nested_lists(transform, indices) + + @classmethod + def _push_indices_up(cls, diagonal_indices, indices, rank): + positions, shape = cls._get_positions_shape(range(rank), diagonal_indices) + + def transform(x): + for i, e in enumerate(positions): + if (isinstance(e, int) and x == e) or (isinstance(e, (tuple, Tuple)) and (x in e)): + return i + + return _apply_recursively_over_nested_lists(transform, indices) + + @classmethod + def _get_positions_shape(cls, shape, diagonal_indices): + data1 = tuple((i, shp) for i, shp in enumerate(shape) if not any(i in j for j in diagonal_indices)) + pos1, shp1 = zip(*data1) if data1 else ((), ()) + data2 = tuple((i, shape[i[0]]) for i in diagonal_indices) + pos2, shp2 = zip(*data2) if data2 else ((), ()) + positions = pos1 + pos2 + shape = shp1 + shp2 + return positions, shape + + def as_explicit(self): + expr = self.expr + if hasattr(expr, "as_explicit"): + expr = expr.as_explicit() + return tensordiagonal(expr, *self.diagonal_indices) + + +class ArrayElementwiseApplyFunc(_CodegenArrayAbstract): + + def __new__(cls, function, element): + + if not isinstance(function, Lambda): + d = Dummy('d') + function = Lambda(d, function(d)) + + obj = _CodegenArrayAbstract.__new__(cls, function, element) + obj._subranks = _get_subranks(element) + return obj + + @property + def function(self): + return self.args[0] + + @property + def expr(self): + return self.args[1] + + @property + def shape(self): + return self.expr.shape + + def _get_function_fdiff(self): + d = Dummy("d") + function = self.function(d) + fdiff = function.diff(d) + if isinstance(fdiff, Function): + fdiff = type(fdiff) + else: + fdiff = Lambda(d, fdiff) + return fdiff + + def as_explicit(self): + expr = self.expr + if hasattr(expr, "as_explicit"): + expr = expr.as_explicit() + return expr.applyfunc(self.function) + + +class ArrayContraction(_CodegenArrayAbstract): + r""" + This class is meant to represent contractions of arrays in a form easily + processable by the code printers. + """ + + def __new__(cls, expr, *contraction_indices, **kwargs): + contraction_indices = _sort_contraction_indices(contraction_indices) + expr = _sympify(expr) + + canonicalize = kwargs.get("canonicalize", False) + + obj = Basic.__new__(cls, expr, *contraction_indices) + obj._subranks = _get_subranks(expr) + obj._mapping = _get_mapping_from_subranks(obj._subranks) + + free_indices_to_position = {i: i for i in range(sum(obj._subranks)) if all(i not in cind for cind in contraction_indices)} + obj._free_indices_to_position = free_indices_to_position + + shape = get_shape(expr) + cls._validate(expr, *contraction_indices) + if shape: + shape = tuple(shp for i, shp in enumerate(shape) if not any(i in j for j in contraction_indices)) + obj._shape = shape + if canonicalize: + return obj._canonicalize() + return obj + + def _canonicalize(self): + expr = self.expr + contraction_indices = self.contraction_indices + + if len(contraction_indices) == 0: + return expr + + if isinstance(expr, ArrayContraction): + return self._ArrayContraction_denest_ArrayContraction(expr, *contraction_indices) + + if isinstance(expr, (ZeroArray, ZeroMatrix)): + return self._ArrayContraction_denest_ZeroArray(expr, *contraction_indices) + + if isinstance(expr, PermuteDims): + return self._ArrayContraction_denest_PermuteDims(expr, *contraction_indices) + + if isinstance(expr, ArrayTensorProduct): + expr, contraction_indices = self._sort_fully_contracted_args(expr, contraction_indices) + expr, contraction_indices = self._lower_contraction_to_addends(expr, contraction_indices) + if len(contraction_indices) == 0: + return expr + + if isinstance(expr, ArrayDiagonal): + return self._ArrayContraction_denest_ArrayDiagonal(expr, *contraction_indices) + + if isinstance(expr, ArrayAdd): + return self._ArrayContraction_denest_ArrayAdd(expr, *contraction_indices) + + # Check single index contractions on 1-dimensional axes: + contraction_indices = [i for i in contraction_indices if len(i) > 1 or get_shape(expr)[i[0]] != 1] + if len(contraction_indices) == 0: + return expr + + return self.func(expr, *contraction_indices, canonicalize=False) + + def __mul__(self, other): + if other == 1: + return self + else: + raise NotImplementedError("Product of N-dim arrays is not uniquely defined. Use another method.") + + def __rmul__(self, other): + if other == 1: + return self + else: + raise NotImplementedError("Product of N-dim arrays is not uniquely defined. Use another method.") + + @staticmethod + def _validate(expr, *contraction_indices): + shape = get_shape(expr) + if shape is None: + return + + # Check that no contraction happens when the shape is mismatched: + for i in contraction_indices: + if len({shape[j] for j in i if shape[j] != -1}) != 1: + raise ValueError("contracting indices of different dimensions") + + @classmethod + def _push_indices_down(cls, contraction_indices, indices): + flattened_contraction_indices = [j for i in contraction_indices for j in i] + flattened_contraction_indices.sort() + transform = _build_push_indices_down_func_transformation(flattened_contraction_indices) + return _apply_recursively_over_nested_lists(transform, indices) + + @classmethod + def _push_indices_up(cls, contraction_indices, indices): + flattened_contraction_indices = [j for i in contraction_indices for j in i] + flattened_contraction_indices.sort() + transform = _build_push_indices_up_func_transformation(flattened_contraction_indices) + return _apply_recursively_over_nested_lists(transform, indices) + + @classmethod + def _lower_contraction_to_addends(cls, expr, contraction_indices): + if isinstance(expr, ArrayAdd): + raise NotImplementedError() + if not isinstance(expr, ArrayTensorProduct): + return expr, contraction_indices + subranks = expr.subranks + cumranks = list(accumulate([0] + subranks)) + contraction_indices_remaining = [] + contraction_indices_args = [[] for i in expr.args] + backshift = set() + for contraction_group in contraction_indices: + for j in range(len(expr.args)): + if not isinstance(expr.args[j], ArrayAdd): + continue + if all(cumranks[j] <= k < cumranks[j+1] for k in contraction_group): + contraction_indices_args[j].append([k - cumranks[j] for k in contraction_group]) + backshift.update(contraction_group) + break + else: + contraction_indices_remaining.append(contraction_group) + if len(contraction_indices_remaining) == len(contraction_indices): + return expr, contraction_indices + total_rank = get_rank(expr) + shifts = list(accumulate([1 if i in backshift else 0 for i in range(total_rank)])) + contraction_indices_remaining = [Tuple.fromiter(j - shifts[j] for j in i) for i in contraction_indices_remaining] + ret = _array_tensor_product(*[ + _array_contraction(arg, *contr) for arg, contr in zip(expr.args, contraction_indices_args) + ]) + return ret, contraction_indices_remaining + + def split_multiple_contractions(self): + """ + Recognize multiple contractions and attempt at rewriting them as paired-contractions. + + This allows some contractions involving more than two indices to be + rewritten as multiple contractions involving two indices, thus allowing + the expression to be rewritten as a matrix multiplication line. + + Examples: + + * `A_ij b_j0 C_jk` ===> `A*DiagMatrix(b)*C` + + Care for: + - matrix being diagonalized (i.e. `A_ii`) + - vectors being diagonalized (i.e. `a_i0`) + + Multiple contractions can be split into matrix multiplications if + not more than two arguments are non-diagonals or non-vectors. + Vectors get diagonalized while diagonal matrices remain diagonal. + The non-diagonal matrices can be at the beginning or at the end + of the final matrix multiplication line. + """ + + editor = _EditArrayContraction(self) + + contraction_indices = self.contraction_indices + + onearray_insert = [] + + for indl, links in enumerate(contraction_indices): + if len(links) <= 2: + continue + + # Check multiple contractions: + # + # Examples: + # + # * `A_ij b_j0 C_jk` ===> `A*DiagMatrix(b)*C \otimes OneArray(1)` with permutation (1 2) + # + # Care for: + # - matrix being diagonalized (i.e. `A_ii`) + # - vectors being diagonalized (i.e. `a_i0`) + + # Multiple contractions can be split into matrix multiplications if + # not more than three arguments are non-diagonals or non-vectors. + # + # Vectors get diagonalized while diagonal matrices remain diagonal. + # The non-diagonal matrices can be at the beginning or at the end + # of the final matrix multiplication line. + + positions = editor.get_mapping_for_index(indl) + + # Also consider the case of diagonal matrices being contracted: + current_dimension = self.expr.shape[links[0]] + + not_vectors = [] + vectors = [] + for arg_ind, rel_ind in positions: + arg = editor.args_with_ind[arg_ind] + mat = arg.element + abs_arg_start, abs_arg_end = editor.get_absolute_range(arg) + other_arg_pos = 1-rel_ind + other_arg_abs = abs_arg_start + other_arg_pos + if ((1 not in mat.shape) or + ((current_dimension == 1) is True and mat.shape != (1, 1)) or + any(other_arg_abs in l for li, l in enumerate(contraction_indices) if li != indl) + ): + not_vectors.append((arg, rel_ind)) + else: + vectors.append((arg, rel_ind)) + if len(not_vectors) > 2: + # If more than two arguments in the multiple contraction are + # non-vectors and non-diagonal matrices, we cannot find a way + # to split this contraction into a matrix multiplication line: + continue + # Three cases to handle: + # - zero non-vectors + # - one non-vector + # - two non-vectors + for v, rel_ind in vectors: + v.element = diagonalize_vector(v.element) + vectors_to_loop = not_vectors[:1] + vectors + not_vectors[1:] + first_not_vector, rel_ind = vectors_to_loop[0] + new_index = first_not_vector.indices[rel_ind] + + for v, rel_ind in vectors_to_loop[1:-1]: + v.indices[rel_ind] = new_index + new_index = editor.get_new_contraction_index() + assert v.indices.index(None) == 1 - rel_ind + v.indices[v.indices.index(None)] = new_index + onearray_insert.append(v) + + last_vec, rel_ind = vectors_to_loop[-1] + last_vec.indices[rel_ind] = new_index + + for v in onearray_insert: + editor.insert_after(v, _ArgE(OneArray(1), [None])) + + return editor.to_array_contraction() + + def flatten_contraction_of_diagonal(self): + if not isinstance(self.expr, ArrayDiagonal): + return self + contraction_down = self.expr._push_indices_down(self.expr.diagonal_indices, self.contraction_indices) + new_contraction_indices = [] + diagonal_indices = self.expr.diagonal_indices[:] + for i in contraction_down: + contraction_group = list(i) + for j in i: + diagonal_with = [k for k in diagonal_indices if j in k] + contraction_group.extend([l for k in diagonal_with for l in k]) + diagonal_indices = [k for k in diagonal_indices if k not in diagonal_with] + new_contraction_indices.append(sorted(set(contraction_group))) + + new_contraction_indices = ArrayDiagonal._push_indices_up(diagonal_indices, new_contraction_indices) + return _array_contraction( + _array_diagonal( + self.expr.expr, + *diagonal_indices + ), + *new_contraction_indices + ) + + @staticmethod + def _get_free_indices_to_position_map(free_indices, contraction_indices): + free_indices_to_position = {} + flattened_contraction_indices = [j for i in contraction_indices for j in i] + counter = 0 + for ind in free_indices: + while counter in flattened_contraction_indices: + counter += 1 + free_indices_to_position[ind] = counter + counter += 1 + return free_indices_to_position + + @staticmethod + def _get_index_shifts(expr): + """ + Get the mapping of indices at the positions before the contraction + occurs. + + Examples + ======== + + >>> from sympy.tensor.array import tensorproduct, tensorcontraction + >>> from sympy import MatrixSymbol + >>> M = MatrixSymbol("M", 3, 3) + >>> N = MatrixSymbol("N", 3, 3) + >>> cg = tensorcontraction(tensorproduct(M, N), [1, 2]) + >>> cg._get_index_shifts(cg) + [0, 2] + + Indeed, ``cg`` after the contraction has two dimensions, 0 and 1. They + need to be shifted by 0 and 2 to get the corresponding positions before + the contraction (that is, 0 and 3). + """ + inner_contraction_indices = expr.contraction_indices + all_inner = [j for i in inner_contraction_indices for j in i] + all_inner.sort() + # TODO: add API for total rank and cumulative rank: + total_rank = _get_subrank(expr) + inner_rank = len(all_inner) + outer_rank = total_rank - inner_rank + shifts = [0 for i in range(outer_rank)] + counter = 0 + pointer = 0 + for i in range(outer_rank): + while pointer < inner_rank and counter >= all_inner[pointer]: + counter += 1 + pointer += 1 + shifts[i] += pointer + counter += 1 + return shifts + + @staticmethod + def _convert_outer_indices_to_inner_indices(expr, *outer_contraction_indices): + shifts = ArrayContraction._get_index_shifts(expr) + outer_contraction_indices = tuple(tuple(shifts[j] + j for j in i) for i in outer_contraction_indices) + return outer_contraction_indices + + @staticmethod + def _flatten(expr, *outer_contraction_indices): + inner_contraction_indices = expr.contraction_indices + outer_contraction_indices = ArrayContraction._convert_outer_indices_to_inner_indices(expr, *outer_contraction_indices) + contraction_indices = inner_contraction_indices + outer_contraction_indices + return _array_contraction(expr.expr, *contraction_indices) + + @classmethod + def _ArrayContraction_denest_ArrayContraction(cls, expr, *contraction_indices): + return cls._flatten(expr, *contraction_indices) + + @classmethod + def _ArrayContraction_denest_ZeroArray(cls, expr, *contraction_indices): + contraction_indices_flat = [j for i in contraction_indices for j in i] + shape = [e for i, e in enumerate(expr.shape) if i not in contraction_indices_flat] + return ZeroArray(*shape) + + @classmethod + def _ArrayContraction_denest_ArrayAdd(cls, expr, *contraction_indices): + return _array_add(*[_array_contraction(i, *contraction_indices) for i in expr.args]) + + @classmethod + def _ArrayContraction_denest_PermuteDims(cls, expr, *contraction_indices): + permutation = expr.permutation + plist = permutation.array_form + new_contraction_indices = [tuple(permutation(j) for j in i) for i in contraction_indices] + new_plist = [i for i in plist if not any(i in j for j in new_contraction_indices)] + new_plist = cls._push_indices_up(new_contraction_indices, new_plist) + return _permute_dims( + _array_contraction(expr.expr, *new_contraction_indices), + Permutation(new_plist) + ) + + @classmethod + def _ArrayContraction_denest_ArrayDiagonal(cls, expr: 'ArrayDiagonal', *contraction_indices): + diagonal_indices = list(expr.diagonal_indices) + down_contraction_indices = expr._push_indices_down(expr.diagonal_indices, contraction_indices, get_rank(expr.expr)) + # Flatten diagonally contracted indices: + down_contraction_indices = [[k for j in i for k in (j if isinstance(j, (tuple, Tuple)) else [j])] for i in down_contraction_indices] + new_contraction_indices = [] + for contr_indgrp in down_contraction_indices: + ind = contr_indgrp[:] + for j, diag_indgrp in enumerate(diagonal_indices): + if diag_indgrp is None: + continue + if any(i in diag_indgrp for i in contr_indgrp): + ind.extend(diag_indgrp) + diagonal_indices[j] = None + new_contraction_indices.append(sorted(set(ind))) + + new_diagonal_indices_down = [i for i in diagonal_indices if i is not None] + new_diagonal_indices = ArrayContraction._push_indices_up(new_contraction_indices, new_diagonal_indices_down) + return _array_diagonal( + _array_contraction(expr.expr, *new_contraction_indices), + *new_diagonal_indices + ) + + @classmethod + def _sort_fully_contracted_args(cls, expr, contraction_indices): + if expr.shape is None: + return expr, contraction_indices + cumul = list(accumulate([0] + expr.subranks)) + index_blocks = [list(range(cumul[i], cumul[i+1])) for i in range(len(expr.args))] + contraction_indices_flat = {j for i in contraction_indices for j in i} + fully_contracted = [all(j in contraction_indices_flat for j in range(cumul[i], cumul[i+1])) for i, arg in enumerate(expr.args)] + new_pos = sorted(range(len(expr.args)), key=lambda x: (0, default_sort_key(expr.args[x])) if fully_contracted[x] else (1,)) + new_args = [expr.args[i] for i in new_pos] + new_index_blocks_flat = [j for i in new_pos for j in index_blocks[i]] + index_permutation_array_form = _af_invert(new_index_blocks_flat) + new_contraction_indices = [tuple(index_permutation_array_form[j] for j in i) for i in contraction_indices] + new_contraction_indices = _sort_contraction_indices(new_contraction_indices) + return _array_tensor_product(*new_args), new_contraction_indices + + def _get_contraction_tuples(self): + r""" + Return tuples containing the argument index and position within the + argument of the index position. + + Examples + ======== + + >>> from sympy import MatrixSymbol + >>> from sympy.abc import N + >>> from sympy.tensor.array import tensorproduct, tensorcontraction + >>> A = MatrixSymbol("A", N, N) + >>> B = MatrixSymbol("B", N, N) + + >>> cg = tensorcontraction(tensorproduct(A, B), (1, 2)) + >>> cg._get_contraction_tuples() + [[(0, 1), (1, 0)]] + + Notes + ===== + + Here the contraction pair `(1, 2)` meaning that the 2nd and 3rd indices + of the tensor product `A\otimes B` are contracted, has been transformed + into `(0, 1)` and `(1, 0)`, identifying the same indices in a different + notation. `(0, 1)` is the second index (1) of the first argument (i.e. + 0 or `A`). `(1, 0)` is the first index (i.e. 0) of the second + argument (i.e. 1 or `B`). + """ + mapping = self._mapping + return [[mapping[j] for j in i] for i in self.contraction_indices] + + @staticmethod + def _contraction_tuples_to_contraction_indices(expr, contraction_tuples): + # TODO: check that `expr` has `.subranks`: + ranks = expr.subranks + cumulative_ranks = [0] + list(accumulate(ranks)) + return [tuple(cumulative_ranks[j]+k for j, k in i) for i in contraction_tuples] + + @property + def free_indices(self): + return self._free_indices[:] + + @property + def free_indices_to_position(self): + return dict(self._free_indices_to_position) + + @property + def expr(self): + return self.args[0] + + @property + def contraction_indices(self): + return self.args[1:] + + def _contraction_indices_to_components(self): + expr = self.expr + if not isinstance(expr, ArrayTensorProduct): + raise NotImplementedError("only for contractions of tensor products") + ranks = expr.subranks + mapping = {} + counter = 0 + for i, rank in enumerate(ranks): + for j in range(rank): + mapping[counter] = (i, j) + counter += 1 + return mapping + + def sort_args_by_name(self): + """ + Sort arguments in the tensor product so that their order is lexicographical. + + Examples + ======== + + >>> from sympy.tensor.array.expressions.from_matrix_to_array import convert_matrix_to_array + >>> from sympy import MatrixSymbol + >>> from sympy.abc import N + >>> A = MatrixSymbol("A", N, N) + >>> B = MatrixSymbol("B", N, N) + >>> C = MatrixSymbol("C", N, N) + >>> D = MatrixSymbol("D", N, N) + + >>> cg = convert_matrix_to_array(C*D*A*B) + >>> cg + ArrayContraction(ArrayTensorProduct(A, D, C, B), (0, 3), (1, 6), (2, 5)) + >>> cg.sort_args_by_name() + ArrayContraction(ArrayTensorProduct(A, D, B, C), (0, 3), (1, 4), (2, 7)) + """ + expr = self.expr + if not isinstance(expr, ArrayTensorProduct): + return self + args = expr.args + sorted_data = sorted(enumerate(args), key=lambda x: default_sort_key(x[1])) + pos_sorted, args_sorted = zip(*sorted_data) + reordering_map = {i: pos_sorted.index(i) for i, arg in enumerate(args)} + contraction_tuples = self._get_contraction_tuples() + contraction_tuples = [[(reordering_map[j], k) for j, k in i] for i in contraction_tuples] + c_tp = _array_tensor_product(*args_sorted) + new_contr_indices = self._contraction_tuples_to_contraction_indices( + c_tp, + contraction_tuples + ) + return _array_contraction(c_tp, *new_contr_indices) + + def _get_contraction_links(self): + r""" + Returns a dictionary of links between arguments in the tensor product + being contracted. + + See the example for an explanation of the values. + + Examples + ======== + + >>> from sympy import MatrixSymbol + >>> from sympy.abc import N + >>> from sympy.tensor.array.expressions.from_matrix_to_array import convert_matrix_to_array + >>> A = MatrixSymbol("A", N, N) + >>> B = MatrixSymbol("B", N, N) + >>> C = MatrixSymbol("C", N, N) + >>> D = MatrixSymbol("D", N, N) + + Matrix multiplications are pairwise contractions between neighboring + matrices: + + `A_{ij} B_{jk} C_{kl} D_{lm}` + + >>> cg = convert_matrix_to_array(A*B*C*D) + >>> cg + ArrayContraction(ArrayTensorProduct(B, C, A, D), (0, 5), (1, 2), (3, 6)) + + >>> cg._get_contraction_links() + {0: {0: (2, 1), 1: (1, 0)}, 1: {0: (0, 1), 1: (3, 0)}, 2: {1: (0, 0)}, 3: {0: (1, 1)}} + + This dictionary is interpreted as follows: argument in position 0 (i.e. + matrix `A`) has its second index (i.e. 1) contracted to `(1, 0)`, that + is argument in position 1 (matrix `B`) on the first index slot of `B`, + this is the contraction provided by the index `j` from `A`. + + The argument in position 1 (that is, matrix `B`) has two contractions, + the ones provided by the indices `j` and `k`, respectively the first + and second indices (0 and 1 in the sub-dict). The link `(0, 1)` and + `(2, 0)` respectively. `(0, 1)` is the index slot 1 (the 2nd) of + argument in position 0 (that is, `A_{\ldot j}`), and so on. + """ + args, dlinks = _get_contraction_links([self], self.subranks, *self.contraction_indices) + return dlinks + + def as_explicit(self): + expr = self.expr + if hasattr(expr, "as_explicit"): + expr = expr.as_explicit() + return tensorcontraction(expr, *self.contraction_indices) + + +class Reshape(_CodegenArrayAbstract): + """ + Reshape the dimensions of an array expression. + + Examples + ======== + + >>> from sympy.tensor.array.expressions import ArraySymbol, Reshape + >>> A = ArraySymbol("A", (6,)) + >>> A.shape + (6,) + >>> Reshape(A, (3, 2)).shape + (3, 2) + + Check the component-explicit forms: + + >>> A.as_explicit() + [A[0], A[1], A[2], A[3], A[4], A[5]] + >>> Reshape(A, (3, 2)).as_explicit() + [[A[0], A[1]], [A[2], A[3]], [A[4], A[5]]] + + """ + + def __new__(cls, expr, shape): + expr = _sympify(expr) + if not isinstance(shape, Tuple): + shape = Tuple(*shape) + if Equality(Mul.fromiter(expr.shape), Mul.fromiter(shape)) == False: + raise ValueError("shape mismatch") + obj = Expr.__new__(cls, expr, shape) + obj._shape = tuple(shape) + obj._expr = expr + return obj + + @property + def shape(self): + return self._shape + + @property + def expr(self): + return self._expr + + def doit(self, *args, **kwargs): + if kwargs.get("deep", True): + expr = self.expr.doit(*args, **kwargs) + else: + expr = self.expr + if isinstance(expr, (MatrixBase, NDimArray)): + return expr.reshape(*self.shape) + return Reshape(expr, self.shape) + + def as_explicit(self): + ee = self.expr + if hasattr(ee, "as_explicit"): + ee = ee.as_explicit() + if isinstance(ee, MatrixBase): + from sympy import Array + ee = Array(ee) + elif isinstance(ee, MatrixExpr): + return self + return ee.reshape(*self.shape) + + +class _ArgE: + """ + The ``_ArgE`` object contains references to the array expression + (``.element``) and a list containing the information about index + contractions (``.indices``). + + Index contractions are numbered and contracted indices show the number of + the contraction. Uncontracted indices have ``None`` value. + + For example: + ``_ArgE(M, [None, 3])`` + This object means that expression ``M`` is part of an array contraction + and has two indices, the first is not contracted (value ``None``), + the second index is contracted to the 4th (i.e. number ``3``) group of the + array contraction object. + """ + indices: List[Optional[int]] + + def __init__(self, element, indices: Optional[List[Optional[int]]] = None): + self.element = element + if indices is None: + self.indices = [None for i in range(get_rank(element))] + else: + self.indices = indices + + def __str__(self): + return "_ArgE(%s, %s)" % (self.element, self.indices) + + __repr__ = __str__ + + +class _IndPos: + """ + Index position, requiring two integers in the constructor: + + - arg: the position of the argument in the tensor product, + - rel: the relative position of the index inside the argument. + """ + def __init__(self, arg: int, rel: int): + self.arg = arg + self.rel = rel + + def __str__(self): + return "_IndPos(%i, %i)" % (self.arg, self.rel) + + __repr__ = __str__ + + def __iter__(self): + yield from [self.arg, self.rel] + + +class _EditArrayContraction: + """ + Utility class to help manipulate array contraction objects. + + This class takes as input an ``ArrayContraction`` object and turns it into + an editable object. + + The field ``args_with_ind`` of this class is a list of ``_ArgE`` objects + which can be used to easily edit the contraction structure of the + expression. + + Once editing is finished, the ``ArrayContraction`` object may be recreated + by calling the ``.to_array_contraction()`` method. + """ + + def __init__(self, base_array: typing.Union[ArrayContraction, ArrayDiagonal, ArrayTensorProduct]): + + expr: Basic + diagonalized: tTuple[tTuple[int, ...], ...] + contraction_indices: List[tTuple[int]] + if isinstance(base_array, ArrayContraction): + mapping = _get_mapping_from_subranks(base_array.subranks) + expr = base_array.expr + contraction_indices = base_array.contraction_indices + diagonalized = () + elif isinstance(base_array, ArrayDiagonal): + + if isinstance(base_array.expr, ArrayContraction): + mapping = _get_mapping_from_subranks(base_array.expr.subranks) + expr = base_array.expr.expr + diagonalized = ArrayContraction._push_indices_down(base_array.expr.contraction_indices, base_array.diagonal_indices) + contraction_indices = base_array.expr.contraction_indices + elif isinstance(base_array.expr, ArrayTensorProduct): + mapping = {} + expr = base_array.expr + diagonalized = base_array.diagonal_indices + contraction_indices = [] + else: + mapping = {} + expr = base_array.expr + diagonalized = base_array.diagonal_indices + contraction_indices = [] + + elif isinstance(base_array, ArrayTensorProduct): + expr = base_array + contraction_indices = [] + diagonalized = () + else: + raise NotImplementedError() + + if isinstance(expr, ArrayTensorProduct): + args = list(expr.args) + else: + args = [expr] + + args_with_ind: List[_ArgE] = [_ArgE(arg) for arg in args] + for i, contraction_tuple in enumerate(contraction_indices): + for j in contraction_tuple: + arg_pos, rel_pos = mapping[j] + args_with_ind[arg_pos].indices[rel_pos] = i + self.args_with_ind: List[_ArgE] = args_with_ind + self.number_of_contraction_indices: int = len(contraction_indices) + self._track_permutation: Optional[List[List[int]]] = None + + mapping = _get_mapping_from_subranks(base_array.subranks) + + # Trick: add diagonalized indices as negative indices into the editor object: + for i, e in enumerate(diagonalized): + for j in e: + arg_pos, rel_pos = mapping[j] + self.args_with_ind[arg_pos].indices[rel_pos] = -1 - i + + def insert_after(self, arg: _ArgE, new_arg: _ArgE): + pos = self.args_with_ind.index(arg) + self.args_with_ind.insert(pos + 1, new_arg) + + def get_new_contraction_index(self): + self.number_of_contraction_indices += 1 + return self.number_of_contraction_indices - 1 + + def refresh_indices(self): + updates = {} + for arg_with_ind in self.args_with_ind: + updates.update({i: -1 for i in arg_with_ind.indices if i is not None}) + for i, e in enumerate(sorted(updates)): + updates[e] = i + self.number_of_contraction_indices = len(updates) + for arg_with_ind in self.args_with_ind: + arg_with_ind.indices = [updates.get(i, None) for i in arg_with_ind.indices] + + def merge_scalars(self): + scalars = [] + for arg_with_ind in self.args_with_ind: + if len(arg_with_ind.indices) == 0: + scalars.append(arg_with_ind) + for i in scalars: + self.args_with_ind.remove(i) + scalar = Mul.fromiter([i.element for i in scalars]) + if len(self.args_with_ind) == 0: + self.args_with_ind.append(_ArgE(scalar)) + else: + from sympy.tensor.array.expressions.from_array_to_matrix import _a2m_tensor_product + self.args_with_ind[0].element = _a2m_tensor_product(scalar, self.args_with_ind[0].element) + + def to_array_contraction(self): + + # Count the ranks of the arguments: + counter = 0 + # Create a collector for the new diagonal indices: + diag_indices = defaultdict(list) + + count_index_freq = Counter() + for arg_with_ind in self.args_with_ind: + count_index_freq.update(Counter(arg_with_ind.indices)) + + free_index_count = count_index_freq[None] + + # Construct the inverse permutation: + inv_perm1 = [] + inv_perm2 = [] + # Keep track of which diagonal indices have already been processed: + done = set() + + # Counter for the diagonal indices: + counter4 = 0 + + for arg_with_ind in self.args_with_ind: + # If some diagonalization axes have been removed, they should be + # permuted in order to keep the permutation. + # Add permutation here + counter2 = 0 # counter for the indices + for i in arg_with_ind.indices: + if i is None: + inv_perm1.append(counter4) + counter2 += 1 + counter4 += 1 + continue + if i >= 0: + continue + # Reconstruct the diagonal indices: + diag_indices[-1 - i].append(counter + counter2) + if count_index_freq[i] == 1 and i not in done: + inv_perm1.append(free_index_count - 1 - i) + done.add(i) + elif i not in done: + inv_perm2.append(free_index_count - 1 - i) + done.add(i) + counter2 += 1 + # Remove negative indices to restore a proper editor object: + arg_with_ind.indices = [i if i is not None and i >= 0 else None for i in arg_with_ind.indices] + counter += len([i for i in arg_with_ind.indices if i is None or i < 0]) + + inverse_permutation = inv_perm1 + inv_perm2 + permutation = _af_invert(inverse_permutation) + + # Get the diagonal indices after the detection of HadamardProduct in the expression: + diag_indices_filtered = [tuple(v) for v in diag_indices.values() if len(v) > 1] + + self.merge_scalars() + self.refresh_indices() + args = [arg.element for arg in self.args_with_ind] + contraction_indices = self.get_contraction_indices() + expr = _array_contraction(_array_tensor_product(*args), *contraction_indices) + expr2 = _array_diagonal(expr, *diag_indices_filtered) + if self._track_permutation is not None: + permutation2 = _af_invert([j for i in self._track_permutation for j in i]) + expr2 = _permute_dims(expr2, permutation2) + + expr3 = _permute_dims(expr2, permutation) + return expr3 + + def get_contraction_indices(self) -> List[List[int]]: + contraction_indices: List[List[int]] = [[] for i in range(self.number_of_contraction_indices)] + current_position: int = 0 + for arg_with_ind in self.args_with_ind: + for j in arg_with_ind.indices: + if j is not None: + contraction_indices[j].append(current_position) + current_position += 1 + return contraction_indices + + def get_mapping_for_index(self, ind) -> List[_IndPos]: + if ind >= self.number_of_contraction_indices: + raise ValueError("index value exceeding the index range") + positions: List[_IndPos] = [] + for i, arg_with_ind in enumerate(self.args_with_ind): + for j, arg_ind in enumerate(arg_with_ind.indices): + if ind == arg_ind: + positions.append(_IndPos(i, j)) + return positions + + def get_contraction_indices_to_ind_rel_pos(self) -> List[List[_IndPos]]: + contraction_indices: List[List[_IndPos]] = [[] for i in range(self.number_of_contraction_indices)] + for i, arg_with_ind in enumerate(self.args_with_ind): + for j, ind in enumerate(arg_with_ind.indices): + if ind is not None: + contraction_indices[ind].append(_IndPos(i, j)) + return contraction_indices + + def count_args_with_index(self, index: int) -> int: + """ + Count the number of arguments that have the given index. + """ + counter: int = 0 + for arg_with_ind in self.args_with_ind: + if index in arg_with_ind.indices: + counter += 1 + return counter + + def get_args_with_index(self, index: int) -> List[_ArgE]: + """ + Get a list of arguments having the given index. + """ + ret: List[_ArgE] = [i for i in self.args_with_ind if index in i.indices] + return ret + + @property + def number_of_diagonal_indices(self): + data = set() + for arg in self.args_with_ind: + data.update({i for i in arg.indices if i is not None and i < 0}) + return len(data) + + def track_permutation_start(self): + permutation = [] + perm_diag = [] + counter = 0 + counter2 = -1 + for arg_with_ind in self.args_with_ind: + perm = [] + for i in arg_with_ind.indices: + if i is not None: + if i < 0: + perm_diag.append(counter2) + counter2 -= 1 + continue + perm.append(counter) + counter += 1 + permutation.append(perm) + max_ind = max(max(i) if i else -1 for i in permutation) if permutation else -1 + perm_diag = [max_ind - i for i in perm_diag] + self._track_permutation = permutation + [perm_diag] + + def track_permutation_merge(self, destination: _ArgE, from_element: _ArgE): + index_destination = self.args_with_ind.index(destination) + index_element = self.args_with_ind.index(from_element) + self._track_permutation[index_destination].extend(self._track_permutation[index_element]) # type: ignore + self._track_permutation.pop(index_element) # type: ignore + + def get_absolute_free_range(self, arg: _ArgE) -> typing.Tuple[int, int]: + """ + Return the range of the free indices of the arg as absolute positions + among all free indices. + """ + counter = 0 + for arg_with_ind in self.args_with_ind: + number_free_indices = len([i for i in arg_with_ind.indices if i is None]) + if arg_with_ind == arg: + return counter, counter + number_free_indices + counter += number_free_indices + raise IndexError("argument not found") + + def get_absolute_range(self, arg: _ArgE) -> typing.Tuple[int, int]: + """ + Return the absolute range of indices for arg, disregarding dummy + indices. + """ + counter = 0 + for arg_with_ind in self.args_with_ind: + number_indices = len(arg_with_ind.indices) + if arg_with_ind == arg: + return counter, counter + number_indices + counter += number_indices + raise IndexError("argument not found") + + +def get_rank(expr): + if isinstance(expr, (MatrixExpr, MatrixElement)): + return 2 + if isinstance(expr, _CodegenArrayAbstract): + return len(expr.shape) + if isinstance(expr, NDimArray): + return expr.rank() + if isinstance(expr, Indexed): + return expr.rank + if isinstance(expr, IndexedBase): + shape = expr.shape + if shape is None: + return -1 + else: + return len(shape) + if hasattr(expr, "shape"): + return len(expr.shape) + return 0 + + +def _get_subrank(expr): + if isinstance(expr, _CodegenArrayAbstract): + return expr.subrank() + return get_rank(expr) + + +def _get_subranks(expr): + if isinstance(expr, _CodegenArrayAbstract): + return expr.subranks + else: + return [get_rank(expr)] + + +def get_shape(expr): + if hasattr(expr, "shape"): + return expr.shape + return () + + +def nest_permutation(expr): + if isinstance(expr, PermuteDims): + return expr.nest_permutation() + else: + return expr + + +def _array_tensor_product(*args, **kwargs): + return ArrayTensorProduct(*args, canonicalize=True, **kwargs) + + +def _array_contraction(expr, *contraction_indices, **kwargs): + return ArrayContraction(expr, *contraction_indices, canonicalize=True, **kwargs) + + +def _array_diagonal(expr, *diagonal_indices, **kwargs): + return ArrayDiagonal(expr, *diagonal_indices, canonicalize=True, **kwargs) + + +def _permute_dims(expr, permutation, **kwargs): + return PermuteDims(expr, permutation, canonicalize=True, **kwargs) + + +def _array_add(*args, **kwargs): + return ArrayAdd(*args, canonicalize=True, **kwargs) + + +def _get_array_element_or_slice(expr, indices): + return ArrayElement(expr, indices) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/arrayexpr_derivatives.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/arrayexpr_derivatives.py new file mode 100644 index 0000000000000000000000000000000000000000..ab44a6fbf715ac7f2b8c287dcc84a49289f2dd76 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/arrayexpr_derivatives.py @@ -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() diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/conv_array_to_indexed.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/conv_array_to_indexed.py new file mode 100644 index 0000000000000000000000000000000000000000..1929c3401e131cca0a83080131ead9198b37bcbb --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/conv_array_to_indexed.py @@ -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) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/conv_array_to_matrix.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/conv_array_to_matrix.py new file mode 100644 index 0000000000000000000000000000000000000000..2708e74aaa98d6ee38eae46d97d4483a546e0776 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/conv_array_to_matrix.py @@ -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) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/conv_indexed_to_array.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/conv_indexed_to_array.py new file mode 100644 index 0000000000000000000000000000000000000000..6058b31f20778834ea23a01553d594b7965eb6bb --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/conv_indexed_to_array.py @@ -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) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/conv_matrix_to_array.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/conv_matrix_to_array.py new file mode 100644 index 0000000000000000000000000000000000000000..46469df60703c237527c0b2834235309640afe7c --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/conv_matrix_to_array.py @@ -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) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/from_array_to_indexed.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/from_array_to_indexed.py new file mode 100644 index 0000000000000000000000000000000000000000..9eb86e7cfbe31ebfe7c9649803d9cb5e34b98276 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/from_array_to_indexed.py @@ -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) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/from_array_to_matrix.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/from_array_to_matrix.py new file mode 100644 index 0000000000000000000000000000000000000000..46bd3bb2dfca6d2ad7f27eef9a11af7e1e3ac3cc --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/from_array_to_matrix.py @@ -0,0 +1,1003 @@ +import itertools +from collections import defaultdict +from typing import Tuple as tTuple, Union as tUnion, FrozenSet, Dict as tDict, List, Optional +from functools import singledispatch +from itertools import accumulate + +from sympy import MatMul, Basic, Wild, KroneckerProduct +from sympy.assumptions.ask import (Q, ask) +from sympy.core.mul import Mul +from sympy.core.singleton import S +from sympy.matrices.expressions.diagonal import DiagMatrix +from sympy.matrices.expressions.hadamard import hadamard_product, HadamardPower +from sympy.matrices.expressions.matexpr import MatrixExpr +from sympy.matrices.expressions.special import (Identity, ZeroMatrix, OneMatrix) +from sympy.matrices.expressions.trace import Trace +from sympy.matrices.expressions.transpose import Transpose +from sympy.combinatorics.permutations import _af_invert, Permutation +from sympy.matrices.matrixbase import MatrixBase +from sympy.matrices.expressions.applyfunc import ElementwiseApplyFunction +from sympy.matrices.expressions.matexpr import MatrixElement +from sympy.tensor.array.expressions.array_expressions import PermuteDims, ArrayDiagonal, \ + ArrayTensorProduct, OneArray, get_rank, _get_subrank, ZeroArray, ArrayContraction, \ + ArrayAdd, _CodegenArrayAbstract, get_shape, ArrayElementwiseApplyFunc, _ArrayExpr, _EditArrayContraction, _ArgE, \ + ArrayElement, _array_tensor_product, _array_contraction, _array_diagonal, _array_add, _permute_dims +from sympy.tensor.array.expressions.utils import _get_mapping_from_subranks + + +def _get_candidate_for_matmul_from_contraction(scan_indices: List[Optional[int]], remaining_args: List[_ArgE]) -> tTuple[Optional[_ArgE], bool, int]: + + scan_indices_int: List[int] = [i for i in scan_indices if i is not None] + if len(scan_indices_int) == 0: + return None, False, -1 + + transpose: bool = False + candidate: Optional[_ArgE] = None + candidate_index: int = -1 + for arg_with_ind2 in remaining_args: + if not isinstance(arg_with_ind2.element, MatrixExpr): + continue + for index in scan_indices_int: + if candidate_index != -1 and candidate_index != index: + # A candidate index has already been selected, check + # repetitions only for that index: + continue + if index in arg_with_ind2.indices: + if set(arg_with_ind2.indices) == {index}: + # Index repeated twice in arg_with_ind2 + candidate = None + break + if candidate is None: + candidate = arg_with_ind2 + candidate_index = index + transpose = (index == arg_with_ind2.indices[1]) + else: + # Index repeated more than twice, break + candidate = None + break + return candidate, transpose, candidate_index + + +def _insert_candidate_into_editor(editor: _EditArrayContraction, arg_with_ind: _ArgE, candidate: _ArgE, transpose1: bool, transpose2: bool): + other = candidate.element + other_index: Optional[int] + if transpose2: + other = Transpose(other) + other_index = candidate.indices[0] + else: + other_index = candidate.indices[1] + new_element = (Transpose(arg_with_ind.element) if transpose1 else arg_with_ind.element) * other + editor.args_with_ind.remove(candidate) + new_arge = _ArgE(new_element) + return new_arge, other_index + + +def _support_function_tp1_recognize(contraction_indices, args): + if len(contraction_indices) == 0: + return _a2m_tensor_product(*args) + + ac = _array_contraction(_array_tensor_product(*args), *contraction_indices) + editor = _EditArrayContraction(ac) + editor.track_permutation_start() + + while True: + flag_stop = True + for i, arg_with_ind in enumerate(editor.args_with_ind): + if not isinstance(arg_with_ind.element, MatrixExpr): + continue + + first_index = arg_with_ind.indices[0] + second_index = arg_with_ind.indices[1] + + first_frequency = editor.count_args_with_index(first_index) + second_frequency = editor.count_args_with_index(second_index) + + if first_index is not None and first_frequency == 1 and first_index == second_index: + flag_stop = False + arg_with_ind.element = Trace(arg_with_ind.element)._normalize() + arg_with_ind.indices = [] + break + + scan_indices = [] + if first_frequency == 2: + scan_indices.append(first_index) + if second_frequency == 2: + scan_indices.append(second_index) + + candidate, transpose, found_index = _get_candidate_for_matmul_from_contraction(scan_indices, editor.args_with_ind[i+1:]) + if candidate is not None: + flag_stop = False + editor.track_permutation_merge(arg_with_ind, candidate) + transpose1 = found_index == first_index + new_arge, other_index = _insert_candidate_into_editor(editor, arg_with_ind, candidate, transpose1, transpose) + if found_index == first_index: + new_arge.indices = [second_index, other_index] + else: + new_arge.indices = [first_index, other_index] + set_indices = set(new_arge.indices) + if len(set_indices) == 1 and set_indices != {None}: + # This is a trace: + new_arge.element = Trace(new_arge.element)._normalize() + new_arge.indices = [] + editor.args_with_ind[i] = new_arge + # TODO: is this break necessary? + break + + if flag_stop: + break + + editor.refresh_indices() + return editor.to_array_contraction() + + +def _find_trivial_matrices_rewrite(expr: ArrayTensorProduct): + # If there are matrices of trivial shape in the tensor product (i.e. shape + # (1, 1)), try to check if there is a suitable non-trivial MatMul where the + # expression can be inserted. + + # For example, if "a" has shape (1, 1) and "b" has shape (k, 1), the + # expressions "_array_tensor_product(a, b*b.T)" can be rewritten as + # "b*a*b.T" + + trivial_matrices = [] + pos: Optional[int] = None + first: Optional[MatrixExpr] = None + second: Optional[MatrixExpr] = None + removed: List[int] = [] + counter: int = 0 + args: List[Optional[Basic]] = list(expr.args) + for i, arg in enumerate(expr.args): + if isinstance(arg, MatrixExpr): + if arg.shape == (1, 1): + trivial_matrices.append(arg) + args[i] = None + removed.extend([counter, counter+1]) + elif pos is None and isinstance(arg, MatMul): + margs = arg.args + for j, e in enumerate(margs): + if isinstance(e, MatrixExpr) and e.shape[1] == 1: + pos = i + first = MatMul.fromiter(margs[:j+1]) + second = MatMul.fromiter(margs[j+1:]) + break + counter += get_rank(arg) + if pos is None: + return expr, [] + args[pos] = (first*MatMul.fromiter(i for i in trivial_matrices)*second).doit() + return _array_tensor_product(*[i for i in args if i is not None]), removed + + +def _find_trivial_kronecker_products_broadcast(expr: ArrayTensorProduct): + newargs: List[Basic] = [] + removed = [] + count_dims = 0 + for arg in expr.args: + count_dims += get_rank(arg) + shape = get_shape(arg) + current_range = [count_dims-i for i in range(len(shape), 0, -1)] + if (shape == (1, 1) and len(newargs) > 0 and 1 not in get_shape(newargs[-1]) and + isinstance(newargs[-1], MatrixExpr) and isinstance(arg, MatrixExpr)): + # KroneckerProduct object allows the trick of broadcasting: + newargs[-1] = KroneckerProduct(newargs[-1], arg) + removed.extend(current_range) + elif 1 not in shape and len(newargs) > 0 and get_shape(newargs[-1]) == (1, 1): + # Broadcast: + newargs[-1] = KroneckerProduct(newargs[-1], arg) + prev_range = [i for i in range(min(current_range)) if i not in removed] + removed.extend(prev_range[-2:]) + else: + newargs.append(arg) + return _array_tensor_product(*newargs), removed + + +@singledispatch +def _array2matrix(expr): + return expr + + +@_array2matrix.register(ZeroArray) +def _(expr: ZeroArray): + if get_rank(expr) == 2: + return ZeroMatrix(*expr.shape) + else: + return expr + + +@_array2matrix.register(ArrayTensorProduct) +def _(expr: ArrayTensorProduct): + return _a2m_tensor_product(*[_array2matrix(arg) for arg in expr.args]) + + +@_array2matrix.register(ArrayContraction) +def _(expr: ArrayContraction): + expr = expr.flatten_contraction_of_diagonal() + expr = identify_removable_identity_matrices(expr) + expr = expr.split_multiple_contractions() + expr = identify_hadamard_products(expr) + if not isinstance(expr, ArrayContraction): + return _array2matrix(expr) + subexpr = expr.expr + contraction_indices: tTuple[tTuple[int]] = expr.contraction_indices + if contraction_indices == ((0,), (1,)) or ( + contraction_indices == ((0,),) and subexpr.shape[1] == 1 + ) or ( + contraction_indices == ((1,),) and subexpr.shape[0] == 1 + ): + shape = subexpr.shape + subexpr = _array2matrix(subexpr) + if isinstance(subexpr, MatrixExpr): + return OneMatrix(1, shape[0])*subexpr*OneMatrix(shape[1], 1) + if isinstance(subexpr, ArrayTensorProduct): + newexpr = _array_contraction(_array2matrix(subexpr), *contraction_indices) + contraction_indices = newexpr.contraction_indices + if any(i > 2 for i in newexpr.subranks): + addends = _array_add(*[_a2m_tensor_product(*j) for j in itertools.product(*[i.args if isinstance(i, + ArrayAdd) else [i] for i in expr.expr.args])]) + newexpr = _array_contraction(addends, *contraction_indices) + if isinstance(newexpr, ArrayAdd): + ret = _array2matrix(newexpr) + return ret + assert isinstance(newexpr, ArrayContraction) + ret = _support_function_tp1_recognize(contraction_indices, list(newexpr.expr.args)) + return ret + elif not isinstance(subexpr, _CodegenArrayAbstract): + ret = _array2matrix(subexpr) + if isinstance(ret, MatrixExpr): + assert expr.contraction_indices == ((0, 1),) + return _a2m_trace(ret) + else: + return _array_contraction(ret, *expr.contraction_indices) + + +@_array2matrix.register(ArrayDiagonal) +def _(expr: ArrayDiagonal): + pexpr = _array_diagonal(_array2matrix(expr.expr), *expr.diagonal_indices) + pexpr = identify_hadamard_products(pexpr) + if isinstance(pexpr, ArrayDiagonal): + pexpr = _array_diag2contr_diagmatrix(pexpr) + if expr == pexpr: + return expr + return _array2matrix(pexpr) + + +@_array2matrix.register(PermuteDims) +def _(expr: PermuteDims): + if expr.permutation.array_form == [1, 0]: + return _a2m_transpose(_array2matrix(expr.expr)) + elif isinstance(expr.expr, ArrayTensorProduct): + ranks = expr.expr.subranks + inv_permutation = expr.permutation**(-1) + newrange = [inv_permutation(i) for i in range(sum(ranks))] + newpos = [] + counter = 0 + for rank in ranks: + newpos.append(newrange[counter:counter+rank]) + counter += rank + newargs = [] + newperm = [] + scalars = [] + for pos, arg in zip(newpos, expr.expr.args): + if len(pos) == 0: + scalars.append(_array2matrix(arg)) + elif pos == sorted(pos): + newargs.append((_array2matrix(arg), pos[0])) + newperm.extend(pos) + elif len(pos) == 2: + newargs.append((_a2m_transpose(_array2matrix(arg)), pos[0])) + newperm.extend(reversed(pos)) + else: + raise NotImplementedError() + newargs = [i[0] for i in newargs] + return _permute_dims(_a2m_tensor_product(*scalars, *newargs), _af_invert(newperm)) + elif isinstance(expr.expr, ArrayContraction): + mat_mul_lines = _array2matrix(expr.expr) + if not isinstance(mat_mul_lines, ArrayTensorProduct): + return _permute_dims(mat_mul_lines, expr.permutation) + # TODO: this assumes that all arguments are matrices, it may not be the case: + permutation = Permutation(2*len(mat_mul_lines.args)-1)*expr.permutation + permuted = [permutation(i) for i in range(2*len(mat_mul_lines.args))] + args_array = [None for i in mat_mul_lines.args] + for i in range(len(mat_mul_lines.args)): + p1 = permuted[2*i] + p2 = permuted[2*i+1] + if p1 // 2 != p2 // 2: + return _permute_dims(mat_mul_lines, permutation) + if p1 > p2: + args_array[i] = _a2m_transpose(mat_mul_lines.args[p1 // 2]) + else: + args_array[i] = mat_mul_lines.args[p1 // 2] + return _a2m_tensor_product(*args_array) + else: + return expr + + +@_array2matrix.register(ArrayAdd) +def _(expr: ArrayAdd): + addends = [_array2matrix(arg) for arg in expr.args] + return _a2m_add(*addends) + + +@_array2matrix.register(ArrayElementwiseApplyFunc) +def _(expr: ArrayElementwiseApplyFunc): + subexpr = _array2matrix(expr.expr) + if isinstance(subexpr, MatrixExpr): + if subexpr.shape != (1, 1): + d = expr.function.bound_symbols[0] + w = Wild("w", exclude=[d]) + p = Wild("p", exclude=[d]) + m = expr.function.expr.match(w*d**p) + if m is not None: + return m[w]*HadamardPower(subexpr, m[p]) + return ElementwiseApplyFunction(expr.function, subexpr) + else: + return ArrayElementwiseApplyFunc(expr.function, subexpr) + + +@_array2matrix.register(ArrayElement) +def _(expr: ArrayElement): + ret = _array2matrix(expr.name) + if isinstance(ret, MatrixExpr): + return MatrixElement(ret, *expr.indices) + return ArrayElement(ret, expr.indices) + + +@singledispatch +def _remove_trivial_dims(expr): + return expr, [] + + +@_remove_trivial_dims.register(ArrayTensorProduct) +def _(expr: ArrayTensorProduct): + # Recognize expressions like [x, y] with shape (k, 1, k, 1) as `x*y.T`. + # The matrix expression has to be equivalent to the tensor product of the + # matrices, with trivial dimensions (i.e. dim=1) dropped. + # That is, add contractions over trivial dimensions: + + removed = [] + newargs = [] + cumul = list(accumulate([0] + [get_rank(arg) for arg in expr.args])) + pending = None + prev_i = None + for i, arg in enumerate(expr.args): + current_range = list(range(cumul[i], cumul[i+1])) + if isinstance(arg, OneArray): + removed.extend(current_range) + continue + if not isinstance(arg, (MatrixExpr, MatrixBase)): + rarg, rem = _remove_trivial_dims(arg) + removed.extend(rem) + newargs.append(rarg) + continue + elif getattr(arg, "is_Identity", False) and arg.shape == (1, 1): + if arg.shape == (1, 1): + # Ignore identity matrices of shape (1, 1) - they are equivalent to scalar 1. + removed.extend(current_range) + continue + elif arg.shape == (1, 1): + arg, _ = _remove_trivial_dims(arg) + # Matrix is equivalent to scalar: + if len(newargs) == 0: + newargs.append(arg) + elif 1 in get_shape(newargs[-1]): + if newargs[-1].shape[1] == 1: + newargs[-1] = newargs[-1]*arg + else: + newargs[-1] = arg*newargs[-1] + removed.extend(current_range) + else: + newargs.append(arg) + elif 1 in arg.shape: + k = [i for i in arg.shape if i != 1][0] + if pending is None: + pending = k + prev_i = i + newargs.append(arg) + elif pending == k: + prev = newargs[-1] + if prev.shape[0] == 1: + d1 = cumul[prev_i] + prev = _a2m_transpose(prev) + else: + d1 = cumul[prev_i] + 1 + if arg.shape[1] == 1: + d2 = cumul[i] + 1 + arg = _a2m_transpose(arg) + else: + d2 = cumul[i] + newargs[-1] = prev*arg + pending = None + removed.extend([d1, d2]) + else: + newargs.append(arg) + pending = k + prev_i = i + else: + newargs.append(arg) + pending = None + newexpr, newremoved = _a2m_tensor_product(*newargs), sorted(removed) + if isinstance(newexpr, ArrayTensorProduct): + newexpr, newremoved2 = _find_trivial_matrices_rewrite(newexpr) + newremoved = _combine_removed(-1, newremoved, newremoved2) + if isinstance(newexpr, ArrayTensorProduct): + newexpr, newremoved2 = _find_trivial_kronecker_products_broadcast(newexpr) + newremoved = _combine_removed(-1, newremoved, newremoved2) + return newexpr, newremoved + + +@_remove_trivial_dims.register(ArrayAdd) +def _(expr: ArrayAdd): + rec = [_remove_trivial_dims(arg) for arg in expr.args] + newargs, removed = zip(*rec) + if len({get_shape(i) for i in newargs}) > 1: + return expr, [] + if len(removed) == 0: + return expr, removed + removed1 = removed[0] + return _a2m_add(*newargs), removed1 + + +@_remove_trivial_dims.register(PermuteDims) +def _(expr: PermuteDims): + subexpr, subremoved = _remove_trivial_dims(expr.expr) + p = expr.permutation.array_form + pinv = _af_invert(expr.permutation.array_form) + shift = list(accumulate([1 if i in subremoved else 0 for i in range(len(p))])) + premoved = [pinv[i] for i in subremoved] + p2 = [e - shift[e] for e in p if e not in subremoved] + # TODO: check if subremoved should be permuted as well... + newexpr = _permute_dims(subexpr, p2) + premoved = sorted(premoved) + if newexpr != expr: + newexpr, removed2 = _remove_trivial_dims(_array2matrix(newexpr)) + premoved = _combine_removed(-1, premoved, removed2) + return newexpr, premoved + + +@_remove_trivial_dims.register(ArrayContraction) +def _(expr: ArrayContraction): + new_expr, removed0 = _array_contraction_to_diagonal_multiple_identity(expr) + if new_expr != expr: + new_expr2, removed1 = _remove_trivial_dims(_array2matrix(new_expr)) + removed = _combine_removed(-1, removed0, removed1) + return new_expr2, removed + rank1 = get_rank(expr) + expr, removed1 = remove_identity_matrices(expr) + if not isinstance(expr, ArrayContraction): + expr2, removed2 = _remove_trivial_dims(expr) + return expr2, _combine_removed(rank1, removed1, removed2) + newexpr, removed2 = _remove_trivial_dims(expr.expr) + shifts = list(accumulate([1 if i in removed2 else 0 for i in range(get_rank(expr.expr))])) + new_contraction_indices = [tuple(j for j in i if j not in removed2) for i in expr.contraction_indices] + # Remove possible empty tuples "()": + new_contraction_indices = [i for i in new_contraction_indices if len(i) > 0] + contraction_indices_flat = [j for i in expr.contraction_indices for j in i] + removed2 = [i for i in removed2 if i not in contraction_indices_flat] + new_contraction_indices = [tuple(j - shifts[j] for j in i) for i in new_contraction_indices] + # Shift removed2: + removed2 = ArrayContraction._push_indices_up(expr.contraction_indices, removed2) + removed = _combine_removed(rank1, removed1, removed2) + return _array_contraction(newexpr, *new_contraction_indices), list(removed) + + +def _remove_diagonalized_identity_matrices(expr: ArrayDiagonal): + assert isinstance(expr, ArrayDiagonal) + editor = _EditArrayContraction(expr) + mapping = {i: {j for j in editor.args_with_ind if i in j.indices} for i in range(-1, -1-editor.number_of_diagonal_indices, -1)} + removed = [] + counter: int = 0 + for i, arg_with_ind in enumerate(editor.args_with_ind): + counter += len(arg_with_ind.indices) + if isinstance(arg_with_ind.element, Identity): + if None in arg_with_ind.indices and any(i is not None and (i < 0) == True for i in arg_with_ind.indices): + diag_ind = [j for j in arg_with_ind.indices if j is not None][0] + other = [j for j in mapping[diag_ind] if j != arg_with_ind][0] + if not isinstance(other.element, MatrixExpr): + continue + if 1 not in other.element.shape: + continue + if None not in other.indices: + continue + editor.args_with_ind[i].element = None + none_index = other.indices.index(None) + other.element = DiagMatrix(other.element) + other_range = editor.get_absolute_range(other) + removed.extend([other_range[0] + none_index]) + editor.args_with_ind = [i for i in editor.args_with_ind if i.element is not None] + removed = ArrayDiagonal._push_indices_up(expr.diagonal_indices, removed, get_rank(expr.expr)) + return editor.to_array_contraction(), removed + + +@_remove_trivial_dims.register(ArrayDiagonal) +def _(expr: ArrayDiagonal): + newexpr, removed = _remove_trivial_dims(expr.expr) + shifts = list(accumulate([0] + [1 if i in removed else 0 for i in range(get_rank(expr.expr))])) + new_diag_indices_map = {i: tuple(j for j in i if j not in removed) for i in expr.diagonal_indices} + for old_diag_tuple, new_diag_tuple in new_diag_indices_map.items(): + if len(new_diag_tuple) == 1: + removed = [i for i in removed if i not in old_diag_tuple] + new_diag_indices = [tuple(j - shifts[j] for j in i) for i in new_diag_indices_map.values()] + rank = get_rank(expr.expr) + removed = ArrayDiagonal._push_indices_up(expr.diagonal_indices, removed, rank) + removed = sorted(set(removed)) + # If there are single axes to diagonalize remaining, it means that their + # corresponding dimension has been removed, they no longer need diagonalization: + new_diag_indices = [i for i in new_diag_indices if len(i) > 0] + if len(new_diag_indices) > 0: + newexpr2 = _array_diagonal(newexpr, *new_diag_indices, allow_trivial_diags=True) + else: + newexpr2 = newexpr + if isinstance(newexpr2, ArrayDiagonal): + newexpr3, removed2 = _remove_diagonalized_identity_matrices(newexpr2) + removed = _combine_removed(-1, removed, removed2) + return newexpr3, removed + else: + return newexpr2, removed + + +@_remove_trivial_dims.register(ElementwiseApplyFunction) +def _(expr: ElementwiseApplyFunction): + subexpr, removed = _remove_trivial_dims(expr.expr) + if subexpr.shape == (1, 1): + # TODO: move this to ElementwiseApplyFunction + return expr.function(subexpr), removed + [0, 1] + return ElementwiseApplyFunction(expr.function, subexpr), [] + + +@_remove_trivial_dims.register(ArrayElementwiseApplyFunc) +def _(expr: ArrayElementwiseApplyFunc): + subexpr, removed = _remove_trivial_dims(expr.expr) + return ArrayElementwiseApplyFunc(expr.function, subexpr), removed + + +def convert_array_to_matrix(expr): + r""" + Recognize matrix expressions in codegen objects. + + If more than one matrix multiplication line have been detected, return a + list with the matrix expressions. + + Examples + ======== + + >>> from sympy.tensor.array.expressions.from_indexed_to_array import convert_indexed_to_array + >>> from sympy.tensor.array import tensorcontraction, tensorproduct + >>> from sympy import MatrixSymbol, Sum + >>> from sympy.abc import i, j, k, l, N + >>> from sympy.tensor.array.expressions.from_matrix_to_array import convert_matrix_to_array + >>> from sympy.tensor.array.expressions.from_array_to_matrix import convert_array_to_matrix + >>> A = MatrixSymbol("A", N, N) + >>> B = MatrixSymbol("B", N, N) + >>> C = MatrixSymbol("C", N, N) + >>> D = MatrixSymbol("D", N, N) + + >>> expr = Sum(A[i, j]*B[j, k], (j, 0, N-1)) + >>> cg = convert_indexed_to_array(expr) + >>> convert_array_to_matrix(cg) + A*B + >>> cg = convert_indexed_to_array(expr, first_indices=[k]) + >>> convert_array_to_matrix(cg) + B.T*A.T + + Transposition is detected: + + >>> expr = Sum(A[j, i]*B[j, k], (j, 0, N-1)) + >>> cg = convert_indexed_to_array(expr) + >>> convert_array_to_matrix(cg) + A.T*B + >>> cg = convert_indexed_to_array(expr, first_indices=[k]) + >>> convert_array_to_matrix(cg) + B.T*A + + Detect the trace: + + >>> expr = Sum(A[i, i], (i, 0, N-1)) + >>> cg = convert_indexed_to_array(expr) + >>> convert_array_to_matrix(cg) + Trace(A) + + Recognize some more complex traces: + + >>> expr = Sum(A[i, j]*B[j, i], (i, 0, N-1), (j, 0, N-1)) + >>> cg = convert_indexed_to_array(expr) + >>> convert_array_to_matrix(cg) + Trace(A*B) + + More complicated expressions: + + >>> expr = Sum(A[i, j]*B[k, j]*A[l, k], (j, 0, N-1), (k, 0, N-1)) + >>> cg = convert_indexed_to_array(expr) + >>> convert_array_to_matrix(cg) + A*B.T*A.T + + Expressions constructed from matrix expressions do not contain literal + indices, the positions of free indices are returned instead: + + >>> expr = A*B + >>> cg = convert_matrix_to_array(expr) + >>> convert_array_to_matrix(cg) + A*B + + If more than one line of matrix multiplications is detected, return + separate matrix multiplication factors embedded in a tensor product object: + + >>> cg = tensorcontraction(tensorproduct(A, B, C, D), (1, 2), (5, 6)) + >>> convert_array_to_matrix(cg) + ArrayTensorProduct(A*B, C*D) + + The two lines have free indices at axes 0, 3 and 4, 7, respectively. + """ + rec = _array2matrix(expr) + rec, removed = _remove_trivial_dims(rec) + return rec + + +def _array_diag2contr_diagmatrix(expr: ArrayDiagonal): + if isinstance(expr.expr, ArrayTensorProduct): + args = list(expr.expr.args) + diag_indices = list(expr.diagonal_indices) + mapping = _get_mapping_from_subranks([_get_subrank(arg) for arg in args]) + tuple_links = [[mapping[j] for j in i] for i in diag_indices] + contr_indices = [] + total_rank = get_rank(expr) + replaced = [False for arg in args] + for i, (abs_pos, rel_pos) in enumerate(zip(diag_indices, tuple_links)): + if len(abs_pos) != 2: + continue + (pos1_outer, pos1_inner), (pos2_outer, pos2_inner) = rel_pos + arg1 = args[pos1_outer] + arg2 = args[pos2_outer] + if get_rank(arg1) != 2 or get_rank(arg2) != 2: + if replaced[pos1_outer]: + diag_indices[i] = None + if replaced[pos2_outer]: + diag_indices[i] = None + continue + pos1_in2 = 1 - pos1_inner + pos2_in2 = 1 - pos2_inner + if arg1.shape[pos1_in2] == 1: + if arg1.shape[pos1_inner] != 1: + darg1 = DiagMatrix(arg1) + else: + darg1 = arg1 + args.append(darg1) + contr_indices.append(((pos2_outer, pos2_inner), (len(args)-1, pos1_inner))) + total_rank += 1 + diag_indices[i] = None + args[pos1_outer] = OneArray(arg1.shape[pos1_in2]) + replaced[pos1_outer] = True + elif arg2.shape[pos2_in2] == 1: + if arg2.shape[pos2_inner] != 1: + darg2 = DiagMatrix(arg2) + else: + darg2 = arg2 + args.append(darg2) + contr_indices.append(((pos1_outer, pos1_inner), (len(args)-1, pos2_inner))) + total_rank += 1 + diag_indices[i] = None + args[pos2_outer] = OneArray(arg2.shape[pos2_in2]) + replaced[pos2_outer] = True + diag_indices_new = [i for i in diag_indices if i is not None] + cumul = list(accumulate([0] + [get_rank(arg) for arg in args])) + contr_indices2 = [tuple(cumul[a] + b for a, b in i) for i in contr_indices] + tc = _array_contraction( + _array_tensor_product(*args), *contr_indices2 + ) + td = _array_diagonal(tc, *diag_indices_new) + return td + return expr + + +def _a2m_mul(*args): + if not any(isinstance(i, _CodegenArrayAbstract) for i in args): + from sympy.matrices.expressions.matmul import MatMul + return MatMul(*args).doit() + else: + return _array_contraction( + _array_tensor_product(*args), + *[(2*i-1, 2*i) for i in range(1, len(args))] + ) + + +def _a2m_tensor_product(*args): + scalars = [] + arrays = [] + for arg in args: + if isinstance(arg, (MatrixExpr, _ArrayExpr, _CodegenArrayAbstract)): + arrays.append(arg) + else: + scalars.append(arg) + scalar = Mul.fromiter(scalars) + if len(arrays) == 0: + return scalar + if scalar != 1: + if isinstance(arrays[0], _CodegenArrayAbstract): + arrays = [scalar] + arrays + else: + arrays[0] *= scalar + return _array_tensor_product(*arrays) + + +def _a2m_add(*args): + if not any(isinstance(i, _CodegenArrayAbstract) for i in args): + from sympy.matrices.expressions.matadd import MatAdd + return MatAdd(*args).doit() + else: + return _array_add(*args) + + +def _a2m_trace(arg): + if isinstance(arg, _CodegenArrayAbstract): + return _array_contraction(arg, (0, 1)) + else: + from sympy.matrices.expressions.trace import Trace + return Trace(arg) + + +def _a2m_transpose(arg): + if isinstance(arg, _CodegenArrayAbstract): + return _permute_dims(arg, [1, 0]) + else: + from sympy.matrices.expressions.transpose import Transpose + return Transpose(arg).doit() + + +def identify_hadamard_products(expr: tUnion[ArrayContraction, ArrayDiagonal]): + + editor: _EditArrayContraction = _EditArrayContraction(expr) + + map_contr_to_args: tDict[FrozenSet, List[_ArgE]] = defaultdict(list) + map_ind_to_inds: tDict[Optional[int], int] = defaultdict(int) + for arg_with_ind in editor.args_with_ind: + for ind in arg_with_ind.indices: + map_ind_to_inds[ind] += 1 + if None in arg_with_ind.indices: + continue + map_contr_to_args[frozenset(arg_with_ind.indices)].append(arg_with_ind) + + k: FrozenSet[int] + v: List[_ArgE] + for k, v in map_contr_to_args.items(): + make_trace: bool = False + if len(k) == 1 and next(iter(k)) >= 0 and sum(next(iter(k)) in i for i in map_contr_to_args) == 1: + # This is a trace: the arguments are fully contracted with only one + # index, and the index isn't used anywhere else: + make_trace = True + first_element = S.One + elif len(k) != 2: + # Hadamard product only defined for matrices: + continue + if len(v) == 1: + # Hadamard product with a single argument makes no sense: + continue + for ind in k: + if map_ind_to_inds[ind] <= 2: + # There is no other contraction, skip: + continue + + def check_transpose(x): + x = [i if i >= 0 else -1-i for i in x] + return x == sorted(x) + + # Check if expression is a trace: + if all(map_ind_to_inds[j] == len(v) and j >= 0 for j in k) and all(j >= 0 for j in k): + # This is a trace + make_trace = True + first_element = v[0].element + if not check_transpose(v[0].indices): + first_element = first_element.T + hadamard_factors = v[1:] + else: + hadamard_factors = v + + # This is a Hadamard product: + + hp = hadamard_product(*[i.element if check_transpose(i.indices) else Transpose(i.element) for i in hadamard_factors]) + hp_indices = v[0].indices + if not check_transpose(hadamard_factors[0].indices): + hp_indices = list(reversed(hp_indices)) + if make_trace: + hp = Trace(first_element*hp.T)._normalize() + hp_indices = [] + editor.insert_after(v[0], _ArgE(hp, hp_indices)) + for i in v: + editor.args_with_ind.remove(i) + + return editor.to_array_contraction() + + +def identify_removable_identity_matrices(expr): + editor = _EditArrayContraction(expr) + + flag = True + while flag: + flag = False + for arg_with_ind in editor.args_with_ind: + if isinstance(arg_with_ind.element, Identity): + k = arg_with_ind.element.shape[0] + # Candidate for removal: + if arg_with_ind.indices == [None, None]: + # Free identity matrix, will be cleared by _remove_trivial_dims: + continue + elif None in arg_with_ind.indices: + ind = [j for j in arg_with_ind.indices if j is not None][0] + counted = editor.count_args_with_index(ind) + if counted == 1: + # Identity matrix contracted only on one index with itself, + # transform to a OneArray(k) element: + editor.insert_after(arg_with_ind, OneArray(k)) + editor.args_with_ind.remove(arg_with_ind) + flag = True + break + elif counted > 2: + # Case counted = 2 is a matrix multiplication by identity matrix, skip it. + # Case counted > 2 is a multiple contraction, + # this is a case where the contraction becomes a diagonalization if the + # identity matrix is dropped. + continue + elif arg_with_ind.indices[0] == arg_with_ind.indices[1]: + ind = arg_with_ind.indices[0] + counted = editor.count_args_with_index(ind) + if counted > 1: + editor.args_with_ind.remove(arg_with_ind) + flag = True + break + else: + # This is a trace, skip it as it will be recognized somewhere else: + pass + elif ask(Q.diagonal(arg_with_ind.element)): + if arg_with_ind.indices == [None, None]: + continue + elif None in arg_with_ind.indices: + pass + elif arg_with_ind.indices[0] == arg_with_ind.indices[1]: + ind = arg_with_ind.indices[0] + counted = editor.count_args_with_index(ind) + if counted == 3: + # A_ai B_bi D_ii ==> A_ai D_ij B_bj + ind_new = editor.get_new_contraction_index() + other_args = [j for j in editor.args_with_ind if j != arg_with_ind] + other_args[1].indices = [ind_new if j == ind else j for j in other_args[1].indices] + arg_with_ind.indices = [ind, ind_new] + flag = True + break + + return editor.to_array_contraction() + + +def remove_identity_matrices(expr: ArrayContraction): + editor = _EditArrayContraction(expr) + removed: List[int] = [] + + permutation_map = {} + + free_indices = list(accumulate([0] + [sum(i is None for i in arg.indices) for arg in editor.args_with_ind])) + free_map = dict(zip(editor.args_with_ind, free_indices[:-1])) + + update_pairs = {} + + for ind in range(editor.number_of_contraction_indices): + args = editor.get_args_with_index(ind) + identity_matrices = [i for i in args if isinstance(i.element, Identity)] + number_identity_matrices = len(identity_matrices) + # If the contraction involves a non-identity matrix and multiple identity matrices: + if number_identity_matrices != len(args) - 1 or number_identity_matrices == 0: + continue + # Get the non-identity element: + non_identity = [i for i in args if not isinstance(i.element, Identity)][0] + # Check that all identity matrices have at least one free index + # (otherwise they would be contractions to some other elements) + if any(None not in i.indices for i in identity_matrices): + continue + # Mark the identity matrices for removal: + for i in identity_matrices: + i.element = None + removed.extend(range(free_map[i], free_map[i] + len([j for j in i.indices if j is None]))) + last_removed = removed.pop(-1) + update_pairs[last_removed, ind] = non_identity.indices[:] + # Remove the indices from the non-identity matrix, as the contraction + # no longer exists: + non_identity.indices = [None if i == ind else i for i in non_identity.indices] + + removed.sort() + + shifts = list(accumulate([1 if i in removed else 0 for i in range(get_rank(expr))])) + for (last_removed, ind), non_identity_indices in update_pairs.items(): + pos = [free_map[non_identity] + i for i, e in enumerate(non_identity_indices) if e == ind] + assert len(pos) == 1 + for j in pos: + permutation_map[j] = last_removed + + editor.args_with_ind = [i for i in editor.args_with_ind if i.element is not None] + ret_expr = editor.to_array_contraction() + permutation = [] + counter = 0 + counter2 = 0 + for j in range(get_rank(expr)): + if j in removed: + continue + if counter2 in permutation_map: + target = permutation_map[counter2] + permutation.append(target - shifts[target]) + counter2 += 1 + else: + while counter in permutation_map.values(): + counter += 1 + permutation.append(counter) + counter += 1 + counter2 += 1 + ret_expr2 = _permute_dims(ret_expr, _af_invert(permutation)) + return ret_expr2, removed + + +def _combine_removed(dim: int, removed1: List[int], removed2: List[int]) -> List[int]: + # Concatenate two axis removal operations as performed by + # _remove_trivial_dims, + removed1 = sorted(removed1) + removed2 = sorted(removed2) + i = 0 + j = 0 + removed = [] + while True: + if j >= len(removed2): + while i < len(removed1): + removed.append(removed1[i]) + i += 1 + break + elif i < len(removed1) and removed1[i] <= i + removed2[j]: + removed.append(removed1[i]) + i += 1 + else: + removed.append(i + removed2[j]) + j += 1 + return removed + + +def _array_contraction_to_diagonal_multiple_identity(expr: ArrayContraction): + editor = _EditArrayContraction(expr) + editor.track_permutation_start() + removed: List[int] = [] + diag_index_counter: int = 0 + for i in range(editor.number_of_contraction_indices): + identities = [] + args = [] + for j, arg in enumerate(editor.args_with_ind): + if i not in arg.indices: + continue + if isinstance(arg.element, Identity): + identities.append(arg) + else: + args.append(arg) + if len(identities) == 0: + continue + if len(args) + len(identities) < 3: + continue + new_diag_ind = -1 - diag_index_counter + diag_index_counter += 1 + # Variable "flag" to control whether to skip this contraction set: + flag: bool = True + for i1, id1 in enumerate(identities): + if None not in id1.indices: + flag = True + break + free_pos = list(range(*editor.get_absolute_free_range(id1)))[0] + editor._track_permutation[-1].append(free_pos) # type: ignore + id1.element = None + flag = False + break + if flag: + continue + for arg in identities[:i1] + identities[i1+1:]: + arg.element = None + removed.extend(range(*editor.get_absolute_free_range(arg))) + for arg in args: + arg.indices = [new_diag_ind if j == i else j for j in arg.indices] + for j, e in enumerate(editor.args_with_ind): + if e.element is None: + editor._track_permutation[j] = None # type: ignore + editor._track_permutation = [i for i in editor._track_permutation if i is not None] # type: ignore + # Renumber permutation array form in order to deal with deleted positions: + remap = {e: i for i, e in enumerate(sorted({k for j in editor._track_permutation for k in j}))} + editor._track_permutation = [[remap[j] for j in i] for i in editor._track_permutation] + editor.args_with_ind = [i for i in editor.args_with_ind if i.element is not None] + new_expr = editor.to_array_contraction() + return new_expr, removed diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/from_indexed_to_array.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/from_indexed_to_array.py new file mode 100644 index 0000000000000000000000000000000000000000..c219a205c4305bd7070e5117978146224521c58c --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/from_indexed_to_array.py @@ -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, () diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/from_matrix_to_array.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/from_matrix_to_array.py new file mode 100644 index 0000000000000000000000000000000000000000..8f66961727f6338318d65876a7768802773e4f2d --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/from_matrix_to_array.py @@ -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 diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__init__.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/__init__.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d5e773e82b8fa8d1cb6337b3cc5017598a72d85 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/__init__.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_array_expressions.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_array_expressions.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f5a97f6352dea2a13ff1992832486edf351f69c6 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_array_expressions.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_arrayexpr_derivatives.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_arrayexpr_derivatives.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..97a55fd5f684071013307b7bed94d1e17d83e3ab Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_arrayexpr_derivatives.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_as_explicit.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_as_explicit.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8506657e4f122158a2b54ddda8ed2ec8a5419383 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_as_explicit.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_convert_array_to_indexed.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_convert_array_to_indexed.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7eb43df10ab40709dcdea1cb3120cf512036fc00 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_convert_array_to_indexed.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_convert_array_to_matrix.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_convert_array_to_matrix.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6cd3f6edc714a141879b8dd15668cf1779b277aa Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_convert_array_to_matrix.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_convert_indexed_to_array.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_convert_indexed_to_array.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e45ee855cb4890300b3ddcc0e8ef9ffb7ff07429 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_convert_indexed_to_array.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_convert_matrix_to_array.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_convert_matrix_to_array.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d471eea5fc6bf1a5f794a39e20a124f2b07b8e7f Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_convert_matrix_to_array.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_deprecated_conv_modules.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_deprecated_conv_modules.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..06cc10084eee0b4584d5908444b9619757acaa27 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/__pycache__/test_deprecated_conv_modules.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_array_expressions.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_array_expressions.py new file mode 100644 index 0000000000000000000000000000000000000000..63fb79ab7ced7bff5ecb55b1764f43e29f98609d --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_array_expressions.py @@ -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) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_arrayexpr_derivatives.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_arrayexpr_derivatives.py new file mode 100644 index 0000000000000000000000000000000000000000..bc0fcf63f2607b23feb38758e4f0994de4f0384b --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_arrayexpr_derivatives.py @@ -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)) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_as_explicit.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_as_explicit.py new file mode 100644 index 0000000000000000000000000000000000000000..30cc61b1ee651ca032e165cd67926fa33c71354f --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_as_explicit.py @@ -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() diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_convert_array_to_indexed.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_convert_array_to_indexed.py new file mode 100644 index 0000000000000000000000000000000000000000..a6b713fbec94ab7808c5a8a778b3313402d9d0c7 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_convert_array_to_indexed.py @@ -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] diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_convert_array_to_matrix.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_convert_array_to_matrix.py new file mode 100644 index 0000000000000000000000000000000000000000..26839d5e7cec0554948c6b726482f9d8ca250b1c --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_convert_array_to_matrix.py @@ -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]) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_convert_indexed_to_array.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_convert_indexed_to_array.py new file mode 100644 index 0000000000000000000000000000000000000000..258062eadeca041ae3c864dabeefd5165f1cef11 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_convert_indexed_to_array.py @@ -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])) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_convert_matrix_to_array.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_convert_matrix_to_array.py new file mode 100644 index 0000000000000000000000000000000000000000..142585882588df6aa0e4648d9d8881ea755f42a0 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_convert_matrix_to_array.py @@ -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)) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_deprecated_conv_modules.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_deprecated_conv_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..b41b6105410a308e7774fce760b235497d0303bb --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/tests/test_deprecated_conv_modules.py @@ -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) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/utils.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e55c0e6ed47cdc9ff1c24cc92f006998aeb86822 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/utils.py @@ -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 diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/mutable_ndim_array.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/mutable_ndim_array.py new file mode 100644 index 0000000000000000000000000000000000000000..e1eaaf7241bc3b4a48234178d18da3aa5736e189 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/mutable_ndim_array.py @@ -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() diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/ndim_array.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/ndim_array.py new file mode 100644 index 0000000000000000000000000000000000000000..b49faf0fb4c45f6de5bd2a702701d2b453508129 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/ndim_array.py @@ -0,0 +1,600 @@ +from sympy.core.basic import Basic +from sympy.core.containers import (Dict, Tuple) +from sympy.core.expr import Expr +from sympy.core.kind import Kind, NumberKind, UndefinedKind +from sympy.core.numbers import Integer +from sympy.core.singleton import S +from sympy.core.sympify import sympify +from sympy.external.gmpy import SYMPY_INTS +from sympy.printing.defaults import Printable + +import itertools +from collections.abc import Iterable + + +class ArrayKind(Kind): + """ + Kind for N-dimensional array in SymPy. + + This kind represents the multidimensional array that algebraic + operations are defined. Basic class for this kind is ``NDimArray``, + but any expression representing the array can have this. + + Parameters + ========== + + element_kind : Kind + Kind of the element. Default is :obj:NumberKind ``, + which means that the array contains only numbers. + + Examples + ======== + + Any instance of array class has ``ArrayKind``. + + >>> from sympy import NDimArray + >>> NDimArray([1,2,3]).kind + ArrayKind(NumberKind) + + Although expressions representing an array may be not instance of + array class, it will have ``ArrayKind`` as well. + + >>> from sympy import Integral + >>> from sympy.tensor.array import NDimArray + >>> from sympy.abc import x + >>> intA = Integral(NDimArray([1,2,3]), x) + >>> isinstance(intA, NDimArray) + False + >>> intA.kind + ArrayKind(NumberKind) + + Use ``isinstance()`` to check for ``ArrayKind` without specifying + the element kind. Use ``is`` with specifying the element kind. + + >>> from sympy.tensor.array import ArrayKind + >>> from sympy.core import NumberKind + >>> boolA = NDimArray([True, False]) + >>> isinstance(boolA.kind, ArrayKind) + True + >>> boolA.kind is ArrayKind(NumberKind) + False + + See Also + ======== + + shape : Function to return the shape of objects with ``MatrixKind``. + + """ + def __new__(cls, element_kind=NumberKind): + obj = super().__new__(cls, element_kind) + obj.element_kind = element_kind + return obj + + def __repr__(self): + return "ArrayKind(%s)" % self.element_kind + + @classmethod + def _union(cls, kinds) -> 'ArrayKind': + elem_kinds = {e.kind for e in kinds} + if len(elem_kinds) == 1: + elemkind, = elem_kinds + else: + elemkind = UndefinedKind + return ArrayKind(elemkind) + + +class NDimArray(Printable): + """N-dimensional array. + + Examples + ======== + + Create an N-dim array of zeros: + + >>> from sympy import MutableDenseNDimArray + >>> a = MutableDenseNDimArray.zeros(2, 3, 4) + >>> a + [[[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]], [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]] + + Create an N-dim array from a list; + + >>> a = MutableDenseNDimArray([[2, 3], [4, 5]]) + >>> a + [[2, 3], [4, 5]] + + >>> b = MutableDenseNDimArray([[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]]) + >>> b + [[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]] + + Create an N-dim array from a flat list with dimension shape: + + >>> a = MutableDenseNDimArray([1, 2, 3, 4, 5, 6], (2, 3)) + >>> a + [[1, 2, 3], [4, 5, 6]] + + Create an N-dim array from a matrix: + + >>> from sympy import Matrix + >>> a = Matrix([[1,2],[3,4]]) + >>> a + Matrix([ + [1, 2], + [3, 4]]) + >>> b = MutableDenseNDimArray(a) + >>> b + [[1, 2], [3, 4]] + + Arithmetic operations on N-dim arrays + + >>> a = MutableDenseNDimArray([1, 1, 1, 1], (2, 2)) + >>> b = MutableDenseNDimArray([4, 4, 4, 4], (2, 2)) + >>> c = a + b + >>> c + [[5, 5], [5, 5]] + >>> a - b + [[-3, -3], [-3, -3]] + + """ + + _diff_wrt = True + is_scalar = False + + def __new__(cls, iterable, shape=None, **kwargs): + from sympy.tensor.array import ImmutableDenseNDimArray + return ImmutableDenseNDimArray(iterable, shape, **kwargs) + + def __getitem__(self, index): + raise NotImplementedError("A subclass of NDimArray should implement __getitem__") + + def _parse_index(self, index): + if isinstance(index, (SYMPY_INTS, Integer)): + if index >= self._loop_size: + raise ValueError("Only a tuple index is accepted") + return index + + if self._loop_size == 0: + raise ValueError("Index not valid with an empty array") + + if len(index) != self._rank: + raise ValueError('Wrong number of array axes') + + real_index = 0 + # check if input index can exist in current indexing + for i in range(self._rank): + if (index[i] >= self.shape[i]) or (index[i] < -self.shape[i]): + raise ValueError('Index ' + str(index) + ' out of border') + if index[i] < 0: + real_index += 1 + real_index = real_index*self.shape[i] + index[i] + + return real_index + + def _get_tuple_index(self, integer_index): + index = [] + for sh in reversed(self.shape): + index.append(integer_index % sh) + integer_index //= sh + index.reverse() + return tuple(index) + + def _check_symbolic_index(self, index): + # Check if any index is symbolic: + tuple_index = (index if isinstance(index, tuple) else (index,)) + if any((isinstance(i, Expr) and (not i.is_number)) for i in tuple_index): + for i, nth_dim in zip(tuple_index, self.shape): + if ((i < 0) == True) or ((i >= nth_dim) == True): + raise ValueError("index out of range") + from sympy.tensor import Indexed + return Indexed(self, *tuple_index) + return None + + def _setter_iterable_check(self, value): + from sympy.matrices.matrixbase import MatrixBase + if isinstance(value, (Iterable, MatrixBase, NDimArray)): + raise NotImplementedError + + @classmethod + def _scan_iterable_shape(cls, iterable): + def f(pointer): + if not isinstance(pointer, Iterable): + return [pointer], () + + if len(pointer) == 0: + return [], (0,) + + result = [] + elems, shapes = zip(*[f(i) for i in pointer]) + if len(set(shapes)) != 1: + raise ValueError("could not determine shape unambiguously") + for i in elems: + result.extend(i) + return result, (len(shapes),)+shapes[0] + + return f(iterable) + + @classmethod + def _handle_ndarray_creation_inputs(cls, iterable=None, shape=None, **kwargs): + from sympy.matrices.matrixbase import MatrixBase + from sympy.tensor.array import SparseNDimArray + + if shape is None: + if iterable is None: + shape = () + iterable = () + # Construction of a sparse array from a sparse array + elif isinstance(iterable, SparseNDimArray): + return iterable._shape, iterable._sparse_array + + # Construct N-dim array from another N-dim array: + elif isinstance(iterable, NDimArray): + shape = iterable.shape + + # Construct N-dim array from an iterable (numpy arrays included): + elif isinstance(iterable, Iterable): + iterable, shape = cls._scan_iterable_shape(iterable) + + # Construct N-dim array from a Matrix: + elif isinstance(iterable, MatrixBase): + shape = iterable.shape + + else: + shape = () + iterable = (iterable,) + + if isinstance(iterable, (Dict, dict)) and shape is not None: + new_dict = iterable.copy() + for k in new_dict: + if isinstance(k, (tuple, Tuple)): + new_key = 0 + for i, idx in enumerate(k): + new_key = new_key * shape[i] + idx + iterable[new_key] = iterable[k] + del iterable[k] + + if isinstance(shape, (SYMPY_INTS, Integer)): + shape = (shape,) + + if not all(isinstance(dim, (SYMPY_INTS, Integer)) for dim in shape): + raise TypeError("Shape should contain integers only.") + + return tuple(shape), iterable + + def __len__(self): + """Overload common function len(). Returns number of elements in array. + + Examples + ======== + + >>> from sympy import MutableDenseNDimArray + >>> a = MutableDenseNDimArray.zeros(3, 3) + >>> a + [[0, 0, 0], [0, 0, 0], [0, 0, 0]] + >>> len(a) + 9 + + """ + return self._loop_size + + @property + def shape(self): + """ + Returns array shape (dimension). + + Examples + ======== + + >>> from sympy import MutableDenseNDimArray + >>> a = MutableDenseNDimArray.zeros(3, 3) + >>> a.shape + (3, 3) + + """ + return self._shape + + def rank(self): + """ + Returns rank of array. + + Examples + ======== + + >>> from sympy import MutableDenseNDimArray + >>> a = MutableDenseNDimArray.zeros(3,4,5,6,3) + >>> a.rank() + 5 + + """ + return self._rank + + def diff(self, *args, **kwargs): + """ + Calculate the derivative of each element in the array. + + Examples + ======== + + >>> from sympy import ImmutableDenseNDimArray + >>> from sympy.abc import x, y + >>> M = ImmutableDenseNDimArray([[x, y], [1, x*y]]) + >>> M.diff(x) + [[1, 0], [0, y]] + + """ + from sympy.tensor.array.array_derivatives import ArrayDerivative + kwargs.setdefault('evaluate', True) + return ArrayDerivative(self.as_immutable(), *args, **kwargs) + + def _eval_derivative(self, base): + # Types are (base: scalar, self: array) + return self.applyfunc(lambda x: base.diff(x)) + + def _eval_derivative_n_times(self, s, n): + return Basic._eval_derivative_n_times(self, s, n) + + def applyfunc(self, f): + """Apply a function to each element of the N-dim array. + + Examples + ======== + + >>> from sympy import ImmutableDenseNDimArray + >>> m = ImmutableDenseNDimArray([i*2+j for i in range(2) for j in range(2)], (2, 2)) + >>> m + [[0, 1], [2, 3]] + >>> m.applyfunc(lambda i: 2*i) + [[0, 2], [4, 6]] + """ + from sympy.tensor.array import SparseNDimArray + from sympy.tensor.array.arrayop import Flatten + + if isinstance(self, SparseNDimArray) and f(S.Zero) == 0: + return type(self)({k: f(v) for k, v in self._sparse_array.items() if f(v) != 0}, self.shape) + + return type(self)(map(f, Flatten(self)), self.shape) + + def _sympystr(self, printer): + def f(sh, shape_left, i, j): + if len(shape_left) == 1: + return "["+", ".join([printer._print(self[self._get_tuple_index(e)]) for e in range(i, j)])+"]" + + sh //= shape_left[0] + return "[" + ", ".join([f(sh, shape_left[1:], i+e*sh, i+(e+1)*sh) for e in range(shape_left[0])]) + "]" # + "\n"*len(shape_left) + + if self.rank() == 0: + return printer._print(self[()]) + + return f(self._loop_size, self.shape, 0, self._loop_size) + + def tolist(self): + """ + Converting MutableDenseNDimArray to one-dim list + + Examples + ======== + + >>> from sympy import MutableDenseNDimArray + >>> a = MutableDenseNDimArray([1, 2, 3, 4], (2, 2)) + >>> a + [[1, 2], [3, 4]] + >>> b = a.tolist() + >>> b + [[1, 2], [3, 4]] + """ + + def f(sh, shape_left, i, j): + if len(shape_left) == 1: + return [self[self._get_tuple_index(e)] for e in range(i, j)] + result = [] + sh //= shape_left[0] + for e in range(shape_left[0]): + result.append(f(sh, shape_left[1:], i+e*sh, i+(e+1)*sh)) + return result + + return f(self._loop_size, self.shape, 0, self._loop_size) + + def __add__(self, other): + from sympy.tensor.array.arrayop import Flatten + + if not isinstance(other, NDimArray): + return NotImplemented + + if self.shape != other.shape: + raise ValueError("array shape mismatch") + result_list = [i+j for i,j in zip(Flatten(self), Flatten(other))] + + return type(self)(result_list, self.shape) + + def __sub__(self, other): + from sympy.tensor.array.arrayop import Flatten + + if not isinstance(other, NDimArray): + return NotImplemented + + if self.shape != other.shape: + raise ValueError("array shape mismatch") + result_list = [i-j for i,j in zip(Flatten(self), Flatten(other))] + + return type(self)(result_list, self.shape) + + def __mul__(self, other): + from sympy.matrices.matrixbase import MatrixBase + from sympy.tensor.array import SparseNDimArray + from sympy.tensor.array.arrayop import Flatten + + if isinstance(other, (Iterable, NDimArray, MatrixBase)): + raise ValueError("scalar expected, use tensorproduct(...) for tensorial product") + + other = sympify(other) + if isinstance(self, SparseNDimArray): + if other.is_zero: + return type(self)({}, self.shape) + return type(self)({k: other*v for (k, v) in self._sparse_array.items()}, self.shape) + + result_list = [i*other for i in Flatten(self)] + return type(self)(result_list, self.shape) + + def __rmul__(self, other): + from sympy.matrices.matrixbase import MatrixBase + from sympy.tensor.array import SparseNDimArray + from sympy.tensor.array.arrayop import Flatten + + if isinstance(other, (Iterable, NDimArray, MatrixBase)): + raise ValueError("scalar expected, use tensorproduct(...) for tensorial product") + + other = sympify(other) + if isinstance(self, SparseNDimArray): + if other.is_zero: + return type(self)({}, self.shape) + return type(self)({k: other*v for (k, v) in self._sparse_array.items()}, self.shape) + + result_list = [other*i for i in Flatten(self)] + return type(self)(result_list, self.shape) + + def __truediv__(self, other): + from sympy.matrices.matrixbase import MatrixBase + from sympy.tensor.array import SparseNDimArray + from sympy.tensor.array.arrayop import Flatten + + if isinstance(other, (Iterable, NDimArray, MatrixBase)): + raise ValueError("scalar expected") + + other = sympify(other) + if isinstance(self, SparseNDimArray) and other != S.Zero: + return type(self)({k: v/other for (k, v) in self._sparse_array.items()}, self.shape) + + result_list = [i/other for i in Flatten(self)] + return type(self)(result_list, self.shape) + + def __rtruediv__(self, other): + raise NotImplementedError('unsupported operation on NDimArray') + + def __neg__(self): + from sympy.tensor.array import SparseNDimArray + from sympy.tensor.array.arrayop import Flatten + + if isinstance(self, SparseNDimArray): + return type(self)({k: -v for (k, v) in self._sparse_array.items()}, self.shape) + + result_list = [-i for i in Flatten(self)] + return type(self)(result_list, self.shape) + + def __iter__(self): + def iterator(): + if self._shape: + for i in range(self._shape[0]): + yield self[i] + else: + yield self[()] + + return iterator() + + def __eq__(self, other): + """ + NDimArray instances can be compared to each other. + Instances equal if they have same shape and data. + + Examples + ======== + + >>> from sympy import MutableDenseNDimArray + >>> a = MutableDenseNDimArray.zeros(2, 3) + >>> b = MutableDenseNDimArray.zeros(2, 3) + >>> a == b + True + >>> c = a.reshape(3, 2) + >>> c == b + False + >>> a[0,0] = 1 + >>> b[0,0] = 2 + >>> a == b + False + """ + from sympy.tensor.array import SparseNDimArray + if not isinstance(other, NDimArray): + return False + + if not self.shape == other.shape: + return False + + if isinstance(self, SparseNDimArray) and isinstance(other, SparseNDimArray): + return dict(self._sparse_array) == dict(other._sparse_array) + + return list(self) == list(other) + + def __ne__(self, other): + return not self == other + + def _eval_transpose(self): + if self.rank() != 2: + raise ValueError("array rank not 2") + from .arrayop import permutedims + return permutedims(self, (1, 0)) + + def transpose(self): + return self._eval_transpose() + + def _eval_conjugate(self): + from sympy.tensor.array.arrayop import Flatten + + return self.func([i.conjugate() for i in Flatten(self)], self.shape) + + def conjugate(self): + return self._eval_conjugate() + + def _eval_adjoint(self): + return self.transpose().conjugate() + + def adjoint(self): + return self._eval_adjoint() + + def _slice_expand(self, s, dim): + if not isinstance(s, slice): + return (s,) + start, stop, step = s.indices(dim) + return [start + i*step for i in range((stop-start)//step)] + + def _get_slice_data_for_array_access(self, index): + sl_factors = [self._slice_expand(i, dim) for (i, dim) in zip(index, self.shape)] + eindices = itertools.product(*sl_factors) + return sl_factors, eindices + + def _get_slice_data_for_array_assignment(self, index, value): + if not isinstance(value, NDimArray): + value = type(self)(value) + sl_factors, eindices = self._get_slice_data_for_array_access(index) + slice_offsets = [min(i) if isinstance(i, list) else None for i in sl_factors] + # TODO: add checks for dimensions for `value`? + return value, eindices, slice_offsets + + @classmethod + def _check_special_bounds(cls, flat_list, shape): + if shape == () and len(flat_list) != 1: + raise ValueError("arrays without shape need one scalar value") + if shape == (0,) and len(flat_list) > 0: + raise ValueError("if array shape is (0,) there cannot be elements") + + def _check_index_for_getitem(self, index): + if isinstance(index, (SYMPY_INTS, Integer, slice)): + index = (index,) + + if len(index) < self.rank(): + index = tuple(index) + \ + tuple(slice(None) for i in range(len(index), self.rank())) + + if len(index) > self.rank(): + raise ValueError('Dimension of index greater than rank of array') + + return index + + +class ImmutableNDimArray(NDimArray, Basic): + _op_priority = 11.0 + + def __hash__(self): + return Basic.__hash__(self) + + def as_immutable(self): + return self + + def as_mutable(self): + raise NotImplementedError("abstract method") diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/sparse_ndim_array.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/sparse_ndim_array.py new file mode 100644 index 0000000000000000000000000000000000000000..f11aa95be8ec9d10a9104d48fb28f406fe43845e --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/sparse_ndim_array.py @@ -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} diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/__init__.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/__pycache__/test_array_comprehension.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/__pycache__/test_array_comprehension.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..90abb4a1c913dd0965e5e01b137facc06d65bae4 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/__pycache__/test_array_comprehension.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/__pycache__/test_ndim_array_conversions.cpython-311.pyc b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/__pycache__/test_ndim_array_conversions.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9f70f770f21c9cc0fff246484f610e97934558f0 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/__pycache__/test_ndim_array_conversions.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_array_comprehension.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_array_comprehension.py new file mode 100644 index 0000000000000000000000000000000000000000..510e068f287fa04419712e5e9a16a314e522a62d --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_array_comprehension.py @@ -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()) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_array_derivatives.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_array_derivatives.py new file mode 100644 index 0000000000000000000000000000000000000000..7f6c777c55a9170704f309bf74387d140bf2ec32 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_array_derivatives.py @@ -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) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_arrayop.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_arrayop.py new file mode 100644 index 0000000000000000000000000000000000000000..de56e81e0064f1e303a7a58e41932d15f2d0b41e --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_arrayop.py @@ -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])) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_immutable_ndim_array.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_immutable_ndim_array.py new file mode 100644 index 0000000000000000000000000000000000000000..c6bed4b605c424284b4752592b03b13a9178aac8 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_immutable_ndim_array.py @@ -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 diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_mutable_ndim_array.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_mutable_ndim_array.py new file mode 100644 index 0000000000000000000000000000000000000000..9a232f399bbc0639d326217975fb0a12e645a984 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_mutable_ndim_array.py @@ -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]]) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_ndim_array.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_ndim_array.py new file mode 100644 index 0000000000000000000000000000000000000000..7ff9b032631c01272c00478e4cdf0dcbc6997990 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_ndim_array.py @@ -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 diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_ndim_array_conversions.py b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_ndim_array_conversions.py new file mode 100644 index 0000000000000000000000000000000000000000..f43260ccc636ac461ba0c06dbfcf3fe3a8d5338d --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/array/tests/test_ndim_array_conversions.py @@ -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 diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/functions.py b/.venv/lib/python3.11/site-packages/sympy/tensor/functions.py new file mode 100644 index 0000000000000000000000000000000000000000..f14599d69152db1713f21c9dd785683901c5eeb9 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/functions.py @@ -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 diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/index_methods.py b/.venv/lib/python3.11/site-packages/sympy/tensor/index_methods.py new file mode 100644 index 0000000000000000000000000000000000000000..12f707b60b4ad0bcadc35a222d9abe0cc5e033fc --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/index_methods.py @@ -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)) diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/indexed.py b/.venv/lib/python3.11/site-packages/sympy/tensor/indexed.py new file mode 100644 index 0000000000000000000000000000000000000000..feddad21e52bbab2e1243beafdb11f30b2eded4d --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/indexed.py @@ -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) + + + 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]) + + + 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 diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/tensor.py b/.venv/lib/python3.11/site-packages/sympy/tensor/tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..67e287db8625f09ebeb813b89a1424b1787bc5fb --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/tensor.py @@ -0,0 +1,4979 @@ +""" +This module defines tensors with abstract index notation. + +The abstract index notation has been first formalized by Penrose. + +Tensor indices are formal objects, with a tensor type; there is no +notion of index range, it is only possible to assign the dimension, +used to trace the Kronecker delta; the dimension can be a Symbol. + +The Einstein summation convention is used. +The covariant indices are indicated with a minus sign in front of the index. + +For instance the tensor ``t = p(a)*A(b,c)*q(-c)`` has the index ``c`` +contracted. + +A tensor expression ``t`` can be called; called with its +indices in sorted order it is equal to itself: +in the above example ``t(a, b) == t``; +one can call ``t`` with different indices; ``t(c, d) == p(c)*A(d,a)*q(-a)``. + +The contracted indices are dummy indices, internally they have no name, +the indices being represented by a graph-like structure. + +Tensors are put in canonical form using ``canon_bp``, which uses +the Butler-Portugal algorithm for canonicalization using the monoterm +symmetries of the tensors. + +If there is a (anti)symmetric metric, the indices can be raised and +lowered when the tensor is put in canonical form. +""" + +from __future__ import annotations +from typing import Any +from functools import reduce +from math import prod + +from abc import abstractmethod, ABC +from collections import defaultdict +import operator +import itertools +from sympy.core.numbers import (Integer, Rational) +from sympy.combinatorics import Permutation +from sympy.combinatorics.tensor_can import get_symmetric_group_sgs, \ + bsgs_direct_product, canonicalize, riemann_bsgs +from sympy.core import Basic, Expr, sympify, Add, Mul, S +from sympy.core.cache import clear_cache +from sympy.core.containers import Tuple, Dict +from sympy.core.sorting import default_sort_key +from sympy.core.symbol import Symbol, symbols +from sympy.core.sympify import CantSympify, _sympify +from sympy.core.operations import AssocOp +from sympy.external.gmpy import SYMPY_INTS +from sympy.matrices import eye +from sympy.utilities.exceptions import (sympy_deprecation_warning, + SymPyDeprecationWarning, + ignore_warnings) +from sympy.utilities.decorator import memoize_property, deprecated +from sympy.utilities.iterables import sift + + +def deprecate_data(): + sympy_deprecation_warning( + """ + The data attribute of TensorIndexType is deprecated. Use The + replace_with_arrays() method instead. + """, + deprecated_since_version="1.4", + active_deprecations_target="deprecated-tensorindextype-attrs", + stacklevel=4, + ) + +def deprecate_fun_eval(): + sympy_deprecation_warning( + """ + The Tensor.fun_eval() method is deprecated. Use + Tensor.substitute_indices() instead. + """, + deprecated_since_version="1.5", + active_deprecations_target="deprecated-tensor-fun-eval", + stacklevel=4, + ) + + +def deprecate_call(): + sympy_deprecation_warning( + """ + Calling a tensor like Tensor(*indices) is deprecated. Use + Tensor.substitute_indices() instead. + """, + deprecated_since_version="1.5", + active_deprecations_target="deprecated-tensor-fun-eval", + stacklevel=4, + ) + + +class _IndexStructure(CantSympify): + """ + This class handles the indices (free and dummy ones). It contains the + algorithms to manage the dummy indices replacements and contractions of + free indices under multiplications of tensor expressions, as well as stuff + related to canonicalization sorting, getting the permutation of the + expression and so on. It also includes tools to get the ``TensorIndex`` + objects corresponding to the given index structure. + """ + + def __init__(self, free, dum, index_types, indices, canon_bp=False): + self.free = free + self.dum = dum + self.index_types = index_types + self.indices = indices + self._ext_rank = len(self.free) + 2*len(self.dum) + self.dum.sort(key=lambda x: x[0]) + + @staticmethod + def from_indices(*indices): + """ + Create a new ``_IndexStructure`` object from a list of ``indices``. + + Explanation + =========== + + ``indices`` ``TensorIndex`` objects, the indices. Contractions are + detected upon construction. + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_indices, _IndexStructure + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> m0, m1, m2, m3 = tensor_indices('m0,m1,m2,m3', Lorentz) + >>> _IndexStructure.from_indices(m0, m1, -m1, m3) + _IndexStructure([(m0, 0), (m3, 3)], [(1, 2)], [Lorentz, Lorentz, Lorentz, Lorentz]) + """ + + free, dum = _IndexStructure._free_dum_from_indices(*indices) + index_types = [i.tensor_index_type for i in indices] + indices = _IndexStructure._replace_dummy_names(indices, free, dum) + return _IndexStructure(free, dum, index_types, indices) + + @staticmethod + def from_components_free_dum(components, free, dum): + index_types = [] + for component in components: + index_types.extend(component.index_types) + indices = _IndexStructure.generate_indices_from_free_dum_index_types(free, dum, index_types) + return _IndexStructure(free, dum, index_types, indices) + + @staticmethod + def _free_dum_from_indices(*indices): + """ + Convert ``indices`` into ``free``, ``dum`` for single component tensor. + + Explanation + =========== + + ``free`` list of tuples ``(index, pos, 0)``, + where ``pos`` is the position of index in + the list of indices formed by the component tensors + + ``dum`` list of tuples ``(pos_contr, pos_cov, 0, 0)`` + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_indices, \ + _IndexStructure + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> m0, m1, m2, m3 = tensor_indices('m0,m1,m2,m3', Lorentz) + >>> _IndexStructure._free_dum_from_indices(m0, m1, -m1, m3) + ([(m0, 0), (m3, 3)], [(1, 2)]) + """ + n = len(indices) + if n == 1: + return [(indices[0], 0)], [] + + # find the positions of the free indices and of the dummy indices + free = [True]*len(indices) + index_dict = {} + dum = [] + for i, index in enumerate(indices): + name = index.name + typ = index.tensor_index_type + contr = index.is_up + if (name, typ) in index_dict: + # found a pair of dummy indices + is_contr, pos = index_dict[(name, typ)] + # check consistency and update free + if is_contr: + if contr: + raise ValueError('two equal contravariant indices in slots %d and %d' %(pos, i)) + else: + free[pos] = False + free[i] = False + else: + if contr: + free[pos] = False + free[i] = False + else: + raise ValueError('two equal covariant indices in slots %d and %d' %(pos, i)) + if contr: + dum.append((i, pos)) + else: + dum.append((pos, i)) + else: + index_dict[(name, typ)] = index.is_up, i + + free = [(index, i) for i, index in enumerate(indices) if free[i]] + free.sort() + return free, dum + + def get_indices(self): + """ + Get a list of indices, creating new tensor indices to complete dummy indices. + """ + return self.indices[:] + + @staticmethod + def generate_indices_from_free_dum_index_types(free, dum, index_types): + indices = [None]*(len(free)+2*len(dum)) + for idx, pos in free: + indices[pos] = idx + + generate_dummy_name = _IndexStructure._get_generator_for_dummy_indices(free) + for pos1, pos2 in dum: + typ1 = index_types[pos1] + indname = generate_dummy_name(typ1) + indices[pos1] = TensorIndex(indname, typ1, True) + indices[pos2] = TensorIndex(indname, typ1, False) + + return _IndexStructure._replace_dummy_names(indices, free, dum) + + @staticmethod + def _get_generator_for_dummy_indices(free): + cdt = defaultdict(int) + # if the free indices have names with dummy_name, start with an + # index higher than those for the dummy indices + # to avoid name collisions + for indx, ipos in free: + if indx.name.split('_')[0] == indx.tensor_index_type.dummy_name: + cdt[indx.tensor_index_type] = max(cdt[indx.tensor_index_type], int(indx.name.split('_')[1]) + 1) + + def dummy_name_gen(tensor_index_type): + nd = str(cdt[tensor_index_type]) + cdt[tensor_index_type] += 1 + return tensor_index_type.dummy_name + '_' + nd + + return dummy_name_gen + + @staticmethod + def _replace_dummy_names(indices, free, dum): + dum.sort(key=lambda x: x[0]) + new_indices = list(indices) + assert len(indices) == len(free) + 2*len(dum) + generate_dummy_name = _IndexStructure._get_generator_for_dummy_indices(free) + for ipos1, ipos2 in dum: + typ1 = new_indices[ipos1].tensor_index_type + indname = generate_dummy_name(typ1) + new_indices[ipos1] = TensorIndex(indname, typ1, True) + new_indices[ipos2] = TensorIndex(indname, typ1, False) + return new_indices + + def get_free_indices(self) -> list[TensorIndex]: + """ + Get a list of free indices. + """ + # get sorted indices according to their position: + free = sorted(self.free, key=lambda x: x[1]) + return [i[0] for i in free] + + def __str__(self): + return "_IndexStructure({}, {}, {})".format(self.free, self.dum, self.index_types) + + def __repr__(self): + return self.__str__() + + def _get_sorted_free_indices_for_canon(self): + sorted_free = self.free[:] + sorted_free.sort(key=lambda x: x[0]) + return sorted_free + + def _get_sorted_dum_indices_for_canon(self): + return sorted(self.dum, key=lambda x: x[0]) + + def _get_lexicographically_sorted_index_types(self): + permutation = self.indices_canon_args()[0] + index_types = [None]*self._ext_rank + for i, it in enumerate(self.index_types): + index_types[permutation(i)] = it + return index_types + + def _get_lexicographically_sorted_indices(self): + permutation = self.indices_canon_args()[0] + indices = [None]*self._ext_rank + for i, it in enumerate(self.indices): + indices[permutation(i)] = it + return indices + + def perm2tensor(self, g, is_canon_bp=False): + """ + Returns a ``_IndexStructure`` instance corresponding to the permutation ``g``. + + Explanation + =========== + + ``g`` permutation corresponding to the tensor in the representation + used in canonicalization + + ``is_canon_bp`` if True, then ``g`` is the permutation + corresponding to the canonical form of the tensor + """ + sorted_free = [i[0] for i in self._get_sorted_free_indices_for_canon()] + lex_index_types = self._get_lexicographically_sorted_index_types() + lex_indices = self._get_lexicographically_sorted_indices() + nfree = len(sorted_free) + rank = self._ext_rank + dum = [[None]*2 for i in range((rank - nfree)//2)] + free = [] + + index_types = [None]*rank + indices = [None]*rank + for i in range(rank): + gi = g[i] + index_types[i] = lex_index_types[gi] + indices[i] = lex_indices[gi] + if gi < nfree: + ind = sorted_free[gi] + assert index_types[i] == sorted_free[gi].tensor_index_type + free.append((ind, i)) + else: + j = gi - nfree + idum, cov = divmod(j, 2) + if cov: + dum[idum][1] = i + else: + dum[idum][0] = i + dum = [tuple(x) for x in dum] + + return _IndexStructure(free, dum, index_types, indices) + + def indices_canon_args(self): + """ + Returns ``(g, dummies, msym, v)``, the entries of ``canonicalize`` + + See ``canonicalize`` in ``tensor_can.py`` in combinatorics module. + """ + # to be called after sorted_components + from sympy.combinatorics.permutations import _af_new + n = self._ext_rank + g = [None]*n + [n, n+1] + + # Converts the symmetry of the metric into msym from .canonicalize() + # method in the combinatorics module + def metric_symmetry_to_msym(metric): + if metric is None: + return None + sym = metric.symmetry + if sym == TensorSymmetry.fully_symmetric(2): + return 0 + if sym == TensorSymmetry.fully_symmetric(-2): + return 1 + return None + + # ordered indices: first the free indices, ordered by types + # then the dummy indices, ordered by types and contravariant before + # covariant + # g[position in tensor] = position in ordered indices + for i, (indx, ipos) in enumerate(self._get_sorted_free_indices_for_canon()): + g[ipos] = i + pos = len(self.free) + j = len(self.free) + dummies = [] + prev = None + a = [] + msym = [] + for ipos1, ipos2 in self._get_sorted_dum_indices_for_canon(): + g[ipos1] = j + g[ipos2] = j + 1 + j += 2 + typ = self.index_types[ipos1] + if typ != prev: + if a: + dummies.append(a) + a = [pos, pos + 1] + prev = typ + msym.append(metric_symmetry_to_msym(typ.metric)) + else: + a.extend([pos, pos + 1]) + pos += 2 + if a: + dummies.append(a) + + return _af_new(g), dummies, msym + + +def components_canon_args(components): + numtyp = [] + prev = None + for t in components: + if t == prev: + numtyp[-1][1] += 1 + else: + prev = t + numtyp.append([prev, 1]) + v = [] + for h, n in numtyp: + if h.comm in (0, 1): + comm = h.comm + else: + comm = TensorManager.get_comm(h.comm, h.comm) + v.append((h.symmetry.base, h.symmetry.generators, n, comm)) + return v + + +class _TensorDataLazyEvaluator(CantSympify): + """ + EXPERIMENTAL: do not rely on this class, it may change without deprecation + warnings in future versions of SymPy. + + Explanation + =========== + + This object contains the logic to associate components data to a tensor + expression. Components data are set via the ``.data`` property of tensor + expressions, is stored inside this class as a mapping between the tensor + expression and the ``ndarray``. + + Computations are executed lazily: whereas the tensor expressions can have + contractions, tensor products, and additions, components data are not + computed until they are accessed by reading the ``.data`` property + associated to the tensor expression. + """ + _substitutions_dict: dict[Any, Any] = {} + _substitutions_dict_tensmul: dict[Any, Any] = {} + + def __getitem__(self, key): + dat = self._get(key) + if dat is None: + return None + + from .array import NDimArray + if not isinstance(dat, NDimArray): + return dat + + if dat.rank() == 0: + return dat[()] + elif dat.rank() == 1 and len(dat) == 1: + return dat[0] + return dat + + def _get(self, key): + """ + Retrieve ``data`` associated with ``key``. + + Explanation + =========== + + This algorithm looks into ``self._substitutions_dict`` for all + ``TensorHead`` in the ``TensExpr`` (or just ``TensorHead`` if key is a + TensorHead instance). It reconstructs the components data that the + tensor expression should have by performing on components data the + operations that correspond to the abstract tensor operations applied. + + Metric tensor is handled in a different manner: it is pre-computed in + ``self._substitutions_dict_tensmul``. + """ + if key in self._substitutions_dict: + return self._substitutions_dict[key] + + if isinstance(key, TensorHead): + return None + + if isinstance(key, Tensor): + # special case to handle metrics. Metric tensors cannot be + # constructed through contraction by the metric, their + # components show if they are a matrix or its inverse. + signature = tuple([i.is_up for i in key.get_indices()]) + srch = (key.component,) + signature + if srch in self._substitutions_dict_tensmul: + return self._substitutions_dict_tensmul[srch] + array_list = [self.data_from_tensor(key)] + return self.data_contract_dum(array_list, key.dum, key.ext_rank) + + if isinstance(key, TensMul): + tensmul_args = key.args + if len(tensmul_args) == 1 and len(tensmul_args[0].components) == 1: + # special case to handle metrics. Metric tensors cannot be + # constructed through contraction by the metric, their + # components show if they are a matrix or its inverse. + signature = tuple([i.is_up for i in tensmul_args[0].get_indices()]) + srch = (tensmul_args[0].components[0],) + signature + if srch in self._substitutions_dict_tensmul: + return self._substitutions_dict_tensmul[srch] + #data_list = [self.data_from_tensor(i) for i in tensmul_args if isinstance(i, TensExpr)] + data_list = [self.data_from_tensor(i) if isinstance(i, Tensor) else i.data for i in tensmul_args if isinstance(i, TensExpr)] + coeff = prod([i for i in tensmul_args if not isinstance(i, TensExpr)]) + if all(i is None for i in data_list): + return None + if any(i is None for i in data_list): + raise ValueError("Mixing tensors with associated components "\ + "data with tensors without components data") + data_result = self.data_contract_dum(data_list, key.dum, key.ext_rank) + return coeff*data_result + + if isinstance(key, TensAdd): + data_list = [] + free_args_list = [] + for arg in key.args: + if isinstance(arg, TensExpr): + data_list.append(arg.data) + free_args_list.append([x[0] for x in arg.free]) + else: + data_list.append(arg) + free_args_list.append([]) + if all(i is None for i in data_list): + return None + if any(i is None for i in data_list): + raise ValueError("Mixing tensors with associated components "\ + "data with tensors without components data") + + sum_list = [] + from .array import permutedims + for data, free_args in zip(data_list, free_args_list): + if len(free_args) < 2: + sum_list.append(data) + else: + free_args_pos = {y: x for x, y in enumerate(free_args)} + axes = [free_args_pos[arg] for arg in key.free_args] + sum_list.append(permutedims(data, axes)) + return reduce(lambda x, y: x+y, sum_list) + + return None + + @staticmethod + def data_contract_dum(ndarray_list, dum, ext_rank): + from .array import tensorproduct, tensorcontraction, MutableDenseNDimArray + arrays = list(map(MutableDenseNDimArray, ndarray_list)) + prodarr = tensorproduct(*arrays) + return tensorcontraction(prodarr, *dum) + + def data_tensorhead_from_tensmul(self, data, tensmul, tensorhead): + """ + This method is used when assigning components data to a ``TensMul`` + object, it converts components data to a fully contravariant ndarray, + which is then stored according to the ``TensorHead`` key. + """ + if data is None: + return None + + return self._correct_signature_from_indices( + data, + tensmul.get_indices(), + tensmul.free, + tensmul.dum, + True) + + def data_from_tensor(self, tensor): + """ + This method corrects the components data to the right signature + (covariant/contravariant) using the metric associated with each + ``TensorIndexType``. + """ + tensorhead = tensor.component + + if tensorhead.data is None: + return None + + return self._correct_signature_from_indices( + tensorhead.data, + tensor.get_indices(), + tensor.free, + tensor.dum) + + def _assign_data_to_tensor_expr(self, key, data): + if isinstance(key, TensAdd): + raise ValueError('cannot assign data to TensAdd') + # here it is assumed that `key` is a `TensMul` instance. + if len(key.components) != 1: + raise ValueError('cannot assign data to TensMul with multiple components') + tensorhead = key.components[0] + newdata = self.data_tensorhead_from_tensmul(data, key, tensorhead) + return tensorhead, newdata + + def _check_permutations_on_data(self, tens, data): + from .array import permutedims + from .array.arrayop import Flatten + + if isinstance(tens, TensorHead): + rank = tens.rank + generators = tens.symmetry.generators + elif isinstance(tens, Tensor): + rank = tens.rank + generators = tens.components[0].symmetry.generators + elif isinstance(tens, TensorIndexType): + rank = tens.metric.rank + generators = tens.metric.symmetry.generators + + # Every generator is a permutation, check that by permuting the array + # by that permutation, the array will be the same, except for a + # possible sign change if the permutation admits it. + for gener in generators: + sign_change = +1 if (gener(rank) == rank) else -1 + data_swapped = data + last_data = data + permute_axes = list(map(gener, range(rank))) + # the order of a permutation is the number of times to get the + # identity by applying that permutation. + for i in range(gener.order()-1): + data_swapped = permutedims(data_swapped, permute_axes) + # if any value in the difference array is non-zero, raise an error: + if any(Flatten(last_data - sign_change*data_swapped)): + raise ValueError("Component data symmetry structure error") + last_data = data_swapped + + def __setitem__(self, key, value): + """ + Set the components data of a tensor object/expression. + + Explanation + =========== + + Components data are transformed to the all-contravariant form and stored + with the corresponding ``TensorHead`` object. If a ``TensorHead`` object + cannot be uniquely identified, it will raise an error. + """ + data = _TensorDataLazyEvaluator.parse_data(value) + self._check_permutations_on_data(key, data) + + # TensorHead and TensorIndexType can be assigned data directly, while + # TensMul must first convert data to a fully contravariant form, and + # assign it to its corresponding TensorHead single component. + if not isinstance(key, (TensorHead, TensorIndexType)): + key, data = self._assign_data_to_tensor_expr(key, data) + + if isinstance(key, TensorHead): + for dim, indextype in zip(data.shape, key.index_types): + if indextype.data is None: + raise ValueError("index type {} has no components data"\ + " associated (needed to raise/lower index)".format(indextype)) + if not indextype.dim.is_number: + continue + if dim != indextype.dim: + raise ValueError("wrong dimension of ndarray") + self._substitutions_dict[key] = data + + def __delitem__(self, key): + del self._substitutions_dict[key] + + def __contains__(self, key): + return key in self._substitutions_dict + + def add_metric_data(self, metric, data): + """ + Assign data to the ``metric`` tensor. The metric tensor behaves in an + anomalous way when raising and lowering indices. + + Explanation + =========== + + A fully covariant metric is the inverse transpose of the fully + contravariant metric (it is meant matrix inverse). If the metric is + symmetric, the transpose is not necessary and mixed + covariant/contravariant metrics are Kronecker deltas. + """ + # hard assignment, data should not be added to `TensorHead` for metric: + # the problem with `TensorHead` is that the metric is anomalous, i.e. + # raising and lowering the index means considering the metric or its + # inverse, this is not the case for other tensors. + self._substitutions_dict_tensmul[metric, True, True] = data + inverse_transpose = self.inverse_transpose_matrix(data) + # in symmetric spaces, the transpose is the same as the original matrix, + # the full covariant metric tensor is the inverse transpose, so this + # code will be able to handle non-symmetric metrics. + self._substitutions_dict_tensmul[metric, False, False] = inverse_transpose + # now mixed cases, these are identical to the unit matrix if the metric + # is symmetric. + m = data.tomatrix() + invt = inverse_transpose.tomatrix() + self._substitutions_dict_tensmul[metric, True, False] = m * invt + self._substitutions_dict_tensmul[metric, False, True] = invt * m + + @staticmethod + def _flip_index_by_metric(data, metric, pos): + from .array import tensorproduct, tensorcontraction + + mdim = metric.rank() + ddim = data.rank() + + if pos == 0: + data = tensorcontraction( + tensorproduct( + metric, + data + ), + (1, mdim+pos) + ) + else: + data = tensorcontraction( + tensorproduct( + data, + metric + ), + (pos, ddim) + ) + return data + + @staticmethod + def inverse_matrix(ndarray): + m = ndarray.tomatrix().inv() + return _TensorDataLazyEvaluator.parse_data(m) + + @staticmethod + def inverse_transpose_matrix(ndarray): + m = ndarray.tomatrix().inv().T + return _TensorDataLazyEvaluator.parse_data(m) + + @staticmethod + def _correct_signature_from_indices(data, indices, free, dum, inverse=False): + """ + Utility function to correct the values inside the components data + ndarray according to whether indices are covariant or contravariant. + + It uses the metric matrix to lower values of covariant indices. + """ + # change the ndarray values according covariantness/contravariantness of the indices + # use the metric + for i, indx in enumerate(indices): + if not indx.is_up and not inverse: + data = _TensorDataLazyEvaluator._flip_index_by_metric(data, indx.tensor_index_type.data, i) + elif not indx.is_up and inverse: + data = _TensorDataLazyEvaluator._flip_index_by_metric( + data, + _TensorDataLazyEvaluator.inverse_matrix(indx.tensor_index_type.data), + i + ) + return data + + @staticmethod + def _sort_data_axes(old, new): + from .array import permutedims + + new_data = old.data.copy() + + old_free = [i[0] for i in old.free] + new_free = [i[0] for i in new.free] + + for i in range(len(new_free)): + for j in range(i, len(old_free)): + if old_free[j] == new_free[i]: + old_free[i], old_free[j] = old_free[j], old_free[i] + new_data = permutedims(new_data, (i, j)) + break + return new_data + + @staticmethod + def add_rearrange_tensmul_parts(new_tensmul, old_tensmul): + def sorted_compo(): + return _TensorDataLazyEvaluator._sort_data_axes(old_tensmul, new_tensmul) + + _TensorDataLazyEvaluator._substitutions_dict[new_tensmul] = sorted_compo() + + @staticmethod + def parse_data(data): + """ + Transform ``data`` to array. The parameter ``data`` may + contain data in various formats, e.g. nested lists, SymPy ``Matrix``, + and so on. + + Examples + ======== + + >>> from sympy.tensor.tensor import _TensorDataLazyEvaluator + >>> _TensorDataLazyEvaluator.parse_data([1, 3, -6, 12]) + [1, 3, -6, 12] + + >>> _TensorDataLazyEvaluator.parse_data([[1, 2], [4, 7]]) + [[1, 2], [4, 7]] + """ + from .array import MutableDenseNDimArray + + if not isinstance(data, MutableDenseNDimArray): + if len(data) == 2 and hasattr(data[0], '__call__'): + data = MutableDenseNDimArray(data[0], data[1]) + else: + data = MutableDenseNDimArray(data) + return data + +_tensor_data_substitution_dict = _TensorDataLazyEvaluator() + + +class _TensorManager: + """ + Class to manage tensor properties. + + Notes + ===== + + Tensors belong to tensor commutation groups; each group has a label + ``comm``; there are predefined labels: + + ``0`` tensors commuting with any other tensor + + ``1`` tensors anticommuting among themselves + + ``2`` tensors not commuting, apart with those with ``comm=0`` + + Other groups can be defined using ``set_comm``; tensors in those + groups commute with those with ``comm=0``; by default they + do not commute with any other group. + """ + def __init__(self): + self._comm_init() + + def _comm_init(self): + self._comm = [{} for i in range(3)] + for i in range(3): + self._comm[0][i] = 0 + self._comm[i][0] = 0 + self._comm[1][1] = 1 + self._comm[2][1] = None + self._comm[1][2] = None + self._comm_symbols2i = {0:0, 1:1, 2:2} + self._comm_i2symbol = {0:0, 1:1, 2:2} + + @property + def comm(self): + return self._comm + + def comm_symbols2i(self, i): + """ + Get the commutation group number corresponding to ``i``. + + ``i`` can be a symbol or a number or a string. + + If ``i`` is not already defined its commutation group number + is set. + """ + if i not in self._comm_symbols2i: + n = len(self._comm) + self._comm.append({}) + self._comm[n][0] = 0 + self._comm[0][n] = 0 + self._comm_symbols2i[i] = n + self._comm_i2symbol[n] = i + return n + return self._comm_symbols2i[i] + + def comm_i2symbol(self, i): + """ + Returns the symbol corresponding to the commutation group number. + """ + return self._comm_i2symbol[i] + + def set_comm(self, i, j, c): + """ + Set the commutation parameter ``c`` for commutation groups ``i, j``. + + Parameters + ========== + + i, j : symbols representing commutation groups + + c : group commutation number + + Notes + ===== + + ``i, j`` can be symbols, strings or numbers, + apart from ``0, 1`` and ``2`` which are reserved respectively + for commuting, anticommuting tensors and tensors not commuting + with any other group apart with the commuting tensors. + For the remaining cases, use this method to set the commutation rules; + by default ``c=None``. + + The group commutation number ``c`` is assigned in correspondence + to the group commutation symbols; it can be + + 0 commuting + + 1 anticommuting + + None no commutation property + + Examples + ======== + + ``G`` and ``GH`` do not commute with themselves and commute with + each other; A is commuting. + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_indices, TensorHead, TensorManager, TensorSymmetry + >>> Lorentz = TensorIndexType('Lorentz') + >>> i0,i1,i2,i3,i4 = tensor_indices('i0:5', Lorentz) + >>> A = TensorHead('A', [Lorentz]) + >>> G = TensorHead('G', [Lorentz], TensorSymmetry.no_symmetry(1), 'Gcomm') + >>> GH = TensorHead('GH', [Lorentz], TensorSymmetry.no_symmetry(1), 'GHcomm') + >>> TensorManager.set_comm('Gcomm', 'GHcomm', 0) + >>> (GH(i1)*G(i0)).canon_bp() + G(i0)*GH(i1) + >>> (G(i1)*G(i0)).canon_bp() + G(i1)*G(i0) + >>> (G(i1)*A(i0)).canon_bp() + A(i0)*G(i1) + """ + if c not in (0, 1, None): + raise ValueError('`c` can assume only the values 0, 1 or None') + + i = sympify(i) + j = sympify(j) + + if i not in self._comm_symbols2i: + n = len(self._comm) + self._comm.append({}) + self._comm[n][0] = 0 + self._comm[0][n] = 0 + self._comm_symbols2i[i] = n + self._comm_i2symbol[n] = i + if j not in self._comm_symbols2i: + n = len(self._comm) + self._comm.append({}) + self._comm[0][n] = 0 + self._comm[n][0] = 0 + self._comm_symbols2i[j] = n + self._comm_i2symbol[n] = j + ni = self._comm_symbols2i[i] + nj = self._comm_symbols2i[j] + self._comm[ni][nj] = c + self._comm[nj][ni] = c + + """ + Cached sympy functions (e.g. expand) may have cached the results of + expressions involving tensors, but those results may not be valid after + changing the commutation properties. To stay on the safe side, we clear + the cache of all functions. + """ + clear_cache() + + def set_comms(self, *args): + """ + Set the commutation group numbers ``c`` for symbols ``i, j``. + + Parameters + ========== + + args : sequence of ``(i, j, c)`` + """ + for i, j, c in args: + self.set_comm(i, j, c) + + def get_comm(self, i, j): + """ + Return the commutation parameter for commutation group numbers ``i, j`` + + see ``_TensorManager.set_comm`` + """ + return self._comm[i].get(j, 0 if i == 0 or j == 0 else None) + + def clear(self): + """ + Clear the TensorManager. + """ + self._comm_init() + + +TensorManager = _TensorManager() + + +class TensorIndexType(Basic): + """ + A TensorIndexType is characterized by its name and its metric. + + Parameters + ========== + + name : name of the tensor type + dummy_name : name of the head of dummy indices + dim : dimension, it can be a symbol or an integer or ``None`` + eps_dim : dimension of the epsilon tensor + metric_symmetry : integer that denotes metric symmetry or ``None`` for no metric + metric_name : string with the name of the metric tensor + + Attributes + ========== + + ``metric`` : the metric tensor + ``delta`` : ``Kronecker delta`` + ``epsilon`` : the ``Levi-Civita epsilon`` tensor + ``data`` : (deprecated) a property to add ``ndarray`` values, to work in a specified basis. + + Notes + ===== + + The possible values of the ``metric_symmetry`` parameter are: + + ``1`` : metric tensor is fully symmetric + ``0`` : metric tensor possesses no index symmetry + ``-1`` : metric tensor is fully antisymmetric + ``None``: there is no metric tensor (metric equals to ``None``) + + The metric is assumed to be symmetric by default. It can also be set + to a custom tensor by the ``.set_metric()`` method. + + If there is a metric the metric is used to raise and lower indices. + + In the case of non-symmetric metric, the following raising and + lowering conventions will be adopted: + + ``psi(a) = g(a, b)*psi(-b); chi(-a) = chi(b)*g(-b, -a)`` + + From these it is easy to find: + + ``g(-a, b) = delta(-a, b)`` + + where ``delta(-a, b) = delta(b, -a)`` is the ``Kronecker delta`` + (see ``TensorIndex`` for the conventions on indices). + For antisymmetric metrics there is also the following equality: + + ``g(a, -b) = -delta(a, -b)`` + + If there is no metric it is not possible to raise or lower indices; + e.g. the index of the defining representation of ``SU(N)`` + is 'covariant' and the conjugate representation is + 'contravariant'; for ``N > 2`` they are linearly independent. + + ``eps_dim`` is by default equal to ``dim``, if the latter is an integer; + else it can be assigned (for use in naive dimensional regularization); + if ``eps_dim`` is not an integer ``epsilon`` is ``None``. + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> Lorentz.metric + metric(Lorentz,Lorentz) + """ + + def __new__(cls, name, dummy_name=None, dim=None, eps_dim=None, + metric_symmetry=1, metric_name='metric', **kwargs): + if 'dummy_fmt' in kwargs: + dummy_fmt = kwargs['dummy_fmt'] + sympy_deprecation_warning( + f""" + The dummy_fmt keyword to TensorIndexType is deprecated. Use + dummy_name={dummy_fmt} instead. + """, + deprecated_since_version="1.5", + active_deprecations_target="deprecated-tensorindextype-dummy-fmt", + ) + dummy_name = dummy_fmt + + if isinstance(name, str): + name = Symbol(name) + + if dummy_name is None: + dummy_name = str(name)[0] + if isinstance(dummy_name, str): + dummy_name = Symbol(dummy_name) + + if dim is None: + dim = Symbol("dim_" + dummy_name.name) + else: + dim = sympify(dim) + + if eps_dim is None: + eps_dim = dim + else: + eps_dim = sympify(eps_dim) + + metric_symmetry = sympify(metric_symmetry) + + if isinstance(metric_name, str): + metric_name = Symbol(metric_name) + + if 'metric' in kwargs: + SymPyDeprecationWarning( + """ + The 'metric' keyword argument to TensorIndexType is + deprecated. Use the 'metric_symmetry' keyword argument or the + TensorIndexType.set_metric() method instead. + """, + deprecated_since_version="1.5", + active_deprecations_target="deprecated-tensorindextype-metric", + ) + metric = kwargs.get('metric') + if metric is not None: + if metric in (True, False, 0, 1): + metric_name = 'metric' + #metric_antisym = metric + else: + metric_name = metric.name + #metric_antisym = metric.antisym + + if metric: + metric_symmetry = -1 + else: + metric_symmetry = 1 + + obj = Basic.__new__(cls, name, dummy_name, dim, eps_dim, + metric_symmetry, metric_name) + + obj._autogenerated = [] + return obj + + @property + def name(self): + return self.args[0].name + + @property + def dummy_name(self): + return self.args[1].name + + @property + def dim(self): + return self.args[2] + + @property + def eps_dim(self): + return self.args[3] + + @memoize_property + def metric(self): + metric_symmetry = self.args[4] + metric_name = self.args[5] + if metric_symmetry is None: + return None + + if metric_symmetry == 0: + symmetry = TensorSymmetry.no_symmetry(2) + elif metric_symmetry == 1: + symmetry = TensorSymmetry.fully_symmetric(2) + elif metric_symmetry == -1: + symmetry = TensorSymmetry.fully_symmetric(-2) + + return TensorHead(metric_name, [self]*2, symmetry) + + @memoize_property + def delta(self): + return TensorHead('KD', [self]*2, TensorSymmetry.fully_symmetric(2)) + + @memoize_property + def epsilon(self): + if not isinstance(self.eps_dim, (SYMPY_INTS, Integer)): + return None + symmetry = TensorSymmetry.fully_symmetric(-self.eps_dim) + return TensorHead('Eps', [self]*self.eps_dim, symmetry) + + def set_metric(self, tensor): + self._metric = tensor + + def __lt__(self, other): + return self.name < other.name + + def __str__(self): + return self.name + + __repr__ = __str__ + + # Everything below this line is deprecated + + @property + def data(self): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + return _tensor_data_substitution_dict[self] + + @data.setter + def data(self, data): + deprecate_data() + # This assignment is a bit controversial, should metric components be assigned + # to the metric only or also to the TensorIndexType object? The advantage here + # is the ability to assign a 1D array and transform it to a 2D diagonal array. + from .array import MutableDenseNDimArray + + data = _TensorDataLazyEvaluator.parse_data(data) + if data.rank() > 2: + raise ValueError("data have to be of rank 1 (diagonal metric) or 2.") + if data.rank() == 1: + if self.dim.is_number: + nda_dim = data.shape[0] + if nda_dim != self.dim: + raise ValueError("Dimension mismatch") + + dim = data.shape[0] + newndarray = MutableDenseNDimArray.zeros(dim, dim) + for i, val in enumerate(data): + newndarray[i, i] = val + data = newndarray + dim1, dim2 = data.shape + if dim1 != dim2: + raise ValueError("Non-square matrix tensor.") + if self.dim.is_number: + if self.dim != dim1: + raise ValueError("Dimension mismatch") + _tensor_data_substitution_dict[self] = data + _tensor_data_substitution_dict.add_metric_data(self.metric, data) + with ignore_warnings(SymPyDeprecationWarning): + delta = self.get_kronecker_delta() + i1 = TensorIndex('i1', self) + i2 = TensorIndex('i2', self) + with ignore_warnings(SymPyDeprecationWarning): + delta(i1, -i2).data = _TensorDataLazyEvaluator.parse_data(eye(dim1)) + + @data.deleter + def data(self): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + if self in _tensor_data_substitution_dict: + del _tensor_data_substitution_dict[self] + if self.metric in _tensor_data_substitution_dict: + del _tensor_data_substitution_dict[self.metric] + + @deprecated( + """ + The TensorIndexType.get_kronecker_delta() method is deprecated. Use + the TensorIndexType.delta attribute instead. + """, + deprecated_since_version="1.5", + active_deprecations_target="deprecated-tensorindextype-methods", + ) + def get_kronecker_delta(self): + sym2 = TensorSymmetry(get_symmetric_group_sgs(2)) + delta = TensorHead('KD', [self]*2, sym2) + return delta + + @deprecated( + """ + The TensorIndexType.get_epsilon() method is deprecated. Use + the TensorIndexType.epsilon attribute instead. + """, + deprecated_since_version="1.5", + active_deprecations_target="deprecated-tensorindextype-methods", + ) + def get_epsilon(self): + if not isinstance(self._eps_dim, (SYMPY_INTS, Integer)): + return None + sym = TensorSymmetry(get_symmetric_group_sgs(self._eps_dim, 1)) + epsilon = TensorHead('Eps', [self]*self._eps_dim, sym) + return epsilon + + def _components_data_full_destroy(self): + """ + EXPERIMENTAL: do not rely on this API method. + + This destroys components data associated to the ``TensorIndexType``, if + any, specifically: + + * metric tensor data + * Kronecker tensor data + """ + if self in _tensor_data_substitution_dict: + del _tensor_data_substitution_dict[self] + + def delete_tensmul_data(key): + if key in _tensor_data_substitution_dict._substitutions_dict_tensmul: + del _tensor_data_substitution_dict._substitutions_dict_tensmul[key] + + # delete metric data: + delete_tensmul_data((self.metric, True, True)) + delete_tensmul_data((self.metric, True, False)) + delete_tensmul_data((self.metric, False, True)) + delete_tensmul_data((self.metric, False, False)) + + # delete delta tensor data: + delta = self.get_kronecker_delta() + if delta in _tensor_data_substitution_dict: + del _tensor_data_substitution_dict[delta] + + +class TensorIndex(Basic): + """ + Represents a tensor index + + Parameters + ========== + + name : name of the index, or ``True`` if you want it to be automatically assigned + tensor_index_type : ``TensorIndexType`` of the index + is_up : flag for contravariant index (is_up=True by default) + + Attributes + ========== + + ``name`` + ``tensor_index_type`` + ``is_up`` + + Notes + ===== + + Tensor indices are contracted with the Einstein summation convention. + + An index can be in contravariant or in covariant form; in the latter + case it is represented prepending a ``-`` to the index name. Adding + ``-`` to a covariant (is_up=False) index makes it contravariant. + + Dummy indices have a name with head given by + ``tensor_inde_type.dummy_name`` with underscore and a number. + + Similar to ``symbols`` multiple contravariant indices can be created + at once using ``tensor_indices(s, typ)``, where ``s`` is a string + of names. + + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, TensorIndex, TensorHead, tensor_indices + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> mu = TensorIndex('mu', Lorentz, is_up=False) + >>> nu, rho = tensor_indices('nu, rho', Lorentz) + >>> A = TensorHead('A', [Lorentz, Lorentz]) + >>> A(mu, nu) + A(-mu, nu) + >>> A(-mu, -rho) + A(mu, -rho) + >>> A(mu, -mu) + A(-L_0, L_0) + """ + def __new__(cls, name, tensor_index_type, is_up=True): + if isinstance(name, str): + name_symbol = Symbol(name) + elif isinstance(name, Symbol): + name_symbol = name + elif name is True: + name = "_i{}".format(len(tensor_index_type._autogenerated)) + name_symbol = Symbol(name) + tensor_index_type._autogenerated.append(name_symbol) + else: + raise ValueError("invalid name") + + is_up = sympify(is_up) + return Basic.__new__(cls, name_symbol, tensor_index_type, is_up) + + @property + def name(self): + return self.args[0].name + + @property + def tensor_index_type(self): + return self.args[1] + + @property + def is_up(self): + return self.args[2] + + def _print(self): + s = self.name + if not self.is_up: + s = '-%s' % s + return s + + def __lt__(self, other): + return ((self.tensor_index_type, self.name) < + (other.tensor_index_type, other.name)) + + def __neg__(self): + t1 = TensorIndex(self.name, self.tensor_index_type, + (not self.is_up)) + return t1 + + +def tensor_indices(s, typ): + """ + Returns list of tensor indices given their names and their types. + + Parameters + ========== + + s : string of comma separated names of indices + + typ : ``TensorIndexType`` of the indices + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_indices + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> a, b, c, d = tensor_indices('a,b,c,d', Lorentz) + """ + if isinstance(s, str): + a = [x.name for x in symbols(s, seq=True)] + else: + raise ValueError('expecting a string') + + tilist = [TensorIndex(i, typ) for i in a] + if len(tilist) == 1: + return tilist[0] + return tilist + + +class TensorSymmetry(Basic): + """ + Monoterm symmetry of a tensor (i.e. any symmetric or anti-symmetric + index permutation). For the relevant terminology see ``tensor_can.py`` + section of the combinatorics module. + + Parameters + ========== + + bsgs : tuple ``(base, sgs)`` BSGS of the symmetry of the tensor + + Attributes + ========== + + ``base`` : base of the BSGS + ``generators`` : generators of the BSGS + ``rank`` : rank of the tensor + + Notes + ===== + + A tensor can have an arbitrary monoterm symmetry provided by its BSGS. + Multiterm symmetries, like the cyclic symmetry of the Riemann tensor + (i.e., Bianchi identity), are not covered. See combinatorics module for + information on how to generate BSGS for a general index permutation group. + Simple symmetries can be generated using built-in methods. + + See Also + ======== + + sympy.combinatorics.tensor_can.get_symmetric_group_sgs + + Examples + ======== + + Define a symmetric tensor of rank 2 + + >>> from sympy.tensor.tensor import TensorIndexType, TensorSymmetry, get_symmetric_group_sgs, TensorHead + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> sym = TensorSymmetry(get_symmetric_group_sgs(2)) + >>> T = TensorHead('T', [Lorentz]*2, sym) + + Note, that the same can also be done using built-in TensorSymmetry methods + + >>> sym2 = TensorSymmetry.fully_symmetric(2) + >>> sym == sym2 + True + """ + def __new__(cls, *args, **kw_args): + if len(args) == 1: + base, generators = args[0] + elif len(args) == 2: + base, generators = args + else: + raise TypeError("bsgs required, either two separate parameters or one tuple") + + if not isinstance(base, Tuple): + base = Tuple(*base) + if not isinstance(generators, Tuple): + generators = Tuple(*generators) + + return Basic.__new__(cls, base, generators, **kw_args) + + @property + def base(self): + return self.args[0] + + @property + def generators(self): + return self.args[1] + + @property + def rank(self): + return self.generators[0].size - 2 + + @classmethod + def fully_symmetric(cls, rank): + """ + Returns a fully symmetric (antisymmetric if ``rank``<0) + TensorSymmetry object for ``abs(rank)`` indices. + """ + if rank > 0: + bsgs = get_symmetric_group_sgs(rank, False) + elif rank < 0: + bsgs = get_symmetric_group_sgs(-rank, True) + elif rank == 0: + bsgs = ([], [Permutation(1)]) + return TensorSymmetry(bsgs) + + @classmethod + def direct_product(cls, *args): + """ + Returns a TensorSymmetry object that is being a direct product of + fully (anti-)symmetric index permutation groups. + + Notes + ===== + + Some examples for different values of ``(*args)``: + ``(1)`` vector, equivalent to ``TensorSymmetry.fully_symmetric(1)`` + ``(2)`` tensor with 2 symmetric indices, equivalent to ``.fully_symmetric(2)`` + ``(-2)`` tensor with 2 antisymmetric indices, equivalent to ``.fully_symmetric(-2)`` + ``(2, -2)`` tensor with the first 2 indices commuting and the last 2 anticommuting + ``(1, 1, 1)`` tensor with 3 indices without any symmetry + """ + base, sgs = [], [Permutation(1)] + for arg in args: + if arg > 0: + bsgs2 = get_symmetric_group_sgs(arg, False) + elif arg < 0: + bsgs2 = get_symmetric_group_sgs(-arg, True) + else: + continue + base, sgs = bsgs_direct_product(base, sgs, *bsgs2) + + return TensorSymmetry(base, sgs) + + @classmethod + def riemann(cls): + """ + Returns a monotorem symmetry of the Riemann tensor + """ + return TensorSymmetry(riemann_bsgs) + + @classmethod + def no_symmetry(cls, rank): + """ + TensorSymmetry object for ``rank`` indices with no symmetry + """ + return TensorSymmetry([], [Permutation(rank+1)]) + + +@deprecated( + """ + The tensorsymmetry() function is deprecated. Use the TensorSymmetry + constructor instead. + """, + deprecated_since_version="1.5", + active_deprecations_target="deprecated-tensorsymmetry", +) +def tensorsymmetry(*args): + """ + Returns a ``TensorSymmetry`` object. This method is deprecated, use + ``TensorSymmetry.direct_product()`` or ``.riemann()`` instead. + + Explanation + =========== + + One can represent a tensor with any monoterm slot symmetry group + using a BSGS. + + ``args`` can be a BSGS + ``args[0]`` base + ``args[1]`` sgs + + Usually tensors are in (direct products of) representations + of the symmetric group; + ``args`` can be a list of lists representing the shapes of Young tableaux + + Notes + ===== + + For instance: + ``[[1]]`` vector + ``[[1]*n]`` symmetric tensor of rank ``n`` + ``[[n]]`` antisymmetric tensor of rank ``n`` + ``[[2, 2]]`` monoterm slot symmetry of the Riemann tensor + ``[[1],[1]]`` vector*vector + ``[[2],[1],[1]`` (antisymmetric tensor)*vector*vector + + Notice that with the shape ``[2, 2]`` we associate only the monoterm + symmetries of the Riemann tensor; this is an abuse of notation, + since the shape ``[2, 2]`` corresponds usually to the irreducible + representation characterized by the monoterm symmetries and by the + cyclic symmetry. + """ + from sympy.combinatorics import Permutation + + def tableau2bsgs(a): + if len(a) == 1: + # antisymmetric vector + n = a[0] + bsgs = get_symmetric_group_sgs(n, 1) + else: + if all(x == 1 for x in a): + # symmetric vector + n = len(a) + bsgs = get_symmetric_group_sgs(n) + elif a == [2, 2]: + bsgs = riemann_bsgs + else: + raise NotImplementedError + return bsgs + + if not args: + return TensorSymmetry(Tuple(), Tuple(Permutation(1))) + + if len(args) == 2 and isinstance(args[1][0], Permutation): + return TensorSymmetry(args) + base, sgs = tableau2bsgs(args[0]) + for a in args[1:]: + basex, sgsx = tableau2bsgs(a) + base, sgs = bsgs_direct_product(base, sgs, basex, sgsx) + return TensorSymmetry(Tuple(base, sgs)) + +@deprecated( + "TensorType is deprecated. Use tensor_heads() instead.", + deprecated_since_version="1.5", + active_deprecations_target="deprecated-tensortype", +) +class TensorType(Basic): + """ + Class of tensor types. Deprecated, use tensor_heads() instead. + + Parameters + ========== + + index_types : list of ``TensorIndexType`` of the tensor indices + symmetry : ``TensorSymmetry`` of the tensor + + Attributes + ========== + + ``index_types`` + ``symmetry`` + ``types`` : list of ``TensorIndexType`` without repetitions + """ + is_commutative = False + + def __new__(cls, index_types, symmetry, **kw_args): + assert symmetry.rank == len(index_types) + obj = Basic.__new__(cls, Tuple(*index_types), symmetry, **kw_args) + return obj + + @property + def index_types(self): + return self.args[0] + + @property + def symmetry(self): + return self.args[1] + + @property + def types(self): + return sorted(set(self.index_types), key=lambda x: x.name) + + def __str__(self): + return 'TensorType(%s)' % ([str(x) for x in self.index_types]) + + def __call__(self, s, comm=0): + """ + Return a TensorHead object or a list of TensorHead objects. + + Parameters + ========== + + s : name or string of names. + + comm : Commutation group. + + see ``_TensorManager.set_comm`` + """ + if isinstance(s, str): + names = [x.name for x in symbols(s, seq=True)] + else: + raise ValueError('expecting a string') + if len(names) == 1: + return TensorHead(names[0], self.index_types, self.symmetry, comm) + else: + return [TensorHead(name, self.index_types, self.symmetry, comm) for name in names] + + +@deprecated( + """ + The tensorhead() function is deprecated. Use tensor_heads() instead. + """, + deprecated_since_version="1.5", + active_deprecations_target="deprecated-tensorhead", +) +def tensorhead(name, typ, sym=None, comm=0): + """ + Function generating tensorhead(s). This method is deprecated, + use TensorHead constructor or tensor_heads() instead. + + Parameters + ========== + + name : name or sequence of names (as in ``symbols``) + + typ : index types + + sym : same as ``*args`` in ``tensorsymmetry`` + + comm : commutation group number + see ``_TensorManager.set_comm`` + """ + if sym is None: + sym = [[1] for i in range(len(typ))] + with ignore_warnings(SymPyDeprecationWarning): + sym = tensorsymmetry(*sym) + return TensorHead(name, typ, sym, comm) + + +class TensorHead(Basic): + """ + Tensor head of the tensor. + + Parameters + ========== + + name : name of the tensor + index_types : list of TensorIndexType + symmetry : TensorSymmetry of the tensor + comm : commutation group number + + Attributes + ========== + + ``name`` + ``index_types`` + ``rank`` : total number of indices + ``symmetry`` + ``comm`` : commutation group + + Notes + ===== + + Similar to ``symbols`` multiple TensorHeads can be created using + ``tensorhead(s, typ, sym=None, comm=0)`` function, where ``s`` + is the string of names and ``sym`` is the monoterm tensor symmetry + (see ``tensorsymmetry``). + + A ``TensorHead`` belongs to a commutation group, defined by a + symbol on number ``comm`` (see ``_TensorManager.set_comm``); + tensors in a commutation group have the same commutation properties; + by default ``comm`` is ``0``, the group of the commuting tensors. + + Examples + ======== + + Define a fully antisymmetric tensor of rank 2: + + >>> from sympy.tensor.tensor import TensorIndexType, TensorHead, TensorSymmetry + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> asym2 = TensorSymmetry.fully_symmetric(-2) + >>> A = TensorHead('A', [Lorentz, Lorentz], asym2) + + Examples with ndarray values, the components data assigned to the + ``TensorHead`` object are assumed to be in a fully-contravariant + representation. In case it is necessary to assign components data which + represents the values of a non-fully covariant tensor, see the other + examples. + + >>> from sympy.tensor.tensor import tensor_indices + >>> from sympy import diag + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> i0, i1 = tensor_indices('i0:2', Lorentz) + + Specify a replacement dictionary to keep track of the arrays to use for + replacements in the tensorial expression. The ``TensorIndexType`` is + associated to the metric used for contractions (in fully covariant form): + + >>> repl = {Lorentz: diag(1, -1, -1, -1)} + + Let's see some examples of working with components with the electromagnetic + tensor: + + >>> from sympy import symbols + >>> Ex, Ey, Ez, Bx, By, Bz = symbols('E_x E_y E_z B_x B_y B_z') + >>> c = symbols('c', positive=True) + + Let's define `F`, an antisymmetric tensor: + + >>> F = TensorHead('F', [Lorentz, Lorentz], asym2) + + Let's update the dictionary to contain the matrix to use in the + replacements: + + >>> repl.update({F(-i0, -i1): [ + ... [0, Ex/c, Ey/c, Ez/c], + ... [-Ex/c, 0, -Bz, By], + ... [-Ey/c, Bz, 0, -Bx], + ... [-Ez/c, -By, Bx, 0]]}) + + Now it is possible to retrieve the contravariant form of the Electromagnetic + tensor: + + >>> F(i0, i1).replace_with_arrays(repl, [i0, i1]) + [[0, -E_x/c, -E_y/c, -E_z/c], [E_x/c, 0, -B_z, B_y], [E_y/c, B_z, 0, -B_x], [E_z/c, -B_y, B_x, 0]] + + and the mixed contravariant-covariant form: + + >>> F(i0, -i1).replace_with_arrays(repl, [i0, -i1]) + [[0, E_x/c, E_y/c, E_z/c], [E_x/c, 0, B_z, -B_y], [E_y/c, -B_z, 0, B_x], [E_z/c, B_y, -B_x, 0]] + + Energy-momentum of a particle may be represented as: + + >>> from sympy import symbols + >>> P = TensorHead('P', [Lorentz], TensorSymmetry.no_symmetry(1)) + >>> E, px, py, pz = symbols('E p_x p_y p_z', positive=True) + >>> repl.update({P(i0): [E, px, py, pz]}) + + The contravariant and covariant components are, respectively: + + >>> P(i0).replace_with_arrays(repl, [i0]) + [E, p_x, p_y, p_z] + >>> P(-i0).replace_with_arrays(repl, [-i0]) + [E, -p_x, -p_y, -p_z] + + The contraction of a 1-index tensor by itself: + + >>> expr = P(i0)*P(-i0) + >>> expr.replace_with_arrays(repl, []) + E**2 - p_x**2 - p_y**2 - p_z**2 + """ + is_commutative = False + + def __new__(cls, name, index_types, symmetry=None, comm=0): + if isinstance(name, str): + name_symbol = Symbol(name) + elif isinstance(name, Symbol): + name_symbol = name + else: + raise ValueError("invalid name") + + if symmetry is None: + symmetry = TensorSymmetry.no_symmetry(len(index_types)) + else: + assert symmetry.rank == len(index_types) + + obj = Basic.__new__(cls, name_symbol, Tuple(*index_types), symmetry, sympify(comm)) + return obj + + @property + def name(self): + return self.args[0].name + + @property + def index_types(self): + return list(self.args[1]) + + @property + def symmetry(self): + return self.args[2] + + @property + def comm(self): + return TensorManager.comm_symbols2i(self.args[3]) + + @property + def rank(self): + return len(self.index_types) + + def __lt__(self, other): + return (self.name, self.index_types) < (other.name, other.index_types) + + def commutes_with(self, other): + """ + Returns ``0`` if ``self`` and ``other`` commute, ``1`` if they anticommute. + + Returns ``None`` if ``self`` and ``other`` neither commute nor anticommute. + """ + r = TensorManager.get_comm(self.comm, other.comm) + return r + + def _print(self): + return '%s(%s)' %(self.name, ','.join([str(x) for x in self.index_types])) + + def __call__(self, *indices, **kw_args): + """ + Returns a tensor with indices. + + Explanation + =========== + + There is a special behavior in case of indices denoted by ``True``, + they are considered auto-matrix indices, their slots are automatically + filled, and confer to the tensor the behavior of a matrix or vector + upon multiplication with another tensor containing auto-matrix indices + of the same ``TensorIndexType``. This means indices get summed over the + same way as in matrix multiplication. For matrix behavior, define two + auto-matrix indices, for vector behavior define just one. + + Indices can also be strings, in which case the attribute + ``index_types`` is used to convert them to proper ``TensorIndex``. + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_indices, TensorSymmetry, TensorHead + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> a, b = tensor_indices('a,b', Lorentz) + >>> A = TensorHead('A', [Lorentz]*2, TensorSymmetry.no_symmetry(2)) + >>> t = A(a, -b) + >>> t + A(a, -b) + + """ + + updated_indices = [] + for idx, typ in zip(indices, self.index_types): + if isinstance(idx, str): + idx = idx.strip().replace(" ", "") + if idx.startswith('-'): + updated_indices.append(TensorIndex(idx[1:], typ, + is_up=False)) + else: + updated_indices.append(TensorIndex(idx, typ)) + else: + updated_indices.append(idx) + + updated_indices += indices[len(updated_indices):] + + tensor = Tensor(self, updated_indices, **kw_args) + return tensor.doit() + + # Everything below this line is deprecated + + def __pow__(self, other): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + if self.data is None: + raise ValueError("No power on abstract tensors.") + from .array import tensorproduct, tensorcontraction + metrics = [_.data for _ in self.index_types] + + marray = self.data + marraydim = marray.rank() + for metric in metrics: + marray = tensorproduct(marray, metric, marray) + marray = tensorcontraction(marray, (0, marraydim), (marraydim+1, marraydim+2)) + + return marray ** (other * S.Half) + + @property + def data(self): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + return _tensor_data_substitution_dict[self] + + @data.setter + def data(self, data): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + _tensor_data_substitution_dict[self] = data + + @data.deleter + def data(self): + deprecate_data() + if self in _tensor_data_substitution_dict: + del _tensor_data_substitution_dict[self] + + def __iter__(self): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + return self.data.__iter__() + + def _components_data_full_destroy(self): + """ + EXPERIMENTAL: do not rely on this API method. + + Destroy components data associated to the ``TensorHead`` object, this + checks for attached components data, and destroys components data too. + """ + # do not garbage collect Kronecker tensor (it should be done by + # ``TensorIndexType`` garbage collection) + deprecate_data() + if self.name == "KD": + return + + # the data attached to a tensor must be deleted only by the TensorHead + # destructor. If the TensorHead is deleted, it means that there are no + # more instances of that tensor anywhere. + if self in _tensor_data_substitution_dict: + del _tensor_data_substitution_dict[self] + + +def tensor_heads(s, index_types, symmetry=None, comm=0): + """ + Returns a sequence of TensorHeads from a string `s` + """ + if isinstance(s, str): + names = [x.name for x in symbols(s, seq=True)] + else: + raise ValueError('expecting a string') + + thlist = [TensorHead(name, index_types, symmetry, comm) for name in names] + if len(thlist) == 1: + return thlist[0] + return thlist + + +class TensExpr(Expr, ABC): + """ + Abstract base class for tensor expressions + + Notes + ===== + + A tensor expression is an expression formed by tensors; + currently the sums of tensors are distributed. + + A ``TensExpr`` can be a ``TensAdd`` or a ``TensMul``. + + ``TensMul`` objects are formed by products of component tensors, + and include a coefficient, which is a SymPy expression. + + + In the internal representation contracted indices are represented + by ``(ipos1, ipos2, icomp1, icomp2)``, where ``icomp1`` is the position + of the component tensor with contravariant index, ``ipos1`` is the + slot which the index occupies in that component tensor. + + Contracted indices are therefore nameless in the internal representation. + """ + + _op_priority = 12.0 + is_commutative = False + + def __neg__(self): + return self*S.NegativeOne + + def __abs__(self): + raise NotImplementedError + + def __add__(self, other): + return TensAdd(self, other).doit() + + def __radd__(self, other): + return TensAdd(other, self).doit() + + def __sub__(self, other): + return TensAdd(self, -other).doit() + + def __rsub__(self, other): + return TensAdd(other, -self).doit() + + def __mul__(self, other): + """ + Multiply two tensors using Einstein summation convention. + + Explanation + =========== + + If the two tensors have an index in common, one contravariant + and the other covariant, in their product the indices are summed + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_indices, tensor_heads + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> m0, m1, m2 = tensor_indices('m0,m1,m2', Lorentz) + >>> g = Lorentz.metric + >>> p, q = tensor_heads('p,q', [Lorentz]) + >>> t1 = p(m0) + >>> t2 = q(-m0) + >>> t1*t2 + p(L_0)*q(-L_0) + """ + return TensMul(self, other).doit() + + def __rmul__(self, other): + return TensMul(other, self).doit() + + def __truediv__(self, other): + other = _sympify(other) + if isinstance(other, TensExpr): + raise ValueError('cannot divide by a tensor') + return TensMul(self, S.One/other).doit() + + def __rtruediv__(self, other): + raise ValueError('cannot divide by a tensor') + + def __pow__(self, other): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + if self.data is None: + raise ValueError("No power without ndarray data.") + from .array import tensorproduct, tensorcontraction + free = self.free + marray = self.data + mdim = marray.rank() + for metric in free: + marray = tensorcontraction( + tensorproduct( + marray, + metric[0].tensor_index_type.data, + marray), + (0, mdim), (mdim+1, mdim+2) + ) + return marray ** (other * S.Half) + + def __rpow__(self, other): + raise NotImplementedError + + @property + @abstractmethod + def nocoeff(self): + raise NotImplementedError("abstract method") + + @property + @abstractmethod + def coeff(self): + raise NotImplementedError("abstract method") + + @abstractmethod + def get_indices(self): + raise NotImplementedError("abstract method") + + @abstractmethod + def get_free_indices(self) -> list[TensorIndex]: + raise NotImplementedError("abstract method") + + @abstractmethod + def _replace_indices(self, repl: dict[TensorIndex, TensorIndex]) -> TensExpr: + raise NotImplementedError("abstract method") + + def fun_eval(self, *index_tuples): + deprecate_fun_eval() + return self.substitute_indices(*index_tuples) + + def get_matrix(self): + """ + DEPRECATED: do not use. + + Returns ndarray components data as a matrix, if components data are + available and ndarray dimension does not exceed 2. + """ + from sympy.matrices.dense import Matrix + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + if 0 < self.rank <= 2: + rows = self.data.shape[0] + columns = self.data.shape[1] if self.rank == 2 else 1 + if self.rank == 2: + mat_list = [] * rows + for i in range(rows): + mat_list.append([]) + for j in range(columns): + mat_list[i].append(self[i, j]) + else: + mat_list = [None] * rows + for i in range(rows): + mat_list[i] = self[i] + return Matrix(mat_list) + else: + raise NotImplementedError( + "missing multidimensional reduction to matrix.") + + @staticmethod + def _get_indices_permutation(indices1, indices2): + return [indices1.index(i) for i in indices2] + + def expand(self, **hints): + return _expand(self, **hints).doit() + + def _expand(self, **kwargs): + return self + + def _get_free_indices_set(self): + indset = set() + for arg in self.args: + if isinstance(arg, TensExpr): + indset.update(arg._get_free_indices_set()) + return indset + + def _get_dummy_indices_set(self): + indset = set() + for arg in self.args: + if isinstance(arg, TensExpr): + indset.update(arg._get_dummy_indices_set()) + return indset + + def _get_indices_set(self): + indset = set() + for arg in self.args: + if isinstance(arg, TensExpr): + indset.update(arg._get_indices_set()) + return indset + + @property + def _iterate_dummy_indices(self): + dummy_set = self._get_dummy_indices_set() + + def recursor(expr, pos): + if isinstance(expr, TensorIndex): + if expr in dummy_set: + yield (expr, pos) + elif isinstance(expr, (Tuple, TensExpr)): + for p, arg in enumerate(expr.args): + yield from recursor(arg, pos+(p,)) + + return recursor(self, ()) + + @property + def _iterate_free_indices(self): + free_set = self._get_free_indices_set() + + def recursor(expr, pos): + if isinstance(expr, TensorIndex): + if expr in free_set: + yield (expr, pos) + elif isinstance(expr, (Tuple, TensExpr)): + for p, arg in enumerate(expr.args): + yield from recursor(arg, pos+(p,)) + + return recursor(self, ()) + + @property + def _iterate_indices(self): + def recursor(expr, pos): + if isinstance(expr, TensorIndex): + yield (expr, pos) + elif isinstance(expr, (Tuple, TensExpr)): + for p, arg in enumerate(expr.args): + yield from recursor(arg, pos+(p,)) + + return recursor(self, ()) + + @staticmethod + def _contract_and_permute_with_metric(metric, array, pos, dim): + # TODO: add possibility of metric after (spinors) + from .array import tensorcontraction, tensorproduct, permutedims + + array = tensorcontraction(tensorproduct(metric, array), (1, 2+pos)) + permu = list(range(dim)) + permu[0], permu[pos] = permu[pos], permu[0] + return permutedims(array, permu) + + @staticmethod + def _match_indices_with_other_tensor(array, free_ind1, free_ind2, replacement_dict): + from .array import permutedims + + index_types1 = [i.tensor_index_type for i in free_ind1] + + # Check if variance of indices needs to be fixed: + pos2up = [] + pos2down = [] + free2remaining = free_ind2[:] + for pos1, index1 in enumerate(free_ind1): + if index1 in free2remaining: + pos2 = free2remaining.index(index1) + free2remaining[pos2] = None + continue + if -index1 in free2remaining: + pos2 = free2remaining.index(-index1) + free2remaining[pos2] = None + free_ind2[pos2] = index1 + if index1.is_up: + pos2up.append(pos2) + else: + pos2down.append(pos2) + else: + index2 = free2remaining[pos1] + if index2 is None: + raise ValueError("incompatible indices: %s and %s" % (free_ind1, free_ind2)) + free2remaining[pos1] = None + free_ind2[pos1] = index1 + if index1.is_up ^ index2.is_up: + if index1.is_up: + pos2up.append(pos1) + else: + pos2down.append(pos1) + + if len(set(free_ind1) & set(free_ind2)) < len(free_ind1): + raise ValueError("incompatible indices: %s and %s" % (free_ind1, free_ind2)) + + # Raise indices: + for pos in pos2up: + index_type_pos = index_types1[pos] + if index_type_pos not in replacement_dict: + raise ValueError("No metric provided to lower index") + metric = replacement_dict[index_type_pos] + metric_inverse = _TensorDataLazyEvaluator.inverse_matrix(metric) + array = TensExpr._contract_and_permute_with_metric(metric_inverse, array, pos, len(free_ind1)) + # Lower indices: + for pos in pos2down: + index_type_pos = index_types1[pos] + if index_type_pos not in replacement_dict: + raise ValueError("No metric provided to lower index") + metric = replacement_dict[index_type_pos] + array = TensExpr._contract_and_permute_with_metric(metric, array, pos, len(free_ind1)) + + if free_ind1: + permutation = TensExpr._get_indices_permutation(free_ind2, free_ind1) + array = permutedims(array, permutation) + + if hasattr(array, "rank") and array.rank() == 0: + array = array[()] + + return free_ind2, array + + def replace_with_arrays(self, replacement_dict, indices=None): + """ + Replace the tensorial expressions with arrays. The final array will + correspond to the N-dimensional array with indices arranged according + to ``indices``. + + Parameters + ========== + + replacement_dict + dictionary containing the replacement rules for tensors. + indices + the index order with respect to which the array is read. The + original index order will be used if no value is passed. + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_indices + >>> from sympy.tensor.tensor import TensorHead + >>> from sympy import symbols, diag + + >>> L = TensorIndexType("L") + >>> i, j = tensor_indices("i j", L) + >>> A = TensorHead("A", [L]) + >>> A(i).replace_with_arrays({A(i): [1, 2]}, [i]) + [1, 2] + + Since 'indices' is optional, we can also call replace_with_arrays by + this way if no specific index order is needed: + + >>> A(i).replace_with_arrays({A(i): [1, 2]}) + [1, 2] + + >>> expr = A(i)*A(j) + >>> expr.replace_with_arrays({A(i): [1, 2]}) + [[1, 2], [2, 4]] + + For contractions, specify the metric of the ``TensorIndexType``, which + in this case is ``L``, in its covariant form: + + >>> expr = A(i)*A(-i) + >>> expr.replace_with_arrays({A(i): [1, 2], L: diag(1, -1)}) + -3 + + Symmetrization of an array: + + >>> H = TensorHead("H", [L, L]) + >>> a, b, c, d = symbols("a b c d") + >>> expr = H(i, j)/2 + H(j, i)/2 + >>> expr.replace_with_arrays({H(i, j): [[a, b], [c, d]]}) + [[a, b/2 + c/2], [b/2 + c/2, d]] + + Anti-symmetrization of an array: + + >>> expr = H(i, j)/2 - H(j, i)/2 + >>> repl = {H(i, j): [[a, b], [c, d]]} + >>> expr.replace_with_arrays(repl) + [[0, b/2 - c/2], [-b/2 + c/2, 0]] + + The same expression can be read as the transpose by inverting ``i`` and + ``j``: + + >>> expr.replace_with_arrays(repl, [j, i]) + [[0, -b/2 + c/2], [b/2 - c/2, 0]] + """ + from .array import Array + + indices = indices or [] + remap = {k.args[0] if k.is_up else -k.args[0]: k for k in self.get_free_indices()} + for i, index in enumerate(indices): + if isinstance(index, (Symbol, Mul)): + if index in remap: + indices[i] = remap[index] + else: + indices[i] = -remap[-index] + + replacement_dict = {tensor: Array(array) for tensor, array in replacement_dict.items()} + + # Check dimensions of replaced arrays: + for tensor, array in replacement_dict.items(): + if isinstance(tensor, TensorIndexType): + expected_shape = [tensor.dim for i in range(2)] + else: + expected_shape = [index_type.dim for index_type in tensor.index_types] + if len(expected_shape) != array.rank() or (not all(dim1 == dim2 if + dim1.is_number else True for dim1, dim2 in zip(expected_shape, + array.shape))): + raise ValueError("shapes for tensor %s expected to be %s, "\ + "replacement array shape is %s" % (tensor, expected_shape, + array.shape)) + + ret_indices, array = self._extract_data(replacement_dict) + + last_indices, array = self._match_indices_with_other_tensor(array, indices, ret_indices, replacement_dict) + return array + + def _check_add_Sum(self, expr, index_symbols): + from sympy.concrete.summations import Sum + indices = self.get_indices() + dum = self.dum + sum_indices = [ (index_symbols[i], 0, + indices[i].tensor_index_type.dim-1) for i, j in dum] + if sum_indices: + expr = Sum(expr, *sum_indices) + return expr + + def _expand_partial_derivative(self): + # simply delegate the _expand_partial_derivative() to + # its arguments to expand a possibly found PartialDerivative + return self.func(*[ + a._expand_partial_derivative() + if isinstance(a, TensExpr) else a + for a in self.args]) + + +class TensAdd(TensExpr, AssocOp): + """ + Sum of tensors. + + Parameters + ========== + + free_args : list of the free indices + + Attributes + ========== + + ``args`` : tuple of addends + ``rank`` : rank of the tensor + ``free_args`` : list of the free indices in sorted order + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_heads, tensor_indices + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> a, b = tensor_indices('a,b', Lorentz) + >>> p, q = tensor_heads('p,q', [Lorentz]) + >>> t = p(a) + q(a); t + p(a) + q(a) + + Examples with components data added to the tensor expression: + + >>> from sympy import symbols, diag + >>> x, y, z, t = symbols("x y z t") + >>> repl = {} + >>> repl[Lorentz] = diag(1, -1, -1, -1) + >>> repl[p(a)] = [1, 2, 3, 4] + >>> repl[q(a)] = [x, y, z, t] + + The following are: 2**2 - 3**2 - 2**2 - 7**2 ==> -58 + + >>> expr = p(a) + q(a) + >>> expr.replace_with_arrays(repl, [a]) + [x + 1, y + 2, z + 3, t + 4] + """ + + def __new__(cls, *args, **kw_args): + args = [_sympify(x) for x in args if x] + args = TensAdd._tensAdd_flatten(args) + args.sort(key=default_sort_key) + if not args: + return S.Zero + if len(args) == 1: + return args[0] + + return Basic.__new__(cls, *args, **kw_args) + + @property + def coeff(self): + return S.One + + @property + def nocoeff(self): + return self + + def get_free_indices(self) -> list[TensorIndex]: + return self.free_indices + + def _replace_indices(self, repl: dict[TensorIndex, TensorIndex]) -> TensExpr: + newargs = [arg._replace_indices(repl) if isinstance(arg, TensExpr) else arg for arg in self.args] + return self.func(*newargs) + + @memoize_property + def rank(self): + if isinstance(self.args[0], TensExpr): + return self.args[0].rank + else: + return 0 + + @memoize_property + def free_args(self): + if isinstance(self.args[0], TensExpr): + return self.args[0].free_args + else: + return [] + + @memoize_property + def free_indices(self): + if isinstance(self.args[0], TensExpr): + return self.args[0].get_free_indices() + else: + return set() + + def doit(self, **hints): + deep = hints.get('deep', True) + if deep: + args = [arg.doit(**hints) for arg in self.args] + else: + args = self.args + + # if any of the args are zero (after doit), drop them. Otherwise, _tensAdd_check will complain about non-matching indices, even though the TensAdd is correctly formed. + args = [arg for arg in args if arg != S.Zero] + + if len(args) == 0: + return S.Zero + elif len(args) == 1: + return args[0] + + # now check that all addends have the same indices: + TensAdd._tensAdd_check(args) + + # Collect terms appearing more than once, differing by their coefficients: + args = TensAdd._tensAdd_collect_terms(args) + + # collect canonicalized terms + def sort_key(t): + if not isinstance(t, TensExpr): + return [], [], [] + if hasattr(t, "_index_structure") and hasattr(t, "components"): + x = get_index_structure(t) + return t.components, x.free, x.dum + return [], [], [] + args.sort(key=sort_key) + + if not args: + return S.Zero + # it there is only a component tensor return it + if len(args) == 1: + return args[0] + + obj = self.func(*args) + return obj + + @staticmethod + def _tensAdd_flatten(args): + # flatten TensAdd, coerce terms which are not tensors to tensors + a = [] + for x in args: + if isinstance(x, (Add, TensAdd)): + a.extend(list(x.args)) + else: + a.append(x) + args = [x for x in a if x.coeff] + return args + + @staticmethod + def _tensAdd_check(args): + # check that all addends have the same free indices + + def get_indices_set(x: Expr) -> set[TensorIndex]: + if isinstance(x, TensExpr): + return set(x.get_free_indices()) + return set() + + indices0 = get_indices_set(args[0]) + list_indices = [get_indices_set(arg) for arg in args[1:]] + if not all(x == indices0 for x in list_indices): + raise ValueError('all tensors must have the same indices') + + @staticmethod + def _tensAdd_collect_terms(args): + # collect TensMul terms differing at most by their coefficient + terms_dict = defaultdict(list) + scalars = S.Zero + if isinstance(args[0], TensExpr): + free_indices = set(args[0].get_free_indices()) + else: + free_indices = set() + + for arg in args: + if not isinstance(arg, TensExpr): + if free_indices != set(): + raise ValueError("wrong valence") + scalars += arg + continue + if free_indices != set(arg.get_free_indices()): + raise ValueError("wrong valence") + # TODO: what is the part which is not a coeff? + # needs an implementation similar to .as_coeff_Mul() + terms_dict[arg.nocoeff].append(arg.coeff) + + new_args = [TensMul(Add(*coeff), t).doit() for t, coeff in terms_dict.items() if Add(*coeff) != 0] + if isinstance(scalars, Add): + new_args = list(scalars.args) + new_args + elif scalars != 0: + new_args = [scalars] + new_args + return new_args + + def get_indices(self): + indices = [] + for arg in self.args: + indices.extend([i for i in get_indices(arg) if i not in indices]) + return indices + + def _expand(self, **hints): + return TensAdd(*[_expand(i, **hints) for i in self.args]) + + def __call__(self, *indices): + deprecate_call() + free_args = self.free_args + indices = list(indices) + if [x.tensor_index_type for x in indices] != [x.tensor_index_type for x in free_args]: + raise ValueError('incompatible types') + if indices == free_args: + return self + index_tuples = list(zip(free_args, indices)) + a = [x.func(*x.substitute_indices(*index_tuples).args) for x in self.args] + res = TensAdd(*a).doit() + return res + + def canon_bp(self): + """ + Canonicalize using the Butler-Portugal algorithm for canonicalization + under monoterm symmetries. + """ + expr = self.expand() + args = [canon_bp(x) for x in expr.args] + res = TensAdd(*args).doit() + return res + + def equals(self, other): + other = _sympify(other) + if isinstance(other, TensMul) and other.coeff == 0: + return all(x.coeff == 0 for x in self.args) + if isinstance(other, TensExpr): + if self.rank != other.rank: + return False + if isinstance(other, TensAdd): + if set(self.args) != set(other.args): + return False + else: + return True + t = self - other + if not isinstance(t, TensExpr): + return t == 0 + else: + if isinstance(t, TensMul): + return t.coeff == 0 + else: + return all(x.coeff == 0 for x in t.args) + + def __getitem__(self, item): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + return self.data[item] + + def contract_delta(self, delta): + args = [x.contract_delta(delta) for x in self.args] + t = TensAdd(*args).doit() + return canon_bp(t) + + def contract_metric(self, g): + """ + Raise or lower indices with the metric ``g``. + + Parameters + ========== + + g : metric + + contract_all : if True, eliminate all ``g`` which are contracted + + Notes + ===== + + see the ``TensorIndexType`` docstring for the contraction conventions + """ + + args = [contract_metric(x, g) for x in self.args] + t = TensAdd(*args).doit() + return canon_bp(t) + + def substitute_indices(self, *index_tuples): + new_args = [] + for arg in self.args: + if isinstance(arg, TensExpr): + arg = arg.substitute_indices(*index_tuples) + new_args.append(arg) + return TensAdd(*new_args).doit() + + def _print(self): + a = [] + args = self.args + for x in args: + a.append(str(x)) + s = ' + '.join(a) + s = s.replace('+ -', '- ') + return s + + def _extract_data(self, replacement_dict): + from sympy.tensor.array import Array, permutedims + args_indices, arrays = zip(*[ + arg._extract_data(replacement_dict) if + isinstance(arg, TensExpr) else ([], arg) for arg in self.args + ]) + arrays = [Array(i) for i in arrays] + ref_indices = args_indices[0] + for i in range(1, len(args_indices)): + indices = args_indices[i] + array = arrays[i] + permutation = TensMul._get_indices_permutation(indices, ref_indices) + arrays[i] = permutedims(array, permutation) + return ref_indices, sum(arrays, Array.zeros(*array.shape)) + + @property + def data(self): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + return _tensor_data_substitution_dict[self.expand()] + + @data.setter + def data(self, data): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + _tensor_data_substitution_dict[self] = data + + @data.deleter + def data(self): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + if self in _tensor_data_substitution_dict: + del _tensor_data_substitution_dict[self] + + def __iter__(self): + deprecate_data() + if not self.data: + raise ValueError("No iteration on abstract tensors") + return self.data.flatten().__iter__() + + def _eval_rewrite_as_Indexed(self, *args, **kwargs): + return Add.fromiter(args) + + def _eval_partial_derivative(self, s): + # Evaluation like Add + list_addends = [] + for a in self.args: + if isinstance(a, TensExpr): + list_addends.append(a._eval_partial_derivative(s)) + # do not call diff if s is no symbol + elif s._diff_wrt: + list_addends.append(a._eval_derivative(s)) + + return self.func(*list_addends) + + +class Tensor(TensExpr): + """ + Base tensor class, i.e. this represents a tensor, the single unit to be + put into an expression. + + Explanation + =========== + + This object is usually created from a ``TensorHead``, by attaching indices + to it. Indices preceded by a minus sign are considered contravariant, + otherwise covariant. + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_indices, TensorHead + >>> Lorentz = TensorIndexType("Lorentz", dummy_name="L") + >>> mu, nu = tensor_indices('mu nu', Lorentz) + >>> A = TensorHead("A", [Lorentz, Lorentz]) + >>> A(mu, -nu) + A(mu, -nu) + >>> A(mu, -mu) + A(L_0, -L_0) + + It is also possible to use symbols instead of inidices (appropriate indices + are then generated automatically). + + >>> from sympy import Symbol + >>> x = Symbol('x') + >>> A(x, mu) + A(x, mu) + >>> A(x, -x) + A(L_0, -L_0) + + """ + + is_commutative = False + + _index_structure = None # type: _IndexStructure + args: tuple[TensorHead, Tuple] + + def __new__(cls, tensor_head, indices, *, is_canon_bp=False, **kw_args): + indices = cls._parse_indices(tensor_head, indices) + obj = Basic.__new__(cls, tensor_head, Tuple(*indices), **kw_args) + obj._index_structure = _IndexStructure.from_indices(*indices) + obj._free = obj._index_structure.free[:] + obj._dum = obj._index_structure.dum[:] + obj._ext_rank = obj._index_structure._ext_rank + obj._coeff = S.One + obj._nocoeff = obj + obj._component = tensor_head + obj._components = [tensor_head] + if tensor_head.rank != len(indices): + raise ValueError("wrong number of indices") + obj.is_canon_bp = is_canon_bp + obj._index_map = Tensor._build_index_map(indices, obj._index_structure) + return obj + + @property + def free(self): + return self._free + + @property + def dum(self): + return self._dum + + @property + def ext_rank(self): + return self._ext_rank + + @property + def coeff(self): + return self._coeff + + @property + def nocoeff(self): + return self._nocoeff + + @property + def component(self): + return self._component + + @property + def components(self): + return self._components + + @property + def head(self): + return self.args[0] + + @property + def indices(self): + return self.args[1] + + @property + def free_indices(self): + return set(self._index_structure.get_free_indices()) + + @property + def index_types(self): + return self.head.index_types + + @property + def rank(self): + return len(self.free_indices) + + @staticmethod + def _build_index_map(indices, index_structure): + index_map = {} + for idx in indices: + index_map[idx] = (indices.index(idx),) + return index_map + + def doit(self, **hints): + args, indices, free, dum = TensMul._tensMul_contract_indices([self]) + return args[0] + + @staticmethod + def _parse_indices(tensor_head, indices): + if not isinstance(indices, (tuple, list, Tuple)): + raise TypeError("indices should be an array, got %s" % type(indices)) + indices = list(indices) + for i, index in enumerate(indices): + if isinstance(index, Symbol): + indices[i] = TensorIndex(index, tensor_head.index_types[i], True) + elif isinstance(index, Mul): + c, e = index.as_coeff_Mul() + if c == -1 and isinstance(e, Symbol): + indices[i] = TensorIndex(e, tensor_head.index_types[i], False) + else: + raise ValueError("index not understood: %s" % index) + elif not isinstance(index, TensorIndex): + raise TypeError("wrong type for index: %s is %s" % (index, type(index))) + return indices + + def _set_new_index_structure(self, im, is_canon_bp=False): + indices = im.get_indices() + return self._set_indices(*indices, is_canon_bp=is_canon_bp) + + def _set_indices(self, *indices, is_canon_bp=False, **kw_args): + if len(indices) != self.ext_rank: + raise ValueError("indices length mismatch") + return self.func(self.args[0], indices, is_canon_bp=is_canon_bp).doit() + + def _get_free_indices_set(self): + return {i[0] for i in self._index_structure.free} + + def _get_dummy_indices_set(self): + dummy_pos = set(itertools.chain(*self._index_structure.dum)) + return {idx for i, idx in enumerate(self.args[1]) if i in dummy_pos} + + def _get_indices_set(self): + return set(self.args[1].args) + + @property + def free_in_args(self): + return [(ind, pos, 0) for ind, pos in self.free] + + @property + def dum_in_args(self): + return [(p1, p2, 0, 0) for p1, p2 in self.dum] + + @property + def free_args(self): + return sorted([x[0] for x in self.free]) + + def commutes_with(self, other): + """ + :param other: + :return: + 0 commute + 1 anticommute + None neither commute nor anticommute + """ + if not isinstance(other, TensExpr): + return 0 + elif isinstance(other, Tensor): + return self.component.commutes_with(other.component) + return NotImplementedError + + def perm2tensor(self, g, is_canon_bp=False): + """ + Returns the tensor corresponding to the permutation ``g``. + + For further details, see the method in ``TIDS`` with the same name. + """ + return perm2tensor(self, g, is_canon_bp) + + def canon_bp(self): + if self.is_canon_bp: + return self + expr = self.expand() + g, dummies, msym = expr._index_structure.indices_canon_args() + v = components_canon_args([expr.component]) + can = canonicalize(g, dummies, msym, *v) + if can == 0: + return S.Zero + tensor = self.perm2tensor(can, True) + return tensor + + def split(self): + return [self] + + def _expand(self, **kwargs): + return self + + def sorted_components(self): + return self + + def get_indices(self) -> list[TensorIndex]: + """ + Get a list of indices, corresponding to those of the tensor. + """ + return list(self.args[1]) + + def get_free_indices(self) -> list[TensorIndex]: + """ + Get a list of free indices, corresponding to those of the tensor. + """ + return self._index_structure.get_free_indices() + + def _replace_indices(self, repl: dict[TensorIndex, TensorIndex]) -> TensExpr: + # TODO: this could be optimized by only swapping the indices + # instead of visiting the whole expression tree: + return self.xreplace(repl) + + def as_base_exp(self): + return self, S.One + + def substitute_indices(self, *index_tuples): + """ + Return a tensor with free indices substituted according to ``index_tuples``. + + ``index_types`` list of tuples ``(old_index, new_index)``. + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_indices, tensor_heads, TensorSymmetry + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> i, j, k, l = tensor_indices('i,j,k,l', Lorentz) + >>> A, B = tensor_heads('A,B', [Lorentz]*2, TensorSymmetry.fully_symmetric(2)) + >>> t = A(i, k)*B(-k, -j); t + A(i, L_0)*B(-L_0, -j) + >>> t.substitute_indices((i, k),(-j, l)) + A(k, L_0)*B(-L_0, l) + """ + indices = [] + for index in self.indices: + for ind_old, ind_new in index_tuples: + if (index.name == ind_old.name and index.tensor_index_type == + ind_old.tensor_index_type): + if index.is_up == ind_old.is_up: + indices.append(ind_new) + else: + indices.append(-ind_new) + break + else: + indices.append(index) + return self.head(*indices) + + def _get_symmetrized_forms(self): + """ + Return a list giving all possible permutations of self that are allowed by its symmetries. + """ + comp = self.component + gens = comp.symmetry.generators + rank = comp.rank + + old_perms = None + new_perms = {self} + while new_perms != old_perms: + old_perms = new_perms.copy() + for tens in old_perms: + for gen in gens: + inds = tens.get_indices() + per = [gen.apply(i) for i in range(0,rank)] + sign = (-1)**(gen.apply(rank) - rank) + ind_map = dict(zip(inds, [inds[i] for i in per])) + new_perms.add( sign * tens._replace_indices(ind_map) ) + + return new_perms + + def matches(self, expr, repl_dict=None, old=False): + expr = sympify(expr) + + if repl_dict is None: + repl_dict = {} + else: + repl_dict = repl_dict.copy() + + #simple checks + if self == expr: + return repl_dict + if not isinstance(expr, Tensor): + return None + if self.head != expr.head: + return None + + #Now consider all index symmetries of expr, and see if any of them allow a match. + for new_expr in expr._get_symmetrized_forms(): + m = self._matches(new_expr, repl_dict, old=old) + if m is not None: + repl_dict.update(m) + return repl_dict + + return None + + def _matches(self, expr, repl_dict=None, old=False): + """ + This does not account for index symmetries of expr + """ + expr = sympify(expr) + + if repl_dict is None: + repl_dict = {} + else: + repl_dict = repl_dict.copy() + + #simple checks + if self == expr: + return repl_dict + if not isinstance(expr, Tensor): + return None + if self.head != expr.head: + return None + + s_indices = self.get_indices() + e_indices = expr.get_indices() + + if len(s_indices) != len(e_indices): + return None + + for i in range(len(s_indices)): + s_ind = s_indices[i] + m = s_ind.matches(e_indices[i]) + if m is None: + return None + elif -s_ind in repl_dict.keys() and -repl_dict[-s_ind] != m[s_ind]: + return None + else: + repl_dict.update(m) + + return repl_dict + + def __call__(self, *indices): + deprecate_call() + free_args = self.free_args + indices = list(indices) + if [x.tensor_index_type for x in indices] != [x.tensor_index_type for x in free_args]: + raise ValueError('incompatible types') + if indices == free_args: + return self + t = self.substitute_indices(*list(zip(free_args, indices))) + + # object is rebuilt in order to make sure that all contracted indices + # get recognized as dummies, but only if there are contracted indices. + if len({i if i.is_up else -i for i in indices}) != len(indices): + return t.func(*t.args) + return t + + # TODO: put this into TensExpr? + def __iter__(self): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + return self.data.__iter__() + + # TODO: put this into TensExpr? + def __getitem__(self, item): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + return self.data[item] + + def _extract_data(self, replacement_dict): + from .array import Array + for k, v in replacement_dict.items(): + if isinstance(k, Tensor) and k.args[0] == self.args[0]: + other = k + array = v + break + else: + raise ValueError("%s not found in %s" % (self, replacement_dict)) + + # TODO: inefficient, this should be done at root level only: + replacement_dict = {k: Array(v) for k, v in replacement_dict.items()} + array = Array(array) + + dum1 = self.dum + dum2 = other.dum + + if len(dum2) > 0: + for pair in dum2: + # allow `dum2` if the contained values are also in `dum1`. + if pair not in dum1: + raise NotImplementedError("%s with contractions is not implemented" % other) + # Remove elements in `dum2` from `dum1`: + dum1 = [pair for pair in dum1 if pair not in dum2] + if len(dum1) > 0: + indices1 = self.get_indices() + indices2 = other.get_indices() + repl = {} + for p1, p2 in dum1: + repl[indices2[p2]] = -indices2[p1] + for pos in (p1, p2): + if indices1[pos].is_up ^ indices2[pos].is_up: + metric = replacement_dict[indices1[pos].tensor_index_type] + if indices1[pos].is_up: + metric = _TensorDataLazyEvaluator.inverse_matrix(metric) + array = self._contract_and_permute_with_metric(metric, array, pos, len(indices2)) + other = other.xreplace(repl).doit() + array = _TensorDataLazyEvaluator.data_contract_dum([array], dum1, len(indices2)) + + free_ind1 = self.get_free_indices() + free_ind2 = other.get_free_indices() + + return self._match_indices_with_other_tensor(array, free_ind1, free_ind2, replacement_dict) + + @property + def data(self): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + return _tensor_data_substitution_dict[self] + + @data.setter + def data(self, data): + deprecate_data() + # TODO: check data compatibility with properties of tensor. + with ignore_warnings(SymPyDeprecationWarning): + _tensor_data_substitution_dict[self] = data + + @data.deleter + def data(self): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + if self in _tensor_data_substitution_dict: + del _tensor_data_substitution_dict[self] + if self.metric in _tensor_data_substitution_dict: + del _tensor_data_substitution_dict[self.metric] + + def _print(self): + indices = [str(ind) for ind in self.indices] + component = self.component + if component.rank > 0: + return ('%s(%s)' % (component.name, ', '.join(indices))) + else: + return ('%s' % component.name) + + def equals(self, other): + if other == 0: + return self.coeff == 0 + other = _sympify(other) + if not isinstance(other, TensExpr): + assert not self.components + return S.One == other + + def _get_compar_comp(self): + t = self.canon_bp() + r = (t.coeff, tuple(t.components), \ + tuple(sorted(t.free)), tuple(sorted(t.dum))) + return r + + return _get_compar_comp(self) == _get_compar_comp(other) + + def contract_metric(self, g): + # if metric is not the same, ignore this step: + if self.component != g: + return self + # in case there are free components, do not perform anything: + if len(self.free) != 0: + return self + + #antisym = g.index_types[0].metric_antisym + if g.symmetry == TensorSymmetry.fully_symmetric(-2): + antisym = 1 + elif g.symmetry == TensorSymmetry.fully_symmetric(2): + antisym = 0 + elif g.symmetry == TensorSymmetry.no_symmetry(2): + antisym = None + else: + raise NotImplementedError + sign = S.One + typ = g.index_types[0] + + if not antisym: + # g(i, -i) + sign = sign*typ.dim + else: + # g(i, -i) + sign = sign*typ.dim + + dp0, dp1 = self.dum[0] + if dp0 < dp1: + # g(i, -i) = -D with antisymmetric metric + sign = -sign + + return sign + + def contract_delta(self, metric): + return self.contract_metric(metric) + + def _eval_rewrite_as_Indexed(self, tens, indices, **kwargs): + from sympy.tensor.indexed import Indexed + # TODO: replace .args[0] with .name: + index_symbols = [i.args[0] for i in self.get_indices()] + expr = Indexed(tens.args[0], *index_symbols) + return self._check_add_Sum(expr, index_symbols) + + def _eval_partial_derivative(self, s): # type: (Tensor) -> Expr + + if not isinstance(s, Tensor): + return S.Zero + else: + + # @a_i/@a_k = delta_i^k + # @a_i/@a^k = g_ij delta^j_k + # @a^i/@a^k = delta^i_k + # @a^i/@a_k = g^ij delta_j^k + # TODO: if there is no metric present, the derivative should be zero? + + if self.head != s.head: + return S.Zero + + # if heads are the same, provide delta and/or metric products + # for every free index pair in the appropriate tensor + # assumed that the free indices are in proper order + # A contravariante index in the derivative becomes covariant + # after performing the derivative and vice versa + + kronecker_delta_list = [1] + + # not guarantee a correct index order + + for (count, (iself, iother)) in enumerate(zip(self.get_free_indices(), s.get_free_indices())): + if iself.tensor_index_type != iother.tensor_index_type: + raise ValueError("index types not compatible") + else: + tensor_index_type = iself.tensor_index_type + tensor_metric = tensor_index_type.metric + dummy = TensorIndex("d_" + str(count), tensor_index_type, + is_up=iself.is_up) + if iself.is_up == iother.is_up: + kroneckerdelta = tensor_index_type.delta(iself, -iother) + else: + kroneckerdelta = ( + TensMul(tensor_metric(iself, dummy), + tensor_index_type.delta(-dummy, -iother)) + ) + kronecker_delta_list.append(kroneckerdelta) + return TensMul.fromiter(kronecker_delta_list).doit() + # doit necessary to rename dummy indices accordingly + + +class TensMul(TensExpr, AssocOp): + """ + Product of tensors. + + Parameters + ========== + + coeff : SymPy coefficient of the tensor + args + + Attributes + ========== + + ``components`` : list of ``TensorHead`` of the component tensors + ``types`` : list of nonrepeated ``TensorIndexType`` + ``free`` : list of ``(ind, ipos, icomp)``, see Notes + ``dum`` : list of ``(ipos1, ipos2, icomp1, icomp2)``, see Notes + ``ext_rank`` : rank of the tensor counting the dummy indices + ``rank`` : rank of the tensor + ``coeff`` : SymPy coefficient of the tensor + ``free_args`` : list of the free indices in sorted order + ``is_canon_bp`` : ``True`` if the tensor in in canonical form + + Notes + ===== + + ``args[0]`` list of ``TensorHead`` of the component tensors. + + ``args[1]`` list of ``(ind, ipos, icomp)`` + where ``ind`` is a free index, ``ipos`` is the slot position + of ``ind`` in the ``icomp``-th component tensor. + + ``args[2]`` list of tuples representing dummy indices. + ``(ipos1, ipos2, icomp1, icomp2)`` indicates that the contravariant + dummy index is the ``ipos1``-th slot position in the ``icomp1``-th + component tensor; the corresponding covariant index is + in the ``ipos2`` slot position in the ``icomp2``-th component tensor. + + """ + identity = S.One + + _index_structure = None # type: _IndexStructure + + def __new__(cls, *args, **kw_args): + is_canon_bp = kw_args.get('is_canon_bp', False) + args = list(map(_sympify, args)) + + """ + If the internal dummy indices in one arg conflict with the free indices + of the remaining args, we need to rename those internal dummy indices. + """ + free = [get_free_indices(arg) for arg in args] + free = set(itertools.chain(*free)) #flatten free + newargs = [] + for arg in args: + dum_this = set(get_dummy_indices(arg)) + dum_other = [get_dummy_indices(a) for a in newargs] + dum_other = set(itertools.chain(*dum_other)) #flatten dum_other + free_this = set(get_free_indices(arg)) + if len(dum_this.intersection(free)) > 0: + exclude = free_this.union(free, dum_other) + newarg = TensMul._dedupe_indices(arg, exclude) + else: + newarg = arg + newargs.append(newarg) + + args = newargs + + # Flatten: + args = [i for arg in args for i in (arg.args if isinstance(arg, (TensMul, Mul)) else [arg])] + + args, indices, free, dum = TensMul._tensMul_contract_indices(args, replace_indices=False) + + # Data for indices: + index_types = [i.tensor_index_type for i in indices] + index_structure = _IndexStructure(free, dum, index_types, indices, canon_bp=is_canon_bp) + + obj = TensExpr.__new__(cls, *args) + obj._indices = indices + obj._index_types = index_types[:] + obj._index_structure = index_structure + obj._free = index_structure.free[:] + obj._dum = index_structure.dum[:] + obj._free_indices = {x[0] for x in obj.free} + obj._rank = len(obj.free) + obj._ext_rank = len(obj._index_structure.free) + 2*len(obj._index_structure.dum) + obj._coeff = S.One + obj._is_canon_bp = is_canon_bp + return obj + + index_types = property(lambda self: self._index_types) + free = property(lambda self: self._free) + dum = property(lambda self: self._dum) + free_indices = property(lambda self: self._free_indices) + rank = property(lambda self: self._rank) + ext_rank = property(lambda self: self._ext_rank) + + @staticmethod + def _indices_to_free_dum(args_indices): + free2pos1 = {} + free2pos2 = {} + dummy_data = [] + indices = [] + + # Notation for positions (to better understand the code): + # `pos1`: position in the `args`. + # `pos2`: position in the indices. + + # Example: + # A(i, j)*B(k, m, n)*C(p) + # `pos1` of `n` is 1 because it's in `B` (second `args` of TensMul). + # `pos2` of `n` is 4 because it's the fifth overall index. + + # Counter for the index position wrt the whole expression: + pos2 = 0 + + for pos1, arg_indices in enumerate(args_indices): + + for index in arg_indices: + if not isinstance(index, TensorIndex): + raise TypeError("expected TensorIndex") + if -index in free2pos1: + # Dummy index detected: + other_pos1 = free2pos1.pop(-index) + other_pos2 = free2pos2.pop(-index) + if index.is_up: + dummy_data.append((index, pos1, other_pos1, pos2, other_pos2)) + else: + dummy_data.append((-index, other_pos1, pos1, other_pos2, pos2)) + indices.append(index) + elif index in free2pos1: + raise ValueError("Repeated index: %s" % index) + else: + free2pos1[index] = pos1 + free2pos2[index] = pos2 + indices.append(index) + pos2 += 1 + + free = list(free2pos2.items()) + free_names = [i.name for i in free2pos2.keys()] + + dummy_data.sort(key=lambda x: x[3]) + return indices, free, free_names, dummy_data + + @staticmethod + def _dummy_data_to_dum(dummy_data): + return [(p2a, p2b) for (i, p1a, p1b, p2a, p2b) in dummy_data] + + @staticmethod + def _tensMul_contract_indices(args, replace_indices=True): + replacements = [{} for _ in args] + + #_index_order = all(_has_index_order(arg) for arg in args) + + args_indices = [get_indices(arg) for arg in args] + indices, free, free_names, dummy_data = TensMul._indices_to_free_dum(args_indices) + + cdt = defaultdict(int) + + def dummy_name_gen(tensor_index_type): + nd = str(cdt[tensor_index_type]) + cdt[tensor_index_type] += 1 + return tensor_index_type.dummy_name + '_' + nd + + if replace_indices: + for old_index, pos1cov, pos1contra, pos2cov, pos2contra in dummy_data: + index_type = old_index.tensor_index_type + while True: + dummy_name = dummy_name_gen(index_type) + if dummy_name not in free_names: + break + dummy = TensorIndex(dummy_name, index_type, True) + replacements[pos1cov][old_index] = dummy + replacements[pos1contra][-old_index] = -dummy + indices[pos2cov] = dummy + indices[pos2contra] = -dummy + args = [ + arg._replace_indices(repl) if isinstance(arg, TensExpr) else arg + for arg, repl in zip(args, replacements)] + + dum = TensMul._dummy_data_to_dum(dummy_data) + return args, indices, free, dum + + @staticmethod + def _get_components_from_args(args): + """ + Get a list of ``Tensor`` objects having the same ``TIDS`` if multiplied + by one another. + """ + components = [] + for arg in args: + if not isinstance(arg, TensExpr): + continue + if isinstance(arg, TensAdd): + continue + components.extend(arg.components) + return components + + @staticmethod + def _rebuild_tensors_list(args, index_structure): + indices = index_structure.get_indices() + #tensors = [None for i in components] # pre-allocate list + ind_pos = 0 + for i, arg in enumerate(args): + if not isinstance(arg, TensExpr): + continue + prev_pos = ind_pos + ind_pos += arg.ext_rank + args[i] = Tensor(arg.component, indices[prev_pos:ind_pos]) + + def doit(self, **hints): + is_canon_bp = self._is_canon_bp + deep = hints.get('deep', True) + if deep: + args = [arg.doit(**hints) for arg in self.args] + + """ + There may now be conflicts between dummy indices of different args + (each arg's doit method does not have any information about which + dummy indices are already used in the other args), so we + deduplicate them. + """ + rule = dict(zip(self.args, args)) + rule = self._dedupe_indices_in_rule(rule) + args = [rule[a] for a in self.args] + + else: + args = self.args + + args = [arg for arg in args if arg != self.identity] + + # Extract non-tensor coefficients: + coeff = reduce(lambda a, b: a*b, [arg for arg in args if not isinstance(arg, TensExpr)], S.One) + args = [arg for arg in args if isinstance(arg, TensExpr)] + + if len(args) == 0: + return coeff + + if coeff != self.identity: + args = [coeff] + args + if coeff == 0: + return S.Zero + + if len(args) == 1: + return args[0] + + args, indices, free, dum = TensMul._tensMul_contract_indices(args) + + # Data for indices: + index_types = [i.tensor_index_type for i in indices] + index_structure = _IndexStructure(free, dum, index_types, indices, canon_bp=is_canon_bp) + + obj = self.func(*args) + obj._index_types = index_types + obj._index_structure = index_structure + obj._ext_rank = len(obj._index_structure.free) + 2*len(obj._index_structure.dum) + obj._coeff = coeff + obj._is_canon_bp = is_canon_bp + return obj + + # TODO: this method should be private + # TODO: should this method be renamed _from_components_free_dum ? + @staticmethod + def from_data(coeff, components, free, dum, **kw_args): + return TensMul(coeff, *TensMul._get_tensors_from_components_free_dum(components, free, dum), **kw_args).doit() + + @staticmethod + def _get_tensors_from_components_free_dum(components, free, dum): + """ + Get a list of ``Tensor`` objects by distributing ``free`` and ``dum`` indices on the ``components``. + """ + index_structure = _IndexStructure.from_components_free_dum(components, free, dum) + indices = index_structure.get_indices() + tensors = [None for i in components] # pre-allocate list + + # distribute indices on components to build a list of tensors: + ind_pos = 0 + for i, component in enumerate(components): + prev_pos = ind_pos + ind_pos += component.rank + tensors[i] = Tensor(component, indices[prev_pos:ind_pos]) + return tensors + + def _get_free_indices_set(self): + return {i[0] for i in self.free} + + def _get_dummy_indices_set(self): + dummy_pos = set(itertools.chain(*self.dum)) + return {idx for i, idx in enumerate(self._index_structure.get_indices()) if i in dummy_pos} + + def _get_position_offset_for_indices(self): + arg_offset = [None for i in range(self.ext_rank)] + counter = 0 + for arg in self.args: + if not isinstance(arg, TensExpr): + continue + for j in range(arg.ext_rank): + arg_offset[j + counter] = counter + counter += arg.ext_rank + return arg_offset + + @property + def free_args(self): + return sorted([x[0] for x in self.free]) + + @property + def components(self): + return self._get_components_from_args(self.args) + + @property + def free_in_args(self): + arg_offset = self._get_position_offset_for_indices() + argpos = self._get_indices_to_args_pos() + return [(ind, pos-arg_offset[pos], argpos[pos]) for (ind, pos) in self.free] + + @property + def coeff(self): + # return Mul.fromiter([c for c in self.args if not isinstance(c, TensExpr)]) + return self._coeff + + @property + def nocoeff(self): + return self.func(*[t for t in self.args if isinstance(t, TensExpr)]).doit() + + @property + def dum_in_args(self): + arg_offset = self._get_position_offset_for_indices() + argpos = self._get_indices_to_args_pos() + return [(p1-arg_offset[p1], p2-arg_offset[p2], argpos[p1], argpos[p2]) for p1, p2 in self.dum] + + def equals(self, other): + if other == 0: + return self.coeff == 0 + other = _sympify(other) + if not isinstance(other, TensExpr): + assert not self.components + return self.coeff == other + + return self.canon_bp() == other.canon_bp() + + def get_indices(self): + """ + Returns the list of indices of the tensor. + + Explanation + =========== + + The indices are listed in the order in which they appear in the + component tensors. + The dummy indices are given a name which does not collide with + the names of the free indices. + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_indices, tensor_heads + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> m0, m1, m2 = tensor_indices('m0,m1,m2', Lorentz) + >>> g = Lorentz.metric + >>> p, q = tensor_heads('p,q', [Lorentz]) + >>> t = p(m1)*g(m0,m2) + >>> t.get_indices() + [m1, m0, m2] + >>> t2 = p(m1)*g(-m1, m2) + >>> t2.get_indices() + [L_0, -L_0, m2] + """ + return self._indices + + def get_free_indices(self) -> list[TensorIndex]: + """ + Returns the list of free indices of the tensor. + + Explanation + =========== + + The indices are listed in the order in which they appear in the + component tensors. + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_indices, tensor_heads + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> m0, m1, m2 = tensor_indices('m0,m1,m2', Lorentz) + >>> g = Lorentz.metric + >>> p, q = tensor_heads('p,q', [Lorentz]) + >>> t = p(m1)*g(m0,m2) + >>> t.get_free_indices() + [m1, m0, m2] + >>> t2 = p(m1)*g(-m1, m2) + >>> t2.get_free_indices() + [m2] + """ + return self._index_structure.get_free_indices() + + def _replace_indices(self, repl: dict[TensorIndex, TensorIndex]) -> TensExpr: + return self.func(*[arg._replace_indices(repl) if isinstance(arg, TensExpr) else arg for arg in self.args]) + + def split(self): + """ + Returns a list of tensors, whose product is ``self``. + + Explanation + =========== + + Dummy indices contracted among different tensor components + become free indices with the same name as the one used to + represent the dummy indices. + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_indices, tensor_heads, TensorSymmetry + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> a, b, c, d = tensor_indices('a,b,c,d', Lorentz) + >>> A, B = tensor_heads('A,B', [Lorentz]*2, TensorSymmetry.fully_symmetric(2)) + >>> t = A(a,b)*B(-b,c) + >>> t + A(a, L_0)*B(-L_0, c) + >>> t.split() + [A(a, L_0), B(-L_0, c)] + """ + if self.args == (): + return [self] + splitp = [] + res = 1 + for arg in self.args: + if isinstance(arg, Tensor): + splitp.append(res*arg) + res = 1 + else: + res *= arg + return splitp + + def _expand(self, **hints): + # TODO: temporary solution, in the future this should be linked to + # `Expr.expand`. + args = [_expand(arg, **hints) for arg in self.args] + args1 = [arg.args if isinstance(arg, (Add, TensAdd)) else (arg,) for arg in args] + return TensAdd(*[ + TensMul(*i) for i in itertools.product(*args1)] + ) + + def __neg__(self): + return TensMul(S.NegativeOne, self, is_canon_bp=self._is_canon_bp).doit() + + def __getitem__(self, item): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + return self.data[item] + + def _get_args_for_traditional_printer(self): + args = list(self.args) + if self.coeff.could_extract_minus_sign(): + # expressions like "-A(a)" + sign = "-" + if args[0] == S.NegativeOne: + args = args[1:] + else: + args[0] = -args[0] + else: + sign = "" + return sign, args + + def _sort_args_for_sorted_components(self): + """ + Returns the ``args`` sorted according to the components commutation + properties. + + Explanation + =========== + + The sorting is done taking into account the commutation group + of the component tensors. + """ + cv = [arg for arg in self.args if isinstance(arg, TensExpr)] + sign = 1 + n = len(cv) - 1 + for i in range(n): + for j in range(n, i, -1): + c = cv[j-1].commutes_with(cv[j]) + # if `c` is `None`, it does neither commute nor anticommute, skip: + if c not in (0, 1): + continue + typ1 = sorted(set(cv[j-1].component.index_types), key=lambda x: x.name) + typ2 = sorted(set(cv[j].component.index_types), key=lambda x: x.name) + if (typ1, cv[j-1].component.name) > (typ2, cv[j].component.name): + cv[j-1], cv[j] = cv[j], cv[j-1] + # if `c` is 1, the anticommute, so change sign: + if c: + sign = -sign + + coeff = sign * self.coeff + if coeff != 1: + return [coeff] + cv + return cv + + def sorted_components(self): + """ + Returns a tensor product with sorted components. + """ + return TensMul(*self._sort_args_for_sorted_components()).doit() + + def perm2tensor(self, g, is_canon_bp=False): + """ + Returns the tensor corresponding to the permutation ``g`` + + For further details, see the method in ``TIDS`` with the same name. + """ + return perm2tensor(self, g, is_canon_bp=is_canon_bp) + + def canon_bp(self): + """ + Canonicalize using the Butler-Portugal algorithm for canonicalization + under monoterm symmetries. + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_indices, TensorHead, TensorSymmetry + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> m0, m1, m2 = tensor_indices('m0,m1,m2', Lorentz) + >>> A = TensorHead('A', [Lorentz]*2, TensorSymmetry.fully_symmetric(-2)) + >>> t = A(m0,-m1)*A(m1,-m0) + >>> t.canon_bp() + -A(L_0, L_1)*A(-L_0, -L_1) + >>> t = A(m0,-m1)*A(m1,-m2)*A(m2,-m0) + >>> t.canon_bp() + 0 + """ + if self._is_canon_bp: + return self + expr = self.expand() + if isinstance(expr, TensAdd): + return expr.canon_bp() + if not expr.components: + return expr + t = expr.sorted_components() + g, dummies, msym = t._index_structure.indices_canon_args() + v = components_canon_args(t.components) + can = canonicalize(g, dummies, msym, *v) + if can == 0: + return S.Zero + tmul = t.perm2tensor(can, True) + return tmul + + def contract_delta(self, delta): + t = self.contract_metric(delta) + return t + + def _get_indices_to_args_pos(self): + """ + Get a dict mapping the index position to TensMul's argument number. + """ + pos_map = {} + pos_counter = 0 + for arg_i, arg in enumerate(self.args): + if not isinstance(arg, TensExpr): + continue + assert isinstance(arg, Tensor) + for i in range(arg.ext_rank): + pos_map[pos_counter] = arg_i + pos_counter += 1 + return pos_map + + def contract_metric(self, g): + """ + Raise or lower indices with the metric ``g``. + + Parameters + ========== + + g : metric + + Notes + ===== + + See the ``TensorIndexType`` docstring for the contraction conventions. + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_indices, tensor_heads + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> m0, m1, m2 = tensor_indices('m0,m1,m2', Lorentz) + >>> g = Lorentz.metric + >>> p, q = tensor_heads('p,q', [Lorentz]) + >>> t = p(m0)*q(m1)*g(-m0, -m1) + >>> t.canon_bp() + metric(L_0, L_1)*p(-L_0)*q(-L_1) + >>> t.contract_metric(g).canon_bp() + p(L_0)*q(-L_0) + """ + expr = self.expand() + if self != expr: + expr = canon_bp(expr) + return contract_metric(expr, g) + pos_map = self._get_indices_to_args_pos() + args = list(self.args) + + #antisym = g.index_types[0].metric_antisym + if g.symmetry == TensorSymmetry.fully_symmetric(-2): + antisym = 1 + elif g.symmetry == TensorSymmetry.fully_symmetric(2): + antisym = 0 + elif g.symmetry == TensorSymmetry.no_symmetry(2): + antisym = None + else: + raise NotImplementedError + + # list of positions of the metric ``g`` inside ``args`` + gpos = [i for i, x in enumerate(self.args) if isinstance(x, Tensor) and x.component == g] + if not gpos: + return self + + # Sign is either 1 or -1, to correct the sign after metric contraction + # (for spinor indices). + sign = 1 + dum = self.dum[:] + free = self.free[:] + elim = set() + for gposx in gpos: + if gposx in elim: + continue + free1 = [x for x in free if pos_map[x[1]] == gposx] + dum1 = [x for x in dum if pos_map[x[0]] == gposx or pos_map[x[1]] == gposx] + if not dum1: + continue + elim.add(gposx) + # subs with the multiplication neutral element, that is, remove it: + args[gposx] = 1 + if len(dum1) == 2: + if not antisym: + dum10, dum11 = dum1 + if pos_map[dum10[1]] == gposx: + # the index with pos p0 contravariant + p0 = dum10[0] + else: + # the index with pos p0 is covariant + p0 = dum10[1] + if pos_map[dum11[1]] == gposx: + # the index with pos p1 is contravariant + p1 = dum11[0] + else: + # the index with pos p1 is covariant + p1 = dum11[1] + + dum.append((p0, p1)) + else: + dum10, dum11 = dum1 + # change the sign to bring the indices of the metric to contravariant + # form; change the sign if dum10 has the metric index in position 0 + if pos_map[dum10[1]] == gposx: + # the index with pos p0 is contravariant + p0 = dum10[0] + if dum10[1] == 1: + sign = -sign + else: + # the index with pos p0 is covariant + p0 = dum10[1] + if dum10[0] == 0: + sign = -sign + if pos_map[dum11[1]] == gposx: + # the index with pos p1 is contravariant + p1 = dum11[0] + sign = -sign + else: + # the index with pos p1 is covariant + p1 = dum11[1] + + dum.append((p0, p1)) + + elif len(dum1) == 1: + if not antisym: + dp0, dp1 = dum1[0] + if pos_map[dp0] == pos_map[dp1]: + # g(i, -i) + typ = g.index_types[0] + sign = sign*typ.dim + + else: + # g(i0, i1)*p(-i1) + if pos_map[dp0] == gposx: + p1 = dp1 + else: + p1 = dp0 + + ind, p = free1[0] + free.append((ind, p1)) + else: + dp0, dp1 = dum1[0] + if pos_map[dp0] == pos_map[dp1]: + # g(i, -i) + typ = g.index_types[0] + sign = sign*typ.dim + + if dp0 < dp1: + # g(i, -i) = -D with antisymmetric metric + sign = -sign + else: + # g(i0, i1)*p(-i1) + if pos_map[dp0] == gposx: + p1 = dp1 + if dp0 == 0: + sign = -sign + else: + p1 = dp0 + ind, p = free1[0] + free.append((ind, p1)) + dum = [x for x in dum if x not in dum1] + free = [x for x in free if x not in free1] + + # shift positions: + shift = 0 + shifts = [0]*len(args) + for i in range(len(args)): + if i in elim: + shift += 2 + continue + shifts[i] = shift + free = [(ind, p - shifts[pos_map[p]]) for (ind, p) in free if pos_map[p] not in elim] + dum = [(p0 - shifts[pos_map[p0]], p1 - shifts[pos_map[p1]]) for p0, p1 in dum if pos_map[p0] not in elim and pos_map[p1] not in elim] + + res = sign*TensMul(*args).doit() + if not isinstance(res, TensExpr): + return res + im = _IndexStructure.from_components_free_dum(res.components, free, dum) + return res._set_new_index_structure(im) + + def _set_new_index_structure(self, im, is_canon_bp=False): + indices = im.get_indices() + return self._set_indices(*indices, is_canon_bp=is_canon_bp) + + def _set_indices(self, *indices, is_canon_bp=False, **kw_args): + if len(indices) != self.ext_rank: + raise ValueError("indices length mismatch") + args = list(self.args)[:] + pos = 0 + for i, arg in enumerate(args): + if not isinstance(arg, TensExpr): + continue + assert isinstance(arg, Tensor) + ext_rank = arg.ext_rank + args[i] = arg._set_indices(*indices[pos:pos+ext_rank]) + pos += ext_rank + return TensMul(*args, is_canon_bp=is_canon_bp).doit() + + @staticmethod + def _index_replacement_for_contract_metric(args, free, dum): + for arg in args: + if not isinstance(arg, TensExpr): + continue + assert isinstance(arg, Tensor) + + def substitute_indices(self, *index_tuples): + new_args = [] + for arg in self.args: + if isinstance(arg, TensExpr): + arg = arg.substitute_indices(*index_tuples) + new_args.append(arg) + return TensMul(*new_args).doit() + + def __call__(self, *indices): + deprecate_call() + free_args = self.free_args + indices = list(indices) + if [x.tensor_index_type for x in indices] != [x.tensor_index_type for x in free_args]: + raise ValueError('incompatible types') + if indices == free_args: + return self + t = self.substitute_indices(*list(zip(free_args, indices))) + + # object is rebuilt in order to make sure that all contracted indices + # get recognized as dummies, but only if there are contracted indices. + if len({i if i.is_up else -i for i in indices}) != len(indices): + return t.func(*t.args) + return t + + def _extract_data(self, replacement_dict): + args_indices, arrays = zip(*[arg._extract_data(replacement_dict) for arg in self.args if isinstance(arg, TensExpr)]) + coeff = reduce(operator.mul, [a for a in self.args if not isinstance(a, TensExpr)], S.One) + indices, free, free_names, dummy_data = TensMul._indices_to_free_dum(args_indices) + dum = TensMul._dummy_data_to_dum(dummy_data) + ext_rank = self.ext_rank + free.sort(key=lambda x: x[1]) + free_indices = [i[0] for i in free] + return free_indices, coeff*_TensorDataLazyEvaluator.data_contract_dum(arrays, dum, ext_rank) + + @property + def data(self): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + dat = _tensor_data_substitution_dict[self.expand()] + return dat + + @data.setter + def data(self, data): + deprecate_data() + raise ValueError("Not possible to set component data to a tensor expression") + + @data.deleter + def data(self): + deprecate_data() + raise ValueError("Not possible to delete component data to a tensor expression") + + def __iter__(self): + deprecate_data() + with ignore_warnings(SymPyDeprecationWarning): + if self.data is None: + raise ValueError("No iteration on abstract tensors") + return self.data.__iter__() + + @staticmethod + def _dedupe_indices(new, exclude): + """ + exclude: set + new: TensExpr + + If ``new`` has any dummy indices that are in ``exclude``, return a version + of new with those indices replaced. If no replacements are needed, + return None + + """ + exclude = set(exclude) + dums_new = set(get_dummy_indices(new)) + free_new = set(get_free_indices(new)) + + conflicts = dums_new.intersection(exclude) + if len(conflicts) == 0: + return None + + """ + ``exclude_for_gen`` is to be passed to ``_IndexStructure._get_generator_for_dummy_indices()``. + Since the latter does not use the index position for anything, we just + set it as ``None`` here. + """ + exclude.update(dums_new) + exclude.update(free_new) + exclude_for_gen = [(i, None) for i in exclude] + gen = _IndexStructure._get_generator_for_dummy_indices(exclude_for_gen) + repl = {} + for d in conflicts: + if -d in repl.keys(): + continue + newname = gen(d.tensor_index_type) + new_d = d.func(newname, *d.args[1:]) + repl[d] = new_d + repl[-d] = -new_d + + if len(repl) == 0: + return None + + new_renamed = new._replace_indices(repl) + return new_renamed + + def _dedupe_indices_in_rule(self, rule): + """ + rule: dict + + This applies TensMul._dedupe_indices on all values of rule. + + """ + index_rules = {k:v for k,v in rule.items() if isinstance(k, TensorIndex)} + other_rules = {k:v for k,v in rule.items() if k not in index_rules.keys()} + exclude = set(self.get_indices()) + + newrule = {} + newrule.update(index_rules) + exclude.update(index_rules.keys()) + exclude.update(index_rules.values()) + for old, new in other_rules.items(): + new_renamed = TensMul._dedupe_indices(new, exclude) + if old == new or new_renamed is None: + newrule[old] = new + else: + newrule[old] = new_renamed + exclude.update(get_indices(new_renamed)) + return newrule + + def _eval_rewrite_as_Indexed(self, *args, **kwargs): + from sympy.concrete.summations import Sum + index_symbols = [i.args[0] for i in self.get_indices()] + args = [arg.args[0] if isinstance(arg, Sum) else arg for arg in args] + expr = Mul.fromiter(args) + return self._check_add_Sum(expr, index_symbols) + + def _eval_partial_derivative(self, s): + # Evaluation like Mul + terms = [] + for i, arg in enumerate(self.args): + # checking whether some tensor instance is differentiated + # or some other thing is necessary, but ugly + if isinstance(arg, TensExpr): + d = arg._eval_partial_derivative(s) + else: + # do not call diff is s is no symbol + if s._diff_wrt: + d = arg._eval_derivative(s) + else: + d = S.Zero + if d: + terms.append(TensMul.fromiter(self.args[:i] + (d,) + self.args[i + 1:])) + return TensAdd.fromiter(terms) + + +class TensorElement(TensExpr): + """ + Tensor with evaluated components. + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, TensorHead, TensorSymmetry + >>> from sympy import symbols + >>> L = TensorIndexType("L") + >>> i, j, k = symbols("i j k") + >>> A = TensorHead("A", [L, L], TensorSymmetry.fully_symmetric(2)) + >>> A(i, j).get_free_indices() + [i, j] + + If we want to set component ``i`` to a specific value, use the + ``TensorElement`` class: + + >>> from sympy.tensor.tensor import TensorElement + >>> te = TensorElement(A(i, j), {i: 2}) + + As index ``i`` has been accessed (``{i: 2}`` is the evaluation of its 3rd + element), the free indices will only contain ``j``: + + >>> te.get_free_indices() + [j] + """ + + def __new__(cls, expr, index_map): + if not isinstance(expr, Tensor): + # remap + if not isinstance(expr, TensExpr): + raise TypeError("%s is not a tensor expression" % expr) + return expr.func(*[TensorElement(arg, index_map) for arg in expr.args]) + expr_free_indices = expr.get_free_indices() + name_translation = {i.args[0]: i for i in expr_free_indices} + index_map = {name_translation.get(index, index): value for index, value in index_map.items()} + index_map = {index: value for index, value in index_map.items() if index in expr_free_indices} + if len(index_map) == 0: + return expr + free_indices = [i for i in expr_free_indices if i not in index_map.keys()] + index_map = Dict(index_map) + obj = TensExpr.__new__(cls, expr, index_map) + obj._free_indices = free_indices + return obj + + @property + def free(self): + return [(index, i) for i, index in enumerate(self.get_free_indices())] + + @property + def dum(self): + # TODO: inherit dummies from expr + return [] + + @property + def expr(self): + return self._args[0] + + @property + def index_map(self): + return self._args[1] + + @property + def coeff(self): + return S.One + + @property + def nocoeff(self): + return self + + def get_free_indices(self): + return self._free_indices + + def _replace_indices(self, repl: dict[TensorIndex, TensorIndex]) -> TensExpr: + # TODO: can be improved: + return self.xreplace(repl) + + def get_indices(self): + return self.get_free_indices() + + def _extract_data(self, replacement_dict): + ret_indices, array = self.expr._extract_data(replacement_dict) + index_map = self.index_map + slice_tuple = tuple(index_map.get(i, slice(None)) for i in ret_indices) + ret_indices = [i for i in ret_indices if i not in index_map] + array = array.__getitem__(slice_tuple) + return ret_indices, array + + +class WildTensorHead(TensorHead): + """ + A wild object that is used to create ``WildTensor`` instances + + Explanation + =========== + + Examples + ======== + >>> from sympy.tensor.tensor import TensorHead, TensorIndex, WildTensorHead, TensorIndexType + >>> R3 = TensorIndexType('R3', dim=3) + >>> p = TensorIndex('p', R3) + >>> q = TensorIndex('q', R3) + + A WildTensorHead can be created without specifying a ``TensorIndexType`` + + >>> W = WildTensorHead("W") + + Calling it with a ``TensorIndex`` creates a ``WildTensor`` instance. + + >>> type(W(p)) + + + The ``TensorIndexType`` is automatically detected from the index that is passed + + >>> W(p).component + W(R3) + + Calling it with no indices returns an object that can match tensors with any number of indices. + + >>> K = TensorHead('K', [R3]) + >>> Q = TensorHead('Q', [R3, R3]) + >>> W().matches(K(p)) + {W: K(p)} + >>> W().matches(Q(p,q)) + {W: Q(p, q)} + + If you want to ignore the order of indices while matching, pass ``unordered_indices=True``. + + >>> U = WildTensorHead("U", unordered_indices=True) + >>> W(p,q).matches(Q(q,p)) + >>> U(p,q).matches(Q(q,p)) + {U(R3,R3): _WildTensExpr(Q(q, p))} + + Parameters + ========== + name : name of the tensor + unordered_indices : whether the order of the indices matters for matching + (default: False) + + See also + ======== + ``WildTensor`` + ``TensorHead`` + + """ + def __new__(cls, name, index_types=None, symmetry=None, comm=0, unordered_indices=False): + if isinstance(name, str): + name_symbol = Symbol(name) + elif isinstance(name, Symbol): + name_symbol = name + else: + raise ValueError("invalid name") + + if index_types is None: + index_types = [] + + if symmetry is None: + symmetry = TensorSymmetry.no_symmetry(len(index_types)) + else: + assert symmetry.rank == len(index_types) + + if symmetry != TensorSymmetry.no_symmetry(len(index_types)): + raise NotImplementedError("Wild matching based on symmetry is not implemented.") + + obj = Basic.__new__(cls, name_symbol, Tuple(*index_types), sympify(symmetry), sympify(comm), sympify(unordered_indices)) + + return obj + + @property + def unordered_indices(self): + return self.args[4] + + def __call__(self, *indices, **kwargs): + tensor = WildTensor(self, indices, **kwargs) + return tensor.doit() + + +class WildTensor(Tensor): + """ + A wild object which matches ``Tensor`` instances + + Explanation + =========== + This is instantiated by attaching indices to a ``WildTensorHead`` instance. + + Examples + ======== + >>> from sympy.tensor.tensor import TensorHead, TensorIndex, WildTensorHead, TensorIndexType + >>> W = WildTensorHead("W") + >>> R3 = TensorIndexType('R3', dim=3) + >>> p = TensorIndex('p', R3) + >>> q = TensorIndex('q', R3) + >>> K = TensorHead('K', [R3]) + >>> Q = TensorHead('Q', [R3, R3]) + + Matching also takes the indices into account + >>> W(p).matches(K(p)) + {W(R3): _WildTensExpr(K(p))} + >>> W(p).matches(K(q)) + >>> W(p).matches(K(-p)) + + If you want to match objects with any number of indices, just use a ``WildTensor`` with no indices. + >>> W().matches(K(p)) + {W: K(p)} + >>> W().matches(Q(p,q)) + {W: Q(p, q)} + + See Also + ======== + ``WildTensorHead`` + ``Tensor`` + + """ + def __new__(cls, tensor_head, indices, **kw_args): + is_canon_bp = kw_args.pop("is_canon_bp", False) + + if tensor_head.func == TensorHead: + """ + If someone tried to call WildTensor by supplying a TensorHead (not a WildTensorHead), return a normal tensor instead. This is helpful when using subs on an expression to replace occurrences of a WildTensorHead with a TensorHead. + """ + return Tensor(tensor_head, indices, is_canon_bp=is_canon_bp, **kw_args) + elif tensor_head.func == _WildTensExpr: + return tensor_head(*indices) + + indices = cls._parse_indices(tensor_head, indices) + index_types = [ind.tensor_index_type for ind in indices] + tensor_head = tensor_head.func( + tensor_head.name, + index_types, + symmetry=None, + comm=tensor_head.comm, + unordered_indices=tensor_head.unordered_indices, + ) + + obj = Basic.__new__(cls, tensor_head, Tuple(*indices)) + obj.name = tensor_head.name + obj._index_structure = _IndexStructure.from_indices(*indices) + obj._free = obj._index_structure.free[:] + obj._dum = obj._index_structure.dum[:] + obj._ext_rank = obj._index_structure._ext_rank + obj._coeff = S.One + obj._nocoeff = obj + obj._component = tensor_head + obj._components = [tensor_head] + if tensor_head.rank != len(indices): + raise ValueError("wrong number of indices") + obj.is_canon_bp = is_canon_bp + obj._index_map = obj._build_index_map(indices, obj._index_structure) + + return obj + + + def matches(self, expr, repl_dict=None, old=False): + if not isinstance(expr, TensExpr) and expr != S(1): + return None + + if repl_dict is None: + repl_dict = {} + else: + repl_dict = repl_dict.copy() + + if len(self.indices) > 0: + if not hasattr(expr, "get_free_indices"): + return None + expr_indices = expr.get_free_indices() + if len(expr_indices) != len(self.indices): + return None + if self._component.unordered_indices: + m = self._match_indices_ignoring_order(expr) + if m is None: + return None + else: + repl_dict.update(m) + else: + for i in range(len(expr_indices)): + m = self.indices[i].matches(expr_indices[i]) + if m is None: + return None + else: + repl_dict.update(m) + + repl_dict[self.component] = _WildTensExpr(expr) + else: + #If no indices were passed to the WildTensor, it may match tensors with any number of indices. + repl_dict[self] = expr + + return repl_dict + + def _match_indices_ignoring_order(self, expr, repl_dict=None, old=False): + """ + Helper method for matches. Checks if the indices of self and expr + match disregarding index ordering. + """ + if repl_dict is None: + repl_dict = {} + else: + repl_dict = repl_dict.copy() + + def siftkey(ind): + if isinstance(ind, WildTensorIndex): + if ind.ignore_updown: + return "wild, updown" + else: + return "wild" + else: + return "nonwild" + + indices_sifted = sift(self.indices, siftkey) + + matched_indices = [] + expr_indices_remaining = expr.get_indices() + for ind in indices_sifted["nonwild"]: + matched_this_ind = False + for e_ind in expr_indices_remaining: + if e_ind in matched_indices: + continue + m = ind.matches(e_ind) + if m is not None: + matched_this_ind = True + repl_dict.update(m) + matched_indices.append(e_ind) + break + if not matched_this_ind: + return None + + expr_indices_remaining = [i for i in expr_indices_remaining if i not in matched_indices] + for ind in indices_sifted["wild"]: + matched_this_ind = False + for e_ind in expr_indices_remaining: + m = ind.matches(e_ind) + if m is not None: + if -ind in repl_dict.keys() and -repl_dict[-ind] != m[ind]: + return None + matched_this_ind = True + repl_dict.update(m) + matched_indices.append(e_ind) + break + if not matched_this_ind: + return None + + expr_indices_remaining = [i for i in expr_indices_remaining if i not in matched_indices] + for ind in indices_sifted["wild, updown"]: + matched_this_ind = False + for e_ind in expr_indices_remaining: + m = ind.matches(e_ind) + if m is not None: + if -ind in repl_dict.keys() and -repl_dict[-ind] != m[ind]: + return None + matched_this_ind = True + repl_dict.update(m) + matched_indices.append(e_ind) + break + if not matched_this_ind: + return None + + if len(matched_indices) < len(self.indices): + return None + else: + return repl_dict + +class WildTensorIndex(TensorIndex): + """ + A wild object that matches TensorIndex instances. + + Examples + ======== + >>> from sympy.tensor.tensor import TensorIndex, TensorIndexType, WildTensorIndex + >>> R3 = TensorIndexType('R3', dim=3) + >>> p = TensorIndex("p", R3) + + By default, covariant indices only match with covariant indices (and + similarly for contravariant) + + >>> q = WildTensorIndex("q", R3) + >>> (q).matches(p) + {q: p} + >>> (q).matches(-p) + + If you want matching to ignore whether the index is co/contra-variant, set + ignore_updown=True + + >>> r = WildTensorIndex("r", R3, ignore_updown=True) + >>> (r).matches(-p) + {r: -p} + >>> (r).matches(p) + {r: p} + + Parameters + ========== + name : name of the index (string), or ``True`` if you want it to be + automatically assigned + tensor_index_type : ``TensorIndexType`` of the index + is_up : flag for contravariant index (is_up=True by default) + ignore_updown : bool, Whether this should match both co- and contra-variant + indices (default:False) + """ + def __new__(cls, name, tensor_index_type, is_up=True, ignore_updown=False): + if isinstance(name, str): + name_symbol = Symbol(name) + elif isinstance(name, Symbol): + name_symbol = name + elif name is True: + name = "_i{}".format(len(tensor_index_type._autogenerated)) + name_symbol = Symbol(name) + tensor_index_type._autogenerated.append(name_symbol) + else: + raise ValueError("invalid name") + + is_up = sympify(is_up) + ignore_updown = sympify(ignore_updown) + return Basic.__new__(cls, name_symbol, tensor_index_type, is_up, ignore_updown) + + @property + def ignore_updown(self): + return self.args[3] + + def __neg__(self): + t1 = WildTensorIndex(self.name, self.tensor_index_type, + (not self.is_up), self.ignore_updown) + return t1 + + def matches(self, expr, repl_dict=None, old=False): + if not isinstance(expr, TensorIndex): + return None + if self.tensor_index_type != expr.tensor_index_type: + return None + if not self.ignore_updown: + if self.is_up != expr.is_up: + return None + + if repl_dict is None: + repl_dict = {} + else: + repl_dict = repl_dict.copy() + + repl_dict[self] = expr + return repl_dict + + +class _WildTensExpr(Basic): + """ + INTERNAL USE ONLY + + This is an object that helps with replacement of WildTensors in expressions. + When this object is set as the tensor_head of a WildTensor, it replaces the + WildTensor by a TensExpr (passed when initializing this object). + + Examples + ======== + >>> from sympy.tensor.tensor import WildTensorHead, TensorIndex, TensorHead, TensorIndexType + >>> W = WildTensorHead("W") + >>> R3 = TensorIndexType('R3', dim=3) + >>> p = TensorIndex('p', R3) + >>> q = TensorIndex('q', R3) + >>> K = TensorHead('K', [R3]) + >>> print( ( K(p) ).replace( W(p), W(q)*W(-q)*W(p) ) ) + K(R_0)*K(-R_0)*K(p) + + """ + def __init__(self, expr): + if not isinstance(expr, TensExpr): + raise TypeError("_WildTensExpr expects a TensExpr as argument") + self.expr = expr + + def __call__(self, *indices): + return self.expr._replace_indices(dict(zip(self.expr.get_free_indices(), indices))) + + def __neg__(self): + return self.func(self.expr*S.NegativeOne) + + def __abs__(self): + raise NotImplementedError + + def __add__(self, other): + if other.func != self.func: + raise TypeError(f"Cannot add {self.func} to {other.func}") + return self.func(self.expr+other.expr) + + def __radd__(self, other): + if other.func != self.func: + raise TypeError(f"Cannot add {self.func} to {other.func}") + return self.func(other.expr+self.expr) + + def __sub__(self, other): + return self + (-other) + + def __rsub__(self, other): + return other + (-self) + + def __mul__(self, other): + raise NotImplementedError + + def __rmul__(self, other): + raise NotImplementedError + + def __truediv__(self, other): + raise NotImplementedError + + def __rtruediv__(self, other): + raise NotImplementedError + + def __pow__(self, other): + raise NotImplementedError + + def __rpow__(self, other): + raise NotImplementedError + + +def canon_bp(p): + """ + Butler-Portugal canonicalization. See ``tensor_can.py`` from the + combinatorics module for the details. + """ + if isinstance(p, TensExpr): + return p.canon_bp() + return p + + +def tensor_mul(*a): + """ + product of tensors + """ + if not a: + return TensMul.from_data(S.One, [], [], []) + t = a[0] + for tx in a[1:]: + t = t*tx + return t + + +def riemann_cyclic_replace(t_r): + """ + replace Riemann tensor with an equivalent expression + + ``R(m,n,p,q) -> 2/3*R(m,n,p,q) - 1/3*R(m,q,n,p) + 1/3*R(m,p,n,q)`` + + """ + free = sorted(t_r.free, key=lambda x: x[1]) + m, n, p, q = [x[0] for x in free] + t0 = t_r*Rational(2, 3) + t1 = -t_r.substitute_indices((m,m),(n,q),(p,n),(q,p))*Rational(1, 3) + t2 = t_r.substitute_indices((m,m),(n,p),(p,n),(q,q))*Rational(1, 3) + t3 = t0 + t1 + t2 + return t3 + +def riemann_cyclic(t2): + """ + Replace each Riemann tensor with an equivalent expression + satisfying the cyclic identity. + + This trick is discussed in the reference guide to Cadabra. + + Examples + ======== + + >>> from sympy.tensor.tensor import TensorIndexType, tensor_indices, TensorHead, riemann_cyclic, TensorSymmetry + >>> Lorentz = TensorIndexType('Lorentz', dummy_name='L') + >>> i, j, k, l = tensor_indices('i,j,k,l', Lorentz) + >>> R = TensorHead('R', [Lorentz]*4, TensorSymmetry.riemann()) + >>> t = R(i,j,k,l)*(R(-i,-j,-k,-l) - 2*R(-i,-k,-j,-l)) + >>> riemann_cyclic(t) + 0 + """ + t2 = t2.expand() + if isinstance(t2, (TensMul, Tensor)): + args = [t2] + else: + args = t2.args + a1 = [x.split() for x in args] + a2 = [[riemann_cyclic_replace(tx) for tx in y] for y in a1] + a3 = [tensor_mul(*v) for v in a2] + t3 = TensAdd(*a3).doit() + if not t3: + return t3 + else: + return canon_bp(t3) + + +def get_lines(ex, index_type): + """ + Returns ``(lines, traces, rest)`` for an index type, + where ``lines`` is the list of list of positions of a matrix line, + ``traces`` is the list of list of traced matrix lines, + ``rest`` is the rest of the elements of the tensor. + """ + def _join_lines(a): + i = 0 + while i < len(a): + x = a[i] + xend = x[-1] + xstart = x[0] + hit = True + while hit: + hit = False + for j in range(i + 1, len(a)): + if j >= len(a): + break + if a[j][0] == xend: + hit = True + x.extend(a[j][1:]) + xend = x[-1] + a.pop(j) + continue + if a[j][0] == xstart: + hit = True + a[i] = reversed(a[j][1:]) + x + x = a[i] + xstart = a[i][0] + a.pop(j) + continue + if a[j][-1] == xend: + hit = True + x.extend(reversed(a[j][:-1])) + xend = x[-1] + a.pop(j) + continue + if a[j][-1] == xstart: + hit = True + a[i] = a[j][:-1] + x + x = a[i] + xstart = x[0] + a.pop(j) + continue + i += 1 + return a + + arguments = ex.args + dt = {} + for c in ex.args: + if not isinstance(c, TensExpr): + continue + if c in dt: + continue + index_types = c.index_types + a = [] + for i in range(len(index_types)): + if index_types[i] is index_type: + a.append(i) + if len(a) > 2: + raise ValueError('at most two indices of type %s allowed' % index_type) + if len(a) == 2: + dt[c] = a + #dum = ex.dum + lines = [] + traces = [] + traces1 = [] + #indices_to_args_pos = ex._get_indices_to_args_pos() + # TODO: add a dum_to_components_map ? + for p0, p1, c0, c1 in ex.dum_in_args: + if arguments[c0] not in dt: + continue + if c0 == c1: + traces.append([c0]) + continue + ta0 = dt[arguments[c0]] + ta1 = dt[arguments[c1]] + if p0 not in ta0: + continue + if ta0.index(p0) == ta1.index(p1): + # case gamma(i,s0,-s1) in c0, gamma(j,-s0,s2) in c1; + # to deal with this case one could add to the position + # a flag for transposition; + # one could write [(c0, False), (c1, True)] + raise NotImplementedError + # if p0 == ta0[1] then G in pos c0 is mult on the right by G in c1 + # if p0 == ta0[0] then G in pos c1 is mult on the right by G in c0 + ta0 = dt[arguments[c0]] + b0, b1 = (c0, c1) if p0 == ta0[1] else (c1, c0) + lines1 = lines[:] + for line in lines: + if line[-1] == b0: + if line[0] == b1: + n = line.index(min(line)) + traces1.append(line) + traces.append(line[n:] + line[:n]) + else: + line.append(b1) + break + elif line[0] == b1: + line.insert(0, b0) + break + else: + lines1.append([b0, b1]) + + lines = [x for x in lines1 if x not in traces1] + lines = _join_lines(lines) + rest = [] + for line in lines: + for y in line: + rest.append(y) + for line in traces: + for y in line: + rest.append(y) + rest = [x for x in range(len(arguments)) if x not in rest] + + return lines, traces, rest + + +def get_free_indices(t): + if not isinstance(t, TensExpr): + return () + return t.get_free_indices() + + +def get_indices(t): + if not isinstance(t, TensExpr): + return () + return t.get_indices() + +def get_dummy_indices(t): + if not isinstance(t, TensExpr): + return () + inds = t.get_indices() + free = t.get_free_indices() + return [i for i in inds if i not in free] + +def get_index_structure(t): + if isinstance(t, TensExpr): + return t._index_structure + return _IndexStructure([], [], [], []) + + +def get_coeff(t): + if isinstance(t, Tensor): + return S.One + if isinstance(t, TensMul): + return t.coeff + if isinstance(t, TensExpr): + raise ValueError("no coefficient associated to this tensor expression") + return t + +def contract_metric(t, g): + if isinstance(t, TensExpr): + return t.contract_metric(g) + return t + + +def perm2tensor(t, g, is_canon_bp=False): + """ + Returns the tensor corresponding to the permutation ``g`` + + For further details, see the method in ``TIDS`` with the same name. + """ + if not isinstance(t, TensExpr): + return t + elif isinstance(t, (Tensor, TensMul)): + nim = get_index_structure(t).perm2tensor(g, is_canon_bp=is_canon_bp) + res = t._set_new_index_structure(nim, is_canon_bp=is_canon_bp) + if g[-1] != len(g) - 1: + return -res + + return res + raise NotImplementedError() + + +def substitute_indices(t, *index_tuples): + if not isinstance(t, TensExpr): + return t + return t.substitute_indices(*index_tuples) + + +def _expand(expr, **kwargs): + if isinstance(expr, TensExpr): + return expr._expand(**kwargs) + else: + return expr.expand(**kwargs) + + +def get_postprocessor(cls): + def _postprocessor(expr): + tens_class = {Mul: TensMul, Add: TensAdd}[cls] + if any(isinstance(a, TensExpr) for a in expr.args): + return tens_class(*expr.args) + else: + return expr + + return _postprocessor + +Basic._constructor_postprocessor_mapping[TensExpr] = { + "Mul": [get_postprocessor(Mul)], +} diff --git a/.venv/lib/python3.11/site-packages/sympy/tensor/toperators.py b/.venv/lib/python3.11/site-packages/sympy/tensor/toperators.py new file mode 100644 index 0000000000000000000000000000000000000000..1bdd67c4f4a7e86b9821ee55b1d2f9bde29c96a8 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/sympy/tensor/toperators.py @@ -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 diff --git a/.venv/lib/python3.11/site-packages/transformers/utils/__pycache__/dummy_pt_objects.cpython-311.pyc b/.venv/lib/python3.11/site-packages/transformers/utils/__pycache__/dummy_pt_objects.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..31f3dc9700b79ddcb40c46172575be2a7743bab4 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/transformers/utils/__pycache__/dummy_pt_objects.cpython-311.pyc @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f2cc98d5da92926a5126a38d884acf481a1b0d37481dd038a611ed9ae01a7cd2 +size 551815 diff --git a/.venv/lib/python3.11/site-packages/transformers/utils/__pycache__/dummy_tf_objects.cpython-311.pyc b/.venv/lib/python3.11/site-packages/transformers/utils/__pycache__/dummy_tf_objects.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fd11ef7cbbec9f035918fef169005e82d2cafc2f --- /dev/null +++ b/.venv/lib/python3.11/site-packages/transformers/utils/__pycache__/dummy_tf_objects.cpython-311.pyc @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:454eca0262d400d4bbea810a0da1a8de2343fb73941c730246ef89d9c59a39a6 +size 149669