Coverage for /usr/lib/python3/dist-packages/sympy/tensor/array/dense_ndim_array.py: 37%

1import functools 

2from typing import List 

3 

4from sympy.core.basic import Basic 

5from sympy.core.containers import Tuple 

6from sympy.core.singleton import S 

7from sympy.core.sympify import _sympify 

8from sympy.tensor.array.mutable_ndim_array import MutableNDimArray 

9from sympy.tensor.array.ndim_array import NDimArray, ImmutableNDimArray, ArrayKind 

10from sympy.utilities.iterables import flatten 

11 

12 

13class DenseNDimArray(NDimArray): 

14 

15 _array: List[Basic] 

16 

17 def __new__(self, *args, **kwargs): 

18 return ImmutableDenseNDimArray(*args, **kwargs) 

19 

20 @property 

21 def kind(self) -> ArrayKind: 

22 return ArrayKind._union(self._array) 

23 

24 def __getitem__(self, index): 

25 """ 

26 Allows to get items from N-dim array. 

27 

28 Examples 

29 ======== 

30 

31 >>> from sympy import MutableDenseNDimArray 

32 >>> a = MutableDenseNDimArray([0, 1, 2, 3], (2, 2)) 

33 >>> a 

34 [[0, 1], [2, 3]] 

35 >>> a[0, 0] 

36 0 

37 >>> a[1, 1] 

38 3 

39 >>> a[0] 

40 [0, 1] 

41 >>> a[1] 

42 [2, 3] 

43 

44 

45 Symbolic index: 

46 

47 >>> from sympy.abc import i, j 

48 >>> a[i, j] 

49 [[0, 1], [2, 3]][i, j] 

50 

51 Replace `i` and `j` to get element `(1, 1)`: 

52 

53 >>> a[i, j].subs({i: 1, j: 1}) 

54 3 

55 

56 """ 

57 syindex = self._check_symbolic_index(index) 

58 if syindex is not None: 

59 return syindex 

60 

61 index = self._check_index_for_getitem(index) 

62 

63 if isinstance(index, tuple) and any(isinstance(i, slice) for i in index): 

64 sl_factors, eindices = self._get_slice_data_for_array_access(index) 

65 array = [self._array[self._parse_index(i)] for i in eindices] 

66 nshape = [len(el) for i, el in enumerate(sl_factors) if isinstance(index[i], slice)] 

67 return type(self)(array, nshape) 

68 else: 

69 index = self._parse_index(index) 

70 return self._array[index] 

71 

72 @classmethod 

73 def zeros(cls, *shape): 

74 list_length = functools.reduce(lambda x, y: x*y, shape, S.One) 

75 return cls._new(([0]*list_length,), shape) 

76 

77 def tomatrix(self): 

78 """ 

79 Converts MutableDenseNDimArray to Matrix. Can convert only 2-dim array, else will raise error. 

80 

81 Examples 

82 ======== 

83 

84 >>> from sympy import MutableDenseNDimArray 

85 >>> a = MutableDenseNDimArray([1 for i in range(9)], (3, 3)) 

86 >>> b = a.tomatrix() 

87 >>> b 

88 Matrix([ 

89 [1, 1, 1], 

90 [1, 1, 1], 

91 [1, 1, 1]]) 

92 

93 """ 

94 from sympy.matrices import Matrix 

95 

96 if self.rank() != 2: 

97 raise ValueError('Dimensions must be of size of 2') 

98 

99 return Matrix(self.shape[0], self.shape[1], self._array) 

100 

101 def reshape(self, *newshape): 

102 """ 

103 Returns MutableDenseNDimArray instance with new shape. Elements number 

104 must be suitable to new shape. The only argument of method sets 

105 new shape. 

106 

107 Examples 

108 ======== 

109 

110 >>> from sympy import MutableDenseNDimArray 

111 >>> a = MutableDenseNDimArray([1, 2, 3, 4, 5, 6], (2, 3)) 

112 >>> a.shape 

113 (2, 3) 

114 >>> a 

115 [[1, 2, 3], [4, 5, 6]] 

116 >>> b = a.reshape(3, 2) 

117 >>> b.shape 

118 (3, 2) 

119 >>> b 

120 [[1, 2], [3, 4], [5, 6]] 

121 

122 """ 

