search

apax-hub / apax

A flexible and performant framework for training machine learning potentials.
MIT License
12 stars 1 forks source link

Piecewise linear MD temperature schedule #270

Open Chronum94 opened 2 months ago

Chronum94 commented 2 months ago

This is a feature request for piecewise linear temperature schedule for the JaxMD calculator.

This is a common MD temperature schedule that I don't think is possible right now with the JaxMD calculator in apax.

Could in principle be done by adding a temperature_schedule config keyword (and corresponding class) with the current one being a specialization liek ConstantTemperatureSchedule, and new one being PiecewiseLinearTemperatureSchedule (piecewise_linear keyword?)

For the PLTS case, one could let the temperatures be a list of floats and have a list of floats for the endpoint temperatures of each linearly interpolated segments

So something like:

temperature_schedule: piecewise_linear
    temperatures: [100, 3000, 2500, 300]
    times_in_ps: [0.0, 100.0, 400.0, 500.0]

Would start at 100K go to 3k K over 100 ps, etc etc.

M-R-Schaefer commented 2 months ago

Hi Chronum,

thanks for reaching out. This should be relatively easy to implement. It is also something that could be useful for our LotF applications since sampling with jaxmd is quite a bit more efficient. I'll see if I can get it done some time this week.

Best