Skip to content

Commit

Permalink
Merge pull request python-graphblas#94 from eriknw/fix/head_iso
Browse files Browse the repository at this point in the history 
Fix `head` on iso-valued objects, which fixes formatting of large objects
  • Loading branch information
@eriknw
eriknw authored Jul 12, 2021
2 parents 6480b3d + 4308498 commit 7095fd1
Show file tree
Hide file tree
Showing 3 changed files with 109 additions and 61 deletions.
73 changes: 43 additions & 30 deletions grblas/_ss/matrix.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,34 +21,44 @@


@njit
def _head_matrix_full(values, nrows, ncols, dtype, n): # pragma: no cover
def _head_matrix_full(values, nrows, ncols, dtype, n, is_iso): # pragma: no cover
rows = np.empty(n, dtype=np.uint64)
cols = np.empty(n, dtype=np.uint64)
vals = np.empty(n, dtype=dtype)
if is_iso:
vals = np.empty(1, dtype=dtype)
vals[0] = values[0]
else:
vals = np.empty(n, dtype=dtype)
k = 0
for i in range(nrows):
for j in range(ncols):
rows[k] = i
cols[k] = j
vals[k] = values[i * ncols + j]
if not is_iso:
vals[k] = values[i * ncols + j]
k += 1
if k == n:
return rows, cols, vals
return rows, cols, vals


@njit
def _head_matrix_bitmap(bitmap, values, nrows, ncols, dtype, n): # pragma: no cover
def _head_matrix_bitmap(bitmap, values, nrows, ncols, dtype, n, is_iso): # pragma: no cover
rows = np.empty(n, dtype=np.uint64)
cols = np.empty(n, dtype=np.uint64)
vals = np.empty(n, dtype=dtype)
if is_iso:
vals = np.empty(1, dtype=dtype)
vals[0] = values[0]
else:
vals = np.empty(n, dtype=dtype)
k = 0
for i in range(nrows):
for j in range(ncols):
if bitmap[i * ncols + j]:
rows[k] = i
cols[k] = j
vals[k] = values[i * ncols + j]
if not is_iso:
vals[k] = values[i * ncols + j]
k += 1
if k == n:
return rows, cols, vals
Expand Down Expand Up @@ -110,24 +120,25 @@ def head(matrix, n=10, *, sort=False, dtype=None):
np.empty(0, dtype=np.uint64),
np.empty(0, dtype=dtype.np_type),
)
is_iso = matrix.ss.is_iso
d = matrix.ss.unpack(raw=True, sort=sort)
try:
fmt = d["format"]
if fmt == "fullr":
rows, cols, vals = _head_matrix_full(
d["values"], d["nrows"], d["ncols"], dtype.np_type, n
d["values"], d["nrows"], d["ncols"], dtype.np_type, n, is_iso
)
elif fmt == "fullc":
cols, rows, vals = _head_matrix_full(
d["values"], d["ncols"], d["nrows"], dtype.np_type, n
d["values"], d["ncols"], d["nrows"], dtype.np_type, n, is_iso
)
elif fmt == "bitmapr":
rows, cols, vals = _head_matrix_bitmap(
d["bitmap"], d["values"], d["nrows"], d["ncols"], dtype.np_type, n
d["bitmap"], d["values"], d["nrows"], d["ncols"], dtype.np_type, n, is_iso
)
elif fmt == "bitmapc":
cols, rows, vals = _head_matrix_bitmap(
d["bitmap"], d["values"], d["ncols"], d["nrows"], dtype.np_type, n
d["bitmap"], d["values"], d["ncols"], d["nrows"], dtype.np_type, n, is_iso
)
elif fmt == "csr":
vals = d["values"][:n].astype(dtype.np_type)
Expand All @@ -149,6 +160,8 @@ def head(matrix, n=10, *, sort=False, dtype=None):
raise RuntimeError(f"Invalid format: {fmt}")
finally:
matrix.ss.pack_any(take_ownership=True, **d)
if is_iso:
vals = np.broadcast_to(vals[:1], (n,))
return rows, cols, vals


