Add Shaing corrections to Angioni-Sauter model

- Add smooth transition between models
- Fix conversion from psi with 2pi factor
- Tune default parameters for an 'ok' match with NCLASS on STEP and ITER cases

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,42 @@
 # 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_ion_correction: array_typing.BoolScalar
+  shaing_ion_multiplier: array_typing.FloatScalar
+  shaing_blend_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_ion_correction: bool = False
+  shaing_ion_multiplier: float = 1.8
+  shaing_blend_start: float = 0.2
+  shaing_blend_rate: float = 5.0
 
   @override
   def build_model(self) -> 'AngioniSauterModel':
     return AngioniSauterModel()
 
   @override
-  def build_runtime_params(self) -> transport_runtime_params.RuntimeParams:
-    return super().build_runtime_params()
+  def build_runtime_params(self) -> RuntimeParams:
+    base_kwargs = dataclasses.asdict(super().build_runtime_params())
+    return RuntimeParams(
+        use_shaing_ion_correction=self.use_shaing_ion_correction,
+        shaing_ion_multiplier=self.shaing_ion_multiplier,
+        shaing_blend_start=self.shaing_blend_start,
+        shaing_blend_rate=self.shaing_blend_rate,
+        **base_kwargs
+    )
 
 
 class AngioniSauterModel(base.NeoclassicalTransportModel):
@@ -65,13 +89,48 @@ 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, chi_ion is smoothly blended 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 sigmoid blend weight for Angioni-Sauter (alpha)
+    # If correction disabled: alpha = 1 (pure Angioni-Sauter)
+    # If correction enabled: alpha varies smoothly with rho_norm
+    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_start,
+            rate=runtime_params.neoclassical.transport.shaing_blend_rate,
+        ),
+        1.0,  # Pure Angioni-Sauter when correction disabled
+    )
+
+    return base.NeoclassicalTransport(
+        # Ion transport blend: (1-alpha)*Shaing + alpha*Angioni-Sauter
+        chi_neo_i=(1.0 - alpha) * shaing.chi_neo_i
+        + alpha * angioni_sauter.chi_neo_i,
+        # Electron transport: pure Angioni-Sauter
+        chi_neo_e=angioni_sauter.chi_neo_e,
+        D_neo_e=angioni_sauter.D_neo_e,
+        V_neo_e=angioni_sauter.V_neo_e,
+        V_neo_ware_e=angioni_sauter.V_neo_ware_e,
+    )
+
   def __hash__(self) -> int:
     return hash(self.__class__.__name__)
 
@@ -587,3 +646,111 @@ def _calculate_Lmn(
   )
 
   return Lmn_e, Lmn_i
+
+
+def _calculate_shaing_transport(
+    runtime_params: runtime_params_lib.RuntimeParams,
+    geometry: geometry_lib.Geometry,
+    core_profiles: state.CoreProfiles,
+) -> base.NeoclassicalTransport:
+  """JIT-compatible implementation of the Shaing transport model.
+
+  Currently only implements ion thermal transport. Other contributions are
+  negligible.
+
+  Args:
+    runtime_params: Runtime parameters.
+    geometry: Geometry object.
+    core_profiles: Core profiles object.
+
+  Returns:
+    - Neoclassical transport coefficients.
+    - Radius of validity (in terms of rho_norm) of Shaing model for electrons.
+    - Radius of validity (in terms of rho_norm) of Shaing model for ions.
+  """
+  # Aliases for readability
+  m_ion = core_profiles.A_i * constants.CONSTANTS.m_amu
+  q = core_profiles.q_face
+  kappa = geometry.elongation_face  # Note: denoted delta in Shaing
+  F = geometry.F_face  # Note: denoted I in Shaing
+  R = geometry.R_major_profile_face
+  T_i_J = core_profiles.T_i.face_value() * constants.CONSTANTS.keV_to_J
+
+  # Collisionality
+  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,
+  )
+  tau_ii = collisions.calculate_tau_ii(
+      A_i=core_profiles.A_i,
+      Z_i=core_profiles.Z_i_face,
+      T_i=core_profiles.T_i.face_value(),
+      n_i=core_profiles.n_i.face_value(),
+      ln_Lambda_ii=ln_Lambda_ii,
+  )
+  nu_ii = 1 / tau_ii # Ion-ion collision frequency
+
+  # Thermal velocity
+  v_t_ion = jnp.sqrt(2 * T_i_J / m_ion)
+
+  # Larmor (gyro)frequency
+  Omega_0_ion = (
+      constants.CONSTANTS.q_e * core_profiles.Z_i_face * geometry.B_0 / m_ion
+  )
+
+  # Large aspect ratio approximation (Equation 3, Shaing March 1997)
+  C_1 = (2 * q / (kappa * F * R)) ** (1 / 2)
+
+  # Conversion from flux^2/s -> m^2/s
+  # TODO: make a more informed choice for dpsi_drhon near the axis (currently
+  # we simply copy the value at i=1). This is ok as chi[0] is never used.
+  dpsi_drhon = core_profiles.psi.face_grad()
+  dpsi_drhon = dpsi_drhon.at[0].set(dpsi_drhon[1])
+  conversion_factor = 1 / (dpsi_drhon / (2*jnp.pi*geometry.rho_b)) ** 2
+
+  # Trapped particle fraction (Equation 46, Shaing March 1997)
+  f_t_ion = (F * v_t_ion * C_1**2 / Omega_0_ion) ** (1 / 3)
+
+  # Orbit width in psi coordinates (Equation 73, Shaing March 1997)
+  Delta_psi_ion = (F**2 * v_t_ion**2 * C_1 / Omega_0_ion**2) ** (2 / 3)
+
+  # Chi i term (Equation 74, Shaing March 1997)
+  # psi normalization difference accounted for in conversion_factor
+  chi_i = (nu_ii * Delta_psi_ion**2 / f_t_ion) * conversion_factor
+
+  return base.NeoclassicalTransport(
+      chi_neo_i=runtime_params.neoclassical.transport.shaing_ion_multiplier
+      * chi_i,
+      chi_neo_e=jnp.zeros_like(geometry.rho_face),
+      D_neo_e=jnp.zeros_like(geometry.rho_face),
+      V_neo_e=jnp.zeros_like(geometry.rho_face),
+      V_neo_ware_e=jnp.zeros_like(geometry.rho_face),
+  )
+
+
+def _calculate_blend_alpha(
+    rho_norm: array_typing.FloatVectorFace,
+    start: array_typing.FloatScalar,
+    rate: array_typing.FloatScalar,
+) -> array_typing.FloatVectorFace:
+  """Calculate blending weight between Angioni-Sauter and Shaing models.
+
+  The blend is:
+    result = (1-alpha)*Shaing + alpha*Angioni-Sauter
+  where alpha = 1 / (1 + exp(-2*rate*(rho_norm - start))).
+
+  This gives:
+    - At axis (rho_norm = 0 << start): alpha ~ 0 (pure Shaing)
+    - At start: alpha = 0.5 (equal blend)
+    - 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 is centered
+    rate: Controls transition steepness (higher = sharper transition)
+
+  Returns:
+    Blend factor alpha in range [0, 1]
+  """
+  return 1.0 / (1.0 + jnp.exp(-2.0 * rate * (rho_norm - start)))

