From c72da87968e4db386b99317abab8f86a04304f8f Mon Sep 17 00:00:00 2001 From: Alejandro Acosta <127198532+aacostadiaz@users.noreply.github.com> Date: Fri, 11 Sep 2026 12:36:03 +0100 Subject: [PATCH 01/12] Add typed model outputs and a declarative observable specification 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__d_`, 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. --- docs/reforge/output_surface.md | 66 +++ docs/reforge/target_layout.md | 28 +- packages/mace-core/pyproject.toml | 6 +- packages/mace-core/src/mace_core/__init__.py | 23 +- .../src/mace_core/defaults/observables.yaml | 49 ++ .../src/mace_core/observables/__init__.py | 49 ++ .../src/mace_core/observables/defaults.py | 50 ++ .../src/mace_core/observables/derivatives.py | 65 +++ .../src/mace_core/observables/grammar.py | 117 +++++ .../src/mace_core/observables/spec.py | 320 ++++++++++++ packages/mace-core/src/mace_core/outputs.py | 141 +++++ packages/mace-core/tests/test_observables.py | 410 +++++++++++++++ packages/mace-core/tests/test_outputs.py | 113 ++++ tests/architecture/observable_coverage.py | 483 ++++++++++++++++++ .../test_observable_completeness.py | 265 ++++++++++ 15 files changed, 2168 insertions(+), 17 deletions(-) create mode 100644 docs/reforge/output_surface.md create mode 100644 packages/mace-core/src/mace_core/defaults/observables.yaml create mode 100644 packages/mace-core/src/mace_core/observables/__init__.py create mode 100644 packages/mace-core/src/mace_core/observables/defaults.py create mode 100644 packages/mace-core/src/mace_core/observables/derivatives.py create mode 100644 packages/mace-core/src/mace_core/observables/grammar.py create mode 100644 packages/mace-core/src/mace_core/observables/spec.py create mode 100644 packages/mace-core/src/mace_core/outputs.py create mode 100644 packages/mace-core/tests/test_observables.py create mode 100644 packages/mace-core/tests/test_outputs.py create mode 100644 tests/architecture/observable_coverage.py create mode 100644 tests/architecture/test_observable_completeness.py diff --git a/docs/reforge/output_surface.md b/docs/reforge/output_surface.md new file mode 100644 index 000000000..8814c72c2 --- /dev/null +++ b/docs/reforge/output_surface.md @@ -0,0 +1,66 @@ +# The user-observable output surface + +"v1 has the same functionality as develop" is a claim about the names a user +can read out, not about the names one function happens to return. A key that +reaches the user only through the ase calculator, or only through +`mace_eval_configs`, is functionality just the same, and a completeness gate +keyed on the model `forward` alone cannot see it. + +The surface is therefore the union of three layers. + +| layer | read from | owner | keys | new at this layer | +|---|---|---|---|---| +| (a) model forward | `mace/modules/models.py`, `mace/modules/extensions.py` | CORE-1 (#1555) | 43 | 43 | +| (b) ase calculator `results` | `mace/calculators/mace.py` | DEP-1 (#1583) for the energy family, DEP-1a (#1634) for dipole / dielectric / polar / LES / magnetic | 31 | 15 | +| (c) `mace_eval_configs` | `mace/cli/eval_configs.py` | CLI-1 (#1579) | 13 | 3 | +| **union** | | | | **61** | + +61, not 43, is the number "the output surface survived the rewrite" is measured +against. + +## Deriving it, rather than trusting this table + +Every number above is extracted from the frozen tree by +`tests/golden/surface_scan.py`, the same mechanical scan run at three sites, +and `tests/architecture/test_observable_completeness.py` re-derives all four +and fails if this table disagrees. Each owning ticket runs the extraction over +its own layer rather than copying a number from here: a hand-kept list that +looks complete and silently is not is the defect this whole exercise exists to +remove. + +Two traps the scan already accounts for, both of which shrink the surface +quietly when missed: + +- Layer (a) must follow keys **assigned onto the returned object**, not only + dict literals. The self-consistent magnetic model assigns three diagnostics + after building its output, so an extraction that stops at return literals + stops at 40. +- Layer (b)'s committee keys must be read off the code, not off + `implemented_properties`. The loop emits both a `_comm` and a `_var` suffix + for all four members of the ensemble store, while `implemented_properties` + advertises four of those eight: `forces_var`, `stress_comm` and + `dipole_comm` are produced and never declared. + +## What is new at each layer + +**(b) exists only at the calculator**, 15 names: `free_energy` and `energies` +(the aliases and the E0-inclusive per-atom energy), `stresses` (the Voigt +per-atom rename of the model's `atomic_stresses`), `LES_alphas`, `LES_kappas`, +`bec` (the lower-cased mean of the model's `BEC`), `MACE_magmoms`, and the +eight committee keys `{energy,forces,stress,dipole}_{comm,var}`. + +**(c) exists only at the evaluation CLI**, 3 names, and each is a *rename* of a +model key, which is exactly why they are easy to lose: `BO_contributions` (model +`contributions`), `descriptors` (model `node_feats`, after invariant extraction +and layer truncation), `node_energies` (model `node_energy`). Only `energy`, +`forces` and `stress` are shared with the calculator; the other ten names that +layer writes reach the user through this CLI alone. + +## Layer (a), key by key + +CORE-1 owns layer (a) and classifies all 43 in +`tests/architecture/observable_coverage.py`: each key becomes a declared +`ObservableSpec`, a derivative of one under the `d__d_` rule, or a row +saying explicitly that it is not an observable and naming the mechanism that +owns it instead. The test beside that file fails on a key with no row, so a key +added to a legacy forward cannot pass unclassified. diff --git a/docs/reforge/target_layout.md b/docs/reforge/target_layout.md index 6eafa5f17..f0952986a 100644 --- a/docs/reforge/target_layout.md +++ b/docs/reforge/target_layout.md @@ -30,7 +30,8 @@ packages/mace-core/ │ ├── __init__.py # re-exports public types/config/registries; does NOT import heavy submodules │ ├── _version.py # contract version (semver of the weights format/spec) │ │ -│ ├── types.py # MACEOutputs (typed dataclass, replaces the get_outputs dict); GRAPH_SCHEMA + GraphInfo/GraphView (rfc-03 flat-dict contract) +│ ├── outputs.py # MACEOutput (typed dataclass generic over the tensor type, replaces the get_outputs dict) +│ ├── types.py # GRAPH_SCHEMA + GraphInfo/GraphView (rfc-03 flat-dict contract) │ ├── graph.py # flat-dict contract {node_attrs,edge_index,positions,batch,cell,shifts,...} + shape/dtype validation │ │ │ ├── config/ @@ -46,15 +47,14 @@ packages/mace-core/ │ │ ├── number_table.py # AtomicNumberTable (reimplemented; mirror of tools/utils.py, without dragging in train.py) │ │ └── default_keys.py # DefaultKeys (reimplemented; mirror of tools/default_keys.py) │ │ -│ ├── observables/ -│ │ ├── __init__.py # OBSERVABLE_REGISTRY (declarative: name → spec) -│ │ ├── base.py # Observable protocol: output irreps, how it's derived (readout | autograd | grad-strain) -│ │ ├── energy.py # Energy, SiteEnergy (readout+scatter_sum) -│ │ ├── forces.py # Forces (=-dE/dx via autograd) — spec only, the physics is executed by mace_torch/mace_jax -│ │ ├── stress.py # Stress, Virials (grad w.r.t. strain; sign convention PINNED here) -│ │ ├── dipole.py # Dipole, AtomicDipole -│ │ ├── polarizability.py # Polarizability (dielectric/polar) -│ │ └── hessian.py # Hessian (second derivative) +│ ├── observables/ # a property is a row in a declarations file, not a module per property +│ │ ├── __init__.py # the public surface of the package +│ │ ├── spec.py # InputSpec, ObservableSpec, DerivativeSpec, ObservableCatalogue (pydantic) +│ │ ├── grammar.py # the irreps string grammar: parse + validate + dimension (no algebra) +│ │ ├── derivatives.py # d__d_ naming, and the three special cases with their signs +│ │ └── defaults.py # loader for a declarations file +│ ├── defaults/ +│ │ └── observables.yaml # energy + its position and cell derivatives: the row every observable copies │ │ │ ├── kernels/ │ │ ├── protocol.py # KernelBackend Protocol, generic over TensorT: make_* factories + capabilities (§3.1) @@ -493,10 +493,10 @@ silently wrong forces); it stays usable for inference (`supports_double_backward ### 3.2 A new observable (config only) -- **Extender touches:** a `mace_core/observables/myobs.py` file with an `Observable` (declares output irreps and derivation mode: `readout` | `autograd(energy, wrt=positions)` | `grad_strain`) + `@register_observable("myobs")`. -- **Core touched:** zero existing files (only the new module is added). The model exposes it automatically because `BaseMACE` iterates over `config.observables`; `MACEOutputs` is a dataclass with optional fields populated by name. -- **Enabling it:** `ModelConfig(observables=["energy","forces","myobs"])`. -- **Test:** `mace_core/tests/test_observable_registry.py` validates irreps/derivation consistency (pure); if it is autograd-derived, `tests/parity` verifies finite-diff. +- **Extender touches:** a declarations file: one `ObservableSpec` row giving `name`, `irreps`, `per_atom`, `units`, `normalization`, `default_loss_weight`, and the declared inputs to differentiate against. No module, no decorator. A derivative is named by the rule `d__d_`, with `forces`, `stress` and `magforces` as the three special cases, so asking for a derivative against a newly declared input needs no code either. +- **Core touched:** zero files. The model exposes the row automatically because `BaseMACE` iterates over the declared observables; `MACEOutput` carries the six core fields and everything else by name in `extras`. +- **Enabling it:** list it in the model config's observables, or point the config at a declarations file that extends `defaults/observables.yaml`. +- **Test:** `packages/mace-core/tests/test_observables.py` validates the grammar and the derivative naming (pure); if it is autograd-derived, `tests/parity` verifies finite-diff. ### 3.3 A new loss / transform (plugin registry) diff --git a/packages/mace-core/pyproject.toml b/packages/mace-core/pyproject.toml index 7caf4c664..ec0fe9967 100644 --- a/packages/mace-core/pyproject.toml +++ b/packages/mace-core/pyproject.toml @@ -17,7 +17,11 @@ classifiers = [ "Programming Language :: Python :: 3.13", "Operating System :: OS Independent", ] -dependencies = [] +dependencies = [ + "numpy>=1.23", + "pydantic>=2.7", + "pyyaml>=6.0", +] [project.urls] Homepage = "https://github.com/ACEsuit/mace" diff --git a/packages/mace-core/src/mace_core/__init__.py b/packages/mace-core/src/mace_core/__init__.py index 4122cd7d8..e35dee6f9 100644 --- a/packages/mace-core/src/mace_core/__init__.py +++ b/packages/mace-core/src/mace_core/__init__.py @@ -1,11 +1,30 @@ """Framework-agnostic contract and pure math for MACE v1. -Scaffold only. The public surface arrives with the tickets that build it. +This package imports no framework. Everything here is expressed over plain +Python, numpy and pydantic, so the same types carry ``torch.Tensor`` in +``mace_torch`` and ``jax.Array`` in ``mace_jax``. """ from importlib.metadata import PackageNotFoundError, version -__all__ = ["__version__"] +from mace_core.observables import ( + DerivativeSpec, + InputSpec, + ObservableCatalogue, + ObservableSpec, + load_default_catalogue, +) +from mace_core.outputs import MACEOutput + +__all__ = [ + "DerivativeSpec", + "InputSpec", + "MACEOutput", + "ObservableCatalogue", + "ObservableSpec", + "__version__", + "load_default_catalogue", +] #: Version of the installed `mace-core` distribution. Read from installed metadata #: rather than hardcoded, so it cannot drift from what pip resolved. diff --git a/packages/mace-core/src/mace_core/defaults/observables.yaml b/packages/mace-core/src/mace_core/defaults/observables.yaml new file mode 100644 index 000000000..9cbc7a369 --- /dev/null +++ b/packages/mace-core/src/mace_core/defaults/observables.yaml @@ -0,0 +1,49 @@ +# The canonical observable declarations: energy, and its derivatives with +# respect to the two inputs every model has. +# +# This file is the example every other observable copies. A new property is a +# row here (or in a project's own file) and needs no code: the spec drives the +# head, the loss term, the derivative names and the per-atom/per-graph padding. +# +# Units follow the project convention: eV and Å. The strain is dimensionless, +# so the cell input carries "1". + +inputs: + # Atomic positions. A polar vector: it changes sign under inversion, which is + # what makes the energy gradient taken against it a 1o quantity as well. + - name: pos + irreps: "1o" + per_atom: true + units: "Å" + + # The symmetric strain the cell derivative is actually taken against, not the + # nine cell entries. A symmetric rank-2 tensor is a scalar plus an l=2 part. + - name: cell + irreps: "0e+2e" + per_atom: false + units: "1" + +observables: + - name: energy + irreps: "0e" + per_atom: false + units: "eV" + # The std of the target, the successor of the legacy `std_scaling` entry. + # The legacy default was `rms_forces_scaling`, which scales the energy + # readout by the RMS of the *force* targets; that couples two observables + # through one number and cannot be written as a per-observable field. Which + # of the two v1 defaults to is a decision for the head, not for this file. + normalization: "std" + default_loss_weight: 1.0 + derivatives: + # Named `forces`, reported as -dE/dpos. + - wrt: pos + units: "eV/Å" + normalization: "rms" + default_loss_weight: 100.0 + # Named `stress`, reported as +dE/dstrain, divided by the cell volume by + # whatever computes it. The division is not part of the sign. + - wrt: cell + units: "eV/Å^3" + normalization: "none" + default_loss_weight: 1.0 diff --git a/packages/mace-core/src/mace_core/observables/__init__.py b/packages/mace-core/src/mace_core/observables/__init__.py new file mode 100644 index 000000000..d5d3f4503 --- /dev/null +++ b/packages/mace-core/src/mace_core/observables/__init__.py @@ -0,0 +1,49 @@ +"""Declarative observables: what a model computes, declared rather than coded.""" + +from mace_core.observables.defaults import ( + DEFAULTS_RESOURCE, + load_catalogue, + load_default_catalogue, +) +from mace_core.observables.derivatives import ( + SPECIAL_CASES, + derivative_name, + derivative_sign, +) +from mace_core.observables.grammar import ( + IRREPS_GRAMMAR, + IrrepsGrammarError, + IrrepTerm, + irreps_dimension, + parse_irreps, +) +from mace_core.observables.spec import ( + NORMALIZATIONS, + DerivativeRequest, + DerivativeSpec, + InputSpec, + Normalization, + ObservableCatalogue, + ObservableSpec, +) + +__all__ = [ + "DEFAULTS_RESOURCE", + "IRREPS_GRAMMAR", + "NORMALIZATIONS", + "SPECIAL_CASES", + "DerivativeRequest", + "DerivativeSpec", + "InputSpec", + "IrrepTerm", + "IrrepsGrammarError", + "Normalization", + "ObservableCatalogue", + "ObservableSpec", + "derivative_name", + "derivative_sign", + "irreps_dimension", + "load_catalogue", + "load_default_catalogue", + "parse_irreps", +] diff --git a/packages/mace-core/src/mace_core/observables/defaults.py b/packages/mace-core/src/mace_core/observables/defaults.py new file mode 100644 index 000000000..37eb96293 --- /dev/null +++ b/packages/mace-core/src/mace_core/observables/defaults.py @@ -0,0 +1,50 @@ +"""Loading a catalogue from a declarations file. + +The shipped file is packaged data rather than a Python literal, because the +whole point of the declarative spec is that a property can be added without +touching code -- including the code that holds the defaults. +""" + +from __future__ import annotations + +from importlib.resources import files +from pathlib import Path +from typing import Any + +import yaml + +from mace_core.observables.spec import ObservableCatalogue + +__all__ = [ + "DEFAULTS_RESOURCE", + "load_catalogue", + "load_default_catalogue", +] + +#: Where the shipped declarations live, relative to the package root. +DEFAULTS_RESOURCE = "defaults/observables.yaml" + + +def _catalogue_from_text(text: str, source: str) -> ObservableCatalogue: + document: Any = yaml.safe_load(text) + if document is None: + document = {} + if not isinstance(document, dict): + raise ValueError( + f"{source}: an observable declarations file must be a mapping with " + f"`inputs` and `observables` keys, not a " + f"{type(document).__name__}." + ) + return ObservableCatalogue.model_validate(document) + + +def load_catalogue(path: str | Path) -> ObservableCatalogue: + """Load and validate a declarations file from disk.""" + path = Path(path) + return _catalogue_from_text(path.read_text(encoding="utf-8"), str(path)) + + +def load_default_catalogue() -> ObservableCatalogue: + """The shipped declarations: energy plus its position and cell derivatives.""" + resource = files("mace_core").joinpath(DEFAULTS_RESOURCE) + return _catalogue_from_text(resource.read_text(encoding="utf-8"), DEFAULTS_RESOURCE) diff --git a/packages/mace-core/src/mace_core/observables/derivatives.py b/packages/mace-core/src/mace_core/observables/derivatives.py new file mode 100644 index 000000000..50633d5b6 --- /dev/null +++ b/packages/mace-core/src/mace_core/observables/derivatives.py @@ -0,0 +1,65 @@ +"""How a derivative of a declared quantity is named and signed. + +The rule is one line: the derivative of a declared quantity ``q`` with respect +to a declared input ``x`` is called ``d__d_``. Three pairs have a name of +their own, and they are data in the table below rather than branches spread +through the consumers that need them. + +The third special case is the reason this grammar is written over declared +inputs rather than over positions and the cell. ``magforces`` is +``-dE/d(magmom)``, computed in the same autograd call as the forces, trained +with its own loss term, and used by the magnetic self-consistent model to drive +its fixed point. A grammar that knew only ``d__d_pos`` and ``d__d_cell`` +could not express it, and the magnetic work would have had to go around the +abstraction that exists to prevent exactly that. + +The sign is the one the reported quantity carries, so that +``reported = sign * d(quantity)/d(input)``. The volume division that turns the +strain derivative into a stress is not a sign and is not here: it belongs to +whatever computes the stress. +""" + +from __future__ import annotations + +__all__ = [ + "SPECIAL_CASES", + "derivative_name", + "derivative_sign", +] + +#: ``(quantity, input) -> (name, sign)`` for the three pairs whose name is not +#: the ``d__d_`` default. Everything else follows the rule. +SPECIAL_CASES: dict[tuple[str, str], tuple[str, int]] = { + ("energy", "pos"): ("forces", -1), + ("energy", "cell"): ("stress", +1), + ("energy", "magmom"): ("magforces", -1), +} + + +def derivative_name(quantity: str, wrt: str) -> str: + """The canonical name of ``d(quantity)/d(wrt)``. + + Args: + quantity: The name of the differentiated observable. + wrt: The name of the declared input it is differentiated against. + + Returns: + The special-cased name if the pair has one, otherwise + ``f"d_{quantity}_d_{wrt}"``. + """ + special = SPECIAL_CASES.get((quantity, wrt)) + if special is not None: + return special[0] + return f"d_{quantity}_d_{wrt}" + + +def derivative_sign(quantity: str, wrt: str) -> int: + """The sign the reported derivative carries: ``reported = sign * dq/dx``. + + ``+1`` unless the pair is one of the two negated special cases, forces and + magnetic forces, which are both the negative gradient of the energy. + """ + special = SPECIAL_CASES.get((quantity, wrt)) + if special is not None: + return special[1] + return 1 diff --git a/packages/mace-core/src/mace_core/observables/grammar.py b/packages/mace-core/src/mace_core/observables/grammar.py new file mode 100644 index 000000000..d5c088412 --- /dev/null +++ b/packages/mace-core/src/mace_core/observables/grammar.py @@ -0,0 +1,117 @@ +"""The irreps string grammar, and nothing beyond it. + +An observable declares the spherical-tensor shape of its values as a string: +``"0e"`` for a scalar, ``"1o"`` for a polar vector, ``"1e"`` for an axial one, +``"0e+2e"`` for a symmetric rank-2 tensor, ``"128x0e+128x1o+128x2e"`` for a +block of hidden features. This module parses and validates that string. + +It is grammar only. No tensor products, no Clebsch-Gordan coefficients, no +simplification or sorting of terms: the algebra lands in its own module with +the reduced basis, and a half-implementation here would be the version every +later caller had to work around. What a caller gets from this module is the +guarantee that a declaration is well formed, the terms it names, and the +dimension they add up to. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass + +__all__ = [ + "IRREPS_GRAMMAR", + "IrrepTerm", + "IrrepsGrammarError", + "irreps_dimension", + "parse_irreps", +] + +#: Stated in every error this module raises, because an error that says only +#: "invalid" leaves the reader to guess between four plausible spellings. +IRREPS_GRAMMAR = ( + "a '+'-separated sum of terms, each written '' or " + "'x', where is a non-negative integer and " + " is 'e' (even) or 'o' (odd). Examples: '0e' (a scalar), '1o' (a " + "polar vector), '1e' (an axial vector), '0e+2e' (a symmetric rank-2 " + "tensor), '128x0e+128x1o+128x2e'." +) + +_TERM = re.compile(r"^(?:(\d+)x)?(\d+)([eo])$") + + +class IrrepsGrammarError(ValueError): + """A declaration that is not a well-formed irreps string.""" + + +@dataclass(frozen=True) +class IrrepTerm: + """One ``x`` term of a declaration. + + Attributes: + multiplicity: How many copies of the irrep the term declares. + degree: The rotation order ``l``. A value of ``l`` spans ``2l + 1`` + components. + parity: ``"e"`` or ``"o"``, the behaviour under inversion. The + distinction is load-bearing rather than decorative: a force is + ``1o`` and a magnetic moment is ``1e``, and a model that confuses + them is wrong under inversion while looking right under rotation. + """ + + multiplicity: int + degree: int + parity: str + + @property + def dimension(self) -> int: + """The number of components this term contributes.""" + return self.multiplicity * (2 * self.degree + 1) + + def __str__(self) -> str: + return f"{self.multiplicity}x{self.degree}{self.parity}" + + +def parse_irreps(text: str, *, observable: str | None = None) -> tuple[IrrepTerm, ...]: + """Parse an irreps declaration into its terms. + + Args: + text: The declaration, for example ``"128x0e+128x1o"``. + observable: The observable the declaration belongs to. Named in the + error message, because a validation failure reported without it + tells the reader which grammar was violated and not which of their + declarations violated it. + + Returns: + The terms, in the order they were written. The order is preserved + rather than sorted: it is the layout of the values themselves. + + Raises: + IrrepsGrammarError: If the declaration is empty or any term is + malformed. + """ + where = f"observable {observable!r}: " if observable else "" + if not text or not text.strip(): + raise IrrepsGrammarError( + f"{where}the irreps declaration is empty. Expected {IRREPS_GRAMMAR}" + ) + terms = [] + for piece in text.split("+"): + match = _TERM.match(piece.strip()) + if match is None: + raise IrrepsGrammarError( + f"{where}{piece.strip()!r} is not a valid irreps term in " + f"{text!r}. Expected {IRREPS_GRAMMAR}" + ) + multiplicity, degree, parity = match.groups() + terms.append( + IrrepTerm( + multiplicity=1 if multiplicity is None else int(multiplicity), + degree=int(degree), + parity=parity, + ) + ) + return tuple(terms) + + +def irreps_dimension(text: str, *, observable: str | None = None) -> int: + """The total number of components a declaration spans.""" + return sum(term.dimension for term in parse_irreps(text, observable=observable)) diff --git a/packages/mace-core/src/mace_core/observables/spec.py b/packages/mace-core/src/mace_core/observables/spec.py new file mode 100644 index 000000000..f6458a0ba --- /dev/null +++ b/packages/mace-core/src/mace_core/observables/spec.py @@ -0,0 +1,320 @@ +"""Declaring an observable, an input, and the derivatives between them. + +Which quantities a legacy model computes is decided by a ladder of string +comparisons on the model class name that sets six boolean flags, plus ten more +``compute_*`` arguments on the forward signatures. Adding one property means +editing several files. Here a property is a row: name it, say what shape it +has, say whether there is one per atom or one per structure, and it is declared. + +Three objects, and they do different jobs: + +``InputSpec`` + something the model is given -- positions, the cell, a magnetic moment, an + electronic temperature. Declared so that a derivative can be taken with + respect to it without new code. + +``ObservableSpec`` + something the model produces and a loss can be written against. + +``DerivativeSpec`` + derived, never declared: the result of asking an observable for its + derivative with respect to an input. Its name and sign come from + :mod:`mace_core.observables.derivatives`. + +``ObservableCatalogue`` holds a set of them and is what the defaults file loads +into. Validation lives there rather than in the individual specs because the +interesting errors are between rows: a derivative asked against an input nobody +declared, or two rows whose derived names collide. +""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from mace_core.observables.derivatives import derivative_name, derivative_sign +from mace_core.observables.grammar import ( + IrrepTerm, + irreps_dimension, + parse_irreps, +) + +__all__ = [ + "NORMALIZATIONS", + "DerivativeRequest", + "DerivativeSpec", + "InputSpec", + "Normalization", + "ObservableCatalogue", + "ObservableSpec", +] + +#: How a target is scaled before it reaches a head and a loss term. The +#: per-observable successor of the legacy scaling registry, whose three entries +#: were a global choice for the whole model. This package stores and validates +#: the value; the head applies it in the output layer and the loss applies the +#: matching term, both reading this one field, so there is never a second +#: scaling mechanism to keep in step with it. +Normalization = Literal["none", "std", "rms"] + +#: The accepted values, for error messages and for callers that enumerate them. +NORMALIZATIONS: tuple[str, ...] = ("none", "std", "rms") + +_SCALAR = (IrrepTerm(multiplicity=1, degree=0, parity="e"),) + + +def _check_name(value: str, kind: str) -> str: + if not value.isidentifier(): + raise ValueError( + f"{kind} name {value!r} is not usable: a name must be a valid " + f"Python identifier, because it is also the key the value is " + f"stored under and part of any derivative name derived from it." + ) + return value + + +class InputSpec(BaseModel): + """Something the model is given, and can be differentiated against. + + ``pos`` and ``cell`` are the two every model has. Anything else is declared + the same way, which is what makes ``d_energy_d_`` reachable for a + new feature without touching code. + """ + + model_config = ConfigDict(extra="forbid", frozen=True) + + name: str + irreps: str + #: ``True`` for one value per atom (positions, magnetic moments), ``False`` + #: for one per structure (the cell, a total charge). This is what decides + #: whether a derivative taken against the input is padded per node or per + #: graph. + per_atom: bool + units: str = Field(min_length=1) + + @model_validator(mode="after") + def _validate(self) -> InputSpec: + _check_name(self.name, "input") + parse_irreps(self.irreps, observable=self.name) + return self + + +class DerivativeRequest(BaseModel): + """A derivative an observable asks for, and the loss settings it carries. + + The name and the sign are not here: they are derived, and letting a + declaration override them would reintroduce the per-consumer naming this + abstraction removes. What a declaration does own is what a loss needs -- + its weight, its normalization, and the unit string to report. + """ + + model_config = ConfigDict(extra="forbid", frozen=True) + + #: The name of the declared input to differentiate against. + wrt: str + default_loss_weight: float = Field(default=1.0, ge=0.0) + normalization: Normalization = "none" + #: Left to the declaration. Deriving it would mean unit algebra over the + #: quantity and the input, which this ticket does not own. + units: str | None = None + + @model_validator(mode="before") + @classmethod + def _accept_bare_name(cls, value: object) -> object: + """``derivatives: [pos, cell]`` is the same as spelling out ``wrt``.""" + if isinstance(value, str): + return {"wrt": value} + return value + + +class DerivativeSpec(BaseModel): + """A derivative, as resolved by the catalogue. Never declared directly.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + name: str + #: The observable being differentiated. + of: str + #: The input it is differentiated against. + wrt: str + #: ``reported = sign * d(of)/d(wrt)``. + sign: int + #: Inherited from the input: a derivative against a per-atom input has one + #: value per atom, whatever the differentiated quantity is. + per_atom: bool + #: Known when the differentiated quantity is a single scalar, in which case + #: the gradient carries the input's irreps. ``None`` otherwise, because the + #: general case is a tensor product and the algebra is not this module's. + irreps: str | None + units: str | None + normalization: Normalization + default_loss_weight: float + + +class ObservableSpec(BaseModel): + """One declared property: what it is, what shape it has, how it is scaled. + + Any atomic or total spherical-tensor property declared here becomes + trainable with no new code: the spec drives the head, the loss term and the + padding classification. + """ + + model_config = ConfigDict(extra="forbid", frozen=True) + + name: str + #: The spherical-tensor shape, in the grammar of + #: :mod:`mace_core.observables.grammar`. + irreps: str + #: ``True`` for one value per atom, ``False`` for one per structure. Every + #: output carries this, which is what lets padding be added and removed + #: without a consumer keeping its own list of which names are which. + per_atom: bool + #: Project convention: eV, Å. + units: str = Field(min_length=1) + normalization: Normalization + default_loss_weight: float = Field(default=1.0, ge=0.0) + #: The derivatives this observable asks for. Naming works for any declared + #: input whether or not it is listed here; listing it is what says the + #: model should compute it. + derivatives: tuple[DerivativeRequest, ...] = () + + @model_validator(mode="after") + def _validate(self) -> ObservableSpec: + _check_name(self.name, "observable") + parse_irreps(self.irreps, observable=self.name) + seen: set[str] = set() + for request in self.derivatives: + if request.wrt in seen: + raise ValueError( + f"observable {self.name!r} asks for the derivative with " + f"respect to {request.wrt!r} twice. Declare it once." + ) + seen.add(request.wrt) + return self + + @property + def dimension(self) -> int: + """The number of components one value of this observable spans.""" + return irreps_dimension(self.irreps, observable=self.name) + + @property + def is_scalar(self) -> bool: + """Whether the declaration is a single ``0e``.""" + return parse_irreps(self.irreps, observable=self.name) == _SCALAR + + def derivative_name(self, wrt: str) -> str: + """The canonical name of this observable's derivative against ``wrt``.""" + return derivative_name(self.name, wrt) + + def derivative_sign(self, wrt: str) -> int: + """The sign that derivative is reported with.""" + return derivative_sign(self.name, wrt) + + def requested_derivatives(self) -> tuple[str, ...]: + """The inputs this observable asked to be differentiated against.""" + return tuple(request.wrt for request in self.derivatives) + + +class ObservableCatalogue(BaseModel): + """A set of declared inputs and observables, validated against each other.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + inputs: tuple[InputSpec, ...] = () + observables: tuple[ObservableSpec, ...] = () + + @model_validator(mode="after") + def _validate(self) -> ObservableCatalogue: + self._reject_duplicates([spec.name for spec in self.inputs], "input") + self._reject_duplicates([spec.name for spec in self.observables], "observable") + declared = {spec.name for spec in self.inputs} + taken = {spec.name for spec in self.observables} + for observable in self.observables: + for request in observable.derivatives: + if request.wrt not in declared: + raise ValueError( + f"observable {observable.name!r} asks for its " + f"derivative with respect to {request.wrt!r}, which is " + f"not a declared input. Declare it under `inputs`, or " + f"use one of {sorted(declared)}." + ) + name = observable.derivative_name(request.wrt) + if name in taken: + raise ValueError( + f"the derivative of {observable.name!r} with respect " + f"to {request.wrt!r} is named {name!r}, which is " + f"already taken. Rename the observable that holds it, " + f"or drop the derivative." + ) + taken.add(name) + return self + + @staticmethod + def _reject_duplicates(names: list[str], kind: str) -> None: + seen: set[str] = set() + for name in names: + if name in seen: + raise ValueError( + f"{kind} {name!r} is declared twice. Every {kind} name is " + f"a key, so it has to be unique." + ) + seen.add(name) + + def input(self, name: str) -> InputSpec: + """The declared input called ``name``.""" + for spec in self.inputs: + if spec.name == name: + return spec + raise KeyError( + f"{name!r} is not a declared input. The declared inputs are " + f"{sorted(spec.name for spec in self.inputs)}." + ) + + def observable(self, name: str) -> ObservableSpec: + """The declared observable called ``name``.""" + for spec in self.observables: + if spec.name == name: + return spec + raise KeyError( + f"{name!r} is not a declared observable. The declared observables " + f"are {sorted(spec.name for spec in self.observables)}." + ) + + def derivative(self, observable: str, wrt: str) -> DerivativeSpec: + """Resolve one derivative, whether or not the observable asked for it. + + Naming and signing are properties of the pair, not of the request, so a + consumer can ask what a derivative *would* be called without the + declaration having listed it. + """ + spec = self.observable(observable) + input_spec = self.input(wrt) + request = next( + (r for r in spec.derivatives if r.wrt == wrt), DerivativeRequest(wrt=wrt) + ) + return DerivativeSpec( + name=spec.derivative_name(wrt), + of=spec.name, + wrt=wrt, + sign=spec.derivative_sign(wrt), + per_atom=input_spec.per_atom, + irreps=input_spec.irreps if spec.is_scalar else None, + units=request.units, + normalization=request.normalization, + default_loss_weight=request.default_loss_weight, + ) + + def requested_derivatives(self) -> tuple[DerivativeSpec, ...]: + """Every derivative the declarations actually asked for.""" + return tuple( + self.derivative(spec.name, request.wrt) + for spec in self.observables + for request in spec.derivatives + ) + + def names(self) -> tuple[str, ...]: + """Every name this catalogue puts on a model output.""" + return tuple(spec.name for spec in self.observables) + tuple( + spec.name for spec in self.requested_derivatives() + ) diff --git a/packages/mace-core/src/mace_core/outputs.py b/packages/mace-core/src/mace_core/outputs.py new file mode 100644 index 000000000..18a340e60 --- /dev/null +++ b/packages/mace-core/src/mace_core/outputs.py @@ -0,0 +1,141 @@ +"""The typed object a model returns instead of a dictionary of tensors. + +Every legacy model ``forward`` returns ``Dict[str, Optional[torch.Tensor]]``, +and between the eleven of them they emit 43 distinct string keys. Nothing +checks a key against anything, so each consumer keeps its own hand-written +list of the names it knows: the ase calculator classifies 22 of the 43 and +returns the other 21 with their padding rows still in them. This module +replaces the dictionary with a typed object whose core fields are named once. + +Six fields are core and everything else goes through :attr:`MACEOutput.extras`. +That split is deliberate rather than a compromise: the electrostatic, magnetic +and dielectric families are large, model-specific and still moving, so pinning +them as attributes would mean a core type that changes shape every time a new +model lands. What classifies an entry of ``extras`` as per-atom or per-graph is +its :class:`~mace_core.observables.ObservableSpec`, not a second key list kept +somewhere else. + +The class is generic over the tensor type. ``mace_core`` imports no framework, +so ``TensorT`` is bound to ``torch.Tensor`` in ``mace_torch``, to ``jax.Array`` +in ``mace_jax``, and to ``numpy.ndarray`` in this package's own tests. The same +pattern carries the kernel Protocol, so it has to work with no framework +installed at all. + +Units follow the project convention: eV, Å, and eV/Å for a force. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field, fields +from typing import Generic, TypeVar + +__all__ = [ + "CORE_FIELD_NAMES", + "FIELD_BY_OBSERVABLE", + "MACEOutput", + "TensorT", +] + +#: The array type a framework binds. Deliberately unbound: a bound of, say, +#: "something with a .shape" would be a structural claim about torch and jax +#: that this package cannot check and does not need. +TensorT = TypeVar("TensorT") + + +@dataclass +class MACEOutput(Generic[TensorT]): + """What a model computed, in one typed object. + + A field left at ``None`` was not computed. That is different from an entry + of ``extras`` that is missing: a core field is part of the type whether or + not this model produces it, while ``extras`` carries only what was asked + for. + + Attributes: + total_energy: Total energy per graph, in eV. Shape ``(n_graphs,)``. + The observable is named ``energy``; see :data:`FIELD_BY_OBSERVABLE`. + node_energies: Per-atom energy, in eV, shape ``(n_atoms,)``. Whether + the isolated-atom reference is included is the model's business + and is stated by the model, not here. + forces: ``-d(energy)/d(positions)``, in eV/Å, shape ``(n_atoms, 3)``. + stress: ``+d(energy)/d(strain) / volume``, in eV/ų, shape + ``(n_graphs, 3, 3)``. + virials: The same derivative before the volume division, in eV, shape + ``(n_graphs, 3, 3)``. + dipole: Total dipole per graph, shape ``(n_graphs, 3)``. + extras: Every other declared observable, keyed by its + :class:`~mace_core.observables.ObservableSpec` name. + + The object is mutable on purpose. Forces and stress are computed by a + derivative engine *around* the model call rather than inside a module's + forward, so something has to fill those fields in after the model returned, + and a frozen object would mean copying the whole thing to do it. + """ + + total_energy: TensorT | None = None + node_energies: TensorT | None = None + forces: TensorT | None = None + stress: TensorT | None = None + virials: TensorT | None = None + dipole: TensorT | None = None + extras: dict[str, TensorT] = field(default_factory=dict) + + def __post_init__(self) -> None: + """Reject an ``extras`` key that a core field already owns. + + Writing ``extras["forces"]`` is otherwise silent: the value is stored, + ``output.forces`` stays ``None``, and every consumer that reads the + field sees nothing. That is the shape of bug this class exists to + remove, so it is an error at construction instead. + """ + shadowed = sorted( + name + for name in self.extras + if FIELD_BY_OBSERVABLE.get(name, name) in CORE_FIELD_NAMES + ) + if shadowed: + raise ValueError( + f"{shadowed} are core fields of MACEOutput and cannot also be " + f"keys of `extras`: a consumer reading the field would see " + f"nothing. Assign them as fields instead." + ) + + def get(self, name: str) -> TensorT | None: + """The value stored under ``name``, or ``None`` if there is none. + + ``name`` is an observable name or a core field name. Without this, a + consumer that iterates over declared observables has to branch on + which of them happen to be core fields, and that branch is the + hand-kept key list this class exists to remove. + """ + field_name = FIELD_BY_OBSERVABLE.get(name, name) + if field_name in CORE_FIELD_NAMES: + return getattr(self, field_name) + return self.extras.get(name) + + def names(self) -> tuple[str, ...]: + """Every name that carries a value, core fields first, then ``extras``. + + A core field holding ``None`` was not computed and is left out, so this + is what the model actually produced rather than what it could produce. + """ + present = [name for name in CORE_FIELD_NAMES if getattr(self, name) is not None] + present.extend(self.extras) + return tuple(present) + + def __contains__(self, name: str) -> bool: + return self.get(name) is not None + + +#: The six fields that are part of the type. Derived from the dataclass rather +#: than written out again, so the two cannot disagree. +CORE_FIELD_NAMES: tuple[str, ...] = tuple( + f.name for f in fields(MACEOutput) if f.name != "extras" +) + +#: The one place an observable's name and its storage field differ. The field +#: says "total" because the type also carries per-atom energies, while the +#: observable is named ``energy`` because that is the name the derivative +#: grammar's special cases are keyed on (``energy`` + positions -> ``forces``). +#: Written down as one entry rather than left to each consumer to remember. +FIELD_BY_OBSERVABLE: dict[str, str] = {"energy": "total_energy"} diff --git a/packages/mace-core/tests/test_observables.py b/packages/mace-core/tests/test_observables.py new file mode 100644 index 000000000..31491360a --- /dev/null +++ b/packages/mace-core/tests/test_observables.py @@ -0,0 +1,410 @@ +"""The declarative observable specification. + +The acceptance bar is "zero new code", so most of these tests declare something +in YAML text and assert what comes out. If any of them needed a new branch in +``mace_core`` to pass, the abstraction would not be doing its job. +""" + +import pytest +from mace_core.observables import ( + DerivativeRequest, + InputSpec, + IrrepsGrammarError, + ObservableCatalogue, + ObservableSpec, + derivative_name, + derivative_sign, + irreps_dimension, + load_catalogue, + load_default_catalogue, + parse_irreps, +) +from pydantic import ValidationError + +# A catalogue with the two inputs every model has, written out so that the +# tests below can add one row at a time to it. +BASE_INPUTS = """ +inputs: + - name: pos + irreps: "1o" + per_atom: true + units: "Å" + - name: cell + irreps: "0e+2e" + per_atom: false + units: "1" +""" + + +def catalogue_from(text, tmp_path): + path = tmp_path / "observables.yaml" + path.write_text(BASE_INPUTS + text, encoding="utf-8") + return load_catalogue(path) + + +# --------------------------------------------------------------------------- +# The irreps grammar +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("text", "dimension"), + [ + ("0e", 1), + ("1o", 3), + ("1e", 3), + ("0e+2e", 6), + ("128x0e+128x1o+128x2e", 128 * (1 + 3 + 5)), + (" 0e + 1o ", 4), + ], +) +def test_the_grammar_accepts_a_sum_of_multiplied_irreps(text, dimension): + assert irreps_dimension(text) == dimension + + +def test_a_term_keeps_its_written_order_and_its_parity(): + terms = parse_irreps("2x1o+0e") + assert [(t.multiplicity, t.degree, t.parity) for t in terms] == [ + (2, 1, "o"), + (1, 0, "e"), + ] + + +@pytest.mark.parametrize("text", ["", " ", "1", "1x", "x0e", "0e+", "1u", "-1o"]) +def test_a_malformed_declaration_is_rejected(text): + with pytest.raises(IrrepsGrammarError): + parse_irreps(text) + + +def test_a_grammar_error_names_the_observable_and_the_grammar(): + with pytest.raises(IrrepsGrammarError) as caught: + parse_irreps("1u", observable="quadrupole") + message = str(caught.value) + assert "quadrupole" in message + assert "'1u'" in message + assert "multiplicity" in message and "parity" in message + + +def test_a_malformed_spec_names_the_observable_and_the_grammar(): + """The same contract through pydantic, which is how a user meets it.""" + with pytest.raises(ValidationError) as caught: + ObservableSpec( + name="quadrupole", + irreps="rank2", + per_atom=True, + units="e*Å^2", + normalization="none", + ) + message = str(caught.value) + assert "quadrupole" in message + assert "'rank2'" in message + assert "Examples: '0e'" in message + + +# --------------------------------------------------------------------------- +# Spec validation +# --------------------------------------------------------------------------- + + +def test_normalization_is_required_and_closed(): + # Both calls are rejected by the type checker as well, which is the point: + # the field is a closed Literal, so a wrong value is caught statically and + # at runtime. The ignores are what let the runtime half be tested. + with pytest.raises(ValidationError): + ObservableSpec(name="q", irreps="0e", per_atom=False, units="eV") # ty: ignore[missing-argument] + with pytest.raises(ValidationError): + ObservableSpec( + name="q", + irreps="0e", + per_atom=False, + units="eV", + normalization="minmax", # ty: ignore[invalid-argument-type] + ) + + +def test_an_unknown_field_is_an_error_rather_than_ignored(): + with pytest.raises(ValidationError): + ObservableSpec( + name="q", + irreps="0e", + per_atom=False, + units="eV", + normalization="none", + weight=3.0, # ty: ignore[unknown-argument] + ) + + +def test_a_name_that_is_not_an_identifier_is_rejected(): + with pytest.raises(ValidationError) as caught: + ObservableSpec( + name="latent charges", + irreps="0e", + per_atom=True, + units="e", + normalization="none", + ) + assert "identifier" in str(caught.value) + + +def test_units_may_not_be_empty(): + with pytest.raises(ValidationError): + ObservableSpec( + name="q", irreps="0e", per_atom=False, units="", normalization="none" + ) + + +# --------------------------------------------------------------------------- +# Derivative naming and signs +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("quantity", "wrt", "name", "sign"), + [ + ("energy", "pos", "forces", -1), + ("energy", "cell", "stress", +1), + ("energy", "magmom", "magforces", -1), + ], +) +def test_the_three_special_cases_keep_their_names_and_signs(quantity, wrt, name, sign): + assert derivative_name(quantity, wrt) == name + assert derivative_sign(quantity, wrt) == sign + + +@pytest.mark.parametrize( + ("quantity", "wrt", "name"), + [ + ("dipole", "pos", "d_dipole_d_pos"), + ("dipole", "cell", "d_dipole_d_cell"), + ("polarizability", "pos", "d_polarizability_d_pos"), + ("energy", "elec_temp", "d_energy_d_elec_temp"), + ("quadrupole", "magmom", "d_quadrupole_d_magmom"), + ], +) +def test_everything_else_follows_the_rule(quantity, wrt, name): + assert derivative_name(quantity, wrt) == name + assert derivative_sign(quantity, wrt) == +1 + + +# --------------------------------------------------------------------------- +# The shipped defaults +# --------------------------------------------------------------------------- + + +def test_the_defaults_declare_energy_and_its_two_derivatives(): + catalogue = load_default_catalogue() + assert catalogue.names() == ("energy", "forces", "stress") + assert [spec.name for spec in catalogue.inputs] == ["pos", "cell"] + + +def test_the_default_forces_row_is_the_negative_position_gradient(): + forces = load_default_catalogue().derivative("energy", "pos") + assert forces.name == "forces" + assert forces.sign == -1 + assert forces.per_atom is True + assert forces.irreps == "1o" + assert forces.units == "eV/Å" + assert forces.default_loss_weight == 100.0 + + +def test_the_default_stress_row_is_the_positive_strain_gradient(): + stress = load_default_catalogue().derivative("energy", "cell") + assert stress.name == "stress" + assert stress.sign == +1 + assert stress.per_atom is False + assert stress.irreps == "0e+2e" + + +# --------------------------------------------------------------------------- +# The two "zero new code" acceptance cases +# --------------------------------------------------------------------------- + + +def test_a_new_rank_two_per_atom_observable_is_a_row_in_yaml(tmp_path): + catalogue = catalogue_from( + """ +observables: + - name: quadrupole + irreps: "0e+2e" + per_atom: true + units: "e*Å^2" + normalization: "rms" + default_loss_weight: 2.5 + derivatives: [pos, cell] +""", + tmp_path, + ) + quadrupole = catalogue.observable("quadrupole") + assert quadrupole.per_atom is True + assert quadrupole.dimension == 6 + assert quadrupole.default_loss_weight == 2.5 + assert catalogue.names() == ( + "quadrupole", + "d_quadrupole_d_pos", + "d_quadrupole_d_cell", + ) + + +def test_a_new_input_feature_makes_its_derivative_declarable(tmp_path): + """`magmom` is the case that pays for the grammar being written over + declared inputs rather than over positions and the cell.""" + catalogue = catalogue_from( + """ + - name: magmom + irreps: "1e" + per_atom: true + units: "muB" + +observables: + - name: energy + irreps: "0e" + per_atom: false + units: "eV" + normalization: "std" + derivatives: + - wrt: magmom + units: "eV/muB" + default_loss_weight: 10.0 +""", + tmp_path, + ) + magforces = catalogue.derivative("energy", "magmom") + assert magforces.name == "magforces" + assert magforces.sign == -1 + assert magforces.per_atom is True + # A magnetic moment is an axial vector, so its conjugate force is too. + assert magforces.irreps == "1e" + assert magforces.default_loss_weight == 10.0 + assert catalogue.names() == ("energy", "magforces") + + +def test_the_bare_string_form_and_the_mapping_form_agree(tmp_path): + shorthand = catalogue_from( + """ +observables: + - name: energy + irreps: "0e" + per_atom: false + units: "eV" + normalization: "std" + derivatives: [pos] +""", + tmp_path, + ) + assert shorthand.observable("energy").derivatives == (DerivativeRequest(wrt="pos"),) + + +# --------------------------------------------------------------------------- +# Catalogue-level validation: the errors that live between rows +# --------------------------------------------------------------------------- + + +def test_a_derivative_against_an_undeclared_input_is_an_error(tmp_path): + with pytest.raises(ValidationError) as caught: + catalogue_from( + """ +observables: + - name: energy + irreps: "0e" + per_atom: false + units: "eV" + normalization: "std" + derivatives: [elec_temp] +""", + tmp_path, + ) + message = str(caught.value) + assert "energy" in message + assert "elec_temp" in message + assert "['cell', 'pos']" in message + + +def test_a_derived_name_may_not_collide_with_a_declared_observable(tmp_path): + with pytest.raises(ValidationError) as caught: + catalogue_from( + """ +observables: + - name: forces + irreps: "1o" + per_atom: true + units: "eV/Å" + normalization: "rms" + - name: energy + irreps: "0e" + per_atom: false + units: "eV" + normalization: "std" + derivatives: [pos] +""", + tmp_path, + ) + assert "'forces'" in str(caught.value) + + +def test_a_name_declared_twice_is_an_error(tmp_path): + with pytest.raises(ValidationError) as caught: + catalogue_from( + """ +observables: + - name: energy + irreps: "0e" + per_atom: false + units: "eV" + normalization: "std" + - name: energy + irreps: "0e" + per_atom: false + units: "eV" + normalization: "none" +""", + tmp_path, + ) + assert "declared twice" in str(caught.value) + + +def test_asking_for_the_same_derivative_twice_is_an_error(tmp_path): + with pytest.raises(ValidationError): + catalogue_from( + """ +observables: + - name: energy + irreps: "0e" + per_atom: false + units: "eV" + normalization: "std" + derivatives: [pos, pos] +""", + tmp_path, + ) + + +def test_an_unknown_observable_or_input_says_what_is_declared(): + catalogue = load_default_catalogue() + with pytest.raises(KeyError) as caught: + catalogue.observable("dipole") + assert "['energy']" in str(caught.value) + with pytest.raises(KeyError) as caught: + catalogue.input("magmom") + assert "['cell', 'pos']" in str(caught.value) + + +def test_a_derivative_can_be_named_without_having_been_requested(): + """Naming is a property of the pair; requesting is what says "compute it".""" + catalogue = ObservableCatalogue( + inputs=[InputSpec(name="pos", irreps="1o", per_atom=True, units="Å")], + observables=[ + ObservableSpec( + name="dipole", + irreps="1o", + per_atom=False, + units="Debye", + normalization="rms", + ) + ], + ) + assert catalogue.names() == ("dipole",) + derived = catalogue.derivative("dipole", "pos") + assert derived.name == "d_dipole_d_pos" + # The differentiated quantity is not a scalar, so the gradient's irreps are + # a tensor product this package deliberately does not compute. + assert derived.irreps is None diff --git a/packages/mace-core/tests/test_outputs.py b/packages/mace-core/tests/test_outputs.py new file mode 100644 index 000000000..6da3aae16 --- /dev/null +++ b/packages/mace-core/tests/test_outputs.py @@ -0,0 +1,113 @@ +"""MACEOutput: the typed replacement for the forward's dictionary of tensors. + +The tests run on numpy arrays throughout. That is the point rather than a +convenience: ``MACEOutput`` is generic over the tensor type so that one class +serves torch, jax and neither, and a test suite that only ever exercised it +with torch would not notice the day it stopped being framework-free. +""" + +import subprocess +import sys + +import numpy as np +import pytest +from mace_core.outputs import CORE_FIELD_NAMES, FIELD_BY_OBSERVABLE, MACEOutput + + +def test_the_six_core_fields_are_the_declared_ones(): + """The type's shape, pinned. `extras` is not one of them.""" + assert CORE_FIELD_NAMES == ( + "total_energy", + "node_energies", + "forces", + "stress", + "virials", + "dipole", + ) + + +def test_an_empty_output_holds_nothing(): + output = MACEOutput[np.ndarray]() + assert output.names() == () + assert output.get("forces") is None + assert "forces" not in output + + +def test_core_fields_round_trip_numpy_arrays(): + forces = np.zeros((4, 3)) + output = MACEOutput(total_energy=np.array([-1.5]), forces=forces) + assert output.get("total_energy") is output.total_energy + assert output.get("forces") is forces + assert output.names() == ("total_energy", "forces") + + +def test_energy_reaches_the_total_energy_field_under_either_name(): + """The one place an observable name and a field name differ.""" + assert FIELD_BY_OBSERVABLE == {"energy": "total_energy"} + output = MACEOutput(total_energy=np.array([2.0])) + assert output.get("energy") is output.get("total_energy") + assert "energy" in output + + +def test_extras_carries_what_the_core_fields_do_not(): + """The escape hatch is load-bearing: 43 legacy keys, six core fields.""" + output = MACEOutput(extras={"latent_charges": np.zeros(4)}) + assert "latent_charges" in output + latent_charges = output.get("latent_charges") + assert latent_charges is not None + assert latent_charges.shape == (4,) + assert output.names() == ("latent_charges",) + + +def test_a_name_nobody_wrote_is_absent_rather_than_an_error(): + """A consumer iterating declared observables asks for names that a given + model did not compute; that is not an error, it is a `None`.""" + output = MACEOutput(extras={"charges": np.zeros(2)}) + assert output.get("magforces") is None + assert "magforces" not in output + + +def test_a_core_field_left_none_is_not_reported_as_present(): + output = MACEOutput(forces=None, extras={}) + assert output.names() == () + + +def test_two_outputs_do_not_share_an_extras_dictionary(): + """`extras` has a default factory; a shared mutable default would make one + model's outputs appear in another's.""" + first = MACEOutput[np.ndarray]() + second = MACEOutput[np.ndarray]() + first.extras["dipole_moment"] = np.zeros(3) + assert second.extras == {} + + +def test_the_class_is_generic_over_the_tensor_type(): + """Subscripting has to work with no framework installed at all: the kernel + Protocol reuses this pattern, so it cannot depend on torch being there.""" + assert MACEOutput[np.ndarray] is not None + assert isinstance(MACEOutput[np.ndarray](), MACEOutput) + + +@pytest.mark.parametrize("framework", ["torch", "jax", "e3nn"]) +def test_importing_mace_core_imports_no_framework(framework): + """Run in a fresh interpreter on purpose. Asserting this in-process would + pass whenever some earlier test in the session had already imported torch, + which is the case the assertion exists for.""" + probe = ( + "import sys, mace_core, mace_core.observables, mace_core.outputs\n" + f"assert {framework!r} not in sys.modules, " + f"'importing mace_core pulled in {framework}'\n" + ) + subprocess.run([sys.executable, "-c", probe], check=True) + + +@pytest.mark.parametrize("key", ["forces", "energy", "total_energy", "dipole"]) +def test_extras_may_not_shadow_a_core_field(key): + """A value written into `extras` under a core field's name would be stored + and never read: the field it shadows stays `None`.""" + with pytest.raises(ValueError, match="cannot also be"): + MACEOutput(extras={key: np.zeros(3)}) + + +def test_a_name_that_is_not_a_core_field_is_fine_in_extras(): + assert MACEOutput(extras={"node_energy": np.zeros(3)}).names() == ("node_energy",) diff --git a/tests/architecture/observable_coverage.py b/tests/architecture/observable_coverage.py new file mode 100644 index 000000000..84e406ebc --- /dev/null +++ b/tests/architecture/observable_coverage.py @@ -0,0 +1,483 @@ +"""Where every legacy model output goes in the declarative specification. + +The frozen models emit 43 distinct keys from their ``forward`` methods. This +module says, for each one, what it becomes: a declared observable, a derivative +of one, or a row saying explicitly that it is not an observable and naming the +mechanism that owns it instead. A key with no row fails the test beside this +file, so a key added to a legacy forward cannot slip through unclassified. + +The 43 are **not listed here**. They are read out of ``mace/modules/models.py`` +and ``mace/modules/extensions.py`` by ``tests/golden/surface_scan.py``, which +follows keys assigned onto the returned object as well as dict literals -- the +self-consistent model assigns its three diagnostics after the fact, so an +extraction that stopped at return literals would stop at 40 and lose them +silently. + +Two pieces of metadata are likewise derived rather than retyped. The golden +harness already declares a kind and a unit for every one of these keys, so the +per-atom/per-graph classification and the unit string come from there and the +rows below carry only what a schema cannot know: the spherical-tensor shape, +and the decision. + +One convention worth stating once. A rank-2 Cartesian quantity -- a stress, a +virial, a polarizability -- is stored by the legacy models as a full 3x3, nine +numbers, while its irreps declaration ``0e+2e`` spans six. That is not a +contradiction: the declaration says what the quantity *is* under rotation, and +the 3x3 is a layout with three redundant entries. Reconciling the two is the +head's job, and it is written down here because the mismatch otherwise looks +like an error in these rows. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Union + +from tests.golden import harness, surface_scan + +REPO_ROOT = Path(__file__).resolve().parents[2] + +#: The two files the ticket's count is defined over. The scanner discovers +#: forwards across the whole package -- the LAMMPS and torchsim wrappers define +#: one too -- and those are deployment layers with their own tickets, so this +#: surface is named rather than discovered. +LEGACY_MODEL_SOURCES = ( + REPO_ROOT / "mace" / "modules" / "models.py", + REPO_ROOT / "mace" / "modules" / "extensions.py", +) + +#: The ase calculator, layer (b) of the three-layer surface. +LEGACY_CALCULATOR_SOURCES = (REPO_ROOT / "mace" / "calculators" / "mace.py",) + +PER_ATOM_KINDS = frozenset( + { + harness.PER_ATOM_SCALAR, + harness.PER_ATOM_VECTOR, + harness.PER_ATOM_TENSOR, + harness.PER_ATOM_MATRIX, + } +) +PER_GRAPH_KINDS = frozenset( + { + harness.GRAPH_SCALAR, + harness.GRAPH_VECTOR, + harness.GRAPH_TENSOR, + harness.GRAPH_ARRAY, + } +) + +#: The inputs a derivative row may be taken against. ``pos`` and ``cell`` are +#: the two every model has and are declared in the shipped defaults; ``magmom`` +#: is the third the frozen tree actually differentiates against, and is +#: declared by whichever configuration turns the magnetic model on. +DECLARED_INPUTS = frozenset({"pos", "cell", "magmom"}) + + +@dataclass(frozen=True) +class Spec: + """The key becomes a declared observable. + + An irreps declaration is two independent facts, and lumping them together + is what made a third of these rows look unanswerable. What the quantity is + under rotation is a property of the quantity. How wide it is can be a + property of the *model*: the number of readout layers, the maximum multipole + order, whether an anisotropic readout was declared. A row that names both is + complete even when it cannot write a single literal, because the model fills + the rest in when it is declared, exactly as the ticket's own + ``128x0e+128x1o+128x2e`` example does. + + So a row states either ``irreps`` or ``irreps_pattern`` plus ``set_by``, and + there is no third state. Nothing here is deferred to another ticket: every + one of the 43 was worked out, and a future row that cannot be fails the test + rather than acquiring a TODO. + + Attributes: + irreps: The declaration, when the quantity fixes it outright. + irreps_pattern: The declaration with the model-dependent part named. + Prose, not the grammar: it is read by people, and the grammar would + have to grow placeholders to hold it. + set_by: What the model supplies. Required with ``irreps_pattern``, + because "it depends on the model" is not an answer until it says on + what. + note: Anything about the row a reader would otherwise have to rederive. + """ + + irreps: str | None = None + irreps_pattern: str | None = None + set_by: str = "" + note: str = "" + + +@dataclass(frozen=True) +class Derivative: + """The key is a derivative of another quantity, and is renamed by the rule. + + Attributes: + of: The differentiated quantity. + wrt: The declared input it is differentiated against. + sign: The sign the frozen tree reports, so that + ``legacy value = sign * d(of)/d(wrt)``. Read off the source and, + where the note says so, measured. It is stated per row rather than + taken from the rule because a row that silently agreed with the + rule and a row nobody checked look identical. The test asserts the + two agree, with no way to annotate a disagreement: a pair the rule + gets wrong is either a misclassified row or a real gap in the + grammar, and both have to be resolved rather than recorded. + note: Why the row is worth a second look, where it is. + """ + + of: str + wrt: str + sign: int + note: str = "" + + +@dataclass(frozen=True) +class Drop: + """The key is not an observable. ``reason`` names what owns it instead.""" + + reason: str + + +Disposition = Union[Spec, Derivative, Drop] + + +#: Every key the two frozen model modules emit, and what it becomes. +DISPOSITIONS: dict[str, Disposition] = { + # --- energies ---------------------------------------------------------- + "energy": Spec(irreps="0e"), + "node_energy": Spec(irreps="0e"), + "interaction_energy": Spec(irreps="0e"), + "les_energy": Spec(irreps="0e"), + "electrostatic_energy": Spec(irreps="0e"), + "electron_energy": Spec(irreps="0e"), + "fermi_level": Spec(irreps="0e"), + "contributions": Spec( + irreps_pattern="<2 + num_interactions>x0e", + set_by="the number of energy terms the model sums", + note=( + "the energy decomposed into its terms. `energies = [e0, " + "pair_energy]` and one entry is appended per interaction layer " + "(mace/modules/models.py:361, :398), so the extent is the " + "isolated-atom reference plus the ZBL pair term plus one per " + "layer. Measured on the tiny_mace anchor, which has two layers: " + "shape (1, 4)." + ), + ), + # --- charges and potentials ------------------------------------------- + "charges": Spec(irreps="0e"), + "latent_charges": Spec(irreps="0e"), + "spins": Spec( + irreps="0e", + note=( + "the per-atom spin population, a scalar. Not to be confused with a " + "magnetic moment, which is an axial vector ('1e')." + ), + ), + "electrostatic_potentials": Spec(irreps="0e"), + # --- dipoles ----------------------------------------------------------- + "dipole": Spec(irreps="1o"), + "atomic_dipoles": Spec(irreps="1o"), + "latent_dipoles": Spec(irreps="1o"), + # --- rank-2 quantities ------------------------------------------------- + "polarizability": Spec(irreps="0e+2e"), + "polarizability_sh": Spec( + irreps="0e+2e", + note=( + "already the spherical form, six components. The Cartesian " + "`polarizability` above is the same quantity in the 3x3 layout, " + "which is what makes the pair a useful check on the head." + ), + ), + "virials": Spec( + irreps="0e+2e", + note=( + "the NEGATED cell derivative, and the sign is not a detail: " + "`compute_forces_virials` computes the stress from the raw " + "gradient and negates the virial only in its return statement " + "(mace/modules/utils.py:107-115), so the frozen tree reports " + "`virials = -dE/dstrain` and `stress = +dE/dstrain / V`. Measured " + "on the tiny_scaleshift anchor over an fcc cell: " + "max|stress * V + virials| = 1.2e-35, exactly zero, while " + "max|stress * V - virials| = 6.5e-3. The two therefore differ by a " + "sign as well as by the volume. `stress` is the row that carries " + "the derivative naming; `virials` stays an observable of its own " + "because the spec has no way to say 'the same derivative in " + "another normalization and the opposite sign', which is a question " + "for the head that produces them." + ), + ), + "atomic_virials": Spec(irreps="0e+2e"), + "atomic_stresses": Spec(irreps="0e+2e"), + # --- families whose shape is a model hyperparameter --------------------- + "density_coefficients": Spec( + irreps_pattern="two concatenated ladders 0e+1o+2e+...+", + set_by="atomic_multipoles_max_l", + note=( + "the spin-summed electron density in the model's multipole basis. " + "`self.charges_irreps = 2 * o3.Irreps.spherical_harmonics(" + "atomic_multipoles_max_l)` (mace/modules/extensions.py:789), and " + "in e3nn that product concatenates the ladder twice rather than " + "doubling each multiplicity, so the layout is two ladders end to " + "end even though the content simplifies to 2x0e+2x1o+2x2e+... " + "Dimension 8, 18, 32 for max_l 1, 2, 3." + ), + ), + "spin_density": Spec( + irreps_pattern="two concatenated ladders 0e+1o+2e+...+", + set_by="atomic_multipoles_max_l", + note=( + "the alpha minus beta difference of the same basis as " + "`density_coefficients` (mace/modules/extensions.py:1250), so it " + "carries identical irreps." + ), + ), + "spin_charge_density": Spec( + irreps_pattern=( + "<2 spin channels> x two concatenated ladders " + "0e+1o+2e+...+" + ), + set_by="atomic_multipoles_max_l, and the fixed pair of spin channels", + note=( + "the density before the spin channels are summed or subtracted: " + "`spin_charge_density.view(shape[0], 2, -1)` " + "(mace/modules/extensions.py:1094). `density_coefficients` is its " + "sum over that axis and `spin_density` its difference." + ), + ), + "fukui_functions": Spec( + irreps="2x0e", + note=( + "two scalars per atom, one per spin channel, and not model " + "dependent at all: `self.fukui_source_map` is a readout whose " + "output irreps are the literal `o3.Irreps(\"2x0e\")` " + "(mace/modules/extensions.py:838-842). They are added to the l=0 " + "component of each spin channel of the density." + ), + ), + "BEC": Spec( + irreps_pattern="<1 or 2>x(0e+1e+2e)", + set_by=( + "whether the model passes latent dipoles to LES alongside the " + "latent charges, which adds a second channel" + ), + note=( + "NOT the same quantity as `dmu_dr`, which is the tempting reading " + "and the wrong one. LES builds a polarization from its own latent " + "charges with the mean removed and an epsilon^(1/2) factor, takes " + "it through a Berry phase under periodic boundary conditions, and " + "differentiates *that* against the positions, dephasing and " + "projecting the result by the cell " + "(les/module/bec.py:56-93). `dmu_dr` differentiates the " + "dielectric model's `dipole` readout, a different quantity in a " + "different gauge and a different unit (e against Debye/Ang). " + "Measured on the tiny_maceles anchor, the shape is " + "(n_atoms, 2, 3, 3), not (n_atoms, 3, 3): the leading pair is the " + "charge-derived tensor and the latent-dipole-derived one, and " + "whether that axis is there at all depends on whether the model " + "passes latent dipoles. Each 3x3 block is a general rank-2 tensor " + "rather than a symmetric one (max|B - B^T| = 0.20 on that anchor), " + "so a block is 0e+1e+2e and the channel axis is the multiplicity " + "DEP-1a has to fix." + ), + ), + "latent_kappas": Spec( + irreps="0e", + note=( + "one scalar per atom. The LES signature documents it as " + "[n_atoms, ] (les/les.py:89) and `les_kappa_readouts` is a scalar " + "readout; measured (3,) on the tiny_maceles anchor." + ), + ), + "latent_alphas": Spec( + irreps_pattern="0e, or 0e+2e", + set_by="whether the model declares the anisotropic alpha readouts", + note=( + "an atomic polarizability, isotropic by default and anisotropic " + "when `use_induced_dipoles` brings in `les_alpha_2e_readouts`. " + "That branch reads out a spherical 0e+2e and expands it through " + "`spherical_to_cartesian` (mace/modules/extensions.py:485-494), " + "then squares it as A A^T, which is symmetric, so the anisotropic " + "form is 0e+2e and not the general 0e+1e+2e. LES accepts both " + "shapes explicitly (les/les.py:133-136). Measured (3,) on the " + "tiny_maceles anchor, which is the isotropic path." + ), + ), + "latent_quads": Spec( + irreps="2e", + note=( + "an atomic quadrupole. Stored as a Cartesian 3x3 " + "(les/les.py:88), but the model subtracts the trace explicitly " + "(mace/modules/extensions.py:549-552), so it is symmetric and " + "traceless, which is exactly 2e and five numbers rather than nine. " + "Measured on the tiny_maceles anchor: max|A - A^T| = 2.7e-20 and " + "|trace| = 5.4e-20." + ), + ), + # --- derivatives ------------------------------------------------------- + "forces": Derivative( + of="energy", + wrt="pos", + sign=-1, + note="mace/modules/utils.py:115 returns `-1 * forces`.", + ), + "stress": Derivative( + of="energy", + wrt="cell", + sign=+1, + note=( + "positive, and only because the stress is built from the raw " + "gradient before the virial is negated. See the `virials` row: the " + "two are not the same sign." + ), + ), + "magforces": Derivative( + of="energy", + wrt="magmom", + sign=-1, + note=( + "mace/modules/utils.py returns `-mag_forces` from both " + "`compute_forces_virials_magforces` and `compute_forces_magforces`." + ), + ), + "dmu_dr": Derivative( + of="dipole", + wrt="pos", + sign=+1, + note=( + "mace/modules/models.py:1169 differentiates the model's own " + "`dipole` key, and `compute_dielectric_gradients` returns the " + "gradient with no sign flip." + ), + ), + "dalpha_dr": Derivative( + of="polarizability", + wrt="pos", + sign=+1, + note=( + "the Cartesian polarizability flattened to nine components " + "(mace/modules/models.py:1173-1176), differentiated the same way " + "as `dmu_dr`." + ), + ), + # --- not observables --------------------------------------------------- + "hessian": Drop( + reason=( + "the second derivative of the energy with respect to the " + "positions, and the derivative grammar is first order by design: " + "it names d(quantity)/d(input) for a declared quantity and a " + "declared input, and all three of its special cases are first " + "order. Reading it instead as the first derivative of the forces " + "does not rescue it, and the sign is how that shows: " + "`compute_hessians_vmap` differentiates `-1 * forces` " + "(mace/modules/utils.py:168), so the key holds +d2E/dpos2, which " + "is MINUS d(forces)/d(pos). Measured on the tiny_scaleshift anchor " + "against a central difference of the forces, " + "max|hessian[:, 0] + dF/dx| = 1.8e-10 against " + "max|hessian[:, 0] - dF/dx| = 0.32. It is a real output and it is " + "not lost: the derivative engine owns second derivatives, the same " + "way the export path owns `edge_forces`." + ) + ), + "displacement": Drop( + reason=( + "the symmetric strain handle the cell derivative is taken against. " + "It is created as zeros to attach the cell to the autograd graph " + "and nothing ever writes to it, so its value is identically zero " + "on every structure. It belongs to the derivative engine, not to " + "the output surface." + ) + ), + "edge_forces": Drop( + reason=( + "indexed by the neighbour list, so it is neither per-atom nor " + "per-graph and the spec's classification cannot express it. It is " + "the per-edge decomposition the LAMMPS pair style sums into its " + "virial, and belongs to the export path." + ) + ), + "node_feats": Drop( + reason=( + "the backbone's node features: what every head reads, not what one " + "produces. No dataset carries a target for it, so declaring it " + "would make the head-creation check unsatisfiable. The evaluation " + "CLI exposes it as `descriptors`, which is CLI-1's layer." + ) + ), + "external_field": Drop( + reason=( + "an input echoed back, which is how the golden harness classifies " + "it too. It reaches v1 as a declared input feature, which is also " + "what makes a derivative against it expressible." + ) + ), + "total_charge": Drop( + reason=( + "an input echoed back, like `external_field`. Declared as an input " + "feature rather than as an observable." + ) + ), + "charges_history": Drop( + reason=( + "the iterates of the electrostatic self-consistency loop: how a " + "fixed point was reached, not the fixed point. Run telemetry, and " + "it belongs in the stage's log." + ) + ), + "scf_energy_history": Drop( + reason="run telemetry of the self-consistent loop, like `charges_history`." + ), + "scf_steps": Drop( + reason="run telemetry of the self-consistent loop, like `charges_history`." + ), + "equilibrated_magmom": Drop( + reason=( + "an input the self-consistent stage converged, exposed through the " + "stage result rather than as a model observable. Note the golden " + "harness classifies it as an output channel, because from where it " + "sits it is one; the distinction is which mechanism owns it." + ) + ), +} + + +def legacy_model_keys() -> set[str]: + """The keys the two frozen model modules can return, read from the source.""" + return surface_scan.scan_model_surface(list(LEGACY_MODEL_SOURCES)).all_keys + + +def legacy_calculator_keys() -> set[str]: + """The keys the ase calculators can write into ``results``.""" + return surface_scan.scan_calculator_surface( + list(LEGACY_CALCULATOR_SOURCES) + ).all_keys + + +def legacy_eval_keys() -> set[str]: + """The names the evaluation CLI writes onto its structures, unprefixed.""" + scan, _stores = surface_scan.scan_eval_surface() + return scan.all_keys + + +def channel_of(key: str) -> harness.Channel | None: + """The golden harness's declaration for ``key`` on the model surface.""" + name = harness.resolve_channel(key, harness.SURFACE_MODEL) + return None if name is None else harness.CHANNELS[name] + + +def per_atom_of(key: str) -> bool | None: + """Whether ``key`` is per-atom, derived from its harness kind. + + ``None`` for the kinds that are neither -- a per-edge quantity, a hessian, + a gradient whose atom axis is not the leading one. A key that lands here + cannot be a plain observable, and the test asserts exactly that. + """ + channel = channel_of(key) + if channel is None: + return None + if channel.kind in PER_ATOM_KINDS: + return True + if channel.kind in PER_GRAPH_KINDS: + return False + return None diff --git a/tests/architecture/test_observable_completeness.py b/tests/architecture/test_observable_completeness.py new file mode 100644 index 000000000..f8705038e --- /dev/null +++ b/tests/architecture/test_observable_completeness.py @@ -0,0 +1,265 @@ +"""Every legacy model output is accounted for by the declarative spec. + +Two failure modes this suite is built against. The first is a key nobody +classified, which is how the ase calculator ended up returning 21 of the 43 +model keys with their padding rows still in them. The second, quieter one is a +test that measures nothing: an extraction that finds no keys reports perfect +coverage, so the counts are asserted as numbers rather than left to a set +comparison that an empty set would satisfy. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest +from mace_core.observables import ObservableSpec, derivative_name, derivative_sign + +from tests.architecture.observable_coverage import ( + DECLARED_INPUTS, + DISPOSITIONS, + LEGACY_MODEL_SOURCES, + Derivative, + Drop, + Spec, + channel_of, + legacy_calculator_keys, + legacy_eval_keys, + legacy_model_keys, + per_atom_of, +) +from tests.golden import surface_scan + +SURFACE_DOC = ( + Path(__file__).resolve().parents[2] / "docs" / "reforge" / "output_surface.md" +) + + +def test_the_scan_resolves_every_write_it_finds(): + """The honesty check. A write whose key cannot be computed would shrink the + surface silently, so it fails here instead.""" + scan = surface_scan.scan_model_surface(list(LEGACY_MODEL_SOURCES)) + assert surface_scan.unexplained(scan) == [] + + +def test_the_model_forward_surface_is_forty_three_keys(): + keys = legacy_model_keys() + assert len(keys) == 43, sorted(keys) + + +def test_every_legacy_key_has_exactly_one_disposition(): + keys = legacy_model_keys() + missing = sorted(keys - set(DISPOSITIONS)) + assert not missing, ( + f"{missing} are returned by a frozen model forward and have no row in " + f"observable_coverage.DISPOSITIONS. Add a Spec, a Derivative or a Drop " + f"row; an unclassified output is one nothing downstream can pad, " + f"unpad or train." + ) + stale = sorted(set(DISPOSITIONS) - keys) + assert not stale, ( + f"{stale} have a disposition row and are no longer returned by any " + f"frozen model forward. Remove the rows." + ) + + +@pytest.mark.parametrize( + "key", sorted(k for k, d in DISPOSITIONS.items() if isinstance(d, Spec)) +) +def test_a_spec_row_builds_a_valid_observable_spec(key): + """The classification the padding depends on has to exist for every row, + and a row states its irreps either outright or with the model-dependent + part named. There is no third state: no row is allowed to defer the + question to a later ticket.""" + row = DISPOSITIONS[key] + per_atom = per_atom_of(key) + assert per_atom is not None, ( + f"{key!r} is declared as an observable but the golden harness gives it " + f"a kind that is neither per-atom nor per-graph. It cannot be a Spec " + f"row; make it a Derivative or a Drop." + ) + assert bool(row.irreps) != bool(row.irreps_pattern), ( + f"{key!r} must state exactly one of `irreps` and `irreps_pattern`. " + f"Neither means the row was never worked out, and both means it is " + f"unclear which one a reader should believe." + ) + if row.irreps_pattern: + assert row.set_by.strip(), ( + f"{key!r} says its shape depends on the model without saying on " + f"what. Name the parameter in `set_by`." + ) + return + channel = channel_of(key) + assert channel is not None + spec = ObservableSpec( + name=key, + irreps=row.irreps, + per_atom=per_atom, + # The legacy unit, as the golden harness records it. This ticket does + # not canonicalise unit strings, so `Ang` is left as it is written + # there rather than rewritten to `Å`. + units=channel.unit, + # Neutral: which normalization each family wants is decided by the + # ticket that builds its head, not by this coverage table. + normalization="none", + ) + assert spec.per_atom == per_atom + assert spec.dimension >= 1 + + +def test_no_observable_row_defers_its_irreps_to_another_ticket(): + """The whole set was worked out, and this asserts it stays that way. A row + that cannot say what its irreps are is a row nobody has read the model + for, and it should fail here rather than sit as a TODO.""" + unresolved = sorted( + key + for key, row in DISPOSITIONS.items() + if isinstance(row, Spec) and not row.irreps and not row.irreps_pattern + ) + assert not unresolved, unresolved + + +@pytest.mark.parametrize( + "key", sorted(k for k, d in DISPOSITIONS.items() if isinstance(d, Derivative)) +) +def test_a_derivative_row_resolves_through_the_rule(key): + row = DISPOSITIONS[key] + assert row.wrt in DECLARED_INPUTS, ( + f"{key!r} is differentiated against {row.wrt!r}, which is not one of " + f"the declared inputs {sorted(DECLARED_INPUTS)}." + ) + parent = DISPOSITIONS.get(row.of) + assert isinstance(parent, (Spec, Derivative)), ( + f"{key!r} is the derivative of {row.of!r}, which is not itself a " + f"declared observable or a derivative of one. A derivative chain has " + f"to ground out in something declared." + ) + name = derivative_name(row.of, row.wrt) + assert name.isidentifier() + derived_sign = derivative_sign(row.of, row.wrt) + assert row.sign == derived_sign, ( + f"{key!r} is reported with sign {row.sign:+d} by the frozen tree and " + f"the rule derives {derived_sign:+d}. There is no way to annotate that " + f"away, deliberately: either the row differentiates the wrong thing, " + f"or the grammar has a real gap. `hessian` was the first candidate for " + f"such a gap and turned out to be the former: it is a second " + f"derivative of the energy, not a first derivative of the forces." + ) + + +def test_the_three_named_derivatives_keep_their_legacy_names(): + """forces, stress and magforces are the pairs that have a name of their own, + and the two negated ones are the two the frozen tree negates.""" + assert derivative_name("energy", "pos") == "forces" + assert derivative_sign("energy", "pos") == -1 + assert derivative_name("energy", "cell") == "stress" + assert derivative_sign("energy", "cell") == +1 + assert derivative_name("energy", "magmom") == "magforces" + assert derivative_sign("energy", "magmom") == -1 + for name in ("forces", "stress", "magforces"): + assert isinstance(DISPOSITIONS[name], Derivative) + + +def test_the_renamed_derivatives_are_renamed_and_not_lost(): + """Two legacy spellings become their rule-derived names. A rename is a + breaking change and is fine; losing one is not. + + `BEC` is deliberately absent. It looks like a third, and treating it as one + would have merged it into `dmu_dr`; see its row for why the two are + different quantities. `hessian` is absent for a different reason: it is a + second derivative and carries a Drop row. + """ + renamed = { + "dmu_dr": "d_dipole_d_pos", + "dalpha_dr": "d_polarizability_d_pos", + } + for key, expected in renamed.items(): + row = DISPOSITIONS[key] + assert derivative_name(row.of, row.wrt) == expected + assert isinstance(DISPOSITIONS["BEC"], Spec) + + +def test_no_two_derivative_rows_resolve_to_one_name(): + """Two legacy keys collapsing onto a single canonical name is either the + abstraction working or a silent merge of two different quantities. It has + to be looked at rather than discovered later, so it fails here.""" + seen: dict[str, str] = {} + for key, row in DISPOSITIONS.items(): + if not isinstance(row, Derivative): + continue + name = derivative_name(row.of, row.wrt) + assert name not in seen, ( + f"{key!r} and {seen[name]!r} both resolve to {name!r}. If they are " + f"one quantity, say so and drop one row; if they are not, one of " + f"them is not a derivative of what it claims." + ) + seen[name] = key + + +def test_a_key_the_classification_cannot_place_is_never_a_spec_row(): + """`edge_forces` is per-edge and a hessian is a square over 3N degrees of + freedom. Neither is per-atom or per-graph, so neither can be an observable + under a classification that only has those two.""" + for key, row in DISPOSITIONS.items(): + if per_atom_of(key) is None: + assert isinstance(row, (Derivative, Drop)), key + + +@pytest.mark.parametrize( + "key", sorted(k for k, d in DISPOSITIONS.items() if isinstance(d, Drop)) +) +def test_a_drop_row_says_what_owns_the_key_instead(key): + reason = DISPOSITIONS[key].reason + assert reason.strip() + assert len(reason.split()) >= 8, ( + f"the reason for dropping {key!r} is too short to be a decision " + f"anybody can review later: {reason!r}" + ) + + +def test_the_scf_trio_is_reached_at_all(): + """The three keys an extraction that stops at return literals never sees. + They are assigned onto the output after it is built.""" + keys = legacy_model_keys() + for key in ("scf_energy_history", "scf_steps", "equilibrated_magmom"): + assert key in keys + assert isinstance(DISPOSITIONS[key], Drop) + + +# --------------------------------------------------------------------------- +# The three-layer surface +# --------------------------------------------------------------------------- + + +def documented_counts() -> dict[str, tuple[int, int]]: + """The per-layer `(keys, new here)` pairs the surface document records.""" + rows = {} + for line in SURFACE_DOC.read_text(encoding="utf-8").splitlines(): + match = re.match(r"^\|\s*\((a|b|c)\)[^|]*\|(.*)$", line) + if match is None: + continue + cells = [cell.strip() for cell in match.group(2).split("|")] + numbers = [int(cell) for cell in cells if cell.isdigit()] + assert len(numbers) == 2, line + rows[match.group(1)] = (numbers[0], numbers[1]) + return rows + + +def test_the_surface_document_records_the_derived_counts(): + model = legacy_model_keys() + calculator = legacy_calculator_keys() + evaluation = legacy_eval_keys() + derived = { + "a": (len(model), len(model)), + "b": (len(calculator), len(calculator - model)), + "c": (len(evaluation), len(evaluation - model - calculator)), + } + assert documented_counts() == derived + assert derived == {"a": (43, 43), "b": (31, 15), "c": (13, 3)} + + +def test_the_union_is_sixty_one_names(): + union = legacy_model_keys() | legacy_calculator_keys() | legacy_eval_keys() + assert len(union) == 61 + assert "**61**" in SURFACE_DOC.read_text(encoding="utf-8") From 0065314f3d7e3b9cd494cd4d10ca6a6d97dfe238 Mon Sep 17 00:00:00 2001 From: Alejandro Acosta <127198532+aacostadiaz@users.noreply.github.com> Date: Tue, 15 Sep 2026 10:10:46 +0100 Subject: [PATCH 02/12] Keep the whole-tree sweeps out of tests/architecture 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 #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. --- .github/gitlab/ci.yml | 14 ++++++++++++++ .github/workflows/nightly.yaml | 8 ++++++++ 2 files changed, 22 insertions(+) diff --git a/.github/gitlab/ci.yml b/.github/gitlab/ci.yml index caa1493ed..00a236165 100644 --- a/.github/gitlab/ci.yml +++ b/.github/gitlab/ci.yml @@ -113,10 +113,24 @@ stages: # and the GPU cases are the by-hand reproduction recipe recorded in that # file's module docstring. Removing this clause is not the way to fix that; # a separate, non-gating benchmark job would be. + # `tests/architecture` is ignored, and the reason is an install rather than + # a marker. That suite needs BOTH trees importable: the import contracts + # resolve each root package on the filesystem, and the observable + # completeness test imports `mace_core`. This job installs the legacy + # distribution only, so the directory fails at *collection*, before any + # marker expression gets to deselect it -- a marker on those tests would + # therefore not help. It is not lost coverage: `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 + # it is gpu-marked. + # + # Any later suite that needs the v1 stack belongs on this line too. + # `tests/parity` is the next one. - > MACE_REQUIRE_CAPS="$REQUIRE_CAPS" MACE_CI_ALLOW_NETWORK=1 python -m pytest tests + --ignore=tests/architecture -m "($MARKER_EXPR) and not benchmark" -n 2 --timeout 1200 -ra diff --git a/.github/workflows/nightly.yaml b/.github/workflows/nightly.yaml index 62d03c580..fe4d29ca5 100644 --- a/.github/workflows/nightly.yaml +++ b/.github/workflows/nightly.yaml @@ -316,6 +316,11 @@ jobs: - uses: ./.github/actions/run-tests with: tests: tests + # See the same flag in .github/gitlab/ci.yml: `tests/architecture` + # needs both trees importable and this job installs the legacy one, + # so the directory fails at collection rather than being deselected. + # It is gated per pull request by ci-core's `architecture` job. + extra-args: --ignore=tests/architecture markers: not gpu and not benchmark require-caps: network,cueq,polar,les,torchsim,schedulefree,magnetic,wandb allow-network: "true" @@ -504,6 +509,9 @@ jobs: - uses: ./.github/actions/run-tests with: tests: tests + # Same reason as coverage-full above: this job installs `dev` alone, + # so `tests/architecture` cannot be collected here. + extra-args: --ignore=tests/architecture allow-network: "true" store-durations: "true" timeout: "1200" From 6114c25b8d7f0122ab77483d320e563090b3cd53 Mon Sep 17 00:00:00 2001 From: Alejandro Acosta <127198532+aacostadiaz@users.noreply.github.com> Date: Tue, 15 Sep 2026 10:24:03 +0100 Subject: [PATCH 03/12] Close two gaps found in review of the observable coverage 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. --- .github/workflows/ci-extensions.yaml | 7 ++++ packages/mace-core/tests/test_outputs.py | 7 +++- tests/architecture/observable_coverage.py | 34 +++++++++++++++++- .../test_observable_completeness.py | 35 ++++++++++++++++++- 4 files changed, 80 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci-extensions.yaml b/.github/workflows/ci-extensions.yaml index 0f265b201..e690e5832 100644 --- a/.github/workflows/ci-extensions.yaml +++ b/.github/workflows/ci-extensions.yaml @@ -263,5 +263,12 @@ jobs: - uses: ./.github/actions/run-tests with: tests: tests + # See the same flag in .github/gitlab/ci.yml: `tests/architecture` + # needs both trees importable and this job installs the legacy one, + # so the directory fails at collection rather than being deselected + # by the marker expression. Gated per pull request by ci-core's + # `architecture` job. This job is paths-filtered, so the breakage + # would have waited for the first pull request touching backends. + extra-args: --ignore=tests/architecture markers: cueq and not gpu and not polar require-caps: cueq diff --git a/packages/mace-core/tests/test_outputs.py b/packages/mace-core/tests/test_outputs.py index 6da3aae16..b4bb927f5 100644 --- a/packages/mace-core/tests/test_outputs.py +++ b/packages/mace-core/tests/test_outputs.py @@ -110,4 +110,9 @@ def test_extras_may_not_shadow_a_core_field(key): def test_a_name_that_is_not_a_core_field_is_fine_in_extras(): - assert MACEOutput(extras={"node_energy": np.zeros(3)}).names() == ("node_energy",) + """Deliberately not a near-miss of a field name. `node_energy` would be a + bad example here: it is the legacy spelling of `node_energies`, so using it + would read as an endorsement of putting the per-atom energy in `extras` + while the field that owns it stays `None`. v1 renames that key instead.""" + out = MACEOutput(extras={"latent_charges": np.zeros(3)}) + assert out.names() == ("latent_charges",) diff --git a/tests/architecture/observable_coverage.py b/tests/architecture/observable_coverage.py index 84e406ebc..9e9691905 100644 --- a/tests/architecture/observable_coverage.py +++ b/tests/architecture/observable_coverage.py @@ -34,6 +34,8 @@ from pathlib import Path from typing import Union +from mace_core.observables import derivative_name + from tests.golden import harness, surface_scan REPO_ROOT = Path(__file__).resolve().parents[2] @@ -100,12 +102,19 @@ class Spec: set_by: What the model supplies. Required with ``irreps_pattern``, because "it depends on the model" is not an answer until it says on what. + renamed_to: The v1 name, when it differs from the legacy key. Empty + means the key survives as it is. It matters beyond tidiness: a + legacy spelling that almost matches a core field of ``MACEOutput`` + is the shape of a silent dual storage, where the value sits in + ``extras`` under the old name while the field it belongs in stays + ``None``. note: Anything about the row a reader would otherwise have to rederive. """ irreps: str | None = None irreps_pattern: str | None = None set_by: str = "" + renamed_to: str = "" note: str = "" @@ -147,7 +156,20 @@ class Drop: DISPOSITIONS: dict[str, Disposition] = { # --- energies ---------------------------------------------------------- "energy": Spec(irreps="0e"), - "node_energy": Spec(irreps="0e"), + "node_energy": Spec( + irreps="0e", + renamed_to="node_energies", + note=( + "the core field of MACEOutput is plural, so v1 renames the key " + "rather than carrying both spellings. Keeping the singular would " + "leave `extras['node_energy']` able to sit beside a `node_energies` " + "field holding the same quantity, which is exactly the dual " + "storage MACEOutput refuses for the names it does know. Note the " + "model's spelling is the E0-inclusive quantity; the ase calculator " + "uses the same word for the E0-subtracted one, which is why the " + "golden harness aliases this key onto its `energies` channel." + ), + ), "interaction_energy": Spec(irreps="0e"), "les_energy": Spec(irreps="0e"), "electrostatic_energy": Spec(irreps="0e"), @@ -460,6 +482,16 @@ def legacy_eval_keys() -> set[str]: return scan.all_keys +def v1_name(key: str) -> str: + """The name ``key`` carries in v1: renamed, rule-derived, or unchanged.""" + row = DISPOSITIONS[key] + if isinstance(row, Derivative): + return derivative_name(row.of, row.wrt) + if isinstance(row, Spec) and row.renamed_to: + return row.renamed_to + return key + + def channel_of(key: str) -> harness.Channel | None: """The golden harness's declaration for ``key`` on the model surface.""" name = harness.resolve_channel(key, harness.SURFACE_MODEL) diff --git a/tests/architecture/test_observable_completeness.py b/tests/architecture/test_observable_completeness.py index f8705038e..a5d33c23b 100644 --- a/tests/architecture/test_observable_completeness.py +++ b/tests/architecture/test_observable_completeness.py @@ -16,6 +16,8 @@ import pytest from mace_core.observables import ObservableSpec, derivative_name, derivative_sign +from mace_core.outputs import CORE_FIELD_NAMES, FIELD_BY_OBSERVABLE + from tests.architecture.observable_coverage import ( DECLARED_INPUTS, DISPOSITIONS, @@ -28,6 +30,7 @@ legacy_eval_keys, legacy_model_keys, per_atom_of, + v1_name, ) from tests.golden import surface_scan @@ -93,7 +96,7 @@ def test_a_spec_row_builds_a_valid_observable_spec(key): channel = channel_of(key) assert channel is not None spec = ObservableSpec( - name=key, + name=v1_name(key), irreps=row.irreps, per_atom=per_atom, # The legacy unit, as the golden harness records it. This ticket does @@ -263,3 +266,33 @@ def test_the_union_is_sixty_one_names(): union = legacy_model_keys() | legacy_calculator_keys() | legacy_eval_keys() assert len(union) == 61 assert "**61**" in SURFACE_DOC.read_text(encoding="utf-8") + + +def test_every_core_field_of_the_output_is_claimed_by_exactly_one_key(): + """`MACEOutput` has six named fields, and each has to be what some legacy + key becomes. A field nothing claims is a field no model fills. A field + claimed twice, or claimed under a near-miss spelling, is the silent dual + storage the type refuses for the names it knows: the value sits in `extras` + under the old name while the field stays `None`. + + The per-atom energy is why this test exists. The field is `node_energies` + and the legacy key is `node_energy`, so without a recorded rename the row + would have specified an observable whose name misses the field by one + letter. + """ + claims: dict[str, list[str]] = {name: [] for name in CORE_FIELD_NAMES} + for key in DISPOSITIONS: + if isinstance(DISPOSITIONS[key], Drop): + continue + name = v1_name(key) + field = FIELD_BY_OBSERVABLE.get(name, name) + if field in claims: + claims[field].append(key) + unclaimed = sorted(f for f, keys in claims.items() if not keys) + assert not unclaimed, ( + f"{unclaimed} are fields of MACEOutput that no legacy key becomes. " + f"Either a row needs `renamed_to` pointing at the field, or the field " + f"is one nothing fills." + ) + contested = {f: keys for f, keys in claims.items() if len(keys) > 1} + assert not contested, contested From 2c9d536877c62289b6824f933376a49c8392511f Mon Sep 17 00:00:00 2001 From: Alejandro Acosta <127198532+aacostadiaz@users.noreply.github.com> Date: Wed, 16 Sep 2026 11:27:47 +0100 Subject: [PATCH 04/12] Guard the v1 import in the tested tree, not in the pipeline 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. --- .github/gitlab/ci.yml | 14 ---------- .github/workflows/ci-extensions.yaml | 7 ----- .github/workflows/nightly.yaml | 8 ------ .../test_observable_completeness.py | 28 +++++++++++++++++++ 4 files changed, 28 insertions(+), 29 deletions(-) diff --git a/.github/gitlab/ci.yml b/.github/gitlab/ci.yml index 00a236165..caa1493ed 100644 --- a/.github/gitlab/ci.yml +++ b/.github/gitlab/ci.yml @@ -113,24 +113,10 @@ stages: # and the GPU cases are the by-hand reproduction recipe recorded in that # file's module docstring. Removing this clause is not the way to fix that; # a separate, non-gating benchmark job would be. - # `tests/architecture` is ignored, and the reason is an install rather than - # a marker. That suite needs BOTH trees importable: the import contracts - # resolve each root package on the filesystem, and the observable - # completeness test imports `mace_core`. This job installs the legacy - # distribution only, so the directory fails at *collection*, before any - # marker expression gets to deselect it -- a marker on those tests would - # therefore not help. It is not lost coverage: `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 - # it is gpu-marked. - # - # Any later suite that needs the v1 stack belongs on this line too. - # `tests/parity` is the next one. - > MACE_REQUIRE_CAPS="$REQUIRE_CAPS" MACE_CI_ALLOW_NETWORK=1 python -m pytest tests - --ignore=tests/architecture -m "($MARKER_EXPR) and not benchmark" -n 2 --timeout 1200 -ra diff --git a/.github/workflows/ci-extensions.yaml b/.github/workflows/ci-extensions.yaml index e690e5832..0f265b201 100644 --- a/.github/workflows/ci-extensions.yaml +++ b/.github/workflows/ci-extensions.yaml @@ -263,12 +263,5 @@ jobs: - uses: ./.github/actions/run-tests with: tests: tests - # See the same flag in .github/gitlab/ci.yml: `tests/architecture` - # needs both trees importable and this job installs the legacy one, - # so the directory fails at collection rather than being deselected - # by the marker expression. Gated per pull request by ci-core's - # `architecture` job. This job is paths-filtered, so the breakage - # would have waited for the first pull request touching backends. - extra-args: --ignore=tests/architecture markers: cueq and not gpu and not polar require-caps: cueq diff --git a/.github/workflows/nightly.yaml b/.github/workflows/nightly.yaml index fe4d29ca5..62d03c580 100644 --- a/.github/workflows/nightly.yaml +++ b/.github/workflows/nightly.yaml @@ -316,11 +316,6 @@ jobs: - uses: ./.github/actions/run-tests with: tests: tests - # See the same flag in .github/gitlab/ci.yml: `tests/architecture` - # needs both trees importable and this job installs the legacy one, - # so the directory fails at collection rather than being deselected. - # It is gated per pull request by ci-core's `architecture` job. - extra-args: --ignore=tests/architecture markers: not gpu and not benchmark require-caps: network,cueq,polar,les,torchsim,schedulefree,magnetic,wandb allow-network: "true" @@ -509,9 +504,6 @@ jobs: - uses: ./.github/actions/run-tests with: tests: tests - # Same reason as coverage-full above: this job installs `dev` alone, - # so `tests/architecture` cannot be collected here. - extra-args: --ignore=tests/architecture allow-network: "true" store-durations: "true" timeout: "1200" diff --git a/tests/architecture/test_observable_completeness.py b/tests/architecture/test_observable_completeness.py index a5d33c23b..a359b05ba 100644 --- a/tests/architecture/test_observable_completeness.py +++ b/tests/architecture/test_observable_completeness.py @@ -10,10 +10,38 @@ from __future__ import annotations +import importlib.util import re from pathlib import Path import pytest + +# This module is the only thing under `tests/architecture` that imports the v1 +# stack, and four jobs run `pytest tests` over the whole tree with the legacy +# distribution alone: the two GPU jobs of the MPCDF pipeline, and nightly's +# coverage-full and durations-refresh. A module-level import of `mace_core` +# fails there at *collection*, before any marker expression can deselect it, so +# a capability marker on the tests below would not help. +# +# The guard has to live here rather than in the jobs, and the GPU pipeline is +# why. 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, deliberately, so that a fork cannot choose what runs on MPCDF hardware. +# An `--ignore` added to `.github/gitlab/ci.yml` is therefore invisible to the +# pull request that adds it. The tested tree is the only lever a fork has. +# +# `find_spec` and not `pytest.importorskip`: it resolves the module without +# executing it, so "mace_core is not installed" skips while "mace_core is +# installed and broken" still raises at the real import below. And the skip +# cannot quietly hide the tests from the job that owes them, because the +# `architecture` job runs `lint-imports` first and that step fails outright +# when a root package is missing from the filesystem. +if importlib.util.find_spec("mace_core") is None: # pragma: no cover + pytest.skip( + "needs the v1 packages installed; run this suite from the architecture " + "job, or `pip install -e packages/mace-core`", + allow_module_level=True, + ) from mace_core.observables import ObservableSpec, derivative_name, derivative_sign from mace_core.outputs import CORE_FIELD_NAMES, FIELD_BY_OBSERVABLE From 00ae645f8f051908385966d46392d9a0a465c0b3 Mon Sep 17 00:00:00 2001 From: Alejandro Acosta <127198532+aacostadiaz@users.noreply.github.com> Date: Fri, 18 Sep 2026 11:16:40 +0100 Subject: [PATCH 05/12] Correct the sign the virial field is documented with 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. --- packages/mace-core/src/mace_core/outputs.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/packages/mace-core/src/mace_core/outputs.py b/packages/mace-core/src/mace_core/outputs.py index 18a340e60..9efe392d7 100644 --- a/packages/mace-core/src/mace_core/outputs.py +++ b/packages/mace-core/src/mace_core/outputs.py @@ -60,8 +60,12 @@ class MACEOutput(Generic[TensorT]): forces: ``-d(energy)/d(positions)``, in eV/Å, shape ``(n_atoms, 3)``. stress: ``+d(energy)/d(strain) / volume``, in eV/ų, shape ``(n_graphs, 3, 3)``. - virials: The same derivative before the volume division, in eV, shape - ``(n_graphs, 3, 3)``. + virials: ``-d(energy)/d(strain)``, in eV, shape ``(n_graphs, 3, 3)``. + Note the sign: this is the **negative** of the derivative the + stress is built from, so ``stress * volume == -virials`` rather + than ``+virials``. The two are easy to state as one quantity up to + a volume, and they are not; the `virials` row of the observable + inventory carries the measurement. dipole: Total dipole per graph, shape ``(n_graphs, 3)``. extras: Every other declared observable, keyed by its :class:`~mace_core.observables.ObservableSpec` name. From dee462c17818db0e9986bad2fbe00ede3478f453 Mon Sep 17 00:00:00 2001 From: Alejandro Acosta <127198532+aacostadiaz@users.noreply.github.com> Date: Fri, 18 Sep 2026 12:35:56 +0100 Subject: [PATCH 06/12] Name the input the stress is differentiated against strain 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__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. --- docs/reforge/target_layout.md | 2 +- .../src/mace_core/defaults/observables.yaml | 13 ++++++---- .../src/mace_core/observables/defaults.py | 2 +- .../src/mace_core/observables/derivatives.py | 6 ++--- .../src/mace_core/observables/spec.py | 24 ++++++++++++------- packages/mace-core/tests/test_observables.py | 20 ++++++++-------- tests/architecture/observable_coverage.py | 19 ++++++++------- .../test_observable_completeness.py | 4 ++-- 8 files changed, 50 insertions(+), 40 deletions(-) diff --git a/docs/reforge/target_layout.md b/docs/reforge/target_layout.md index f0952986a..d7df0d741 100644 --- a/docs/reforge/target_layout.md +++ b/docs/reforge/target_layout.md @@ -54,7 +54,7 @@ packages/mace-core/ │ │ ├── derivatives.py # d__d_ naming, and the three special cases with their signs │ │ └── defaults.py # loader for a declarations file │ ├── defaults/ -│ │ └── observables.yaml # energy + its position and cell derivatives: the row every observable copies +│ │ └── observables.yaml # energy + its position and strain derivatives: the row every observable copies │ │ │ ├── kernels/ │ │ ├── protocol.py # KernelBackend Protocol, generic over TensorT: make_* factories + capabilities (§3.1) diff --git a/packages/mace-core/src/mace_core/defaults/observables.yaml b/packages/mace-core/src/mace_core/defaults/observables.yaml index 9cbc7a369..005ebf20c 100644 --- a/packages/mace-core/src/mace_core/defaults/observables.yaml +++ b/packages/mace-core/src/mace_core/defaults/observables.yaml @@ -6,7 +6,7 @@ # head, the loss term, the derivative names and the per-atom/per-graph padding. # # Units follow the project convention: eV and Å. The strain is dimensionless, -# so the cell input carries "1". +# so its input row carries "1". inputs: # Atomic positions. A polar vector: it changes sign under inversion, which is @@ -16,9 +16,12 @@ inputs: per_atom: true units: "Å" - # The symmetric strain the cell derivative is actually taken against, not the - # nine cell entries. A symmetric rank-2 tensor is a scalar plus an l=2 part. - - name: cell + # The symmetric strain the stress is the derivative against, not the nine + # cell entries: a symmetric rank-2 tensor is a scalar plus an l=2 part, six + # components. Unlike `pos` it is not read from the data. The derivative + # engine materialises it as zeros around the model call and applies it to + # the positions and the cell, which is how the frozen tree does it too. + - name: strain irreps: "0e+2e" per_atom: false units: "1" @@ -43,7 +46,7 @@ observables: default_loss_weight: 100.0 # Named `stress`, reported as +dE/dstrain, divided by the cell volume by # whatever computes it. The division is not part of the sign. - - wrt: cell + - wrt: strain units: "eV/Å^3" normalization: "none" default_loss_weight: 1.0 diff --git a/packages/mace-core/src/mace_core/observables/defaults.py b/packages/mace-core/src/mace_core/observables/defaults.py index 37eb96293..f3ced638b 100644 --- a/packages/mace-core/src/mace_core/observables/defaults.py +++ b/packages/mace-core/src/mace_core/observables/defaults.py @@ -45,6 +45,6 @@ def load_catalogue(path: str | Path) -> ObservableCatalogue: def load_default_catalogue() -> ObservableCatalogue: - """The shipped declarations: energy plus its position and cell derivatives.""" + """The shipped declarations: energy plus its position and strain derivatives.""" resource = files("mace_core").joinpath(DEFAULTS_RESOURCE) return _catalogue_from_text(resource.read_text(encoding="utf-8"), DEFAULTS_RESOURCE) diff --git a/packages/mace-core/src/mace_core/observables/derivatives.py b/packages/mace-core/src/mace_core/observables/derivatives.py index 50633d5b6..f0817b916 100644 --- a/packages/mace-core/src/mace_core/observables/derivatives.py +++ b/packages/mace-core/src/mace_core/observables/derivatives.py @@ -6,10 +6,10 @@ through the consumers that need them. The third special case is the reason this grammar is written over declared -inputs rather than over positions and the cell. ``magforces`` is +inputs rather than over positions and the strain. ``magforces`` is ``-dE/d(magmom)``, computed in the same autograd call as the forces, trained with its own loss term, and used by the magnetic self-consistent model to drive -its fixed point. A grammar that knew only ``d__d_pos`` and ``d__d_cell`` +its fixed point. A grammar that knew only ``d__d_pos`` and ``d__d_strain`` could not express it, and the magnetic work would have had to go around the abstraction that exists to prevent exactly that. @@ -31,7 +31,7 @@ #: the ``d__d_`` default. Everything else follows the rule. SPECIAL_CASES: dict[tuple[str, str], tuple[str, int]] = { ("energy", "pos"): ("forces", -1), - ("energy", "cell"): ("stress", +1), + ("energy", "strain"): ("stress", +1), ("energy", "magmom"): ("magforces", -1), } diff --git a/packages/mace-core/src/mace_core/observables/spec.py b/packages/mace-core/src/mace_core/observables/spec.py index f6458a0ba..7d7460a93 100644 --- a/packages/mace-core/src/mace_core/observables/spec.py +++ b/packages/mace-core/src/mace_core/observables/spec.py @@ -9,9 +9,9 @@ Three objects, and they do different jobs: ``InputSpec`` - something the model is given -- positions, the cell, a magnetic moment, an - electronic temperature. Declared so that a derivative can be taken with - respect to it without new code. + a leaf a derivative can be taken against: positions, the strain, a magnetic + moment, an electronic temperature. Declaring one is what makes its + derivative reachable without new code. ``ObservableSpec`` something the model produces and a loss can be written against. @@ -75,11 +75,17 @@ def _check_name(value: str, kind: str) -> str: class InputSpec(BaseModel): - """Something the model is given, and can be differentiated against. + """Something a derivative can be taken against. - ``pos`` and ``cell`` are the two every model has. Anything else is declared - the same way, which is what makes ``d_energy_d_`` reachable for a - new feature without touching code. + ``pos`` and ``strain`` are the two every model has. Anything else is + declared the same way, which is what makes ``d_energy_d_`` + reachable for a new feature without touching code. + + An input is a leaf of the derivative graph, not necessarily a field read + from the data. ``pos`` is both. ``strain`` is only the first: the + derivative engine materialises it as zeros around the model call and + applies it to the positions and the cell, so nothing reads a strain from a + dataset and none is stored. """ model_config = ConfigDict(extra="forbid", frozen=True) @@ -87,7 +93,7 @@ class InputSpec(BaseModel): name: str irreps: str #: ``True`` for one value per atom (positions, magnetic moments), ``False`` - #: for one per structure (the cell, a total charge). This is what decides + #: for one per structure (the strain, a total charge). This is what decides #: whether a derivative taken against the input is padded per node or per #: graph. per_atom: bool @@ -122,7 +128,7 @@ class DerivativeRequest(BaseModel): @model_validator(mode="before") @classmethod def _accept_bare_name(cls, value: object) -> object: - """``derivatives: [pos, cell]`` is the same as spelling out ``wrt``.""" + """``derivatives: [pos, strain]`` is the same as spelling out ``wrt``.""" if isinstance(value, str): return {"wrt": value} return value diff --git a/packages/mace-core/tests/test_observables.py b/packages/mace-core/tests/test_observables.py index 31491360a..df19c881e 100644 --- a/packages/mace-core/tests/test_observables.py +++ b/packages/mace-core/tests/test_observables.py @@ -29,7 +29,7 @@ irreps: "1o" per_atom: true units: "Å" - - name: cell + - name: strain irreps: "0e+2e" per_atom: false units: "1" @@ -162,7 +162,7 @@ def test_units_may_not_be_empty(): ("quantity", "wrt", "name", "sign"), [ ("energy", "pos", "forces", -1), - ("energy", "cell", "stress", +1), + ("energy", "strain", "stress", +1), ("energy", "magmom", "magforces", -1), ], ) @@ -175,7 +175,7 @@ def test_the_three_special_cases_keep_their_names_and_signs(quantity, wrt, name, ("quantity", "wrt", "name"), [ ("dipole", "pos", "d_dipole_d_pos"), - ("dipole", "cell", "d_dipole_d_cell"), + ("dipole", "strain", "d_dipole_d_strain"), ("polarizability", "pos", "d_polarizability_d_pos"), ("energy", "elec_temp", "d_energy_d_elec_temp"), ("quadrupole", "magmom", "d_quadrupole_d_magmom"), @@ -194,7 +194,7 @@ def test_everything_else_follows_the_rule(quantity, wrt, name): def test_the_defaults_declare_energy_and_its_two_derivatives(): catalogue = load_default_catalogue() assert catalogue.names() == ("energy", "forces", "stress") - assert [spec.name for spec in catalogue.inputs] == ["pos", "cell"] + assert [spec.name for spec in catalogue.inputs] == ["pos", "strain"] def test_the_default_forces_row_is_the_negative_position_gradient(): @@ -208,7 +208,7 @@ def test_the_default_forces_row_is_the_negative_position_gradient(): def test_the_default_stress_row_is_the_positive_strain_gradient(): - stress = load_default_catalogue().derivative("energy", "cell") + stress = load_default_catalogue().derivative("energy", "strain") assert stress.name == "stress" assert stress.sign == +1 assert stress.per_atom is False @@ -230,7 +230,7 @@ def test_a_new_rank_two_per_atom_observable_is_a_row_in_yaml(tmp_path): units: "e*Å^2" normalization: "rms" default_loss_weight: 2.5 - derivatives: [pos, cell] + derivatives: [pos, strain] """, tmp_path, ) @@ -241,13 +241,13 @@ def test_a_new_rank_two_per_atom_observable_is_a_row_in_yaml(tmp_path): assert catalogue.names() == ( "quadrupole", "d_quadrupole_d_pos", - "d_quadrupole_d_cell", + "d_quadrupole_d_strain", ) def test_a_new_input_feature_makes_its_derivative_declarable(tmp_path): """`magmom` is the case that pays for the grammar being written over - declared inputs rather than over positions and the cell.""" + declared inputs rather than over positions and the strain.""" catalogue = catalogue_from( """ - name: magmom @@ -316,7 +316,7 @@ def test_a_derivative_against_an_undeclared_input_is_an_error(tmp_path): message = str(caught.value) assert "energy" in message assert "elec_temp" in message - assert "['cell', 'pos']" in message + assert "['pos', 'strain']" in message def test_a_derived_name_may_not_collide_with_a_declared_observable(tmp_path): @@ -385,7 +385,7 @@ def test_an_unknown_observable_or_input_says_what_is_declared(): assert "['energy']" in str(caught.value) with pytest.raises(KeyError) as caught: catalogue.input("magmom") - assert "['cell', 'pos']" in str(caught.value) + assert "['pos', 'strain']" in str(caught.value) def test_a_derivative_can_be_named_without_having_been_requested(): diff --git a/tests/architecture/observable_coverage.py b/tests/architecture/observable_coverage.py index 9e9691905..bb918ad11 100644 --- a/tests/architecture/observable_coverage.py +++ b/tests/architecture/observable_coverage.py @@ -69,11 +69,11 @@ } ) -#: The inputs a derivative row may be taken against. ``pos`` and ``cell`` are -#: the two every model has and are declared in the shipped defaults; ``magmom`` -#: is the third the frozen tree actually differentiates against, and is -#: declared by whichever configuration turns the magnetic model on. -DECLARED_INPUTS = frozenset({"pos", "cell", "magmom"}) +#: The inputs a derivative row may be taken against. ``pos`` and ``strain`` +#: are the two every model has and are declared in the shipped defaults; +#: ``magmom`` is the third the frozen tree actually differentiates against, and +#: is declared by whichever configuration turns the magnetic model on. +DECLARED_INPUTS = frozenset({"pos", "strain", "magmom"}) @dataclass(frozen=True) @@ -215,7 +215,7 @@ class Drop: "virials": Spec( irreps="0e+2e", note=( - "the NEGATED cell derivative, and the sign is not a detail: " + "the NEGATED strain derivative, and the sign is not a detail: " "`compute_forces_virials` computes the stress from the raw " "gradient and negates the virial only in its return statement " "(mace/modules/utils.py:107-115), so the frozen tree reports " @@ -346,7 +346,7 @@ class Drop: ), "stress": Derivative( of="energy", - wrt="cell", + wrt="strain", sign=+1, note=( "positive, and only because the stress is built from the raw " @@ -404,8 +404,9 @@ class Drop: ), "displacement": Drop( reason=( - "the symmetric strain handle the cell derivative is taken against. " - "It is created as zeros to attach the cell to the autograd graph " + "the symmetric strain the stress is the derivative against, the " + "`strain` input of the shipped declarations. It is created as " + "zeros to attach the positions and the cell to the autograd graph " "and nothing ever writes to it, so its value is identically zero " "on every structure. It belongs to the derivative engine, not to " "the output surface." diff --git a/tests/architecture/test_observable_completeness.py b/tests/architecture/test_observable_completeness.py index a359b05ba..946eca75a 100644 --- a/tests/architecture/test_observable_completeness.py +++ b/tests/architecture/test_observable_completeness.py @@ -184,8 +184,8 @@ def test_the_three_named_derivatives_keep_their_legacy_names(): and the two negated ones are the two the frozen tree negates.""" assert derivative_name("energy", "pos") == "forces" assert derivative_sign("energy", "pos") == -1 - assert derivative_name("energy", "cell") == "stress" - assert derivative_sign("energy", "cell") == +1 + assert derivative_name("energy", "strain") == "stress" + assert derivative_sign("energy", "strain") == +1 assert derivative_name("energy", "magmom") == "magforces" assert derivative_sign("energy", "magmom") == -1 for name in ("forces", "stress", "magforces"): From 7b59ca89760b115f6b0d87f0b7297097473bf63d Mon Sep 17 00:00:00 2001 From: Alejandro Acosta <127198532+aacostadiaz@users.noreply.github.com> Date: Mon, 21 Sep 2026 09:21:26 +0100 Subject: [PATCH 07/12] Move scaling and loss weights out of the observable spec 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. --- docs/reforge/extending_mace.md | 13 +++---- docs/reforge/target_layout.md | 6 +-- .../src/mace_core/defaults/observables.yaml | 11 ------ .../src/mace_core/observables/__init__.py | 4 -- .../src/mace_core/observables/spec.py | 31 ++------------- packages/mace-core/tests/test_observables.py | 38 +------------------ .../test_observable_completeness.py | 3 -- 7 files changed, 13 insertions(+), 93 deletions(-) diff --git a/docs/reforge/extending_mace.md b/docs/reforge/extending_mace.md index 55871b51c..79fe9189f 100644 --- a/docs/reforge/extending_mace.md +++ b/docs/reforge/extending_mace.md @@ -92,7 +92,6 @@ magforces: derivation: "autograd(energy, wrt=magmom)" # -dE/dmagmom per_atom: true irreps: "1o" # a 3-vector, same convention as magmom - default_loss_weight: 1.0 ``` ```toml @@ -118,8 +117,6 @@ magnetic_moment: derivation: readout # a learned equivariant readout over node features per_atom: true irreps: "1o" # same convention as the magmom input - normalization: "component" # scale-only; a 1o vector can be scaled but not shifted - default_loss_weight: 1.0 ``` ```toml @@ -129,9 +126,9 @@ observables = ["energy", "forces", "stress", "magforces", "magnetic_moment"] ``` `MACEOutputs` would build the equivariant `1o` readout head automatically; the result appears as -`output.extras["magnetic_moment"]`. **Zero code** — the head, its typed output, its `normalization` -and its loss term are all derived from this one row. (`normalization` is a user knob: a non-scalar like -`1o` can be scaled but not shifted; only scalars such as energy take the classic **scale-shift**.) That +`output.extras["magnetic_moment"]`. **Zero code**: the head, its typed output and its loss term are +all derived from this one row. The head's scaling is set in the model config, where a non-scalar like +`1o` can be scaled but not shifted; only scalars such as energy take the classic **scale-shift**. That is the payoff of the declarative table: a genuinely new *predicted* property is a row, not a model change. @@ -164,8 +161,8 @@ transforms = ["rotate_magmom"] # legacy --data_aug_magmom ## 4. Loss — a term for `magforces` (config only) -Because `magforces` is a declared observable, its loss term is **generated automatically** with its -`default_loss_weight`; you only override the weights in config. Tuning the loss is never new code: +Because `magforces` is a declared observable, its loss term is **generated automatically**, and its +weight is a field in `LossConfig`. Tuning the loss is never new code: ```toml # config.toml diff --git a/docs/reforge/target_layout.md b/docs/reforge/target_layout.md index d7df0d741..b21d3e75f 100644 --- a/docs/reforge/target_layout.md +++ b/docs/reforge/target_layout.md @@ -400,7 +400,7 @@ that come up a lot: code at all — its loss term appears automatically with a default weight you can override. - **A new property is a table row, not an add-on.** The framework is property-agnostic by construction: the observable table maps a property name to its mathematical structure (irreps, - per-atom vs total, units, normalization), and everything downstream (head, loss, derivatives) is + per-atom vs total, units), and everything downstream (head, loss, derivatives) is derived from that row. The sections below give the worked examples for each row of the ladder; for a single **end-to-end @@ -493,14 +493,14 @@ silently wrong forces); it stays usable for inference (`supports_double_backward ### 3.2 A new observable (config only) -- **Extender touches:** a declarations file: one `ObservableSpec` row giving `name`, `irreps`, `per_atom`, `units`, `normalization`, `default_loss_weight`, and the declared inputs to differentiate against. No module, no decorator. A derivative is named by the rule `d__d_`, with `forces`, `stress` and `magforces` as the three special cases, so asking for a derivative against a newly declared input needs no code either. +- **Extender touches:** a declarations file: one `ObservableSpec` row giving `name`, `irreps`, `per_atom`, `units`, and the declared inputs to differentiate against. The scaling of the head that produces it is set in the model config and its loss weight in `LossConfig`, both keyed by this name. No module, no decorator. A derivative is named by the rule `d__d_`, with `forces`, `stress` and `magforces` as the three special cases, so asking for a derivative against a newly declared input needs no code either. - **Core touched:** zero files. The model exposes the row automatically because `BaseMACE` iterates over the declared observables; `MACEOutput` carries the six core fields and everything else by name in `extras`. - **Enabling it:** list it in the model config's observables, or point the config at a declarations file that extends `defaults/observables.yaml`. - **Test:** `packages/mace-core/tests/test_observables.py` validates the grammar and the derivative naming (pure); if it is autograd-derived, `tests/parity` verifies finite-diff. ### 3.3 A new loss / transform (plugin registry) -- **Tuning an existing loss is config, not a new loss.** `LossConfig` carries the per-observable **weights** *and* the loss's own **parameters** (e.g. `params={"huber_delta": 0.02}`) — this preserves the legacy `--energy_weight`/`--forces_weight`/`--huber_delta` knobs as config fields. Changing a coefficient never needs a `@register_loss`. And a well-defined spherical-tensor observable needs **no** loss code at all — its term is generated from the observable table with `default_loss_weight`. +- **Tuning an existing loss is config, not a new loss.** `LossConfig` carries the per-observable **weights** *and* the loss's own **parameters** (e.g. `params={"huber_delta": 0.02}`) — this preserves the legacy `--energy_weight`/`--forces_weight`/`--huber_delta` knobs as config fields. Changing a coefficient never needs a `@register_loss`. And a well-defined spherical-tensor observable needs **no** loss code at all — its term is generated from the observable table, with its weight read from `LossConfig`. - **A genuinely new loss:** `mace_torch/train/loss.py` (or an external package) with `@register_loss("myloss")` on a `torch.nn.Module`; select via `LossConfig(name="myloss", weights=..., params=...)`. - **Data transform:** `@register_transform("mytransform")` in `mace_torch/data/`; chained via `DataConfig(transforms=[...])`. - **Core touched:** the registries (`LOSS_REGISTRY`, `TRANSFORM_REGISTRY`) live in `mace_core.registries` as specs; **adding one does not edit the registry**, only the decorator populates it at import time. Zero core edits. diff --git a/packages/mace-core/src/mace_core/defaults/observables.yaml b/packages/mace-core/src/mace_core/defaults/observables.yaml index 005ebf20c..64263a56d 100644 --- a/packages/mace-core/src/mace_core/defaults/observables.yaml +++ b/packages/mace-core/src/mace_core/defaults/observables.yaml @@ -31,22 +31,11 @@ observables: irreps: "0e" per_atom: false units: "eV" - # The std of the target, the successor of the legacy `std_scaling` entry. - # The legacy default was `rms_forces_scaling`, which scales the energy - # readout by the RMS of the *force* targets; that couples two observables - # through one number and cannot be written as a per-observable field. Which - # of the two v1 defaults to is a decision for the head, not for this file. - normalization: "std" - default_loss_weight: 1.0 derivatives: # Named `forces`, reported as -dE/dpos. - wrt: pos units: "eV/Å" - normalization: "rms" - default_loss_weight: 100.0 # Named `stress`, reported as +dE/dstrain, divided by the cell volume by # whatever computes it. The division is not part of the sign. - wrt: strain units: "eV/Å^3" - normalization: "none" - default_loss_weight: 1.0 diff --git a/packages/mace-core/src/mace_core/observables/__init__.py b/packages/mace-core/src/mace_core/observables/__init__.py index d5d3f4503..83067c767 100644 --- a/packages/mace-core/src/mace_core/observables/__init__.py +++ b/packages/mace-core/src/mace_core/observables/__init__.py @@ -18,11 +18,9 @@ parse_irreps, ) from mace_core.observables.spec import ( - NORMALIZATIONS, DerivativeRequest, DerivativeSpec, InputSpec, - Normalization, ObservableCatalogue, ObservableSpec, ) @@ -30,14 +28,12 @@ __all__ = [ "DEFAULTS_RESOURCE", "IRREPS_GRAMMAR", - "NORMALIZATIONS", "SPECIAL_CASES", "DerivativeRequest", "DerivativeSpec", "InputSpec", "IrrepTerm", "IrrepsGrammarError", - "Normalization", "ObservableCatalogue", "ObservableSpec", "derivative_name", diff --git a/packages/mace-core/src/mace_core/observables/spec.py b/packages/mace-core/src/mace_core/observables/spec.py index 7d7460a93..e268f72dd 100644 --- a/packages/mace-core/src/mace_core/observables/spec.py +++ b/packages/mace-core/src/mace_core/observables/spec.py @@ -29,8 +29,6 @@ from __future__ import annotations -from typing import Literal - from pydantic import BaseModel, ConfigDict, Field, model_validator from mace_core.observables.derivatives import derivative_name, derivative_sign @@ -41,26 +39,13 @@ ) __all__ = [ - "NORMALIZATIONS", "DerivativeRequest", "DerivativeSpec", "InputSpec", - "Normalization", "ObservableCatalogue", "ObservableSpec", ] -#: How a target is scaled before it reaches a head and a loss term. The -#: per-observable successor of the legacy scaling registry, whose three entries -#: were a global choice for the whole model. This package stores and validates -#: the value; the head applies it in the output layer and the loss applies the -#: matching term, both reading this one field, so there is never a second -#: scaling mechanism to keep in step with it. -Normalization = Literal["none", "std", "rms"] - -#: The accepted values, for error messages and for callers that enumerate them. -NORMALIZATIONS: tuple[str, ...] = ("none", "std", "rms") - _SCALAR = (IrrepTerm(multiplicity=1, degree=0, parity="e"),) @@ -107,20 +92,18 @@ def _validate(self) -> InputSpec: class DerivativeRequest(BaseModel): - """A derivative an observable asks for, and the loss settings it carries. + """A derivative an observable asks for. The name and the sign are not here: they are derived, and letting a declaration override them would reintroduce the per-consumer naming this - abstraction removes. What a declaration does own is what a loss needs -- - its weight, its normalization, and the unit string to report. + abstraction removes. What a declaration owns is which input to + differentiate against, and the unit string to report. """ model_config = ConfigDict(extra="forbid", frozen=True) #: The name of the declared input to differentiate against. wrt: str - default_loss_weight: float = Field(default=1.0, ge=0.0) - normalization: Normalization = "none" #: Left to the declaration. Deriving it would mean unit algebra over the #: quantity and the input, which this ticket does not own. units: str | None = None @@ -154,12 +137,10 @@ class DerivativeSpec(BaseModel): #: general case is a tensor product and the algebra is not this module's. irreps: str | None units: str | None - normalization: Normalization - default_loss_weight: float class ObservableSpec(BaseModel): - """One declared property: what it is, what shape it has, how it is scaled. + """One declared property: what it is and what shape it has. Any atomic or total spherical-tensor property declared here becomes trainable with no new code: the spec drives the head, the loss term and the @@ -178,8 +159,6 @@ class ObservableSpec(BaseModel): per_atom: bool #: Project convention: eV, Å. units: str = Field(min_length=1) - normalization: Normalization - default_loss_weight: float = Field(default=1.0, ge=0.0) #: The derivatives this observable asks for. Naming works for any declared #: input whether or not it is listed here; listing it is what says the #: model should compute it. @@ -307,8 +286,6 @@ def derivative(self, observable: str, wrt: str) -> DerivativeSpec: per_atom=input_spec.per_atom, irreps=input_spec.irreps if spec.is_scalar else None, units=request.units, - normalization=request.normalization, - default_loss_weight=request.default_loss_weight, ) def requested_derivatives(self) -> tuple[DerivativeSpec, ...]: diff --git a/packages/mace-core/tests/test_observables.py b/packages/mace-core/tests/test_observables.py index df19c881e..5935662cc 100644 --- a/packages/mace-core/tests/test_observables.py +++ b/packages/mace-core/tests/test_observables.py @@ -93,7 +93,6 @@ def test_a_malformed_spec_names_the_observable_and_the_grammar(): irreps="rank2", per_atom=True, units="e*Å^2", - normalization="none", ) message = str(caught.value) assert "quadrupole" in message @@ -106,22 +105,6 @@ def test_a_malformed_spec_names_the_observable_and_the_grammar(): # --------------------------------------------------------------------------- -def test_normalization_is_required_and_closed(): - # Both calls are rejected by the type checker as well, which is the point: - # the field is a closed Literal, so a wrong value is caught statically and - # at runtime. The ignores are what let the runtime half be tested. - with pytest.raises(ValidationError): - ObservableSpec(name="q", irreps="0e", per_atom=False, units="eV") # ty: ignore[missing-argument] - with pytest.raises(ValidationError): - ObservableSpec( - name="q", - irreps="0e", - per_atom=False, - units="eV", - normalization="minmax", # ty: ignore[invalid-argument-type] - ) - - def test_an_unknown_field_is_an_error_rather_than_ignored(): with pytest.raises(ValidationError): ObservableSpec( @@ -129,7 +112,6 @@ def test_an_unknown_field_is_an_error_rather_than_ignored(): irreps="0e", per_atom=False, units="eV", - normalization="none", weight=3.0, # ty: ignore[unknown-argument] ) @@ -141,16 +123,13 @@ def test_a_name_that_is_not_an_identifier_is_rejected(): irreps="0e", per_atom=True, units="e", - normalization="none", ) assert "identifier" in str(caught.value) def test_units_may_not_be_empty(): with pytest.raises(ValidationError): - ObservableSpec( - name="q", irreps="0e", per_atom=False, units="", normalization="none" - ) + ObservableSpec(name="q", irreps="0e", per_atom=False, units="") # --------------------------------------------------------------------------- @@ -204,7 +183,6 @@ def test_the_default_forces_row_is_the_negative_position_gradient(): assert forces.per_atom is True assert forces.irreps == "1o" assert forces.units == "eV/Å" - assert forces.default_loss_weight == 100.0 def test_the_default_stress_row_is_the_positive_strain_gradient(): @@ -228,8 +206,6 @@ def test_a_new_rank_two_per_atom_observable_is_a_row_in_yaml(tmp_path): irreps: "0e+2e" per_atom: true units: "e*Å^2" - normalization: "rms" - default_loss_weight: 2.5 derivatives: [pos, strain] """, tmp_path, @@ -237,7 +213,6 @@ def test_a_new_rank_two_per_atom_observable_is_a_row_in_yaml(tmp_path): quadrupole = catalogue.observable("quadrupole") assert quadrupole.per_atom is True assert quadrupole.dimension == 6 - assert quadrupole.default_loss_weight == 2.5 assert catalogue.names() == ( "quadrupole", "d_quadrupole_d_pos", @@ -260,11 +235,9 @@ def test_a_new_input_feature_makes_its_derivative_declarable(tmp_path): irreps: "0e" per_atom: false units: "eV" - normalization: "std" derivatives: - wrt: magmom units: "eV/muB" - default_loss_weight: 10.0 """, tmp_path, ) @@ -274,7 +247,6 @@ def test_a_new_input_feature_makes_its_derivative_declarable(tmp_path): assert magforces.per_atom is True # A magnetic moment is an axial vector, so its conjugate force is too. assert magforces.irreps == "1e" - assert magforces.default_loss_weight == 10.0 assert catalogue.names() == ("energy", "magforces") @@ -286,7 +258,6 @@ def test_the_bare_string_form_and_the_mapping_form_agree(tmp_path): irreps: "0e" per_atom: false units: "eV" - normalization: "std" derivatives: [pos] """, tmp_path, @@ -308,7 +279,6 @@ def test_a_derivative_against_an_undeclared_input_is_an_error(tmp_path): irreps: "0e" per_atom: false units: "eV" - normalization: "std" derivatives: [elec_temp] """, tmp_path, @@ -328,12 +298,10 @@ def test_a_derived_name_may_not_collide_with_a_declared_observable(tmp_path): irreps: "1o" per_atom: true units: "eV/Å" - normalization: "rms" - name: energy irreps: "0e" per_atom: false units: "eV" - normalization: "std" derivatives: [pos] """, tmp_path, @@ -350,12 +318,10 @@ def test_a_name_declared_twice_is_an_error(tmp_path): irreps: "0e" per_atom: false units: "eV" - normalization: "std" - name: energy irreps: "0e" per_atom: false units: "eV" - normalization: "none" """, tmp_path, ) @@ -371,7 +337,6 @@ def test_asking_for_the_same_derivative_twice_is_an_error(tmp_path): irreps: "0e" per_atom: false units: "eV" - normalization: "std" derivatives: [pos, pos] """, tmp_path, @@ -398,7 +363,6 @@ def test_a_derivative_can_be_named_without_having_been_requested(): irreps="1o", per_atom=False, units="Debye", - normalization="rms", ) ], ) diff --git a/tests/architecture/test_observable_completeness.py b/tests/architecture/test_observable_completeness.py index 946eca75a..9cc2bcb9b 100644 --- a/tests/architecture/test_observable_completeness.py +++ b/tests/architecture/test_observable_completeness.py @@ -131,9 +131,6 @@ def test_a_spec_row_builds_a_valid_observable_spec(key): # not canonicalise unit strings, so `Ang` is left as it is written # there rather than rewritten to `Å`. units=channel.unit, - # Neutral: which normalization each family wants is decided by the - # ticket that builds its head, not by this coverage table. - normalization="none", ) assert spec.per_atom == per_atom assert spec.dimension >= 1 From e2f728a723b726badd72a9d9578d93b1870b92df Mon Sep 17 00:00:00 2001 From: Alejandro Acosta <127198532+aacostadiaz@users.noreply.github.com> Date: Mon, 21 Sep 2026 21:28:00 +0100 Subject: [PATCH 08/12] Declare a derivative's name and sign in the file rather than in a table 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__d_` 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__d_`. 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. --- .../src/mace_core/defaults/observables.yaml | 16 +- .../src/mace_core/observables/__init__.py | 12 +- .../src/mace_core/observables/derivatives.py | 76 +++++----- .../src/mace_core/observables/spec.py | 75 +++++++++- packages/mace-core/tests/test_observables.py | 138 +++++++++++++++--- tests/architecture/observable_coverage.py | 13 +- .../test_observable_completeness.py | 73 ++++++--- 7 files changed, 303 insertions(+), 100 deletions(-) diff --git a/packages/mace-core/src/mace_core/defaults/observables.yaml b/packages/mace-core/src/mace_core/defaults/observables.yaml index 64263a56d..3da29c2ea 100644 --- a/packages/mace-core/src/mace_core/defaults/observables.yaml +++ b/packages/mace-core/src/mace_core/defaults/observables.yaml @@ -5,6 +5,12 @@ # row here (or in a project's own file) and needs no code: the spec drives the # head, the loss term, the derivative names and the per-atom/per-graph padding. # +# A derivative with a name of its own declares it here, together with the sign +# it is reported with, rather than being special-cased in code. Leave both out +# and it is called `d__d_` and carries the gradient's own +# sign; give a name and the sign becomes required, because a renamed quantity +# that silently inherited +1 is a model trained on inverted forces that runs. +# # Units follow the project convention: eV and Å. The strain is dimensionless, # so its input row carries "1". @@ -32,10 +38,14 @@ observables: per_atom: false units: "eV" derivatives: - # Named `forces`, reported as -dE/dpos. - wrt: pos + name: forces + sign: -1 units: "eV/Å" - # Named `stress`, reported as +dE/dstrain, divided by the cell volume by - # whatever computes it. The division is not part of the sign. + # The volume division that turns the strain derivative into a stress is + # not a sign, and does not belong here. It is applied by whatever + # computes the stress. - wrt: strain + name: stress + sign: +1 units: "eV/Å^3" diff --git a/packages/mace-core/src/mace_core/observables/__init__.py b/packages/mace-core/src/mace_core/observables/__init__.py index 83067c767..db656dd7c 100644 --- a/packages/mace-core/src/mace_core/observables/__init__.py +++ b/packages/mace-core/src/mace_core/observables/__init__.py @@ -6,9 +6,9 @@ load_default_catalogue, ) from mace_core.observables.derivatives import ( - SPECIAL_CASES, - derivative_name, - derivative_sign, + DEFAULT_SIGN, + default_derivative_name, + is_default_shaped_name, ) from mace_core.observables.grammar import ( IRREPS_GRAMMAR, @@ -27,8 +27,8 @@ __all__ = [ "DEFAULTS_RESOURCE", + "DEFAULT_SIGN", "IRREPS_GRAMMAR", - "SPECIAL_CASES", "DerivativeRequest", "DerivativeSpec", "InputSpec", @@ -36,9 +36,9 @@ "IrrepsGrammarError", "ObservableCatalogue", "ObservableSpec", - "derivative_name", - "derivative_sign", + "default_derivative_name", "irreps_dimension", + "is_default_shaped_name", "load_catalogue", "load_default_catalogue", "parse_irreps", diff --git a/packages/mace-core/src/mace_core/observables/derivatives.py b/packages/mace-core/src/mace_core/observables/derivatives.py index f0817b916..f63708190 100644 --- a/packages/mace-core/src/mace_core/observables/derivatives.py +++ b/packages/mace-core/src/mace_core/observables/derivatives.py @@ -1,65 +1,59 @@ -"""How a derivative of a declared quantity is named and signed. +"""How a derivative of a declared quantity is named by default. The rule is one line: the derivative of a declared quantity ``q`` with respect -to a declared input ``x`` is called ``d__d_``. Three pairs have a name of -their own, and they are data in the table below rather than branches spread -through the consumers that need them. - -The third special case is the reason this grammar is written over declared -inputs rather than over positions and the strain. ``magforces`` is -``-dE/d(magmom)``, computed in the same autograd call as the forces, trained -with its own loss term, and used by the magnetic self-consistent model to drive -its fixed point. A grammar that knew only ``d__d_pos`` and ``d__d_strain`` -could not express it, and the magnetic work would have had to go around the -abstraction that exists to prevent exactly that. - -The sign is the one the reported quantity carries, so that -``reported = sign * d(quantity)/d(input)``. The volume division that turns the -strain derivative into a stress is not a sign and is not here: it belongs to -whatever computes the stress. +to a declared input ``x`` is called ``d__d_`` and is reported with the +gradient's own sign. + +A quantity whose derivative has a name of its own, and a sign of its own, says +so **in the declarations file**. It used to say so in a table here, and that +table was the same fact written a third time: the machine-readable pair lived +here, the prose lived in :mod:`mace_core.units`, and the declarations file +carried it in a comment 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, +holding a row for ``("energy", "magmom")`` that no shipped catalogue declares. + +The consequence that matters is not tidiness. With the names in code, a fourth +one, torques as ``-dE/d(orientation)`` or a polarizability as +``d(dipole)/d(field)``, meant editing this package, which is exactly what a +declarative grammar exists to prevent. """ from __future__ import annotations +import re + __all__ = [ - "SPECIAL_CASES", - "derivative_name", - "derivative_sign", + "DEFAULT_SIGN", + "default_derivative_name", + "is_default_shaped_name", ] -#: ``(quantity, input) -> (name, sign)`` for the three pairs whose name is not -#: the ``d__d_`` default. Everything else follows the rule. -SPECIAL_CASES: dict[tuple[str, str], tuple[str, int]] = { - ("energy", "pos"): ("forces", -1), - ("energy", "strain"): ("stress", +1), - ("energy", "magmom"): ("magforces", -1), -} +#: A derivative is reported with the gradient's own sign unless its declaration +#: says otherwise. ``reported = sign * d(quantity)/d(input)``. +DEFAULT_SIGN: int = 1 + +_DEFAULT_SHAPE = re.compile(r"^d_.+_d_.+$") -def derivative_name(quantity: str, wrt: str) -> str: - """The canonical name of ``d(quantity)/d(wrt)``. +def default_derivative_name(quantity: str, wrt: str) -> str: + """The name of ``d(quantity)/d(wrt)`` when the declaration gives none. Args: quantity: The name of the differentiated observable. wrt: The name of the declared input it is differentiated against. Returns: - The special-cased name if the pair has one, otherwise ``f"d_{quantity}_d_{wrt}"``. """ - special = SPECIAL_CASES.get((quantity, wrt)) - if special is not None: - return special[0] return f"d_{quantity}_d_{wrt}" -def derivative_sign(quantity: str, wrt: str) -> int: - """The sign the reported derivative carries: ``reported = sign * dq/dx``. +def is_default_shaped_name(name: str) -> bool: + """Whether ``name`` is spelled like one this rule would generate. - ``+1`` unless the pair is one of the two negated special cases, forces and - magnetic forces, which are both the negative gradient of the energy. + A declared name that is shaped like the grammar's own but belongs to a + different pair reads as a fact about which quantity was differentiated, and + is not one. The catalogue refuses it rather than letting the two spellings + mean different things. """ - special = SPECIAL_CASES.get((quantity, wrt)) - if special is not None: - return special[1] - return 1 + return bool(_DEFAULT_SHAPE.match(name)) diff --git a/packages/mace-core/src/mace_core/observables/spec.py b/packages/mace-core/src/mace_core/observables/spec.py index e268f72dd..d683b275e 100644 --- a/packages/mace-core/src/mace_core/observables/spec.py +++ b/packages/mace-core/src/mace_core/observables/spec.py @@ -31,7 +31,11 @@ from pydantic import BaseModel, ConfigDict, Field, model_validator -from mace_core.observables.derivatives import derivative_name, derivative_sign +from mace_core.observables.derivatives import ( + DEFAULT_SIGN, + default_derivative_name, + is_default_shaped_name, +) from mace_core.observables.grammar import ( IrrepTerm, irreps_dimension, @@ -94,16 +98,28 @@ def _validate(self) -> InputSpec: class DerivativeRequest(BaseModel): """A derivative an observable asks for. - The name and the sign are not here: they are derived, and letting a - declaration override them would reintroduce the per-consumer naming this - abstraction removes. What a declaration owns is which input to - differentiate against, and the unit string to report. + Both the name and the sign belong to the declaration. Deriving them from a + table in code was the same fact written three times, and it meant a + quantity with a name of its own could not be added without editing this + package. The name is still decided **once**, here, and every consumer reads + it off the resolved spec, so nothing about this lets a consumer invent one. + + Giving a ``name`` obliges the declaration to give a ``sign`` too. A renamed + derivative that silently inherited ``+1`` is a model trained on inverted + forces that runs perfectly well, which is a failure nothing downstream can + see. """ model_config = ConfigDict(extra="forbid", frozen=True) #: The name of the declared input to differentiate against. wrt: str + #: What the derivative is called. ``None`` takes the grammar's own + #: ``d__d_``. + name: str | None = None + #: ``reported = sign * d(quantity)/d(input)``, either ``+1`` or ``-1``. + #: ``None`` takes the gradient's own sign. + sign: int | None = None #: Left to the declaration. Deriving it would mean unit algebra over the #: quantity and the input, which this ticket does not own. units: str | None = None @@ -116,6 +132,34 @@ def _accept_bare_name(cls, value: object) -> object: return {"wrt": value} return value + @model_validator(mode="after") + def _validate(self) -> DerivativeRequest: + if self.sign is not None and self.sign not in (1, -1): + raise ValueError( + f"the derivative with respect to {self.wrt!r} declares sign " + f"{self.sign!r}. A sign is +1 or -1; a scale factor is not a " + f"sign and belongs to whatever computes the quantity." + ) + if self.name is not None: + _check_name(self.name, "derivative") + if self.sign is None: + raise ValueError( + f"the derivative with respect to {self.wrt!r} is named " + f"{self.name!r} but declares no sign. A name of its own " + f"means a convention of its own, so state it: `sign: -1` " + f"for a quantity reported as the negative gradient, " + f"`sign: +1` otherwise." + ) + if is_default_shaped_name(self.name): + raise ValueError( + f"the derivative with respect to {self.wrt!r} is named " + f"{self.name!r}, which is spelled like the grammar's own " + f"`d__d_`. That spelling states which " + f"quantity was differentiated, so a custom name must not " + f"use it. Drop `name` to get the generated one." + ) + return self + class DerivativeSpec(BaseModel): """A derivative, as resolved by the catalogue. Never declared directly.""" @@ -188,13 +232,30 @@ def is_scalar(self) -> bool: """Whether the declaration is a single ``0e``.""" return parse_irreps(self.irreps, observable=self.name) == _SCALAR + def _request(self, wrt: str) -> DerivativeRequest | None: + """This observable's declared request against ``wrt``, if it made one. + + Naming works for any declared input whether or not it was requested, so + this is allowed to find nothing. + """ + for request in self.derivatives: + if request.wrt == wrt: + return request + return None + def derivative_name(self, wrt: str) -> str: """The canonical name of this observable's derivative against ``wrt``.""" - return derivative_name(self.name, wrt) + request = self._request(wrt) + if request is not None and request.name is not None: + return request.name + return default_derivative_name(self.name, wrt) def derivative_sign(self, wrt: str) -> int: """The sign that derivative is reported with.""" - return derivative_sign(self.name, wrt) + request = self._request(wrt) + if request is not None and request.sign is not None: + return request.sign + return DEFAULT_SIGN def requested_derivatives(self) -> tuple[str, ...]: """The inputs this observable asked to be differentiated against.""" diff --git a/packages/mace-core/tests/test_observables.py b/packages/mace-core/tests/test_observables.py index 5935662cc..72d6d85f9 100644 --- a/packages/mace-core/tests/test_observables.py +++ b/packages/mace-core/tests/test_observables.py @@ -5,15 +5,18 @@ ``mace_core`` to pass, the abstraction would not be doing its job. """ +from importlib.resources import files + import pytest +import yaml from mace_core.observables import ( + DEFAULTS_RESOURCE, DerivativeRequest, InputSpec, IrrepsGrammarError, ObservableCatalogue, ObservableSpec, - derivative_name, - derivative_sign, + default_derivative_name, irreps_dimension, load_catalogue, load_default_catalogue, @@ -137,19 +140,6 @@ def test_units_may_not_be_empty(): # --------------------------------------------------------------------------- -@pytest.mark.parametrize( - ("quantity", "wrt", "name", "sign"), - [ - ("energy", "pos", "forces", -1), - ("energy", "strain", "stress", +1), - ("energy", "magmom", "magforces", -1), - ], -) -def test_the_three_special_cases_keep_their_names_and_signs(quantity, wrt, name, sign): - assert derivative_name(quantity, wrt) == name - assert derivative_sign(quantity, wrt) == sign - - @pytest.mark.parametrize( ("quantity", "wrt", "name"), [ @@ -160,9 +150,62 @@ def test_the_three_special_cases_keep_their_names_and_signs(quantity, wrt, name, ("quadrupole", "magmom", "d_quadrupole_d_magmom"), ], ) -def test_everything_else_follows_the_rule(quantity, wrt, name): - assert derivative_name(quantity, wrt) == name - assert derivative_sign(quantity, wrt) == +1 +def test_an_undeclared_name_follows_the_rule(quantity, wrt, name): + assert default_derivative_name(quantity, wrt) == name + spec = ObservableSpec(name=quantity, irreps="0e", per_atom=False, units="eV") + assert spec.derivative_name(wrt) == name + assert spec.derivative_sign(wrt) == +1 + + +def test_a_declared_name_and_sign_win_over_the_rule(): + """The whole point: a quantity with a name of its own says so in the file.""" + spec = ObservableSpec( + name="energy", + irreps="0e", + per_atom=False, + units="eV", + derivatives=[{"wrt": "magmom", "name": "magforces", "sign": -1}], + ) + assert spec.derivative_name("magmom") == "magforces" + assert spec.derivative_sign("magmom") == -1 + + +def test_an_input_that_was_not_requested_still_has_a_name(): + """Naming works for any declared input, requested or not.""" + spec = ObservableSpec( + name="energy", + irreps="0e", + per_atom=False, + units="eV", + derivatives=[{"wrt": "pos", "name": "forces", "sign": -1}], + ) + assert spec.derivative_name("strain") == "d_energy_d_strain" + assert spec.derivative_sign("strain") == +1 + + +def test_a_name_without_a_sign_is_refused(): + """A renamed derivative inheriting +1 in silence trains inverted forces.""" + with pytest.raises(ValidationError, match="declares no sign"): + DerivativeRequest(wrt="pos", name="forces") + + +@pytest.mark.parametrize("sign", [0, 2, -3]) +def test_a_sign_that_is_not_plus_or_minus_one_is_refused(sign): + with pytest.raises(ValidationError, match="A sign is \\+1 or -1"): + DerivativeRequest(wrt="pos", name="forces", sign=sign) + + +def test_a_custom_name_may_not_imitate_the_generated_spelling(): + """`d__d_` states which quantity was differentiated. A custom name + wearing that spelling would be stating something that is not true.""" + with pytest.raises(ValidationError, match="spelled like the grammar"): + DerivativeRequest(wrt="pos", name="d_dipole_d_pos", sign=+1) + + +def test_a_sign_alone_needs_no_name(): + request = DerivativeRequest(wrt="pos", sign=-1) + assert request.name is None + assert request.sign == -1 # --------------------------------------------------------------------------- @@ -173,6 +216,30 @@ def test_everything_else_follows_the_rule(quantity, wrt, name): def test_the_defaults_declare_energy_and_its_two_derivatives(): catalogue = load_default_catalogue() assert catalogue.names() == ("energy", "forces", "stress") + + +def test_the_shipped_names_and_signs_come_from_the_file_and_not_from_code(): + """The reason the table was removed: this is now a property of the data. + + If these were still special-cased in code, the assertion would pass with + the declarations file saying nothing at all. + """ + catalogue = load_default_catalogue() + resolved = {spec.name: spec.sign for spec in catalogue.requested_derivatives()} + assert resolved == {"forces": -1, "stress": +1} + + stripped = yaml.safe_load( + files("mace_core").joinpath(DEFAULTS_RESOURCE).read_text(encoding="utf-8") + ) + for observable in stripped["observables"]: + for request in observable.get("derivatives", []): + request.pop("name", None) + request.pop("sign", None) + bare = ObservableCatalogue.model_validate(stripped) + assert {spec.name for spec in bare.requested_derivatives()} == { + "d_energy_d_pos", + "d_energy_d_strain", + } assert [spec.name for spec in catalogue.inputs] == ["pos", "strain"] @@ -222,7 +289,12 @@ def test_a_new_rank_two_per_atom_observable_is_a_row_in_yaml(tmp_path): def test_a_new_input_feature_makes_its_derivative_declarable(tmp_path): """`magmom` is the case that pays for the grammar being written over - declared inputs rather than over positions and the strain.""" + declared inputs rather than over positions and the strain. + + It is also the case that pays for the name and the sign living in the file. + `magforces` used to be a row in a table inside this package, and this + catalogue reaches it with no code at all. + """ catalogue = catalogue_from( """ - name: magmom @@ -237,6 +309,8 @@ def test_a_new_input_feature_makes_its_derivative_declarable(tmp_path): units: "eV" derivatives: - wrt: magmom + name: magforces + sign: -1 units: "eV/muB" """, tmp_path, @@ -294,7 +368,7 @@ def test_a_derived_name_may_not_collide_with_a_declared_observable(tmp_path): catalogue_from( """ observables: - - name: forces + - name: d_energy_d_pos irreps: "1o" per_atom: true units: "eV/Å" @@ -303,6 +377,30 @@ def test_a_derived_name_may_not_collide_with_a_declared_observable(tmp_path): per_atom: false units: "eV" derivatives: [pos] +""", + tmp_path, + ) + assert "'d_energy_d_pos'" in str(caught.value) + + +def test_a_declared_name_may_not_collide_with_a_declared_observable(tmp_path): + """The same guard, now reachable through a name the file chose itself.""" + with pytest.raises(ValidationError) as caught: + catalogue_from( + """ +observables: + - name: forces + irreps: "1o" + per_atom: true + units: "eV/Å" + - name: energy + irreps: "0e" + per_atom: false + units: "eV" + derivatives: + - wrt: pos + name: forces + sign: -1 """, tmp_path, ) diff --git a/tests/architecture/observable_coverage.py b/tests/architecture/observable_coverage.py index bb918ad11..559136f6b 100644 --- a/tests/architecture/observable_coverage.py +++ b/tests/architecture/observable_coverage.py @@ -34,7 +34,7 @@ from pathlib import Path from typing import Union -from mace_core.observables import derivative_name +from mace_core.observables import default_derivative_name from tests.golden import harness, surface_scan @@ -133,12 +133,18 @@ class Derivative: two agree, with no way to annotate a disagreement: a pair the rule gets wrong is either a misclassified row or a real gap in the grammar, and both have to be resolved rather than recorded. + name: What v1 calls it, when that is not the rule's own + ``d__d_``. Stated here rather than read out of the + declarations file on purpose: this table is the independent + statement of what the frozen tree does, and a test that derived + both sides from the same source would agree with itself. note: Why the row is worth a second look, where it is. """ of: str wrt: str sign: int + name: str = "" note: str = "" @@ -339,12 +345,14 @@ class Drop: ), # --- derivatives ------------------------------------------------------- "forces": Derivative( + name="forces", of="energy", wrt="pos", sign=-1, note="mace/modules/utils.py:115 returns `-1 * forces`.", ), "stress": Derivative( + name="stress", of="energy", wrt="strain", sign=+1, @@ -355,6 +363,7 @@ class Drop: ), ), "magforces": Derivative( + name="magforces", of="energy", wrt="magmom", sign=-1, @@ -487,7 +496,7 @@ def v1_name(key: str) -> str: """The name ``key`` carries in v1: renamed, rule-derived, or unchanged.""" row = DISPOSITIONS[key] if isinstance(row, Derivative): - return derivative_name(row.of, row.wrt) + return row.name or default_derivative_name(row.of, row.wrt) if isinstance(row, Spec) and row.renamed_to: return row.renamed_to return key diff --git a/tests/architecture/test_observable_completeness.py b/tests/architecture/test_observable_completeness.py index 9cc2bcb9b..a34638b2b 100644 --- a/tests/architecture/test_observable_completeness.py +++ b/tests/architecture/test_observable_completeness.py @@ -42,7 +42,12 @@ "job, or `pip install -e packages/mace-core`", allow_module_level=True, ) -from mace_core.observables import ObservableSpec, derivative_name, derivative_sign +from mace_core.observables import ( + DEFAULT_SIGN, + ObservableSpec, + default_derivative_name, + load_default_catalogue, +) from mace_core.outputs import CORE_FIELD_NAMES, FIELD_BY_OBSERVABLE @@ -163,30 +168,55 @@ def test_a_derivative_row_resolves_through_the_rule(key): f"declared observable or a derivative of one. A derivative chain has " f"to ground out in something declared." ) - name = derivative_name(row.of, row.wrt) + name = row.name or default_derivative_name(row.of, row.wrt) assert name.isidentifier() - derived_sign = derivative_sign(row.of, row.wrt) - assert row.sign == derived_sign, ( - f"{key!r} is reported with sign {row.sign:+d} by the frozen tree and " - f"the rule derives {derived_sign:+d}. There is no way to annotate that " - f"away, deliberately: either the row differentiates the wrong thing, " - f"or the grammar has a real gap. `hessian` was the first candidate for " - f"such a gap and turned out to be the former: it is a second " - f"derivative of the energy, not a first derivative of the forces." + if row.name: + # A row with a name of its own also carries the sign that goes with it, + # and a declaration is what supplies both. What the rule can still be + # checked against is that the pair is expressible: an `ObservableSpec` + # built from this row has to resolve to exactly these two values. + spec = ObservableSpec( + name=row.of, + irreps="0e", + per_atom=False, + units="1", + derivatives=[{"wrt": row.wrt, "name": row.name, "sign": row.sign}], + ) + assert spec.derivative_name(row.wrt) == row.name + assert spec.derivative_sign(row.wrt) == row.sign + return + assert row.sign == DEFAULT_SIGN, ( + f"{key!r} is reported with sign {row.sign:+d} by the frozen tree, and " + f"it carries no name of its own, so the rule derives " + f"{DEFAULT_SIGN:+d}. A negated quantity has a convention of its own " + f"and therefore a name of its own: give the row a `name`, or the row " + f"differentiates the wrong thing. `hessian` was the first candidate " + f"for a real gap in the grammar and turned out to be the latter: it " + f"is a second derivative of the energy, not a first derivative of the " + f"forces." ) -def test_the_three_named_derivatives_keep_their_legacy_names(): +def test_the_named_derivatives_keep_their_legacy_names(): """forces, stress and magforces are the pairs that have a name of their own, - and the two negated ones are the two the frozen tree negates.""" - assert derivative_name("energy", "pos") == "forces" - assert derivative_sign("energy", "pos") == -1 - assert derivative_name("energy", "strain") == "stress" - assert derivative_sign("energy", "strain") == +1 - assert derivative_name("energy", "magmom") == "magforces" - assert derivative_sign("energy", "magmom") == -1 + and the two negated ones are the two the frozen tree negates. + + The shipped declarations are the other side of this comparison, and they are + a genuinely separate source: this table is read off the frozen tree, and + that file is authored. Where both name a pair they have to agree, and a pair + the file does not ship yet, `magforces`, is simply absent rather than wrong. + """ + declared = { + spec.name: spec for spec in load_default_catalogue().requested_derivatives() + } for name in ("forces", "stress", "magforces"): - assert isinstance(DISPOSITIONS[name], Derivative) + row = DISPOSITIONS[name] + assert isinstance(row, Derivative) + assert row.name == name + if name not in declared: + continue + assert (declared[name].of, declared[name].wrt) == (row.of, row.wrt) + assert declared[name].sign == row.sign def test_the_renamed_derivatives_are_renamed_and_not_lost(): @@ -204,7 +234,8 @@ def test_the_renamed_derivatives_are_renamed_and_not_lost(): } for key, expected in renamed.items(): row = DISPOSITIONS[key] - assert derivative_name(row.of, row.wrt) == expected + assert not row.name + assert default_derivative_name(row.of, row.wrt) == expected assert isinstance(DISPOSITIONS["BEC"], Spec) @@ -216,7 +247,7 @@ def test_no_two_derivative_rows_resolve_to_one_name(): for key, row in DISPOSITIONS.items(): if not isinstance(row, Derivative): continue - name = derivative_name(row.of, row.wrt) + name = row.name or default_derivative_name(row.of, row.wrt) assert name not in seen, ( f"{key!r} and {seen[name]!r} both resolve to {name!r}. If they are " f"one quantity, say so and drop one row; if they are not, one of " From cceed846d0bc4e3000256db6ed3e4e34e4a46a6d Mon Sep 17 00:00:00 2001 From: Alejandro Acosta <127198532+aacostadiaz@users.noreply.github.com> Date: Tue, 22 Sep 2026 06:44:02 +0100 Subject: [PATCH 09/12] Make names() and membership speak one vocabulary `__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. --- packages/mace-core/src/mace_core/outputs.py | 30 ++++++++++++-- packages/mace-core/tests/test_outputs.py | 43 ++++++++++++++++++++- 2 files changed, 68 insertions(+), 5 deletions(-) diff --git a/packages/mace-core/src/mace_core/outputs.py b/packages/mace-core/src/mace_core/outputs.py index 9efe392d7..3d2346374 100644 --- a/packages/mace-core/src/mace_core/outputs.py +++ b/packages/mace-core/src/mace_core/outputs.py @@ -32,6 +32,7 @@ __all__ = [ "CORE_FIELD_NAMES", "FIELD_BY_OBSERVABLE", + "OBSERVABLE_BY_FIELD", "MACEOutput", "TensorT", ] @@ -118,12 +119,25 @@ def get(self, name: str) -> TensorT | None: return self.extras.get(name) def names(self) -> tuple[str, ...]: - """Every name that carries a value, core fields first, then ``extras``. + """Every **observable** name that carries a value, core fields first. A core field holding ``None`` was not computed and is left out, so this is what the model actually produced rather than what it could produce. + + The names are the ones a declaration uses, not the storage fields, so + ``total_energy`` appears here as ``energy``. That is what makes + ``set(catalogue.names()) & set(output.names())`` mean what it reads as: + yielding the field name instead put the two vocabularies one alias + apart, and the intersection dropped the energy in silence while + ``"energy" in output`` was `True` the whole time. Use :meth:`get` to + reach a value, which accepts either spelling; ``getattr`` over these is + the one thing they are not for. """ - present = [name for name in CORE_FIELD_NAMES if getattr(self, name) is not None] + present = [ + OBSERVABLE_BY_FIELD.get(name, name) + for name in CORE_FIELD_NAMES + if getattr(self, name) is not None + ] present.extend(self.extras) return tuple(present) @@ -132,7 +146,10 @@ def __contains__(self, name: str) -> bool: #: The six fields that are part of the type. Derived from the dataclass rather -#: than written out again, so the two cannot disagree. +#: than written out again, so the two cannot disagree. These are **storage** +#: names: `total_energy` appears here and `energy` does not. Correlating them +#: with a catalogue's observable names needs :data:`OBSERVABLE_BY_FIELD`, or +#: :meth:`MACEOutput.names`, which has already applied it. CORE_FIELD_NAMES: tuple[str, ...] = tuple( f.name for f in fields(MACEOutput) if f.name != "extras" ) @@ -143,3 +160,10 @@ def __contains__(self, name: str) -> bool: #: grammar's special cases are keyed on (``energy`` + positions -> ``forces``). #: Written down as one entry rather than left to each consumer to remember. FIELD_BY_OBSERVABLE: dict[str, str] = {"energy": "total_energy"} + +#: The same map read the other way, for going from storage back to the name a +#: declaration uses. Inverted here rather than written out, so an entry added +#: to one direction cannot be missing from the other. +OBSERVABLE_BY_FIELD: dict[str, str] = { + field: observable for observable, field in FIELD_BY_OBSERVABLE.items() +} diff --git a/packages/mace-core/tests/test_outputs.py b/packages/mace-core/tests/test_outputs.py index b4bb927f5..19f62fb24 100644 --- a/packages/mace-core/tests/test_outputs.py +++ b/packages/mace-core/tests/test_outputs.py @@ -11,7 +11,13 @@ import numpy as np import pytest -from mace_core.outputs import CORE_FIELD_NAMES, FIELD_BY_OBSERVABLE, MACEOutput +from mace_core.observables import load_default_catalogue +from mace_core.outputs import ( + CORE_FIELD_NAMES, + FIELD_BY_OBSERVABLE, + OBSERVABLE_BY_FIELD, + MACEOutput, +) def test_the_six_core_fields_are_the_declared_ones(): @@ -38,12 +44,13 @@ def test_core_fields_round_trip_numpy_arrays(): output = MACEOutput(total_energy=np.array([-1.5]), forces=forces) assert output.get("total_energy") is output.total_energy assert output.get("forces") is forces - assert output.names() == ("total_energy", "forces") + assert output.names() == ("energy", "forces") def test_energy_reaches_the_total_energy_field_under_either_name(): """The one place an observable name and a field name differ.""" assert FIELD_BY_OBSERVABLE == {"energy": "total_energy"} + assert OBSERVABLE_BY_FIELD == {"total_energy": "energy"} output = MACEOutput(total_energy=np.array([2.0])) assert output.get("energy") is output.get("total_energy") assert "energy" in output @@ -116,3 +123,35 @@ def test_a_name_that_is_not_a_core_field_is_fine_in_extras(): while the field that owns it stays `None`. v1 renames that key instead.""" out = MACEOutput(extras={"latent_charges": np.zeros(3)}) assert out.names() == ("latent_charges",) + + +def test_membership_and_listing_agree_on_every_name(): + """The two accessors have to use one vocabulary. + + They did not: `"energy" in output` resolved the alias and `names()` yielded + the storage field, so a consumer intersecting a catalogue's names with an + output's dropped the energy while membership said it was there. Asserted + over both directions rather than over the one alias, so a second entry in + `FIELD_BY_OBSERVABLE` cannot reopen it. + """ + output = MACEOutput( + total_energy=np.array([-1.5]), + forces=np.zeros((4, 3)), + extras={"latent_charges": np.zeros(4)}, + ) + for name in output.names(): + assert name in output, name + assert output.get(name) is not None, name + for observable, field in FIELD_BY_OBSERVABLE.items(): + assert field not in output.names() + if getattr(output, field) is not None: + assert observable in output.names() + + +def test_a_catalogue_name_finds_its_value_in_an_output(): + """The correlation the two vocabularies exist to allow.""" + catalogue = load_default_catalogue() + output = MACEOutput(total_energy=np.array([-1.5]), forces=np.zeros((4, 3))) + shared = set(catalogue.names()) & set(output.names()) + assert shared == {"energy", "forces"} + assert all(output.get(name) is not None for name in shared) From 86080f5e1e425e3247113be75c2520072fb64b97 Mon Sep 17 00:00:00 2001 From: Alejandro Acosta <127198532+aacostadiaz@users.noreply.github.com> Date: Tue, 22 Sep 2026 06:57:27 +0100 Subject: [PATCH 10/12] Pin the tensor parameter in the new output test `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, ]`. 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. --- packages/mace-core/tests/test_outputs.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/packages/mace-core/tests/test_outputs.py b/packages/mace-core/tests/test_outputs.py index 19f62fb24..eb0d8e9e2 100644 --- a/packages/mace-core/tests/test_outputs.py +++ b/packages/mace-core/tests/test_outputs.py @@ -134,7 +134,13 @@ def test_membership_and_listing_agree_on_every_name(): over both directions rather than over the one alias, so a second entry in `FIELD_BY_OBSERVABLE` cannot reopen it. """ - output = MACEOutput( + # Pinned to `np.ndarray` rather than left to inference. numpy types an + # array's shape, so inferring the parameter from a 1-D energy beside a 2-D + # force array makes it a union, and `dict` is invariant in its value type: + # the `extras` literal is then unassignable. Nothing about the class, and + # widening `extras` to a `Mapping` to quiet it would make a result object + # mutable-looking to fix a test. + output: MACEOutput[np.ndarray] = MACEOutput( total_energy=np.array([-1.5]), forces=np.zeros((4, 3)), extras={"latent_charges": np.zeros(4)}, From b1d0af7740b1cedc6d21218264eee9517bdc9af8 Mon Sep 17 00:00:00 2001 From: Alejandro Acosta <127198532+aacostadiaz@users.noreply.github.com> Date: Tue, 22 Sep 2026 06:58:33 +0100 Subject: [PATCH 11/12] Refuse a retired spelling as an extras key 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. --- packages/mace-core/src/mace_core/outputs.py | 22 +++++++++++ packages/mace-core/tests/test_outputs.py | 41 +++++++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/packages/mace-core/src/mace_core/outputs.py b/packages/mace-core/src/mace_core/outputs.py index 3d2346374..62fbe3fe9 100644 --- a/packages/mace-core/src/mace_core/outputs.py +++ b/packages/mace-core/src/mace_core/outputs.py @@ -33,6 +33,7 @@ "CORE_FIELD_NAMES", "FIELD_BY_OBSERVABLE", "OBSERVABLE_BY_FIELD", + "RETIRED_NAMES", "MACEOutput", "TensorT", ] @@ -104,6 +105,17 @@ def __post_init__(self) -> None: f"keys of `extras`: a consumer reading the field would see " f"nothing. Assign them as fields instead." ) + retired = sorted(name for name in self.extras if name in RETIRED_NAMES) + if retired: + replacements = ", ".join( + f"{name} -> {RETIRED_NAMES[name]}" for name in retired + ) + raise ValueError( + f"{retired} are spellings this type retired and cannot be keys " + f"of `extras`: the value would sit beside the field holding the " + f"same quantity, which is the dual storage the rename exists to " + f"prevent. Assign the field instead ({replacements})." + ) def get(self, name: str) -> TensorT | None: """The value stored under ``name``, or ``None`` if there is none. @@ -154,6 +166,16 @@ def __contains__(self, name: str) -> bool: f.name for f in fields(MACEOutput) if f.name != "extras" ) +#: A legacy spelling this type does not carry, and the field that replaced it. +#: These are **not** resolvable names: :meth:`MACEOutput.get` does not find a +#: field through them, because carrying both spellings is exactly what the +#: rename decided against. What they are is refused as ``extras`` keys, which +#: is the half a rename cannot enforce on its own: without this, `extras` +#: ``node_energy`` sits happily beside a filled ``node_energies`` field and the +#: two hold the same quantity, which is the case the shadowing guard above +#: exists to make impossible for the names it does know. +RETIRED_NAMES: dict[str, str] = {"node_energy": "node_energies"} + #: The one place an observable's name and its storage field differ. The field #: says "total" because the type also carries per-atom energies, while the #: observable is named ``energy`` because that is the name the derivative diff --git a/packages/mace-core/tests/test_outputs.py b/packages/mace-core/tests/test_outputs.py index eb0d8e9e2..68462b2af 100644 --- a/packages/mace-core/tests/test_outputs.py +++ b/packages/mace-core/tests/test_outputs.py @@ -16,6 +16,7 @@ CORE_FIELD_NAMES, FIELD_BY_OBSERVABLE, OBSERVABLE_BY_FIELD, + RETIRED_NAMES, MACEOutput, ) @@ -161,3 +162,43 @@ def test_a_catalogue_name_finds_its_value_in_an_output(): shared = set(catalogue.names()) & set(output.names()) assert shared == {"energy", "forces"} assert all(output.get(name) is not None for name in shared) + + +def test_a_retired_spelling_cannot_sit_beside_the_field_that_replaced_it(): + """The half a rename cannot enforce on its own. + + `node_energy` is the frozen tree's spelling and v1 carries `node_energies` + instead, deliberately: keeping both would let the same quantity be stored + twice. The rename alone does not stop it, because `extras` accepts any key + the guard does not know, so `extras["node_energy"]` sat happily beside a + filled `node_energies` field. + """ + with pytest.raises(ValueError, match="retired"): + MACEOutput(node_energies=np.zeros(4), extras={"node_energy": np.ones(4)}) + + +def test_a_retired_spelling_is_still_not_a_name_this_type_answers_to(): + """Refusing it as an `extras` key is not the same as reviving it. + + Carrying both spellings is what the rename decided against, so `get` must + keep missing it. Otherwise the guard above would have quietly turned the + retired name back into a working alias. + """ + output = MACEOutput(node_energies=np.zeros(4)) + assert output.get("node_energy") is None + assert "node_energy" not in output + assert output.names() == ("node_energies",) + assert set(RETIRED_NAMES) & set(CORE_FIELD_NAMES) == set() + assert set(RETIRED_NAMES) & set(FIELD_BY_OBSERVABLE) == set() + + +def test_a_retired_spelling_is_free_as_an_extras_key_when_nothing_replaced_it(): + """The guard is about the collision, not about the word. + + A file that carries the legacy spelling and no `node_energies` is still + refused, because the point is that the two names mean one quantity and only + one of them is this type's. Stated as its own case so a later relaxation to + "only when the field is filled" is a deliberate change rather than a slip. + """ + with pytest.raises(ValueError, match="retired"): + MACEOutput(extras={"node_energy": np.ones(4)}) From 1cb4fed4648b2fe05b1299df4965247929d51bbe Mon Sep 17 00:00:00 2001 From: Alejandro Acosta <127198532+aacostadiaz@users.noreply.github.com> Date: Thu, 24 Sep 2026 07:56:44 +0100 Subject: [PATCH 12/12] Declare the default observables in Python instead of a packaged file 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 --- docs/reforge/extending_mace.md | 31 ++- docs/reforge/extending_plugin.md | 9 +- docs/reforge/target_layout.md | 12 +- packages/mace-core/pyproject.toml | 1 - packages/mace-core/src/mace_core/__init__.py | 4 +- .../src/mace_core/defaults/observables.yaml | 51 ---- .../src/mace_core/observables/__init__.py | 10 +- .../src/mace_core/observables/defaults.py | 101 ++++---- packages/mace-core/tests/test_observables.py | 243 +++++++----------- packages/mace-core/tests/test_outputs.py | 4 +- .../test_observable_completeness.py | 15 +- 11 files changed, 181 insertions(+), 300 deletions(-) delete mode 100644 packages/mace-core/src/mace_core/defaults/observables.yaml diff --git a/docs/reforge/extending_mace.md b/docs/reforge/extending_mace.md index 79fe9189f..aea12b383 100644 --- a/docs/reforge/extending_mace.md +++ b/docs/reforge/extending_mace.md @@ -86,12 +86,13 @@ scalars). The magnetic output is the moment's **conjugate force**, `magforces = by autograd exactly like `forces = −dE/dpositions`. It is a per-atom `1o` vector, declared as a derivative observable: -```yaml -# mace_torch/extras/magnetic/observables.yaml — the extra ships its own observable rows -magforces: - derivation: "autograd(energy, wrt=magmom)" # -dE/dmagmom - per_atom: true - irreps: "1o" # a 3-vector, same convention as magmom +```python +# mace_torch/extras/magnetic/observables.py — the extra ships its own declarations +from mace_core.observables import DerivativeRequest, InputSpec + +MAGMOM = InputSpec(name="magmom", irreps="1e", per_atom=True, units="muB") +# -dE/dmagmom, declared on the energy: a per-atom vector like the moment itself +MAGFORCES = DerivativeRequest(wrt="magmom", name="magforces", sign=-1, units="eV/muB") ``` ```toml @@ -111,12 +112,14 @@ readout is a one-row change in the observable table, so here is what a *predicte like. **This block is illustrative** — it is not in the current implementation; it is shown only to demonstrate the readout mechanism. -```yaml -# mace_torch/extras/magnetic/observables.yaml — illustrative, NOT in #1244 -magnetic_moment: - derivation: readout # a learned equivariant readout over node features - per_atom: true - irreps: "1o" # same convention as the magmom input +```python +# mace_torch/extras/magnetic/observables.py — illustrative, NOT in #1244 +from mace_core.observables import ObservableSpec + +# a learned equivariant readout over node features, same convention as magmom +MAGNETIC_MOMENT = ObservableSpec( + name="magnetic_moment", irreps="1e", per_atom=True, units="muB" +) ``` ```toml @@ -314,7 +317,7 @@ directory, wired by the `__init__.py` above: mace_torch/extras/magnetic/ # everything the feature owns lives here ├── __init__.py # the registration entry point (the mace.plugins target) ├── embedding.py # MagmomEmbedding (§1) -├── observables.yaml # the magforces row (§2) +├── observables.py # the magforces declaration (§2) ├── transforms.py # RotateMagmom (§3) └── model.py # MagneticScaleShiftMACE / MagneticSCFMACE (§5) ``` @@ -364,7 +367,7 @@ unchanged — only where the files sit and a couple of packaging details drop aw - The same modules move from `mace_torch/extras/magnetic/` into the main `mace-torch` tree (`mace_torch/nn/`, `mace_torch/data/`, `mace_torch/models/`), still registered by the same decorators. -- The observable row goes in the shared `defaults/observables.yaml` instead of a feature-local file. +- The observable declarations go in the shared `observables/defaults.py` instead of a feature-local module. - No `mace.plugins` entry point, no `[magnetic]` optional-dependency group, no capability marker — its dependencies (here `sphericart-torch`) would be base dependencies and its tests run unconditionally. diff --git a/docs/reforge/extending_plugin.md b/docs/reforge/extending_plugin.md index 02be106bc..cc2de3fbd 100644 --- a/docs/reforge/extending_plugin.md +++ b/docs/reforge/extending_plugin.md @@ -38,7 +38,7 @@ mace-magnetic/ # a separate repo / PyPI package — not in ├── src/mace_magnetic/ │ ├── __init__.py # the mace.plugins target: runs the @register_* decorators │ ├── embedding.py # MagmomEmbedding (extending_mace.md §1) -│ ├── observables.yaml # the magforces row (§2) +│ ├── observables.py # the magforces declaration (§2) │ ├── transforms.py # RotateMagmom (§3) │ └── model.py # MagneticScaleShiftMACE / MagneticSCFMACE (§5) └── tests/ @@ -63,16 +63,17 @@ edit to any MACE file. The module code is what [Extending MACE](extending_mace.md) §1, §3 and §5 already showed; **only the import root changes** (`mace_magnetic` instead of `mace_torch.extras.magnetic`). The `__init__.py` runs the -decorators and loads the observable rows: +decorators and registers the observable declarations: ```python # src/mace_magnetic/__init__.py from mace_torch.models import register_model from mace_torch.data import register_transform from mace_torch.nn import register_input_embedding -from mace_torch.observables import register_observables_yaml +from mace_torch.observables import register_observables from .embedding import MagmomEmbedding +from .observables import MAGMOM, MAGFORCES from .transforms import RotateMagmom from .model import MagneticScaleShiftMACE, MagneticSCFMACE @@ -80,7 +81,7 @@ register_input_embedding("magmom")(MagmomEmbedding) register_transform("rotate_magmom")(RotateMagmom) register_model("MagneticScaleShiftMACE")(MagneticScaleShiftMACE) register_model("MagneticSCFMACE")(MagneticSCFMACE) -register_observables_yaml(__file__, "observables.yaml") # magforces +register_observables(inputs=[MAGMOM], derivatives=[MAGFORCES]) # magforces ``` `config.toml` is byte-for-byte the one from the in-tree example — it refers to the feature only by diff --git a/docs/reforge/target_layout.md b/docs/reforge/target_layout.md index b21d3e75f..0526e18be 100644 --- a/docs/reforge/target_layout.md +++ b/docs/reforge/target_layout.md @@ -47,14 +47,12 @@ packages/mace-core/ │ │ ├── number_table.py # AtomicNumberTable (reimplemented; mirror of tools/utils.py, without dragging in train.py) │ │ └── default_keys.py # DefaultKeys (reimplemented; mirror of tools/default_keys.py) │ │ -│ ├── observables/ # a property is a row in a declarations file, not a module per property +│ ├── observables/ # a property is a declaration, not a module per property │ │ ├── __init__.py # the public surface of the package │ │ ├── spec.py # InputSpec, ObservableSpec, DerivativeSpec, ObservableCatalogue (pydantic) │ │ ├── grammar.py # the irreps string grammar: parse + validate + dimension (no algebra) │ │ ├── derivatives.py # d__d_ naming, and the three special cases with their signs -│ │ └── defaults.py # loader for a declarations file -│ ├── defaults/ -│ │ └── observables.yaml # energy + its position and strain derivatives: the row every observable copies +│ │ └── defaults.py # DEFAULT_CATALOGUE: energy + its position and strain derivatives, the declaration every observable copies │ │ │ ├── kernels/ │ │ ├── protocol.py # KernelBackend Protocol, generic over TensorT: make_* factories + capabilities (§3.1) @@ -383,7 +381,7 @@ in the forward is a real code change: | You want to… | How | New code? | |---|---|---| | **Tune a parameter** — a loss weight, a Huber `delta`, a cutoff, a schedule, any hyperparameter | set a field in config | **none** | -| **Train a new property** — any well-defined spherical-tensor observable (a dipole, a rank-2 tensor, spectra, a magnetic moment) | add a **row to the observable table** (`ObservableSpec` in config, canonical defaults in `defaults/observables.yaml`) — it auto-creates the head, the loss term, and the derivative names | **none** | +| **Train a new property** — any well-defined spherical-tensor observable (a dipole, a rank-2 tensor, spectra, a magnetic moment) | add a **row to the observable table** (`ObservableSpec` in config, canonical defaults in `observables/defaults.py`) — it auto-creates the head, the loss term, and the derivative names | **none** | | **A new loss** — a non-standard reduction, or a data/relative-energy/mask transform | `@register_loss` / `@register_transform` + select it in config | a small module | | **A new readout / head** | `@register_readout` + config | a small module | | **A new backend** — kernel, data format, neighbour list, electrostatics solver | ship a wheel with one entry-point line (`mace.kernel_backends.torch`, `mace.data_backends`, `mace.neighbor_backends`, `mace.electrostatics_backends.torch`) | a backend module, **zero core edits** | @@ -493,9 +491,9 @@ silently wrong forces); it stays usable for inference (`supports_double_backward ### 3.2 A new observable (config only) -- **Extender touches:** a declarations file: one `ObservableSpec` row giving `name`, `irreps`, `per_atom`, `units`, and the declared inputs to differentiate against. The scaling of the head that produces it is set in the model config and its loss weight in `LossConfig`, both keyed by this name. No module, no decorator. A derivative is named by the rule `d__d_`, with `forces`, `stress` and `magforces` as the three special cases, so asking for a derivative against a newly declared input needs no code either. +- **Extender touches:** one `ObservableSpec` declaration giving `name`, `irreps`, `per_atom`, `units`, and the declared inputs to differentiate against. The scaling of the head that produces it is set in the model config and its loss weight in `LossConfig`, both keyed by this name. No module, no decorator. A derivative is named by the rule `d__d_`, with `forces`, `stress` and `magforces` as the three special cases, so asking for a derivative against a newly declared input needs no code either. - **Core touched:** zero files. The model exposes the row automatically because `BaseMACE` iterates over the declared observables; `MACEOutput` carries the six core fields and everything else by name in `extras`. -- **Enabling it:** list it in the model config's observables, or point the config at a declarations file that extends `defaults/observables.yaml`. +- **Enabling it:** list it in the model config's observables, in a catalogue that extends `DEFAULT_CATALOGUE`. - **Test:** `packages/mace-core/tests/test_observables.py` validates the grammar and the derivative naming (pure); if it is autograd-derived, `tests/parity` verifies finite-diff. ### 3.3 A new loss / transform (plugin registry) diff --git a/packages/mace-core/pyproject.toml b/packages/mace-core/pyproject.toml index ec0fe9967..2a798e310 100644 --- a/packages/mace-core/pyproject.toml +++ b/packages/mace-core/pyproject.toml @@ -20,7 +20,6 @@ classifiers = [ dependencies = [ "numpy>=1.23", "pydantic>=2.7", - "pyyaml>=6.0", ] [project.urls] diff --git a/packages/mace-core/src/mace_core/__init__.py b/packages/mace-core/src/mace_core/__init__.py index e35dee6f9..7dace87f8 100644 --- a/packages/mace-core/src/mace_core/__init__.py +++ b/packages/mace-core/src/mace_core/__init__.py @@ -8,22 +8,22 @@ from importlib.metadata import PackageNotFoundError, version from mace_core.observables import ( + DEFAULT_CATALOGUE, DerivativeSpec, InputSpec, ObservableCatalogue, ObservableSpec, - load_default_catalogue, ) from mace_core.outputs import MACEOutput __all__ = [ + "DEFAULT_CATALOGUE", "DerivativeSpec", "InputSpec", "MACEOutput", "ObservableCatalogue", "ObservableSpec", "__version__", - "load_default_catalogue", ] #: Version of the installed `mace-core` distribution. Read from installed metadata diff --git a/packages/mace-core/src/mace_core/defaults/observables.yaml b/packages/mace-core/src/mace_core/defaults/observables.yaml deleted file mode 100644 index 3da29c2ea..000000000 --- a/packages/mace-core/src/mace_core/defaults/observables.yaml +++ /dev/null @@ -1,51 +0,0 @@ -# The canonical observable declarations: energy, and its derivatives with -# respect to the two inputs every model has. -# -# This file is the example every other observable copies. A new property is a -# row here (or in a project's own file) and needs no code: the spec drives the -# head, the loss term, the derivative names and the per-atom/per-graph padding. -# -# A derivative with a name of its own declares it here, together with the sign -# it is reported with, rather than being special-cased in code. Leave both out -# and it is called `d__d_` and carries the gradient's own -# sign; give a name and the sign becomes required, because a renamed quantity -# that silently inherited +1 is a model trained on inverted forces that runs. -# -# Units follow the project convention: eV and Å. The strain is dimensionless, -# so its input row carries "1". - -inputs: - # Atomic positions. A polar vector: it changes sign under inversion, which is - # what makes the energy gradient taken against it a 1o quantity as well. - - name: pos - irreps: "1o" - per_atom: true - units: "Å" - - # The symmetric strain the stress is the derivative against, not the nine - # cell entries: a symmetric rank-2 tensor is a scalar plus an l=2 part, six - # components. Unlike `pos` it is not read from the data. The derivative - # engine materialises it as zeros around the model call and applies it to - # the positions and the cell, which is how the frozen tree does it too. - - name: strain - irreps: "0e+2e" - per_atom: false - units: "1" - -observables: - - name: energy - irreps: "0e" - per_atom: false - units: "eV" - derivatives: - - wrt: pos - name: forces - sign: -1 - units: "eV/Å" - # The volume division that turns the strain derivative into a stress is - # not a sign, and does not belong here. It is applied by whatever - # computes the stress. - - wrt: strain - name: stress - sign: +1 - units: "eV/Å^3" diff --git a/packages/mace-core/src/mace_core/observables/__init__.py b/packages/mace-core/src/mace_core/observables/__init__.py index db656dd7c..95a5c6cb1 100644 --- a/packages/mace-core/src/mace_core/observables/__init__.py +++ b/packages/mace-core/src/mace_core/observables/__init__.py @@ -1,10 +1,6 @@ """Declarative observables: what a model computes, declared rather than coded.""" -from mace_core.observables.defaults import ( - DEFAULTS_RESOURCE, - load_catalogue, - load_default_catalogue, -) +from mace_core.observables.defaults import DEFAULT_CATALOGUE from mace_core.observables.derivatives import ( DEFAULT_SIGN, default_derivative_name, @@ -26,7 +22,7 @@ ) __all__ = [ - "DEFAULTS_RESOURCE", + "DEFAULT_CATALOGUE", "DEFAULT_SIGN", "IRREPS_GRAMMAR", "DerivativeRequest", @@ -39,7 +35,5 @@ "default_derivative_name", "irreps_dimension", "is_default_shaped_name", - "load_catalogue", - "load_default_catalogue", "parse_irreps", ] diff --git a/packages/mace-core/src/mace_core/observables/defaults.py b/packages/mace-core/src/mace_core/observables/defaults.py index f3ced638b..05b7ba000 100644 --- a/packages/mace-core/src/mace_core/observables/defaults.py +++ b/packages/mace-core/src/mace_core/observables/defaults.py @@ -1,50 +1,59 @@ -"""Loading a catalogue from a declarations file. - -The shipped file is packaged data rather than a Python literal, because the -whole point of the declarative spec is that a property can be added without -touching code -- including the code that holds the defaults. +"""The default catalogue: energy, and its derivatives against the two inputs. + +It is built from the same objects as any other catalogue, which is what makes +it the example every other declaration copies. A new property is one more +:class:`~mace_core.observables.ObservableSpec` in the catalogue a model is +given, and needs no edit to this package: the spec drives the head, the loss +term, the derivative names and the per-atom or per-graph padding. + +A derivative with a name of its own declares it here, together with the sign +it is reported with, rather than being special-cased in code. Leave both out +and it is called ``d__d_`` and carries the gradient's own +sign; give a name and the sign becomes required, because a renamed quantity +that silently inherited +1 is a model trained on inverted forces that runs. + +Units follow the project convention: eV and Angstrom. The strain is +dimensionless, so its input carries ``"1"``. """ from __future__ import annotations -from importlib.resources import files -from pathlib import Path -from typing import Any - -import yaml - -from mace_core.observables.spec import ObservableCatalogue - -__all__ = [ - "DEFAULTS_RESOURCE", - "load_catalogue", - "load_default_catalogue", -] - -#: Where the shipped declarations live, relative to the package root. -DEFAULTS_RESOURCE = "defaults/observables.yaml" - - -def _catalogue_from_text(text: str, source: str) -> ObservableCatalogue: - document: Any = yaml.safe_load(text) - if document is None: - document = {} - if not isinstance(document, dict): - raise ValueError( - f"{source}: an observable declarations file must be a mapping with " - f"`inputs` and `observables` keys, not a " - f"{type(document).__name__}." - ) - return ObservableCatalogue.model_validate(document) - - -def load_catalogue(path: str | Path) -> ObservableCatalogue: - """Load and validate a declarations file from disk.""" - path = Path(path) - return _catalogue_from_text(path.read_text(encoding="utf-8"), str(path)) - - -def load_default_catalogue() -> ObservableCatalogue: - """The shipped declarations: energy plus its position and strain derivatives.""" - resource = files("mace_core").joinpath(DEFAULTS_RESOURCE) - return _catalogue_from_text(resource.read_text(encoding="utf-8"), DEFAULTS_RESOURCE) +from mace_core.observables.spec import ( + DerivativeRequest, + InputSpec, + ObservableCatalogue, + ObservableSpec, +) + +__all__ = ["DEFAULT_CATALOGUE"] + +#: Atomic positions. A polar vector: it changes sign under inversion, which is +#: what makes the energy gradient taken against it a ``1o`` quantity as well. +_POSITIONS = InputSpec(name="pos", irreps="1o", per_atom=True, units="Å") + +#: The symmetric strain the stress is the derivative against, not the nine cell +#: entries: a symmetric rank-2 tensor is a scalar plus an l=2 part, six +#: components. Unlike ``pos`` it is not read from the data. The derivative +#: engine materialises it as zeros around the model call and applies it to the +#: positions and the cell, which is how the frozen tree does it too. +_STRAIN = InputSpec(name="strain", irreps="0e+2e", per_atom=False, units="1") + +_ENERGY = ObservableSpec( + name="energy", + irreps="0e", + per_atom=False, + units="eV", + derivatives=( + DerivativeRequest(wrt="pos", name="forces", sign=-1, units="eV/Å"), + # The volume division that turns the strain derivative into a stress is + # not a sign, and does not belong here. It is applied by whatever + # computes the stress. + DerivativeRequest(wrt="strain", name="stress", sign=+1, units="eV/Å^3"), + ), +) + +#: Energy plus its position and strain derivatives, named ``forces`` and +#: ``stress``. Every spec in it is frozen, so it is shared rather than rebuilt. +DEFAULT_CATALOGUE = ObservableCatalogue( + inputs=(_POSITIONS, _STRAIN), observables=(_ENERGY,) +) diff --git a/packages/mace-core/tests/test_observables.py b/packages/mace-core/tests/test_observables.py index 72d6d85f9..6fff72be2 100644 --- a/packages/mace-core/tests/test_observables.py +++ b/packages/mace-core/tests/test_observables.py @@ -1,16 +1,14 @@ """The declarative observable specification. The acceptance bar is "zero new code", so most of these tests declare something -in YAML text and assert what comes out. If any of them needed a new branch in -``mace_core`` to pass, the abstraction would not be doing its job. +as plain data, the rows a configuration would hold, and assert what comes out. +If any of them needed a new branch in ``mace_core`` to pass, the abstraction +would not be doing its job. """ -from importlib.resources import files - import pytest -import yaml from mace_core.observables import ( - DEFAULTS_RESOURCE, + DEFAULT_CATALOGUE, DerivativeRequest, InputSpec, IrrepsGrammarError, @@ -18,31 +16,24 @@ ObservableSpec, default_derivative_name, irreps_dimension, - load_catalogue, - load_default_catalogue, parse_irreps, ) from pydantic import ValidationError -# A catalogue with the two inputs every model has, written out so that the -# tests below can add one row at a time to it. -BASE_INPUTS = """ -inputs: - - name: pos - irreps: "1o" - per_atom: true - units: "Å" - - name: strain - irreps: "0e+2e" - per_atom: false - units: "1" -""" +# The two inputs every model has, written out so that the tests below can add +# one row at a time to them. +BASE_INPUTS = [ + {"name": "pos", "irreps": "1o", "per_atom": True, "units": "Å"}, + {"name": "strain", "irreps": "0e+2e", "per_atom": False, "units": "1"}, +] +ENERGY_ROW = {"name": "energy", "irreps": "0e", "per_atom": False, "units": "eV"} -def catalogue_from(text, tmp_path): - path = tmp_path / "observables.yaml" - path.write_text(BASE_INPUTS + text, encoding="utf-8") - return load_catalogue(path) + +def catalogue_from(observables, inputs=()): + return ObservableCatalogue.model_validate( + {"inputs": [*BASE_INPUTS, *inputs], "observables": list(observables)} + ) # --------------------------------------------------------------------------- @@ -209,42 +200,40 @@ def test_a_sign_alone_needs_no_name(): # --------------------------------------------------------------------------- -# The shipped defaults +# The default catalogue # --------------------------------------------------------------------------- def test_the_defaults_declare_energy_and_its_two_derivatives(): - catalogue = load_default_catalogue() - assert catalogue.names() == ("energy", "forces", "stress") + assert DEFAULT_CATALOGUE.names() == ("energy", "forces", "stress") -def test_the_shipped_names_and_signs_come_from_the_file_and_not_from_code(): +def test_the_default_names_and_signs_come_from_the_declaration_and_not_from_code(): """The reason the table was removed: this is now a property of the data. If these were still special-cased in code, the assertion would pass with - the declarations file saying nothing at all. + the declaration saying nothing at all. """ - catalogue = load_default_catalogue() - resolved = {spec.name: spec.sign for spec in catalogue.requested_derivatives()} + resolved = { + spec.name: spec.sign for spec in DEFAULT_CATALOGUE.requested_derivatives() + } assert resolved == {"forces": -1, "stress": +1} - stripped = yaml.safe_load( - files("mace_core").joinpath(DEFAULTS_RESOURCE).read_text(encoding="utf-8") - ) + stripped = DEFAULT_CATALOGUE.model_dump() for observable in stripped["observables"]: - for request in observable.get("derivatives", []): - request.pop("name", None) - request.pop("sign", None) + for request in observable["derivatives"]: + request.pop("name") + request.pop("sign") bare = ObservableCatalogue.model_validate(stripped) assert {spec.name for spec in bare.requested_derivatives()} == { "d_energy_d_pos", "d_energy_d_strain", } - assert [spec.name for spec in catalogue.inputs] == ["pos", "strain"] + assert [spec.name for spec in DEFAULT_CATALOGUE.inputs] == ["pos", "strain"] def test_the_default_forces_row_is_the_negative_position_gradient(): - forces = load_default_catalogue().derivative("energy", "pos") + forces = DEFAULT_CATALOGUE.derivative("energy", "pos") assert forces.name == "forces" assert forces.sign == -1 assert forces.per_atom is True @@ -253,7 +242,7 @@ def test_the_default_forces_row_is_the_negative_position_gradient(): def test_the_default_stress_row_is_the_positive_strain_gradient(): - stress = load_default_catalogue().derivative("energy", "strain") + stress = DEFAULT_CATALOGUE.derivative("energy", "strain") assert stress.name == "stress" assert stress.sign == +1 assert stress.per_atom is False @@ -265,17 +254,17 @@ def test_the_default_stress_row_is_the_positive_strain_gradient(): # --------------------------------------------------------------------------- -def test_a_new_rank_two_per_atom_observable_is_a_row_in_yaml(tmp_path): +def test_a_new_rank_two_per_atom_observable_is_one_declaration(): catalogue = catalogue_from( - """ -observables: - - name: quadrupole - irreps: "0e+2e" - per_atom: true - units: "e*Å^2" - derivatives: [pos, strain] -""", - tmp_path, + [ + { + "name": "quadrupole", + "irreps": "0e+2e", + "per_atom": True, + "units": "e*Å^2", + "derivatives": ["pos", "strain"], + } + ] ) quadrupole = catalogue.observable("quadrupole") assert quadrupole.per_atom is True @@ -287,33 +276,29 @@ def test_a_new_rank_two_per_atom_observable_is_a_row_in_yaml(tmp_path): ) -def test_a_new_input_feature_makes_its_derivative_declarable(tmp_path): +def test_a_new_input_feature_makes_its_derivative_declarable(): """`magmom` is the case that pays for the grammar being written over declared inputs rather than over positions and the strain. - It is also the case that pays for the name and the sign living in the file. - `magforces` used to be a row in a table inside this package, and this - catalogue reaches it with no code at all. + It is also the case that pays for the name and the sign living in the + declaration. `magforces` used to be a row in a table inside this package, + and this catalogue reaches it with no code at all. """ catalogue = catalogue_from( - """ - - name: magmom - irreps: "1e" - per_atom: true - units: "muB" - -observables: - - name: energy - irreps: "0e" - per_atom: false - units: "eV" - derivatives: - - wrt: magmom - name: magforces - sign: -1 - units: "eV/muB" -""", - tmp_path, + [ + { + **ENERGY_ROW, + "derivatives": [ + { + "wrt": "magmom", + "name": "magforces", + "sign": -1, + "units": "eV/muB", + } + ], + } + ], + inputs=[{"name": "magmom", "irreps": "1e", "per_atom": True, "units": "muB"}], ) magforces = catalogue.derivative("energy", "magmom") assert magforces.name == "magforces" @@ -324,18 +309,8 @@ def test_a_new_input_feature_makes_its_derivative_declarable(tmp_path): assert catalogue.names() == ("energy", "magforces") -def test_the_bare_string_form_and_the_mapping_form_agree(tmp_path): - shorthand = catalogue_from( - """ -observables: - - name: energy - irreps: "0e" - per_atom: false - units: "eV" - derivatives: [pos] -""", - tmp_path, - ) +def test_the_bare_string_form_and_the_mapping_form_agree(): + shorthand = catalogue_from([{**ENERGY_ROW, "derivatives": ["pos"]}]) assert shorthand.observable("energy").derivatives == (DerivativeRequest(wrt="pos"),) @@ -344,105 +319,59 @@ def test_the_bare_string_form_and_the_mapping_form_agree(tmp_path): # --------------------------------------------------------------------------- -def test_a_derivative_against_an_undeclared_input_is_an_error(tmp_path): +def test_a_derivative_against_an_undeclared_input_is_an_error(): with pytest.raises(ValidationError) as caught: - catalogue_from( - """ -observables: - - name: energy - irreps: "0e" - per_atom: false - units: "eV" - derivatives: [elec_temp] -""", - tmp_path, - ) + catalogue_from([{**ENERGY_ROW, "derivatives": ["elec_temp"]}]) message = str(caught.value) assert "energy" in message assert "elec_temp" in message assert "['pos', 'strain']" in message -def test_a_derived_name_may_not_collide_with_a_declared_observable(tmp_path): +def test_a_derived_name_may_not_collide_with_a_declared_observable(): with pytest.raises(ValidationError) as caught: catalogue_from( - """ -observables: - - name: d_energy_d_pos - irreps: "1o" - per_atom: true - units: "eV/Å" - - name: energy - irreps: "0e" - per_atom: false - units: "eV" - derivatives: [pos] -""", - tmp_path, + [ + { + "name": "d_energy_d_pos", + "irreps": "1o", + "per_atom": True, + "units": "eV/Å", + }, + {**ENERGY_ROW, "derivatives": ["pos"]}, + ] ) assert "'d_energy_d_pos'" in str(caught.value) -def test_a_declared_name_may_not_collide_with_a_declared_observable(tmp_path): - """The same guard, now reachable through a name the file chose itself.""" +def test_a_declared_name_may_not_collide_with_a_declared_observable(): + """The same guard, now reachable through a name the declaration chose.""" with pytest.raises(ValidationError) as caught: catalogue_from( - """ -observables: - - name: forces - irreps: "1o" - per_atom: true - units: "eV/Å" - - name: energy - irreps: "0e" - per_atom: false - units: "eV" - derivatives: - - wrt: pos - name: forces - sign: -1 -""", - tmp_path, + [ + {"name": "forces", "irreps": "1o", "per_atom": True, "units": "eV/Å"}, + { + **ENERGY_ROW, + "derivatives": [{"wrt": "pos", "name": "forces", "sign": -1}], + }, + ] ) assert "'forces'" in str(caught.value) -def test_a_name_declared_twice_is_an_error(tmp_path): +def test_a_name_declared_twice_is_an_error(): with pytest.raises(ValidationError) as caught: - catalogue_from( - """ -observables: - - name: energy - irreps: "0e" - per_atom: false - units: "eV" - - name: energy - irreps: "0e" - per_atom: false - units: "eV" -""", - tmp_path, - ) + catalogue_from([ENERGY_ROW, ENERGY_ROW]) assert "declared twice" in str(caught.value) -def test_asking_for_the_same_derivative_twice_is_an_error(tmp_path): +def test_asking_for_the_same_derivative_twice_is_an_error(): with pytest.raises(ValidationError): - catalogue_from( - """ -observables: - - name: energy - irreps: "0e" - per_atom: false - units: "eV" - derivatives: [pos, pos] -""", - tmp_path, - ) + catalogue_from([{**ENERGY_ROW, "derivatives": ["pos", "pos"]}]) def test_an_unknown_observable_or_input_says_what_is_declared(): - catalogue = load_default_catalogue() + catalogue = DEFAULT_CATALOGUE with pytest.raises(KeyError) as caught: catalogue.observable("dipole") assert "['energy']" in str(caught.value) diff --git a/packages/mace-core/tests/test_outputs.py b/packages/mace-core/tests/test_outputs.py index 68462b2af..9a7afafdf 100644 --- a/packages/mace-core/tests/test_outputs.py +++ b/packages/mace-core/tests/test_outputs.py @@ -11,7 +11,7 @@ import numpy as np import pytest -from mace_core.observables import load_default_catalogue +from mace_core.observables import DEFAULT_CATALOGUE from mace_core.outputs import ( CORE_FIELD_NAMES, FIELD_BY_OBSERVABLE, @@ -157,7 +157,7 @@ def test_membership_and_listing_agree_on_every_name(): def test_a_catalogue_name_finds_its_value_in_an_output(): """The correlation the two vocabularies exist to allow.""" - catalogue = load_default_catalogue() + catalogue = DEFAULT_CATALOGUE output = MACEOutput(total_energy=np.array([-1.5]), forces=np.zeros((4, 3))) shared = set(catalogue.names()) & set(output.names()) assert shared == {"energy", "forces"} diff --git a/tests/architecture/test_observable_completeness.py b/tests/architecture/test_observable_completeness.py index a34638b2b..9c92568bc 100644 --- a/tests/architecture/test_observable_completeness.py +++ b/tests/architecture/test_observable_completeness.py @@ -43,10 +43,10 @@ allow_module_level=True, ) from mace_core.observables import ( + DEFAULT_CATALOGUE, DEFAULT_SIGN, ObservableSpec, default_derivative_name, - load_default_catalogue, ) from mace_core.outputs import CORE_FIELD_NAMES, FIELD_BY_OBSERVABLE @@ -201,14 +201,13 @@ def test_the_named_derivatives_keep_their_legacy_names(): """forces, stress and magforces are the pairs that have a name of their own, and the two negated ones are the two the frozen tree negates. - The shipped declarations are the other side of this comparison, and they are - a genuinely separate source: this table is read off the frozen tree, and - that file is authored. Where both name a pair they have to agree, and a pair - the file does not ship yet, `magforces`, is simply absent rather than wrong. + The default catalogue is the other side of this comparison, and it is a + genuinely separate source: this table is read off the frozen tree, and that + catalogue is authored. Where both name a pair they have to agree, and a pair + the catalogue does not declare yet, `magforces`, is simply absent rather + than wrong. """ - declared = { - spec.name: spec for spec in load_default_catalogue().requested_derivatives() - } + declared = {spec.name: spec for spec in DEFAULT_CATALOGUE.requested_derivatives()} for name in ("forces", "stress", "magforces"): row = DISPOSITIONS[name] assert isinstance(row, Derivative)