diff --git a/CHANGELOG.md b/CHANGELOG.md index 85ab3a47a..ccf156230 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,13 @@ Changelog New Features - Add ``use_signed_distance`` flag to ``PlasmaVesselDistance`` which will use a signed distance as the target, which is positive when the plasma is inside of the vessel surface and negative if the plasma is outside of the vessel surface, to allow optimizer to distinguish if the equilbrium surface exits the vessel surface and guard against it by targeting a positive signed distance. +- Add ``VectorPotentialField`` class to allow calculation of magnetic fields from a user-specified + vector potential function. +- Add ``compute_magnetic_vector_potential`` methods to most ``MagneticField`` objects to allow vector potential + computation. +- Add ability to save and load vector potential information from ``mgrid`` files. +- Changes ``ToroidalFlux`` objective to default using a 1D loop integral of the vector potential +to compute the toroidal flux when possible, as opposed to a 2D surface integral of the magnetic field dotted with ``n_zeta``. - Allow specification of Nyquist spectrum maximum modenumbers when using ``VMECIO.save`` to save a DESC .h5 file as a VMEC-format wout file Bug Fixes diff --git a/desc/coils.py b/desc/coils.py index f18436591..2bfe3d9fa 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -82,6 +82,53 @@ def biot_savart_hh(eval_pts, coil_pts_start, coil_pts_end, current): return B +@jit +def biot_savart_vector_potential_hh(eval_pts, coil_pts_start, coil_pts_end, current): + """Biot-Savart law for vector potential for filamentary coils following [1]. + + The coil is approximated by a series of straight line segments + and an analytic expression is used to evaluate the vector potential from each + segment. This expression assumes the Coulomb gauge. + + Parameters + ---------- + eval_pts : array-like shape(n,3) + Evaluation points in cartesian coordinates + coil_pts_start, coil_pts_end : array-like shape(m,3) + Points in cartesian space defining the start and end of each segment. + Should be a closed curve, such that coil_pts_start[0] == coil_pts_end[-1] + though this is not checked. + current : float + Current through the coil (in Amps). + + Returns + ------- + A : ndarray, shape(n,3) + Magnetic vector potential in cartesian components at specified points + + [1] Hanson & Hirshman, "Compact expressions for the Biot-Savart + fields of a filamentary segment" (2002) + """ + d_vec = coil_pts_end - coil_pts_start + L = jnp.linalg.norm(d_vec, axis=-1) + d_vec_over_L = ((1 / L) * d_vec.T).T + + Ri_vec = eval_pts[jnp.newaxis, :] - coil_pts_start[:, jnp.newaxis, :] + Ri = jnp.linalg.norm(Ri_vec, axis=-1) + Rf = jnp.linalg.norm( + eval_pts[jnp.newaxis, :] - coil_pts_end[:, jnp.newaxis, :], axis=-1 + ) + Ri_p_Rf = Ri + Rf + + eps = L[:, jnp.newaxis] / (Ri_p_Rf) + + A_mag = 1.0e-7 * current * jnp.log((1 + eps) / (1 - eps)) # 1.0e-7 == mu_0/(4 pi) + + # Now just need to multiply by e^ = d_vec/L = (x_f - x_i)/L + A = jnp.sum(A_mag[:, :, jnp.newaxis] * d_vec_over_L[:, jnp.newaxis, :], axis=0) + return A + + @jit def biot_savart_quad(eval_pts, coil_pts, tangents, current): """Biot-Savart law for filamentary coil using numerical quadrature. @@ -123,6 +170,42 @@ def biot_savart_quad(eval_pts, coil_pts, tangents, current): return B +@jit +def biot_savart_vector_potential_quad(eval_pts, coil_pts, tangents, current): + """Biot-Savart law (for A) for filamentary coil using numerical quadrature. + + This expression assumes the Coulomb gauge. + + Parameters + ---------- + eval_pts : array-like shape(n,3) + Evaluation points in cartesian coordinates + coil_pts : array-like shape(m,3) + Points in cartesian space defining coil + tangents : array-like, shape(m,3) + Tangent vectors to the coil at coil_pts. If the curve is given + by x(s) with curve parameter s, coil_pts = x, tangents = dx/ds*ds where + ds is the spacing between points. + current : float + Current through the coil (in Amps). + + Returns + ------- + A : ndarray, shape(n,3) + Magnetic vector potential in cartesian components at specified points. + """ + dl = tangents + R_vec = eval_pts[jnp.newaxis, :] - coil_pts[:, jnp.newaxis, :] + R_mag = jnp.linalg.norm(R_vec, axis=-1) + + vec = dl[:, jnp.newaxis, :] + denom = R_mag + + # 1e-7 == mu_0/(4 pi) + A = jnp.sum(1.0e-7 * current * vec / denom[:, :, None], axis=0) + return A + + class _Coil(_MagneticField, Optimizable, ABC): """Base class representing a magnetic field coil. @@ -187,10 +270,16 @@ def _compute_position(self, params=None, grid=None, **kwargs): x = x.at[:, :, 1].set(jnp.mod(x[:, :, 1], 2 * jnp.pi)) return x - def compute_magnetic_field( - self, coords, params=None, basis="rpz", source_grid=None, transforms=None + def _compute_A_or_B( + self, + coords, + params=None, + basis="rpz", + source_grid=None, + transforms=None, + compute_A_or_B="B", ): - """Compute magnetic field at a set of points. + """Compute magnetic field or vector potential at a set of points. The coil current may be overridden by including `current` in the `params` dictionary. @@ -208,6 +297,9 @@ def compute_magnetic_field( points. Should NOT include endpoint at 2pi. transforms : dict of Transform or array-like Transforms for R, Z, lambda, etc. Default is to build from grid. + compute_A_or_B: {"A", "B"}, optional + whether to compute the magnetic vector potential "A" or the magnetic field + "B". Defaults to "B" Returns @@ -223,6 +315,14 @@ def compute_magnetic_field( may not be zero if not fully converged. """ + errorif( + compute_A_or_B not in ["A", "B"], + ValueError, + f'Expected "A" or "B" for compute_A_or_B, instead got {compute_A_or_B}', + ) + op = {"B": biot_savart_quad, "A": biot_savart_vector_potential_quad}[ + compute_A_or_B + ] assert basis.lower() in ["rpz", "xyz"] coords = jnp.atleast_2d(jnp.asarray(coords)) if basis.lower() == "rpz": @@ -256,13 +356,87 @@ def compute_magnetic_field( data["x_s"] = rpz2xyz_vec(data["x_s"], phi=data["x"][:, 1]) data["x"] = rpz2xyz(data["x"]) - B = biot_savart_quad( - coords, data["x"], data["x_s"] * data["ds"][:, None], current - ) + AB = op(coords, data["x"], data["x_s"] * data["ds"][:, None], current) if basis.lower() == "rpz": - B = xyz2rpz_vec(B, phi=phi) - return B + AB = xyz2rpz_vec(AB, phi=phi) + return AB + + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic field at a set of points. + + The coil current may be overridden by including `current` + in the `params` dictionary. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict, optional + Parameters to pass to Curve. + basis : {"rpz", "xyz"} + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None, optional + Grid used to discretize coil. If an integer, uses that many equally spaced + points. Should NOT include endpoint at 2pi. + transforms : dict of Transform or array-like + Transforms for R, Z, lambda, etc. Default is to build from grid. + + + Returns + ------- + field : ndarray, shape(n,3) + magnetic field at specified points, in either rpz or xyz coordinates + + Notes + ----- + Uses direct quadrature of the Biot-Savart integral for filamentary coils with + 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. + + """ + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "B") + + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic vector potential at a set of points. + + The coil current may be overridden by including `current` + in the `params` dictionary. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict, optional + Parameters to pass to Curve. + basis : {"rpz", "xyz"} + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None, optional + Grid used to discretize coil. If an integer, uses that many equally spaced + points. Should NOT include endpoint at 2pi. + transforms : dict of Transform or array-like + Transforms for R, Z, lambda, etc. Default is to build from grid. + + Returns + ------- + vector_potential : ndarray, shape(n,3) + Magnetic vector potential at specified points, in either rpz or + xyz coordinates. + + Notes + ----- + Uses direct quadrature of the Biot-Savart integral for filamentary coils with + 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. + + """ + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") def __repr__(self): """Get the string form of the object.""" @@ -783,10 +957,16 @@ def __init__( ): super().__init__(current, X, Y, Z, knots, method, name) - def compute_magnetic_field( - self, coords, params=None, basis="rpz", source_grid=None, transforms=None + def _compute_A_or_B( + self, + coords, + params=None, + basis="rpz", + source_grid=None, + transforms=None, + compute_A_or_B="B", ): - """Compute magnetic field at a set of points. + """Compute magnetic field or vector potential at a set of points. The coil current may be overridden by including `current` in the `params` dictionary. @@ -804,6 +984,9 @@ def compute_magnetic_field( points. Should NOT include endpoint at 2pi. transforms : dict of Transform or array-like Transforms for R, Z, lambda, etc. Default is to build from grid. + compute_A_or_B: {"A", "B"}, optional + whether to compute the magnetic vector potential "A" or the magnetic field + "B". Defaults to "B" Returns ------- @@ -817,6 +1000,12 @@ def compute_magnetic_field( is approximately quadratic in the number of coil points. """ + errorif( + compute_A_or_B not in ["A", "B"], + ValueError, + f'Expected "A" or "B" for compute_A_or_B, instead got {compute_A_or_B}', + ) + op = {"B": biot_savart_hh, "A": biot_savart_vector_potential_hh}[compute_A_or_B] assert basis.lower() in ["rpz", "xyz"] coords = jnp.atleast_2d(jnp.asarray(coords)) if basis == "rpz": @@ -826,7 +1015,9 @@ def compute_magnetic_field( else: current = params.pop("current", self.current) - data = self.compute(["x"], grid=source_grid, params=params, basis="xyz") + data = self.compute( + ["x"], grid=source_grid, params=params, basis="xyz", transforms=transforms + ) # 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) @@ -837,11 +1028,85 @@ def compute_magnetic_field( # coils curvature which is a 2nd derivative of the position, and doing that # with only possibly c1 cubic splines is inaccurate, so we don't do it # (for now, maybe in the future?) - B = biot_savart_hh(coords, coil_pts_start, coil_pts_end, current) + AB = op(coords, coil_pts_start, coil_pts_end, current) if basis == "rpz": - B = xyz2rpz_vec(B, x=coords[:, 0], y=coords[:, 1]) - return B + AB = xyz2rpz_vec(AB, x=coords[:, 0], y=coords[:, 1]) + return AB + + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic field at a set of points. + + The coil current may be overridden by including `current` + in the `params` dictionary. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict, optional + Parameters to pass to Curve. + basis : {"rpz", "xyz"} + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None, optional + Grid used to discretize coil. If an integer, uses that many equally spaced + points. Should NOT include endpoint at 2pi. + transforms : dict of Transform or array-like + Transforms for R, Z, lambda, etc. Default is to build from grid. + + Returns + ------- + field : ndarray, shape(n,3) + magnetic field at specified points, in either rpz or xyz coordinates + + Notes + ----- + 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. + + """ + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "B") + + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic vector potential at a set of points. + + The coil current may be overridden by including `current` + in the `params` dictionary. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate magnetic vector potential at in [R,phi,Z] + or [X,Y,Z] coordinates. + params : dict, optional + Parameters to pass to Curve. + basis : {"rpz", "xyz"} + Basis for input coordinates and returned magnetic vector potential. + source_grid : Grid, int or None, optional + Grid used to discretize coil. If an integer, uses that many equally spaced + points. Should NOT include endpoint at 2pi. + transforms : dict of Transform or array-like + Transforms for R, Z, lambda, etc. Default is to build from grid. + + Returns + ------- + A : ndarray, shape(n,3) + Magnetic vector potential at specified points, in either + rpz or xyz coordinates + + Notes + ----- + Discretizes the coil into straight segments between grid points, and uses the + Hanson-Hirshman expression for exact vector potential from a straight segment. + Convergence is approximately quadratic in the number of coil points. + + """ + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") @classmethod def from_values( @@ -1153,8 +1418,14 @@ def _compute_position(self, params=None, grid=None, **kwargs): x = rpz return x - def compute_magnetic_field( - self, coords, params=None, basis="rpz", source_grid=None, transforms=None + def _compute_A_or_B( + self, + coords, + params=None, + basis="rpz", + source_grid=None, + transforms=None, + compute_A_or_B="B", ): """Compute magnetic field at a set of points. @@ -1171,13 +1442,22 @@ def compute_magnetic_field( points. Should NOT include endpoint at 2pi. transforms : dict of Transform or array-like Transforms for R, Z, lambda, etc. Default is to build from grid. + compute_A_or_B: {"A", "B"}, optional + whether to compute the magnetic vector potential "A" or the magnetic field + "B". Defaults to "B" Returns ------- field : ndarray, shape(n,3) - Magnetic field at specified nodes, in [R,phi,Z] or [X,Y,Z] coordinates. + Magnetic field or vector potential at specified nodes, in [R,phi,Z] + or [X,Y,Z] coordinates. """ + errorif( + compute_A_or_B not in ["A", "B"], + ValueError, + f'Expected "A" or "B" for compute_A_or_B, instead got {compute_A_or_B}', + ) assert basis.lower() in ["rpz", "xyz"] coords = jnp.atleast_2d(jnp.asarray(coords)) if params is None: @@ -1207,31 +1487,89 @@ def compute_magnetic_field( # field period rotation is easiest in [R,phi,Z] coordinates coords_rpz = xyz2rpz(coords_xyz) + op = { + "B": self[0].compute_magnetic_field, + "A": self[0].compute_magnetic_vector_potential, + }[compute_A_or_B] # sum the magnetic fields from each field period - def nfp_loop(k, B): + def nfp_loop(k, AB): coords_nfp = coords_rpz + jnp.array([0, 2 * jnp.pi * k / self.NFP, 0]) - def body(B, x): - B += self[0].compute_magnetic_field( - coords_nfp, params=x, basis="rpz", source_grid=source_grid - ) - return B, None + def body(AB, x): + AB += op(coords_nfp, params=x, basis="rpz", source_grid=source_grid) + return AB, None - B += scan(body, jnp.zeros(coords_nfp.shape), tree_stack(params))[0] - return B + AB += scan(body, jnp.zeros(coords_nfp.shape), tree_stack(params))[0] + return AB - B = fori_loop(0, self.NFP, nfp_loop, jnp.zeros_like(coords_rpz)) + AB = fori_loop(0, self.NFP, nfp_loop, jnp.zeros_like(coords_rpz)) - # sum the magnetic fields from both halves of the symmetric field period + # sum the magnetic field/potential from both halves of + # the symmetric field period if self.sym: - B = B[: coords.shape[0], :] + B[coords.shape[0] :, :] * jnp.array( + AB = AB[: coords.shape[0], :] + AB[coords.shape[0] :, :] * jnp.array( [-1, 1, 1] ) if basis.lower() == "xyz": - B = rpz2xyz_vec(B, x=coords[:, 0], y=coords[:, 1]) - return B + AB = rpz2xyz_vec(AB, x=coords[:, 0], y=coords[:, 1]) + return AB + + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic field at a set of points. + + Parameters + ---------- + 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 coils, either the same for all coils or one for each. + basis : {"rpz", "xyz"} + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None, optional + Grid used to discretize coils. If an integer, uses that many equally spaced + points. Should NOT include endpoint at 2pi. + transforms : dict of Transform or array-like + Transforms for R, Z, lambda, etc. Default is to build from grid. + + Returns + ------- + field : ndarray, shape(n,3) + Magnetic field at specified nodes, in [R,phi,Z] or [X,Y,Z] coordinates. + + """ + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "B") + + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic vector potential at a set of points. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate potential at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional + 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. + source_grid : Grid, int or None, optional + Grid used to discretize coils. If an integer, uses that many equally spaced + points. Should NOT include endpoint at 2pi. + transforms : dict of Transform or array-like + Transforms for R, Z, lambda, etc. Default is to build from grid. + + Returns + ------- + vector_potential : ndarray, shape(n,3) + magnetic vector potential at specified points, in either rpz + or xyz coordinates + + """ + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") @classmethod def linspaced_angular( @@ -2002,6 +2340,65 @@ def _compute_position(self, params=None, grid=None, **kwargs): ) return x + def _compute_A_or_B( + self, + coords, + params=None, + basis="rpz", + source_grid=None, + transforms=None, + compute_A_or_B="B", + ): + """Compute magnetic field or vector potential at a set of points. + + Parameters + ---------- + 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 coils, either the same for all coils or one for each. + If array-like, should be 1 value per coil. + basis : {"rpz", "xyz"} + 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. + If array-like, should be 1 value per coil. + transforms : dict of Transform or array-like + Transforms for R, Z, lambda, etc. Default is to build from grid. + compute_A_or_B: {"A", "B"}, optional + whether to compute the magnetic vector potential "A" or the magnetic field + "B". Defaults to "B" + + Returns + ------- + field : ndarray, shape(n,3) + magnetic field or vector potential at specified points, in either rpz + or xyz coordinates + + """ + errorif( + compute_A_or_B not in ["A", "B"], + ValueError, + f'Expected "A" or "B" for compute_A_or_B, instead got {compute_A_or_B}', + ) + params = self._make_arraylike(params) + source_grid = self._make_arraylike(source_grid) + transforms = self._make_arraylike(transforms) + + AB = 0 + if compute_A_or_B == "B": + for coil, par, grd, tr in zip(self.coils, params, source_grid, transforms): + AB += coil.compute_magnetic_field( + coords, par, basis, grd, transforms=tr + ) + elif compute_A_or_B == "A": + for coil, par, grd, tr in zip(self.coils, params, source_grid, transforms): + AB += coil.compute_magnetic_vector_potential( + coords, par, basis, grd, transforms=tr + ) + return AB + def compute_magnetic_field( self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): @@ -2029,15 +2426,37 @@ def compute_magnetic_field( magnetic field at specified points, in either rpz or xyz coordinates """ - params = self._make_arraylike(params) - source_grid = self._make_arraylike(source_grid) - transforms = self._make_arraylike(transforms) + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "B") - B = 0 - for coil, par, grd, tr in zip(self.coils, params, source_grid, transforms): - B += coil.compute_magnetic_field(coords, par, basis, grd, transforms=tr) + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic vector potential at a set of points. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate potential at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional + Parameters to pass to coils, either the same for all coils or one for each. + If array-like, should be 1 value per coil. + basis : {"rpz", "xyz"} + 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. + If array-like, should be 1 value per coil. + transforms : dict of Transform or array-like + Transforms for R, Z, lambda, etc. Default is to build from grid. - return B + Returns + ------- + vector_potential : ndarray, shape(n,3) + magnetic vector potential at specified points, in either rpz + or xyz coordinates + + """ + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") def to_FourierPlanar( self, N=10, grid=None, basis="xyz", name="", check_intersection=False diff --git a/desc/magnetic_fields/__init__.py b/desc/magnetic_fields/__init__.py index 0a8f18abd..173b04e7e 100644 --- a/desc/magnetic_fields/__init__.py +++ b/desc/magnetic_fields/__init__.py @@ -9,6 +9,7 @@ SplineMagneticField, SumMagneticField, ToroidalMagneticField, + VectorPotentialField, VerticalMagneticField, _MagneticField, field_line_integrate, diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 7b32a1217..760f4e372 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -15,7 +15,7 @@ DoubleFourierSeries, ) from desc.compute import compute as compute_fun -from desc.compute import rpz2xyz, rpz2xyz_vec, xyz2rpz +from desc.compute import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec from desc.compute.utils import get_params, get_transforms from desc.derivatives import Derivative from desc.equilibrium import EquilibriaFamily, Equilibrium @@ -62,6 +62,40 @@ def body(i, B): return 1e-7 * fori_loop(0, J.shape[0], body, B) +def biot_savart_general_vector_potential(re, rs, J, dV): + """Biot-Savart law for arbitrary sources for vector potential. + + Parameters + ---------- + re : ndarray, shape(n_eval_pts, 3) + evaluation points to evaluate B at, in cartesian. + rs : ndarray, shape(n_src_pts, 3) + source points for current density J, in cartesian. + J : ndarray, shape(n_src_pts, 3) + current density vector at source points, in cartesian. + dV : ndarray, shape(n_src_pts) + volume element at source points + + Returns + ------- + A : ndarray, shape(n,3) + magnetic vector potential in cartesian components at specified points + """ + re, rs, J, dV = map(jnp.asarray, (re, rs, J, dV)) + assert J.shape == rs.shape + JdV = J * dV[:, None] + A = jnp.zeros_like(re) + + def body(i, A): + r = re - rs[i, :] + num = JdV[i, :] + den = jnp.linalg.norm(r, axis=-1) + A = A + jnp.where(den[:, None] == 0, 0, num / den[:, None]) + return A + + return 1e-7 * fori_loop(0, J.shape[0], body, A) + + def read_BNORM_file(fname, surface, eval_grid=None, scale_by_curpol=True): """Read BNORM-style .txt file containing Bnormal Fourier coefficients. @@ -193,6 +227,8 @@ def compute_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. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid Returns ------- @@ -205,6 +241,33 @@ def __call__(self, grid, params=None, basis="rpz"): """Compute magnetic field at a set of points.""" return self.compute_magnetic_field(grid, params, basis) + @abstractmethod + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic vector potential at a set of points. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate vector potential 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 vector potential. + source_grid : Grid, int or None or array-like, optional + Grid used to discretize MagneticField object if calculating A from + Biot-Savart. Should NOT include endpoint at 2pi. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + + Returns + ------- + A : ndarray, shape(N,3) + magnetic vector potential at specified points + + """ + def compute_Bnormal( self, surface, @@ -410,6 +473,7 @@ def save_mgrid( nR=101, nZ=101, nphi=90, + save_vector_potential=True, ): """Save the magnetic field to an mgrid NetCDF file in "raw" format. @@ -431,6 +495,9 @@ def save_mgrid( Number of grid points in the Z coordinate (default = 101). nphi : int, optional Number of grid points in the toroidal angle (default = 90). + save_vector_potential : bool, optional + Whether or not to save the magnetic vector potential to the mgrid + file, in addition to the magnetic field. Defaults to True. Returns ------- @@ -451,6 +518,15 @@ def save_mgrid( B_phi = field[:, 1].reshape(nphi, nZ, nR) B_Z = field[:, 2].reshape(nphi, nZ, nR) + # evaluate magnetic vector potential on grid + if save_vector_potential: + field = self.compute_magnetic_vector_potential(grid, basis="rpz") + A_R = field[:, 0].reshape(nphi, nZ, nR) + A_phi = field[:, 1].reshape(nphi, nZ, nR) + A_Z = field[:, 2].reshape(nphi, nZ, nR) + else: + A_R = None + # write mgrid file file = Dataset(path, mode="w", format="NETCDF3_64BIT_OFFSET") @@ -537,6 +613,28 @@ def save_mgrid( ) bz_001[:] = B_Z + if save_vector_potential: + ar_001 = file.createVariable("ar_001", np.float64, ("phi", "zee", "rad")) + ar_001.long_name = ( + "A_R = radial component of magnetic vector potential " + "in lab frame (T/m)." + ) + ar_001[:] = A_R + + ap_001 = file.createVariable("ap_001", np.float64, ("phi", "zee", "rad")) + ap_001.long_name = ( + "A_phi = toroidal component of magnetic vector potential " + "in lab frame (T/m)." + ) + ap_001[:] = A_phi + + az_001 = file.createVariable("az_001", np.float64, ("phi", "zee", "rad")) + az_001.long_name = ( + "A_Z = vertical component of magnetic vector potential " + "in lab frame (T/m)." + ) + az_001[:] = A_Z + file.close() @@ -618,6 +716,33 @@ def compute_magnetic_field( B = rpz2xyz_vec(B, phi=coords[:, 1]) return B + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None + ): + """Compute magnetic vector potential at a set of points. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate vector potential 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 vector potential. + source_grid : Grid, int or None or array-like, optional + Unused by this MagneticField class. + + Returns + ------- + A : ndarray, shape(N,3) + magnetic vector potential at specified points + + """ + raise NotImplementedError( + "MagneticFieldFromUser does not have vector potential calculation " + "implemented." + ) + class ScaledMagneticField(_MagneticField, Optimizable): """Magnetic field scaled by a scalar value. @@ -703,6 +828,35 @@ def compute_magnetic_field( coords, params, basis, source_grid ) + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic vector potential at a set of points. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate vector potential 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 vector potential. + source_grid : Grid, int or None or array-like, optional + Grid used to discretize MagneticField object if calculating A from + Biot-Savart. Should NOT include endpoint at 2pi. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + + Returns + ------- + A : ndarray, shape(N,3) + scaled magnetic vector potential at specified points + + """ + return self._scale * self._field.compute_magnetic_vector_potential( + coords, params, basis, source_grid + ) + class SumMagneticField(_MagneticField, MutableSequence, OptimizableCollection): """Sum of two or more magnetic field sources. @@ -724,10 +878,16 @@ def __init__(self, *fields): ) self._fields = fields - def compute_magnetic_field( - self, coords, params=None, basis="rpz", source_grid=None, transforms=None + def _compute_A_or_B( + self, + coords, + params=None, + basis="rpz", + source_grid=None, + transforms=None, + compute_A_or_B="B", ): - """Compute magnetic field at a set of points. + """Compute magnetic field or vector potential at a set of points. Parameters ---------- @@ -742,6 +902,9 @@ def compute_magnetic_field( Biot-Savart. Should NOT include endpoint at 2pi. transforms : dict of Transform Transforms for R, Z, lambda, etc. Default is to build from source_grid + compute_A_or_B: {"A", "B"}, optional + whether to compute the magnetic vector potential "A" or the magnetic field + "B". Defaults to "B" Returns ------- @@ -749,6 +912,11 @@ def compute_magnetic_field( scaled magnetic field at specified points """ + errorif( + compute_A_or_B not in ["A", "B"], + ValueError, + f'Expected "A" or "B" for compute_A_or_B, instead got {compute_A_or_B}', + ) if params is None: params = [None] * len(self._fields) if isinstance(params, dict): @@ -770,13 +938,74 @@ def compute_magnetic_field( # zip does not terminate early transforms = transforms * len(self._fields) - B = 0 + op = {"B": "compute_magnetic_field", "A": "compute_magnetic_vector_potential"}[ + compute_A_or_B + ] + + AB = 0 for i, (field, g, tr) in enumerate(zip(self._fields, source_grid, transforms)): - B += field.compute_magnetic_field( + AB += getattr(field, op)( coords, params[i % len(params)], basis, source_grid=g, transforms=tr ) + return AB - return B + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic field at a set of points. + + Parameters + ---------- + 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. + 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. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + + Returns + ------- + field : ndarray, shape(N,3) + sum magnetic field at specified points + + """ + return self._compute_A_or_B( + coords, params, basis, source_grid, transforms, compute_A_or_B="B" + ) + + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic vector potential at a set of points. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate vector potential 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 vector potential. + source_grid : Grid, int or None or array-like, optional + Grid used to discretize MagneticField object if calculating A from + Biot-Savart. Should NOT include endpoint at 2pi. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + + Returns + ------- + A : ndarray, shape(N,3) + sum magnetic vector potential at specified points + + """ + return self._compute_A_or_B( + coords, params, basis, source_grid, transforms, compute_A_or_B="A" + ) # dunder methods required by MutableSequence def __getitem__(self, i): @@ -886,10 +1115,54 @@ def compute_magnetic_field( return B + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic vector potential at a set of points. + + The vector potential is specified assuming the Coulomb Gauge. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate vector potential 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 vector potential. + source_grid : Grid, int or None or array-like, optional + Unused by this MagneticField class. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + Unused by this MagneticField class. + + Returns + ------- + A : ndarray, shape(N,3) + magnetic vector potential at specified points + + """ + params = setdefault(params, {}) + B0 = params.get("B0", self.B0) + R0 = params.get("R0", self.R0) + + assert basis.lower() in ["rpz", "xyz"] + coords = jnp.atleast_2d(jnp.asarray(coords)) + if basis == "xyz": + coords = xyz2rpz(coords) + az = -B0 * R0 * jnp.log(coords[:, 0]) + arp = jnp.zeros_like(az) + A = jnp.array([arp, arp, az]).T + # b/c it only has a nonzero z component, no need + # to switch bases back if xyz is given + return A + class VerticalMagneticField(_MagneticField, Optimizable): """Uniform magnetic field purely in the vertical (Z) direction. + The vector potential is specified assuming the Coulomb Gauge. + Parameters ---------- B0 : float @@ -940,18 +1213,63 @@ def compute_magnetic_field( params = setdefault(params, {}) B0 = params.get("B0", self.B0) - assert basis.lower() in ["rpz", "xyz"] coords = jnp.atleast_2d(jnp.asarray(coords)) - if basis == "xyz": - coords = xyz2rpz(coords) bz = B0 * jnp.ones_like(coords[:, 2]) brp = jnp.zeros_like(bz) B = jnp.array([brp, brp, bz]).T - if basis == "xyz": - B = rpz2xyz_vec(B, phi=coords[:, 1]) + # b/c it only has a nonzero z component, no need + # to switch bases back if xyz is given return B + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic vector potential at a set of points. + + The vector potential is specified assuming the Coulomb Gauge. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate vector potential 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 vector potential. + source_grid : Grid, int or None or array-like, optional + Unused by this MagneticField class. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + Unused by this MagneticField class. + + Returns + ------- + A : ndarray, shape(N,3) + magnetic vector potential at specified points + + """ + params = setdefault(params, {}) + B0 = params.get("B0", self.B0) + + coords = jnp.atleast_2d(jnp.asarray(coords)) + + if basis == "xyz": + coords_xyz = coords + coords_rpz = xyz2rpz(coords) + else: + coords_rpz = coords + coords_xyz = rpz2xyz(coords) + ax = B0 / 2 * coords_xyz[:, 1] + ay = -B0 / 2 * coords_xyz[:, 0] + + az = jnp.zeros_like(ax) + A = jnp.array([ax, ay, az]).T + if basis == "rpz": + A = xyz2rpz_vec(A, phi=coords_rpz[:, 1]) + + return A + class PoloidalMagneticField(_MagneticField, Optimizable): """Pure poloidal magnetic field (ie in theta direction). @@ -1062,6 +1380,36 @@ def compute_magnetic_field( return B + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic vector potential at a set of points. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate vector potential 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 vector potential. + source_grid : Grid, int or None or array-like, optional + Unused by this MagneticField class. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + Unused by this MagneticField class. + + Returns + ------- + A : ndarray, shape(N,3) + magnetic vector potential at specified points + + """ + raise NotImplementedError( + "PoloidalMagneticField has nonzero divergence, therefore it can't be " + "represented with a vector potential." + ) + class SplineMagneticField(_MagneticField, Optimizable): """Magnetic field from precomputed values on a grid. @@ -1080,6 +1428,12 @@ class SplineMagneticField(_MagneticField, Optimizable): toroidal magnetic field on grid BZ : array-like, shape(NR,Nphi,NZ,Ngroups) vertical magnetic field on grid + AR : array-like, shape(NR,Nphi,NZ,Ngroups) + radial magnetic vector potential on grid, optional + aphi : array-like, shape(NR,Nphi,NZ,Ngroups) + toroidal magnetic vector potential on grid, optional + AZ : array-like, shape(NR,Nphi,NZ,Ngroups) + vertical magnetic vector potential on grid, optional currents : array-like, shape(Ngroups) Currents or scaling factors for each field group. NFP : int, optional @@ -1098,6 +1452,9 @@ class SplineMagneticField(_MagneticField, Optimizable): "_BR", "_Bphi", "_BZ", + "_AR", + "_Aphi", + "_AZ", "_method", "_extrap", "_derivs", @@ -1110,7 +1467,20 @@ class SplineMagneticField(_MagneticField, Optimizable): _static_attrs = ["_extrap", "_period"] def __init__( - self, R, phi, Z, BR, Bphi, BZ, currents=1.0, NFP=1, method="cubic", extrap=False + self, + R, + phi, + Z, + BR, + Bphi, + BZ, + AR=None, + Aphi=None, + AZ=None, + currents=1.0, + NFP=1, + method="cubic", + extrap=False, ): R, phi, Z, currents = map( lambda x: jnp.atleast_1d(jnp.asarray(x)), (R, phi, Z, currents) @@ -1152,6 +1522,17 @@ def _atleast_4d(x): self._derivs["BR"] = self._approx_derivs(self._BR) self._derivs["Bphi"] = self._approx_derivs(self._Bphi) self._derivs["BZ"] = self._approx_derivs(self._BZ) + if AR is not None and Aphi is not None and AZ is not None: + AR, Aphi, AZ = map(_atleast_4d, (AR, Aphi, AZ)) + assert AR.shape == Aphi.shape == AZ.shape == shape + self._AR = AR + self._Aphi = Aphi + self._AZ = AZ + self._derivs["AR"] = self._approx_derivs(self._AR) + self._derivs["Aphi"] = self._approx_derivs(self._Aphi) + self._derivs["AZ"] = self._approx_derivs(self._AZ) + else: + self._AR = self._Aphi = self._AZ = None @property def NFP(self): @@ -1190,10 +1571,16 @@ def _approx_derivs(self, Bi): tempdict[key] = val[:, 0, :] return tempdict - def compute_magnetic_field( - self, coords, params=None, basis="rpz", source_grid=None, transforms=None + def _compute_A_or_B( + self, + coords, + params=None, + basis="rpz", + source_grid=None, + transforms=None, + compute_A_or_B="B", ): - """Compute magnetic field at a set of points. + """Compute magnetic field or magnetic vector potential at a set of points. Parameters ---------- @@ -1208,107 +1595,185 @@ def compute_magnetic_field( transforms : dict of Transform Transforms for R, Z, lambda, etc. Default is to build from source_grid Unused by this MagneticField class. + compute_A_or_B: {"A", "B"}, optional + whether to compute the magnetic vector potential "A" or the magnetic field + "B". Defaults to "B" Returns ------- field : ndarray, shape(N,3) - magnetic field at specified points, in cylindrical form [BR, Bphi,BZ] + magnetic field or vector potential at specified points, + in cylindrical form [BR, Bphi,BZ] """ + errorif( + compute_A_or_B not in ["A", "B"], + ValueError, + f'Expected "A" or "B" for compute_A_or_B, instead got {compute_A_or_B}', + ) + errorif( + compute_A_or_B == "A" and self._AR is None, + ValueError, + "Cannot calculate vector potential" + " as no vector potential spline values exist.", + ) assert basis.lower() in ["rpz", "xyz"] currents = self.currents if params is None else params["currents"] coords = jnp.atleast_2d(jnp.asarray(coords)) if basis == "xyz": coords = xyz2rpz(coords) Rq, phiq, Zq = coords.T + if compute_A_or_B == "B": + A_or_B_R = self._BR + A_or_B_phi = self._Bphi + A_or_B_Z = self._BZ + elif compute_A_or_B == "A": + A_or_B_R = self._AR + A_or_B_phi = self._Aphi + A_or_B_Z = self._AZ + if self._axisym: - BRq = interp2d( + ABRq = interp2d( Rq, Zq, self._R, self._Z, - self._BR[:, 0, :], + A_or_B_R[:, 0, :], self._method, (0, 0), self._extrap, (None, None), - **self._derivs["BR"], + **self._derivs[compute_A_or_B + "R"], ) - Bphiq = interp2d( + ABphiq = interp2d( Rq, Zq, self._R, self._Z, - self._Bphi[:, 0, :], + A_or_B_phi[:, 0, :], self._method, (0, 0), self._extrap, (None, None), - **self._derivs["Bphi"], + **self._derivs[compute_A_or_B + "phi"], ) - BZq = interp2d( + ABZq = interp2d( Rq, Zq, self._R, self._Z, - self._BZ[:, 0, :], + A_or_B_Z[:, 0, :], self._method, (0, 0), self._extrap, (None, None), - **self._derivs["BZ"], + **self._derivs[compute_A_or_B + "Z"], ) else: - BRq = interp3d( + ABRq = interp3d( Rq, phiq, Zq, self._R, self._phi, self._Z, - self._BR, + A_or_B_R, self._method, (0, 0, 0), self._extrap, (None, 2 * np.pi / self.NFP, None), - **self._derivs["BR"], + **self._derivs[compute_A_or_B + "R"], ) - Bphiq = interp3d( + ABphiq = interp3d( Rq, phiq, Zq, self._R, self._phi, self._Z, - self._Bphi, + A_or_B_phi, self._method, (0, 0, 0), self._extrap, (None, 2 * np.pi / self.NFP, None), - **self._derivs["Bphi"], + **self._derivs[compute_A_or_B + "phi"], ) - BZq = interp3d( + ABZq = interp3d( Rq, phiq, Zq, self._R, self._phi, self._Z, - self._BZ, + A_or_B_Z, self._method, (0, 0, 0), self._extrap, (None, 2 * np.pi / self.NFP, None), - **self._derivs["BZ"], + **self._derivs[compute_A_or_B + "Z"], ) - # BRq etc shape(nq, ngroups) - B = jnp.stack([BRq, Bphiq, BZq], axis=1) - # B shape(nq, 3, ngroups) - B = jnp.sum(B * currents, axis=-1) + # ABRq etc shape(nq, ngroups) + AB = jnp.stack([ABRq, ABphiq, ABZq], axis=1) + # AB shape(nq, 3, ngroups) + AB = jnp.sum(AB * currents, axis=-1) if basis == "xyz": - B = rpz2xyz_vec(B, phi=coords[:, 1]) - return B + AB = rpz2xyz_vec(AB, phi=coords[:, 1]) + return AB + + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic field at a set of points. + + Parameters + ---------- + 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. + source_grid : Grid, int or None + Unused by this MagneticField class. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + Unused by this MagneticField class. + + Returns + ------- + field : ndarray, shape(N,3) + magnetic field at specified points, in cylindrical form [BR, Bphi,BZ] + + """ + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "B") + + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic vector potential at a set of points. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate vector potential 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 vector potential. + source_grid : Grid, int or None or array-like, optional + Unused by this MagneticField class. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + Unused by this MagneticField class. + + Returns + ------- + A : ndarray, shape(N,3) + magnetic vector potential at specified points + + """ + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") @classmethod def from_mgrid(cls, mgrid_file, extcur=None, method="cubic", extrap=False): @@ -1366,8 +1831,40 @@ def from_mgrid(cls, mgrid_file, extcur=None, method="cubic", extrap=False): bp = np.moveaxis(bp, (0, 1, 2), (1, 2, 0)) bz = np.moveaxis(bz, (0, 1, 2), (1, 2, 0)) + # sum magnetic vector potentials from each coil + ar = np.zeros([kp, jz, ir, nextcur]) + ap = np.zeros([kp, jz, ir, nextcur]) + az = np.zeros([kp, jz, ir, nextcur]) + try: + for i in range(nextcur): + coil_id = "%03d" % (i + 1,) + ar[:, :, :, i] += mgrid["ar_" + coil_id][ + () + ] # A_R radial mag. vec. potential + ap[:, :, :, i] += mgrid["ap_" + coil_id][ + () + ] # A_phi toroidal mag. vec. potential + az[:, :, :, i] += mgrid["az_" + coil_id][ + () + ] # A_Z vertical mag. vec. potential + + # shift axes to correct order + ar = np.moveaxis(ar, (0, 1, 2), (1, 2, 0)) + ap = np.moveaxis(ap, (0, 1, 2), (1, 2, 0)) + az = np.moveaxis(az, (0, 1, 2), (1, 2, 0)) + except IndexError: + warnif( + True, + UserWarning, + "mgrid does not appear to contain vector potential information." + " Vector potential will not be computable.", + ) + ar = ap = az = None + mgrid.close() - return cls(Rgrid, pgrid, Zgrid, br, bp, bz, extcur, nfp, method, extrap) + return cls( + Rgrid, pgrid, Zgrid, br, bp, bz, ar, ap, az, extcur, nfp, method, extrap + ) @classmethod def from_field( @@ -1397,6 +1894,15 @@ def from_field( shp = rr.shape coords = np.array([rr.flatten(), pp.flatten(), zz.flatten()]).T BR, BP, BZ = field.compute_magnetic_field(coords, params, basis="rpz").T + try: + AR, AP, AZ = field.compute_magnetic_vector_potential( + coords, params, basis="rpz" + ).T + AR = AR.reshape(shp) + AP = AP.reshape(shp) + AZ = AZ.reshape(shp) + except NotImplementedError: + AR = AP = AZ = None return cls( R, phi, @@ -1404,6 +1910,9 @@ def from_field( BR.reshape(shp), BP.reshape(shp), BZ.reshape(shp), + AR=AR, + Aphi=AP, + AZ=AZ, currents=1.0, NFP=NFP, method=method, @@ -1474,6 +1983,187 @@ def compute_magnetic_field( B = rpz2xyz_vec(B, phi=coords[:, 1]) return B + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic vector potential at a set of points. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate vector potential 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 vector potential. + source_grid : Grid, int or None or array-like, optional + Unused by this MagneticField class. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + Unused by this MagneticField class. + + Returns + ------- + A : ndarray, shape(N,3) + magnetic vector potential at specified points + + """ + raise NotImplementedError( + "ScalarPotentialField does not have vector potential calculation " + "implemented." + ) + + +class VectorPotentialField(_MagneticField): + """Magnetic field due to a vector magnetic potential in cylindrical coordinates. + + Parameters + ---------- + potential : callable + function to compute the vector potential. Should have a signature of + the form potential(R,phi,Z,*params) -> ndarray. + R,phi,Z are arrays of cylindrical coordinates. + params : dict, optional + default parameters to pass to potential function + + """ + + def __init__(self, potential, params=None): + self._potential = potential + self._params = params + + def _compute_A_or_B( + self, + coords, + params=None, + basis="rpz", + source_grid=None, + transforms=None, + compute_A_or_B="B", + ): + """Compute magnetic field or vector potential at a set of points. + + Parameters + ---------- + 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. + source_grid : Grid, int or None + Unused by this MagneticField class. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + Unused by this MagneticField class. + compute_A_or_B: {"A", "B"}, optional + whether to compute the magnetic vector potential "A" or the magnetic field + "B". Defaults to "B" + + Returns + ------- + field : ndarray, shape(N,3) + magnetic field at specified points + + """ + errorif( + compute_A_or_B not in ["A", "B"], + ValueError, + f'Expected "A" or "B" for compute_A_or_B, instead got {compute_A_or_B}', + ) + assert basis.lower() in ["rpz", "xyz"] + coords = jnp.atleast_2d(jnp.asarray(coords)) + coords = coords.astype(float) + if basis == "xyz": + coords = xyz2rpz(coords) + + if params is None: + params = self._params + r, p, z = coords.T + + if compute_A_or_B == "B": + funR = lambda x: self._potential(x, p, z, **params) + funP = lambda x: self._potential(r, x, z, **params) + funZ = lambda x: self._potential(r, p, x, **params) + + ap = self._potential(r, p, z, **params)[:, 1] + + # these are the gradients of each component of A + dAdr = Derivative.compute_jvp(funR, 0, (jnp.ones_like(r),), r) + dAdp = Derivative.compute_jvp(funP, 0, (jnp.ones_like(p),), p) + dAdz = Derivative.compute_jvp(funZ, 0, (jnp.ones_like(z),), z) + + # form the B components with the appropriate combinations + B = jnp.array( + [ + dAdp[:, 2] / r - dAdz[:, 1], + dAdz[:, 0] - dAdr[:, 2], + dAdr[:, 1] + (ap - dAdp[:, 0]) / r, + ] + ).T + if basis == "xyz": + B = rpz2xyz_vec(B, phi=coords[:, 1]) + return B + elif compute_A_or_B == "A": + A = self._potential(r, p, z, **params) + if basis == "xyz": + A = rpz2xyz_vec(A, phi=coords[:, 1]) + return A + + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic field at a set of points. + + Parameters + ---------- + 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. + source_grid : Grid, int or None + Unused by this MagneticField class. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + Unused by this MagneticField class. + + Returns + ------- + field : ndarray, shape(N,3) + magnetic field at specified points + + """ + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "B") + + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic vector potential at a set of points. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate vector potential 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 vector potential. + source_grid : Grid, int or None or array-like, optional + Unused by this MagneticField class. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + Unused by this MagneticField class. + + Returns + ------- + A : ndarray, shape(N,3) + magnetic vector potential at specified points + + """ + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") + def field_line_integrate( r0, diff --git a/desc/magnetic_fields/_current_potential.py b/desc/magnetic_fields/_current_potential.py index ab8a42d90..a8759155e 100644 --- a/desc/magnetic_fields/_current_potential.py +++ b/desc/magnetic_fields/_current_potential.py @@ -11,7 +11,11 @@ from desc.optimizable import Optimizable, optimizable_parameter from desc.utils import copy_coeffs, errorif, setdefault, warnif -from ._core import _MagneticField, biot_savart_general +from ._core import ( + _MagneticField, + biot_savart_general, + biot_savart_general_vector_potential, +) class CurrentPotentialField(_MagneticField, FourierRZToroidalSurface): @@ -177,10 +181,16 @@ 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", source_grid=None, transforms=None + def _compute_A_or_B( + self, + coords, + params=None, + basis="rpz", + source_grid=None, + transforms=None, + compute_A_or_B="B", ): - """Compute magnetic field at a set of points. + """Compute magnetic field or vector potential at a set of points. Parameters ---------- @@ -194,11 +204,14 @@ def compute_magnetic_field( Source grid upon which to evaluate the surface current density K. transforms : dict of Transform Transforms for R, Z, lambda, etc. Default is to build from source_grid + compute_A_or_B: {"A", "B"}, optional + whether to compute the magnetic vector potential "A" or the magnetic field + "B". Defaults to "B" Returns ------- field : ndarray, shape(N,3) - magnetic field at specified points + magnetic field or vector potential at specified points """ source_grid = source_grid or LinearGrid( @@ -206,15 +219,70 @@ def compute_magnetic_field( N=30 + 2 * self.N, NFP=self.NFP, ) - return _compute_magnetic_field_from_CurrentPotentialField( + return _compute_A_or_B_from_CurrentPotentialField( field=self, coords=coords, params=params, basis=basis, source_grid=source_grid, transforms=transforms, + compute_A_or_B=compute_A_or_B, ) + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic field at a set of points. + + Parameters + ---------- + 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. + source_grid : Grid, int or None or array-like, optional + Source grid upon which to evaluate the surface current density K. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + + Returns + ------- + field : ndarray, shape(N,3) + magnetic field at specified points + + """ + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "B") + + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic vector potential at a set of points. + + This assumes the Coulomb gauge. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate vector potential 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 vector potential. + source_grid : Grid, int or None or array-like, optional + Source grid upon which to evaluate the surface current density K. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + + Returns + ------- + A : ndarray, shape(N,3) + Magnetic vector potential at specified points. + + """ + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") + @classmethod def from_surface( cls, @@ -496,10 +564,16 @@ 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", source_grid=None, transforms=None + def _compute_A_or_B( + self, + coords, + params=None, + basis="rpz", + source_grid=None, + transforms=None, + compute_A_or_B="B", ): - """Compute magnetic field at a set of points. + """Compute magnetic field or vector potential at a set of points. Parameters ---------- @@ -513,11 +587,14 @@ def compute_magnetic_field( Source grid upon which to evaluate the surface current density K. transforms : dict of Transform Transforms for R, Z, lambda, etc. Default is to build from source_grid + compute_A_or_B: {"A", "B"}, optional + whether to compute the magnetic vector potential "A" or the magnetic field + "B". Defaults to "B" Returns ------- field : ndarray, shape(N,3) - magnetic field at specified points + magnetic field or vector potential at specified points """ source_grid = source_grid or LinearGrid( @@ -525,15 +602,70 @@ def compute_magnetic_field( N=30 + 2 * max(self.N, self.N_Phi), NFP=self.NFP, ) - return _compute_magnetic_field_from_CurrentPotentialField( + return _compute_A_or_B_from_CurrentPotentialField( field=self, coords=coords, params=params, basis=basis, source_grid=source_grid, transforms=transforms, + compute_A_or_B=compute_A_or_B, ) + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic field at a set of points. + + Parameters + ---------- + 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. + source_grid : Grid, int or None or array-like, optional + Source grid upon which to evaluate the surface current density K. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + + Returns + ------- + field : ndarray, shape(N,3) + magnetic field at specified points + + """ + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "B") + + def compute_magnetic_vector_potential( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None + ): + """Compute magnetic vector potential at a set of points. + + This assumes the Coulomb gauge. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate vector potential 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 vector potential. + source_grid : Grid, int or None or array-like, optional + Source grid upon which to evaluate the surface current density K. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid + + Returns + ------- + A : ndarray, shape(N,3) + Magnetic vector potential at specified points. + + """ + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") + @classmethod def from_surface( cls, @@ -613,10 +745,16 @@ def from_surface( ) -def _compute_magnetic_field_from_CurrentPotentialField( - field, coords, source_grid, params=None, basis="rpz", transforms=None +def _compute_A_or_B_from_CurrentPotentialField( + field, + coords, + source_grid, + params=None, + basis="rpz", + transforms=None, + compute_A_or_B="B", ): - """Compute magnetic field at a set of points. + """Compute magnetic field or vector potential at a set of points. Parameters ---------- @@ -631,25 +769,36 @@ def _compute_magnetic_field_from_CurrentPotentialField( should include the potential basis : {"rpz", "xyz"} basis for input coordinates and returned magnetic field + compute_A_or_B: {"A", "B"}, optional + whether to compute the magnetic vector potential "A" or the magnetic field + "B". Defaults to "B" Returns ------- field : ndarray, shape(N,3) - magnetic field at specified points + magnetic field or vector potential at specified points """ + errorif( + compute_A_or_B not in ["A", "B"], + ValueError, + f'Expected "A" or "B" for compute_A_or_B, instead got {compute_A_or_B}', + ) assert basis.lower() in ["rpz", "xyz"] coords = jnp.atleast_2d(jnp.asarray(coords)) if basis == "rpz": coords = rpz2xyz(coords) - + op = {"B": biot_savart_general, "A": biot_savart_general_vector_potential}[ + compute_A_or_B + ] # compute surface current, and store grid quantities # needed for integration in class if not params or not transforms: data = field.compute( ["K", "x"], grid=source_grid, + basis="rpz", params=params, transforms=transforms, jitable=True, @@ -680,7 +829,7 @@ def nfp_loop(j, f): rs = jnp.vstack((_rs[:, 0], phi, _rs[:, 2])).T rs = rpz2xyz(rs) K = rpz2xyz_vec(_K, phi=phi) - fj = biot_savart_general( + fj = op( coords, rs, K, diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index 58049ccc9..c7575b136 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -1310,6 +1310,14 @@ class ToroidalFlux(_Objective): by making the coil currents zero. Instead, this objective ensures the coils create the necessary toroidal flux for the equilibrium field. + Will try to use the vector potential method to calculate the toroidal flux + (Φ = ∮ 𝐀 ⋅ 𝐝𝐥 over the perimeter of a constant zeta plane) + instead of the brute force method using the magnetic field + (Φ = ∯ 𝐁 ⋅ 𝐝𝐒 over a constant zeta XS). The vector potential method + is much more efficient, however not every ``MagneticField`` object + has a vector potential available to compute, so in those cases + the magnetic field method is used. + Parameters ---------- eq : Equilibrium @@ -1355,6 +1363,7 @@ class ToroidalFlux(_Objective): name : str, optional Name of the objective function. + """ _coordinates = "rtz" @@ -1382,6 +1391,7 @@ def __init__( self._field_grid = field_grid self._eval_grid = eval_grid self._eq = eq + # TODO: add eq_fixed option so this can be used in single stage super().__init__( things=[field], @@ -1407,9 +1417,17 @@ def build(self, use_jit=True, verbose=1): """ eq = self._eq + self._use_vector_potential = True + try: + self._field.compute_magnetic_vector_potential([0, 0, 0]) + except (NotImplementedError, ValueError): + self._use_vector_potential = False if self._eval_grid is None: eval_grid = LinearGrid( - L=eq.L_grid, M=eq.M_grid, zeta=jnp.array(0.0), NFP=eq.NFP + L=eq.L_grid if not self._use_vector_potential else 0, + M=eq.M_grid, + zeta=jnp.array(0.0), + NFP=eq.NFP, ) self._eval_grid = eval_grid eval_grid = self._eval_grid @@ -1444,10 +1462,12 @@ def build(self, use_jit=True, verbose=1): if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - - data = eq.compute( - ["R", "phi", "Z", "|e_rho x e_theta|", "n_zeta"], grid=eval_grid - ) + data_keys = ["R", "phi", "Z"] + if self._use_vector_potential: + data_keys += ["e_theta"] + else: + data_keys += ["|e_rho x e_theta|", "n_zeta"] + data = eq.compute(data_keys, grid=eval_grid) plasma_coords = jnp.array([data["R"], data["phi"], data["Z"]]).T @@ -1489,22 +1509,32 @@ def compute(self, field_params=None, constants=None): data = constants["equil_data"] plasma_coords = constants["plasma_coords"] - - B = constants["field"].compute_magnetic_field( - plasma_coords, - basis="rpz", - source_grid=constants["field_grid"], - params=field_params, - ) grid = constants["eval_grid"] - B_dot_n_zeta = jnp.sum(B * data["n_zeta"], axis=1) + if self._use_vector_potential: + A = constants["field"].compute_magnetic_vector_potential( + plasma_coords, + basis="rpz", + source_grid=constants["field_grid"], + params=field_params, + ) - Psi = jnp.sum( - grid.spacing[:, 0] - * grid.spacing[:, 1] - * data["|e_rho x e_theta|"] - * B_dot_n_zeta - ) + A_dot_e_theta = jnp.sum(A * data["e_theta"], axis=1) + Psi = jnp.sum(grid.spacing[:, 1] * A_dot_e_theta) + else: + B = constants["field"].compute_magnetic_field( + plasma_coords, + basis="rpz", + source_grid=constants["field_grid"], + params=field_params, + ) + + B_dot_n_zeta = jnp.sum(B * data["n_zeta"], axis=1) + Psi = jnp.sum( + grid.spacing[:, 0] + * grid.spacing[:, 1] + * data["|e_rho x e_theta|"] + * B_dot_n_zeta + ) return Psi diff --git a/tests/test_coils.py b/tests/test_coils.py index 704ad5f76..d379200df 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -4,7 +4,9 @@ import numpy as np import pytest +import scipy +from desc.backend import jnp from desc.coils import ( CoilSet, FourierPlanarCoil, @@ -13,9 +15,10 @@ MixedCoilSet, SplineXYZCoil, ) -from desc.compute import get_params, get_transforms, xyz2rpz, xyz2rpz_vec +from desc.compute import get_params, get_transforms, rpz2xyz, xyz2rpz, xyz2rpz_vec +from desc.compute.utils import dot from desc.examples import get -from desc.geometry import FourierRZCurve, FourierRZToroidalSurface +from desc.geometry import FourierRZCurve, FourierRZToroidalSurface, FourierXYZCurve from desc.grid import Grid, LinearGrid from desc.io import load from desc.magnetic_fields import SumMagneticField, VerticalMagneticField @@ -149,6 +152,198 @@ def test_biot_savart_all_coils(self): B_true_rpz_phi, B_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierRZCoil" ) + @pytest.mark.unit + def test_biot_savart_vector_potential_all_coils(self): + """Test biot-savart vec potential implementation against analytic formula.""" + coil_grid = LinearGrid(zeta=100, endpoint=False) + + R = 2 + y = 1 + I = 1e7 + + A_true = np.atleast_2d([0, 0, 0]) + grid_xyz = np.atleast_2d([10, y, 0]) + grid_rpz = xyz2rpz(grid_xyz) + + def test(coil, grid_xyz, grid_rpz): + A_xyz = coil.compute_magnetic_vector_potential( + grid_xyz, basis="xyz", source_grid=coil_grid + ) + A_rpz = coil.compute_magnetic_vector_potential( + grid_rpz, basis="rpz", source_grid=coil_grid + ) + np.testing.assert_allclose( + A_true, A_xyz, rtol=1e-3, atol=1e-10, err_msg=f"Using {coil}" + ) + np.testing.assert_allclose( + A_true, A_rpz, rtol=1e-3, atol=1e-10, err_msg=f"Using {coil}" + ) + np.testing.assert_allclose( + A_true, A_rpz, rtol=1e-3, atol=1e-10, err_msg=f"Using {coil}" + ) + + # FourierXYZCoil + coil = FourierXYZCoil(I) + test(coil, grid_xyz, grid_rpz) + + # SplineXYZCoil + x = coil.compute("x", grid=coil_grid, basis="xyz")["x"] + coil = SplineXYZCoil(I, X=x[:, 0], Y=x[:, 1], Z=x[:, 2]) + test(coil, grid_xyz, grid_rpz) + + # FourierPlanarCoil + coil = FourierPlanarCoil(I) + test(coil, grid_xyz, grid_rpz) + + grid_xyz = np.atleast_2d([0, 0, y]) + grid_rpz = xyz2rpz(grid_xyz) + + # FourierRZCoil + coil = FourierRZCoil(I, R_n=np.array([R]), modes_R=np.array([0])) + test(coil, grid_xyz, grid_rpz) + # test in a CoilSet + coil2 = CoilSet(coil) + test(coil2, grid_xyz, grid_rpz) + # test in a MixedCoilSet + coil3 = MixedCoilSet(coil2, coil, check_intersection=False) + coil3[1].current = 0 + test(coil3, grid_xyz, grid_rpz) + + @pytest.mark.unit + def test_biot_savart_vector_potential_integral_all_coils(self): + """Test analytic expression of flux integral for all coils.""" + # taken from analytic benchmark in + # "A Magnetic Diagnostic Code for 3D Fusion Equilibria", Lazerson 2013 + # find flux for concentric loops of varying radii to a circular coil + + coil_grid = LinearGrid(zeta=1000, endpoint=False) + + R = 1 + I = 1e7 + + # analytic eqn for "A_phi" (phi is in dl direction for loop) + def _A_analytic(r): + # elliptic integral arguments must be k^2, not k, + # error in original paper and apparently in Jackson EM book too. + theta = np.pi / 2 + arg = R**2 + r**2 + 2 * r * R * np.sin(theta) + term_1_num = 4.0e-7 * I * R + term_1_den = np.sqrt(arg) + k_sqd = 4 * r * R * np.sin(theta) / arg + term_2_num = (2 - k_sqd) * scipy.special.ellipk( + k_sqd + ) - 2 * scipy.special.ellipe(k_sqd) + term_2_den = k_sqd + return term_1_num * term_2_num / term_1_den / term_2_den + + # we only evaluate it at theta=np.pi/2 (b/c it is in spherical coords) + rs = np.linspace(0.1, 3, 10, endpoint=True) + N = 200 + curve_grid = LinearGrid(zeta=N) + + def test( + coil, grid_xyz, grid_rpz, A_true_rpz, correct_flux, rtol=1e-10, atol=1e-12 + ): + """Test that we compute the correct flux for the given coil.""" + A_xyz = coil.compute_magnetic_vector_potential( + grid_xyz, basis="xyz", source_grid=coil_grid + ) + A_rpz = coil.compute_magnetic_vector_potential( + grid_rpz, basis="rpz", source_grid=coil_grid + ) + flux_xyz = jnp.sum( + dot(A_xyz, curve_data["x_s"], axis=-1) * curve_grid.spacing[:, 2] + ) + flux_rpz = jnp.sum( + dot(A_rpz, curve_data_rpz["x_s"], axis=-1) * curve_grid.spacing[:, 2] + ) + + np.testing.assert_allclose( + correct_flux, flux_xyz, rtol=rtol, err_msg=f"Using {coil}" + ) + np.testing.assert_allclose( + correct_flux, flux_rpz, rtol=rtol, err_msg=f"Using {coil}" + ) + np.testing.assert_allclose( + A_true_rpz, + A_rpz, + rtol=rtol, + atol=atol, + err_msg=f"Using {coil}", + ) + + for r in rs: + # A_phi is constant around the loop (no phi dependence) + A_true_phi = _A_analytic(r) * np.ones(N) + A_true_rpz = np.vstack( + (np.zeros_like(A_true_phi), A_true_phi, np.zeros_like(A_true_phi)) + ).T + correct_flux = np.sum(r * A_true_phi * 2 * np.pi / N) + + curve = FourierXYZCurve( + X_n=[-r, 0, 0], Y_n=[0, 0, r], Z_n=[0, 0, 0] + ) # flux loop to integrate A over + + curve_data = curve.compute(["x", "x_s"], grid=curve_grid, basis="xyz") + curve_data_rpz = curve.compute(["x", "x_s"], grid=curve_grid, basis="rpz") + + grid_rpz = np.vstack( + [ + curve_data_rpz["x"][:, 0], + curve_data_rpz["x"][:, 1], + curve_data_rpz["x"][:, 2], + ] + ).T + grid_xyz = rpz2xyz(grid_rpz) + # FourierXYZCoil + coil = FourierXYZCoil(I, X_n=[-R, 0, 0], Y_n=[0, 0, R], Z_n=[0, 0, 0]) + test( + coil, + grid_xyz, + grid_rpz, + A_true_rpz, + correct_flux, + rtol=1e-8, + atol=1e-12, + ) + + # SplineXYZCoil + x = coil.compute("x", grid=coil_grid, basis="xyz")["x"] + coil = SplineXYZCoil(I, X=x[:, 0], Y=x[:, 1], Z=x[:, 2]) + test( + coil, + grid_xyz, + grid_rpz, + A_true_rpz, + correct_flux, + rtol=1e-4, + atol=1e-12, + ) + + # FourierPlanarCoil + coil = FourierPlanarCoil(I, center=[0, 0, 0], normal=[0, 0, -1], r_n=R) + test( + coil, + grid_xyz, + grid_rpz, + A_true_rpz, + correct_flux, + rtol=1e-8, + atol=1e-12, + ) + + # FourierRZCoil + coil = FourierRZCoil(I, R_n=np.array([R]), modes_R=np.array([0])) + test( + coil, + grid_xyz, + grid_rpz, + A_true_rpz, + correct_flux, + rtol=1e-8, + atol=1e-12, + ) + @pytest.mark.unit def test_properties(self): """Test getting/setting attributes for Coil class.""" diff --git a/tests/test_examples.py b/tests/test_examples.py index 6bae9dd16..a03e364c3 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1077,9 +1077,12 @@ def test_freeb_axisym(): -6.588300858364606e04, -3.560589388468855e05, ] - ext_field = SplineMagneticField.from_mgrid( - r"tests/inputs/mgrid_solovev.nc", extcur=extcur - ) + with pytest.warns(UserWarning, match="Vector potential"): + # the mgrid file does not have the vector potential + # saved so we will ignore the thrown warning + ext_field = SplineMagneticField.from_mgrid( + r"tests/inputs/mgrid_solovev.nc", extcur=extcur + ) pres = PowerSeriesProfile([1.25e-1, 0, -1.25e-1]) iota = PowerSeriesProfile([-4.9e-1, 0, 3.0e-1]) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 06e7b8380..ed0ae4317 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -6,10 +6,11 @@ from desc.backend import jit, jnp from desc.basis import DoubleFourierSeries -from desc.compute import rpz2xyz_vec, xyz2rpz_vec -from desc.compute.utils import get_params, get_transforms +from desc.compute import rpz2xyz, rpz2xyz_vec, xyz2rpz_vec +from desc.compute.utils import dot, get_params, get_transforms +from desc.derivatives import FiniteDiffDerivative as Derivative from desc.examples import get -from desc.geometry import FourierRZToroidalSurface +from desc.geometry import FourierRZToroidalSurface, FourierXYZCurve from desc.grid import LinearGrid from desc.io import load from desc.magnetic_fields import ( @@ -22,6 +23,7 @@ ScalarPotentialField, SplineMagneticField, ToroidalMagneticField, + VectorPotentialField, VerticalMagneticField, field_line_integrate, read_BNORM_file, @@ -59,8 +61,41 @@ def test_basic_fields(self): tfield = ToroidalMagneticField(2, 1) vfield = VerticalMagneticField(1) pfield = PoloidalMagneticField(2, 1, 2) + + def tfield_A(R, phi, Z, B0=2, R0=1): + az = -B0 * R0 * jnp.log(R) + arp = jnp.zeros_like(az) + A = jnp.array([arp, arp, az]).T + return A + + tfield_from_A = VectorPotentialField(tfield_A, params={"B0": 2, "R0": 1}) + + def vfield_A(R, phi, Z, B0=None): + coords_rpz = jnp.vstack([R, phi, Z]).T + coords_xyz = rpz2xyz(coords_rpz) + ax = B0 / 2 * coords_xyz[:, 1] + ay = -B0 / 2 * coords_xyz[:, 0] + + az = jnp.zeros_like(ax) + A = jnp.array([ax, -ay, az]).T + A = xyz2rpz_vec(A, phi=coords_rpz[:, 1]) + return A + + vfield_params = {"B0": 1} + vfield_from_A = VectorPotentialField(vfield_A, params=vfield_params) + np.testing.assert_allclose(tfield([1, 0, 0]), [[0, 2, 0]]) np.testing.assert_allclose((4 * tfield)([2, 0, 0]), [[0, 4, 0]]) + np.testing.assert_allclose(tfield_from_A([1, 0, 0]), [[0, 2, 0]]) + np.testing.assert_allclose( + tfield_A(1, 0, 0), + tfield_from_A.compute_magnetic_vector_potential([1, 0, 0]).squeeze(), + ) + np.testing.assert_allclose( + vfield_A(1, 0, 0, **vfield_params), + vfield_from_A.compute_magnetic_vector_potential([1, 0, 0]), + ) + np.testing.assert_allclose((tfield + vfield)([1, 0, 0]), [[0, 2, 1]]) np.testing.assert_allclose( (tfield + vfield - pfield)([1, 0, 0.1]), [[0.4, 2, 1]] @@ -104,17 +139,40 @@ def test_combined_fields(self): assert scaled_field.B0 == 2 assert scaled_field.scale == 3.1 np.testing.assert_allclose(scaled_field([1.0, 0, 0]), np.array([[0, 6.2, 0]])) + np.testing.assert_allclose( + scaled_field.compute_magnetic_vector_potential([2.0, 0, 0]), + np.array([[0, 0, -3.1 * 2 * 1 * np.log(2)]]), + ) + scaled_field.R0 = 1.3 scaled_field.scale = 1.0 np.testing.assert_allclose(scaled_field([1.3, 0, 0]), np.array([[0, 2, 0]])) + np.testing.assert_allclose( + scaled_field.compute_magnetic_vector_potential([2.0, 0, 0]), + np.array([[0, 0, -2 * 1.3 * np.log(2)]]), + ) assert scaled_field.optimizable_params == ["B0", "R0", "scale"] assert hasattr(scaled_field, "B0") sum_field = vfield + pfield + tfield + sum_field_tv = vfield + tfield # to test A since pfield does not have A assert len(sum_field) == 3 + assert len(sum_field_tv) == 2 + np.testing.assert_allclose( sum_field([1.3, 0, 0.0]), [[0.0, 2, 3.2 + 2 * 1.2 * 0.3]] ) + + tfield_A = np.array([[0, 0, -tfield.B0 * tfield.R0 * np.log(tfield.R0)]]) + x = tfield.R0 * np.cos(np.pi / 4) + y = tfield.R0 * np.sin(np.pi / 4) + vfield_A = np.array([[vfield.B0 * y, -vfield.B0 * x, 0]]) / 2 + + np.testing.assert_allclose( + sum_field_tv.compute_magnetic_vector_potential([x, y, 0.0], basis="xyz"), + tfield_A + vfield_A, + ) + assert sum_field.optimizable_params == [ ["B0"], ["B0", "R0", "iota"], @@ -304,6 +362,87 @@ def test_current_potential_field(self): with pytest.raises(AssertionError): field.potential_dzeta = 1 + @pytest.mark.unit + def test_current_potential_vector_potential(self): + """Test current potential field vector potential against analytic result.""" + R0 = 10 + a = 1 + surface = FourierRZToroidalSurface( + R_lmn=jnp.array([R0, a]), + Z_lmn=jnp.array([0, -a]), + modes_R=jnp.array([[0, 0], [1, 0]]), + modes_Z=jnp.array([[0, 0], [-1, 0]]), + NFP=10, + ) + # make a current potential corresponding a purely poloidal current + G = 100 # net poloidal current + potential = lambda theta, zeta, G: G * zeta / 2 / jnp.pi + potential_dtheta = lambda theta, zeta, G: jnp.zeros_like(theta) + potential_dzeta = lambda theta, zeta, G: G * jnp.ones_like(theta) / 2 / jnp.pi + + params = {"G": -G} + + field = CurrentPotentialField( + potential, + R_lmn=surface.R_lmn, + Z_lmn=surface.Z_lmn, + modes_R=surface._R_basis.modes[:, 1:], + modes_Z=surface._Z_basis.modes[:, 1:], + params=params, + potential_dtheta=potential_dtheta, + potential_dzeta=potential_dzeta, + NFP=surface.NFP, + ) + # test the loop integral of A around a curve encompassing the torus + # against the analytic result for flux in an ideal toroidal solenoid + prefactors = mu_0 * G / 2 / jnp.pi + correct_flux = -2 * np.pi * prefactors * (np.sqrt(R0**2 - a**2) - R0) + + curve = FourierXYZCurve() # curve to integrate A over + curve_grid = LinearGrid(zeta=20) + curve_data = curve.compute(["x", "x_s"], grid=curve_grid, basis="xyz") + curve_data_rpz = curve.compute(["x", "x_s"], grid=curve_grid, basis="rpz") + + surface_grid = LinearGrid(M=60, N=60, NFP=10) + + A_xyz = field.compute_magnetic_vector_potential( + curve_data["x"], basis="xyz", source_grid=surface_grid + ) + A_rpz = field.compute_magnetic_vector_potential( + curve_data_rpz["x"], basis="rpz", source_grid=surface_grid + ) + + # integrate to get the flux + flux_xyz = jnp.sum( + dot(A_xyz, curve_data["x_s"], axis=-1) * curve_grid.spacing[:, 2] + ) + flux_rpz = jnp.sum( + dot(A_rpz, curve_data_rpz["x_s"], axis=-1) * curve_grid.spacing[:, 2] + ) + + np.testing.assert_allclose(correct_flux, flux_xyz, rtol=1e-8) + np.testing.assert_allclose(correct_flux, flux_rpz, rtol=1e-8) + + field.params["G"] = -2 * field.params["G"] + + A_xyz = field.compute_magnetic_vector_potential( + curve_data["x"], basis="xyz", source_grid=surface_grid + ) + A_rpz = field.compute_magnetic_vector_potential( + curve_data_rpz["x"], basis="rpz", source_grid=surface_grid + ) + + # integrate to get the flux + flux_xyz = jnp.sum( + dot(A_xyz, curve_data["x_s"], axis=-1) * curve_grid.spacing[:, 2] + ) + flux_rpz = jnp.sum( + dot(A_rpz, curve_data_rpz["x_s"], axis=-1) * curve_grid.spacing[:, 2] + ) + + np.testing.assert_allclose(-2 * correct_flux, flux_xyz, rtol=1e-8) + np.testing.assert_allclose(-2 * correct_flux, flux_rpz, rtol=1e-8) + @pytest.mark.unit def test_fourier_current_potential_field(self): """Test Fourier current potential magnetic field against analytic result.""" @@ -416,6 +555,124 @@ def test_fourier_current_potential_field(self): atol=1e-16, ) + @pytest.mark.unit + def test_fourier_current_potential_vector_potential(self): + """Test Fourier current potential vector potential against analytic result.""" + R0 = 10 + a = 1 + surface = FourierRZToroidalSurface( + R_lmn=jnp.array([R0, a]), + Z_lmn=jnp.array([0, -a]), + modes_R=jnp.array([[0, 0], [1, 0]]), + modes_Z=jnp.array([[0, 0], [-1, 0]]), + NFP=10, + ) + + basis = DoubleFourierSeries(M=2, N=2, sym="sin") + phi_mn = np.ones((basis.num_modes,)) + # make a current potential corresponding a purely poloidal current + G = 100 # net poloidal current + + # test the loop integral of A around a curve encompassing the torus + # against the analytic result for flux in an ideal toroidal solenoid + ## expression for flux inside of toroidal solenoid of radius a + prefactors = mu_0 * G / 2 / jnp.pi + correct_flux = -2 * np.pi * prefactors * (np.sqrt(R0**2 - a**2) - R0) + + curve = FourierXYZCurve() # curve to integrate A over + curve_grid = LinearGrid(zeta=20) + curve_data = curve.compute(["x", "x_s"], grid=curve_grid) + curve_data_rpz = curve.compute(["x", "x_s"], grid=curve_grid, basis="rpz") + + field = FourierCurrentPotentialField( + Phi_mn=phi_mn, + modes_Phi=basis.modes[:, 1:], + I=0, + G=-G, # to get a positive B_phi, we must put G negative + # since -G is the net poloidal current on the surface + # ( with G=-(net_current) meaning that we have net_current + # flowing poloidally (in clockwise direction) around torus) + sym_Phi="sin", + R_lmn=surface.R_lmn, + Z_lmn=surface.Z_lmn, + modes_R=surface._R_basis.modes[:, 1:], + modes_Z=surface._Z_basis.modes[:, 1:], + NFP=10, + ) + surface_grid = LinearGrid(M=60, N=60, NFP=10) + + phi_mn = np.zeros((basis.num_modes,)) + + field.Phi_mn = phi_mn + + field.change_resolution(3, 3) + field.change_Phi_resolution(2, 2) + + A_xyz = field.compute_magnetic_vector_potential( + curve_data["x"], basis="xyz", source_grid=surface_grid + ) + A_rpz = field.compute_magnetic_vector_potential( + curve_data_rpz["x"], basis="rpz", source_grid=surface_grid + ) + + # integrate to get the flux + flux_xyz = jnp.sum( + dot(A_xyz, curve_data["x_s"], axis=-1) * curve_grid.spacing[:, 2] + ) + flux_rpz = jnp.sum( + dot(A_rpz, curve_data_rpz["x_s"], axis=-1) * curve_grid.spacing[:, 2] + ) + + np.testing.assert_allclose(correct_flux, flux_xyz, rtol=1e-8) + np.testing.assert_allclose(correct_flux, flux_rpz, rtol=1e-8) + + field.G = -2 * field.G + field.I = 0 + + A_xyz = field.compute_magnetic_vector_potential( + curve_data["x"], basis="xyz", source_grid=surface_grid + ) + A_rpz = field.compute_magnetic_vector_potential( + curve_data_rpz["x"], basis="rpz", source_grid=surface_grid + ) + + # integrate to get the flux + flux_xyz = jnp.sum( + dot(A_xyz, curve_data["x_s"], axis=-1) * curve_grid.spacing[:, 2] + ) + flux_rpz = jnp.sum( + dot(A_rpz, curve_data_rpz["x_s"], axis=-1) * curve_grid.spacing[:, 2] + ) + + np.testing.assert_allclose(-2 * correct_flux, flux_xyz, rtol=1e-8) + np.testing.assert_allclose(-2 * correct_flux, flux_rpz, rtol=1e-8) + + field = FourierCurrentPotentialField.from_surface( + surface=surface, + Phi_mn=phi_mn, + modes_Phi=basis.modes[:, 1:], + I=0, + G=-G, + ) + + A_xyz = field.compute_magnetic_vector_potential( + curve_data["x"], basis="xyz", source_grid=surface_grid + ) + A_rpz = field.compute_magnetic_vector_potential( + curve_data_rpz["x"], basis="rpz", source_grid=surface_grid + ) + + # integrate to get the flux + flux_xyz = jnp.sum( + dot(A_xyz, curve_data["x_s"], axis=-1) * curve_grid.spacing[:, 2] + ) + flux_rpz = jnp.sum( + dot(A_rpz, curve_data_rpz["x_s"], axis=-1) * curve_grid.spacing[:, 2] + ) + + np.testing.assert_allclose(correct_flux, flux_xyz, rtol=1e-8) + np.testing.assert_allclose(correct_flux, flux_rpz, rtol=1e-8) + @pytest.mark.unit def test_fourier_current_potential_field_symmetry(self): """Test Fourier current potential magnetic field Phi symmetry logic.""" @@ -644,7 +901,7 @@ def test_init_Phi_mn_fourier_current_field(self): @pytest.mark.slow @pytest.mark.unit - def test_spline_field(self): + def test_spline_field(self, tmpdir_factory): """Test accuracy of spline magnetic field.""" field1 = ScalarPotentialField(phi_lm, args) R = np.linspace(0.5, 1.5, 20) @@ -659,10 +916,65 @@ def test_spline_field(self): extcur = [4700.0, 1000.0] mgrid = "tests/inputs/mgrid_test.nc" field3 = SplineMagneticField.from_mgrid(mgrid, extcur) + # test saving and loading from mgrid + tmpdir = tmpdir_factory.mktemp("spline_mgrid_with_A") + path = tmpdir.join("spline_mgrid_with_A.nc") + field3.save_mgrid( + path, + Rmin=np.min(field3._R), + Rmax=np.max(field3._R), + Zmin=np.min(field3._Z), + Zmax=np.max(field3._Z), + nR=field3._R.size, + nZ=field3._Z.size, + nphi=field3._phi.size, + ) + # no need for extcur b/c is saved in "raw" format, no need to scale again + field4 = SplineMagneticField.from_mgrid(path) + attrs_4d = ["_AR", "_Aphi", "_AZ", "_BR", "_Bphi", "_BZ"] + for attr in attrs_4d: + np.testing.assert_allclose( + (getattr(field3, attr) * np.array(extcur)).sum(axis=-1), + getattr(field4, attr).squeeze(), + err_msg=attr, + ) + attrs_3d = ["_R", "_phi", "_Z"] + for attr in attrs_3d: + np.testing.assert_allclose(getattr(field3, attr), getattr(field4, attr)) + + r = 0.70 + p = 0 + z = 0 + # use finite diff derivatives to check A accuracy + tfield_A = lambda R, phi, Z: field3.compute_magnetic_vector_potential( + jnp.vstack([R, phi, Z]).T + ) + funR = lambda x: tfield_A(x, p, z) + funP = lambda x: tfield_A(r, x, z) + funZ = lambda x: tfield_A(r, p, x) + + ap = tfield_A(r, p, z)[:, 1] + + # these are the gradients of each component of A + dAdr = Derivative.compute_jvp(funR, 0, (jnp.ones_like(r),), r) + dAdp = Derivative.compute_jvp(funP, 0, (jnp.ones_like(p),), p) + dAdz = Derivative.compute_jvp(funZ, 0, (jnp.ones_like(z),), z) + + # form the B components with the appropriate combinations + B2 = jnp.array( + [ + dAdp[:, 2] / r - dAdz[:, 1], + dAdz[:, 0] - dAdr[:, 2], + dAdr[:, 1] + (ap - dAdp[:, 0]) / r, + ] + ).T np.testing.assert_allclose( field3([0.70, 0, 0]), np.array([[0, -0.671, 0.0858]]), rtol=1e-3, atol=1e-8 ) + + np.testing.assert_allclose(field3([0.70, 0, 0]), B2, rtol=1e-3, atol=5e-3) + field3.currents *= 2 np.testing.assert_allclose( field3([0.70, 0, 0]), @@ -697,14 +1009,20 @@ def test_spline_field_axisym(self): -2.430716e04, -2.380229e04, ] - field = SplineMagneticField.from_mgrid( - "tests/inputs/mgrid_d3d.nc", extcur=extcur - ) + with pytest.warns(UserWarning): + # user warning because saved mgrid no vector potential + field = SplineMagneticField.from_mgrid( + "tests/inputs/mgrid_d3d.nc", extcur=extcur + ) # make sure field is invariant to shift in phi B1 = field.compute_magnetic_field(np.array([1.75, 0.0, 0.0])) B2 = field.compute_magnetic_field(np.array([1.75, 1.0, 0.0])) np.testing.assert_allclose(B1, B2) + # test the error when no vec pot values exist + with pytest.raises(ValueError, match="no vector potential"): + field.compute_magnetic_vector_potential(np.array([1.75, 0.0, 0.0])) + @pytest.mark.unit def test_field_line_integrate(self): """Test field line integration.""" @@ -842,8 +1160,15 @@ def test_mgrid_io(self, tmpdir_factory): Rmax = 7 Zmin = -2 Zmax = 2 - save_field.save_mgrid(path, Rmin, Rmax, Zmin, Zmax) - load_field = SplineMagneticField.from_mgrid(path) + with pytest.raises(NotImplementedError): + # Raises error because poloidal field has no vector potential + # and so cannot save the vector potential + save_field.save_mgrid(path, Rmin, Rmax, Zmin, Zmax) + save_field.save_mgrid(path, Rmin, Rmax, Zmin, Zmax, save_vector_potential=False) + with pytest.warns(UserWarning): + # user warning because saved mgrid has no vector potential + # and so cannot load the vector potential + load_field = SplineMagneticField.from_mgrid(path) # check that the fields are the same num_nodes = 50 diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 1619e5c84..4711907ab 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -30,6 +30,7 @@ from desc.magnetic_fields import ( FourierCurrentPotentialField, OmnigenousField, + PoloidalMagneticField, SplineMagneticField, ToroidalMagneticField, VerticalMagneticField, @@ -1160,10 +1161,14 @@ def test_quadratic_flux(self): @pytest.mark.unit def test_toroidal_flux(self): """Test calculation of toroidal flux from coils.""" - grid1 = LinearGrid(L=10, M=10, zeta=np.array(0.0)) + grid1 = LinearGrid(L=0, M=40, zeta=np.array(0.0)) def test(eq, field, correct_value, rtol=1e-14, grid=None): - obj = ToroidalFlux(eq=eq, field=field, eval_grid=grid) + obj = ToroidalFlux( + eq=eq, + field=field, + eval_grid=grid, + ) obj.build(verbose=2) torflux = obj.compute_unscaled(*obj.xs(field)) np.testing.assert_allclose(torflux, correct_value, rtol=rtol) @@ -1173,22 +1178,20 @@ def test(eq, field, correct_value, rtol=1e-14, grid=None): field = ToroidalMagneticField(B0=1, R0=1) # calc field Psi - data = eq.compute(["R", "phi", "Z", "|e_rho x e_theta|", "n_zeta"], grid=grid1) - field_B = field.compute_magnetic_field( + data = eq.compute(["R", "phi", "Z", "e_theta"], grid=grid1) + field_A = field.compute_magnetic_vector_potential( np.vstack([data["R"], data["phi"], data["Z"]]).T ) - B_dot_n_zeta = jnp.sum(field_B * data["n_zeta"], axis=1) + A_dot_e_theta = jnp.sum(field_A * data["e_theta"], axis=1) - psi_from_field = np.sum( - grid1.spacing[:, 0] - * grid1.spacing[:, 1] - * data["|e_rho x e_theta|"] - * B_dot_n_zeta - ) - eq.change_resolution(L_grid=10, M_grid=10) + psi_from_field = np.sum(grid1.spacing[:, 1] * A_dot_e_theta) + eq.change_resolution(L_grid=20, M_grid=20) test(eq, field, psi_from_field) + test(eq, field, psi_from_field, rtol=1e-3) + # test on field with no vector potential + test(eq, PoloidalMagneticField(1, 1, 1), 0.0) @pytest.mark.unit def test_signed_plasma_vessel_distance(self): @@ -2275,7 +2278,9 @@ def test_compute_scalar_resolution_heating_power(self): @pytest.mark.regression def test_compute_scalar_resolution_boundary_error(self): """BoundaryError.""" - ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") + with pytest.warns(UserWarning): + # user warning because saved mgrid no vector potential + ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") pres = PowerSeriesProfile([1.25e-1, 0, -1.25e-1]) iota = PowerSeriesProfile([-4.9e-1, 0, 3.0e-1]) @@ -2301,7 +2306,9 @@ def test_compute_scalar_resolution_boundary_error(self): @pytest.mark.regression def test_compute_scalar_resolution_vacuum_boundary_error(self): """VacuumBoundaryError.""" - ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") + with pytest.warns(UserWarning): + # user warning because saved mgrid no vector potential + ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") pres = PowerSeriesProfile([1.25e-1, 0, -1.25e-1]) iota = PowerSeriesProfile([-4.9e-1, 0, 3.0e-1]) @@ -2328,7 +2335,8 @@ def test_compute_scalar_resolution_vacuum_boundary_error(self): @pytest.mark.regression def test_compute_scalar_resolution_quadratic_flux(self): """VacuumBoundaryError.""" - ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") + with pytest.warns(UserWarning): + ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") pres = PowerSeriesProfile([1.25e-1, 0, -1.25e-1]) iota = PowerSeriesProfile([-4.9e-1, 0, 3.0e-1]) @@ -2352,7 +2360,25 @@ def test_compute_scalar_resolution_quadratic_flux(self): np.testing.assert_allclose(f, f[-1], rtol=5e-2) @pytest.mark.regression - def test_compute_scalar_resolution_toroidal_flux(self): + def test_compute_scalar_resolution_toroidal_flux_A(self): + """ToroidalFlux.""" + ext_field = ToroidalMagneticField(1, 1) + eq = get("precise_QA") + with pytest.warns(UserWarning, match="Reducing radial"): + eq.change_resolution(4, 4, 4, 8, 8, 8) + + f = np.zeros_like(self.res_array, dtype=float) + for i, res in enumerate(self.res_array): + eq.change_resolution( + L_grid=int(eq.L * res), M_grid=int(eq.M * res), N_grid=int(eq.N * res) + ) + obj = ObjectiveFunction(ToroidalFlux(eq, ext_field), use_jit=False) + obj.build(verbose=0) + f[i] = obj.compute_scalar(obj.x()) + np.testing.assert_allclose(f, f[-1], rtol=5e-2) + + @pytest.mark.regression + def test_compute_scalar_resolution_toroidal_flux_B(self): """ToroidalFlux.""" ext_field = ToroidalMagneticField(1, 1) eq = get("precise_QA") @@ -2626,7 +2652,9 @@ def test_objective_no_nangrad_heating_power(self): @pytest.mark.unit def test_objective_no_nangrad_boundary_error(self): """BoundaryError.""" - ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") + with pytest.warns(UserWarning): + # user warning because saved mgrid no vector potential + ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") pres = PowerSeriesProfile([1.25e-1, 0, -1.25e-1]) iota = PowerSeriesProfile([-4.9e-1, 0, 3.0e-1]) @@ -2647,7 +2675,9 @@ def test_objective_no_nangrad_boundary_error(self): @pytest.mark.unit def test_objective_no_nangrad_vacuum_boundary_error(self): """VacuumBoundaryError.""" - ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") + with pytest.warns(UserWarning): + # user warning because saved mgrid no vector potential + ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") pres = PowerSeriesProfile([1.25e-1, 0, -1.25e-1]) iota = PowerSeriesProfile([-4.9e-1, 0, 3.0e-1]) @@ -2670,7 +2700,9 @@ def test_objective_no_nangrad_vacuum_boundary_error(self): @pytest.mark.unit def test_objective_no_nangrad_quadratic_flux(self): """QuadraticFlux.""" - ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") + with pytest.warns(UserWarning): + # user warning because saved mgrid no vector potential + ext_field = SplineMagneticField.from_mgrid(r"tests/inputs/mgrid_solovev.nc") pres = PowerSeriesProfile([1.25e-1, 0, -1.25e-1]) iota = PowerSeriesProfile([-4.9e-1, 0, 3.0e-1]) @@ -2701,7 +2733,12 @@ def test_objective_no_nangrad_toroidal_flux(self): obj = ObjectiveFunction(ToroidalFlux(eq, ext_field), use_jit=False) obj.build() g = obj.grad(obj.x(ext_field)) - assert not np.any(np.isnan(g)), "toroidal flux" + assert not np.any(np.isnan(g)), "toroidal flux A" + + obj = ObjectiveFunction(ToroidalFlux(eq, ext_field), use_jit=False) + obj.build() + g = obj.grad(obj.x(ext_field)) + assert not np.any(np.isnan(g)), "toroidal flux B" @pytest.mark.unit @pytest.mark.parametrize(