diff --git a/README.md b/README.md index 562aac0..57137c7 100644 --- a/README.md +++ b/README.md @@ -167,8 +167,10 @@ model, les_module, results = train_ecenet_xyz( 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. +alternative read-outs and the related options — latent dipoles +(`les_dipole`), latent polarizabilities (`les_alpha='iso'|'aniso'`, the +induced-dipole term of the polarizable-multipole LES), 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, @@ -185,8 +187,12 @@ print(atoms.get_potential_energy()) # E_sr + E_lr, eV ``` 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. +(and dipoles / polarizabilities as `calc.results['les_dipoles']` / +`calc.results['les_alphas']`); `calc.compute_bec(atoms)` returns Born +effective charges `Z* = ∂P/∂r` with charge-flow terms included, and for a +`les_alpha` checkpoint `calc.compute_polarizability(atoms)` returns the +molecular polarizability tensor `Σᵢ αᵢ` (latent units; physical for isolated +molecules, bulk needs the ε∞ unscaling of the LES paper). `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 diff --git a/ecenet/calculator.py b/ecenet/calculator.py index cc27dea..a6e7f66 100644 --- a/ecenet/calculator.py +++ b/ecenet/calculator.py @@ -505,8 +505,9 @@ class ECENetLESCalculator(ECENetCalculator): ECENetCalculator refusing LES ones). The upstream charge head is materialised and its trained state loaded here; edge-mode read-outs (``les_readout='edge'/'edge_basis'``) carry the charge (and with - ``les_dipole`` the bond dipoles) inside the model itself, so the LES - module is then parameter-free. + ``les_dipole`` the bond dipoles, with ``les_alpha`` the latent + polarizabilities) inside the model itself, so the LES module is then + parameter-free. """ _uses_cell = True @@ -552,10 +553,12 @@ def from_checkpoint(cls, checkpoint_path, device=None, dtype=None, # ── Energy seams: add E_lr on the same graph ──────────────────────────── # # Both seams also stash the per-atom latent charges (and, with - # ``les_dipole``, the latent atomic dipoles) into ``self.results`` as a - # side effect — every force call computes them anyway, so exposing them - # is free. ``atoms.get_charges()`` reads ``results['charges']`` (e); - # ``results['les_dipoles']`` holds u (N, 3) in e·Å. Stashing is safe in + # ``les_dipole`` / ``les_alpha``, the latent atomic dipoles and + # polarizabilities) into ``self.results`` as a side effect — every force + # call computes them anyway, so exposing them is free. + # ``atoms.get_charges()`` reads ``results['charges']`` (e); + # ``results['les_dipoles']`` holds u (N, 3) in e·Å; ``results['les_alphas']`` + # holds α^les as (N,) or (N, 3, 3) in e·Å/(V/Å). Stashing is safe in # the strain (stress) pass too: it evaluates at ε = 0, i.e. the # unstrained geometry, so it stores the same numbers. The global sign of # the latent charges is arbitrary (E_lr is quadratic in q) — consistent @@ -563,8 +566,13 @@ def from_checkpoint(cls, checkpoint_path, device=None, dtype=None, def _stash_charges(self, q, l0): self.results['charges'] = q.detach().cpu().numpy().reshape(-1) - if self.model.les_dipole: - self.results['les_dipoles'] = l0[:, 1:4].detach().cpu().numpy() + if self.model.les_dipole or self.model.les_alpha: + from ecenet.les import unpack_l0 + _, u, alpha = unpack_l0(l0, **self.les_flags) + if u is not None: + self.results['les_dipoles'] = u.detach().cpu().numpy() + if alpha is not None: + self.results['les_alphas'] = alpha.detach().cpu().numpy() def _energy_free(self, pos, types): e_sr, l0 = self.model.forward(pos, types, return_embeddings=True, @@ -595,7 +603,9 @@ def compute_bec(self, atoms): (Berry-phase-style for periodic cells, direct sum otherwise), mean-charge removal, and the √ε∞ normalisation as configured in the checkpoint's ``les_arguments``; with ``les_dipole`` the charge and - dipole parts are summed. + dipole parts are summed, and with ``les_alpha`` the induced dipoles + Δu_i = α_i·E_i are part of the polarization (see + ``LESLongRange.born_charges``). Unlike the latent-charge stash this cannot ride along ``calculate()`` — the force backward frees the graph, and Z* needs gradients of the @@ -623,21 +633,48 @@ def compute_bec(self, atoms): _, l0 = self.model.forward(pos, types, return_embeddings=True, l0_only=True) cell_t = None - if self.les_flags['l0_is_charge']: - q = l0[:, 0] - u = l0[:, 1:4] if self.les_flags['les_dipole'] else None - else: - # upstream's atomwise head maps the l0 descriptor to charges - q = self.les_module.les.atomwise( - l0.reshape(l0.shape[0], -1), - torch.zeros(len(symbols), dtype=torch.long, - device=self.device)) - u = None - bec = self.les_module.les.bec(q=q, r=pos, cell=cell_t, u=u) - if bec.dim() == 4: # (N, 2, 3, 3): charge + dipole parts - bec = bec.sum(dim=1) + bec = self.les_module.born_charges(l0, pos, cell=cell_t, + **self.les_flags) return bec.detach().cpu().numpy() + def compute_polarizability(self, atoms): + """Polarizability tensor α = Σ_i α_i for one structure, (3, 3). + + Needs a ``les_alpha`` checkpoint. The induced dipoles respond + linearly and non-self-consistently to the field, and the fixed + multipoles don't respond to an external field at all, so the + polarization response is exactly the sum of the atomic tensors + (paper Eq. 29) — one forward, no gradients. An isotropic ``'iso'`` + model gives a multiple of the identity. Units e·Å/(V/Å) = e·Å²/V + (multiply by 1/(4πε₀) = 14.3996 eV·Å/e² for ų). This is the latent + α^les: physical for an isolated molecule (ε_e = 1); for bulk it needs + the ε_e = ε_∞/(1 + χ^les) unscaling of the paper's Eqs. 23–24, which + needs ε_∞ and is left to the caller. Unlike the latent charges the + sign is physical: E_lr is linear in α. + """ + if self.model.les_alpha is None: + raise ValueError("compute_polarizability needs a checkpoint " + "trained with les_alpha ('iso' or 'aniso')") + from ecenet.les import total_polarizability, unpack_l0 + symbols = atoms.get_chemical_symbols() + types = self._types(symbols) + pos = torch.tensor(atoms.get_positions(), dtype=self.dtype, + device=self.device) + with torch.no_grad(): + if atoms.pbc.any(): + cell_np = atoms.get_cell().array + (edge_i, edge_j, she, + nb_src, nb_dst, shn) = self._neighbor_lists(pos, cell_np) + _, l0 = self.model.forward_pbc( + pos, types, edge_i, edge_j, she, nb_src, nb_dst, shn, + return_embeddings=True, l0_only=True) + else: + _, l0 = self.model.forward(pos, types, return_embeddings=True, + l0_only=True) + _, _, alpha = unpack_l0(l0, **self.les_flags) + pol = total_polarizability(alpha)[0] + return pol.cpu().numpy() + def load_calculator(checkpoint_path, verbose=True, **kwargs): """Load the right calculator for a checkpoint, whatever it was trained with. diff --git a/ecenet/les.py b/ecenet/les.py index 7d7898c..38d2585 100644 --- a/ecenet/les.py +++ b/ecenet/les.py @@ -61,6 +61,57 @@ def _upstream_les(): return Les +def unpack_l0(l0, l0_is_charge=True, les_dipole=False, les_alpha=None): + """Split an edge-mode packed ``l0`` (N, 1 + 3·dipole + α) into (q, u, α). + + The one place that knows the packed layout ``[q | u_xyz | α]`` the + model's edge head emits (see ``ECENet.__init__``); call with + ``**model.les_flags``. Returns q (N,), u (N, 3) or None, and α as (N,) + for ``'iso'``, (N, 3, 3) for ``'aniso'``, or None — the shapes upstream's + ``Les.forward`` takes for ``latent_dipoles`` / ``latent_alphas``. + """ + if not l0_is_charge: + raise ValueError("unpack_l0 needs an edge-mode l0 (l0_is_charge=True); " + "atomwise read-outs carry a descriptor, not [q | u | α]") + n = l0.shape[0] + q = l0[:, 0] + k = 1 + u = None + if les_dipole: + u = l0[:, 1:4] + k = 4 + alpha = None + if les_alpha == 'iso': + alpha = l0[:, k] + elif les_alpha == 'aniso': + alpha = l0[:, k:k + 9].reshape(n, 3, 3) + elif les_alpha is not None: + raise ValueError(f"unknown les_alpha {les_alpha!r}") + return q, u, alpha + + +def total_polarizability(alpha, batch=None, n_struct=None): + """Per-structure polarizability tensor Σ_i α_i, (B, 3, 3). + + The model's induced dipoles respond linearly and non-self-consistently + to the field, and the fixed multipoles don't respond at all, so the + polarization response to a uniform external field is exactly the sum + of the atomic tensors (paper Eq. 29) — no autograd needed. An isotropic + α (N,) is broadcast to α·I. This is the *latent* α^les; the physical + tensor is ε_e·α^les with ε_e = 1 in vacuum (isolated molecules) and + ε_e = ε_∞/(1 + χ^les) for bulk (paper Eqs. 23–24). + """ + if alpha.dim() == 1: + eye = torch.eye(3, device=alpha.device, dtype=alpha.dtype) + alpha = alpha[:, None, None] * eye + if batch is None: + return alpha.sum(dim=0, keepdim=True) + if n_struct is None: + n_struct = int(batch.max().item()) + 1 + out = torch.zeros(n_struct, 3, 3, device=alpha.device, dtype=alpha.dtype) + return out.index_add_(0, batch, alpha) + + class LESLongRange(nn.Module): """Long-range electrostatic energy from per-atom invariant embeddings. @@ -95,7 +146,8 @@ def forward(self, l0: torch.Tensor, positions: torch.Tensor, return_charges: bool = False, n_struct: int | None = None, l0_is_charge: bool = False, - les_dipole: bool = False): + les_dipole: bool = False, + les_alpha: str | None = None): """Long-range energy for one structure or a packed batch. l0 (N, C) per-atom invariant descriptor (any flattenable @@ -106,7 +158,12 @@ def forward(self, l0: torch.Tensor, positions: torch.Tensor, module holds no parameters. With ``les_dipole=True`` (requires ``l0_is_charge``), l0 is the model's packed (N, 4) = [q | u] and the latent atomic dipoles u are passed - to upstream's charge–dipole/dipole–dipole terms. + to upstream's charge–dipole/dipole–dipole terms. With + ``les_alpha`` ('iso' / 'aniso', requires ``l0_is_charge``) + the packed l0 also carries the latent polarizability α + (see ``unpack_l0``), passed to upstream's induced-dipole + term −½ E_i·α_i·E_i (E_i the field of the fixed + multipoles at atom i). positions (N, 3) in Å, on the same autograd graph as the SR energy cell (B, 3, 3) or (3, 3); None → isolated / non-periodic, served by the vectorized batched path below (verified equal to @@ -120,20 +177,24 @@ def forward(self, l0: torch.Tensor, positions: torch.Tensor, ``return_charges=True``, also the per-atom latent charges — the q column only under ``les_dipole``; take u from l0 directly). """ - if les_dipole and not l0_is_charge: - raise ValueError("les_dipole=True requires l0_is_charge=True " - "(the packed [q | u] comes from the model's " + if (les_dipole or les_alpha) and not l0_is_charge: + raise ValueError("les_dipole / les_alpha require l0_is_charge=True " + "(the packed [q | u | α] comes from the model's " "edge head).") if batch is None: batch = torch.zeros(positions.shape[0], dtype=torch.long, device=positions.device) if n_struct is None: n_struct = 1 - # unpack once; the branches below carry u alongside the charges - q, u = (l0[:, 0], l0[:, 1:4]) if les_dipole else (l0, None) + # unpack once; the branches below carry u and α alongside the charges + if les_dipole or les_alpha: + q, u, alpha = unpack_l0(l0, True, les_dipole, les_alpha) + else: + q, u, alpha = l0, None, None if cell is None: return self._isolated_batched(q, positions, batch, n_struct, - return_charges, l0_is_charge, u=u) + return_charges, l0_is_charge, + u=u, alpha=alpha) # Scope the default dtype to the input's so upstream's lazily built # charge MLP (and any default-dtype internals) match float64 inputs. prev_dtype = torch.get_default_dtype() @@ -143,6 +204,7 @@ def forward(self, l0: torch.Tensor, positions: torch.Tensor, result = self.les( latent_charges=q.reshape(-1), latent_dipoles=u, + latent_alphas=alpha, positions=positions, cell=cell.view(-1, 3, 3), batch=batch, @@ -165,7 +227,7 @@ def forward(self, l0: torch.Tensor, positions: torch.Tensor, return result["E_lr"] def _isolated_batched(self, l0, positions, batch, n_struct, return_charges, - l0_is_charge=False, u=None): + l0_is_charge=False, u=None, alpha=None): """Isolated (non-periodic) long-range energy, vectorized over the batch. Upstream's ``Les.forward`` loops over structures in Python — masked @@ -185,9 +247,15 @@ def _isolated_batched(self, l0, positions, batch, n_struct, return_charges, E_b = ½ qᵀf_qq q + (u·f_qu)ᵀq − ½ uᵀf_uu u (i,j ∈ b) (the qu coefficient is 1, not ½, and the self-interaction is removed - by ``make_kernels`` — both upstream's conventions). Verified equal to + by ``make_kernels`` — both upstream's conventions). With latent + polarizabilities ``alpha`` ((N,) isotropic or (N, 3, 3)) the field + of the fixed multipoles at each atom is formed from the same masked + kernels, E_j = Σ_i q_i f_qu[i,j] + Σ_i u_i·f_uu[i,j], and the + induced-dipole energy −½ Σ_j E_j·α_j·E_j is added, mirroring + upstream's ``_get_induced_u`` (non-self-consistent: the field + excludes the induced dipoles themselves). Verified equal to upstream's loop (energies, charges, and position gradients) in - tests/test_les.py, charge-only and with dipoles. + tests/test_les.py, charge-only, with dipoles, and with alphas. Tradeoff: the dense kernel spans ALL atom pairs, so this does (ΣN)² pair work where the loop does Σ(N_b²) — ~batch_size× redundant @@ -237,7 +305,8 @@ def _isolated_batched(self, l0, positions, batch, n_struct, return_charges, ew = self.les.ewald f_qq, f_qu, f_uu, _, _ = make_kernels(positions + shift, ew.sigma, ew.norm_factor / ew.twopi, - compute_u=u is not None, + compute_u=(u is not None + or alpha is not None), compute_Q=False) same = (batch.unsqueeze(0) == batch.unsqueeze(1)).to(f_qq.dtype) e_phi = torch.einsum('iq,ij->jq', q, f_qq * same) @@ -255,6 +324,19 @@ def _isolated_batched(self, l0, positions, batch, n_struct, return_charges, G = torch.einsum('ic,ijcd->ijd', u_, f_uu) # (N, N, 3) T = torch.einsum('ijd,jd->ij', G, u_) per_atom = per_atom - 0.5 * (T * same).sum(dim=0) + if alpha is not None: + # field of the fixed multipoles at j (upstream's e_field: charges + # via f_qu, dipoles via f_uu — the latter is G above), masked to + # the structure; then Δu_j = α_j·E_j and U^iu_j = −½ E_j·Δu_j + e_field = (q[:, :, None] * f_qu * same[:, :, None]).sum(dim=0) + if u is not None: + e_field = e_field + (G * same[:, :, None]).sum(dim=0) + alpha_ = alpha.to(positions.dtype) + if alpha_.dim() == 1: + u_ind = e_field * alpha_[:, None] + else: + u_ind = torch.einsum('jc,jcd->jd', e_field, alpha_) + per_atom = per_atom - 0.5 * (e_field * u_ind).sum(dim=-1) e_lr = torch.zeros(n_struct, dtype=per_atom.dtype, device=per_atom.device ).scatter_add(0, batch, per_atom) @@ -262,6 +344,46 @@ def _isolated_batched(self, l0, positions, batch, n_struct, return_charges, return e_lr, charges return e_lr + def born_charges(self, l0, positions, cell=None, l0_is_charge=False, + les_dipole=False, les_alpha=None): + """Born effective charges Z* = ∂P/∂r for ONE structure, (N, 3, 3). + + ``positions`` must be on the autograd graph that produced ``l0``, so + the latent variables stay functions of the positions and upstream's + BEC module delivers the charge-flow terms. Upstream handles the + polarization (Berry-phase-style for a periodic ``cell``, direct sum + for ``None``), mean-charge removal, and the √ε∞ normalisation from + ``les_arguments``; charge and dipole parts are summed. With + ``les_alpha`` the induced dipoles Δu_i = α_i·E_i are part of the + polarization, so the call goes through upstream's full forward (one + extra field evaluation); otherwise straight to the BEC module. The + single implementation behind the calculator and the eval tools. + """ + prev_dtype = torch.get_default_dtype() + torch.set_default_dtype(positions.dtype) + try: + cell_t = None if cell is None else cell.view(-1, 3, 3) + if l0_is_charge: + q, u, alpha = unpack_l0(l0, True, les_dipole, les_alpha) + if alpha is not None: + res = self.les(latent_charges=q, latent_dipoles=u, + latent_alphas=alpha, positions=positions, + cell=cell_t, compute_energy=True, + compute_bec=True) + bec = res['BEC'] + else: + bec = self.les.bec(q=q, r=positions, cell=cell_t, u=u) + else: + batch = torch.zeros(l0.shape[0], dtype=torch.long, + device=l0.device) + q = self.les.atomwise(l0.reshape(l0.shape[0], -1), batch) + bec = self.les.bec(q=q, r=positions, cell=cell_t) + finally: + torch.set_default_dtype(prev_dtype) + if bec.dim() == 4: # (N, 2, 3, 3): charge + dipole parts + bec = bec.sum(dim=1) + return bec + def load_les_module(ckpt_les, model, device, dtype, load_state=True): """Rebuild a ready-to-use ``LESLongRange`` from a checkpoint's ``les`` dict. diff --git a/ecenet/model.py b/ecenet/model.py index bad19bd..26c025b 100644 --- a/ecenet/model.py +++ b/ecenet/model.py @@ -45,6 +45,11 @@ # head bypassed via l0_is_charge). The single definition — consumers read the # derived facts off the model via `les_flags` rather than re-spelling this. _LES_EDGE_MODES = ('edge', 'edge_basis') +# les_alpha modes → width of the α slot in the packed l0 (1 scalar, or the +# full 3×3 tensor flattened) and number of per-edge head scalars it needs +# (a_e for the isotropic part; a_e + b_e for isotropic + bond-axial parts). +_LES_ALPHA_WIDTH = {None: 0, 'iso': 1, 'aniso': 9} +_LES_ALPHA_BLOCKS = {None: 0, 'iso': 1, 'aniso': 2} # --------------------------------------------------------------------------- # Main model @@ -170,6 +175,7 @@ def __init__( les_charge_scale: float = 1.0, les_dipole: bool = False, les_charges: bool = True, + les_alpha: str | None = None, ): super().__init__() if mp_type == 'transformer': @@ -286,7 +292,42 @@ def __init__( "for the LES energy.") self.les_dipole = bool(les_dipole) self.les_charges = bool(les_charges) - self._l0_dim = (4 if les_dipole else 1) if _edge_mode else 2 * embed_dim + # les_alpha: the edge head also emits a per-atom latent polarizability + # α_i (Allegro-LES's construction; Kim, King, Park et al. 2026, arXiv + # 2605.05746) for upstream's induced-dipole term — each atom's dipole + # responds linearly to the field of the fixed multipoles, Δu_i = α_i·E_i, + # lowering the energy by ½E_i·α_i·E_i (non-self-consistent: induced + # dipoles don't see each other). Appended to the packed l0 after [q | u]: + # 'iso': one scalar a_e per edge, scatter-summed → α_i (1 column) + # 'aniso': two scalars per edge, a_e·I + b_e·(r̂r̂ᵀ − I/3) — an + # isotropic part plus a traceless bond-axial part, symmetric + # and equivariant by construction (the Cartesian form of an + # l=0 + l=2 edge harmonic; a lone bond gives exactly the + # uniaxial α_∥/α_⊥ of a diatomic) — summed → α_i (9 columns, + # the flattened 3×3) + # Head slots are zero-init: α ≡ 0 at init leaves an existing model + # untouched at step 0, and is NOT a saddle (E_lr is linear in α with + # gradient −½|E_i|², nonzero once charges exist). Unlike q, the sign of + # α is physical (E_lr is linear in it), and les_charge_scale does not + # touch it (a response coefficient, not a source). Physical α needs the + # ε_e unscaling of the paper's Eqs. 23–24 (ε_e = 1 in vacuum). + if les_alpha not in _LES_ALPHA_WIDTH: + raise ValueError("les_alpha must be None, 'iso' or 'aniso', " + f"got {les_alpha!r}") + if les_alpha is not None and not _edge_mode: + raise ValueError( + f"les_alpha={les_alpha!r} requires les_readout='edge' or " + f"'edge_basis' (got {les_readout!r}): the polarizability is " + "emitted by the per-edge charge head, which the atomwise " + "read-outs don't have.") + self.les_alpha = les_alpha + self._n_alpha = _LES_ALPHA_WIDTH[les_alpha] + self._l0_dim = ((1 + (3 if les_dipole else 0) + self._n_alpha) + if _edge_mode else 2 * embed_dim) + # head slot counts (q, d_e, α scalars): each slot is one n_output_out + # block of the 'edge_basis' MLP's last layer, or one 'edge' linear row + self._les_head_blocks = (int(les_charges), int(les_dipole), + _LES_ALPHA_BLOCKS[les_alpha]) # les_charge_scale: fixed multiplier on the edge-mode latent charge # — the whole packed [q | u] when les_dipole is on, keeping the q–u # sign coupling intact (MACELES's output_scale; they ship 0.1). With standard head init, @@ -312,12 +353,14 @@ def __init__( nn.init.zeros_(self.les_score.weight) nn.init.zeros_(self.les_score.bias) elif les_readout == 'edge': - n_head_out = (2 if les_dipole else 1) if les_charges else 1 - self.les_edge_charge = nn.Linear(2 * embed_dim, n_head_out, + n_q, n_u, n_a = self._les_head_blocks + self.les_edge_charge = nn.Linear(2 * embed_dim, n_q + n_u + n_a, bias=False) - if les_dipole and les_charges: - with torch.no_grad(): + with torch.no_grad(): + if les_dipole and les_charges: self.les_edge_charge.weight[1].zero_() # dipole slot + if n_a: + self.les_edge_charge.weight[n_q + n_u:].zero_() # α slots # les_charges=False: the single (standard-init) row IS the dipole # scalar — zero-init would sit on the uu-quadratic saddle (above) # (the 'edge_basis' head is built with the output MLP below) @@ -448,15 +491,19 @@ def __init__( # last layer widens to a second n_output_out block (dipole channels, # dotted with the same radial basis), zero-init per the note above. if les_readout == 'edge_basis': - n_blocks = (2 if les_dipole else 1) if les_charges else 1 - q_dims = mlp_dims[:-1] + [n_output_out * n_blocks] + n_q, n_u, n_a = self._les_head_blocks + q_dims = mlp_dims[:-1] + [n_output_out * (n_q + n_u + n_a)] self.les_edge_charge = OutputMLP(q_dims, activation=act(), zero_init_last=False) - if les_dipole and les_charges: - with torch.no_grad(): - last = self.les_edge_charge.linears[-1] - last.weight[n_output_out:].zero_() # dipole block - last.bias[n_output_out:].zero_() + with torch.no_grad(): + last = self.les_edge_charge.linears[-1] + if les_dipole and les_charges: + last.weight[n_output_out:2 * n_output_out].zero_() # dipole + last.bias[n_output_out:2 * n_output_out].zero_() + if n_a: # α blocks + k = (n_q + n_u) * n_output_out + last.weight[k:].zero_() + last.bias[k:].zero_() # les_charges=False: the single (standard-init) block IS the # dipole — see the saddle note at the les_charges check above @@ -470,14 +517,16 @@ def __init__( def les_flags(self): """The l0 convention as ``LESLongRange.forward`` kwargs. - ``{'l0_is_charge': ..., 'les_dipole': ...}`` — the single source of - truth for how this model's ``l0`` read-out is to be interpreted - (edge modes: l0 IS the charge, packed [q | u] under ``les_dipole``). + ``{'l0_is_charge': ..., 'les_dipole': ..., 'les_alpha': ...}`` — the + single source of truth for how this model's ``l0`` read-out is to be + interpreted (edge modes: l0 IS the charge, packed [q | u | α] under + ``les_dipole`` / ``les_alpha``; ``ecenet.les.unpack_l0`` splits it). Call sites do ``les_module(l0, pos, ..., **model.les_flags)`` instead of re-deriving the flags from hparams or ``les_readout`` literals. """ return {'l0_is_charge': self.les_readout in _LES_EDGE_MODES, - 'les_dipole': self.les_dipole} + 'les_dipole': self.les_dipole, + 'les_alpha': self.les_alpha} def _compute_ace_basis(self, pos_batch, nb_src, nb_dst, types, shift_vecs_nb=None): """Compute ACE atomic basis: (B, N, n_types, n_max, n_sph).""" @@ -651,7 +700,9 @@ def _aggregate_lr_embeddings(self, A_cos, A_sin, r_hat, edge_j, n_atoms, Returns: l0: (n_atoms, 2*embed_dim) per-atom invariant scalar embeddings — edge modes: (n_atoms, 1), the latent charge itself; with - les_dipole, packed (n_atoms, 4) = [q | u_x u_y u_z] + les_dipole, packed (n_atoms, 4) = [q | u_x u_y u_z]; with + les_alpha, a further 1 ('iso') or 9 ('aniso', flattened 3×3) + columns of latent polarizability (see __init__) l1: (n_atoms, 2*embed_dim, 3) per-atom equivariant vector embeddings with_l1=False skips the l=1 Wigner rotation + scatter and returns @@ -710,17 +761,30 @@ def _aggregate_lr_embeddings(self, A_cos, A_sin, r_hat, edge_j, n_atoms, out = out * env[:, None] else: out = self.les_edge_charge(h_l0) # (E, 1 | 2) - if self.les_dipole: - if self.les_charges: - h_l0 = torch.cat([out[:, :1], out[:, 1:2] * r_hat], dim=1) - else: - # dipoles-only: the head's single block is the dipole - # scalar; the q column stays hard zero so the packed - # [q | u] layout — and every consumer of it — is unchanged - h_l0 = torch.cat([torch.zeros_like(out[:, :1]), - out[:, :1] * r_hat], dim=1) + # Pack the per-edge contributions as [q | d·r̂ | α] (see __init__). + # dipoles-only (les_charges=False): the q column stays hard zero + # so the packed layout — and every consumer of it — is unchanged. + k = 0 + if self.les_charges: + cols = [out[:, :1]] + k = 1 else: - h_l0 = out + cols = [torch.zeros_like(out[:, :1])] + if self.les_dipole: + cols.append(out[:, k:k + 1] * r_hat) + k += 1 + if self.les_alpha == 'iso': + cols.append(out[:, k:k + 1]) + k += 1 + elif self.les_alpha == 'aniso': + # a_e·I + b_e·(r̂r̂ᵀ − I/3): isotropic + traceless bond-axial + eye = torch.eye(3, device=device, dtype=dtype) + tr0 = r_hat[:, :, None] * r_hat[:, None, :] - eye / 3.0 + alpha_e = (out[:, k, None, None] * eye + + out[:, k + 1, None, None] * tr0) # (E, 3, 3) + cols.append(alpha_e.reshape(n_e, 9)) + k += 2 + h_l0 = cols[0] if len(cols) == 1 else torch.cat(cols, dim=1) elif self.les_readout == 'softmax': # Same shape as the MP layers' softmax path (one score slot): # a_e = exp(s_e)·f_cut_e / (Σ_{e'→j} exp(s_e')·f_cut_e' + eps) · f_cut_e @@ -750,7 +814,13 @@ def _aggregate_lr_embeddings(self, A_cos, A_sin, r_hat, edge_j, n_atoms, # l1 keeps the plain unweighted sum) if (self.les_charge_scale != 1.0 and self.les_readout in _LES_EDGE_MODES): - l0 = l0 * self.les_charge_scale + if self._n_alpha: + # the α slot is a response coefficient, not a source: unscaled + n_qu = l0.shape[1] - self._n_alpha + l0 = torch.cat([l0[:, :n_qu] * self.les_charge_scale, + l0[:, n_qu:]], dim=1) + else: + l0 = l0 * self.les_charge_scale if not with_l1: return l0, None diff --git a/scripts/train_ecenet.py b/scripts/train_ecenet.py index a28a136..744b6bf 100644 --- a/scripts/train_ecenet.py +++ b/scripts/train_ecenet.py @@ -184,6 +184,7 @@ def train_ecenet( 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 + les_alpha=None, # None | 'iso' | 'aniso': edge head also emits polarizabilities α (induced dipoles) # Joint LES long-range training: E = E_sr + E_lr on one autograd graph # (isolated pairwise path — rMD17/MD22 molecules have no cell). NOTE on # units: these datasets are in kcal/mol while the LES Coulomb constant is @@ -312,6 +313,7 @@ def loss_fn(pred, tgt): les_charge_scale=les_charge_scale, les_dipole=les_dipole, les_charges=les_charges, + les_alpha=les_alpha, ) if dtype == torch.float64: model = model.double() @@ -491,6 +493,7 @@ def save_checkpoint(epoch): les_charge_scale=les_charge_scale, les_dipole=les_dipole, les_charges=les_charges, + les_alpha=les_alpha, ), # molecule-specific element mapping: {symbol: type_index} 'element_to_type': elements.to_element_to_type(type_to_idx), diff --git a/scripts/train_ecenet_mptrj.py b/scripts/train_ecenet_mptrj.py index cf4c78e..a98f102 100644 --- a/scripts/train_ecenet_mptrj.py +++ b/scripts/train_ecenet_mptrj.py @@ -710,6 +710,7 @@ def train_ecenet_mptrj( 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 + les_alpha=None, # None | 'iso' | 'aniso': edge head also emits polarizabilities α (induced dipoles) # Joint LES long-range training: E = E_sr + E_lr on one autograd graph. # Periodic structures use upstream's reciprocal-space Ewald, so every # frame needs its cell tensor — prepared shards written before cells were @@ -993,6 +994,7 @@ def train_ecenet_mptrj( les_charge_scale=les_charge_scale, les_dipole=les_dipole, les_charges=les_charges, + les_alpha=les_alpha, ) if dtype == torch.float64: model = model.double() @@ -1161,6 +1163,7 @@ def save_checkpoint(epoch): les_charge_scale=les_charge_scale, les_dipole=les_dipole, les_charges=les_charges, + les_alpha=les_alpha, ), 'element_to_type': elements.to_element_to_type(type_map), # {symbol: type_idx} 'e_ref': e_ref, # per-element reference energies (eV) diff --git a/scripts/train_ecenet_spice.py b/scripts/train_ecenet_spice.py index f9dc04c..1130da2 100644 --- a/scripts/train_ecenet_spice.py +++ b/scripts/train_ecenet_spice.py @@ -344,6 +344,7 @@ def train_ecenet_spice( 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 + les_alpha=None, # None | 'iso' | 'aniso': edge head also emits polarizabilities α (induced dipoles) # Long-range (LES): E = E_sr + E_lr on one autograd graph. Needs the # optional `les` package (see ecenet/les.py for install + licensing). use_les=False, @@ -510,6 +511,7 @@ def train_ecenet_spice( les_charge_scale=les_charge_scale, les_dipole=les_dipole, les_charges=les_charges, + les_alpha=les_alpha, ) if dtype == torch.float64: model = model.double() @@ -717,6 +719,7 @@ def save_checkpoint(epoch): les_charge_scale=les_charge_scale, les_dipole=les_dipole, les_charges=les_charges, + les_alpha=les_alpha, ), 'e_ref': e_ref, # per-element reference energies (eV/atom) # Self-describing metadata for the calculator (no dataset coupling). diff --git a/scripts/train_ecenet_xyz.py b/scripts/train_ecenet_xyz.py index 388cc29..94aecf0 100644 --- a/scripts/train_ecenet_xyz.py +++ b/scripts/train_ecenet_xyz.py @@ -126,6 +126,7 @@ def train_ecenet_xyz( 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 + les_alpha=None, # None | 'iso' | 'aniso': edge head also emits polarizabilities (induced dipoles) # Geometry r_cut_edge=5.0, r_cut_neighbor=4.0, @@ -295,6 +296,7 @@ def train_ecenet_xyz( les_charge_scale=les_charge_scale, les_dipole=les_dipole, les_charges=les_charges, + les_alpha=les_alpha, ) if dtype == torch.float64: model = model.double() @@ -419,6 +421,7 @@ def save_checkpoint(epoch): les_charge_scale=les_charge_scale, les_dipole=les_dipole, les_charges=les_charges, + les_alpha=les_alpha, ), 'element_to_type': elements.to_element_to_type(type_map), 'e_ref': e_ref, diff --git a/tests/test_les.py b/tests/test_les.py index 89c5138..b5183a4 100644 --- a/tests/test_les.py +++ b/tests/test_les.py @@ -11,6 +11,10 @@ works and constructing `LESLongRange` raises an ImportError carrying the install hint; with it installed, a smoke forward returns a finite energy. +With les_alpha ('iso'/'aniso'): the packed [q | u | α] layout, α's +equivariance and uniaxial single-bond form, and the wrapper's induced-dipole +energy against upstream's loop and finite differences. + Run: python tests/test_les.py """ @@ -497,6 +501,241 @@ def test_edge_dipole_les_energy(): f"({dc:.1e}); flag misuse rejected") +def test_edge_alpha(): + """les_alpha: l0 packed [q | u | α]; α exactly 0 at init (zero-init + slots); 'iso' an invariant scalar, 'aniso' a symmetric rank-2 tensor + transforming as α' = RαRᵀ; a lone bond gives a uniaxial tensor with r̂ + as principal axis (the diatomic α_∥/α_⊥); les_charge_scale scales [q | u] + but not α; batched variants consistent; non-edge read-outs rejected.""" + from ecenet.les import total_polarizability, unpack_l0 + for mode in ('edge', 'edge_basis'): + for kind, width in (('iso', 1), ('aniso', 9)): + m = make_model(seed=0, les_readout=mode, les_dipole=True, + les_alpha=kind) + pos, types = random_structure() + _, l0 = m(pos, types, return_embeddings=True, l0_only=True) + assert l0.shape == (len(types), 4 + width), \ + f"{mode}/{kind}: {tuple(l0.shape)}" + assert l0[:, 4:].abs().max() == 0.0, f"{mode}/{kind}: α≠0 at init" + assert l0[:, 0].abs().max() > 0, f"{mode}/{kind}: q=0 at init" + + # perturb the zero-init α slots, then check SO(3) behaviour + with torch.no_grad(): + if mode == 'edge': + m.les_edge_charge.weight[2:].normal_(std=0.5) + else: + last = m.les_edge_charge.linears[-1] + last.weight[2 * m.n_max_d:].normal_(std=0.5) + Q = rand_rotation() + _, l0a = m(pos, types, return_embeddings=True, l0_only=True) + _, l0b = m(pos @ Q.T, types, return_embeddings=True, l0_only=True) + qa, ua, aa = unpack_l0(l0a, **m.les_flags) + qb, ub, ab = unpack_l0(l0b, **m.les_flags) + assert aa.abs().max() > 0, f"{mode}/{kind}: α still 0" + dq = (qa - qb).abs().max() + du = (ua @ Q.T - ub).abs().max() + if kind == 'iso': + assert aa.shape == (len(types),) + da = (aa - ab).abs().max() + else: + assert aa.shape == (len(types), 3, 3) + sym = (aa - aa.transpose(1, 2)).abs().max() + assert sym == 0.0, f"{mode}: α not symmetric: {sym:.3e}" + da = (Q @ aa @ Q.T - ab).abs().max() + assert dq < TOL and du < TOL, f"{mode}/{kind}: q/u broke" + assert da < TOL, f"{mode}/{kind}: α not equivariant: {da:.3e}" + P = total_polarizability(aa) + assert P.shape == (1, 3, 3) + print(f" {mode}+α({kind}): packed (N,{4 + width}), α=0 at init, " + f"equivariant ({da:.1e})") + + # a single bond: α_i = a·I + b·(r̂r̂ᵀ − I/3) is uniaxial about r̂ — r̂ is + # an eigenvector and the two perpendicular eigenvalues coincide + m = make_model(seed=0, les_readout='edge_basis', les_dipole=False, + les_alpha='aniso') + with torch.no_grad(): + m.les_edge_charge.linears[-1].weight[m.n_max_d:].normal_(std=0.5) + pos2 = torch.tensor([[0.0, 0.0, 0.0], [0.7, 0.9, -0.5]], dtype=DTYPE) + types2 = torch.tensor([0, 1]) + _, l0 = m(pos2, types2, return_embeddings=True, l0_only=True) + _, _, a2 = unpack_l0(l0, **m.les_flags) + r_hat = pos2[1] - pos2[0] + r_hat = r_hat / r_hat.norm() + for i in range(2): + A = a2[i] + lam = r_hat @ A @ r_hat + d_axis = (A @ r_hat - lam * r_hat).abs().max() + assert d_axis < TOL, f"atom {i}: r̂ not an eigenvector ({d_axis:.1e})" + evals = torch.linalg.eigvalsh(A) + # two of the three eigenvalues coincide (perpendicular pair) + gaps = sorted([(evals[1] - evals[0]).abs(), (evals[2] - evals[1]).abs()]) + assert gaps[0] < 1e-9, f"atom {i}: not uniaxial: {evals.tolist()}" + assert (a2[0] - a2[0].trace() / 3 * torch.eye(3, dtype=DTYPE)).abs().max() > 0, \ + "traceless part vanished" + print(" lone bond → uniaxial α with r̂ as principal axis") + + # les_charge_scale scales [q | u], leaves α alone + ms = make_model(seed=0, les_readout='edge_basis', les_dipole=True, + les_alpha='iso', les_charge_scale=0.1) + m1 = make_model(seed=0, les_readout='edge_basis', les_dipole=True, + les_alpha='iso', les_charge_scale=1.0) + for mm in (ms, m1): + with torch.no_grad(): + mm.les_edge_charge.linears[-1].weight[2 * mm.n_max_d:].normal_( + std=0.5, generator=torch.Generator().manual_seed(4)) + mm.les_edge_charge.linears[-1].weight[mm.n_max_d:2 * mm.n_max_d]\ + .normal_(std=0.5, generator=torch.Generator().manual_seed(5)) + pos, types = random_structure() + _, l0s = ms(pos, types, return_embeddings=True, l0_only=True) + _, l01 = m1(pos, types, return_embeddings=True, l0_only=True) + dqu = (l0s[:, :4] - 0.1 * l01[:, :4]).abs().max() + dal = (l0s[:, 4:] - l01[:, 4:]).abs().max() + assert dqu < TOL and dal < TOL, f"scale: dqu={dqu:.1e}, dα={dal:.1e}" + print(" les_charge_scale: [q | u] scaled, α untouched") + + # batched slicing of the packed l0 matches per-structure forwards + structs = [random_structure(5, seed=1), random_structure(7, seed=3)] + m = make_model(seed=0, les_readout='edge_basis', les_dipole=True, + les_alpha='aniso') + with torch.no_grad(): + m.les_edge_charge.linears[-1].weight[m.n_max_d:].normal_(std=0.5) + _, l0_list = m.forward_batch_multi([p for p, _ in structs], + [t for _, t in structs], + return_embeddings=True, l0_only=True) + for b, (pos_b, types_b) in enumerate(structs): + _, l0_ref = m(pos_b, types_b, return_embeddings=True, l0_only=True) + dl = (l0_list[b] - l0_ref).abs().max() + assert dl < TOL, f"structure {b}: dl0={dl:.3e}" + + for bad in (dict(les_readout='sum', les_alpha='iso'), + dict(les_readout='edge_basis', les_alpha='full')): + try: + ECENet(**COMMON, **bad) + raise AssertionError(f"{bad} should have raised") + except ValueError as e: + assert 'les_alpha' in str(e) + print(" α read-out: batched variants consistent; bad configs rejected") + + +def test_edge_alpha_les_energy(): + """Wrapper with les_alpha: the vectorized isolated path (field of the + fixed multipoles from the masked f_qu/f_uu kernels, then −½E·α·E) equals + upstream's per-structure loop in energies and position gradients for + 'iso'/'aniso' with and without dipoles; forces pass a finite-difference + check independent of upstream; α=0 reduces exactly to the no-α energy; + the periodic path passes α through; born_charges with α goes through + upstream's induced dipoles; flag misuse is rejected.""" + if not HAVE_LES: + print(" skipped (`les` not installed)") + return + from ecenet.les import total_polarizability + lr = LESLongRange().double() + g = torch.Generator().manual_seed(12) + sizes = [4, 6] + N = sum(sizes) + pos = torch.randn(N, 3, generator=g, dtype=DTYPE) * 2.0 + q = torch.randn(N, generator=g, dtype=DTYPE) * 0.3 + u = torch.randn(N, 3, generator=g, dtype=DTYPE) * 0.2 + a_iso = torch.rand(N, generator=g, dtype=DTYPE) * 0.5 + S = torch.randn(N, 3, 3, generator=g, dtype=DTYPE) + a_ani = 0.3 * (S @ S.transpose(1, 2)) # symmetric positive + batch = torch.cat([torch.full((n,), b, dtype=torch.long) + for b, n in enumerate(sizes)]) + + worst = 0.0 + for dip in (False, True): + for kind, alpha in (('iso', a_iso), ('aniso', a_ani)): + cols = [q[:, None]] + ([u] if dip else []) + [alpha.reshape(N, -1)] + packed = torch.cat(cols, dim=1) + p_a = pos.clone().requires_grad_(True) + e = lr(packed, p_a, batch=batch, n_struct=2, l0_is_charge=True, + les_dipole=dip, les_alpha=kind) + f_a = torch.autograd.grad(e.sum(), p_a)[0] + p_b = pos.clone().requires_grad_(True) + res = lr.les(latent_charges=q, latent_dipoles=u if dip else None, + latent_alphas=alpha, positions=p_b, cell=None, + batch=batch, compute_energy=True) + f_b = torch.autograd.grad(res['E_lr'].sum(), p_b)[0] + de = (e - res['E_lr'].reshape(e.shape)).abs().max() + df = (f_a - f_b).abs().max() + assert de < 1e-10 and df < 1e-10, \ + f"{kind}/dip={dip}: != upstream loop: dE={de:.3e}, dF={df:.3e}" + worst = max(worst, float(de), float(df)) + # α=0 reduces to the no-α energy; α≠0 changes it + cols0 = cols[:-1] + [torch.zeros_like(cols[-1])] + e0 = lr(torch.cat(cols0, dim=1), pos, batch=batch, n_struct=2, + l0_is_charge=True, les_dipole=dip, les_alpha=kind) + base = (torch.cat(cols[:-1], dim=1) if dip else q[:, None]) + eb = lr(base, pos, batch=batch, n_struct=2, l0_is_charge=True, + les_dipole=dip) + d0 = (e0 - eb).abs().max() + assert d0 < 1e-12, f"{kind}/dip={dip}: α=0 ≠ no-α: {d0:.3e}" + assert (e - e0).abs().max() > 1e-3, "α changed nothing" + + # finite-difference forces (aniso + dipoles), independent of upstream + packed = torch.cat([q[:, None], u, a_ani.reshape(N, 9)], dim=1) + flags = dict(l0_is_charge=True, les_dipole=True, les_alpha='aniso') + p_a = pos.clone().requires_grad_(True) + grad = torch.autograd.grad(lr(packed, p_a, batch=batch, n_struct=2, + **flags).sum(), p_a)[0] + h = 1e-5 + max_fd = 0.0 + for (i, c) in [(0, 0), (3, 2), (7, 1)]: + es = [] + for sgn in (+1, -1): + pp = pos.clone() + pp[i, c] += sgn * h + es.append(lr(packed, pp, batch=batch, n_struct=2, **flags).sum()) + fd = (es[0] - es[1]) / (2 * h) + max_fd = max(max_fd, abs(float(fd - grad[i, c]))) + assert max_fd < 1e-7, f"FD gradient mismatch with α: {max_fd:.2e}" + + # periodic path: α passes straight through to upstream's Ewald + cell = (torch.eye(3, dtype=DTYPE) * 8.0).expand(2, 3, 3).contiguous() + e_p = lr(packed, pos, cell=cell, batch=batch, **flags) + res_p = lr.les(latent_charges=q, latent_dipoles=u, latent_alphas=a_ani, + positions=pos, cell=cell, batch=batch) + dp = (e_p - res_p['E_lr']).abs().max() + assert dp < 1e-12, f"periodic α path != upstream: {dp:.3e}" + + # born_charges: α=0 through upstream's forward equals the direct BEC + # module call (no-α path); α≠0 adds the induced dipoles and stays finite + # (latents must be functions of the positions, as they are through the + # model: upstream differentiates the polarization w.r.t. r) + n1 = sizes[0] + p1 = pos[:n1].clone().requires_grad_(True) + q1 = q[:n1] + 0.1 * p1.norm(dim=1) + u1 = u[:n1] + 0.05 * p1 + a1 = a_ani[:n1] + 0.05 * p1[:, :, None] * p1[:, None, :] + l0_1 = torch.cat([q1[:, None], u1, a1.reshape(n1, 9)], dim=1) + Z_a = lr.born_charges(l0_1, p1, cell=None, **flags) + l0_0 = torch.cat([l0_1[:, :4], torch.zeros(n1, 9, dtype=DTYPE)], dim=1) + Z_0 = lr.born_charges(l0_0, p1, cell=None, **flags) + Z_ref = lr.born_charges(l0_1[:, :4], p1, cell=None, l0_is_charge=True, + les_dipole=True) + dz0 = (Z_0 - Z_ref).abs().max() + assert Z_a.shape == (n1, 3, 3) and torch.isfinite(Z_a).all() + assert dz0 < 1e-12, f"born_charges α=0 != no-α path: {dz0:.3e}" + assert (Z_a - Z_0).abs().max() > 1e-6, "induced dipoles absent from Z*" + + # total polarizability: per-structure sums, iso → multiple of identity + P = total_polarizability(a_ani, batch=batch, n_struct=2) + assert P.shape == (2, 3, 3) + assert (P[0] - a_ani[:n1].sum(0)).abs().max() < 1e-12 + P_iso = total_polarizability(a_iso, batch=batch) + assert (P_iso[1] - a_iso[n1:].sum() * torch.eye(3, dtype=DTYPE)).abs().max() < 1e-12 + + try: + lr(packed, pos, batch=batch, n_struct=2, les_dipole=True, + les_alpha='aniso') + raise AssertionError("les_alpha without l0_is_charge should raise") + except ValueError as err: + assert 'l0_is_charge' in str(err) + print(f" vectorized α path == upstream loop (worst {worst:.1e}); FD " + f"forces {max_fd:.1e}; periodic {dp:.1e}; born_charges α=0 " + f"{dz0:.1e}; polarizability sums OK; flag misuse rejected") + + def test_edge_readout_les_energy(): """l0_is_charge=True: isolated fast path and upstream's latent_charges path agree, and the LES module holds no parameters (head bypassed).""" @@ -681,6 +920,8 @@ def test_isolated_batched_coincident_cross_atoms(): test_edge_dipole() test_dipoles_only() test_edge_dipole_les_energy() + test_edge_alpha() + test_edge_alpha_les_energy() test_edge_readout_les_energy() test_les_readout_validation() test_lazy_import() diff --git a/tests/test_trainer_les.py b/tests/test_trainer_les.py index 70a6c33..9fdf3af 100644 --- a/tests/test_trainer_les.py +++ b/tests/test_trainer_les.py @@ -4,7 +4,7 @@ Requires the optional `les` package (skips cleanly when absent). Covers: 1. rMD17 trainer: end-to-end use_les smoke on a synthetic npz ('sum' head - and parameter-free 'edge_basis'+dipole), checkpoint `les` key, and the + and parameter-free 'edge_basis'+dipole+iso-α), checkpoint `les` key, and the use_les resume-mismatch guard; 2. MPtrj trainer: end-to-end use_les smoke (periodic Ewald, stress on), resume continues, mismatch guard; @@ -73,8 +73,9 @@ def test_rmd17_use_les(): common = dict(molecule='ethanol', data_dir=tmp, n_train=8, n_val=2, n_test=2, n_epochs=2, batch_size=4, eval_every=1, dtype=DTYPE, device=DEVICE, seed=0, verbose=False, **TINY) - for ro, dip in (('sum', False), ('edge_basis', True)): + for ro, dip, alp in (('sum', False, None), ('edge_basis', True, 'iso')): _, res = train_ecenet(use_les=True, les_readout=ro, les_dipole=dip, + les_alpha=alp, checkpoint_path=ckpt if ro == 'sum' else None, **common) assert np.isfinite(res['val_force_mae']), ro @@ -91,7 +92,7 @@ def test_rmd17_use_les(): raise AssertionError("use_les mismatch not caught") except ValueError as e: assert 'use_les' in str(e) - print(" smoke (sum, edge_basis+dipole), 'les' key, resume + guard OK\n") + print(" smoke (sum, edge_basis+dipole+α), 'les' key, resume + guard OK\n") def test_mptrj_use_les(): diff --git a/tests/test_xyz_trainer.py b/tests/test_xyz_trainer.py index d4b5992..7afefbe 100644 --- a/tests/test_xyz_trainer.py +++ b/tests/test_xyz_trainer.py @@ -420,6 +420,154 @@ def pol(pos_np): f"FD {d_bec:.1e}; isolated path {dfree:.1e}; SR refused\n") +def test_les_calculator_alpha(): + """les_alpha end to end: an xyz-trained 'aniso' checkpoint loads through + load_calculator; α moved off its zero init; energy/forces equal the + manual joint graph with the induced-dipole term; stress FD holds through + the periodic α path; results['les_alphas'] is the unpacked l0 slot; + compute_polarizability is the symmetric sum Σα_i (isolated and periodic) + and equals dP/dE_ext via upstream's e_ext; compute_bec is finite and + differs from the α-less Z*; an SR/α-less checkpoint refuses it.""" + if not _has_les(): + print("=== SKIP: LES α calculator (optional `les` package not installed) ===\n") + return + print("=== ECENetLESCalculator + les_alpha: forces, stress FD, α, BEC ===") + from ase import Atoms + + from ecenet.calculator import load_calculator + from ecenet.les import total_polarizability, unpack_l0 + + with tempfile.TemporaryDirectory() as td: + ckpt = os.path.join(td, 'les_alpha.mdl') + train_ecenet_xyz( + train_structures=make_structures(8, seed=21), n_val=2, + use_les=True, les_readout='edge_basis', les_dipole=True, + les_alpha='aniso', checkpoint_path=ckpt, n_epochs=2, + batch_size=4, lr=5e-3, **COMMON) + calc = load_calculator(ckpt, device='cpu', verbose=False) + assert calc.les_flags['les_alpha'] == 'aniso' + flags = calc.les_flags + + s = make_structures(1, seed=22, box=(8.5, 9.0))[0] + atoms = Atoms(numbers=s['numbers'], positions=s['positions'], + cell=s['cell'], pbc=True) + atoms.calc = calc + e_calc = atoms.get_potential_energy() + f_calc = atoms.get_forces() + na = len(atoms) + + # manual joint graph with the same weights/topology + types = torch.tensor( + [calc.element_to_type[sym] for sym in atoms.get_chemical_symbols()], + dtype=torch.long) + pos = torch.tensor(s['positions'], dtype=DTYPE).requires_grad_(True) + ei, ej, she = calc._gpu_neighbor_list(pos.detach(), s['cell'], + calc.model.r_cut_edge) + ni, nj, shn = calc._gpu_neighbor_list(pos.detach(), s['cell'], + calc.model.r_cut_neighbor) + e_sr, l0 = calc.model.forward_pbc(pos, types, ei, ej, she, ni, nj, shn, + return_embeddings=True, l0_only=True) + cell_t = torch.tensor(s['cell'], dtype=DTYPE) + e_man = e_sr + calc.les_module(l0, pos, cell=cell_t, **flags).sum() + f_man = -torch.autograd.grad(e_man, pos)[0].numpy() + e_ref_sum = sum(calc.energy_reference[sym] + for sym in atoms.get_chemical_symbols()) + de = abs(e_calc - (e_man.item() + e_ref_sum)) + df = np.abs(f_calc - f_man).max() + assert de < 1e-10 and df < 1e-10, f"joint graph: dE={de:.1e} dF={df:.1e}" + + # α trained off its zero init, and it changes E_lr + q_l, u_l, a_l = unpack_l0(l0.detach(), **flags) + assert a_l.shape == (na, 3, 3) and a_l.abs().max() > 0, "α stayed 0" + l0_noa = l0.detach()[:, :4] + e_lr_a = calc.les_module(l0.detach(), pos.detach(), cell=cell_t, + **flags).sum() + e_lr_0 = calc.les_module(l0_noa, pos.detach(), cell=cell_t, + l0_is_charge=True, les_dipole=True).sum() + assert abs(float(e_lr_a - e_lr_0)) > 0, "α term contributes nothing" + + # stash: α slot exposed on every force call + da = np.abs(calc.results['les_alphas'] - a_l.numpy()).max() + du = np.abs(calc.results['les_dipoles'] - u_l.numpy()).max() + assert da < 1e-12 and du < 1e-12, f"stash: dα={da:.1e} du={du:.1e}" + + # stress FD through the periodic α path (one normal, one shear) + stress_v = atoms.get_stress() + V = abs(np.linalg.det(s['cell'])) + eps = 1e-6 + max_err = 0.0 + for (a, b), vi in [((1, 1), 1), ((0, 2), 4)]: + E = np.zeros((3, 3)); E[a, b] = eps + es = [] + for sign in (+1, -1): + F = np.eye(3) + sign * E + at = Atoms(numbers=s['numbers'], positions=s['positions'] @ F, + cell=s['cell'] @ F, pbc=True) + at.calc = calc + es.append(at.get_potential_energy()) + fd = (es[0] - es[1]) / (2 * eps) / V + max_err = max(max_err, abs(fd - stress_v[vi])) + assert max_err < 1e-7, f"stress FD with α: {max_err:.2e}" + + # polarizability: Σα_i, symmetric, same on the isolated path, and + # equal to dP/dE_ext through upstream (the induced dipoles' response + # to a uniform external field IS Σα for this non-self-consistent model) + P = calc.compute_polarizability(atoms) + assert P.shape == (3, 3) + dP = np.abs(P - total_polarizability(a_l)[0].numpy()).max() + assert dP < 1e-12 and np.abs(P - P.T).max() < 1e-12 + atoms_free = Atoms(numbers=s['numbers'], positions=s['positions']) + atoms_free.calc = calc + P_free = calc.compute_polarizability(atoms_free) + pos_f = torch.tensor(s['positions'], dtype=DTYPE) + with torch.no_grad(): + _, l0_f = calc.model(pos_f, types, return_embeddings=True, + l0_only=True) + _, _, a_f = unpack_l0(l0_f, **flags) + dPf = np.abs(P_free - a_f.sum(0).numpy()).max() + assert dPf < 1e-12, f"isolated polarizability: {dPf:.1e}" + e_ext = torch.zeros(3, dtype=DTYPE, requires_grad=True) + q_f, u_f, _ = unpack_l0(l0_f, **flags) + res = calc.les_module.les(latent_charges=q_f, latent_dipoles=u_f, + latent_alphas=a_f, positions=pos_f, + cell=None, e_ext=e_ext, compute_energy=True) + P_tot = res['latent_dipoles'].reshape(-1, 3).sum(0) # Σ(u + Δu) + dPdE = torch.stack([torch.autograd.grad(P_tot[c], e_ext, + retain_graph=True)[0] + for c in range(3)]).numpy() + dE = np.abs(dPdE - P_free).max() + assert dE < 1e-10, f"Σα != dP/dE_ext: {dE:.1e}" + + # BEC with the induced dipoles: finite, periodic and isolated, and + # not the α-less tensor + Z = calc.compute_bec(atoms) + Z_free = calc.compute_bec(atoms_free) + assert Z.shape == (na, 3, 3) and np.isfinite(Z).all() + assert Z_free.shape == (na, 3, 3) and np.isfinite(Z_free).all() + pos_g = torch.tensor(s['positions'], dtype=DTYPE, requires_grad=True) + _, l0_g = calc.model(pos_g, types, return_embeddings=True, l0_only=True) + Z_noa = calc.les_module.born_charges(l0_g[:, :4], pos_g, cell=None, + l0_is_charge=True, les_dipole=True) + dz = np.abs(Z_free - Z_noa.detach().numpy()).max() + assert dz > 1e-9, "induced dipoles missing from Z*" + + # an α-less checkpoint refuses compute_polarizability + ckpt0 = os.path.join(td, 'les_noalpha.mdl') + train_ecenet_xyz( + train_structures=make_structures(6, seed=23), n_val=2, + use_les=True, les_readout='edge_basis', checkpoint_path=ckpt0, + n_epochs=1, batch_size=4, lr=5e-3, **COMMON) + calc0 = load_calculator(ckpt0, device='cpu', verbose=False) + try: + calc0.compute_polarizability(atoms_free) + raise AssertionError("α-less checkpoint should refuse") + except ValueError as e: + assert 'les_alpha' in str(e) + print(f" joint graph (dE={de:.1e}, dF={df:.1e}); stress FD {max_err:.1e}; " + f"α stashed ({da:.1e}); Σα == dP/dE_ext ({dE:.1e}); BEC finite, " + f"induced part {dz:.1e}; α-less refused\n") + + def test_tensorize_keeps_cell(): print("=== tensorize: cell kept for periodic, None otherwise ===") structs = make_structures(2, seed=5) @@ -441,4 +589,5 @@ def test_tensorize_keeps_cell(): test_smoke_train_les() test_calculator_rejects_les_checkpoint() test_les_calculator() + test_les_calculator_alpha() print("ALL TESTS PASSED") diff --git a/tools/eval_spice_bec.py b/tools/eval_spice_bec.py index c389e6c..c17dfa7 100644 --- a/tools/eval_spice_bec.py +++ b/tools/eval_spice_bec.py @@ -58,8 +58,6 @@ def predict_becs(checkpoint_path, xyz_files, device='cpu', max_frames=None): device = torch.device(device) model, les_module, hp, elem_to_type, dtype = load_les_model( checkpoint_path, device) - is_charge = model.les_flags['l0_is_charge'] - les_dip = model.les_dipole records, skipped = [], {} for path in xyz_files: @@ -88,17 +86,10 @@ def predict_becs(checkpoint_path, xyz_files, device='cpu', max_frames=None): requires_grad=True) _, l0 = model.forward_pbc(pos, types, ei, ej, she, ni, nj, shn, return_embeddings=True, l0_only=True) - if is_charge: - q = l0[:, 0] - u = l0[:, 1:4] if les_dip else None - else: - q = les_module.les.atomwise( - l0.reshape(l0.shape[0], -1), - torch.zeros(len(symbols), dtype=torch.long, device=device)) - u = None - bec = les_module.les.bec(q=q, r=pos, cell=None, u=u) - if bec.dim() == 4: # (N, 2, 3, 3): charge + dipole parts - bec = bec.sum(dim=1) + # same implementation as ECENetLESCalculator.compute_bec (with + # les_alpha the induced dipoles are part of the polarization) + bec = les_module.born_charges(l0, pos, cell=None, + **model.les_flags) records.append({ 'subset': subset, 'frame': fi, 'symbols': np.array(symbols), 'bec_pred': bec.detach().cpu().numpy(), diff --git a/tools/predict_charges.py b/tools/predict_charges.py index 432d089..310a822 100644 --- a/tools/predict_charges.py +++ b/tools/predict_charges.py @@ -103,8 +103,13 @@ def _predict_frame(model, les_module, hp, elem_to_type, dtype, atoms, device): 'charges': q.cpu().numpy().reshape(-1), 'e_lr': float(e_lr.sum()), } - if model.les_dipole: - out['dipoles'] = l0[:, 1:4].cpu().numpy() + if model.les_dipole or model.les_alpha: + from ecenet.les import unpack_l0 + _, u, alpha = unpack_l0(l0, **model.les_flags) + if u is not None: + out['dipoles'] = u.cpu().numpy() + if alpha is not None: # (N,) 'iso' or (N, 3, 3) 'aniso' + out['alphas'] = alpha.cpu().numpy() if 'q' in atoms.arrays: out['charges_ref'] = np.asarray(atoms.arrays['q'], dtype=np.float64) return out