diff --git a/pixi.lock b/pixi.lock index a016539..fa8e3d0 100644 --- a/pixi.lock +++ b/pixi.lock @@ -5595,6 +5595,8 @@ packages: - pandas>=2.0 - h5py>=3.8 - pyarrow>=14 + - anndata + - mudata - pysceptre - pytest>=7 ; extra == 'dev' - ruff>=0.5 ; extra == 'dev' diff --git a/src/watteg/simulate.py b/src/watteg/simulate.py index 11313f2..3b498a2 100644 --- a/src/watteg/simulate.py +++ b/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 diff --git a/tests/test_simulation_core.py b/tests/test_simulation_core.py index 3849a5c..2b20586 100644 --- a/tests/test_simulation_core.py +++ b/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():