Repository: cogent3/cogent3
Pull request #2041 (2024-10-11 05:50:57)

ENH: new class method for constructing ASD and now supporting slicing Aligned/Alignments with FeatureMaps  


Patch:
diff --git a/src/cogent3/core/alignment.py b/src/cogent3/core/alignment.py
index 6ce4dc63d..a32646e8a 100644
--- a/src/cogent3/core/alignment.py
+++ b/src/cogent3/core/alignment.py
@@ -2459,7 +2459,9 @@ class Aligned:
     def remapped_to(self, map):
         return Aligned(map[self.map.inverse()].inverse(), self.data)
 
-    def make_feature(self, feature: FeatureDataType, alignment: "Alignment") -> Feature:
+    def make_feature(
+        self, feature: FeatureDataType, alignment: "Alignment"
+    ) -> Feature:  # ported
         """returns a feature, not written into annotation_db"""
         annot = self.data.make_feature(feature)
         inverted = self.map.to_feature_map().inverse()
diff --git a/src/cogent3/core/annotation.py b/src/cogent3/core/annotation.py
index af8f5aa4c..083a31aae 100644
--- a/src/cogent3/core/annotation.py
+++ b/src/cogent3/core/annotation.py
@@ -159,7 +159,7 @@ class Feature:
             # due to separation of IndelMap and Map, change class
             gmap = gmap.to_feature_map()
 
