Repository navigation
Add the configuration type, the key convention and the units module - #1738
aacostadiaz wants to merge 5 commits into
Conversation
| def _extract_isolated_atom_energies( | ||
| atoms_list: Sequence[Atoms], energy_key: str | ||
| ) -> dict[int, float]: | ||
| energies: dict[int, float] = {} | ||
| for index, atoms in enumerate(atoms_list): | ||
| if not _is_isolated_atom(atoms): | ||
| continue | ||
| atomic_number = int(atoms.get_atomic_numbers()[0]) | ||
| energy = atoms.info.get(energy_key) | ||
| if energy is None: | ||
| logger.warning( | ||
| "Structure %d is marked as an isolated atom but carries no " | ||
| "energy under %r. Recording zero for element %d.", | ||
| index, | ||
| energy_key, | ||
| atomic_number, | ||
| ) | ||
| energies[atomic_number] = 0.0 | ||
| else: | ||
| energies[atomic_number] = float(energy) | ||
| return energies |
There was a problem hiding this comment.
This function keeps the last occurence of an isolated atom per element as its E0. For elements that have differrent spin states for isolated atoms, eg. oxygen, the selection of the spin state of the isolated atoms here depends solely on the (random) order of structures in the dataset. So either, we might deprecate getting E0s from a dataset, or we should take the lowest energy isolated atom from the dataset (or another well-motivated, non random choice maybe ? )
There was a problem hiding this comment.
But I guess if it requires the user to manually label the strcutures anyways, then we should check that they didnt label multiple ioslated atoms as ISOLATED_ATOM_CONFIG_TYPE
There was a problem hiding this comment.
I agree that it should only be done if explicitly labeled, and the code should fail if there are multiple configs labelled, or maybe just if they conflict?
There was a problem hiding this comment.
Fixed. It now raises an error. Same energies are fine.
| """The thirteen convention names, in declaration order.""" | ||
| return tuple(member.convention_name for member in cls) | ||
|
|
||
| @classmethod |
There was a problem hiding this comment.
This should raise if storage is not in {"atom" , "graph"} . Type checker should catch it though.
>>> DefaultKeys.names_stored_per("atoms") # with a typo atoms
frozenset()
There was a problem hiding this comment.
Fixed. It now raises for anything that is not "graph" or "atom".
| STRESS_SIGN_CONVENTION = ( | ||
| "stress = (1/V) dE/d(strain), in eV/Ang^3, with V = |det(cell)|." | ||
| ) | ||
|
|
There was a problem hiding this comment.
i guess some other conventions that we use is
- cell vectors are rows.
- stress should be 3x3 (?) not voigt layout as in ase?
There was a problem hiding this comment.
The shapes should also be checked, im sure ppl will try to train on (9,) stress tensors.
There was a problem hiding this comment.
Do you want to auto-convert (9,) to (3,3), implicitly assuming it came from an old extxyz format?
There was a problem hiding this comment.
no, from my view it should be rejected, if we clearly state/ communicate what format a dataset should be in.
There was a problem hiding this comment.
Added. Cell vectors are rows and stress is a 3x3 matrix. The shape is now checked 3x3 is kept, Voigt (6,) is converted anything else raises.
| if isinstance(spec, EmbeddingFeatureSpec) | ||
| else EmbeddingFeatureSpec.from_mapping(name, spec) | ||
| ) | ||
| target = self.atom_keys if feature.per == "atom" else self.graph_keys |
There was a problem hiding this comment.
User can declare a new Embedding spec with {"charges": {"per": "graph"}} puts charges in both graph_keys and atom_keys. Configuration.properties checks for membership in graph_keys.items() first. so charges doesnt live on atoms anymore.
There was a problem hiding this comment.
Fixed. A feature with the same name as an existing property now raises an error.
| f"Add per: atom for a per-atom array or per: graph for one " | ||
| f"value per structure." | ||
| ) from None | ||
| if per not in ("atom", "graph"): |
There was a problem hiding this comment.
This validation only happens if from_mapping is used as constructor. Building with bond_feature = EmbeddingFeatureSpec(per="bond") isnt validated, but should probably raise. In add_emebdding_features() L:175 bond feature is then added to the graph_keys, because it only checks is not per atom.
There was a problem hiding this comment.
Fixed. The check now runs for every spec not only for from_mapping.
| key_spec: KeySpecification, | ||
| *, | ||
| config_type_weights: Mapping[str, float] | None = None, | ||
| head_name: str = DEFAULT_HEAD, |
There was a problem hiding this comment.
head name is defined here, and also in per_graph properties. Here, the property is ignored for defining´ config.head, but lives on in properies.
Also head passed through the arguments, isnt written into properties.
There should only be one place where head is defined. Either in properties, or in Configuration, not both
There was a problem hiding this comment.
Fixed. The head is now only in Configuration.head
| @@ -0,0 +1,884 @@ | |||
| """Parsing labelled structure files into configurations. | |||
There was a problem hiding this comment.
Duplicates and tests that can't fail
-
Storage tests that re-check their own definition (test_data_configuration.py):
- test_the_two_halves_are_read_off_the_key_table (:133) repeats names_stored_per, which is how the two sets are built in the first place. It cannot fail.
- test_the_magnetic_arrays_resolve_with_no_magnetic_flag_set (:169) compares atom_keys against ATOM_CONVENTION_NAMES, which is also built from the table. It overlaps :118.
- test_the_graph_level_inputs_are_per_structure_keys (:177) pins four names by hand.
The result is that where each property lives ("graph" or "atom") is never checked against a written-out table, unlike the file keys, which DEFAULT_KEY_TABLE pins for exactly that reason. The only thing that would catch, say, virials moving to "atom" is the round-trip weight check at :378, and only indirectly. Suggestion: write out DEFAULT_KEY_TABLE as name → (file_key, storage) and drop :133, :169 and :177.
-
The sign conventions are written out a third time in test_units.py:67, which restates all three strings. tests/architecture/test_convention_statements.py already checks them word for word against the finite-difference suite. units.py itself says that a restated copy is a second copy that can drift, and this test is one: rewording now means editing three places. test_the_stress_is_the_one_quantity_that_is_not_negated (:81) only does substring checks, which the exact-match test already covers.
-
Weight composition is tested twice: test_the_structure_weight_is_the_product_of_both_weights (:312) is a subset of test_weights_compose_through_a_real_file (:449). It's cheap, but redundant.
-
test_a_declared_feature_and_a_default_key_use_the_same_two_words (:154): its first assert holds by construction, and its second repeats :220.
-
test_core_purity.py starts the subprocess probe twice. A module-scoped fixture would halve its runtime.
-
Misleading docstring and file layout:
- test_the_constants_are_ase_s_and_not_a_second_copy claims "a constant added above without a line here fails". It doesn't: both tests loop over PINNED, so a new constant in units.py goes unchecked. To make that true, assert that the float constants in units.all match PINNED's keys.
- The is_labelled and _RESERVED_KEYS tests (:841–884) sit under the "element table" heading.
There was a problem hiding this comment.
Done, all six points.
| pbc: tuple[bool, bool, bool] | None = None | ||
| weight: float = 1.0 | ||
| config_type: str = DEFAULT_CONFIG_TYPE | ||
| head: str = DEFAULT_HEAD |
There was a problem hiding this comment.
Duplicate storage of head in Configuration.properties["head"] and Configuration.head
mace_core.data carries the boundary object of the whole data layer: a format backend yields Configurations and graph construction consumes them, so the type has to exist before either side can be written. It is numpy and plain Python, no tensors and no neighbour list. Property keys are convention names everywhere downstream. The mapping from a file's own spelling to a convention name lives in KeySpecification, is resolved at the parse, and goes no further, so nothing after the parser has to be told which key spec produced it. The cell is the physical cell as parsed: an artificial box for a neighbour search belongs where the search happens, and a configuration carrying a synthetic one would report a stress divided by an invented volume with nothing in the object to say so. AtomicNumberTable and DefaultKeys are reimplemented rather than imported. The element table keeps the legacy split between the class, which adopts the order it is given because a checkpoint's order is not negotiable, and the factory, which sorts and de-duplicates for the dataset path. DefaultKeys keeps all thirteen spellings: they are the on-disk data contract, and REF_magmom and REF_magforces resolve on the default path with no magnetic flag set. mace_core.units is the single statement of each sign convention. The three derivative statements are reproduced character for character from the characterization suite that pinned them against finite differences, and tests/architecture asserts the two texts stay identical, so a convention cannot be re-worded without the re-derivation being visible. The unit factors are read from ase rather than written out, with the current values pinned so a CODATA revision is a reviewed change rather than a silent one. ase becomes an explicit mace-core dependency. The extended-XYZ contract is defined in terms of it: which of info/arrays a key lands in, which keys are reserved for calculator results, how a cell and its boundary flags are spelled. It pulls in no framework, so the purity rule is untouched. Three deliberate departures from the legacy parser, each asserted by a test: The reserved-key rewrite works on a copy of the key specification instead of mutating the caller's and restoring it at the end. A missed restore turns a specification shared between two datasets into a one-shot object, and a copy cannot miss it. An isolated atom's reference energy is read through the rewritten key. Legacy captured the energy key before the rewrite and then looked the pre-rewrite spelling up, so asking for the reserved key recorded every reference energy as zero. This is the only divergence that changes a number. An override naming no known property raises instead of being dropped. A dropped override leaves the property unparsed with nothing anywhere saying why the labels never arrived. The thirteen names the CLI exposes are exactly the ones accepted, so nothing that works today starts failing. The grouping function is named group_by_config_type. Under its legacy name pytest collects it as a test wherever it is imported. It no longer writes an empty head back onto the configuration either: a function that reports should not edit what it is reporting on. read_configurations returns a typed result carrying the configurations and the isolated-atom energies, rather than a bare list or an untyped tuple.
Twenty-seven rows now carry a second pin: the thirteen default property keys, the eleven property-key flags that resolve them, and the embedding_specs runtime extension. The legacy pin stays first and the v1 one follows it, so the row says both that the frozen tree is protected and that the capability exists on the other side. Which of the two a reader is looking at is stated once in the pin vocabulary rather than tagged onto every cell. The gate had to learn the layout first, or the exercise would have added exactly the pin these rules exist to reject. check_pins resolved only spans under tests/ and skipped everything else, so a pin naming a v1 test was never looked for on disk: a renamed or never-written test would have read as coverage. Both halves now apply to packages/<distribution>/tests/ as well, the path and the node id. The specificity floor moves with it. Depth is counted from the suite root rather than from the repository root, because a package's whole suite sits three parts deep and would otherwise pass a rule written to reject precisely that: it is the packages/ analogue of tests/unit, not of a per-family directory. The four package suites are named in the coarse list too, which keeps that list redundant with the depth rule instead of load-bearing. A v1 pin may not open a cell, and the checker says so. This file records what the frozen tree does, and a test in the rewrite cannot report whether the rewrite still matches it; a cell that led with one would turn a claim the oracle checks into a claim the rewrite makes about itself.
Four things the CORE-1 review found, in the same shapes, checked for here before opening a pull request. **The thirteen names lived in two modules.** `DefaultKeys` held them, and `data/keys.py` wrote them out again as two frozensets that also carried the half each belongs to. A test bound the membership, so they could not drift in silence, but the split was still a second statement of the same table. The storage side moves onto the enum, as `"graph"` or `"atom"`, and the two frozensets are read off it. **One distinction had two names.** A user writes `per: atom` on an embedding feature; the key specification called the same halves `info_keys` and `arrays_keys`, and a ternary translated between them. `info` and `arrays` are ase's words for ase's two stores, and `KeySpecification` is format-neutral, so it now says `graph_keys` and `atom_keys` and ase's spelling survives only in the xyz backend, where ase is what is being touched. **A zero weight meant two things.** `property_weights` carried the per-property loss weight and also marked an absent label by zeroing it, and the docstring said the zero was what marked absence. It is not: a file may write `config_forces_weight=0.0` for a structure whose forces are present, and the two are the same number. The zeroing stays, as a safety net for a consumer that reads only the weight, but absence now has its own answer in `Configuration.is_labelled`, and the test pins that the weight cannot tell the two apart while that can. **`units.py` claimed to be the single source for the sign conventions.** It is the single source for the *prose*, which cannot be derived from a number. The machine-readable sign belongs with the declaration of the quantity, and the observables work puts it there, so this module says what it owns and where the other half is rather than claiming both. Also: the reserved-key rewrite spelled `REF_energy` a second time instead of reading it off the key table, so a rewrite could have gone on landing where the parser no longer looks.
A storage other than "graph" or "atom" now raises in names_stored_per, and an EmbeddingFeatureSpec validates per on construction, so a spec built directly is held to the same rule as one built from a mapping. An embedding feature named like a property the specification already resolves raises: in the other half it sat in both, and the parser kept whichever it read last. The head is no longer a property. It lives on Configuration.head only, and the parse no longer stamps it into the structure's info. A stress or virial is stored as a 3x3 matrix. Six Voigt components are expanded, since that is what the reserved stress key recovers from the calculator. Any other shape raises, naming the structure and the file. Isolated atoms of one element that give different energies raise instead of the last one winning, so the E0 no longer depends on the order of the file. Agreeing duplicates are accepted, and a marked atom without an energy counts only when no other atom of its element has one. The key table test now pins where each property is stored, which makes the tests that restated that derivation redundant; they are removed and the inventory rows that cited them point at the table. The sign convention strings are no longer restated in test_units, a test checks that every exported factor is pinned, and the purity probe runs once per module.
3128bff to
733c1a8
Compare
GitHub Actions
Error
GitHub Actions
Error
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 2 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 733c1a8. Configure here.
A dataclass equality compared the numpy fields with ==, which raises for any array longer than one, so comparing two distinct configurations failed instead of returning a boolean. The shape error named a structure by its position after the isolated atoms were dropped, while the isolated-atom errors used the position in the file. Both now use the position in the file.
This branch was successfully deployed
3 participants
Configurationis the boundary object of the data layer: numpy arrays and plain Python, no tensors and no neighbour list, so a format backend produces one and graph construction consumes it without either side naming a framework. Itscellis the physical cell exactly as parsed. An aperiodic system needs an artificial box to bin a neighbour list into, and that box belongs to the neighbour builder; a configuration carrying a synthetic cell would report a stress divided by an invented volume, with nothing in the object saying the volume was invented.KeySpecificationis the only place a file key appears. Everything downstream of parsing is keyed by convention name, so nothing after the parser has to be told which key spec produced it. The thirteen default keys are a data contract rather than a default anyone is free to adjust: every labelled dataset on disk was written against them, and renaming one does not break a build, it stops reading somebody's forces.REF_magmomandREF_magforcesare on the default path, not behind a magnetic switch, so a magnetically labelled file reads with no flag set.embedding_specsis the runtime extension: declaring a feature is what makes a quantity that is neither a position nor a label nameable, which is what CORE-1 (#1555)'s derivative grammar then differentiates against.A property's storage is stated once, on the key table, as
graphoratom. Those are the same two words a user writes in an embedding feature'sper:.infoandarraysare ase's names for ase's two stores and appear only in the xyz backend, where ase is what is being touched.units.pyreads its factors fromase.unitsrather than writing them out, so a CODATA revision reaches both stacks together;test_units.pypins the values such a revision would move. The three sign statements are reproduced character for character from the characterization suite that pinned them against finite differences, andtests/architecture/test_convention_statements.pyasserts the two texts are still identical by reading both files withast, so it needs neither tree installed. The asymmetry is the thing a port normalises away: of the three, the stress is the only one not negated, and the virial is the negative of the very quantity the stress is built from.Two conflations worth naming, both found by checking this branch against the review of CORE-1 (#1555) before opening it:
property_weightsdoes not mean a label is absent. An absent label is zeroed, and a file is free to writeconfig_forces_weight=0.0for a structure whose forces are present.Configuration.is_labelledis the answer that can tell them apart, and a test pins that the weight cannot.units.pyis the single source for the prose. The machine-readable sign belongs with the declaration of the quantity it governs, which is CORE-1 (CORE-1 — Typed outputs + declarative observable specification #1555)'s.81 tests over the new modules: 65 on the parsing semantics ported from the characterization suite, 14 on the units and the conventions, 2 asserting
mace_coreimports neithermacenor torch nor jax.aseis an explicit dependency ofmace-core, with the reason in the module docstrings: the XYZ backend reads throughase.ioand the unit factors are ase's, so taking the dependency is what keeps a second copy of either from existing.Implements CORE-3 (#1557).
Note
Medium Risk
New parsing and key-resolution logic sits on the training data path and changes some legacy semantics (E0 conflicts, head storage, shared key_spec copies); broad test coverage mitigates but downstream adopters must align with the new contract.
Overview
Introduces the framework-agnostic data contract for MACE v1 in
mace-core: a numpy/Configurationboundary type, convention-based property keys resolved only at parse time, and ASE-backed XYZ reading.Configurationholds structures with labels keyed by convention names (energy,forces, …), not file keys;KeySpecificationmaps those names to file keys split intographvsatomstorage, withDefaultKeysfreezing the thirteen on-disk contract names. Parsing viaread_configurationshandles Voigt→3×3 stress/virials, ASE 3.23 reserved-key rewrites (on a copied spec), isolated-atom E0 extraction (conflicting E0s now error), andis_labelledto distinguish absent labels from zero loss weights. Also addsunits(factors fromase.units, verbatim sign-convention prose),AtomicNumberTable, and train/valid splitting with reproducible index logging.Dependencies: explicit
numpyandase>=3.23onmace-core. Tests: large characterization port (test_data_configuration.py), subprocess import-purity guard, pinned unit constants, and architecture check that sign prose matches the legacy characterization suite. Inventory tooling now accepts dual legacy +packages/mace-core/tests/…pins and rejects v1-only pins.Reviewed by Cursor Bugbot for commit 7cf3a81. Bugbot is set up for automated code reviews on this repo. Configure here.