torax/_src/physics/collisions.py

@@ -27,6 +27,7 @@
       electron-ion collisions.
     - calculate_log_lambda_ii: Calculates the Coulomb logarithm for ion-ion
       collisions.
+    - calculate_tau_ii: Calculates the ion-ion collision time.
     - _calculate_weighted_Z_eff: Calculates ion mass weighted Z_eff used in
       the equipartion calculation.
     - _calculate_log_tau_e_Z1: Calculates log of electron-ion collision time for
@@ -244,6 +245,41 @@ def calculate_log_lambda_ii(
   return 30.0 - 0.5 * jnp.log(n_i) + 1.5 * jnp.log(T_i_ev) - 3.0 * jnp.log(Z_i)
 
 
+def calculate_tau_ii(
+  A_i: jax.Array,
+  Z_i: jax.Array,
+  T_i: jax.Array,
+  n_i: jax.Array,
+  ln_Lambda_ii: jax.Array,
+) -> jax.Array:
+  """Calculates ion-ion (self) collision time for a single ion species.
+
+  See Wesson 3rd edition p730.
+
+  Args:
+    A_i: Ion mass in amu.
+    Z_i: Ion charge.
+    T_i: Ion temperature in keV.
+    n_i: Ion density in m^-3.
+    ln_Lambda_ii: Coulomb logarithm for ion-ion collisions.
+
+  Returns:
+    Ion-ion collision time in seconds.
+  """
+  log_tau_ii = (
+    jnp.log(12)
+    + 1.5 * jnp.log(jnp.pi)
+    + 2 * jnp.log(constants.CONSTANTS.epsilon_0)
+    + 0.5 * (jnp.log(A_i) + jnp.log(constants.CONSTANTS.m_amu))
+    + 1.5 * (jnp.log(T_i) + jnp.log(constants.CONSTANTS.keV_to_J))
+    - jnp.log(n_i)
+    - 4 * jnp.log(Z_i)
+    - 4 * jnp.log(constants.CONSTANTS.q_e)
+    - jnp.log(ln_Lambda_ii)
+  )
+  return jnp.exp(log_tau_ii)
+
+
 # TODO(b/377225415): generalize to arbitrary number of ions.
 def _calculate_weighted_Z_eff(
     core_profiles: state.CoreProfiles,

torax/examples/step_flattop_bgb.py

@@ -32,9 +32,9 @@
 
 import imas
 import numpy as np
+
 from torax._src import path_utils
-from torax._src.imas_tools.input import core_profiles
-from torax._src.imas_tools.input import loader
+from torax._src.imas_tools.input import core_profiles, loader
 
 # Load IDSs
 path = (
@@ -145,20 +145,16 @@
         "D_e_max": 100.0,
         "V_e_min": -50.0,
         "V_e_max": 50.0,
-        # Patching
-        # Replaces neoclassical in the core, pending potato orbit correction
-        # https://github.com/google-deepmind/torax/issues/1406
-        "apply_inner_patch": True,
-        "rho_inner": 0.05,
-        "chi_e_inner": 1.0,
-        "chi_i_inner": 15.0,
         # Smoothing
         "smooth_everywhere": True,
         "smoothing_width": 0.05,
     },
     "neoclassical": {
         "bootstrap_current": {"model_name": "sauter"},
-        "transport": {"model_name": "angioni_sauter"},
+        "transport": {
+            "model_name": "angioni_sauter",
+            "use_shaing_ion_correction": True,
+        },
     },
     "numerics": {
         "t_initial": 0.0,