From dd562816abc7802d34312f423d09d9c615a61aea Mon Sep 17 00:00:00 2001 From: unalmis Date: Mon, 12 Aug 2024 01:18:21 -0400 Subject: [PATCH] Undo mod by 2pi on phi --- desc/compute/_core.py | 2 +- tests/test_equilibrium.py | 18 +++++++++++++++--- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/desc/compute/_core.py b/desc/compute/_core.py index 6e5a137ed3..f5ae9a2552 100644 --- a/desc/compute/_core.py +++ b/desc/compute/_core.py @@ -2873,7 +2873,7 @@ def _zeta(params, transforms, profiles, data, **kwargs): data=["zeta", "omega"], ) def _phi(params, transforms, profiles, data, **kwargs): - data["phi"] = (data["zeta"] + data["omega"]) % (2 * jnp.pi) + data["phi"] = data["zeta"] + data["omega"] return data diff --git a/tests/test_equilibrium.py b/tests/test_equilibrium.py index 834775ecae..52958a6a9b 100644 --- a/tests/test_equilibrium.py +++ b/tests/test_equilibrium.py @@ -15,6 +15,7 @@ from desc.io import InputReader from desc.objectives import ForceBalance, ObjectiveFunction, get_equilibrium_objective from desc.profiles import PowerSeriesProfile +from desc.utils import errorif from .utils import area_difference, compute_coords @@ -55,16 +56,27 @@ def test_map_coordinates(): iota = grid.compress(eq.compute("iota", grid=grid)["iota"]) iota = np.broadcast_to(iota, shape=(n,)) - out_1 = eq.map_coordinates(coords, inbasis=["rho", "alpha", "zeta"], iota=iota) + tol = 1e-5 + out_1 = eq.map_coordinates( + coords, inbasis=["rho", "alpha", "zeta"], iota=iota, tol=tol + ) assert np.isfinite(out_1).all() out_2 = eq.map_coordinates( coords, inbasis=["rho", "alpha", "zeta"], period=(np.inf, 2 * np.pi, np.inf), + tol=tol, ) assert np.isfinite(out_2).all() - diff = (out_1 - out_2) % (2 * np.pi) - assert np.all(np.isclose(diff, 0) | np.isclose(np.abs(diff), 2 * np.pi)) + diff = out_1 - out_2 + errorif( + not np.all( + np.isclose(diff, 0, atol=2 * tol) + | np.isclose(np.abs(diff), 2 * np.pi, atol=2 * tol) + ), + AssertionError, + f"diff: {diff}", + ) eq = get("DSHAPE")