A minor bug when encountering an empty graph

[bhcao]

https://github.com/instadeepai/mlip/blob/main/src/mlip/data/helpers/neighborhood.py#L107

It should be np.zeros((len(senders), 3)) instead of np.array([[0] * 3] * len(senders)). Although they are mostly equal, errors may occur in certain situations as follows.

When len(senders) == 0, which means the graph is empty, the shape of shifts will be (0,) instead of (0,3), which will cause shape mismatching when concat the empty graph with others (this will be called by jraph.batch_np in dynamically_batch). This situation does not occur during training because the empty graphs have been filtered, but it may still occur when using run_batched_inference.

[chrbrunk]

Hello @bhcao,

Thank you very much for pointing this out. We will take this as an opportunity to investigate our error handling for these "empty" graphs again. The change you are proposing also makes sense to us. We will adapt this line with our next release, however, I cannot provide an exact timeline yet, as we don't see this as something very urgent to patch by itself.

Thank you again and have a great day!

[Hellboy28D]

hi @bhcao
Good catch 👍

You’re right — np.array([[0] * 3] * len(senders)) does not behave correctly when len(senders) == 0, resulting in a shape of (0,) instead of (0, 3).

Using np.zeros((len(senders), 3)) is safer and guarantees the correct shape even for empty graphs. This would prevent shape mismatches during batching (e.g. in jraph.batch_np when using run_batched_inference).

I agree this change would make the code more robust. I can open a PR with this fix if that’s helpful.

[chrbrunk]

A new release has just been published on this repo and to PyPI, which should solve the mentioned issue. The corresponding PR is #38. I will therefore close this issue.

Thanks again to both of you for providing this very valuable feedback to us.

PS to @Hellboy28D: We currently, unfortunately, cannot easily accept external PRs on the repo, which is due to our automated release pipeline that releases changes from a private internal repository. However, we plan to extend this pipeline in the future to allow for external contributions more easily. Thanks for your understanding!

[bhcao]

Hi @chrbrunk, I have an additional question about PR. I have implemented some additional functions in my fork of MLIP repository, including:

1. Load the dataset while training, rather than loading it all into memory at once. This is important for training extremely large-scale data such as Omol25.
2. EquiformerV2, UMA (equivalent to eSEN+MoE) and LiTEN implemented using NNX.
3. Code for converting various data formats, including LMDB, extxyz, OpenQDC, to HDF5.
4. Partial support for charge, spin, multi-head, multi-task, etc.

Though, these functions have not undergone complete unit testing and have only been tested locally on my end. And we plan to implement the following functions:

1. Support for more thermostats and ensembles (including NPT, NVE, Nose Hoover, etc.).
2. Support for more models, such as SO3krates, Allegro, etc.
3. Support for larger systems during training and MD calculations through graph partitioning.

Would you be willing to incorporate some of them, and if so, how should I submit them.

[chrbrunk]

Hi @bhcao,

This all sounds very interesting. Unfortunately, we don't yet have an established process in place for taking external contributions in this repository. However, I just reached out to you by email to discuss possibilities how we could potentially collaborate on some of these topics. Thank you very much!