Include vloop_lcfs updates in updates to core_profiles

torax/core_profile_setters.py

@@ -772,9 +772,9 @@ def initial_core_profiles(
   s_face = jnp.zeros_like(geo.rho_face)
   currents = state.Currents.zeros(geo)
 
-  # Set vloop_lcfs to nan for the first time step if not provided
+  # Set vloop_lcfs to 0 for the first time step if not provided
   vloop_lcfs = (
-      jnp.array(jnp.nan)
+      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
@@ -821,7 +821,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
   core_profiles = physics.update_jtot_q_face_s_face(

torax/sim.py

@@ -586,7 +586,7 @@ def finalize_output(
 
     # Update psidot based on the new core profiles.
     # Will include the phibdot calculation since geo=geo_t_plus_dt.
-    output_state.core_profiles = _update_psidot(
+    output_state.core_profiles = _update_psidot_and_vloop_lcfs(
         source_models=self._stepper_fn.source_models,
         dynamic_runtime_params_slice=dynamic_runtime_params_slice_t_plus_dt,
         static_runtime_params_slice=static_runtime_params_slice,
@@ -1380,7 +1380,7 @@ def update_current_distribution(
   return new_core_profiles
 
 
-def _update_psidot(
+def _update_psidot_and_vloop_lcfs(
     static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
     dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
     geo: geometry.Geometry,
@@ -1403,6 +1403,7 @@ def _update_psidot(
   new_core_profiles = dataclasses.replace(
       core_profiles,
       psidot=psidot,
+      vloop_lcfs=psidot.face_value()[-1],
   )
   return new_core_profiles