adding _annotation_mapping in AnalysisMixin

src/moscot/base/problems/_mixins.py

@@ -272,6 +272,47 @@ def _cell_transition_online(
             forward=forward,
         )
 
+    def _celltype_mapping(
+        self: AnalysisMixinProtocol[K, B],
+        mapping_mode: Literal["sum", "max"],
+        key: Optional[str],
+        source: K,
+        target: K,
+        source_groups: Str_Dict_t,
+        target_groups: Str_Dict_t,
+        forward: bool = False,  # return value will be row-stochastic if forward=True, else column-stochastic
+        aggregation_mode: Literal["annotation", "cell"] = "annotation",
+        other_key: Optional[str] = None,
+        other_adata: Optional[str] = None,
+        batch_size: Optional[int] = None,
+        normalize: bool = True,
+        scale_by_marginals: bool = True,
+    ):
+        if mapping_mode == "sum":
+            return self._cell_transition_online(
+                key=key,
+                source=source,
+                target=target,
+                source_groups=source_groups,
+                target_groups=target_groups,
+                forward=forward,
+                aggregation_mode=aggregation_mode,
+                other_key=other_key,
+                other_adata=other_adata,
+                batch_size=batch_size,
+                normalize=normalize,
+            )
+        if mapping_mode == "max":
+            source_annotation_key, source_annotations, source_annotations_ordered = _validate_args_cell_transition(
+                self.adata, source_groups
+            )
+            target_annotation_key, target_annotations, target_annotations_ordered = _validate_args_cell_transition(
+                self.adata if other_adata is None else other_adata, target_groups
+            )
+            dummy = pd.get_dummies(source_annotations)
+            out = self.pull(dummy, scale_by_marginals=scale_by_marginals)
+            return pd.Categorical([dummy.columns[i] for i in np.array(out.argmax(1))])
+
     def _sample_from_tmap(
         self: AnalysisMixinProtocol[K, B],
         source: K,

src/moscot/problems/space/_mixins.py

@@ -63,6 +63,13 @@ def _cell_transition(
     ) -> pd.DataFrame:
         ...
 
+    def _celltype_mapping(
+        self: AnalysisMixinProtocol[K, B],
+        *args: Any,
+        **kwargs: Any,
+    ) -> pd.DataFrame:
+        ...
+
 
 class SpatialMappingMixinProtocol(AnalysisMixinProtocol[K, B]):
     """Protocol class."""
@@ -82,6 +89,9 @@ def _filter_vars(
     def _cell_transition(self: AnalysisMixinProtocol[K, B], *args: Any, **kwargs: Any) -> pd.DataFrame:
         ...
 
+    def _celltype_mapping(self: AnalysisMixinProtocol[K, B], *args: Any, **kwargs: Any) -> pd.DataFrame:
+        ...
+
 
 class SpatialAlignmentMixin(AnalysisMixin[K, B]):
     """Spatial alignment mixin class."""
@@ -273,6 +283,38 @@ def cell_transition(  # type: ignore[misc]
             key_added=key_added,
         )
 
+    def celltype_mapping(
+        self: SpatialAlignmentMixinProtocol[K, B],
+        mapping_mode: Literal["sum", "max"],
+        key: Optional[str],
+        source: K,
+        target: K,
+        source_groups: Str_Dict_t,
+        target_groups: Str_Dict_t,
+        forward: bool = False,  # return value will be row-stochastic if forward=True, else column-stochastic
+        aggregation_mode: Literal["annotation", "cell"] = "annotation",
+        other_key: Optional[str] = None,
+        other_adata: Optional[str] = None,
+        batch_size: Optional[int] = None,
+        normalize: bool = True,
+        scale_by_marginals: bool = True,
+    ) -> pd.DataFrame:
+        return self._celltype_mapping(
+            mapping_mode=mapping_mode,
+            key=key,
+            source=source,
+            target=target,
+            source_groups=source_groups,
+            target_groups=target_groups,
+            forward=forward,
+            aggregation_mode=aggregation_mode,
+            other_key=other_key,
+            other_adata=other_adata,
+            batch_size=batch_size,
+            normalize=normalize,
+            scale_by_marginals=scale_by_marginals,
+        )
+
     @property
     def spatial_key(self) -> Optional[str]:
         """Spatial key in :attr:`~anndata.AnnData.obsm`."""
@@ -562,6 +604,38 @@ def cell_transition(  # type: ignore[misc]
             key_added=key_added,
         )
 
+    def celltype_mapping(
+        self: SpatialMappingMixinProtocol[K, B],
+        mapping_mode: Literal["sum", "max"],
+        key: Optional[str],
+        source: K,
+        target: K,
+        source_groups: Str_Dict_t,
+        target_groups: Str_Dict_t,
+        forward: bool = False,  # return value will be row-stochastic if forward=True, else column-stochastic
+        aggregation_mode: Literal["annotation", "cell"] = "annotation",
+        other_key: Optional[str] = None,
+        other_adata: Optional[str] = None,
+        batch_size: Optional[int] = None,
+        normalize: bool = True,
+        scale_by_marginals: bool = True,
+    ) -> pd.DataFrame:
+        return self._celltype_mapping(
+            mapping_mode=mapping_mode,
+            key=key,
+            source=source,
+            target=target,
+            source_groups=source_groups,
+            target_groups=target_groups,
+            forward=forward,
+            aggregation_mode=aggregation_mode,
+            other_key=other_key,
+            other_adata=other_adata,
+            batch_size=batch_size,
+            normalize=normalize,
+            scale_by_marginals=scale_by_marginals,
+        )
+
     @property
     def batch_key(self) -> Optional[str]:
         """Batch key in :attr:`~anndata.AnnData.obs`."""