PyTorch-based Libcint Interface

Compatibility:
Availability:
Status:

Content

This package contains a PyTorch-based interface to the libcint high-performance integral library (check version here). The interface facilitates automatic differentiation with custom backward functions that call libcint's C backend for derivative calculation.

It is based on PySCF's lib directory. Modifications include

• only the driver for two-centre integral matrices (GTOint2c) and its helpers are kept
• a new CMake build that compiles the libcint submodule directly and suffixes all libraries with _tad

The project was heavily inspired by diffqc/dqc and diffqc/dqclibs.

This interface was mainly written for the calculation of integrals within grimme-lab/dxtb.

Usage

A Basis holds the shells of every element: their angular momenta, exponents and contraction coefficients. The atoms are passed to every integral call as a tad-mctc Structure (atomic numbers and positions in Bohr), so one basis serves all conformers of a molecule as well as different molecules. Contraction coefficients are used as given, i.e., they must include the normalization in libcint's convention (the radial part times a real spherical harmonic with unit norm). Coefficients that normalize the Cartesian function $x^l e^{-\alpha r^2}$ differ by a factor of $\sqrt{4\pi/(2l+1)}$; basis.normalized() scales every shell to unit self-overlap instead.

import torch
from tad_mctc.io.structure import Structure
from tad_libcint import Basis, Shell, int1e, overlap

dd = {"dtype": torch.double}


def shell(angmom, alphas, coeffs):
    return Shell(
        angmom=angmom,
        alphas=torch.tensor(alphas, **dd),
        coeffs=torch.tensor(coeffs, **dd),
    )


# STO-3G shells of each element: angular momentum, exponents, coefficients
# (coefficients include the normalization in libcint's convention)
shells = {
    1: [
        shell(0, [3.42525, 0.623914, 0.168855], [0.981707, 0.949464, 0.295906]),
    ],
    8: [
        shell(0, [130.709, 23.8089, 6.44361], [15.0727, 14.5777, 4.54323]),
        shell(0, [5.03315, 1.1696, 0.380389], [-0.848697, 1.1352, 0.856753]),
        shell(1, [5.03315, 1.1696, 0.380389], [3.42907, 2.15629, 0.341592]),
    ],
}

water = Structure(
    numbers=torch.tensor([8, 1, 1]),
    positions=torch.tensor([[0.0, 0.0, 0.22], [0.0, 1.43, -0.89], [0.0, -1.43, -0.89]], **dd),
)

basis = Basis(shells=shells)
s = overlap(basis, water)  # (nao, nao)
dipole = int1e("j", basis, water)  # (3, nao, nao), about each ket's atom

Gradients with respect to positions, exponents (basis.alphas) and coefficients (basis.coeffs) are available through autograd and are computed with libcint's derivative integrals. Shell and Basis are tad-mctc Nodes, like Structure: basis.to(device) converts them, and they pass through torch.func transforms (e.g. jacrev with respect to a whole basis). The exponents and coefficients are shared by all atoms of an element.

Periodic structures

For a structure with a lattice (vectors as rows, in Bohr), the integrals are lattice sums over the periodic images of the ket, $A_{\mu\nu}(\mathbf{k}) = \sum_{\mathbf{T}} e^{i\mathbf{k}\cdot\mathbf{T}} \langle \mu | \hat{O} | \nu(\mathbf{T}) \rangle$: real at the Γ point (the default), complex at k-points. Ket-centred moments ("j", "jj", "jjj") are taken about the ket's image, as in tblite; moments about the origin ("r0", ...) and the nuclear attraction ("nuc") have no periodic integral. The images run along the periodic axes of structure.periodic only, so 1D and 2D systems need no special treatment.

from tad_mctc.kpoints import make_kpoints

cell = water.replace(lattice=6.0 * torch.eye(3, **dd))  # periodic along all axes
s = overlap(basis, cell)  # (nao, nao), Gamma point
kpoints = make_kpoints(cell.lattice, (2, 2, 2))  # Gamma-centred grid
dipole_k = int1e("j", basis, cell, kpoints=kpoints)  # (8, 3, nao, nao), complex

k-points are Cartesian (1/Bohr), as in PySCF, either a tensor or a tad-mctc KPoints with weights and a mask; sets of different size in a batch are packed with tad_mctc.kpoints.pack_kpoints (padding k-points give zeros). The images are screened to the precision (default 1e-10), an estimate of the largest neglected contribution to an element.

Gradients with respect to the lattice are available like those with respect to the positions; the strain derivative follows as $\sigma = (\partial E/\partial \mathbf{A})^T \mathbf{A} + (\partial E/\partial \mathbf{R})^T \mathbf{R}$. They are taken at fixed fractional k-points, as in PySCF's stress code: the k-points are construction data and carry no gradient, however they were built.

Batches and torch.func

The integrals are pure functions of the structure, exponents and coefficients and compose with torch.vmap, torch.func.grad and torch.func.jacrev, to any derivative order (forward-mode AD is not supported). A batched Structure (one leading batch dimension) is evaluated structure by structure, and torch.vmap maps over a Structure directly:

• Conformers repeat the atomic numbers, e.g. a multi-frame XYZ file read with tad_mctc.io.read.
• Different structures are zero-padded to a common number of atoms (atomic number 0 marks padding), e.g. with tad_mctc.io.structure.pack_structures. The integrals of each structure are zero-padded to the largest number of AOs in the batch, with the structure's AOs first. Padding atoms get zero gradients.

