Getting started with DMFF: A comprehensive beginner's tutorial.

In the simulation of molecular systems, the underlying force field (FF) model plays an extremely important role, determining the reliability of the simulation. However, the quality of the state-of-the-art molecular force fields is still unsatisfactory in many cases, and the FF parameterization process largely relies on human experience, which is not scalable. To address this issue, we introduce DMFF (Differentiable Molecular Force Field), an open-source molecular FF development platform based on automatic differentiation technique. Using DMFF, both energies/forces and thermodynamic quantities such as ensemble averages and free energies can be evaluated in a differentiable way, realizing an automatic, yet highly flexible force field optimization workflow

For more details, please refer to this link: https://mp.weixin.qq.com/s/eVXTr1eU1-dGbC5UFybr0g

Note: JAX, DMFF, and other related tools are still undergoing rapid development and iteration. This notebook has been successfully executed using DMFF 0.2.0.

Choose dmff:0.2.0-notebook and you can directly run this notebook on Bohrium

This notebook is primarily organized by Wei Feng.

Before You Run

Molecular systems are a highly important category of systems, encompassing various organic molecules such as biomolecules, drug molecules, various coatings used in industries, porous materials like COFs, polymers and small molecule electrolytes in batteries, and more. Problems of interest in industries, such as structure prediction of biomolecules, drug screening, screening and design of materials, heavily rely on molecular dynamics (MD) simulations of molecular systems. Although the applications of MD have extended beyond molecular systems, simulating molecular systems remains one of the core applications of MD.

MD simulations usually rely on a potential energy surface (PES), which describes the energies and the forces of the simulated atoms. For the sake of computational efficiency, the PES is typically approximated using a classical model, namely molecular force field, instead of being computed ab initioly on-the-fly. Therefore, the quality of the underlying force field limits the accuracy and the predictive power of the simulation.

So it's obvious that the core of simulation is fitting microscale interactions, or potential functions, to accurately predict macroscopic properties. In classical molecular dynamics, the potential function is expressed mathematically in terms of a molecular force field.

What Is Molecular Force Field

As we mentioned above, one of the most important aspects in simulating molecular systems is characterizing the system's potential energy function and describing the interactions within the system. In Classical Molecular Dynamics (CMD) simulations, the potential energy function follows a fixed mathematical form, and the force field provides the mathematical functions for intermolecular and intramolecular interaction potentials. The traditional force fields commonly used in the industry, such as OPLS and GAFF, generally share a similar form:

The total energy can be divided into bonded and nonbonded components. The bonded component naturally depends on the intramolecular coordinates (bond lengths, bond angles, dihedral angles), while the nonbonded component (further decomposable into van der Waals interactions and electrostatic interactions) naturally depends on the interatomic distances. This form is almost universally used as the standard form in the entire industry. However, the molecular force fields currently employed in industrial applications face several challenges:

• Lack of portability and predictive capability: When studying new systems, it is often uncertain which force field will yield better results until actual simulations are performed. Apart from relying on "experience," it can be challenging to determine the appropriate force field selection criteria.
• Parameter's non-uniqueness and inconsistency: It is common to encounter multiple sets of completely different parameters that yield similar macroscopic predictions. As a result, it is not possible to solely rely on macroscopic data to determine which set of parameters is more reasonable at the microscale. Moreover, similar molecular systems often have entirely different force field parameters, making it challenging to combine force fields developed by different research groups.

Over the past few decades, the development of empirical force fields and the improvement of computational accuracy have been both an age-old and emerging topic. The quest to accurately and efficiently describe atomic interactions and establish a technical roadmap for developing a molecular force field applicable to the majority of systems has been the focus of researchers in the field of molecular force fields and molecular dynamics.

The optimization of parameters in force fields has long relied on manual intervention and "empirical" parameter tuning methods. This reliance raises concerns about the reliability and efficiency of empirical force field fitting. However, with the advent of the artificial intelligence era, one underlying technology, automatic differentiation, offers a new solution. This technology has paved the way for the development of Differentiable Molecular Force Field (DMFF), which holds the promise of becoming a powerful tool for molecular force field developers.

Automatic DIfferentiation & DMFF

Automatic differentiation is a high-precision and versatile method for computing derivatives in computer programs. It follows the mathematical chain rule to compute derivatives of composite functions. By tracing the derivative chains of every data point based on the computation graph, it can calculate the differentiation of the output with respect to the input variables.

