diff --git a/desc/compute/_surface.py b/desc/compute/_surface.py index 7ed57bf84..729760578 100644 --- a/desc/compute/_surface.py +++ b/desc/compute/_surface.py @@ -462,7 +462,7 @@ def _Phi_z_CurrentPotentialField(params, transforms, profiles, data, **kwargs): ], ) def _K_sup_theta_CurrentPotentialField(params, transforms, profiles, data, **kwargs): - data["K^theta"] = -data["Phi_z"] * (1 / data["|e_theta x e_zeta|"]) + data["K^theta"] = -data["Phi_z"] / data["|e_theta x e_zeta|"] return data @@ -484,7 +484,7 @@ def _K_sup_theta_CurrentPotentialField(params, transforms, profiles, data, **kwa ], ) def _K_sup_zeta_CurrentPotentialField(params, transforms, profiles, data, **kwargs): - data["K^zeta"] = data["Phi_t"] * (1 / data["|e_theta x e_zeta|"]) + data["K^zeta"] = data["Phi_t"] / data["|e_theta x e_zeta|"] return data @@ -507,9 +507,10 @@ def _K_sup_zeta_CurrentPotentialField(params, transforms, profiles, data, **kwar ], ) def _K_CurrentPotentialField(params, transforms, profiles, data, **kwargs): - data["K"] = (data["K^zeta"] * data["e_zeta"].T).T + ( - data["K^theta"] * data["e_theta"].T - ).T + data["K"] = ( + data["K^zeta"][:, jnp.newaxis] * data["e_zeta"] + + data["K^theta"][:, jnp.newaxis] * data["e_theta"] + ) return data diff --git a/desc/integrals/singularities.py b/desc/integrals/singularities.py index ab2371a83..3326b756a 100644 --- a/desc/integrals/singularities.py +++ b/desc/integrals/singularities.py @@ -5,13 +5,16 @@ import numpy as np import scipy from interpax import fft_interp2d +from jax import jacfwd +from scipy.constants import mu_0 from desc.backend import fori_loop, jnp, put, vmap from desc.basis import DoubleFourierSeries from desc.compute.geom_utils import rpz2xyz, rpz2xyz_vec, xyz2rpz_vec from desc.grid import LinearGrid from desc.io import IOAble -from desc.utils import isalmostequal, islinspaced, safediv, safenorm +from desc.transform import Transform +from desc.utils import dot, isalmostequal, islinspaced, safediv, safenorm, setdefault def _get_quadrature_nodes(q): @@ -704,11 +707,55 @@ def _kernel_biot_savart_A(eval_data, source_data, diag=False): _kernel_biot_savart_A.keys = ["R", "phi", "Z", "K_vc"] +def _kernel_Bn_over_r(eval_data, source_data, diag=False): + # B dot n' / |dx| + source_x = jnp.atleast_2d( + rpz2xyz(jnp.array([source_data["R"], source_data["phi"], source_data["Z"]]).T) + ) + eval_x = jnp.atleast_2d( + rpz2xyz(jnp.array([eval_data["R"], eval_data["phi"], eval_data["Z"]]).T) + ) + if diag: + dx = eval_x - source_x + else: + dx = eval_x[:, None] - source_x[None] + return safediv(source_data["Bn"], safenorm(dx, axis=-1)) + + +_kernel_Bn_over_r.ndim = 1 +_kernel_Bn_over_r.keys = ["R", "phi", "Z", "Bn"] + + +def _kernel_Phi_dG_dn(eval_data, source_data, diag=False): + # Phi(x') * dG(x,x')/dn' = Phi' * n' dot dx / |dx|^3 + # where Phi has units tesla-meters. + source_x = jnp.atleast_2d( + rpz2xyz(jnp.array([source_data["R"], source_data["phi"], source_data["Z"]]).T) + ) + eval_x = jnp.atleast_2d( + rpz2xyz(jnp.array([eval_data["R"], eval_data["phi"], eval_data["Z"]]).T) + ) + if diag: + dx = eval_x - source_x + else: + dx = eval_x[:, None] - source_x[None] + n = rpz2xyz_vec(source_data["n_rho"], phi=source_data["phi"]) + return safediv( + source_data["Phi"] * dot(n, dx), + safenorm(dx, axis=-1) ** 3, + ) + + +_kernel_Phi_dG_dn.ndim = 1 +_kernel_Phi_dG_dn.keys = ["R", "phi", "Z", "Phi", "n_rho"] + kernels = { "1_over_r": _kernel_1_over_r, "nr_over_r3": _kernel_nr_over_r3, "biot_savart": _kernel_biot_savart, "biot_savart_A": _kernel_biot_savart_A, + "_kernel_Bn_over_r": _kernel_Bn_over_r, + "_kernel_Phi_dG_dn": _kernel_Phi_dG_dn, } @@ -856,3 +903,544 @@ def compute_B_plasma(eq, eval_grid, source_grid=None, normal_only=False): if normal_only: Bplasma = jnp.sum(Bplasma * eval_data["n_rho"], axis=1) return Bplasma + + +def compute_B_laplace( + eq, + B0, + eval_grid, + source_grid=None, + Phi_grid=None, + Phi_M=None, + Phi_N=None, + sym=False, +): + """Compute magnetic field in interior of plasma due to vacuum potential. + + Let D, D^∁ denote the interior, exterior of a toroidal region with + boundary βˆ‚D. Computes the magnetic field 𝐁 in units of Tesla such that + + - 𝐁 = 𝐁₀ + βˆ‡Ξ¦ on D + - βˆ‡ Γ— 𝐁 = βˆ‡ Γ— 𝐁₀ on D βˆͺ D^∁ (i.e. βˆ‡Ξ¦ is single-valued or periodic) + - βˆ‡ Γ— 𝐁₀ = ΞΌβ‚€ 𝐉 on D βˆͺ D^∁ + - 𝐁 * βˆ‡Ο = 0 on βˆ‚D + - βˆ‡Β²Ξ¦ = 0 on D + + Examples + -------- + In a vacuum, the magnetic field may be written 𝐁 = βˆ‡π›·. + The solution to βˆ‡Β²π›· = 0, under a homogenous boundary + condition 𝐁 * βˆ‡Ο = 0, is 𝛷 = 0. To obtain a non-trivial solution, + the boundary condition may be modified. + Let 𝐁 = 𝐁₀ + βˆ‡Ξ¦. + If 𝐁₀ β‰  0 and satisfies βˆ‡ Γ— 𝐁₀ = 0, then βˆ‡Β²Ξ¦ = 0 solved + under an inhomogeneous boundary condition yields a non-trivial solution. + If 𝐁₀ β‰  -βˆ‡Ξ¦, then 𝐁 β‰  0. + + Parameters + ---------- + eq : Equilibrium + Configuration with surface geometry defining βˆ‚D. + B0 : MagneticField + Magnetic field such that βˆ‡ Γ— 𝐁₀ = ΞΌβ‚€ 𝐉 + where 𝐉 is the current in amperes everywhere. + eval_grid : Grid + Evaluation points on D for the magnetic field. + source_grid : Grid + Source points on βˆ‚D for quadrature of kernels. + Resolution determines the accuracy of the boundary condition, + and evaluation of the magnetic field. + Phi_grid : Grid + Source points on βˆ‚D. + Resolution determines accuracy of the spectral coefficients. + Should have resolution at least ``M=Phi_M`` and ``N=Phi_N``. + Phi_M : int + Number of poloidal Fourier modes. + Default is ``eq.M``. + Phi_N : int + Number of toroidal Fourier modes. + Default is ``eq.N``. + sym : str + ``DoubleFourierSeries`` basis symmetry. + Default is no symmetry. + + Returns + ------- + B : jnp.ndarray + Magnetic field evaluated on ``eval_grid``. + + """ + from desc.magnetic_fields import FourierCurrentPotentialField + + if source_grid is None: + source_grid = LinearGrid( + rho=np.array([1.0]), + M=eq.M_grid, + N=eq.N_grid, + NFP=eq.NFP if eq.N > 0 else 64, + sym=False, + ) + Bn, _ = B0.compute_Bnormal( + eq.surface, eval_grid=source_grid, source_grid=source_grid + ) + Phi_mn, Phi_transform = compute_Phi_mn( + eq, Bn, source_grid, Phi_grid, Phi_M, Phi_N, sym + ) + # 𝐁 - 𝐁₀ = βˆ‡Ξ¦ = 𝐁_vacuum in the interior. + # Merkel eq. 1.4 is the Green's function solution to βˆ‡Β²Ξ¦ = 0 in the interior. + # Note that 𝐁₀′ in eq. 3.5 has the wrong sign. + grad_Phi = FourierCurrentPotentialField.from_surface( + eq.surface, Phi_mn / mu_0, Phi_transform.basis.modes[:, 1:] + ) + data = eq.compute(["R", "phi", "Z"], grid=eval_grid) + coords = jnp.column_stack([data["R"], data["phi"], data["Z"]]) + B = (B0 + grad_Phi).compute_magnetic_field(coords, source_grid=source_grid) + return B + + +def compute_Phi_mn( + eq, + B0n, + source_grid, + Phi_grid=None, + Phi_M=None, + Phi_N=None, + sym=False, +): + """Compute Fourier coefficients of vacuum potential Ξ¦ on βˆ‚D. + + Let D, D^∁ denote the interior, exterior of a toroidal region with + boundary βˆ‚D. Computes the magnetic field 𝐁 in units of Tesla such that + + - 𝐁 = 𝐁₀ + βˆ‡Ξ¦ on D + - βˆ‡ Γ— 𝐁 = βˆ‡ Γ— 𝐁₀ on D βˆͺ D^∁ (i.e. βˆ‡Ξ¦ is single-valued or periodic) + - βˆ‡ Γ— 𝐁₀ = ΞΌβ‚€ 𝐉 on D βˆͺ D^∁ + - 𝐁 * βˆ‡Ο = 0 on βˆ‚D + - βˆ‡Β²Ξ¦ = 0 on D + + Parameters + ---------- + eq : Equilibrium + Configuration with surface geometry defining βˆ‚D. + B0n : MagneticField + 𝐁₀ * βˆ‡Ο / |βˆ‡Ο| evaluated on ``source_grid`` of magnetic field + such that βˆ‡ Γ— 𝐁₀ = ΞΌβ‚€ 𝐉 where 𝐉 is the current in amperes everywhere. + source_grid : Grid + Source points on βˆ‚D for quadrature of kernels. + Resolution determines the accuracy of the boundary condition. + Phi_grid : Grid + Source points on βˆ‚D. + Resolution determines accuracy of the spectral coefficients. + Should have resolution at least ``M=Phi_M`` and ``N=Phi_N``. + Phi_M : int + Number of poloidal Fourier modes. + Default is ``eq.M``. + Phi_N : int + Number of toroidal Fourier modes. + Default is ``eq.N``. + sym : str + ``DoubleFourierSeries`` basis symmetry. + Default is no symmetry. + + Returns + ------- + Phi_mn, Phi_transform : jnp.ndarray, Transform + Fourier coefficients of Ξ¦ on βˆ‚D. + + """ + basis = DoubleFourierSeries( + M=setdefault(Phi_M, eq.M), + N=setdefault(Phi_N, eq.N), + NFP=eq.NFP, + sym=sym, + ) + if Phi_grid is None: + Phi_grid = LinearGrid( + rho=np.array([1.0]), + M=2 * basis.M, + N=2 * basis.N, + NFP=basis.NFP if basis.N > 0 else 64, + sym=False, + ) + assert source_grid.is_meshgrid + assert Phi_grid.is_meshgrid + + # Malhotra recommends s = q = N⁰ᐧ²⁡ where NΒ² is num theta * num zeta*NFP, + # but using same defaults as above. + k = min(source_grid.num_theta, source_grid.num_zeta * source_grid.NFP) + s = k - 1 + q = k // 2 + int(np.sqrt(k)) + try: + interpolator = FFTInterpolator(Phi_grid, source_grid, s, q) + except AssertionError as e: + print( + f"Unable to create FFTInterpolator, got error {e}," + "falling back to DFT method which is much slower" + ) + interpolator = DFTInterpolator(Phi_grid, source_grid, s, q) + + names = ["R", "phi", "Z"] + Phi_data = eq.compute(names, grid=Phi_grid) + src_data = eq.compute(names + ["n_rho"], grid=source_grid) + src_data["Bn"] = B0n + src_transform = Transform(source_grid, basis) + Phi_transform = Transform(Phi_grid, basis) + + def LHS(Phi_mn): + # After Fourier transform, the LHS is linear in the spectral coefficients Ξ¦β‚˜β‚™. + # We approximate this as finite-dimensional, which enables writing the left + # hand side as A @ Ξ¦β‚˜β‚™. Then Ξ¦β‚˜β‚™ is found by solving LHS(Ξ¦β‚˜β‚™) = A @ Ξ¦β‚˜β‚™ = RHS. + src_data["Phi"] = src_transform.transform(Phi_mn) + I = singular_integral( + Phi_data, + src_data, + _kernel_Phi_dG_dn, + interpolator, + loop=True, + ).squeeze() + Phi = Phi_transform.transform(Phi_mn) + return Phi + I / (2 * jnp.pi) + + # LHS is expensive, so it is better to construct full Jacobian once + # rather than iterative solves like jax.scipy.sparse.linalg.cg. + A = jacfwd(LHS)(jnp.ones(basis.num_modes)) + RHS = -singular_integral( + Phi_data, + src_data, + _kernel_Bn_over_r, + interpolator, + loop=True, + ).squeeze() / (2 * jnp.pi) + # Fourier coefficients of Ξ¦ on boundary + Phi_mn, _, _, _ = jnp.linalg.lstsq(A, RHS) + # np.testing.assert_allclose(LHS(Phi_mn), A @ Phi_mn, atol=1e-12) # noqa: E801 + return Phi_mn, Phi_transform + + +def compute_dPhi_dn(eq, eval_grid, source_grid, Phi_mn, basis): + """Computes vacuum field βˆ‡Ξ¦ β‹… n on βˆ‚D. + + Let D, D^∁ denote the interior, exterior of a toroidal region with + boundary βˆ‚D. Computes the magnetic field 𝐁 in units of Tesla such that + + - 𝐁 = 𝐁₀ + βˆ‡Ξ¦ on D + - βˆ‡ Γ— 𝐁 = βˆ‡ Γ— 𝐁₀ on D βˆͺ D^∁ (i.e. βˆ‡Ξ¦ is single-valued or periodic) + - βˆ‡ Γ— 𝐁₀ = ΞΌβ‚€ 𝐉 on D βˆͺ D^∁ + - 𝐁 * βˆ‡Ο = 0 on βˆ‚D + - βˆ‡Β²Ξ¦ = 0 on D + + Parameters + ---------- + eq : Equilibrium + Configuration with surface geometry defining βˆ‚D. + eval_grid : Grid + Evaluation points on βˆ‚D. + source_grid : Grid + Source points on βˆ‚D for quadrature of kernels. + Phi_mn : jnp.ndarray + Fourier coefficients of Ξ¦ on the boundary. + basis : DoubleFourierSeries + Basis for Ξ¦β‚˜β‚™. + + Returns + ------- + dPhi_dn : jnp.ndarray + Shape (``eval_grid.grid.num_nodes``, 3). + Vacuum field βˆ‡Ξ¦ β‹… n on βˆ‚D. + + """ + k = min(source_grid.num_theta, source_grid.num_zeta * source_grid.NFP) + s = k - 1 + q = k // 2 + int(np.sqrt(k)) + try: + interpolator = FFTInterpolator(eval_grid, source_grid, s, q) + except AssertionError as e: + print( + f"Unable to create FFTInterpolator, got error {e}," + "falling back to DFT method which is much slower" + ) + interpolator = DFTInterpolator(eval_grid, source_grid, s, q) + + names = ["R", "phi", "Z"] + evl_data = eq.compute(names + ["n_rho"], grid=eval_grid) + transform = Transform(source_grid, basis, derivs=1) + src_data = { + "Phi_t": transform.transform(Phi_mn, dt=1), + "Phi_z": transform.transform(Phi_mn, dz=1), + } + src_data = eq.compute( + names + ["|e_theta x e_zeta|", "e_theta", "e_zeta"], + grid=source_grid, + data=src_data, + ) + src_data["K^theta"] = -src_data["Phi_z"] / src_data["|e_theta x e_zeta|"] + src_data["K^zeta"] = src_data["Phi_t"] / src_data["|e_theta x e_zeta|"] + src_data["K_vc"] = ( + src_data["K^theta"][:, jnp.newaxis] * src_data["e_theta"] + + src_data["K^zeta"][:, jnp.newaxis] * src_data["e_zeta"] + ) + + # βˆ‡Ξ¦ = βˆ‚Ξ¦/βˆ‚Ο βˆ‡Ο + βˆ‚Ξ¦/βˆ‚ΞΈ βˆ‡ΞΈ + βˆ‚Ξ¦/βˆ‚ΞΆ βˆ‡ΞΆ + # but we can not obtain βˆ‚Ξ¦/βˆ‚Ο from Ξ¦β‚˜β‚™. Biot-Savart gives + # K_vc = n Γ— βˆ‡Ξ¦ where Ξ¦ has units Tesla-meters + # βˆ‡Ξ¦(x ∈ βˆ‚D) dot n = [1/2Ο€ ∫_βˆ‚D df' K_vc Γ— βˆ‡G(x,x')] dot n + # (Same instructions but divide by 2 for x ∈ D). + # Biot-Savart kernel assumes Ξ¦ in amperes, so we account for that. + dPhi_dn = (2 / mu_0) * dot( + singular_integral( + evl_data, src_data, _kernel_biot_savart, interpolator, loop=True + ), + evl_data["n_rho"], + ) + return dPhi_dn + + +# TODO: surface integral correctness validation: should match output of compute_dPhi_dn. +def _dPhi_dn_triple_layer( + eq, + B0n, + eval_grid, + source_grid, + Phi_mn, + basis, +): + """Compute βˆ‡Ξ¦ β‹… n on βˆ‚D. + + Parameters + ---------- + eq : Equilibrium + Configuration with surface geometry defining βˆ‚D. + B0n : MagneticField + 𝐁₀ * βˆ‡Ο / |βˆ‡Ο| evaluated on ``source_grid`` of magnetic field + such that βˆ‡ Γ— 𝐁₀ = ΞΌβ‚€ 𝐉 where 𝐉 is the current in amperes everywhere. + eval_grid : Grid + Evaluation points on βˆ‚D. + source_grid : Grid + Source points on βˆ‚D for quadrature of kernels. + Phi_mn : jnp.ndarray + Fourier coefficients of Ξ¦ on the boundary. + basis : DoubleFourierSeries + Basis for Ξ¦β‚˜β‚™. + + Returns + ------- + dPhi_dn : jnp.ndarray + Shape (``Phi_trans.grid.num_nodes``, ). + βˆ‡Ξ¦ β‹… n on βˆ‚D. + + """ + k = min(source_grid.num_theta, source_grid.num_zeta * source_grid.NFP) + s = k - 1 + q = k // 2 + int(np.sqrt(k)) + try: + interpolator = FFTInterpolator(eval_grid, source_grid, s, q) + except AssertionError as e: + print( + f"Unable to create FFTInterpolator, got error {e}," + "falling back to DFT method which is much slower" + ) + interpolator = DFTInterpolator(eval_grid, source_grid, s, q) + + names = ["R", "phi", "Z", "n_rho"] + evl_data = eq.compute(names, grid=eval_grid) + src_data = eq.compute(names, grid=source_grid) + src_data["Bn"] = B0n + I2 = -singular_integral( + evl_data, + src_data, + _kernel_Bn_grad_G_dot_n, + interpolator, + loop=True, + ) + src_data["Phi"] = Transform(source_grid, basis).transform(Phi_mn) + I1 = singular_integral( + evl_data, + src_data, + # triple layer kernel may need more resolution + _kernel_Phi_grad_dG_dn_dot_m, + interpolator, + loop=True, + ) + dPhi_dn = -I1 + I2 + return dPhi_dn + + +def _kernel_Phi_grad_dG_dn_dot_m(eval_data, source_data, diag=False): + # Phi(x') * grad dG(x,x')/dn' dot n + # = Phi' * n' dot (grad(dx) / |dx|^3 - 3 dx transpose(dx) / |dx|^5) dot n + # = Phi' * n' dot (n / |dx|^3 - 3 dx (dx dot n) / |dx|^5) + # = Phi' * n' dot [n |dx|^2 - 3 dx (dx dot n)] / |dx|^5 + # where Phi has units tesla-meters. + source_x = jnp.atleast_2d( + rpz2xyz(jnp.array([source_data["R"], source_data["phi"], source_data["Z"]]).T) + ) + eval_x = jnp.atleast_2d( + rpz2xyz(jnp.array([eval_data["R"], eval_data["phi"], eval_data["Z"]]).T) + ) + if diag: + dx = eval_x - source_x + else: + dx = eval_x[:, None] - source_x[None] + # this is n' + nn = rpz2xyz_vec(source_data["n_rho"], phi=source_data["phi"]) + # this is n + n = rpz2xyz_vec(eval_data["n_rho"], phi=eval_data["phi"]) + dx_norm = safenorm(dx, axis=-1) + return safediv( + source_data["Phi"] * (dot(nn, n) * dx_norm**2 - 3 * dot(nn, dx) * dot(dx, n)), + dx_norm**5, + ) + + +def _kernel_Bn_grad_G_dot_n(eval_data, source_data, diag=False): + # Bn(x') * dG(x,x')/dn = - Bn * n dot dx / |dx|^3 + source_x = jnp.atleast_2d( + rpz2xyz(jnp.array([source_data["R"], source_data["phi"], source_data["Z"]]).T) + ) + eval_x = jnp.atleast_2d( + rpz2xyz(jnp.array([eval_data["R"], eval_data["phi"], eval_data["Z"]]).T) + ) + if diag: + dx = eval_x - source_x + else: + dx = eval_x[:, None] - source_x[None] + n = rpz2xyz_vec(eval_data["n_rho"], phi=eval_data["phi"]) + return safediv( + -source_data["Bn"] * dot(n, dx), + safenorm(dx, axis=-1) ** 3, + ) + + +def compute_K_mn( + eq, + G, + grid=None, + K_M=None, + K_N=None, + sym=False, +): + """Compute Fourier coefficients of surface current on βˆ‚D. + + Parameters + ---------- + eq : Equilibrium + Configuration with surface geometry defining βˆ‚D. + G : float + Secular term of poloidal current in amperes. + Should be ``2*np.pi/mu_0*data["G"]``. + grid : Grid + Points on βˆ‚D. + K_M : int + Number of poloidal Fourier modes. + Default is ``eq.M``. + K_N : int + Number of toroidal Fourier modes. + Default is ``eq.N``. + sym : str + ``DoubleFourierSeries`` basis symmetry. + Default is no symmetry. + + Returns + ------- + K_mn, K_sec, K_transform : jnp.ndarray, Transform + Fourier coefficients of surface current on βˆ‚D. + + """ + from desc.magnetic_fields import FourierCurrentPotentialField + + K_sec = FourierCurrentPotentialField.from_surface(eq.surface, G=G) + basis = DoubleFourierSeries( + M=setdefault(K_M, eq.M), + N=setdefault(K_N, eq.N), + NFP=eq.NFP, + sym=sym, + ) + if grid is None: + grid = LinearGrid( + rho=np.array([1.0]), + # 3x higher than typical since we need to fit a vector + M=6 * basis.M, + N=6 * basis.N, + NFP=basis.NFP if basis.N > 0 else 64, + sym=False, + ) + assert grid.is_meshgrid + transform = Transform(grid, basis) + K_secular = K_sec.compute("K", grid=grid)["K"] + n = eq.compute("n_rho", grid=grid)["n_rho"] + + def LHS(K_mn): + # After Fourier transform, the LHS is linear in the spectral coefficients Ξ¦β‚˜β‚™. + # We approximate this as finite-dimensional, which enables writing the left + # hand side as A @ Ξ¦β‚˜β‚™. Then Ξ¦β‚˜β‚™ is found by solving LHS(Ξ¦β‚˜β‚™) = A @ Ξ¦β‚˜β‚™ = RHS. + num_coef = K_mn.size // 3 + K_R = transform.transform(K_mn[:num_coef]) + K_phi = transform.transform(K_mn[num_coef : 2 * num_coef]) + K_Z = transform.transform(K_mn[2 * num_coef :]) + K_fourier = jnp.column_stack([K_R, K_phi, K_Z]) + return dot(K_fourier + K_secular, n) + + A = jacfwd(LHS)(jnp.ones(basis.num_modes * 3)) + K_mn, _, _, _ = jnp.linalg.lstsq(A, jnp.zeros(grid.num_nodes)) + # np.testing.assert_allclose(LHS(K_mn), A @ K_mn, atol=1e-8) # noqa: E801 + return K_mn, K_sec, transform + + +def compute_B_dot_n_from_K(eq, eval_grid, source_grid, K_mn, K_sec, basis): + """Computes B β‹… n on βˆ‚D from surface current. + + Parameters + ---------- + eq : Equilibrium + Configuration with surface geometry defining βˆ‚D. + eval_grid : Grid + Evaluation points on βˆ‚D. + source_grid : Grid + Source points on βˆ‚D for quadrature of kernels. + K_mn : jnp.ndarray + Fourier coefficients of surface current on βˆ‚D. + K_sec : FourierCurrentPotentialField + Secular part of Fourier current potential field. + basis : DoubleFourierSeries + Basis for Kβ‚˜β‚™. + + Returns + ------- + B_dot_n : jnp.ndarray + Shape (``eval_grid.grid.num_nodes``, 3). + + """ + k = min(source_grid.num_theta, source_grid.num_zeta * source_grid.NFP) + s = k - 1 + q = k // 2 + int(np.sqrt(k)) + try: + interpolator = FFTInterpolator(eval_grid, source_grid, s, q) + except AssertionError as e: + print( + f"Unable to create FFTInterpolator, got error {e}," + "falling back to DFT method which is much slower" + ) + interpolator = DFTInterpolator(eval_grid, source_grid, s, q) + + names = ["R", "phi", "Z"] + evl_data = eq.compute(names + ["n_rho"], grid=eval_grid) + transform = Transform(source_grid, basis) + src_data = eq.compute(names + ["|e_theta x e_zeta|"], grid=source_grid) + num_coef = K_mn.size // 3 + K_R = transform.transform(K_mn[:num_coef]) + K_phi = transform.transform(K_mn[num_coef : 2 * num_coef]) + K_Z = transform.transform(K_mn[2 * num_coef :]) + K_fourier = jnp.column_stack([K_R, K_phi, K_Z]) + K_secular = K_sec.compute("K", grid=source_grid)["K"] + src_data["K_vc"] = K_fourier + K_secular + + # Biot-Savart gives K_vc = n Γ— B + # B(x ∈ βˆ‚D) dot n = [ΞΌβ‚€/2Ο€ ∫_βˆ‚D df' K_vc Γ— βˆ‡G(x,x')] dot n + # (Same instructions but divide by 2 for x ∈ D). + Bn = 2 * dot( + singular_integral( + evl_data, src_data, _kernel_biot_savart, interpolator, loop=True + ), + evl_data["n_rho"], + ) + return Bn diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index b8af799d8..e90744bfc 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -33,7 +33,7 @@ from desc.io import IOAble from desc.optimizable import Optimizable, OptimizableCollection, optimizable_parameter from desc.transform import Transform -from desc.utils import copy_coeffs, errorif, flatten_list, setdefault, warnif +from desc.utils import copy_coeffs, dot, errorif, flatten_list, setdefault, warnif from desc.vmec_utils import ptolemy_identity_fwd, ptolemy_identity_rev @@ -344,7 +344,7 @@ def compute_Bnormal( B = self.compute_magnetic_field( coords, basis="rpz", source_grid=source_grid, params=params ) - Bnorm = jnp.sum(B * surf_normal, axis=-1) + Bnorm = dot(B, surf_normal) if calc_Bplasma: Bplasma = compute_B_plasma(eq, eval_grid, vc_source_grid, normal_only=True) diff --git a/tests/baseline/test_laplace_bdotn_eq0.png b/tests/baseline/test_laplace_bdotn_eq0.png new file mode 100644 index 000000000..6aea57944 Binary files /dev/null and b/tests/baseline/test_laplace_bdotn_eq0.png differ diff --git a/tests/baseline/test_laplace_bdotn_eq1.png b/tests/baseline/test_laplace_bdotn_eq1.png new file mode 100644 index 000000000..7bccf717c Binary files /dev/null and b/tests/baseline/test_laplace_bdotn_eq1.png differ diff --git a/tests/baseline/test_laplace_bdotn_from_K_eq0.png b/tests/baseline/test_laplace_bdotn_from_K_eq0.png new file mode 100644 index 000000000..04794fab5 Binary files /dev/null and b/tests/baseline/test_laplace_bdotn_from_K_eq0.png differ diff --git a/tests/baseline/test_laplace_bdotn_from_K_eq1.png b/tests/baseline/test_laplace_bdotn_from_K_eq1.png new file mode 100644 index 000000000..741b12985 Binary files /dev/null and b/tests/baseline/test_laplace_bdotn_from_K_eq1.png differ diff --git a/tests/baseline/test_laplace_bdotn_from_K_eq2.png b/tests/baseline/test_laplace_bdotn_from_K_eq2.png new file mode 100644 index 000000000..e97b3b61e Binary files /dev/null and b/tests/baseline/test_laplace_bdotn_from_K_eq2.png differ diff --git a/tests/baseline/test_laplace_dommaschk.png b/tests/baseline/test_laplace_dommaschk.png new file mode 100644 index 000000000..5b490b445 Binary files /dev/null and b/tests/baseline/test_laplace_dommaschk.png differ diff --git a/tests/test_integrals.py b/tests/test_integrals.py index 62e364946..5213c8f3d 100644 --- a/tests/test_integrals.py +++ b/tests/test_integrals.py @@ -9,6 +9,7 @@ from numpy.polynomial.chebyshev import chebinterpolate, chebroots from numpy.polynomial.legendre import leggauss from scipy import integrate +from scipy.constants import mu_0 from scipy.interpolate import CubicHermiteSpline from scipy.special import ellipe, ellipkm1 from tests.test_interp_utils import _f_1d, _f_1d_nyquist_freq @@ -19,6 +20,7 @@ from desc.equilibrium import Equilibrium from desc.equilibrium.coords import get_rtz_grid from desc.examples import get +from desc.geometry import FourierRZToroidalSurface from desc.grid import ConcentricGrid, Grid, LinearGrid, QuadratureGrid from desc.integrals import ( Bounce1D, @@ -54,8 +56,16 @@ leggauss_lob, tanh_sinh, ) -from desc.integrals.singularities import _get_quadrature_nodes +from desc.integrals.singularities import ( + _get_quadrature_nodes, + compute_B_dot_n_from_K, + compute_dPhi_dn, + compute_K_mn, + compute_Phi_mn, +) from desc.integrals.surface_integral import _get_grid_surface +from desc.magnetic_fields import ToroidalMagneticField +from desc.plotting import plot_boundary from desc.transform import Transform from desc.utils import dot, errorif, safediv @@ -724,6 +734,294 @@ def test_biest_interpolators(self): np.testing.assert_allclose(g1, g2) np.testing.assert_allclose(g1, ff) + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=False, tolerance=tol_1d) + @pytest.mark.parametrize("eq", [get("DSHAPE"), get("ESTELL")]) + def test_laplace_bdotn(self, eq): + """Test that Laplace solution satisfies boundary condition.""" + MN = 40 + src_grid = LinearGrid(M=MN, N=MN, sym=False, NFP=eq.NFP) + data = eq.compute(["G", "R0"], grid=src_grid) + + def test(G): + B0 = ToroidalMagneticField(B0=G / data["R0"], R0=data["R0"]) + B0n, _ = B0.compute_Bnormal( + eq.surface, eval_grid=src_grid, source_grid=src_grid + ) + Phi_mn, Phi_transform = compute_Phi_mn( + eq=eq, B0n=B0n, source_grid=src_grid, Phi_N=max(eq.N, 1) + ) + + evl_grid = Phi_transform.grid + dPhi_dn = compute_dPhi_dn( + eq=eq, + eval_grid=evl_grid, + source_grid=src_grid, + Phi_mn=Phi_mn, + basis=Phi_transform.basis, + ) + B0n, _ = B0.compute_Bnormal( + eq.surface, eval_grid=evl_grid, source_grid=src_grid + ) + + # Get data as function of theta, zeta on the only flux surface (LCFS). + Bn = evl_grid.meshgrid_reshape(B0n + dPhi_dn, "rtz")[0] + theta = evl_grid.meshgrid_reshape(evl_grid.nodes[:, 1], "rtz")[0] + zeta = evl_grid.meshgrid_reshape(evl_grid.nodes[:, 2], "rtz")[0] + + fig, ax = plt.subplots() + contour = ax.contourf(theta, zeta, Bn) + ax.set_title(r"$(\nabla \Phi + B_0) \cdot n$ on $\partial D$") + fig.colorbar(contour, ax=ax) + # FIXME: Doesn't pass unless G = 0 for stellarators. + try: + np.testing.assert_allclose(B0n + dPhi_dn, 0, err_msg=f"G = {G}") + except AssertionError as e: + print(e) + return fig, ax + + test(0) + fig, ax = test(src_grid.compress(data["G"])[-1]) + return fig + + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=False, tolerance=tol_1d) + @pytest.mark.parametrize("eq", [get("DSHAPE"), get("HELIOTRON"), get("ESTELL")]) + def test_laplace_bdotn_from_K(self, eq): + """Test that Laplace solution from fit of K satisfies boundary condition.""" + MN = 40 + grid = LinearGrid(M=MN, N=MN, sym=False, NFP=eq.NFP) + data = eq.compute(["G"], grid=grid) + + def test(G): + G_amp = 2 * np.pi * G / mu_0 + K_mn, K_sec, K_transform = compute_K_mn( + eq=eq, G=G_amp, grid=grid, K_N=max(eq.N, 1) + ) + Bn = compute_B_dot_n_from_K( + eq=eq, + eval_grid=grid, + source_grid=grid, + K_mn=K_mn, + K_sec=K_sec, + basis=K_transform.basis, + ) + # Get data as function of theta, zeta on the only flux surface (LCFS). + Bn = grid.meshgrid_reshape(Bn, "rtz")[0] + theta = grid.meshgrid_reshape(grid.nodes[:, 1], "rtz")[0] + zeta = grid.meshgrid_reshape(grid.nodes[:, 2], "rtz")[0] + + fig, ax = plt.subplots() + contour = ax.contourf(theta, zeta, Bn) + ax.set_title(r"$B \cdot n$ on $\partial D$") + fig.colorbar(contour, ax=ax) + # FIXME: Doesn't pass unless G = 0 for stellarators. + try: + np.testing.assert_allclose(Bn, 0, err_msg=f"G = {G}") + except AssertionError as e: + print(e) + return fig, ax + + test(0) + fig, ax = test(grid.compress(data["G"])[-1]) + return fig + + @staticmethod + def manual_transform(coef, m, n, theta, zeta): + """Evaluates Double Fourier Series of form G_n^m at theta and zeta pts.""" + op_four = np.where( + ((m < 0) & (n < 0))[:, np.newaxis], + np.sin(np.abs(m)[:, np.newaxis] * theta) + * np.sin(np.abs(n)[:, np.newaxis] * zeta), + n[:, np.newaxis] * zeta * np.nan, + ) + op_three = np.where( + ((m < 0) & (n >= 0))[:, np.newaxis], + np.sin(np.abs(m)[:, np.newaxis] * theta) * np.cos(n[:, np.newaxis] * zeta), + op_four, + ) + op_two = np.where( + ((m >= 0) & (n < 0))[:, np.newaxis], + np.cos(m[:, np.newaxis] * theta) * np.sin(np.abs(n)[:, np.newaxis] * zeta), + op_three, + ) + op_one = np.where( + ((m >= 0) & (n >= 0))[:, np.newaxis], + np.cos(m[:, np.newaxis] * theta) * np.cos(n[:, np.newaxis] * zeta), + op_two, + ) + return np.sum(coef[:, np.newaxis] * op_one, axis=0) + + @staticmethod + def merkel_transform(coef, m, n, theta, zeta, fun=np.cos): + """Evaluates double Fourier series of form cos(m theta + n zeta).""" + return np.sum( + coef[:, np.newaxis] + * fun(m[:, np.newaxis] * theta + n[:, np.newaxis] * zeta), + axis=0, + ) + + @staticmethod + def _merkel_surf(C_r, C_z): + """Convert merkel coefficients to DESC coefficients and check.""" + m_b = max( + max(abs(mn[0]) for mn in C_r.keys()), max(abs(mn[0]) for mn in C_z.keys()) + ) + n_b = max( + max(abs(mn[1]) for mn in C_r.keys()), max(abs(mn[1]) for mn in C_z.keys()) + ) + R_lmn = { + (m, n): 0.0 for m in range(-m_b, m_b + 1) for n in range(-n_b, n_b + 1) + } + Z_lmn = { + (m, n): 0.0 for m in range(-m_b, m_b + 1) for n in range(-n_b, n_b + 1) + } + for m in range(0, m_b + 1): + for n in range(-n_b, n_b + 1): + if (m, n) in C_r: + R_lmn[(m, abs(n))] += C_r[(m, n)] + if n < 0: + R_lmn[(-m, n)] += C_r[(m, n)] + elif n > 0: + R_lmn[(-m, -n)] -= C_r[(m, n)] + if (m, n) in C_z: + Z_lmn[(-m, abs(n))] += C_z[(m, n)] + if n < 0: + Z_lmn[(m, n)] -= C_z[(m, n)] + elif n > 0: + Z_lmn[(m, -n)] += C_z[(m, n)] + + grid = LinearGrid(rho=1, M=5, N=5) + R_bench = TestSingularities.manual_transform( + np.array(list(R_lmn.values())), + np.array([mn[0] for mn in R_lmn.keys()]), + np.array([mn[1] for mn in R_lmn.keys()]), + -grid.nodes[:, 1], # theta is flipped + grid.nodes[:, 2], + ) + R_merk = TestSingularities.merkel_transform( + np.array(list(C_r.values())), + np.array([mn[0] for mn in C_r.keys()]), + np.array([mn[1] for mn in C_r.keys()]), + -grid.nodes[:, 1], # theta is flipped + grid.nodes[:, 2], + ) + Z_bench = TestSingularities.manual_transform( + np.array(list(Z_lmn.values())), + np.array([mn[0] for mn in Z_lmn.keys()]), + np.array([mn[1] for mn in Z_lmn.keys()]), + -grid.nodes[:, 1], # theta is flipped + grid.nodes[:, 2], + ) + Z_merk = TestSingularities.merkel_transform( + np.array(list(C_z.values())), + np.array([mn[0] for mn in C_z.keys()]), + np.array([mn[1] for mn in C_z.keys()]), + -grid.nodes[:, 1], # theta is flipped + grid.nodes[:, 2], + fun=np.sin, + ) + np.testing.assert_allclose(R_bench, R_merk) + np.testing.assert_allclose(Z_bench, Z_merk) + surf = FourierRZToroidalSurface( + R_lmn=list(R_lmn.values()), + Z_lmn=list(Z_lmn.values()), + modes_R=list(R_lmn.keys()), + modes_Z=list(Z_lmn.keys()), + ) + surf_data = surf.compute(["R", "Z"], grid=grid) + np.testing.assert_allclose(surf_data["R"], R_merk) + np.testing.assert_allclose(surf_data["Z"], Z_merk) + return surf + + @pytest.mark.unit + @pytest.mark.mpl_image_compare(remove_text=False, tolerance=tol_1d) + def test_laplace_dommaschk(self): + """Use Dommaschk potentials to generate benchmark test and compare.""" + C_r = { + (0, -2): 0.000056, + (0, -1): -0.000921, + (0, 0): 0.997922, + (0, 1): -0.000921, + (0, 2): 0.000056, + (1, -2): -0.000067, + (1, -1): -0.034645, + (1, 0): 0.093260, + (1, 1): 0.000880, + (1, 2): 0.000178, + (2, -2): 0.000373, + (2, -1): 0.000575, + (2, 0): 0.002916, + (2, 1): -0.000231, + (2, 2): 0.000082, + (3, -2): 0.000462, + (3, -1): -0.001509, + (3, 0): 0.001748, + (3, 1): -0.000239, + (3, 2): 0.000052, + } + C_z = { + (0, -2): 0.000076, + (0, -1): 0.000923, + (0, 0): 0.000000, + (0, 1): -0.000923, + (0, 2): -0.000076, + (1, -2): 0.000069, + (1, -1): 0.035178, + (1, 0): 0.099830, + (1, 1): 0.000860, + (1, 2): -0.000179, + (2, -2): -0.000374, + (2, -1): 0.000257, + (2, 0): 0.003096, + (2, 1): 0.003021, + (2, 2): 0.000007, + (3, -2): -0.000518, + (3, -1): 0.002233, + (3, 0): 0.001828, + (3, 1): 0.000257, + (3, 2): 0.000035, + } + surf = self._merkel_surf(C_r, C_z) + eq = Equilibrium(surface=surf) + plot_boundary(eq, phi=[0, 1, 2]) + plt.show() + MN = 40 + grid = LinearGrid(M=MN, N=MN, sym=False, NFP=eq.NFP) + + # About half the error of the below test. + def test(G): + G_amp = 2 * np.pi * G / mu_0 + K_mn, K_sec, K_transform = compute_K_mn( + eq=eq, G=G_amp, grid=grid, K_N=max(eq.N, 1) + ) + Bn = compute_B_dot_n_from_K( + eq=eq, + eval_grid=grid, + source_grid=grid, + K_mn=K_mn, + K_sec=K_sec, + basis=K_transform.basis, + ) + # Get data as function of theta, zeta on the only flux surface (LCFS). + Bn = grid.meshgrid_reshape(Bn, "rtz")[0] + theta = grid.meshgrid_reshape(grid.nodes[:, 1], "rtz")[0] + zeta = grid.meshgrid_reshape(grid.nodes[:, 2], "rtz")[0] + + fig, ax = plt.subplots() + contour = ax.contourf(theta, zeta, Bn) + ax.set_title(r"$B \cdot n$ on $\partial D$") + fig.colorbar(contour, ax=ax) + # FIXME: Doesn't pass unless G = 0 for stellarators. + try: + np.testing.assert_allclose(Bn, 0, err_msg=f"G = {G}") + except AssertionError as e: + print(e) + return fig, ax + + fig, ax = test(1) + return fig + class TestBouncePoints: """Test that bounce points are computed correctly."""