-        seqid = grandparent.name or f"from {self.seqid!r}"
+        seqid = getattr(grandparent, "name", None) or f"from {self.seqid!r}"
         kwargs = {
             **self._serialisable,
             **{"map": gmap[self.map], "parent": grandparent, "seqid": seqid},
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index 5516863c0..4cb8a8f34 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -23,7 +23,12 @@ from cogent3.core.annotation_db import (
     SupportsFeatures,
 )
 from cogent3.core.info import Info as InfoClass
-from cogent3.core.location import IndelMap, _input_vals_neg_step, _input_vals_pos_step
+from cogent3.core.location import (
+    FeatureMap,
+    IndelMap,
+    _input_vals_neg_step,
+    _input_vals_pos_step,
+)
 from cogent3.core.profile import PSSM, MotifCountsArray, MotifFreqsArray, load_pssm
 from cogent3.format.alignment import save_to_filename
 from cogent3.format.fasta import seqs_to_fasta
@@ -180,6 +185,11 @@ class SeqsDataABC(ABC):
 
     __slots__ = ()
 
+    # refactor: design
+    # SeqsData needs a class method for creating instances from a dict, which will need to be
+    # on AlignedSeqsData as well. If we rename AlignedSeqsData.from_aligned_seqs to from_seqs
+    # (or something similar), we can use the same method name for both classes.
+
     @abstractmethod
     def seq_lengths(self) -> dict[str, int]: ...
 
@@ -2494,6 +2504,13 @@ class Aligned:
         """Returns gap_vector of GappedSeq, for omit_gap_pos."""
         return self.gapped_seq.gap_vector()
 
+    def make_feature(self, feature: FeatureDataType, alignment: "Alignment") -> Feature:
+        """returns a feature, not written into annotation_db"""
+        annot = self.seq.make_feature(feature)
+        inverted = self.map.to_feature_map().inverse()
+        # todo should indicate whether tidy or not
+        return annot.remapped_to(alignment, inverted)
+
     def __str__(self) -> str:
         return str(self.gapped_seq)
 
@@ -2519,6 +2536,40 @@ class Aligned:
     def _(self, span: slice):
         return self.__class__(data=self.data[span], moltype=self.moltype)
 
+    @__getitem__.register
+    def _(self, span: FeatureMap):
+        # we assume the feature map is in align coordinates
+        data, gaps = self._slice_with_map(span)
+        seqid = self.data.seqid
+        seqs_data = self.data.parent.from_seqs_and_gaps(
+            seqs={seqid: data},
+            gaps={seqid: gaps},
+            alphabet=self.moltype.most_degen_alphabet(),
+        )
+        view = seqs_data.get_view(seqid)
+
+        return Aligned(view, self.moltype)
+
+    def _slice_with_map(self, span: FeatureMap) -> tuple[numpy.ndarray, numpy.ndarray]:
+        start, end = span.start, span.end
+        if span.useful and len(list(span.spans)) == 1:
+            im = self.map[start:end]
+            seq_start = self.map.get_seq_index(start)
+            seq_end = self.map.get_seq_index(end)
+            data = self.data.array_value[seq_start:seq_end]
+        elif not span.useful:
+            im = self.map[start:end]
+            data = self.data.array_value[:0]
+        else:
+            # multiple spans
+            align_coords = span.get_coordinates()
+            im = self.map.joined_segments(align_coords)
+            seq_map = self.map.make_seq_feature_map(span)
+            data = numpy.array(self.seq.gapped_by_map(seq_map))
+
+        gaps = numpy.array([im.gap_pos, im.cum_gap_lengths]).T
+        return data, gaps
+
     def parent_coordinates(self, seq_coords=False):
         """returns seqid, start, stop, strand on the parent sequence
 
@@ -2564,6 +2615,16 @@ class AlignedSeqsDataABC(SeqsDataABC):
         alphabet: new_alphabet.AlphabetABC,
     ): ...
 
+    @classmethod
+    @abstractmethod
+    def from_seqs_and_gaps(
+        cls,
+        *,
+        seqs: dict[str, StrORBytesORArray],
+        gaps: dict[str, numpy.ndarray],
+        alphabet: new_alphabet.AlphabetABC,
+    ): ...
+
     @property
     @abstractmethod
     def align_len(self) -> int: ...
@@ -2705,6 +2766,34 @@ class AlignedSeqsData(AlignedSeqsDataABC):
             **kwargs,
         )
 
+    @classmethod
+    def from_seqs_and_gaps(
+        cls,
+        *,
+        seqs: dict[str, StrORBytesORArray],
+        gaps: dict[str, numpy.ndarray],
+        alphabet: new_alphabet.AlphabetABC,
+        **kwargs,
+    ):
+        """Construct an AlignedSeqsData object from a dict of ungapped sequences
+        and a corresponding dict of gap data.
+
+        Parameters
+        ----------
+        seqs
+            dict of ungapped sequences {name: seq, ...}
+        gaps
+            gap data {name: [[seq gap position, cumulative gap length], ...], ...}
+        alphabet
+            alphabet object for the sequences
+        """
+        return cls(
+            seqs=seqs,
+            gaps=gaps,
+            alphabet=alphabet,
+            **kwargs,
+        )
+
     @property
     def names(self) -> tuple[str]:
         return list(self._seqs.keys())
@@ -3317,6 +3406,10 @@ class Alignment(SequenceCollection):
             info=self.info,
         )
 
+    @__getitem__.register
+    def _(self, index: FeatureMap):
+        return self._mapped(index)
+
     def __repr__(self):
         seqs = []
         limit = 10
@@ -3906,7 +3999,7 @@ class Alignment(SequenceCollection):
             seq = self.seqs[seqname]
             # we use parent seqid
             parent_id, start, stop, _ = seq.parent_coordinates()
-            offset = seq.annotation_offset
+            offset = seq.data.offset
 
             for feature in self.annotation_db.get_features_matching(
                 seqid=parent_id,
@@ -4238,6 +4331,20 @@ class Alignment(SequenceCollection):
             wrap=settings["wrap"],
         )
 
+    def _mapped(self, slicemap):
+        seqs = {}
+        maps = {}
+        for aligned in self.seqs:
+            seq, map_data = aligned._slice_with_map(slicemap)
+            seqs[aligned.name] = seq
+            maps[aligned.name] = map_data
+
+        data = self._seqs_data.from_seqs_and_gaps(
+            seqs=seqs, gaps=maps, alphabet=self.moltype.most_degen_alphabet()
+        )
+
+        return self.__class__(seqs_data=data, moltype=self.moltype, info=self.info)
+
 
 @singledispatch
 def make_aligned_seqs(


Test patch:
diff --git a/tests/test_core/test_alignment.py b/tests/test_core/test_alignment.py
index 007eacf14..99f17a095 100644
--- a/tests/test_core/test_alignment.py
+++ b/tests/test_core/test_alignment.py
@@ -2000,7 +2000,7 @@ class IntegrationTests(TestCase):
     "raw_seq,coords",
     (("ACGGTAAAG", ((2, 4), (5, 8))), ("CCC---CCC", ((0, 3), (6, 9)))),
 )
-def test_featuremap_slice_aligned(raw_seq, coords):
+def test_featuremap_slice_aligned(raw_seq, coords):  # ported
     from cogent3.core.location import FeatureMap, Span
 
     im, seq = DNA.make_seq(seq=raw_seq).parse_out_gaps()
@@ -3761,7 +3761,7 @@ def test_slice_aligned(raw):  # ported
     assert str(sliced) == raw[:-3]
 
 
-def test_slice_aligned_featuremap_allgap():
+def test_slice_aligned_featuremap_allgap():  # ported
     from cogent3.core.location import FeatureMap, LostSpan
 
     imap, seq = DNA.make_seq(seq="AAAGGGGGAACCCT", name="x").parse_out_gaps()
@@ -3771,7 +3771,7 @@ def test_slice_aligned_featuremap_allgap():
     assert not sliced
 
 
-def test_slice_aligned_featuremap_multi_spans():
+def test_slice_aligned_featuremap_multi_spans():  # ported
     from cogent3.core.location import FeatureMap
 
     #                    1111111
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index da7c05108..0c397ee3c 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -12,6 +12,7 @@ from cogent3._version import __version__
 from cogent3.core import new_alignment, new_alphabet, new_moltype, new_sequence
 from cogent3.core.annotation import Feature
 from cogent3.core.annotation_db import GffAnnotationDb, load_annotations
+from cogent3.core.location import FeatureMap, LostSpan, Span
 from cogent3.util.deserialise import deserialise_object
 from cogent3.util.misc import get_object_provenance
 
@@ -2454,6 +2455,58 @@ def test_aligned_getitem_int(gapped_seqs_dict, seqid, i):
     assert str(got) == expect
 
 
+@pytest.mark.parametrize(
+    "raw_seq,coords",
+    (("ACGGTAAAG", ((2, 4), (5, 8))), ("CCC---CCC", ((0, 3), (6, 9)))),
+)
+def test_aligned_getitem_featuremap(raw_seq, coords):
+    dna = new_moltype.get_moltype("dna")
+    im, seq = dna.make_seq(seq=raw_seq).parse_out_gaps()
+    gaps = numpy.array([im.gap_pos, im.cum_gap_lengths]).T
+    asd = new_alignment.AlignedSeqsData.from_seqs_and_gaps(
+        seqs={"seq1": numpy.array(seq)},
+        gaps={"seq1": gaps},
+        alphabet=dna.most_degen_alphabet(),
+    )
+    aln = new_alignment.make_aligned_seqs(asd, moltype=dna)
+    ia = aln.seqs["seq1"]
+    length = len(raw_seq)
+    fmap = FeatureMap(spans=[Span(s, e) for s, e in coords], parent_length=length)
+    expect = "".join(raw_seq[s:e] for s, e in fmap.get_coordinates())
+    got = ia[fmap]
+    assert str(got) == expect
+
+
+@pytest.fixture
+def aligned():
+    data = {
+        "seq1": "AAAGG--GGG-AACCCT",
+    }
+    aln = new_alignment.make_aligned_seqs(data, moltype="dna")
+    return aln.seqs["seq1"]
+
+
+def test_aligned_getitem_featuremap_allgap(aligned):
+    fmap = FeatureMap(spans=[LostSpan(4)], parent_length=0)
+    sliced = aligned[fmap]
+    assert not sliced
+
+
+def test_aligned_getitem_featuremap_multi_spans(aligned):
+    #                      1111111
+    #            01234567890123456
+    #             ***   **    ***
+    # raw_seq = "AAAGG--GGG-AACCCT"
+    #            01234  567 890123
+    #                         1111
+
+    fmap = FeatureMap.from_locations(
+        locations=[(1, 4), (7, 9), (13, 16)], parent_length=17
+    )
+    sliced = aligned[fmap]
+    assert str(sliced) == "AAGGGCCC"
+
+
 def test_aligned_getitem_raises(gapped_seqs_dict):
     aln = new_alignment.make_aligned_seqs(gapped_seqs_dict, moltype="dna")
     al = aln.seqs["seq1"]
@@ -2537,6 +2590,37 @@ def test_aligned_seqs_data_unequal_seqlens_raises(data_type, dna_alphabet, dna_m
         )
 
 
+def test_from_seqs_and_maps(dna_alphabet):
+    # AlignedSeqsData should be able to be constructed from sequences and gap maps
+    seqs = {"seq1": "ACCTA", "seq2": ""}
+    gaps = {"seq1": numpy.array([[0, 1]]), "seq2": numpy.array([[0, 6]])}
+    asd = new_alignment.AlignedSeqsData.from_seqs_and_gaps(
+        seqs=seqs, gaps=gaps, alphabet=dna_alphabet
+    )
+    assert asd.get_seq_str(seqid="seq1") == "ACCTA"
+    assert asd.get_gapped_seq_str(seqid="seq1") == "-ACCTA"
+    assert asd.get_seq_str(seqid="seq2") == ""
+    assert asd.get_gapped_seq_str(seqid="seq2") == "------"
+
+
+def test_from_seqs_and_maps_diff_seq_lens_raises(dna_alphabet):
+    seqs = {"seq1": "ACCTA", "seq2": "A"}
+    gaps = {"seq1": numpy.array([[0, 1]]), "seq2": numpy.array([[0, 1]])}
+    with pytest.raises(ValueError):
+        _ = new_alignment.AlignedSeqsData.from_seqs_and_gaps(
+            seqs=seqs, gaps=gaps, alphabet=dna_alphabet
+        )
+
+
+def test_from_seqs_and_maps_diff_keys_raises(dna_alphabet):
+    seqs = {"seq1": "ACCTA", "seq2": "A"}
+    gaps = {"seq1": numpy.array([[0, 1]]), "seq3": numpy.array([[0, 1]])}
+    with pytest.raises(ValueError):
+        _ = new_alignment.AlignedSeqsData.from_seqs_and_gaps(
+            seqs=seqs, gaps=gaps, alphabet=dna_alphabet
+        )
+
+
 def test_aligned_seqs_data_diff_keys_raises(dna_alphabet):
     """AlignedSeqsData expect identical keys in seqs and gaps"""
     seqs = dict(seq1=numpy.array([2, 1]), seq2=numpy.array([2, 0, 3, 1]))
@@ -3498,7 +3582,6 @@ def test_sliding_windows():
     assert result[3].to_dict() == {"seq3": "TACGT", "seq2": "TACGT", "seq1": "TACGT"}
 
 
-@pytest.mark.xfail(reason="todo: get_features not properly implemented so test hangs")
 def test_get_feature():
     aln = new_alignment.make_aligned_seqs(
         {"x": "-AAAAAAAAA", "y": "TTTT--CCCT"}, moltype="dna"