def mult_ab(a_h, b_h):
"""
Multiply matrices A and B.
Args:
a_h: the handle of matrix A
b_h: the handle of matrix B
Returns:
the handle of the product; it must be released when no longer needed.
"""
assert a_h.ncols == b_h.nrows
c_rp = np.zeros(a_h.nrows + 1, np.intc)
# step 1: symbolic multiplication
c_ci = _sym_mm(a_h, b_h, c_rp)
c_nnz = c_rp[a_h.nrows]
# step 2: numeric multiplication
c_vs = _num_mm(a_h, b_h, c_rp, c_ci)
# build the result
return CSR(a_h.nrows, b_h.ncols, c_nnz, c_rp, c_ci, c_vs)
def test_large_mult_vec():
# 10M * 500 = 2.5B >= INT_MAX
nrows = 10000000
ncols = 500
dense = 250
nnz = nrows * dense
rowptrs = np.arange(0, nnz + 1, dense, dtype=np.int64)
assert len(rowptrs) == nrows + 1
assert rowptrs[-1] == nnz
try:
_log.info('allocating indexes')
colinds = np.empty(nnz, dtype=np.intc)
_log.info('allocating values')
values = np.zeros(nnz)
except MemoryError:
skip('insufficient memory')
_log.info('randomizing array contents')
fill_rows(values, colinds, nrows, ncols, dense)
csr = CSR(nrows, ncols, nnz, rowptrs, colinds, values)
v = np.random.randn(ncols)
res = csr.mult_vec(v)
assert res.shape == (nrows, )
assert np.all(~np.isnan(res))
def mkh(csr):
vs = csr._required_values().astype(np.float64)
csr2 = CSR(csr.nrows, csr.ncols, csr.nnz, csr.rowptrs, csr.colinds, vs)
if csr.nnz == 0:
return mkl_h(0, csr.nrows, csr.ncols, csr2)
return _make_handle(csr2)
def test_csr64_pickle(csr):
csr = CSR(csr.nrows, csr.ncols, csr.nnz,
csr.rowptrs.astype(np.int64), csr.colinds, csr.values,
_cast=False)
data = pickle.dumps(csr)
csr2 = pickle.loads(data)
assert csr2.nrows == csr.nrows
assert csr2.ncols == csr.ncols
assert csr2.nnz == csr.nnz
assert all(csr2.rowptrs == csr.rowptrs)
assert csr2.rowptrs.dtype == np.int64
assert all(csr2.colinds == csr.colinds)
if csr.values is not None:
assert all(csr2.values == csr.values)
else:
assert csr2.values is None
def test_large_init():
# 10M * 500 = 5B >= INT_MAX
nrows = 10000000
ncols = 500
nnz = nrows * 250
rowptrs = np.arange(0, nnz + 1, 250, dtype=np.int64)
assert len(rowptrs) == nrows + 1
assert rowptrs[-1] == nnz
try:
colinds = np.empty(nnz, dtype=np.intc)
except MemoryError:
pytest.skip('insufficient memory')
csr = CSR(nrows, ncols, nnz, rowptrs, colinds, None)
assert csr.nrows == nrows
assert csr.ncols == ncols
assert csr.nnz == nnz
assert csr.rowptrs.dtype == np.dtype('i8')
def from_handle(h: mkl_h) -> CSR:
if not h.H:
return create_empty(h.nrows, h.ncols)
rvp = lk_mkl_spexport_p(h.H)
if rvp is None:
return None
nrows = lk_mkl_spe_nrows(rvp)
ncols = lk_mkl_spe_ncols(rvp)
sp = lk_mkl_spe_row_sp(rvp)
ep = lk_mkl_spe_row_ep(rvp)
cis = lk_mkl_spe_colinds(rvp)
vs = lk_mkl_spe_values(rvp)
rowptrs = np.zeros(nrows + 1, dtype=np.intc)
nnz = 0
for i in range(nrows):
nnz += ep[i] - sp[i]
rowptrs[i + 1] = nnz
assert nnz == ep[nrows - 1]
colinds = np.zeros(nnz, dtype=np.intc)
values = np.zeros(nnz)
for i in range(nrows):
rs = rowptrs[i]
re = rowptrs[i + 1]
ss = sp[i]
for j in range(re - rs):
colinds[rs + j] = cis[ss + j]
values[rs + j] = vs[ss + j]
lk_mkl_spe_free(rvp)
return CSR(nrows, ncols, nnz, rowptrs, colinds, values)
def from_handle(h):
m: csr_matrix = h.tocsr()
nr, nc = m.shape
return CSR(nr, nc, m.nnz, m.indptr, m.indices, m.data)
