StiffnessPreconditioner for PoissonAdiabaticGyrokinetic: general mid/stab, weighted average in the stabilization

[spossann]

Summary

Makes StiffnessPreconditioner work for PoissonAdiabaticGyrokinetic (and other ImplicitDiffusion operators with a general diffusion matrix and stabilization), and makes the toroidal adiabatic gyrokinetic Poisson operator symmetric.

Follow-up of #739 (merged), rebased onto devel. Addresses the review comment #739 (comment) properly; #739 itself only rejected multigrid for PoissonAdiabaticGyrokinetic.

Points the feectools submodule to the branch of struphy-hub/feectools#102 (KroneckerSumSolver with stiffness coefficients and additional Kronecker terms, feectools 0.8.0). The feectools pin in pyproject.toml is raised once 0.8.0 is on PyPI.

StiffnessPreconditioner: mid, stab, stab_average

Preconditioner for $\mathbb d^T \mathbb M_{\textrm{mid}} \mathbb d + \sigma , \mathbb S (\mathbb I - \mathbb A)$, all new inputs as strings with the previous behavior as default:

• mid: name of a mass operator in mass_ops (e.g. "M1gyro", "M1perp") or the operator itself (WeightedMassOperator, or sums/multiples such as M1gyro = M1_MHDeq - M1para_MHDeq), with the domain/codomain of $\mathbb M_{k+1}$; default $\mathbb M_{k+1}$. Blocks with zero mean (e.g. the parallel direction of a perpendicular operator) have no stiffness term; their 1d stiffness still defines the eigenbasis (stiffness_coeffs in feectools#102).
• stab: name or operator with the domain/codomain of $\mathbb M_k$ (e.g. "M0ad_withT"), or "Id"; default $\mathbb M_k$.
• stab_average (gradient only): directions averaged in the stabilization, e.g. "eta3", "eta2 eta3". The term is approximated by $\bar s,(M_1 \otimes M_2 \otimes M_3 - F_1 \otimes F_2 \otimes F_3)$ with $F_d = c_d c_d^T/(\mathbb 1^T c_d)$, $c_d = M_d \mathbb 1$, which is exact in the eigenbasis of periodic directions (mass_terms in feectools#102). The average couples whole lines and is left out of the diagonal scaling (probing cannot resolve it; its diagonal is O(1/n)), which keeps the preconditioner exact on the logical cube.

WeightedAverageProjection (new, feec/mass.py)

The $\mathbb S$-orthogonal projection onto functions constant in the averaged directions, $\Pi = \mathbb E (\mathbb E^T \mathbb S \mathbb E)^{-1} \mathbb E^T \mathbb S$; then $\mathbb S(\mathbb I - \Pi)$ is symmetric. For $\mathbb M^0$ this is the flux-surface average with $\sqrt g$; it equals the plain AverageOperator when the weight of $\mathbb S$ does not depend on the averaged directions (cube, cylinder: identical to round-off). The reduced matrix is assembled once by colored probing and factorized (sparse LU); one application costs one application of $\mathbb S$.

Why: with the plain (unweighted) AverageOperator, the toroidal stabilization $\mathbb M^0 (\mathbb I - \mathbb A)$ is not symmetric ($\sqrt g = rR$ depends on $\eta_2$; $|x^T S y - y^T S x| / |x^T S y| \approx 0.3$), but PoissonAdiabaticGyrokinetic solves it with pcg. Toroidal results change (the average is now weighted with $\sqrt g$, and with $n/T$ of M0ad_withT, a flux function for the usual equilibria); cylindrical results are unchanged.

ImplicitDiffusion, PoissonSolve, PoissonAdiabaticGyrokinetic

• New option stab_average; ImplicitDiffusion builds stab_mat @ (I - WeightedAverageProjection) before the preconditioners and passes diffusion_mat, stab_mat and stab_average to the stiffness preconditioner. precond="MultiGrid" is rejected with stab_average.
• PoissonAdiabaticGyrokinetic: stab_average follows from which_geometry (cylindrical: "eta3", toroidal: "eta2 eta3") instead of patching _stab_mat after allocate(); multigrid rejected.

Results

PCG iterations, PoissonAdiabaticGyrokinetic (M1gyro, M0ad_withT, epsilon = 0.1, 12x12x8 elements):

geometry	no preconditioner	stiffness preconditioner
cylindrical (HollowCylinder)	34	10
toroidal (HollowTorus)	109 (429 before, non-symmetric)	20

On the unit cube the preconditioner is exact (2 iterations) for constant anisotropic or perpendicular mid, a scaled stab or stab="Id" (whose average is Euclidean, $F_d = \mathbb 1 \mathbb 1^T/n$), and stab_average "eta3"/"eta2 eta3". stab_average="eta1 eta2 eta3" (global weighted mean) is supported as well.

Tests

• test_stiffness_preconditioner_mid_stab_average (exactness on the cube), test_stiffness_preconditioner_mid_stab_options
• test_weighted_average_projection.py (new file in feec/tests, already in the CI shards): symmetry, projection, equality with AverageOperator (Cuboid, HollowCylinder, HollowTorus x M0, M0ad_withT)
• test_poisson_adiabatic_gyrokinetic_stiffness_preconditioner (cylindrical, toroidal): same solution as without preconditioner, fewer than half the iterations; multigrid rejected

Locally: stiffness, weighted average, gyrokinetic Poisson and Poisson tests pass serially and on 2 ranks; default-params test of DriftKineticElectrostaticAdiabatic passes; feectools kron tests pass serially and on 2 ranks.

🤖 Generated with Claude Code

[Copilot]

🟡 Changes recommended

Identity-weighted averaging is constructed incorrectly, and the new preconditioner API is incompatible with the currently pinned feectools version.

3 open findings

• Update feectools dependency for new keyword arguments · New
• Validate unsupported global stab_average configuration · New
• Identity stabilization uses incorrect mass-weighted factors · New

What changed in this PR

Extends diffusion and gyrokinetic Poisson operators with weighted stabilization averages and generalized stiffness preconditioning.

Changes:

• Adds WeightedAverageProjection for symmetric weighted averaging.
• Generalizes StiffnessPreconditioner for custom diffusion, stabilization, and averaging operators.
• Adds validation and numerical tests for cylindrical and toroidal geometries.

File	Description
src/​struphy/​propagators/​tests/​test_gyrokinetic_poisson.py	Tests gyrokinetic Poisson preconditioning.
src/​struphy/​propagators/​poisson_solve.py	Exposes stabilization averaging options.
src/​struphy/​propagators/​poisson_adiabatic_gyrokinetic.py	Selects geometry-dependent weighted averages.
src/​struphy/​propagators/​implicit_diffusion.py	Builds weighted stabilization and generalized preconditioners.
src/​struphy/​feec/​tests/​test_weighted_average_projection.py	Tests projection, symmetry, and plain-average equivalence.
src/​struphy/​feec/​tests/​test_stiffness_preconditioner.py	Tests generalized stiffness configurations.
src/​struphy/​feec/​preconditioner.py	Implements generalized Kronecker approximations.
src/​struphy/​feec/​mass.py	Adds weighted-average projection.

🧠 Review effort: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.