readme: condense to an overview

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
@@ -115,10 +116,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
@@ -153,12 +155,12 @@ 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,
         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,
@@ -235,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

scripts/train_ecenet.py

@@ -168,19 +168,19 @@ 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,
     mp_l_attention=False,
     # FiLM gate
-    element_film=False,
+    element_film=True,
     film_embed_dim=16,
     film_n_rbf=0,
     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:

scripts/train_ecenet_mptrj.py

@@ -694,19 +694,19 @@ 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,
     mp_l_attention=False,
     # FiLM gate
-    element_film=False,
+    element_film=True,
     film_embed_dim=16,
     film_n_rbf=0,
     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

scripts/train_ecenet_spice.py

@@ -328,19 +328,19 @@ 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,
     mp_l_attention=False,
     # FiLM gate
-    element_film=False,
+    element_film=True,
     film_embed_dim=16,
     film_n_rbf=0,
     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)

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
@@ -145,13 +145,13 @@ 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,
     mp_l_attention=False,
     # FiLM gate
-    element_film=False,
+    element_film=True,
     film_embed_dim=16,
     film_n_rbf=0,
     film_hidden=None,
@@ -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]

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,20 +67,20 @@ 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')
     except ValueError as e:
         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()

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")

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}")

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:

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: