test_pick_rep_set.py

#!/usr/bin/env python

"""Tests of code for representative set picking"""

__author__ = "Rob Knight"
__copyright__ = "Copyright 2011, The QIIME Project"
# remember to add yourself if you make changes
__credits__ = ["Rob Knight", "Kyle Bittinger", "Greg Caporaso"]
__license__ = "GPL"
__version__ = "1.9.1"
__maintainer__ = "Daniel McDonald"
__email__ = "wasade@gmail.com"

from os import remove, close
from tempfile import mkstemp
from unittest import TestCase, main

from skbio.util import remove_files
from skbio.parse.sequences import parse_fasta
from skbio.alignment import SequenceCollection
from skbio.sequence import DNA

from qiime.pick_rep_set import (RepSetPicker, GenericRepSetPicker, first_id,
                                first, random_id, longest_id, unique_id_map, label_to_name,
                                make_most_abundant, parse_fasta, ReferenceRepSetPicker)


class RepSetPickerTests(TestCase):

    """Tests of the abstract RepSetPicker class"""

    def test_init(self):
        """Abstract RepSetPicker __init__ should store name, params"""
        p = RepSetPicker({})
        self.assertEqual(p.Name, 'RepSetPicker')
        self.assertEqual(p.Params, {})

    def test_call(self):
        """Abstract RepSetPicker __call__ should raise NotImplementedError"""
        p = RepSetPicker({})
        self.assertRaises(NotImplementedError, p, '/path/to/seqs',
                          '/path/to/otus')


class SharedSetupTestCase(TestCase):

    """Wrapper for shared setup stuff"""

    def setUp(self):
        # create the temporary input files
        fd, self.tmp_seq_filepath = mkstemp(prefix='GenericRepSetPickerTest_',
                                           suffix='.fasta')
        close(fd)
        seq_file = open(self.tmp_seq_filepath, 'w')
        seq_file.write(dna_seqs)
        seq_file.close()

        fd, self.tmp_otu_filepath = mkstemp(prefix='GenericRepSetPickerTest_',
                                           suffix='.otu')
        close(fd)
        otu_file = open(self.tmp_otu_filepath, 'w')
        otu_file.write(otus)
        otu_file.close()

        self.files_to_remove = [self.tmp_seq_filepath, self.tmp_otu_filepath]

        self.params = {'Algorithm': 'first', 'ChoiceF': first_id}

    def tearDown(self):
        remove_files(self.files_to_remove)


