diff --git a/desc/plotting.py b/desc/plotting.py index 430c9ad53..4c0e6d497 100644 --- 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 index e43ba54cd..f3927bec6 100644 Binary files a/tests/baseline/test_bounce1d_checks.png and b/tests/baseline/test_bounce1d_checks.png differ diff --git a/tests/baseline/test_plot_b_mag.png b/tests/baseline/test_plot_b_mag.png index f86a3b183..c81f89f03 100644 Binary files a/tests/baseline/test_plot_b_mag.png and b/tests/baseline/test_plot_b_mag.png differ diff --git a/tests/baseline/test_plot_boozer_modes.png b/tests/baseline/test_plot_boozer_modes.png index b72e30839..a2d1c6e37 100644 Binary files a/tests/baseline/test_plot_boozer_modes.png and b/tests/baseline/test_plot_boozer_modes.png differ diff --git a/tests/baseline/test_plot_boozer_modes_breaking_only.png b/tests/baseline/test_plot_boozer_modes_breaking_only.png index 424ba1f66..06f04654f 100644 Binary files a/tests/baseline/test_plot_boozer_modes_breaking_only.png and b/tests/baseline/test_plot_boozer_modes_breaking_only.png differ diff --git a/tests/baseline/test_plot_boozer_modes_max.png b/tests/baseline/test_plot_boozer_modes_max.png index ac59a48ed..b3d52bbe2 100644 Binary files a/tests/baseline/test_plot_boozer_modes_max.png and b/tests/baseline/test_plot_boozer_modes_max.png differ diff --git a/tests/baseline/test_plot_boozer_modes_no_norm.png b/tests/baseline/test_plot_boozer_modes_no_norm.png index b2fe1e29b..a5ccc0e32 100644 Binary files a/tests/baseline/test_plot_boozer_modes_no_norm.png and b/tests/baseline/test_plot_boozer_modes_no_norm.png differ diff --git a/tests/baseline/test_plot_boundaries.png b/tests/baseline/test_plot_boundaries.png index 44d4792c1..9e730b85a 100644 Binary files a/tests/baseline/test_plot_boundaries.png and b/tests/baseline/test_plot_boundaries.png differ diff --git a/tests/baseline/test_plot_boundary.png b/tests/baseline/test_plot_boundary.png index 8f2dd2769..ab4e97d73 100644 Binary files a/tests/baseline/test_plot_boundary.png and b/tests/baseline/test_plot_boundary.png differ diff --git a/tests/baseline/test_plot_boundary_surface.png b/tests/baseline/test_plot_boundary_surface.png index 5dfe95191..1e4f60d91 100644 Binary files a/tests/baseline/test_plot_boundary_surface.png and b/tests/baseline/test_plot_boundary_surface.png differ diff --git a/tests/baseline/test_plot_qs_error.png b/tests/baseline/test_plot_qs_error.png index c90e59fb8..10485feb2 100644 Binary files a/tests/baseline/test_plot_qs_error.png and b/tests/baseline/test_plot_qs_error.png differ diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 6048a7d60..33f3e1152 100644 --- 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