Gradients

A most fundamental feature of this framework is its ability to rapidly compute the gradients of physical properties with respect to the force field parameters used to estimate them.

Note

Prior to v0.3.0 of this framework a combination of re-weighting and the central finite difference was employed to estimate the gradients of observables. From v0.3.0 onwards the fluctuation method [1] is instead used. The change was made to, in future, enable better integration with automatic differentiation libraries such as jax, and differentiable simulation engines such as timemachine.

Theory

The framework currently employs the fluctuation approach [1] to compute gradients of observables with respect to the force field parameters used to estimate them.

This approach may be derived by direct differentiation of the ensemble average an observable \(X\):

\[\left<X\left(\theta\right)\right> = \dfrac{1}{Q\left(\theta\right)} \int X\left(\theta\right) \exp \left[ - \beta \left(U \left(\vec{r}, V; \theta \right) + pV \right) \right] \mathrm{d}\vec{r} \mathrm{d}V\]

where

\[Q\left(\theta\right) = \int \exp \left[ - \beta \left(U \left(\vec{r}, V; \theta \right) + pV \right) \right] \mathrm{d}\vec{r} \mathrm{d}V\]

is the isothermal-isobaric partion function, \(\theta\) are the force field parameters being used to estimate the observable, \(U\) the systems potential energy, \(\beta \equiv k_b T\), \(k_b\) the Boltzmann constant, \(T\) the temperature, \(p\) the pressure and \(V\) the volume.

The derivative of the ensemble average defined above with respect to a particular force field parameter of interest \(\theta\) is given by:

\[ \begin{align}\begin{aligned}\dfrac{ \mathrm{d} \left<X\right> }{ \mathrm{d} \theta_i } = \left< \dfrac{ \mathrm{d} X } { \mathrm{d} \theta_i } \right> - \beta \left[\\ \left< X \dfrac{ \mathrm{d} U } { \mathrm{d} \theta_i } \right> - \left< \dfrac{ \mathrm{d} U } { \mathrm{d} \theta_i } \right> \left< X \right>\\ \right]\end{aligned}\end{align} \]

Computing \(\mathrm{d} U / \mathrm{d} \theta_i\)

While future integrations with differentiable simulation engines such as timemachine will allow \(\mathrm{d} U / \mathrm{d} \theta_i\) to be computed directly from molecular simulation runs, currently most common simulation engines do not directly support computing this quantity.

Until such an integration is complete, the framework currently employs a central finite difference approach, whereby

\[\dfrac{\mathrm{d} U}{ \mathrm{d} \theta_i } \approx \dfrac { U \left( \theta_i + h \right) - U \left( \theta_i - h \right) }{ 2 h}\]

Although more expensive than computing either the forward or backwards derivative, the central difference method should give a more accurate estimate of the gradient at the minima, maxima and transition points. By default a value of \(h = \theta_i \times 10^{-4}\) is used. This has been found to yield finite differences which do not suffer from precision issues, while being sufficiently small so as to yield an accurate estimate.

In practice the derivatives obtained by re-evaluating the energies of each configuration in a trajectory generated by a molecular simulation (either after a simulation or after loading one from disk) at each of the perturbed parameters.

While there is an expense associated with extra evaluations of the potential energy function for each configuration, this is mitigated by only computing those terms which depend upon (or may depend upon) \(\theta_i\). As an example, when computing derivatives with respect to a bond length the electrostatic and van der Waal contributions are not computed. This significantly speeds up the computation of these derivatives.

The final derivatives are stored in ObservableArray objects for convenience and for easy propagation of gradients through workflows. See the observables documentation for more information.

References

1(1,2)

Lee-Ping Wang, Teresa Head-Gordon, Jay W Ponder, Pengyu Ren, John D Chodera, Peter K Eastman, Todd J Martinez, and Vijay S Pande. Systematic improvement of a classical molecular model of water. The Journal of Physical Chemistry B, 117(34):9956–9972, 2013.