class GenericRepSetPickerTests(SharedSetupTestCase):

    """ Tests of the generic RepSet picker """

    def test_call_default_params(self):
        """GenericRepSetPicker.__call__ returns expected clusters default params"""

        # adapted from test_app.test_cd_hit.test_cdhit_clusters_from_seqs

        exp = {'0': 'R27DLI_4812',
               '1': 'U1PLI_7889',
               '2': 'W3Cecum_4858',
               '3': 'R27DLI_3243',
               }
        app = GenericRepSetPicker(params={'Algorithm': 'first',
                                          'ChoiceF': first_id})
        obs = app(self.tmp_seq_filepath, self.tmp_otu_filepath)
        self.assertEqual(obs, exp)

    def test_call_wrapped_function(self):
        """GenericRepSetPicker.__call__ returns expected clusters default params"""

        # adapted from test_app.test_cd_hit.test_cdhit_clusters_from_seqs

        exp = {'0': 'R27DLI_4812',
               '1': 'U1PLI_7889',
               '2': 'W3Cecum_4858',
               '3': 'R27DLI_3243',
               }
        app = GenericRepSetPicker(params={'Algorithm': 'most_abundant',
                                          'ChoiceF': make_most_abundant, 'ChoiceFRequiresSeqs': True})
        obs = app(self.tmp_seq_filepath, self.tmp_otu_filepath)
        self.assertEqual(obs, exp)

    def test_call_output_to_file(self):
        """GenericRepSetPicker.__call__ output to file functions as expected
        """

        fd, tmp_result_filepath = mkstemp(
            prefix='GenericRepSetPickerTest.test_call_output_to_file_',
            suffix='.txt')
        close(fd)

        app = GenericRepSetPicker(params=self.params)
        obs = app(self.tmp_seq_filepath, self.tmp_otu_filepath,
                  result_path=tmp_result_filepath)

        result_file = open(tmp_result_filepath)
        result_file_str = result_file.read()
        result_file.close()
        # remove the result file before running the test, so in
        # case it fails the temp file is still cleaned up
        remove(tmp_result_filepath)

        # compare data in result file to fake expected file
        self.assertEqual(result_file_str, rep_seqs_result_file_exp)
        # confirm that nothing is returned when result_path is specified
        self.assertEqual(obs, None)

    def test_call_output_to_file_sorted(self):
        """GenericRepSetPicker.__call__ output to file sorts when requested
        """

        fd, tmp_result_filepath = mkstemp(
            prefix='GenericRepSetPickerTest.test_call_output_to_file_',
            suffix='.txt')
        close(fd)

        app = GenericRepSetPicker(params=self.params)
        obs = app(self.tmp_seq_filepath, self.tmp_otu_filepath,
                  result_path=tmp_result_filepath, sort_by='seq_id')

        result_file = open(tmp_result_filepath)
        result_file_str = result_file.read()
        result_file.close()
        # remove the result file before running the test, so in
        # case it fails the temp file is still cleaned up
        remove(tmp_result_filepath)

        # compare data in result file to fake expected file
        self.assertEqual(result_file_str, rep_seqs_result_file_sorted_exp)
        # confirm that nothing is returned when result_path is specified
        self.assertEqual(obs, None)

    def test_call_log_file(self):
        """GenericRepSetPicker.__call__ writes log when expected
        """

        fd, tmp_log_filepath = mkstemp(
            prefix='GenericRepSetPickerTest.test_call_output_to_file_l_',
            suffix='.txt')
        close(fd)
        fd, tmp_result_filepath = mkstemp(
            prefix='GenericRepSetPickerTest.test_call_output_to_file_r_',
            suffix='.txt')
        close(fd)

        app = GenericRepSetPicker(params=self.params)
        obs = app(self.tmp_seq_filepath, self.tmp_otu_filepath,
                  result_path=tmp_result_filepath, log_path=tmp_log_filepath)

        log_file = open(tmp_log_filepath)
        log_file_str = log_file.read()
        log_file.close()
        # remove the temp files before running the test, so in
        # case it fails the temp file is still cleaned up
        remove(tmp_log_filepath)
        remove(tmp_result_filepath)

        log_file_exp = ["GenericRepSetPicker parameters:",
                        'Algorithm:first',
                        "Application:None",
                        'ChoiceF:first',
                        'ChoiceFRequiresSeqs:False',
                        "Result path: %s" % tmp_result_filepath, ]
        # compare data in log file to fake expected log file
        for i, j in zip(log_file_str.splitlines(), log_file_exp):
            if not i.startswith('ChoiceF:'):  # can't test, different each time
                self.assertEqual(i, j)


