Repository: cogent3/cogent3
Pull request #2119 (2024-11-20 01:03:49)

ENH: plus attrs now are adjusted to be in modulo of step
## Summary by Sourcery

Enhance sequence slicing by adjusting plus attributes to be in modulo of step, ensuring consistent handling of sequence views. Add tests to verify correct behavior when stop values do not align with step.

Enhancements:
- Adjust sequence slicing to ensure plus attributes are in modulo of step, improving consistency in sequence handling.

Tests:
- Add tests for sequence view attributes when stop value does not align with step, ensuring correct plus_start and plus_stop values.


Patch:
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index 93c3076e4..4b9c2a046 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -185,19 +185,13 @@ class SeqDataView(new_sequence.SeqView):
     @property
     def array_value(self) -> numpy.ndarray:
         """returns the sequence as a numpy array"""
-        # todo: kath, in ADV, the .get_seq_str method gets passed the step and
-        # the returned sequence is sliced by the step. In SDV, the step is not
-        # applied in the get_seq_str method, but is applied in this method here.
-        # can we make this consistent?
-
-        # also, keep using parent_start and parent_stop or to use start and stop?
-        # this is another inconsistency between SDV and ADV
         raw = self.parent.get_seq_array(
             seqid=self.seqid,
             start=self.slice_record.plus_start,
             stop=self.slice_record.plus_stop,
+            step=self.slice_record.plus_step,
         )
-        return raw if self.slice_record.step == 1 else raw[:: self.slice_record.step]
+        return raw[::-1] if self.slice_record.is_reversed else raw
 
     @property
     def bytes_value(self) -> bytes:
@@ -390,21 +384,50 @@ class SeqsData(SeqsDataABC):
         return self._data[seqid].shape[0]
 
     def get_seq_array(
-        self, *, seqid: str, start: OptInt = None, stop: OptInt = None
+        self,
+        *,
+        seqid: str,
+        start: OptInt = None,
+        stop: OptInt = None,
+        step: OptInt = None,
     ) -> numpy.ndarray:
-        return self._data[seqid][start:stop]
+        start = start or 0
+        stop = stop if stop is not None else self.get_seq_length(seqid)
+        step = step or 1
+
+        assert (
+            start >= 0 and stop >= 0 and step > 0
+        ), "start, stop, and step must be >= 0"
+
+        out_len = (stop - start + step - 1) // step
+        out = numpy.empty(out_len, dtype=self.alphabet.dtype)
+
+        out[:] = self._data[seqid][start:stop:step]
+        return out
 
     def get_seq_str(
-        self, *, seqid: str, start: OptInt = None, stop: OptInt = None
+        self,
+        *,
+        seqid: str,
+        start: OptInt = None,
+        stop: OptInt = None,
+        step: OptInt = None,
     ) -> str:
         return self._alphabet.from_indices(
-            self.get_seq_array(seqid=seqid, start=start, stop=stop)
+            self.get_seq_array(seqid=seqid, start=start, stop=stop, step=step)
         )
 
     def get_seq_bytes(
-        self, *, seqid: str, start: OptInt = None, stop: OptInt = None
+        self,
+        *,
+        seqid: str,
+        start: OptInt = None,
+        stop: OptInt = None,
+        step: OptInt = None,
     ) -> bytes:
-        return self.get_seq_str(seqid=seqid, start=start, stop=stop).encode("utf8")
+        return self.get_seq_str(seqid=seqid, start=start, stop=stop, step=step).encode(
+            "utf8"
+        )
 
     def get_view(self, seqid: str) -> SeqDataView:
         seq_len = len(self._data[seqid])
@@ -2878,8 +2901,9 @@ class Aligned:
 
 
 class AlignedSeqsDataABC(SeqsDataABC):
-    # all methods that are from SeqsDataABC should work in sequence coordinates
-    # all methods unique to AlignedSeqsDataABC should work in aligned coordinates
+    # all methods that are from SeqsDataABC should work in sequence coordinates,
+    # all methods unique to AlignedSeqsDataABC should work in aligned coordinates.
+    # all indices provided to AlignedSeqsDataABC should be on the plus strand.
     __slots__ = ()
 
     @classmethod
@@ -3302,8 +3326,20 @@ class AlignedSeqsData(AlignedSeqsDataABC):
         """Return ungapped sequence corresponding to seqid as an array of indices.
         assumes start/stop are in sequence coordinates. Excludes gaps.
         """
