MCPcopy Create free account
hub / github.com/DedalusProject/dedalus / apply_sparse

Function apply_sparse

dedalus/tools/array.py:171–203  ·  view source on GitHub ↗

Apply sparse matrix along any axis of an array. Must be out of place if ouptut is specified.

(matrix, array, axis, out=None, check_shapes=False, num_threads=1)

Source from the content-addressed store, hash-verified

169
170
171def apply_sparse(matrix, array, axis, out=None, check_shapes=False, num_threads=1):
172 """
173 Apply sparse matrix along any axis of an array.
174 Must be out of place if ouptut is specified.
175 """
176 # Check matrix
177 if not isinstance(matrix, sparse.csr_matrix):
178 raise ValueError("Matrix must be in CSR format.")
179 # Check output
180 if out is None:
181 out_shape = list(array.shape)
182 out_shape[axis] = matrix.shape[0]
183 out = np.empty(out_shape, dtype=array.dtype)
184 elif out is array:
185 raise ValueError("Cannot apply in place")
186 # Check shapes
187 if check_shapes:
188 if not (0 <= axis < array.ndim):
189 raise ValueError("Axis out of bounds.")
190 if matrix.shape[1] != array.shape[axis] or matrix.shape[0] != out.shape[axis]:
191 raise ValueError("Matrix shape mismatch.")
192 # Old way if requested
193 if OLD_CSR_MATVECS and array.ndim == 2 and axis == 0:
194 out.fill(0)
195 return csr_matvecs(matrix, array, out)
196 # Promote datatypes
197 # TODO: find way to optimize this with fused types
198 matrix_data = matrix.data
199 if matrix_data.dtype != out.dtype:
200 matrix_data = matrix_data.astype(out.dtype)
201 # Call cython routine
202 cython_linalg.apply_csr(matrix.indptr, matrix.indices, matrix_data, array, out, axis, num_threads)
203 return out
204
205
206def solve_upper_sparse(matrix, rhs, axis, out=None, check_shapes=False, num_threads=1):

Callers 8

apply_matrixFunction · 0.85
stepMethod · 0.85
stepMethod · 0.85
gather_inputsMethod · 0.85
gather_outputsMethod · 0.85
scatter_inputsMethod · 0.85
scatter_outputsMethod · 0.85

Calls 2

csr_matvecsFunction · 0.85
emptyMethod · 0.80

Tested by

no test coverage detected