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)	Perform imports for libraries with APIs matching numpy, scipy.integrate.ivp, scipy.sparse.
reimport_numerical_libs(context=None)	Reimport xp, xp_sparse, xp_ivp from the global context (in case they’ve been update to cupy).
sync_numerical_libs(func)	Decorator that pulls xp, xp_sparse, xp_ivp from the global bucky context into the wrapped function.
to_cpu_noop(x, stream=None, order='C', out=None)

bucky.numerical_libs.ExperimentalWarning

bucky.numerical_libs._qr

bucky.numerical_libs.empty_like_pinned

bucky.numerical_libs.empty_pinned

bucky.numerical_libs.optimize_kernels

bucky.numerical_libs.reimport_cache

bucky.numerical_libs.scatter_add

bucky.numerical_libs.special

bucky.numerical_libs.to_cpu

bucky.numerical_libs.xp

bucky.numerical_libs.xp_ivp

bucky.numerical_libs.xp_sparse

bucky.numerical_libs.zeros_like_pinned

bucky.numerical_libs.zeros_pinned

exception bucky.numerical_libs.MockExperimentalWarning

Bases: Warning

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

bucky.numerical_libs.enable_cupy(optimize=False)

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.reimport_numerical_libs(context=None)

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

bucky.numerical_libs.sync_numerical_libs(func)

Decorator that pulls 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)

bucky.make_input_graph bucky.postprocess