diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py index 9a3ed29..4fc0269 100644 --- a/src/vsparse/_base.py +++ b/src/vsparse/_base.py @@ -515,31 +515,37 @@ def _select_major(self, key: Any) -> _VCSBase: 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. + A key may name the same index more than once, so this is a fan-out + rather than a remap and one stored element can appear in several + output positions. """ 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) + # Output positions grouped by the original index they came from. + fanout = np.bincount(idx, minlength=self.n_minor) + offsets = np.zeros(self.n_minor + 1, dtype=np.int64) + np.cumsum(fanout, out=offsets[1:]) + positions = np.argsort(idx, kind="stable").astype(np.int64, copy=False) - keep = remap[self.indices] >= 0 - new_indices = remap[self.indices[keep]].astype(self.indices.dtype, copy=False) + slot_counts = _ops.minor_select_counts(self.value_ptr, self.indices, fanout) + kept_slots = np.flatnonzero(slot_counts) - 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[kept_slots] + new_value_ptr = np.zeros(kept_slots.shape[0] + 1, dtype=np.int64) + np.cumsum(slot_counts[kept_slots], out=new_value_ptr[1:]) - 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:]) + # Keep the parent's index dtype, as every other structural op does. + new_indices = np.empty(int(new_value_ptr[-1]), dtype=self.indices.dtype) + _ops.minor_select_fill( + self.value_ptr, self.indices, offsets, positions, + kept_slots, new_value_ptr, new_indices, + ) - 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) + # Slots keep their original order, so each major slice owns a + # contiguous run of them. + major_of_slot = np.searchsorted(self.major_ptr, kept_slots, side="right") - 1 + major_counts = np.bincount(major_of_slot, 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:]) diff --git a/src/vsparse/_ops.py b/src/vsparse/_ops.py index 6aded8a..20f8d26 100644 --- a/src/vsparse/_ops.py +++ b/src/vsparse/_ops.py @@ -5,7 +5,14 @@ import numba import numpy as np -__all__ = ["major_matmat", "major_matvec", "minor_matmat", "minor_matvec"] +__all__ = [ + "major_matmat", + "major_matvec", + "minor_matmat", + "minor_matvec", + "minor_select_counts", + "minor_select_fill", +] @numba.njit(cache=True) @@ -91,3 +98,46 @@ def major_matmat(major_ptr, values, value_ptr, indices, b, n_major, n_minor): def minor_matmat(major_ptr, values, value_ptr, indices, b, n_major): b = np.ascontiguousarray(b) return _minor_matmat(major_ptr, values, value_ptr, indices, b, n_major) + + +# -- minor-axis selection ---------------------------------------------------- + + +@numba.njit(cache=True, parallel=True) +def _minor_select_counts(value_ptr, indices, fanout, out_counts): + # A selection may name one index several times, so each stored element + # contributes ``fanout`` output entries rather than one. + n_slots = value_ptr.shape[0] - 1 + for u in numba.prange(n_slots): # ty: ignore[not-iterable] + c = 0 + for k in range(value_ptr[u], value_ptr[u + 1]): + c += fanout[indices[k]] + out_counts[u] = c + + +@numba.njit(cache=True, parallel=True) +def _minor_select_fill( + value_ptr, indices, offsets, positions, kept_slots, new_value_ptr, out_indices +): + for s in numba.prange(kept_slots.shape[0]): # ty: ignore[not-iterable] + u = kept_slots[s] + pos = new_value_ptr[s] # each slot fills a disjoint range + for k in range(value_ptr[u], value_ptr[u + 1]): + o = indices[k] + for j in range(offsets[o], offsets[o + 1]): # every destination of o + out_indices[pos] = positions[j] + pos += 1 + + +def minor_select_counts(value_ptr, indices, fanout): + """Output element count per stored (major, value) slot.""" + counts = np.empty(value_ptr.shape[0] - 1, dtype=np.int64) + _minor_select_counts(value_ptr, indices, fanout, counts) + return counts + + +def minor_select_fill(value_ptr, indices, offsets, positions, kept_slots, new_value_ptr, out_indices): + """Write the remapped minor indices for the surviving slots, in place.""" + _minor_select_fill( + value_ptr, indices, offsets, positions, kept_slots, new_value_ptr, out_indices + ) diff --git a/tests/test_select_minor_fanout.py b/tests/test_select_minor_fanout.py new file mode 100644 index 0000000..16ed23d --- /dev/null +++ b/tests/test_select_minor_fanout.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import tracemalloc + +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 _scipy_for(vcls, dense): + return sp.csc_array(dense) if vcls is VCSCArray else sp.csr_array(dense) + + +def _grid(n_rows: int = 4, n_cols: int = 6) -> np.ndarray: + """Distinct values everywhere, so any misplacement is visible.""" + return np.arange(1, n_rows * n_cols + 1, dtype=float).reshape(n_rows, n_cols) + + +# -- the regression ---------------------------------------------------------- + + +@pytest.mark.parametrize( + "cols", + [ + pytest.param([1, 1, 3], id="one_repeat"), + pytest.param([0, 0, 0, 0], id="one_index_only"), + pytest.param([0, 1, 1, 2, 2, 2], id="mixed_multiplicities"), + ], +) +def test_duplicate_minor_indices_fan_out(vcls, cols): + dense = _grid() + v = vcls.from_scipy(_scipy_for(vcls, dense)) + np.testing.assert_allclose(v[:, cols].toarray(), dense[:, cols]) + + +def test_duplicate_indices_on_the_other_axis_too(vcls): + """Whichever axis is the minor one for this format takes the same path.""" + dense = _grid() + v = vcls.from_scipy(_scipy_for(vcls, dense)) + rows = [1, 1, 2, 0, 0] + np.testing.assert_allclose(v[rows, :].toarray(), dense[rows, :]) + + +def test_duplicates_on_both_axes_at_once(vcls): + dense = _grid() + v = vcls.from_scipy(_scipy_for(vcls, dense)) + rows, cols = [0, 0, 2], [1, 1, 4] + np.testing.assert_allclose(v[rows, cols].toarray(), dense[np.ix_(rows, cols)]) + + +def test_matches_scipy_for_a_random_selection_with_repeats(vcls, rng): + """Against scipy, which fans duplicate indices out correctly.""" + dense = rng.integers(0, 4, size=(12, 9)).astype(np.float64) + v = vcls.from_scipy(_scipy_for(vcls, dense)) + reference = _scipy_for(vcls, dense) + + for _ in range(25): + cols = rng.integers(0, dense.shape[1], size=rng.integers(1, 15)).tolist() + np.testing.assert_allclose( + v[:, cols].toarray(), np.asarray(reference[:, cols].todense()) + ) + + +# -- selections that already worked, kept working ---------------------------- + + +@pytest.mark.parametrize( + "cols", + [ + pytest.param([0, 2], id="sorted"), + pytest.param([2, 0], id="reversed"), + pytest.param(np.array([], dtype=int), id="empty"), + pytest.param([-1, -2], id="negative"), + pytest.param(slice(1, 4), id="slice"), + pytest.param(slice(None, None, 2), id="strided_slice"), + pytest.param([True, False, True, False, True, False], id="boolean_mask"), + ], +) +def test_non_duplicate_selections_unchanged(vcls, cols): + dense = _grid() + v = vcls.from_scipy(_scipy_for(vcls, dense)) + expected = dense[:, cols] if isinstance(cols, slice) else dense[:, np.asarray(cols)] + np.testing.assert_allclose(v[:, cols].toarray(), expected) + + +def test_empty_selection_gives_a_zero_width_array(vcls): + dense = _grid() + v = vcls.from_scipy(_scipy_for(vcls, dense)) + out = v[:, []] + assert out.shape[1] == 0 + assert out.nnz == 0 + + +def test_selection_of_an_all_zero_array(vcls): + v = vcls.from_scipy(_scipy_for(vcls, np.zeros((4, 5)))) + out = v[:, [1, 1, 2]] + assert out.nnz == 0 + np.testing.assert_allclose(out.toarray(), np.zeros((4, 3))) + + +def test_sparse_columns_and_duplicates_together(vcls): + """A duplicated index whose column is entirely implicit zeros.""" + dense = np.array([[1.0, 0.0, 2.0], [3.0, 0.0, 0.0]]) + v = vcls.from_scipy(_scipy_for(vcls, dense)) + cols = [1, 1, 0, 2] + np.testing.assert_allclose(v[:, cols].toarray(), dense[:, cols]) + + +def test_round_trip_through_scipy_after_a_duplicate_selection(vcls): + """The result has to be a structurally valid VCS array.""" + dense = _grid() + v = vcls.from_scipy(_scipy_for(vcls, dense)) + out = v[:, [1, 1, 3]] + + rebuilt = vcls.from_scipy(out.to_scipy()) + np.testing.assert_allclose(rebuilt.toarray(), dense[:, [1, 1, 3]]) + assert out.value_ptr[-1] == out.indices.shape[0] + assert out.major_ptr[-1] == out.values.shape[0] + + +# -- memory ------------------------------------------------------------------ + + +def test_minor_selection_allocates_nothing_nnz_sized(): + """Bounds the scratch, not the output, which is legitimately nnz-scale.""" + rng = np.random.default_rng(0) + n_rows, n_cols = 2_000, 500 + dense = rng.integers(1, 5, size=(n_rows, n_cols)).astype(np.float64) + v = VCSRArray.from_scipy(sp.csr_array(dense)) + cols = np.arange(0, n_cols, 2) + + v[:, cols] # warm up numba's JIT before measuring + + tracemalloc.start() + try: + before = tracemalloc.get_traced_memory()[0] + tracemalloc.reset_peak() + out = v[:, cols] + peak = tracemalloc.get_traced_memory()[1] + finally: + tracemalloc.stop() + + assert peak - before < 16 << 20, f"allocated {(peak - before) / 1e6:.1f} MB" + np.testing.assert_allclose(out.toarray(), dense[:, cols])