Repository: cogent3/cogent3
Pull request #1424 (2023-05-09 03:55:24)

ENH: sketch implementation of SequenceCollection.distance_matrix()


Patch:
diff --git a/src/cogent3/app/composable.py b/src/cogent3/app/composable.py
index 8777a83ce..63abeb8c0 100644
--- a/src/cogent3/app/composable.py
+++ b/src/cogent3/app/composable.py
@@ -460,7 +460,7 @@ def _class_from_func(func):
 
 def define_app(
     klass=None, *, app_type: AppType = GENERIC, skip_not_completed: bool = True
-):
+) -> type:
     """decorator for building callable apps
 
     Parameters
diff --git a/src/cogent3/app/dist.py b/src/cogent3/app/dist.py
index bb3426235..f30a54dfb 100644
--- a/src/cogent3/app/dist.py
+++ b/src/cogent3/app/dist.py
@@ -2,15 +2,23 @@ import itertools
 
 from typing import Union
 
+from numpy import log, polyval, triu_indices, zeros_like
+
 from cogent3 import get_moltype, make_tree
 from cogent3.evolve.fast_distance import (
     DistanceMatrix,
     get_distance_calculator,
 )
 from cogent3.evolve.models import get_model
+from cogent3.maths.distance_transform import jaccard
 
 from .composable import define_app
-from .typing import AlignedSeqsType, PairwiseDistanceType, SerialisableType
+from .typing import (
+    AlignedSeqsType,
+    PairwiseDistanceType,
+    SerialisableType,
+    UnalignedSeqsType,
+)
 
 
 __author__ = "Gavin Huttley"
@@ -23,6 +31,25 @@ __email__ = "Gavin.Huttley@anu.edu.au"
 __status__ = "Alpha"
 
 
+# The following coefficients are derived from a polynomial fit between Jaccard distance
+# and the proportion of different sites for mammalian DNA sequences. NOTE: the Jaccard
+# distance used kmers where k=10.
+JACCARD_PDIST_POLY_COEFFS = [
+    2271.7714153914335,
+    -11998.34362001251,
+    27525.142573445955,
+    -35922.0159776342,
+    29337.5102940838,
+    -15536.693064681693,
+    5346.929667208838,
+    -1165.616998965176,
+    151.8581396241204,
+    -10.489082251524346,
+    0.3334853259953467,
+    0.0,
+]
+
+
 @define_app
 class fast_slow_dist:
     """Pairwise distance calculation for aligned sequences.
@@ -125,3 +152,126 @@ class fast_slow_dist:
 def get_fast_slow_calc(distance, **kwargs):
     """returns FastSlow instance for a given distance name"""
     return fast_slow_dist(distance, **kwargs)
