Add n_e_right_bc to EdgeModelOutputs and implement _update_density.

Complete the base core boundary condition interface by adding electron density (n_e_right_bc) to EdgeModelOutputs.
In ExtendedLengyelOutputs, default n_e_right_bc to jnp.nan as Lengyel does not predict core density.
Implement _update_density in torax/_src/edge/updaters.py to update runtime_params.profile_conditions.n_e_right_bc, and add unit test coverage in edge_updaters_test.py.

PiperOrigin-RevId: 983998350

Author: tamaranorman

torax/_src/edge/base.py

@@ -15,6 +15,7 @@
 """Base classes for edge models."""
 
 import abc
+from collections.abc import Mapping
 import dataclasses
 import chex
 import jax
@@ -40,44 +41,22 @@ class EdgeModelOutputs:
   Attributes:
     T_e_right_bc: Electron temperature boundary condition at LCFS [keV].
     T_i_right_bc: Ion temperature boundary condition at LCFS [keV].
-    q_parallel: Parallel heat flux upstream [W/m^2].
-    q_perpendicular_target: Heat flux perpendicular to the target [W/m^2].
-    T_e_separatrix: Electron temperature at the separatrix [keV].
-    T_e_target: Electron temperature at sheath entrance [eV].
-    pressure_neutral_divertor: Neutral pressure in the divertor [Pa].
+    n_e_right_bc: Electron density boundary condition at LCFS [m^-3].
+    impurity_right_bc: Mapping from impurity symbol to its right boundary
+      condition (n_e_ratio at LCFS).
   """
 
   T_e_right_bc: jax.Array
   T_i_right_bc: jax.Array
-  q_parallel: jax.Array
-  q_perpendicular_target: jax.Array
-  T_e_separatrix: jax.Array
-  T_e_target: jax.Array
-  pressure_neutral_divertor: jax.Array
+  n_e_right_bc: jax.Array
+  impurity_right_bc: Mapping[str, jax.Array]
 
   def to_output_dict(
       self, context: output_grid_context.OutputGridContext
   ) -> dict[str, output_grid_context.OutputVar]:
     """Returns a dictionary of standard edge output variable tuples."""
-    outputs = {
-        output_keys.Q_PARALLEL: context.pack(
-            output_keys.Q_PARALLEL, self.q_parallel
-        ),
-        output_keys.Q_PERPENDICULAR_TARGET: context.pack(
-            output_keys.Q_PERPENDICULAR_TARGET, self.q_perpendicular_target
-        ),
-        output_keys.T_E_SEPARATRIX: context.pack(
-            output_keys.T_E_SEPARATRIX, self.T_e_separatrix
-        ),
-        output_keys.T_E_TARGET: context.pack(
-            output_keys.T_E_TARGET, self.T_e_target
-        ),
-        output_keys.PRESSURE_NEUTRAL_DIVERTOR: context.pack(
-            output_keys.PRESSURE_NEUTRAL_DIVERTOR,
-            self.pressure_neutral_divertor,
-        ),
-    }
-    return {k: v for k, v in outputs.items() if v is not None}
+    del context
+    return {}
 
   def to_xr_datatree(
       self, context: output_grid_context.OutputGridContext

torax/_src/edge/extended_lengyel/extended_lengyel_enums.py

@@ -42,3 +42,27 @@ class SolverMode(enum.StrEnum):
   FIXED_POINT = 'fixed_point'
   NEWTON_RAPHSON = 'newton_raphson'
   HYBRID = 'hybrid'
+
+
+class FixedImpuritySourceOfTruth(enum.StrEnum):
+  """Source of truth for fixed impurity concentrations when using an edge model.
+
+  Determines how impurity concentrations are handled between the core plasma
+  simulation and the edge model.
+
+  Attributes:
+    CORE: * The core impurity profiles are the source of truth. * The edge
+      model's impurity concentrations are derived from the core values at the
+      last closed flux surface: `c_edge = c_core_face[-1] * enrichment_factor`.
+    EDGE: * The edge model's `fixed_impurity_concentrations` are the source of
+      truth. * The core impurity profiles (n_e_ratios) are scaled to match the
+      values determined by the edge model. runtime_params still sets the profile
+      shape: `c_core = c_core / c_core_face[-1] * c_edge / enrichment_factor`.
+
+  Note: For seeded impurities in the extended Lengyel edge model, the source of
+  truth is always the edge model, regardless of this setting. This enum only
+  controls the behavior for fixed impurities in that case.
+  """
+
+  CORE = 'core'
+  EDGE = 'edge'

torax/_src/edge/extended_lengyel/extended_lengyel_model.py

@@ -15,7 +15,6 @@
 """Implementation of extended_lengyel instance of EdgeModel."""
 
 import dataclasses
-import enum
 import logging
 from typing import Mapping
 import jax
@@ -32,6 +31,7 @@
 from torax._src.edge.extended_lengyel import divertor_sol_1d as divertor_sol_1d_lib
 from torax._src.edge.extended_lengyel import extended_lengyel_defaults
 from torax._src.edge.extended_lengyel import extended_lengyel_enums
+from torax._src.edge.extended_lengyel import extended_lengyel_formulas
 from torax._src.edge.extended_lengyel import extended_lengyel_solvers
 from torax._src.edge.extended_lengyel import extended_lengyel_standalone
 from torax._src.geometry import geometry
@@ -42,28 +42,6 @@
 
 
 # pylint: disable=invalid-name
-class FixedImpuritySourceOfTruth(enum.StrEnum):
-  """Source of truth for fixed impurity concentrations when using an edge model.
-
-  Determines how impurity concentrations are handled between the core plasma
-  simulation and the edge model.
-
-  Attributes:
-    CORE: * The core impurity profiles are the source of truth. * The edge
-      model's impurity concentrations are derived from the core values at the
-      last closed flux surface: `c_edge = c_core_face[-1] * enrichment_factor`.
-    EDGE: * The edge model's `fixed_impurity_concentrations` are the source of
-      truth. * The core impurity profiles (n_e_ratios) are scaled to match the
-      values determined by the edge model. runtime_params still sets the profile
-        shape: `c_core = c_core / c_core_face[-1] * c_edge / enrichment_factor`.
-
-  Note: For seeded impurities in the extended Lengyel edge model, the source of
-  truth is always the edge model, regardless of this setting. This enum only
-  controls the behavior for fixed impurities in that case.
-  """
-
-  CORE = 'core'
-  EDGE = 'edge'
 
 
 @jax.tree_util.register_dataclass
@@ -87,7 +65,7 @@ class InitialGuessRuntimeParams:
 
 
 @jax.tree_util.register_dataclass
-@dataclasses.dataclass(frozen=True)
+@dataclasses.dataclass(frozen=True, kw_only=True)
 class RuntimeParams(edge_runtime_params.RuntimeParams):
   """Runtime parameters for the extended Lengyel edge model."""
 
@@ -100,11 +78,12 @@ class RuntimeParams(edge_runtime_params.RuntimeParams):
   solver_mode: extended_lengyel_enums.SolverMode = dataclasses.field(
       metadata={'static': True}
   )
-  impurity_sot: FixedImpuritySourceOfTruth = dataclasses.field(
-      metadata={'static': True}
+  impurity_sot: extended_lengyel_enums.FixedImpuritySourceOfTruth = (
+      dataclasses.field(metadata={'static': True})
   )
   # Not static to allow rapid sensitivity checking of edge-model impact.
   update_temperatures: array_typing.BoolScalar
+  update_density: array_typing.BoolScalar
   update_impurities: array_typing.BoolScalar
   fixed_point_iterations: int
   newton_raphson_iterations: int
@@ -139,7 +118,7 @@ class RuntimeParams(edge_runtime_params.RuntimeParams):
   # --- Impurity parameters ---
   seed_impurity_weights: Mapping[str, array_typing.FloatScalar] | None
   fixed_impurity_concentrations: Mapping[str, array_typing.FloatScalar]
-  enrichment_factor: Mapping[str, array_typing.FloatScalar]
+  enrichment_factor: Mapping[str, array_typing.FloatScalar] | None
   use_enrichment_model: bool = dataclasses.field(metadata={'static': True})
   enrichment_model_multiplier: array_typing.FloatScalar
 
@@ -238,7 +217,10 @@ def __call__(
     fixed_impurity_concentrations = edge_params.fixed_impurity_concentrations
     # If the source of truth for fixed impurities is the core, calculate the
     # edge concentrations from the core ratios.
-    if edge_params.impurity_sot == FixedImpuritySourceOfTruth.CORE:
+    if (
+        edge_params.impurity_sot
+        == extended_lengyel_enums.FixedImpuritySourceOfTruth.CORE
+    ):
       # Initialization
       fixed_impurity_concentrations = {}
       impurity_params = runtime_params.plasma_composition.impurity
@@ -255,10 +237,26 @@ def __call__(
           continue
 
         # Calculate edge concentration: c_edge = c_core_lcfs * enrichment_factor
-        # Enrichment factor exists for all species (validated in config)
-        fixed_impurity_concentrations[species] = (
-            ratio_face[-1] * edge_params.enrichment_factor[species]
-        )
+        if edge_params.use_enrichment_model:
+          if previous_edge_outputs is not None:
+            assert isinstance(
+                previous_edge_outputs,
+                extended_lengyel_standalone.ExtendedLengyelOutputs,
+            )
+            enrichment = previous_edge_outputs.calculated_enrichment[species]
+          else:
+            # For initial timestep when previous_edge_outputs is None
+            enrichment = extended_lengyel_formulas.calc_enrichment_kallenbach(
+                1.0, species, edge_params.enrichment_model_multiplier
+            )
+        elif edge_params.enrichment_factor is not None:
+          enrichment = edge_params.enrichment_factor[species]
+        else:
+          raise ValueError(
+              'enrichment_factor must be provided when use_enrichment_model is'
+              ' False.'
+          )
+        fixed_impurity_concentrations[species] = ratio_face[-1] * enrichment
 
     # Determine initial guesses
     initial_guess = _get_initial_guess(edge_params, previous_edge_outputs)
@@ -308,6 +306,9 @@ def __call__(
         multistart_num_guesses=edge_params.multistart_num_guesses,
         enrichment_model_multiplier=edge_params.enrichment_model_multiplier,
         diverted=diverted,
+        use_enrichment_model=edge_params.use_enrichment_model,
+        enrichment_factor=edge_params.enrichment_factor,
+        impurity_sot=edge_params.impurity_sot,
         initial_guess=initial_guess,
     )