123 new_total_size = functools.reduce(lambda x,y: x*y, newshape) 

124 if new_total_size != self._loop_size: 

125 raise ValueError('Expecting reshape size to %d but got prod(%s) = %d' % ( 

126 self._loop_size, str(newshape), new_total_size)) 

127 

128 # there is no `.func` as this class does not subtype `Basic`: 

129 return type(self)(self._array, newshape) 

130 

131 

132class ImmutableDenseNDimArray(DenseNDimArray, ImmutableNDimArray): # type: ignore 

133 def __new__(cls, iterable, shape=None, **kwargs): 

134 return cls._new(iterable, shape, **kwargs) 

135 

136 @classmethod 

137 def _new(cls, iterable, shape, **kwargs): 

138 shape, flat_list = cls._handle_ndarray_creation_inputs(iterable, shape, **kwargs) 

139 shape = Tuple(*map(_sympify, shape)) 

140 cls._check_special_bounds(flat_list, shape) 

141 flat_list = flatten(flat_list) 

142 flat_list = Tuple(*flat_list) 

143 self = Basic.__new__(cls, flat_list, shape, **kwargs) 

144 self._shape = shape 

145 self._array = list(flat_list) 

146 self._rank = len(shape) 

147 self._loop_size = functools.reduce(lambda x,y: x*y, shape, 1) 

148 return self 

149 

150 def __setitem__(self, index, value): 

151 raise TypeError('immutable N-dim array') 

152 

153 def as_mutable(self): 

154 return MutableDenseNDimArray(self) 

155 

156 def _eval_simplify(self, **kwargs): 

157 from sympy.simplify.simplify import simplify 

158 return self.applyfunc(simplify) 

159 

160class MutableDenseNDimArray(DenseNDimArray, MutableNDimArray): 

161 

162 def __new__(cls, iterable=None, shape=None, **kwargs): 

163 return cls._new(iterable, shape, **kwargs) 

164 

165 @classmethod 

166 def _new(cls, iterable, shape, **kwargs): 

167 shape, flat_list = cls._handle_ndarray_creation_inputs(iterable, shape, **kwargs) 

168 flat_list = flatten(flat_list) 

169 self = object.__new__(cls) 

170 self._shape = shape 

171 self._array = list(flat_list) 

172 self._rank = len(shape) 

173 self._loop_size = functools.reduce(lambda x,y: x*y, shape) if shape else len(flat_list) 

174 return self 

175 

176 def __setitem__(self, index, value): 

177 """Allows to set items to MutableDenseNDimArray. 

178 

179 Examples 

180 ======== 

181 

182 >>> from sympy import MutableDenseNDimArray 

183 >>> a = MutableDenseNDimArray.zeros(2, 2) 

184 >>> a[0,0] = 1 

185 >>> a[1,1] = 1 

186 >>> a 

187 [[1, 0], [0, 1]] 

188 

189 """ 

190 if isinstance(index, tuple) and any(isinstance(i, slice) for i in index): 

191 value, eindices, slice_offsets = self._get_slice_data_for_array_assignment(index, value) 

192 for i in eindices: 

193 other_i = [ind - j for ind, j in zip(i, slice_offsets) if j is not None] 

194 self._array[self._parse_index(i)] = value[other_i] 

195 else: 

196 index = self._parse_index(index) 

197 self._setter_iterable_check(value) 

198 value = _sympify(value) 

199 self._array[index] = value 

200 

201 def as_immutable(self): 

202 return ImmutableDenseNDimArray(self) 

203 

204 @property 

205 def free_symbols(self): 

206 return {i for j in self._array for i in j.free_symbols}