Source code for strongcoca.compute_config

from dataclasses import dataclass
from typing import Optional

_VALID_BACKENDS = ('none', 'torch', 'cupy')
_VALID_PRECISIONS = ('float32', 'float64')




@dataclass(frozen=True)
class ComputeConfig:
    """Global runtime configuration for the O(N³) linear-algebra operations
    in :class:`~strongcoca.calculators.PolarizabilityCalculator`.

    Parameters
    ----------
    backend
        ``'none'`` (CPU only), ``'torch'`` (PyTorch/CUDA), or ``'cupy'`` (CuPy/CUDA).
    precision
        Floating-point precision for GPU operations. ``'float32'`` casts to
        complex64 before the GPU kernel and back to complex128 afterwards
        (roughly 10x faster on consumer GPUs where FP64 throughput is
        restricted); ``'float64'`` keeps complex128 throughout.
    max_solve_mem
        Memory budget in MiB for intermediate matrices in the CPU linear
        solve step. The frequency axis is chunked so that the A matrix
        stays below this limit.
    gpu_batched_max
        Maximum value of N3 = 3 x N_particles for which the cuBLAS batched
        solver is used in GPU kernels. Above this threshold, the code loops
        over the batch dimension to invoke the non-batched cuSOLVER path,
        which is faster for large matrices on consumer GPUs. Increase this
        value (e.g. to 1024 or higher) on A100/H100/H200 cards, where the
        batched cuSOLVER path performs well at large N.
    """

    backend: str
    precision: str
    max_solve_mem: float
    gpu_batched_max: int



_config = ComputeConfig(backend='none', precision='float32',
                        max_solve_mem=80.0, gpu_batched_max=32)




def set_compute_config(backend: Optional[str] = None,
                       precision: Optional[str] = None,
                       max_solve_mem: Optional[float] = None,
                       gpu_batched_max: Optional[int] = None) -> None:
    """Update one or more fields of the global compute configuration.

    Any argument left as ``None`` keeps its current value.
    All arguments are validated before anything is written, so an invalid
    call leaves the configuration completely unchanged.

    Parameters
    ----------
    backend
        See :class:`ComputeConfig`.
    precision
        See :class:`ComputeConfig`.
    max_solve_mem
        See :class:`ComputeConfig`.
    gpu_batched_max
        See :class:`ComputeConfig`.

    Examples
    --------
        >>> from strongcoca import set_compute_config, get_compute_config
        >>> set_compute_config(max_solve_mem=200)
        >>> get_compute_config().max_solve_mem
        200.0
        >>> set_compute_config(max_solve_mem=80)  # reset to the default
    """
    global _config

    new_backend = _config.backend if backend is None else backend
    new_precision = _config.precision if precision is None else precision
    new_max_solve_mem = _config.max_solve_mem if max_solve_mem is None else float(max_solve_mem)
    new_gpu_batched_max = _config.gpu_batched_max
    if gpu_batched_max is not None:
        new_gpu_batched_max = int(gpu_batched_max)

    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 new_max_solve_mem <= 0:
        raise ValueError(f'max_solve_mem must be positive, got {new_max_solve_mem}')
    if new_gpu_batched_max <= 0:
        raise ValueError(f'gpu_batched_max must be positive, got {new_gpu_batched_max}')

    _config = ComputeConfig(backend=new_backend,
                            precision=new_precision,
                            max_solve_mem=new_max_solve_mem,
                            gpu_batched_max=new_gpu_batched_max)





def get_compute_config() -> ComputeConfig:
    """Return the current global compute configuration."""
    return _config