Coverage for /usr/lib/python3/dist-packages/sympy/matrices/sparsetools.py: 7%
109 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1from sympy.core.containers import Dict
2from sympy.core.symbol import Dummy
3from sympy.utilities.iterables import is_sequence
4from sympy.utilities.misc import as_int, filldedent
6from .sparse import MutableSparseMatrix as SparseMatrix
9def _doktocsr(dok):
10 """Converts a sparse matrix to Compressed Sparse Row (CSR) format.
12 Parameters
13 ==========
15 A : contains non-zero elements sorted by key (row, column)
16 JA : JA[i] is the column corresponding to A[i]
17 IA : IA[i] contains the index in A for the first non-zero element
18 of row[i]. Thus IA[i+1] - IA[i] gives number of non-zero
19 elements row[i]. The length of IA is always 1 more than the
20 number of rows in the matrix.
22 Examples
23 ========
25 >>> from sympy.matrices.sparsetools import _doktocsr
26 >>> from sympy import SparseMatrix, diag
27 >>> m = SparseMatrix(diag(1, 2, 3))
28 >>> m[2, 0] = -1
29 >>> _doktocsr(m)
30 [[1, 2, -1, 3], [0, 1, 0, 2], [0, 1, 2, 4], [3, 3]]
32 """
33 row, JA, A = [list(i) for i in zip(*dok.row_list())]
34 IA = [0]*((row[0] if row else 0) + 1)
35 for i, r in enumerate(row):
36 IA.extend([i]*(r - row[i - 1])) # if i = 0 nothing is extended
37 IA.extend([len(A)]*(dok.rows - len(IA) + 1))
38 shape = [dok.rows, dok.cols]
39 return [A, JA, IA, shape]
42def _csrtodok(csr):
43 """Converts a CSR representation to DOK representation.
45 Examples
46 ========
48 >>> from sympy.matrices.sparsetools import _csrtodok
49 >>> _csrtodok([[5, 8, 3, 6], [0, 1, 2, 1], [0, 0, 2, 3, 4], [4, 3]])
50 Matrix([
51 [0, 0, 0],
52 [5, 8, 0],
53 [0, 0, 3],
54 [0, 6, 0]])
56 """
57 smat = {}
58 A, JA, IA, shape = csr
59 for i in range(len(IA) - 1):
60 indices = slice(IA[i], IA[i + 1])
61 for l, m in zip(A[indices], JA[indices]):
62 smat[i, m] = l
63 return SparseMatrix(*shape, smat)
66def banded(*args, **kwargs):
67 """Returns a SparseMatrix from the given dictionary describing
68 the diagonals of the matrix. The keys are positive for upper
69 diagonals and negative for those below the main diagonal. The
70 values may be:
72 * expressions or single-argument functions,
74 * lists or tuples of values,
76 * matrices
78 Unless dimensions are given, the size of the returned matrix will
79 be large enough to contain the largest non-zero value provided.
81 kwargs
82 ======
84 rows : rows of the resulting matrix; computed if
85 not given.
87 cols : columns of the resulting matrix; computed if
88 not given.
90 Examples
91 ========
93 >>> from sympy import banded, ones, Matrix
94 >>> from sympy.abc import x
96 If explicit values are given in tuples,
97 the matrix will autosize to contain all values, otherwise
98 a single value is filled onto the entire diagonal:
100 >>> banded({1: (1, 2, 3), -1: (4, 5, 6), 0: x})
101 Matrix([
102 [x, 1, 0, 0],
103 [4, x, 2, 0],
104 [0, 5, x, 3],
105 [0, 0, 6, x]])
107 A function accepting a single argument can be used to fill the
108 diagonal as a function of diagonal index (which starts at 0).
109 The size (or shape) of the matrix must be given to obtain more
110 than a 1x1 matrix:
112 >>> s = lambda d: (1 + d)**2
113 >>> banded(5, {0: s, 2: s, -2: 2})
114 Matrix([
115 [1, 0, 1, 0, 0],
116 [0, 4, 0, 4, 0],
117 [2, 0, 9, 0, 9],
118 [0, 2, 0, 16, 0],
119 [0, 0, 2, 0, 25]])
121 The diagonal of matrices placed on a diagonal will coincide
122 with the indicated diagonal:
124 >>> vert = Matrix([1, 2, 3])
125 >>> banded({0: vert}, cols=3)
126 Matrix([
127 [1, 0, 0],
128 [2, 1, 0],
129 [3, 2, 1],
130 [0, 3, 2],
131 [0, 0, 3]])
133 >>> banded(4, {0: ones(2)})
134 Matrix([
135 [1, 1, 0, 0],
136 [1, 1, 0, 0],
137 [0, 0, 1, 1],
138 [0, 0, 1, 1]])
140 Errors are raised if the designated size will not hold
141 all values an integral number of times. Here, the rows
142 are designated as odd (but an even number is required to
143 hold the off-diagonal 2x2 ones):
145 >>> banded({0: 2, 1: ones(2)}, rows=5)
146 Traceback (most recent call last):
147 ...
148 ValueError:
149 sequence does not fit an integral number of times in the matrix
151 And here, an even number of rows is given...but the square
152 matrix has an even number of columns, too. As we saw
153 in the previous example, an odd number is required:
155 >>> banded(4, {0: 2, 1: ones(2)}) # trying to make 4x4 and cols must be odd
156 Traceback (most recent call last):
157 ...
158 ValueError:
159 sequence does not fit an integral number of times in the matrix
161 A way around having to count rows is to enclosing matrix elements
162 in a tuple and indicate the desired number of them to the right:
164 >>> banded({0: 2, 2: (ones(2),)*3})
165 Matrix([
166 [2, 0, 1, 1, 0, 0, 0, 0],
167 [0, 2, 1, 1, 0, 0, 0, 0],
168 [0, 0, 2, 0, 1, 1, 0, 0],
169 [0, 0, 0, 2, 1, 1, 0, 0],
170 [0, 0, 0, 0, 2, 0, 1, 1],
171 [0, 0, 0, 0, 0, 2, 1, 1]])
173 An error will be raised if more than one value
174 is written to a given entry. Here, the ones overlap
175 with the main diagonal if they are placed on the
176 first diagonal:
178 >>> banded({0: (2,)*5, 1: (ones(2),)*3})
179 Traceback (most recent call last):
180 ...
181 ValueError: collision at (1, 1)
183 By placing a 0 at the bottom left of the 2x2 matrix of
184 ones, the collision is avoided:
186 >>> u2 = Matrix([
187 ... [1, 1],
188 ... [0, 1]])
189 >>> banded({0: [2]*5, 1: [u2]*3})
190 Matrix([
191 [2, 1, 1, 0, 0, 0, 0],
192 [0, 2, 1, 0, 0, 0, 0],
193 [0, 0, 2, 1, 1, 0, 0],
194 [0, 0, 0, 2, 1, 0, 0],
195 [0, 0, 0, 0, 2, 1, 1],
196 [0, 0, 0, 0, 0, 0, 1]])
197 """
198 try:
199 if len(args) not in (1, 2, 3):
200 raise TypeError
201 if not isinstance(args[-1], (dict, Dict)):
202 raise TypeError
203 if len(args) == 1:
204 rows = kwargs.get('rows', None)
205 cols = kwargs.get('cols', None)
206 if rows is not None:
207 rows = as_int(rows)
208 if cols is not None:
209 cols = as_int(cols)
210 elif len(args) == 2:
211 rows = cols = as_int(args[0])
212 else:
213 rows, cols = map(as_int, args[:2])
214 # fails with ValueError if any keys are not ints
215 _ = all(as_int(k) for k in args[-1])
216 except (ValueError, TypeError):
217 raise TypeError(filldedent(
218 '''unrecognized input to banded:
219 expecting [[row,] col,] {int: value}'''))
220 def rc(d):
221 # return row,col coord of diagonal start
222 r = -d if d < 0 else 0
223 c = 0 if r else d
224 return r, c
225 smat = {}
226 undone = []
227 tba = Dummy()
228 # first handle objects with size
229 for d, v in args[-1].items():
230 r, c = rc(d)
231 # note: only list and tuple are recognized since this
232 # will allow other Basic objects like Tuple
233 # into the matrix if so desired
234 if isinstance(v, (list, tuple)):
235 extra = 0
236 for i, vi in enumerate(v):
237 i += extra
238 if is_sequence(vi):
239 vi = SparseMatrix(vi)
240 smat[r + i, c + i] = vi
241 extra += min(vi.shape) - 1
242 else:
243 smat[r + i, c + i] = vi
244 elif is_sequence(v):
245 v = SparseMatrix(v)
246 rv, cv = v.shape
247 if rows and cols:
248 nr, xr = divmod(rows - r, rv)
249 nc, xc = divmod(cols - c, cv)
250 x = xr or xc
251 do = min(nr, nc)
252 elif rows:
253 do, x = divmod(rows - r, rv)
254 elif cols:
255 do, x = divmod(cols - c, cv)
256 else:
257 do = 1
258 x = 0
259 if x:
260 raise ValueError(filldedent('''
261 sequence does not fit an integral number of times
262 in the matrix'''))
263 j = min(v.shape)
264 for i in range(do):
265 smat[r, c] = v
266 r += j
267 c += j
268 elif v:
269 smat[r, c] = tba
270 undone.append((d, v))
271 s = SparseMatrix(None, smat) # to expand matrices
272 smat = s.todok()
273 # check for dim errors here
274 if rows is not None and rows < s.rows:
275 raise ValueError('Designated rows %s < needed %s' % (rows, s.rows))
276 if cols is not None and cols < s.cols:
277 raise ValueError('Designated cols %s < needed %s' % (cols, s.cols))
278 if rows is cols is None:
279 rows = s.rows
280 cols = s.cols
281 elif rows is not None and cols is None:
282 cols = max(rows, s.cols)
283 elif cols is not None and rows is None:
284 rows = max(cols, s.rows)
285 def update(i, j, v):
286 # update smat and make sure there are
287 # no collisions
288 if v:
289 if (i, j) in smat and smat[i, j] not in (tba, v):
290 raise ValueError('collision at %s' % ((i, j),))
291 smat[i, j] = v
292 if undone:
293 for d, vi in undone:
294 r, c = rc(d)
295 v = vi if callable(vi) else lambda _: vi
296 i = 0
297 while r + i < rows and c + i < cols:
298 update(r + i, c + i, v(i))
299 i += 1
300 return SparseMatrix(rows, cols, smat)