Add general fancy/boolean indexing, per-axis getnnz, and astype

src/vsparse/_anndata_class.py

@@ -48,9 +48,10 @@ def _subset_2d(v: Any, oidx: Any, vidx: Any) -> Any:
  return None
  if isinstance(v, _VCS_TYPES):
  result = v[oidx, vidx]
- # VCSCArray/VCSRArray fall back to a plain scipy array for general
- # (both-axes) indexing (see _VCSBase.__getitem__); re-wrap so X/raw_X
- # stay VCS-backed the way the rest of this class requires.
+ # VCSCArray/VCSRArray.__getitem__ only converts to a plain scipy
+ # array when both axes collapse to a scalar (oidx/vidx are always
+ # slices/arrays here, never bare ints -- see _as_slice_index above),
+ # but re-wrap defensively so X/raw_X always stay VCS-backed.
  if not isinstance(result, _VCS_TYPES):
  result = type(v).from_scipy(result)
  return result

src/vsparse/_base.py

@@ -148,6 +148,19 @@ def to_csr(self) -> sp.csr_array:
  def toarray(self) -> np.ndarray:
  return self.to_scipy().toarray()
 
+ def astype(self, dtype: Any, copy: bool = True) -> _VCSBase:
+ """Cast the stored values to ``dtype``. Structural zeros stay zero implicitly."""
+ dtype = np.dtype(dtype)
+ if not copy and dtype == self.dtype:
+ return self
+ return type(self)(
+ self.shape,
+ self.major_ptr.copy(),
+ self.values.astype(dtype),
+ self.value_ptr.copy(),
+ self.indices.copy(),
+ )
+
  # -- structural ops ----------------------------------------------------
 
  @property
@@ -225,6 +238,30 @@ def sum(self, axis: int | None = None) -> np.ndarray | float:
  major_axis = 0 if self._format == "csc" else 1
  return self._major_sums() if axis == major_axis else self._minor_sums()
 
+ def _major_nnz(self) -> np.ndarray:
+ """Per-major-slice stored-element counts -- doesn't touch ``indices``."""
+ return (self.value_ptr[self.major_ptr[1:]] - self.value_ptr[self.major_ptr[:-1]]).astype(
+ np.int64
+ )
+
+ def _minor_nnz(self) -> np.ndarray:
+ """Per-minor-index stored-element counts -- a scatter-add over every nonzero."""
+ return np.bincount(self.indices, minlength=self.n_minor).astype(np.int64)
+
+ def getnnz(self, axis: int | None = None) -> np.ndarray | int:
+ """Count of stored elements along ``axis``, or overall if ``None``."""
+ if axis is None:
+ return self.nnz
+ if axis not in (0, 1):
+ raise ValueError(f"axis must be None, 0, or 1, got {axis!r}")
+ major_axis = 0 if self._format == "csc" else 1
+ return self._major_nnz() if axis == major_axis else self._minor_nnz()
+
+ def count_nonzero(self) -> int:
+ """Count of stored elements that are actually nonzero (unlike :attr:`nnz`/``getnnz``)."""
+ group_sizes = np.diff(self.value_ptr)
+ return int(np.sum(group_sizes[self.values != 0]))
+
  def mean(self, axis: int | None = None) -> np.ndarray | float:
  """Mean of (structural + implicit-zero) values along ``axis``, or overall if ``None``."""
  if axis is None:
@@ -475,6 +512,42 @@ def _select_major(self, key: Any) -> _VCSBase:
  )
  return type(self)(new_shape, new_major_ptr, new_values, new_value_ptr, new_indices)
 
+ def _select_minor(self, key: Any) -> _VCSBase:
+ """Select along the minor axis (rows for VCSC, columns for VCSR).
+
+ Unlike :meth:`_select_major`, the kept elements aren't already
+ contiguous per major slice, so this filters/remaps ``indices`` and
+ drops any (major, unique-value) slot that no longer has any kept
+ index, shrinking ``major_ptr``/``value_ptr`` accordingly.
+ """
+ idx = _normalize_major_idx(key, self.n_minor)
+ n_minor_new = idx.shape[0]
+
+ remap = np.full(self.n_minor, -1, dtype=np.int64)
+ remap[idx] = np.arange(n_minor_new, dtype=np.int64)
+
+ keep = remap[self.indices] >= 0
+ new_indices = remap[self.indices[keep]].astype(self.indices.dtype, copy=False)
+
+ n_unique = self.values.shape[0]
+ value_slot_of_index = np.repeat(np.arange(n_unique, dtype=np.int64), np.diff(self.value_ptr))
+ kept_per_slot = np.bincount(value_slot_of_index[keep], minlength=n_unique)
+ surviving = kept_per_slot > 0
+
+ new_values = self.values[surviving]
+ new_value_ptr = np.zeros(int(surviving.sum()) + 1, dtype=np.int64)
+ np.cumsum(kept_per_slot[surviving], out=new_value_ptr[1:])
+
+ group_of_major = np.repeat(np.arange(self.n_major, dtype=np.int64), np.diff(self.major_ptr))
+ major_counts = np.bincount(group_of_major[surviving], minlength=self.n_major)
+ new_major_ptr = np.zeros(self.n_major + 1, dtype=np.int64)
+ np.cumsum(major_counts, out=new_major_ptr[1:])
+
+ new_shape = (
+ (n_minor_new, self.n_major) if self._format == "csc" else (self.n_major, n_minor_new)
+ )
+ return type(self)(new_shape, new_major_ptr, new_values, new_value_ptr, new_indices)
+
  def __getitem__(self, key):
  if isinstance(key, tuple):
  if len(key) != 2:
