def __add__(self, other): from sympy.tensor.array.arrayop import Flatten if not isinstance(other, NDimArray): raise TypeError(str(other)) 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): if not isinstance(other, NewArray): result_list = [i - other for i in Flatten(self)] return type(self)(result_list, self.shape) 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 __rdiv__(self, other): if isinstance(other, (np.ndarray, NewArray)): if self.shape != other.shape: raise ValueError("array shape mismatch") else: result_list = [ i / j for i, j in zip(Flatten(self), Flatten(other)) ] return type(self)(result_list, self.shape) else: result_list = [other / i for i in Flatten(self)] return type(self)(result_list, self.shape)
def test_flatten(): from sympy import Matrix for ArrayType in [ ImmutableDenseNDimArray, ImmutableSparseNDimArray, Matrix ]: A = ArrayType(range(24)).reshape(4, 6) assert [i for i in 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)): i == v
def __rsub__(self, other): if not isinstance(other, NewArray): result_list = [other - i for i in Flatten(self)] return type(self)(result_list, self.shape) else: raise NotImplementedError("array shape mismatch")
def my_cos(x): if isinstance(x, (np.ndarray, NewArray)): result_list = [sympy.cos(i) for i in Flatten(x)] return NewArray(result_list, x.shape) else: return sympy.cos(x)
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 __div__(self, other): from sympy.matrices.matrices 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 my_der(self, other): if isinstance(self, (np.ndarray, NewArray)): if isinstance(other, (np.ndarray, NewArray)): if self.shape != other.shape: raise ValueError("array shape mismatch") else: result_list = [ sympy.diff(i, j, evaluate=False) for i, j in zip(Flatten(self), Flatten(other)) ] return type(self)(result_list, self.shape) else: result_list = [ sympy.diff(i, other, evaluate=False) for i in Flatten(self) ] return type(self)(result_list, self.shape) else: if isinstance(other, (np.ndarray, NewArray)): result_list = [ sympy.diff(self, i, evaluate=False) for i in Flatten(other) ] return type(self)(result_list, other.shape) else: return sympy.diff(self, other, evaluate=False)
def __rmul__(self, other): from sympy.matrices.matrices 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 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 _eval_conjugate(self): from sympy.tensor.array.arrayop import Flatten return self.func([i.conjugate() for i in Flatten(self)], self.shape)
def __abs__(self): result_list = [sympy.Abs(i) for i in Flatten(self)] return type(self)(result_list, self.shape)
def _eval_power(self, other): if isinstance(other, (numbers.Real, sympy.Rational, sympy.Float)): result_list = [i**other for i in Flatten(self)] return type(self)(result_list, self.shape) else: raise ValueError("array shape mismatch")