diff --git a/docs/installation.rst b/docs/installation.rst index 917e53290..14189e7b3 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -117,3 +117,32 @@ download the QLKNN dependencies at a location of your choice: To use QLKNN10D , you then need to set ``model_path`` in the ``transport`` section of your TORAX config to the path of the cloned repository. See :ref:`configuration` for more details. + + +(Optional) Install gyaradax +--------------------------- + +The ``gyaradax-ql`` transport model is backed by gyaradax, a pure-JAX +reimplementation of the GKW gyrokinetic solver. gyaradax is an optional +dependency: TORAX runs without it, and the ``gyaradax-ql`` model only becomes +selectable once gyaradax is importable in the same environment. + +install gyaradax directly from GitHub into your TORAX virtual environment: + +.. code-block:: console + + pip install git+https://github.com/gerkone/gyaradax + +or clone the repository and install it: + +.. code-block:: console + + git clone https://github.com/gerkone/gyaradax.git + pip install -e gyaradax + +The optional calibration heads selected via ``cn_calibration_path`` +additionally require ``scikit-learn``, the default ITG saturation rule does not. + +Once gyaradax is installed, set ``model_name: 'gyaradax-ql'`` in the +``transport`` section of your TORAX config. See :ref:`configuration` and the +gyaradax-QL entry under :ref:`physics_models` for the available options. diff --git a/docs/physics_models.rst b/docs/physics_models.rst index 858d2367c..d8fb4ab26 100644 --- a/docs/physics_models.rst +++ b/docs/physics_models.rst @@ -291,6 +291,19 @@ nonlinearity in the PDE system. TORAX currently offers five transport models: use TORAX as the same framework for both ML-surrogate and high-fidelity simulations. + - **gyaradax-QL:** A quasilinear transport model backed by `gyaradax + `_, a pure-JAX reimplementation of the + GKW gyrokinetic solver. At each of a user-chosen set of flux-tube radii + (``rho_match``) it runs a linear gyrokinetic eigenmode solve, applies a + quasilinear saturation rule, and interpolates the resulting fluxes onto the + TORAX face grid. Because it is pure JAX and end-to-end differentiable it can + run inside the ``newton_raphson`` solver without disabling JAX compilation. + gyaradax is an external dependency that is not on PyPI and must be installed + separately from GitHub (``pip install git+https://github.com/gerkone/gyaradax``). + When gyaradax is importable the ``model_name: 'gyaradax-ql'`` transport model + becomes available, otherwise it is silently skipped. An optional calibration + head (``cn_calibration_path``) rescales the saturated fluxes. + For all transport models, optional spatial smoothing of the transport coefficients using a Gaussian convolution kernel is implemented, to improve solver convergence rates, an issue which can arise with stiff transport diff --git a/experiments/run_adiabatic_iterhybrid_predictor_corrector.py b/experiments/run_adiabatic_iterhybrid_predictor_corrector.py new file mode 100644 index 000000000..a19ec2b86 --- /dev/null +++ b/experiments/run_adiabatic_iterhybrid_predictor_corrector.py @@ -0,0 +1,488 @@ +# Copyright 2026 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Generate the PR baseline: TORAX + gyaradax-QL vs TORAX + QLKNN. + +Runs the ITER hybrid predictor-corrector scenario (CHEASE equilibrium) under a +deliberately simplified, like-for-like setup -- only the ion-heat channel is +evolved, for *both* transport models -- and writes the artifacts behind the +PR's "early results": a headline panel, a profile overlay, a JSON summary, and +the raw arrays. + +The gyaradax-QL model here is adiabatic-electrostatic, so only q_i (hence T_i) +is a genuine prediction; the plugin sets q_e = q_i and the particle flux to +zero. Evolving T_i alone for both models keeps the overlay honest. + +Must run under the TORAX venv, which imports TORAX from its source tree: + + /system/user/publicwork/galletti/git/torax/.venv/bin/python \ + experiments/run_adiabatic_iterhybrid_predictor_corrector.py --device 0 +""" + +from __future__ import annotations + +import argparse +import copy +import dataclasses +import json +import logging +import os +import sys +import time + +# This script lives in /experiments/; the repo root is its parent. +# Putting it on sys.path lets `import torax` resolve from the source tree +# regardless of the current working directory. +TORAX_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DEFAULT_RHO_MATCH = (0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9) +GYARADAX_LABEL, QLKNN_LABEL = "gyaradax-QL", "QLKNN" + +# Clip / patch / domain keys copied from the stock qlknn transport block so both +# models see identical edge handling and the comparison isolates the turbulent +# core flux. +_INHERITED_TRANSPORT_KEYS = ( + "chi_min", "chi_max", "D_e_min", + "apply_inner_patch", "apply_outer_patch", + "chi_i_inner", "chi_e_inner", "chi_i_outer", "chi_e_outer", + "D_e_inner", "D_e_outer", "V_e_inner", "V_e_outer", + "rho_inner", "rho_outer", +) + +# Representative ITG flux-tube for the standalone per-radius diagnostic. +_DIAG_FLUX_TUBE = dict(rlt=8.0, rln=2.5, q=2.0, shat=1.0, eps=0.18) +_GRAD_N_STEPS = 60 # short linear solve for the AD smoke test + +log = logging.getLogger("baseline") + + +@dataclasses.dataclass(frozen=True) +class ExperimentConfig: + """Knobs for one baseline run. Grid defaults match the calibrated Cn head.""" + + out_dir: str + t_final_rel: float + rho_match: tuple[float, ...] + cn_calibration_path: str + run_diagnostics: bool + nvpar: int = 32 + nmu: int = 8 + ns: int = 16 + nkx: int = 43 + nky: int = 16 + ikxspace: int = 5 + n_steps_linear: int = 200 + + @property + def linear_grid(self) -> dict: + return dict( + nvpar=self.nvpar, nmu=self.nmu, ns=self.ns, nkx=self.nkx, + nky=self.nky, ikxspace=self.ikxspace, n_steps_linear=self.n_steps_linear, + ) + + +@dataclasses.dataclass +class Comparison: + """Both simulation outputs plus their wall-clock times.""" + + gyaradax: object + qlknn: object + t_gyaradax: float + t_qlknn: float + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--device", default="0", help="GPU index (CUDA_VISIBLE_DEVICES)") + parser.add_argument("--out-dir", default=os.path.dirname(os.path.abspath(__file__))) + parser.add_argument("--t-final", type=float, default=1.0, + help="seconds of T_i relaxation past t_initial") + parser.add_argument("--rho-match", default=",".join(map(str, DEFAULT_RHO_MATCH)), + help="comma-separated rho_match radii") + parser.add_argument("--cn", default="auto", + help="cn_calibration_path: 'auto', a registry name, " + "'none', or a path to a pickled head") + parser.add_argument("--skip-diagnostics", action="store_true") + return parser.parse_args() + + +def _configure_runtime(args: argparse.Namespace) -> None: + """Pin the GPU and locate TORAX. Must run before JAX is imported.""" + os.environ.setdefault("CUDA_VISIBLE_DEVICES", args.device) + os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") + if TORAX_ROOT not in sys.path: + sys.path.insert(0, TORAX_ROOT) + + +# Runtime is configured at import time because JAX reads CUDA_VISIBLE_DEVICES on +# first import; the heavy imports below must follow it. +_ARGS = _parse_args() +_configure_runtime(_ARGS) + +import jax # noqa: E402 +import matplotlib # noqa: E402 + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 + +import torax # noqa: E402 +from torax.examples import iterhybrid_predictor_corrector # noqa: E402 +from torax._src.transport_model import pydantic_model # noqa: E402 + + +def _evolve_ion_heat_only(config_dict: dict, t_final_rel: float) -> dict: + """Restrict a TORAX config to a fixed-duration ion-heat-only relaxation.""" + numerics = config_dict["numerics"] + t_initial = float(numerics.get("t_initial", 0.0)) + numerics.update( + evolve_ion_heat=True, evolve_electron_heat=False, + evolve_density=False, evolve_current=False, + t_final=t_initial + t_final_rel, + ) + return config_dict + + +def build_configs(exp: ExperimentConfig): + """Build the gyaradax-QL and QLKNN ToraxConfigs for the same scenario.""" + base = iterhybrid_predictor_corrector.CONFIG + inherited = {key: base["transport"][key] for key in _INHERITED_TRANSPORT_KEYS} + + gyaradax_dict = _evolve_ion_heat_only(copy.deepcopy(base), exp.t_final_rel) + gyaradax_dict["transport"] = { + "model_name": "gyaradax-ql", + "backend": "jax", + "rho_match": exp.rho_match, + "cn_calibration_path": exp.cn_calibration_path, + **exp.linear_grid, + **inherited, + "rho_min": inherited["rho_inner"], + "rho_max": inherited["rho_outer"], + } + qlknn_dict = _evolve_ion_heat_only(copy.deepcopy(base), exp.t_final_rel) + return ( + torax.ToraxConfig.from_dict(gyaradax_dict), + torax.ToraxConfig.from_dict(qlknn_dict), + gyaradax_dict, + ) + + +def run_timed(torax_config): + """Run one simulation, returning (data_tree, wall_clock_seconds).""" + start = time.perf_counter() + data_tree, _ = torax.run_simulation(torax_config) + return data_tree, time.perf_counter() - start + + +def _final_scalar(data_tree, name: str) -> float: + return float(getattr(data_tree.scalars, name).values[-1]) + + +def _final_profile(data_tree, name: str): + return getattr(data_tree.profiles, name).values[-1] + + +def summarize(comparison, gyaradax_config, gyaradax_dict, exp): + """Collect the headline numbers and run metadata into a JSON-able dict.""" + transport = gyaradax_config.transport.build_transport_model() + + def channel(name: str) -> dict: + g = _final_scalar(comparison.gyaradax, name) + q = _final_scalar(comparison.qlknn, name) + return {"gyaradax_ql": g, "qlknn": q, "rel_delta_pct": 100.0 * (g - q) / q} + + return { + "scenario": "iterhybrid_predictor_corrector (CHEASE, T_i-only, adiabatic-ES)", + "rho_match": list(exp.rho_match), + "t_final_s": float(gyaradax_dict["numerics"]["t_final"]), + "linear_grid": exp.linear_grid, + "cn_head": type(transport.cn_head).__name__, + "cn_scalar": float(transport.cn_scalar), + "Q_fusion": channel("Q_fusion"), + "tau_E": channel("tau_E"), + "wallclock_s": { + "gyaradax_ql": comparison.t_gyaradax, + "qlknn": comparison.t_qlknn, + }, + "jax_devices": [str(d) for d in jax.devices()], + } + + +def save_arrays(out_dir: str, comparison) -> None: + """Persist the arrays behind the figures for downstream re-plotting.""" + g, q = comparison.gyaradax, comparison.qlknn + arrays = { + "rho_cell": g.profiles.rho_norm.values, + "rho_face": g.profiles.rho_face_norm.values, + "time_gyaradax": g.scalars.time.values, + "time_qlknn": q.scalars.time.values, + "Q_fusion_gyaradax": g.scalars.Q_fusion.values, + "Q_fusion_qlknn": q.scalars.Q_fusion.values, + } + for key in ("T_i", "T_e", "n_e", "chi_turb_i"): + arrays[f"{key}_gyaradax"] = _final_profile(g, key) + arrays[f"{key}_qlknn"] = _final_profile(q, key) + np.savez(os.path.join(out_dir, "profiles.npz"), **arrays) + + +def _overlay(ax, x, comparison, key, ylabel, title=None) -> None: + """Overlay the QLKNN reference and the gyaradax-QL prediction for one field.""" + ax.plot(x, _final_profile(comparison.qlknn, key), "C0-", lw=2.2, label=QLKNN_LABEL) + ax.plot(x, _final_profile(comparison.gyaradax, key), "C3--", lw=2.2, label=GYARADAX_LABEL) + ax.set_xlabel(r"$\rho_{\mathrm{norm}}$") + ax.set_ylabel(ylabel) + ax.grid(alpha=0.3) + if title: + ax.set_title(title, fontsize=11) + + +def plot_profiles(out_dir: str, comparison) -> None: + """Four-panel final-profile overlay (chi_i, T_i, T_e, n_e).""" + rho_face = comparison.gyaradax.profiles.rho_face_norm.values + rho_cell = comparison.gyaradax.profiles.rho_norm.values + panels = [ + ("chi_turb_i", r"$\chi_i$ [m$^2$/s]", rho_face), + ("T_i", r"$T_i$ [keV]", rho_cell), + ("T_e", r"$T_e$ [keV]", rho_cell), + ("n_e", r"$n_e$ [$10^{20}$ m$^{-3}$]", rho_cell), + ] + fig, axes = plt.subplots(2, 2, figsize=(11, 7.5)) + for ax, (key, ylabel, x) in zip(axes.ravel(), panels): + _overlay(ax, x, comparison, key, ylabel) + ax.legend(fontsize=9) + fig.tight_layout() + fig.savefig(os.path.join(out_dir, "profiles.png"), dpi=140, bbox_inches="tight") + plt.close(fig) + + +def plot_headline(out_dir: str, comparison) -> None: + """Headline row: ion temperature, ion heat diffusivity, fusion gain.""" + g, q = comparison.gyaradax, comparison.qlknn + rho_cell = g.profiles.rho_norm.values + rho_face = g.profiles.rho_face_norm.values + + fig, axes = plt.subplots(1, 3, figsize=(15, 4.5)) + _overlay(axes[0], rho_cell, comparison, "T_i", r"$T_i$ [keV]", "Ion temperature") + axes[0].legend(fontsize=10) + _overlay(axes[1], rho_face, comparison, "chi_turb_i", + r"$\chi_{\mathrm{turb},i}$ [m$^2$/s]", "Ion heat diffusivity") + + ax = axes[2] + ax.plot(q.scalars.time.values, q.scalars.Q_fusion.values, "C0-o", lw=2.2, ms=5) + ax.plot(g.scalars.time.values, g.scalars.Q_fusion.values, "C3--s", lw=2.2, ms=5) + ax.set_xlabel("t [s]") + ax.set_ylabel(r"$Q_{\mathrm{fusion}}$") + ax.set_title("Fusion gain", fontsize=11) + ax.grid(alpha=0.3) + + fig.tight_layout() + fig.savefig(os.path.join(out_dir, "headline.png"), dpi=160, bbox_inches="tight") + plt.close(fig) + + +def save_datatrees(out_dir: str, comparison) -> None: + """Persist both runs as NetCDF so they can be re-plotted without re-running.""" + comparison.gyaradax.to_netcdf(os.path.join(out_dir, "gyaradax_ql.nc")) + comparison.qlknn.to_netcdf(os.path.join(out_dir, "qlknn.nc")) + + +def plot_torax_native(out_dir: str, comparison) -> None: + """Save TORAX's own multipanel comparison (overview + transport plot configs). + + `plot_run_from_data_tree` returns a plotly figure; HTML is always written, + PNG only if `kaleido` is installed. + """ + from torax.plotting.configs import default_plot_config, transport_plot_config + + data_trees = {GYARADAX_LABEL: comparison.gyaradax, QLKNN_LABEL: comparison.qlknn} + panels = (("torax_overview", default_plot_config.PLOT_CONFIG), + ("torax_transport", transport_plot_config.PLOT_CONFIG)) + for name, plot_config in panels: + fig = torax.plot_run_from_data_tree( + plot_config=plot_config, data_trees=data_trees, interactive=False, + fig_title="ITER hybrid: gyaradax-QL vs QLKNN") + fig.write_html(os.path.join(out_dir, f"{name}.html")) + try: + fig.write_image(os.path.join(out_dir, f"{name}.png"), scale=2) + except Exception as exc: # noqa: BLE001 - kaleido is an optional dependency + log.info("%s.png skipped (%s: %s); HTML written", + name, type(exc).__name__, exc) + + +def plot_profile_evolution(out_dir: str, comparison) -> None: + """2D (rho, time) heatmaps of the time-evolving channels for each model. + + Only T_i and chi_turb,i vary in time in the T_i-only setup (T_e / n_e are + prescribed). This is the static analogue of TORAX's interactive time slider. + """ + g, q = comparison.gyaradax, comparison.qlknn + grids = {"cell": g.profiles.rho_norm.values, "face": g.profiles.rho_face_norm.values} + channels = (("T_i", "cell", r"$T_i$ [keV]"), + ("chi_turb_i", "face", r"$\chi_{\mathrm{turb},i}$ [m$^2$/s]")) + models = ((QLKNN_LABEL, q), (GYARADAX_LABEL, g)) + + fig, axes = plt.subplots(len(channels), len(models), figsize=(11, 7), squeeze=False) + for row, (key, grid, clabel) in enumerate(channels): + rho = grids[grid] + fields = [(label, getattr(m.profiles, key)) for label, m in models] + vmax = float(np.percentile( + np.concatenate([f.values.ravel() for _, f in fields]), 98)) + for col, (label, field) in enumerate(fields): + t = field["time"].values + z = field.values + if z.shape != (t.size, rho.size): + z = z.T + ax = axes[row, col] + mesh = ax.pcolormesh(rho, t, z, shading="auto", cmap="magma", vmin=0, vmax=vmax) + ax.set_xlabel(r"$\rho_{\mathrm{norm}}$") + ax.set_ylabel("t [s]") + ax.set_title(f"{label}: {clabel}", fontsize=10) + fig.colorbar(mesh, ax=ax) + fig.tight_layout() + fig.savefig(os.path.join(out_dir, "profile_evolution.png"), dpi=140, bbox_inches="tight") + plt.close(fig) + + +def per_radius_diagnostic(gyaradax_config) -> dict: + """Run one linear gyaradax solve at a representative ITG flux-tube. + + Exercises the plugin's per-radius path directly to report the per-ky growth + rate spectrum and the calibrated QL heat flux for a known operating point. + """ + import jax.numpy as jnp + from gyaradax.integrals import calculate_fluxes, geom_tensors + from gyaradax.params import GKParams + from gyaradax.quasilinear.saturation import ql_flux_diagnostics + from gyaradax.solver import default_state, gksolve, linear_precompute + from torax._src.transport_model.gyaradax_based_transport_model import gyaradax_geometry_at + from torax._src.transport_model.gyaradax_ql_transport_model import GyaradaxQLTransportModel + + transport = gyaradax_config.transport + model = GyaradaxQLTransportModel.from_config(transport) + params = dataclasses.replace( + GKParams(), beta=jnp.asarray(0.0), + adiabatic_electrons=True, non_linear=False, disable_per_ky_norm=True, dt=0.005, + **{k: jnp.asarray(v) for k, v in _DIAG_FLUX_TUBE.items()}) + geom = gyaradax_geometry_at(q=params.q, shat=params.shat, eps=params.eps, + config=transport, topology=model.topology) + df, (phi, _), state = gksolve( + model._initial_df(), geom, params, default_state(nky=transport.nky), + n_steps=transport.n_steps_linear, pre=linear_precompute(geom, params)) + df.block_until_ready() + + _, eflux_kxy, _ = calculate_fluxes(geom_tensors(geom), df, phi, reduce=False) + weights = jnp.asarray(geom["ints"]) + phi2 = jnp.abs(phi) ** 2 + little_g = jnp.asarray(geom["little_g"]) + diag = ql_flux_diagnostics( + growth_rate=state.last_growth_rate, phi2=phi2, + phi2_kxy=jnp.sum(phi2 * weights[:, None, None], axis=0), flux_kxy=eflux_kxy, + krho=jnp.asarray(geom["krho"]), kxrh=jnp.asarray(geom["kxrh"]), + little_g=little_g if little_g.shape[0] == 3 else little_g.T, ds=jnp.mean(weights)) + return { + "flux_tube": _DIAG_FLUX_TUBE, + "gamma_ky": np.asarray(state.last_growth_rate).tolist(), + "Q_QL": float(diag["Q_QL"]), + } + + +def gradient_check(gyaradax_config) -> dict: + """Differentiate q_i w.r.t. R/L_T through the pure-JAX plugin path.""" + import jax.numpy as jnp + from gyaradax.params import GKParams + from torax._src.transport_model.gyaradax_based_transport_model import gyaradax_geometry_at + from torax._src.transport_model.gyaradax_ql_transport_model import ( + GyaradaxQLConfig, GyaradaxQLTransportModel) + + # early_stop=False -> the differentiated solve is a plain lax.scan. + config = GyaradaxQLConfig( + rho_match=gyaradax_config.transport.rho_match, backend="jax", + n_steps_linear=_GRAD_N_STEPS, early_stop=False, cn_calibration_path="auto") + model = GyaradaxQLTransportModel.from_config(config) + fixed = {k: jnp.asarray(v) for k, v in _DIAG_FLUX_TUBE.items() if k != "rlt"} + + def qi_of_rlt(rlt): + params = dataclasses.replace( + GKParams(), rlt=rlt, beta=jnp.asarray(0.0), adiabatic_electrons=True, + non_linear=False, disable_per_ky_norm=True, dt=0.005, **fixed) + geom = gyaradax_geometry_at(q=params.q, shat=params.shat, eps=params.eps, + config=config, topology=model.topology) + qi, _, _ = model._gyaradax_ql_at_radius(params, geom) + return qi + + derivative = float(jax.jit(jax.grad(qi_of_rlt))(jnp.asarray(_DIAG_FLUX_TUBE["rlt"]))) + return {"d_qi_d_rlt": derivative, "n_steps_linear": _GRAD_N_STEPS} + + +def run_diagnostics(gyaradax_config) -> dict: + """Optional standalone checks; failures are recorded, not fatal.""" + results = {} + for name, fn in (("per_radius", per_radius_diagnostic), + ("gradient", gradient_check)): + try: + results[name] = fn(gyaradax_config) + except Exception as exc: # noqa: BLE001 - diagnostics must not abort the run + log.warning("%s diagnostic failed: %s: %s", name, type(exc).__name__, exc) + results[name] = {"error": f"{type(exc).__name__}: {exc}"} + return results + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(message)s") + exp = ExperimentConfig( + out_dir=_ARGS.out_dir, + t_final_rel=_ARGS.t_final, + rho_match=tuple(float(x) for x in _ARGS.rho_match.split(",")), + cn_calibration_path=_ARGS.cn, + run_diagnostics=not _ARGS.skip_diagnostics, + ) + os.makedirs(exp.out_dir, exist_ok=True) + + log.info("JAX devices: %s", jax.devices()) + log.info("gyaradax-ql registered: %s", + "GyaradaxQLConfig" in str(pydantic_model.CombinedCompatibleTransportModel)) + + gyaradax_config, qlknn_config, gyaradax_dict = build_configs(exp) + log.info("rho_match=%s t_final=%.2fs (T_i only, both models)", + exp.rho_match, gyaradax_dict["numerics"]["t_final"]) + + log.info("running gyaradax-QL (first JIT pass is a few minutes) ...") + data_g, t_g = run_timed(gyaradax_config) + log.info("gyaradax-QL done in %.1fs; running QLKNN reference ...", t_g) + data_q, t_q = run_timed(qlknn_config) + comparison = Comparison(data_g, data_q, t_g, t_q) + + summary = summarize(comparison, gyaradax_config, gyaradax_dict, exp) + with open(os.path.join(exp.out_dir, "summary.json"), "w") as fh: + json.dump(summary, fh, indent=2) + save_arrays(exp.out_dir, comparison) + plot_profiles(exp.out_dir, comparison) + plot_headline(exp.out_dir, comparison) + save_datatrees(exp.out_dir, comparison) + plot_torax_native(exp.out_dir, comparison) + plot_profile_evolution(exp.out_dir, comparison) + log.info("\n%s", json.dumps(summary, indent=2)) + + if exp.run_diagnostics: + log.info("running diagnostics (per-radius QL solve + AD smoke test) ...") + diagnostics = run_diagnostics(gyaradax_config) + with open(os.path.join(exp.out_dir, "diagnostics.json"), "w") as fh: + json.dump(diagnostics, fh, indent=2) + log.info("%s", json.dumps(diagnostics, indent=2)) + + log.info("artifacts written to %s", exp.out_dir) + + +if __name__ == "__main__": + main() diff --git a/torax/_src/transport_model/gyaradax_based_transport_model.py b/torax/_src/transport_model/gyaradax_based_transport_model.py new file mode 100644 index 000000000..ad3bc567f --- /dev/null +++ b/torax/_src/transport_model/gyaradax_based_transport_model.py @@ -0,0 +1,264 @@ +"""Base class and utils for gyaradax-based transport models. + +Same pattern as `qualikiz_based_transport_model` and `tglf_based_transport_model` +in TORAX: a `RuntimeParams` dataclass, a shared `GyaradaxBasedTransportModel` +parent that subclasses `QuasilinearTransportModel`, and a `_prepare_gyaradax_inputs` +factory that builds the per-face physics inputs (drives + geometry) for +gyaradax's linear solver. + +The subclass (`GyaradaxQLTransportModel`) implements +`_per_radius(params, geom) -> (qi, qe, pfe)` and inherits the vmap-over-rho_match ++ interp-onto-face-grid call_implementation. +""" + +import abc +import dataclasses +import math +from functools import lru_cache +from typing import Any, Dict, Tuple + +from gyaradax.geometry import build_topology +from gyaradax.geometry import compute_continuous_geometry +from gyaradax.params import GKParams +import jax +import jax.numpy as jnp +from torax._src import state +from torax._src.config import runtime_params as runtime_params_lib +from torax._src.geometry import geometry as geometry_lib +from torax._src.pedestal_model import pedestal_model_output as pedestal_model_output_lib +from torax._src.physics import psi_calculations +from torax._src.transport_model import quasilinear_transport_model +from torax._src.transport_model import runtime_params as transport_runtime_params_lib +from torax._src.transport_model import transport_model as transport_model_lib +from torax._src.transport_model.quasilinear_transport_model import calculate_chiGB +from torax._src.transport_model.quasilinear_transport_model import NormalizedLogarithmicGradients +from torax._src.transport_model.quasilinear_transport_model import QuasilinearInputs +from torax._src.transport_model.quasilinear_transport_model import QuasilinearTransportModel + +# safe-operating clip ranges (same philosophy as QLKNN clip_inputs) +_RLT_MIN, _RLT_MAX = 0.0, 30.0 +_RLN_MIN, _RLN_MAX = -15.0, 15.0 +_Q_MIN, _Q_MAX = 0.5, 10.0 +_SHAT_MIN, _SHAT_MAX = -3.0, 6.0 +_EPS_MIN, _EPS_MAX = 0.02, 0.5 + +# gyaradax geometry / linear-solve defaults +_GKPARAMS_DT = 0.005 +_VPAR_MAX = 3.0 +_KRHOMAX = 1.4 +_RREF = 100.0 +_SIGNB = 1.0 +_NPERIOD = 1 +_KXMAX = 0.0 +_GEOM_TYPE = "circ" + +# initial-df seed amplitude (cosine along s, every non-zonal ky) +_DF_SEED_AMPLITUDE = 1e-3 + +# gyaradax gyrobohm flux unit is 2*sqrt(2) larger than TORAX +_GYARADAX_GB_FLUX_FACTOR = 2.0 * math.sqrt(2.0) + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class RuntimeParams(quasilinear_transport_model.RuntimeParams): + """Runtime parameters shared by all gyaradax-based transport models.""" + + +@lru_cache(maxsize=8) +def _get_topology_cached(nkx: int, nky: int, ikxspace: int, ns: int): + """Topology dict keyed on static grid sizes (cached across `from_config`).""" + return build_topology(nkx=nkx, nky=nky, ikxspace=ikxspace, ns=ns) + + +def build_quasilinear_inputs(core_profiles, geo) -> QuasilinearInputs: + """Build TORAX's QuasilinearInputs on the full face grid.""" + # gyaradax normalizes the gyrobohm flux to the MAJOR radius + chi_gb = calculate_chiGB( + reference_temperature=core_profiles.T_i.face_value(), + reference_magnetic_field=geo.B_0, + reference_mass=core_profiles.A_i, + reference_length=geo.R_major, + ) + log_grads = NormalizedLogarithmicGradients.from_profiles( + core_profiles=core_profiles, + radial_coordinate=geo.r_mid, + radial_face_coordinate=geo.r_mid_face, + reference_length=geo.R_major, + ) + return QuasilinearInputs( + chiGB=chi_gb, + Rmaj=geo.R_major, + Rmin=geo.a_minor, + lref_over_lti=log_grads.lref_over_lti, + lref_over_lte=log_grads.lref_over_lte, + lref_over_lne=log_grads.lref_over_lne, + lref_over_lni0=log_grads.lref_over_lni0, + lref_over_lni1=log_grads.lref_over_lni1, + ) + + +def face_indices_for_radii(geo, rho_match: Tuple[float, ...]) -> jnp.ndarray: + """Pick the face index closest to each rho_match value. Shape (K,).""" + rho_face = geo.rho_face_norm + return jnp.argmin( + jnp.abs(rho_face[:, None] - jnp.asarray(rho_match)[None, :]), axis=0 + ) + + +def gkparams_for_radius( + rho_idx, + ql_inputs: QuasilinearInputs, + core_profiles, + geo, + config, +) -> GKParams: + """Build a GKParams instance for a single flux-tube radius.""" + rlt = jnp.clip(ql_inputs.lref_over_lti[rho_idx], _RLT_MIN, _RLT_MAX) + rln = jnp.clip(ql_inputs.lref_over_lne[rho_idx], _RLN_MIN, _RLN_MAX) + q = jnp.clip(core_profiles.q_face[rho_idx], _Q_MIN, _Q_MAX) + smag_face = psi_calculations.calc_s_rmid(geo, core_profiles.psi) + shat = jnp.clip(smag_face[rho_idx], _SHAT_MIN, _SHAT_MAX) + eps = jnp.clip(geo.epsilon_face[rho_idx], _EPS_MIN, _EPS_MAX) + beta = jnp.asarray(0.0) + backend = getattr(config, "backend", "jax") + return GKParams( + rlt=rlt, + rln=rln, + q=q, + shat=shat, + eps=eps, + beta=beta, + adiabatic_electrons=True, + non_linear=False, # subclass may flip via dataclasses.replace + disable_per_ky_norm=True, # subclass may flip via dataclasses.replace + dt=_GKPARAMS_DT, + backend=backend, + ) + + +def gyaradax_geometry_at( + q, shat, eps, config, topology: Dict[str, Any] +) -> Dict[str, Any]: + """Build a gyaradax geometry dict at one radius (jit/AD safe over q,shat,eps).""" + return compute_continuous_geometry( + q=q, + shat=shat, + eps=eps, + ns=config.ns, + nkx=config.nkx, + nky=config.nky, + nvpar=config.nvpar, + nmu=config.nmu, + vpar_max=_VPAR_MAX, + nperiod=_NPERIOD, + kxmax=_KXMAX, + krhomax=_KRHOMAX, + ikxspace=config.ikxspace, + signB=_SIGNB, + Rref=_RREF, + geom_type=_GEOM_TYPE, + topology=topology, + ) + + +def precompute_topology(config) -> Dict[str, Any]: + """Build the static topology dict from grid sizes.""" + return _get_topology_cached( + config.nkx, config.nky, config.ikxspace, config.ns + ) + + +def initial_df(config) -> jnp.ndarray: + """Initial df: cosine in s, every non-zonal ky seeded. Bypasses init_f's .item().""" + nv, nmu, ns, nkx, nky = ( + config.nvpar, + config.nmu, + config.ns, + config.nkx, + config.nky, + ) + s_grid = (jnp.arange(ns) + 0.5) / ns - 0.5 + seed_s = _DF_SEED_AMPLITUDE * (jnp.cos(2 * jnp.pi * s_grid) + 1.0) + df = jnp.zeros((nv, nmu, ns, nkx, nky), dtype=jnp.complex128) + seed = jnp.broadcast_to( + seed_s[None, None, :, None, None] / (nkx * (nky - 1)), + (nv, nmu, ns, nkx, nky), + ) + ky_mask = jnp.arange(nky) > 0 + return df + seed * ky_mask[None, None, None, None, :] + + +@dataclasses.dataclass(kw_only=True, frozen=True, eq=False) +class GyaradaxBasedTransportModel(QuasilinearTransportModel, abc.ABC): + """Shared frozen dataclass + call_implementation for all gyaradax models.""" + + rho_match: Tuple[float, ...] = (0.35, 0.55, 0.75, 0.875) + backend: str = "jax" + nvpar: int = 32 + nmu: int = 8 + ns: int = 16 + nkx: int = 43 + nky: int = 16 + ikxspace: int = 5 + + @property + def topology(self): + return _get_topology_cached(self.nkx, self.nky, self.ikxspace, self.ns) + + def _initial_df(self) -> jnp.ndarray: + return initial_df(self) + + @abc.abstractmethod + def _per_radius( + self, params: GKParams, geom: Dict[str, Any] + ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Compute (qi, qe, pfe) at one radius in gyroBohm units.""" + + def call_implementation( + self, + transport_runtime_params: transport_runtime_params_lib.RuntimeParams, + runtime_params: runtime_params_lib.RuntimeParams, + geo: geometry_lib.Geometry, + core_profiles: state.CoreProfiles, + pedestal_model_output: pedestal_model_output_lib.PedestalModelOutput, + ) -> transport_model_lib.TurbulentTransport: + del pedestal_model_output, runtime_params + + ql_inputs = build_quasilinear_inputs(core_profiles, geo) + match_idx = face_indices_for_radii(geo, self.rho_match) + + def per_radius(idx): + params = gkparams_for_radius(idx, ql_inputs, core_profiles, geo, self) + geom = gyaradax_geometry_at( + q=params.q, + shat=params.shat, + eps=params.eps, + config=self, + topology=self.topology, + ) + return self._per_radius(params, geom) + + qi_m, qe_m, pfe_m = jax.vmap(per_radius)(match_idx) + + rho_face = geo.rho_face_norm + rho_match_arr = jnp.asarray(self.rho_match) + qi_face = jnp.interp(rho_face, rho_match_arr, qi_m) + qe_face = jnp.interp(rho_face, rho_match_arr, qe_m) + pfe_face = jnp.interp(rho_face, rho_match_arr, pfe_m) + + qi_face = qi_face * _GYARADAX_GB_FLUX_FACTOR + qe_face = qe_face * _GYARADAX_GB_FLUX_FACTOR + pfe_face = pfe_face * _GYARADAX_GB_FLUX_FACTOR + + return self._make_core_transport( + qi=qi_face, + qe=qe_face, + pfe=pfe_face, + quasilinear_inputs=ql_inputs, + transport=transport_runtime_params, + geo=geo, + core_profiles=core_profiles, + gradient_reference_length=geo.R_major, + gyrobohm_flux_reference_length=geo.R_major, + ) diff --git a/torax/_src/transport_model/gyaradax_ql_transport_model.py b/torax/_src/transport_model/gyaradax_ql_transport_model.py new file mode 100644 index 000000000..35de2cba2 --- /dev/null +++ b/torax/_src/transport_model/gyaradax_ql_transport_model.py @@ -0,0 +1,324 @@ +"""QL gyaradax as a TORAX transport model. + +Pure-JAX linear gyrokinetic solve + saturation rule + calibration head. +Inherits the vmap/interp/_make_core_transport machinery from +`GyaradaxBasedTransportModel` and only specifies what happens at one radius. +""" + +import dataclasses +from functools import lru_cache +import pickle +from typing import Annotated, Any, Dict, Literal, Optional, Tuple +import warnings + +import chex +from gyaradax.integrals import calculate_fluxes +from gyaradax.integrals import geom_tensors +from gyaradax.params import GKParams +from gyaradax.quasilinear.saturation import ql_flux +from gyaradax.solver import default_state +from gyaradax.solver import gksolve +from gyaradax.solver import linear_precompute +import jax +import jax.numpy as jnp +from torax._src.torax_pydantic import torax_pydantic +from torax._src.transport_model import pydantic_model_base +from torax._src.transport_model.gyaradax_based_transport_model import _get_topology_cached +from torax._src.transport_model.gyaradax_based_transport_model import GyaradaxBasedTransportModel +from torax._src.transport_model.gyaradax_based_transport_model import RuntimeParams as _BaseRuntimeParams + +# nan-guard / clip on per-radius q_i (poisons TORAX Newton otherwise) +_QI_CLIP_ABS = 1e3 + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class RuntimeParams(_BaseRuntimeParams): + """Runtime parameters for the gyaradax-QL transport model.""" + + +@lru_cache(maxsize=8) +def _get_cn_head(path: str): + """Resolve a Cn calibration head (the cn-version parametric/polynomial weights). + + Resolution mirrors fusion_surrogates' name-or-path scheme: + '' / falsy -> no head (the basic-ql scalar Cn is used instead). + 'auto' -> the default head bundled with gyaradax. + a registry name (gyaradax.quasilinear.models.registry.MODELS) -> that + bundled head, resolved by name. + anything else -> a path to a user pickle (calibrate your own via + gyaradax.quasilinear.fit_cn_heads). + Returns the head (with `cn_jax`) or None. + """ + if not path: + return None + if path == "auto": + from gyaradax.quasilinear import load_default_cn_weights + + obj = load_default_cn_weights() + else: + from gyaradax.quasilinear.models import registry + + if path in registry.MODELS: + from gyaradax.quasilinear import load_cn_weights_from_name + + obj = load_cn_weights_from_name(path) + else: + try: + with open(path, "rb") as f: + obj = pickle.load(f) + except FileNotFoundError: + warnings.warn( + f"gyaradax-QL: cn_calibration_path='{path}' is neither a bundled " + "model name nor an existing file; falling back to the basic-ql " + "scalar Cn.", + RuntimeWarning, + stacklevel=2, + ) + return None + if isinstance(obj, dict) and "polynomial" in obj: + return obj["polynomial"] + return obj + + +@lru_cache(maxsize=1) +def _default_cn_scalar() -> float: + """Calibrated scalar Cn bundled with gyaradax (the basic-ql amplitude).""" + try: + from gyaradax.quasilinear import load_default_cn_weights + + return float(load_default_cn_weights().get("scalar", 1.0)) + except Exception: # pylint: disable=broad-except + return 1.0 + + +@dataclasses.dataclass(kw_only=True, frozen=True, eq=False) +class GyaradaxQLTransportModel(GyaradaxBasedTransportModel): + """QL gyaradax transport model.""" + + n_steps_linear: int = 200 + ncv_eigensolve: int = 0 + cn_calibration_path: str = "auto" + cn_scalar: float = 1.0 + # early-stop knobs: skip remaining gksolve steps once per-ky growth rates + # stop moving. solver itself is untouched -- we chunk the call and check + # convergence between chunks via lax.while_loop. + early_stop: bool = True + early_stop_block: int = 25 + early_stop_atol: float = 1e-4 + early_stop_rtol: float = 1e-3 + early_stop_min_steps: int = 50 + + @classmethod + def from_config(cls, cfg) -> "GyaradaxQLTransportModel": + # warm caches outside any jit; build_topology allocates int8 arrays + # that would otherwise become tracers inside torax's jit + _get_topology_cached(cfg.nkx, cfg.nky, cfg.ikxspace, cfg.ns) + _get_cn_head(cfg.cn_calibration_path or "") + return cls( + rho_match=tuple(cfg.rho_match), + backend=cfg.backend, + n_steps_linear=cfg.n_steps_linear, + ncv_eigensolve=cfg.ncv_eigensolve, + nvpar=cfg.nvpar, + nmu=cfg.nmu, + ns=cfg.ns, + nkx=cfg.nkx, + nky=cfg.nky, + ikxspace=cfg.ikxspace, + cn_calibration_path=cfg.cn_calibration_path or "", + cn_scalar=_default_cn_scalar(), + early_stop=cfg.early_stop, + early_stop_block=cfg.early_stop_block, + early_stop_atol=cfg.early_stop_atol, + early_stop_rtol=cfg.early_stop_rtol, + early_stop_min_steps=cfg.early_stop_min_steps, + ) + + @property + def cn_head(self): + return _get_cn_head(self.cn_calibration_path) + + def _per_radius( + self, params: GKParams, geom: Dict[str, Any] + ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + return self._gyaradax_ql_at_radius(params, geom) + + def _linear_with_early_stop(self, df, geom, params, sim_state, pre): + """Chunked linear gksolve with per-ky growth-rate convergence check. + + Runs gksolve in blocks of `early_stop_block` and exits early once the + per-ky growth-rate vector stops moving within (atol, rtol). Pure + wrapper -- never touches the solver internals. Hard-caps the total + iteration count at `n_steps_linear`. Compatible with jit + vmap + (lax.while_loop runs until *all* batch elements converge). + """ + block = int(self.early_stop_block) + max_blocks = max(int(self.n_steps_linear) // block, 1) + min_blocks = max(int(self.early_stop_min_steps) // block, 1) + atol = float(self.early_stop_atol) + rtol = float(self.early_stop_rtol) + + # bootstrap one block so `phi_last` has a concrete dtype/shape to carry + df1, (phi1, _flx), sim1 = gksolve( + df, geom, params, sim_state, n_steps=block, pre=pre + ) + g0 = sim1.last_growth_rate + + def cond(state): + i, _df, _phi, _sim, _prev, conv = state + return jnp.logical_and(i < max_blocks, jnp.logical_not(conv)) + + def body(state): + i, df_c, _phi_c, sim_c, prev_g, _ = state + df_n, (phi_n, _f), sim_n = gksolve( + df_c, geom, params, sim_c, n_steps=block, pre=pre + ) + new_g = sim_n.last_growth_rate + delta = jnp.max(jnp.abs(new_g - prev_g)) + scale = atol + rtol * jnp.max(jnp.abs(new_g)) + conv = jnp.logical_and(i + 1 >= min_blocks, delta <= scale) + return (i + 1, df_n, phi_n, sim_n, new_g, conv) + + init = (jnp.asarray(1), df1, phi1, sim1, g0, jnp.array(False)) + _i, df_f, phi_f, sim_f, _g, _c = jax.lax.while_loop(cond, body, init) + return df_f, phi_f, sim_f + + def _gyaradax_ql_at_radius(self, params, geom): + """Linear gksolve + QL saturation rule + Cn head at one radius.""" + df = self._initial_df() + sim_state = default_state(nky=self.nky) + pre = linear_precompute(geom, params) + if self.early_stop: + df_final, phi, sim_state_final = self._linear_with_early_stop( + df, + geom, + params, + sim_state, + pre, + ) + else: + df_final, (phi, _fluxes), sim_state_final = gksolve( + df, + geom, + params, + sim_state, + n_steps=self.n_steps_linear, + pre=pre, + ) + + gt = geom_tensors(geom) + _pflux, eflux_kxy, _vflux = calculate_fluxes( + gt, df_final, phi, reduce=False + ) + ints = jnp.asarray(geom["ints"]) + ds = jnp.mean(ints) + phi2 = jnp.abs(phi) ** 2 + phi2_kxy = jnp.sum(phi2 * ints[:, None, None], axis=0) + lg = jnp.asarray(geom["little_g"]) + little_g = lg.T if lg.shape[0] != 3 else lg + krho = jnp.asarray(geom["krho"], dtype=jnp.float64) + kxrh = jnp.asarray(geom["kxrh"], dtype=jnp.float64) + gamma = sim_state_final.last_growth_rate + + # FEATURE_NAMES = (rlt_i, rln_i, rlt_e, rln_e, shat, q, eps, beta) + head = self.cn_head + if head is not None and hasattr(head, "cn_jax"): + features = jnp.array([[ + params.rlt, + params.rln, + params.rlt, + params.rln, + params.shat, + params.q, + params.eps, + params.beta, + ]]) + cn = head.cn_jax(features)[0] + else: + cn = jnp.asarray(self.cn_scalar) + + q_i = ql_flux( + growth_rate=gamma, + phi2=phi2, + phi2_kxy=phi2_kxy, + flux_kxy=eflux_kxy, + krho=krho, + kxrh=kxrh, + little_g=little_g, + ds=ds, + cn=cn, + ) + q_i = jnp.where( + jnp.isfinite(q_i), jnp.clip(q_i, -_QI_CLIP_ABS, _QI_CLIP_ABS), 0.0 + ) + # qe = qi, pfe = 0 placeholder (ITG-adiabatic) + return q_i, q_i, jnp.asarray(0.0) + + +class GyaradaxQLConfig(pydantic_model_base.TransportBase): + """Config for the gyaradax-QL transport model. + + Attributes: + model_name: transport model selector. Hardcoded to 'gyaradax-ql'. + rho_match: normalized-radius flux tubes where gyaradax is actually run; + fluxes are interpolated from these onto the full face grid. + backend: gyaradax compute backend, 'jax' (AD-clean) or 'cuda' (no AD). + nvpar: parallel-velocity grid points. + nmu: magnetic-moment grid points. + ns: parallel (field-line) grid points. + nkx: radial wavenumber modes. + nky: binormal wavenumber modes. + ikxspace: kx mode spacing (parallel boundary connection). + n_steps_linear: hard cap on RK4 steps per linear gyaradax run. + ncv_eigensolve: 0 uses the IVP growth rate; >0 uses the JAX-Arnoldi + eigensolver with this many Krylov vectors. + cn_calibration_path: selects the Cn calibration (cn version). 'auto' + (default) uses the head bundled with gyaradax; a registry name + (gyaradax.quasilinear.models.registry.MODELS) selects a named bundled + head; any other value is a path to your own pickled head (produce one + with gyaradax.quasilinear.fit_cn_heads on your (X_QL, Y_NL, F) dataset); + None uses the basic-ql scalar Cn (also bundled, calibrated). + early_stop: stop the linear solve once per-ky growth rates converge. + early_stop_block: gksolve steps per convergence-check block. + early_stop_atol: absolute tolerance on the growth-rate change. + early_stop_rtol: relative tolerance on the growth-rate change. + early_stop_min_steps: minimum steps before early-stop can trigger. + """ + + model_name: Annotated[Literal["gyaradax-ql"], torax_pydantic.JAX_STATIC] = ( + "gyaradax-ql" + ) + + rho_match: Annotated[Tuple[float, ...], torax_pydantic.JAX_STATIC] = ( + 0.35, + 0.55, + 0.75, + 0.875, + ) + backend: Annotated[str, torax_pydantic.JAX_STATIC] = "jax" + + nvpar: Annotated[int, torax_pydantic.JAX_STATIC] = 32 + nmu: Annotated[int, torax_pydantic.JAX_STATIC] = 8 + ns: Annotated[int, torax_pydantic.JAX_STATIC] = 16 + nkx: Annotated[int, torax_pydantic.JAX_STATIC] = 43 + nky: Annotated[int, torax_pydantic.JAX_STATIC] = 16 + ikxspace: Annotated[int, torax_pydantic.JAX_STATIC] = 5 + + n_steps_linear: Annotated[int, torax_pydantic.JAX_STATIC] = 200 + ncv_eigensolve: Annotated[int, torax_pydantic.JAX_STATIC] = 0 + cn_calibration_path: Annotated[Optional[str], torax_pydantic.JAX_STATIC] = ( + "auto" + ) + early_stop: Annotated[bool, torax_pydantic.JAX_STATIC] = True + early_stop_block: Annotated[int, torax_pydantic.JAX_STATIC] = 25 + early_stop_atol: Annotated[float, torax_pydantic.JAX_STATIC] = 1e-4 + early_stop_rtol: Annotated[float, torax_pydantic.JAX_STATIC] = 1e-3 + early_stop_min_steps: Annotated[int, torax_pydantic.JAX_STATIC] = 50 + + def build_transport_model(self) -> "GyaradaxQLTransportModel": + return GyaradaxQLTransportModel.from_config(self) + + def build_runtime_params(self, t: chex.Numeric) -> RuntimeParams: + base_kwargs = dataclasses.asdict(super().build_runtime_params(t)) + return RuntimeParams(DV_effective=True, An_min=0.05, **base_kwargs) diff --git a/torax/_src/transport_model/pydantic_model.py b/torax/_src/transport_model/pydantic_model.py index 58a890e66..11c22043c 100644 --- a/torax/_src/transport_model/pydantic_model.py +++ b/torax/_src/transport_model/pydantic_model.py @@ -442,6 +442,19 @@ def build_runtime_params( | tglf_transport_model.TGLFTransportModelConfig ) +# gyaradax-QL transport model (optional dependency: the gyaradax gyrokinetic solver). +# the model module imports gyaradax at top, so this import fails when gyaradax is +# not installed and the config is simply left unregistered (same as qualikiz/tglf). +try: + from torax._src.transport_model import gyaradax_ql_transport_model # pylint: disable=g-import-not-at-top + + CombinedCompatibleTransportModel = ( + CombinedCompatibleTransportModel + | gyaradax_ql_transport_model.GyaradaxQLConfig + ) +except ImportError: + pass + class CombinedTransportModel(pydantic_model_base.TransportBase): """Model for the Combined transport model.