Update factors (#44)

* Add lables for patients

* Update figure A1

---------

Co-authored-by: Andrew Ramirez <aramirez@aretha.seas.ucla.edu>

pf2/figures/commonFuncs/plotFactors.py

@@ -21,45 +21,40 @@ def plot_condition_factors(
     pd.set_option("display.max_rows", None)
     yt = pd.Series(np.unique(data.obs[cond]))
     X = np.array(data.uns["Pf2_A"])
-
-    # X = np.log10(X)
+
     XX = X
-
     X -= np.median(XX, axis=0)
     X /= np.std(XX, axis=0)
 
-    # ind = reorder_table(X)
-    # X = X[ind]
-    # yt = yt.iloc[ind]
-
-    # if cond_group_labels is not None:
-    #     cond_group_labels = cond_group_labels.iloc[ind]
-    #     ax.tick_params(axis="y", which="major", pad=20, length=0)
-    #     # extra padding to leave room for the row colors
-    #     # get list of colors for each label:
-    #     colors = sns.color_palette(
-    #         n_colors=pd.Series(cond_group_labels).nunique()
-    #     ).as_hex()
-    #     lut = {}
-    #     legend_elements = []
-    #     for index, group in enumerate(pd.Series(cond_group_labels).unique()):
-    #         lut[group] = colors[index]
-    #         legend_elements.append(Patch(color=colors[index], label=group))
-    #     row_colors = pd.Series(cond_group_labels).map(lut)
-    #     for iii, color in enumerate(row_colors):
-    #         ax.add_patch(
-    #             plt.Rectangle(
-    #                 xy=(-0.05, iii),
-    #                 width=0.05,
-    #                 height=1,
-    #                 color=color,
-    #                 lw=0,
-    #                 transform=ax.get_yaxis_transform(),
-    #                 clip_on=False,
-    #             )
-    #         )
-    #     # add a little legend
-    #     ax.legend(handles=legend_elements, bbox_to_anchor=(0.18, 1.07))
+    ind = reorder_table(X)
+    X = X[ind]
+    yt = yt.iloc[ind]
+
+    if cond_group_labels is not None:
+        cond_group_labels = cond_group_labels.iloc[ind]
+        ax.tick_params(axis="y", which="major", pad=20, length=0)
+        colors = sns.color_palette(
+            n_colors=pd.Series(cond_group_labels).nunique()
+        ).as_hex()
+        lut = {}
+        legend_elements = []
+        for index, group in enumerate(pd.Series(cond_group_labels).unique()):
+            lut[group] = colors[index]
+            legend_elements.append(Patch(color=colors[index], label=group))
+        row_colors = pd.Series(cond_group_labels).map(lut)
+        for iii, color in enumerate(row_colors):
+            ax.add_patch(
+                plt.Rectangle(
+                    xy=(-0.05, iii),
+                    width=0.05,
+                    height=1,
+                    color=color,
+                    lw=0,
+                    transform=ax.get_yaxis_transform(),
+                    clip_on=False,
+                )
+            )
+        ax.legend(handles=legend_elements, bbox_to_anchor=(0.18, 1.07))
 
     xticks = np.arange(1, X.shape[1] + 1)
     sns.heatmap(

pf2/figures/commonFuncs/plotGeneral.py

@@ -51,6 +51,19 @@ def plot_avegene_per_status(
 
     return df
 
+def bal_combine_bo_covid(df, status1: str = "binary_outcome", status2: str = "patient_category"):
+    """Combines binary outcome and covid status columns"""
+    df = df.replace({status1: {0: "L-", 
+                                1: "D-"}})
+
+    df = df.replace({status2: {"COVID-19": "C19",
+                                "Non-Pneumonia Control": "nC19", 
+                                "Other Pneumonia": "nC19",
+                                "Other Viral Pneumonia": "nC19"}})
+    df["Status"] = df[status1] + df[status2]
+
+    return df
+
 
 def rotate_xaxis(ax, rotation=90):
     """Rotates text by 90 degrees for x-axis"""

pf2/figures/figureA1.py

@@ -2,6 +2,7 @@
 along with PaCMAP labeled by cell type"""
 
 import anndata
+import pandas as pd
 from pf2.figures.common import getSetup, subplotLabel
 from pf2.tensor import correct_conditions
 from pf2.figures.commonFuncs.plotFactors import (
@@ -10,23 +11,63 @@
     plot_eigenstate_factors,
 )
 from pf2.figures.commonFuncs.plotPaCMAP import plot_labels_pacmap
+from pf2.figures.commonFuncs.plotGeneral import bal_combine_bo_covid
 from pf2.data_import import combine_cell_types, add_obs
 
 
 def makeFigure():
-    ax, f = getSetup((50, 50), (2, 3))
+    ax, f = getSetup((20, 20), (3, 3))
     subplotLabel(ax)
 
     X = anndata.read_h5ad("/opt/northwest_bal/full_fitted.h5ad")
     X.uns["Pf2_A"] = correct_conditions(X)
-    plot_condition_factors(X, ax[0], cond="sample_id")
+    add_obs(X, "patient_category")
+    add_obs(X, "binary_outcome")
+    plot_condition_factors(X, ax[0], cond="sample_id", cond_group_labels=pd.Series(label_all_samples(X)))
+    ax[0].yaxis.set_ticklabels([])
     plot_eigenstate_factors(X, ax[1])
     plot_gene_factors(X, ax[2])
-    plot_labels_pacmap(X, "cell_type", ax[3])
-    add_obs(X, "patient_category")
-    plot_labels_pacmap(X, "patient_category", ax[4])
+    ax[2].yaxis.set_ticklabels([])
+
+    df = X.obs[["patient_category", "binary_outcome"]].reset_index(drop=True)
+    df = bal_combine_bo_covid(df)
+    X.obs["Status"] = df["Status"].to_numpy()
+    plot_labels_pacmap(X, "Status", ax[3])
 
     combine_cell_types(X)
+    plot_labels_pacmap(X, "cell_type", ax[4])
     plot_labels_pacmap(X, "combined_cell_type", ax[5])
 
     return f
+
+
+def label_all_samples(X: anndata.AnnData):
+    """Label all patient samples by C19 and lived status"""
+    bo_only = ["" for x in range(len(pd.unique(X.obs["sample_id"])))]
+    pc_only = ["" for x in range(len(pd.unique(X.obs["sample_id"])))]
+    labels_samples = ["" for x in range(len(pd.unique(X.obs["sample_id"])))]
+
+    for i, sample in enumerate(pd.unique(X.obs["sample_id"])):
+        bo = pd.unique(X[X.obs.sample_id.isin([sample])].obs["binary_outcome"])
+        if bo == 0:
+            bo = "L-" 
+        else:
+            bo = "D-"
+        bo_only[i] = bo
+
+        pc = pd.unique(X[X.obs.sample_id.isin([sample])].obs["patient_category"]) 
+        if pc == "COVID-19":
+            pc = "C19"
+        else: 
+            pc = "nC19"
+        pc_only[i] = pc
+
+    for i in range(len(labels_samples)):
+        labels_samples[i] = bo_only[i]+pc_only[i]
+
+
+    return labels_samples
+
+
+
+

pf2/figures/figureA6.py

@@ -26,7 +26,7 @@ def makeFigure():
     return f
 
 
-def plot_cell_count(X: anndata.AnnData, ax: Axes, cond: str = "batch"):
+def plot_cell_count(X: anndata.AnnData, ax: Axes, cond: str = "sample_id"):
     """Plots overall cell count for Chen et al."""
     df = X.obs[[cond]].reset_index(drop=True)
     dfCond = df.groupby([cond], observed=True).size().reset_index(name="Cell Count")