From f75e4a2b3e6b39c0ccde11e8028c5484b8cbe478 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 23 Jul 2024 15:37:00 -0400 Subject: [PATCH 1/9] remove ds from compute index in favor of using grad spacing directly --- desc/compute/_curve.py | 34 ++++++++-------------------------- 1 file changed, 8 insertions(+), 26 deletions(-) diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index 818da18092..41cd2e44e3 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -26,25 +26,6 @@ def _s(params, transforms, profiles, data, **kwargs): return data -@register_compute_fun( - name="ds", - label="ds", - units="~", - units_long="None", - description="Spacing of curve parameter", - dim=1, - params=[], - transforms={"grid": []}, - profiles=[], - coordinates="s", - data=[], - parameterization="desc.geometry.core.Curve", -) -def _ds(params, transforms, profiles, data, **kwargs): - data["ds"] = transforms["grid"].spacing[:, 2] - return data - - @register_compute_fun( name="X", label="X", @@ -980,17 +961,17 @@ def _torsion(params, transforms, profiles, data, **kwargs): description="Length of the curve", dim=0, params=[], - transforms={}, + transforms={"grid": []}, profiles=[], coordinates="", - data=["ds", "x_s"], + data=["x_s"], parameterization=["desc.geometry.core.Curve"], ) def _length(params, transforms, profiles, data, **kwargs): T = jnp.linalg.norm(data["x_s"], axis=-1) # this is equivalent to jnp.trapz(T, s) for a closed curve, # but also works if grid.endpoint is False - data["length"] = jnp.sum(T * data["ds"]) + data["length"] = jnp.sum(T * transforms["grid"].spacing[:, 2]) return data @@ -1002,10 +983,10 @@ def _length(params, transforms, profiles, data, **kwargs): description="Length of the curve", dim=0, params=[], - transforms={}, + transforms={"grid": []}, profiles=[], coordinates="", - data=["ds", "x", "x_s"], + data=["x", "x_s"], parameterization="desc.geometry.curve.SplineXYZCurve", method="Interpolation type, Default 'cubic'. See SplineXYZCurve docs for options.", ) @@ -1015,7 +996,8 @@ def _length_SplineXYZCurve(params, transforms, profiles, data, **kwargs): if kwargs.get("basis", "rpz").lower() == "rpz": coords = rpz2xyz(coords) # ensure curve is closed - # if it's already closed this doesn't add any length since ds will be zero + # if it's already closed this doesn't add any length since + # grid spacing will be zero at the duplicate point coords = jnp.concatenate([coords, coords[:1]]) X = coords[:, 0] Y = coords[:, 1] @@ -1026,5 +1008,5 @@ def _length_SplineXYZCurve(params, transforms, profiles, data, **kwargs): T = jnp.linalg.norm(data["x_s"], axis=-1) # this is equivalent to jnp.trapz(T, s) for a closed curve # but also works if grid.endpoint is False - data["length"] = jnp.sum(T * data["ds"]) + data["length"] = jnp.sum(T * transforms["grid"].spacing[:, 2]) return data From b9ba0beaf15f1e992a89f4d8004fafce4bf6c96c Mon Sep 17 00:00:00 2001 From: YigitElma Date: Tue, 13 Aug 2024 00:57:59 -0400 Subject: [PATCH 2/9] fix the sign error in wide QR factorization --- desc/optimize/tr_subproblems.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/optimize/tr_subproblems.py b/desc/optimize/tr_subproblems.py index c9149d2bb3..b107bca0a4 100644 --- a/desc/optimize/tr_subproblems.py +++ b/desc/optimize/tr_subproblems.py @@ -421,7 +421,7 @@ def trust_region_step_exact_qr( p_newton = solve_triangular_regularized(R, -Q.T @ f) else: Q, R = qr(J.T, mode="economic") - p_newton = Q @ solve_triangular_regularized(R.T, f, lower=True) + p_newton = Q @ solve_triangular_regularized(R.T, -f, lower=True) def truefun(*_): return p_newton, False, 0.0 From 9549b7d0d59149b1bd8d02fc55deff9a080bf297 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 13 Aug 2024 20:25:21 -0400 Subject: [PATCH 3/9] 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 4/9] 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 550b20d5c190ba610449d342eaed1216857d13fb Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 14 Aug 2024 11:04:06 -0400 Subject: [PATCH 5/9] Revert "remove ds from compute index in favor of using grad spacing directly" This reverts commit f75e4a2b3e6b39c0ccde11e8028c5484b8cbe478. --- desc/compute/_curve.py | 34 ++++++++++++++++++++++++++-------- 1 file changed, 26 insertions(+), 8 deletions(-) diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index 41cd2e44e3..818da18092 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -26,6 +26,25 @@ def _s(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="ds", + label="ds", + units="~", + units_long="None", + description="Spacing of curve parameter", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.core.Curve", +) +def _ds(params, transforms, profiles, data, **kwargs): + data["ds"] = transforms["grid"].spacing[:, 2] + return data + + @register_compute_fun( name="X", label="X", @@ -961,17 +980,17 @@ def _torsion(params, transforms, profiles, data, **kwargs): description="Length of the curve", dim=0, params=[], - transforms={"grid": []}, + transforms={}, profiles=[], coordinates="", - data=["x_s"], + data=["ds", "x_s"], parameterization=["desc.geometry.core.Curve"], ) def _length(params, transforms, profiles, data, **kwargs): T = jnp.linalg.norm(data["x_s"], axis=-1) # this is equivalent to jnp.trapz(T, s) for a closed curve, # but also works if grid.endpoint is False - data["length"] = jnp.sum(T * transforms["grid"].spacing[:, 2]) + data["length"] = jnp.sum(T * data["ds"]) return data @@ -983,10 +1002,10 @@ def _length(params, transforms, profiles, data, **kwargs): description="Length of the curve", dim=0, params=[], - transforms={"grid": []}, + transforms={}, profiles=[], coordinates="", - data=["x", "x_s"], + data=["ds", "x", "x_s"], parameterization="desc.geometry.curve.SplineXYZCurve", method="Interpolation type, Default 'cubic'. See SplineXYZCurve docs for options.", ) @@ -996,8 +1015,7 @@ def _length_SplineXYZCurve(params, transforms, profiles, data, **kwargs): if kwargs.get("basis", "rpz").lower() == "rpz": coords = rpz2xyz(coords) # ensure curve is closed - # if it's already closed this doesn't add any length since - # grid spacing will be zero at the duplicate point + # if it's already closed this doesn't add any length since ds will be zero coords = jnp.concatenate([coords, coords[:1]]) X = coords[:, 0] Y = coords[:, 1] @@ -1008,5 +1026,5 @@ def _length_SplineXYZCurve(params, transforms, profiles, data, **kwargs): T = jnp.linalg.norm(data["x_s"], axis=-1) # this is equivalent to jnp.trapz(T, s) for a closed curve # but also works if grid.endpoint is False - data["length"] = jnp.sum(T * transforms["grid"].spacing[:, 2]) + data["length"] = jnp.sum(T * data["ds"]) return data From b369c5ea6e1ee2c1e9782f6e3672f98198de6daf Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 14 Aug 2024 11:06:05 -0400 Subject: [PATCH 6/9] just change description of ds --- desc/compute/_curve.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index 818da18092..d1bd231f6c 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -31,7 +31,10 @@ def _s(params, transforms, profiles, data, **kwargs): label="ds", units="~", units_long="None", - description="Spacing of curve parameter", + description=( + "Quadrature weights for integration along the curve, " + + " i.e. an alias for ``grid.spacing[:,2]``" + ), dim=1, params=[], transforms={"grid": []}, From 98190a31839b5ad3a0d02093b1ddcb8294d903b7 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 14 Aug 2024 16:59:18 -0400 Subject: [PATCH 7/9] remove double space --- desc/compute/_curve.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index d1bd231f6c..2e96e7a767 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -32,7 +32,7 @@ def _s(params, transforms, profiles, data, **kwargs): units="~", units_long="None", description=( - "Quadrature weights for integration along the curve, " + "Quadrature weights for integration along the curve," + " i.e. an alias for ``grid.spacing[:,2]``" ), dim=1, From c5a49f66ae146418c545e4585953c9e30e2ab664 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 15 Aug 2024 13:34:07 -0400 Subject: [PATCH 8/9] 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 9/9] 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(