openmm / NNPOps

High-performance operations for neural network potentials
Other
81 stars 17 forks source link

Allow `getNeighborPairs` cutoff to be larger than half the box width? #119

Open orionarcher opened 12 hours ago

orionarcher commented 12 hours ago

I am exploring small crystalline systems with both MLIPs and DFT. These systems can be quite small, less than 0.8 nm. I'd like to use OpenMM to evolve these systems with MLIPs, but because the box size must be twice the neighbor cutoff, that's not currently possible. I can't decrease the cutoff because it would be unfaithful to the MLIP and I can't increase the box size because then the systems would be too large for DFT.

My (perhaps naive) understanding is that the box size limitation is implementation choice, rather than a theoretical limit. Can the getNeighborPairs function be modified to support smaller boxes? Or are there other torch-scriptable implementations I could use that wouldn't share this limitation?

peastman commented 11 hours ago

The reason for the limitation is that it allows you to use a single periodic image convention. When looking at the neighbors of an atom, you only need to consider one periodic copy of each other atom. It's guaranteed that only one copy will be within the cutoff distance, so any copy other than the nearest one is guaranteed to be outside the cutoff.

If you allow the cutoff to be more than half the box size, it's possible for multiple copies of a single atom to be within the cutoff distance. That makes the neighbor search much more expensive, and also requires lots of other parts of the model to become a lot more complicated.

orionarcher commented 11 hours ago

Got it, thanks for the response. The increase in neighbor list cost makes sense. I imagine different implementations could be called depending on box size, e.g.

if box_size > 2 * cutoff: 
   # run fast implementation
else: 
   # run slow implementation

I don't quite understand why other parts of the model must become more complicated. Some models, like the MACE implementation in the ASE calculator, already support small box sizes (call to matscipy in mace, neighbor list implementation in matscipy).

Would such a small-box neighbor list be in scope for NNPops? I could take a swing at it, if so (though in PyTorch, not with a custom kernel). Circumventing the box size limitation would be valuable for materials science use cases.