Merge pull request #3058 from nspope/fix-coalrate-timedim

Correct output dimension for `ts.pair_coalescence_counts` with extra time windows

python/CHANGELOG.rst

@@ -2,6 +2,12 @@
 [0.6.1] - 2024-XX-XX
 --------------------
 
+**Bufixes**
+
+- Fix to ``TreeSequence.pair_coalescence_counts`` output dimension when
+  provided with time windows containing no nodes (:user:`nspope`,
+  :issue:`3046`, :pr:`3058`)
+
 
 --------------------
 [0.6.0] - 2024-10-16

python/tests/test_coalrate.py

@@ -1509,6 +1509,25 @@ def test_output_dim(self):
         )
         assert implm.shape == (2, ts.num_nodes)
 
+    def test_extra_time_windows(self):
+        """
+        test that output dimensions match number of time windows
+        and windows without nodes have zero counts
+        """
+        ts = self.example_ts()
+        ss = [[0, 1, 2], [3, 4, 5]]
+        max_time = ts.nodes_time.max()
+        time_windows = np.linspace(0, max_time * 2, 10)
+        implm = ts.pair_coalescence_counts(
+            sample_sets=ss,
+            windows=None,
+            indexes=None,
+            time_windows=time_windows,
+        )
+        assert implm.shape == (time_windows.size - 1,)
+        max_idx = np.searchsorted(time_windows, max_time, side="right")
+        np.testing.assert_allclose(implm[max_idx:], 0.0)
+
 
 class TestPairCoalescenceQuantiles:
     """
@@ -1747,3 +1766,22 @@ def test_long_sequence(self):
         np.testing.assert_allclose(
             rates.flatten(), np.repeat(ck_rates, windows.size - 1)
         )
+
+    def test_extra_time_windows(self):
+        """
+        test that output dimensions match number of time windows
+        and windows without nodes have NaN rates
+        """
+        ts = self.example_ts()
+        ss = [[0, 1, 2], [3, 4, 5]]
+        max_time = ts.nodes_time.max()
+        time_windows = np.append(np.linspace(0, max_time * 2, 10), np.inf)
+        implm = ts.pair_coalescence_rates(
+            time_windows,
+            sample_sets=ss,
+            windows=None,
+            indexes=None,
+        )
+        assert implm.shape == (time_windows.size - 1,)
+        max_idx = np.searchsorted(time_windows, max_time, side="right")
+        assert np.all(np.isnan(implm[max_idx:]))

python/tskit/trees.py

@@ -9739,32 +9739,40 @@ def pair_coalescence_counts(
                 raise ValueError(
                     "Must specify indexes if there are more than two sample sets"
                 )
+        num_indexes = len(indexes)
 
         drop_left_dimension = False
         if windows is None:
             drop_left_dimension = True
             windows = np.array([0.0, self.sequence_length])
+        num_windows = len(windows) - 1
 
         if isinstance(time_windows, str) and time_windows == "nodes":
-            node_bin_map = np.arange(self.num_nodes, dtype=np.int32)
+            num_time_windows = self.num_nodes
+            node_bin_map = np.arange(num_time_windows, dtype=np.int32)
         else:
             if self.time_units == tskit.TIME_UNITS_UNCALIBRATED:
                 raise ValueError("Time windows require calibrated node times")
+            num_time_windows = len(time_windows) - 1
             node_bin_map = np.digitize(self.nodes_time, time_windows) - 1
-            node_bin_map[node_bin_map == time_windows.size - 1] = tskit.NULL
+            node_bin_map[node_bin_map == num_time_windows] = tskit.NULL
             node_bin_map = node_bin_map.astype(np.int32)
+        num_bins = node_bin_map.max() + 1
 
         sample_set_sizes = np.array([len(s) for s in sample_sets], dtype=np.uint32)
         sample_sets = util.safe_np_int_cast(np.hstack(sample_sets), np.int32)
 
-        coalescing_pairs = self.ll_tree_sequence.pair_coalescence_counts(
-            sample_sets=sample_sets,
-            sample_set_sizes=sample_set_sizes,
-            windows=windows,
-            indexes=indexes,
-            node_bin_map=node_bin_map,
-            span_normalise=span_normalise,
-            pair_normalise=pair_normalise,
+        coalescing_pairs = np.zeros((num_windows, num_indexes, num_time_windows))
+        coalescing_pairs[..., :num_bins] = (
+            self.ll_tree_sequence.pair_coalescence_counts(
+                sample_sets=sample_sets,
+                sample_set_sizes=sample_set_sizes,
+                windows=windows,
+                indexes=indexes,
+                node_bin_map=node_bin_map,
+                span_normalise=span_normalise,
+                pair_normalise=pair_normalise,
+            )
         )
 
         if drop_middle_dimension: