Repository: cogent3/cogent3
Pull request #2024 (2024-09-18 00:52:41)

ENH: IndelMap gets methods for identifying and removing common gaps
[NEW] IndelMap.shared_gaps() takes either another IndelMap instance,
     or a numpy array which consists of start, end coordinates that
     define gap locations in alignment coordinates. It returns a
     numpy array with start, end coordinates for shared gaps. This
     method relies on two new functions, coords_intersect() and
     span_and_span() the latter defining the intersection for two
     spans.
     The result of this method can be applied to multiple IndelMap
     instances from an alignment so that all common gaps can be
     identified and then removed for each case.

[NEW] IndelMap.minus_gaps() takes either another IndelMap instance,
     or a numpy array which consists of start, end coordinates that
     define gap locations in alignment coordinates. It returns a
     new IndelMap instance after removing shared gaps. It uses
     a new coords_minus_coords() function.


Patch:
diff --git a/src/cogent3/core/location.py b/src/cogent3/core/location.py
index dbae426d5..70bd1dc31 100644
--- a/src/cogent3/core/location.py
+++ b/src/cogent3/core/location.py
@@ -6,7 +6,6 @@ import json
 from abc import ABC, abstractmethod
 from bisect import bisect_left, bisect_right
 from functools import total_ordering
-from itertools import chain
 from typing import Any, Iterator, Optional, Sequence, Union
 
 import numpy
@@ -15,13 +14,7 @@ from numpy.typing import NDArray
 from cogent3._version import __version__
 from cogent3.util import warning as c3warn
 from cogent3.util.deserialise import register_deserialiser
-from cogent3.util.misc import (
-    ClassChecker,
-    ConstrainedList,
-    FunctionWrapper,
-    get_object_provenance,
-    iterable,
-)
+from cogent3.util.misc import get_object_provenance
 
 strip = str.strip
 
@@ -1458,7 +1451,9 @@ class IndelMap(MapABC):
             result = result.reshape((0, 2))
         return result
 
