Set TORAX equations to use rhonorm form

Previously, the equations were in mixed form:
temp_ion, temp_el, and dens were in the "rho form", meaning that the vpr derivative and discretization matrix derivatives were with respect to rho. The psi equation was in the "rhonorm form", meaning that these derivatives were with respect to rhonorm. Also, TORAX native derivatives were mixed: fvm gradients are with respect to rhonorm, the mesh is in rhonorm, but vpr was previously with respect to rho. geo.rmax was used to convert the equations to rho or rhonorm forms.

This PR makes things consistent. All governing equations are now posed in rhonorm form, as in the published paper and documentation. vpr and spr are modified to be with respect to rhonorm. These changes had the following impact on the equation solutions and necessitated newly generated reference cases for the CHEASE cases:

1. numerics.largeValueT and numerics.largeValuen, used for the adaptive source to set the pedestal, are not multiplied by vpr like all the other sources. Due to the change in vpr definition, this means that the ratio of the standard sources to this adaptive source changed. The defaults were modified such that the ratio is exactly maintained for the circular geometry cases. However, for the CHEASE cases, due to different geo.rmax, the change in vpr resulted in a slightly different ratio of  largeValueT and largeValuen to the other sources, slightly changing the dynamics of pedestal formation.

2. In CHEASE, vpr = gradient(volume, rho)*rho[-1] is not exactly the same as vpr = gradient(volume, rhon), leading to slight differences

Overall the impact of the above is minor. The average rel differences in the new test/sim references compared to before, are in the 1e-3 - 1e-4 range.

PiperOrigin-RevId: 655641372

torax/calc_coeffs.py

@@ -58,18 +58,16 @@ def calculate_pereverzev_flux(
       * true_ni_face
       * consts.keV2J
       * dynamic_runtime_params_slice.stepper.chi_per
-      / geo.rmax**2
   )
 
   chi_face_per_el = (
       geo.g1_over_vpr_face
       * true_ne_face
       * consts.keV2J
       * dynamic_runtime_params_slice.stepper.chi_per
-      / geo.rmax**2
   )
 
-  d_face_per_el = dynamic_runtime_params_slice.stepper.d_per / geo.rmax
+  d_face_per_el = dynamic_runtime_params_slice.stepper.d_per
   v_face_per_el = (
       core_profiles.ne.face_grad()
       / core_profiles.ne.face_value()
@@ -116,7 +114,7 @@ def calculate_pereverzev_flux(
           > dynamic_runtime_params_slice.profile_conditions.Ped_top,
       ),
       0.0,
-      d_face_per_el * geo.g1_over_vpr_face / geo.rmax,
+      d_face_per_el * geo.g1_over_vpr_face,
   )
 
   v_face_per_el = jnp.where(
@@ -126,7 +124,7 @@ def calculate_pereverzev_flux(
           > dynamic_runtime_params_slice.profile_conditions.Ped_top,
       ),
       0.0,
-      v_face_per_el * geo.g0_face / geo.rmax,
+      v_face_per_el * geo.g0_face,
   )
 
   chi_face_per_ion = chi_face_per_ion.at[0].set(chi_face_per_ion[1])
