make confusion matrix diaply more flexible

Author: Marius1311

src/cellmapper/model/evaluate.py

@@ -7,13 +7,7 @@
 from scipy.sparse import issparse
 from scipy.spatial.distance import jensenshannon
 from scipy.stats import pearsonr, spearmanr
-from sklearn.metrics import (
-    accuracy_score,
-    classification_report,
-    f1_score,
-    precision_score,
-    recall_score,
-)
+from sklearn.metrics import accuracy_score, classification_report, f1_score, precision_score, recall_score
 
 from cellmapper._docs import d
 from cellmapper.logging import logger
@@ -386,6 +380,8 @@ def evaluate_label_transfer(
     def plot_confusion_matrix(
         self,
         label_key: str,
+        *,
+        true_key: str | None = None,
         subset: np.ndarray | pd.Series | None = None,
         figsize: tuple[int, int] = (10, 8),
         cmap: str = "viridis",
@@ -412,7 +408,11 @@ def plot_confusion_matrix(
         Parameters
         ----------
         label_key
-            Key in .obs storing ground-truth cell type annotations.
+            Key in .obs storing predicted labels (from map_obs). The column
+            ``f"{label_key}{prediction_postfix}"`` is used as the x-axis (predicted).
+        true_key
+            Key in .obs to use for the y-axis (true labels). If None, uses ``label_key``.
+            This allows comparing arbitrary columns, e.g., source_time vs mapped_time.
         subset
             Boolean mask to select a subset of cells for the confusion matrix.
             Must have the same length as query.obs or be a pandas Series indexed by obs_names.
@@ -466,7 +466,8 @@ def plot_confusion_matrix(
             raise ValueError("Label transfer has not been performed. Call map_obs() first.")
 
         # Extract true and predicted labels
-        y_true = self.query.obs[label_key].copy()
+        true_col = true_key if true_key is not None else label_key
+        y_true = self.query.obs[true_col].copy()
         y_pred = self.query.obs[f"{label_key}{self.prediction_postfix}"].copy()
 
         # Drop NaNs