Extrapolate chi_neo_e and V_neo_ware_e

Author: theo-brown

torax/_src/neoclassical/transport/angioni_sauter.py

@@ -284,12 +284,14 @@ 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_e_axis = jnp.atleast_1d(chi_neo_e_bulk[0])
+  chi_neo_e = jnp.concatenate([chi_neo_e_axis, chi_neo_e_bulk])
   chi_neo_i = -Bi2 / (
       core_profiles.n_i.face_value()
       * dlnti_dpsi
@@ -320,16 +322,18 @@ def _calculate_angioni_sauter_transport(
 
   # 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/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": {