From e1cf9f6929b8b5cb5a260ab243df3fe7e4ca8149 Mon Sep 17 00:00:00 2001 From: alacour Date: Fri, 28 Aug 2026 10:22:31 -0700 Subject: [PATCH 1/7] readme: condense to an overview MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 562 → ~310 lines. The README had accumulated design-log-level detail (FiLM option tables and symmetry derivations, MP envelope measurements, LES init/DDP internals, batching alignment mechanics). Feature sections now state what each option does and how to invoke it, and point to the class/trainer docstrings — where all of that detail already lives — for the reasoning. Kept intact: layout, install, quickstart (now includes the SPICE/torchrun and calculator examples that were buried in the batching section), the LES install + IP/licensing note, the calculator refusal semantics, the force-training safety split between the two fused kernels, and the test index. Co-Authored-By: Claude Fable 5 --- README.md | 500 ++++++++++++++---------------------------------------- 1 file changed, 126 insertions(+), 374 deletions(-) diff --git a/README.md b/README.md index 47c38ba..a0ba2a8 100644 --- a/README.md +++ b/README.md @@ -75,14 +75,13 @@ are deliberately conservative. ## Quickstart -Run everything from the repo root so `import ecenet` resolves. All three trainers -are **import-and-call** — every option is a keyword argument of the training -function. The multi-GPU trainers (`train_ecenet_spice`, `train_ecenet_mptrj`) -additionally keep a `__main__` entry point so they launch directly under -`torchrun`; set hyperparameters in the call at the bottom of the script (or -import the function from your own driver). +Run everything from the repo root so `import ecenet` resolves. All trainers are +**import-and-call** — every option is a keyword argument of the training +function (see each function's docstring for the full list). The multi-GPU +trainers (`train_ecenet_spice`, `train_ecenet_mptrj`) additionally keep a +`__main__` entry point so they launch directly under `torchrun`. -Train on an rMD17 / MD22 molecule (import-and-call): +Train on an rMD17 / MD22 molecule: ```python from scripts.train_ecenet import train_ecenet @@ -91,110 +90,80 @@ model, results = train_ecenet(molecule='ethanol', n_train=950, n_mp=2) # n_mp ≥ 2 turns on message passing ``` -Optional low-rank equivariant layers (off by default, available on every trainer -as well as `ecenet.ECENet(...)`) — down → nonlinearity at `r` → up, with a -zero-init up-projection, so each layer is the identity at initialisation: +Train on SPICE (10 elements), single process or DDP: ```python -train_ecenet(..., bottleneck_dim=16) +from scripts.train_ecenet_spice import train_ecenet_spice +model, results = train_ecenet_spice(l_max=3, n_max=4, embed_dim=32, n_layers=2) ``` -### FiLM gate +```bash +torchrun --nproc_per_node=4 scripts/train_ecenet_spice.py # 4-GPU DDP +``` -`element_film=True` modulates the edge features once — right after they are -built and rotated into the bond frame, before the equivariant layer stack. A -small MLP on `[embed(type_i), embed(type_j), φ(r_ij)]` predicts a scale `γ` -(and optionally a shift `β`), applied as `γ⊙A + β`: +Use a trained model from Python / ASE: ```python -train_ecenet(..., element_film=True, film_n_rbf=8) # element + distance -train_ecenet(..., element_film=True) # element-only (film_n_rbf=0) +from ase.io import read +from ecenet.calculator import ECENetCalculator + +atoms = read('molecule.xyz') +atoms.calc = ECENetCalculator.from_checkpoint('model.mdl') +print(atoms.get_potential_energy()) # eV +print(atoms.get_forces()) # eV/Å +print(atoms.get_stress()) # eV/ų (periodic systems) +``` + +```python +import ecenet +model = ecenet.ECENet(n_types=10, l_max=3, n_max=4, embed_dim=16) +energy = model(positions, types) # positions (N,3), types (N,) ``` -The gate MLP's last layer is zero-init, so `γ=1`, `β=0`, and the model is -unchanged at initialisation — the gate learns away from the identity. +## Model options + +All options below are keyword arguments of `ecenet.ECENet(...)` and of every +trainer; the class and trainer docstrings document the details. -| option | effect | +**Low-rank layers** — `bottleneck_dim=16` replaces each equivariant layer with +down → nonlinearity → up (zero-init up, so each layer is the identity at +initialisation). + +**FiLM gate** — `element_film=True` modulates the freshly built edge features +with a scale (and optionally a shift) predicted by a small MLP on the two +element types and, with `film_n_rbf > 0`, the bond length. Identity at +initialisation; equivariance-safe by construction. Sub-options: +`film_embed_dim`, `film_hidden`, `film_per_m`, `film_shift` (see +`ecenet/film.py`). + +**Message passing** — with `n_mp >= 2`, each MP layer computes a per-edge +message and an invariant score, aggregates messages at the receiving atom, and +applies a receiver transform. `mp_type` selects the weighting: + +| `mp_type` | behaviour | | --- | --- | -| `film_n_rbf` | size of the radial leg `φ(r_ij)`; `0` (default) makes the gate element-only, so `γ` does not vary with bond length | -| `film_embed_dim` | width of each element embedding (default 16) | -| `film_hidden` | gate-MLP hidden width(s); `None` → `[max(2*C, 32)]` | -| `film_per_m` | emit a scale per `(channel, m)` instead of one per channel broadcast over `m` | -| `film_shift` | also predict a shift `β`, as an extra head on the *same* MLP | - -Both extras keep the symmetry intact, which constrains where they may act. -`A_cos` and `A_sin` of a given `(channel, m)` share `γ`, so a per-`m` scale still -commutes with the bond frame's per-mode SO(2) rotation; the structural-zero slots -(`m > l` of that channel) are masked to `γ=1`. An additive shift does *not* -commute with that rotation, so `β` lands on the `m=0` slot of `A_cos` only — the -rotation-invariant mode, and the one place a shift is exactly equivariant. - -### Message passing - -With `n_mp >= 2`, every message-passing layer computes, per edge, a low-rank -**message** and an invariant scalar **score**, aggregates the messages at each -receiver atom in the common global frame, and passes the result through a -**receiver** transform back in the bond frame. `mp_type` selects how the -per-edge weight is formed: +| `'softmax'` (default) | attention over the receiver's incoming edges — a weighted average, intensive in coordination | +| `'sum'` | signed score × cutoff envelope — extensive, and a neighbour can contribute negatively | -```python -train_ecenet(..., n_mp=2, mp_type='softmax', mp_n_heads=4) # default -train_ecenet(..., n_mp=2, mp_type='sum') -``` +Both are smooth as edges cross `r_cut_edge`. `mp_dim` sets the message trunk's +bottleneck width, `mp_n_heads` the number of attention heads, and +`mp_l_attention` gives each head one score per degree `l`. Zero-init at the +trunk output, so message passing is a no-op at initialisation. (The older +`mp_type='edge'` has been removed; its checkpoints are rejected with an +explicit error.) + +**Nonlinearity** — `activation` selects the pointwise nonlinearity +(`'silu'` default); `activation='identity'` linearizes the full equivariant +stack (ablation), and `use_nonlinearity=False` skips it in the main layer +stack only. -| `mp_type` | weight | behaviour | -| --- | --- | --- | -| `'softmax'` (default) | `exp(s)·f_cut / Σ_{e→j}(exp(s)·f_cut)` | softmax over the receiver's incoming edges — a weighted *average*, intensive in coordination | -| `'sum'` | `s·f_cut` | raw signed score × cutoff envelope — *extensive* in coordination, and signed, so a neighbour can contribute negatively | - -Either way the smooth cutoff envelope keeps the energy continuous as an edge -crosses `r_cut_edge`. - -`mp_msg_envelope` (on by default) makes the aggregated message decay with -*absolute* distance. It matters only for `'softmax'`: the softmax normalizer -divides the absolute `f_cut` back out, leaving only the relative cutoff across a -receiver's in-edges — so without it a lone neighbour near `r_cut` still gets -weight ≈ 1 and the message is essentially flat in distance. Multiplying `f_cut` -back in fixes that (measured on a dimer, the weight then equals `f_cut` exactly, -falling ~32× from 1.5 Å to 4.5 Å instead of staying at 1.000). `'sum'` is already -enveloped by construction, so the flag is a no-op there — and setting it `False` -warns rather than silently doing nothing. - -Message and scores share **one fused trunk**: a low-rank block (down → -nonlinearity at `mp_dim` → up) whose up-projection emits `2*embed_dim*(l_max+1)` -message channels plus one score channel per head, the score being that channel's -`m=0` (rotation-invariant) component. Sharing a trunk is cheaper than a separate -message block and score head. The up-projection is zero-init, so at -initialisation the message residual and every score are 0 — which makes `'sum'` -an exact no-op and leaves `'softmax'` with uniform attention (`exp(0) = 1`). - -`mp_dim` sets that trunk's bottleneck width (and the receiver's); `mp_n_heads` -splits the value channels (`2*embed_dim`) into that many attention heads, so it -must divide them evenly. - -`mp_l_attention` gives each head one score **per degree `l`** instead of one -overall, so a neighbour can be weighted differently for `l=1` than for `l=2`. -Each `(head, l)` then runs its own independent softmax over the receiver's -in-edges, and the fused trunk widens to `n_ch + n_heads*(l_max+1)`. This stays -equivariant because the Wigner-D block is `l`-diagonal: an invariant scalar -applied uniformly across one `l`'s whole `m`-block commutes with the rotation. -Splitting *within* an `l`, across `m`, would not — which is why the weight is -expanded through a fixed `l_of_s` map rather than being free per spherical -index. - -> **Note.** The older distance/type-weighted `mp_type='edge'` message passing has -> been removed. Checkpoints trained with it (identifiable by `W_msg` weights) are -> rejected with an explicit error rather than silently loading — retrain with -> `'softmax'` or `'sum'`. - -### Long-range electrostatics (LES, optional) - -ECENet's message passing sees only atoms within `r_cut_edge`. The optional -**LES** add-on (Latent Ewald Summation) closes that gap: a head predicts a -scalar latent charge per atom from the model's invariant `l0` embedding, and -the long-range energy is the smeared-Coulomb interaction between those charges -(reciprocal-space Ewald for periodic systems), summed into the total energy on -one shared autograd graph so forces and stress need no extra code. +## Long-range electrostatics (LES, optional) + +The optional **LES** add-on (Latent Ewald Summation) extends ECENet beyond +`r_cut_edge`: a head predicts a latent charge per atom, and the smeared-Coulomb +interaction between those charges (reciprocal-space Ewald for periodic systems) +joins the total energy on one autograd graph, so forces and stress need no +extra code. The implementation is **not vendored** — `ecenet.les.LESLongRange` wraps the inventors' reference package, installed separately (pinned; it is not on PyPI): @@ -204,27 +173,8 @@ pip install -e ".[les]" # or directly: pip install "les @ git+https://github.com/ChengUCB/les@c8063fad18e3d59cb4d783e0ed5a1efea8d55b8d" ``` -`LESLongRange()(l0, positions, cell=None, batch=None)` returns the long-range -energy; the upstream package's own head maps `l0` to latent charges, so the -wrapper's state dict is exactly the upstream module's. The model exposes `l0` -on every forward variant — `l0_only=True` skips the `l=1` work (`l0` is -rotation-invariant, so it needs no frame change; a charge is a scalar, so the -head never sees `l1`): - -```python -lr = ecenet.LESLongRange() -E_sr, l0 = model(pos, types, return_embeddings=True, l0_only=True) -E = E_sr + lr(l0, pos).sum() # one autograd graph → forces via autograd -``` - -The batched paths (`forward_batch`, `forward_batch_multi`) return `l0` as a -per-structure list. - -**Joint training** (`use_les=True`) is available on all four trainers — -`E = E_sr + E_lr` minimised on one autograd graph, forces from the same -graph. `scripts/train_ecenet_xyz.py` is the single-process, in-memory -reference for any ASE-readable file (extxyz etc.); stress strain-transforms -the cell alongside positions and shifts so the Ewald part is covered too: +**Joint training** (`use_les=True`) is available on all four trainers, DDP +included: ```python from scripts.train_ecenet_xyz import train_ecenet_xyz @@ -234,87 +184,29 @@ model, les_module, results = train_ecenet_xyz( use_les=True, n_epochs=200) ``` -`les_readout` selects how the per-atom `l0` fed to the charge head is -aggregated from the final edge invariants (available on all trainers and -`ecenet.ECENet(...)`): +`les_readout` selects how the latent charge is produced: | `les_readout` | aggregation | | --- | --- | -| `'sum'` (default) | parameter-free scatter-sum over the atom's in-edges — extensive in coordination | -| `'softmax'` | attention: a zero-init linear score on each edge's invariants, segment-softmax over the receiver's in-edges with `f_cut` as a multiplicative log-bias, envelope multiplied back in (exactly the MP layers' softmax + `mp_msg_envelope` recipe) — intensive, and decaying with absolute distance | -| `'edge'` | Allegro-LES-style per-edge charge decomposition: a linear scalar head on each edge's invariants, scatter-summed per atom — the width-1 `l0` **is** the latent charge, so upstream's atomwise head is bypassed (the wrapper is called with `l0_is_charge=True`; standard init, since a zero-init charge head would sit on the quadratic energy's gradient-free saddle) | -| `'edge_basis'` | `'edge'` upgraded to mirror the per-edge **energy** readout end to end: an MLP with the energy head's architecture (same full invariant input set, hidden widths, and activation — but standard init, since near-zero charges would start on the quadratic energy's saddle) emits `n_max_d` channels dotted with the cutoff-enveloped radial basis of the bond length, so each bond's charge contribution has a learnable distance profile and vanishes exactly at `r_cut` | - -The softmax weight is an invariant scalar shared by the `l0` and `l1` -messages, so SO(3) behaviour is untouched (verified in `tests/test_les.py`, -including the closed-form dimer weight `f_cut²/(f_cut+ε)`). - -For the edge modes, `les_charge_scale` (default 1.0) multiplies the emitted -latent charge by a fixed factor (MACE-LES ships 0.1 as `output_scale`): with -a standard-init head the charges then start small but nonzero — clear of the -quadratic energy's q = 0 saddle, while E_lr (quadratic in q) is suppressed -~scale² early, so the short-range fit leads and the charges grow gently -relative to it. It is not a parameter and is recorded in the checkpoint's -hparams; for `'sum'`/`'softmax'` it cannot apply (the charge is produced -inside the upstream head) and setting it warns. - -`les_dipole=True` (edge modes only; `train_ecenet_xyz`, `train_ecenet_spice`, -and the tools) additionally gives every atom a **latent dipole**: the edge head -emits a second block of channels, reduced exactly like the charge, whose -scalar `d_e` contributes the bond dipole `d_e·r̂_e` at the receiver. The -model's `l0` is then packed `(n_atoms, 4) = [q | u]`, and the wrapper feeds -`u` to upstream's charge–dipole and dipole–dipole Ewald terms -(`E_lr = ½qᵀf_qq q + qᵀf_qu u + ½uᵀf_uu u`), so polarization the fixed -point charges cannot express — the physics behind the extended-LES `u` -channel — joins the same autograd graph. Because `u` is an invariant scalar -times a true polar vector, parity is exact by construction: a planar -molecule (e.g. water) cannot acquire an out-of-plane dipole, matching its -mirror symmetry — and that bond-direction span coincides with the -symmetry-allowed subspace whenever it binds (coplanar or collinear -neighbourhoods), so nothing expressible is lost. The dipole block is -zero-init — safe here, unlike the charge head, because the `qᵀf_qu u` -cross-term supplies a gradient at u = 0 — so enabling the flag changes -nothing at initialisation. The molecular dipole becomes -`μ = Σᵢ qᵢrᵢ + Σᵢ uᵢ` (`tools/eval_spice_dipoles.py` handles this -automatically). The SPICE trainer's one-batched-call path covers the dipole -terms too — the vectorized isolated path masks upstream's `f_qu`/`f_uu` -kernels exactly as it does `f_qq` (verified against upstream's per-structure -loop) — but note the dipole–dipole kernel is a `(ΣN)²·3·3` tensor, 9× the -charge kernel's memory, so dipole runs want smaller atom budgets. - -`les_charges=False` (requires `les_dipole=True`; all four trainers and -`ecenet.ECENet(...)`) is the **dipoles-only** ablation: the head emits only -the dipole block and the q column of the packed l0 is exactly zero, so -`E_lr = ½uᵀf_uu u` alone — no monopole–monopole term and no q–u cross-term, -which also means inter-molecular charge transfer is inexpressible by -construction. Because that cross-term is what makes the zero-init dipole -slot trainable, the dipoles-only head gets **standard init** instead (a -zero-init u would sit on the uu-quadratic energy's gradient-free saddle, -exactly the charge head's original problem), with `les_charge_scale` now -acting on u. Everything downstream — wrapper, calculators, dumped `les_q` -(all zeros), BECs (now pure dipole flow) — works unchanged, since the l0 -layout is untouched. - -The SPICE trainer takes the same flags (`use_les=True`, `les_arguments`) and -trains jointly under DDP: the LES head lives inside the DDP-wrapped forward -module (so its gradients join the bucket reduction and run on every step, -keeping `find_unused_parameters=False` valid), and the long-range energy is -computed in **one batched LES call** per step — concatenated atoms plus a -structure-index vector, zero cells → the isolated pairwise path (verified -bit-identical to per-structure calls in `tests/test_spice_trainer.py`). - -Upstream builds its charge head lazily on the first forward, so both trainers -materialise it with one throwaway forward before the DDP wrap / optimiser / -checkpoint restore. LES checkpoints carry a top-level `les` dict; -`ECENetCalculator.from_checkpoint` **refuses** them rather than silently -dropping the long-range term (`ignore_les=True` loads the short-range part -deliberately). For MD and single-point use, **`ECENetLESCalculator`** loads a -joint checkpoint and evaluates `E = E_sr + E_lr` on one autograd graph — -forces from the joint backward, and stress from a strain pass that strains -positions, shift vectors, *and the cell*, so the Ewald term's explicit cell -dependence is included (verified against finite differences). Non-periodic -systems use the isolated pairwise path, periodic ones reciprocal-space -Ewald; `examples/run_md_xyz.py` picks the right calculator automatically: +| `'sum'` (default) | parameter-free scatter-sum of edge invariants; upstream's head maps it to charges | +| `'softmax'` | attention-weighted read-out (intensive, distance-decaying) | +| `'edge'` | Allegro-LES-style: a linear per-edge charge, summed per atom | +| `'edge_basis'` | per-edge charge head mirroring the energy readout (learnable distance profile, vanishes at `r_cut`) | + +Further knobs (docstrings have the reasoning): `les_charge_scale` (fixed +multiplier on the edge-mode latent charge, à la MACE-LES's `output_scale`), +`les_dipole=True` (edge modes: every atom also gets a latent dipole `u`, fed +to upstream's charge–dipole and dipole–dipole Ewald terms — polarization the +fixed charges cannot express; molecular dipole `μ = Σ qᵢrᵢ + Σ uᵢ`), and +`les_charges=False` (dipoles-only ablation). + +**MD and evaluation.** `ECENetLESCalculator` loads a joint checkpoint and +evaluates `E = E_sr + E_lr` on one graph — forces from the joint backward, +stress from a strain pass that covers the Ewald term's cell dependence +(verified against finite differences). It refuses short-range checkpoints, +and `ECENetCalculator` symmetrically refuses LES ones (`ignore_les=True` +overrides). `examples/run_md_xyz.py` picks the right calculator +automatically: ```python from ecenet.calculator import ECENetLESCalculator @@ -322,38 +214,14 @@ atoms.calc = ECENetLESCalculator.from_checkpoint('water_les.mdl') print(atoms.get_potential_energy()) # E_sr + E_lr, eV ``` -It refuses short-range checkpoints (symmetric with `ECENetCalculator` -refusing LES ones). - -Every force call also exposes the per-atom latent charges — they are computed -on the way to `E_lr` anyway — via `atoms.get_charges()` (and, for `les_dipole` -checkpoints, the latent atomic dipoles as `calc.results['les_dipoles']`). -`run_md_xyz --dump_charges` writes them onto every dumped frame as `les_q` / -`les_u` extxyz columns, giving charge/dipole trajectories along MD for free -(extxyz output only — ASE's `.traj` format drops custom per-atom arrays). The -global sign of the latent charges is arbitrary (`E_lr` is quadratic in `q`): -consistent within a checkpoint, not physically pinned. - -`calc.compute_bec(atoms)` returns the Born effective charges `Z* = ∂P/∂r` -as an `(N, 3, 3)` array, charge-flow terms included (upstream's BEC module -does the differentiation — Berry-phase polarization for periodic cells, -direct sum otherwise; verified against finite differences of the -polarization on a periodic box, charge flow included, in -`tests/test_xyz_trainer.py`). Unlike the charges this cannot ride along the -force call — it runs its own forward plus three backward passes (≈ 4 force -calls). `run_md_xyz --dump_bec` writes it per dumped frame as a per-atom -`bec` column of 9 row-major components (the `les_fit` reference layout). - -The other two trainers take `use_les=True` as well. `train_ecenet_mptrj` -runs the **periodic** path: one batched LES call per step over the -concatenated atoms with the stacked cells (reciprocal-space Ewald per -structure), the LES head inside the DDP-wrapped forward module (as in the -SPICE trainer), and the stress strain pass covering the cell. Every frame -needs its cell tensor, which `prepare_mptrj.py` now stores — prepared dirs -written before that must be re-prepared for LES runs (a clear error says so). -`train_ecenet` (rMD17/MD22) uses the isolated pairwise path; note those -datasets are in kcal/mol while the LES Coulomb constant is eV·Å-based, so -the latent charges absorb the unit scale. +Every force call also exposes the latent charges via `atoms.get_charges()` +(and dipoles as `calc.results['les_dipoles']`); `calc.compute_bec(atoms)` +returns Born effective charges `Z* = ∂P/∂r` with charge-flow terms included. +`run_md_xyz --dump_charges` / `--dump_bec` write them onto every dumped frame +as extxyz columns (`les_q`, `les_u`, `bec`), giving charge/dipole/BEC +trajectories along MD. The global sign of the latent charges is arbitrary +(`E_lr` is quadratic in `q`): consistent within a checkpoint, not physically +pinned. > **IP / licensing.** The `les` package is CC BY-NC 4.0 (**non-commercial**); > it is an optional dependency and none of its code is included in this @@ -362,163 +230,47 @@ the latent charges absorb the unit scale. > (academic use unrestricted). This repository's own license covers only the > code in this repository and grants no rights to either. -### Fused kernels (optional) - -Two opt-in fused paths trade nothing numerically for memory (and, with Triton -on CUDA, HBM traffic). Both are runtime toggles on the model, off by default: - -```python -model.set_edge_frame_fused(True) # gather→Wigner-rotate→reshape as one op, - # + the MP layers' pack/unrotate (e2n=True) -model.set_activation_fused(True) # nonlinearity grid recomputed in backward -``` - -`set_edge_frame_fused` re-gathers in the backward instead of saving the -`(n_edges, 2C, n_sph)` intermediates; its backward is built from differentiable -ops, so it is safe for double-backward force-loss training. -`set_activation_fused` drops the `(n_edges, F, n_grid)` grid transient from the -saved-for-backward set; it is single-backward oriented — leave it off for -force-loss training. On CUDA + float32 both dispatch to Triton kernels -(`ecenet/edge_frame_kernel.py`, `ecenet/realspace_kernel.py`); elsewhere they -run equivalent pure-PyTorch fallbacks. Verified bit-identical (CPU) / -fp32-accurate (kernels) in `tests/test_edge_frame_kernel.py` and -`tests/test_realspace_kernel.py`. - -### Learning-rate schedule +## Trainer options -All three trainers take `lr_schedule`, defaulting to `'plateau'` -(`ReduceLROnPlateau` on the validation metric, as before): +**Learning-rate schedules** — `lr_schedule='plateau'` (default) | +`'multistep'` (`lr_milestones`, `lr_gamma`) | `'cosine'` (`lr_min_factor`), +plus `warmup_epochs` for the latter two. `multistep` and `cosine` are pure +functions of the epoch index — resume-exact, nothing in the checkpoint, and +every DDP rank computes the same LR independently. ```python -train_ecenet(..., lr_schedule='multistep', - lr_milestones=[80, 130, 170, 190], lr_gamma=0.5) - train_ecenet(..., lr_schedule='cosine', warmup_epochs=5, lr_min_factor=0.01) ``` -| option | applies to | effect | -| --- | --- | --- | -| `lr_milestones` | `multistep` | epochs at which the LR is multiplied by `lr_gamma` | -| `lr_gamma` | `multistep` | decay factor at each milestone (default `0.1`) | -| `warmup_epochs` | `multistep`, `cosine` | linear ramp `0 → lr` over the first N epochs | -| `lr_min_factor` | `cosine` | LR floor as a fraction of `lr`, reached exactly on the last epoch | -| `scheduler_patience` | `plateau` | unchanged | - -`multistep` and `cosine` are computed as **pure functions of the epoch index** -rather than through torch's stateful schedulers. That has three consequences -worth knowing: a resumed run lands on exactly the LR a fresh run would have at -that epoch (torch's `MultiStepLR` counts `.step()` calls, so it would replay from -the initial LR); nothing needs to go in the checkpoint; and every DDP rank -computes the same LR independently, with no state to keep in sync. The -`multistep` curve is verified identical to `torch.optim.lr_scheduler.MultiStepLR` -over a full run. - -### Batching and precision (SPICE / MPtrj trainers) - -Two size-aware batching modes for the SPICE trainer, both off by default: +**Size-aware batching** (SPICE trainer, and the MPtrj prepared-shard mode) — +`bucket=True` batches similar-sized structures; `max_atoms_per_batch=250` +packs to a total-atom budget so per-step memory/compute is roughly uniform; +`max_batch_count` caps structures per batch; `bucket_sort=False` trades a +little load balance for batch diversity. All modes keep every DDP rank on the +same batch count (the collective in backward deadlocks otherwise) via a +deterministic round-alignment scheme — see the trainer docstrings. -```python -train_ecenet_spice(..., bucket=True) # size-bucketed, fixed batch_size -train_ecenet_spice(..., max_atoms_per_batch=250) # atom budget; implies bucket -train_ecenet_spice(..., max_atoms_per_batch=250, max_batch_count=16) -``` - -`bucket=True` sorts the epoch's structures by atom count and batches consecutive -ones, so a batch holds similar-sized molecules rather than one giant molecule -alongside several small ones (measured: mean within-batch atom spread 0.2 vs 44.1 -unsorted). `max_atoms_per_batch` goes further and packs to a total-atom budget, so -memory and compute per step are roughly uniform and a batch of several large -molecules can no longer OOM; `max_batch_count` optionally caps structures per -batch, bounding the per-structure Python overhead when a batch is all tiny -molecules. - -`bucket_sort=True` (the default) sorts by atom count before packing. -`bucket_sort=False` greedy-packs the already-shuffled order instead, trading a -little DDP balance for batch diversity: with sorting, the largest, rare-size -structures keep the *same* batch-mates every epoch (measured: they retain 84% of -their batch-mates across epochs, vs 10% unsorted), which costs gradient diversity -exactly where there is least data. The atom budget still bounds per-batch cost -either way, so the per-rank load spread at `world_size=8` only loosens from 0.9% -to 2.4%. Mainly meaningful together with `max_atoms_per_batch` — without a budget, -unsorted packing is just fixed-size batches plus the round alignment. - -Both share the cross-rank alignment, which is what makes them useful under DDP. -Every rank must run the same *number* of batches or the collective in backward -deadlocks, so the assignment is derived identically on every rank rather than -communicated: sort by atom count, form batches, group into rounds of -`world_size`, drop any partial final round, shuffle the round order with a shared -seed, and give rank r the r-th batch of each round. Adjacent (similar-cost) -batches share a round, so per-step work is aligned across ranks and the -molecule-size straggler disappears — the biggest win multi-node. Measured spread -in per-rank total atoms at `world_size=8`: 3.4% bucketed, 0.9% atom-budget. - -The MPtrj trainer's prepared-shard mode takes the atom budget too: - -```python -train_ecenet_mptrj(..., prepared_dir='mptrj_prepared', max_atoms_per_batch=250) -``` - -Shards stream with no random access, so there is no global sort; instead each -shard (~10k frames — an i.i.d. sample of the dataset, thanks to the -prepare-time global shuffle) is packed independently, with `bucket_sort` and -`max_batch_count` meaning the same as above. The DDP invariant is restored by -round alignment at the shard level: shards are grouped into rounds of -`world_size` (rank r owns the r-th shard of each round), every rank *plans* -every shard's packing from per-frame atom counts alone — packing depends only -on counts and the shared seed, never on the tensors — and all ranks truncate -to the round's minimum batch count. A seeded permutation decides which batches -are dropped, so with `bucket_sort` the largest-structure tail is not -systematically the part lost; the truncation loss is well under 1% at 10k -frames/shard. The counts live in an `atom_counts.pt` sidecar written by -`prepare_mptrj.py`; prepared dirs that predate it are back-filled -automatically on first use (one pass over the shards, then cached). - -`precompute_topology=True` (SPICE trainer) builds every training structure's -neighbour lists once at startup and reuses them each step. Training positions -are fixed, so the topology never changes — yet the on-the-fly path recomputes -the O(N²) distance matrix and calls `nonzero` twice per structure per step, -each a GPU→CPU sync. Skipping them is numerics-identical (verified to 0 in -`tests/test_spice_trainer.py`); evaluation still builds topology on the fly. - -`tf32=True` (both trainers) routes float32 matmuls to TF32 tensor cores on -Ampere+. TF32 keeps ~10 mantissa bits, so A/B the validation MAE before trusting -it. It is a float32-only mode: under `dtype=torch.float64` it warns and changes -nothing. - -Train on SPICE dataset (10 elements): - -```python -from scripts.train_ecenet_spice import train_ecenet_spice -model, results = train_ecenet_spice(l_max=3, n_max=4, embed_dim=32, n_layers=2) -``` - -Multi-GPU via `torchrun` -(`LOCAL_RANK`/`RANK`/`WORLD_SIZE` are read from the environment for DDP): - -```bash -python scripts/train_ecenet_spice.py # single process +**Other** — `precompute_topology=True` (SPICE) builds neighbour lists once at +startup (numerics-identical, skips per-step GPU syncs); `tf32=True` routes +float32 matmuls to TF32 tensor cores (A/B the validation MAE before trusting +it; float64 warns and changes nothing). -torchrun --nproc_per_node=4 scripts/train_ecenet_spice.py # 4-GPU DDP -``` +## Fused kernels (optional) -Use a trained model from Python / ASE: +Two opt-in fused paths trade nothing numerically for memory (and, with Triton +on CUDA, HBM traffic). Both are runtime toggles on the model, off by default: ```python -from ase.io import read -from ecenet.calculator import ECENetCalculator - -atoms = read('molecule.xyz') -atoms.calc = ECENetCalculator.from_checkpoint('model.mdl') -print(atoms.get_potential_energy()) # eV -print(atoms.get_forces()) # eV/Å -print(atoms.get_stress()) # eV/ų (periodic systems) +model.set_edge_frame_fused(True) # gather→Wigner-rotate→reshape as one op, + # + the MP layers' pack/unrotate (e2n=True) +model.set_activation_fused(True) # nonlinearity grid recomputed in backward ``` -```python -import ecenet -model = ecenet.ECENet(n_types=10, l_max=3, n_max=4, embed_dim=16) -energy = model(positions, types) # positions (N,3), types (N,) -``` +`set_edge_frame_fused` is safe for double-backward force-loss training; +`set_activation_fused` is single-backward oriented — leave it off when +training with a force loss. On CUDA + float32 both dispatch to Triton kernels; +elsewhere they run equivalent pure-PyTorch fallbacks. Verified bit-identical +(CPU) / fp32-accurate (kernels) in the kernel test files. ## Tests @@ -528,7 +280,7 @@ The test suite is pure PyTorch and runs on CPU. Each file is runnable as a scrip python tests/test_ecenet.py # ECENet integration: SO(3) invariance, forces, MP python tests/test_bottleneck.py # low-rank layers: identity at init, SO(3) python tests/test_element_film.py # FiLM gate: identity at init, SO(3), per-m, shift -python tests/test_spice_trainer.py # SPICE trainer: atom-budget batching, DDP invariant +python tests/test_spice_trainer.py # SPICE trainer: atom-budget batching, DDP invariant python tests/test_attention_mp.py # attention MP: SO(3), cutoff continuity, sum vs softmax python tests/test_les.py # LES: l0/l1 read-out SO(3) + batch/PBC consistency; wrapper lazy import python tests/test_edge_frame_kernel.py # fused edge-frame/e2n: gradchecks, model on/off equality (Triton legs need CUDA) From 4f946063cc0b07cb19ddad12f91105c46f2c2b98 Mon Sep 17 00:00:00 2001 From: alacour Date: Fri, 28 Aug 2026 10:28:39 -0700 Subject: [PATCH 2/7] element_film on by default; readme mention shrinks to one line MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The FiLM gate becomes part of the default architecture: element_film=True in ECENet and all four trainers. Old checkpoints are unaffected — from_checkpoint rebuilds from stored hparams, which have always recorded element_film. The gate is identity at init (zero-init MLP), so a default-built model's initial energies are unchanged; parameter counts and state-dict keys grow by the gate. Test fallout, both assumption fixes rather than breaks: - test_ignored_flags_warn now passes element_film=False explicitly (with the gate on by default, the flags it sets are genuinely used — no warning is correct behavior). - test_les's cell=None vs explicit-zero-cell check asserted bitwise equality between two independent implementations; they only agreed exactly by luck of the old seeded values (the gate's parameters shift the RNG stream). Now asserts < 1e-14 with the reason documented. Co-Authored-By: Claude Fable 5 --- README.md | 8 ++------ ecenet/model.py | 11 ++++++----- scripts/train_ecenet.py | 2 +- scripts/train_ecenet_mptrj.py | 2 +- scripts/train_ecenet_spice.py | 2 +- scripts/train_ecenet_xyz.py | 2 +- tests/test_element_film.py | 5 +++-- tests/test_les.py | 6 ++++-- 8 files changed, 19 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index a0ba2a8..6e3e0f7 100644 --- a/README.md +++ b/README.md @@ -129,12 +129,8 @@ trainer; the class and trainer docstrings document the details. down → nonlinearity → up (zero-init up, so each layer is the identity at initialisation). -**FiLM gate** — `element_film=True` modulates the freshly built edge features -with a scale (and optionally a shift) predicted by a small MLP on the two -element types and, with `film_n_rbf > 0`, the bond length. Identity at -initialisation; equivariance-safe by construction. Sub-options: -`film_embed_dim`, `film_hidden`, `film_per_m`, `film_shift` (see -`ecenet/film.py`). +An element-conditioned FiLM gate on the edge features is on by default +(`element_film=False` disables it; sub-options in `ecenet/film.py`). **Message passing** — with `n_mp >= 2`, each MP layer computes a per-edge message and an invariant score, aggregates messages at the receiving atom, and diff --git a/ecenet/model.py b/ecenet/model.py index b8edfa2..0d8201b 100644 --- a/ecenet/model.py +++ b/ecenet/model.py @@ -115,10 +115,11 @@ class ECENet(nn.Module): (without it, a lone neighbour near r_cut still gets weight ≈ 1). mp_type='sum' is already enveloped by construction, so the flag is a no-op there — and cannot be turned off. - element_film: if True, modulate the edge features once — right after - they are built and rotated into the bond frame, before - the layer stack — by an element(+distance)-conditioned - FiLM gate (see ecenet/film.py). Identity at init. + element_film: modulate the edge features once — right after they are + built and rotated into the bond frame, before the layer + stack — by an element(+distance)-conditioned FiLM gate + (see ecenet/film.py). Identity at init. ON by default; + element_film=False disables the gate. film_embed_dim: width of each element embedding in the FiLM gate (default 16) film_n_rbf: radial-basis size φ(r) for the FiLM gate (0 → element-only, so the gate depends on the element pair but not the bond @@ -158,7 +159,7 @@ def __init__( mp_n_heads: int = 1, mp_msg_envelope: bool = True, mp_l_attention: bool = False, - element_film: bool = False, + element_film: bool = True, film_embed_dim: int = 16, film_n_rbf: int = 0, film_hidden=None, diff --git a/scripts/train_ecenet.py b/scripts/train_ecenet.py index 2ee5d33..861def0 100644 --- a/scripts/train_ecenet.py +++ b/scripts/train_ecenet.py @@ -174,7 +174,7 @@ def train_ecenet( mp_msg_envelope=True, mp_l_attention=False, # FiLM gate - element_film=False, + element_film=True, film_embed_dim=16, film_n_rbf=0, film_hidden=None, diff --git a/scripts/train_ecenet_mptrj.py b/scripts/train_ecenet_mptrj.py index b5b9cac..dd70314 100644 --- a/scripts/train_ecenet_mptrj.py +++ b/scripts/train_ecenet_mptrj.py @@ -700,7 +700,7 @@ def train_ecenet_mptrj( mp_msg_envelope=True, mp_l_attention=False, # FiLM gate - element_film=False, + element_film=True, film_embed_dim=16, film_n_rbf=0, film_hidden=None, diff --git a/scripts/train_ecenet_spice.py b/scripts/train_ecenet_spice.py index 6755aff..00452e1 100644 --- a/scripts/train_ecenet_spice.py +++ b/scripts/train_ecenet_spice.py @@ -334,7 +334,7 @@ def train_ecenet_spice( mp_msg_envelope=True, mp_l_attention=False, # FiLM gate - element_film=False, + element_film=True, film_embed_dim=16, film_n_rbf=0, film_hidden=None, diff --git a/scripts/train_ecenet_xyz.py b/scripts/train_ecenet_xyz.py index 45d1663..9cf75f8 100644 --- a/scripts/train_ecenet_xyz.py +++ b/scripts/train_ecenet_xyz.py @@ -151,7 +151,7 @@ def train_ecenet_xyz( mp_msg_envelope=True, mp_l_attention=False, # FiLM gate - element_film=False, + element_film=True, film_embed_dim=16, film_n_rbf=0, film_hidden=None, diff --git a/tests/test_element_film.py b/tests/test_element_film.py index 5153ac7..cfcebea 100644 --- a/tests/test_element_film.py +++ b/tests/test_element_film.py @@ -221,10 +221,11 @@ def test_gate_runs_on_every_forward_path(): def test_ignored_flags_warn(): - """FiLM knobs with element_film=False configure a gate that is never built.""" + """FiLM knobs with element_film=False configure a gate that is never built + (element_film defaults to True, so the off state must be explicit here).""" with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") - ECENet(**COMMON, film_n_rbf=6, film_per_m=True) + ECENet(**COMMON, element_film=False, film_n_rbf=6, film_per_m=True) assert any('film_n_rbf' in str(x.message) for x in w), "expected an ignored-flag warning" with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") diff --git a/tests/test_les.py b/tests/test_les.py index 65d4d87..89c5138 100644 --- a/tests/test_les.py +++ b/tests/test_les.py @@ -575,12 +575,14 @@ def test_smoke_forward(): e = e_sr + lr(l0, p).sum() assert torch.isfinite(e).all(), f"non-finite total energy: {e}" # cell=None (fast path, no per-structure det check) must equal an explicit - # zero cell (upstream's det<1e-6 branch) — both mean isolated. + # zero cell (upstream's det<1e-6 branch) — both mean isolated. The two are + # independent implementations with different summation orders, so agreement + # is to float64 rounding, not bitwise. with torch.no_grad(): e_none = lr(l0, p) e_zero = lr(l0, p, cell=torch.zeros(1, 3, 3, dtype=p.dtype)) dz = (e_none - e_zero).abs().max() - assert dz == 0.0, f"cell=None != zero cell: {dz:.3e}" + assert dz < 1e-14, f"cell=None != zero cell: {dz:.3e}" f = -torch.autograd.grad(e, p)[0] assert f.shape == pos.shape and torch.isfinite(f).all() print(f" smoke: E={e.item():.6f} eV, |F|max={f.abs().max():.3f}") From 2e590fa1abd93408f3b917baae1627ce4562b8dc Mon Sep 17 00:00:00 2001 From: alacour Date: Fri, 28 Aug 2026 10:29:18 -0700 Subject: [PATCH 3/7] readme: drop the per-file test index The Layout block already points at tests/ (runnable scripts, repo root); the file-by-file list belongs in the files' own docstrings. Co-Authored-By: Claude Fable 5 --- README.md | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/README.md b/README.md index 6e3e0f7..7aa8580 100644 --- a/README.md +++ b/README.md @@ -268,26 +268,6 @@ training with a force loss. On CUDA + float32 both dispatch to Triton kernels; elsewhere they run equivalent pure-PyTorch fallbacks. Verified bit-identical (CPU) / fp32-accurate (kernels) in the kernel test files. -## Tests - -The test suite is pure PyTorch and runs on CPU. Each file is runnable as a script: - -```bash -python tests/test_ecenet.py # ECENet integration: SO(3) invariance, forces, MP -python tests/test_bottleneck.py # low-rank layers: identity at init, SO(3) -python tests/test_element_film.py # FiLM gate: identity at init, SO(3), per-m, shift -python tests/test_spice_trainer.py # SPICE trainer: atom-budget batching, DDP invariant -python tests/test_attention_mp.py # attention MP: SO(3), cutoff continuity, sum vs softmax -python tests/test_les.py # LES: l0/l1 read-out SO(3) + batch/PBC consistency; wrapper lazy import -python tests/test_edge_frame_kernel.py # fused edge-frame/e2n: gradchecks, model on/off equality (Triton legs need CUDA) -python tests/test_realspace_kernel.py # fused nonlinearity: backward equivalence, DFT precision (Triton legs need CUDA) -python tests/test_mptrj_trainer.py # end-to-end MPtrj trainer smoke (synthetic) -python tests/test_mptrj_shard_batching.py # shard atom-budget batching: DDP count alignment, sidecar -python tests/test_xyz_trainer.py # small-dataset trainer: smoke, LES resume, force-FD through E_lr -python tests/test_trainer_les.py # use_les in the rMD17 + MPtrj trainers; stress-FD through the Ewald cell strain -python tests/test_wbm_eval.py # WBM relax+score pipeline: slicing/resume, e_form + hull metrics exact on a perfect model -``` - ## License Copyright ©2026. The Regents of the University of California (Regents). All From d4ff76a4c1a5196e669591d4531771e4bf2d225c Mon Sep 17 00:00:00 2001 From: alacour Date: Fri, 28 Aug 2026 10:35:16 -0700 Subject: [PATCH 4/7] mp_type default 'softmax' -> 'sum'; readme drops the MP/nonlinearity detail MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 'sum' — the raw signed score times the cutoff envelope, extensive in coordination and an exact no-op at init — becomes the default aggregation in ECENet and all four trainers. Old checkpoints are unaffected (mp_type is in hparams). Consequence of the default pairing: a default model's msg_envelope flag is now False, since 'sum' is enveloped structurally and never sets it. Tests updated to the new default: test_default_is_softmax becomes test_default_is_sum, and test_msg_envelope_defaults_and_flag pins mp_type='softmax' explicitly where it exercises the softmax-only flag. README: the mp_type comparison table and the nonlinearity paragraph go; message passing is now one short paragraph deferring to the docstring. Co-Authored-By: Claude Fable 5 --- README.md | 23 ++++------------------- ecenet/model.py | 17 +++++++++-------- scripts/train_ecenet.py | 2 +- scripts/train_ecenet_mptrj.py | 2 +- scripts/train_ecenet_spice.py | 2 +- scripts/train_ecenet_xyz.py | 2 +- tests/test_attention_mp.py | 29 ++++++++++++++++------------- 7 files changed, 33 insertions(+), 44 deletions(-) diff --git a/README.md b/README.md index 7aa8580..bbb2c1b 100644 --- a/README.md +++ b/README.md @@ -133,25 +133,10 @@ An element-conditioned FiLM gate on the edge features is on by default (`element_film=False` disables it; sub-options in `ecenet/film.py`). **Message passing** — with `n_mp >= 2`, each MP layer computes a per-edge -message and an invariant score, aggregates messages at the receiving atom, and -applies a receiver transform. `mp_type` selects the weighting: - -| `mp_type` | behaviour | -| --- | --- | -| `'softmax'` (default) | attention over the receiver's incoming edges — a weighted average, intensive in coordination | -| `'sum'` | signed score × cutoff envelope — extensive, and a neighbour can contribute negatively | - -Both are smooth as edges cross `r_cut_edge`. `mp_dim` sets the message trunk's -bottleneck width, `mp_n_heads` the number of attention heads, and -`mp_l_attention` gives each head one score per degree `l`. Zero-init at the -trunk output, so message passing is a no-op at initialisation. (The older -`mp_type='edge'` has been removed; its checkpoints are rejected with an -explicit error.) - -**Nonlinearity** — `activation` selects the pointwise nonlinearity -(`'silu'` default); `activation='identity'` linearizes the full equivariant -stack (ablation), and `use_nonlinearity=False` skips it in the main layer -stack only. +message weighted by an invariant score, aggregates at the receiving atom, and +applies a receiver transform; smooth as edges cross `r_cut_edge`, and a no-op +at initialisation. `mp_type='sum'` by default (`'softmax'`, `mp_dim`, +`mp_n_heads`, `mp_l_attention` in the docstring). ## Long-range electrostatics (LES, optional) diff --git a/ecenet/model.py b/ecenet/model.py index 0d8201b..5b20cce 100644 --- a/ecenet/model.py +++ b/ecenet/model.py @@ -88,13 +88,14 @@ class ECENet(nn.Module): — a fused message/score trunk and a receiver transform — and differ only in the weight applied to each incoming message: - 'softmax' (default): softmax over the receiver's - incoming edges, so the aggregate is a weighted - *average* (intensive in coordination). Zero-init - scores make the attention uniform at init. - 'sum': the raw signed score times the cutoff envelope, - summed (extensive in coordination). Zero-init scores - make the layer an exact no-op at init. + 'sum' (default): the raw signed score times the cutoff + envelope, summed (extensive in coordination). + Zero-init scores make the layer an exact no-op at + init. + 'softmax': softmax over the receiver's incoming edges, + so the aggregate is a weighted *average* (intensive + in coordination). Zero-init scores make the + attention uniform at init. mp_dim: bottleneck width of the fused message/score trunk and of the receiver block (default: n_features_per_m // 4) mp_n_heads: number of attention heads; the value channels @@ -154,7 +155,7 @@ def __init__( output_hidden_dims: list = None, m_max: int = None, bottleneck_dim: int = None, - mp_type: str = 'softmax', + mp_type: str = 'sum', mp_dim: int = None, mp_n_heads: int = 1, mp_msg_envelope: bool = True, diff --git a/scripts/train_ecenet.py b/scripts/train_ecenet.py index 861def0..60a5a57 100644 --- a/scripts/train_ecenet.py +++ b/scripts/train_ecenet.py @@ -168,7 +168,7 @@ def train_ecenet( bottleneck_dim=None, # Message passing n_mp=1, - mp_type='softmax', + mp_type='sum', mp_dim=None, mp_n_heads=1, mp_msg_envelope=True, diff --git a/scripts/train_ecenet_mptrj.py b/scripts/train_ecenet_mptrj.py index dd70314..d0716e7 100644 --- a/scripts/train_ecenet_mptrj.py +++ b/scripts/train_ecenet_mptrj.py @@ -694,7 +694,7 @@ def train_ecenet_mptrj( bottleneck_dim=None, # Message passing n_mp=1, - mp_type='softmax', + mp_type='sum', mp_dim=None, mp_n_heads=1, mp_msg_envelope=True, diff --git a/scripts/train_ecenet_spice.py b/scripts/train_ecenet_spice.py index 00452e1..3a633e2 100644 --- a/scripts/train_ecenet_spice.py +++ b/scripts/train_ecenet_spice.py @@ -328,7 +328,7 @@ def train_ecenet_spice( bottleneck_dim=None, # Message passing n_mp=1, - mp_type='softmax', + mp_type='sum', mp_dim=None, mp_n_heads=1, mp_msg_envelope=True, diff --git a/scripts/train_ecenet_xyz.py b/scripts/train_ecenet_xyz.py index 9cf75f8..c8148f5 100644 --- a/scripts/train_ecenet_xyz.py +++ b/scripts/train_ecenet_xyz.py @@ -145,7 +145,7 @@ def train_ecenet_xyz( bottleneck_dim=None, # Message passing n_mp=1, - mp_type='softmax', + mp_type='sum', mp_dim=None, mp_n_heads=1, mp_msg_envelope=True, diff --git a/tests/test_attention_mp.py b/tests/test_attention_mp.py index 51f34e6..190a4b3 100644 --- a/tests/test_attention_mp.py +++ b/tests/test_attention_mp.py @@ -5,8 +5,8 @@ global frame, then a per-edge receiver residual. Two aggregations share that structure and differ only in the weight: - 'softmax' (default) — softmax over the receiver's in-edges (intensive) - 'sum' — raw signed score × cutoff envelope (extensive) + 'sum' (default) — raw signed score × cutoff envelope (extensive) + 'softmax' — softmax over the receiver's in-edges (intensive) Message and scores come from ONE fused trunk whose zero-init up-projection emits n_ch message channels plus one score channel per head. @@ -67,12 +67,12 @@ def _activate_scores(layer, std=0.5, bias=None): layer.msg_up.bias[-layer.n_scores:].fill_(bias) -def test_default_is_softmax(): +def test_default_is_sum(): m = ECENet(**COMMON, n_mp=2).double() assert isinstance(m.mp_layers[0], ECENetAttentionMPLayer) - assert m.mp_layers[0].aggregation == 'softmax' - m_s = ECENet(**COMMON, n_mp=2, mp_type='sum').double() - assert m_s.mp_layers[0].aggregation == 'sum' + assert m.mp_layers[0].aggregation == 'sum' + m_s = ECENet(**COMMON, n_mp=2, mp_type='softmax').double() + assert m_s.mp_layers[0].aggregation == 'softmax' # the removed 'edge' MP is rejected, with a message that says so try: ECENet(**COMMON, n_mp=2, mp_type='edge') @@ -80,7 +80,7 @@ def test_default_is_softmax(): assert 'edge' in str(e) and 'removed' in str(e) else: raise AssertionError("expected mp_type='edge' to be rejected") - print(" default mp_type='softmax'; 'sum' selects the weighted-sum aggregation; " + print(" default mp_type='sum'; 'softmax' selects the attention aggregation; " "'edge' is rejected") @@ -316,11 +316,14 @@ def weight(aggregation, envelope, r): def test_msg_envelope_defaults_and_flag(): """On by default for 'softmax'; structurally already on (and not - disableable) for 'sum', which warns rather than silently ignoring.""" - assert ECENet(**COMMON, n_mp=2).mp_layers[0].msg_envelope is True - assert ECENet(**COMMON, n_mp=2, mp_msg_envelope=False).mp_layers[0].msg_envelope is False + disableable) for 'sum' — the default — which warns rather than silently + ignoring.""" + assert ECENet(**COMMON, n_mp=2, + mp_type='softmax').mp_layers[0].msg_envelope is True + assert ECENet(**COMMON, n_mp=2, mp_type='softmax', + mp_msg_envelope=False).mp_layers[0].msg_envelope is False # 'sum' never sets the flag — its weight is s*f_cut, so f_cut twice would be f_cut² - assert ECENet(**COMMON, n_mp=2, mp_type='sum').mp_layers[0].msg_envelope is False + assert ECENet(**COMMON, n_mp=2).mp_layers[0].msg_envelope is False with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") ECENet(**COMMON, n_mp=2, mp_type='sum', mp_msg_envelope=False) @@ -330,7 +333,7 @@ def test_msg_envelope_defaults_and_flag(): warnings.simplefilter("always") ECENet(**COMMON, n_mp=2, mp_type='sum') assert not any('envelope' in str(x.message) for x in w), "plain sum should be quiet" - print(" mp_msg_envelope: default on for transformer, structural for sum, warns if disabled there") + print(" mp_msg_envelope: default on for softmax, structural for sum, warns if disabled there") def test_sum_is_extensive(): @@ -506,7 +509,7 @@ def test_ignored_flags_warn(): if __name__ == "__main__": print("Attention message-passing tests (mp_type='softmax' / 'sum')") - test_default_is_softmax() + test_default_is_sum() test_so3_invariance() test_cutoff_continuity() test_forces_finite() From 9ddb27a07706ee1ab614b133a5fdf69a0ca801bf Mon Sep 17 00:00:00 2001 From: alacour Date: Fri, 28 Aug 2026 10:37:57 -0700 Subject: [PATCH 5/7] readme: plainer phrasing for the LES wrapper intro The 'not vendored' jargon goes; the IP note below still carries the load-bearing statement that none of the les package's code is included. Co-Authored-By: Claude Fable 5 --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index bbb2c1b..1ed6e49 100644 --- a/README.md +++ b/README.md @@ -146,8 +146,8 @@ interaction between those charges (reciprocal-space Ewald for periodic systems) joins the total energy on one autograd graph, so forces and stress need no extra code. -The implementation is **not vendored** — `ecenet.les.LESLongRange` wraps the -inventors' reference package, installed separately (pinned; it is not on PyPI): +`ecenet.les.LESLongRange` wraps the inventors' reference package, installed +separately (pinned; it is not on PyPI): ```bash pip install -e ".[les]" # or directly: From c1d2bec0eabd2dc062c4089446022b4f2d8f3e53 Mon Sep 17 00:00:00 2001 From: alacour Date: Fri, 28 Aug 2026 10:39:06 -0700 Subject: [PATCH 6/7] readme: drop the Trainer options section LR schedules, size-aware batching, precompute_topology, and tf32 are all documented in the trainer docstrings the Quickstart already points at. Co-Authored-By: Claude Fable 5 --- README.md | 25 ------------------------- 1 file changed, 25 deletions(-) diff --git a/README.md b/README.md index 1ed6e49..a22ab14 100644 --- a/README.md +++ b/README.md @@ -211,31 +211,6 @@ pinned. > (academic use unrestricted). This repository's own license covers only the > code in this repository and grants no rights to either. -## Trainer options - -**Learning-rate schedules** — `lr_schedule='plateau'` (default) | -`'multistep'` (`lr_milestones`, `lr_gamma`) | `'cosine'` (`lr_min_factor`), -plus `warmup_epochs` for the latter two. `multistep` and `cosine` are pure -functions of the epoch index — resume-exact, nothing in the checkpoint, and -every DDP rank computes the same LR independently. - -```python -train_ecenet(..., lr_schedule='cosine', warmup_epochs=5, lr_min_factor=0.01) -``` - -**Size-aware batching** (SPICE trainer, and the MPtrj prepared-shard mode) — -`bucket=True` batches similar-sized structures; `max_atoms_per_batch=250` -packs to a total-atom budget so per-step memory/compute is roughly uniform; -`max_batch_count` caps structures per batch; `bucket_sort=False` trades a -little load balance for batch diversity. All modes keep every DDP rank on the -same batch count (the collective in backward deadlocks otherwise) via a -deterministic round-alignment scheme — see the trainer docstrings. - -**Other** — `precompute_topology=True` (SPICE) builds neighbour lists once at -startup (numerics-identical, skips per-step GPU syncs); `tf32=True` routes -float32 matmuls to TF32 tensor cores (A/B the validation MAE before trusting -it; float64 warns and changes nothing). - ## Fused kernels (optional) Two opt-in fused paths trade nothing numerically for memory (and, with Triton From 1ae0bcd45d61a4a6d103c352c9372d9dfcadc1d9 Mon Sep 17 00:00:00 2001 From: alacour Date: Fri, 28 Aug 2026 10:47:45 -0700 Subject: [PATCH 7/7] les_readout: default 'edge_basis' under use_les; readme drops the readout table MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The trainers' les_readout now defaults to None, resolved at run start to 'edge_basis' when use_les=True and 'sum' otherwise. LES runs get the per-edge charge head (the readout behind the dipole/BEC results) by default, while short-range runs keep the parameter-free 'sum' read-out — a bare-model or SR run must not carry an unused charge head, which under DDP's find_unused_parameters=False would hang the gradient reduction (the existing spice/mptrj guards for explicit misconfigurations stay). The bare-model default stays 'sum' for the same reason, documented in the readout comment block. Explicit les_readout settings and old checkpoints are unaffected (the resolved value is what lands in hparams). Tests: the xyz/rmd17 LES smokes pin les_readout='sum' — they specifically exercise upstream's lazy atomwise head and resume against 'sum' checkpoints, and a resume must match the checkpoint's readout like any other architecture hparam (the strict load rejects the mismatch loudly). README: the les_readout table and the knobs paragraph collapse to two sentences deferring to the model docstring. Co-Authored-By: Claude Fable 5 --- README.md | 19 ++++--------------- ecenet/model.py | 4 +++- scripts/train_ecenet.py | 7 ++++++- scripts/train_ecenet_mptrj.py | 7 ++++++- scripts/train_ecenet_spice.py | 7 ++++++- scripts/train_ecenet_xyz.py | 7 ++++++- tests/test_trainer_les.py | 5 +++-- tests/test_xyz_trainer.py | 6 +++--- 8 files changed, 37 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index a22ab14..562aac0 100644 --- a/README.md +++ b/README.md @@ -165,21 +165,10 @@ model, les_module, results = train_ecenet_xyz( use_les=True, n_epochs=200) ``` -`les_readout` selects how the latent charge is produced: - -| `les_readout` | aggregation | -| --- | --- | -| `'sum'` (default) | parameter-free scatter-sum of edge invariants; upstream's head maps it to charges | -| `'softmax'` | attention-weighted read-out (intensive, distance-decaying) | -| `'edge'` | Allegro-LES-style: a linear per-edge charge, summed per atom | -| `'edge_basis'` | per-edge charge head mirroring the energy readout (learnable distance profile, vanishes at `r_cut`) | - -Further knobs (docstrings have the reasoning): `les_charge_scale` (fixed -multiplier on the edge-mode latent charge, à la MACE-LES's `output_scale`), -`les_dipole=True` (edge modes: every atom also gets a latent dipole `u`, fed -to upstream's charge–dipole and dipole–dipole Ewald terms — polarization the -fixed charges cannot express; molecular dipole `μ = Σ qᵢrᵢ + Σ uᵢ`), and -`les_charges=False` (dipoles-only ablation). +The latent charge comes from a per-edge charge head mirroring the energy +readout (`les_readout='edge_basis'`, the default with `use_les=True`); +alternative read-outs and the related options — latent dipoles, charge +scaling, dipoles-only — are documented in the model docstring. **MD and evaluation.** `ECENetLESCalculator` loads a joint checkpoint and evaluates `E = E_sr + E_lr` on one graph — forces from the joint backward, diff --git a/ecenet/model.py b/ecenet/model.py index 5b20cce..bad19bd 100644 --- a/ecenet/model.py +++ b/ecenet/model.py @@ -237,7 +237,9 @@ def __init__( # saddle a zero-init head could never leave. Smoothness at r_cut is # inherited from the edge features' own radial envelope, as in # Allegro-LES's EdgewiseReduce. - # 'edge_basis': 'edge' upgraded to mirror the energy readout end to + # 'edge_basis' (the trainers' default when use_les=True — the + # bare-model default stays 'sum' so short-range models carry no + # unused charge-head parameters): 'edge' upgraded to mirror the energy readout end to # end — an MLP with output_net's architecture (same input: the full # n_features_per_m m=0 invariant set of _contract, not the l'-summed # h_l0; same hidden widths and activation) emits n_max_d channels diff --git a/scripts/train_ecenet.py b/scripts/train_ecenet.py index 60a5a57..9156cbc 100644 --- a/scripts/train_ecenet.py +++ b/scripts/train_ecenet.py @@ -180,7 +180,7 @@ def train_ecenet( film_hidden=None, film_per_m=False, film_shift=False, - les_readout='sum', # (l0,l1) read-out for LES: 'sum' | 'softmax' | 'edge' | 'edge_basis' + les_readout=None, # None -> 'edge_basis' if use_les else 'sum' les_charge_scale=1.0, # fixed multiplier on the edge-mode latent charge (MACELES: 0.1) les_dipole=False, # edge head also emits bond dipoles; l0 packed [q | u] les_charges=True, # False (needs les_dipole): dipoles-only — q hard zero, standard-init dipole head @@ -236,6 +236,11 @@ def loss_fn(pred, tgt): if loss_type == 'huber': return nn.functional.huber_loss(pred, tgt, delta=huber_delta) return ((pred - tgt) ** 2).mean() + if les_readout is None: + # 'edge_basis' is the LES default; short-range runs take 'sum' so the + # model carries no unused charge-head parameters (which would also + # break DDP's find_unused_parameters=False). + les_readout = 'edge_basis' if use_les else 'sum' if optimizer_type not in ('adamw', 'adam', 'sgd'): raise ValueError(f"optimizer_type must be 'adamw', 'adam', or 'sgd'; got {optimizer_type!r}") if device is None: diff --git a/scripts/train_ecenet_mptrj.py b/scripts/train_ecenet_mptrj.py index d0716e7..cf4c78e 100644 --- a/scripts/train_ecenet_mptrj.py +++ b/scripts/train_ecenet_mptrj.py @@ -706,7 +706,7 @@ def train_ecenet_mptrj( film_hidden=None, film_per_m=False, film_shift=False, - les_readout='sum', # (l0,l1) read-out for LES: 'sum' | 'softmax' | 'edge' | 'edge_basis' + les_readout=None, # None -> 'edge_basis' if use_les else 'sum' les_charge_scale=1.0, # fixed multiplier on the edge-mode latent charge (MACELES: 0.1) les_dipole=False, # edge head also emits bond dipoles; l0 packed [q | u] les_charges=True, # False (needs les_dipole): dipoles-only — q hard zero, standard-init dipole head @@ -753,6 +753,11 @@ def train_ecenet_mptrj( world_size=1, local_rank=0, ): + if les_readout is None: + # 'edge_basis' is the LES default; short-range runs take 'sum' so the + # model carries no unused charge-head parameters (which would also + # break DDP's find_unused_parameters=False). + les_readout = 'edge_basis' if use_les else 'sum' is_ddp = world_size > 1 is_main = (rank == 0) verbose = verbose and is_main diff --git a/scripts/train_ecenet_spice.py b/scripts/train_ecenet_spice.py index 3a633e2..f9dc04c 100644 --- a/scripts/train_ecenet_spice.py +++ b/scripts/train_ecenet_spice.py @@ -340,7 +340,7 @@ def train_ecenet_spice( film_hidden=None, film_per_m=False, film_shift=False, - les_readout='sum', # (l0,l1) read-out for LES: 'sum' | 'softmax' | 'edge' | 'edge_basis' + les_readout=None, # None -> 'edge_basis' if use_les else 'sum' les_charge_scale=1.0, # fixed multiplier on the edge-mode latent charge (MACELES: 0.1) les_dipole=False, # edge head also emits bond dipoles; l0 packed [q | u] les_charges=True, # False (needs les_dipole): dipoles-only — q hard zero, standard-init dipole head @@ -439,6 +439,11 @@ def train_ecenet_spice( if verbose: print_flush(f" Loaded {len(test_raw):,} test structures") + if les_readout is None: + # 'edge_basis' is the LES default; short-range runs take 'sum' so the + # model carries no unused charge-head parameters (which would also + # break DDP's find_unused_parameters=False). + les_readout = 'edge_basis' if use_les else 'sum' # ── Split train → train + val ───────────────────────────────────────── idx = np.random.permutation(len(train_raw)) n_val_actual = min(n_val, len(train_raw) // 10) diff --git a/scripts/train_ecenet_xyz.py b/scripts/train_ecenet_xyz.py index c8148f5..388cc29 100644 --- a/scripts/train_ecenet_xyz.py +++ b/scripts/train_ecenet_xyz.py @@ -122,7 +122,7 @@ def train_ecenet_xyz( # Long-range (LES) use_les=False, les_arguments=None, # extra kwargs for upstream les.Les (see ecenet/les.py) - les_readout='sum', # (l0,l1) read-out: 'sum' | 'softmax' | 'edge' | 'edge_basis' + les_readout=None, # None -> 'edge_basis' if use_les else 'sum' les_charge_scale=1.0, # fixed multiplier on the edge-mode latent charge (MACELES: 0.1) les_dipole=False, # edge head also emits bond dipoles; l0 packed [q | u] les_charges=True, # False (needs les_dipole): dipoles-only — q hard zero, standard-init dipole head @@ -247,6 +247,11 @@ def train_ecenet_xyz( type_map = elements.build_type_map( z for s in (train_raw + test_raw) for z in s['numbers']) + if les_readout is None: + # 'edge_basis' is the LES default; short-range runs take 'sum' so the + # model carries no unused charge-head parameters (which would also + # break DDP's find_unused_parameters=False). + les_readout = 'edge_basis' if use_les else 'sum' n_types = len(type_map) if verbose: n_atoms_list = [s['n_atoms'] for s in train_use] diff --git a/tests/test_trainer_les.py b/tests/test_trainer_les.py index 8684a49..70a6c33 100644 --- a/tests/test_trainer_les.py +++ b/tests/test_trainer_les.py @@ -81,8 +81,9 @@ def test_rmd17_use_les(): assert res['les_module'] is not None ck = torch.load(ckpt, map_location='cpu', weights_only=False) assert 'les' in ck and ck['les']['state_dict'] is not None - # resume with use_les continues from the checkpoint... - train_ecenet(use_les=True, checkpoint_path=ckpt, + # resume with use_les continues from the checkpoint (les_readout must + # match the checkpoint's — the default now resolves to 'edge_basis')... + train_ecenet(use_les=True, les_readout='sum', checkpoint_path=ckpt, **{**common, 'n_epochs': 3}) # ...and a short-range run against an LES checkpoint is refused. try: diff --git a/tests/test_xyz_trainer.py b/tests/test_xyz_trainer.py index b297741..d4b5992 100644 --- a/tests/test_xyz_trainer.py +++ b/tests/test_xyz_trainer.py @@ -104,7 +104,7 @@ def test_smoke_train_les(): ckpt = os.path.join(td, 'xyz_les.mdl') _, les_module, results = train_ecenet_xyz( train_structures=[dict(s) for s in structs], n_val=2, - use_les=True, checkpoint_path=ckpt, + use_les=True, les_readout='sum', checkpoint_path=ckpt, n_epochs=2, batch_size=4, lr=5e-3, **COMMON, ) assert les_module is not None @@ -116,7 +116,7 @@ def test_smoke_train_les(): # Resume: fresh call restores model + LES + optimizer and continues. _, les2, results2 = train_ecenet_xyz( train_structures=[dict(s) for s in structs], n_val=2, - use_les=True, checkpoint_path=ckpt, + use_les=True, les_readout='sum', checkpoint_path=ckpt, n_epochs=4, batch_size=4, lr=5e-3, **COMMON, ) assert np.isfinite(results2['val_force_mae']) @@ -224,7 +224,7 @@ def test_calculator_rejects_les_checkpoint(): ckpt = os.path.join(td, 'xyz_les.mdl') train_ecenet_xyz( train_structures=make_structures(8, seed=9), n_val=2, - use_les=True, checkpoint_path=ckpt, + use_les=True, les_readout='sum', checkpoint_path=ckpt, n_epochs=1, batch_size=4, lr=5e-3, **COMMON, ) try: