Add gyaradax plugin

[gerkone]

hey Torax team
this PR adds the gyaradax-ql plugin (quasilinear model based on the gyaradax local gyrokinetics code) as an optional transport model. this is still WIP and more results to follow, mostly here to get feedback.

• gyaradax-ql model: per-radius linear solve => QL saturation => flux, over a few equidistant rho_match then interpolated.
• basic calibration of gyaradax-ql on ~1000 runs (rlt, rln, q, shat scanned, eps scan left out!, beta=0)
• early results on ITER hybrid (simplified iterhybrid_predictor_corrector case, ES adiabatic), below.
• proper validation on the ITER case. requires: broader NL saturation scan for calibration (specifically: eps scan, beta scan).
• (optional) nonlinear transport variant (using full NL gyaradax).
• (gyaradax side) validate electromagnetic Bpar and Apar + Bpar.
• (gyaradax side) investigate high beta instabilities (Apar + Bpar).

Open questions

• TORAX normalization: TORAX flux minor-radius, gyaradax major-radius + differences in the velocity normalization (should be taken care of).
• radii: QLKNN runs a batch, gyaradax-QL currently evals a rho_match subset and interpolates.
• what can we do with AD? advantages in transport modeling?
• tips on how to calibrate the QL model. would the QLKNN dataset be helpful?
• what is expected in validation?
• is the NL model relevant or interesting? QL is already ~100x slower to QLKNN

Early results (iterhybrid_predictor_corrector)

• calibration set is adiabatic-ES, no eps scan. does this explain the discrepancies?
• gyaradax-QL evaluated at 8 rho_match radii, interpolated to the face grid.
• t_final 1s (was 5s).
• ideas on the diffusivity spike at ~0.25?

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[jcitrin]

TORAX normalization: TORAX flux minor-radius, gyaradax major-radius + differences in the velocity normalization (should be taken care of).

You mean this normalization here?

torax/torax/_src/transport_model/quasilinear_transport_model.py

Line 407 in 0d9546f

def _make_core_transport(

You can choose the reference lengths for chiGB and the heat flux in these methods, so TORAX flux is not fixed to be normalized to minor radius (unless I misunderstood something).

radii: QLKNN runs a batch, gyaradax-QL currently evals a rho_match subset and interpolates.

That sounds fine.

what can we do with AD? advantages in transport modeling?

e.g. flux matching with nonlinear solvers will be faster (once we have a stationary state flux matching solver in TORAX). And for time-dependent can use Newton in the inner loop. And of course can build optimization wrappers around all these for various constrained optimization applications (e.g. scenario optimization).

tips on how to calibrate the QL model. would the QLKNN dataset be helpful?
what is expected in validation?

Not sure how relevant the QLKNN dataset is. Fundamentally you need 3 ingredients:
i) A fast linear model for growth rates, QL weights, mode structures
ii) Analytical priors for the saturation rule
iii) A database of NL simulations to tune the saturation rule against.

There is a lot of literature. See e.g. https://iopscience.iop.org/article/10.1088/1741-4326/ad6ba5/meta . fcasson knows a lot about this.

is the NL model relevant or interesting? QL is already ~100x slower to QLKNN

Yes it's interesting, e.g. for ultimate ground-truth validation or for training set generation for surrogates or QL saturation rules.