From c63e95aec136b5ca7886259c50ffcfde08c864ea Mon Sep 17 00:00:00 2001
From: Armaan Abraham <armaanabraham@g.ucla.edu>
Date: Thu, 1 Feb 2024 20:04:16 -0800
Subject: [PATCH] figure S2

---
 ddmc/figures/common.py    |  6 ++++--
 ddmc/figures/figureMS2.py | 36 ++++++++++++++++++------------------
 2 files changed, 22 insertions(+), 20 deletions(-)

diff --git a/ddmc/figures/common.py b/ddmc/figures/common.py
index 7a2de1f3..d2fed372 100644
--- a/ddmc/figures/common.py
+++ b/ddmc/figures/common.py
@@ -131,10 +131,12 @@ def genFigure():
 def plot_motifs(pssm, ax: axes.Axes, titles=False, yaxis=False):
     """Generate logo plots of a list of PSSMs"""
     pssm = pssm.T
+    pssm = pd.DataFrame(pssm)
     if pssm.shape[0] == 11:
         pssm.index = [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5]
     elif pssm.shape[0] == 9:
         pssm.index = [-5, -4, -3, -2, -1, 1, 2, 3, 4]
+    pssm.columns = AAlist
     logo = lm.Logo(
         pssm,
         font_name="Arial",
@@ -202,7 +204,7 @@ def plot_cluster_kinase_distances(
 
 
 def get_pvals_across_clusters(
-    label: pd.Series[bool] | np.ndarray[bool], centers: pd.DataFrame | np.ndarray
+    label: pd.Series | np.ndarray[bool], centers: pd.DataFrame | np.ndarray
 ) -> np.ndarray[float]:
     pvals = []
     centers_pos = centers[label]
@@ -213,7 +215,7 @@ def get_pvals_across_clusters(
 
 
 def plot_p_signal_across_clusters_and_binary_feature(
-    label: pd.Series[bool] | np.ndarray[bool],
+    label: pd.Series | np.ndarray[bool],
     centers: pd.DataFrame | np.ndarray,
     label_name: str,
     ax,
diff --git a/ddmc/figures/figureMS2.py b/ddmc/figures/figureMS2.py
index a84b50c8..6ed9c950 100644
--- a/ddmc/figures/figureMS2.py
+++ b/ddmc/figures/figureMS2.py
@@ -3,29 +3,29 @@
 """
 
 import numpy as np
-from .common import getSetup, getDDMC_CPTAC
-from .common import plot_motifs
 
+from ddmc.clustering import DDMC
+from ddmc.datasets import CPTAC, select_peptide_subset
+from ddmc.figures.common import getSetup, plot_motifs
 
-def makeFigure():
-    """Get a list of the axis objects and create a figure"""
-    # Get list of axis objects
-    ax, f = getSetup((9, 9), (5, 5))
 
-    # Fit DDMC
-    model, _ = getDDMC_CPTAC(n_components=30, SeqWeight=100.0)
+def makeFigure():
+    # Increase number of peptides and components for actual figure
+    p_signal = select_peptide_subset(CPTAC().get_p_signal(), keep_num=500)
+    model = DDMC(n_components=5, seq_weight=100).fit(p_signal)
 
-    pssms, cl_num = model.get_pssms(PsP_background=False)
+    ax, f = getSetup((9, 9), (5, 5))
+    clusters, pssms = model.get_pssms(PsP_background=False)
     ylabels = np.arange(0, 21, 5)
     xlabels = [20, 21, 22, 23, 24, 25]
-    for ii, cc in enumerate(cl_num):
-        cluster = "Cluster " + str(cc)
-        plot_motifs(pssms[ii], ax=ax[ii], titles=cluster, yaxis=[0, 10])
-        if ii not in ylabels:
-            ax[ii].set_ylabel("")
-            ax[ii].get_yaxis().set_visible(False)
-        if ii not in xlabels:
-            ax[ii].set_xlabel("")
-            ax[ii].get_xaxis().set_visible(False)
+    for cluster in clusters:
+        cluster_label = "Cluster " + str(cluster)
+        plot_motifs(pssms[cluster], ax=ax[cluster], titles=cluster_label, yaxis=[0, 10])
+        if cluster not in ylabels:
+            ax[cluster].set_ylabel("")
+            ax[cluster].get_yaxis().set_visible(False)
+        if cluster not in xlabels:
+            ax[cluster].set_xlabel("")
+            ax[cluster].get_xaxis().set_visible(False)
 
     return f