Merge branch 'master' into felicia

PlasmaControl · Oct 2, 2024 · 53d2ebb · 53d2ebb
2 parents fdd9336 + f09eaa7
commit 53d2ebb
Show file tree

Hide file tree

Showing 13 changed files with 23 additions and 14 deletions.
diff --git a/desc/objectives/_omnigenity.py b/desc/objectives/_omnigenity.py
@@ -163,7 +163,7 @@ def build(self, use_jit=True, verbose=1):
             M_booz=M_booz,
             N_booz=N_booz,
         )
-        matrix, modes, idx = ptolemy_linear_transform(
+        matrix, _, idx = ptolemy_linear_transform(
             transforms["B"].basis.modes,
             helicity=self.helicity,
             NFP=transforms["B"].basis.NFP,
@@ -180,7 +180,7 @@ def build(self, use_jit=True, verbose=1):
         if verbose > 1:
             timer.disp("Precomputing transforms")
 
-        self._dim_f = idx.size
+        self._dim_f = idx.size * grid.num_rho
 
         if self._normalize:
             scales = compute_scaling_factors(eq)

diff --git a/desc/plotting.py b/desc/plotting.py
@@ -2016,7 +2016,7 @@ def plot_boundary(eq, phi=None, plot_axis=True, ax=None, return_data=False, **kw
     ax.set_ylabel(_AXIS_LABELS_RPZ[2], fontsize=ylabel_fontsize)
     ax.tick_params(labelbottom=True, labelleft=True)
 
-    fig.legend(**legend_kw)
+    ax.legend(**legend_kw)
     _set_tight_layout(fig)
 
     plot_data = {}
@@ -2120,9 +2120,7 @@ def plot_boundaries(
 
     phi = (1 if eqs[-1].N == 0 else 4) if phi is None else phi
     if isinstance(phi, numbers.Integral):
-        phi = np.linspace(
-            0, 2 * np.pi / eqs[-1].NFP, phi + 1
-        )  # +1 to include pi and 2pi
+        phi = np.linspace(0, 2 * np.pi / eqs[-1].NFP, phi, endpoint=False)
     phi = np.atleast_1d(phi)
 
     neq = len(eqs)
@@ -2173,7 +2171,7 @@ def plot_boundaries(
         plot_data["R"].append(R)
         plot_data["Z"].append(Z)
 
-        for j in range(nz - 1):
+        for j in range(nz):
             (line,) = ax.plot(
                 R[:, -1, j], Z[:, -1, j], color=colors[i], linestyle=ls[i], lw=lw[i]
             )
@@ -2190,7 +2188,7 @@ def plot_boundaries(
     ax.tick_params(labelbottom=True, labelleft=True)
 
     if any(labels) and kwargs.pop("legend", True):
-        fig.legend(**kwargs.pop("legend_kw", {}))
+        ax.legend(**kwargs.pop("legend_kw", {}))
     _set_tight_layout(fig)
 
     assert (
@@ -2706,7 +2704,7 @@ def plot_boozer_modes(  # noqa: C901
     ax.set_ylabel(ylabel, fontsize=ylabel_fontsize)
 
     if kwargs.pop("legend", True):
-        fig.legend(**kwargs.pop("legend_kw", {"loc": "lower right"}))
+        ax.legend(**kwargs.pop("legend_kw", {"loc": "lower right"}))
 
     assert (
         len(kwargs) == 0
@@ -3119,7 +3117,7 @@ def plot_qs_error(  # noqa: 16 fxn too complex
         ax.set_ylabel(ylabel, fontsize=ylabel_fontsize)
 
     if kwargs.pop("legend", True):
-        fig.legend(**kwargs.pop("legend_kw", {"loc": "center right"}))
+        ax.legend(**kwargs.pop("legend_kw", {"loc": "center right"}))
 
     assert (
         len(kwargs) == 0

diff --git a/tests/baseline/test_bounce1d_checks.png b/tests/baseline/test_bounce1d_checks.png
diff --git a/tests/baseline/test_plot_b_mag.png b/tests/baseline/test_plot_b_mag.png
diff --git a/tests/baseline/test_plot_boozer_modes.png b/tests/baseline/test_plot_boozer_modes.png
diff --git a/tests/baseline/test_plot_boozer_modes_breaking_only.png b/tests/baseline/test_plot_boozer_modes_breaking_only.png
diff --git a/tests/baseline/test_plot_boozer_modes_max.png b/tests/baseline/test_plot_boozer_modes_max.png
diff --git a/tests/baseline/test_plot_boozer_modes_no_norm.png b/tests/baseline/test_plot_boozer_modes_no_norm.png
diff --git a/tests/baseline/test_plot_boundaries.png b/tests/baseline/test_plot_boundaries.png
diff --git a/tests/baseline/test_plot_boundary.png b/tests/baseline/test_plot_boundary.png
diff --git a/tests/baseline/test_plot_boundary_surface.png b/tests/baseline/test_plot_boundary_surface.png
diff --git a/tests/baseline/test_plot_qs_error.png b/tests/baseline/test_plot_qs_error.png
diff --git a/tests/test_plotting.py b/tests/test_plotting.py
@@ -47,9 +47,9 @@
 from desc.utils import isalmostequal
 from desc.vmec import VMECIO
 
-tol_1d = 7.8
-tol_2d = 15
-tol_3d = 15
+tol_1d = 4.5
+tol_2d = 10
+tol_3d = 10
 
 
 @pytest.mark.unit
@@ -507,7 +507,7 @@ def test_plot_boundary_surface(self):
         return fig
 
     @pytest.mark.unit
-    @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d)
+    @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d)
     def test_plot_boundaries(self):
         """Test plotting boundaries."""
         eq1 = get("SOLOVEV")
@@ -516,6 +516,11 @@ def test_plot_boundaries(self):
         eq4 = get("ESTELL")
         with pytest.raises(ValueError, match="differing field periods"):
             fig, ax = plot_boundaries([eq3, eq4], theta=0)
+        _, _, data1 = plot_boundaries(
+            (eq1, eq2, eq3),
+            phi=4,
+            return_data=True,
+        )
         fig, ax, data = plot_boundaries(
             (eq1, eq2, eq3),
             phi=np.linspace(0, 2 * np.pi / eq3.NFP, 4, endpoint=False),
@@ -525,6 +530,12 @@ def test_plot_boundaries(self):
         assert "Z" in data.keys()
         assert len(data["R"]) == 3
         assert len(data["Z"]) == 3
+        assert (
+            data["R"][-1].shape == data1["R"][-1].shape
+        ), "Passing phi as an integer or array results in different behavior"
+        assert (
+            data["Z"][-1].shape == data1["Z"][-1].shape
+        ), "Passing phi as an integer or array results in different behavior"
 
         return fig