Open jeevster opened 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.
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.