quantui 0.5.1__py3-none-any.whl

quantui/freq_ir_workers.py

@@ -0,0 +1,229 @@
+"""ProcessPoolExecutor workers for the IR-intensity displacement loop.
+
+The Frequency calculation's IR-intensity step requires ``6N`` SCFs over
+finite-difference geometries (one per Cartesian displacement of each atom,
++Δ and −Δ). The default path in :mod:`quantui.freq_calc` runs them
+serially with each SCF internally parallelized via BLAS + libcint OpenMP.
+
+When the user opts in via ``QUANTUI_FREQ_PARALLEL=1`` AND no GPU is
+available AND the host has ``>= 4`` cores AND the molecule has
+``>= 2`` atoms (i.e. ``>= 6`` displacements), the freq_calc driver hands
+this loop off to a ``ProcessPoolExecutor`` whose workers each call
+:func:`run_displaced_scf` on one displaced geometry. Each worker process
+re-imports PySCF, rebuilds the ``gto.Mole`` from the same atom string /
+basis / charge / spin as the parent, applies the displacement, and runs
+the SCF. The initial guess ``dm0`` is shared once per worker via a temp
+pickle file (the path is passed through ``initargs``) so we don't pay
+per-task IPC for a 100×100 matrix.
+
+The functions in this module are intentionally top-level (not nested in
+``freq_calc.py``) because ``ProcessPoolExecutor`` requires picklable
+references for both ``initializer`` and the task callable. Nested
+functions cannot be pickled.
+
+POSIX-first design note: on Linux/macOS the parent process has already
+imported NumPy + PySCF by the time we spawn workers. We use
+``multiprocessing.get_context("spawn")`` so each worker starts with a
+fresh Python interpreter, reads the BLAS-thread env vars BEFORE NumPy is
+imported, and therefore actually honors the configured thread budget.
+Without ``spawn``, on Linux the default ``fork`` would inherit the
+parent's NumPy thread pool and ignore any env-var changes the worker
+makes.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Dict
+
+# Process-global state, populated by :func:`init_worker` once per worker.
+# Kept as a module-level dict (not class state) so workers don't need to
+# import any container class to access it.
+_WORKER_STATE: Dict[str, Any] = {}
+
+
+def init_worker(
+    atom_str: str,
+    basis: str,
+    charge: int,
+    spin: int,
+    xc: str | None,
+    dm0_pickle_path: str,
+    omp_threads: int,
+) -> None:
+    """ProcessPoolExecutor worker initializer.
+
+    Runs once per worker process. **Sets BLAS-thread env vars BEFORE
+    importing NumPy** — this is the whole point of the ``spawn`` start
+    method: each worker reads the env vars on its fresh interpreter
+    startup, NOT on the parent's already-imported NumPy state. Then loads
+    the shared initial-guess density matrix from the parent's tempfile
+    into ``_WORKER_STATE`` so per-task IPC stays tiny.
+
+    Parameters
+    ----------
+    atom_str:
+        Pyscf-format atom string ("O 0 0 0; H 0.96 0 0; ..."). Used to
+        rebuild the Mole in the worker.
+    basis:
+        Basis set name (e.g. ``"STO-3G"``).
+    charge, spin:
+        Molecular charge and 2S (spin) for the Mole.
+    xc:
+        DFT functional name when running a KS calculation; ``None`` for
+        plain HF.
+    dm0_pickle_path:
+        Path to a tempfile containing the parent's converged density
+        matrix as a NumPy array, used as the SCF initial guess in every
+        displaced calculation. Read once here, then kept in
+        ``_WORKER_STATE`` for all subsequent task calls.
+    omp_threads:
+        BLAS thread budget for this worker. Set as ``OMP_NUM_THREADS`` /
+        ``MKL_NUM_THREADS`` / ``OPENBLAS_NUM_THREADS`` / ``PYSCF_NUM_THREADS``.
+    """
+    # Order matters: set env vars before any NumPy / PySCF import.
+    threads = str(int(omp_threads))
+    os.environ["OMP_NUM_THREADS"] = threads
+    os.environ["OPENBLAS_NUM_THREADS"] = threads
+    os.environ["MKL_NUM_THREADS"] = threads
+    os.environ["PYSCF_NUM_THREADS"] = threads
+
+    import pickle
+
+    with open(dm0_pickle_path, "rb") as fh:
+        dm0 = pickle.load(fh)
+
+    _WORKER_STATE.update(
+        atom_str=atom_str,
+        basis=basis,
+        charge=int(charge),
+        spin=int(spin),
+        xc=xc,
+        dm0=dm0,
+    )
+
+
+def run_displaced_scf(coords_bohr_flat) -> Any:
+    """Run one SCF at the displaced geometry; return the dipole as ndarray.
+
+    Called by :class:`concurrent.futures.ProcessPoolExecutor` once per
+    submitted displacement task. ``coords_bohr_flat`` is the displaced
+    geometry packed as a flat Python list (``[x0, y0, z0, x1, y1, z1, ...]``)
+    for cheap pickling — reshaped to ``(N_atoms, 3)`` inside the worker.
+
+    Uses ``_WORKER_STATE`` populated by :func:`init_worker` for the
+    invariant inputs (atom string, basis, etc.) + the shared initial-guess
+    density matrix.
+
+    Returns
+    -------
+    np.ndarray
+        Three-component dipole moment in Debye.
+
+    Notes
+    -----
+    Any exception raised here propagates to the parent via the
+    ``Future.result()`` call. The freq_calc driver catches such failures
+    and falls back to the serial loop so the user's calc still completes.
+    """
+    import numpy as np
+    from pyscf import dft, gto, scf
+
+    state = _WORKER_STATE
+    coords = np.asarray(coords_bohr_flat, dtype=float).reshape(-1, 3)
+
+    mol = gto.Mole()
+    mol.atom = state["atom_str"]
+    mol.basis = state["basis"]
+    mol.charge = state["charge"]
+    mol.spin = state["spin"]
+    mol.verbose = 0
+    mol.build()
+    mol.set_geom_(coords, unit="Bohr")
+
+    # M5 audit fix (2026-07-14): whether this displaced SCF needs an
+    # unrestricted (UHF/UKS) object is determined by the shared dm0's
+    # actual shape -- (2, nao, nao) for UHF/UKS/ROHF, (nao, nao) for
+    # RHF/RKS -- NOT by mol.spin == 0. Those two signals only agree when
+    # the user's method choice matches the molecule's natural spin state.
+    # They diverge when a user explicitly selects UHF for a closed-shell
+    # molecule (mol.spin == 0 but the parent mf, and therefore dm0, is
+    # still UHF-shaped): building RHF from mol.spin == 0 and then feeding
+    # it the UHF-shaped dm0 raises a shape-mismatch ValueError inside
+    # PySCF. Mirrors the serial-path fix in freq_calc.py.
+    dm0 = state.get("dm0")
+    dm0_is_unrestricted = dm0 is not None and np.asarray(dm0).ndim == 3
+
+    xc = state.get("xc")
+    if xc is not None:
+        mf = dft.UKS(mol) if dm0_is_unrestricted else dft.RKS(mol)
+        mf.xc = xc
+    else:
+        mf = scf.UHF(mol) if dm0_is_unrestricted else scf.RHF(mol)
+    mf.verbose = 0
+    mf.kernel(dm0=dm0)
+    return np.array(mf.dip_moment(verbose=0))
+
+
+def parallel_enabled_for_run(
+    cpu_count: int,
+    displacement_count: int,
+    gpu_available: bool,
+) -> bool:
+    """Decide whether the freq_calc IR loop should use the parallel path.
+
+    Centralised in this module so both the driver and the tests can
+    consult the same predicate. The current rules:
+
+    - **Opt-in**: ``QUANTUI_FREQ_PARALLEL`` env var must be truthy
+      (``"1"`` / ``"true"`` / ``"True"``). Shipping this off-by-default
+      while the parallel path matures.
+    - **No GPU**: if gpu4pyscf is doing the offload, each SCF is already
+      ~10× faster; running multiple in parallel would compete for one
+      GPU's VRAM and is not worth the complexity. Stay serial.
+    - **Cores threshold**: at least 4 cores. Below that, the BLAS
+      oversubscription tradeoff doesn't pay off.
+    - **Displacement threshold**: at least 6 (i.e. ``>= 2`` atoms). For a
+      diatomic the serial loop is 12 SCFs at most and parallel overhead
+      dominates.
+    """
+    if not _truthy(os.environ.get("QUANTUI_FREQ_PARALLEL", "")):
+        return False
+    if gpu_available:
+        return False
+    if cpu_count < 4:
+        return False
+    if displacement_count < 6:
+        return False
+    return True
+
+
+def pick_worker_count(cpu_count: int, displacement_count: int) -> int:
+    """Pick a worker count that balances parallelism vs BLAS oversubscription.
+
+    Heuristic: use half the available cores, capped by the number of
+    displacement tasks. This leaves room for each worker to have ``>= 2``
+    BLAS threads on common 4/8/16-core configurations:
+
+    - 4 cores, 18 displacements → 2 workers × 2 threads each.
+    - 8 cores, 60 displacements → 4 workers × 2 threads each.
+    - 16 cores, 60 displacements → 8 workers × 2 threads each.
+    """
+    half = max(1, cpu_count // 2)
+    return min(half, displacement_count)
+
+
+def threads_per_worker(cpu_count: int, n_workers: int) -> int:
+    """How many BLAS threads each worker process should get.
+
+    Floors to 1 to avoid setting ``OMP_NUM_THREADS=0`` (which BLAS
+    interprets as "use the runtime default" — defeating the budgeting).
+    """
+    if n_workers <= 0:
+        return 1
+    return max(1, cpu_count // n_workers)
+
+
+def _truthy(value: str) -> bool:
+    """Match the truthy convention used by ``QUANTUI_DISABLE_GPU`` etc."""
+    return str(value).strip().lower() in ("1", "true", "yes", "on")

quantui/gpu_offload.py

@@ -0,0 +1,278 @@
+"""GPU offload helpers.
+
+Wraps the runtime decision "should this SCF object be migrated to GPU?".
+Detection probes ``gpu4pyscf`` + ``cupy`` for a CUDA-capable device; if
+anything is missing or broken the helpers silently report "no GPU" so the
+caller falls back to CPU. This means GPU integration is safe to leave
+enabled by default on every platform — Windows users without CUDA, WSL
+users without gpu4pyscf installed, and remote machines with broken NVIDIA
+drivers all converge to the same "CPU" outcome with no exception leakage.
+
+The companion ``log_utils._detect_gpu`` reports system-level GPU info for
+the run banner (nvidia-smi name + memory). This module's job is narrower:
+"can QuantUI's PySCF dispatcher offload to that GPU right now?".
+
+Method coverage (verified against the gpu4pyscf README 2026-05):
+
+- RHF / UHF / RKS / UKS — fully supported, ``mf.to_gpu()`` is canonical.
+- MP2, CCSD — listed as experimental; ``.to_gpu()`` may succeed but the
+  post-HF kernel may still fall back to CPU. ``try_to_gpu`` honours the
+  user's intent (offload the SCF; let gpu4pyscf decide the rest).
+- CCSD(T), double hybrids — explicitly not supported. ``try_to_gpu`` skips
+  GPU for these methods so the SCF + (T) step stays on CPU.
+
+User opt-out, two ways:
+
+- ``QUANTUI_DISABLE_GPU=1`` in the environment — process-wide, wins over
+  everything. Useful for benchmarks, regression debugging, and "first run
+  as student" comparisons, and it propagates to subprocess workers.
+- ``compute.gpu_enabled = false`` in user settings — the persistent UI
+  toggle (Status tab → Settings). Survives restarts.
+
+Not every CUDA device is worth using: double-precision throughput on
+consumer cards is a small fraction of their single-precision, and PySCF is
+FP64 throughout. See ``is_low_fp64_device`` — QuantUI reports an advisory
+but never overrides the user's choice.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from functools import lru_cache
+from typing import Any, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+# Methods for which gpu4pyscf has zero or known-broken support.
+#
+# - ``CCSD(T)`` is documented as unsupported in the gpu4pyscf README.
+# - ``MP2`` and ``CCSD`` are labelled "experimental" by gpu4pyscf and
+#   were observed (2026-05-25) to fail
+#   immediately after a successful RHF reference on GPU — the failure
+#   fingerprint was "step completed in RHF wall time + small delta,
+#   then errored", which fits the post-HF code choking on a
+#   GPU-migrated mf object. Until the upstream support matures, route
+#   these through CPU so calibration data accrues reliably. The RHF
+#   reference still benefits from GPU because ``try_to_gpu`` only
+#   short-circuits BEFORE the migration.
+# - Double-hybrids would belong here too, but QuantUI doesn't expose
+#   any double-hybrid methods today.
+_GPU_UNSUPPORTED_METHODS: frozenset = frozenset({"MP2", "CCSD", "CCSD(T)"})
+
+# Datacenter GPU families with usable double-precision throughput (FP64 at
+# roughly 1/2 of FP32). Everything else — GeForce, RTX/Quadro workstation, the
+# inference-oriented T4/L4/L40/A10/A40 line — gates FP64 to about 1/32–1/64.
+#
+# This matters because PySCF/gpu4pyscf SCF is FP64 throughout, so a consumer
+# card can be genuinely *slower* than a many-core CPU. Measured 2026-07-29 on an
+# RTX 5060 Ti: 346 GFLOP/s FP64 versus a 20-core CPU's 792 (0.44×), with real
+# B3LYP single points landing at 0.44–0.91× of CPU wall time.
+#
+# Matching is by device-name substring. Anything unrecognised is reported as
+# low-FP64 **deliberately**: consumer hardware is the common case for students,
+# and a spurious advisory on a brand-new datacenter card costs far less than
+# silently halving someone's throughput. This only ever drives an advisory
+# string — it never changes whether GPU offload runs.
+_STRONG_FP64_MARKERS: tuple = (
+    "A100",
+    "A30",
+    "H100",
+    "H200",
+    "B100",
+    "B200",
+    "GB200",
+    "GH200",
+    "V100",
+    "P100",
+    "GV100",
+    "GP100",
+    "TITAN V",
+)
+
+# Reason strings for the "GPU not in use" cases, surfaced by `quantui gpu check`
+# and the Status tab. Kept here so the CLI and the UI can't drift from the
+# probe's actual logic (they previously re-derived it, which is how a broken
+# CUDA install came to be reported as "gpu4pyscf not installed").
+_REASON_OK = ""
+_REASON_ENV_DISABLED = "QUANTUI_DISABLE_GPU is set in the environment"
+_REASON_SETTINGS_DISABLED = (
+    "GPU offload is switched off in QuantUI settings "
+    "(Status tab → Settings → GPU offload)"
+)
+_REASON_NOT_INSTALLED = (
+    "gpu4pyscf is not installed — install the extra matching your driver's "
+    "CUDA version, e.g. pip install 'quantui[gpu-cuda13x]' "
+    "(see README → 'Optional: GPU acceleration')"
+)
+_REASON_NO_DEVICE = "cupy reports 0 CUDA devices"
+
+
+def is_low_fp64_device(name: Optional[str]) -> bool:
+    """Return True when *name* is a GPU with crippled double-precision.
+
+    See ``_STRONG_FP64_MARKERS`` for the rationale, the measured numbers, and
+    why an unknown device is treated as low-FP64 rather than assumed fast.
+    """
+    if not name:
+        return False
+    upper = name.upper()
+    return not any(marker in upper for marker in _STRONG_FP64_MARKERS)
+
+
+def _gpu_enabled_in_settings() -> bool:
+    """Read the persistent ``compute.gpu_enabled`` preference.
+
+    Any failure reading settings returns True — a broken or unreadable
+    settings file must never be what silently disables the GPU.
+    """
+    try:
+        from quantui.user_settings import UserSettings
+
+        return bool(UserSettings.load().compute.gpu_enabled)
+    except Exception as exc:  # noqa: BLE001 — settings must never gate compute
+        logger.debug("could not read compute.gpu_enabled, assuming True: %s", exc)
+        return True
+
+
+@lru_cache(maxsize=1)
+def _probe_gpu() -> Tuple[bool, Optional[str], str]:
+    """Probe for a usable GPU. Returns ``(available, name, reason)``.
+
+    ``reason`` is empty when available and otherwise a human-readable
+    explanation suitable for showing a user directly.
+
+    The check sequence:
+
+    1. ``QUANTUI_DISABLE_GPU=1`` → disabled.
+    2. ``compute.gpu_enabled = false`` in user settings → disabled.
+    3. ``import gpu4pyscf`` — distinguishing "package absent"
+       (``ModuleNotFoundError``) from "present but its import chain is broken"
+       (any other ``ImportError``, typically missing CUDA math libraries).
+       These are very different problems and must not share a message.
+    4. ``import cupy`` + ``cupy.cuda.runtime.getDeviceCount()``.
+    5. Read device 0's properties for a friendly name.
+
+    Failures at any step are swallowed and logged; this never raises.
+    """
+    if os.environ.get("QUANTUI_DISABLE_GPU", "").strip() in ("1", "true", "True"):
+        return (False, None, _REASON_ENV_DISABLED)
+    if not _gpu_enabled_in_settings():
+        return (False, None, _REASON_SETTINGS_DISABLED)
+    try:
+        import gpu4pyscf  # noqa: F401
+    except ModuleNotFoundError:
+        # The package genuinely isn't installed — the common "user didn't opt
+        # into the extra" path. Not an error worth logging at warning level.
+        logger.debug("gpu4pyscf is not installed")
+        return (False, None, _REASON_NOT_INSTALLED)
+    except ImportError as exc:
+        # gpu4pyscf IS installed but its import chain is broken — nearly always
+        # missing NVIDIA CUDA math libraries (e.g. libnvJitLink.so,
+        # libcublasLt.so), which the gpu4pyscf wheels do not declare as
+        # dependencies. Surfacing the real exception is the whole point: the old
+        # code reported this as "not installed" and sent users back to an
+        # install step they had already completed.
+        logger.warning("gpu4pyscf is installed but failed to import: %s", exc)
+        return (
+            False,
+            None,
+            f"gpu4pyscf is installed but failed to import: {exc}. This usually "
+            "means its CUDA libraries are missing — see README → 'Optional: "
+            "GPU acceleration'",
+        )
+    except Exception as exc:  # noqa: BLE001 — any other import-chain breakage
+        logger.warning("gpu4pyscf import raised %s: %s", type(exc).__name__, exc)
+        return (
+            False,
+            None,
+            f"gpu4pyscf import raised {type(exc).__name__}: {exc}",
+        )
+
+    try:
+        import cupy as _cupy
+
+        n = int(_cupy.cuda.runtime.getDeviceCount())
+        if n < 1:
+            logger.debug("cupy reports 0 CUDA devices")
+            return (False, None, _REASON_NO_DEVICE)
+        props = _cupy.cuda.runtime.getDeviceProperties(0)
+        name_raw = props.get("name", b"GPU")
+        if isinstance(name_raw, bytes):
+            name = name_raw.decode("utf-8", errors="replace")
+        else:
+            name = str(name_raw)
+        return (True, name, _REASON_OK)
+    except Exception as exc:  # noqa: BLE001 — fall-back to CPU on probe failure
+        logger.warning("cupy device probe failed: %s", exc)
+        return (False, None, f"cupy device probe failed: {exc}")
+
+
+def probe_gpu() -> Tuple[bool, Optional[str], str]:
+    """Return ``(available, gpu_name, reason)`` — the full probe result.
+
+    Prefer this over :func:`is_gpu_available` when you need to tell the user
+    *why* the GPU isn't being used.
+    """
+    return _probe_gpu()
+
+
+def is_gpu_available() -> Tuple[bool, Optional[str]]:
+    """Return ``(available, gpu_name)`` for the current process.
+
+    Thin view over :func:`probe_gpu` that drops the reason, kept because it is
+    the long-standing call used by the run dispatcher and the Status tab.
+    Cached for the process lifetime — the answer doesn't change once the kernel
+    is up. Callers that need a re-check (a settings toggle, or a test
+    simulating driver loss) call ``is_gpu_available.cache_clear()``.
+    """
+    available, name, _reason = _probe_gpu()
+    return (available, name)
+
+
+# The cache lives on ``_probe_gpu`` now that two accessors share it. Forward the
+# lru_cache surface so the historical ``is_gpu_available.cache_clear()`` /
+# ``.cache_info()`` keep working — without this, existing callers would silently
+# clear nothing and read a stale result.
+is_gpu_available.cache_clear = _probe_gpu.cache_clear  # type: ignore[attr-defined]
+is_gpu_available.cache_info = _probe_gpu.cache_info  # type: ignore[attr-defined]
+
+
+def try_to_gpu(mf: Any, method_upper: str) -> Tuple[Any, bool, Optional[str]]:
+    """Attempt to migrate a PySCF SCF object to GPU. Safe CPU fallback.
+
+    Parameters
+    ----------
+    mf:
+        A constructed PySCF mean-field object (``scf.RHF(mol)``,
+        ``dft.RKS(mol)``, …) BEFORE ``mf.kernel()`` is called. ``to_gpu``
+        on a converged object is undefined behaviour in current gpu4pyscf.
+    method_upper:
+        Upper-cased method name (e.g. ``"RHF"``, ``"B3LYP"``, ``"CCSD(T)"``).
+        Used only to skip GPU for methods that gpu4pyscf doesn't support.
+
+    Returns
+    -------
+    tuple ``(maybe_gpu_mf, used_gpu, gpu_name)``:
+        - ``maybe_gpu_mf`` is the (possibly converted) SCF object the
+          caller should use for ``.kernel()``. Always usable — the
+          original ``mf`` is returned unchanged on any failure.
+        - ``used_gpu`` is ``True`` only when conversion succeeded.
+        - ``gpu_name`` is the device name when ``used_gpu`` is True,
+          ``None`` otherwise.
+    """
+    if method_upper in _GPU_UNSUPPORTED_METHODS:
+        return (mf, False, None)
+    available, gpu_name = is_gpu_available()
+    if not available:
+        return (mf, False, None)
+    try:
+        mf_gpu = mf.to_gpu()
+        return (mf_gpu, True, gpu_name)
+    except Exception as exc:
+        # gpu4pyscf migration can fail for many reasons (unsupported method
+        # variant, density-fitting requirement, basis-set quirk). On any
+        # failure we fall back to CPU — the calc still runs. Log so the
+        # user can `quantui log tail` and see why offload didn't happen.
+        logger.warning("mf.to_gpu() migration failed, falling back to CPU: %s", exc)
+        return (mf, False, None)