Added methods to sort and squash segments in the IbdFinder output.

python/tests/ibd.py

@@ -1,6 +1,6 @@
 # MIT License
 #
-# Copyright (c) 2020-2021 Tskit Developers
+# Copyright (c) 2020-2022 Tskit Developers
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -138,6 +138,37 @@ def add_segment(self, a, b, seg):
         key = (a, b) if a < b else (b, a)
         self.segments[key].append(tskit.IdentitySegment(seg.left, seg.right, seg.node))
 
+    def sort(self):
+        """
+        In-place sorting of each segment list.
+        """
+        for pair in self.segments:
+            num_segs = len(self.segments[pair])
+            if num_segs > 1:
+                self.segments[pair] = sorted(self.segments[pair])
+
+    def squash(self):
+        """
+        Squashes each segment list.
+        """
+        self.sort()
+        for pair in self.segments:
+            num_segs = len(self.segments[pair])
+            if num_segs < 2:
+                continue
+            new_list = []
+            prev_n = self.segments[pair][0].node
+            prev_l = self.segments[pair][0].left
+            prev_r = self.segments[pair][0].right
+            for seg in self.segments[pair]:
+                if seg.node > prev_n and seg.left > prev_r:
+                    new_list.append(Segment(prev_l, prev_r, prev_n))
+                    prev_l = seg.left
+                    prev_n = seg.node
+                prev_r = seg.right
+            new_list.append(Segment(prev_l, prev_r, prev_n))
+            self.segments[pair] = new_list
+
 
 class IbdFinder:
     """
@@ -177,7 +208,7 @@ def print_state(self):
         for u, a in enumerate(self.A):
             print(u, self.sample_set_id[u], a, sep="\t")
 
-    def run(self):
+    def run(self, squash=False):
         node_times = self.tables.nodes.time
         for e in self.ts.edges():
             time = node_times[e.parent]
@@ -197,6 +228,8 @@ def run(self):
                 s = s.next
             self.record_ibd(e.parent, child_segs)
             self.A[e.parent].extend(child_segs)
+        if squash:
+            self.result.squash()
         return self.result.segments
 
     def record_ibd(self, current_parent, child_segs):

python/tests/test_ibd.py

@@ -38,6 +38,7 @@ def ibd_segments(
     compare_lib=True,
     print_c=False,
     print_py=False,
+    squash=False
 ):
     """
     Calculates IBD segments using Python and converts output to lists of segments.
@@ -46,7 +47,7 @@ def ibd_segments(
     ibd_f = ibd.IbdFinder(
         ts, within=within, between=between, max_time=max_time, min_span=min_span
     )
-    ibd_segs = ibd_f.run()
+    ibd_segs = ibd_f.run(squash=squash)
     if print_py:
         print("Python output:\n")
         print(ibd_segs)
@@ -581,6 +582,11 @@ def test_length(self):
         true_segs = {(0, 1): [tskit.IdentitySegment(0.2, 0.7, 4)]}
         assert_ibd_equal(ibd_segs, true_segs)
 
+    def test_squash(self):
+        ibd_segments(
+            self.ts(), within=[0, 1], squash=True, compare_lib=False, print_py=True
+        )
+
 
 class TestIbdDifferentPaths2:
     #
@@ -632,6 +638,11 @@ def test_defaults(self):
         }
         assert_ibd_equal(ibd_segs, true_segs)
 
+    def test_squash(self):
+        ibd_segments(
+            self.ts(), within=[1, 2], squash=True, compare_lib=False, print_py=True
+        )
+
 
 class TestIbdDifferentPaths3:
     # 2.00┊   4   ┊   4   ┊
@@ -663,6 +674,16 @@ def test_defaults(self):
         }
         assert_ibd_equal(ibd_segs, true_segs)
 
+    def test_squash(self):
+        ibd_segments(
+            self.ts(),
+            within=[0, 1],
+            squash=True,
+            compare_lib=False,
+            print_py=True,
+            print_c=True,
+        )
+
 
 class TestIbdPolytomies:
     #