added functions get_starting_nodes, get_sink_nodes, get_contigs, save_contigs

debruijn/debruijn.py

@@ -2,6 +2,7 @@
 Small assembly module based on de bruijn graphs
 """
 import networkx as nx
+from networkx import algorithms
 
 def read_fastq(fichier):
     """
@@ -58,10 +59,107 @@ def build_graph(hash_table):
     Arguments:
         hash_table, dict: dictionnary obtained with build_kmer_dict() function
     Return:
-        graph, nx.Graph: the de Bruijn tree corresponding to hash_table
+        graph, nx.DiGraph: the de Bruijn tree corresponding to hash_table
     """
     graph = nx.DiGraph()
     for key in hash_table:
         graph.add_edge(key[:-1], key[1:], weight=hash_table[key])
 
     return graph
+
+def get_starting_nodes(graph):
+    """
+    Arguments:
+        graph, nx.DiGraph: de Bruijn tree
+
+    Return:
+        starting_nodes, list of strings: list of starting nodes 
+    """
+    starting_nodes = []
+    for node in graph:
+        if graph.in_degree(node) == 0:
+            starting_nodes.append(node)
+
+    return starting_nodes
+
+def std():
+    pass
+
+
+def get_sink_nodes(graph):
+    """
+    Arguments:
+        graph, nx.DiGraph: de Bruijn tree
+
+    Return:
+        sink_nodes, list of strings: list of terminal nodes 
+    """
+    sink_nodes = []
+    for node in graph:
+        if graph.out_degree(node) == 0:
+            sink_nodes.append(node)
+
+    return sink_nodes
+
+
+def path_average_weight():
+    pass
+
+
+def remove_paths():
+    pass
+
+
+def select_best_path():
+    pass
+
+
+def save_contigs(tuples, outname):
+    """
+    Arguments:
+        tuples, tuple: Obtained from get_contigs()
+        outname, str: name of the file to be written
+    """
+    with open(outname, "w") as outfile:
+        for duo in tuples:
+            outfile.write("{} {}".format(duo[0], duo[1]))
+
+    return
+
+
+def get_contigs(graph, starting_nodes, sink_nodes):
+    """
+    Arguments:
+        graph, nx.DiGraph: de Bruijn tree
+        starting_nodes, list of strings: list of starting nodes 
+        sink_nodes, list of strings: list of terminal nodes 
+
+    Return:
+        contigs, list of tupple: list of tupple (contigs, len(contigs))
+    """
+    contigs = []
+    for starting_node in starting_nodes:
+        for sink_node in sink_nodes:
+            if algorithms.has_path(graph, starting_node, sink_node) == True:
+                path = algorithms.shortest_path(graph, starting_node, sink_node)
+                contig = path[0]
+                for i in range(len(path)-1):
+                    contig += path[i+1][-1]
+                contigs.append((contig, len(contig)))
+
+    return contigs
+
+def solve_bubble():
+    pass
+
+
+def simplify_bubbles():
+    pass
+
+
+def solve_entry_tips():
+    pass
+
+
+def solve_out_tips():
+    pass