diff --git a/ecenet/model.py b/ecenet/model.py index 578af41..bad19bd 100644 --- a/ecenet/model.py +++ b/ecenet/model.py @@ -39,15 +39,7 @@ from ecenet.equivariant import EquivariantLinear, RealSpaceNonlinearity from ecenet.film import ElementFiLM from ecenet.radial import find_edges, get_cutoff_fn, radial_basis -from ecenet.spherical import ( - RotationWeights, - analytic_rotation_weights, - build_D1_from_rhat, - build_D_block, - build_D_slice, - kept_offsets, - spherical_harmonics_float64, -) +from ecenet.spherical import build_D1_from_rhat, build_D_block, spherical_harmonics_float64, wigner_rotate # les_readout modes whose l0 IS the latent charge itself (upstream's atomwise # head bypassed via l0_is_charge). The single definition — consumers read the @@ -77,15 +69,7 @@ class ECENet(nn.Module): (n_mp_steps=K-1, n_final_layers=n_layers) layout. n_max_d: if set, outer-product the invariants with f_d(r_ij) of this rank m_max: max angular mode |m| kept after the equivariant layers - (default: l_max); lower it to cut cost at large l_max. - With m_max < l_max the eager frame changes use only the - kept COLUMNS of each Wigner-D block (an exact - rectangular slice — see spherical.build_D_slice), so - the truncation also shrinks the rotation, not just the - layer stack. For m_max <= 1, set_analytic_wigner(True) - additionally builds those columns in closed form from - one sphericart call (single-backward only — MD, not - force-loss training). + (default: l_max); lower it to cut cost at large l_max cutoff_type: 'cosine' or 'poly' activation: pointwise activation in the realspace nonlinearity ('silu', 'tanh', 'relu', 'gelu'); 'identity' turns every RealSpaceNonlinearity — the @@ -794,81 +778,25 @@ def set_edge_frame_fused(self, enabled: bool = True, e2n: bool = True): double-backward force-loss training. Numerically identical to the unfused ops (see tests/test_edge_frame_kernel.py). Returns self. """ - if enabled and getattr(self, '_analytic_wigner', False): - raise ValueError( - "edge_frame_fused and analytic Wigner-D are mutually " - "exclusive: the fused kernels index the full D block. " - "set_analytic_wigner(False) first.") self._edge_frame_fused = enabled for layer in getattr(self, 'mp_layers', []): layer.edge_frame_fused = enabled layer.edge_frame_fused_e2n = e2n return self - def set_analytic_wigner(self, enabled: bool = True): - """Toggle the analytic, MATRIX-FREE Wigner rotation (m_max <= 1 - only): no D tensor is built at all — the rotation is applied as - three weight vectors from one sphericart call (Y for m=0, - tangential gradients along the gauge frame for m=±1; see - spherical.analytic_rotation_weights), used by the edge frame and - both MP frame crossings. SINGLE-BACKWARD - ONLY: forces for MD/inference and the stress strain pass are fine, - but force-loss training (double backward) raises — like - set_activation_fused, leave it off for training. Returns self.""" - if enabled and self.m_max > 1: - raise ValueError( - f"analytic Wigner-D needs m_max <= 1 (closed forms for " - f"|m| <= 1), got m_max={self.m_max}") - if enabled and getattr(self, '_edge_frame_fused', False): - raise ValueError( - "edge_frame_fused and analytic Wigner-D are mutually " - "exclusive: the fused kernels index the full D block. " - "set_edge_frame_fused(False) first.") - self._analytic_wigner = enabled - return self - - def _d_slice_active(self): - """Whether the eager frame changes use the kept-column D slice. - - Active when the truncation actually drops columns (m_max < l_max) or - the analytic build is requested, and the fused edge-frame path is off - (the Triton kernels index the full block — they already read only the - kept entries, so they get the application-side saving for free).""" - return (getattr(self, '_use_d_slice', True) - and (self.m_max < self.l_max - or getattr(self, '_analytic_wigner', False)) - and not getattr(self, '_edge_frame_fused', False)) - - def _build_D(self, r_hat): - """Wigner-D for this forward: kept-column slice when the truncated - path is active (analytic when requested), full block otherwise.""" - if self._d_slice_active(): - if getattr(self, '_analytic_wigner', False): - # matrix-free: the rotation as three weight vectors — see - # spherical.analytic_rotation_weights (single-backward) - return analytic_rotation_weights(r_hat, self.l_max, self.m_max) - return build_D_slice(r_hat, self.l_max, self.m_max) - return build_D_block(r_hat, self.l_max) - def _edge_frame(self, A_emb, edge_i, edge_j, r_hat): """Steps 3-4: gather endpoint features, rotate into the bond frame, - reshape to (A_cos, A_sin). Returns (A_cos, A_sin, D); D (full block, - or the kept-column slice when m_max < l_max) is built here so callers - can reuse it (MP layers, node aggregation).""" - D = self._build_D(r_hat) - if isinstance(D, RotationWeights): - A_both = torch.cat([A_emb[edge_i], A_emb[edge_j]], dim=1) - A_cos, A_sin = _matrix_free_to_angular( - A_both, D, self.l_max, self.sph_to_angular) - elif getattr(self, '_edge_frame_fused', False): - A_cos, A_sin = edge_frame_fused(A_emb, edge_i, edge_j, D, + reshape to (A_cos, A_sin). Returns (A_cos, A_sin, D_block); D_block is + built here so callers can reuse it (MP layers, node aggregation).""" + D_block = build_D_block(r_hat, self.l_max) + if getattr(self, '_edge_frame_fused', False): + A_cos, A_sin = edge_frame_fused(A_emb, edge_i, edge_j, D_block, self.sph_to_angular) else: A_both = torch.cat([A_emb[edge_i], A_emb[edge_j]], dim=1) - A_rot = torch.bmm(A_both, D) - A_cos, A_sin = self.sph_to_angular(A_rot, - sliced=D.shape[-1] != self.n_sph) - return A_cos, A_sin, D + A_rot = wigner_rotate(A_both, D_block) + A_cos, A_sin = self.sph_to_angular(A_rot) + return A_cos, A_sin, D_block def set_activation_fused(self, enabled: bool = True): """Toggle the fused recompute-in-backward path on all RealSpaceNonlinearity @@ -1327,15 +1255,11 @@ def forward_batch(self, positions_list, types, topology=None, r_hat_flat = r_hat.reshape(B * n_edges, 3) A_both_flat = A_both.reshape(B * n_edges, 2 * self.embed_dim, self.n_sph) - D_block = self._build_D(r_hat_flat) - if isinstance(D_block, RotationWeights): - A_cos_flat, A_sin_flat = _matrix_free_to_angular( - A_both_flat, D_block, self.l_max, self.sph_to_angular) - else: - A_rot_flat = torch.bmm(A_both_flat, D_block) - # ── Step 4: Reshape to A_cos / A_sin ───────────────────────── - A_cos_flat, A_sin_flat = self.sph_to_angular( - A_rot_flat, sliced=D_block.shape[-1] != self.n_sph) + D_block = build_D_block(r_hat_flat, self.l_max) + A_rot_flat = wigner_rotate(A_both_flat, D_block) + + # ── Step 4: Reshape to A_cos / A_sin ───────────────────────────── + A_cos_flat, A_sin_flat = self.sph_to_angular(A_rot_flat) # shapes: (B*n_edges, n_features_per_m, n_angular) # ── Step 5: Equivariant layers ──────────────────────────────────── @@ -1510,117 +1434,6 @@ def _unpack_sph_to_angular(v_rot, n_base, l_max, m_max, n_angular, n_sph): return d_cos, d_sin -@functools.lru_cache(maxsize=None) -def _sph_pack_index_sliced(l_max, m_max, n_angular, device): - """Kept-column-layout twin of ``_sph_pack_index``: the packed features are - zero outside |m| <= min(l, m_max), so pack/unpack against a D *slice* - (spherical.build_D_slice) target the narrow kept layout instead of the full - SH grid — same maps with the block base l²+l replaced by - kept_offset(l) + min(l, m_max). Returns the same 8-tuple, widths n_kept.""" - lp1 = l_max + 1 - off = kept_offsets(l_max, m_max) - n_kept = off[-1] - pack_c = torch.zeros(n_kept, dtype=torch.long) - pack_s = torch.zeros(n_kept, dtype=torch.long) - pack_cm = torch.zeros(n_kept, dtype=torch.bool) - pack_sm = torch.zeros(n_kept, dtype=torch.bool) - up_c = torch.zeros(lp1 * n_angular, dtype=torch.long) - up_s = torch.zeros(lp1 * n_angular, dtype=torch.long) - up_cm = torch.zeros(lp1 * n_angular, dtype=torch.bool) - up_sm = torch.zeros(lp1 * n_angular, dtype=torch.bool) - for l in range(lp1): - m_out = min(l, m_max) - base_k = off[l] + m_out - for m in range(0, m_out + 1): - s = base_k + m - pack_c[s] = l * n_angular + m - pack_cm[s] = True - up_c[l * n_angular + m] = s - up_cm[l * n_angular + m] = True - for k in range(1, m_out + 1): - s = base_k - k - pack_s[s] = l * n_angular + k - pack_sm[s] = True - up_s[l * n_angular + k] = s - up_sm[l * n_angular + k] = True - dev = torch.device(device) - return tuple(t.to(dev) for t in (pack_c, pack_cm, pack_s, pack_sm, - up_c, up_cm, up_s, up_sm)) - - -def _pack_angular_to_kept(A_cos, A_sin, n_base, l_max, m_max, n_angular): - """(n_e, n_base*lp1, n_angular) → kept layout (n_e, n_base, n_kept).""" - n_e = A_cos.shape[0] - lp1 = l_max + 1 - pc, pcm, ps, psm, *_ = _sph_pack_index_sliced(l_max, m_max, n_angular, - A_cos.device) - n_kept = pc.shape[0] - ac = A_cos.reshape(n_e, n_base, lp1 * n_angular) - asn = A_sin.reshape(n_e, n_base, lp1 * n_angular) - ic = pc.view(1, 1, n_kept).expand(n_e, n_base, n_kept) - isn = ps.view(1, 1, n_kept).expand(n_e, n_base, n_kept) - return (ac.gather(2, ic) * pcm.to(A_cos.dtype) - + asn.gather(2, isn) * psm.to(A_cos.dtype)) - - -def _unpack_kept_to_angular(v_rot, n_base, l_max, m_max, n_angular): - """Kept layout (n_e, n_base, n_kept) → cos/sin (n_e, n_base, lp1, n_angular).""" - n_e = v_rot.shape[0] - lp1 = l_max + 1 - _, _, _, _, uc, ucm, us, usm = _sph_pack_index_sliced(l_max, m_max, - n_angular, v_rot.device) - L = lp1 * n_angular - ic = uc.view(1, 1, L).expand(n_e, n_base, L) - isn = us.view(1, 1, L).expand(n_e, n_base, L) - d_cos = (v_rot.gather(2, ic) * ucm.to(v_rot.dtype)).view(n_e, n_base, lp1, n_angular) - d_sin = (v_rot.gather(2, isn) * usm.to(v_rot.dtype)).view(n_e, n_base, lp1, n_angular) - return d_cos, d_sin - - -@functools.lru_cache(maxsize=None) -def _l_index_of_s(l_max, device): - """(n_sph,) map from flat SH index to its degree l.""" - return torch.cat([torch.full((2 * l + 1,), l, dtype=torch.long) - for l in range(l_max + 1)]).to(torch.device(device)) - - -def _matrix_free_segments(X, W, l_max): - """Per-l weighted segment sums: X (E, C, n_sph) ⊙ each weight row of - W.w (E, k, n_sph), summed within each degree → (E, C, k, l_max+1). - - This IS the bond-frame rotation for m_max <= 1 (the kept components are - the channel functions' values / tangential derivatives at r̂); a small - k-loop keeps the transient at one X-sized tensor per row.""" - E, C, S = X.shape - l_idx = _l_index_of_s(l_max, X.device) - outs = [] - for k in range(W.w.shape[1]): - P = X * W.w[:, k].unsqueeze(1) # (E, C, S) - o = torch.zeros(E, C, l_max + 1, dtype=X.dtype, device=X.device) - outs.append(o.index_add_(2, l_idx, P)) - return torch.stack(outs, dim=2) # (E, C, k, l+1) - - -def _matrix_free_to_angular(A_both, W, l_max, sph_to_angular): - """Matrix-free rotate-in straight to the angular layout: (E, Cb, n_sph) → - A_cos/A_sin (E, Cb·(l_max+1), n_angular). Row order of W.w is - [W0 → cos m=0, W+ → cos m=1, W- → sin m=1]; the l=0 rows of W± are - exactly zero (b_0 = 0), so the structural slots stay zero — the valid - masks are applied anyway for bit-parity with the sliced path.""" - E, Cb, _ = A_both.shape - seg = _matrix_free_segments(A_both, W, l_max) # (E, Cb, k, l+1) - n_ch = Cb * (l_max + 1) - if W.w.shape[1] == 1: # m_max = 0 - A_cos = seg[:, :, 0].reshape(E, n_ch, 1) - A_sin = torch.zeros_like(A_cos) - else: - A_cos = torch.stack([seg[:, :, 0], seg[:, :, 1]], - dim=-1).reshape(E, n_ch, 2) - A_sin = torch.stack([torch.zeros_like(seg[:, :, 2]), seg[:, :, 2]], - dim=-1).reshape(E, n_ch, 2) - return A_cos * sph_to_angular.cos_valid, A_sin * sph_to_angular.sin_valid - - class ECENetAttentionMPLayer(nn.Module): """Attention-style message passing for ECENet. @@ -1806,36 +1619,9 @@ def forward(self, A_cos, A_sin, r_hat, dist_ij, edge_i, edge_j, # materializes); the weighting below acts on h_global either way — the # gate is applied in the node frame, so 'sum'/'softmax'/l_attention all # compose with the fusion unchanged. - mfree = isinstance(D_block, RotationWeights) - d_sliced = (not mfree) and D_block.shape[-1] != self.n_sph if getattr(self, 'edge_frame_fused_e2n', False): - assert not (d_sliced or mfree), "fused e2n needs the full D block" h_global = pack_unrotate_fused(m_cos, m_sin, D_block, self.l_max, self.m_max) - elif mfree: - # transpose of the matrix-free rotate: broadcast each kept - # component over its l-block and sum the weighted copies. The - # l=0 rows of W± are zero, so any structural-scratch content in - # those input slots is annihilated exactly. - lp1 = self.l_max + 1 - l_idx = _l_index_of_s(self.l_max, m_cos.device) - mc = m_cos.view(n_e, self.n_base, lp1, self.n_angular) - h_global = (mc[..., 0].index_select(2, l_idx) - * D_block.w[:, 0].unsqueeze(1)) - if D_block.w.shape[1] == 3: - ms = m_sin.view(n_e, self.n_base, lp1, self.n_angular) - h_global = (h_global - + mc[..., 1].index_select(2, l_idx) - * D_block.w[:, 1].unsqueeze(1) - + ms[..., 1].index_select(2, l_idx) - * D_block.w[:, 2].unsqueeze(1)) - elif d_sliced: - # packed features are zero outside the kept |m| <= min(l, m_max) - # slots, so packing into the kept layout and multiplying the D - # slice's transpose gives the identical full-frame h_global - h_k = _pack_angular_to_kept(m_cos, m_sin, self.n_base, self.l_max, - self.m_max, self.n_angular) - h_global = torch.bmm(h_k, D_block.transpose(-1, -2)) else: h = self._pack(m_cos, m_sin) h_global = torch.bmm(h, D_block.transpose(-1, -2)) # transposed view, no copy @@ -1894,27 +1680,8 @@ def forward(self, A_cos, A_sin, r_hat, dist_ij, edge_i, edge_j, # 6. Gather to edges (source atom), rotate back to the edge frame. if getattr(self, 'edge_frame_fused', False): - assert not (d_sliced or mfree), \ - "fused rotate-back needs the full D block" d_cos, d_sin = edge_frame_fused_single( Delta, edge_i, D_block, self.l_max, self.m_max) - elif mfree: - seg = _matrix_free_segments(Delta[edge_i], D_block, self.l_max) - if D_block.w.shape[1] == 1: - d_cos = seg[:, :, 0].reshape(n_e, self.n_ch, 1) - d_sin = torch.zeros_like(d_cos) - else: - d_cos = torch.stack([seg[:, :, 0], seg[:, :, 1]], - dim=-1).reshape(n_e, self.n_ch, 2) - d_sin = torch.stack( - [torch.zeros_like(seg[:, :, 2]), seg[:, :, 2]], - dim=-1).reshape(n_e, self.n_ch, 2) - elif d_sliced: - v_rot = torch.bmm(Delta[edge_i], D_block) # (n_e, n_base, n_kept) - dk_cos, dk_sin = _unpack_kept_to_angular( - v_rot, self.n_base, self.l_max, self.m_max, self.n_angular) - d_cos = dk_cos.reshape(n_e, self.n_ch, self.n_angular) - d_sin = dk_sin.reshape(n_e, self.n_ch, self.n_angular) else: v_rot = torch.bmm(Delta[edge_i], D_block) # (n_e, n_base, n_sph) d_cos, d_sin = self._unpack(v_rot, n_e) @@ -1996,36 +1763,10 @@ def __init__(self, embed_dim: int, l_max: int, m_max: int = None): self.register_buffer('cos_flat_idx', cos_flat.reshape(-1), persistent=False) self.register_buffer('sin_flat_idx', sin_flat.reshape(-1), persistent=False) - # Kept-column layout (m_max < l_max): the rotate only produces the - # columns |m| <= min(l, m_max) (see spherical.build_D_slice), so a - # second index set gathers from that narrower layout — same maps with - # the block base l²+l replaced by kept_offset(l) + min(l, m_max). - off = kept_offsets(l_max, m_max) - self.n_kept = off[-1] - cos_idx_k = torch.zeros(self.n_ch, self.n_angular, dtype=torch.long) - sin_idx_k = torch.zeros(self.n_ch, self.n_angular, dtype=torch.long) - c = 0 - for _ in range(n_ch_base): - for l in range(l_max + 1): - base_k = off[l] + min(l, m_max) - for m in range(self.n_angular): - if m <= l: - cos_idx_k[c, m] = base_k + m - if m > 0: - sin_idx_k[c, m] = base_k - m - c += 1 - cos_flat_k = ch_src[:, None] * self.n_kept + cos_idx_k - sin_flat_k = ch_src[:, None] * self.n_kept + sin_idx_k - self.register_buffer('cos_flat_idx_k', cos_flat_k.reshape(-1), - persistent=False) - self.register_buffer('sin_flat_idx_k', sin_flat_k.reshape(-1), - persistent=False) - - def forward(self, A_rot, sliced=False): + def forward(self, A_rot): """ Args: - A_rot: (n_edges, 2*embed_dim, n_sph) — or, with ``sliced=True``, - (n_edges, 2*embed_dim, n_kept) from a D-slice rotation + A_rot: (n_edges, 2*embed_dim, n_sph) Returns: A_cos, A_sin: (n_edges, 2*embed_dim*(l_max+1), l_max+1) """ @@ -2033,12 +1774,10 @@ def forward(self, A_rot, sliced=False): # Gather +m (cos) and −m (sin) components straight from A_rot. The flat # indices already encode the (embed, l) channel repeat, so there is no # (l_max+1)×-larger A_exp intermediate (and no grad buffer for it). - ci = self.cos_flat_idx_k if sliced else self.cos_flat_idx - si = self.sin_flat_idx_k if sliced else self.sin_flat_idx - A_flat = A_rot.reshape(n_edges, -1) # (n_edges, n_ch_base * width) - A_cos = (A_flat.index_select(1, ci) + A_flat = A_rot.reshape(n_edges, -1) # (n_edges, n_ch_base * n_sph) + A_cos = (A_flat.index_select(1, self.cos_flat_idx) .view(n_edges, self.n_ch, self.n_angular)) * self.cos_valid - A_sin = (A_flat.index_select(1, si) + A_sin = (A_flat.index_select(1, self.sin_flat_idx) .view(n_edges, self.n_ch, self.n_angular)) * self.sin_valid return A_cos, A_sin diff --git a/ecenet/spherical.py b/ecenet/spherical.py index 9b1ad46..bf259ca 100644 --- a/ecenet/spherical.py +++ b/ecenet/spherical.py @@ -270,194 +270,3 @@ def wigner_rotate(A_flat: torch.Tensor, D_block: torch.Tensor) -> torch.Tensor: (N, C, n_sph) rotated features. """ return torch.bmm(A_flat, D_block) - - -# ───────────────────────────────────────────────────────────────────────────── -# Truncated (m_max < l_max) Wigner-D: only the columns the model consumes -# ───────────────────────────────────────────────────────────────────────────── -# -# With the angular layout truncated at m_max, the layers only ever read the -# bond-frame components with |m| <= min(l, m_max) — and, in the reverse -# direction, the packed features are exactly zero outside those slots. So both -# the rotate (A @ D) and the unrotate (h @ D^T) consume only the corresponding -# COLUMNS of each D^l block: a rectangular (n_sph, n_kept) slice of the -# block-diagonal, n_kept = Σ_l (2·min(l, m_max)+1). The kept layout keeps the -# per-l blocks contiguous in the SH ordering (m = -mk..mk), so index maps -# mirror the full-layout ones with the block base l²+l replaced by -# kept_offset(l) + min(l, m_max). - - -def n_kept_columns(l_max: int, m_max: int) -> int: - """Width of the kept-column layout: Σ_l (2·min(l, m_max)+1).""" - return sum(2 * min(l, m_max) + 1 for l in range(l_max + 1)) - - -def kept_offsets(l_max: int, m_max: int): - """Start of each l's block in the kept layout (list of l_max+2 ints; - the last entry is n_kept).""" - off = [0] - for l in range(l_max + 1): - off.append(off[-1] + 2 * min(l, m_max) + 1) - return off - - -def build_D_slice(r_hat: torch.Tensor, l_max: int, m_max: int) -> torch.Tensor: - """Kept columns of the block-diagonal Wigner-D: (N, n_sph, n_kept). - - ``A @ build_D_slice(...)`` equals the kept columns of - ``A @ build_D_block(...)`` exactly (same recursion, column slice) — pure - differentiable torch, safe for double-backward force-loss training. - """ - D_list = recursive_wigner_D(r_hat, l_max) - N = r_hat.shape[0] - n_sph = (l_max + 1) ** 2 - blocks = [] - for l, Dl in enumerate(D_list): - mk = min(l, m_max) - blk = torch.zeros(N, n_sph, 2 * mk + 1, - dtype=r_hat.dtype, device=r_hat.device) - blk[:, l * l:(l + 1) * (l + 1), :] = Dl[:, :, l - mk:l + mk + 1] - blocks.append(blk) - return torch.cat(blocks, dim=2) - - -class _SphYGrad(torch.autograd.Function): - """Y and its Cartesian gradient as *values*, with a first-order backward - via sphericart's Hessians. once_differentiable: a second backward (force- - loss training) raises rather than silently returning wrong gradients — - this is the single-backward leg of the analytic D-slice below.""" - - @staticmethod - def forward(ctx, r_hat, l_max): - sph = _get_sph_cache(l_max) - Y, dY = sph.compute_with_gradients(r_hat.detach().contiguous()) - ctx.save_for_backward(r_hat) - ctx.l_max = l_max - return Y, dY - - @staticmethod - def backward(ctx, gY, gdY): - # Grad mode is enabled inside backward only under create_graph — i.e. - # when someone is building a force graph to differentiate again - # (force-loss training). once_differentiable would be SILENT here: - # the analytic D mixes this Function with differentiable ops (the D1 - # block), so a double backward would quietly drop this branch's - # second derivative instead of erroring. Raise loudly instead. - if torch.is_grad_enabled(): - raise RuntimeError( - "analytic Wigner-D is single-backward only (its gradient " - "runs through sphericart Hessians; no third derivatives). " - "For force-loss training use the recursion path: " - "model.set_analytic_wigner(False).") - (r_hat,) = ctx.saved_tensors - sph = _get_sph_cache(ctx.l_max) - _, dY, H = sph.compute_with_hessians(r_hat.detach().contiguous()) - grad = torch.einsum('ns,ncs->nc', gY, dY) - grad = grad + torch.einsum('nds,ncds->nc', gdY, H) - return grad, None - - -def build_D_slice_analytic(r_hat: torch.Tensor, l_max: int, - m_max: int) -> torch.Tensor: - """Analytic kept columns for m_max <= 1 — no CG recursion. - - Closed forms (discovered numerically against ``recursive_wigner_D`` and - verified to ~1e-15 in tests/test_wigner_slice.py; the code's gauge is the - two-chart frame of ``build_D1_from_rhat``, whose ±1 columns are exactly - the tangent frame e_±): - - D^l[:, m=0] = sqrt(4π/(2l+1)) · Y_l(r̂) - D^l[:, m=±1] = sqrt(4π/(2l+1)) · sqrt(2/(l(l+1))) · (∇Y_l(r̂) · e_±) - - One sphericart call replaces the recursion. SINGLE-BACKWARD ONLY (the - gradient path runs through sphericart Hessians): fine for MD / inference - forces and the stress strain pass, but a force-loss training step (double - backward) raises. Use ``build_D_slice`` for training. - """ - if m_max > 1: - raise ValueError(f"build_D_slice_analytic supports m_max <= 1 " - f"(closed forms for |m| <= 1), got m_max={m_max}") - N = r_hat.shape[0] - n_sph = (l_max + 1) ** 2 - Y, dY = _SphYGrad.apply(r_hat, l_max) - if m_max >= 1 and l_max >= 1: - D1 = build_D1_from_rhat(r_hat) # differentiable; defines gauge - # frame vectors in Cartesian from the (y, z, x) real-SH basis rows - e_p = torch.stack([D1[:, 2, 2], D1[:, 0, 2], D1[:, 1, 2]], dim=1) - e_m = torch.stack([D1[:, 2, 0], D1[:, 0, 0], D1[:, 1, 0]], dim=1) - g_p = torch.einsum('nc,ncs->ns', e_p, dY) - g_m = torch.einsum('nc,ncs->ns', e_m, dY) - blocks = [] - for l in range(l_max + 1): - s, e = l * l, (l + 1) * (l + 1) - a = math.sqrt(4.0 * math.pi / (2 * l + 1)) - col0 = a * Y[:, s:e] - if l == 0 or m_max == 0: - cols = col0.unsqueeze(2) - elif l == 1: - cols = D1 # exact, differentiable, cheap - else: - b = a * math.sqrt(2.0 / (l * (l + 1))) - cols = torch.stack([b * g_m[:, s:e], col0, b * g_p[:, s:e]], dim=2) - blk = torch.zeros(N, n_sph, cols.shape[2], - dtype=r_hat.dtype, device=r_hat.device) - blk[:, s:e, :] = cols - blocks.append(blk) - return torch.cat(blocks, dim=2) - - -class RotationWeights: - """Matrix-free bond-frame rotation for m_max <= 1. - - Wraps the (N, k, n_sph) weight tensor of ``analytic_rotation_weights`` - (k = 1 for m_max=0, 3 for m_max=1; rows [W0, W+, W-]). A distinct type — - not a bare tensor — so consumers (the model's edge frame, the MP layers' - frame crossings) can dispatch on it unambiguously instead of sniffing - shapes that can collide with D-slice widths. - """ - __slots__ = ('w',) - - def __init__(self, w): - self.w = w - - -def analytic_rotation_weights(r_hat, l_max: int, m_max: int) -> RotationWeights: - """The m_max <= 1 bond-frame rotation as three weight VECTORS — no matrix. - - Contracting the analytic D columns with the features shows the rotation is - a point evaluation: per channel and degree l, the m=0 output is the - channel's spherical function evaluated at r̂ and the m=±1 outputs are its - two tangential derivatives there. So all any consumer needs is - - W0[s] = a_{l(s)} · Y[s] (→ cos m=0) - W+[s] = b_{l(s)} · (∇Y[s] · e_+) (→ cos m=1) - W-[s] = b_{l(s)} · (∇Y[s] · e_-) (→ sin m=1) - - with a_l, b_l the coefficients of ``build_D_slice_analytic`` and e_± the - gauge frame. Rotate-in / unrotate become elementwise-multiply + per-l - segment sums — Σ_l (2l+1)·k flops per channel instead of the slice bmm's - n_sph·n_kept, and no D tensor is materialized. Same single-backward - contract as the analytic slice (∇Y values; Hessians in backward). - """ - if m_max > 1: - raise ValueError(f"analytic_rotation_weights supports m_max <= 1, " - f"got m_max={m_max}") - Y, dY = _SphYGrad.apply(r_hat, l_max) - a = torch.cat([torch.full((2 * l + 1,), - math.sqrt(4.0 * math.pi / (2 * l + 1)), - dtype=r_hat.dtype, device=r_hat.device) - for l in range(l_max + 1)]) - W0 = a * Y - if m_max == 0: - return RotationWeights(W0.unsqueeze(1)) - b = torch.cat([torch.full((2 * l + 1,), - (math.sqrt(4.0 * math.pi / (2 * l + 1)) - * math.sqrt(2.0 / (l * (l + 1)))) if l > 0 else 0.0, - dtype=r_hat.dtype, device=r_hat.device) - for l in range(l_max + 1)]) - D1 = build_D1_from_rhat(r_hat) - e_p = torch.stack([D1[:, 2, 2], D1[:, 0, 2], D1[:, 1, 2]], dim=1) - e_m = torch.stack([D1[:, 2, 0], D1[:, 0, 0], D1[:, 1, 0]], dim=1) - Wp = b * torch.einsum('nc,ncs->ns', e_p, dY) - Wm = b * torch.einsum('nc,ncs->ns', e_m, dY) - return RotationWeights(torch.stack([W0, Wp, Wm], dim=1)) diff --git a/tests/test_wigner_slice.py b/tests/test_wigner_slice.py deleted file mode 100644 index 146b391..0000000 --- a/tests/test_wigner_slice.py +++ /dev/null @@ -1,280 +0,0 @@ -# Prototype, mainly implemented by Claude -"""Truncated Wigner-D: the kept-column slice (m_max < l_max) and the analytic -small-m build. - -With the angular layout truncated at m_max, the layers only consume bond-frame -components with |m| <= min(l, m_max), and the packed features are zero outside -those slots — so both frame changes need only the corresponding COLUMNS of -each D^l block. Checks: - - 1. build_D_slice == the kept columns of build_D_block, exactly. - 2. build_D_slice_analytic (m_max <= 1; Y for m=0, tangential gradients along - the gauge frame for m=±1) == the recursion to fp rounding, on both gauge - charts and on-axis edges; gradients match through the normalize chain - (direct-r̂ gradients differ by a pure radial component — the two builds - extend D off the unit sphere differently, and normalize projects that - out); double backward raises (single-backward contract); m_max > 1 - rejected. - 3. Model equality: with m_max < l_max the sliced path (automatic) is - BIT-IDENTICAL to the full-block path (escape hatch _use_d_slice=False) - in energies, forces, and (l0, l1) embeddings, with and without MP; - forward_batch_multi included; SO(3) invariance holds on the sliced path. - 4. The analytic toggle: forces match the recursion path; force-loss double - backward raises; fused/m_max guards. - -Run: python tests/test_wigner_slice.py -""" - -import os -import sys # repo root on path for `import ecenet` when run as a script - -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - - -import torch -import torch.nn.functional as F - -from ecenet import ECENet -from ecenet.spherical import ( - build_D_block, - build_D_slice, - build_D_slice_analytic, - kept_offsets, - n_kept_columns, -) - -torch.manual_seed(0) -DTYPE = torch.float64 -N_TYPES = 4 -COMMON = dict(n_types=N_TYPES, r_cut_edge=5.0, r_cut_neighbor=4.0, - l_max=3, n_max=3, embed_dim=8, n_layers=1, n_max_d=4) - - -def hard_directions(n=24, seed=0): - """Random unit vectors plus the awkward ones: chart-B (|rx| >= 0.9) and - exactly on-axis (the gauge charts' dead branches).""" - g = torch.Generator().manual_seed(seed) - r = torch.nn.functional.normalize( - torch.randn(n, 3, generator=g, dtype=DTYPE), dim=-1) - r[0] = F.normalize(torch.tensor([0.95, 0.2, 0.05], dtype=DTYPE), dim=-1) - r[1] = torch.tensor([1.0, 0.0, 0.0], dtype=DTYPE) - r[2] = torch.tensor([0.0, 1.0, 0.0], dtype=DTYPE) - r[3] = torch.tensor([0.0, 0.0, 1.0], dtype=DTYPE) - return r - - -def kept_cols(l_max, m_max): - return [i for l in range(l_max + 1) - for i in range(l * l + l - min(l, m_max), - l * l + l + min(l, m_max) + 1)] - - -def random_structure(n=7, seed=0): - g = torch.Generator().manual_seed(seed) - pos = torch.randn(n, 3, generator=g, dtype=DTYPE) * 1.8 - types = torch.randint(0, N_TYPES, (n,), generator=g) - return pos, types - - -def rand_rotation(seed=1): - g = torch.Generator().manual_seed(seed) - Q, R = torch.linalg.qr(torch.randn(3, 3, generator=g, dtype=DTYPE)) - Q = Q * torch.sign(torch.diag(R)) - if torch.det(Q) < 0: - Q[:, 0] = -Q[:, 0] - return Q - - -def test_slice_matches_block_columns(): - r = hard_directions() - for l_max in (2, 3, 4): - Dfull = build_D_block(r, l_max) - for m_max in (0, 1, 2, l_max): - cols = kept_cols(l_max, m_max) - assert len(cols) == n_kept_columns(l_max, m_max) - assert kept_offsets(l_max, m_max)[-1] == len(cols) - d = (build_D_slice(r, l_max, m_max) - Dfull[:, :, cols]).abs().max() - assert d == 0.0, f"l_max={l_max} m_max={m_max}: slice != cols {d:.2e}" - print(" build_D_slice == kept columns of build_D_block (exact)") - - -def test_analytic_matches_recursion(): - r = hard_directions() - for l_max in (2, 3, 5): - for m_max in (0, 1): - ref = build_D_slice(r, l_max, m_max) - d = (build_D_slice_analytic(r, l_max, m_max) - ref).abs().max() - assert d < 1e-13, f"l_max={l_max} m_max={m_max}: analytic {d:.2e}" - # gradients through the model's actual chain (positions → normalize → D): - # comparing at fixed r̂ would be ill-posed — the two builds extend D off - # the unit sphere differently (pure radial disagreement, projected out - # by normalize's backward) - raw = hard_directions()[:10] * 2.0 - r1 = raw.clone().requires_grad_(True) - g1 = torch.autograd.grad( - build_D_slice(F.normalize(r1, dim=-1), 3, 1).pow(2).sum(), r1)[0] - r2 = raw.clone().requires_grad_(True) - g2 = torch.autograd.grad( - build_D_slice_analytic(F.normalize(r2, dim=-1), 3, 1).pow(2).sum(), r2)[0] - dg = (g1 - g2).abs().max() - assert dg < 1e-12, f"analytic gradient mismatch: {dg:.2e}" - - # single-backward contract: building a force graph (create_graph=True) - # raises at the FIRST backward — loud and early, because a silent - # fallback would drop this branch's second derivative (the analytic D - # mixes the Y-based Function with the differentiable D1 block) - r3 = hard_directions()[:4].clone().requires_grad_(True) - try: - torch.autograd.grad(build_D_slice_analytic(r3, 3, 1).pow(2).sum(), - r3, create_graph=True) - raise AssertionError("create_graph through analytic D should raise") - except RuntimeError as e: - assert 'single-backward' in str(e) - # ...while a plain (single) backward is fine, even inside enable_grad - with torch.enable_grad(): - r4 = hard_directions()[:4].clone().requires_grad_(True) - g = torch.autograd.grad(build_D_slice_analytic(r4, 3, 1).pow(2).sum(), - r4)[0] - assert torch.isfinite(g).all() - try: - build_D_slice_analytic(hard_directions()[:4], 3, 2) - raise AssertionError("m_max=2 should have been rejected") - except ValueError as e: - assert 'm_max' in str(e) - print(f" analytic == recursion (grad via normalize {dg:.1e}); " - "double backward + m_max>1 rejected") - - -def test_model_slice_bit_identical(): - pos, types = random_structure() - for n_mp, m_max in ((1, 1), (2, 1), (2, 2)): - torch.manual_seed(0) - m = ECENet(**COMMON, n_mp=n_mp, m_max=m_max).double() - if n_mp >= 2: - for L in m.mp_layers: - with torch.no_grad(): - L.msg_up.weights.normal_(std=0.3) - L.msg_up.bias.normal_(std=0.1) - assert m._d_slice_active(), "slice should be active for m_max < l_max" - - def ef(model, p0): - p = p0.clone().requires_grad_(True) - e, l0, l1 = model(p, types, return_embeddings=True) - f = torch.autograd.grad(e, p)[0] - return e.detach(), f, l0.detach(), l1.detach() - - e1, f1, l01, l11 = ef(m, pos) - m._use_d_slice = False # escape hatch → full-block path - e2, f2, l02, l12 = ef(m, pos) - m._use_d_slice = True - d = max((e1 - e2).abs().item(), (f1 - f2).abs().max().item(), - (l01 - l02).abs().max().item(), (l11 - l12).abs().max().item()) - assert d == 0.0, f"n_mp={n_mp} m_max={m_max}: slice != full ({d:.2e})" - - # SO(3) invariance on the sliced path. The bit-identity assertion - # above already proves slice == full, so any residual here is the - # random-weight model's own float64 noise floor (the std=0.3 MP - # perturbation amplifies rounding), not a slice-induced break. - Q = rand_rotation() - de = (m(pos, types) - m(pos @ Q.T, types)).abs().item() - assert de < 1e-8, f"SO(3) broken on sliced path: {de:.2e}" - - # batched path - _, l0_list = m.forward_batch_multi([pos, pos + 0.05], [types, types], - return_embeddings=True, l0_only=True) - m._use_d_slice = False - _, l0_ref = m.forward_batch_multi([pos, pos + 0.05], [types, types], - return_embeddings=True, l0_only=True) - m._use_d_slice = True - db = max((a - b).abs().max().item() for a, b in zip(l0_list, l0_ref)) - assert db == 0.0, f"batched slice != full: {db:.2e}" - print(f" n_mp={n_mp}, m_max={m_max}: sliced == full (E/F/l0/l1/batched " - f"exact), SO(3) {de:.1e}") - - -def test_matrix_free_weights(): - """analytic_rotation_weights: the weight-vector contraction equals the - analytic D-slice columns exactly (it is the same rotation, matrix-free), - for m_max 0 and 1, hard directions included.""" - from ecenet.spherical import analytic_rotation_weights - - r = hard_directions(16) - A = torch.randn(16, 6, 16, dtype=DTYPE) - for m_max in (0, 1): - ref = torch.bmm(A, build_D_slice_analytic(r, 3, m_max)) - W = analytic_rotation_weights(r, 3, m_max).w - off = kept_offsets(3, m_max) - worst = 0.0 - for l in range(4): - s, e = l * l, (l + 1) * (l + 1) - mk = min(l, m_max) - def seg(w): - return torch.einsum('ecs,es->ec', A[:, :, s:e], w[:, s:e]) - rows = ([seg(W[:, 2])] if mk >= 1 else []) + [seg(W[:, 0])] \ - + ([seg(W[:, 1])] if mk >= 1 else []) - for k, gcol in enumerate(rows): - worst = max(worst, - (gcol - ref[:, :, off[l] + k]).abs().max().item()) - assert worst < 1e-14, f"m_max={m_max}: weights != slice ({worst:.2e})" - print(" matrix-free weights == analytic slice columns (m_max 0 and 1)") - - -def test_model_analytic(): - pos, types = random_structure() - torch.manual_seed(0) - m = ECENet(**COMMON, n_mp=2, m_max=1).double() - for L in m.mp_layers: - with torch.no_grad(): - L.msg_up.weights.normal_(std=0.3) - p1 = pos.clone().requires_grad_(True) - e1 = m(p1, types) - f1 = torch.autograd.grad(e1, p1)[0] - m.set_analytic_wigner(True) - p2 = pos.clone().requires_grad_(True) - e2 = m(p2, types) - f2 = torch.autograd.grad(e2, p2)[0] - de = (e1 - e2).abs().item() - df = (f1 - f2).abs().max().item() - assert de < 1e-12 and df < 1e-11, f"analytic model mismatch: {de:.2e}/{df:.2e}" - - # force-loss training must raise, not silently mis-train: the error fires - # already at the create_graph force computation - p3 = pos.clone().requires_grad_(True) - try: - torch.autograd.grad(m(p3, types), p3, create_graph=True) - raise AssertionError("force-graph build through analytic D should raise") - except RuntimeError as e: - assert 'single-backward' in str(e) - - # guards: fused is mutually exclusive; m_max > 1 rejected - try: - m.set_edge_frame_fused(True) - raise AssertionError("fused + analytic should be rejected") - except ValueError as e: - assert 'analytic' in str(e) - m.set_analytic_wigner(False) - m.set_edge_frame_fused(True) - try: - m.set_analytic_wigner(True) - raise AssertionError("analytic + fused should be rejected") - except ValueError as e: - assert 'fused' in str(e).lower() or 'edge_frame' in str(e) - m.set_edge_frame_fused(False) - m2 = ECENet(**COMMON, m_max=2).double() - try: - m2.set_analytic_wigner(True) - raise AssertionError("m_max=2 analytic should be rejected") - except ValueError as e: - assert 'm_max' in str(e) - print(f" analytic model: E/F match (dE={de:.1e}, dF={df:.1e}); " - "force-loss raises; guards OK") - - -if __name__ == "__main__": - print("Truncated Wigner-D (kept-column slice + analytic small-m build)") - test_slice_matches_block_columns() - test_analytic_matches_recursion() - test_model_slice_bit_identical() - test_matrix_free_weights() - test_model_analytic() - print("All tests passed.") diff --git a/tools/profile_step.py b/tools/profile_step.py index c5bba9f..7392e54 100644 --- a/tools/profile_step.py +++ b/tools/profile_step.py @@ -15,9 +15,8 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) # repo root → import ecenet from ecenet.calculator import load_calculator -from ecenet.model import _matrix_free_to_angular from ecenet.radial import radial_basis -from ecenet.spherical import RotationWeights, build_D_block_from_list, recursive_wigner_D, wigner_rotate +from ecenet.spherical import build_D_block_from_list, recursive_wigner_D, wigner_rotate parser = argparse.ArgumentParser() parser.add_argument('--checkpoint', required=True) @@ -36,12 +35,6 @@ help='Enable the fused gather+Wigner-rotate+reshape edge-frame op ' '(set_edge_frame_fused; Triton on CUDA+float32, else eager ' 'recompute-in-backward).') -parser.add_argument('--analytic_wigner', action='store_true', - help='Matrix-free analytic Wigner rotation ' - '(set_analytic_wigner; m_max <= 1 checkpoints only, ' - 'single-backward — forces fine, no force-loss ' - 'training; mutually exclusive with ' - '--edge_frame_fused).') parser.add_argument('--edge_frame_e2n', action='store_true', help="With --edge_frame_fused: ALSO fuse the MP layers' " 'pack+unrotate (their step 3).') @@ -83,13 +76,6 @@ e2n = 'on' if args.edge_frame_e2n else 'off' print(f"[edge_frame_fused] fused gather+rotate+reshape ({path}, e2n {e2n})") -if args.analytic_wigner: - # raises for m_max > 1 or together with --edge_frame_fused (the model's - # own guards) — surface those as-is rather than pre-checking here - model.set_analytic_wigner(True) - print(f"[analytic_wigner] matrix-free rotation (m_max={model.m_max}; " - "single-backward — forces profiled, no force-loss training)") - # ── Utilities ───────────────────────────────────────────────────────────────── def sync(): @@ -179,9 +165,7 @@ def fresh_pos(): A_emb = model._embed(A, types) A_both = torch.cat([A_emb[edge_i], A_emb[edge_j]], dim=1) -# Wigner rotation breakdown — reference full-block decomposition first, -# then the ACTIVE path if it differs (kept-column slice for m_max < l_max, -# or matrix-free weights under set_analytic_wigner). +# Wigner rotation breakdown _, D_list = time_fn(" Wigner D (recursive_wigner_D)", lambda: recursive_wigner_D(r_hat, model.l_max)) _, D_block_main = time_fn(" Wigner D (build_D_block_from_list)", @@ -192,19 +176,6 @@ def fresh_pos(): time_fn("Wigner rotate total (main, w/ cached D)", lambda: wigner_rotate(A_both, D_block_main)) -D_active = model._build_D(r_hat) -if isinstance(D_active, RotationWeights): - time_fn("Wigner ACTIVE: matrix-free weights (build)", - lambda: model._build_D(r_hat)) - time_fn("Wigner ACTIVE: matrix-free rotate+angular", - lambda: _matrix_free_to_angular(A_both, D_active, model.l_max, - model.sph_to_angular)) -elif D_active.shape[-1] != model.n_sph: - time_fn("Wigner ACTIVE: D slice (build)", - lambda: model._build_D(r_hat)) - time_fn("Wigner ACTIVE: bmm (A_both @ D_slice)", - lambda: torch.bmm(A_both, D_active)) - A_rot = wigner_rotate(A_both, D_block_main) # sph_to_angular