fix bug in QSBoozer

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,

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

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

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):