Fix some plot issues when NFP differs from one for objects, or when passed-in phi exceeds 2pi/nfp

desc/plotting.py

@@ -1349,10 +1349,9 @@
     phi = np.atleast_1d(phi)
     nphi = len(phi)
     if grid is None:
-        nfp = eq.NFP
         grid_kwargs = {
             "L": 25,
-            "NFP": nfp,
+            "NFP": 1,
             "axis": False,
             "theta": np.linspace(0, 2 * np.pi, 91, endpoint=True),
             "zeta": phi,
@@ -1607,9 +1606,14 @@
     phi = np.atleast_1d(phi)
     nphi = len(phi)
 
+    # do not need NFP supplied to these grids as
+    # the above logic takes care of the correct phi range
+    # if defaults are requested. Setting NFP here instead
+    # can create reshaping issues when phi is supplied and gets
+    # truncated by 2pi/NFP. See PR #1204
     grid_kwargs = {
         "rho": rho,
-        "NFP": nfp,
+        "NFP": 1,
         "theta": np.linspace(0, 2 * np.pi, NT, endpoint=True),
         "zeta": phi,
     }
@@ -1628,7 +1632,7 @@
     )
     grid_kwargs = {
         "rho": np.linspace(0, 1, NR),
-        "NFP": nfp,
+        "NFP": 1,
         "theta": theta,
         "zeta": phi,
     }
@@ -1957,7 +1961,7 @@
     plot_axis = plot_axis and eq.L > 0
     rho = np.array([0.0, 1.0]) if plot_axis else np.array([1.0])
 
-    grid_kwargs = {"NFP": eq.NFP, "rho": rho, "theta": 100, "zeta": phi}
+    grid_kwargs = {"NFP": 1, "rho": rho, "theta": 100, "zeta": phi}
     grid = _get_grid(**grid_kwargs)
     nr, nt, nz = grid.num_rho, grid.num_theta, grid.num_zeta
     grid = Grid(
@@ -2027,6 +2031,13 @@
 ):
     """Plot stellarator boundaries at multiple toroidal coordinates.
 
+    NOTE: supplied objects must have either all the same NFP, or if
+    there are differing NFPs, the non-axisymmetric objects must have the
+    same NFP and the rest of the objects must be axisymmetric. i.e
+    can plot a tokamak and an NFP=2 stellarator, but cannot plot
+    a NFP=2 and NFP=3 stellarator as there is some ambiguity on the
+    choice of phi
+
     Parameters
     ----------
     eqs : array-like of Equilibrium, Surface or EquilibriaFamily
@@ -2082,6 +2093,28 @@
         fig, ax = plot_boundaries((eq1, eq2, eq3))
 
     """
+    NFPs = np.array([thing.NFP for thing in eqs])
+    Ns = np.array([thing.N for thing in eqs])
+    if not np.allclose(NFPs, NFPs[0]) and np.any(Ns == 0):
+        # if all NFPs are not equal, maybe there are some axisymmetric
+        # objects. We can try to change those to match the NFP of the first
+        # of the nonaxisymmetric objects
+        eqs = [thing.copy() for thing in eqs]  # make copy so we dont modify originals
+        NFP_nonax = int(NFPs[NFPs > 1][0])
+        [
+            thing.change_resolution(NFP=NFP_nonax if thing.N == 0 else thing.NFP)
+            for thing in eqs
+        ]
+    # if after above, the NFPs are still not all equal, means there are multiple
+    # nonaxisymmetric objects with differing NFPs, which it is not clear
+    # how to choose the phis for by default, so we will throw an error.
+    errorif(
+        not np.allclose([thing.NFP for thing in eqs], eqs[0].NFP),
+        ValueError,
+        "supplied objects must have the same number of field periods, "
+        "or if there are differing field periods, the ones which differ must be"
+        " axisymmetric.",
+    )
     phi = parse_argname_change(phi, kwargs, "zeta", "phi")
 
     figsize = kwargs.pop("figsize", None)