Automatic differentiation plays a crucial role in optimizing neural network models. During model training, it is necessary to compute the gradients of the output with respect to the input variables through backpropagation and utilize gradient descent to optimize the model parameters. Automatic differentiation frameworks excel at achieving this task. With their efficient optimization capabilities for high-dimensional parameters, automatic differentiation techniques can be applied not only to neural network models but also to any framework that follows the "input model parameters → apply model computation → obtain computed results" structure. Molecular dynamics simulations precisely follow this workflow.

Therefore, leveraging automatic differentiation techniques and utilizing experimental or first-principles computed data as references, it is possible to optimize force field parameters by calculating the differentiation of the output results with respect to the input parameters through backpropagation. This enables the optimization of force field parameters in molecular dynamics simulations, just as in the optimization of neural network models.

Several attempts have been made to harness automatic differentiation techniques for molecular dynamics simulations, such as TorchMD, JAX-MD, and SPONGE. However, the deep computational graph involved in molecular dynamics simulations often introduces additional challenges. The backpropagation process, which calculates the differentiation from the output results to the input variables, can be computationally expensive. Additionally, there is still a lack of simulation engines specifically designed for rapid implementation and parameter optimization of force fields. Molecular force field developers urgently require comprehensive support for a wider range of force field functional forms and various types of objective functions. Now DMFF comes.

DMFF provides a comprehensive and rapid implementation of force field models and offers differentiable estimators of system energy, forces, and thermodynamic quantities. These differentiable estimators allow for the definition of corresponding object functions, enabling an automatic optimization process.

As mentioned above, the deep computational graph spanning the entire trajectory in molecular dynamics simulations would be computationally expensive and time-consuming. However, this limitation can be mitigated by employing a reweighting scheme for the trajectory. In DMFF, the reweighting algorithm is incorporated into the MBAR method, and extends the differentiable estimators for average properties and free energy calculations. We will demonstrate this method with practical examples to showcase its effectiveness.

DMFF gives a solution to two major problems in molecular force field development:

• How to ensure the rapid implementation and iteration of complex force fields in molecular dynamics simulations?
• How to improve the efficiency of optimizing parameters in high-dimensional functions and automate this process, while increasing the transferability of the parameters?

In our latest release, DMFF 0.2.0, we have explored a pathway for automatic optimization of force field parameters. The corresponding workflow has been validated in both simple small molecular systems and more complex electrolyte systems. We will guide you through a step-by-step notebook experience, demonstrating the complete workflow of DMFF for automatic optimization of force field parameters. Furthermore, using the notebook as a teaching template, we will provide a real hands-on experience showcasing the impact of DMFF on force field development.

Table of Contents:

• 1. Quick Start Guide to DMFF
  • 1.0 Import dependencies and prepare files
  • 1.1 Load existing force field parameters and topology
  • 1.2 Calculation
  • 1.3 [Review] Basic interface and key usage points of DMFF
• 2. Generating and extending potential functions
  • 2.1 Multipolar Polarizable Force Field
  • 2.2 Extending the potential function using backend functions
• 3. Bottom-Up Fitting
  • 3.1 Problem introduction
  • 3.2 Definition of potential function
  • 3.3 Preparing inputs for the potential function
  • 3.4 Load the data
  • 3.5 Definition of loss function
  • 3.6 Optimize the energy
• 4.Using DMFF for force field optimization of liquid dimethyl carbonate (DMC)
  • 4.1 Introduction
  • 4.2 Definition of potential function
  • 4.3 Definition of OpenMM sampler
  • 4.4 Initial MD sampling
  • 4.5 Definition of property calculation functions
  • 4.6 Read the data and perform the comparison
  • 4.7 Estimator initialization
  • 4.8 Definition of target ensemble
  • 4.9 Definition of loss function
  • 4.10 Optimizer setup and optimization loop
• 5. Summary & Outlook

1. Quick Start Guide to DMFF

This case is primarily contributed by Yanbo Han.

1.0 Import dependencies and prepare files

First, we set up the runtime environment and import the potential functions needed.

In addition to DMFF, we also need to use JAX and OpenMM.

• OpenMM: Manages the force field files and parameter data
• JAX：framework of differentiation

At the same time, we will also use other libraries and trajectory analysis software such as mdtraj in the subsequent examples. Let's import them as well.

