Add typed model outputs and a declarative observable specification

docs/reforge/extending_mace.md

@@ -86,13 +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
-  default_loss_weight: 1.0
+```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
@@ -112,14 +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
-  normalization: "component" # scale-only; a 1o vector can be scaled but not shifted
-  default_loss_weight: 1.0
+```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
@@ -129,9 +129,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 +164,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
@@ -317,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)
 ```
@@ -367,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.
 

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,24 +63,25 @@ 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
 
 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

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_<q>_d_<x>` 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.

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,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/
-│   │   ├── __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 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_<q>_d_<x> naming, and the three special cases with their signs
+│   │   └── 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** |
@@ -400,7 +398,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 +491,14 @@ 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:** 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_<q>_d_<x>`, 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, 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)
 
-- **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.

packages/mace-core/pyproject.toml

@@ -17,7 +17,10 @@ classifiers = [
     "Programming Language :: Python :: 3.13",
     "Operating System :: OS Independent",
 ]
-dependencies = []
+dependencies = [
+    "numpy>=1.23",
+    "pydantic>=2.7",
+]
 
 [project.urls]
 Homepage = "https://github.com/ACEsuit/mace"

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 (
+    DEFAULT_CATALOGUE,
+    DerivativeSpec,
+    InputSpec,
+    ObservableCatalogue,
+    ObservableSpec,
+)
+from mace_core.outputs import MACEOutput
+
+__all__ = [
+    "DEFAULT_CATALOGUE",
+    "DerivativeSpec",
+    "InputSpec",
+    "MACEOutput",
+    "ObservableCatalogue",
+    "ObservableSpec",
+    "__version__",
+]
 
 #: Version of the installed `mace-core` distribution. Read from installed metadata
 #: rather than hardcoded, so it cannot drift from what pip resolved.