abtem.core.backend

`abtem.core.backend`#

Module for handling the array backend (NumPy, CuPy, Dask, etc.) of the library.

Module Contents#

Functions#

`asnumpy`	Convert an array to NumPy.
`check_cupy_is_installed`	Check if CuPy is installed, raise an error if not.
`copy_to_device`	Copy an array to a different device (CPU or GPU) using CuPy.
`device_name_from_array_module`	Get the device string from the array module. The array module must be either NumPy or CuPy.
`ensure_cuda_cluster`	Start a dask-cuda cluster spanning all visible GPUs and return its client.
`get_array_module`	Get the array module (NumPy or CuPy) for a given array or string.
`get_cuda_cluster_client`	Return the dask-cuda cluster client started by `ensure_cuda_cluster`.
`get_ndimage_module`	Get the ndimage module for a given array or device string.
`get_scipy_module`	Get the SciPy module for a given array or device string.
`is_gpu_dask_client`	Check whether a distributed client can safely execute CuPy computations.
`push_config_to_workers`	Mirror this process’s abTEM configuration onto the client’s workers.
`validate_device`	Validate the device string.

Data#

`ArrayModule`
`logger`

API#

abtem.core.backend.ArrayModule#: None

abtem.core.backend.asnumpy(...)[source]#

Convert an array to NumPy.

Parameters:: array (ndarray, dask.array.Array) – The array to convert.
Returns:: The array converted to NumPy.
Return type:: ndarray

abtem.core.backend.check_cupy_is_installed()[source]#: Check if CuPy is installed, raise an error if not.

abtem.core.backend.copy_to_device(...)[source]#

Copy an array to a different device (CPU or GPU) using CuPy.

Parameters:

array (ndarray) – The array to copy.
device (str) – The device to copy to. Either ‘cpu’ or ‘gpu’.

Returns:

The array copied to the specified device.

Return type:

ndarray or cupy.ndarray

abtem.core.backend.device_name_from_array_module(...) → str[source]#

Get the device string from the array module. The array module must be either NumPy or CuPy.

Parameters:: xp (numpy or cupy) – The array module.
Returns:: The device string.
Return type:: str

abtem.core.backend.ensure_cuda_cluster()[source]#

Start a dask-cuda cluster spanning all visible GPUs and return its client.

The cluster assigns one worker process to each GPU, allowing dask to distribute computations across all of them. It is created once per process and reused on subsequent calls. Requires the optional dask-cuda package.

Returns:: The client connected to the dask-cuda cluster.
Return type:: distributed.Client

abtem.core.backend.get_array_module(...) → ModuleType[source]#

Get the array module (NumPy or CuPy) for a given array or string.

Parameters:: x (ndarray, cupy.ndarray, dask.array.Array, str, None) – The array or string to get the array module for. If None, the default device is used.
Returns:: The array module.
Return type:: numpy or cupy

abtem.core.backend.get_cuda_cluster_client()[source]#

Return the dask-cuda cluster client started by ensure_cuda_cluster.

Returns the running client, or None when no cluster has been started or the previous one was shut down. Unlike ensure_cuda_cluster this never starts a cluster, which makes it suitable for inspecting whether multi-GPU execution is active (e.g. from benchmark or verification scripts).

Returns:: The client connected to the running dask-cuda cluster, if any.
Return type:: distributed.Client or None

abtem.core.backend.get_ndimage_module(...) → ModuleType[source]#

Get the ndimage module for a given array or device string.

Parameters:: x (ndarray, cupy.ndarray, dask.array.Array, str, None) – The array or string to get the ndimage module for. If None, the default device is used.
Returns:: The ndimage module.
Return type:: scipy.ndimage or cupyx.ndimage

abtem.core.backend.get_scipy_module(...)[source]#

Get the SciPy module for a given array or device string.

Parameters:: x (ndarray, cupy.ndarray, dask.array.Array, str, None) – The array or string to get the SciPy module for. If None, the default device is used.
Returns:: The SciPy module.
Return type:: scipy or cupyx.scipy

abtem.core.backend.is_gpu_dask_client(...) → bool[source]#

Check whether a distributed client can safely execute CuPy computations.

Only a client whose workers are each single-threaded — as produced by dask_cuda.LocalCUDACluster, which additionally pins one GPU per worker — is considered suitable. The threaded scheduler and multi-threaded workers share a single CUDA context per process, which cannot be used with CuPy.

Parameters:: client (distributed.Client or None) – The client to check.
Returns:: True if the client is running and all of its workers are single-threaded.
Return type:: bool

abtem.core.backend.logger#: ‘getLogger(…)’

abtem.core.backend.push_config_to_workers(...)[source]#

Mirror this process’s abTEM configuration onto the client’s workers.

abTEM resolves configuration inside tasks, in the worker process – get_dtype reads precision at call time, for example – but worker processes only ever see the defaults: abtem.config.set in the client does not reach them, silently changing results (a float64 computation dispatched to default-configured workers runs in float32).

The snapshot is carried by a named worker plugin, so workers that join or restart later also receive it; when the configuration changes, re-registering under the same name replaces the plugin and re-runs its setup on all workers. Repeated pushes of an unchanged configuration to the same client (keyed on client.id) are skipped.

abtem.core.backend.validate_device(...) → str[source]#

Validate the device string.

Parameters:: device (str, None) – The device string to validate. Must be either ‘cpu’ or ‘gpu’. If None, the device from the configuration is used.
Returns:: The validated device string.
Return type:: str