Fix VCSCAnnData.copy() silently dropping X/raw_X

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:

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