Add per-axis sum/mean/max/min and elementwise arithmetic

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

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

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