class ReferenceRepSetPickerTests(SharedSetupTestCase):

    """Tests of the ReferenceRepSetPickerclass """

    def setUp(self):
        # create the temporary input files
        fd, self.tmp_seq_filepath = mkstemp(
            prefix='ReferenceRepSetPickerTest_',
            suffix='.fasta')
        close(fd)
        seq_file = open(self.tmp_seq_filepath, 'w')
        seq_file.write(dna_seqs)
        seq_file.close()

        fd, self.ref_seq_filepath = mkstemp(
            prefix='ReferenceRepSetPickerTest_',
            suffix='.fasta')
        close(fd)
        seq_file = open(self.ref_seq_filepath, 'w')
        seq_file.write(reference_seqs)
        seq_file.close()

        fd, self.tmp_otu_filepath = mkstemp(
            prefix='ReferenceRepSetPickerTest_',
            suffix='.otu')
        close(fd)
        otu_file = open(self.tmp_otu_filepath, 'w')
        otu_file.write(otus_w_ref)
        otu_file.close()

        fd, self.result_filepath = mkstemp(
            prefix='ReferenceRepSetPickerTest_',
            suffix='.fasta')
        close(fd)
        otu_file = open(self.result_filepath, 'w')
        otu_file.write(otus_w_ref)
        otu_file.close()

        self.files_to_remove = [self.tmp_seq_filepath,
                                self.tmp_otu_filepath,
                                self.ref_seq_filepath,
                                self.result_filepath]

        self.params = {'Algorithm': 'first', 'ChoiceF': first_id}

    def test_call_default_params(self):
        """ReferenceRepSetPicker.__call__ expected clusters default params"""

        exp = {'0': ('R27DLI_4812', 'CTGGGCCGTATCTC'),
               'ref1': ('ref1', 'GGGGGGGAAAAAAAAAAAAA'),
               '2': ('W3Cecum_4858', 'TTGGGCCGTGTCTCAGT'),
               'ref0': ('ref0', 'CCCAAAAAAATTTTTT'),
               }
        app = ReferenceRepSetPicker(params={'Algorithm': 'first',
                                            'ChoiceF': first_id})
        obs = app(self.tmp_seq_filepath,
                  self.tmp_otu_filepath,
                  self.ref_seq_filepath)
        self.assertEqual(obs, exp)

    def test_call_write_to_file(self):
        """ReferenceRepSetPicker.__call__ otu map correctly written to file"""
        app = ReferenceRepSetPicker(params={'Algorithm': 'first',
                                            'ChoiceF': first_id})
        app(self.tmp_seq_filepath,
            self.tmp_otu_filepath,
            self.ref_seq_filepath,
            result_path=self.result_filepath)
        with open(self.result_filepath) as f:
            actual = SequenceCollection.from_fasta_records(parse_fasta(f), DNA)
        expected = SequenceCollection.from_fasta_records(
            parse_fasta(rep_seqs_reference_result_file_exp.split('\n')), DNA)
        # we don't care about order in the results
        self.assertEqual(set(actual), set(expected))

    def test_non_ref_otus(self):
        """ReferenceRepSetPicker.__call__ same result as Generic when no ref otus
        """
        exp = {'0': ('R27DLI_4812', 'CTGGGCCGTATCTC'),
               '1': ('U1PLI_7889', 'TTGGACCGTG'),
               '2': ('W3Cecum_4858', 'TTGGGCCGTGTCTCAGT'),
               '3': ('R27DLI_3243', 'CTGGACCGTGTCT')}
        fd, tmp_otu_filepath = mkstemp(
            prefix='ReferenceRepSetPickerTest_',
            suffix='.otu')
        close(fd)
        otu_file = open(tmp_otu_filepath, 'w')
        otu_file.write(otus)
        otu_file.close()

        self.files_to_remove.append(tmp_otu_filepath)

        app = ReferenceRepSetPicker(params={'Algorithm': 'first',
                                            'ChoiceF': first_id})
        obs = app(self.tmp_seq_filepath,
                  tmp_otu_filepath,
                  self.ref_seq_filepath)
        self.assertEqual(obs, exp)

    def test_call_invalid_id(self):
        """ReferenceRepSetPicker.__call__ expected clusters default params"""
        app = ReferenceRepSetPicker(params={'Algorithm': 'first',
                                            'ChoiceF': first_id})

        fd, tmp_otu_filepath = mkstemp(
            prefix='ReferenceRepSetPickerTest_',
            suffix='.otu')
        close(fd)
        otu_file = open(tmp_otu_filepath, 'w')
        # replace a valid sequence identifier with an invalid
        # sequence identifier (i.e., one that we don't have a sequence for)
        otu_file.write(otus_w_ref.replace('R27DLI_4812', 'bad_seq_identifier'))
        otu_file.close()
        self.files_to_remove.append(tmp_otu_filepath)

        # returning in dict
        self.assertRaises(KeyError,
                          app,
                          self.tmp_seq_filepath,
                          tmp_otu_filepath,
                          self.ref_seq_filepath)
        # writing to file
        self.assertRaises(KeyError,
                          app,
                          self.tmp_seq_filepath,
                          tmp_otu_filepath,
                          self.ref_seq_filepath,
                          result_path=self.result_filepath)

    def test_call_ref_only(self):
        """ReferenceRepSetPicker.__call__ functions with no non-refseqs"""

        fd, tmp_otu_filepath = mkstemp(
            prefix='ReferenceRepSetPickerTest_',
            suffix='.otu')
        close(fd)
        otu_file = open(tmp_otu_filepath, 'w')
        otu_file.write(otus_all_ref)
        otu_file.close()
        self.files_to_remove.append(tmp_otu_filepath)

        exp = {'ref1': ('ref1', 'GGGGGGGAAAAAAAAAAAAA'),
               'ref0': ('ref0', 'CCCAAAAAAATTTTTT')}

        # passing only reference (not input seqs)
        app = ReferenceRepSetPicker(params={'Algorithm': 'first',
                                            'ChoiceF': first_id})
        obs = app(None,
                  tmp_otu_filepath,
                  self.ref_seq_filepath)
        self.assertEqual(obs, exp)

        # passing reference and input seqs
        app = ReferenceRepSetPicker(params={'Algorithm': 'first',
                                            'ChoiceF': first_id})
        obs = app(self.tmp_seq_filepath,
                  tmp_otu_filepath,
                  self.ref_seq_filepath)
        self.assertEqual(obs, exp)

    def test_call_alt_non_ref_picker(self):
        """ReferenceRepSetPicker.__call__ handles alt non-ref picking method"""

        exp = {'0': ('U1PLI_9526', 'CTGGGCCGTATCTCAGTCCCAATGTGGCCGGTCG'
                     'GTCTCTCAACCCGGCTACCCATCGCGGGCTAGGTGGGCCGTT'
                     'ACCCCGCCTACTACCTAATGGGCCGCGACCCCATCCCTTGCCGTCTGGGC'
                     'TTTCCCGGGCCCCCCAGGAGGGGGGCGAGGAGTATCCGGTATTAGCCTCGGTT'
                     'TCCCAAGGTTGTCCCGGAGCAAGGGGCAGGTTGGTCACGTGTTACTCACCCGT'
                     'TCGCCACTTCATGTCCGCCCGAGGGCGGTTTCATCG'),
               'ref1': ('ref1', 'GGGGGGGAAAAAAAAAAAAA'),
               '2': ('W3Cecum_4858', 'TTGGGCCGTGTCTCAGT'),
               'ref0': ('ref0', 'CCCAAAAAAATTTTTT'),
               }
        app = ReferenceRepSetPicker(params={'Algorithm': 'longest',
                                            'ChoiceF': longest_id})
        obs = app(self.tmp_seq_filepath,
                  self.tmp_otu_filepath,
                  self.ref_seq_filepath)
        self.assertEqual(obs, exp)


