Fix NequIP atom type mapping for deployed models

[susumu-fujii]

Fix NequIP atom type mapping for deployed models

Summary

This PR fixes the atom type mapping used by the NequIP solver during sampling.

Currently, NequipSolver.Input.from_directory() reads the atom type order from:

yaml_dic["chemical_symbols"]

in input.yaml.

However, with NequIP 0.6.2, the atom type order used internally by the trained and deployed model is stored in type_names and may differ from the order of chemical_symbols.

For example, the training input may contain:

chemical_symbols:
  - Cu
  - V
  - Sn

while the trained config.yaml contains:

type_names:
  - V
  - Cu
  - Sn

type_to_chemical_symbol:
  0: V
  1: Cu
  2: Sn

The deployed model metadata also contains:

type_names: V Cu Sn

Therefore, the deployed model expects:

V  -> type 0
Cu -> type 1
Sn -> type 2

but the current abICS implementation constructs the type indices from chemical_symbols and instead uses:

Cu -> type 0
V  -> type 1
Sn -> type 2

As a result, Cu and V are evaluated using the wrong atom types, which can produce incorrect energies during sampling.

Fix

Use type_names and r_max stored in the metadata of deployed.pth instead of reconstructing the atom type mapping from input.yaml["chemical_symbols"].

This ensures that the atom type indices used by abICS are consistent with the mapping actually used by the deployed NequIP model.

Test

A regression test was added for the case where the order of chemical_symbols differs from the deployed model's type_names.

The test uses:

chemical_symbols:
  - Cu
  - V
  - Sn

and mocked deployed-model metadata:

type_names: V Cu Sn
r_max: 8.0

and verifies that NequipSolver.Input.from_directory() uses:

["V", "Cu", "Sn"]

rather than the order in chemical_symbols.

The test fails with the previous implementation and passes with this fix.

Additional verification

The fix was also checked using an actual NequIP 0.6.2 deployed model.

After the change, the energy calculated through the abICS NequipSolver agrees with the energy obtained directly using NequIPCalculator.from_deployed_model() for the same structure.