Improved NL computation - Githubissues

Exscientia / physicsml

A package for all physics based/related models

MIT License

40 stars 1 forks source link

Improved NL computation #22

Closed wardhaddadin1 closed 4 months ago

wardhaddadin1 commented 4 months ago

Here is a little benchmark:

On a pos = 5 * torch.randn(1000, 3), I get:

Using primitive_neighbor_list_torch from physicsml => time ~0.39s (mean of 10 iters)
Using a simple n^2 implementation, time ~ 0.03s (mean of 10)
Using the from torch_nl import compute_neighborlist, time ~ 0.07s (mean of 10)
Using the from matscipy.neighbours import neighbour_list, time ~ 0.05s (mean of 10)

If we include PBCs, it's:

Using primitive_neighbor_list_torch from physicsml => time ~0.4s (mean of 10 iters)
Using the from torch_nl import compute_neighborlist, time ~ 0.09s (mean of 10)
Using the from matscipy.neighbours import neighbour_list, time ~ 0.09s (mean of 10)

To fix this:

Change to the torch-nl implementation when we have pbcs. (From here: https://github.com/felixmusil/torch_nl)
Use the n^2 implementation for when there are no pbcs.

peastman commented 4 months ago

That looks great. Thanks for the quick fix!