Introduce thermal ion and impurity densities.

This change adds `n_i_thermal` and `n_impurity_thermal` properties to the state, which represent the main ion and impurity densities after subtracting the contributions from fast ions. The calculation of `pressure_thermal_i` is updated to use these new thermal densities. New variables are also added to the output tools skip list.

PiperOrigin-RevId: 878655712

Author: hamelphi

torax/_src/output_tools/output.py

@@ -574,6 +574,7 @@ def _save_core_profiles(
           "charge_state_info_face",
           "impurity_density_scaling",
           "fast_ions",
+          "n_impurity_thermal",
       ):
         continue
 

torax/_src/physics/formulas.py

@@ -69,20 +69,20 @@ def calc_pprime(
   psi = core_profiles.psi.face_value()
   n_e = core_profiles.n_e.face_value()
   n_i = core_profiles.n_i.face_value()
-  n_impurity = core_profiles.n_impurity.face_value()
+  n_impurity_thermal = core_profiles.n_impurity_thermal.face_value()
   T_i = core_profiles.T_i.face_value()
   T_e = core_profiles.T_e.face_value()
   dne_drhon = core_profiles.n_e.face_grad()
   dni_drhon = core_profiles.n_i.face_grad()
-  dnimp_drhon = core_profiles.n_impurity.face_grad()
+  dnimp_drhon = core_profiles.n_impurity_thermal.face_grad()
   dti_drhon = core_profiles.T_i.face_grad()
   dte_drhon = core_profiles.T_e.face_grad()
   dpsi_drhon = core_profiles.psi.face_grad()
 
   dptot_drhon = constants.CONSTANTS.keV_to_J * (
       n_e * dte_drhon
       + n_i * dti_drhon
-      + n_impurity * dti_drhon
+      + n_impurity_thermal * dti_drhon
       + dne_drhon * T_e
       + dni_drhon * T_i
       + dnimp_drhon * T_i

torax/_src/state.py

@@ -172,19 +172,45 @@ def pressure_thermal_e(self) -> cell_variable.CellVariable:
         right_face_grad_constraint=None,
     )
 
+  @functools.cached_property
+  def n_impurity_thermal(self) -> cell_variable.CellVariable:
+    """True thermal impurity density without fast ions [m^-3].
+
+    TODO(b/868852029): This assumes fast ions subtracted from thermal impurities
+    are originally from thermal species. This needs to be updated when we
+    introduce other sources of fast ions that do not deplete thermal species
+    (e.g. fusion).
+    """
+    n_impurity_thermal_value = self.n_impurity.value
+    n_impurity_thermal_right = self.n_impurity.right_face_constraint
+    for fast_ion in self.fast_ions:
+      if fast_ion.species in self.impurity_fractions:
+        n_impurity_thermal_value -= fast_ion.n.value
+        if (
+            n_impurity_thermal_right is not None
+            and fast_ion.n.right_face_constraint is not None
+        ):
+          n_impurity_thermal_right -= fast_ion.n.right_face_constraint
+    return cell_variable.CellVariable(
+        value=n_impurity_thermal_value,
+        face_centers=self.n_impurity.face_centers,
+        right_face_constraint=n_impurity_thermal_right,
+        right_face_grad_constraint=None,
+    )
+
   @functools.cached_property
   def pressure_thermal_i(self) -> cell_variable.CellVariable:
     """Ion thermal pressure [Pa]."""
     return cell_variable.CellVariable(
         value=self.T_i.value
         * constants.CONSTANTS.keV_to_J
-        * (self.n_i.value + self.n_impurity.value),
+        * (self.n_i.value + self.n_impurity_thermal.value),
         face_centers=self.n_i.face_centers,
         right_face_constraint=self.T_i.right_face_constraint
         * constants.CONSTANTS.keV_to_J
         * (
             self.n_i.right_face_constraint
-            + self.n_impurity.right_face_constraint
+            + self.n_impurity_thermal.right_face_constraint
         ),
         right_face_grad_constraint=None,
     )

torax/_src/tests/state_test.py

@@ -29,6 +29,7 @@
 from torax._src.core_profiles.plasma_composition import plasma_composition
 from torax._src.geometry import circular_geometry
 from torax._src.orchestration import run_simulation
+from torax._src.physics import fast_ions as fast_ions_lib
 from torax._src.test_utils import core_profile_helpers
 from torax._src.test_utils import default_configs
 from torax._src.torax_pydantic import model_config
@@ -356,6 +357,40 @@ def test_pressure_thermal_total(self):
         8.75e20 * constants.CONSTANTS.keV_to_J,
     )
 
+  def test_thermal_densities_with_fast_ions(self):
+    geo = circular_geometry.CircularConfig(n_rho=10).build_geometry()
+    base_core_profiles = core_profile_helpers.make_zero_core_profiles(geo)
+
+    fast_ion_he3 = fast_ions_lib.FastIon(
+        species='He3',
+        source='ICRH',
+        n=core_profile_helpers.make_constant_core_profile(geo, 0.05e20),
+        T=core_profile_helpers.make_constant_core_profile(geo, 100.0),
+    )
+
+    core_profiles = dataclasses.replace(
+        base_core_profiles,
+        T_i=core_profile_helpers.make_constant_core_profile(geo, 1.0),
+        T_e=core_profile_helpers.make_constant_core_profile(geo, 2.0),
+        n_e=core_profile_helpers.make_constant_core_profile(geo, 3.0e20),
+        n_i=core_profile_helpers.make_constant_core_profile(geo, 2.5e20),
+        n_impurity=core_profile_helpers.make_constant_core_profile(
+            geo, 0.25e20
+        ),
+        fast_ions=(fast_ion_he3,),
+        main_ion_fractions={'D': 1.0},
+        impurity_fractions={'He3': jnp.ones(10)},
+    )
+
+    np.testing.assert_allclose(
+        core_profiles.n_impurity_thermal.value,
+        0.20e20,
+    )
+    np.testing.assert_allclose(
+        core_profiles.pressure_thermal_i.value,
+        2.7e20 * constants.CONSTANTS.keV_to_J,
+    )
+
   def test_cached_properties_with_jit(self):
     """Tests cached properties behave correctly within a jitted function."""
 

torax/_src/transport_model/quasilinear_transport_model.py

@@ -60,7 +60,7 @@ def from_profiles(
         "lref_over_lte": core_profiles.T_e,
         "lref_over_lne": core_profiles.n_e,
         "lref_over_lni0": core_profiles.n_i,
-        "lref_over_lni1": core_profiles.n_impurity,
+        "lref_over_lni1": core_profiles.n_impurity_thermal,
     }.items():
       gradients[name] = calculate_normalized_logarithmic_gradient(
           var=profile,
@@ -167,7 +167,7 @@ def calculate_alpha(
           normalized_logarithmic_gradients.lref_over_lti
           + normalized_logarithmic_gradients.lref_over_lni0
       )
-      + core_profiles.n_impurity.face_value()
+      + core_profiles.n_impurity_thermal.face_value()
       * core_profiles.T_i.face_value()
       * (
           normalized_logarithmic_gradients.lref_over_lti

torax/_src/transport_model/tests/quasilinear_transport_model_test.py

@@ -204,6 +204,7 @@ def _get_cell_variable(value):
         n_impurity=_get_cell_variable(1.0e20),
         fast_ions=(),
         pressure_fast_i=_get_cell_variable(0.0),
+        n_impurity_thermal=_get_cell_variable(1.0e20),
     )
 
     normalized_logarithmic_gradients = (