LES extension (dipole, quadrupoles, polarizability, etc.) merge clean

[djkim-kr]

Add MACELES: long-range electrostatics implementation with added features (dipoles, quadrupoles, polarizability, etc.)

Hi,

This is Dongjin Kim from UC Berkeley working with Prof. Bingqing Cheng.
Here's a new updated version of our LES implementation for MACE.
While the extended LES method is well described in our preprint: https://arxiv.org/abs/2605.05746, I briefly describe the main changes in this update below.

Now, this branch is mergeable with your most recent main and is fully backward compatible with our main branch. Please check and make changes if needed.

Thank you very much for your consideration.
Best,
Dongjin

mace/modules/extensions.py — MACELES class

class MACELES (line 123)

MACELES augments the model wiht per-atom electrostatic
degrees of freedom predicted from the equivariant node features and
passing them to the les library to calculate long-range part.

This PR adds extended features including dipoles, quadrupoles, polarizability, etc.

Readout architecture (__init__, lines 124–226)

Each quantity uses a dedicated readout module cloned from the MACE readout
stack via two helper functions:

• _copy_mace_readout(readout, change_irrep_out=...) (line 50): extended
  from upstream's version with an optional change_irrep_out argument.
  Used to clone the existing readout and retarget its output irreps —
  e.g. change_irrep_out="1x1o" for dipoles, "1x2e" for quadrupoles.

• _copy_mace_readout_tp(readout, ...) (line 77, new): builds a
  LinearLesReadoutBlock or NonLinearLesReadoutBlock for quantities
  that require an outer product of l=1 features (see blocks.py below).

The per-quantity readouts are:

Quantity	Approach	Irreps
Atomic charges (les_readouts)	scalar readout, cloned from MACE readout	0e
Atomic dipoles (les_u_readouts)	cloned with change_irrep_out="1x1o"	1o
Quadrupoles via l=2 (les_quad_2e_readouts)	cloned with change_irrep_out="1x2e", then converted via change_of_basis_quads (ReducedTensorProducts('ij=ji', i='1o')) to Cartesian 3×3	2e
Quadrupoles via l=1 (les_quad_1o_readouts)	outer product v⊗v using _copy_mace_readout_tp → LinearLesReadoutBlock	1o ⊗ 1o
Isotropic polarizability (les_alpha_readouts)	scalar readout	0e
Anisotropic polarizability via l=2 (les_alpha_2e_readouts)	cloned with change_irrep_out="1x0e + 1x2e", converted via CartesianTensor("ij=ji").reduced_tensor_products().change_of_basis	0e + 2e
Anisotropic polarizability via l=1 (les_alpha_1o_readouts)	outer product v⊗v using _copy_mace_readout_tp → LinearLesReadoutBlock	1o ⊗ 1o
Screening charges κ (les_kappa_readouts)	scalar readout	0e

The l=1 outer-product path (les_quad_1o_readouts, les_alpha_1o_readouts)
exists because many trained MACE models use only l≤1 hidden irreps and
therefore have no 2e features — the quadrupole/polarizability tensor
must then be formed as v⊗v from the available 1o channel.

keep_last_layer_irreps requirement (line 126)

kwargs["keep_last_layer_irreps"] = True

MACELES forces this flag before calling super().__init__(). Without it,
MACE.__init__ would collapse the last interaction layer to scalar-only
output (see models.py fix below), making vector/tensor features
unavailable at readout time.

---

mace/modules/blocks.py — LES readout blocks (lines 1400–1555)

LinearLesReadoutBlock (line 1400)

Predicts a per-atom 3×3 symmetric polarizability tensor from equivariant
features. For each node:

1. Apply an equivariant Linear layer mixing features within the input
   irreps.
2. Extract scalar (0e) weights αᵢₖ and vector (1o/1e) channels vᵢₖ.
3. Map scalar weights to a mixing matrix via a learned linear
   scalar_to_weight_1o / scalar_to_weight_1e layer.
