diff --git a/binning_utils/__init__.py b/binning_utils/__init__.py index 478e12d..8aaa729 100644 --- a/binning_utils/__init__.py +++ b/binning_utils/__init__.py @@ -271,3 +271,39 @@ def query_ball(bin_edges, start, stop): ee = np.arange(bin_start, bin_stop + 1, 1) mask = np.logical_and(ee >= 0, ee < num_bins) return ee[mask] + + +def draw_random_bin(prng, bin_apertures, size=None): + """ + Draws a random bin from multiple bins with different apertures. + The relative apertures of the bins is proportional to the probability + of the bins being drawn. + Defining bins by their apertures is useful in case one dimensional + bin_edges are not applicable in case the bins represent areas, solid angles + or other apertures of higher dimensions. + + Parameters + ---------- + prng : numpy.random.Generator + Pseudo random number generator. + bin_apertures : array_like, floats + The apertues of the bins. If the bins are one dimensional, this + is simply the widths of the bins. Apertures must be >= 0. + size : int or None (default None) + Number of bins to be drawn. Adopted from numpy.random. + If None, the return value is scalar. + + Returns + ------- + bin index : int + Index of bin + """ + bin_apertures = np.asarray(bin_apertures) + assert np.all(bin_apertures >= 0.0) + assert len(bin_apertures) > 0 + bin_order = np.arange(len(bin_apertures)) + prng.shuffle(bin_order) + cc = np.cumsum(bin_apertures[bin_order]) + cc_max = cc[-1] + c = prng.uniform(low=0.0, high=cc_max, size=size) + return bin_order[np.digitize(c, bins=cc)] diff --git a/binning_utils/tests/test_draw_random_bin.py b/binning_utils/tests/test_draw_random_bin.py new file mode 100644 index 0000000..937e1db --- /dev/null +++ b/binning_utils/tests/test_draw_random_bin.py @@ -0,0 +1,40 @@ +import binning_utils as bu +import numpy as np + + +def test_draw_random_bin_size(): + prng = np.random.Generator(np.random.PCG64(132)) + + r = bu.draw_random_bin(prng=prng, bin_apertures=[1, 2, 3]) + assert len(np.shape(r)) == 0 + + r = bu.draw_random_bin(prng=prng, bin_apertures=[1, 2, 3], size=None) + assert len(np.shape(r)) == 0 + + r = bu.draw_random_bin(prng=prng, bin_apertures=[1, 2, 3], size=1) + assert len(np.shape(r)) == 1 + assert r.shape[0] == 1 + + r = bu.draw_random_bin(prng=prng, bin_apertures=[1, 2, 3], size=10) + assert len(np.shape(r)) == 1 + assert r.shape[0] == 10 + + +def test_draw_bin(): + prng = np.random.Generator(np.random.PCG64(132)) + bin_apertures = [2, 100, 1, 5, 1, 8] + + bins = bu.draw_random_bin( + prng=prng, bin_apertures=bin_apertures, size=1000 * 1000 + ) + + buni, bcounts = np.unique(bins, return_counts=True) + + bin_counts_normalized = bcounts / np.sum(bcounts) + bin_apertures_normalized = bin_apertures / np.sum(bin_apertures) + + np.testing.assert_array_almost_equal( + bin_counts_normalized, + bin_apertures_normalized, + decimal=3, + ) diff --git a/binning_utils/version.py b/binning_utils/version.py index d62d967..1f658a4 100644 --- a/binning_utils/version.py +++ b/binning_utils/version.py @@ -1 +1 @@ -__version__ = "0.0.16" +__version__ = "0.0.17"