diff --git a/CMakeLists.txt b/CMakeLists.txt index d9eaa534e..5860c8223 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -132,6 +132,10 @@ pybind11_add_module(${PROJECT_NAME} src/simsoptpp/magneticfield_wireframe.cpp src/simsoptpp/boozerradialinterpolant.cpp src/simsoptpp/wireframe_optimization.cpp + src/simsoptpp/python_currentpotential.cpp + src/simsoptpp/currentpotential.cpp + src/simsoptpp/currentpotentialfourier.cpp + src/simsoptpp/winding_surface.cpp ) set_target_properties(${PROJECT_NAME} diff --git a/examples/2_Intermediate/winding_surface.py b/examples/2_Intermediate/winding_surface.py new file mode 100755 index 000000000..d920a6f67 --- /dev/null +++ b/examples/2_Intermediate/winding_surface.py @@ -0,0 +1,439 @@ +#! /usr/bin/env python3 +from matplotlib import pyplot as plt +import numpy as np +from simsopt.objectives import SquaredFlux +from simsopt.field.magneticfieldclasses import WindingSurfaceField +from simsopt.field import CurrentPotentialFourier, CurrentPotentialSolve +from simsopt.geo import SurfaceRZFourier +from simsoptpp import WindingSurfaceBn_REGCOIL +from simsopt.util import in_github_actions +from pathlib import Path +import os +import time + +""" +In this example, we optimize a winding surface coil optimization to generate a specific target normal field on a user-provided plasma boundary. +We use the REGCOIL (Tikhonov regularization) and Lasso (L1 regularization) variants of the problem. +""" + +TEST_DIR = Path(__file__).parent / ".." / ".." / "tests" / "test_files" + + +def make_Bnormal_plots(cpst, OUT_DIR, filename): + """ + Plot Bnormal on the full torus plasma surface using the optimized + CurrentPotentialFourier object (cpst). + + Args: + cpst: CurrentPotentialSolve object + OUT_DIR: Output directory + filename: Filename for the output plots + + Returns: + None + """ + + # redefine the plasma surface to the full torus + s_plasma = cpst.plasma_surface + nfp = s_plasma.nfp + mpol = s_plasma.mpol + ntor = s_plasma.ntor + nphi = len(s_plasma.quadpoints_phi) + ntheta = len(s_plasma.quadpoints_theta) + quadpoints_phi = np.linspace(0, 1, nfp * nphi, endpoint=True) + quadpoints_theta = np.linspace(0, 1, ntheta, endpoint=True) + s_plasma_full = SurfaceRZFourier(nfp=nfp, mpol=mpol, ntor=ntor, stellsym=s_plasma.stellsym, quadpoints_phi=quadpoints_phi, quadpoints_theta=quadpoints_theta) + s_plasma_full.x = s_plasma.local_full_x + + # Compute the BiotSavart fields from the optimized current potential + Bfield_opt = WindingSurfaceField(cpst.current_potential) + Bfield_opt.set_points(s_plasma_full.gamma().reshape((-1, 3))) + Bn_coil = np.sum(Bfield_opt.B().reshape((nfp * nphi, ntheta, 3)) * s_plasma_full.unitnormal(), axis=2) + Bn_ext = - cpst.Bnormal_plasma.reshape(nphi, ntheta) + + # interpolate the known Bnormal_plasma onto new grid + if nfp > 1: + Bn_ext_full = np.vstack((Bn_ext, Bn_ext)) + for i in range(nfp - 2): + Bn_ext_full = np.vstack((Bn_ext_full, Bn_ext)) + else: + Bn_ext_full = Bn_ext + + # Save all the data to file + pointData = {"Bn": Bn_coil[:, :, None], "Bn_ext": Bn_ext_full[:, :, None], "Bn_total": (Bn_ext_full + Bn_coil)[:, :, None]} + s_plasma_full.to_vtk(OUT_DIR + filename, extra_data=pointData) + + +OUT_DIR = 'simsopt_winding_surface_example/' +os.makedirs(OUT_DIR, exist_ok=True) +files = ['regcoil_out.hsx.nc'] + + +def run_scan(): + """ + Run the REGCOIL and Lasso-regularized winding surface + problems across a wide range of regularization values, + and generate comparison plots. + """ + mpol = 8 + ntor = 8 + for file in files: + filename = TEST_DIR / file + + # Load current potential from NCSX configuration from REGCOIL + cpst = CurrentPotentialSolve.from_netcdf(filename) + cp = CurrentPotentialFourier.from_netcdf(filename) + + # Overwrite low-resolution NCSX config with higher-resolution + # current potential now with mpol = 16, ntor = 16. + cp = CurrentPotentialFourier( + cpst.winding_surface, mpol=mpol, ntor=ntor, + net_poloidal_current_amperes=cp.net_poloidal_current_amperes, + net_toroidal_current_amperes=cp.net_toroidal_current_amperes, + stellsym=True) + cpst = CurrentPotentialSolve(cp, cpst.plasma_surface, cpst.Bnormal_plasma) + + # define a number of geometric quantities from the plasma and coil surfaces + s_coil = cpst.winding_surface + s_plasma = cpst.plasma_surface + normal_coil = s_coil.normal().reshape(-1, 3) + normN = np.linalg.norm(normal_coil, axis=-1) + _nfp = s_plasma.nfp + nphi = len(s_plasma.quadpoints_phi) + ntheta = len(s_plasma.quadpoints_theta) + points = s_plasma.gamma().reshape(-1, 3) + normal = s_plasma.normal().reshape(-1, 3) + normN_plasma = np.linalg.norm(normal, axis=-1) + ws_points = s_coil.gamma().reshape(-1, 3) + dtheta_coil = s_coil.quadpoints_theta[1] + dzeta_coil = s_coil.quadpoints_phi[1] + + # function needed for saving to vtk after optimizing + contig = np.ascontiguousarray + + # Loop through wide range of regularization values + lambdas = np.logspace(-20, -10, 5) + fB_tikhonov = np.zeros(len(lambdas)) + fB_lasso = np.zeros(len(lambdas)) + fK_tikhonov = np.zeros(len(lambdas)) + fK_lasso = np.zeros(len(lambdas)) + fK_l1_lasso = np.zeros(len(lambdas)) + Kmax_tikhonov = np.zeros(len(lambdas)) + Kmean_tikhonov = np.zeros(len(lambdas)) + Kmax_lasso = np.zeros(len(lambdas)) + Kmean_lasso = np.zeros(len(lambdas)) + Bmax_tikhonov = np.zeros(len(lambdas)) + Bmean_tikhonov = np.zeros(len(lambdas)) + Bmax_lasso = np.zeros(len(lambdas)) + Bmean_lasso = np.zeros(len(lambdas)) + for i, lambda_reg in enumerate(lambdas): + print(i, lambda_reg) + + # Solve the REGCOIL problem that uses Tikhonov regularization (L2 norm) + optimized_phi_mn, f_B, _ = cpst.solve_tikhonov(lam=lambda_reg) + fB_tikhonov[i] = f_B + cp_opt = cpst.current_potential + K = cp_opt.K() + K2 = np.sum(K ** 2, axis=2) + f_K_direct = 0.5 * np.sum(np.ravel(K2) * normN) / (normal_coil.shape[0]) + fK_tikhonov[i] = f_K_direct + K = np.ascontiguousarray(K) + Kmax_tikhonov[i] = np.max(np.linalg.norm(K, axis=-1)) + Kmean_tikhonov[i] = np.mean(np.linalg.norm(K, axis=-1)) + + Bfield_opt = WindingSurfaceField(cp_opt) + Bfield_opt.set_points(s_plasma.gamma().reshape((-1, 3))) + + # For agreement at low regularization, + # Bnormal MUST be computed with the function below + # Since using Biot Savart will be discretized differently + Bnormal_REGCOIL = WindingSurfaceBn_REGCOIL( + points, + ws_points, + normal_coil, + cp_opt.Phi(), + normal + ) * dtheta_coil * dzeta_coil + Bnormal_REGCOIL += cpst.B_GI + Bn = Bnormal_REGCOIL + cpst.Bnormal_plasma + Bmax_tikhonov[i] = np.max(abs(Bn)) + Bmean_tikhonov[i] = np.mean(abs(Bn)) + res = (np.ravel(Bn) ** 2) @ normN_plasma + f_B_manual = 0.5 * res / (nphi * ntheta) + + # Check that fB calculations are consistent + print('f_B from least squares = ', f_B) + print('f_B direct = ', f_B_manual) + f_B_sf = SquaredFlux( + s_plasma, + Bfield_opt, + -contig(cpst.Bnormal_plasma.reshape(nphi, ntheta)) + ).J() + print('f_B from plasma surface = ', f_B_sf) + + # Repeat with the L1 instead of the L2 norm! + optimized_phi_mn, f_B, f_K, fB_history, fK_history = cpst.solve_lasso(lam=lambda_reg, max_iter=10000, acceleration=True) + + # Make plots of the history so we can see convergence was achieved + plt.figure(100) + lamstr = r'$\lambda$={0:.2e}'.format(lambda_reg) + color = f'C{i}' + plt.semilogy(fB_history, '-', color=color, label=lamstr + r': $f_B$') + plt.semilogy(lambda_reg * np.array(fK_history), '--', color=color, label=lamstr + r': $\lambda f_K$') + plt.grid(True) + + # repeat computing the metrics we defined + fB_lasso[i] = f_B + fK_l1_lasso[i] = f_K + K = cp_opt.K() + K2 = np.sum(K ** 2, axis=2) + f_K_direct = 0.5 * np.sum(np.ravel(K2) * normN) / (normal_coil.shape[0]) + fK_lasso[i] = f_K_direct + cp_opt = cpst.current_potential + K = np.ascontiguousarray(cp_opt.K()) + Kmax_lasso[i] = np.max(np.linalg.norm(K, axis=-1)) + Kmean_lasso[i] = np.mean(np.linalg.norm(K, axis=-1)) + Bfield_opt = WindingSurfaceField(cp_opt) + Bfield_opt.set_points(s_plasma.gamma().reshape((-1, 3))) + Bnormal_REGCOIL = WindingSurfaceBn_REGCOIL( + points, + ws_points, + normal_coil, + cp_opt.Phi(), + normal + ) * dtheta_coil * dzeta_coil + Bnormal_REGCOIL += cpst.B_GI + Bn = Bnormal_REGCOIL + cpst.Bnormal_plasma + Bmax_lasso[i] = np.max(abs(Bn)) + Bmean_lasso[i] = np.mean(abs(Bn)) + res = (np.ravel(Bn) ** 2) @ normN_plasma + f_B_manual = 0.5 * res / (nphi * ntheta) + + print('Results from Lasso: ') + print('f_B from least squares = ', f_B) + print('f_B direct = ', f_B_manual) + f_B_sf = SquaredFlux( + s_plasma, + Bfield_opt, + -contig(cpst.Bnormal_plasma.reshape(nphi, ntheta)) + ).J() + print('f_B from plasma surface = ', f_B_sf) + + # Finalize and save combined figure + plt.figure(100) + plt.title(r'Lasso convergence: $f_B$ and $\lambda f_K$') + plt.xlabel('Iteration') + plt.ylabel('Value') + plt.legend(fontsize=14, ncol=2) + plt.savefig(OUT_DIR + 'fB_fK_f_history.jpg') + + # plot cost function results + plt.figure(figsize=(14, 4)) + plt.subplot(1, 2, 1) + plt.suptitle(file) + plt.plot(lambdas, fB_tikhonov, 'b', label='f_B Tikhonov') + plt.plot(lambdas, fK_tikhonov / 1e14, 'r', label='f_K Tikhonov / 1e14') + plt.plot(lambdas, fK_tikhonov / 1e14 + fB_tikhonov, 'm', label='Total f Tikhonov') + plt.plot(lambdas, fB_lasso, 'b--', label='f_B Lasso') + plt.plot(lambdas, fK_lasso / 1e14, 'r--', label='f_K Lasso / 1e14') + plt.plot(lambdas, fK_lasso / 1e14 + fB_lasso, 'm--', label='Total f Lasso') + plt.xscale('log') + plt.yscale('log') + plt.grid(True) + plt.xlabel(r'$\lambda$') + plt.ylabel(r'$f_B$, $f_K$/1e14, Total $f$') + plt.legend() + plt.subplot(1, 2, 2) + plt.plot(lambdas, Kmean_tikhonov / 1e6, 'c', label='Kmean (MA) Tikhonov') + plt.plot(lambdas, Kmax_tikhonov / 1e6, 'k', label='Kmax (MA) Tikhonov') + plt.plot(lambdas, Kmean_lasso / 1e6, 'c--', label='Kmean (MA) Lasso') + plt.plot(lambdas, Kmax_lasso / 1e6, 'k--', label='Kmax (MA) Lasso') + plt.xscale('log') + plt.yscale('log') + plt.grid(True) + plt.xlabel(r'$\lambda$') + plt.ylabel(r'$K$ (MA)') + plt.legend() + plt.savefig(OUT_DIR + file + '_lambda_scan.jpg') + + plt.figure() + plt.suptitle(file) + plt.plot(fK_tikhonov, fB_tikhonov, 'r', label='L2') + plt.plot(fK_lasso, fB_lasso, 'b', label='L1 (same term as L2)') + # plt.plot(fK_l1_lasso, fB_lasso, 'm', label='L1 (using the L1 fK)') + plt.xlabel(r'$f_K$') + plt.ylabel(r'$f_B$') + plt.grid(True) + plt.legend() + plt.xscale('log') + plt.yscale('log') + plt.savefig(OUT_DIR + file + '_fK_fB.jpg') + + plt.figure() + plt.suptitle(file) + plt.plot(Kmax_tikhonov, Bmax_tikhonov, 'r', label='L2') + plt.plot(Kmax_lasso, Bmax_lasso, 'b', label='L1') + plt.ylabel(r'$max(|Bn|)$ on plasma surface') + plt.xlabel(r'$max(K)$ on coil surface') + plt.xscale('log') + plt.yscale('log') + plt.grid(True) + plt.legend() + plt.savefig(OUT_DIR + file + '_Kmax_Bmax.jpg') + + # Save results, useful if the run was intensive + np.savetxt( + OUT_DIR + file + '_metrics.txt', + np.array( + [lambdas, + fB_tikhonov, fK_tikhonov, + Bmax_tikhonov, Bmean_tikhonov, + Kmax_tikhonov, Kmean_tikhonov, + fB_lasso, fK_lasso, fK_l1_lasso, + Bmax_lasso, Bmean_lasso, + Kmax_lasso, Kmean_lasso, + ]).T + ) + + +def run_target(): + """ + Run REGCOIL (L2 regularization) and Lasso (L1 regularization) + starting from high regularization to low. When fB < fB_target + is achieved, the algorithms quit. This allows one to compare + L2 and L1 results at comparable levels of fB, which seems + like the fairest way to compare them. + """ + + fB_target = 1e-4 + mpol = 4 + ntor = 4 + coil_ntheta_res = 1 + coil_nzeta_res = coil_ntheta_res + plasma_ntheta_res = coil_ntheta_res + plasma_nzeta_res = coil_ntheta_res + + for file in files: + filename = TEST_DIR / file + + # Load in low-resolution NCSX file from REGCOIL + cpst = CurrentPotentialSolve.from_netcdf( + filename, plasma_ntheta_res, plasma_nzeta_res, coil_ntheta_res, coil_nzeta_res + ) + cp = CurrentPotentialFourier.from_netcdf(filename, coil_ntheta_res, coil_nzeta_res) + + # Overwrite low-resolution file with more mpol and ntor modes + cp = CurrentPotentialFourier( + cpst.winding_surface, mpol=mpol, ntor=ntor, + net_poloidal_current_amperes=cp.net_poloidal_current_amperes, + net_toroidal_current_amperes=cp.net_toroidal_current_amperes, + stellsym=True) + cpst = CurrentPotentialSolve(cp, cpst.plasma_surface, cpst.Bnormal_plasma) + s_coil = cpst.winding_surface + + nfp = s_coil.nfp + mpol = s_coil.mpol + ntor = s_coil.ntor + nphi = len(s_coil.quadpoints_phi) + ntheta = len(s_coil.quadpoints_theta) + quadpoints_phi = np.linspace(0, 1, nphi + 1, endpoint=True) + quadpoints_theta = np.linspace(0, 1, ntheta + 1, endpoint=True) + s_coil_full = SurfaceRZFourier(nfp=nfp, mpol=mpol, ntor=ntor, stellsym=s_coil.stellsym, quadpoints_phi=quadpoints_phi, quadpoints_theta=quadpoints_theta) + s_coil_full.x = s_coil.local_full_x + G = cp.net_poloidal_current_amperes + I = cp.net_toroidal_current_amperes + phi_secular, theta_secular = np.meshgrid(quadpoints_phi, quadpoints_theta, indexing='ij') + Phi_secular = G * phi_secular + I * theta_secular + # function needed for saving to vtk after optimizing + contig = np.ascontiguousarray + + # Loop through regularization values (avoid extreme 1e-22: ill-conditioned, SVD can fail on CI) + lambdas = np.flip(np.logspace(-20, -10, 2)) + for i, lambda_reg in enumerate(lambdas): + # Solve the REGCOIL problem that uses Tikhonov regularization (L2 norm) + optimized_phi_mn, f_B, _ = cpst.solve_tikhonov(lam=lambda_reg) + print(i, lambda_reg, f_B) + cp_opt = cpst.current_potential + + if f_B < fB_target: + K = cp_opt.K() + print('fB < fB_target has been achieved: ') + print('f_B from least squares = ', f_B) + print('lambda = ', lambda_reg) + make_Bnormal_plots( + cpst, + OUT_DIR, + file + "_tikhonov_fBtarget_Bnormal_lambda{0:.2e}".format(lambda_reg) + ) + Phi = cp_opt.Phi() + Phi = np.hstack((Phi, Phi[:, 0:1])) + Phi = np.vstack((Phi, Phi[0, :])) + Phi_secular + K = np.hstack((K, K[:, 0:1, :])) + K = np.vstack((K, K[0:1, :, :])) + pointData = {"phi": contig(Phi[:, :, None]), + "K": (contig(K[..., 0]), contig(K[..., 1]), contig(K[..., 2])) + } + s_coil_full.to_vtk( + OUT_DIR + file + "_tikhonov_fBtarget_winding_surface_lambda{0:.2e}".format(lambda_reg), + extra_data=pointData + ) + break + print('Now repeating for Lasso: ') + for i, lambda_reg in enumerate(lambdas): + # Solve the REGCOIL problem with the Lasso + optimized_phi_mn, f_B, _, fB_history, fK_history = cpst.solve_lasso(lam=lambda_reg, max_iter=5000, acceleration=True) + print(i, lambda_reg, f_B) + cp_opt = cpst.current_potential + + # Always plot Lasso convergence (even if target not achieved) so user sees the figure + plt.figure(100, figsize=(10, 6)) + lamstr = r'$\lambda$={0:.2e}'.format(lambda_reg) + color = f'C{i}' + plt.semilogy(fB_history, '-', color=color, linewidth=2, + label=lamstr + r': $f_B$ — squared flux residual') + plt.semilogy(lambda_reg * np.array(fK_history), '--', color=color, linewidth=2, + label=lamstr + r': $\lambda f_K$ — current density penalty') + plt.grid(True, alpha=0.5) + status = 'converged' if f_B < fB_target else 'did not converge' + plt.title(r'Winding surface Lasso optimization: $f_B$ and $\lambda f_K$ vs iteration ' + r'(target $f_B < {:.0e}$; {})'.format(fB_target, status), fontsize=12) + plt.xlabel('Iteration', fontsize=12) + plt.ylabel(r'Objective value (log scale)', fontsize=12) + plt.legend(fontsize=11, loc='upper right', framealpha=0.9) + plt.tight_layout() + plt.savefig(OUT_DIR + 'run_target_lasso_convergence.jpg', dpi=150, bbox_inches='tight') + + if f_B < fB_target: + K = contig(cp_opt.K()) + print('fB < fB_target has been achieved: ') + print('f_B from Lasso = ', f_B) + print('lambda = ', lambda_reg) + make_Bnormal_plots( + cpst, + OUT_DIR, + file + "_lasso_fBtarget_Bnormal_lambda{0:.2e}".format(lambda_reg) + ) + Phi = cp_opt.Phi() + Phi = np.hstack((Phi, Phi[:, 0:1])) + Phi = np.vstack((Phi, Phi[0, :])) + K = np.hstack((K, K[:, 0:1, :])) + K = np.vstack((K, K[0:1, :, :])) + pointData = {"phi": contig(Phi[:, :, None]), + "K": (contig(K[..., 0]), contig(K[..., 1]), contig(K[..., 2])) + } + s_coil_full.to_vtk( + OUT_DIR + file + "_lasso_fBtarget_winding_surface_lambda{0:.2e}".format(lambda_reg), + extra_data=pointData + ) + break + cpst.write_regcoil_out(filename='simsopt_' + file) + + +# Run one of the functions and time it +t1 = time.time() +# run_scan() +run_target() +t2 = time.time() +print('Total run time = ', t2 - t1) +if not in_github_actions: + plt.show() diff --git a/examples/3_Advanced/winding_surface_advanced.py b/examples/3_Advanced/winding_surface_advanced.py new file mode 100755 index 000000000..f29da55fd --- /dev/null +++ b/examples/3_Advanced/winding_surface_advanced.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 +""" +Cut Coils: Extract discrete coils from a current potential on a winding surface. + +This example demonstrates how to extract coil contours from a regularized current +potential (from REGCOIL or simsopt-regcoil) and convert them into simsopt Coil +objects. The workflow includes: + +1. Loading current potential data from a NetCDF file (simsopt-regcoil or legacy REGCOIL) +2. Computing the current potential φ and |∇φ| on a (θ, ζ) grid +3. Selecting contours by: + - Interactive double-click selection + - Specifying (θ, ζ) points that lie on desired contours + - Specifying contour level values + - Choosing N contours per field period +4. Classifying contours as window-pane (closed), modular, or helical +5. Computing currents for each contour +6. Converting contours to 3D curves and Coil objects +7. Applying stellarator symmetry to get the full coil set + +Usage +----- +From the simsopt root directory: + + python winding_surface_advanced.py --surface regcoil_out.hsx.nc + +With custom options: + + python winding_surface_advanced.py \\ + --surface /path/to/simsopt_regcoil_out.nc \\ + --ilambda 2 \\ + --points "0.5,1.0" "1.0,0.5" \\ + --output my_output_dir + +For interactive mode (double-click to select contours): + + python winding_surface_advanced.py --surface file.nc --interactive + +Requirements +------------ +- A simsopt-regcoil or REGCOIL NetCDF output file +- For legacy REGCOIL: a surface file (nescin format) for 3D mapping +- matplotlib for contour visualization +""" + +import argparse +from pathlib import Path + +# Add parent to path for standalone run +import sys +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from simsopt.util import run_cut_coils + +# Default test file +TEST_DIR = (Path(__file__).parent / ".." / ".." / "tests" / "test_files").resolve() +DEFAULT_REGCOIL = TEST_DIR / "regcoil_out.hsx.nc" + + +def main(): + parser = argparse.ArgumentParser( + description="Extract coils from a current potential (REGCOIL / simsopt-regcoil output).", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + parser.add_argument( + "--surface", + type=Path, + default=DEFAULT_REGCOIL, + help="REGCOIL or simsopt-regcoil NetCDF file (relative paths are resolved against tests/test_files).", + ) + parser.add_argument( + "--ilambda", + type=int, + default=-1, + help="Lambda index (0-based) to use.", + ) + parser.add_argument( + "--points", + nargs="*", + type=str, + default=None, + help='(θ,ζ) points as "theta,zeta" for contour selection.', + ) + parser.add_argument( + "--no-sv", + action="store_false", + help="Use full (multi-valued) current potential with net toroidal/poloidal currents.", + ) + parser.add_argument( + "--interactive", + action="store_true", + help="Use interactive double-click contour selection.", + ) + parser.add_argument( + "--output", + type=Path, + default=None, + help="Output directory for VTK and JSON (default: winding_surface_/).", + ) + parser.add_argument( + "--no-plot", + action="store_true", + help="Do not show final coil plot.", + ) + parser.add_argument( + "--save", + action="store_true", + help="Save coils to JSON file.", + ) + args = parser.parse_args() + + points = None + if args.points: + points = [] + for s in args.points: + parts = s.split(",") + if len(parts) != 2: + parser.error(f"Invalid point format: {s}. Use theta,zeta") + points.append([float(parts[0]), float(parts[1])]) + + surface_filename = args.surface if args.surface.is_absolute() else TEST_DIR / args.surface + output_path = args.output or Path(f"winding_surface_{surface_filename.stem}") + run_cut_coils( + surface_filename=surface_filename, + ilambda=args.ilambda, + points=points, + single_valued=not args.no_sv, + interactive=args.interactive, + show_final_coilset=not args.no_plot, + write_coils_to_file=args.save, + output_path=output_path, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/run_serial_examples b/examples/run_serial_examples index cb30cf97b..cac4550bd 100755 --- a/examples/run_serial_examples +++ b/examples/run_serial_examples @@ -20,6 +20,7 @@ set -ex ./2_Intermediate/stage_two_optimization_stochastic.py ./2_Intermediate/stage_two_optimization_finite_beta.py ./2_Intermediate/strain_optimization.py +./2_Intermediate/winding_surface.py ./2_Intermediate/permanent_magnet_MUSE.py ./2_Intermediate/permanent_magnet_QA.py ./2_Intermediate/permanent_magnet_PM4Stell.py @@ -30,3 +31,7 @@ set -ex ./3_Advanced/stage_two_optimization_finitebuild.py ./3_Advanced/coil_forces.py ./3_Advanced/wireframe_gsco_multistep.py +./3_Advanced/winding_surface_advanced.py --surface regcoil_out.hsx.nc +./3_Advanced/winding_surface_advanced.py --surface regcoil_out.w7x.nc --ilambda 0 +./3_Advanced/winding_surface_advanced.py --surface regcoil_out.near_axis.nc --points "0.5,0.5" "1.0,0.3" +./3_Advanced/winding_surface_advanced.py --surface regcoil_out.near_axis_asym.nc --no-sv --output winding_surface_near_axis_asym diff --git a/src/simsopt/field/__init__.py b/src/simsopt/field/__init__.py index 67a6f4113..a77153060 100644 --- a/src/simsopt/field/__init__.py +++ b/src/simsopt/field/__init__.py @@ -7,6 +7,7 @@ from .mgrid import * from .normal_field import * from .tracing import * +from .currentpotential import * from .wireframefield import * from .selffield import * from .magnetic_axis_helpers import * @@ -16,6 +17,7 @@ biotsavart.__all__ + boozermagneticfield.__all__ + coil.__all__ + + currentpotential.__all__ + coilset.__all__ + magneticfield.__all__ + magneticfieldclasses.__all__ diff --git a/src/simsopt/field/coil.py b/src/simsopt/field/coil.py index f334d5d5e..54d36762b 100644 --- a/src/simsopt/field/coil.py +++ b/src/simsopt/field/coil.py @@ -559,6 +559,7 @@ def coils_to_vtk(coils, filename, close=False, extra_data=None): ppl = np.asarray([c.gamma().shape[0]+1 for c in curves]) else: ppl = np.asarray([c.gamma().shape[0] for c in curves]) + ppl_cumsum = np.concatenate([[0], np.cumsum(ppl)]) # get the current data, which is the same at every point on a given coil contig = np.ascontiguousarray @@ -566,7 +567,7 @@ def coils_to_vtk(coils, filename, close=False, extra_data=None): data = np.concatenate([i*np.ones((ppl[i], )) for i in range(len(curves))]) coil_data = np.zeros(data.shape) for i in range(len(currents)): - coil_data[i * ppl[i]: (i + 1) * ppl[i]] = currents[i] + coil_data[ppl_cumsum[i]:ppl_cumsum[i+1]] = currents[i] coil_data = np.ascontiguousarray(coil_data) pointData['I'] = coil_data pointData['I_mag'] = contig(np.abs(coil_data)) @@ -593,20 +594,20 @@ def coils_to_vtk(coils, filename, close=False, extra_data=None): if close: coil_force_temp = np.vstack((coil_force_temp, coil_force_temp[0, :])) coil_torque_temp = np.vstack((coil_torque_temp, coil_torque_temp[0, :])) - coil_forces[i * ppl[i]: (i + 1) * ppl[i], :] = coil_force_temp - coil_torques[i * ppl[i]: (i + 1) * ppl[i], :] = coil_torque_temp + coil_forces[ppl_cumsum[i]:ppl_cumsum[i+1], :] = coil_force_temp + coil_torques[ppl_cumsum[i]:ppl_cumsum[i+1], :] = coil_torque_temp # copy force and torque data over to pointwise data on a coil curve coil_data = np.zeros((data.shape[0], 3)) for i in range(len(coils)): - coil_data[i * ppl[i]: (i + 1) * ppl[i], :] = net_forces[i, :] + coil_data[ppl_cumsum[i]:ppl_cumsum[i+1], :] = net_forces[i, :] coil_data = np.ascontiguousarray(coil_data) pointData['NetForces'] = (contig(coil_data[:, 0]), contig(coil_data[:, 1]), contig(coil_data[:, 2])) coil_data = np.zeros((data.shape[0], 3)) for i in range(len(coils)): - coil_data[i * ppl[i]: (i + 1) * ppl[i], :] = net_torques[i, :] + coil_data[ppl_cumsum[i]:ppl_cumsum[i+1], :] = net_torques[i, :] coil_data = np.ascontiguousarray(coil_data) # Add pointwise force and torque data to the dictionary diff --git a/src/simsopt/field/currentpotential.py b/src/simsopt/field/currentpotential.py new file mode 100644 index 000000000..21b8fe0df --- /dev/null +++ b/src/simsopt/field/currentpotential.py @@ -0,0 +1,1298 @@ +from __future__ import annotations + +from typing import Optional, Union, List, Tuple +import numpy as np +import warnings +from scipy.io import netcdf_file +from scipy.interpolate import RegularGridInterpolator +from .._core.optimizable import DOFs, Optimizable +from .._core.json import GSONDecoder +import simsoptpp as sopp +from simsopt.geo import SurfaceRZFourier +from .magneticfieldclasses import WindingSurfaceField + +__all__ = ['CurrentPotentialFourier', 'CurrentPotential', 'CurrentPotentialSolve'] + + +class CurrentPotential(Optimizable): + """ + Current Potential base object, not necessarily assuming + that the current potential will be represented by a + Fourier expansion in the toroidal and poloidal modes. + + Args: + winding_surface: SurfaceRZFourier object representing the coil surface. + kwargs: Additional keyword arguments passed to the Optimizable base class. + """ + + def __init__(self, winding_surface: SurfaceRZFourier, **kwargs) -> None: + super().__init__(**kwargs) + self.winding_surface = winding_surface + + def K(self) -> np.ndarray: + """ + Get the K vector of the CurrentPotential object. + + Returns: + np.ndarray: The K vector of the CurrentPotential object. + """ + data = np.zeros((len(self.quadpoints_phi), len(self.quadpoints_theta), 3)) + dg1 = self.winding_surface.gammadash1() + dg2 = self.winding_surface.gammadash2() + normal = self.winding_surface.normal() + self.K_impl_helper(data, dg1, dg2, normal) + return data + + def K_matrix(self) -> np.ndarray: + """ + Get the K matrix of the CurrentPotential object. + + Returns: + np.ndarray: The K matrix of the CurrentPotential object. + """ + data = np.zeros((self.num_dofs(), self.num_dofs())) + dg1 = self.winding_surface.gammadash1() + dg2 = self.winding_surface.gammadash2() + normal = self.winding_surface.normal() + self.K_matrix_impl_helper(data, dg1, dg2, normal) + return data + + def num_dofs(self) -> int: + """ + Get the number of dofs of the CurrentPotential object. + + Returns: + int: The number of dofs of the CurrentPotential object. + """ + return len(self.get_dofs()) + + +class CurrentPotentialFourier(sopp.CurrentPotentialFourier, CurrentPotential): + """ + Current Potential Fourier object is designed for initializing + the winding surface problem, assuming that the current potential + will be represented by a Fourier expansion in the toroidal + and poloidal modes. + + Args: + winding_surface: SurfaceRZFourier object representing the coil surface. + net_poloidal_current_amperes: Net poloidal current in amperes, needed + to compute the B_GI contributions to the Bnormal part of the optimization. + net_toroidal_current_amperes: Net toroidal current in amperes, needed + to compute the B_GI contributions to the Bnormal part of the optimization. + nfp: The number of field periods. + stellsym: Whether the surface is stellarator-symmetric, i.e. + symmetry under rotation by :math:`\pi` about the x-axis. + mpol: Maximum poloidal mode number included. + ntor: Maximum toroidal mode number included, divided by ``nfp``. + quadpoints_phi: Set this to a list or 1D array to set the :math:`\phi_j` grid points directly. + quadpoints_theta: Set this to a list or 1D array to set the :math:`\theta_j` grid points directly. + """ + + def __init__( + self, + winding_surface: SurfaceRZFourier, + net_poloidal_current_amperes: float = 1, + net_toroidal_current_amperes: float = 0, + nfp: Optional[int] = None, + stellsym: Optional[bool] = None, + mpol: Optional[int] = None, + ntor: Optional[int] = None, + quadpoints_phi: Optional[Union[np.ndarray, List[float]]] = None, + quadpoints_theta: Optional[Union[np.ndarray, List[float]]] = None, + ) -> None: + + if nfp is None: + nfp = winding_surface.nfp + if stellsym is None: + stellsym = winding_surface.stellsym + if mpol is None: + mpol = winding_surface.mpol + if ntor is None: + ntor = winding_surface.ntor + + if nfp > 1 and np.max(winding_surface.quadpoints_phi) <= 1/nfp: + raise AttributeError('winding_surface must contain all field periods.') + + # quadpoints_phi or quadpoints_theta has to be the same as those + # in the winding surface. Otherwise, it can cause issues during + # CurrentPotential.K(), CurrentPotential.K_matrix and + # in WindingSurfaceField.__init__(). + if quadpoints_theta is None: + quadpoints_theta = winding_surface.quadpoints_theta + if quadpoints_phi is None: + quadpoints_phi = winding_surface.quadpoints_phi + + sopp.CurrentPotentialFourier.__init__(self, mpol, ntor, nfp, stellsym, + quadpoints_phi, quadpoints_theta, + net_poloidal_current_amperes, + net_toroidal_current_amperes) + + CurrentPotential.__init__(self, winding_surface, x0=self.get_dofs(), + external_dof_setter=CurrentPotentialFourier.set_dofs_impl, + names=self._make_names()) + + self._make_mn() + phi_secular, theta_secular = np.meshgrid( + self.winding_surface.quadpoints_phi, + self.winding_surface.quadpoints_theta, + indexing='ij' + ) + self.current_potential_secular = ( + phi_secular * net_poloidal_current_amperes + theta_secular * net_toroidal_current_amperes + ) + + def _make_names(self) -> List[str]: + """ + Create the dof names for the CurrentPotentialFourier object. + + Returns: + List[str]: The names of the coefficients. + """ + if self.stellsym: + names = self._make_names_helper('Phis') + else: + names = self._make_names_helper('Phis') \ + + self._make_names_helper('Phic') + return names + + def _make_names_helper(self, prefix: str) -> List[str]: + """ + Helper function for _make_names() method to format the strings. + + Args: + prefix: The prefix for the name of the coefficients. + + Returns: + List[str]: The names of the coefficients. + """ + names = [] + + start = 1 + names += [prefix + '(0,' + str(n) + ')' for n in range(start, self.ntor + 1)] + for m in range(1, self.mpol + 1): + names += [prefix + '(' + str(m) + ',' + str(n) + ')' for n in range(-self.ntor, self.ntor + 1)] + return names + + def get_dofs(self) -> np.ndarray: + """ + Get the dofs of the CurrentPotentialFourier object. + + Returns: + np.ndarray: The dofs of the CurrentPotentialFourier object. + """ + return np.asarray(sopp.CurrentPotentialFourier.get_dofs(self)) + + def change_resolution(self, mpol: int, ntor: int) -> None: + """ + Modeled after SurfaceRZFourier + Change the values of `mpol` and `ntor`. Any new Fourier amplitudes + will have a magnitude of zero. Any previous nonzero Fourier + amplitudes that are not within the new range will be + discarded. + + Args: + mpol: New poloidal mode number. + ntor: New toroidal mode number. + """ + old_mpol = self.mpol + old_ntor = self.ntor + old_phis = self.phis + if not self.stellsym: + old_phic = self.phic + self.mpol = mpol + self.ntor = ntor + self.allocate() + if mpol < old_mpol or ntor < old_ntor: + self.invalidate_cache() + + min_mpol = np.min((mpol, old_mpol)) + min_ntor = np.min((ntor, old_ntor)) + for m in range(min_mpol + 1): + for n in range(-min_ntor, min_ntor + 1): + self.phis[m, n + ntor] = old_phis[m, n + old_ntor] + if not self.stellsym: + self.phic[m, n + ntor] = old_phic[m, n + old_ntor] + + # Update the dofs object + self._dofs = DOFs(self.get_dofs(), self._make_names()) + + # The following methods of graph Optimizable framework need to be called + Optimizable.update_free_dof_size_indices(self) + Optimizable._update_full_dof_size_indices(self) + Optimizable.set_recompute_flag(self) + + def get_phic(self, m: int, n: int) -> float: + """ + Return a particular `phic` parameter. + + Args: + m: Poloidal mode number. + n: Toroidal mode number. + + Raises: + ValueError: If `stellsym` is True (phic does not exist for this stellarator-symmetric current potential). + IndexError: If `m` is less than 0, `m` is greater than `mpol`, `n` is greater than `ntor`, or `n` is less than -`ntor`. + """ + if self.stellsym: + raise ValueError( + 'phic does not exist for this stellarator-symmetric current potential.') + self._validate_mn(m, n) + return self.phic[m, n + self.ntor] + + def get_phis(self, m: int, n: int) -> float: + """ + Return a particular `phis` parameter. + + Args: + m: Poloidal mode number. + n: Toroidal mode number. + + Raises: + IndexError: If `m` is less than 0, `m` is greater than `mpol`, `n` is greater than `ntor`, or `n` is less than -`ntor`. + """ + self._validate_mn(m, n) + return self.phis[m, n + self.ntor] + + def set_phic(self, m: int, n: int, val: float) -> None: + """ + Set a particular `phic` Parameter. + + Args: + m: Poloidal mode number. + n: Toroidal mode number. + val: Value to set the `phic` parameter to. + + Raises: + ValueError: If `stellsym` is True (phic does not exist for this stellarator-symmetric current potential). + IndexError: If `m` is less than 0, `m` is greater than `mpol`, `n` is greater than `ntor`, or `n` is less than -`ntor`. + """ + if self.stellsym: + raise ValueError( + 'phic does not exist for this stellarator-symmetric current potential.') + self._validate_mn(m, n) + self.phic[m, n + self.ntor] = val + self.local_full_x = self.get_dofs() + self.invalidate_cache() + + def set_phis(self, m: int, n: int, val: float) -> None: + """ + Set a particular `phis` Parameter. + + Args: + m: Poloidal mode number. + n: Toroidal mode number. + val: Value to set the `phis` parameter to. + + Raises: + IndexError: If `m` is less than 0, `m` is greater than `mpol`, `n` is greater than `ntor`, or `n` is less than -`ntor`. + """ + self._validate_mn(m, n) + self.phis[m, n + self.ntor] = val + self.local_full_x = self.get_dofs() + self.invalidate_cache() + + def set_net_toroidal_current_amperes(self, val: float) -> None: + """ + Set the net toroidal current in Amperes. + + Args: + val: Value to set the net toroidal current to. + """ + self.net_toroidal_current_amperes = val + self.invalidate_cache() + + def set_net_poloidal_current_amperes(self, val: float) -> None: + """ + Set the net poloidal current in Amperes. + + Args: + val: Value to set the net poloidal current to. + """ + self.net_poloidal_current_amperes = val + self.invalidate_cache() + + def fixed_range( + self, mmin: int, mmax: int, nmin: int, nmax: int, fixed: bool = True + ) -> None: + """ + Modeled after SurfaceRZFourier + Set the 'fixed' property for a range of `m` and `n` values. + + All modes with `m` in the interval [`mmin`, `mmax`] and `n` in the + interval [`nmin`, `nmax`] will have their fixed property set to + the value of the `fixed` parameter. Note that `mmax` and `nmax` + are included (unlike the upper bound in python's range(min, + max).) + + Args: + mmin: Minimum poloidal mode number. + mmax: Maximum poloidal mode number. + nmin: Minimum toroidal mode number. + nmax: Maximum toroidal mode number. + fixed: Whether to fix the modes. + + Raises: + IndexError: If `mmin` is less than 0, `mmax` is greater than `mpol`, `nmin` is less than -`ntor`, or `nmax` is greater than `ntor`. + """ + # TODO: This will be slow because free dof indices are evaluated all + # TODO: the time in the loop + fn = self.fix if fixed else self.unfix + for m in range(mmin, mmax + 1): + this_nmin = nmin + if m == 0 and nmin < 0: + this_nmin = 1 + for n in range(this_nmin, nmax + 1): + if m > 0 or n != 0: + fn(f'Phis({m},{n})') + if not self.stellsym: + fn(f'Phic({m},{n})') + + def _validate_mn(self, m: int, n: int) -> None: + """ + Copied from SurfaceRZFourier + Check whether `m` and `n` are in the allowed range. + + Args: + m: Poloidal mode number. + n: Toroidal mode number. + + Raises: + IndexError: If `m` is less than 0, `m` is greater than `mpol`, `n` is greater than `ntor`, or `n` is less than -`ntor`. + """ + if m < 0: + raise IndexError('m must be >= 0') + if m > self.mpol: + raise IndexError('m must be <= mpol') + if n > self.ntor: + raise IndexError('n must be <= ntor') + if n < -self.ntor: + raise IndexError('n must be >= -ntor') + + def _make_mn(self) -> None: + """ + Make the list of m and n values. + """ + m1d = np.arange(self.mpol + 1) + n1d = np.arange(-self.ntor, self.ntor + 1) + n2d, m2d = np.meshgrid(n1d, m1d) + m0 = m2d.flatten()[self.ntor:] + n0 = n2d.flatten()[self.ntor:] + self.m = m0[1::] + self.n = n0[1::] + + if not self.stellsym: + self.m = np.append(self.m, self.m) + self.n = np.append(self.n, self.n) + + def set_current_potential_from_regcoil(self, filename: str, ilambda: int): + """ + Set phic and phis based on a regcoil netcdf file. + + Args: + filename: Name of the ``regcoil_out.*.nc`` file to read. + ilambda: 0-based index for the lambda array, indicating which current + potential solution to use + """ + f = netcdf_file(filename, 'r', mmap=False) + nfp = f.variables['nfp'][()] + mpol_potential = f.variables['mpol_potential'][()] + ntor_potential = f.variables['ntor_potential'][()] + _xm_potential = f.variables['xm_potential'][()] + _xn_potential = f.variables['xn_potential'][()] + symmetry_option = f.variables['symmetry_option'][()] + single_valued_current_potential_mn = f.variables['single_valued_current_potential_mn'][()][ilambda, :] + f.close() + + # Check that correct shape of arrays are being used + if mpol_potential != self.mpol: + raise ValueError('Incorrect mpol_potential') + if ntor_potential != self.ntor: + raise ValueError('Incorrect ntor_potential') + if nfp != self.nfp: + raise ValueError('Incorrect nfp') + if symmetry_option == 1: + stellsym = True + else: + stellsym = False + if stellsym != self.stellsym: + raise ValueError('Incorrect stellsym') + + self.set_dofs(single_valued_current_potential_mn) + + def as_dict(self, serial_objs_dict=None) -> dict: + """Sync Python _dofs with C++ state before serialization (set_dofs is C++-only).""" + if len(self.local_full_x): + self.local_full_x = self.get_dofs() + return super().as_dict(serial_objs_dict) + + @classmethod + def from_netcdf( + cls, + filename: str, + coil_ntheta_res: float = 1.0, + coil_nzeta_res: float = 1.0, + ) -> "CurrentPotentialFourier": + """ + Initialize a CurrentPotentialFourier object from a regcoil netcdf output file. + + Args: + filename: Name of the ``regcoil_out.*.nc`` file to read. + coil_ntheta_res: The resolution of the coil surface in the theta direction. + coil_nzeta_res: The resolution of the coil surface in the zeta direction. + + Returns: + cls: The CurrentPotentialFourier object. + """ + f = netcdf_file(filename, 'r', mmap=False) + nfp = f.variables['nfp'][()] + mpol_potential = f.variables['mpol_potential'][()] + ntor_potential = f.variables['ntor_potential'][()] + net_poloidal_current_amperes = f.variables['net_poloidal_current_Amperes'][()] + net_toroidal_current_amperes = f.variables['net_toroidal_current_Amperes'][()] + _xm_potential = f.variables['xm_potential'][()] + _xn_potential = f.variables['xn_potential'][()] + symmetry_option = f.variables['symmetry_option'][()] + if symmetry_option == 1: + stellsym = True + else: + stellsym = False + rmnc_coil = f.variables['rmnc_coil'][()] + zmns_coil = f.variables['zmns_coil'][()] + if ('rmns_coil' in f.variables and 'zmnc_coil' in f.variables): + rmns_coil = f.variables['rmns_coil'][()] + zmnc_coil = f.variables['zmnc_coil'][()] + if np.all(zmnc_coil == 0) and np.all(rmns_coil == 0): + stellsym_surf = True + else: + stellsym_surf = False + else: + rmns_coil = np.zeros_like(rmnc_coil) + zmnc_coil = np.zeros_like(zmns_coil) + stellsym_surf = True + xm_coil = f.variables['xm_coil'][()] + xn_coil = f.variables['xn_coil'][()] + ntheta_coil = int(f.variables['ntheta_coil'][()] * coil_ntheta_res) + nzeta_coil = int(f.variables['nzeta_coil'][()] * coil_nzeta_res) + f.close() + mpol_coil = int(np.max(xm_coil)) + ntor_coil = int(np.max(xn_coil)/nfp) + s_coil = SurfaceRZFourier(nfp=nfp, mpol=mpol_coil, ntor=ntor_coil, stellsym=stellsym_surf) + s_coil = s_coil.from_nphi_ntheta(nfp=nfp, ntheta=ntheta_coil, nphi=nzeta_coil * nfp, + mpol=mpol_coil, ntor=ntor_coil, stellsym=stellsym_surf, range='full torus') + + s_coil.set_dofs(0*s_coil.get_dofs()) + + for im in range(len(xm_coil)): + s_coil.set_rc(xm_coil[im], int(xn_coil[im]/nfp), rmnc_coil[im]) + s_coil.set_zs(xm_coil[im], int(xn_coil[im]/nfp), zmns_coil[im]) + _m = int(xm_coil[im]) + _n = int(xn_coil[im] / nfp) + + if not stellsym_surf: + s_coil.set_rs(xm_coil[im], int(xn_coil[im]/nfp), rmns_coil[im]) + s_coil.set_zc(xm_coil[im], int(xn_coil[im]/nfp), zmnc_coil[im]) + + + s_coil.local_full_x = s_coil.get_dofs() + + cp = cls(s_coil, mpol=mpol_potential, ntor=ntor_potential, + net_poloidal_current_amperes=net_poloidal_current_amperes, + net_toroidal_current_amperes=net_toroidal_current_amperes, + stellsym=stellsym) + + return cp + + @classmethod + def from_dict(cls, d, serial_objs_dict, recon_objs): + """ + Reconstruct CurrentPotentialFourier from serialized dict. + Excludes dofs from __init__ and sets them after construction. + """ + decoder = GSONDecoder() + init_kwargs = dict(d) + dofs_raw = init_kwargs.pop("dofs", None) + decoded = {k: decoder.process_decoded(v, serial_objs_dict, recon_objs) + for k, v in init_kwargs.items()} + obj = cls(**decoded) + if dofs_raw is not None: + dofs = decoder.process_decoded(dofs_raw, serial_objs_dict, recon_objs) + obj.set_dofs(np.asarray(dofs.full_x)) + return obj + + +class CurrentPotentialSolve: + """ + Current Potential Solve object is designed for performing + the winding surface coil optimization. We provide functionality + for the REGCOIL (tikhonov regularization) and L1 norm (Lasso) + variants of the problem. + + Args: + cp: CurrentPotential class object containing the winding surface. + plasma_surface: The plasma surface to optimize Bnormal over. + Bnormal_plasma: Bnormal coming from plasma currents. + B_GI: Bnormal coming from the net coil currents. + """ + + def __init__( + self, + cp: CurrentPotentialFourier, + plasma_surface: SurfaceRZFourier, + Bnormal_plasma: Union[np.ndarray, float], + ) -> None: + + if np.max(plasma_surface.quadpoints_phi) >= 1/plasma_surface.nfp: + raise AttributeError('winding_surface must contain only one field period.') + + self.current_potential = cp + self.winding_surface = self.current_potential.winding_surface + self.ndofs = self.current_potential.num_dofs() + self.plasma_surface = plasma_surface + self.ntheta_plasma = len(self.plasma_surface.quadpoints_theta) + self.nzeta_plasma = len(self.plasma_surface.quadpoints_phi) + self.ntheta_coil = len(self.current_potential.quadpoints_theta) + self.nzeta_coil = len(self.current_potential.quadpoints_phi) + # Calculating B_GI + cp_no_phi_sv = CurrentPotentialFourier( + cp.winding_surface, mpol=cp.mpol, ntor=cp.ntor, + net_poloidal_current_amperes=cp.net_poloidal_current_amperes, + net_toroidal_current_amperes=cp.net_toroidal_current_amperes, + quadpoints_phi=cp.quadpoints_phi, + quadpoints_theta=cp.quadpoints_theta, + stellsym=cp.stellsym + ) + Bfield = WindingSurfaceField(cp_no_phi_sv) + points = plasma_surface.gamma().reshape(-1, 3) + Bfield.set_points(points) + B_GI_vector = Bfield.B() + normal = plasma_surface.unitnormal().reshape(-1, 3) + B_GI_winding_surface = np.sum(B_GI_vector*normal, axis=1) + # Permitting Bnormal_plasma to be a scalar + if np.isscalar(Bnormal_plasma): + Bnormal_plasma = Bnormal_plasma * np.ones(normal.shape[0]) + else: + # If Bnormal is not a scalar, try reshaping it into + # the proper shape + try: + Bnormal_plasma = Bnormal_plasma.reshape(normal.shape[0]) + except Exception: + raise ValueError('The shape of Bnormal_plasma does not match with the quadrature points of plasma_surface.') + self.Bnormal_plasma = Bnormal_plasma + self.B_GI = B_GI_winding_surface + # Save list of results for each L2 or L1 winding surface + # optimization performed with this class object + self.ilambdas_l2 = [] + self.dofs_l2 = [] + self.current_potential_l2 = [] + self.K2s_l2 = [] + self.fBs_l2 = [] + self.fKs_l2 = [] + self.ilambdas_l1 = [] + self.dofs_l1 = [] + self.current_potential_l1 = [] + self.K2s_l1 = [] + self.fBs_l1 = [] + self.fKs_l1 = [] + # Reference xm/xn_potential from source file (set by from_netcdf) for REGCOIL write + self._ref_xm_potential = None + self._ref_xn_potential = None + warnings.warn( + "Beware: the f_B (also called chi^2_B) computed from the " + "CurrentPotentialSolve class will be slightly different than " + "the f_B computed using SquaredFlux with the BiotSavart law " + "implemented in WindingSurfaceField. This is because the " + "optimization formulation and the full BiotSavart calculation " + "are discretized in different ways. This disagreement will " + "worsen at low regularization, but improve with higher " + "resolution on the plasma and coil surfaces. " + ) + + @classmethod + def from_netcdf( + cls, + filename: str, + plasma_ntheta_res: float = 1.0, + plasma_nzeta_res: float = 1.0, + coil_ntheta_res: float = 1.0, + coil_nzeta_res: float = 1.0, + ) -> "CurrentPotentialSolve": + """ + Initialize a CurrentPotentialSolve using a CurrentPotentialFourier + from a regcoil netcdf output file. The single_valued_current_potential_mn + are set to zero. + + Args: + filename: Name of the ``regcoil_out.*.nc`` file to read. + plasma_ntheta_res: The resolution of the plasma surface in the theta direction. + plasma_nzeta_res: The resolution of the plasma surface in the zeta direction. + coil_ntheta_res: The resolution of the coil surface in the theta direction. + coil_nzeta_res: The resolution of the coil surface in the zeta direction. + + Returns: + cls: The CurrentPotentialSolve object. + + Notes: + This function initializes a CurrentPotentialSolve object from a regcoil netcdf output file. + The CurrentPotentialFourier object is initialized from the file, and the Bnormal_from_plasma_current + is read from the file. The plasma surface is initialized from the file, and the coil surface is + initialized from the file. The CurrentPotentialSolve object is returned. + """ + f = netcdf_file(filename, 'r', mmap=False) + nfp = f.variables['nfp'][()] + Bnormal_from_plasma_current = f.variables['Bnormal_from_plasma_current'][()] + rmnc_plasma = f.variables['rmnc_plasma'][()] + zmns_plasma = f.variables['zmns_plasma'][()] + xm_plasma = f.variables['xm_plasma'][()] + xn_plasma = f.variables['xn_plasma'][()] + mpol_plasma = int(np.max(xm_plasma)) + ntor_plasma = int(np.max(xn_plasma)/nfp) + ntheta_plasma = int(f.variables['ntheta_plasma'][()] * plasma_ntheta_res) + nzeta_plasma = int(f.variables['nzeta_plasma'][()] * plasma_nzeta_res) + if ('rmns_plasma' in f.variables and 'zmnc_plasma' in f.variables): + rmns_plasma = f.variables['rmns_plasma'][()] + zmnc_plasma = f.variables['zmnc_plasma'][()] + if np.all(zmnc_plasma == 0) and np.all(rmns_plasma == 0): + stellsym_plasma_surf = True + else: + stellsym_plasma_surf = False + else: + rmns_plasma = np.zeros_like(rmnc_plasma) + zmnc_plasma = np.zeros_like(zmns_plasma) + stellsym_plasma_surf = True + + cp = CurrentPotentialFourier.from_netcdf(filename, coil_ntheta_res, coil_nzeta_res) + ref_xm_potential = f.variables['xm_potential'][()] + ref_xn_potential = f.variables['xn_potential'][()] + + s_plasma = SurfaceRZFourier( + nfp=nfp, + mpol=mpol_plasma, + ntor=ntor_plasma, + stellsym=stellsym_plasma_surf + ) + s_plasma = s_plasma.from_nphi_ntheta( + nfp=nfp, ntheta=ntheta_plasma, + nphi=nzeta_plasma, + mpol=mpol_plasma, ntor=ntor_plasma, + stellsym=stellsym_plasma_surf, range="field period" + ) + + # Need to interpolate Bnormal_from_plasma if increasing resolution + if plasma_ntheta_res > 1.0 or plasma_nzeta_res > 1.0: + warnings.warn( + "User specified to increase the plasma surface resolution, but is " + "reading the CurrentPotential object from a netcdf file. Therefore, " + "the Bnormal_from_plasma_current will be interpolated to the " + "high resolution grid, which may not be very accurate!" + ) + # Source grid: REGCOIL file convention (matches Bnormal_from_plasma_current layout) + quadpoints_phi_1d = np.linspace( + 0, 1 / ((int(stellsym_plasma_surf) + 1) * nfp), + f.variables['nzeta_plasma'][()] + 1, endpoint=True + )[:-1] + quadpoints_theta_1d = np.linspace( + 0, 1, f.variables['ntheta_plasma'][()] + 1, endpoint=True + )[:-1] + Bnormal_interp = RegularGridInterpolator( + (quadpoints_phi_1d, quadpoints_theta_1d), + Bnormal_from_plasma_current, + method='cubic', + bounds_error=False, + fill_value=None + ) + phi_grid, theta_grid = np.meshgrid( + s_plasma.quadpoints_phi, s_plasma.quadpoints_theta, indexing='ij' + ) + Bnormal_from_plasma_current = Bnormal_interp( + np.column_stack([phi_grid.ravel(), theta_grid.ravel()]) + ).reshape(phi_grid.shape) + f.close() + s_plasma.set_dofs(0 * s_plasma.get_dofs()) + for im in range(len(xm_plasma)): + s_plasma.set_rc(xm_plasma[im], int(xn_plasma[im] / nfp), rmnc_plasma[im]) + s_plasma.set_zs(xm_plasma[im], int(xn_plasma[im] / nfp), zmns_plasma[im]) + if not stellsym_plasma_surf: + s_plasma.set_rs(xm_plasma[im], int(xn_plasma[im] / nfp), rmns_plasma[im]) + s_plasma.set_zc(xm_plasma[im], int(xn_plasma[im] / nfp), zmnc_plasma[im]) + cps = cls(cp, s_plasma, np.ravel(Bnormal_from_plasma_current)) + cps._ref_xm_potential = ref_xm_potential + cps._ref_xn_potential = ref_xn_potential + return cps + + def write_regcoil_out(self, filename: str) -> None: + """ + Take optimized CurrentPotentialSolve class and save it to a regcoil-style + outfile for backwards compatability with other stellarator codes. + + Args: + filename: Name of the ``regcoil_out.*.nc`` file to read. + """ + f = netcdf_file(filename, 'w') + f.history = 'Created for writing a SIMSOPT-optimized winding surface and current potential to a regcoil-style output file' + + scalars = ['nfp', 'mpol_plasma', 'ntor_plasma', 'mpol_potential', 'ntor_potential', 'symmetry_option', 'net_poloidal_current_Amperes', 'net_toroidal_current_Amperes', 'ntheta_plasma', 'nzeta_plasma', 'ntheta_coil', 'nzeta_coil'] + s = self.plasma_surface + w = self.winding_surface + G = self.current_potential.net_poloidal_current_amperes + I = self.current_potential.net_toroidal_current_amperes + scalar_variables = [s.nfp, s.mpol, s.ntor, w.mpol, w.ntor, s.stellsym + 1, G, I, self.ntheta_plasma, self.nzeta_plasma, self.ntheta_coil, self.nzeta_coil / s.nfp] + for i_scalar, scalar_name in enumerate(scalars): + f.createDimension(scalar_name, 1) + if 'amperes' not in scalar_name: + var = f.createVariable(scalar_name, 'i', (scalar_name,)) + var.units = 'dimensionless' + else: + var = f.createVariable(scalar_name, 'f', (scalar_name,)) + var.units = 'Amperes' + var[:] = scalar_variables[i_scalar] + + # go through and compute all the rmnc, rmns for the plasma surface + nfp = s.nfp + xn_plasma = s.n * nfp + xm_plasma = s.m + rmnc = np.zeros(len(xm_plasma)) + zmns = np.zeros(len(xm_plasma)) + zmnc = np.zeros(len(xm_plasma)) + rmns = np.zeros(len(xm_plasma)) + for im in range(len(xm_plasma)): + rmnc[im] = s.get_rc(xm_plasma[im], int(xn_plasma[im]/nfp)) + zmns[im] = s.get_zs(xm_plasma[im], int(xn_plasma[im]/nfp)) + if not s.stellsym: + rmns[im] = s.get_rs(xm_plasma[im], int(xn_plasma[im]/nfp)) + zmnc[im] = s.get_zc(xm_plasma[im], int(xn_plasma[im]/nfp)) + + rmnc_plasma = np.copy(rmnc) + rmns_plasma = np.copy(rmns) + zmns_plasma = np.copy(zmns) + zmnc_plasma = np.copy(zmnc) + + # go through and compute all the rmnc, rmns for the coil surface + xn_potential = self.current_potential.n * w.nfp + xm_potential = self.current_potential.m + if self._ref_xm_potential is not None and self._ref_xn_potential is not None: + xm_potential = self._ref_xm_potential + xn_potential = self._ref_xn_potential + else: + xm_potential = xm_potential if self.current_potential.stellsym else xm_potential[:len(xm_potential) // 2] + xn_potential = xn_potential if self.current_potential.stellsym else xn_potential[:len(xn_potential) // 2] + xn_coil = w.n * w.nfp + xm_coil = w.m + rmnc = np.zeros(len(xm_coil)) + rmns = np.zeros(len(xm_coil)) + zmnc = np.zeros(len(xm_coil)) + zmns = np.zeros(len(xm_coil)) + nfp = w.nfp + for im in range(len(xm_coil)): + rmnc[im] = w.get_rc(xm_coil[im], int(xn_coil[im]/nfp)) + zmns[im] = w.get_zs(xm_coil[im], int(xn_coil[im]/nfp)) + if not w.stellsym: + rmns[im] = w.get_rs(xm_coil[im], int(xn_coil[im]/nfp)) + zmnc[im] = w.get_zc(xm_coil[im], int(xn_coil[im]/nfp)) + + # get the RHS b vector in the optimization + RHS_B, _ = self.B_matrix_and_rhs() + + # Define geometric objects and then compute all the Bnormals + points = s.gamma().reshape(-1, 3) + normal = s.normal().reshape(-1, 3) + ws_points = w.gamma().reshape(-1, 3) + ws_normal = w.normal().reshape(-1, 3) + dtheta_coil = w.quadpoints_theta[1] + dzeta_coil = w.quadpoints_phi[1] + Bnormal_totals = [] + Bnormal_totals_l1 = [] + if len(self.ilambdas_l2) > 0: + for i, ilambda in enumerate(self.ilambdas_l2): + # current_potential_l2 stores 1 field period; tile to full torus for WindingSurfaceBn_REGCOIL + cp_full = np.tile(self.current_potential_l2[i], (nfp, 1)) + Bnormal_regcoil_sv = sopp.WindingSurfaceBn_REGCOIL(points, ws_points, ws_normal, cp_full, normal) * dtheta_coil * dzeta_coil + Bnormal_totals.append((Bnormal_regcoil_sv + self.B_GI + self.Bnormal_plasma).reshape(self.ntheta_plasma, self.nzeta_plasma)) + if len(self.ilambdas_l1) > 0: + for i, ilambda in enumerate(self.ilambdas_l1): + # current_potential_l1 stores 1 field period; tile to full torus for WindingSurfaceBn_REGCOIL + cp_full = np.tile(self.current_potential_l1[i], (nfp, 1)) + Bnormal_regcoil_sv = sopp.WindingSurfaceBn_REGCOIL(points, ws_points, ws_normal, cp_full, normal) * dtheta_coil * dzeta_coil + Bnormal_totals_l1.append((Bnormal_regcoil_sv + self.B_GI + self.Bnormal_plasma).reshape(self.ntheta_plasma, self.nzeta_plasma)) + + vectors = ['Bnormal_from_plasma_current', 'Bnormal_from_net_coil_currents', + 'rmnc_plasma', 'rmns_plasma', 'zmns_plasma', 'zmnc_plasma', + 'rmnc_coil', 'rmns_coil', 'zmns_coil', 'zmnc_coil', + 'xm_plasma', 'xn_plasma', 'xm_coil', 'xn_coil', + 'xm_potential', 'xn_potential', + 'r_plasma', 'r_coil', + 'theta_coil', 'zeta_coil', + 'RHS_B', 'RHS_regularization', + 'norm_normal_plasma', 'norm_normal_coil', + 'single_valued_current_potential_mn', 'single_valued_current_potential_thetazeta', + 'current_potential', + 'K2', 'lambda', 'chi2_B', 'chi2_K', 'Bnormal_total', + 'single_valued_current_potential_mn_l1', 'single_valued_current_potential_thetazeta_l1', + 'current_potential_l1', + 'K2_l1', 'lambda_l1', 'chi2_B_l1', 'chi2_K_l1', 'Bnormal_total_l1' + ] + + # Define the full plasma surface and few other geometric quantities + quadpoints_phi = np.linspace(0, 1, self.nzeta_plasma * nfp + 1, endpoint=True) + quadpoints_theta = np.linspace(0, 1, self.ntheta_plasma + 1, endpoint=True) + quadpoints_phi = quadpoints_phi[:-1] + quadpoints_theta = quadpoints_theta[:-1] + sf = SurfaceRZFourier( + nfp=s.nfp, + mpol=s.mpol, + ntor=s.ntor, + stellsym=s.stellsym, + quadpoints_phi=quadpoints_phi, + quadpoints_theta=quadpoints_theta + ) + sf.set_dofs(0 * sf.get_dofs()) + for im in range(len(xm_plasma)): + sf.set_rc(xm_plasma[im], int(xn_plasma[im] / s.nfp), rmnc_plasma[im]) + sf.set_zs(xm_plasma[im], int(xn_plasma[im] / s.nfp), zmns_plasma[im]) + if not sf.stellsym: + sf.set_rs(xm_plasma[im], int(xn_plasma[im] / s.nfp), rmns_plasma[im]) + sf.set_zc(xm_plasma[im], int(xn_plasma[im] / s.nfp), zmnc_plasma[im]) + + norm_normal_plasma = np.linalg.norm(s.normal(), axis=-1) / (2 * np.pi * 2 * np.pi) + norm_normal_coil = np.linalg.norm(w.normal(), axis=-1) / (2 * np.pi * 2 * np.pi) + + current_potential = [] + for i in range(len(self.current_potential_l2)): + current_potential.append(self.current_potential_l2[i] + self.current_potential.current_potential_secular[:self.nzeta_coil // nfp, :]) + + current_potential_l1 = [] + for i in range(len(self.current_potential_l1)): + current_potential_l1.append(self.current_potential_l1[i] + self.current_potential.current_potential_secular[:self.nzeta_coil // nfp, :]) + + # Define all the vectors we need to save + vector_variables = [self.Bnormal_plasma.reshape(self.ntheta_plasma, self.nzeta_plasma), + self.B_GI.reshape(self.ntheta_plasma, self.nzeta_plasma), + rmnc_plasma, rmns_plasma, zmns_plasma, zmnc_plasma, + rmnc, rmns, zmns, zmnc, + xm_plasma[:(len(xm_plasma) // 2) + 1 if s.stellsym else (len(xm_plasma) // 4) + 1], + xn_plasma[:(len(xn_plasma) // 2) + 1 if s.stellsym else (len(xn_plasma) // 4) + 1], + xm_coil[:(len(xm_coil) // 2) + 1 if w.stellsym else (len(xm_coil) // 4) + 1], + xn_coil[:(len(xn_coil) // 2) + 1 if w.stellsym else (len(xn_coil) // 4) + 1], + xm_potential, + xn_potential, + sf.gamma(), + w.gamma(), + w.quadpoints_theta * 2 * np.pi, + w.quadpoints_phi[:self.nzeta_coil // w.nfp] * 2 * np.pi, + RHS_B, self.K_rhs(), norm_normal_plasma, norm_normal_coil, + np.array(self.dofs_l2), np.array(self.current_potential_l2), + np.array(current_potential), + np.array(self.K2s_l2)[:, :self.nzeta_coil // w.nfp, :], + np.array(self.ilambdas_l2), + 2 * np.array(self.fBs_l2), 2 * np.array(self.fKs_l2), + np.array(Bnormal_totals), + np.array(self.dofs_l1), np.array(self.current_potential_l1), + np.array(current_potential_l1), + np.array(self.K2s_l1)[:, :self.nzeta_coil // w.nfp, :], np.array(self.ilambdas_l1), + 2 * np.array(self.fBs_l1), 2 * np.array(self.fKs_l1), np.array(Bnormal_totals_l1) + ] + + # Loop through and save all the vector variables + for i_vector, vector_name in enumerate(vectors): + vector_shape = vector_variables[i_vector].shape + shape_tuple = (vector_name + '0', ) + for j, vshape in enumerate(vector_shape): + f.createDimension(vector_name + str(j), vshape) + if j > 0: + shape_tuple = shape_tuple + (vector_name + str(j),) + var = f.createVariable(vector_name, 'f', shape_tuple) + if 'Bnormal' in vector_name: + var.units = 'Tesla' + elif 'chi2_B' in vector_name: + var.units = 'Tesla^2 m^2' + elif 'chi2_K' in vector_name: + var.units = 'Ampere^2 / m^2' + elif 'thetazeta' in vector_name: + var.units = 'Ampere / m' + elif 'rmn' in vector_name or 'zmn' in vector_name or 'r_' in vector_name or 'norm_' in vector_name: + var.units = 'm' + elif 'lambda' in vector_name: # lambda * chi2_K must have same units as chi2_B + var.units = 'Tesla^2 * m^4 / Ampere^2' + else: + var.units = 'dimensionless' + var[:] = vector_variables[i_vector] + + f.close() + + def K_rhs_impl(self, K_rhs: np.ndarray) -> None: + """ + Implied function for the K rhs for the REGCOIL problem. + + Args: + K_rhs: The K rhs to be computed. + """ + dg1 = self.winding_surface.gammadash1() + dg2 = self.winding_surface.gammadash2() + normal = self.winding_surface.normal() + self.current_potential.K_rhs_impl_helper(K_rhs, dg1, dg2, normal) + K_rhs *= self.winding_surface.quadpoints_theta[1] * self.winding_surface.quadpoints_phi[1] / self.winding_surface.nfp + + def K_rhs(self) -> np.ndarray: + """ + Compute the K rhs for the REGCOIL problem. + + Args: + None + + Returns: + K_rhs: The K rhs for the REGCOIL problem. + """ + K_rhs = np.zeros((self.current_potential.num_dofs(),)) + self.K_rhs_impl(K_rhs) + return K_rhs + + def K_matrix_impl(self, K_matrix: np.ndarray) -> None: + """ + Implied function for the K matrix for the REGCOIL problem. + + Args: + K_matrix: The K matrix to be computed. + """ + dg1 = self.winding_surface.gammadash1() + dg2 = self.winding_surface.gammadash2() + normal = self.winding_surface.normal() + self.current_potential.K_matrix_impl_helper(K_matrix, dg1, dg2, normal) + K_matrix *= self.winding_surface.quadpoints_theta[1] * self.winding_surface.quadpoints_phi[1] / self.winding_surface.nfp + + def K_matrix(self) -> np.ndarray: + """ + Compute the K matrix for the REGCOIL problem. + + Args: + None + + Returns: + K_matrix: The K matrix for the REGCOIL problem. + """ + K_matrix = np.zeros((self.current_potential.num_dofs(), self.current_potential.num_dofs())) + self.K_matrix_impl(K_matrix) + return K_matrix + + def B_matrix_and_rhs(self) -> Tuple[np.ndarray, np.ndarray]: + """ + Compute the matrices and right-hand-side corresponding the Bnormal part of + the optimization, for both the Tikhonov and Lasso optimizations. + """ + plasma_surface = self.plasma_surface + normal = self.winding_surface.normal().reshape(-1, 3) + Bnormal_plasma = self.Bnormal_plasma + normal_plasma = plasma_surface.normal().reshape(-1, 3) + points_plasma = plasma_surface.gamma().reshape(-1, 3) + points_coil = self.winding_surface.gamma().reshape(-1, 3) + theta = self.winding_surface.quadpoints_theta + phi_mesh, theta_mesh = np.meshgrid(self.winding_surface.quadpoints_phi, theta, indexing='ij') + zeta_coil = np.ravel(phi_mesh) + theta_coil = np.ravel(theta_mesh) + + if self.winding_surface.stellsym: + ndofs_half = self.current_potential.num_dofs() + else: + ndofs_half = self.current_potential.num_dofs() // 2 + + # Compute terms for the REGCOIL (L2) problem + contig = np.ascontiguousarray + gj, B_matrix = sopp.winding_surface_field_Bn(contig(points_plasma), contig(points_coil), contig(normal_plasma), contig(normal), self.winding_surface.stellsym, contig(zeta_coil), contig(theta_coil), self.current_potential.num_dofs(), contig(self.current_potential.m[:ndofs_half]), contig(self.current_potential.n[:ndofs_half]), self.winding_surface.nfp) + B_GI = self.B_GI + + # set up RHS of optimization + b_rhs = - np.ravel(B_GI + Bnormal_plasma) @ gj + dzeta_plasma = (plasma_surface.quadpoints_phi[1] - plasma_surface.quadpoints_phi[0]) + dtheta_plasma = (plasma_surface.quadpoints_theta[1] - plasma_surface.quadpoints_theta[0]) + dzeta_coil = (self.winding_surface.quadpoints_phi[1] - self.winding_surface.quadpoints_phi[0]) + dtheta_coil = (self.winding_surface.quadpoints_theta[1] - self.winding_surface.quadpoints_theta[0]) + + # scale bmatrix and b_rhs by factors of the grid spacing + b_rhs = b_rhs * dzeta_plasma * dtheta_plasma * dzeta_coil * dtheta_coil + B_matrix = B_matrix * dzeta_plasma * dtheta_plasma * dzeta_coil ** 2 * dtheta_coil ** 2 + normN = np.linalg.norm(self.plasma_surface.normal().reshape(-1, 3), axis=-1) + self.gj = gj * np.sqrt(dzeta_plasma * dtheta_plasma * dzeta_coil ** 2 * dtheta_coil ** 2) + self.b_e = - np.sqrt(normN * dzeta_plasma * dtheta_plasma) * (B_GI + Bnormal_plasma) + + normN = np.linalg.norm(self.winding_surface.normal().reshape(-1, 3), axis=-1) + dr_dzeta = self.winding_surface.gammadash1().reshape(-1, 3) + dr_dtheta = self.winding_surface.gammadash2().reshape(-1, 3) + G = self.current_potential.net_poloidal_current_amperes + I = self.current_potential.net_toroidal_current_amperes + + normal_coil = self.winding_surface.normal().reshape(-1, 3) + m = self.current_potential.m[:ndofs_half] + n = self.current_potential.n[:ndofs_half] + nfp = self.winding_surface.nfp + + contig = np.ascontiguousarray + + # Compute terms for the Lasso (L1) problem + d, fj = sopp.winding_surface_field_K2_matrices( + contig(dr_dzeta), contig(dr_dtheta), contig(normal_coil), self.winding_surface.stellsym, + contig(zeta_coil), contig(theta_coil), self.ndofs, contig(m), contig(n), nfp, G, I + ) + self.fj = fj * 2 * np.pi * np.sqrt(dzeta_coil * dtheta_coil) + self.d = d * 2 * np.pi * np.sqrt(dzeta_coil * dtheta_coil) + return b_rhs, B_matrix + + def solve_tikhonov( + self, + lam: float = 0, + record_history: bool = True, + ) -> Tuple[np.ndarray, float, float]: + """ + Solve the REGCOIL problem -- winding surface optimization with + the L2 norm. This is tested against REGCOIL runs extensively in + tests/field/test_regcoil.py. + + Args: + lam: Regularization parameter for the Tikhonov regularization. + record_history: Whether to record the history of the optimization. + + Returns: + phi_mn_opt: The optimized winding surface potential. + f_B: The value of the Bnormal loss term. + f_K: The value of the K loss term. + + Notes: + This function solves the REGCOIL problem with Tikhonov regularization. + The regularization term is added to the Bnormal and K matrices to + prevent overfitting. The optimization is performed using a least-squares + solve. The history of the optimization is recorded if record_history is True. + The history is recorded in the ilambdas_l2, dofs_l2, current_potential_l2, + fBs_l2, and fKs_l2 lists. + """ + K_matrix = self.K_matrix() + K_rhs = self.K_rhs() + b_rhs, B_matrix = self.B_matrix_and_rhs() + + # least-squares solve + phi_mn_opt = np.linalg.solve(B_matrix + lam * K_matrix, b_rhs + lam * K_rhs) + self.current_potential.set_dofs(phi_mn_opt) + + # Get other matrices for direct computation of fB and fK loss terms + nfp = self.plasma_surface.nfp + normN = np.linalg.norm(self.plasma_surface.normal().reshape(-1, 3), axis=-1) + A_times_phi = self.gj @ phi_mn_opt / np.sqrt(normN) + b_e = self.b_e + Ak_times_phi = self.fj @ phi_mn_opt + f_B = 0.5 * np.linalg.norm(A_times_phi - b_e) ** 2 * nfp + # extra normN factor needed here because fj and d don't have it + # K^2 has 1/normn^2 factor, the sum over the winding surface has factor of normn, + # for total factor of 1/normn + normN = np.linalg.norm(self.winding_surface.normal().reshape(-1, 3), axis=-1) + f_K = 0.5 * np.linalg.norm((Ak_times_phi - self.d) / np.sqrt(normN[:, None])) ** 2 + + if record_history: + self.ilambdas_l2.append(lam) + self.dofs_l2.append(phi_mn_opt) + # REGCOIL only uses 1 / 2 nfp of the winding surface + self.current_potential_l2.append(np.copy(self.current_potential.Phi()[:self.nzeta_coil // nfp, :])) + self.fBs_l2.append(f_B) + self.fKs_l2.append(f_K) + K2 = np.sum(self.current_potential.K() ** 2, axis=2) + self.K2s_l2.append(K2) + return phi_mn_opt, f_B, f_K + + def solve_lasso( + self, + lam: float = 0, + max_iter: int = 1000, + acceleration: bool = True, + ) -> Tuple[np.ndarray, float, float, List[float], List[float]]: + """ + Solve the Lasso problem -- winding surface optimization with + the L1 norm, which should tend to allow stronger current + filaments to form than the L2. There are a couple changes to make: + + 1. Need to define new optimization variable z = A_k * phi_mn - b_k + so that optimization becomes + ||AA_k^{-1} * z - (b - A * A_k^{-1} * b_k)||_2^2 + alpha * ||z||_1 + which is the form required to use the Lasso pre-built optimizer + from sklearn (which actually works poorly) or the optimizer used + here (proximal gradient descent for LASSO, also called ISTA or FISTA). + 2. The alpha term should be similar amount of regularization as the L2 + so we rescale lam -> sqrt(lam) since lam is used for the (L2 norm)^2 + loss term used for Tikhonov regularization. + + + We use the FISTA algorithm but you could use scikit-learn's Lasso + optimizer too. Like any gradient-based algorithm, both FISTA + and Lasso will work very poorly at low regularization since convergence + goes like the condition number of the fB matrix. In both cases, + this can be addressed by using the exact Tikhonov solution (obtained + with a matrix inverse instead of a gradient-based optimization) as + an initial guess to the optimizers. + """ + # Set up some matrices + _, _ = self.B_matrix_and_rhs() + normN = np.linalg.norm(self.plasma_surface.normal().reshape(-1, 3), axis=-1) + ws_normN = np.linalg.norm(self.winding_surface.normal().reshape(-1, 3), axis=-1) + A_matrix = self.gj + for i in range(self.gj.shape[0]): + A_matrix[i, :] *= (1.0 / np.sqrt(normN[i])) + b_e = self.b_e + fj = self.fj / np.sqrt(ws_normN)[:, None, None] + d = self.d / np.sqrt(ws_normN)[:, None] + Ak_matrix = fj.reshape(fj.shape[0] * 3, fj.shape[-1]) + d = np.ravel(d) + nfp = self.plasma_surface.nfp + + # Ak is non-square so pinv required. Careful with rcond parameter. + # SVD can fail on ill-conditioned matrices (e.g. at very low lambda); use + # scipy fallback which may handle edge cases better on some platforms. + try: + Ak_inv = np.linalg.pinv(Ak_matrix, rcond=1e-10) + except np.linalg.LinAlgError: + from scipy.linalg import pinv as scipy_pinv + Ak_inv = scipy_pinv(Ak_matrix, rtol=1e-8) + A_new = A_matrix @ Ak_inv + b_new = b_e - A_new @ d + + # rescale the l1 regularization + l1_reg = lam + # l1_reg = np.sqrt(lam) + + # if alpha << 1, want to use initial guess from the Tikhonov solve, + # which is exact since it comes from a matrix inverse. + phi0, _, _, = self.solve_tikhonov(lam=lam, record_history=False) + + # L1 norm here should already include the contributions from the winding surface discretization + # and factor of 1 / ws_normN from the K, cancelling the factor of ws_normN from the surface + z0 = np.ravel((Ak_matrix @ phi0 - d).reshape(-1, 3) * np.sqrt(ws_normN)[:, None]) + z_opt, z_history = self._FISTA(A=A_new, b=b_new, alpha=l1_reg, max_iter=max_iter, acceleration=acceleration, xi0=z0) + # Need to put back in the 1 / ws_normN dependence in K + + # Compute the history of values from the optimizer + phi_history = [] + fB_history = [] + fK_history = [] + for i in range(len(z_history)): + phi_history.append(Ak_inv @ (z_history[i] + d)) + fB_history.append(0.5 * np.linalg.norm(A_matrix @ phi_history[i] - b_e) ** 2 * nfp) + fK_history.append(np.linalg.norm(z_history[i], ord=1)) + # fK_history.append(np.linalg.norm(Ak_matrix @ phi_history[i] - d, ord=1)) + + # Remember, Lasso solved for z = A_k * phi_mn - b_k so need to convert back + phi_mn_opt = Ak_inv @ (z_opt + d) + self.current_potential.set_dofs(phi_mn_opt) + f_B = 0.5 * np.linalg.norm(A_matrix @ phi_mn_opt - b_e) ** 2 * nfp + f_K = np.linalg.norm(Ak_matrix @ phi_mn_opt - d, ord=1) + self.ilambdas_l1.append(lam) + self.dofs_l1.append(phi_mn_opt) + # REGCOIL only uses 1 / 2 nfp of the winding surface + self.current_potential_l1.append(np.copy(self.current_potential.Phi()[:self.nzeta_coil // nfp, :])) + self.fBs_l1.append(f_B) + self.fKs_l1.append(f_K) + K2 = np.sum(self.current_potential.K() ** 2, axis=2) + self.K2s_l1.append(K2) + return phi_mn_opt, f_B, f_K, fB_history, fK_history + + def _FISTA( + self, + A: np.ndarray, + b: np.ndarray, + alpha: float = 0.0, + max_iter: int = 1000, + acceleration: bool = True, + xi0: Optional[np.ndarray] = None, + ) -> Tuple[np.ndarray, List[np.ndarray]]: + """ + This function uses Nesterov's accelerated proximal + gradient descent algorithm to solve the Lasso + (L1-regularized) winding surface problem. This is + usually called fast iterative soft-thresholding algorithm + (FISTA). If acceleration = False, it will use the ISTA + algorithm, which tends to converge much slower than FISTA + but is a true descent algorithm, unlike FISTA. Here n = the + number of plasma quadrature points and + m = the number of winding surface DOFs. The algorithm is: + + .. math:: + x_{k+1} = \\text{prox}_{\\alpha \\| \\cdot \\|_1}(x_k + \\frac{1}{L} (b - A x_k)) + x_0 = \\text{prox}_{\\alpha \\| \\cdot \\|_1}(x_0) + + Args: + A (shape (n, m)): The A matrix. + b (shape (n,)): The b vector. + alpha (float): The regularization parameter. + max_iter (int): The maximum number of iterations. + acceleration (bool): Whether to use acceleration. + xi0 (shape (m,)): The initial guess. + + Returns: + z_opt: The optimized z vector. + z_history: The history of the z vector. + """ + from scipy.sparse.linalg import svds + + # Tolerance for algorithm to consider itself converged + tol = 1e-5 + + # pre-compute/load some stuff so for loops (below) are faster + AT = A.T + ATb = AT @ b + prox = self._prox_l1 + + # L = largest eigenvalue of A^T A = (largest singular value of A)^2. + # svds uses iterative methods (ARPACK) with only A@v and A.T@v matvecs, + # avoiding O(n*m^2) cost of forming A^T A explicitly. + try: + sigma_max = svds(A, k=1, which='LM', return_singular_vectors=False)[0] + L = sigma_max ** 2 + except Exception: + ATA = AT @ A + L = np.max(np.sum(np.abs(ATA), axis=1)) # Gershgorin upper bound + + # initial step size should be just smaller than 1 / L + # which for most of these problems L ~ 1e-13 or smaller + # so the step size is enormous + ti = 1.0 / L + + # initialize current potential to random values in [5e-4, 5e4] + if xi0 is None: + xi0 = (np.random.rand(A.shape[1]) - 0.5) * 1e5 + x_history = [xi0] + if acceleration: # FISTA algorithm + # first iteration do ISTA + x_prev = xi0 + x = prox(xi0 + ti * (ATb - AT @ (A @ xi0)), ti * alpha) + for i in range(1, max_iter): + vi = x + i / (i + 3) * (x - x_prev) + x_prev = x + # note l1 'threshold' is rescaled here + x = prox(vi + ti * (ATb - AT @ (A @ vi)), ti * alpha) + ti = (1 + np.sqrt(1 + 4 * ti ** 2)) / 2.0 + if (i % 100) == 0: + x_history.append(x) + if np.all(abs(x_history[-1] - x_history[-2]) / max(1e-10, np.mean(np.abs(x_history[-2]))) < tol): + break + else: # ISTA algorithm (forms ATA only when needed) + alpha = ti * alpha # ti does not vary in ISTA algorithm + AT_ti = ti * AT + # I_ATA = np.eye(ATA.shape[0]) - ATA + ATb_scaled = ti * ATb + for i in range(max_iter): + x_history.append(prox(ATb_scaled + x_history[i] - (AT_ti @ (A @ x_history[i])), alpha)) + if (i % 100) == 0: + if np.all(abs(x_history[i + 1] - x_history[i]) / max(1e-10, np.mean(np.abs(x_history[i]))) < tol): + break + xi = x_history[-1] + return xi, x_history + + def _prox_l1(self, x: np.ndarray, threshold: float) -> np.ndarray: + """ + Proximal operator for L1 regularization, + which is often called soft-thresholding. + + .. math:: + \\text{prox}_{\\lambda \\| \\cdot \\|_1}(x) = \\text{sign}(x) \\max(0, |x| - \\lambda) + + Args: + x (shape (n,)): The x vector. + threshold (float): The threshold for the L1 regularization. + + Returns: + prox_x (shape (n,)): The proximal operator of the L1 regularization. + """ + return np.sign(x) * np.maximum(np.abs(x) - threshold, 0) diff --git a/src/simsopt/field/magneticfieldclasses.py b/src/simsopt/field/magneticfieldclasses.py index 215b6bb3e..45387e92b 100644 --- a/src/simsopt/field/magneticfieldclasses.py +++ b/src/simsopt/field/magneticfieldclasses.py @@ -18,7 +18,7 @@ __all__ = ['ToroidalField', 'PoloidalField', 'ScalarPotentialRZMagneticField', 'CircularCoil', 'Dommaschk', 'Reiman', 'InterpolatedField', 'DipoleField', - 'MirrorModel'] + 'MirrorModel', 'WindingSurfaceField'] class ToroidalField(MagneticField): @@ -569,6 +569,63 @@ def wrap(data): polyLinesToVTK(str(filename), x, y, z, pointsPerLine=ppl) +class WindingSurfaceField(MagneticField): + """ + Magnetic field object associated with a winding surface coil + optimization with a surface potential. + + Args: + current_potential: CurrentPotential class object containing + the winding surface and the surface current needed for + fast computation of the magnetic field and vector potential. + """ + + def __init__(self, current_potential): + MagneticField.__init__(self, depends_on=[current_potential]) + self.current_potential = current_potential + self.ws_points = current_potential.winding_surface.gamma().reshape((-1, 3)) + self.ws_normal = current_potential.winding_surface.normal().reshape((-1, 3)) + self.K = current_potential.K().reshape((self.ws_points.shape[0], 3)) + self.nphi = len(current_potential.winding_surface.quadpoints_phi) + self.ntheta = len(current_potential.winding_surface.quadpoints_theta) + + def _B_impl(self, B): + points = self.get_points_cart_ref() + B[:] = sopp.WindingSurfaceB(points, self.ws_points, self.ws_normal, self.K) / self.nphi / self.ntheta + + def _A_impl(self, A): + points = self.get_points_cart_ref() + A[:] = sopp.WindingSurfaceA(points, self.ws_points, self.ws_normal, self.K) / self.nphi / self.ntheta + + def _dB_by_dX_impl(self, dB): + points = self.get_points_cart_ref() + dB[:] = sopp.WindingSurfacedB(points, self.ws_points, self.ws_normal, self.K) / self.nphi / self.ntheta + + def _dA_by_dX_impl(self, dA): + points = self.get_points_cart_ref() + dA[:] = sopp.WindingSurfacedA(points, self.ws_points, self.ws_normal, self.K) / self.nphi / self.ntheta + + def as_dict(self, serial_objs_dict) -> dict: + d = super().as_dict(serial_objs_dict=serial_objs_dict) + name = getattr(self.current_potential, "name", str(id(self.current_potential))) + if name not in serial_objs_dict: + serial_objs_dict[name] = self.current_potential.as_dict(serial_objs_dict) + d["current_potential"] = {"$type": "ref", "value": name} + d["points"] = self.get_points_cart() + return d + + @classmethod + def from_dict(cls, d, serial_objs_dict, recon_objs): + decoder = GSONDecoder() + current_potential = decoder.process_decoded( + d["current_potential"], serial_objs_dict, recon_objs + ) + field = cls(current_potential) + xyz = decoder.process_decoded(d["points"], serial_objs_dict, recon_objs) + field.set_points_cart(xyz) + return field + + class DipoleField(MagneticField): r""" Computes the MagneticField induced by N dipoles. The field is given by diff --git a/src/simsopt/util/__init__.py b/src/simsopt/util/__init__.py index fd3549af1..afac2d31f 100644 --- a/src/simsopt/util/__init__.py +++ b/src/simsopt/util/__init__.py @@ -6,6 +6,7 @@ from .polarization_project import * from .permanent_magnet_helper_functions import * from .coil_optimization_helper_functions import * +from .winding_surface_helper_functions import * """Boolean indicating if we are in the GitHub actions CI""" in_github_actions = "CI" in os.environ and os.environ['CI'].lower() in ['1', 'true'] @@ -17,5 +18,6 @@ + polarization_project.__all__ + permanent_magnet_helper_functions.__all__ + coil_optimization_helper_functions.__all__ + + winding_surface_helper_functions.__all__ + ['in_github_actions'] ) diff --git a/src/simsopt/util/winding_surface_helper_functions.py b/src/simsopt/util/winding_surface_helper_functions.py new file mode 100644 index 000000000..3baeb6057 --- /dev/null +++ b/src/simsopt/util/winding_surface_helper_functions.py @@ -0,0 +1,1888 @@ +""" +Winding surface and current potential contour utilities for coil extraction. + +This module provides helper functions for extracting coil contours from a current +potential defined on a winding surface. It supports both REGCOIL (legacy) and +simsopt-regcoil NetCDF output formats, as well as loading directly from +:class:`simsopt.field.CurrentPotentialSolve` objects. + +The workflow typically involves: +1. Loading current potential data (from file or object) +2. Computing current potential and |∇φ| (surface current density) on a grid +3. Selecting contours (interactively or by level/points) +4. Classifying contours as window-pane (closed), modular, or helical +5. Computing currents for each contour +6. Converting contours to 3D curves and simsopt Coil objects + +See the cut_coils example in examples/3_Advanced/ for a complete usage. +""" + +from __future__ import annotations + +__all__ = ["run_cut_coils"] + +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple, Union + +import numpy as np +from matplotlib import path +from matplotlib import pyplot as plt +from scipy.io import netcdf_file + + +# ============================================================================= +# Current potential evaluation +# ============================================================================= + + +def _current_potential_at_point( + x: Union[np.ndarray, List[float]], args: Tuple[Any, ...] +) -> float: + """ + Evaluate the current potential φ at a point (θ, ζ) in flux coordinates. + + The potential has the form: + φ = Σ (φ_cos[m,n] cos(mθ - nζ) + φ_sin[m,n] sin(mθ - nζ)) + I_t θ/(2π) + I_p ζ/(2π) + + where the sum is over Fourier modes (m, n), I_t is net toroidal current, + and I_p is net poloidal current. + + Args: + x: Point [θ, ζ] (poloidal, toroidal angle in radians). Can be list or array. + args: Tuple (Ip, It, xm_potential, xn_potential, phi_cos, phi_sin, nfp) + with phi_cos, phi_sin arrays of same length as xm_potential, xn_potential. + + Returns: + Current potential value at the point (float). + """ + Ip, It, xm_potential, xn_potential, phi_cos, phi_sin, nfp = args + t, z = x[0], x[1] + angle = xm_potential * t - xn_potential * z + MV = It * t / (2 * np.pi) + Ip * z / (2 * np.pi) + phi = np.sum(phi_cos * np.cos(angle) + phi_sin * np.sin(angle)) + MV + return float(phi[0]) if isinstance(phi, np.ndarray) else float(phi) + + +def _grad_current_potential_at_point( + x: Union[np.ndarray, List[float]], args: Tuple[Any, ...] +) -> float: + """ + Evaluate |∇φ| (magnitude of current potential gradient) at (θ, ζ). + + This approximates the surface current density magnitude |K|. + + Args: + x: Array [θ, ζ]. + args: Tuple (Ip, It, xm_potential, xn_potential, phi_cos, phi_sin, nfp). + + Returns: + |∇φ| at the point. + """ + Ip, It, xm_potential, xn_potential, phi_cos, phi_sin, nfp = args + t, z = x[0], x[1] + angle = xm_potential * t - xn_potential * z + dMV_dth = It / (2 * np.pi) + dMV_dze = Ip / (2 * np.pi) + dphi_dth = ( + np.sum( + -xm_potential * phi_cos * np.sin(angle) + + xm_potential * phi_sin * np.cos(angle) + ) + + dMV_dth + ) + dphi_dze = ( + np.sum( + xn_potential * phi_cos * np.sin(angle) + - xn_potential * phi_sin * np.cos(angle) + ) + + dMV_dze + ) + normK = np.sqrt(dphi_dth**2 + dphi_dze**2) + return float(normK[0]) if isinstance(normK, np.ndarray) else float(normK) + + +def _genCPvals( + thetaVals: Tuple[float, float], + zetaVals: Tuple[float, float], + numbers: Tuple[int, int], + args: Tuple[Any, ...], +) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, Tuple[Any, ...]]: + """ + Compute current potential on a (θ, ζ) grid. + + Evaluates full, single-valued, and non-single-valued components. + + Args: + thetaVals: (θ_min, θ_max). + zetaVals: (ζ_min, ζ_max). + numbers: (n_theta, n_zeta) grid sizes. + args: (Ip, It, xm, xn, phi_cos, phi_sin, nfp). + + Returns: + (thetas, zetas, phi_full, phi_SV, phi_NSV, (args, args_SV, args_NSV)). + """ + Ip, It, xm_potential, xn_potential, phi_cos, phi_sin, nfp = args + args_SV = (0 * Ip, 0 * It, xm_potential, xn_potential, phi_cos, phi_sin, nfp) + args_NSV = (Ip, It, xm_potential, xn_potential, 0 * phi_cos, 0 * phi_sin, nfp) + ARGS = (args, args_SV, args_NSV) + nT, nZ = numbers + thetas = np.linspace(thetaVals[0], thetaVals[1], nT, endpoint=False) + zetas = np.linspace(zetaVals[0], zetaVals[1], nZ, endpoint=False) + foo = np.zeros((nT, nZ)) + foo_SV = np.zeros((nT, nZ)) + foo_NSV = np.zeros((nT, nZ)) + for i in range(nT): + for j in range(nZ): + foo[i, j] = _current_potential_at_point([thetas[i], zetas[j]], args) + foo_SV[i, j] = _current_potential_at_point([thetas[i], zetas[j]], args_SV) + foo_NSV[i, j] = _current_potential_at_point([thetas[i], zetas[j]], args_NSV) + return (thetas, zetas, foo.T, foo_SV.T, foo_NSV.T, ARGS) + + +def _genKvals( + thetaVals: Tuple[float, float], + zetaVals: Tuple[float, float], + numbers: Tuple[int, int], + args: Tuple[Any, ...], +) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, Tuple[Any, ...]]: + """ + Compute |∇φ| on a (θ, ζ) grid. + + Args: + thetaVals: (θ_min, θ_max). + zetaVals: (ζ_min, ζ_max). + numbers: (n_theta, n_zeta). + args: (Ip, It, xm, xn, phi_cos, phi_sin, nfp). + + Returns: + (thetas, zetas, K_full, K_SV, K_NSV, (args, args_SV, args_NSV)). + """ + Ip, It, xm_potential, xn_potential, phi_cos, phi_sin, nfp = args + args_SV = (0 * Ip, 0 * It, xm_potential, xn_potential, phi_cos, phi_sin, nfp) + args_NSV = (Ip, It, xm_potential, xn_potential, 0 * phi_cos, 0 * phi_sin, nfp) + ARGS = (args, args_SV, args_NSV) + nT, nZ = numbers + thetas = np.linspace(thetaVals[0], thetaVals[1], nT, endpoint=False) + zetas = np.linspace(zetaVals[0], zetaVals[1], nZ, endpoint=False) + foo = np.zeros((nT, nZ)) + foo_SV = np.zeros((nT, nZ)) + foo_NSV = np.zeros((nT, nZ)) + for i in range(nT): + for j in range(nZ): + foo[i, j] = _grad_current_potential_at_point([thetas[i], zetas[j]], args) + foo_SV[i, j] = _grad_current_potential_at_point( + [thetas[i], zetas[j]], args_SV + ) + foo_NSV[i, j] = _grad_current_potential_at_point( + [thetas[i], zetas[j]], args_NSV + ) + return (thetas, zetas, foo.T, foo_SV.T, foo_NSV.T, ARGS) + + +# ============================================================================= +# Data loading +# ============================================================================= + + +def _load_regcoil_data(regcoilName: str) -> List[Any]: + """ + Load current potential data from a legacy REGCOIL NetCDF file. + + Args: + regcoilName: Path to the .nc file. + + Returns: + List of [ntheta_coil, nzeta_coil, theta_coil, zeta_coil, r_coil, + xm_coil, xn_coil, xm_potential, xn_potential, nfp, Ip, It, + current_potential_allLambda, single_valued_..., phi_mn_allLambda, + lambdas, chi2_B, chi2_K, K2]. + """ + with netcdf_file(regcoilName, "r", mmap=False) as f: + nfp = f.variables["nfp"][()] + ntheta_coil = f.variables["ntheta_coil"][()] + nzeta_coil = f.variables["nzeta_coil"][()] + theta_coil = f.variables["theta_coil"][()] + zeta_coil = f.variables["zeta_coil"][()] + r_coil = f.variables["r_coil"][()] + xm_coil = f.variables["xm_coil"][()] + xn_coil = f.variables["xn_coil"][()] + xm_potential = f.variables["xm_potential"][()] + xn_potential = f.variables["xn_potential"][()] + single_valued_current_potential_thetazeta_allLambda = f.variables[ + "single_valued_current_potential_thetazeta" + ][()] + current_potential_allLambda = f.variables["current_potential"][()] + net_poloidal_current_Amperes = f.variables["net_poloidal_current_Amperes"][()] + net_toroidal_current_Amperes = f.variables["net_toroidal_current_Amperes"][()] + phi_mn_allLambda = f.variables["single_valued_current_potential_mn"][()] + lambdas = f.variables["lambda"][()] + K2 = f.variables["K2"][()] + chi2_B = f.variables["chi2_B"][()] + chi2_K = f.variables["chi2_K"][()] + + nfp = np.array([nfp]) + return [ + ntheta_coil, + nzeta_coil, + theta_coil, + zeta_coil, + r_coil, + xm_coil, + xn_coil, + xm_potential, + xn_potential, + nfp, + net_poloidal_current_Amperes, + net_toroidal_current_Amperes, + current_potential_allLambda, + single_valued_current_potential_thetazeta_allLambda, + phi_mn_allLambda, + lambdas, + chi2_B, + chi2_K, + K2, + ] + + +def _load_simsopt_regcoil_data(regcoilName: str, sparse: bool = False) -> List[Any]: + """ + Load current potential data from a simsopt-regcoil NetCDF file. + + Args: + regcoilName: Path to the .nc file. + sparse: If True, load L1/sparse solution variables. + + Returns: + Same structure as load_regcoil_data. + """ + with netcdf_file(regcoilName, "r", mmap=False) as f: + nfp = f.variables["nfp"][()] + ntheta_coil = f.variables["ntheta_coil"][()] + nzeta_coil = f.variables["nzeta_coil"][()] + theta_coil = f.variables["theta_coil"][()] + zeta_coil = f.variables["zeta_coil"][()] + r_coil = f.variables["r_coil"][()] + xm_coil = f.variables["xm_coil"][()] + xn_coil = f.variables["xn_coil"][()] + xm_potential = f.variables["xm_potential"][()] + xn_potential = f.variables["xn_potential"][()] + net_poloidal_current_Amperes = f.variables["net_poloidal_current_amperes"][()] + net_toroidal_current_Amperes = f.variables["net_toroidal_current_amperes"][()] + + if sparse: + phi_mn_allLambda = f.variables["single_valued_current_potential_mn_l1"][()] + single_valued_current_potential_thetazeta_allLambda = f.variables[ + "single_valued_current_potential_thetazeta_l1" + ][()] + current_potential_allLambda = np.zeros_like( + single_valued_current_potential_thetazeta_allLambda + ) + lambdas = f.variables["lambda"][()] + K2 = f.variables["K2_l1"][()] + chi2_B = f.variables["chi2_B_l1"][()] + chi2_K = f.variables["chi2_K_l1"][()] + else: + phi_mn_allLambda = f.variables["single_valued_current_potential_mn"][()] + single_valued_current_potential_thetazeta_allLambda = f.variables[ + "single_valued_current_potential_thetazeta" + ][()] + current_potential_allLambda = np.zeros_like( + single_valued_current_potential_thetazeta_allLambda + ) + lambdas = f.variables["lambda"][()] + K2 = f.variables["K2"][()] + chi2_B = f.variables["chi2_B"][()] + chi2_K = f.variables["chi2_K"][()] + + return [ + ntheta_coil, + nzeta_coil, + theta_coil, + zeta_coil, + r_coil, + xm_coil, + xn_coil, + xm_potential, + xn_potential, + nfp, + net_poloidal_current_Amperes[0], + net_toroidal_current_Amperes[0], + np.asarray(current_potential_allLambda), + np.asarray(single_valued_current_potential_thetazeta_allLambda), + np.array(phi_mn_allLambda), + lambdas, + chi2_B, + chi2_K, + K2, + ] + + +def _load_surface_dofs_properly(s: Any, s_new: Any) -> Any: + """ + Copy surface Fourier coefficients from s to s_new with correct indexing. + + Handles the mapping between full-torus and field-period surface + representations (quadpoint ordering differs). + + Args: + s: Source SurfaceRZFourier (full torus). + s_new: Target SurfaceRZFourier (field period). + + Returns: + s_new with DOFs set from s. + """ + xm_coil = s.m + xn_coil = s.n + s_new.set_dofs(0 * s.get_dofs()) + rc_ = s.rc.copy() + rc_[0, :] = rc_[0, ::-1] + rs_ = s.rs.copy() + rs_[0, :] = rs_[0, ::-1] + zc_ = s.zc.copy() + zc_[0, :] = zc_[0, ::-1] + zs_ = s.zs.copy() + zs_[0, :] = zs_[0, ::-1] + for im in range(len(xm_coil)): + m = xm_coil[im] + n = xn_coil[im] + rc = rc_[m, n + s.ntor] + rs = rs_[m, n + s.ntor] + zc = zc_[m, n + s.ntor] + zs = zs_[m, n + s.ntor] + s_new.set_rc(xm_coil[im], xn_coil[im], rc) + s_new.set_zs(xm_coil[im], xn_coil[im], zs) + if not s.stellsym: + s_new.set_rs(xm_coil[im], xn_coil[im], rs) + s_new.set_zc(xm_coil[im], xn_coil[im], zc) + return s_new + + +def _load_CP_and_geometries( + filename: str, + plot_flags: Tuple[int, int, int, int] = (0, 0, 0, 0), + plasma_resolution_mult: Tuple[float, float] = (1.0, 1.0), + wsurf_resolution_mult: Tuple[float, float] = (1.0, 1.0), + loadDOFsProperly: bool = True, +) -> List[Any]: + """ + Load CurrentPotentialSolve and plasma/winding surfaces from a simsopt-regcoil NetCDF. + + Args: + filename: Path to the .nc file. + plot_flags: (plot coil fp, plot coil full, plot plasma fp, plot plasma full). + plasma_resolution_mult: Resolution multipliers for plasma surface. + wsurf_resolution_mult: Resolution multipliers for winding surface. + loadDOFsProperly: Whether to copy surface DOFs with correct indexing. + + Returns: + [cpst, s_coil_fp, s_coil_full, s_plasma_fp, s_plasma_full]. + """ + from simsopt.geo import SurfaceRZFourier, plot + from simsopt.field import CurrentPotentialFourier, CurrentPotentialSolve + + plasma_ntheta_res, plasma_nzeta_res = plasma_resolution_mult + coil_ntheta_res, coil_nzeta_res = wsurf_resolution_mult + + cpst = CurrentPotentialSolve.from_netcdf( + filename, plasma_ntheta_res, plasma_nzeta_res, coil_ntheta_res, coil_nzeta_res + ) + cp = CurrentPotentialFourier.from_netcdf(filename, coil_ntheta_res, coil_nzeta_res) + cp = CurrentPotentialFourier( + cpst.winding_surface, + mpol=cp.mpol, + ntor=cp.ntor, + net_poloidal_current_amperes=cp.net_poloidal_current_amperes, + net_toroidal_current_amperes=cp.net_toroidal_current_amperes, + stellsym=True, + ) + cpst = CurrentPotentialSolve(cp, cpst.plasma_surface, cpst.Bnormal_plasma) + + s_coil_full = cpst.winding_surface + s_plasma_fp = cpst.plasma_surface + + s_plasma_full = SurfaceRZFourier( + nfp=s_plasma_fp.nfp, + stellsym=s_plasma_fp.stellsym, + mpol=s_plasma_fp.mpol, + ntor=s_plasma_fp.ntor, + ) + s_plasma_full = s_plasma_full.from_nphi_ntheta( + nfp=s_plasma_fp.nfp, + ntheta=len(s_plasma_fp.quadpoints_theta), + nphi=len(s_plasma_fp.quadpoints_phi) * s_plasma_fp.nfp, + mpol=s_plasma_fp.mpol, + ntor=s_plasma_fp.ntor, + stellsym=s_plasma_fp.stellsym, + range="full torus", + ) + s_plasma_full.set_dofs(s_plasma_fp.x) + + s_coil_fp = SurfaceRZFourier( + nfp=s_coil_full.nfp, + stellsym=s_coil_full.stellsym, + mpol=s_coil_full.mpol, + ntor=s_coil_full.ntor, + ) + s_coil_fp = s_coil_fp.from_nphi_ntheta( + nfp=s_coil_full.nfp, + ntheta=len(s_coil_full.quadpoints_theta), + nphi=len(s_coil_full.quadpoints_phi) // s_coil_full.nfp, + mpol=s_coil_full.mpol, + ntor=s_coil_full.ntor, + stellsym=s_coil_full.stellsym, + range="field period", + ) + s_coil_fp.set_dofs(s_coil_full.x) + + if loadDOFsProperly: + s_coil_fp = _load_surface_dofs_properly(s=s_coil_full, s_new=s_coil_fp) + + if plot_flags[0]: + plot([s_coil_fp], alpha=0.5) + if plot_flags[1]: + plot([s_coil_full], alpha=0.5) + if plot_flags[2]: + plot([s_plasma_fp], alpha=0.5) + if plot_flags[3]: + plot([s_plasma_full], alpha=0.5) + + return [cpst, s_coil_fp, s_coil_full, s_plasma_fp, s_plasma_full] + + +# ============================================================================= +# Contour utilities +# ============================================================================= + + +def is_periodic_lines(line: np.ndarray, tol: float = 0.01) -> bool: + """ + Check if a contour is closed (start and end points coincide within tolerance). + + Args: + line: Nx2 array of (θ, ζ) points. + tol: Distance tolerance for closure. + + Returns: + True if closed (window-pane type). + """ + p_sta = line[0, :] + p_end = line[-1, :] + return bool(np.linalg.norm(p_sta - p_end) < tol) + + +def minDist(pt: np.ndarray, line: np.ndarray) -> float: + """ + Minimum distance from a point to the nearest vertex in a polyline. + + Args: + pt: Point [θ, ζ] or (x, y). + line: Nx2 array of polyline vertices. + + Returns: + Minimum Euclidean distance from pt to any vertex in line. + """ + diffs = pt[np.newaxis, :] - line + dists = np.linalg.norm(diffs, axis=1) + return float(np.min(dists)) + + +def _removearray(L: List[np.ndarray], arr: np.ndarray) -> None: + """ + Remove the first list element that equals arr (by value). + + Args: + L: List of arrays (modified in place). + arr: Array to remove (matched by np.array_equal). + + Raises: + ValueError: If arr is not found in L. + + Returns: + None. Modifies L in place. + """ + for ind in range(len(L)): + if np.array_equal(L[ind], arr): + L.pop(ind) + return + raise ValueError("array not found in list.") + + +def sortLevels(levels: List[float], points: List[Any]) -> Tuple[List[float], List[np.ndarray]]: + """ + Sort levels and points by level value (ascending). + + Args: + levels: List of contour level values. + points: List of (θ, ζ) points corresponding to each level. + + Returns: + Tuple of (levels_sorted, points_sorted) with same lengths as inputs. + """ + indexing = np.argsort(levels) + levels = [levels[i] for i in indexing] + points = [np.asarray(points[i]) for i in indexing] + return (levels, points) + + +def chooseContours_matching_coilType(lines: List[np.ndarray], ctype: str) -> List[np.ndarray]: + """ + Filter contours by coil type: 'free' (all), 'wp' (closed), 'mod'/'hel' (open). + + Args: + lines: List of contour arrays. + ctype: 'free', 'wp', 'mod', or 'hel'. + + Returns: + Filtered list of contours. + """ + if ctype == "free": + return lines + closed = [is_periodic_lines(line) for line in lines] + ret = [] + if ctype == "wp": + for i, c in enumerate(closed): + if c: + ret.append(lines[i]) + if ctype in ("mod", "hel"): + for i, c in enumerate(closed): + if not c: + ret.append(lines[i]) + return ret + + +def _points_in_polygon( + coords: Tuple[np.ndarray, np.ndarray], + polygon: np.ndarray, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Find grid points inside a polygon. + + Args: + coords: (t_values, z_values) 1D arrays defining the grid. + polygon: Nx2 array of vertices defining the polygon (θ, ζ or x, y). + + Returns: + Tuple (i1_ret, i2_ret, BOOL) where i1_ret, i2_ret are indices into + coords[0] and coords[1] for points inside the polygon, and BOOL is + a flattened mask of shape (len(coords[0])*len(coords[1]),). + """ + p = path.Path(polygon) + XX, YY = np.meshgrid(coords[0], coords[1]) + I1, I2 = np.meshgrid(np.arange(len(coords[0])), np.arange(len(coords[1]))) + xx = XX.flatten() + yy = YY.flatten() + i1 = I1.flatten() + i2 = I2.flatten() + pts = np.vstack([xx, yy]).T + BOOL = p.contains_points(pts) + i1_ret = i1[BOOL] + i2_ret = i2[BOOL] + return (i1_ret, i2_ret, BOOL) + + +def _map_data_full_torus_3x3( + xdata: np.ndarray, + ydata: np.ndarray, + zdata: np.ndarray, + Ip: float, + It: float, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Tile full-torus (θ, ζ ∈ [0, 2π]) SV data 3×3 and add secular (NSV) contribution. + + Used when the base grid already covers the full torus [0, 2π] × [0, 2π]. + Tiles with period 2π in both θ and ζ (not 2π/nfp). + + Args: + xdata: 1D array of poloidal angles θ. + ydata: 1D array of toroidal angles ζ. + zdata: 2D single-valued current potential on the base grid. + Ip: Net poloidal current (for secular term). + It: Net toroidal current (for secular term NSV = (It/2π)θ + (Ip/2π)ζ). + + Returns: + Tuple (xN, yN, ret) of extended grid and data. + """ + xf = 2 * np.pi + yf = 2 * np.pi + xN = np.hstack([xdata - xf, xdata, xdata + xf]) + yN = np.hstack([ydata - yf, ydata, ydata + yf]) + XX, YY = np.meshgrid(xN, yN) + NSV = (It / (2 * np.pi)) * XX + (Ip / (2 * np.pi)) * YY + z0_tiled = np.tile(zdata, (3, 3)) + return (xN, yN, z0_tiled + NSV) + + +def _map_data_full_torus_3x1( + xdata: np.ndarray, + ydata: np.ndarray, + zdata: np.ndarray, + Ip: float, + It: float, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Tile full-torus data 3×1 in ζ (period 2π) for halfway-contour extraction. + + One period in θ, three periods in ζ by 2π. Used for single-valued + current potential when finding halfway contours between open contours. + + Args: + xdata: 1D array of poloidal angles θ. + ydata: 1D array of toroidal angles ζ. + zdata: 2D single-valued current potential on the base grid. + Ip: Net poloidal current (for secular term). + It: Net toroidal current (for secular term NSV = (It/2π)θ + (Ip/2π)ζ). + + Returns: + Tuple (xN, yN, ret) of extended grid and data. + """ + xN = xdata + yf = 2 * np.pi + yN = np.hstack([ydata - yf, ydata, ydata + yf]) + XX, YY = np.meshgrid(xN, yN) + NSV = (It / (2 * np.pi)) * XX + (Ip / (2 * np.pi)) * YY + z0_tiled = np.tile(zdata, (3, 1)) + return (xN, yN, z0_tiled + NSV) + + +def _tile_helical_contours( + contours: List[np.ndarray], types_contours: List[int], nfp: int +) -> None: + """ + Tile helical contours (type 2 or 3) nfp times to span full torus. + + Modifies contours in place. Each helical contour at a given level spans + one field period in ζ; we tile nfp copies at shifted ζ to build the + full helical coil. + + Args: + contours: List of contour arrays (modified in place). + types_contours: Contour type codes (1=modular, 2/3=helical). + nfp: Number of field periods. + + Returns: + None. Modifies contours in place. + """ + d_zeta = 2 * np.pi / nfp + for i in range(len(contours)): + if types_contours[i] in (2, 3): + contour = contours[i] + L = len(contour[:, 0]) + Contour = np.zeros((L * nfp, 2)) + for j in range(nfp): + j1, j2 = j * L, (j + 1) * L + Contour[j1:j2, 0] = contour[:, 0] + Contour[j1:j2, 1] = contour[:, 1] - j * d_zeta + contours[i] = Contour + + +def _ID_mod_hel( + contours: List[np.ndarray], tol: float = 0.05 +) -> Tuple[List[str], List[int]]: + """ + Classify open contours as modular (1), helical (2), or vacuum-field (3). + + Args: + contours: List of open contour arrays (Nx2 θ, ζ). + tol: Tolerance for cos(θ) and ζ matching at contour endpoints. + + Returns: + (names, ints) e.g. (['mod','hel'], [1,2]). + """ + ret = [] + ints = [] + for contour in contours: + th0 = contour[0, 0] + thf = contour[-1, 0] + ze0 = contour[0, 1] + zef = contour[-1, 1] + start = np.array([np.cos(th0), ze0]) + end = np.array([np.cos(thf), zef]) + if np.linalg.norm(start - end) < tol: + ret.append("mod") + ints.append(1) + elif np.isclose(np.abs(th0 - thf), 2 * np.pi, rtol=1e-2): + ret.append("hel") + ints.append(2) + else: + ret.append("vf") + ints.append(3) + return (ret, ints) + + +def ID_and_cut_contour_types( + contours: List[np.ndarray], +) -> Tuple[List[np.ndarray], List[np.ndarray], List[int]]: + """ + Classify contours as closed (type 0) or open, and sub-classify open as mod(1)/hel(2)/vf(3). + + Args: + contours: List of contour arrays (each Nx2 θ, ζ). + + Returns: + (open_contours, closed_contours, type_array). + """ + closed = [is_periodic_lines(c) for c in contours] + surfInt = [] + lineInt = [] + type_array = [] + for i in range(len(contours)): + if closed[i]: + surfInt.append(i) + type_array.append(0) + else: + lineInt.append(i) + type_array.append(1) + Open_contours = [contours[i] for i in lineInt] + Closed_contours = [contours[i] for i in surfInt] + names, idcs = _ID_mod_hel(Open_contours) + for i in range(len(Open_contours)): + type_array[lineInt[i]] = idcs[i] + return (Open_contours, Closed_contours, type_array) + + +def ID_halfway_contour( + contours: List[np.ndarray], + data: Tuple[np.ndarray, np.ndarray, np.ndarray], + do_plot: bool, + args: Tuple[Any, ...], +) -> List[np.ndarray]: + """ + Find "halfway" contours between consecutive open contours for current integration. + + Args: + contours: List of open contour arrays (Nx2 θ, ζ). + data: Tuple (theta, zeta, current_potential) on extended grid. + do_plot: Whether to show debug plot. + args: Tuple (Ip, It, xm, xn, phi_cos, phi_sin, nfp) for potential evaluation. + + Returns: + List of halfway contour arrays. + """ + theta, zeta, cpd = data + halfway_contours = [] + fig = plt.figure() + ax = fig.add_subplot(111) + if do_plot: + ax.set_title('Identification of "halfway" contours', fontsize=25) + ax.set_xlabel(r"$\theta$", fontsize=45) + ax.set_ylabel(r"$\zeta$", fontsize=45) + ax.tick_params(axis="both", which="major", labelsize=30) + for icoil in contours: + plt.plot( + icoil[:, 0], icoil[:, 1], linestyle="-", color="orangered", picker=True + ) + for i in range(len(contours) - 1): + eye = i % len(contours) + eyep1 = (i + 1) % len(contours) + coil1 = contours[eye] + coil2 = contours[eyep1] + Lvl1 = _current_potential_at_point([coil1[0, 0], coil1[0, 1]], args) + Lvl2 = _current_potential_at_point([coil2[0, 0], coil2[0, 1]], args) + Lvlnew = (Lvl1 + Lvl2) / 2 + CS = ax.contour(theta, zeta, cpd, [Lvlnew], linestyles="dashed") + paths = _contour_paths(CS, 0) + halfway_contours.append(paths[0]) + if not do_plot: + plt.close() + else: + ax.figure.canvas.draw() + plt.show() + return halfway_contours + + +def compute_baseline_WP_currents( + closed_contours: List[np.ndarray], + args: Tuple[Any, ...], + plot: bool = False, +) -> Tuple[List[float], List[Any], List[float]]: + """ + Compute currents for window-pane (closed) contours from potential difference. + + For each closed contour, samples points inside and uses max potential difference + to estimate the current enclosed. + + Args: + closed_contours: List of closed contour arrays (Nx2 θ, ζ). + args: Tuple (Ip, It, xm, xn, phi_cos, phi_sin, nfp) for potential evaluation. + plot: Whether to plot sampling points. + + Returns: + (wp_currents, max_closed_contours_func_val, closed_contours_func_vals). + """ + COORDS = [] + res = 20 + closed_contours_func_vals = [] + max_closed_contours_func_val = [] + for cont in closed_contours: + t = np.linspace(np.min(cont[:, 0]), np.max(cont[:, 0]), res) + z = np.linspace(np.min(cont[:, 1]), np.max(cont[:, 1]), res) + i1_ret, i2_ret, BOOL = _points_in_polygon([t, z], cont) + coordinates = np.vstack([t[i1_ret], z[i2_ret]]).T + COORDS.append(coordinates) + func_vals = [_current_potential_at_point(pt, args) for pt in coordinates] + contour_value = _current_potential_at_point(cont[0, :], args) + Z = np.abs(contour_value - np.array(func_vals)) + idx = np.argmax(Z) + max_closed_contours_func_val.append([coordinates[idx, :], np.max(Z)]) + closed_contours_func_vals.append(contour_value) + + if plot: + fig = plt.figure() + ax = fig.gca() + ax.set_xlabel(r"$\theta$", fontsize=35) + ax.set_ylabel(r"$\zeta$", fontsize=35) + cs = ["r", "g", "b", "m"] + ss = ["s", "o", "x", "^"] + for i in range(len(closed_contours)): + ax.plot( + closed_contours[i][:, 0], + closed_contours[i][:, 1], + marker="", + linestyle="--", + color=cs[i % len(cs)], + ) + ax.plot( + COORDS[i][:, 0], + COORDS[i][:, 1], + marker=ss[i % len(ss)], + linestyle="", + color=cs[i % len(cs)], + ) + ax.plot( + max_closed_contours_func_val[i][0][0], + max_closed_contours_func_val[i][0][1], + marker="o", + linestyle="", + color="c", + ) + plt.show() + + wp_currents = [] + for i in range(len(closed_contours_func_vals)): + wp_currents.append( + max_closed_contours_func_val[i][1] - closed_contours_func_vals[i] + ) + return (wp_currents, max_closed_contours_func_val, closed_contours_func_vals) + + +def check_and_compute_nested_WP_currents( + closed_contours: List[np.ndarray], + wp_currents: List[float], + args: Tuple[Any, ...], + plot: bool = False, +) -> Tuple[List[float], Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]: + """ + Adjust window-pane currents when contours are nested (one inside another). + + Args: + closed_contours: List of closed contour arrays (Nx2 θ, ζ). + wp_currents: Initial current estimates (one per contour). + args: Tuple (Ip, It, xm, xn, phi_cos, phi_sin, nfp) for potential evaluation. + plot: Whether to plot. + + Returns: + (wp_currents, nested_BOOL, func_vals, numContsContained). + """ + nested_BOOL = np.zeros((len(closed_contours), len(closed_contours)), dtype=bool) + func_vals = np.zeros((len(closed_contours), len(closed_contours), 2)) + for i, cont1 in enumerate(closed_contours): + for j, cont2 in enumerate(closed_contours): + if i == j: + continue + test_pt = cont2[0, :] + p = path.Path(cont1) + BOOL = p.contains_points([test_pt]) + if np.sum(BOOL) > 0: + nested_BOOL[i, j] = True + func_vals[i, j, :] = np.array( + [ + _current_potential_at_point(cont1[0, :], args), + _current_potential_at_point(test_pt, args), + ] + ) + numContsContained = np.sum(nested_BOOL, axis=1) + if np.sum(numContsContained) == 0: + return (wp_currents, None, None, None) + for i in range(len(closed_contours)): + for j in range(len(closed_contours)): + if nested_BOOL[i, j]: + wp_currents[i] = func_vals[i, j, 0] - func_vals[i, j, 1] + if plot: + fig = plt.figure() + ax = fig.add_subplot(111) + ax.set_title(r"$\kappa$ chosen contours", fontsize=25) + ax.tick_params(axis="both", which="major", labelsize=30) + ax.set_xlabel(r"$\theta$", fontsize=45) + ax.set_ylabel(r"$\zeta$", fontsize=45) + for i in range(len(closed_contours)): + M = f"${_current_potential_at_point(closed_contours[i][0, :], args):.0f}$" + ax.plot( + closed_contours[i][0, 0], + closed_contours[i][0, 1], + marker=M, + linestyle="--", + color="r", + markersize=50, + ) + ax.plot( + closed_contours[i][:, 0], + closed_contours[i][:, 1], + marker="", + linestyle="--", + color="k", + ) + plt.show() + return (wp_currents, nested_BOOL, func_vals, numContsContained) + + +def _SIMSOPT_line_XYZ_RZ( + surf: Any, coords: Tuple[np.ndarray, np.ndarray] +) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """ + Map (θ, ζ) to (X, Y, R, Z) on a simsopt SurfaceRZFourier. + + Uses the surface's gamma_lin for correct evaluation. Contour angles (theta, zeta) + are in radians; theta is poloidal [0, 2π), zeta is toroidal. + + Args: + surf: SurfaceRZFourier instance. + coords: (theta_array, zeta_array) in radians. + + Returns: + (X, Y, R, Z) Cartesian and cylindrical coordinates on the surface. + """ + coords_the, coords_ze = np.asarray(coords[0]), np.asarray(coords[1]) + # Surface quadpoints are in [0, 1); physical angles = 2π * quadpoints + # theta (poloidal): quadpoints_theta = theta / (2π) + # phi (toroidal): SurfaceRZFourier uses phi_phys = 2π * quadpoints_phi (one turn = 2π) + # so quadpoints_phi = zeta / (2π). zeta in [0, 2π] maps to full torus. + # Wrap into [0, 1) for closing points with zeta > 2π. + quadpoints_theta = np.mod(coords_the / (2 * np.pi), 1.0) + quadpoints_phi = np.mod(coords_ze / (2 * np.pi), 1.0) + data = np.zeros((len(coords_the), 3)) + surf.gamma_lin(data, quadpoints_phi, quadpoints_theta) + XX = data[:, 0] + YY = data[:, 1] + RR = np.sqrt(XX**2 + YY**2) + ZZ = data[:, 2] + return (XX, YY, RR, ZZ) + + +def _writeToCurve( + contour: np.ndarray, + contour_xyz: List[Any], + fourier_trunc: int, + plotting_args: Tuple[int, ...] = (0,), + winding_surface: Optional[Any] = None, + fix_stellarator_symmetry: Optional[bool] = None, +) -> Any: + """ + Convert a (θ, ζ) contour to a simsopt CurveXYZFourier. + + Args: + contour: Nx2 (θ, ζ) points. + contour_xyz: [X, Y, R, Z] from _SIMSOPT_line_XYZ_RZ or real_space. + fourier_trunc: Number of Fourier modes. + plotting_args: (plot_comparison,). + winding_surface: If contour_xyz empty, compute from this surface. + fix_stellarator_symmetry: If True, fix xs, yc, zc so the curve stays + stellarator-symmetric during optimization. If None, auto-detect for + helical contours (3D closed but not θζ closed). + + Returns: + CurveXYZFourier instance. + """ + from simsopt.geo import CurveXYZFourier + + contour_theta = contour[:, 0] + contour_zeta = contour[:, 1] + if len(contour_xyz) == 0 and winding_surface is not None: + X, Y, R, Z = _SIMSOPT_line_XYZ_RZ( + winding_surface, [contour_theta, contour_zeta] + ) + contour_xyz = [X, Y, R, Z] + X, Y, R, Z = contour_xyz + XYZ = np.column_stack([X, Y, Z]) + + # Arc-length parametrization along the 3D contour (not theta-zeta space) + ds = np.sqrt(np.sum(np.diff(XYZ, axis=0) ** 2, axis=1)) + ds_sum = np.sum(ds) + 1e-12 + closed_theta_zeta = np.allclose(contour[0], contour[-1]) + closed_3d = np.linalg.norm(XYZ[-1] - XYZ[0]) < 1e-6 * ds_sum + + if closed_theta_zeta or closed_3d: + ds = np.append(ds, np.linalg.norm(XYZ[-1] - XYZ[0])) + else: + ds = np.append(ds, 0) + S = np.cumsum(ds) + S = S / S[-1] if S[-1] > 0 else np.linspace(0, 1, len(S), endpoint=False) + # Enforce periodicity for closed curves (smooth closure for Fourier fit) + if closed_theta_zeta or closed_3d: + XYZ = np.asarray(XYZ, dtype=float) + XYZ[-1] = XYZ[0] + # Resample to uniform arc-length spacing for least-squares fit + from scipy.interpolate import interp1d + + n_quad = max(4 * len(contour_theta), 128) + s_uniform = np.linspace(0, 1, n_quad, endpoint=False) + target_xyz = np.column_stack( + [ + interp1d(S, XYZ[:, 0], kind="linear", fill_value="extrapolate")(s_uniform), + interp1d(S, XYZ[:, 1], kind="linear", fill_value="extrapolate")(s_uniform), + interp1d(S, XYZ[:, 2], kind="linear", fill_value="extrapolate")(s_uniform), + ] + ) + quadpoints = list(s_uniform) + curve = CurveXYZFourier(quadpoints, fourier_trunc) + curve.least_squares_fit(target_xyz) + ORD = fourier_trunc + + # For helical coils: fix stellarator-symmetric coefficients (xs, yc, zc) + if fix_stellarator_symmetry is None: + fix_stellarator_symmetry = closed_3d and not closed_theta_zeta + if fix_stellarator_symmetry: + for m in range(1, ORD + 1): + curve.fix(f"xs({m})") + for m in range(ORD + 1): + curve.fix(f"yc({m})") + curve.fix(f"zc({m})") + + if plotting_args[0]: + ax = curve.plot(show=False) + ax.plot( + target_xyz[:, 0], + target_xyz[:, 1], + target_xyz[:, 2], + marker="", + color="r", + markersize=10, + ) + XYZcurve = curve.gamma() + ax.plot( + XYZcurve[:, 0], + XYZcurve[:, 1], + XYZcurve[:, 2], + marker="", + color="k", + markersize=10, + linestyle="--", + ) + plt.show() + return curve + + +def set_axes_equal(ax: Any) -> None: + """ + Make 3D axes have equal scale (spheres appear spherical). + + Workaround for matplotlib's set_aspect('equal') not working in 3D. + + Args: + ax: Matplotlib 3D axes instance. + + Returns: + None. Modifies ax in place. + """ + x_limits = ax.get_xlim3d() + y_limits = ax.get_ylim3d() + z_limits = ax.get_zlim3d() + x_range = abs(x_limits[1] - x_limits[0]) + x_middle = np.mean(x_limits) + y_range = abs(y_limits[1] - y_limits[0]) + y_middle = np.mean(y_limits) + z_range = abs(z_limits[1] - z_limits[0]) + z_middle = np.mean(z_limits) + plot_radius = 0.5 * max([x_range, y_range, z_range]) + ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius]) + ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius]) + ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius]) + + +def _contour_paths(cdata: Any, level_idx: int = 0) -> List[np.ndarray]: + """ + Get contour path vertices (matplotlib version-agnostic). + + Args: + cdata: Matplotlib contour object (ContourSet). + level_idx: Index of the contour level. + + Returns: + List of Nx2 arrays, one per contour path at that level. + """ + if hasattr(cdata, "allsegs") and level_idx < len(cdata.allsegs): + return list(cdata.allsegs[level_idx]) + return [p.vertices for p in cdata.collections[level_idx].get_paths()] + + +# Interactive contour selection callbacks (used by CutCoils example) +def _make_onclick( + ax: Any, + args: Tuple[Any, ...], + contours: List[np.ndarray], + theta_coil: np.ndarray, + zeta_coil: np.ndarray, + current_potential: np.ndarray, +) -> Callable[..., None]: + """ + Factory for double-click contour selection callback. + + Args: + ax: Matplotlib axes instance. + args: Tuple (Ip, It, xm, xn, phi_cos, phi_sin, nfp) for potential evaluation. + contours: List to append selected contour to (modified in place). + theta_coil: 1D array of poloidal angles for the contour grid. + zeta_coil: 1D array of toroidal angles for the contour grid. + current_potential: 2D current potential array. + + Returns: + Callback function for connect('button_press_event'). + """ + + def onclick(event): + if event.dblclick: + xdata, ydata = event.xdata, event.ydata + click = np.array([xdata, ydata]) + tmp_value = _current_potential_at_point(click, args) + tmp_level = np.array([tmp_value]) + print( + f"click coords theta,zeta=({xdata:.2f},{ydata:.2f}) with contour value {tmp_value:.3f}" + ) + cdata = plt.contour( + theta_coil, zeta_coil, current_potential, tmp_level, linestyles="dashed" + ) + lines = _contour_paths(cdata, 0) + minDists = [minDist(click, line) for line in lines] + chosen_line_idx = np.argmin(minDists) + chosen_line = [lines[chosen_line_idx]] + for icoil in chosen_line: + plt.plot( + icoil[:, 0], icoil[:, 1], linestyle="--", color="r", picker=True + ) + ax.figure.canvas.draw() + contours.append(chosen_line[0]) + + return onclick + + +def _make_onpick(contours: List[np.ndarray]) -> Callable[..., None]: + """ + Factory for pick event (store picked artist). + + Args: + contours: Unused; kept for API compatibility with _make_on_key. + + Returns: + Callback function for connect('pick_event'). + """ + + def onpick(event): + plt.gca().picked_object = event.artist + + return onpick + + +def _make_on_key(contours: List[np.ndarray]) -> Callable[..., None]: + """ + Factory for key press (delete picked contour). + + Args: + contours: List of contours; removes the picked contour on 'delete' key. + + Returns: + Callback function for connect('key_press_event'). + """ + + def on_key(event): + if event.key == "delete": + ax = plt.gca() + if hasattr(ax, "picked_object") and ax.picked_object is not None: + arr = ax.picked_object.get_xydata() + _removearray(contours, arr) + ax.picked_object.remove() + ax.picked_object = None + ax.figure.canvas.draw() + + return on_key + + +def _contour_paths(cdata: Any, i: int) -> List[np.ndarray]: + """ + Get contour path vertices from contour set (matplotlib version-agnostic). + + Args: + cdata: Matplotlib QuadContourSet or contour object. + i: Index of the contour level. + + Returns: + List of Nx2 arrays, one per contour path at that level. + """ + if hasattr(cdata, "allsegs") and i < len(cdata.allsegs): + return list(cdata.allsegs[i]) + return [p.vertices for p in cdata.collections[i].get_paths()] + + +def _run_cut_coils_select_contours_non_interactive( + theta_coil: np.ndarray, + zeta_coil: np.ndarray, + current_potential_: np.ndarray, + args: Tuple[Any, ...], + coil_type: str, + points: Optional[List[List[float]]], + levels: Optional[List[float]], + contours_per_period: Optional[int], + ax: Any, +) -> List[np.ndarray]: + """ + Select coil contours non-interactively from the current potential. + + Uses one of four selection modes (in priority order): + 1. **points**: (θ, ζ) coordinates on desired contours. For each point, finds the + nearest contour at that level and filters by coil_type (wp/mod/hel). + 2. **levels**: Explicit contour level values. Extracts all contours at those levels + and filters by coil_type. + 3. **contours_per_period**: Number of contours to distribute uniformly between + min and max potential. Extracts all and filters by coil_type. + 4. **default**: Uses three fixed (θ, ζ) points for demonstration. + + Plots selected contours on the given axes in red dashed lines. + + Args: + theta_coil: 1D array of poloidal angles (θ) for the grid. + zeta_coil: 1D array of toroidal angles (ζ) for the grid. + current_potential_: 2D current potential on the grid. + args: Tuple (Ip, It, xm, xn, phi_cos, phi_sin, nfp) for potential evaluation. + coil_type: 'free', 'wp', 'mod', or 'hel' to filter contour types. + points: Optional list of [θ, ζ] points; None to use levels or default. + levels: Optional list of contour level values; None to use contours_per_period. + contours_per_period: Optional int; None to use default points. + ax: Matplotlib axes instance to plot contours on. + + Returns: + List of contour arrays, each Nx2 (θ, ζ) in radians. + """ + contours = [] + lines = [] + if points is not None: + Pts = [np.array(pt) for pt in points] + LvlsTemp = [_current_potential_at_point(pt, args) for pt in Pts] + Lvls, Pts = sortLevels(LvlsTemp, Pts) + for i in range(len(Lvls)): + cdata = plt.contour( + theta_coil, zeta_coil, current_potential_, Lvls, linestyles="dashed" + ) + paths = _contour_paths(cdata, i) + minDists = [minDist(Pts[i], line) for line in paths] + chosen_line_idx = np.argmin(minDists) + chosen_line = paths[chosen_line_idx] + if coil_type == "wp" and not is_periodic_lines(chosen_line): + print(f"Skipping non-window-pane contour at {Pts[i]}") + continue + if coil_type in ("mod", "hel") and is_periodic_lines(chosen_line): + print(f"Skipping window-pane contour at {Pts[i]}") + continue + lines.append(chosen_line) + contours = lines + elif levels is not None: + Lvls = levels + Pts = [None] * len(Lvls) + Lvls, Pts = sortLevels(Lvls, Pts) + for i in range(len(Lvls)): + cdata = plt.contour( + theta_coil, zeta_coil, current_potential_, Lvls, linestyles="dashed" + ) + paths = _contour_paths(cdata, i) + for p in paths: + lines.append(p) + contours = chooseContours_matching_coilType(lines, coil_type) + elif contours_per_period is not None: + Lvls = np.linspace( + np.min(current_potential_), np.max(current_potential_), contours_per_period + ) + Lvls, Pts = sortLevels(list(Lvls), list(Lvls)) + for i in range(len(Lvls)): + cdata = plt.contour( + theta_coil, zeta_coil, current_potential_, Lvls, linestyles="dashed" + ) + paths = _contour_paths(cdata, i) + for p in paths: + lines.append(p) + contours = chooseContours_matching_coilType(lines, coil_type) + else: + Pts = [[0.5, 0.5], [1.0, 0.3], [1.5, 0.5]] + LvlsTemp = [_current_potential_at_point(np.array(pt), args) for pt in Pts] + Lvls, Pts = sortLevels(LvlsTemp, Pts) + for i in range(len(Lvls)): + cdata = plt.contour( + theta_coil, zeta_coil, current_potential_, Lvls, linestyles="dashed" + ) + paths = _contour_paths(cdata, i) + minDists = [minDist(Pts[i], line) for line in paths] + chosen_line_idx = np.argmin(minDists) + chosen_line = paths[chosen_line_idx] + lines.append(chosen_line) + contours = lines + + for icoil in contours: + ax.plot( + icoil[:, 0], + icoil[:, 1], + linestyle="--", + color="r", + picker=True, + linewidth=1, + ) + return contours + + +def _run_cut_coils_refine_contours_3x3( + contours: List[np.ndarray], + theta_coil: np.ndarray, + zeta_coil: np.ndarray, + current_potential_SV: np.ndarray, + args: Tuple[Any, ...], + nfp_val: int, +) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray], List[int]]: + """ + Refine contours using a 3×3 periodic extension for proper marching-squares extraction. + + The base grid [0, 2π] × [0, 2π] can produce ambiguous contours at boundaries. + Tiling to 3×3 in both θ and ζ yields clean closed/open contours. This function: + + 1. Tiles the single-valued current potential 3×3 and adds the secular (NSV) term. + 2. Re-extracts contours at the same levels as the input contours. + 3. Replaces input contours with the 3×3 versions when types differ (better classification). + 4. Cuts modular (type 1) and helical (type 2) contours to one θ-period [0, 2π]. + 5. Cuts vacuum-field (type 3) contours to one ζ field-period [0, 2π/nfp]. + 6. Deletes repeated contours (marching squares duplicates for periodic data). + 7. Deletes degenerate single-point contours. + + Modifies contours in place. + + Args: + contours: List of Nx2 (θ, ζ) contour arrays; modified in place. + theta_coil: 1D poloidal grid. + zeta_coil: 1D toroidal grid. + current_potential_SV: 2D single-valued current potential. + args: (Ip, It, xm, xn, phi_cos, phi_sin, nfp) for potential evaluation. + nfp_val: Number of field periods (int). + + Returns: + Tuple (contours, open_contours, closed_contours, types_contours) where + types_contours[i] is 0 (closed), 1 (modular), 2 (helical), or 3 (vacuum-field). + """ + Ip, It = args[0], args[1] + theta_coil_3x3, zeta_coil_3x3, current_potential_3x3 = _map_data_full_torus_3x3( + theta_coil, zeta_coil, current_potential_SV, Ip, It + ) + Pts = [contour[0, :] for contour in contours] + LvlsTemp = [_current_potential_at_point(pt, args) for pt in Pts] + Lvls, Pts = sortLevels(LvlsTemp, Pts) + lines_3x3 = [] + for i in range(len(Lvls)): + cdata = plt.contour( + theta_coil_3x3, + zeta_coil_3x3, + current_potential_3x3, + [Lvls[i]], + linestyles="dashed", + ) + paths = _contour_paths(cdata, 0) + minDists = [minDist(Pts[i], line) for line in paths] + chosen_line_idx = np.argmin(minDists) + lines_3x3.append(paths[chosen_line_idx]) + plt.close("all") + + _, _, types_lines_3x3 = ID_and_cut_contour_types(lines_3x3) + open_contours, closed_contours, types_contours = ID_and_cut_contour_types(contours) + + for i in range(len(types_contours)): + if types_contours[i] in (1, 2, 3) and types_contours[i] != types_lines_3x3[i]: + contours[i] = lines_3x3[i] + for i in range(len(types_contours)): + if types_lines_3x3[i] in (1, 2, 3): + icoil = lines_3x3[i] + contour_theta = icoil[:, 0] + contour_zeta = icoil[:, 1] + if types_lines_3x3[i] in (1, 2): + arg_theta0 = np.argmin(np.abs(contour_theta)) + arg_theta2pi = np.argmin(np.abs(contour_theta - 2 * np.pi)) + if arg_theta0 > arg_theta2pi: + contours[i] = lines_3x3[i][arg_theta2pi:arg_theta0, :] + else: + contours[i] = lines_3x3[i][arg_theta0:arg_theta2pi, :] + if types_lines_3x3[i] == 3: + zeta_fp = 2 * np.pi / nfp_val + arg_zeta0 = np.argmin(np.abs(contour_zeta)) + arg_zeta2pi_nfp = np.argmin(np.abs(contour_zeta - zeta_fp)) + if arg_zeta0 > arg_zeta2pi_nfp: + contours[i] = lines_3x3[i][arg_zeta2pi_nfp:arg_zeta0, :] + else: + contours[i] = lines_3x3[i][arg_zeta0:arg_zeta2pi_nfp, :] + + open_contours, closed_contours, types_contours = ID_and_cut_contour_types(contours) + + # Delete repeated contours + delete_indices = [] + marching_squares_tol = 0.001 + for i in range(len(types_contours) - 1): + if (types_contours[i] == 1 and types_contours[i + 1] == 1) or ( + types_contours[i] == 2 and types_contours[i + 1] == 2 + ): + contour1 = contours[i] + contour2 = contours[i + 1] + minDists = [minDist(pt, contour1) for pt in contour2] + if np.min(minDists) < marching_squares_tol: + delete_indices.append(i + 1) + for idx in sorted(delete_indices, reverse=True): + print("Repeated contour found. Deleting.") + del contours[idx] + open_contours, closed_contours, types_contours = ID_and_cut_contour_types(contours) + + # Delete degenerate single-point contours + for i in range(len(contours) - 1, -1, -1): + if len(contours[i]) <= 1: + print("Degenerate single-point contour found. Deleting.") + del contours[i] + open_contours, closed_contours, types_contours = ID_and_cut_contour_types(contours) + + return contours, open_contours, closed_contours, types_contours + + +def _run_cut_coils_compute_NWP_currents( + open_contours: List[np.ndarray], + single_valued: bool, + totalCurrent: float, + theta_coil: np.ndarray, + zeta_coil: np.ndarray, + current_potential_SV: np.ndarray, + args: Tuple[Any, ...], +) -> List[float]: + """ + Compute currents for non-window-pane (open) contours. + + For a single open contour, assigns the full total current. For multiple open + contours, distributes current based on the potential difference between + adjacent contours: + + - **Multi-valued (single_valued=False)**: Current between contours i and i+1 + is proportional to the potential difference. Contours are sorted by level; + the last contour gets the remainder to conserve total current. + - **Single-valued (single_valued=True)**: Uses a 3×1 periodic extension in ζ + to find halfway contours between consecutive open contours. Current in each + band is the potential difference between halfway contours. + + Args: + open_contours: List of open contour arrays (Nx2 θ, ζ). + single_valued: If True, use halfway-contour method; else use level differences. + totalCurrent: Total net current (Ip + It/nfp) for normalization. + theta_coil: 1D poloidal grid. + zeta_coil: 1D toroidal grid. + current_potential_SV: 2D single-valued current potential. + args: (Ip, It, xm, xn, phi_cos, phi_sin, nfp) for potential evaluation. + + Returns: + List of currents, one per open contour, in the same order as open_contours. + Empty list if no open contours. + """ + if len(open_contours) == 1: + return [totalCurrent] + if len(open_contours) == 0: + return [] + + if not single_valued: + open_levels = np.array( + [_current_potential_at_point(cont[0, :], args) for cont in open_contours] + ) + sort_idx = np.argsort(open_levels) + open_levels_sorted = open_levels[sort_idx] + NWP_from_diff = np.diff(open_levels_sorted) + last_current = totalCurrent - np.sum(NWP_from_diff) + NWP_currents = np.concatenate([NWP_from_diff, [last_current]]) + inv_sort = np.argsort(sort_idx) + return list(NWP_currents[inv_sort]) + + Ip, It = args[0], args[1] + theta_3x1, zeta_3x1, current_potential_3x1 = _map_data_full_torus_3x1( + theta_coil, zeta_coil, current_potential_SV, Ip, It + ) + open_Pts = [contour[0, :] for contour in open_contours] + pt_offset = np.array([0, 2 * np.pi]) + Pts_3x1 = [open_Pts[0] - pt_offset] + open_Pts + [open_Pts[-1] + pt_offset] + LvlsTemp_3x1 = [_current_potential_at_point(pt, args) for pt in Pts_3x1] + Lvls_3x1, Pts_3x1 = sortLevels(LvlsTemp_3x1, Pts_3x1) + lines_3x1 = [] + for i in range(len(Lvls_3x1)): + cdata_3x1 = plt.contour( + theta_3x1, + zeta_3x1, + current_potential_3x1, + Lvls_3x1, + linestyles="dashed", + alpha=0.01, + ) + paths = _contour_paths(cdata_3x1, i) + minDists = [minDist(Pts_3x1[i], line) for line in paths] + chosen_line_idx = np.argmin(minDists) + lines_3x1.append(paths[chosen_line_idx]) + open_contours_3x1, _, _ = ID_and_cut_contour_types(lines_3x1) + halfway_contours_3x1 = ID_halfway_contour( + open_contours_3x1, [theta_3x1, zeta_3x1, current_potential_3x1], False, args + ) + halfway_contours_Lvls = np.array( + [_current_potential_at_point(cont[0, :], args) for cont in halfway_contours_3x1] + ) + return list(np.diff(halfway_contours_Lvls)) + + +def run_cut_coils( + surface_filename: Union[str, Path], + ilambda: int = -1, + num_of_contours: int = 100, + ntheta: int = 128, + nzeta: int = 128, + single_valued: bool = False, + coil_type: str = "free", + points: Optional[List[List[float]]] = None, + levels: Optional[List[float]] = None, + contours_per_period: Optional[int] = None, + interactive: bool = False, + curve_fourier_cutoff: int = 20, + map_to_full_torus: bool = True, + modular_coils_via_symmetries: bool = True, + helical_coils_via_symmetries: bool = True, + show_final_coilset: bool = True, + show_plots: bool = True, + write_coils_to_file: bool = False, + output_path: Optional[Union[str, Path]] = None, +) -> List[Any]: + """ + Run the cut coils workflow to extract coils from a current potential. + + Args: + surface_filename: Path to the REGCOIL or simsopt-regcoil NetCDF file. + ilambda: Index of regularization parameter to use (0-based). + num_of_contours: Number of contour levels for visualization. + ntheta: Grid resolution in poloidal direction. + nzeta: Grid resolution in toroidal direction. + single_valued: If True, use single-valued current potential (zero net currents). + coil_type: 'free' (all), 'wp' (window-pane only), 'mod' or 'hel' (open only). + points: Optional list of [θ, ζ] points on desired contours. Overrides levels. + levels: Optional list of contour level values. Overrides contours_per_period. + contours_per_period: Optional number of contours to distribute per field period. + interactive: If True, use double-click to select contours; press Delete to remove. + curve_fourier_cutoff: Fourier mode cutoff for curve representation. + map_to_full_torus: If True, use nfp symmetry to map coils to full torus. + modular_coils_via_symmetries: If True, duplicate modular/WP coils via stellarator symmetry. + helical_coils_via_symmetries: If True, also duplicate helical coils via symmetry. + show_final_coilset: If True, show 3D plot of final coils. + show_plots: If True, display contour and coil plots; if False, close figures. + write_coils_to_file: If True, save coils to JSON. + output_path: Output directory for VTK, JSON. Default: winding_surface_/. + + Returns: + List of extracted Coil objects. + """ + from pathlib import Path + from simsopt.field import ( + RegularizedCoil, + Current, + coils_to_vtk, + regularization_circ, + ) + from simsopt.field import BiotSavart, coils_via_symmetries + from simsopt.field.magneticfieldclasses import WindingSurfaceField + from simsopt.geo import plot + from simsopt.geo.curveobjectives import CurveLength + from simsopt import save + from simsopt.util import in_github_actions + + surface_filename = Path(surface_filename) + output_path = Path(output_path) if output_path is not None else Path(f"winding_surface_{surface_filename.stem}") + output_path.mkdir(parents=True, exist_ok=True) + + # Load current potential data (auto-detect format) + try: + current_potential_vars = _load_simsopt_regcoil_data( + str(surface_filename), sparse=False + ) + except KeyError: + current_potential_vars = _load_regcoil_data(str(surface_filename)) + + # Load geometries for 3D mapping (works for both formats) + cpst, s_coil_fp, s_coil_full, s_plasma_fp, s_plasma_full = _load_CP_and_geometries( + str(surface_filename), plot_flags=(0, 0, 0, 0) + ) + + ( + ntheta_coil, + nzeta_coil, + theta_coil, + zeta_coil, + r_coil, + xm_coil, + xn_coil, + xm_potential, + xn_potential, + nfp, + Ip, + It, + current_potential_allLambda, + single_valued_current_potential_thetazeta_allLambda, + phi_mn_allLambda, + lambdas, + chi2_B, + chi2_K, + K2, + ) = current_potential_vars + + nfp_val = nfp[0] + print( + f"ilambda={ilambda + 1} of {len(lambdas)}, λ={lambdas[ilambda]:.2e}, χ²_B={chi2_B[ilambda]:.2e}" + ) + + # Build args for current potential evaluation + mn_max = len(xm_potential) + phi_mn = phi_mn_allLambda[ilambda] + phi_cos = np.zeros(mn_max) + phi_sin = phi_mn[:mn_max] + current_potential_ = current_potential_allLambda[ilambda] + args = (Ip, It, xm_potential, xn_potential, phi_cos, phi_sin, nfp) + + if single_valued: + Ip, It = 0, 0 + args = (0, 0, xm_potential, xn_potential, phi_cos, phi_sin, nfp) + + # Base grid: full winding surface [0, 2π] × [0, 2π] for contour selection. + theta_coil = np.linspace(0, 2 * np.pi, ntheta, endpoint=False) + zeta_coil = np.linspace(0, 2 * np.pi, nzeta * nfp_val, endpoint=False) + theta_range, zeta_range = [0, 2 * np.pi], [0, 2 * np.pi] + grid_shape = (ntheta, nzeta * nfp_val) + + # Compute current potential and |K| + _, _, current_potential_, current_potential_SV, _, _ = _genCPvals( + theta_range, zeta_range, grid_shape, args + ) + _, _, K_full, K_SV, _, ARGS = _genKvals(theta_range, zeta_range, grid_shape, args) + args, args_SV, _ = ARGS + if single_valued: + current_potential_ = current_potential_SV.copy() + args = args_SV + K_vals = K_SV + else: + K_vals = K_full + + # Initialize contour selection + contours = [] + totalCurrent = np.sum(Ip + (It / nfp_val)) + + # Contour plot: full winding surface (θ, ζ ∈ [0, 2π]) with current potential + fig = plt.figure() + ax = fig.add_subplot(111) + ax.set_xlabel(r"$\theta$", fontsize=14) + ax.set_ylabel(r"$\zeta$", fontsize=14) + cf = ax.contourf(theta_coil, zeta_coil, K_vals, levels=50, cmap="viridis") + ax.contour( + theta_coil, + zeta_coil, + current_potential_, + num_of_contours, + linewidths=0.5, + colors="k", + alpha=0.5, + ) + _ = fig.colorbar(cf, ax=ax, label=r"Current density $|\nabla\varphi|$") + ax.set_xlim([0, 2 * np.pi]) + ax.set_ylim([0, 2 * np.pi]) + ax.set_aspect("equal") + ax.set_title( + rf"Full winding surface $\lambda$={lambdas[ilambda]:.2e} $\chi^2_B$={chi2_B[ilambda]:.2e}" + ) + + # Contour selection + if interactive: + print( + "Interactive mode: Double-click to select contours. Press Delete to remove. Close figure when done." + ) + _ = fig.canvas.mpl_connect( + "button_press_event", + _make_onclick( + ax, args, contours, theta_coil, zeta_coil, current_potential_ + ), + ) + fig.canvas.mpl_connect("pick_event", _make_onpick(contours)) + fig.canvas.mpl_connect("key_press_event", _make_on_key(contours)) + (plt.show() if show_plots and not in_github_actions else plt.close("all")) + else: + contours = _run_cut_coils_select_contours_non_interactive( + theta_coil, + zeta_coil, + current_potential_, + args, + coil_type, + points, + levels, + contours_per_period, + ax, + ) + fig.canvas.mpl_connect("pick_event", _make_onpick(contours)) + fig.canvas.mpl_connect("key_press_event", _make_on_key(contours)) + (plt.show() if show_plots and not in_github_actions else plt.close("all")) + + if len(contours) == 0: + print("No contours selected. Exiting.") + return [] + + # Refine contours via 3×3 periodic extension + contours, open_contours, closed_contours, types_contours = ( + _run_cut_coils_refine_contours_3x3( + contours, theta_coil, zeta_coil, current_potential_SV, args, nfp_val + ) + ) + + # Compute currents + WP_currents, _, _ = compute_baseline_WP_currents(closed_contours, args, plot=False) + WP_currents, _, _, _ = check_and_compute_nested_WP_currents( + closed_contours, WP_currents, args, plot=False + ) + + NWP_currents = _run_cut_coils_compute_NWP_currents( + open_contours, + single_valued, + totalCurrent, + theta_coil, + zeta_coil, + current_potential_SV, + args, + ) + + # Tile helical contours (type 2 or 3) nfp times to span full torus. + _tile_helical_contours(contours, types_contours, nfp_val) + + # Map to Cartesian + contours_xyz = [] + for contour in contours: + contour_theta, contour_zeta = contour[:, 0], contour[:, 1] + X, Y, R, Z = _SIMSOPT_line_XYZ_RZ(s_coil_full, [contour_theta, contour_zeta]) + contours_xyz.append([X, Y, R, Z]) + + # Build curves and coils + curves = [] + for i in range(len(contours)): + curve = _writeToCurve( + contours[i], contours_xyz[i], curve_fourier_cutoff, [0], s_coil_full + ) + curves.append(curve) + + currents = [] + j, k = 0, 0 + for i in range(len(contours)): + if types_contours[i] in (1, 2, 3): + currents.append(NWP_currents[j]) + j += 1 + else: + currents.append(WP_currents[k]) + k += 1 + + # Filter out zero-current contours + zero_current_tol = 1e-14 + keep = [i for i in range(len(contours)) if abs(currents[i]) > zero_current_tol] + if len(keep) < len(contours): + for i in sorted(range(len(contours)), reverse=True): + if i not in keep: + ctype = ( + "mod" + if types_contours[i] == 1 + else "hel" + if types_contours[i] in (2, 3) + else "wp" + ) + print( + f"Removing contour {i} (type={ctype}): zero current ({currents[i]:.2e})" + ) + contours = [contours[i] for i in keep] + contours_xyz = [contours_xyz[i] for i in keep] + types_contours = [types_contours[i] for i in keep] + currents = [currents[i] for i in keep] + curves = [] + for i in range(len(contours)): + curve = _writeToCurve( + contours[i], contours_xyz[i], curve_fourier_cutoff, [0], s_coil_full + ) + curves.append(curve) + + Currents = [Current(c) for c in currents] + coils = [ + RegularizedCoil(curve, curr, regularization_circ(0.05)) + for curve, curr in zip(curves, Currents) + ] + + if map_to_full_torus and modular_coils_via_symmetries: + foo = [] + for i in range(len(contours)): + if types_contours[i] < 2 or helical_coils_via_symmetries: + coils_symm = coils_via_symmetries( + [curves[i]], + [Currents[i]], + nfp_val, + stellsym=True, + regularizations=[regularization_circ(0.05)], + ) + for c in coils_symm: + foo.append(c) + if types_contours[i] >= 2 and not helical_coils_via_symmetries: + foo.append(coils[i]) + curves = [c.curve for c in foo] + Currents = [c.current for c in foo] + currents = [c.get_value() for c in Currents] + coils = foo + + print(f"\n{len(coils)} coils extracted.") + + print("\nCoil currents (A) and lengths (m):") + for i, c in enumerate(coils): + curr = c.current.get_value() + length = float(CurveLength(c.curve).J()) + print(f" Coil {i + 1:3d}: current = {curr:14.6e} length = {length:.6f}") + + if show_final_coilset: + ax = plt.figure().add_subplot(projection="3d") + Plot = [] + if s_plasma_full is not None: + Plot.append(s_plasma_full) + elif s_plasma_fp is not None: + Plot.append(s_plasma_fp) + Plot.extend(curves) + ax = plot(Plot, ax=ax, alpha=1, color="b", close=True) + set_axes_equal(ax) + (plt.show() if show_plots and not in_github_actions else plt.close("all")) + + if write_coils_to_file: + coils_path = output_path / "coils.json" + save(coils, str(coils_path)) + print(f"Coils saved to {coils_path}") + + # Bn from WindingSurfaceField (current potential solution) + phi_mn = np.atleast_2d(phi_mn_allLambda)[ilambda] + cpst.current_potential.set_dofs(phi_mn) + wsfield = WindingSurfaceField(cpst.current_potential) + points_plasma = s_plasma_full.gamma().reshape(-1, 3) + wsfield.set_points_cart(points_plasma) + B_ws = wsfield.B().reshape(s_plasma_full.gamma().shape) + Bn_ws = np.sum(B_ws * s_plasma_full.unitnormal(), axis=2) + s_plasma_full.to_vtk( + str(output_path / "plasma_Bn_WindingSurfaceField"), + extra_data={"B_N": Bn_ws[:, :, None]}, + ) + + # Bn from BiotSavart (extracted coils) + bs = BiotSavart(coils) + bs.set_points_cart(points_plasma) + B_bs = bs.B().reshape(s_plasma_full.gamma().shape) + Bn_bs = np.sum(B_bs * s_plasma_full.unitnormal(), axis=2) + s_plasma_full.to_vtk( + str(output_path / "plasma_Bn_BiotSavart"), extra_data={"B_N": Bn_bs[:, :, None]} + ) + s_coil_full.to_vtk(str(output_path / "winding_surface")) + coils_to_vtk(coils, str(output_path / "final_coils"), close=True) + return coils diff --git a/src/simsoptpp/currentpotential.cpp b/src/simsoptpp/currentpotential.cpp new file mode 100644 index 000000000..5e3fe0c8a --- /dev/null +++ b/src/simsoptpp/currentpotential.cpp @@ -0,0 +1,103 @@ +#include "currentpotential.h" + +// template class Surface, class Array> +template +void CurrentPotential::K_impl_helper(Array& data, Array& dg1, Array& dg2, Array& normal) { + auto dphid1 = this->Phidash1(); + auto dphid2 = this->Phidash2(); + // auto dg1 = this->winding_surface->gammadash1(); + // auto dg2 = this->winding_surface->gammadash2(); + // auto normal = this->winding_surface->normal(); + // K = n \times Phi + // N \times \nabla \theta = - dr/dzeta + // N \times \nabla \zeta = dr/dtheta + // K = (- dPhidtheta dr/dzeta + dPhidzeta dr/dtheta)/N + for (int i = 0; i < numquadpoints_phi; ++i) { + for (int j = 0; j < numquadpoints_theta; ++j) { + double normn = std::sqrt(normal(i, j, 0)*normal(i, j, 0) + normal(i, j, 1)*normal(i, j, 1) + normal(i, j, 2)*normal(i, j, 2)); + data(i, j, 0) = (- dg1(i,j,0) * (dphid2(i,j) + this->net_toroidal_current_amperes) + dg2(i,j,0) * (dphid1(i,j) + this->net_poloidal_current_amperes))/normn; + data(i, j, 1) = (- dg1(i,j,1) * (dphid2(i,j) + this->net_toroidal_current_amperes) + dg2(i,j,1) * (dphid1(i,j) + this->net_poloidal_current_amperes))/normn; + data(i, j, 2) = (- dg1(i,j,2) * (dphid2(i,j) + this->net_toroidal_current_amperes) + dg2(i,j,2) * (dphid1(i,j) + this->net_poloidal_current_amperes))/normn; + } + } +} + +template +void CurrentPotential::K_GI_impl_helper(Array& data, Array& dg1, Array& dg2, Array& normal) { + // K = n \times Phi + // N \times \nabla \theta = - dr/dzeta + // N \times \nabla \zeta = dr/dtheta + // K = (- dPhidtheta dr/dzeta + dPhidzeta dr/dtheta)/N + for (int i = 0; i < numquadpoints_phi; ++i) { + for (int j = 0; j < numquadpoints_theta; ++j) { + double normn = std::sqrt(normal(i, j, 0)*normal(i, j, 0) + normal(i, j, 1)*normal(i, j, 1) + normal(i, j, 2)*normal(i, j, 2)); + data(i, j, 0) = (- dg1(i,j,0) * this->net_toroidal_current_amperes + dg2(i,j,0) * this->net_poloidal_current_amperes)/normn; + data(i, j, 1) = (- dg1(i,j,1) * this->net_toroidal_current_amperes + dg2(i,j,1) * this->net_poloidal_current_amperes)/normn; + data(i, j, 2) = (- dg1(i,j,2) * this->net_toroidal_current_amperes + dg2(i,j,2) * this->net_poloidal_current_amperes)/normn; + } + } +} + +template +void CurrentPotential::K_by_dcoeff_impl_helper(Array& data, Array& dg1, Array& dg2, Array& normal) { + auto dphid1_by_dcoeff = this->dPhidash1_by_dcoeff(); + auto dphid2_by_dcoeff = this->dPhidash2_by_dcoeff(); + // K = n \times Phi + // N \times \nabla \theta = - dr/dzeta + // N \times \nabla \zeta = dr/dtheta + // K = (- dPhidtheta dr/dzeta + dPhidzeta dr/dtheta)/N + int ndofs = num_dofs(); + for (int i = 0; i < numquadpoints_phi; ++i) { + for (int j = 0; j < numquadpoints_theta; ++j) { + double normn = std::sqrt(normal(i, j, 0)*normal(i, j, 0) + normal(i, j, 1)*normal(i, j, 1) + normal(i, j, 2)*normal(i, j, 2)); + for (int m = 0; m < ndofs; ++m ) { + data(i, j, 0, m) = (- dg1(i,j,0) * dphid2_by_dcoeff(i,j,m) + dg2(i,j,0) * dphid1_by_dcoeff(i,j,m))/normn; + data(i, j, 1, m) = (- dg1(i,j,1) * dphid2_by_dcoeff(i,j,m) + dg2(i,j,1) * dphid1_by_dcoeff(i,j,m))/normn; + data(i, j, 2, m) = (- dg1(i,j,2) * dphid2_by_dcoeff(i,j,m) + dg2(i,j,2) * dphid1_by_dcoeff(i,j,m))/normn; + } + } + } +} + +template +void CurrentPotential::K_rhs_impl_helper(Array& data, Array& dg1, Array& dg2, Array& normal) { + int ndofs = num_dofs(); + Array K_by_dcoeff = xt::zeros({numquadpoints_phi, numquadpoints_theta, 3, ndofs}); + this->K_by_dcoeff_impl_helper(K_by_dcoeff, dg1, dg2, normal); + Array K_GI = xt::zeros({numquadpoints_phi, numquadpoints_theta, 3}); + this->K_GI_impl_helper(K_GI, dg1, dg2, normal); + for (int i = 0; i < numquadpoints_phi; ++i) { + for (int j = 0; j < numquadpoints_theta; ++j) { + double normn = std::sqrt(normal(i, j, 0)*normal(i, j, 0) + normal(i, j, 1)*normal(i, j, 1) + normal(i, j, 2)*normal(i, j, 2)); + for (int m = 0; m < ndofs; ++m) { + double K_GI_dot_K_by_dcoeff = K_GI(i,j,0)*K_by_dcoeff(i,j,0,m) + K_GI(i,j,1)*K_by_dcoeff(i,j,1,m) + K_GI(i,j,2)*K_by_dcoeff(i,j,2,m); + data(m) += -K_GI_dot_K_by_dcoeff*normn; + } + } + } +}; + +template +void CurrentPotential::K_matrix_impl_helper(Array& data, Array& dg1, Array& dg2, Array& normal) { + int ndofs = num_dofs(); + Array K_by_dcoeff = xt::zeros({numquadpoints_phi, numquadpoints_theta, 3, ndofs}); + this->K_by_dcoeff_impl_helper(K_by_dcoeff, dg1, dg2, normal); + + for (int i = 0; i < numquadpoints_phi; ++i) { + for (int j = 0; j < numquadpoints_theta; ++j) { + double normn = std::sqrt(normal(i, j, 0)*normal(i, j, 0) + normal(i, j, 1)*normal(i, j, 1) + normal(i, j, 2)*normal(i, j, 2)); + + for (int m = 0; m < ndofs; ++m ) { + for (int n = 0; n < ndofs; ++ n ) { + data(m,n) += (K_by_dcoeff(i,j,0,m)*K_by_dcoeff(i,j,0,n) + \ + + K_by_dcoeff(i,j,1,m)*K_by_dcoeff(i,j,1,n) + \ + + K_by_dcoeff(i,j,2,m)*K_by_dcoeff(i,j,2,n))*normn; + } + } + } + } +}; + +#include "xtensor-python/pyarray.hpp" // Numpy bindings +typedef xt::pyarray Array; +template class CurrentPotential; diff --git a/src/simsoptpp/currentpotential.h b/src/simsoptpp/currentpotential.h new file mode 100644 index 000000000..b3e71cdd3 --- /dev/null +++ b/src/simsoptpp/currentpotential.h @@ -0,0 +1,127 @@ +#pragma once +#include +using std::vector; + +#include "xtensor-python/pyarray.hpp" // Numpy bindings +typedef xt::pyarray Array; + +#include +using std::string; + +#include +using std::map; +#include +using std::logic_error; + +#include "xtensor/xlayout.hpp" + +#include "xtensor/xarray.hpp" +#include "cachedarray.h" +#include "surface.h" +#include +using std::shared_ptr; + +template +class CurrentPotential { + private: + map> cache; + map> cache_persistent; + + Array& check_the_cache(string key, vector dims, std::function impl){ + auto loc = cache.find(key); + if(loc == cache.end()){ // Key not found --> allocate array + loc = cache.insert(std::make_pair(key, CachedArray(xt::zeros(dims)))).first; + } + if(!((loc->second).status)){ // needs recomputing + impl((loc->second).data); + (loc->second).status = true; + } + return (loc->second).data; + } + + Array& check_the_persistent_cache(string key, vector dims, std::function impl){ + auto loc = cache_persistent.find(key); + if(loc == cache_persistent.end()){ // Key not found --> allocate array + loc = cache_persistent.insert(std::make_pair(key, CachedArray(xt::zeros(dims)))).first; + } + if(!((loc->second).status)){ // needs recomputing + impl((loc->second).data); + (loc->second).status = true; + } + return (loc->second).data; + } + + public: + int numquadpoints_phi; + int numquadpoints_theta; + Array quadpoints_phi; + Array quadpoints_theta; + double net_poloidal_current_amperes; + double net_toroidal_current_amperes; + + public: + + CurrentPotential( + vector _quadpoints_phi, vector _quadpoints_theta, + double _net_poloidal_current_amperes, double _net_toroidal_current_amperes) + { + numquadpoints_phi = _quadpoints_phi.size(); + numquadpoints_theta = _quadpoints_theta.size(); + net_poloidal_current_amperes = _net_poloidal_current_amperes; + net_toroidal_current_amperes = _net_toroidal_current_amperes; + + quadpoints_phi = xt::zeros({numquadpoints_phi}); + for (int i = 0; i < numquadpoints_phi; ++i) { + quadpoints_phi[i] = _quadpoints_phi[i]; + } + quadpoints_theta = xt::zeros({numquadpoints_theta}); + for (int i = 0; i < numquadpoints_theta; ++i) { + quadpoints_theta[i] = _quadpoints_theta[i]; + } + } + + void invalidate_cache() { + for (auto it = cache.begin(); it != cache.end(); ++it) { + (it->second).status = false; + } + } + + virtual void set_dofs(const vector& _dofs) { + this->set_dofs_impl(_dofs); + this->invalidate_cache(); + } + + void K_impl_helper(Array& data, Array& dg1, Array& dg2, Array& normal); + void K_GI_impl_helper(Array& data, Array& dg1, Array& dg2, Array& normal); + void K_by_dcoeff_impl_helper(Array& data, Array& dg1, Array& dg2, Array& normal); + void K_matrix_impl_helper(Array& data, Array& dg1, Array& dg2, Array& normal); + void K_rhs_impl_helper(Array& data, Array& dg1, Array& dg2, Array& normal); + + virtual int num_dofs() { throw logic_error("num_dofs was not implemented"); }; + virtual void set_dofs_impl(const vector& _dofs) { throw logic_error("set_dofs_impl was not implemented"); }; + virtual vector get_dofs() { throw logic_error("get_dofs was not implemented"); }; + + virtual void Phi_impl(Array& data, Array& quadpoints_phi, Array& quadpoints_theta) { throw logic_error("Phi_impl was not implemented"); }; + virtual void Phidash1_impl(Array& data) { throw logic_error("Phidash1_impl was not implemented"); }; + virtual void Phidash2_impl(Array& data) { throw logic_error("Phidash2_impl was not implemented"); }; + virtual void dPhidash1_by_dcoeff_impl(Array& data) { throw logic_error("dPhidash1_by_dcoeff_impl was not implemented"); }; + virtual void dPhidash2_by_dcoeff_impl(Array& data) { throw logic_error("dPhidash2_by_dcoeff_impl was not implemented"); }; + + Array& Phi() { + return check_the_cache("Phi", {numquadpoints_phi, numquadpoints_theta}, [this](Array& A) { return Phi_impl(A, this->quadpoints_phi, this->quadpoints_theta);}); + } + Array& Phidash1() { + return check_the_cache("Phidash1", {numquadpoints_phi, numquadpoints_theta}, [this](Array& A) { return Phidash1_impl(A);}); + } + Array& Phidash2() { + return check_the_cache("Phidash2", {numquadpoints_phi, numquadpoints_theta}, [this](Array& A) { return Phidash2_impl(A);}); + } + Array& dPhidash1_by_dcoeff() { + return check_the_persistent_cache("dPhidash1_by_dcoeff", {numquadpoints_phi, numquadpoints_theta,num_dofs()}, [this](Array& A) { return dPhidash1_by_dcoeff_impl(A);}); + } + Array& dPhidash2_by_dcoeff() { + return check_the_persistent_cache("dPhidash2_by_dcoeff", {numquadpoints_phi, numquadpoints_theta,num_dofs()}, [this](Array& A) { return dPhidash2_by_dcoeff_impl(A);}); + } + + virtual ~CurrentPotential() = default; +}; diff --git a/src/simsoptpp/currentpotentialfourier.cpp b/src/simsoptpp/currentpotentialfourier.cpp new file mode 100644 index 000000000..5e3a77e4d --- /dev/null +++ b/src/simsoptpp/currentpotentialfourier.cpp @@ -0,0 +1,345 @@ +#include "currentpotentialfourier.h" +#include "simdhelpers.h" + +#define ANGLE_RECOMPUTE 5 + +#if defined(USE_XSIMD) + +// template class Surface, class Array> +template +void CurrentPotentialFourier::Phi_impl(Array& data, Array& quadpoints_phi, Array& quadpoints_theta) { + int numquadpoints_phi = quadpoints_phi.size(); + int numquadpoints_theta = quadpoints_theta.size(); + constexpr int simd_size = xsimd::simd_type::size; +#pragma omp parallel for + for (int k1 = 0; k1 < numquadpoints_phi; ++k1) { + double phi = 2*M_PI*quadpoints_phi[k1]; + for(int k2 = 0; k2 < numquadpoints_theta; k2 += simd_size) { + simd_t theta; + for (int l = 0; l < simd_size; ++l) { + if(k2 + l >= numquadpoints_theta) + break; + theta[l] = 2*M_PI * quadpoints_theta[k2+l]; + } + simd_t Phi(0.); + double sin_nfpphi = sin(-nfp*phi); + double cos_nfpphi = cos(-nfp*phi); + for (int m = 0; m <= mpol; ++m) { + simd_t sinterm, costerm; + for (int i = 0; i < 2*ntor+1; ++i) { + int n = i - ntor; + // recompute the angle from scratch every so often, to + // avoid accumulating floating point error + if(i % ANGLE_RECOMPUTE == 0) + xsimd::sincos(m*theta-n*nfp*phi, sinterm, costerm); + if (! (m == 0 && n <= 0)) { + Phi += phis(m, i) * sinterm; + if(!stellsym) { + Phi += phic(m, i) * costerm; + } + } + if(i % ANGLE_RECOMPUTE != ANGLE_RECOMPUTE - 1){ + simd_t sinterm_old = sinterm; + simd_t costerm_old = costerm; + sinterm = cos_nfpphi * sinterm_old + costerm_old * sin_nfpphi; + costerm = costerm_old * cos_nfpphi - sinterm_old * sin_nfpphi; + } + } + } + for (int l = 0; l < simd_size; ++l) { + if(k2 + l >= numquadpoints_theta) + break; + data(k1, k2+l) = Phi[l]; + } + } + } +} + +#else + +template +void CurrentPotentialFourier::Phi_impl(Array& data, Array& quadpoints_phi, Array& quadpoints_theta) { + int numquadpoints_phi = quadpoints_phi.size(); + int numquadpoints_theta = quadpoints_theta.size(); + constexpr int simd_size = 1; +#pragma omp parallel for + for (int k1 = 0; k1 < numquadpoints_phi; ++k1) { + double phi = 2*M_PI*quadpoints_phi[k1]; + for(int k2 = 0; k2 < numquadpoints_theta; k2 += simd_size) { + double theta = 2*M_PI * quadpoints_theta[k2]; + double Phi = 0.; + double sin_nfpphi = sin(-nfp*phi); + double cos_nfpphi = cos(-nfp*phi); + for (int m = 0; m <= mpol; ++m) { + double sinterm, costerm; + for (int i = 0; i < 2*ntor+1; ++i) { + int n = i - ntor; + if(i % ANGLE_RECOMPUTE == 0) { + sinterm = sin(m*theta - n*nfp*phi); + costerm = cos(m*theta - n*nfp*phi); + } + if (! (m == 0 && n <= 0)) { + Phi += phis(m, i) * sinterm; + if(!stellsym) { + Phi += phic(m, i) * costerm; + } + } + if(i % ANGLE_RECOMPUTE != ANGLE_RECOMPUTE - 1){ + double sinterm_old = sinterm; + double costerm_old = costerm; + sinterm = cos_nfpphi * sinterm_old + costerm_old * sin_nfpphi; + costerm = costerm_old * cos_nfpphi - sinterm_old * sin_nfpphi; + } + } + } + data(k1, k2) = Phi; + } + } +} + +#endif + +#if defined(USE_XSIMD) +// template class T, class Array> +template +void CurrentPotentialFourier::Phidash1_impl(Array& data) { + constexpr int simd_size = xsimd::simd_type::size; +#pragma omp parallel for + for (int k1 = 0; k1 < numquadpoints_phi; ++k1) { + double phi = 2*M_PI*quadpoints_phi[k1]; + for(int k2 = 0; k2 < numquadpoints_theta; k2 += simd_size) { + simd_t theta; + for (int l = 0; l < simd_size; ++l) { + if(k2 + l >= numquadpoints_theta) + break; + theta[l] = 2*M_PI * quadpoints_theta[k2+l]; + } + simd_t Phidash1(0.); + double sin_nfpphi = sin(-nfp*phi); + double cos_nfpphi = cos(-nfp*phi); + for (int m = 0; m <= mpol; ++m) { + simd_t sinterm, costerm; + for (int i = 0; i < 2*ntor+1; ++i) { + int n = i - ntor; + // recompute the angle from scratch every so often, to + // avoid accumulating floating point error + if(i % ANGLE_RECOMPUTE == 0) + xsimd::sincos(m*theta-n*nfp*phi, sinterm, costerm); + if (! (m == 0 && n <= 0)) { + Phidash1 += (-n*nfp) * phis(m, i) * costerm; + if(!stellsym) { + Phidash1 += - (-n*nfp) * phic(m, i) * sinterm; + } + } + if(i % ANGLE_RECOMPUTE != ANGLE_RECOMPUTE - 1){ + simd_t sinterm_old = sinterm; + simd_t costerm_old = costerm; + sinterm = cos_nfpphi * sinterm_old + costerm_old * sin_nfpphi; + costerm = costerm_old * cos_nfpphi - sinterm_old * sin_nfpphi; + } + } + } + for (int l = 0; l < simd_size; ++l) { + if(k2 + l >= numquadpoints_theta) + break; + data(k1, k2+l) = Phidash1[l]*2*M_PI; + } + } + } +} + +#else + +template +void CurrentPotentialFourier::Phidash1_impl(Array& data) { + constexpr int simd_size = 1; +#pragma omp parallel for + for (int k1 = 0; k1 < numquadpoints_phi; ++k1) { + double phi = 2*M_PI*quadpoints_phi[k1]; + for(int k2 = 0; k2 < numquadpoints_theta; k2 += simd_size) { + double theta = 2*M_PI * quadpoints_theta[k2]; + double Phidash1 = 0.; + double sin_nfpphi = sin(-nfp*phi); + double cos_nfpphi = cos(-nfp*phi); + for (int m = 0; m <= mpol; ++m) { + double sinterm, costerm; + for (int i = 0; i < 2*ntor+1; ++i) { + int n = i - ntor; + if(i % ANGLE_RECOMPUTE == 0) { + sinterm = sin(m*theta - n*nfp*phi); + costerm = cos(m*theta - n*nfp*phi); + } + if (! (m == 0 && n <= 0)) { + Phidash1 += (-n*nfp) * phis(m, i) * costerm; + if(!stellsym) { + Phidash1 += - (-n*nfp) * phic(m, i) * sinterm; + } + } + if(i % ANGLE_RECOMPUTE != ANGLE_RECOMPUTE - 1){ + double sinterm_old = sinterm; + double costerm_old = costerm; + sinterm = cos_nfpphi * sinterm_old + costerm_old * sin_nfpphi; + costerm = costerm_old * cos_nfpphi - sinterm_old * sin_nfpphi; + } + } + } + data(k1, k2) = Phidash1 * 2*M_PI; + } + } +} + +#endif + +#if defined(USE_XSIMD) +// template class Surface, class Array> +// template class T> +template +void CurrentPotentialFourier::Phidash2_impl(Array& data) { + constexpr int simd_size = xsimd::simd_type::size; +#pragma omp parallel for + for (int k1 = 0; k1 < numquadpoints_phi; ++k1) { + double phi = 2*M_PI*quadpoints_phi[k1]; + for(int k2 = 0; k2 < numquadpoints_theta; k2 += simd_size) { + simd_t theta; + for (int l = 0; l < simd_size; ++l) { + if(k2 + l >= numquadpoints_theta) + break; + theta[l] = 2*M_PI * quadpoints_theta[k2+l]; + } + simd_t Phidash2(0.); + double sin_nfpphi = sin(-nfp*phi); + double cos_nfpphi = cos(-nfp*phi); + for (int m = 0; m <= mpol; ++m) { + simd_t sinterm, costerm; + for (int i = 0; i < 2*ntor+1; ++i) { + int n = i - ntor; + // recompute the angle from scratch every so often, to + // avoid accumulating floating point error + if(i % ANGLE_RECOMPUTE == 0) + xsimd::sincos(m*theta-n*nfp*phi, sinterm, costerm); + if (! (m == 0 && n <= 0)) { + Phidash2 += m * phis(m, i) * costerm; + if(!stellsym) { + Phidash2 += - m * phic(m, i) * sinterm; + } + } + if(i % ANGLE_RECOMPUTE != ANGLE_RECOMPUTE - 1){ + simd_t sinterm_old = sinterm; + simd_t costerm_old = costerm; + sinterm = cos_nfpphi * sinterm_old + costerm_old * sin_nfpphi; + costerm = costerm_old * cos_nfpphi - sinterm_old * sin_nfpphi; + } + } + } + for (int l = 0; l < simd_size; ++l) { + if(k2 + l >= numquadpoints_theta) + break; + data(k1, k2+l) = Phidash2[l] * 2*M_PI; + } + } + } +} + +#else + +template +void CurrentPotentialFourier::Phidash2_impl(Array& data) { + constexpr int simd_size = 1; +#pragma omp parallel for + for (int k1 = 0; k1 < numquadpoints_phi; ++k1) { + double phi = 2*M_PI*quadpoints_phi[k1]; + for(int k2 = 0; k2 < numquadpoints_theta; k2 += simd_size) { + double theta = 2*M_PI * quadpoints_theta[k2]; + double Phidash2 = 0.; + double sin_nfpphi = sin(-nfp*phi); + double cos_nfpphi = cos(-nfp*phi); + for (int m = 0; m <= mpol; ++m) { + double sinterm, costerm; + for (int i = 0; i < 2*ntor+1; ++i) { + int n = i - ntor; + if(i % ANGLE_RECOMPUTE == 0) { + sinterm = sin(m*theta - n*nfp*phi); + costerm = cos(m*theta - n*nfp*phi); + } + if (! (m == 0 && n <= 0)) { + Phidash2 += m * phis(m, i) * costerm; + if(!stellsym) { + Phidash2 += - m * phic(m, i) * sinterm; + } + } + if(i % ANGLE_RECOMPUTE != ANGLE_RECOMPUTE - 1){ + double sinterm_old = sinterm; + double costerm_old = costerm; + sinterm = cos_nfpphi * sinterm_old + costerm_old * sin_nfpphi; + costerm = costerm_old * cos_nfpphi - sinterm_old * sin_nfpphi; + } + } + } + data(k1, k2) = Phidash2 * 2*M_PI; + } + } +} + +#endif + +template +void CurrentPotentialFourier::dPhidash2_by_dcoeff_impl(Array& data) { +#pragma omp parallel for + for (int k1 = 0; k1 < numquadpoints_phi; ++k1) { + double phi = 2*M_PI*quadpoints_phi[k1]; + for (int k2 = 0; k2 < numquadpoints_theta; ++k2) { + double theta = 2*M_PI*quadpoints_theta[k2]; + int counter = 0; + for (int m = 0; m <= mpol; ++m) { + for (int n = -ntor; n <= ntor; ++n) { + if(m==0 && n<=0) continue; + data(k1, k2, counter) = 2*M_PI*m * cos(m*theta-n*nfp*phi); + counter++; + } + } + if (!stellsym) { + for (int m = 0; m <= mpol; ++m) { + for (int n = -ntor; n <= ntor; ++n) { + if(m==0 && n<=0) continue; + data(k1, k2, counter) = 2*M_PI*(-m) * sin(m*theta-n*nfp*phi); + counter++; + } + } + } + } + } +} + +template +void CurrentPotentialFourier::dPhidash1_by_dcoeff_impl(Array& data) { +#pragma omp parallel for + for (int k1 = 0; k1 < numquadpoints_phi; ++k1) { + double phi = 2*M_PI*quadpoints_phi[k1]; + for (int k2 = 0; k2 < numquadpoints_theta; ++k2) { + double theta = 2*M_PI*quadpoints_theta[k2]; + int counter = 0; + for (int m = 0; m <= mpol; ++m) { + for (int n = -ntor; n <= ntor; ++n) { + if(m==0 && n<=0) continue; + data(k1, k2, counter) = 2*M_PI*(-n*nfp) * cos(m*theta-n*nfp*phi); + counter++; + } + } + if (!stellsym) { + for (int m = 0; m <= mpol; ++m) { + for (int n = -ntor; n <= ntor; ++n) { + if(m==0 && n<=0) continue; + data(k1, k2, counter) = 2*M_PI*(n*nfp) * sin(m*theta-n*nfp*phi); + counter++; + } + } + } + } + } +} + +#include "xtensor-python/pyarray.hpp" +typedef xt::pyarray Array; +template class CurrentPotentialFourier; +// template class T> +// template class Surface, class Array> diff --git a/src/simsoptpp/currentpotentialfourier.h b/src/simsoptpp/currentpotentialfourier.h new file mode 100644 index 000000000..b91bb1a08 --- /dev/null +++ b/src/simsoptpp/currentpotentialfourier.h @@ -0,0 +1,80 @@ +#pragma once + +#include "currentpotential.h" + +template +class CurrentPotentialFourier : public CurrentPotential { + + public: + using CurrentPotential::quadpoints_phi; + using CurrentPotential::quadpoints_theta; + using CurrentPotential::numquadpoints_phi; + using CurrentPotential::numquadpoints_theta; + Array phic; + Array phis; + int nfp; + int mpol; + int ntor; + bool stellsym; + using CurrentPotential::net_poloidal_current_amperes; + using CurrentPotential::net_toroidal_current_amperes; + + CurrentPotentialFourier( + int _mpol, int _ntor, int _nfp, bool _stellsym, + vector _quadpoints_phi, vector _quadpoints_theta, + double net_poloidal_current_amperes, double net_toroidal_current_amperes) + : CurrentPotential(_quadpoints_phi, _quadpoints_theta, net_poloidal_current_amperes, net_toroidal_current_amperes), mpol(_mpol), ntor(_ntor), nfp(_nfp), stellsym(_stellsym) { + this->allocate(); + } + + void allocate() { + phic = xt::zeros({mpol+1, 2*ntor+1}); + phis = xt::zeros({mpol+1, 2*ntor+1}); + } + + int num_dofs() override { + if(stellsym) + // does not include a dof for phis(0, 0) + return mpol*(2*ntor + 1) + (ntor + 1) - 1; + else + // does not include a dof for phic(0, 0) or phis(0, 0) + return 2*(mpol*(2*ntor + 1) + (ntor + 1) - 1); + } + + void set_dofs_impl(const vector& dofs) override { + int shift = (mpol+1)*(2*ntor+1); + int counter = 0; + if(stellsym) { + for (int i = ntor+1; i < shift; ++i) + phis.data()[i] = dofs[counter++]; + } else { + for (int i = ntor+1; i < shift; ++i) + phis.data()[i] = dofs[counter++]; + for (int i = ntor+1; i < shift; ++i) + phic.data()[i] = dofs[counter++]; + } + } + + vector get_dofs() override { + auto res = vector(num_dofs(), 0.); + int shift = (mpol+1)*(2*ntor+1); + int counter = 0; + if(stellsym) { + for (int i = ntor+1; i < shift; ++i) + res[counter++] = phis.data()[i]; + } else { + for (int i = ntor+1; i < shift; ++i) + res[counter++] = phis.data()[i]; + for (int i = ntor+1; i < shift; ++i) + res[counter++] = phic.data()[i]; + } + return res; + } + + void Phi_impl(Array& data, Array& quadpoints_phi, Array& quadpoints_theta) override; + void Phidash1_impl(Array& data) override; + void Phidash2_impl(Array& data) override; + void dPhidash1_by_dcoeff_impl(Array& data) override; + void dPhidash2_by_dcoeff_impl(Array& data) override; + +}; diff --git a/src/simsoptpp/pycurrentpotential.h b/src/simsoptpp/pycurrentpotential.h new file mode 100644 index 000000000..0af9987f9 --- /dev/null +++ b/src/simsoptpp/pycurrentpotential.h @@ -0,0 +1,34 @@ +#pragma once + +#include "currentpotential.h" +#include "xtensor-python/pyarray.hpp" +typedef xt::pyarray PyArray; + +typedef CurrentPotential PyCurrentPotential; + +template class PyCurrentPotentialTrampoline : public CurrentPotentialBase { + public: + using CurrentPotentialBase::CurrentPotentialBase; + + virtual int num_dofs() override { + PYBIND11_OVERLOAD(int, CurrentPotentialBase, num_dofs); + } + virtual void set_dofs_impl(const vector& _dofs) override { + PYBIND11_OVERLOAD(void, CurrentPotentialBase, set_dofs_impl, _dofs); + } + virtual void set_dofs(const vector& _dofs) override { + PYBIND11_OVERLOAD(void, CurrentPotentialBase, set_dofs, _dofs); + } + virtual vector get_dofs() override { + PYBIND11_OVERLOAD(vector, CurrentPotentialBase, get_dofs); + } + virtual void Phi_impl(PyArray& data, PyArray& quadpoints_phi, PyArray& quadpoints_theta) override { + PYBIND11_OVERLOAD(void, CurrentPotentialBase, Phi_impl, data, quadpoints_phi, quadpoints_theta); + } + virtual void Phidash1_impl(PyArray& data) override { + PYBIND11_OVERLOAD(void, CurrentPotentialBase, Phidash1_impl, data); + } + virtual void Phidash2_impl(PyArray& data) override { + PYBIND11_OVERLOAD(void, CurrentPotentialBase, Phidash2_impl, data); + } +}; diff --git a/src/simsoptpp/python.cpp b/src/simsoptpp/python.cpp index c3b2c23dd..13dff0f43 100644 --- a/src/simsoptpp/python.cpp +++ b/src/simsoptpp/python.cpp @@ -22,6 +22,7 @@ typedef xt::pytensor PyTensor; #include "wireframe_optimization.h" #include "reiman.h" #include "simdhelpers.h" +#include "winding_surface.h" #include "boozerresidual_py.h" namespace py = pybind11; @@ -35,8 +36,7 @@ void init_magneticfields(py::module_ &); void init_boozermagneticfields(py::module_ &); void init_tracing(py::module_ &); void init_distance(py::module_ &); - - +void init_currentpotential(py::module_ &); PYBIND11_MODULE(simsoptpp, m) { xt::import_numpy(); @@ -47,6 +47,7 @@ PYBIND11_MODULE(simsoptpp, m) { init_boozermagneticfields(m); init_tracing(m); init_distance(m); + init_currentpotential(m); #if defined(USE_XSIMD) m.attr("using_xsimd") = true; @@ -60,6 +61,7 @@ PYBIND11_MODULE(simsoptpp, m) { m.def("biot_savart_vjp_graph", &biot_savart_vjp_graph); m.def("biot_savart_vector_potential_vjp_graph", &biot_savart_vector_potential_vjp_graph); + // Functions below are implemented for permanent magnet optimization m.def("dipole_field_B" , &dipole_field_B); m.def("dipole_field_A" , &dipole_field_A); @@ -88,6 +90,15 @@ PYBIND11_MODULE(simsoptpp, m) { m.def("DommaschkB" , &DommaschkB); m.def("DommaschkdB", &DommaschkdB); + m.def("WindingSurfaceBn_REGCOIL", &WindingSurfaceBn_REGCOIL); + m.def("WindingSurfaceB", &WindingSurfaceB); + m.def("WindingSurfacedB", &WindingSurfacedB); + m.def("WindingSurfaceA", &WindingSurfaceA); + m.def("WindingSurfacedA", &WindingSurfacedA); + m.def("winding_surface_field_Bn", &winding_surface_field_Bn); + m.def("winding_surface_field_K2_matrices", &winding_surface_field_K2_matrices); + m.def("winding_surface_field_Bn_GI", &winding_surface_field_Bn_GI); + m.def("integral_BdotN", &integral_BdotN); m.def("ReimanB" , &ReimanB); diff --git a/src/simsoptpp/python_currentpotential.cpp b/src/simsoptpp/python_currentpotential.cpp new file mode 100644 index 000000000..d4c79e94a --- /dev/null +++ b/src/simsoptpp/python_currentpotential.cpp @@ -0,0 +1,75 @@ +#include "pybind11/pybind11.h" +#include "pybind11/stl.h" +#include "pybind11/functional.h" +#include "xtensor-python/pyarray.hpp" +typedef xt::pyarray PyArray; +using std::shared_ptr; +using std::vector; + +#include "currentpotential.h" +#include "pycurrentpotential.h" +#include "pysurface.h" +#include "surface.h" +#include "currentpotentialfourier.h" +typedef CurrentPotentialFourier PyCurrentPotentialFourier; + +template class PyCurrentPotentialFourierTrampoline : public PyCurrentPotentialTrampoline { + public: + using PyCurrentPotentialTrampoline::PyCurrentPotentialTrampoline; + using PyCurrentPotentialFourierBase::mpol; + using PyCurrentPotentialFourierBase::ntor; + using PyCurrentPotentialFourierBase::nfp; + using PyCurrentPotentialFourierBase::stellsym; + + int num_dofs() override { + return PyCurrentPotentialFourierBase::num_dofs(); + } + + void set_dofs_impl(const vector& _dofs) override { + PyCurrentPotentialFourierBase::set_dofs_impl(_dofs); + } + + vector get_dofs() override { + return PyCurrentPotentialFourierBase::get_dofs(); + } + + void Phi_impl(PyArray& data, PyArray& quadpoints_phi, PyArray& quadpoints_theta) override { + PyCurrentPotentialFourierBase::Phi_impl(data, quadpoints_phi, quadpoints_theta); + } +}; + +template void register_common_currentpotential_methods(S &s) { + s.def("Phi", pybind11::overload_cast<>(&T::Phi)) + .def("K_impl_helper", &T::K_impl_helper) + .def("K_matrix_impl_helper", &T::K_matrix_impl_helper) + .def("K_rhs_impl_helper", &T::K_rhs_impl_helper) + .def("set_dofs_impl", &T::set_dofs_impl) + .def("Phidash1", pybind11::overload_cast<>(&T::Phidash1)) + .def("Phidash2", pybind11::overload_cast<>(&T::Phidash2)) + .def("Phidash1_impl", &T::Phidash1_impl) + .def("Phidash2_impl", &T::Phidash2_impl) + .def("invalidate_cache", &T::invalidate_cache) + .def("set_dofs", &T::set_dofs) + .def("get_dofs", &T::get_dofs) + .def_readonly("quadpoints_phi", &T::quadpoints_phi) + .def_readonly("quadpoints_theta", &T::quadpoints_theta); +} + +void init_currentpotential(pybind11::module_ &m){ + auto pycurrentpotential = pybind11::class_, PyCurrentPotentialTrampoline>(m, "CurrentPotential") + .def(pybind11::init,vector, double, double>()); + register_common_currentpotential_methods(pycurrentpotential); + + auto pycurrentpotentialfourier = pybind11::class_, PyCurrentPotentialFourierTrampoline>(m, "CurrentPotentialFourier") + .def(pybind11::init, vector, double, double>()) + .def_readwrite("phic", &PyCurrentPotentialFourier::phic) + .def_readwrite("phis", &PyCurrentPotentialFourier::phis) + .def_readwrite("mpol", &PyCurrentPotentialFourier::mpol) + .def_readwrite("ntor", &PyCurrentPotentialFourier::ntor) + .def_readwrite("nfp", &PyCurrentPotentialFourier::nfp) + .def_readwrite("stellsym", &PyCurrentPotentialFourier::stellsym) + .def_readwrite("net_poloidal_current_amperes", &PyCurrentPotentialFourier::net_poloidal_current_amperes) + .def_readwrite("net_toroidal_current_amperes", &PyCurrentPotentialFourier::net_toroidal_current_amperes) + .def("allocate", &PyCurrentPotentialFourier::allocate); + register_common_currentpotential_methods(pycurrentpotentialfourier); +} diff --git a/src/simsoptpp/winding_surface.cpp b/src/simsoptpp/winding_surface.cpp new file mode 100644 index 000000000..b3074d031 --- /dev/null +++ b/src/simsoptpp/winding_surface.cpp @@ -0,0 +1,721 @@ +#include "winding_surface.h" +#include "simdhelpers.h" +#include "vec3dsimd.h" +#include + +// Compute Bnormal using equation A8 in the REGCOIL paper. This is implemented +// because computing Bnormal using the the normal, discretized BiotSavart +// will give a slightly different answer if the resolution is relatively low. +// This is just because the integrals being discretized are a bit different! +Array WindingSurfaceBn_REGCOIL(Array& points, Array& ws_points, Array& ws_normal, Array& current_potential, Array& plasma_normal) +{ + int num_points = points.shape(0); + int num_ws_points = ws_points.shape(0); + Array Bn = xt::zeros({points.shape(0)}); + double fak = 1e-7; // mu0 divided by 4 * pi factor + + #pragma omp parallel for schedule(static) + for(int i = 0; i < num_points; ++i) { + double x = points(i, 0); + double y = points(i, 1); + double z = points(i, 2); + double nx = plasma_normal(i, 0); + double ny = plasma_normal(i, 1); + double nz = plasma_normal(i, 2); + double nmag = sqrt(nx * nx + ny * ny + nz * nz); + double Bi_normal = 0.0; + + // Sum contributions from all the winding surface points + // i.e. do the surface integral over the winding surface + for (int j = 0; j < num_ws_points; ++j) { + double xx = ws_points(j, 0); + double yy = ws_points(j, 1); + double zz = ws_points(j, 2); + double nxx = ws_normal(j, 0); + double nyy = ws_normal(j, 1); + double nzz = ws_normal(j, 2); + double phi = current_potential(j); + double rx = x - xx; + double ry = y - yy; + double rz = z - zz; + double rmag_inv = 1.0 / sqrt(rx * rx + ry * ry + rz * rz); + double rmag_inv_3 = rmag_inv * (rmag_inv * rmag_inv); + double rmag_inv_5 = rmag_inv * rmag_inv * rmag_inv_3; + double NdotNprime = nx * nxx + ny * nyy + nz * nzz; + double RdotN = rx * nx + ry * ny + rz * nz; + double RdotNprime = rx * nxx + ry * nyy + rz * nzz; + double integrand = phi * (NdotNprime * rmag_inv_3 - 3.0 * RdotN * RdotNprime * rmag_inv_5); + Bi_normal += integrand; + } + Bn(i) = fak * Bi_normal / nmag; + } + return Bn; +} + +#if defined(USE_XSIMD) +Array WindingSurfaceB(Array& points, Array& ws_points, Array& ws_normal, Array& K) +{ + // warning: row_major checks below do NOT throw an error correctly on a compute node on Cori + if(points.layout() != xt::layout_type::row_major) + throw std::runtime_error("points needs to be in row-major storage order"); + if(ws_points.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface points needs to be in row-major storage order"); + if(ws_normal.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface normal vector needs to be in row-major storage order"); + if(K.layout() != xt::layout_type::row_major) + throw std::runtime_error("surface_current needs to be in row-major storage order"); + + int num_points = points.shape(0); + int num_ws_points = ws_points.shape(0); + constexpr int simd_size = xsimd::simd_type::size; + Array B = xt::zeros({points.shape(0), points.shape(1)}); + + // initialize pointer to the beginning of ws_points + double* ws_points_ptr = &(ws_points(0, 0)); + double* ws_normal_ptr = &(ws_normal(0, 0)); + double* K_ptr = &(K(0, 0)); + double fak = 1e-7; // mu0 divided by 4 * pi factor + + // Loop through the evaluation points by chunks of simd_size + #pragma omp parallel for schedule(static) + for(int i = 0; i < num_points; i += simd_size) { + auto point_i = Vec3dSimd(); + auto B_i = Vec3dSimd(); + + // check that i + k isn't bigger than num_points + int klimit = std::min(simd_size, num_points - i); + for(int k = 0; k < klimit; k++){ + for (int d = 0; d < 3; ++d) { + point_i[d][k] = points(i + k, d); + } + } + // Sum contributions from all the winding surface points + // i.e. do the surface integral over the winding surface + for (int j = 0; j < num_ws_points; ++j) { + Vec3dSimd r_j = Vec3dSimd(ws_points_ptr[3 * j + 0], ws_points_ptr[3 * j + 1], ws_points_ptr[3 * j + 2]); + Vec3dSimd n_j = Vec3dSimd(ws_normal_ptr[3 * j + 0], ws_normal_ptr[3 * j + 1], ws_normal_ptr[3 * j + 2]); + Vec3dSimd K_j = Vec3dSimd(K_ptr[3 * j + 0], K_ptr[3 * j + 1], K_ptr[3 * j + 2]); + Vec3dSimd r = point_i - r_j; + simd_t rmag_2 = normsq(r); + simd_t rmag_inv = rsqrt(rmag_2); + simd_t rmag_inv_3 = rmag_inv * (rmag_inv * rmag_inv); + simd_t nmag = sqrt(normsq(n_j)); + Vec3dSimd Kcrossr = cross(K_j, r); + B_i.x += nmag * Kcrossr.x * rmag_inv_3; + B_i.y += nmag * Kcrossr.y * rmag_inv_3; + B_i.z += nmag * Kcrossr.z * rmag_inv_3; + } + for(int k = 0; k < klimit; k++){ + B(i + k, 0) = fak * B_i.x[k]; + B(i + k, 1) = fak * B_i.y[k]; + B(i + k, 2) = fak * B_i.z[k]; + } + } + return B; +} + +#else + +Array WindingSurfaceB(Array& points, Array& ws_points, Array& ws_normal, Array& K) +{ + if(points.layout() != xt::layout_type::row_major) + throw std::runtime_error("points needs to be in row-major storage order"); + if(ws_points.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface points needs to be in row-major storage order"); + if(ws_normal.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface normal vector needs to be in row-major storage order"); + if(K.layout() != xt::layout_type::row_major) + throw std::runtime_error("surface_current needs to be in row-major storage order"); + + int num_points = points.shape(0); + int num_ws_points = ws_points.shape(0); + Array B = xt::zeros({points.shape(0), points.shape(1)}); + double fak = 1e-7; + + #pragma omp parallel for schedule(static) + for(int i = 0; i < num_points; ++i) { + double Bx = 0., By = 0., Bz = 0.; + double px = points(i, 0), py = points(i, 1), pz = points(i, 2); + for (int j = 0; j < num_ws_points; ++j) { + double xx = ws_points(j, 0), yy = ws_points(j, 1), zz = ws_points(j, 2); + double nx = ws_normal(j, 0), ny = ws_normal(j, 1), nz = ws_normal(j, 2); + double Kx = K(j, 0), Ky = K(j, 1), Kz = K(j, 2); + double rx = px - xx, ry = py - yy, rz = pz - zz; + double rmag2 = rx*rx + ry*ry + rz*rz; + double rmag_inv = 1.0 / std::sqrt(rmag2); + double rmag_inv_3 = rmag_inv * rmag_inv * rmag_inv; + double nmag = std::sqrt(nx*nx + ny*ny + nz*nz); + double Kcrx = Ky*rz - Kz*ry, Kcry = Kz*rx - Kx*rz, Kcrz = Kx*ry - Ky*rx; + Bx += nmag * Kcrx * rmag_inv_3; + By += nmag * Kcry * rmag_inv_3; + Bz += nmag * Kcrz * rmag_inv_3; + } + B(i, 0) = fak * Bx; + B(i, 1) = fak * By; + B(i, 2) = fak * Bz; + } + return B; +} + +#endif + +#if defined(USE_XSIMD) +Array WindingSurfacedB(Array& points, Array& ws_points, Array& ws_normal, Array& K) +{ + // warning: row_major checks below do NOT throw an error correctly on a compute node on Cori + if(points.layout() != xt::layout_type::row_major) + throw std::runtime_error("points needs to be in row-major storage order"); + if(ws_points.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface points needs to be in row-major storage order"); + if(ws_normal.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface normal vector needs to be in row-major storage order"); + if(K.layout() != xt::layout_type::row_major) + throw std::runtime_error("surface_current needs to be in row-major storage order"); + + int num_points = points.shape(0); + int num_ws_points = ws_points.shape(0); + constexpr int simd_size = xsimd::simd_type::size; + Array dB = xt::zeros({points.shape(0), points.shape(1), points.shape(1)}); + + // initialize pointer to the beginning of ws_points + double* ws_points_ptr = &(ws_points(0, 0)); + double* ws_normal_ptr = &(ws_normal(0, 0)); + double* K_ptr = &(K(0, 0)); + double fak = 1e-7; // mu0 divided by 4 * pi factor + + // Loop through the evaluation points by chunks of simd_size + #pragma omp parallel for schedule(static) + for(int i = 0; i < num_points; i += simd_size) { + auto point_i = Vec3dSimd(); + auto dB_i1 = Vec3dSimd(); + auto dB_i2 = Vec3dSimd(); + auto dB_i3 = Vec3dSimd(); + + // check that i + k isn't bigger than num_points + int klimit = std::min(simd_size, num_points - i); + for(int k = 0; k < klimit; k++){ + for (int d = 0; d < 3; ++d) { + point_i[d][k] = points(i + k, d); + } + } + // Sum contributions from all the winding surface points + // i.e. do the surface integral over the winding surface + for (int j = 0; j < num_ws_points; ++j) { + Vec3dSimd r_j = Vec3dSimd(ws_points_ptr[3 * j + 0], ws_points_ptr[3 * j + 1], ws_points_ptr[3 * j + 2]); + Vec3dSimd n_j = Vec3dSimd(ws_normal_ptr[3 * j + 0], ws_normal_ptr[3 * j + 1], ws_normal_ptr[3 * j + 2]); + Vec3dSimd K_j = Vec3dSimd(K_ptr[3 * j + 0], K_ptr[3 * j + 1], K_ptr[3 * j + 2]); + Vec3dSimd r = point_i - r_j; + simd_t rmag_2 = normsq(r); + simd_t rmag_inv = rsqrt(rmag_2); + simd_t rmag_inv_3 = rmag_inv * (rmag_inv * rmag_inv); + simd_t rmag_inv_5 = rmag_inv_3 * rmag_inv * rmag_inv; + simd_t nmag = sqrt(normsq(n_j)); + Vec3dSimd Kcrossr = cross(K_j, r); + Vec3dSimd ex = Vec3dSimd(1, 0, 0); + Vec3dSimd ey = Vec3dSimd(0, 1, 0); + Vec3dSimd ez = Vec3dSimd(0, 0, 1); + Vec3dSimd Kcrossex = cross(K_j, ex); + Vec3dSimd Kcrossey = cross(K_j, ey); + Vec3dSimd Kcrossez = cross(K_j, ez); + dB_i1.x += nmag * (Kcrossex.x * rmag_inv_3 - 3.0 * Kcrossr.x * rmag_inv_5 * r.x); + dB_i1.y += nmag * (Kcrossex.y * rmag_inv_3 - 3.0 * Kcrossr.y * rmag_inv_5 * r.x); + dB_i1.z += nmag * (Kcrossex.z * rmag_inv_3 - 3.0 * Kcrossr.z * rmag_inv_5 * r.x); + dB_i2.x += nmag * (Kcrossey.x * rmag_inv_3 - 3.0 * Kcrossr.x * rmag_inv_5 * r.y); + dB_i2.y += nmag * (Kcrossey.y * rmag_inv_3 - 3.0 * Kcrossr.y * rmag_inv_5 * r.y); + dB_i2.z += nmag * (Kcrossey.z * rmag_inv_3 - 3.0 * Kcrossr.z * rmag_inv_5 * r.y); + dB_i3.x += nmag * (Kcrossez.x * rmag_inv_3 - 3.0 * Kcrossr.x * rmag_inv_5 * r.z); + dB_i3.y += nmag * (Kcrossez.y * rmag_inv_3 - 3.0 * Kcrossr.y * rmag_inv_5 * r.z); + dB_i3.z += nmag * (Kcrossez.z * rmag_inv_3 - 3.0 * Kcrossr.z * rmag_inv_5 * r.z); + } + for(int k = 0; k < klimit; k++){ + dB(i + k, 0, 0) = fak * dB_i1.x[k]; + dB(i + k, 0, 1) = fak * dB_i1.y[k]; + dB(i + k, 0, 2) = fak * dB_i1.z[k]; + dB(i + k, 1, 0) = fak * dB_i2.x[k]; + dB(i + k, 1, 1) = fak * dB_i2.y[k]; + dB(i + k, 1, 2) = fak * dB_i2.z[k]; + dB(i + k, 2, 0) = fak * dB_i3.x[k]; + dB(i + k, 2, 1) = fak * dB_i3.y[k]; + dB(i + k, 2, 2) = fak * dB_i3.z[k]; + } + } + return dB; +} + +#else + +Array WindingSurfacedB(Array& points, Array& ws_points, Array& ws_normal, Array& K) +{ + if(points.layout() != xt::layout_type::row_major) + throw std::runtime_error("points needs to be in row-major storage order"); + if(ws_points.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface points needs to be in row-major storage order"); + if(ws_normal.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface normal vector needs to be in row-major storage order"); + if(K.layout() != xt::layout_type::row_major) + throw std::runtime_error("surface_current needs to be in row-major storage order"); + + int num_points = points.shape(0); + int num_ws_points = ws_points.shape(0); + Array dB = xt::zeros({points.shape(0), points.shape(1), points.shape(1)}); + double fak = 1e-7; + + #pragma omp parallel for schedule(static) + for(int i = 0; i < num_points; ++i) { + double dB11 = 0., dB12 = 0., dB13 = 0.; + double dB21 = 0., dB22 = 0., dB23 = 0.; + double dB31 = 0., dB32 = 0., dB33 = 0.; + double px = points(i, 0), py = points(i, 1), pz = points(i, 2); + for (int j = 0; j < num_ws_points; ++j) { + double xx = ws_points(j, 0), yy = ws_points(j, 1), zz = ws_points(j, 2); + double nx = ws_normal(j, 0), ny = ws_normal(j, 1), nz = ws_normal(j, 2); + double Kx = K(j, 0), Ky = K(j, 1), Kz = K(j, 2); + double rx = px - xx, ry = py - yy, rz = pz - zz; + double rmag2 = rx*rx + ry*ry + rz*rz; + double rmag_inv = 1.0 / std::sqrt(rmag2); + double rmag_inv_3 = rmag_inv * rmag_inv * rmag_inv; + double rmag_inv_5 = rmag_inv_3 * rmag_inv * rmag_inv; + double nmag = std::sqrt(nx*nx + ny*ny + nz*nz); + double Kcrx = Ky*rz - Kz*ry, Kcry = Kz*rx - Kx*rz, Kcrz = Kx*ry - Ky*rx; + double Kcrossex_y = Kz, Kcrossex_z = -Ky; + double Kcrossey_x = -Kz, Kcrossey_z = Kx; + double Kcrossez_x = Ky, Kcrossez_y = -Kx; + dB11 += nmag * (-3.0 * Kcrx * rmag_inv_5 * rx); + dB12 += nmag * (Kcrossex_y * rmag_inv_3 - 3.0 * Kcry * rmag_inv_5 * rx); + dB13 += nmag * (Kcrossex_z * rmag_inv_3 - 3.0 * Kcrz * rmag_inv_5 * rx); + dB21 += nmag * (Kcrossey_x * rmag_inv_3 - 3.0 * Kcrx * rmag_inv_5 * ry); + dB22 += nmag * (-3.0 * Kcry * rmag_inv_5 * ry); + dB23 += nmag * (Kcrossey_z * rmag_inv_3 - 3.0 * Kcrz * rmag_inv_5 * ry); + dB31 += nmag * (Kcrossez_x * rmag_inv_3 - 3.0 * Kcrx * rmag_inv_5 * rz); + dB32 += nmag * (Kcrossez_y * rmag_inv_3 - 3.0 * Kcry * rmag_inv_5 * rz); + dB33 += nmag * (-3.0 * Kcrz * rmag_inv_5 * rz); + } + dB(i, 0, 0) = fak * dB11; + dB(i, 0, 1) = fak * dB12; + dB(i, 0, 2) = fak * dB13; + dB(i, 1, 0) = fak * dB21; + dB(i, 1, 1) = fak * dB22; + dB(i, 1, 2) = fak * dB23; + dB(i, 2, 0) = fak * dB31; + dB(i, 2, 1) = fak * dB32; + dB(i, 2, 2) = fak * dB33; + } + return dB; +} + +#endif + +#if defined(USE_XSIMD) +Array WindingSurfaceA(Array& points, Array& ws_points, Array& ws_normal, Array& K) +{ + // warning: row_major checks below do NOT throw an error correctly on a compute node on Cori + if(points.layout() != xt::layout_type::row_major) + throw std::runtime_error("points needs to be in row-major storage order"); + if(ws_points.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface points needs to be in row-major storage order"); + if(ws_normal.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface normal vector needs to be in row-major storage order"); + if(K.layout() != xt::layout_type::row_major) + throw std::runtime_error("surface_current needs to be in row-major storage order"); + + int num_points = points.shape(0); + int num_ws_points = ws_points.shape(0); + constexpr int simd_size = xsimd::simd_type::size; + Array A = xt::zeros({points.shape(0), points.shape(1)}); + + // initialize pointer to the beginning of ws_points + double* ws_points_ptr = &(ws_points(0, 0)); + double* ws_normal_ptr = &(ws_normal(0, 0)); + double* K_ptr = &(K(0, 0)); + double fak = 1e-7; // mu0 divided by 4 * pi factor + + // Loop through the evaluation points by chunks of simd_size + #pragma omp parallel for schedule(static) + for(int i = 0; i < num_points; i += simd_size) { + auto point_i = Vec3dSimd(); + auto A_i = Vec3dSimd(); + + // check that i + k isn't bigger than num_points + int klimit = std::min(simd_size, num_points - i); + for(int k = 0; k < klimit; k++){ + for (int d = 0; d < 3; ++d) { + point_i[d][k] = points(i + k, d); + } + } + // Sum contributions from all the winding surface points + // i.e. do the surface integral over the winding surface + for (int j = 0; j < num_ws_points; ++j) { + Vec3dSimd r_j = Vec3dSimd(ws_points_ptr[3 * j + 0], ws_points_ptr[3 * j + 1], ws_points_ptr[3 * j + 2]); + Vec3dSimd n_j = Vec3dSimd(ws_normal_ptr[3 * j + 0], ws_normal_ptr[3 * j + 1], ws_normal_ptr[3 * j + 2]); + Vec3dSimd K_j = Vec3dSimd(K_ptr[3 * j + 0], K_ptr[3 * j + 1], K_ptr[3 * j + 2]); + Vec3dSimd r = point_i - r_j; + simd_t rmag_2 = normsq(r); + simd_t rmag_inv = rsqrt(rmag_2); + simd_t nmag = sqrt(normsq(n_j)); + A_i.x += nmag * K_j.x * rmag_inv; + A_i.y += nmag * K_j.y * rmag_inv; + A_i.z += nmag * K_j.z * rmag_inv; + } + for(int k = 0; k < klimit; k++){ + A(i + k, 0) = fak * A_i.x[k]; + A(i + k, 1) = fak * A_i.y[k]; + A(i + k, 2) = fak * A_i.z[k]; + } + } + return A; +} + +#else + +Array WindingSurfaceA(Array& points, Array& ws_points, Array& ws_normal, Array& K) +{ + if(points.layout() != xt::layout_type::row_major) + throw std::runtime_error("points needs to be in row-major storage order"); + if(ws_points.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface points needs to be in row-major storage order"); + if(ws_normal.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface normal vector needs to be in row-major storage order"); + if(K.layout() != xt::layout_type::row_major) + throw std::runtime_error("surface_current needs to be in row-major storage order"); + + int num_points = points.shape(0); + int num_ws_points = ws_points.shape(0); + Array A = xt::zeros({points.shape(0), points.shape(1)}); + double fak = 1e-7; + + #pragma omp parallel for schedule(static) + for(int i = 0; i < num_points; ++i) { + double Ax = 0., Ay = 0., Az = 0.; + double px = points(i, 0), py = points(i, 1), pz = points(i, 2); + for (int j = 0; j < num_ws_points; ++j) { + double xx = ws_points(j, 0), yy = ws_points(j, 1), zz = ws_points(j, 2); + double nx = ws_normal(j, 0), ny = ws_normal(j, 1), nz = ws_normal(j, 2); + double Kx = K(j, 0), Ky = K(j, 1), Kz = K(j, 2); + double rx = px - xx, ry = py - yy, rz = pz - zz; + double rmag2 = rx*rx + ry*ry + rz*rz; + double rmag_inv = 1.0 / std::sqrt(rmag2); + double nmag = std::sqrt(nx*nx + ny*ny + nz*nz); + Ax += nmag * Kx * rmag_inv; + Ay += nmag * Ky * rmag_inv; + Az += nmag * Kz * rmag_inv; + } + A(i, 0) = fak * Ax; + A(i, 1) = fak * Ay; + A(i, 2) = fak * Az; + } + return A; +} + +#endif + +#if defined(USE_XSIMD) +Array WindingSurfacedA(Array& points, Array& ws_points, Array& ws_normal, Array& K) +{ + // warning: row_major checks below do NOT throw an error correctly on a compute node on Cori + if(points.layout() != xt::layout_type::row_major) + throw std::runtime_error("points needs to be in row-major storage order"); + if(ws_points.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface points needs to be in row-major storage order"); + if(ws_normal.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface normal vector needs to be in row-major storage order"); + if(K.layout() != xt::layout_type::row_major) + throw std::runtime_error("surface_current needs to be in row-major storage order"); + + int num_points = points.shape(0); + int num_ws_points = ws_points.shape(0); + constexpr int simd_size = xsimd::simd_type::size; + Array dA = xt::zeros({points.shape(0), points.shape(1), points.shape(1)}); + + // initialize pointer to the beginning of ws_points + double* ws_points_ptr = &(ws_points(0, 0)); + double* ws_normal_ptr = &(ws_normal(0, 0)); + double* K_ptr = &(K(0, 0)); + double fak = 1e-7; // mu0 divided by 4 * pi factor + + // Loop through the evaluation points by chunks of simd_size + #pragma omp parallel for schedule(static) + for(int i = 0; i < num_points; i += simd_size) { + auto point_i = Vec3dSimd(); + auto dA_i1 = Vec3dSimd(); + auto dA_i2 = Vec3dSimd(); + auto dA_i3 = Vec3dSimd(); + + // check that i + k isn't bigger than num_points + int klimit = std::min(simd_size, num_points - i); + for(int k = 0; k < klimit; k++){ + for (int d = 0; d < 3; ++d) { + point_i[d][k] = points(i + k, d); + } + } + // Sum contributions from all the winding surface points + // i.e. do the surface integral over the winding surface + for (int j = 0; j < num_ws_points; ++j) { + Vec3dSimd r_j = Vec3dSimd(ws_points_ptr[3 * j + 0], ws_points_ptr[3 * j + 1], ws_points_ptr[3 * j + 2]); + Vec3dSimd n_j = Vec3dSimd(ws_normal_ptr[3 * j + 0], ws_normal_ptr[3 * j + 1], ws_normal_ptr[3 * j + 2]); + Vec3dSimd K_j = Vec3dSimd(K_ptr[3 * j + 0], K_ptr[3 * j + 1], K_ptr[3 * j + 2]); + Vec3dSimd r = point_i - r_j; + simd_t rmag_2 = normsq(r); + simd_t rmag_inv = rsqrt(rmag_2); + simd_t rmag_inv_3 = rmag_inv * (rmag_inv * rmag_inv); + simd_t nmag = sqrt(normsq(n_j)); + dA_i1.x += - nmag * K_j.x * r.x * rmag_inv_3; + dA_i1.y += - nmag * K_j.y * r.x * rmag_inv_3; + dA_i1.z += - nmag * K_j.z * r.x * rmag_inv_3; + dA_i2.x += - nmag * K_j.x * r.y * rmag_inv_3; + dA_i2.y += - nmag * K_j.y * r.y * rmag_inv_3; + dA_i2.z += - nmag * K_j.z * r.y * rmag_inv_3; + dA_i3.x += - nmag * K_j.x * r.z * rmag_inv_3; + dA_i3.y += - nmag * K_j.y * r.z * rmag_inv_3; + dA_i3.z += - nmag * K_j.z * r.z * rmag_inv_3; + } + for(int k = 0; k < klimit; k++){ + dA(i + k, 0, 0) = fak * dA_i1.x[k]; + dA(i + k, 0, 1) = fak * dA_i1.y[k]; + dA(i + k, 0, 2) = fak * dA_i1.z[k]; + dA(i + k, 1, 0) = fak * dA_i2.x[k]; + dA(i + k, 1, 1) = fak * dA_i2.y[k]; + dA(i + k, 1, 2) = fak * dA_i2.z[k]; + dA(i + k, 2, 0) = fak * dA_i3.x[k]; + dA(i + k, 2, 1) = fak * dA_i3.y[k]; + dA(i + k, 2, 2) = fak * dA_i3.z[k]; + } + } + return dA; +} + +#else + +Array WindingSurfacedA(Array& points, Array& ws_points, Array& ws_normal, Array& K) +{ + if(points.layout() != xt::layout_type::row_major) + throw std::runtime_error("points needs to be in row-major storage order"); + if(ws_points.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface points needs to be in row-major storage order"); + if(ws_normal.layout() != xt::layout_type::row_major) + throw std::runtime_error("winding surface normal vector needs to be in row-major storage order"); + if(K.layout() != xt::layout_type::row_major) + throw std::runtime_error("surface_current needs to be in row-major storage order"); + + int num_points = points.shape(0); + int num_ws_points = ws_points.shape(0); + Array dA = xt::zeros({points.shape(0), points.shape(1), points.shape(1)}); + double fak = 1e-7; + + #pragma omp parallel for schedule(static) + for(int i = 0; i < num_points; ++i) { + double dA11 = 0., dA12 = 0., dA13 = 0.; + double dA21 = 0., dA22 = 0., dA23 = 0.; + double dA31 = 0., dA32 = 0., dA33 = 0.; + double px = points(i, 0), py = points(i, 1), pz = points(i, 2); + for (int j = 0; j < num_ws_points; ++j) { + double xx = ws_points(j, 0), yy = ws_points(j, 1), zz = ws_points(j, 2); + double nx = ws_normal(j, 0), ny = ws_normal(j, 1), nz = ws_normal(j, 2); + double Kx = K(j, 0), Ky = K(j, 1), Kz = K(j, 2); + double rx = px - xx, ry = py - yy, rz = pz - zz; + double rmag2 = rx*rx + ry*ry + rz*rz; + double rmag_inv = 1.0 / std::sqrt(rmag2); + double rmag_inv_3 = rmag_inv * rmag_inv * rmag_inv; + double nmag = std::sqrt(nx*nx + ny*ny + nz*nz); + dA11 += -nmag * Kx * rx * rmag_inv_3; + dA12 += -nmag * Ky * rx * rmag_inv_3; + dA13 += -nmag * Kz * rx * rmag_inv_3; + dA21 += -nmag * Kx * ry * rmag_inv_3; + dA22 += -nmag * Ky * ry * rmag_inv_3; + dA23 += -nmag * Kz * ry * rmag_inv_3; + dA31 += -nmag * Kx * rz * rmag_inv_3; + dA32 += -nmag * Ky * rz * rmag_inv_3; + dA33 += -nmag * Kz * rz * rmag_inv_3; + } + dA(i, 0, 0) = fak * dA11; + dA(i, 0, 1) = fak * dA12; + dA(i, 0, 2) = fak * dA13; + dA(i, 1, 0) = fak * dA21; + dA(i, 1, 1) = fak * dA22; + dA(i, 1, 2) = fak * dA23; + dA(i, 2, 0) = fak * dA31; + dA(i, 2, 1) = fak * dA32; + dA(i, 2, 2) = fak * dA33; + } + return dA; +} + +#endif + +// Calculate the geometric factor needed for the A^B term in winding surface optimization +std::tuple winding_surface_field_Bn(Array& points_plasma, Array& points_coil, Array& normal_plasma, Array& normal_coil, bool stellsym, Array& zeta_coil, Array& theta_coil, int ndofs, Array& m, Array& n, int nfp) +{ + // warning: row_major checks below do NOT throw an error correctly on a compute node on Cori + if(points_plasma.layout() != xt::layout_type::row_major) + throw std::runtime_error("points needs to be in row-major storage order"); + if(points_coil.layout() != xt::layout_type::row_major) + throw std::runtime_error("points needs to be in row-major storage order"); + if(normal_plasma.layout() != xt::layout_type::row_major) + throw std::runtime_error("normal_plasma needs to be in row-major storage order"); + if(normal_coil.layout() != xt::layout_type::row_major) + throw std::runtime_error("normal_winding_surface needs to be in row-major storage order"); + if(zeta_coil.layout() != xt::layout_type::row_major) + throw std::runtime_error("phi needs to be in row-major storage order"); + if(theta_coil.layout() != xt::layout_type::row_major) + throw std::runtime_error("theta needs to be in row-major storage order"); + + int num_plasma = normal_plasma.shape(0); + int num_coil = normal_coil.shape(0); + Array gij = xt::zeros({num_plasma, num_coil}); + Array gj = xt::zeros({num_plasma, ndofs}); + Array Ajk = xt::zeros({ndofs, ndofs}); + + // initialize pointer to the beginning of the coil quadrature points + //double* coil_points_ptr = &(points_coil(0, 0)); + //double* normal_coil_ptr = &(normal_coil(0, 0)); + double fak = 1e-7; // mu0 divided by 4 * pi factor + + #pragma omp parallel for schedule(static) + for(int i = 0; i < num_plasma; i++) { + double npx = normal_plasma(i, 0); + double npy = normal_plasma(i, 1); + double npz = normal_plasma(i, 2); + + // Loop through the coil quadrature points, using all the symmetries + for (int j = 0; j < num_coil; ++j) { + double ncx = normal_coil(j, 0); + double ncy = normal_coil(j, 1); + double ncz = normal_coil(j, 2); + double rx = points_plasma(i, 0) - points_coil(j, 0); + double ry = points_plasma(i, 1) - points_coil(j, 1); + double rz = points_plasma(i, 2) - points_coil(j, 2); + double rmag2 = rx * rx + ry * ry + rz * rz; + double rmag_inv = 1.0 / std::sqrt(rmag2); + double rmag_inv_3 = rmag_inv * rmag_inv * rmag_inv; + double rmag_inv_5 = rmag_inv_3 * rmag_inv * rmag_inv; + double npdotnc = npx * ncx + npy * ncy + npz * ncz; + double rdotnp = rx * npx + ry * npy + rz * npz; + double rdotnc = rx * ncx + ry * ncy + rz * ncz; + double G_i = npdotnc * rmag_inv_3 - 3.0 * rdotnp * rdotnc * rmag_inv_5; + gij(i, j) = fak * G_i; + } + } + // Precompute cos/sin: angle depends only on (j,k), not on plasma point i + int num_dofs_half = m.size(); + Array sin_phi = xt::zeros({num_dofs_half, num_coil}); + Array cos_phi = xt::zeros({num_dofs_half, num_coil}); + #pragma omp parallel for schedule(static) collapse(2) + for (int j = 0; j < num_dofs_half; j++) { + for (int k = 0; k < num_coil; k++) { + double angle = 2 * M_PI * m(j) * theta_coil(k) - 2 * M_PI * n(j) * zeta_coil(k) * nfp; + sin_phi(j, k) = std::sin(angle); + cos_phi(j, k) = std::cos(angle); + } + } + #pragma omp parallel for schedule(static) + for(int i = 0; i < num_plasma; i++) { + // now take gij and loop over the dofs (Eq. A10 in REGCOIL paper) + for (int j = 0; j < num_dofs_half; j++) { + for(int k = 0; k < num_coil; k++){ + gj(i, j) += sin_phi(j, k) * gij(i, k); + if (!stellsym) { + gj(i, j + num_dofs_half) += cos_phi(j, k) * gij(i, k); + } + } + } + } + + // Precompute 1/n_norm for each plasma point (was computed ndofs^2 times per point) + Array n_norm_inv = xt::zeros({num_plasma}); + #pragma omp parallel for schedule(static) + for(int i = 0; i < num_plasma; i++) { + double npx = normal_plasma(i, 0); + double npy = normal_plasma(i, 1); + double npz = normal_plasma(i, 2); + n_norm_inv(i) = 1.0 / std::sqrt(npx * npx + npy * npy + npz * npz); + } + // j outer for parallelization (each thread owns distinct Ajk rows) + #pragma omp parallel for schedule(static) + for(int j = 0; j < ndofs; j++) { + for(int i = 0; i < num_plasma; i++) { + double n_inv = n_norm_inv(i); + for(int k = 0; k < ndofs; k++) { + Ajk(j, k) += gj(i, j) * gj(i, k) * n_inv; + } + } + } + return std::make_tuple(gj, Ajk); +} + +// Compute GI part of Bnormal +Array winding_surface_field_Bn_GI(Array& points_plasma, Array& points_coil, Array& normal_plasma, Array& zeta_coil, Array& theta_coil, double G, double I, Array& gammadash1, Array& gammadash2) +{ + int num_plasma = normal_plasma.shape(0); + int num_coil = points_coil.shape(0); + double fak = 1e-7; // mu0 divided by 8 * pi^2 factor + Array B_GI = xt::zeros({num_plasma}); + #pragma omp parallel for schedule(static) + for(int i = 0; i < num_plasma; i++) { + double nx = normal_plasma(i, 0); + double ny = normal_plasma(i, 1); + double nz = normal_plasma(i, 2); + double nmag = std::sqrt(nx * nx + ny * ny + nz * nz); + nx = nx / nmag; + ny = ny / nmag; + nz = nz / nmag; + for(int j = 0; j < num_coil; j++) { + double rx = points_plasma(i, 0) - points_coil(j, 0); + double ry = points_plasma(i, 1) - points_coil(j, 1); + double rz = points_plasma(i, 2) - points_coil(j, 2); + double rmag2 = rx * rx + ry * ry + rz * rz; + double rmag_inv = 1.0 / std::sqrt(rmag2); + double rmag_inv_3 = rmag_inv * rmag_inv * rmag_inv; + double GIx = G * gammadash2(j, 0) - I * gammadash1(j, 0); + double GIy = G * gammadash2(j, 1) - I * gammadash1(j, 1); + double GIz = G * gammadash2(j, 2) - I * gammadash1(j, 2); + double GIcrossr_dotn = nx * (GIy * rz - GIz * ry) + ny * (GIz * rx - GIx * rz) + nz * (GIx * ry - GIy * rx); + B_GI(i) += fak * GIcrossr_dotn * rmag_inv_3; + } + } + return B_GI; +} + +// Compute the Ak matrix associated with ||K||_2^2 = ||Ak * phi_mn - d||_2^2 term in REGCOIL +std::tuple winding_surface_field_K2_matrices(Array& dr_dzeta_coil, Array& dr_dtheta_coil, Array& normal_coil, bool stellsym, Array& zeta_coil, Array& theta_coil, int ndofs, Array& m, Array& n, int nfp, double G, double I) +{ + // warning: row_major checks below do NOT throw an error correctly on a compute node on Cori + if(dr_dzeta_coil.layout() != xt::layout_type::row_major) + throw std::runtime_error("dr_dzeta_coil needs to be in row-major storage order"); + if(dr_dtheta_coil.layout() != xt::layout_type::row_major) + throw std::runtime_error("dr_dtheta_coil needs to be in row-major storage order"); + if(normal_coil.layout() != xt::layout_type::row_major) + throw std::runtime_error("normal_winding_surface needs to be in row-major storage order"); + if(zeta_coil.layout() != xt::layout_type::row_major) + throw std::runtime_error("phi needs to be in row-major storage order"); + if(theta_coil.layout() != xt::layout_type::row_major) + throw std::runtime_error("theta needs to be in row-major storage order"); + + int num_coil = normal_coil.shape(0); + Array d = xt::zeros({num_coil, 3}); + Array fj = xt::zeros({num_coil, 3, ndofs}); + // Loop through the coil quadrature points, using all the symmetries + #pragma omp parallel for schedule(static) + for (int j = 0; j < num_coil; ++j) { + double nx = normal_coil(j, 0); + double ny = normal_coil(j, 1); + double nz = normal_coil(j, 2); + double normN = sqrt(nx * nx + ny * ny + nz * nz); + d(j, 0) = (G * dr_dtheta_coil(j, 0) - I * dr_dzeta_coil(j, 0)) / (2 * M_PI); + d(j, 1) = (G * dr_dtheta_coil(j, 1) - I * dr_dzeta_coil(j, 1)) / (2 * M_PI); + d(j, 2) = (G * dr_dtheta_coil(j, 2) - I * dr_dzeta_coil(j, 2)) / (2 * M_PI); + + for (int k = 0; k < m.size(); k++) { + double angle = 2 * M_PI * m(k) * theta_coil(j) - 2 * M_PI * n(k) * zeta_coil(j) * nfp; + double cphi = std::cos(angle); + double sphi = std::sin(angle); + fj(j, 0, k) = cphi * (m(k) * dr_dzeta_coil(j, 0) + nfp * n(k) * dr_dtheta_coil(j, 0)); + fj(j, 1, k) = cphi * (m(k) * dr_dzeta_coil(j, 1) + nfp * n(k) * dr_dtheta_coil(j, 1)); + fj(j, 2, k) = cphi * (m(k) * dr_dzeta_coil(j, 2) + nfp * n(k) * dr_dtheta_coil(j, 2)); + if (! stellsym) { + fj(j, 0, k+m.size()) = -sphi * (m(k) * dr_dzeta_coil(j, 0) + nfp * n(k) * dr_dtheta_coil(j, 0)); + fj(j, 1, k+m.size()) = -sphi * (m(k) * dr_dzeta_coil(j, 1) + nfp * n(k) * dr_dtheta_coil(j, 1)); + fj(j, 2, k+m.size()) = -sphi * (m(k) * dr_dzeta_coil(j, 2) + nfp * n(k) * dr_dtheta_coil(j, 2)); + } + } + } + return std::make_tuple(d, fj); +} diff --git a/src/simsoptpp/winding_surface.h b/src/simsoptpp/winding_surface.h new file mode 100644 index 000000000..3449ca770 --- /dev/null +++ b/src/simsoptpp/winding_surface.h @@ -0,0 +1,27 @@ +#include // c++ tuples +#include "xtensor-python/pyarray.hpp" +typedef xt::pyarray Array; + +// Compute the Bnormal using the REGCOIL method. +Array WindingSurfaceBn_REGCOIL(Array& points, Array& ws_points, Array& ws_normal, Array& current_potential, Array& plasma_normal); + +// Compute the B field using the Biot-Savart law from the winding surface current K. +Array WindingSurfaceB(Array& points, Array& ws_points, Array& ws_normal, Array& K); + +// Compute the dB/dX field using the Biot-Savart law from the winding surface current K. +Array WindingSurfacedB(Array& points, Array& ws_points, Array& ws_normal, Array& K); + +// Compute the A field using the Biot-Savart law from the winding surface current K. +Array WindingSurfaceA(Array& points, Array& ws_points, Array& ws_normal, Array& K); + +// Compute the dA/dX field using the Biot-Savart law from the winding surface current K. +Array WindingSurfacedA(Array& points, Array& ws_points, Array& ws_normal, Array& K); + +// Compute the Bn field using the winding surface current K. +std::tuple winding_surface_field_Bn(Array& points_plasma, Array& points_coil, Array& normal_plasma, Array& normal_coil, bool stellsym, Array& zeta_coil, Array& theta_coil, int ndofs, Array& m, Array& n, int nfp); + +// Compute part of the Bn field using the winding surface current K and the plasma surface normal. +Array winding_surface_field_Bn_GI(Array& points_plasma, Array& points_coil, Array& normal_plasma, Array& zeta_coil, Array& theta_coil, double G, double I, Array& gammadash1, Array& gammadash2); + +// Compute the K2 matrices using the winding surface current K and the plasma surface normal. +std::tuple winding_surface_field_K2_matrices(Array& dr_dzeta_coil, Array& dr_dtheta_coil, Array& normal_coil, bool stellsym, Array& zeta_coil, Array& theta_coil, int ndofs, Array& m, Array& n, int nfp, double G, double I); diff --git a/tests/field/test_currentpotential.py b/tests/field/test_currentpotential.py new file mode 100644 index 000000000..213a56fa4 --- /dev/null +++ b/tests/field/test_currentpotential.py @@ -0,0 +1,578 @@ +import unittest +from pathlib import Path + +import numpy as np + +from simsopt.field import CurrentPotentialFourier +from simsopt.field import CurrentPotentialSolve +from simsopt.geo import SurfaceRZFourier +from scipy.io import netcdf_file +from simsopt.util import in_github_actions + +TEST_DIR = Path(__file__).parent / ".." / "test_files" + +stellsym_list = [True, False] + +try: + import pyevtk # noqa: F401 + pyevtk_found = True +except ImportError: + pyevtk_found = False + + +class CurrentPotentialTests(unittest.TestCase): + def test_compare_K_with_regcoil(self): + # axisymmetric case with no Phimnc, non-axisymmetric, axisymmetric with Phimnc + for filename in ['regcoil_out.w7x.nc', 'regcoil_out.near_axis_asym.nc']: + filename = TEST_DIR / filename + cp = CurrentPotentialFourier.from_netcdf(filename) + cp.set_current_potential_from_regcoil(filename, -1) + f = netcdf_file(filename, 'r', mmap=False) + cp_regcoil = f.variables['single_valued_current_potential_mn'][()][-1, :] + _xm_regcoil = f.variables['xm_potential'][()] + stellsym = f.variables['symmetry_option'][()] + if stellsym == 1: + stellsym = True + else: + stellsym = False + + # need to add phic(0, 0) term to cp_regcoil to compare them + assert np.allclose(cp.get_dofs(), cp_regcoil) + K2_regcoil = f.variables['K2'][()][-1, :, :] + K = cp.K() + K2 = np.sum(K*K, axis=2) + K2_average = np.mean(K2, axis=(0, 1)) + print(K2[0:int(len(cp.quadpoints_phi)/cp.nfp), :]/K2_average, K2.shape, K2_average) + print(K2_regcoil[0:int(len(cp.quadpoints_phi)/cp.nfp), :]/K2_average, K2_regcoil.shape) + assert np.allclose(K2[0:int(len(cp.quadpoints_phi)/cp.nfp), :]/K2_average, K2_regcoil/K2_average) + + +def get_currentpotential(cptype, stellsym, phis=None, thetas=None): + np.random.seed(2) + mpol = 4 + ntor = 3 + nfp = 2 + phis = phis if phis is not None else np.linspace(0, 1, 31, endpoint=False) + thetas = thetas if thetas is not None else np.linspace(0, 1, 31, endpoint=False) + + if cptype == "CurrentPotentialFourier": + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, nfp=nfp, stellsym=stellsym, mpol=mpol, ntor=ntor, + quadpoints_phi=phis, quadpoints_theta=thetas) + cp.x = cp.x * 0. + cp.phis[1, ntor + 0] = 0.3 + else: + assert False + + dofs = cp.get_dofs() + np.random.seed(2) + rand_scale = 0.01 + cp.x = dofs + rand_scale * np.random.rand(len(dofs)) # .reshape(dofs.shape) + return cp + + +class CurrentPotentialFourierTests(unittest.TestCase): + + def test_init(self): + """ + Modeled after test_init in test_surface_rzfourier + """ + + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, nfp=2, mpol=3, ntor=2) + self.assertEqual(cp.phis.shape, (4, 5)) + self.assertEqual(cp.phic.shape, (4, 5)) + + cp = CurrentPotentialFourier(s, nfp=10, mpol=1, ntor=3, stellsym=False) + self.assertEqual(cp.phis.shape, (2, 7)) + self.assertEqual(cp.phic.shape, (2, 7)) + + def test_get_dofs(self): + """ + Modeled after test_get_dofs from test_surface_rzfourier + Test that we can convert the degrees of freedom into a 1D vector + """ + + # axisymmetric current potential + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=3, ntor=0) + cp.phis[1, 0] = 0.2 + cp.phis[2, 0] = 1.3 + cp.phis[3, 0] = 0.7 + dofs = cp.get_dofs() + self.assertEqual(dofs.shape, (3,)) + self.assertAlmostEqual(dofs[0], 0.2) + self.assertAlmostEqual(dofs[1], 1.3) + self.assertAlmostEqual(dofs[2], 0.7) + + # non-axisymmetric current potential + cp = CurrentPotentialFourier(s, mpol=3, ntor=1) + cp.phis[:, :] = [[100, 101, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]] + dofs = cp.get_dofs() + self.assertEqual(dofs.shape, (3*(2*1 + 1) + 1,)) + for j in range(len(dofs)): + self.assertAlmostEqual(dofs[j], j + 3) + + def test_set_dofs(self): + """ + Modeled after test_get_dofs from test_surface_rzfourier + Test that we can set the shape from a 1D vector + """ + + # First try an axisymmetric current potential for simplicity: + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=3, ntor=0) + cp.set_dofs([-1.1, 0.7, 0.5]) + self.assertAlmostEqual(cp.phis[1, 0], -1.1) + self.assertAlmostEqual(cp.phis[2, 0], 0.7) + self.assertAlmostEqual(cp.phis[3, 0], 0.5) + + # Now try a nonaxisymmetric shape: + cp = CurrentPotentialFourier(s, mpol=3, ntor=1) + ndofs = len(cp.get_dofs()) + cp.set_dofs(np.array(list(range(ndofs))) + 1) + self.assertAlmostEqual(cp.phis[0, 0], 0) + self.assertAlmostEqual(cp.phis[0, 1], 0) + self.assertAlmostEqual(cp.phis[0, 2], 1) + self.assertAlmostEqual(cp.phis[1, 0], 2) + self.assertAlmostEqual(cp.phis[1, 1], 3) + self.assertAlmostEqual(cp.phis[1, 2], 4) + self.assertAlmostEqual(cp.phis[2, 0], 5) + self.assertAlmostEqual(cp.phis[2, 1], 6) + self.assertAlmostEqual(cp.phis[2, 2], 7) + self.assertAlmostEqual(cp.phis[3, 0], 8) + self.assertAlmostEqual(cp.phis[3, 1], 9) + self.assertAlmostEqual(cp.phis[3, 2], 10) + + # Now try a nonaxisymmetric shape without stell sym + cp = CurrentPotentialFourier(s, mpol=3, ntor=1, stellsym=False) + ndofs = len(cp.get_dofs()) + cp.set_dofs(np.array(list(range(ndofs))) + 1) + self.assertAlmostEqual(cp.phis[0, 0], 0) + self.assertAlmostEqual(cp.phis[0, 1], 0) + self.assertAlmostEqual(cp.phis[0, 2], 1) + self.assertAlmostEqual(cp.phis[1, 0], 2) + self.assertAlmostEqual(cp.phis[1, 1], 3) + self.assertAlmostEqual(cp.phis[1, 2], 4) + self.assertAlmostEqual(cp.phis[2, 0], 5) + self.assertAlmostEqual(cp.phis[2, 1], 6) + self.assertAlmostEqual(cp.phis[2, 2], 7) + self.assertAlmostEqual(cp.phis[3, 0], 8) + self.assertAlmostEqual(cp.phis[3, 1], 9) + self.assertAlmostEqual(cp.phis[3, 2], 10) + self.assertAlmostEqual(cp.phic[0, 0], 0) + self.assertAlmostEqual(cp.phic[0, 1], 0) + self.assertAlmostEqual(cp.phic[0, 2], 11) + self.assertAlmostEqual(cp.phic[1, 0], 12) + self.assertAlmostEqual(cp.phic[1, 1], 13) + self.assertAlmostEqual(cp.phic[1, 2], 14) + self.assertAlmostEqual(cp.phic[2, 0], 15) + self.assertAlmostEqual(cp.phic[2, 1], 16) + self.assertAlmostEqual(cp.phic[2, 2], 17) + self.assertAlmostEqual(cp.phic[3, 0], 18) + self.assertAlmostEqual(cp.phic[3, 1], 19) + self.assertAlmostEqual(cp.phic[3, 2], 20) + + def test_change_resolution(self): + """ + Modeled after test_get_dofs from test_surface_rzfourier + Check that we can change mpol and ntor. + """ + for mpol in [1, 2]: + for ntor in [0, 1]: + for stellsym in [True, False]: + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=mpol, ntor=ntor) + ndofs = len(cp.get_dofs()) + cp.set_dofs((np.random.rand(ndofs) - 0.5) * 0.01) + cp.set_phis(1, 0, 0.13) + phi1 = cp.Phi() + + cp.change_resolution(mpol+1, ntor) + s.recalculate = True + phi2 = cp.Phi() + self.assertTrue(np.allclose(phi1, phi2)) + + def test_get_phic(self): + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=1, ntor=0, stellsym=False) + cp.x = [0.7, -1.1] + self.assertAlmostEqual(cp.get_phis(1, 0), 0.7) + self.assertAlmostEqual(cp.get_phic(1, 0), -1.1) + self.assertAlmostEqual(cp.get_phis(0, 0), 0.0) + self.assertAlmostEqual(cp.get_phic(0, 0), 0.0) + + def test_get_phis(self): + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=3, ntor=1) + ndofs = len(cp.get_dofs()) + cp.x = np.array(list(range(ndofs))) + 1 + + self.assertAlmostEqual(cp.get_phis(0, -1), 0) + self.assertAlmostEqual(cp.get_phis(0, 0), 0) + self.assertAlmostEqual(cp.get_phis(0, 1), 1) + self.assertAlmostEqual(cp.get_phis(1, -1), 2) + self.assertAlmostEqual(cp.get_phis(1, 0), 3) + self.assertAlmostEqual(cp.get_phis(1, 1), 4) + self.assertAlmostEqual(cp.get_phis(2, -1), 5) + self.assertAlmostEqual(cp.get_phis(2, 0), 6) + self.assertAlmostEqual(cp.get_phis(2, 1), 7) + self.assertAlmostEqual(cp.get_phis(3, -1), 8) + self.assertAlmostEqual(cp.get_phis(3, 0), 9) + self.assertAlmostEqual(cp.get_phis(3, 1), 10) + + def test_set_phic(self): + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=1, ntor=0, stellsym=False) + cp.x = [2.9, -1.1] + cp.set_phic(1, 0, 3.1) + self.assertAlmostEqual(cp.x[1], 3.1) + self.assertAlmostEqual(cp.get_phic(1, 0), 3.1) + + def test_set_phis(self): + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=1, ntor=0, stellsym=True) + s.x = [2.9, -1.1, 0.7] + cp.set_phis(1, 0, 1.4) + self.assertAlmostEqual(cp.x[0], 1.4) + + def test_names_order(self): + """ + Verify that the order of phis and phic in the dof names is + correct. This requires that the order of these four arrays in + ``_make_names()`` matches the order in the C++ functions + ``set_dofs_impl()`` and ``get_dofs()`` in + ``src/simsoptpp/currentpotentialfourier.h``. + """ + mpol = 1 + ntor = 1 + nfp = 4 + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, nfp=nfp, mpol=mpol, ntor=ntor, stellsym=False) + cp.set_phic(0, 1, 100.0) + cp.set_phis(0, 1, 200.0) + self.assertAlmostEqual(cp.get('Phic(0,1)'), 100.0) + self.assertAlmostEqual(cp.get('Phis(0,1)'), 200.0) + + def test_mn(self): + """ + Test the arrays of mode numbers m and n. + """ + mpol = 3 + ntor = 2 + nfp = 4 + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, nfp=nfp, mpol=mpol, ntor=ntor) + m_correct = [0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3] + n_correct = [1, 2, -2, -1, 0, 1, 2, -2, -1, 0, 1, 2, -2, -1, 0, 1, 2] + np.testing.assert_array_equal(cp.m, m_correct) + np.testing.assert_array_equal(cp.n, n_correct) + + m_correct = [0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3] + n_correct = [1, 2, -2, -1, 0, 1, 2, -2, -1, 0, 1, 2, -2, -1, 0, 1, 2] + + cp = CurrentPotentialFourier(s, nfp=nfp, mpol=mpol, ntor=ntor, stellsym=False) + np.testing.assert_array_equal(cp.m, m_correct + m_correct) + np.testing.assert_array_equal(cp.n, n_correct + n_correct) + + def test_mn_matches_names(self): + """ + Verify that the m and n attributes match the dof names. + """ + mpol = 2 + ntor = 3 + nfp = 5 + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, nfp=nfp, mpol=mpol, ntor=ntor, stellsym=True) + names = [name[4:] for name in cp.local_dof_names] + names2 = [f'({m},{n})' for m, n in zip(cp.m, cp.n)] + self.assertEqual(names, names2) + + # Now try a non-stellarator-symmetric case: + cp = CurrentPotentialFourier(s, nfp=nfp, mpol=mpol, ntor=ntor, stellsym=False) + assert 'Phic(0,0)' not in cp.local_dof_names + names = [name[4:] for name in cp.local_dof_names] + names2 = [f'({m},{n})' for m, n in zip(cp.m, cp.n)] + self.assertEqual(names, names2) + + def test_get_phic_raises_for_stellsym(self): + """get_phic raises ValueError when stellsym=True (phic does not exist).""" + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=2, ntor=1, stellsym=True) + with self.assertRaises(ValueError) as cm: + cp.get_phic(1, 0) + self.assertIn('phic does not exist', str(cm.exception)) + + def test_set_phic_raises_for_stellsym(self): + """set_phic raises ValueError when stellsym=True (phic does not exist).""" + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=2, ntor=1, stellsym=True) + with self.assertRaises(ValueError) as cm: + cp.set_phic(1, 0, 1.0) + self.assertIn('phic does not exist', str(cm.exception)) + + def test_validate_mn_raises_index_error_m_negative(self): + """_validate_mn raises IndexError when m < 0.""" + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=2, ntor=1) + with self.assertRaises(IndexError) as cm: + cp.get_phis(-1, 0) + self.assertIn('m must be >= 0', str(cm.exception)) + + def test_validate_mn_raises_index_error_m_gt_mpol(self): + """_validate_mn raises IndexError when m > mpol.""" + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=2, ntor=1) + with self.assertRaises(IndexError) as cm: + cp.get_phis(3, 0) + self.assertIn('m must be <= mpol', str(cm.exception)) + + def test_validate_mn_raises_index_error_n_gt_ntor(self): + """_validate_mn raises IndexError when n > ntor.""" + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=2, ntor=1) + with self.assertRaises(IndexError) as cm: + cp.get_phis(1, 2) + self.assertIn('n must be <= ntor', str(cm.exception)) + + def test_validate_mn_raises_index_error_n_lt_ntor(self): + """_validate_mn raises IndexError when n < -ntor.""" + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=2, ntor=1) + with self.assertRaises(IndexError) as cm: + cp.get_phis(1, -2) + self.assertIn('n must be >= -ntor', str(cm.exception)) + + def test_set_current_potential_from_regcoil_raises_incorrect_mpol(self): + """set_current_potential_from_regcoil raises ValueError for incorrect mpol.""" + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=2, ntor=1) + filename = TEST_DIR / 'regcoil_out.w7x_infty.nc' + with self.assertRaises(ValueError) as cm: + cp.set_current_potential_from_regcoil(filename, 0) + self.assertIn('mpol_potential', str(cm.exception)) + + def test_set_current_potential_from_regcoil_raises_incorrect_ntor(self): + """set_current_potential_from_regcoil raises ValueError for incorrect ntor.""" + f = netcdf_file(TEST_DIR / 'regcoil_out.w7x_infty.nc', 'r', mmap=False) + mpol = int(f.variables['mpol_potential'][()]) + f.close() + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=mpol, ntor=0) + with self.assertRaises(ValueError) as cm: + cp.set_current_potential_from_regcoil(TEST_DIR / 'regcoil_out.w7x_infty.nc', 0) + self.assertIn('ntor_potential', str(cm.exception)) + + def test_set_current_potential_from_regcoil_raises_incorrect_nfp(self): + """set_current_potential_from_regcoil raises ValueError for incorrect nfp.""" + f = netcdf_file(TEST_DIR / 'regcoil_out.w7x_infty.nc', 'r', mmap=False) + mpol = int(f.variables['mpol_potential'][()]) + ntor = int(f.variables['ntor_potential'][()]) + nfp_file = int(f.variables['nfp'][()]) + f.close() + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=mpol, ntor=ntor, nfp=nfp_file + 1) + with self.assertRaises(ValueError) as cm: + cp.set_current_potential_from_regcoil(TEST_DIR / 'regcoil_out.w7x_infty.nc', 0) + self.assertIn('nfp', str(cm.exception)) + + def test_set_current_potential_from_regcoil_raises_incorrect_stellsym(self): + """set_current_potential_from_regcoil raises ValueError for incorrect stellsym.""" + f = netcdf_file(TEST_DIR / 'regcoil_out.w7x_infty.nc', 'r', mmap=False) + mpol = int(f.variables['mpol_potential'][()]) + ntor = int(f.variables['ntor_potential'][()]) + nfp = int(f.variables['nfp'][()]) + f.close() + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=mpol, ntor=ntor, nfp=nfp, stellsym=False) + with self.assertRaises(ValueError) as cm: + cp.set_current_potential_from_regcoil(TEST_DIR / 'regcoil_out.w7x_infty.nc', 0) + self.assertIn('stellsym', str(cm.exception)) + + def test_fixed_range_fix(self): + """fixed_range with fixed=True fixes the specified modes.""" + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=2, ntor=1, stellsym=True) + self.assertTrue(np.all(cp.local_dofs_free_status)) + cp.fixed_range(mmin=1, mmax=2, nmin=-1, nmax=1, fixed=True) + # Phis(1,-1), Phis(1,0), Phis(1,1), Phis(2,-1), Phis(2,0), Phis(2,1) should be fixed + for name in ['Phis(1,-1)', 'Phis(1,0)', 'Phis(1,1)', 'Phis(2,-1)', 'Phis(2,0)', 'Phis(2,1)']: + idx = cp.local_full_dof_names.index(name) + self.assertFalse(cp.local_dofs_free_status[idx], msg=f"{name} should be fixed") + # Phis(0,1) should remain free + idx = cp.local_full_dof_names.index('Phis(0,1)') + self.assertTrue(cp.local_dofs_free_status[idx]) + + def test_fixed_range_unfix(self): + """fixed_range with fixed=False unfixes the specified modes.""" + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=2, ntor=1, stellsym=True) + cp.fix_all() + cp.fixed_range(mmin=0, mmax=1, nmin=-1, nmax=1, fixed=False) + # Phis(0,1) and Phis(1,-1), Phis(1,0), Phis(1,1) should be unfixed + for name in ['Phis(0,1)', 'Phis(1,-1)', 'Phis(1,0)', 'Phis(1,1)']: + idx = cp.local_full_dof_names.index(name) + self.assertTrue(cp.local_dofs_free_status[idx], msg=f"{name} should be free") + # Phis(2,-1) should remain fixed + idx = cp.local_full_dof_names.index('Phis(2,-1)') + self.assertFalse(cp.local_dofs_free_status[idx]) + + def test_fixed_range_stellsym_false(self): + """fixed_range works for non-stellarator-symmetric (fixes both Phis and Phic).""" + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=1, ntor=1, stellsym=False) + cp.fixed_range(mmin=1, mmax=1, nmin=0, nmax=1, fixed=True) + for name in ['Phis(1,0)', 'Phis(1,1)', 'Phic(1,0)', 'Phic(1,1)']: + idx = cp.local_full_dof_names.index(name) + self.assertFalse(cp.local_dofs_free_status[idx], msg=f"{name} should be fixed") + + def test_K_matrix(self): + """CurrentPotential.K_matrix returns symmetric matrix of correct shape.""" + s = SurfaceRZFourier() + cp = CurrentPotentialFourier(s, mpol=2, ntor=1, stellsym=True) + cp.set_dofs(np.random.rand(cp.num_dofs()) * 0.01) + K_mat = cp.K_matrix() + self.assertEqual(K_mat.shape, (cp.num_dofs(), cp.num_dofs())) + np.testing.assert_allclose(K_mat, K_mat.T, atol=1e-14, err_msg="K_matrix should be symmetric") + # K_matrix should be positive semi-definite (Gram matrix) + eigs = np.linalg.eigvalsh(K_mat) + self.assertTrue(np.all(eigs >= -1e-10), msg="K_matrix should be PSD") + + def test_target(self): + """ + Test that the current potential solve works. Similar to winding_surface.py example. + """ + from matplotlib import pyplot as plt + def run_target_test( + filename='regcoil_out.hsx.nc', + lambda_reg=0.0, + ): + """ + Run REGCOIL (L2 regularization) and Lasso (L1 regularization) + starting from high regularization to low. When fB < fB_target + is achieved, the algorithms quit. This allows one to compare + L2 and L1 results at comparable levels of fB, which seems + like the fairest way to compare them. + + Args: + filename (str): Path to the REGCOIL netcdf file + lambda_reg (float): Regularization parameter + """ + _fB_target = 1e-2 + mpol = 4 + ntor = 4 + coil_ntheta_res = 1 + coil_nzeta_res = coil_ntheta_res + plasma_ntheta_res = coil_ntheta_res + plasma_nzeta_res = coil_ntheta_res + + # Load in low-resolution NCSX file from REGCOIL + cpst = CurrentPotentialSolve.from_netcdf( + TEST_DIR / filename, plasma_ntheta_res, plasma_nzeta_res, coil_ntheta_res, coil_nzeta_res + ) + cp = CurrentPotentialFourier.from_netcdf(TEST_DIR / filename, coil_ntheta_res, coil_nzeta_res) + + # Overwrite low-resolution file with more mpol and ntor modes + cp = CurrentPotentialFourier( + cpst.winding_surface, mpol=mpol, ntor=ntor, + net_poloidal_current_amperes=cp.net_poloidal_current_amperes, + net_toroidal_current_amperes=cp.net_toroidal_current_amperes, + stellsym=True) + cpst = CurrentPotentialSolve(cp, cpst.plasma_surface, cpst.Bnormal_plasma) #, cpst.B_GI) + + optimized_phi_mn, f_B, _ = cpst.solve_tikhonov(lam=lambda_reg) + cp_opt = cpst.current_potential + return(cp_opt, cp, cpst, optimized_phi_mn, f_B) + + + cp_opt, _, _, _, _ = run_target_test(lambda_reg=0.1) + + theta_study1d, phi_study1d = cp_opt.quadpoints_theta, cp_opt.quadpoints_phi + theta_study2d, phi_study2d = np.meshgrid(theta_study1d, phi_study1d) + + # Plotting K + K_mag_study = np.sum(cp_opt.K()**2, axis=2) + plt.figure() + plt.pcolor(phi_study1d, theta_study1d, K_mag_study.T, shading='nearest') + plt.colorbar() + plt.title('|K|') + plt.figure() + plt.pcolor(phi_study2d.T, theta_study2d.T, K_mag_study.T, shading='nearest') + plt.ylabel(r'$\theta$') + plt.xlabel(r'$\phi$') + plt.colorbar() + plt.title('|K|') + + # Contours for Phi + Phi_study = cp_opt.Phi() + plt.figure() + plt.pcolor(phi_study1d, theta_study1d, Phi_study.T, shading='nearest') + plt.colorbar() + plt.ylabel(r'$\theta$') + plt.xlabel(r'$\phi$') + plt.title(r'$\Phi$') + + # Contours for normal component + plt.figure() + plt.pcolor(cp_opt.winding_surface.normal()[:,:,0].T) + plt.colorbar() + plt.ylabel(r'$\theta$') + plt.xlabel(r'$\phi$') + plt.title(r'$\Phi_{normal}$') + if not in_github_actions: + plt.show() + + +class CurrentPotentialTaylorTests(unittest.TestCase): + cptypes = ["CurrentPotentialFourier"] + + def subtest_currentpotential_phi_derivative(self, cptype, stellsym): + epss = [0.5**i for i in range(10, 15)] + phis = np.asarray([0.6] + [0.6 + eps for eps in epss]) + cp = get_currentpotential(cptype, stellsym, phis=phis) + + f0 = cp.Phi()[0, 0] + deriv = cp.Phidash1()[0, 0] + err_old = 1e6 + for i in range(len(epss)): + fh = cp.Phi()[i+1, 0] + deriv_est = (fh-f0)/epss[i] + err = np.linalg.norm(deriv_est-deriv) + assert err < 0.55 * err_old + err_old = err + + def test_currentpotential_phi_derivative(self): + """ + Taylor test to verify that the surface tangent in the phi direction + """ + for cptype in self.cptypes: + for stellsym in [True, False]: + with self.subTest(cptype=cptype, stellsym=stellsym): + self.subtest_currentpotential_phi_derivative(cptype, stellsym) + + def subtest_currentpotential_theta_derivative(self, cptype, stellsym): + epss = [0.5**i for i in range(10, 15)] + thetas = np.asarray([0.6] + [0.6 + eps for eps in epss]) + cp = get_currentpotential(cptype, stellsym, thetas=thetas) + f0 = cp.Phi()[0, 0] + deriv = cp.Phidash2()[0, 0] + err_old = 1e6 + for i in range(len(epss)): + fh = cp.Phi()[0, i+1] + deriv_est = (fh-f0)/epss[i] + err = np.linalg.norm(deriv_est-deriv) + assert err < 0.55 * err_old + err_old = err + + def test_currentpotential_theta_derivative(self): + """ + Taylor test to verify that the surface tangent in the theta direction + """ + for cptype in self.cptypes: + for stellsym in [True, False]: + with self.subTest(cptype=cptype, stellsym=stellsym): + self.subtest_currentpotential_theta_derivative(cptype, stellsym) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/field/test_fieldline.py b/tests/field/test_fieldline.py index 588f0b1d8..75d01d648 100644 --- a/tests/field/test_fieldline.py +++ b/tests/field/test_fieldline.py @@ -105,6 +105,9 @@ def test_poincare_plot(self): surf=surf) except ImportError: pass + except np.linalg.LinAlgError: + # Matplotlib Bezier extents can fail (eigvals) on some CI platforms + pass def test_poincare_ncsx_known(self): base_curves, base_currents, ma, nfp, bs = get_data("ncsx") diff --git a/tests/field/test_regcoil.py b/tests/field/test_regcoil.py new file mode 100644 index 000000000..58eceb94d --- /dev/null +++ b/tests/field/test_regcoil.py @@ -0,0 +1,941 @@ +import json +import unittest +import warnings +from simsopt.geo import SurfaceRZFourier +from matplotlib import pyplot as plt +import numpy as np +from simsoptpp import WindingSurfaceBn_REGCOIL +from simsopt.field.magneticfieldclasses import WindingSurfaceField +from simsopt.objectives import SquaredFlux +from simsopt.field import CurrentPotentialFourier, CurrentPotentialSolve +from simsopt.util import in_github_actions +from simsopt._core.json import SIMSON, GSONEncoder, GSONDecoder +from scipy.special import ellipk, ellipe +from pathlib import Path +from scipy.io import netcdf_file +np.random.seed(100) + +TEST_DIR = Path(__file__).parent / ".." / "test_files" + + +class Testing(unittest.TestCase): + + def test_windingsurface_exact(self): + """ + Make an infinitesimally thin current loop in the Z = 0 plane + Following approximate analytic solution in Jackson 5.37 for the + vector potential A. From this, we can also check calculations for + B, dA/dX and dB/dX using the WindingSurface class. + """ + nphi = 128 + ntheta = 16 + + # Make winding surface with major radius = 1, no minor radius + winding_surface = SurfaceRZFourier() + winding_surface = winding_surface.from_nphi_ntheta(nphi=nphi, ntheta=ntheta) + for i in range(winding_surface.mpol + 1): + for j in range(-winding_surface.ntor, winding_surface.ntor + 1): + winding_surface.set_rc(i, j, 0.0) + winding_surface.set_zs(i, j, 0.0) + winding_surface.set_rc(0, 0, 1.0) + eps = 1e-12 + winding_surface.set_rc(1, 0, eps) # current loop must have finite width for simsopt + winding_surface.set_zs(1, 0, eps) # current loop must have finite width for simsopt + + # Make CurrentPotential class from this winding surface with 1 amp toroidal current + current_potential = CurrentPotentialFourier(winding_surface, net_poloidal_current_amperes=0, net_toroidal_current_amperes=-1) + + # compute the Bfield from this current loop at some points + Bfield = WindingSurfaceField(current_potential) + N = 2000 + _phi = winding_surface.quadpoints_phi + + # Check that the full expression is correct + points = (np.random.rand(N, 3) - 0.5) * 10 + Bfield.set_points(np.ascontiguousarray(points)) + B_predict = Bfield.B() + dB_predict = Bfield.dB_by_dX() + A_predict = Bfield.A() + _dA_predict = Bfield.dA_by_dX() + + # calculate the Bfield analytically in spherical coordinates + mu_fac = 1e-7 + + # See Jackson 5.37 for the vector potential in terms of the elliptic integrals + r = np.sqrt(points[:, 0] ** 2 + points[:, 1] ** 2 + points[:, 2] ** 2) + theta = np.arctan2(np.sqrt(points[:, 0] ** 2 + points[:, 1] ** 2), points[:, 2]) + k = np.sqrt(4 * r * np.sin(theta) / (1 + r ** 2 + 2 * r * np.sin(theta))) + + # Note scipy is very annoying... scipy function ellipk(k^2) + # is equivalent to what Jackson calls ellipk(k) so call it with k^2 + Aphi = mu_fac * (4 / np.sqrt(1 + r ** 2 + 2 * r * np.sin(theta))) * ((2 - k ** 2) * ellipk(k ** 2) - 2 * ellipe(k ** 2)) / k ** 2 + + # convert A_analytic to Cartesian + Ax = np.zeros(len(Aphi)) + Ay = np.zeros(len(Aphi)) + phi_points = np.arctan2(points[:, 1], points[:, 0]) + theta_points = np.arctan2(np.sqrt(points[:, 0] ** 2 + points[:, 1] ** 2), points[:, 2]) + for i in range(N): + Ax[i] = - np.sin(phi_points[i]) * Aphi[i] + Ay[i] = np.cos(phi_points[i]) * Aphi[i] + A_analytic_elliptic = np.array([Ax, Ay, np.zeros(len(Aphi))]).T + + np.testing.assert_allclose(A_predict, A_analytic_elliptic, rtol=1e-2, atol=1e-12, err_msg="A_predict != A_analytic (near-field elliptic)") + + # now check the Bfield and shape derivatives using the analytic + # expressions that can be derived by hand or found here + # https://ntrs.nasa.gov/citations/20010038494 + C = 4 * mu_fac + alpha2 = 1 + r ** 2 - 2 * r * np.sin(theta) + beta2 = 1 + r ** 2 + 2 * r * np.sin(theta) + k2 = 1 - alpha2 / beta2 + Br = C * np.cos(theta) * ellipe(k2) / (alpha2 * np.sqrt(beta2)) + Btheta = C * ((r ** 2 + np.cos(2 * theta)) * ellipe(k2) - alpha2 * ellipk(k2)) / (2 * alpha2 * np.sqrt(beta2) * np.sin(theta)) + + # convert B_analytic to Cartesian + Bx = np.zeros(len(Br)) + By = np.zeros(len(Br)) + Bz = np.zeros(len(Br)) + for i in range(N): + Bx[i] = np.sin(theta_points[i]) * np.cos(phi_points[i]) * Br[i] + np.cos(theta_points[i]) * np.cos(phi_points[i]) * Btheta[i] + By[i] = np.sin(theta_points[i]) * np.sin(phi_points[i]) * Br[i] + np.cos(theta_points[i]) * np.sin(phi_points[i]) * Btheta[i] + Bz[i] = np.cos(theta_points[i]) * Br[i] - np.sin(theta_points[i]) * Btheta[i] + B_analytic = np.array([Bx, By, Bz]).T + + np.testing.assert_allclose(B_predict, B_analytic, rtol=1e-2, atol=1e-12, err_msg="B_predict != B_analytic (near-field)") + + x = points[:, 0] + y = points[:, 1] + gamma = x ** 2 - y ** 2 + z = points[:, 2] + rho = np.sqrt(points[:, 0] ** 2 + points[:, 1] ** 2) + Bx_dx = C * z * (((- gamma * (3 * z ** 2 + 1) + rho ** 2 * (8 * x ** 2 - y ** 2)) - (rho ** 4 * (5 * x ** 2 + y ** 2) - 2 * rho ** 2 * z ** 2 * (2 * x ** 2 + y ** 2) + 3 * z ** 4 * gamma) - r ** 4 * (2 * x ** 4 + gamma * (y ** 2 + z ** 2))) * ellipe(k2) + (gamma * (1 + 2 * z ** 2) - rho ** 2 * (3 * x ** 2 - 2 * y ** 2) + r ** 2 * (2 * x ** 4 + gamma * (y ** 2 + z ** 2))) * alpha2 * ellipk(k2)) / (2 * alpha2 ** 2 * beta2 ** (3 / 2) * rho ** 4) + + Bx_dy = C * x * y * z * ((3 * (3 * rho ** 2 - 2 * z ** 2) - r ** 4 * (2 * r ** 2 + rho ** 2) - 2 - 2 * (2 * rho ** 4 - rho ** 2 * z ** 2 + 3 * z ** 4)) * ellipe(k2) + (r ** 2 * (2 * r ** 2 + rho ** 2) - (5 * rho ** 2 - 4 * z ** 2) + 2) * alpha2 * ellipk(k2)) / (2 * alpha2 ** 2 * beta2 ** (3 / 2) * rho ** 4) + Bx_dz = C * x * (((rho ** 2 - 1) ** 2 * (rho ** 2 + 1) + 2 * z ** 2 * (1 - 6 * rho ** 2 + rho ** 4) + z ** 4 * (1 + rho ** 2)) * ellipe(k2) - ((rho ** 2 - 1) ** 2 + z ** 2 * (rho ** 2 + 1)) * alpha2 * ellipk(k2)) / (2 * alpha2 ** 2 * beta2 ** (3 / 2) * rho ** 2) + By_dx = Bx_dy + By_dy = C * z * (((gamma * (3 * z ** 2 + 1) + rho ** 2 * (8 * y ** 2 - x ** 2)) - (rho ** 4 * (5 * y ** 2 + x ** 2) - 2 * rho ** 2 * z ** 2 * (2 * y ** 2 + x ** 2) - 3 * z ** 4 * gamma) - r ** 4 * (2 * y ** 4 - gamma * (x ** 2 + z ** 2))) * ellipe(k2) + ((- gamma * (1 + 2 * z ** 2) - rho ** 2 * (3 * y ** 2 - 2 * x ** 2)) + r ** 2 * (2 * y ** 4 - gamma * (x ** 2 + z ** 2))) * alpha2 * ellipk(k2)) / (2 * alpha2 ** 2 * beta2 ** (3 / 2) * rho ** 4) + By_dz = y / x * Bx_dz + Bz_dx = Bx_dz + Bz_dy = By_dz + Bz_dz = C * z * ((6 * (rho ** 2 - z ** 2) - 7 + r ** 4) * ellipe(k2) + alpha2 * (1 - r ** 2) * ellipk(k2)) / (2 * alpha2 ** 2 * beta2 ** (3 / 2)) + dB_analytic = np.transpose(np.array([[Bx_dx, Bx_dy, Bx_dz], + [By_dx, By_dy, By_dz], + [Bz_dx, Bz_dy, Bz_dz]]), [2, 0, 1]) + + np.testing.assert_allclose(dB_predict, dB_analytic, rtol=1e-2, atol=1e-12, err_msg="dB_predict != dB_analytic") + + # Now check that the far-field looks like a dipole + points = (np.random.rand(N, 3) + 1) * 1000 + gamma = winding_surface.gamma().reshape((-1, 3)) + + Bfield.set_points(np.ascontiguousarray(points)) + B_predict = Bfield.B() + A_predict = Bfield.A() + + # calculate the Bfield analytically in spherical coordinates + mu_fac = 1e-7 + + # See Jackson 5.37 for the vector potential in terms of the elliptic integrals + r = np.sqrt(points[:, 0] ** 2 + points[:, 1] ** 2 + points[:, 2] ** 2) + theta = np.arctan2(np.sqrt(points[:, 0] ** 2 + points[:, 1] ** 2), points[:, 2]) + k = np.sqrt(4 * r * np.sin(theta) / (1 + r ** 2 + 2 * r * np.sin(theta))) + Aphi = mu_fac * (4 / np.sqrt(1 + r ** 2 + 2 * r * np.sin(theta))) * ((2 - k ** 2) * ellipk(k ** 2) - 2 * ellipe(k ** 2)) / k ** 2 + + # convert A_analytic to Cartesian + Ax = np.zeros(len(Aphi)) + Ay = np.zeros(len(Aphi)) + phi_points = np.arctan2(points[:, 1], points[:, 0]) + theta_points = np.arctan2(np.sqrt(points[:, 0] ** 2 + points[:, 1] ** 2), points[:, 2]) + for i in range(N): + Ax[i] = - np.sin(phi_points[i]) * Aphi[i] + Ay[i] = np.cos(phi_points[i]) * Aphi[i] + A_analytic_elliptic = np.array([Ax, Ay, np.zeros(len(Aphi))]).T + + np.testing.assert_allclose(A_predict, A_analytic_elliptic, rtol=1e-2, atol=1e-12, err_msg="A_predict != A_analytic (far-field elliptic, pass 1)") + + # Now check that the far-field looks like a dipole + points = (np.random.rand(N, 3) + 1) * 1000 + gamma = winding_surface.gamma().reshape((-1, 3)) + + Bfield.set_points(np.ascontiguousarray(points)) + B_predict = Bfield.B() + A_predict = Bfield.A() + + # calculate the Bfield analytically in spherical coordinates + mu_fac = 1e-7 + + # See Jackson 5.37 for the vector potential in terms of the elliptic integrals + r = np.sqrt(points[:, 0] ** 2 + points[:, 1] ** 2 + points[:, 2] ** 2) + theta = np.arctan2(np.sqrt(points[:, 0] ** 2 + points[:, 1] ** 2), points[:, 2]) + k = np.sqrt(4 * r * np.sin(theta) / (1 + r ** 2 + 2 * r * np.sin(theta))) + + # Note scipy is very annoying... scipy function ellipk(k^2) + # is equivalent to what Jackson calls ellipk(k) so call it with k^2 + Aphi = mu_fac * (4 / np.sqrt(1 + r ** 2 + 2 * r * np.sin(theta))) * ((2 - k ** 2) * ellipk(k ** 2) - 2 * ellipe(k ** 2)) / k ** 2 + + # convert A_analytic to Cartesian + Ax = np.zeros(len(Aphi)) + Ay = np.zeros(len(Aphi)) + phi_points = np.arctan2(points[:, 1], points[:, 0]) + theta_points = np.arctan2(np.sqrt(points[:, 0] ** 2 + points[:, 1] ** 2), points[:, 2]) + for i in range(N): + Ax[i] = - np.sin(phi_points[i]) * Aphi[i] + Ay[i] = np.cos(phi_points[i]) * Aphi[i] + A_analytic_elliptic = np.array([Ax, Ay, np.zeros(len(Aphi))]).T + + np.testing.assert_allclose(A_predict, A_analytic_elliptic, rtol=1e-3, atol=1e-12, err_msg="A_predict != A_analytic (far-field elliptic, pass 2)") + + # double check with vector potential of a dipole + Aphi = np.pi * mu_fac * np.sin(theta) / r ** 2 + + # convert A_analytic to Cartesian + Ax = np.zeros(len(Aphi)) + Ay = np.zeros(len(Aphi)) + for i in range(N): + Ax[i] = - np.sin(phi_points[i]) * Aphi[i] + Ay[i] = np.cos(phi_points[i]) * Aphi[i] + A_analytic = np.array([Ax, Ay, np.zeros(len(Aphi))]).T + + np.testing.assert_allclose(A_predict, A_analytic, rtol=1e-3, atol=1e-12, err_msg="A_predict != A_analytic (far-field dipole)") + + def test_regcoil_K_solve(self): + """ + This function tests the Tikhonov solve with lambda -> infinity + and extensively checks the SIMSOPT grid, currents, solution, etc. + against the REGCOIL variables. + """ + for filename in ['regcoil_out.w7x_infty.nc', 'regcoil_out.li383_infty.nc']: + print(filename) + filename = TEST_DIR / filename + cpst = CurrentPotentialSolve.from_netcdf(filename) + # initialize a solver object for the cp CurrentPotential + s_plasma = cpst.plasma_surface + s_coil = cpst.winding_surface + # Check B and K RHS's -> these are independent of lambda + b_rhs_simsopt, _ = cpst.B_matrix_and_rhs() + k_rhs = cpst.K_rhs() + + for ilambda in range(1, 3): + # Load in big list of variables from REGCOIL to check agree with SIMSOPT + f = netcdf_file(filename, 'r', mmap=False) + Bnormal_regcoil_total = f.variables['Bnormal_total'][()][ilambda, :, :] + Bnormal_from_plasma_current = f.variables['Bnormal_from_plasma_current'][()] + Bnormal_from_net_coil_currents = f.variables['Bnormal_from_net_coil_currents'][()] + r_plasma = f.variables['r_plasma'][()] + r_coil = f.variables['r_coil'][()] + nzeta_plasma = f.variables['nzeta_plasma'][()] + nzeta_coil = f.variables['nzeta_coil'][()] + _ntheta_coil = f.variables['ntheta_coil'][()] + _nfp = f.variables['nfp'][()] + _ntheta_plasma = f.variables['ntheta_plasma'][()] + K2_regcoil = f.variables['K2'][()][ilambda, :, :] + lambda_regcoil = f.variables['lambda'][()][ilambda] + b_rhs_regcoil = f.variables['RHS_B'][()] + k_rhs_regcoil = f.variables['RHS_regularization'][()] + single_valued_current_potential_mn = f.variables['single_valued_current_potential_mn'][()][ilambda, :] + _xm_potential = f.variables['xm_potential'][()] + _xn_potential = f.variables['xn_potential'][()] + _theta_coil = f.variables['theta_coil'][()] + _zeta_coil = f.variables['zeta_coil'][()] + f_B_regcoil = 0.5 * f.variables['chi2_B'][()][ilambda] + f_K_regcoil = 0.5 * f.variables['chi2_K'][()][ilambda] + norm_normal_plasma = f.variables['norm_normal_plasma'][()] + current_potential_thetazeta = f.variables['single_valued_current_potential_thetazeta'][()][ilambda, :, :] + f.close() + Bnormal_single_valued = Bnormal_regcoil_total - Bnormal_from_plasma_current - Bnormal_from_net_coil_currents + print('ilambda index = ', ilambda, lambda_regcoil) + + np.testing.assert_allclose(b_rhs_regcoil, b_rhs_simsopt, rtol=1e-3, atol=1e-12, err_msg="b_rhs_regcoil != b_rhs_simsopt") + np.testing.assert_allclose(k_rhs, k_rhs_regcoil, rtol=1e-3, atol=1e-12, err_msg="k_rhs != k_rhs_regcoil") + + # Compare Bnormal from plasma + np.testing.assert_allclose(cpst.Bnormal_plasma, Bnormal_from_plasma_current.flatten(), rtol=1e-3, atol=1e-12, err_msg="Bnormal_plasma mismatch") + + # Compare optimized dofs + cp = cpst.current_potential + + # when lambda -> infinity, the L1 and L2 regularized problems should agree + optimized_phi_mn_lasso, f_B_lasso, f_K_lasso, _, _ = cpst.solve_lasso(lam=lambda_regcoil) + optimized_phi_mn_lasso_ista, f_B_lasso_ista, f_K_lasso_ista, _, _ = cpst.solve_lasso(lam=lambda_regcoil, acceleration=False, max_iter=5000) + optimized_phi_mn, f_B, f_K = cpst.solve_tikhonov(lam=lambda_regcoil) + np.testing.assert_allclose(single_valued_current_potential_mn, optimized_phi_mn, rtol=1e-3, atol=1e-12, err_msg="single_valued_current_potential_mn != optimized_phi_mn (Tikhonov)") + print(optimized_phi_mn_lasso, optimized_phi_mn) + print(f_B, f_B_lasso, f_B_regcoil) + np.testing.assert_allclose(f_B, f_B_regcoil, rtol=1e-3, atol=1e-12, err_msg="f_B (Tikhonov) != f_B_regcoil") + # assert np.isclose(f_K_lasso, f_K) + np.testing.assert_allclose(optimized_phi_mn_lasso, optimized_phi_mn, rtol=1e-3, atol=1e-12, err_msg="optimized_phi_mn_lasso != optimized_phi_mn") + np.testing.assert_allclose(optimized_phi_mn_lasso_ista, optimized_phi_mn, rtol=1e-3, atol=1e-12, err_msg="optimized_phi_mn_lasso (ISTA) != optimized_phi_mn") + + # Compare plasma surface position + np.testing.assert_allclose(r_plasma[0:nzeta_plasma, :, :], s_plasma.gamma(), rtol=1e-3, atol=1e-12, err_msg="plasma surface position mismatch") + + # Compare plasma surface normal + np.testing.assert_allclose( + norm_normal_plasma[0:nzeta_plasma, :], + np.linalg.norm(s_plasma.normal(), axis=2) / (2 * np.pi * 2 * np.pi), + rtol=1e-3, atol=1e-12, + err_msg="plasma surface normal mismatch" + ) + + # Compare winding surface position + s_coil = cp.winding_surface + np.testing.assert_allclose(r_coil, s_coil.gamma(), rtol=1e-3, atol=1e-12, err_msg="winding surface position mismatch") + + # Compare field from net coil currents + cp_GI = CurrentPotentialFourier.from_netcdf(filename) + Bfield = WindingSurfaceField(cp_GI) + points = s_plasma.gamma().reshape(-1, 3) + Bfield.set_points(points) + B = Bfield.B() + _norm_normal = np.linalg.norm(s_plasma.normal(), axis=2) / (2 * np.pi * 2 * np.pi) + normal = s_plasma.unitnormal().reshape(-1, 3) + B_GI_winding_surface = np.sum(B * normal, axis=1) + np.testing.assert_allclose(B_GI_winding_surface, np.ravel(Bnormal_from_net_coil_currents), rtol=1e-3, atol=1e-12, err_msg="B_GI from WindingSurfaceField != Bnormal_from_net_coil_currents") + np.testing.assert_allclose(cpst.B_GI, np.ravel(Bnormal_from_net_coil_currents), rtol=1e-3, atol=1e-12, err_msg="cpst.B_GI != Bnormal_from_net_coil_currents") + + # Compare single-valued current potential + # Initialization not from netcdf + cp_no_GI = CurrentPotentialFourier( + cp_GI.winding_surface, + net_poloidal_current_amperes=0.0, + net_toroidal_current_amperes=0.0, + mpol=cp_GI.mpol, # critical line here + ntor=cp_GI.ntor, # critical line here + ) + cp_no_GI.set_dofs(optimized_phi_mn) + np.testing.assert_allclose(cp_no_GI.Phi()[0:nzeta_coil, :], current_potential_thetazeta, rtol=1e-3, atol=1e-12, err_msg="single-valued current potential Phi mismatch") + + # Check that f_B from SquaredFlux and f_B from least-squares agree + Bfield_opt = WindingSurfaceField(cp) + Bfield_opt.set_points(s_plasma.gamma().reshape(-1, 3)) + B = Bfield_opt.B() + normal = s_plasma.unitnormal().reshape(-1, 3) + _Bn_opt = np.sum(B * normal, axis=1) + _nfp = cpst.plasma_surface.nfp + nphi = len(cpst.plasma_surface.quadpoints_phi) + ntheta = len(cpst.plasma_surface.quadpoints_theta) + f_B_sq = SquaredFlux( + s_plasma, + Bfield_opt, + -np.ascontiguousarray(cpst.Bnormal_plasma.reshape(nphi, ntheta)) + ).J() + + # These will not exactly agree + # assert np.isclose(f_B, f_B_sq, rtol=1e-2) + + # These will not exactly agree because using different integral discretizations + # assert np.isclose(f_B, f_B_regcoil, rtol=1e-2) + + # These should agree much better + np.testing.assert_allclose(f_B_regcoil, f_B_sq, rtol=1e-3, atol=1e-12, err_msg="f_B_regcoil != f_B from SquaredFlux") + + # Compare current density + cp.set_dofs(optimized_phi_mn) + K = cp.K() + K2 = np.sum(K ** 2, axis=2) + K2_average = np.mean(K2, axis=(0, 1)) + np.testing.assert_allclose(K2[0:nzeta_coil, :] / K2_average, K2_regcoil / K2_average, rtol=1e-3, atol=1e-12, err_msg="K2 mismatch") + + # Compare values of f_K computed in three different ways + normal = s_coil.normal().reshape(-1, 3) + normN = np.linalg.norm(normal, axis=-1) + f_K_direct = 0.5 * np.sum(np.ravel(K2) * normN) / (normal.shape[0]) + # print(f_K_regcoil, f_K_direct, f_K) + np.testing.assert_allclose(f_K_regcoil, f_K_direct, rtol=1e-3, atol=1e-12, err_msg="f_K_regcoil != f_K_direct") + np.testing.assert_allclose(f_K_regcoil, f_K, rtol=1e-3, atol=1e-12, err_msg="f_K_regcoil != f_K (from solve)") + + # Check normal field + Bfield_opt = WindingSurfaceField(cp) + Bfield_opt.set_points(s_plasma.gamma().reshape(-1, 3)) + B_opt = Bfield_opt.B() + normal = s_plasma.unitnormal().reshape(-1, 3) + Bnormal = np.sum(B_opt*normal, axis=1).reshape(np.shape(s_plasma.gamma()[:, :, 0])) + Bnormal_regcoil = Bnormal_regcoil_total - Bnormal_from_plasma_current + np.testing.assert_allclose(np.sum(Bnormal), 0, rtol=1e-3, atol=1e-12, err_msg="sum(Bnormal) != 0") + np.testing.assert_allclose(np.sum(Bnormal_regcoil), 0, rtol=1e-3, atol=1e-12, err_msg="sum(Bnormal_regcoil) != 0") + + # B computed from inductance, i.e. equation A8 in REGCOIL paper + normal_plasma = s_plasma.normal().reshape(-1, 3) + r_plasma = s_plasma.gamma().reshape(-1, 3) + normal_coil = s_coil.normal().reshape(-1, 3) + r_coil = s_coil.gamma().reshape(-1, 3) + rdiff = r_plasma[None, :, :] - r_coil[:, None, :] + rdiff_norm = np.linalg.norm(rdiff, axis=2) + n_dot_nprime = np.sum(normal_plasma[None, :, :] * normal_coil[:, None, :], axis=2) + rdiff_dot_n = np.sum(rdiff * normal_plasma[None, :, :], axis=2) + rdiff_dot_nprime = np.sum(rdiff * normal_coil[:, None, :], axis=2) + inductance_simsopt = (n_dot_nprime / rdiff_norm ** 3 - 3 * rdiff_dot_n * rdiff_dot_nprime / rdiff_norm ** 5) * 1e-7 + dtheta_coil = s_coil.quadpoints_theta[1] + dzeta_coil = s_coil.quadpoints_phi[1] + Bnormal_g = (np.sum(inductance_simsopt * cp.Phi().reshape(-1)[:, None], axis=0) * dtheta_coil * dzeta_coil / np.linalg.norm(normal_plasma, axis=1)).reshape(np.shape(s_plasma.gamma()[:, :, 0])) + + # REGCOIL calculation in c++ + points = s_plasma.gamma().reshape(-1, 3) + normal = s_plasma.normal().reshape(-1, 3) + ws_points = s_coil.gamma().reshape(-1, 3) + ws_normal = s_coil.normal().reshape(-1, 3) + Bnormal_REGCOIL = WindingSurfaceBn_REGCOIL(points, ws_points, ws_normal, cp.Phi(), normal) * dtheta_coil * dzeta_coil + np.testing.assert_allclose(Bnormal_REGCOIL, np.ravel(Bnormal_single_valued), rtol=1e-3, atol=1e-12, err_msg="Bnormal_REGCOIL (C++) != Bnormal_single_valued") + normN = np.linalg.norm(normal, axis=-1) + res = (np.ravel(Bnormal_regcoil_total) ** 2) @ normN + f_B_manual = 0.5 * res / (nphi * ntheta) + np.testing.assert_allclose(f_B_regcoil, f_B_manual, rtol=1e-3, atol=1e-12, err_msg="f_B_regcoil != f_B_manual") + + Bnormal_g += B_GI_winding_surface.reshape(np.shape(s_plasma.gamma()[:, :, 0])) + Bnormal_REGCOIL += B_GI_winding_surface + + # Check that Bnormal calculations using the REGCOIL discretization all agree + np.testing.assert_allclose(np.ravel(Bnormal_g), Bnormal_REGCOIL, rtol=1e-3, atol=1e-12, err_msg="Bnormal_g != Bnormal_REGCOIL") + np.testing.assert_allclose(np.ravel(Bnormal_g), np.ravel(Bnormal_regcoil), rtol=1e-3, atol=1e-12, err_msg="Bnormal_g != Bnormal_regcoil") + + # will be some substantial disagreement here because of the different discretizations, + # although it should improve with higher resolution + # assert np.allclose(Bnormal / np.mean(np.abs(Bnormal_regcoil)), Bnormal_regcoil / np.mean(np.abs(Bnormal_regcoil)), atol=1e-2) + + def test_winding_surface_regcoil(self): + """ + Extensive tests are done for multiple stellarators (stellarator symmetric and + stellarator asymmetric) to verify that REGCOIL, as implemented in SIMSOPT, + agrees with the REGCOIL test file solutions. + """ + for filename in ['regcoil_out.near_axis_asym.nc', 'regcoil_out.near_axis.nc', 'regcoil_out.w7x.nc', 'regcoil_out.li383.nc']: + print(filename) + + # Load big list of variables from REGCOIL to check against SIMSOPT implementation + filename = TEST_DIR / filename + f = netcdf_file(filename, 'r', mmap=False) + Bnormal_regcoil_total = f.variables['Bnormal_total'][()] + Bnormal_from_plasma_current = f.variables['Bnormal_from_plasma_current'][()] + Bnormal_from_net_coil_currents = f.variables['Bnormal_from_net_coil_currents'][()] + r_plasma = f.variables['r_plasma'][()] + r_coil = f.variables['r_coil'][()] + nzeta_plasma = f.variables['nzeta_plasma'][()] + nzeta_coil = f.variables['nzeta_coil'][()] + K2_regcoil = f.variables['K2'][()] + lambda_regcoil = f.variables['lambda'][()] + f_B_regcoil = 0.5 * f.variables['chi2_B'][()] + f_K_regcoil = 0.5 * f.variables['chi2_K'][()] + b_rhs_regcoil = f.variables['RHS_B'][()] + k_rhs_regcoil = f.variables['RHS_regularization'][()] + _xm_potential = f.variables['xm_potential'][()] + _xn_potential = f.variables['xn_potential'][()] + _nfp = f.variables['nfp'][()] + single_valued_current_potential_mn = f.variables['single_valued_current_potential_mn'][()] + norm_normal_plasma = f.variables['norm_normal_plasma'][()] + current_potential_thetazeta = f.variables['single_valued_current_potential_thetazeta'][()] + norm_normal_coil = f.variables['norm_normal_coil'][()] + f.close() + + # Compare K and B RHS's -> these are independent of lambda + cpst = CurrentPotentialSolve.from_netcdf(filename) + + b_rhs_simsopt, _ = cpst.B_matrix_and_rhs() + + np.testing.assert_allclose(b_rhs_regcoil, b_rhs_simsopt, rtol=1e-3, atol=1e-12, err_msg=f"{filename}: b_rhs mismatch") + + k_rhs = cpst.K_rhs() + np.testing.assert_allclose(k_rhs, k_rhs_regcoil, rtol=1e-3, atol=1e-12, err_msg=f"{filename}: k_rhs mismatch") + + # Compare plasma current + np.testing.assert_allclose(cpst.Bnormal_plasma, Bnormal_from_plasma_current.flatten(), rtol=1e-3, atol=1e-12, err_msg=f"{filename}: Bnormal_plasma mismatch") + + # Compare Bnormal from net coil currents + np.testing.assert_allclose(cpst.B_GI, np.ravel(Bnormal_from_net_coil_currents), rtol=1e-3, atol=1e-12, err_msg=f"{filename}: B_GI mismatch") + + cp = cpst.current_potential + s_plasma = cpst.plasma_surface + + s_plasma_full = SurfaceRZFourier( + nfp=s_plasma.nfp, + mpol=s_plasma.mpol, + ntor=s_plasma.ntor, + stellsym=s_plasma.stellsym + ) + s_plasma_full = s_plasma_full.from_nphi_ntheta( + nfp=s_plasma.nfp, ntheta=len(s_plasma.quadpoints_theta), + nphi=len(s_plasma.quadpoints_phi)*s_plasma.nfp, + mpol=s_plasma.mpol, ntor=s_plasma.ntor, + stellsym=s_plasma.stellsym, range="full torus" + ) + s_plasma_full.set_dofs(s_plasma.get_dofs()) + # Compare plasma surface position + np.testing.assert_allclose(r_plasma, s_plasma_full.gamma(), rtol=1e-3, atol=1e-12, err_msg=f"{filename}: plasma surface position mismatch") + + # Compare plasma surface normal + norm_normal_plasma_simsopt = np.linalg.norm(s_plasma_full.normal(), axis=-1) + np.testing.assert_allclose(norm_normal_plasma*2*np.pi*2*np.pi, norm_normal_plasma_simsopt[0:nzeta_plasma, :], rtol=1e-3, atol=1e-12, err_msg=f"{filename}: plasma surface normal mismatch") + + # Compare winding surface position + s_coil = cp.winding_surface + np.testing.assert_allclose(r_coil, s_coil.gamma(), rtol=1e-3, atol=1e-12, err_msg=f"{filename}: winding surface position mismatch") + + # Compare winding surface normal + norm_normal_coil_simsopt = np.linalg.norm(s_coil.normal(), axis=-1) + np.testing.assert_allclose(norm_normal_coil*2*np.pi*2*np.pi, norm_normal_coil_simsopt[0:nzeta_coil, :], rtol=1e-3, atol=1e-12, err_msg=f"{filename}: winding surface normal mismatch") + + # Compare two different ways of computing K() + K = cp.K().reshape(-1, 3) + winding_surface = cp.winding_surface + normal_vec = winding_surface.normal().reshape(-1, 3) + dzeta_coil = (winding_surface.quadpoints_phi[1] - winding_surface.quadpoints_phi[0]) + dtheta_coil = (winding_surface.quadpoints_theta[1] - winding_surface.quadpoints_theta[0]) + normn = np.sqrt(np.sum(normal_vec**2, axis=-1)) # |N| + K_2 = -(cpst.fj @ cp.get_dofs() - cpst.d) / \ + (np.sqrt(dzeta_coil * dtheta_coil) * normn[:, None]) + np.testing.assert_allclose(K, K_2, rtol=1e-3, atol=1e-12, err_msg=f"{filename}: K from cp.K() != K from matrix computation") + + # Compare field from net coil currents + cp_GI = CurrentPotentialFourier.from_netcdf(filename) + Bfield = WindingSurfaceField(cp_GI) + points = s_plasma.gamma().reshape(-1, 3) + Bfield.set_points(points) + B = Bfield.B() + normal = s_plasma.unitnormal().reshape(-1, 3) + B_GI_winding_surface = np.sum(B * normal, axis=1) + np.testing.assert_allclose(B_GI_winding_surface, np.ravel(Bnormal_from_net_coil_currents), rtol=1e-3, atol=1e-12, err_msg=f"{filename}: B_GI from WindingSurfaceField mismatch") + # Make sure single-valued part of current potential is working + cp_no_GI = CurrentPotentialFourier.from_netcdf(filename) + cp_no_GI.set_net_toroidal_current_amperes(0) + cp_no_GI.set_net_poloidal_current_amperes(0) + + # Now loop over all the regularization values in the REGCOIL solution + for i, lambda_reg in enumerate(lambda_regcoil): + + # Set current potential Fourier harmonics from regcoil file + cp.set_current_potential_from_regcoil(filename, i) + + # Compare current potential Fourier harmonics + np.testing.assert_allclose(cp.get_dofs(), single_valued_current_potential_mn[i, :], rtol=1e-3, atol=1e-12, err_msg=f"{filename} lambda[{i}]: current potential dofs mismatch") + + # Compare current density + K = cp.K() + K2 = np.sum(K ** 2, axis=2) + K2_average = np.mean(K2, axis=(0, 1)) + np.testing.assert_allclose(K2[0:nzeta_plasma, :] / K2_average, K2_regcoil[i, :, :] / K2_average, rtol=1e-3, atol=1e-12, err_msg=f"{filename} lambda[{i}]: K2 mismatch") + + f_B_REGCOIL = f_B_regcoil[i] + f_K_REGCOIL = f_K_regcoil[i] + + cp_no_GI.set_current_potential_from_regcoil(filename, i) + + # Compare single-valued current potential + np.testing.assert_allclose(cp_no_GI.Phi()[0:nzeta_plasma, :], current_potential_thetazeta[i, :, :], rtol=1e-3, atol=1e-10, err_msg=f"{filename} lambda[{i}]: single-valued Phi mismatch") + + f_K_direct = 0.5 * np.sum(K2 * norm_normal_coil_simsopt) / (norm_normal_coil_simsopt.shape[0]*norm_normal_coil_simsopt.shape[1]) + np.testing.assert_allclose(f_K_direct/np.abs(f_K_REGCOIL), f_K_REGCOIL/np.abs(f_K_REGCOIL), rtol=1e-3, atol=1e-10, err_msg=f"{filename} lambda[{i}]: f_K_direct/|f_K_REGCOIL| != f_K_REGCOIL/|f_K_REGCOIL|, got {f_K_direct} vs {f_K_REGCOIL}") + + normal = s_plasma.unitnormal().reshape(-1, 3) + norm_normal_plasma_simsopt = np.linalg.norm(s_plasma.normal(), axis=-1) + Bnormal_regcoil = Bnormal_regcoil_total[i, :, :] - Bnormal_from_plasma_current + + # REGCOIL calculation in c++ + points = s_plasma.gamma().reshape(-1, 3) + normal = s_plasma.normal().reshape(-1, 3) + ws_points = s_coil.gamma().reshape(-1, 3) + ws_normal = s_coil.normal().reshape(-1, 3) + dtheta_coil = s_coil.quadpoints_theta[1] + dzeta_coil = s_coil.quadpoints_phi[1] + Bnormal = WindingSurfaceBn_REGCOIL(points, ws_points, ws_normal, cp.Phi(), normal) * dtheta_coil * dzeta_coil + Bnormal += cpst.B_GI + Bnormal = Bnormal.reshape(Bnormal_regcoil.shape) + + np.testing.assert_allclose(Bnormal, Bnormal_regcoil, rtol=1e-3, atol=1e-12, err_msg=f"{filename} lambda[{i}]: Bnormal (C++) mismatch") + + # check Bnormal and Bnormal_regcoil integrate over the surface to zero + # This is only true in the stellarator symmetric case! + if s_plasma.stellsym: + np.testing.assert_allclose(np.sum(Bnormal*norm_normal_plasma_simsopt), 0, rtol=1e-3, atol=1e-10, err_msg=f"{filename} lambda[{i}]: sum(Bnormal*norm) != 0") + np.testing.assert_allclose(np.sum(Bnormal_regcoil*norm_normal_plasma_simsopt[0:nzeta_plasma, :]), 0, rtol=1e-3, atol=1e-10, err_msg=f"{filename} lambda[{i}]: sum(Bnormal_regcoil*norm) != 0") + + # Check that L1 optimization agrees if lambda = 0 + # With lambda=0, FISTA converges slowly (ill-conditioned); need many iterations + if lambda_reg == 0.0: + optimized_phi_mn_lasso, f_B_lasso, _, _, _ = cpst.solve_lasso(lam=lambda_reg, max_iter=10000, acceleration=True) + optimized_phi_mn_lasso_ista, f_B_lasso_ista, _, _, _ = cpst.solve_lasso(lam=lambda_reg, max_iter=10000, acceleration=False) + + # Check the optimization in SIMSOPT is working + optimized_phi_mn, f_B, f_K = cpst.solve_tikhonov(lam=lambda_reg) + np.testing.assert_allclose(single_valued_current_potential_mn[i, :], optimized_phi_mn, rtol=1e-3, atol=1e-12, err_msg=f"{filename} lambda[{i}]: Tikhonov optimized_phi_mn mismatch") + if lambda_reg == 0.0: + np.testing.assert_allclose(f_B, f_B_lasso, rtol=1e-3, atol=1e-12, err_msg=f"{filename} lambda[{i}]: f_B (Tikhonov) != f_B (Lasso) at lambda=0") + # np.testing.assert_allclose(f_B, f_B_lasso_ista, rtol=1e-1, atol=1e-12, err_msg=f"{filename} lambda[{i}]: f_B (Tikhonov) != f_B (Lasso ISTA) at lambda=0") + + # Check f_B from SquaredFlux and f_B from least-squares agree + Bfield_opt = WindingSurfaceField(cpst.current_potential) + Bfield_opt.set_points(s_plasma.gamma().reshape(-1, 3)) + _nfp = cpst.plasma_surface.nfp + nphi = len(cpst.plasma_surface.quadpoints_phi) + ntheta = len(cpst.plasma_surface.quadpoints_theta) + _f_B_sq = SquaredFlux( + s_plasma, + Bfield_opt, + -np.ascontiguousarray(cpst.Bnormal_plasma.reshape(nphi, ntheta)) + ).J() + + # These do not agree well when lambda >> 1 + # or other situations where the exact plasma surface + # locations are critical, so the REGCOIL Bnormal + # calculation must be used + #print(f_B, f_B_sq) + #assert np.isclose(f_B, f_B_sq, rtol=1e-1) + np.testing.assert_allclose(f_B, f_B_REGCOIL, rtol=1e-3, atol=1e-12, err_msg=f"{filename} lambda[{i}]: f_B != f_B_REGCOIL") + np.testing.assert_allclose(f_K, f_K_REGCOIL, rtol=1e-3, atol=1e-12, err_msg=f"{filename} lambda[{i}]: f_K != f_K_REGCOIL") + + # check the REGCOIL Bnormal calculation in c++ """ + points = s_plasma.gamma().reshape(-1, 3) + normal = s_plasma.normal().reshape(-1, 3) + ws_points = s_coil.gamma().reshape(-1, 3) + ws_normal = s_coil.normal().reshape(-1, 3) + dtheta_coil = s_coil.quadpoints_theta[1] + dzeta_coil = s_coil.quadpoints_phi[1] + Bnormal_REGCOIL = WindingSurfaceBn_REGCOIL(points, ws_points, ws_normal, cp.Phi(), normal) * dtheta_coil * dzeta_coil + + normN = np.linalg.norm(normal, axis=-1) + res = (np.ravel(Bnormal_regcoil_total[i, :, :]) ** 2) @ normN + f_B_manual = 0.5 * res / (nphi * ntheta) + np.testing.assert_allclose(f_B_REGCOIL, f_B_manual, rtol=1e-3, atol=1e-12, err_msg=f"{filename} lambda[{i}]: f_B_REGCOIL != f_B_manual") + + Bnormal_REGCOIL += B_GI_winding_surface + np.testing.assert_allclose(Bnormal_REGCOIL, np.ravel(Bnormal_regcoil), rtol=1e-3, atol=1e-12, err_msg=f"{filename} lambda[{i}]: Bnormal_REGCOIL (C++) != Bnormal_regcoil") + + def test_Bnormal_interpolation_from_netcdf(self): + """ + Test the branch that interpolates Bnormal_from_plasma when increasing + plasma surface resolution via from_netcdf(plasma_ntheta_res > 1 or plasma_nzeta_res > 1). + """ + filename = TEST_DIR / 'regcoil_out.w7x_infty.nc' + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + cpst = CurrentPotentialSolve.from_netcdf( + filename, plasma_ntheta_res=2.0, plasma_nzeta_res=1.0 + ) + self.assertGreater(len(w), 0, msg="Expected interpolation accuracy warning") + self.assertTrue(any("interpolated" in str(warning.message).lower() for warning in w)) + + # Bnormal_plasma should have shape matching higher-resolution grid + nzeta = cpst.nzeta_plasma + ntheta = cpst.ntheta_plasma + self.assertEqual(len(cpst.Bnormal_plasma), nzeta * ntheta, + msg="Bnormal_plasma length should match nzeta*ntheta") + + # solve_tikhonov and B_matrix_and_rhs should work + b_rhs, B_matrix = cpst.B_matrix_and_rhs() + self.assertEqual(len(b_rhs), cpst.ndofs) + optimized_phi_mn, f_B, f_K = cpst.solve_tikhonov(lam=1e-6) + self.assertEqual(len(optimized_phi_mn), cpst.ndofs) + + def test_Bnormal_interpolation_plasma_nzeta_res(self): + """Cover plasma_nzeta_res > 1 branch (plasma_ntheta_res=1, plasma_nzeta_res=2).""" + filename = TEST_DIR / 'regcoil_out.w7x_infty.nc' + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + cpst = CurrentPotentialSolve.from_netcdf( + filename, plasma_ntheta_res=1.0, plasma_nzeta_res=2.0 + ) + self.assertEqual(len(cpst.Bnormal_plasma), cpst.nzeta_plasma * cpst.ntheta_plasma) + optimized_phi_mn, f_B, f_K = cpst.solve_tikhonov(lam=1e-6) + self.assertEqual(len(optimized_phi_mn), cpst.ndofs) + + def test_CurrentPotentialSolve_Bnormal_plasma_shape_mismatch(self): + """CurrentPotentialSolve raises ValueError when Bnormal_plasma shape mismatches.""" + cp = CurrentPotentialFourier.from_netcdf(TEST_DIR / 'regcoil_out.w7x_infty.nc') + s_plasma = SurfaceRZFourier( + nfp=cp.nfp, mpol=4, ntor=4, stellsym=True + ).from_nphi_ntheta(nfp=cp.nfp, ntheta=32, nphi=32, mpol=4, ntor=4, stellsym=True, range="field period") + # Bnormal with wrong size (e.g. 10 elements instead of 32*32) + bad_Bnormal = np.ones(10) + with self.assertRaises(ValueError) as cm: + CurrentPotentialSolve(cp, s_plasma, bad_Bnormal) + self.assertIn("shape", str(cm.exception).lower()) + + def test_WindingSurfaceField_as_dict_from_dict(self): + """Test WindingSurfaceField serialization via as_dict and from_dict. + + Uses surfaces from winding_surface_test.json (created via set_dofs so _dofs + are in sync) since SurfaceRZFourier from from_netcdf uses set_rc/set_zs + which does not sync _dofs for serialization. + """ + from simsopt import load + winding_surface, _ = load(TEST_DIR / 'winding_surface_test.json') + cp = CurrentPotentialFourier( + winding_surface, mpol=4, ntor=4, + net_poloidal_current_amperes=11884578.094260072, + net_toroidal_current_amperes=0, + stellsym=True) + cp.set_dofs(np.array([ + 235217.63668779, -700001.94517193, 1967024.36417348, + -1454861.01406576, -1021274.81793687, 1657892.17597651, + -784146.17389912, 136356.84602536, -670034.60060171, + 194549.6432583, 1006169.72177152, -1677003.74430119, + 1750470.54137804, 471941.14387043, -1183493.44552104, + 1046707.62318593, -334620.59690486, 658491.14959397, + -1169799.54944824, -724954.843765, 1143998.37816758, + -2169655.54190455, -106677.43308896, 761983.72021537, + -986348.57384563, 532788.64040937, -600463.7957275, + 1471477.22666607, 1009422.80860728, -2000273.40765417, + 2179458.3105468, -55263.14222144, -315581.96056445, + 587702.35409154, -637943.82177418, 609495.69135857, + -1050960.33686344, -970819.1808181, 1467168.09965404, + -198308.0580687 + ])) + bfield = WindingSurfaceField(cp) + points = np.ascontiguousarray(np.random.RandomState(42).rand(20, 3) * 2) + bfield.set_points(points) + B_orig = bfield.B() + A_orig = bfield.A() + + field_json_str = json.dumps(SIMSON(bfield), cls=GSONEncoder) + bfield_regen = json.loads(field_json_str, cls=GSONDecoder) + bfield_regen.set_points(points) + np.testing.assert_allclose(bfield_regen.B(), B_orig, rtol=1e-3, atol=1e-12, + err_msg="WindingSurfaceField B() mismatch after load") + np.testing.assert_allclose(bfield_regen.A(), A_orig, rtol=1e-3, atol=1e-12, + err_msg="WindingSurfaceField A() mismatch after load") + + def test_K_calculations(self): + from simsopt import load + winding_surface, plasma_surface = load(TEST_DIR / 'winding_surface_test.json') + cp = CurrentPotentialFourier( + winding_surface, mpol=4, ntor=4, + net_poloidal_current_amperes=11884578.094260072, + net_toroidal_current_amperes=0, + stellsym=True) + cp.set_dofs(np.array([ + 235217.63668779, -700001.94517193, 1967024.36417348, + -1454861.01406576, -1021274.81793687, 1657892.17597651, + -784146.17389912, 136356.84602536, -670034.60060171, + 194549.6432583, 1006169.72177152, -1677003.74430119, + 1750470.54137804, 471941.14387043, -1183493.44552104, + 1046707.62318593, -334620.59690486, 658491.14959397, + -1169799.54944824, -724954.843765, 1143998.37816758, + -2169655.54190455, -106677.43308896, 761983.72021537, + -986348.57384563, 532788.64040937, -600463.7957275, + 1471477.22666607, 1009422.80860728, -2000273.40765417, + 2179458.3105468, -55263.14222144, -315581.96056445, + 587702.35409154, -637943.82177418, 609495.69135857, + -1050960.33686344, -970819.1808181, 1467168.09965404, + -198308.0580687 + ])) + cpst = CurrentPotentialSolve(cp, plasma_surface, np.zeros(1024)) + np.testing.assert_allclose(cpst.current_potential.get_dofs(), cp.get_dofs(), rtol=1e-3, atol=1e-12, err_msg="CurrentPotentialSolve dofs != original cp dofs") + # Pre-compute some important matrices + cpst.B_matrix_and_rhs() + + # Copied over from the packaged grid K operator. + winding_surface = cp.winding_surface + normal_vec = winding_surface.normal() + normn = np.sqrt(np.sum(normal_vec**2, axis=-1)) # |N| + + test_K_1 = ( + cp.winding_surface.gammadash2() + * (cp.Phidash1()+cp.net_poloidal_current_amperes)[:, :, None] + - cp.winding_surface.gammadash1() + * (cp.Phidash2()+cp.net_toroidal_current_amperes)[:, :, None])/normn[:, :, None] + + test_K_3 = cp.K() + + normn = normn.reshape(-1) + dzeta_coil = (winding_surface.quadpoints_phi[1] - winding_surface.quadpoints_phi[0]) + dtheta_coil = (winding_surface.quadpoints_theta[1] - winding_surface.quadpoints_theta[0]) + + # Notice Equation A.13 for the current in Matt L's regcoil paper has factor of 1/nnorm in it + # But cpst.fj and cpst.d have factor of only 1/sqrt(normn) + test_K_2 = -(cpst.fj @ cp.get_dofs() - cpst.d) / \ + (np.sqrt(dzeta_coil * dtheta_coil) * normn[:, None]) + nzeta_coil = cpst.nzeta_coil + test_K_2 = test_K_2.reshape(nzeta_coil, nzeta_coil // cp.nfp, 3) + + # Figure 1: Surface current density K (A/m) from three equivalent formulations. + # Row 1: K from analytic formula (Eq. A.13 REGCOIL). + # Row 2: K from matrix computation (fj @ phi - d) / (sqrt(dzeta*dtheta) * |N|). + # Row 3: K from cp.K() C++ implementation. + # Columns: x, y, z components of K on (zeta, theta) winding surface grid. + fig1, axes1 = plt.subplots(3, 3, figsize=(12, 10), squeeze=True) + fig1.suptitle(r'Surface current density $\mathbf{K}$ (A/m): comparison of three formulations', + fontsize=12) + for j in range(3): + ax = axes1[0, j] + im = ax.pcolor(test_K_1[:, :, j]) + fig1.colorbar(im, ax=ax) + ax.set_title(r'Analytic: $K_{}$'.format('xyz'[j])) + if j == 0: + ax.set_ylabel(r'Analytic (Eq. A.13)') + for j in range(3): + ax = axes1[1, j] + im = ax.pcolor(test_K_2[:, :, j]) + fig1.colorbar(im, ax=ax) + ax.set_title(r'Matrix: $K_{}$'.format('xyz'[j])) + if j == 0: + ax.set_ylabel(r'Matrix (fj@phi-d)') + for j in range(3): + ax = axes1[2, j] + im = ax.pcolor(test_K_3[:, :, j]) + fig1.colorbar(im, ax=ax) + ax.set_title(r'cp.K(): $K_{}$'.format('xyz'[j])) + ax.set_xlabel(r'$\theta$ (poloidal)') + if j == 0: + ax.set_ylabel(r'cp.K() (C++)') + fig1.tight_layout() + + # Figure 2: Differences between formulations (should be ~0 if implementations agree). + # Row 1: analytic - cp.K() for x, y, z. + # Row 2: matrix - cp.K() for x, y, z. + fig2, axes2 = plt.subplots(2, 3, figsize=(12, 8), squeeze=True) + fig2.suptitle(r'Difference in $\mathbf{K}$: analytic vs cp.K() and matrix vs cp.K() ' + r'(should be ~0)', fontsize=12) + # Row 0: analytic - cp.K() + for j, comp in enumerate(['$K_x$', '$K_y$', '$K_z$']): + ax = axes2[0, j] + im = ax.pcolor(test_K_1[:, :, j] - test_K_3[:, :, j]) + fig2.colorbar(im, ax=ax) + ax.set_title(r'Analytic $-$ cp.K(): ' + comp) + if j == 0: + ax.set_ylabel(r'Analytic $-$ cp.K()') + # Row 1: matrix - cp.K() + for j, comp in enumerate(['$K_x$', '$K_y$', '$K_z$']): + ax = axes2[1, j] + im = ax.pcolor(test_K_2[:, :, j] - test_K_3[:, :, j]) + fig2.colorbar(im, ax=ax) + ax.set_title(r'Matrix $-$ cp.K(): ' + comp) + ax.set_xlabel(r'$\theta$ (poloidal)') + if j == 0: + ax.set_ylabel(r'Matrix $-$ cp.K()') + fig2.tight_layout() + if not in_github_actions: + plt.show() + + np.testing.assert_allclose(test_K_1, test_K_2, rtol=1e-3, atol=1e-12, err_msg="K from analytic formula != K from matrix computation") + np.testing.assert_allclose(test_K_1, test_K_3, rtol=1e-3, atol=1e-12, err_msg="K from analytic formula != K from cp.K()") + + def test_regcoil_write(self): + """ + This function tests the SIMSOPT routine that writes + REGCOIL outfiles for backwards compatability. + """ + for fname in ['regcoil_out.w7x_infty.nc', 'regcoil_out.li383_infty.nc', 'regcoil_out.near_axis_asym.nc', 'regcoil_out.near_axis.nc', 'regcoil_out.w7x.nc', 'regcoil_out.li383.nc']: + filename = TEST_DIR / fname + cpst = CurrentPotentialSolve.from_netcdf(filename) + + f = netcdf_file(filename, 'r', mmap=False) + n_lambda = f.variables['Bnormal_total'][()].shape[0] + f.close() + # Run solve for at least one lambda (need ilambdas_l2/l1 non-empty for write_regcoil_out) + ilambda_range = range(max(0, min(1, n_lambda - 1)), min(3, n_lambda)) + for ilambda in ilambda_range: + # Load in big list of variables from REGCOIL to check agree with SIMSOPT + f = netcdf_file(filename, 'r', mmap=False) + Bnormal_regcoil_total = f.variables['Bnormal_total'][()][ilambda, :, :] + Bnormal_from_plasma_current = f.variables['Bnormal_from_plasma_current'][()] + Bnormal_from_net_coil_currents = f.variables['Bnormal_from_net_coil_currents'][()] + r_plasma = f.variables['r_plasma'][()] + r_coil = f.variables['r_coil'][()] + nzeta_plasma = f.variables['nzeta_plasma'][()] + nzeta_coil = f.variables['nzeta_coil'][()] + ntheta_coil = f.variables['ntheta_coil'][()] + nfp = f.variables['nfp'][()] + _stellsym = f.variables['symmetry_option'][()] + ntheta_plasma = f.variables['ntheta_plasma'][()] + K2_regcoil = f.variables['K2'][()][ilambda, :, :] + lambda_regcoil = f.variables['lambda'][()][ilambda] + b_rhs_regcoil = f.variables['RHS_B'][()] + k_rhs_regcoil = f.variables['RHS_regularization'][()] + single_valued_current_potential_mn = f.variables['single_valued_current_potential_mn'][()][ilambda, :] + xm_potential = f.variables['xm_potential'][()] + xn_potential = f.variables['xn_potential'][()] + theta_coil = f.variables['theta_coil'][()] + zeta_coil = f.variables['zeta_coil'][()] + f_B_regcoil = 0.5 * f.variables['chi2_B'][()][ilambda] + f_K_regcoil = 0.5 * f.variables['chi2_K'][()][ilambda] + norm_normal_plasma = f.variables['norm_normal_plasma'][()] + current_potential_thetazeta = f.variables['single_valued_current_potential_thetazeta'][()][ilambda, :, :] + f.close() + + # Compare optimized dofs + _cp = cpst.current_potential + + # only when lambda -> infinity or lambda -> 0, the L1 and L2 regularized problems should agree + _, _, _, _, _ = cpst.solve_lasso(lam=lambda_regcoil) + optimized_phi_mn, _, _ = cpst.solve_tikhonov(lam=lambda_regcoil) + np.testing.assert_allclose(single_valued_current_potential_mn, optimized_phi_mn, rtol=1e-2, atol=1e-12, err_msg=f"{fname} ilambda={ilambda}: phi_mn (Tikhonov) mismatch") + # np.testing.assert_allclose(f_B_lasso, f_B, rtol=1e-2, atol=1e-12, err_msg=f"{fname} ilambda={ilambda}: f_B (Lasso) != f_B (Tikhonov)") + # np.testing.assert_allclose(optimized_phi_mn_lasso, optimized_phi_mn, rtol=1e-2, atol=1e-12, err_msg=f"{fname} ilambda={ilambda}: phi_mn (Lasso) != phi_mn (Tikhonov)") + + # Test that current potential solve class correctly writes REGCOIL outfiles + cpst.write_regcoil_out(filename='simsopt_' + fname) + g = netcdf_file('simsopt_' + fname, 'r', mmap=False) + f = netcdf_file(filename, 'r', mmap=False) + n_lambda_ref = f.variables['Bnormal_total'][()].shape[0] + n_lambda_written = g.variables['Bnormal_total'][()].shape[0] + for ilambda in range(min(2, n_lambda_written, n_lambda_ref - 1)): + print(filename, ilambda) + Bnormal_regcoil_total = f.variables['Bnormal_total'][()][ilambda + 1, :, :] + Bnormal_from_plasma_current = f.variables['Bnormal_from_plasma_current'][()] + Bnormal_from_net_coil_currents = f.variables['Bnormal_from_net_coil_currents'][()] + r_plasma = f.variables['r_plasma'][()] + r_coil = f.variables['r_coil'][()] + nzeta_plasma = f.variables['nzeta_plasma'][()] + nzeta_coil = f.variables['nzeta_coil'][()] + ntheta_coil = f.variables['ntheta_coil'][()] + nfp = f.variables['nfp'][()] + ntheta_plasma = f.variables['ntheta_plasma'][()] + K2_regcoil = f.variables['K2'][()][ilambda + 1, :, :] + b_rhs_regcoil = f.variables['RHS_B'][()] + k_rhs_regcoil = f.variables['RHS_regularization'][()] + lambda_regcoil = f.variables['lambda'][()][ilambda + 1] + b_rhs_regcoil = f.variables['RHS_B'][()] + k_rhs_regcoil = f.variables['RHS_regularization'][()] + single_valued_current_potential_mn = f.variables['single_valued_current_potential_mn'][()][ilambda + 1, :] + xm_plasma = f.variables['xm_plasma'][()] + xn_plasma = f.variables['xn_plasma'][()] + xm_coil = f.variables['xm_coil'][()] + _xn_coil = f.variables['xn_coil'][()] + xm_potential = f.variables['xm_potential'][()] + xn_potential = f.variables['xn_potential'][()] + theta_coil = f.variables['theta_coil'][()] + zeta_coil = f.variables['zeta_coil'][()] + f_B_regcoil = f.variables['chi2_B'][()][ilambda + 1] + f_K_regcoil = f.variables['chi2_K'][()][ilambda + 1] + norm_normal_plasma = f.variables['norm_normal_plasma'][()] + current_potential_thetazeta = f.variables['single_valued_current_potential_thetazeta'][()][ilambda + 1, :, :] + np.testing.assert_allclose(single_valued_current_potential_mn, g.variables['single_valued_current_potential_mn'][()][ilambda, :], atol=1e-12, err_msg=f"{fname} ilambda={ilambda}: written phi_mn mismatch") + np.testing.assert_allclose(Bnormal_regcoil_total, g.variables['Bnormal_total'][()][ilambda, :, :], atol=1e-12, err_msg=f"{fname} ilambda={ilambda}: written Bnormal_total mismatch") + np.testing.assert_allclose(Bnormal_from_plasma_current, g.variables['Bnormal_from_plasma_current'][()], atol=1e-12, err_msg=f"{fname} ilambda={ilambda}: written Bnormal_from_plasma_current mismatch") + np.testing.assert_allclose(Bnormal_from_net_coil_currents, g.variables['Bnormal_from_net_coil_currents'][()], atol=1e-12, err_msg=f"{fname} ilambda={ilambda}: written Bnormal_from_net_coil_currents mismatch") + np.testing.assert_allclose(nzeta_plasma, g.variables['nzeta_plasma'][()], atol=1e-12, err_msg=f"{fname}: written nzeta_plasma mismatch") + np.testing.assert_allclose(nzeta_coil, g.variables['nzeta_coil'][()], atol=1e-12, err_msg=f"{fname}: written nzeta_coil mismatch") + np.testing.assert_allclose(ntheta_coil, g.variables['ntheta_coil'][()], atol=1e-12, err_msg=f"{fname}: written ntheta_coil mismatch") + np.testing.assert_allclose(nfp, g.variables['nfp'][()], atol=1e-12, err_msg=f"{fname}: written nfp mismatch") + np.testing.assert_allclose(ntheta_plasma, g.variables['ntheta_plasma'][()], atol=1e-12, err_msg=f"{fname}: written ntheta_plasma mismatch") + np.testing.assert_allclose(xm_plasma, g.variables['xm_plasma'][()], atol=1e-12, err_msg=f"{fname}: written xm_plasma mismatch") + np.testing.assert_allclose(xn_plasma, g.variables['xn_plasma'][()], atol=1e-12, err_msg=f"{fname}: written xn_plasma mismatch") + np.testing.assert_allclose(xm_coil, g.variables['xm_coil'][()], atol=1e-12, err_msg=f"{fname}: written xm_coil mismatch") + np.testing.assert_allclose(xm_potential, g.variables['xm_potential'][()], atol=1e-12, err_msg=f"{fname}: written xm_potential mismatch") + np.testing.assert_allclose(xn_potential, g.variables['xn_potential'][()], atol=1e-12, err_msg=f"{fname}: written xn_potential mismatch") + np.testing.assert_allclose(theta_coil, g.variables['theta_coil'][()], atol=1e-12, err_msg=f"{fname}: written theta_coil mismatch") + np.testing.assert_allclose(zeta_coil, g.variables['zeta_coil'][()], atol=1e-12, err_msg=f"{fname}: written zeta_coil mismatch") + np.testing.assert_allclose(r_coil, g.variables['r_coil'][()], atol=1e-12, err_msg=f"{fname}: written r_coil mismatch") + np.testing.assert_allclose(r_plasma, g.variables['r_plasma'][()], atol=1e-12, err_msg=f"{fname}: written r_plasma mismatch") + np.testing.assert_allclose(K2_regcoil, g.variables['K2'][()][ilambda, :, :], atol=1e-12, err_msg=f"{fname} ilambda={ilambda}: written K2 mismatch") + assert (K2_regcoil.shape == g.variables['K2_l1'][()][ilambda, :, :].shape), f"{fname} ilambda={ilambda}: K2_l1 shape mismatch: {K2_regcoil.shape} vs {g.variables['K2_l1'][()][ilambda, :, :].shape}" + def _lambda_close(a, b): + a, b = np.float64(a), np.float64(b) + if np.isfinite(a) and np.isfinite(b): + return np.isclose(a, b, atol=1e-12) + return (a > 1e99 or np.isposinf(a)) and (b > 1e99 or np.isposinf(b)) + self.assertTrue(_lambda_close(lambda_regcoil, g.variables['lambda'][()][ilambda]), msg=f"{fname} ilambda={ilambda}: written lambda mismatch") + self.assertTrue(_lambda_close(lambda_regcoil, g.variables['lambda_l1'][()][ilambda]), msg=f"{fname} ilambda={ilambda}: written lambda_l1 mismatch") + np.testing.assert_allclose(b_rhs_regcoil, g.variables['RHS_B'][()], atol=1e-12, err_msg=f"{fname}: written RHS_B mismatch") + np.testing.assert_allclose(k_rhs_regcoil, g.variables['RHS_regularization'][()], atol=1e-12, err_msg=f"{fname}: written RHS_regularization mismatch") + np.testing.assert_allclose(f_B_regcoil, g.variables['chi2_B'][()][ilambda], atol=1e-12, err_msg=f"{fname} ilambda={ilambda}: written chi2_B mismatch") + # chi2_B_l1 (L1) can differ significantly from chi2_B (L2) - reference files may only have L2 + np.testing.assert_allclose(f_B_regcoil, g.variables['chi2_B_l1'][()][ilambda], rtol=1e1, atol=1.0, err_msg=f"{fname} ilambda={ilambda}: written chi2_B_l1 mismatch") + np.testing.assert_allclose(f_K_regcoil, g.variables['chi2_K'][()][ilambda], rtol=1e-2, atol=1e-7, err_msg=f"{fname} ilambda={ilambda}: written chi2_K mismatch") + np.testing.assert_allclose(norm_normal_plasma, g.variables['norm_normal_plasma'][()], atol=1e-12, err_msg=f"{fname}: written norm_normal_plasma mismatch") + np.testing.assert_allclose(current_potential_thetazeta, g.variables['single_valued_current_potential_thetazeta'][()][ilambda, :, :], atol=1e-6, err_msg=f"{fname} ilambda={ilambda}: written current_potential_thetazeta mismatch") + # np.testing.assert_allclose(current_potential_thetazeta, g.variables['single_valued_current_potential_thetazeta_l1'][()][ilambda, :, :], atol=1e-12, err_msg=f"{fname} ilambda={ilambda}: written current_potential_thetazeta_l1 mismatch") + g.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_files/regcoil_out.hsx.nc b/tests/test_files/regcoil_out.hsx.nc new file mode 100644 index 000000000..580edf3fb Binary files /dev/null and b/tests/test_files/regcoil_out.hsx.nc differ diff --git a/tests/test_files/regcoil_out.li383.nc b/tests/test_files/regcoil_out.li383.nc new file mode 100644 index 000000000..47e84bbcc Binary files /dev/null and b/tests/test_files/regcoil_out.li383.nc differ diff --git a/tests/test_files/regcoil_out.li383_infty.nc b/tests/test_files/regcoil_out.li383_infty.nc new file mode 100644 index 000000000..4a489196b Binary files /dev/null and b/tests/test_files/regcoil_out.li383_infty.nc differ diff --git a/tests/test_files/regcoil_out.near_axis.nc b/tests/test_files/regcoil_out.near_axis.nc new file mode 100644 index 000000000..985273f93 Binary files /dev/null and b/tests/test_files/regcoil_out.near_axis.nc differ diff --git a/tests/test_files/regcoil_out.near_axis_asym.nc b/tests/test_files/regcoil_out.near_axis_asym.nc new file mode 100644 index 000000000..8d92f5b81 Binary files /dev/null and b/tests/test_files/regcoil_out.near_axis_asym.nc differ diff --git a/tests/test_files/regcoil_out.w7x.nc b/tests/test_files/regcoil_out.w7x.nc new file mode 100644 index 000000000..6d3c63157 Binary files /dev/null and b/tests/test_files/regcoil_out.w7x.nc differ diff --git a/tests/test_files/regcoil_out.w7x_infty.nc b/tests/test_files/regcoil_out.w7x_infty.nc new file mode 100644 index 000000000..e6911aa5e Binary files /dev/null and b/tests/test_files/regcoil_out.w7x_infty.nc differ diff --git a/tests/test_files/winding_surface_test.json b/tests/test_files/winding_surface_test.json new file mode 100644 index 000000000..5ae800824 --- /dev/null +++ b/tests/test_files/winding_surface_test.json @@ -0,0 +1,2207 @@ +{ + "@module": "simsopt._core.json", + "@class": "SIMSON", + "@version": "0.19.0.post217+g54251e8b.d20240725.dirty", + "graph": [ + { + "$type": "ref", + "value": "SurfaceRZFourier7" + }, + { + "$type": "ref", + "value": "SurfaceRZFourier4" + } + ], + "simsopt_objs": { + "5807918224": { + "@module": "simsopt._core.optimizable", + "@class": "DOFs", + "@name": "5807918224", + "@version": "0.19.0.post217+g54251e8b.d20240725.dirty", + "x": { + "@module": "numpy", + "@class": "array", + "dtype": "float64", + "data": [ + 1.3991480988737919, + 0.008413011404644715, + 0.002275465937287521, + 0.0025527743567711485, + -0.00342678545600683, + -0.0008118092202128168, + -0.0012814373279649576, + 0.003972349462239593, + 0.005464298389595823, + -0.01014609558331318, + 0.0769425757897101, + 0.9721236896786388, + -0.1012566811325677, + 0.006440310384220321, + -0.007695991643153007, + -0.0026193978221245944, + 0.0036462234650491395, + -0.00040468456419894574, + -0.0008786819403620429, + -0.0009598311085319094, + 0.0021560451918240255, + 0.000646442903933475, + 0.005328893082357635, + 0.022344562448847003, + 0.01616329830663549, + 0.001843319068373139, + 0.00031736197400336147, + 0.0007384352695652626, + 0.0001177333913004869, + 0.0007428149479712462, + 0.0013745907572829633, + -0.0033845300588809042, + 0.0021877155261840185, + -0.010955471535789175, + -0.028559616532423535, + 0.004639746735344535, + -0.0065274863344688275, + 0.0014938337841373217, + 0.0009376015869931356, + -9.841977659788292e-06, + -0.000481808272668778, + -0.0009002894720259349, + 0.0020703030378239197, + -0.00041276080705539997, + -0.0023890646837338853, + 0.010324785086546152, + 0.003219750646770263, + -0.0021229912196773183, + 0.0017975707779160951, + -0.0006194325712509404, + -0.0004543300349236303, + -0.0002868182533726062, + 0.0006643387233529266, + -0.0017354384415934402, + -0.0004484692659537939, + 0.002921411800146266, + -0.0034993812098134593, + -5.746610615618304e-05, + -0.0018971019883876328, + 0.002081090481415879, + -0.00016366253479650742, + -4.901125252378672e-05, + 3.150277175423273e-05, + -0.0006547485112516678, + 0.00021555401586171655, + 4.31236271479316e-06, + -0.0016518095439264416, + 0.0005398348888687889, + -0.0025510190757651293, + 0.00034245880645158907, + 0.0017957525139819155, + -0.001460159452424121, + -0.00015221357510537432, + 0.00029467917464015363, + 0.0003138927938376339, + -0.00027479344429587517, + 0.00047938018776191974, + 0.0004062842240844888, + -0.0005455686976354522, + 0.0005322384246015745, + 0.001555084566962813, + 0.000571991330253481, + -0.0009365758696000555, + 6.734974747663757e-05, + 6.943403496677737e-05, + 0.00010001347556806527, + -5.2307923893423455e-05, + -0.00029047532977290104, + 0.00100151797047199, + 0.0016103990119311724, + 0.0005131796319630577, + 0.0006683269865756838, + 7.4002604237059946e-06, + -0.0005371554047770387, + -0.02214346474013617, + 0.016794094826306304, + -0.004025967359185296, + -0.0002905570360920155, + 0.0008524102874805489, + -0.003055386112916023, + 0.0020392079452522744, + 0.0039805712172386164, + -0.0037469648159867562, + 0.05925939973889675, + 1.070308346404483, + 0.08214063457945922, + -0.0016342301614208203, + -0.00040091719819116886, + -0.0010341038896804673, + 0.002012760535148478, + 0.0004674808225432368, + -0.0002897211693410111, + 0.0010418576175709074, + -0.000245041865598247, + -0.0017931964476003307, + -0.0009668209111811646, + -0.014796236030209822, + -0.013323822561725514, + -0.000260048573573255, + -0.0008526332137507434, + -0.001040368745589977, + 0.0009942752482566845, + 0.00102587793054071, + -0.001725767888196689, + 0.0005655091926815686, + 0.005147901721102176, + -0.024536479730933224, + -0.02791650016070407, + 0.0016857463130560359, + 0.001470249209345316, + 9.47361168958422e-05, + -8.670246967642241e-06, + -0.0007963640703317124, + -2.9888904550718314e-06, + 0.00037142990104026795, + -0.00037725960333987764, + 0.00029072433591034833, + 0.007858072068337867, + 0.008497944623367807, + 0.008291864820601617, + 0.0024489338383286803, + -0.0009700681018906904, + 0.0013525713011347748, + -0.0006296355756249766, + 0.00011627128751750781, + 0.0009710694089331748, + -0.0008896644310515731, + -0.000963072959015437, + 0.003016541924641319, + 0.0050186828141548494, + 0.0041987858305337, + -0.003492176205171512, + -0.00021037277684570166, + 0.00026469445135235376, + 0.0008049705788839383, + -0.0005377904737783495, + -9.962378954430911e-05, + 0.001193378744771027, + -0.0010248549151275633, + -0.004698416283810698, + -0.0018825676106537228, + -0.004056053451607298, + -0.0030486110453807882, + 0.000354808816131495, + -0.0007940595368987249, + -1.9742540907698645e-05, + -0.0001622040484190238, + -0.0010358887075598408, + -0.00013535552922282407, + 0.0007016910773292311, + -0.00028619009586306566, + -0.004206430520364198, + -0.0015817885446743197, + 0.00044728723142764895, + 0.001431468977255503, + -0.0002540133454853355, + -0.00032496382969797596, + 0.00030701965517004816, + 0.0002077056879574103, + -0.000672664468788692, + 0.0006857654414789853, + 0.0018255106560533013, + 0.0009446893547110723, + 0.0005720123883304623, + 0.002232993882601502, + 0.00040510075217493327, + -0.0002331998940398098 + ] + }, + "names": [ + "rc(0,0)", + "rc(0,1)", + "rc(0,2)", + "rc(0,3)", + "rc(0,4)", + "rc(0,5)", + "rc(1,-5)", + "rc(1,-4)", + "rc(1,-3)", + "rc(1,-2)", + "rc(1,-1)", + "rc(1,0)", + "rc(1,1)", + "rc(1,2)", + "rc(1,3)", + "rc(1,4)", + "rc(1,5)", + "rc(2,-5)", + "rc(2,-4)", + "rc(2,-3)", + "rc(2,-2)", + "rc(2,-1)", + "rc(2,0)", + "rc(2,1)", + "rc(2,2)", + "rc(2,3)", + "rc(2,4)", + "rc(2,5)", + "rc(3,-5)", + "rc(3,-4)", + "rc(3,-3)", + "rc(3,-2)", + "rc(3,-1)", + "rc(3,0)", + "rc(3,1)", + "rc(3,2)", + "rc(3,3)", + "rc(3,4)", + "rc(3,5)", + "rc(4,-5)", + "rc(4,-4)", + "rc(4,-3)", + "rc(4,-2)", + "rc(4,-1)", + "rc(4,0)", + "rc(4,1)", + "rc(4,2)", + "rc(4,3)", + "rc(4,4)", + "rc(4,5)", + "rc(5,-5)", + "rc(5,-4)", + "rc(5,-3)", + "rc(5,-2)", + "rc(5,-1)", + "rc(5,0)", + "rc(5,1)", + "rc(5,2)", + "rc(5,3)", + "rc(5,4)", + "rc(5,5)", + "rc(6,-5)", + "rc(6,-4)", + "rc(6,-3)", + "rc(6,-2)", + "rc(6,-1)", + "rc(6,0)", + "rc(6,1)", + "rc(6,2)", + "rc(6,3)", + "rc(6,4)", + "rc(6,5)", + "rc(7,-5)", + "rc(7,-4)", + "rc(7,-3)", + "rc(7,-2)", + "rc(7,-1)", + "rc(7,0)", + "rc(7,1)", + "rc(7,2)", + "rc(7,3)", + "rc(7,4)", + "rc(7,5)", + "rc(8,-5)", + "rc(8,-4)", + "rc(8,-3)", + "rc(8,-2)", + "rc(8,-1)", + "rc(8,0)", + "rc(8,1)", + "rc(8,2)", + "rc(8,3)", + "rc(8,4)", + "rc(8,5)", + "zs(0,1)", + "zs(0,2)", + "zs(0,3)", + "zs(0,4)", + "zs(0,5)", + "zs(1,-5)", + "zs(1,-4)", + "zs(1,-3)", + "zs(1,-2)", + "zs(1,-1)", + "zs(1,0)", + "zs(1,1)", + "zs(1,2)", + "zs(1,3)", + "zs(1,4)", + "zs(1,5)", + "zs(2,-5)", + "zs(2,-4)", + "zs(2,-3)", + "zs(2,-2)", + "zs(2,-1)", + "zs(2,0)", + "zs(2,1)", + "zs(2,2)", + "zs(2,3)", + "zs(2,4)", + "zs(2,5)", + "zs(3,-5)", + "zs(3,-4)", + "zs(3,-3)", + "zs(3,-2)", + "zs(3,-1)", + "zs(3,0)", + "zs(3,1)", + "zs(3,2)", + "zs(3,3)", + "zs(3,4)", + "zs(3,5)", + "zs(4,-5)", + "zs(4,-4)", + "zs(4,-3)", + "zs(4,-2)", + "zs(4,-1)", + "zs(4,0)", + "zs(4,1)", + "zs(4,2)", + "zs(4,3)", + "zs(4,4)", + "zs(4,5)", + "zs(5,-5)", + "zs(5,-4)", + "zs(5,-3)", + "zs(5,-2)", + "zs(5,-1)", + "zs(5,0)", + "zs(5,1)", + "zs(5,2)", + "zs(5,3)", + "zs(5,4)", + "zs(5,5)", + "zs(6,-5)", + "zs(6,-4)", + "zs(6,-3)", + "zs(6,-2)", + "zs(6,-1)", + "zs(6,0)", + "zs(6,1)", + "zs(6,2)", + "zs(6,3)", + "zs(6,4)", + "zs(6,5)", + "zs(7,-5)", + "zs(7,-4)", + "zs(7,-3)", + "zs(7,-2)", + "zs(7,-1)", + "zs(7,0)", + "zs(7,1)", + "zs(7,2)", + "zs(7,3)", + "zs(7,4)", + "zs(7,5)", + "zs(8,-5)", + "zs(8,-4)", + "zs(8,-3)", + "zs(8,-2)", + "zs(8,-1)", + "zs(8,0)", + "zs(8,1)", + "zs(8,2)", + "zs(8,3)", + "zs(8,4)", + "zs(8,5)" + ], + "free": { + "@module": "numpy", + "@class": "array", + "dtype": "bool", + "data": [ + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true + ] + }, + "lower_bounds": { + "@module": "numpy", + "@class": "array", + "dtype": "float64", + "data": [ + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity + ] + }, + "upper_bounds": { + "@module": "numpy", + "@class": "array", + "dtype": "float64", + "data": [ + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity + ] + } + }, + "SurfaceRZFourier7": { + "@module": "simsopt.geo.surfacerzfourier", + "@class": "SurfaceRZFourier", + "@name": "SurfaceRZFourier7", + "@version": "0.19.0.post217+g54251e8b.d20240725.dirty", + "nfp": 3, + "stellsym": true, + "mpol": 8, + "ntor": 5, + "quadpoints_phi": { + "@module": "numpy", + "@class": "array", + "dtype": "float64", + "data": [ + 0.0, + 0.010416666666666666, + 0.020833333333333332, + 0.03125, + 0.041666666666666664, + 0.052083333333333336, + 0.0625, + 0.07291666666666667, + 0.08333333333333333, + 0.09375, + 0.10416666666666667, + 0.11458333333333333, + 0.125, + 0.13541666666666666, + 0.14583333333333334, + 0.15625, + 0.16666666666666666, + 0.17708333333333334, + 0.1875, + 0.19791666666666666, + 0.20833333333333334, + 0.21875, + 0.22916666666666666, + 0.23958333333333334, + 0.25, + 0.2604166666666667, + 0.2708333333333333, + 0.28125, + 0.2916666666666667, + 0.3020833333333333, + 0.3125, + 0.3229166666666667, + 0.3333333333333333, + 0.34375, + 0.3541666666666667, + 0.3645833333333333, + 0.375, + 0.3854166666666667, + 0.3958333333333333, + 0.40625, + 0.4166666666666667, + 0.4270833333333333, + 0.4375, + 0.4479166666666667, + 0.4583333333333333, + 0.46875, + 0.4791666666666667, + 0.4895833333333333, + 0.5, + 0.5104166666666666, + 0.5208333333333334, + 0.53125, + 0.5416666666666666, + 0.5520833333333334, + 0.5625, + 0.5729166666666666, + 0.5833333333333334, + 0.59375, + 0.6041666666666666, + 0.6145833333333334, + 0.625, + 0.6354166666666666, + 0.6458333333333334, + 0.65625, + 0.6666666666666666, + 0.6770833333333334, + 0.6875, + 0.6979166666666666, + 0.7083333333333334, + 0.71875, + 0.7291666666666666, + 0.7395833333333334, + 0.75, + 0.7604166666666666, + 0.7708333333333334, + 0.78125, + 0.7916666666666666, + 0.8020833333333334, + 0.8125, + 0.8229166666666666, + 0.8333333333333334, + 0.84375, + 0.8541666666666666, + 0.8645833333333334, + 0.875, + 0.8854166666666666, + 0.8958333333333334, + 0.90625, + 0.9166666666666666, + 0.9270833333333334, + 0.9375, + 0.9479166666666666, + 0.9583333333333334, + 0.96875, + 0.9791666666666666, + 0.9895833333333334 + ] + }, + "quadpoints_theta": { + "@module": "numpy", + "@class": "array", + "dtype": "float64", + "data": [ + 0.0, + 0.03125, + 0.0625, + 0.09375, + 0.125, + 0.15625, + 0.1875, + 0.21875, + 0.25, + 0.28125, + 0.3125, + 0.34375, + 0.375, + 0.40625, + 0.4375, + 0.46875, + 0.5, + 0.53125, + 0.5625, + 0.59375, + 0.625, + 0.65625, + 0.6875, + 0.71875, + 0.75, + 0.78125, + 0.8125, + 0.84375, + 0.875, + 0.90625, + 0.9375, + 0.96875 + ] + }, + "dofs": { + "$type": "ref", + "value": "5807918224" + } + }, + "5807726320": { + "@module": "simsopt._core.optimizable", + "@class": "DOFs", + "@name": "5807726320", + "@version": "0.19.0.post217+g54251e8b.d20240725.dirty", + "x": { + "@module": "numpy", + "@class": "array", + "dtype": "float64", + "data": [ + 1.3782, + -0.0041452, + -0.0041374, + 0.0036116, + -0.00033285, + 0.0, + 0.0, + -0.00045103, + 0.00047123, + -0.0026155999999999996, + 0.020869000000000002, + 0.27073, + -0.135, + 0.0048519, + 0.00048906, + -0.0009469400000000001, + 0.0, + 0.0, + 0.00044897000000000003, + -0.00046166000000000013, + -0.00019987999999999956, + 0.00893389999999999, + 0.091094, + 0.090684, + 0.027208, + -0.0044665, + -0.00014812, + 0.0, + 0.0, + -2.4029999999999978e-05, + 0.0005943299999999994, + 5.247300000000081e-05, + 0.0031952, + -0.014174, + -0.0017796, + -0.010529, + -0.0058013, + 0.0014257, + 0.0, + 0.0, + 2.589299999999993e-05, + -0.00032664999999999994, + -4.8038999999999734e-05, + -0.0005296600000000005, + 0.0045451, + -0.0046998, + 0.0064586, + 0.00017810999999999996, + 0.0017237, + 0.0, + 0.0, + -2.6970999999999987e-06, + 7.652000000000002e-05, + 1.002199999999994e-05, + 0.00040483, + -0.0016816999999999997, + 9.484100000000015e-06, + -0.0013733, + -0.00075237, + 0.00018598, + 0.0, + 0.0, + 3.4246999999999953e-06, + -5.801699999999999e-05, + -7.060900000000014e-06, + -8.104099999999992e-05, + 0.00043087, + -0.0014165, + 0.00093906, + 7.109599999999998e-05, + 0.00034674, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0071634, + 0.0075995, + -0.0029503, + 0.00051481, + -0.0, + 0.0, + 7.048600000000004e-05, + -0.00015952000000000002, + 0.0014115999999999996, + 0.0092873, + 0.46465, + 0.16516, + -0.006055900000000001, + -0.00017736, + 0.00056642, + 0.0, + 0.0, + 0.00017521, + 0.00059225, + 0.00019349999999999928, + 0.012753999999999998, + 0.015451, + 0.011617999999999998, + -0.027337, + 0.0013515, + 0.0008653500000000001, + 0.0, + 0.0, + 1.759200000000009e-05, + -0.0005865399999999997, + -0.00010220000000000023, + -0.0030523, + 0.011779, + -0.001926, + 0.010663, + 0.0058094999999999996, + -0.0014443, + 0.0, + 0.0, + -8.7822e-06, + 0.00026077999999999995, + 2.345300000000011e-05, + 0.00029088000000000036, + 0.00018259999999999997, + 0.0096416, + -0.0030495, + -0.0005417299999999999, + -0.001787, + 0.0, + 0.0, + -2.6970999999999987e-06, + 7.652000000000002e-05, + 1.002199999999994e-05, + 0.00040483, + -0.0016816999999999997, + 9.484100000000015e-06, + -0.0013733, + -0.00075237, + 0.00018598, + 0.0, + 0.0, + 3.4246999999999953e-06, + -5.801699999999999e-05, + -7.060900000000014e-06, + -8.104099999999992e-05, + 0.00043087, + -0.0014165, + 0.00093906, + 7.109599999999998e-05, + 0.00034674, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ] + }, + "names": [ + "rc(0,0)", + "rc(0,1)", + "rc(0,2)", + "rc(0,3)", + "rc(0,4)", + "rc(0,5)", + "rc(1,-5)", + "rc(1,-4)", + "rc(1,-3)", + "rc(1,-2)", + "rc(1,-1)", + "rc(1,0)", + "rc(1,1)", + "rc(1,2)", + "rc(1,3)", + "rc(1,4)", + "rc(1,5)", + "rc(2,-5)", + "rc(2,-4)", + "rc(2,-3)", + "rc(2,-2)", + "rc(2,-1)", + "rc(2,0)", + "rc(2,1)", + "rc(2,2)", + "rc(2,3)", + "rc(2,4)", + "rc(2,5)", + "rc(3,-5)", + "rc(3,-4)", + "rc(3,-3)", + "rc(3,-2)", + "rc(3,-1)", + "rc(3,0)", + "rc(3,1)", + "rc(3,2)", + "rc(3,3)", + "rc(3,4)", + "rc(3,5)", + "rc(4,-5)", + "rc(4,-4)", + "rc(4,-3)", + "rc(4,-2)", + "rc(4,-1)", + "rc(4,0)", + "rc(4,1)", + "rc(4,2)", + "rc(4,3)", + "rc(4,4)", + "rc(4,5)", + "rc(5,-5)", + "rc(5,-4)", + "rc(5,-3)", + "rc(5,-2)", + "rc(5,-1)", + "rc(5,0)", + "rc(5,1)", + "rc(5,2)", + "rc(5,3)", + "rc(5,4)", + "rc(5,5)", + "rc(6,-5)", + "rc(6,-4)", + "rc(6,-3)", + "rc(6,-2)", + "rc(6,-1)", + "rc(6,0)", + "rc(6,1)", + "rc(6,2)", + "rc(6,3)", + "rc(6,4)", + "rc(6,5)", + "rc(7,-5)", + "rc(7,-4)", + "rc(7,-3)", + "rc(7,-2)", + "rc(7,-1)", + "rc(7,0)", + "rc(7,1)", + "rc(7,2)", + "rc(7,3)", + "rc(7,4)", + "rc(7,5)", + "rc(8,-5)", + "rc(8,-4)", + "rc(8,-3)", + "rc(8,-2)", + "rc(8,-1)", + "rc(8,0)", + "rc(8,1)", + "rc(8,2)", + "rc(8,3)", + "rc(8,4)", + "rc(8,5)", + "zs(0,1)", + "zs(0,2)", + "zs(0,3)", + "zs(0,4)", + "zs(0,5)", + "zs(1,-5)", + "zs(1,-4)", + "zs(1,-3)", + "zs(1,-2)", + "zs(1,-1)", + "zs(1,0)", + "zs(1,1)", + "zs(1,2)", + "zs(1,3)", + "zs(1,4)", + "zs(1,5)", + "zs(2,-5)", + "zs(2,-4)", + "zs(2,-3)", + "zs(2,-2)", + "zs(2,-1)", + "zs(2,0)", + "zs(2,1)", + "zs(2,2)", + "zs(2,3)", + "zs(2,4)", + "zs(2,5)", + "zs(3,-5)", + "zs(3,-4)", + "zs(3,-3)", + "zs(3,-2)", + "zs(3,-1)", + "zs(3,0)", + "zs(3,1)", + "zs(3,2)", + "zs(3,3)", + "zs(3,4)", + "zs(3,5)", + "zs(4,-5)", + "zs(4,-4)", + "zs(4,-3)", + "zs(4,-2)", + "zs(4,-1)", + "zs(4,0)", + "zs(4,1)", + "zs(4,2)", + "zs(4,3)", + "zs(4,4)", + "zs(4,5)", + "zs(5,-5)", + "zs(5,-4)", + "zs(5,-3)", + "zs(5,-2)", + "zs(5,-1)", + "zs(5,0)", + "zs(5,1)", + "zs(5,2)", + "zs(5,3)", + "zs(5,4)", + "zs(5,5)", + "zs(6,-5)", + "zs(6,-4)", + "zs(6,-3)", + "zs(6,-2)", + "zs(6,-1)", + "zs(6,0)", + "zs(6,1)", + "zs(6,2)", + "zs(6,3)", + "zs(6,4)", + "zs(6,5)", + "zs(7,-5)", + "zs(7,-4)", + "zs(7,-3)", + "zs(7,-2)", + "zs(7,-1)", + "zs(7,0)", + "zs(7,1)", + "zs(7,2)", + "zs(7,3)", + "zs(7,4)", + "zs(7,5)", + "zs(8,-5)", + "zs(8,-4)", + "zs(8,-3)", + "zs(8,-2)", + "zs(8,-1)", + "zs(8,0)", + "zs(8,1)", + "zs(8,2)", + "zs(8,3)", + "zs(8,4)", + "zs(8,5)" + ], + "free": { + "@module": "numpy", + "@class": "array", + "dtype": "bool", + "data": [ + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true + ] + }, + "lower_bounds": { + "@module": "numpy", + "@class": "array", + "dtype": "float64", + "data": [ + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity, + -Infinity + ] + }, + "upper_bounds": { + "@module": "numpy", + "@class": "array", + "dtype": "float64", + "data": [ + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity, + Infinity + ] + } + }, + "SurfaceRZFourier4": { + "@module": "simsopt.geo.surfacerzfourier", + "@class": "SurfaceRZFourier", + "@name": "SurfaceRZFourier4", + "@version": "0.19.0.post217+g54251e8b.d20240725.dirty", + "nfp": 3, + "stellsym": true, + "mpol": 8, + "ntor": 5, + "quadpoints_phi": { + "@module": "numpy", + "@class": "array", + "dtype": "float64", + "data": [ + 0.0, + 0.010416666666666666, + 0.020833333333333332, + 0.03125, + 0.041666666666666664, + 0.05208333333333333, + 0.0625, + 0.07291666666666666, + 0.08333333333333333, + 0.09375, + 0.10416666666666666, + 0.11458333333333333, + 0.125, + 0.13541666666666666, + 0.14583333333333331, + 0.15625, + 0.16666666666666666, + 0.17708333333333331, + 0.1875, + 0.19791666666666666, + 0.20833333333333331, + 0.21875, + 0.22916666666666666, + 0.23958333333333331, + 0.25, + 0.26041666666666663, + 0.2708333333333333, + 0.28125, + 0.29166666666666663, + 0.3020833333333333, + 0.3125, + 0.32291666666666663 + ] + }, + "quadpoints_theta": { + "@module": "numpy", + "@class": "array", + "dtype": "float64", + "data": [ + 0.0, + 0.03125, + 0.0625, + 0.09375, + 0.125, + 0.15625, + 0.1875, + 0.21875, + 0.25, + 0.28125, + 0.3125, + 0.34375, + 0.375, + 0.40625, + 0.4375, + 0.46875, + 0.5, + 0.53125, + 0.5625, + 0.59375, + 0.625, + 0.65625, + 0.6875, + 0.71875, + 0.75, + 0.78125, + 0.8125, + 0.84375, + 0.875, + 0.90625, + 0.9375, + 0.96875 + ] + }, + "dofs": { + "$type": "ref", + "value": "5807726320" + } + } + } +} \ No newline at end of file diff --git a/tests/util/test_winding_surface_helper_functions.py b/tests/util/test_winding_surface_helper_functions.py new file mode 100644 index 000000000..76cc79dec --- /dev/null +++ b/tests/util/test_winding_surface_helper_functions.py @@ -0,0 +1,1177 @@ +""" +Unit tests for simsopt.util.winding_surface_helper_functions. + +These tests cover current potential evaluation, contour utilities, geometry +mapping, and data loading. Tests use synthetic data where possible to avoid +file dependencies; a few tests require the regcoil_out.hsx.nc test file. +""" + +import unittest +import numpy as np +from pathlib import Path + +from simsopt.util.winding_surface_helper_functions import ( + _load_simsopt_regcoil_data, + _load_CP_and_geometries, + _current_potential_at_point, + _grad_current_potential_at_point, + _genCPvals, + _genKvals, + is_periodic_lines, + minDist, + sortLevels, + chooseContours_matching_coilType, + _points_in_polygon, + _ID_mod_hel, + ID_and_cut_contour_types, + ID_halfway_contour, + _removearray, + compute_baseline_WP_currents, + check_and_compute_nested_WP_currents, + _SIMSOPT_line_XYZ_RZ, + _writeToCurve, + _load_regcoil_data, + set_axes_equal, + _load_surface_dofs_properly, + _make_onclick, + _make_onpick, + _make_on_key, + _map_data_full_torus_3x1, + _run_cut_coils_compute_NWP_currents, + _run_cut_coils_select_contours_non_interactive, +) +# Import private function for testing fallback path +from simsopt.util.winding_surface_helper_functions import _contour_paths + +TEST_DIR = Path(__file__).resolve().parent.parent / "test_files" + + +def _make_args(Ip=0, It=0, nfp=4): + """Create minimal args tuple for current potential evaluation.""" + xm = np.array([0, 1, 1]) + xn = np.array([0, 0, 1]) * nfp + phi_cos = np.array([0.0, 0.5, 0.2]) + phi_sin = np.array([0.0, 0.1, -0.1]) + return (Ip, It, xm, xn, phi_cos, phi_sin, np.array([nfp])) + + +class TestCurrentPotentialAtPoint(unittest.TestCase): + """Tests for current_potential_at_point.""" + + def test_zero_modes(self): + """With phi_cos=phi_sin=0 and Ip=It=0, potential is zero.""" + args = (0, 0, np.array([0]), np.array([0]), np.array([0.0]), np.array([0.0]), np.array([4])) + val = _current_potential_at_point(np.array([0.5, 0.3]), args) + self.assertAlmostEqual(val, 0.0) + + def test_multi_valued_part(self): + """Multi-valued part It*θ/(2π) + Ip*ζ/(2π) is correct.""" + args = (2 * np.pi, 4 * np.pi, np.array([0]), np.array([0]), np.array([0.0]), np.array([0.0]), np.array([4])) + val = _current_potential_at_point(np.array([0.5, 0.25]), args) + # MV = It*θ/(2π) + Ip*ζ/(2π) = 4π*0.5/(2π) + 2π*0.25/(2π) = 1 + 0.25 = 1.25 + self.assertAlmostEqual(val, 1.25) + + def test_fourier_part(self): + """Single mode cos(mθ - nζ) gives correct value.""" + args = (0, 0, np.array([1]), np.array([0]), np.array([1.0]), np.array([0.0]), np.array([4])) + val = _current_potential_at_point(np.array([0.0, 0.0]), args) + self.assertAlmostEqual(val, 1.0) + + def test_scalar_input(self): + """x can be list or array; returns float.""" + args = _make_args() + v1 = _current_potential_at_point([0.5, 0.3], args) + v2 = _current_potential_at_point(np.array([0.5, 0.3]), args) + self.assertIsInstance(v1, float) + self.assertAlmostEqual(v1, v2) + + +class TestGradCurrentPotentialAtPoint(unittest.TestCase): + """Tests for _grad_current_potential_at_point.""" + + def test_constant_potential_zero_grad(self): + """Constant potential has zero gradient.""" + args = (0, 0, np.array([0]), np.array([0]), np.array([0.0]), np.array([0.0]), np.array([4])) + val = _grad_current_potential_at_point(np.array([0.5, 0.3]), args) + self.assertAlmostEqual(val, 0.0) + + def test_returns_positive(self): + """|∇φ| is non-negative.""" + args = _make_args() + val = _grad_current_potential_at_point(np.array([0.5, 0.3]), args) + self.assertGreaterEqual(val, 0) + + +class TestGenCPvals(unittest.TestCase): + """Tests for genCPvals.""" + + def test_output_shapes(self): + """Output arrays have correct shapes.""" + args = _make_args() + thetas, zetas, phi, phi_SV, phi_NSV, ARGS = _genCPvals( + (0, 2 * np.pi), (0, 2 * np.pi / 4), (8, 4), args + ) + self.assertEqual(len(thetas), 8) + self.assertEqual(len(zetas), 4) + self.assertEqual(phi.shape, (4, 8)) # .T so (nZ, nT) + self.assertEqual(phi_SV.shape, (4, 8)) + self.assertEqual(phi_NSV.shape, (4, 8)) + + def test_single_valued_vs_full(self): + """With Ip=It=0, phi matches phi_SV.""" + args = _make_args(Ip=0, It=0) + _, _, phi, phi_SV, _, _ = _genCPvals( + (0, 2 * np.pi), (0, 2 * np.pi / 4), (4, 4), args + ) + np.testing.assert_allclose(phi, phi_SV) + + +class TestGenKvals(unittest.TestCase): + """Tests for genKvals.""" + + def test_output_shapes(self): + """Output arrays have correct shapes.""" + args = _make_args() + thetas, zetas, K, K_SV, K_NSV, ARGS = _genKvals( + (0, 2 * np.pi), (0, 2 * np.pi / 4), (6, 3), args + ) + self.assertEqual(K.shape, (3, 6)) + self.assertGreaterEqual(np.min(K), 0) + + +class TestIsPeriodicLines(unittest.TestCase): + """Tests for is_periodic_lines.""" + + def test_closed_contour(self): + """Contour that starts and ends at same point is periodic.""" + line = np.array([[0, 0], [1, 0], [1, 1], [0, 0]]) + self.assertTrue(is_periodic_lines(line, tol=0.1)) + + def test_open_contour(self): + """Contour with different endpoints is not periodic.""" + line = np.array([[0, 0], [1, 0], [1, 1], [0.5, 0.5]]) + self.assertFalse(is_periodic_lines(line, tol=0.01)) + + def test_nearly_closed(self): + """Within tolerance counts as closed.""" + line = np.array([[0, 0], [1, 0], [1, 1], [0, 0.001]]) + self.assertTrue(is_periodic_lines(line, tol=0.1)) + + +class TestMinDist(unittest.TestCase): + """Tests for minDist.""" + + def test_point_on_line(self): + """Distance from point on line is zero.""" + pt = np.array([1.0, 0.0]) + line = np.array([[0, 0], [1, 0], [2, 0]]) + self.assertAlmostEqual(minDist(pt, line), 0.0) + + def test_point_off_line(self): + """Distance from point to nearest point in line (discretized path).""" + pt = np.array([1.0, 1.0]) + # Line from (0,0) to (2,0) - include (1,0) so closest point gives distance 1 + line = np.array([[0, 0], [1, 0], [2, 0]]) + self.assertAlmostEqual(minDist(pt, line), 1.0) + + +class TestSortLevels(unittest.TestCase): + """Tests for sortLevels.""" + + def test_ascending_order(self): + """Levels and points are sorted by level ascending.""" + levels = [3, 1, 2] + points = [np.array([3, 3]), np.array([1, 1]), np.array([2, 2])] + L, P = sortLevels(levels, points) + self.assertEqual(L, [1, 2, 3]) + np.testing.assert_array_almost_equal(P[0], [1, 1]) + np.testing.assert_array_almost_equal(P[1], [2, 2]) + np.testing.assert_array_almost_equal(P[2], [3, 3]) + + +class TestChooseContoursMatchingCoilType(unittest.TestCase): + """Tests for chooseContours_matching_coilType.""" + + def test_free_returns_all(self): + """ctype='free' returns all lines.""" + closed = np.array([[0, 0], [1, 0], [0, 0]]) + open_ = np.array([[0, 0], [1, 1]]) + lines = [closed, open_] + result = chooseContours_matching_coilType(lines, 'free') + self.assertEqual(len(result), 2) + + def test_wp_returns_only_closed(self): + """ctype='wp' returns only closed contours.""" + closed = np.array([[0, 0], [1, 0], [0, 0]]) + open_ = np.array([[0, 0], [1, 1]]) + lines = [closed, open_] + result = chooseContours_matching_coilType(lines, 'wp') + self.assertEqual(len(result), 1) + np.testing.assert_array_almost_equal(result[0], closed) + + def test_mod_returns_only_open(self): + """ctype='mod' returns only open contours.""" + closed = np.array([[0, 0], [1, 0], [0, 0]]) + open_ = np.array([[0, 0], [1, 1]]) + lines = [closed, open_] + result = chooseContours_matching_coilType(lines, 'mod') + self.assertEqual(len(result), 1) + np.testing.assert_array_almost_equal(result[0], open_) + + def test_hel_returns_only_open(self): + """ctype='hel' returns only open contours (same as mod).""" + closed = np.array([[0, 0], [1, 0], [0, 0]]) + open_ = np.array([[0, 0], [1, 1]]) + lines = [closed, open_] + result = chooseContours_matching_coilType(lines, 'hel') + self.assertEqual(len(result), 1) + np.testing.assert_array_almost_equal(result[0], open_) + + +class TestPointsInPolygon(unittest.TestCase): + """Tests for _points_in_polygon.""" + + def test_triangle_contains_inside(self): + """Points inside triangle are identified.""" + # Triangle (0,0), (1,0), (0.5,1) + polygon = np.array([[0, 0], [1, 0], [0.5, 1], [0, 0]]) + t = np.linspace(0, 1, 5) + z = np.linspace(0, 1, 5) + i1, i2, BOOL = _points_in_polygon((t, z), polygon) + # Center (0.5, 0.33) should be inside + self.assertGreater(np.sum(BOOL), 0) + + def test_returns_correct_lengths(self): + """Returned indices match BOOL mask.""" + polygon = np.array([[0, 0], [1, 0], [1, 1], [0, 1], [0, 0]]) + t = np.linspace(0.1, 0.9, 3) + z = np.linspace(0.1, 0.9, 3) + i1, i2, BOOL = _points_in_polygon((t, z), polygon) + self.assertEqual(len(i1), np.sum(BOOL)) + self.assertEqual(len(i2), np.sum(BOOL)) + + +class TestIDModHel(unittest.TestCase): + """Tests for _ID_mod_hel.""" + + def test_modular_contour(self): + """Contour with same cos(θ) and ζ at ends is modular.""" + # Start (0, 0), end (2π, 0) -> cos(0)=cos(2π)=1, zeta same + contour = np.array([[0, 0], [np.pi, 0.5], [2 * np.pi, 0]]) + names, ints = _ID_mod_hel([contour], tol=0.1) + self.assertEqual(names[0], 'mod') + self.assertEqual(ints[0], 1) + + def test_helical_contour(self): + """Contour with θ spanning 2π is helical.""" + contour = np.array([[0, 0], [np.pi, 0.25], [2 * np.pi, 0.5]]) + names, ints = _ID_mod_hel([contour], tol=0.05) + self.assertEqual(names[0], 'hel') + self.assertEqual(ints[0], 2) + + def test_vacuum_field_contour(self): + """Contour with neither mod nor hel endpoints is vacuum-field (vf).""" + # θ spans ~π, ζ differs; neither cos(θ) match nor |θ0-θf|≈2π + contour = np.array([[0, 0], [np.pi / 2, 0.5], [np.pi, 1.0]]) + names, ints = _ID_mod_hel([contour], tol=0.05) + self.assertEqual(names[0], 'vf') + self.assertEqual(ints[0], 3) + + +class TestIDAndCutContourTypes(unittest.TestCase): + """Tests for ID_and_cut_contour_types.""" + + def test_mixed_contours(self): + """Closed and open contours are classified correctly.""" + closed = np.array([[0, 0], [1, 0], [1, 1], [0, 0]]) + open_ = np.array([[0, 0], [1, 1]]) + open_contours, closed_contours, types = ID_and_cut_contour_types([closed, open_]) + self.assertEqual(len(closed_contours), 1) + self.assertEqual(len(open_contours), 1) + self.assertEqual(types[0], 0) + self.assertIn(types[1], (1, 2, 3)) + + +class TestIDHalfwayContour(unittest.TestCase): + """Tests for ID_halfway_contour.""" + + def test_halfway_contours_between_open(self): + """Halfway contours are found between consecutive open contours.""" + args = _make_args() + theta = np.linspace(0, 2 * np.pi, 20, endpoint=False) + zeta = np.linspace(0, 2 * np.pi / 4, 15, endpoint=False) + _, _, cpd, _, _, _ = _genCPvals((0, 2 * np.pi), (0, 2 * np.pi / 4), (20, 15), args) + c1 = np.array([[0.5, 0.3], [1.5, 0.4]]) + c2 = np.array([[1.0, 0.5], [2.0, 0.6]]) + contours = [c1, c2] + data = (theta, zeta, cpd) + halfway = ID_halfway_contour(contours, data, do_plot=False, args=args) + self.assertEqual(len(halfway), 1) + self.assertGreater(len(halfway[0]), 0) + + def test_halfway_contours_do_plot_true(self): + """ID_halfway_contour with do_plot=True sets up axes and returns contours.""" + from unittest.mock import patch + args = _make_args() + theta = np.linspace(0, 2 * np.pi, 20, endpoint=False) + zeta = np.linspace(0, 2 * np.pi / 4, 15, endpoint=False) + _, _, cpd, _, _, _ = _genCPvals((0, 2 * np.pi), (0, 2 * np.pi / 4), (20, 15), args) + c1 = np.array([[0.5, 0.3], [1.5, 0.4]]) + c2 = np.array([[1.0, 0.5], [2.0, 0.6]]) + contours = [c1, c2] + data = (theta, zeta, cpd) + with patch('simsopt.util.winding_surface_helper_functions.plt.show'): + halfway = ID_halfway_contour(contours, data, do_plot=True, args=args) + self.assertEqual(len(halfway), 1) + + +class TestRemoveArray(unittest.TestCase): + """Tests for _removearray.""" + + def test_remove_existing(self): + """Removing existing array modifies list in place.""" + a = np.array([1, 2]) + b = np.array([3, 4]) + L = [a, b] + _removearray(L, a) + self.assertEqual(len(L), 1) + np.testing.assert_array_equal(L[0], b) + + def test_remove_raises_if_not_found(self): + """Removing non-existent array raises ValueError.""" + L = [np.array([1, 2])] + with self.assertRaises(ValueError): + _removearray(L, np.array([9, 9])) + + +class TestComputeBaselineWPCurrents(unittest.TestCase): + """Tests for compute_baseline_WP_currents.""" + + def test_single_contour(self): + """Single closed contour returns one current.""" + args = _make_args() + contour = np.array([[0.5, 0.5], [1.5, 0.5], [1.5, 1.5], [0.5, 1.5], [0.5, 0.5]]) + wp_currents, max_vals, func_vals = compute_baseline_WP_currents([contour], args, plot=False) + self.assertEqual(len(wp_currents), 1) + self.assertGreater(wp_currents[0], 0) + + def test_single_contour_plot_true(self): + """compute_baseline_WP_currents with plot=True runs without error.""" + from unittest.mock import patch + args = _make_args() + contour = np.array([[0.5, 0.5], [1.5, 0.5], [1.5, 1.5], [0.5, 1.5], [0.5, 0.5]]) + with patch('simsopt.util.winding_surface_helper_functions.plt.show'): + wp_currents, _, _ = compute_baseline_WP_currents([contour], args, plot=True) + self.assertEqual(len(wp_currents), 1) + + +class TestCheckAndComputeNestedWPCurrents(unittest.TestCase): + """Tests for check_and_compute_nested_WP_currents.""" + + def test_no_nesting(self): + """Non-nested contours return unchanged currents.""" + args = _make_args() + c1 = np.array([[0.5, 0.5], [1.0, 0.5], [1.0, 1.0], [0.5, 1.0], [0.5, 0.5]]) + c2 = np.array([[2.0, 2.0], [2.5, 2.0], [2.5, 2.5], [2.0, 2.5], [2.0, 2.0]]) + wp = [1.0, 1.0] + result, nested, fv, nc = check_and_compute_nested_WP_currents([c1, c2], wp, args, plot=False) + np.testing.assert_array_almost_equal(result, wp) + + def test_nested_contours(self): + """Nested contours get adjusted currents from potential difference.""" + args = _make_args() + # Outer square contains inner square + outer = np.array([[0.0, 0.0], [3.0, 0.0], [3.0, 3.0], [0.0, 3.0], [0.0, 0.0]]) + inner = np.array([[1.0, 1.0], [2.0, 1.0], [2.0, 2.0], [1.0, 2.0], [1.0, 1.0]]) + wp = [1.0, 1.0] # Will be overwritten for outer + result, nested, fv, nc = check_and_compute_nested_WP_currents([outer, inner], wp, args, plot=False) + self.assertIsNotNone(nested) + self.assertTrue(nested[0, 1]) # outer contains inner + self.assertEqual(len(result), 2) + + def test_nested_contours_plot_true(self): + """check_and_compute_nested_WP_currents with plot=True runs without error.""" + from unittest.mock import patch + args = _make_args() + outer = np.array([[0.0, 0.0], [3.0, 0.0], [3.0, 3.0], [0.0, 3.0], [0.0, 0.0]]) + inner = np.array([[1.0, 1.0], [2.0, 1.0], [2.0, 2.0], [1.0, 2.0], [1.0, 1.0]]) + wp = [1.0, 1.0] + with patch('simsopt.util.winding_surface_helper_functions.plt.show'): + result, nested, _, _ = check_and_compute_nested_WP_currents( + [outer, inner], wp, args, plot=True + ) + self.assertEqual(len(result), 2) + self.assertTrue(nested[0, 1]) + + +class TestSIMSOPTLineXYZRZ(unittest.TestCase): + """Tests for SIMSOPT_line_XYZ_RZ.""" + + def test_output_shape(self): + """Output arrays have correct length for SurfaceRZFourier.""" + from simsopt.geo import SurfaceRZFourier + surf = SurfaceRZFourier(nfp=4, mpol=2, ntor=2, stellsym=True) + surf.set_dofs(np.zeros(len(surf.get_dofs()))) + npts = 8 + theta = np.linspace(0, 2 * np.pi, npts, endpoint=False) + zeta = np.linspace(0, 2 * np.pi / 4, npts, endpoint=False) + X, Y, R, Z = _SIMSOPT_line_XYZ_RZ(surf, (theta, zeta)) + self.assertEqual(len(X), npts) + self.assertEqual(len(R), npts) + + def test_points_on_surface(self): + """Points from SIMSOPT_line_XYZ_RZ lie on the surface (uses gamma_lin).""" + from simsopt.geo import SurfaceRZFourier + surf = SurfaceRZFourier(nfp=4, mpol=2, ntor=2, stellsym=True) + npts = 12 + theta = np.linspace(0, 2 * np.pi, npts, endpoint=False) + zeta = np.linspace(0, 2 * np.pi / 4, npts, endpoint=False) + X, Y, R, Z = _SIMSOPT_line_XYZ_RZ(surf, (theta, zeta)) + quadpoints_theta = np.mod(theta / (2 * np.pi), 1.0) + quadpoints_phi = np.mod(zeta / (2 * np.pi), 1.0) + gamma = np.zeros((npts, 3)) + surf.gamma_lin(gamma, quadpoints_phi, quadpoints_theta) + np.testing.assert_allclose(X, gamma[:, 0], atol=1e-14) + np.testing.assert_allclose(Y, gamma[:, 1], atol=1e-14) + np.testing.assert_allclose(Z, gamma[:, 2], atol=1e-14) + + +class TestWriteToCurve(unittest.TestCase): + """Tests for writeToCurve.""" + + def test_curve_created(self): + """writeToCurve creates a CurveXYZFourier that approximates the contour.""" + from simsopt.geo import SurfaceRZFourier + surf = SurfaceRZFourier(nfp=4, mpol=2, ntor=2, stellsym=True) + npts = 32 + theta = np.linspace(0, 2 * np.pi, npts, endpoint=False) + zeta = np.linspace(0, 2 * np.pi / 4, npts, endpoint=False) + contour = np.column_stack([theta, zeta]) + curve = _writeToCurve(contour, [], fourier_trunc=10, winding_surface=surf) + gamma = curve.gamma() + self.assertEqual(gamma.shape[1], 3) + # Curve should be in a reasonable range (surface has R~1) + self.assertGreater(np.max(np.linalg.norm(gamma, axis=1)), 0.5) + self.assertLess(np.max(np.linalg.norm(gamma, axis=1)), 2.0) + + def test_writeToCurve_plotting_and_fix_stellarator_symmetry(self): + """writeToCurve with plotting_args=(1,) and fix_stellarator_symmetry=True.""" + from simsopt.geo import SurfaceRZFourier + from unittest.mock import patch + surf = SurfaceRZFourier(nfp=4, mpol=2, ntor=2, stellsym=True) + n_pts = 50 + t = np.linspace(0, 2 * np.pi, n_pts, endpoint=False) + theta, zeta = t, t # Helical: closes in 3D + contour = np.column_stack([theta, zeta]) + contour_xyz = _SIMSOPT_line_XYZ_RZ(surf, [theta, zeta]) + with patch('simsopt.util.winding_surface_helper_functions.plt.show'): + curve = _writeToCurve( + contour, contour_xyz, fourier_trunc=10, + plotting_args=(1,), fix_stellarator_symmetry=True + ) + gamma = curve.gamma() + self.assertEqual(gamma.shape[1], 3) + + +def _max_distance_to_contour(points, contour_xyz): + """Max distance from any point to nearest point on contour (one-sided Hausdorff).""" + # contour_xyz is [X, Y, R, Z] or [X, Y, Z] + X, Y = contour_xyz[0], contour_xyz[1] + Z = contour_xyz[3] if len(contour_xyz) > 3 else contour_xyz[2] + xyz = np.column_stack([X, Y, Z]) + dists = np.min(np.linalg.norm(points[:, None, :] - xyz[None, :, :], axis=2), axis=1) + return np.max(dists) + + +class TestWriteToCurveHelicalClosed(unittest.TestCase): + """Test writeToCurve with helical contour (3D closure, not θζ closure).""" + + def test_writeToCurve_helical_contour_approximates(self): + """writeToCurve with 3D-closed helical contour approximates the input.""" + from simsopt.geo import SurfaceRZFourier + surf = SurfaceRZFourier(nfp=4, mpol=2, ntor=2, stellsym=True) + n_pts = 150 + # Helical: θ and ζ increase together, contour closes in 3D but not in (θ,ζ) + t = np.linspace(0, 2 * np.pi, n_pts, endpoint=False) + theta = t + zeta = t + contour = np.column_stack([theta, zeta]) + contour_xyz = _SIMSOPT_line_XYZ_RZ(surf, [theta, zeta]) + curve = _writeToCurve(contour, contour_xyz, fourier_trunc=15, winding_surface=surf) + gamma = curve.gamma() + X, Y, Z = contour_xyz[0], contour_xyz[1], contour_xyz[3] + err = _max_distance_to_contour(gamma, [X, Y, Z]) + self.assertLess(err, 0.25, msg=f"Curve should approximate helical input; max dist={err:.4f}") + + +class TestHelicalExtensionCloses(unittest.TestCase): + """Sanity checks: extended helical contour closes in 3D; curve returns to start.""" + + def test_extended_helical_contour_closes_in_3d(self): + """Helical contour extended by +j*2π/nfp and closing point closes in 3D.""" + from simsopt.geo import SurfaceRZFourier + surf = SurfaceRZFourier(nfp=4, mpol=2, ntor=2, stellsym=True) + nfp = 4 + n_pts = 50 + zeta_1fp = np.linspace(0, 2 * np.pi / nfp, n_pts, endpoint=False) + theta_1fp = 0.5 * zeta_1fp + contour_1fp = np.column_stack([theta_1fp, zeta_1fp]) + L = len(contour_1fp) + d = 2 * np.pi / nfp + Contour = np.zeros((L * nfp, 2)) + for j in range(nfp): + j1, j2 = j * L, (j + 1) * L + Contour[j1:j2, 0] = contour_1fp[:, 0] + Contour[j1:j2, 1] = contour_1fp[:, 1] + j * d + # Force closure: append first point with ζ+2π*nfp (surface period 2π*nfp in zeta) + first_pt = np.array([[contour_1fp[0, 0], contour_1fp[0, 1] + 2 * np.pi * nfp]]) + Contour = np.vstack([Contour, first_pt]) + X, Y, R, Z = _SIMSOPT_line_XYZ_RZ(surf, [Contour[:, 0], Contour[:, 1]]) + gap = np.linalg.norm(np.array([X[-1], Y[-1], Z[-1]]) - np.array([X[0], Y[0], Z[0]])) + self.assertLess(gap, 0.01, msg=f"Helical contour must close in 3D; gap={gap:.4f}") + + def test_helical_curve_returns_to_start(self): + """Curve from closed helical contour has gamma(0) ≈ gamma(1).""" + from simsopt.geo import SurfaceRZFourier + surf = SurfaceRZFourier(nfp=4, mpol=2, ntor=2, stellsym=True) + nfp = 4 + n_pts = 40 + zeta_1fp = np.linspace(0, 2 * np.pi / nfp, n_pts, endpoint=False) + theta_1fp = 0.5 * zeta_1fp + contour_1fp = np.column_stack([theta_1fp, zeta_1fp]) + L = len(contour_1fp) + d = 2 * np.pi / nfp + Contour = np.zeros((L * nfp, 2)) + for j in range(nfp): + j1, j2 = j * L, (j + 1) * L + Contour[j1:j2, 0] = contour_1fp[:, 0] + Contour[j1:j2, 1] = contour_1fp[:, 1] + j * d + first_pt = np.array([[contour_1fp[0, 0], contour_1fp[0, 1] + 2 * np.pi * nfp]]) + Contour = np.vstack([Contour, first_pt]) + contour_xyz = list(_SIMSOPT_line_XYZ_RZ(surf, [Contour[:, 0], Contour[:, 1]])) + curve = _writeToCurve(Contour, contour_xyz, fourier_trunc=12, winding_surface=surf) + gamma = curve.gamma() + closure_err = np.linalg.norm(gamma[0] - gamma[-1]) + self.assertLess(closure_err, 0.05, msg=f"Helical curve must close; |gamma[0]-gamma[-1]|={closure_err:.4f}") + + +class TestCutCoilsHelicalClosure(unittest.TestCase): + """Integration test: cut_coils helical coils close (run with regcoil data).""" + + def test_cut_coils_helical_coils_close(self): + """Helical coils from cut_coils close: first point ≈ last point.""" + fpath = TEST_DIR / "regcoil_out.hsx.nc" + if not fpath.exists(): + self.skipTest(f"Test file not found: {fpath}") + from simsopt.util import run_cut_coils + coils = run_cut_coils( + surface_filename=fpath, + ilambda=6, + single_valued=False, + show_final_coilset=False, + show_plots=False, + write_coils_to_file=False, + ) + # Check each coil's curve closes (helical coils are typically last) + for i, coil in enumerate(coils): + gamma = coil.curve.gamma() + closure_err = np.linalg.norm(gamma[0] - gamma[-1]) + self.assertLess( + closure_err, 0.1, + msg=f"Coil {i+1} must close; |gamma[0]-gamma[-1]|={closure_err:.4f}" + ) + + def test_cut_coils_no_center_crossing(self): + """Coils must not cross through the center (R_min > threshold); detects X artifact.""" + fpath = TEST_DIR / "regcoil_out.hsx.nc" + if not fpath.exists(): + self.skipTest(f"Test file not found: {fpath}") + from simsopt.util import run_cut_coils + coils = run_cut_coils( + surface_filename=fpath, + ilambda=6, + single_valued=False, + show_final_coilset=False, + show_plots=False, + write_coils_to_file=False, + curve_fourier_cutoff=40, + ) + # Coils lie on winding surface; R = sqrt(x^2+y^2) must stay away from axis + R_min_acceptable = 0.3 # Winding surface is typically R > 0.5; allow some margin + for i, coil in enumerate(coils): + gamma = coil.curve.gamma() + R = np.sqrt(gamma[:, 0]**2 + gamma[:, 1]**2) + R_min = np.min(R) + self.assertGreater( + R_min, R_min_acceptable, + msg=f"Coil {i+1} crosses center: R_min={R_min:.4f} (expected >{R_min_acceptable})" + ) + + def test_cut_coils_on_winding_surface(self): + """Extracted coils must lie on the winding surface within tolerance (helical stitching).""" + fpath = TEST_DIR / "regcoil_out.hsx.nc" + if not fpath.exists(): + self.skipTest(f"Test file not found: {fpath}") + from scipy.spatial import cKDTree + from simsopt.util import run_cut_coils + cpst, s_coil_fp, s_coil_full, s_plasma_fp, s_plasma_full = _load_CP_and_geometries( + str(fpath), plot_flags=(0, 0, 0, 0) + ) + coils = run_cut_coils( + surface_filename=fpath, + ilambda=6, + single_valued=False, + show_final_coilset=False, + show_plots=False, + write_coils_to_file=False, + curve_fourier_cutoff=40, + ) + # Dense surface sampling for distance check + ntheta, nphi = 64, 64 * s_coil_full.nfp + theta = np.linspace(0, 1, ntheta, endpoint=False) + phi = np.linspace(0, 1, nphi, endpoint=False) + data = np.zeros((ntheta * nphi, 3)) + s_coil_full.gamma_lin(data, np.repeat(phi, ntheta), np.tile(theta, nphi)) + surf_pts = data + tree = cKDTree(surf_pts) + dist_tol = 0.08 # Coils within 8cm of winding surface (Fourier fit tolerance) + for i, coil in enumerate(coils): + gamma = coil.curve.gamma() + d, _ = tree.query(gamma, k=1) + max_d = np.max(d) + argmax_d = np.argmax(d) + self.assertLess( + max_d, dist_tol, + msg=f"Coil {i+1} has points off surface: max_dist={max_d:.4f} at pt {argmax_d} " + f"(tol={dist_tol}); xyz={gamma[argmax_d]}" + ) + + +class TestRunCutCoilsFlags(unittest.TestCase): + """Tests for run_cut_coils with different flags: show_final_coilset, write_coils_to_file.""" + + def test_run_cut_coils_show_final_coilset_true(self): + """run_cut_coils with show_final_coilset=True runs without error.""" + fpath = TEST_DIR / "regcoil_out.hsx.nc" + if not fpath.exists(): + self.skipTest(f"Test file not found: {fpath}") + from simsopt.util import run_cut_coils + coils = run_cut_coils( + surface_filename=fpath, + ilambda=6, + single_valued=False, + show_final_coilset=True, + show_plots=False, + write_coils_to_file=False, + ) + self.assertGreater(len(coils), 0) + + def test_run_cut_coils_write_coils_to_file_true(self): + """run_cut_coils with write_coils_to_file=True creates coils.json.""" + from monty.tempfile import ScratchDir + fpath = TEST_DIR / "regcoil_out.hsx.nc" + if not fpath.exists(): + self.skipTest(f"Test file not found: {fpath}") + from simsopt.util import run_cut_coils + with ScratchDir("."): + run_cut_coils( + surface_filename=str(fpath), + ilambda=6, + single_valued=False, + show_final_coilset=False, + show_plots=False, + write_coils_to_file=True, + output_path="cut_coils_output", + ) + coils_path = Path("cut_coils_output") / "coils.json" + self.assertTrue(coils_path.exists(), f"Expected {coils_path} to exist") + self.assertGreater(coils_path.stat().st_size, 0) + + def test_run_cut_coils_both_flags_true(self): + """run_cut_coils with show_final_coilset=True and write_coils_to_file=True.""" + from monty.tempfile import ScratchDir + fpath = TEST_DIR / "regcoil_out.hsx.nc" + if not fpath.exists(): + self.skipTest(f"Test file not found: {fpath}") + from simsopt.util import run_cut_coils + with ScratchDir("."): + coils = run_cut_coils( + surface_filename=str(fpath), + ilambda=6, + single_valued=False, + show_final_coilset=True, + show_plots=False, + write_coils_to_file=True, + output_path="cut_coils_output", + ) + self.assertGreater(len(coils), 0) + self.assertTrue((Path("cut_coils_output") / "coils.json").exists()) + + def test_run_cut_coils_single_valued_true(self): + """run_cut_coils with single_valued=True exercises _map_data_full_torus_3x1 path.""" + fpath = TEST_DIR / "regcoil_out.hsx.nc" + if not fpath.exists(): + self.skipTest(f"Test file not found: {fpath}") + from simsopt.util import run_cut_coils + coils = run_cut_coils( + surface_filename=fpath, + ilambda=6, + single_valued=True, + show_final_coilset=False, + show_plots=False, + write_coils_to_file=False, + ) + self.assertGreaterEqual(len(coils), 0) + + def test_run_cut_coils_interactive_true(self): + """run_cut_coils with interactive=True runs without error (returns [] when no contours selected).""" + fpath = TEST_DIR / "regcoil_out.hsx.nc" + if not fpath.exists(): + self.skipTest(f"Test file not found: {fpath}") + import matplotlib + matplotlib.use("Agg") + from simsopt.util import run_cut_coils + coils = run_cut_coils( + surface_filename=fpath, + ilambda=6, + interactive=True, + show_final_coilset=False, + show_plots=False, + write_coils_to_file=False, + ) + # In non-interactive CI, no contours are selected so returns [] + self.assertIsInstance(coils, list) + + +class TestLoadSimsoptRegcoilData(unittest.TestCase): + """Tests for load_simsopt_regcoil_data.""" + + def test_load_simsopt_regcoil_file(self): + """load_simsopt_regcoil_data loads simsopt-regcoil format.""" + from unittest.mock import MagicMock, patch + + def _mock_var(val): + m = MagicMock() + m.__getitem__ = MagicMock(return_value=val) + return m + + mock_f = MagicMock() + mock_f.__enter__ = MagicMock(return_value=mock_f) + mock_f.__exit__ = MagicMock(return_value=False) + mock_f.variables = { + 'nfp': _mock_var(4), + 'ntheta_coil': _mock_var(32), + 'nzeta_coil': _mock_var(16), + 'theta_coil': _mock_var(np.linspace(0, 2*np.pi, 32, endpoint=False)), + 'zeta_coil': _mock_var(np.linspace(0, 2*np.pi/4, 16, endpoint=False)), + 'r_coil': _mock_var(np.ones(32*16)), + 'xm_coil': _mock_var(np.array([0, 1])), + 'xn_coil': _mock_var(np.array([0, 1])*4), + 'xm_potential': _mock_var(np.array([0, 1])), + 'xn_potential': _mock_var(np.array([0, 1])*4), + 'net_poloidal_current_amperes': _mock_var(np.array([0.0])), + 'net_toroidal_current_amperes': _mock_var(np.array([0.0])), + 'single_valued_current_potential_mn': _mock_var(np.zeros((1, 2))), + 'single_valued_current_potential_thetazeta': _mock_var(np.zeros((1, 16, 32))), + 'lambda': _mock_var(np.array([1e-10])), + 'K2': _mock_var(np.zeros((1, 16, 32))), + 'chi2_B': _mock_var(np.array([0.1])), + 'chi2_K': _mock_var(np.array([0.01])), + } + with patch('simsopt.util.winding_surface_helper_functions.netcdf_file', return_value=mock_f): + data = _load_simsopt_regcoil_data('/fake/path.nc', sparse=False) + self.assertEqual(len(data), 19) + self.assertEqual(data[0], 32) + self.assertEqual(data[1], 16) + + def test_load_simsopt_regcoil_sparse(self): + """load_simsopt_regcoil_data with sparse=True loads L1 variables.""" + from unittest.mock import MagicMock, patch + + def _mock_var(val): + m = MagicMock() + m.__getitem__ = MagicMock(return_value=val) + return m + + mock_f = MagicMock() + mock_f.__enter__ = MagicMock(return_value=mock_f) + mock_f.__exit__ = MagicMock(return_value=False) + mock_f.variables = { + 'nfp': _mock_var(4), 'ntheta_coil': _mock_var(16), 'nzeta_coil': _mock_var(8), + 'theta_coil': _mock_var(np.linspace(0, 2*np.pi, 16, endpoint=False)), + 'zeta_coil': _mock_var(np.linspace(0, 2*np.pi/4, 8, endpoint=False)), + 'r_coil': _mock_var(np.ones(128)), 'xm_coil': _mock_var(np.array([0, 1])), + 'xn_coil': _mock_var(np.array([0, 1])*4), 'xm_potential': _mock_var(np.array([0, 1])), + 'xn_potential': _mock_var(np.array([0, 1])*4), + 'net_poloidal_current_amperes': _mock_var(np.array([0.0])), + 'net_toroidal_current_amperes': _mock_var(np.array([0.0])), + 'single_valued_current_potential_mn_l1': _mock_var(np.zeros((1, 2))), + 'single_valued_current_potential_thetazeta_l1': _mock_var(np.zeros((1, 8, 16))), + 'lambda': _mock_var(np.array([1e-10])), 'K2_l1': _mock_var(np.zeros((1, 8, 16))), + 'chi2_B_l1': _mock_var(np.array([0.1])), 'chi2_K_l1': _mock_var(np.array([0.01])), + } + with patch('simsopt.util.winding_surface_helper_functions.netcdf_file', return_value=mock_f): + data = _load_simsopt_regcoil_data('/fake/path.nc', sparse=True) + self.assertEqual(len(data), 19) + self.assertEqual(data[0], 16) + + +class TestLoadCPAndGeometries(unittest.TestCase): + """Tests for load_CP_and_geometries.""" + + def test_load_CP_and_geometries_loadDOFsProperly_false(self): + """load_CP_and_geometries with loadDOFsProperly=False skips DOF copy.""" + fpath = TEST_DIR / "regcoil_out.w7x_infty.nc" + if not fpath.exists(): + fpath = TEST_DIR / "regcoil_out.hsx.nc" + if not fpath.exists(): + self.skipTest("No regcoil test file found") + result = _load_CP_and_geometries(str(fpath), plot_flags=(0, 0, 0, 0), loadDOFsProperly=False) + self.assertEqual(len(result), 5) + cpst, s_coil_fp, s_coil_full, s_plasma_fp, s_plasma_full = result + self.assertIsNotNone(cpst) + self.assertIsNotNone(s_coil_fp) + + def test_load_CP_and_geometries_loadDOFsProperly_true(self): + """load_CP_and_geometries with loadDOFsProperly=True copies surface DOFs.""" + fpath = TEST_DIR / "regcoil_out.w7x_infty.nc" + if not fpath.exists(): + fpath = TEST_DIR / "regcoil_out.hsx.nc" + if not fpath.exists(): + self.skipTest("No regcoil test file found") + result = _load_CP_and_geometries(str(fpath), plot_flags=(0, 0, 0, 0), loadDOFsProperly=True) + self.assertEqual(len(result), 5) + cpst, s_coil_fp, s_coil_full, s_plasma_fp, s_plasma_full = result + self.assertIsNotNone(s_coil_fp) + + def test_load_CP_and_geometries_with_plot_flags(self): + """load_CP_and_geometries with plot_flags calls plot (mocked).""" + fpath = TEST_DIR / "regcoil_out.w7x_infty.nc" + if not fpath.exists(): + fpath = TEST_DIR / "regcoil_out.hsx.nc" + if not fpath.exists(): + self.skipTest("No regcoil test file found") + from unittest.mock import patch + with patch('simsopt.geo.plot') as mock_plot: + result = _load_CP_and_geometries(str(fpath), plot_flags=(1, 1, 1, 1)) + self.assertEqual(len(result), 5) + mock_plot.assert_called() + + +class TestLoadRegcoilData(unittest.TestCase): + """Tests for load_regcoil_data with real file.""" + + def test_load_regcoil_file(self): + """Load legacy REGCOIL file returns expected structure.""" + fpath = TEST_DIR / "regcoil_out.hsx.nc" + if not fpath.exists(): + self.skipTest(f"Test file not found: {fpath}") + data = _load_regcoil_data(str(fpath)) + self.assertEqual(len(data), 19) + ntheta, nzeta, theta, zeta, r, xm, xn, xmp, xnp, nfp, Ip, It = data[:12] + self.assertGreater(ntheta, 0) + self.assertGreater(nzeta, 0) + self.assertEqual(len(theta), ntheta) + self.assertEqual(len(zeta), nzeta) + lambdas, chi2_B, chi2_K, K2 = data[15:19] + self.assertGreater(len(lambdas), 0) + + +class TestSetAxesEqual(unittest.TestCase): + """Tests for set_axes_equal.""" + + def test_set_axes_equal_3d(self): + """set_axes_equal adjusts 3D axes to equal scale.""" + import matplotlib.pyplot as plt + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.plot([0, 1], [0, 1], [0, 1]) + set_axes_equal(ax) + xlim = ax.get_xlim3d() + ylim = ax.get_ylim3d() + zlim = ax.get_zlim3d() + xr = abs(xlim[1] - xlim[0]) + yr = abs(ylim[1] - ylim[0]) + zr = abs(zlim[1] - zlim[0]) + self.assertAlmostEqual(xr, yr) + self.assertAlmostEqual(yr, zr) + plt.close() + + +class TestContourPaths(unittest.TestCase): + """Tests for _contour_paths (matplotlib version-agnostic).""" + + def test_contour_paths_allsegs(self): + """_contour_paths returns paths from contour object.""" + import matplotlib.pyplot as plt + x = np.linspace(0, 2 * np.pi, 20) + y = np.linspace(0, 2 * np.pi / 4, 15) + z = np.outer(np.sin(y), np.cos(x)) + cs = plt.contour(x, y, z, levels=[0]) + paths = _contour_paths(cs, 0) + self.assertIsInstance(paths, list) + if len(paths) > 0: + self.assertIsInstance(paths[0], np.ndarray) + plt.close() + + def test_contour_paths_fallback_collections(self): + """_contour_paths uses collections.get_paths when allsegs not available.""" + from unittest.mock import MagicMock + mock_path = MagicMock() + mock_path.vertices = np.array([[0, 0], [1, 1]]) + mock_collection = MagicMock() + mock_collection.get_paths.return_value = [mock_path] + mock_cdata = type('CData', (), {'collections': [mock_collection]})() + paths = _contour_paths(mock_cdata, 0) + self.assertEqual(len(paths), 1) + np.testing.assert_array_equal(paths[0], np.array([[0, 0], [1, 1]])) + + +class TestMapDataFullTorus3x1(unittest.TestCase): + """Tests for _map_data_full_torus_3x1.""" + + def test_output_shapes(self): + """_map_data_full_torus_3x1 returns extended grid with 3x1 tiling in zeta.""" + xdata = np.linspace(0, 2 * np.pi, 8, endpoint=False) + ydata = np.linspace(0, 2 * np.pi / 4, 4, endpoint=False) + zdata = np.outer(np.sin(ydata), np.cos(xdata)) + xN, yN, ret = _map_data_full_torus_3x1(xdata, ydata, zdata, Ip=0, It=0) + self.assertEqual(len(xN), len(xdata)) + self.assertEqual(len(yN), 3 * len(ydata)) + self.assertEqual(ret.shape, (3 * len(ydata), len(xdata))) + + def test_secular_term_added(self): + """With Ip, It non-zero, secular term is added (data differs from tiled zdata).""" + xdata = np.linspace(0, 2 * np.pi, 4, endpoint=False) + ydata = np.linspace(0, 2 * np.pi / 4, 2, endpoint=False) + zdata = np.zeros((2, 4)) + xN, yN, ret = _map_data_full_torus_3x1(xdata, ydata, zdata, Ip=2 * np.pi, It=4 * np.pi) + # With NSV = (It/2π)*θ + (Ip/2π)*ζ, ret should be non-zero + self.assertTrue(np.any(ret != 0)) + + +class TestRunCutCoilsComputeNWPCurrents(unittest.TestCase): + """Tests for _run_cut_coils_compute_NWP_currents.""" + + def test_empty_open_contours(self): + """Empty open_contours returns empty list.""" + args = _make_args() + theta = np.linspace(0, 2 * np.pi, 8, endpoint=False) + zeta = np.linspace(0, 2 * np.pi / 4, 4, endpoint=False) + _, _, cp_sv, _, _, _ = _genCPvals((0, 2 * np.pi), (0, 2 * np.pi / 4), (8, 4), args) + result = _run_cut_coils_compute_NWP_currents( + [], True, 1.0, theta, zeta, cp_sv, args + ) + self.assertEqual(result, []) + + def test_single_open_contour(self): + """Single open contour returns [totalCurrent].""" + args = _make_args() + theta = np.linspace(0, 2 * np.pi, 8, endpoint=False) + zeta = np.linspace(0, 2 * np.pi / 4, 4, endpoint=False) + _, _, cp_sv, _, _, _ = _genCPvals((0, 2 * np.pi), (0, 2 * np.pi / 4), (8, 4), args) + contour = np.array([[0.5, 0.3], [1.5, 0.4]]) + result = _run_cut_coils_compute_NWP_currents( + [contour], True, 10.0, theta, zeta, cp_sv, args + ) + self.assertEqual(len(result), 1) + self.assertAlmostEqual(result[0], 10.0) + + + def test_multi_open_contours_multi_valued(self): + """Multiple open contours with single_valued=False uses level diff.""" + args = _make_args(Ip=2 * np.pi, It=4 * np.pi) + theta = np.linspace(0, 2 * np.pi, 20, endpoint=False) + zeta = np.linspace(0, 2 * np.pi / 4, 15, endpoint=False) + _, _, cp_full, _, _, _ = _genCPvals((0, 2 * np.pi), (0, 2 * np.pi / 4), (20, 15), args) + c1 = np.array([[0.5, 0.3], [1.5, 0.4]]) + c2 = np.array([[1.0, 0.5], [2.0, 0.6]]) + result = _run_cut_coils_compute_NWP_currents( + [c1, c2], False, 1.0, theta, zeta, cp_full, args + ) + self.assertEqual(len(result), 2) + self.assertAlmostEqual(sum(result), 1.0) + + def test_multi_open_contours_single_valued(self): + """Multiple open contours with single_valued=True uses halfway-contour method.""" + import matplotlib + matplotlib.use("Agg") + args = _make_args(Ip=2 * np.pi, It=4 * np.pi) + theta = np.linspace(0, 2 * np.pi, 30, endpoint=False) + zeta = np.linspace(0, 2 * np.pi / 4, 20, endpoint=False) + _, _, cp_sv, _, _, _ = _genCPvals((0, 2 * np.pi), (0, 2 * np.pi / 4), (30, 20), args) + c1 = np.array([[0.5, 0.3], [1.5, 0.4]]) + c2 = np.array([[1.5, 0.5], [2.5, 0.6]]) + c3 = np.array([[2.5, 0.7], [3.5, 0.8]]) + result = _run_cut_coils_compute_NWP_currents( + [c1, c2, c3], True, 1.0, theta, zeta, cp_sv, args + ) + self.assertEqual(len(result), 3) + for r in result: + self.assertGreater(r, 0, "Each contour current should be positive") + + +class TestRunCutCoilsSelectContoursNonInteractive(unittest.TestCase): + """Tests for _run_cut_coils_select_contours_non_interactive.""" + + def test_select_by_points(self): + """Selection by points returns contours near those points.""" + import matplotlib.pyplot as plt + args = _make_args() + theta = np.linspace(0, 2 * np.pi, 20, endpoint=False) + zeta = np.linspace(0, 2 * np.pi / 4, 15, endpoint=False) + _, _, cp, _, _, _ = _genCPvals((0, 2 * np.pi), (0, 2 * np.pi / 4), (20, 15), args) + fig, ax = plt.subplots() + points = [[1.0, 0.5], [2.0, 0.8]] + contours = _run_cut_coils_select_contours_non_interactive( + theta, zeta, cp, args, "free", points=points, levels=None, + contours_per_period=None, ax=ax + ) + self.assertGreater(len(contours), 0) + plt.close("all") + + def test_select_by_levels(self): + """Selection by levels returns contours at those levels.""" + import matplotlib.pyplot as plt + args = _make_args() + theta = np.linspace(0, 2 * np.pi, 20, endpoint=False) + zeta = np.linspace(0, 2 * np.pi / 4, 15, endpoint=False) + _, _, cp, _, _, _ = _genCPvals((0, 2 * np.pi), (0, 2 * np.pi / 4), (20, 15), args) + fig, ax = plt.subplots() + levels = [np.min(cp) + 0.2 * (np.max(cp) - np.min(cp))] + contours = _run_cut_coils_select_contours_non_interactive( + theta, zeta, cp, args, "free", points=None, levels=levels, + contours_per_period=None, ax=ax + ) + self.assertGreaterEqual(len(contours), 0) + plt.close("all") + + def test_select_by_contours_per_period(self): + """Selection by contours_per_period distributes contours.""" + import matplotlib.pyplot as plt + args = _make_args() + theta = np.linspace(0, 2 * np.pi, 20, endpoint=False) + zeta = np.linspace(0, 2 * np.pi / 4, 15, endpoint=False) + _, _, cp, _, _, _ = _genCPvals((0, 2 * np.pi), (0, 2 * np.pi / 4), (20, 15), args) + fig, ax = plt.subplots() + contours = _run_cut_coils_select_contours_non_interactive( + theta, zeta, cp, args, "free", points=None, levels=None, + contours_per_period=3, ax=ax + ) + self.assertGreaterEqual(len(contours), 0) + plt.close("all") + + def test_select_default_points(self): + """Default (no points/levels/contours_per_period) uses demo points.""" + import matplotlib.pyplot as plt + args = _make_args() + theta = np.linspace(0, 2 * np.pi, 20, endpoint=False) + zeta = np.linspace(0, 2 * np.pi / 4, 15, endpoint=False) + _, _, cp, _, _, _ = _genCPvals((0, 2 * np.pi), (0, 2 * np.pi / 4), (20, 15), args) + fig, ax = plt.subplots() + contours = _run_cut_coils_select_contours_non_interactive( + theta, zeta, cp, args, "free", points=None, levels=None, + contours_per_period=None, ax=ax + ) + self.assertEqual(len(contours), 3) + plt.close("all") + + +class TestLoadSurfaceDofsProperly(unittest.TestCase): + """Tests for _load_surface_dofs_properly.""" + + def test_load_surface_dofs_properly(self): + """_load_surface_dofs_properly copies surface DOFs and returns s_new.""" + from simsopt.geo import SurfaceRZFourier + s = SurfaceRZFourier(nfp=4, mpol=2, ntor=2, stellsym=True) + dofs = s.get_dofs() + s.set_dofs(np.random.randn(len(dofs)) * 0.01) + s_new = SurfaceRZFourier(nfp=4, mpol=2, ntor=2, stellsym=True) + s_new = s_new.from_nphi_ntheta(nfp=4, ntheta=8, nphi=16, mpol=2, ntor=2, stellsym=True, range='field period') + result = _load_surface_dofs_properly(s, s_new) + self.assertIsNotNone(result) + self.assertIs(result, s_new) + # Result should have non-zero DOFs where s has them (at least rc, zs) + self.assertTrue(np.any(result.get_dofs() != 0)) + + +class TestMakeOnclickOnpickOnKey(unittest.TestCase): + """Tests for _make_onclick, _make_onpick, _make_on_key.""" + + def test_make_onclick_returns_callable(self): + """_make_onclick returns a callable.""" + import matplotlib.pyplot as plt + fig, ax = plt.subplots() + args = _make_args() + contours = [] + theta = np.linspace(0, 2 * np.pi, 10, endpoint=False) + zeta = np.linspace(0, 2 * np.pi / 4, 8, endpoint=False) + _, _, cp, _, _, _ = _genCPvals((0, 2 * np.pi), (0, 2 * np.pi / 4), (10, 8), args) + handler = _make_onclick(ax, args, contours, theta, zeta, cp) + self.assertTrue(callable(handler)) + plt.close() + + def test_make_onpick_returns_callable(self): + """_make_onpick returns a callable.""" + contours = [] + handler = _make_onpick(contours) + self.assertTrue(callable(handler)) + + def test_make_on_key_returns_callable(self): + """_make_on_key returns a callable.""" + contours = [] + handler = _make_on_key(contours) + self.assertTrue(callable(handler)) + + def test_make_onclick_dblclick_appends_contour(self): + """_make_onclick on dblclick finds contour and appends to list.""" + import matplotlib.pyplot as plt + from unittest.mock import MagicMock + fig, ax = plt.subplots() + args = _make_args() + contours = [] + theta = np.linspace(0, 2 * np.pi, 15, endpoint=False) + zeta = np.linspace(0, 2 * np.pi / 4, 10, endpoint=False) + _, _, cp, _, _, _ = _genCPvals((0, 2 * np.pi), (0, 2 * np.pi / 4), (15, 10), args) + handler = _make_onclick(ax, args, contours, theta, zeta, cp) + event = MagicMock() + event.dblclick = True + event.xdata, event.ydata = 1.0, 0.3 + handler(event) + self.assertGreater(len(contours), 0) + self.assertEqual(contours[0].shape[1], 2) + plt.close() + + def test_make_onpick_sets_picked_object(self): + """_make_onpick stores picked artist on axes.""" + import matplotlib.pyplot as plt + from unittest.mock import MagicMock + fig, ax = plt.subplots() + ax.plot([0, 1], [0, 1]) + handler = _make_onpick([]) + event = MagicMock() + event.artist = ax.lines[0] + handler(event) + self.assertIs(plt.gca().picked_object, event.artist) + plt.close() + + def test_make_on_key_delete_removes_picked(self): + """_make_on_key on 'delete' removes picked contour from list.""" + import matplotlib.pyplot as plt + from unittest.mock import MagicMock + fig, ax = plt.subplots() + line_data = np.array([[0.5, 0.3], [1.0, 0.4]]) + line, = ax.plot(line_data[:, 0], line_data[:, 1]) + contours = [line_data.copy()] + handler = _make_on_key(contours) + ax.picked_object = line + event = MagicMock() + event.key = 'delete' + handler(event) + self.assertEqual(len(contours), 0) + self.assertIsNone(ax.picked_object) + plt.close() + + +if __name__ == "__main__": + unittest.main()