Setting the Backend

GWPopulation provides a unified interface to a number of numpy/scipy like APIs.

The backend can be set using gwpopulation.set_backend(backend), where backend is one of numpy, cupy, jax.

Downstream packages can automatically track the active backend using entry_points. With this set up, packages can use xp and scs in specified modules. Additionally, users can provide a full arbitrary scipy object to be used if anything beyond scipy.special is needed. An example of how to set numpy, scipy.special, and the toeplitz function from scipy.linalg via the setup.cfg file is shown below. Specification using pyproject.toml and setup.py follows slightly different syntax documentation for which can be found online.

[options.entry_points]
gwpopulation.xp =
    mypackage_foo = mypackage.foo
gwpopulation.scs =
    mypackage_foo = mypackage.foo
    mypackage_bar = mypackage.bar
gwpopulation.other =
    mypackage_baz_toeplitz = mypackage.baz:scipy.linalg.toeplitz

Note

Each module that wants to use the GWPopulation backend must be specified independently for the automatic propagation to work.

If there is a backend that you would like to use that is not currently supported, please open an issue.

Precision

With jax, 64-bit floats are enabled the first time the backend is set. For single precision, which can be several times faster on GPUs, especially those with limited double-precision throughput, use gwpopulation.set_backend("jax", enable_x64=False). Arrays keep the precision they were created with, so set it before creating any arrays.

Densities below about 1e-38 underflow to zero in single precision, so use it with models that provide a log density (see Population models). The likelihood and selection functions form their Monte Carlo sums in log space, and the sampling priors can be given as ln_prior. A model without a log density still works, but its log density is -inf wherever the density underflows.