Repository: cogent3/cogent3
Pull request #2120 (2024-11-20 01:39:16)

ENH: added apply_scaled_gaps() method
## Summary by Sourcery

Add the apply_scaled_gaps() method to enable gap scaling between codon and amino acid alignments, and enhance the IndelMap class with a scaled() method for scaling indel maps. Include comprehensive tests for these new functionalities.

New Features:
- Introduce the apply_scaled_gaps() method to handle gap scaling between codon and amino acid alignments.

Enhancements:
- Add the scaled() method to the IndelMap class for scaling indel maps by a specified factor, facilitating conversions between codon and amino acid alignments.

Tests:
- Add tests for the apply_scaled_gaps() method to ensure correct gap scaling between codon and amino acid alignments.
- Add tests for the IndelMap.scaled() method to verify correct scaling behavior and error handling for invalid scales.


Patch:
diff --git a/src/cogent3/core/location.py b/src/cogent3/core/location.py
index d8c1b4f23..aefb1d4a0 100644
--- a/src/cogent3/core/location.py
+++ b/src/cogent3/core/location.py
@@ -1198,6 +1198,11 @@ def align_to_seq_index(
         return int(gap_pos[index])
 
 
+def all_gaps_modulo_factor(cum_lengths: numpy.ndarray, factor: int) -> bool:
+    """returns True if all gap lengths are multiples of factor"""
+    return numpy.all(cum_lengths % factor == 0)
+
+
 @dataclasses.dataclass
 class IndelMap(MapABC):
     """store locations of deletions in a Aligned sequence
@@ -1906,6 +1911,46 @@ class IndelMap(MapABC):
             parent_length=self.parent_length,
         )
 
+    def scaled(self, scale_factor: float) -> "IndelMap":
+        """returns indel map scaled by scale_factor
+
+        Notes
+        -----
+        For usage with interconverting between codon and amino-acid alignments.
+        Set scale_factor to 3 for converting amino-acid alignment to codon
+        alignment. Set scale_factor to 1/3 for converting codon alignment to
+        amino-acid alignment.
+        """
+        if scale_factor == 1:
+            return self
+
+        # need to check that, if scale_factor is < 1, it's inverse is an integer
+        if scale_factor < 1:
+            if not numpy.allclose(1 / scale_factor, int(1 / scale_factor)):
+                raise ValueError(
+                    f"scale_factor {scale_factor} must be an integer or 1/integer"
+                )
+
+            if not all_gaps_modulo_factor(self.cum_gap_lengths, int(1 / scale_factor)):
+                raise ValueError(
+                    f"gap lengths must be multiples of {int(1 / scale_factor)}"
+                )
+
+        # make sure gap lengths are modulo factor
+        parent_len = int(self.parent_length * scale_factor)
+        gap_pos = (self.gap_pos * scale_factor).astype(self.gap_pos.dtype)
+        cum_gap_lengths = (self.cum_gap_lengths * scale_factor).astype(
+            self.cum_gap_lengths.dtype
+        )
+        return self.__class__(
+            gap_pos=gap_pos, cum_gap_lengths=cum_gap_lengths, parent_length=parent_len
+        )
+
+    @property
+    def array(self):
+        """returns 2D numpy array with columns gap position and cum gap lengths"""
+        return numpy.array([self.gap_pos, self.cum_gap_lengths]).T
+
 
 _empty = None, None
 
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index 93c3076e4..deff17534 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -5981,6 +5981,50 @@ class Alignment(SequenceCollection):
             info=self.info,
         )
 
+    def apply_scaled_gaps(
+        self, other: SequenceCollection, aa_to_codon: bool | None = None
+    ):
+        """applies gaps in self to unagpped sequences"""
+        assert set(other.names) == set(self.names), "Names must match"
+        if aa_to_codon and not all(
+            (not self.moltype.is_nucleic, other.moltype.is_nucleic)
+        ):
+            raise ValueError(
+                "aa_to_codon True requires nucleic moltype destination not {self.moltype.name!r}"
+            )
+        elif not aa_to_codon and not all(
+            (self.moltype.is_nucleic, not other.moltype.is_nucleic)
+        ):
+            raise ValueError(
+                f"aa_to_codon False requires protein moltype destination not {other.moltype.name!r}"
+            )
+
+        assert aa_to_codon is not None
+        gaps = {}
+        seqs = {}
+        scale = 3 if aa_to_codon else 1 / 3
+
+        for seq in self.seqs:
+            parent_name = self._name_map[seq.name]
+            gaps[parent_name] = seq.map.scaled(scale).array
+            ungapped = numpy.array(other.seqs[seq.name])
+            seqs[parent_name] = ungapped
+            seq_len = self._seqs_data.get_seq_length(parent_name)
+            if len(ungapped) != int(seq_len * scale):
+                raise ValueError(
+                    f"Destination sequence for {seq.name!r} != {scale:.2f} x {seq_len} sequence"
+                )
+
+        seq_data = self._seqs_data.from_seqs_and_gaps(
+            seqs=seqs, gaps=gaps, alphabet=other.moltype.most_degen_alphabet()
+        )
+        init_args = self._get_init_kwargs()
+        init_args.pop("slice_record")  # slice is realised
+        init_args["moltype"] = other.moltype
+        init_args["seqs_data"] = seq_data
+        init_args["annotation_db"] = other.annotation_db
+        return self.__class__(**init_args)
+
 
 @singledispatch
 def make_aligned_seqs(


Test patch:
diff --git a/tests/test_core/test_location.py b/tests/test_core/test_location.py
index 7a8e19f7e..9a87fdb07 100644
--- a/tests/test_core/test_location.py
+++ b/tests/test_core/test_location.py
@@ -1434,3 +1434,43 @@ def test_indelmap_negative_step_variant_slices(start, stop, step, data):
     expect, _ = new_moltype.DNA.make_seq(seq=data[start:stop:step]).parse_out_gaps()
     assert (got.gap_pos == expect.gap_pos).all()
     assert (got.cum_gap_lengths == expect.cum_gap_lengths).all()
+
+
+@pytest.mark.parametrize("scale", [1 / 2, 1 / 1.1])
+def test_indelmap_scaled_invalid(scale):
+    im = IndelMap(
+        gap_pos=numpy.array([1, 3]), gap_lengths=numpy.array([2, 1]), parent_length=5
+    )
+    with pytest.raises(ValueError):
+        im.scaled(scale)
+
+
+@pytest.mark.parametrize("scale", [1, 1.0])
+def test_indelmap_scaled_one(scale):
+    im = IndelMap(
+        gap_pos=numpy.array([1, 3]), gap_lengths=numpy.array([2, 1]), parent_length=5
+    )
+    got = im.scaled(scale)
+    assert got is im
+
+
+@pytest.fixture
+def codon_and_aa_maps():
+    import cogent3
+
+    data = dict(s1="ATG --- --- GAT --- AAA", s2="ATG CAA TCG AAT GAA ATA")
+    dna = cogent3.make_aligned_seqs(
+        {n: s.replace(" ", "") for n, s in data.items()}, moltype="dna", new_type=True
+    )
+    aa = dna.get_translation()
+    dna_1 = dna.seqs["s1"].map
+    aa_1 = aa.seqs["s1"].map
+    return dna_1, aa_1
+
+
+def test_indelmap_scaled(codon_and_aa_maps):
+    cmap, amap = codon_and_aa_maps
+    cmap_to_amap = cmap.scaled(1 / 3)
+    amap_to_cmap = amap.scaled(3)
+    numpy.testing.assert_allclose(cmap_to_amap.array, amap.array)
+    numpy.testing.assert_allclose(amap_to_cmap.array, cmap.array)
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index 8a6009005..97a7ae495 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -4892,3 +4892,53 @@ def test_make_gap_filter():
     assert f3(s1) == False
     assert f1(s3) == False
     assert f3(s4) == True
+
+
+@pytest.fixture(scope="session")
+def codon_and_aa_alns():
+    import cogent3
+
+    data = dict(s1="ATG --- --- GAT --- AAA", s2="ATG CAA TCG AAT GAA ATA")
+    dna = cogent3.make_aligned_seqs(
+        {n: s.replace(" ", "") for n, s in data.items()}, moltype="dna", new_type=True
+    )
+    aa = dna.get_translation()
+    return dna, aa
+
+
+def test_alignment_apply_scaled_gaps_aa_to_codon(codon_and_aa_alns):
+    codon, aa = codon_and_aa_alns
+    ungapped = codon.degap()
+    scaled = aa.apply_scaled_gaps(ungapped, aa_to_codon=True)
+    assert scaled.to_dict() == codon.to_dict()
+
+
+def test_alignment_apply_scaled_gaps_codon_to_aa(codon_and_aa_alns):
+    codon, aa = codon_and_aa_alns
+    ungapped = aa.degap()
+    scaled = codon.apply_scaled_gaps(ungapped, aa_to_codon=False)
+    assert scaled.to_dict() == aa.to_dict()
+
+
+def test_alignment_apply_scaled_gaps_invalid_seqlen(codon_and_aa_alns):
+    codon, aa = codon_and_aa_alns
+    codon = codon[:-2]
+    ungapped = codon.degap()
+    with pytest.raises(ValueError):
+        aa.apply_scaled_gaps(ungapped, aa_to_codon=True)
+
+
+def test_alignment_apply_scaled_gaps_aa2codon_invalid_moltype(codon_and_aa_alns):
+    codon, aa = codon_and_aa_alns
+    codon = codon.to_moltype("text")
+    ungapped = codon.degap()
+    with pytest.raises(ValueError):
+        aa.apply_scaled_gaps(ungapped, aa_to_codon=True)
+
+
+def test_alignment_apply_scaled_gaps_codon2aa_invalid_moltype(codon_and_aa_alns):
+    codon, aa = codon_and_aa_alns
+    ungapped = aa.degap()
+    codon = codon.to_moltype("text")
+    with pytest.raises(ValueError):
+        codon.apply_scaled_gaps(ungapped, aa_to_codon=False)