##### https://github.com/fenderglass/Ragout

Tip revision:

**40119043878d212d67385e23d0cfe1e27e99ca87**authored by**fenderglass**on**08 August 2014, 06:20:32 UTC****version 1.0** Tip revision:

**4011904** assembly_refine.py

```
#(c) 2013-2014 by Authors
#This file is a part of Ragout program.
#Released under the BSD license (see LICENSE file)
"""
This module performs refinement with the assembly (overlap) graph
"""
import networkx as nx
import re
import logging
from collections import namedtuple
try:
import Queue
except ImportError:
import queue as Queue
from ragout.shared import config
from ragout.shared.datatypes import Contig, Scaffold
logger = logging.getLogger()
def refine_scaffolds(graph_file, scaffolds, contigs_fasta):
"""
Does the job
"""
max_path_len = config.vals["overlap"]["max_path_len"]
logger.info("Refining with assembly graph")
logger.debug("Max path len = {0}".format(max_path_len))
graph = _load_dot(graph_file)
_check_overaps_number(graph, contigs_fasta)
new_scaffolds = _insert_from_graph(graph, scaffolds, max_path_len)
_reestimate_distances(graph, new_scaffolds, max_path_len, contigs_fasta)
return new_scaffolds
def _load_dot(filename):
"""
Loads dot file (ignore heavy python-graphviz)
"""
graph = nx.DiGraph()
pattern = re.compile("\"(.+)\"\s*\->\s*\"(.+)\"\s*\[.*=.*\"(.+)\".*\];")
for line in open(filename, "r").read().splitlines():
m = pattern.match(line)
if not m:
continue
v1, v2 = m.group(1), m.group(2)
assert not graph.has_edge(v1, v2)
graph.add_edge(v1, v2, label=m.group(3))
return graph
def _check_overaps_number(graph, contigs_fasta):
rate = float(len(graph.edges())) / len(contigs_fasta)
if rate < config.vals["min_overlap_rate"]:
logger.warning("Too few overlaps ({0}) between contigs were detected "
"-- refine procedure will be useless. Possible reasons:"
"\n\n1. Some contigs output by assembler are missing\n"
"2. Contigs overlap not on a constant value "
"(like k-mer for assemblers which use debruijn graph)\n"
"3. Contigs ends are trimmed/postprocessed\n"
.format(len(graph.edges())))
def _insert_from_graph(graph, scaffolds_in, max_path_len):
"""
Inserts contigs from the assembly graph into scaffolds
"""
new_scaffolds = []
ordered_contigs = set()
for scf in scaffolds_in:
ordered_contigs |= set(map(lambda s: s.name, scf.contigs))
reverse_graph = graph.reverse()
for scf in scaffolds_in:
new_scaffolds.append(Scaffold(scf.name))
for prev_cont, new_cont in zip(scf.contigs[:-1], scf.contigs[1:]):
new_scaffolds[-1].contigs.append(prev_cont)
#find contigs to insert
path_nodes = _get_cut_vertices(graph, reverse_graph, prev_cont,
new_cont, max_path_len,
ordered_contigs)
if not path_nodes:
continue
#insert contigs along the path
supp_genomes = prev_cont.link.supporting_genomes
for node in path_nodes:
sign = 1 if node[0] == "+" else -1
name = node[1:]
new_scaffolds[-1].contigs.append(Contig(name, sign))
new_scaffolds[-1].contigs[-2].link.supporting_assembly = True
new_scaffolds[-1].contigs[-1].link.supporting_assembly = True
(new_scaffolds[-1].contigs[-1].link
.supporting_genomes) = supp_genomes
new_scaffolds[-1].contigs.append(new_cont)
return new_scaffolds
def _get_cut_vertices(graph, reverse_graph, prev_cont, next_cont,
max_path_len, ordered_contigs):
"""
Finds cut vertices on a subgraph of all possible paths from one
node to another. Corresponding contigs will be inserted into scaffolds
between src and dst. This is a generalized version of what we have in paper
"""
src, dst = str(prev_cont), str(next_cont)
if not (graph.has_node(src) and graph.has_node(dst)):
logger.debug("contigs {0} / {1} are not in the graph"
.format(prev_cont, next_cont))
return None
if graph.has_edge(src, dst):
logger.debug("adjacent contigs {0} -- {1}".format(prev_cont, next_cont))
return None
restricted_nodes = set()
for contig in ordered_contigs:
restricted_nodes.add("+" + contig)
restricted_nodes.add("-" + contig)
induced_subgraph = _get_induced_subgraph(graph, reverse_graph, src, dst,
max_path_len, restricted_nodes)
if (not induced_subgraph.has_node(src) or
not induced_subgraph.has_node(dst) or
not nx.has_path(induced_subgraph, src, dst)):
return []
path = _shortest_path(induced_subgraph, src, dst, restricted_nodes)
assert path is not None
cut_vertices = set()
for node in path[1:-1]:
restricted_nodes.add(node)
if (not _test_connectivity(induced_subgraph, src, dst,
max_path_len, restricted_nodes)):
cut_vertices.add(node)
restricted_nodes.remove(node)
ordered_cut_vertices = [p for p in path if p in cut_vertices]
if len(ordered_cut_vertices):
logger.debug("found {0} cut vertixes between {1} -- {2}"
.format(len(ordered_cut_vertices), prev_cont, next_cont))
return ordered_cut_vertices
def _get_induced_subgraph(input_graph, reverse_graph, src, dst,
max_path_len, restricted_nodes):
"""
Finds subgraphs in which all possible paths between two nodes lie
"""
def dfs(graph, vertex, end_vertex, depth, visited):
visited.add(vertex)
if depth == max_path_len:
return
for _, u in graph.edges(vertex):
if u == end_vertex:
visited.add(u)
continue
if u not in visited and u not in restricted_nodes:
dfs(graph, u, end_vertex, depth + 1, visited)
visited_fwd = set()
dfs(input_graph, src, dst, 0, visited_fwd)
visited_back = set()
dfs(reverse_graph, dst, src, 0, visited_back)
result = list(visited_fwd.intersection(visited_back))
induced_digraph = nx.DiGraph()
for node in result:
for u, v in input_graph.edges(node):
if v in result:
induced_digraph.add_edge(u, v)
return induced_digraph
def _reestimate_distances(graph, scaffolds, max_path_len, contigs_fasta):
"""
Estimates distances between contigs using overlap graph
"""
restricted_nodes = set()
for scf in scaffolds:
for contig in scf.contigs:
restricted_nodes.add("+" + contig.name)
restricted_nodes.add("-" + contig.name)
for scf in scaffolds:
for prev_cont, next_cont in zip(scf.contigs[:-1], scf.contigs[1:]):
src, dst = str(prev_cont), str(next_cont)
if graph.has_edge(src, dst):
overlap = graph[src][dst]["label"]
prev_cont.link.gap = -int(overlap)
else:
path = _shortest_path(graph, src, dst, restricted_nodes)
if not path:
continue
path_len = 0
for node in path[1:-1]:
path_len += len(contigs_fasta[node[1:]])
for n1, n2 in zip(path[:-1], path[1:]):
overlap = graph[n1][n2]["label"]
path_len -= int(overlap)
prev_cont.link.gap = path_len
def _shortest_path(graph, src, dst, restricted_nodes):
"""
Finds shortest path wrt to restricted nodes
"""
queue = Queue.Queue()
queue.put(src)
visited = set([src])
parent = {src : src}
found = False
while not queue.empty():
node = queue.get()
for _, u in graph.edges(node):
if u == dst:
parent[u] = node
found = True
break
if u not in visited and u not in restricted_nodes:
visited.add(u)
queue.put(u)
parent[u] = node
if not found:
return None
path = [dst]
cur_node = dst
while parent[cur_node] != cur_node:
path.append(parent[cur_node])
cur_node = parent[cur_node]
return path[::-1]
def _test_connectivity(graph, start, end, max_path_len, restricted_nodes):
"""
Quickly tests if there is a path between two nodes
"""
class ExitSuccess(Exception):
pass
def dfs(node, depth):
visited.add(node)
if depth == max_path_len:
return
for _, u in graph.edges(node):
if u == end:
raise ExitSuccess
if u not in visited and u not in restricted_nodes:
dfs(u, depth + 1)
visited = set()
try:
dfs(start, 0)
except ExitSuccess:
return True
return False
```