fixes bug in PoissonIntervals distribution where if last rate is zero, it cycles back to beginning too early

geraintpalmer · geraintpalmer · commit 88944647a8b1 · 2025-03-25T10:42:40.000Z
diff --git a/ciw/dists/distributions.py b/ciw/dists/distributions.py
@@ -561,9 +561,7 @@ def __init__(self, rates, endpoints, max_sample_date):
         self.max_sample_date = max_sample_date
         self.get_intervals()
         self.get_dates()
-        self.inter_arrivals = [t - s for s, t in zip(self.dates, self.dates[1:])]
-        if self.inter_arrivals == []:
-            self.inter_arrivals = [float("inf")]
+        self.inter_arrivals = [t - s for s, t in zip(self.dates, self.dates[1:])] + [float('inf')]
         super().__init__(self.inter_arrivals)
 
     def __repr__(self):
diff --git a/ciw/tests/test_sampling.py b/ciw/tests/test_sampling.py
@@ -1434,7 +1434,7 @@ def test_poissoninterval_dist_object(self):
         expected_dates = [0]
         for t in Pi.inter_arrivals:
             expected_dates.append(expected_dates[-1] + t)
-        self.assertEqual(Pi.dates, expected_dates)
+        self.assertEqual(Pi.dates, expected_dates[:-1])
         self.assertLessEqual(Pi.dates[-1], Pi.max_sample_date)
 
         self.assertRaises(
@@ -1483,19 +1483,38 @@ def test_poissoninterval_dist_object(self):
     def test_poissoninterval_rate_zero(self):
         ciw.seed(5)
         Pi = ciw.dists.PoissonIntervals(
-            rates=[10, 0], endpoints=[1, 2], max_sample_date=15
+            rates=[10, 0], endpoints=[1, 2], max_sample_date=15.5
         )
         arrivals_when_rate_is_zero = [date for date in Pi.dates if int(date) % 2 == 1]
         self.assertEqual(arrivals_when_rate_is_zero, [])
+        final_date = Pi.dates[-1]
+        self.assertTrue(final_date < 15)
 
         ciw.seed(5)
         Pi = ciw.dists.PoissonIntervals(
-            rates=[0, 0], endpoints=[1, 2], max_sample_date=15
+            rates=[0, 0], endpoints=[1, 2], max_sample_date=15.5
         )
         arrivals_when_rate_is_zero = [date for date in Pi.dates if int(date) % 2 == 1]
         self.assertEqual(arrivals_when_rate_is_zero, [])
         self.assertEqual(Pi.sample(), float("inf"))
 
+        ciw.seed(5)
+        Pi = ciw.dists.PoissonIntervals(
+            rates=[3, 2, 0], endpoints=[1, 2, 10], max_sample_date=9
+        )
+        final_date = Pi.dates[-1]
+        self.assertTrue(final_date < 2)
+        
+        t = 0
+        arrivals = []
+        while t < 8:
+            t += Pi.sample()
+            arrivals.append(t)
+        final_sample = arrivals[-1]
+        final_arrival = arrivals[-2]
+        self.assertEqual(final_sample, float('inf'))
+        self.assertEqual(final_arrival, final_date)
+
     def test_poissoninterval_against_theory(self):
         """
         rates = [5, 1.5, 3]
