From 9f8175718e81bdb0f67a25fb1e03d2237f9e93a1 Mon Sep 17 00:00:00 2001 From: Philippe Hamel Date: Wed, 3 Dec 2025 14:22:36 -0800 Subject: [PATCH] Add runtime_param to control the poloidal_rotation_multiplier. PiperOrigin-RevId: 839914032 --- docs/configuration.rst | 3 +++ torax/_src/neoclassical/formulas.py | 10 ++++++---- torax/_src/neoclassical/pydantic_model.py | 3 +++ torax/_src/neoclassical/runtime_params.py | 2 ++ torax/_src/output_tools/post_processing.py | 1 + torax/_src/physics/rotation.py | 7 ++++--- 6 files changed, 19 insertions(+), 7 deletions(-) diff --git a/docs/configuration.rst b/docs/configuration.rst index 665cfa906..968709811 100644 --- a/docs/configuration.rst +++ b/docs/configuration.rst @@ -1979,6 +1979,9 @@ transport Note that clipping to the desired range will be applied to :math:`V_\mathrm{neo}` and :math:`V_\mathrm{neo, ware}` separately. +``poloidal_velocity_multiplier`` (float [default = 1.0]) + Multiplier for the poloidal velocity. + restart ------- diff --git a/torax/_src/neoclassical/formulas.py b/torax/_src/neoclassical/formulas.py index 4c27efcb8..20808f3aa 100644 --- a/torax/_src/neoclassical/formulas.py +++ b/torax/_src/neoclassical/formulas.py @@ -237,7 +237,7 @@ def calculate_poloidal_velocity( B_tor: array_typing.FloatVectorFace, B_total_squared: array_typing.FloatVectorFace, geo: geometry_lib.Geometry, - rotation_multiplier: array_typing.FloatScalar = 1.0, + poloidal_velocity_multiplier: array_typing.FloatScalar = 1.0, ) -> cell_variable.CellVariable: """Computes the neoclassical ion poloidal velocity profile. @@ -259,8 +259,10 @@ def calculate_poloidal_velocity( B_tor: Toroidal magnetic field on the face grid [T]. B_total_squared: Total magnetic field (toroidal + poloidal) on the face grid [T]. - geo : Geometry - rotation_multiplier: A multiplier to apply to the poloidal velocity. + geo: Geometry + poloidal_velocity_multiplier: A multiplier to apply to the poloidal + velocity. + Returns: v_pol : Poloidal velocity profile [m/s]. """ @@ -298,7 +300,7 @@ def calculate_poloidal_velocity( / (constants.CONSTANTS.q_e * Z_i) ) - v_pol = rotation_multiplier * v_pol + v_pol = poloidal_velocity_multiplier * v_pol return cell_variable.CellVariable( value=geometry_lib.face_to_cell(v_pol), diff --git a/torax/_src/neoclassical/pydantic_model.py b/torax/_src/neoclassical/pydantic_model.py index 768fc8fbc..8f610c6ab 100644 --- a/torax/_src/neoclassical/pydantic_model.py +++ b/torax/_src/neoclassical/pydantic_model.py @@ -17,6 +17,7 @@ from typing import Any import pydantic +from torax._src import array_typing from torax._src.neoclassical import neoclassical_models from torax._src.neoclassical import runtime_params as runtime_params_lib from torax._src.neoclassical.bootstrap_current import sauter as sauter_current @@ -40,6 +41,7 @@ class Neoclassical(torax_pydantic.BaseModelFrozen): transport: ( transport_zeros.ZerosModelConfig | angioni_sauter.AngioniSauterModelConfig ) = pydantic.Field(discriminator="model_name") + poloidal_velocity_multiplier: array_typing.FloatScalar = 1.0 @pydantic.model_validator(mode="before") @classmethod @@ -63,6 +65,7 @@ def build_runtime_params(self) -> runtime_params_lib.RuntimeParams: bootstrap_current=self.bootstrap_current.build_runtime_params(), conductivity=self.conductivity.build_runtime_params(), transport=self.transport.build_runtime_params(), + poloidal_velocity_multiplier=self.poloidal_velocity_multiplier, ) def build_models(self) -> neoclassical_models.NeoclassicalModels: diff --git a/torax/_src/neoclassical/runtime_params.py b/torax/_src/neoclassical/runtime_params.py index 570ef4b2b..3ca922a99 100644 --- a/torax/_src/neoclassical/runtime_params.py +++ b/torax/_src/neoclassical/runtime_params.py @@ -15,6 +15,7 @@ import dataclasses import jax +from torax._src import array_typing from torax._src.neoclassical.bootstrap_current import runtime_params as bootstrap_current_runtime_params from torax._src.neoclassical.conductivity import runtime_params as conductivity_runtime_params from torax._src.neoclassical.transport import runtime_params as transport_runtime_params @@ -27,3 +28,4 @@ class RuntimeParams: bootstrap_current: bootstrap_current_runtime_params.RuntimeParams conductivity: conductivity_runtime_params.RuntimeParams transport: transport_runtime_params.RuntimeParams + poloidal_velocity_multiplier: array_typing.FloatScalar diff --git a/torax/_src/output_tools/post_processing.py b/torax/_src/output_tools/post_processing.py index 0a62e7025..f907ba1df 100644 --- a/torax/_src/output_tools/post_processing.py +++ b/torax/_src/output_tools/post_processing.py @@ -821,6 +821,7 @@ def cumulative_values(): toroidal_velocity=sim_state.core_profiles.toroidal_velocity, pressure_thermal_i=sim_state.core_profiles.pressure_thermal_i, geo=sim_state.geometry, + poloidal_velocity_multiplier=runtime_params.neoclassical.poloidal_velocity_multiplier, ) return PostProcessedOutputs( diff --git a/torax/_src/physics/rotation.py b/torax/_src/physics/rotation.py index 9498d3661..8995764ef 100644 --- a/torax/_src/physics/rotation.py +++ b/torax/_src/physics/rotation.py @@ -87,7 +87,7 @@ def calculate_rotation( toroidal_velocity: cell_variable.CellVariable, pressure_thermal_i: cell_variable.CellVariable, geo: geometry.Geometry, - rotation_multiplier: float = 1.0, + poloidal_velocity_multiplier: array_typing.FloatScalar = 1.0, ): """Calculates quantities related to the rotation of the plasma. @@ -101,7 +101,8 @@ def calculate_rotation( toroidal_velocity: Toroidal velocity profile as a cell variable. pressure_thermal_i: Pressure profile as a cell variable. geo: Geometry object. - rotation_multiplier: A multiplier to apply to the poloidal velocity. + poloidal_velocity_multiplier: A multiplier to apply to the poloidal + velocity. Returns: v_ExB: ExB velocity profile on the face grid [m/s]. @@ -129,7 +130,7 @@ def calculate_rotation( B_tor=B_tor_face, B_total_squared=B_total_squared_face, geo=geo, - rotation_multiplier=rotation_multiplier, + poloidal_velocity_multiplier=poloidal_velocity_multiplier, ) Er = _calculate_radial_electric_field(