Repository: cogent3/cogent3
Pull request #2128 (2024-11-21 23:59:34)

ENH: implementation of new Alignment.with_masked_annotations()
## Summary by Sourcery

Implement the Alignment.with_masked_annotations() method to enable masking of annotated regions in sequences. Refactor and port related tests to a new test file for better organization and clarity.

New Features:
- Introduce the Alignment.with_masked_annotations() method to allow masking of annotated regions in sequences with a specified character.

Enhancements:
- Refactor tests related to annotated region masking by porting them from test_features.py to test_new_aln_annotation.py and breaking them into multiple tests for clarity.


Patch:
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index 786c55933..4d1df8582 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -2757,7 +2757,7 @@ class Aligned:
         else:
             # gapped_array_value gives the reverse of the plus strand
             # so we need to complement it. We do that here because with a
-            # step != 1 we cannot retain a connection to the underlying
+            # step != 1, we cannot retain a connection to the underlying
             # annotations
             seq = self.moltype.degap(self.data.gapped_array_value)
             seq = self.moltype.complement(seq)
@@ -6139,6 +6139,44 @@ class Alignment(SequenceCollection):
             **kwargs,
         )
 
+
+    def with_masked_annotations(
+        self, biotypes: PySeqStr, mask_char: str = "?", shadow: bool = False
+    ):
+        """returns an alignment with regions replaced by mask_char
+
+        Parameters
+        ----------
+        biotypes
+            annotation type(s)
+        mask_char
+            must be a character valid for the moltype. The default value is
+            the most ambiguous character, eg. '?' for DNA
+        shadow
+            If True, masks everything but the biotypes
+        """
+        # by doing numpy.array(seq), this method applies the slice_record
+        # and modifies the underlying sequence data. We therefore split
+        # gapped sequences into their gaps and seq components and discard
+        # the current slice_record.
+        gaps = {}
+        ungapped = {}
+        for seq in self.seqs:
+            parent_name = self._name_map[seq.name]
+            gaps[parent_name] = seq.map.array
+            ungapped[parent_name] = numpy.array(
+                seq.seq.with_masked_annotations(biotypes, mask_char, shadow)
+            )
+
+        seq_data = self._seqs_data.from_seqs_and_gaps(
+            seqs=ungapped, gaps=gaps, alphabet=self.moltype.most_degen_alphabet()
+        )
+        init = self._get_init_kwargs()
+        init["seqs_data"] = seq_data
+        # we have to drop the slice record since we have changed the data
+        init["slice_record"] = None
+        return self.__class__(**init)
+
     def to_rich_dict(self) -> dict[str, str | dict[str, str]]:
         """returns a json serialisable dict"""
         kwargs = self._get_init_kwargs()
@@ -6171,6 +6209,7 @@ def deserialise_alignment(data: dict[str, str | dict[str, str]]) -> Alignment:
     return Alignment.from_rich_dict(data)
 
 
+
 @singledispatch
 def make_aligned_seqs(
     data: Union[dict[str, StrORBytesORArray], list, AlignedSeqsDataABC],


Test patch:
diff --git a/tests/test_core/test_features.py b/tests/test_core/test_features.py
index 1bad66029..f058cb1d5 100755
--- a/tests/test_core/test_features.py
+++ b/tests/test_core/test_features.py
@@ -287,7 +287,7 @@ def test_terminal_gaps():
     assert aln_exons[0].get_slice().to_dict() == dict(x="AAAAA", y="--T--")
 
 
-def test_annotated_region_masks():
+def test_annotated_region_masks():  # ported to test_new_aln_annotation.py
     """masking a sequence with specific features"""
 
     # Annotated regions can be masked (observed sequence characters
@@ -388,7 +388,7 @@ def test_annotated_region_masks():
     }
 
 
-def test_nested_annotated_region_masks():
+def test_nested_annotated_region_masks():  # ported to test_new_aln_annotation.py
     """masking a sequence with specific features when nested annotations"""
     db = GffAnnotationDb()
     db.add_feature(seqid="x", biotype="gene", name="norwegian", spans=[(0, 4)])
