Repository navigation
Add gyaradax plugin - #2189
Add gyaradax plugin#2189gerkone wants to merge 9 commits into
Conversation
|
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. |
You mean this normalization here? 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).
That sounds fine.
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).
Not sure how relevant the QLKNN dataset is. Fundamentally you need 3 ingredients: 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.
Yes it's interesting, e.g. for ultimate ground-truth validation or for training set generation for surrogates or QL saturation rules. |
Two transport models backed by gyaradax, a pure-JAX reimplementation of the GKW gyrokinetic solver, plus the shared GKW<->TORAX normalization and an optional host-side diagnostics sink. gyaradax-QL runs a linear gyrokinetic eigenmode solve at each rho_match radius, applies a quasilinear saturation rule and interpolates the resulting fluxes onto the face grid. gyaradax-NL reuses the same profile mapping and unit conversion but integrates the full nonlinear turbulence to saturation at each radius; it is for offline ground-truth generation, not transport loops. Both register themselves through register_transport_model, so no upstream file is modified: the only edits outside the new files are the two docs sections describing how to install gyaradax and what the models do. gyaradax is an optional dependency. Without it the modules are not imported and TORAX is unaffected. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Every gyaradax transport model subclassed GyaradaxQLTransportModel, so each one inherited the quasilinear machinery whether or not it used it. The NL model's own docstring said as much -- "the head itself is unused here", "the linear early-stop diagnostics don't apply" -- while still carrying `rule`, `cn_override`, `cn_scalar` and six `early_stop_*` fields. The gyroflow model in neugk-jax inherited the same set, plus `n_steps_linear`. None of that was quasilinear in the first place. `call_implementation` -- build the quasilinear inputs, map rho_match onto face indices, vmap over radii, convert GKW gyroBohm fluxes to TORAX units -- contains no QL physics, and neither do the grid fields, the per-radius GKParams/geometry construction, `topology` or `sink`. They now live on BaseGKWGyaradaxPlugin, and QL, NL and gyroflow are siblings that supply only a flux model, by overriding `_per_radius` or `_per_radius_from_profiles`. Grid resolution no longer reads the Cn calibration head. `_resolve_grid` takes a fallback grid and each subclass provides one through `_grid_fallback`, so the base does not depend on QL calibration. NL keeps its head-derived grid defaults, now declared for what they are -- grid metadata only, no Cn is applied in a nonlinear run -- so NL configs that leave grid fields unset behave as before. The config needed two tiers rather than one. BaseGKWGyaradaxConfig carries what all three share (rho_match, DV_effective, An_min, diagnostics_path); SolverGKWGyaradaxConfig adds the grid, backend, mixed_precision and em for the models that actually run the solver. Collapsing them into one would have handed gyroflow seven fields it ignores, since its grid is fixed by the GyroSwin training data -- the same problem in a new place. Incidentally fixes a leak the old layout hid: the QL model's `sink` reached for gyroflow's `latent_dump_*` through getattr, so the QL base was already quietly gyroflow-aware. That hook is now on the shared base, where it is honest. QL goes from 512 to 175 lines; the base is 424. Two tests referenced `face_indices_for_radii` and `gkparams_for_radius` on the QL module and now use the base. gyaradax_transport_model_test: 14 passed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Merging 41 commits of upstream into qlgyx applies cleanly as text but breaks the gyaradax plugins at runtime, in two places. The neoclassical model was renamed from plural to singular upstream: `neoclassical.build_models()` is now `build_model()`, `initial_core_profiles` takes `neoclassical_model=`, and `build_source_profiles` no longer takes a neoclassical argument at all (it takes optional conductivity and bootstrap_current instead). Only the test helper touched any of this; no plugin references neoclassical. `quasilinear_transport_model.RuntimeParams` gained a required `DV_effective_smooth_width`, which blends the effective D/V split for small down-gradient flux. BaseGKWGyaradaxConfig now declares it with the same default upstream uses for its own quasilinear configs (0.01) and passes it through build_runtime_params. gyaradax_transport_model_test: 14 passed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
gkparams_for_radius grows a kinetic branch: rlt = [lti, lte], one rln = lne shared by both species (the data are pure D and 2-species gyaradax needs gradient quasineutrality, so n_e rather than n_i0), tmp = [1, Te/Ti] against the GKW ion reference, mas = [1, 2.72e-4], signz = [1, -1], and vthrat = sqrt(tmp/mas), which the species precompute needs per species rather than as the scalar default. The branch is off by default and refuses to build without an explicit cn_override: no Cn head fitted on kinetic data exists yet, and silently reusing the adiabatic one would be wrong rather than merely uncalibrated. The QL rule docstring named six rules that no longer exist; the set is now canonical / sat0 / sat2, with sat2 reproducing GACODE TGLF SAT2's intensity spectrum on the full linear eigenmode. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Upstream removed the shared quasilinear base: quasilinear_transport_model is now quasilinear_utils and keeps only the helpers, while QuasilinearInputs, QuasilinearTransportModel and the shared RuntimeParams are gone, each model carrying its own (see TGLFInputs). The gyaradax plugins follow suit. QuasilinearInputs moves into gyaradax_normalization, which is the lower-level of the two modules and already builds it. The base class derives from component.ComponentTransportModel directly, RuntimeParams declares DV_effective / An_min / DV_effective_smooth_width itself, and _make_core_transport is vendored unchanged onto the base. PedestalTransitionState.empty_L_mode is now empty(), defaulting to L mode. Verified on the merge: gyaradax 16, transport_model 28, pydantic_model 39, quasilinear_utils 17, all passing. Needs jax >= 0.11.1, as upstream does. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…tion Nothing about the linear solve should need a source edit to change. The settings that were module constants are now config fields defaulting to the values the calibration scans used: vpar_max, nperiod, geom_type, gkparams_dt, rref, electron_mass_ratio, and the clip box. Grids come from a named preset -- adiabatic 85x32x16, kinetic 192x32x18 -- rather than from whatever metadata a Cn head carries. The bundled head still advertises 43x16, which is under the resolution floor, so inheriting from it was never a safe default. Explicit config fields still win over a preset. rule_options forwards arbitrary keywords to the selected rule, so sat2's flux_exponent and kx_width_model, and anything a future rule adds, are reachable from a TORAX config without touching the plugin. ikxspace can vary per rho_match radius. It sizes the mode connectivity and so has to be static, while shat is traced and evolves, so the values are fixed at build time from a reference profile rather than tracking the live shear -- ikxspace_for_shat implements floor(6.53*|shat|), the rule the stored runs follow. Radii then need one topology and one compilation each, so that path unrolls instead of sharing a vmap. Five tests cover the new surface, including that every name in rule_options is a real parameter of the selected rule, and that every rule registered in gyaradax is selectable. 21 pass. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
gyaradax builds its kx grid symmetric about kx = 0, so an even nkx mismatches the field it allocates. The kinetic preset said 192 while every stored kinetic run is 191, and the mismatch surfaced only at the first transport call as a bare shape error inside the solver. The preset is now 191 and _resolve_grid rejects any even nkx with a message naming the fix, so a bad grid fails at config time. Tests cover both, plus a vmap regression for dlnpdr: it is evaluated per radius inside the vmap and had been using Python min/max on tracers. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Is this something that you re still working on? If so what about doing something like https://github.com/aaronkho/epednn_mit/tree/main/src/epednn_mit and using the registry while you experiment etc. Then it can live in your own repo but people can still install etc as needed? If you do let us know as we are planning a docs page of all the registered models that are open-source and we can link to this here. Feel free to still use this for discussion etc in the meantime |
|
@tamaranorman thanks, good to know. |
3 participants
hey Torax team
this PR adds the
gyaradax-qlplugin (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-qlmodel: per-radius linear solve => QL saturation => flux, over a few equidistantrho_matchthen interpolated.gyaradax-qlon ~1000 runs (rlt, rln, q, shat scanned, eps scan left out!, beta=0)iterhybrid_predictor_correctorcase, ES adiabatic), below.gyaradax).Open questions
rho_matchsubset and interpolates.Early results (
iterhybrid_predictor_corrector)rho_matchradii, interpolated to the face grid.t_final1s (was 5s).