Merge pull request #1059 from broadinstitute/dp/rasusa-coverage-normalization

Add rasusa-based coverage normalization for assembly polishing

Author: dpark01

src/viral_ngs/assemble/__init__.py

@@ -13,5 +13,6 @@
 from . import skani
 from . import spades
 from . import freebayes
+from . import rasusa
 from . import vcf
 from . import wgsim

src/viral_ngs/assemble/rasusa.py

@@ -0,0 +1,48 @@
+'''
+    Rasusa — randomly subsample sequencing reads or alignments to a
+    specified coverage depth.
+
+    https://github.com/mbhall88/rasusa
+'''
+
+import logging
+import shutil
+import subprocess
+
+from ..core import Tool, PrexistingUnixCommand
+
+_log = logging.getLogger(__name__)
+
+TOOL_NAME = 'rasusa'
+
+
+class RasusaTool(Tool):
+
+    def __init__(self):
+        install_methods = [PrexistingUnixCommand(shutil.which(TOOL_NAME), require_executability=True)]
+        Tool.__init__(self, install_methods=install_methods)
+
+    def downsample_bam(self, inBam, outBam, coverage, seed=None):
+        """Downsample aligned reads in a BAM to a target coverage depth
+        using ``rasusa aln``.
+
+        The input BAM must be coordinate-sorted.  The output BAM is
+        **not** guaranteed to be sorted — the caller is responsible for
+        sorting and indexing if downstream tools require it.
+
+        Args:
+            inBam:    Input BAM file (coordinate-sorted).
+            outBam:   Output BAM file path.
+            coverage: Target coverage depth (integer).
+            seed:     Random seed for reproducibility (optional).
+        """
+        tool_cmd = [
+            self.install_and_get_path(), 'aln',
+            '--coverage', str(coverage),
+            '-o', outBam,
+            inBam,
+        ]
+        if seed is not None:
+            tool_cmd.extend(['--seed', str(seed)])
+        _log.debug(' '.join(tool_cmd))
+        subprocess.check_call(tool_cmd)

src/viral_ngs/assembly.py

@@ -31,6 +31,7 @@
 import viral_ngs.core.samtools
 import viral_ngs.core.novoalign
 import viral_ngs.assemble.freebayes
+import viral_ngs.assemble.rasusa
 
 # intra-module (assembly tool wrappers)
 import viral_ngs.assemble.vcf