diff --git a/tests/test_core/test_new_aln_annotation.py b/tests/test_core/test_new_aln_annotation.py
index 9cdd81b64..fa17e3b54 100644
--- a/tests/test_core/test_new_aln_annotation.py
+++ b/tests/test_core/test_new_aln_annotation.py
@@ -529,3 +529,167 @@ def test_project_features_no_features_for_specified_seqid():
     projected_features = aln.get_projected_features(seqid="seq2", biotype="exon")
 
     assert len(projected_features) == 0
+
+
+def test_annotated_region_masks():
+    """masking a sequence with specific features"""
+    # ported from test_features.py
+    # refactor: break into multiple tests
+
+    # Annotated regions can be masked (observed sequence characters
+    # replaced by another), either through the sequence on which they
+    # reside or by projection from the alignment. Note that mask_char must
+    # be a valid character for the sequence MolType. Either the features
+    # (multiple can be named), or their shadow, can be masked.
+
+    # We create an alignment with a sequence that has two different annotation types.
+    orig_data = {"x": "C-CCCAAAAAGGGAA", "y": "-T----TTTTG-GTT"}
+    db = GffAnnotationDb()
+    db.add_feature(seqid="x", biotype="exon", name="norwegian", spans=[(0, 4)])
+    db.add_feature(
+        biotype="repeat",
+        name="blue",
+        spans=[(9, 12)],
+        seqid="x",
+    )
+    db.add_feature(seqid="y", biotype="repeat", name="frog", spans=[(5, 7)])
+    aln = new_alignment.make_aligned_seqs(orig_data, moltype="dna")
+    aln.annotation_db = db
+
+    assert aln.to_dict() == {"x": "C-CCCAAAAAGGGAA", "y": "-T----TTTTG-GTT"}
+    x = aln.get_seq("x")
+    y = aln.get_seq("y")
+    exon = list(x.get_features(biotype="exon"))[0]
+    assert str(exon.get_slice()) == "CCCC"
+    repeat_x = list(x.get_features(biotype="repeat"))[0]
+    assert str(repeat_x.get_slice()) == "GGG"
+    repeat_y = list(y.get_features(biotype="repeat"))[0]
+    assert str(repeat_y.get_slice()) == "GG"
+
+    # Each sequence should correctly mask either the single feature,
+    # it's shadow, or the multiple features, or shadow.
+
+    assert (
+        str(aln.get_seq("x").with_masked_annotations("exon", mask_char="?"))
+        == "????AAAAAGGGAA"
+    )
+    assert (
+        str(
+            aln.get_seq("x").with_masked_annotations("exon", mask_char="?", shadow=True)
+        )
+        == "CCCC??????????"
+    )
+    assert (
+        str(aln.get_seq("x").with_masked_annotations(["exon", "repeat"], mask_char="?"))
+        == "????AAAAA???AA"
+    )
+    assert (
+        str(
+            aln.get_seq("x").with_masked_annotations(
+                ["exon", "repeat"], mask_char="?", shadow=True
+            )
+        )
+        == "CCCC?????GGG??"
+    )
+    assert (
+        str(aln.get_seq("y").with_masked_annotations("exon", mask_char="?"))
+        == "TTTTTGGTT"
+    )
+    assert (
+        str(aln.get_seq("y").with_masked_annotations("repeat", mask_char="?"))
+        == "TTTTT??TT"
+    )
+    assert (
+        str(
+            aln.get_seq("y").with_masked_annotations(
+                "repeat", mask_char="?", shadow=True
+            )
+        )
+        == "?????GG??"
+    )
+
+    # The same methods can be applied to annotated Alignment's.
+
+    assert aln.with_masked_annotations("exon", mask_char="?").to_dict() == {
+        "x": "?-???AAAAAGGGAA",
+        "y": "-T----TTTTG-GTT",
+    }
+    assert aln.with_masked_annotations(
+        "exon", mask_char="?", shadow=True
+    ).to_dict() == {"x": "C-CCC??????????", "y": "-?----?????-???"}
+    assert aln.with_masked_annotations("repeat", mask_char="?").to_dict() == {
+        "x": "C-CCCAAAAA???AA",
+        "y": "-T----TTTT?-?TT",
+    }
+    assert aln.with_masked_annotations(
+        "repeat", mask_char="?", shadow=True
+    ).to_dict() == {"x": "?-????????GGG??", "y": "-?----????G-G??"}
+    assert aln.with_masked_annotations(["repeat", "exon"], mask_char="?").to_dict() == {
+        "x": "?-???AAAAA???AA",
+        "y": "-T----TTTT?-?TT",
+    }
+    assert aln.with_masked_annotations(["repeat", "exon"], shadow=True).to_dict() == {
+        "x": "C-CCC?????GGG??",
+        "y": "-?----????G-G??",
+    }
+
+
+def test_masking_strand_agnostic_aln():
+    # ported from test_features.py
+    db = GffAnnotationDb()
+    db.add_feature(
+        seqid="x", biotype="CDS", name="gene", spans=[(2, 6), (10, 15), (25, 35)]
+    )
+    # Annotations should be correctly masked,
+    # whether the sequence has been reverse complemented or not.
+    # We use the plus/minus strand CDS containing sequences created above.
+
+    aln = new_alignment.make_aligned_seqs(
+        {
+            "x": "AAGGGGAAAACCCCCAAAAAAAAAATTTTTTTTTTAAA",
+            "y": "AAGGGGAAAACCCCCGGGGGGGGGGTTTTTTTTTTAAA",
+        },
+        moltype="dna",
+        annotation_db=db,
+    )
+    masked = aln.with_masked_annotations("CDS")
+    assert masked.to_dict() == {
+        "x": "AA????AAAA?????AAAAAAAAAA??????????AAA",
+        "y": str(aln.seqs["y"]),
+    }
+    rc = aln.rc()
+    masked = rc.with_masked_annotations("CDS")
+    assert masked.to_dict() == {
+        "x": "TTT??????????TTTTTTTTTT?????TTTT????TT",
+        "y": str(rc.seqs["y"]),
+    }
+
+
+def test_nested_annotated_region_masks():
+    """masking a sequence with specific features when nested annotations"""
+    # ported from test_features.py
+    db = GffAnnotationDb()
+    db.add_feature(seqid="x", biotype="gene", name="norwegian", spans=[(0, 4)])
+    db.add_feature(seqid="x", biotype="repeat", name="blue", spans=[(1, 3)])
+    db.add_feature(seqid="y", biotype="repeat", name="frog", spans=[(1, 4)])
+    aln = new_alignment.make_aligned_seqs(
+        [["x", "C-GGCAAAAATTTAA"], ["y", "-T----TTTTG-GTT"]],
+        annotation_db=db,
+        moltype="text",
+    )
+    gene = list(aln.get_seq("x").get_features(biotype="gene"))[0]
+    assert str(gene.get_slice()) == "CGGC"
+
+    # evaluate the sequence directly
+    masked = str(aln.get_seq("x").with_masked_annotations("repeat", mask_char="?"))
+    assert masked == "C??CAAAAATTTAA"
+
+    exon = list(aln.get_seq("y").get_features(biotype="repeat", name="frog"))[0]
+    assert str(exon.get_slice()) == "TTT"
+    # evaluate the sequence directly
+    masked = str(aln.get_seq("y").with_masked_annotations("repeat", mask_char="?"))
+    assert masked == "T???TGGTT"
+    masked = aln.with_masked_annotations("gene", mask_char="?")
+    got = masked.to_dict()
+    assert got["x"] == "?-???AAAAATTTAA"
+    assert got["y"] == "-T----TTTTG-GTT"