4. Form the outer product: T = Σₖ wₖ · (vₖ ⊗ vₖ) — this is
   automatically symmetric and equivariant (SO(3)-covariant under
   rotations of the input features).
5. Optionally enforce positive-definite trace via softplus on the weights
   (make_w_pos=True).

The outer product lives entirely in standard PyTorch (no e3nn CG
coefficients needed for the contraction), keeping compilation overhead low.

NonLinearLesReadoutBlock (line 1502)

Nonlinear variant: passes the equivariant features through a gated
nonlinearity (e3nn.nn.Gate) before the same outer-product readout.
Used when the linear model undershoots the polarizability dynamic range.

---

mace/modules/models.py (line 169)

# before
if num_interactions == 1:
    hidden_irreps_out = str(hidden_irreps[0])

# after
if num_interactions == 1 and not keep_last_layer_irreps:
    hidden_irreps_out = str(hidden_irreps[0])

When num_interactions=1, upstream unconditionally reduced the hidden
irreps to scalars only, even when keep_last_layer_irreps=True. MACELES
sets this flag to preserve vector/tensor features; without the fix,
hidden_irreps_out becomes e.g. "128x0e" instead of the full irreps,
and the downstream LES readout blocks raise error due to lack of L=1 or L=2 node features.

---

mace/data/atomic_data.py — atomic_numbers field (line 402)

cls_kwargs = dict(
    ...
+   atomic_numbers=torch.tensor(config.atomic_numbers, dtype=torch.long),
    ...
)

Required for future use of element-specific fixed charge options in the LES
library (e.g. use_fixed_atomic_charges).
AtomicData.from_config was not populating this field, causing a
KeyError when those LES options are active.

---

mace/calculators/mace.py — BEC forces under external electric fields and LES outputs

MACECalculator.__init__ (line ~112)

Five new parameters: compute_bec, external_field, eps_infty,
electric_field_unit, keep_neutral. When compute_bec=True, "bec"
is added to implemented_properties.

MACECalculator.calculate — forward pass (lines 644–683)

compute_bec=True is forwarded via forward_kwargs.

The LES latent outputs (latent_alphas, latent_kappas, BEC) are
collected into Python lists (ret_tensors.setdefault(..., []).append(...)),
so they are accessible as ASE calculator results during MD.

MACECalculator.calculate — post-processing (lines 760–803)

After the results_map block, the list-accumulated LES quantities are
averaged over ensemble models and stored:

self.results["LES_alphas"] = torch.mean(torch.stack(...), dim=0).cpu().numpy()
self.results["LES_kappas"] = torch.mean(torch.stack(...), dim=0).cpu().numpy()
self.results["bec"]        = torch.mean(torch.stack(...), dim=0).cpu().numpy()

When external_field is set, atomic forces are augmented by the BEC
contribution:

F_i  +=  Z*_i · E_eff
E_eff = E_ext · sqrt(ε_r)
ε_r  = ε_∞ / (1 + χ),   χ = Σᵢ α_i / (V ε₀)

keep_neutral=True subtracts the mean BEC before the force application
to enforce acoustic sum rule neutrality.
More details are described in our preprint.

---

mace/cli/eval_configs.py — latent quantity output

run(): added collection lists us_list, kappas_list, alphas_list
alongside the existing qs_list. Per-structure latent quantities are
split from the batch using batch.ptr and written to atoms.arrays as
{prefix}latent_dipoles, {prefix}latent_kappas, {prefix}latent_alphas.
This enables inspection of the internal LES state without modifying the
model forward pass.

---

Backward compatibility

