import itertools from tempfile import TemporaryDirectory from unittest import TestCase import pytest from numpy import log, polyval from numpy.testing import assert_allclose from cogent3 import ( DNA, PROTEIN, get_app, load_aligned_seqs, make_aligned_seqs, make_unaligned_seqs, open_data_store, ) from cogent3.app.composable import WRITER, NotCompleted from cogent3.app.dist import ( JACCARD_PDIST_POLY_COEFFS, approx_jc69, approx_pdist, jaccard_dist, ) from cogent3.evolve.fast_distance import DistanceMatrix, HammingPair, TN93Pair from cogent3.maths.distance_transform import jaccard _seqs1 = { "Human": "GCCAGCTCATTACAGCATGAGAACAGCAGTTTATTACTCACT", "Bandicoot": "NACTCATTAATGCTTGAAACCAGCAGTTTATTGTCCAAC", "Rhesus": "GCCAGCTCATTACAGCATGAGAACAGTTTGTTACTCACT", "FlyingFox": "GCCAGCTCTTTACAGCATGAGAACAGTTTATTATACACT", } _seqs2 = { "Human": "ATGCGGCTCGCGGAGGCCGCGCTCGCGGAG", "Mouse": "ATGCCCGGCGCCAAGGCAGCGCTGGCGGAG", "Opossum": "ATGCCAGTGAAAGTGGCGGCGGTGGCTGAG", } _seqs3 = { "Human": "ATGCGGCTCGCGGAGGCCGCGCTCGCGGAG", "Mouse": "ATGCCCGGCGCCAAGGCAGCGCTGGCGGAG", } _seqs4 = { "Human": "ATGCGGCTCGCGGAGGCCGCGCTCGCGGAG", "Opossum": "ATGCCAGTGAAAGTGGCGGCGGTGGCTGAG", } _seqs5 = {"Human": "ASSLQHENSSLLLT", "Bandicoot": "XSLMLETSSLLSN"} @pytest.fixture(scope="function") def _seqs1_collection(): return make_unaligned_seqs(data=_seqs1, moltype="dna") @pytest.fixture(scope="function") def _seqs2_collection(): return make_unaligned_seqs(data=_seqs2, moltype="dna") def _get_all_composable_apps(): out_dstore = open_data_store(":memory:", mode="w") return [ get_app("align_to_ref"), get_app("progressive_align", model="GY94"), get_app("fixed_length", 100), get_app("sample.min_length", 100), get_app("write_seqs", out_dstore), get_app( "omit_bad_seqs", ), get_app( "omit_degenerates", ), get_app("take_codon_positions", 1), get_app( "take_named_seqs", ), get_app("trim_stop_codons", gc=1), ] class FastSlowDistTests(TestCase): seqs1 = make_unaligned_seqs(_seqs1, moltype=DNA) seqs2 = make_unaligned_seqs(_seqs2, moltype=DNA) seqs3 = make_unaligned_seqs(_seqs3, moltype=DNA) seqs4 = make_unaligned_seqs(_seqs4, moltype=DNA) seqs5 = make_unaligned_seqs(_seqs5, moltype=PROTEIN) def test_init(self): """tests if fast_slow_dist can be initialised correctly""" fast_slow_dist = get_app("fast_slow_dist", fast_calc="hamming", moltype="dna") self.assertIsInstance(fast_slow_dist.fast_calc, HammingPair) self.assertIsNone(fast_slow_dist._sm) fast_slow_dist = get_app("fast_slow_dist", distance="TN93") self.assertIsInstance(fast_slow_dist.fast_calc, TN93Pair) self.assertEqual(fast_slow_dist._sm.name, "TN93") fast_slow_dist = get_app("fast_slow_dist", distance="GTR") self.assertEqual(fast_slow_dist._sm.name, "GTR") fast_slow_dist = get_app("fast_slow_dist", slow_calc="TN93") self.assertEqual(fast_slow_dist._sm.name, "TN93") self.assertIsNone(fast_slow_dist.fast_calc) with self.assertRaises(ValueError): fast_slow_dist = get_app( "fast_slow_dist", distance="TN93", fast_calc="TN93", slow_calc="TN93" ) with self.assertRaises(ValueError): fast_slow_dist = get_app("fast_slow_dist", fast_calc="GTR") with self.assertRaises(ValueError): fast_slow_dist = get_app("fast_slow_dist", slow_calc="hamming") def test_compatible_parameters(self): """tests if the input parameters are compatible with fast_slow_dist initialisation""" for kwargs in ( dict(fast_calc="hamming", moltype="dna"), dict(fast_calc="TN93"), dict(slow_calc="GTR"), dict(fast_calc="TN93"), ): _ = get_app("fast_slow_dist", **kwargs) def test_incompatible_parameters(self): """tests incompatible input parameters with fast_slow_dist initialisation""" for kwargs in ( dict(fast_calc="hamming"), dict(slow_calc="paralinear"), dict(fast_calc="GTR"), dict(slow_calc="hamming", moltype="dna"), ): with self.assertRaises(ValueError): _ = get_app("fast_slow_dist", **kwargs) def test_composable_apps(self): """tests two composable apps""" composable_apps = _get_all_composable_apps() calc_dist = get_app("fast_slow_dist", fast_calc="hamming", moltype="dna") for app in composable_apps: if app.app_type is WRITER: # cannot have a WRITER before a GENERIC continue # Compose two composable applications, there should not be exceptions. got = app + calc_dist self.assertIsInstance(got, type(calc_dist)) self.assertIs(got.input, app) self.assertIsInstance(got._data_types, frozenset) self.assertIsInstance(got._return_types, frozenset) self.assertIs(got.input, app) app.disconnect() calc_dist.disconnect() def test_est_dist_pair_slow(self): """tests the distance between seq pairs in aln""" aligner = get_app( "align_to_ref", ) aln3 = aligner(self.seqs3) fast_slow_dist = get_app("fast_slow_dist", slow_calc="GTR") got = fast_slow_dist(aln3).to_dict() assert_allclose(got[("Human", "Mouse")], got[("Mouse", "Human")]) self.assertTrue(got[("Mouse", "Human")] >= 0) fast_slow_dist = get_app("fast_slow_dist", slow_calc="TN93") got = fast_slow_dist(aln3).to_dict() assert_allclose(got[("Human", "Mouse")], got[("Mouse", "Human")]) self.assertTrue(got[("Mouse", "Human")] >= 0) aligner = get_app("align_to_ref", ref_seq="Human") aln3 = aligner(self.seqs3) fast_slow_dist = get_app("fast_slow_dist", slow_calc="GTR") got = fast_slow_dist(aln3).to_dict() assert_allclose(got[("Human", "Mouse")], got[("Mouse", "Human")]) fast_slow_dist = get_app("fast_slow_dist", slow_calc="TN93") got = fast_slow_dist(aln3).to_dict() assert_allclose(got[("Human", "Mouse")], got[("Mouse", "Human")]) self.assertTrue(got[("Mouse", "Human")] >= 0) aligner = get_app("align_to_ref", ref_seq="Mouse") aln3 = aligner(self.seqs3) fast_slow_dist = get_app("fast_slow_dist", slow_calc="GTR") got = fast_slow_dist(aln3).to_dict() self.assertTrue(got[("Mouse", "Human")] >= 0) fast_slow_dist = get_app("fast_slow_dist", slow_calc="TN93") got = fast_slow_dist(aln3).to_dict() self.assertTrue(got[("Mouse", "Human")] >= 0) aligner = get_app( "align_to_ref", ) aln3 = aligner(self.seqs4) fast_slow_dist = get_app("fast_slow_dist", slow_calc="GTR") got = fast_slow_dist(aln3).to_dict() self.assertTrue(got[("Human", "Opossum")] >= 0) fast_slow_dist = get_app("fast_slow_dist", slow_calc="TN93") got = fast_slow_dist(aln3).to_dict() self.assertTrue(got[("Human", "Opossum")] >= 0) aligner = get_app("align_to_ref", ref_seq="Human") aln3 = aligner(self.seqs4) fast_slow_dist = get_app("fast_slow_dist", slow_calc="GTR") got = fast_slow_dist(aln3).to_dict() self.assertTrue(got[("Human", "Opossum")] >= 0) fast_slow_dist = get_app("fast_slow_dist", slow_calc="TN93") got = fast_slow_dist(aln3).to_dict() self.assertTrue(got[("Human", "Opossum")] >= 0) aligner = get_app("align_to_ref", ref_seq="Opossum") aln3 = aligner(self.seqs4) fast_slow_dist = get_app("fast_slow_dist", slow_calc="GTR") got = fast_slow_dist(aln3).to_dict() self.assertTrue(got[("Human", "Opossum")] >= 0) fast_slow_dist = get_app("fast_slow_dist", slow_calc="TN93") got = fast_slow_dist(aln3).to_dict() self.assertTrue(got[("Human", "Opossum")] >= 0) # now as a process proc = get_app( "align_to_ref", ) + get_app("fast_slow_dist", fast_calc="hamming", moltype="dna") got = proc(self.seqs1) self.assertEqual(got[("Human", "Rhesus")], 1) treestring = "(Human:0.2,Bandicoot:0.2)" aligner = get_app("progressive_align", model="WG01", guide_tree=treestring) _ = aligner(self.seqs5) def test_composes_with_write_tabular(self): """correctly links to tabular""" with TemporaryDirectory(dir=".") as dirname: out_dstore = open_data_store(dirname, suffix="tsv", mode="w") writer = get_app("write_tabular", out_dstore) dist_calc = get_app("fast_slow_dist", distance="hamming", moltype="protein") _ = dist_calc + writer def test_functions_as_composable(self): """works as a composable app""" from pathlib import Path loader = get_app("load_aligned", moltype="dna", format="paml") dist = get_app("fast_slow_dist", "hamming", moltype="dna") with TemporaryDirectory(dir=".") as dirname: dirname = Path(dirname) out_dstore = open_data_store(dirname, suffix="tsv", mode="w") writer = get_app("write_tabular", out_dstore) proc = loader + dist + writer _ = proc("data/brca1_5.paml") output = dirname / "brca1_5.tsv" self.assertTrue(output.exists()) @pytest.mark.parametrize("moltype", ("dna", "rna")) def test_jaccard_dist(moltype): """jaccard_dist app should work for the simple case ("s1", "ACGTA"), ("s2", "----C"), with k=2 s1 kmers = "AC", "CG", "GT", "TA" s2 kmers = "AC", "CG", "GT", "TC" J(A,B) = 1 - |A ∩ B| / |A ∪ B| J(s1, s2) = 1 - |{"AC", "CG", "GT"}| / |{"AC", "CG", "GT", "TA", "TC"}| J(s1, s2) = 1 - 3 / 5 J(s1, s2) = 0.4 """ data = dict([("s1", "ACGTA"), ("s2", "ACGTC")]) collection = make_unaligned_seqs(data=data, moltype=moltype) jdist_k2 = jaccard_dist(k=2) dists = jdist_k2(collection) assert dists[("s1", "s2")] == 0.4 assert dists[("s2", "s1")] == 0.4 assert dists[("s1", "s1")] == 0.0 assert dists[("s2", "s2")] == 0.0 def test_approx_pdist(): """approx_pdist should work for the simple case y = polyval(JACCARD_PDIST_POLY_COEFFS, x) """ data = dict( [ (("s1", "s1"), 0.0), (("s1", "s2"), 0.4), (("s2", "s1"), 0.4), (("s2", "s2"), 0.0), ] ) dm = DistanceMatrix(data) pdist_app = approx_pdist() pdists = pdist_app(dm) expect_diff = polyval(JACCARD_PDIST_POLY_COEFFS, 0.4) expect_same = polyval(JACCARD_PDIST_POLY_COEFFS, 0.0) assert pdists[("s1", "s2")] == expect_diff assert pdists[("s2", "s1")] == expect_diff assert pdists[("s1", "s1")] == expect_same assert pdists[("s2", "s2")] == expect_same @pytest.mark.parametrize("moltype", ("dna", "rna")) def test_approx_jc69(moltype): """approx_jc69 should work the same as exact jc69 when given exact pdist""" seq_data = dict([("s1", "ACGAA"), ("s2", "ACGAC")]) aln = make_aligned_seqs(data=seq_data, moltype=moltype) expected = aln.distance_matrix(calc="jc69") data = dict( [ (("s1", "s1"), 0.0), (("s1", "s2"), 1 / 5), (("s2", "s1"), 1 / 5), (("s2", "s2"), 0.0), ] ) dm = DistanceMatrix(data) jc_dist_app = approx_jc69() got = jc_dist_app(dm) assert got[("s1", "s2")] == expected[("s1", "s2")] assert got[("s2", "s1")] == expected[("s2", "s1")] assert got[("s1", "s1")] == expected[("s1", "s1")] assert got[("s2", "s2")] == expected[("s2", "s2")] @pytest.mark.parametrize("moltype", ("dna", "rna")) def test_approx_pdist_same_diff(moltype): """comparisons between seqs with the same position different should be equal. comparison between seqs with more positions different should yield a higher measure than comparisons between seqs with fewer positions different. NOTE: the coefficients used in Jaccard to Pdist fit were generated using k=10, here I used k=3 ("s1", "ACGTA"), ("s2", "----C"), ("s3", "----T"), ("s4", "---AT"), """ data = dict( [ ("s1", "ACGTA"), ("s2", "ACGTC"), ("s3", "ACGTT"), ("s4", "ACGAT"), ] ) pdist_app = jaccard_dist(k=3) + approx_pdist() collection = make_unaligned_seqs(data=data, moltype=moltype) dists = pdist_app(collection) # comparisons with one position different should be smaller than those with two assert dists[("s1", "s2")] < dists[("s1", "s4")] assert dists[("s1", "s3")] < dists[("s1", "s4")] assert dists[("s1", "s2")] < dists[("s1", "s4")] assert dists[("s2", "s3")] < dists[("s2", "s4")] # both (s1 and s2) and (s2 and s3) have the same position different assert dists[("s1", "s2")] == dists[("s1", "s3")] def test_jaccard_dist_vals(_seqs1_collection): """values in the DistanceMatrix should match individually calculating the jaccard distance for pairs of sequence. """ seqs = _seqs1_collection jaccard_dist_app = jaccard_dist(k=10) jdists = jaccard_dist_app(seqs) names = jdists.names for i, j in itertools.combinations(range(len(names)), 2): seq1, seq2 = names[i], names[j] got = jdists[(seq1, seq2)] expect = jaccard( set(seqs.get_seq(seq1).get_kmers(k=10, strict=True)), set(seqs.get_seq(seq2).get_kmers(k=10, strict=True)), ) assert_allclose(got, expect) def test_jaccard_dist_one_seq(_seqs1_collection): """values in the DistanceMatrix should match individually calculating the jaccard distance for pairs of sequence. """ seqs = _seqs1_collection.take_seqs(_seqs1_collection.names[0]) jaccard_dist_app = jaccard_dist(k=10) got = jaccard_dist_app(seqs) assert isinstance(got, NotCompleted) assert got.origin == "jaccard_dist" def test_approx_pdist_vals(_seqs1_collection): """values in the DistanceMatrix should match individually calculating the pdist for pairs of sequence. testing integration of jaccard_dist() + approx_pdist() is identical to step-by-step calculation """ seqs = _seqs1_collection jaccard_dist_app = jaccard_dist(k=10) jdists = jaccard_dist_app(seqs) pdist_app = jaccard_dist(k=10) + approx_pdist() pdists = pdist_app(seqs) names = pdists.names for i, j in itertools.combinations(range(len(names)), 2): seq1, seq2 = names[i], names[j] got = pdists[(seq1, seq2)] expect = polyval(JACCARD_PDIST_POLY_COEFFS, jdists[(seq1, seq2)]) assert got == expect def test_approx_jc69_vals(_seqs1_collection): """values in the DistanceMatrix should match individually calculating the jc distance for pairs of sequence. testing integration of jaccard_dist() + approx_pdist() + approx_jc69() is identical to step-by-step calculation """ seqs = _seqs1_collection jaccard_dist_app = jaccard_dist(k=10) jdists = jaccard_dist_app(seqs) names = jdists.names pdist_app = approx_pdist() pdists = pdist_app(jdists) jc_app = jaccard_dist(k=10) + approx_pdist() + approx_jc69() jc_dists = jc_app(seqs) for i, j in itertools.combinations(range(len(names)), 2): seq1, seq2 = names[i], names[j] got = jc_dists[(seq1, seq2)] expect = -3 / 4 * log(1 - 4 / 3 * pdists[(seq1, seq2)]) assert got == expect def test_symmetry_of_dists(DATA_DIR): """distances are symmetric""" seqs = load_aligned_seqs(DATA_DIR / "primate_brca1.fasta", moltype="dna") dists = seqs.distance_matrix(calc="pdist") app = approx_jc69() got = app(dists) assert_allclose(got.array, got.array.T) def test_gap_dist(): app = get_app("gap_dist", gap_insert=10, gap_extend=1) # two sequences share a gap data = { "a": "TG----AATATGT------GAAAGAG", "b": "TTGAAGAATATGT------GAAAGAG", "c": "CTGAAGAACCTGTGAAAGTGAAAGAG", } aln = make_aligned_seqs(data, moltype="dna", array_align=True) expect = { ("a", "b"): 14.0, # one gap diff of size 4 ("a", "c"): 30.0, ("b", "c"): 16.0, } expect = DistanceMatrix(expect) dmat = app.main(aln) assert dmat.to_dict() == expect.to_dict() # shared gap actually not shared, 3 events data = { "a": "TG----AATATGTA-----GAAAGAG", "b": "TTGAAGAATATGTA------AAAGAG", "c": "CTGAAGAACCTGTGAAAGTGAAAGAG", } aln = make_aligned_seqs(data, moltype="dna", array_align=True) expect = { ("a", "b"): 45.0, # 3 gaps diff of size 15 ("a", "c"): 29.0, ("b", "c"): 16.0, } expect = DistanceMatrix(expect) dmat = app.main(aln) assert dmat.to_dict() == expect.to_dict() # additional gaps on either side of shared gap is two events data = { "a": "G--AG----A", "b": "TGGAGT--GA", "c": "TGGAGTGTGA", } aln = make_aligned_seqs(data, moltype="dna", array_align=True) expect = { ("a", "b"): 38, # 3 gaps diff of size 8 ("a", "c"): 26.0, ("b", "c"): 12.0, } expect = DistanceMatrix(expect) dmat = app.main(aln) assert dmat.to_dict() == expect.to_dict() data = {"a": "AAGAA-A", "b": "-ATAATG", "c": "C-TGG-G"} aln = make_aligned_seqs(data, moltype="dna", array_align=True) expect = { ("a", "b"): 22.0, # 2 gaps diff of size 2 ("a", "c"): 11.0, ("b", "c"): 33.0, # 3 gaps diff of size 2 } expect = DistanceMatrix(expect) dmat = app.main(aln) assert dmat.to_dict() == expect.to_dict()