ARCH-5 — Conditioned node inputs, non-positional derivatives, and the fixed-point re-entrancy contract

[aacostadiaz]

Depends on: ARCH-3 (#1562), ARCH-4 (#1563), DATA-3 (#1571) · Blocks: MAG-1 (#1584), ELEC-2 (#1592), ELEC-3 (#1593), TRN-2 (#1576)

Context: ARCH-2 (#1561)/3/4 describe a model whose only differentiable inputs are positions and the strain displacement. Develop already ships two families that break that assumption, and without this ticket v1 would drop both by omission.

• Magnetic. magmom is a [n_nodes, 3] input node feature that is also a grad leaf: MagneticScaleShiftMACE.forward marks it requires_grad_(True) (mace/modules/extensions.py:1757), feeds it through SHModule/ChebyshevBasisGeneral into EquivariantProductBasisWithSelfMagmomBlock (mace/modules/blocks.py:516) and the two spin-orbit-coupled interaction blocks (blocks.py:1525,1678), and reports magforces = −∂E/∂magmom as a first-class output (extensions.py:1920-1933, emitted at :1951). Legacy computes it in the same autograd.grad call as forces and virials — compute_forces_virials_magforces (mace/modules/utils.py:178-226) and compute_forces_magforces (:229-266), selected inside get_outputs by compute_magforces (:270-357), which raises if magmoms is None. The data contract already carries it: REF_magmom / REF_magforces are DefaultKeys members (mace/tools/default_keys.py:18-19), magmom/magforces/magforces_weight are AtomicData fields (mace/data/atomic_data.py:43-44,58), and --compute_magforces (mace/tools/arg_parser.py:444), --magmom_key/--magforces_key (:708,714) and --magforces_weight (:810) are live flags.
• Electrostatics. external_field is a [n_graphs, 3] graph-level input (mace/data/atomic_data.py:63,100) that MACELES reads at mace/modules/extensions.py:313-316 (treating an all-zero field as absent) and passes to the LES solver as e_ext; Born effective charges come back as BEC in the output dict (extensions.py:654, beside latent_charges/latent_dipoles/latent_quads/latent_alphas/latent_kappas at :649-653), gated by compute_bec (:290, :600). The ASE calculator exposes the surrounding knobs — external_field, eps_infty, keep_neutral (mace/calculators/mace.py:123-143, applied at :816 and :824-846) — and tests/extensions/les/test_maceles.py:269,434 already pins the BEC and keep_neutral behaviour.
• Fixed point. MagneticSCFMACE (mace/modules/extensions.py:1968) wraps a magnetic model and relaxes magmom to a fixed point with torch.optim.LBFGS inside forward (:2040-2082), setting magmom.grad = -output["magforces"].detach() by hand each closure, optionally zeroing the transverse components for a collinear SCF, caching the equilibrated moments on the module (self.cache_magmom), and returning scf_energy_history / scf_steps / equilibrated_magmom alongside the ordinary outputs. Its own comment states it does not differentiate through the SCF; it re-evaluates the inner model once at the converged moments. Two consequences develop has already had to encode: the wrapper delegates unknown attributes to the wrapped model (__getattr__, :1991-2011) because otherwise every consumer must reach through magmom_mace, and hessians are refused outright for SCF-wrapped models (pinned by tests/extensions/magnetic/test_magmace.py:922, with the unwrapped case still working at :951).

This ticket generalizes those three shapes into the v1 architecture as one mechanism each — a declared differentiable input, a derivation mode, and a re-entrancy contract — so the magnetic and electrostatic model tickets (MAG-1 (#1584), ELEC-2 (#1592), ELEC-3 (#1593)) build models, not engine extensions. ELEC-1 (#1591) is not one of them: the reference solver and its dispatch layer have neither a fixed point nor a non-positional derivative, and the LES side takes external_field as a plain forward input and reads BEC straight out of the external solver's result (mace/modules/extensions.py:654) rather than from the derivative engine — so it is not blocked by this contract.

Interface & constraints:

1. Declared differentiable inputs. The model config declares zero or more InputSpecs beyond positions: name, shape (per: node|graph), irreps, differentiable: bool. A differentiable input becomes a grad leaf in the derivative engine's phase A — never inside the model forward, which keeps ARCH-4 (ARCH-4 — Two-phase forward and derivative engine (forces via autograd, stress via strain) #1563)'s "no requires_grad_ mutation in forward" intact. magmom ([n_nodes, 3], differentiable) and external_field ([n_graphs, 3], differentiable) are the two develop instances; total_charge, total_spin, elec_temp, fermi_level are declared but not differentiable.

2. The fourth derivation mode: grad_input. ObservableSpec gains derivation: grad_input with a named target, auto-named d_<q>_d_<input> and aliased where develop has a name of its own — magforces for −∂E/∂magmom, BEC for the field derivative the LES solver returns. All grad targets requested for one energy go into a single torch.autograd.grad call, exactly as legacy's fused magforce paths do; separate calls would need retain_graph and would not reproduce legacy's numbers under create_graph=True. The sign is the legacy one: −grad, asserted, not assumed.

3. Missing input is a hard error, not a zero. Legacy already raises when compute_magforces is requested without magmoms (mace/modules/utils.py:296-297); v1 raises at build time when an observable derived from an input the config does not declare, and at call time when the declared input is absent from the graph — naming both. The all-zero-field-means-absent shortcut MACELES uses (extensions.py:315-316) is a recorded DROP: absence is expressed by the key being absent, never by its value, because a genuinely zero external field and no external field must be distinguishable.

4. Fixed-point solvers wrap the engine, never the model. A FixedPointSpec (variable, solver, max_iter, tol, step_size, and the projection develop calls use_collinear) describes an outer loop that repeatedly invokes the engine — model + grad_input — and updates the variable from the returned derivative. Contract clauses, each matching develop:

   • the loop does not differentiate through its own iterations; the reported outputs come from one final evaluation at the converged variable;
   • convergence telemetry (scf_energy_history, scf_steps, equilibrated_magmom) is part of the typed output's extras, not the core fields;
   • the warm-start cache is explicit state on the solver object, with a documented reset, not an attribute quietly written onto a nn.Module during forward;
   • second derivatives through a fixed point are refused with a message naming the wrapper, reproducing develop's hessian refusal rather than returning a plausible wrong number;
   • the wrapper is transparent by construction: v1 composes solver + model rather than shadowing the model's attributes, so develop's __getattr__ delegation has no v1 counterpart and is a recorded DROP (the need for it was an artefact of the wrapper being a nn.Module that stood in for the model).
     TRN-2 (TRN-2 — Composable losses from observable specs, per-stage schedules, transform registry, and the SCF model-transform hook #1576) owns the training-time model-transform hook that installs such a solver; this ticket owns the contract it installs against.

5. Graph keys. No new schema surface is invented here: magmom, magforces, magforces_weight, external_field, fermi_level are DATA-3 (DATA-3 — Graph construction: neighbor-list backends, graph schema, and collation without torch_geometric #1571) schema entries, and the padding rules for them are DATA-3 (DATA-3 — Graph construction: neighbor-list backends, graph schema, and collation without torch_geometric #1571)'s. What this ticket adds is the statement that a differentiable input is padded inert like any other node/graph field and that its derivative on padding slots is discarded by unpad_outputs.

6. The sphericart second-derivative decision (from ARCH-1 (ARCH-1 — Radial bases, cutoffs, node embedding, and native spherical harmonics (reference-only ops) #1559)) is resolved here, because this is where the magmom second-derivative path is specified. Either v1 reproduces develop bit-for-bit with backward_second_derivatives=False and the defect is a recorded, documented DROP-with-successor, or v1 sets it True and the deviation is declared with its measured cost (45 % on fwd+bwd). The decision is written into the PR with the measurement that justifies it; the flag is never left implicit.

Task:

1. mace_core.config: InputSpec, the grad_input derivation mode with its naming/alias rules, and FixedPointSpec — torch-free.
2. Extend the ARCH-4 (ARCH-4 — Two-phase forward and derivative engine (forces via autograd, stress via strain) #1563) engine: differentiable declared inputs become phase-A grad leaves; phase C issues one autograd.grad over the full target list (positions, displacement, declared inputs) and maps None gradients to zeros exactly as legacy does; results land in the typed output under their aliases.
3. Implement the fixed-point driver per §4 in mace_torch/physics/fixed_point.py: solver protocol, LBFGS instance reproducing develop's configuration (max_iter=n_scf_step, tolerance_grad=scf_tol, line_search_fn="strong_wolfe", lr=scf_step_size), the collinear projection, explicit warm-start cache, telemetry into extras, and the second-derivative refusal.
4. BaseMACE/MACEOutputs plumbing for declared node-input streams (no special-casing of magmom by name anywhere).
5. Parity tests vs the frozen legacy mace/ under tests/parity/: magforces on a tiny magnetic fixture at fp64, forces+virials+magforces from one grad call, and the SCF wrapper's converged moments and energy history on a fixed seed.
6. Finite-difference validation of −∂E/∂magmom against central differences in the magmom components, fp64 CPU — the magforce analogue of ARCH-4 (ARCH-4 — Two-phase forward and derivative engine (forces via autograd, stress via strain) #1563)'s force check, which no test in develop performs.

Out of scope: the magnetic blocks, calculator, data augmentation, and the 13 magnetic train flags (MAG-1 (#1584)); the LES solver, latent readouts, and the electrostatics config section (ELEC-1 (#1591)); PolarMACE's own field machinery in mace/modules/field_blocks.py (ELEC-2 (#1592)); the split-charge SCF models (ELEC-3 (#1593)); the training-time transform hook that installs a solver (TRN-2 (#1576)); loss terms over magforces (TRN-2 (#1576)).

Acceptance criteria:

• A declared differentiable input produces its derivative through the ordinary observable machinery, with zero occurrences of magmom, external_field or BEC as literals in engine or backbone code (asserted by grep).
• Forces, virials and magforces for one energy are produced by a single autograd.grad call; the result matches frozen legacy get_outputs(compute_magforces=True) at fp64 on the tiny magnetic fixture.
• −∂E/∂magmom matches 5-point central differences in the magmom components at fp64 (h chosen per the P0-6 methodology).
• Requesting a grad_input observable whose input is undeclared fails at build time; a declared input missing from the graph fails at call time; both messages name the input and the observable. An all-zero external field is present, not absent.
• The fixed-point driver reproduces develop's converged moments, energy history and step count on a fixed seed; it does not backpropagate through its iterations; requesting a hessian through it raises, naming the wrapper.
• The warm-start cache is explicit state with a tested reset; no attribute is written onto a model module during forward (asserted by the ARCH-4 (ARCH-4 — Two-phase forward and derivative engine (forces via autograd, stress via strain) #1563) read-only test extended to modules).
• Padding slots' input derivatives are dropped by unpad_outputs; real-slot derivatives are identical padded vs unpadded.
• The backward_second_derivatives decision from ARCH-1 (ARCH-1 — Radial bases, cutoffs, node embedding, and native spherical harmonics (reference-only ops) #1559) is recorded in the PR with its measurement, and the chosen value is explicit in the code.

Inventory gaps assigned here:

• --compute_magforces (mace/tools/arg_parser.py:444) has no numeric test today — tests/extensions/magnetic/test_magmace.py covers training, eval, equivariance and the SCF/hessian refusals, but never checks dE/dm against a reference. This ticket adds the finite-difference case.
• use_collinear (mace/modules/extensions.py:1977) is reachable only as a constructor argument, has no flag and no test; it becomes a FixedPointSpec field with a case.
• compute_bec / bec_output_index (mace/modules/extensions.py:156-157) arrive through the stringly --les_arguments file (arg_parser.py:531); they become declared observable fields here, with tests/extensions/les/test_maceles.py:269 as the parity oracle.

Verify:

python -m pytest packages/mace-torch/tests/physics/test_grad_inputs.py -v       # single-call multi-target grads, finite differences on magmom
python -m pytest packages/mace-torch/tests/physics/test_fixed_point.py -v       # convergence, no-backprop-through-iterations, hessian refusal
python -m pytest tests/parity -m "not slow" -k "magforce or bec or scf"         # vs frozen legacy mace/

Review focus: that the derivative engine gained a mode and not a family of special cases (no capability names in engine code); the single-autograd.grad construction, since splitting it silently changes create_graph=True numerics; and the fixed-point contract's refusals — a fixed point that silently returns a second derivative is the worst failure available here. Needs a physics reviewer.

---