From 68f6dea65b2dc4effe697550f8639808525fd144 Mon Sep 17 00:00:00 2001 From: Zijian Zeng <39647839+ArtificialZeng@users.noreply.github.com> Date: Tue, 27 Oct 2026 14:09:30 +0800 Subject: [PATCH 1/2] Evaluate coalescence rates at the terminal augmented time --- msprime/demography.py | 21 +++++------ tests/test_demography.py | 77 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 86 insertions(+), 12 deletions(-) diff --git a/msprime/demography.py b/msprime/demography.py index 6e4f016aa..9c518e09b 100644 --- a/msprime/demography.py +++ b/msprime/demography.py @@ -4357,16 +4357,17 @@ def _calculate_coalescence_rate_trajectory(self, steps, num_samples, min_pop_siz # recall that steps_b[0] = 0.0 r = np.zeros(num_steps) p_t = np.zeros(num_steps) - for j in range(num_steps - 1): - time = steps_b[j] - dt = steps_b[j + 1] - steps_b[j] + for j, time in enumerate(steps_b): N, M = self._pop_size_and_migration_at_t(time) C = np.zeros([num_pops**2, num_pops**2]) for idx in range(num_pops): C[IA[idx, idx], IA[idx, idx]] = 1 / (2 * max(min_pop_size, N[idx])) dM = np.diag([sum(s) for s in M]) + # Preserve the existing event handling when the augmented grid + # ends with two copies of time zero. + final_duplicate = j == num_steps - 1 and j > 0 and time == steps_b[j - 1] for mmt, mmo in zip(mass_migration_times, mass_migration_objects): - if mmt == time: + if mmt == time and not final_duplicate: a = mmo.source b = mmo.dest p = mmo.proportion @@ -4386,14 +4387,10 @@ def _calculate_coalescence_rate_trajectory(self, steps, num_samples, min_pop_siz r[j] = np.sum(np.matmul(P, C)) / np.sum(P) else: r[j] = np.nan - G = (np.kron(M - dM, Identity) + np.kron(Identity, M - dM)) - C - P = np.matmul(P, _matrix_exponential(dt * G)) - p_notcoal = np.sum(P) - p_t[num_steps - 1] = p_notcoal - if p_notcoal > 0: - r[num_steps - 1] = np.sum(np.matmul(P, C)) / p_notcoal - else: - r[num_steps - 1] = np.nan + if j + 1 < num_steps: + dt = steps_b[j + 1] - time + G = (np.kron(M - dM, Identity) + np.kron(Identity, M - dM)) - C + P = np.matmul(P, _matrix_exponential(dt * G)) return r[keep_steps], p_t[keep_steps] def _pop_size_and_migration_at_t(self, t): diff --git a/tests/test_demography.py b/tests/test_demography.py index 7d14256e6..373eb165c 100644 --- a/tests/test_demography.py +++ b/tests/test_demography.py @@ -1826,6 +1826,83 @@ def test_ooa(self): assert np.allclose(0.5 / R[0], 29725.3435, rtol=1e-3) +class TestTerminalCoalescenceRate: + @pytest.mark.parametrize("validation", [False, True]) + @pytest.mark.parametrize("steps", [[100], [0, 100], [0, 50, 100]]) + @pytest.mark.parametrize("sizes", [(1000, 10000), (10000, 1000), (0.5, 0.25)]) + def test_size_change(self, validation, steps, sizes): + old_size, new_size = sizes + demography = msprime.Demography() + demography.add_population(name="A", initial_size=old_size) + demography.add_population_parameters_change( + time=100, population="A", initial_size=new_size + ) + debugger = demography.debug() + rates, survival = debugger.coalescence_rate_trajectory( + steps, {"A": 2}, double_step_validation=validation + ) + extended_rates, extended_survival = debugger.coalescence_rate_trajectory( + steps + [101], {"A": 2}, double_step_validation=validation + ) + np.testing.assert_allclose(rates, extended_rates[:-1]) + np.testing.assert_allclose(survival, extended_survival[:-1]) + assert rates[-1] == pytest.approx( + 1 / (2 * max(1, new_size)), rel=1e-12, abs=0 + ) + assert survival[-1] == pytest.approx(np.exp(-100 / (2 * max(1, old_size)))) + + @pytest.mark.parametrize("validation", [False, True]) + @pytest.mark.parametrize("proportion", [0.25, 0.5, 1]) + def test_lineage_movement(self, validation, proportion): + demography = msprime.Demography() + demography.add_population(name="A", initial_size=1000) + demography.add_population(name="B", initial_size=10000) + demography.add_mass_migration( + time=100, source="A", dest="B", proportion=proportion + ) + debugger = demography.debug() + rates, survival = debugger.coalescence_rate_trajectory( + [0, 100], {"A": 2}, double_step_validation=validation + ) + extended_rates, extended_survival = debugger.coalescence_rate_trajectory( + [0, 100, 101], {"A": 2}, double_step_validation=validation + ) + expected_rate = (1 - proportion) ** 2 / 2000 + proportion**2 / 20000 + assert rates[-1] == pytest.approx(expected_rate) + assert survival[-1] == pytest.approx(np.exp(-0.05)) + np.testing.assert_allclose(rates, extended_rates[:-1]) + np.testing.assert_allclose(survival, extended_survival[:-1]) + + @pytest.mark.parametrize("validation", [False, True]) + def test_multiple_lineage_movements(self, validation): + demography = msprime.Demography() + for name, size in [("A", 1000), ("B", 10000), ("C", 100000)]: + demography.add_population(name=name, initial_size=size) + demography.add_mass_migration(time=100, source="A", dest="B", proportion=0.5) + demography.add_mass_migration(time=100, source="B", dest="C", proportion=0.5) + rates, survival = demography.debug().coalescence_rate_trajectory( + [0, 100], {"A": 2}, double_step_validation=validation + ) + # At the event time each surviving lineage has probabilities + # (1/2, 1/4, 1/4) of being in A, B, C, respectively. + expected_rate = 0.5**2 / 2000 + 0.25**2 / 20000 + 0.25**2 / 200000 + assert rates[-1] == pytest.approx(expected_rate) + assert survival[-1] == pytest.approx(np.exp(-0.05)) + + @pytest.mark.parametrize("steps", [[0], [100], [0, 100]]) + def test_constant_size_singleton_and_zero(self, steps): + debugger = msprime.Demography.isolated_model([1000]).debug() + rates, survival = debugger.coalescence_rate_trajectory(steps, {0: 2}) + np.testing.assert_allclose(rates, 1 / 2000) + np.testing.assert_allclose(survival, np.exp(-np.array(steps) / 2000)) + + def test_zero_survival(self): + debugger = msprime.Demography.isolated_model([1]).debug() + rates, survival = debugger.coalescence_rate_trajectory([0, 2000], {0: 2}) + assert survival[-1] == 0 + assert np.isnan(rates[-1]) + + class TestEventTimes: """ Tests that demographic events occur when they should. From 56427173f6c013dc11903a9e0d641f83e514ea51 Mon Sep 17 00:00:00 2001 From: Zijian Zeng <39647839+ArtificialZeng@users.noreply.github.com> Date: Tue, 27 Oct 2026 18:19:49 +0800 Subject: [PATCH 2/2] Apply Ruff formatting to terminal-rate regression --- tests/test_demography.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_demography.py b/tests/test_demography.py index 373eb165c..6940de108 100644 --- a/tests/test_demography.py +++ b/tests/test_demography.py @@ -1846,9 +1846,7 @@ def test_size_change(self, validation, steps, sizes): ) np.testing.assert_allclose(rates, extended_rates[:-1]) np.testing.assert_allclose(survival, extended_survival[:-1]) - assert rates[-1] == pytest.approx( - 1 / (2 * max(1, new_size)), rel=1e-12, abs=0 - ) + assert rates[-1] == pytest.approx(1 / (2 * max(1, new_size)), rel=1e-12, abs=0) assert survival[-1] == pytest.approx(np.exp(-100 / (2 * max(1, old_size)))) @pytest.mark.parametrize("validation", [False, True])