From a94782eb7eea0dd21422fc00a215ec8369aeb346 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 10 Sep 2024 19:32:18 -0600 Subject: [PATCH 1/2] add flip_theta --- desc/compat.py | 40 +++++++++++++++++++ docs/api.rst | 1 + docs/api_equilibrium.rst | 1 + tests/test_compat.py | 86 ++++++++++++++++++++++++++++++++++++++-- 4 files changed, 125 insertions(+), 3 deletions(-) diff --git a/desc/compat.py b/desc/compat.py index bea90d470e..a7e7732e87 100644 --- a/desc/compat.py +++ b/desc/compat.py @@ -111,6 +111,46 @@ def flip_helicity(eq): return eq +def flip_theta(eq): + """Change the gauge freedom of the poloidal angle of an Equilibrium. + + Equivalent to redefining theta_new = theta_old + π + + Parameters + ---------- + eq : Equilibrium or iterable of Equilibrium + Equilibria to flip the helicity of. + + Returns + ------- + eq : Equilibrium or iterable of Equilibrium + Same as input, but with the poloidal angle redefined. + + """ + # maybe it's iterable: + if hasattr(eq, "__len__"): + for e in eq: + flip_theta(e) + return eq + + rone = np.ones_like(eq.R_lmn) + rone[eq.R_basis.modes[:, 1] % 2 == 1] *= -1 + eq.R_lmn *= rone + + zone = np.ones_like(eq.Z_lmn) + zone[eq.Z_basis.modes[:, 1] % 2 == 1] *= -1 + eq.Z_lmn *= zone + + lone = np.ones_like(eq.L_lmn) + lone[eq.L_basis.modes[:, 1] % 2 == 1] *= -1 + eq.L_lmn *= lone + + eq.axis = eq.get_axis() + eq.surface = eq.get_surface_at(rho=1) + + return eq + + def rescale( eq, L=("R0", None), B=("B0", None), scale_pressure=True, copy=False, verbose=0 ): diff --git a/docs/api.rst b/docs/api.rst index c5147d4077..02c7cc8c73 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -41,6 +41,7 @@ Compatibility desc.compat.ensure_positive_jacobian desc.compat.flip_helicity + desc.compat.flip_theta desc.compat.rescale Continuation diff --git a/docs/api_equilibrium.rst b/docs/api_equilibrium.rst index 2adc6296c8..5b349d79e5 100644 --- a/docs/api_equilibrium.rst +++ b/docs/api_equilibrium.rst @@ -66,4 +66,5 @@ equilibria to a given size and/or field strength. desc.compat.ensure_positive_jacobian desc.compat.flip_helicity + desc.compat.flip_theta desc.compat.rescale diff --git a/tests/test_compat.py b/tests/test_compat.py index 7a13387fce..68e14748ad 100644 --- a/tests/test_compat.py +++ b/tests/test_compat.py @@ -3,7 +3,7 @@ import numpy as np import pytest -from desc.compat import flip_helicity, rescale +from desc.compat import flip_helicity, flip_theta, rescale from desc.examples import get from desc.grid import Grid, LinearGrid, QuadratureGrid @@ -47,7 +47,6 @@ def test_flip_helicity_axisym(): @pytest.mark.unit -@pytest.mark.solve def test_flip_helicity_iota(): """Test flip_helicity on an Equilibrium with an iota profile.""" eq = get("HELIOTRON") @@ -90,7 +89,6 @@ def test_flip_helicity_iota(): @pytest.mark.unit -@pytest.mark.solve def test_flip_helicity_current(): """Test flip_helicity on an Equilibrium with a current profile.""" eq = get("HSX") @@ -135,6 +133,88 @@ def test_flip_helicity_current(): np.testing.assert_allclose(data_old["f_C"], data_flip["f_C"], atol=1e-8) +@pytest.mark.unit +def test_flip_theta_axisym(): + """Test flip_theta on an axisymmetric Equilibrium.""" + eq = get("DSHAPE") + + grid = LinearGrid( + L=eq.L_grid, + theta=2 * eq.M_grid, + N=eq.N_grid, + NFP=eq.NFP, + sym=eq.sym, + axis=False, + ) + data_keys = ["current", "|F|", "D_Mercier"] + + data_old = eq.compute(data_keys, grid=grid) + eq = flip_theta(eq) + data_new = eq.compute(data_keys, grid=grid) + + # check that Jacobian and force balance did not change + np.testing.assert_allclose( + data_old["sqrt(g)"].reshape((grid.num_rho, grid.num_theta)), + np.fliplr(data_new["sqrt(g)"].reshape((grid.num_rho, grid.num_theta))), + ) + np.testing.assert_allclose( + data_old["|F|"].reshape((grid.num_rho, grid.num_theta)), + np.fliplr(data_new["|F|"].reshape((grid.num_rho, grid.num_theta))), + rtol=2e-5, + ) + + # check that profiles did not change + np.testing.assert_allclose( + grid.compress(data_old["iota"]), grid.compress(data_new["iota"]) + ) + np.testing.assert_allclose( + grid.compress(data_old["current"]), grid.compress(data_new["current"]) + ) + np.testing.assert_allclose( + grid.compress(data_old["D_Mercier"]), grid.compress(data_new["D_Mercier"]) + ) + + +@pytest.mark.unit +def test_flip_theta_nonaxisym(): + """Test flip_theta on a non-axisymmetric Equilibrium.""" + eq = get("HSX") + + grid = QuadratureGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP) + nodes = grid.nodes.copy() + nodes[:, 1] = np.mod(nodes[:, 1] + np.pi, 2 * np.pi) + grid_flip = Grid(nodes) # grid with flipped theta values + data_keys = ["current", "|F|", "D_Mercier", "f_C"] + + data_old = eq.compute(data_keys, grid=grid, helicity=(1, eq.NFP)) + eq = flip_theta(eq) + data_new = eq.compute(data_keys, grid=grid, helicity=(1, eq.NFP)) + data_flip = eq.compute(data_keys, grid=grid_flip, helicity=(1, eq.NFP)) + + # check that basis vectors did not change + np.testing.assert_allclose(data_old["e_rho"], data_flip["e_rho"], atol=1e-15) + np.testing.assert_allclose(data_old["e_theta"], data_flip["e_theta"], atol=1e-15) + np.testing.assert_allclose(data_old["e^zeta"], data_flip["e^zeta"], atol=1e-15) + + # check that Jacobian is still positive + np.testing.assert_array_less(0, grid.compress(data_new["sqrt(g)"])) + + # check that stability did not change + np.testing.assert_allclose( + grid.compress(data_old["D_Mercier"]), + grid.compress(data_new["D_Mercier"]), + rtol=2e-2, + ) + + # check that the total force balance error on each surface did not change + # (equivalent collocation points now corresond to theta + pi) + np.testing.assert_allclose(data_old["|F|"], data_flip["|F|"], rtol=1e-3) + + # check that the QH helicity did not change + # (equivalent collocation points now corresond to theta + pi) + np.testing.assert_allclose(data_old["f_C"], data_flip["f_C"], atol=1e-8) + + @pytest.mark.unit def test_rescale(): """Test rescale function.""" From 098f7f2379b9f2ea29c365ea27dfdad0b5c7675d Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Wed, 11 Sep 2024 18:57:02 -0600 Subject: [PATCH 2/2] simplify test --- desc/compat.py | 2 +- tests/test_compat.py | 13 ++++++------- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/desc/compat.py b/desc/compat.py index a7e7732e87..38f3d9520a 100644 --- a/desc/compat.py +++ b/desc/compat.py @@ -119,7 +119,7 @@ def flip_theta(eq): Parameters ---------- eq : Equilibrium or iterable of Equilibrium - Equilibria to flip the helicity of. + Equilibria to redefine the poloidal angle of. Returns ------- diff --git a/tests/test_compat.py b/tests/test_compat.py index 68e14748ad..ca1a0c55da 100644 --- a/tests/test_compat.py +++ b/tests/test_compat.py @@ -188,13 +188,12 @@ def test_flip_theta_nonaxisym(): data_old = eq.compute(data_keys, grid=grid, helicity=(1, eq.NFP)) eq = flip_theta(eq) - data_new = eq.compute(data_keys, grid=grid, helicity=(1, eq.NFP)) - data_flip = eq.compute(data_keys, grid=grid_flip, helicity=(1, eq.NFP)) + data_new = eq.compute(data_keys, grid=grid_flip, helicity=(1, eq.NFP)) # check that basis vectors did not change - np.testing.assert_allclose(data_old["e_rho"], data_flip["e_rho"], atol=1e-15) - np.testing.assert_allclose(data_old["e_theta"], data_flip["e_theta"], atol=1e-15) - np.testing.assert_allclose(data_old["e^zeta"], data_flip["e^zeta"], atol=1e-15) + np.testing.assert_allclose(data_old["e_rho"], data_new["e_rho"], atol=1e-15) + np.testing.assert_allclose(data_old["e_theta"], data_new["e_theta"], atol=1e-15) + np.testing.assert_allclose(data_old["e^zeta"], data_new["e^zeta"], atol=1e-15) # check that Jacobian is still positive np.testing.assert_array_less(0, grid.compress(data_new["sqrt(g)"])) @@ -208,11 +207,11 @@ def test_flip_theta_nonaxisym(): # check that the total force balance error on each surface did not change # (equivalent collocation points now corresond to theta + pi) - np.testing.assert_allclose(data_old["|F|"], data_flip["|F|"], rtol=1e-3) + np.testing.assert_allclose(data_old["|F|"], data_new["|F|"], rtol=1e-3) # check that the QH helicity did not change # (equivalent collocation points now corresond to theta + pi) - np.testing.assert_allclose(data_old["f_C"], data_flip["f_C"], atol=1e-8) + np.testing.assert_allclose(data_old["f_C"], data_new["f_C"], atol=1e-8) @pytest.mark.unit