diff --git a/desc/objectives/nae_utils.py b/desc/objectives/nae_utils.py index f4e9973ce0..68d5adf9d8 100644 --- a/desc/objectives/nae_utils.py +++ b/desc/objectives/nae_utils.py @@ -9,7 +9,7 @@ from .linear_objectives import FixSumModesR, FixSumModesZ -def _calc_1st_order_NAE_coeffs(qsc, desc_eq): +def _calc_1st_order_NAE_coeffs(qsc, desc_eq, N=None): """Calculate 1st order NAE coefficients' toroidal Fourier representations. Uses the passed-in qsc object, and the desc_eq's stellarator symmetry is used. @@ -20,6 +20,9 @@ def _calc_1st_order_NAE_coeffs(qsc, desc_eq): Qsc object to use as the NAE constraints on the DESC equilibrium. desc_eq : Equilibrium desc equilibrium to constrain. + N : int, + max toroidal resolution to constrain. + If None, defaults to equilibrium's toroidal resolution Returns ------- @@ -43,6 +46,12 @@ def _calc_1st_order_NAE_coeffs(qsc, desc_eq): R0 = qsc.R0_func(phi) dR0_dphi = qsc.R0p dZ0_dphi = qsc.Z0p + if N is None: + N = desc_eq.N + else: + N = np.min([desc_eq.N, N]) + assert N == int(N), "Toroidal Resolution must be an integer!" + N = int(N) # normal and binormal vector components # Spline interpolants for the cylindrical components of the Frenet-Serret frame: # these are functions of phi (toroidal cylindrical angle) @@ -76,15 +85,15 @@ def _calc_1st_order_NAE_coeffs(qsc, desc_eq): nfp = qsc.nfp if desc_eq.sym: - Rbasis = FourierSeries(N=desc_eq.N, NFP=nfp, sym="cos") - Zbasis = FourierSeries(N=desc_eq.N, NFP=nfp, sym="cos") - Rbasis_sin = FourierSeries(N=desc_eq.N, NFP=nfp, sym="sin") - Zbasis_sin = FourierSeries(N=desc_eq.N, NFP=nfp, sym="sin") + Rbasis = FourierSeries(N=N, NFP=nfp, sym="cos") + Zbasis = FourierSeries(N=N, NFP=nfp, sym="cos") + Rbasis_sin = FourierSeries(N=N, NFP=nfp, sym="sin") + Zbasis_sin = FourierSeries(N=N, NFP=nfp, sym="sin") else: - Rbasis = FourierSeries(N=desc_eq.N, NFP=nfp, sym=False) - Zbasis = FourierSeries(N=desc_eq.N, NFP=nfp, sym=False) - Rbasis_sin = FourierSeries(N=desc_eq.N, NFP=nfp, sym=False) - Zbasis_sin = FourierSeries(N=desc_eq.N, NFP=nfp, sym=False) + Rbasis = FourierSeries(N=N, NFP=nfp, sym=False) + Zbasis = FourierSeries(N=N, NFP=nfp, sym=False) + Rbasis_sin = FourierSeries(N=N, NFP=nfp, sym=False) + Zbasis_sin = FourierSeries(N=N, NFP=nfp, sym=False) grid = LinearGrid(M=0, L=0, zeta=phi, NFP=nfp) Rtrans = Transform(grid, Rbasis, build_pinv=True, method="auto") @@ -217,7 +226,7 @@ def _make_RZ_cons_order_rho(qsc, desc_eq, coeffs, bases): return Rconstraints, Zconstraints -def make_RZ_cons_1st_order(qsc, desc_eq): +def make_RZ_cons_1st_order(qsc, desc_eq, N=None): """Make the first order NAE constraints for a DESC equilibrium. Parameters @@ -226,6 +235,9 @@ def make_RZ_cons_1st_order(qsc, desc_eq): Qsc object to use as the NAE constraints on the DESC equilibrium. desc_eq : Equilibrium desc equilibrium to constrain. + N : int, + max toroidal resolution to constrain. + If None, defaults to equilibrium's toroidal resolution Returns ------- @@ -239,7 +251,7 @@ def make_RZ_cons_1st_order(qsc, desc_eq): Rconstraints = () Zconstraints = () - coeffs, bases = _calc_1st_order_NAE_coeffs(qsc, desc_eq) + coeffs, bases = _calc_1st_order_NAE_coeffs(qsc, desc_eq, N=N) Rconstraints, Zconstraints = _make_RZ_cons_order_rho(qsc, desc_eq, coeffs, bases) return Rconstraints + Zconstraints diff --git a/desc/objectives/utils.py b/desc/objectives/utils.py index acd889b82f..a521703758 100644 --- a/desc/objectives/utils.py +++ b/desc/objectives/utils.py @@ -174,7 +174,14 @@ def get_fixed_axis_constraints( def get_NAE_constraints( - desc_eq, qsc_eq, order=1, profiles=True, iota=False, kinetic=False, normalize=True + desc_eq, + qsc_eq, + order=1, + profiles=True, + iota=False, + kinetic=False, + normalize=True, + N=None, ): """Get the constraints necessary for fixing NAE behavior in an equilibrium problem. # noqa D205 @@ -195,6 +202,9 @@ def get_NAE_constraints( Whether to also fix kinetic profiles. normalize : bool Whether to apply constraints in normalized units. + N : int, + max toroidal resolution to constrain. + If None, defaults to equilibrium's toroidal resolution Returns ------- @@ -238,7 +248,7 @@ def get_NAE_constraints( FixCurrent(eq=desc_eq, normalize=normalize, normalize_target=normalize), ) if order >= 1: # first order constraints - constraints += make_RZ_cons_1st_order(qsc=qsc_eq, desc_eq=desc_eq) + constraints += make_RZ_cons_1st_order(qsc=qsc_eq, desc_eq=desc_eq, N=N) if order >= 2: # 2nd order constraints raise NotImplementedError("NAE constraints only implemented up to O(rho) ") diff --git a/tests/test_examples.py b/tests/test_examples.py index ee06928edf..07251bcfb6 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -542,7 +542,7 @@ def test_NAE_QSC_solve(): # this has all the constraints we need, # iota=False specifies we want to fix current instead of iota - cs = get_NAE_constraints(eq, qsc, iota=False, order=1) + cs = get_NAE_constraints(eq, qsc, iota=False, order=1, N=eq.N) objectives = ForceBalance(eq=eq) obj = ObjectiveFunction(objectives)