class TopLevelTests(SharedSetupTestCase):

    """Tests of top-level functions"""

    def test_first(self):
        """first should always return first item"""
        vals = [3, 4, 2]
        self.assertEqual(first(vals), 3)
        vals.reverse()
        self.assertEqual(first(vals), 2)

    def test_first_id(self):
        """first_id should return first id from list"""
        ids = \
            "R27DLI_4812 R27DLI_600  R27DLI_727  U1PLI_403   U1PLI_8969".split(
            )
        self.assertEqual(first_id(ids, {}), 'R27DLI_4812')

    def test_random_id(self):
        """random_id should return random id from list"""
        ids = \
            "R27DLI_4812 R27DLI_600  R27DLI_727  U1PLI_403   U1PLI_8969".split(
            )
        assert random_id(ids, {}) in ids
        # just test we got something from the list, don't add stochastic test

    def test_longest_id(self):
        """longest_id should return id associated with longest seq"""
        ids = \
            "R27DLI_4812 R27DLI_600  R27DLI_727  U1PLI_403   U1PLI_8969".split(
            )
        seqs = dict(parse_fasta(dna_seqs.splitlines(),
                                       label_to_name=label_to_name))
        self.assertEqual(longest_id(ids, seqs), 'U1PLI_403')

    def test_unique_id_map(self):
        """unique_id_map should return map of seqs:unique representatives"""
        seqs = {'a': 'AG', 'b': 'AG', 'c': 'CC', 'd': 'CT'}
        obs = unique_id_map(seqs)
        exp = {'c': ['c'], 'd': ['d'], 'a': ['a', 'b'], 'b': ['a', 'b']}
        # can't predict if a or b
        for k in obs:
            assert obs[k] in exp[k]

    def test_make_most_abundant(self):
        """make_most_abundant should return function with correct behavior"""
        ids = \
            "R27DLI_4812 R27DLI_600  R27DLI_727  U1PLI_403   U1PLI_8969".split(
            )
        seqs = dict(parse_fasta(dna_seqs.splitlines(),
                                       label_to_name=label_to_name))
        f = make_most_abundant(seqs)
        result = f(ids, seqs)
        assert result in ['R27DLI_4812', 'R27DLI_727', 'U1PLI_8969']


