add shape scores

molcloud/main.py

@@ -28,6 +28,8 @@ def smiles(smiles_file, output_file, width, background_color, node_size, quiet,
 @click.argument("fasta-file", type=click.Path(exists=True))
 @click.option("--output-file", default="rnacloud.png", type=click.Path(exists=False), help="Path to the output file")
 @click.option("--width", type=int, default=10, help="Width of the output image")
+@click.option("--use-shape", is_flag=True, default=False, help="Use shape information from the fasta file")
+@click.option("--cmap", type=str, default="Reds", help="Color map for the shape information")
 @click.option("--background-color", default="#f5f4e9", help="Background color of the output image")
 @click.option("--node-size", type=int, default=10, help="Size of the nodes")
 @click.option("--quiet", is_flag=True, default=False, help="Don't show progress")
@@ -44,6 +46,7 @@ def rna(fasta_file, output_file, width, background_color, node_size, quiet):
             else:
                 fasta_text += line
     plt.figure(figsize=(width, width))
-    plot_rnacloud(fasta_texts, background_color=background_color,
+    plot_rnacloud(fasta_texts, use_shape=use_shape, cmap=cmap,
+                  background_color=background_color,
                   node_size=node_size, quiet=quiet)
     plt.savefig(output_file)

molcloud/rna.py

@@ -1,4 +1,5 @@
 import forgi.graph.bulge_graph as fgb
+import matplotlib.pyplot as plt
 from .lib import *
 
 _rna_colors = {
@@ -12,16 +13,30 @@
     "hydrogen": "#999999"
 }
 
+unknown_color = "#AAAAAA"
+
+
+def scale(colors):
+    return (colors - colors.min())/(colors.max() - colors.min())
+
+def _fasta_text2graph(fasta_text, use_shape):
+    if use_shape:
+        seq, ss, nucleotide_colors = fasta_text.strip().split("\n")
+        fasta_text = "\n".join([seq, ss])
+        nucleotide_colors = scale(np.array(list(map(float, nucleotide_colors.split(",")))))
 
-def _fasta_text2graph(fasta_text):
     bg = fgb.BulgeGraph.from_fasta_text(fasta_text)[0]
     G = nx.Graph()
     residues = []
     for d in bg.defines:
         prev = None
 
         for r in bg.define_residue_num_iterator(d):
-            G.add_node(r, nucleotide=str(bg._seq)[r-1])
+            if use_shape:
+                # print(r, str(bg._seq)[r-1], nucleotide_colors[r-1])
+                G.add_node(r, nucleotide=str(bg._seq)[r-1], index = r-1, color_scale=nucleotide_colors[r-1])
+            else:
+                G.add_node(r, nucleotide=str(bg._seq)[r-1], index = r-1)
             residues += [r]
 
     # Add links along the backbone
@@ -40,11 +55,12 @@ def _fasta_text2graph(fasta_text):
     return G
 
 
-def _colors(G):
+def _colors(G, use_shape, cmap):
+    if use_shape: cmap = plt.get_cmap(cmap)
     node_colors = []
     for n, d in G.nodes(data=True):
         try:
-            c = _rna_colors[d["nucleotide"]]
+            c = cmap(d["color_scale"]) if use_shape else _rna_colors[d["nucleotide"]]
         except KeyError as e:
             c = unknown_color
         node_colors.append(c)
@@ -59,16 +75,16 @@ def _colors(G):
     return node_colors, edge_colors
 
 
-def plot_rnacloud(fasta_texts, background_color=background_color, node_size=10, quiet=False):
+def plot_rnacloud(fasta_texts, use_shape, cmap, background_color=background_color, node_size=10, quiet=False):
     G = None
 
     for fasta_text in tqdm.tqdm(fasta_texts, disable=quiet):
-        g = _fasta_text2graph(fasta_text)
+        g = _fasta_text2graph(fasta_text, use_shape)
         if g is None:
             continue
         if G is None:
             G = g
         else:
             G = nx.disjoint_union(g, G)
-    node_colors, edge_colors = _colors(G)
+    node_colors, edge_colors = _colors(G, use_shape, cmap)
     plot_graphs(G, node_colors, edge_colors, background_color, node_size)