Handle subgraphs and subgraphs output to fasta

Author: winni2k

.coveragerc

@@ -6,12 +6,16 @@ omit =
     cortexpy/test/*
 
 [report]
+skip_covered = True
 exclude_lines =
     no cov
     no qa
     noqa
     pragma: no cover
-    if __name__ == .__main__.:
+
+    # Don't complain if tests don't hit defensive assertion code:
+    raise AssertionError
+    raise NotImplementedError
 
 [html]
 directory = coverage_html_report

.gitignore

@@ -104,3 +104,4 @@ ENV/
 .idea/
 
 scratch/
+coverage_html_report/

Makefile

@@ -9,6 +9,7 @@ UNIT_TEST_COMMAND = $(RUN_IN_ENV) pytest \
 	--flake8 \
 	--cov=cortexpy \
 	--cov-report term-missing \
+	--cov-report html \
 	--hypothesis-profile $(HYPOTHESIS_PROFILE)
 TEST_COMMAND = BIN_DIR=$(BIN_DIR) $(UNIT_TEST_COMMAND)
 BENCHMARK_DIR := cortex_tools_benchmark

cortexpy/__init__.py

@@ -1,2 +1,2 @@
-__version__ = '0.10.0'
+__version__ = '0.11.0'
 VERSION_STRING = 'cortexpy version {}'.format(__version__)

cortexpy/command/prune.py

@@ -13,6 +13,11 @@
     graph            A cortexpy graph.  '-' redirects to or from stdout.
 """
 from cortexpy.graph import Interactor
+import logging
+
+from cortexpy.utils import get_graph_stream_iterator
+
+logger = logging.getLogger(__name__)
 
 
 def validate_prune(args):
@@ -41,10 +46,10 @@ def prune(argv):
         output = open(args['--out'], 'wb')
 
     if args['<graph>'] == '-':
-        graph = nx.read_gpickle(sys.stdin.buffer)
+        graphs = get_graph_stream_iterator(sys.stdin.buffer)
     else:
-        graph = nx.read_gpickle(args['<graph>'])
-
-    Interactor(graph, colors=None).prune_tips_less_than(args['--remove-tips'])
+        graphs = get_graph_stream_iterator(open(args['<graph>'], 'rb'))
 
-    nx.write_gpickle(graph, output)
+    for graph in graphs:
+        Interactor(graph, colors=None).prune_tips_less_than(args['--remove-tips'])
+        nx.write_gpickle(graph, output)

cortexpy/command/traverse.py

@@ -13,6 +13,7 @@
            The traverser will follow all colors specified.
     --max-nodes <n>               Maximum number of nodes to traverse [default: 1000].
     --initial-fasta               Treat <initial-contigs> as fasta
+    --subgraphs                  Emit traversal as sequence of networkx subgraphs
 
 Description:
     Traverse a cortex graph starting from each k-mer in an initial_contig and return the subgraph as
@@ -63,4 +64,9 @@ def traverse(argv):
     else:
         engine.traverse_from_each_kmer_in(args['<initial-contigs>'])
 
-    nx.write_gpickle(engine.graph, output)
+    output_graph = engine.graph
+    if args['--subgraphs'] and len(output_graph) > 0:
+        for subgraph in nx.weakly_connected_component_subgraphs(output_graph):
+            nx.write_gpickle(subgraph, output)
+    else:
+        nx.write_gpickle(output_graph, output)

cortexpy/command/view.py

@@ -18,9 +18,12 @@
                 <traversal> is '-'.
 
 """
+import networkx as nx
 from docopt import docopt
 import logging
 
+from cortexpy.utils import get_graph_stream_iterator
+
 logger = logging.getLogger('cortexpy.view')
 
 
@@ -61,24 +64,22 @@ def view_traversal(args):
 
     args = validate_view_traversal(args)
 
-    import networkx as nx
-
     if args['<traversal>'] == '-':
-        graph = nx.read_gpickle(sys.stdin.buffer)
+        graphs = get_graph_stream_iterator(sys.stdin.buffer)
     else:
-        graph = nx.read_gpickle(args['<traversal>'])
+        graphs = get_graph_stream_iterator(open(args['<traversal>'], 'rb'))
 
     if args['--to-json']:
+        graph = nx.compose_all(list(graphs))
         print(serializer.Serializer(graph).to_json())
     else:
-        kmer_serializer = serializer.Kmers(graph)
-        if args['--kmers']:
-            seq_record_generator = kmer_serializer.to_seq_records()
-        else:
-            seq_record_generator = interactor.Contigs(
-                graph, args['--color']
-            ).all_simple_paths()
-        SeqIO.write(seq_record_generator, sys.stdout, 'fasta')
+        for graph_idx, graph in enumerate(graphs):
+            if args['--kmers']:
+                seq_record_generator = serializer.Kmers(graph).to_seq_records()
+            else:
+                seq_record_generator = interactor.Contigs(graph, args['--color']).all_simple_paths()
+            seq_record_generator = annotated_seq_records(seq_record_generator, graph_idx=graph_idx)
+            SeqIO.write(seq_record_generator, sys.stdout, 'fasta')
 
 
 def view_graph(args):
@@ -120,3 +121,9 @@ def kmer_to_cortex_jdk_print_string(kmer, alt_kmer_string=None):
     to_print.append(' ' + ' '.join(map(str, kmer.coverage)))
     to_print.append(' ' + ' '.join(edge_set_strings))
     return ''.join(to_print)
+
+
+def annotated_seq_records(seq_record_generator, *, graph_idx):
+    for rec in seq_record_generator:
+        rec.id = 'g{}_p{}'.format(graph_idx, rec.id)
+        yield rec

cortexpy/graph/interactor.py

@@ -35,6 +35,7 @@ def prune_tips_less_than(self, n):
                 tip = find_tip_from(source, n=n, graph=graph, next_node_func=next_node_func)
                 if tip:
                     nodes_to_prune |= tip
+        logger.info('Pruning {} nodes'.format(len(nodes_to_prune)))
         self.graph.remove_nodes_from(nodes_to_prune)
         return self
 
@@ -68,7 +69,6 @@ def make_copy_of_color(graph, color, include_self_refs=False):
 def convert_kmer_path_to_contig(path):
     if len(path) == 0:
         return ''
-    logger.info(path)
     contig = [path[0]]
     for kmer in path[1:]:
         contig.append(kmer[-1])
@@ -95,12 +95,11 @@ class Contigs(object):
     color = attr.ib(None)
 
     def all_simple_paths(self):
-        graph = self.graph
         if self.color is not None:
             graph = make_copy_of_color(self.graph, self.color, include_self_refs=False)
+        else:
+            graph = self.graph
         idx = 0
-        logger.info(graph.nodes())
-        logger.info(graph.edges())
         for source in in_nodes_of(graph):
             for target in out_nodes_of(graph):
                 if source == target:

cortexpy/graph/serializer.py

@@ -41,7 +41,6 @@ def to_json(self):
         return json.dumps(serializable)
 
     def _collapse_graph_to_unitigs(self):
-        # self.graph = annotate_kmer_graph_edges(self.graph, self.colors)
         self._collapse_kmer_graph()
         self._make_unitig_graph_json_representable()
 

cortexpy/test/builder/graph/body.py

@@ -120,8 +120,5 @@ def edge_set_string_to_array(edge_set_string):
     return np.array(edge_set, dtype=np.uint8)
 
 
-def print_kmer(kmer):
-    print(kmer)
-    print(kmer.kmer)
-    print(kmer.coverage)
-    print(kmer.edges)
+# def print_kmer(kmer):
+#     print(kmer)

cortexpy/test/builder/mccortex.py

@@ -37,8 +37,10 @@ def build(self, tmpdir):
         output_graph = str(tmpdir.join('output.ctx'))
         mccortex_args.append(output_graph)
 
-        runner.Mccortex(self.kmer_size).run(mccortex_args)
+        ret = runner.Mccortex(self.kmer_size).run(mccortex_args)
+        logger.debug('\n' + ret.stderr.decode())
 
-        runner.Mccortex(self.kmer_size).view(output_graph)
+        ret = runner.Mccortex(self.kmer_size).view(output_graph)
+        logger.debug('\n' + ret.stdout.decode())
 
         return output_graph

cortexpy/test/driver/command.py

@@ -1,7 +1,10 @@
 from pathlib import Path
 
+import Bio
 import attr
 import os
+
+import networkx as nx
 from Bio import SeqIO
 from Bio.Seq import Seq
 from Bio.SeqRecord import SeqRecord
@@ -58,13 +61,13 @@ class Prune(object):
     tmpdir = attr.ib()
     mccortex_builder = attr.ib(attr.Factory(builder.Mccortex))
     min_tip_length = attr.ib(None)
-    last_record = attr.ib(None)
+    records = attr.ib(attr.Factory(list))
     kmer_size = attr.ib(None)
 
     def with_records(self, *records):
         for rec in records:
             self.mccortex_builder.with_dna_sequence(rec)
-            self.last_record = rec
+            self.records.append(SeqRecord(Seq(rec)))
         return self
 
     def prune_tips_less_than(self, n):
@@ -77,18 +80,110 @@ def with_kmer_size(self, size):
         return self
 
     def run(self):
-        import networkx as nx
         mccortex_graph = self.mccortex_builder.build(self.tmpdir)
 
         cortexpy_graph = self.tmpdir / 'cortexpy_graph.pickle'
-        initial_contig = self.last_record[0:self.kmer_size]
+        contig_fasta = self.tmpdir / 'initial_contigs.fa'
+        with open(str(contig_fasta), 'w') as fh:
+            SeqIO.write(self.records, fh, 'fasta')
         ctp_runner = runner.Cortexpy(SPAWN_PROCESS)
-        ctp_runner.traverse(graph=mccortex_graph, out=cortexpy_graph, contig=initial_contig)
+        ctp_runner.traverse(graph=mccortex_graph, out=cortexpy_graph, contig=contig_fasta,
+                            contig_fasta=True, subgraphs=True)
 
         pruned_graph = Path(cortexpy_graph).with_suffix('.pruned.pickle')
         completed_process = ctp_runner.prune(graph=cortexpy_graph, out=pruned_graph,
                                              remove_tips=self.min_tip_length)
 
         assert completed_process.returncode == 0, completed_process
 
-        return expectation.KmerGraphExpectation(nx.read_gpickle(str(pruned_graph)))
+        subgraphs = list(load_graph_stream(str(pruned_graph)))
+        return expectation.graph.KmerGraphsExpectation(subgraphs)
+
+
+@attr.s(slots=True)
+class Traverse(object):
+    """Runner for traverse acceptance tests"""
+    tmpdir = attr.ib()
+    mccortex_builder = attr.ib(attr.Factory(builder.Mccortex))
+    traversal_contigs = attr.ib(None)
+    added_records = attr.ib(attr.Factory(list))
+    output_subgraphs = attr.ib(False)
+    cortexpy_graph = attr.ib(init=False)
+
+    def with_records(self, *records):
+        for rec in records:
+            self.mccortex_builder.with_dna_sequence(rec)
+            self.added_records.append(rec)
+        return self
+
+    def with_initial_contigs(self, *contigs):
+        self.traversal_contigs = contigs
+
+    def with_subgraph_output(self):
+        self.output_subgraphs = True
+        return self
+
+    def with_kmer_size(self, size):
+        self.mccortex_builder.with_kmer_size(size)
+        return self
+
+    def run(self):
+        mccortex_graph = self.mccortex_builder.build(self.tmpdir)
+        contig_fasta = self.tmpdir / 'cortexpy_initial_contigs.fa'
+        if self.traversal_contigs:
+            initial_contigs = self.traversal_contigs
+        else:
+            initial_contigs = self.added_records
+        with open(str(contig_fasta), 'w') as fh:
+            Bio.SeqIO.write([SeqRecord(Seq(s)) for s in initial_contigs], fh, 'fasta')
+
+        self.cortexpy_graph = self.tmpdir / 'cortexpy_graph.pickle'
+        ctp_runner = runner.Cortexpy(SPAWN_PROCESS)
+        completed_process = ctp_runner.traverse(graph=mccortex_graph,
+                                                out=self.cortexpy_graph,
+                                                contig=contig_fasta,
+                                                contig_fasta=True,
+                                                subgraphs=self.output_subgraphs)
+
+        subgraphs = list(load_graph_stream(self.cortexpy_graph))
+        assert completed_process.returncode == 0, completed_process
+
+        return expectation.graph.KmerGraphsExpectation(subgraphs)
+
+
+@attr.s(slots=True)
+class ViewTraversal(object):
+    """Runner for view of traversal acceptance tests"""
+    tmpdir = attr.ib()
+    traverse_driver = attr.ib(init=False)
+
+    def __attrs_post_init__(self):
+        self.traverse_driver = Traverse(self.tmpdir)
+
+    def with_records(self, *records):
+        self.traverse_driver.with_records(*records)
+        return self
+
+    def with_subgraph_output(self):
+        self.traverse_driver.with_subgraph_output()
+        return self
+
+    def with_kmer_size(self, size):
+        self.traverse_driver.with_kmer_size(size)
+        return self
+
+    def run(self):
+        self.traverse_driver.run()
+        ret = runner.Cortexpy(SPAWN_PROCESS).view(
+            cortexpy_graph=self.traverse_driver.cortexpy_graph)
+        assert ret.returncode == 0, ret
+        return expectation.Fasta(ret.stdout)
+
+
+def load_graph_stream(path):
+    with open(str(path), 'rb') as fh:
+        while True:
+            try:
+                yield nx.read_gpickle(fh)
+            except EOFError:
+                break

cortexpy/test/expectation/fasta.py

@@ -1,8 +1,9 @@
 import io
-
-import sys
 from Bio import SeqIO
 import attr
+import logging
+
+logger = logging.getLogger(__name__)
 
 
 @attr.s(slots=True)
@@ -21,12 +22,16 @@ class Fasta(object):
     fasta_records = attr.ib(init=False)
 
     def __attrs_post_init__(self):
-        print(self.fasta_string, file=sys.stderr)
+        logger.info(self.fasta_string)
         self.fasta_records = [rec for rec in SeqIO.parse(io.StringIO(self.fasta_string), "fasta")]
-        print(self.fasta_records, file=sys.stderr)
+        logger.info(self.fasta_records)
 
     def has_n_records(self, n):
-        assert len(self.fasta_records) == n
+        assert n == len(self.fasta_records)
+        return self
+
+    def has_record_ids(self, *ids):
+        assert set(ids) == {rec.id for rec in self.fasta_records}
         return self
 
     def has_record(self, expected_record_seq):