diff --git a/desc/coils.py b/desc/coils.py index 4fcb7248b1..214734aa1e 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -6,8 +6,9 @@ import numpy as np -from desc.backend import jit, jnp, scan, tree_stack, tree_unstack, vmap -from desc.compute import get_params, rpz2xyz, xyz2rpz_vec +from desc.backend import fori_loop, jit, jnp, scan, tree_stack, tree_unstack, vmap +from desc.compute import get_params, rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec +from desc.compute.geom_utils import reflection_matrix from desc.geometry import ( FourierPlanarCurve, FourierRZCurve, @@ -17,7 +18,7 @@ from desc.grid import LinearGrid from desc.magnetic_fields import _MagneticField from desc.optimizable import Optimizable, OptimizableCollection, optimizable_parameter -from desc.utils import equals, errorif, flatten_list +from desc.utils import equals, errorif, flatten_list, warnif @jit @@ -146,7 +147,9 @@ def current(self, new): assert jnp.isscalar(new) or new.size == 1 self._current = float(new) - def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): """Compute magnetic field at a set of points. The coil current may be overridden by including `current` @@ -154,13 +157,13 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): Parameters ---------- - coords : array-like shape(n,3) or Grid - coordinates to evaluate field at [R,phi,Z] or [x,y,z] - params : dict, optional - parameters to pass to curve + coords : array-like shape(n,3) + 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 Curve. basis : {"rpz", "xyz"} - basis for input coordinates and returned magnetic field - grid : Grid, int or None + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional Grid used to discretize coil. If an integer, uses that many equally spaced points. Should NOT include endpoint at 2pi. @@ -175,25 +178,27 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): tangents provided by the underlying curve class. Convergence should be exponential in the number of points used to discretize the curve, though curl(B) may not be zero if not fully converged. + """ assert basis.lower() in ["rpz", "xyz"] - if hasattr(coords, "nodes"): - coords = coords.nodes coords = jnp.atleast_2d(coords) - if basis == "rpz": + if basis.lower() == "rpz": + phi = coords[:, 1] coords = rpz2xyz(coords) if params is None: current = self.current else: current = params.pop("current", self.current) - data = self.compute(["x", "x_s", "ds"], grid=grid, params=params, basis="xyz") + data = self.compute( + ["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 ) - if basis == "rpz": - B = xyz2rpz_vec(B, x=coords[:, 0], y=coords[:, 1]) + if basis.lower() == "rpz": + B = xyz2rpz_vec(B, phi=phi) return B def __repr__(self): @@ -563,7 +568,9 @@ def __init__( ): super().__init__(current, X, Y, Z, knots, method, name) - def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): """Compute magnetic field at a set of points. The coil current may be overridden by including `current` @@ -571,13 +578,13 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): Parameters ---------- - coords : array-like shape(n,3) or Grid - coordinates to evaluate field at [R,phi,Z] or [x,y,z] - params : dict, optional - parameters to pass to curve + coords : array-like shape(n,3) + 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 Curve. basis : {"rpz", "xyz"} - basis for input coordinates and returned magnetic field - grid : Grid, int or None + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional Grid used to discretize coil. If an integer, uses that many equally spaced points. Should NOT include endpoint at 2pi. @@ -591,10 +598,9 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): Discretizes the coil into straight segments between grid points, and uses the Hanson-Hirshman expression for exact field from a straight segment. Convergence is approximately quadratic in the number of coil points. + """ assert basis.lower() in ["rpz", "xyz"] - if hasattr(coords, "nodes"): - coords = coords.nodes coords = jnp.atleast_2d(coords) if basis == "rpz": coords = rpz2xyz(coords) @@ -603,7 +609,7 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): else: current = params.pop("current", self.current) - data = self.compute(["x"], grid=grid, params=params, basis="xyz") + data = self.compute(["x"], grid=source_grid, params=params, basis="xyz") # need to make sure the curve is closed. If it's already closed, this doesn't # do anything (effectively just adds a segment of zero length which has no # effect on the overall result) @@ -716,21 +722,30 @@ class CoilSet(OptimizableCollection, _Coil, MutableSequence): Parameters ---------- coils : Coil or array-like of Coils - collection of coils. Must all be the same type and resolution. - currents : float or array-like of float - currents in each coil, or a single current shared by all coils in the set + Collection of coils. Must all be the same type and resolution. + NFP : int (optional) + Number of field periods for enforcing field period symmetry. + If NFP > 1, only include the unique coils in the first field period, + and the magnetic field will be computed assuming 'virtual' coils from the other + field periods. Default = 1. + sym : bool (optional) + Whether to enforce stellarator symmetry. If sym = True, only include the + unique coils in a half field period, and the magnetic field will be computed + assuming 'virtual' coils from the other half field period. Default = False. name : str - name of this CoilSet + Name of this CoilSet. """ - _io_attrs_ = _Coil._io_attrs_ + ["_coils"] + _io_attrs_ = _Coil._io_attrs_ + ["_coils", "_NFP", "_sym"] - def __init__(self, *coils, name=""): + def __init__(self, *coils, NFP=1, sym=False, name=""): coils = flatten_list(coils, flatten_tuple=True) assert all([isinstance(coil, (_Coil)) for coil in coils]) [_check_type(coil, coils[0]) for coil in coils] self._coils = list(coils) + self._NFP = NFP + self._sym = sym self._name = str(name) @property @@ -747,6 +762,16 @@ def coils(self): """list: coils in the coilset.""" return self._coils + @property + def NFP(self): + """int: Number of (toroidal) field periods.""" + return self._NFP + + @property + def sym(self): + """bool: Whether this coil set is stellarator symmetric.""" + return self._sym + @property def current(self): """list: currents in each coil.""" @@ -829,43 +854,80 @@ 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, coords, params=None, basis="rpz", source_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] + coords : array-like shape(n,3) + 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 - 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. + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional + Grid used to discretize coils. If an integer, uses that many equally spaced + points. Should NOT include endpoint at 2pi. 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. + """ + assert basis.lower() in ["rpz", "xyz"] + coords = jnp.atleast_2d(coords) 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) - def body(B, x): - B += self[0].compute_magnetic_field( - coords, params=x, basis=basis, grid=grid + # stellarator symmetry is easiest in [X,Y,Z] coordinates + if basis.lower() == "rpz": + coords_xyz = rpz2xyz(coords) + else: + coords_xyz = coords + + # if stellarator symmetric, add reflected nodes from the other half field period + if self.sym: + normal = jnp.array( + [-jnp.sin(jnp.pi / self.NFP), jnp.cos(jnp.pi / self.NFP), 0] ) - return B, None + coords_sym = ( + coords_xyz + @ reflection_matrix(normal).T + @ reflection_matrix([0, 0, 1]).T + ) + coords_xyz = jnp.vstack((coords_xyz, coords_sym)) + + # field period rotation is easiest in [R,phi,Z] coordinates + coords_rpz = xyz2rpz(coords_xyz) + + # sum the magnetic fields from each field period + def nfp_loop(k, B): + coords_nfp = coords_rpz + jnp.array([0, 2 * jnp.pi * k / self.NFP, 0]) - B = scan(body, jnp.zeros(coords.shape), tree_stack(params))[0] + def body(B, x): + B += self[0].compute_magnetic_field( + coords_nfp, params=x, basis="rpz", source_grid=source_grid + ) + return B, None + + B += scan(body, jnp.zeros(coords_nfp.shape), tree_stack(params))[0] + return B + + B = fori_loop(0, self.NFP, nfp_loop, jnp.zeros_like(coords_rpz)) + + # sum the magnetic fields from both halves of the symmetric field period + if self.sym: + B = B[: coords.shape[0], :] + B[coords.shape[0] :, :] * jnp.array( + [-1, 1, 1] + ) + if basis.lower() == "xyz": + B = rpz2xyz_vec(B, x=coords[:, 0], y=coords[:, 1]) return B @classmethod @@ -938,23 +1000,33 @@ def linspaced_linear( return cls(*coils) @classmethod - def from_symmetry(cls, coils, NFP, sym=False): + def from_symmetry(cls, coils, NFP=1, sym=False): """Create a coil group by reflection and symmetry. Given coils over one field period, repeat coils NFP times between 0 and 2pi to form full coil set. - Or, give coils over 1/2 of a field period, repeat coils 2*NFP times + Or, given coils over 1/2 of a field period, repeat coils 2*NFP times between 0 and 2pi to form full stellarator symmetric coil set. Parameters ---------- coils : Coil, CoilGroup, Coilset - base coil or collection of coils to repeat - NFP : int - number of field periods - sym : bool - whether coils should be stellarator symmetric + Coil or collection of coils in one field period or half field period. + NFP : int (optional) + Number of field periods for enforcing field period symmetry. + The coils will be duplicated NFP times. Default = 1. + sym : bool (optional) + Whether to enforce stellarator symmetry. + If True, the coils will be duplicated 2*NFP times. Default = False. + + Returns + ------- + coilset : CoilSet + A new coil set with NFP=1 and sym=False that is equivalent to the unique + coils with field period symmetry and stellarator symmetry. + The total number of coils in the new coil set is: + len(coilset) = len(coils) * NFP * (int(sym) + 1) """ if not isinstance(coils, CoilSet): @@ -962,12 +1034,27 @@ def from_symmetry(cls, coils, NFP, sym=False): [_check_type(coil, coils[0]) for coil in coils] + # check toroidal extent of coils to be repeated + maxphi = 2 * np.pi / NFP / (sym + 1) + data = coils.compute("phi") + for i, cdata in enumerate(data): + errorif( + np.any(cdata["phi"] > maxphi), + ValueError, + f"coil {i} exceeds the toroidal extent for NFP={NFP} and sym={sym}", + ) + warnif( + sym and np.any(cdata["phi"] < np.finfo(cdata["phi"].dtype).eps), + UserWarning, + f"coil {i} is on the symmetry plane phi=0", + ) + coilset = [] if sym: # first reflect/flip original coilset - # ie, given coils [1,2,3] at angles [0, pi/6, 2pi/6] - # we want a new set like [1,2,3,flip(3),flip(2),flip(1)] - # at [0, pi/6, 2pi/6, 3pi/6, 4pi/6, 5pi/6] + # ie, given coils [1, 2, 3] at angles [pi/6, pi/2, 5pi/6] + # we want a new set like [1, 2, 3, flip(3), flip(2), flip(1)] + # at [pi/6, pi/2, 5pi/6, 7pi/6, 3pi/2, 11pi/6] flipped_coils = [] normal = jnp.array([-jnp.sin(jnp.pi / NFP), jnp.cos(jnp.pi / NFP), 0]) for coil in coils[::-1]: @@ -977,10 +1064,11 @@ def from_symmetry(cls, coils, NFP, sym=False): fcoil.current = -1 * coil.current flipped_coils.append(fcoil) coils = coils + flipped_coils + # next rotate the coilset for each field period for k in range(0, NFP): - coil = coils.copy() - coil.rotate(axis=[0, 0, 1], angle=2 * jnp.pi * k / NFP) - coilset.append(coil) + rotated_coils = coils.copy() + rotated_coils.rotate(axis=[0, 0, 1], angle=2 * jnp.pi * k / NFP) + coilset += rotated_coils return cls(*coilset) @@ -1297,11 +1385,9 @@ class MixedCoilSet(CoilSet): Parameters ---------- coils : Coil or array-like of Coils - collection of coils - currents : float or array-like of float - currents in each coil, or a single current shared by all coils in the set + Collection of coils. name : str - name of this CoilSet + Name of this CoilSet. """ @@ -1309,6 +1395,8 @@ def __init__(self, *coils, name=""): coils = flatten_list(coils, flatten_tuple=True) assert all([isinstance(coil, (_Coil)) for coil in coils]) self._coils = list(coils) + self._NFP = 1 + self._sym = False self._name = str(name) def compute( @@ -1361,33 +1449,34 @@ def compute( ) ] - def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_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] + coords : array-like shape(n,3) + 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 - grid : Grid, int or None or array-like, optional - Grid used to discretize coil, either the same for all coils or one for each - member of the coilset. If an integer, uses that many equally spaced - points. + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional + Grid used to discretize coils. If an integer, uses that many equally spaced + points. Should NOT include endpoint at 2pi. Returns ------- field : ndarray, shape(n,3) magnetic field at specified points, in either rpz or xyz coordinates + """ params = self._make_arraylike(params) - grid = self._make_arraylike(grid) + source_grid = self._make_arraylike(source_grid) B = 0 - for coil, par, grd in zip(self.coils, params, grid): + for coil, par, grd in zip(self.coils, params, source_grid): B += coil.compute_magnetic_field(coords, par, basis, grd) return B diff --git a/desc/compute/geom_utils.py b/desc/compute/geom_utils.py index 15c62cb4ce..fc5e1dab83 100644 --- a/desc/compute/geom_utils.py +++ b/desc/compute/geom_utils.py @@ -65,6 +65,7 @@ def xyz2rpz(pts): ------- pts : ndarray, shape(...,3) points in polar (R,phi,Z) coordinates + """ x, y, z = pts.T r = jnp.sqrt(x**2 + y**2) @@ -84,6 +85,7 @@ def rpz2xyz(pts): ------- pts : ndarray, shape(...,3) points in cartesian (X,Y,Z) coordinates + """ r, p, z = pts.T x = r * jnp.cos(p) @@ -105,6 +107,7 @@ def xyz2rpz_vec(vec, x=None, y=None, phi=None): ------- vec : ndarray, shape(...,3) vectors, in polar (R,phi,Z) form + """ if x is not None and y is not None: phi = jnp.arctan2(y, x) @@ -139,6 +142,7 @@ def rpz2xyz_vec(vec, x=None, y=None, phi=None): ------- vec : ndarray, shape(n,3) vectors, in cartesian (X,Y,Z) form + """ if x is not None and y is not None: phi = jnp.arctan2(y, x) diff --git a/desc/equilibrium/equilibrium.py b/desc/equilibrium/equilibrium.py index 666478384f..e506821e6a 100644 --- a/desc/equilibrium/equilibrium.py +++ b/desc/equilibrium/equilibrium.py @@ -1975,7 +1975,7 @@ def _optimize( # noqa: C901 print("Trust-Region ratio = {:9.3e}".format(tr_ratio[0])) # perturb + solve - (_, predicted_reduction, dc_opt, dc, c_norm, bound_hit,) = optimal_perturb( + (_, predicted_reduction, dc_opt, dc, c_norm, bound_hit) = optimal_perturb( self, constraint, objective, diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 0800a719d1..b9899fce35 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -170,21 +170,22 @@ def __sub__(self, x): return self.__add__(-x) @abstractmethod - def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): """Compute magnetic field at a set of points. Parameters ---------- - coords : array-like shape(N,3) or Grid - cylindrical or cartesian coordinates - params : dict, optional - parameters to pass to scalar potential function + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional + Dictionary of optimizable parameters, eg field.params_dict. basis : {"rpz", "xyz"} - basis for input coordinates and returned magnetic field - grid : Grid, int or None - Grid used to discretize MagneticField object if calculating - B from biot savart. If an integer, uses that many equally spaced - points. + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional + Grid used to discretize MagneticField object if calculating B from + Biot-Savart. Should NOT include endpoint at 2pi. Returns ------- @@ -193,9 +194,9 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): """ - def __call__(self, coords, params=None, basis="rpz"): + def __call__(self, grid, params=None, basis="rpz"): """Compute magnetic field at a set of points.""" - return self.compute_magnetic_field(coords, params, basis) + return self.compute_magnetic_field(grid, params, basis) def compute_Bnormal( self, surface, eval_grid=None, source_grid=None, params=None, basis="rpz" @@ -213,9 +214,8 @@ def compute_Bnormal( the surface poloidal and toroidal resolutions points are in surface angular coordinates i.e theta and zeta source_grid : Grid, int or None - Grid used to discretize MagneticField object if calculating - B from biot savart. If an integer, uses that many equally spaced - points. + Grid used to discretize MagneticField object if calculating B from + Biot-Savart. Should NOT include endpoint at 2pi. params : list or tuple of dict, optional parameters to pass to underlying field's compute_magnetic_field function. If None, uses the default parameters for each field. @@ -244,7 +244,7 @@ def compute_Bnormal( coords = data["x"] surf_normal = data["n_rho"] B = self.compute_magnetic_field( - coords, basis="xyz", grid=source_grid, params=params + coords, basis="xyz", source_grid=source_grid, params=params ) Bnorm = jnp.sum(B * surf_normal, axis=-1) @@ -287,9 +287,8 @@ def save_BNORM_file( the surface poloidal and toroidal resolutions points are in surface angular coordinates i.e theta and zeta source_grid : Grid, int or None - Grid used to discretize MagneticField object if calculating - B from biot savart. If an integer, uses that many equally spaced - points. + Grid used to discretize MagneticField object if calculating B from + Biot-Savart. Should NOT include endpoint at 2pi. params : list or tuple of dict, optional parameters to pass to underlying field's compute_magnetic_field function. If None, uses the default parameters for each field. @@ -433,29 +432,31 @@ def __setattr__(self, name, value): def __hasattr__(self, attr): return hasattr(self, attr) or hasattr(self._field, attr) - def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): """Compute magnetic field at a set of points. Parameters ---------- - coords : array-like shape(N,3) or Grid - cylindrical or cartesian coordinates - params : tuple, optional - parameters to pass to underlying field + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional + Dictionary of optimizable parameters, eg field.params_dict. basis : {"rpz", "xyz"} - basis for input coordinates and returned magnetic field - grid : Grid, int or None - Grid used to discretize MagneticField object if calculating - B from biot savart. If an integer, uses that many equally spaced - points. + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional + Grid used to discretize MagneticField object if calculating B from + Biot-Savart. Should NOT include endpoint at 2pi. Returns ------- field : ndarray, shape(N,3) scaled magnetic field at specified points + """ return self._scale * self._field.compute_magnetic_field( - coords, params, basis, grid + coords, params, basis, source_grid ) @@ -479,47 +480,46 @@ def __init__(self, *fields): ) self._fields = fields - def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): """Compute magnetic field at a set of points. Parameters ---------- - coords : array-like shape(N,3) or Grid - cylindrical or cartesian coordinates - params : list or tuple of dict, optional - parameters to pass to underlying fields. If None, - uses the default parameters for each field. If a list or tuple, should have - one entry for each component field. + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional + Dictionary of optimizable parameters, eg field.params_dict. basis : {"rpz", "xyz"} - basis for input coordinates and returned magnetic field - grid : list or tuple of Grid, int or None, optional - Grid used to discretize MagneticField object if calculating - B from biot savart. If an integer, uses that many equally spaced - points. + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional + Grid used to discretize MagneticField object if calculating B from + Biot-Savart. Should NOT include endpoint at 2pi. Returns ------- field : ndarray, shape(N,3) scaled magnetic field at specified points + """ if params is None: params = [None] * len(self._fields) if isinstance(params, dict): params = [params] - if grid is None: - grid = [None] * len(self._fields) - if not isinstance(grid, (list, tuple)): - grid = [grid] - if len(grid) != len(self._fields): - # ensure that if grid is shorter, that - # it is simply repeated so that zip - # does not terminate early - grid = grid * len(self._fields) + if source_grid is None: + source_grid = [None] * len(self._fields) + if not isinstance(source_grid, (list, tuple)): + source_grid = [source_grid] + if len(source_grid) != len(self._fields): + # ensure that if source_grid is shorter, that it is simply repeated so that + # zip does not terminate early + source_grid = source_grid * len(self._fields) B = 0 - for i, (field, g) in enumerate(zip(self._fields, grid)): + for i, (field, g) in enumerate(zip(self._fields, source_grid)): B += field.compute_magnetic_field( - coords, params[i % len(params)], basis, grid=g + coords, params[i % len(params)], basis, source_grid=g ) return B @@ -594,22 +594,22 @@ def B0(self, new): assert float(new) == new, "B0 must be a scalar" self._B0 = new - def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): """Compute magnetic field at a set of points. Parameters ---------- - coords : array-like shape(N,3) or Grid - cylindrical or cartesian coordinates - params : dict, optional + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional Dict of values for R0 and B0. basis : {"rpz", "xyz"} - basis for input coordinates and returned magnetic field - grid : Grid, int or None - Grid used to discretize MagneticField object if calculating - B from biot savart. If an integer, uses that many equally spaced - points. - Unused by this MagneticField class + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional + Unused by this MagneticField class. + Returns ------- field : ndarray, shape(N,3) @@ -621,8 +621,6 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): R0 = params.get("R0", self.R0) assert basis.lower() in ["rpz", "xyz"] - if hasattr(coords, "nodes"): - coords = coords.nodes coords = jnp.atleast_2d(coords) if basis == "xyz": coords = xyz2rpz(coords) @@ -661,22 +659,21 @@ def B0(self, new): assert float(new) == new, "B0 must be a scalar" self._B0 = new - def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): """Compute magnetic field at a set of points. Parameters ---------- - coords : array-like shape(N,3) or Grid - cylindrical or cartesian coordinates - params : dict, optional - Dict of value for B0. + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional + Dict of values for B0. basis : {"rpz", "xyz"} - basis for input coordinates and returned magnetic field - grid : Grid, int or None - Grid used to discretize MagneticField object if calculating - B from biot savart. If an integer, uses that many equally spaced - points. - Unused by this MagneticField class + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional + Unused by this MagneticField class. Returns ------- @@ -688,8 +685,6 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): B0 = params.get("B0", self.B0) assert basis.lower() in ["rpz", "xyz"] - if hasattr(coords, "nodes"): - coords = coords.nodes coords = jnp.atleast_2d(coords) if basis == "xyz": coords = xyz2rpz(coords) @@ -766,22 +761,21 @@ def iota(self, new): assert float(new) == new, "iota must be a scalar" self._iota = new - def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): """Compute magnetic field at a set of points. Parameters ---------- - coords : array-like shape(N,3) or Grid - cylindrical or cartesian coordinates - params : dict, optional - Dict of values for B0, R0, iota. + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional + Dict of values for R0, B0, and iota. basis : {"rpz", "xyz"} - basis for input coordinates and returned magnetic field - grid : Grid, int or None - Grid used to discretize MagneticField object if calculating - B from biot savart. If an integer, uses that many equally spaced - points. - Unused by this MagneticField class + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional + Unused by this MagneticField class. Returns ------- @@ -795,8 +789,6 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): iota = params.get("iota", self.iota) assert basis.lower() in ["rpz", "xyz"] - if hasattr(coords, "nodes"): - coords = coords.nodes coords = jnp.atleast_2d(coords) if basis == "xyz": coords = xyz2rpz(coords) @@ -937,22 +929,21 @@ def _approx_derivs(self, Bi): tempdict[key] = val[:, 0, :] return tempdict - def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): """Compute magnetic field at a set of points. Parameters ---------- - coords : array-like shape(N,3) or Grid - cylindrical or cartesian coordinates - params : dict, optional + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional Dictionary of optimizable parameters, eg field.params_dict. basis : {"rpz", "xyz"} - basis for input coordinates and returned magnetic field - grid : Grid, int or None - Grid used to discretize MagneticField object if calculating - B from biot savart. If an integer, uses that many equally spaced - points. - Unused by this MagneticField class + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None + Unused by this MagneticField class. Returns ------- @@ -962,8 +953,6 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): """ assert basis.lower() in ["rpz", "xyz"] currents = self.currents if params is None else params["currents"] - if hasattr(coords, "nodes"): - coords = coords.nodes coords = jnp.atleast_2d(coords) if basis == "xyz": coords = xyz2rpz(coords) @@ -1190,22 +1179,21 @@ def __init__(self, potential, params=None): self._potential = potential self._params = params - def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): """Compute magnetic field at a set of points. Parameters ---------- - coords : array-like shape(N,3) or Grid - cylindrical or cartesian coordinates - params : dict, optional - parameters to pass to scalar potential function + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional + Dictionary of optimizable parameters, eg field.params_dict. basis : {"rpz", "xyz"} - basis for input coordinates and returned magnetic field - grid : Grid, int or None - Grid used to discretize MagneticField object if calculating - B from biot savart. If an integer, uses that many equally spaced - points. - Unused by this MagneticField class + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None + Unused by this MagneticField class. Returns ------- @@ -1214,8 +1202,6 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): """ assert basis.lower() in ["rpz", "xyz"] - if hasattr(coords, "nodes"): - coords = coords.nodes coords = jnp.atleast_2d(coords) if basis == "xyz": coords = xyz2rpz(coords) @@ -1237,7 +1223,15 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): def field_line_integrate( - r0, z0, phis, field, params=None, grid=None, rtol=1e-8, atol=1e-8, maxstep=1000 + r0, + z0, + phis, + field, + params=None, + source_grid=None, + rtol=1e-8, + atol=1e-8, + maxstep=1000, ): """Trace field lines by integration. @@ -1253,7 +1247,7 @@ def field_line_integrate( source of magnetic field to integrate params: dict parameters passed to field - grid : Grid, optional + source_grid : Grid, optional Collocation points used to discretize source field. rtol, atol : float relative and absolute tolerances for ode integration @@ -1277,7 +1271,9 @@ def field_line_integrate( def odefun(rpz, s): rpz = rpz.reshape((3, -1)).T r = rpz[:, 0] - br, bp, bz = field.compute_magnetic_field(rpz, params, basis="rpz", grid=grid).T + br, bp, bz = field.compute_magnetic_field( + rpz, params, basis="rpz", source_grid=source_grid + ).T return jnp.array( [r * br / bp * jnp.sign(bp), jnp.sign(bp), r * bz / bp * jnp.sign(bp)] ).squeeze() @@ -1446,20 +1442,21 @@ def save(self, file_name, file_format=None, file_mode="w"): " as the potential function cannot be serialized." ) - def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): """Compute magnetic field at a set of points. Parameters ---------- - coords : array-like shape(N,3) or Grid - cylindrical or cartesian coordinates - params : dict, optional - parameters to pass to compute function - should include the potential + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional + Dictionary of optimizable parameters, eg field.params_dict. basis : {"rpz", "xyz"} - basis for input coordinates and returned magnetic field - grid : Grid, - source grid upon which to evaluate the surface current density K + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional + Source grid upon which to evaluate the surface current density K. Returns ------- @@ -1468,7 +1465,11 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): """ return _compute_magnetic_field_from_CurrentPotentialField( - field=self, coords=coords, params=params, basis=basis, grid=grid + field=self, + coords=coords, + params=params, + basis=basis, + source_grid=source_grid, ) @classmethod @@ -1746,20 +1747,21 @@ def change_Phi_resolution(self, M=None, N=None, NFP=None, sym_Phi=None): NFP=NFP ) # make sure surface and Phi basis NFP are the same - def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): """Compute magnetic field at a set of points. Parameters ---------- - coords : array-like shape(N,3) or Grid - cylindrical or cartesian coordinates - params : dict, optional - parameters to pass to compute function - should include the potential + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional + Dictionary of optimizable parameters, eg field.params_dict. basis : {"rpz", "xyz"} - basis for input coordinates and returned magnetic field - grid : Grid, - grid upon which to evaluate the surface current density K + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional + Source grid upon which to evaluate the surface current density K. Returns ------- @@ -1767,11 +1769,15 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): magnetic field at specified points """ - grid = grid or LinearGrid( + coords = coords or LinearGrid( M=self._M_Phi * 4 + 1, N=self._N_Phi * 4 + 1, NFP=self.NFP ) return _compute_magnetic_field_from_CurrentPotentialField( - field=self, coords=coords, params=params, basis=basis, grid=grid + field=self, + coords=coords, + params=params, + basis=basis, + source_grid=source_grid, ) @classmethod @@ -1847,7 +1853,7 @@ def from_surface( def _compute_magnetic_field_from_CurrentPotentialField( - field, coords, params=None, basis="rpz", grid=None + field, coords, params=None, basis="rpz", source_grid=None ): """Compute magnetic field at a set of points. @@ -1855,14 +1861,14 @@ def _compute_magnetic_field_from_CurrentPotentialField( ---------- field : CurrentPotentialField or FourierCurrentPotentialField current potential field object from which to compute magnetic field. - coords : array-like shape(N,3) or Grid + coords : array-like shape(N,3) cylindrical or cartesian coordinates params : dict, optional parameters to pass to compute function should include the potential basis : {"rpz", "xyz"} basis for input coordinates and returned magnetic field - grid : Grid, + source_grid : Grid, source grid upon which to evaluate the surface current density K Returns @@ -1872,12 +1878,10 @@ def _compute_magnetic_field_from_CurrentPotentialField( """ assert basis.lower() in ["rpz", "xyz"] - if hasattr(coords, "nodes"): - coords = coords.nodes coords = jnp.atleast_2d(coords) if basis == "rpz": coords = rpz2xyz(coords) - surface_grid = grid or LinearGrid(M=30, N=30, NFP=field.NFP) + surface_grid = source_grid or LinearGrid(M=30, N=30, NFP=field.NFP) # compute surface current, and store grid quantities # needed for integration in class diff --git a/tests/test_coils.py b/tests/test_coils.py index e9393f045a..6b6dc8df50 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -13,8 +13,10 @@ MixedCoilSet, SplineXYZCoil, ) +from desc.compute import rpz2xyz_vec, xyz2rpz, xyz2rpz_vec +from desc.examples import get from desc.geometry import FourierRZCurve, FourierRZToroidalSurface -from desc.grid import Grid, LinearGrid +from desc.grid import LinearGrid from desc.magnetic_fields import SumMagneticField, VerticalMagneticField @@ -24,51 +26,106 @@ class TestCoil: @pytest.mark.unit def test_biot_savart_all_coils(self): """Test biot-savart implementation against analytic formula.""" + coil_grid = LinearGrid(zeta=100, endpoint=False) + R = 2 y = 1 I = 1e7 + By_true = 1e-7 * 2 * np.pi * R**2 * I / (y**2 + R**2) ** (3 / 2) - B_true = np.array([0, By_true, 0]) + Bz_true = 1e-7 * 2 * np.pi * R**2 * I / (y**2 + R**2) ** (3 / 2) + + B_true_xyz = np.atleast_2d([0, By_true, 0]) + grid_xyz = np.atleast_2d([10, y, 0]) + grid_rpz = xyz2rpz(grid_xyz) + B_true_rpz_xy = xyz2rpz_vec(B_true_xyz, x=grid_xyz[:, 0], y=grid_xyz[:, 1]) + B_true_rpz_phi = xyz2rpz_vec(B_true_xyz, phi=grid_rpz[:, 1]) # FourierXYZCoil coil = FourierXYZCoil(I) - grid = LinearGrid(zeta=100, endpoint=False) - B_approx = coil.compute_magnetic_field( - Grid([[10, y, 0], [10, -y, 0]]), basis="xyz", grid=grid - )[0] + B_xyz = coil.compute_magnetic_field( + grid_xyz, basis="xyz", source_grid=coil_grid + ) + B_rpz = coil.compute_magnetic_field( + grid_rpz, basis="rpz", source_grid=coil_grid + ) np.testing.assert_allclose( - B_true, B_approx, rtol=1e-3, atol=1e-10, err_msg="Using FourierXYZCoil" + B_true_xyz, B_xyz, rtol=1e-3, atol=1e-10, err_msg="Using FourierXYZCoil" + ) + np.testing.assert_allclose( + B_true_rpz_xy, B_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierXYZCoil" + ) + np.testing.assert_allclose( + B_true_rpz_phi, B_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierXYZCoil" ) # SplineXYZCoil - coords = coil.compute("x", grid=grid, basis="xyz")["x"] - coil = SplineXYZCoil(I, X=coords[:, 0], Y=coords[:, 1], Z=coords[:, 2]) - B_approx = coil.compute_magnetic_field( - Grid([[10, y, 0], [10, -y, 0]]), basis="xyz", grid=grid - )[0] + x = coil.compute("x", grid=coil_grid, basis="xyz")["x"] + coil = SplineXYZCoil(I, X=x[:, 0], Y=x[:, 1], Z=x[:, 2]) + B_xyz = coil.compute_magnetic_field( + grid_xyz, basis="xyz", source_grid=coil_grid + ) + B_rpz = coil.compute_magnetic_field( + grid_rpz, basis="rpz", source_grid=coil_grid + ) np.testing.assert_allclose( - B_true, B_approx, rtol=1e-3, atol=1e-10, err_msg="Using SplineXYZCoil" + B_true_xyz, B_xyz, rtol=1e-3, atol=1e-10, err_msg="Using SplineXYZCoil" + ) + np.testing.assert_allclose( + B_true_rpz_xy, B_rpz, rtol=1e-3, atol=1e-10, err_msg="Using SplineXYZCoil" + ) + np.testing.assert_allclose( + B_true_rpz_phi, B_rpz, rtol=1e-3, atol=1e-10, err_msg="Using SplineXYZCoil" ) # FourierPlanarCoil coil = FourierPlanarCoil(I) - grid = LinearGrid(zeta=100, endpoint=False) - B_approx = coil.compute_magnetic_field( - Grid([[10, y, 0], [10, -y, 0]]), basis="xyz", grid=grid - )[0] + B_xyz = coil.compute_magnetic_field( + grid_xyz, basis="xyz", source_grid=coil_grid + ) + B_rpz = coil.compute_magnetic_field( + grid_rpz, basis="rpz", source_grid=coil_grid + ) np.testing.assert_allclose( - B_true, B_approx, rtol=1e-3, atol=1e-10, err_msg="Using FourierPlanarCoil" + B_true_xyz, B_xyz, rtol=1e-3, atol=1e-10, err_msg="Using FourierPlanarCoil" + ) + np.testing.assert_allclose( + B_true_rpz_xy, + B_rpz, + rtol=1e-3, + atol=1e-10, + err_msg="Using FourierPlanarCoil", + ) + np.testing.assert_allclose( + B_true_rpz_phi, + B_rpz, + rtol=1e-3, + atol=1e-10, + err_msg="Using FourierPlanarCoil", ) + B_true_xyz = np.atleast_2d([0, 0, Bz_true]) + grid_xyz = np.atleast_2d([0, 0, y]) + grid_rpz = xyz2rpz(grid_xyz) + B_true_rpz_xy = xyz2rpz_vec(B_true_xyz, x=grid_xyz[:, 0], y=grid_xyz[:, 1]) + B_true_rpz_phi = xyz2rpz_vec(B_true_xyz, phi=grid_rpz[:, 1]) + # FourierRZCoil - Bz_true = 1e-7 * 2 * np.pi * R**2 * I / (y**2 + R**2) ** (3 / 2) - B_true = np.array([0, 0, Bz_true]) coil = FourierRZCoil(I, R_n=np.array([R]), modes_R=np.array([0])) - B_approx = coil.compute_magnetic_field( - Grid([[0, 0, y], [0, 0, -y]]), basis="xyz", grid=grid - )[0] + B_xyz = coil.compute_magnetic_field( + grid_xyz, basis="xyz", source_grid=coil_grid + ) + B_rpz = coil.compute_magnetic_field( + grid_rpz, basis="rpz", source_grid=coil_grid + ) + np.testing.assert_allclose( + B_true_xyz, B_xyz, rtol=1e-3, atol=1e-10, err_msg="Using FourierRZCoil" + ) + np.testing.assert_allclose( + B_true_rpz_xy, B_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierRZCoil" + ) np.testing.assert_allclose( - B_true, B_approx, rtol=1e-3, atol=1e-10, err_msg="Using FourierRZCoil" + B_true_rpz_phi, B_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierRZCoil" ) @pytest.mark.unit @@ -94,7 +151,7 @@ def test_SumMagneticField_with_Coil(self): field = SumMagneticField(coil, VerticalMagneticField(B_Z)) B_approx = field.compute_magnetic_field( - Grid([[10, y, 0], [10, -y, 0]]), basis="xyz", grid=100 + np.array([[10, y, 0], [10, -y, 0]]), basis="xyz", source_grid=100 )[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) @@ -121,7 +178,7 @@ def test_adding_MagneticField_with_Coil_or_CoilSet(self): for i, field in enumerate([field1, field2, field3, field4, field5, field6]): B_approx = field.compute_magnetic_field( - Grid([[10, y, 0], [10, -y, 0]]), basis="xyz", grid=100 + np.array([[10, y, 0], [10, -y, 0]]), basis="xyz", source_grid=100 )[0] np.testing.assert_allclose( B_true, B_approx, rtol=1e-3, atol=1e-10, err_msg=f"field {i}" @@ -138,9 +195,15 @@ def test_converting_coil_types(self): x1 = coil1.compute("x", grid=grid, basis="xyz")["x"] x2 = coil2.compute("x", grid=grid, basis="xyz")["x"] x3 = coil3.compute("x", grid=grid, basis="xyz")["x"] - B1 = coil1.compute_magnetic_field(np.zeros((1, 3)), grid=grid, basis="xyz") - B2 = coil2.compute_magnetic_field(np.zeros((1, 3)), grid=grid, basis="xyz") - B3 = coil3.compute_magnetic_field(np.zeros((1, 3)), grid=grid, basis="xyz") + B1 = coil1.compute_magnetic_field( + np.zeros((1, 3)), source_grid=grid, basis="xyz" + ) + B2 = coil2.compute_magnetic_field( + np.zeros((1, 3)), source_grid=grid, basis="xyz" + ) + B3 = coil3.compute_magnetic_field( + np.zeros((1, 3)), source_grid=grid, basis="xyz" + ) np.testing.assert_allclose(x1, x2, atol=1e-12) np.testing.assert_allclose(x1, x3, atol=1e-12) np.testing.assert_allclose(B1, B2, rtol=1e-8, atol=1e-8) @@ -165,7 +228,9 @@ def test_linspaced_linear(self): ) coils.current = I np.testing.assert_allclose(coils.current, I) - B_approx = coils.compute_magnetic_field([0, 0, z[-1]], basis="xyz", grid=32)[0] + B_approx = coils.compute_magnetic_field( + [0, 0, z[-1]], basis="xyz", source_grid=32 + )[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) @pytest.mark.unit @@ -179,7 +244,9 @@ def test_linspaced_angular(self): coil = FourierPlanarCoil() coil.current = I coils = CoilSet.linspaced_angular(coil, n=N) - B_approx = coils.compute_magnetic_field([10, 0, 0], basis="rpz", grid=32)[0] + B_approx = coils.compute_magnetic_field( + [10, 0, 0], basis="rpz", source_grid=32 + )[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) surf = FourierRZToroidalSurface( @@ -203,7 +270,9 @@ def test_from_symmetry(self): coil = FourierPlanarCoil(I) coils = CoilSet.linspaced_angular(coil, angle=np.pi / 2, n=N // 4) coils = MixedCoilSet.from_symmetry(coils, NFP=4) - B_approx = coils.compute_magnetic_field([10, 0, 0], basis="rpz", grid=32)[0] + B_approx = coils.compute_magnetic_field( + [10, 0, 0], basis="rpz", source_grid=32 + )[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) # with stellarator symmetry @@ -214,9 +283,56 @@ def test_from_symmetry(self): coil, I, [0, 0, 1], np.pi / NFP, N // NFP // 2 ) coils2 = MixedCoilSet.from_symmetry(coils, NFP, True) - B_approx = coils2.compute_magnetic_field([10, 0, 0], basis="rpz", grid=32)[0] + B_approx = coils2.compute_magnetic_field( + [10, 0, 0], basis="rpz", source_grid=32 + )[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) + @pytest.mark.unit + def test_symmetry_magnetic_field(self): + """Tests that compute magnetic field is correct from symmetry.""" + eq = get("precise_QH") + minor_radius = eq.compute("a")["a"] + + # 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)) + with pytest.warns(UserWarning): # because eq.NFP != grid.NFP + 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=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) + + # equivalent CoilSet without symmetry + asym_coilset = CoilSet.from_symmetry(sym_coilset, NFP=eq.NFP, sym=eq.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): """Test getting/setting of CoilSet attributes.""" @@ -372,7 +488,7 @@ def test_dunder_methods(self): del coils1[-2] assert len(coils1) == 4 assert coils1[-1] is coil2 - assert coils1[-2][0].__class__ is coil1.__class__ + assert coils1[-2].__class__ is coil1.__class__ coils2 = CoilSet.linspaced_angular(coil1) assert coils2[0].eq(coil1) and not (coils2[0] is coil1) @@ -399,8 +515,8 @@ def test_coilset_convert(self): np.testing.assert_allclose( [xi["x"] for xi in x1], [xi["x"] for xi in x2], atol=1e-12 ) - B1 = coils1.compute_magnetic_field(np.array([[10, 2, 1]]), grid=grid) - B2 = coils2.compute_magnetic_field(np.array([[10, 2, 1]]), grid=grid) + B1 = coils1.compute_magnetic_field(np.array([[10, 2, 1]]), source_grid=grid) + B2 = coils2.compute_magnetic_field(np.array([[10, 2, 1]]), source_grid=grid) np.testing.assert_allclose(B1, B2, rtol=1e-2) coils3 = CoilSet.linspaced_angular(coil2, n=12) @@ -412,8 +528,8 @@ def test_coilset_convert(self): np.testing.assert_allclose( [xi["x"] for xi in x3], [xi["x"] for xi in x4], atol=1e-12 ) - B3 = coils3.compute_magnetic_field(np.array([[10, 2, 1]]), grid=grid) - B4 = coils4.compute_magnetic_field(np.array([[10, 2, 1]]), grid=grid) + B3 = coils3.compute_magnetic_field(np.array([[10, 2, 1]]), source_grid=grid) + B4 = coils4.compute_magnetic_field(np.array([[10, 2, 1]]), source_grid=grid) np.testing.assert_allclose(B3, B4, rtol=1e-2) @@ -461,9 +577,11 @@ def test_load_and_save_makegrid_coils(tmpdir_factory): # check magnetic field from both, check that matches grid = LinearGrid(N=200, endpoint=False) - B1 = coilset.compute_magnetic_field(np.array([[0.7, 0, 0]]), basis="xyz", grid=grid) + B1 = coilset.compute_magnetic_field( + np.array([[0.7, 0, 0]]), basis="xyz", source_grid=grid + ) B2 = coilset2.compute_magnetic_field( - np.array([[0.7, 0, 0]]), basis="xyz", grid=grid + np.array([[0.7, 0, 0]]), basis="xyz", source_grid=grid ) np.testing.assert_allclose(B1, B2, atol=1e-7) @@ -536,8 +654,12 @@ def test_save_and_load_makegrid_coils_rotated(tmpdir_factory): np.testing.assert_allclose(B_normal2, 0, atol=1e-16) # check B btwn the two coilsets - B1 = coilset.compute_magnetic_field(np.array([[10, 0, 0]]), basis="xyz", grid=32) - B2 = coilset2.compute_magnetic_field(np.array([[10, 0, 0]]), basis="xyz", grid=1000) + B1 = coilset.compute_magnetic_field( + np.array([[10, 0, 0]]), basis="xyz", source_grid=32 + ) + B2 = coilset2.compute_magnetic_field( + np.array([[10, 0, 0]]), basis="xyz", source_grid=1000 + ) # coilset uses fourier discretization so biot savart is more accurate # coilset2 uses hanson hirshman which is only 2nd order @@ -611,8 +733,12 @@ def test_save_and_load_makegrid_coils_rotated_int_grid(tmpdir_factory): np.testing.assert_allclose(B_normal2, 0, atol=1e-16) # check B btwn the two coilsets - B1 = coilset.compute_magnetic_field(np.array([[10, 0, 0]]), basis="xyz", grid=grid) - B2 = coilset2.compute_magnetic_field(np.array([[10, 0, 0]]), basis="xyz", grid=grid) + B1 = coilset.compute_magnetic_field( + np.array([[10, 0, 0]]), basis="xyz", source_grid=grid + ) + B2 = coilset2.compute_magnetic_field( + np.array([[10, 0, 0]]), basis="xyz", source_grid=grid + ) np.testing.assert_allclose(B1, B2, atol=1e-10) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 11c7af655a..b67e58ebba 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -255,7 +255,8 @@ def test_current_potential_field(self): np.testing.assert_allclose( sumfield.compute_magnetic_field( - [10.0, 0, 0], grid=[None, LinearGrid(M=30, N=30, NFP=surface.NFP)] + [10.0, 0, 0], + source_grid=[None, LinearGrid(M=30, N=30, NFP=surface.NFP)], ), correct_field(10.0, 0, 0) + B_TF([10.0, 0, 0]), atol=1e-16, @@ -313,13 +314,15 @@ def test_fourier_current_potential_field(self): field.change_Phi_resolution(2, 2) np.testing.assert_allclose( - field.compute_magnetic_field([10.0, 0, 0], grid=surface_grid), + field.compute_magnetic_field([10.0, 0, 0], source_grid=surface_grid), correct_field(10.0, 0, 0), atol=1e-16, rtol=1e-8, ) np.testing.assert_allclose( - field.compute_magnetic_field([10.0, np.pi / 4, 0], grid=surface_grid), + field.compute_magnetic_field( + [10.0, np.pi / 4, 0], source_grid=surface_grid + ), correct_field(10.0, np.pi / 4, 0), atol=1e-16, rtol=1e-8, @@ -329,14 +332,14 @@ def test_fourier_current_potential_field(self): field.I = 0 np.testing.assert_allclose( - field.compute_magnetic_field([10.0, 0, 0], grid=surface_grid), + field.compute_magnetic_field([10.0, 0, 0], source_grid=surface_grid), correct_field(10.0, 0, 0) * 2, atol=1e-16, rtol=1e-8, ) # use default grid np.testing.assert_allclose( - field.compute_magnetic_field([10.0, np.pi / 4, 0], grid=None), + field.compute_magnetic_field([10.0, np.pi / 4, 0], source_grid=None), correct_field(10.0, np.pi / 4, 0) * 2, atol=1e-12, rtol=1e-8, @@ -351,13 +354,15 @@ def test_fourier_current_potential_field(self): ) np.testing.assert_allclose( - field.compute_magnetic_field([10.0, 0, 0], grid=surface_grid), + field.compute_magnetic_field([10.0, 0, 0], source_grid=surface_grid), correct_field(10.0, 0, 0), atol=1e-16, rtol=1e-8, ) np.testing.assert_allclose( - field.compute_magnetic_field([10.0, np.pi / 4, 0], grid=surface_grid), + field.compute_magnetic_field( + [10.0, np.pi / 4, 0], source_grid=surface_grid + ), correct_field(10.0, np.pi / 4, 0), atol=1e-16, rtol=1e-8, @@ -719,7 +724,7 @@ def test_Bnormal_save_and_load_DSHAPE(self, tmpdir_factory): @pytest.mark.unit def test_Bnormal_save_and_load_HELIOTRON(self, tmpdir_factory): """Tests Bnormal save/load for simple toroidal field with HELIOTRON.""" - ### test on simple field with stellarator + # test on simple field with stellarator tmpdir = tmpdir_factory.mktemp("BNORM_files") path = tmpdir.join("BNORM_desc_heliotron.txt") tfield = ToroidalMagneticField(2, 1)