add color to clustermap

KoslickiLab · Sep 27, 2023 · 2b0ed52 · 2b0ed52
1 parent 2744f8c
commit 2b0ed52
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 9 deletions.
diff --git a/fununifrac/reproducibility/_filter_files_in_dir.py b/fununifrac/reproducibility/_filter_files_in_dir.py
@@ -17,7 +17,6 @@ def parsearg():
     args = parsearg()
     all_files = os.listdir(args.source_dir)
     meta_data = pd.read_table(args.metadata_file)
-    print(all_files[:5])
     ids = meta_data['f_uid']
     selected_files = [f'sourmash_gather_out_scale1000_k_11_{i}.csv' for i in ids]
     for i, file in enumerate(all_files):

diff --git a/fununifrac/reproducibility/plot_clustermap.py b/fununifrac/reproducibility/plot_clustermap.py
@@ -20,11 +20,22 @@ def parsearg():
 
 if __name__ == "__main__":
     args = parsearg()
-    from sklearn.datasets import load_iris
-    iris = load_iris()
-    X, y = iris.data, iris.target
-    DF = pd.DataFrame(X, index=["iris_%d" % (i) for i in range(X.shape[0])], columns=iris.feature_names)
-    DF_corr = DF.T.corr()
-    DF_dism = 1 - DF_corr  # distance matrix
-    linkage = hc.linkage(sp.distance.squareform(DF_dism), method='average')
-    sns.clustermap(DF_dism, row_linkage=linkage, col_linkage=linkage)
+    pw_distance = np.load(args.pairwise_distance)
+    print(pw_distance)
+    linkage = hc.linkage(pw_distance, method='average')
+    metadata = pd.read_table(args.metadata_file)
+
+    lut = dict(zip(metadata['study_full_name'].unique(), "rbg"))
+    row_colors = metadata['study_full_name'].map(lut)
+    sns.clustermap(pw_distance, row_linkage=linkage, col_linkage=linkage)
+    plt.show()
+
+
+    # from sklearn.datasets import load_iris
+    # iris = load_iris()
+    # X, y = iris.data, iris.target
+    # DF = pd.DataFrame(X, index=["iris_%d" % (i) for i in range(X.shape[0])], columns=iris.feature_names)
+    # DF_corr = DF.T.corr()
+    # DF_dism = 1 - DF_corr  # distance matrix
+    # linkage = hc.linkage(sp.distance.squareform(DF_dism), method='average')
+    # sns.clustermap(DF_dism, row_linkage=linkage, col_linkage=linkage)