diff --git a/model/examples/cpp/cpp_integration/prepare_inputs.py b/model/examples/cpp/cpp_integration/prepare_inputs.py index 6eda4b34..f08564e4 100755 --- a/model/examples/cpp/cpp_integration/prepare_inputs.py +++ b/model/examples/cpp/cpp_integration/prepare_inputs.py @@ -6,7 +6,8 @@ import torch from pyscf.dft import gen_grid from skala.functional.traditional import LDA -from skala.pyscf.features import generate_features +from skala.pyscf.evaluation import FeatureSpec +from skala.pyscf.model_chunking import evaluate_model_features from skala_model import SkalaFunctional from pyscf import dft, gto @@ -41,8 +42,8 @@ def main() -> None: grid = gen_grid.Grids(molecule) grid.level = 3 grid.build(sort_grids=False) - features = generate_features( - molecule, dm, grid, features=set(SkalaFunctional.features) + features = evaluate_model_features( + molecule, dm, grid, FeatureSpec(SkalaFunctional.features) ) # Save all features as individual .pt files. diff --git a/pixi.lock b/pixi.lock index dd83ec99..cb3664c9 100644 --- a/pixi.lock +++ b/pixi.lock @@ -599,6 +599,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/prometheus-cpp-1.3.0-ha5d0236_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.2-py312h5253ce2_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb03c661_1003.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/py-spy-0.4.2-hbe0e3b8_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pyarrow-25.0.0-py312h7900ff3_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pyarrow-core-25.0.0-py312h2054cf2_0_cpu.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.12.13-h8ab3286_1_cpython.conda @@ -3791,6 +3792,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/prometheus-cpp-1.3.0-ha5d0236_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.2-py312h5253ce2_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb03c661_1003.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/py-spy-0.4.2-hbe0e3b8_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pyarrow-25.0.0-py312h7900ff3_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pyarrow-core-25.0.0-py312h2054cf2_0_cpu.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.12.13-h8ab3286_1_cpython.conda @@ -4255,6 +4257,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/pixman-0.46.4-h54a6638_3.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/prometheus-cpp-1.3.0-ha5d0236_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb03c661_1003.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/py-spy-0.4.2-hbe0e3b8_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pyarrow-25.0.0-py312h7900ff3_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pyarrow-core-25.0.0-py312h2054cf2_0_cpu.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.12.13-h8ab3286_1_cpython.conda @@ -4920,6 +4923,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/pixman-0.46.4-h54a6638_3.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/prometheus-cpp-1.3.0-ha5d0236_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb03c661_1003.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/py-spy-0.4.2-hbe0e3b8_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pyarrow-25.0.0-py312h7900ff3_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pyarrow-core-25.0.0-py312h2054cf2_0_cpu.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.12.13-h8ab3286_1_cpython.conda @@ -10464,6 +10468,22 @@ packages: run_exports: {} size: 9115 timestamp: 1786067714761 +- conda: https://conda.anaconda.org/conda-forge/linux-64/py-spy-0.4.2-hbe0e3b8_1.conda + sha256: cb23242015d1b3fb9d009e98edf673900811b73a1e16a5e3d74104c96eb172dc + md5: a9b890615f46d689cbadb145415d333a + depends: + - libgcc >=15 + - __glibc >=2.17,<3.0.a0 + - libunwind >=1.8.3,<1.9.0a0 + - liblzma >=5.8.3,<6.0a0 + constrains: + - __glibc >=2.17 + license: MIT + license_family: MIT + purls: [] + run_exports: {} + size: 1909369 + timestamp: 1787924738130 - conda: https://conda.anaconda.org/conda-forge/linux-64/pyarrow-25.0.0-py312h7900ff3_0.conda sha256: 9b6c4aed5da2a3b8f649e1b73596e0867a5506bed9c3aa45dbaedfb674822191 md5: 6f1918b7565c3cc5ed047a97ea510e55 diff --git a/pixi.toml b/pixi.toml index a4c99fe8..b0055110 100644 --- a/pixi.toml +++ b/pixi.toml @@ -125,6 +125,7 @@ types-PyYAML = "*" [feature.profiling.target.linux-64.dependencies] memray = "*" +py-spy = "*" [feature.benchmark.dependencies] jinja2 = "*" diff --git a/skala/src/skala/features.py b/skala/src/skala/features.py index 04baf73d..9c0187f2 100644 --- a/skala/src/skala/features.py +++ b/skala/src/skala/features.py @@ -2,15 +2,15 @@ """Names of built-in molecular features.""" -from enum import Enum -from typing import TYPE_CHECKING, TypeAlias +from collections.abc import Iterable, Iterator +from enum import StrEnum +from typing import TypeAlias -if TYPE_CHECKING: - from torch import Tensor +from torch import Tensor -class Feature(str, Enum): # noqa: UP042 - Python 3.10-compatible StrEnum - """String-compatible names of features understood by Skala.""" +class Feature(StrEnum): + """features understood by Skala.""" DENSITY = "density" GRAD = "grad" @@ -23,7 +23,73 @@ class Feature(str, Enum): # noqa: UP042 - Python 3.10-compatible StrEnum ATOMIC_GRID_SIZE_BOUND_SHAPE = "atomic_grid_size_bound_shape" COARSE_0_ATOMIC_COORDS = "coarse_0_atomic_coords" - __str__ = str.__str__ +AO_FEATURES = frozenset( + { + Feature.DENSITY, + Feature.GRAD, + Feature.KIN, + Feature.LAPL, + } +) -FeatureMap: TypeAlias = dict[Feature, "Tensor"] + +class AOFeatureSpec: + """Normalized non-empty set of AO-derived features.""" + + def __init__(self, features: Iterable[Feature]) -> None: + self._features = frozenset(features) + unsupported = self._features - AO_FEATURES + if unsupported: + unsupported_names = ", ".join( + sorted(str(feature) for feature in unsupported) + ) + raise ValueError(f"Unsupported AO features: {unsupported_names}") + if not self._features: + raise ValueError("At least one AO-derived feature must be selected.") + + self._feature_slices: dict[Feature, slice] = {} + feature_index = 0 + for feature, width in ( + (Feature.DENSITY, 1), + (Feature.GRAD, 3), + (Feature.KIN, 1), + (Feature.LAPL, 1), + ): + if feature in self._features: + self._feature_slices[feature] = slice( + feature_index, feature_index + width + ) + feature_index += width + self._nfeats = feature_index + + def __contains__(self, feature: object) -> bool: + return feature in self._features + + def __iter__(self) -> Iterator[tuple[Feature, slice]]: + return iter(self._feature_slices.items()) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, AOFeatureSpec): + return NotImplemented + return self._features == other._features + + def __hash__(self) -> int: + return hash(self._features) + + @property + def nderiv(self) -> int: + """Return the required AO derivative order.""" + if Feature.LAPL in self._features: + return 2 + if self._features & {Feature.GRAD, Feature.KIN}: + return 1 + return 0 + + @property + def nfeats(self) -> int: + """Return the number of packed scalar feature channels.""" + return self._nfeats + + +FeatureMap: TypeAlias = dict[Feature, Tensor] diff --git a/skala/src/skala/gpu4pyscf/gradients.py b/skala/src/skala/gpu4pyscf/gradients.py index 17830bb3..5566c956 100644 --- a/skala/src/skala/gpu4pyscf/gradients.py +++ b/skala/src/skala/gpu4pyscf/gradients.py @@ -2,7 +2,7 @@ """Modification of PySCF nuclear gradient object to work with Skala functional.""" -import logging +from collections.abc import Iterator from typing import Any import cupy as cp @@ -15,26 +15,24 @@ from gpu4pyscf.scf.hf import SCF from torch.utils.dlpack import from_dlpack -import skala.pyscf.features as feature from pyscf import gto from skala.dispersion import DFTD3Dispersion -from skala.features import Feature, FeatureMap +from skala.features import Feature from skala.functional.base import ExcFunctionalBase from skala.pyscf.gradient_core import ( - contract_ao_derivative_block, - feature_derivatives, - grid_derivative_block, + assemble_nuclear_gradient, + evaluate_nuclear_feature_derivatives, ) -LOG = logging.getLogger(__name__) - -def veff_and_expl_nuc_grad( +def _veff_and_expl_nuc_grad( functional: ExcFunctionalBase, mol: gto.Mole, grid: dft.Grids, rdm1: torch.Tensor, nuc_grad_feats: set[Feature] | None = None, + *, + max_memory_in_mb: int, ) -> tuple[torch.Tensor, torch.Tensor]: """ returns: @@ -42,41 +40,6 @@ def veff_and_expl_nuc_grad( - 2nd tuple argument: explicit contributions to the nuclear gradient """ - SUPPORTED_FEATS = { - Feature.DENSITY, - Feature.GRAD, - Feature.KIN, - Feature.GRID_COORDS, - Feature.GRID_WEIGHTS, - Feature.ATOMIC_GRID_WEIGHTS, - Feature.COARSE_0_ATOMIC_COORDS, - } - - if nuc_grad_feats is None: # generate feature list from functional features - nuc_grad_feats = set(functional.features) - - # Integer-valued features have no nuclear gradient — always discard them - nuc_grad_feats.discard(Feature.ATOMIC_GRID_SIZES) - nuc_grad_feats.discard(Feature.ATOMIC_GRID_SIZE_BOUND_SHAPE) - - # check for unsupported features - unsupported_feats = {feat for feat in nuc_grad_feats if feat not in SUPPORTED_FEATS} - if unsupported_feats != set(): - raise NotImplementedError( - f"Not supported features for nuclear gradient: {unsupported_feats}" - ) - - LOG.debug("nuc_grad_feats = %s", nuc_grad_feats) - - # determine the maximum ao derivative needed - if Feature.GRAD in nuc_grad_feats or Feature.KIN in nuc_grad_feats: - ao_deriv = 2 - elif Feature.DENSITY in nuc_grad_feats: - ao_deriv = 1 - else: - ao_deriv = 0 - - # Get the derivatives of the weights per atom and make sure the grid is blocked per atom, no padding, etc. coord_list = [] weight_list = [] for coords, weight in grids_noresponse_cc(grid): @@ -86,79 +49,31 @@ def veff_and_expl_nuc_grad( grid_ = grid.copy() grid_.coords = cp.concatenate(coord_list) grid_.weights = cp.concatenate(weight_list) - mol_feats = feature.generate_features( - mol, rdm1, grid_, set(functional.features), gpu=True + ao_deriv, derivatives = evaluate_nuclear_feature_derivatives( + functional, + mol, + grid_, + rdm1, + nuc_grad_feats, + max_memory_in_mb=max_memory_in_mb, ) - # Discard atomic_grid_weights from VJP features: d(atomic_grid_weights)/dR = 0 - # because they are raw quadrature weights that depend only on the radial/angular - # grid rule, not on nuclear positions. They still pass through as other_feats. - nuc_grad_feats.discard(Feature.ATOMIC_GRID_WEIGHTS) - - # Get required derivatives - nuc_feats = {feat: mol_feats[feat] for feat in nuc_grad_feats} - other_feats = { - feat: mol_feats[feat] for feat in mol_feats if feat not in nuc_grad_feats - } - - def exc_feat_func(differentiable_features: FeatureMap) -> torch.Tensor: - return functional.get_exc(differentiable_features | other_feats) - - dExc = feature_derivatives(exc_feat_func, nuc_feats) - - LOG.debug("autograd gradients for nuclear features done") - - nao = rdm1.shape[-1] - veff = torch.zeros((2, 3, nao, nao), dtype=rdm1.dtype, device=rdm1.device) - nuc_grad = torch.zeros((mol.natm, 3), dtype=rdm1.dtype, device=rdm1.device) - - atm_start = 0 - for atm_id, (coords, weight, weight1) in enumerate(grids_response_cc(grid)): - mask = dft.gen_grid.make_mask(mol, coords) - ao = from_dlpack( - dft.numint.eval_ao( - mol, - coords, - deriv=ao_deriv, - non0tab=mask, # cutoff=grid.cutoff + def atom_grid_blocks() -> Iterator[tuple[torch.Tensor, int, torch.Tensor]]: + for coords, weight, weight1 in grids_response_cc(grid): + mask = dft.gen_grid.make_mask(mol, coords) + ao = from_dlpack( + dft.numint.eval_ao( + mol, + coords, + deriv=ao_deriv, + non0tab=mask, + ) ) - ) - if ao_deriv == 0: - ao = ao[None, ...] - atm_end = atm_start + weight.shape[0] - dExc_atm = grid_derivative_block(dExc, atm_start, atm_end) - - # Calculate the contribution to veff for this atomic grid - veff_atm = contract_ao_derivative_block(ao, dExc_atm) - - if Feature.GRID_COORDS in dExc_atm: - # also add the explicit grid coordinate dependence - nuc_grad[atm_id] += dExc_atm[Feature.GRID_COORDS].sum(dim=0) - - if Feature.GRID_WEIGHTS in dExc_atm: - Exc_dgw = dExc_atm[Feature.GRID_WEIGHTS] - nuc_grad += from_dlpack(weight1) @ Exc_dgw - # add the grid coordinate dependence via the density-like quantities to the nuclear gradient - # we get those from the veff block. This tends to largely cancel with the grid_weights derivative, - # so that's why we include it here. - if len(rdm1.shape) == 2: - nuc_grad[atm_id] += torch.einsum("sxpq,qp->x", veff_atm, rdm1) - else: - nuc_grad[atm_id] += torch.einsum("sxpq,sqp->x", veff_atm, rdm1) * 2 - - veff += veff_atm - atm_start = atm_end - - if Feature.COARSE_0_ATOMIC_COORDS in nuc_grad_feats: - nuc_grad += dExc[Feature.COARSE_0_ATOMIC_COORDS] - - # finalize - if len(rdm1.shape) == 2: - veff = veff.sum(0) / 2 - - LOG.debug("veff and explicit components for nuclear gradient calculated") + if ao_deriv == 0: + ao = ao[None, ...] + yield ao, weight.shape[0], from_dlpack(weight1) - return -veff, nuc_grad + return assemble_nuclear_gradient(derivatives, rdm1, mol.natm, atom_grid_blocks()) def nuc_grad_from_veff( @@ -208,12 +123,13 @@ def energy_ee( if dm is None: dm = self.base.make_rdm1() - veff, self.veff_nuc_grad_ = veff_and_expl_nuc_grad( + veff, self.veff_nuc_grad_ = _veff_and_expl_nuc_grad( self.functional, mol=mol, grid=self.grids, rdm1=from_dlpack(dm), nuc_grad_feats=self.nuc_grad_feats, + max_memory_in_mb=int(self.base.max_memory), ) veff_grad = ( 2 * nuc_grad_from_veff(mol, veff, from_dlpack(dm)).detach().cpu().numpy() @@ -240,12 +156,13 @@ def grad_elec( if self.veff_nuc_grad_ is None: dm = self.base.make_rdm1() - _, self.veff_nuc_grad_ = veff_and_expl_nuc_grad( + _, self.veff_nuc_grad_ = _veff_and_expl_nuc_grad( self.functional, mol=self.mol, grid=self.grids, rdm1=from_dlpack(dm), nuc_grad_feats=self.nuc_grad_feats, + max_memory_in_mb=int(self.base.max_memory), ) veff_nuc_grad = self.veff_nuc_grad_ if veff_nuc_grad is None: @@ -314,12 +231,13 @@ def energy_ee( if dm is None: dm = self.base.make_rdm1() - veff, self.veff_nuc_grad_ = veff_and_expl_nuc_grad( + veff, self.veff_nuc_grad_ = _veff_and_expl_nuc_grad( self.functional, mol=mol, grid=self.grids, rdm1=from_dlpack(dm), nuc_grad_feats=self.nuc_grad_feats, + max_memory_in_mb=int(self.base.max_memory), ) veff_grad = ( 2 * nuc_grad_from_veff(mol, veff, from_dlpack(dm)).detach().cpu().numpy() @@ -346,12 +264,13 @@ def grad_elec( if self.veff_nuc_grad_ is None: dm = self.base.make_rdm1() - _, self.veff_nuc_grad_ = veff_and_expl_nuc_grad( + _, self.veff_nuc_grad_ = _veff_and_expl_nuc_grad( self.functional, mol=self.mol, grid=self.grids, rdm1=from_dlpack(dm), nuc_grad_feats=self.nuc_grad_feats, + max_memory_in_mb=int(self.base.max_memory), ) veff_nuc_grad = self.veff_nuc_grad_ if veff_nuc_grad is None: diff --git a/skala/src/skala/pyscf/ao_evaluation.py b/skala/src/skala/pyscf/ao_evaluation.py index bbb581f1..ccd2b5ac 100644 --- a/skala/src/skala/pyscf/ao_evaluation.py +++ b/skala/src/skala/pyscf/ao_evaluation.py @@ -16,7 +16,6 @@ from torch.utils.dlpack import from_dlpack from pyscf import dft, gto -from skala.features import FeatureMap from skala.pyscf import feature_math from skala.pyscf.backend import ( Array, @@ -575,7 +574,7 @@ def _resolve_ao_block_size( return result -def auto_chunk( +def evaluate_raw_features_auto_chunk( dm: torch.Tensor, mol: gto.Mole, grids: Grid, @@ -583,8 +582,8 @@ def auto_chunk( block_size: int | None = None, max_memory: int = 2000, gpu: bool = False, -) -> FeatureMap: - """Evaluate raw features with a memory-derived or explicit AO block size.""" +) -> Tensor: + """Evaluate packed raw features with a memory-derived AO block size.""" if gpu: check_gpu_imports_were_successful() if dm.device.type != "cuda": @@ -598,4 +597,4 @@ def auto_chunk( features = evaluate_ao_features_blockwise( dm.double(), mol, grids, feature_function, blksize ) - return feature_function.to_dict(features) + return features diff --git a/skala/src/skala/pyscf/evaluation.py b/skala/src/skala/pyscf/evaluation.py index 3c1e7ca6..e8f02850 100644 --- a/skala/src/skala/pyscf/evaluation.py +++ b/skala/src/skala/pyscf/evaluation.py @@ -2,19 +2,13 @@ """Feature requirements and numerical-evaluation policy.""" -from collections.abc import Iterable +from __future__ import annotations + +from collections.abc import Iterable, Iterator from dataclasses import dataclass -from skala.features import Feature +from skala.features import AO_FEATURES, AOFeatureSpec, Feature -_AO_FEATURES = frozenset( - { - Feature.DENSITY, - Feature.GRAD, - Feature.KIN, - Feature.LAPL, - } -) _ATOMIC_LAYOUT_FEATURES = frozenset( { Feature.ATOMIC_GRID_WEIGHTS, @@ -25,76 +19,51 @@ class FeatureSpec: - """Normalized feature names and their evaluation requirements.""" + """Normalized set of named molecular features.""" - def __init__(self, names: Iterable[Feature]) -> None: - self._names = frozenset(names) + def __init__(self, features: Iterable[Feature]) -> None: + self._features = frozenset(features) + ao_features = self._features & AO_FEATURES + self._ao_features = AOFeatureSpec(ao_features) if ao_features else None - @property - def names(self) -> frozenset[Feature]: - """Return the normalized feature names.""" - return self._names + def __iter__(self) -> Iterator[Feature]: + return iter(self._features) def __eq__(self, other: object) -> bool: if not isinstance(other, FeatureSpec): return NotImplemented - return self.names == other.names + return self._features == other._features def __hash__(self) -> int: - return hash(self.names) + return hash(self._features) def requests(self, feature: Feature) -> bool: """Return whether a feature is requested.""" - return feature in self.names - - @property - def with_density(self) -> bool: - """Return whether density is requested.""" - return self.requests(Feature.DENSITY) + return feature in self._features - @property - def with_grad(self) -> bool: - """Return whether the density gradient is requested.""" - return self.requests(Feature.GRAD) - - @property - def with_kin(self) -> bool: - """Return whether kinetic-energy density is requested.""" - return self.requests(Feature.KIN) + def __or__(self, other: FeatureSpec | Iterable[Feature]) -> FeatureSpec: + """Return a feature specification containing both operands.""" + return FeatureSpec(self._features | frozenset(other)) @property - def with_lapl(self) -> bool: - """Return whether the density Laplacian is requested.""" - return self.requests(Feature.LAPL) + def ao_features(self) -> AOFeatureSpec | None: + """Return the requested AO-derived features.""" + return self._ao_features @property def requires_ao_evaluation(self) -> bool: """Return whether AO-derived features are requested.""" - return bool(self.names & _AO_FEATURES) - - @property - def mgga_feature_count(self) -> int: - """Return the scalar width of the requested meta-GGA features.""" - return self.with_density + 3 * self.with_grad + self.with_kin + self.with_lapl - - @property - def ao_derivative_order(self) -> int: - """Return the highest AO derivative order needed by the features.""" - if Feature.LAPL in self.names: - return 2 - if self.names & {Feature.GRAD, Feature.KIN}: - return 1 - return 0 + return self._ao_features is not None @property def requires_atomic_layout(self) -> bool: """Return whether grid points must retain per-atom ordering.""" - return bool(self.names & _ATOMIC_LAYOUT_FEATURES) + return bool(self._features & _ATOMIC_LAYOUT_FEATURES) @property def supports_spatial_decomposition(self) -> bool: """Return whether spatial decomposition is supported.""" - return Feature.ATOMIC_GRID_SIZES in self.names + return Feature.ATOMIC_GRID_SIZES in self._features @dataclass(frozen=True) diff --git a/skala/src/skala/pyscf/feature_math.py b/skala/src/skala/pyscf/feature_math.py index deb8a223..cca32ef7 100644 --- a/skala/src/skala/pyscf/feature_math.py +++ b/skala/src/skala/pyscf/feature_math.py @@ -3,12 +3,148 @@ """Raw density-feature mathematics and model formatting.""" from abc import ABC, abstractmethod +from collections.abc import Callable, Iterable, Iterator +from enum import IntEnum +from typing import ClassVar import torch from torch import nn -from skala.features import Feature, FeatureMap -from skala.pyscf.evaluation import FeatureSpec +from skala.features import AOFeatureSpec, Feature, FeatureMap + + +def _disconnected_features_error(features: Iterable[Feature]) -> RuntimeError: + feature_names = ", ".join(sorted(feature.value for feature in features)) + return RuntimeError( + f"XC energy is disconnected from requested features: {feature_names}" + ) + + +def feature_derivatives( + exc_func: Callable[[FeatureMap], torch.Tensor], features: FeatureMap +) -> FeatureMap: + """Differentiate a scalar XC energy with respect to molecular features. + + Ordinary autograd tensors are used instead of ``torch.func.vjp`` functional + tensors because traced TorchScript models may require accessible backing + storage. + + Args: + exc_func: Callable accepting the differentiable features and returning + scalar XC energy. + features: Molecular features to differentiate, keyed by feature name. + + Returns: + XC energy derivatives keyed by feature name. + + Raises: + RuntimeError: If the XC energy is disconnected from a requested feature. + """ + if not features: + return {} + + differentiable_features = { + feature: tensor.detach().requires_grad_(True) + for feature, tensor in features.items() + } + exc = exc_func(differentiable_features) + if not exc.requires_grad: + raise _disconnected_features_error(differentiable_features) + + gradients = torch.autograd.grad( + exc, + tuple(differentiable_features.values()), + create_graph=False, + retain_graph=False, + allow_unused=True, + ) + derivatives: FeatureMap = { + feature: gradient.detach() + for feature, gradient in zip(differentiable_features, gradients, strict=True) + if gradient is not None + } + if len(derivatives) != len(differentiable_features): + disconnected_features = differentiable_features.keys() - derivatives.keys() + raise _disconnected_features_error(disconnected_features) + return derivatives + + +class AODirection(IntEnum): + """Cartesian direction indices used by feature and potential tensors.""" + + X = 0 + Y = 1 + Z = 2 + + +class PackedAO: + """Component-axis views over packed PySCF AO derivative values. + + PySCF stores AO values and Cartesian derivatives along the leading tensor + dimension in the order ``value, x, y, z, xx, xy, xz, yy, yz, zz``. This + wrapper normalizes rank-two value-only input by inserting a leading + component dimension. For packed input, it interprets only the leading + dimension and leaves all trailing dimensions unchanged. It therefore + supports both the ``(component, nao, ngrid)`` layout used by density-feature + evaluation and the ``(component, ngrid, nao)`` layout used by + nuclear-gradient contractions. + + Valid packed component counts are one for AO values, four through first + derivatives, and ten through second derivatives. Accessors require the + corresponding components to be available: :attr:`gradient` requires first + derivatives, while :attr:`diagonal_hessian` and :meth:`hessian` require + second derivatives. + + Args: + tensor: AO values, either unpacked with rank two or with PySCF derivative + components on the leading axis. + + Raises: + ValueError: If a packed tensor has an unsupported component count. + """ + + # PySCF's ``eval_ao(..., deriv=2)`` order is + # value, x, y, z, xx, xy, xz, yy, yz, zz. Mixed derivatives are reused + # across the symmetric off-diagonal entries of the Cartesian Hessian. + _HESSIAN_COMPONENTS: ClassVar[tuple[tuple[int, int, int], ...]] = ( + (4, 5, 6), + (5, 7, 8), + (6, 8, 9), + ) + + def __init__(self, tensor: torch.Tensor) -> None: + if tensor.dim() == 2: + tensor = tensor.unsqueeze(0) + + ncomponents = tensor.shape[0] + if ncomponents not in (1, 4, 10): + raise ValueError( + "Packed AO values must contain 1, 4, or 10 derivative components; " + f"got {ncomponents}." + ) + self._tensor = tensor + + @property + def value(self) -> torch.Tensor: + """Return the AO values.""" + return self._tensor[0] + + @property + def gradient(self) -> torch.Tensor: + """Return the three Cartesian AO gradients.""" + return self._tensor[1:4] + + @property + def diagonal_hessian(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return the three diagonal Cartesian AO Hessian components.""" + return self._tensor[4], self._tensor[7], self._tensor[9] + + def hessian(self) -> Iterator[tuple[AODirection, AODirection, torch.Tensor]]: + """Yield both Cartesian directions and each AO Hessian component.""" + for ao_direction in AODirection: + for feature_direction in AODirection: + component = self._HESSIAN_COMPONENTS[ao_direction][feature_direction] + yield ao_direction, feature_direction, self._tensor[component] def maybe_expand_and_divide( @@ -40,33 +176,19 @@ def to_dict(self, features: torch.Tensor) -> FeatureMap: ... class MGGAFeatureFunction(LinearFeature): """Evaluate the requested linear meta-GGA density features.""" - def __init__(self, feature_spec: FeatureSpec) -> None: + def __init__(self, feature_spec: AOFeatureSpec) -> None: super().__init__() - if not feature_spec.requires_ao_evaluation: - raise ValueError("At least one AO-derived feature must be selected.") - self.feature_spec = feature_spec - self.deriv = feature_spec.ao_derivative_order - self.nfeats = feature_spec.mgga_feature_count + self._feature_spec = feature_spec + self.deriv = feature_spec.nderiv + self.nfeats = feature_spec.nfeats def to_dict(self, features: torch.Tensor) -> FeatureMap: """Convert a packed feature tensor to its named feature tensors.""" - feature_index = 0 - feature_dict: FeatureMap = {} - if self.feature_spec.with_density: - feature_dict[Feature.DENSITY] = features[..., feature_index, :] - feature_index += 1 - if self.feature_spec.with_grad: - feature_dict[Feature.GRAD] = features[ - ..., feature_index : feature_index + 3, : - ] - feature_index += 3 - if self.feature_spec.with_kin: - feature_dict[Feature.KIN] = features[..., feature_index, :] - feature_index += 1 - if self.feature_spec.with_lapl: - feature_dict[Feature.LAPL] = features[..., feature_index, :] - return feature_dict + return { + feature: features[..., feature_slice, :].squeeze(-2) + for feature, feature_slice in self._feature_spec + } def forward(self, dm: torch.Tensor, ao: torch.Tensor) -> torch.Tensor: dm_view = dm.view(-1, dm.shape[-2], dm.shape[-1]) @@ -78,50 +200,36 @@ def forward(self, dm: torch.Tensor, ao: torch.Tensor) -> torch.Tensor: dtype=dm.dtype, ) - if self.deriv == 0: - c0 = dm_view @ ao - features[..., 0, :] = torch.sum(c0 * ao[None, :, :], dim=-2) - if len(dm.shape) == 2: - return features.reshape((self.nfeats, -1)) - return features.reshape((*dm.shape[:-2], self.nfeats, -1)) - - c0 = dm_view @ ao[0] - - feature_index = 0 - if self.feature_spec.with_density: - features[..., feature_index, :] = torch.sum(c0 * ao[0, None, :, :], dim=-2) - feature_index += 1 - - if self.feature_spec.with_grad: - for component in range(3): - features[..., feature_index, :] = 2 * torch.sum( - c0 * ao[component + 1, None, :, :], dim=-2 + packed_ao = PackedAO(ao) + phi = packed_ao.value + c0 = dm_view @ phi + gradient_contraction = None + if Feature.KIN in self._feature_spec or Feature.LAPL in self._feature_spec: + gradient_contraction = features.new_zeros((dm_view.shape[0], ao.shape[-1])) + for ao_gradient in packed_ao.gradient: + ci = dm_view @ ao_gradient + gradient_contraction += torch.sum(ci * ao_gradient[None], dim=-2) + + for feature, feature_slice in self._feature_spec: + if feature == Feature.DENSITY: + feature_values = torch.sum(c0 * phi[None], dim=-2).unsqueeze(-2) + elif feature == Feature.GRAD: + feature_values = 2 * torch.sum( + c0[:, None] * packed_ao.gradient[None], dim=-2 ) - feature_index += 1 - - if self.feature_spec.with_kin or self.feature_spec.with_lapl: - for component in range(3): - ci = dm_view @ ao[component + 1] - features[..., feature_index, :] += 0.5 * torch.sum( - ci * ao[component + 1, None, :, :], dim=-2 - ) - - if self.feature_spec.with_kin: - feature_index += 1 - if self.feature_spec.with_lapl: - features[..., feature_index, :] = ( - 4 * features[..., feature_index - 1, :] - ) + elif feature == Feature.KIN: + assert gradient_contraction is not None + feature_values = (0.5 * gradient_contraction).unsqueeze(-2) else: - features[..., feature_index, :] *= 4.0 - - if self.feature_spec.with_lapl: - for component in (4, 7, 9): - features[..., feature_index, :] += 2 * torch.sum( - c0 * ao[component, None, :, :], dim=-2 - ) - - if len(dm.shape) == 2: + assert feature == Feature.LAPL + assert gradient_contraction is not None + laplacian = 2 * gradient_contraction + for ao_hessian in packed_ao.diagonal_hessian: + laplacian += 2 * torch.sum(c0 * ao_hessian[None], dim=-2) + feature_values = laplacian.unsqueeze(-2) + features[..., feature_slice, :] = feature_values + + if dm.dim() == 2: return features.reshape((self.nfeats, -1)) return features.reshape((*dm.shape[:-2], self.nfeats, -1)) @@ -130,44 +238,33 @@ def vjp(self, ao: torch.Tensor, cotangent: torch.Tensor) -> torch.Tensor: batch_shape = cotangent.shape[:-2] ngrids = cotangent.shape[-1] weights = cotangent.reshape(-1, self.nfeats, ngrids) - phi = ao if self.deriv == 0 else ao[0] + packed_ao = PackedAO(ao) + phi = packed_ao.value nao = phi.shape[-2] - if self.deriv == 0: - result = (weights[:, 0, None, :] * phi) @ phi.transpose(-1, -2) - return result.reshape(*batch_shape, nao, nao) - left = weights.new_zeros((weights.shape[0], nao, ngrids)) - feature_index = 0 - if self.feature_spec.with_density: - left += weights[:, feature_index, None, :] * phi - feature_index += 1 - - if self.feature_spec.with_grad: - for component in range(3): - left.addcmul_( - weights[:, feature_index + component, None, :], - ao[component + 1], - value=2, + derivative_weight = weights.new_zeros((weights.shape[0], 1, ngrids)) + for feature, feature_slice in self._feature_spec: + feature_weight = weights[:, feature_slice, :] + if feature == Feature.DENSITY: + left += feature_weight * phi + elif feature == Feature.GRAD: + left += 2 * torch.sum( + feature_weight[:, :, None] * packed_ao.gradient[None], dim=1 ) - feature_index += 3 - - derivative_weight = weights.new_zeros((weights.shape[0], ngrids)) - if self.feature_spec.with_kin: - derivative_weight += 0.5 * weights[:, feature_index] - feature_index += 1 - - if self.feature_spec.with_lapl: - laplacian_weight = weights[:, feature_index] - derivative_weight += 2 * laplacian_weight - for component in (4, 7, 9): - left.addcmul_(laplacian_weight[:, None, :], ao[component], value=2) + elif feature == Feature.KIN: + derivative_weight += 0.5 * feature_weight + else: + assert feature == Feature.LAPL + derivative_weight += 2 * feature_weight + for ao_hessian in packed_ao.diagonal_hessian: + left.addcmul_(feature_weight, ao_hessian, value=2) result = left @ phi.transpose(-1, -2) - if self.feature_spec.with_kin or self.feature_spec.with_lapl: - for component in range(1, 4): - weighted_derivative = derivative_weight[:, None, :] * ao[component] - result += weighted_derivative @ ao[component].transpose(-1, -2) + if Feature.KIN in self._feature_spec or Feature.LAPL in self._feature_spec: + for ao_gradient in packed_ao.gradient: + weighted_derivative = derivative_weight * ao_gradient + result += weighted_derivative @ ao_gradient.transpose(-1, -2) result = 0.5 * (result + result.transpose(-1, -2)) return result.reshape(*batch_shape, nao, nao) diff --git a/skala/src/skala/pyscf/features.py b/skala/src/skala/pyscf/features.py index 40b64f6a..68ae1250 100644 --- a/skala/src/skala/pyscf/features.py +++ b/skala/src/skala/pyscf/features.py @@ -10,83 +10,8 @@ from pyscf import gto from skala.features import Feature, FeatureMap -from skala.pyscf import ao_evaluation, feature_math from skala.pyscf.backend import Grid, from_numpy_or_cupy -from skala.pyscf.evaluation import EvaluationPolicy, FeatureSpec - -DEFAULT_FEATURES = [ - Feature.DENSITY, - Feature.KIN, - Feature.GRAD, - Feature.GRID_COORDS, - Feature.GRID_WEIGHTS, -] -DEFAULT_FEATURES_SET = set(DEFAULT_FEATURES) - - -def generate_features( - mol: gto.Mole, - dm: Tensor, - grids: Grid, - features: set[Feature] | None = None, - chunk_size: int | None = None, - max_memory: int = 2000, - gpu: bool = False, -) -> FeatureMap: - """Generate density features for a given molecule. The density features are stored in a dictionary - with the keys matching the requested features. - - Parameters - ---------- - mol: gto.Mole - the molecule - dm: Tensor - the density matrix - grids: Grid - the grid - features: set[str] | None - the requested features - chunk_size: int | None - a manually specified chunk size for processing the grids, if None the chunk size is determined automatically - max_memory: int - the maximum memory to use for calculating the features - gpu: bool - whether to use the GPU(4pyscf) for calculations - - Returns - ------- - dict[str, Tensor] - A dictionary containing the requested features. The keys are the feature names, - and the values are the corresponding tensors. - """ - feature_spec = FeatureSpec(DEFAULT_FEATURES_SET if features is None else features) - evaluation_policy = EvaluationPolicy(ao_block_size=chunk_size) - - # if dm is a 3D tensor, then we have a spin-polarized system - is_spin_polarized = len(dm.shape) == 3 - - if gpu and dm.device.type != "cuda": - raise ValueError("Density matrix must be on the GPU when gpu=True.") - - mol_features = get_grid_features(mol, dm, grids, feature_spec) - - if feature_spec.requires_ao_evaluation: - mgga_features = ao_evaluation.auto_chunk( - dm, - mol, - grids, - feature_math.MGGAFeatureFunction(feature_spec), - block_size=evaluation_policy.ao_block_size, - max_memory=max_memory, - gpu=gpu, - ) - - for feature in mgga_features: - mol_features[feature] = feature_math.maybe_expand_and_divide( - mgga_features[feature], not is_spin_polarized, 2 - ) - - return mol_features +from skala.pyscf.evaluation import FeatureSpec def get_grid_features( diff --git a/skala/src/skala/pyscf/gradient_core.py b/skala/src/skala/pyscf/gradient_core.py index 48ea2a22..4ab0a888 100644 --- a/skala/src/skala/pyscf/gradient_core.py +++ b/skala/src/skala/pyscf/gradient_core.py @@ -2,118 +2,40 @@ """Backend-independent PyTorch operations for PySCF nuclear gradients.""" -from collections.abc import Callable, Iterable, Iterator -from dataclasses import dataclass -from enum import IntEnum +from collections.abc import Iterable from types import EllipsisType -from typing import ClassVar, TypeAlias +from typing import TypeAlias import torch +from torch import Tensor +from pyscf import gto from skala.features import Feature, FeatureMap - - -class _Direction(IntEnum): - """Cartesian direction indices used by feature and potential tensors.""" - - X = 0 - Y = 1 - Z = 2 - - -@dataclass(frozen=True) -class _PackedAO: - """Views into PySCF's packed AO derivative dimension.""" - - tensor: torch.Tensor - - # PySCF's ``eval_ao(..., deriv=2)`` order is - # value, x, y, z, xx, xy, xz, yy, yz, zz. Mixed derivatives are reused - # across the symmetric off-diagonal entries of the Cartesian Hessian. - _HESSIAN_COMPONENTS: ClassVar[tuple[tuple[int, int, int], ...]] = ( - (4, 5, 6), - (5, 7, 8), - (6, 8, 9), - ) - - @property - def value(self) -> torch.Tensor: - """AO values with shape ``(npoints, nao)``.""" - return self.tensor[0] - - @property - def gradient(self) -> torch.Tensor: - """AO gradients with shape ``(3, npoints, nao)``.""" - return self.tensor[1:4] - - def hessian(self) -> Iterator[tuple[_Direction, _Direction, torch.Tensor]]: - """Yield both Cartesian directions and each AO Hessian component.""" - for ao_direction in _Direction: - for feature_direction in _Direction: - component = self._HESSIAN_COMPONENTS[ao_direction][feature_direction] - yield ao_direction, feature_direction, self.tensor[component] +from skala.functional.base import ExcFunctionalBase +from skala.pyscf.backend import Grid +from skala.pyscf.evaluation import FeatureSpec +from skala.pyscf.feature_math import PackedAO +from skala.pyscf.model_chunking import ( + evaluate_chunked_feature_gradients, + evaluate_model_features, +) + +SUPPORTED_NUCLEAR_GRADIENT_FEATURES = frozenset( + { + Feature.DENSITY, + Feature.GRAD, + Feature.KIN, + Feature.GRID_COORDS, + Feature.GRID_WEIGHTS, + Feature.ATOMIC_GRID_WEIGHTS, + Feature.COARSE_0_ATOMIC_COORDS, + } +) _FeatureSlice: TypeAlias = tuple[slice] | tuple[EllipsisType, slice] -def _disconnected_features_error(features: Iterable[Feature]) -> RuntimeError: - feature_names = ", ".join(sorted(feature.value for feature in features)) - return RuntimeError( - f"XC energy is disconnected from requested features: {feature_names}" - ) - - -def feature_derivatives( - exc_func: Callable[[FeatureMap], torch.Tensor], features: FeatureMap -) -> FeatureMap: - """Differentiate a scalar XC energy with respect to molecular features. - - Ordinary autograd tensors are used instead of ``torch.func.vjp`` functional - tensors because traced TorchScript models may require accessible backing - storage. - - Args: - exc_func: Callable accepting the differentiable features and returning - scalar XC energy. - features: Molecular features to differentiate, keyed by feature name. - - Returns: - XC energy derivatives keyed by feature name. - - Raises: - RuntimeError: If the XC energy is disconnected from a requested feature. - - """ - if not features: - return {} - - differentiable_features = { - feature: tensor.detach().requires_grad_(True) - for feature, tensor in features.items() - } - exc = exc_func(differentiable_features) - if not exc.requires_grad: - raise _disconnected_features_error(differentiable_features) - - gradients = torch.autograd.grad( - exc, - tuple(differentiable_features.values()), - create_graph=False, - retain_graph=False, - allow_unused=True, - ) - derivatives: FeatureMap = { - feature: gradient.detach() - for feature, gradient in zip(differentiable_features, gradients, strict=True) - if gradient is not None - } - if len(derivatives) != len(differentiable_features): - disconnected_features = differentiable_features.keys() - derivatives.keys() - raise _disconnected_features_error(disconnected_features) - return derivatives - - def grid_derivative_block( derivatives: FeatureMap, grid_start: int, grid_end: int ) -> FeatureMap: @@ -146,6 +68,117 @@ def grid_derivative_block( } +def evaluate_nuclear_feature_derivatives( + functional: ExcFunctionalBase, + mol: gto.Mole, + grid: Grid, + rdm1: Tensor, + features: set[Feature] | None = None, + *, + max_memory_in_mb: int, +) -> tuple[int, FeatureMap]: + """Evaluate model derivatives needed for a nuclear gradient. + + Args: + functional: XC functional whose energy is differentiated. + mol: Molecule associated with the density matrix and grid. + grid: Atom-major integration grid used for feature evaluation. + rdm1: One-particle density matrix. + features: Features to differentiate, defaulting to the functional inputs. + max_memory_in_mb: Memory limit for AO evaluation and CPU model-gradient + chunking. CUDA model-gradient chunking probes available device memory. + + Returns: + Required AO derivative order and XC derivatives by feature. + + Raises: + NotImplementedError: If a requested feature has no nuclear-gradient rule. + """ + differentiable_features = set(functional.features if features is None else features) + differentiable_features -= { + Feature.ATOMIC_GRID_SIZES, + Feature.ATOMIC_GRID_SIZE_BOUND_SHAPE, + } + unsupported = differentiable_features - SUPPORTED_NUCLEAR_GRADIENT_FEATURES + if unsupported: + raise NotImplementedError( + f"Not supported features for nuclear gradient: {unsupported}" + ) + + ao_feature_spec = FeatureSpec(differentiable_features).ao_features + # Nuclear gradients need one order beyond feature evaluation; non-AO paths need none. + ao_deriv = 0 if ao_feature_spec is None else ao_feature_spec.nderiv + 1 + differentiable_features.discard(Feature.ATOMIC_GRID_WEIGHTS) + model_features = evaluate_model_features( + mol, + rdm1, + grid, + FeatureSpec((*functional.features, Feature.ATOMIC_GRID_SIZES)), + max_memory_in_mb=max_memory_in_mb, + ) + derivatives = evaluate_chunked_feature_gradients( + functional, + rdm1, + model_features, + differentiable_features, + max_memory_in_mb=max_memory_in_mb if rdm1.device.type == "cpu" else None, + ) + return ao_deriv, derivatives + + +def assemble_nuclear_gradient( + derivatives: FeatureMap, + rdm1: Tensor, + natm: int, + atom_grid_blocks: Iterable[tuple[Tensor, int, Tensor]], +) -> tuple[Tensor, Tensor]: + """Assemble effective-potential and explicit nuclear-gradient contributions. + + Args: + derivatives: XC energy derivatives keyed by molecular feature. + rdm1: One-particle density matrix. + natm: Number of atoms in the molecule. + atom_grid_blocks: Tuples containing AO values, grid-point count, and + grid-weight derivatives for each atom. + + Returns: + Negative effective potential and explicit nuclear gradient. + """ + nao = rdm1.shape[-1] + veff = rdm1.new_zeros((2, 3, nao, nao)) + nuclear_gradient = rdm1.new_zeros((natm, 3)) + + grid_start = 0 + for atom, (ao, grid_size, weight_derivatives) in enumerate(atom_grid_blocks): + grid_end = grid_start + grid_size + atom_derivatives = grid_derivative_block(derivatives, grid_start, grid_end) + atom_veff = contract_ao_derivative_block(ao, atom_derivatives) + + if Feature.GRID_COORDS in atom_derivatives: + nuclear_gradient[atom] += atom_derivatives[Feature.GRID_COORDS].sum(dim=0) + + if Feature.GRID_WEIGHTS in atom_derivatives: + grid_weight_derivative = atom_derivatives[Feature.GRID_WEIGHTS] + nuclear_gradient += weight_derivatives @ grid_weight_derivative + if rdm1.ndim == 2: + nuclear_gradient[atom] += torch.einsum("sxpq,qp->x", atom_veff, rdm1) + else: + nuclear_gradient[atom] += ( + torch.einsum("sxpq,sqp->x", atom_veff, rdm1) * 2 + ) + + veff += atom_veff + grid_start = grid_end + + if Feature.COARSE_0_ATOMIC_COORDS in derivatives: + nuclear_gradient += derivatives[Feature.COARSE_0_ATOMIC_COORDS] + + if rdm1.ndim == 2: + veff = veff.sum(0) / 2 + + return -veff, nuclear_gradient + + def contract_ao_derivative_block( ao: torch.Tensor, feature_derivatives: FeatureMap ) -> torch.Tensor: @@ -174,45 +207,42 @@ def contract_ao_derivative_block( indexed by spin, Cartesian direction, and the two AO indices. The result has the same device and dtype as ``ao``. """ - packed_ao = _PackedAO(ao) + packed_ao = PackedAO(ao) nao = ao.shape[-1] potential_derivatives = ao.new_zeros((2, 3, nao, nao)) - if Feature.DENSITY in feature_derivatives: - potential_derivatives += torch.einsum( - "si, xip, iq -> sxpq", - feature_derivatives[Feature.DENSITY], - packed_ao.gradient, - packed_ao.value, - ) - - if Feature.GRAD in feature_derivatives: - exc_dgrad = feature_derivatives[Feature.GRAD] - potential_derivatives += torch.einsum( - "syi, xip, yiq -> sxpq", - exc_dgrad, - packed_ao.gradient, - packed_ao.gradient, - ) - for ao_direction, feature_direction, ao_hessian in packed_ao.hessian(): - potential_derivatives[:, ao_direction] += torch.einsum( - "si, ip, iq -> spq", - exc_dgrad[:, feature_direction], - ao_hessian, + for feature, feature_derivative in feature_derivatives.items(): + if feature == Feature.DENSITY: + potential_derivatives += torch.einsum( + "si, xip, iq -> sxpq", + feature_derivative, + packed_ao.gradient, packed_ao.value, ) - - if Feature.KIN in feature_derivatives: - exc_dkin = feature_derivatives[Feature.KIN] - for ao_direction, feature_direction, ao_hessian in packed_ao.hessian(): - potential_derivatives[:, ao_direction] += ( - torch.einsum( + elif feature == Feature.GRAD: + potential_derivatives += torch.einsum( + "syi, xip, yiq -> sxpq", + feature_derivative, + packed_ao.gradient, + packed_ao.gradient, + ) + for ao_direction, feature_direction, ao_hessian in packed_ao.hessian(): + potential_derivatives[:, ao_direction] += torch.einsum( "si, ip, iq -> spq", - exc_dkin, + feature_derivative[:, feature_direction], ao_hessian, - packed_ao.gradient[feature_direction], + packed_ao.value, + ) + elif feature == Feature.KIN: + for ao_direction, feature_direction, ao_hessian in packed_ao.hessian(): + potential_derivatives[:, ao_direction] += ( + torch.einsum( + "si, ip, iq -> spq", + feature_derivative, + ao_hessian, + packed_ao.gradient[feature_direction], + ) + / 2 ) - / 2 - ) return potential_derivatives diff --git a/skala/src/skala/pyscf/gradients.py b/skala/src/skala/pyscf/gradients.py index ed2d405f..f8ffacea 100644 --- a/skala/src/skala/pyscf/gradients.py +++ b/skala/src/skala/pyscf/gradients.py @@ -2,7 +2,7 @@ """Modification of PySCF nuclear gradient object to work with Skala functional.""" -import logging +from collections.abc import Iterator from typing import Any import numpy as np @@ -12,26 +12,24 @@ from pyscf.grad.uks import Gradients as UHFGradient from pyscf.scf.hf import SCF -import skala.pyscf.features as feature from pyscf import dft, gto from skala.dispersion import DFTD3Dispersion -from skala.features import Feature, FeatureMap +from skala.features import Feature from skala.functional.base import ExcFunctionalBase from skala.pyscf.gradient_core import ( - contract_ao_derivative_block, - feature_derivatives, - grid_derivative_block, + assemble_nuclear_gradient, + evaluate_nuclear_feature_derivatives, ) -LOG = logging.getLogger(__name__) - -def veff_and_expl_nuc_grad( +def _veff_and_expl_nuc_grad( functional: ExcFunctionalBase, mol: gto.Mole, grid: dft.Grids, rdm1: torch.Tensor, nuc_grad_feats: set[Feature] | None = None, + *, + max_memory_in_mb: int, ) -> tuple[torch.Tensor, torch.Tensor]: """ returns: @@ -39,41 +37,6 @@ def veff_and_expl_nuc_grad( - 2nd tuple argument: explicit contributions to the nuclear gradient """ - SUPPORTED_FEATS = { - Feature.DENSITY, - Feature.GRAD, - Feature.KIN, - Feature.GRID_COORDS, - Feature.GRID_WEIGHTS, - Feature.ATOMIC_GRID_WEIGHTS, - Feature.COARSE_0_ATOMIC_COORDS, - } - - if nuc_grad_feats is None: # generate feature list from functional features - nuc_grad_feats = set(functional.features) - - # Integer-valued features have no nuclear gradient — always discard them - nuc_grad_feats.discard(Feature.ATOMIC_GRID_SIZES) - nuc_grad_feats.discard(Feature.ATOMIC_GRID_SIZE_BOUND_SHAPE) - - # check for unsupported features - unsupported_feats = {feat for feat in nuc_grad_feats if feat not in SUPPORTED_FEATS} - if unsupported_feats != set(): - raise NotImplementedError( - f"Not supported features for nuclear gradient: {unsupported_feats}" - ) - - LOG.debug("nuc_grad_feats = %s", nuc_grad_feats) - - # determine the maximum ao derivative needed - if Feature.GRAD in nuc_grad_feats or Feature.KIN in nuc_grad_feats: - ao_deriv = 2 - elif Feature.DENSITY in nuc_grad_feats: - ao_deriv = 1 - else: - ao_deriv = 0 - - # Get the derivatives of the weights per atom and make sure the grid is blocked per atom, no padding, etc. coord_list = [] weight_list = [] for coords, weight in grids_noresponse_cc(grid): @@ -83,74 +46,28 @@ def veff_and_expl_nuc_grad( grid_ = grid.copy() grid_.coords = np.concatenate(coord_list) grid_.weights = np.concatenate(weight_list) - mol_feats = feature.generate_features(mol, rdm1, grid_, set(functional.features)) - - # Discard atomic_grid_weights from VJP features: d(atomic_grid_weights)/dR = 0 - # because they are raw quadrature weights that depend only on the radial/angular - # grid rule, not on nuclear positions. They still pass through as other_feats. - nuc_grad_feats.discard(Feature.ATOMIC_GRID_WEIGHTS) - - # Get required derivatives - nuc_feats = {feat: mol_feats[feat] for feat in nuc_grad_feats} - other_feats = { - feat: mol_feats[feat] for feat in mol_feats if feat not in nuc_grad_feats - } - - def exc_feat_func(differentiable_features: FeatureMap) -> torch.Tensor: - return functional.get_exc(differentiable_features | other_feats) - - dExc = feature_derivatives(exc_feat_func, nuc_feats) - - LOG.debug("autograd gradients for nuclear features done") - - nao = rdm1.shape[-1] - veff = torch.zeros((2, 3, nao, nao), dtype=rdm1.dtype) - nuc_grad = torch.zeros((mol.natm, 3), dtype=rdm1.dtype) - - atm_start = 0 - for atm_id, (coords, weight, weight1) in enumerate(grids_response_cc(grid)): - mask = dft.gen_grid.make_mask(mol, coords) - ao = torch.from_numpy( - dft.numint.eval_ao( - mol, coords, deriv=ao_deriv, non0tab=mask, cutoff=grid.cutoff + ao_deriv, derivatives = evaluate_nuclear_feature_derivatives( + functional, + mol, + grid_, + rdm1, + nuc_grad_feats, + max_memory_in_mb=max_memory_in_mb, + ) + + def atom_grid_blocks() -> Iterator[tuple[torch.Tensor, int, torch.Tensor]]: + for coords, weight, weight1 in grids_response_cc(grid): + mask = dft.gen_grid.make_mask(mol, coords) + ao = torch.from_numpy( + dft.numint.eval_ao( + mol, coords, deriv=ao_deriv, non0tab=mask, cutoff=grid.cutoff + ) ) - ) - if ao_deriv == 0: - ao = ao[None, ...] - atm_end = atm_start + weight.shape[0] - dExc_atm = grid_derivative_block(dExc, atm_start, atm_end) - - # Calculate the contribution to veff for this atomic grid - veff_atm = contract_ao_derivative_block(ao, dExc_atm) - - if Feature.GRID_COORDS in dExc_atm: - # also add the explicit grid coordinate dependence - nuc_grad[atm_id] += dExc_atm[Feature.GRID_COORDS].sum(dim=0) - - if Feature.GRID_WEIGHTS in dExc_atm: - Exc_dgw = dExc_atm[Feature.GRID_WEIGHTS] - nuc_grad += torch.from_numpy(weight1) @ Exc_dgw - # add the grid coordinate dependence via the density-like quantities to the nuclear gradient - # we get those from the veff block. This tends to largely cancel with the grid_weights derivative, - # so that's why we include it here. - if len(rdm1.shape) == 2: - nuc_grad[atm_id] += torch.einsum("sxpq,qp->x", veff_atm, rdm1) - else: - nuc_grad[atm_id] += torch.einsum("sxpq,sqp->x", veff_atm, rdm1) * 2 - - veff += veff_atm - atm_start = atm_end - - if Feature.COARSE_0_ATOMIC_COORDS in nuc_grad_feats: - nuc_grad += dExc[Feature.COARSE_0_ATOMIC_COORDS] - - # finalize - if len(rdm1.shape) == 2: - veff = veff.sum(0) / 2 - - LOG.debug("veff and explicit components for nuclear gradient calculated") + if ao_deriv == 0: + ao = ao[None, ...] + yield ao, weight.shape[0], torch.from_numpy(weight1) - return (-veff, nuc_grad) + return assemble_nuclear_gradient(derivatives, rdm1, mol.natm, atom_grid_blocks()) class SkalaRKSGradient(RHFGradient): # type: ignore[misc] @@ -186,12 +103,13 @@ def get_veff( if dm is None: dm = self.base.make_rdm1() - veff, self.veff_nuc_grad_ = veff_and_expl_nuc_grad( + veff, self.veff_nuc_grad_ = _veff_and_expl_nuc_grad( self.functional, mol=mol, grid=self.grids, rdm1=torch.from_numpy(dm), nuc_grad_feats=self.nuc_grad_feats, + max_memory_in_mb=int(self.base.max_memory), ) self.veff_nuc_grad_.detach_() result = veff.detach_().numpy() + self.get_j(mol, dm) @@ -269,12 +187,13 @@ def get_veff( if dm is None: dm = self.base.make_rdm1() - veff, self.veff_nuc_grad_ = veff_and_expl_nuc_grad( + veff, self.veff_nuc_grad_ = _veff_and_expl_nuc_grad( self.functional, mol=mol, grid=self.grids, rdm1=torch.from_numpy(dm), nuc_grad_feats=self.nuc_grad_feats, + max_memory_in_mb=int(self.base.max_memory), ) result = veff.detach_().numpy() + self.get_j(mol, dm).sum(0) assert isinstance(result, np.ndarray) diff --git a/skala/src/skala/pyscf/model_chunking.py b/skala/src/skala/pyscf/model_chunking.py index 85035080..c6e258d6 100644 --- a/skala/src/skala/pyscf/model_chunking.py +++ b/skala/src/skala/pyscf/model_chunking.py @@ -16,8 +16,10 @@ from pyscf import gto from skala.features import Feature, FeatureMap -from skala.pyscf import feature_math +from skala.functional.base import ExcFunctionalBase +from skala.pyscf import ao_evaluation, feature_math from skala.pyscf.backend import Grid +from skala.pyscf.evaluation import FeatureSpec from skala.pyscf.features import get_grid_features from skala.pyscf.memory_estimators import ( estimate_max_model_atoms_per_chunk, @@ -34,6 +36,43 @@ Feature.COARSE_0_ATOMIC_COORDS, Feature.ATOMIC_GRID_SIZES, ) +_AO_DERIVED_FEATURES = ( + Feature.DENSITY, + Feature.GRAD, + Feature.KIN, + Feature.LAPL, +) + + +def evaluate_model_features( + mol: gto.Mole, + dm: Tensor, + grids: Grid, + feature_spec: FeatureSpec, + max_memory_in_mb: int = 2000, +) -> FeatureMap: + """Evaluate a model's named features without running the model.""" + model_features = get_grid_features(mol, dm, grids, feature_spec) + ao_features = feature_spec.ao_features + if ao_features is None: + return model_features + + feature_function = feature_math.MGGAFeatureFunction(ao_features) + raw_features = ao_evaluation.evaluate_raw_features_auto_chunk( + dm, + mol, + grids, + feature_function, + block_size=None, + max_memory=max_memory_in_mb, + gpu=dm.device.type == "cuda", + ) + is_spin_polarized = dm.ndim == 3 + for feature_name, feature_value in feature_function.to_dict(raw_features).items(): + model_features[feature_name] = feature_math.maybe_expand_and_divide( + feature_value, not is_spin_polarized, 2 + ) + return model_features @dataclass(frozen=True) @@ -112,6 +151,56 @@ def _make_atom_grid_order(atomic_grid_sizes: Tensor) -> AtomGridOrder: return AtomGridOrder(atom_indices=atom_indices, grid_indices=grid_indices) +@dataclass(frozen=True) +class ModelChunkIndices: + """Original atom and grid-point indices for one model evaluation.""" + + atom_indices: Tensor + grid_indices: Tensor + + +def _make_model_chunk_indices( + dm: Tensor, + atomic_grid_sizes: Tensor, + nfeatures: int, + deriv_order: int, + max_memory_in_mb: int | None, + safety_fraction: float, +) -> list[ModelChunkIndices]: + """Build memory-sized homogeneous chunks in a stable atom ordering.""" + atom_grid_order = _make_atom_grid_order(atomic_grid_sizes) + sorted_atomic_grid_sizes = atomic_grid_sizes.index_select( + 0, atom_grid_order.atom_indices + ) + max_atoms_per_grid_size = estimate_max_model_atoms_per_chunk( + dm=dm, + atomic_grid_sizes=sorted_atomic_grid_sizes, + nfeatures=nfeatures, + max_memory_in_mb=max_memory_in_mb, + safety_fraction=safety_fraction, + func_deriv=deriv_order, + ) + for grid_size, max_atoms in max_atoms_per_grid_size.items(): + if max_atoms < 1: + LOG.warning( + "Adjusted model chunk capacity for atomic grid size %d from %d " + "to one atom. Hope for no OOM.", + grid_size, + max_atoms, + ) + max_atoms_per_grid_size[grid_size] = 1 + + return [ + ModelChunkIndices( + atom_indices=atom_grid_order.atom_indices[layout.atom_slice], + grid_indices=atom_grid_order.grid_indices[layout.grid_slice], + ) + for layout in _make_atom_grid_chunks( + sorted_atomic_grid_sizes, max_atoms_per_grid_size + ) + ] + + @dataclass(frozen=True) class ModelFeatureChunk: """Chunk-local raw features and the corresponding model input dictionary.""" @@ -121,12 +210,57 @@ class ModelFeatureChunk: model_features: FeatureMap -class ModelFeatureChunker: - """Prepared atom-aligned partition of raw and model features. +@dataclass(frozen=True) +class ModelFeaturePlan: + """Describe feature evaluation and model inputs for model chunking. + + The evaluation specification contains every feature needed by the calculation, + including runtime bookkeeping. The model specification defines the smaller + feature set exposed to each functional invocation. + + Examples: + Evaluation may include features that are not passed to the model: + + >>> evaluation = FeatureSpec([Feature.DENSITY, Feature.ATOMIC_GRID_SIZES]) + >>> model = FeatureSpec([Feature.DENSITY]) + >>> ModelFeaturePlan(evaluation, model).model_feature_spec == model + True + + Model features that are unavailable from evaluation are rejected: + + >>> ModelFeaturePlan(model, FeatureSpec([Feature.GRID_WEIGHTS])) + Traceback (most recent call last): + ... + ValueError: Model features missing from evaluation: grid_weights - The chunker is a snapshot of the supplied density matrix, grid, and raw - features. It can be iterated repeatedly while those inputs represent the - same calculation, but must not be reused after their state changes. + Attributes: + evaluation_feature_spec: Complete feature set evaluated for the calculation. + model_feature_spec: Features exposed to each functional invocation. + """ + + evaluation_feature_spec: FeatureSpec + model_feature_spec: FeatureSpec + + def __post_init__(self) -> None: + missing_features = set(self.model_feature_spec) - set( + self.evaluation_feature_spec + ) + if missing_features: + missing_names = ", ".join( + sorted(feature.value for feature in missing_features) + ) + raise ValueError(f"Model features missing from evaluation: {missing_names}") + + +class ModelFeatureChunker: + """Prepare atom-aligned model inputs from globally evaluated AO features. + + Iteration yields detached chunk-local raw tensors requiring gradients so model + cotangents can be assembled and propagated through the global AO evaluation. + The chunker is a snapshot and may only be reused while its density matrix, + grid, and raw features still describe the same calculation. When spatial + decomposition is unsupported, it yields one full-grid chunk to preserve + non-additive model semantics. """ def __init__( @@ -135,94 +269,200 @@ def __init__( dm: Tensor, grids: Grid, atom_major_raw_features: Tensor, - feature_function: feature_math.MGGAFeatureFunction, + feature_plan: ModelFeaturePlan, deriv_order: int, max_memory_in_mb: int | None = None, safety_fraction: float = 0.8, ) -> None: - feature_spec = feature_function.feature_spec - if not feature_spec.supports_spatial_decomposition: - raise ValueError( - f"Atom-aligned model chunking requires " - f"{Feature.ATOMIC_GRID_SIZES.value!r}." - ) - - grid_features = get_grid_features(mol, dm, grids, feature_spec) + evaluation_feature_spec = feature_plan.evaluation_feature_spec + ao_features = evaluation_feature_spec.ao_features + assert ao_features is not None + feature_function = feature_math.MGGAFeatureFunction(ao_features) + grid_features = get_grid_features(mol, dm, grids, evaluation_feature_spec) atomic_grid_sizes = grid_features[Feature.ATOMIC_GRID_SIZES] - atom_grid_order = _make_atom_grid_order(atomic_grid_sizes) - sorted_atomic_grid_sizes = atomic_grid_sizes.index_select( - 0, atom_grid_order.atom_indices - ) - - max_atoms_per_grid_size = estimate_max_model_atoms_per_chunk( - dm=dm, - atomic_grid_sizes=sorted_atomic_grid_sizes, - nfeatures=feature_function.nfeats, - max_memory_in_mb=max_memory_in_mb, - safety_fraction=safety_fraction, - func_deriv=deriv_order, - ) - for grid_size, max_atoms in max_atoms_per_grid_size.items(): - if max_atoms < 1: - LOG.warning( - "Adjusted model chunk capacity for atomic grid size %d from %d " - "to one atom. Hope for no OOM.", - grid_size, - max_atoms, + if feature_plan.model_feature_spec.supports_spatial_decomposition: + self._chunk_indices = _make_model_chunk_indices( + dm, + atomic_grid_sizes, + nfeatures=feature_function.nfeats, + deriv_order=deriv_order, + max_memory_in_mb=max_memory_in_mb, + safety_fraction=safety_fraction, + ) + else: + self._chunk_indices = [ + ModelChunkIndices( + atom_indices=torch.arange(mol.natm, device=dm.device), + grid_indices=torch.arange( + int(atomic_grid_sizes.sum().item()), device=dm.device + ), ) - max_atoms_per_grid_size[grid_size] = 1 + ] self._atom_major_raw_features = atom_major_raw_features self._grid_features = grid_features self._feature_function = feature_function - self._chunk_layouts = _make_atom_grid_chunks( - sorted_atomic_grid_sizes, max_atoms_per_grid_size - ) - self._atom_order = atom_grid_order.atom_indices - self._grid_order = atom_grid_order.grid_indices + self._model_feature_spec = feature_plan.model_feature_spec self._is_spin_polarized = dm.ndim == 3 def __iter__(self) -> Iterator[ModelFeatureChunk]: """Yield detached raw features paired with atom-aligned model inputs.""" - feature_spec = self._feature_function.feature_spec - for layout in self._chunk_layouts: - atom_indices = self._atom_order[layout.atom_slice] - grid_indices = self._grid_order[layout.grid_slice] + for chunk_indices in self._chunk_indices: + atom_indices = chunk_indices.atom_indices + grid_indices = chunk_indices.grid_indices raw_features = ( self._atom_major_raw_features.index_select(-1, grid_indices) .detach() .requires_grad_() ) + decoded_features = { + feature: feature_math.maybe_expand_and_divide( + value, not self._is_spin_polarized, 2 + ) + for feature, value in self._feature_function.to_dict( + raw_features + ).items() + } model_features: FeatureMap = {} - for feature_name in _GRID_POINT_FEATURES: - if feature_spec.requests(feature_name): - model_features[feature_name] = self._grid_features[ - feature_name - ].index_select(0, grid_indices) - - for feature_name in _ATOM_FEATURES: - if feature_spec.requests(feature_name): - model_features[feature_name] = self._grid_features[ - feature_name + for feature in self._model_feature_spec: + if feature in _AO_DERIVED_FEATURES: + model_features[feature] = decoded_features[feature] + elif feature in _GRID_POINT_FEATURES: + model_features[feature] = self._grid_features[feature].index_select( + 0, grid_indices + ) + elif feature in _ATOM_FEATURES: + model_features[feature] = self._grid_features[feature].index_select( + 0, atom_indices + ) + elif feature is Feature.ATOMIC_GRID_SIZE_BOUND_SHAPE: + atomic_grid_sizes = self._grid_features[ + Feature.ATOMIC_GRID_SIZES ].index_select(0, atom_indices) + model_features[feature] = torch.zeros( + int(atomic_grid_sizes.max().item()), + 0, + dtype=torch.long, + device=raw_features.device, + ) + else: + raise ValueError(f"Unsupported model feature: {feature}") + yield ModelFeatureChunk( + grid_indices=grid_indices, + raw_features=raw_features, + model_features=model_features, + ) + + +def _evaluate_feature_gradients( + functional: ExcFunctionalBase, + model_features: FeatureMap, + differentiable_features: set[Feature], +) -> FeatureMap: + inputs = { + feature_name: model_features[feature_name] + for feature_name in differentiable_features + } + if not inputs: + return {} + + return feature_math.feature_derivatives( + lambda input_features: functional.get_exc(model_features | input_features), + inputs, + ) + - if feature_spec.requests(Feature.ATOMIC_GRID_SIZE_BOUND_SHAPE): - max_size = int(model_features[Feature.ATOMIC_GRID_SIZES].max().item()) - model_features[Feature.ATOMIC_GRID_SIZE_BOUND_SHAPE] = torch.zeros( - max_size, +def evaluate_chunked_feature_gradients( + functional: ExcFunctionalBase, + dm: Tensor, + model_features: FeatureMap, + differentiable_features: set[Feature], + max_memory_in_mb: int | None = None, + safety_fraction: float = 0.8, +) -> FeatureMap: + """Evaluate and assemble model gradients from precomputed feature chunks.""" + atomic_grid_sizes = model_features[Feature.ATOMIC_GRID_SIZES] + feature_spec = FeatureSpec(functional.features) + if not feature_spec.supports_spatial_decomposition: + full_features = { + feature_name: feature_value + for feature_name, feature_value in model_features.items() + if feature_spec.requests(feature_name) + } + full_features |= { + feature_name: model_features[feature_name] + for feature_name in differentiable_features + } + return _evaluate_feature_gradients( + functional, full_features, differentiable_features + ) + + ao_features = feature_spec.ao_features + nfeatures = ao_features.nfeats if ao_features is not None else 0 + chunk_indices = _make_model_chunk_indices( + dm, + atomic_grid_sizes, + nfeatures=nfeatures, + deriv_order=1, + max_memory_in_mb=max_memory_in_mb, + safety_fraction=safety_fraction, + ) + gradients = { + feature_name: torch.zeros_like(model_features[feature_name]) + for feature_name in differentiable_features + } + if not differentiable_features: + return gradients + + for indices in chunk_indices: + chunk_features: FeatureMap = {} + for feature_name, feature_value in model_features.items(): + if not feature_spec.requests(feature_name): + continue + if feature_name in _AO_DERIVED_FEATURES: + chunk_features[feature_name] = feature_value.index_select( + -1, indices.grid_indices + ) + elif feature_name in _GRID_POINT_FEATURES: + chunk_features[feature_name] = feature_value.index_select( + 0, indices.grid_indices + ) + elif feature_name in _ATOM_FEATURES: + chunk_features[feature_name] = feature_value.index_select( + 0, indices.atom_indices + ) + elif feature_name is Feature.ATOMIC_GRID_SIZE_BOUND_SHAPE: + max_grid_size = int( + atomic_grid_sizes.index_select(0, indices.atom_indices).max().item() + ) + chunk_features[feature_name] = torch.zeros( + max_grid_size, 0, dtype=torch.long, - device=raw_features.device, + device=feature_value.device, ) + else: + raise ValueError(f"Unsupported model feature: {feature_name}") - for feature_name, feature in self._feature_function.to_dict( - raw_features - ).items(): - model_features[feature_name] = feature_math.maybe_expand_and_divide( - feature, not self._is_spin_polarized, 2 + local_gradients = _evaluate_feature_gradients( + functional, chunk_features, differentiable_features + ) + for feature_name, local_gradient in local_gradients.items(): + if feature_name in _AO_DERIVED_FEATURES: + gradients[feature_name].index_copy_( + -1, indices.grid_indices, local_gradient.detach() ) - yield ModelFeatureChunk( - grid_indices=grid_indices, - raw_features=raw_features, - model_features=model_features, - ) + elif feature_name in _GRID_POINT_FEATURES: + gradients[feature_name].index_copy_( + 0, indices.grid_indices, local_gradient.detach() + ) + elif feature_name in _ATOM_FEATURES: + gradients[feature_name].index_copy_( + 0, indices.atom_indices, local_gradient.detach() + ) + else: + raise ValueError( + f"Unsupported differentiable model feature: {feature_name}" + ) + + return gradients diff --git a/skala/src/skala/pyscf/xc_integrator.py b/skala/src/skala/pyscf/xc_integrator.py index 8d34c5f4..999a1ede 100644 --- a/skala/src/skala/pyscf/xc_integrator.py +++ b/skala/src/skala/pyscf/xc_integrator.py @@ -11,14 +11,13 @@ from torch import Tensor from pyscf import gto -from skala.features import Feature +from skala.features import AOFeatureSpec, Feature from skala.functional.base import ExcFunctionalBase from skala.pyscf import ao_evaluation, feature_math from skala.pyscf.backend import Grid, check_gpu_imports_were_successful from skala.pyscf.evaluation import EvaluationPolicy, FeatureSpec -from skala.pyscf.features import generate_features from skala.pyscf.grids import SkalaGrids as PySCFSkalaGrids -from skala.pyscf.model_chunking import ModelFeatureChunker +from skala.pyscf.model_chunking import ModelFeatureChunker, ModelFeaturePlan from skala.pyscf.spatial_grid_layout import SpatialGridLayout if TYPE_CHECKING: @@ -87,16 +86,22 @@ def density( max_memory: int = 2000, ) -> Tensor: """Evaluate the total density on each grid point.""" - mol_features = generate_features( + self._validate_device(dm) + _assert_skala_grid(grids, self.device) + feature_function = feature_math.MGGAFeatureFunction( + AOFeatureSpec([Feature.DENSITY]) + ) + evaluate_raw_features = self._raw_feature_evaluator( mol, - dm, grids, - features={Feature.DENSITY}, - chunk_size=self.evaluation_policy.ao_block_size, - max_memory=max_memory, - gpu=self.device.type == "cuda", + feature_function, + max_memory, + screened=_should_screen_aos(mol), ) - return mol_features[Feature.DENSITY].sum(0) + density = feature_function.to_dict(evaluate_raw_features(dm))[Feature.DENSITY] + if dm.dim() == 2: + return density + return density.sum(0) def __call__( self, @@ -108,84 +113,45 @@ def __call__( """Evaluate electron count, XC energy, and XC potential.""" self._validate_device(dm) _assert_skala_grid(grids, self.device) - if self.feature_spec.supports_spatial_decomposition and _should_screen_aos(mol): - return self._integrate_screened(mol, grids, dm, max_memory) - return self._integrate_dense(mol, grids, dm, max_memory) - - def gen_response( - self, - mol: gto.Mole, - grids: Grid, - dm0: Tensor, - max_memory: int = 2000, - safety_fraction: float | None = None, - ) -> Callable[[Tensor], Tensor]: - """Build an XC-only Hessian-vector product callable.""" - self._validate_device(dm0) - _assert_skala_grid(grids, self.device) - if self.feature_spec.supports_spatial_decomposition and _should_screen_aos(mol): - return self._gen_response_screened( - mol, - grids, - dm0, - max_memory=max_memory, - safety_fraction=( - self.evaluation_policy.safety_fraction - if safety_fraction is None - else safety_fraction - ), - ) - return self._gen_response_dense(mol, grids, dm0, max_memory=max_memory) - - def _validate_device(self, dm: Tensor) -> None: - if self.device != dm.device: - raise ValueError( - f"Density matrix device {dm.device} does not match functional device {self.device}" - ) - - def _get_spatial_grid_layout(self, mol: gto.Mole, grids: Grid) -> SpatialGridLayout: - if _assert_skala_grid(grids, self.device): - return grids.prepare_spatial_grid_layout(mol, self.device) - raise AssertionError("unreachable") - - def _integrate_screened( - self, - mol: gto.Mole, - grids: Grid, - dm: Tensor, - max_memory: int, - ) -> XCResult: dm = dm.detach().requires_grad_() dm_eval = dm.double() electron_count = torch.zeros(2, device=self.device, dtype=dm_eval.dtype) energy = torch.tensor(0.0, device=self.device, dtype=dm_eval.dtype) - integration_feature_spec = FeatureSpec( - self.feature_spec.names | {Feature.DENSITY, Feature.GRID_WEIGHTS} + evaluation_feature_spec = FeatureSpec( + self.feature_spec + | { + Feature.DENSITY, + Feature.GRID_WEIGHTS, + Feature.ATOMIC_GRID_SIZES, + } ) - feature_function = feature_math.MGGAFeatureFunction(integration_feature_spec) - spatial_grid_layout = self._get_spatial_grid_layout(mol, grids) - sorted_raw_features = ao_evaluation.evaluate_ao_features_blockwise( - dm_eval, + model_feature_spec = self.feature_spec | { + Feature.DENSITY, + Feature.GRID_WEIGHTS, + } + feature_plan = ModelFeaturePlan(evaluation_feature_spec, model_feature_spec) + ao_features = evaluation_feature_spec.ao_features + assert ao_features is not None + feature_function = feature_math.MGGAFeatureFunction(ao_features) + evaluate_raw_features = self._raw_feature_evaluator( mol, - spatial_grid_layout.sorted_grids, + grids, feature_function, - spatial_grid_layout.block_size, - False, - ) - atom_major_raw_features = sorted_raw_features.index_select( - -1, spatial_grid_layout.inverse_permutation + max_memory, + screened=_should_screen_aos(mol), ) + raw_features = evaluate_raw_features(dm_eval) model_chunks = ModelFeatureChunker( mol, dm, grids, - atom_major_raw_features=atom_major_raw_features, - feature_function=feature_function, + atom_major_raw_features=raw_features, + feature_plan=feature_plan, deriv_order=1, max_memory_in_mb=max_memory if dm.device.type == "cpu" else None, safety_fraction=self.evaluation_policy.safety_fraction, ) - atom_major_cotangent = torch.zeros_like(atom_major_raw_features) + raw_cotangent = torch.zeros_like(raw_features) for chunk in model_chunks: local_raw_features = chunk.raw_features mol_features = chunk.model_features @@ -195,9 +161,7 @@ def _integrate_screened( local_raw_features, torch.ones_like(energy_chunk), ) - atom_major_cotangent.index_copy_( - -1, chunk.grid_indices, local_cotangent.detach() - ) + raw_cotangent.index_copy_(-1, chunk.grid_indices, local_cotangent.detach()) electron_count += ( (mol_features[Feature.DENSITY] * mol_features[Feature.GRID_WEIGHTS]) .sum(dim=-1) @@ -206,94 +170,53 @@ def _integrate_screened( energy += energy_chunk.detach() del energy_chunk, local_cotangent, local_raw_features, mol_features - sorted_cotangent = atom_major_cotangent.index_select( - -1, spatial_grid_layout.forward_permutation - ) - (potential,) = torch.autograd.grad( - sorted_raw_features, - dm, - sorted_cotangent, - ) - return XCResult(electron_count, energy, potential) - - def _integrate_dense( - self, - mol: gto.Mole, - grids: Grid, - dm: Tensor, - max_memory: int, - *, - create_graph: bool = False, - ) -> XCResult: - dm = dm.requires_grad_() - mol_features = generate_features( - mol, - dm, - grids, - set(self.feature_spec.names) | {Feature.DENSITY, Feature.GRID_WEIGHTS}, - chunk_size=self.evaluation_policy.ao_block_size, - max_memory=max_memory, - gpu=self.device.type == "cuda", - ) - energy = self.functional.get_exc(mol_features) - (potential,) = torch.autograd.grad( - energy, - dm, - torch.ones_like(energy), - retain_graph=create_graph, - create_graph=create_graph, - ) - electron_count = ( - mol_features[Feature.DENSITY] * mol_features[Feature.GRID_WEIGHTS] - ).sum(dim=-1) + (potential,) = torch.autograd.grad(raw_features, dm, raw_cotangent) return XCResult(electron_count, energy, potential) - def _gen_response_screened( + def gen_response( self, mol: gto.Mole, grids: Grid, dm0: Tensor, - *, - max_memory: int, - safety_fraction: float, + max_memory: int = 2000, + safety_fraction: float | None = None, ) -> Callable[[Tensor], Tensor]: + """Build an XC-only Hessian-vector product callable.""" + self._validate_device(dm0) + _assert_skala_grid(grids, self.device) + response_safety_fraction = ( + self.evaluation_policy.safety_fraction + if safety_fraction is None + else safety_fraction + ) dm0 = dm0.requires_grad_() - feature_function = feature_math.MGGAFeatureFunction(self.feature_spec) - spatial_grid_layout = self._get_spatial_grid_layout(mol, grids) - sorted_raw_features = ao_evaluation.evaluate_ao_features_blockwise( - dm0.double(), + evaluation_feature_spec = self.feature_spec | {Feature.ATOMIC_GRID_SIZES} + feature_plan = ModelFeaturePlan(evaluation_feature_spec, self.feature_spec) + ao_features = evaluation_feature_spec.ao_features + assert ao_features is not None + feature_function = feature_math.MGGAFeatureFunction(ao_features) + evaluate_raw_features = self._raw_feature_evaluator( mol, - spatial_grid_layout.sorted_grids, + grids, feature_function, - spatial_grid_layout.block_size, - False, - ) - atom_major_raw_features = sorted_raw_features.index_select( - -1, spatial_grid_layout.inverse_permutation + max_memory, + screened=_should_screen_aos(mol), ) + raw_features = evaluate_raw_features(dm0) model_chunks = ModelFeatureChunker( mol, dm0, grids, - atom_major_raw_features=atom_major_raw_features, - feature_function=feature_function, + atom_major_raw_features=raw_features, + feature_plan=feature_plan, deriv_order=2, max_memory_in_mb=max_memory if dm0.device.type == "cpu" else None, - safety_fraction=safety_fraction, + safety_fraction=response_safety_fraction, ) def hessian_vector_product(dm1: Tensor) -> Tensor: - sorted_tangent = ao_evaluation.evaluate_ao_features_blockwise( - dm1, - mol, - spatial_grid_layout.sorted_grids, - feature_function, - spatial_grid_layout.block_size, - ) - atom_major_tangent = sorted_tangent.index_select( - -1, spatial_grid_layout.inverse_permutation - ).detach() - atom_major_hessian_action = torch.zeros_like(atom_major_raw_features) + raw_tangent = evaluate_raw_features(dm1).detach() + raw_hessian_action = torch.zeros_like(raw_features) for chunk in model_chunks: local_raw_features = chunk.raw_features mol_features = chunk.model_features @@ -312,10 +235,10 @@ def hessian_vector_product(dm1: Tensor) -> Tensor: (local_hessian_action,) = torch.autograd.grad( local_gradient, local_raw_features, - atom_major_tangent.index_select(-1, chunk.grid_indices), + raw_tangent.index_select(-1, chunk.grid_indices), ) - atom_major_hessian_action.index_copy_( + raw_hessian_action.index_copy_( -1, chunk.grid_indices, local_hessian_action.detach() ) del ( @@ -326,46 +249,63 @@ def hessian_vector_product(dm1: Tensor) -> Tensor: mol_features, ) - sorted_hessian_action = atom_major_hessian_action.index_select( - -1, spatial_grid_layout.forward_permutation - ) (hvp_total,) = torch.autograd.grad( - sorted_raw_features, + raw_features, dm0, - sorted_hessian_action, + raw_hessian_action, retain_graph=True, ) return hvp_total return hessian_vector_product - def _gen_response_dense( + def _validate_device(self, dm: Tensor) -> None: + if self.device != dm.device: + raise ValueError( + f"Density matrix device {dm.device} does not match functional device {self.device}" + ) + + def _get_spatial_grid_layout(self, mol: gto.Mole, grids: Grid) -> SpatialGridLayout: + if _assert_skala_grid(grids, self.device): + return grids.prepare_spatial_grid_layout(mol, self.device) + raise AssertionError("unreachable") + + def _raw_feature_evaluator( self, mol: gto.Mole, grids: Grid, - dm0: Tensor, - *, + feature_function: feature_math.MGGAFeatureFunction, max_memory: int, + *, + screened: bool, ) -> Callable[[Tensor], Tensor]: - dm0 = dm0.requires_grad_() - potential = self._integrate_dense( - mol, - grids, - dm0, - max_memory, - create_graph=True, - ).potential + """Build an evaluator returning raw features in atom-major grid order.""" + if screened: + spatial_grid_layout = self._get_spatial_grid_layout(mol, grids) + + def evaluate_screened(value: Tensor) -> Tensor: + sorted_features = ao_evaluation.evaluate_ao_features_blockwise( + value.double(), + mol, + spatial_grid_layout.sorted_grids, + feature_function, + spatial_grid_layout.block_size, + ) + return sorted_features.index_select( + -1, spatial_grid_layout.inverse_permutation + ) - def hessian_vector_product(dm1: Tensor) -> Tensor: - # A potential independent of dm0 is constant, so its directional - # derivative (the density-matrix Hessian action) is zero. - if not potential.requires_grad: - return torch.zeros_like(dm0) - return torch.autograd.grad( - potential, - dm0, - dm1, - retain_graph=True, - )[0] + return evaluate_screened - return hessian_vector_product + def evaluate_dense(value: Tensor) -> Tensor: + return ao_evaluation.evaluate_raw_features_auto_chunk( + value, + mol, + grids, + feature_function, + block_size=self.evaluation_policy.ao_block_size, + max_memory=max_memory, + gpu=self.device.type == "cuda", + ) + + return evaluate_dense diff --git a/skala/tests/test_ao_screening.py b/skala/tests/test_ao_screening.py index 5a052866..0b5a434c 100644 --- a/skala/tests/test_ao_screening.py +++ b/skala/tests/test_ao_screening.py @@ -5,7 +5,7 @@ import numpy as np import pytest import torch -from skala.features import Feature, FeatureMap +from skala.features import AOFeatureSpec, Feature, FeatureMap from skala.functional.base import ExcFunctionalBase from skala.pyscf import ao_evaluation as ao_evaluation_module from skala.pyscf import grids as grids_module @@ -23,7 +23,10 @@ from skala.pyscf.evaluation import FeatureSpec from skala.pyscf.feature_math import MGGAFeatureFunction from skala.pyscf.grids import SkalaGrids -from skala.pyscf.model_chunking import ModelFeatureChunk +from skala.pyscf.model_chunking import ( + ModelFeatureChunk, + ModelFeaturePlan, +) from skala.pyscf.numint import SkalaNumInt from skala.pyscf.spatial_grid_layout import ( SpatialGridLayout, @@ -80,8 +83,7 @@ def test_mgga_supported_features_are_linear_in_density_matrix( Linearity requires the first JVP to equal direct feature evaluation on the tangent and the second JVP to vanish. """ - feature_spec = FeatureSpec(feature_names) - feature_function = MGGAFeatureFunction(feature_spec) + feature_function = MGGAFeatureFunction(AOFeatureSpec(feature_names)) ncomp = (expected_deriv + 1) * (expected_deriv + 2) * (expected_deriv + 3) // 6 ao = torch.arange(1, ncomp * 2 * 3 + 1, dtype=torch.float64).reshape(ncomp, 2, 3) if expected_deriv == 0: @@ -113,7 +115,6 @@ def first_jvp(value: torch.Tensor) -> torch.Tensor: assert feature_function.deriv == expected_deriv assert feature_function.nfeats == expected_nfeats - assert feature_function.feature_spec is feature_spec assert features.shape == (expected_nfeats, 3) assert set(feature_function.to_dict(features)) == feature_names torch.testing.assert_close(feature_jvp, feature_function(tangent, ao)) @@ -129,7 +130,7 @@ def test_mgga_analytic_vjp_matches_autograd( feature_names: tuple[Feature, ...], spin_channels: int | None ) -> None: """Match the analytic MGGA VJP to autograd for every feature and spin layout.""" - feature_function = MGGAFeatureFunction(FeatureSpec(feature_names)) + feature_function = MGGAFeatureFunction(AOFeatureSpec(feature_names)) ncomp = ( (feature_function.deriv + 1) * (feature_function.deriv + 2) @@ -158,7 +159,7 @@ def test_feature_block_compiled_vjp_matches_eager( monkeypatch: pytest.MonkeyPatch, ) -> None: """Match compiled blockwise feature evaluation and VJP to eager execution.""" - feature_function = MGGAFeatureFunction(FeatureSpec(_MGGA_FEATURES)) + feature_function = MGGAFeatureFunction(AOFeatureSpec(_MGGA_FEATURES)) generator = torch.Generator().manual_seed(0) block = _AOBlock( ao_values=torch.randn((10, 3, 5), dtype=torch.float64, generator=generator), @@ -202,15 +203,19 @@ def test_feature_block_compiled_vjp_matches_eager( torch.testing.assert_close(compiled_vjp, eager_vjp) -@pytest.mark.parametrize("feature_names", [[], [Feature.GRID_WEIGHTS]]) +@pytest.mark.parametrize( + ("feature_names", "message"), + [ + ([], "At least one AO-derived feature must be selected"), + ([Feature.GRID_WEIGHTS], "Unsupported AO features: grid_weights"), + ], +) def test_mgga_requires_at_least_one_ao_derived_feature( - feature_names: list[Feature], + feature_names: list[Feature], message: str ) -> None: """Reject MGGA feature functions that contain no AO-derived feature.""" - with pytest.raises( - ValueError, match="At least one AO-derived feature must be selected" - ): - MGGAFeatureFunction(FeatureSpec(feature_names)) + with pytest.raises(ValueError, match=message): + MGGAFeatureFunction(AOFeatureSpec(feature_names)) def test_patch_ao_screening_restores_previous_decision(carbon: gto.Mole) -> None: @@ -248,9 +253,41 @@ def test_active_cpu_ao_indices(carbon: gto.Mole) -> None: assert empty.size == 0 +def test_dense_model_evaluation_uses_model_chunks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Keep model chunking independent from the AO screening decision.""" + mol = gto.M(atom="H 0 0 0; H 0 0 0.74", basis="sto-3g", verbose=0) + grids = _minimal_atom_grid(mol) + atom_grid_size = grids.weights.size // mol.natm + monkeypatch.setattr( + model_chunking_module, + "estimate_max_model_atoms_per_chunk", + lambda *args, **kwargs: {atom_grid_size: 1}, + ) + + class CountingFunctional(QuadraticFunctional): + def __init__(self) -> None: + super().__init__() + self.calls = 0 + + def get_exc(self, mol: FeatureMap) -> torch.Tensor: + self.calls += 1 + return super().get_exc(mol) + + functional = CountingFunctional() + integrator = XCIntegrator(functional) + dm = torch.as_tensor(dft.RKS(mol).get_init_guess()) + + with force_ao_screening(False): + integrator(mol, grids, dm) + + assert functional.calls == mol.natm + + def test_resolve_ao_block_size_modes(carbon: gto.Mole) -> None: """Resolve aligned CPU block sizes and reject explicit GPU block sizing.""" - feature_function = MGGAFeatureFunction(FeatureSpec([Feature.DENSITY])) + feature_function = MGGAFeatureFunction(AOFeatureSpec([Feature.DENSITY])) backend_block_size = dft.gen_grid.BLKSIZE # CPU sizes are aligned locally; GPU sizing is delegated unless explicitly invalid. @@ -450,7 +487,9 @@ def fake_prepare_spatial_grid_layout( "prepare_spatial_grid_layout", fake_prepare_spatial_grid_layout, ) - numint: _NumPyNumInt = SkalaNumInt(QuadraticFunctional()) + functional = QuadraticFunctional([Feature.DENSITY, Feature.GRID_WEIGHTS]) + assert not FeatureSpec(functional.features).supports_spatial_decomposition + numint: _NumPyNumInt = SkalaNumInt(functional) other_numint: _NumPyNumInt = SkalaNumInt(QuadraticFunctional()) layout = numint.integrator._get_spatial_grid_layout(carbon, grids) @@ -518,20 +557,13 @@ def test_first_and_second_order_use_same_screening_decision( routes: list[str] = [] safety_fractions: list[float] = [] - def fake_generate_features( - mol: gto.Mole, + def fake_evaluate_raw_features_auto_chunk( dm: torch.Tensor, - grids: object, - features: set[Feature] | None = None, + *args: object, **kwargs: object, - ) -> FeatureMap: + ) -> torch.Tensor: routes.append("dense") - density = dm.square().sum().reshape(1).expand(2, 1) / 2 - return { - Feature.ATOMIC_GRID_SIZES: torch.tensor([1]), - Feature.DENSITY: density, - Feature.GRID_WEIGHTS: torch.ones(1, dtype=dm.dtype), - } + return dm.sum().reshape(1, 1) class FakeSpatialGridLayout: block_size = 1 @@ -577,17 +609,20 @@ def fake_model_feature_chunker( dm: torch.Tensor, grids: object, atom_major_raw_features: torch.Tensor, - feature_function: MGGAFeatureFunction, + feature_plan: ModelFeaturePlan, deriv_order: int, **kwargs: object, ) -> FakeModelFeatureChunks: + assert feature_plan.evaluation_feature_spec.ao_features safety_fraction = kwargs["safety_fraction"] assert isinstance(safety_fraction, float) safety_fractions.append(safety_fraction) return FakeModelFeatureChunks(atom_major_raw_features) monkeypatch.setattr( - xc_integrator_module, "generate_features", fake_generate_features + ao_evaluation_module, + "evaluate_raw_features_auto_chunk", + fake_evaluate_raw_features_auto_chunk, ) monkeypatch.setattr( grids_module, @@ -628,15 +663,11 @@ def fake_model_feature_chunker( assert result.shape == (carbon.nao_nr(), carbon.nao_nr()) expected_route = "screened" if expected else "dense" - expected_route_count = 3 if expected else 2 - assert routes == [expected_route] * expected_route_count - if expected: - assert safety_fractions == [ - 0.8, - 0.8 if response_safety_fraction is None else response_safety_fraction, - ] - else: - assert safety_fractions == [] + assert routes == [expected_route] * 3 + assert safety_fractions == [ + 0.8, + 0.8 if response_safety_fraction is None else response_safety_fraction, + ] def test_feature_block_helper_localizes_derivative_vectors() -> None: @@ -647,7 +678,7 @@ def test_feature_block_helper_localizes_derivative_vectors() -> None: adjoint calculation must select only this block's grid cotangent. Comparing both operations with direct local formulas catches mixing up AO and grid localization. """ - feature_function = MGGAFeatureFunction(FeatureSpec([Feature.DENSITY])) + feature_function = MGGAFeatureFunction(AOFeatureSpec([Feature.DENSITY])) block = _AOBlock( ao_values=torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float64), active_ao_indices=torch.tensor([0, 2]), @@ -688,7 +719,7 @@ def test_blockwise_ao_feature_transforms_follow_linear_operator( ) -> None: """Check spin-resolved first and second JVPs and the adjoint JVP.""" grids = _minimal_atom_grid(carbon) - feature_function = MGGAFeatureFunction(FeatureSpec([Feature.DENSITY])) + feature_function = MGGAFeatureFunction(AOFeatureSpec([Feature.DENSITY])) identity = torch.eye(carbon.nao_nr(), dtype=torch.float64) dm = torch.stack((identity, 2 * identity)) tangent = torch.arange(1, dm.numel() + 1, dtype=dm.dtype).reshape(dm.shape) @@ -778,7 +809,7 @@ def block_loop( monkeypatch.setattr(dft.numint, "NumInt", FakeNumInt) - feature_function = MGGAFeatureFunction(FeatureSpec([Feature.DENSITY])) + feature_function = MGGAFeatureFunction(AOFeatureSpec([Feature.DENSITY])) dm = torch.diag( torch.arange(1, carbon.nao_nr() + 1, dtype=torch.float64) ).requires_grad_() @@ -850,7 +881,7 @@ def block_loop( ) monkeypatch.setattr(dft.numint, "NumInt", FakeNumInt) - feature_function = MGGAFeatureFunction(FeatureSpec([Feature.DENSITY])) + feature_function = MGGAFeatureFunction(AOFeatureSpec([Feature.DENSITY])) blocks = list(_CPUAOBlockLoop(carbon, grids, feature_function, block_size)) @@ -884,7 +915,7 @@ def block_loop( yield ao, None, grids.weights, grids.coords monkeypatch.setattr(dft.numint, "NumInt", FakeNumInt) - feature_function = MGGAFeatureFunction(FeatureSpec([Feature.DENSITY])) + feature_function = MGGAFeatureFunction(AOFeatureSpec([Feature.DENSITY])) dm = torch.eye(carbon.nao_nr(), dtype=torch.float64).requires_grad_() features = evaluate_ao_features_blockwise( @@ -1036,6 +1067,23 @@ def test_cpu_quadratic_dense_screened_equivalence_heteronuclear() -> None: ) +def test_cpu_density_dense_screened_equivalence() -> None: + """Match dense and screened density in the original atom-major grid order.""" + mol = gto.M(atom="H 0 0 0; F 0 0 0.92", basis="sto-3g", spin=0, verbose=0) + grids = _minimal_atom_grid(mol) + numint: _NumPyNumInt = SkalaNumInt(QuadraticFunctional()) + dm = dft.RKS(mol).get_init_guess() + + with force_ao_screening(False): + dense = numint.get_rho(mol, dm, grids) + with force_ao_screening(True): + screened = numint.get_rho(mol, dm, grids) + + assert dense.shape == grids.weights.shape + assert screened.shape == grids.weights.shape + np.testing.assert_allclose(screened, dense, rtol=1e-10, atol=1e-11) + + @pytest.mark.parametrize( ("features", "energy_feature"), [ diff --git a/skala/tests/test_evaluation.py b/skala/tests/test_evaluation.py index 6ec89daf..522f53b0 100644 --- a/skala/tests/test_evaluation.py +++ b/skala/tests/test_evaluation.py @@ -1,62 +1,174 @@ from dataclasses import FrozenInstanceError import pytest -from skala.features import Feature +import torch +from skala.features import AOFeatureSpec, Feature from skala.pyscf.evaluation import EvaluationPolicy, FeatureSpec +from skala.pyscf.feature_math import AODirection, MGGAFeatureFunction, PackedAO def test_feature_name_parses_model_metadata_string() -> None: assert Feature("density") is Feature.DENSITY +def test_packed_ao_views_preserve_trailing_dimensions() -> None: + components = torch.arange(60).reshape(10, 2, 3) + + for tensor in (components, components.transpose(-1, -2)): + packed_ao = PackedAO(tensor) + + torch.testing.assert_close(packed_ao.value, tensor[0]) + torch.testing.assert_close(packed_ao.gradient, tensor[1:4]) + for actual, component in zip( + packed_ao.diagonal_hessian, (4, 7, 9), strict=True + ): + torch.testing.assert_close(actual, tensor[component]) + + +def test_packed_ao_normalizes_unpacked_values() -> None: + values = torch.arange(6).reshape(2, 3) + + packed_ao = PackedAO(values) + + assert packed_ao._tensor.shape == (1, 2, 3) + torch.testing.assert_close(packed_ao.value, values) + + +def test_packed_ao_rejects_unsupported_component_count() -> None: + with pytest.raises(ValueError, match="1, 4, or 10"): + PackedAO(torch.zeros((3, 2, 3))) + + +def test_packed_ao_hessian_uses_pyscf_component_order() -> None: + packed_ao = PackedAO(torch.arange(10).reshape(10, 1, 1)) + + components = [ + (ao_direction, feature_direction, int(value.item())) + for ao_direction, feature_direction, value in packed_ao.hessian() + ] + + assert components == [ + (AODirection.X, AODirection.X, 4), + (AODirection.X, AODirection.Y, 5), + (AODirection.X, AODirection.Z, 6), + (AODirection.Y, AODirection.X, 5), + (AODirection.Y, AODirection.Y, 7), + (AODirection.Y, AODirection.Z, 8), + (AODirection.Z, AODirection.X, 6), + (AODirection.Z, AODirection.Y, 8), + (AODirection.Z, AODirection.Z, 9), + ] + + +def test_ao_feature_spec_normalizes_and_derives_requirements() -> None: + spec = AOFeatureSpec([Feature.DENSITY, Feature.DENSITY, Feature.GRAD, Feature.LAPL]) + + assert list(spec) == [ + (Feature.DENSITY, slice(0, 1)), + (Feature.GRAD, slice(1, 4)), + (Feature.LAPL, slice(4, 5)), + ] + assert spec.nderiv == 2 + assert spec.nfeats == 5 + + +def test_ao_feature_spec_rejects_grid_features() -> None: + with pytest.raises(ValueError, match="grid_weights"): + AOFeatureSpec([Feature.DENSITY, Feature.GRID_WEIGHTS]) + + +def test_ao_feature_spec_rejects_empty_features() -> None: + with pytest.raises(ValueError, match="At least one"): + AOFeatureSpec([]) + + +def test_feature_spec_normalizes_model_contract() -> None: + spec = FeatureSpec([Feature.DENSITY, Feature.DENSITY]) + + assert set(spec) == {Feature.DENSITY} + assert spec.requests(Feature.DENSITY) + assert not spec.requests(Feature.GRAD) + + @pytest.mark.parametrize( - ("features", "expected_order"), + ("feature", "supports_screened_evaluation"), + [ + (Feature.ATOMIC_GRID_WEIGHTS, False), + (Feature.ATOMIC_GRID_SIZES, True), + (Feature.ATOMIC_GRID_SIZE_BOUND_SHAPE, False), + ], +) +def test_feature_spec_derives_spatial_decomposition_support( + feature: Feature, supports_screened_evaluation: bool +) -> None: + spec = FeatureSpec([feature, feature]) + + assert set(spec) == {feature} + assert spec.supports_spatial_decomposition is supports_screened_evaluation + + +def test_feature_spec_derives_evaluation_needs() -> None: + feature_spec = FeatureSpec( + [Feature.DENSITY, Feature.DENSITY, Feature.ATOMIC_GRID_WEIGHTS] + ) + + assert set(feature_spec) == {Feature.DENSITY, Feature.ATOMIC_GRID_WEIGHTS} + assert feature_spec.ao_features == AOFeatureSpec([Feature.DENSITY]) + assert feature_spec.requires_ao_evaluation + assert feature_spec.requires_atomic_layout + + grid_feature_spec = FeatureSpec([Feature.GRID_WEIGHTS]) + assert grid_feature_spec.ao_features is None + assert not grid_feature_spec.requires_ao_evaluation + + +@pytest.mark.parametrize( + ("features", "expected"), [ - ([], 0), ([Feature.DENSITY], 0), ([Feature.GRAD], 1), ([Feature.KIN], 1), ([Feature.LAPL], 2), - ( - [ - Feature.DENSITY, - Feature.GRAD, - Feature.KIN, - Feature.LAPL, - ], - 2, - ), ], ) -def test_feature_spec_derives_mgga_requirements( - features: list[Feature], expected_order: int +def test_ao_feature_spec_derivative_order( + features: list[Feature], expected: int ) -> None: - spec = FeatureSpec(features) - - assert spec.requires_ao_evaluation is bool(features) - assert spec.ao_derivative_order == expected_order - assert spec.with_density is (Feature.DENSITY in features) - assert spec.with_grad is (Feature.GRAD in features) - assert spec.with_kin is (Feature.KIN in features) - assert spec.with_lapl is (Feature.LAPL in features) + assert AOFeatureSpec(features).nderiv == expected @pytest.mark.parametrize( - ("feature", "supports_screened_evaluation"), + ("features", "expected_derivative_order", "expected_feature_count"), [ - (Feature.ATOMIC_GRID_WEIGHTS, False), - (Feature.ATOMIC_GRID_SIZES, True), - (Feature.ATOMIC_GRID_SIZE_BOUND_SHAPE, False), + ([Feature.DENSITY], 0, 1), + ([Feature.GRAD], 1, 3), + ([Feature.KIN], 1, 1), + ([Feature.LAPL], 2, 1), + ([Feature.DENSITY, Feature.GRAD, Feature.KIN, Feature.LAPL], 2, 6), ], ) -def test_feature_spec_derives_atomic_layout_requirements( - feature: Feature, supports_screened_evaluation: bool +def test_mgga_feature_function_derives_channel_requirements( + features: list[Feature], + expected_derivative_order: int, + expected_feature_count: int, ) -> None: - spec = FeatureSpec([feature, feature]) + feature_function = MGGAFeatureFunction(AOFeatureSpec(features)) - assert spec.names == frozenset({feature}) - assert spec.requires_atomic_layout - assert spec.supports_spatial_decomposition is supports_screened_evaluation + assert feature_function.deriv == expected_derivative_order + assert feature_function.nfeats == expected_feature_count + + +def test_mgga_feature_function_to_dict_preserves_public_feature_shapes() -> None: + feature_function = MGGAFeatureFunction( + AOFeatureSpec([Feature.DENSITY, Feature.GRAD, Feature.KIN]) + ) + packed_features = torch.zeros((2, 5, 7)) + + features = feature_function.to_dict(packed_features) + + assert features[Feature.DENSITY].shape == (2, 7) + assert features[Feature.GRAD].shape == (2, 3, 7) + assert features[Feature.KIN].shape == (2, 7) def test_evaluation_policy_defaults_and_is_immutable() -> None: diff --git a/skala/tests/test_gpu4pyscf_ao_screening.py b/skala/tests/test_gpu4pyscf_ao_screening.py index 35983962..df2f062c 100644 --- a/skala/tests/test_gpu4pyscf_ao_screening.py +++ b/skala/tests/test_gpu4pyscf_ao_screening.py @@ -16,7 +16,7 @@ cupy = require_gpu() -from skala.features import Feature # noqa: E402 +from skala.features import AOFeatureSpec, Feature # noqa: E402 from skala.functional.base import ExcFunctionalBase # noqa: E402 from skala.gpu4pyscf import SkalaKS # noqa: E402 from skala.gpu4pyscf.grids import SkalaGrids as GPU4PySCFSkalaGrids # noqa: E402 @@ -25,7 +25,6 @@ evaluate_full_grid, ) from skala.pyscf.backend import Array, dft_gpu # noqa: E402 -from skala.pyscf.evaluation import FeatureSpec # noqa: E402 from skala.pyscf.feature_math import MGGAFeatureFunction # noqa: E402 from skala.pyscf.grids import SkalaGrids as PySCFSkalaGrids # noqa: E402 from skala.pyscf.numint import SkalaNumInt # noqa: E402 @@ -348,7 +347,7 @@ def test_gpu_empty_ao_block_matches_dense_reference() -> None: grids._non0ao_idx = None feature_function = MGGAFeatureFunction( - FeatureSpec([Feature.DENSITY, Feature.GRAD, Feature.KIN]) + AOFeatureSpec([Feature.DENSITY, Feature.GRAD, Feature.KIN]) ) dm = torch.eye(mol.nao_nr(), dtype=torch.float64, device="cuda", requires_grad=True) screened = evaluate_ao_features_blockwise( @@ -402,7 +401,7 @@ def block_loop( monkeypatch.setattr(dft_gpu.numint, "NumInt", FakeGpuNumInt) - feature_function = MGGAFeatureFunction(FeatureSpec([Feature.DENSITY])) + feature_function = MGGAFeatureFunction(AOFeatureSpec([Feature.DENSITY])) dm = torch.diag( torch.arange(1, mol.nao_nr() + 1, dtype=torch.float64, device="cuda") ).requires_grad_() diff --git a/skala/tests/test_gpu4pyscf_gradients.py b/skala/tests/test_gpu4pyscf_gradients.py index 8bedf27a..23d845cd 100644 --- a/skala/tests/test_gpu4pyscf_gradients.py +++ b/skala/tests/test_gpu4pyscf_gradients.py @@ -19,12 +19,13 @@ from skala.gpu4pyscf.gradients import ( # noqa: E402 SkalaRKSGradient, SkalaUKSGradient, + _veff_and_expl_nuc_grad, nuc_grad_from_veff, - veff_and_expl_nuc_grad, ) from skala.pyscf import SkalaKS as CpuSkalaKS # noqa: E402 -from skala.pyscf.features import generate_features # noqa: E402 +from skala.pyscf.evaluation import FeatureSpec # noqa: E402 from skala.pyscf.gradients import SkalaRKSGradient as CpuSkalaRKSGradient # noqa: E402 +from skala.pyscf.model_chunking import evaluate_model_features # noqa: E402 from skala.utils import torch_allocator # noqa: E402 from pyscf import gto # noqa: E402 @@ -155,6 +156,23 @@ def get_grid_and_rdm1(mol: gto.Mole) -> tuple[dft.Grids, torch.Tensor]: return mf.grids, rdm1 # maybe_expand_and_divide(rdm1, len(rdm1.shape) == 2, 2) +def _evaluate_nuclear_gradient( + functional: ExcFunctionalBase, + mol: gto.Mole, + grid: dft.Grids, + rdm1: torch.Tensor, + nuc_grad_feats: set[Feature] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + return _veff_and_expl_nuc_grad( + functional, + mol, + grid, + rdm1, + nuc_grad_feats, + max_memory_in_mb=int(mol.max_memory), + ) + + def test_grid_coords_gradient(mol_name: str) -> None: class TestFunc(ExcFunctionalBase): def __init__(self) -> None: @@ -168,7 +186,7 @@ def get_exc(self, mol: FeatureMap) -> torch.Tensor: mol = get_mol(mol_name) grid, rdm1 = get_grid_and_rdm1(mol) exc_test = TestFunc() - ana_grad = veff_and_expl_nuc_grad(exc_test, mol, grid, rdm1)[1] + ana_grad = _evaluate_nuclear_gradient(exc_test, mol, grid, rdm1)[1] # calculate exact result atom_grids_tab = grid.gen_atomic_grids( @@ -195,7 +213,7 @@ def get_exc(self, mol: FeatureMap) -> torch.Tensor: mol = get_mol(mol_name) grid, rdm1 = get_grid_and_rdm1(mol) exc_test = TestFunc() - ana_grad = veff_and_expl_nuc_grad(exc_test, mol, grid, rdm1)[1] + ana_grad = _evaluate_nuclear_gradient(exc_test, mol, grid, rdm1)[1] # calculate exact result exact_grad = torch.ones_like(ana_grad) @@ -220,8 +238,8 @@ def finite_difference_nuc_grad( ) -> tuple[torch.Tensor, torch.Tensor]: """Calculates the gradient in Exc w.r.t. nuclear coordinates numerically""" # mol_.verbose = 2 - mol_feats = generate_features( - mol, rdm1, minimal_grid(mol), set(weight_sum.features), gpu=True + mol_feats = evaluate_model_features( + mol, rdm1, minimal_grid(mol), FeatureSpec(weight_sum.features) ) def weight_sum_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: @@ -236,7 +254,7 @@ def weight_sum_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: grid, rdm1 = get_grid_and_rdm1(mol) exc_test = TestFunc() - ana_grad = veff_and_expl_nuc_grad(exc_test, mol, grid, rdm1)[1] + ana_grad = _evaluate_nuclear_gradient(exc_test, mol, grid, rdm1)[1] # calculate numerical derivative as accurate as possible num_grad, num_err = finite_difference_nuc_grad(exc_test, mol, rdm1) @@ -278,8 +296,8 @@ def finite_difference_nuc_grad( def dens_sum_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: """Exc wrapper for the finite difference""" mol_.set_geom_(nuc_coords.cpu().numpy(), "bohr", symmetry=None) - mol_feats = generate_features( - mol_, rdm1, grid, set(dens_sum.features), gpu=True + mol_feats = evaluate_model_features( + mol_, rdm1, grid, FeatureSpec(dens_sum.features) ) return dens_sum.get_exc(mol_feats) @@ -294,7 +312,7 @@ def dens_sum_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: num_grad, num_err = finite_difference_nuc_grad(exc_test, mol, rdm1) # calculate analytic result - veff = veff_and_expl_nuc_grad( + veff = _evaluate_nuclear_gradient( exc_test, mol, grid, rdm1, nuc_grad_feats={Feature.DENSITY} )[0] ana_grad = 2 * nuc_grad_from_veff(mol, veff, rdm1) @@ -339,8 +357,8 @@ def finite_difference_nuc_grad( def grad_func_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: """Exc wrapper for the finite difference""" mol_.set_geom_(nuc_coords.cpu().numpy(), "bohr", symmetry=None) - mol_feats = generate_features( - mol_, rdm1, grid, set(grad_func.features), gpu=True + mol_feats = evaluate_model_features( + mol_, rdm1, grid, FeatureSpec(grad_func.features) ) return grad_func.get_exc(mol_feats) @@ -356,7 +374,7 @@ def grad_func_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: num_grad, num_err = finite_difference_nuc_grad(exc_test, mol, rdm1) # calculate analytic result - veff = veff_and_expl_nuc_grad( + veff = _evaluate_nuclear_gradient( exc_test, mol, grid, rdm1, nuc_grad_feats={Feature.GRAD} )[0] ana_grad = 2 * nuc_grad_from_veff(mol, veff, rdm1) @@ -396,8 +414,8 @@ def finite_difference_nuc_grad( def kin_func_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: """Exc wrapper for the finite difference""" mol_.set_geom_(nuc_coords.cpu().numpy(), "bohr", symmetry=None) - mol_feats = generate_features( - mol_, rdm1, grid, set(kin_func.features), gpu=True + mol_feats = evaluate_model_features( + mol_, rdm1, grid, FeatureSpec(kin_func.features) ) return kin_func.get_exc(mol_feats) @@ -413,7 +431,7 @@ def kin_func_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: num_grad, num_err = finite_difference_nuc_grad(exc_test, mol, rdm1) # calculate analytic result - veff = veff_and_expl_nuc_grad( + veff = _evaluate_nuclear_gradient( exc_test, mol, grid, rdm1, nuc_grad_feats={Feature.KIN} )[0] ana_grad = 2 * nuc_grad_from_veff(mol, veff, rdm1) @@ -579,7 +597,7 @@ def get_exc(self, mol: FeatureMap) -> torch.Tensor: # Warmup to avoid counting one-time allocations from CUDA runtime/libraries. for _ in range(2): - veff = veff_and_expl_nuc_grad( + veff = _evaluate_nuclear_gradient( exc_test, mol, grid, rdm1, nuc_grad_feats={Feature.GRAD} )[0] _ = 2 * nuc_grad_from_veff(mol, veff, rdm1) @@ -590,7 +608,7 @@ def get_exc(self, mol: FeatureMap) -> torch.Tensor: allocations: list[int] = [] torch.cuda.reset_peak_memory_stats() for _ in range(5): - veff = veff_and_expl_nuc_grad( + veff = _evaluate_nuclear_gradient( exc_test, mol, grid, rdm1, nuc_grad_feats={Feature.GRAD} )[0] _ = 2 * nuc_grad_from_veff(mol, veff, rdm1) diff --git a/skala/tests/test_gradient_core.py b/skala/tests/test_gradient_core.py index 214d217b..2e5d9261 100644 --- a/skala/tests/test_gradient_core.py +++ b/skala/tests/test_gradient_core.py @@ -2,15 +2,89 @@ """Tests for backend-independent nuclear-gradient operations.""" +from typing import cast + import pytest import torch -from skala.features import Feature +from skala.features import Feature, FeatureMap +from skala.functional.base import ExcFunctionalBase +from skala.pyscf import gradient_core +from skala.pyscf.backend import Grid +from skala.pyscf.feature_math import feature_derivatives from skala.pyscf.gradient_core import ( contract_ao_derivative_block, - feature_derivatives, grid_derivative_block, ) +from pyscf import gto + + +@pytest.mark.parametrize( + ("device_type", "expected_budgets"), + [ + pytest.param("cpu", [321, 321], id="cpu"), + pytest.param("cuda", [321, None], id="cuda"), + ], +) +def test_nuclear_feature_derivatives_select_memory_budget( + monkeypatch: pytest.MonkeyPatch, + device_type: str, + expected_budgets: list[int | None], +) -> None: + observed_budgets: list[int | None] = [] + density = torch.ones(1, dtype=torch.float64) + + class TestDensityMatrix: + device = torch.device(device_type) + + class TestFunctional(ExcFunctionalBase): + features = [Feature.DENSITY] + + def get_exc(self, mol: FeatureMap) -> torch.Tensor: + return mol[Feature.DENSITY].sum() + + def fake_evaluate_model_features( + mol: gto.Mole, + dm: torch.Tensor, + grid: Grid, + feature_spec: object, + max_memory_in_mb: int = 2000, + ) -> FeatureMap: + observed_budgets.append(max_memory_in_mb) + return { + Feature.DENSITY: density, + Feature.ATOMIC_GRID_SIZES: torch.ones(1, dtype=torch.long), + } + + def fake_evaluate_chunked_feature_gradients( + functional: ExcFunctionalBase, + dm: torch.Tensor, + model_features: FeatureMap, + differentiable_features: set[Feature], + max_memory_in_mb: int | None = None, + ) -> FeatureMap: + observed_budgets.append(max_memory_in_mb) + return {Feature.DENSITY: density} + + monkeypatch.setattr( + gradient_core, "evaluate_model_features", fake_evaluate_model_features + ) + monkeypatch.setattr( + gradient_core, + "evaluate_chunked_feature_gradients", + fake_evaluate_chunked_feature_gradients, + ) + + gradient_core.evaluate_nuclear_feature_derivatives( + TestFunctional(), + cast(gto.Mole, object()), + cast(Grid, object()), + cast(torch.Tensor, TestDensityMatrix()), + max_memory_in_mb=321, + ) + + assert observed_budgets == expected_budgets + def test_feature_derivatives() -> None: density = torch.tensor([1.0, 2.0], dtype=torch.float64) @@ -96,6 +170,8 @@ def test_contract_ao_derivative_block_matches_reference() -> None: dtype=torch.float64, ), Feature.KIN: torch.tensor([[0.0, 67.0], [0.0, 71.0]], dtype=torch.float64), + Feature.GRID_COORDS: torch.ones((2, 3), dtype=torch.float64), + Feature.GRID_WEIGHTS: torch.ones(2, dtype=torch.float64), } actual = contract_ao_derivative_block(ao, grid_derivative_block(derivatives, 1, 2)) expected = torch.tensor( diff --git a/skala/tests/test_model_chunking.py b/skala/tests/test_model_chunking.py index 6bd7d59b..b86ce19c 100644 --- a/skala/tests/test_model_chunking.py +++ b/skala/tests/test_model_chunking.py @@ -5,14 +5,65 @@ import pytest import torch from skala.features import Feature, FeatureMap +from skala.functional.base import ExcFunctionalBase from skala.pyscf import model_chunking from skala.pyscf.backend import Grid from skala.pyscf.evaluation import FeatureSpec -from skala.pyscf.feature_math import MGGAFeatureFunction +from skala.pyscf.feature_math import feature_derivatives from pyscf import gto +def test_model_feature_plan_rejects_features_missing_from_evaluation() -> None: + with pytest.raises( + ValueError, match="Model features missing from evaluation: grid_weights" + ): + model_chunking.ModelFeaturePlan( + evaluation_feature_spec=FeatureSpec([Feature.DENSITY]), + model_feature_spec=FeatureSpec([Feature.DENSITY, Feature.GRID_WEIGHTS]), + ) + + +def test_feature_derivatives() -> None: + density = torch.tensor([1.0, 2.0], dtype=torch.float64) + weights = torch.tensor([3.0, 4.0], dtype=torch.float64) + + derivatives = feature_derivatives( + lambda features: ( + features[Feature.DENSITY].square() * features[Feature.GRID_WEIGHTS] + ).sum(), + {Feature.DENSITY: density, Feature.GRID_WEIGHTS: weights}, + ) + + torch.testing.assert_close(derivatives[Feature.DENSITY], 2 * density * weights) + torch.testing.assert_close(derivatives[Feature.GRID_WEIGHTS], density.square()) + assert feature_derivatives(lambda _: torch.tensor(0.0), {}) == {} + + +def test_feature_derivatives_rejects_disconnected_feature() -> None: + used = torch.tensor(2.0) + unused = torch.tensor(3.0) + + with pytest.raises( + RuntimeError, + match="XC energy is disconnected from requested features: grid_weights", + ): + feature_derivatives( + lambda features: features[Feature.DENSITY].square(), + {Feature.DENSITY: used, Feature.GRID_WEIGHTS: unused}, + ) + + +def test_feature_derivatives_rejects_constant_energy() -> None: + feature = torch.tensor(2.0) + + with pytest.raises( + RuntimeError, + match="XC energy is disconnected from requested features: density", + ): + feature_derivatives(lambda _: torch.tensor(1.0), {Feature.DENSITY: feature}) + + def test_model_feature_chunker_sorts_complete_atomic_grids( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -43,24 +94,27 @@ def fake_estimate_max_model_atoms_per_chunk( fake_estimate_max_model_atoms_per_chunk, ) - feature_spec = FeatureSpec( - { - Feature.DENSITY, - Feature.GRID_COORDS, - Feature.GRID_WEIGHTS, - Feature.ATOMIC_GRID_WEIGHTS, - Feature.COARSE_0_ATOMIC_COORDS, - Feature.ATOMIC_GRID_SIZES, - Feature.ATOMIC_GRID_SIZE_BOUND_SHAPE, - } - ) + features = { + Feature.DENSITY, + Feature.GRID_COORDS, + Feature.GRID_WEIGHTS, + Feature.ATOMIC_GRID_WEIGHTS, + Feature.COARSE_0_ATOMIC_COORDS, + Feature.ATOMIC_GRID_SIZES, + Feature.ATOMIC_GRID_SIZE_BOUND_SHAPE, + } + evaluation_feature_spec = FeatureSpec(features) + feature_spec = FeatureSpec(features) raw_features = point_ids.reshape(1, -1) chunker = model_chunking.ModelFeatureChunker( mol=cast(gto.Mole, object()), dm=torch.eye(1, dtype=torch.float64), grids=cast(Grid, object()), atom_major_raw_features=raw_features, - feature_function=MGGAFeatureFunction(feature_spec), + feature_plan=model_chunking.ModelFeaturePlan( + evaluation_feature_spec=evaluation_feature_spec, + model_feature_spec=feature_spec, + ), deriv_order=1, ) @@ -79,6 +133,145 @@ def fake_estimate_max_model_atoms_per_chunk( assert torch.equal(chunks[2].grid_indices, torch.tensor([0, 1, 2])) +def test_model_feature_chunker_builds_bound_shape_from_internal_grid_sizes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + atomic_grid_sizes = torch.tensor([3, 1, 2, 1]) + + monkeypatch.setattr( + model_chunking, + "get_grid_features", + lambda *args, **kwargs: {Feature.ATOMIC_GRID_SIZES: atomic_grid_sizes}, + ) + + evaluation_feature_spec = FeatureSpec( + { + Feature.DENSITY, + Feature.ATOMIC_GRID_SIZES, + Feature.ATOMIC_GRID_SIZE_BOUND_SHAPE, + } + ) + chunker = model_chunking.ModelFeatureChunker( + mol=cast(gto.Mole, type("FakeMole", (), {"natm": 4})()), + dm=torch.eye(1, dtype=torch.float64), + grids=cast(Grid, object()), + atom_major_raw_features=torch.arange(7, dtype=torch.float64).reshape(1, -1), + feature_plan=model_chunking.ModelFeaturePlan( + evaluation_feature_spec=evaluation_feature_spec, + model_feature_spec=FeatureSpec({Feature.ATOMIC_GRID_SIZE_BOUND_SHAPE}), + ), + deriv_order=1, + ) + + (chunk,) = tuple(chunker) + assert set(chunk.model_features) == {Feature.ATOMIC_GRID_SIZE_BOUND_SHAPE} + assert chunk.model_features[Feature.ATOMIC_GRID_SIZE_BOUND_SHAPE].shape == (3, 0) + + +def test_chunked_feature_gradients_match_unchunked( + monkeypatch: pytest.MonkeyPatch, +) -> None: + atomic_grid_sizes = torch.tensor([3, 1, 2, 1]) + point_ids = torch.arange(7, dtype=torch.float64) + atom_ids = torch.arange(4, dtype=torch.float64) + differentiable_features: FeatureMap = { + Feature.DENSITY: point_ids.reshape(1, -1).requires_grad_(), + Feature.GRID_WEIGHTS: (point_ids + 10).requires_grad_(), + Feature.COARSE_0_ATOMIC_COORDS: atom_ids[:, None] + .expand(-1, 3) + .clone() + .requires_grad_(), + } + model_features: FeatureMap = { + **differentiable_features, + Feature.ATOMIC_GRID_SIZES: atomic_grid_sizes, + } + + class TestFunctional(ExcFunctionalBase): + features = list(model_features) + + def __init__(self) -> None: + super().__init__() + self.calls = 0 + + def get_exc(self, mol: FeatureMap) -> torch.Tensor: + self.calls += 1 + return ( + mol[Feature.DENSITY].square() * mol[Feature.GRID_WEIGHTS] + ).sum() + mol[Feature.COARSE_0_ATOMIC_COORDS].square().sum() + + monkeypatch.setattr( + model_chunking, + "estimate_max_model_atoms_per_chunk", + lambda **kwargs: {1: 2, 2: 1, 3: 1}, + ) + functional = TestFunctional() + reference = torch.autograd.grad( + functional.get_exc(model_features), tuple(differentiable_features.values()) + ) + + actual = model_chunking.evaluate_chunked_feature_gradients( + functional, + dm=torch.eye(1, dtype=torch.float64), + model_features=model_features, + differentiable_features=set(differentiable_features), + max_memory_in_mb=100, + ) + + for feature_name, reference_gradient in zip( + differentiable_features, reference, strict=True + ): + torch.testing.assert_close(actual[feature_name], reference_gradient) + assert functional.calls == 4 + + +@pytest.mark.parametrize("supports_spatial_decomposition", [False, True]) +@pytest.mark.parametrize( + ("constant_energy", "disconnected_features"), + [(False, "grid_weights"), (True, "density, grid_weights")], +) +def test_chunked_feature_gradients_reject_disconnected_features( + monkeypatch: pytest.MonkeyPatch, + supports_spatial_decomposition: bool, + constant_energy: bool, + disconnected_features: str, +) -> None: + density = torch.tensor([[2.0, 3.0]], dtype=torch.float64) + grid_weights = torch.tensor([4.0, 5.0], dtype=torch.float64) + model_features: FeatureMap = { + Feature.DENSITY: density, + Feature.GRID_WEIGHTS: grid_weights, + Feature.ATOMIC_GRID_SIZES: torch.tensor([2]), + } + + class TestFunctional(ExcFunctionalBase): + features = [Feature.DENSITY, Feature.GRID_WEIGHTS] + if supports_spatial_decomposition: + features.append(Feature.ATOMIC_GRID_SIZES) + + def get_exc(self, mol: FeatureMap) -> torch.Tensor: + if constant_energy: + return mol[Feature.DENSITY].new_tensor(1.0) + return mol[Feature.DENSITY].square().sum() + + monkeypatch.setattr( + model_chunking, + "estimate_max_model_atoms_per_chunk", + lambda **kwargs: {2: 1}, + ) + + with pytest.raises( + RuntimeError, + match=f"XC energy is disconnected from requested features: {disconnected_features}", + ): + model_chunking.evaluate_chunked_feature_gradients( + TestFunctional(), + dm=torch.eye(1, dtype=torch.float64), + model_features=model_features, + differentiable_features={Feature.DENSITY, Feature.GRID_WEIGHTS}, + ) + + def test_atom_grid_chunks_pack_equal_sizes_up_to_cap() -> None: chunks = model_chunking._make_atom_grid_chunks( torch.tensor([2, 2, 2, 2, 2]), max_atoms_per_grid_size={2: 2} diff --git a/skala/tests/test_pyscf_gradients.py b/skala/tests/test_pyscf_gradients.py index b3d064a7..a114821e 100644 --- a/skala/tests/test_pyscf_gradients.py +++ b/skala/tests/test_pyscf_gradients.py @@ -6,12 +6,13 @@ from skala.features import Feature, FeatureMap from skala.functional.base import ExcFunctionalBase from skala.pyscf import SkalaKS -from skala.pyscf.features import generate_features +from skala.pyscf.evaluation import FeatureSpec from skala.pyscf.gradients import ( SkalaRKSGradient, SkalaUKSGradient, - veff_and_expl_nuc_grad, + _veff_and_expl_nuc_grad, ) +from skala.pyscf.model_chunking import evaluate_model_features from pyscf import dft, gto, scf from tests.ridders import num_grad_ridders @@ -57,6 +58,23 @@ def get_grid_and_rdm1(mol: gto.Mole) -> tuple[dft.Grids, torch.Tensor]: return mf.grids, rdm1 # maybe_expand_and_divide(rdm1, len(rdm1.shape) == 2, 2) +def _evaluate_nuclear_gradient( + functional: ExcFunctionalBase, + mol: gto.Mole, + grid: dft.Grids, + rdm1: torch.Tensor, + nuc_grad_feats: set[Feature] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + return _veff_and_expl_nuc_grad( + functional, + mol, + grid, + rdm1, + nuc_grad_feats, + max_memory_in_mb=int(mol.max_memory), + ) + + def test_grid_coords_gradient(mol_name: str) -> None: class TestFunc(ExcFunctionalBase): def __init__(self) -> None: @@ -70,7 +88,7 @@ def get_exc(self, mol: FeatureMap) -> torch.Tensor: mol = get_mol(mol_name) grid, rdm1 = get_grid_and_rdm1(mol) exc_test = TestFunc() - ana_grad = veff_and_expl_nuc_grad(exc_test, mol, grid, rdm1)[1] + ana_grad = _evaluate_nuclear_gradient(exc_test, mol, grid, rdm1)[1] # calculate exact result atom_grids_tab = grid.gen_atomic_grids( @@ -97,7 +115,7 @@ def get_exc(self, mol: FeatureMap) -> torch.Tensor: mol = get_mol(mol_name) grid, rdm1 = get_grid_and_rdm1(mol) exc_test = TestFunc() - ana_grad = veff_and_expl_nuc_grad(exc_test, mol, grid, rdm1)[1] + ana_grad = _evaluate_nuclear_gradient(exc_test, mol, grid, rdm1)[1] # calculate exact result exact_grad = torch.ones_like(ana_grad) @@ -122,8 +140,8 @@ def finite_difference_nuc_grad( ) -> tuple[torch.Tensor, torch.Tensor]: """Calculates the gradient in Exc w.r.t. nuclear coordinates numerically""" # mol_.verbose = 2 - mol_feats = generate_features( - mol, rdm1, minimal_grid(mol), set(weight_sum.features) + mol_feats = evaluate_model_features( + mol, rdm1, minimal_grid(mol), FeatureSpec(weight_sum.features) ) def weight_sum_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: @@ -140,7 +158,7 @@ def weight_sum_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: grid, rdm1 = get_grid_and_rdm1(mol) exc_test = TestFunc() - ana_grad = veff_and_expl_nuc_grad(exc_test, mol, grid, rdm1)[1] + ana_grad = _evaluate_nuclear_gradient(exc_test, mol, grid, rdm1)[1] # calculate numerical derivative as accurate as possible num_grad, num_err = finite_difference_nuc_grad(exc_test, mol, rdm1) @@ -193,7 +211,9 @@ def finite_difference_nuc_grad( def dens_sum_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: """Exc wrapper for the finite difference""" mol_.set_geom_(nuc_coords.numpy(), "bohr", symmetry=None) - mol_feats = generate_features(mol_, rdm1, grid, set(dens_sum.features)) + mol_feats = evaluate_model_features( + mol_, rdm1, grid, FeatureSpec(dens_sum.features) + ) return dens_sum.get_exc(mol_feats) @@ -207,7 +227,7 @@ def dens_sum_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: num_grad, num_err = finite_difference_nuc_grad(exc_test, mol, rdm1) # calculate analytic result - veff = veff_and_expl_nuc_grad( + veff = _evaluate_nuclear_gradient( exc_test, mol, grid, rdm1, nuc_grad_feats={Feature.DENSITY} )[0] ana_grad = nuc_grad_from_veff(mol, veff, rdm1) @@ -248,7 +268,9 @@ def finite_difference_nuc_grad( def grad_func_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: """Exc wrapper for the finite difference""" mol_.set_geom_(nuc_coords.numpy(), "bohr", symmetry=None) - mol_feats = generate_features(mol_, rdm1, grid, set(grad_func.features)) + mol_feats = evaluate_model_features( + mol_, rdm1, grid, FeatureSpec(grad_func.features) + ) return grad_func.get_exc(mol_feats) @@ -263,7 +285,7 @@ def grad_func_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: num_grad, num_err = finite_difference_nuc_grad(exc_test, mol, rdm1) # calculate analytic result - veff = veff_and_expl_nuc_grad( + veff = _evaluate_nuclear_gradient( exc_test, mol, grid, rdm1, nuc_grad_feats={Feature.GRAD} )[0] ana_grad = nuc_grad_from_veff(mol, veff, rdm1) @@ -311,7 +333,9 @@ def finite_difference_nuc_grad( def kin_func_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: """Exc wrapper for the finite difference""" mol_.set_geom_(nuc_coords.numpy(), "bohr", symmetry=None) - mol_feats = generate_features(mol_, rdm1, grid, set(kin_func.features)) + mol_feats = evaluate_model_features( + mol_, rdm1, grid, FeatureSpec(kin_func.features) + ) return kin_func.get_exc(mol_feats) @@ -326,7 +350,7 @@ def kin_func_as_nuc_coords_func(nuc_coords: torch.Tensor) -> torch.Tensor: num_grad, num_err = finite_difference_nuc_grad(exc_test, mol, rdm1) # calculate analytic result - veff = veff_and_expl_nuc_grad( + veff = _evaluate_nuclear_gradient( exc_test, mol, grid, rdm1, nuc_grad_feats={Feature.KIN} )[0] ana_grad = nuc_grad_from_veff(mol, veff, rdm1) @@ -426,7 +450,7 @@ def get_exc(self, mol: FeatureMap) -> torch.Tensor: grid = minimal_grid(mol, sort_grids=False) exc_test = TestFunc() - _, nuc_grad = veff_and_expl_nuc_grad(exc_test, mol, grid, rdm1) + _, nuc_grad = _evaluate_nuclear_gradient(exc_test, mol, grid, rdm1) assert torch.allclose(nuc_grad, torch.zeros_like(nuc_grad), atol=1e-15), ( f"atomic_grid_weights gradient should be zero, got {nuc_grad}" @@ -464,7 +488,7 @@ def get_exc(self, mol: FeatureMap) -> torch.Tensor: exc_test = TestFunc() # This should not raise NotImplementedError - veff, nuc_grad = veff_and_expl_nuc_grad(exc_test, mol, grid, rdm1) + veff, nuc_grad = _evaluate_nuclear_gradient(exc_test, mol, grid, rdm1) if mol.spin == 0: assert veff.shape == (3, mol.nao, mol.nao) @@ -496,7 +520,7 @@ def get_exc(self, mol: FeatureMap) -> torch.Tensor: exc_test = TestFunc() # Explicitly pass all features including integer ones — should auto-discard them - _vexc, nuc_grad = veff_and_expl_nuc_grad( + _vexc, nuc_grad = _evaluate_nuclear_gradient( exc_test, mol, grid, rdm1, nuc_grad_feats=set(exc_test.features) ) assert nuc_grad.shape == (mol.natm, 3) diff --git a/skala/tests/test_xc_integrator.py b/skala/tests/test_xc_integrator.py index c4049a66..c1518432 100644 --- a/skala/tests/test_xc_integrator.py +++ b/skala/tests/test_xc_integrator.py @@ -1,8 +1,12 @@ import pytest import torch -from skala.features import Feature, FeatureMap +from skala.features import Feature +from skala.pyscf import ao_evaluation as ao_evaluation_module +from skala.pyscf import feature_math from skala.pyscf import xc_integrator as xc_integrator_module +from skala.pyscf.backend import Grid from skala.pyscf.grids import SkalaGrids +from skala.pyscf.model_chunking import ModelFeatureChunk from skala.pyscf.xc_integrator import XCIntegrator, XCResult from pyscf import dft, gto @@ -87,22 +91,45 @@ def test_xc_integrator_returns_tensors_and_xc_only_response( mol = gto.M(atom="H 0 0 0; H 0 0 0.74", basis="sto-3g", verbose=0) grids = SkalaGrids(mol) - def fake_generate_features( - mol: gto.Mole, + def fake_evaluate_raw_features( dm: torch.Tensor, - grids: object, - features: set[Feature], - **kwargs: object, - ) -> FeatureMap: - return { - Feature.DENSITY: dm.sum().reshape(1), - Feature.GRID_WEIGHTS: torch.tensor([2.0], dtype=dm.dtype), - } + mol: gto.Mole, + grids: Grid, + feature_function: feature_math.LinearFeature, + block_size: int | None = None, + max_memory: int = 2000, + gpu: bool = False, + ) -> torch.Tensor: + return dm.sum().reshape(1) + + class FakeModelFeatureChunker: + def __init__( + self, + mol: gto.Mole, + dm: torch.Tensor, + grids: object, + atom_major_raw_features: torch.Tensor, + **kwargs: object, + ) -> None: + self.raw_features = atom_major_raw_features + + def __iter__(self) -> object: + yield ModelFeatureChunk( + grid_indices=torch.zeros(1, dtype=torch.int64), + raw_features=self.raw_features, + model_features={ + Feature.DENSITY: self.raw_features, + Feature.GRID_WEIGHTS: self.raw_features.new_tensor([2.0]), + }, + ) monkeypatch.setattr( - xc_integrator_module, - "generate_features", - fake_generate_features, + ao_evaluation_module, + "evaluate_raw_features_auto_chunk", + fake_evaluate_raw_features, + ) + monkeypatch.setattr( + xc_integrator_module, "ModelFeatureChunker", FakeModelFeatureChunker ) integrator = XCIntegrator(QuadraticFunctional()) dm = torch.tensor([[1.0, 2.0], [2.0, 3.0]], dtype=torch.float64) @@ -111,7 +138,7 @@ def fake_generate_features( response = integrator.gen_response(mol, grids, dm.detach().clone()) assert isinstance(result, XCResult) - torch.testing.assert_close(result.electron_count, dm.new_tensor(16.0)) + torch.testing.assert_close(result.electron_count, dm.new_tensor([16.0, 16.0])) torch.testing.assert_close(result.energy, dm.new_tensor(128.0)) torch.testing.assert_close(result.potential, torch.full_like(dm, 32.0)) torch.testing.assert_close(response(torch.ones_like(dm)), torch.full_like(dm, 16.0))