diff --git a/src/vsparse/_anndata_class.py b/src/vsparse/_anndata_class.py index 789b172..69e479a 100644 --- a/src/vsparse/_anndata_class.py +++ b/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 diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py index 5402638..9a3ed29 100644 --- a/src/vsparse/_base.py +++ b/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): diff --git a/tests/test_general_indexing_nnz_astype.py b/tests/test_general_indexing_nnz_astype.py new file mode 100644 index 0000000..a263d52 --- /dev/null +++ b/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) diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 1440aa8..e1becc3 100644 --- a/tests/test_indexing.py +++ b/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])