Density Functional Theory (DFT) calculations using PyTorch

All the calculations in DFT are fully implemented with PyTorch, which allows for automatic differentiation and GPU acceleration. Note that the code is not optimized for speed, but rather for readability and ease of understanding. The motivation of this work is to accelarate the development of deep learning methods for DFT calculations. Since all the parts of the code are implemented with PyTorch, it is easy to use the code in combination with deep learning methods.

Installation

# Install PyTorch and PyTorch Geometric
conda install pytorch==2.2.2 pytorch-cuda=12.1 -c pytorch -c nvidia
pip install torch_geometric==2.5.3
pip install torch_scatter torch_cluster -f https://data.pyg.org/whl/torch-2.2.0+cu121.html

# Install other dependencies
conda install -c rdkit rdkit
pip install torchtyping torchcubicspline xitorch pyscf pylebedev basis_set_exchange

# Install torchdft
git clone https://github.com/fate1997/torchdft.git
cd torchdft
pip install -e .

Usage

from torchdft.mol import Mol
from torchdft.solver.scf import SCFSolver


mol = Mol.from_file('../tests/example/ntf2.sdf')

solver = SCFSolver(
    basis='sto-3g', 
    num_radial=24, 
    num_angular=110, 
    l_max=8,
    device='cuda:0', 
    integral_method='pyqint', # or pytorch
    verbose=False
)
energy = solver.solve(mol)

Benchmark

We compare the torchdft with PySCF, the results are shown below:

Molecule/Ion	PySCF	torchdft
H₂O	-73	-73
SO₂	-534	-535
HBr	-2540	-2495
C₂H₆	-73	-73
NTF₂⁻	-1750	-1776
Imidazolium	-207	-213
Pyridinium	-227	-245
Amino	-265	-288

For the computation time, we compare the torchdft with PySCF and pydft on NTF₂⁻ molecule:

Package	Time (s)
PySCF	3.8
torchdft + pyqint integral + GPU	7.1
torchdft + pyqint integral + CPU	275.0
torchdft + pytorch integral + GPU	45.7

Note that the torchdft with pytorch integral is not faster than torchdft with pyqint integral. This is reasonable because pyqint is implemented with C++. Despite the slower computation time, the torchdft with pytorch integral is useful for investigating how deep learning can accelarate or approximate these integrals.

Acknowledgement

This work is highly dependent on pydft and pyqint. We thank the author (@ifilot, Ivo Filot) of these packages for their great work. And we also thank the author's book Elements of Electronic Structure Theory, which is a great resource for understanding the DFT calculations.