Use Vloop_bound_right in compute_boundary_conditions

google-deepmind · Nov 26, 2024 · 93fd0aa · 93fd0aa
1 parent 90098a5
commit 93fd0aa
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/torax/core_profile_setters.py b/torax/core_profile_setters.py
@@ -564,6 +564,7 @@ def _calculate_psi_grad_constraint_from_Ip_tot(
       / (geo.g2g3_over_rhon_face[-1] * geo.F_face[-1])
   )
 
+
 def _psi_value_constraint_from_Vloop(
   dt: jax.Array,
   dynamic_runtime_params_slice_t: runtime_params_slice.DynamicRuntimeParamsSlice,
@@ -576,6 +577,7 @@ def _psi_value_constraint_from_Vloop(
     + dynamic_runtime_params_slice_t.profile_conditions.Vloop_bound_right * dt
   )
 
+
 def _init_psi_and_current(
     dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
     geo: Geometry,
@@ -965,10 +967,15 @@ def compute_boundary_conditions(
           right_face_constraint=jnp.array(nimp_bound_right),
       ),
       'psi': dict(
-          right_face_grad_constraint=_calculate_psi_grad_constraint_from_Ip_tot(
+          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_bound_right
+            is None
+            else None
+          ),
           right_face_constraint=(
             _psi_value_constraint_from_Vloop(
               dynamic_runtime_params_slice_t,
@@ -980,7 +987,7 @@ def compute_boundary_conditions(
             else None
           ),
         ),
-  }
+    }
 
 
 # pylint: disable=invalid-name