From f6e9bbbd9f9fbd8b7355ceeaba6daad385d1e256 Mon Sep 17 00:00:00 2001 From: jenswehner Date: Tue, 1 Sep 2026 12:58:28 +0200 Subject: [PATCH 1/6] refactor gradient computation --- pixi.lock | 46 ++++---- pixi.toml | 2 +- skala/src/skala/gpu4pyscf/gradients.py | 115 ++---------------- skala/src/skala/pyscf/gradient_core.py | 157 +++++++++++++++++++++++++ skala/src/skala/pyscf/gradients.py | 103 +++------------- skala/tests/test_gradient_core.py | 102 ++++++++++++++++ 6 files changed, 310 insertions(+), 215 deletions(-) create mode 100644 skala/src/skala/pyscf/gradient_core.py create mode 100644 skala/tests/test_gradient_core.py diff --git a/pixi.lock b/pixi.lock index 5733fe21..dd83ec99 100644 --- a/pixi.lock +++ b/pixi.lock @@ -6,18 +6,7 @@ platforms: - __linux=4.18 - __glibc=2.28 - __archspec=0=x86_64 -- name: linux-aarch64 - virtual-packages: - - __unix=0=0 - - __linux=4.18 - - __glibc=2.28 - - __archspec=0=aarch64 -- name: osx-arm64 - virtual-packages: - - __unix=0=0 - - __osx=13.0 - - __archspec=0=m1 -- name: p1 +- name: linux-64-cuda12 subdir: linux-64 virtual-packages: - __cuda=12 @@ -25,7 +14,7 @@ platforms: - __linux=4.18 - __glibc=2.28 - __archspec=0=x86_64 -- name: p2 +- name: linux-64-cuda13 subdir: linux-64 virtual-packages: - __cuda=13 @@ -33,6 +22,17 @@ platforms: - __linux=4.18 - __glibc=2.28 - __archspec=0=x86_64 +- name: linux-aarch64 + virtual-packages: + - __unix=0=0 + - __linux=4.18 + - __glibc=2.28 + - __archspec=0=aarch64 +- name: osx-arm64 + virtual-packages: + - __unix=0=0 + - __osx=13.0 + - __archspec=0=m1 environments: assets: channels: @@ -270,7 +270,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/40/2f/10a2e9ad81bd5c2479a255d18c4f02c734a3a21bf4b36a236ba615919592/pyscf_dispersion-1.5.0-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - pypi: https://files.pythonhosted.org/packages/88/30/8e48717aff32f11fd78f981192d0b567971eb285ea048414b7dac2a211cf/pyscf-2.14.0-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - pypi: https://files.pythonhosted.org/packages/9e/e9/1a19e42cd43cc1365e127db6aae85e1c671da1d9a5d746f4d34a50edb577/h5py-3.16.0-cp312-cp312-manylinux_2_28_x86_64.whl - p1: + linux-64-cuda12: - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-7_kmp_llvm.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/binutils_impl_linux-64-2.46.1-default_hfdba357_102.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-1.2.0-h505cf86_3.conda @@ -1618,7 +1618,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.16.0-h69aa097_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.16.0-pyhcf101f3_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2026c-h151e31d_0.conda - p1: + linux-64-cuda12: - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-7_kmp_llvm.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/binutils_impl_linux-64-2.46.1-default_hfdba357_102.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-hda65f42_10.conda @@ -1735,7 +1735,7 @@ environments: channels: - url: https://conda.anaconda.org/conda-forge/ packages: - p1: + linux-64-cuda12: - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-7_kmp_llvm.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-auth-0.10.4-h4610da3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-cal-0.9.15-h6bbde05_1.conda @@ -2023,7 +2023,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.16.0-h69aa097_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.16.0-pyhcf101f3_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2026c-h151e31d_0.conda - p1: + linux-64-cuda12: - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-7_kmp_llvm.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-auth-0.10.4-h4610da3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-cal-0.9.15-h6bbde05_1.conda @@ -2272,7 +2272,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.16.0-h69aa097_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.16.0-pyhcf101f3_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2026c-h151e31d_0.conda - p1: + linux-64-cuda12: - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-7_kmp_llvm.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-auth-0.10.4-h4610da3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-cal-0.9.15-h6bbde05_1.conda @@ -2507,7 +2507,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.16.0-h69aa097_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.16.0-pyhcf101f3_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2026c-h151e31d_0.conda - p1: + linux-64-cuda12: - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-7_kmp_llvm.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-auth-0.10.4-h4610da3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-cal-0.9.15-h6bbde05_1.conda @@ -2728,7 +2728,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.16.0-h69aa097_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.16.0-pyhcf101f3_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2026c-h151e31d_0.conda - p1: + linux-64-cuda12: - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-7_kmp_llvm.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-auth-0.10.4-h4610da3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-cal-0.9.15-h6bbde05_1.conda @@ -2844,7 +2844,7 @@ environments: indexes: - https://pypi.org/simple packages: - p1: + linux-64-cuda12: - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-7_kmp_llvm.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-auth-0.10.4-hb7a77c6_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-cal-0.9.14-h2aa3ae6_4.conda @@ -3111,7 +3111,7 @@ environments: indexes: - https://pypi.org/simple packages: - p1: + linux-64-cuda12: - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-7_kmp_llvm.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-auth-0.10.4-hb7a77c6_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-cal-0.9.14-h2aa3ae6_4.conda @@ -3378,7 +3378,7 @@ environments: indexes: - https://pypi.org/simple packages: - p2: + linux-64-cuda13: - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-7_kmp_llvm.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-auth-0.10.4-hb7a77c6_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-cal-0.9.14-h2aa3ae6_4.conda diff --git a/pixi.toml b/pixi.toml index 728c55f2..a4c99fe8 100644 --- a/pixi.toml +++ b/pixi.toml @@ -10,7 +10,7 @@ platforms = [ { name = "linux-64-cuda12", platform = "linux-64", cuda = "12" }, { name = "linux-64-cuda13", platform = "linux-64", cuda = "13" }, ] -requires-pixi = ">=0.75,<0.76" +requires-pixi = ">=0.78,<0.79" preview = ["pixi-build"] [workspace.conda-pypi-map] diff --git a/skala/src/skala/gpu4pyscf/gradients.py b/skala/src/skala/gpu4pyscf/gradients.py index 38ff3224..616ee5f1 100644 --- a/skala/src/skala/gpu4pyscf/gradients.py +++ b/skala/src/skala/gpu4pyscf/gradients.py @@ -20,6 +20,11 @@ from skala.dispersion import DFTD3Dispersion from skala.features import Feature, FeatureMap from skala.functional.base import ExcFunctionalBase +from skala.pyscf.gradient_core import ( + contract_veff_block, + feature_derivatives, + grid_derivative_block, +) LOG = logging.getLogger(__name__) @@ -103,25 +108,8 @@ def exc_feat_func(*nuc_feat_tensors: torch.Tensor) -> torch.Tensor: ) return functional.get_exc(exc_mol_feats) - # torch.func.vjp wraps primals in functional tensors that may not expose - # backing storage, which TorchScript traced models can reject. - diff_nuc_feat_tensors = [ - feat.detach().requires_grad_(True) for feat in nuc_feat_tensors - ] - if len(diff_nuc_feat_tensors) > 0: - exc = exc_feat_func(*diff_nuc_feat_tensors) - dExc_tuple = torch.autograd.grad( - exc, - tuple(diff_nuc_feat_tensors), - create_graph=False, - retain_graph=False, - allow_unused=False, - ) - else: - dExc_tuple = () - dExc: FeatureMap = {} - for i in range(len(dExc_tuple)): - dExc[nuc_feat_names[i]] = dExc_tuple[i].detach() + dExc_tuple = feature_derivatives(exc_feat_func, nuc_feat_tensors) + dExc: FeatureMap = dict(zip(nuc_feat_names, dExc_tuple, strict=True)) LOG.debug("autograd gradients for nuclear features done") @@ -143,94 +131,17 @@ def exc_feat_func(*nuc_feat_tensors: torch.Tensor) -> torch.Tensor: 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 = torch.zeros((2, 3, nao, nao), dtype=rdm1.dtype, device=rdm1.device) - - if Feature.DENSITY in nuc_grad_feats: - veff_atm += torch.einsum( - "si, xip, iq -> sxpq", - dExc[Feature.DENSITY][:, atm_start:atm_end], - ao[1:4], - ao[0], - ) - - if Feature.GRAD in nuc_grad_feats: - Exc_dgrad_atm = dExc[Feature.GRAD][:, :, atm_start:atm_end] - - veff_atm += torch.einsum( - "syi, xip, yiq -> sxpq", Exc_dgrad_atm, ao[1:4], ao[1:4] - ) - # XX, XY, XZ = 4, 5, 6 - veff_atm[:, 0] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 0], ao[4], ao[0] - ) - veff_atm[:, 0] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 1], ao[5], ao[0] - ) - veff_atm[:, 0] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 2], ao[6], ao[0] - ) - # YX, YY, YZ = 5, 7, 8 - veff_atm[:, 1] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 0], ao[5], ao[0] - ) - veff_atm[:, 1] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 1], ao[7], ao[0] - ) - veff_atm[:, 1] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 2], ao[8], ao[0] - ) - # ZX, ZY, ZZ = 6, 8, 9 - veff_atm[:, 2] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 0], ao[6], ao[0] - ) - veff_atm[:, 2] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 1], ao[8], ao[0] - ) - veff_atm[:, 2] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 2], ao[9], ao[0] - ) - - if Feature.KIN in nuc_grad_feats: - Exc_dkin_atm = dExc[Feature.KIN][:, atm_start:atm_end] - # XX, XY, XZ = 4, 5, 6 - veff_atm[:, 0] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[4], ao[1]) / 2 - ) - veff_atm[:, 0] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[5], ao[2]) / 2 - ) - veff_atm[:, 0] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[6], ao[3]) / 2 - ) - # YX, YY, YZ = 5, 7, 8 - veff_atm[:, 1] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[5], ao[1]) / 2 - ) - veff_atm[:, 1] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[7], ao[2]) / 2 - ) - veff_atm[:, 1] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[8], ao[3]) / 2 - ) - # ZX, ZY, ZZ = 6, 8, 9 - veff_atm[:, 2] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[6], ao[1]) / 2 - ) - veff_atm[:, 2] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[8], ao[2]) / 2 - ) - veff_atm[:, 2] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[9], ao[3]) / 2 - ) + veff_atm = contract_veff_block(ao, dExc_atm) - if Feature.GRID_COORDS in nuc_grad_feats: + if Feature.GRID_COORDS in dExc_atm: # also add the explicit grid coordinate dependence - nuc_grad[atm_id] += dExc[Feature.GRID_COORDS][atm_start:atm_end].sum(dim=0) + nuc_grad[atm_id] += dExc_atm[Feature.GRID_COORDS].sum(dim=0) - if Feature.GRID_WEIGHTS in nuc_grad_feats: - Exc_dgw = dExc[Feature.GRID_WEIGHTS][atm_start:atm_end] + 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, diff --git a/skala/src/skala/pyscf/gradient_core.py b/skala/src/skala/pyscf/gradient_core.py new file mode 100644 index 00000000..2d5797c0 --- /dev/null +++ b/skala/src/skala/pyscf/gradient_core.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: MIT + +"""Backend-independent PyTorch operations for PySCF nuclear gradients.""" + +from collections.abc import Callable + +import torch + +from skala.features import Feature, FeatureMap + +# PySCF packs AO derivatives as value, x, y, z, xx, xy, xz, yy, yz, zz. +# Each row below selects the Hessian components for one nuclear displacement +# direction, with columns ordered by the x, y, and z density responses. +_HESSIAN_COMPONENTS = ((4, 5, 6), (5, 7, 8), (6, 8, 9)) + +# Grid metadata is laid out with one row per point: coordinates have shape +# (npoints, 3), while weights have shape (npoints,). +_POINT_FIRST_FEATURES = {Feature.GRID_COORDS, Feature.GRID_WEIGHTS} + +# Electronic features keep spin and, for GRAD, Cartesian components before the +# grid dimension: DENSITY and KIN have shape (2, npoints), and GRAD has shape +# (2, 3, npoints). +_POINT_LAST_FEATURES = {Feature.DENSITY, Feature.GRAD, Feature.KIN} + + +def feature_derivatives( + exc_func: Callable[..., torch.Tensor], + feature_tensors: list[torch.Tensor], +) -> tuple[torch.Tensor, ...]: + """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 feature tensors and returning scalar XC energy. + feature_tensors: Feature tensors in the order expected by ``exc_func``. + + Returns: + Feature derivatives in the same order as ``feature_tensors``. + + """ + if not feature_tensors: + return () + + differentiable_features = tuple( + tensor.detach().requires_grad_(True) for tensor in feature_tensors + ) + exc = exc_func(*differentiable_features) + if not exc.requires_grad: + return tuple(torch.zeros_like(feature) for feature in differentiable_features) + + gradients = torch.autograd.grad( + exc, + tuple(differentiable_features), + create_graph=False, + retain_graph=False, + allow_unused=True, + ) + return tuple( + torch.zeros_like(feature) if gradient is None else gradient.detach() + for feature, gradient in zip(differentiable_features, gradients, strict=True) + ) + + +def grid_derivative_block( + derivatives: FeatureMap, grid_start: int, grid_end: int +) -> FeatureMap: + """Slice grid-resolved feature derivatives to one atom's grid block. + + Density-like features store the grid-point dimension last, whereas grid + coordinates and weights store it first. Non-grid features, such as atomic + coordinates, are intentionally omitted. + + Args: + derivatives: XC energy derivatives keyed by molecular feature. + grid_start: Start of this block in the full molecular feature tensors. + grid_end: End of this block in the full molecular feature tensors. + + Returns: + Grid-resolved derivatives restricted to ``[grid_start:grid_end]``. + """ + block = { + feature: derivative[grid_start:grid_end] + for feature, derivative in derivatives.items() + if feature in _POINT_FIRST_FEATURES + } + block.update( + { + feature: derivative[..., grid_start:grid_end] + for feature, derivative in derivatives.items() + if feature in _POINT_LAST_FEATURES + } + ) + return block + + +def contract_veff_block( + ao: torch.Tensor, + derivatives: FeatureMap, +) -> torch.Tensor: + """Contract one atom's AO values with XC feature derivatives. + + Args: + ao: Component-major AO values for this atom's grid block, shaped + ``(ncomponents, npoints, nao)``. Here ``npoints`` must equal + ``grid_end - grid_start`` and ``nao`` is the number of molecular + AOs. PySCF supplies 1, 4, or 10 components for derivative orders + 0, 1, or 2, respectively, packed as ``value, x, y, z, xx, xy, + xz, yy, yz, zz``. The requested features determine how many of + these leading components the contraction reads. + derivatives: XC energy derivatives already sliced to this grid block. + + Returns: + Spin- and Cartesian-resolved effective-potential contribution with shape + ``(2, 3, nao, nao)`` on the same device and with the same dtype as ``ao``. + """ + nao = ao.shape[-1] + veff = ao.new_zeros((2, 3, nao, nao)) + + if Feature.DENSITY in derivatives: + veff += torch.einsum( + "si, xip, iq -> sxpq", + derivatives[Feature.DENSITY], + ao[1:4], + ao[0], + ) + + if Feature.GRAD in derivatives: + exc_dgrad = derivatives[Feature.GRAD] + veff += torch.einsum("syi, xip, yiq -> sxpq", exc_dgrad, ao[1:4], ao[1:4]) + for force_direction, components in enumerate(_HESSIAN_COMPONENTS): + for response_direction, component in enumerate(components): + veff[:, force_direction] += torch.einsum( + "si, ip, iq -> spq", + exc_dgrad[:, response_direction], + ao[component], + ao[0], + ) + + if Feature.KIN in derivatives: + exc_dkin = derivatives[Feature.KIN] + for force_direction, components in enumerate(_HESSIAN_COMPONENTS): + for response_direction, component in enumerate(components): + veff[:, force_direction] += ( + torch.einsum( + "si, ip, iq -> spq", + exc_dkin, + ao[component], + # AO components 1, 2, and 3 are the x, y, and z derivatives. + ao[response_direction + 1], + ) + / 2 + ) + + return veff diff --git a/skala/src/skala/pyscf/gradients.py b/skala/src/skala/pyscf/gradients.py index c7af698b..029f824d 100644 --- a/skala/src/skala/pyscf/gradients.py +++ b/skala/src/skala/pyscf/gradients.py @@ -17,6 +17,11 @@ from skala.dispersion import DFTD3Dispersion from skala.features import Feature, FeatureMap from skala.functional.base import ExcFunctionalBase +from skala.pyscf.gradient_core import ( + contract_veff_block, + feature_derivatives, + grid_derivative_block, +) LOG = logging.getLogger(__name__) @@ -98,13 +103,10 @@ def exc_feat_func(*nuc_feat_tensors: torch.Tensor) -> torch.Tensor: ) return functional.get_exc(exc_mol_feats) - dExc_func = torch.func.vjp(exc_feat_func, *nuc_feat_tensors)[1] - dExc_tuple = dExc_func(torch.tensor(1.0, dtype=rdm1.dtype)) - dExc: FeatureMap = {} - for i in range(len(dExc_tuple)): - dExc[nuc_feat_names[i]] = dExc_tuple[i].detach() + dExc_tuple = feature_derivatives(exc_feat_func, nuc_feat_tensors) + dExc: FeatureMap = dict(zip(nuc_feat_names, dExc_tuple, strict=True)) - LOG.debug("torch.func.vjp done") + LOG.debug("autograd gradients for nuclear features done") nao = rdm1.shape[-1] veff = torch.zeros((2, 3, nao, nao), dtype=rdm1.dtype) @@ -121,94 +123,17 @@ def exc_feat_func(*nuc_feat_tensors: torch.Tensor) -> torch.Tensor: 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 = torch.zeros((2, 3, nao, nao), dtype=rdm1.dtype) - - if Feature.DENSITY in nuc_grad_feats: - veff_atm += torch.einsum( - "si, xip, iq -> sxpq", - dExc[Feature.DENSITY][:, atm_start:atm_end], - ao[1:4], - ao[0], - ) - - if Feature.GRAD in nuc_grad_feats: - Exc_dgrad_atm = dExc[Feature.GRAD][:, :, atm_start:atm_end] - - veff_atm += torch.einsum( - "syi, xip, yiq -> sxpq", Exc_dgrad_atm, ao[1:4], ao[1:4] - ) - # XX, XY, XZ = 4, 5, 6 - veff_atm[:, 0] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 0], ao[4], ao[0] - ) - veff_atm[:, 0] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 1], ao[5], ao[0] - ) - veff_atm[:, 0] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 2], ao[6], ao[0] - ) - # YX, YY, YZ = 5, 7, 8 - veff_atm[:, 1] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 0], ao[5], ao[0] - ) - veff_atm[:, 1] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 1], ao[7], ao[0] - ) - veff_atm[:, 1] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 2], ao[8], ao[0] - ) - # ZX, ZY, ZZ = 6, 8, 9 - veff_atm[:, 2] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 0], ao[6], ao[0] - ) - veff_atm[:, 2] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 1], ao[8], ao[0] - ) - veff_atm[:, 2] += torch.einsum( - "si, ip, iq -> spq", Exc_dgrad_atm[:, 2], ao[9], ao[0] - ) - - if Feature.KIN in nuc_grad_feats: - Exc_dkin_atm = dExc[Feature.KIN][:, atm_start:atm_end] - # XX, XY, XZ = 4, 5, 6 - veff_atm[:, 0] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[4], ao[1]) / 2 - ) - veff_atm[:, 0] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[5], ao[2]) / 2 - ) - veff_atm[:, 0] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[6], ao[3]) / 2 - ) - # YX, YY, YZ = 5, 7, 8 - veff_atm[:, 1] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[5], ao[1]) / 2 - ) - veff_atm[:, 1] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[7], ao[2]) / 2 - ) - veff_atm[:, 1] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[8], ao[3]) / 2 - ) - # ZX, ZY, ZZ = 6, 8, 9 - veff_atm[:, 2] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[6], ao[1]) / 2 - ) - veff_atm[:, 2] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[8], ao[2]) / 2 - ) - veff_atm[:, 2] += ( - torch.einsum("si, ip, iq -> spq", Exc_dkin_atm, ao[9], ao[3]) / 2 - ) + veff_atm = contract_veff_block(ao, dExc_atm) - if Feature.GRID_COORDS in nuc_grad_feats: + if Feature.GRID_COORDS in dExc_atm: # also add the explicit grid coordinate dependence - nuc_grad[atm_id] += dExc[Feature.GRID_COORDS][atm_start:atm_end].sum(dim=0) + nuc_grad[atm_id] += dExc_atm[Feature.GRID_COORDS].sum(dim=0) - if Feature.GRID_WEIGHTS in nuc_grad_feats: - Exc_dgw = dExc[Feature.GRID_WEIGHTS][atm_start:atm_end] + 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, diff --git a/skala/tests/test_gradient_core.py b/skala/tests/test_gradient_core.py new file mode 100644 index 00000000..870b3744 --- /dev/null +++ b/skala/tests/test_gradient_core.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: MIT + +"""Tests for backend-independent nuclear-gradient operations.""" + +import torch +from skala.features import Feature +from skala.pyscf.gradient_core import ( + contract_veff_block, + feature_derivatives, + grid_derivative_block, +) + + +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 rho, grid_weights: (rho.square() * grid_weights).sum(), + [density, weights], + ) + + torch.testing.assert_close(derivatives[0], 2 * density * weights) + torch.testing.assert_close(derivatives[1], density.square()) + assert feature_derivatives(lambda: torch.tensor(0.0), []) == () + + +def test_feature_derivatives_returns_zero_for_disconnected_feature() -> None: + used = torch.tensor(2.0) + unused = torch.tensor(3.0) + + derivatives = feature_derivatives(lambda value, _: value.square(), [used, unused]) + + torch.testing.assert_close(derivatives[0], 2 * used) + torch.testing.assert_close(derivatives[1], torch.zeros_like(unused)) + + +def test_feature_derivatives_returns_zero_for_constant_energy() -> None: + feature = torch.tensor(2.0) + + (derivative,) = feature_derivatives(lambda _: torch.tensor(1.0), [feature]) + + torch.testing.assert_close(derivative, torch.zeros_like(feature)) + + +def test_grid_derivative_block_slices_each_grid_dimension() -> None: + derivatives = { + Feature.DENSITY: torch.arange(8).reshape(2, 4), + Feature.GRAD: torch.arange(24).reshape(2, 3, 4), + Feature.GRID_COORDS: torch.arange(12).reshape(4, 3), + Feature.GRID_WEIGHTS: torch.arange(4), + Feature.COARSE_0_ATOMIC_COORDS: torch.arange(6).reshape(2, 3), + } + + block = grid_derivative_block(derivatives, grid_start=1, grid_end=3) + + assert set(block) == { + Feature.DENSITY, + Feature.GRAD, + Feature.GRID_COORDS, + Feature.GRID_WEIGHTS, + } + torch.testing.assert_close( + block[Feature.DENSITY], derivatives[Feature.DENSITY][..., 1:3] + ) + torch.testing.assert_close(block[Feature.GRAD], derivatives[Feature.GRAD][..., 1:3]) + torch.testing.assert_close( + block[Feature.GRID_COORDS], derivatives[Feature.GRID_COORDS][1:3] + ) + torch.testing.assert_close( + block[Feature.GRID_WEIGHTS], derivatives[Feature.GRID_WEIGHTS][1:3] + ) + + +def test_contract_veff_block_matches_reference() -> None: + ao = torch.tensor( + [2.0, 3.0, 5.0, 7.0, 11.0, 13.0, 17.0, 19.0, 23.0, 29.0], + dtype=torch.float64, + ).reshape(10, 1, 1) + derivatives = { + Feature.DENSITY: torch.tensor([[0.0, 31.0], [0.0, 37.0]], dtype=torch.float64), + Feature.GRAD: torch.tensor( + [ + [[0.0, 41.0], [0.0, 43.0], [0.0, 47.0]], + [[0.0, 53.0], [0.0, 59.0], [0.0, 61.0]], + ], + dtype=torch.float64, + ), + Feature.KIN: torch.tensor([[0.0, 67.0], [0.0, 71.0]], dtype=torch.float64), + } + actual = contract_veff_block(ao, grid_derivative_block(derivatives, 1, 2)) + expected = torch.tensor( + [ + [[[13074.5]], [[18389.5]], [[23562.5]]], + [[[15342.5]], [[21673.5]], [[27838.5]]], + ], + dtype=torch.float64, + ) + + torch.testing.assert_close(actual, expected) + assert actual.dtype == ao.dtype + assert actual.device == ao.device From 6aafccdc8dc3a67e1e4e3a3db25424bdd443c5b6 Mon Sep 17 00:00:00 2001 From: jenswehner Date: Tue, 1 Sep 2026 13:45:32 +0200 Subject: [PATCH 2/6] update pixi workflows --- .github/workflows/benchmark-test.yml | 2 +- .github/workflows/docs.yml | 2 +- .github/workflows/examples.yml | 4 ++-- .github/workflows/gauxc-test.yml | 2 +- .github/workflows/lint.yml | 2 +- .github/workflows/model-benchmark.yml | 4 ++-- .github/workflows/model-examples.yml | 8 ++++---- .github/workflows/model-test.yml | 2 +- .github/workflows/pypi.yml | 4 ++-- .github/workflows/skala-test.yml | 6 +++--- 10 files changed, 18 insertions(+), 18 deletions(-) diff --git a/.github/workflows/benchmark-test.yml b/.github/workflows/benchmark-test.yml index 1ca715bf..4460bef8 100644 --- a/.github/workflows/benchmark-test.yml +++ b/.github/workflows/benchmark-test.yml @@ -35,7 +35,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 environments: default activate-environment: true locked: true diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index d302aa47..95bf9bb4 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -40,7 +40,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 environments: docs activate-environment: true locked: true diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index da7f84ea..9b05e006 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -29,7 +29,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 environments: assets activate-environment: true locked: true @@ -88,7 +88,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 working-directory: skala environments: ${{ matrix.environment }} activate-environment: true diff --git a/.github/workflows/gauxc-test.yml b/.github/workflows/gauxc-test.yml index cde33362..3fc1d8e3 100644 --- a/.github/workflows/gauxc-test.yml +++ b/.github/workflows/gauxc-test.yml @@ -31,7 +31,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 environments: default activate-environment: true locked: true diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index b97f764b..f787a558 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -19,7 +19,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 environments: lint activate-environment: true locked: true diff --git a/.github/workflows/model-benchmark.yml b/.github/workflows/model-benchmark.yml index ccc881ed..b34f5c03 100644 --- a/.github/workflows/model-benchmark.yml +++ b/.github/workflows/model-benchmark.yml @@ -37,7 +37,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 pixi-bin-path: ${{ runner.temp }}/bin/pixi environments: default activate-environment: true @@ -65,7 +65,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 pixi-bin-path: ${{ runner.temp }}/bin/pixi environments: gpu-cuda12-torch213 activate-environment: true diff --git a/.github/workflows/model-examples.yml b/.github/workflows/model-examples.yml index e0e513a1..5042bd89 100644 --- a/.github/workflows/model-examples.yml +++ b/.github/workflows/model-examples.yml @@ -31,7 +31,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 environments: assets activate-environment: true locked: true @@ -58,7 +58,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 environments: default activate-environment: true locked: true @@ -90,7 +90,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 environments: cpp-integration activate-environment: true locked: true @@ -133,7 +133,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 environments: ftorch activate-environment: true locked: true diff --git a/.github/workflows/model-test.yml b/.github/workflows/model-test.yml index a712118c..f25117a7 100644 --- a/.github/workflows/model-test.yml +++ b/.github/workflows/model-test.yml @@ -32,7 +32,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 environments: default activate-environment: true locked: true diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml index 8c4e2161..7ff4b7c2 100644 --- a/.github/workflows/pypi.yml +++ b/.github/workflows/pypi.yml @@ -25,7 +25,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 environments: release activate-environment: true locked: true @@ -54,7 +54,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 environments: release activate-environment: true locked: true diff --git a/.github/workflows/skala-test.yml b/.github/workflows/skala-test.yml index c2ff1e72..14f7aac7 100644 --- a/.github/workflows/skala-test.yml +++ b/.github/workflows/skala-test.yml @@ -50,7 +50,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 environments: ${{ matrix.environment }} activate-environment: true locked: true @@ -90,7 +90,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 pixi-bin-path: ${{ runner.temp }}/bin/pixi environments: ${{ matrix.environment }} activate-environment: true @@ -124,7 +124,7 @@ jobs: - name: Setup Pixi uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 with: - pixi-version: v0.75.0 + pixi-version: v0.78.0 environments: test-py312-pyscf214-torch213 activate-environment: true locked: true From 50b96e7eca48e29862e639e7349ef6e69802bc29 Mon Sep 17 00:00:00 2001 From: jenswehner Date: Tue, 1 Sep 2026 14:13:32 +0200 Subject: [PATCH 3/6] improve some readability --- skala/src/skala/pyscf/gradient_core.py | 108 ++++++++++++++++--------- 1 file changed, 72 insertions(+), 36 deletions(-) diff --git a/skala/src/skala/pyscf/gradient_core.py b/skala/src/skala/pyscf/gradient_core.py index 2d5797c0..f6eefc23 100644 --- a/skala/src/skala/pyscf/gradient_core.py +++ b/skala/src/skala/pyscf/gradient_core.py @@ -2,30 +2,69 @@ """Backend-independent PyTorch operations for PySCF nuclear gradients.""" -from collections.abc import Callable +from collections.abc import Callable, Iterator, Sequence +from enum import IntEnum import torch from skala.features import Feature, FeatureMap -# PySCF packs AO derivatives as value, x, y, z, xx, xy, xz, yy, yz, zz. -# Each row below selects the Hessian components for one nuclear displacement -# direction, with columns ordered by the x, y, and z density responses. -_HESSIAN_COMPONENTS = ((4, 5, 6), (5, 7, 8), (6, 8, 9)) + +class _AOComponent(IntEnum): + """Indices in PySCF's packed AO derivative dimension.""" + + VALUE = 0 + X = 1 + Y = 2 + Z = 3 + XX = 4 + XY = 5 + XZ = 6 + YY = 7 + YZ = 8 + ZZ = 9 + + +class _Direction(IntEnum): + """Cartesian direction indices used by feature and potential tensors.""" + + X = 0 + Y = 1 + Z = 2 + + +_AO_GRADIENT = slice(_AOComponent.X, _AOComponent.XX) +_GRADIENT_COMPONENTS = (_AOComponent.X, _AOComponent.Y, _AOComponent.Z) +_HESSIAN_COMPONENTS = ( + (_AOComponent.XX, _AOComponent.XY, _AOComponent.XZ), + (_AOComponent.XY, _AOComponent.YY, _AOComponent.YZ), + (_AOComponent.XZ, _AOComponent.YZ, _AOComponent.ZZ), +) + + +def _hessian_components() -> Iterator[tuple[_Direction, _Direction, _AOComponent]]: + """Yield force direction, response direction, and packed AO component.""" + for force_direction in _Direction: + for response_direction in _Direction: + yield ( + force_direction, + response_direction, + _HESSIAN_COMPONENTS[force_direction][response_direction], + ) + # Grid metadata is laid out with one row per point: coordinates have shape # (npoints, 3), while weights have shape (npoints,). -_POINT_FIRST_FEATURES = {Feature.GRID_COORDS, Feature.GRID_WEIGHTS} +_POINT_FIRST_FEATURES = (Feature.GRID_COORDS, Feature.GRID_WEIGHTS) # Electronic features keep spin and, for GRAD, Cartesian components before the # grid dimension: DENSITY and KIN have shape (2, npoints), and GRAD has shape # (2, 3, npoints). -_POINT_LAST_FEATURES = {Feature.DENSITY, Feature.GRAD, Feature.KIN} +_POINT_LAST_FEATURES = (Feature.DENSITY, Feature.GRAD, Feature.KIN) def feature_derivatives( - exc_func: Callable[..., torch.Tensor], - feature_tensors: list[torch.Tensor], + exc_func: Callable[..., torch.Tensor], feature_tensors: Sequence[torch.Tensor] ) -> tuple[torch.Tensor, ...]: """Differentiate a scalar XC energy with respect to molecular features. @@ -53,7 +92,7 @@ def feature_derivatives( gradients = torch.autograd.grad( exc, - tuple(differentiable_features), + differentiable_features, create_graph=False, retain_graph=False, allow_unused=True, @@ -96,10 +135,7 @@ def grid_derivative_block( return block -def contract_veff_block( - ao: torch.Tensor, - derivatives: FeatureMap, -) -> torch.Tensor: +def contract_veff_block(ao: torch.Tensor, derivatives: FeatureMap) -> torch.Tensor: """Contract one atom's AO values with XC feature derivatives. Args: @@ -123,35 +159,35 @@ def contract_veff_block( veff += torch.einsum( "si, xip, iq -> sxpq", derivatives[Feature.DENSITY], - ao[1:4], - ao[0], + ao[_AO_GRADIENT], + ao[_AOComponent.VALUE], ) if Feature.GRAD in derivatives: exc_dgrad = derivatives[Feature.GRAD] - veff += torch.einsum("syi, xip, yiq -> sxpq", exc_dgrad, ao[1:4], ao[1:4]) - for force_direction, components in enumerate(_HESSIAN_COMPONENTS): - for response_direction, component in enumerate(components): - veff[:, force_direction] += torch.einsum( - "si, ip, iq -> spq", - exc_dgrad[:, response_direction], - ao[component], - ao[0], - ) + ao_gradient = ao[_AO_GRADIENT] + veff += torch.einsum( + "syi, xip, yiq -> sxpq", exc_dgrad, ao_gradient, ao_gradient + ) + for force_direction, response_direction, component in _hessian_components(): + veff[:, force_direction] += torch.einsum( + "si, ip, iq -> spq", + exc_dgrad[:, response_direction], + ao[component], + ao[_AOComponent.VALUE], + ) if Feature.KIN in derivatives: exc_dkin = derivatives[Feature.KIN] - for force_direction, components in enumerate(_HESSIAN_COMPONENTS): - for response_direction, component in enumerate(components): - veff[:, force_direction] += ( - torch.einsum( - "si, ip, iq -> spq", - exc_dkin, - ao[component], - # AO components 1, 2, and 3 are the x, y, and z derivatives. - ao[response_direction + 1], - ) - / 2 + for force_direction, response_direction, component in _hessian_components(): + veff[:, force_direction] += ( + torch.einsum( + "si, ip, iq -> spq", + exc_dkin, + ao[component], + ao[_GRADIENT_COMPONENTS[response_direction]], ) + / 2 + ) return veff From 96004c92fa239a3a65d026b5324a665a7725b03a Mon Sep 17 00:00:00 2001 From: Jens Date: Thu, 3 Sep 2026 15:02:43 +0200 Subject: [PATCH 4/6] Update skala/src/skala/pyscf/gradient_core.py Co-authored-by: Thijs Vogels <840825+tvogels@users.noreply.github.com> --- skala/src/skala/pyscf/gradient_core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/skala/src/skala/pyscf/gradient_core.py b/skala/src/skala/pyscf/gradient_core.py index f6eefc23..2bf0969c 100644 --- a/skala/src/skala/pyscf/gradient_core.py +++ b/skala/src/skala/pyscf/gradient_core.py @@ -106,7 +106,7 @@ def feature_derivatives( def grid_derivative_block( derivatives: FeatureMap, grid_start: int, grid_end: int ) -> FeatureMap: - """Slice grid-resolved feature derivatives to one atom's grid block. + """Select a slice of the feature derivatives from `grid_start` to `grid_end`. Density-like features store the grid-point dimension last, whereas grid coordinates and weights store it first. Non-grid features, such as atomic From 936c11729c61f948a10684c885786edddfa46766 Mon Sep 17 00:00:00 2001 From: jenswehner Date: Thu, 3 Sep 2026 15:17:09 +0200 Subject: [PATCH 5/6] fix actions --- .github/actions/setup-pixi/action.yml | 34 +++++++++++++++++++++++++++ .github/workflows/benchmark-test.yml | 9 +++---- .github/workflows/docs.yml | 9 +++---- .github/workflows/examples.yml | 17 ++++---------- .github/workflows/gauxc-test.yml | 9 +++---- .github/workflows/lint.yml | 8 ++----- .github/workflows/model-benchmark.yml | 17 ++++---------- .github/workflows/model-examples.yml | 33 +++++++------------------- .github/workflows/model-test.yml | 9 +++---- .github/workflows/pypi.yml | 16 ++++--------- .github/workflows/skala-test.yml | 25 ++++++-------------- CONTRIBUTING.md | 2 +- 12 files changed, 79 insertions(+), 109 deletions(-) create mode 100644 .github/actions/setup-pixi/action.yml diff --git a/.github/actions/setup-pixi/action.yml b/.github/actions/setup-pixi/action.yml new file mode 100644 index 00000000..35494b81 --- /dev/null +++ b/.github/actions/setup-pixi/action.yml @@ -0,0 +1,34 @@ +name: Setup Pixi +description: Install Pixi and activate a locked project environment. + +inputs: + environment: + description: Pixi environment to install and activate. + required: true + working-directory: + description: Directory containing the Pixi project. + required: false + default: '' + pixi-bin-path: + description: Optional path at which to install the Pixi binary. + required: false + default: '' + post-cleanup: + description: Whether to remove Pixi and its environment after the job. + required: false + default: 'false' + +runs: + using: composite + steps: + - name: Setup Pixi + uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + with: + pixi-version: v0.78.0 + working-directory: ${{ inputs.working-directory }} + pixi-bin-path: ${{ inputs.pixi-bin-path }} + environments: ${{ inputs.environment }} + activate-environment: true + locked: true + cache: false + post-cleanup: ${{ inputs.post-cleanup }} \ No newline at end of file diff --git a/.github/workflows/benchmark-test.yml b/.github/workflows/benchmark-test.yml index 4460bef8..0dcf3205 100644 --- a/.github/workflows/benchmark-test.yml +++ b/.github/workflows/benchmark-test.yml @@ -5,6 +5,7 @@ on: branches: [main] paths: &benchmark-test-paths - '.github/workflows/benchmark-test.yml' + - '.github/actions/setup-pixi/**' - 'benchmark/**' - 'website/_ext/benchmark_report.py' - 'pixi.lock' @@ -33,13 +34,9 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 - environments: default - activate-environment: true - locked: true - cache: false + environment: default - name: Run benchmark unit tests run: pytest -v benchmark/tests/ diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 95bf9bb4..36d7b8c5 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -5,6 +5,7 @@ on: branches: [main] paths: &docs-paths - '.github/workflows/docs.yml' + - '.github/actions/setup-pixi/**' - 'benchmark/pyproject.toml' - 'pixi.lock' - 'pixi.toml' @@ -38,13 +39,9 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 - environments: docs - activate-environment: true - locked: true - cache: false + environment: docs - name: Build documentation run: | diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 9b05e006..cf192fb8 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -5,6 +5,7 @@ on: branches: [main] paths: &gauxc-paths - '.github/workflows/examples.yml' + - '.github/actions/setup-pixi/**' - 'gauxc/**' - 'pixi.lock' - 'pixi.toml' @@ -27,13 +28,9 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 - environments: assets - activate-environment: true - locked: true - cache: false + environment: assets - name: Download checkpoint run: >- @@ -86,14 +83,10 @@ jobs: uses: ./skala/.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./skala/.github/actions/setup-pixi with: - pixi-version: v0.78.0 working-directory: skala - environments: ${{ matrix.environment }} - activate-environment: true - locked: true - cache: false + environment: ${{ matrix.environment }} - name: Checkout development GauXC if: ${{ matrix.source == 'development' }} diff --git a/.github/workflows/gauxc-test.yml b/.github/workflows/gauxc-test.yml index 3fc1d8e3..9dc85ec8 100644 --- a/.github/workflows/gauxc-test.yml +++ b/.github/workflows/gauxc-test.yml @@ -5,6 +5,7 @@ on: branches: [main] paths: &gauxc-test-paths - '.github/workflows/gauxc-test.yml' + - '.github/actions/setup-pixi/**' - 'gauxc/**' - 'pixi.lock' - 'pixi.toml' @@ -29,13 +30,9 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 - environments: default - activate-environment: true - locked: true - cache: false + environment: default - name: Run GauXC unit tests run: pytest -v gauxc/tests/ \ No newline at end of file diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index f787a558..80b2d370 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -17,13 +17,9 @@ jobs: - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4 - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 - environments: lint - activate-environment: true - locked: true - cache: false + environment: lint - name: Run pre-commit hooks run: pre-commit run --all-files diff --git a/.github/workflows/model-benchmark.yml b/.github/workflows/model-benchmark.yml index b34f5c03..68414b54 100644 --- a/.github/workflows/model-benchmark.yml +++ b/.github/workflows/model-benchmark.yml @@ -5,6 +5,7 @@ on: branches: [main] paths: &model-benchmark-paths - '.github/workflows/model-benchmark.yml' + - '.github/actions/setup-pixi/**' - 'conftest.py' - 'pixi.lock' - 'pixi.toml' @@ -35,14 +36,10 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 + environment: default pixi-bin-path: ${{ runner.temp }}/bin/pixi - environments: default - activate-environment: true - locked: true - cache: false post-cleanup: true - name: Run CPU model benchmarks @@ -63,14 +60,10 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 + environment: gpu-cuda12-torch213 pixi-bin-path: ${{ runner.temp }}/bin/pixi - environments: gpu-cuda12-torch213 - activate-environment: true - locked: true - cache: false post-cleanup: true - name: Verify CUDA is available diff --git a/.github/workflows/model-examples.yml b/.github/workflows/model-examples.yml index 5042bd89..5b87f863 100644 --- a/.github/workflows/model-examples.yml +++ b/.github/workflows/model-examples.yml @@ -5,6 +5,7 @@ on: branches: [main] paths: &model-example-paths - '.github/workflows/model-examples.yml' + - '.github/actions/setup-pixi/**' - 'model/examples/**' - 'model/src/**' - 'pixi.lock' @@ -29,13 +30,9 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 - environments: assets - activate-environment: true - locked: true - cache: false + environment: assets - name: Download checkpoint run: >- @@ -56,13 +53,9 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 - environments: default - activate-environment: true - locked: true - cache: false + environment: default - name: Generate features run: >- @@ -88,13 +81,9 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 - environments: cpp-integration - activate-environment: true - locked: true - cache: false + environment: cpp-integration - name: Configure and build project run: | @@ -131,13 +120,9 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 - environments: ftorch - activate-environment: true - locked: true - cache: false + environment: ftorch - name: Configure, build, and install project run: | diff --git a/.github/workflows/model-test.yml b/.github/workflows/model-test.yml index f25117a7..b67fa70e 100644 --- a/.github/workflows/model-test.yml +++ b/.github/workflows/model-test.yml @@ -5,6 +5,7 @@ on: branches: [main] paths: &model-test-paths - '.github/workflows/model-test.yml' + - '.github/actions/setup-pixi/**' - 'model/**' - 'pixi.lock' - 'pixi.toml' @@ -30,13 +31,9 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 - environments: default - activate-environment: true - locked: true - cache: false + environment: default - name: Run model unit tests run: pytest -v model/tests/test_model.py model/tests/test_utils.py \ No newline at end of file diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml index 7ff4b7c2..05ba35a4 100644 --- a/.github/workflows/pypi.yml +++ b/.github/workflows/pypi.yml @@ -23,13 +23,9 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 - environments: release - activate-environment: true - locked: true - cache: false + environment: release - name: Build a binary wheel and a source tarball run: python tools/build_release.py ${{ matrix.package }} @@ -52,13 +48,9 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 - environments: release - activate-environment: true - locked: true - cache: false + environment: release - name: Build conda artifact run: pixi publish --path skala --target-dir dist-conda diff --git a/.github/workflows/skala-test.yml b/.github/workflows/skala-test.yml index 14f7aac7..5fe50e35 100644 --- a/.github/workflows/skala-test.yml +++ b/.github/workflows/skala-test.yml @@ -5,6 +5,7 @@ on: branches: [main] paths: &skala-test-paths - '.github/workflows/skala-test.yml' + - '.github/actions/setup-pixi/**' - 'pixi.lock' - 'pixi.toml' - 'pyproject.toml' @@ -48,13 +49,9 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 - environments: ${{ matrix.environment }} - activate-environment: true - locked: true - cache: false + environment: ${{ matrix.environment }} - name: Run tests with coverage run: >- @@ -88,14 +85,10 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 + environment: ${{ matrix.environment }} pixi-bin-path: ${{ runner.temp }}/bin/pixi - environments: ${{ matrix.environment }} - activate-environment: true - locked: true - cache: false post-cleanup: true - name: Verify CUDA is available @@ -122,13 +115,9 @@ jobs: uses: ./.github/actions/cpu-diagnostics - name: Setup Pixi - uses: prefix-dev/setup-pixi@f00437f565399d418b0acc85936d12c1fb668347 # v0.10.1 + uses: ./.github/actions/setup-pixi with: - pixi-version: v0.78.0 - environments: test-py312-pyscf214-torch213 - activate-environment: true - locked: true - cache: false + environment: test-py312-pyscf214-torch213 - name: Run profiling tests run: pytest -v -m profiling skala/tests/ \ No newline at end of file diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 73664794..9eece202 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -16,7 +16,7 @@ or contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any addi ## Development setup -Install Pixi 0.75, then create the default locked development environment from the repository root: +Install Pixi 0.78, then create the default locked development environment from the repository root: ```bash pixi install --locked -e default From 2e055a8b35d7f0c189f1cedf874a28656a7845ca Mon Sep 17 00:00:00 2001 From: jenswehner Date: Thu, 3 Sep 2026 16:39:13 +0200 Subject: [PATCH 6/6] refactor gradient_core --- skala/src/skala/gpu4pyscf/gradients.py | 17 +- skala/src/skala/pyscf/gradient_core.py | 227 ++++++++++++++----------- skala/src/skala/pyscf/gradients.py | 17 +- skala/tests/test_gradient_core.py | 43 +++-- 4 files changed, 164 insertions(+), 140 deletions(-) diff --git a/skala/src/skala/gpu4pyscf/gradients.py b/skala/src/skala/gpu4pyscf/gradients.py index 616ee5f1..17830bb3 100644 --- a/skala/src/skala/gpu4pyscf/gradients.py +++ b/skala/src/skala/gpu4pyscf/gradients.py @@ -21,7 +21,7 @@ from skala.features import Feature, FeatureMap from skala.functional.base import ExcFunctionalBase from skala.pyscf.gradient_core import ( - contract_veff_block, + contract_ao_derivative_block, feature_derivatives, grid_derivative_block, ) @@ -96,20 +96,15 @@ def veff_and_expl_nuc_grad( nuc_grad_feats.discard(Feature.ATOMIC_GRID_WEIGHTS) # Get required derivatives - nuc_feat_names = list(nuc_grad_feats) # ensure specific order - nuc_feat_tensors = [mol_feats[feat] for feat in nuc_feat_names] + 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(*nuc_feat_tensors: torch.Tensor) -> torch.Tensor: - exc_mol_feats = ( - dict(zip(nuc_feat_names, nuc_feat_tensors, strict=True)) | other_feats - ) - return functional.get_exc(exc_mol_feats) + def exc_feat_func(differentiable_features: FeatureMap) -> torch.Tensor: + return functional.get_exc(differentiable_features | other_feats) - dExc_tuple = feature_derivatives(exc_feat_func, nuc_feat_tensors) - dExc: FeatureMap = dict(zip(nuc_feat_names, dExc_tuple, strict=True)) + dExc = feature_derivatives(exc_feat_func, nuc_feats) LOG.debug("autograd gradients for nuclear features done") @@ -134,7 +129,7 @@ def exc_feat_func(*nuc_feat_tensors: torch.Tensor) -> torch.Tensor: dExc_atm = grid_derivative_block(dExc, atm_start, atm_end) # Calculate the contribution to veff for this atomic grid - veff_atm = contract_veff_block(ao, dExc_atm) + veff_atm = contract_ao_derivative_block(ao, dExc_atm) if Feature.GRID_COORDS in dExc_atm: # also add the explicit grid coordinate dependence diff --git a/skala/src/skala/pyscf/gradient_core.py b/skala/src/skala/pyscf/gradient_core.py index 2bf0969c..48ea2a22 100644 --- a/skala/src/skala/pyscf/gradient_core.py +++ b/skala/src/skala/pyscf/gradient_core.py @@ -2,29 +2,17 @@ """Backend-independent PyTorch operations for PySCF nuclear gradients.""" -from collections.abc import Callable, Iterator, Sequence +from collections.abc import Callable, Iterable, Iterator +from dataclasses import dataclass from enum import IntEnum +from types import EllipsisType +from typing import ClassVar, TypeAlias import torch from skala.features import Feature, FeatureMap -class _AOComponent(IntEnum): - """Indices in PySCF's packed AO derivative dimension.""" - - VALUE = 0 - X = 1 - Y = 2 - Z = 3 - XX = 4 - XY = 5 - XZ = 6 - YY = 7 - YZ = 8 - ZZ = 9 - - class _Direction(IntEnum): """Cartesian direction indices used by feature and potential tensors.""" @@ -33,39 +21,52 @@ class _Direction(IntEnum): Z = 2 -_AO_GRADIENT = slice(_AOComponent.X, _AOComponent.XX) -_GRADIENT_COMPONENTS = (_AOComponent.X, _AOComponent.Y, _AOComponent.Z) -_HESSIAN_COMPONENTS = ( - (_AOComponent.XX, _AOComponent.XY, _AOComponent.XZ), - (_AOComponent.XY, _AOComponent.YY, _AOComponent.YZ), - (_AOComponent.XZ, _AOComponent.YZ, _AOComponent.ZZ), -) +@dataclass(frozen=True) +class _PackedAO: + """Views into PySCF's packed AO derivative dimension.""" + tensor: torch.Tensor -def _hessian_components() -> Iterator[tuple[_Direction, _Direction, _AOComponent]]: - """Yield force direction, response direction, and packed AO component.""" - for force_direction in _Direction: - for response_direction in _Direction: - yield ( - force_direction, - response_direction, - _HESSIAN_COMPONENTS[force_direction][response_direction], - ) + # 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] -# Grid metadata is laid out with one row per point: coordinates have shape -# (npoints, 3), while weights have shape (npoints,). -_POINT_FIRST_FEATURES = (Feature.GRID_COORDS, Feature.GRID_WEIGHTS) + @property + def gradient(self) -> torch.Tensor: + """AO gradients with shape ``(3, npoints, nao)``.""" + return self.tensor[1:4] -# Electronic features keep spin and, for GRAD, Cartesian components before the -# grid dimension: DENSITY and KIN have shape (2, npoints), and GRAD has shape -# (2, 3, npoints). -_POINT_LAST_FEATURES = (Feature.DENSITY, Feature.GRAD, Feature.KIN) + 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] + + +_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[..., torch.Tensor], feature_tensors: Sequence[torch.Tensor] -) -> tuple[torch.Tensor, ...]: + 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 @@ -73,34 +74,44 @@ def feature_derivatives( storage. Args: - exc_func: Callable accepting the feature tensors and returning scalar XC energy. - feature_tensors: Feature tensors in the order expected by ``exc_func``. + exc_func: Callable accepting the differentiable features and returning + scalar XC energy. + features: Molecular features to differentiate, keyed by feature name. Returns: - Feature derivatives in the same order as ``feature_tensors``. + XC energy derivatives keyed by feature name. + + Raises: + RuntimeError: If the XC energy is disconnected from a requested feature. """ - if not feature_tensors: - return () + if not features: + return {} - differentiable_features = tuple( - tensor.detach().requires_grad_(True) for tensor in feature_tensors - ) - exc = exc_func(*differentiable_features) + differentiable_features = { + feature: tensor.detach().requires_grad_(True) + for feature, tensor in features.items() + } + exc = exc_func(differentiable_features) if not exc.requires_grad: - return tuple(torch.zeros_like(feature) for feature in differentiable_features) + raise _disconnected_features_error(differentiable_features) gradients = torch.autograd.grad( exc, - differentiable_features, + tuple(differentiable_features.values()), create_graph=False, retain_graph=False, allow_unused=True, ) - return tuple( - torch.zeros_like(feature) if gradient is None else gradient.detach() + 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( @@ -120,74 +131,88 @@ def grid_derivative_block( Returns: Grid-resolved derivatives restricted to ``[grid_start:grid_end]``. """ - block = { - feature: derivative[grid_start:grid_end] + grid_slice = slice(grid_start, grid_end) + feature_slices: dict[Feature, _FeatureSlice] = { + Feature.GRID_COORDS: (grid_slice,), + Feature.GRID_WEIGHTS: (grid_slice,), + Feature.DENSITY: (..., grid_slice), + Feature.GRAD: (..., grid_slice), + Feature.KIN: (..., grid_slice), + } + return { + feature: derivative[feature_slice] for feature, derivative in derivatives.items() - if feature in _POINT_FIRST_FEATURES + if (feature_slice := feature_slices.get(feature)) is not None } - block.update( - { - feature: derivative[..., grid_start:grid_end] - for feature, derivative in derivatives.items() - if feature in _POINT_LAST_FEATURES - } - ) - return block -def contract_veff_block(ao: torch.Tensor, derivatives: FeatureMap) -> torch.Tensor: - """Contract one atom's AO values with XC feature derivatives. +def contract_ao_derivative_block( + ao: torch.Tensor, feature_derivatives: FeatureMap +) -> torch.Tensor: + """Contract XC feature derivatives with one-sided spatial AO derivatives. + + For each Cartesian direction, form one grid block's contribution to the + spatial derivative of the AO-basis XC potential. The AO associated with + the first matrix index is differentiated, while the AO associated with the + second index is held fixed. Consequently, the returned matrices are not + generally symmetric in their AO indices. + + Density, density-gradient, and kinetic-energy-density contributions are + included when their corresponding feature derivatives are present. Args: - ao: Component-major AO values for this atom's grid block, shaped - ``(ncomponents, npoints, nao)``. Here ``npoints`` must equal - ``grid_end - grid_start`` and ``nao`` is the number of molecular - AOs. PySCF supplies 1, 4, or 10 components for derivative orders - 0, 1, or 2, respectively, packed as ``value, x, y, z, xx, xy, - xz, yy, yz, zz``. The requested features determine how many of - these leading components the contraction reads. - derivatives: XC energy derivatives already sliced to this grid block. + ao: Packed component-major AO values with shape + ``(ncomponents, npoints, nao)``. Components follow PySCF's ordering: + ``value, x, y, z, xx, xy, xz, yy, yz, zz``. + feature_derivatives: XC energy derivatives with respect to features on + this grid block. Density and kinetic derivatives have shape + ``(2, npoints)``; gradient derivatives have shape + ``(2, 3, npoints)``. Returns: - Spin- and Cartesian-resolved effective-potential contribution with shape - ``(2, 3, nao, nao)`` on the same device and with the same dtype as ``ao``. + One-sided AO-potential derivatives with shape ``(2, 3, nao, nao)``, + indexed by spin, Cartesian direction, and the two AO indices. The result + has the same device and dtype as ``ao``. """ + packed_ao = _PackedAO(ao) nao = ao.shape[-1] - veff = ao.new_zeros((2, 3, nao, nao)) + potential_derivatives = ao.new_zeros((2, 3, nao, nao)) - if Feature.DENSITY in derivatives: - veff += torch.einsum( + if Feature.DENSITY in feature_derivatives: + potential_derivatives += torch.einsum( "si, xip, iq -> sxpq", - derivatives[Feature.DENSITY], - ao[_AO_GRADIENT], - ao[_AOComponent.VALUE], + feature_derivatives[Feature.DENSITY], + packed_ao.gradient, + packed_ao.value, ) - if Feature.GRAD in derivatives: - exc_dgrad = derivatives[Feature.GRAD] - ao_gradient = ao[_AO_GRADIENT] - veff += torch.einsum( - "syi, xip, yiq -> sxpq", exc_dgrad, ao_gradient, ao_gradient + 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 force_direction, response_direction, component in _hessian_components(): - veff[:, force_direction] += torch.einsum( + for ao_direction, feature_direction, ao_hessian in packed_ao.hessian(): + potential_derivatives[:, ao_direction] += torch.einsum( "si, ip, iq -> spq", - exc_dgrad[:, response_direction], - ao[component], - ao[_AOComponent.VALUE], + exc_dgrad[:, feature_direction], + ao_hessian, + packed_ao.value, ) - if Feature.KIN in derivatives: - exc_dkin = derivatives[Feature.KIN] - for force_direction, response_direction, component in _hessian_components(): - veff[:, force_direction] += ( + 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( "si, ip, iq -> spq", exc_dkin, - ao[component], - ao[_GRADIENT_COMPONENTS[response_direction]], + ao_hessian, + packed_ao.gradient[feature_direction], ) / 2 ) - return veff + return potential_derivatives diff --git a/skala/src/skala/pyscf/gradients.py b/skala/src/skala/pyscf/gradients.py index 029f824d..ed2d405f 100644 --- a/skala/src/skala/pyscf/gradients.py +++ b/skala/src/skala/pyscf/gradients.py @@ -18,7 +18,7 @@ from skala.features import Feature, FeatureMap from skala.functional.base import ExcFunctionalBase from skala.pyscf.gradient_core import ( - contract_veff_block, + contract_ao_derivative_block, feature_derivatives, grid_derivative_block, ) @@ -91,20 +91,15 @@ def veff_and_expl_nuc_grad( nuc_grad_feats.discard(Feature.ATOMIC_GRID_WEIGHTS) # Get required derivatives - nuc_feat_names = list(nuc_grad_feats) # ensure specific order - nuc_feat_tensors = [mol_feats[feat] for feat in nuc_feat_names] + 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(*nuc_feat_tensors: torch.Tensor) -> torch.Tensor: - exc_mol_feats = ( - dict(zip(nuc_feat_names, nuc_feat_tensors, strict=True)) | other_feats - ) - return functional.get_exc(exc_mol_feats) + def exc_feat_func(differentiable_features: FeatureMap) -> torch.Tensor: + return functional.get_exc(differentiable_features | other_feats) - dExc_tuple = feature_derivatives(exc_feat_func, nuc_feat_tensors) - dExc: FeatureMap = dict(zip(nuc_feat_names, dExc_tuple, strict=True)) + dExc = feature_derivatives(exc_feat_func, nuc_feats) LOG.debug("autograd gradients for nuclear features done") @@ -126,7 +121,7 @@ def exc_feat_func(*nuc_feat_tensors: torch.Tensor) -> torch.Tensor: dExc_atm = grid_derivative_block(dExc, atm_start, atm_end) # Calculate the contribution to veff for this atomic grid - veff_atm = contract_veff_block(ao, dExc_atm) + veff_atm = contract_ao_derivative_block(ao, dExc_atm) if Feature.GRID_COORDS in dExc_atm: # also add the explicit grid coordinate dependence diff --git a/skala/tests/test_gradient_core.py b/skala/tests/test_gradient_core.py index 870b3744..214d217b 100644 --- a/skala/tests/test_gradient_core.py +++ b/skala/tests/test_gradient_core.py @@ -2,10 +2,11 @@ """Tests for backend-independent nuclear-gradient operations.""" +import pytest import torch from skala.features import Feature from skala.pyscf.gradient_core import ( - contract_veff_block, + contract_ao_derivative_block, feature_derivatives, grid_derivative_block, ) @@ -16,31 +17,39 @@ def test_feature_derivatives() -> None: weights = torch.tensor([3.0, 4.0], dtype=torch.float64) derivatives = feature_derivatives( - lambda rho, grid_weights: (rho.square() * grid_weights).sum(), - [density, weights], + lambda features: ( + features[Feature.DENSITY].square() * features[Feature.GRID_WEIGHTS] + ).sum(), + {Feature.DENSITY: density, Feature.GRID_WEIGHTS: weights}, ) - torch.testing.assert_close(derivatives[0], 2 * density * weights) - torch.testing.assert_close(derivatives[1], density.square()) - assert feature_derivatives(lambda: torch.tensor(0.0), []) == () + 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_returns_zero_for_disconnected_feature() -> None: +def test_feature_derivatives_rejects_disconnected_feature() -> None: used = torch.tensor(2.0) unused = torch.tensor(3.0) - derivatives = feature_derivatives(lambda value, _: value.square(), [used, unused]) + 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}, + ) - torch.testing.assert_close(derivatives[0], 2 * used) - torch.testing.assert_close(derivatives[1], torch.zeros_like(unused)) - -def test_feature_derivatives_returns_zero_for_constant_energy() -> None: +def test_feature_derivatives_rejects_constant_energy() -> None: feature = torch.tensor(2.0) - (derivative,) = feature_derivatives(lambda _: torch.tensor(1.0), [feature]) - - torch.testing.assert_close(derivative, torch.zeros_like(feature)) + 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_grid_derivative_block_slices_each_grid_dimension() -> None: @@ -72,7 +81,7 @@ def test_grid_derivative_block_slices_each_grid_dimension() -> None: ) -def test_contract_veff_block_matches_reference() -> None: +def test_contract_ao_derivative_block_matches_reference() -> None: ao = torch.tensor( [2.0, 3.0, 5.0, 7.0, 11.0, 13.0, 17.0, 19.0, 23.0, 29.0], dtype=torch.float64, @@ -88,7 +97,7 @@ def test_contract_veff_block_matches_reference() -> None: ), Feature.KIN: torch.tensor([[0.0, 67.0], [0.0, 71.0]], dtype=torch.float64), } - actual = contract_veff_block(ao, grid_derivative_block(derivatives, 1, 2)) + actual = contract_ao_derivative_block(ao, grid_derivative_block(derivatives, 1, 2)) expected = torch.tensor( [ [[[13074.5]], [[18389.5]], [[23562.5]]],