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
