[spossann]

Summary

With polar splines, the implicit shear Alfvén step (ShearAlfvenPropagator) fell back from the Kronecker Schur preconditioner to PolarMassMatrixPreconditioner, which ignores the dt**2/4 T^T C^T M2 C T part of the Schur operator. On the low-mode TAE example with polar splines (IGAPolarTorus, 16x18x8, dt = 0.5) this needed about 160 PCG iterations per step, against about 21 with tensor-product splines.

This PR adds a polar version of ShearAlfvenSchurPreconditioner that does not assume axisymmetry and is now the default with polar splines. It also adds an exact solver for axisymmetric operators as an explicit option, and fixes a crash of the energy scalars with polar splines.

Changes

PolarKroneckerBlockPreconditioner (feec/preconditioner.py): a polar 2D x 1D preconditioner for sum_i L_i^T M_i L_i + sigma S, where L_i are KroneckerBlockOperators on the parent tensor-product spaces (here L = C T~). It builds on the machinery of #747 and #758.

• Poloidal plane: everything is kept. The poloidal Kronecker factors kron(F1, F2) of L are combined with the poloidal mass matrices (weights reduced in eta_3, as in Polar splines: StiffnessPreconditioner (grad, curl) and stab_average for GyrokineticPoisson #758), and the polar extraction is applied. The components (u1, u2), which are coupled at the pole, are solved together.
• eta_3: the 1D factors F3^T M3 G3 are replaced by their diagonals in the eigenbasis of a summary stiffness, as in the tensor-product preconditioner. This uses PolarKronSum2Dx1DSolver, with one poloidal factorization per eta_3 mode.
• What is approximate: the eta_3 dependence of the weights, the antisymmetric parts of the eta_3 factors (the poloidal/toroidal cross terms), the coupling between the groups {u1, u2} and {u3}, and the separable weights of T~. It works for any geometry and does not assume axisymmetry.

ShearAlfvenSchurPreconditioner uses it with polar splines, and resolve() no longer falls back for polar.

• The E x B auxiliary-space correction is not used with polar splines. With polar splines its 0-form part (sigma = 0) was indefinite (eigenvalues of ±1e12 on a 6x18x4 grid, PCG hits maxiter), and it is not needed: coupling u1 and u2 in the poloidal plane already handles the E x B flows.
• kronecker_block.basis_projection_operator accepts polar operators (it approximates them on the tensor-product FEM spaces).
• derivative_operator now uses derham.{grad,curl,div}_tp. These are the same objects as *_bcfree without polar splines.
• New helper mass_matrix_1d.

ToroidalFourierSolver (new module feec/toroidal_fourier.py), option precond="ToroidalFourierPreconditioner":

• It is the exact inverse of an operator that is invariant under shifts in a periodic eta_3, such as an axisymmetric equilibrium in a torus or cylinder, with tensor-product or polar splines.
• It works by probing (one operator application per poloidal unknown), an rfft along eta_3 and one dense Cholesky or LU factorization per toroidal mode number. The modes are distributed over the processes.
• It checks itself on a random vector and logs the residual.
• It is meant for moderate poloidal resolution, and it is not the default because it relies on axisymmetry (stellarators).
• ShearAlfvenCurrentCoupling5D rejects this option.

Bug fix: en_B_eq and en_B_tot of ShearAlfven and of the linear MHD models (models/utils.MagneticEnergies) paired the tensor-product projected_equil.b2 with the polar mass matrix, which raised an AssertionError with polar splines. The new ProjectedFluidEquilibriumWithB.b1_polar/b2_polar return the projection in the polar space; without polar splines they are identical to b1/b2, so baselines are unchanged.

Results

Low-mode TAE (examples/ShearAlfven/lowmode_tae, IGAPolarTorus 16x18x8, degree 3, dt = 0.5, tol 1e-8, 2 ranks):

preconditioner	PCG iterations per step	time per step	setup
PolarMassMatrixPreconditioner (before)	155-165	3.2 s (on 4 ranks)	–
StiffnessPreconditioner → PolarKroneckerBlockPreconditioner (new default)	14	0.69 s	~6 s
ToroidalFourierPreconditioner (exact, axisymmetric)	2	–	920 probes; whole 4-step run 22 s
tensor-product Kronecker, HollowTorus 16x16x8 (for comparison)	21	–	–

Dense check on 6x18x4 (polar, serial), condition number of P S:

• polar mass preconditioner: 558;
• new preconditioner without E x B correction: 3.3;
• with the E x B correction: indefinite, hence disabled.

ToroidalFourierSolver residual on a random vector: 2e-13 on the TAE case, and 1e-15 in the unit tests, which cover tensor-product and polar splines and grids decomposed along eta_3 and eta_1.

Tests

• propagators/tests/test_shear_alfven.py: resolve() is updated (polar keeps StiffnessPreconditioner; checks of the Fourier option). The new test_shear_alfven_polar_preconditioners checks that the polar mass preconditioner, the new default and the Fourier preconditioner give the same step (to 1e-8), that the new ones need fewer iterations (Fourier at most 3), and that energy is conserved.
• New feec/tests/test_toroidal_fourier.py: HollowCylinder and IGAPolarCylinder; M2 and M1 + G M0 G^T (Cholesky); M2 + 0.3 M2B (non-symmetric, LU); grid decompositions along eta_3 and eta_1.
• Results: these two files pass serially (12) and on 2 ranks (11 passed, 1 mpi_skip). test_kronecker_block.py and test_stiffness_preconditioner.py pass serially (67) and on 2 ranks.
• Tensor-product path: unchanged (same objects and code path), so no baseline changes are expected.
• Not tested: the polar path of ShearAlfvenCurrentCoupling5D. With polar splines it now also gets the polar preconditioner, but the 5D models cannot run with polar splines yet (Polar splines in LinearMHDDriftkineticCC: 5D current-coupling propagators index PolarVectors as BlockVectors #917).

Follow-ups

• Polar splines in LinearMHDDriftkineticCC: 5D current-coupling propagators index PolarVectors as BlockVectors #917: polar splines in the 5D current-coupling propagators (LinearMHDDriftkineticCC).
• A sparse alternative for ToroidalFourierSolver at higher poloidal resolution.
• A complex-Fourier variant of the eta_3 diagonalization in the 2D x 1D preconditioner. It would keep the poloidal/toroidal cross terms exactly for periodic eta_3, which matter for small k_par.

🤖 Generated with Claude Code

[Copilot]

🟡 Changes recommended

The Fourier solver fails on CuPy due to implicit device-to-NumPy conversions, and the preconditioner configuration accepts CG even though it does not support preconditioning.

2 open issues

• Reject CG when propagator configures a preconditioner · New
• CuPy Fourier setup fails on implicit device-to-NumPy conversion · New

What changed in this PR

This PR extends Struphy’s shear Alfvén solvers to improve polar-spline preconditioning, building on the existing 2D × 1D solver infrastructure. It also adds an optional Fourier solver for shift-invariant operators and fixes polar magnetic-energy calculations.

Changes:

• Makes the polar Kronecker Schur preconditioner the default for supported shear Alfvén configurations.
• Adds a toroidal Fourier solver with serial/MPI tests and propagator integration.
• Uses polar-space equilibrium projections for magnetic-energy diagnostics.

File	Description
src/​struphy/​propagators/​tests/​test_shear_alfven.py	Tests polar preconditioner agreement, iteration counts, and energy conservation.
src/​struphy/​propagators/​shear_alfven_propagator.py	Integrates the Fourier option and documents polar preconditioning.
src/​struphy/​propagators/​shear_alfven_current_coupling_5d.py	Updates polar documentation and rejects the unsupported Fourier option.
src/​struphy/​models/​utils.py	Uses polar equilibrium coefficients in shared magnetic-energy diagnostics.
src/​struphy/​models/​shear_alfven.py	Corrects equilibrium coefficients used for magnetic energies.
src/​struphy/​io/​options.py	Registers the Fourier preconditioner option.
src/​struphy/​fields_background/​projected_equils.py	Exposes magnetic-field projections in polar spaces.
src/​struphy/​feec/​toroidal_fourier.py	Adds Fourier-mode factorization and distributed solves.
src/​struphy/​feec/​tests/​test_toroidal_fourier.py	Tests symmetric and nonsymmetric operators with both spline types.
src/​struphy/​feec/​preconditioner.py	Adds polar Kronecker block preconditioning and Schur dispatch.
src/​struphy/​feec/​kronecker_block.py	Supports parent-space polar approximations and adds 1D mass assembly.
doc/​sections/​subsections/​dev-feec.rst	Includes Fourier solver API documentation.

_{Review effort: Balanced}