add new feature

Author: hutaobo

src/pyXenium/analysis/differential.py

@@ -0,0 +1,93 @@
+# pyXenium/analysis/differential.py
+
+import pandas as pd
+import numpy as np
+from scipy.stats import ttest_ind, spearmanr
+from statsmodels.stats import multitest
+
+def get_rna_expr_df(adata, layer_key="rna"):
+    expr = adata.layers.get(layer_key)
+    if expr is None:
+        raise KeyError(f"adata.layers does not have {layer_key}")
+    try:
+        arr = expr.toarray()
+    except:
+        arr = expr
+    return pd.DataFrame(arr, index=adata.obs.index, columns=adata.var.index)
+
+def analyze_one_score(adata, rna_expr, cluster, protein, cluster_key="rna_cluster", score_prefix="score", min_cells=3):
+    score_col = f"{score_prefix}:{cluster}:{protein}"
+    if score_col not in adata.obs.columns:
+        return None
+
+    mask = (adata.obs[cluster_key] == cluster)
+    cells = adata.obs.index[mask]
+    if len(cells) < min_cells:
+        return None
+
+    s = adata.obs.loc[cells, score_col].astype(float)
+    median = s.median()
+    high = s[s >= median].index
+    low  = s[s <  median].index
+    if len(high) < min_cells or len(low) < min_cells:
+        return None
+
+    # 差异表达
+    de_res = []
+    for gene in rna_expr.columns:
+        gh = rna_expr.loc[high, gene].dropna()
+        gl = rna_expr.loc[low, gene].dropna()
+        if len(gh) < 3 or len(gl) < 3:
+            continue
+        t, p = ttest_ind(gh, gl, equal_var=False)
+        de_res.append((gene, t, p, gh.mean() - gl.mean()))
+    de_df = pd.DataFrame(de_res, columns=["gene","tstat","pval","mean_diff"])
+    if not de_df.empty:
+        de_df["adj_pval"] = multitest.multipletests(de_df["pval"], method="fdr_bh")[1]
+
+    # 相关性
+    corr_res = []
+    for gene in rna_expr.columns:
+        x = rna_expr.loc[cells, gene].fillna(0).values
+        y = s.values
+        r, p = spearmanr(x, y)
+        corr_res.append((gene, r, p))
+    corr_df = pd.DataFrame(corr_res, columns=["gene","spearman_r","pval"])
+    if not corr_df.empty:
+        corr_df["adj_pval"] = multitest.multipletests(corr_df["pval"], method="fdr_bh")[1]
+
+    return {
+        "cluster": cluster,
+        "protein": protein,
+        "n_cells": len(cells),
+        "de": de_df,
+        "corr": corr_df
+    }
+
+def run_all_clusters_proteins(adata, rna_expr, cluster_label, protein_names, score_prefix="score", min_cells=3):
+    results = []
+    for cl in adata.obs[cluster_label].unique():
+        for p in protein_names:
+            rec = analyze_one_score(
+                adata, rna_expr, cl, p,
+                cluster_key=cluster_label, score_prefix=score_prefix, min_cells=min_cells
+            )
+            if rec is not None:
+                results.append(rec)
+    return results
+
+def summarize_results(results):
+    # 把结果字典列表拆成两个 DataFrame
+    de_list = []
+    corr_list = []
+    for rec in results:
+        c = rec["cluster"]; p = rec["protein"]
+        df_de = rec["de"].copy()
+        df_de["cluster"] = c; df_de["protein"] = p
+        de_list.append(df_de)
+        df_corr = rec["corr"].copy()
+        df_corr["cluster"] = c; df_corr["protein"] = p
+        corr_list.append(df_corr)
+    all_de = pd.concat(de_list, ignore_index=True) if de_list else pd.DataFrame()
+    all_corr = pd.concat(corr_list, ignore_index=True) if corr_list else pd.DataFrame()
+    return all_de, all_corr

src/pyXenium/analysis/plotting.py

