ELEC-4 — `MACELES` on the new architecture: the `[les]` extra, BEC, latent multipoles, and the external-field path

[aacostadiaz]

Depends on: ELEC-1 (#1591), ARCH-3 (#1562), TRN-1 (#1575), DEP-1 (#1583) · Blocks: RET-4 (#1600)

Context: LES is a full second electrostatics family in develop and it needs its own ticket: folding it into ELEC-1 (#1591) as one bullet ("port --les_arguments") understated it by an order of magnitude, and v1 without it is a regression. MACELES (mace/modules/extensions.py:142, decorated @compile_mode("script") at :141) subclasses ScaleShiftMACE, forces keep_last_layer_irreps=True so the LES readouts can see vector features, and wraps the external les package (https://github.com/ChengUCB/les, pinned by commit fadfb0ca in requirements/les.txt, import name les, marker les, probe _module_available("les") at tests/conftest.py:65). It is not the same extra as [electrostatics] and must not be merged into it: the two externals are independent, separately pinned, and separately CI'd (ci-extensions.yaml runs a les job with require-caps: les, filtered over the shared set plus tests/extensions/les/** and requirements/les.txt).

What the port actually has to carry, all of it live on develop:

• A 16-key les_arguments dict, not three flags. MACELES.__init__ reads use_atomwise (the only key in the default {"use_atomwise": False}), compute_bec, bec_output_index, use_dipole, use_quad, use_induced_charge, use_induced_dipole, use_anisotropic_polarizability, alpha_irreps (default "0e+1o+2e"), alpha_1o_nonlinear_readout, alpha_1o_linear_w_pos (default True), make_alpha_positive, make_kappa_positive, output_scale (0.1), kappa_scale (0.01), alpha_scale (0.01). It reaches the trainer as one flag, --les_arguments (mace/tools/arg_parser.py:531).
• Eight parallel readout ModuleLists cloned off the base readouts — les_readouts, les_u_readouts, les_quad_2e_readouts, les_quad_1o_readouts, les_alpha_readouts, les_alpha_1o_readouts, les_alpha_2e_readouts, les_kappa_readouts — built by _copy_mace_readout / _copy_mace_readout_tp (extensions.py:70,98) over LinearLesReadoutBlock / NonLinearLesReadoutBlock (exported from mace/modules/__init__.py).
• Five latent outputs the eval CLI writes to atoms.arrays: latent_charges, latent_dipoles, latent_kappas, latent_alphas, latent_quads (mace/cli/eval_configs.py:252-286 reading them, :411-421 writing them), behind --compute_bec (:59) which refuses on a non-MACELES model at :161.
• Born effective charges and the external-field force. The ASE calculator takes compute_bec, external_field, eps_infty, keep_neutral (mace/calculators/mace.py:122-126), surfaces LES_alphas (:798), LES_kappas (:802) and bec (:806), and — when an external field is set — adds a field force into results["forces"] (:808-857): BEC of shape [N,2,3,3] is summed over axis 1, keep_neutral subtracts the per-atom mean, eps_infty plus the LES polarizabilities give chi = tr(alpha)/3/V/eps_0 (eps_0 = 5.52635e-3) and eps_r = eps_infty/(1+chi), and F_nj = sum_i Z*_nij · E_i · sqrt(eps_r) · electric_field_unit.
• Trainer plumbing: model_script_utils.py:431-435 builds MACELES, and :131-132 rewrites --model MACELES to the internal FoundationMACELES when a foundation model is supplied (built at :367-371); scripts_utils.py:233,346 lists MACELES in the config-extraction allowlist and its config branch, and :1000,1047-1050 documents that the external LES Atomwise MLP exists only when use_atomwise=True and therefore slips past the optimizer's parameter-group construction.

Two develop behaviours are fixes, not defects, and are the target. keep_neutral subtracts the mean not in place (mace/calculators/mace.py:816-822): the older in-place -= aliased self.results["bec"] on the 3-D path and silently neutralised the stored BEC, while the 4-D path escaped because np.sum copies — the same flag behaved differently per BEC layout. And --compute_bec refuses up front on a non-LES model rather than failing inside the forward. Reproduce both.

Interface & constraints:

1. MACELES becomes a registered model in models/electrostatics.py: a BaseMACE subclass returning a typed MACEOutputs, with the LES coupling as a model-transform hook (TRN-1 (TRN-1 — Staged training pipeline with typed stage contracts and an explicit training loop #1575)), never a ScaleShiftMACE.forward override. No @compile_mode, no e3nn, no TorchScript.
2. keep_last_layer_irreps is a model requirement, not a user flag. Legacy silently overwrites the kwarg; v1 declares it as a constraint of the model and errors loudly if the config contradicts it, rather than overwriting the user's value in __init__.
3. les_arguments becomes a typed config section, one field per key of §Context, with the same defaults. It is validated at config time (CFG-1 (CFG-1 — Full training config schema with cross-section validation #1574)), not read out of an untyped dict with .get() at construction.
4. The five latent quantities and BEC are declared observables (CORE-1 (CORE-1 — Typed outputs + declarative observable specification #1555)), produced by this model and consumed by mace eval and the calculator; they are not forward kwargs threaded through the model. --compute_bec's "only for LES models" refusal becomes an observable-availability check that names the model.
5. The external-field BEC force is part of the calculator contract, not a post-hoc patch of results["forces"]. In v1 the field contribution is a declared term of the force observable so that its presence is visible in the outputs schema; the numeric formula (including the eps_infty/chi correction and electric_field_unit) reproduces develop exactly, and keep_neutral is non-mutating.
6. Separate [les] extra. Distinct from [electrostatics], pinning the les package by commit (migrated from requirements/les.txt), with its own paths-filtered CI job carrying require-caps: les. The les capability marker keeps its name so MACE_REQUIRE_CAPS and the skip-o-fail contract are unchanged.
7. The foundation-model path is part of the port. --model MACELES with a foundation checkpoint routes to a distinct construction path in legacy (FoundationMACELES); v1 expresses that as ordinary fine-tuning config over the registered LES model, and the equivalence is tested, not assumed.
8. The optimizer parameter-group hazard is designed out. Legacy's own comment records that LES's external Atomwise MLP does not yet exist when parameter groups are built, so its parameters slip past both the group construction and the completeness check. v1 builds parameter groups after the model is fully constructed and asserts every parameter lands in exactly one group.

Task:

1. Port MACELES to models/electrostatics.py as a registered BaseMACE subclass + typed MACEOutputs + model-transform hook; drop @compile_mode and e3nn.
2. Create the [les] extra (pin by commit, migrated from requirements/les.txt) and its paths-filtered CI job; keep the les marker name.
3. Turn les_arguments into a typed, validated config section covering all 16 keys with develop's defaults; make keep_last_layer_irreps a declared model constraint.
4. Declare BEC and the five latent quantities (latent_charges, latent_dipoles, latent_kappas, latent_alphas, latent_quads) as observables; wire mace eval to write them to atoms.arrays under the same names, and reproduce the "not a LES model" refusal.
5. Port the external-field force path (BEC × field, eps_infty/chi, electric_field_unit, non-mutating keep_neutral) as a declared force term in the DEP-1 (DEP-1 — ASE calculator for the v1 engine (committee, Hessians, descriptors, key pass-through) #1583) calculator.
6. Port the foundation-model construction path and prove it equals the fine-tuning-config route.
7. Fix the parameter-group ordering so no LES parameter is unassigned; assert it.
8. Port the 9 cases in tests/extensions/les/test_maceles.py, including its 3 xfails — each xfail is re-examined and either fixed or carried forward with its reason recorded, never silently dropped.

Out of scope: the solver dispatch layer and reference solver (ELEC-1 (#1591)); PolarMACE (ELEC-2 (#1592)); the split-charge/SCF families (ELEC-3 (#1593)); the generic ASE calculator (DEP-1 (#1583)); deleting legacy mace/modules/extensions.py (RET-4 (#1600)).

Acceptance criteria:

• MACELES is a registered v1 model instantiable from config, with the LES coupling as a model-transform hook; no @compile_mode, no e3nn.
• All 16 les_arguments keys are typed config fields with develop's defaults; an unknown key is a config error, not a silently ignored .get().
• BEC and the five latent quantities are declared observables; mace eval writes latent_{charges,dipoles,kappas,alphas,quads} to atoms.arrays under the legacy names, and requesting BEC from a non-LES model fails with an error naming the model.
• The external-field force reproduces develop numerically on a fixture with external_field, eps_infty and keep_neutral set, for both BEC layouts ([N,2,3,3] and [N,3,3]); keep_neutral does not mutate the reported BEC on either layout.
• keep_last_layer_irreps is enforced as a model constraint with a loud error on a contradicting config, not silently overwritten.
• The [les] extra exists, pins les by commit, and is separate from [electrostatics]; the les marker still selects the tests and MACE_REQUIRE_CAPS=les still fails a broken install.
• A LES model built from a foundation checkpoint through the v1 fine-tuning config matches the legacy FoundationMACELES construction path.
• Every LES parameter (including the Atomwise MLP when use_atomwise=True) belongs to exactly one optimizer parameter group — asserted.
• The 9 test_maceles.py cases are ported; each of the 3 xfails is resolved or carried with a written reason.

Verify:

pip install -e "packages/mace-torch[les]"
python -m pytest packages/mace-torch/tests/extensions/les/test_maceles.py -v   # ported model + xfail dispositions
python -m pytest packages/mace-torch/tests/extensions/les/test_les_bec.py -v   # BEC, both layouts, keep_neutral non-mutating, external field
python -m pytest packages/mace-torch/tests -m les                              # marker + skip-without-extra discipline

Review focus: that all 16 les_arguments keys and all five latent outputs survive with their legacy names (this is the surface v1 is most likely to quietly shrink), that the external-field force reproduces develop for both BEC layouts with keep_neutral non-mutating, and that the LES coupling is a hook rather than a forward override.

---