Navigate to the working directory

1.1 Load existing force field parameters and topology | OpenMM Frontend

If you are using a force field in OpenMM, you can load the force field parameters using the ForceField class, like:

app.Topology.loadBondDefinitions("lig-top.xml")
pdb = app.PDBFile('lig.pdb')
pdb_system = pdb.topology

forcefield = ForceField('someforcefield.xml')

In DMFF, there is a functionality class called Hamiltonian that is similar to OpenMM's ForceField class. It allows you to read force field parameters and define a more generalized system potential energy function. At the same time, it is compatible with reading existing force field parameters.

• We can use OpenMM to read PDB and topology files：
  • topology（log-top.xml）
  • PDB（lig.pdb）
• Load the force field parameters using DMFF's Hamiltonian class and create a differentiable potential energy function:
  • GAFF force field file （gaff-2.11.xml）
  • Additional parameters for the corresponding molecule (assigned charges) （lig-prm.xml）

The DMFF potential function, apart from the name "Hamiltonian," is used in a similar way to OpenMM. The XML files of OpenMM force fields can also be directly reused.

In DMFF, the parameters and calculations of the potential function are managed by JAX. For example, the DMFF potential function includes the re-implemented HarmonicBondForce, HarmonicAngleForce, PeriodicTorsionForce, and NonbondedForce. The parameters in DMFF are JAX-wrapped Array objects. For example, if we define:

jnp.array([1.])

You will obtain a DeviceArray type. This type has a similar interface to numpy.ndarray(), as both are high-performance arrays/matrices. However, unlike numpy, JAX's jax.numpy.array() can be extended to work with GPUs/TPUs, making it suitable for efficient automatic differentiation.

1.2 Calculation | JAX Differentiable Backend

In the calculation of the defined potential functions, we require the following parameters:

• coordinates: just use the coordinates in PDB file

• box: We need to add the definition of the box since our PDB file does not contain this information. However, it is worth mentioning that if we use None for the box size, we can still obtain results because our system does not involve periodic boundaries.

• pairs: The GAFF2 potential function requires the input of NeighborList to calculate nonbonded forces. Therefore, we can use the NeighborList class to obtain the pairs for the energy calculation.

Then, we can pass these to the get_energy function, which is generated by the generator from the parsed XML force field parameters in potentials.dmff_potentials. For example, we can calculate the energy associated with the NonbondedForce interactions as follows:

For the system we defined above, the total energy can be calculated as follows:

$$E_{\mathrm{total}}^{GAFF2}=E_{\mathrm{bond}}+E_{\mathrm{angle}}+E_{\mathrm{torsion}}+E_{\mathrm{nonbond}}$$

To calculate the total energy using the previously defined potential, you can use the getPotentialFunc() method, which will return the function for calculating the total energy.

So far, the methods we have used to calculate the system energy are similar to those in OpenMM.

The main advantage of using JAX as a computational backend is that we can use the jax.grad function to obtain the gradients of a function. Its syntax is jax.grad(func, argnums), where it calculates the (partial) derivatives of the function with respect to the argument specified by argnums.

The interface of the total energy calculation function efunc that we obtained is [coordinates, box, bond pairs, force field parameters].

By applying the "differentiation of the function" operation, specifically taking the partial derivatives of the total energy with respect to the coordinates, we can compute the partial derivatives of the total energy with respect to the coordinates. These derivatives can then be used to calculate the forces acting on the atoms in the molecule.

$$\frac{\partial E_{\mathrm{total}}}{\partial {\mathbf{Z}}_{\mathbf{i}}}=-{\mathbf{F}}_{\mathbf{i}},i=x,y,z$$

Similarly, we can also take derivatives with respect to the force field parameters. For example, if we want to optimize the parameters related to the nonbonded interactions, we can adjust the argnum parameter accordingly. By doing so, we can obtain the derivative of the nonbonded energy with respect to the charge parameter.

1.3 [Review] Basic interface and key usage points of DMFF

1.3.1 Applications of DMFF：

Molecular Force Field Optimization Platform

• DMFF is compatible with OpenMM's XML format for molecular force field parameters
• Automatic differentiation and GPU support enable fast implementation of complex molecular dynamics force fields in DMFF
• flexible force field parameter optimization capability of DMFF：
  • Improving a specific component among them×
  • Complex parameter optimization√ Rapid implementation through automatic differentiation frameworks

