KeyError: 'edge_index' when running Allegro model on previous NequIP dataset

[wangxuan1586]

Dear NequIP-allegro developers:

I am trying to train an Allegro model using the same training dataset that previously worked with NequIP. However, I encounter a KeyError: 'edge_index' during the validation step.

The traceback is:
File "./.conda/envs/allegro_env/lib/python3.10/site-packages/nequip/data/AtomicDataDict.py", line 267, in num_edges
return data[_keys.EDGE_INDEX_KEY].size(1)
KeyError: 'edge_index'

The environments are as list as follows:

Python: 3.10
PyTorch: 2.6.0+cu118
Torchvision: 0.15.1+cu118
NequIP: 0.17.1
nequip-allegro: 0.8.2
LAMMPS: CPU only
OS: Linux

My questions are as follows:
Why does Allegro require edge_index even when using the same dataset that worked for NequIP?
Is there a recommended way to precompute edge_index for Allegro, similar to NequIP's previous pipeline?
Are there any changes in Allegro 0.8.x compared to NequIP that require special data formatting?

Thank you for your help!

[cw-tan]

If you don't report specific details of what you have done, we have to guess to give you advice.

I am guessing that you did not include the neighborlist transform: https://nequip.readthedocs.io/en/latest/guide/configuration/data.html#data-transforms

e.g.

allegro/configs/tutorial.yaml

Line 15 in 6d05192

- _target_: nequip.data.transforms.NeighborListTransform

Please go through the docs and tutorial config https://github.com/mir-group/nequip/blob/main/configs/tutorial.yaml

[wangxuan1586]

Thank you for your help. I followed your suggestion and updated my config.

The issue was that my original config did not explicitly include the required data transforms. I have now added the following transforms to the ASEDataModule:

transforms:
  - _target_: nequip.data.transforms.ChemicalSpeciesToAtomTypeMapper
    model_type_names: ${model_type_names}

  - _target_: nequip.data.transforms.NeighborListTransform
    r_max: ${cutoff_radius}

I also enabled the statistics manager because my model uses training-data statistics such as per_atom_energy_mean, forces_rms, and per_type_num_neighbors_mean:

stats_manager:
  _target_: nequip.data.CommonDataStatisticsManager
  type_names: ${model_type_names}

After these changes, the minimal example with small_20frames.xyz runs successfully for 5 epochs. The previous edge_index error is resolved.

I also changed compile_mode: compile to compile_mode: eager to avoid a local Torch/Inductor/GCC compilation issue on my cluster.

So I think the original problem was due to missing explicit data transforms and statistics configuration in the newer NequIP/Allegro config style.

Thank you again for pointing me in the right direction.

I also upload the minimal config if that would still be useful.
tutorial_minimal.yaml