From 0dca8a43c37d2fa6f204cc1e93eee7254c1999ee Mon Sep 17 00:00:00 2001
From: dpanici <dpanici@princeton.edu>
Date: Thu, 16 Apr 2026 15:57:42 -0400
Subject: [PATCH] add parameter to shift the max B theta to Omnigenity
 objective

---
 desc/objectives/_omnigenity.py | 20 +++++++++++---
 tests/test_objective_funs.py   | 49 ++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 3 deletions(-)

diff --git a/desc/objectives/_omnigenity.py b/desc/objectives/_omnigenity.py
index b497d46cf4..db8fd5bd8b 100644
--- a/desc/objectives/_omnigenity.py
+++ b/desc/objectives/_omnigenity.py
@@ -503,8 +503,13 @@ class Omnigenity(_Objective):
     Errors are relative to a target field that is perfectly omnigenous,
     and are computed on a collocation grid in (ρ,η,α) coordinates.
 
-    This objective assumes that the collocation point (θ=0,ζ=0) lies on the contour of
-    maximum field strength ||B||=B_max.
+    By default this objective assumes that the collocation point (θ=0,ζ=0) of the
+    equilibrium lies on the contour of maximum field strength ||B||=B_max. The
+    ``B_max_theta_location`` argument relaxes this assumption: the field's (θ=0) point
+    (which corresponds to the field's B_max) is compared against the equilibrium's |B|
+    evaluated at θ = ``B_max_theta_location``.
+
+    A single well in |B| per field period is assumed for this objective.
 
     Parameters
     ----------
@@ -541,6 +546,11 @@ class Omnigenity(_Objective):
         computation time during optimization and only ``eq`` is allowed to change.
         If False, the field is allowed to change during the optimization and its
         associated data are re-computed at every iteration (Default).
+    B_max_theta_location : float, optional
+        Boozer poloidal angle θ (in radians) at which the equilibrium's |B| attains its
+        maximum. The field's theta_B (which assumes the maximum is at θ=0) is shifted
+        by this value when evaluating the equilibrium's |B| for the comparison.
+        Default = 0.0.
 
     """
 
@@ -580,6 +590,7 @@ def __init__(
         field_fixed=False,
         name="omnigenity",
         jac_chunk_size=None,
+        B_max_theta_location=0.0,
     ):
         if target is None and bounds is None:
             target = 0
@@ -593,6 +604,7 @@ def __init__(
         self.eta_weight = eta_weight
         self._eq_fixed = eq_fixed
         self._field_fixed = field_fixed
+        self._B_max_theta_location = B_max_theta_location
         if not eq_fixed and not field_fixed:
             things = [eq, field]
         elif eq_fixed and not field_fixed:
@@ -824,11 +836,13 @@ def compute(self, params_1=None, params_2=None, constants=None):
 
         # additional computations that cannot be part of the regular compute API
 
+        B_max_theta_location = self._B_max_theta_location
+
         def _compute_B_eta_alpha(theta_B, zeta_B, B_mn):
             nodes = jnp.vstack(
                 (
                     jnp.zeros_like(theta_B),
-                    theta_B,
+                    theta_B + B_max_theta_location,
                     zeta_B,
                 )
             ).T
diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py
index 58078349bb..ab8441feb0 100644
--- a/tests/test_objective_funs.py
+++ b/tests/test_objective_funs.py
@@ -25,6 +25,7 @@
     MixedCoilSet,
     initialize_modular_coils,
 )
+from desc.compat import flip_theta
 from desc.compute import get_transforms
 from desc.equilibrium import Equilibrium
 from desc.examples import get
@@ -1942,6 +1943,54 @@ def test_linking_current(self):
         f = obj.compute(coilset4.params_dict, eq.params_dict)
         np.testing.assert_allclose(f, -0.5 * G / 8, rtol=1e-7)
 
+    @pytest.mark.unit
+    def test_omnigenity_flipped_theta(self):
+        """Test omnigenity transform when Bmax is not at theta=0."""
+        surf = FourierRZToroidalSurface.from_qp_model(
+            major_radius=1,
+            aspect_ratio=20,
+            elongation=6,
+            mirror_ratio=0.2,
+            torsion=0.1,
+            NFP=1,
+            sym=True,
+        )
+        eq = Equilibrium(
+            Psi=6e-3,
+            M=4,
+            N=4,
+            surface=surf,
+            iota=PowerSeriesProfile(1, 0, -1),  # ensure diff surfs have diff iota
+        )
+        field = OmnigenousField(
+            L_B=1,
+            M_B=3,
+            L_x=1,
+            M_x=1,
+            N_x=1,
+            NFP=eq.NFP,
+            helicity=(1, 1),
+            B_lm=np.array(
+                [
+                    [0.8, 1.0, 1.2],
+                    [-0.4, 0.0, 0.6],  # radially varying B
+                ]
+            ).flatten(),
+        )
+        grid1 = LinearGrid(rho=0.5, M=eq.M_grid, N=eq.N_grid)
+        obj1 = Omnigenity(eq=eq, field=field, eq_grid=grid1)
+        obj1.build()
+
+        f1 = obj1.compute(*obj1.xs(eq, field))
+        # now flip where Bmaxn is
+        eq = flip_theta(eq)
+        obj2 = Omnigenity(eq=eq, field=field, eq_grid=grid1, B_max_theta_location=np.pi)
+        obj2.build()
+        f2 = obj2.compute(*obj2.xs(eq, field))
+
+        # values should be exact the same because all we did was re-define theta
+        np.testing.assert_allclose(f1, f2, atol=1e-14)
+
     @pytest.mark.unit
     def test_omnigenity_multiple_surfaces(self):
         """Test omnigenity transform vectorized over multiple surfaces."""