fix deleting modalities

Author: ilia-kats

src/mudata/_core/mudata.py

@@ -452,6 +452,8 @@ def _check_changed_attr_names(self, attr: str, columns: bool = False):
         attr_names_changed, attr_columns_changed = False, False
         if not hasattr(self, attrhash):
             attr_names_changed, attr_columns_changed = True, True
+        elif len(self.mod) < len(getattr(self, attrhash)):
+            attr_names_changed, attr_columns_changed = True, None
         else:
             for m, mod in self.mod.items():
                 if m in getattr(self, attrhash):
@@ -607,7 +609,13 @@ def _update_attr(
         _attrhash = f"_{attr}hash"
         attr_changed = self._check_changed_attr_names(attr)
 
-        attr_duplicated = self._check_duplicated_attr_names(attr)
+        if not any(attr_changed):
+            # Nothing to update
+            return
+
+        data_global = getattr(self, attr)
+
+        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:
@@ -619,12 +627,6 @@ def _update_attr(
                     f"Behaviour is not defined with axis=-1, {attr}_names need to be made unique first."
                 )
 
-        if not any(attr_changed):
-            # Nothing to update
-            return
-
-        data_global = getattr(self, attr)
-
         # Generate unique colnames
         (rowcol,) = self._find_unique_colnames(attr, 1)
 
@@ -646,7 +648,6 @@ def _update_attr(
         # Join modality .obs/.var tables
         #
         # Main case: no duplicates and no intersection if the axis is not shared
-        #
         if not attr_duplicated:
             # Shared axis
             if axis == (1 - self._axis) or self._axis == -1:
@@ -665,7 +666,7 @@ def _update_attr(
             data_mod = _make_index_unique(data_mod, force=attr_intersecting)
             data_global = _make_index_unique(data_global, force=attr_intersecting)
             if data_global.shape[1] > 0:
-                data_mod = data_global.join(data_mod, how="left", sort=False)
+                data_mod = data_mod.join(data_global, how="left", sort=False)
 
             if data_global.shape[0] > 0:
                 # reorder new index to conform to the old index as much as possible
@@ -728,8 +729,11 @@ def _update_attr(
                     )
                 # after inserting a new modality with duplicates, but no duplicates before:
                 # data_mod.index is not unique
-                data_global = _make_index_unique(data_global, force=not data_mod.index.is_unique)
-                data_mod = _make_index_unique(data_mod)
+                # after deleting a modality with duplicates: data_global.index is not unique, but
+                # data_mod.index is unique
+                need_unique = data_mod.index.is_unique | data_global.index.is_unique
+                data_global = _make_index_unique(data_global, force=need_unique)
+                data_mod = _make_index_unique(data_mod, force=need_unique)
                 data_mod = data_mod.join(data_global, how="left", sort=False)
 
                 # reorder new index to conform to the old index as much as possible
@@ -750,6 +754,10 @@ def _update_attr(
                     > data_global.shape[0]  # new modality added and concacenated
                 )
 
+                if need_unique:
+                    data_mod = _restore_index(data_mod)
+                    data_global = _restore_index(data_global)
+
             data_mod.reset_index(level=list(range(1, data_mod.index.nlevels)), inplace=True)
             data_global.reset_index(level=list(range(1, data_global.index.nlevels)), inplace=True)
             data_mod.index.set_names(None, inplace=True)

tests/test_update.py

@@ -161,6 +161,38 @@ def test_update_add_modality(self, modalities, axis):
                 == old_oattrnames.append(getattr(modalities[modnames[i]], f"{oattr}_names"))
             ).all()
 
+    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"
+
+        fullbatch = getattr(mdata, attr)["batch"]
+        fullobatch = getattr(mdata, oattr)["batch"]
+        keptmask = (getattr(mdata, f"{attr}map")[modnames[1]].reshape(-1) > 0) | (
+            getattr(mdata, f"{attr}map")[modnames[2]].reshape(-1) > 0
+        )
+        keptomask = (getattr(mdata, f"{oattr}map")[modnames[1]].reshape(-1) > 0) | (
+            getattr(mdata, f"{oattr}map")[modnames[2]].reshape(-1) > 0
+        )
+
+        del mdata.mod[modnames[0]]
+        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()
+
+        fullbatch = getattr(mdata, attr)["batch"]
+        fullobatch = getattr(mdata, oattr)["batch"]
+        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()
+
     def test_update_intersecting(self, modalities, axis):
         """
         Update should work when