more comprehensive tests and fixes for _update_attr()

Author: ilia-kats

src/mudata/_core/mudata.py

@@ -614,19 +614,11 @@ def _update_attr(
             return
 
         data_global = getattr(self, attr)
+        prev_index = data_global.index
 
         attr_duplicated = not data_global.index.is_unique or self._check_duplicated_attr_names(attr)
         attr_intersecting = self._check_intersecting_attr_names(attr)
 
-        if attr_duplicated:
-            warnings.warn(
-                f"{attr}_names are not unique. To make them unique, call `.{attr}_names_make_unique`."
-            )
-            if self._axis == -1:
-                warnings.warn(
-                    f"Behaviour is not defined with axis=-1, {attr}_names need to be made unique first."
-                )
-
         # Generate unique colnames
         (rowcol,) = self._find_unique_colnames(attr, 1)
 
@@ -708,16 +700,19 @@ def _update_attr(
                 col.replace(np.nan, 0, inplace=True)
                 col = col.astype(np.uint32)
                 data_mod[colname] = col
-                if mod in attrmap and (
-                    col.shape[0] != data_global.shape[0]
-                    and np.sum(attrmap[mod] > 0)
-                    == getattr(amod, attr).shape[0]  # added/removed observations
-                    or col.shape[0] == data_global.shape[0]
-                    and np.array_equal(attrmap[mod], col)  # reordered
-                ):
-                    data_mod.set_index(colname, append=True, inplace=True)
-                    data_global.set_index(attrmap[mod].reshape(-1), append=True, inplace=True)
-                    data_global.index.set_names(colname, level=-1, inplace=True)
+                if mod in attrmap:
+                    modmap = attrmap[mod].reshape(-1)
+                    modmask = modmap > 0
+                    # only use unchanged modalities for ordering
+                    if (
+                        modmask.sum() == getattr(amod, attr).shape[0]
+                        and (
+                            getattr(amod, attr).index[modmap[modmask] - 1] == prev_index[modmask]
+                        ).all()
+                    ):
+                        data_mod.set_index(colname, append=True, inplace=True)
+                        data_global.set_index(attrmap[mod].reshape(-1), append=True, inplace=True)
+                        data_global.index.set_names(colname, level=-1, inplace=True)
 
             if data_global.shape[0] > 0:
                 if not data_global.index.is_unique:
@@ -749,9 +744,11 @@ def _update_attr(
                     == data_global.shape[
                         0
                     ]  # renamed (since new_idx.shape[0] > 0 and kept_idx.shape[0] < data_global.shape[0])
-                    or axis == self._axis
-                    and data_mod.shape[0]
-                    > data_global.shape[0]  # new modality added and concacenated
+                    or (
+                        axis == self._axis
+                        and axis != -1
+                        and data_mod.shape[0] > data_global.shape[0]
+                    )  # new modality added and concacenated
                 )
 
                 if need_unique:
@@ -770,6 +767,15 @@ def _update_attr(
             mdict[m] = data_mod[colname].to_numpy()
             data_mod.drop(colname, axis=1, inplace=True)
 
+        if not data_mod.index.is_unique:
+            warnings.warn(
+                f"{attr}_names are not unique. To make them unique, call `.{attr}_names_make_unique`."
+            )
+            if self._axis == -1:
+                warnings.warn(
+                    f"Behaviour is not defined with axis=-1, {attr}_names need to be made unique first."
+                )
+
         setattr(
             self,
             "_" + attr,

tests/test_update.py

@@ -34,15 +34,16 @@ def modalities(request, obs_n, obs_across, obs_mod):
 
     if obs_mod:
         if obs_mod == "duplicated":
-            for m in ["mod1", "mod2"]:
-                # Index does not support mutable operations
-                obs_names = mods[m].obs_names.values.copy()
-                obs_names[1] = obs_names[0]
-                mods[m].obs_names = obs_names
-
-                var_names = mods[m].var_names.values.copy()
-                var_names[1] = var_names[0]
-                mods[m].var_names = var_names
+            obsnames2 = mods["mod2"].obs_names.to_numpy()
+            obsnames3 = mods["mod3"].obs_names.to_numpy()
+            varnames2 = mods["mod2"].var_names.to_numpy()
+            varnames3 = mods["mod3"].var_names.to_numpy()
+            obsnames2[0] = obsnames2[1] = obsnames3[1] = "testobs"
+            varnames2[0] = varnames2[1] = varnames3[1] = "testvar"
+            mods["mod2"].obs_names = obsnames2
+            mods["mod3"].obs_names = obsnames3
+            mods["mod2"].var_names = varnames2
+            mods["mod3"].var_names = varnames3
 
     return mods
 
@@ -64,8 +65,8 @@ def mdata(modalities, axis):
     md.obs["batch"] = np.random.choice(["a", "b", "c"], size=md.shape[0], replace=True)
     md.var["batch"] = np.random.choice(["d", "e", "f"], size=md.shape[1], replace=True)
 
-    md.obsm["test_obsm"] = np.random.normal(size=(md.n_obs, 2))
-    md.varm["test_varm"] = np.random.normal(size=(md.n_var, 2))
+    md.obsm["test"] = np.random.normal(size=(md.n_obs, 2))
+    md.varm["test"] = np.random.normal(size=(md.n_var, 2))
 
     return md
 
@@ -82,6 +83,14 @@ def new_update(self):
         yield
         set_options(pull_on_update=None)
 
+    @staticmethod
+    def get_attrm_values(mdata, attr, key, names):
+        attrm = getattr(mdata, f"{attr}m")
+        index = getattr(mdata, f"{attr}_names")
+        return np.concatenate(
+            [np.atleast_1d(attrm[key][np.nonzero(index == name)[0]]) for name in names]
+        )
+
     def test_update_simple(self, mdata, axis):
         """
         Update should work when
@@ -134,9 +143,25 @@ def test_update_add_modality(self, modalities, axis):
             old_attrnames = getattr(mdata, f"{attr}_names")
             old_oattrnames = getattr(mdata, f"{oattr}_names")
 
+            some_obs_names = mdata.obs_names[:2]
+            mdata.obsm["test"] = np.random.normal(size=(mdata.n_obs, 1))
+            true_obsm_values = self.get_attrm_values(mdata, "obs", "test", some_obs_names)
+
             mdata.mod[modnames[i]] = modalities[modnames[i]]
             mdata.update()
 
+            test_obsm_values = self.get_attrm_values(mdata, "obs", "test", some_obs_names)
+            if axis == 1:
+                assert np.isnan(mdata.obsm["test"]).sum() == modalities[modnames[i]].n_obs
+                assert np.all(np.isnan(mdata.obsm["test"][-modalities[modnames[i]].n_obs :]))
+                assert np.all(~np.isnan(mdata.obsm["test"][: -modalities[modnames[i]].n_obs]))
+                assert (
+                    test_obsm_values[~np.isnan(test_obsm_values)].reshape(-1)
+                    == true_obsm_values.reshape(-1)
+                ).all()
+            else:
+                assert (test_obsm_values == true_obsm_values).all()
+
             attrnames = getattr(mdata, f"{attr}_names")
             oattrnames = getattr(mdata, f"{oattr}_names")
             assert (attrnames[: old_attrnames.size] == old_attrnames).all()
@@ -157,9 +182,13 @@ def test_update_delete_modality(self, mdata, axis):
         modnames = list(mdata.mod.keys())
         attr = "obs" if axis == 0 else "var"
         oattr = "var" if axis == 0 else "obs"
+        attrm = f"{attr}m"
+        oattrm = f"{oattr}m"
 
         fullbatch = getattr(mdata, attr)["batch"]
         fullobatch = getattr(mdata, oattr)["batch"]
+        fulltestm = getattr(mdata, attrm)["test"]
+        fullotestm = getattr(mdata, oattrm)["test"]
         keptmask = (getattr(mdata, f"{attr}map")[modnames[1]].reshape(-1) > 0) | (
             getattr(mdata, f"{attr}map")[modnames[2]].reshape(-1) > 0
         )
@@ -171,19 +200,26 @@ def test_update_delete_modality(self, mdata, axis):
         mdata.update()
 
         assert mdata.shape[1 - axis] == sum(mod.shape[1 - axis] for mod in mdata.mod.values())
-        assert (getattr(mdata, oattr)["batch"] == fullobatch[keptomask]).all()
         assert (getattr(mdata, attr)["batch"] == fullbatch[keptmask]).all()
+        assert (getattr(mdata, oattr)["batch"] == fullobatch[keptomask]).all()
+        assert (getattr(mdata, attrm)["test"] == fulltestm[keptmask, :]).all()
+        assert (getattr(mdata, oattrm)["test"] == fullotestm[keptomask, :]).all()
 
         fullbatch = getattr(mdata, attr)["batch"]
         fullobatch = getattr(mdata, oattr)["batch"]
+        fulltestm = getattr(mdata, attrm)["test"]
+        fullotestm = getattr(mdata, oattrm)["test"]
         keptmask = getattr(mdata, f"{attr}map")[modnames[1]].reshape(-1) > 0
         keptomask = getattr(mdata, f"{oattr}map")[modnames[1]].reshape(-1) > 0
+
         del mdata.mod[modnames[2]]
         mdata.update()
 
         assert mdata.shape[1 - axis] == sum(mod.shape[1 - axis] for mod in mdata.mod.values())
         assert (getattr(mdata, oattr)["batch"] == fullobatch[keptomask]).all()
         assert (getattr(mdata, attr)["batch"] == fullbatch[keptmask]).all()
+        assert (getattr(mdata, attrm)["test"] == fulltestm[keptmask, :]).all()
+        assert (getattr(mdata, oattrm)["test"] == fullotestm[keptomask, :]).all()
 
     def test_update_intersecting(self, modalities, axis):
         """
@@ -216,22 +252,14 @@ def test_update_intersecting(self, modalities, axis):
             )
         ).all()
 
-        # names along axis are intersected
+        # names along axis are unioned
         axisnames = reduce(
             lambda x, y: x.union(y, sort=False),
             (getattr(mod, f"{attr}_names") for mod in modalities.values()),
         )
         assert mdata.shape[axis] == axisnames.shape[0]
         assert (getattr(mdata, f"{attr}_names") == axisnames).all()
 
-        # Variables are different across modalities
-        for m, mod in modalities.items():
-            # Columns are intact in individual modalities
-            assert "mod" in mod.obs.columns
-            assert all(mod.obs["mod"] == m)
-            assert "mod" in mod.var.columns
-            assert all(mod.var["mod"] == m)
-
     def test_update_after_filter_obs_adata(self, mdata, axis):
         """
         Check for muon issue #44.
@@ -242,6 +270,14 @@ def test_update_after_filter_obs_adata(self, mdata, axis):
         old_obsnames = mdata.obs_names
         old_varnames = mdata.var_names
 
+        filtermask = mdata["mod3"].obs["min_count"] < -2
+        fullfiltermask = mdata.obsmap["mod3"].copy() > 0
+        fullfiltermask[fullfiltermask] = filtermask
+        keptmask = (mdata.obsmap["mod1"] > 0) | (mdata.obsmap["mod2"] > 0) | fullfiltermask
+
+        some_obs_names = mdata[keptmask, :].obs_names.values[:2]
+        true_obsm_values = self.get_attrm_values(mdata[keptmask], "obs", "test", some_obs_names)
+
         mdata.mod["mod3"] = mdata["mod3"][mdata["mod3"].obs["min_count"] < -2].copy()
         mdata.update()
         assert mdata.obs["batch"].isna().sum() == 0
@@ -251,28 +287,23 @@ def test_update_after_filter_obs_adata(self, mdata, axis):
             # check if the order is preserved
             assert (mdata.obs_names == old_obsnames[old_obsnames.isin(mdata.obs_names)]).all()
 
+        test_obsm_values = self.get_attrm_values(mdata, "obs", "test", some_obs_names)
+        assert (true_obsm_values == test_obsm_values).all()
+
     def test_update_after_obs_reordered(self, mdata):
         """
         Update should work if obs are reordered.
         """
         some_obs_names = mdata.obs_names.values[:2]
 
-        true_obsm_values = [
-            mdata.obsm["test_obsm"][np.where(mdata.obs_names.values == name)[0][0]]
-            for name in some_obs_names
-        ]
+        true_obsm_values = self.get_attrm_values(mdata, "obs", "test", some_obs_names)
 
         mdata.mod["mod1"] = mdata["mod1"][::-1].copy()
         mdata.update()
 
-        test_obsm_values = [
-            mdata.obsm["test_obsm"][np.where(mdata.obs_names == name)[0][0]]
-            for name in some_obs_names
-        ]
+        test_obsm_values = self.get_attrm_values(mdata, "obs", "test", some_obs_names)
 
-        assert all(
-            [all(true_obsm_values[i] == test_obsm_values[i]) for i in range(len(true_obsm_values))]
-        )
+        assert (true_obsm_values == test_obsm_values).all()
 
 
 @pytest.mark.usefixtures("filepath_h5mu")