Neighborlist api to handle batch

[abhijeetgangan]

Summary

• Neighborlist api with system idx
• Swaps vesin for portable torch neighborlist

Checklist

Before a pull request can be merged, the following items must be checked:

• Doc strings have been added in the Google docstring format.
• Run ruff on your code.
• Tests have been added for any new functionality or bug fixes.

[orionarcher]

Nice work! The API makes sense to me. It would be helpful to have a specific example though. How about creating a native neighborlist that handles this API and include it in the PR? As an intermediate solution, that could be as simple as pushing the loop inside the vesin_nl_ts to fake batching.

[abhijeetgangan]

Yeah, probably we would wrap the non batched neighborlist with this api. I think making a neighborls_ext.py to have external neighborlist be implemented be an option or make it like a model where we have a neighbors/vesin_nl.py etc. This way it will be easier to handle the deps and CI. The fallback would be the torch_nl neighborlists which is hardware portable.

We can have a DEFAULT_NL which uses vesin if it is installed or the fallback option.

Looking at the models which use this we could have the following:

edge_idx, unit_shifts, cartesian_shifts = neighborlist_fn(state or state info)

This will also make model integration simple.

[orionarcher]

We can have a DEFAULT_NL which uses vesin if it is installed or the fallback option.

Looking at the models which use this we could have the following:

edge_idx, unit_shifts, cartesian_shifts = neighborlist_fn(state or state info)

This will also make model integration simple.

Do you want to do these things before merging in? Maybe this PR could clean up neighbors.py a bit before merging in. A couple things come to mind:

• I'd prefer to not have two types of neighborlists with different APIs where one will fail and the other will work.
• I like the DEFAULT_NL idea but let's make sure that follows the batched API and get rid of/make private the vesin_nl_ts and/or wrap it in a for loop so it conforms with the API
• What is the linked cell algorithm? Is that just a batched API or is it something else entirely?

[abhijeetgangan]

Yes, I will add this to this PR. I think I would just wrap the neighborlist implementations even when they do not have explicit batch handling (just need to for loop here) as a generic function which returns the above.

Both the torch_nl versions api handle system_idx but the others do not. linked cell is just linear scaling neighborlist.

Performance wise, I am removing the standard_nl because it is too slow and doesn't handle batching. The linked cell is a good portable default without external dependencies. I also need to update vesin to use cuda with the recent update.

[orionarcher]

Note that #347 makes some changes to the neighborlist that cover some of what we discuss here.

[abhijeetgangan]

Yeah, will keep those in mind when I have the API changes