chemtrain: Learning Deep Potential Models via Automatic Differentiation and Statistical Physics

Abstract

Neural Networks (NNs) are effective models for refining the accuracy of molecular dynamics, opening up new fields of application. Typically trained bottom-up, atomistic NN potential models can reach first-principle accuracy, while coarse-grained implicit solvent NN potentials surpass classical continuum solvent models. However, overcoming the limitations of costly generation of accurate reference data and data inefficiency of common bottom-up training demands efficient incorporation of data from many sources. This paper introduces the framework chemtrain to learn sophisticated NN potential models through customizable training routines and advanced training algorithms. These routines can combine multiple top-down and bottom-up algorithms, e.g., to incorporate both experimental and simulation data or pre-train potentials with less costly algorithms. chemtrain provides an object-oriented high-level interface to simplify the creation of custom routines. On the lower level, chemtrain relies on JAX to compute gradients and scale the computations to use available resources. We demonstrate the simplicity and importance of combining multiple algorithms in the examples of parametrizing an all-atomistic model of titanium and a coarse-grained implicit solvent model of alanine dipeptide.

Package source code published at http://github.com/tummfm/chemtrain

References in corpus (11)

• ANI-1: An extensible neural network potential with DFT accuracy at force field computational cost
• DeePMD-kit v2: A software package for Deep Potential models
• Accurate sampling using Langevin dynamics
• Uncertainty estimation for molecular dynamics and sampling
• Machine Learning Implicit Solvation for Molecular Dynamics
• Towards empirical force fields that match experimental observables
• Accurate machine learning force fields via experimental and simulation data fusion
• PhysNet Meets CHARMM: A Framework for Routine Machine Learning / Molecular Mechanics Simulations
• FeNNol: an Efficient and Flexible Library for Building Force-field-enhanced Neural Network Potentials
• Guest Editorial: Special Topic on Software for Atomistic Machine Learning
• JaxSGMC: Modular stochastic gradient MCMC in JAX

Cited by in corpus (4)

• Refining Coarse-Grained Molecular Topologies: A Bayesian Optimization Approach
• Energy-Based Coarse-Graining in Molecular Dynamics: A Flow-Based Framework without Data
• MDRefine: a Python package for refining Molecular Dynamics trajectories with experimental data
• Differentiable hybrid force fields support scalable autonomous electrolyte discovery