[Feature Request] Support parameter optimization for dynamic properties based on trajectory

[ice-hanbin]

Summary

Expand the function of parameter optimization in DMFF for dynamic properties, such as diffusion coefficient, viscosity, etc. This feature requests,

1. An efficient algorithm for evaluating the gradient.
2. A robust optimization strategy.
3. A general interface for computing all kinds of dynamic properties.

Motivation

In v0.2.0, DMFF already support parameter optimization for some thermodynamic properties, such as free energy, RDF, etc. However, for more complicated dynamic properties, DMFF can't get the gradient from the trajectory. Expanding the function of parameter optimization for dynamic properties helps DMFF to be a more powerful engine to accelerate forcefield development in real applications.

Suggested Solutions

Many dynamic properties can be generally evaluated using time-correlation function, which utilizes the information of every state in the trajectory. If we simply use jax.grad on the whole trajectory, the memory may be a huge problem. Inspired by Neural ODE, we can evaluate the gradient in backpropagation process , which utilizes the time reversibility of NVE trajectory. As we get the gradient, we can use modern AI optimization algorithm to fit the dynamic properties.

Any suggestions and comments on this solution are welcomed !

Further Information, Files, and Links

No response

[Chenghao-Wu]

Great project.

Here is some information may be relevant:
https://arxiv.org/abs/2003.00868
https://frankschae.github.io/post/shadowing/

basically, the automatic differentiation through a trajectory generated by a stiff differential equation (MD simulations can be considered as this kind), chaos can not be avoided which causing the gradient becoming useless. It gives huge difficulty to optimize parameters based on MD trajectory.

Another thing can be done is probably to add stochastic terms (langevin thermostat) in the MD system, some tests can be found in this paper https://arxiv.org/abs/2301.03480.

[ice-hanbin]

Thanks for your sharing !!!
I am currently dealing with this problem, and I will pay attention to these works. Hope I can discuss with you when I have some thoughts on this.