Repository: cogent3/cogent3
Pull request #2090 (2024-11-01 05:10:44)

ENH: implement degap for alignment
NEW: Introduced a new abstract method get_ungapped in AlignedSeqsDataABC and its implementation in AlignedSeqsData to support ungapped sequence retrieval

## Summary by Sourcery

Implement the degap method to remove gaps and missing characters from sequences and introduce the get_ungapped method for ungapped sequence retrieval. Enhance the testing suite to cover these new functionalities.

New Features:
- Introduce a new abstract method get_ungapped in AlignedSeqsDataABC and its implementation in AlignedSeqsData to support ungapped sequence retrieval.

Enhancements:
- Enhance the degap method to remove both gaps and missing characters from sequences.

Tests:
- Add tests for the degap method to ensure it correctly removes gaps and missing characters from sequences, including tests for sliced alignments.


Patch:
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index 291a2ea4e..481f445ca 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -40,6 +40,7 @@ from cogent3.util.deserialise import deserialise_object, register_deserialiser
 from cogent3.util.dict_array import DictArray, DictArrayTemplate
 from cogent3.util.io import atomic_write, get_format_suffixes
 from cogent3.util.misc import (
+    extend_docstring_from,
     get_object_provenance,
     get_setting_from_environ,
     negate_condition,
@@ -901,16 +902,13 @@ class SequenceCollection:
         return json.dumps(self.to_rich_dict())
 
     def degap(self) -> SequenceCollection:
-        """Returns new collection in which sequences have no gaps.
+        """Returns new collection in which sequences have no gaps or missing
+        characters.
 
         Notes
         -----
         The returned collection will not retain an annotation_db if present.
         """
-        # refactor: design
-        # todo: this needs to be implemented for Alignment - need to think about the design.
-        # in the old implementation, this method removed all gaps and then returned a
-        # SequenceCollection (even for alignemnts).
         seqs = {
             name: self.moltype.degap(self._seqs_data.get_seq_array(seqid=name))
             for name in self._name_map.values()
@@ -2785,6 +2783,43 @@ class AlignedSeqsDataABC(SeqsDataABC):
         step: OptInt = None,
     ) -> bytes: ...
 
+    @abstractmethod
+    def get_ungapped(
+        self,
+        name_map: dict[str, str],
+        start: OptInt = None,
+        stop: OptInt = None,
+        step: OptInt = None,
+    ) -> tuple[dict, dict]:
+        """
+        Returns a dictionary of sequence data with no gaps or missing characters and
+        a dictionary with information to construct a new SequenceCollection via
+        make_unaligned_seqs.
+
+        Parameters
+        ----------
+        name_map
+            A dict of {aln_name: data_name, ...} indicating the mapping between
+            names in the encompassing Alignment (aln_name) and the names in self
+            (data_name).
+        start
+            The starting position for slicing the sequences.
+        stop
+            The stopping position for slicing the sequences.
+        step
+            The step size for slicing the sequences.
+
+        Returns
+        -------
+        tuple
+            A tuple containing the following:
+            - seqs (dict): A dictionary of {name: seq, ...} where the sequences have no gaps
+              or missing characters.
+            - kwargs (dict): A dictionary of keyword arguments for make_unaligned_seqs, e.g.,
+              {"offset": self.offset, "name_map": name_map}.
+        """
+        ...
+
 
 def _gapped_seq_len(seq: numpy.ndarray, gap_map: numpy.ndarray) -> int:
     """calculate the gapped sequence length from a ungapped sequence and gap map
@@ -3116,6 +3151,28 @@ class AlignedSeqsData(AlignedSeqsDataABC):
             seqid=seqid, start=start, stop=stop, step=step
         ).encode("utf8")
 
