FixOmniBmax bug fix (#1268)

Resolves #1266

desc/objectives/linear_objectives.py

@@ -3065,7 +3065,7 @@ def build(self, use_jit=True, verbose=1):
                     np.logical_and((basis.modes[:, 1] % 2 == 0), basis.modes[:, 1] > 0),
                 )
             )[0]
-            mm = basis.modes[idx_m, 2]
+            mm = basis.modes[idx_m, 1]
             self._A[i, idx_0] = 1
             self._A[i, idx_m] = (mm % 2 - 1) * (mm % 4 - 1)
 

tests/test_linear_objectives.py

@@ -36,6 +36,7 @@
     FixNearAxisLambda,
     FixNearAxisR,
     FixNearAxisZ,
+    FixOmniBmax,
     FixOmniMap,
     FixOmniWell,
     FixParameters,
@@ -915,6 +916,26 @@ def test_fix_omni_indices():
     assert constraint.dim_f == indices.size
 
 
+@pytest.mark.unit
+def test_fix_omni_Bmax():
+    """Test that omnigenity parameters are constrained for B_max to be a straight line.
+
+    Test for GH issue #1266.
+    """
+
+    def _test(M_x, N_x, NFP, sum):
+        field = OmnigenousField(L_x=2, M_x=M_x, N_x=N_x, NFP=NFP, helicity=(1, NFP))
+        constraint = FixOmniBmax(field=field)
+        constraint.build()
+        assert constraint.dim_f == (2 * field.N_x + 1) * (field.L_x + 1)
+        # 0 - 2 + 4 - 6 + 8 ...
+        np.testing.assert_allclose(constraint._A @ field.x_basis.modes[:, 1], sum)
+
+    _test(M_x=6, N_x=3, NFP=1, sum=-4)
+    _test(M_x=9, N_x=4, NFP=2, sum=4)
+    _test(M_x=12, N_x=5, NFP=3, sum=6)
+
+
 @pytest.mark.unit
 def test_fix_parameters_input_order(DummyStellarator):
     """Test that FixParameters preserves the input indices and target ordering."""