Expand Down Expand Up @@ -1068,9 +1081,9 @@ def _import_csr(
values, dtype = values_to_numpy_buffer(values, dtype, copy=copy, ownable=True)
if col_indices is values:
values = np.copy(values)
Ap = ffi_new("GrB_Index**", ffi.cast("GrB_Index*", ffi.from_buffer(indptr)))
Aj = ffi_new("GrB_Index**", ffi.cast("GrB_Index*", ffi.from_buffer(col_indices)))
Ax = ffi_new("void**", ffi.cast("void**", ffi.from_buffer(values)))
Ap = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", indptr))
Aj = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", col_indices))
Ax = ffi_new("void**", ffi.from_buffer("void*", values))
if method == "import":
mhandle = ffi_new("GrB_Matrix*")
args = (dtype._carg, nrows, ncols)
Expand Down Expand Up @@ -1243,9 +1256,9 @@ def _import_csc(
values, dtype = values_to_numpy_buffer(values, dtype, copy=copy, ownable=True)
if row_indices is values:
values = np.copy(values)
Ap = ffi_new("GrB_Index**", ffi.cast("GrB_Index*", ffi.from_buffer(indptr)))
Ai = ffi_new("GrB_Index**", ffi.cast("GrB_Index*", ffi.from_buffer(row_indices)))
Ax = ffi_new("void**", ffi.cast("void**", ffi.from_buffer(values)))
Ap = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", indptr))
Ai = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", row_indices))
Ax = ffi_new("void**", ffi.from_buffer("void*", values))
if method == "import":
mhandle = ffi_new("GrB_Matrix*")
args = (dtype._carg, nrows, ncols)
Expand Down Expand Up @@ -1433,10 +1446,10 @@ def _import_hypercsr(
values, dtype = values_to_numpy_buffer(values, dtype, copy=copy, ownable=True)
if col_indices is values:
values = np.copy(values)
Ap = ffi_new("GrB_Index**", ffi.cast("GrB_Index*", ffi.from_buffer(indptr)))
Ah = ffi_new("GrB_Index**", ffi.cast("GrB_Index*", ffi.from_buffer(rows)))
Aj = ffi_new("GrB_Index**", ffi.cast("GrB_Index*", ffi.from_buffer(col_indices)))
Ax = ffi_new("void**", ffi.cast("void**", ffi.from_buffer(values)))
Ap = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", indptr))
Ah = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", rows))
Aj = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", col_indices))
Ax = ffi_new("void**", ffi.from_buffer("void*", values))
if nvec is None:
nvec = rows.size
if method == "import":
Expand Down Expand Up @@ -1629,10 +1642,10 @@ def _import_hypercsc(
values, dtype = values_to_numpy_buffer(values, dtype, copy=copy, ownable=True)
if row_indices is values:
values = np.copy(values)
Ap = ffi_new("GrB_Index**", ffi.cast("GrB_Index*", ffi.from_buffer(indptr)))
Ah = ffi_new("GrB_Index**", ffi.cast("GrB_Index*", ffi.from_buffer(cols)))
Ai = ffi_new("GrB_Index**", ffi.cast("GrB_Index*", ffi.from_buffer(row_indices)))
Ax = ffi_new("void**", ffi.cast("void**", ffi.from_buffer(values)))
Ap = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", indptr))
Ah = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", cols))
Ai = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", row_indices))
Ax = ffi_new("void**", ffi.from_buffer("void*", values))
if nvec is None:
nvec = cols.size
if method == "import":
Expand Down Expand Up @@ -1813,8 +1826,8 @@ def _import_bitmapr(
nrows, ncols = get_shape(nrows, ncols, values=values, bitmap=bitmap)
else:
nrows, ncols = matrix.shape
Ab = ffi_new("int8_t**", ffi.cast("int8_t*", ffi.from_buffer(bitmap)))
Ax = ffi_new("void**", ffi.cast("void**", ffi.from_buffer(values)))
Ab = ffi_new("int8_t**", ffi.from_buffer("int8_t*", bitmap))
Ax = ffi_new("void**", ffi.from_buffer("void*", values))
if nvals is None:
if bitmap.size == nrows * ncols:
nvals = np.count_nonzero(bitmap)
Expand Down Expand Up @@ -1991,8 +2004,8 @@ def _import_bitmapc(
nrows, ncols = get_shape(nrows, ncols, values=values, bitmap=bitmap)
else:
nrows, ncols = matrix.shape
Ab = ffi_new("int8_t**", ffi.cast("int8_t*", ffi.from_buffer(bitmap.T)))
Ax = ffi_new("void**", ffi.cast("void**", ffi.from_buffer(values.T)))
Ab = ffi_new("int8_t**", ffi.from_buffer("int8_t*", bitmap.T))
Ax = ffi_new("void**", ffi.from_buffer("void*", values.T))
if nvals is None:
if bitmap.size == nrows * ncols:
nvals = np.count_nonzero(bitmap)
Expand Down Expand Up @@ -2150,7 +2163,7 @@ def _import_fullr(
else:
nrows, ncols = matrix.shape

Ax = ffi_new("void**", ffi.cast("void**", ffi.from_buffer(values)))
Ax = ffi_new("void**", ffi.from_buffer("void*", values))
if method == "import":
mhandle = ffi_new("GrB_Matrix*")
args = (dtype._carg, nrows, ncols)
Expand Down Expand Up @@ -2299,7 +2312,7 @@ def _import_fullc(
nrows, ncols = get_shape(nrows, ncols, values=values)
else:
nrows, ncols = matrix.shape
Ax = ffi_new("void**", ffi.cast("void**", ffi.from_buffer(values.T)))
Ax = ffi_new("void**", ffi.from_buffer("void*", values.T))
if method == "import":
mhandle = ffi_new("GrB_Matrix*")
args = (dtype._carg, nrows, ncols)
Expand Down
26 changes: 17 additions & 9 deletions grblas/_ss/vector.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,19 @@


@njit
def _head_indices_vector_bitmap(bitmap, values, size, dtype, n): # pragma: no cover
def _head_indices_vector_bitmap(bitmap, values, size, dtype, n, is_iso): # pragma: no cover
indices = np.empty(n, dtype=np.uint64)
vals = np.empty(n, dtype=dtype)
if is_iso:
vals = np.empty(1, dtype=dtype)
vals[0] = values[0]
else:
vals = np.empty(n, dtype=dtype)
j = 0
for i in range(size):
if bitmap[i]:
indices[j] = i
vals[j] = values[i]
if not is_iso:
vals[j] = values[i]
j += 1
if j == n:
break
Expand All @@ -42,6 +47,7 @@ def head(vector, n=10, *, sort=False, dtype=None):
n = min(n, vector._nvals)
if n == 0:
return (np.empty(0, dtype=np.uint64), np.empty(0, dtype=dtype.np_type))
is_iso = vector.ss.is_iso
d = vector.ss.unpack(raw=True, sort=sort)
fmt = d["format"]
try:
Expand All @@ -50,7 +56,7 @@ def head(vector, n=10, *, sort=False, dtype=None):
vals = d["values"][:n].astype(dtype.np_type)
elif fmt == "bitmap":
indices, vals = _head_indices_vector_bitmap(
d["bitmap"], d["values"], d["size"], dtype.np_type, n
d["bitmap"], d["values"], d["size"], dtype.np_type, n, is_iso
)
elif fmt == "sparse":
indices = d["indices"][:n].copy()
Expand All @@ -59,6 +65,8 @@ def head(vector, n=10, *, sort=False, dtype=None):
raise RuntimeError(f"Invalid format: {fmt}")
finally:
vector.ss.pack_any(take_ownership=True, **d)
if is_iso:
vals = np.broadcast_to(vals[:1], (n,))
return indices, vals


Expand Down Expand Up @@ -647,8 +655,8 @@ def _import_sparse(
values, dtype = values_to_numpy_buffer(values, dtype, copy=copy, ownable=True)
if indices is values:
values = np.copy(values)
vi = ffi_new("GrB_Index**", ffi.cast("GrB_Index*", ffi.from_buffer(indices)))
vx = ffi_new("void**", ffi.cast("void**", ffi.from_buffer(values)))
vi = ffi_new("GrB_Index**", ffi.from_buffer("GrB_Index*", indices))
vx = ffi_new("void**", ffi.from_buffer("void*", values))
if nvals is None:
if is_iso:
nvals = indices.size
Expand Down Expand Up @@ -813,8 +821,8 @@ def _import_bitmap(
if bitmap is values:
values = np.copy(values)
vhandle = ffi_new("GrB_Vector*")
vb = ffi_new("int8_t**", ffi.cast("int8_t*", ffi.from_buffer(bitmap)))
vx = ffi_new("void**", ffi.cast("void**", ffi.from_buffer(values)))
vb = ffi_new("int8_t**", ffi.from_buffer("int8_t*", bitmap))
vx = ffi_new("void**", ffi.from_buffer("void*", values))
if size is None:
if is_iso:
size = bitmap.size
Expand Down Expand Up @@ -966,7 +974,7 @@ def _import_full(
size = vector._size
values, dtype = values_to_numpy_buffer(values, dtype, copy=copy, ownable=True)
vhandle = ffi_new("GrB_Vector*")
vx = ffi_new("void**", ffi.cast("void**", ffi.from_buffer(values)))
vx = ffi_new("void**", ffi.from_buffer("void*", values))
if size is None:
size = values.size
if method == "import":
Expand Down
Loading

0 comments on commit 7095fd1

Please sign in to comment.