+    @extend_docstring_from(AlignedSeqsDataABC.get_ungapped)
+    def get_ungapped(
+        self,
+        name_map: dict[str, str],
+        start: OptInt = None,
+        stop: OptInt = None,
+        step: OptInt = None,
+    ) -> tuple[dict, dict]:
+        seqs = {}
+        for name in name_map.values():
+            # use get_gapped_seq_array because we want to slice in alignment coordinates
+            seq = self.get_gapped_seq_array(
+                seqid=name, start=start, stop=stop, step=step
+            )
+            # we need to remove the index of gap and missing characters
+            mask = (seq != self.alphabet.gap_index) & (
+                seq != self.alphabet.missing_index
+            )
+            seqs[name] = seq[mask]
+
+        return seqs, {"offset": self.offset, "name_map": name_map}
+
     def add_seqs(
         self,
         seqs: dict[str, StrORArray],
@@ -4132,6 +4189,37 @@ class Alignment(SequenceCollection):
         names = [name for name, count in gap_counts.items() if count <= cutoff]
         return self.take_seqs(names)
 
+    def degap(self) -> SequenceCollection:
+        """Returns new SequenceCollection in which sequences have no gaps or
+        missing characters.
+
+        Notes
+        -----
+        The returned collection will not retain an annotation_db if present.
+        """
+        # refactor: design
+        # add optional argument to retain annotation_db
+
+        # because SequenceCollection does not track slice operations, we need
+        # to apply any slice record to the underlying data
+        start, stop, step = (
+            self._slice_record.start,
+            self._slice_record.stop,
+            self._slice_record.step,
+        )
+        data, kwargs = self._seqs_data.get_ungapped(
+            name_map=self._name_map, start=start, stop=stop, step=step
+        )
+        # the SeqsData classes will return the data corresponding to the slice,
+        # however, will not complement the data if the step is negative. We do
+        # this here.
+        data = (
+            {name: self.moltype.complement(seq) for name, seq in data.items()}
+            if step < 0
+            else data
+        )
+        return make_unaligned_seqs(data, moltype=self.moltype, info=self.info, **kwargs)
+
     def matching_ref(self, ref_name, gap_fraction, gap_run):
         """Returns new alignment with seqs well aligned with a reference.
 


Test patch:
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index 2efb5f5f3..9b5b7d621 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -1041,24 +1041,67 @@ def test_sequence_collection_get_seq():
         seqs.get_seq("seqx")
 
 
-def test_sequence_collection_degap():
+@pytest.fixture
+def gap_ambig_seqs():
+    return {"s1": "ATGRY?", "s2": "T-AG??"}
+
+
+@pytest.mark.parametrize("rced", [False, True])
+@pytest.mark.parametrize(
+    "mk_cls",
+    [new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs],
+)
+def test_sequence_collection_degap(mk_cls, gap_ambig_seqs, rced):
     """SequenceCollection.degap should strip gaps from each seq"""
-    data = {"s1": "ATGRY?", "s2": "T-AG??"}
-    seqs = new_alignment.make_unaligned_seqs(data, moltype="dna")
-    got = seqs.degap().to_dict()
-    expect = {"s1": "ATGRY", "s2": "TAG"}
+
+    seqs = mk_cls(gap_ambig_seqs, moltype="dna")
+    # Test normal case
+    degapped = seqs.rc().degap() if rced else seqs.degap()
+    got = degapped.to_dict()
+
+    expect = {"s1": "RYCAT", "s2": "CTA"} if rced else {"s1": "ATGRY", "s2": "TAG"}
     assert got == expect
 
+    # Test empty sequences case
+    empty_seqs = mk_cls({"empty1": "", "empty2": ""}, moltype="dna")
+    got_empty = empty_seqs.degap().to_dict()
+    expect_empty = {"empty1": "", "empty2": ""}
+    assert got_empty == expect_empty
+
 
-def test_sequence_collection_degap_info():
+@pytest.mark.parametrize(
+    "mk_cls",
+    [new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs],
+)
+def test_sequence_collection_degap_info(mk_cls, gap_ambig_seqs):
     """.degap should preserve info attributes"""
-    data = {"s1": "ATGRY?", "s2": "T-AG??"}
-    aln = new_alignment.make_unaligned_seqs(data, moltype="dna")
+    aln = mk_cls(gap_ambig_seqs, moltype="dna")
     aln.info.path = "blah"
     got = aln.degap()
     assert got.info.path == "blah"
 
 
+def test_alignemnt_degap_sliced(gap_ambig_seqs):
+    """degap should apply slice_record to alignment"""
+    aln = new_alignment.make_aligned_seqs(gap_ambig_seqs, moltype="dna")
+    sliced = aln[:3]
+    got = sliced.degap()
+    expect = {"s1": "ATG", "s2": "TA"}
+    assert got.to_dict() == expect
+
+    # ATGRY?
+    # T-AG??
+    #  * * *
+    sliced = aln[1::2]
+    got = sliced.degap()
+    expect = {"s1": "TR", "s2": "G"}
+    assert got.to_dict() == expect
+
+    sliced = aln[::-1]
+    got = sliced.degap()
+    expect = {"s1": "YRGTA", "s2": "GAT"}
+
+
 def test_sequence_collection_to_fasta():
     """SequenceCollection should return correct FASTA string"""
     data = {"seq_0": "AAA", "seq_1": "CCC"}