[BUG] Periodic PET long-range output depends on neighbor-list row order

[MichelangeloDomina]

Summary

On main@2618bb09143672b63024a7f668e76190ef684869, grouping a periodic neighbor list by center atom changes PET's long-range output, even though the geometry and the neighbor-list rows are otherwise unchanged. The local PET features agree to float64 roundoff, which isolates the difference to the periodic long-range path.

The original, ungrouped order is produced naturally by the public vesin.metatomic.NeighborList calculator.
The likely cause is that PET uses center-sorted distances together with atom-pair identities in the original neighbor-list order. When those orders differ, a distance can be assigned to the wrong pair.

Expected behavior

Grouping valid neighbor-list rows by center atom while keeping every row's sample labels and values together should not change a prediction. Each distance passed to the periodic long-range calculator should remain aligned with its corresponding (first_atom, second_atom, cell_shift) row.

Actual behavior

With a small deterministic PET model, the attached reproducer prints:

Natural vesin center order: [1, 0, 2, 2]
Equivalent grouped center order: [0, 1, 2, 2]
Misaligned natural long-range distance/pair rows: 2 / 4, max assigned-distance error=1.17157287525 Angstrom
Local PET feature difference: max=8.32667268469e-17
P3M long-range difference: RMS=0.00186113016113, max=0.00466070211158, relative RMS=0.914408%
Ewald long-range difference: RMS=0.00186113645151, max=0.00466069618501, relative RMS=0.914409%
P3M end-to-end absolute energy difference: 2.25876356974e-06

It then fails the row-order-invariance assertion:

AssertionError: Tensor-likes are not close!

Mismatched elements: 16 / 24 (66.7%)
Greatest absolute difference: 0.004660702111581705 at index (0, 2)
Greatest relative difference: 0.08813623274876492 at index (1, 4)

The model is deliberately small and randomly initialized. These values demonstrate a broken ordering invariant, but they do not estimate the error in a trained model or a production simulation.
The end-to-end energy difference in this synthetic example is small so the direct distance-to-pair mismatch is the primary evidence of the bug.

Version

main@2618bb09143672b63024a7f668e76190ef684869

Steps to reproduce

1. Check out the affected revision and install it in a virtual environment:

   git clone https://github.com/metatensor/metatrain.git
   cd metatrain
   git checkout 2618bb09143672b63024a7f668e76190ef684869
   python -m venv .venv
   .venv/bin/python -m pip install --editable .
   .venv/bin/python -m pip install torch-pme==0.4.0 vesin==0.5.8

2. Save the attached f009_long_range_distance_order.py in the repository root.
3. Run:

   .venv/bin/python f009_long_range_distance_order.py

The script creates a periodic three-atom system. Vesin naturally returns the center order [1, 0, 2, 2]. Then it compares the order with an equivalent stable center sort, keeping every row's identity, cell shift, and displacement together. It checks local PET, the Ewald and P3M long-range paths, and an end-to-end P3M energy. The final row-order-invariance assertion fails on the affected revision.

f009_long_range_distance_order.py

Further information, files, and links

Why this can occur in normal use

The metatomic ASE calculator uses vesin for CPU neighbor lists. Vesin treats sorting as an explicit option, and the public vesin.metatomic.NeighborList path used by the reproducer does not request it.
Neighbor-list rows are also not documented as requiring center grouping. By contrast,
metatrain's training/evaluation helper
does request sorting, so that path can mask the problem.

Technical analysis

At the affected revision:

• compute_batch_tensors creates a center-sorted layout and derives distances in that
  order.
• PET.forward sends those distances to the long-range
  path,
  which flattens them in the sorted
  order.
• LongRangeFeaturizer reads pair identities from the original neighbor-list
  order.

Torch-PME uses the supplied neighbor_distances with the supplied neighbor_indices
in its real-space calculation.
It does not reconstruct those distances from the atomic positions, so the reciprocal
Ewald or P3M calculation cannot correct a wrong distance-to-pair association.

Scope and limitations

• This reproducer covers periodic PET with long-range features enabled and a fixed cutoff.
• Default short-range PET is not affected because it does not call this long-range path.
• The non-periodic direct calculator constructs its own all-pair indices and distances.
• No trained checkpoint or representative production trajectory was used, so the practical magnitude in such a model has not been measured here.

Tested on macOS arm64 with Python 3.14.3, PyTorch 2.10.0,
metatensor-torch==0.10.1, metatomic-torch==0.1.15,
metatomic-ase==0.1.1, vesin==0.5.8, and torch-pme==0.4.0.

[sirmarcel]

We should really remove the long range component until the release of the "actual" long-range model. But generally, do we know where this bug comes from?