Merge pull request #658 from google-deepmind:add-ne-greenwald-fraction

PiperOrigin-RevId: 726452906

torax/post_processing.py

@@ -314,6 +314,32 @@ def _calculate_q95(
   return q95
 
 
+def _calculate_greenwald_fraction_from_ne_vol_avg(
+    ne_volume_avg: array_typing.ScalarFloat,
+    core_profiles: state.CoreProfiles,
+    geo: geometry.Geometry,
+) -> array_typing.ScalarFloat:
+  """Calculates the Greenwald fraction from the volume-average electron density.
+
+  Args:
+    ne_volume_avg: Volume-averaged electron density [nref m^-3]
+    core_profiles: CoreProfiles object containing information on currents
+      and densities.
+    geo: Geometry object
+
+  Returns:
+    fgw_ne_vol_avg: Volume-averaged electron density Greenwald fraction
+  """
+  # gw_limit is in units of 10^20 m^-3 when Ip is in MA and rmid is in m.
+  gw_limit = (
+      core_profiles.currents.Ip_profile_face[-1]
+      * 1e-6
+      / (jnp.pi * geo.Rmin ** 2)
+  )
+  fgw_ne_vol_avg = ne_volume_avg * core_profiles.nref / (gw_limit * 1e20)
+  return fgw_ne_vol_avg
+
+
 @jax_utils.jit
 def make_outputs(
     sim_state: state.ToraxSimState,
@@ -465,6 +491,9 @@ def make_outputs(
       / geo.volume[-1]
   )
 
+  fgw_ne_volume_avg = _calculate_greenwald_fraction_from_ne_vol_avg(
+      ne_volume_avg, sim_state.core_profiles, geo
+      )
   Wpol = physics.calc_Wpol(geo, sim_state.core_profiles.psi)
   li3 = physics.calc_li3(
       geo.Rmaj, Wpol, sim_state.core_profiles.currents.Ip_profile_face[-1]
@@ -499,6 +528,7 @@ def make_outputs(
       ti_volume_avg=ti_volume_avg,
       ne_volume_avg=ne_volume_avg,
       ni_volume_avg=ni_volume_avg,
+      fgw_ne_volume_avg=fgw_ne_volume_avg,
       q95=q95,
       Wpol=Wpol,
       li3=li3,

torax/state.py

@@ -327,6 +327,8 @@ class PostProcessedOutputs:
     ti_volume_avg: Volume average ion temperature [keV]
     ne_volume_avg: Volume average electron density [nref m^-3]
     ni_volume_avg: Volume average main ion density [nref m^-3]
+    fgw_ne_volume_avg: Greenwald fraction from volume-averaged electron density
+      [dimensionless]
     q95: q at 95% of the normalized poloidal flux
     Wpol: Total magnetic energy [J]
     li3: Normalized plasma internal inductance, ITER convention [dimensionless]
@@ -385,6 +387,7 @@ class PostProcessedOutputs:
   ti_volume_avg: array_typing.ScalarFloat
   ne_volume_avg: array_typing.ScalarFloat
   ni_volume_avg: array_typing.ScalarFloat
+  fgw_ne_volume_avg: array_typing.ScalarFloat
   q95: array_typing.ScalarFloat
   Wpol: array_typing.ScalarFloat
   li3: array_typing.ScalarFloat
@@ -443,6 +446,7 @@ def zeros(cls, geo: geometry.Geometry) -> PostProcessedOutputs:
         ti_volume_avg=jnp.array(0.0),
         ne_volume_avg=jnp.array(0.0),
         ni_volume_avg=jnp.array(0.0),
+        fgw_ne_volume_avg=jnp.array(0.0),
         q95=jnp.array(0.0),
         Wpol=jnp.array(0.0),
         li3=jnp.array(0.0),

torax/tests/post_processing.py

@@ -15,6 +15,7 @@
 """Tests for post_processing.py."""
 
 import dataclasses
+from unittest import mock
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
@@ -29,6 +30,7 @@
 from torax.config import runtime_params_slice
 from torax.fvm import cell_variable
 from torax.geometry import circular_geometry
+from torax.geometry import geometry
 from torax.geometry import geometry_provider
 from torax.sources import source_profiles as source_profiles_lib
 from torax.tests.test_lib import default_sources
@@ -184,6 +186,40 @@ def test_compute_stored_thermal_energy(self):
     np.testing.assert_allclose(wth_ion, 1.5 * p_ion[0] * volume)
     np.testing.assert_allclose(wth_tot, 1.5 * p_tot[0] * volume)
 
+  def test_calculate_greenwald_fraction_from_ne_vol_avg(self):
+    """Test that Greenwald fraction is calculated correctly."""
+    ne_volume_avg = 1.0
+
+    core_profiles = mock.create_autospec(
+        state.CoreProfiles,
+        instance=True,
+        nref=1e20,
+        currents=mock.create_autospec(
+            state.Currents,
+            instance=True,
+            Ip_profile_face=np.array([0, np.pi * 1e6]),
+        ),
+    )
+    geo = mock.create_autospec(
+        geometry.Geometry,
+        instance=True,
+        Rmin=1.0,
+    )
+
+    # pylint: disable=protected-access
+    fgw_ne_volume_avg_calculated = (
+        post_processing._calculate_greenwald_fraction_from_ne_vol_avg(
+            ne_volume_avg, core_profiles, geo
+        )
+    )
+
+    fgw_ne_volume_avg_expected = 1.0
+
+    np.testing.assert_allclose(
+        fgw_ne_volume_avg_calculated, fgw_ne_volume_avg_expected
+    )
+    # pylint: enable=protected-access
+
   def test_calculate_integrated_sources(self):
     """Checks integrated quantities match expectations."""
     # pylint: disable=protected-access