diff --git a/examples/3_Advanced/analysis_tools.py b/examples/3_Advanced/analysis_tools.py new file mode 100644 index 000000000..03ab67900 --- /dev/null +++ b/examples/3_Advanced/analysis_tools.py @@ -0,0 +1,645 @@ +import desc +import glob +import imageio +import json +import matplotlib +import matplotlib.pyplot as plt +import numpy as np +import os +import pandas as pd +import plotly.express as px +import scipy +import seaborn as sns +import shutil +import simsopt +import subprocess +from desc.grid import LinearGrid +from desc.geometry import FourierRZToroidalSurface +from desc.equilibrium import Equilibrium +from desc.plotting import plot_boozer_surface +from desc.vmec_utils import ptolemy_identity_fwd +from matplotlib import colors +from numbers import Number +from paretoset import paretoset +from scipy.spatial.distance import cdist +from simsopt._core.optimizable import load +from simsopt.field import (BiotSavart, + InterpolatedField, + SurfaceClassifier, + particles_to_vtk, + compute_fieldlines, + LevelsetStoppingCriterion) +from simsopt.field.force import coil_force, self_force +from simsopt.field.selffield import B_regularized, regularization_circ +from simsopt.geo import QfmResidual, QfmSurface, SurfaceRZFourier, Volume +from simsopt.util import comm_world + +π = np.pi + +############################################################################### +# I) DATA ANALYSIS +############################################################################### + +def get_dfs(INPUT_DIR='./output/QA/with-force-penalty/1/optimizations/', OUTPUT_DIR=None): + """Returns DataFrames for the raw, filtered, and Pareto data.""" + ### STEP 1: Import raw data + inputs=f"{INPUT_DIR}**/results.json" + results = glob.glob(inputs, recursive=True) + dfs = [] + for results_file in results: + with open(results_file, "r") as f: + data = json.load(f) + # Wrap lists in another list + for key, value in data.items(): + if isinstance(value, list): + data[key] = [value] + dfs.append(pd.DataFrame(data)) + df = pd.concat(dfs, ignore_index=True) + + ### STEP 2: Filter the data + margin_up = 1.05 + margin_low = 0.95 + + df_filtered = df.query( + # ENGINEERING CONSTRAINTS: + f"max_length < {5 * margin_up}" + f"and max_max_κ < {12.00 * margin_up}" + f"and max_MSC < {6.00 * margin_up}" + f"and coil_coil_distance > {0.083 * margin_low}" + f"and coil_surface_distance > {0.166 * margin_low}" + f"and mean_AbsB > 0.22" #prevent coils from becoming detached from LCFS + # FILTERING OUT BAD/UNNECESSARY DATA: + f"and max_arclength_variance < 1e-2" + f"and coil_surface_distance < 0.375" + f"and coil_coil_distance < 0.14" + f"and max_length > 4.0" + f"and normalized_BdotN < {4 * 1e-3}" + f"and max_max_force<50000" + ) + + ### STEP 3: Generate Pareto front and export UUIDs as .txt + pareto_mask = paretoset(df_filtered[["normalized_BdotN", "max_max_force"]], sense=[min, min]) + df_pareto = df_filtered[pareto_mask] + + # Copy pareto fronts to a separate folder + if OUTPUT_DIR is not None: + if os.path.exists(OUTPUT_DIR): + shutil.rmtree(OUTPUT_DIR) + os.makedirs(OUTPUT_DIR, exist_ok=True) + for UUID in df_pareto['UUID']: + SOURCE_DIR = glob.glob(f"/**/{UUID}/", recursive=True)[0] + DEST_DIR = f"{OUTPUT_DIR}{UUID}/" + shutil.copytree(SOURCE_DIR, DEST_DIR) + + ### Return statement + return df, df_filtered, df_pareto + + +def parameter_correlations(df, sort_by='normalized_BdotN', matrix=False, columns_to_drop=None, fontsize=10, annot=True): + matplotlib.rcParams.update({'font.size': fontsize}) + df_sorted = df.sort_values(by=[sort_by]) + if columns_to_drop is None: + columns_to_drop = ['BdotN', 'gradient_norm', 'Jf', 'JF', 'mean_AbsB', + 'max_arclength_variance', 'iterations', 'linking_number', + 'function_evaluations', 'success', 'arclength_weight', + 'R0', 'ntheta', 'nphi', 'ncoils', 'nfp','MSCs', + 'max_forces', 'arclength_variances', 'max_κ', + 'UUID_init', 'message', 'coil_currents', 'UUID', + 'lengths', 'eval_time', 'order', 'dx', 'RMS_forces', 'min_forces'] + df_sorted = df_sorted.drop(columns=columns_to_drop) + + df_correlation = pd.DataFrame({'Parameter': [], 'R^2': [], 'P': [], 'Equation': []}) + for i in range(len(df_sorted.columns)): + series_name = df_sorted.columns[i] + series = (np.array(df_sorted)[:, i]).astype(float) + bdotn_series = np.array(df_sorted[sort_by]) + result = scipy.stats.linregress(bdotn_series, series) + + r = result.rvalue**2 + p = result.pvalue + slope = result.slope + intercept = result.intercept + equation = f"y = {slope:.2e} * x + {intercept:.2e}" + + df_row = {'Parameter': series_name, 'R^2': r, 'P': p, 'Equation': equation} + df_correlation = df_correlation._append(df_row, ignore_index = True) + + if(matrix): + plt.figure(figsize=(9,9)) + matrix = df_sorted.corr() ** 2 + matrix = np.round(matrix,2) + ax = sns.heatmap(matrix, cmap="Blues",annot=annot, cbar_kws={'label': r'$R^2$'}) + ax.set_xticklabels(ax.get_xticklabels(), rotation=45) + plt.title("Corrrelation Matrix") + plt.show() + + return df_correlation.sort_values(by=['R^2'], ascending=False) + + + +############################################################################### +# II) PLOTTING +############################################################################### + +def create_movies(images, OUT_NAME, OUT_PATH): + writer = imageio.get_writer(OUT_PATH + OUT_NAME + ".mp4") + for image in images: + writer.append_data(image) + writer.close() + imageio.mimsave(OUT_PATH + OUT_NAME + ".gif", images, loop=0) + +def force_reduction_plots(UUIDs=["6266c8d4bb25499b899d86e9e3dd2ee2", + "881b166b6ddb4f8ebe4a2b63d6cb9bc1"], + IN_DIR="./output/QA/with-force-penalty/4/pareto/", + regularization=regularization_circ(0.05), + ncoils=5): + fig, axs = plt.subplots(7, 5, figsize=(8,9)) + + for UUID in UUIDs: + biotsavart_path = IN_DIR + UUID + "/biot_savart.json" + bs = load(biotsavart_path) + coils = bs.coils + base_coils = coils[:ncoils] + for i, c in enumerate(base_coils): + phis = 2 * np.pi * c.curve.quadpoints + gammadash = c.curve.gammadash() + gammadash_norm = np.linalg.norm(gammadash, axis=1)[:, None] + tangent = gammadash / gammadash_norm + current = c.current.get_value() * np.ones_like(phis) + + mutual_coils = [coil for coil in coils if coil is not c] + b_mutual = BiotSavart(mutual_coils).set_points(c.curve.gamma()).B() + b_reg = B_regularized(c, regularization) + b_tot = b_reg + b_mutual + b_tot_norm = np.linalg.norm(b_tot, axis=1)[:, None] + b_tot_tangent = b_tot / b_tot_norm + force = np.linalg.norm(coil_force(c, coils, regularization), axis=1)[:, None] + + # plot force + axs[0, i].plot(phis, force) + axs[0, i].set_title(f"Coil {i+1}") + axs[0, i].set(xlabel = r"$\phi$", ylabel=r"$|d\bf{F}$$/d\ell|$") + axs[0, i].set_xticks([0, 2*np.pi]) + + # plot $t\times b$ + t_cross_b = np.cross(tangent, b_tot_tangent, axis=1) #good thru here! + t_cross_b = np.linalg.norm(t_cross_b, axis=1) + axs[1, i].plot(phis, t_cross_b) + axs[1, i].set(xlabel = r"$\phi$", ylabel=r"$|\bf{t}$$\times\bf{B}|$$/B$") + axs[1, i].set_ylim([0, 1]) + axs[1, i].set_xticks([0, 2*np.pi]) + + # plot |B| + axs[2, i].plot(phis, b_tot_norm) + axs[2, i].set(xlabel = r"$\phi$", ylabel=r"$|\bf{B}$$|$") + axs[2, i].set_xticks([0, 2*np.pi]) + + # plot |B_mut|^2 + b_mut_squared = np.linalg.norm(b_mutual, axis=1)[:, None] ** 2 + axs[3, i].plot(phis, b_mut_squared) + axs[3, i].set(xlabel = r"$\phi$", ylabel=r"$|\bf{B}$$_{mutual}|^2$") + axs[3, i].set_xticks([0, 2*np.pi]) + + # plot 2B_mut\cdot B_reg + B_mut_dot_B_reg = 2 * np.einsum('ij,ij->i', b_mutual, b_reg) + axs[4, i].plot(phis, B_mut_dot_B_reg) + axs[4, i].set(xlabel = r"$\phi$", ylabel=r"$2\bf{B}$$_{mutual}\cdot\bf{B}$$_{reg}$") + axs[4, i].set_xticks([0, 2*np.pi]) + + # plot |B_reg|^2 + b_reg_squared = np.linalg.norm(b_reg, axis=1)[:, None] ** 2 + axs[5, i].plot(phis, b_reg_squared) + axs[5, i].set(xlabel = r"$\phi$", ylabel=r"$|\bf{B}$$_{reg}|^2$") + axs[5, i].set_xticks([0, 2*np.pi]) + axs[5, i].set_xticklabels(["0", r"$2\pi$"]) + + # plot current + b_reg_squared = np.linalg.norm(b_reg, axis=1)[:, None] ** 2 + axs[6, i].plot(phis, current) + axs[6, i].set(xlabel = r"$\phi$", ylabel=r"$I$") + axs[6, i].set_xticks([0, 2*np.pi]) + axs[6, i].set_xticklabels(["0", r"$2\pi$"]) + axs[6, i].set_ylim(bottom=0, top=100000) + + return fig, axs + +def pareto_interactive_plt(df, color='coil_surface_distance'): + """Creates an interactive plot of the Pareto front.""" + fig = px.scatter( + df, + x="normalized_BdotN", + y="max_max_force", + color=color, + log_x=True, + width=500, + height=400, + hover_data={ + 'UUID':True, + 'max_max_force':':.2e', + 'coil_surface_distance':':.2f', + 'coil_coil_distance':':.3f', + 'length_target':':.2f', + 'force_threshold':':.2e', + 'cc_threshold':':.2e', + 'cs_threshold':':.2e', + 'length_weight':':.2e', + 'msc_weight':':.2e', + 'cc_weight':':.2e', + 'cs_weight':':.2e', + 'force_weight':':.2e', + 'normalized_BdotN':':.2e', + 'max_MSC':':.2e', + 'max_max_κ':':.2e' + } + ) + + fig.update_layout( + margin=dict(l=20, r=20, t=40, b=20), + plot_bgcolor='white' + ) + fig.update_xaxes( + mirror=True, + ticks='outside', + showline=True, + linecolor='black', + tickformat='.1e', + dtick=0.25 + ) + fig.update_yaxes( + mirror=True, + ticks='outside', + showline=True, + linecolor='black', + tickformat = "000" + ) + return fig + +def plot_losses(INPUT_PATH, s=0.3): + # INPUT_PATH points to "confined_fraction.dat" + + with open(INPUT_PATH, "r") as f: ar = np.loadtxt(f) + + time = ar[:, 0] + confined_passing = ar[:, 1] + confined_trapped = ar[:, 2] + number = ar[:, 3] + confined_total = confined_passing + confined_trapped + loss_frac = 1 - confined_total + + fig = plt.figure() + plt.plot(time, loss_frac) + plt.title(f"Collisionless particle losses, s={s}, N={int(number[0])}") + plt.xscale('log') + plt.xlabel("time [s]") + plt.ylabel("loss fraction") + return fig, time, loss_frac + + +def success_plt(df, df_filtered): + fig = plt.figure(1, figsize=(14.5, 11)) + nrows = 5 + ncols = 5 + + def plot_2d_hist(field, log=False, subplot_index=0): + plt.subplot(nrows, ncols, subplot_index) + nbins = 20 + if log: + data = df[field] + bins = np.logspace(np.log10(data.min()), np.log10(data.max()), nbins) + else: + bins = nbins + n,bins,patchs = plt.hist(df[field], bins=bins, label="before filtering") + plt.hist(df_filtered[field], bins=bins, alpha=1, label="after filtering") + plt.xlabel(field) + plt.legend(loc=0, fontsize=6) + if log: + plt.xscale("log") + + + # 2nd entry of each tuple is True if the field should be plotted on a log x-scale. + fields = ( + ("R1", False), + ("order", False), + ("max_max_force", False), + ("max_length", False), + ("max_max_κ", False), + ("max_MSC", False), + ("coil_coil_distance", False), + ("coil_surface_distance", False), + ("length_target", False), + ("force_threshold", False), + ("max_κ_threshold", False), + ("msc_threshold", False), + ("cc_threshold", False), + ("cs_threshold", False), + ("length_weight", True), + ("max_κ_weight", True), + ("msc_weight", True), + ("cc_weight", True), + ("cs_weight", True), + ("force_weight", True), + # ("linking_number", False), # TODO: uncomment later + ('ncoils', False) + ) + + i=1 + for field, log in fields: + plot_2d_hist(field, log, i) + i += 1 + + plt.tight_layout() + return fig + + + +############################################################################### +# III) PHYSICS STUFF +############################################################################### + +def closest_approach(coils, φ_min=0, φ_max=2*π): + """Closest approach within a set of coils.""" + ncoils = len(coils) + phis = coils[0].curve.quadpoints * 2 * np.pi + + index_min = np.absolute(phis-φ_min).argmin() + if phis[index_min] < φ_min: index_min += 1 + index_max = np.absolute(phis-φ_max).argmin() + if phis[index_max] > φ_max: index_max -= 1 + + gammas = [coil.curve.gamma() for coil in coils] + gammas_truncated = [gamma[index_min:index_max+1] for gamma in gammas] + + return min([np.min(cdist(gammas_truncated[i], gammas_truncated[j])) for i in range(ncoils) for j in range(i)]) + +def poincare(UUID, OUT_DIR='./output/QA/with-force-penalty/1/poincare/', + INPUT_FILE="./inputs/input.LandremanPaul2021_QA", phis=[0.0], + nfieldlines=10, tmax_fl=20000, degree=4, R0_min=1.2125346, + R0_max=1.295, interpolate=True, debug=False): + """Compute Poincare plots.""" + + # Directory for output + if debug: os.makedirs(OUT_DIR, exist_ok=True) + + # Load in the boundary surface: + surf = SurfaceRZFourier.from_vmec_input(INPUT_FILE, nphi=200, ntheta=30, range="full torus") + nfp = surf.nfp + + # Load in the optimized coils from stage_two_optimization.py: + coils_filename = glob.glob(f"./**/{UUID}/biot_savart.json", recursive=True)[0] + bs = simsopt.load(coils_filename) + + sc_fieldline = SurfaceClassifier(surf, h=0.03, p=2) + if debug: + surf.to_vtk(OUT_DIR + 'surface') + sc_fieldline.to_vtk(OUT_DIR + 'levelset', h=0.02) + + def trace_fieldlines(bfield, label): + R0 = np.linspace(R0_min, R0_max, nfieldlines) + Z0 = np.zeros(nfieldlines) + fieldlines_tys, fieldlines_phi_hits = compute_fieldlines( + bfield, R0, Z0, tmax=tmax_fl, tol=1e-16, comm=comm_world, + phis=phis, stopping_criteria=[LevelsetStoppingCriterion(sc_fieldline.dist)]) + if debug: particles_to_vtk(fieldlines_tys, OUT_DIR + f'fieldlines_{label}') + + if len(phis) > 1: + raise NotImplementedError("Not yet implemented!") + # plot_poincare_data(fieldlines_phi_hits, phis, OUT_DIR + f'poincare_fieldline_{label}.png', dpi=150) + + fig = plt.figure() + for j in range(len(fieldlines_phi_hits)): + data_this_phi = fieldlines_phi_hits[j][np.where(fieldlines_phi_hits[j][:, 1] == 0)[0], :] + if data_this_phi.size == 0: + continue + r = np.sqrt(data_this_phi[:, 2]**2+data_this_phi[:, 3]**2) + plt.scatter(r, data_this_phi[:, 4], marker='o', s=1, linewidths=0, color='blue') + + return fig + + # Bounds for the interpolated magnetic field chosen so that the surface is + # entirely contained in it + n = 20 + rs = np.linalg.norm(surf.gamma()[:, :, 0:2], axis=2) + zs = surf.gamma()[:, :, 2] + rrange = (np.min(rs), np.max(rs), n) + phirange = (0, 2*np.pi/nfp, n*2) + zrange = (0, np.max(zs), n//2) + + def skip(rs, phis, zs): + rphiz = np.asarray([rs, phis, zs]).T.copy() + dists = sc_fieldline.evaluate_rphiz(rphiz) + skip = list((dists < -0.05).flatten()) + return skip + + bsh = InterpolatedField( + bs, degree, rrange, phirange, zrange, True, nfp=nfp, stellsym=True, skip=skip + ) + bsh.set_points(surf.gamma().reshape((-1, 3))) + bs.set_points(surf.gamma().reshape((-1, 3))) + image = trace_fieldlines(bsh, 'bsh') if(interpolate) else trace_fieldlines(bs, 'bs') + return image + +def poincare_with_target(UUID="6266c8d4bb25499b899d86e9e3dd2ee2", phi=0.1, + INPUT_FILE="./inputs/input.LandremanPaul2021_QA", R0_max=1.295, + nfieldlines=10, tmax_fl=20000): + """ + Checks that the target LCFS, QFM surface, and Poincaré plots all agree. + """ + + def get_crosssection(s): + thetas = np.linspace(0, 1, 200) + cs = s.cross_section(phi, thetas=thetas) + r = np.sqrt(cs[:, 0] ** 2 + cs[:, 1] ** 2) + r = np.concatenate((r, [r[0]])) + z = cs[:, 2] + z = np.concatenate((z, [z[0]])) + return r, z + + # show poincaré plots + poincare(UUID, INPUT_FILE=INPUT_FILE, tmax_fl=tmax_fl, + degree=4, phis=[phi], R0_max=R0_max, nfieldlines=nfieldlines) + + # show target LCFS + LCFS = SurfaceRZFourier.from_vmec_input(INPUT_FILE, range="full torus", nphi=32, ntheta=32) + r, z = get_crosssection(LCFS) + plt.plot(r, z, linewidth=1, c='r', label="target LCFS") + plt.xlabel('R [m]') + plt.ylabel('Z [m]') + + +def qfm(UUID, INPUT_FILE="./inputs/input.LandremanPaul2021_QA", vol_frac=1.00): + """Generated quadratic flux minimizing surfaces, adapted from + https://github.com/hiddenSymmetries/simsopt/blob/master/examples/1_Simple/qfm.py""" + + # We start with an initial guess that is just the stage I LCFS + s = SurfaceRZFourier.from_vmec_input(INPUT_FILE, range="full torus", nphi=32, ntheta=32) + + # Optimize at fixed volume: + path = glob.glob("/**/" + UUID + "/biot_savart.json", recursive=True)[0] + bs = load(path) + qfm = QfmResidual(s, bs) + qfm.J() + vol = Volume(s) + vol_target = vol_frac * vol.J() + qfm_surface = QfmSurface(bs, s, vol, vol_target) + + constraint_weight = 1e0 + res = qfm_surface.minimize_qfm_penalty_constraints_LBFGS(tol=1e-12, maxiter=1000, + constraint_weight=constraint_weight) + vol_err = abs((s.volume()-vol_target)/vol_target) + residual = np.linalg.norm(qfm.J()) + # print(f"||vol constraint||={0.5*(s.volume()-vol_target)**2:.8e}, ||residual||={np.linalg.norm(qfm.J()):.8e}") + return qfm_surface.surface, vol_err, residual + + +def run_SIMPLE(UUID, trace_time=1e-1, s=0.3, n_test_part=1024, vmec_name="eq_scaled.nc", + BUILD_DIR="/Users/sienahurwitz/Documents/Physics/Codes/SIMPLE/build/", + suppress_output=False): + + # STEP 1: generate the input files and save to the run directory + RUN_DIR = glob.glob(f"../**/{UUID}/", recursive=True)[0] + with open(RUN_DIR + "simple.in", "w") as f: + f.write(f"&config\n") + f.write(f"trace_time = {trace_time}d0\n") + f.write(f"sbeg = {s}d0\n") + f.write(f"ntestpart = {n_test_part}\n") + f.write(f"netcdffile = '{vmec_name}'\n") + f.write(f"/\n") + + # STEP 2: run SIMPLE + command = BUILD_DIR + "simple.x" + if suppress_output: + subprocess.run(command, shell=True, check=True, stdout=subprocess.DEVNULL) + else: + subprocess.run(command, shell=True, check=True) + + # STEP 3: delete irrelevant files + files = ["simple.in", "times_lost.dat", vmec_name] + for file in files: + os.remove(RUN_DIR + file) + + +def surf_to_desc(simsopt_surf, LMN=8): + """Returns a DESC equilibrium from a simsopt surface, adapted from code + written by Matt Landreman. Note that LMN is a DESC resolution parameter. """ + nfp = simsopt_surf.nfp + ndofs = len(simsopt_surf.x) + index = int((ndofs + 1) / 2) + xm = simsopt_surf.m[:index] + xn = simsopt_surf.n[:index] + rmnc = simsopt_surf.x[:index] + zmns = np.concatenate(([0], simsopt_surf.x[index:])) + + rmns = np.zeros_like(rmnc) + zmnc = np.zeros_like(zmns) + + # Adapted from desc's VMECIO.load around line 126: + m, n, Rb_lmn = ptolemy_identity_fwd(xm, xn, s=rmns, c=rmnc) + m, n, Zb_lmn = ptolemy_identity_fwd(xm, xn, s=zmns, c=zmnc) + surface = np.vstack((np.zeros_like(m), m, n, Rb_lmn, Zb_lmn)).T + desc_surface = FourierRZToroidalSurface( + surface[:, 3], + surface[:, 4], + surface[:, 1:3].astype(int), + surface[:, 1:3].astype(int), + nfp, + simsopt_surf.stellsym, + check_orientation=False, + ) + + # To avoid warning message, flip the orientation manually: + if desc_surface._compute_orientation() == -1: + desc_surface._flip_orientation() + assert desc_surface._compute_orientation() == 1 + + eq = Equilibrium( + surface=desc_surface, + L=LMN, + M=LMN, + N=LMN, + L_grid=2 * LMN, + M_grid=2 * LMN, + N_grid=2 * LMN, + sym=True, + NFP=nfp, + Psi=np.pi * simsopt_surf.minor_radius()**2, + ensure_nested=False, + ) + + # Check that the desc surface matches the input surface. + # Grid resolution for testing the surfaces match: + ntheta = len(simsopt_surf.quadpoints_theta) + nphi = len(simsopt_surf.quadpoints_phi) + simsopt_surf2 = SurfaceRZFourier( + mpol=simsopt_surf.mpol, + ntor=simsopt_surf.ntor, + nfp=simsopt_surf.nfp, + stellsym=simsopt_surf.stellsym, + dofs=simsopt_surf.dofs, + quadpoints_phi=np.linspace(0, 1 / simsopt_surf.nfp, nphi, endpoint=False), + quadpoints_theta=np.linspace(0, 1, ntheta, endpoint=False), + ) + gamma = simsopt_surf2.gamma() + + grid = LinearGrid( + rho=1, + theta=np.linspace(0, 2 * np.pi, ntheta, endpoint=False), + zeta=np.linspace(0, 2 * np.pi / nfp, nphi, endpoint=False), + NFP=nfp, + ) + data = eq.compute(["X", "Y", "Z"], grid=grid) + + def compare_simsopt_desc(simsopt_data, desc_data): + desc_arr = desc_data.reshape((ntheta, nphi), order="F") + desc_arr = np.vstack((desc_arr[:1, :], np.flipud(desc_arr[1:, :]))) + np.testing.assert_allclose(simsopt_data, desc_arr, atol=1e-14) + + compare_simsopt_desc(gamma[:, :, 0].T, data["X"]) + compare_simsopt_desc(gamma[:, :, 1].T, data["Y"]) + compare_simsopt_desc(gamma[:, :, 2].T, data["Z"]) + + # If tests are passed: + # eq.solve() + desc.continuation.solve_continuation_automatic(eq, verbose=0) + return eq + + +############################################################################### +# IV) SANITY CHECKS +############################################################################### + +def check_poincare_and_qfm(UUID="6266c8d4bb25499b899d86e9e3dd2ee2", phi=0.1, + INPUT_FILE="./inputs/input.LandremanPaul2021_QA"): + """ + Checks that the target LCFS, QFM surface, and Poincaré plots all agree. + """ + + def get_crosssection(s): + cs = s.cross_section(phi) + r = np.sqrt(cs[:, 0] ** 2 + cs[:, 1] ** 2) + r = np.concatenate((r, [r[0]])) + z = cs[:, 2] + z = np.concatenate((z, [z[0]])) + return r, z + + # show poincaré plots + poincare(UUID, INPUT_FILE=INPUT_FILE, nfieldlines=10, tmax_fl=20000, + degree=4, phis=[phi]) + + # show target LCFS + LCFS = SurfaceRZFourier.from_vmec_input(INPUT_FILE, range="full torus", nphi=32, ntheta=32) + r, z = get_crosssection(LCFS) + plt.plot(r, z, linewidth=2, c='r', label="target LCFS") + + # show QFMs at various radial distances + N=2 + for i in range(N): + vol_frac = 1.0 - 0.02 * i / (N - 1 + 1e-100) + qfm_surf, _, _ = qfm(UUID, INPUT_FILE=INPUT_FILE, vol_frac=vol_frac) + r, z = get_crosssection(qfm_surf) + if i != N - 1: + plt.plot(r, z, '--', linewidth=2, c='k') + else: + plt.plot(r, z, '--', linewidth=2, c='k', label="QFMs") + + plt.legend() + plt.show() \ No newline at end of file diff --git a/examples/3_Advanced/coil_force_pareto_scans.py b/examples/3_Advanced/coil_force_pareto_scans.py new file mode 100644 index 000000000..79eaa5a2d --- /dev/null +++ b/examples/3_Advanced/coil_force_pareto_scans.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python + +""" +Example script for the force metric in a stage-two coil optimization +""" +from analysis_tools import * +from optimization_tools import * +import imageio +import matplotlib +import matplotlib.pyplot as plt +import os + +# plt.ioff() +initial_optimizations(N=10, with_force=False, MAXITER=400) +df, df_filtered, df_pareto = get_dfs() +success_plt(df, df_filtered).show() + +df, df_filtered, df_pareto = get_dfs() + +y_axes = ["max_max_force", "mean_RMS_force"] +labels = ["max force [N/m]", "mean force [N/m]"] +y_lims = [(8500, 25000), (6000, 10000)] +for y_axis, label, y_lim in zip(y_axes, labels, y_lims): + + fig = plt.figure(figsize=(6.5, 6.5)) + # plt.rc("font", size=13) + plt.rc("font", size=17) + markersize = 5 + n_pareto = df_pareto.shape[0] + n_filtered = df_filtered.shape[0] - n_pareto + color="coil_surface_distance" + color_label="coil-surface distance [m]" + norm = plt.Normalize(min(df_filtered[color]), max(df_filtered[color])) + plt.scatter( + df_filtered["normalized_BdotN"], + df_filtered[y_axis], + c=df_filtered[color], + s=markersize, + label=f'all optimizations, N={n_filtered}', + norm=norm + ) + plt.scatter( + df_pareto["normalized_BdotN"], + df_pareto[y_axis], + c=df_pareto[color], + marker="+", + label=f'Pareto front, N={n_pareto}', + norm=norm, + ) + plt.xlabel(r'$\langle|\mathbf{B}\cdot\mathbf{n}|\rangle/\langle B \rangle$ [unitless]') + plt.ylabel(label) + plt.xlim(0.7 * min(df_filtered["normalized_BdotN"]), max(df_filtered["normalized_BdotN"])) + plt.ylim(y_lim) + plt.xscale("log") + plt.colorbar(label=color_label) + plt.clim(0.17, 0.31) + plt.legend(loc='upper right', fontsize='11') + # plt.title('Pareto Front') + plt.savefig(f"./output/QA/with-force-penalty/4/pareto_{y_axis}.pdf", bbox_inches='tight') + plt.show() \ No newline at end of file diff --git a/examples/3_Advanced/gen_pareto.py b/examples/3_Advanced/gen_pareto.py new file mode 100644 index 000000000..44ab27b9a --- /dev/null +++ b/examples/3_Advanced/gen_pareto.py @@ -0,0 +1,17 @@ +import sys +# sys.path.append("/global/homes/s/shurwitz/Force-Optimizations/" ) + +from analysis_tools import * + +iteration = int(sys.argv[1]) +INPUT_DIR = f"../output/QA/without-force-penalty/{iteration-1}/optimizations/" +OUTPUT_DIR = f"../output/QA/without-force-penalty/{iteration-1}/pareto/" +_, df_filtered, _ = get_dfs(INPUT_DIR=INPUT_DIR) +df_sorted = df_filtered.sort_values(by=["normalized_BdotN"]) + +if not os.path.exists(OUTPUT_DIR): + os.makedirs(OUTPUT_DIR) + for UUID in df_sorted[:100]['UUID']: + SOURCE_DIR = glob.glob(f"../**/{UUID}/", recursive=True)[0] + DEST_DIR = f"{OUTPUT_DIR}{UUID}/" + shutil.copytree(SOURCE_DIR, DEST_DIR) \ No newline at end of file diff --git a/examples/3_Advanced/inputs/LandremanPaul2021_QA.h5 b/examples/3_Advanced/inputs/LandremanPaul2021_QA.h5 new file mode 100644 index 000000000..339bfe8a3 Binary files /dev/null and b/examples/3_Advanced/inputs/LandremanPaul2021_QA.h5 differ diff --git a/examples/3_Advanced/inputs/input.LandremanPaul2021_QA b/examples/3_Advanced/inputs/input.LandremanPaul2021_QA new file mode 100644 index 000000000..266361a13 --- /dev/null +++ b/examples/3_Advanced/inputs/input.LandremanPaul2021_QA @@ -0,0 +1,104 @@ +&INDATA +!----- Runtime Parameters ----- + DELT = 9.00E-01 + NITER = 10000 + NSTEP = 200 + TCON0 = 2.00E+00 +! NS_ARRAY = 16 50 +! NITER_ARRAY = 600 3000 +! FTOL_ARRAY = 1.0E-16 1.0E-11 + + NS_ARRAY = 12 25 50 75 100 150 201 + NITER_ARRAY = 1200 2000 3000 4000 5000 6000 20000 + FTOL_ARRAY = 1.0E-17 1.0e-17 1.0e-17 1.0e-17 1.0e-17 1.0e-17 2.0e-17 + + PRECON_TYPE = 'none' + PREC2D_THRESHOLD = 1.000000E-19 +!----- Grid Parameters ----- + LASYM = F + NFP = 0002 + MPOL = 00016 + NTOR = 00012 + PHIEDGE = 0.08385727554 +!----- Free Boundary Parameters ----- + LFREEB = F + NVACSKIP = 6 +!----- Pressure Parameters ----- + GAMMA = 0.000000000000E+000 + BLOAT = 1.000000000000E+000 + SPRES_PED = 1.000000000000E+000 + PRES_SCALE = 1.000000000000E+000 + PMASS_TYPE = 'power_series' + AM = 000000000E+00 +!----- Current/Iota Parameters ----- + CURTOR = 0 + NCURR = 1 + PIOTA_TYPE = 'power_series' + PCURR_TYPE = 'power_series' +!----- Axis Parameters ----- + RAXIS_CC = 0 + ZAXIS_CS = 0 +!----- Boundary Parameters ----- +! n comes before m! +RBC( 0, 0) = 1.000000000000000e+00, ZBS( 0, 0) = -0.000000000000000e+00 +RBC( 1, 0) = 1.635736469509085e-01, ZBS( 1, 0) = -1.253658594586452e-01 +RBC( 2, 0) = 1.153464403265224e-02, ZBS( 2, 0) = -1.143100828466809e-02 +RBC( 3, 0) = 1.660825542716507e-04, ZBS( 3, 0) = -1.836705489992535e-04 +RBC( 4, 0) = -4.289839778893124e-05, ZBS( 4, 0) = 6.379980554947428e-05 +RBC( 5, 0) = -7.088365208037713e-06, ZBS( 5, 0) = 6.424749751809887e-06 +RBC( -5, 1) = -1.234892163278594e-06, ZBS( -5, 1) = 1.650698547431063e-06 +RBC( -4, 1) = 5.885906939720765e-05, ZBS( -4, 1) = 7.747703517376636e-05 +RBC( -3, 1) = 1.177705088726893e-03, ZBS( -3, 1) = 8.768242973159689e-04 +RBC( -2, 1) = 6.117449935928349e-03, ZBS( -2, 1) = 6.917463886154144e-03 +RBC( -1, 1) = 3.695998128952078e-02, ZBS( -1, 1) = 3.357703036173842e-02 +RBC( 0, 1) = 1.658776713857751e-01, ZBS( 0, 1) = 2.153938070783234e-01 +RBC( 1, 1) = -1.088789497913388e-01, ZBS( 1, 1) = 8.004202421496563e-02 +RBC( 2, 1) = -2.128857374166514e-02, ZBS( 2, 1) = 1.941805377930779e-02 +RBC( 3, 1) = -2.300157572693968e-03, ZBS( 3, 1) = 2.043310207615721e-03 +RBC( 4, 1) = -2.627971451624495e-05, ZBS( 4, 1) = 9.515817827817248e-05 +RBC( 5, 1) = -3.624748429446710e-06, ZBS( 5, 1) = 2.237608189340047e-06 +RBC( -5, 2) = 3.429971812740527e-07, ZBS( -5, 2) = 3.653065246955757e-06 +RBC( -4, 2) = 4.850684989433037e-05, ZBS( -4, 2) = 2.581277522578798e-05 +RBC( -3, 2) = -1.629464724791329e-05, ZBS( -3, 2) = 9.255771223110816e-05 +RBC( -2, 2) = 6.490349497379012e-04, ZBS( -2, 2) = 3.952687065376200e-04 +RBC( -1, 2) = 3.445618605450382e-05, ZBS( -1, 2) = -1.331440902105533e-03 +RBC( 0, 2) = 1.582563423506581e-02, ZBS( 0, 2) = 1.561378330710932e-02 +RBC( 1, 2) = 2.546788563643479e-02, ZBS( 1, 2) = 8.409380325679782e-05 +RBC( 2, 2) = 3.589988783905243e-03, ZBS( 2, 2) = -7.162905352448011e-03 +RBC( 3, 2) = 2.023292416997494e-03, ZBS( 3, 2) = -1.623525685575538e-03 +RBC( 4, 2) = 1.334125199187403e-04, ZBS( 4, 2) = -2.144823958959020e-04 +RBC( 5, 2) = 7.495835691917963e-06, ZBS( 5, 2) = -3.041995113869397e-06 +RBC( -5, 3) = 2.595971377554895e-06, ZBS( -5, 3) = 7.546398657171724e-07 +RBC( -4, 3) = -2.752586825793088e-05, ZBS( -4, 3) = -1.293908425979845e-05 +RBC( -3, 3) = -6.881215658169502e-05, ZBS( -3, 3) = -1.146639773380277e-04 +RBC( -2, 3) = 1.851061354443457e-04, ZBS( -2, 3) = 2.034757712275398e-04 +RBC( -1, 3) = -9.499158045035575e-04, ZBS( -1, 3) = -1.194121521590374e-03 +RBC( 0, 3) = 6.789213877319417e-04, ZBS( 0, 3) = 2.074912419441758e-03 +RBC( 1, 3) = 1.278256592588691e-04, ZBS( 1, 3) = 1.453779212226518e-03 +RBC( 2, 3) = 1.216640222421739e-04, ZBS( 2, 3) = 1.751939663468610e-03 +RBC( 3, 3) = -8.824623427313230e-04, ZBS( 3, 3) = 4.283901601890926e-04 +RBC( 4, 3) = -8.290455528140711e-05, ZBS( 4, 3) = 1.361192934036727e-04 +RBC( 5, 3) = -7.724473257867983e-06, ZBS( 5, 3) = 8.518581946796281e-06 +RBC( -5, 4) = -4.525265993447933e-07, ZBS( -5, 4) = -4.025242836970244e-07 +RBC( -4, 4) = -3.765525591663771e-06, ZBS( -4, 4) = -5.091039705513993e-06 +RBC( -3, 4) = 1.302918803102485e-05, ZBS( -3, 4) = 2.938909520741841e-06 +RBC( -2, 4) = -9.479666265169178e-05, ZBS( -2, 4) = -2.244078515714136e-05 +RBC( -1, 4) = 2.264194360369701e-04, ZBS( -1, 4) = -2.520850733246328e-05 +RBC( 0, 4) = -3.151153789937265e-04, ZBS( 0, 4) = 9.320140252262431e-05 +RBC( 1, 4) = 3.296873412220537e-04, ZBS( 1, 4) = -2.042044185495455e-05 +RBC( 2, 4) = 3.702513915656681e-04, ZBS( 2, 4) = 4.064784880640345e-05 +RBC( 3, 4) = 1.391783522713434e-04, ZBS( 3, 4) = -1.361278177512208e-05 +RBC( 4, 4) = 6.688738405012066e-06, ZBS( 4, 4) = -6.608364805280040e-05 +RBC( 5, 4) = 1.339256039281701e-05, ZBS( 5, 4) = -4.386750839649032e-06 +RBC( -5, 5) = -4.605122542160908e-07, ZBS( -5, 5) = -1.918844832075774e-07 +RBC( -4, 5) = 4.705334806798627e-06, ZBS( -4, 5) = 4.045952798129327e-06 +RBC( -3, 5) = -3.907227055013670e-06, ZBS( -3, 5) = 1.473135558324792e-06 +RBC( -2, 5) = 1.564971051544039e-06, ZBS( -2, 5) = -4.697232422054022e-06 +RBC( -1, 5) = 1.512725785055272e-05, ZBS( -1, 5) = 2.716682965888606e-05 +RBC( 0, 5) = 2.685611719517791e-05, ZBS( 0, 5) = 4.010578314605041e-05 +RBC( 1, 5) = -6.524317753178391e-05, ZBS( 1, 5) = 1.623461963012973e-05 +RBC( 2, 5) = -3.537817812851590e-05, ZBS( 2, 5) = -3.511150206275442e-05 +RBC( 3, 5) = 3.342938698620821e-05, ZBS( 3, 5) = 2.352419744940536e-05 +RBC( 4, 5) = -9.745754270559217e-06, ZBS( 4, 5) = 1.637637084327691e-06 +RBC( 5, 5) = -1.976517490730610e-06, ZBS( 5, 5) = 4.498757881413679e-08 +/ diff --git a/examples/3_Advanced/inputs/input.LandremanPaul2021_QH_magwell b/examples/3_Advanced/inputs/input.LandremanPaul2021_QH_magwell new file mode 100644 index 000000000..c18d6811b --- /dev/null +++ b/examples/3_Advanced/inputs/input.LandremanPaul2021_QH_magwell @@ -0,0 +1,145 @@ +!----- Runtime Parameters ----- +&INDATA + DELT = 9.000000000000E-001 + NITER = 10000 + NSTEP = 200 + TCON0 = 2.000000000000E+000 + NS_ARRAY = 12 25 50 75 + 100 150 201 + FTOL_ARRAY = 1.000000E-17 1.000000E-17 1.000000E-17 1.000000E-17 + 1.000000E-17 1.000000E-17 2.000000E-17 + NITER_ARRAY = 1200 2000 3000 4000 + 5000 6000 20000 + PRECON_TYPE = 'none' + PREC2D_THRESHOLD = 1.000000E-19 +!----- Grid Parameters ----- + LASYM = F + NFP = 0004 + MPOL = 0008 + NTOR = 0008 + PHIEDGE = 4.551294381756E+001 + NTHETA = 0000 + NZETA = 0000 +!----- Free Boundary Parameters ----- + LFREEB = F +!----- Pressure Parameters ----- + GAMMA = 0.000000000000E+000 + BLOAT = 1.000000000000E+000 + SPRES_PED = 1.000000000000E+000 + PRES_SCALE = 3.439366905306E+001 + PMASS_TYPE = 'power_series' + AM = 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 +!----- ANI/FLOW Parameters ----- + BCRIT = 1.000000000000E+000 + PT_TYPE = 'power_series' + AT = 1.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 + PH_TYPE = 'power_series' + AH = 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 +!----- Current/Iota Parameters ----- + CURTOR = 0.000000000000E+000 + NCURR = 1 + PIOTA_TYPE = 'power_series' + AI = 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 + PCURR_TYPE = 'power_series' + AC = 0.000000000000E+000 0.000000000000E+000 0.000000000000E+000 0.000000000000E+000 + 0.000000000000E+000 0.000000000000E+000 0.000000000000E+000 0.000000000000E+000 + 0.000000000000E+000 0.000000000000E+000 0.000000000000E+000 0.000000000000E+000 + 0.000000000000E+000 0.000000000000E+000 0.000000000000E+000 0.000000000000E+000 + 0.000000000000E+000 0.000000000000E+000 0.000000000000E+000 0.000000000000E+000 + 0.000000000000E+000 +!----- Axis Parameters ----- + RAXIS_CC = 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 0.00000000000000E+00 + 0.00000000000000E+00 + ZAXIS_CS = 0.000000000000E+000 0.000000000000E+000 0.000000000000E+000 0.000000000000E+000 + 0.000000000000E+000 0.000000000000E+000 0.000000000000E+000 0.000000000000E+000 + 0.000000000000E+000 +!----- Boundary Parameters ----- + RBC( 000,000) = 1.378480264360E+001 ZBS( 000,000) = 0.000000000000E+000 + RBC( 001,000) = 1.726187404079E+000 ZBS( 001,000) = 1.100163466578E+000 + RBC( 002,000) = 6.784619411470E-002 ZBS( 002,000) = 7.875463181795E-002 + RBC( 003,000) = -8.610603893077E-003 ZBS( 003,000) = -4.473489744658E-003 + RBC( 004,000) = -6.954768654741E-005 ZBS( 004,000) = 1.588395888575E-004 + RBC( 005,000) = 9.681155117252E-006 ZBS( 005,000) = 4.138060458560E-005 + RBC(-005,001) = -7.090062438846E-005 ZBS(-005,001) = -5.640105757046E-005 + RBC(-004,001) = 2.462338732367E-003 ZBS(-004,001) = -2.148554902649E-003 + RBC(-003,001) = -1.896117167297E-002 ZBS(-003,001) = 1.594981036419E-002 + RBC(-002,001) = -2.137920688773E-001 ZBS(-002,001) = 2.506646350870E-001 + RBC(-001,001) = -1.362873009923E+000 ZBS(-001,001) = 8.962984591705E-001 + RBC( 000,001) = 1.733237139203E+000 ZBS( 000,001) = 2.081327204381E+000 + RBC( 001,001) = 4.417119515646E-001 ZBS( 001,001) = 4.553330033140E-001 + RBC( 002,001) = 6.695419953370E-002 ZBS( 002,001) = 7.309171552376E-002 + RBC( 003,001) = 4.586829621480E-003 ZBS( 003,001) = 4.687286129107E-003 + RBC( 004,001) = 4.667077346695E-004 ZBS( 004,001) = -1.027361195085E-004 + RBC( 005,001) = 1.175176479306E-005 ZBS( 005,001) = -2.552650176539E-005 + RBC(-005,002) = 1.575377035260E-005 ZBS(-005,002) = 9.901363685120E-005 + RBC(-004,002) = 2.261424757931E-003 ZBS(-004,002) = -3.956732991220E-003 + RBC(-003,002) = 4.309259932575E-002 ZBS(-003,002) = -2.805088743419E-002 + RBC(-002,002) = 1.315953940480E-002 ZBS(-002,002) = -1.704646792977E-001 + RBC(-001,002) = 4.304332319616E-001 ZBS(-001,002) = 3.215570994146E-001 + RBC( 000,002) = 2.868901853821E-001 ZBS( 000,002) = 1.667695990457E-001 + RBC( 001,002) = 1.931383998087E-002 ZBS( 001,002) = -1.404506982196E-003 + RBC( 002,002) = 1.565755517097E-003 ZBS( 002,002) = 3.373921721361E-003 + RBC( 003,002) = 1.111978236837E-003 ZBS( 003,002) = 8.029270102691E-004 + RBC( 004,002) = 3.212364014056E-004 ZBS( 004,002) = 3.123283352231E-004 + RBC( 005,002) = 3.930745657250E-005 ZBS( 005,002) = 9.979673774622E-005 + RBC(-005,003) = -2.253965955942E-004 ZBS(-005,003) = 2.851378679531E-004 + RBC(-004,003) = -8.964050388804E-004 ZBS(-004,003) = 3.120148794567E-003 + RBC(-003,003) = -2.917431937647E-002 ZBS(-003,003) = 7.844357333328E-003 + RBC(-002,003) = -1.581028170367E-003 ZBS(-002,003) = 5.005227155127E-002 + RBC(-001,003) = 2.494149802424E-002 ZBS(-001,003) = 2.636750589380E-002 + RBC( 000,003) = -3.851998943789E-003 ZBS( 000,003) = 5.133128097390E-003 + RBC( 001,003) = 1.563504795759E-003 ZBS( 001,003) = 9.062801189268E-003 + RBC( 002,003) = 6.042310892453E-003 ZBS( 002,003) = 8.373946146862E-003 + RBC( 003,003) = 1.935609824006E-003 ZBS( 003,003) = 1.991723637988E-003 + RBC( 004,003) = 2.583033337260E-004 ZBS( 004,003) = 2.930369647971E-004 + RBC( 005,003) = 7.771615403367E-005 ZBS( 005,003) = 1.801603873809E-005 + RBC(-005,004) = 1.989970159442E-004 ZBS(-005,004) = -7.617610661490E-005 + RBC(-004,004) = 3.133474921973E-003 ZBS(-004,004) = -3.926686806018E-003 + RBC(-003,004) = 1.688513373469E-003 ZBS(-003,004) = 1.588962033956E-003 + RBC(-002,004) = 9.542511784421E-003 ZBS(-002,004) = 5.383886804322E-003 + RBC(-001,004) = 1.581562725838E-002 ZBS(-001,004) = 4.929187684486E-004 + RBC( 000,004) = 3.741454918911E-003 ZBS( 000,004) = 1.375986819706E-003 + RBC( 001,004) = 3.967477935721E-003 ZBS( 001,004) = 3.679633131334E-003 + RBC( 002,004) = 2.956934985275E-003 ZBS( 002,004) = 2.135542216902E-003 + RBC( 003,004) = 4.356320120811E-004 ZBS( 003,004) = -5.186710313568E-005 + RBC( 004,004) = -9.630558365277E-005 ZBS( 004,004) = -1.889282254155E-004 + RBC( 005,004) = -6.633879325437E-005 ZBS( 005,004) = -5.671974511566E-005 + RBC(-005,005) = -4.133502120774E-005 ZBS(-005,005) = -7.005722734239E-005 + RBC(-004,005) = -2.658222009382E-004 ZBS(-004,005) = 2.116648875803E-004 + RBC(-003,005) = -1.756431084094E-005 ZBS(-003,005) = 1.776829639892E-004 + RBC(-002,005) = 3.799362725283E-004 ZBS(-002,005) = -1.795387962455E-004 + RBC(-001,005) = -3.390501030100E-005 ZBS(-001,005) = -2.995903236632E-004 + RBC( 000,005) = -3.278942957178E-004 ZBS( 000,005) = -2.507480133573E-004 + RBC( 001,005) = 1.127886913857E-004 ZBS( 001,005) = -2.617033500753E-004 + RBC( 002,005) = -9.475120829442E-005 ZBS( 002,005) = -2.234820548655E-004 + RBC( 003,005) = -2.122004096876E-004 ZBS( 003,005) = -9.817830003961E-005 + RBC( 004,005) = 2.786329922830E-005 ZBS( 004,005) = 7.264324435394E-005 + RBC( 005,005) = 5.820410564626E-006 ZBS( 005,005) = 1.186994399716E-005 +/ +! RMIN = 9.416741239278133 +! RMAX = 17.062045101562568 +! ZMIN = -3.974708748860022 +! ZMAX = 3.974708748860023 +! PHIMIN = 0.0 +! PHIMAX = 1.5707963267948966 diff --git a/examples/3_Advanced/inputs/input.LandremanPaul2021_QH_magwell_R0=1 b/examples/3_Advanced/inputs/input.LandremanPaul2021_QH_magwell_R0=1 new file mode 100644 index 000000000..930923a81 --- /dev/null +++ b/examples/3_Advanced/inputs/input.LandremanPaul2021_QH_magwell_R0=1 @@ -0,0 +1,161 @@ +! +&INDATA +!---- Free-Boundary Parameters ---- + LFREEB = F + MGRID_FILE = 'none' + NVACSKIP = 5 +!---- Runtime Parameters ---- + NSTEP = 250 + DELT = 0.9 + NS_ARRAY = 17 33 65 129 257 + NITER_ARRAY = 10000 20000 30000 40000 50000 + FTOL_ARRAY = 1E-12 1E-12 1E-12 1E-12 1E-12 +!---- Grid Parameters ---- + LASYM = F + NFP = 4 + MPOL = 8 + NTOR = 8 + PHIEDGE = +2.44787652E-01 +!---- Pressure Parameters ---- + GAMMA = 0 + PRES_SCALE = 1 + AM = +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 + PMASS_TYPE = 'power_series' +!---- Current/Iota Parameters ---- + NCURR = 1 + CURTOR = +0.00000000E+00 + AC = +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 +0.00000000E+00 + PCURR_TYPE = 'power_series_I' +!---- Axis Parameters ---- + RAXIS_CC = +1.01315885E+00 +1.64722857E-01 +2.17525471E-02 +2.33532827E-03 +1.10553158E-04 -2.83866804E-05 -7.80383825E-06 +4.75116130E-07 -1.68554045E-07 + ZAXIS_CS = +0.00000000E+00 +1.26119027E-01 +1.90222172E-02 +2.31597239E-03 +1.51575318E-04 -2.00571500E-05 -7.26595895E-06 -7.43269026E-07 -4.15650694E-07 +!---- Boundary Parameters ---- + RBC( 0, 0) = +1.01094573E+00 ZBS( 0, 0) = +0.00000000E+00 + RBC( 1, 0) = +1.26594615E-01 ZBS( 1, 0) = +8.06834594E-02 + RBC( 2, 0) = +4.97568388E-03 ZBS( 2, 0) = +5.77568362E-03 + RBC( 3, 0) = -6.31481892E-04 ZBS( 3, 0) = -3.28075452E-04 + RBC( 4, 0) = -5.10046742E-06 ZBS( 4, 0) = +1.16489302E-05 + RBC( 5, 0) = +7.09993657E-07 ZBS( 5, 0) = +3.03475840E-06 + RBC( 6, 0) = -1.31501865E-19 ZBS( 6, 0) = -2.31832918E-18 + RBC( 7, 0) = +5.25160427E-20 ZBS( 7, 0) = +1.15725876E-19 + RBC( 8, 0) = +7.72917564E-21 ZBS( 8, 0) = -5.00808230E-20 + RBC( -8, 1) = +1.86022200E-17 ZBS( -8, 1) = -1.94724934E-19 + RBC( -7, 1) = -1.27369403E-17 ZBS( -7, 1) = +2.42780819E-19 + RBC( -6, 1) = -4.76032516E-19 ZBS( -6, 1) = -5.16690098E-19 + RBC( -5, 1) = +8.61847411E-07 ZBS( -5, 1) = +1.87205495E-06 + RBC( -4, 1) = +3.42272722E-05 ZBS( -4, 1) = +7.53443083E-06 + RBC( -3, 1) = +3.36387537E-04 ZBS( -3, 1) = -3.43754789E-04 + RBC( -2, 1) = +4.91026705E-03 ZBS( -2, 1) = -5.36037836E-03 + RBC( -1, 1) = +3.23941389E-02 ZBS( -1, 1) = -3.33930755E-02 + RBC( 0, 1) = +1.27111627E-01 ZBS( 0, 1) = -1.52639751E-01 + RBC( 1, 1) = -9.99499729E-02 ZBS( 1, 1) = -6.57324681E-02 + RBC( 2, 1) = -1.56790187E-02 ZBS( 2, 1) = -1.83831680E-02 + RBC( 3, 1) = -1.39056874E-03 ZBS( 3, 1) = -1.16972242E-03 + RBC( 4, 1) = +1.80582261E-04 ZBS( 4, 1) = +1.57570076E-04 + RBC( 5, 1) = -5.19968878E-06 ZBS( 5, 1) = +4.13632388E-06 + RBC( 6, 1) = +5.47183284E-19 ZBS( 6, 1) = +6.48827238E-19 + RBC( 7, 1) = +1.31994203E-17 ZBS( 7, 1) = -2.66285983E-19 + RBC( 8, 1) = -1.61532359E-17 ZBS( 8, 1) = +1.73019714E-19 + RBC( -8, 2) = +2.79567587E-17 ZBS( -8, 2) = +1.00091735E-19 + RBC( -7, 2) = +6.04662972E-19 ZBS( -7, 2) = +4.17766047E-19 + RBC( -6, 2) = -1.21139971E-18 ZBS( -6, 2) = -3.10561018E-19 + RBC( -5, 2) = +2.88271849E-06 ZBS( -5, 2) = -7.31886331E-06 + RBC( -4, 2) = +2.35587390E-05 ZBS( -4, 2) = -2.29054420E-05 + RBC( -3, 2) = +8.15499282E-05 ZBS( -3, 2) = -5.88848201E-05 + RBC( -2, 2) = +1.14828908E-04 ZBS( -2, 2) = -2.47435645E-04 + RBC( -1, 2) = +1.41643262E-03 ZBS( -1, 2) = +1.03003314E-04 + RBC( 0, 2) = +2.10398665E-02 ZBS( 0, 2) = -1.22304990E-02 + RBC( 1, 2) = +3.15669836E-02 ZBS( 1, 2) = -2.35822583E-02 + RBC( 2, 2) = +9.65090363E-04 ZBS( 2, 2) = +1.25014877E-02 + RBC( 3, 2) = +3.16031216E-03 ZBS( 3, 2) = +2.05718759E-03 + RBC( 4, 2) = +1.65847692E-04 ZBS( 4, 2) = +2.90177700E-04 + RBC( 5, 2) = +1.15534530E-06 ZBS( 5, 2) = -7.26143249E-06 + RBC( 6, 2) = +7.56343264E-19 ZBS( 6, 2) = +1.97812741E-19 + RBC( 7, 2) = -1.03726352E-19 ZBS( 7, 2) = +1.24900212E-20 + RBC( 8, 2) = -2.76379185E-17 ZBS( 8, 2) = +5.31242075E-20 + RBC( -8, 3) = +5.83446882E-18 ZBS( -8, 3) = +1.91217688E-19 + RBC( -7, 3) = -1.18601553E-17 ZBS( -7, 3) = -1.45689667E-19 + RBC( -6, 3) = +2.07048734E-17 ZBS( -6, 3) = -2.06676039E-19 + RBC( -5, 3) = +5.69952405E-06 ZBS( -5, 3) = -1.32125488E-06 + RBC( -4, 3) = +1.89433725E-05 ZBS( -4, 3) = -2.14906573E-05 + RBC( -3, 3) = +1.41953172E-04 ZBS( -3, 3) = -1.46068432E-04 + RBC( -2, 3) = +4.43129188E-04 ZBS( -2, 3) = -6.14125957E-04 + RBC( -1, 3) = +1.14663847E-04 ZBS( -1, 3) = -6.64645003E-04 + RBC( 0, 3) = -2.82496746E-04 ZBS( 0, 3) = -3.76451813E-04 + RBC( 1, 3) = +1.82915212E-03 ZBS( 1, 3) = -1.93373226E-03 + RBC( 2, 3) = -1.15948971E-04 ZBS( 2, 3) = -3.67071851E-03 + RBC( 3, 3) = -2.13957751E-03 ZBS( 3, 3) = -5.75287131E-04 + RBC( 4, 3) = -6.57402844E-05 ZBS( 4, 3) = -2.28824539E-04 + RBC( 5, 3) = -1.65300681E-05 ZBS( 5, 3) = -2.09113556E-05 + RBC( 6, 3) = -2.00204707E-17 ZBS( 6, 3) = -2.27004830E-19 + RBC( 7, 3) = +1.02694275E-17 ZBS( 7, 3) = +3.79470760E-19 + RBC( 8, 3) = -6.37847573E-18 ZBS( 8, 3) = -1.97993951E-19 + RBC( -8, 4) = -2.03835832E-18 ZBS( -8, 4) = +4.91464398E-19 + RBC( -7, 4) = -1.24937360E-17 ZBS( -7, 4) = +8.76679100E-20 + RBC( -6, 4) = -1.82959117E-19 ZBS( -6, 4) = -1.74488787E-19 + RBC( -5, 4) = -4.86513455E-06 ZBS( -5, 4) = +4.15969571E-06 + RBC( -4, 4) = -7.06283018E-06 ZBS( -4, 4) = +1.38555616E-05 + RBC( -3, 4) = +3.19482504E-05 ZBS( -3, 4) = +3.80381423E-06 + RBC( -2, 4) = +2.16854814E-04 ZBS( -2, 4) = -1.56615756E-04 + RBC( -1, 4) = +2.90965711E-04 ZBS( -1, 4) = -2.69855833E-04 + RBC( 0, 4) = +2.74389701E-04 ZBS( 0, 4) = -1.00911710E-04 + RBC( 1, 4) = +1.15988174E-03 ZBS( 1, 4) = -3.61495293E-05 + RBC( 2, 4) = +6.99825875E-04 ZBS( 2, 4) = -3.94841880E-04 + RBC( 3, 4) = +1.23831689E-04 ZBS( 3, 4) = -1.16530822E-04 + RBC( 4, 4) = +2.29801846E-04 ZBS( 4, 4) = +2.87974182E-04 + RBC( 5, 4) = +1.45939836E-05 ZBS( 5, 4) = +5.58658052E-06 + RBC( 6, 4) = -8.80914265E-20 ZBS( 6, 4) = +1.74488787E-19 + RBC( 7, 4) = +1.19330002E-17 ZBS( 7, 4) = -8.34327453E-20 + RBC( 8, 4) = +3.05050975E-18 ZBS( 8, 4) = -7.80064405E-20 + RBC( -8, 5) = +4.76456033E-20 ZBS( -8, 5) = -5.17325373E-19 + RBC( -7, 5) = -7.48353609E-19 ZBS( -7, 5) = +2.88838235E-19 + RBC( -6, 5) = +3.16620916E-18 ZBS( -6, 5) = -3.39448454E-19 + RBC( -5, 5) = +4.26855531E-07 ZBS( -5, 5) = -8.70514401E-07 + RBC( -4, 5) = +2.04343028E-06 ZBS( -4, 5) = -5.32748853E-06 + RBC( -3, 5) = -1.55622902E-05 ZBS( -3, 5) = +7.20017079E-06 + RBC( -2, 5) = -6.94883578E-06 ZBS( -2, 5) = +1.63896601E-05 + RBC( -1, 5) = +8.27166331E-06 ZBS( -1, 5) = +1.91927220E-05 + RBC( 0, 5) = -2.40470138E-05 ZBS( 0, 5) = +1.83892827E-05 + RBC( 1, 5) = -2.48651551E-06 ZBS( 1, 5) = +2.19712655E-05 + RBC( 2, 5) = +2.78636527E-05 ZBS( 2, 5) = +1.31669625E-05 + RBC( 3, 5) = -1.28812618E-06 ZBS( 3, 5) = -1.30308600E-05 + RBC( 4, 5) = -1.94947890E-05 ZBS( 4, 5) = -1.55230162E-05 + RBC( 5, 5) = -3.03141542E-06 ZBS( 5, 5) = +5.13783597E-06 + RBC( 6, 5) = -2.94089839E-18 ZBS( 6, 5) = -2.96673290E-19 + RBC( 7, 5) = -6.21298667E-19 ZBS( 7, 5) = -2.80367906E-19 + RBC( 8, 5) = +3.28225267E-20 ZBS( 8, 5) = +6.44380315E-19 + RBC( -8, 6) = +0.00000000E+00 ZBS( -8, 6) = +0.00000000E+00 + RBC( -7, 6) = +0.00000000E+00 ZBS( -7, 6) = +0.00000000E+00 + RBC( -6, 6) = +0.00000000E+00 ZBS( -6, 6) = +0.00000000E+00 + RBC( -5, 6) = +0.00000000E+00 ZBS( -5, 6) = +0.00000000E+00 + RBC( -4, 6) = +0.00000000E+00 ZBS( -4, 6) = +0.00000000E+00 + RBC( -3, 6) = +0.00000000E+00 ZBS( -3, 6) = +0.00000000E+00 + RBC( -2, 6) = +0.00000000E+00 ZBS( -2, 6) = +0.00000000E+00 + RBC( -1, 6) = +0.00000000E+00 ZBS( -1, 6) = +0.00000000E+00 + RBC( 0, 6) = +0.00000000E+00 ZBS( 0, 6) = +0.00000000E+00 + RBC( 1, 6) = +0.00000000E+00 ZBS( 1, 6) = +0.00000000E+00 + RBC( 2, 6) = +0.00000000E+00 ZBS( 2, 6) = +0.00000000E+00 + RBC( 3, 6) = +0.00000000E+00 ZBS( 3, 6) = +0.00000000E+00 + RBC( 4, 6) = +0.00000000E+00 ZBS( 4, 6) = +0.00000000E+00 + RBC( 5, 6) = +0.00000000E+00 ZBS( 5, 6) = +0.00000000E+00 + RBC( 6, 6) = +0.00000000E+00 ZBS( 6, 6) = +0.00000000E+00 + RBC( 7, 6) = +0.00000000E+00 ZBS( 7, 6) = +0.00000000E+00 + RBC( 8, 6) = +0.00000000E+00 ZBS( 8, 6) = +0.00000000E+00 + RBC( -8, 7) = +0.00000000E+00 ZBS( -8, 7) = +0.00000000E+00 + RBC( -7, 7) = +0.00000000E+00 ZBS( -7, 7) = +0.00000000E+00 + RBC( -6, 7) = +0.00000000E+00 ZBS( -6, 7) = +0.00000000E+00 + RBC( -5, 7) = +0.00000000E+00 ZBS( -5, 7) = +0.00000000E+00 + RBC( -4, 7) = +0.00000000E+00 ZBS( -4, 7) = +0.00000000E+00 + RBC( -3, 7) = +0.00000000E+00 ZBS( -3, 7) = +0.00000000E+00 + RBC( -2, 7) = +0.00000000E+00 ZBS( -2, 7) = +0.00000000E+00 + RBC( -1, 7) = +0.00000000E+00 ZBS( -1, 7) = +0.00000000E+00 + RBC( 0, 7) = +0.00000000E+00 ZBS( 0, 7) = +0.00000000E+00 + RBC( 1, 7) = +0.00000000E+00 ZBS( 1, 7) = +0.00000000E+00 + RBC( 2, 7) = +0.00000000E+00 ZBS( 2, 7) = +0.00000000E+00 + RBC( 3, 7) = +0.00000000E+00 ZBS( 3, 7) = +0.00000000E+00 + RBC( 4, 7) = +0.00000000E+00 ZBS( 4, 7) = +0.00000000E+00 + RBC( 5, 7) = +0.00000000E+00 ZBS( 5, 7) = +0.00000000E+00 + RBC( 6, 7) = +0.00000000E+00 ZBS( 6, 7) = +0.00000000E+00 + RBC( 7, 7) = +0.00000000E+00 ZBS( 7, 7) = +0.00000000E+00 + RBC( 8, 7) = +0.00000000E+00 ZBS( 8, 7) = +0.00000000E+00 +/ \ No newline at end of file diff --git a/examples/3_Advanced/optimization_tools.py b/examples/3_Advanced/optimization_tools.py new file mode 100644 index 000000000..8a8c7a872 --- /dev/null +++ b/examples/3_Advanced/optimization_tools.py @@ -0,0 +1,443 @@ +import glob +import json +import numpy as np +import os +import pandas as pd +import time +import uuid +from scipy.optimize import minimize +from simsopt._core.optimizable import load +from simsopt.field import Current, coils_via_symmetries, BiotSavart +from simsopt.geo import ( + CurveLength, + CurveCurveDistance, + CurveSurfaceDistance, + MeanSquaredCurvature, + LpCurveCurvature, + ArclengthVariation, + curves_to_vtk, + create_equally_spaced_curves, + SurfaceRZFourier, + LinkingNumber) +from simsopt.objectives import SquaredFlux, QuadraticPenalty +from simsopt.field.force import coil_force, LpCurveForce +from simsopt.field.selffield import regularization_circ + + +def continuation(N=10000, dx=0.05, + INPUT_DIR="./output/QA/with-force-penalty/1/pareto/", + OUTPUT_DIR="./output/QA/with-force-penalty/2/optimizations/", + INPUT_FILE="./inputs/input.LandremanPaul2021_QA", + MAXITER=14000): + """Performs a continuation method on a set of previous optimizations.""" + # Read in input optimizations + results = glob.glob(f"{INPUT_DIR}*/results.json") + df = pd.DataFrame() + for results_file in results: + with open(results_file, "r") as f: + data = json.load(f) + # Wrap lists in another list + for key, value in data.items(): + if isinstance(value, list): + data[key] = [value] + df = pd.concat([df, pd.DataFrame(data)], ignore_index=True) + + def perturb(init, parameter): + "Perturbs a parameter from an initial optimization." + return rand(1-dx, 1+dx) * init[parameter].iloc[0] + + for i in range(N): + init = df.sample() + + # FIXED PARAMETERS + ARCLENGTH_WEIGHT = 0.01 + UUID_init_from = init['UUID'].iloc[0] + ncoils = init['ncoils'].iloc[0] + order = init['order'].iloc[0] + R1 = init['R1'].iloc[0] + + # RANDOM PARAMETERS + CURVATURE_THRESHOLD = perturb(init, 'max_κ_threshold') + MSC_THRESHOLD = perturb(init, 'msc_threshold') + CS_THRESHOLD = perturb(init, 'cs_threshold') + CC_THRESHOLD = perturb(init, 'cc_threshold') + FORCE_THRESHOLD = perturb(init, 'force_threshold') + LENGTH_TARGET = perturb(init, 'length_target') + + LENGTH_WEIGHT = perturb(init, 'length_weight') + CURVATURE_WEIGHT = perturb(init, 'max_κ_weight') + MSC_WEIGHT = perturb(init, 'msc_weight') + CS_WEIGHT = perturb(init, 'cs_weight') + CC_WEIGHT = perturb(init, 'cc_weight') + FORCE_WEIGHT = perturb(init, 'force_weight') + + # RUNNING THE JOBS + res, results, coils = optimization( + OUTPUT_DIR, + INPUT_FILE, + R1, + order, + ncoils, + UUID_init_from, + LENGTH_TARGET, + LENGTH_WEIGHT, + CURVATURE_THRESHOLD, + CURVATURE_WEIGHT, + MSC_THRESHOLD, + MSC_WEIGHT, + CC_THRESHOLD, + CC_WEIGHT, + CS_THRESHOLD, + CS_WEIGHT, + FORCE_THRESHOLD, + FORCE_WEIGHT, + ARCLENGTH_WEIGHT, + dx=dx, + MAXITER=MAXITER) + + print(f"Job {i+1} completed with UUID={results['UUID']}") + + +def initial_optimizations(N=10000, with_force=True, MAXITER=14000, + OUTPUT_DIR="./output/QA/with-force-penalty/1/optimizations/", + INPUT_FILE="./inputs/input.LandremanPaul2021_QA", + ncoils=5): + + """Performs a set of initial optimizations by scanning over parameters.""" + for i in range(N): + # FIXED PARAMETERS + ARCLENGTH_WEIGHT = 0.01 + UUID_init_from = None # not starting from prev. optimization + order = 16 + + # RANDOM PARAMETERS + R1 = rand(0.35, 0.75) + CURVATURE_THRESHOLD = rand(5, 12) + MSC_THRESHOLD = rand(4,6) + CS_THRESHOLD = rand(0.166, 0.300) + CC_THRESHOLD = rand(0.083, 0.120) + FORCE_THRESHOLD = rand(0, 5e+04) + LENGTH_TARGET = rand(4.9,5.0) + + LENGTH_WEIGHT = 10.0 ** rand(-4, -2) + CURVATURE_WEIGHT = 10.0 ** rand(-9, -5) + MSC_WEIGHT = 10.0 ** rand(-7, -3) + CS_WEIGHT = 10.0 ** rand(-1, 4) + CC_WEIGHT = 10.0 ** rand(2, 5) + + if with_force: + FORCE_WEIGHT = 10.0 ** rand(-14, -9) + else: + FORCE_WEIGHT = 0 + + # RUNNING THE JOBS + res, results, coils = optimization( + OUTPUT_DIR, + INPUT_FILE, + R1, + order, + ncoils, + UUID_init_from, + LENGTH_TARGET, + LENGTH_WEIGHT, + CURVATURE_THRESHOLD, + CURVATURE_WEIGHT, + MSC_THRESHOLD, + MSC_WEIGHT, + CC_THRESHOLD, + CC_WEIGHT, + CS_THRESHOLD, + CS_WEIGHT, + FORCE_THRESHOLD, + FORCE_WEIGHT, + ARCLENGTH_WEIGHT, + MAXITER=MAXITER) + + print(f"Job {i+1} completed with UUID={results['UUID']}") + + +def initial_optimizations_QH(N=10000, with_force=True, MAXITER=14000, + OUTPUT_DIR="./output/QA/with-force-penalty/1/optimizations/", + INPUT_FILE="./inputs/input.LandremanPaul2021_QH_magwell_R0=1", + ncoils=3): + + """Performs a set of initial optimizations by scanning over parameters.""" + for i in range(N): + # FIXED PARAMETERS + ARCLENGTH_WEIGHT = 0.01 + UUID_init_from = None # not starting from prev. optimization + order = 16 + + # RANDOM PARAMETERS + R1 = rand(0.35, 0.75) + CURVATURE_THRESHOLD = rand(5, 12) + MSC_THRESHOLD = rand(4,6) + CS_THRESHOLD = rand(0.166, 0.300) + CC_THRESHOLD = rand(0.083, 0.120) + FORCE_THRESHOLD = rand(0, 5e+04) + LENGTH_TARGET = rand(4.9,5.0) + + LENGTH_WEIGHT = 10.0 ** rand(-3, -1) + CURVATURE_WEIGHT = 10.0 ** rand(-9, -5) + MSC_WEIGHT = 10.0 ** rand(-5, -1) + CS_WEIGHT = 10.0 ** rand(-1, 4) + CC_WEIGHT = 10.0 ** rand(2, 5) + + if with_force: + FORCE_WEIGHT = 10.0 ** rand(-14, -9) + else: + FORCE_WEIGHT = 0 + + # RUNNING THE JOBS + res, results, coils = optimization( + OUTPUT_DIR, + INPUT_FILE, + R1, + order, + ncoils, + UUID_init_from, + LENGTH_TARGET, + LENGTH_WEIGHT, + CURVATURE_THRESHOLD, + CURVATURE_WEIGHT, + MSC_THRESHOLD, + MSC_WEIGHT, + CC_THRESHOLD, + CC_WEIGHT, + CS_THRESHOLD, + CS_WEIGHT, + FORCE_THRESHOLD, + FORCE_WEIGHT, + ARCLENGTH_WEIGHT, + MAXITER=MAXITER) + + print(f"Job {i+1} completed with UUID={results['UUID']}") + + +def optimization( + OUTPUT_DIR="./output/QA/with-force-penalty/1/optimizations/", + INPUT_FILE="./inputs/input.LandremanPaul2021_QA", + R1 = 0.5, + order = 5, + ncoils = 5, + UUID_init_from=None, + LENGTH_TARGET=5.00, + LENGTH_WEIGHT=1e-03, + CURVATURE_THRESHOLD=12.0, + CURVATURE_WEIGHT=1e-08, + MSC_THRESHOLD=5.00, + MSC_WEIGHT=1e-04, + CC_THRESHOLD=0.083, + CC_WEIGHT=1e+03, + CS_THRESHOLD=0.166, + CS_WEIGHT=1e+03, + FORCE_THRESHOLD=2e+04, + FORCE_WEIGHT=1e-10, + ARCLENGTH_WEIGHT=1e-2, + dx=None, + MAXITER=14000): + """Performs a stage II force optimization based on specified criteria. """ + start_time = time.perf_counter() + + # Initialize the boundary magnetic surface: + nphi = 32 + ntheta = 32 + s = SurfaceRZFourier.from_vmec_input(INPUT_FILE, range="half period", nphi=nphi, ntheta=ntheta) + nfp = s.nfp + R0 = s.get_rc(0, 0) + + # Create a copy of the surface that is closed in theta and phi, and covers the + # full torus toroidally. This is nice for visualization. + nphi_big = nphi * 2 * nfp + 1 + ntheta_big = ntheta + 1 + quadpoints_theta = np.linspace(0, 1, ntheta_big) + quadpoints_phi = np.linspace(0, 1, nphi_big) + surf_big = SurfaceRZFourier( + dofs=s.dofs, + nfp=nfp, + mpol=s.mpol, + ntor=s.ntor, + quadpoints_phi=quadpoints_phi, + quadpoints_theta=quadpoints_theta, + ) + + def initial_base_curves(R0, R1, order, ncoils): + return create_equally_spaced_curves( + ncoils, + nfp, + stellsym=True, + R0=R0, + R1=R1, + order=order, + ) + + if UUID_init_from is None: + base_curves = initial_base_curves(R0, R1, order, ncoils) + total_current = 3e5 + # Since we know the total sum of currents, we only optimize for ncoils-1 + # currents, and then pick the last one so that they all add up to the correct + # value. + base_currents = [Current(total_current / ncoils * 1e-5) * 1e5 for _ in range(ncoils-1)] + # Above, the factors of 1e-5 and 1e5 are included so the current + # degrees of freedom are O(1) rather than ~ MA. The optimization + # algorithm may not perform well if the dofs are scaled badly. + total_current = Current(total_current) + total_current.fix_all() + base_currents += [total_current - sum(base_currents)] + + coils = coils_via_symmetries(base_curves, base_currents, nfp, True) + base_coils = coils[:ncoils] + curves = [c.curve for c in coils] + + bs = BiotSavart(coils) + bs.set_points(s.gamma().reshape((-1, 3))) + else: + path = glob.glob(f"../**/{UUID_init_from}/biot_savart.json", recursive=True)[0] + print("2") + bs = load(path) + coils = bs.coils + curves = [c.curve for c in coils] + base_coils = coils[:ncoils] + base_curves = [base_coils[i].curve for i in range(ncoils)] + base_currents = [base_coils[i].current for i in range(ncoils)] + bs.set_points(s.gamma().reshape((-1, 3))) + + ########################################################################### + ## FORM THE OBJECTIVE FUNCTION ############################################ + + # Define the individual terms objective function: + Jf = SquaredFlux(s, bs) + Jls = [CurveLength(c) for c in base_curves] + Jccdist = CurveCurveDistance(curves, CC_THRESHOLD, num_basecurves=ncoils) + Jcsdist = CurveSurfaceDistance(curves, s, CS_THRESHOLD) + Jcs = [LpCurveCurvature(c, 2, CURVATURE_THRESHOLD) for c in base_curves] + Jmscs = [MeanSquaredCurvature(c) for c in base_curves] + Jforce = [LpCurveForce(c, coils, regularization_circ(0.05), p=2, threshold=FORCE_THRESHOLD) for c in base_coils] + Jals = [ArclengthVariation(c) for c in base_curves] + + # Form the total objective function. + JF = Jf \ + + LENGTH_WEIGHT * sum(QuadraticPenalty(Jl, LENGTH_TARGET, "max") for Jl in Jls) \ + + CC_WEIGHT * Jccdist \ + + CS_WEIGHT * Jcsdist \ + + CURVATURE_WEIGHT * sum(Jcs) \ + + MSC_WEIGHT * sum(QuadraticPenalty(J, MSC_THRESHOLD, "max") for J in Jmscs) \ + + FORCE_WEIGHT * sum(Jforce) \ + + ARCLENGTH_WEIGHT * sum(Jals) + + ########################################################################### + ## PERFORM OPTIMIZATION ################################################### + + def fun(dofs): + JF.x = dofs + J = JF.J() + grad = JF.dJ() + return J, grad + + res = minimize(fun, JF.x, jac=True, method='L-BFGS-B', + options={'maxiter': MAXITER, 'maxcor': 300}, tol=1e-15) + JF.x = res.x + + ########################################################################### + ## EXPORT OPTIMIZATION DATA ############################################### + + # MAKE DIRECTORY FOR EXPORT + + UUID = uuid.uuid4().hex # unique id for each optimization + OUTPUT_DIR = OUTPUT_DIR + UUID + "/" # Directory for output + os.makedirs(OUTPUT_DIR, exist_ok=True) + + + # EXPORT VTKS + + forces = [] + for c in coils: + force = np.linalg.norm(coil_force(c, coils, regularization_circ(0.05)), axis=1) + force = np.append(force, force[0]) + forces = np.concatenate([forces, force]) + pointData_forces = {"F": forces} + curves_to_vtk(curves, OUTPUT_DIR + "curves_opt", close=True, extra_point_data=pointData_forces) + + bs_big = BiotSavart(coils) + bs_big.set_points(surf_big.gamma().reshape((-1, 3))) + pointData = { + "B_N": np.sum( + bs_big.B().reshape((nphi_big, ntheta_big, 3)) * surf_big.unitnormal(), + axis=2, + )[:, :, None] + } + surf_big.to_vtk(OUTPUT_DIR + "surf_opt", extra_data=pointData) + + + # SAVE DATA TO JSON + + BdotN = np.mean(np.abs(np.sum(bs.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2))) + mean_AbsB = np.mean(bs.AbsB()) + max_forces = [np.max(np.linalg.norm(coil_force(c, coils, regularization_circ(0.05)), axis=1)) for c in base_coils] + min_forces = [np.min(np.linalg.norm(coil_force(c, coils, regularization_circ(0.05)), axis=1)) for c in base_coils] + RMS_forces = [np.sqrt(np.mean(np.square(np.linalg.norm(coil_force(c, coils, regularization_circ(0.05)), axis=1)))) for c in base_coils] + results = { + "nfp": nfp, + "ncoils": int(ncoils), + "order": int(order), + "nphi": nphi, + "ntheta": ntheta, + "R0": R0, # if initialized from circles, else None + "R1": R1, # if initialized from circles, else None + "UUID_init": UUID_init_from, # if initialized from optimization, else None + "length_target": LENGTH_TARGET, + "length_weight": LENGTH_WEIGHT, + "max_κ_threshold": CURVATURE_THRESHOLD, + "max_κ_weight": CURVATURE_WEIGHT, + "msc_threshold": MSC_THRESHOLD, + "msc_weight": MSC_WEIGHT, + "cc_threshold": CC_THRESHOLD, + "cc_weight": CC_WEIGHT, + "cs_threshold": CS_THRESHOLD, + "cs_weight": CS_WEIGHT, + "force_threshold": FORCE_THRESHOLD, + "force_weight": FORCE_WEIGHT, + "arclength_weight": ARCLENGTH_WEIGHT, + "JF": float(JF.J()), + "Jf": float(Jf.J()), + "gradient_norm": np.linalg.norm(JF.dJ()), + # "linking_number": LinkingNumber(curves).J(), #TODO: UNCOMMENT!!!!!!! + "lengths": [float(J.J()) for J in Jls], + "max_length": max(float(J.J()) for J in Jls), + "max_κ": [np.max(c.kappa()) for c in base_curves], + "max_max_κ": max(np.max(c.kappa()) for c in base_curves), + "MSCs": [float(J.J()) for J in Jmscs], + "max_MSC": max(float(J.J()) for J in Jmscs), + "max_forces": [float(f) for f in max_forces], + "max_max_force": max(float(f) for f in max_forces), + "min_forces": [float(f) for f in min_forces], + "min_min_force": min(float(f) for f in min_forces), + "RMS_forces": [float(f) for f in RMS_forces], + "mean_RMS_force": float(np.mean([f for f in RMS_forces])), + "arclength_variances": [float(J.J()) for J in Jals], + "max_arclength_variance": max(float(J.J()) for J in Jals), + "BdotN": BdotN, + "mean_AbsB": mean_AbsB, + "normalized_BdotN": BdotN/mean_AbsB, + "coil_coil_distance": Jccdist.shortest_distance(), + "coil_surface_distance": Jcsdist.shortest_distance(), + "message": res.message, + "success": res.success, + "iterations": res.nit, + "function_evaluations": res.nfev, + "coil_currents": [c.get_value() for c in base_currents], + "UUID": UUID, + "eval_time": time.perf_counter() - start_time, + "dx": dx + } + + with open(OUTPUT_DIR + "results.json", "w") as outfile: + json.dump(results , outfile, indent=2) + bs.save(OUTPUT_DIR + f"biot_savart.json") # save the optimized coil shapes and currents + + return res, results, base_coils + + +def rand(min, max): + """Generate a random float between min and max.""" + return np.random.rand() * (max - min) + min \ No newline at end of file