1.3.2 General Operations

Generating Potential Functions

• Defining force parameters in an XML file (OpenMM interface)
• Description of the system in PDB and topology in XML format

pdb = app.PDBFile("lig.pdb")
ff = Hamiltonian("gaff-2.11.xml", "lig-prm.xml")
potentials = ff.createPotential(pdb.topology)

Calculating System's Energy and Forces

efunc = potentials.getPotentialFunc()
params = ff.getParameters()

# energy function
totene = efunc(positions, box, pairs, params)
print(totene)

# 力
pos_grad_func = jax.grad(efunc, argnums=0)
force = -pos_grad_func(positions, box, pairs, params)
print(force)

1.3.3 [More] Why do we need "differentiable" JAX: The first step towards infinite possibilities for new force fields.

Why JAX？

Here we take the example of the implementation of the HarmonicBondForce in DMFF to build a simple model to illustrate what we can do with the JAX backend of DMFF.

In this part we will use some customized files (modifying these files might help you understand what the DMFF frontend does):

we an create a new file dummy-prm.xml

<ForceField>
    <Residues>
        <Residue name="DUM">
            <Atom name="C1" type="a" />
            <Atom name="C2" type="a" />
            <Bond atomName1="C1" atomName2="C2"/>
        </Residue>
    </Residues>
</ForceField>

dummy-top.xml：

<Residues>
    <Residue name='DUM'>
        <Bond from='C1' to='C2'/>
    </Residue>
</Residues>

dummy.pdb：

REMARK   1 CREATED WITH MANUAL
HETATM    1  C1  DUM A   1       0.000   0.000   0.000  1.00  0.00           C  
HETATM    2  C2  DUM A   1       1.000   0.000   0.000  1.00  0.00           C  
TER       3      DUM A   1
END

<ForceField>
  <AtomTypes>
    <Type element="C" name="a" class="a" mass="12.01"/>
    </AtomTypes>
  <HarmonicBondForce>
    <Bond type1="a" type2="a" length="0.10" k="100"/>
  </HarmonicBondForce>
</ForceField>

• By taking derivatives with respect to coordinates, we can perform molecular dynamics (MD) simulations.
• By taking derivatives with respect to the box, we can perform NPT ensemble molecular dynamics (MD) simulations.
• By taking derivatives with respect to the parameters, we will be able to optimize the parameters of our defined force field using techniques such as gradient descent.

1.4 Code

https://gitee.com/deepmodeling/DMFF https://github.com/deepmodeling/DMFF

2. Generating and extending potential functions

This case is primarily contributed by Professor Kuang Yu

In the following example, using water as an illustration, we will demonstrate: 1. How to define a potential function using the DMFF frontend; 2. How to extend the backend function of a multipolar polarizable force field to introduce charge fluctuations with respect to the structure.

Navigate to the working directory

2.1 Multipolar Polarizable Force Field

Let's briefly introduce the concept of a multipolar polarizable force field. In such force fields (e.g., AMOEBA, MPID), each atom is associated with a local coordinate system defined by its neighboring atoms. Within this local coordinate system, the atom can have a permanent multipole moment ($Q$). In addition to the permanent multipole moment, each atom also carries an induced dipole moment ($\mu$), which responds to the electric field from the surrounding water molecules. The energy of the system, given a certain induced dipole moment, can be expressed as:

$$E(\mathbf{R},Q_i,{\mu }_i)=\sum _{i\ne j}{Q}_i{\hat{T}}_{ij}^{\text{perm}}{Q}_j+Q_i{\hat{T}}_{ij}^{\text{perm-ind}}{\mu }_j+{\mu }_i{\hat{T}}_{ij}^{\text{ind-ind}}{\mu }_j$$

To compute the actual potential energy surface, it is necessary to adjust the induced dipole moments to minimize the energy:

$$U(\mathbf{R},Q_i)=\underset{\mu }{\min }U(\mathbf{R},Q_i,{\mu }_i)$$

In the DMFF ADMP module, we provide full support for such potential energy models. We start by writing the corresponding forcefield.xml file:

