Extrapolate chi_neo_e and V_neo_ware_e

Author: theo-brown

torax/_src/neoclassical/transport/angioni_sauter.py

@@ -284,52 +284,63 @@ def _calculate_angioni_sauter_transport(
 
   # Q_e = - chi_e * n_e * dT/drho = B_e2 * T_e / dpsi/drho
   # Q_i = - chi_i * n_i * dT/drho = B_i2 * T_i / dpsi/drho
-  chi_neo_e = -Be2 / (
-      core_profiles.n_e.face_value()
-      * dlnte_dpsi
-      * (dpsi_drhon / geometry.rho_b) ** 2
+  chi_neo_e_bulk = -Be2[1:] / (
+      core_profiles.n_e.face_value()[1:]
+      * dlnte_dpsi[1:]
+      * (dpsi_drhon[1:] / geometry.rho_b) ** 2
       + constants.CONSTANTS.eps
   )
-  chi_neo_i = -Bi2 / (
-      core_profiles.n_i.face_value()
-      * dlnti_dpsi
-      * (dpsi_drhon / geometry.rho_b) ** 2
+  chi_neo_e_axis = jnp.array([chi_neo_e_bulk[0]])
+  chi_neo_e = jnp.concatenate([chi_neo_e_axis, chi_neo_e_bulk])
+
+  chi_neo_i_bulk = -Bi2[1:] / (
+      core_profiles.n_i.face_value()[1:]
+      * dlnti_dpsi[1:]
+      * (dpsi_drhon[1:] / geometry.rho_b) ** 2
       + constants.CONSTANTS.eps
   )
+  chi_neo_i_axis = jnp.array([chi_neo_i_bulk[0]])
+  chi_neo_i = jnp.concatenate([chi_neo_i_axis, chi_neo_i_bulk])
 
   # Decomposition of particle flux Be1 = Gamma * dpsi_drho. Page 1232+1233.
 
   # Diffusive part of particle flux
   # D_e * dn_e/drho  = - L00 *dlog(n_e)/dpsi / dpsi/drho
-  D_neo_e = -Lmn_e[:, 0, 0] / (
-      core_profiles.n_e.face_value() * (dpsi_drhon / geometry.rho_b) ** 2
+  D_neo_e_bulk = -Lmn_e[1:, 0, 0] / (
+      core_profiles.n_e.face_value()[1:] * (dpsi_drhon[1:] / geometry.rho_b) ** 2
       + constants.CONSTANTS.eps
   )
+  D_neo_e_axis = jnp.array([D_neo_e_bulk[0]])
+  D_neo_e = jnp.concatenate([D_neo_e_axis, D_neo_e_bulk])
 
   # Convective part of particle flux, apart from the Ware Pinch term
   # V*n*dpsi/rho = (L00+L01)*dlog(Te)/dpsi + (1-Rpe)/Rpe*L00*dlog(ni)/dpsi +
   # (1-Rpe)/Rpe * (L00+alpha*L03) *dlog(Ti)/dpsi
-  V_neo_e = (
-      (Lmn_e[:, 0, 0] + Lmn_e[:, 0, 1]) * dlnte_dpsi
-      + (1 - Rpe) / Rpe * Lmn_e[:, 0, 0] * dlnni_dpsi
-      + (1 - Rpe) / Rpe * (Lmn_e[:, 0, 0] + alpha * Lmn_e[:, 0, 3]) * dlnti_dpsi
+  V_neo_e_bulk = (
+      (Lmn_e[1:, 0, 0] + Lmn_e[1:, 0, 1]) * dlnte_dpsi[1:]
+      + (1 - Rpe[1:]) / Rpe[1:] * Lmn_e[1:, 0, 0] * dlnni_dpsi[1:]
+      + (1 - Rpe[1:]) / Rpe[1:] * (Lmn_e[1:, 0, 0] + alpha[1:] * Lmn_e[1:, 0, 3]) * dlnti_dpsi[1:]
   ) / (
-      dpsi_drhon / geometry.rho_b * core_profiles.n_e.face_value()
+      dpsi_drhon[1:] / geometry.rho_b * core_profiles.n_e.face_value()[1:]
       + constants.CONSTANTS.eps
   )
+  V_neo_e_axis = jnp.array([V_neo_e_bulk[0]])
+  V_neo_e = jnp.concatenate([V_neo_e_axis, V_neo_e_bulk])
 
   # Ware pinch term component of particle convection
   # V_ware*n*dpsi/rho = L02*<E_parallel * B>/<B^2>
-  V_neo_ware_e = (
-      Lmn_e[:, 0, 2]
-      * E_parallel
+  V_neo_ware_e_bulk = (
+      Lmn_e[1:, 0, 2]
+      * E_parallel[1:]
       / (
           geometry.B_0
-          * (dpsi_drhon / geometry.rho_b)
-          * core_profiles.n_e.face_value()
+          * (dpsi_drhon[1:] / geometry.rho_b)
+          * core_profiles.n_e.face_value()[1:]
           + constants.CONSTANTS.eps
       )
   )
+  V_neo_ware_e_axis = jnp.atleast_1d(V_neo_ware_e_bulk[0])
+  V_neo_ware_e = jnp.concatenate([V_neo_ware_e_axis, V_neo_ware_e_bulk])
 
   return base.NeoclassicalTransport(
       chi_neo_i=chi_neo_i,

torax/_src/neoclassical/transport/tests/angioni_sauter_test.py

@@ -186,7 +186,7 @@ def test_angioni_sauter_with_shaing_against_reference_values(self):
 # generation of the Kmn matrix.
 _ANGIONI_SAUTER_REFERENCE_VALUES = base.NeoclassicalTransport(
   chi_neo_i=np.array([
-      -0.0,
+      0.01220085,
       0.01220085,
       0.02223608,
       0.03117304,
@@ -199,7 +199,7 @@ def test_angioni_sauter_with_shaing_against_reference_values(self):
       0.0591895,
   ]),
   chi_neo_e=np.array([
-      -0.0,
+      -0.00210023,
       -0.00210023,
       -0.0030792,
       -0.00388683,
@@ -212,7 +212,7 @@ def test_angioni_sauter_with_shaing_against_reference_values(self):
       -0.00750323,
   ]),
   D_neo_e=np.array([
-      0.0,
+      0.00011698,
       0.00011698,
       0.00021105,
       0.00028474,
@@ -225,7 +225,7 @@ def test_angioni_sauter_with_shaing_against_reference_values(self):
       0.0002924,
   ]),
   V_neo_e=np.array([
-      0.0,
+      1.07951440e-05,
       1.07951440e-05,
       1.11015003e-05,
       1.54065751e-05,
@@ -238,7 +238,7 @@ def test_angioni_sauter_with_shaing_against_reference_values(self):
       1.18868626e-03,
   ]),
   V_neo_ware_e=np.array([
-      -0.0,
+      -0.00038114,
       -0.00038114,
       -0.00041759,
       -0.00037123,
@@ -255,7 +255,7 @@ def test_angioni_sauter_with_shaing_against_reference_values(self):
 # Shaing correction only affects ions, so we can reuse the other values
 _ANGIONI_SAUTER_SHAING_REFERENCE_VALUES = base.NeoclassicalTransport(
   chi_neo_i=np.array([
-    0.20237419,
+    0.20382857,
     0.17130245,
     0.03031974,
     0.02593766,

torax/examples/step_flattop_bgb.py

@@ -164,8 +164,6 @@
         "transport": {
             "model_name": "angioni_sauter",
             "use_shaing_ion_correction": True,
-            "V_e_min": -1e-6,
-            "V_e_max": 1e-6,
         },
     },
     "numerics": {