Compute jxb_gradp (and related quantities) even if VMEC did not conve…

…rge.
proximafusion · Mar 5, 2025 · b5409d4 · b5409d4
1 parent 6e5e281
commit b5409d4
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 3 deletions.
diff --git a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.cc b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.cc
@@ -1251,7 +1251,8 @@ vmecpp::OutputQuantities vmecpp::ComputeOutputQuantities(
       models_from_threads, radial_profiles);
 
   if (vmec_status == VmecStatus::NORMAL_TERMINATION ||
-      vmec_status == VmecStatus::SUCCESSFUL_TERMINATION) {
+      vmec_status == VmecStatus::SUCCESSFUL_TERMINATION ||
+      indata.return_outputs_even_if_not_converged) {
     MeshBledingBSubZeta(
         s, fc,
         /*m_vmec_internal_results=*/output_quantities.vmec_internal_results);
@@ -1314,6 +1315,7 @@ vmecpp::OutputQuantities vmecpp::ComputeOutputQuantities(
     output_quantities.jxbout = ComputeJxBOutputFileContents(
         s, fc, output_quantities.vmec_internal_results,
         output_quantities.bsubs_full, output_quantities.covariant_b_derivatives,
+        indata.return_outputs_even_if_not_converged,
         vmec_status);
 
     if (checkpoint == VmecCheckpoint::JXBOUT) {
@@ -2360,6 +2362,7 @@ vmecpp::JxBOutFileContents vmecpp::ComputeJxBOutputFileContents(
     const VmecInternalResults& vmec_internal_results,
     const BSubSFull& bsubs_full,
     const CovariantBDerivatives& covariant_b_derivatives,
+    const bool return_outputs_even_if_not_converged,
     VmecStatus vmec_status) {
   JxBOutFileContents jxbout;
 
@@ -2598,7 +2601,7 @@ vmecpp::JxBOutFileContents vmecpp::ComputeJxBOutputFileContents(
     jxbout.jperp2[jF] = dnorm1 * tjnorm * average_jperp2;
 
     // Some quantities are only computed if VMEC++ actually converged.
-    if (vmec_status == VmecStatus::SUCCESSFUL_TERMINATION) {
+    if (vmec_status == VmecStatus::SUCCESSFUL_TERMINATION || return_outputs_even_if_not_converged) {
       // normalized toroidal magnetic flux
       jxbout.phin[jF] = vmec_internal_results.phiF[jF] /
                         vmec_internal_results.phiF[fc.ns - 1];
@@ -2635,7 +2638,7 @@ vmecpp::JxBOutFileContents vmecpp::ComputeJxBOutputFileContents(
   }  // jF
 
   // Some quantities are only computed if VMEC++ actually converged.
-  if (vmec_status == VmecStatus::SUCCESSFUL_TERMINATION) {
+  if (vmec_status == VmecStatus::SUCCESSFUL_TERMINATION || return_outputs_even_if_not_converged) {
     // The loop in jxbforce.f90:594 goes over js=2,ns1,
     // which means that the last half-grid point is not touched.
     for (int jH = 0; jH < vmec_internal_results.num_half - 1; ++jH) {

diff --git a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.h b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.h
@@ -1331,6 +1331,7 @@ JxBOutFileContents ComputeJxBOutputFileContents(
     const VmecInternalResults& vmec_internal_results,
     const BSubSFull& bsubs_full,
     const CovariantBDerivatives& covariant_b_derivatives,
+    const bool return_outputs_even_if_not_converged,
     VmecStatus vmec_status);
 
 MercierStabilityIntermediateQuantities ComputeIntermediateMercierQuantities(

diff --git a/tests/test_init.py b/tests/test_init.py
@@ -55,6 +55,11 @@ def test_get_outputs_if_non_converged_if_wanted():
     assert vmec_output.wout is not None
     assert vmec_output.wout.niter == 2
 
+    # actually check that some arrays,
+    # which were previously only filled if VMEC converged,
+    # also get populated now
+    assert not np.all(vmec_output.jxbout.jxb_gradp == 0.0)
+
 
 # We trust the C++ tests to cover the hot restart functionality properly,
 # here we just want to test that the Python API for it works.