Merge pull request #626 from PlasmaControl/dp/fix-QS-B

fix bug in QSBoozer
PlasmaControl · Aug 17, 2023 · 6b366f5 · 6b366f5
2 parents 7cbd4eb + 9fa9d69
commit 6b366f5
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 3 deletions.
diff --git a/desc/objectives/_qs.py b/desc/objectives/_qs.py
@@ -143,6 +143,7 @@ def build(self, eq=None, use_jit=True, verbose=1):
             helicity=self.helicity,
             NFP=self._transforms["B"].basis.NFP,
         )
+
         self._constants = {
             "transforms": self._transforms,
             "profiles": self._profiles,

diff --git a/desc/vmec_utils.py b/desc/vmec_utils.py
@@ -267,7 +267,7 @@ def ptolemy_linear_transform(desc_modes, vmec_modes=None, helicity=None, NFP=Non
         else:
             idx_MN = np.nonzero(vmec_modes[:, 1] * N == vmec_modes[:, 2] * M)[0]
         idx[idx_MN] = False
-
+        idx = np.nonzero(idx)[0]
         return matrix, vmec_modes, idx
 
     return matrix, vmec_modes

diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py
@@ -197,6 +197,20 @@ def test(eq):
         test(Equilibrium(L=2, M=2, N=1, iota=PowerSeriesProfile(0)))
         test(Equilibrium(L=2, M=2, N=1, current=PowerSeriesProfile(0)))
 
+    @pytest.mark.unit
+    def test_jax_compile_boozer(self):
+        """Test compilation of Boozer QA metric in ObjectiveFunction."""
+        # making sure that compiles without any errors from JAX
+        # Related to issue #625
+        def test(eq):
+            obj = ObjectiveFunction(QuasisymmetryBoozer(eq=eq))
+            obj.build()
+            obj.compile()
+            fb = obj.compute_unscaled(obj.x(eq))
+            np.testing.assert_allclose(fb, 0, atol=1e-12)
+
+        test(Equilibrium(L=2, M=2, N=1, current=PowerSeriesProfile(0)))
+
     @pytest.mark.unit
     def test_qh_boozer(self):
         """Test calculation of Boozer QH metric."""
@@ -231,7 +245,7 @@ def test_qh_boozer(self):
         idx_B = np.argsort(np.abs(B_mn))
 
         # check that largest amplitudes are the QH modes
-        np.testing.assert_allclose(B_mn[idx_B[-3:]], np.flip(B_mn[~idx][:3]))
+        np.testing.assert_allclose(B_mn[idx_B[-3:]], np.flip(np.delete(B_mn, idx)[:3]))
         # check that these QH modes are not returned by the objective
         assert [b not in f for b in B_mn[idx_B[-3:]]]
         # check that the objective returns the lowest amplitudes

diff --git a/tests/test_vmec.py b/tests/test_vmec.py
@@ -182,7 +182,7 @@ def test_ptolemy_linear_transform(self):
                 [1, 3, 3],
             ]
         )
-        np.testing.assert_allclose(modes[~idx], sym_modes)
+        np.testing.assert_allclose(np.delete(modes, idx, axis=0), sym_modes)
 
     @pytest.mark.unit
     def test_fourier_to_zernike(self):