Repository: acorg/dark-matter

Issue:
Add file of all reads for a pathogen to HTML output
Julia would like a link to download all reads for a pathogen (i.e., across all samples) in the protein / pathogen HTML output.

Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1795010..4c3b542 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,8 @@
+## 3.0.77 Aug 2, 2019
+
+Added link to per-pathogen reads in protein-to-pathogens HTML output for
+Julia.
+
 ## 3.0.76 Jul 31, 2019
 
 Slightly adjust appearance of HTML links for pathogens.
diff --git a/bin/proteins-to-pathogens.py b/bin/proteins-to-pathogens.py
index d77ccf1..252d025 100755
--- a/bin/proteins-to-pathogens.py
+++ b/bin/proteins-to-pathogens.py
@@ -147,7 +147,12 @@ if __name__ == '__main__':
         '--assetDir', default='out',
         help=('The output directory where noninteractive-alignment-panel.py '
               'puts its HTML, plots and FASTA or FASTQ files, needed for '
-              'using --html'))
+              'using --html.'))
+
+    parser.add_argument(
+        '--pathogenDataDir', default='pathogen-data',
+        help=('The directory where per-pathogen information (e.g., collected '
+              'reads across all samples) should be written.'))
 
     parser.add_argument(
         '--title', default='Summary of pathogens',
@@ -202,7 +207,8 @@ if __name__ == '__main__':
                              proteinFastaFilenames=proteinFastaFilenames,
                              saveReadLengths=args.showReadLengths,
                              titleRegex=args.titleRegex,
-                             negativeTitleRegex=args.negativeTitleRegex)
+                             negativeTitleRegex=args.negativeTitleRegex,
+                             pathogenDataDir=args.pathogenDataDir)
 
     if args.filenames:
         filenames = args.filenames
diff --git a/dark/__init__.py b/dark/__init__.py
index 3f08533..b86ecac 100644
--- a/dark/__init__.py
+++ b/dark/__init__.py
@@ -8,4 +8,4 @@ if sys.version_info < (2, 7):
 #
 # Remember to update ../CHANGELOG.md describing what's new in each version.
 
-__version__ = '3.0.76'
+__version__ = '3.0.77'
diff --git a/dark/proteins.py b/dark/proteins.py
index f14719e..0ef3d1b 100644
--- a/dark/proteins.py
+++ b/dark/proteins.py
@@ -1,5 +1,6 @@
 from __future__ import division, print_function
 
+import os
 from collections import defaultdict, Counter
 import numpy as np
 from os.path import dirname, join
@@ -88,9 +89,12 @@ def getPathogenProteinCounts(filenames):
 
 class PathogenSampleFiles(object):
     """
-    Maintain a cache of pathogen/sample FASTA/FASTQ file names, creating
-    de-duplicated (by read id) FASTA/FASTQ files (from reads for all proteins
-    of a pathogen that a sample has), on demand.
+    Maintain a cache of FASTA/FASTQ file names for the samples that contain a
+    given pathogen, create de-duplicated (by read id) FASTA/FASTQ files
+    for each pathogen/sample pair, provide functions to write out index files
+    of sample and pathogen numbers (which are generated here in C{self.add}),
+    and provide a filename lookup function for pathogen/sample combinations
+    or just pathogen names by themselves.
 
     @param proteinGrouper: An instance of C{ProteinGrouper}.
     @param format_: A C{str}, either 'fasta' or 'fastq' indicating the format
@@ -186,6 +190,17 @@ class PathogenSampleFiles(object):
             sorted((index, name) for (name, index) in self._pathogens.items())
         ), file=fp)
 
+    def pathogenIndex(self, pathogenName):
+        """
+        Get the index for a pathogen.
+
+        @param pathogenName: A C{str} pathogen name.
+        @raise KeyError: If the named pathogen is unknown.
+        @return: An C{int} giving the index for the pathogen (as set in
+            self.add).
+        """
+        return self._pathogens[pathogenName]
+
 
 class ProteinGrouper(object):
     """
@@ -220,16 +235,20 @@ class ProteinGrouper(object):
         Note that this matching is done on the final part of the protein title
         in square brackets, according to the convention used by the NCBI viral
         refseq database and RVDB.
