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

1from sympy.core.basic import Basic 

2from sympy.core.containers import (Dict, Tuple) 

3from sympy.core.singleton import S 

4from sympy.core.sympify import _sympify 

5from sympy.tensor.array.mutable_ndim_array import MutableNDimArray 

6from sympy.tensor.array.ndim_array import NDimArray, ImmutableNDimArray 

7from sympy.utilities.iterables import flatten 

8 

9import functools 

10 

11class SparseNDimArray(NDimArray): 

12 

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

14 return ImmutableSparseNDimArray(*args, **kwargs) 

15 

16 def __getitem__(self, index): 

17 """ 

18 Get an element from a sparse N-dim array. 

19 

20 Examples 

21 ======== 

22 

23 >>> from sympy import MutableSparseNDimArray 

24 >>> a = MutableSparseNDimArray(range(4), (2, 2)) 

25 >>> a 

26 [[0, 1], [2, 3]] 

27 >>> a[0, 0] 

28 0 

29 >>> a[1, 1] 

30 3 

31 >>> a[0] 

32 [0, 1] 

33 >>> a[1] 

34 [2, 3] 

35 

36 Symbolic indexing: 

37 

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

39 >>> a[i, j] 

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

41 

42 Replace `i` and `j` to get element `(0, 0)`: 

43 

44 >>> a[i, j].subs({i: 0, j: 0}) 

45 0 

46 

47 """ 

48 syindex = self._check_symbolic_index(index) 

49 if syindex is not None: 

50 return syindex 

51 

52 index = self._check_index_for_getitem(index) 

53 

54 # `index` is a tuple with one or more slices: 

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

56 sl_factors, eindices = self._get_slice_data_for_array_access(index) 

57 array = [self._sparse_array.get(self._parse_index(i), S.Zero) for i in eindices] 

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

59 return type(self)(array, nshape) 

60 else: 

61 index = self._parse_index(index) 

62 return self._sparse_array.get(index, S.Zero) 

63 

64 @classmethod 

65 def zeros(cls, *shape): 

66 """ 

67 Return a sparse N-dim array of zeros. 

68 """ 

69 return cls({}, shape) 

70 

71 def tomatrix(self): 

72 """ 

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

74 

75 Examples 

76 ======== 

77 

78 >>> from sympy import MutableSparseNDimArray 

79 >>> a = MutableSparseNDimArray([1 for i in range(9)], (3, 3)) 

80 >>> b = a.tomatrix() 

81 >>> b 

82 Matrix([ 

83 [1, 1, 1], 

84 [1, 1, 1], 

85 [1, 1, 1]]) 

86 """ 

87 from sympy.matrices import SparseMatrix 

88 if self.rank() != 2: 

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

90 

91 mat_sparse = {} 

92 for key, value in self._sparse_array.items(): 

93 mat_sparse[self._get_tuple_index(key)] = value 

94 

95 return SparseMatrix(self.shape[0], self.shape[1], mat_sparse) 

96 

97 def reshape(self, *newshape): 

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

99 if new_total_size != self._loop_size: 

100 raise ValueError("Invalid reshape parameters " + newshape) 

101 

102 return type(self)(self._sparse_array, newshape) 

103 

104class ImmutableSparseNDimArray(SparseNDimArray, ImmutableNDimArray): # type: ignore 

105 

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

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

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

109 cls._check_special_bounds(flat_list, shape) 

110 loop_size = functools.reduce(lambda x,y: x*y, shape) if shape else len(flat_list) 

111 

112 # Sparse array: 

113 if isinstance(flat_list, (dict, Dict)): 

114 sparse_array = Dict(flat_list) 

115 else: 

116 sparse_array = {} 

117 for i, el in enumerate(flatten(flat_list)): 

118 if el != 0: 

119 sparse_array[i] = _sympify(el) 

120 

121 sparse_array = Dict(sparse_array) 

122 

123 self = Basic.__new__(cls, sparse_array, shape, **kwargs) 

124 self._shape = shape 

125 self._rank = len(shape) 

126 self._loop_size = loop_size 

127 self._sparse_array = sparse_array 

128 

129 return self 

130 

131 def __setitem__(self, index, value): 

132 raise TypeError("immutable N-dim array") 

133 

134 def as_mutable(self): 

135 return MutableSparseNDimArray(self) 

136 

137 

138class MutableSparseNDimArray(MutableNDimArray, SparseNDimArray): 

139 

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

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

142 self = object.__new__(cls) 

143 self._shape = shape 

144 self._rank = len(shape) 

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

146 

147 # Sparse array: 

148 if isinstance(flat_list, (dict, Dict)): 

149 self._sparse_array = dict(flat_list) 

150 return self 

151 

152 self._sparse_array = {} 

153 

154 for i, el in enumerate(flatten(flat_list)): 

155 if el != 0: 

156 self._sparse_array[i] = _sympify(el) 

157 

158 return self 

159 

160 def __setitem__(self, index, value): 

161 """Allows to set items to MutableDenseNDimArray. 

162 

163 Examples 

164 ======== 

165 

166 >>> from sympy import MutableSparseNDimArray 

167 >>> a = MutableSparseNDimArray.zeros(2, 2) 

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

169 >>> a[1, 1] = 1 

170 >>> a 

171 [[1, 0], [0, 1]] 

172 """ 

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

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

175 for i in eindices: 

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

177 other_value = value[other_i] 

178 complete_index = self._parse_index(i) 

179 if other_value != 0: 

180 self._sparse_array[complete_index] = other_value 

181 elif complete_index in self._sparse_array: 

182 self._sparse_array.pop(complete_index) 

183 else: 

184 index = self._parse_index(index) 

185 value = _sympify(value) 

186 if value == 0 and index in self._sparse_array: 

187 self._sparse_array.pop(index) 

188 else: 

189 self._sparse_array[index] = value 

190 

191 def as_immutable(self): 

192 return ImmutableSparseNDimArray(self) 

193 

194 @property 

195 def free_symbols(self): 

196 return {i for j in self._sparse_array.values() for i in j.free_symbols}