With this branch, training scripts and pre-trained MACELES potentials from ChengUCB/mace:main load and run without modification. (The training scripts and potentials are
shared in https://github.com/ChengUCB/extended_les_fit/tree/main/MLIPs/MACE-LES)

MACELES.forward wraps attribute accesses that were
added in later versions (les_output_scale, les_kappa_scale,
les_alpha_scale, use_anisotropic_polarizability) with hasattr guards
to handle older checkpoints.

[aacostadiaz]

Thanks for the work here! Two things before merge: there are no tests yet, and I think the new tensor outputs give wrong numbers on the cueq backend.

[aacostadiaz]

Could we add tests for the new outputs? A rotation/equivariance check on the new tensor readouts would be ideal

[djkim-kr]

Thank you. We now added the test for the new tensor outputs.

[aacostadiaz]

With cueq the features are in ir_mul layout not mul_ir, so this reshape mixes up the vector components

[djkim-kr]

Thank you. We didn't realize that they used a different layout. Now modified the code to deal with cueq as well without problem.

[aacostadiaz]

I think this breaks with more than one model. These are python lists and the val[i] = ... loop above does list[1] = ... on the second model we will get indexerror. Preallocate them like the other outputs.

[djkim-kr]

Thank you. Now we preallocate them like the ohter outputs

[aacostadiaz]

This looks fragile, for a 2-atom system shape[0] == 2 matches the atom axis and we sum over atoms. From what I can see les always stacks the contributions on dim 1, so isn't it simpler to just sum over dim 1 when ndim == 4?

[djkim-kr]

True. We followed your suggestion. Thank you.

[aacostadiaz]

This assumes all irreps have the same multiplicity and come ordered [0e, 1o, ...]. Instead of hard-coding cdim, loop over irreps_in and record the (start, end) of each 0e and 1o block, like LinearLesReadoutBlock does in its init.

[djkim-kr]

Thank you. Now it is more flexible as we loop over irreps_in and record the (start, end) of each 0e and 1o block instead of hard-coding them.

[gabor1]

@cursor review

[cursor]

Bugbot couldn't run

Bugbot is not enabled for your user on this team.

Ask your team administrator to increase your team's hard limit for Bugbot seats or add you to the allowlist in the Cursor dashboard.

[djkim-kr]

Changes addressing review feedback

Thank you for the prompt review!
I've committed a few changes to the existing code.
I hope these changes address the review comments on the LES tensor readouts, calculator, and tests for new outputs.

mace/calculators/mace.py

• Fixed an IndexError with multi-model committees: the LES outputs (latent_alphas, latent_kappas, BEC) are now preallocated per model instead of collected in Python lists.
• Fixed BEC force handling for small systems: a 4-D BEC is now summed over axis 1 (where les stacks the charge/dipole contributions).

mace/modules/blocks.py

• LinearLesReadoutBlock / NonLinearLesReadoutBlock: handle the cueq ir_mul layout correctly. Vectors are now reshaped according to the active layout, so the 1o components are no longer mixed up under cueq. The `mul_ir (e3nn) path is unchanged.
• NonLinearLesReadoutBlock: stop hard-coding cdim from the irreps string (which assumed uniform multiplicity and [0e, 1o, ...] ordering). It now records the 0e/1o block slices by iterating irreps_in, mirroring LinearLesReadoutBlock.
• Both blocks lazily backfill the new attributes in forward, so checkpoints saved before this change still load and run (backward compatible).

tests/test_maceles.py

• Add test_les_readout_equivariance: checks the predicted 3x3 tensor transforms as R T R^T under random input rotations (SO(3)-equivariance), for both readout blocks.
• Add test_les_readout_layout_invariance: checks the ir_mul path recovers the same vectors as mul_ir, using non-uniform irreps to also cover the slice-based multiplicity handling.
• Update test_run_eval_no_bec to reflect that latent charges are always inferred by MACELES (independent of --compute_bec).

+++
Cleaned up blocks.py and extensions.py to pass the pre-commit checks
(isort + pylint) e.g. removed unused imports, dropped an unused loop
variable, and simplified if/else loop. No functional changes.
I think it should pass the pre-commit test without a problem.