Resolve behavior with NAE constraints and lambda when equilibrium is asymmetric

[dpanici]

When NAE constraints are used with an asymmetric NAE solution, the NAE lambda may not have a 0 average (which is enforced by LambdaGauge in DESC for asymmetric equilibria). This causes a clash between the FixNearAxisLambda constraint and the LambdaGauge constraint in DESC.

We can remove the average of the NAE lambda, but then we would be effectively changing the origin of the NAE $\theta_{Boozer}$. We should change the coefficients we calculate, which are based off of terms like $R_{1,1}\rho\cos(\theta_{Boozer})$, since we've modified the angle.

[dpanici]

The below test works if LambdaGauge is disabled (on #1036 ). The equilibrium is pretty high aspect ratio though, so I should test this with a higher aspect ratio and see if it still works, and also with a more asymmetric equilirbium

"""Test near-axis 1st/2nd order constraints for asym qsc equilibrium."""
qsc = Qsc.from_paper("precise QA", rs=[1e-3, 1e-2])
ntheta = 75
r = 0.075
N = 9
eq_fit = Equilibrium.from_near_axis(qsc, r=r, L=6, M=8, N=N, ntheta=ntheta)
eq = eq_fit.copy()
eq_2nd_order = eq_fit.copy()
orig_Rax_val = eq.axis.R_n
orig_Zax_val = eq.axis.Z_n

cs = get_NAE_constraints(eq, qsc_eq=qsc, order=1, fix_lambda=0, N=eq.N)
cs_2nd_order = get_NAE_constraints(
    eq_2nd_order, qsc_eq=qsc, order=2, fix_lambda=0, N=eq.N
)
# with pytest.raises(NotImplementedError):
#     get_NAE_constraints(eq, qsc_eq=qsc, order=1, fix_lambda=0, N=eq.N)
# for c in cs:
#     # should be no FixSumModeslambda in the fix_lambda=False constraint
#     assert not isinstance(c, FixSumModesLambda)

for eqq, constraints in zip(
    [eq, eq_2nd_order],
    [cs, cs_2nd_order],
):
    objectives = ForceBalance(eq=eqq)
    obj = ObjectiveFunction(objectives)

    eqq.solve(
        verbose=3,
        ftol=1e-3,
        objective=obj,
        maxiter=100,
        xtol=1e-6,
        gtol=1e-8,
        constraints=constraints,
    )
grid_axis = LinearGrid(rho=0.0, theta=0.0, N=eq.N_grid, NFP=eq.NFP)
# Make sure axis is same
for eqq, string in zip(
    [eq, eq_2nd_order],
    ["RZ O(rho)", "RZ O(rho^2)"],
):
    np.testing.assert_array_almost_equal(orig_Rax_val, eqq.axis.R_n, err_msg=string)
    np.testing.assert_array_almost_equal(orig_Zax_val, eqq.axis.Z_n, err_msg=string)

    # Make sure iota of solved equilibrium is same on axis as QSC

    iota = eqq.compute("iota", grid=LinearGrid(rho=0.0))["iota"]

    np.testing.assert_allclose(iota[0], qsc.iota, rtol=2e-4, err_msg=string)

    ### check lambda to match on axis
    # Evaluate lambda on the axis
    data_nae = eqq.compute(["lambda", "|B|"], grid=grid_axis)
    lam_nae = data_nae["lambda"]

    phi = np.squeeze(grid_axis.nodes[:, 2])
    np.testing.assert_allclose(
        lam_nae,
        -qsc.iota * qsc.nu_spline(phi),
        atol=4e-3,
        rtol=1e-3,
        err_msg=string,
    )

    # check |B| on axis
    np.testing.assert_allclose(
        data_nae["|B|"], np.ones(np.size(phi)) * qsc.B0, atol=3e-4, err_msg=string
    )

[dpanici]

Have LambdaGauge fix to whatever the current average is for each surface