From bc6c94707ae3e9db7044cd13a76aa247ed7e4afc Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Wed, 2 Sep 2026 11:19:36 -0700 Subject: [PATCH 1/2] Add per-axis sum/mean/max/min and elementwise arithmetic VCSCArray/VCSRArray only supported scalar-only mul/truediv and a global sum(); everything else (per-axis reductions, add/sub, and array-array/array-dense multiply) fell back to converting to scipy and losing the VCS type. This wires up mean/max/min per-axis (matching scipy's implicit-zero-aware semantics) and add/sub/multiply against scalars, dense arrays, and other VCS/scipy sparse arrays, re-wrapping sparse results back into the same VCS class. Co-Authored-By: Claude Sonnet 5 --- src/vsparse/_base.py | 155 +++++++++++++++++++++---- tests/test_ops.py | 10 +- tests/test_reductions_and_arith.py | 179 +++++++++++++++++++++++++++++ 3 files changed, 321 insertions(+), 23 deletions(-) create mode 100644 tests/test_reductions_and_arith.py diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py index eb89f05..5402638 100644 --- a/src/vsparse/_base.py +++ b/src/vsparse/_base.py @@ -225,6 +225,121 @@ 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 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: + total = self.shape[0] * self.shape[1] + return self.sum() / total if total else float("nan") + if axis not in (0, 1): + raise ValueError(f"axis must be None, 0, or 1, got {axis!r}") + denom = self.shape[0] if axis == 0 else self.shape[1] + return self.sum(axis=axis) / denom if denom else np.full(0, float("nan")) + + def _reduce_initial(self, kind: str) -> Any: + """Identity element for a max/min reduction over ``self.values.dtype``.""" + dt = self.values.dtype + if np.issubdtype(dt, np.integer): + return np.iinfo(dt).min if kind == "max" else np.iinfo(dt).max + return -np.inf if kind == "max" else np.inf + + def _major_reduce(self, ufunc: np.ufunc, initial: Any) -> np.ndarray: + """Per-major-slice max/min, accounting for implicit zeros in sparse slices.""" + out = np.full(self.n_major, initial, dtype=self.values.dtype) + group_of_major = np.repeat(np.arange(self.n_major, dtype=np.int64), np.diff(self.major_ptr)) + ufunc.at(out, group_of_major, self.values) + nnz_per_major = self.value_ptr[self.major_ptr[1:]] - self.value_ptr[self.major_ptr[:-1]] + not_dense = nnz_per_major < self.n_minor + out[not_dense] = ufunc(out[not_dense], 0) + return out + + def _minor_reduce(self, ufunc: np.ufunc, initial: Any) -> np.ndarray: + """Per-minor-index max/min, accounting for implicit zeros in sparse slices.""" + group_sizes = np.diff(self.value_ptr) + expanded = np.repeat(self.values, group_sizes) + out = np.full(self.n_minor, initial, dtype=self.values.dtype) + ufunc.at(out, self.indices, expanded) + counts = np.bincount(self.indices, minlength=self.n_minor) + not_dense = counts < self.n_major + out[not_dense] = ufunc(out[not_dense], 0) + return out + + def _reduce(self, ufunc: np.ufunc, kind: str, axis: int | None) -> np.ndarray | Any: + initial = self._reduce_initial(kind) + if axis is None: + m = ufunc.reduce(self.values, initial=initial) if self.values.size else initial + if self.nnz < self.shape[0] * self.shape[1]: + m = ufunc(m, 0) + return m.item() if hasattr(m, "item") else m + 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_reduce(ufunc, initial) + if axis == major_axis + else self._minor_reduce(ufunc, initial) + ) + + def max(self, axis: int | None = None) -> np.ndarray | Any: + """Maximum value (including implicit zeros) along ``axis``, or overall if ``None``.""" + return self._reduce(np.maximum, "max", axis) + + def min(self, axis: int | None = None) -> np.ndarray | Any: + """Minimum value (including implicit zeros) along ``axis``, or overall if ``None``.""" + return self._reduce(np.minimum, "min", axis) + + # -- elementwise arithmetic ---------------------------------------------- + + def _elementwise(self, other: Any, op_name: str) -> _VCSBase | np.ndarray: + """Fall back to scipy to compute an elementwise binary op, re-wrapping a sparse result.""" + other_arg = other.to_scipy() if isinstance(other, _VCSBase) else other + self_scipy = self.to_scipy() + if op_name == "multiply": + result = self_scipy.multiply(other_arg) + else: + result = getattr(self_scipy, op_name)(other_arg) + if result is NotImplemented: + return NotImplemented + if sp.issparse(result): + return type(self).from_scipy(result) + return np.asarray(result) + + def __add__(self, other): + if np.isscalar(other): + if other == 0: + return self.copy() + raise NotImplementedError( + "adding a nonzero scalar to a sparse array is not supported" + ) + return self._elementwise(other, "__add__") + + __radd__ = __add__ + + def __sub__(self, other): + if np.isscalar(other): + if other == 0: + return self.copy() + raise NotImplementedError( + "subtracting a nonzero scalar from a sparse array is not supported" + ) + return self._elementwise(other, "__sub__") + + def __rsub__(self, other): + if np.isscalar(other): + if other == 0: + return -self + raise NotImplementedError( + "subtracting a sparse array from a nonzero scalar is not supported" + ) + other_arg = other.to_scipy() if isinstance(other, _VCSBase) else other + result = other_arg - self.to_scipy() + if sp.issparse(result): + return type(self).from_scipy(result) + return np.asarray(result) + + def multiply(self, other) -> _VCSBase | np.ndarray: + """Elementwise multiplication (matches scipy's sparse-array ``.multiply``).""" + return self * other + # -- scalar arithmetic -------------------------------------------------- def _empty_like(self) -> _VCSBase: @@ -238,30 +353,30 @@ def _empty_like(self) -> _VCSBase: ) def __mul__(self, other): - if not np.isscalar(other): - return NotImplemented - if other == 0: - return self._empty_like() - return type(self)( - self.shape, - self.major_ptr.copy(), - self.values * other, - self.value_ptr.copy(), - self.indices.copy(), - ) + if np.isscalar(other): + if other == 0: + return self._empty_like() + return type(self)( + self.shape, + self.major_ptr.copy(), + self.values * other, + self.value_ptr.copy(), + self.indices.copy(), + ) + return self._elementwise(other, "multiply") __rmul__ = __mul__ def __truediv__(self, other): - if not np.isscalar(other): - return NotImplemented - return type(self)( - self.shape, - self.major_ptr.copy(), - self.values / other, - self.value_ptr.copy(), - self.indices.copy(), - ) + if np.isscalar(other): + return type(self)( + self.shape, + self.major_ptr.copy(), + self.values / other, + self.value_ptr.copy(), + self.indices.copy(), + ) + return self._elementwise(other, "__truediv__") def __neg__(self): return self * -1 diff --git a/tests/test_ops.py b/tests/test_ops.py index 2b9f92c..04df920 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -96,10 +96,14 @@ def test_scalar_mul_zero_returns_empty_like(dense, vcls): def test_unsupported_scalar_operands_raise(dense, vcls): - """Verify that non-scalar operand types in arithmetic operators return NotImplemented or raise TypeError.""" + """Non-scalar operands are now elementwise: non-broadcastable shapes/bad types raise, they don't silently no-op.""" v = _make(vcls, dense) - assert v.__mul__([1, 2]) is NotImplemented - assert v.__truediv__([1, 2]) is NotImplemented + if dense.shape != (1, 1): # a (1, 1) array broadcasts against any shape + bad_shape = np.ones((dense.shape[0] + 3, dense.shape[1] + 3)) + with pytest.raises(ValueError, match="inconsistent shapes"): + _ = v * bad_shape + with pytest.raises(ValueError, match="inconsistent shapes"): + _ = v / bad_shape with pytest.raises(TypeError): _ = v * {"a": 1} with pytest.raises(TypeError): diff --git a/tests/test_reductions_and_arith.py b/tests/test_reductions_and_arith.py new file mode 100644 index 0000000..d1fe920 --- /dev/null +++ b/tests/test_reductions_and_arith.py @@ -0,0 +1,179 @@ +"""Tests for per-axis sum/mean/max/min and elementwise arithmetic on VCSCArray/VCSRArray.""" + +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 + + +def make_signed_dense(rng: np.random.Generator, shape: tuple[int, int]) -> np.ndarray: + """Dense matrix with negative, positive, and structural-zero entries.""" + dense = rng.integers(-5, 6, size=shape).astype(np.float64) + mask = rng.random(shape) < 0.4 + dense[mask] = 0.0 + return dense + + +@pytest.fixture(params=[(1, 1), (5, 1), (1, 7), (8, 6), (25, 40), (50, 3)]) +def shape(request) -> tuple[int, int]: + return request.param + + +@pytest.fixture +def signed_dense(shape) -> np.ndarray: + return make_signed_dense(np.random.default_rng(7), shape) + + +# -- sum / mean ------------------------------------------------------------- + + +def test_sum_per_axis(dense, vcls): + v = vcls.from_scipy(sp.csr_array(dense)) + assert v.sum() == pytest.approx(dense.sum()) + np.testing.assert_allclose(v.sum(axis=0), dense.sum(axis=0)) + np.testing.assert_allclose(v.sum(axis=1), dense.sum(axis=1)) + + +def test_sum_invalid_axis(dense, vcls): + v = vcls.from_scipy(sp.csr_array(dense)) + with pytest.raises(ValueError): + v.sum(axis=2) + + +def test_mean_per_axis(dense, vcls): + v = vcls.from_scipy(sp.csr_array(dense)) + assert v.mean() == pytest.approx(dense.mean()) + np.testing.assert_allclose(v.mean(axis=0), dense.mean(axis=0)) + np.testing.assert_allclose(v.mean(axis=1), dense.mean(axis=1)) + + +# -- max / min ---------------------------------------------------------------- + + +def test_max_per_axis(signed_dense, vcls): + v = vcls.from_scipy(sp.csr_array(signed_dense)) + assert v.max() == pytest.approx(signed_dense.max()) + np.testing.assert_allclose(v.max(axis=0), signed_dense.max(axis=0)) + np.testing.assert_allclose(v.max(axis=1), signed_dense.max(axis=1)) + + +def test_min_per_axis(signed_dense, vcls): + v = vcls.from_scipy(sp.csr_array(signed_dense)) + assert v.min() == pytest.approx(signed_dense.min()) + np.testing.assert_allclose(v.min(axis=0), signed_dense.min(axis=0)) + np.testing.assert_allclose(v.min(axis=1), signed_dense.min(axis=1)) + + +def test_max_min_invalid_axis(dense, vcls): + v = vcls.from_scipy(sp.csr_array(dense)) + with pytest.raises(ValueError): + v.max(axis=2) + with pytest.raises(ValueError): + v.min(axis=2) + + +def test_max_min_all_negative_column(): + """A column of only negative values must still report 0 if it has a structural zero.""" + dense = np.array([[-1.0, -2.0], [0.0, -3.0]]) + for vcls in (VCSCArray, VCSRArray): + v = vcls.from_scipy(sp.csr_array(dense)) + np.testing.assert_allclose(v.max(axis=0), dense.max(axis=0)) + np.testing.assert_allclose(v.min(axis=0), dense.min(axis=0)) + + +def test_max_min_fully_dense_negative(): + """A fully-dense negative row/column must not spuriously include 0.""" + dense = np.array([[-1.0, -2.0], [-4.0, -3.0]]) + for vcls in (VCSCArray, VCSRArray): + v = vcls.from_scipy(sp.csr_array(dense)) + assert v.max() == pytest.approx(-1.0) + np.testing.assert_allclose(v.max(axis=0), dense.max(axis=0)) + np.testing.assert_allclose(v.max(axis=1), dense.max(axis=1)) + + +# -- elementwise arithmetic --------------------------------------------------- + + +def test_add_sub_vcs_vcs(dense, vcls): + v = vcls.from_scipy(sp.csr_array(dense)) + other_dense = dense * 2 + other = vcls.from_scipy(sp.csr_array(other_dense)) + + added = v + other + assert isinstance(added, vcls) + np.testing.assert_allclose(added.toarray(), dense + other_dense) + + subbed = v - other + assert isinstance(subbed, vcls) + np.testing.assert_allclose(subbed.toarray(), dense - other_dense) + + +def test_add_sub_dense(dense, vcls): + v = vcls.from_scipy(sp.csr_array(dense)) + other_dense = np.ones_like(dense) + + added = v + other_dense + np.testing.assert_allclose(np.asarray(added), dense + other_dense) + + radded = other_dense + v + np.testing.assert_allclose(np.asarray(radded), other_dense + dense) + + subbed = v - other_dense + np.testing.assert_allclose(np.asarray(subbed), dense - other_dense) + + rsubbed = other_dense - v + np.testing.assert_allclose(np.asarray(rsubbed), other_dense - dense) + + +def test_add_sub_zero_scalar(dense, vcls): + v = vcls.from_scipy(sp.csr_array(dense)) + np.testing.assert_allclose((v + 0).toarray(), dense) + np.testing.assert_allclose((v - 0).toarray(), dense) + np.testing.assert_allclose((0 - v).toarray(), -dense) + + +def test_add_nonzero_scalar_raises(dense, vcls): + v = vcls.from_scipy(sp.csr_array(dense)) + with pytest.raises(NotImplementedError): + v + 5 + with pytest.raises(NotImplementedError): + v - 5 + with pytest.raises(NotImplementedError): + 5 - v + + +def test_multiply_elementwise(dense, vcls): + v = vcls.from_scipy(sp.csr_array(dense)) + other_dense = dense + 1 # avoid trivially all-zero result + other = vcls.from_scipy(sp.csr_array(other_dense)) + + prod = v.multiply(other) + assert isinstance(prod, vcls) + np.testing.assert_allclose(prod.toarray(), dense * other_dense) + + prod_star = v * other + assert isinstance(prod_star, vcls) + np.testing.assert_allclose(prod_star.toarray(), dense * other_dense) + + prod_dense = v.multiply(other_dense) + assert isinstance(prod_dense, vcls) + np.testing.assert_allclose(prod_dense.toarray(), dense * other_dense) + + +def test_scalar_mul_div_unaffected(dense, vcls): + """Existing scalar multiply/divide behavior must be preserved.""" + v = vcls.from_scipy(sp.csr_array(dense)) + np.testing.assert_allclose((v * 3).toarray(), dense * 3) + np.testing.assert_allclose((3 * v).toarray(), dense * 3) + with np.errstate(invalid="ignore", divide="ignore"): + np.testing.assert_allclose((v / 2).toarray(), dense / 2) + assert isinstance(v * 0, vcls) + np.testing.assert_allclose((v * 0).toarray(), np.zeros_like(dense)) From 9a1d55d81df6e4bdfbbeed8e731c7d1724363349 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Wed, 2 Sep 2026 11:29:07 -0700 Subject: [PATCH 2/2] Add general fancy/boolean indexing, per-axis getnnz, and astype __getitem__ previously only had a fast native path for a full minor slice + arbitrary major selection; everything else (minor-only selection, or both axes at once) fell back to converting through to_scipy(). Add _select_minor (mirrors _select_major but filters and remaps `indices`, dropping now-empty unique-value slots) and compose it with _select_major so only the true both-axes-scalar case still needs to convert. Also add getnnz(axis)/count_nonzero() (reusing the major/minor reduction split from sum/max/min) and astype(dtype, copy), rounding out more of scipy's csr_array/csc_array surface natively. Co-Authored-By: Claude Sonnet 5 --- src/vsparse/_anndata_class.py | 7 +- src/vsparse/_base.py | 91 +++++++++++- tests/test_general_indexing_nnz_astype.py | 172 ++++++++++++++++++++++ tests/test_indexing.py | 4 +- 4 files changed, 266 insertions(+), 8 deletions(-) create mode 100644 tests/test_general_indexing_nnz_astype.py 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])