Add Shaing corrections to Angioni-Sauter model

Author: theo-brown

torax/_src/neoclassical/transport/angioni_sauter.py

@@ -21,8 +21,10 @@
 https://gitlab.epfl.ch/spc/public/neos [O. Sauter et al]
 """
 
+import dataclasses
 from typing import Annotated, Literal
 
+import jax
 from jax import numpy as jnp
 from torax._src import array_typing
 from torax._src import constants
@@ -39,20 +41,74 @@
 # pylint: disable=invalid-name
 
 
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True)
+class RuntimeParams(transport_runtime_params.RuntimeParams):
+  """RuntimeParams for the Angioni-Sauter neoclassical transport model."""
+
+  use_shaing_correction: array_typing.BoolScalar
+  shaing_chi_i_multiplier: array_typing.FloatScalar
+  shaing_chi_e_multiplier: array_typing.FloatScalar
+  shaing_D_e_multiplier: array_typing.FloatScalar
+  shaing_blend_rho_norm_start: array_typing.FloatScalar
+  shaing_blend_rate: array_typing.FloatScalar
+
+
 class AngioniSauterModelConfig(base.NeoclassicalTransportModelConfig):
   """Pydantic model for the Angioni-Sauter neoclassical transport model."""
 
   model_name: Annotated[
       Literal['angioni_sauter'], torax_pydantic.JAX_STATIC
   ] = 'angioni_sauter'
+  use_shaing_correction: bool = False
+  shaing_chi_i_multiplier: float = 0.8
+  shaing_chi_e_multiplier: float = 1.0
+  shaing_D_e_multiplier: float = 1.0
+  shaing_blend_rho_norm_start: float = 0.2
+  shaing_blend_rate: float = 0.005
 
   @override
   def build_model(self) -> 'AngioniSauterModel':
     return AngioniSauterModel()
 
   @override
-  def build_runtime_params(self) -> transport_runtime_params.RuntimeParams:
-    return super().build_runtime_params()
+  # TODO: Why does this not accept t?
+  def build_runtime_params(self) -> RuntimeParams:
+
+    base_kwargs = dataclasses.asdict(super().build_runtime_params())
+    return RuntimeParams(
+        use_shaing_correction=self.use_shaing_correction,
+        shaing_chi_i_multiplier=self.shaing_chi_i_multiplier,
+        shaing_chi_e_multiplier=self.shaing_chi_e_multiplier,
+        shaing_D_e_multiplier=self.shaing_D_e_multiplier,
+        shaing_blend_rho_norm_start=self.shaing_blend_rho_norm_start,
+        shaing_blend_rate=self.shaing_blend_rate,
+        **base_kwargs
+    )
+
+
+def _calculate_blend_alpha(
+    rho_norm: array_typing.FloatVectorFace,
+    start: array_typing.FloatScalar,
+    rate: array_typing.FloatScalar,
+) -> array_typing.FloatVectorFace:
+  """Calculate exponential blend factor for smooth transition between models.
+
+  The blend function is: alpha = 1 - exp(-rate * max(0, rho_norm - start))
+
+  This gives:
+    - Near axis (rho_norm < start): alpha ~ 0 (full Shaing correction)
+    - Far from axis (rho_norm >> start): alpha ~ 1 (pure Angioni-Sauter)
+
+  Args:
+    rho_norm: Normalized toroidal flux coordinate (face grid)
+    start: Rho norm value where blend transition begins
+    rate: Controls transition steepness (higher = sharper transition)
+
+  Returns:
+    Blend factor alpha in range [0, 1]
+  """
+  return 1.0 - jnp.exp(-rate * jnp.maximum(0.0, rho_norm - start))
 
 
 class AngioniSauterModel(base.NeoclassicalTransportModel):
@@ -65,12 +121,46 @@ def _call_implementation(
       geometry: geometry_lib.Geometry,
       core_profiles: state.CoreProfiles,
   ) -> base.NeoclassicalTransport:
-    """Calculates neoclassical transport coefficients."""
-    return _calculate_angioni_sauter_transport(
+    """Calculates neoclassical transport coefficients with smooth blend.
+
+    When use_shaing_correction is enabled, smoothly blends between Shaing
+    (near axis) and Angioni-Sauter (far from axis) models using an exponential
+    transition function.
+    """
+    angioni_sauter = _calculate_angioni_sauter_transport(
         runtime_params=runtime_params,
         geometry=geometry,
         core_profiles=core_profiles,
     )
+    shaing = _calculate_shaing_transport(
+        runtime_params=runtime_params,
+        geometry=geometry,
+        core_profiles=core_profiles,
+    )
+
+    # Calculate blend factor: alpha = 0 near axis (Shaing), alpha = 1 far (A-S)
+    # If correction disabled, alpha = 1 everywhere (pure Angioni-Sauter)
+    alpha = jnp.where(
+        runtime_params.neoclassical.transport.use_shaing_correction,
+        _calculate_blend_alpha(
+            rho_norm=geometry.rho_face_norm,
+            start=runtime_params.neoclassical.transport.shaing_blend_rho_norm_start,
+            rate=runtime_params.neoclassical.transport.shaing_blend_rate,
+        ),
+        1.0,  # Pure Angioni-Sauter when correction disabled
+    )
+
+    # Blend: (1-alpha)*Shaing + alpha*Angioni-Sauter
+    return base.NeoclassicalTransport(
+        chi_neo_i=(1.0 - alpha) * shaing.chi_neo_i
+        + alpha * angioni_sauter.chi_neo_i,
+        chi_neo_e=(1.0 - alpha) * shaing.chi_neo_e
+        + alpha * angioni_sauter.chi_neo_e,
+        D_neo_e=(1.0 - alpha) * shaing.D_neo_e + alpha * angioni_sauter.D_neo_e,
+        V_neo_e=(1.0 - alpha) * shaing.V_neo_e + alpha * angioni_sauter.V_neo_e,
+        V_neo_ware_e=(1.0 - alpha) * shaing.V_neo_ware_e
+        + alpha * angioni_sauter.V_neo_ware_e,
+    )
 
   def __hash__(self) -> int:
     return hash(self.__class__.__name__)
@@ -595,3 +685,100 @@ def _calculate_Lmn(
   )
 
   return Lmn_e, Lmn_i
+
+
+def _calculate_shaing_transport(
+    runtime_params: runtime_params_slice.RuntimeParams,
+    geometry: geometry_lib.Geometry,
+    core_profiles: state.CoreProfiles,
+) -> base.NeoclassicalTransport:
+  """JIT-compatible implementation of the Shaing transport model.
+
+  Args:
+    runtime_params: Runtime parameters.
+    geometry: Geometry object.
+    core_profiles: Core profiles object.
+
+  Returns:
+    Neoclassical transport coefficients.
+  """
+  # Aliases for readability
+  m_ion = core_profiles.A_i * constants.CONSTANTS.m_amu
+  Z_ion = core_profiles.Z_i_face
+  q = core_profiles.q_face
+  delta = geometry.elongation_face
+  F = geometry.F_face  # Note: denoted I in Shaing
+  R = geometry.R_major
+  T_i_J = core_profiles.T_i.face_value() * constants.CONSTANTS.keV_to_J
+  T_e_J = core_profiles.T_e.face_value() * constants.CONSTANTS.keV_to_J
+  ln_Lambda_ii = collisions.calculate_log_lambda_ii(
+      core_profiles.T_i.face_value(),
+      core_profiles.n_i.face_value(),
+      core_profiles.Z_i_face,
+  )
+  ln_Lambda_ei = collisions.calculate_log_lambda_ei(
+      core_profiles.T_e.face_value(), core_profiles.n_e.face_value()
+  )
+  tau_ii = (
+      12
+      * jnp.pi ** (3 / 2)
+      * constants.CONSTANTS.epsilon_0**2
+      * m_ion ** (1 / 2)
+      * T_i_J ** (3 / 2)
+      / (
+          core_profiles.n_i.face_value()
+          * Z_ion**4
+          * constants.CONSTANTS.q_e**4
+          * ln_Lambda_ii
+      )
+  )
+  tau_ei = (
+      3
+      * (2 * jnp.pi) ** (3 / 2)
+      * constants.CONSTANTS.epsilon_0**2
+      * constants.CONSTANTS.m_e ** (1 / 2)
+      * T_e_J ** (3 / 2)
+      / (
+          core_profiles.n_i.face_value()
+          * Z_ion**4
+          * constants.CONSTANTS.q_e**4
+          * ln_Lambda_ei
+      )
+  )
+  nu_ii = 1 / tau_ii
+  nu_ei = 1 / tau_ei
+  v_t_ion = jnp.sqrt(2 * T_i_J / m_ion)
+  v_t_electron = jnp.sqrt(2 * T_e_J / constants.CONSTANTS.m_e)
+  Omega_0_ion = (
+      constants.CONSTANTS.q_e * core_profiles.Z_i_face * geometry.B_0 / m_ion
+  )
+  Omega_0_electron = (
+      constants.CONSTANTS.q_e * geometry.B_0 / constants.CONSTANTS.m_e
+  )
+
+  # Common terms
+  # Large aspect ratio approximation
+  # (See Equation 3, Shaing March 1997)
+  C_1 = (2 * q / (delta * F * R)) ** (1 / 2)
+
+  # Chi i term
+  # Equation 74, Shaing March 1997
+  chi_i = nu_ii * (F * v_t_ion / Omega_0_ion) ** (7 / 3) * C_1 ** (2 / 3)
+
+  # Chi e term
+  # Equation 31, Shaing May 1997
+  chi_e = (
+      nu_ei * (F * v_t_electron / Omega_0_electron) ** (7 / 3) * C_1 ** (2 / 3)
+  )
+
+  return base.NeoclassicalTransport(
+      chi_neo_i=runtime_params.neoclassical.transport.shaing_chi_i_multiplier
+      * chi_i,
+      chi_neo_e=runtime_params.neoclassical.transport.shaing_chi_e_multiplier
+      * chi_e,
+      D_neo_e=runtime_params.neoclassical.transport.shaing_D_e_multiplier
+      * chi_e,
+      # TODO
+      V_neo_e=jnp.zeros_like(geometry.rho_face),
+      V_neo_ware_e=jnp.zeros_like(geometry.rho_face),
+  )