@@ -0,0 +1,71 @@
+# pyXenium/analysis/plotting.py
+
+import seaborn as sns
+import matplotlib.pyplot as plt
+import pandas as pd
+
+def plot_auc_heatmap(summary: pd.DataFrame, figsize=(10,8)):
+    mat = summary.pivot(index="cluster", columns="protein", values="test_auc")
+    mat = mat.apply(pd.to_numeric, errors="coerce")
+    g = sns.clustermap(mat, cmap="viridis", linewidths=.3, figsize=figsize)
+    g.ax_heatmap.set_xlabel("Protein"); g.ax_heatmap.set_ylabel("Cluster")
+    return g
+
+def plot_topk_per_cluster(summary: pd.DataFrame, k=5, metric="test_auc"):
+    topk = (summary.sort_values(["cluster", metric], ascending=[True, False])
+                   .groupby("cluster").head(k))
+    fig, ax = plt.subplots(figsize=(max(10, k * 1.2), 6))
+    labels = []
+    vals = []
+    for cl, sub in topk.groupby("cluster"):
+        for _, r in sub.iterrows():
+            labels.append(f"{cl}:{r['protein']}")
+            vals.append(r[metric])
+    ax.bar(labels, vals)
+    ax.set_ylabel(metric)
+    ax.set_xticklabels(labels, rotation=90)
+    plt.tight_layout()
+    return fig
+
+def plot_DE_volcano(de_df: pd.DataFrame, title="DE Volcano",
+                    logfc_col="mean_diff", pval_col="pval", adj_col="adj_pval",
+                    fdr_thresh=0.05):
+    df = de_df.copy()
+    df["-log10p"] = -np.log10(df[pval_col])
+    plt.figure(figsize=(6,5))
+    sns.scatterplot(data=df, x=logfc_col, y="-log10p",
+                    hue=df[adj_col] < fdr_thresh,
+                    palette={True: "red", False: "gray"}, legend=False)
+    plt.axhline(-np.log10(0.05), ls="--", color="black")
+    plt.title(title)
+    plt.xlabel("Mean difference (High vs Low)")
+    plt.ylabel("-log10(p)")
+    plt.tight_layout()
+    plt.show()
+
+def plot_model_diagnostics(adata, models, cluster, protein, feature_key="X_rna_pca"):
+    from sklearn.metrics import RocCurveDisplay, PrecisionRecallDisplay
+    from sklearn.calibration import calibration_curve
+
+    res = models[cluster][protein]
+    clf, scaler = res.model, res.scaler
+    thr = getattr(res, "threshold", None)
+
+    mask = (adata.obs["rna_cluster"] == cluster)
+    X = scaler.transform(adata.obsm[feature_key][mask, :])
+    # y 真值需要你自己定义：可能 adata.obs[f"protein:{protein}"] ≥ thr
+    y = (adata.obs.loc[mask, f"protein:{protein}"] >= thr).astype(int).to_numpy()
+    y_prob = clf.predict_proba(X)[:, 1]
+
+    RocCurveDisplay.from_predictions(y, y_prob)
+    plt.title(f"ROC — {cluster}:{protein}")
+    PrecisionRecallDisplay.from_predictions(y, y_prob)
+    plt.title(f"PR — {cluster}:{protein}")
+    prob_true, prob_pred = calibration_curve(y, y_prob, n_bins=10, strategy="quantile")
+    plt.figure()
+    plt.plot(prob_pred, prob_true, marker="o")
+    plt.plot([0,1],[0,1], "--")
+    plt.xlabel("Predicted prob"); plt.ylabel("Empirical freq")
+    plt.title(f"Calibration — {cluster}:{protein}")
+    plt.tight_layout()
+    plt.show()

src/pyXenium/analysis/scoring.py

@@ -0,0 +1,38 @@
+# pyXenium/analysis/scoring.py
+
+import numpy as np
+import pandas as pd
+
+def write_model_scores(adata, models, feature_key="X_rna_pca", cluster_key="rna_cluster"):
+    """
+    对 models 中每个 (cluster, protein) 模型，在对应簇的细胞上预测概率，
+    并把结果写入 adata.obs 作为 score:cluster:protein 列。
+    """
+    # 为避免 obs 碎片化，建议先收集所有新列数据，再一次性 assign
+    new_cols = {}
+    for cluster, protodict in models.items():
+        mask = (adata.obs[cluster_key] == cluster)
+        if mask.sum() == 0:
+            continue
+        X_all = adata.obsm.get(feature_key)
+        if X_all is None:
+            raise KeyError(f"Feature key {feature_key} not in adata.obsm.")
+        X_sub = X_all[mask, :]
+        idx = adata.obs.index[mask]
+
+        for protein, res in protodict.items():
+            clf = res.model
+            scaler = res.scaler
+            X_scaled = scaler.transform(X_sub)
+            y_prob = clf.predict_proba(X_scaled)[:, 1]
+            col_name = f"score:{cluster}:{protein}"
+            # 创建一个全体 NaN 列，然后填入子集
+            col_ser = pd.Series(np.nan, index=adata.obs.index)
+            col_ser.loc[idx] = y_prob
+            new_cols[col_name] = col_ser
+
+    # 批量添加到 adata.obs
+    for col_name, col_ser in new_cols.items():
+        adata.obs[col_name] = col_ser
+
+    return adata