add qsep score

Author: mffrank

grassp/plotting/__init__.py

@@ -1,4 +1,4 @@
 # from .heatmaps import grouped_heatmap
-from .heatmaps import protein_clustermap, sample_heatmap
+from .heatmaps import protein_clustermap, sample_heatmap, qsep_heatmap, qsep_boxplot
 from .integration import aligned_umap, remodeling_sankey, remodeling_score
 from .qc import bait_volcano_plots, highly_variable_proteins

grassp/plotting/heatmaps.py

@@ -203,3 +203,167 @@ def sample_heatmap(
 #         swap_axes=True,
 #         show=False,
 #     )
+
+
+def qsep_heatmap(
+    data: AnnData,
+    normalize: bool = True,
+    ax: plt.Axes = None,
+    cmap: str = "RdBu_r",
+    **kwargs,
+) -> plt.Axes:
+    """Plot QSep cluster distance heatmap.
+
+    Parameters
+    ----------
+    data : AnnData
+        Annotated data matrix containing QSep results.
+    normalize : bool, optional
+        If True, normalize distances by diagonal values.
+        Defaults to True.
+    ax : matplotlib.axes.Axes, optional
+        Axes to plot on. If None, current axes will be used.
+    cmap : str, optional
+        Colormap to use. Defaults to "RdBu_r".
+    **kwargs
+        Additional arguments passed to sns.heatmap.
+
+    Returns
+    -------
+    matplotlib.axes.Axes
+        The axes object with the plot.
+    """
+    if ax is None:
+        ax = plt.gca()
+
+    try:
+        distances = data.uns["cluster_distances"]["distances"]
+        clusters = data.uns["cluster_distances"]["clusters"]
+    except KeyError:
+        raise ValueError(
+            "Cluster distances not found in data.uns['cluster_distances'], run gr.tl.qsep_score first"
+        )
+
+    if normalize:
+        # Normalize by diagonal values
+        norm_distances = distances / np.diag(distances)[:, np.newaxis]
+        plot_data = norm_distances[::-1]
+        vmin = 1.0
+        vmax = np.max(norm_distances)
+    else:
+        plot_data = distances[::-1]
+        vmin = None
+        vmax = None
+
+    # Create heatmap
+    sns.heatmap(
+        plot_data,
+        xticklabels=clusters,
+        yticklabels=clusters[::-1],  # Reverse the y-axis labels
+        cmap=cmap,
+        vmin=vmin,
+        vmax=vmax,
+        ax=ax,
+        **kwargs,
+    )
+
+    ax.set_title("QSep Cluster Distances" + (" (Normalized)" if normalize else ""))
+
+    return ax
+
+
+def qsep_boxplot(
+    data: AnnData,
+    normalize: bool = True,
+    ax: plt.Axes = None,
+    palette: str = "Set2",
+    **kwargs,
+) -> plt.Axes:
+    """Plot QSep cluster distances as boxplots.
+
+    Parameters
+    ----------
+    data : AnnData
+        Annotated data matrix containing QSep results.
+    normalize : bool, optional
+        If True, normalize distances by diagonal values.
+        Defaults to True.
+    ax : matplotlib.axes.Axes, optional
+        Axes to plot on. If None, current axes will be used.
+    palette : str, optional
+        Color palette for the boxplots. Defaults to "Set2".
+    **kwargs
+        Additional arguments passed to sns.boxplot.
+
+    Returns
+    -------
+    matplotlib.axes.Axes
+        The axes object with the plot.
+    """
+    if ax is None:
+        ax = plt.gca()
+
+    try:
+        distances = data.uns["cluster_distances"]["distances"]
+        clusters = data.uns["cluster_distances"]["clusters"]
+    except KeyError:
+        raise ValueError(
+            "Cluster distances not found in data.uns['cluster_distances'], run gr.tl.qsep_score first"
+        )
+
+    if normalize:
+        # Normalize by diagonal values
+        distances = distances / np.diag(distances)[:, np.newaxis]
+
+    # Create DataFrame for plotting
+    plot_data = []
+    for i, ref_cluster in enumerate(clusters):
+        for j, target_cluster in enumerate(clusters):
+            plot_data.append(
+                {
+                    "Reference Cluster": ref_cluster,
+                    "Target Cluster": target_cluster,
+                    "Distance": distances[i, j],
+                    "color": "grey" if i == j else "red",
+                }
+            )
+    plot_df = pd.DataFrame(plot_data)
+    print(plot_df)
+
+    # Create boxplot
+    sns.boxplot(
+        data=plot_df,
+        x="Distance",
+        y="Reference Cluster",
+        color="grey",
+        orient="h",
+        ax=ax,
+        legend=False,
+        showfliers=False,
+        **kwargs,
+    )
+
+    # Add individual points
+    sns.stripplot(
+        data=plot_df,
+        x="Distance",
+        y="Reference Cluster",
+        hue="color",
+        # hue="Target Cluster",
+        orient="h",
+        size=4,
+        # color=".3",
+        alpha=0.6,
+        ax=ax,
+        legend=False,
+    )
+
+    # Customize plot
+    if normalize:
+        ax.axvline(x=1 if normalize else 0, color="gray", linestyle="--", alpha=0.5)
+    ax.set_xlabel("QSep Cluster Distances" + (" (Normalized)" if normalize else ""))
+
+    # Move legend outside
+    # ax.legend(bbox_to_anchor=(1.05, 1), loc="upper left", borderaxespad=0.0)
+
+    return ax

grassp/tools/__init__.py

@@ -7,4 +7,9 @@
 )
 from .enrichment import calculate_cluster_enrichment, rank_proteins_groups
 from .integration import align_adatas, aligned_umap, remodeling_score
