Modify input of gyrobohm function for versatility.

The previous version of the gyrobohm function took as input the CoreProfiles object. This worked in the Qualikiz transport model, but for other transport models like TGLF, we want to be able to specify the gyrobohm unit using different reference values for temperature and mass.

PiperOrigin-RevId: 709083009

torax/transport_model/qualikiz_based_transport_model.py

@@ -124,8 +124,9 @@ def _prepare_qualikiz_inputs(
     # gyrobohm diffusivity
     # (defined here with Lref=Rmin due to QLKNN training set normalization)
     chiGB = quasilinear_transport_model.calculate_chiGB(
-        core_profiles=core_profiles,
-        b_unit=geo.B0,
+        reference_temperature=core_profiles.temp_ion.face_value(),
+        reference_magnetic_field=geo.B0,
+        reference_mass=core_profiles.Ai,
         reference_length=geo.Rmin,
     )
 
@@ -180,7 +181,7 @@ def _prepare_qualikiz_inputs(
         core_profiles=core_profiles,
         nref=nref,
         q=q,
-        b_unit=geo.B0,
+        reference_magnetic_field=geo.B0,
         normalized_logarithmic_gradients=normalized_logarithmic_gradients,
     )
 

torax/transport_model/quasilinear_transport_model.py

@@ -26,26 +26,30 @@
 
 
 def calculate_chiGB(  # pylint: disable=invalid-name
-    core_profiles: state.CoreProfiles,
-    b_unit: chex.Numeric,
+    reference_temperature: chex.Array,
+    reference_magnetic_field: chex.Numeric,
+    reference_mass: chex.Numeric,
     reference_length: chex.Numeric,
 ) -> chex.Array:
   """Calculates the gyrobohm diffusivity.
 
+  Different transport models make different choices for the reference
+  temperature, magnetic field, and mass used for gyrobohm normalization.
+
   Args:
-    core_profiles: CoreProfiles object containing plasma profiles.
-    b_unit: Magnetic field strength [T]. Different transport models have
-      different definitions of the specific magnetic field input.
+    reference_temperature: Reference temperature on the face grid [keV].
+    reference_magnetic_field: Magnetic field strength [T].
+    reference_mass: Reference ion mass [amu].
     reference_length: Reference length for normalization [m].
 
   Returns:
     Gyrobohm diffusivity as a chex.Array [dimensionless].
   """
   constants = constants_module.CONSTANTS
   return (
-      (core_profiles.Ai * constants.mp) ** 0.5
-      / (constants.qe * b_unit) ** 2
-      * (core_profiles.temp_ion.face_value() * constants.keV2J) ** 1.5
+      (reference_mass * constants.mp) ** 0.5
+      / (reference_magnetic_field * constants.qe) ** 2
+      * (reference_temperature * constants.keV2J) ** 1.5
       / reference_length
   )
 
@@ -54,7 +58,7 @@ def calculate_alpha(
     core_profiles: state.CoreProfiles,
     nref: chex.Numeric,
     q: chex.Array,
-    b_unit: chex.Numeric,
+    reference_magnetic_field: chex.Numeric,
     normalized_logarithmic_gradients: NormalizedLogarithmicGradients,
 ) -> chex.Array:
   """Calculates the alpha_MHD parameter.
@@ -68,8 +72,8 @@ def calculate_alpha(
     core_profiles: CoreProfiles object containing plasma profiles.
     nref: Reference density.
     q: Safety factor.
-    b_unit: Magnetic field strength. Different transport models have different
-      definitions of the specific magnetic field input.
+    reference_magnetic_field: Magnetic field strength. Different transport
+      models have different definitions of the specific magnetic field input.
     normalized_logarithmic_gradients: Normalized logarithmic gradients of plasma
       profiles.
 
@@ -78,7 +82,14 @@ def calculate_alpha(
   """
   constants = constants_module.CONSTANTS
 
-  factor_0 = 2 * constants.keV2J * nref / b_unit**2 * constants.mu0 * q**2
+  factor_0 = (
+      2
+      * constants.keV2J
+      * nref
+      / reference_magnetic_field**2
+      * constants.mu0
+      * q**2
+  )
   alpha = factor_0 * (
       core_profiles.temp_el.face_value()
       * core_profiles.ne.face_value()

torax/transport_model/tests/quasilinear_transport_model.py

@@ -154,8 +154,9 @@ def test_calculate_chiGB(self):
         value=jnp.array([1.0]), right_face_constraint=jnp.array(1.0)
     )
     chiGB = quasilinear_transport_model.calculate_chiGB(
-        core_profiles=core_profiles,
-        b_unit=1.0,
+        reference_temperature=core_profiles.temp_ion.face_value(),
+        reference_magnetic_field=1.0,
+        reference_mass=1.0,
         reference_length=1.0,
     )
     chi_GB_expected = (
@@ -183,7 +184,7 @@ def test_calculate_alpha(self):
         core_profiles=core_profiles,
         nref=1e20,
         q=np.array(1.0),
-        b_unit=1.0,
+        reference_magnetic_field=1.0,
         normalized_logarithmic_gradients=normalized_logarithmic_gradients,
     )