CORE-3 — Framework-agnostic `Configuration` type + units/conventions module

[aacostadiaz]

Depends on: INF-1 (#1550) · Blocks: DATA-1 (#1569), DATA-3 (#1571)

Prerequisite: the Phase 0 characterization suite is committed — goldens, fixtures and the single-source tolerance table under tests/golden/, characterization tests under tests/unit/. References of the form P0-N name that work.

Context: The raw parsed structure and the numeric conventions belong in mace_core, shared by both frameworks and every data backend: backends (DATA-1 (#1569)/DATA-2 (#1570)) yield Configurations — never graph dicts — and graph construction (DATA-3 (#1571)) consumes them, so this type is the boundary object of the whole data layer. Legacy already has the right shape in the wrong place: Configuration is a numpy-backed dataclass at mace/data/utils.py:92-102, KeySpecification maps convention names to file keys at mace/data/utils.py:28-46 (with the command-line-style converter update_keyspec_from_kwargs at :49-86), AtomicNumberTable lives at mace/tools/utils.py:94-108, and DefaultKeys at mace/tools/default_keys.py:6-27 — all tangled into the torch-importing mace package. The P0-7 parsing tests define the parsing behavior to preserve; the P0-6 test docstrings pin the sign conventions. This ticket reimplements the types and the XYZ-parsing pure functions in mace_core.data, and creates mace_core.units as the single source of truth for units and physics sign conventions. Legacy is read in place as the behavioural oracle (mace/data/utils.py, mace/tools/utils.py); reproduce its parsing semantics, not its structure.

Interface & constraints:

1. mace_core.data.Configuration — typed, numpy-backed dataclass (no torch, no jax), field-for-field the legacy shape at mace/data/utils.py:92-102:

@dataclass
class Configuration:
    atomic_numbers: np.ndarray          # [n_atoms] int
    positions: np.ndarray               # [n_atoms, 3] Angstrom
    properties: dict[str, Any]          # keyed by CONVENTION names ("energy", "forces", …;
                                        #  multi-level-of-theory names like "pbe_energy",
                                        #  "r2scan_forces" are legal keys)
    property_weights: dict[str, float]  # per-property weights
    cell: np.ndarray | None = None      # [3, 3] — the PHYSICAL cell as parsed
    pbc: tuple | None = None            # (3,)
    weight: float = 1.0                 # weight of config in loss
    config_type: str = "Default"
    head: str = "Default"

Property keys in properties are always convention names, never raw file keys — the key_spec mapping is resolved during parsing, so nothing downstream ever sees REF_energy-style file keys. Document this convention in the module docstring. cell is the cell as it appears in the file and nothing else: the neighbour list's artificial construction cell is DATA-3 (#1571)'s concern, and develop deliberately returns it only for fully aperiodic systems, returning the physical cell whenever any axis is periodic and patching only all-zero rows of non-periodic axes from the extended one (mace/data/neighborhood.py:32-60) — a Configuration never carries a synthetic box.

2. KeySpecification — the mapping from convention names to file keys, split into info_keys (per-config: energy, stress, virials, dipole, head, elec_temp, total_charge, polarizability, total_spin) and arrays_keys (per-atom: forces, charges, magmom, magforces), with a from-defaults constructor. Reimplement from the behavior at mace/data/utils.py:28-86. Two parts of that behaviour are easy to miss and both are required:

   • magmom and magforces are arrays keys on the default path, not a magnetic-model extra — they are in DefaultKeys and are resolved for every parse.
   • embedding_specs extends the key spec at runtime (mace/data/utils.py:74-85): each declared feature registers as an arrays key when per == "atom" and an info key when per == "graph", using spec["key"] or the feature name, and any other per value is a hard ValueError. This is the user-declared graph-level/atom-level input-feature channel (--embedding_specs, mace/tools/arg_parser.py:743), and CORE-1 (CORE-1 — Typed outputs + declarative observable specification #1555)'s derivative grammar depends on those features being nameable here.

3. AtomicNumberTable and DefaultKeys reimplemented from scratch in mace_core — never imported from mace (packages/ importing mace/ is banned; a test asserts it). AtomicNumberTable (legacy mace/tools/utils.py:94-108) holds an ordered zs sequence with index_to_z / z_to_index / __len__; note that sorting and de-duplication live in the factory, get_atomic_number_table_from_zs (:111-115), not in the class — reimplement both, and keep the split, because the class is also constructed directly from an already-ordered list. DefaultKeys — the default file keys the from-defaults KeySpecification resolves to (legacy mace/tools/default_keys.py:6-27, 13 members, exposed as <name>_key → value by DefaultKeys.keydict()):

Convention name	Default file key
energy	REF_energy
forces	REF_forces
stress	REF_stress
virials	REF_virials
dipole	dipole
polarizability	polarizability
head	head
charges	REF_charges
total_charge	total_charge
total_spin	total_spin
elec_temp	elec_temp
magmom	REF_magmom
magforces	REF_magforces

4. Parsing — pure functions read_configurations(path, key_spec) -> list[Configuration] porting the behavior pinned by the P0-7 tests (legacy oracle: load_from_xyz / config_from_atoms, mace/data/utils.py:239-358,172-217). Pinned semantics:

   • Property-key variants: default keys (table above) and custom keys (pbe_energy-style) both work; a property whose key is absent from the file is stored as None and its property_weights[name] is forced to 0.0 (:197-205).
   • The ASE 3.23 key rewrite. If the configured key for energy, forces or stress is literally "energy", "forces" or "stress", legacy logs a warning, rewrites the key spec to the corresponding REF_* name, and back-fills that key on every ase.Atoms from the calculator (atoms.get_potential_energy() / get_forces() / get_stress()), storing None on failure; the original key strings are restored on the key spec before returning, so the caller's object is not mutated (:256-289, :355-357). Reproduce all three legs — the rewrite, the calculator back-fill with its per-quantity failure handling, and the restore. Silently skipping the restore turns a shared key spec into a one-shot object.
   • Presence check and no_data_ok. After the rewrite, if none of the final energy, forces or dipole keys is present in any frame, legacy raises ValueError naming all three keys and the file — unless no_data_ok=True, which downgrades it to a warning and continues; a missing energy or missing forces alone is always a warning (:291-314). Keep the hard error as the default and keep the escape hatch explicit.
   • Weights: weight = info["config_weight"] (default 1.0) × config_type_weights[config_type] (default 1.0); per-property weights read from info["config_<name>_weight"] (default 1.0), for the union of info and arrays key names.
   • config_type: read from info["config_type"], default "Default".
   • Isolated-atom detection: a config with exactly 1 atom and config_type == "IsolatedAtom" yields an E0 for its element; marked-but-energyless isolated atoms warn and record 0.0; a keep-isolated-atoms flag controls whether they stay in the returned configurations (legacy mace/data/utils.py:320-346). The head info key is stamped on every frame, isolated atoms included, before the split.

5. Two more pure functions belong here, and have no other home in the plan. Both are numpy/pure-python and operate on Configuration lists, so putting them anywhere else re-splits the boundary object:

   • random_train_valid_split(items, valid_fraction, seed, work_dir, prefix) (mace/data/utils.py:108-146) — np.random.default_rng(seed) shuffle, train_size = min(size - int(valid_fraction * size), size - 1) so at least one validation item always exists, and the side effect that makes it reproducible in practice: when the validation set has 10 or more items the chosen indices are written to <prefix>_valid_indices_<seed>.txt in the work dir, otherwise they are logged inline. The index file is part of the observable behaviour, not a debug aid.
   • test_config_types(test_configs) (:220-236) — the config_type + "_" + head grouping that TRN-3 (TRN-3 — Multi-dataloader balancing and metrics/logging (error tables, optional wandb) #1577)'s per-config-type error tables are built on, with the legacy normalisation of a None head to "".

6. mace_core.units — unit constants aligned with ase.units (energies in eV, lengths in Å); the physics conventions as documented constants, e.g. STRESS_SIGN_CONVENTION, stating: forces = −∂E/∂positions; stress = (1/V) ∂E/∂strain; virials = −stress·V; and, for the magnetic family, magforces = −∂E/∂magmom (mace/modules/utils.py:262-299). The physics statements must match the P0-6 test docstrings verbatim — those docstrings pinned the virial/stress sign convention against finite differences and are the authority; do not re-derive or re-word them. This module is the single source for conventions: nothing else restates them.

7. Dependencies: ase is an explicit mace-core dependency (it is framework-agnostic, and the parsing contract above is defined in terms of ase.Atoms info/arrays and the calculator back-fill); declare it in packages/mace-core/pyproject.toml with a one-line rationale. No torch, no jax, no e3nn anywhere in mace_core.

Task:

1. mace_core.data: Configuration, KeySpecification (+ from-defaults constructor + the embedding_specs extension), reimplemented AtomicNumberTable (class + sorting factory) and DefaultKeys (13 members), and read_configurations(path, key_spec) implementing the pinned parsing semantics above.
2. mace_core.data: random_train_valid_split and test_config_types per constraint 5, including the validation-index file.
3. mace_core.units: constants aligned with ase.units; convention constants with the P0-6 physics statements verbatim, plus the magforces sign.
4. Port the P0-7 parsing test cases into packages/mace-core/tests/ (property-key variants, weights, config_type, isolated-atom detection, the ASE 3.23 rewrite + restore, the presence check and its no_data_ok downgrade) and add a test asserting mace_core imports neither mace nor torch/jax.
5. Document the ase-dependency decision and the property-key naming convention in the module docstrings.

Out of scope: neighbor lists and the construction-cell regimes (DATA-3 (#1571)), torch tensors / graph construction (DATA-3 (#1571)), HDF5/LMDB backends (DATA-2 (#1570)), the E0 least-squares solvers compute_average_E0s / estimate_e0s_from_foundation (DATA-1 (#1569) statistics + CFG-1 (#1574)'s E0Spec).

Acceptance criteria:

• The ported P0-7 parsing cases pass identically against mace_core.data (same properties, weights, config types, isolated-atom results as legacy on the same inputs).
• DefaultKeys has exactly the 13 members above, asserted against the table rather than spot-checked; magmom/magforces resolve on the default path with no magnetic flag set.
• An embedding_specs declaration with per: atom lands in atom_keys and one with per: graph in graph_keys, using the declared key when present; any other per value raises. The two halves carry the same two words a user writes in per:; info and arrays are ase's names for ase's two stores and stay inside the xyz backend.
• A file whose energy key is literally energy is rewritten to REF_energy, back-filled from the calculator, and the caller's KeySpecification still reads energy afterwards.
• A file with none of energy/forces/dipole raises naming all three keys and the path; the same file with no_data_ok=True warns and returns.
• AtomicNumberTable and DefaultKeys are reimplemented in mace_core; a test asserts mace_core does not import mace (nor torch/jax).
• random_train_valid_split reproduces legacy's split for a given seed, including the ≥10-item index file and its name.
• The property-key convention (convention names in properties, file keys only inside KeySpecification) is documented in the module docstring.
• mace_core.units convention statements match the P0-6 docstrings verbatim.
• ase is declared as an explicit mace-core dependency with its rationale.

Inventory gaps assigned here:

• Graph-level input keys (elec_temp, total_spin, total_charge) and the full 13-key default property set parse through the v1 key convention — including REF_magmom/REF_magforces, which extend the data contract every labelled XYZ depends on, and the embedding_specs runtime extension, which is the only way a user-declared input feature becomes nameable to CORE-1 (CORE-1 — Typed outputs + declarative observable specification #1555)'s derivative grammar.

Verify:

python -m pytest packages/mace-core/tests/test_data_configuration.py packages/mace-core/tests/test_units.py -v

Review focus: code + physics conventions — the sign-convention constants and their verbatim match to the P0-6 docstrings need physics expertise; parsing edge cases (missing keys → weight 0.0, the ASE key rewrite/restore, no_data_ok, isolated-atom handling) need a line-by-line diff against the legacy oracle.

---