-    def merge_maps(self, other, parent_length: Optional[int] = None) -> "IndelMap":
+    def merge_maps(
+        self, other: "IndelMap", parent_length: Optional[int] = None
+    ) -> "IndelMap":
         """merge gaps of other with self
 
         Parameters
@@ -1596,6 +1591,202 @@ class IndelMap(MapABC):
 
         return FeatureMap(spans=spans, parent_length=self.parent_length)
 
+    @functools.singledispatchmethod
+    def shared_gaps(self, other: Union["IndelMap", numpy.ndarray]) -> numpy.ndarray:
+        """returns a numpy array of the shared [(gap start, gap end), ...]
+
+        Notes
+        -----
+        The result is in alignment coordinates
+        """
+        if len(self) != len(other):
+            raise AssertionError(
+                f"{len(other)=} != {len(self)=}, from a different alignment?"
+            )
+        if self.num_gaps == 0 or other.num_gaps == 0:
+            return numpy.array([], dtype=_DEFAULT_GAP_DTYPE)
+
+        return self.shared_gaps(other.get_gap_align_coordinates())
+
+    @shared_gaps.register
+    def _(self, other: numpy.ndarray) -> numpy.ndarray:
+        """returns a numpy array of the shared [(gap start, gap end), ...]
+
+        Notes
+        -----
+        The result is in alignment coordinates
+        """
+        other_gaps = other
+        if not len(other_gaps):
+            return numpy.array([], dtype=_DEFAULT_GAP_DTYPE)
+
+        if other_gaps[-1][-1] > len(self):
+            raise AssertionError(
+                f"other_gaps {other_gaps!r} from a different alignment?"
+            )
+        self_gaps = self.get_gap_align_coordinates()
+
+        result = coords_intersect(self_gaps, other_gaps)
+        return numpy.array(result, dtype=_DEFAULT_GAP_DTYPE)
+
+    @functools.singledispatchmethod
+    def minus_gaps(self, other_gaps: Union["IndelMap", numpy.ndarray]):
+        """returns new map with gaps in other_gaps removed from self
+
+        Parameters
+        ----------
+        other_gaps
+            IndelMap instance, or numpy array in alignment coordinates
+
+        Returns
+        -------
+        New instance with gap spans represented by other_gaps removed from self
+        """
+        if len(self) != len(other_gaps):
+            raise AssertionError(
+                f"{len(other_gaps)=} != {len(self)=}, from a different alignment?"
+            )
+
+        return self.minus_gaps(other_gaps.get_gap_align_coordinates())
+
+    @minus_gaps.register
+    def _(self, other_gaps: numpy.ndarray):
+        if not len(other_gaps):
+            return self
+
+        if other_gaps[-1][-1] > len(self):
+            raise AssertionError(
+                f"other_gaps {other_gaps!r} from a different alignment?"
+            )
+        self_gaps = self.get_gap_align_coordinates()
+        unique = coords_minus_coords(self_gaps, other_gaps)
+        new_gaps = numpy.empty((len(unique), 2), dtype=_DEFAULT_GAP_DTYPE)
+        for i, (start, end) in enumerate(unique):
+            new_gaps[i] = self.get_seq_index(start), end - start
+        gap_pos, gap_lengths = new_gaps.T
+        return self.__class__(
+            gap_pos=gap_pos,
+            gap_lengths=gap_lengths,
+            parent_length=self.parent_length,
+        )
+
+
+_empty = None, None
+
+
+def coords_minus_coords(
+    coords1: numpy.ndarray, coords2: numpy.ndarray
+) -> numpy.ndarray:
+    """returns the coords1 minus any overlaps with coords2
+
+    Parameters
+    ----------
+    coords1, coords2
+        coordinates defining segments as [(gap start, gap end)..]
+
+    Notes
+    ------
+    Assumes both coordinate series are sorted.
+    """
+    dtype = getattr(coords1, "dtype", _DEFAULT_GAP_DTYPE)
+    unique_segments = []
+    for a1, a2 in coords1:
+        total_intersect = None
+        for b1, b2 in coords2:
+            if b2 < a1:
+                continue
+
+            if a2 <= b1:
+                # no intersection
+                break
+
+            result = span_and_span((a1, a2), (b1, b2))
+            if result[0] is not None:
+                total_intersect = result[1] - result[0] + (total_intersect or 0)
+
+        end = a2 - (total_intersect or 0)
+        if end < 0:
+            raise ValueError(
+                f"new length negative segment length {a1=} {a2=} {total_intersect=}"
+            )
+        elif a2 - a1 != total_intersect:
+            unique_segments.append((a1, end))
+
+    return numpy.array(unique_segments, dtype=dtype)
+
+
+def span_and_span(
+    spans1: tuple[int, int], spans2: tuple[int, int]
+) -> Union[tuple[int, int], tuple[None, None]]:
+    """returns the intersection of two spans
+
+    Parameters
+    ----------
+    spans1, span2
+        each span is a tuple of start, end coordinates
+
+    Returns
+    -------
+        If no intersection, returns (None, None)
+
+    Raises
+    ------
+    ValueError
+        if a span is not in span[0] < span[1] order
+    """
+    a1, a2 = spans1
+    b1, b2 = spans2
+    # return intersection of the spans
+    if a1 >= a2 or b1 >= b2:
+        raise ValueError("coordinates must be start < end")
+
+    if a1 < b1 and a2 > b2:
+        # span1 contains span2
+        return b1, b2
+    elif a1 >= b1 and a2 <= b2:
+        # span1 equal to or within span2
+        return a1, a2
+    elif a1 == b1:
+        # same start, different end
+        return a1, min(a2, b2)
+    elif a2 == b2:
+        # different start, same end
+        return max(a1, b1), a2
+    elif a1 < b1 < a2:
+        # span1 overlaps start of span2
+        return b1, min(a2, b2)
+    elif a1 < b2 < a2:
+        # span1 overlaps end of span2
+        return max(a1, b1), b2
+    else:
+        return _empty
+
+
+def coords_intersect(
+    coords1: numpy.ndarray, coords2: numpy.ndarray
+) -> list[tuple[int, int]]:
+    """returns the intersecting spans between two sets of coordinates
+
+    Parameters
+    ----------
+    coords1, coords2
+        coordinates defining segments as [(gap start, gap end)..]
+
+    Notes
+    ------
+    Assumes both coordinate series are sorted.
+    """
+    intersects = []
+    for a1, a2 in coords1:
+        for b1, b2 in coords2:
+            if a1 <= b2 and b1 <= a2:
+                i1, i2 = span_and_span((a1, a2), (b1, b2))
+                if i1 is not None:
+                    intersects.append((i1, i2))
+            elif a2 < b1:
+                break
+    return intersects
+
 
 @dataclasses.dataclass
 class FeatureMap(MapABC):
diff --git a/src/cogent3/draw/dotplot.py b/src/cogent3/draw/dotplot.py
index 0ed277422..4495205a3 100644
--- a/src/cogent3/draw/dotplot.py
+++ b/src/cogent3/draw/dotplot.py
@@ -6,6 +6,7 @@ from cogent3.align.pycompare import (
     _calc_seed_size,
     find_matched_paths,
 )
+from cogent3.core import location
 from cogent3.core.moltype import get_moltype
 from cogent3.draw.drawable import Drawable
 from cogent3.util.union_dict import UnionDict
@@ -66,13 +67,17 @@ def _convert_input(seq, moltype):
     return gap_map, seq
 
 
-def get_align_coords(map1, map2, aligned=False) -> MatchedSeqPaths:
+def get_align_coords(
+    map1: location.IndelMap, map2: location.IndelMap, aligned=False
+) -> MatchedSeqPaths:
     """sequence coordinates of aligned segments"""
     from cogent3.align.pycompare import segment
 
     if aligned:
         assert len(map1) == len(map2), "aligned sequences not equal length"
 
+    # we reduce the maps to only gaps that are unique to each
+
     # we get the gap coordinates in alignment indices for both sequences
     # sorting this allows us to trivially identify the alignment indices of
     # ungapped segments,
diff --git a/src/cogent3/util/misc.py b/src/cogent3/util/misc.py
index 22629dfe6..f50b44f72 100644
--- a/src/cogent3/util/misc.py
+++ b/src/cogent3/util/misc.py
@@ -8,7 +8,7 @@ import os
 import re
 import typing
 import warnings
-from random import randint
+from random import choice
 from typing import Any, Callable, Tuple
 from urllib.parse import urlparse
 from warnings import warn
@@ -832,7 +832,7 @@ def get_independent_coords(spans, random_tie_breaker=False):
         curr = spans[i]
         if curr[0] < last[1]:
             if random_tie_breaker:
-                result[-1] = [last, curr][randint(0, 1)]
+                result[-1] = choice((last, curr))
             else:
                 result[-1] = last
             continue


Test patch:
diff --git a/tests/test_core/test_location.py b/tests/test_core/test_location.py
index 5119494cc..c37db2a2b 100644
--- a/tests/test_core/test_location.py
+++ b/tests/test_core/test_location.py
@@ -1,5 +1,6 @@
 """Unit tests for Span classes."""
 
+import re
 from itertools import combinations
 from unittest import TestCase
 
@@ -13,7 +14,10 @@ from cogent3.core.location import (
     LostSpan,
     Span,
     TerminalPadding,
+    coords_intersect,
+    coords_minus_coords,
     gap_coords_to_map,
+    span_and_span,
 )
 
 
@@ -1173,6 +1177,192 @@ def test_featuremap_div():
     assert fm_1.parent_length == 6 / 3
 
 
+# Constants
+EMPTY = None, None
+
+
+@pytest.mark.parametrize(
+    "span1, span2, expected",
+    [
+        ((1, 5), (2, 4), (2, 4)),  # span1 contains span2
+        ((2, 4), (1, 5), (2, 4)),  # span1 within span2
+        ((1, 5), (1, 3), (1, 3)),  # same start, different end
+        ((1, 3), (1, 5), (1, 3)),  # same start, different end
+        ((1, 5), (2, 5), (2, 5)),  # different start, same end
+        ((2, 5), (1, 5), (2, 5)),  # different start, same end
+        ((1, 4), (2, 5), (2, 4)),  # span1 overlaps start of span2
+        ((1, 9), (6, 10), (6, 9)),  # span1 overlaps start of span2
+        ((2, 5), (1, 4), (2, 4)),  # span1 overlaps end of span2
+    ],
+)
+def test_span_and_span(span1, span2, expected):
+    result = span_and_span(span1, span2)
+    assert result == expected
+
+
+@pytest.mark.parametrize(
+    "span1, span2",
+    [
+        ((5, 1), (2, 4)),
+        ((1, 5), (4, 2)),
+    ],
+)
+def test_span_and_span_error_cases(span1, span2):
+    with pytest.raises(ValueError):
+        span_and_span(span1, span2)
+
+
+@pytest.mark.parametrize(
+    "arr1, arr2, expected",
+    [
+        # unique or equal
+        ([[1, 3], [5, 7]], [[2, 4], [6, 8]], [[1, 2], [5, 6]]),
+        ([[1, 2], [3, 4], [5, 6]], [[5, 6], [7, 8]], [[1, 2], [3, 4]]),
+        # unique or overlap
+        ([[1, 4], [7, 9]], [[3, 4]], [[1, 3], [7, 9]]),
+        ([[1, 2], [5, 9]], [[2, 6], [8, 10]], [[1, 2], [5, 7]]),
+    ],
+)
+def test_coords_sub_coords(arr1, arr2, expected):
+    arr1 = numpy.array(arr1)
+    arr2 = numpy.array(arr2)
+    result = coords_minus_coords(arr1, arr2)
+    assert numpy.array_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "arr1, arr2, expected",
+    [
+        ([(1, 5)], [(2, 3)], [(2, 3)]),
+        ([(1, 5)], [(0, 2), (3, 4)], [(1, 2), (3, 4)]),
+        ([(1, 5), (6, 10)], [(2, 3), (7, 8)], [(2, 3), (7, 8)]),
+        ([(1, 5), (6, 10), (12, 14)], [(2, 3), (7, 8)], [(2, 3), (7, 8)]),
+        ([(1, 9), (12, 14)], [(2, 4), (6, 10)], [(2, 4), (6, 9)]),
+        ([(1, 5)], [(5, 6)], []),
+        ([(1, 5)], [(0, 1)], []),
+        ([(1, 5)], [(4, 5)], [(4, 5)]),
+        ([(1, 5)], [(1, 5)], [(1, 5)]),
+        ([(1, 5)], [(0, 6)], [(1, 5)]),
+        ([], [(1, 2)], []),
+        ([(1, 2)], [], []),
+        ([], [], []),
+    ],
+)
+def test_coord_intersect(arr1, arr2, expected):
+    result = coords_intersect(arr1, arr2)
+    assert result == expected
+
+
+RAW_SEQS = (
+    "----GTA-TG",
+    "AC--GTA---",
+    "AC--GTA-TG",
+    "A-C-G-T-A-",
+    "-A-C-G-T-A",
+    "ACGTAACGTA",
+    "AC--AACGTA",
+    "ACGTAA--TA",
+    "GTATG-----",
+    "-----GTATG",
+    "----------",
+)
+
+gap_run = re.compile("-+")
+
+
+@pytest.fixture(scope="session", params=combinations(RAW_SEQS, 2))
+def seq_pairs(request):
+    raw1, raw2 = request.param
+    return make_seq(raw1, moltype="dna"), make_seq(raw2, moltype="dna")
+
+
+@pytest.mark.parametrize("as_array", (True, False))
+def test_indelmap_shared_gaps(seq_pairs, as_array):
+    # we check against a sequence modified such that only the shared
+    # gaps are retained
+    s1, s2 = seq_pairs
+    common = "".join(
+        "-" if f"{a}{b}" == "--" else "X" for a, b in zip(str(s1), str(s2))
+    )
+    expect = [(m.start(), m.end()) for m in gap_run.finditer(common)]
+    ig1, s1 = s1.parse_out_gaps()
+    ig2, s2 = s2.parse_out_gaps()
+    arg = ig2.get_gap_align_coordinates() if as_array else ig2
+    got = ig1.shared_gaps(arg)
+    numpy.testing.assert_allclose(got, expect)
+
+
+@pytest.mark.parametrize("as_array", (True, False))
+def test_indelmap_subtraction(seq_pairs, as_array):
+    s1, s2 = seq_pairs
+    unique_gaps = "".join(a for a, b in zip(str(s1), str(s2)) if f"{a}{b}" != "--")
+    ig1, s1 = s1.parse_out_gaps()
+    ig2, s2 = s2.parse_out_gaps()
+    arg = ig2.get_gap_align_coordinates() if as_array else ig2
+    got1 = ig1.minus_gaps(arg)
+    expect, _ = make_seq(unique_gaps, moltype="dna").parse_out_gaps()
+    numpy.testing.assert_allclose(
+        got1.get_gap_coordinates(), expect.get_gap_coordinates()
+    )
+
+
+def test_indelmap_subtraction_build_aligned(seq_pairs):
+    from cogent3.core.alignment import Aligned
+
+    s1, s2 = seq_pairs
+
+    unique_gaps = "".join(a for a, b in zip(str(s1), str(s2)) if f"{a}{b}" != "--")
+    ig1, s1 = s1.parse_out_gaps()
+    ig2, s2 = s2.parse_out_gaps()
+    got1 = ig1.minus_gaps(ig2.get_gap_align_coordinates())
+    got = Aligned(got1, s1)
+    assert str(got) == unique_gaps
+
+
+def test_indelmap_combined_and_subtract(seq_pairs):
+    s1, s2 = seq_pairs
+    ig1, s1 = s1.parse_out_gaps()
+    ig2, s2 = s2.parse_out_gaps()
+    expect = ig1.minus_gaps(ig2.get_gap_align_coordinates())
+    intersect = ig1.shared_gaps(ig2)
+    got = ig1.minus_gaps(intersect)
+    numpy.testing.assert_allclose(
+        got.get_gap_coordinates(), expect.get_gap_coordinates()
+    )
+
+
+def test_indelmap_shared_gap_error():
+    ig1 = IndelMap(
+        gap_pos=numpy.array([1, 3]), gap_lengths=numpy.array([2, 1]), parent_length=5
+    )
+    ig2 = IndelMap(
+        gap_pos=numpy.array([1, 3]), gap_lengths=numpy.array([2, 1]), parent_length=10
+    )
+    with pytest.raises(AssertionError):
+        ig1.shared_gaps(ig2)
+
+
+def test_indelmap_minus_gap_error():
+    ig1 = IndelMap(
+        gap_pos=numpy.array([1, 3]), gap_lengths=numpy.array([2, 1]), parent_length=5
+    )
+    ig2 = IndelMap(
+        gap_pos=numpy.array([1, 3]), gap_lengths=numpy.array([2, 100]), parent_length=10
+    )
+    with pytest.raises(AssertionError):
+        ig1.minus_gaps(ig2.get_gap_align_coordinates())
+
+
+def test_indelmap_invalid_slice_type():
+    imap = IndelMap(
+        gap_pos=numpy.array([10], dtype=int),
+        gap_lengths=numpy.array([2], dtype=int),
+        parent_length=10,
+    )
+    with pytest.raises(NotImplementedError):
+        imap[()]  # pylint: disable=pointless-statement
+
+
 def test_indelmap_make_seq_feature_map():
     #           1
     # 01234567890