@@ -404,9 +402,7 @@ def _calc_coeffs_full(
   chi_face_el = transport_coeffs.chi_face_el
   d_face_el = transport_coeffs.d_face_el
   v_face_el = transport_coeffs.v_face_el
-  # TODO(b/351356977): remove rmax from TORAX by normalizing other quantities
-  # like g2g3_over_rho_face, vpr, spr
-  d_face_psi = geo.g2g3_over_rho_face * geo.rmax
+  d_face_psi = geo.g2g3_over_rhon_face
 
   if static_runtime_params_slice.dens_eq:
     if d_face_el is None or v_face_el is None:
@@ -533,19 +529,17 @@ def _calc_coeffs_full(
       * true_ni_face
       * consts.keV2J
       * chi_face_ion
-      / geo.rmax**2
   )
   full_chi_face_el = (
       geo.g1_over_vpr_face
       * true_ne_face
       * consts.keV2J
       * chi_face_el
-      / geo.rmax**2
   )
 
   # entire coefficient preceding dne/dr in particle equation
-  full_d_face_el = geo.g1_over_vpr_face * d_face_el / geo.rmax**2
-  full_v_face_el = geo.g0_face * v_face_el / geo.rmax
+  full_d_face_el = geo.g1_over_vpr_face * d_face_el
+  full_v_face_el = geo.g0_face * v_face_el
 
   # density source terms. Initialize source matrix to zero
   source_mat_nn = jnp.zeros_like(geo.r)
@@ -624,7 +618,6 @@ def _calc_coeffs_full(
       * geo.vpr_face
       * true_ni_face
       * consts.keV2J
-      / geo.rmax
   )
 
   v_heat_face_el += (
@@ -636,7 +629,6 @@ def _calc_coeffs_full(
       * geo.vpr_face
       * true_ne_face
       * consts.keV2J
-      / geo.rmax
   )
 
   # Add phibdot terms to particle transport convection
@@ -647,7 +639,6 @@ def _calc_coeffs_full(
       / geo.Phib
       * geo.r_face_norm
       * geo.vpr_face
-      / geo.rmax
   )
 
   # Add phibdot terms to poloidal flux convection
@@ -660,7 +651,6 @@ def _calc_coeffs_full(
       * sigma_face
       * geo.r_face_norm**2
       / geo.F_face**2
-      / geo.rmax
   )
 
   # Ion and electron heat sources.
@@ -748,7 +738,7 @@ def _calc_coeffs_full(
   # Add effective phibdot heat source terms
 
   # second derivative of volume profile with respect to r_norm
-  vprpr_norm = math_utils.gradient(geo.vpr, geo.r) / geo.rmax**2
+  vprpr_norm = math_utils.gradient(geo.vpr, geo.r_norm)
 
   source_i += (
       1.0

torax/config/runtime_params.py

@@ -185,10 +185,10 @@ class Numerics:
   # numerical (e.g. no. of grid points, other info needed by solver)
   # effective source to dominate PDE in internal boundary condtion location
   # if T != Tped
-  largeValue_T: float = 1.0e10
+  largeValue_T: float = 2.0e10
   # effective source to dominate density PDE in internal boundary condtion
   # location if n != neped
-  largeValue_n: float = 1.0e8
+  largeValue_n: float = 2.0e8
 
 
 # NOMUTANTS -- It's expected for the tests to pass with different defaults.

torax/core_profile_setters.py

@@ -235,7 +235,7 @@ def _prescribe_currents_no_bootstrap(
       1 - geo.r_face_norm**2
   ) ** dynamic_runtime_params_slice.profile_conditions.nu
   # calculate total and Ohmic current profiles
-  denom = _trapz(jformula_face * geo.spr_face, geo.r_face)
+  denom = _trapz(jformula_face * geo.spr_face, geo.r_face_norm)
   if dynamic_runtime_params_slice.profile_conditions.initial_j_is_total_current:
     Ctot = Ip * 1e6 / denom
     jtot_face = jformula_face * Ctot
@@ -345,7 +345,7 @@ def _prescribe_currents_with_bootstrap(
   jformula_face = (
       1 - geo.r_face_norm**2
   ) ** dynamic_runtime_params_slice.profile_conditions.nu
-  denom = _trapz(jformula_face * geo.spr_face, geo.r_face)
+  denom = _trapz(jformula_face * geo.spr_face, geo.r_face_norm)
   # calculate total and Ohmic current profiles
   if dynamic_runtime_params_slice.profile_conditions.initial_j_is_total_current:
     Ctot = Ip * 1e6 / denom
@@ -495,21 +495,21 @@ def _update_psi_from_j(
   psi_constraint = (
       dynamic_runtime_params_slice.profile_conditions.Ip
       * 1e6
-      * (16 * jnp.pi**4 * constants.CONSTANTS.mu0 * geo.B0)
-      / (geo.g2g3_over_rho_face[-1] * geo.F_face[-1])
+      * (16 * jnp.pi**4 * constants.CONSTANTS.mu0 * geo.B0 * geo.rmax)
+      / (geo.g2g3_over_rhon_face[-1] * geo.F_face[-1])
       * geo.rmax
   )
 
   y = currents.jtot_hires * geo.vpr_hires
   assert y.ndim == 1
   assert geo.r_hires.ndim == 1
   integrated = math_utils.cumulative_trapezoid(
-      geo.r_hires, y, initial=jnp.zeros(())
+      geo.r_hires_norm, y, initial=jnp.zeros(())
   )
   scale = jnp.concatenate((
       jnp.zeros((1,)),
-      (8 * jnp.pi**3 * constants.CONSTANTS.mu0 * geo.B0)
-      / (geo.F_hires[1:] * geo.Rmaj * geo.g2g3_over_rho_hires[1:]),
+      (8 * jnp.pi**3 * constants.CONSTANTS.mu0 * geo.B0 * geo.rmax)
+      / (geo.F_hires[1:] * geo.Rmaj * geo.g2g3_over_rhon_hires[1:]),
   ))
   # dpsi_dr on the cell grid
   dpsi_dr_hires = scale * integrated
@@ -612,8 +612,8 @@ def initial_core_profiles(
     psi_constraint = (
         dynamic_runtime_params_slice.profile_conditions.Ip
         * 1e6
-        * (16 * jnp.pi**4 * constants.CONSTANTS.mu0 * geo.B0)
-        / (geo.g2g3_over_rho_face[-1] * geo.F_face[-1])
+        * (16 * jnp.pi**4 * constants.CONSTANTS.mu0 * geo.B0 * geo.rmax)
+        / (geo.g2g3_over_rhon_face[-1] * geo.F_face[-1])
         * geo.rmax
     )
     psi = cell_variable.CellVariable(
@@ -873,8 +873,8 @@ def compute_boundary_conditions(
       'psi': dict(
           right_face_grad_constraint=Ip
           * 1e6
-          * (16 * jnp.pi**4 * constants.CONSTANTS.mu0 * geo.B0)
-          / (geo.g2g3_over_rho_face[-1] * geo.F_face[-1])
+          * (16 * jnp.pi**4 * constants.CONSTANTS.mu0 * geo.B0 * geo.rmax)
+          / (geo.g2g3_over_rhon_face[-1] * geo.F_face[-1])
           * geo.rmax,
           right_face_constraint=None,
       ),
@@ -906,7 +906,7 @@ def _get_jtot_hires(
   jformula_hires = (
       1 - geo.r_hires_norm**2
   ) ** dynamic_runtime_params_slice.profile_conditions.nu
-  denom = _trapz(jformula_hires * geo.spr_hires, geo.r_hires)
+  denom = _trapz(jformula_hires * geo.spr_hires, geo.r_hires_norm)
   if dynamic_runtime_params_slice.profile_conditions.initial_j_is_total_current:
     Ctot_hires = (
         dynamic_runtime_params_slice.profile_conditions.Ip * 1e6 / denom

torax/examples/iterhybrid_predictor_corrector.py

@@ -68,12 +68,6 @@
             # likely be increased further beyond this default.
             'dtmult': 50,
             'dt_reduction_factor': 3,
-            # effective source to dominate PDE in internal boundary condtion
-            # location if T != Tped
-            'largeValue_T': 1.0e10,
-            # effective source to dominate density PDE in internal boundary\
-            # condtion location if n != neped
-            'largeValue_n': 1.0e8,
         },
     },
     'geometry': {

torax/examples/iterhybrid_rampup.py

@@ -77,12 +77,6 @@
             # likely be increased further beyond this default.
             'dtmult': 30,
             'dt_reduction_factor': 3,
-            # effective source to dominate PDE in internal boundary condtion
-            # location if T != Tped
-            'largeValue_T': 1.0e10,
-            # effective source to dominate density PDE in internal boundary
-            # condtion location if n != neped
-            'largeValue_n': 1.0e8,
         },
     },
     'geometry': {