+    @param pathogenDataDir: The C{str} directory where per-pathogen information
+        (e.g., collected reads across all samples) should be written. Will be
+        created (in C{self.toHTML}) if it doesn't exist.
     @raise ValueError: If C{format_} is unknown.
     """
 
     VIRALZONE = 'https://viralzone.expasy.org/search?query='
     ICTV = 'https://talk.ictvonline.org/search-124283882/?q='
+    READCOUNT_MARKER = '*READ-COUNT*'
 
     def __init__(self, assetDir='out', sampleName=None, sampleNameRegex=None,
                  format_='fasta', proteinFastaFilenames=None,
                  saveReadLengths=False, titleRegex=None,
-                 negativeTitleRegex=None):
+                 negativeTitleRegex=None, pathogenDataDir='pathogen-data'):
         self._assetDir = assetDir
         self._sampleName = sampleName
         self._sampleNameRegex = (re.compile(sampleNameRegex) if sampleNameRegex
@@ -246,6 +265,8 @@ class ProteinGrouper(object):
         else:
             self.titleFilter = None
 
+        self._pathogenDataDir = pathogenDataDir
+
         self._pathogenProteinCount = getPathogenProteinCounts(
             proteinFastaFilenames)
 
@@ -451,6 +472,9 @@ class ProteinGrouper(object):
                 "Unrecognized pathogenType argument: %r. Value must be either "
                 "'bacterial' or 'viral'." % pathogenType)
 
+        if not os.exists(self._pathogenDataDir):
+            os.mkdir(self._pathogenDataDir)
+
         highlightSymbol = '&starf;'
         self._computeUniqueReadCounts()
 
@@ -465,22 +489,26 @@ class ProteinGrouper(object):
             with open(pathogenIndexFilename, 'w') as fp:
                 self.pathogenSampleFiles.writePathogenIndex(fp)
 
-        toDelete = defaultdict(list)
-        for pathogenName in self.pathogenNames:
-            proteinCount = self._pathogenProteinCount[pathogenName]
-            for s in self.pathogenNames[pathogenName]:
-                if proteinCount:
-                    sampleProteinFraction = (
-                        len(self.pathogenNames[pathogenName][s]['proteins']) /
-                        proteinCount)
-                else:
-                    sampleProteinFraction = 1.0
-                if sampleProteinFraction < minProteinFraction:
-                    toDelete[pathogenName].append(s)
-
-        for pathogenName in toDelete:
-            for sample in toDelete[pathogenName]:
-                del self.pathogenNames[pathogenName][sample]
+        # Figure out if we have to delete some pathogens because the
+        # fraction of their proteins that we have matches for is too low.
+        if minProteinFraction > 0.0:
+            toDelete = defaultdict(list)
+            for pathogenName in self.pathogenNames:
+                proteinCount = self._pathogenProteinCount[pathogenName]
+                for s in self.pathogenNames[pathogenName]:
+                    if proteinCount:
+                        sampleProteinFraction = (
+                            len(self.pathogenNames[
+                                pathogenName][s]['proteins']) /
+                            proteinCount)
+                    else:
+                        sampleProteinFraction = 1.0
+                    if sampleProteinFraction < minProteinFraction:
+                        toDelete[pathogenName].append(s)
+
+            for pathogenName in toDelete:
+                for sample in toDelete[pathogenName]:
+                    del self.pathogenNames[pathogenName][sample]
 
         pathogenNames = sorted(
             pathogenName for pathogenName in self.pathogenNames
@@ -667,27 +695,66 @@ class ProteinGrouper(object):
             if pathogenType == 'viral':
                 quoted = quote(pathogenName)
                 pathogenLinksHTML = (
-                    '(<a href="%s%s">ICTV</a>, <a href="%s%s">ViralZone</a>)'
+                    ' (<a href="%s%s">ICTV</a>, <a href="%s%s">ViralZone</a>)'
                 ) % (self.ICTV, quoted, self.VIRALZONE, quoted)
             else:
                 pathogenLinksHTML = ''
+
+            if pathogenProteinCount:
+                withStr = (' with %d protein%s' %
+                           (pathogenProteinCount,
+                            '' if pathogenProteinCount == 1 else 's'))
+            else:
+                withStr = ''
+
+            pathogenIndex = self.pathogenSampleFiles.pathogenIndex(
+                pathogenName)
+
+            pathogenReadsFilename = join(
+                self._pathogenDataDir,
+                'pathogen-%d.%s' % (pathogenIndex, self._format))
+
+            pathogenReadsFp = open(pathogenReadsFilename, 'w')
+            pathogenReadCount = 0
+
             append(
                 '<a id="pathogen-%s"></a>'
-                '<p class="pathogen"><span class="pathogen-name">%s</span>'
-                '%s %s, was matched by %d sample%s:</p>' %
-                (pathogenName, pathogenName, pathogenLinksHTML,
-                 ((' with %d protein%s' %
-                   (pathogenProteinCount,
-                    '' if pathogenProteinCount == 1 else 's'))
-                  if pathogenProteinCount else ''),
-                 sampleCount,
-                 '' if sampleCount == 1 else 's'))
+                '<p class="pathogen">'
+                '<span class="pathogen-name">%s</span>'
+                '%s %s, '
+                'was matched by %d sample%s '
+                '(<a href="%s">%s</a> in total):'
+                '</p>' %
+                (pathogenName,
+                 pathogenName,
+                 pathogenLinksHTML, withStr,
+                 sampleCount, '' if sampleCount == 1 else 's',
+                 pathogenReadsFilename, self.READCOUNT_MARKER))
+
+            # Remember where we are in the output result so we can fill in
+            # the total read count once we have processed all samples for
+            # this pathogen. Not nice, I know.
+            pathogenReadCountLineIndex = len(result) - 1
+
             for sampleName in sorted(samples):
                 readsFileName = self.pathogenSampleFiles.lookup(
                     pathogenName, sampleName)
+
+                # Copy the read data from the per-sample reads for this
+                # pathogen into the per-pathogen file of reads.
+                with open(readsFileName) as readsFp:
+                    while True:
+                        data = readsFp.read(4096)
+                        if data:
+                            pathogenReadsFp.write(data)
+                        else:
+                            break
+
                 proteins = samples[sampleName]['proteins']
                 proteinCount = len(proteins)
                 uniqueReadCount = samples[sampleName]['uniqueReadCount']
+                pathogenReadCount += uniqueReadCount
+
                 if pathogenProteinCount and (
                         proteinCount / pathogenProteinCount >=
                         minProteinFraction):
@@ -750,6 +817,26 @@ class ProteinGrouper(object):
 
                 append('</ul>')
 
+            pathogenReadsFp.close()
+
+            # Sanity check there's a read count marker text in our output
+            # where we expect it.
+            readCountLine = result[pathogenReadCountLineIndex]
+            if readCountLine.find(self.READCOUNT_MARKER) == -1:
+                raise ValueError(
+                    'Could not find pathogen read count marker (%s) in result '
+                    'index %d text (%s).' %
+                    (self.READCOUNT_MARKER, pathogenReadCountLineIndex,
+                     readCountLine))
+
+            # Put the read count into the pathogen summary line we wrote
+            # earlier, replacing the read count marker with the correct
+            # text.
+            result[pathogenReadCountLineIndex] = readCountLine.replace(
+                self.READCOUNT_MARKER,
+                '%d read%s' % (pathogenReadCount,
+                               '' if pathogenReadCount == 1 else 's'))
+
         # Write all samples (with pathogens (with proteins)).
         append('<hr>')
         append('<h1>Samples by pathogen</h1>')


Test patch:
diff --git a/test/test_proteins.py b/test/test_proteins.py
index f68397a..ba46f58 100644
--- a/test/test_proteins.py
+++ b/test/test_proteins.py
@@ -844,7 +844,7 @@ class TestPathogenSampleFiles(TestCase):
         being raised.
         """
         pg = ProteinGrouper()
