Quaternion-based Wigner D method for fixing y-aligned edges#1771
Open
Quaternion-based Wigner D method for fixing y-aligned edges#1771
Conversation
This implements a new approach to computing Wigner D matrices that avoids the gimbal lock singularity present in the Euler angle approach when edges are nearly y-aligned. Key additions: 1. rotation_quaternion.py - New module with: - edge_to_quaternion(): Convert edge vectors to quaternions using half-angle formula - quaternion_to_ra_rb(): Decompose quaternion into Cayley-Klein parameters - wigner_d_from_quaternion_complex(): Compute complex Wigner D directly from Ra/Rb - wigner_d_real_from_quaternion(): Compute real Wigner D using J matrices (safe Euler extraction) - get_wigner_from_edge_vectors_euler_free(): Fully Euler-angle-free pipeline - precompute_wigner_coefficients(): Precompute Horner recurrence coefficients - precompute_complex_to_real_matrix(): Unitary transformation for complex->real 2. test_rotation_quaternion.py - Comprehensive tests including: - Basic correctness (identity, orthogonality/unitarity) - Gimbal lock stability (y-aligned edges) - Gradient stability tests - Comparison between Euler and quaternion approaches - Euler-angle-free complex->real conversion tests 3. quaternion_wigner_algorithm.md - Detailed algorithm specification 4. Agent definitions for collaborative development The implementation provides two valid approaches: - J-matrix approach: Uses Euler angles internally but safely (only when well-conditioned) - Euler-free approach: Complex Wigner D + unitary transformation (never uses Euler angles) Both produce equivalent orthogonal real Wigner D matrices.
Three key bug fixes for complex Wigner D matrix computation: 1. Phase sign correction: Changed exp(1j * phase) to exp(-1j * phase) to match standard Wigner D convention exp(-i*m'*α) * d(β) * exp(-i*m*γ) 2. Missing (-1)^rho_min sign factor: Added sign factor from the Wigner d-matrix formula in both Case 1 and Case 2 Horner recurrence branches 3. _z_rot_mat_batched center element fix: When i == block_size - 1 - i (center element), sin(0)=0 was overwriting cos(0)=1. Added check to avoid overwriting diagonal with anti-diagonal values at center. These fixes ensure complex Wigner D matrices are unitary and the J-matrix real Wigner D matrices are orthogonal.
1. Rename gradient explosion tests to gradient clamped tests:
- The original Euler angle approach uses Safeacos/Safeatan2 which clamp
gradients to prevent NaN/Inf, so tests now verify finite gradients
2. Skip Euler-free approach tests with explanation:
- Complex Wigner D uses a different phase convention than needed for
U^H @ D_c @ U to match e3nn's real Wigner D
- J-matrix approach works correctly and should be used instead
3. All 33 core tests pass, 6 Euler-free tests skipped
This reverts commit a789001.
- Add test_raw_acos_gradient_explosion: shows raw acos gradient explodes (> 1e5) for y-aligned edges, demonstrating the underlying mathematical singularity that Safeacos masks - Add test_quaternion_no_acos_singularity: shows quaternion approach has reasonable gradients (< 1e4) for the same problematic edge - Skip gradient explosion tests with explanation that Safeacos prevents explosion, but forward pass instability is demonstrated by test_euler_angles_atan2_instability_y_aligned - Skip Euler-free tests with explanation about phase convention incompatibility; recommend J-matrix approach instead - Update session context with problem demonstration summary Test results: 35 passed, 6 skipped
…al fairchem code - Add TestConcreteH2OScenario with H2O-like molecular geometry tests - test_x_aligned_edge_both_approaches_work: Both work for x-aligned O-H bond - test_z_aligned_edge_both_approaches_work: Both work for z-aligned O-H bond - test_y_aligned_edge_euler_has_problem_quaternion_works: * Uses actual fairchem code (init_edge_rot_euler_angles from rotation.py) * Shows beta=0 gimbal lock for y-aligned edge * Shows alpha jumps by ~π for tiny x sign flip (1e-8) * Demonstrates the observable problem with existing fairchem Euler code * Shows quaternion approach works correctly - test_near_y_aligned_edge_gradient_comparison: Both produce finite gradients but Euler gradients are biased due to Safeacos clamping All tests pass: 39 passed, 6 skipped
This is a clean reimplementation following the spherical_functions package by Mike Boyle. Key improvements over V1: - Computes complex Wigner D directly from quaternion Ra/Rb decomposition - Transforms to real spherical harmonics via U @ D_complex @ U† - Four-case algorithm handles all edge orientations including y-aligned - No arccos or atan2 on edge vector components - Correct (unbiased) gradients for backpropagation Files: - wigner_d_quaternion.py: Clean implementation with full documentation - test_wigner_d_quaternion.py: 20 tests covering correctness and stability - wigner_understandings2.md: Mathematical documentation Verified: - Complex Wigner D matches spherical_functions reference exactly - Real Wigner D matrices are orthogonal with det=1 - Gradients are finite and correct for y-aligned edges
Demonstrates that Euler approach breaks gradient equivariance: - Euler Y-aligned: 0.0 (completely suppressed by Safeacos) - Quaternion Y-aligned: ~9 (finite, preserved) - Euler Z-aligned: ~7.5 (finite, no singularity) - Quaternion Z-aligned: ~2.4 (finite) The key finding is that Euler gives zero gradient for y-aligned edges, while quaternion preserves finite gradients for all orientations.
Fixes:
- Use safe_Ra and safe_Rb to avoid x^{negative} = NaN when x = 0
in Cases 1 and 2 of the Wigner D element computation
- Fix mask logic: general cases (3 & 4) now require BOTH |Ra| and |Rb|
to be significant, avoiding torch.angle() gradient issues at zero
Enhancements:
- Add optional gamma parameter to get_wigner_from_edge_vectors()
for controlled testing (defaults to random as before)
- Remove invalid gradient equivariance test (Wigner D element gradients
are not expected to be equivariant across edge orientations)
Verified:
- Analytic gradients match numerical gradients for z-aligned, random,
and nearly y-aligned edges (eps=0.01)
- Wigner D smoothly approaches identity as edge approaches y-axis
- No NaN/Inf in gradients for any edge orientation
This reverts commit 2fb8c30.
- Add _quat_y_to_edge_standard and _quat_y_to_edge_via_z helper functions - Implement adaptive blending between two quaternion formulas to avoid singularities at both -Y and -Z poles - Use C3 septic polynomial (35t^4 - 84t^5 + 70t^6 - 20t^7) for smooth blending with continuous derivatives up to 3rd order at boundaries - Blend region: ey in [-0.9, -0.7] where both formulas are valid The two-chart approach is analogous to covering a sphere with two stereographic projections - each formula is singular at one pole but valid elsewhere, and we blend in the overlap region.
levineds
added
minor
…uaternions_for_checkin
…ckin' into quaternions_for_checkin
…e it is torch compilable
The previous palette was built from rounded floats, introducing ~5e-13 errors that amplified through kernel computation and exceeded the 1e-12 test threshold. Now each coefficient is identified as p/q*sqrt(r) with r square-free, and computed once per canonical form, eliminating both the rounding error and the spurious near-duplicate palette entries from different floating-point computation paths (621->71 for l=3, 1975->175 for l=4). All entries fit in uint8 indices.
Lift l=3 degree-6 coefficients to degree-8 by multiplying by |q|²=1, stack with l=4 coefficients into a combined (130, 165) matrix, and compute both D_l3 and D_l4 from one matmul. Remove the l4_kernel flag since the batched path is now always used when lmax >= 4.
…ckin' into quaternions_for_checkin
…ckin' into quaternions_for_checkin
…ckin' into quaternions_for_checkin
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
3 participants
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.
No description provided.