Merge branch 'master' into change_resolution_grid_attribute

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

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

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)