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))