Skip to content

Smooth $C^1$ DV_effective particle transport splitting via DV_effective_smooth_width (default 0.01). #2522

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Sep 29, 2026
Merged

Smooth $C^1$ DV_effective particle transport splitting via DV_effective_smooth_width (default 0.01). #2522

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions docs/configuration.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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}}`.

Expand Down Expand Up @@ -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
^^^^
Expand Down Expand Up @@ -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.

Expand Down
28 changes: 26 additions & 2 deletions torax/_src/transport_model/pydantic_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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] = (
Expand All @@ -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')

Expand All @@ -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,
Expand Down
10 changes: 9 additions & 1 deletion torax/_src/transport_model/qualikiz_transport_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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] = (
Expand All @@ -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
Expand All @@ -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,
Expand Down
163 changes: 124 additions & 39 deletions torax/_src/transport_model/quasilinear_transport_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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 = (
Expand All @@ -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,
Expand Down
4 changes: 3 additions & 1 deletion torax/_src/transport_model/tests/qualikiz_based_transport_model_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand Down
Loading
Loading