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

reimport_numerical_libs(context=None)

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

use_cupy(optimize=False)

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

bucky.numerical_libs.ExperimentalWarning[source]
bucky.numerical_libs.optimize_kernels[source]
bucky.numerical_libs.reimport_cache[source]
bucky.numerical_libs.scatter_add[source]
bucky.numerical_libs.special[source]
bucky.numerical_libs.to_cpu[source]
bucky.numerical_libs.xp[source]
bucky.numerical_libs.xp_ivp[source]
bucky.numerical_libs.xp_sparse[source]
exception bucky.numerical_libs.MockExperimentalWarning[source]

Bases: Warning

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

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

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

bucky.numerical_libs.use_cupy(optimize=False)[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.