CoilSet working with multiple field periods

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
 

tests/test_coils.py

@@ -222,32 +222,48 @@ 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"]
         normals = rpz2xyz_vec(data_normal["e^zeta"], phi=grid.nodes[:, 2])
         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):