kyonofx / MDsim

[TMLR 2023] Training and simulating MD with ML force fields
MIT License
101 stars 11 forks source link

Simulate multiple initial conditions simultaneously? #6

Open jeevster opened 11 months ago

jeevster commented 11 months ago

Is there a way in the codebase to simulate MD with the trained potentials on multiple initial conditions in parallel (e.g using the existing Simulator class from mdsim/md/ase_utils)? Would be great if this could be done in a way that takes advantage of Pytorch vectorization.

kyonofx commented 11 months ago

Thanks for your question!

The current code does not support batching. The simulation code is currently under ASE, which does not support batched integration/thermostat control.

One can implement an integrator/thermostat, and then the code can be made parallelizable by keeping separate atom states and thermostats. I feel this will be nontrivial but not too much work. Unfortunately, I currently don't have the bandwidth to work on this, but it is certainly something very useful to have. I think the torch thermostat implemented here https://github.com/torchmd/mdgrad is a good place to start.