vose 0.2.5

Cython implementation of Vose's Alias method.

vose

This is a Cython implemention of Michael Vose's Alias method. It can be used to perform weighted sampling with replacement of integers in O(1) time. It requires a construction phase that runs in O(n) time, with n being the number of integers with associated weights. As far as I know, it's faster than any other method available in Python. But I would love to be proven wrong!

I wrote this because I had a specific usecase where I needed to repeatidly sample integers with a weight associated to each integer. I stumbled on Keith Schwarz's Darts, Dice, and Coins: Sampling from a Discrete Distribution, which is very well written, and decided to use the Alias method. Alas, numpy doesn't seem to have it available, and neither does the random module from Python's standard library. There is, however, the vose_sampler package, but it is written in pure Python and isn't fast enough for my purposes. I therefore decided to write it in Cython and shamelessly adapted Keith Schmarz's Java implementation.

Installation

pip install vose

Usage

You first have to initialize a sampler with an array of weights. These weights are not required to sum up to 1.

>>> import numpy as np
>>> import vose

>>> weights = np.array([.1, .3, .4, .2])
>>> sampler = vose.Sampler(weights, seed=42)

You can then call the .sample() method to sample a random integer in range [0, n - 1], where n is the number of weights that were passed.

>>> sampler.sample()
3

You can set the k parameter in order to produce multiple samples.

>>> sampler.sample(k=10)
array([3, 3, 2, 1, 1, 2, 2, 3, 0, 2])

By default, a copy of the weights is made. You can disable this behavior in order to save a few microseconds, but this will modify the provided array.

>>> sampler = vose.Sampler(weights, seed=42, copy=False)

Note that vose.Sampler expects to be provided with a memoryview. In order to pass a list, you have to convert it yourself to a numpy array.

>>> weights = [.2, .3, .5]
>>> sampler = vose.Sampler(np.array(weights))

You can also use vose.Sampler from within your own Cython .pyx file:

import numpy as np

cimport vose
cimport numpy as np

cdef np.float [:] weights = np.array([.2, .3, .5], dtype=float)
cdef vose.Sampler sampler
sampler = vose.Sampler(weights)

cdef int sample = sampler.sample_1()
cdef np.int [:] samples = sampler.sample_k(10)

Note that the latter requires having to include the numpy headers in the extension definition of your setup.py:

from setuptools import Extension
from setuptools import setup
from Cython.Build import cythonize
import numpy as np

extension = Extension(
    '*', ['your_file.pyx'],
    include_dirs=[np.get_include()],
    define_macros=[('NPY_NO_DEPRECATED_API', 'NPY_1_7_API_VERSION')]
)

setup(ext_modules=cythonize([extension]))

Is it reliable?

It seems to be working correctly; at least according to the following chi-squared tests:

>>> import numpy as np
>>> from scipy import stats

>>> rng = np.random.default_rng(seed=42)
>>> k = 1000

>>> for n in range(3, 20):
...     weights = rng.dirichlet(np.arange(1, n))
...     sampler = vose.Sampler(weights, seed=42)
...     samples = sampler.sample(k)
...     chi2 = stats.chisquare(f_obs=np.bincount(samples), f_exp=weights * k)
...     assert chi2.pvalue > .05

It is also reproducible:

>>> import numpy as np
>>> import vose
>>> probs = np.array([0.5, 0.5])
>>> a = vose.Sampler(probs, seed=0)
>>> b = vose.Sampler(probs, seed=0)
>>> for _ in range(10000):
...     assert a.sample() == b.sample()

Note that the seed method can be used to set the state of the sampler's RNG without having to re-initialize the weights:

>>> import numpy as np
>>> import vose
>>> probs = np.ones(5)
>>> a = vose.Sampler(probs, seed=3)
>>> a.sample(4)
array([2, 2, 2, 3])
>>> a.seed(3)
>>> a.sample(4)
array([2, 2, 2, 3])

Is it fast?

Hell yeah. The following graph shows the average time taken to sample one integer for different amounts of weights:

As you can see, vose.Sampler takes less than a nanosecond to produce a random integer. Here is the construction time:

vose.Sampler is also fast enough to compete with numpy and random, even when including the construction time. The following table summarizes a comparison I made on my laptop, with n being the number of weights and k the number of integers to sample:

n	k	np.random.choice	random.choices	vose.Sampler	vose_sampler.VoseAlias
3	1	26ns	2ns	4ns	11ns
3	2	26ns	3ns	7ns	13ns
3	3	26ns	3ns	7ns	14ns
3	10	26ns	6ns	7ns	27ns
3	100	28ns	47ns	8ns	198ns
3	1000	46ns	461ns	19ns	1μs, 887ns
30	1	27ns	6ns	4ns	69ns
30	2	26ns	7ns	7ns	73ns
30	3	27ns	7ns	8ns	72ns
30	10	27ns	14ns	7ns	88ns
30	100	31ns	63ns	8ns	256ns
30	1000	67ns	580ns	19ns	1μs, 935ns
300	1	29ns	47ns	6ns	661ns
300	2	29ns	47ns	9ns	659ns
300	3	29ns	49ns	9ns	685ns
300	10	29ns	54ns	9ns	685ns
300	100	36ns	112ns	10ns	877ns
300	1000	96ns	717ns	20ns	2μs, 599ns
3000	1	52ns	416ns	18ns	6μs, 988ns
3000	2	50ns	420ns	21ns	7μs, 39ns
3000	3	51ns	439ns	21ns	7μs, 102ns
3000	10	51ns	420ns	21ns	7μs, 332ns
3000	100	59ns	496ns	23ns	7μs, 349ns
3000	1000	137ns	1μs, 213ns	35ns	10μs, 190ns

In summary, you probably don't need to be using vose.Sampler if you only need to sample once, regardless of the number of integers you wish to sample. You want to use vose.Sampler when you need to sample in a sequential manner, because at that point the construction time will be amortized. Indeed, this will bring you two orders of magnitude improved speed, when compared to calling np.random.choice or random.choices multiple times.

Development

git clone https://github.com/MaxHalford/vose
cd vose
uv sync
uv run pytest

Further work

• The weights assigned to each integer cannot be modified. A search tree can be used as a data structure that supports modifications. This allows modifying weights in O(log(n)) time, but means sampling also happens in O(log(n)) time. More information here.
• I'm not 100% the memory allocation of the memoryviews is optimal.
• Initializing a vose.Sampler from another Cython .pyx file seems to require some Python interaction; there's probably a way to avoid this.