Add Vloop boundary condition for the psi equation

Vloop (= dpsi_edge/dt) can be a useful BC for scenarios where the voltage swing of the central solenoid is directly controlled (e.g. in a fully non-inductive scenario, Vloop = 0).

Summary of changes:
- Added new BC for the psi equation.
  - Can be combined with any of the psi initialisation methods:
    - Prescribed psi + Ip_tot rescaling
    - Psi from geometry file
    - Psi from geometry file + Ip_tot rescaling
    - Psi from nu formula + Ip_tot rescaling
  - Added vloop_lcfs to state.CoreProfiles, which is a direct copy from psidot.face_value()[-1].
  - Required adding the core profiles from the previous timestep to the boundary condition computation
- Added psi_right_bc and vloop_lcfs to output.

torax/config/profile_conditions.py

@@ -35,11 +35,18 @@ class ProfileConditions(
 ):
   """Prescribed values and boundary conditions for the core profiles."""
 
-  # total plasma current in MA
+  # 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
+  # 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
+
   # Temperature boundary conditions at r=Rmin. If this is `None` the boundary
   # condition will instead be taken from `Ti` and `Te` at rhon=1.
   Ti_bound_right: interpolated_param.TimeInterpolatedInput | None = None
@@ -175,6 +182,7 @@ class ProfileConditionsProvider(
 
   runtime_params_config: ProfileConditions
   Ip_tot: interpolated_param.InterpolatedVarSingleAxis
+  vloop_lcfs: interpolated_param.InterpolatedVarSingleAxis | None
   Ti_bound_right: (
       interpolated_param.InterpolatedVarSingleAxis
       | interpolated_param.InterpolatedVarTimeRho
@@ -208,6 +216,7 @@ class DynamicProfileConditions:
   """Prescribed values and boundary conditions for the core profiles."""
 
   Ip_tot: array_typing.ScalarFloat
+  vloop_lcfs: array_typing.ScalarFloat | None
   Ti_bound_right: array_typing.ScalarFloat
   Te_bound_right: array_typing.ScalarFloat
   # Temperature profiles defined on the cell grid.

torax/core_profile_setters.py

@@ -34,7 +34,7 @@
 from torax.config import runtime_params_slice
 from torax.fvm import cell_variable
 from torax.geometry import geometry
-from torax.sources import ohmic_heat_source
+from torax.sources import ohmic_heat_source 
 from torax.sources import source_models as source_models_lib
 from torax.sources import source_profiles as source_profiles_lib
 
@@ -529,11 +529,6 @@ def _update_psi_from_j(
   Returns:
     psi: Poloidal flux cell variable.
   """
-  psi_grad_constraint = _calculate_psi_grad_constraint(
-      dynamic_runtime_params_slice,
-      geo,
-  )
-
   y = jtot_hires * geo.spr_hires
   assert y.ndim == 1
   assert geo.rho_hires.ndim == 1
@@ -545,30 +540,48 @@ def _update_psi_from_j(
       (16 * jnp.pi**3 * constants.CONSTANTS.mu0 * geo.Phib)
       / (geo.F_hires[1:] * geo.g2g3_over_rhon_hires[1:]),
   ))
-  # dpsi_dr on the cell grid
+  # dpsi_dr on hires cell grid
   dpsi_drhon_hires = scale * Ip_profile
 
-  # psi on cell grid
+  # psi on hires cell grid
   psi_hires = math_utils.cumulative_trapezoid(
       y=dpsi_drhon_hires, x=geo.rho_hires_norm, initial=0.0
   )
 
   psi_value = jnp.interp(geo.rho_norm, geo.rho_hires_norm, psi_hires)
 
+  # Set the BCs for psi to ensure the correct Ip_tot
+  dpsi_drhonorm_edge = _calculate_psi_grad_constraint_from_Ip_tot(
+            dynamic_runtime_params_slice,
+            geo,
+      )
+
+  if dynamic_runtime_params_slice.profile_conditions.vloop_lcfs is not None:
+    right_face_grad_constraint = None
+    right_face_constraint = (
+        psi_value[-1] + dpsi_drhonorm_edge * geo.drho_norm / 2
+      )
+  else:
+    # Use the dpsi/drho calculated above as the right face gradient constraint
+    right_face_grad_constraint = dpsi_drhonorm_edge
+    right_face_constraint = None
+
+
   psi = cell_variable.CellVariable(
       value=psi_value,
       dr=geo.drho_norm,
-      right_face_grad_constraint=psi_grad_constraint,
+      right_face_grad_constraint=right_face_grad_constraint,
+      right_face_constraint=right_face_constraint,
   )
 
   return psi
 
 
-def _calculate_psi_grad_constraint(
+def _calculate_psi_grad_constraint_from_Ip_tot(
     dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
     geo: geometry.Geometry,
 ) -> jax.Array:
-  """Calculates the constraint on the poloidal flux (psi)."""
+  """Calculates the gradient constraint on the poloidal flux (psi) from Ip."""
   return (
       dynamic_runtime_params_slice.profile_conditions.Ip_tot
       * 1e6
@@ -577,6 +590,19 @@ def _calculate_psi_grad_constraint(
   )
 
 
+def _psi_value_constraint_from_vloop(
+  dt: jax.Array,
+  dynamic_runtime_params_slice_t: runtime_params_slice.DynamicRuntimeParamsSlice,
+  core_profiles_t_minus_dt: state.CoreProfiles,
+  geo: geometry.Geometry,
+) -> jax.Array:
+  """Calculates the value constraint on the poloidal flux (psi) from loop voltage at LCFS."""
+  return (
+    core_profiles_t_minus_dt.psi.face_value()[-1]
+    + dynamic_runtime_params_slice_t.profile_conditions.vloop_lcfs * dt
+  )
+
+
 def _init_psi_and_current(
     static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
     dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
@@ -604,16 +630,37 @@ def _init_psi_and_current(
   Returns:
     Refined core profiles.
   """
-  # Retrieving psi from the profile conditions.
+  use_vloop_bc = (
+        dynamic_runtime_params_slice.profile_conditions.vloop_lcfs is not None
+  )
+
+  # Case 1: retrieving psi from the profile conditions, using the prescribed profile and Ip_tot
   if dynamic_runtime_params_slice.profile_conditions.psi is not None:
+    # Calculate the dpsi/drho necessary to achieve the given Ip_tot
+    dpsi_drhonorm_edge = _calculate_psi_grad_constraint_from_Ip_tot(
+          dynamic_runtime_params_slice,
+          geo,
+    )
+
+    # Set the BCs to ensure the correct Ip_tot
+    if use_vloop_bc:
+      right_face_grad_constraint = None
+      right_face_constraint = (
+          dynamic_runtime_params_slice.profile_conditions.psi[-1]
+          + dpsi_drhonorm_edge * geo.drho_norm / 2
+        )
+    else:
+      # Use the dpsi/drho calculated above as the right face gradient constraint
+      right_face_grad_constraint = dpsi_drhonorm_edge
+      right_face_constraint = None
+
     psi = cell_variable.CellVariable(
-        value=dynamic_runtime_params_slice.profile_conditions.psi,
-        right_face_grad_constraint=_calculate_psi_grad_constraint(
-            dynamic_runtime_params_slice,
-            geo,
-        ),
-        dr=geo.drho_norm,
+      value=dynamic_runtime_params_slice.profile_conditions.psi,
+      right_face_grad_constraint=right_face_grad_constraint,
+      right_face_constraint=right_face_constraint,
+      dr=geo.drho_norm,
     )
+
     core_profiles = dataclasses.replace(core_profiles, psi=psi)
     currents = _calculate_currents_from_psi(
         dynamic_runtime_params_slice=dynamic_runtime_params_slice,
@@ -622,31 +669,53 @@ def _init_psi_and_current(
         core_profiles=core_profiles,
         source_models=source_models,
     )
-  # Retrieving psi from the standard geometry input.
+
+  # Case 2: retrieving psi from the standard geometry input.
   elif (
       isinstance(geo, geometry.StandardGeometry)
       and not dynamic_runtime_params_slice.profile_conditions.initial_psi_from_j
   ):
     # psi is already provided from a numerical equilibrium, so no need to
     # first calculate currents. However, non-inductive currents are still
     # calculated and used in current diffusion equation.
+
+    # Calculate the dpsi/drho necessary to achieve the given Ip_tot
+    dpsi_drhonorm_edge = _calculate_psi_grad_constraint_from_Ip_tot(
+          dynamic_runtime_params_slice,
+          geo,
+    )
+
+    # Set the BCs, with the option of scaling to the given Ip_tot
+    if use_vloop_bc and geo.Ip_from_parameters:
+      right_face_grad_constraint = None
+      right_face_constraint = (
+        geo.psi_from_Ip[-1] + dpsi_drhonorm_edge * geo.drho_norm / 2
+      )
+    elif use_vloop_bc:
+      right_face_grad_constraint = None
+      right_face_constraint = geo.psi_from_Ip[-1]
+    else:
+      right_face_grad_constraint = dpsi_drhonorm_edge
+      right_face_constraint = None
+
     psi = cell_variable.CellVariable(
-        value=geo.psi_from_Ip,
-        right_face_grad_constraint=_calculate_psi_grad_constraint(
-            dynamic_runtime_params_slice,
-            geo,
-        ),
-        dr=geo.drho_norm,
+      value=geo.psi_from_Ip,  # Use psi from equilibrium
+      right_face_grad_constraint=right_face_grad_constraint,
+      right_face_constraint=right_face_constraint,
+      dr=geo.drho_norm,
     )
     core_profiles = dataclasses.replace(core_profiles, psi=psi)
+
+    # Calculate non-inductive currents
     currents = _calculate_currents_from_psi(
         dynamic_runtime_params_slice=dynamic_runtime_params_slice,
         static_runtime_params_slice=static_runtime_params_slice,
         geo=geo,
         core_profiles=core_profiles,
         source_models=source_models,
     )
-  # Calculating j according to nu formula and psi from j.
+
+  # Case 3: calculating j according to nu formula and psi from j.
   elif (
       isinstance(geo, geometry.CircularAnalyticalGeometry)
       or dynamic_runtime_params_slice.profile_conditions.initial_psi_from_j
@@ -686,7 +755,12 @@ def _init_psi_and_current(
   else:
     raise ValueError('Cannot compute psi for given config.')
 
-  core_profiles = dataclasses.replace(core_profiles, psi=psi, currents=currents)
+  core_profiles = dataclasses.replace(
+    core_profiles,
+    psi=psi,
+    currents=currents,
+    vloop_lcfs=dynamic_runtime_params_slice.profile_conditions.vloop_lcfs,
+  )
 
   return core_profiles
 
@@ -743,6 +817,15 @@ def initial_core_profiles(
   s_face = jnp.zeros_like(geo.rho_face)
   currents = state.Currents.zeros(geo)
 
+  # Set vloop_lcfs to 0 for the first time step if not provided
+  vloop_lcfs = (
+      jnp.array(0.0)
+      if dynamic_runtime_params_slice.profile_conditions.vloop_lcfs is None
+      else jnp.array(
+          dynamic_runtime_params_slice.profile_conditions.vloop_lcfs
+      )
+  )
+
   core_profiles = state.CoreProfiles(
       temp_ion=temp_ion,
       temp_el=temp_el,
@@ -761,6 +844,7 @@ def initial_core_profiles(
       q_face=q_face,
       s_face=s_face,
       nref=jnp.asarray(dynamic_runtime_params_slice.numerics.nref),
+      vloop_lcfs=vloop_lcfs,
   )
 
   core_profiles = _init_psi_and_current(
@@ -784,7 +868,11 @@ def initial_core_profiles(
           source_models,
       ),
   )
-  core_profiles = dataclasses.replace(core_profiles, psidot=psidot)
+  core_profiles = dataclasses.replace(
+    core_profiles,
+    psidot=psidot,
+    vloop_lcfs=psidot.face_value()[-1],
+  )
 
   # Set psi as source of truth and recalculate jtot, q, s
   return physics.update_jtot_q_face_s_face(
@@ -923,8 +1011,10 @@ def get_update(x_new, var):
 
 
 def compute_boundary_conditions(
+    dt: jax.Array,
     static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
-    dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
+    dynamic_runtime_params_slice_t: runtime_params_slice.DynamicRuntimeParamsSlice,
+    core_profiles_t_minus_dt: state.CoreProfiles,
     geo: geometry.Geometry,
 ) -> dict[str, dict[str, jax.Array | None]]:
   """Computes boundary conditions for time t and returns updates to State.
@@ -940,38 +1030,38 @@ def compute_boundary_conditions(
     values in a State object.
   """
   Ti_bound_right = jax_utils.error_if_not_positive(
-      dynamic_runtime_params_slice.profile_conditions.Ti_bound_right,
+      dynamic_runtime_params_slice_t.profile_conditions.Ti_bound_right,
       'Ti_bound_right',
   )
 
   Te_bound_right = jax_utils.error_if_not_positive(
-      dynamic_runtime_params_slice.profile_conditions.Te_bound_right,
+      dynamic_runtime_params_slice_t.profile_conditions.Te_bound_right,
       'Te_bound_right',
   )
   # TODO(b/390143606): Separate out the boundary condition calculation from the
   # core profile calculation.
   ne = _get_ne(
-      dynamic_runtime_params_slice.numerics,
-      dynamic_runtime_params_slice.profile_conditions,
+      dynamic_runtime_params_slice_t.numerics,
+      dynamic_runtime_params_slice_t.profile_conditions,
       geo,
   )
   ne_bound_right = ne.right_face_constraint
 
   Zi_edge = charge_states.get_average_charge_state(
       static_runtime_params_slice.main_ion_names,
-      ion_mixture=dynamic_runtime_params_slice.plasma_composition.main_ion,
+      ion_mixture=dynamic_runtime_params_slice_t.plasma_composition.main_ion,
       Te=Te_bound_right,
   )
   Zimp_edge = charge_states.get_average_charge_state(
       static_runtime_params_slice.impurity_names,
-      ion_mixture=dynamic_runtime_params_slice.plasma_composition.impurity,
+      ion_mixture=dynamic_runtime_params_slice_t.plasma_composition.impurity,
       Te=Te_bound_right,
   )
 
   dilution_factor_edge = physics.get_main_ion_dilution_factor(
       Zi_edge,
       Zimp_edge,
-      dynamic_runtime_params_slice.plasma_composition.Zeff_face[-1],
+      dynamic_runtime_params_slice_t.plasma_composition.Zeff_face[-1],
   )
 
   ni_bound_right = ne_bound_right * dilution_factor_edge
@@ -1004,9 +1094,24 @@ def compute_boundary_conditions(
           right_face_constraint=jnp.array(nimp_bound_right),
       ),
       'psi': dict(
-          right_face_grad_constraint=_calculate_psi_grad_constraint(
-              dynamic_runtime_params_slice,
+          right_face_grad_constraint=(
+            _calculate_psi_grad_constraint_from_Ip_tot(
+              dynamic_runtime_params_slice_t,
+              geo,
+            )
+            if dynamic_runtime_params_slice_t.profile_conditions.vloop_lcfs
+            is None
+            else None
+          ),
+          right_face_constraint=(
+            _psi_value_constraint_from_vloop(
+              dt,
+              dynamic_runtime_params_slice_t,
+              core_profiles_t_minus_dt,
               geo,
+            )
+            if dynamic_runtime_params_slice_t.profile_conditions.vloop_lcfs is not None
+            else None
           ),
       ),
       'Zi_edge': Zi_edge,