GitHub - lowdanie/hartree-fock-solver: A quantum chemisty library in jax.

Slaterform

slaterform is a differentiable Hartree-Fock engine written in jax. It includes a native implementation of the necessary electron integrals and supports standard basis sets from basis set exchange.

Example: Butane Geometry Optimization

Because slaterform is written in pure jax, it can easily be used to define a differentiable molecular energy function. This function can be minimized in a standard jax optimization loop to optimize the molecular geometry.

import jax

import slaterform as sf
import slaterform.hartree_fock.scf as scf

def total_energy(molecule: sf.Molecule):
    """Compute the total energy of the molecule with Hartree-Fock"""
    options = scf.Options(
        max_iterations=20,
        execution_mode=scf.ExecutionMode.FIXED,
        integral_strategy=scf.IntegralStrategy.CACHED,
        perturbation=1e-10,
    )
    result = scf.solve(mol, options)

    return result.total_energy

# Add gradients and JIT compile.
total_energy_and_grad = jax.jit(jax.value_and_grad(total_energy))

In this colab notebook you can select a molecule, optimize the nuclear positions with optax, and finally visualize the trajectory of the nuclei and electron density using 3dmol. Here is a sample output for butane. We initialize the carbon chain to lie in a straight line, and the optimizer moves it into the classic zig-zag configuration. The blue cloud is rendered by sampling the electron density returned by scf.solve.

Example: Aspirin

Here is a rendering of the electron density generated by slaterform for aspirin. See the Quick Start section for more details.

Accuracy Benchmarks

Here is a benchmark of total energies against PySCF using the sto-3g basis set. The table was generated using this colab notebook

Molecule	Atoms	f32 Error (Ha)	f64 Error (Ha)
Water ($H_2O$)	3	`1.80e-6`	`2.32e-8`
Methane ($CH_4$)	5	`4.59e-8`	`2.40e-9`
Ammonia ($NH_3$)	4	`2.87e-7`	`1.66e-8`
Ethylene ($C_2H_4$)	6	`9.30e-8`	`9.00e-9`
Benzene ($C_6H_6$)	12	`2.37e-8`	`1.83e-8`
Butane ($C_4H_{10}$)	14	`1.33e-7`	`2.90e-9`
Aspirin ($C_9H_8O_4$)	21	`7.13e-6`	`6.04e-8`

Throughput Benchmarks

Here is a performance benchmark of the solver on NVIDIA T4 and A100 GPUs. The benchmarks were run on Benzene ($C_6H_6$) with the sto-3g basis set (N=36). The mixed precision strategy computes integrals in float32 but runs SCF in float64.

GPU	Mode	Batch Size	Throughput (mol/s)	Time per Mol (ms)
T4	Forward Only	16	4.25	235
T4	Forward + Grad	8	1.38	725
A100	Forward Only	32	23.22	43
A100	Forward + Grad	16	8.37	119

Quick Start

Here is an example which estimates the electronic ground state of water using the STO-3G basis set from basis set exchange.

import jax
import jax.numpy as jnp

import slaterform as sf
import slaterform.hartree_fock.scf as scf

# Build the H2O molecule with nuclear positions from pubchem and the sto-3g basis set.
mol = sf.Molecule.from_geometry(
    [
        sf.Atom("O", 8, jnp.array([0.0, 0.0, 0.0], dtype=jnp.float64)),
        sf.Atom("H", 1, jnp.array([0.52421003, 1.68733646, 0.48074633], dtype=jnp.float64)),
        sf.Atom("H", 1, jnp.array([1.14668581, -0.45032174, -1.35474466], dtype=jnp.float64)),
    ], basis_name="sto-3g",
)

# Jit compile and run SCF to solve for the energy.
result = jax.jit(scf.solve)(mol)

print(f"Total Energy: {result.total_energy} H")
print(f"Electronic Energy: {result.electronic_energy} H")

Output:

Total Energy: -74.96444760738025 H
Electronic Energy: -84.04881211726543 H

We can now evaluate the electron density on the points of a grid and save the result to a cube file so that we can render it with tools like 3dmol.

grid = sf.analysis_grid.build_bounding_grid(mol)
density_data = sf.analysis.evaluate(result.basis.basis_blocks, result.density, grid)

with open('density.cube', 'w') as f:
    sf.analysis.write_cube_data(
        mol=mol, grid=grid, data=density_data,
        description="density", f=f
    )

Here is what the result looks like rendered by 3dmol:

Installation

git clone https://github.com/lowdanie/hartree-fock-solver.git
cd hartree-fock-solver

pip install -e .

Tests

This project uses pytest for testing and pytest-cov for coverage.

# Run the standard test suite (skips slow tests by default)
pytest

# Run all tests (including slow ones)
pytest -m ""

# Check code coverage
pytest --cov=slaterform

Theory

For the details of the math, physics and algorithms behind this library see: