- DFT-D3
Currently, within the code, we use TAD-DFTD3 for the VDW contributions. While this package works fine, it will recalculate neighbor lists internally during the computation. Computing this a second time obviously takes additional time which could further slow things down when in production. Additionally, I'm not quite sure how efficient this neighbor listing is (I assume just brute force) which will further slow things down (e.g., not being able to take advantage of the verlet neighbor listing). However, I'll note that one advantage of this package is that it is on condo-forge.
https://github.com/NVIDIA/nvalchemi-toolkit-ops/tree/main
The nvalchemi toolkit provides a DFTD3 implementation (which is essentially just the implementation from aimnet2, cleaned up a bit). This package takes as input the interacting pairs from the neighbor list. As such, it will avoid additional recompilation. I'll note that for periodic systems, it will require the shifts, which is not something our current neighbor list provides (periodicity and, in theory, different shaped boxes, are handled via a "space" class). It should be straight forward to have this class return the shifts as well.
I'll note one down side is that nvvalchemi is not on conda-forge at the current time (something we should look into).
I'll note, we can easily support both methods; lazy-imports and the delayed import class I added (which will give warnings and tell you where to get the package) should be useful in this context. Aside from the recomputing of neighbor lists, nvalchemi is billed as being "fast" but I have not yet tested that (documentation is poor and provides examples that do not actually work and are missing libraries/data).
- Neighborlists
I'll note that the neighbor lists in modelforge do not necessarily return the interacting pairs in numerical order (where we could just count up the number of repeats in the first dimension of the tensor to get the number of neighbors... e.g., it could be[ 0, 0, 0, 1, 1, 2, 2, 1, 0, 2]). This is because to calculate the "full" neighbor list, I found it was faster to just simply copy the half neighbor list (applying the appropriate sign change). I think it would make sense to simply sort based on the first index in the [2, npairs] tensor. This would be more general anyway.
Regarding neighbor lists, I think I need to shift around some of the logic: while for local interactions the "unique_pairs_only" is an important toggle, when it comes to VDW the routines will require the full neighbor list regardless; electrostatics (in current implementation, require half, and it has a function in there to remove non-unique pairs). We should just switch electrostatics to always calculate half, rather than full and then. making it half again; this was certainly an oversight in terms of efficiency.
I'll note, the toolkit does provide some of their own internal neighbor lists (specifically a cell list based one). However, I'll note that these operate as traditional neighbor lists, simply providing pairs, rather than also returning r_ij and d_ij as well, as our code expected (this was a specific design choice which has some advantages). However, the design of the current neighbor lists does allow us to relatively easily swap in different backends (e.g., previously setting up versions that use nnp_ops). This is something worth looking into; my previous testing (of other implementations) didn't really show much benefit for a cell-list based scheme as compared to a nsq^2-verlet until system size was really large (the overhead of dealing with cells canceled out the cost savings given the ability of current GPUs to handle a lot of simple calculations during rebuild).
- Electrostatics
The nvalchemi toolkit also has various electrostatic methods (including Ewald summation, which I'd prefer not to have to figure out how to implement in pytorch). This should again be straightforward to implement once other changes to the neighbor list are accomplished.
I'll make sub issues for the various parts.
Currently, within the code, we use TAD-DFTD3 for the VDW contributions. While this package works fine, it will recalculate neighbor lists internally during the computation. Computing this a second time obviously takes additional time which could further slow things down when in production. Additionally, I'm not quite sure how efficient this neighbor listing is (I assume just brute force) which will further slow things down (e.g., not being able to take advantage of the verlet neighbor listing). However, I'll note that one advantage of this package is that it is on condo-forge.
https://github.com/NVIDIA/nvalchemi-toolkit-ops/tree/main
The nvalchemi toolkit provides a DFTD3 implementation (which is essentially just the implementation from aimnet2, cleaned up a bit). This package takes as input the interacting pairs from the neighbor list. As such, it will avoid additional recompilation. I'll note that for periodic systems, it will require the shifts, which is not something our current neighbor list provides (periodicity and, in theory, different shaped boxes, are handled via a "space" class). It should be straight forward to have this class return the shifts as well.
I'll note one down side is that nvvalchemi is not on conda-forge at the current time (something we should look into).
I'll note, we can easily support both methods; lazy-imports and the delayed import class I added (which will give warnings and tell you where to get the package) should be useful in this context. Aside from the recomputing of neighbor lists, nvalchemi is billed as being "fast" but I have not yet tested that (documentation is poor and provides examples that do not actually work and are missing libraries/data).
I'll note that the neighbor lists in modelforge do not necessarily return the interacting pairs in numerical order (where we could just count up the number of repeats in the first dimension of the tensor to get the number of neighbors... e.g., it could be[ 0, 0, 0, 1, 1, 2, 2, 1, 0, 2]). This is because to calculate the "full" neighbor list, I found it was faster to just simply copy the half neighbor list (applying the appropriate sign change). I think it would make sense to simply sort based on the first index in the [2, npairs] tensor. This would be more general anyway.
Regarding neighbor lists, I think I need to shift around some of the logic: while for local interactions the "unique_pairs_only" is an important toggle, when it comes to VDW the routines will require the full neighbor list regardless; electrostatics (in current implementation, require half, and it has a function in there to remove non-unique pairs). We should just switch electrostatics to always calculate half, rather than full and then. making it half again; this was certainly an oversight in terms of efficiency.
I'll note, the toolkit does provide some of their own internal neighbor lists (specifically a cell list based one). However, I'll note that these operate as traditional neighbor lists, simply providing pairs, rather than also returning r_ij and d_ij as well, as our code expected (this was a specific design choice which has some advantages). However, the design of the current neighbor lists does allow us to relatively easily swap in different backends (e.g., previously setting up versions that use nnp_ops). This is something worth looking into; my previous testing (of other implementations) didn't really show much benefit for a cell-list based scheme as compared to a nsq^2-verlet until system size was really large (the overhead of dealing with cells canceled out the cost savings given the ability of current GPUs to handle a lot of simple calculations during rebuild).
The nvalchemi toolkit also has various electrostatic methods (including Ewald summation, which I'd prefer not to have to figure out how to implement in pytorch). This should again be straightforward to implement once other changes to the neighbor list are accomplished.
I'll make sub issues for the various parts.