bucky.numerical_libs#

Provides an interface to import numerical libraries using the GPU (if available).

The main goal of this is to smooth over the differences between numpy and cupy so that the rest of the code can use them interchangably. We also need to monkey patch scipy’s ivp solver to work on cupy arrays.

Notes

Linters HATE this module because it’s really abusing the import system (by design).

Module Contents#

Functions#

enable_cupy(optimize=False, cache_dir=None)

Perform imports for libraries with APIs matching numpy, scipy.integrate.ivp, scipy.sparse.

ensure_batched_qr(np)

Monkey patch older versions of numpy to provide vectorized qr decomp.

reimport_numerical_libs(context=None)

Reimport xp, xp_sparse, xp_ivp from the global context (in case they've been updated to cupy).

sync_numerical_libs(func)

Decorator pullng xp, xp_sparse, xp_ivp from the global bucky context into the wrapped function.

to_cpu_noop(x, stream=None, order='C', out=None)

NOOP function that accounts for possible args of to_cpu().

Attributes#

CUPY_FORCE_FP32

ExperimentalWarning

device_count

empty_like_pinned

empty_pinned

is_cupy

optimize_kernels

reimport_cache

scatter_add

signal

special

to_cpu

xp

xp_ivp

xp_qr

xp_sparse

zeros_like_pinned

zeros_pinned

bucky.numerical_libs.CUPY_FORCE_FP32 = False#
bucky.numerical_libs.ExperimentalWarning#
bucky.numerical_libs.device_count = 0#
bucky.numerical_libs.empty_like_pinned#
bucky.numerical_libs.empty_pinned#
bucky.numerical_libs.is_cupy = False#
bucky.numerical_libs.optimize_kernels#
bucky.numerical_libs.reimport_cache#
bucky.numerical_libs.scatter_add#
bucky.numerical_libs.signal#
bucky.numerical_libs.special#
bucky.numerical_libs.to_cpu#
bucky.numerical_libs.xp[source]#
bucky.numerical_libs.xp_ivp[source]#
bucky.numerical_libs.xp_qr[source]#
bucky.numerical_libs.xp_sparse[source]#
bucky.numerical_libs.zeros_like_pinned#
bucky.numerical_libs.zeros_pinned#
exception bucky.numerical_libs.MockExperimentalWarning[source]#

Bases: Warning

Simple class to mock the optuna warning if we don’t have optuna.

bucky.numerical_libs.enable_cupy(optimize=False, cache_dir=None)[source]#

Perform imports for libraries with APIs matching numpy, scipy.integrate.ivp, scipy.sparse.

These imports will use a monkey-patched version of these modules that has had all it’s numpy references replaced with CuPy.

if optimize is True, place the kernel optimization context in xp.optimize_kernels, otherwise make it a nullcontext (noop)

returns nothing but imports a version of ‘xp’, ‘ivp’, and ‘sparse’ to the global scope of this module

Parameters:

optimize (bool) – Enable kernel optimization in cupy >=v8.0.0. This will slow down initial function call (mostly reduction operations) but will offer better performance for repeated calls (e.g. in the RHS call of an integrator).

Returns:

exit_code – Non-zero value indicates error code, or zero on success.

Return type:

int

Raises:

NotImplementedError – If the user calls a monkeypatched function of the libs that isn’t fully implemented.

bucky.numerical_libs.ensure_batched_qr(np)[source]#

Monkey patch older versions of numpy to provide vectorized qr decomp.

bucky.numerical_libs.reimport_numerical_libs(context=None)[source]#

Reimport xp, xp_sparse, xp_ivp from the global context (in case they’ve been updated to cupy).

bucky.numerical_libs.sync_numerical_libs(func)[source]#

Decorator pullng xp, xp_sparse, xp_ivp from the global bucky context into the wrapped function.

bucky.numerical_libs.to_cpu_noop(x, stream=None, order='C', out=None)[source]#

NOOP function that accounts for possible args of to_cpu().