<ForceField>
 <AtomTypes>
  <Type name="380" class="OW" element="O" mass="15.999"/>
  <Type name="381" class="HW" element="H" mass="1.008"/>
 </AtomTypes>
 <Residues>
  <Residue name="HOH">
   <Atom name="H1" type="381"/>
   <Atom name="H2" type="381"/>
   <Atom name="O" type="380"/>
   <Bond from="0" to="2"/>
   <Bond from="1" to="2"/>
 </Residue>
 </Residues>

 <ADMPPmeForce lmax="2"
     mScale12="0.00" mScale13="0.00" mScale14="0.00" mScale15="1.00" mScale16="1.00"
     pScale12="0.00" pScale13="0.00" pScale14="0.00" pScale15="1.00" pScale16="1.00"
     dScale12="1.00" dScale13="1.00" dScale14="1.00" dScale15="1.00" dScale16="1.00">
   <Atom type="380" kz="381" kx="-381"
             c0="-0.803721"
             dX="0.0" dY="0.0"  dZ="-0.00784325"
             qXX="0.000366476" qXY="0.0" qYY="-0.000381799" qXZ="0.0" qYZ="0.0" qZZ="1.53231e-05"
             />
   <Atom type="381" kz="380" kx="381"
             c0="0.401876"
             dX="-0.00121713" dY="0.0"  dZ="-0.00095895"
             qXX="6.7161e-06" qXY="0.0" qYY="-3.37874e-05" qXZ="1.25905e-05" qYZ="0.0" qZZ="2.70713e-05"
             />
   <Polarize type="380" polarizabilityXX="1.1249e-03" polarizabilityYY="1.1249e-03" polarizabilityZZ="1.1249e-03" thole="0.33"/>
   <Polarize type="381" polarizabilityXX="2.6906e-04" polarizabilityYY="2.6906e-04" polarizabilityZZ="2.6906e-04" thole="0.33"/>
 </ADMPPmeForce>
</ForceField>

As we can see, we have defined two atom types, O(380) and H(381), and defined the connectivity topology template for water molecules. In the <ADMPPmeForce> section, we have defined the charges, dipoles, quadrupoles, and polarizabilities for the two atom types. Next, we create a Hamiltonian and use the molecular topology to build the potential energy surface:

efunc represents the newly generated potential function. It's important to note that the createPotential function returns a Potential object. From this object, you can obtain the total potential energy function using pots.getPotentialFunc(), and you can also access each individual force field term defined in the XML using pots.dmff_potentials.

To use the newly obtained potential function, you can follow this approach:

Due to the JIT compilation, the first call may be slower, but the subsequent calls will be significantly faster.

Up to this point, we have completed the definition of the potential function in the frontend. The definition and usage of other potential functions are similar. Next, let's take a look at how we can extend this potential function.

2.2 Extending the potential function using backend functions

In this case, we aim to introduce charge fluctuations into the existing multipolar polarizable force field. In the previous model, the permanent multipoles of all atoms were considered constant. However, we observe that the charges of water molecules vary with the changes in bond lengths and bond angles due to molecular vibrations. In other words, atomic parameters, including charges and dispersion coefficients, should ideally be functions of the molecular geometry. Although the induced dipoles from other molecules are already considered in the potential function, the charge fluctuations resulting from the structural changes of the molecule itself are not taken into account. To incorporate this into our model, we can define a new potential function based on the existing one.

First and foremost, we need to acknowledge that the previously defined frontend function is not suitable for such extensions. This is because the efunc function's input is based on force field parameters rather than atomic parameters. We can examine the charge values in the params variable within the efunc function's input:

It is evident that efunc and gfunc only accept charges for the O and H atoms, and they do not allow specifying individual charges for each atom (i.e., by default assigning the same charge to all O atoms). This is a characteristic of the DMFF frontend functions: they handle the typification process internally and do not expose per-atom parameters as inputs. However, in our example, we want to specify charges for each atom individually. Therefore, we need to call the backend functions, which are typically stored in the generator object:

It can be seen that the backend function's API design is more complex. This function is not responsible for typification and accepts the local multipole moments (Q_local), polarizability (pol), and the thole damping width (thole) for each atom as inputs. Therefore, it offers greater flexibility for extension. Additionally, the generator object also contains other information that we will need later, such as the atom_type that records the atom types for each atom (we will use this variable for expanding the force field parameters):

We would like to implement the following computation process:

That is, introduce structural dependence into Q based on the positions: $$U=U(\mathbf{R},Q_i(\mathbf{R}))$$

