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.
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 methane. We initialize the nuclei to lie flat on a plane, and the optimizer moves them into the classic tetrahedral configuration. The blue cloud is rendered by sampling the electron density returned by scf.solve.
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:
git clone https://github.com/lowdanie/hartree-fock-solver.git
cd hartree-fock-solver
pip install -e .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=slaterformFor the details of the math, physics and algorithms behind this library see: