diff --git a/src/vsparse/_anndata_class.py b/src/vsparse/_anndata_class.py index f514b83..ee284b0 100644 --- a/src/vsparse/_anndata_class.py +++ b/src/vsparse/_anndata_class.py @@ -2,6 +2,7 @@ from __future__ import annotations +import copy as _copy from types import MappingProxyType from typing import TYPE_CHECKING, Any, cast @@ -69,6 +70,16 @@ def _subset_2d(v: Any, oidx: Any, vidx: Any) -> Any: return np.asarray(v)[oidx][:, vidx] +def _copy_value(v: Any) -> Any: + """A deep-enough copy of one obs/var/obsm/varm/obsp/varp/layers value. + + Every value anndata can hold there -- a DataFrame, ndarray, scipy sparse + array, or a :class:`~vsparse.VCSCArray`/:class:`~vsparse.VCSRArray` -- + implements its own ``.copy()``. + """ + return None if v is None else v.copy() + + def _check_vcs_type(value: Any, name: str) -> None: if value is not None and not isinstance(value, _VCS_TYPES): raise TypeError( @@ -176,6 +187,11 @@ def __getitem__(self, index: Any) -> VCSCAnnData: # ty: ignore[invalid-method-o this class doesn't use (see the class docstring), so it can't be reused here. ``X``/``raw_X`` stay VCSC/VCSR-backed either way, via that array type's own indexing. + + Returns ``type(self)``, not a bare ``VCSCAnnData``, so a subclass + overriding ``X`` (e.g. one that always hands back a normalized view + computed fresh from ``_vcs_X``, as with :meth:`copy`/:meth:`to_memory`) + keeps that behavior after slicing. """ oidx, vidx = self._normalize_indices(index) oidx = _as_slice_index(oidx, self.n_obs) @@ -197,7 +213,7 @@ def __getitem__(self, index: Any) -> VCSCAnnData: # ty: ignore[invalid-method-o if _VSPARSE_UNS_KEY in uns: uns = {**uns, _VSPARSE_UNS_KEY: {**uns[_VSPARSE_UNS_KEY], "stale": True}} - return VCSCAnnData( + return type(self)( X=_subset_2d(self._vcs_X, oidx, vidx), raw_X=_subset_2d(self._vcs_raw_X, oidx, vidx), obs=obs, @@ -210,6 +226,49 @@ def __getitem__(self, index: Any) -> VCSCAnnData: # ty: ignore[invalid-method-o layers={k: _subset_2d(v, oidx, vidx) for k, v in self.layers.items() if k is not None}, ) + def copy(self) -> VCSCAnnData: # ty: ignore[invalid-method-override] + """A deep copy, preserving the VCSC/VCSR-backed ``X``/``raw_X``. + + The inherited :meth:`anndata.AnnData.copy` only knows how to copy the + standard private ``_X`` attribute, which this class never sets (see + the class docstring: ``X``/``raw_X`` live in ``_vcs_X``/``_vcs_raw_X`` + instead) -- so it silently drops them (``X`` comes back ``None``) and + returns a plain ``AnnData`` rather than this class. This copies every + field explicitly instead, including a real ``VCSCArray``/ + ``VCSRArray`` copy of ``X``/``raw_X``, and returns ``type(self)`` so a + subclass (e.g. one overriding the ``X`` property) round-trips too. + """ + return type(self)( + X=_copy_value(self._vcs_X), + raw_X=_copy_value(self._vcs_raw_X), + obs=cast(pd.DataFrame, self.obs).copy(), + var=cast(pd.DataFrame, self.var).copy(), + uns=_copy.deepcopy(dict(self.uns)), + obsm={k: _copy_value(v) for k, v in self.obsm.items() if k is not None}, + varm={k: _copy_value(v) for k, v in self.varm.items() if k is not None}, + obsp={k: _copy_value(v) for k, v in self.obsp.items() if k is not None}, + varp={k: _copy_value(v) for k, v in self.varp.items() if k is not None}, + layers={k: _copy_value(v) for k, v in self.layers.items() if k is not None}, + ) + + def to_memory(self, *, copy: bool = False) -> VCSCAnnData: + """Return this object with its data loaded into memory. + + The inherited :meth:`anndata.AnnData.to_memory` has the same problem + as the inherited ``copy()`` (see above): it iterates the object's + *standard* attributes, which never includes this class's ``X``/ + ``raw_X`` (held in ``_vcs_X``/``_vcs_raw_X`` instead), and reconstructs + a plain ``AnnData`` -- so ``X`` silently comes back ``None``. + + This class never actually supports a lazily backed ``X``/``raw_X`` + (they're always eagerly-held ``VCSCArray``/``VCSRArray`` instances), + so there is never anything to load -- this always returns a full + :meth:`copy` instead, regardless of ``copy`` (unlike plain + ``AnnData``, where ``copy=False`` can skip copying arrays already in + memory; here everything already is, so the distinction doesn't apply). + """ + return self.copy() + # -- normalization ---------------------------------------------------------- def normalized(self, view: str | Recipe = DEFAULT_RECIPE, *, recalculate: bool = True) -> Any: diff --git a/tests/test_anndata_class.py b/tests/test_anndata_class.py index 1c4c06d..141ef17 100644 --- a/tests/test_anndata_class.py +++ b/tests/test_anndata_class.py @@ -251,3 +251,126 @@ def test_getitem_single_int_row(base_adata, dense): assert sub.shape == (1, dense.shape[1]) assert isinstance(sub.X, VCSCArray) np.testing.assert_allclose(sub.X.toarray(), dense[0:1]) + + +def test_getitem_preserves_a_subclass_overriding_x(base_adata, dense): + """A subclass overriding the `X` getter (as BAL-Pf2's own + lazy-normalized-view AnnData does) must keep that behavior after + slicing, not silently fall back to the raw (un-normalized) array.""" + if dense.shape[0] < 2 or dense.sum() == 0: + pytest.skip("shape too small or all-zero matrix") + + class _NormalizedView(VCSCAnnData): + @property + def X(self): + return self.normalized("parafac2") + + nv = _NormalizedView.from_anndata(base_adata, include_raw=False) + sub = nv[0:2, :] + assert type(sub) is _NormalizedView + sub_x = sub.X + assert sub_x is not None + assert not isinstance(sub_x, VCSCArray) # the normalized view, not the raw array + + expected = VCSCArray.from_scipy(sp.csc_array(dense[0:2])).normalized("parafac2") + np.testing.assert_allclose(sub_x.toarray(), expected.toarray()) + + +# -- copy() -- the inherited anndata.AnnData.copy() only knows how to copy +# the standard private _X attribute, which this class never sets (X/raw_X +# live in _vcs_X/_vcs_raw_X instead), so it silently drops them. + + +def test_copy_preserves_x_and_raw_x(base_adata, dense): + va = VCSCAnnData.from_anndata(base_adata) + c = va.copy() + assert type(c) is VCSCAnnData + assert isinstance(c.X, VCSCArray) + assert isinstance(c.raw_X, VCSCArray) + np.testing.assert_allclose(c.X.toarray(), dense) + np.testing.assert_allclose(c.raw_X.toarray(), dense) + + +def test_copy_is_independent_of_the_original(base_adata, dense): + if dense.shape[0] == 0 or dense.shape[1] == 0: + pytest.skip("shape too small") + va = VCSCAnnData.from_anndata(base_adata, include_raw=False) + c = va.copy() + assert c.X is not va.X + assert c.obs is not va.obs + + c.obs["grp"] = "mutated" + assert list(va.obs["grp"]) != list(c.obs["grp"]) + + +def test_copy_preserves_obs_var_uns(base_adata, dense): + base_adata.uns["note"] = {"k": "v"} + va = VCSCAnnData.from_anndata(base_adata, include_raw=False) + c = va.copy() + assert list(c.obs["grp"]) == list(va.obs["grp"]) + assert list(c.var["gene"]) == list(va.var["gene"]) + assert c.uns["note"] == {"k": "v"} + + +def test_copy_preserves_a_normalized_x_subclass(base_adata, dense): + """A subclass overriding the `X` getter (as BAL-Pf2's own + lazy-normalized-view AnnData does) must round-trip through copy().""" + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + + class _NormalizedView(VCSCAnnData): + @property + def X(self): + return self.normalized("parafac2") + + nv = _NormalizedView.from_anndata(base_adata, include_raw=False) + c = nv.copy() + assert type(c) is _NormalizedView + c_x, nv_x = c.X, nv.X + assert c_x is not None + assert nv_x is not None + np.testing.assert_allclose(c_x.toarray(), nv_x.toarray()) + + +def test_copy_of_a_slice_round_trips_x(base_adata, dense): + """The exact pattern a caller like `parafac2`'s BiCV split uses: + `adata[obs_mask][:, var_mask].copy()`.""" + if dense.shape[0] < 2 or dense.shape[1] < 2: + pytest.skip("shape too small") + va = VCSCAnnData.from_anndata(base_adata, include_raw=False) + sub = va[0:2, 0:2] + c = sub.copy() + assert c.X is not None + np.testing.assert_allclose(c.X.toarray(), dense[0:2, 0:2]) + + +# -- to_memory() -- same underlying problem as copy(): the inherited +# anndata.AnnData.to_memory() doesn't know about _vcs_X/_vcs_raw_X either. + + +def test_to_memory_preserves_x_and_raw_x(base_adata, dense): + va = VCSCAnnData.from_anndata(base_adata) + m = va.to_memory() + assert type(m) is VCSCAnnData + assert isinstance(m.X, VCSCArray) + assert isinstance(m.raw_X, VCSCArray) + np.testing.assert_allclose(m.X.toarray(), dense) + np.testing.assert_allclose(m.raw_X.toarray(), dense) + + +def test_to_memory_preserves_a_normalized_x_subclass(base_adata, dense): + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + + class _NormalizedView(VCSCAnnData): + @property + def X(self): + return self.normalized("parafac2") + + nv = _NormalizedView.from_anndata(base_adata, include_raw=False) + m = nv.to_memory() + assert type(m) is _NormalizedView + m_x, nv_x = m.X, nv.X + assert m_x is not None + assert nv_x is not None + np.testing.assert_allclose(m_x.toarray(), nv_x.toarray())