Why does ASEDataset object appends noise to coordinates?

[ipcamit]

I am trying to construct a Nequip model manually on my own and compare it with model generated by Nequip. While my results do match Nequip's within 10 eV, I can't get exact numerical match (lets say, upto single precision) because of what I believe small random noise in input data.
Example, Given below is an XYZ file coordinates:

64
Properties=species:S:1:pos:R:3:forces:R:3 pbc="1 1 1" Energy="-367.78660021" 
Si 4.76119081246 3.57093928276 0.326120749998 0.002089 0.000255 -0.000188
Si 2.38061518754 3.57086671724 0.978862250003 0.000634 0.002208 -0.002732
Si 7.14176981246 5.95151828276 1.63135375 0.000091 -0.001699 -0.005444
...

Here is the output of when I read these coordinates:

tensor([[ 4.761190812460000,  3.570939282760000,  0.326120749998000],
        [ 2.380615187540000,  3.570866717240000,  0.978862250003000],
        [ 7.141769812460000,  5.951518282760000,  1.631353750000000],

Below is the output of ASEDataset class:

input_dict["pos"]
tensor([[ 4.761190891265869,  3.570939302444458,  0.326120764017105],
        [ 2.380615234375000,  3.570866823196411,  0.978862226009369],
        [ 7.141769886016846,  5.951518058776855,  1.631353735923767],

I generated the ASEDataset by passing the following config directly to dataset_from_config function

{'root': 'results/toluene',
 'run_name': 'example-run-toluene',
 'default_dtype': 'float64',
 'dataset_seed': 456,
 'r_max': 4.0,
 'dataset': 'ase',
 'dataset_file_name': './toy_dataset.xyz',
 'dataset_key_mapping': {'Energy': 'total_energy'},
 'chemical_symbols': ['Si'],
 'dataset_include_frames': range(0, 4),
 'dataset_extra_fixed_fields': {'r_max': 4.0}}

It seems like it reading the data in 32 bits and then use it as 64 bit. Is there a way overcome this behavior?

[Linux-cpp-lisp]

Hi @ipcamit ,

ase reads using Python floats and is always using double, as far as I know. What you are getting out of ASEDataset is the float32 castdown of the float64 read-in:

In [4]: torch.as_tensor(4.761190812460000, dtype=torch.float32).item()
Out[4]: 4.761190891265869

When you call dataset_from_config with 'default_dtype': 'float64',, that option is being ignored because it isn't handled by dataset_from_config. You need to call _set_global_options, see train.py for an example, which should resolve this discrepancy.

However, I think your inability to reproduce below this threshold of numerics (note the final scalings, and how they relate numerical noise to real units) is also probably largely due to GPU nondeterminism, which can have a meaningful effect (all scatter sums in the model are unavoidably nondeterministic, which can have somewhat noticeable numeric effects in larger systems). For practical purposes, none of this seems to be an issue in our experience--- if you are worried about reproduction between these two things you are running, I would just check it in float64.

[ipcamit]

Thank you for the clarification. Reg. GPU nondeterminism, I am using single-threaded CPU for this test.

Actually I also wanted to provide an FYI and ask for comments.
I am trying to port NEQUIP model to OpenKIM (https://openkim.org/), which would enable the archiving of the nequip model and reuse in LAMMPS, ASE etc. We have come up with our own flexible framework to that (https://github.com/ipcamit/colabfit-model-driver) which would be able to port nearly all ML models in distributed memory framework, where it can be used in LAMMPS etc out of the box. It does so by staged graph convolutions, where you provide separate graphs of an "unrolled" ghost atom kind of atomic configuration (instead of minimum image kind of configuration that is usually used). We would submit manuscript shortly explaining the details (by the end of next week).

To demonstrate it, I have trained a Nequip model using the example.yaml file (3 layers). Then modified Nequip code slightly for our framework, and assigned the trained weights directly. Barring the numerical noise I mentioned above, it does seem to work. We compute forces explicitly by calling backward method.

---

Could you suggest any pitfalls I should be aware of? Any comments to make this process easier? I can provide more details if needed.

---

biggest pro of this approach is that Nequip can work in distributed environment like allegro, and for a trained model, end-user would just need to add one line in LAMMPS script:

# Initialize KIM Model
kim init NEQUIP_KIM_ID metal

As OpenKIM API is part of LAMMPS core, user need not install any extra plugins. Users will also able to run same model on all supported MD softwares like ASE, DL_POLY, GULP etc. with support for OpenMM coming up in future.

[Linux-cpp-lisp]

I see @ipcamit , very nice! Glad to hear that you are working to integrate NequIP here!

I'd certainly be curious to read about your graph unrolling approach; would you post a link to the preprint when it is public? I'm also happy to discuss by email (listed in my GitHub profile here). I'm assuming you're also writing an integration for Allegro since that comes "free" and without all the graph-unrolling necessary?

You would still have to build LAMMPS against libtorch, right?

[ipcamit]

Nequip was chosen specifically for two reasons: 1) more well behaved MD simulations, 2) specific challenges of non-parallel to parallel conversion, therefore it would make any error in our assumptions and implementations more obvious. Using this I could get energies and forces to match < 10^-10 for simple enough configurations. But for large configs, forces seems to have more noise. I believe we could match forces more accurately by giving more weight to forces during training (but will test it later).

For allegro I would do the same in near future, but given that it is parallel, I expect it to be lot less involved.

We don't build LAMMPS against libtorch, but rather our own "model driver" against the libtorch. Now LAMMPS just call KIM_API for forces and energy (which now comes as part of core LAMMPS distribution), which in turn is executes the model driver we wrote with libtorch. This gives us much more flexibility as we can have several different drivers as per our convenience and LAMMPS and core infra remains the same. That way our model driver becomes in essence a universal model driver that can run any TorchScript model, provided it follows certain calling conventions.

I would definitely send you the link of the preprint. And thank you for the offer for further assistance.

[Linux-cpp-lisp]

But for large configs

How large is "large"? You may want to try the latest develop, which improves the numerics for large systems.

For allegro I would do the same in near future, but given that it is parallel, I expect it to be lot less involved.

👍

I would definitely send you the link of the preprint.

Thanks!

[ipcamit]

How large is "large"? You may want to try the latest develop, which improves the numerics for large systems.

I am testing against 100000 atoms. I believe the numerical difference that I do see in forces might be insignificant when compared to inherent limitation of the LAMMPS timestep integration. Will definitely try the develop branch. It would give us more confidence too.

[ipcamit]

@Linux-cpp-lisp Sorry for raising such an old thread, but I thought you might be interested in followup, We were able to parallelize and deploy nequip models. Currently a very modest model for Si is already up (trained on small GAP Si dataset). But soon others will follow! It can run on multiple GPUs and nodes without any issue and performance wise matches pair_nequip implementation within 90%.

the model, and its benchmark results (predicted lattice constants etc) can be accessed here:
https://openkim.org/id/TorchML_NequIP_GuptaTadmorMartiniani_2024_Si__MO_196181738937_000

We also provide easy to use guide to port existing trained models to KIM framework:
https://kim-torchml-port.readthedocs.io/en/latest/kim_nequip.html

Paper will be out "soon"! Thank you for your support.