Promote a simulated count too big for int16 instead of refusing it

src/watteg/simulate.py

@@ -16,10 +16,14 @@
 
 import numpy as np
 
-# int16 holds counts to 32,767. The largest single count in a real screen
-# measured here is 2,717, but a simulated draw has a tail, so the promotion is
-# checked rather than assumed.
+# int16 holds counts to 32,767, which covers almost every draw. Not all of them: sceptre's fitted
+# mean for a very highly expressed gene can reach far past anything observed in the cells with
+# extreme covariates -- HBA2 in DC-TAP K562 has an observed maximum of 6,583 and a fitted mean of
+# 20,104 in its most extreme cell, and a draw from NB(mean 20,104, theta 2.1) passes 32,767 about
+# one time in seven. Those draws are legitimate draws from the model the test assumes, so a draw
+# that does not fit is promoted to the next integer width rather than refused.
 _COUNT_DTYPE = np.int16
+_WIDER = (np.int16, np.int32, np.int64)
 
 
 def draw_counts(
@@ -38,9 +42,11 @@ def draw_counts(
     `n(1-p)/p`, so `n = theta` and `p = theta / (theta + mu)` give mean `mu`
     and variance `mu + mu^2/theta` -- R's `rnbinom(mu=, size=)`.
 
-    Returned as `int16` by default. These are counts; holding them as float64
-    costs four times the memory for no information, and the simulation's whole
-    shape depends on how many replicates fit in one process.
+    Returned as `int16` by default, or the narrowest of int16/int32/int64 at
+    least as wide as `dtype` that holds the largest draw. These are counts;
+    holding them as float64 costs four times the memory for no information, and
+    the simulation's whole shape depends on how many replicates fit in one
+    process. A draw is never wrapped or clipped.
     """
     baseline = np.asarray(baseline, dtype=float)
     if baseline.shape != effect_size.shape:
@@ -60,10 +66,11 @@ def draw_counts(
 
     if dtype is None:
         return counts
-    info = np.iinfo(dtype)
-    if counts.max(initial=0) > info.max:
-        raise OverflowError(
-            f"a simulated count exceeded {dtype.__name__}'s range ({info.max}); pass "
-            "dtype=None to keep the draw at full width"
-        )
-    return counts.astype(dtype)
+    largest = counts.max(initial=0)
+    for candidate in _WIDER:
+        if (
+            np.dtype(candidate).itemsize >= np.dtype(dtype).itemsize
+            and largest <= np.iinfo(candidate).max
+        ):
+            return counts.astype(candidate)
+    return counts

tests/test_simulation_core.py

@@ -510,18 +510,27 @@ def test_a_fully_knocked_down_gene_draws_zeros_rather_than_failing():
     assert (counts == 0).all()
 
 
-def test_counts_come_back_as_int16_and_overflow_is_refused_not_wrapped():
+def test_counts_come_back_as_int16_and_a_draw_too_big_for_it_is_promoted_not_wrapped():
     counts = draw_counts(
         np.full((1, 20), 3.0), np.ones((1, 20)), np.array([5.0]), np.random.default_rng(8)
     )
     assert counts.dtype == np.int16
-    with pytest.raises(OverflowError, match="int16"):
-        draw_counts(
-            np.full((1, 200), 1e5),
-            np.ones((1, 200)),
-            np.array([1e6]),
-            np.random.default_rng(9),
-        )
+    big = draw_counts(
+        np.full((1, 200), 1e5),
+        np.ones((1, 200)),
+        np.array([1e6]),
+        np.random.default_rng(9),
+    )
+    assert big.dtype == np.int32
+    exact = draw_counts(
+        np.full((1, 200), 1e5),
+        np.ones((1, 200)),
+        np.array([1e6]),
+        np.random.default_rng(9),
+        dtype=None,
+    )
+    np.testing.assert_array_equal(big, exact)
+    assert big.min() > np.iinfo(np.int16).max
 
 
 def test_the_draw_rejects_inputs_that_do_not_line_up():