To do this, let's define a function that calculates the charges based on the molecular structure:($Q(\mathbf{R})$)：

Then, let's define a function generator that combines the compute_charges function and the pme_force.get_energy function:

Next, we can calculate the corresponding energy and forces:

By comparing the forces before and after, we can clearly see the effect of geometric fluctuations of atomic charges on the atomic forces.

3. Bottom-Up Fitting

This case is primarily contributed by Junmin Chen

In this case, we will demonstrate how to use DMFF in a bottom-up fitting approach.

Navigate to the working directory

3.1 Problem introduction:

The main objective of bottom-up fitting is energies and forces calculated from first-principle calculations. The definition of this object function is relatively straightforward. In this example, we will focus on fitting the Pauli Exchange interaction between two PEG dimers.

The intermolecular exchange repulsion term for this system can be written in the following form:

$$\begin{array}{r}\{\begin{array}{ll}{E}_{ex}& ={\sum }_{i<j}{A}_{ij}\left(\frac{1}{3}(B_{ij}r{)}^2+B_{ij}r+1\right)\exp (-B_{ij}r)\\ A_{ij}& =A_i\cdot A_j\\ B_{ij}& =\sqrt{B_i\cdot B_j}\end{array}\end{array}$$

In this fitting process, the target is the exchange energy calculated using the SAPT (Symmetry-Adapted Perturbation Theory) first-principle method. Based on past experience, there is a strong correlation between the prefactor（$A$）and exponents（$B$）in this fitting, which can lead to overfitting and yield unrealistic parameters. Here, we aim to utilize the ADAM algorithm in machine learning for optimization, with the hope of obtaining physically reasonable parameters. The system consists of five atom types, which are defined in the following order according to the XML file:

C1, H1, O, C2, H2

All atomic parameters ($A_i$&$B_i$) will also be arranged in the same order as mentioned above.

3.2 Definition of potential function

We define the relevant Hamiltonian (we will define different Hamiltonians and potential functions for the dimer and the two monomers separately):

And obtain the corresponding potential functions:

Note that in forcefield.xml, there is a complete force field that contains many components. Here, we are only interested in the exchange part, so we extract the relevant potential function from it.

In the subsequent calculation of the loss function, we will process minibatches of structures as input. To speed up the computation, we can use jax.vmap to vectorize the single potential calculation function into a batched energy calculation function.

During this process, we specify vectorization along the first axis of the first variable (positions).

3.3 Preparing inputs for the potential function (such as neighbor lists, etc)

Next, we prepare the inputs required for these potential functions, with the most important one, the neighbor list. For a simple system with only two molecules, we can set rc to a relatively large value to include all possible interactions. This way, we don't need to update the neighbor list due to structural changes, which will speed up the calculations.

First, we obtain the initial structure and box size:

（WARNING: All input functions in the ADMP module use Angstrom as the unit, which is different from the classical module where the default unit is nm.）

Then, construct a list of all atom pairs using the initial structure:

Next, obtain the initial parameters as the starting point for optimization:

3.4 Load the data

The dataset is obtained from first-principle (SAPT) calculations and consists of 50 dimer scans. Each scan is identified by a unique scan ID and contains 12 structures. We can take a look at the storage structure of the dataset:

We only use the exchange energy data from the dataset.

In the subsequent fitting process, we will shuffle the order of all scans and use each scan's data as a minibatch input. We will use ADAM to train the parameters.

3.5 Definition of loss function

To define the energy loss function for the intermolecular exchange repulsion interaction, we can use the following formulation:

In this function, we have added a statistical weight to each data point:

$$W(E_{tot})=\{\begin{array}{l}\begin{array}{rlrl}& 1.0& & \text{ if }{E}_{tot}<25\\ & \exp \left(-\frac{E_{tot}-25}{kT}\right)& & \text{ otherwise}\end{array}\end{array}$$

By assigning weights to the data points, we can effectively ignore those with significantly higher total energies, which may not be important for the actual simulations.

3.6 Optimize the energy

Then initialize our optimizer:

To start our optimization loop, we can iterate over the data for a specified number of epochs and update the parameters using the optimizer. After each epoch, we can save the parameters as a pickle file.

We will first perform 50 epochs of optimization, but of course, we can continue to optimize further to achieve better results. We can then examine whether the optimized $B$ parameters are physically reasonable.

Finally, we can evaluate the fitting performance:

4.Using DMFF for force field optimization of liquid dimethyl carbonate (DMC)

In the following case, we will provide a practical application of DMFF and build the corresponding workflow.

This case is primarily contributed by Wei Feng

4.1 Introduction

The initial step in building a workflow using DMFF is to generate the potential function. The process and API are similar to OpenMM and generally follow the following steps:

1. The first step in building a force field using DMFF is to define the force field object, known as the Hamiltonian. It encapsulates the typification rules and all the force field parameters. It can be seen as the equivalent of an XML file in OpenMM or an ITP file in GROMACS, but represented as a Python object.
2. Obtain the molecular bond connectivity information for the system you want to simulate (you can use a PDB file).
3. Based on the system's connectivity topology, you can use the Hamiltonian to classify and parameterize each atom, bond length, bond angle, dihedral angle, etc., and create a potential function. The function takes inputs such as atomic positions, box size, and force field parameters, making it suitable for subsequent derivative calculations.
4. If necessary, the potential function can be modified and extended as needed. For example, using jax.grad, one can obtain a function to compute forces and the Virial matrix, which can be directly used in molecular dynamics (MD) simulations.
5. Estimator and loss functions can be defined based on the potential function to compute desired properties. The optimization of parameters can be performed using a gradient descent optimizer.

In DMFF, we provide JAX implementations of various potential functions commonly used in traditional physics force fields, such as PME, pairwise interactions, multipole moments, and polarizable force fields. With the help of JAX's transformation tools, users can differentiate, vectorize, compile, and modify these potential functions. They can also create new potential functions by re-encapsulating and recombining existing ones. Additionally, users can develop property estimators based on these potential functions, and these new functions can also be further differentiated or vectorized. This function-centric programming paradigm in DMFF offers unbelievable flexibility in defining new force fields and objective functions.

To provide users with a comprehensive understanding of the main modules and functionalities in DMFF and how these modules work together in force field development, we will demonstrate the complete process of model definition, optimization, and deployment using a coarse-grained model of dimethyl carbonate (DMC). The goal of this task is to optimize the interactions between three sites representing the DMC chain (3-site model) and reproduce the radial distribution function (RDF) of liquid DMC obtained from X-ray diffraction (XRD) experiments. We will also compare the DMC coarse-grained model force field with the all-atom OPLS-AA force field.

Navigate to the working directory

4.2 Definition of potential function

Load the topology of the DMC system and construct the potential function for the simulation box:

From the prm.xml force field file, we can see that the function includes the following terms for intramolecular atomic interactions: bond, angle, and electrostatics, as well as intermolecular interactions represented by the Lennard-Jones potential:

$$\begin{array}{rl}V(\mathbf{R})& =V_{\mathrm{bond}}+V_{\mathrm{angle}}+V_{\mathrm{elec}}+V_{\mathrm{vdW}}\\ & =\sum _{\mathrm{bonds}}{k}_r(r-r_0{)}^2+\sum _{\mathrm{angles}}{k}_{\theta }(\theta -{\theta }_{eq}{)}^2+\sum _{\mathrm{i}<\mathrm{j}}\frac{q_i{q}_j}{4\pi {ϵ}_0{r}_{ij}}\\ & +\sum _{ij}4{\epsilon }_{ij}\left[{\left(\frac{{\sigma }_{ij}}{r_{ij}}\right)}^{12}-{\left(\frac{{\sigma }_{ij}}{r_{ij}}\right)}^6\right]\end{array}$$

And we will optimize some parameters.

4.3 Definition of OpenMM sampler

The function generates an NPT ensemble using OpenMM based on the input parameters and saves the generated samples in the specified trajectory file. This function will be repeatedly called for resampling purposes during subsequent optimization steps.

4.4 Initial MD sampling

4.5 Definition of property calculation functions

In the fitting process, we need to simultaneously fit the Radial Distribution Function (RDF) and density. Therefore, we will define separate calculation functions for RDF and density:

When calculating the ensemble average of an observable quantity $A$, the following statistical quantities need to be computed: $⟨A{⟩}_p={\sum }_{n}W(x_n;\theta )A(x_n;\theta )$, code above has defined the $A(x_n;\theta )$ function.

In general, we need to encapsulate $⟨A{⟩}_p$ in a function and then take its overall derivative during each optimization iteration. However, in this case, since the RDF and density are purely structural properties that are independent of $\theta$ and do not participate in the derivative calculation, we can precompute the RDF for each sample ($A(x_n)$) and store it. It is not necessary to recalculate it in every iteration.

In the code above, we use the existing tools in mdtraj to compute the RDF for each frame, rather than using a differentiable implementation with jax.。

4.6 Read the data and perform the comparison

Read the experimental data and pre-compute the RDF for the DMC full-atom model in the OPLS-AA force field (generated by LigPargen) as comparison:

We will compare the RDF results obtained from the RDF calculation function with the experimental data and the RDF results from the full-atom model in the OPLS-AA force field. Let's start by looking at the RDF generated from the initial parameter sampling, which serves as our starting point for optimization.

Up to now, we have completed the preparation for force field optimization. Before starting the optimization loop, we need to prepare some few tools. Here, we introduce the distinctive feature of DMFF version 0.2.0 - Property Estimator.

The computational cost of computing deep computational graphs that span the entire trajectory in molecular dynamics simulations is significantly high in terms of both time and computation. However, this drawback can be mitigated by a trajectory reweighting approach. In DMFF, the reweighting algorithm is introduced in the MBAR method, which extends the differentiable estimators for ensemble averages and free energies.

In the gradient descent training of parameter optimization, the parameters are only perturbed slightly in each training epoch, allowing us to continue using the samples from the previous epoch without significant deviation from the sampling ensemble. This eliminates the need for re-sampling until the optimized ensemble significantly deviates from the sampling ensemble, resulting in a substantial reduction in optimization time and computational cost.

In the reweighted MBAR Estimator, we define two ensembles: the sampling ensemble from which all samples are extracted, and the target ensemble that we want to optimize. The sampling ensemble is updated only when necessary and does not require differentiability. Its data can be generated by external samplers such as OpenMM. The Estimator provides weighting factors, allowing us to estimate the target ensemble by the weighted average of samples from the sampling ensemble.

The workflow can be understood using the following flowchart:

The Effective Sample Size is a measure of the degree to which the sampling ensemble deviates from the target ensemble. When the deviation becomes too large, the sampling ensemble is updated.

4.7 Estimator initialization

Here we create MBAREstimator and initialize it. There are two parts of variables that need to be initialized in MBAREstimator:

• all sampling state information（state）and samples $x_n$（sample） ；
• The partition function of the sampling state, estimated based on the samples and sampling state information. ( ${\hat{f}}_m$)

It can be observed that in the above code, we define an OpenMMSampleState and add this state and the samples generated from this state to the estimator. Then, the optimize_mbar function is called (which internally calls the external pymbar tool) to estimate the partition function of all the sampling states. It is important to note that these operations do not involve differentiation and do not depend on JAX. After initialization, the estimator can directly provide $W(x_n)$ in a differentiable manner.

4.8 Definition of target ensemble

Here, we define our target ensemble target_state using the differentiable 3-site model potential function efunc that was defined above. This target_state is used as an input parameter for subsequent reweighting. In the computation process, the estimator will repeatedly call the efunc from the target_state to calculate the effective energy of the samples in the target ensemble. For the NVT ensemble, it is calculated as $u=U(x_n)$, and for the NPT ensemble, it is calculated as $u=U(x_n)-pV$.

4.9 Definition of loss function

Here, we use the estimator.estimate_weight function to generate the $W(x_n;{\theta }_p)$ corresponding to the current target_state. We then calculate the average RDF value under the current target_state. Finally, we compare the calculated RDF with the experimental reference values to obtain the RDF object function:

$L_{\text{RDF}}=\text{log}\Vert {\sum }_{n=1}^N{W}_n{\mathbf{v}}_n-{\mathbf{v}}_{\text{ref}}\Vert$

At the same time, we have noticed that when optimizing RDF alone, the optimizer may explore non-physical parameter regions, leading to a decrease in system density and eventually crossing the phase boundary to become gas. To avoid this phenomenon, we also include density error in the loss function to control the system density. Therefore, we separately define the loss function for density:

$L_{\rho }={\left({\sum }_{n=1}^N{W}_n{\rho }_n-{\rho }_{\text{ref}}\right)}^2$

In the following optimization loop, we will first perform several rounds of optimization for density and then proceed with the RDF.