-        seq = self._get_ungapped(seqid)
-        return seq[start:stop:step]
+        start = start or 0
+        stop = stop if stop is not None else self.get_seq_length(seqid)
+        step = step or 1
+
+        assert (
+            start >= 0 and stop >= 0 and step > 0
+        ), "start, stop, and step must be >= 0"
+
+        out_len = (stop - start + step - 1) // step
+
+        seq = numpy.empty(out_len, dtype=self.alphabet.dtype)
+        seq[:] = self._get_ungapped(seqid)[start:stop:step]
+
+        return seq
 
     def get_gapped_seq_array(
         self,
@@ -3316,15 +3352,19 @@ class AlignedSeqsData(AlignedSeqsDataABC):
         """Return sequence data corresponding to seqid as an array of indices.
         start/stop are in alignment coordinates. Includes gaps.
         """
-        # reminder: design
-        # this is recreating the orignal seq, so will be memory and compute inefficient
-        # better to build the smallest gapped seq represented by the start/stop
-        # determine the seq coordinates for slicing the ungapped sequence then apply the step
-        # note with the latter approach, if the step is negative, need to apply the step as
-        # a positive integer than return the reversed array
+        start = start or 0
+        stop = stop if stop is not None else self.align_len
+        step = step or 1
+        assert (
+            start >= 0 and stop >= 0 and step > 0
+        ), "start, stop, and step must be >= 0"
+
         index = self._name_to_index[seqid]
-        gapped = self._gapped[index]
-        return gapped[start:stop:step]
+        out_len = (stop - start + step - 1) // step
+        gapped = numpy.empty(out_len, dtype=self.alphabet.dtype)
+        gapped[:] = self._gapped[index][start:stop:step]
+
+        return gapped
 
     def get_seq_str(
         self,
@@ -3625,57 +3665,39 @@ class AlignedDataView(new_sequence.SeqViewABC):
 
     @property
     def str_value(self) -> str:
-        return self.parent.get_seq_str(
-            seqid=self.seqid,
-            start=self.slice_record.start,
-            stop=self.slice_record.stop,
-            step=self.slice_record.step,
-        )
+        return self.alphabet.from_indices(self.array_value)
 
     @property
     def gapped_str_value(self) -> str:
-        return self.parent.get_gapped_seq_str(
-            seqid=self.seqid,
-            start=self.slice_record.start,
-            stop=self.slice_record.stop,
-            step=self.slice_record.step,
-        )
+        return self.alphabet.from_indices(self.gapped_array_value)
 
     @property
     def array_value(self) -> numpy.ndarray:
-        return self.parent.get_seq_array(
+        value = self.parent.get_seq_array(
             seqid=self.seqid,
-            start=self.slice_record.start,
-            stop=self.slice_record.stop,
-            step=self.slice_record.step,
+            start=self.map.get_seq_index(self.slice_record.plus_start),
+            stop=self.map.get_seq_index(self.slice_record.plus_stop),
+            step=self.map.get_seq_index(self.slice_record.plus_step),
         )
+        return value[::-1] if self.slice_record.is_reversed else value
 
     @property
     def gapped_array_value(self) -> numpy.ndarray:
-        return self.parent.get_gapped_seq_array(
+        value = self.parent.get_gapped_seq_array(
             seqid=self.seqid,
-            start=self.slice_record.start,
-            stop=self.slice_record.stop,
-            step=self.slice_record.step,
+            start=self.slice_record.plus_start,
+            stop=self.slice_record.plus_stop,
+            step=self.slice_record.plus_step,
         )
+        return value[::-1] if self.slice_record.is_reversed else value
 
     @property
     def bytes_value(self) -> bytes:
-        return self.parent.get_seq_bytes(
-            seqid=self.seqid,
-            start=self.slice_record.start,
-            stop=self.slice_record.stop,
-            step=self.slice_record.step,
-        )
+        return self.str_value.encode("utf8")
 
     @property
     def gapped_bytes_value(self) -> bytes:
-        return self.parent.get_gapped_seq_bytes(
-            seqid=self.seqid,
-            start=self.slice_record.start,
-            stop=self.slice_record.stop,
-            step=self.slice_record.step,
-        )
+        return self.gapped_str_value.encode("utf8")
 
     def __str__(self) -> str:
         return self.str_value
diff --git a/src/cogent3/core/new_sequence.py b/src/cogent3/core/new_sequence.py
index ecd82a331..265743935 100644
--- a/src/cogent3/core/new_sequence.py
+++ b/src/cogent3/core/new_sequence.py
@@ -2113,9 +2113,17 @@ class SliceRecordABC(ABC):
     def plus_start(self) -> int:
         """start on plus strand"""
         if self.is_reversed:
-            # self.stop becomes the start, self.stop will be negative
+            # all indices are negative, so the stop becomes the start. To convert
+            # to positive indices, we add the length of the parent sequence to
+            # the stop. Note that when abs(self.step) > 1, we instead calculate
+            # the "true stop", i.e., the index that immediately follows (in the
+            # direction of the step) the last selected index in the slice. We then
+            # add the abs(self.step) to account for the fact that the stop uses
+            # exclusive indexing, and the start is inclusive.
+
             assert self.stop < 0, "expected stop on reverse strand SeqView < 0"
-            start = self.stop + self.parent_len + 1
+            stop = self.stop if self.step == -1 else self.start + len(self) * self.step
+            start = self.parent_len + stop + abs(self.step)
         else:
             start = self.start
         return start
@@ -2128,7 +2136,13 @@ class SliceRecordABC(ABC):
             assert self.start < 0, "expected start on reverse strand SeqView < 0"
             stop = self.start + self.parent_len + 1
         else:
-            stop = self.stop
+            # we ensure that the plus_stop is the index that immediately follows
+            # the last selected index in the slice.
+            stop = (
+                self.stop
+                if self.step == 1
+                else (self.start + len(self) * self.step - self.step) + 1
+            )
         return stop
 
     @property
@@ -2441,12 +2455,10 @@ class SliceRecordABC(ABC):
     def _get_reverse_slice_from_reverse_seqview_(
         self, slice_start: int, slice_stop: int, step: int, **kwargs
     ):
-        # max start is "true stop" + abs(step), because stop is not inclusive
-        # "true stop" adjust for if abs(stop-start) % step != 0
+        # refactor: simplify
+        # are there other places in slice_record where we can use plus_start/stop/step
         if slice_start >= len(self):
-            start = (
-                self.parent_len + self.start + len(self) * self.step + abs(self.step)
-            )
+            start = self.plus_start
         elif slice_start >= 0:
             start = self.parent_len + (self.start + slice_start * self.step)
         else:


Test patch:
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index 8a6009005..6368ee6a6 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -236,8 +236,8 @@ def test_seqs_data_get_seq_view(str_seqs_dict, dna_alphabet, seqid):
 
 
 @pytest.mark.parametrize("seq", ("seq1", "seq2"))
-@pytest.mark.parametrize("start", (None, -1, 0, 1, 4))
-@pytest.mark.parametrize("stop", (None, -1, 0, 1, 4))
+@pytest.mark.parametrize("start", (None, 0, 1))
+@pytest.mark.parametrize("stop", (None, 3, 4))
 def test_seqs_data_get_seq_str(str_seqs_dict, dna_alphabet, seq, start, stop):
     # slicing should be tested in test_get_seq_array
     expect = str_seqs_dict[seq][start:stop]
@@ -623,6 +623,30 @@ def test_make_seqs_from_sequences(mk_cls):
     assert coll.get_seq("s1") == seq1
 
 
+@pytest.mark.parametrize(
+    "mk_cls",
+    [new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs],
+)
+@pytest.mark.parametrize(
+    "seq_name, parent_name", [("seq_1", "parent_1"), ("seq_2", "parent_2")]
+)
+def test_make_seqs_renamed_seqs(mk_cls, seq_name, parent_name, dna_alphabet):
+    # the parent_name should persist from sequence object to SequenceCollection
+    seq_view_1 = new_sequence.SeqView(
+        parent="AAAA", parent_len=4, seqid="parent_1", alphabet=dna_alphabet
+    )
+    seq_view_2 = new_sequence.SeqView(
+        parent="TTTT", parent_len=4, seqid="parent_2", alphabet=dna_alphabet
+    )
+
+    seq_1 = new_moltype.DNA.make_seq(seq=seq_view_1, name="seq_1")
+    seq_2 = new_moltype.DNA.make_seq(seq=seq_view_2, name="seq_2")
+
+    seqs = mk_cls([seq_1, seq_2], moltype="dna")
+    assert set(seqs.names) == {"seq_1", "seq_2"}
+    assert seqs.get_seq(seq_name)._seq.seqid == parent_name
+
+
 @pytest.mark.parametrize(
     "mk_cls, data_cls",
     [
@@ -631,7 +655,7 @@ def test_make_seqs_from_sequences(mk_cls):
     ],
 )
 @pytest.mark.parametrize("seq", ("a", "b"))
-def test_make_seqs_offset(mk_cls, data_cls, seq):
+def test_make_seqs_offset(mk_cls, data_cls, seq, dna_alphabet):
     """SequenceCollection and Alignment constructor functions should handle
     offset argument"""
     data = {"a": "AGGCCC", "b": "AGAAAA"}
@@ -645,6 +669,24 @@ def test_make_seqs_offset(mk_cls, data_cls, seq):
     with pytest.raises(ValueError):
         _ = mk_cls(data, moltype="dna", offset=offset)
 
+    # if provided with sequence objects with offsets, they should be propogated
+    seq_view_1 = new_sequence.SeqView(
+        parent="AAAA", parent_len=4, alphabet=dna_alphabet, offset=1
+    )
+    seq_view_2 = new_sequence.SeqView(
+        parent="TTTT", parent_len=4, alphabet=dna_alphabet, offset=2
+    )
+    seq_1 = new_moltype.DNA.make_seq(seq=seq_view_1, name="seq_1")
+    seq_2 = new_moltype.DNA.make_seq(seq=seq_view_2, name="seq_2")
+
+    seqs = mk_cls([seq_1, seq_2], moltype="dna")
+
+    got = seqs.get_seq("seq_1")
+    assert got._seq.offset == 1
+
+    got = seqs.get_seq("seq_2")
+    assert got._seq.offset == 2
+
 
 @pytest.mark.parametrize(
     "mk_cls",
@@ -1047,7 +1089,8 @@ def test_sequence_collection_degap(mk_cls, gap_ambig_seqs, rced):
 
     # Test empty sequences case
     empty_seqs = mk_cls({"empty1": "", "empty2": ""}, moltype="dna")
-    got_empty = empty_seqs.degap().to_dict()
+    empty_seqs = empty_seqs.degap()
+    got_empty = empty_seqs.to_dict()
     expect_empty = {"empty1": "", "empty2": ""}
     assert got_empty == expect_empty
 
diff --git a/tests/test_core/test_new_sequence.py b/tests/test_core/test_new_sequence.py
index 15579526d..e583c46b4 100644
--- a/tests/test_core/test_new_sequence.py
+++ b/tests/test_core/test_new_sequence.py
@@ -2670,7 +2670,7 @@ def test_seqview_seq_len_init(start, stop, step, length, dna_alphabet):
 @pytest.mark.parametrize("step", (1, 3))
 @pytest.mark.parametrize("rev", (False, True))
 def test_seqview_plus_attrs(rev, step):
-    # always equal values from plkus strand
+    # always equal values from plus strand
     start = 2
     stop = 12
     sr = new_sequence.SliceRecord(start=start, stop=stop, step=step, parent_len=30)
@@ -2678,6 +2678,65 @@ def test_seqview_plus_attrs(rev, step):
     assert (sr.plus_start, sr.plus_stop, sr.plus_step) == (start, stop, step)
 
 
+@pytest.mark.parametrize("rev", (False, True))
+def test_seqview_attrs_non_modulo(rev):
+    # for when the stop value does not align with the step, the plus_[stop/start]
+    # is adjusted
+
+    # start = 1, stop = 10, step = 2
+    #   [  -->  )
+    # 0123456789
+    #   * * * *
+    # we expect the plus_stop to be 9, as this is the index immediately following
+    # the last selected index
+    sr = new_sequence.SliceRecord(start=2, stop=10, step=2, parent_len=10)
+    sr = sr[::-1] if rev else sr
+    assert sr.plus_start == 2
+    assert sr.plus_stop == 9
+
+    # start = -4, stop = -10, step = -2
+    # ( <-- ]
+    # 0123456789
+    #   * * *
+    # although a stop of -10 would be a start of 1 for step == 1, we expect
+    # in this case the plus_start to be 2
+    sr = new_sequence.SliceRecord(start=-4, stop=-10, step=-2, parent_len=10)
+    sr = sr[::-1] if rev else sr
+    assert sr.plus_start == 2
+    assert sr.plus_stop == 7
+
+    # start = 1, stop = 7, step = 3
+    #  [ --> )
+    # 0123456789
+    #  *  *
+    sr = new_sequence.SliceRecord(start=1, stop=7, step=3, parent_len=10)
+    sr = sr[::-1] if rev else sr
+    assert sr.plus_start == 1
+    assert sr.plus_stop == 5
+
+    # start = -3, stop = -9, step = -3
+    #  (     ]
+    # 0123456789
+    #     *  *
+    sr = new_sequence.SliceRecord(start=-3, stop=-9, step=-3, parent_len=10)
+    sr = sr[::-1] if rev else sr
+    assert sr.plus_start == 4
+    assert sr.plus_stop == 8
+
+    sr = new_sequence.SliceRecord(start=0, stop=0, step=1, parent_len=10)
+    sr = sr[::-1] if rev else sr
+    assert sr.plus_start == 0
+    assert sr.plus_stop == 0
+
+
+@pytest.mark.parametrize("rev", (False, True))
+def test_seqview_attrs_zero_slice(rev):
+    sr = new_sequence.SliceRecord(start=0, stop=0, step=1, parent_len=0)
+    sr = sr[::-1] if rev else sr
+    assert sr.plus_start == 0
+    assert sr.plus_stop == 0
+
+
 def test_seqview_copy_propagates_seq_len(dna_alphabet):
     seq = "ACGGTGGGAC"
     sv = new_sequence.SeqView(parent=seq, parent_len=len(seq), alphabet=dna_alphabet)