Merge branch 'master' into dp/order_rho_sq_NAE

CHANGELOG.md

@@ -49,7 +49,8 @@ Bug Fixes
 - Fixes the coil currents in ``desc.coils.initialize_modular_coils`` to now give the correct expected linking current.
 - ``desc.objectives.PlasmaVesselDistance`` now correctly accounts for multiple field periods on both the equilibrium and the vessel surface. Previously it only considered distances within a single field period.
 - Sets ``os.environ["JAX_PLATFORMS"] = "cpu"`` instead of ``os.environ["JAX_PLATFORM_NAME"] = "cpu"`` when doing ``set_device("cpu")``.
-
+- Fixes bug in ``desc.input_reader.desc_output_to_input`` utility function for asymmetric equilibria profiles, where the full profile resolution was not being saved.
+- Fixes bug when passing only `sym` into `.change_resolution` for ``FourierRZToroidalSurface``, ``ZernikeRZToroidalSection`` and ``FourierRZCurve``.
 
 Performance Improvements
 

desc/basis.py

@@ -539,7 +539,7 @@ def change_resolution(self, N, NFP=None, sym=None):
         """
         NFP = check_posint(NFP, "NFP")
         self._NFP = NFP if NFP is not None else self.NFP
-        if N != self.N:
+        if N != self.N or (sym is not None and sym != self.sym):
             self._N = check_nonnegint(N, "N", False)
             self._sym = sym if sym is not None else self.sym
             self._modes = self._get_modes(self.N)
@@ -690,7 +690,7 @@ def change_resolution(self, M, N, NFP=None, sym=None):
         """
         NFP = check_posint(NFP, "NFP")
         self._NFP = NFP if NFP is not None else self.NFP
-        if M != self.M or N != self.N or sym != self.sym:
+        if M != self.M or N != self.N or (sym is not None and sym != self.sym):
             self._M = check_nonnegint(M, "M", False)
             self._N = check_nonnegint(N, "N", False)
             self._sym = sym if sym is not None else self.sym
@@ -898,7 +898,7 @@ def change_resolution(self, L, M, sym=None):
         None
 
         """
-        if L != self.L or M != self.M or sym != self.sym:
+        if L != self.L or M != self.M or (sym is not None and sym != self.sym):
             self._L = check_nonnegint(L, "L", False)
             self._M = check_nonnegint(M, "M", False)
             self._sym = sym if sym is not None else self.sym
@@ -1073,7 +1073,12 @@ def change_resolution(self, L, M, N, NFP=None, sym=None):
         """
         NFP = check_posint(NFP, "NFP")
         self._NFP = NFP if NFP is not None else self.NFP
-        if L != self.L or M != self.M or N != self.N or sym != self.sym:
+        if (
+            L != self.L
+            or M != self.M
+            or N != self.N
+            or (sym is not None and sym != self.sym)
+        ):
             self._L = check_nonnegint(L, "L", False)
             self._M = check_nonnegint(M, "M", False)
             self._N = check_nonnegint(N, "N", False)
