def test_upcast(): a0 = csr_matrix([[np.pi, np.pi * 1j], [3, 4]], dtype=complex) b0 = np.array([256 + 1j, 2**32], dtype=complex) for a_dtype in supported_dtypes: for b_dtype in supported_dtypes: msg = "(%r, %r)" % (a_dtype, b_dtype) if np.issubdtype(a_dtype, np.complexfloating): a = a0.copy().astype(a_dtype) else: a = a0.real.copy().astype(a_dtype) if np.issubdtype(b_dtype, np.complexfloating): b = b0.copy().astype(b_dtype) else: b = b0.real.copy().astype(b_dtype) if not (a_dtype == np.bool_ and b_dtype == np.bool_): c = np.zeros((2, ), dtype=np.bool_) assert_raises(ValueError, _sparsetools.csr_matvec, 2, 2, a.indptr, a.indices, a.data, b, c) if ((np.issubdtype(a_dtype, np.complexfloating) and not np.issubdtype(b_dtype, np.complexfloating)) or (not np.issubdtype(a_dtype, np.complexfloating) and np.issubdtype(b_dtype, np.complexfloating))): c = np.zeros((2, ), dtype=np.float64) assert_raises(ValueError, _sparsetools.csr_matvec, 2, 2, a.indptr, a.indices, a.data, b, c) c = np.zeros((2, ), dtype=np.result_type(a_dtype, b_dtype)) _sparsetools.csr_matvec(2, 2, a.indptr, a.indices, a.data, b, c) assert_allclose(c, np.dot(a.toarray(), b), err_msg=msg)
def test_upcast(): a0 = csr_matrix([[np.pi, np.pi*1j], [3, 4]], dtype=complex) b0 = np.array([256+1j, 2**32], dtype=complex) for a_dtype in supported_dtypes: for b_dtype in supported_dtypes: msg = "(%r, %r)" % (a_dtype, b_dtype) if np.issubdtype(a_dtype, np.complexfloating): a = a0.copy().astype(a_dtype) else: a = a0.real.copy().astype(a_dtype) if np.issubdtype(b_dtype, np.complexfloating): b = b0.copy().astype(b_dtype) else: b = b0.real.copy().astype(b_dtype) if not (a_dtype == np.bool_ and b_dtype == np.bool_): c = np.zeros((2,), dtype=np.bool_) assert_raises(ValueError, _sparsetools.csr_matvec, 2, 2, a.indptr, a.indices, a.data, b, c) if ((np.issubdtype(a_dtype, np.complexfloating) and not np.issubdtype(b_dtype, np.complexfloating)) or (not np.issubdtype(a_dtype, np.complexfloating) and np.issubdtype(b_dtype, np.complexfloating))): c = np.zeros((2,), dtype=np.float64) assert_raises(ValueError, _sparsetools.csr_matvec, 2, 2, a.indptr, a.indices, a.data, b, c) c = np.zeros((2,), dtype=np.result_type(a_dtype, b_dtype)) _sparsetools.csr_matvec(2, 2, a.indptr, a.indices, a.data, b, c) assert_allclose(c, np.dot(a.toarray(), b), err_msg=msg)
def fast_csr_matvec(A_csr, x_vec, out_vec): """ Fast CSR matvec skipping type and shape checks. The result is added to the specificed output array, so the output should be manually zeroed prior to calling this routine, if necessary. """ # Check format for don't convert if A_csr.format != "csr": raise ValueError("Matrix must be in CSR format.") M, N = A_csr._shape _sparsetools.csr_matvec(M, N, A_csr.indptr, A_csr.indices, A_csr.data, x_vec, out_vec) return out_vec
def csr_matvec(A_csr, x_vec, out_vec): """ Fast CSR matvec with dense vector skipping output allocation. The result is added to the specificed output array, so the output should be manually zeroed prior to calling this routine, if necessary. """ # Check format but don't convert if A_csr.format != "csr": raise ValueError("Matrix must be in CSR format.") # Check shapes M, N = A_csr.shape m, n = out_vec.size, x_vec.size if x_vec.ndim > 1 or out_vec.ndim > 1: raise ValueError("Only vectors allowed for input and output.") if M != m or N != n: raise ValueError( f"Matrix shape {(M,N)} does not match input {(m,)} and output {(n,)} shapes." ) # Apply matvec _sparsetools.csr_matvec(M, N, A_csr.indptr, A_csr.indices, A_csr.data, x_vec, out_vec) return out_vec