Atom type different from element

[peastman]

It looks to me like NequIP assumes a 1:1 correspondence between elements and atom types. I want to train a model where that isn't true. For example, I might have multiple atom types corresponding to the same element in different chemical contexts. Is there a way to do that?

[Linux-cpp-lisp]

Hi @peastman ,

Actually we do not assume any such correspondence, and the entire NequIP type system is designed to make this possible. In general in nequip atom type is just an arbitrary per-atom index ranging from 0 to n_types - 1 stored in the atom_types field (AtomicDataDict.ATOM_TYPE_KEY in Python).

Inside of the network itself, this is the only relevant data key; the data loading/creation code (either AtomicDataset via TypeMapper in the Python code, or pair_nequip/pair_allegro in LAMMPS) is responsible for converting atomic numbers or LAMMPS type information into atom type indexes.

When setting up a YAML config, instead of providing chemical_symbols or chemical_symbol_to_type, you can specify type_names, where you provide arbitrary string names for the atom types in your system in order and implicitly set the number of types. For example,

type_names:
   - H
   - C1
   - C2
   - fancyatom

will create a model with num_types = 4 where type 0 is named H, type 1 is named C1, and so on. Then you just need to make sure that your dataset contains a per-atom field atom_types containing the appropriate type indexes for each atom. (You may need to add something like include_keys: ["atom_types"] in your YAML, I'm not sure.)

(For .xyz dataset files, in ASE this can be achieved by setting the species/atomic numbers to whatever you please and adding an atom_types array to the ase.Atoms object before writing out the .xyz: atoms.arrays["atom_types"] = np.array([0, 0, 1, ...]).)

Type setup in LAMMPS works exactly like it would before; you just use the type names you defined in the YAML when defining the mapping from LAMMPS types to nequip types.

[peastman]

Perfect, that's exactly what I was looking for. Thanks!

[Linux-cpp-lisp]

Great, just let us know if there are any issues or other questions.