diff --git a/docs/configuration.rst b/docs/configuration.rst index 64bbddf5a..863c1c024 100644 --- a/docs/configuration.rst +++ b/docs/configuration.rst @@ -1868,6 +1868,15 @@ It is recommended to not set ``qlknn_model_name``, or :math:`|R/L_{ne}|` value below which :math:`V_{eff}` is used instead of :math:`D_{eff}`, if ``DV_effective==True``. +``DV_effective_smooth_width`` (float [default = 0.01]) + Width in dimensionless GyroBohm-normalized particle flux units + (:math:`\Gamma_e / \Gamma_{GB}`) over which down-gradient transport + transitions smoothly from :math:`V_{eff}` to :math:`D_{eff}`. If ``0.0``, uses + a sharp step-function transition. Note that QuaLiKiz normalizes using major + radius :math:`R_{major}` rather than minor radius :math:`a`, so the default + (:math:`0.01`) corresponds roughly to the TGLF default (:math:`0.001`) in SI + units. + ``rotation_multiplier`` (float [default = 1.0]) Multiplier for :math:`v_{E\times B}` in the rotation correction factor. @@ -1911,6 +1920,15 @@ Runtime parameters for the TGLFNN-UKAEA model. If you use this model, please cit :math:`|R/L_{ne}|` value below which :math:`V_{eff}` is used instead of :math:`D_{eff}`, if ``DV_effective==True``. +``DV_effective_smooth_width`` (float [default = 0.001]) + Width in dimensionless GyroBohm-normalized particle flux units + (:math:`\Gamma_e / \Gamma_{GB}`) over which down-gradient transport + transitions smoothly from :math:`V_{eff}` to :math:`D_{eff}`. If ``0.0``, uses + a sharp step-function transition. Note that TGLF normalizes using minor + radius :math:`a` rather than :math:`R_{major}`, so the default + (:math:`0.001`) corresponds roughly to the QuaLiKiz default (:math:`0.01`) in + SI units. + ``rotation_multiplier`` (float [default = 1.0]) Multiplier for :math:`v_{E\times B}^{\text{shear}}`. @@ -1978,6 +1996,15 @@ Runtime parameters for the QuaLiKiz model. :math:`|R/L_{ne}|` value below which :math:`V_{eff}` is used instead of :math:`D_{eff}`, if ``DV_effective==True``. +``DV_effective_smooth_width`` (float [default = 0.01]) + Width in dimensionless GyroBohm-normalized particle flux units + (:math:`\Gamma_e / \Gamma_{GB}`) over which down-gradient transport + transitions smoothly from :math:`V_{eff}` to :math:`D_{eff}`. If ``0.0``, uses + a sharp step-function transition. Note that QuaLiKiz normalizes using major + radius :math:`R_{major}` rather than minor radius :math:`a`, so the default + (:math:`0.01`) corresponds roughly to the TGLF default (:math:`0.001`) in SI + units. + tglf ^^^^ @@ -2015,6 +2042,15 @@ Runtime parameters for the TGLF model. If you want to use TORAX with TGLF, see :math:`|R/L_{ne}|` value below which :math:`V_{eff}` is used instead of :math:`D_{eff}`, if ``DV_effective==True``. +``DV_effective_smooth_width`` (float [default = 0.001]) + Width in dimensionless GyroBohm-normalized particle flux units + (:math:`\Gamma_e / \Gamma_{GB}`) over which down-gradient transport + transitions smoothly from :math:`V_{eff}` to :math:`D_{eff}`. If ``0.0``, uses + a sharp step-function transition. Note that TGLF normalizes using minor + radius :math:`a` rather than :math:`R_{major}`, so the default + (:math:`0.001`) corresponds roughly to the QuaLiKiz default (:math:`0.01`) in + SI units. + ``collisionality_multiplier`` (float [default = 1.0]) Collisionality multiplier. diff --git a/torax/_src/transport_model/pydantic_model.py b/torax/_src/transport_model/pydantic_model.py index f0aaea768..f76d69c87 100644 --- a/torax/_src/transport_model/pydantic_model.py +++ b/torax/_src/transport_model/pydantic_model.py @@ -109,6 +109,12 @@ class QLKNNTransportModel(pydantic_model_base.ComponentTransportBase): DV_effective: Effective D / effective V approach for particle transport. An_min: Minimum |R/Lne| below which effective V is used instead of effective D. + DV_effective_smooth_width: Particle flux width in dimensionless + GyroBohm-normalized units (Gamma_e / Gamma_GB) over which down-gradient + transport transitions smoothly from effective V to effective D. If 0.0, + uses a sharp step transition. Note that QuaLiKiz normalizes with major + radius R_major rather than minor radius a, so the default (0.01) is chosen + to be consistent in SI units with TGLF (0.001). rotation_multiplier: Multiplier for rotation. rotation_mode: Mode for rotation, either HALF_RADIUS, FULL_RADIUS or OFF. shear_suppression_alpha: Alpha parameter for Waltz rule applied to @@ -131,8 +137,9 @@ class QLKNNTransportModel(pydantic_model_base.ComponentTransportBase): avoid_big_negative_s: bool = True smag_alpha_correction: bool = True q_sawtooth_proxy: bool = True - DV_effective: bool = False + DV_effective: Annotated[bool, torax_pydantic.JAX_STATIC] = False An_min: pydantic.PositiveFloat = 0.05 + DV_effective_smooth_width: pydantic.NonNegativeFloat = 0.01 rotation_multiplier: pydantic.NonNegativeFloat = 1.0 rotation_mode: Annotated[ qualikiz_based_transport_model.RotationMode, torax_pydantic.JAX_STATIC @@ -195,6 +202,7 @@ def build_runtime_params( q_sawtooth_proxy=self.q_sawtooth_proxy, DV_effective=self.DV_effective, An_min=self.An_min, + DV_effective_smooth_width=self.DV_effective_smooth_width, rotation_multiplier=self.rotation_multiplier, rotation_mode=self.rotation_mode, shear_suppression_alpha=self.shear_suppression_alpha, @@ -209,6 +217,20 @@ class TGLFNNukaeaTransportModel(pydantic_model_base.ComponentTransportBase): Attributes: model_name: The transport model to use. Hardcoded to 'tglfnn-ukaea'. machine: The machine type to use. Either 'step' or 'multimachine'. + rotation_multiplier: Multiplier for rotation. + use_rotation: Whether to use rotation shear in the model. + DV_effective: Effective D / effective V approach for particle transport. + An_min: Minimum |R/Lne| below which effective V is used instead of effective + D. + DV_effective_smooth_width: Particle flux width in dimensionless + GyroBohm-normalized units (Gamma_e / Gamma_GB) over which down-gradient + transport transitions smoothly from effective V to effective D. If 0.0, + uses a sharp step transition. Note that TGLF normalizes with minor radius + a rather than R_major, so the default (0.001) is chosen to be consistent + in SI units with QuaLiKiz (0.01). + collisionality_multiplier: Collisionality multiplier. + max_normalized_collisionality: Maximum normalized collisionality passed to + the model. """ model_name: Annotated[Literal['tglfnn-ukaea'], torax_pydantic.JAX_STATIC] = ( @@ -220,8 +242,9 @@ class TGLFNNukaeaTransportModel(pydantic_model_base.ComponentTransportBase): rotation_multiplier: pydantic.NonNegativeFloat = 1.0 use_rotation: Annotated[bool, torax_pydantic.JAX_STATIC] = False # Quasilinear transport options - DV_effective: bool = False + DV_effective: Annotated[bool, torax_pydantic.JAX_STATIC] = False An_min: pydantic.PositiveFloat = 0.05 + DV_effective_smooth_width: pydantic.NonNegativeFloat = 0.001 collisionality_multiplier: float = 1.0 max_normalized_collisionality: pydantic.PositiveFloat = float('inf') @@ -239,6 +262,7 @@ def build_runtime_params( return tglfnn_ukaea_transport_model.RuntimeParams( DV_effective=self.DV_effective, An_min=self.An_min, + DV_effective_smooth_width=self.DV_effective_smooth_width, rotation_multiplier=self.rotation_multiplier, use_rotation=self.use_rotation, collisionality_multiplier=self.collisionality_multiplier, diff --git a/torax/_src/transport_model/qualikiz_transport_model.py b/torax/_src/transport_model/qualikiz_transport_model.py index e88ec5d64..473d0dc8a 100644 --- a/torax/_src/transport_model/qualikiz_transport_model.py +++ b/torax/_src/transport_model/qualikiz_transport_model.py @@ -473,6 +473,12 @@ class QualikizTransportModelConfig(pydantic_model_base.ComponentTransportBase): DV_effective: Effective D / effective V approach for particle transport. An_min: Minimum |R/Lne| below which effective V is used instead of effective D. + DV_effective_smooth_width: Particle flux width in dimensionless + GyroBohm-normalized units (Gamma_e / Gamma_GB) over which down-gradient + transport transitions smoothly from effective V to effective D. If 0.0, + uses a sharp step transition. Note that QuaLiKiz normalizes with major + radius R_major rather than minor radius a, so the default (0.01) is chosen + to be consistent in SI units with TGLF (0.001). """ model_name: Annotated[Literal['qualikiz'], torax_pydantic.JAX_STATIC] = ( @@ -485,8 +491,9 @@ class QualikizTransportModelConfig(pydantic_model_base.ComponentTransportBase): avoid_big_negative_s: bool = True smag_alpha_correction: bool = True q_sawtooth_proxy: bool = True - DV_effective: bool = False + DV_effective: Annotated[bool, torax_pydantic.JAX_STATIC] = False An_min: pydantic.PositiveFloat = 0.05 + DV_effective_smooth_width: pydantic.NonNegativeFloat = 0.01 rotation_multiplier: pydantic.NonNegativeFloat = 1.0 rotation_mode: Annotated[ qualikiz_based_transport_model.RotationMode, torax_pydantic.JAX_STATIC @@ -507,6 +514,7 @@ def build_runtime_params(self, t: chex.Numeric) -> RuntimeParams: q_sawtooth_proxy=self.q_sawtooth_proxy, DV_effective=self.DV_effective, An_min=self.An_min, + DV_effective_smooth_width=self.DV_effective_smooth_width, rotation_multiplier=self.rotation_multiplier, rotation_mode=self.rotation_mode, **base_kwargs, diff --git a/torax/_src/transport_model/quasilinear_transport_model.py b/torax/_src/transport_model/quasilinear_transport_model.py index 2adccafcc..ea7164e02 100644 --- a/torax/_src/transport_model/quasilinear_transport_model.py +++ b/torax/_src/transport_model/quasilinear_transport_model.py @@ -204,8 +204,9 @@ def calculate_alpha( class RuntimeParams(runtime_params_lib.ComponentRuntimeParams): """Shared parameters for Quasilinear models.""" - DV_effective: bool + DV_effective: bool = dataclasses.field(metadata={"static": True}) An_min: float + DV_effective_smooth_width: float @jax.jit @@ -242,6 +243,110 @@ def calculate_normalized_logarithmic_gradient( return result +def calculate_dv_effective( + particle_flux_SI: jax.Array, + normalized_particle_flux: jax.Array, + n_e: cell_variable.CellVariable, + geo: geometry.Geometry, + gradient_reference_length: array_typing.FloatScalar, + An_min: array_typing.FloatScalar, + DV_effective_smooth_width: array_typing.FloatScalar, + two_point_mask: array_typing.BoolVectorFace | None = None, +) -> tuple[jax.Array, jax.Array]: + """Calculates effective diffusivity and convectivity for DV_effective mode. + + Splits the particle flux between pure effective diffusion (`D_eff`) and pure + effective convection (`V_eff`): + d_face_el = diffusion_weight * D_eff + v_face_el = (1 - diffusion_weight) * V_eff + where `diffusion_weight = grad_weight * flux_weight` in [0, 1] assigns + transport to convection for small density gradients (|L_ref / L_ne| < An_min) + or up-gradient flux, and to diffusion otherwise. + + When `DV_effective_smooth_width` == 0.0, `grad_weight` and `flux_weight` are + step functions. When `DV_effective_smooth_width` > 0.0, they use polynomial + splines chosen so that `d_face_el` and `v_face_el` have continuous first + derivatives (C1) across both transitions: + 1. `grad_weight`: Quartic spline `x^3 * (4 - 3*x)` on + `x = clip(|L_ref / L_ne| / An_min, 0, 1)`. Because `D_eff` scales as + `1 / (L_ref / L_ne)`, the `x^3` factor makes `grad_weight * D_eff` scale + as `x^2`, giving zero value and zero derivative at zero density gradient + along with zero derivative at `An_min`. + 2. `flux_weight`: Quadratic spline `x * (2 - x)` on + `x = clip(downgrad_flux / DV_effective_smooth_width, 0, 1)`. Because + `D_eff` and `V_eff` are already linear in flux, a linear factor `x` near + zero flux makes `flux_weight * D_eff` quadratic in flux, giving zero + derivative across flux reversal along with zero derivative at + `DV_effective_smooth_width`. + + Args: + particle_flux_SI: Electron particle flux in SI units [m^-2 s^-1]. + normalized_particle_flux: Dimensionless GyroBohm-normalized particle flux. + n_e: Electron density CellVariable [m^-3]. + geo: Torus geometry. + gradient_reference_length: Reference length L_ref for An_min [m]. + An_min: Normalized logarithmic gradient threshold |L_ref / L_ne| below which + transport transitions from diffusion to convection. + DV_effective_smooth_width: Particle flux width in dimensionless + GyroBohm-normalized units (Gamma_e / Gamma_GB) over which down-gradient + transport transitions from convection to diffusion. If 0.0, uses a sharp + step-function transition. + two_point_mask: Optional boolean mask for 2-point face gradients. + + Returns: + Tuple of (d_face_el [m^2/s], v_face_el [m/s]). + """ + n_e_face = n_e.face_value() + dn_e_drhon = n_e.face_grad(two_point_mask=two_point_mask) + dn_e_drhon_safe = jnp.where(dn_e_drhon == 0.0, 1.0, dn_e_drhon) + + # Effective diffusivity (pure D) and convectivity (pure V) that each + # individually reproduce the full particle flux. + D_eff = jnp.where( + dn_e_drhon == 0.0, + 0.0, + -particle_flux_SI / (dn_e_drhon_safe * geo.g1_over_vpr2_face * geo.rho_b), + ) + V_eff = particle_flux_SI / (n_e_face * geo.g0_over_vpr_face * geo.rho_b) + + dn_e_drmid = n_e.face_grad( + x=geo.r_mid, + x_left=geo.r_mid_face[0], + x_right=geo.r_mid_face[-1], + two_point_mask=two_point_mask, + ) + lref_over_lne = -dn_e_drmid * gradient_reference_length / n_e_face + + # 1. Smooth density-gradient weight: 0 at zero gradient, 1 for + # |lref_over_lne| >= An_min. + norm_grad = jnp.clip(jnp.abs(lref_over_lne) / An_min, 0.0, 1.0) + grad_weight = norm_grad**3 * (4.0 - 3.0 * norm_grad) + + # 2. Smooth down-gradient flux weight: 0 for up-gradient flux, 1 for + # down-gradient flux >= DV_effective_smooth_width. + sign_lref_over_lne = jnp.where(lref_over_lne >= 0.0, 1.0, -1.0) + downgrad_flux = normalized_particle_flux * sign_lref_over_lne + # Avoid division by zero if DV_effective_smooth_width == 0.0, where + # sharp_diffusion_mask is used instead. + smooth_width_safe = jnp.where( + DV_effective_smooth_width > 0.0, DV_effective_smooth_width, 1.0 + ) + norm_flux = jnp.clip(downgrad_flux / smooth_width_safe, 0.0, 1.0) + flux_weight = norm_flux * (2.0 - norm_flux) + + sharp_diffusion_mask = (jnp.abs(lref_over_lne) >= An_min) & ( + downgrad_flux >= 0.0 + ) + diffusion_weight = jnp.where( + DV_effective_smooth_width == 0.0, + sharp_diffusion_mask, + grad_weight * flux_weight, + ) + d_face_el = diffusion_weight * D_eff + v_face_el = (1.0 - diffusion_weight) * V_eff + return d_face_el, v_face_el + + @jax.tree_util.register_dataclass @dataclasses.dataclass(frozen=True) class QuasilinearInputs: @@ -422,13 +527,11 @@ def _make_core_transport( transport: RuntimeParams, geo: geometry.Geometry, core_profiles: state.CoreProfiles, - gradient_reference_length: chex.Numeric, - gyrobohm_flux_reference_length: chex.Numeric, + gradient_reference_length: array_typing.FloatScalar, + gyrobohm_flux_reference_length: array_typing.FloatScalar, two_point_mask: array_typing.BoolVectorFace | None = None, ) -> transport_coeffs.TransportCoeffs: """Converts model output to TransportCoeffs.""" - constants = constants_module.CONSTANTS - # conversion to SI units (note that n is normalized here) # Convert the electron particle flux from GB (pfe) to SI units. @@ -457,35 +560,23 @@ def _make_core_transport( / quasilinear_inputs.lref_over_lte ) * quasilinear_inputs.chiGB - # Effective D / Effective V approach. - # For small density gradients or up-gradient transport, set pure effective - # convection. Otherwise pure effective diffusion. - def DV_effective_approach() -> tuple[jax.Array, jax.Array]: - # The geo.rho_b is to unnormalize the face_grad. - Deff = -pfe_SI / ( - core_profiles.n_e.face_grad(two_point_mask=two_point_mask) - * geo.g1_over_vpr2_face - * geo.rho_b - + constants.eps - ) - Veff = pfe_SI / ( - core_profiles.n_e.face_value() * geo.g0_over_vpr_face * geo.rho_b + if transport.DV_effective: + d_face_el, v_face_el = calculate_dv_effective( + particle_flux_SI=pfe_SI, + normalized_particle_flux=pfe, + n_e=core_profiles.n_e, + geo=geo, + gradient_reference_length=gradient_reference_length, + An_min=transport.An_min, + DV_effective_smooth_width=transport.DV_effective_smooth_width, + two_point_mask=two_point_mask, ) - Deff_mask = ( - ((pfe >= 0) & (quasilinear_inputs.lref_over_lne >= 0)) - | ((pfe < 0) & (quasilinear_inputs.lref_over_lne < 0)) - ) & (abs(quasilinear_inputs.lref_over_lne) >= transport.An_min) - Veff_mask = jnp.invert(Deff_mask) - # Veff_mask is where to use effective V only, so zero out D there. - d_face_el = jnp.where(Veff_mask, 0.0, Deff) - # And vice versa - v_face_el = jnp.where(Deff_mask, 0.0, Veff) - return d_face_el, v_face_el - - # Scaled D approach. Scale electron diffusivity to electron heat - # conductivity (this has some physical motivations), - # and set convection to then match total particle transport - def Dscaled_approach() -> tuple[jax.Array, jax.Array]: + else: + # Scaled D approach. Scale electron diffusivity to electron heat + # conductivity (this has some physical motivations), + # and set convection to then match total particle transport. + # TODO(b/567403838): Create a helper function, calculate_d_scaled, and + # fix gradient to use dn_e_drhon for consistency with TGLF in follow-up. chex.assert_rank(pfe, 1) d_face_el = chi_face_el v_face_el = ( @@ -496,13 +587,7 @@ def Dscaled_approach() -> tuple[jax.Array, jax.Array]: * geo.g1_over_vpr2_face * geo.rho_b**2 ) / (geo.g0_over_vpr_face * geo.rho_b) - return d_face_el, v_face_el # pyrefly: ignore[bad-return] - d_face_el, v_face_el = jax.lax.cond( - transport.DV_effective, - DV_effective_approach, - Dscaled_approach, - ) return transport_coeffs.TransportCoeffs( chi_face_ion=chi_face_ion, chi_face_el=chi_face_el, diff --git a/torax/_src/transport_model/tests/qualikiz_based_transport_model_test.py b/torax/_src/transport_model/tests/qualikiz_based_transport_model_test.py index fd2fe2637..e93fb2c87 100644 --- a/torax/_src/transport_model/tests/qualikiz_based_transport_model_test.py +++ b/torax/_src/transport_model/tests/qualikiz_based_transport_model_test.py @@ -374,8 +374,9 @@ class QualikizBasedTransportModelConfig( avoid_big_negative_s: bool = True smag_alpha_correction: bool = True q_sawtooth_proxy: bool = True - DV_effective: bool = False + DV_effective: Annotated[bool, torax_pydantic.JAX_STATIC] = False An_min: pydantic.PositiveFloat = 0.05 + DV_effective_smooth_width: pydantic.NonNegativeFloat = 0.01 rotation_multiplier: pydantic.NonNegativeFloat = 1.0 rotation_mode: Annotated[ qualikiz_based_transport_model.RotationMode, torax_pydantic.JAX_STATIC @@ -397,6 +398,7 @@ def build_runtime_params(self, t: chex.Numeric): q_sawtooth_proxy=self.q_sawtooth_proxy, DV_effective=self.DV_effective, An_min=self.An_min, + DV_effective_smooth_width=self.DV_effective_smooth_width, rotation_multiplier=self.rotation_multiplier, rotation_mode=self.rotation_mode, **base_kwargs, diff --git a/torax/_src/transport_model/tests/quasilinear_transport_model_test.py b/torax/_src/transport_model/tests/quasilinear_transport_model_test.py index d7ea44997..034019156 100644 --- a/torax/_src/transport_model/tests/quasilinear_transport_model_test.py +++ b/torax/_src/transport_model/tests/quasilinear_transport_model_test.py @@ -158,15 +158,80 @@ def test_quasilinear_transport_model_dveff( }, }) core_transport = model(*model_inputs) + # On the magnetic axis, the density gradient is zero by symmetry, so + # effective D/V mode always assigns pure convection there regardless of the + # minimum gradient threshold. Check the axis and the bulk separately. + self.assertNotEqual(core_transport.total.v_face_el[0], 0.0) + self.assertEqual(core_transport.total.d_face_el[0] == 0.0, DV_effective) self.assertEqual( - (np.sum(np.abs(core_transport.total.v_face_el)) == 0.0), + (np.sum(np.abs(core_transport.total.v_face_el[1:])) == 0.0), expected_zero_v_face_el, ) self.assertEqual( - (np.sum(np.abs(core_transport.total.d_face_el)) == 0.0), + (np.sum(np.abs(core_transport.total.d_face_el[1:])) == 0.0), expected_zero_d_face_el, ) + @parameterized.product( + A_n_and_pfe=[ + # Zero density gradient (singular point for 1/A_n; purely convective). + (0.0, 0.5), + # Down-gradient within both |A_n| < An_min and pfe < smooth_width + # transition zones (blended D and V in smooth mode). + (0.02, 0.01), + # Up-gradient transport (A_n * pfe < 0; purely convective). + (0.5, -0.5), + # Down-gradient above both thresholds (purely diffusive). + (0.5, 0.5), + ], + DV_effective_smooth_width=[0.0, 0.02], + ) + def test_dv_effective_conserves_flux_and_non_negative_d( + self, A_n_and_pfe, DV_effective_smooth_width + ): + A_n_target, pfe_val = A_n_and_pfe + _, model_inputs = _get_model_and_model_inputs({ + 'core_transport_models': {'quasilinear': {'model_name': 'quasilinear'}}, + }) + _, geo, _, _, two_point_mask = model_inputs + L_ref = 3.0 + An_min = 0.05 + n_e_0 = 1.0e20 + slope = -A_n_target * n_e_0 * geo.rho_b / L_ref + n_e = cell_variable.CellVariable( + value=n_e_0 + slope * geo.rho_norm, + face_centers=geo.rho_face_norm, + left_face_grad_constraint=jnp.asarray(slope), + right_face_constraint=n_e_0 + slope, + right_face_grad_constraint=None, + ) + pfe = jnp.full_like(geo.rho_face_norm, pfe_val) + pfe_SI = pfe * (n_e.face_value() / geo.a_minor) * 4.0 + + d_face_el, v_face_el = quasilinear_transport_model.calculate_dv_effective( + particle_flux_SI=pfe_SI, + normalized_particle_flux=pfe, + n_e=n_e, + geo=geo, + gradient_reference_length=L_ref, + An_min=An_min, + DV_effective_smooth_width=DV_effective_smooth_width, + two_point_mask=two_point_mask, + ) + reconstructed_flux = ( + -d_face_el + * n_e.face_grad(two_point_mask=two_point_mask) + * geo.g1_over_vpr2_face + * geo.rho_b + + v_face_el * n_e.face_value() * geo.g0_over_vpr_face * geo.rho_b + ) + np.testing.assert_allclose( + reconstructed_flux, pfe_SI, rtol=1e-12, atol=1e-8 + ) + self.assertTrue(np.all(d_face_el >= 0.0)) + if A_n_target * pfe_val <= 0.0: + np.testing.assert_allclose(d_face_el, 0.0, atol=1e-15) + def test_calculate_chiGB(self): """Tests that chiGB is calculated correctly.""" @@ -555,8 +620,9 @@ class QuasilinearTransportConfig( model_name: Annotated[Literal['quasilinear'], torax_pydantic.JAX_STATIC] = ( 'quasilinear' ) - DV_effective: bool = False + DV_effective: Annotated[bool, torax_pydantic.JAX_STATIC] = False An_min: pydantic.PositiveFloat = 0.05 + DV_effective_smooth_width: pydantic.NonNegativeFloat = 0.0 def build_transport_model(self) -> FakeQuasilinearTransportModel: return FakeQuasilinearTransportModel() @@ -568,6 +634,7 @@ def build_runtime_params( return quasilinear_transport_model.RuntimeParams( DV_effective=self.DV_effective, An_min=self.An_min, + DV_effective_smooth_width=self.DV_effective_smooth_width, **base_kwargs, ) diff --git a/torax/_src/transport_model/tests/tglf_based_transport_model_test.py b/torax/_src/transport_model/tests/tglf_based_transport_model_test.py index 1a57aa33c..28aeedaf7 100644 --- a/torax/_src/transport_model/tests/tglf_based_transport_model_test.py +++ b/torax/_src/transport_model/tests/tglf_based_transport_model_test.py @@ -226,6 +226,43 @@ def test_max_normalized_collisionality_caps_xnue(self): uncapped.XNUE[~above_cap], ) + def test_tglf_based_transport_model_dv_effective(self): + """Tests that DV_effective switches between effective D/V and scaled D.""" + torax_config_false, model_inputs_false = _get_config_and_model_inputs({ + "core_transport_models": { + "tglf_based": { + "model_name": "tglf_based", + "DV_effective": False, + }, + }, + }) + transport_model_false = torax_config_false.transport.build_transport_model() + core_transport_false = transport_model_false(*model_inputs_false) + + torax_config_true, model_inputs_true = _get_config_and_model_inputs({ + "core_transport_models": { + "tglf_based": { + "model_name": "tglf_based", + "DV_effective": True, + }, + }, + }) + transport_model_true = torax_config_true.transport.build_transport_model() + core_transport_true = transport_model_true(*model_inputs_true) + + # With DV_effective=False (scaled D), d_face_el is set equal to chi_face_el. + np.testing.assert_allclose( + core_transport_false.total.d_face_el, + core_transport_false.total.chi_face_el, + ) + # With DV_effective=True, effective D/V differs from scaled D. + self.assertFalse( + np.allclose( + core_transport_true.total.d_face_el, + core_transport_false.total.d_face_el, + ) + ) + @dataclasses.dataclass(frozen=True, eq=False) class FakeTGLFBasedTransportModel( @@ -288,6 +325,7 @@ class TGLFBasedTransportModelConfig( model_name: Annotated[Literal["tglf_based"], torax_pydantic.JAX_STATIC] = ( "tglf_based" ) + DV_effective: Annotated[bool, torax_pydantic.JAX_STATIC] = False max_normalized_collisionality: float = float("inf") # pylint: disable=undefined-variable @@ -300,8 +338,9 @@ def build_runtime_params(self, t: chex.Numeric): base_kwargs = dataclasses.asdict(super().build_runtime_params(t)) return tglf_based_transport_model.RuntimeParams( # DV_effective and An_min are inherited from QuasilinearTransportModel - DV_effective=False, + DV_effective=self.DV_effective, An_min=0.05, + DV_effective_smooth_width=0.01, use_rotation=True, rotation_multiplier=1.0, collisionality_multiplier=1.0, diff --git a/torax/_src/transport_model/tglf/tglf_transport_model.py b/torax/_src/transport_model/tglf/tglf_transport_model.py index 64069cf33..2a2e45dc1 100644 --- a/torax/_src/transport_model/tglf/tglf_transport_model.py +++ b/torax/_src/transport_model/tglf/tglf_transport_model.py @@ -329,6 +329,12 @@ class TGLFTransportModelConfig(pydantic_model_base.ComponentTransportBase): DV_effective: Effective D / effective V approach for particle transport. An_min: Minimum |R/Lne| below which effective V is used instead of effective D. + DV_effective_smooth_width: Particle flux width in dimensionless + GyroBohm-normalized units (Gamma_e / Gamma_GB) over which down-gradient + transport transitions smoothly from effective V to effective D. If 0.0, + uses a sharp step transition. Note that TGLF normalizes with minor radius + a rather than R_major, so the default (0.001) is chosen to be consistent + in SI units with QuaLiKiz (0.01). collisionality_multiplier: Collisionality multiplier. max_normalized_collisionality: Maximum normalized collisionality passed to the model. Acts as a ceiling to mitigate unreliable transport predictions @@ -352,6 +358,7 @@ class TGLFTransportModelConfig(pydantic_model_base.ComponentTransportBase): rotation_multiplier: pydantic.NonNegativeFloat = 1.0 DV_effective: Annotated[bool, torax_pydantic.JAX_STATIC] = False An_min: pydantic.PositiveFloat = 0.05 + DV_effective_smooth_width: pydantic.NonNegativeFloat = 0.001 collisionality_multiplier: float = 1.0 max_normalized_collisionality: pydantic.PositiveFloat = float('inf') tglf_settings: Annotated[ @@ -435,6 +442,7 @@ def build_runtime_params(self, t: chex.Numeric) -> RuntimeParams: collisionality_multiplier=self.collisionality_multiplier, max_normalized_collisionality=self.max_normalized_collisionality, An_min=self.An_min, + DV_effective_smooth_width=self.DV_effective_smooth_width, tglf_settings=self.tglf_settings, **base_kwargs, ) diff --git a/torax/_src/transport_model/tglf_based_transport_model.py b/torax/_src/transport_model/tglf_based_transport_model.py index 64e92ce7a..81f93cb8b 100644 --- a/torax/_src/transport_model/tglf_based_transport_model.py +++ b/torax/_src/transport_model/tglf_based_transport_model.py @@ -511,11 +511,10 @@ def _make_core_transport( # pyrefly: ignore[bad-override] Q_i = ion_heat_flux_GB * tglf_inputs.Q_GB # [W/m^2] Gamma_e = electron_particle_flux_GB * tglf_inputs.GAMMA_GB # [s^-1/m^2] - # Total thermal power and particle rate. + # Total thermal power. dV_drho = geo.vpr_face / geo.rho_b P_e = Q_e * dV_drho # [W] P_i = Q_i * dV_drho # [W] - S_e = Gamma_e * dV_drho # [s^-1] # Convert from power to chi. # Note: g1/vpr = ⟨(∇ρₙ)²⟩ ∂V/∂ρₙ, and has units [m]. @@ -542,39 +541,31 @@ def _make_core_transport( # pyrefly: ignore[bad-override] eps=1e-7, ) - # Convert from particle rate to D, V using effective - # diffusivity/convectivity method. This sets purely diffusive transport in - # regions where the flux is with the temperature gradient, otherwise it - # sets purely convective transport. - D_eff = math_utils.safe_divide( - num=-S_e, - denom=core_profiles.n_e.face_grad(two_point_mask=two_point_mask) - * geo.g1_over_vpr_face, - eps=1e-7, - ) - V_eff = math_utils.safe_divide( - num=S_e, - denom=core_profiles.n_e.face_value() * geo.g0_face, - eps=1e-7, - ) - D_eff = jnp.where(jnp.isfinite(D_eff), D_eff, 0.0) - V_eff = jnp.where(jnp.isfinite(V_eff), V_eff, 0.0) - D_eff_mask = ((S_e >= 0) & (tglf_inputs.lref_over_lne >= 0)) | ( - (S_e < 0) & (tglf_inputs.lref_over_lne < 0) - ) - # For stability, we also set purely diffusive transport at some minimum - # threshold of the density gradient. - D_eff_mask &= ( - abs(tglf_inputs.lref_over_lne) - >= transport.An_min * geo.a_minor / geo.R_major - ) - V_eff_mask = jnp.logical_not(D_eff_mask) - d_face_el = jnp.where(D_eff_mask, D_eff, 0.0) - v_face_el = jnp.where(V_eff_mask, V_eff, 0.0) + if transport.DV_effective: + d_face_el, v_face_el = quasilinear_transport_model.calculate_dv_effective( + particle_flux_SI=Gamma_e, + normalized_particle_flux=electron_particle_flux_GB, + n_e=core_profiles.n_e, + geo=geo, + gradient_reference_length=geo.R_major, + An_min=transport.An_min, + DV_effective_smooth_width=transport.DV_effective_smooth_width, + two_point_mask=two_point_mask, + ) + else: + # Scaled D approach. Scale electron diffusivity to electron heat + # conductivity (this has some physical motivations), + # and set convection to then match total particle transport. + # TODO(b/567403838): Create a helper function, calculate_d_scaled. + dn_e_drhon = core_profiles.n_e.face_grad(two_point_mask=two_point_mask) + d_face_el = chi_e + v_face_el = ( + Gamma_e + d_face_el * dn_e_drhon * geo.g1_over_vpr2_face * geo.rho_b + ) / (core_profiles.n_e.face_value() * geo.g0_over_vpr_face * geo.rho_b) return transport_coeffs.TransportCoeffs( - chi_face_ion=chi_i, # pyrefly: ignore[bad-argument-type] - chi_face_el=chi_e, # pyrefly: ignore[bad-argument-type] + chi_face_ion=chi_i, + chi_face_el=chi_e, d_face_el=d_face_el, v_face_el=v_face_el, ) diff --git a/torax/tests/test_data/test_changing_config_after.nc b/torax/tests/test_data/test_changing_config_after.nc index 1b47fd9aa..fa6b34fd5 100644 Binary files a/torax/tests/test_data/test_changing_config_after.nc and b/torax/tests/test_data/test_changing_config_after.nc differ diff --git a/torax/tests/test_data/test_changing_config_before.nc b/torax/tests/test_data/test_changing_config_before.nc index 874b5c054..2bebc435e 100644 Binary files a/torax/tests/test_data/test_changing_config_before.nc and b/torax/tests/test_data/test_changing_config_before.nc differ diff --git a/torax/tests/test_data/test_imas_profiles_and_geo.nc b/torax/tests/test_data/test_imas_profiles_and_geo.nc index 9e326e438..2e5263e0e 100644 Binary files a/torax/tests/test_data/test_imas_profiles_and_geo.nc and b/torax/tests/test_data/test_imas_profiles_and_geo.nc differ diff --git a/torax/tests/test_data/test_iterbaseline_mockup.nc b/torax/tests/test_data/test_iterbaseline_mockup.nc index efeafac63..7a85d27a7 100644 Binary files a/torax/tests/test_data/test_iterbaseline_mockup.nc and b/torax/tests/test_data/test_iterbaseline_mockup.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_mockup.nc b/torax/tests/test_data/test_iterhybrid_mockup.nc index d3d69e0fa..96740bd56 100644 Binary files a/torax/tests/test_data/test_iterhybrid_mockup.nc and b/torax/tests/test_data/test_iterhybrid_mockup.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc index a242e1cfe..adfb104b4 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_Lmode_combined.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_Lmode_combined.nc index a4b6c6e7b..ef2242734 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_Lmode_combined.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_Lmode_combined.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc index 24dba3aba..80b308295 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc index 824986d53..3d55880fc 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc index 6091c996f..21ef21fc8 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc index 9ab2d64f1..9483f1ade 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc index bb7bd659e..d64a7e728 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_imas.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_imas.nc index 4a407b825..66ff890e5 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_imas.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_imas.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc index 3ae73b4ab..979101302 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios.nc index fe22d2bda..6814801b0 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.nc index 84894974c..fff2f5a42 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_z_eff.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_z_eff.nc index 051ba0a47..3e39dbf68 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_z_eff.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_z_eff.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_neoclassical.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_neoclassical.nc index e0a5dd1f2..b39f9a8be 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_neoclassical.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_neoclassical.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_rotation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_rotation.nc index 51ee918b9..63602a1a7 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_rotation.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_rotation.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc index 691e369f9..081064291 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea.nc index 0e591fcef..bec91c4cb 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea_rotation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea_rotation.nc index 3969aa6df..be2e5c839 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea_rotation.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea_rotation.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc index 0d37b0f99..35f9d80d7 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc index 5338e0bae..a9f0bf775 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc index cfd02cfdf..dd3e97742 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_rampup.nc b/torax/tests/test_data/test_iterhybrid_rampup.nc index 111c0f1ab..6708453c5 100644 Binary files a/torax/tests/test_data/test_iterhybrid_rampup.nc and b/torax/tests/test_data/test_iterhybrid_rampup.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_rampup_sawtooth.nc b/torax/tests/test_data/test_iterhybrid_rampup_sawtooth.nc index 058408d23..3021369cd 100644 Binary files a/torax/tests/test_data/test_iterhybrid_rampup_sawtooth.nc and b/torax/tests/test_data/test_iterhybrid_rampup_sawtooth.nc differ diff --git a/torax/tests/test_data/test_ne_qlknn_deff_veff.nc b/torax/tests/test_data/test_ne_qlknn_deff_veff.nc index 93e66344b..faa367c0d 100644 Binary files a/torax/tests/test_data/test_ne_qlknn_deff_veff.nc and b/torax/tests/test_data/test_ne_qlknn_deff_veff.nc differ