Alternative neighborlists for big systems with PBCs

[MorrowChem]

I found I was running out of memory trying to use AIMNet2 for systems above about 10,000 atoms and with PBCs. It turned out it was possible to increase accessible system sizes using the torch-cluster radius function instead of building a (possibly very big) connectivity matrix directly (i.e. conn_mat ), similar to how the non-PBC implementation works. Is there any interest in merging this into main, as I found it very useful for my application?

Here's the timing info for ASE MD simulations of water on an 80GB A100. Left to right shows no LR, cutoff_lr=10Å, and cutoff_lr=20Å respectively. Lines terminate when I run out of GPU memory. The best performance increase by doing this is gained when the molecular graph is sparse (i.e. not using or using 'shorter' long-range interactions). You can get out close to 100k atoms for just the short-range interaction case, which is close to an order-of-magnitude increase. It's also a bit faster this way for medium-sized systems of >5k atoms. With LR interactions, it's slower, so I made a default setting of if cutoff < 20 and size > 10000 then use new implementation.

[zubatyuk]

Thanks a lot for the contribution. I will carefully review it. Actually, I had a plan to completely avoid torch_cluster in favor of numba kernels.

[MorrowChem]

ah I see, guessing that's because torch_cluster has quite strict compatibility requirements (at least, I found it didn't play nice with other packages when installing). In any case, a fast way of finding the neighbours of atoms would be important for big systems, e.g. with a KDTree. I couldn't find any standard GPU-implementation of that available though

[zubatyuk]

Hi Joe,

Please check out https://github.com/isayevlab/aimnetcentral
It has numba implementation of PBC neighbor list. It is O(N^2), not linear scaling cell list version. But it is reasonably fast up to 50k, maybe 100k atoms.
For any optimizations, please do PR to https://github.com/isayevlab/aimnetcentral