"""Global configuration for the computational backend used by dynasor's
compute functions.
Keeping this outside the compute functions' own signatures means
``compute_dynamic_structure_factors``/``compute_static_structure_factors``
don't need ``backend``/``precision``/``batch_size`` kwargs of their own --
set the backend once, and every subsequent call uses it::
from dynasor import set_compute_config, compute_dynamic_structure_factors
set_compute_config(backend='torch', precision='float32')
sample = compute_dynamic_structure_factors(traj, q_points, dt=25.0, window_size=2000)
See :ref:`the backends reference page <backends>` for the full picture,
including which parts of a calculation this actually affects.
"""
import importlib
from dataclasses import dataclass
from typing import Optional
_VALID_BACKENDS = ('numba', 'torch', 'cupy')
_VALID_PRECISIONS = ('float64', 'float32')
_DEFAULT_BATCH_SIZE = 32
@dataclass(frozen=True)
class ComputeConfig:
"""Immutable snapshot of the current compute configuration.
Attributes
----------
backend
One of ``'numba'``, ``'torch'``, or ``'cupy'``.
precision
One of ``'float64'`` or ``'float32'``. Only affects the GPU backends;
ignored for ``backend='numba'``.
batch_size
Number of trajectory frames stacked into a single GPU kernel call.
Only affects the GPU backends; ignored for ``backend='numba'``. Larger
values make better use of GPU memory bandwidth for small systems, but
scale up GPU memory usage roughly linearly (``batch_size * N_atoms *
N_qpoints``) -- reduce it for large systems/q-point grids to avoid an
out-of-memory error. See :ref:`the backends reference page <backends>`.
"""
backend: str
precision: str
batch_size: int
_config = ComputeConfig(backend='numba', precision='float64', batch_size=_DEFAULT_BATCH_SIZE)
[docs]
def set_compute_config(
backend: Optional[str] = None,
precision: Optional[str] = None,
batch_size: Optional[int] = None,
) -> None:
"""Set the computational backend, precision, and/or GPU batch size used
by all subsequent calls to
``compute_dynamic_structure_factors``/``compute_static_structure_factors``.
Parameters
----------
backend
One of ``'numba'`` (CPU, default), ``'torch'``, or ``'cupy'`` (both CUDA
GPU, optional dependencies). Leave unset to keep the current value.
precision
One of ``'float64'`` (default) or ``'float32'``. Only affects the GPU
backends. Leave unset to keep the current value.
batch_size
Number of trajectory frames stacked into a single GPU kernel call
(default 32). Only affects the GPU backends. Leave unset to keep the
current value; see :attr:`ComputeConfig.batch_size` for tuning guidance.
Examples
--------
>>> from dynasor import set_compute_config
>>> set_compute_config(backend='torch', precision='float32')
"""
global _config
new_backend = _config.backend if backend is None else backend
new_precision = _config.precision if precision is None else precision
new_batch_size = _config.batch_size if batch_size is None else batch_size
if new_backend not in _VALID_BACKENDS:
raise ValueError(f'backend must be one of {_VALID_BACKENDS}, got {new_backend!r}')
if new_precision not in _VALID_PRECISIONS:
raise ValueError(f'precision must be one of {_VALID_PRECISIONS}, got {new_precision!r}')
if not isinstance(new_batch_size, int) or isinstance(new_batch_size, bool) \
or new_batch_size < 1:
raise ValueError(f'batch_size must be a positive integer, got {new_batch_size!r}')
_config = ComputeConfig(backend=new_backend, precision=new_precision,
batch_size=new_batch_size)
_GPU_BACKEND_PACKAGES = {'torch': ('torch', 'gpu-torch'), 'cupy': ('cupy', 'gpu-cupy')}
def check_backend_is_available(backend: str) -> None:
"""Check that *backend* can run on this machine, and raise a message
naming what is missing if it cannot.
Selecting a backend only records a name, so the optional dependency and
the CUDA device it needs are looked at here, where the calculation is
about to start. Without this the first thing the user sees is an error
from torch or cupy about a missing device, raised from inside dynasor
after the trajectory has already been opened.
Parameters
----------
backend
One of the backends accepted by :func:`set_compute_config`.
``'numba'`` needs nothing and always passes.
Raises
------
RuntimeError
If the package the backend needs is not installed, or if it is
installed but reports no CUDA device.
Examples
--------
>>> from dynasor.compute_config import check_backend_is_available
>>> check_backend_is_available('numba')
"""
if backend not in _GPU_BACKEND_PACKAGES:
return
package, extra = _GPU_BACKEND_PACKAGES[backend]
try:
module = importlib.import_module(package)
except ImportError as error:
raise RuntimeError(
f'backend {backend!r} requires the {package} package, which is not installed; '
f'install it with "pip install dynasor[{extra}]", or use backend="numba" to run '
'on the CPU') from error
is_available = module.cuda.is_available if backend == 'torch' else module.is_available
if not is_available():
raise RuntimeError(
f'backend {backend!r} requires a CUDA device, but {package} reports none; use '
'backend="numba" to run on the CPU. See the backends page of the dynasor '
'documentation for how to select a device.')
[docs]
def get_compute_config() -> ComputeConfig:
"""Return the current compute configuration.
Returns
-------
ComputeConfig
The current ``backend``/``precision``/``batch_size`` setting;
``ComputeConfig(backend='numba', precision='float64', batch_size=32)``
if :func:`set_compute_config` has never been called.
"""
return _config
