From 152933fef9193ac3fd3016a0cc58e877db47f231 Mon Sep 17 00:00:00 2001
From: Andrew Ramirez <aramirez@aretha.seas.ucla.edu>
Date: Fri, 6 Jun 2025 14:55:25 -0700
Subject: [PATCH] Add nC19 scores plot breaking down by viral/bacterial

---
 pf2/figures/figureS3a.py                  | 47 +++++++++++++++++++++++
 pf2/figures/{figureS3.py => figureS3b.py} |  0
 2 files changed, 47 insertions(+)
 create mode 100644 pf2/figures/figureS3a.py
 rename pf2/figures/{figureS3.py => figureS3b.py} (100%)

diff --git a/pf2/figures/figureS3a.py b/pf2/figures/figureS3a.py
new file mode 100644
index 0000000..9ecdf2d
--- /dev/null
+++ b/pf2/figures/figureS3a.py
@@ -0,0 +1,47 @@
+"""
+Figure S3: PLSR scores for patients with non-COVID-19 conditions
+"""
+
+import pandas as pd
+import anndata
+import seaborn as sns
+from ..data_import import condition_factors_meta
+from ..predict import plsr_acc
+from .common import subplotLabel, getSetup
+
+def makeFigure():
+    """Get a list of the axis objects and create a figure."""
+    ax, f = getSetup((3, 2), (1, 1))
+    subplotLabel(ax)
+
+    X = anndata.read_h5ad("/opt/northwest_bal/full_fitted.h5ad")
+
+    cond_fact_meta_df = condition_factors_meta(X)
+    
+    labels, plsr_results_both = plsr_acc(X, cond_fact_meta_df, n_components=1)
+    
+    overlapping_idx = cond_fact_meta_df.index.intersection(labels.index)
+    cond_fact_meta_df_filtered = cond_fact_meta_df.loc[overlapping_idx]
+    cond_fact_meta_df_filtered = cond_fact_meta_df_filtered[cond_fact_meta_df_filtered.loc[:, "patient_category"] != "COVID-19"]
+    cond_fact_meta_df_filtered = cond_fact_meta_df_filtered["patient_category"]
+
+    plot_plsr_scores_extra(plsr_results_both, cond_fact_meta_df, cond_fact_meta_df_filtered, ax[0])
+    ax[0].set(xlim=[-9.5, 9.5])
+
+    return f
+        
+def plot_plsr_scores_extra(plsr_results, cond_fact_meta_df, labels, ax):
+    """Runs PLSR and plots ROC AUC based on actual and prediction labels"""
+    type_of_data = ["nC19"]
+    score_labels = labels.loc[
+                cond_fact_meta_df.loc[:, "patient_category"] != "COVID-19"
+            ]
+    x_scores = plsr_results[1].x_scores_[:, 0]
+    df_xscores = pd.DataFrame(data=x_scores, columns=["PLSR 1"])
+    sns.swarmplot(
+        data=df_xscores,
+        x="PLSR 1",
+        ax=ax,
+        hue=score_labels.to_numpy(),
+    )
+    ax.set(xlabel="PLSR 1", ylabel="Samples", title=f"{type_of_data}-scores")
\ No newline at end of file
diff --git a/pf2/figures/figureS3.py b/pf2/figures/figureS3b.py
similarity index 100%
rename from pf2/figures/figureS3.py
rename to pf2/figures/figureS3b.py