@@ -483,13 +556,23 @@ def __getitem__(self, key):
  else:
  row_key, col_key = key, slice(None)
 
+ # A bare int on *both* axes must collapse to a scalar, which a 2-D
+ # VCSC/VCSR array can't represent -- only that case needs to convert.
+ if isinstance(row_key, int | np.integer) and isinstance(col_key, int | np.integer):
+ return self.to_scipy()[row_key, col_key]
+
  major_key, minor_key = (
  (col_key, row_key) if self._format == "csc" else (row_key, col_key)
  )
- if _is_full_slice(minor_key) and not _is_full_slice(major_key):
- return self._select_major(major_key)
-
- return self.to_scipy()[row_key, col_key]
+ if _is_full_slice(major_key) and _is_full_slice(minor_key):
+ return self.copy()
+
+ result = self
+ if not _is_full_slice(major_key):
+ result = result._select_major(major_key)
+ if not _is_full_slice(minor_key):
+ result = result._select_minor(minor_key)
+ return result
 
 
 class VCSCArray(_VCSBase):

tests/test_general_indexing_nnz_astype.py

@@ -0,0 +1,172 @@
+"""Tests for general (both-axes) indexing, getnnz/count_nonzero, and astype."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import scipy.sparse as sp
+
+from vsparse import VCSCArray, VCSRArray
+
+
+@pytest.fixture(params=[VCSCArray, VCSRArray])
+def vcls(request):
+ return request.param
+
+
+# -- general (both-axes) indexing --------------------------------------------
+
+
+def test_general_slice_both_axes_native(dense, vcls):
+ v = vcls.from_scipy(sp.csr_array(dense))
+ if dense.shape[0] < 2 or dense.shape[1] < 2:
+ pytest.skip("shape too small")
+ sub = v[0:2, 0:2]
+ assert isinstance(sub, vcls)
+ np.testing.assert_allclose(sub.toarray(), dense[0:2, 0:2])
+
+
+def test_slice_and_fancy_combo(dense, vcls):
+ if dense.shape[0] < 3 or dense.shape[1] < 3:
+ pytest.skip("shape too small")
+ v = vcls.from_scipy(sp.csr_array(dense))
+ sub = v[1:3, [0, 2]]
+ assert isinstance(sub, vcls)
+ np.testing.assert_allclose(sub.toarray(), dense[1:3][:, [0, 2]])
+
+
+def test_boolean_mask_both_axes(dense, vcls):
+ if dense.shape[0] < 2 or dense.shape[1] < 2:
+ pytest.skip("shape too small")
+ v = vcls.from_scipy(sp.csr_array(dense))
+ row_mask = np.zeros(dense.shape[0], dtype=bool)
+ row_mask[::2] = True
+ col_mask = np.zeros(dense.shape[1], dtype=bool)
+ col_mask[1::2] = True
+ sub = v[row_mask, :][:, col_mask]
+ assert isinstance(sub, vcls)
+ np.testing.assert_allclose(sub.toarray(), dense[row_mask][:, col_mask])
+
+ sub2 = v[row_mask][:, col_mask]
+ np.testing.assert_allclose(sub2.toarray(), dense[row_mask][:, col_mask])
+
+
+def test_minor_axis_only_selection(dense, vcls):
+ """Selecting only along the minor axis (major key is a full slice) now stays native."""
+ v = vcls.from_scipy(sp.csr_array(dense))
+ n_minor = dense.shape[0] if vcls is VCSCArray else dense.shape[1]
+ if n_minor < 2:
+ pytest.skip("axis too small")
+ picks = [n_minor - 1, 0]
+ if vcls is VCSCArray:
+ sub, expected = v[picks, :], dense[picks, :]
+ else:
+ sub, expected = v[:, picks], dense[:, picks]
+ assert isinstance(sub, vcls)
+ np.testing.assert_allclose(sub.toarray(), expected)
+
+
+def test_minor_axis_boolean(dense, vcls):
+ v = vcls.from_scipy(sp.csr_array(dense))
+ n_minor = dense.shape[0] if vcls is VCSCArray else dense.shape[1]
+ mask = np.zeros(n_minor, dtype=bool)
+ mask[::2] = True
+ if vcls is VCSCArray:
+ sub, expected = v[mask, :], dense[mask, :]
+ else:
+ sub, expected = v[:, mask], dense[:, mask]
+ assert isinstance(sub, vcls)
+ np.testing.assert_allclose(sub.toarray(), expected)
+
+
+def test_minor_axis_empty_selection(dense, vcls):
+ v = vcls.from_scipy(sp.csr_array(dense))
+ if vcls is VCSCArray:
+ sub = v[[], :]
+ assert sub.shape == (0, dense.shape[1])
+ else:
+ sub = v[:, []]
+ assert sub.shape == (dense.shape[0], 0)
+ assert isinstance(sub, vcls)
+ assert sub.nnz == 0
+ assert sub.n_unique == 0
+
+
+def test_both_full_slice_returns_copy(dense, vcls):
+ v = vcls.from_scipy(sp.csr_array(dense))
+ sub = v[:, :]
+ assert isinstance(sub, vcls)
+ assert sub is not v
+ np.testing.assert_allclose(sub.toarray(), dense)
+
+
+def test_both_int_still_returns_scalar(dense, vcls):
+ v = vcls.from_scipy(sp.csr_array(dense))
+ val = v[0, 0]
+ assert np.isscalar(val) or isinstance(val, np.generic)
+ assert val == dense[0, 0]
+
+
+# -- getnnz / count_nonzero ----------------------------------------------------
+
+
+def test_getnnz_overall(dense, vcls):
+ v = vcls.from_scipy(sp.csr_array(dense))
+ assert v.getnnz() == np.count_nonzero(dense)
+
+
+def test_getnnz_per_axis(dense, vcls):
+ v = vcls.from_scipy(sp.csr_array(dense))
+ np.testing.assert_array_equal(v.getnnz(axis=0), np.count_nonzero(dense, axis=0))
+ np.testing.assert_array_equal(v.getnnz(axis=1), np.count_nonzero(dense, axis=1))
+
+
+def test_getnnz_invalid_axis(dense, vcls):
+ v = vcls.from_scipy(sp.csr_array(dense))
+ with pytest.raises(ValueError):
+ v.getnnz(axis=2)
+
+
+def test_count_nonzero(dense, vcls):
+ v = vcls.from_scipy(sp.csr_array(dense))
+ assert v.count_nonzero() == np.count_nonzero(dense)
+
+
+def test_count_nonzero_excludes_explicit_zero_values():
+ """count_nonzero must exclude stored-but-zero values, unlike nnz/getnnz."""
+ dense = np.array([[1.0, 0.0], [0.0, 2.0]])
+ for vcls in (VCSCArray, VCSRArray):
+ v = vcls.from_scipy(sp.csr_array(dense))
+ # Manually inject an explicit zero into the stored (unique) values,
+ # which nnz/getnnz still counts as a stored element.
+ zeroed_values = v.values.copy()
+ zeroed_values[0] = 0.0
+ v2 = vcls(v.shape, v.major_ptr, zeroed_values, v.value_ptr, v.indices)
+ assert v2.getnnz() == v.nnz
+ assert v2.count_nonzero() < v2.getnnz()
+
+
+# -- astype ---------------------------------------------------------------------
+
+
+def test_astype_casts_values(dense, vcls):
+ v = vcls.from_scipy(sp.csr_array(dense))
+ out = v.astype(np.float32)
+ assert isinstance(out, vcls)
+ assert out.dtype == np.float32
+ np.testing.assert_allclose(out.toarray(), dense.astype(np.float32))
+ # original is untouched
+ assert v.dtype == dense.dtype
+
+
+def test_astype_no_copy_same_dtype_returns_self(dense, vcls):
+ v = vcls.from_scipy(sp.csr_array(dense))
+ out = v.astype(v.dtype, copy=False)
+ assert out is v
+
+
+def test_astype_copy_true_same_dtype_returns_new_object(dense, vcls):
+ v = vcls.from_scipy(sp.csr_array(dense))
+ out = v.astype(v.dtype, copy=True)
+ assert out is not v
+ np.testing.assert_allclose(out.toarray(), dense)

tests/test_indexing.py

@@ -51,11 +51,13 @@ def test_major_axis_boolean(dense, vcls):
  np.testing.assert_allclose(sub.to_scipy().toarray(), expected)
 
 
-def test_general_2d_indexing_falls_back(dense, vcls):
+def test_general_2d_indexing(dense, vcls):
+ """Slicing both axes at once now stays VCS-native (see test_general_indexing_nnz_astype.py)."""
  v = vcls.from_scipy(sp.csr_array(dense))
  if dense.shape[0] < 2 or dense.shape[1] < 2:
  pytest.skip("shape too small")
  result = v[0:2, 0:2]
+ assert isinstance(result, vcls)
  np.testing.assert_allclose(_as_dense(result), dense[0:2, 0:2])