-from .scoring import calinski_habarasz_score, silhouette_score
+from .scoring import (
+    calinski_habarasz_score,
+    class_balance,
+    qsep_score,
+    silhouette_score,
+)

grassp/tools/scoring.py

@@ -83,3 +83,78 @@ def calinski_habarasz_score(
         data.uns[key_added] = ch
     else:
         return ch
+
+
+def qsep_score(
+    data: AnnData,
+    label_key: str,
+    use_rep: str = "X",
+    distance_key: str = "full_distances",
+    inplace: bool = True,
+) -> None | np.ndarray:
+    """Calculate QSep scores for spatial proteomics data.
+
+    Parameters
+    ----------
+    data : AnnData
+        Annotated data matrix.
+    label_key : str
+        Key in data.obs containing cluster labels.
+    use_rep : str, optional
+        Key for representation to use for distance calculation.
+        Either 'X' or a key in data.obsm. Defaults to 'X'.
+    distance_key : str, optional
+        Key under which to store the full distances in data.obs.
+        Defaults to 'full_distances'.
+    inplace : bool, optional
+        If True, store results in data, else return matrices.
+        Defaults to True.
+
+    Returns
+    -------
+    None or np.ndarray
+        If inplace=True, returns None and stores results in data.
+        If inplace=False, returns cluster_distances.
+    """
+    # Get data matrix
+    if use_rep == "X":
+        X = data.X
+    else:
+        X = data.obsm[use_rep]
+
+    # Calculate pairwise distances between all points
+    full_distances = sklearn.metrics.pairwise_distances(X)
+
+    # Get valid clusters (non-NA)
+    mask = data.obs[label_key].notna()
+    valid_clusters = data.obs[label_key][mask].unique()
+
+    # Calculate cluster distances
+    cluster_distances = np.zeros((len(valid_clusters), len(valid_clusters)))
+    cluster_indices = {
+        cluster: np.where(data.obs[label_key] == cluster)[0]
+        for cluster in valid_clusters
+    }
+
+    for i, cluster1 in enumerate(valid_clusters):
+        for j, cluster2 in enumerate(valid_clusters):
+            idx1 = cluster_indices[cluster1]
+            idx2 = cluster_indices[cluster2]
+
+            # Get submatrix of distances between clusters
+            submatrix = full_distances[np.ix_(idx1, idx2)]
+            cluster_distances[i, j] = np.mean(submatrix)
+
+    if inplace:
+        # Store full distances
+        data.obs[distance_key] = pd.Series(
+            np.mean(full_distances, axis=1), index=data.obs.index
+        )
+
+        # Store cluster distances and metadata
+        data.uns["cluster_distances"] = {
+            "distances": cluster_distances,
+            "clusters": valid_clusters.tolist(),
+        }
+    else:
+        return cluster_distances