ASE example

[schiotz]

Hi,

I followed the tutorial yesterday at the MRS meeting, and was playing a bit with the colab notebook. I wanted to try to use the ASE interface instead of LAMMPS. Do you have an example of how to do that? Basically, how do I load a model created with nequip-train nequip/configs/example.yaml and create an ASE calculator. I found the nequip.ase.NequIPCalculator, but I do not know how to create an instance of it from the output of nequip-train. I suspect it is easy :-)

Best regards
Jakob

[Linux-cpp-lisp]

Hi Jakob,

Thanks for coming to our tutorial and for reaching out!

Once you have deployed your model using nequip-deploy (just as in the colab notebook we did for LAMMPS, the process is the same), you can later load that model for use in ASE in Python with:

from nequip.ase import NequIPCalculator
atoms.calc = NequIPCalculator.from_deployed_model(
    model_path="path/to/deployed.pth",
    species_to_type_name = {
        "C": "NequIPTypeNameForCarbon",
        "H": "NequIPTypeNameForHydrogen",
    }
)

The keys in species_to_type_name are chemical symbols and the values are the names of the corresponding types in your NequIP model, which you can find in your original YAML config by looking at either chemical_symbol_to_type (in which case they are also chemical symbols) or at type_names, whichever was provided. We ask the user to provide this explicitly to prevent any unpleasant issues with silently but incorrectly mapping species, which would only manifest in strange predicted forces.

Two very important other parameters are energy_units_to_eV and length_units_to_A. ASE always works in eV and Angstrom; by providing these values you can tell NequIP how to convert whatever energy/length units your training data is in to eV/A. These default to 1.0, which is correct if your data is already in ASE units (which is true if you use the ASE dataset loader, of course).

You can also pass a device option (a PyTorch device string, values like "cuda" or "cpu") to choose where to run the model. And finally, if you want to silence any resulting warnings about "global state", you can set set_global_optionsto True— from_deployed_model has to set some global options in PyTorch but we don't want to set global state silently by default, but you can explicitly silence those warnings by setting that option.

[schiotz]

Thank you very much! I’ll try it out.

[schiotz]

And I can confirm that it works! Thank you very much indeed.

[Linux-cpp-lisp]

Great, glad to hear!