Automatic-differentiation Gaussian-process surrogates for molecular geometry optimization
Chong Teng · Google Scholar · ORCID · Primary publication
Portfolio snapshot. This repository is frozen at commit
dafac58from 30 June 2025. It preserves a research codebase for technical review and provenance; it is not presented as a current production release.
MadGP is a collaborative scientific machine-learning codebase for constructing Gaussian-process (GP) surrogate potential-energy surfaces (PESs). Its central idea is simple: a researcher specifies a scalar kernel and, optionally, a physics-based prior model; MadGP uses automatic differentiation to assemble the value, gradient, and mixed-derivative covariance blocks needed for energy, force, or joint energy–force learning. This removes the need to derive and maintain every kernel derivative by hand.
The models are integrated with the Atomic Simulation Environment (ASE) to form an on-the-fly optimization loop:
flowchart LR
A[ASE structure and reference calculator] --> B[Energy and force evaluation]
B --> C[Automatic-differentiation GP update]
C --> D[Relaxation on the surrogate PES]
D --> E[Reference calculation and convergence check]
E -->|continue| B
E -->|converged| F[Optimized geometry]
- Differentiable GP construction. JAX and PyTorch implementations generate the derivative covariance operators for the core energy–force models from a user-defined scalar kernel.
- Multiple observation models. The optimizer factory supports energy-only, force-only, joint energy–force, and internal-coordinate variants.
- Physics-informed modeling. Constant, force-field, electronic-structure, mixed, and GP-as-prior models can be combined with Cartesian, Coulombic, and internal-coordinate representations.
- Atomistic workflow integration. ASE connects surrogate training and relaxation to reference calculators; the included end-to-end check uses GFN2-xTB.
- Research scaling prototypes. The snapshot includes committee/multi-PES, local, sparse, conjugate-gradient, and distributed-kernel experiments. These modules are research prototypes rather than a uniform production API.
MadGP is collaborative work. This snapshot highlights my contributions as a principal maintainer across:
- physical, mixed, and adaptive prior models;
- curvilinear kernels and internal-coordinate geometry optimization;
- ASE and electronic-structure calculator integration;
- committee, local, sparse, and distributed GP experiments; and
- numerical reliability, test coverage, and research benchmark support.
The publication author lists and repository history provide the complete credit for the project. Key collaborators on the original software paper were Daniel Huang, Junwei Lucas Bao, and Jean-Baptiste Tristan.
D. Huang, C. Teng, J. L. Bao, and J.-B. Tristan, “mad-GP: automatic differentiation of Gaussian processes for molecules and materials,” Journal of Mathematical Chemistry 60(6), 969–1000 (2022).
The paper introduces the automatic-differentiation design and evaluates
energy-, force-, and joint energy–force GP surrogates for small-molecule
geometry optimization. The core implementation is in
madgp/gplib/ and madgp/geoopt/.
| Research direction | Publication | Representative code in this snapshot |
|---|---|---|
| Physical and posterior-type priors | Dual-Level Training of Gaussian Processes with Physically Inspired Priors for Geometry Optimizations (2022) | madgp/gplib/prior/ |
| Adaptive ab initio priors and curvilinear optimization | A Spur to Molecular Geometry Optimization (2023) | madgp/geoopt/internal_coord/ and energy_force_internal_geo_opt.py |
| Reusing learned PESs across conformers | Exploring Torsional Conformer Space with Physical Prior Mean Function-Driven Meta-Gaussian Processes (2023) | gppes_prior.py and adaptive-prior logic |
| Kernel/coordinate choices for 20 oligopeptides | First-Principle Oligopeptide Structural Optimization with Physical Prior Mean-Driven Gaussian Processes (2025) | periodic/internal-coordinate kernels and madgp/data/polypeptide/ |
| Physical-prior committee GP dynamics | Physical Prior Mean-Driven Bayesian Committee Molecular Dynamics (BCMD) (2025) | ef_bcm_multi.py and multi-PES/variance components |
The snapshot contains shared infrastructure and research implementations developed across these studies. Paper-specific production workflows, complete benchmark inputs, or dynamics drivers may live in supplementary material or separate repositories; this repository should not be treated as a one-command reproduction package for every paper. A related downstream application to minimum-energy reaction paths is described in the 2024 CI-NEB study, but its NEB driver is not included here.
The historical environment is best reproduced with Python 3.10. Install an ASE-compatible reference calculator before installing the package; the bundled check uses xTB.
git clone https://github.com/SingletC/madgp_snapshot.git
cd madgp_snapshot
conda create --name madgp-snapshot python=3.10
conda activate madgp-snapshot
conda install -c conda-forge xtb-python
python -m pip install -e .Run the end-to-end check:
python check_install.pyThis is a scientific integration check, not a lightweight import test: it runs JAX- and PyTorch-backed geometry optimizations of a Baker-set molecule with GFN2-xTB. Runtime therefore depends on the local numerical and quantum-chemistry environment.
from ase.io import read
from xtb.ase.calculator import XTB
from madgp.geoopt.optim import create_optimizer
from madgp.gplib.energy_force_gp_pes import EnergyForceGPPES
from madgp.gplib.kernel.matern_kernel import jax_matern_kernel
from madgp.gplib.kernel.structure.direct import cartesian
from madgp.gplib.numlib import JaxNumLib
def prior(atoms, coordinates):
return 0.0
def kernel(atoms_x, atoms_y, x, y, hyperparameters):
return jax_matern_kernel(
x,
y,
weight=hyperparameters[0],
length_scale=hyperparameters[1],
)
model = EnergyForceGPPES(
prior,
cartesian(kernel),
numlib=JaxNumLib,
hyper_parameters=(1.0, 0.4),
)
atoms = read("madgp/data/Baker_25/01.xyz")
atoms.calc = XTB(method="GFN2-xTB")
optimizer = create_optimizer(model)
converged, optimized_atoms = optimizer.run(atoms, fmax=0.05)
print(f"converged={converged}")| Path | Purpose |
|---|---|
madgp/gplib/ |
GP posteriors, kernels, priors, committee and local models |
madgp/geoopt/ |
ASE-driven surrogate optimizers and coordinate transformations |
madgp/data/ |
Molecular structures used by examples and research tests |
tests/ |
Unit and scientific integration tests |
examples/ |
Introductory notebooks and scripts |
experiments/ |
Historical research notebooks preserved for provenance |
setup.pyrecords the later June 2025 package constraints;requirements.txtis a legacy environment file and contains older pins.- Several tests require external programs, reference data, or substantial numerical work. A passing import does not validate a chemistry workflow.
- JAX and PyTorch cover the core GP models, but experimental modules do not provide feature-for-feature backend parity.
- The notebooks and Sphinx files preserve historical research context and may use older APIs or paths.
For a compact introduction, start with
examples/example1.ipynb. For implementation-level
behavior, the tests are the more reliable reference for this frozen revision.
If this software contributes to published work, please cite the primary paper:
@article{huang2022madgp,
author = {Huang, Daniel and Teng, Chong and Bao, Junwei Lucas and Tristan, Jean-Baptiste},
title = {mad-GP: Automatic Differentiation of Gaussian Processes for Molecules and Materials},
journal = {Journal of Mathematical Chemistry},
volume = {60},
number = {6},
pages = {969--1000},
year = {2022},
doi = {10.1007/s10910-022-01334-x}
}The code is available under the Apache License 2.0.