+
+
+@define_app
+def jaccard_dist(seq_coll: UnalignedSeqsType, k: int = 10) -> PairwiseDistanceType:
+    """Calculates the pairwise Jaccard distance between the sets of kmers generated from
+    sequences in the collection. A measure of distance for unaligned sequences.
+
+    Parameters
+    ----------
+    seq_coll: UnalignedSeqsType
+        a collection of unaligned sequences
+    k:  int
+        size of kmer to use for
+
+    Returns
+    -------
+    DistanceMatrix
+        returns pairwise Jaccard distance between sequences in the collection.
+    """
+
+    kmers = {name: set(seq.get_kmers(k)) for name, seq in seq_coll.named_seqs.items()}
+    seq_names = sorted(kmers.keys())
+    num_seqs = len(seq_names)
+
+    jaccard_dists = {}
+
+    for i in range(num_seqs):
+        for j in range(i):
+            name1, name2 = seq_names[i], seq_names[j]
+            dist = jaccard(kmers[name1], kmers[name2])
+            jaccard_dists[(name1, name2)] = dist
+            jaccard_dists[(name2, name1)] = dist
+
+    return DistanceMatrix(jaccard_dists)
+
+
+@define_app()
+def approx_pdist(jaccard_dists: PairwiseDistanceType) -> PairwiseDistanceType:
+    """Converts Jaccard distances to approximate pairwise distances using coefficient from
+    a pre-determined polynomial fit.
+
+    NOTE: coefficients are derived from a polynomial fit between Jaccard distance of kmers
+    with k=10 and the proportion of sites different for mammalian DNA sequences.
+
+    Parameters
+    ----------
+    jaccard_dists : DistanceMatrix
+    The pairwise Jaccard distance matrix
+
+    Returns
+    -------
+    DistanceMatrix
+    The pairwise approximate PDist matrix
+    """
+    j_dists = jaccard_dists.array
+
+    # Initialise an array of the same size as j_dists with all values = 0.0
+    p_dists = zeros_like(j_dists)
+
+    # The matrix is symmetric across the diagonal, and we only want
+    # to do calculations once, so grab the indices of the upper triangle,
+    # setting k=1 will exclude the diagonal
+    upper_indices = triu_indices(n=j_dists.shape[0], k=1)
+
+    # Convert only the upper indices from Jaccard distance to approximate PDist
+    upper_vals = polyval(JACCARD_PDIST_POLY_COEFFS, j_dists[upper_indices])
+    p_dists[upper_indices] = upper_vals
+
+    # Reflect the upper triangle to the lower triangle
+    lower_indices = (upper_indices[1], upper_indices[0])
+    p_dists[lower_indices] = upper_vals
+
+    # add dists to dictionary where {(seq, seq) : dist}
+    names = jaccard_dists.names
+    named_dists = {
+        (names[i], names[j]): p_dists[i, j]
+        for i, j in itertools.combinations(range(len(names)), 2)
+    }
+
+    return DistanceMatrix(named_dists)
+
+
+@define_app
+def approx_jc69(
+    pdist_predicted: PairwiseDistanceType,
+) -> PairwiseDistanceType:
+    """takes pairwise predicted p-distances and returns pairwise JC69 distances
+
+    Parameters
+    ----------
+    pdist_predicted
+        The pairwise approximate PDist matrix
+
+    Returns
+    -------
+    DistanceMatrix of pairwise JC69 distances
+
+    """
+    pdists = pdist_predicted.array
+    jc_dists = zeros_like(pdists)
+
+    # calculate approx jc dist from approx pdist for upper triangle of matrix
+    upper_indices = triu_indices(n=pdists.shape[0], k=1)
+    upper_vals = _jc69_from_pdist(pdists[upper_indices])
+    jc_dists[upper_indices] = upper_vals
+
+    # reflect into lower triangle of matrix
+    lower_indices = (upper_indices[1], upper_indices[0])
+    jc_dists[lower_indices] = upper_vals
+
+    # add dists to dictionary where {(seq, seq) : dist}, so can be wrapped in DistanceMatrix
+    names = pdist_predicted.names
+    named_dists = {
+        (names[i], names[j]): jc_dists[i, j]
+        for i, j in itertools.combinations(range(len(names)), 2)
+    }
+
+    return DistanceMatrix(named_dists)
+
+
+def _jc69_from_pdist(p):
+    """convert proportion of sites different to Jukes Cantor distance"""
+    return -3.0 * log(1 - (4 / 3) * p) / 4
diff --git a/src/cogent3/core/alignment.py b/src/cogent3/core/alignment.py
index 34153228e..76f458495 100644
--- a/src/cogent3/core/alignment.py
+++ b/src/cogent3/core/alignment.py
@@ -2285,6 +2285,51 @@ class SequenceCollection(_SequenceCollectionBase):
                 # passing self only used when self is an Alignment
                 yield seq.make_feature(feature, self)
 
+    def distance_matrix(self, calc="pdist"):
+        """estimated pairwise distance between sequences. Distance calculation can be either:
+
+        pdist:  an approximation of the Proportional Sites Different, estimated by calculating
+                the Jaccard distance using kmers with k=10, then transforming with coefficients
+                from a pre-determined polynomial fit between Jaccard distance and pdist.
+        jc69:   an approximation of the Jukes Cantor distance using the approx pdist. i.e.,
+                a transformation of the above using Jukes Cantor distance.
+
+        Parameters
+        ----------
+        calc : str
+            The distance calculation method to use, either "pdist" or "jc69"
+
+        Returns
+        -------
+        DistanceMatrix
+            Estimated pairwise distances between sequences in the collection
+        """
+
+        # check moltype
+        if not len(self.moltype.alphabet) == 4:
+            raise NotImplementedError("only defined for DNA/RNA molecular types")
+
+        # assert we have more than one sequence in the SequenceCollection
+        if self.num_seqs == 1:
+            raise ValueError(
+                "Pairwise distance cannot be computed for a single sequence. Please provide at least two sequences."
+            )
+
+        jdist = cogent3.get_app("jaccard_dist")
+        pdist = cogent3.get_app("approx_pdist")
+
+        if calc == "pdist":
+            dist_calc_app = jdist + pdist
+        elif calc == "jc69":
+            jc69dist = cogent3.get_app("approx_jc69")
+            dist_calc_app = jdist + pdist + jc69dist
+        else:
+            raise ValueError(
+                f"No support for calc={calc}. Use either 'pdist' or 'jc69'"
+            )
+
+        return dist_calc_app(self)
+
 
 @total_ordering
 class Aligned:


Test patch:
diff --git a/tests/test_app/test_dist.py b/tests/test_app/test_dist.py
index d7ab5679d..b67b490bf 100644
--- a/tests/test_app/test_dist.py
+++ b/tests/test_app/test_dist.py
@@ -1,11 +1,31 @@
+import itertools
+
 from tempfile import TemporaryDirectory
 from unittest import TestCase, main
 
+import pytest
+
+from numpy import polyval
 from numpy.testing import assert_allclose
 
-from cogent3 import DNA, PROTEIN, get_app, make_unaligned_seqs, open_data_store
+from cogent3 import (
+    DNA,
+    PROTEIN,
+    get_app,
+    make_aligned_seqs,
+    make_unaligned_seqs,
+    open_data_store,
+)
 from cogent3.app.composable import WRITER
-from cogent3.evolve.fast_distance import HammingPair, TN93Pair
+from cogent3.app.dist import (
+    JACCARD_PDIST_POLY_COEFFS,
+    _jc69_from_pdist,
+    approx_jc69,
+    approx_pdist,
+    jaccard_dist,
+)
+from cogent3.evolve.fast_distance import DistanceMatrix, HammingPair, TN93Pair
+from cogent3.maths.distance_transform import jaccard
 
 
 __author__ = "Gavin Huttley"
@@ -43,6 +63,16 @@ _seqs4 = {
 _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 [
@@ -242,3 +272,192 @@ class FastSlowDistTests(TestCase):
 
 if __name__ == "__main__":
     main()
+
+
+@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)),
+            set(seqs.get_seq(seq2).get_kmers(k=10)),
+        )
+        assert got == expect
+
+
+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 = _jc69_from_pdist(pdists[(seq1, seq2)])
+        assert got == expect
diff --git a/tests/test_core/test_alignment.py b/tests/test_core/test_alignment.py
index d102c9817..de62255ad 100644
--- a/tests/test_core/test_alignment.py
+++ b/tests/test_core/test_alignment.py
@@ -6,7 +6,7 @@ import re
 
 from os import remove
 from tempfile import TemporaryDirectory, mktemp
-from unittest import TestCase, main
+from unittest import TestCase
 
 import numpy
 import pytest
@@ -21,6 +21,7 @@ from cogent3 import (
     load_unaligned_seqs,
     make_aligned_seqs,
     make_seq,
+    make_unaligned_seqs,
     open_,
 )
 from cogent3.core.alignment import (
@@ -3428,3 +3429,45 @@ def test_get_seq_entropy(cls):
     entropy = a.entropy_per_seq()
     e = 0.81127812445913283  # sum(p log_2 p) for p = 0.25, 0.75
     assert_allclose(entropy, array([1, 0, e]))
+
+
+@pytest.mark.parametrize("moltype", ("dna", "rna"))
+def test_distance_matrix_singleton_collection(moltype):
+    """SequenceCollection.distance_matrix() should raise error if collection
+    only contains a single sequence"""
+    collection = make_unaligned_seqs(data={"s1": "ACGTACGTAGTCGCG"}, moltype=moltype)
+    with pytest.raises(ValueError):
+        _ = collection.distance_matrix()
+
+
+@pytest.mark.parametrize("moltype", ("dna", "rna"))
+def test_collection_distance_matrix_same_seq(moltype):
+    """Identical seqs should return distance measure of 0.0"""
+    data = dict(
+        [("s1", "ACGTACGTAGTCGCG"), ("s2", "GTGTACGTATCGCG"), ("s3", "GTGTACGTATCGCG")]
+    )
+    collection = make_unaligned_seqs(data=data, moltype=moltype)
+    dists = collection.distance_matrix(calc="pdist")
+
+    # all comparison of a sequence to itself should be zero
+    for seq in collection.names:
+        assert dists[(seq, seq)] == 0.0
+
+    # s2 and s3 are identical, so should be zero
+    assert dists[("s2", "s3")] == 0.0
+    assert dists[("s3", "s2")] == 0.0
+
+
+@pytest.mark.parametrize("moltype", ("protein", "text", "bytes"))
+def test_distance_matrix_fails_wrong_moltype(moltype):
+    data = [("s1", "ACGTA"), ("s2", "ACGTA")]
+    seqs = make_unaligned_seqs(data=data, moltype=moltype)
+    with pytest.raises(NotImplementedError):
+        seqs.distance_matrix()
+
+
+@pytest.mark.parametrize("moltype", ("dna", "rna"))
+def test_distance_matrix_passes_correct_moltype(moltype):
+    data = [("s1", "ACGTA"), ("s2", "ACGTA")]
+    seqs = make_unaligned_seqs(data=data, moltype=moltype)
+    seqs.distance_matrix()