Repository: cogent3/cogent3
Pull request #2105 (2024-11-12 00:11:49)

ENH: add quick_tree, distance_matrix to new Alignment
## Summary by Sourcery

Add new methods 'distance_matrix' and 'quick_tree' to the Alignment class for calculating pairwise distances and generating phylogenetic trees, respectively. Enhance the 'get_translation' method with a 'name_map' parameter. Introduce comprehensive tests for the new methods and sequence renaming functionality.

New Features:
- Introduce the 'distance_matrix' method to calculate pairwise distances between sequences in an alignment, with options for different calculators and handling of invalid data.
- Add the 'quick_tree' method to generate a phylogenetic tree from an alignment, with support for bootstrapping and handling of invalid data.

Enhancements:
- Improve the 'get_translation' method by adding a 'name_map' parameter to better handle sequence naming.

Tests:
- Add tests for the 'distance_matrix' method to ensure correct calculation of pairwise distances, including handling of invalid data.
- Introduce tests for the 'quick_tree' method to verify the generation of phylogenetic trees and the handling of bootstrap replicates.
- Enhance tests for sequence renaming in the 'iter_seqs' method and the 'get_translation' method to ensure correct functionality.


Patch:
diff --git a/src/cogent3/core/moltype.py b/src/cogent3/core/moltype.py
index eff2c3f7c..86810ec66 100644
--- a/src/cogent3/core/moltype.py
+++ b/src/cogent3/core/moltype.py
@@ -1541,13 +1541,15 @@ def get_moltype(name, new_type: bool = False):
         The default will be changed to True in 2024.12. Support for the old
         style will be removed as of 2025.6.
     """
-    if new_type or "COGENT3_NEW_TYPE" in os.environ:
-        from cogent3.core.new_moltype import get_moltype as new_get_moltype
-
-        return new_get_moltype(name=name)
+    from cogent3.core import new_moltype
 
+    if new_type or "COGENT3_NEW_TYPE" in os.environ:
+        return new_moltype.get_moltype(name=name)
     if isinstance(name, MolType):
         return name
+    if isinstance(name, new_moltype.MolType):
+        return name
+
     name = name.lower()
     if name not in moltypes:
         raise ValueError(f"unknown moltype {name!r}")
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index d4b1ee6e9..f3448fe2e 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -1017,7 +1017,7 @@ class SequenceCollection:
                 include_stop=include_stop,
                 trim_stop=trim_stop,
             )
-            translated[seq.name] = numpy.array(pep)
+            translated[seq._seq.seqid] = numpy.array(pep)
 
         pep_moltype = new_moltype.get_moltype(
             "protein_with_stop" if include_stop else "protein"
@@ -1030,6 +1030,7 @@ class SequenceCollection:
         return self.__class__(
             seqs_data=seqs_data,
             moltype=pep_moltype,
+            name_map=self._name_map,
             info=self.info,
             source=self.source,
             is_reversed=False,  # reversed not meaningful for a protein
@@ -4426,12 +4427,18 @@ class Alignment(SequenceCollection):
             info=self.info,
         )
 
-    def matching_ref(self, ref_name, gap_fraction, gap_run):
+    def matching_ref(self, ref_name: str, gap_fraction: float, gap_run: int):
         """Returns new alignment with seqs well aligned with a reference.
 
-        gap_fraction = fraction of positions that either have a gap in the
+        Parameters
+        ----------
+        ref_name
+            name of the sequence to use as the reference
+        gap_fraction
+            fraction of positions that either have a gap in the
             template but not in the seq or in the seq but not in the template
