Faster setup of the Hiptmair-Xu preconditioner (TAE example): exact stiffness diagonal, preconditioners built in the setup

[spossann]

Shortens the setup of the ITPA TAE example with ShearAlfvenCurlCurl (Hiptmair-Xu preconditioner). Bottlenecks were found with scope-profiler regions in the HX setup (first two commits).

TAE --test-run, 2 ranks, one time step: profiled session per rank 7.6 s → 3.0 s; HX build 3.8 s → 0.5 s, same PCG iterations.

Changes

• LocalProjectionMatrix / LocalRotationMatrix evaluate each component once (were 18× / 6×). Weight evaluation of M2para_Bn: 0.97 → 0.18 s/rank.
• Exact diagonal of dᵀ M d (_stiffness_local_diagonal in feec/preconditioner.py), used for the diagonal scaling of StiffnessPreconditioner, instead of probing with colored unit vectors (about 160 colors per block, each a full operator application). It is read from the band of the mid mass matrix and the 1d factors of the derivative; rows owned by the neighbouring process come from ghost regions, and band entries wrapping around in short periodic directions are summed. Probing remains the fallback for matrix-free mass operators. Without diagonal scaling HX needs 10–25 % more iterations (6 vs 8 at dt = 0.1, 24 vs 26 at dt = 1), so the scaling is kept.
• StiffnessPreconditioner resolves the default mid M{k+1} only if mid is not given (the HX smoother no longer assembles M2).
• New hook Propagator.prepare_time_step(dt), called by Simulation.run after allocate (profiling region setup: prepare time step) with the dt each propagator gets from integrate (dt/2 for all but the last with Strang). ShearAlfvenCurlCurlPropagator builds HX there and ShearAlfvenPropagator its multigrid / Kronecker preconditioner, instead of in the first time step. __call__ still rebuilds when dt changes.
• scope-profiler upper bound 0.7.3.

Tests

• New test_stiffness_local_diagonal: exact vs probed diagonal for grad/curl/div on a HollowTorus, a sum of mass terms, Dirichlet bcs and short periodic directions with p = 3.
• New test_prepare_time_step: HX is built ahead of time, reused for the same dt and rebuilt for a new one.
• Locally on 1 and 2 ranks: stiffness preconditioner, Hiptmair-Xu, Kronecker block, preconditioner transpose, ShearAlfven, ShearAlfvenCurlCurl and curl-curl propagator tests pass. A separate review also checked the exact diagonal on 3 ranks (each decomposition direction, degrees 1 and 3, M2gyro / M1perp mids, identity and summed stabilizations).

Not included: assembling M2gyro in one pass instead of M2_Bn - M2para_Bn. Mass matrix assembly (about 1.2 s/rank) is now the largest remaining setup cost.

🤖 Generated with Claude Code