Feature Request: Time-dependent temperature

[ellamichelleking]

Having the energy be time-dependent is really useful, but it seems like it may not be able to interfacing with temperature changing with time. It would be nice to be able to have temperature schedules be time-dependent in addition to the energy.

[sschoenholz]

This should already work. If you pass a function to the temperature in any of the simulation environments it will naturally behave correctly.

For example see this colab: https://colab.sandbox.google.com/github/google/jax-md/blob/master/notebooks/nvt_simulation.ipynb and check the fourth cell where it defines a function T_schedule.

Let me know if you run into any issues with this or if it's not behaving as it should!