@@ -670,6 +671,8 @@ def refine_assembly(
     major_cutoff=0.5,
     chr_names=None,
     keep_all_reads=False,
+    max_coverage=None,
+    rasusa_seed=None,
     already_realigned_bam=None,
     JVMmemory=None,
     threads=None,
@@ -737,36 +740,50 @@ def refine_assembly(
             shutil.copyfile(realignBam, outBam)
 
     # Modify original assembly with VCF calls from FreeBayes
-    tmpVcf = viral_ngs.core.file.mkstempfname('.vcf.gz')
-    tmpFasta = viral_ngs.core.file.mkstempfname('.fasta')
-    fb.call(realignBam, deambigFasta, tmpVcf)
-    if already_realigned_bam is None:
-        os.unlink(realignBam)
-    os.unlink(deambigFasta)
-    name_opts = []
-    if chr_names:
-        name_opts = ['--name'] + chr_names
-    main_vcf_to_fasta(
-        parser_vcf_to_fasta(argparse.ArgumentParser(
-        )).parse_args([
-            tmpVcf,
-            tmpFasta,
-            '--trim_ends',
-            '--min_coverage',
-            str(min_coverage),
-            '--major_cutoff',
-            str(major_cutoff)
-        ] + name_opts)
-    )
-    if outVcf:
-        shutil.copyfile(tmpVcf, outVcf)
-        if outVcf.endswith('.gz'):
-            shutil.copyfile(tmpVcf + '.tbi', outVcf + '.tbi')
-    os.unlink(tmpVcf)
-    if os.path.exists(tmpVcf + '.tbi'):
-        os.unlink(tmpVcf + '.tbi')
-    shutil.copyfile(tmpFasta, outFasta)
-    os.unlink(tmpFasta)
+    with viral_ngs.core.file.tempfname('.vcf.gz') as tmpVcf, \
+         viral_ngs.core.file.tempfname('.fasta') as tmpFasta:
+
+        # Optionally downsample coverage for variant calling only
+        if max_coverage:
+            rasusa_tool = viral_ngs.assemble.rasusa.RasusaTool()
+            with viral_ngs.core.file.tempfname('.downsampled.unsorted.bam') as downsampledUnsortedBam, \
+                 viral_ngs.core.file.tempfname('.downsampled.sorted.bam') as downsampledBam:
+                rasusa_tool.downsample_bam(realignBam, downsampledUnsortedBam, max_coverage, seed=rasusa_seed)
+                samtools.sort(downsampledUnsortedBam, downsampledBam, threads=threads)
+                samtools.index(downsampledBam, threads=threads)
+                fb.call(downsampledBam, deambigFasta, tmpVcf)
+                # clean up samtools index sidecar
+                if os.path.isfile(downsampledBam + '.bai'):
+                    os.unlink(downsampledBam + '.bai')
+        else:
+            fb.call(realignBam, deambigFasta, tmpVcf)
+
+        if already_realigned_bam is None:
+            os.unlink(realignBam)
+        os.unlink(deambigFasta)
+        name_opts = []
+        if chr_names:
+            name_opts = ['--name'] + chr_names
+        main_vcf_to_fasta(
+            parser_vcf_to_fasta(argparse.ArgumentParser(
+            )).parse_args([
+                tmpVcf,
+                tmpFasta,
+                '--trim_ends',
+                '--min_coverage',
+                str(min_coverage),
+                '--major_cutoff',
+                str(major_cutoff)
+            ] + name_opts)
+        )
+        if outVcf:
+            shutil.copyfile(tmpVcf, outVcf)
+            if outVcf.endswith('.gz'):
+                shutil.copyfile(tmpVcf + '.tbi', outVcf + '.tbi')
+        shutil.copyfile(tmpFasta, outFasta)
+        # clean up tabix index sidecar created by FreeBayes/tabix
+        if os.path.isfile(tmpVcf + '.tbi'):
+            os.unlink(tmpVcf + '.tbi')
 
     # Index final output FASTA for Picard, Samtools, and Novoalign
     picard_index.execute(outFasta, overwrite=True)
@@ -837,6 +854,22 @@ def parser_refine_assembly(parser=argparse.ArgumentParser()):
         action="store_true",
         dest="keep_all_reads"
     )
+    parser.add_argument(
+        '--max_coverage',
+        default=None,
+        type=int,
+        help="""Maximum read coverage depth for variant calling. If specified,
+            reads are downsampled to this coverage using rasusa before
+            FreeBayes variant calling. The downsampled BAM is internal only
+            and not included in any output. [default: %(default)s]"""
+    )
+    parser.add_argument(
+        '--rasusa_seed',
+        default=None,
+        type=int,
+        help="""Random seed for rasusa downsampling reproducibility.
+            Only used when --max_coverage is set. [default: %(default)s]"""
+    )
     parser.add_argument(
         '--JVMmemory',
         default='2g',
@@ -856,6 +889,39 @@ def parser_refine_assembly(parser=argparse.ArgumentParser()):
 __commands__.append(('refine_assembly', parser_refine_assembly))
 
 
+def normalize_coverage(inBam, outBam, max_coverage, seed=None, threads=None):
+    '''Downsample an aligned BAM to a maximum coverage depth using rasusa.
+       The output BAM is coordinate-sorted and indexed.
+    '''
+    samtools = viral_ngs.core.samtools.SamtoolsTool()
+    rasusa_tool = viral_ngs.assemble.rasusa.RasusaTool()
+
+    with viral_ngs.core.file.tempfname('.rasusa.unsorted.bam') as tmpBam:
+        rasusa_tool.downsample_bam(inBam, tmpBam, max_coverage, seed=seed)
+        samtools.sort(tmpBam, outBam, threads=threads)
+
+    samtools.index(outBam, threads=threads)
+    return outBam
+
+
+def parser_normalize_coverage(parser=argparse.ArgumentParser()):
+    parser.add_argument('inBam', help='Input aligned BAM file, coordinate-sorted.')
+    parser.add_argument('outBam', help='Output downsampled BAM file, coordinate-sorted and indexed.')
+    parser.add_argument('max_coverage', type=int, help='Maximum coverage depth to downsample to.')
+    parser.add_argument(
+        '--seed',
+        default=None,
+        type=int,
+        help='Random seed for rasusa reproducibility. [default: %(default)s]'
+    )
+    viral_ngs.core.cmd.common_args(parser, (('threads', None), ('loglevel', None), ('version', None), ('tmp_dir', None)))
+    viral_ngs.core.cmd.attach_main(parser, normalize_coverage, split_args=True)
+    return parser
+
+
+__commands__.append(('normalize_coverage', parser_normalize_coverage))
+
+
 def parser_filter_short_seqs(parser=argparse.ArgumentParser()):
     parser.add_argument("inFile", help="input sequence file")
     parser.add_argument("minLength", help="minimum length for contig", type=int)

tests/unit/assemble/test_assembly.py

@@ -123,6 +123,39 @@ def test_empty_input_succeed(self):
         self.assertTrue(os.path.getsize(outFasta) == 0)
 
 
+class TestNormalizeCoverage(TestCaseWithTmp):
+    '''Test the normalize_coverage command'''
+
+    def test_help_parser(self):
+        parser = viral_ngs.assembly.parser_normalize_coverage(argparse.ArgumentParser())
+        helpstring = parser.format_help()
+        self.assertIn('max_coverage', helpstring)
+        self.assertIn('seed', helpstring)
+
+    def test_normalize_coverage_on_aligned_bam(self):
+        '''normalize_coverage should downsample and produce a sorted, indexed BAM'''
+        mm2 = viral_ngs.core.minimap2.Minimap2()
+        samtools = viral_ngs.core.samtools.SamtoolsTool()
+
+        inDir = viral_ngs.core.file.get_test_input_path()
+        inBam = os.path.join(inDir, 'G5012.3.testreads.bam')
+        refFasta = os.path.join(inDir, 'ebov-makona.fasta')
+
+        with viral_ngs.core.file.tempfname('.aligned.bam') as alignedBam, \
+             viral_ngs.core.file.tempfname('.normalized.bam') as outBam:
+            mm2.align_bam(inBam, refFasta, alignedBam, options=['-x', 'sr'])
+            samtools.index(alignedBam)
+
+            viral_ngs.assembly.normalize_coverage(alignedBam, outBam, max_coverage=10, seed=42)
+
+            self.assertTrue(os.path.isfile(outBam))
+            self.assertGreater(os.path.getsize(outBam), 0)
+            # index file should also exist
+            self.assertTrue(os.path.isfile(outBam + '.bai'))
+            # output should have fewer or equal reads compared to input
+            self.assertLessEqual(samtools.count(outBam), samtools.count(alignedBam))
+
+
 class TestAssembleSpades(TestCaseWithTmp):
     ''' Test the assemble_spades command (no validation of output) '''