diff --git a/desc/coils.py b/desc/coils.py index 4d0f5f4b72..d82e3cabb0 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -224,7 +224,7 @@ def compute_magnetic_field( current = params.pop("current", self.current) data = self.compute( - ["x", "x_s", "ds"], grid=source_grid, params=params, basis="xyz" + ["phi", "x", "x_s", "ds"], grid=source_grid, params=params, basis="xyz" ) B = biot_savart_quad( coords, data["x"], data["x_s"] * data["ds"][:, None], current diff --git a/desc/compute/_omnigenity.py b/desc/compute/_omnigenity.py index e8d74dd796..4c8ec74731 100644 --- a/desc/compute/_omnigenity.py +++ b/desc/compute/_omnigenity.py @@ -443,15 +443,25 @@ def _omni_map(params, transforms, profiles, data, **kwargs): iota = kwargs.get("iota", 1) # coordinate mapping matrix from (alpha,h) to (theta_B,zeta_B) + # need a bunch of wheres to avoid division by zero causing NaN in backward pass + # this is fine since the incorrect values get ignored later, except in OT or OH + # where fieldlines are exactly parallel to |B| contours, but this is a degenerate + # case of measure 0 so this kludge shouldn't affect things too much. + mat_OP = jnp.array( + [[N, iota / jnp.where(N == 0, 1, N)], [0, 1 / jnp.where(N == 0, 1, N)]] + ) + mat_OT = jnp.array([[0, -1], [M, -1 / jnp.where(iota == 0, 1.0, iota)]]) + den = jnp.where((N - M * iota) == 0, 1.0, (N - M * iota)) + mat_OH = jnp.array([[N, M * iota / den], [M, M / den]]) matrix = jnp.where( M == 0, - jnp.array([N, iota / N, 0, 1 / N]), # OP + mat_OP, jnp.where( N == 0, - jnp.array([0, -1, M, -1 / iota]), # OT - jnp.array([N, M * iota / (N - M * iota), M, M / (N - M * iota)]), # OH + mat_OT, + mat_OH, ), - ).reshape((2, 2)) + ) # solve for (theta_B,zeta_B) corresponding to (eta,alpha) booz = matrix @ jnp.vstack((data["alpha"], data["h"])) diff --git a/desc/compute/utils.py b/desc/compute/utils.py index 545dd123c0..0d7d0a0f5a 100644 --- a/desc/compute/utils.py +++ b/desc/compute/utils.py @@ -9,6 +9,7 @@ from desc.backend import cond, fori_loop, jnp, put from desc.grid import ConcentricGrid, Grid, LinearGrid +from desc.utils import errorif from .data_index import allowed_kwargs, data_index @@ -105,7 +106,12 @@ def compute(parameterization, names, params, transforms, profiles, data=None, ** if name == "x": data[name] = rpz2xyz(data[name]) else: - # TODO: check if phi in data? + errorif( + "phi" not in data.keys(), + ValueError, + "'phi' must be included in the compute data " + + "to convert to 'xyz' basis.", + ) data[name] = rpz2xyz_vec(data[name], phi=data["phi"]) return data diff --git a/desc/objectives/_omnigenity.py b/desc/objectives/_omnigenity.py index a1acb54056..8653d4a770 100644 --- a/desc/objectives/_omnigenity.py +++ b/desc/objectives/_omnigenity.py @@ -666,7 +666,7 @@ def __init__( normalize=True, normalize_target=True, loss_function=None, - deriv_mode="fwd", # FIXME: get it working with rev mode (see GH issue #943) + deriv_mode="auto", eq_grid=None, field_grid=None, M_booz=None, @@ -891,20 +891,26 @@ def compute(self, params_1=None, params_2=None, constants=None): # update theta_B and zeta_B with new iota from the equilibrium M, N = constants["helicity"] iota = jnp.mean(eq_data["iota"]) + # see comment in desc.compute._omnigenity for the explanation of these + # wheres + mat_OP = jnp.array( + [[N, iota / jnp.where(N == 0, 1, N)], [0, 1 / jnp.where(N == 0, 1, N)]] + ) + mat_OT = jnp.array([[0, -1], [M, -1 / jnp.where(iota == 0, 1.0, iota)]]) + den = jnp.where((N - M * iota) == 0, 1.0, (N - M * iota)) + mat_OH = jnp.array([[N, M * iota / den], [M, M / den]]) matrix = jnp.where( M == 0, - jnp.array([N, iota / N, 0, 1 / N]), # OP + mat_OP, jnp.where( N == 0, - jnp.array([0, -1, M, -1 / iota]), # OT - jnp.array( - [N, M * iota / (N - M * iota), M, M / (N - M * iota)] # OH - ), + mat_OT, + mat_OH, ), - ).reshape((2, 2)) + ) booz = matrix @ jnp.vstack((field_data["alpha"], field_data["h"])) - field_data["theta_B"] = booz[0, :] - field_data["zeta_B"] = booz[1, :] + theta_B = booz[0, :] + zeta_B = booz[1, :] else: field_data = compute_fun( "desc.magnetic_fields._core.OmnigenousField", @@ -915,13 +921,15 @@ def compute(self, params_1=None, params_2=None, constants=None): helicity=constants["helicity"], iota=jnp.mean(eq_data["iota"]), ) + theta_B = field_data["theta_B"] + zeta_B = field_data["zeta_B"] # additional computations that cannot be part of the regular compute API nodes = jnp.vstack( ( - jnp.zeros_like(field_data["theta_B"]), - field_data["theta_B"], - field_data["zeta_B"], + jnp.zeros_like(theta_B), + theta_B, + zeta_B, ) ).T B_eta_alpha = jnp.matmul( diff --git a/tests/test_curves.py b/tests/test_curves.py index b0fbdba8dd..92244ee344 100644 --- a/tests/test_curves.py +++ b/tests/test_curves.py @@ -55,9 +55,7 @@ def test_torsion(self): def test_frenet(self): """Test frenet-serret frame of circular curve.""" c = FourierRZCurve() - data = c.compute( - ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="rpz", grid=0 - ) + data = c.compute(["frenet_tangent", "frenet_normal", "frenet_binormal"], grid=0) T, N, B = data["frenet_tangent"], data["frenet_normal"], data["frenet_binormal"] np.testing.assert_allclose(T, np.array([[0, 1, 0]]), atol=1e-12) np.testing.assert_allclose(N, np.array([[-1, 0, 0]]), atol=1e-12) @@ -65,9 +63,7 @@ def test_frenet(self): c.rotate(angle=np.pi) c.flip([0, 1, 0]) c.translate([1, 1, 1]) - data = c.compute( - ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="xyz", grid=0 - ) + data = c.compute(["frenet_tangent", "frenet_normal", "frenet_binormal"], grid=0) T, N, B = data["frenet_tangent"], data["frenet_normal"], data["frenet_binormal"] np.testing.assert_allclose(T, np.array([[0, 1, 0]]), atol=1e-12) np.testing.assert_allclose(N, np.array([[1, 0, 0]]), atol=1e-12) @@ -470,9 +466,9 @@ def test_rotation(self): datax = cx.compute("x", grid=20, basis="xyz") datay = cy.compute("x", grid=20, basis="xyz") dataz = cz.compute("x", grid=20, basis="xyz") - np.testing.assert_allclose(datax["x"][:, 0], 0, atol=1e-16) # only in Y-Z plane - np.testing.assert_allclose(datay["x"][:, 1], 0, atol=1e-16) # only in X-Z plane - np.testing.assert_allclose(dataz["x"][:, 2], 0, atol=1e-16) # only in X-Y plane + np.testing.assert_allclose(datax["x"][:, 0], 0, atol=2e-16) # only in Y-Z plane + np.testing.assert_allclose(datay["x"][:, 1], 0, atol=2e-16) # only in X-Z plane + np.testing.assert_allclose(dataz["x"][:, 2], 0, atol=2e-16) # only in X-Y plane @pytest.mark.unit def test_length(self): diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 962cb4356e..af6eefb251 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -2418,6 +2418,35 @@ def test_objective_no_nangrad_coils(self, objective): g = obj.grad(obj.x()) assert not np.any(np.isnan(g)), str(objective) + @pytest.mark.unit + @pytest.mark.parametrize("helicity", [(1, 0), (1, 1), (0, 1)]) + def test_objective_no_nangrad_omnigenity(self, helicity): + """Omnigenity.""" + surf = FourierRZToroidalSurface.from_qp_model( + major_radius=1, + aspect_ratio=20, + elongation=6, + mirror_ratio=0.2, + torsion=0.1, + NFP=1, + sym=True, + ) + eq = Equilibrium(Psi=6e-3, M=4, N=4, surface=surf) + field = OmnigenousField( + L_B=0, + M_B=2, + L_x=1, + M_x=1, + N_x=1, + NFP=eq.NFP, + helicity=helicity, + B_lm=np.array([0.8, 1.2]), + ) + obj = ObjectiveFunction(Omnigenity(eq=eq, field=field)) + obj.build() + g = obj.grad(obj.x()) + assert not np.any(np.isnan(g)), str(helicity) + @pytest.mark.unit def test_asymmetric_normalization():