Help integrate AlphaNet into chemtrain-deploy

[zmyybc]

First, I want to express my sincere respect for your work. I truly believe that MLFF (machine learning force fields) play a crucial role in simulating large-scale systems, and your contributions in this field are highly valuable.

I have developed an MLFF model ​​https://github.com/zmyybc/AlphaNet and I would greatly appreciate your guidance on integrating it into ​​chemtrain-deploy​​. Could you kindly provide some advice or best practices for deployment?

Looking forward to your insights.

[pl-fuchs]

Hi, thank you for your interest in using chemtrain-deploy.

The first step to integrate AlphaNet into chemtrain would involve making it compatible with JAX.
Tools like torch2jax can automatically port from PyTorch to JAX.
However, we have not yet tested such automatic porting of PyTorch models, and it will likely not work out of the box.
Thus, we recommend devising a JAX-based implementation of AlphaNet, e.g., similar to the JAX-based implementations of MACE and Allegro.

Once you have a JAX-compatible version of AlphaNet, the next step would be to interface the model with chemtrain and chemtrain-deploy. This step should be straightforward if the graph only requires information about neighbors, node species, and node displacements.
For example, we interfaced the PaiNN model through jraph.

[zmyybc]

Thanks! And does PBC condition considered in your code?

[pl-fuchs]

In JAX, M.D, different boundary conditions are realized by choosing the correct displacement function.

In LAMMPS, boundary conditions are realized through ghost atoms, e.g., as described here: https://docs.lammps.org/Developer_par_comm.html.
Therefore, the model does not distinguish between free space and a periodic space, and we always export the model with displacement_fn, _ = jax.space.free().

[zmyybc]

Appreciate it! I'm handling this now but may need to revisit with some questions later, so I'll keep this open.