diff --git a/desc/coils.py b/desc/coils.py index a8a7c7b05e..9553502dd6 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -8,6 +8,7 @@ from desc.backend import jit, jnp, scan, tree_stack, tree_unstack, vmap from desc.compute import get_params, rpz2xyz, xyz2rpz_vec +from desc.compute.geom_utils import rotation_matrix from desc.geometry import ( FourierPlanarCurve, FourierRZCurve, @@ -847,43 +848,58 @@ def flip(self, *args, **kwargs): """Flip the coils across a plane.""" [coil.flip(*args, **kwargs) for coil in self.coils] - def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): + def compute_magnetic_field(self, nodes, params=None, basis="rpz", grid=None): """Compute magnetic field at a set of points. Parameters ---------- - coords : array-like shape(n,3) or Grid - coordinates to evaluate field at [R,phi,Z] or [x,y,z] + nodes : array-like shape(n,3) or Grid + Nodes to evaluate field at in [R, phi, Z] or [X, Y, Z] coordinates. params : dict or array-like of dict, optional - parameters to pass to curves, either the same for all curves, - or one for each member + Parameters to pass to coils, either the same for all coils or one for each. basis : {"rpz", "xyz"} - basis for input coordinates and returned magnetic field + Basis for input coordinates and returned magnetic field. grid : Grid, int or None or array-like, optional - Grid used to discretize coil, the same for all coils. If an integer, uses - that many equally spaced points. + Grid used to discretize coil, the same for all coils. + If an integer, uses that many equally spaced points. Returns ------- field : ndarray, shape(n,3) - magnetic field at specified points, in either rpz or xyz coordinates + Magnetic field at specified nodes, in [R, phi, Z] or [X, Y, Z] coordinates. """ + if hasattr(nodes, "nodes"): + nodes = nodes.nodes + nodes = jnp.atleast_2d(nodes) if params is None: params = [get_params(["x_s", "x", "s", "ds"], coil) for coil in self] for par, coil in zip(params, self): par["current"] = coil.current - if hasattr(coords, "nodes"): - coords = coords.nodes - coords = jnp.atleast_2d(coords) + assert basis in ["rpz", "xyz"] + + # sum B contributions from each field period + B = jnp.zeros_like(nodes) + for k in range(self.NFP): + if basis == "rpz": + nodes_nfp = jnp.array( + [ + nodes[:, 0], + nodes[:, 1] + 2 * jnp.pi * k / self.NFP, + nodes[:, 2], + ] + ).T + else: + R = rotation_matrix(axis=[0, 0, 1], angle=2 * jnp.pi * k / self.NFP) + nodes_nfp = (R @ nodes.T).T - def body(B, x): - B += self[0].compute_magnetic_field( - coords, params=x, basis=basis, grid=grid - ) - return B, None + def body(B, x): + B += self[0].compute_magnetic_field( + nodes_nfp, params=x, basis=basis, grid=grid + ) + return B, None - B = scan(body, jnp.zeros(coords.shape), tree_stack(params))[0] + B += scan(body, jnp.zeros(nodes_nfp.shape), tree_stack(params))[0] return B diff --git a/tests/test_coils.py b/tests/test_coils.py index f96393b409..d24fd933ea 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -222,12 +222,14 @@ def test_from_symmetry(self): @pytest.mark.unit def test_symmetry_magnetic_field(self): """Tests that compute magnetic field is correct from symmetry.""" - eq = get("precise_QA") + eq = get("precise_QA") # TODO: use precise QH for higher NFP minor_radius = eq.compute("a")["a"] + sym = False # FIXME: make work with stellarator symmetry and change to eq.sym + # initialize CoilSet with symmetry num_coils = 3 # number of unique coils per half field period - grid = LinearGrid(rho=[0.0], M=0, zeta=2 * num_coils, NFP=eq.NFP * (eq.sym + 1)) + grid = LinearGrid(rho=[0.0], M=0, zeta=2 * num_coils, NFP=eq.NFP * (sym + 1)) data_center = eq.axis.compute("x", grid=grid, basis="xyz") data_normal = eq.compute("e^zeta", grid=grid) centers = data_center["x"] @@ -235,19 +237,33 @@ def test_symmetry_magnetic_field(self): coils = [] for k in range(1, 2 * num_coils + 1, 2): coil = FourierPlanarCoil( - current=1, + current=1e6, center=centers[k, :], normal=normals[k, :], r_n=[0, minor_radius + 0.5, 0], ) coils.append(coil) - sym_coilset = CoilSet(coils, NFP=eq.NFP, sym=eq.sym) + sym_coilset = CoilSet(coils, NFP=eq.NFP, sym=sym) # equivalent CoilSet without symmetry - asym_coilset = CoilSet.from_symmetry(sym_coilset, NFP=eq.NFP, sym=eq.sym) - - # TODO: check that sym_coilset and asym_coilset compute same field - assert asym_coilset + asym_coilset = CoilSet.from_symmetry(sym_coilset, NFP=eq.NFP, sym=sym) + + # test that both coil sets compute the same field on the plasma surface + grid = LinearGrid(rho=[1.0], M=eq.M_grid, N=eq.N_grid, NFP=1, sym=False) + with pytest.warns(UserWarning): # because eq.NFP != grid.NFP + data = eq.compute(["phi", "R", "X", "Y", "Z"], grid) + + # test in (R, phi, Z) coordinates + nodes_rpz = np.array([data["R"], data["phi"], data["Z"]]).T + B_sym_rpz = sym_coilset.compute_magnetic_field(nodes_rpz, basis="rpz") + B_asym_rpz = asym_coilset.compute_magnetic_field(nodes_rpz, basis="rpz") + np.testing.assert_allclose(B_sym_rpz, B_asym_rpz, atol=1e-14) + + # test in (X, Y, Z) coordinates + nodes_xyz = np.array([data["X"], data["Y"], data["Z"]]).T + B_sym_xyz = sym_coilset.compute_magnetic_field(nodes_xyz, basis="xyz") + B_asym_xyz = asym_coilset.compute_magnetic_field(nodes_xyz, basis="xyz") + np.testing.assert_allclose(B_sym_xyz, B_asym_xyz, atol=1e-14) @pytest.mark.unit def test_properties(self):