From a0ded344ef438ac0dbe9d93f3650ac7acb1d6b33 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 20 May 2024 12:22:03 -0400 Subject: [PATCH 01/43] add biot savart vector potential for general and coils and add test for coil vector potential --- desc/coils.py | 97 +++++++++++++++++++++++++++++++++++ desc/magnetic_fields/_core.py | 34 ++++++++++++ tests/test_coils.py | 97 +++++++++++++++++++++++++++++++++++ 3 files changed, 228 insertions(+) diff --git a/desc/coils.py b/desc/coils.py index ab7f0dc2b3..b1fe04eb62 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -112,6 +112,48 @@ 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. + + 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. + + Notes + ----- + #FIXME: what is sacrificed for A? that curl(A) != B exactly? + This method does not give curl(B) == 0 exactly. The error in the curl + scales the same as the error in B itself, so will only be zero when fully + converged. However in practice, for smooth curves described by Fourier series, + this method converges exponentially in the number of coil 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. @@ -201,6 +243,61 @@ def compute_magnetic_field( B = xyz2rpz_vec(B, phi=phi) 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. + + 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. + + 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. + + """ + assert basis.lower() in ["rpz", "xyz"] + coords = jnp.atleast_2d(jnp.asarray(coords)) + if basis.lower() == "rpz": + phi = coords[:, 1] + coords = rpz2xyz(coords) + if params is None: + current = self.current + else: + current = params.pop("current", self.current) + + data = self.compute( + ["x", "x_s", "ds"], grid=source_grid, params=params, basis="xyz" + ) + A = biot_savart_vector_potential_quad( + coords, data["x"], data["x_s"] * data["ds"][:, None], current + ) + + if basis.lower() == "rpz": + A = xyz2rpz_vec(A, phi=phi) + return A + def __repr__(self): """Get the string form of the object.""" return ( diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 6900d9cbd5..0f730dc432 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -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. diff --git a/tests/test_coils.py b/tests/test_coils.py index 5e1518bbf8..21e7bfa847 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -128,6 +128,103 @@ 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) + + # FourierXYZCoil + coil = FourierXYZCoil(I) + 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="Using FourierXYZCoil" + ) + np.testing.assert_allclose( + A_true, A_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierXYZCoil" + ) + np.testing.assert_allclose( + A_true, A_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierXYZCoil" + ) + + # SplineXYZCoil + x = coil.compute("x", grid=coil_grid, basis="xyz")["x"] + coil = SplineXYZCoil(I, X=x[:, 0], Y=x[:, 1], Z=x[:, 2]) + 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="Using SplineXYZCoil" + ) + np.testing.assert_allclose( + A_true, A_rpz, rtol=1e-3, atol=1e-10, err_msg="Using SplineXYZCoil" + ) + np.testing.assert_allclose( + A_true, A_rpz, rtol=1e-3, atol=1e-10, err_msg="Using SplineXYZCoil" + ) + + # FourierPlanarCoil + coil = FourierPlanarCoil(I) + 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="Using FourierPlanarCoil" + ) + np.testing.assert_allclose( + A_true, + A_rpz, + rtol=1e-3, + atol=1e-10, + err_msg="Using FourierPlanarCoil", + ) + np.testing.assert_allclose( + A_true, + A_rpz, + rtol=1e-3, + atol=1e-10, + err_msg="Using FourierPlanarCoil", + ) + + 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])) + 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="Using FourierRZCoil" + ) + np.testing.assert_allclose( + A_true, A_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierRZCoil" + ) + np.testing.assert_allclose( + A_true, A_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierRZCoil" + ) + @pytest.mark.unit def test_properties(self): """Test getting/setting attributes for Coil class.""" From 483ddec60590b04e18226e9af848b2284ce6c910 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 20 May 2024 15:01:45 -0400 Subject: [PATCH 02/43] add hanson hirshman vector pot calc for splinexyzcoil --- desc/coils.py | 120 ++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 112 insertions(+), 8 deletions(-) diff --git a/desc/coils.py b/desc/coils.py index b1fe04eb62..6c64dfc2eb 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -71,6 +71,52 @@ 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. + + 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) + + 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 = x_f - x_i + A = jnp.sum(A_mag[:, :, jnp.newaxis] * d_vec[:, 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. @@ -133,14 +179,6 @@ def biot_savart_vector_potential_quad(eval_pts, coil_pts, tangents, current): ------- A : ndarray, shape(n,3) Magnetic vector potential in cartesian components at specified points. - - Notes - ----- - #FIXME: what is sacrificed for A? that curl(A) != B exactly? - This method does not give curl(B) == 0 exactly. The error in the curl - scales the same as the error in B itself, so will only be zero when fully - converged. However in practice, for smooth curves described by Fourier series, - this method converges exponentially in the number of coil points. """ dl = tangents R_vec = eval_pts[jnp.newaxis, :] - coil_pts[:, jnp.newaxis, :] @@ -726,6 +764,72 @@ def compute_magnetic_field( B = xyz2rpz_vec(B, x=coords[:, 0], y=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. + + 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. + + 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. + + """ + assert basis.lower() in ["rpz", "xyz"] + coords = jnp.atleast_2d(jnp.asarray(coords)) + if basis == "rpz": + coords = rpz2xyz(coords) + if params is None: + current = self.current + else: + current = params.pop("current", self.current) + + data = self.compute(["x"], grid=source_grid, params=params, basis="xyz") + # need to make sure the curve is closed. If it's already closed, this doesn't + # do anything (effectively just adds a segment of zero length which has no + # effect on the overall result) + coil_pts_start = data["x"] + coil_pts_end = jnp.concatenate([data["x"][1:], data["x"][:1]]) + # could get up to 4th order accuracy by shifting points outward as in + # (McGreivy, Zhu, Gunderson, Hudson 2021), however that requires knowing the + # 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_vector_potential_hh( + coords, coil_pts_start, coil_pts_end, current + ) + + if basis == "rpz": + B = xyz2rpz_vec(B, x=coords[:, 0], y=coords[:, 1]) + return B + @classmethod def from_values( cls, current, coords, knots=None, method="cubic", name="", basis="xyz" From c5e2c3f3490f860cc2dcbfd6163396573ea9d3c7 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 20 May 2024 17:29:49 -0400 Subject: [PATCH 03/43] add vector potential function for fourier current potential and add test against analytic flux --- desc/magnetic_fields/_current_potential.py | 113 ++++++++++++++++++- tests/test_magnetic_fields.py | 121 ++++++++++++++++++++- 2 files changed, 232 insertions(+), 2 deletions(-) diff --git a/desc/magnetic_fields/_current_potential.py b/desc/magnetic_fields/_current_potential.py index 7624cd21de..402d6b7dbb 100644 --- a/desc/magnetic_fields/_current_potential.py +++ b/desc/magnetic_fields/_current_potential.py @@ -10,7 +10,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): @@ -527,6 +531,41 @@ def compute_magnetic_field( source_grid=source_grid, ) + 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 + 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. + + Returns + ------- + A : ndarray, shape(N,3) + Magnetic vector potential at specified points. + + """ + source_grid = source_grid or LinearGrid( + M=30 + 2 * max(self.M, self.M_Phi), + N=30 + 2 * max(self.N, self.N_Phi), + NFP=self.NFP, + ) + return _compute_vector_potential_from_CurrentPotentialField( + field=self, + coords=coords, + params=params, + basis=basis, + source_grid=source_grid, + ) + @classmethod def from_surface( cls, @@ -676,3 +715,75 @@ def nfp_loop(j, f): if basis == "rpz": B = xyz2rpz_vec(B, x=coords[:, 0], y=coords[:, 1]) return B + + +def _compute_vector_potential_from_CurrentPotentialField( + field, + coords, + source_grid, + params=None, + basis="rpz", +): + """Compute magnetic vector potential at a set of points. + + Parameters + ---------- + field : CurrentPotentialField or FourierCurrentPotentialField + current potential field object from which to compute magnetic vector potential. + coords : array-like shape(N,3) + cylindrical or cartesian coordinates + source_grid : Grid, + source grid upon which to evaluate the surface current density K + params : dict, optional + parameters to pass to compute function + should include the potential + basis : {"rpz", "xyz"} + basis for input coordinates and returned magnetic vector potential + + + Returns + ------- + A : ndarray, shape(N,3) + magnetic vector potential at specified points + + """ + assert basis.lower() in ["rpz", "xyz"] + coords = jnp.atleast_2d(jnp.asarray(coords)) + if basis == "rpz": + coords = rpz2xyz(coords) + + # compute surface current, and store grid quantities + # needed for integration in class + # TODO: does this have to be xyz, or can it be computed in rpz as well? + data = field.compute(["K", "x"], grid=source_grid, basis="xyz", params=params) + + _rs = xyz2rpz(data["x"]) + _K = xyz2rpz_vec(data["K"], phi=source_grid.nodes[:, 2]) + + # surface element, must divide by NFP to remove the NFP multiple on the + # surface grid weights, as we account for that when doing the for loop + # over NFP + _dV = source_grid.weights * data["|e_theta x e_zeta|"] / source_grid.NFP + + def nfp_loop(j, f): + # calculate (by rotating) rs, rs_t, rz_t + phi = (source_grid.nodes[:, 2] + j * 2 * jnp.pi / source_grid.NFP) % ( + 2 * jnp.pi + ) + # new coords are just old R,Z at a new phi (bc of discrete NFP symmetry) + rs = jnp.vstack((_rs[:, 0], phi, _rs[:, 2])).T + rs = rpz2xyz(rs) + K = rpz2xyz_vec(_K, phi=phi) + fj = biot_savart_general_vector_potential( + coords, + rs, + K, + _dV, + ) + f += fj + return f + + A = fori_loop(0, source_grid.NFP, nfp_loop, jnp.zeros_like(coords)) + if basis == "rpz": + A = xyz2rpz_vec(A, x=coords[:, 0], y=coords[:, 1]) + return A diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 9e4103f0dc..a763d6bb4f 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -7,8 +7,9 @@ from desc.backend import jit, jnp from desc.basis import DoubleFourierSeries from desc.compute import rpz2xyz_vec, xyz2rpz_vec +from desc.compute.utils import dot 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 ( @@ -407,6 +408,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.""" From a10028fcff34c459875f18591a3b91b6d7954efc Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 20 May 2024 18:08:48 -0400 Subject: [PATCH 04/43] add vector potential computation to the current potential field class --- desc/magnetic_fields/_current_potential.py | 35 +++++++++ tests/test_magnetic_fields.py | 82 ++++++++++++++++++++++ 2 files changed, 117 insertions(+) diff --git a/desc/magnetic_fields/_current_potential.py b/desc/magnetic_fields/_current_potential.py index 402d6b7dbb..b7cde66f8f 100644 --- a/desc/magnetic_fields/_current_potential.py +++ b/desc/magnetic_fields/_current_potential.py @@ -215,6 +215,41 @@ def compute_magnetic_field( source_grid=source_grid, ) + 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 + 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. + + Returns + ------- + A : ndarray, shape(N,3) + Magnetic vector potential at specified points. + + """ + source_grid = source_grid or LinearGrid( + M=30 + 2 * self.M, + N=30 + 2 * self.N, + NFP=self.NFP, + ) + return _compute_vector_potential_from_CurrentPotentialField( + field=self, + coords=coords, + params=params, + basis=basis, + source_grid=source_grid, + ) + @classmethod def from_surface( cls, diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index a763d6bb4f..c558c4114f 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -304,6 +304,88 @@ 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 + ## 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") + + 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.""" From 7bf6915cfd6a5047655dc8d9f3ff6251185ea76b Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 21 May 2024 00:28:57 -0400 Subject: [PATCH 05/43] add test for coils integrating over a flux loop, has revealed a bug in HH vector potential fxn --- desc/coils.py | 6 +- tests/test_coils.py | 168 +++++++++++++++++++++++++++++++++- tests/test_magnetic_fields.py | 2 +- 3 files changed, 170 insertions(+), 6 deletions(-) diff --git a/desc/coils.py b/desc/coils.py index 6c64dfc2eb..34f7e222e7 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -822,13 +822,13 @@ def compute_magnetic_vector_potential( # 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_vector_potential_hh( + A = biot_savart_vector_potential_hh( coords, coil_pts_start, coil_pts_end, current ) if basis == "rpz": - B = xyz2rpz_vec(B, x=coords[:, 0], y=coords[:, 1]) - return B + A = xyz2rpz_vec(A, x=coords[:, 0], y=coords[:, 1]) + return A @classmethod def from_values( diff --git a/tests/test_coils.py b/tests/test_coils.py index 21e7bfa847..4ecf9e5bbc 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 rpz2xyz_vec, xyz2rpz, xyz2rpz_vec +from desc.compute import rpz2xyz, rpz2xyz_vec, 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 LinearGrid from desc.magnetic_fields import SumMagneticField, VerticalMagneticField @@ -225,6 +228,167 @@ def test_biot_savart_vector_potential_all_coils(self): A_true, A_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierRZCoil" ) + @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=200, endpoint=False) + + R = 1 + I = 1e7 + + # analytic eqn for "A_phi" (I assume 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 (it is in spherical coords) + rs = np.linspace(0.1, 3, 10, endpoint=True) + N = 200 + curve_grid = LinearGrid(zeta=N) + for r in rs: + # A_phi is constant around the loop + 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=[0, 0, -r], Y_n=[r, 0, 0], 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=[0, 0, R], Y_n=[-R, 0, 0], Z_n=[0, 0, 0]) + 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=1e-8, err_msg="Using FourierXYZCoil" + ) + np.testing.assert_allclose( + correct_flux, flux_rpz, rtol=1e-8, err_msg="Using FourierXYZCoil" + ) + np.testing.assert_allclose( + A_true_rpz, + np.abs(A_rpz), + rtol=1e-8, + atol=1e-12, + err_msg="Using FourierXYZCoil", + ) + + # SplineXYZCoil + x = coil.compute("x", grid=coil_grid, basis="xyz")["x"] + coil = SplineXYZCoil(I, X=x[:, 0], Y=x[:, 1], Z=x[:, 2]) + 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] + ) + # FIXME: errors here, probably something is incorrect in the + # HH vector integral function + np.testing.assert_allclose( + correct_flux, flux_xyz, rtol=1e-8, err_msg="Using SplineXYZCoil" + ) + np.testing.assert_allclose( + correct_flux, flux_rpz, rtol=1e-8, err_msg="Using SplineXYZCoil" + ) + np.testing.assert_allclose( + A_true_rpz, + np.abs(A_rpz), + rtol=1e-8, + atol=1e-12, + err_msg="Using SplineXYZCoil", + ) + # FourierPlanarCoil + coil = FourierPlanarCoil(I, center=[0, 0, 0], normal=[0, 0, -1], r_n=R) + 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 + ) + # FIXME: have to negate even if normal is negative z^ such that the + # dl of the coil is the same as the above coils?? + 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, err_msg="Using FourierPlanarCoil" + ) + np.testing.assert_allclose( + correct_flux, flux_rpz, rtol=1e-8, err_msg="Using FourierPlanarCoil" + ) + + # FourierRZCoil + coil = FourierRZCoil(I, R_n=np.array([R]), modes_R=np.array([0])) + 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 + ) + # FIXME: have to negate even if normal is negative z^ such that the + # dl of the coil is the same as the above coils?? + 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, err_msg="Using FourierRZCoil" + ) + np.testing.assert_allclose( + correct_flux, flux_rpz, rtol=1e-8, err_msg="Using FourierRZCoil" + ) + @pytest.mark.unit def test_properties(self): """Test getting/setting attributes for Coil class.""" diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index c558c4114f..274ebe4267 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -343,7 +343,7 @@ def test_current_potential_vector_potential(self): curve = FourierXYZCurve() # curve to integrate A over curve_grid = LinearGrid(zeta=20) - curve_data = curve.compute(["x", "x_s"], grid=curve_grid) + 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) From 34ed85f53dba815eee67e19b5a59096aaea205cd Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 21 May 2024 00:34:23 -0400 Subject: [PATCH 06/43] fix error in HH vector potential integral function --- desc/coils.py | 5 +++-- tests/test_coils.py | 24 ++++++++++++++++++------ 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/desc/coils.py b/desc/coils.py index 34f7e222e7..2722b25aaa 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -100,6 +100,7 @@ def biot_savart_vector_potential_hh(eval_pts, coil_pts_start, coil_pts_end, curr """ 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) @@ -112,8 +113,8 @@ def biot_savart_vector_potential_hh(eval_pts, coil_pts_start, coil_pts_end, curr 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 = x_f - x_i - A = jnp.sum(A_mag[:, :, jnp.newaxis] * d_vec[:, jnp.newaxis, :], axis=0) + # 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 diff --git a/tests/test_coils.py b/tests/test_coils.py index 4ecf9e5bbc..d8934b3eee 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -235,7 +235,7 @@ def test_biot_savart_vector_potential_integral_all_coils(self): # "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=200, endpoint=False) + coil_grid = LinearGrid(zeta=1000, endpoint=False) R = 1 I = 1e7 @@ -326,18 +326,16 @@ def _A_analytic(r): flux_rpz = jnp.sum( dot(A_rpz, curve_data_rpz["x_s"], axis=-1) * curve_grid.spacing[:, 2] ) - # FIXME: errors here, probably something is incorrect in the - # HH vector integral function np.testing.assert_allclose( - correct_flux, flux_xyz, rtol=1e-8, err_msg="Using SplineXYZCoil" + correct_flux, flux_xyz, rtol=1e-4, err_msg="Using SplineXYZCoil" ) np.testing.assert_allclose( - correct_flux, flux_rpz, rtol=1e-8, err_msg="Using SplineXYZCoil" + correct_flux, flux_rpz, rtol=1e-4, err_msg="Using SplineXYZCoil" ) np.testing.assert_allclose( A_true_rpz, np.abs(A_rpz), - rtol=1e-8, + rtol=1e-4, atol=1e-12, err_msg="Using SplineXYZCoil", ) @@ -364,6 +362,13 @@ def _A_analytic(r): np.testing.assert_allclose( correct_flux, flux_rpz, rtol=1e-8, err_msg="Using FourierPlanarCoil" ) + np.testing.assert_allclose( + A_true_rpz, + np.abs(A_rpz), + rtol=1e-8, + atol=1e-12, + err_msg="Using FourierPlanarCoil", + ) # FourierRZCoil coil = FourierRZCoil(I, R_n=np.array([R]), modes_R=np.array([0])) @@ -388,6 +393,13 @@ def _A_analytic(r): np.testing.assert_allclose( correct_flux, flux_rpz, rtol=1e-8, err_msg="Using FourierRZCoil" ) + np.testing.assert_allclose( + A_true_rpz, + np.abs(A_rpz), + rtol=1e-8, + atol=1e-12, + err_msg="Using FourierRZCoil", + ) @pytest.mark.unit def test_properties(self): From 94fce679a8eaa34340fb827bd894606c6e241ac9 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 21 May 2024 00:41:05 -0400 Subject: [PATCH 07/43] fix comment in test --- tests/test_coils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_coils.py b/tests/test_coils.py index d8934b3eee..9af7265cb6 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -378,7 +378,7 @@ def _A_analytic(r): A_rpz = coil.compute_magnetic_vector_potential( grid_rpz, basis="rpz", source_grid=coil_grid ) - # FIXME: have to negate even if normal is negative z^ such that the + # FIXME: have to negate even if oriented such that the # dl of the coil is the same as the above coils?? flux_xyz = -jnp.sum( dot(A_xyz, curve_data["x_s"], axis=-1) * curve_grid.spacing[:, 2] From 170b550e174943f4089a02286eecbda022b89dc2 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 21 May 2024 16:39:27 -0400 Subject: [PATCH 08/43] fix orientation confusion --- tests/test_coils.py | 32 +++++++++++++++----------------- 1 file changed, 15 insertions(+), 17 deletions(-) diff --git a/tests/test_coils.py b/tests/test_coils.py index 9af7265cb6..6a23e06934 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -240,7 +240,7 @@ def test_biot_savart_vector_potential_integral_all_coils(self): R = 1 I = 1e7 - # analytic eqn for "A_phi" (I assume phi is in dl direction for loop) + # 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. @@ -255,12 +255,12 @@ def _A_analytic(r): 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 (it is in spherical coords) + # 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) for r in rs: - # A_phi is constant around the loop + # 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)) @@ -268,7 +268,7 @@ def _A_analytic(r): correct_flux = np.sum(r * A_true_phi * 2 * np.pi / N) curve = FourierXYZCurve( - X_n=[0, 0, -r], Y_n=[r, 0, 0], Z_n=[0, 0, 0] + 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") @@ -283,7 +283,7 @@ def _A_analytic(r): ).T grid_xyz = rpz2xyz(grid_rpz) # FourierXYZCoil - coil = FourierXYZCoil(I, X_n=[0, 0, R], Y_n=[-R, 0, 0], Z_n=[0, 0, 0]) + coil = FourierXYZCoil(I, X_n=[-R, 0, 0], Y_n=[0, 0, R], Z_n=[0, 0, 0]) A_xyz = coil.compute_magnetic_vector_potential( grid_xyz, basis="xyz", source_grid=coil_grid ) @@ -305,7 +305,7 @@ def _A_analytic(r): ) np.testing.assert_allclose( A_true_rpz, - np.abs(A_rpz), + A_rpz, rtol=1e-8, atol=1e-12, err_msg="Using FourierXYZCoil", @@ -334,7 +334,7 @@ def _A_analytic(r): ) np.testing.assert_allclose( A_true_rpz, - np.abs(A_rpz), + A_rpz, rtol=1e-4, atol=1e-12, err_msg="Using SplineXYZCoil", @@ -347,12 +347,11 @@ def _A_analytic(r): A_rpz = coil.compute_magnetic_vector_potential( grid_rpz, basis="rpz", source_grid=coil_grid ) - # FIXME: have to negate even if normal is negative z^ such that the - # dl of the coil is the same as the above coils?? - flux_xyz = -jnp.sum( + + flux_xyz = jnp.sum( dot(A_xyz, curve_data["x_s"], axis=-1) * curve_grid.spacing[:, 2] ) - flux_rpz = -jnp.sum( + flux_rpz = jnp.sum( dot(A_rpz, curve_data_rpz["x_s"], axis=-1) * curve_grid.spacing[:, 2] ) @@ -364,7 +363,7 @@ def _A_analytic(r): ) np.testing.assert_allclose( A_true_rpz, - np.abs(A_rpz), + A_rpz, rtol=1e-8, atol=1e-12, err_msg="Using FourierPlanarCoil", @@ -378,12 +377,11 @@ def _A_analytic(r): A_rpz = coil.compute_magnetic_vector_potential( grid_rpz, basis="rpz", source_grid=coil_grid ) - # FIXME: have to negate even if oriented such that the - # dl of the coil is the same as the above coils?? - flux_xyz = -jnp.sum( + + flux_xyz = jnp.sum( dot(A_xyz, curve_data["x_s"], axis=-1) * curve_grid.spacing[:, 2] ) - flux_rpz = -jnp.sum( + flux_rpz = jnp.sum( dot(A_rpz, curve_data_rpz["x_s"], axis=-1) * curve_grid.spacing[:, 2] ) @@ -395,7 +393,7 @@ def _A_analytic(r): ) np.testing.assert_allclose( A_true_rpz, - np.abs(A_rpz), + A_rpz, rtol=1e-8, atol=1e-12, err_msg="Using FourierRZCoil", From acc63545525a7f01e849611ab406f2b7f5c1e9e8 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 21 May 2024 21:54:54 -0400 Subject: [PATCH 09/43] add vector potential calculation for analytic fields --- desc/magnetic_fields/_core.py | 204 +++++++++++++++++++++++++++++++++- tests/test_magnetic_fields.py | 21 +++- 2 files changed, 219 insertions(+), 6 deletions(-) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 0f730dc432..346ca9407d 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -239,6 +239,31 @@ 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 + ): + """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. + + Returns + ------- + A : ndarray, shape(N,3) + magnetic vector potential at specified points + + """ + def compute_Bnormal( self, surface, @@ -734,6 +759,33 @@ def compute_magnetic_field( coords, params, basis, source_grid ) + 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 + 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. + + 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. @@ -799,6 +851,50 @@ def compute_magnetic_field( 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 + 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. + + Returns + ------- + A : ndarray, shape(N,3) + scaled magnetic vector potential at specified points + + """ + if params is None: + params = [None] * len(self._fields) + if isinstance(params, dict): + params = [params] + if source_grid is None: + source_grid = [None] * len(self._fields) + if not isinstance(source_grid, (list, tuple)): + source_grid = [source_grid] + if len(source_grid) != len(self._fields): + # ensure that if source_grid is shorter, that it is simply repeated so that + # zip does not terminate early + source_grid = source_grid * len(self._fields) + + A = 0 + for i, (field, g) in enumerate(zip(self._fields, source_grid)): + A += field.compute_magnetic_vector_potential( + coords, params[i % len(params)], basis, source_grid=g + ) + + return A + # dunder methods required by MutableSequence def __getitem__(self, i): return self._fields[i] @@ -904,6 +1000,43 @@ def compute_magnetic_field( 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 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. + + 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. @@ -955,18 +1088,52 @@ 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 + ): + """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 + + """ + 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) + axy = (B0 / 2) * jnp.ones_like(coords[:, 2]) + az = jnp.zeros_like(axy) + A = jnp.array([axy, -axy, az]).T + if basis == "xyz": + A = rpz2xyz_vec(A, phi=coords[:, 1]) + + return A + class PoloidalMagneticField(_MagneticField, Optimizable): """Pure poloidal magnetic field (ie in theta direction). @@ -1074,6 +1241,33 @@ def compute_magnetic_field( 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( + "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. diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 274ebe4267..4ce4238e5c 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -104,17 +104,37 @@ 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]] ) + R = 1.3 + tfield_A = np.array([[0, 0, -2 * 1.3 * np.log(R)]]) + vfield_A = np.array([[vfield.B0, -vfield.B0, 0]]) / 2 + np.testing.assert_allclose( + sum_field_tv.compute_magnetic_vector_potential([R, 0, 0.0], basis="xyz"), + tfield_A + vfield_A, + ) + assert sum_field.optimizable_params == [ ["B0"], ["B0", "R0", "iota"], @@ -337,7 +357,6 @@ def test_current_potential_vector_potential(self): ) # 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) From dcf73f1a3fe536620fa657c45e36d0eb74960365 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 21 May 2024 22:08:01 -0400 Subject: [PATCH 10/43] add vector potential method with not implemented error for the fields that it does not make sense to compute A --- desc/magnetic_fields/_core.py | 85 +++++++++++++++++++++++++++++++++++ 1 file changed, 85 insertions(+) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 346ca9407d..4debf69aea 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -677,6 +677,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. @@ -1513,6 +1540,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( + "SplineMagneticField does not have vector potential calculation " + "implemented." + ) + @classmethod def from_mgrid(cls, mgrid_file, extcur=None, method="cubic", extrap=False): """Create a SplineMagneticField from an "mgrid" file from MAKEGRID. @@ -1674,6 +1728,37 @@ 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( + "ScalarPotentialField does not have vector potential calculation " + "implemented." + ) + + +# TODO: Implement a VectorPotentialField that uses AD to do curl(A) on the input +# magnetic vector potential function + def field_line_integrate( r0, From 89f5200350671c26a716977da0c7b92391cbed4a Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sun, 9 Jun 2024 10:03:31 -0400 Subject: [PATCH 11/43] add VectorPotentialField, ability to save/load vector potential for mgrids ans SplineMagneticField --- CHANGELOG.md | 6 +- desc/magnetic_fields/__init__.py | 1 + desc/magnetic_fields/_core.py | 320 ++++++++++++++++++++++++++++++- tests/test_magnetic_fields.py | 63 +++++- 4 files changed, 379 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 518ff5545f..92e1c968b5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,11 @@ New Features - Add method ``from_values`` to ``FourierRZCurve`` to allow fitting of data points to a ``FourierRZCurve`` object, and ``to_FourierRZCurve`` methods to ``Curve`` class. - +- 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. v0.11.1 ------- [Github Commits](https://github.com/PlasmaControl/DESC/compare/v0.11.0...v0.11.1) diff --git a/desc/magnetic_fields/__init__.py b/desc/magnetic_fields/__init__.py index 0a8f18abd8..173b04e7ee 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 4debf69aea..8816511f54 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -510,6 +510,22 @@ 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 + try: + 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) + except Exception as e: + warnif( + True, + UserWarning, + "Encountered error:" + f"{e} \nwhile attempting to compute vector magnetic potential." + " Vector potential will not be saved.", + ) + A_R = None + # write mgrid file file = Dataset(path, mode="w", format="NETCDF3_64BIT_OFFSET") @@ -596,6 +612,28 @@ def save_mgrid( ) bz_001[:] = B_Z + if A_R: + 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() @@ -1313,6 +1351,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 @@ -1331,6 +1375,9 @@ class SplineMagneticField(_MagneticField, Optimizable): "_BR", "_Bphi", "_BZ", + "_AR", + "_Aphi", + "_AZ", "_method", "_extrap", "_derivs", @@ -1343,7 +1390,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) @@ -1386,6 +1446,16 @@ def _atleast_4d(x): self._derivs["Bphi"] = self._approx_derivs(self._Bphi) self._derivs["BZ"] = self._approx_derivs(self._BZ) + if np.any(AR) and np.any(Aphi) and np.any(AZ): + 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) + @property def NFP(self): """int: Number of toroidal field periods.""" @@ -1562,10 +1632,106 @@ def compute_magnetic_vector_potential( magnetic vector potential at specified points """ - raise NotImplementedError( - "SplineMagneticField does not have vector potential calculation " - "implemented." + errorif( + not hasattr(self, "_AR"), + 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 self._axisym: + ARq = interp2d( + Rq, + Zq, + self._R, + self._Z, + self._AR[:, 0, :], + self._method, + (0, 0), + self._extrap, + (None, None), + **self._derivs["AR"], + ) + Aphiq = interp2d( + Rq, + Zq, + self._R, + self._Z, + self._Aphi[:, 0, :], + self._method, + (0, 0), + self._extrap, + (None, None), + **self._derivs["Aphi"], + ) + AZq = interp2d( + Rq, + Zq, + self._R, + self._Z, + self._AZ[:, 0, :], + self._method, + (0, 0), + self._extrap, + (None, None), + **self._derivs["AZ"], + ) + + else: + ARq = interp3d( + Rq, + phiq, + Zq, + self._R, + self._phi, + self._Z, + self._AR, + self._method, + (0, 0, 0), + self._extrap, + (None, 2 * np.pi / self.NFP, None), + **self._derivs["AR"], + ) + Aphiq = interp3d( + Rq, + phiq, + Zq, + self._R, + self._phi, + self._Z, + self._Aphi, + self._method, + (0, 0, 0), + self._extrap, + (None, 2 * np.pi / self.NFP, None), + **self._derivs["Aphi"], + ) + AZq = interp3d( + Rq, + phiq, + Zq, + self._R, + self._phi, + self._Z, + self._AZ, + self._method, + (0, 0, 0), + self._extrap, + (None, 2 * np.pi / self.NFP, None), + **self._derivs["AZ"], + ) + # ARq etc shape(nq, ngroups) + A = jnp.stack([ARq, Aphiq, AZq], axis=1) + # A shape(nq, 3, ngroups) + A = jnp.sum(A * currents, axis=-1) + if basis == "xyz": + A = rpz2xyz_vec(A, phi=coords[:, 1]) + return A @classmethod def from_mgrid(cls, mgrid_file, extcur=None, method="cubic", extrap=False): @@ -1623,8 +1789,41 @@ 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 as e: + warnif( + True, + UserWarning, + "Encountered error:" + f"{e} \nwhile attempting to read vector magnetic potential from mgrid." + " 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( @@ -1756,6 +1955,117 @@ def compute_magnetic_vector_potential( ) +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_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): + """Compute magnetic field at a set of points. + + Parameters + ---------- + coords : array-like shape(n,3) + 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. + + Returns + ------- + field : ndarray, shape(N,3) + magnetic field at specified points + + """ + 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 + 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 + + 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 + + """ + 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 + + A = self._potential(r, p, z, **params) + + if basis == "xyz": + A = rpz2xyz_vec(A, phi=coords[:, 1]) + return A + + # TODO: Implement a VectorPotentialField that uses AD to do curl(A) on the input # magnetic vector potential function diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 4ce4238e5c..624f7f1df9 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -8,6 +8,7 @@ from desc.basis import DoubleFourierSeries from desc.compute import rpz2xyz_vec, xyz2rpz_vec from desc.compute.utils import dot +from desc.derivatives import FiniteDiffDerivative as Derivative from desc.examples import get from desc.geometry import FourierRZToroidalSurface, FourierXYZCurve from desc.grid import LinearGrid @@ -22,6 +23,7 @@ ScalarPotentialField, SplineMagneticField, ToroidalMagneticField, + VectorPotentialField, VerticalMagneticField, field_line_integrate, read_BNORM_file, @@ -59,8 +61,23 @@ 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}) + 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((tfield + vfield)([1, 0, 0]), [[0, 2, 1]]) np.testing.assert_allclose( (tfield + vfield - pfield)([1, 0, 0.1]), [[0.4, 2, 1]] @@ -871,9 +888,38 @@ def test_spline_field(self): mgrid = "tests/inputs/mgrid_test.nc" field3 = SplineMagneticField.from_mgrid(mgrid, extcur) + 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]), @@ -908,9 +954,10 @@ 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): + 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])) @@ -1053,8 +1100,14 @@ 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.warns(UserWarning): + # user warning because poloidal field has no vector potential + # and so cannot save the vector potential + save_field.save_mgrid(path, Rmin, Rmax, Zmin, Zmax) + 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 From 08a6b3cd026301c833a49729f381d0d210bd3cd3 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sun, 9 Jun 2024 15:07:58 -0400 Subject: [PATCH 12/43] fix vfield vector potential --- desc/magnetic_fields/_core.py | 24 +++++++++++++++++------- tests/test_magnetic_fields.py | 19 ++++++++++++++++++- 2 files changed, 35 insertions(+), 8 deletions(-) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 8816511f54..c434e75f14 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 @@ -1070,6 +1070,8 @@ def compute_magnetic_vector_potential( ): """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) @@ -1167,6 +1169,8 @@ def compute_magnetic_vector_potential( ): """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) @@ -1190,12 +1194,18 @@ def compute_magnetic_vector_potential( assert basis.lower() in ["rpz", "xyz"] coords = jnp.atleast_2d(jnp.asarray(coords)) if basis == "xyz": - coords = xyz2rpz(coords) - axy = (B0 / 2) * jnp.ones_like(coords[:, 2]) - az = jnp.zeros_like(axy) - A = jnp.array([axy, -axy, az]).T - if basis == "xyz": - A = rpz2xyz_vec(A, phi=coords[:, 1]) + 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 diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 624f7f1df9..7f5f0d4c88 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -6,7 +6,7 @@ from desc.backend import jit, jnp from desc.basis import DoubleFourierSeries -from desc.compute import rpz2xyz_vec, xyz2rpz_vec +from desc.compute import rpz2xyz, rpz2xyz_vec, xyz2rpz_vec from desc.compute.utils import dot from desc.derivatives import FiniteDiffDerivative as Derivative from desc.examples import get @@ -70,6 +70,19 @@ def tfield_A(R, phi, Z, B0=2, R0=1): tfield_from_A = VectorPotentialField(tfield_A, params={"B0": 2, "R0": 1}) + def vfield_A(R, phi, Z, B0): + 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_from_A = VectorPotentialField(vfield_A, params={"B0": 1}) + 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]]) @@ -77,6 +90,10 @@ def tfield_A(R, phi, Z, B0=2, R0=1): 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_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( From e5d2302ca38954bff13a84abba36a1a1df92ddf2 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 10 Jun 2024 09:50:44 -0400 Subject: [PATCH 13/43] add catch for warnings in test --- tests/test_objective_funs.py | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 9dd4ec5fee..62f3d4bd41 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -1807,7 +1807,9 @@ def test_compute_scalar_resolution_bootstrap(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]) @@ -1833,7 +1835,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]) @@ -1860,7 +1864,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]) @@ -2121,7 +2126,9 @@ def test_objective_no_nangrad_bootstrap(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]) @@ -2142,7 +2149,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]) @@ -2165,7 +2174,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]) From 877e5ad3d2111b96ef212ab06ae5790b0933e3d0 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 11 Jun 2024 15:00:23 -0400 Subject: [PATCH 14/43] add to changelog and add comment --- CHANGELOG.md | 1 + tests/test_magnetic_fields.py | 1 + 2 files changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 92e1c968b5..fb17141313 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ to a ``FourierRZCurve`` object, and ``to_FourierRZCurve`` methods to ``Curve`` c - 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. + v0.11.1 ------- [Github Commits](https://github.com/PlasmaControl/DESC/compare/v0.11.0...v0.11.1) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 7f5f0d4c88..92d3559f45 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -972,6 +972,7 @@ def test_spline_field_axisym(self): -2.380229e04, ] with pytest.warns(UserWarning): + # user warning because saved mgrid no vector potential field = SplineMagneticField.from_mgrid( "tests/inputs/mgrid_d3d.nc", extcur=extcur ) From 5c8b56abc45781a5618c176cb4c0a488b9dcddfc Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 16 Jul 2024 15:24:44 -0400 Subject: [PATCH 15/43] fix basic_fields test --- tests/test_magnetic_fields.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 92d3559f45..6689b6854c 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -70,7 +70,7 @@ def tfield_A(R, phi, Z, B0=2, R0=1): tfield_from_A = VectorPotentialField(tfield_A, params={"B0": 2, "R0": 1}) - def vfield_A(R, phi, Z, B0): + 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] @@ -81,7 +81,8 @@ def vfield_A(R, phi, Z, B0): A = xyz2rpz_vec(A, phi=coords_rpz[:, 1]) return A - vfield_from_A = VectorPotentialField(vfield_A, params={"B0": 1}) + 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]]) @@ -91,7 +92,7 @@ def vfield_A(R, phi, Z, B0): tfield_from_A.compute_magnetic_vector_potential([1, 0, 0]).squeeze(), ) np.testing.assert_allclose( - vfield_A(1, 0, 0), + vfield_A(1, 0, 0, **vfield_params), vfield_from_A.compute_magnetic_vector_potential([1, 0, 0]), ) From ec576d67efecfd83b064c8a1cc8080361aa68ea2 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 16 Jul 2024 15:40:58 -0400 Subject: [PATCH 16/43] fix bug in vector magnetic potential in spline magnetic field --- desc/magnetic_fields/_core.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 3601e331bb..a9e82ef560 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -1477,8 +1477,7 @@ 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 np.any(AR) and np.any(Aphi) and np.any(AZ): + 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 @@ -1888,6 +1887,12 @@ 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 + if hasattr(field, "compute_magnetic_vector_potential"): + AR, AP, AZ = field.compute_magnetic_vector_potential( + coords, params, basis="rpz" + ).T + else: + AR = AP = AZ = None return cls( R, phi, @@ -1895,6 +1900,9 @@ def from_field( BR.reshape(shp), BP.reshape(shp), BZ.reshape(shp), + AR=AR.reshape(shp), + Aphi=AP.reshape(shp), + AZ=AZ.reshape(shp), currents=1.0, NFP=NFP, method=method, From 08f74456f5e917b39bb33de75bf0086a72320be6 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 16 Jul 2024 15:48:00 -0400 Subject: [PATCH 17/43] fix freeb test --- tests/test_examples.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index 56753905c0..bfe3e95c4e 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1066,9 +1066,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]) From 3d4e820186c5f903e6437918335e6d5ec2628715 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 16 Jul 2024 23:10:32 -0400 Subject: [PATCH 18/43] fix tests --- desc/magnetic_fields/_core.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index a9e82ef560..435a813a9a 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -1887,11 +1887,14 @@ 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 - if hasattr(field, "compute_magnetic_vector_potential"): + try: AR, AP, AZ = field.compute_magnetic_vector_potential( coords, params, basis="rpz" ).T - else: + AR = AR.reshape(shp) + AP = AP.reshape(shp) + AZ = AZ.reshape(shp) + except NotImplementedError: AR = AP = AZ = None return cls( R, @@ -1900,9 +1903,9 @@ def from_field( BR.reshape(shp), BP.reshape(shp), BZ.reshape(shp), - AR=AR.reshape(shp), - Aphi=AP.reshape(shp), - AZ=AZ.reshape(shp), + AR=AR, + Aphi=AP, + AZ=AZ, currents=1.0, NFP=NFP, method=method, From dc5aa01b3f064141992f37ba8586a09be731f704 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 17 Jul 2024 12:08:25 -0400 Subject: [PATCH 19/43] fix tests --- desc/magnetic_fields/_core.py | 2 +- tests/test_magnetic_fields.py | 11 +++++++---- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 435a813a9a..002a09f179 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -1222,7 +1222,7 @@ def compute_magnetic_vector_potential( ay = -B0 / 2 * coords_xyz[:, 0] az = jnp.zeros_like(ax) - A = jnp.array([ax, -ay, az]).T + A = jnp.array([ax, ay, az]).T if basis == "rpz": A = xyz2rpz_vec(A, phi=coords_rpz[:, 1]) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 6689b6854c..4cbbf83436 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -162,11 +162,14 @@ def test_combined_fields(self): np.testing.assert_allclose( sum_field([1.3, 0, 0.0]), [[0.0, 2, 3.2 + 2 * 1.2 * 0.3]] ) - R = 1.3 - tfield_A = np.array([[0, 0, -2 * 1.3 * np.log(R)]]) - vfield_A = np.array([[vfield.B0, -vfield.B0, 0]]) / 2 + + 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([R, 0, 0.0], basis="xyz"), + sum_field_tv.compute_magnetic_vector_potential([x, y, 0.0], basis="xyz"), tfield_A + vfield_A, ) From c6b80906c313e6897d941cca1d1af856be5ee8a4 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 19 Jul 2024 14:21:57 -0400 Subject: [PATCH 20/43] make test more compact --- tests/test_coils.py | 142 +++++++++++++++++--------------------------- 1 file changed, 53 insertions(+), 89 deletions(-) diff --git a/tests/test_coils.py b/tests/test_coils.py index 39fccf6a60..cc0f79ca9d 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -275,6 +275,38 @@ def _A_analytic(r): 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) @@ -300,119 +332,51 @@ def _A_analytic(r): grid_xyz = rpz2xyz(grid_rpz) # FourierXYZCoil coil = FourierXYZCoil(I, X_n=[-R, 0, 0], Y_n=[0, 0, R], Z_n=[0, 0, 0]) - 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=1e-8, err_msg="Using FourierXYZCoil" - ) - np.testing.assert_allclose( - correct_flux, flux_rpz, rtol=1e-8, err_msg="Using FourierXYZCoil" - ) - np.testing.assert_allclose( + test( + coil, + grid_xyz, + grid_rpz, A_true_rpz, - A_rpz, + correct_flux, rtol=1e-8, atol=1e-12, - err_msg="Using FourierXYZCoil", ) # SplineXYZCoil x = coil.compute("x", grid=coil_grid, basis="xyz")["x"] coil = SplineXYZCoil(I, X=x[:, 0], Y=x[:, 1], Z=x[:, 2]) - 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=1e-4, err_msg="Using SplineXYZCoil" - ) - np.testing.assert_allclose( - correct_flux, flux_rpz, rtol=1e-4, err_msg="Using SplineXYZCoil" - ) - np.testing.assert_allclose( + test( + coil, + grid_xyz, + grid_rpz, A_true_rpz, - A_rpz, + correct_flux, rtol=1e-4, atol=1e-12, - err_msg="Using SplineXYZCoil", ) + # FourierPlanarCoil coil = FourierPlanarCoil(I, center=[0, 0, 0], normal=[0, 0, -1], r_n=R) - 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=1e-8, err_msg="Using FourierPlanarCoil" - ) - np.testing.assert_allclose( - correct_flux, flux_rpz, rtol=1e-8, err_msg="Using FourierPlanarCoil" - ) - np.testing.assert_allclose( + test( + coil, + grid_xyz, + grid_rpz, A_true_rpz, - A_rpz, + correct_flux, rtol=1e-8, atol=1e-12, - err_msg="Using FourierPlanarCoil", ) # FourierRZCoil coil = FourierRZCoil(I, R_n=np.array([R]), modes_R=np.array([0])) - 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=1e-8, err_msg="Using FourierRZCoil" - ) - np.testing.assert_allclose( - correct_flux, flux_rpz, rtol=1e-8, err_msg="Using FourierRZCoil" - ) - np.testing.assert_allclose( + test( + coil, + grid_xyz, + grid_rpz, A_true_rpz, - A_rpz, + correct_flux, rtol=1e-8, atol=1e-12, - err_msg="Using FourierRZCoil", ) @pytest.mark.unit From c5464a92bdd963546c0d898346cea28713f01118 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 19 Jul 2024 14:24:42 -0400 Subject: [PATCH 21/43] make test more compact --- tests/test_coils.py | 89 +++++++++++---------------------------------- 1 file changed, 21 insertions(+), 68 deletions(-) diff --git a/tests/test_coils.py b/tests/test_coils.py index cc0f79ca9d..39c56d90cf 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -160,89 +160,42 @@ def test_biot_savart_vector_potential_all_coils(self): 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) - 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="Using FourierXYZCoil" - ) - np.testing.assert_allclose( - A_true, A_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierXYZCoil" - ) - np.testing.assert_allclose( - A_true, A_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierXYZCoil" - ) + 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]) - 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="Using SplineXYZCoil" - ) - np.testing.assert_allclose( - A_true, A_rpz, rtol=1e-3, atol=1e-10, err_msg="Using SplineXYZCoil" - ) - np.testing.assert_allclose( - A_true, A_rpz, rtol=1e-3, atol=1e-10, err_msg="Using SplineXYZCoil" - ) + test(coil, grid_xyz, grid_rpz) # FourierPlanarCoil coil = FourierPlanarCoil(I) - 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="Using FourierPlanarCoil" - ) - np.testing.assert_allclose( - A_true, - A_rpz, - rtol=1e-3, - atol=1e-10, - err_msg="Using FourierPlanarCoil", - ) - np.testing.assert_allclose( - A_true, - A_rpz, - rtol=1e-3, - atol=1e-10, - err_msg="Using FourierPlanarCoil", - ) + 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])) - 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="Using FourierRZCoil" - ) - np.testing.assert_allclose( - A_true, A_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierRZCoil" - ) - np.testing.assert_allclose( - A_true, A_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierRZCoil" - ) + test(coil, grid_xyz, grid_rpz) @pytest.mark.unit def test_biot_savart_vector_potential_integral_all_coils(self): From ac979c5e2134c01a9aadaf233a1d7fe73c7c61b7 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 19 Jul 2024 22:15:20 -0400 Subject: [PATCH 22/43] address comments --- desc/coils.py | 42 +++++++---- desc/magnetic_fields/_core.py | 81 ++++++++++++++-------- desc/magnetic_fields/_current_potential.py | 37 +++++++--- 3 files changed, 113 insertions(+), 47 deletions(-) diff --git a/desc/coils.py b/desc/coils.py index 7a895b3d1d..8db2eb6403 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -333,8 +333,8 @@ def compute_magnetic_field( params=params, transforms=transforms, profiles={}, - basis="xyz", ) + # TODO: need to check if basis is xyz here and swap if so? B = biot_savart_quad( coords, data["x"], data["x_s"] * data["ds"][:, None], current @@ -345,7 +345,7 @@ def compute_magnetic_field( return B def compute_magnetic_vector_potential( - self, coords, params=None, basis="rpz", source_grid=None + self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): """Compute magnetic vector potential at a set of points. @@ -363,6 +363,8 @@ def compute_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 ------- @@ -388,9 +390,23 @@ def compute_magnetic_vector_potential( else: current = params.pop("current", self.current) - data = self.compute( - ["x", "x_s", "ds"], grid=source_grid, params=params, basis="xyz" - ) + if not params or not transforms: + data = self.compute( + ["x", "x_s", "ds"], + grid=source_grid, + params=params, + transforms=transforms, + basis="xyz", + ) + else: + data = compute_fun( + self, + name=["x", "x_s", "ds"], + params=params, + transforms=transforms, + profiles={}, + ) + # TODO: need to check if basis is xyz here and swap if so? A = biot_savart_vector_potential_quad( coords, data["x"], data["x_s"] * data["ds"][:, None], current ) @@ -814,7 +830,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) @@ -832,11 +850,7 @@ def compute_magnetic_field( return B def compute_magnetic_vector_potential( - self, - coords, - params=None, - basis="rpz", - source_grid=None, + self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): """Compute magnetic vector potential at a set of points. @@ -855,6 +869,8 @@ def compute_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 ------- @@ -878,7 +894,9 @@ def compute_magnetic_vector_potential( 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) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 002a09f179..bc822c9b1b 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -227,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 ------- @@ -241,7 +243,7 @@ def __call__(self, grid, params=None, basis="rpz"): @abstractmethod def compute_magnetic_vector_potential( - self, coords, params=None, basis="rpz", source_grid=None + self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): """Compute magnetic vector potential at a set of points. @@ -256,6 +258,8 @@ def compute_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 ------- @@ -469,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. @@ -490,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 ------- @@ -511,19 +519,12 @@ def save_mgrid( B_Z = field[:, 2].reshape(nphi, nZ, nR) # evaluate magnetic vector potential on grid - try: + 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) - except Exception as e: - warnif( - True, - UserWarning, - "Encountered error:" - f"{e} \nwhile attempting to compute vector magnetic potential." - " Vector potential will not be saved.", - ) + else: A_R = None # write mgrid file @@ -828,7 +829,7 @@ def compute_magnetic_field( ) def compute_magnetic_vector_potential( - self, coords, params=None, basis="rpz", source_grid=None + self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): """Compute magnetic vector potential at a set of points. @@ -843,6 +844,8 @@ def compute_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 ------- @@ -930,7 +933,7 @@ def compute_magnetic_field( return B def compute_magnetic_vector_potential( - self, coords, params=None, basis="rpz", source_grid=None + self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): """Compute magnetic vector potential at a set of points. @@ -945,6 +948,8 @@ def compute_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 ------- @@ -960,15 +965,23 @@ def compute_magnetic_vector_potential( source_grid = [None] * len(self._fields) if not isinstance(source_grid, (list, tuple)): source_grid = [source_grid] + if transforms is None: + transforms = [None] * len(self._fields) + if not isinstance(transforms, (list, tuple)): + transforms = [transforms] if len(source_grid) != len(self._fields): # ensure that if source_grid is shorter, that it is simply repeated so that # zip does not terminate early source_grid = source_grid * len(self._fields) + if len(transforms) != len(self._fields): + # ensure that if transforms is shorter, that it is simply repeated so that + # zip does not terminate early + transforms = transforms * len(self._fields) A = 0 - for i, (field, g) in enumerate(zip(self._fields, source_grid)): + for i, (field, g, tr) in enumerate(zip(self._fields, source_grid, transforms)): A += field.compute_magnetic_vector_potential( - coords, params[i % len(params)], basis, source_grid=g + coords, params[i % len(params)], basis, source_grid=g, transforms=tr ) return A @@ -1082,11 +1095,11 @@ def compute_magnetic_field( return B def compute_magnetic_vector_potential( - self, coords, params=None, basis="rpz", source_grid=None + 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. + The vector potential is specified assuming the Coulomb Gauge. Parameters ---------- @@ -1098,6 +1111,9 @@ def compute_magnetic_vector_potential( 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 ------- @@ -1124,6 +1140,8 @@ def compute_magnetic_vector_potential( 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 @@ -1184,7 +1202,7 @@ def compute_magnetic_field( return B def compute_magnetic_vector_potential( - self, coords, params=None, basis="rpz", source_grid=None + self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): """Compute magnetic vector potential at a set of points. @@ -1200,6 +1218,9 @@ def compute_magnetic_vector_potential( 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 ------- @@ -1339,7 +1360,7 @@ def compute_magnetic_field( return B def compute_magnetic_vector_potential( - self, coords, params=None, basis="rpz", source_grid=None + self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): """Compute magnetic vector potential at a set of points. @@ -1353,6 +1374,9 @@ def compute_magnetic_vector_potential( 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 ------- @@ -1486,6 +1510,8 @@ def _atleast_4d(x): 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): @@ -1645,7 +1671,7 @@ def compute_magnetic_field( return B def compute_magnetic_vector_potential( - self, coords, params=None, basis="rpz", source_grid=None + self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): """Compute magnetic vector potential at a set of points. @@ -1659,6 +1685,9 @@ def compute_magnetic_vector_potential( 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 ------- @@ -1844,12 +1873,11 @@ def from_mgrid(cls, mgrid_file, extcur=None, method="cubic", extrap=False): 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 as e: + except IndexError: warnif( True, UserWarning, - "Encountered error:" - f"{e} \nwhile attempting to read vector magnetic potential from mgrid." + "mgrid does not appear to contain vector potential information." " Vector potential will not be computable.", ) ar = ap = az = None @@ -1977,7 +2005,7 @@ def compute_magnetic_field( return B def compute_magnetic_vector_potential( - self, coords, params=None, basis="rpz", source_grid=None + self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): """Compute magnetic vector potential at a set of points. @@ -1991,6 +2019,9 @@ def compute_magnetic_vector_potential( 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 ------- @@ -2115,10 +2146,6 @@ def compute_magnetic_vector_potential( return A -# TODO: Implement a VectorPotentialField that uses AD to do curl(A) on the input -# magnetic vector potential function - - def field_line_integrate( r0, z0, diff --git a/desc/magnetic_fields/_current_potential.py b/desc/magnetic_fields/_current_potential.py index 49f0d7385a..918e900f29 100644 --- a/desc/magnetic_fields/_current_potential.py +++ b/desc/magnetic_fields/_current_potential.py @@ -220,7 +220,7 @@ def compute_magnetic_field( ) def compute_magnetic_vector_potential( - self, coords, params=None, basis="rpz", source_grid=None + self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): """Compute magnetic vector potential at a set of points. @@ -234,6 +234,8 @@ def compute_magnetic_vector_potential( 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 ------- @@ -252,6 +254,7 @@ def compute_magnetic_vector_potential( params=params, basis=basis, source_grid=source_grid, + transforms=transforms, ) @classmethod @@ -574,7 +577,7 @@ def compute_magnetic_field( ) def compute_magnetic_vector_potential( - self, coords, params=None, basis="rpz", source_grid=None + self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): """Compute magnetic vector potential at a set of points. @@ -588,6 +591,8 @@ def compute_magnetic_vector_potential( 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 ------- @@ -606,6 +611,7 @@ def compute_magnetic_vector_potential( params=params, basis=basis, source_grid=source_grid, + transforms=transforms, ) @classmethod @@ -773,11 +779,7 @@ def nfp_loop(j, f): def _compute_vector_potential_from_CurrentPotentialField( - field, - coords, - source_grid, - params=None, - basis="rpz", + field, coords, source_grid, params=None, basis="rpz", transforms=None ): """Compute magnetic vector potential at a set of points. @@ -794,6 +796,8 @@ def _compute_vector_potential_from_CurrentPotentialField( should include the potential basis : {"rpz", "xyz"} basis for input coordinates and returned magnetic vector potential + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from source_grid Returns @@ -810,7 +814,24 @@ def _compute_vector_potential_from_CurrentPotentialField( # compute surface current, and store grid quantities # needed for integration in class # TODO: does this have to be xyz, or can it be computed in rpz as well? - data = field.compute(["K", "x"], grid=source_grid, basis="xyz", params=params) + if not params or not transforms: + data = field.compute( + ["K", "x"], + grid=source_grid, + basis="xyz", + params=params, + transforms=transforms, + jitable=True, + ) + else: + data = compute_fun( + field, + names=["K", "x"], + params=params, + transforms=transforms, + profiles={}, + basis="xyz", + ) _rs = xyz2rpz(data["x"]) _K = xyz2rpz_vec(data["K"], phi=source_grid.nodes[:, 2]) From 911650ae8af03d8c06235ac20463c9ab1c92d4c1 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 19 Jul 2024 23:03:34 -0400 Subject: [PATCH 23/43] fix test --- tests/test_magnetic_fields.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 4cbbf83436..ec721a7c4f 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -1122,10 +1122,11 @@ def test_mgrid_io(self, tmpdir_factory): Rmax = 7 Zmin = -2 Zmax = 2 - with pytest.warns(UserWarning): - # user warning because poloidal field has no vector potential + 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 From d22fa596a9e59b7ca08073b047ad0e284dcb255d Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sun, 21 Jul 2024 01:05:27 -0400 Subject: [PATCH 24/43] refactor magnetic field computes to re-use code --- desc/magnetic_fields/_core.py | 618 ++++++++++++++++++++-------------- 1 file changed, 359 insertions(+), 259 deletions(-) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index bc822c9b1b..9b3547436e 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -878,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 ---------- @@ -896,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 ------- @@ -903,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): @@ -924,13 +938,54 @@ def compute_magnetic_field( # zip does not terminate early transforms = transforms * len(self._fields) - B = 0 - for i, (field, g, tr) in enumerate(zip(self._fields, source_grid, transforms)): - B += field.compute_magnetic_field( - coords, params[i % len(params)], basis, source_grid=g, transforms=tr - ) + if compute_A_or_B == "B": + B = 0 + for i, (field, g, tr) in enumerate( + zip(self._fields, source_grid, transforms) + ): + B += field.compute_magnetic_field( + coords, params[i % len(params)], basis, source_grid=g, transforms=tr + ) + return B + elif compute_A_or_B == "A": + A = 0 + for i, (field, g, tr) in enumerate( + zip(self._fields, source_grid, transforms) + ): + A += field.compute_magnetic_vector_potential( + coords, params[i % len(params)], basis, source_grid=g, transforms=tr + ) - return B + 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 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 @@ -954,37 +1009,12 @@ def compute_magnetic_vector_potential( Returns ------- A : ndarray, shape(N,3) - scaled magnetic vector potential at specified points + sum magnetic vector potential at specified points """ - if params is None: - params = [None] * len(self._fields) - if isinstance(params, dict): - params = [params] - if source_grid is None: - source_grid = [None] * len(self._fields) - if not isinstance(source_grid, (list, tuple)): - source_grid = [source_grid] - if transforms is None: - transforms = [None] * len(self._fields) - if not isinstance(transforms, (list, tuple)): - transforms = [transforms] - if len(source_grid) != len(self._fields): - # ensure that if source_grid is shorter, that it is simply repeated so that - # zip does not terminate early - source_grid = source_grid * len(self._fields) - if len(transforms) != len(self._fields): - # ensure that if transforms is shorter, that it is simply repeated so that - # zip does not terminate early - transforms = transforms * len(self._fields) - - A = 0 - for i, (field, g, tr) in enumerate(zip(self._fields, source_grid, transforms)): - A += field.compute_magnetic_vector_potential( - coords, params[i % len(params)], basis, source_grid=g, transforms=tr - ) - - return A + 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): @@ -1053,10 +1083,16 @@ def B0(self): def B0(self, new): self._B0 = float(np.squeeze(new)) - 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 ---------- @@ -1071,13 +1107,21 @@ 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 + 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}', + ) params = setdefault(params, {}) B0 = params.get("B0", self.B0) R0 = params.get("R0", self.R0) @@ -1086,13 +1130,48 @@ def compute_magnetic_field( coords = jnp.atleast_2d(jnp.asarray(coords)) if basis == "xyz": coords = xyz2rpz(coords) - bp = B0 * R0 / coords[:, 0] - brz = jnp.zeros_like(bp) - B = jnp.array([brz, bp, brz]).T - if basis == "xyz": - B = rpz2xyz_vec(B, phi=coords[:, 1]) + if compute_A_or_B == "B": + bp = B0 * R0 / coords[:, 0] + brz = jnp.zeros_like(bp) + B = jnp.array([brz, bp, brz]).T + if basis == "xyz": + B = rpz2xyz_vec(B, phi=coords[:, 1]) + + return B + elif compute_A_or_B == "A": + 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 - 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 + Dict of values for R0 and B0. + basis : {"rpz", "xyz"} + Basis for input coordinates and returned magnetic field. + 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 + ------- + 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 @@ -1121,20 +1200,7 @@ def compute_magnetic_vector_potential( 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 + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") class VerticalMagneticField(_MagneticField, Optimizable): @@ -1164,6 +1230,71 @@ def B0(self): def B0(self, new): self._B0 = float(np.squeeze(new)) + 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 magnetic 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 + Dict of values for B0. + basis : {"rpz", "xyz"} + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional + 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 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}', + ) + params = setdefault(params, {}) + B0 = params.get("B0", self.B0) + + coords = jnp.atleast_2d(jnp.asarray(coords)) + if compute_A_or_B == "B": + bz = B0 * jnp.ones_like(coords[:, 2]) + brp = jnp.zeros_like(bz) + B = jnp.array([brp, brp, bz]).T + # b/c it only has a nonzero z component, no need + # to switch bases back if xyz is given + + return B + elif compute_A_or_B == "A": + 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 + def compute_magnetic_field( self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): @@ -1189,17 +1320,7 @@ def compute_magnetic_field( magnetic field at specified points """ - params = setdefault(params, {}) - B0 = params.get("B0", self.B0) - - coords = jnp.atleast_2d(jnp.asarray(coords)) - bz = B0 * jnp.ones_like(coords[:, 2]) - brp = jnp.zeros_like(bz) - B = jnp.array([brp, brp, bz]).T - # b/c it only has a nonzero z component, no need - # to switch bases back if xyz is given - - return B + 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 @@ -1228,26 +1349,7 @@ def compute_magnetic_vector_potential( magnetic vector potential at specified points """ - 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_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 + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") class PoloidalMagneticField(_MagneticField, Optimizable): @@ -1550,10 +1652,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 ---------- @@ -1568,107 +1676,158 @@ 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 not hasattr(self, "_AR"), + 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 @@ -1695,106 +1854,7 @@ def compute_magnetic_vector_potential( magnetic vector potential at specified points """ - errorif( - not hasattr(self, "_AR"), - 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 self._axisym: - ARq = interp2d( - Rq, - Zq, - self._R, - self._Z, - self._AR[:, 0, :], - self._method, - (0, 0), - self._extrap, - (None, None), - **self._derivs["AR"], - ) - Aphiq = interp2d( - Rq, - Zq, - self._R, - self._Z, - self._Aphi[:, 0, :], - self._method, - (0, 0), - self._extrap, - (None, None), - **self._derivs["Aphi"], - ) - AZq = interp2d( - Rq, - Zq, - self._R, - self._Z, - self._AZ[:, 0, :], - self._method, - (0, 0), - self._extrap, - (None, None), - **self._derivs["AZ"], - ) - - else: - ARq = interp3d( - Rq, - phiq, - Zq, - self._R, - self._phi, - self._Z, - self._AR, - self._method, - (0, 0, 0), - self._extrap, - (None, 2 * np.pi / self.NFP, None), - **self._derivs["AR"], - ) - Aphiq = interp3d( - Rq, - phiq, - Zq, - self._R, - self._phi, - self._Z, - self._Aphi, - self._method, - (0, 0, 0), - self._extrap, - (None, 2 * np.pi / self.NFP, None), - **self._derivs["Aphi"], - ) - AZq = interp3d( - Rq, - phiq, - Zq, - self._R, - self._phi, - self._Z, - self._AZ, - self._method, - (0, 0, 0), - self._extrap, - (None, 2 * np.pi / self.NFP, None), - **self._derivs["AZ"], - ) - # ARq etc shape(nq, ngroups) - A = jnp.stack([ARq, Aphiq, AZq], axis=1) - # A shape(nq, 3, ngroups) - A = jnp.sum(A * currents, axis=-1) - if basis == "xyz": - A = rpz2xyz_vec(A, phi=coords[:, 1]) - return A + 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): @@ -2053,10 +2113,16 @@ def __init__(self, potential, params=None): self._potential = potential self._params = params - def compute_magnetic_field( - self, coords, params=None, basis="rpz", source_grid=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 ---------- @@ -2068,6 +2134,12 @@ def compute_magnetic_field( 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 ------- @@ -2075,6 +2147,11 @@ def compute_magnetic_field( 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) @@ -2084,31 +2161,65 @@ def compute_magnetic_field( if params is None: params = self._params r, p, z = coords.T - 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] + 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) - # 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) + ap = self._potential(r, p, z, **params)[:, 1] - # 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 + # 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 + self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): """Compute magnetic vector potential at a set of points. @@ -2122,6 +2233,9 @@ def compute_magnetic_vector_potential( 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 ------- @@ -2129,21 +2243,7 @@ def compute_magnetic_vector_potential( magnetic vector potential at specified points """ - 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 - - A = self._potential(r, p, z, **params) - - if basis == "xyz": - A = rpz2xyz_vec(A, phi=coords[:, 1]) - return A + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") def field_line_integrate( From 4d89bde3033465aa89931a57ba97b4d260e4e2c2 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sun, 21 Jul 2024 01:20:08 -0400 Subject: [PATCH 25/43] refactor compute methods of current potential fields --- desc/magnetic_fields/_current_potential.py | 221 +++++++++++++++++---- 1 file changed, 185 insertions(+), 36 deletions(-) diff --git a/desc/magnetic_fields/_current_potential.py b/desc/magnetic_fields/_current_potential.py index 918e900f29..8f574770a2 100644 --- a/desc/magnetic_fields/_current_potential.py +++ b/desc/magnetic_fields/_current_potential.py @@ -181,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 ---------- @@ -198,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( @@ -210,15 +219,42 @@ 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 ): @@ -243,19 +279,7 @@ def compute_magnetic_vector_potential( Magnetic vector potential at specified points. """ - source_grid = source_grid or LinearGrid( - M=30 + 2 * self.M, - N=30 + 2 * self.N, - NFP=self.NFP, - ) - return _compute_vector_potential_from_CurrentPotentialField( - field=self, - coords=coords, - params=params, - basis=basis, - source_grid=source_grid, - transforms=transforms, - ) + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") @classmethod def from_surface( @@ -538,10 +562,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 ---------- @@ -555,11 +585,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( @@ -567,15 +600,42 @@ 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 ): @@ -600,19 +660,7 @@ def compute_magnetic_vector_potential( Magnetic vector potential at specified points. """ - source_grid = source_grid or LinearGrid( - M=30 + 2 * max(self.M, self.M_Phi), - N=30 + 2 * max(self.N, self.N_Phi), - NFP=self.NFP, - ) - return _compute_vector_potential_from_CurrentPotentialField( - field=self, - coords=coords, - params=params, - basis=basis, - source_grid=source_grid, - transforms=transforms, - ) + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") @classmethod def from_surface( @@ -693,6 +741,107 @@ def from_surface( ) +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 or vector potential at a set of points. + + Parameters + ---------- + field : CurrentPotentialField or FourierCurrentPotentialField + current potential field object from which to compute magnetic field. + coords : array-like shape(N,3) + cylindrical or cartesian coordinates + source_grid : Grid, + source grid upon which to evaluate the surface current density K + params : dict, optional + parameters to pass to compute function + 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 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 + # TODO: does this have to be xyz, or can it be computed in rpz as well? + if not params or not transforms: + data = field.compute( + ["K", "x"], + grid=source_grid, + basis="xyz", + params=params, + transforms=transforms, + jitable=True, + ) + else: + data = compute_fun( + field, + names=["K", "x"], + params=params, + transforms=transforms, + profiles={}, + basis="xyz", + ) + + _rs = xyz2rpz(data["x"]) + _K = xyz2rpz_vec(data["K"], phi=source_grid.nodes[:, 2]) + + # surface element, must divide by NFP to remove the NFP multiple on the + # surface grid weights, as we account for that when doing the for loop + # over NFP + _dV = source_grid.weights * data["|e_theta x e_zeta|"] / source_grid.NFP + + def nfp_loop(j, f): + # calculate (by rotating) rs, rs_t, rz_t + phi = (source_grid.nodes[:, 2] + j * 2 * jnp.pi / source_grid.NFP) % ( + 2 * jnp.pi + ) + # new coords are just old R,Z at a new phi (bc of discrete NFP symmetry) + rs = jnp.vstack((_rs[:, 0], phi, _rs[:, 2])).T + rs = rpz2xyz(rs) + K = rpz2xyz_vec(_K, phi=phi) + fj = op( + coords, + rs, + K, + _dV, + ) + f += fj + return f + + B = fori_loop(0, source_grid.NFP, nfp_loop, jnp.zeros_like(coords)) + if basis == "rpz": + B = xyz2rpz_vec(B, x=coords[:, 0], y=coords[:, 1]) + return B + + def _compute_magnetic_field_from_CurrentPotentialField( field, coords, source_grid, params=None, basis="rpz", transforms=None ): From 9d8970960ea095fb909c056a75dd8a3471a303b7 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sun, 21 Jul 2024 01:56:55 -0400 Subject: [PATCH 26/43] refactor coils compute magnetic field to use less code --- desc/coils.py | 397 +++++++++++++++++++++++++++++++++----------- tests/test_coils.py | 7 + 2 files changed, 304 insertions(+), 100 deletions(-) diff --git a/desc/coils.py b/desc/coils.py index 8db2eb6403..9023a976bd 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -268,10 +268,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. @@ -289,6 +295,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 @@ -304,6 +313,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": @@ -336,13 +353,49 @@ def compute_magnetic_field( ) # TODO: need to check if basis is xyz here and swap if so? - 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 @@ -380,40 +433,7 @@ def compute_magnetic_vector_potential( may not be zero if not fully converged. """ - assert basis.lower() in ["rpz", "xyz"] - coords = jnp.atleast_2d(jnp.asarray(coords)) - if basis.lower() == "rpz": - phi = coords[:, 1] - coords = rpz2xyz(coords) - if params is None: - current = self.current - else: - current = params.pop("current", self.current) - - if not params or not transforms: - data = self.compute( - ["x", "x_s", "ds"], - grid=source_grid, - params=params, - transforms=transforms, - basis="xyz", - ) - else: - data = compute_fun( - self, - name=["x", "x_s", "ds"], - params=params, - transforms=transforms, - profiles={}, - ) - # TODO: need to check if basis is xyz here and swap if so? - A = biot_savart_vector_potential_quad( - coords, data["x"], data["x_s"] * data["ds"][:, None], current - ) - - if basis.lower() == "rpz": - A = xyz2rpz_vec(A, phi=phi) - return A + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") def __repr__(self): """Get the string form of the object.""" @@ -787,10 +807,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. @@ -808,6 +834,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 ------- @@ -821,6 +850,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": @@ -843,11 +878,47 @@ 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 @@ -885,35 +956,7 @@ def compute_magnetic_vector_potential( Convergence is approximately quadratic in the number of coil points. """ - assert basis.lower() in ["rpz", "xyz"] - coords = jnp.atleast_2d(jnp.asarray(coords)) - if basis == "rpz": - coords = rpz2xyz(coords) - if params is None: - current = self.current - else: - current = params.pop("current", self.current) - - 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) - coil_pts_start = data["x"] - coil_pts_end = jnp.concatenate([data["x"][1:], data["x"][:1]]) - # could get up to 4th order accuracy by shifting points outward as in - # (McGreivy, Zhu, Gunderson, Hudson 2021), however that requires knowing the - # 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?) - A = biot_savart_vector_potential_hh( - coords, coil_pts_start, coil_pts_end, current - ) - - if basis == "rpz": - A = xyz2rpz_vec(A, x=coords[:, 0], y=coords[:, 1]) - return A + return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") @classmethod def from_values( @@ -1222,8 +1265,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. @@ -1240,13 +1289,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: @@ -1276,31 +1334,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( @@ -1976,6 +2092,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 ): @@ -2003,15 +2178,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. - return B + 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. + + 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_FourierXYZ(self, N=10, grid=None, s=None, name=""): """Convert all coils to FourierXYZCoil representation. diff --git a/tests/test_coils.py b/tests/test_coils.py index 39c56d90cf..80f4b9120c 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -196,6 +196,13 @@ def test(coil, grid_xyz, grid_rpz): # 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): From 1d60c40958737ffd3a3295bfb450db7a620ec5e8 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sun, 21 Jul 2024 02:50:42 -0400 Subject: [PATCH 27/43] fix docstring error --- desc/coils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/coils.py b/desc/coils.py index 9023a976bd..edd0cff457 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -1403,7 +1403,7 @@ def compute_magnetic_vector_potential( 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 + 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 From 35fd2942ca2ed56399938595890c706827238b1d Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 31 Jul 2024 21:23:32 -0400 Subject: [PATCH 28/43] fix vector pot saving for spline field, add test --- desc/magnetic_fields/_core.py | 2 +- tests/test_magnetic_fields.py | 28 +++++++++++++++++++++++++++- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 9b3547436e..06a1a4e1c9 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -613,7 +613,7 @@ def save_mgrid( ) bz_001[:] = B_Z - if A_R: + if A_R is not None: ar_001 = file.createVariable("ar_001", np.float64, ("phi", "zee", "rad")) ar_001.long_name = ( "A_R = radial component of magnetic vector potential " diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index f1ab758a94..3e118683f6 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -901,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) @@ -916,6 +916,31 @@ 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 @@ -947,6 +972,7 @@ def test_spline_field(self): 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 From c54723df4b638fd66c98b67427111a033335f739 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 1 Aug 2024 11:46:39 -0500 Subject: [PATCH 29/43] change toroidal flux obj to use vector potential --- desc/objectives/_coils.py | 65 +++++++++++++++++++++++++----------- tests/test_objective_funs.py | 29 ++++++++-------- 2 files changed, 62 insertions(+), 32 deletions(-) diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index 62ef664622..6a43760581 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -1348,6 +1348,15 @@ class ToroidalFlux(_Objective): zeta=jnp.array(0.0), NFP=eq.NFP). name : str, optional Name of the objective function. + use_vector_potential : True + Whether 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 this option should be + set to False in those cases or if the magnetic field method is preferred. + #TODO: add eq_fixed option so this can be used in single stage """ @@ -1369,6 +1378,7 @@ def __init__( field_grid=None, eval_grid=None, name="toroidal-flux", + use_vector_potential=True, ): if target is None and bounds is None: target = eq.Psi @@ -1376,6 +1386,7 @@ def __init__( self._field_grid = field_grid self._eval_grid = eval_grid self._eq = eq + self._use_vector_potential = use_vector_potential super().__init__( things=[field], @@ -1403,7 +1414,10 @@ def build(self, use_jit=True, verbose=1): eq = self._eq 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 @@ -1438,10 +1452,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 @@ -1483,22 +1499,33 @@ 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) + # TODO: use the line integral compute utilities + 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_objective_funs.py b/tests/test_objective_funs.py index 552f189bfd..bd02033eb5 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -1107,10 +1107,17 @@ 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) + def test( + eq, field, correct_value, rtol=1e-14, grid=None, use_vector_potential=True + ): + obj = ToroidalFlux( + eq=eq, + field=field, + eval_grid=grid, + use_vector_potential=use_vector_potential, + ) obj.build(verbose=2) torflux = obj.compute_unscaled(*obj.xs(field)) np.testing.assert_allclose(torflux, correct_value, rtol=rtol) @@ -1120,22 +1127,18 @@ 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, use_vector_potential=False) @pytest.mark.regression From 4574dc479b8aa4c7ede1238307412abcae4f7e4e Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 1 Aug 2024 11:49:02 -0500 Subject: [PATCH 30/43] add tests to cover both methods of toroidal flux obj --- tests/test_objective_funs.py | 31 +++++++++++++++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index bd02033eb5..1432ae2091 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -2093,7 +2093,7 @@ 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") @@ -2110,6 +2110,26 @@ def test_compute_scalar_resolution_toroidal_flux(self): 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") + 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_vector_potential=False), 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_generic_scalar(self): """Generic objective with scalar qty.""" @@ -2415,7 +2435,14 @@ 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_vector_potential=False), 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( From 00e4b982040a99a52ccf84ba047ee77b77878065 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 1 Aug 2024 11:50:36 -0500 Subject: [PATCH 31/43] address comment --- desc/magnetic_fields/_core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 06a1a4e1c9..9c4047d65d 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -1693,7 +1693,7 @@ def _compute_A_or_B( f'Expected "A" or "B" for compute_A_or_B, instead got {compute_A_or_B}', ) errorif( - compute_A_or_B == "A" and not hasattr(self, "_AR"), + compute_A_or_B == "A" and self._AR is None, ValueError, "Cannot calculate vector potential" " as no vector potential spline values exist.", From 4e4410b0ddec58767c625ff4a402469c33af4283 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 1 Aug 2024 11:52:19 -0500 Subject: [PATCH 32/43] add test --- tests/test_magnetic_fields.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 3e118683f6..ed0ae43178 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -1019,6 +1019,10 @@ def test_spline_field_axisym(self): 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.""" From d337ed993e5ba8a3153036d8bca9153b0047c394 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 1 Aug 2024 12:12:48 -0500 Subject: [PATCH 33/43] address comment --- desc/magnetic_fields/_core.py | 29 ++++++++++------------------- 1 file changed, 10 insertions(+), 19 deletions(-) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 9c4047d65d..c7e68a4152 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -938,25 +938,16 @@ def _compute_A_or_B( # zip does not terminate early transforms = transforms * len(self._fields) - if compute_A_or_B == "B": - B = 0 - for i, (field, g, tr) in enumerate( - zip(self._fields, source_grid, transforms) - ): - B += field.compute_magnetic_field( - coords, params[i % len(params)], basis, source_grid=g, transforms=tr - ) - return B - elif compute_A_or_B == "A": - A = 0 - for i, (field, g, tr) in enumerate( - zip(self._fields, source_grid, transforms) - ): - A += field.compute_magnetic_vector_potential( - coords, params[i % len(params)], basis, source_grid=g, transforms=tr - ) - - return A + 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)): + AB += getattr(field, op)( + coords, params[i % len(params)], basis, source_grid=g, transforms=tr + ) + return AB def compute_magnetic_field( self, coords, params=None, basis="rpz", source_grid=None, transforms=None From bd61ef88ae89e0a0568f498db14e2ed690ae6b89 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 1 Aug 2024 12:15:32 -0500 Subject: [PATCH 34/43] update changelog --- CHANGELOG.md | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5fc10ab694..75d895d8c0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,17 @@ Changelog ========= +New Features + +- 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. +- Add ``use_vector_potential`` flag to ``ToroidalFlux`` objective which defaults to True, which will use a 1D loop integral of the vector potential +to compute the toroidal flux, as opposed to a 2D surface integral of the magnetic field dotted with ``n_zeta`` + + v0.12.1 ------- @@ -18,15 +29,6 @@ to be much faster for these cases. - Adds new ``from_values`` method for ``FourierPlanarCurve`` and ``FourierPlanarCoil`` -New Features - -- 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. - - v0.12.0 ------- From ac7aaa2abf51c29d50f42321ac02e21856cdd214 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 1 Aug 2024 12:18:55 -0500 Subject: [PATCH 35/43] update changelog (again) --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 75d895d8c0..8f63939d19 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ New Features computation. - Add ability to save and load vector potential information from ``mgrid`` files. - Add ``use_vector_potential`` flag to ``ToroidalFlux`` objective which defaults to True, which will use a 1D loop integral of the vector potential -to compute the toroidal flux, as opposed to a 2D surface integral of the magnetic field dotted with ``n_zeta`` +to compute the toroidal flux, as opposed to a 2D surface integral of the magnetic field dotted with ``n_zeta``. v0.12.1 From 008ecbc337a2577793bb77e8e427866338da6ef8 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 2 Aug 2024 01:19:24 -0400 Subject: [PATCH 36/43] put comments assuming coulomb gauge --- desc/coils.py | 4 +++- desc/magnetic_fields/_current_potential.py | 4 ++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/desc/coils.py b/desc/coils.py index 46f64d5628..e30f910cfc 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -88,7 +88,7 @@ def biot_savart_vector_potential_hh(eval_pts, coil_pts_start, coil_pts_end, curr 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. + segment. This expression assumes the Coulomb gauge. Parameters ---------- @@ -174,6 +174,8 @@ def biot_savart_quad(eval_pts, coil_pts, tangents, current): 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) diff --git a/desc/magnetic_fields/_current_potential.py b/desc/magnetic_fields/_current_potential.py index 8ffb0e0cea..dc48ebaf5e 100644 --- a/desc/magnetic_fields/_current_potential.py +++ b/desc/magnetic_fields/_current_potential.py @@ -260,6 +260,8 @@ def compute_magnetic_vector_potential( ): """Compute magnetic vector potential at a set of points. + This assumes the Coulomb gauge. + Parameters ---------- coords : array-like shape(n,3) @@ -641,6 +643,8 @@ def compute_magnetic_vector_potential( ): """Compute magnetic vector potential at a set of points. + This assumes the Coulomb gauge. + Parameters ---------- coords : array-like shape(n,3) From e3b0e44e8c6478a544173ba55a9911913724fde3 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 8 Aug 2024 16:30:26 -0500 Subject: [PATCH 37/43] address comments --- desc/magnetic_fields/_current_potential.py | 1 - desc/objectives/_coils.py | 26 +++++++++++++--------- tests/test_objective_funs.py | 10 ++++----- 3 files changed, 20 insertions(+), 17 deletions(-) diff --git a/desc/magnetic_fields/_current_potential.py b/desc/magnetic_fields/_current_potential.py index dc48ebaf5e..a8759155ee 100644 --- a/desc/magnetic_fields/_current_potential.py +++ b/desc/magnetic_fields/_current_potential.py @@ -794,7 +794,6 @@ def _compute_A_or_B_from_CurrentPotentialField( ] # compute surface current, and store grid quantities # needed for integration in class - # TODO: does this have to be xyz, or can it be computed in rpz as well? if not params or not transforms: data = field.compute( ["K", "x"], diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index 6a43760581..14c4c7165d 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -1304,6 +1304,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 @@ -1348,15 +1356,7 @@ class ToroidalFlux(_Objective): zeta=jnp.array(0.0), NFP=eq.NFP). name : str, optional Name of the objective function. - use_vector_potential : True - Whether 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 this option should be - set to False in those cases or if the magnetic field method is preferred. - #TODO: add eq_fixed option so this can be used in single stage + """ @@ -1378,7 +1378,6 @@ def __init__( field_grid=None, eval_grid=None, name="toroidal-flux", - use_vector_potential=True, ): if target is None and bounds is None: target = eq.Psi @@ -1386,7 +1385,7 @@ def __init__( self._field_grid = field_grid self._eval_grid = eval_grid self._eq = eq - self._use_vector_potential = use_vector_potential + # TODO: add eq_fixed option so this can be used in single stage super().__init__( things=[field], @@ -1412,6 +1411,11 @@ 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: + self._use_vector_potential = False if self._eval_grid is None: eval_grid = LinearGrid( L=eq.L_grid if not self._use_vector_potential else 0, diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 1432ae2091..0eeefa9c35 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, @@ -1109,14 +1110,11 @@ def test_toroidal_flux(self): """Test calculation of toroidal flux from coils.""" grid1 = LinearGrid(L=0, M=40, zeta=np.array(0.0)) - def test( - eq, field, correct_value, rtol=1e-14, grid=None, use_vector_potential=True - ): + def test(eq, field, correct_value, rtol=1e-14, grid=None): obj = ToroidalFlux( eq=eq, field=field, eval_grid=grid, - use_vector_potential=use_vector_potential, ) obj.build(verbose=2) torflux = obj.compute_unscaled(*obj.xs(field)) @@ -1138,7 +1136,9 @@ def test( eq.change_resolution(L_grid=20, M_grid=20) test(eq, field, psi_from_field) - test(eq, field, psi_from_field, rtol=1e-3, use_vector_potential=False) + 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.regression From c45a475af8db711ed8a0afe61d29f0945b10e90f Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 8 Aug 2024 16:32:41 -0500 Subject: [PATCH 38/43] update changelog --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8f63939d19..118ceb2482 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,8 +8,8 @@ New Features - 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. -- Add ``use_vector_potential`` flag to ``ToroidalFlux`` objective which defaults to True, which will use a 1D loop integral of the vector potential -to compute the toroidal flux, as opposed to a 2D surface integral of the magnetic field dotted with ``n_zeta``. +- 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``. v0.12.1 From 68beebcbb9eadc73609bfd950ee756159f1be5c3 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 8 Aug 2024 17:19:58 -0500 Subject: [PATCH 39/43] remove forgotten arg --- tests/test_objective_funs.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 0eeefa9c35..cc5184e03c 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -2123,9 +2123,7 @@ def test_compute_scalar_resolution_toroidal_flux_B(self): 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_vector_potential=False), use_jit=False - ) + 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) @@ -2437,9 +2435,7 @@ def test_objective_no_nangrad_toroidal_flux(self): g = obj.grad(obj.x(ext_field)) assert not np.any(np.isnan(g)), "toroidal flux A" - obj = ObjectiveFunction( - ToroidalFlux(eq, ext_field, use_vector_potential=False), use_jit=False - ) + 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" From 9134103e8cbd31d11fdbfe8dc9df5a6330a08e62 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 9 Aug 2024 14:04:11 -0500 Subject: [PATCH 40/43] fix list of expected errors --- desc/objectives/_coils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index 14c4c7165d..459bd821a3 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -1414,7 +1414,7 @@ def build(self, use_jit=True, verbose=1): self._use_vector_potential = True try: self._field.compute_magnetic_vector_potential([0, 0, 0]) - except NotImplementedError: + except (NotImplementedError, ValueError): self._use_vector_potential = False if self._eval_grid is None: eval_grid = LinearGrid( From 7bb7bd3a4f8b3269cbd370e1e91591d0a40c147c Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 29 Aug 2024 12:22:30 -0400 Subject: [PATCH 41/43] remove _AB fxns for simpler fields --- desc/magnetic_fields/_core.py | 182 ++++++++++------------------------ 1 file changed, 55 insertions(+), 127 deletions(-) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index 3d1c1c85bd..ab5f852bdc 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -1074,16 +1074,10 @@ def B0(self): def B0(self, new): self._B0 = float(np.squeeze(new)) - def _compute_A_or_B( - self, - coords, - params=None, - basis="rpz", - source_grid=None, - transforms=None, - compute_A_or_B="B", + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): - """Compute magnetic field or vector potential at a set of points. + """Compute magnetic field at a set of points. Parameters ---------- @@ -1098,21 +1092,13 @@ def _compute_A_or_B( 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 or vector potential at specified points + 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}', - ) params = setdefault(params, {}) B0 = params.get("B0", self.B0) R0 = params.get("R0", self.R0) @@ -1121,48 +1107,13 @@ def _compute_A_or_B( coords = jnp.atleast_2d(jnp.asarray(coords)) if basis == "xyz": coords = xyz2rpz(coords) - if compute_A_or_B == "B": - bp = B0 * R0 / coords[:, 0] - brz = jnp.zeros_like(bp) - B = jnp.array([brz, bp, brz]).T - if basis == "xyz": - B = rpz2xyz_vec(B, phi=coords[:, 1]) - - return B - elif compute_A_or_B == "A": - 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 - - 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 - Dict of values for R0 and B0. - basis : {"rpz", "xyz"} - Basis for input coordinates and returned magnetic field. - 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 - ------- - field : ndarray, shape(N,3) - magnetic field at specified points + bp = B0 * R0 / coords[:, 0] + brz = jnp.zeros_like(bp) + B = jnp.array([brz, bp, brz]).T + if basis == "xyz": + B = rpz2xyz_vec(B, phi=coords[:, 1]) - """ - return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "B") + return B def compute_magnetic_vector_potential( self, coords, params=None, basis="rpz", source_grid=None, transforms=None @@ -1191,7 +1142,20 @@ def compute_magnetic_vector_potential( magnetic vector potential at specified points """ - return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") + 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): @@ -1221,71 +1185,6 @@ def B0(self): def B0(self, new): self._B0 = float(np.squeeze(new)) - 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 magnetic 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 - Dict of values for B0. - basis : {"rpz", "xyz"} - Basis for input coordinates and returned magnetic field. - source_grid : Grid, int or None or array-like, optional - 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 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}', - ) - params = setdefault(params, {}) - B0 = params.get("B0", self.B0) - - coords = jnp.atleast_2d(jnp.asarray(coords)) - if compute_A_or_B == "B": - bz = B0 * jnp.ones_like(coords[:, 2]) - brp = jnp.zeros_like(bz) - B = jnp.array([brp, brp, bz]).T - # b/c it only has a nonzero z component, no need - # to switch bases back if xyz is given - - return B - elif compute_A_or_B == "A": - 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 - def compute_magnetic_field( self, coords, params=None, basis="rpz", source_grid=None, transforms=None ): @@ -1311,7 +1210,17 @@ def compute_magnetic_field( magnetic field at specified points """ - return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "B") + params = setdefault(params, {}) + B0 = params.get("B0", self.B0) + + coords = jnp.atleast_2d(jnp.asarray(coords)) + bz = B0 * jnp.ones_like(coords[:, 2]) + brp = jnp.zeros_like(bz) + B = jnp.array([brp, brp, bz]).T + # 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 @@ -1340,7 +1249,26 @@ def compute_magnetic_vector_potential( magnetic vector potential at specified points """ - return self._compute_A_or_B(coords, params, basis, source_grid, transforms, "A") + 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): From cd2a04121e7a913efe7fae639a682186f9f84149 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 29 Aug 2024 12:23:36 -0400 Subject: [PATCH 42/43] address comment --- desc/magnetic_fields/_core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index ab5f852bdc..760f4e372b 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -613,7 +613,7 @@ def save_mgrid( ) bz_001[:] = B_Z - if A_R is not None: + 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 " From 2b99cc818f0c7a789b3253edd1cf49cdbea64e54 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 29 Aug 2024 12:40:14 -0400 Subject: [PATCH 43/43] remove todo --- desc/objectives/_coils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index 4b0c117c3a..c7575b136b 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -1520,7 +1520,6 @@ def compute(self, field_params=None, constants=None): ) A_dot_e_theta = jnp.sum(A * data["e_theta"], axis=1) - # TODO: use the line integral compute utilities Psi = jnp.sum(grid.spacing[:, 1] * A_dot_e_theta) else: B = constants["field"].compute_magnetic_field(