diff --git a/nums/core/array/application.py b/nums/core/array/application.py index ea9d85ed..48654617 100644 --- a/nums/core/array/application.py +++ b/nums/core/array/application.py @@ -19,9 +19,11 @@ import opt_einsum as oe import numpy as np +import sparse from nums.core.array import utils as array_utils from nums.core.array.blockarray import BlockArray, Block +from nums.core.array.sparse import SparseBlockArray from nums.experimental.zarrgroup import ZarrGroup from nums.core.array.random import NumsRandomState from nums.core.kernel.kernel_manager import KernelManager @@ -228,8 +230,8 @@ def loadtxt( def scalar(self, value): return BlockArray.from_scalar(value, self.km) - def array(self, array: Union[np.ndarray, List[float]], block_shape: tuple = None): - if not isinstance(array, np.ndarray): + def array(self, array: Union[np.ndarray, sparse.COO, List[float]], block_shape: tuple = None): + if not isinstance(array, (np.ndarray, sparse.COO)): if array_utils.is_array_like(array): array = np.array(array) else: @@ -237,9 +239,14 @@ def array(self, array: Union[np.ndarray, List[float]], block_shape: tuple = None "Unable to instantiate array from type %s" % type(array) ) assert len(array.shape) == len(block_shape) - return BlockArray.from_np( - array, block_shape=block_shape, copy=False, km=self.km - ) + if isinstance(array, np.ndarray): + return BlockArray.from_np( + array, block_shape=block_shape, copy=False, km=self.km + ) + else: + return SparseBlockArray.from_sparse( + array, block_shape=block_shape, copy=False, km=self.km + ) def zeros(self, shape: tuple, block_shape: tuple, dtype: np.dtype = None): return self._new_array("zeros", shape, block_shape, dtype) diff --git a/nums/core/array/sparse.py b/nums/core/array/sparse.py index b12c6f2d..fb71b0bd 100644 --- a/nums/core/array/sparse.py +++ b/nums/core/array/sparse.py @@ -285,6 +285,25 @@ def from_np(cls, arr, block_shape, copy, km, fill_value=0): rarr.blocks[grid_entry].dtype = getattr(np, dtype_str) return rarr + @classmethod + def from_sparse(cls, arr, block_shape, copy, km, fill_value=0): + dtype_str = str(arr.dtype) + grid = ArrayGrid(arr.shape, block_shape, dtype_str) + rarr = SparseBlockArray(grid, km, fill_value) + grid_entry_iterator = grid.get_entry_iterator() + for grid_entry in grid_entry_iterator: + grid_slice = grid.get_slice(grid_entry) + block = arr[grid_slice] + if copy: + block = sparse.COO.copy(block) + # TODO: generalize for different kernels + rarr.blocks[grid_entry].oid = km.put( + block, + syskwargs={"grid_entry": grid_entry, "grid_shape": grid.grid_shape}, + ) + rarr.blocks[grid_entry].dtype = getattr(np, dtype_str) + return rarr + @classmethod def from_scalar(cls, val, km): if not array_utils.is_scalar(val): diff --git a/nums/core/settings.py b/nums/core/settings.py index 37369b46..cc6357c2 100644 --- a/nums/core/settings.py +++ b/nums/core/settings.py @@ -76,7 +76,6 @@ # and converting back to the original block shape. doctest_fallbacks = { "argwhere", - "asscalar", "clip", "compress", "convolve", diff --git a/nums/numpy/api/creation.py b/nums/numpy/api/creation.py index 4633c0ce..a76e815c 100644 --- a/nums/numpy/api/creation.py +++ b/nums/numpy/api/creation.py @@ -16,6 +16,7 @@ import warnings import numpy as np +import sparse from nums.core.application_manager import instance as _instance from nums.core.array.blockarray import BlockArray @@ -225,6 +226,15 @@ def array(object, dtype=None, copy=True, order="K", ndmin=0, subok=False) -> Blo return app.array(result, block_shape) +def from_coo(a: sparse.COO): + assert(isinstance(a, sparse.COO)) + dtype = np.__getattribute__(str(a.dtype)) + shape = a.shape + app = _instance() + block_shape = app.compute_block_shape(shape, dtype) + return app.array(a, block_shape) + + def copy(a: BlockArray, order="K", subok=False): """Return an array copy of the given object. diff --git a/tests/core/array/test_sparse.py b/tests/core/array/test_sparse.py index febf6b4d..4ef9bc06 100644 --- a/tests/core/array/test_sparse.py +++ b/tests/core/array/test_sparse.py @@ -13,11 +13,24 @@ def test_sparse_init(app_inst: ArrayApplication): x_ba = app_inst.array(x1, block_shape=(2, 2)) x_sba = SparseBlockArray.from_ba(x_ba, fill_value=0) assert x_sba.nnz == 8 - assert x_sba.nbytes == 8 * 8 + 2 * 8 * 8 + assert x_sba.nbytes == 8 * 4 + 2 * 8 * 8 y_ba = x_sba.to_ba() assert np.array_equal(x1, y_ba.get()) +def test_from_coo(app_inst: ArrayApplication): + row_coords = [0, 1, 0, 1, 2, 2, 3, 3] + col_coords = [3, 2, 2, 3, 0, 1, 0, 1] + values = [1, 1, 1, 1, 2, 2, 2, 2] + x_sp = sparse.COO([row_coords, col_coords], values) + x_de = x_sp.todense() + x_sba = app_inst.array(x_sp, block_shape=(2, 2)) + assert x_sba.nnz == 8 + assert x_sba.nbytes == 8 * 4 + 2 * 8 * 8 + y_ba = x_sba.to_ba() + assert np.array_equal(x_de, y_ba.get()) + + def test_sparse_random(app_inst: ArrayApplication): rs: NumsRandomState = app_inst.random_state(1337) x_sba = rs.sparse_randint( @@ -185,6 +198,8 @@ def test_sdtd(app_inst): import conftest app_inst = conftest.get_app("serial") + test_from_coo(app_inst) + test_sparse_init(app_inst) test_sdtp(app_inst) # test_sdtd(app_inst) # test_sparse_init(app_inst)