From 9549b7d0d59149b1bd8d02fc55deff9a080bf297 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 13 Aug 2024 20:25:21 -0400 Subject: [PATCH 1/5] reduce test stringency for NAE --- tests/test_examples.py | 103 ++++++----------------------------------- 1 file changed, 15 insertions(+), 88 deletions(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index 8f6e42ccec..430258cb8f 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -440,7 +440,6 @@ def test_NAE_QSC_solve(): orig_Rax_val = eq.axis.R_n orig_Zax_val = eq.axis.Z_n - eq_fit = eq.copy() eq_lambda_fixed_0th_order = eq.copy() eq_lambda_fixed_1st_order = eq.copy() @@ -472,7 +471,6 @@ def test_NAE_QSC_solve(): xtol=1e-6, constraints=constraints, ) - grid = LinearGrid(L=10, M=20, N=20, NFP=eq.NFP, sym=True, axis=False) grid_axis = LinearGrid(rho=0.0, theta=0.0, N=eq.N_grid, NFP=eq.NFP) # Make sure axis is same for eqq, string in zip( @@ -482,24 +480,15 @@ def test_NAE_QSC_solve(): np.testing.assert_array_almost_equal(orig_Rax_val, eqq.axis.R_n, err_msg=string) np.testing.assert_array_almost_equal(orig_Zax_val, eqq.axis.Z_n, err_msg=string) - # Make sure surfaces of solved equilibrium are similar near axis as QSC - rho_err, theta_err = area_difference_desc(eqq, eq_fit) + # Make sure iota of solved equilibrium is same on-axis as QSC - np.testing.assert_allclose(rho_err[:, 0:-6], 0, atol=1.5e-2, err_msg=string) - np.testing.assert_allclose(theta_err[:, 0:-6], 0, atol=1e-3, err_msg=string) - - # Make sure iota of solved equilibrium is same near axis as QSC - - iota = grid.compress(eqq.compute("iota", grid=grid)["iota"]) + iota = eqq.compute("iota", grid=LinearGrid(rho=0.0))["iota"] np.testing.assert_allclose(iota[0], qsc.iota, atol=1e-5, err_msg=string) - np.testing.assert_allclose(iota[1:10], qsc.iota, atol=1e-3, err_msg=string) - # check lambda to match near axis - # Evaluate lambda near the axis + # check lambda to match on-axis data_nae = eqq.compute(["lambda", "|B|"], grid=grid_axis) lam_nae = data_nae["lambda"] - # Reshape to form grids on theta and phi phi = np.squeeze(grid_axis.nodes[:, 2]) np.testing.assert_allclose( @@ -526,7 +515,6 @@ def test_NAE_QSC_solve_near_axis_based_off_eq(): orig_Rax_val = eq.axis.R_n orig_Zax_val = eq.axis.Z_n - eq_fit = eq.copy() eq_lambda_fixed_0th_order = eq.copy() eq_lambda_fixed_1st_order = eq.copy() @@ -558,7 +546,6 @@ def test_NAE_QSC_solve_near_axis_based_off_eq(): xtol=1e-6, constraints=constraints, ) - grid = LinearGrid(L=10, M=20, N=20, NFP=eq.NFP, sym=True, axis=False) grid_axis = LinearGrid(rho=0.0, theta=0.0, N=eq.N_grid, NFP=eq.NFP) # Make sure axis is same for eqq, string in zip( @@ -568,18 +555,11 @@ def test_NAE_QSC_solve_near_axis_based_off_eq(): np.testing.assert_array_almost_equal(orig_Rax_val, eqq.axis.R_n, err_msg=string) np.testing.assert_array_almost_equal(orig_Zax_val, eqq.axis.Z_n, err_msg=string) - # Make sure surfaces of solved equilibrium are similar near axis as QSC - rho_err, theta_err = area_difference_desc(eqq, eq_fit) - - np.testing.assert_allclose(rho_err[:, 0:-6], 0, atol=1.5e-2, err_msg=string) - np.testing.assert_allclose(theta_err[:, 0:-6], 0, atol=1e-3, err_msg=string) - - # Make sure iota of solved equilibrium is same near axis as QSC + # Make sure iota of solved equilibrium is same on axis as QSC - iota = grid.compress(eqq.compute("iota", grid=grid)["iota"]) + iota = eqq.compute("iota", grid=LinearGrid(rho=0.0))["iota"] np.testing.assert_allclose(iota[0], qsc.iota, atol=1e-5, err_msg=string) - np.testing.assert_allclose(iota[1:10], qsc.iota, rtol=5e-3, err_msg=string) ### check lambda to match on axis # Evaluate lambda on the axis @@ -603,18 +583,16 @@ def test_NAE_QSC_solve_near_axis_based_off_eq(): def test_NAE_QIC_solve(): """Test O(rho) NAE QIC constraints solve.""" # get Qic example - qic = Qic.from_paper("QI NFP2 r2", nphi=301, order="r1") + qic = Qic.from_paper("QI NFP2 r2", nphi=199, order="r1") qic.lasym = False # don't need to consider stellarator asym for order 1 constraints ntheta = 75 r = 0.01 - N = 11 + N = 9 eq = Equilibrium.from_near_axis(qic, r=r, L=7, M=7, N=N, ntheta=ntheta) orig_Rax_val = eq.axis.R_n orig_Zax_val = eq.axis.Z_n - eq_fit = eq.copy() - # this has all the constraints we need, cs = get_NAE_constraints(eq, qic, order=1) @@ -629,75 +607,24 @@ def test_NAE_QIC_solve(): np.testing.assert_array_almost_equal(orig_Rax_val, eq.axis.R_n) np.testing.assert_array_almost_equal(orig_Zax_val, eq.axis.Z_n) - # Make sure surfaces of solved equilibrium are similar near axis as QIC - rho_err, theta_err = area_difference_desc(eq, eq_fit) - - np.testing.assert_allclose(rho_err[:, 0:3], 0, atol=5e-2) - # theta error isn't really an indicator of near axis behavior - # since it's computed over the full radius, but just indicates that - # eq is similar to eq_fit - np.testing.assert_allclose(theta_err, 0, atol=5e-2) - # Make sure iota of solved equilibrium is same near axis as QIC - grid = LinearGrid(L=10, M=20, N=20, NFP=eq.NFP, sym=True, axis=False) - iota = grid.compress(eq.compute("iota", grid=grid)["iota"]) + iota = eq.compute("iota", grid=LinearGrid(rho=0.0))["iota"] np.testing.assert_allclose(iota[0], qic.iota, rtol=1e-5) - np.testing.assert_allclose(iota[1:10], qic.iota, rtol=1e-3) - # check lambda to match near axis - grid_2d_05 = LinearGrid(rho=np.array(1e-6), M=50, N=50, NFP=eq.NFP, endpoint=True) + grid_axis = LinearGrid(rho=0.0, theta=0.0, zeta=qic.phi, NFP=eq.NFP) + phi = grid_axis.nodes[:, 2].squeeze() - # Evaluate lambda near the axis - data_nae = eq.compute("lambda", grid=grid_2d_05) - lam_nae = data_nae["lambda"] + # check lambda to match on-axis + lam_nae = eq.compute("lambda", grid=grid_axis)["lambda"] - # Reshape to form grids on theta and phi - zeta = ( - grid_2d_05.nodes[:, 2] - .reshape( - (grid_2d_05.num_theta, grid_2d_05.num_rho, grid_2d_05.num_zeta), order="F" - ) - .squeeze() - ) - - lam_nae = lam_nae.reshape( - (grid_2d_05.num_theta, grid_2d_05.num_rho, grid_2d_05.num_zeta), order="F" - ) - - phi = np.squeeze(zeta[0, :]) - lam_nae = np.squeeze(lam_nae[:, 0, :]) - - lam_av_nae = np.mean(lam_nae, axis=0) np.testing.assert_allclose( - lam_av_nae, -qic.iota * qic.nu_spline(phi), atol=1e-4, rtol=1e-2 + lam_nae, -qic.iota * qic.nu_spline(phi), atol=1e-4, rtol=1e-2 ) # check |B| on axis - - grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, rho=np.array(1e-6)) - # Evaluate B modes near the axis - data_nae = eq.compute(["|B|_mn", "B modes"], grid=grid) - modes = data_nae["B modes"] - B_mn_nae = data_nae["|B|_mn"] - # Evaluate B on an angular grid - theta = np.linspace(0, 2 * np.pi, 150) - phi = np.linspace(0, 2 * np.pi, qic.nphi) - th, ph = np.meshgrid(theta, phi) - B_nae = np.zeros((qic.nphi, 150)) - - for i, (l, m, n) in enumerate(modes): - if m >= 0 and n >= 0: - B_nae += B_mn_nae[i] * np.cos(m * th) * np.cos(n * ph) - elif m >= 0 > n: - B_nae += -B_mn_nae[i] * np.cos(m * th) * np.sin(n * ph) - elif m < 0 <= n: - B_nae += -B_mn_nae[i] * np.sin(m * th) * np.cos(n * ph) - elif m < 0 and n < 0: - B_nae += B_mn_nae[i] * np.sin(m * th) * np.sin(n * ph) - # Eliminate the poloidal angle to focus on the toroidal behavior - B_av_nae = np.mean(B_nae, axis=1) - np.testing.assert_allclose(B_av_nae, np.ones(np.size(phi)) * qic.B0, atol=2e-2) + B_nae = eq.compute(["|B|"], grid=grid_axis) + np.testing.assert_allclose(B_nae, qic.B0, atol=1e-3) @pytest.mark.unit From b397e89a32a0895c6ef9ffa0a3de7e3e5e233908 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 14 Aug 2024 11:00:01 -0400 Subject: [PATCH 2/5] fix test --- tests/test_examples.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index 430258cb8f..8187375caf 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -623,7 +623,7 @@ def test_NAE_QIC_solve(): ) # check |B| on axis - B_nae = eq.compute(["|B|"], grid=grid_axis) + B_nae = eq.compute(["|B|"], grid=grid_axis)["|B|"] np.testing.assert_allclose(B_nae, qic.B0, atol=1e-3) From c5a49f66ae146418c545e4585953c9e30e2ab664 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 15 Aug 2024 13:34:07 -0400 Subject: [PATCH 3/5] increasing grid resolution seems to resolve issue --- tests/test_examples.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index 8f6e42ccec..c336a2f624 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1700,7 +1700,7 @@ def test_signed_PlasmaVesselDistance(): eq = Equilibrium(M=1, N=1) surf = eq.surface.copy() surf.change_resolution(M=1, N=1) - grid = LinearGrid(M=10, N=2, NFP=eq.NFP) + grid = LinearGrid(M=10, N=4, NFP=eq.NFP) obj = PlasmaVesselDistance( surface=surf, @@ -1723,4 +1723,9 @@ def test_signed_PlasmaVesselDistance(): xtol=1e-9, ) - np.testing.assert_allclose(obj.compute(*obj.xs(eq, surf)), target_dist, atol=1e-2) + np.testing.assert_allclose( + obj.compute(*obj.xs(eq, surf)), + target_dist, + atol=1e-2, + err_msg="allowing eq to change", + ) From 4aeb0a73a1913d7a0019ddd4c88a986aad7d110d Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 15 Aug 2024 14:35:39 -0400 Subject: [PATCH 4/5] further increase grid res --- tests/test_examples.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index c336a2f624..9b07014042 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1700,7 +1700,7 @@ def test_signed_PlasmaVesselDistance(): eq = Equilibrium(M=1, N=1) surf = eq.surface.copy() surf.change_resolution(M=1, N=1) - grid = LinearGrid(M=10, N=4, NFP=eq.NFP) + grid = LinearGrid(M=20, N=8, NFP=eq.NFP) obj = PlasmaVesselDistance( surface=surf, @@ -1710,7 +1710,7 @@ def test_signed_PlasmaVesselDistance(): plasma_grid=grid, use_signed_distance=True, ) - objective = ObjectiveFunction((obj,)) + objective = ObjectiveFunction(obj) optimizer = Optimizer("lsq-exact") (eq, surf), _ = optimizer.optimize( From e39d08fc65087c18b8d1b409d55c14cd66cdbb71 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 16 Aug 2024 10:42:54 -0400 Subject: [PATCH 5/5] improve error msg for plot coilis --- desc/plotting.py | 8 ++++++++ tests/test_plotting.py | 2 ++ 2 files changed, 10 insertions(+) diff --git a/desc/plotting.py b/desc/plotting.py index ba822a5b2b..478e11dc66 100644 --- a/desc/plotting.py +++ b/desc/plotting.py @@ -2384,6 +2384,8 @@ def plot_coils(coils, grid=None, fig=None, return_data=False, **kwargs): dictionary of the data plotted, only returned if ``return_data=True`` """ + from desc.coils import _Coil + lw = kwargs.pop("lw", 5) ls = kwargs.pop("ls", "solid") figsize = kwargs.pop("figsize", (10, 10)) @@ -2394,6 +2396,12 @@ def plot_coils(coils, grid=None, fig=None, return_data=False, **kwargs): ValueError, f"plot_coils got unexpected keyword argument: {kwargs.keys()}", ) + errorif( + not isinstance(coils, _Coil), + ValueError, + "Expected `coils` to be object of type `_Coil`, instead got type" + f" {type(coils)}", + ) if not isinstance(lw, (list, tuple)): lw = [lw] diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 8252d349d8..1351b46c93 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -801,6 +801,8 @@ def test_plot_coils(): coil.rotate(angle=np.pi / N) coils = CoilSet.linspaced_angular(coil, I, [0, 0, 1], np.pi / NFP, N // NFP // 2) coils2 = MixedCoilSet.from_symmetry(coils, NFP, True) + with pytest.raises(ValueError, match="Expected `coils`"): + plot_coils("not a coil") fig, data = plot_coils(coils2, return_data=True) def flatten_coils(coilset):