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)