Simplify internal API by adding a config flag for whether to use vloo…

…p or not

This is more in line with the other parameters in the config, e.g. nbar and normalize_to_nbar

torax/config/profile_conditions.py

@@ -38,14 +38,14 @@ class ProfileConditions(
   # Total plasma current in MA
   # Note that if Ip_from_parameters=False in geometry, then this Ip will be
   # overwritten by values from the geometry data.
-  # If vloop_lcfs is not None, then this Ip is used as an
-  # initial condition ONLY.
   Ip_tot: interpolated_param.TimeInterpolatedInput = 15.0
 
   # Boundary condition at LCFS for Vloop ( = dpsi_lcfs/dt )
-  # If this is `None` the boundary condition for the psi equation at each timestep
-  # will instead be taken from `Ip_tot`.
-  vloop_lcfs: interpolated_param.TimeInterpolatedInput | None = None
+  # If use_vloop_lcfs_boundary_condition is True, then the specfied Vloop at
+  # the LCFS is used as the boundary condition for the psi equation; otherwise,
+  # Ip is used as the boundary condition.
+  use_vloop_lcfs_boundary_condition: bool = False
+  vloop_lcfs: interpolated_param.TimeInterpolatedInput = 0.0
 
   # Temperature boundary conditions at r=Rmin. If this is `None` the boundary
   # condition will instead be taken from `Ti` and `Te` at rhon=1.
@@ -216,7 +216,8 @@ class DynamicProfileConditions:
   """Prescribed values and boundary conditions for the core profiles."""
 
   Ip_tot: array_typing.ScalarFloat
-  vloop_lcfs: array_typing.ScalarFloat | None
+  vloop_lcfs: array_typing.ScalarFloat
+  use_vloop_lcfs_as_boundary_condition: bool
   Ti_bound_right: array_typing.ScalarFloat
   Te_bound_right: array_typing.ScalarFloat
   # Temperature profiles defined on the cell grid.

torax/core_profile_setters.py

@@ -639,8 +639,8 @@ def _init_psi_and_current(
   Returns:
     Refined core profiles.
   """
-  use_vloop_bc = (
-      dynamic_runtime_params_slice.profile_conditions.vloop_lcfs is not None
+  use_vloop_as_boundary_condition = (
+      dynamic_runtime_params_slice.profile_conditions.use_vloop_lcfs_as_boundary_condition
   )
 
   # Case 1: retrieving psi from the profile conditions, using the prescribed profile and Ip_tot
@@ -652,7 +652,9 @@ def _init_psi_and_current(
     )
 
     # Set the BCs to ensure the correct Ip_tot
-    if use_vloop_bc:
+    if use_vloop_as_boundary_condition:
+      # Extrapolate the value of psi at the LCFS from the dpsi/drho constraint
+      # to achieve the desired Ip_tot
       right_face_grad_constraint = None
       right_face_constraint = (
           dynamic_runtime_params_slice.profile_conditions.psi[-1]
@@ -695,15 +697,18 @@ def _init_psi_and_current(
     )
 
     # Set the BCs, with the option of scaling to the given Ip_tot
-    if use_vloop_bc and geo.Ip_from_parameters:
+    if use_vloop_as_boundary_condition and geo.Ip_from_parameters:
+      # Extrapolate the value of psi at the LCFS from the dpsi/drho constraint
       right_face_grad_constraint = None
       right_face_constraint = (
           geo.psi_from_Ip[-1] + dpsi_drhonorm_edge * geo.drho_norm / 2
       )
-    elif use_vloop_bc:
+    elif use_vloop_as_boundary_condition:
+      # Use the psi from the equilibrium as the right face constraint
       right_face_grad_constraint = None
       right_face_constraint = geo.psi_from_Ip[-1]
     else:
+      # Use the dpsi/drho calculated above as the right face gradient constraint
       right_face_grad_constraint = dpsi_drhonorm_edge
       right_face_constraint = None
 
@@ -1111,24 +1116,21 @@ def compute_boundary_conditions_for_t_plus_dt(
       'psi': dict(
           right_face_grad_constraint=(
               _calculate_psi_grad_constraint_from_Ip_tot(
-                  dynamic_runtime_params_slice_t_plus_dt,
-                  geo_t_plus_dt,
+                  Ip_tot=dynamic_runtime_params_slice_t_plus_dt.profile_conditions.Ip_tot,
+                  geo=geo_t_plus_dt,
               )
-              if dynamic_runtime_params_slice_t_plus_dt.profile_conditions.vloop_lcfs
-              is None
+              if not dynamic_runtime_params_slice_t.profile_conditions.use_vloop_lcfs_as_boundary_condition
               else None
           ),
           right_face_constraint=(
               _calculate_psi_value_constraint_from_vloop(
-                  dt,
-                  static_runtime_params_slice,
-                  dynamic_runtime_params_slice_t,
-                  dynamic_runtime_params_slice_t_plus_dt,
-                  core_profiles_t,
-                  geo_t_plus_dt,
+                  dt=dt,
+                  vloop_lcfs_t=dynamic_runtime_params_slice_t.vloop_lcfs,
+                  vloop_lcfs_t_plus_dt=dynamic_runtime_params_slice_t_plus_dt.vloop_lcfs,
+                  psi_lcfs_t=core_profiles_t.psi.value[-1],
+                  theta=static_runtime_params_slice.numerics.theta,
               )
-              if dynamic_runtime_params_slice_t_plus_dt.profile_conditions.vloop_lcfs
-              is not None
+              if dynamic_runtime_params_slice_t.profile_conditions.use_vloop_lcfs_as_boundary_condition
               else None
           ),
       ),

torax/tests/test_data/test_psichease_ip_chease_vloop.py

@@ -19,4 +19,7 @@
 
 
 CONFIG = copy.deepcopy(test_psichease_ip_chease.CONFIG)
+CONFIG['runtime_params']['profile_conditions'][
+    'use_vloop_lcfs_boundary_condition'
+] = True
 CONFIG['runtime_params']['profile_conditions']['vloop_lcfs'] = 6.0

torax/tests/test_data/test_psichease_ip_parameters_prescribed_jtot_vloop.py

@@ -19,4 +19,7 @@
 
 
 CONFIG = copy.deepcopy(test_psichease_prescribed_jtot.CONFIG)
+CONFIG['runtime_params']['profile_conditions'][
+    'use_vloop_lcfs_boundary_condition'
+] = True
 CONFIG['runtime_params']['profile_conditions']['vloop_lcfs'] = 6.0

torax/tests/test_data/test_psichease_ip_parameters_vloop.py

@@ -19,4 +19,7 @@
 
 
 CONFIG = copy.deepcopy(test_psichease_ip_parameters.CONFIG)
+CONFIG['runtime_params']['profile_conditions'][
+    'use_vloop_lcfs_boundary_condition'
+] = True
 CONFIG['runtime_params']['profile_conditions']['vloop_lcfs'] = 6.0