perf: sweep-line Timeline.co_iter and Timeline.segmentation (~26x faster DER on long files)

src/pyannote/core/timeline.py

@@ -86,6 +86,7 @@
 See :class:`pyannote.core.Timeline` for the complete reference.
 """
 
+import math
 import warnings
 from collections.abc import Callable, Iterable, Iterator
 from typing import (
@@ -402,12 +403,45 @@ def co_iter(self, other: "Timeline") -> Iterator[tuple[Segment, Segment]]:
             Yields pairs of intersecting segments in chronological order.
         """
 
+        # Sweep over both (sorted) timelines at once instead of re-scanning
+        # every segment of `other` that starts before `segment` ends, which
+        # is quadratic in the number of segments.
+        #
+        # `active` holds candidate segments of `other`, in sorted order:
+        # every segment of `other` up to `pointer` that has not been proven
+        # unable to intersect any later segment of `self`.
+        others = list(other.segments_list_)
+        num_others = len(others)
+        pointer = 0
+        active: list[Segment] = []
+
         for segment in self.segments_list_:
-            # iterate over segments that starts before 'segment' ends
+            # same upper bound as before: segments of `other` sorted
+            # after `temp` can not intersect `segment`
             temp = Segment(start=segment.end, end=segment.end)
-            for other_segment in other.segments_list_.irange(maximum=temp):
+            while pointer < num_others and others[pointer] <= temp:
+                active.append(others[pointer])
+                pointer += 1
+
+            # `active` is sorted, so stop at the first segment beyond `temp`.
+            # Segments of `other` that start before `segment` but do not
+            # intersect it can not intersect any later segment of `self`
+            # either (later segments do not start earlier), so drop them.
+            kept = []
+            index = 0
+            num_active = len(active)
+            while index < num_active:
+                other_segment = active[index]
+                if not other_segment <= temp:
+                    break
+                index += 1
                 if segment.intersects(other_segment):
+                    kept.append(other_segment)
                     yield segment, other_segment
+                elif other_segment.start >= segment.start:
+                    kept.append(other_segment)
+            if len(kept) < index:
+                active[:index] = kept
 
     def crop_iter(
         self,
@@ -1018,13 +1052,36 @@ def segmentation(self) -> "Timeline":
         if len(timestamps) == 0:
             return Timeline(uri=self.uri)
 
+        # Segment middles are increasing, so instead of calling
+        # `support.overlapping(middle)` (which scans every support segment
+        # starting before `middle`), sweep through support segments once
+        # while tracking the largest end time seen so far: `middle` is
+        # covered iff one of the support segments sorted before
+        # Segment(middle, middle) and starting at or before `middle` ends
+        # at or after `middle`. Both conditions select a prefix of the
+        # (sorted) support segments that only grows as `middle` increases.
+        support_segments = support.segments_list_
+        num_support = len(support_segments)
+        pointer = 0
+        max_end = -math.inf
+
         segments = []
         start = timestamps[0]
         for end in timestamps[1:]:
             # only add segments that are covered by original timeline
             segment = Segment(start=start, end=end)
-            if segment and support.overlapping(segment.middle):
-                segments.append(segment)
+            if segment:
+                middle = segment.middle
+                boundary = Segment(start=middle, end=middle)
+                while (
+                    pointer < num_support
+                    and support_segments[pointer] <= boundary
+                    and support_segments[pointer].start <= middle
+                ):
+                    max_end = max(max_end, support_segments[pointer].end)
+                    pointer += 1
+                if max_end >= middle:
+                    segments.append(segment)
             # next segment...
             start = end
 

tests/test_timeline.py

@@ -23,6 +23,9 @@
 # AUTHORS
 # Hervé BREDIN - http://herve.niderb.fr
 
+import itertools
+import random
+
 import pytest
 
 from pyannote.core import Annotation, Segment, Timeline
@@ -269,3 +272,79 @@ def test_consistent_timelines_with_empty_segments():
     second_timeline.add(Segment(8, 10))
 
     assert first_timeline == second_timeline
+
+
+def _naive_co_iter(timeline, other):
+    # reference implementation: for each segment, scan every segment of
+    # `other` sorted before Segment(segment.end, segment.end)
+    for segment in timeline:
+        temp = Segment(segment.end, segment.end)
+        for other_segment in other:
+            if other_segment <= temp and segment.intersects(other_segment):
+                yield segment, other_segment
+
+
+def _naive_segmentation(timeline):
+    # reference implementation: check each elementary segment's middle
+    # against every segment of the support
+    support = timeline.support()
+    timestamps = sorted({t for segment in timeline for t in segment})
+    segments = [Segment(start, end) for start, end in itertools.pairwise(timestamps)]
+    return Timeline(
+        [
+            segment
+            for segment in segments
+            if segment
+            and any(
+                s <= Segment(segment.middle, segment.middle)
+                and s.overlaps(segment.middle)
+                for s in support
+            )
+        ]
+    )
+
+
+def _random_timeline(rng, num_segments):
+    timeline = Timeline()
+    for _ in range(num_segments):
+        # coarse grid so that ties, nested, touching and duplicate
+        # segments are frequent
+        # jitter below and around SEGMENT_PRECISION (1μs) to exercise
+        # the precision-dependent branches of Segment.intersects
+        start = rng.randint(0, 20) / 2 + rng.choice([0, 0, 0, 5e-7, -5e-7, 2e-6])
+        duration = rng.choice([0.5, 1, 2, 5, 10, 1e-7, 2e-6])
+        timeline.add(Segment(start, start + duration))
+    return timeline
+
+
+def test_co_iter_long_segment_followed_by_short_ones():
+    # a long segment makes the sweep add many candidates at once:
+    # later, shorter segments must not be paired with candidates
+    # starting after they end.
+    timeline = Timeline([Segment(0, 10), Segment(1, 2), Segment(3, 4)])
+    other = Timeline([Segment(0.5, 1.5), Segment(2.5, 3.5), Segment(5, 6)])
+    assert list(timeline.co_iter(other)) == [
+        (Segment(0, 10), Segment(0.5, 1.5)),
+        (Segment(0, 10), Segment(2.5, 3.5)),
+        (Segment(0, 10), Segment(5, 6)),
+        (Segment(1, 2), Segment(0.5, 1.5)),
+        (Segment(3, 4), Segment(2.5, 3.5)),
+    ]
+
+
+def test_co_iter_matches_naive_implementation():
+    rng = random.Random(42)
+    for _ in range(500):
+        first = _random_timeline(rng, rng.randint(0, 30))
+        second = _random_timeline(rng, rng.randint(0, 30))
+        for timeline, other in [(first, second), (second, first), (first, first)]:
+            assert list(timeline.co_iter(other)) == list(
+                _naive_co_iter(timeline, other)
+            )
+
+
+def test_segmentation_matches_naive_implementation():
+    rng = random.Random(42)
+    for _ in range(500):
+        timeline = _random_timeline(rng, rng.randint(0, 30))
+        assert list(timeline.segmentation()) == list(_naive_segmentation(timeline))