adding _annotation_mapping in AnalysisMixin

src/moscot/base/problems/_mixins.py

@@ -151,7 +151,6 @@ def _cell_transition(
             )
         if aggregation_mode == "cell" and source_groups is None and target_groups is None:
             raise ValueError("At least one of `source_groups` and `target_group` must be specified.")
-
         _check_argument_compatibility_cell_transition(
             source_annotation=source_groups,
             target_annotation=target_groups,
@@ -208,13 +207,13 @@ def _cell_transition_online(
         )
         df_source = _get_df_cell_transition(
             self.adata,
-            [source_annotation_key, target_annotation_key],
+            [source_annotation_key],
             key,
             source,
         )
         df_target = _get_df_cell_transition(
             self.adata if other_adata is None else other_adata,
-            [source_annotation_key, target_annotation_key],
+            [target_annotation_key],
             key if other_adata is None else other_key,
             target,
         )
@@ -323,7 +322,7 @@ def _annotation_mapping(
         cell_transition_kwargs.setdefault("forward", forward)
         if forward:
             source_df = _get_df_cell_transition(
-                self.adata,
+                self.adata if other_adata is None else other_adata,
                 annotation_keys=[annotation_label],
                 filter_key=key,
                 filter_value=source,
@@ -332,9 +331,9 @@ def _annotation_mapping(
             axis = 1  # columns
             cell_transition_kwargs.setdefault("source_groups", None)
             cell_transition_kwargs.setdefault("target_groups", annotation_label)
-        elif not forward:
+        else:
             target_df = _get_df_cell_transition(
-                self.adata if other_adata is None else other_adata,
+                self.adata,
                 annotation_keys=[annotation_label],
                 filter_key=key,
                 filter_value=target,
@@ -346,7 +345,7 @@ def _annotation_mapping(
         if mapping_mode == "sum":
             out: pd.DataFrame = self._cell_transition(**cell_transition_kwargs)
             return out.idxmax(axis=axis).to_frame(name=annotation_label)
-        elif mapping_mode == "max":
+        if mapping_mode == "max":
             if forward:
                 out: ArrayLike = self[(source, target)].push(dummy, scale_by_marginals=scale_by_marginals)
             else:

src/moscot/base/problems/_utils.py

@@ -107,7 +107,7 @@ def _check_argument_compatibility_cell_transition(
         raise ValueError("Unable to infer distributions, missing `adata` and `key`.")
     if forward and target_annotation is None:
         raise ValueError("No target annotation provided.")
-    if not forward and source_annotation is None:
+    if aggregation_mode == "annotation" and (not forward and source_annotation is None):
         raise ValueError("No source annotation provided.")
     if (aggregation_mode == "annotation") and (source_annotation is None or target_annotation is None):
         raise ValueError(

src/moscot/problems/space/_mixins.py

@@ -605,12 +605,19 @@ def annotation_mapping(
         scale_by_marginals: bool = True,
         cell_transition_kwargs: Mapping[str, Any] = types.MappingProxyType({}),
     ) -> pd.DataFrame:
+        cell_transition_kwargs = dict(cell_transition_kwargs)
+        if forward:
+            cell_transition_kwargs.setdefault("source_groups", None)
+            cell_transition_kwargs.setdefault("target_groups", annotation_label)
+        else:
+            cell_transition_kwargs.setdefault("source_groups", None)
+            cell_transition_kwargs.setdefault("target_groups", annotation_label)
         annotation = self._annotation_mapping(
             mapping_mode=mapping_mode,
             annotation_label=annotation_label,
             source=source,
             target=target,
-            forward=forward,
+            forward=not forward,
             key=self.batch_key,
             other_adata=self.adata_sc,
             scale_by_marginals=scale_by_marginals,