Truncated Wigner-D: kept-column slice + analytic small-m build - #10
Merged
Merged
Conversation
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
This was referenced Aug 28, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Three levels of Wigner-rotation simplification for truncated models (
m_max < l_max); previously the truncation shrank only the layer stack — the rotation always built and applied the full(n_sph, n_sph)block.1. Kept-column slice (
build_D_slice, automatic wheneverm_max < l_max): the layers only consume bond-frame components with|m| <= min(l, m_max), and packed features are exactly zero outside those slots, so every frame change (edge-frame rotate, MP send/receive, batched paths) needs only the corresponding columns of eachD^lblock. Pure differentiable torch, bit-identical to the full-block path (E/F/l0/l1 + batched, exactly 0 intest_wigner_slice.py). Fused Triton kernels keep the full block — they already read only kept entries.2. Analytic small-m columns (
build_D_slice_analytic,m_max <= 1): closed forms replace the CG recursion —D^l[:,m=0] = √(4π/(2l+1))·Y_l(r̂),D^l[:,m=±1] = √(4π/(2l+1))·√(2/(l(l+1)))·(∇Y_l·e_±)with the tangent frame frombuild_D1_from_rhat(verified ~1e-15 on both gauge charts + on-axis edges; gradients to 1e-14 through the normalize chain). Isolated builder benchmark (20k edges, CPU): 9.4 vs 14.7 ms atl_max=3, 39 vs 134 ms atl_max=6— the recursion's cost grows steeply withl_max, the analytic build's doesn't.3. Matrix-free rotation (
model.set_analytic_wigner(True)): the rotation is a point evaluation — m=0 is each channel function's value at r̂, m=±1 its tangential derivatives — so no D tensor is built at all: three weight vectors from one sphericart call, and every frame crossing becomes elementwise-multiply + per-l segment sums (the rotate-in also absorbs the SphToAngular gather). Model E/F match the recursion path at ~1e-16, MP included.Single-backward contract for the analytic/matrix-free mode: the backward runs through sphericart Hessians, and building a force graph raises loudly at the first
create_graphbackward — notably,once_differentiablewould have been silent here (the analytic path mixes the custom Function with the differentiable D1 block, so a double backward would quietly drop part of the second derivative). MD/inference forces and the stress strain pass are fine; force-loss training must use the (default) recursion slice. Mutually exclusive withedge_frame_fused, guarded both ways.New
tests/test_wigner_slice.py(slice exactness, analytic vs recursion incl. gradients and hard directions, matrix-free weight contraction, model bit-identity, single-backward and guard behavior); full suite green, ruff clean.🤖 Generated with Claude Code