"""Pre-flight memory estimator + budget enforcement.
Goal
----
Every compute-heavy vibe-qc driver can estimate its peak memory
*before* starting the SCF, compare against the machine's available
RAM, and abort with a helpful message if the calculation would
otherwise thrash-to-disk or crash the system. The user can opt in to
running anyway by setting ``memory_override=True`` on the relevant
options struct (or by calling the explicit ``check_memory`` API with
``allow_exceed=True``).
Public API
----------
.. autoclass:: MemoryEstimate
.. autoclass:: InsufficientMemoryError
.. autofunction:: estimate_memory
.. autofunction:: estimate_neb_memory
.. autofunction:: check_memory
.. autofunction:: available_memory_bytes
The estimators are deliberately conservative -- they return a peak
upper bound, not a measured footprint. The output feeds two things:
1. A one-line summary printed at the top of every ``run_job`` output
(the "vibe-qc estimates this calculation will require X GB" line).
2. The pre-flight abort -- ``run_job`` calls ``check_memory`` with the
estimate before constructing the SCF driver.
Format of the estimate block (rendered by ``MemoryEstimate.format``):
vibe-qc estimates this calculation will require ~12.4 GB of memory:
ERI tensor 11.6 GB
Fock + density 0.4 GB
AO evaluation 0.3 GB
DIIS history 0.1 GB
scratch 0.0 GB
Available on this machine: 119.8 GB. Proceeding.
"""
from __future__ import annotations
import os
import sys
from dataclasses import dataclass, field
from math import comb
from typing import TYPE_CHECKING, Literal, Optional
from ._vibeqc_core import BasisSet, Functional, Molecule
if TYPE_CHECKING:
from .semiempirical.routes import SemiempiricalRoutePlan
__all__ = [
"MemoryEstimate",
"InsufficientMemoryError",
"estimate_memory",
"estimate_neb_memory",
"check_memory",
"available_memory_bytes",
"format_memory_report",
"estimate_semiempirical_memory",
"estimate_periodic_multik_gdf",
"estimate_periodic_xc_gradient",
"estimate_periodic_gpw_gapw",
"check_periodic_gdf_memory",
]
# ----------------------------------------------------------------------
# Raw probe of the machine
# ----------------------------------------------------------------------
[docs]
def available_memory_bytes() -> int:
"""Best-effort available-RAM probe. Prefers ``psutil`` when
installed (cross-platform, accurate) and falls back to platform-
specific mechanisms otherwise.
Returns ``0`` only when no probe succeeds, in which case the
caller should treat the result as "unknown" rather than "no
memory".
"""
# Preferred: psutil if present (optional dep).
try:
import psutil # type: ignore[import-not-found]
return int(psutil.virtual_memory().available)
except ImportError:
pass
# Linux: /proc/meminfo has MemAvailable since kernel 3.14.
if sys.platform.startswith("linux"):
try:
with open("/proc/meminfo", "r", encoding="ascii") as f:
for line in f:
if line.startswith("MemAvailable:"):
kb = int(line.split()[1])
return kb * 1024
except OSError:
pass
# macOS: sysconf gives total pages; vm_stat would give free/inactive
# but is awkward to parse from Python. Use total as a coarse proxy --
# better than nothing, and the abort logic is meant to catch orders-
# of-magnitude-over-budget runs, not to fine-tune.
try:
if hasattr(os, "sysconf"):
page = os.sysconf("SC_PAGE_SIZE")
n_pages = os.sysconf("SC_PHYS_PAGES")
if page > 0 and n_pages > 0:
return int(page) * int(n_pages)
except (ValueError, OSError):
pass
return 0
# ----------------------------------------------------------------------
# MemoryEstimate -- carrier + formatter
# ----------------------------------------------------------------------
_GB = 1024**3
_DEFAULT_MEMORY_HEADROOM = 1.5
_PYTHON_RUNTIME_FLOOR_BYTES = 100 * 1024**2
_CCSD_NATIVE_RUNTIME_FLOOR_BYTES = 256 * 1024**2
_CCSD_VVVV_INCORE_BUDGET_BYTES = 2 * 1024**3
_CCSD_VVVV_TILE_ROWS = 2048
def _memory_headroom_factor() -> float:
"""Safety factor applied to raw estimates.
HPC jobs routinely carry OS, Python, libint, MPI, filesystem, and allocator
overhead beyond the arrays we can count directly. Keep the default
conservative and let site wrappers tune it through an environment variable.
"""
raw = os.environ.get("VIBEQC_MEMORY_HEADROOM")
if raw is None or raw.strip() == "":
return _DEFAULT_MEMORY_HEADROOM
try:
value = float(raw)
except ValueError as exc:
raise ValueError("VIBEQC_MEMORY_HEADROOM must be a float >= 1.0") from exc
if value < 1.0:
raise ValueError("VIBEQC_MEMORY_HEADROOM must be >= 1.0")
return value
def _add_python_runtime_floor(by_category: dict[str, int]) -> None:
by_category["Python runtime + NumPy overhead"] = max(
int(by_category.get("Python runtime + NumPy overhead", 0)),
_PYTHON_RUNTIME_FLOOR_BYTES,
)
def _memory_estimate(by_category: dict[str, int]) -> "MemoryEstimate":
if by_category:
_add_python_runtime_floor(by_category)
return MemoryEstimate(by_category=by_category)
[docs]
@dataclass
class MemoryEstimate:
"""Peak memory estimate for a calculation.
All byte counts are integers (pre-headroom). ``total_bytes``
multiplies the sum of ``by_category`` by ``headroom_factor`` so
the headline figure already carries a safety margin.
"""
by_category: dict[str, int] = field(default_factory=dict)
headroom_factor: float = field(default_factory=_memory_headroom_factor)
@property
def raw_total_bytes(self) -> int:
return sum(self.by_category.values())
@property
def total_bytes(self) -> int:
return int(self.raw_total_bytes * self.headroom_factor)
@property
def total_gb(self) -> float:
return self.total_bytes / _GB
def __str__(self) -> str:
return self.format()
# ----------------------------------------------------------------------
# Enforcement
# ----------------------------------------------------------------------
[docs]
class InsufficientMemoryError(MemoryError):
"""Raised by ``check_memory`` when the estimate exceeds available
RAM and the caller did not request an override."""
[docs]
def check_memory(
estimate: MemoryEstimate,
*,
allow_exceed: bool = False,
available: Optional[int] = None,
) -> None:
"""Abort if ``estimate.total_bytes > available`` and
``allow_exceed`` is false.
``available`` may be passed in for reproducible tests; leave
``None`` to probe the live machine.
"""
avail = available if available is not None else available_memory_bytes()
if avail <= 0:
# No probe worked -- we can't enforce; the caller's risk.
return
if estimate.total_bytes <= avail:
return
if allow_exceed:
return
msg = estimate.format(available=avail, status="ABORTING") + "\n\n"
msg += (
"InsufficientMemoryError: Set ``memory_override=True`` (on the "
"options or ``run_job``) to proceed anyway. To shrink the estimate: "
"use a smaller basis; enable density fitting (``density_fit=True`` + "
"an ``aux_basis``); or the integral-direct SCF "
"(``scf_mode='direct'``), which avoids the in-core 4-index ERI "
"tensor. (Post-HF / DLPNO methods inherit this from their SCF "
"reference options.)"
)
raise InsufficientMemoryError(msg)
def estimate_neb_memory(
per_image_estimate: MemoryEstimate | int,
*,
n_images: int,
n_jobs: int,
n_atoms: int = 0,
n_basis: int = 0,
open_shell: bool = False,
finite_difference_evaluations: int = 1,
warm_start: bool = True,
) -> MemoryEstimate:
"""Peak-memory estimate for one NEB outer iteration.
``n_images`` is the number of intermediate images. ``per_image_estimate``
is the raw peak for a single image evaluator (usually the molecular
:func:`estimate_memory` result for the SCF method). The estimate charges
the image workers that can be live at once, the coordinate/force band
arrays, density warm-start storage retained across outer iterations, and
the additional finite-difference gradient scratch used by periodic NEB.
"""
image_count = max(1, int(n_images))
if isinstance(per_image_estimate, MemoryEstimate):
per_image_bytes = int(per_image_estimate.raw_total_bytes)
headroom = float(per_image_estimate.headroom_factor)
else:
per_image_bytes = int(per_image_estimate)
headroom = 1.2
per_image_bytes = max(0, per_image_bytes)
requested_jobs = int(n_jobs)
if requested_jobs < 0:
worker_count = min(image_count, max(1, os.cpu_count() or 1))
elif requested_jobs == 0:
worker_count = min(image_count, max(1, os.cpu_count() or 1))
else:
worker_count = min(image_count, requested_jobs)
worker_count = max(1, worker_count)
atoms = max(0, int(n_atoms))
basis_size = max(0, int(n_basis))
spin_channels = 2 if open_shell else 1
fd_evaluations = max(1, int(finite_difference_evaluations))
by_cat: dict[str, int] = {
"NEB parallel image workers": worker_count * per_image_bytes,
"NEB band coordinates/forces": (
(image_count + 2) * max(1, atoms) * 3 * 8 * 6
),
}
if warm_start and basis_size > 0:
by_cat["NEB warm-start density cache"] = (
(image_count + 2) * spin_channels * basis_size * basis_size * 8
)
elif warm_start:
by_cat["NEB warm-start density cache"] = 0
if fd_evaluations > 1:
matrix_bytes = spin_channels * basis_size * basis_size * 8
geometry_bytes = max(1, atoms) * 3 * 8
by_cat["NEB finite-difference gradient scratch"] = (
worker_count * fd_evaluations * (matrix_bytes + geometry_bytes)
)
return MemoryEstimate(by_category=by_cat, headroom_factor=headroom)
def _semiempirical_atoms(system) -> list:
atoms = getattr(system, "atoms", None)
if atoms is None:
atoms = getattr(system, "unit_cell", None)
return list(atoms or [])
def _semiempirical_orbital_count(z: int, method_family: str) -> int:
if z <= 2:
return 1
if method_family == "msindo" and z > 10:
return 9
return 4
def _semiempirical_route_plan(
system,
method: str | SemiempiricalRoutePlan,
*,
nddo: bool,
solvent: str | None,
ccm_options,
) -> SemiempiricalRoutePlan:
from vibeqc.semiempirical.routes import (
BOUNDARY_MOLECULE,
BOUNDARY_PERIODIC_GAMMA,
BOUNDARY_SECCM_DIRECT_TORUS,
SemiempiricalRoutePlan,
plan_periodic_semiempirical_route,
)
is_periodic = hasattr(system, "unit_cell") and hasattr(system, "lattice")
if is_periodic and not nddo and solvent is None and ccm_options is None:
return plan_periodic_semiempirical_route(method, system)
if isinstance(method, SemiempiricalRoutePlan):
plan = method
expected_boundaries = (
{BOUNDARY_PERIODIC_GAMMA}
if is_periodic
else {BOUNDARY_MOLECULE, BOUNDARY_SECCM_DIRECT_TORUS}
)
if plan.boundary not in expected_boundaries:
expected = ", ".join(sorted(expected_boundaries))
raise ValueError(
"semiempirical memory route/system boundary mismatch: "
f"plan has boundary={plan.boundary!r}, expected {expected}."
)
if nddo and plan.variant != "nddo":
raise ValueError("nddo=True does not match the supplied route plan.")
if solvent is not None and plan.variant != "cosmo":
raise ValueError("solvent does not match the supplied route plan.")
if ccm_options is not None and plan.boundary != BOUNDARY_SECCM_DIRECT_TORUS:
raise ValueError("ccm_options does not match the supplied route plan.")
return plan
boundary = BOUNDARY_PERIODIC_GAMMA if is_periodic else None
return SemiempiricalRoutePlan.from_request(
method,
boundary=boundary,
charge=int(getattr(system, "charge", 0) or 0),
multiplicity=int(getattr(system, "multiplicity", 1) or 1),
nddo=nddo,
solvent=solvent,
ccm_options=ccm_options,
)
def estimate_semiempirical_memory(
system,
*,
method: str | SemiempiricalRoutePlan,
nddo: bool = False,
solvent: str | None = None,
optimize: bool = False,
max_steps: int = 0,
ccm_options=None,
) -> MemoryEstimate:
"""Peak-memory estimate for basis-free semiempirical routes.
The estimate validates or consumes a :class:`SemiempiricalRoutePlan`, then
counts valence-sized dense matrices and route-specific scratch without
constructing a Gaussian :class:`BasisSet`. It is intentionally
conservative and cheap enough for dry-run placement.
"""
from vibeqc.semiempirical.routes import (
BOUNDARY_PERIODIC_GAMMA,
BOUNDARY_SECCM_DIRECT_TORUS,
)
plan = _semiempirical_route_plan(
system,
method,
nddo=nddo,
solvent=solvent,
ccm_options=ccm_options,
)
is_periodic = plan.boundary == BOUNDARY_PERIODIC_GAMMA
periodic_dim = max(0, int(getattr(system, "dim", 0) or 0)) if is_periodic else 0
atoms = _semiempirical_atoms(system)
n_atoms = max(1, len(atoms))
z_values = [
int(getattr(atom, "Z", getattr(atom, "number", 0)))
for atom in atoms
]
n_orb = max(
1,
sum(
_semiempirical_orbital_count(z, plan.method_family)
for z in z_values
),
)
spin_channels = 2 if plan.spin == "unrestricted" else 1
matrix_bytes = n_orb * n_orb * 8
vector_bytes = n_orb * 8
coord_bytes = n_atoms * 3 * 8
by_cat: dict[str, int] = {
"Semiempirical valence matrices": (8 + 3 * spin_channels) * matrix_bytes,
"Semiempirical SCF eigensolver scratch": (
6 * matrix_bytes + 4 * vector_bytes
),
"Semiempirical DIIS history": 8 * spin_channels * 2 * matrix_bytes,
"Semiempirical atom-pair scratch": max(1, n_atoms * n_atoms) * 128,
}
if plan.scc in {"atomic", "shell"}:
by_cat["Semiempirical charge/SCC state"] = (
12 * n_atoms * 8 + 4 * matrix_bytes
)
if plan.method_family == "gfn2":
by_cat["GFN2 shell/AES multipole state"] = (
24 * n_atoms * 8 + 6 * matrix_bytes
)
if is_periodic:
image_count = max(1, 3 ** max(1, periodic_dim))
by_cat["Periodic semiempirical image/neighbour buffers"] = (
image_count * (n_atoms * (3 * 8 + 16) + (2 + spin_channels) * matrix_bytes)
)
fd_route = plan.variant in {"pm6", "upm6", "om1", "om2", "om3"} or optimize
if fd_route:
by_cat["Semiempirical finite-difference gradient scratch"] = (
2 * (coord_bytes + (4 + spin_channels) * matrix_bytes)
+ n_atoms * 3 * 8
)
if is_periodic and optimize:
strain_components = max(1, periodic_dim * periodic_dim)
by_cat["Periodic semiempirical stress/strain scratch"] = (
2 * strain_components * (coord_bytes + (4 + spin_channels) * matrix_bytes)
+ 3 * 3 * 8
)
if optimize:
steps = max(1, int(max_steps) if max_steps is not None else 1)
by_cat["Semiempirical optimizer coordinates/history"] = (
min(steps, 50) * n_atoms * 3 * 8 * 4
)
if plan.variant == "cosmo":
n_segments = max(24, 60 * n_atoms)
by_cat["MSINDO COSMO cavity A matrix"] = n_segments * n_segments * 8
by_cat["MSINDO COSMO B matrix"] = n_segments * n_orb * n_orb * 8
by_cat["MSINDO COSMO surface vectors"] = 8 * n_segments * 8
if plan.boundary == BOUNDARY_SECCM_DIRECT_TORUS:
translations = (
getattr(ccm_options, "translations", None) if ccm_options else None
)
dim = len(translations or [])
image_count = max(1, 3 ** max(1, dim))
by_cat["MSINDO SECCM image-cell buffers"] = (
image_count * (n_atoms * (3 * 8 + 16) + 2 * matrix_bytes)
)
if bool(getattr(ccm_options, "madelung", False)):
by_cat["MSINDO SECCM Madelung workspace"] = image_count * matrix_bytes
if plan.variant == "nddo":
by_cat["MSINDO NDDO parameter scratch"] = 2 * matrix_bytes
return _memory_estimate(by_cat)
def estimate_periodic_multik_gdf(
*,
n_basis: int,
n_aux: int,
n_kpoints: int,
need_k_pairs: bool,
open_shell: bool = False,
) -> MemoryEstimate:
"""Peak memory for a dense multi-k periodic GDF SCF.
The dominant allocation is the in-core density-fitting factor cache:
one ``complex128`` ``(n_aux, n_basis, n_basis)`` ``Lpq`` block per stored
``(k_i, k_j)`` pair. HF / hybrid exchange needs every ``n_kpoints**2``
pair; pure DFT (and J-only) keeps just the ``n_kpoints`` diagonal pairs.
The cache is held dense in RAM -- streaming is not yet implemented (see
the :mod:`vibeqc.periodic_k_gdf` module docstring) -- so for a
paper-grade cell this term is what OOM-kills the run before SCF iter 1
(the NiO/def2-SVP KUKS ``(4,4,4)`` exit-137 case).
Parameters are plain integers so the estimate needs no built core and is
unit-testable in isolation. ``complex128`` is 16 bytes/element.
"""
n = int(n_basis)
n_pairs = int(n_kpoints) * int(n_kpoints) if need_k_pairs else int(n_kpoints)
by_cat: dict[str, int] = {}
by_cat["GDF Lpq factor cache"] = n_pairs * int(n_aux) * n * n * 16
# Per-k complex one-electron + SCF-loop buffers (S, Hcore, X, MO, and per
# spin channel a Fock + density). Sub-dominant to the Lpq cache but real.
spin = 2 if open_shell else 1
per_k_matrices = 4 + 2 * spin
by_cat["per-k complex buffers"] = per_k_matrices * int(n_kpoints) * n * n * 16
return _memory_estimate(by_cat)
def estimate_periodic_xc_gradient(
*,
n_basis: int,
n_atoms: int,
n_grid_points: int,
n_cells: int,
functional_kind: str,
open_shell: bool = False,
n_threads: int | None = None,
) -> MemoryEstimate:
"""Peak memory for the current periodic analytic XC-gradient kernel.
The C++ ``xc_lattice_gradient_contribution`` routes are intentionally
unbatched today: they keep home-cell and shifted-cell AO values/gradients
across the full periodic grid, add AO Hessian tables for GGA/MGGA sigma and
tau Pulay terms, and allocate per-thread pair scratch in the cross-cell
bra/ket loop. The estimate is parameter-only so dry-run preflight can place
jobs without running SCF or constructing lattice-integral objects.
"""
n = max(0, int(n_basis))
n_atoms = max(0, int(n_atoms))
n_pts = max(0, int(n_grid_points))
n_cells = max(1, int(n_cells))
if n == 0 or n_pts == 0:
return MemoryEstimate(by_category={})
kind = str(functional_kind or "").strip().upper()
is_mgga = "MGGA" in kind or "META" in kind
# Unknown semilocal kinds are charged as GGA rather than undercounted.
need_hess = is_mgga or kind not in ("", "LDA")
matrix_bytes = n_pts * n * 8
vector_bytes = n_pts * 8
matrix_1e_bytes = n * n * 8
threads = _periodic_xc_gradient_thread_count(n_threads, n_cells)
by_cat: dict[str, int] = {}
# Home-cell reference tables plus one entry per density cell in chi_h /
# dchi_h / hess_h. The home-cell entry is copied into the cell vector, so
# peak memory includes both the reference object and the vector slot.
ao_table_count = 10 if need_hess else 4
by_cat["Periodic-XC gradient AO/gradient/Hessian tables"] = (
(n_cells + 1) * ao_table_count * matrix_bytes
)
# Density blocks are already present when the gradient is called, but they
# contribute to the same peak. UKS carries alpha and beta cell matrices.
spin_channels = 2 if open_shell else 1
by_cat["Periodic-XC gradient density blocks"] = (
spin_channels * n_cells * matrix_1e_bytes
)
if open_shell:
scratch_matrices = 16 if need_hess else 4
else:
scratch_matrices = 8 if need_hess else 2
by_cat["Periodic-XC gradient pair scratch"] = (
threads * scratch_matrices * matrix_bytes
)
if open_shell:
if is_mgga:
grid_vectors = 31
elif need_hess:
grid_vectors = 25
else:
grid_vectors = 7
else:
if is_mgga:
grid_vectors = 14
elif need_hess:
grid_vectors = 11
else:
grid_vectors = 4
by_cat["Periodic-XC gradient grid vectors"] = grid_vectors * vector_bytes
by_cat["Periodic-XC gradient thread buffers"] = (
threads * max(1, n_atoms) * 3 * 8 + n * 8
)
return _memory_estimate(by_cat)
_LEBEDEV_ORDER_POINTS = {
3: 6,
5: 14,
7: 26,
9: 38,
11: 50,
13: 74,
15: 86,
17: 110,
19: 146,
21: 170,
23: 194,
25: 230,
27: 266,
29: 302,
31: 350,
}
def estimate_periodic_gpw_gapw(
*,
n_basis: int,
n_grid_points: int,
route: str,
functional_kind: str | None = None,
open_shell: bool = False,
n_kpoints: int = 1,
compact_multik: bool = False,
n_soft_basis: int | None = None,
n_atoms: int = 0,
augmentation_active: bool | None = None,
lmax: int = 3,
n_radial: int = 80,
lebedev_order: int = 17,
) -> MemoryEstimate:
"""Peak memory for the FFT-grid GPW/GAPW periodic routes.
The estimate mirrors the large iteration-invariant caches in
``periodic_gapw_j`` / ``periodic_gapw_augment``: the full-basis AO
collocation table, reciprocal mesh, density/Poisson/XC grid workspaces,
optional compact multi-k Bloch AO tables, and GAPW's soft-basis plus
analytic augmentation caches. Inputs are plain integers so vq dry-runs can
place jobs without building a plane-wave grid or entering the SCF loop.
"""
route_key = str(route or "").strip().lower()
if route_key not in ("gpw", "gapw"):
raise ValueError("estimate_periodic_gpw_gapw: route must be 'gpw' or 'gapw'")
n = max(0, int(n_basis))
n_pts = max(0, int(n_grid_points))
if n == 0 or n_pts == 0:
return MemoryEstimate(by_category={})
n_k = max(1, int(n_kpoints))
n_soft = n if n_soft_basis is None else max(0, int(n_soft_basis))
n_atoms = max(0, int(n_atoms))
spin = 2 if open_shell else 1
grid_real = n_pts * 8
grid_complex = n_pts * 16
matrix_real = n * n * 8
matrix_complex = n * n * 16
by_cat: dict[str, int] = {}
by_cat["GPW full collocation cache"] = n_pts * n * 8
by_cat["GPW reciprocal mesh cache"] = n_pts * 3 * 8
# The smooth J path keeps a density grid, a Poisson-potential grid and FFT
# work arrays. UKS/compact routes can hold both spin densities locally.
by_cat["GPW density/Poisson grids"] = (
(spin + 3) * grid_real + 2 * grid_complex
)
by_cat["GPW contraction/projection scratch"] = (
spin * matrix_real + 2 * matrix_real + grid_real
)
kind = str(functional_kind or "").strip().upper()
has_xc = bool(kind)
is_mgga = "MGGA" in kind or "META" in kind
needs_gradient = has_xc and (is_mgga or kind != "LDA")
if has_xc:
if open_shell:
grid_vectors = 10 if not needs_gradient else 24
if is_mgga:
grid_vectors = 34
else:
grid_vectors = 6 if not needs_gradient else 15
if is_mgga:
grid_vectors = 22
by_cat["GPW XC grid workspace"] = grid_vectors * grid_real
if is_mgga:
by_cat["GPW meta-GGA AO-gradient scratch"] = (
spin * (3 * n_pts * n * 8 + n_pts * n * 16)
)
if compact_multik and n_k > 1:
by_cat["GPW compact multi-k Bloch AO tables"] = n_k * n_pts * n * 16
by_cat["GPW compact multi-k projection scratch"] = (
n_k * spin * matrix_complex + spin * grid_complex
)
if route_key == "gapw":
active = True if augmentation_active is None else bool(augmentation_active)
if active:
by_cat["GAPW soft collocation cache"] = n_pts * n_soft * 8
n_comp = max(1, int(lmax) + 1) ** 2
n_rad = max(1, int(n_radial))
n_ang = _LEBEDEV_ORDER_POINTS.get(int(lebedev_order), int(lebedev_order))
n_ang = max(1, n_ang)
n_atom_grid = n_rad * n_ang
soft_matrix_real = n_soft * n_soft * 8
by_cat["GAPW analytic augmentation cache"] = n_atoms * (
n_atom_grid * (n + n_soft) * 8
+ n_comp * matrix_real
+ n_comp * n_pts * 8
+ n_comp * n_atom_grid * 8
+ 6 * n_atom_grid * 8
+ soft_matrix_real
)
return _memory_estimate(by_cat)
def _periodic_xc_gradient_thread_count(
n_threads: int | None,
n_cells: int,
) -> int:
if n_threads is None:
try:
from ._vibeqc_core import get_num_threads
n_threads = int(get_num_threads())
except Exception:
n_threads = os.cpu_count() or 1
n_pairs_upper = max(1, int(n_cells) * int(n_cells))
return max(1, min(int(n_threads), n_pairs_upper))
def check_periodic_gdf_memory(
estimate: MemoryEstimate,
*,
n_kpoints: int,
route_label: str,
allow_exceed: bool = False,
available: Optional[int] = None,
) -> None:
"""Fail-early gate for the dense multi-k GDF ``Lpq`` cache.
Raises :class:`InsufficientMemoryError` -- with periodic-specific
remedies rather than the molecular ``check_memory`` text -- when the
estimate exceeds available RAM. A no-op when RAM cannot be probed
(``available`` is 0, treated as "unknown") or ``allow_exceed`` is set,
mirroring :func:`check_memory`'s fail-open-on-unknown contract.
"""
avail = available if available is not None else available_memory_bytes()
if avail <= 0 or estimate.total_bytes <= avail or allow_exceed:
return
msg = estimate.format(available=avail, status="ABORTING") + "\n\n"
msg += (
f"InsufficientMemoryError: the multi-k {route_label} GDF route holds "
f"the density-fitting Lpq factors dense in RAM for {int(n_kpoints)} "
"k-points (streaming is not yet implemented; see the periodic_k_gdf "
"module docstring). To fit the calculation: use a coarser "
"Monkhorst-Pack mesh or the Gamma-only GDF path; choose a smaller "
"orbital / auxiliary basis; or switch to the plane-wave GPW route "
"(jk_method='gpw'), which does not materialise per-k-pair DF factors. "
"Set VIBEQC_GDF_MEMORY_OVERRIDE=1 to attempt the run anyway."
)
raise InsufficientMemoryError(msg)
# ----------------------------------------------------------------------
# Per-driver estimators
# ----------------------------------------------------------------------
def _n_basis(molecule: Molecule, basis: BasisSet) -> int:
return int(basis.nbasis)
def _n_shells(basis: BasisSet) -> int:
value = getattr(basis, "nshells", None)
if value is not None:
return int(value)
try:
return int(len(basis.shells()))
except Exception:
return 0
def _diis_subspace_size(options) -> int:
return int(getattr(options, "diis_subspace_size", 8)) if options is not None else 8
def _enum_name(value) -> str:
"""Return a stable upper-case name for pybind11 enums and strings."""
if value is None:
return ""
name = getattr(value, "name", None)
if isinstance(name, str):
return name.upper()
text = str(value)
return text.rsplit(".", 1)[-1].upper()
def _uses_density_fit(options) -> bool:
return (
bool(getattr(options, "density_fit", False)) if options is not None else False
)
def _uses_cosx(options) -> bool:
return bool(getattr(options, "cosx", False)) if options is not None else False
def _uses_direct_scf(n_basis: int, options) -> bool:
"""Mirror the SCF-mode dispatch used by the molecular drivers.
``options is None`` means the driver builds a default ``RHFOptions`` /
``UHFOptions``, whose ``scf_mode`` default is **AUTO** (threshold 200 bf,
see cpp/include/vibeqc/rhf.hpp). So a large default-options SCF runs the
integral-direct path -- it does NOT materialise the in-core 4-index ERI.
Treat ``None`` / an unset ``scf_mode`` as that AUTO default rather than
assuming the dense ERI tensor; otherwise the estimator over-counts a
large default job by ``n_basis**4`` and spuriously aborts it (e.g. a
DLPNO job's RHF reference on n-octane/cc-pVTZ: ~436 GB phantom ERI
tensor, while the AUTO->direct SCF actually runs in a few GB)."""
if _uses_density_fit(options) or _uses_cosx(options):
return False
mode_name = (
_enum_name(getattr(options, "scf_mode", None)) if options is not None
else ""
)
if mode_name == "DIRECT":
return True
if mode_name == "AUTO" or not mode_name:
# AUTO is the driver default; an absent options object inherits it.
threshold = (
int(getattr(options, "scf_mode_auto_threshold", 200))
if options is not None else 200
)
return n_basis > threshold
return False # explicit CONVENTIONAL: the in-core 4-index path
def _option_value(options, key: str, default=None):
"""Read an estimator option from a dict or an options object."""
if options is None:
return default
if isinstance(options, dict):
return options.get(key, default)
return getattr(options, key, default)
def _scf_estimator_options(options):
"""Return the SCF-options object from an optional estimator bundle."""
if isinstance(options, dict) and "scf_options" in options:
return options.get("scf_options")
return options
def _looks_like_scf_options(options) -> bool:
if options is None or isinstance(options, dict):
return False
if type(options).__name__ in {
"RHFOptions",
"UHFOptions",
"RKSOptions",
"UKSOptions",
"ROHFOptions",
"ROKSOptions",
}:
return True
return any(
hasattr(options, attr)
for attr in ("scf_mode", "scf_mode_auto_threshold", "cosx")
)
def _post_reference_options(options):
"""Return SCF reference options from a post-SCF estimator bundle."""
if isinstance(options, dict):
return options.get("scf_options")
if _looks_like_scf_options(options):
return options
return None
def _post_method_options(options, key: str):
"""Return method-specific post-SCF options from a bundle or bare object."""
if isinstance(options, dict):
return options.get(key)
if _looks_like_scf_options(options):
return None
return options
def _posthf_uses_density_fit(method_options, *, default: bool = False) -> bool:
if method_options is None:
return bool(default)
return bool(getattr(method_options, "density_fit", default))
def _tddft_estimator_options(options) -> Optional[dict]:
"""Return TDDFT post-SCF options when the estimate should include them."""
if not isinstance(options, dict):
return None
if not bool(options.get("tddft", False)):
return None
return options
def _aux_basis_size(molecule: Molecule, n_basis: int, options) -> int:
"""Best-effort auxiliary basis size for DF/COSX estimates.
The estimator is allowed to be conservative, but it should still
reflect the requested algorithm. If an aux basis name is available,
use the same local resolver the drivers use; otherwise fall back to a
typical JKFIT/RIFIT scale rather than pretending a 4-index ERI tensor
will be materialised.
"""
aux_name = str(getattr(options, "aux_basis", "") or "")
if aux_name:
try:
from .aux_basis import make_aux_basis_set
return int(make_aux_basis_set(molecule, aux_name=aux_name).nbasis)
except Exception:
pass
return max(1, 3 * n_basis)
def _mean_field_matrix_bytes(n_basis: int, *, open_shell: bool = False) -> int:
factor = 12 if open_shell else 8
return factor * n_basis * n_basis * 8
def _add_jk_storage(
by_cat: dict[str, int],
molecule: Molecule,
basis: BasisSet,
n_basis: int,
options,
*,
open_shell: bool = False,
) -> None:
"""Add the dominant two-electron storage for the selected JK path."""
if _uses_cosx(options):
n_aux = _aux_basis_size(molecule, n_basis, options)
n_grid = _grid_points(molecule, getattr(options, "cosx_grid", None))
by_cat["RI-J tensors"] = n_aux * n_basis * n_basis * 8 + n_aux * n_aux * 8
by_cat["COSX grid workspace"] = max(1, n_grid) * n_basis * 8
elif _uses_density_fit(options):
n_aux = _aux_basis_size(molecule, n_basis, options)
by_cat["DF three-index tensors"] = n_aux * n_basis * n_basis * 8
by_cat["DF metric/workspace"] = 2 * n_aux * n_aux * 8
elif _uses_direct_scf(n_basis, options):
n_shell = max(0, _n_shells(basis))
shell_pairs = n_shell * (n_shell + 1) // 2
by_cat["Direct-SCF shell-pair scratch"] = (
shell_pairs * 1024
)
else:
by_cat["ERI tensor"] = n_basis**4 * 8
def _grid_dimensions(options) -> tuple[int, int, int]:
"""Rough upper bound on the DFT grid-point count: n_radial x n_theta
x n_phi x n_atoms. Falls back to vibe-qc's default grid dimensions
(75 x 17 x 36) when ``options`` is None or lacks a .grid field --
otherwise the estimator would under-report DFT memory on every
default-options call."""
if isinstance(options, dict) and "scf_options" in options:
options = options.get("scf_options")
n_radial, n_theta, n_phi = 75, 17, 36
if options is not None and hasattr(options, "n_radial"):
n_radial = int(getattr(options, "n_radial", n_radial))
n_theta = int(getattr(options, "n_theta", n_theta))
n_phi = int(getattr(options, "n_phi", n_phi))
elif options is not None and hasattr(options, "grid"):
g = options.grid
n_radial = int(getattr(g, "n_radial", n_radial))
n_theta = int(getattr(g, "n_theta", n_theta))
n_phi = int(getattr(g, "n_phi", n_phi))
return max(0, n_radial), max(0, n_theta), max(0, n_phi)
def _grid_points(molecule: Molecule, options) -> int:
n_radial, n_theta, n_phi = _grid_dimensions(options)
return n_radial * n_theta * n_phi * len(molecule.atoms)
def _functional_kind_name(options, *, open_shell: bool = False) -> str:
name = _option_value(options, "functional", None)
if name is None:
return "GGA"
try:
func = Functional(str(name), 2 if open_shell else 1)
except Exception:
return "GGA"
return _enum_name(getattr(func, "kind", None)) or "GGA"
def _functional_work_vectors(options, *, open_shell: bool = False) -> int:
"""Approximate libxc work-vector count per grid point."""
kind = _functional_kind_name(options, open_shell=open_shell)
spin = 2 if open_shell else 1
if "MGGA" in kind or "META" in kind:
return 28 if spin == 1 else 48
if kind == "LDA":
return 8 if spin == 1 else 14
# Unknown semilocal functionals are charged as GGA rather than LDA.
return 16 if spin == 1 else 28
def _rhf_estimate(molecule: Molecule, basis: BasisSet, options) -> MemoryEstimate:
scf_options = _scf_estimator_options(options)
n = _n_basis(molecule, basis)
by_cat: dict[str, int] = {}
_add_jk_storage(by_cat, molecule, basis, n, scf_options)
# One-electron matrices (S, T, V, Hcore, F, D, X=S^{-1/2}, scratch).
# ~8 * n^2 x 8 covers all core SCF-loop matrices.
by_cat["Fock + density + 1e"] = _mean_field_matrix_bytes(n)
# DIIS extrapolation holds (F, error) pairs across iterations.
diis = _diis_subspace_size(scf_options)
by_cat["DIIS history"] = diis * 2 * n * n * 8
# MO coefficient + eigenvalue arrays + transform buffers.
by_cat["MO workspace"] = 4 * n * n * 8 + n * 8
est = _memory_estimate(by_cat)
_add_tddft_response_estimate(est, molecule, basis, options)
return est
def _uhf_estimate(molecule: Molecule, basis: BasisSet, options) -> MemoryEstimate:
scf_options = _scf_estimator_options(options)
n = _n_basis(molecule, basis)
by_cat: dict[str, int] = {}
_add_jk_storage(by_cat, molecule, basis, n, scf_options, open_shell=True)
by_cat["Fock + density + 1e"] = _mean_field_matrix_bytes(n, open_shell=True)
base = _memory_estimate(by_cat)
diis = _diis_subspace_size(scf_options)
base.by_category["DIIS history"] = diis * 4 * n * n * 8
base.by_category["MO workspace"] = 8 * n * n * 8 + 2 * n * 8
# Open-shell adds D_alpha, D_beta, F_alpha, F_beta; rough 2x on
# the density and Fock buffers.
extra = 4 * n * n * 8
base.by_category["Open-shell UHF buffers"] = extra
_add_tddft_response_estimate(base, molecule, basis, options, open_shell=True)
return base
def _dft_xc_estimate(
molecule: Molecule,
basis: BasisSet,
options,
*,
open_shell: bool = False,
) -> int:
"""Extra memory beyond the HF baseline for XC integration.
Charge AO value/gradient tables plus the runtime scratch that is live with
them: Becke weights, angular-grid metadata, coordinates, libxc work arrays,
and meta-GGA tau buffers. The AO contraction scratch intentionally mirrors
the value/gradient table size because the current grid code keeps multiple
density/Fock contraction buffers live while libxc evaluates a batch.
"""
n = _n_basis(molecule, basis)
n_pts = _grid_points(molecule, options)
if n_pts == 0:
return 0
n_radial, n_theta, n_phi = _grid_dimensions(options)
n_angular = max(1, n_theta * n_phi)
kind = _functional_kind_name(options, open_shell=open_shell)
is_mgga = "MGGA" in kind or "META" in kind
spin = 2 if open_shell else 1
ao_value_gradient = 4 * n_pts * n * 8
ao_contraction_scratch = 4 * n_pts * n * 8
becke_weights = n_pts * 8
angular_points = 4 * max(1, n_radial) * n_angular * 8
coordinates = 3 * n_pts * 8
libxc_work = n_pts * _functional_work_vectors(options, open_shell=open_shell) * 8
tau = spin * n_pts * n * 8 if is_mgga else 0
return (
ao_value_gradient
+ ao_contraction_scratch
+ becke_weights
+ angular_points
+ coordinates
+ libxc_work
+ tau
)
def _clamped_occ(value: int, n_basis: int) -> int:
return max(0, min(int(n_basis), int(value)))
def _add_tddft_response_estimate(
est: MemoryEstimate,
molecule: Molecule,
basis: BasisSet,
options,
*,
open_shell: bool = False,
) -> None:
"""Add dense molecular TDDFT/TDA/CIS response-memory categories.
The current TDDFT implementation materializes the full AO ERI tensor,
transforms dense MO response blocks, builds dense A (and, for Casida, B)
matrices, then diagonalizes dense response matrices. This helper mirrors
those allocation shapes so preflight and vq RAM placement see the post-SCF
peak instead of only the underlying SCF reference.
"""
td_opts = _tddft_estimator_options(options)
if td_opts is None:
return
n = _n_basis(molecule, basis)
if n <= 0:
return
tddft_type = str(td_opts.get("tddft_type", "tda") or "tda").strip().lower()
is_casida = tddft_type == "casida"
n_states = max(1, int(td_opts.get("tddft_n_states", 5) or 5))
functional = td_opts.get("functional")
if open_shell:
n_elec = molecule.n_electrons()
ms2 = int(molecule.multiplicity) - 1
n_alpha = _clamped_occ((n_elec + ms2) // 2, n)
n_beta = _clamped_occ(n_elec - n_alpha, n)
n_virt_a = max(0, n - n_alpha)
n_virt_b = max(0, n - n_beta)
n_pair_a = n_alpha * n_virt_a
n_pair_b = n_beta * n_virt_b
n_pair = n_pair_a + n_pair_b
if n_pair <= 0:
return
# UHF TDA builds same-spin ovov/oovv blocks and one cross-spin block;
# UHF Casida reuses that path and transforms a second set for B.
tda_blocks = 2 * n_pair_a**2 + 2 * n_pair_b**2 + n_pair_a * n_pair_b
casida_extra_blocks = n_pair_a**2 + n_pair_b**2 + n_pair_a * n_pair_b
block_elems = tda_blocks + (casida_extra_blocks if is_casida else 0)
est.by_category["TDDFT ovov/oovv blocks"] = block_elems * 8
pair2 = n_pair**2
else:
n_occ = _clamped_occ(molecule.n_electrons() // 2, n)
n_virt = max(0, n - n_occ)
n_pair = n_occ * n_virt
if n_pair <= 0:
return
pair2 = n_pair**2
est.by_category["TDDFT ovov/oovv blocks"] = 2 * pair2 * 8
est.by_category["TDDFT AO ERI tensor"] = n**4 * 8
response_matrix_count = 2 if is_casida else 1
est.by_category["TDDFT A/B response matrices"] = (
response_matrix_count * pair2 * 8
)
# Dense eigensolvers and Casida square-root products hold several
# additional n_pair x n_pair arrays transiently.
eig_workspace_count = 6 if is_casida else 2
est.by_category["TDDFT eigensolver workspace"] = eig_workspace_count * pair2 * 8
est.by_category["TDDFT state vectors"] = min(n_states, n_pair) * n_pair * 8
if functional is not None:
est.by_category["TDDFT f_xc kernel blocks"] = pair2 * 8
def _rks_estimate(molecule: Molecule, basis: BasisSet, options) -> MemoryEstimate:
est = _rhf_estimate(molecule, basis, options)
extra = _dft_xc_estimate(molecule, basis, options)
if extra:
est.by_category["DFT grid + chi"] = extra
return est
def _uks_estimate(molecule: Molecule, basis: BasisSet, options) -> MemoryEstimate:
est = _uhf_estimate(molecule, basis, options)
extra = _dft_xc_estimate(
molecule,
basis,
options,
open_shell=True,
)
if extra:
est.by_category["DFT grid + chi"] = extra
return est
def _closed_shell_occ_vir(molecule: Molecule, n_basis: int) -> tuple[int, int]:
n_occ = _clamped_occ(molecule.n_electrons() // 2, n_basis)
return max(1, n_occ), max(1, n_basis - n_occ)
def _open_shell_occ_vir(
molecule: Molecule,
n_basis: int,
) -> tuple[int, int, int, int]:
n_elec = molecule.n_electrons()
mult = int(molecule.multiplicity)
n_alpha = _clamped_occ((n_elec + mult - 1) // 2, n_basis)
n_beta = _clamped_occ(n_elec - n_alpha, n_basis)
return (
max(1, n_alpha),
max(1, n_beta),
max(1, n_basis - n_alpha),
max(1, n_basis - n_beta),
)
def _mp2_estimate(molecule: Molecule, basis: BasisSet, options) -> MemoryEstimate:
ref_options = _post_reference_options(options)
mp2_options = _post_method_options(options, "mp2_options")
est = _rhf_estimate(molecule, basis, ref_options)
n = _n_basis(molecule, basis)
n_occ, n_vir = _closed_shell_occ_vir(molecule, n)
ovov_bytes = n_occ * n_vir * n_occ * n_vir * 8
est.by_category["OVOV MO tensor"] = ovov_bytes
uses_df = _posthf_uses_density_fit(mp2_options, default=False)
if uses_df:
n_aux = _aux_basis_size(molecule, n, mp2_options)
est.by_category["DF-MP2 three-index/B tensors"] = (
2 * n_aux * n * n * 8
)
est.by_category["DF-MP2 metric/workspace"] = 2 * n_aux * n_aux * 8
est.by_category["DF-MP2 aux-OV workspace"] = n_aux * n_occ * n_vir * 8
if bool(getattr(mp2_options, "use_float_intermediates", False)):
n_ov = n_occ * n_vir
est.by_category["DF-MP2 float GEMM scratch"] = (
n_aux * n_ov * 4 + n_ov * n_ov * 4
)
if bool(getattr(mp2_options, "report_ri_residual", False)):
est.by_category["MP2 RI-residual AO ERI tensor"] = n**4 * 8
est.by_category["MP2 RI-residual AO->MO I1 scratch"] = (
n_occ * n**3 * 8
)
est.by_category["MP2 RI-residual AO->MO I2 scratch"] = (
n_occ * n_vir * n * n * 8
)
est.by_category["MP2 RI-residual AO->MO I3 scratch"] = (
n_occ * n_vir * n_occ * n * 8
)
est.by_category["MP2 RI-residual OVOV tensor"] = ovov_bytes
else:
est.by_category["MP2 AO ERI tensor"] = n**4 * 8
est.by_category["MP2 AO->MO I1 scratch"] = n_occ * n**3 * 8
est.by_category["MP2 AO->MO I2 scratch"] = (
n_occ * n_vir * n * n * 8
)
est.by_category["MP2 AO->MO I3 scratch"] = (
n_occ * n_vir * n_occ * n * 8
)
return est
def _ump2_estimate(molecule: Molecule, basis: BasisSet, options) -> MemoryEstimate:
ref_options = _post_reference_options(options)
ump2_options = _post_method_options(options, "ump2_options")
est = _uhf_estimate(molecule, basis, ref_options)
n = _n_basis(molecule, basis)
n_alpha, n_beta, n_vir_a, n_vir_b = _open_shell_occ_vir(molecule, n)
ov_aa = n_alpha * n_vir_a * n_alpha * n_vir_a
ov_bb = n_beta * n_vir_b * n_beta * n_vir_b
ov_ab = n_alpha * n_vir_a * n_beta * n_vir_b
largest_ovov = max(ov_aa, ov_bb, ov_ab) * 8
est.by_category["UMP2 largest OVOV channel"] = largest_ovov
uses_df = _posthf_uses_density_fit(ump2_options, default=False)
if uses_df:
n_aux = _aux_basis_size(molecule, n, ump2_options)
est.by_category["DF-UMP2 three-index/B tensors"] = (
2 * n_aux * n * n * 8
)
est.by_category["DF-UMP2 metric/workspace"] = 2 * n_aux * n_aux * 8
est.by_category["DF-UMP2 aux-OV workspace"] = (
n_aux * n_alpha * n_vir_a + n_aux * n_beta * n_vir_b
) * 8
if bool(getattr(ump2_options, "report_ri_residual", False)):
est.by_category["UMP2 RI-residual AO ERI tensor"] = n**4 * 8
same_a = (
n_alpha * n**3
+ n_alpha * n_vir_a * n * n
+ n_alpha * n_vir_a * n_alpha * n
)
same_b = (
n_beta * n**3
+ n_beta * n_vir_b * n * n
+ n_beta * n_vir_b * n_beta * n
)
mixed = (
n_alpha * n**3
+ n_alpha * n_vir_a * n * n
+ n_alpha * n_vir_a * n_beta * n
)
est.by_category["UMP2 RI-residual largest AO->MO scratch"] = (
max(same_a, same_b, mixed) * 8
)
est.by_category["UMP2 RI-residual largest OVOV channel"] = (
largest_ovov
)
else:
est.by_category["UMP2 AO ERI tensor"] = n**4 * 8
same_a = (
n_alpha * n**3
+ n_alpha * n_vir_a * n * n
+ n_alpha * n_vir_a * n_alpha * n
)
same_b = (
n_beta * n**3
+ n_beta * n_vir_b * n * n
+ n_beta * n_vir_b * n_beta * n
)
mixed = (
n_alpha * n**3
+ n_alpha * n_vir_a * n * n
+ n_alpha * n_vir_a * n_beta * n
)
est.by_category["UMP2 largest AO->MO scratch"] = (
max(same_a, same_b, mixed) * 8
)
return est
def _rohf_mp2_estimate(
molecule: Molecule,
basis: BasisSet,
options,
) -> MemoryEstimate:
# ROHF-MP2 uses the restricted-open reference but carries unrestricted
# density/Fock buffers in the estimator.
return _ump2_estimate(molecule, basis, options)
def _ovgf_estimate(molecule: Molecule, basis: BasisSet, options) -> MemoryEstimate:
"""Memory estimate for the dense OVGF/GF2 runner path."""
ref_options = _post_reference_options(options)
n = _n_basis(molecule, basis)
open_shell = int(getattr(molecule, "multiplicity", 1) or 1) != 1
if open_shell:
est = _uhf_estimate(molecule, basis, ref_options)
else:
est = _rhf_estimate(molecule, basis, ref_options)
ao_eri_bytes = n**4 * 8
est.by_category["OVGF AO ERI tensor"] = ao_eri_bytes
if open_shell:
n_spin = 2 * n
n_alpha, n_beta, n_vir_a, n_vir_b = _open_shell_occ_vir(molecule, n)
n_occ_spin = n_alpha + n_beta
n_vir_spin = n_vir_a + n_vir_b
est.by_category["OVGF spin-block MO tensors"] = 3 * ao_eri_bytes
est.by_category["OVGF spin-orbital tensor"] = n_spin**4 * 8
est.by_category["OVGF Dyson spin work slices"] = (
max(
n_occ_spin * n_vir_spin * n_vir_spin,
n_vir_spin * n_occ_spin * n_occ_spin,
)
* 4
* 8
)
return est
n_occ, n_vir = _closed_shell_occ_vir(molecule, n)
est.by_category["OVGF MO ERI tensor"] = ao_eri_bytes
est.by_category["OVGF AO->MO transform scratch"] = ao_eri_bytes
est.by_category["OVGF Dyson work slices"] = (
max(n_occ * n_vir * n_occ, n_vir * n_occ * n_vir) * 4 * 8
)
if 2 * n <= 60:
est.by_category["OVGF renormalized spin-orbital tensor"] = (2 * n) ** 4 * 8
return est
def _ccsd_estimate(
molecule: Molecule,
basis: BasisSet,
options,
*,
open_shell: bool = False,
triples: bool = False,
) -> MemoryEstimate:
ref_options = _post_reference_options(options)
cc_options = _post_method_options(options, "ccsd_options")
est = (
_uhf_estimate(molecule, basis, ref_options)
if open_shell
else _rhf_estimate(molecule, basis, ref_options)
)
n = _n_basis(molecule, basis)
if open_shell:
n_occ = max(1, molecule.n_electrons())
n_vir = max(1, 2 * n - n_occ)
else:
n_occ, n_vir = _closed_shell_occ_vir(molecule, n)
n_ov = n_occ * n_vir
n_oovv = n_occ * n_occ * n_vir * n_vir
est.by_category["CCSD T1/T2 amplitudes"] = (4 * n_ov + 2 * n_oovv) * 8
diis_vectors = max(1, int(getattr(cc_options, "diis_subspace_size", 6) or 6))
est.by_category["CCSD residual/DIIS amplitudes"] = (
2
* diis_vectors
* (n_ov + n_oovv)
* 8
)
est.by_category["CCSD D1/D2 intermediates"] = (
2 * n_oovv # tau and tau-tilde
+ n_occ * n_vir
+ 2 * n_vir * n_vir
+ 2 * n_occ * n_occ
+ n_occ**4
+ 3 * n_ov * n_ov # W1/W2/WX ring intermediates
+ 2 * n_oovv # particle-particle M buffer + pair-exchange half
) * 8
est.by_category["CCSD native solver/runtime scratch"] = (
_CCSD_NATIVE_RUNTIME_FLOOR_BYTES
)
if _posthf_uses_density_fit(cc_options, default=True):
n_aux = _aux_basis_size(molecule, n, cc_options)
est.by_category["DF-CCSD auxiliary integrals"] = (
n_aux
* (n * n + n_occ * n_vir + n_vir * n_vir + n_occ * n_occ)
* 8
)
est.by_category["DF-CCSD MO integral blocks"] = (
n_ov * n_ov
+ n_occ**4
+ n_occ**3 * n_vir
+ n_oovv
+ n_occ * n_vir**3
) * 8
vvvv_elems = n_vir**4
vvvv_bytes = vvvv_elems * 8
if vvvv_bytes <= _CCSD_VVVV_INCORE_BUDGET_BYTES:
# V.vv_vv plus the W_abef work matrix are both resident on the
# in-core path. B_vv itself is counted in the auxiliary category.
est.by_category["DF-CCSD VVVV/tile scratch"] = 2 * vvvv_bytes
else:
rows = max(1, _CCSD_VVVV_TILE_ROWS)
a_block = max(1, min(n_vir, rows // max(1, n_vir)))
est.by_category["DF-CCSD VVVV/tile scratch"] = (
2 * a_block * n_vir**3 + n_aux * a_block * n_vir
) * 8
else:
est.by_category["CCSD MO ERI workspace"] = max(n_oovv, n**4) * 8
est.by_category["CCSD VVVV/tile scratch"] = n_vir**4 * 8
compute_triples = triples or bool(getattr(cc_options, "compute_triples", False))
if compute_triples:
triple_elems = n_occ * n_occ * n_occ * n_vir * n_vir * n_vir
mode = str(
getattr(cc_options, "triples_memory_mode", "fast") or "fast"
).lower()
batch_factor = 1 if mode == "low" else 2
est.by_category["CCSD(T) triples amplitudes"] = triple_elems * 8
est.by_category["CCSD(T) triples intermediates"] = (
batch_factor
* (
n_occ * n_occ * n_vir**4
+ n_occ**4 * n_vir * n_vir
+ triple_elems
)
* 8
)
return est
def _ccsd_t_estimate(molecule: Molecule, basis: BasisSet, options) -> MemoryEstimate:
return _ccsd_estimate(molecule, basis, options, triples=True)
def _uccsd_estimate(molecule: Molecule, basis: BasisSet, options) -> MemoryEstimate:
return _ccsd_estimate(molecule, basis, options, open_shell=True)
def _uccsd_t_estimate(molecule: Molecule, basis: BasisSet, options) -> MemoryEstimate:
return _ccsd_estimate(molecule, basis, options, open_shell=True, triples=True)
def _dlpno_active_counts(
molecule: Molecule,
n_basis: int,
method_options,
*,
open_shell: bool = False,
) -> tuple[int, int, int, int]:
frozen = max(0, int(getattr(method_options, "n_frozen", 0) or 0))
if open_shell:
n_alpha, n_beta, n_vir_a, n_vir_b = _open_shell_occ_vir(molecule, n_basis)
n_occ = max(1, n_alpha + n_beta - 2 * frozen)
n_vir = max(1, n_vir_a + n_vir_b)
else:
n_occ_raw, n_vir = _closed_shell_occ_vir(molecule, n_basis)
n_occ = max(1, n_occ_raw - frozen)
avg_pno = min(n_vir, max(8, (n_vir + 2) // 3))
avg_pao = min(n_basis, max(avg_pno * 2, (2 * n_vir + 2) // 3))
return n_occ, n_vir, avg_pao, avg_pno
def _dlpno_pair_count(n_occ: int, *, open_shell: bool = False) -> int:
if open_shell:
return max(1, n_occ * n_occ)
return max(1, n_occ * (n_occ + 1) // 2)
def _dlpno_mp2_estimate(
molecule: Molecule,
basis: BasisSet,
options,
*,
open_shell: bool = False,
) -> MemoryEstimate:
ref_options = _post_reference_options(options)
method_options = _post_method_options(
options,
"dlpno_options" if not open_shell else "dlpno_ump2_options",
)
if method_options is None and open_shell:
method_options = _post_method_options(options, "dlpno_options")
est = (
_uhf_estimate(molecule, basis, ref_options)
if open_shell
else _rhf_estimate(molecule, basis, ref_options)
)
n = _n_basis(molecule, basis)
n_aux = _aux_basis_size(molecule, n, method_options)
n_occ, n_vir, avg_pao, avg_pno = _dlpno_active_counts(
molecule,
n,
method_options,
open_shell=open_shell,
)
n_pairs = _dlpno_pair_count(n_occ, open_shell=open_shell)
est.by_category["DLPNO DF auxiliary integrals"] = n_aux * n * n * 8
est.by_category["DLPNO PAO coefficients/domains"] = (
n_pairs * (n * avg_pao + avg_pao * avg_pno) * 8
)
est.by_category["DLPNO pair lists"] = n_pairs * 1024
est.by_category["DLPNO-MP2 PNO pair amplitudes"] = (
n_pairs * 4 * avg_pno * avg_pno * 8
)
est.by_category["DLPNO-MP2 aux-OV workspace"] = n_aux * n_occ * n_vir * 8
return est
def _dlpno_ccsd_estimate(
molecule: Molecule,
basis: BasisSet,
options,
*,
open_shell: bool = False,
triples: bool = False,
) -> MemoryEstimate:
ref_options = _post_reference_options(options)
method_options = _post_method_options(options, "dlpno_ccsd_options")
est = (
_uhf_estimate(molecule, basis, ref_options)
if open_shell
else _rhf_estimate(molecule, basis, ref_options)
)
n = _n_basis(molecule, basis)
n_aux = _aux_basis_size(molecule, n, method_options)
n_occ, n_vir, avg_pao, avg_pno = _dlpno_active_counts(
molecule,
n,
method_options,
open_shell=open_shell,
)
n_pairs = _dlpno_pair_count(n_occ, open_shell=open_shell)
est.by_category["DLPNO DF auxiliary integrals"] = n_aux * n * n * 8
est.by_category["DLPNO PAO coefficients/domains"] = (
n_pairs * (n * avg_pao + avg_pao * avg_pno) * 8
)
est.by_category["DLPNO pair lists"] = n_pairs * 1024
est.by_category["DLPNO-CCSD T1/T2 pair amplitudes"] = (
n_pairs * (6 * avg_pno * avg_pno + n_occ * avg_pno) * 8
)
est.by_category["DLPNO-CCSD pair residual/intermediates"] = (
n_pairs * (12 * avg_pno**4 + 8 * n_occ * avg_pno * avg_pno) * 8
)
est.by_category["DLPNO-CCSD auxiliary pair workspace"] = (
n_aux * n_occ * avg_pno + n_pairs * n_aux * avg_pno
) * 8
compute_triples = triples or bool(
getattr(method_options, "compute_triples", False)
)
if compute_triples:
n_triples = max(1, n_occ * (n_occ + 1) * (n_occ + 2) // 6)
avg_tno = min(n_vir, max(avg_pno, (3 * avg_pno + 1) // 2))
est.by_category["DLPNO triples T_ijk^abc workspace"] = (
n_triples * avg_tno**3 * 8
)
est.by_category["DLPNO triples PAO/TNO domains"] = (
n_triples * (n * avg_tno + avg_tno * avg_tno) * 8
)
return est
def _dlpno_mp2_closed_estimate(
molecule: Molecule,
basis: BasisSet,
options,
) -> MemoryEstimate:
return _dlpno_mp2_estimate(molecule, basis, options)
def _dlpno_ump2_estimate(
molecule: Molecule,
basis: BasisSet,
options,
) -> MemoryEstimate:
return _dlpno_mp2_estimate(molecule, basis, options, open_shell=True)
def _dlpno_ccsd_closed_estimate(
molecule: Molecule,
basis: BasisSet,
options,
) -> MemoryEstimate:
return _dlpno_ccsd_estimate(molecule, basis, options)
def _dlpno_ccsd_t_estimate(
molecule: Molecule,
basis: BasisSet,
options,
) -> MemoryEstimate:
return _dlpno_ccsd_estimate(molecule, basis, options, triples=True)
def _dlpno_uccsd_estimate(
molecule: Molecule,
basis: BasisSet,
options,
) -> MemoryEstimate:
return _dlpno_ccsd_estimate(molecule, basis, options, open_shell=True)
def _dlpno_uccsd_t_estimate(
molecule: Molecule,
basis: BasisSet,
options,
) -> MemoryEstimate:
return _dlpno_ccsd_estimate(
molecule,
basis,
options,
open_shell=True,
triples=True,
)
def _wavefunction_estimate(molecule, basis, options=None):
"""Memory estimate for determinant CI-like wavefunction methods."""
n = _n_basis(molecule, basis)
est = _rhf_estimate(molecule, basis, _post_reference_options(options))
n_elec = max(0, molecule.n_electrons())
try:
n_det = max(1, comb(max(0, 2 * n), min(n_elec, max(0, 2 * n))))
except ValueError:
n_det = 1
est.by_category["CI determinant vector"] = 3 * n_det * 8
est.by_category["CI sigma/residual workspace"] = max(n_det * 16, n**4 * 8)
est.by_category["MO integral transform workspace"] = n * n * n * n * 8
return est
def _ccsdt_excitation_count(
norb: int,
nalpha: int,
nbeta: int,
rank: int,
) -> int:
"""Number of spin-preserving determinants at one excitation rank."""
total = 0
for rank_alpha in range(rank + 1):
rank_beta = rank - rank_alpha
if not (
0 <= rank_alpha <= nalpha
and rank_alpha <= norb - nalpha
and 0 <= rank_beta <= nbeta
and rank_beta <= norb - nbeta
):
continue
total += (
comb(nalpha, rank_alpha)
* comb(norb - nalpha, rank_alpha)
* comb(nbeta, rank_beta)
* comb(norb - nbeta, rank_beta)
)
return total
def _ccsdt_estimate(molecule, basis, options=None):
"""Peak estimate for the dense determinant-space full-CCSDT route."""
n_basis = _n_basis(molecule, basis)
reference_estimator = (
_uhf_estimate if int(molecule.multiplicity) > 1 else _rhf_estimate
)
est = reference_estimator(
molecule,
basis,
_post_reference_options(options),
)
method_options = _post_method_options(options, "ccsdt_options")
active_space = _option_value(options, "active_space", None)
n_elec_total = max(0, int(molecule.n_electrons()))
if active_space is not None:
n_active_orb = int(active_space[0])
n_active_elec = int(active_space[1])
else:
n_frozen = _option_value(method_options, "n_frozen_core", None)
if n_frozen is None:
from .cc import chemical_core_orbital_count
n_frozen = chemical_core_orbital_count(molecule)
n_frozen = max(0, int(n_frozen))
n_active_orb = max(1, n_basis - n_frozen)
n_active_elec = max(0, n_elec_total - 2 * n_frozen)
ms2 = max(0, int(molecule.multiplicity) - 1)
if (n_active_elec + ms2) % 2:
nalpha = (n_active_elec + 1) // 2
else:
nalpha = (n_active_elec + ms2) // 2
nbeta = n_active_elec - nalpha
nalpha = min(max(0, nalpha), n_active_orb)
nbeta = min(max(0, nbeta), n_active_orb)
n_amplitudes = sum(
_ccsdt_excitation_count(n_active_orb, nalpha, nbeta, rank)
for rank in (1, 2, 3)
)
max_state_rank = min(
7,
n_active_elec,
2 * n_active_orb - n_active_elec,
)
n_state_determinants = sum(
_ccsdt_excitation_count(n_active_orb, nalpha, nbeta, rank)
for rank in range(max_state_rank + 1)
)
diis_size = int(_option_value(method_options, "diis_subspace_size", 6))
est.by_category["CCSDT MO integral transform workspace"] = (
3 * n_basis**4 * 8
)
est.by_category["CCSDT active MO integrals"] = (
2 * n_active_orb**4 * 8 + n_active_orb**2 * 8
)
est.by_category["CCSDT excitation metadata"] = n_amplitudes * 320
est.by_category["CCSDT amplitudes and DIIS"] = n_amplitudes * (
8 * (8 + 2 * max(1, diis_size)) + 1
)
# Python determinant maps dominate this exact-projection implementation.
# The two-body action can transiently reach two ranks above exp(T)|0>;
# six simultaneous maps at 160 bytes/entry is a conservative peak model.
est.by_category["CCSDT determinant-state dictionaries"] = (
6 * n_state_determinants * 160
)
return est
def _cc3_estimate(molecule, basis, options=None):
"""Conservative peak estimate for dense determinant-space CC3.
CC3 stores DIIS vectors for singles and doubles only, but its transformed
Hamiltonian commutators visit the same determinant ranks as the full
CCSDT projection. Reusing the CCSDT estimate is therefore conservative
and avoids understating the dominant Python determinant-map workspace.
"""
ccsdt_options = {
"scf_options": _post_reference_options(options),
"ccsdt_options": _post_method_options(options, "cc3_options"),
"active_space": _option_value(options, "active_space", None),
}
est = _ccsdt_estimate(molecule, basis, ccsdt_options)
est.by_category = {
label.replace("CCSDT", "CC3"): value
for label, value in est.by_category.items()
}
return est
def _casscf_estimate(molecule, basis, options=None):
n = _n_basis(molecule, basis)
est = _rhf_estimate(molecule, basis, _post_reference_options(options))
active_space = _option_value(options, "active_space", None)
if active_space is None:
active_orb = min(n, max(2, molecule.n_electrons() // 2))
active_elec = min(molecule.n_electrons(), 2 * active_orb)
else:
active_orb, active_elec = int(active_space[0]), int(active_space[1])
n_det = _active_determinant_count(active_orb, active_elec, molecule.multiplicity)
est.by_category["CASSCF CI vector"] = 3 * n_det * 8
est.by_category["CASSCF 4-RDM storage"] = active_orb**8 * 8
est.by_category["CASSCF MO integral transform"] = n**4 * 8 * 3
return est
def _fci_estimate(molecule, basis, options=None):
n = _n_basis(molecule, basis)
est = _rhf_estimate(molecule, basis, _post_reference_options(options))
n_elec = max(0, molecule.n_electrons())
try:
n_det = max(1, comb(max(0, 2 * n), min(n_elec, max(0, 2 * n))))
except ValueError:
n_det = 1
est.by_category["FCI determinant vector"] = 3 * n_det * 8
est.by_category["FCI sigma/residual workspace"] = max(n_det * 16, n**4 * 8)
est.by_category["FCI MO integral transform"] = n**4 * 8 * 3
return est
def _active_determinant_count(
n_active_orb: int,
n_active_elec: int,
multiplicity: int,
) -> int:
"""Determinants in the spin block used by the CAS reference."""
ms2 = int(multiplicity) - 1
if (n_active_elec + ms2) % 2 != 0:
return max(1, comb(2 * n_active_orb, n_active_elec))
n_alpha = (n_active_elec + ms2) // 2
n_beta = n_active_elec - n_alpha
if not (0 <= n_alpha <= n_active_orb and 0 <= n_beta <= n_active_orb):
return max(1, comb(2 * n_active_orb, n_active_elec))
return max(1, comb(n_active_orb, n_alpha) * comb(n_active_orb, n_beta))
def _caspt2_label_count(
n_orb: int,
n_core: int,
n_active_orb: int,
n_frozen: int,
) -> int:
"""Number of raw internally-contracted CASPT2 labels.
Mirrors ``solvers._mrpt._ic_caspt2_label_count`` without importing the
solver stack from the lightweight pre-flight estimator.
"""
inactive = max(0, n_core - max(0, n_frozen))
secondary = max(0, n_orb - n_core - n_active_orb)
occ = inactive + n_active_orb
virt = n_active_orb + secondary
n_exc = occ * virt - n_active_orb
return max(0, n_exc + n_exc * n_exc)
def _mrpt2_orbital_pair_count(
n_orb: int,
n_core: int,
n_active_orb: int,
) -> int:
"""Non-redundant CASSCF/MRPT2 orbital-rotation pair count."""
inactive = max(0, int(n_core))
active = max(0, int(n_active_orb))
secondary = max(0, int(n_orb) - inactive - active)
return inactive * active + inactive * secondary + active * secondary
def _pt2_gradient_requested(options, option_key: str) -> bool:
"""Whether the requested MRPT2 route computes a correlation gradient."""
pt2_options = _option_value(options, option_key, None)
if pt2_options is None:
return False
return bool(_option_value(pt2_options, "compute_corr_grad", False))
def _pt2_fd_worker_bytes(nmo: int, *, use_cpp_transform: bool) -> int:
"""Mirror gradient._parallel_helpers.estimate_pt2_fd_worker_bytes."""
tensor_bytes = max(1, int(nmo)) ** 4 * 8
multiplier = 10 if use_cpp_transform else 7
matrix_bytes = max(1, int(nmo)) ** 2 * 8
return max(256 * 1024**2, multiplier * tensor_bytes + 12 * matrix_bytes)
def _pt2_fd_worker_count(
*,
npr: int,
nmo: int,
use_cpp_transform: bool,
) -> int:
"""Estimate the runtime PT2 FD process count for memory preflight."""
npr = int(npr)
if npr <= 4:
return 1
def _positive_env(name: str) -> int | None:
raw = os.environ.get(name)
if raw is None or raw.strip() == "":
return None
try:
value = int(raw)
except ValueError:
return None
return value if value > 0 else None
forced = _positive_env("VIBEQC_PT2_GRADIENT_N_JOBS")
if forced is not None:
return max(1, min(npr, forced))
cpu_limit = max(1, min(npr, os.cpu_count() or 1))
hard_cap = _positive_env("VIBEQC_PT2_GRADIENT_MAX_JOBS") or 4
worker_bytes = _pt2_fd_worker_bytes(nmo, use_cpp_transform=use_cpp_transform)
available = available_memory_bytes()
if available <= 0:
memory_limit = 1
else:
reserve = max(2 * _GB, available // 4)
usable = max(0, available - reserve)
memory_limit = max(1, usable // worker_bytes)
return max(1, min(cpu_limit, hard_cap, int(memory_limit)))
def _add_pt2_gradient_fd_estimate(
est: MemoryEstimate,
*,
n_orb: int,
n_core: int,
n_active_orb: int,
) -> None:
"""Account for CASPT2/NEVPT2 Z-vector dE2/dk finite-difference workers."""
npr = _mrpt2_orbital_pair_count(n_orb, n_core, n_active_orb)
if npr <= 0:
return
use_cpp_transform = False
try:
from . import _vibeqc_core as _core
use_cpp_transform = hasattr(_core, "transform_4index_mo")
except Exception:
use_cpp_transform = False
n_workers = _pt2_fd_worker_count(
npr=npr,
nmo=n_orb,
use_cpp_transform=use_cpp_transform,
)
worker_bytes = _pt2_fd_worker_bytes(
n_orb,
use_cpp_transform=use_cpp_transform,
)
est.by_category["PT2 gradient FD workers"] = n_workers * worker_bytes
def _caspt2_estimate(molecule: Molecule, basis: BasisSet, options) -> MemoryEstimate:
"""Memory estimate for the internally-contracted CASPT2 pilot route.
The production IC-CASPT2 path materializes active-CI external states and
their generalized-Fock images in Python dictionaries before solving the
signature-block response system. For release-sized active spaces this
dominates the small ``n_basis**4`` MO-ERI estimate by orders of magnitude;
account for it so runs fail before the host watchdog kills them.
"""
est = _rhf_estimate(molecule, basis, None)
n_orb = _n_basis(molecule, basis)
n_elec_total = molecule.n_electrons()
active_space = _option_value(options, "active_space", None)
if active_space is None:
est.by_category["MO ERI tensor + solver"] = (
n_orb * n_orb * n_orb * n_orb * 8 * 3
)
return est
n_active_orb, n_active_elec = (int(active_space[0]), int(active_space[1]))
if n_active_orb <= 0 or n_active_elec < 0:
est.by_category["MO ERI tensor + solver"] = (
n_orb * n_orb * n_orb * n_orb * 8 * 3
)
return est
if (n_elec_total - n_active_elec) % 2 != 0:
est.by_category["MO ERI tensor + solver"] = (
n_orb * n_orb * n_orb * n_orb * 8 * 3
)
return est
caspt2_options = _option_value(options, "caspt2_options", None)
n_core = max(0, (n_elec_total - n_active_elec) // 2)
n_frozen = int(_option_value(caspt2_options, "n_frozen", 0) or 0)
n_labels = _caspt2_label_count(n_orb, n_core, n_active_orb, n_frozen)
n_det = _active_determinant_count(
n_active_orb,
n_active_elec,
molecule.multiplicity,
)
est.by_category["CASPT2 MO integral workspace"] = n_orb**4 * 8 * 3
est.by_category["CASPT2 contracted labels"] = max(1, n_labels) * 256
# Each selected contracted candidate can carry an active-CI dictionary,
# and the current block builder also materializes Fock-applied images.
# Python dict/tuple/int/float overhead dwarfs the raw float payload; 2 KiB
# per active determinant entry is conservative but matches the observed
# N2/cc-pVDZ CAS(6,6) watchdog kill scale.
est.by_category["CASPT2 active-CI state workspace"] = (
max(1, n_labels) * max(1, n_det) * 2048
)
# Dense per-signature overlap/H0/coupling blocks and their
# metric-orthonormalized transforms. The direct active-CI route avoids the
# full n_labels^2 matrix, so keep this term below the state workspace while
# still exposing the response-system cost in the report.
est.by_category["CASPT2 response blocks"] = max(1, n_labels) * 512 * 8
if _pt2_gradient_requested(options, "caspt2_options"):
_add_pt2_gradient_fd_estimate(
est,
n_orb=n_orb,
n_core=n_core,
n_active_orb=n_active_orb,
)
return est
def _nevpt2_estimate(molecule: Molecule, basis: BasisSet, options) -> MemoryEstimate:
"""Memory estimate for NEVPT2, including optional Z-vector FD workers."""
est = _rhf_estimate(molecule, basis, None)
n_orb = _n_basis(molecule, basis)
est.by_category["MO ERI tensor + solver"] = n_orb**4 * 8 * 3
active_space = _option_value(options, "active_space", None)
if active_space is None or not _pt2_gradient_requested(options, "nevpt2_options"):
return est
n_active_orb, n_active_elec = (int(active_space[0]), int(active_space[1]))
n_elec_total = molecule.n_electrons()
if n_active_orb <= 0 or n_active_elec < 0:
return est
if (n_elec_total - n_active_elec) % 2 != 0:
return est
n_core = max(0, (n_elec_total - n_active_elec) // 2)
_add_pt2_gradient_fd_estimate(
est,
n_orb=n_orb,
n_core=n_core,
n_active_orb=n_active_orb,
)
return est
_ESTIMATORS = {
"rhf": _rhf_estimate,
"uhf": _uhf_estimate,
"rks": _rks_estimate,
"uks": _uks_estimate,
"mp2": _mp2_estimate,
"scs-mp2": _mp2_estimate,
"sos-mp2": _mp2_estimate,
"ump2": _ump2_estimate,
"rohf-mp2": _rohf_mp2_estimate,
"ovgf": _ovgf_estimate,
"ccsd": _ccsd_estimate,
"ccsd(t)": _ccsd_t_estimate,
"ccsd[t]": _ccsd_t_estimate,
"a-ccsd(t)": _ccsd_t_estimate,
"uccsd": _uccsd_estimate,
"uccsd(t)": _uccsd_t_estimate,
"bccd": _ccsd_estimate,
"bccd(t)": _ccsd_t_estimate,
"cc2": _ccsd_estimate,
"ccd": _ccsd_estimate,
"lccd": _ccsd_estimate,
"lccsd": _ccsd_estimate,
"cepa(0)": _ccsd_estimate,
"cepa(1)": _ccsd_estimate,
"cepa(2)": _ccsd_estimate,
"cepa(3)": _ccsd_estimate,
"qcisd": _ccsd_estimate,
"qcisd(t)": _ccsd_t_estimate,
"cc3": _cc3_estimate,
"ccsdt": _ccsdt_estimate,
"dlpno-mp2": _dlpno_mp2_closed_estimate,
"dlpno-ump2": _dlpno_ump2_estimate,
"dlpno-ccsd": _dlpno_ccsd_closed_estimate,
"dlpno-ccsd(t)": _dlpno_ccsd_t_estimate,
"dlpno-uccsd": _dlpno_uccsd_estimate,
"dlpno-uccsd(t)": _dlpno_uccsd_t_estimate,
"cisd": _wavefunction_estimate,
"selected_ci": _wavefunction_estimate,
"dmrg": _wavefunction_estimate,
"v2rdm": _wavefunction_estimate,
"transcorrelated_ci": _wavefunction_estimate,
"casci": _casscf_estimate,
"mrci": _wavefunction_estimate,
"casscf": _casscf_estimate,
"nevpt2": _nevpt2_estimate,
"caspt2": _caspt2_estimate,
"fci": _fci_estimate,
}
[docs]
def estimate_memory(
molecule: Molecule,
basis: BasisSet,
*,
method: str,
options=None,
) -> MemoryEstimate:
"""Peak memory estimate for ``method``.
Parameters
----------
molecule
The :class:`Molecule` about to be run.
basis
The :class:`BasisSet` paired with the molecule.
method
Molecular SCF, MP2, CC/CI, DLPNO, CAS, and MRPT2 method label.
Case-insensitive. Periodic drivers use route-specific estimator
helpers rather than this molecular entry point.
options
The matching ``*Options`` struct, or ``None`` for defaults.
DIIS history size and DFT grid dimensions are read from it
when available.
"""
key = method.lower()
if key not in _ESTIMATORS:
raise ValueError(
f"estimate_memory: unknown method {method!r}. Known: {sorted(_ESTIMATORS)}"
)
return _ESTIMATORS[key](molecule, basis, options)