Problems with atom_types

[m-nader]

Hi,

I have been trying to train a model with the 0.10 version of Nequip + Allegro and I am encountering problems with my definition of atom types. This is the error I get:

Traceback (most recent call last):
  File "/opt/lammps/bin/nequip-train", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/opt/lammps/lib/python3.11/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/opt/lammps/lib/python3.11/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/opt/lammps/lib/python3.11/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/opt/lammps/lib/python3.11/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/opt/lammps/lib/python3.11/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
           ^^^^^^
  File "/opt/lammps/lib/python3.11/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
            ^^^^^^^^^^
  File "/opt/lammps/lib/python3.11/site-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
        ^^^^^^^^^^^^^^^^
  File "/opt/lammps/lib/python3.11/site-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/opt/lammps/lib/python3.11/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
                       ^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/lammps/lib/python3.11/site-packages/nequip/scripts/train.py", line 178, in main
    stats_dict = datamodule.get_statistics(dataset="train")
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/lammps/lib/python3.11/site-packages/nequip/data/datamodule/_base_datamodule.py", line 295, in get_statistics
    stats_dict = stats_manager.get_statistics(dloader)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/lammps/lib/python3.11/site-packages/nequip/data/stats_manager.py", line 221, in get_statistics
    self(data)
  File "/opt/pytorch/lib64/python3.11/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/lib64/python3.11/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/lammps/lib/python3.11/site-packages/nequip/data/stats_manager.py", line 131, in forward
    data[AtomicDataDict.ATOM_TYPE_KEY], type_idx
    ~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: 'atom_types'

I have attached my xyz (xyz.txt) and yaml (yaml.txt) files.

It used to work before, can you help me understand why it isn't working anymore?

Thanks!

[m-nader]

I found the problem. In the config file, "include_keys: ['atom_types']" must be included under data.

[cw-tan]

yes, our ASE parser will only parse ASE native keys -- atom_types is not an ASE native key and would have to be registered. That is why we usually have the ChemicalSpeciesToAtomTypeMapper (see https://nequip.readthedocs.io/en/latest/guide/configuration/data.html#data-transforms) -- to convert ASE native chemical symbols to atom_types for the model.