dna_seqs = """>R27DLI_4812 FMSX0OV01EIYV5 orig_bc=CTTGATGCGTAT new_bc=CTTGATGCGTAT bc_diffs=0
CTGGGCCGTATCTC
>R27DLI_600 FMSX0OV01D110Y orig_bc=CTTGATGCGTAT new_bc=CTTGATGCGTAT bc_diffs=0
CTGGGCCGTATCTCA
>R27DLI_727 FMSX0OV01D5X55 orig_bc=CTTGATGCGTAT new_bc=CTTGATGCGTAT bc_diffs=0
CTGGGCCGTATCTC
>U1PLI_403 FMSX0OV01DVG99 orig_bc=TACAGATGGCTC new_bc=TACAGATGGCTC bc_diffs=0
CTGGGCCGTATCTCAGTCCCAA
>U1PLI_8969 FMSX0OV01ARWY7 orig_bc=TACAGATGGCTC new_bc=TACAGATGGCTC bc_diffs=0
CTGGGCCGTATCTC
>U1PLI_9080 FMSX0OV01C9JUX orig_bc=TACAGATGGCTC new_bc=TACAGATGGCTC bc_diffs=0
CTGGGCCG
>U1PLI_9526 FMSX0OV01EUN7B orig_bc=TACAGATGGCTC new_bc=TACAGATGGCTC bc_diffs=0
CTGGGCCGTATCTCAGTCCCAATGTGGCCGGTCGGTCTCTCAACCCGGCTACCCATCGCGGGCTAGGTGGGCCGTTACCCCGCCTACTACCTAATGGGCCGCGACCCCATCCCTTGCCGTCTGGGCTTTCCCGGGCCCCCCAGGAGGGGGGCGAGGAGTATCCGGTATTAGCCTCGGTTTCCCAAGGTTGTCCCGGAGCAAGGGGCAGGTTGGTCACGTGTTACTCACCCGTTCGCCACTTCATGTCCGCCCGAGGGCGGTTTCATCG
>W3Cecum_6642 FMSX0OV01CW7FI orig_bc=GATACGTCCTGA new_bc=GATACGTCCTGA bc_diffs=0
CTGGGCCGTATCTCAGT
>W3Cecum_8992 FMSX0OV01C3YXK orig_bc=GATACGTCCTGA new_bc=GATACGTCCTGA bc_diffs=0
CTGGGCCGTGTCTC
>U1PLI_7889 FMSX0OV01C6HRL orig_bc=TACAGATGGCTC new_bc=TACAGATGGCTC bc_diffs=0
TTGGACCGTG
>W3Cecum_4858 FMSX0OV01BX4KM orig_bc=GATACGTCCTGA new_bc=GATACGTCCTGA bc_diffs=0
TTGGGCCGTGTCTCAGT
>R27DLI_3243 FMSX0OV01DH41R orig_bc=CTTGATGCGTAT new_bc=CTTGATGCGTAT bc_diffs=0
CTGGACCGTGTCT
>R27DLI_4562 FMSX0OV01EJKLT orig_bc=CTTGATGCGTAT new_bc=CTTGATGCGTAT bc_diffs=0
CTGGACCGTGTCT
>R27DLI_6828 FMSX0OV01BCWTL orig_bc=CTTGATGCGTAT new_bc=CTTGATGCGTAT bc_diffs=0
CTGGACCGTGTCT
>R27DLI_9097 FMSX0OV01APUV6 orig_bc=CTTGATGCGTAT new_bc=CTTGATGCGTAT bc_diffs=0
CTGGACCGTGTCT
>U1PLI_2780 FMSX0OV01E2K1S orig_bc=TACAGATGGCTC new_bc=TACAGATGGCTC bc_diffs=0
CTGGACCGTGTCTC
>U1PLI_67 FMSX0OV01DO1NS orig_bc=TACAGATGGCTC new_bc=TACAGATGGCTC bc_diffs=0
CTGGACCGTGT
>U9PSI_10475 FMSX0OV01BB4Q3 orig_bc=GATAGCTGTCTT new_bc=GATAGCTGTCTT bc_diffs=0
CTGGACCGTGTCTC
>U9PSI_4341 FMSX0OV01B8SXV orig_bc=GATAGCTGTCTT new_bc=GATAGCTGTCTT bc_diffs=0
CTGGACCGTGTCT
>W3Cecum_5191 FMSX0OV01BMU6R orig_bc=GATACGTCCTGA new_bc=GATACGTCCTGA bc_diffs=0
CTGGACCGTGTCT
"""

