From 3181a105580929502672f98590fe36db50d8c73b Mon Sep 17 00:00:00 2001 From: alacour Date: Fri, 28 Aug 2026 13:52:53 -0700 Subject: [PATCH 1/2] Truncated Wigner-D: kept-column slice + analytic small-m build MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With the angular layout truncated at m_max < l_max, the layers only consume bond-frame components with |m| <= min(l, m_max), and the packed features are exactly zero outside those slots — so both frame changes (rotate in, unrotate out; MP send and receive included) need only the corresponding COLUMNS of each D^l block. Previously the full (n_sph, n_sph) block was built and multiplied regardless, so lowering m_max never shrank the rotation. build_D_slice (spherical.py) is the recursion + column slice — pure differentiable torch, and the eager paths now use it automatically whenever m_max < l_max (bit-identical to the full block: E/F/l0/l1 and the batched path all at exactly 0 in tests). SphToAngular gains a kept-layout index set; the MP layer packs/unpacks against the narrow layout (_sph_pack_index_sliced). The fused Triton kernels keep the full block — they already read only the kept entries. _use_d_slice is the internal escape hatch the equality tests flip. build_D_slice_analytic (m_max <= 1) replaces the CG recursion with one sphericart call, using closed forms discovered numerically against the recursion and verified to ~1e-15 on both gauge charts and on-axis edges: D^l[:, m=0] = sqrt(4pi/(2l+1)) Y_l(r_hat) D^l[:, m=+-1] = sqrt(4pi/(2l+1)) sqrt(2/(l(l+1))) (grad Y_l . e_+-) with e_+- the tangent frame from build_D1_from_rhat (gauge matches by construction). Gradients agree with the recursion through the normalize chain to 1e-14 (direct-r_hat gradients differ by a pure radial component — the two builds extend D off the unit sphere differently, which normalize's backward projects out). SINGLE-BACKWARD ONLY, exposed as model.set_analytic_wigner(True): the backward runs through sphericart Hessians, and building a force graph raises loudly at the first create_graph backward — once_differentiable would have been SILENT here (the analytic D mixes the Y Function with the differentiable D1 block, so a double backward would quietly drop this branch's second derivative). Mutually exclusive with edge_frame_fused, guarded both ways. Measured (CPU eager, l_max=3 m_max=1, embed 32): sliced 72.5 vs full 76.9 ms/forward; the analytic build is not faster than recursion+slice on CPU at this size (75.9) — its value is the shallower graph, not speed. Co-Authored-By: Claude Fable 5 --- ecenet/model.py | 209 +++++++++++++++++++++++++++--- ecenet/spherical.py | 134 ++++++++++++++++++++ tests/test_wigner_slice.py | 252 +++++++++++++++++++++++++++++++++++++ 3 files changed, 578 insertions(+), 17 deletions(-) create mode 100644 tests/test_wigner_slice.py diff --git a/ecenet/model.py b/ecenet/model.py index bad19bd..83ed7e8 100644 --- a/ecenet/model.py +++ b/ecenet/model.py @@ -39,7 +39,14 @@ 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 build_D1_from_rhat, build_D_block, spherical_harmonics_float64, wigner_rotate +from ecenet.spherical import ( + build_D1_from_rhat, + build_D_block, + build_D_slice, + build_D_slice_analytic, + kept_offsets, + spherical_harmonics_float64, +) # 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 @@ -69,7 +76,15 @@ 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 + (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). cutoff_type: 'cosine' or 'poly' activation: pointwise activation in the realspace nonlinearity ('silu', 'tanh', 'relu', 'gelu'); 'identity' turns every RealSpaceNonlinearity — the @@ -778,25 +793,73 @@ 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 Wigner-D slice (m_max <= 1 only): the kept + columns come from one sphericart call — Y for m=0, tangential + gradients along the gauge frame for m=±1 — instead of the CG + recursion (see spherical.build_D_slice_analytic). 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): + return build_D_slice_analytic(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_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) + 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 getattr(self, '_edge_frame_fused', False): - A_cos, A_sin = edge_frame_fused(A_emb, edge_i, edge_j, D_block, + A_cos, A_sin = edge_frame_fused(A_emb, edge_i, edge_j, D, self.sph_to_angular) else: A_both = torch.cat([A_emb[edge_i], A_emb[edge_j]], dim=1) - 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 + 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 def set_activation_fused(self, enabled: bool = True): """Toggle the fused recompute-in-backward path on all RealSpaceNonlinearity @@ -1255,11 +1318,12 @@ 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 = build_D_block(r_hat_flat, self.l_max) - A_rot_flat = wigner_rotate(A_both_flat, D_block) + D_block = self._build_D(r_hat_flat) + 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) + A_cos_flat, A_sin_flat = self.sph_to_angular( + A_rot_flat, sliced=D_block.shape[-1] != self.n_sph) # shapes: (B*n_edges, n_features_per_m, n_angular) # ── Step 5: Equivariant layers ──────────────────────────────────── @@ -1434,6 +1498,73 @@ 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 + + class ECENetAttentionMPLayer(nn.Module): """Attention-style message passing for ECENet. @@ -1619,9 +1750,18 @@ 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. + d_sliced = D_block.shape[-1] != self.n_sph if getattr(self, 'edge_frame_fused_e2n', False): + assert not d_sliced, "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 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 @@ -1680,8 +1820,15 @@ 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, "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 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) @@ -1763,10 +1910,36 @@ 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) - def forward(self, A_rot): + # 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): """ Args: - A_rot: (n_edges, 2*embed_dim, n_sph) + 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 Returns: A_cos, A_sin: (n_edges, 2*embed_dim*(l_max+1), l_max+1) """ @@ -1774,10 +1947,12 @@ def forward(self, A_rot): # 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). - 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) + 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) .view(n_edges, self.n_ch, self.n_angular)) * self.cos_valid - A_sin = (A_flat.index_select(1, self.sin_flat_idx) + A_sin = (A_flat.index_select(1, si) .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 bf259ca..203991d 100644 --- a/ecenet/spherical.py +++ b/ecenet/spherical.py @@ -270,3 +270,137 @@ 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) diff --git a/tests/test_wigner_slice.py b/tests/test_wigner_slice.py new file mode 100644 index 0000000..e1ef75e --- /dev/null +++ b/tests/test_wigner_slice.py @@ -0,0 +1,252 @@ +# 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_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_model_analytic() + print("All tests passed.") From 7044fce45a9de68e54de80716e043d3909e760db Mon Sep 17 00:00:00 2001 From: alacour Date: Fri, 28 Aug 2026 14:07:44 -0700 Subject: [PATCH 2/2] Matrix-free analytic rotation: point evaluation instead of D matrices MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit For m_max <= 1 the bond-frame rotation is not really a matrix operation: contracting the analytic D columns with the features shows the m=0 output is each channel function's VALUE at r_hat and the m=+-1 outputs are its two tangential derivatives there. set_analytic_wigner(True) now builds no D tensor at all — analytic_rotation_weights returns three (E, n_sph) weight vectors [W0, W+, W-] from one sphericart call, and every frame change becomes elementwise-multiply + per-l segment sums: - rotate-in goes straight to the angular layout (bmm AND the SphToAngular gather both disappear), - MP send broadcasts the kept components over their l-blocks and sums the three weighted copies, - MP receive is the same segment contraction on the gathered aggregate. Flops per channel drop from the slice bmm's n_sph*n_kept to sum_l (2l+1)*k, and the only rotation state is (E, k, n_sph) weights. The l=0 rows of W+- are exactly zero (b_0 = 0), so structural-scratch content in the m=1 slots of l=0 channels is annihilated exactly on the send path — no masking needed there, though the valid masks are applied on outputs for bit-parity. RotationWeights is a distinct type so the MP layer dispatches on isinstance rather than shape-sniffing (kept widths can collide with k). Same single-backward contract as before (the weights carry sphericart gradient values); build_D_slice_analytic stays as the verification reference. Verified: weight contraction == analytic slice columns to 1e-15 (m_max 0 and 1, hard directions); model E/F vs the recursion path at ~1e-16 for n_mp 1/2, m_max 0/1; force-loss still raises; full suite green. Whole-forward CPU (l_max=3, m_max=1, embed 32): matrix-free 118.3 vs slice 120.3 vs full block 127.3 ms — the CPU win is modest because the rotation is a small share at this size; the structural win (no D tensor, one sphericart call, no per-l einsum chain) is aimed at GPU launch latency and memory. Co-Authored-By: Claude Fable 5 --- ecenet/model.py | 116 ++++++++++++++++++++++++++++++++----- ecenet/spherical.py | 57 ++++++++++++++++++ tests/test_wigner_slice.py | 28 +++++++++ 3 files changed, 186 insertions(+), 15 deletions(-) diff --git a/ecenet/model.py b/ecenet/model.py index 83ed7e8..578af41 100644 --- a/ecenet/model.py +++ b/ecenet/model.py @@ -40,10 +40,11 @@ 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, - build_D_slice_analytic, kept_offsets, spherical_harmonics_float64, ) @@ -805,10 +806,12 @@ def set_edge_frame_fused(self, enabled: bool = True, e2n: bool = True): return self def set_analytic_wigner(self, enabled: bool = True): - """Toggle the analytic Wigner-D slice (m_max <= 1 only): the kept - columns come from one sphericart call — Y for m=0, tangential - gradients along the gauge frame for m=±1 — instead of the CG - recursion (see spherical.build_D_slice_analytic). SINGLE-BACKWARD + """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.""" @@ -841,7 +844,9 @@ def _build_D(self, r_hat): path is active (analytic when requested), full block otherwise.""" if self._d_slice_active(): if getattr(self, '_analytic_wigner', False): - return build_D_slice_analytic(r_hat, self.l_max, self.m_max) + # 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) @@ -851,7 +856,11 @@ def _edge_frame(self, A_emb, edge_i, edge_j, r_hat): 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 getattr(self, '_edge_frame_fused', False): + 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, self.sph_to_angular) else: @@ -1319,11 +1328,14 @@ 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) - 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) + 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) # shapes: (B*n_edges, n_features_per_m, n_angular) # ── Step 5: Equivariant layers ──────────────────────────────────── @@ -1565,6 +1577,50 @@ def _unpack_kept_to_angular(v_rot, n_base, l_max, m_max, 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. @@ -1750,11 +1806,29 @@ 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. - d_sliced = D_block.shape[-1] != self.n_sph + 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, "fused e2n needs the full D block" + 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 @@ -1820,9 +1894,21 @@ 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, "fused rotate-back needs the full D block" + 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( diff --git a/ecenet/spherical.py b/ecenet/spherical.py index 203991d..9b1ad46 100644 --- a/ecenet/spherical.py +++ b/ecenet/spherical.py @@ -404,3 +404,60 @@ def build_D_slice_analytic(r_hat: torch.Tensor, l_max: int, 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 index e1ef75e..146b391 100644 --- a/tests/test_wigner_slice.py +++ b/tests/test_wigner_slice.py @@ -192,6 +192,33 @@ def ef(model, p0): 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) @@ -248,5 +275,6 @@ def test_model_analytic(): 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.")