Coverage for dynasor/compute_config.py: 89%
34 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-03 19:46 +0000
1"""Global configuration for the computational backend used by dynasor's
2compute functions.
4Keeping this outside the compute functions' own signatures means
5``compute_dynamic_structure_factors``/``compute_static_structure_factors``
6don't need ``backend``/``precision``/``batch_size`` kwargs of their own --
7set the backend once, and every subsequent call uses it::
9 from dynasor import set_compute_config, compute_dynamic_structure_factors
11 set_compute_config(backend='torch', precision='float32')
12 sample = compute_dynamic_structure_factors(traj, q_points, dt=25.0, window_size=2000)
14See :ref:`the backends reference page <backends>` for the full picture,
15including which parts of a calculation this actually affects.
16"""
18import importlib
19from dataclasses import dataclass
20from typing import Optional
22_VALID_BACKENDS = ('numba', 'torch', 'cupy')
23_VALID_PRECISIONS = ('float64', 'float32')
24_DEFAULT_BATCH_SIZE = 32
27@dataclass(frozen=True)
28class ComputeConfig:
29 """Immutable snapshot of the current compute configuration.
31 Attributes
32 ----------
33 backend
34 One of ``'numba'``, ``'torch'``, or ``'cupy'``.
35 precision
36 One of ``'float64'`` or ``'float32'``. Only affects the GPU backends;
37 ignored for ``backend='numba'``.
38 batch_size
39 Number of trajectory frames stacked into a single GPU kernel call.
40 Only affects the GPU backends; ignored for ``backend='numba'``. Larger
41 values make better use of GPU memory bandwidth for small systems, but
42 scale up GPU memory usage roughly linearly (``batch_size * N_atoms *
43 N_qpoints``) -- reduce it for large systems/q-point grids to avoid an
44 out-of-memory error. See :ref:`the backends reference page <backends>`.
45 """
46 backend: str
47 precision: str
48 batch_size: int
51_config = ComputeConfig(backend='numba', precision='float64', batch_size=_DEFAULT_BATCH_SIZE)
54def set_compute_config(
55 backend: Optional[str] = None,
56 precision: Optional[str] = None,
57 batch_size: Optional[int] = None,
58) -> None:
59 """Set the computational backend, precision, and/or GPU batch size used
60 by all subsequent calls to
61 ``compute_dynamic_structure_factors``/``compute_static_structure_factors``.
63 Parameters
64 ----------
65 backend
66 One of ``'numba'`` (CPU, default), ``'torch'``, or ``'cupy'`` (both CUDA
67 GPU, optional dependencies). Leave unset to keep the current value.
68 precision
69 One of ``'float64'`` (default) or ``'float32'``. Only affects the GPU
70 backends. Leave unset to keep the current value.
71 batch_size
72 Number of trajectory frames stacked into a single GPU kernel call
73 (default 32). Only affects the GPU backends. Leave unset to keep the
74 current value; see :attr:`ComputeConfig.batch_size` for tuning guidance.
76 Examples
77 --------
78 >>> from dynasor import set_compute_config
79 >>> set_compute_config(backend='torch', precision='float32')
80 """
81 global _config
82 new_backend = _config.backend if backend is None else backend
83 new_precision = _config.precision if precision is None else precision
84 new_batch_size = _config.batch_size if batch_size is None else batch_size
85 if new_backend not in _VALID_BACKENDS:
86 raise ValueError(f'backend must be one of {_VALID_BACKENDS}, got {new_backend!r}')
87 if new_precision not in _VALID_PRECISIONS:
88 raise ValueError(f'precision must be one of {_VALID_PRECISIONS}, got {new_precision!r}')
89 if not isinstance(new_batch_size, int) or isinstance(new_batch_size, bool) \
90 or new_batch_size < 1:
91 raise ValueError(f'batch_size must be a positive integer, got {new_batch_size!r}')
92 _config = ComputeConfig(backend=new_backend, precision=new_precision,
93 batch_size=new_batch_size)
96_GPU_BACKEND_PACKAGES = {'torch': ('torch', 'gpu-torch'), 'cupy': ('cupy', 'gpu-cupy')}
99def check_backend_is_available(backend: str) -> None:
100 """Check that *backend* can run on this machine, and raise a message
101 naming what is missing if it cannot.
103 Selecting a backend only records a name, so the optional dependency and
104 the CUDA device it needs are looked at here, where the calculation is
105 about to start. Without this the first thing the user sees is an error
106 from torch or cupy about a missing device, raised from inside dynasor
107 after the trajectory has already been opened.
109 Parameters
110 ----------
111 backend
112 One of the backends accepted by :func:`set_compute_config`.
113 ``'numba'`` needs nothing and always passes.
115 Raises
116 ------
117 RuntimeError
118 If the package the backend needs is not installed, or if it is
119 installed but reports no CUDA device.
121 Examples
122 --------
123 >>> from dynasor.compute_config import check_backend_is_available
124 >>> check_backend_is_available('numba')
125 """
126 if backend not in _GPU_BACKEND_PACKAGES:
127 return
128 package, extra = _GPU_BACKEND_PACKAGES[backend]
129 try:
130 module = importlib.import_module(package)
131 except ImportError as error:
132 raise RuntimeError(
133 f'backend {backend!r} requires the {package} package, which is not installed; '
134 f'install it with "pip install dynasor[{extra}]", or use backend="numba" to run '
135 'on the CPU') from error
136 is_available = module.cuda.is_available if backend == 'torch' else module.is_available
137 if not is_available():
138 raise RuntimeError(
139 f'backend {backend!r} requires a CUDA device, but {package} reports none; use '
140 'backend="numba" to run on the CPU. See the backends page of the dynasor '
141 'documentation for how to select a device.')
144def get_compute_config() -> ComputeConfig:
145 """Return the current compute configuration.
147 Returns
148 -------
149 ComputeConfig
150 The current ``backend``/``precision``/``batch_size`` setting;
151 ``ComputeConfig(backend='numba', precision='float64', batch_size=32)``
152 if :func:`set_compute_config` has never been called.
153 """
154 return _config