Merge remote-tracking branch 'origin/main' into add-vloop-bc

docs/configuration.rst

@@ -517,8 +517,8 @@ total pressure at the pedestal and the ratio of ion to electron temperature.
 geometry
 --------
 
-``geometry_type`` (str = 'chease')
-  Geometry model used, from the following options.
+``geometry_type`` (str)
+  Geometry model used. A string must be provided from the following options.
 
 * ``'circular'``
     An ad-hoc circular geometry model. Includes elongation corrections.

docs/output.rst

@@ -216,6 +216,21 @@ analysis and inspection.
 ``pressure_thermal_tot_face`` (time, rho_face) [Pa]:
   Total thermal pressure on the face grid.
 
+``te_volume_avg`` (time) [keV]:
+  Volume average electron temperature.
+
+``ti_volume_avg`` (time) [keV]:
+  Volume average ion temperature.
+
+``ne_volume_avg`` (time) [nref m^-3]:
+  Volume average electron density.
+
+``ni_volume_avg`` (time) [nref m^-3]:
+  Volume average ion density.
+
+``fgw_ne_volume_avg`` (time) [dimensionless]:
+  Greenwald fraction from volume-averaged electron density.
+
 ``pprime_face`` (time, rho_face) [Pa/Wb]:
   Derivative of total pressure with respect to poloidal flux on the face grid.
 
@@ -228,6 +243,15 @@ analysis and inspection.
 ``W_thermal_tot`` (time) [J]:
   Total thermal stored energy.
 
+``Wpol`` (time) [J]
+  Total magnetic energy
+
+``q95`` (time) [dimensionless]
+  Safety-factor at 95% of the normalized poloidal flux coordinate.
+
+``li3`` (time) [dimensionless]:
+  Normalized plasma internal inductance, ITER convention
+
 ``tauE`` (time) [s]:
   Thermal confinement time defined as ``W_thermal_tot`` / ``P_heating``, where
   ``P_heating`` is the total heating power into the plasma, including external

torax/output.py

@@ -400,7 +400,12 @@ def _save_geometry(
 
     # Get the variables from dataclass fields.
     for field_name, data in dataclasses.asdict(self.geometry).items():
-      if "hires" in field_name or not isinstance(data, jax.Array):
+      if (
+          "hires" in field_name
+          or field_name == "geometry_type"
+          or field_name == "Ip_from_parameters"
+          or not isinstance(data, jax.Array)
+      ):
         continue
       data_array = self._pack_into_data_array(
           field_name,

torax/post_processing.py

@@ -55,7 +55,6 @@
 }
 
 
-@jax_utils.jit
 def _compute_pressure(
     core_profiles: state.CoreProfiles,
 ) -> tuple[array_typing.ArrayFloat, ...]:
@@ -88,7 +87,6 @@ def _compute_pressure(
   )
 
 
-@jax_utils.jit
 def _compute_pprime(
     core_profiles: state.CoreProfiles,
 ) -> array_typing.ArrayFloat:
@@ -137,7 +135,6 @@ def _compute_pprime(
 
 
 # pylint: disable=invalid-name
-@jax_utils.jit
 def _compute_FFprime(
     core_profiles: state.CoreProfiles,
     geo: geometry.Geometry,
@@ -178,7 +175,6 @@ def _compute_FFprime(
 # pylint: enable=invalid-name
 
 
-@jax_utils.jit
 def _compute_stored_thermal_energy(
     p_el: array_typing.ArrayFloat,
     p_ion: array_typing.ArrayFloat,
@@ -205,7 +201,6 @@ def _compute_stored_thermal_energy(
   return wth_el, wth_ion, wth_tot
 
 
-@jax_utils.jit
 def _calculate_integrated_sources(
     geo: geometry.Geometry,
     core_profiles: state.CoreProfiles,
@@ -301,7 +296,6 @@ def _calculate_integrated_sources(
   return integrated
 
 
-@jax_utils.jit
 def _calculate_q95(
     psi_norm_face: array_typing.ArrayFloat,
     core_profiles: state.CoreProfiles,
@@ -320,6 +314,33 @@ 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,
     geo: geometry.Geometry,
@@ -470,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]
@@ -504,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