-        error = "^format_ must be either 'fasta' or 'fastq'\\.$"
+        error = r"^format_ must be either 'fasta' or 'fastq'\.$"
         assertRaisesRegex(self, ValueError, error, PathogenSampleFiles,
                           pg, format_='unknown')
 
@@ -954,6 +954,63 @@ class TestPathogenSampleFiles(TestCase):
         # Make sure all expected filenames were seen by the mocked open.
         self.assertEqual(set(), opener.expectedFilenames)
 
+    def testPathogenIndex(self):
+        """
+        A pathogen index must be retrievable via the pathogenIndex function.
+        """
+        class Open(object):
+            def __init__(self, test, manager):
+                self.test = test
+                self.manager = manager
+                self.expectedFilenames = {'out/0.fasta', 'out/1.fasta',
+                                          'out/pathogen-0-sample-0.fasta'}
+
+            def sideEffect(self, filename, *args, **kwargs):
+                try:
+                    self.expectedFilenames.remove(filename)
+                except KeyError:
+                    self.test.fail(
+                        'Open called with unexpected filename: %r, Args: %r, '
+                        'Keyword args: %r.' % (filename, args, kwargs))
+                else:
+                    if filename == 'out/0.fasta':
+                        return StringIO('>id1\nACTG\n')
+                    elif filename == 'out/1.fasta':
+                        return StringIO('>id1\nACTG\n>id2\nCAGT\n')
+                    else:
+                        return self.manager
+
+        fp = StringIO(
+            '0.63 41.3 44.2 9 9 12 acc|GENBANK|I44.6|GENBANK|J77|VP1 '
+            '[Lausannevirus]\n'
+            '0.77 46.6 48.1 5 6 74 acc|GENBANK|I44.7|GENBANK|J78|VP2 '
+            '[Lausannevirus]\n'
+        )
+        fastaIO = StringIO()
+
+        @contextmanager
+        def manager():
+            yield fastaIO
+
+        pg = ProteinGrouper()
+        pg.addFile('filename-1', fp)
+        pathogenSampleFiles = PathogenSampleFiles(pg)
+
+        opener = Open(self, manager())
+        with patch.object(builtins, 'open') as mockMethod:
+            mockMethod.side_effect = opener.sideEffect
+            filename = pathogenSampleFiles.add('Lausannevirus', 'filename-1')
+
+        # Leaving this line in so we can explicitly see that the first
+        # pathogen has index zero.
+        self.assertEqual('out/pathogen-0-sample-0.fasta', filename)
+
+        # Here's the pathogen index we're actually interested in testing.
+        self.assertEqual(0, pathogenSampleFiles.pathogenIndex('Lausannevirus'))
+
+        # Make sure all expected filenames were seen by the mocked open.
+        self.assertEqual(set(), opener.expectedFilenames)
+
     def testProteinsSavedCorrectly(self):
         """
         Information about proteins must be saved correctly in the