[nicolominichelli]

Summary

Follow-up of #824. In the low-mode TAE the linear variational models were 5-10x slower per step than LinearMHD (about 170 PCG iterations per Schur solve with the mass-matrix preconditioner). This PR reduces the cost of the q model, adds a Kronecker preconditioner for the shear-Alfven Schur operator with an H1vec velocity, and fixes a weight approximation in StiffnessPreconditioner that also speeds up LinearMHD.

All timings below are 4-step profiles of the TAE examples (production grid 16x16x8, hollow torus with inner radius 0.1, dt = 0.5, 4 MPI ranks).

1. Split sound step for the q model

• VariationalQBEvolve(coupling="bq" | "b" | "q"), a constructor argument. "b" ("q") advances only the coupling of the velocity to B (to q), keeping the other variable and its auxiliary variable fixed, with a 2x2 Newton system. The default "bq" is unchanged.
• VariationalPerturbativeMHD_with_q(split_sound=False). With True it uses two substeps, variat_qb (shear-Alfven/fast) and variat_sound (sound), similar to the shear-Alfven/magnetosonic split of LinearMHD.
  • Each substep conserves the energy, so the split step does too.
  • The model checks the propagator variants and the common gamma of both substeps.
• Why not treat q explicitly as the p model treats the pressure: that would break exact energy conservation, because the quadratic thermal energy q^2/(gamma-1) requires q^{n+1/2} in the momentum equation.
• Effect: the Alfven Schur complement of the split q model is the same as the p model's. The sound solve needs 8-9 PCG iterations (beta ~ 0.002). Time per step: 4.2 s -> 2.5 s.
• The q TAE example uses split_sound=True.

2. Kronecker preconditioner of the shear-Alfven Schur operator (H1vec)

• New options schur_precond ("MassMatrixPreconditioner" default, or "StiffnessPreconditioner") and exb_correction of VariationalPBEvolve and VariationalQBEvolve.
  • The preconditioner targets M^v[rho0] + dt^2/4 T0^T M2 T0, built from the analytic equilibrium in prepare_time_step (alfven_schur_preconditioner in feec/variational_utilities.py).
  • It is used for the linear p model and the Alfven substep of the split q model; anywhere else it falls back to the mass-matrix preconditioner with a warning.
• ShearAlfvenSchurPreconditioner(u_space="H1vec"):
  • basis_ops.Tv and the stabilization Mvn;
  • the E x B auxiliary map into H1vec (contravariant components eps_ijk b1_j d_k phi / sqrt(g)).
  • The default u_space="Hdiv" path is unchanged.
• SchurSolverFull and SchurSolverFull3 get a precond setter.

3. Per-component separable weights in StiffnessPreconditioner (general operators)

• Problem: the weights of mid and stab were approximated by one separable shape common to all velocity components. In curvilinear coordinates the component weights have very different profiles. For the H1vec velocity in the torus, G_22 sqrt(g) ~ r^3, and the common shape overestimates the poloidal mass by up to 34x at the inner wall (13.7x for the LinearMHD velocity).
• Fix: each diagonal block now uses its own separable shape (block_shapes=True, the default for general operators). The blocks are inverted independently (block Jacobi), so this costs nothing.
  • It is exact for separable weights with different profiles per component.
  • In Cartesian geometries all components have the same shape, so nothing changes there.
  • The derivative mode of StiffnessPreconditioner (grad/curl/div stiffness) is unchanged.
• Effect on LinearMHD: condition number of the preconditioned shear-Alfven operator 66 -> 19.3 (Lanczos). PCG iterations 23-25 -> 14-17 at the default tolerance.

Results for the variational models

Alfven Schur solve	p: iterations	p: time/step	split q: Alfven iterations	split q: time/step
mass-matrix preconditioner	163-172	2.2 s	163-171	2.5 s
stiffness, block Jacobi	114-124	1.6 s	114-123	1.9 s
stiffness + E x B	93-119	2.1 s	93-117	2.4 s

Block Jacobi (exb_correction=False) is the fastest variant. The examples still use the default mass-matrix preconditioner; the preconditioner changes only the iteration, not the solution.

Analysis: why the H1vec gain is smaller than for LinearMHD

• Lanczos estimates of kappa(P S) for the H1vec operator: 442 (mass), 412 (block Jacobi), 407 (with E x B).
• The variational and the analytic Schur operators agree to 6.5e-5, so the projected equilibrium is not the cause.
• A dense eigen-decomposition on the production grid shows a cluster of slow modes (78 eigenvalues below 0.05). They are:
  • localized at the inner wall;
  • radial and poloidal components only;
  • mass dominated and nearly incompressible;
  • at poloidal mode numbers near the grid scale (m ~ 6-7 of 16 cells).
• For LinearMHD the E x B correction removes the corresponding near-kernel (kappa 1033 -> 66, and 19.3 with the per-component shapes). For the H1vec velocity it does not, even with an exact solve on the auxiliary space (kappa 196). The image of B0 x grad(phi) apparently does not represent the nearly incompressible H1vec fields near the grid scale. This remains open.

Tests

• New propagators/tests/test_variational_schur_precond.py: both preconditioners give the same solution (linear p, split q); fallback for the unsplit q model.
• New in test_verif_VariationalPerturbativeMHD_with_q.py:
  • energy conservation of the split step (drift 3e-9);
  • first-order splitting error of split vs. unsplit (ratio 0.47-0.48 when halving dt).
• Run (MHD only):
  • stiffness/Kronecker, shear-Alfven, magnetosonic and new preconditioner tests (2 ranks: 68 passed);
  • variational propagator tests (serial);
  • default model tests of LinearMHD and the five variational MHD models (2 ranks);
  • q verification (2 ranks: 15 passed);
  • test_examples for the LinearMHD, p and q low-mode TAE examples.
• Not run: the non-MHD users of StiffnessPreconditioner. They use its derivative mode, which is unchanged.

🤖 Generated with Claude Code