Skip to content
This repository was archived by the owner on Nov 10, 2025. It is now read-only.

add from_sparse method #564

Merged
dlzou merged 3 commits into from
Sep 17, 2022
Merged

add from_sparse method #564

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a6c039c
add from_sparse method
trsigg Sep 7, 2022
0e75884
remove asscalar fallback
trsigg Sep 12, 2022
eb24d4e
add new test and outward API, fix old test
trsigg Sep 16, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 12 additions & 5 deletions nums/core/array/application.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -228,18 +230,23 @@ 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:
raise ValueError(
"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)
Expand Down
19 changes: 19 additions & 0 deletions nums/core/array/sparse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
1 change: 0 additions & 1 deletion nums/core/settings.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,6 @@
# and converting back to the original block shape.
doctest_fallbacks = {
"argwhere",
"asscalar",
"clip",
"compress",
"convolve",
Expand Down
10 changes: 10 additions & 0 deletions nums/numpy/api/creation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.

Expand Down
17 changes: 16 additions & 1 deletion tests/core/array/test_sparse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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)
Expand Down