From f00c15c2bda03039af76d3e9fb8fef786db7097c Mon Sep 17 00:00:00 2001
From: Nathaniel Thomas <catchnate@gmail.com>
Date: Thu, 30 May 2024 22:44:21 -0700
Subject: [PATCH] Fixed figure2, figure2 still in progress

---
 spatialpf2/figures/commonFuncs/plotFactors.py |  4 ++--
 spatialpf2/figures/figure2.py                 | 19 ++++++++++++-------
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/spatialpf2/figures/commonFuncs/plotFactors.py b/spatialpf2/figures/commonFuncs/plotFactors.py
index bf217de..97f6c9a 100644
--- a/spatialpf2/figures/commonFuncs/plotFactors.py
+++ b/spatialpf2/figures/commonFuncs/plotFactors.py
@@ -23,11 +23,11 @@ def plot_condition_factors(
     yt = pd.Series(np.unique(data.obs["Condition"]))
     X = np.array(data.uns["Pf2_A"])
 
-    X = np.log10(X)
-
     X -= np.median(X, axis=0)
     X /= np.std(X, axis=0)
 
+    X = np.log10(X)
+
     ind = reorder_table(X)
     X = X[ind]
     yt = yt.iloc[ind]
diff --git a/spatialpf2/figures/figure2.py b/spatialpf2/figures/figure2.py
index bb7c4c4..e6c07b0 100644
--- a/spatialpf2/figures/figure2.py
+++ b/spatialpf2/figures/figure2.py
@@ -6,13 +6,13 @@
 from ..imports import import_HTAN
 from ..factorization import pf2
 from .commonFuncs.plotFactors import (
-    bot_top_genes
+    plot_gene_factors_partial
 )
 import pandas as pd
 
 
 def makeFigure():
-    ax, f = getSetup((10, 12), (4, 5))
+    ax, f = getSetup((10, 12), (8, 5))
 
     rank = 20
     X = import_HTAN()
@@ -20,10 +20,15 @@ def makeFigure():
     X = pf2(X, rank, doEmbedding=False)
 
     for i in range(rank):
-        genes = bot_top_genes(X, i+1)
-        df = pd.DataFrame(
-            data=X.varm["Pf2_C"][:, i], index=X.var_names, columns=["Component"]
-        )
-        ax[i].barh(genes, df.loc[genes, "Component"])
+        plot_gene_factors_partial(i + 1, X, ax[i])
+        ax[i].set(title=f"Component {i+1}")
+        ax[i].set_xlabel("Gene")
+        ax[i].set_ylabel("Weight")
+
+    for i in range(rank, 2 * rank):
+        plot_gene_factors_partial(i - rank + 1, X, ax[i], top=False)
+        ax[i].set(title=f"Component {i - rank + 1}")
+        ax[i].set_xlabel("Gene")
+        ax[i].set_ylabel("Weight")
 
     return f