Generalize the transport_model _make_core_transport method.

torax/transport_model/qlknn_wrapper.py

@@ -356,6 +356,8 @@ def _combined(
         transport=runtime_config_inputs.transport,
         geo=geo,
         core_profiles=core_profiles,
+        gradient_reference_length=geo.Rmaj,
+        gyrobohm_flux_reference_length=geo.Rmin,
     )
 
   def __hash__(self) -> int:

torax/transport_model/qualikiz_based_transport_model.py

@@ -26,6 +26,7 @@
 @chex.dataclass
 class RuntimeParams(quasilinear_transport_model.RuntimeParams):
   """Shared parameters for Qualikiz-based models."""
+
   # Collisionality multiplier.
   coll_mult: float = 1.0
   # ensure that smag - alpha > -0.2 always, to compensate for no slab modes
@@ -44,6 +45,7 @@ def make_provider(
 @chex.dataclass(frozen=True)
 class DynamicRuntimeParams(quasilinear_transport_model.DynamicRuntimeParams):
   """Shared parameters for Qualikiz-based models."""
+
   coll_mult: float
   avoid_big_negative_s: bool
   smag_alpha_correction: bool
@@ -66,8 +68,6 @@ class QualikizInputs(quasilinear_transport_model.QuasilinearInputs):
   """Inputs to Qualikiz-based models."""
 
   Zeff_face: chex.Array
-  Ani0: chex.Array
-  Ani1: chex.Array
   q: chex.Array
   smag: chex.Array
   x: chex.Array
@@ -95,83 +95,30 @@ def _prepare_qualikiz_inputs(
     """Prepare Qualikiz inputs."""
     constants = constants_module.CONSTANTS
 
-    Rmin = geo.Rmin
-    Rmaj = geo.Rmaj
-
     # define radial coordinate as midplane average r
     # (typical assumption for transport models developed in circular geo)
     rmid = (geo.Rout - geo.Rin) * 0.5
     rmid_face = (geo.Rout_face - geo.Rin_face) * 0.5
 
-    temp_ion_var = core_profiles.temp_ion
-    temp_ion_face = temp_ion_var.face_value()
-    temp_ion_face_grad = temp_ion_var.face_grad(rmid)
-    temp_el_var = core_profiles.temp_el
-    temp_electron_face = temp_el_var.face_value()
-    temp_electron_face_grad = temp_el_var.face_grad(rmid)
-    # Careful, these are in n_ref units, not postprocessed to SI units yet
-    raw_ne = core_profiles.ne
-    raw_ne_face = raw_ne.face_value()
-    raw_ne_face_grad = raw_ne.face_grad(rmid)
-    raw_ni = core_profiles.ni
-    raw_ni_face = raw_ni.face_value()
-    raw_ni_face_grad = raw_ni.face_grad(rmid)
-    raw_nimp = core_profiles.nimp
-    raw_nimp_face = raw_nimp.face_value()
-    raw_nimp_face_grad = raw_nimp.face_grad(rmid)
-
-    # True SI value versions
-    true_ne_face = raw_ne_face * nref
-    true_ni_face = raw_ni_face * nref
-    true_nimp_face = raw_nimp_face * nref
-
     # gyrobohm diffusivity
     # (defined here with Lref=Rmin due to QLKNN training set normalization)
-    chiGB = (
-        (core_profiles.Ai * constants.mp) ** 0.5
-        / (constants.qe * geo.B0) ** 2
-        * (temp_ion_face * constants.keV2J) ** 1.5
-        / Rmin
+    chiGB = quasilinear_transport_model.calculate_chiGB(
+        core_profiles=core_profiles,
+        b_unit=geo.B0,
+        reference_length=geo.Rmin,
     )
 
     # transport coefficients from the qlknn-hyper-10D model
     # (K.L. van de Plassche PoP 2020)
 
-    # TODO(b/335581689): make a unit test that tests this function directly
-    # with set_pedestal = False. Currently this is tested only via
-    # sim test7, which has set_pedestal=True. With set_pedestal=True,
-    # mutants of Ati[-1], Ate[-1], An[-1] all affect only chi[-1], but
-    # chi[-1] remains above config.transport.chimin for all mutants.
-    # The pedestal feature then clips chi[-1] to config.transport.chimin, so the
-    # mutants have no effect.
-
     # set up input vectors (all as jax.numpy arrays on face grid)
 
-    # R/LTi profile from current timestep temp_ion
-    Ati = -Rmaj * temp_ion_face_grad / temp_ion_face
-    # to avoid divisions by zero
-    Ati = jnp.where(jnp.abs(Ati) < constants.eps, constants.eps, Ati)
-
-    # R/LTe profile from current timestep temp_el
-    Ate = -Rmaj * temp_electron_face_grad / temp_electron_face
-    # to avoid divisions by zero
-    Ate = jnp.where(jnp.abs(Ate) < constants.eps, constants.eps, Ate)
-
-    # R/Ln profiles from current timestep
-    # OK to use normalized version here, because nref in numer and denom
-    # cancels.
-    Ane = -Rmaj * raw_ne_face_grad / raw_ne_face
-    Ani0 = -Rmaj * raw_ni_face_grad / raw_ni_face
-    # To avoid divisions by zero in cases where Zeff=1.
-    Ani1 = jnp.where(
-        jnp.abs(raw_nimp_face) < constants.eps,
-        0.0,
-        -Rmaj * raw_nimp_face_grad / raw_nimp_face,
+    # Calculate normalized logarithmic gradients
+    normalized_logarithmic_gradients = quasilinear_transport_model.NormalizedLogarithmicGradients.from_profiles(
+        core_profiles=core_profiles,
+        radial_coordinate=rmid,
+        reference_length=geo.Rmaj,
     )
-    # to avoid divisions by zero
-    Ane = jnp.where(jnp.abs(Ane) < constants.eps, constants.eps, Ane)
-    Ani0 = jnp.where(jnp.abs(Ani0) < constants.eps, constants.eps, Ani0)
-    Ani1 = jnp.where(jnp.abs(Ani1) < constants.eps, constants.eps, Ani1)
 
     # Calculate q and s.
     # Need to recalculate since in the nonlinear solver psi has intermediate
@@ -187,13 +134,15 @@ def _prepare_qualikiz_inputs(
     )
 
     # Inverse aspect ratio at LCFS.
-    epsilon_lcfs = rmid_face[-1] / Rmaj
+    epsilon_lcfs = rmid_face[-1] / geo.Rmaj
     # Local normalized radius.
     x = rmid_face / rmid_face[-1]
     x = jnp.where(jnp.abs(x) < constants.eps, constants.eps, x)
 
     # Ion to electron temperature ratio
-    Ti_Te = temp_ion_face / temp_electron_face
+    Ti_Te = (
+        core_profiles.temp_ion.face_value() / core_profiles.temp_el.face_value()
+    )
 
     # logarithm of normalized collisionality
     nu_star = physics.calc_nu_star(
@@ -206,11 +155,12 @@ def _prepare_qualikiz_inputs(
     log_nu_star_face = jnp.log10(nu_star)
 
     # calculate alpha for magnetic shear correction (see S. van Mulders NF 2021)
-    factor_0 = 2 / geo.B0**2 * constants.mu0 * q**2
-    alpha = factor_0 * (
-        temp_electron_face * constants.keV2J * true_ne_face * (Ate + Ane)
-        + true_ni_face * temp_ion_face * constants.keV2J * (Ati + Ani0)
-        + true_nimp_face * temp_ion_face * constants.keV2J * (Ati + Ani1)
+    alpha = quasilinear_transport_model.calculate_alpha(
+        core_profiles=core_profiles,
+        nref=nref,
+        q=q,
+        b_unit=geo.B0,
+        normalized_logarithmic_gradients=normalized_logarithmic_gradients,
     )
 
     # to approximate impact of Shafranov shift. From van Mulders Nucl. Fusion
@@ -248,23 +198,23 @@ def _prepare_qualikiz_inputs(
         alpha - 0.2,
         smag,
     )
-    normni = raw_ni_face / raw_ne_face
+    normni = core_profiles.ni.face_value() / core_profiles.ne.face_value()
     return QualikizInputs(
         Zeff_face=Zeff_face,
-        Ati=Ati,
-        Ate=Ate,
-        Ane=Ane,
-        Ani0=Ani0,
-        Ani1=Ani1,
+        Ati=normalized_logarithmic_gradients.Ati,
+        Ate=normalized_logarithmic_gradients.Ate,
+        Ane=normalized_logarithmic_gradients.Ane,
+        Ani0=normalized_logarithmic_gradients.Ani0,
+        Ani1=normalized_logarithmic_gradients.Ani1,
         q=q,
         smag=smag,
         x=x,
         Ti_Te=Ti_Te,
         log_nu_star_face=log_nu_star_face,
         normni=normni,
         chiGB=chiGB,
-        Rmaj=Rmaj,
-        Rmin=Rmin,
+        Rmaj=geo.Rmaj,
+        Rmin=geo.Rmin,
         alpha=alpha,
         epsilon_lcfs=epsilon_lcfs,
     )

torax/transport_model/qualikiz_wrapper.py

@@ -242,6 +242,8 @@ def _extract_run_data(
         transport=transport,
         geo=geo,
         core_profiles=core_profiles,
+        gradient_reference_length=geo.Rmaj,
+        gyrobohm_flux_reference_length=geo.Rmin,
     )