@@ -1311,7 +1316,12 @@ def change_resolution(self, L, M, N, NFP=None, sym=None):
         """
         NFP = check_posint(NFP, "NFP")
         self._NFP = NFP if NFP is not None else self.NFP
-        if L != self.L or M != self.M or N != self.N or sym != self.sym:
+        if (
+            L != self.L
+            or M != self.M
+            or N != self.N
+            or (sym is not None and sym != self.sym)
+        ):
             self._L = check_nonnegint(L, "L", False)
             self._M = check_nonnegint(M, "M", False)
             self._N = check_nonnegint(N, "N", False)

desc/equilibrium/equilibrium.py

@@ -264,8 +264,8 @@ def __init__(
         self._M_grid = setdefault(M_grid, 2 * self.M)
         self._N_grid = setdefault(N_grid, 2 * self.N)
 
-        self._surface.change_resolution(self.L, self.M, self.N)
-        self._axis.change_resolution(self.N)
+        self._surface.change_resolution(self.L, self.M, self.N, sym=self.sym)
+        self._axis.change_resolution(self.N, sym=self.sym)
 
         # bases
         self._R_basis = FourierZernikeBasis(

desc/geometry/curve.py

@@ -127,8 +127,7 @@ def change_resolution(self, N=None, NFP=None, sym=None):
         if (
             ((N is not None) and (N != self.N))
             or ((NFP is not None) and (NFP != self.NFP))
-            or (sym is not None)
-            and (sym != self.sym)
+            or ((sym is not None) and (sym != self.sym))
         ):
             self._NFP = int(NFP if NFP is not None else self.NFP)
             self._sym = bool(sym) if sym is not None else self.sym

desc/geometry/surface.py

@@ -205,13 +205,14 @@ def change_resolution(self, *args, **kwargs):
         N = check_nonnegint(N, "N")
         NFP = check_posint(NFP, "NFP")
         self._NFP = int(NFP if NFP is not None else self.NFP)
-        self._sym = sym if sym is not None else self.sym
 
         if (
             ((N is not None) and (N != self.N))
             or ((M is not None) and (M != self.M))
             or (NFP is not None)
+            or ((sym is not None) and (sym != self.sym))
         ):
+            self._sym = sym if sym is not None else self.sym
             M = int(M if M is not None else self.M)
             N = int(N if N is not None else self.N)
             R_modes_old = self.R_basis.modes
@@ -975,9 +976,13 @@ def change_resolution(self, *args, **kwargs):
 
         L = check_nonnegint(L, "L")
         M = check_nonnegint(M, "M")
-        self._sym = sym if sym is not None else self.sym
 
-        if ((L is not None) and (L != self.L)) or ((M is not None) and (M != self.M)):
+        if (
+            ((L is not None) and (L != self.L))
+            or ((M is not None) and (M != self.M))
+            or ((sym is not None) and (sym != self.sym))
+        ):
+            self._sym = sym if sym is not None else self.sym
             L = int(L if L is not None else self.L)
             M = int(M if M is not None else self.M)
             R_modes_old = self.R_basis.modes

desc/input_reader.py

@@ -906,12 +906,12 @@ def desc_output_to_input(  # noqa: C901 - fxn too complex
         pres_profile.change_resolution(L=L_profile)
         profile.change_resolution(L=L_profile)
 
-        prof_modes = np.zeros((L_profile, 3))
-        prof_modes[:, 0] = np.arange(L_profile)
+        prof_modes = np.zeros((L_profile + 1, 3))
+        prof_modes[:, 0] = np.arange(L_profile + 1)
         p1 = copy_coeffs(pres_profile.params, pres_profile.basis.modes, prof_modes)
         p2 = copy_coeffs(profile.params, profile.basis.modes, prof_modes)
         f.write("\n# pressure and rotational transform/current profiles\n")
-        for l in range(L_profile):
+        for l in range(L_profile + 1):
             f.write(
                 "l: {:3d}  p = {:15.8E}  {} = {:15.8E}\n".format(
                     int(l), p1[l], char, p2[l]

tests/inputs/iotest_HELIOTRON

@@ -1,45 +1,47 @@
-# This is the DESC input file for a fixed iota heliotron
+# DESC input file generated from the output file:
+# tests/inputs/iotest_HELIOTRON.h5
 
 # global parameters
 sym = 1
 NFP = 19
-Psi = 1.0
+Psi = 1.00000000E+00
 
 # spectral resolution
-L_rad  = 9
-M_pol  =  14
-N_tor  =   3
-
-# continuation parameters
-bdry_ratio = 1
-pres_ratio = 1
-pert_order = 2
+L_rad = 9
+M_pol = 14
+N_tor = 3
+L_grid = 28
+M_grid = 28
+N_grid = 6
 
 # solver tolerances
-ftol = 1e-3
-xtol = 1e-6
-gtol = 1e-6
-maxiter = 50
+ftol = 1.000E-02
+xtol = 1.000E-06
+gtol = 1.000E-06
+maxiter = 100
 
 # solver methods
-optimizer         = lsq-exact
-objective         = force
+optimizer = lsq-exact
+objective = force
 spectral_indexing = ansi
 
 # pressure and rotational transform/current profiles
-l:   0	p =   1E2    i = 0.5
-l:   1	p =   0      i = 0.1
-l:   2	p =   1E2    i = 0
-l:   4	p =   0      i = 0
-
-# magnetic axis initial guess
-n:   0  R0 =  1.0E+1  Z0 =  0.0E+0
+l:   0  p =  1.00000000E+02  i =  5.00000000E-01
+l:   1  p =  1.00000000E-05  i =  1.00000000E-01
+l:   2  p =  1.00000000E+02  i =  0.00000000E+00
+l:   3  p =  0.00000000E+00  i =  0.00000000E+00
+l:   4  p =  0.00000000E+00  i =  0.00000000E+00
+l:   5  p =  0.00000000E+00  i =  0.00000000E+00
+l:   6  p =  0.00000000E+00  i =  0.00000000E+00
+l:   7  p =  0.00000000E+00  i =  0.00000000E+00
+l:   8  p =  0.00000000E+00  i =  0.00000000E+00
+l:   9  p =  0.00000000E+00  i =  0.00000000E+00
 
 # fixed-boundary surface shape
-m:   0  n:   0  R1 =  1.0E+1
-m:   1  n:   0  R1 = -1.0E+0
-m:   1  n:   1  R1 = -3.0E-1
-m:  -1  n:  -1  R1 =  3.0E-1
-m:  -1  n:   0  Z1 =  1.0E+0
-m:  -1  n:   1  Z1 = -3.0E-1
-m:   1  n:  -1  Z1 = -3.0E-1
+l:   0  m:  -1  n:  -1  R1 =  3.00000000E-01  Z1 =  0.00000000E+00
+l:   0  m:   0  n:   0  R1 =  1.00000000E+01  Z1 =  0.00000000E+00
+l:   0  m:   1  n:   0  R1 = -1.00000000E+00  Z1 =  0.00000000E+00
+l:   0  m:   1  n:   1  R1 = -3.00000000E-01  Z1 =  0.00000000E+00
+l:   0  m:   1  n:  -1  R1 =  0.00000000E+00  Z1 = -3.00000000E-01
+l:   0  m:  -1  n:   0  R1 =  0.00000000E+00  Z1 =  1.00000000E+00
+l:   0  m:  -1  n:   1  R1 =  0.00000000E+00  Z1 = -3.00000000E-01

tests/test_curves.py

@@ -327,6 +327,20 @@ def test_to_FourierRZCurve(self):
         with pytest.raises(ValueError):
             xyz.to_FourierRZ(N=1, grid=grid)
 
+    @pytest.mark.unit
+    def test_change_symmetry(self):
+        """Test correct sym changes when only sym is passed to change_resolution."""
+        c = FourierRZCurve(sym=False)
+        c.change_resolution(sym=True)
+        assert c.sym
+        assert c.R_basis.sym == "cos"
+        assert c.Z_basis.sym == "sin"
+
+        c.change_resolution(sym=False)
+        assert c.sym is False
+        assert c.R_basis.sym is False
+        assert c.Z_basis.sym is False
+
 
 class TestFourierXYZCurve:
     """Tests for FourierXYZCurve class."""

tests/test_equilibrium.py

@@ -1,11 +1,11 @@
 """Tests for Equilibrium class."""
 
 import os
-import pickle
 import warnings
 
 import numpy as np
 import pytest
+from qic import Qic
 
 from desc.__main__ import main
 from desc.backend import sign
@@ -263,6 +263,9 @@ def test_eq_change_symmetry():
     assert eq.surface.sym
     assert eq.surface.R_basis.sym == "cos"
     assert eq.surface.Z_basis.sym == "sin"
+    assert eq.axis.sym
+    assert eq.axis.R_basis.sym == "cos"
+    assert eq.axis.Z_basis.sym == "sin"
 
     # undo symmetry
     eq.change_resolution(sym=False)
@@ -276,6 +279,9 @@ def test_eq_change_symmetry():
     assert not eq.surface.sym
     assert not eq.surface.R_basis.sym
     assert not eq.surface.Z_basis.sym
+    assert eq.axis.sym is False
+    assert eq.axis.R_basis.sym is False
+    assert eq.axis.Z_basis.sym is False
 
 
 @pytest.mark.unit
@@ -292,20 +298,58 @@ def test_resolution():
 @pytest.mark.unit
 def test_equilibrium_from_near_axis():
     """Test loading a solution from pyQSC/pyQIC."""
-    qsc_path = "./tests/inputs/qsc_r2section5.5.pkl"
-    file = open(qsc_path, "rb")
-    na = pickle.load(file)
-    file.close()
+    na = Qic.from_paper("r2 section 5.5", rs=[0, 1e-5], zc=[0, 1e-5])
 
     r = 1e-2
     eq = Equilibrium.from_near_axis(na, r=r, M=8, N=8)
     grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym)
     data = eq.compute("|B|", grid=grid)
 
+    # get the sin/cos modes
+    eq_rc = eq.Ra_n[
+        np.where(
+            np.logical_and(
+                eq.axis.R_basis.modes[:, 2] >= 0,
+                eq.axis.R_basis.modes[:, 2] < na.nfourier,
+            )
+        )
+    ]
+    eq_zc = eq.Za_n[
+        np.where(
+            np.logical_and(
+                eq.axis.Z_basis.modes[:, 2] >= 0,
+                eq.axis.Z_basis.modes[:, 2] < na.nfourier,
+            )
+        )
+    ]
+    eq_rs = np.flipud(
+        eq.Ra_n[
+            np.where(
+                np.logical_and(
+                    eq.axis.R_basis.modes[:, 2] < 0,
+                    eq.axis.R_basis.modes[:, 2] > -na.nfourier,
+                )
+            )
+        ]
+    )
+    eq_zs = np.flipud(
+        eq.Za_n[
+            np.where(
+                np.logical_and(
+                    eq.axis.Z_basis.modes[:, 2] < 0,
+                    eq.axis.Z_basis.modes[:, 2] > -na.nfourier,
+                )
+            )
+        ]
+    )
+
     assert eq.is_nested()
     assert eq.NFP == na.nfp
-    np.testing.assert_allclose(eq.Ra_n[:2], na.rc, atol=1e-10)
-    np.testing.assert_allclose(eq.Za_n[-2:], na.zs, atol=1e-10)
+    np.testing.assert_allclose(eq_rc, na.rc, atol=1e-10)
+    # na.zs[0] is always 0, which DESC doesn't include
+    np.testing.assert_allclose(eq_zs, na.zs[1:], atol=1e-10)
+    np.testing.assert_allclose(eq_rs, na.rs[1:], atol=1e-10)
+    np.testing.assert_allclose(eq_zc, na.zc, atol=1e-10)
     np.testing.assert_allclose(data["|B|"][0], na.B_mag(r, 0, 0), rtol=2e-2)
 
 

tests/test_input_output.py

@@ -178,12 +178,14 @@ def test_desc_output_to_input(tmpdir_factory):
     tmp_path = tmpdir.join("input_iotest_HELIOTRON")
     tmpout_path = tmpdir.join("iotest_HELIOTRON.h5")
     shutil.copyfile(outfile_path, tmpout_path)
+    eq = load(outfile_path)[-1]
 
     ir1 = InputReader()
     ir1.desc_output_to_input(str(tmp_path), str(tmpout_path))
     ir1 = InputReader(cl_args=[str(tmp_path)])
     arr1 = ir1.parse_inputs()[-1]["surface"]
     pres1 = ir1.parse_inputs()[-1]["pressure"]
+    iota1 = ir1.parse_inputs()[-1]["iota"]
 
     arr1 = arr1[arr1[:, 1].argsort()]
     arr1mneg = arr1[arr1[:, 1] < 0]
@@ -193,6 +195,8 @@ def test_desc_output_to_input(tmpdir_factory):
     ir2 = InputReader(cl_args=[str(desc_input_truth)])
     arr2 = ir2.parse_inputs()[-1]["surface"]
     pres2 = ir2.parse_inputs()[-1]["pressure"]
+    iota2 = ir2.parse_inputs()[-1]["iota"]
+
     arr2 = arr2[arr2[:, 1].argsort()]
     arr2mneg = arr2[arr2[:, 1] < 0]
     arr2mpos = arr2[arr2[:, 1] >= 0]
@@ -215,7 +219,10 @@ def test_desc_output_to_input(tmpdir_factory):
         atol=1e-8,
     )
 
-    np.testing.assert_allclose(pres1[pres1[:, 1] > 0], pres2[pres2[:, 1] > 0])
+    np.testing.assert_allclose(pres1, pres2)
+    np.testing.assert_allclose(pres1[:, 1], eq.pressure.params)
+    np.testing.assert_allclose(iota1, iota2)
+    np.testing.assert_allclose(iota1[:, 1], eq.iota.params)
 
 
 @pytest.mark.unit