Geometry refactor for clarity and consistency.

1. Remove unnecessary usage of J = F/(BR) 2. Rename intermediate.RBPhi to intermediate.F 3. Make toroidal flux Phi the required "intermediate" geometry input, and define Phib, rho, rho_norm all with respect to Phi. Better since not all geometry codes provide rho directly 4. Reduce usage of rmax. Replace with Phi or Phib 5. Rename r and r_face to rho and rho_face throughout TORAX. "r" is a legacy from when only circular geometry was present, and geometric radius was the same as rho. It is clearer to just call everything rho (toroidal flux coordinate). (3) led to differences in rho and rho_norm for CHEASE geometry test cases due to rho now being derived from Phi, which is interpolated first onto the TORAX grid. For circular geometry there is no interpolation so no differences are seen. The differences led to rel errors of order 1e-3-1e-4 for all CHEASE cases. The references were updated. test_iterhybrid_rampup (and the run_simulation_test based on it) pass locally but needed coarser rel error (now 1e-6) to pass TAP tests. PiperOrigin-RevId: 657021031
google-deepmind · Jul 29, 2024 · 18e99d6 · 18e99d6
1 parent 097a493
commit 18e99d6
Show file tree

Hide file tree

Showing 54 changed files with 351 additions and 414 deletions.
diff --git a/torax/calc_coeffs.py b/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

diff --git a/torax/config/runtime_params.py b/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.

diff --git a/torax/core_profile_setters.py b/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

diff --git a/torax/examples/iterhybrid_predictor_corrector.py b/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': {

diff --git a/torax/examples/iterhybrid_rampup.py b/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': {