otus = """0	R27DLI_4812	R27DLI_600	R27DLI_727	U1PLI_403	U1PLI_8969	U1PLI_9080	U1PLI_9526	W3Cecum_6642	W3Cecum_8992
1	U1PLI_7889
2	W3Cecum_4858
3	R27DLI_3243	R27DLI_4562	R27DLI_6828	R27DLI_9097	U1PLI_2780	U1PLI_67	U9PSI_10475	U9PSI_4341	W3Cecum_5191
"""

rep_seqs_result_file_exp = """>0 R27DLI_4812
CTGGGCCGTATCTC
>1 U1PLI_7889
TTGGACCGTG
>2 W3Cecum_4858
TTGGGCCGTGTCTCAGT
>3 R27DLI_3243
CTGGACCGTGTCT
"""

rep_seqs_result_file_sorted_exp = """>3 R27DLI_3243
CTGGACCGTGTCT
>0 R27DLI_4812
CTGGGCCGTATCTC
>2 W3Cecum_4858
TTGGGCCGTGTCTCAGT
>1 U1PLI_7889
TTGGACCGTG
"""

otus_w_ref = """0	R27DLI_4812	R27DLI_600	R27DLI_727	U1PLI_403	U1PLI_8969	U1PLI_9080	U1PLI_9526	W3Cecum_6642	W3Cecum_8992
ref1	U1PLI_7889
2	W3Cecum_4858
ref0	R27DLI_3243	R27DLI_4562	R27DLI_6828	R27DLI_9097	U1PLI_2780	U1PLI_67	U9PSI_10475	U9PSI_4341	W3Cecum_5191
"""

otus_all_ref = """ref1	U1PLI_7889
ref0	R27DLI_3243	R27DLI_4562	R27DLI_6828	R27DLI_9097	U1PLI_2780	U1PLI_67	U9PSI_10475	U9PSI_4341	W3Cecum_5191
"""

reference_seqs = """>ref0
CCCAAAAAAATTTTTT
>ref1 some comment
GGGGGGGAAAAAAAAAAAAA
>ref2
CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCAAAA
"""

rep_seqs_reference_result_file_exp = """>0 R27DLI_4812
CTGGGCCGTATCTC
>ref1 ref1
GGGGGGGAAAAAAAAAAAAA
>2 W3Cecum_4858
TTGGGCCGTGTCTCAGT
>ref0 ref0
CCCAAAAAAATTTTTT
"""


# run unit tests if run from command-line
if __name__ == '__main__':
    main()