The torch.func transforms differentiate with respect to tensors, not structures (the atomic numbers are integers). Replace the positions inside the differentiated function instead, structure.replace(positions=positions).

from torch.func import jacrev, vmap
from tad_mctc.io import read
from tad_mctc.io.structure import pack_structures

conformers = read.read("conformers.xyz", dtype=torch.double)  # positions: (nconf, nat, 3)
s = overlap(basis, conformers)  # (nconf, nao, nao)

molecules = [read.read(f, dtype=torch.double) for f in ("water.xyz", "methane.xyz")]
batch = pack_structures(molecules)
s = overlap(basis, batch)  # (2, nao_max, nao_max)

def dipole(structure, positions):
    return int1e("j", basis, structure.replace(positions=positions))

# Jacobians w.r.t. the positions, (2, 3, nao_max, nao_max, nat_max, 3)
jac = vmap(jacrev(dipole, argnums=1))(batch, batch.positions)

Since libcint is C code, a batch is evaluated one structure after another. Combining very different structures in one batch wastes memory on padding.

Threads

The integrals are evaluated in parallel over shell pairs with OpenMP. On Linux and macOS, the libraries of the wheels use PyTorch's OpenMP runtime and hence its number of threads, set with torch.set_num_threads or the OMP_NUM_THREADS environment variable (by default, one thread per physical core). On Linux, MKL_NUM_THREADS takes precedence over OMP_NUM_THREADS for PyTorch. On Windows, the libraries have an OpenMP runtime of their own, which only follows OMP_NUM_THREADS (by default, one thread per logical core).

Installation

pip

The tad-libcint package can be obtained from pip.

pip install tad-libcint

Wheels are available for Linux (x86_64), macOS (Apple silicon) and Windows (x86_64), the platforms of PyTorch's wheels. One wheel serves all supported Python versions.

Build from source

You can also build the Python wheels from source.

Get the repository (including libcint submodule).

git clone --recursive git@github.com:tad-mctc/tad-libcint.git
cd tad-libcint

The libcint fork contains additional integrals required for dxtb and removes unneccessary ones. Therefore, downloading from our fork is recommended.

If you already have the repository, you can update the submodule with

git submodule update --remote

The build requires CMake (≥ 3.18) and a C compiler (with OpenMP for parallel integral evaluation). On macOS, Apple's Clang (Xcode Command Line Tools) builds with OpenMP if Homebrew's libomp is installed (brew install libomp). The libraries then use PyTorch's OpenMP runtime, since a second one in the same process aborts. Editable installs find it only with PyTorch in the build environment (pip install --no-build-isolation -e .) and otherwise build without OpenMP. On Windows, libcint needs the GCC of MinGW-w64, since MSVC lacks C99 complex numbers, e.g., the MSYS2 packages mingw-w64-ucrt-x86_64-gcc, mingw-w64-ucrt-x86_64-libgomp (OpenMP) and mingw-w64-ucrt-x86_64-make on the PATH. Its runtime (including OpenMP) is linked statically, so the libraries only depend on system DLLs. Extra CMake options are passed through the CMAKE_ARGS environment variable, e.g., CMAKE_ARGS="-DBUILD_MARCH_NATIVE=ON" pip install ..

The compiler is taken from CC. Besides GCC, the Intel oneAPI compiler (icx) is supported, e.g., after module load intel. The libraries then find Intel's runtime libraries in the compiler installation, so the compiler environment is not needed at runtime.

All shared libraries carry the suffix _tad (libcint_tad.so, libcgto_tad.so, or .dll on Windows), so that they do not clash with PySCF's libraries of the same name when both packages are used in one process.

CIBuildWheel

Building with cibuildwheel is as simple as running

cibuildwheel --platform linux

This builds one wheel (py3-none-manylinux_2_17_x86_64) for all Python versions, as configured in pyproject.toml.

Conda

1. Create an environment (conda) for building the wheels. The example here uses Python 3.10, but we support 3.10-3.14.

mamba create -n wheel-310 --yes python=3.10 auditwheel c-compiler cxx-compiler fortran-compiler cffi cmake git make meson ninja numpy patchelf pkgconfig pip python-build unzip wheel
mamba activate wheel-310

2. Build the wheel on your system.

python -m build --wheel

3. Repair the wheels for cross-distribution packaging.

auditwheel repair -w wheels --plat manylinux_2_17_x86_64 dist/*-py3-none-linux_x86_64.whl

4. Or only install this project with pip in the environment.

pip install .

Compatibility

✅ tested in CI    ✔️ supported, but not tested in CI (should still work)    ❌ not supported

Note that only the latest bug fix version is listed, but all preceding bug fix minor versions are supported. For example, although only version 2.7.1 is listed, version 2.7.0 is also supported.

The table refers to Linux. On macOS and Windows, CI tests the newest PyTorch version with all Python versions.

Development

For development, additionally install the following tools in your environment.

mamba install black covdefaults coverage mypy pre-commit pylint tox

With pip, add the option -e for installing in development mode, and add [dev] for the development dependencies

pip install -e .[dev]

The pre-commit hooks are initialized by running the following command in the root of the repository.

pre-commit install

For testing all Python environments, simply run tox.

tox

Note that this randomizes the order of tests but skips "large" tests. To modify this behavior, tox has to skip the optional posargs.

tox -- test