@@ -2126,7 +2159,7 @@
         plot_axis_i = plot_axis and eqs[i].L > 0
         rho = np.array([0.0, 1.0]) if plot_axis_i else np.array([1.0])
 
-        grid_kwargs = {"NFP": eqs[i].NFP, "theta": 100, "zeta": phi, "rho": rho}
+        grid_kwargs = {"NFP": 1, "theta": 100, "zeta": phi, "rho": rho}
         grid = _get_grid(**grid_kwargs)
         nr, nt, nz = grid.num_rho, grid.num_theta, grid.num_zeta
         grid = Grid(
@@ -2195,6 +2228,13 @@
 ):
     """Plot comparison between flux surfaces of multiple equilibria.
 
+    NOTE: supplied objects must have either all the same NFP, or if
+    there are differing NFPs, the non-axisymmetric objects must have the
+    same NFP and the rest of the objects must be axisymmetric. i.e
+    can plot a tokamak and an NFP=2 stellarator, but cannot plot
+    a NFP=2 and NFP=3 stellarator as there is some ambiguity on the
+    choice of phi
+
     Parameters
     ----------
     eqs : array-like of Equilibrium or EquilibriaFamily
@@ -2263,6 +2303,28 @@
                                  )
 
     """
+    NFPs = np.array([thing.NFP for thing in eqs])
+    Ns = np.array([thing.N for thing in eqs])
+    if not np.allclose(NFPs, NFPs[0]) and np.any(Ns == 0):
+        # if all NFPs are not equal, maybe there are some axisymmetric
+        # objects. We can try to change those to match the NFP of the first
+        # of the nonaxisymmetric objects
+        eqs = [thing.copy() for thing in eqs]  # make copy so we dont modify originals
+        NFP_nonax = int(NFPs[NFPs > 1][0])
+        [
+            thing.change_resolution(NFP=NFP_nonax if thing.N == 0 else thing.NFP)
+            for thing in eqs
+        ]
+    # if after above, the NFPs are still not all equal, means there are multiple
+    # nonaxisymmetric objects with differing NFPs, which it is not clear
+    # how to choose the phis for by default, so we will throw an error.
+    errorif(
+        not np.allclose([thing.NFP for thing in eqs], eqs[0].NFP),
+        ValueError,
+        "supplied objects must have the same number of field periods, "
+        "or if there are differing field periods, the ones which differ must be"
+        " axisymmetric.",
+    )
     phi = parse_argname_change(phi, kwargs, "zeta", "phi")
     color = parse_argname_change(color, kwargs, "colors", "color")
     ls = parse_argname_change(ls, kwargs, "linestyles", "ls")

tests/test_plotting.py

@@ -514,6 +514,9 @@ def test_plot_boundaries(self):
         eq1 = get("SOLOVEV")
         eq2 = get("DSHAPE")
         eq3 = get("W7-X")
+        eq4 = get("ESTELL")
+        with pytest.raises(ValueError, match="differing field periods"):
+            fig, ax = plot_boundaries([eq3, eq4], theta=0)
         fig, ax, data = plot_boundaries((eq1, eq2, eq3), return_data=True)
         assert "R" in data.keys()
         assert "Z" in data.keys()
@@ -551,6 +554,18 @@ def test_plot_comparison_no_theta(self):
         fig, ax = plot_comparison(eqf, theta=0)
         return fig
 
+    @pytest.mark.unit
+    @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d)
+    def test_plot_comparison_different_NFPs(self):
+        """Test plotting comparison of flux surfaces with differing NFPs."""
+        eq = get("SOLOVEV")
+        eq_nonax = get("HELIOTRON")
+        eq_nonax2 = get("ESTELL")
+        with pytest.raises(ValueError, match="differing field periods"):
+            fig, ax = plot_comparison([eq_nonax, eq_nonax2], theta=0)
+        fig, ax = plot_comparison([eq, eq_nonax], theta=0)
+        return fig
+
 
 class TestPlotGrid:
     """Tests for the plot_grid function."""