-        gap_run = number of consecutive gaps tolerated in query relative to
+        gap_run
+            number of consecutive gaps tolerated in query relative to
             sequence or sequence relative to query
         """
         template = self.seqs[ref_name]
@@ -4680,6 +4687,117 @@ class Alignment(SequenceCollection):
             seqs_data=new_seqs_data, info=self.info, moltype=self.moltype
         )
 
+    def distance_matrix(
+        self,
+        calc: str = "pdist",
+        show_progress: bool = False,
+        drop_invalid: bool = False,
+    ):
+        """Returns pairwise distances between sequences.
+
+        Parameters
+        ----------
+        calc
+            a pairwise distance calculator or name of one. For options see
+            cogent3.evolve.fast_distance.available_distances
+        show_progress
+            controls progress display for distance calculation
+        drop_invalid
+            If True, sequences for which a pairwise distance could not be
+            calculated are excluded. If False, an ArithmeticError is raised if
+            a distance could not be computed on observed data.
+        """
+        from cogent3.evolve.fast_distance import get_distance_calculator
+
+        try:
+            calculator = get_distance_calculator(
+                calc,
+                moltype=self.moltype,
+                alignment=self,
+                invalid_raises=not drop_invalid,
+            )
+            calculator.run(show_progress=show_progress)
+        except ArithmeticError as e:
+            msg = "not all pairwise distances could be computed, try drop_invalid=True"
+            raise ArithmeticError(msg) from e
+        result = calculator.get_pairwise_distances()
+        if drop_invalid:
+            result = result.drop_invalid()
+        return result
+
+    @UI.display_wrap
+    def quick_tree(
+        self,
+        calc: str = "pdist",
+        bootstrap: OptInt = None,
+        drop_invalid: bool = False,
+        show_progress: bool = False,
+        ui=None,  # refactor: type hint
+    ):
+        """Returns a phylogenetic tree.
+
+        Parameters
+        ----------
+        calc
+            a pairwise distance calculator or name of one. For options see
+            cogent3.evolve.fast_distance.available_distances
+        bootstrap
+            Number of non-parametric bootstrap replicates. Resamples alignment
+            columns with replacement and builds a phylogeny for each such
+            resampling.
+        drop_invalid
+            If True, sequences for which a pairwise distance could not be
+            calculated are excluded. If False, an ArithmeticError is raised if
+            a distance could not be computed on observed data.
+        show_progress
+            controls progress display for distance calculation
+
+        Returns
+        -------
+        a phylogenetic tree. If bootstrap specified, returns the weighted
+        majority consensus. Support for each node is stored as
+        edge.params['params'].
+
+        Notes
+        -----
+        Sequences in the observed alignment for which distances could not be
+        computed are omitted. Bootstrap replicates are required to have
+        distances for all seqs present in the observed data distance matrix.
+        """
+        from cogent3.phylo.consensus import weighted_majority_rule
+        from cogent3.phylo.nj import gnj
+
+        dm = self.distance_matrix(
+            calc=calc, show_progress=show_progress, drop_invalid=drop_invalid
+        )
+        results = gnj(dm, keep=1, show_progress=show_progress)
+        kept_names = set(dm.template.names[0])
+        if bootstrap:
+            subaln = (
+                self.take_seqs(kept_names) if set(self.names) != kept_names else self
+            )
+            for _ in ui.series(range(bootstrap), count=bootstrap, noun="bootstrap"):
+                b = subaln.sample(with_replacement=True)
+                try:
+                    bdist = b.distance_matrix(
+                        calc=calc, show_progress=show_progress, drop_invalid=False
+                    )
+                except ArithmeticError:
+                    # all sequences must have a pairwise distance to allow
+                    # constructing a consensus tree
+                    continue
+                bresult = gnj(bdist, keep=1, show_progress=show_progress)
+                results.extend(bresult)
+
+            consense = weighted_majority_rule(results)
+            assert len(consense) == 1
+            results = consense
+
+        tree = results[0]
+        if not bootstrap:
+            tree = tree[1]
+        return tree
+
     def trim_stop_codons(self, gc: Any = None, strict: bool = False, **kwargs):
         # refactor: array
         if not self.has_terminal_stop(gc=gc, strict=strict):


Test patch:
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index c0425b81d..987a2a7f7 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -7,7 +7,7 @@ from warnings import catch_warnings, filterwarnings
 import numpy
 import pytest
 
-from cogent3 import get_app, load_unaligned_seqs, open_
+from cogent3 import get_app, load_aligned_seqs, load_unaligned_seqs, open_
 from cogent3._version import __version__
 from cogent3.core import new_alignment, new_alphabet, new_moltype, new_sequence
 from cogent3.core.annotation import Feature
@@ -726,6 +726,14 @@ def test_sequence_collection_iter_seqs_ragged(ragged):
     assert seqs[0] == ragged.seqs["b"]
 
 
+def test_sequence_collection_iter_seqs_renamed(ragged_padded_dict):
+    """SequenceCollection iter_seqs() method should support renaming of seqs"""
+    coll = new_alignment.make_unaligned_seqs(ragged_padded_dict, moltype="dna")
+    renamed = coll.rename_seqs(renamer=lambda x: f"seq_{x}")
+    seqs = list(renamed.iter_seqs())
+    assert seqs == ["AAAAAA", "AAA---", "AAAA--"]
+
+
 def test_sequence_collection_repr():
     data = {
         "ENSMUSG00000056468": "GCCAGGGGGAAAAGGGAGAA",
@@ -1324,6 +1332,17 @@ def test_get_translation(seqs, mk_cls):
     assert got.moltype == new_moltype.PROTEIN
 
 
+@pytest.mark.parametrize(
+    "mk_cls",
+    (new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs),
+)
+def test_get_translation_renamed(seqs, mk_cls):
+    seqs = mk_cls({"s1": "GATTTT", "s2": "GATCTT"}, moltype="dna")
+    renamed = seqs.rename_seqs(lambda x: f"{x}_renamed")
+    got = renamed.get_translation(incomplete_ok=True)
+    assert got.names == ["s1_renamed", "s2_renamed"]
+
+
 @pytest.mark.parametrize(
     "mk_cls",
     (new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs),
@@ -4148,7 +4167,7 @@ def test_sliding_windows():
     assert result[3].to_dict() == {"seq3": "TACGT", "seq2": "TACGT", "seq1": "TACGT"}
 
 
-def test_get_feature():
+def test_alignment_get_feature():
     aln = new_alignment.make_aligned_seqs(
         {"x": "-AAAAAAAAA", "y": "TTTT--CCCT"}, moltype="dna"
     )
@@ -4159,6 +4178,33 @@ def test_get_feature():
     assert feat.get_slice().to_dict() == dict(x="AAA", y="CCT")
 
 
+def test_alignment_distance_matrix():
+    """Alignment distance_matrix should produce correct scores"""
+    data = dict([("s1", "ACGTACGTA"), ("s2", "GTGTACGTA")])
+    aln = new_alignment.make_aligned_seqs(data, moltype="dna")
+    dists = aln.distance_matrix(calc="hamming", show_progress=False)
+    assert dists == {("s1", "s2"): 2.0, ("s2", "s1"): 2.0}
+    # and for protein
+    aa = aln.get_translation()
+    dists = aa.distance_matrix(calc="hamming")
+    assert dists == {("s1", "s2"): 1.0, ("s2", "s1"): 1.0}
+
+    # when there are invalid data
+    data = dict(
+        seq1="AGGGGGGGGGGCCCCCCCCCCCCCCCCCGGGGGGGGGGGGGGGCGGTTTTTTTTTTTTTTTTTT",
+        seq2="TAAAAAAAAAAGGGGGGGGGGGGGGGGGGTTTTTTTTTTTTTTTTTTCCCCCCCCCCCCCCCCC",
+        seq3="TACAAAAAAAAGGGGCGGGGGGGGGGGGGTTTTTTTTTTTTTTTTTTCCCCCCCCCCCCCCCCC",
+    )
+
+    aln = new_alignment.make_aligned_seqs(data, moltype="dna")
+    with pytest.raises(ArithmeticError):
+        # default settings cause an exception
+        _ = aln.distance_matrix(calc="paralinear")
+    # but setting drop_invalid=False allows calc
+    dists = aln.distance_matrix(calc="paralinear", drop_invalid=True)
+    assert dists is None
+
+
 def test_alignment_sample_with_replacement():
     # test with replacement -- just verify that it rnus
     alignment = new_alignment.make_aligned_seqs(
@@ -4466,6 +4512,31 @@ def test_omit_gap_pos():
     assert new_aln.omit_gap_pos(0) is None
 
 
+@pytest.fixture(scope="session")
+def brca1_data():
+    return load_aligned_seqs("data/brca1.fasta", moltype="dna", new_type=True).to_dict()
+
+
+@pytest.mark.parametrize("calc", ("hamming", None))
+def test_alignment_quick_tree(calc, brca1_data):
+    """quick tree method returns tree"""
+    aln = new_alignment.make_aligned_seqs(brca1_data, moltype="dna")[:100]
+    aln = aln.take_seqs(["Human", "Rhesus", "HowlerMon", "Galago", "Mouse"])
+    kwargs = dict(show_progress=False)
+    if calc:
+        kwargs["calc"] = calc
+
+    # bootstrap
+    tree = aln.quick_tree(bootstrap=2, **kwargs)
+    assert set(tree.get_tip_names()) == set(aln.names)
+    types = {
+        type(float(edge.params["support"]))
+        for edge in tree.preorder()
+        if not edge.is_root()
+    }
+    assert types == {float}
+
+
 @pytest.fixture
 def alignment():
     data = {