Add typed model outputs and a declarative observable specification - #1723
Open
aacostadiaz wants to merge 12 commits into
Open
aacostadiaz wants to merge 12 commits into
aacostadiaz wants to merge 12 commits into
Conversation
The two abstractions CORE-1 asks for, and nothing else. `mace_core.outputs.MACEOutput` replaces the dictionary every legacy model forward returns. Six core fields plus a typed `extras` hatch, generic over the tensor type so the same class carries torch.Tensor, jax.Array or numpy arrays without mace_core importing any of them. A plain dataclass rather than a pydantic model: the field values are framework tensors whose type this package cannot name, so there is nothing for a validator to check, and the object is built once per forward pass. The one check it does make is that an `extras` key does not shadow a core field, because that failure is otherwise silent. `mace_core.observables` declares a property as a row. An `ObservableSpec` carries name, irreps, per_atom, units, normalization and a default loss weight; an `InputSpec` declares something the model is given; a derivative of one against the other is named by the rule `d_<q>_d_<x>`, with `forces`, `stress` and `magforces` as the three pairs that keep a name of their own and the two that carry a minus sign. The grammar is written over declared inputs rather than over positions and the cell because the magnetic family already needs the third: `magforces` is `-dE/d(magmom)`, and a grammar that knew only positions and the cell could not express it at all. `defaults/observables.yaml` ships energy with its position and cell derivatives, as packaged data rather than a Python literal. The point of a declarative spec is that a property can be added without touching code, and that has to include the code holding the defaults. `tests/architecture/observable_coverage.py` accounts for all 43 keys the frozen model forwards emit. The 43 are read out of the source by the P0-5 surface scan rather than listed, so a key added to a legacy forward fails the test instead of passing unnoticed, and the per-atom classification and unit of each come from the golden harness's channel declarations. Nineteen become observables with an irreps string, six are observables whose irreps are stated with the model-dependent part named (the number of readout layers, the maximum multipole order, whether an anisotropic readout was declared), five are derivatives under the rule, and ten carry a Drop row naming the mechanism that owns them instead. Nothing is deferred to a later ticket, and a row that cannot say what its irreps are fails the test rather than acquiring a TODO. Every derivative row states the sign the frozen tree actually reports and the test compares it against the rule, which turned up two things worth having in writing. `stress` and `virials` are not the same sign: the stress is built from the raw gradient and the virial is negated in the return statement, so `virials = -dE/dstrain` while `stress = +dE/dstrain / V`. Measured on the tiny_scaleshift anchor, `max|stress * V + virials|` is 1.2e-35. And `hessian` is not a derivative row at all: it holds `+d2E/dpos2`, which is the negative of the `d(forces)/d(pos)` the rule derives, because it is a second derivative of the energy rather than a first derivative of the forces. The grammar is first order by design and all three of its special cases are, so the key carries a Drop row naming the derivative engine, the same way `edge_forces` names the export path. Measured against a central difference of the forces, `max|hessian[:, 0] + dF/dx|` is 1.8e-10. With that settled every derivative row agrees with the rule, and the test asserts it with no way to annotate an exception. `BEC` is filed as its own observable and not as a second spelling of `dmu_dr`. LES builds a polarization from its own mean-removed latent charges with an epsilon factor, takes it through a Berry phase under periodic boundaries and differentiates that; the dielectric model differentiates its `dipole` readout. Different quantity, different gauge, different unit, and measured shape (n_atoms, 2, 3, 3) rather than (n_atoms, 3, 3). `docs/reforge/output_surface.md` writes down the three-layer surface the downstream tickets are measured against: 43 model keys, 15 that exist only at the calculator, 3 only at the evaluation CLI, 61 in the union. Every number there is re-derived by the test, so the document cannot drift from the tree. The target layout is updated to the shape that landed. It described a module per observable, which is the opposite of a declarative table, and it spelled the output type two ways. Verified: the package suites, the architecture suite, tests/golden and tests/unit are green, both toolchains pass, the import contracts hold and the built wheel carries the defaults file. tests/parity does not exist yet.
aacostadiaz
had a problem deploying
to
gpu-internal
September 11, 2026 13:59 — with 
GitHub Actions
Failure
GitHub Actions
Failure
aacostadiaz
had a problem deploying
to
gpu-internal
September 11, 2026 13:59 — with
GitHub Actions
Failure
GitHub Actions
Failure
9 tasks
Three jobs run `pytest tests` over the entire tree: the two GPU jobs in the MPCDF pipeline, and nightly's coverage-full and durations-refresh. None of them installs the v1 packages, and `tests/architecture` needs both trees importable. The import contracts resolve each root package on the filesystem, and the observable completeness test imports `mace_core`. That fails at collection, before any marker expression can deselect it, so a capability marker on those tests would not have helped. The GPU jobs of ACEsuit#1723 reported `ModuleNotFoundError: No module named 'mace_core'` after a suite that was otherwise fully green. No coverage is lost. `tests/architecture` is gated on every pull request by the `architecture` job in ci-core.yaml, which is the one place both trees are installed editable, and nothing in the directory is gpu-marked or worth a coverage number. A later suite that needs the v1 stack has to be listed on those lines too, or be given its own job. `tests/parity` is the next one.
aacostadiaz
had a problem deploying
to
gpu-internal
September 15, 2026 09:11 — with
GitHub Actions
Error
GitHub Actions
Error
aacostadiaz
had a problem deploying
to
gpu-internal
September 15, 2026 09:11 — with
GitHub Actions
Error
GitHub Actions
Error
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
A fourth job sweeps the whole tree: `backends-cpu` in ci-extensions.yaml runs
`pytest tests` with a cueq marker over everything, and installs the legacy
distribution alone. It is paths-filtered, so it skipped on this pull request
and the breakage would have surfaced on the first one touching backends, or in
the nightly that calls this workflow. It gets the same `--ignore` as the other
three. A search for whole-tree selections now finds four, and all four carry it.
The per-atom energy was renamed without the table saying so. The core field of
MACEOutput is `node_energies` and the legacy key is `node_energy`, one letter
apart, so the row specified an observable whose name missed the field it
belongs in. The consequence was live: `get("node_energy")` returned None on an
output whose `node_energies` field was filled, and `extras["node_energy"]` was
free to sit beside that field holding the same quantity, which is precisely the
dual storage the shadowing guard exists to prevent for the names it knows.
The row now records the rename, and a new test asserts that each of the six
core fields is claimed by exactly one legacy key. A field nothing claims is a
field no model fills; a field claimed twice is the same silent dual storage
from the other direction. Both would have passed unnoticed before.
The `extras` test that used `node_energy` as its benign example now uses a name
that is not a near-miss of a field, since the old one read as an endorsement of
the thing v1 renames.
aacostadiaz
had a problem deploying
to
gpu-internal
September 15, 2026 09:27 — with
GitHub Actions
Failure
GitHub Actions
Failure
aacostadiaz
had a problem deploying
to
gpu-internal
September 15, 2026 09:27 — with
GitHub Actions
Failure
GitHub Actions
Failure
The GPU jobs kept failing on `ModuleNotFoundError: No module named 'mace_core'` after the `--ignore=tests/architecture` added for exactly that. The flag was in the wrong tree, and the reason is a deliberate property of the bridge rather than an oversight. For a pull request from a fork, `.github/workflows/ci-gpu-mpcdf.yaml` takes the tested tree from the fork and the pipeline definition from the base ref, so that a fork can change what gets tested and never what runs it on MPCDF hardware. An `--ignore` added to `.github/gitlab/ci.yml` is therefore invisible to the pull request that adds it, and stays invisible until it merges. The tested tree is the only lever a fork has. So the guard moves into the one module that needs it. `find_spec` rather than `pytest.importorskip`, because it resolves the module without executing it: a `mace_core` that is absent skips, while one that is installed and broken still raises at the real import below. Verified both ways, with the module hidden from `sys.path` and with a deliberately broken copy ahead of it on it. The skip cannot quietly hide these tests from the job that owes them. The `architecture` job runs `lint-imports` before the suite, and that step fails outright when a root package is missing from the filesystem, so the install this guard depends on cannot disappear unnoticed. That makes the four `--ignore` flags redundant, and they are reverted. Three of them would only ever have taken effect after merge, and all four removed coverage from sweeps where the rest of `tests/architecture` had been running perfectly well.
The field described itself as the same derivative the stress is built from, before the volume division. It is the negative of it: the virial is -dE/dstrain while the stress is +dE/dstrain over the volume, so stress times volume is minus the virial. The inventory in this same branch already says so, with the measurement that settles it: max|stress * V + virials| is 1.2e-35 while max|stress * V - virials| is 6.5e-3. It was the dataclass that was left behind. The confusion has a source worth naming. The legacy helper builds the stress from the raw gradient and negates it into the virial only in its return statement, so reading that function leaves the impression that the two share a sign. A consumer that takes the docstring at its word gets a virial of the wrong sign, which has the right magnitude and is therefore hard to notice.
The row already declared the symmetric strain: irreps 0e+2e, six components rather than the nine entries of the cell vectors, and dimensionless units. Only the name said cell. A derivative inherits its irreps from the input it is taken against, so the shape stress carried was right only by accident, and any other observable asking for that derivative read d_<q>_d_cell, which is not what is computed. Record what an input is while the name is being corrected. It is a leaf a derivative can be taken against, not necessarily a field read from the data: pos is both, strain is only the first, materialised as zeros by the derivative engine around the model call.
Neither is a fact about a quantity. A loss weight has to change between stages, which legacy already does with a second set of its own, so it belongs to the loss config. Scaling belongs to the head. A derivative has no scale of its own: it is differentiated from what the head produces, so it follows that scaling, and a scaling field on a derivative row has nothing to set. The head's scale can also come from a different target, which a field on one observable cannot express. Today `--scaling` defaults to `rms_forces_scaling`, which scales the energy readout with the RMS of the force targets. What is left describes the quantity: name, irreps, per_atom and units. The declared name is what the model and loss configs key on.
This was referenced Sep 21, 2026
The name and the sign of a derivative were in `SPECIAL_CASES`, three rows
mapping a `(quantity, input)` pair to a name and a sign. The same fact was
written three times in the tree: the machine-readable pair there, the prose in
`mace_core.units.SIGN_CONVENTIONS`, and a comment in the declarations file
saying "Named `forces`, reported as -dE/dpos" because the schema could not
express it. A comment that says what a field would say is a missing field.
It had also already drifted. The table carried a row for `("energy",
"magmom")` and no shipped catalogue declares `magmom`, so the third of its
three rows described something no data asked for.
The consequence is not tidiness. A fourth derivative with a name of its own,
torques as `-dE/d(orientation)` or a polarizability as `d(dipole)/d(field)`,
meant editing `mace_core`, which is the one thing a declarative grammar exists
to prevent. `derivatives.py` made that argument for `magforces` one level down
and then stopped one level short of applying it to itself.
So `DerivativeRequest` gains `name` and `sign`, the table goes, and what
remains is the default rule: `d_<quantity>_d_<input>` with the gradient's own
sign. The shipped catalogue now declares `forces` at -1 and `stress` at +1 as
data.
This reverses a decision the docstring stated deliberately: that letting a
declaration set them "would reintroduce the per-consumer naming this
abstraction removes". That objection is about a *consumer* inventing a name,
and it still cannot: the name is decided once, in the catalogue, and every
consumer reads it off the resolved spec. What changes is only whether the
catalogue's source is Python or YAML.
Two guards, because a free field is a field a typo silently changes:
`name` obliges `sign`. A renamed derivative inheriting +1 in silence is a model
trained on inverted forces that runs perfectly well, and nothing downstream can
see it. A sign must also be exactly +1 or -1; a scale factor is not a sign, and
the volume division that turns a strain derivative into a stress still belongs
to whatever computes the stress.
A declared name may not be spelled `d_<x>_d_<y>`. That spelling states which
quantity was differentiated, so a custom name wearing it would be asserting
something untrue.
The architecture inventory keeps its own copy of the three names on purpose,
and it is now explicit in a `name` field rather than implied by the rule. That
table is read off the frozen tree and the declarations file is authored, so
they are two sources; deriving both from one would have left the test agreeing
with itself.
`units.SIGN_CONVENTIONS` is still a third copy, in prose, which cannot be
derived. Binding it to the declarations needs `mace_core.units`, which arrives
with the configuration work and does not exist on this branch yet.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.
❌ 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 e2f728a. Configure here.
`__contains__` resolved an observable name through `FIELD_BY_OBSERVABLE`, so
`"energy" in output` was `True` when `total_energy` was set, while `names()`
yielded the storage field. The two accessors therefore disagreed about the same
value:
>>> out = MACEOutput(total_energy=..., forces=...)
>>> "energy" in out, "energy" in set(out.names())
(True, False)
A consumer correlating `ObservableCatalogue.names()` with `MACEOutput.names()`
intersects `("energy", "forces", "stress")` against `("total_energy",
"forces")` and gets `{"forces"}`. The energy is dropped, no exception is
raised, and membership says it was there the whole time. Reported by Bugbot on
the review of this branch and reproduced before being believed.
`names()` now yields observable names, since that is the vocabulary a
declaration uses and the one a caller is correlating against. Reaching a value
goes through `get`, which accepts either spelling, and the one consumer in the
tree already did. The inverse map is derived from the forward one rather than
written out, so an entry added to one direction cannot be missing from the
other, and the new test asserts over both directions rather than over the one
alias that exists today.
`CORE_FIELD_NAMES` keeps yielding storage names, which is what it is for, and
now says so and points at the map.
aacostadiaz
had a problem deploying
to
gpu-internal
September 22, 2026 05:46 — with
GitHub Actions
Error
GitHub Actions
Error
aacostadiaz
had a problem deploying
to
gpu-internal
September 22, 2026 05:46 — with
GitHub Actions
Error
GitHub Actions
Error
`packages-lint` failed on the test added with the previous commit, and only there: `ty` runs against the packages installed editable, and the local venv's editable entry points at a different checkout, so the local run reported five unrelated unresolved imports and hid this one. Reproduced in a venv built the way the job builds it, ruff and ty pinned to the same versions. numpy types an array's shape, so inferring `TensorT` from a 1-D energy beside a 2-D force array makes it a union of the two, and `dict` is invariant in its value type: the `extras` literal is then unassignable to `dict[str, <that union>]`. The construction is fine at runtime and says nothing about the class. The parameter is pinned to `np.ndarray` at the construction, which is the spelling the rest of the file already uses. `ty` suggests widening `extras` to a `Mapping` instead, and that is the wrong trade here: it would make a result object look immutable to consumers in order to quiet a test.
The inventory renames the frozen tree's `node_energy` to `node_energies`, and
its note says why: keeping the singular "would leave `extras['node_energy']`
able to sit beside a `node_energies` field holding the same quantity". The
rename did not achieve that. `extras` accepts any key the shadowing guard does
not recognise, and the guard knows only the core fields and the one observable
alias, so the exact construction the note rules out was accepted:
MACEOutput(node_energies=..., extras={"node_energy": ...})
Reported by Bugbot, and only half of it holds. Its first claim is that
`get("node_energy")` should find the field; that is the rename, working as
decided, and resolving it would carry both spellings after all. Its second is
the dual storage above, which is real and reproduced.
So the retired spellings are named and refused at construction, and they are
deliberately *not* resolvable: `get` still misses them and `__contains__` still
says no, with a test on each half so a later change cannot quietly turn a
refusal into an alias. Refusing it is unconditional rather than conditional on
the field being filled, because the two names are one quantity and only one of
them is this type's; the test states that so relaxing it would be a decision.
The default catalogue was a YAML file shipped as package data, with a loader around it. A user could not reach it from a configuration: the model config names observables, and the names resolve against the shipped file, so adding a property meant editing a file inside the package, which is editing code by another name. DEFAULT_CATALOGUE builds the same catalogue from the spec objects, which are frozen, so it is shared rather than reloaded. The validation is unchanged, and the tests that declared rows as YAML text now declare them as the plain data a configuration would hold. mace-core no longer depends on PyYAML. CORE-1
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
MACEOutputreplaces the dict every legacy model forward returns: six core fields plus a typedextrashatch, generic over the tensor type somace_coreimports no framework. A plain dataclass, since the values are tensors whose type this package cannot name and there is nothing to validate.ObservableSpecmakes a property a row instead of a code change. Derivatives are namedd_<q>_d_<x>, withforces,stressandmagforceskeeping their own names. The grammar runs over declared inputs, not just positions and the cell, becausemagforcesalready needs the third.defaults/observables.yamlships energy with its two derivatives.All 43 keys the frozen forwards emit are classified in
tests/architecture/observable_coverage.py, read out of the source by the P0-5 scan rather than listed, so a new legacy key fails the test. 22 carry a literal irreps string, 6 state it with the model parameter that sets the rest (readout count, maximum multipole order, whether an anisotropic readout is declared), 5 are derivatives, 10 are drops naming what owns them instead. Nothing is deferred: a row that cannot say what its irreps are fails the test rather than acquiring a TODO.docs/reforge/output_surface.mdrecords the three-layer surface: 43 + 15 + 3 = 61 names, every number re-derived by the test.Three things the sign checks turned up, all measured on the committed anchors:
virials = -dE/dstrainbutstress = +dE/dstrain / V. They differ by a sign, not only the volume (max|stress*V + virials|= 1.2e-35).hessianis+d2E/dpos2, the negative ofd(forces)/d(pos). It is a second derivative, so it gets a Drop row naming the derivative engine.BECis not another spelling ofdmu_dr. LES differentiates a Berry-phase polarization built from mean-removed latent charges; the dielectric model differentiates itsdipolereadout. Different quantity, different unit, measured shape(n_atoms, 2, 3, 3).Implements CORE-1 (#1555).
Note
Medium Risk
Adds foundational API and validation for all model outputs and derivatives; mistakes in sign/naming or coverage tables could misclassify legacy behavior, though extensive architecture tests mitigate this.
Overview
Introduces
mace-corecontract types for v1: a pydantic observable catalogue (InputSpec,ObservableSpec,DerivativeRequest,DEFAULT_CATALOGUE) with irreps grammar,d_<q>_d_<x>derivative naming, and explicit signs for renamed quantities likeforces/stress; plus genericMACEOutput(six core fields +extras) to replace untyped forward dicts.mace-corenow depends on numpy and pydantic and exports these types from its public API.Docs move extension examples from
observables.yamltoobservables.pyandregister_observables.output_surface.mddefines the 61-name user surface (model forward + ASE calculator + eval CLI).target_layout.mdreflects the consolidatedobservables/layout and config-only observable extension.Architecture gates:
observable_coverage.pyclassifies every legacy forward key (Spec / Derivative / Drop);test_observable_completeness.pyenforces full coverage, derivative/sign consistency, core-field claims, and re-derived layer counts vs the surface doc.Reviewed by Cursor Bugbot for commit 1cb4fed. Bugbot is set up for automated code reviews on this repo. Configure here.