diff --git a/.gitignore b/.gitignore index 5647313e..f66448e6 100644 --- a/.gitignore +++ b/.gitignore @@ -34,8 +34,12 @@ RUN_* *.[0-9]*.[0-9]*.[0-9]* TMP */test_data/simple_*/ -test/test_data/ +test/test_data/* !test/test_data/.gitkeep +!test/test_data/get_test_data.sh +!test/test_data/simple.multiple_fieldline.in +!test/test_data/simple.single_fieldline.in +!test/test_data/simple.volume.in *.geany examples/test1/ examples/test2/ @@ -61,3 +65,5 @@ thirdparty/pyplot_module.F90 .claude docs/_build/ build*/ +artifacts/plots/ +artifacts/notebooks/ diff --git a/requirements.txt b/requirements.txt index 77fb5fe1..c377df85 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,15 +5,16 @@ numpy==2.2.6 scikit-build-core==0.10.0 -# F90wrap - pinned to release version -f90wrap==0.2.16 +# F90wrap - pinned to release version compatible with gvec +f90wrap==0.3.0 # Optional dependencies for testing pytest==7.4.3 pytest-cov==6.2.1 netCDF4==1.7.2 +gvec==1.4.0 # Optional dependencies for examples/plotting matplotlib==3.9.4 scipy==1.15.2 -shapely==2.0.5 +shapely==2.1.2 diff --git a/test/tests/CMakeLists.txt b/test/tests/CMakeLists.txt index 15ae3c15..5859daf1 100644 --- a/test/tests/CMakeLists.txt +++ b/test/tests/CMakeLists.txt @@ -5,6 +5,8 @@ set(TEST_DATA_DIR "${CMAKE_SOURCE_DIR}/test/test_data") set(WOUT_FILE "${TEST_DATA_DIR}/wout.nc") set(WOUT_URL "https://github.com/hiddenSymmetries/simsopt/raw/master/tests/test_files/wout_LandremanPaul2021_QA_reactorScale_lowres_reference.nc") +set(WOUT_QH_FILE "${TEST_DATA_DIR}/wout_qh.nc") +set(WOUT_QH_URL "https://raw.githubusercontent.com/hiddenSymmetries/simsopt/master/tests/test_files/wout_LandremanPaul2021_QH_reactorScale_lowres_reference.nc") # Download if not present if(NOT EXISTS "${WOUT_FILE}") @@ -14,8 +16,20 @@ if(NOT EXISTS "${WOUT_FILE}") message(STATUS "Downloaded VMEC test file") endif() +if(NOT EXISTS "${WOUT_QH_FILE}") + file(MAKE_DIRECTORY "${TEST_DATA_DIR}") + message(STATUS "Downloading QH VMEC file to ${WOUT_QH_FILE}...") + file(DOWNLOAD "${WOUT_QH_URL}" "${WOUT_QH_FILE}") +endif() + # Create symlink in test binary directory for backward compatibility file(CREATE_LINK "${WOUT_FILE}" "${CMAKE_CURRENT_BINARY_DIR}/wout.nc" SYMBOLIC) +file(CREATE_LINK "${WOUT_QH_FILE}" "${CMAKE_CURRENT_BINARY_DIR}/wout_qh.nc" SYMBOLIC) + +set(BOOZER_CHARTMAP_PYTHON "${Python_EXECUTABLE}") +if(EXISTS "${CMAKE_SOURCE_DIR}/.venv/bin/python") + set(BOOZER_CHARTMAP_PYTHON "${CMAKE_SOURCE_DIR}/.venv/bin/python") +endif() # Generate chartmap file from VMEC for testing chartmap coordinate detection # Uses libneo's vmec_to_chartmap tool built as part of dependencies @@ -355,22 +369,46 @@ add_test(NAME test_array_utils COMMAND test_array_utils.x) add_executable(test_boozer_chartmap_roundtrip.x test_boozer_chartmap_roundtrip.f90) target_link_libraries(test_boozer_chartmap_roundtrip.x simple) add_test(NAME test_boozer_chartmap_roundtrip - COMMAND test_boozer_chartmap_roundtrip.x) + COMMAND test_boozer_chartmap_roundtrip.x + ${WOUT_FILE} + ${CMAKE_CURRENT_BINARY_DIR}/roundtrip_test.nc + export + ${CMAKE_CURRENT_BINARY_DIR}/boozer_chartmap_roundtrip + 1.0e-4 + 1.0e-6) set_tests_properties(test_boozer_chartmap_roundtrip PROPERTIES WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} LABELS "slow") + add_test(NAME test_boozer_chartmap_gvec_qa + COMMAND ${BOOZER_CHARTMAP_PYTHON} + ${CMAKE_CURRENT_SOURCE_DIR}/run_gvec_qa_roundtrip.py) + set_tests_properties(test_boozer_chartmap_gvec_qa PROPERTIES + WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} + LABELS "slow;python" + TIMEOUT 7200) + # Export tool for e2e test add_executable(export_boozer_chartmap_tool.x export_boozer_chartmap_tool.f90) target_link_libraries(export_boozer_chartmap_tool.x simple) # E2E test: VMEC-Boozer vs chartmap confined fractions add_test(NAME test_e2e_boozer_chartmap - COMMAND ${Python_EXECUTABLE} + COMMAND ${BOOZER_CHARTMAP_PYTHON} ${CMAKE_CURRENT_SOURCE_DIR}/test_e2e_boozer_chartmap.py) set_tests_properties(test_e2e_boozer_chartmap PROPERTIES WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} - LABELS "slow") + LABELS "slow;plot;python" + TIMEOUT 7200) + + add_test(NAME test_figure8_boozer_chartmap + COMMAND ${BOOZER_CHARTMAP_PYTHON} + ${CMAKE_CURRENT_SOURCE_DIR}/test_figure8_boozer_chartmap.py) + set_tests_properties(test_figure8_boozer_chartmap PROPERTIES + WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} + LABELS "slow;plot;python;golden_record" + TIMEOUT 7200 + DEPENDS test_e2e_boozer_chartmap) if(SIMPLE_ENABLE_CGAL) add_executable(test_stl_wall_intersection.x test_stl_wall_intersection.f90) diff --git a/test/tests/boozer_chartmap_artifacts.py b/test/tests/boozer_chartmap_artifacts.py new file mode 100644 index 00000000..e69a5566 --- /dev/null +++ b/test/tests/boozer_chartmap_artifacts.py @@ -0,0 +1,493 @@ +#!/usr/bin/env python3 +"""Shared artifact helpers for Boozer chartmap integration tests.""" + +from __future__ import annotations + +from pathlib import Path +import urllib.request + +import netCDF4 +import numpy as np + + +STYLES = { + "VMEC-Boozer": ("k", "-", 1.8), + "VMEC chartmap": ("C0", "--", 1.5), + "GVEC chartmap": ("C3", ":", 2.0), + "Figure-8 GVEC chartmap": ("C2", "-", 1.8), +} + + +def ensure_parent(path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + + +def download_if_missing(path: Path, url: str) -> Path: + if path.exists(): + return path + path.parent.mkdir(parents=True, exist_ok=True) + print(f"Downloading {url} -> {path}") + urllib.request.urlretrieve(url, path) + return path + + +def load_confined_fraction(path: Path) -> np.ndarray: + data = np.loadtxt(path) + if data.ndim == 1: + data = data[np.newaxis, :] + return data + + +def confined_metrics(data: np.ndarray) -> dict[str, np.ndarray | float]: + time = np.asarray(data[:, 0], dtype=float) + pass_conf = np.asarray(data[:, 1], dtype=float) + trap_conf = np.asarray(data[:, 2], dtype=float) + npart = float(data[0, 3]) + total_conf = pass_conf + trap_conf + pass_lost = pass_conf[0] - pass_conf + trap_lost = trap_conf[0] - trap_conf + total_lost = total_conf[0] - total_conf + return { + "time": time, + "npart": npart, + "pass_conf": pass_conf, + "trap_conf": trap_conf, + "total_conf": total_conf, + "pass_lost": pass_lost, + "trap_lost": trap_lost, + "total_lost": total_lost, + } + + +def summarize_confined_fraction(data: np.ndarray) -> dict[str, float]: + metrics = confined_metrics(data) + return { + "npart": float(metrics["npart"]), + "initial_total_conf": float(metrics["total_conf"][0]), + "final_total_conf": float(metrics["total_conf"][-1]), + "final_total_lost": float(metrics["total_lost"][-1]), + "final_pass_lost": float(metrics["pass_lost"][-1]), + "final_trap_lost": float(metrics["trap_lost"][-1]), + } + + +def _plot_time_axis(time: np.ndarray) -> np.ndarray: + positive = np.asarray(time[time > 0.0], dtype=float) + if positive.size == 0: + return np.ones_like(time) + floor = positive.min() / 2.0 + return np.where(time > 0.0, time, floor) + + +def _style(label: str) -> tuple[str, str, float]: + return STYLES.get(label, ("0.3", "-", 1.4)) + + +def _set_equal_3d_limits(ax, x: np.ndarray, y: np.ndarray, z: np.ndarray) -> None: + x = np.asarray(x, dtype=float) + y = np.asarray(y, dtype=float) + z = np.asarray(z, dtype=float) + mins = np.array([np.min(x), np.min(y), np.min(z)], dtype=float) + maxs = np.array([np.max(x), np.max(y), np.max(z)], dtype=float) + center = 0.5 * (mins + maxs) + radius = 0.55 * np.max(maxs - mins) + if radius <= 0.0: + radius = 1.0 + ax.set_xlim(center[0] - radius, center[0] + radius) + ax.set_ylim(center[1] - radius, center[1] + radius) + ax.set_zlim(center[2] - radius, center[2] + radius) + ax.set_box_aspect((1.0, 1.0, 1.0)) + + +def _close_curve(curve: np.ndarray) -> np.ndarray: + curve = np.asarray(curve) + if curve.ndim != 1 or curve.size == 0: + return curve + return np.concatenate([curve, curve[:1]]) + + +def _close_surface_toroidally( + xsurf: np.ndarray, ysurf: np.ndarray, zsurf: np.ndarray +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + xsurf = np.asarray(xsurf) + ysurf = np.asarray(ysurf) + zsurf = np.asarray(zsurf) + if xsurf.ndim != 2 or xsurf.shape[1] == 0: + return xsurf, ysurf, zsurf + return ( + np.concatenate([xsurf, xsurf[:, :1]], axis=1), + np.concatenate([ysurf, ysurf[:, :1]], axis=1), + np.concatenate([zsurf, zsurf[:, :1]], axis=1), + ) + + +def _sample_indices(size: int, stride: int) -> list[int]: + indices = list(range(0, size, stride)) + last = size - 1 + if last not in indices: + indices.append(last) + return indices + + +def _plot_surface_lines_3d( + ax, + xsurf: np.ndarray, + ysurf: np.ndarray, + zsurf: np.ndarray, + color: str, + label: str, + axis_x: np.ndarray, + axis_y: np.ndarray, + axis_z: np.ndarray, + linestyle: str = "-", +) -> None: + xsurf, ysurf, zsurf = _close_surface_toroidally(xsurf, ysurf, zsurf) + axis_x = _close_curve(axis_x) + axis_y = _close_curve(axis_y) + axis_z = _close_curve(axis_z) + n_phi = xsurf.shape[1] + n_rho = xsurf.shape[0] + phi_stride = max(1, n_phi // 12) + rho_stride = max(1, n_rho // 12) + for idx in _sample_indices(n_phi, phi_stride): + ax.plot( + xsurf[:, idx], + ysurf[:, idx], + zsurf[:, idx], + color=color, + lw=0.7, + alpha=0.45, + ls=linestyle, + ) + for idx in _sample_indices(n_rho, rho_stride): + ax.plot( + xsurf[idx], + ysurf[idx], + zsurf[idx], + color=color, + lw=0.5, + alpha=0.22, + ls=linestyle, + ) + ax.plot(axis_x, axis_y, axis_z, color=color, lw=1.8, ls=linestyle, label=label) + + +def plot_case_loss_comparison( + output: Path, + case_name: str, + metrics_by_label: dict[str, dict[str, np.ndarray | float]], +) -> None: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + ensure_parent(output) + ref = metrics_by_label.get("VMEC-Boozer") + if ref is None: + ref = next(iter(metrics_by_label.values())) + time_ms = _plot_time_axis(np.asarray(ref["time"])) * 1.0e3 + + panels = [ + ("total_conf", "Total confined", False), + ("total_lost", "Total lost", False), + ("pass_conf", "Passing confined", False), + ("pass_lost", "Passing lost", False), + ("trap_conf", "Trapped confined", False), + ("trap_lost", "Trapped lost", False), + ] + fig, axes = plt.subplots(2, 3, figsize=(15, 8), constrained_layout=True) + for ax, (key, title, semilog_y) in zip(axes.flat, panels): + for label, metrics in metrics_by_label.items(): + color, linestyle, width = _style(label) + y = np.asarray(metrics[key], dtype=float) + if semilog_y: + ax.semilogy(time_ms, y, color=color, ls=linestyle, lw=width, label=label) + else: + ax.plot(time_ms, y, color=color, ls=linestyle, lw=width, label=label) + ax.set_xscale("log") + ax.set_xlabel("time [ms]") + ax.set_ylabel("fraction") + ax.set_ylim(-0.05, 1.05) + ax.set_title(title) + axes[0, 0].legend(fontsize=8) + fig.suptitle(f"{case_name}: confined and lost fractions", fontsize=14) + fig.savefig(output, dpi=180) + print(f"Saved {output}") + + +def plot_all_cases_total_losses( + output: Path, + cases: dict[str, dict[str, dict[str, np.ndarray | float]]], +) -> None: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + ensure_parent(output) + names = list(cases) + panels = [ + ("total_lost", "Total lost"), + ("pass_lost", "Passing lost"), + ("trap_lost", "Trapped lost"), + ] + fig, axes = plt.subplots( + len(names), + len(panels), + figsize=(15, 3.4 * len(names)), + constrained_layout=True, + ) + if len(names) == 1: + axes = np.asarray([axes]) + for row, case_name in enumerate(names): + metrics_by_label = cases[case_name] + ref = next(iter(metrics_by_label.values())) + time_ms = _plot_time_axis(np.asarray(ref["time"])) * 1.0e3 + for col, (key, title) in enumerate(panels): + ax = axes[row, col] + for label, metrics in metrics_by_label.items(): + color, linestyle, width = _style(label) + ax.plot( + time_ms, + np.asarray(metrics[key], dtype=float), + color=color, + ls=linestyle, + lw=width, + label=label, + ) + ax.set_xscale("log") + ax.set_ylim(-0.05, 1.05) + ax.set_xlabel("time [ms]") + if col == 0: + ax.set_ylabel(case_name) + ax.set_title(title) + if row == 0 and col == 0: + ax.legend(fontsize=8) + fig.suptitle("Loss comparison across all benchmark cases", fontsize=14) + fig.savefig(output, dpi=180) + print(f"Saved {output}") + + +def load_chartmap_surface(path: Path) -> dict[str, np.ndarray | int]: + with netCDF4.Dataset(path) as ds: + x = np.asarray(ds["x"][:]).transpose(2, 1, 0) * 1.0e-2 + y = np.asarray(ds["y"][:]).transpose(2, 1, 0) * 1.0e-2 + z = np.asarray(ds["z"][:]).transpose(2, 1, 0) * 1.0e-2 + nfp = int(np.asarray(ds["num_field_periods"][:]).item()) + return { + "boundary_x": x[-1], + "boundary_y": y[-1], + "boundary_z": z[-1], + "axis_x": x[0], + "axis_y": y[0], + "axis_z": z[0], + "nfp": nfp, + } + + +def load_chartmap_fields(path: Path) -> dict[str, np.ndarray]: + with netCDF4.Dataset(path) as ds: + rho = np.asarray(ds["rho"][:], dtype=float) + theta = np.asarray(ds["theta"][:], dtype=float) + zeta = np.asarray(ds["zeta"][:], dtype=float) + bmod = np.asarray(ds["Bmod"][:], dtype=float).transpose(2, 1, 0) + a_phi = np.asarray(ds["A_phi"][:], dtype=float) + b_theta = np.asarray(ds["B_theta"][:], dtype=float) + b_phi = np.asarray(ds["B_phi"][:], dtype=float) + return { + "rho": rho, + "theta": theta, + "zeta": zeta, + "Bmod": bmod, + "A_phi": a_phi, + "B_theta": b_theta, + "B_phi": b_phi, + } + + +def tile_field_periods(xsurf: np.ndarray, ysurf: np.ndarray, zsurf: np.ndarray, nfp: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + x_parts = [] + y_parts = [] + z_parts = [] + for k in range(nfp): + angle = 2.0 * np.pi * k / nfp + cang = np.cos(angle) + sang = np.sin(angle) + x_parts.append(cang * xsurf - sang * ysurf) + y_parts.append(sang * xsurf + cang * ysurf) + z_parts.append(zsurf.copy()) + return np.concatenate(x_parts, axis=1), np.concatenate(y_parts, axis=1), np.concatenate(z_parts, axis=1) + + +def plot_surface_comparison( + output: Path, + case_name: str, + left_path: Path, + right_path: Path, + left_label: str, + right_label: str, +) -> None: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + ensure_parent(output) + left = load_chartmap_surface(left_path) + right = load_chartmap_surface(right_path) + left_x, left_y, left_z = tile_field_periods(left["boundary_x"], left["boundary_y"], left["boundary_z"], int(left["nfp"])) + right_x, right_y, right_z = tile_field_periods(right["boundary_x"], right["boundary_y"], right["boundary_z"], int(right["nfp"])) + left_ax_x, left_ax_y, left_ax_z = tile_field_periods(left["axis_x"], left["axis_y"], left["axis_z"], int(left["nfp"])) + right_ax_x, right_ax_y, right_ax_z = tile_field_periods(right["axis_x"], right["axis_y"], right["axis_z"], int(right["nfp"])) + + fig = plt.figure(figsize=(15, 8.5), constrained_layout=True) + axes = np.empty((2, 3), dtype=object) + for i in range(3): + axes[0, i] = fig.add_subplot(2, 3, i + 1) + for i in range(3): + axes[1, i] = fig.add_subplot(2, 3, i + 4, projection="3d") + + panels = [ + ("x-z", left_x, left_z, right_x, right_z, left_ax_x, left_ax_z, right_ax_x, right_ax_z, "x [m]", "z [m]"), + ("y-z", left_y, left_z, right_y, right_z, left_ax_y, left_ax_z, right_ax_y, right_ax_z, "y [m]", "z [m]"), + ("x-y", left_x, left_y, right_x, right_y, left_ax_x, left_ax_y, right_ax_x, right_ax_y, "x [m]", "y [m]"), + ] + for ax, (title, xl, yl, xr, yr, xal, yal, xar, yar, xlabel, ylabel) in zip(axes[0], panels): + xl, yl, _ = _close_surface_toroidally(xl, yl, yl) + xr, yr, _ = _close_surface_toroidally(xr, yr, yr) + for idx in _sample_indices(xl.shape[1], max(1, xl.shape[1] // 12)): + ax.plot(xl[:, idx], yl[:, idx], color="C0", lw=0.8, alpha=0.45) + for idx in _sample_indices(xr.shape[1], max(1, xr.shape[1] // 12)): + ax.plot(xr[:, idx], yr[:, idx], color="C3", lw=0.8, alpha=0.45) + for idx in _sample_indices(xl.shape[0], max(1, xl.shape[0] // 12)): + ax.plot(xl[idx], yl[idx], color="C0", lw=0.6, alpha=0.3) + for idx in _sample_indices(xr.shape[0], max(1, xr.shape[0] // 12)): + ax.plot(xr[idx], yr[idx], color="C3", lw=0.6, alpha=0.3) + ax.plot(_close_curve(xal.mean(axis=0)), _close_curve(yal.mean(axis=0)), + color="C0", lw=2.0, label=left_label) + ax.plot(_close_curve(xar.mean(axis=0)), _close_curve(yar.mean(axis=0)), + color="C3", lw=1.6, ls="--", label=right_label) + ax.set_aspect("equal", adjustable="box") + ax.set_xlabel(xlabel) + ax.set_ylabel(ylabel) + ax.set_title(title) + axes[0, 0].legend(fontsize=8) + + views = [(22.0, -62.0), (18.0, 22.0), (78.0, -90.0)] + all_x = np.concatenate([left_x.ravel(), right_x.ravel(), left_ax_x.ravel(), right_ax_x.ravel()]) + all_y = np.concatenate([left_y.ravel(), right_y.ravel(), left_ax_y.ravel(), right_ax_y.ravel()]) + all_z = np.concatenate([left_z.ravel(), right_z.ravel(), left_ax_z.ravel(), right_ax_z.ravel()]) + for ax, (elev, azim) in zip(axes[1], views): + _plot_surface_lines_3d( + ax, + left_x, + left_y, + left_z, + "C0", + left_label, + left_ax_x.mean(axis=0), + left_ax_y.mean(axis=0), + left_ax_z.mean(axis=0), + ) + _plot_surface_lines_3d( + ax, + right_x, + right_y, + right_z, + "C3", + right_label, + right_ax_x.mean(axis=0), + right_ax_y.mean(axis=0), + right_ax_z.mean(axis=0), + linestyle="--", + ) + ax.view_init(elev=elev, azim=azim) + _set_equal_3d_limits(ax, all_x, all_y, all_z) + ax.set_xlabel("x [m]") + ax.set_ylabel("y [m]") + ax.set_zlabel("z [m]") + ax.set_title(f"3D view elev={elev:.0f}, azim={azim:.0f}") + axes[1, 0].legend(fontsize=8) + fig.suptitle(f"{case_name}: flux-surface shape comparison", fontsize=14) + fig.savefig(output, dpi=180) + print(f"Saved {output}") + + +def plot_field_component_comparison( + output: Path, + case_name: str, + left_path: Path, + right_path: Path, + left_label: str, + right_label: str, +) -> dict[str, float]: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + ensure_parent(output) + left = load_chartmap_fields(left_path) + right = load_chartmap_fields(right_path) + + rho_common = np.linspace(0.0, 1.0, min(len(left["rho"]), len(right["rho"]))) + left_a = np.interp(rho_common, left["rho"], left["A_phi"]) + right_a = np.interp(rho_common, right["rho"], right["A_phi"]) + left_bt = np.interp(rho_common, left["rho"], left["B_theta"]) + right_bt = np.interp(rho_common, right["rho"], right["B_theta"]) + left_bp = np.interp(rho_common, left["rho"], left["B_phi"]) + right_bp = np.interp(rho_common, right["rho"], right["B_phi"]) + + bmod_left = np.asarray(left["Bmod"]) + bmod_right = np.asarray(right["Bmod"]) + nrho = min(bmod_left.shape[0], bmod_right.shape[0]) + ntheta = min(bmod_left.shape[1], bmod_right.shape[1]) + nzeta = min(bmod_left.shape[2], bmod_right.shape[2]) + bmod_left = bmod_left[:nrho, :ntheta, :nzeta] + bmod_right = bmod_right[:nrho, :ntheta, :nzeta] + mid = nrho // 2 + rel_bmod = np.abs(bmod_right[mid] - bmod_left[mid]) / np.maximum(np.abs(bmod_left[mid]), 1.0e-14) + + metrics = { + "max_rel_A_phi": float(np.max(np.abs(right_a - left_a) / np.maximum(np.abs(left_a), 1.0e-14))), + "max_rel_B_theta": float(np.max(np.abs(right_bt - left_bt) / np.maximum(np.abs(left_bt), 1.0e-14))), + "max_rel_B_phi": float(np.max(np.abs(right_bp - left_bp) / np.maximum(np.abs(left_bp), 1.0e-14))), + "max_rel_Bmod": float(np.max(rel_bmod)), + } + + fig, axes = plt.subplots(2, 3, figsize=(15, 8), constrained_layout=True) + profiles = [ + ("A_phi", left_a, right_a, "A_phi"), + ("B_theta", left_bt, right_bt, "B_theta"), + ("B_phi", left_bp, right_bp, "B_phi"), + ] + for ax, (title, left_vals, right_vals, key) in zip(axes[0], profiles): + ax.plot(rho_common, left_vals, color="C0", lw=1.8, label=left_label) + ax.plot(rho_common, right_vals, color="C3", lw=1.4, ls="--", label=right_label) + ax.set_xlabel("rho") + ax.set_title(f"{title} (max rel {metrics[f'max_rel_{key}']:.2e})") + ax.legend(fontsize=8) + + im0 = axes[1, 0].imshow(bmod_left[mid].T, origin="lower", aspect="auto", cmap="magma") + axes[1, 0].set_title(f"{left_label} |B|") + axes[1, 0].set_xlabel("zeta index") + axes[1, 0].set_ylabel("theta index") + fig.colorbar(im0, ax=axes[1, 0], shrink=0.8) + + im1 = axes[1, 1].imshow(bmod_right[mid].T, origin="lower", aspect="auto", cmap="magma") + axes[1, 1].set_title(f"{right_label} |B|") + axes[1, 1].set_xlabel("zeta index") + axes[1, 1].set_ylabel("theta index") + fig.colorbar(im1, ax=axes[1, 1], shrink=0.8) + + im2 = axes[1, 2].imshow(np.log10(np.maximum(rel_bmod.T, 1.0e-16)), origin="lower", aspect="auto", cmap="viridis") + axes[1, 2].set_title(f"log10 rel |B| diff (max {metrics['max_rel_Bmod']:.2e})") + axes[1, 2].set_xlabel("zeta index") + axes[1, 2].set_ylabel("theta index") + fig.colorbar(im2, ax=axes[1, 2], shrink=0.8) + + fig.suptitle(f"{case_name}: chartmap field comparison", fontsize=14) + fig.savefig(output, dpi=180) + print(f"Saved {output}") + return metrics diff --git a/test/tests/export_boozer_chartmap_tool.f90 b/test/tests/export_boozer_chartmap_tool.f90 index f7b805af..73f770aa 100644 --- a/test/tests/export_boozer_chartmap_tool.f90 +++ b/test/tests/export_boozer_chartmap_tool.f90 @@ -15,10 +15,9 @@ program export_boozer_chartmap_tool implicit none - real(dp), parameter :: twopi = 2.0_dp * 3.14159265358979_dp + real(dp), parameter :: twopi = 8.0_dp * atan(1.0_dp) character(len=1024) :: wout_file, chartmap_file, start_vmec, start_boozer - character(len=1024) :: line - integer :: nargs, ipart, npart, ios + integer :: nargs, ipart, npart, ios, u_in, u_out real(dp) :: s, theta_v, phi_v, v, lam, theta_b, phi_b real(dp) :: RT0, R0i, cbfi, bz0i, bf0, fper, rho integer :: L1i @@ -37,7 +36,7 @@ program export_boozer_chartmap_tool ! Initialize VMEC netcdffile = wout_file - multharm = 3 + multharm = 7 ns_A = 5 ns_s = 5 ns_tp = 5 @@ -57,26 +56,26 @@ program export_boozer_chartmap_tool ! Count particles in start_vmec npart = 0 - open(unit=10, file=trim(start_vmec), status='old', iostat=ios) + open(newunit=u_in, file=trim(start_vmec), status='old', iostat=ios) if (ios /= 0) then print *, 'Cannot open ', trim(start_vmec) error stop end if do - read(10, *, iostat=ios) + read(u_in, *, iostat=ios) if (ios /= 0) exit npart = npart + 1 end do - close(10) + close(u_in) print *, 'Converting', npart, ' particles from VMEC to Boozer coords' ! Convert start.dat coordinates - open(unit=10, file=trim(start_vmec), status='old') - open(unit=11, file=trim(start_boozer), status='replace', recl=1024) + open(newunit=u_in, file=trim(start_vmec), status='old') + open(newunit=u_out, file=trim(start_boozer), status='replace', recl=1024) do ipart = 1, npart - read(10, *) s, theta_v, phi_v, v, lam + read(u_in, *) s, theta_v, phi_v, v, lam ! Transform VMEC angles to Boozer angles call vmec_to_boozer(s, theta_v, phi_v, theta_b, phi_b) @@ -84,11 +83,11 @@ program export_boozer_chartmap_tool ! In chartmap reference coords: x(1) = rho = sqrt(s) rho = sqrt(max(s, 0.0_dp)) - write(11, *) rho, theta_b, phi_b, v, lam + write(u_out, *) rho, theta_b, phi_b, v, lam end do - close(10) - close(11) + close(u_in) + close(u_out) print *, 'Written ', trim(start_boozer) diff --git a/test/tests/figure8_signature_reference.txt b/test/tests/figure8_signature_reference.txt new file mode 100644 index 00000000..4177770e --- /dev/null +++ b/test/tests/figure8_signature_reference.txt @@ -0,0 +1,146 @@ +9.426185672491718126e-01 +8.834420930609492295e-01 +7.037164587517260594e-01 +3.390459631671917107e-01 +-4.051886916911425635e-02 +-4.077811460936539012e-01 +-7.349968908917946520e-01 +-8.781660840032984305e-01 +-9.017458980501251720e-01 +-7.733602345007570822e-01 +-5.341604602490299669e-01 +-1.935346431443295545e-01 +2.661716154896967002e-01 +5.985131155485385879e-01 +8.261361221106177100e-01 +9.402647944726206797e-01 +4.476569126238977292e-19 +3.824612531832787821e-01 +7.276912074623338444e-01 +1.015537700434931967e+00 +1.081017259743941583e+00 +9.637516780356345514e-01 +6.330146290260453279e-01 +2.708007431973673551e-01 +-1.167106833913919145e-01 +-5.646722185319986709e-01 +-8.720781760388673920e-01 +-1.055198731691180392e+00 +-1.043061523929766032e+00 +-8.437356268909649293e-01 +-5.276305057092559592e-01 +-7.752638294572805622e-02 +-1.352234925445357277e-16 +5.188998999248468219e-02 +7.826088774351563915e-02 +5.795718067194196438e-02 +9.358835589908333363e-03 +-3.732432475180267467e-02 +-5.398997371465466544e-02 +-2.931813136881604442e-02 +1.310282952479262936e-02 +5.152290528997031938e-02 +4.881757904556874306e-02 +1.136681578667048109e-02 +-4.965797215754854527e-02 +-7.785476831160662026e-02 +-6.668010487812109366e-02 +-1.126104689796411130e-02 +9.945131977151533409e-01 +9.300901009909475770e-01 +7.372995371893737993e-01 +3.533386702239391663e-01 +-4.299784007251870183e-02 +-4.270122780878211066e-01 +-7.730264309434237147e-01 +-9.273808160325895944e-01 +-9.530906372514760250e-01 +-8.141314671619155341e-01 +-5.599255998669760981e-01 +-2.027568979405206806e-01 +2.771402801799160431e-01 +6.258491581836429507e-01 +8.682033496138347184e-01 +9.919396072357158367e-01 +4.650232945306105592e-02 +1.021287776348148157e-01 +1.380245957416851854e-01 +1.298152562868429838e-01 +8.453660376431243462e-02 +3.212136882814912303e-02 +2.405278915216593760e-03 +1.906032796730530374e-02 +6.020530597401081013e-02 +1.063913612855028473e-01 +1.143792539424904053e-01 +8.510604989321761404e-02 +2.338548540907333176e-02 +-1.396766689285065810e-02 +-1.342583183040768624e-02 +3.536853198940001497e-02 +0.000000000000000000e+00 +-1.631439234561927515e+04 +-6.527235347507253027e+04 +-1.671515065091832657e+05 +-2.880501226730306516e+05 +-4.764321952943700599e+05 +-6.696097976783593185e+05 +-9.450201254796257708e+05 +-1.211014401860102313e+06 +-1.574218781877085567e+06 +-4.683025941858154351e-04 +-4.864125260545248963e+00 +-1.946228343513434211e+01 +-4.985369625704061036e+01 +-8.594098738257116565e+01 +-1.422144884968279257e+02 +-1.999654765252323330e+02 +-2.823561545424180395e+02 +-3.619621765770242519e+02 +-4.706482782084186738e+02 +1.301065901833702177e+08 +1.301067120242614299e+08 +1.301070778394054621e+08 +1.301078399650222361e+08 +1.301087459301267564e+08 +1.301101610650103092e+08 +1.301116165561520010e+08 +1.301136987670112699e+08 +1.301157176730800718e+08 +1.301184866313546151e+08 +1.314282866236967966e+06 +1.314680683571781730e+06 +1.313609436181156430e+06 +1.313487714778708294e+06 +1.314576849644944072e+06 +1.314282866236967966e+06 +1.297289014025866520e+06 +1.296700441367030377e+06 +1.296532885438913945e+06 +1.296090564464562340e+06 +1.296581550345291384e+06 +1.297289014025866054e+06 +1.267848038300013868e+06 +1.268181982878365554e+06 +1.266924437549313996e+06 +1.267404867868238827e+06 +1.267550586654868210e+06 +1.267848038300013868e+06 +1.265138273087826325e+06 +1.265211361547151580e+06 +1.264865994408159750e+06 +1.264238433185068890e+06 +1.265430867158013629e+06 +1.265138273087826325e+06 +1.292883455894635990e+06 +1.292254021328466013e+06 +1.291785495439160382e+06 +1.291955600267246133e+06 +1.292116676961044082e+06 +1.292883455894635757e+06 +1.314282866236967966e+06 +1.314680683571781730e+06 +1.313609436181156430e+06 +1.313487714778709225e+06 +1.314576849644944072e+06 +1.314282866236968432e+06 diff --git a/test/tests/gvec_to_boozer_chartmap.py b/test/tests/gvec_to_boozer_chartmap.py new file mode 100644 index 00000000..33fd34bc --- /dev/null +++ b/test/tests/gvec_to_boozer_chartmap.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +"""Convert a GVEC state to a Boozer chartmap NetCDF file. + +Uses GVEC Python library to evaluate fields in Boozer coordinates +and writes the result in the extended chartmap format that SIMPLE +can read without any GVEC or VMEC library at runtime. + +Usage: + .venv/bin/python gvec_to_boozer_chartmap.py +""" + +import argparse +import sys +import numpy as np + +try: + import gvec +except ImportError: + print("ERROR: gvec Python package required. Install via pip or use the .venv") + sys.exit(1) + +try: + import netCDF4 +except ImportError: + print("ERROR: netCDF4 required. pip install netCDF4") + sys.exit(1) + + +def main(): + parser = argparse.ArgumentParser( + description="Convert GVEC state to Boozer chartmap NetCDF" + ) + parser.add_argument("paramfile", help="GVEC parameter file (.ini)") + parser.add_argument("statefile", help="GVEC state file (.dat)") + parser.add_argument("output", help="Output Boozer chartmap .nc file") + parser.add_argument("--nrho", type=int, default=50) + parser.add_argument("--ntheta", type=int, default=36) + parser.add_argument("--nphi", type=int, default=81) + args = parser.parse_args() + + print(f"Loading GVEC state: {args.paramfile} + {args.statefile}") + state = gvec.load_state(args.paramfile, args.statefile) + nfp = int(state.nfp) + twopi = 2.0 * np.pi + phi_period = twopi / nfp + + n_rho = args.nrho + # Geometry: endpoint-excluded (chartmap validator) + n_theta_geom = args.ntheta + n_phi_geom = args.nphi + # Field: endpoint-included (exact spline reproduction) + n_theta_field = n_theta_geom + 1 + n_phi_field = n_phi_geom + 1 + + rho_grid = np.linspace(0, 1, n_rho) + rho_grid[0] = 1e-3 + theta_geom = np.linspace(0, twopi, n_theta_geom, endpoint=False) + zeta_geom = np.linspace(0, phi_period, n_phi_geom, endpoint=False) + theta_field = np.linspace(0, twopi, n_theta_field) + zeta_field = np.linspace(0, phi_period, n_phi_field) + + # Evaluate field quantities in Boozer coordinates (endpoint-included grid) + print(f"Evaluating field on {n_rho} x {n_theta_field} x {n_phi_field}...") + ev = gvec.EvaluationsBoozer( + rho=rho_grid, theta_B=theta_field, zeta_B=zeta_field, state=state) + gvec.compute(ev, 'mod_B', state=state) + gvec.compute(ev, 'B_theta_B', state=state) + gvec.compute(ev, 'B_zeta_B', state=state) + gvec.compute(ev, 'iota', state=state) + + Bmod_3d = ev['mod_B'].values # (n_rho, n_theta_field, n_phi_field) + B_theta_3d = ev['B_theta_B'].values + B_phi_3d = ev['B_zeta_B'].values + iota_vals = ev['iota'].values + # Surface functions: take mean over angles (should be constant) + if B_theta_3d.ndim == 3: + B_theta = B_theta_3d[:, 0, 0] + B_phi = B_phi_3d[:, 0, 0] + else: + B_theta = B_theta_3d + B_phi = B_phi_3d + iota_prof = iota_vals.ravel()[:n_rho] if iota_vals.ndim > 1 else iota_vals + s_grid = rho_grid**2 + + # Toroidal flux: Phi_edge from evaluations + ev_edge = gvec.Evaluations(rho=np.array([1.0]), + theta=np.array([0.0]), + zeta=np.array([0.0]), state=state) + gvec.compute(ev_edge, 'Phi_edge', state=state) + # GVEC Phi_edge is already Phi/(2*pi) in SI (Wb/(2*pi) = T*m^2/(2*pi)) + torflux_SI = float(ev_edge['Phi_edge'].values.flat[0]) + + # CGS conversion (GVEC: SI, SIMPLE: CGS) + Bmod_3d = Bmod_3d * 1e4 # T -> G + B_theta = B_theta * 1e4 * 1e2 # T*m -> G*cm + B_phi = B_phi * 1e4 * 1e2 + torflux = torflux_SI * 1e8 # T*m^2/(2pi) -> G*cm^2 + + # A_phi: integrate -torflux * iota over s + A_phi = np.zeros(n_rho) + for i in range(1, n_rho): + ds = s_grid[i] - s_grid[i - 1] + A_phi[i] = A_phi[i - 1] - torflux * iota_prof[i] * ds + + # Geometry on endpoint-excluded grid (Boozer angles) + print(f"Evaluating geometry on {n_rho} x {n_theta_geom} x {n_phi_geom}...") + ev_geom = gvec.EvaluationsBoozer( + rho=rho_grid, theta_B=theta_geom, zeta_B=zeta_geom, state=state) + gvec.compute(ev_geom, 'pos', state=state) + pos = ev_geom['pos'].values # (3, n_rho, n_theta_geom, n_phi_geom) + # pos[0]=X, pos[1]=Y, pos[2]=Z in meters + + X = pos[0] * 1e2 # m -> cm + Y = pos[1] * 1e2 + Z_geom = pos[2] * 1e2 + + # Write NetCDF + print(f"Writing {args.output}") + ds = netCDF4.Dataset(args.output, "w", format="NETCDF4") + + ds.createDimension("rho", n_rho) + ds.createDimension("theta", n_theta_geom) + ds.createDimension("zeta", n_phi_geom) + ds.createDimension("theta_field", n_theta_field) + ds.createDimension("zeta_field", n_phi_field) + + v = ds.createVariable("rho", "f8", ("rho",)); v[:] = rho_grid + v = ds.createVariable("theta", "f8", ("theta",)); v[:] = theta_geom + v = ds.createVariable("zeta", "f8", ("zeta",)); v[:] = zeta_geom + + # NetCDF dimensions are (zeta, theta, rho) so NF90 reads as (rho, theta, zeta) + for name, arr in [("x", X), ("y", Y), ("z", Z_geom)]: + v = ds.createVariable(name, "f8", ("zeta", "theta", "rho")) + v[:] = np.transpose(arr, (2, 1, 0)) + v.units = "cm" + + v = ds.createVariable("A_phi", "f8", ("rho",)); v[:] = A_phi + v = ds.createVariable("B_theta", "f8", ("rho",)); v[:] = B_theta + v = ds.createVariable("B_phi", "f8", ("rho",)); v[:] = B_phi + # Bmod uses field dimensions (also reversed for NF90) + v = ds.createVariable("Bmod", "f8", ("zeta_field", "theta_field", "rho")) + v[:] = np.transpose(Bmod_3d, (2, 1, 0)) + + v = ds.createVariable("num_field_periods", "i4"); v[:] = np.int32(nfp) + + ds.rho_convention = "rho_tor" + ds.zeta_convention = "boozer" + ds.rho_lcfs = float(rho_grid[-1]) + ds.boozer_field = np.int32(1) + ds.torflux = torflux + + ds.close() + print(f"Done. nfp={nfp}, torflux={torflux:.6e}") + + +if __name__ == "__main__": + main() diff --git a/test/tests/plot_boozer_chartmap_roundtrip.py b/test/tests/plot_boozer_chartmap_roundtrip.py index 1fd21966..b8be6382 100644 --- a/test/tests/plot_boozer_chartmap_roundtrip.py +++ b/test/tests/plot_boozer_chartmap_roundtrip.py @@ -1,14 +1,16 @@ #!/usr/bin/env python3 -"""Plot Boozer chartmap roundtrip comparison: fields and orbits side by side.""" +"""Plot Boozer chartmap comparison: fields and orbits side by side.""" +import argparse +from pathlib import Path import numpy as np import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt -def plot_field_comparison(): - data = np.loadtxt("/tmp/boozer_field_comparison.dat") +def plot_field_comparison(prefix, direct_label, chartmap_label, title): + data = np.loadtxt(f"{prefix}_field_comparison.dat") s = data[:, 0] theta = data[:, 1] Bmod_ref = data[:, 3] @@ -18,16 +20,16 @@ def plot_field_comparison(): fig, axes = plt.subplots(1, 3, figsize=(15, 4.5)) ax = axes[0] - ax.scatter(s, Bmod_ref, s=15, label="Direct (VMEC)", alpha=0.7) - ax.scatter(s, Bmod_new, s=15, marker="x", label="Chartmap", alpha=0.7) + ax.scatter(s, Bmod_ref, s=15, label=direct_label, alpha=0.7) + ax.scatter(s, Bmod_new, s=15, marker="x", label=chartmap_label, alpha=0.7) ax.set_xlabel("s") ax.set_ylabel("|B| [G]") ax.set_title("Bmod vs s") ax.legend() ax = axes[1] - ax.scatter(theta, Bmod_ref, s=15, label="Direct (VMEC)", alpha=0.7) - ax.scatter(theta, Bmod_new, s=15, marker="x", label="Chartmap", alpha=0.7) + ax.scatter(theta, Bmod_ref, s=15, label=direct_label, alpha=0.7) + ax.scatter(theta, Bmod_new, s=15, marker="x", label=chartmap_label, alpha=0.7) ax.set_xlabel("theta_B") ax.set_ylabel("|B| [G]") ax.set_title("Bmod vs theta") @@ -41,16 +43,17 @@ def plot_field_comparison(): ax.axhline(1e-5, color="r", ls="--", label="tolerance") ax.legend() - fig.suptitle("Boozer Chartmap Field Roundtrip", fontsize=14) + fig.suptitle(f"{title}: field comparison", fontsize=14) fig.tight_layout() - fig.savefig("/tmp/boozer_field_comparison.png", dpi=150) - print("Saved /tmp/boozer_field_comparison.png") + output = f"{prefix}_field_comparison.png" + fig.savefig(output, dpi=150) + print(f"Saved {output}") -def plot_orbit_comparison(): +def plot_orbit_comparison(prefix, direct_label, chartmap_label, title): try: - direct = np.loadtxt("/tmp/orbit_direct.dat") - chartmap = np.loadtxt("/tmp/orbit_chartmap.dat") + direct = np.loadtxt(f"{prefix}_orbit_direct.dat") + chartmap = np.loadtxt(f"{prefix}_orbit_chartmap.dat") except Exception as e: print(f"Could not load orbit data: {e}") return @@ -61,47 +64,61 @@ def plot_orbit_comparison(): t_d, s_d, th_d, ph_d = direct[:, 0], direct[:, 1], direct[:, 2], direct[:, 3] t_c, s_c, th_c, ph_c = chartmap[:, 0], chartmap[:, 1], chartmap[:, 2], chartmap[:, 3] + ds = np.abs(s_d - s_c) + dtheta = th_c - th_d + dphi = ph_c - ph_d fig, axes = plt.subplots(2, 2, figsize=(12, 8)) ax = axes[0, 0] - ax.plot(t_d, s_d, label="Direct", lw=1) - ax.plot(t_c, s_c, "--", label="Chartmap", lw=1) + ax.plot(t_d, s_d, label=direct_label, lw=1) + ax.plot(t_c, s_c, "--", label=chartmap_label, lw=1) ax.set_xlabel("time") ax.set_ylabel("s") ax.set_title("Flux surface label") ax.legend() ax = axes[0, 1] - ax.plot(t_d, th_d, label="Direct", lw=0.5) - ax.plot(t_c, th_c, "--", label="Chartmap", lw=0.5) + ax.semilogy(t_d, ds, "k-", lw=0.8) + ax.axhline(1e-4, color="r", ls="--", label="tolerance") ax.set_xlabel("time") - ax.set_ylabel("theta_B") - ax.set_title("Poloidal angle") + ax.set_ylabel("|s_direct - s_chartmap|") + ax.set_title("Flux-surface difference") ax.legend() ax = axes[1, 0] - ax.plot(t_d, ph_d, label="Direct", lw=0.5) - ax.plot(t_c, ph_c, "--", label="Chartmap", lw=0.5) + ax.plot(t_d, dtheta, color="C2", lw=0.8) + ax.axhline(0.0, color="0.5", ls="--", lw=0.8) ax.set_xlabel("time") - ax.set_ylabel("phi_B") - ax.set_title("Toroidal angle") - ax.legend() + ax.set_ylabel("theta_B chartmap - direct") + ax.set_title("Poloidal gauge offset") ax = axes[1, 1] - ax.semilogy(t_d, np.abs(s_d - s_c), "k-", lw=0.5) + ax.plot(t_d, dphi, color="C3", lw=0.8) + ax.axhline(0.0, color="0.5", ls="--", lw=0.8) ax.set_xlabel("time") - ax.set_ylabel("|s_direct - s_chartmap|") - ax.set_title("Orbit difference") - ax.axhline(1e-4, color="r", ls="--", label="tolerance") - ax.legend() + ax.set_ylabel("phi_B chartmap - direct") + ax.set_title("Toroidal gauge offset") - fig.suptitle("Boozer Chartmap Orbit Roundtrip", fontsize=14) + fig.suptitle(f"{title}: orbit comparison", fontsize=14) fig.tight_layout() - fig.savefig("/tmp/boozer_orbit_comparison.png", dpi=150) - print("Saved /tmp/boozer_orbit_comparison.png") + output = f"{prefix}_orbit_comparison.png" + fig.savefig(output, dpi=150) + print(f"Saved {output}") if __name__ == "__main__": - plot_field_comparison() - plot_orbit_comparison() + parser = argparse.ArgumentParser() + parser.add_argument( + "--prefix", + default="boozer_chartmap_roundtrip", + help="Artifact prefix used by the Fortran comparison test", + ) + parser.add_argument("--direct-label", default="Direct (VMEC)") + parser.add_argument("--chartmap-label", default="Chartmap") + parser.add_argument("--title", default="Boozer chartmap comparison") + args = parser.parse_args() + + prefix = str(Path(args.prefix)) + plot_field_comparison(prefix, args.direct_label, args.chartmap_label, args.title) + plot_orbit_comparison(prefix, args.direct_label, args.chartmap_label, args.title) diff --git a/test/tests/quasr_json_to_boundary.py b/test/tests/quasr_json_to_boundary.py new file mode 100644 index 00000000..be5ebc6b --- /dev/null +++ b/test/tests/quasr_json_to_boundary.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python3 +"""Convert a QUASR simsopt JSON surface to the GVEC boundary NetCDF format.""" + +import argparse +import json +import math +from pathlib import Path + +import netCDF4 +import numpy as np + + +def find_surface(obj): + for key, value in obj["simsopt_objs"].items(): + if key.startswith("SurfaceXYZTensorFourier"): + return value + raise ValueError("No SurfaceXYZTensorFourier object found") + + +def skip_mode(stellsym, dim, m, n, mpol, ntor): + if not stellsym: + return False + if dim == 0: + return (n <= ntor and m > mpol) or (n > ntor and m <= mpol) + return (n <= ntor and m <= mpol) or (n > ntor and m > mpol) + + +def unpack_coefficients(surface, simsopt_objs): + mpol = int(surface["mpol"]) + ntor = int(surface["ntor"]) + stellsym = bool(surface["stellsym"]) + dofs_ref = surface["dofs"]["value"] + dofs = np.array(simsopt_objs[dofs_ref]["x"]["data"], dtype=float) + + coeffs = [np.zeros((2 * mpol + 1, 2 * ntor + 1)) for _ in range(3)] + counter = 0 + for dim in range(3): + for m in range(2 * mpol + 1): + for n in range(2 * ntor + 1): + if skip_mode(stellsym, dim, m, n, mpol, ntor): + continue + coeffs[dim][m, n] = dofs[counter] + counter += 1 + + if counter != len(dofs): + raise ValueError(f"Consumed {counter} dofs, expected {len(dofs)}") + return coeffs + + +def basis_theta(m, theta, mpol): + if m <= mpol: + return math.cos(m * theta) + return math.sin((m - mpol) * theta) + + +def basis_phi(n, phi, ntor, nfp): + if n <= ntor: + return math.cos(nfp * n * phi) + return math.sin(nfp * (n - ntor) * phi) + + +def boundary_enforcer(enabled, phi, theta, nfp): + if not enabled: + return 1.0 + return math.sin(nfp * phi / 2.0) ** 2 + math.sin(theta / 2.0) ** 2 + + +def evaluate_surface(surface, coeffs, ntheta, nzeta): + mpol = int(surface["mpol"]) + ntor = int(surface["ntor"]) + nfp = int(surface["nfp"]) + clamped_dims = surface["clamped_dims"] + + theta_vals = np.linspace(0.0, 1.0, ntheta, endpoint=False) + phi_vals = np.linspace(0.0, 1.0, nzeta * nfp, endpoint=False) + xyz = np.zeros((nzeta * nfp, ntheta, 3)) + + for iz, phi_norm in enumerate(phi_vals): + phi = 2.0 * math.pi * phi_norm + cosphi = math.cos(phi) + sinphi = math.sin(phi) + phi_basis = [basis_phi(n, phi, ntor, nfp) for n in range(2 * ntor + 1)] + for it, theta_norm in enumerate(theta_vals): + theta = 2.0 * math.pi * theta_norm + theta_basis = [basis_theta(m, theta, mpol) for m in range(2 * mpol + 1)] + xhat = 0.0 + yhat = 0.0 + z = 0.0 + for m in range(2 * mpol + 1): + w = theta_basis[m] + for n in range(2 * ntor + 1): + xhat += coeffs[0][m, n] * w * phi_basis[n] * boundary_enforcer( + clamped_dims[0] and n <= ntor and m <= mpol, phi, theta, nfp + ) + yhat += coeffs[1][m, n] * w * phi_basis[n] * boundary_enforcer( + clamped_dims[1] and n <= ntor and m <= mpol, phi, theta, nfp + ) + z += coeffs[2][m, n] * w * phi_basis[n] * boundary_enforcer( + clamped_dims[2] and n <= ntor and m <= mpol, phi, theta, nfp + ) + xyz[iz, it, 0] = xhat * cosphi - yhat * sinphi + xyz[iz, it, 1] = xhat * sinphi + yhat * cosphi + xyz[iz, it, 2] = z + + return theta_vals, phi_vals, xyz + + +def write_boundary(out_path, theta_vals, phi_vals, xyz, nfp): + with netCDF4.Dataset(out_path, "w") as ds: + ds.createDimension("zeta", xyz.shape[0]) + ds.createDimension("theta", xyz.shape[1]) + ds.createDimension("xyz", 3) + ds.createVariable("pos", "f8", ("zeta", "theta", "xyz"))[:] = xyz + ds.createVariable("nfp", "i4")[:] = np.int32(nfp) + ds.createVariable("theta", "f8", ("theta",))[:] = 2.0 * np.pi * theta_vals + ds.createVariable("zeta", "f8", ("zeta",))[:] = 2.0 * np.pi * phi_vals + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("input_json") + parser.add_argument("output_nc") + parser.add_argument("--ntheta", type=int, default=81) + parser.add_argument("--nzeta", type=int, default=81) + args = parser.parse_args() + + obj = json.loads(Path(args.input_json).read_text()) + surface = find_surface(obj) + coeffs = unpack_coefficients(surface, obj["simsopt_objs"]) + theta_vals, phi_vals, xyz = evaluate_surface( + surface, coeffs, args.ntheta, args.nzeta + ) + write_boundary(args.output_nc, theta_vals, phi_vals, xyz, int(surface["nfp"])) + + print(args.output_nc) + print("shape", xyz.shape) + + +if __name__ == "__main__": + main() diff --git a/test/tests/run_gvec_qa_roundtrip.py b/test/tests/run_gvec_qa_roundtrip.py new file mode 100644 index 00000000..35d5a8b1 --- /dev/null +++ b/test/tests/run_gvec_qa_roundtrip.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 +"""Generate the QA GVEC chartmap on the fly and run the roundtrip check.""" + +from __future__ import annotations + +from pathlib import Path +import shutil + +from boozer_chartmap_artifacts import download_if_missing +from test_e2e_boozer_chartmap import ( + build_gvec_chartmap, + assert_boozer_chartmap_file, + run_roundtrip, +) + + +TEST_DATA = Path(__file__).resolve().parent.parent / "test_data" +QA_WOUT = TEST_DATA / "wout.nc" +QA_WOUT_URL = ( + "https://raw.githubusercontent.com/hiddenSymmetries/simsopt/master/" + "tests/test_files/wout_LandremanPaul2021_QA_reactorScale_lowres_reference.nc" +) + + +def main() -> None: + workdir = Path.cwd() / "boozer_chartmap_gvec_qa" + if workdir.exists(): + shutil.rmtree(workdir) + workdir.mkdir(parents=True) + + wout = download_if_missing(QA_WOUT, QA_WOUT_URL) + chartmap = build_gvec_chartmap( + wout, + workdir, + minimize_tol=1.0e-10, + max_iter=4, + total_iter=8, + ) + assert_boozer_chartmap_file(chartmap) + run_roundtrip( + wout, + chartmap, + "external", + workdir / "gvec", + "QA: VMEC direct vs GVEC chartmap", + "GVEC chartmap", + "2.5e-4", + "1.0e-6", + ) + print("QA GVEC roundtrip PASS") + + +if __name__ == "__main__": + main() diff --git a/test/tests/test_boozer_chartmap_roundtrip.f90 b/test/tests/test_boozer_chartmap_roundtrip.f90 index fb7003d9..fddd4ed8 100644 --- a/test/tests/test_boozer_chartmap_roundtrip.f90 +++ b/test/tests/test_boozer_chartmap_roundtrip.f90 @@ -38,23 +38,67 @@ program test_boozer_chartmap_roundtrip real(dp) :: Bmod, dBmod(3), d2Bmod(6), Br, dBr(3), d2Br(6) real(dp) :: phi_period, z(4), z0(5), vartheta, varphi real(dp) :: rel_err, max_err_bmod, max_err_orbit - integer :: i, ierr, nfail, n_steps_done - character(len=256) :: chartmap_file + integer :: i, ierr, nfail, n_steps_done, u_out + character(len=256) :: wout_file, chartmap_file, mode, arg_value + character(len=512) :: artifact_prefix, field_data_file, orbit_direct_file, & + orbit_chartmap_file type(symplectic_integrator_t) :: si type(field_can_t) :: f + real(dp) :: field_tol, orbit_tol nfail = 0 + wout_file = 'wout.nc' chartmap_file = 'roundtrip_test.nc' + mode = 'export' + field_tol = 1.0e-4_dp + orbit_tol = 1.0e-6_dp + artifact_prefix = 'boozer_chartmap_roundtrip' + + if (command_argument_count() >= 1) then + call get_command_argument(1, wout_file) + if (len_trim(wout_file) == 0) wout_file = 'wout.nc' + end if + if (command_argument_count() >= 2) then + call get_command_argument(2, chartmap_file) + if (len_trim(chartmap_file) == 0) chartmap_file = 'roundtrip_test.nc' + end if + if (command_argument_count() >= 3) then + call get_command_argument(3, mode) + if (len_trim(mode) == 0) mode = 'export' + end if + if (command_argument_count() >= 4) then + call get_command_argument(4, artifact_prefix) + if (len_trim(artifact_prefix) == 0) then + artifact_prefix = 'boozer_chartmap_roundtrip' + end if + end if + if (command_argument_count() >= 5) then + call get_command_argument(5, arg_value) + read (arg_value, *) field_tol + end if + if (command_argument_count() >= 6) then + call get_command_argument(6, arg_value) + read (arg_value, *) orbit_tol + end if + + field_data_file = trim(artifact_prefix) // '_field_comparison.dat' + orbit_direct_file = trim(artifact_prefix) // '_orbit_direct.dat' + orbit_chartmap_file = trim(artifact_prefix) // '_orbit_chartmap.dat' print *, 'Starting roundtrip test...' + print *, ' wout_file=', trim(wout_file) + print *, ' chartmap_file=', trim(chartmap_file) + print *, ' mode=', trim(mode) + print *, ' field_tol=', field_tol + print *, ' orbit_tol=', orbit_tol ! ========================================================= ! Step 1: Initialize VMEC and compute Boozer coordinates ! ========================================================= isw_field_type = 2 rmu = 1.0e8_dp - netcdffile = 'wout.nc' - multharm = 5 + netcdffile = trim(wout_file) + multharm = 7 ns_A = 5 ns_s = 5 ns_tp = 5 @@ -97,13 +141,44 @@ program test_boozer_chartmap_roundtrip ! ========================================================= ! Step 3: Export to Boozer chartmap NetCDF ! ========================================================= - call export_boozer_chartmap(chartmap_file) - print *, '=== Step 3: Exported chartmap ===' + if (trim(mode) == 'export') then + call export_boozer_chartmap(chartmap_file) + print *, '=== Step 3: Exported chartmap ===' + else if (trim(mode) == 'external') then + print *, '=== Step 3: Using external chartmap ===' + else + print *, 'Unknown chartmap comparison mode: ', trim(mode) + error stop 1 + end if + + ! ========================================================= + ! Step 4: Run symplectic orbit (direct path) + ! ========================================================= + + call vmec_to_boozer(0.35_dp, 0.33_dp, 0.97_dp, vartheta, varphi) + + dtau = fper * RT0 / 400.0_dp + + z0(1) = 0.35_dp + z0(2) = vartheta + z0(3) = varphi + z0(4) = 1.0_dp + z0(5) = 0.1_dp + + call run_symplectic_orbit(z0, dtau, n_orbit, orbit_direct, n_steps_done) + + open(newunit=u_out, file=trim(orbit_direct_file), status='replace') + write(u_out, '(a)') '# time s theta phi pphi' + do i = 1, n_steps_done + write(u_out, '(5es18.10)') orbit_direct(i, :) + end do + close(u_out) + + print *, '=== Step 4: Direct orbit done, steps=', n_steps_done, ' ===' ! ========================================================= - ! Step 4: Reimport from chartmap and evaluate fields + ! Step 5: Reimport from chartmap and evaluate fields ! ========================================================= - call reset_boozer_batch_splines call load_boozer_from_chartmap(chartmap_file) do i = 1, n_test @@ -114,14 +189,14 @@ program test_boozer_chartmap_roundtrip Bmod_new(i) = Bmod end do - print *, '=== Step 4: Chartmap field evaluation done ===' + print *, '=== Step 5: Chartmap field evaluation done ===' ! ========================================================= - ! Step 5: Compare fields + ! Step 6: Compare fields ! ========================================================= max_err_bmod = 0.0_dp - open(unit=20, file='/tmp/boozer_field_comparison.dat', status='replace') - write(20, '(a)') '# s theta phi Bmod_ref Bmod_chartmap rel_err' + open(newunit=u_out, file=trim(field_data_file), status='replace') + write(u_out, '(a)') '# s theta phi Bmod_ref Bmod_chartmap rel_err' do i = 1, n_test if (abs(Bmod_ref(i)) > 0.0_dp) then @@ -130,62 +205,36 @@ program test_boozer_chartmap_roundtrip rel_err = abs(Bmod_new(i)) end if max_err_bmod = max(max_err_bmod, rel_err) - write(20, '(6es18.10)') s_test(i), th_test(i), ph_test(i), & + write(u_out, '(6es18.10)') s_test(i), th_test(i), ph_test(i), & Bmod_ref(i), Bmod_new(i), rel_err end do - close(20) + close(u_out) print *, ' max relative error Bmod:', max_err_bmod - if (max_err_bmod > 1.0e-4_dp) then - print *, 'FAIL: Bmod roundtrip error too large:', max_err_bmod - nfail = nfail + 1 + if (field_tol >= 0.0_dp) then + if (max_err_bmod > field_tol) then + print *, 'FAIL: Bmod roundtrip error too large:', max_err_bmod + nfail = nfail + 1 + else + print *, 'PASS: Bmod roundtrip error within tolerance' + end if else - print *, 'PASS: Bmod roundtrip error < 1e-4' + print *, 'INFO: skipping Bmod pass/fail threshold, max err=', max_err_bmod end if - ! ========================================================= - ! Step 6: Run symplectic orbit (direct path) - ! ========================================================= - call reset_boozer_batch_splines - call get_boozer_coordinates - call field_can_from_name("boozer") - - call vmec_to_boozer(0.35_dp, 0.33_dp, 0.97_dp, vartheta, varphi) - - dtau = fper * RT0 / 400.0_dp - - z0(1) = 0.35_dp - z0(2) = vartheta - z0(3) = varphi - z0(4) = 1.0_dp - z0(5) = 0.1_dp - - call run_symplectic_orbit(z0, dtau, n_orbit, orbit_direct, n_steps_done) - - open(unit=21, file='/tmp/orbit_direct.dat', status='replace') - write(21, '(a)') '# time s theta phi pphi' - do i = 1, n_steps_done - write(21, '(5es18.10)') orbit_direct(i, :) - end do - close(21) - - print *, '=== Step 6: Direct orbit done, steps=', n_steps_done, ' ===' - ! ========================================================= ! Step 7: Run symplectic orbit (chartmap path) ! ========================================================= - call reset_boozer_batch_splines - call load_boozer_from_chartmap(chartmap_file) call field_can_from_name("boozer") call run_symplectic_orbit(z0, dtau, n_orbit, orbit_chartmap, n_steps_done) - open(unit=22, file='/tmp/orbit_chartmap.dat', status='replace') - write(22, '(a)') '# time s theta phi pphi' + open(newunit=u_out, file=trim(orbit_chartmap_file), status='replace') + write(u_out, '(a)') '# time s theta phi pphi' do i = 1, n_steps_done - write(22, '(5es18.10)') orbit_chartmap(i, :) + write(u_out, '(5es18.10)') orbit_chartmap(i, :) end do - close(22) + close(u_out) print *, '=== Step 7: Chartmap orbit done, steps=', n_steps_done, ' ===' @@ -199,11 +248,11 @@ program test_boozer_chartmap_roundtrip end do print *, ' max |s_direct - s_chartmap|:', max_err_orbit - if (max_err_orbit > 1.0e-6_dp) then + if (max_err_orbit > orbit_tol) then print *, 'FAIL: orbit roundtrip error too large:', max_err_orbit nfail = nfail + 1 else - print *, 'PASS: orbit roundtrip error < 1e-6' + print *, 'PASS: orbit roundtrip error within tolerance' end if ! ========================================================= @@ -225,11 +274,11 @@ subroutine init_test_points(phi_per) real(dp) :: frac do j = 1, n_test - frac = real(j - 1, dp) / real(n_test - 1, dp) + frac = real(j, dp) / real(n_test + 1, dp) s_test(j) = 0.1_dp + 0.7_dp * frac th_test(j) = twopi * frac - ph_test(j) = phi_per * mod(real(j * 7, dp), real(n_test, dp)) & - / real(n_test, dp) + ph_test(j) = phi_per * mod(real(2 * j + 1, dp), real(n_test + 1, dp)) & + / real(n_test + 1, dp) end do end subroutine init_test_points diff --git a/test/tests/test_e2e_boozer_chartmap.py b/test/tests/test_e2e_boozer_chartmap.py index 1a0f2d8e..1a11e72e 100644 --- a/test/tests/test_e2e_boozer_chartmap.py +++ b/test/tests/test_e2e_boozer_chartmap.py @@ -1,50 +1,89 @@ #!/usr/bin/env python3 -"""End-to-end test: VMEC-Boozer vs Boozer chartmap confined fractions. +"""End-to-end benchmark for VMEC and GVEC Boozer chartmaps on common cases.""" -Two fully isolated simple.x runs with identical particles: -1. VMEC path: simple.x reads wout.nc, computes Boozer internally -2. Chartmap path: simple.x reads boozer_chartmap.nc, no VMEC at all - -Particle initial conditions are generated by the VMEC run, then converted -to Boozer/chartmap reference coordinates for the chartmap run. - -Tests multiple equilibria and longer trace times. -""" +from __future__ import annotations +import json import os -import sys +from pathlib import Path import shutil import subprocess +import sys + +import netCDF4 import numpy as np -SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) -BUILD_DIR = os.path.join(os.path.dirname(os.path.dirname(SCRIPT_DIR)), "build") -SIMPLE_X = os.path.join(BUILD_DIR, "simple.x") -TOOL_X = os.path.join(BUILD_DIR, "test", "tests", "export_boozer_chartmap_tool.x") -TEST_DATA = os.path.join(os.path.dirname(SCRIPT_DIR), "test_data") +from boozer_chartmap_artifacts import ( + confined_metrics, + download_if_missing, + load_confined_fraction, + plot_all_cases_total_losses, + plot_case_loss_comparison, + plot_field_component_comparison, + plot_surface_comparison, + summarize_confined_fraction, +) + + +SCRIPT_DIR = Path(__file__).resolve().parent +BUILD_DIR = SCRIPT_DIR.parent.parent / "build" +SIMPLE_X = BUILD_DIR / "simple.x" +TOOL_X = BUILD_DIR / "test" / "tests" / "export_boozer_chartmap_tool.x" +ROUNDTRIP_X = BUILD_DIR / "test" / "tests" / "test_boozer_chartmap_roundtrip.x" +ROUNDTRIP_PLOT = SCRIPT_DIR / "plot_boozer_chartmap_roundtrip.py" +GVEC_TO_CHARTMAP = SCRIPT_DIR / "gvec_to_boozer_chartmap.py" +TEST_DATA = SCRIPT_DIR.parent / "test_data" + EQUILIBRIA = [ { "name": "QA", - "wout": os.path.join(TEST_DATA, "wout.nc"), + "wout": TEST_DATA / "wout.nc", + "url": "https://raw.githubusercontent.com/hiddenSymmetries/simsopt/master/tests/test_files/wout_LandremanPaul2021_QA_reactorScale_lowres_reference.nc", "trace_time": "1d-2", "facE_al": "1.0d0", + "use_gvec": True, + "vmec_roundtrip_field_tol": "1.0e-4", + "gvec_roundtrip_field_tol": "2.5e-4", + "gvec_minimize_tol": 1.0e-10, + "gvec_max_iter": 4, + "gvec_total_iter": 8, + }, + { + "name": "QH", + "wout": TEST_DATA / "wout_qh.nc", + "url": "https://raw.githubusercontent.com/hiddenSymmetries/simsopt/master/tests/test_files/wout_LandremanPaul2021_QH_reactorScale_lowres_reference.nc", + "trace_time": "1d-2", + "facE_al": "1.0d0", + "use_gvec": False, + "vmec_roundtrip_field_tol": "1.0e-4", + "gvec_minimize_tol": 1.0e-10, + "gvec_max_iter": 12, + "gvec_total_iter": 24, }, { "name": "NCSX", - "wout": os.path.join(TEST_DATA, "wout_ncsx.nc"), + "wout": TEST_DATA / "wout_ncsx.nc", + "url": "https://raw.githubusercontent.com/hiddenSymmetries/simsopt/master/tests/test_files/wout_c09r00_fixedBoundary_0.5T_vacuum_ns201.nc", "trace_time": "1d-2", "facE_al": "100.0d0", + "use_gvec": False, + "vmec_roundtrip_field_tol": "1.0e-4", + "gvec_minimize_tol": 1.0e-10, + "gvec_max_iter": 8, + "gvec_total_iter": 16, }, ] -def simple_in_vmec(wout_name, trace_time, facE_al): +def simple_in_vmec(wout_name: str, trace_time: str, facE_al: str) -> str: return f"""\ &config -multharm = 3 +multharm = 5 contr_pp = -1e10 trace_time = {trace_time} +macrostep_time_grid = 'log' +ntimstep = 61 sbeg = 0.3d0 ntestpart = 32 netcdffile = '{wout_name}' @@ -56,16 +95,18 @@ def simple_in_vmec(wout_name, trace_time, facE_al): """ -def simple_in_chartmap(trace_time, facE_al): +def simple_in_chartmap(chartmap_name: str, trace_time: str, facE_al: str) -> str: return f"""\ &config -multharm = 3 +multharm = 5 contr_pp = -1e10 trace_time = {trace_time} +macrostep_time_grid = 'log' +ntimstep = 61 sbeg = 0.3d0 ntestpart = 32 -field_input = 'boozer_chartmap.nc' -coord_input = 'boozer_chartmap.nc' +field_input = '{chartmap_name}' +coord_input = '{chartmap_name}' isw_field_type = 2 deterministic = .True. integmode = 1 @@ -75,168 +116,364 @@ def simple_in_chartmap(trace_time, facE_al): """ -def run_cmd(cmd, cwd, label, timeout=600): - print(f" [{label}] {' '.join(cmd)}") - result = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True, timeout=timeout) - if result.returncode != 0: - print(f" FAILED: {label}") - print(result.stdout[-2000:] if result.stdout else "") - print(result.stderr[-2000:] if result.stderr else "") - return False - return True - - -def run_single_equilibrium(eq, outdir): - name = eq["name"] - wout = eq["wout"] - trace_time = eq["trace_time"] - facE_al = eq["facE_al"] - - if not os.path.exists(wout): - print(f" Skipping {name}: {wout} not found") - return True - - eq_dir = os.path.join(outdir, name) - vmec_dir = os.path.join(eq_dir, "vmec_run") - chartmap_dir = os.path.join(eq_dir, "chartmap_run") - os.makedirs(vmec_dir) - os.makedirs(chartmap_dir) - - wout_name = os.path.basename(wout) - - # Step 1: VMEC-Boozer - print(f" [{name}] Step 1: VMEC-Boozer path...") - shutil.copy(wout, os.path.join(vmec_dir, wout_name)) - with open(os.path.join(vmec_dir, "simple.in"), "w") as f: - f.write(simple_in_vmec(wout_name, trace_time, facE_al)) - - if not run_cmd([SIMPLE_X], cwd=vmec_dir, label=f"{name} VMEC"): - return False - - cf_vmec = os.path.join(vmec_dir, "confined_fraction.dat") - start_vmec = os.path.join(vmec_dir, "start.dat") - assert os.path.exists(cf_vmec), f"{name}: confined_fraction.dat not produced" - assert os.path.exists(start_vmec), f"{name}: start.dat not produced" - - # Step 2: Export chartmap + convert start.dat - print(f" [{name}] Step 2: Exporting chartmap...") - chartmap_nc = os.path.join(chartmap_dir, "boozer_chartmap.nc") - start_boozer = os.path.join(chartmap_dir, "start.dat") - - if not run_cmd([TOOL_X, os.path.join(vmec_dir, wout_name), chartmap_nc, - start_vmec, start_boozer], - cwd=chartmap_dir, label=f"{name} export"): - return False - - # Step 3: Chartmap-only - print(f" [{name}] Step 3: Chartmap-only path...") - with open(os.path.join(chartmap_dir, "simple.in"), "w") as f: - f.write(simple_in_chartmap(trace_time, facE_al)) - - if not run_cmd([SIMPLE_X], cwd=chartmap_dir, label=f"{name} chartmap"): - return False - - cf_chartmap = os.path.join(chartmap_dir, "confined_fraction.dat") - assert os.path.exists(cf_chartmap), f"{name}: confined_fraction.dat not produced" - - # Step 4: Compare - data_vmec = np.loadtxt(cf_vmec) - data_chartmap = np.loadtxt(cf_chartmap) - - n = min(len(data_vmec), len(data_chartmap)) - data_vmec = data_vmec[:n] - data_chartmap = data_chartmap[:n] - - t = data_vmec[:, 0] - pass_v, trap_v = data_vmec[:, 1], data_vmec[:, 2] - pass_c, trap_c = data_chartmap[:, 1], data_chartmap[:, 2] - total_v = pass_v + trap_v - total_c = pass_c + trap_c - - max_diff_total = np.max(np.abs(total_v - total_c)) - max_diff_pass = np.max(np.abs(pass_v - pass_c)) - max_diff_trap = np.max(np.abs(trap_v - trap_c)) - - print(f" [{name}] max |total diff| = {max_diff_total:.6e}") - print(f" [{name}] max |pass diff| = {max_diff_pass:.6e}") - print(f" [{name}] max |trap diff| = {max_diff_trap:.6e}") - - # Plot - plot_comparison(eq_dir, name, t, total_v, total_c, pass_v, pass_c, trap_v, trap_c) - - npart = data_vmec[0, 3] - tol = 1.0 / npart - if max_diff_total > tol: - print(f" [{name}] FAIL: total confined differs by {max_diff_total:.6e} > {tol:.6e}") - return False - - print(f" [{name}] PASS") - return True - - -def plot_comparison(outdir, name, t, total_v, total_c, pass_v, pass_c, trap_v, trap_c): - import matplotlib - matplotlib.use("Agg") - import matplotlib.pyplot as plt - - fig, axes = plt.subplots(1, 3, figsize=(15, 4.5)) - - ax = axes[0] - ax.plot(t * 1e3, total_v, "b-", lw=1.5, label="VMEC-Boozer") - ax.plot(t * 1e3, total_c, "r--", lw=1.5, label="Chartmap") - ax.set_xlabel("time [ms]") - ax.set_ylabel("confined fraction") - ax.set_title("Total confined") - ax.legend() - ax.set_ylim(-0.05, 1.05) - - ax = axes[1] - ax.plot(t * 1e3, pass_v, "b-", lw=1.5, label="VMEC passing") - ax.plot(t * 1e3, pass_c, "r--", lw=1.5, label="Chartmap passing") - ax.plot(t * 1e3, trap_v, "b-", lw=1, alpha=0.5, label="VMEC trapped") - ax.plot(t * 1e3, trap_c, "r--", lw=1, alpha=0.5, label="Chartmap trapped") - ax.set_xlabel("time [ms]") - ax.set_ylabel("confined fraction") - ax.set_title("Passing + Trapped") - ax.legend(fontsize=8) - ax.set_ylim(-0.05, 1.05) - - ax = axes[2] - ax.semilogy(t * 1e3, np.abs(total_v - total_c) + 1e-16, "k-", lw=1) - ax.set_xlabel("time [ms]") - ax.set_ylabel("|difference|") - ax.set_title("Confined fraction difference") - - fig.suptitle(f"E2E: VMEC-Boozer vs Chartmap ({name})", fontsize=14) - fig.tight_layout() - outpath = os.path.join(outdir, f"e2e_{name}.png") - fig.savefig(outpath, dpi=150) - print(f" Saved {outpath}") - - -def main(): - for path, label in [(SIMPLE_X, "simple.x"), (TOOL_X, "export tool")]: - if not os.path.exists(path): - print(f"Missing {label}: {path}") - sys.exit(1) - - outdir = "/tmp/boozer_chartmap_e2e" - if os.path.exists(outdir): +def run_cmd(cmd: list[str], cwd: Path, label: str, timeout: int = 3600) -> None: + print(f"[{label}] {' '.join(cmd)}") + result = subprocess.run( + cmd, + cwd=cwd, + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + if result.returncode == 0: + return + print(result.stdout[-4000:]) + print(result.stderr[-4000:]) + raise RuntimeError(f"{label} failed with exit code {result.returncode}") + + +def build_gvec_chartmap( + wout: Path, + run_dir: Path, + minimize_tol: float, + max_iter: int, + total_iter: int, +) -> Path: + import gvec + from gvec.scripts.convert_wout import convert_vmec_wout + from gvec.util import read_parameters + + convert_dir = run_dir / "convert" + solve_dir = run_dir / "solve" + convert_dir.mkdir(parents=True, exist_ok=True) + solve_dir.mkdir(parents=True, exist_ok=True) + convert_vmec_wout(wout, convert_dir) + + restart_state = next(convert_dir.glob("*State*.dat")) + params = read_parameters(convert_dir / "parameter.toml") + params["minimize_tol"] = float(minimize_tol) + params["maxIter"] = int(max_iter) + params["totalIter"] = int(total_iter) + gvec.run( + params, + restartstate=restart_state, + runpath=solve_dir, + redirect_gvec_stdout=True, + quiet=True, + ) + + param_final = sorted(solve_dir.glob("parameter*_final.ini"))[-1] + state_final = sorted(solve_dir.glob("*State_final.dat"))[-1] + output = run_dir / "boozer_chartmap.nc" + run_cmd( + [ + sys.executable, + str(GVEC_TO_CHARTMAP), + str(param_final), + str(state_final), + str(output), + "--nrho", + "50", + "--ntheta", + "36", + "--nphi", + "81", + ], + cwd=run_dir, + label=f"{wout.stem} GVEC->chartmap", + timeout=3600, + ) + return output + + +def run_roundtrip( + wout: Path, + chartmap: Path, + mode: str, + prefix: Path, + title: str, + chart_label: str, + field_tol: str, + orbit_tol: str, +) -> None: + prefix.parent.mkdir(parents=True, exist_ok=True) + run_cmd( + [ + str(ROUNDTRIP_X), + str(wout), + str(chartmap), + mode, + str(prefix), + field_tol, + orbit_tol, + ], + cwd=prefix.parent, + label=title, + timeout=3600, + ) + run_cmd( + [ + sys.executable, + str(ROUNDTRIP_PLOT), + "--prefix", + str(prefix), + "--direct-label", + "VMEC-Boozer", + "--chartmap-label", + chart_label, + "--title", + title, + ], + cwd=prefix.parent, + label=f"{title} plots", + ) + + +def write_json(path: Path, payload: dict) -> None: + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def assert_boozer_chartmap_file(path: Path) -> None: + with netCDF4.Dataset(path) as ds: + if int(getattr(ds, "boozer_field")) != 1: + raise RuntimeError(f"{path}: boozer_field attribute missing or invalid") + if str(getattr(ds, "zeta_convention")) != "boozer": + raise RuntimeError(f"{path}: expected zeta_convention='boozer'") + if str(getattr(ds, "rho_convention")) != "rho_tor": + raise RuntimeError(f"{path}: expected rho_convention='rho_tor'") + + +def run_single_equilibrium(eq: dict[str, str | Path], outdir: Path) -> tuple[bool, dict[str, dict[str, float]]]: + name = str(eq["name"]) + wout = download_if_missing(Path(eq["wout"]), str(eq["url"])) + trace_time = str(eq["trace_time"]) + facE_al = str(eq["facE_al"]) + use_gvec = bool(eq.get("use_gvec", False)) + vmec_roundtrip_field_tol = str(eq.get("vmec_roundtrip_field_tol", "1.0e-4")) + gvec_roundtrip_field_tol = str(eq.get("gvec_roundtrip_field_tol", "2.5e-4")) + gvec_minimize_tol = float(eq.get("gvec_minimize_tol", 1.0e-10)) + gvec_max_iter = int(eq.get("gvec_max_iter", 4)) + gvec_total_iter = int(eq.get("gvec_total_iter", 8)) + + eq_dir = outdir / name + if eq_dir.exists(): + shutil.rmtree(eq_dir) + vmec_dir = eq_dir / "vmec_run" + chartmap_dir = eq_dir / "vmec_chartmap_run" + gvec_dir = eq_dir / "gvec_chartmap_run" + vmec_dir.mkdir(parents=True) + chartmap_dir.mkdir() + gvec_dir.mkdir() + + wout_name = wout.name + shutil.copy2(wout, vmec_dir / wout_name) + (vmec_dir / "simple.in").write_text(simple_in_vmec(wout_name, trace_time, facE_al), encoding="utf-8") + + print(f"[{name}] Step 1: VMEC Boozer run") + run_cmd([str(SIMPLE_X)], cwd=vmec_dir, label=f"{name} VMEC") + cf_vmec = load_confined_fraction(vmec_dir / "confined_fraction.dat") + start_vmec = vmec_dir / "start.dat" + + print(f"[{name}] Step 2: export VMEC Boozer chartmap") + chartmap_nc = chartmap_dir / "boozer_chartmap.nc" + start_boozer = chartmap_dir / "start.dat" + run_cmd( + [ + str(TOOL_X), + str(vmec_dir / wout_name), + str(chartmap_nc), + str(start_vmec), + str(start_boozer), + ], + cwd=chartmap_dir, + label=f"{name} export", + ) + assert_boozer_chartmap_file(chartmap_nc) + + print(f"[{name}] Step 3: VMEC chartmap SIMPLE run") + (chartmap_dir / "simple.in").write_text( + simple_in_chartmap("boozer_chartmap.nc", trace_time, facE_al), + encoding="utf-8", + ) + run_cmd([str(SIMPLE_X)], cwd=chartmap_dir, label=f"{name} VMEC chartmap") + cf_chartmap = load_confined_fraction(chartmap_dir / "confined_fraction.dat") + + gvec_chartmap = None + cf_gvec = None + if use_gvec: + print(f"[{name}] Step 4: build GVEC chartmap from common VMEC input") + gvec_chartmap = build_gvec_chartmap( + wout, + gvec_dir, + minimize_tol=gvec_minimize_tol, + max_iter=gvec_max_iter, + total_iter=gvec_total_iter, + ) + assert_boozer_chartmap_file(gvec_chartmap) + shutil.copy2(start_boozer, gvec_dir / "start.dat") + (gvec_dir / "simple.in").write_text( + simple_in_chartmap("boozer_chartmap.nc", trace_time, facE_al), + encoding="utf-8", + ) + if list(gvec_dir.glob("wout*")): + raise RuntimeError(f"{name}: GVEC run directory unexpectedly contains VMEC files") + + print(f"[{name}] Step 5: GVEC chartmap SIMPLE run") + run_cmd([str(SIMPLE_X)], cwd=gvec_dir, label=f"{name} GVEC chartmap") + cf_gvec = load_confined_fraction(gvec_dir / "confined_fraction.dat") + + print(f"[{name}] Step 6: field/orbit comparison artifacts") + run_roundtrip( + wout, + chartmap_nc, + "export", + eq_dir / "roundtrip_vmec" / "vmec", + f"{name}: VMEC direct vs VMEC chartmap", + "VMEC chartmap", + vmec_roundtrip_field_tol, + "1.0e-6", + ) + field_metrics: dict[str, float] = {} + if use_gvec: + run_roundtrip( + wout, + gvec_chartmap, + "external", + eq_dir / "roundtrip_gvec" / "gvec", + f"{name}: VMEC direct vs GVEC chartmap", + "GVEC chartmap", + gvec_roundtrip_field_tol, + "1.0e-6", + ) + + field_metrics = plot_field_component_comparison( + eq_dir / f"{name.lower()}_field_components.png", + name, + chartmap_nc, + gvec_chartmap, + "VMEC chartmap", + "GVEC chartmap", + ) + plot_surface_comparison( + eq_dir / f"{name.lower()}_surface_comparison.png", + name, + chartmap_nc, + gvec_chartmap, + "VMEC chartmap", + "GVEC chartmap", + ) + + series = { + "VMEC-Boozer": cf_vmec, + "VMEC chartmap": cf_chartmap, + } + if use_gvec and cf_gvec is not None: + series["GVEC chartmap"] = cf_gvec + min_len = min(len(values) for values in series.values()) + metrics_by_label: dict[str, dict[str, np.ndarray | float]] = {} + max_diffs = {"total": 0.0, "pass": 0.0, "trap": 0.0} + ref_metrics = confined_metrics(series["VMEC-Boozer"][:min_len]) + tol = 1.0 / float(ref_metrics["npart"]) + for label, values in series.items(): + metrics_by_label[label] = confined_metrics(values[:min_len]) + + labels = list(metrics_by_label) + for idx, left in enumerate(labels): + for right in labels[idx + 1:]: + left_metrics = metrics_by_label[left] + right_metrics = metrics_by_label[right] + diff_total = float( + np.max(np.abs(np.asarray(left_metrics["total_conf"]) - np.asarray(right_metrics["total_conf"]))) + ) + diff_pass = float( + np.max(np.abs(np.asarray(left_metrics["pass_conf"]) - np.asarray(right_metrics["pass_conf"]))) + ) + diff_trap = float( + np.max(np.abs(np.asarray(left_metrics["trap_conf"]) - np.asarray(right_metrics["trap_conf"]))) + ) + print(f"[{name}] max |total diff| {left} vs {right} = {diff_total:.6e}") + print(f"[{name}] max |pass diff| {left} vs {right} = {diff_pass:.6e}") + print(f"[{name}] max |trap diff| {left} vs {right} = {diff_trap:.6e}") + max_diffs["total"] = max(max_diffs["total"], diff_total) + max_diffs["pass"] = max(max_diffs["pass"], diff_pass) + max_diffs["trap"] = max(max_diffs["trap"], diff_trap) + + plot_case_loss_comparison(eq_dir / f"e2e_{name}.png", name, metrics_by_label) + + summary = { + label: summarize_confined_fraction(series[label][:min_len]) + for label in series + } + summary["executed_modes"] = list(series) + summary["comparison"] = { + "tol": tol, + "max_total_diff": max_diffs["total"], + "max_pass_diff": max_diffs["pass"], + "max_trap_diff": max_diffs["trap"], + } + summary["field_metrics"] = field_metrics + write_json(eq_dir / "loss_summary.json", summary) + + ok = all(value <= tol for value in max_diffs.values()) + if not ok: + print(f"[{name}] FAIL: confined-fraction mismatch exceeds {tol:.6e}") + else: + print(f"[{name}] PASS") + return ok, summary + + +def main() -> None: + for path, label in [ + (SIMPLE_X, "simple.x"), + (TOOL_X, "export tool"), + (ROUNDTRIP_X, "roundtrip comparison tool"), + (GVEC_TO_CHARTMAP, "GVEC conversion script"), + ]: + if not path.exists(): + raise SystemExit(f"Missing {label}: {path}") + + outdir = Path.cwd() / "boozer_chartmap_e2e" + if outdir.exists(): shutil.rmtree(outdir) - os.makedirs(outdir) - - all_pass = True - for eq in EQUILIBRIA: - print(f"\n=== {eq['name']} (trace_time={eq['trace_time']}) ===") - if not run_single_equilibrium(eq, outdir): - all_pass = False + outdir.mkdir(parents=True) - print() - if all_pass: - print("ALL EQUILIBRIA PASSED") + case_filter = os.environ.get("SIMPLE_BOOZER_CASES", "").strip() + if case_filter: + selected = {item.strip().upper() for item in case_filter.split(",") if item.strip()} + equilibria = [eq for eq in EQUILIBRIA if str(eq["name"]).upper() in selected] else: - print("SOME EQUILIBRIA FAILED") - sys.exit(1) + equilibria = EQUILIBRIA + + all_ok = True + combined_metrics: dict[str, dict[str, dict[str, np.ndarray | float]]] = {} + summaries: dict[str, dict[str, dict[str, float]]] = {} + for eq in equilibria: + print(f"\n=== {eq['name']} ===") + ok, summary = run_single_equilibrium(eq, outdir) + series = {} + eq_dir = outdir / str(eq["name"]) + for label, run_name in [ + ("VMEC-Boozer", "vmec_run"), + ("VMEC chartmap", "vmec_chartmap_run"), + ]: + series[label] = confined_metrics( + load_confined_fraction(eq_dir / run_name / "confined_fraction.dat") + ) + if bool(eq.get("use_gvec", False)): + series["GVEC chartmap"] = confined_metrics( + load_confined_fraction(eq_dir / "gvec_chartmap_run" / "confined_fraction.dat") + ) + combined_metrics[str(eq["name"])] = series + summaries[str(eq["name"])] = summary + all_ok = all_ok and ok + + plot_all_cases_total_losses(outdir / "all_e2e_losses.png", combined_metrics) + write_json(outdir / "common_case_summaries.json", summaries) + + if not all_ok: + raise SystemExit("One or more common-case end-to-end comparisons failed") + print("ALL COMMON-CASE BOOZER CHARTMAP BENCHMARKS PASSED") if __name__ == "__main__": diff --git a/test/tests/test_figure8_boozer_chartmap.py b/test/tests/test_figure8_boozer_chartmap.py new file mode 100644 index 00000000..ac4a32dd --- /dev/null +++ b/test/tests/test_figure8_boozer_chartmap.py @@ -0,0 +1,421 @@ +#!/usr/bin/env python3 +"""Figure-8 QUASR -> GVEC -> Boozer chartmap end-to-end benchmark.""" + +from __future__ import annotations + +import os +from pathlib import Path +import shutil +import subprocess +import sys + +import netCDF4 +import numpy as np + +from boozer_chartmap_artifacts import ( + confined_metrics, + download_if_missing, + load_chartmap_fields, + load_chartmap_surface, + load_confined_fraction, + plot_all_cases_total_losses, + plot_case_loss_comparison, + tile_field_periods, +) + + +SCRIPT_DIR = Path(__file__).resolve().parent +BUILD_DIR = SCRIPT_DIR.parent.parent / "build" +SIMPLE_X = BUILD_DIR / "simple.x" +GVEC_TO_CHARTMAP = SCRIPT_DIR / "gvec_to_boozer_chartmap.py" +QUASR_TO_BOUNDARY = SCRIPT_DIR / "quasr_json_to_boundary.py" +REFERENCE_SIGNATURE = SCRIPT_DIR / "figure8_signature_reference.txt" +QUASR_JSON_URL = ( + "https://raw.githubusercontent.com/gvec-group/gvec/main/" + "test-CI/data/quasr-0112714.json" +) + + +def write_inputs(run_dir: Path) -> None: + starts = [] + for rho in (0.35, 0.55, 0.78): + for theta in (0.0, 2.1, 4.2, 5.4): + starts.append((rho, theta, 0.0, 1.0, 0.35)) + starts.append((rho, theta, 0.2, 1.0, -0.25)) + with (run_dir / "start.dat").open("w", encoding="utf-8") as handle: + for row in starts: + handle.write(" ".join(f"{value:.12g}" for value in row) + "\n") + + with (run_dir / "simple.in").open("w", encoding="utf-8") as handle: + handle.write( + """&config +multharm = 5 +contr_pp = -1e10 +trace_time = 5d-4 +macrostep_time_grid = 'log' +ntimstep = 61 +ntestpart = 24 +field_input = 'figure8.gvec.chartmap.nc' +coord_input = 'figure8.gvec.chartmap.nc' +isw_field_type = 2 +deterministic = .True. +integmode = 1 +startmode = 2 +facE_al = 1.0d0 +/ +""" + ) + + +def run_cmd(cmd: list[str], cwd: Path, label: str, timeout: int = 3600) -> None: + print(f"[{label}] {' '.join(cmd)}") + result = subprocess.run( + cmd, + cwd=cwd, + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + if result.returncode == 0: + return + print(result.stdout[-4000:]) + print(result.stderr[-4000:]) + raise RuntimeError(f"{label} failed with exit code {result.returncode}") + + +def pushd(path: Path): + class _Pushd: + def __enter__(self_inner): + self_inner.old = Path.cwd() + os.chdir(path) + return self_inner + + def __exit__(self_inner, exc_type, exc, tb): + os.chdir(self_inner.old) + return False + + return _Pushd() + + +def build_figure8_chartmap(run_dir: Path) -> tuple[Path, Path]: + import gvec + from gvec.scripts.quasr import load_quasr + from gvec.util import read_parameters + + prepare_dir = run_dir / "prepare" + solve_dir = run_dir / "solve" + prepare_dir.mkdir(parents=True, exist_ok=True) + solve_dir.mkdir(parents=True, exist_ok=True) + source_json = download_if_missing(run_dir / "quasr-0112714.json", QUASR_JSON_URL) + boundary = run_dir / "figure8.quasr.boundary.nc" + run_cmd( + [ + sys.executable, + str(QUASR_TO_BOUNDARY), + str(source_json), + str(boundary), + "--ntheta", + "81", + "--nzeta", + "81", + ], + cwd=run_dir, + label="figure8 QUASR->boundary", + ) + with pushd(prepare_dir): + load_quasr(boundary, filetype="netcdf", quiet=True, name="figure8") + params = read_parameters(prepare_dir / "figure8-parameters.toml") + params["hmap_ncfile"] = str((prepare_dir / params["hmap_ncfile"]).resolve()) + params["minimize_tol"] = 1.0e-8 + params["maxIter"] = 5 + params["totalIter"] = 10 + gvec.run( + params, + runpath=solve_dir, + redirect_gvec_stdout=True, + quiet=True, + ) + + param_final = sorted(solve_dir.glob("parameter*_final.ini"))[-1] + state_final = sorted(solve_dir.glob("*State_final.dat"))[-1] + output = run_dir / "figure8.gvec.chartmap.nc" + run_cmd( + [ + sys.executable, + str(GVEC_TO_CHARTMAP), + str(param_final), + str(state_final), + str(output), + "--nrho", + "50", + "--ntheta", + "36", + "--nphi", + "81", + ], + cwd=run_dir, + label="figure8 GVEC->chartmap", + timeout=3600, + ) + return boundary, output + + +def load_quasr_boundary(path: Path) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + with netCDF4.Dataset(path) as ds: + pos = np.asarray(ds["pos"][:], dtype=float) + return pos[:, :, 0], pos[:, :, 1], pos[:, :, 2] + + +def figure8_signature(path: Path) -> np.ndarray: + surface = load_chartmap_surface(path) + fields = load_chartmap_fields(path) + bx, by, bz = tile_field_periods( + surface["boundary_x"], surface["boundary_y"], surface["boundary_z"], int(surface["nfp"]) + ) + ax, ay, az = tile_field_periods( + surface["axis_x"], surface["axis_y"], surface["axis_z"], int(surface["nfp"]) + ) + axis_x = ax.mean(axis=0) + axis_y = ay.mean(axis=0) + axis_z = az.mean(axis=0) + boundary_mid = bx[bx.shape[0] // 2] + boundary_top = bz[bz.shape[0] // 4] + idx_axis = np.linspace(0, axis_x.size - 1, 16, dtype=int) + idx_bound = np.linspace(0, boundary_mid.size - 1, 16, dtype=int) + rho_idx = np.linspace(0, len(fields["rho"]) - 1, 10, dtype=int) + bmod = np.asarray(fields["Bmod"], dtype=float) + ir = bmod.shape[0] // 2 + it = np.linspace(0, bmod.shape[1] - 1, 6, dtype=int) + ip = np.linspace(0, bmod.shape[2] - 1, 6, dtype=int) + sample_bmod = np.array([bmod[ir, i, j] for i in it for j in ip], dtype=float) + return np.concatenate( + [ + axis_x[idx_axis], + axis_y[idx_axis], + axis_z[idx_axis], + boundary_mid[idx_bound], + boundary_top[idx_bound], + np.asarray(fields["A_phi"])[rho_idx], + np.asarray(fields["B_theta"])[rho_idx], + np.asarray(fields["B_phi"])[rho_idx], + sample_bmod, + ] + ) + + +def assert_figure8_golden(chartmap_path: Path, out_dir: Path) -> None: + generated = figure8_signature(chartmap_path) + reference = np.loadtxt(REFERENCE_SIGNATURE) + np.savetxt(out_dir / "figure8_signature_generated.dat", generated) + np.savetxt(out_dir / "figure8_signature_reference.dat", reference) + if not np.allclose(generated, reference, rtol=1.0e-6, atol=1.0e-8): + diff = np.abs(generated - reference) + idx = int(np.argmax(diff)) + raise RuntimeError( + "Figure-8 golden record mismatch at index " + f"{idx}: generated={generated[idx]:.16e}, reference={reference[idx]:.16e}, " + f"abs_diff={diff[idx]:.3e}" + ) + + +def plot_review(run_dir: Path, boundary_path: Path, chartmap_path: Path) -> None: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + xb_true, yb_true, zb_true = load_quasr_boundary(boundary_path) + surface = load_chartmap_surface(chartmap_path) + fields = load_chartmap_fields(chartmap_path) + xb, yb, zb = tile_field_periods( + surface["boundary_x"], surface["boundary_y"], surface["boundary_z"], int(surface["nfp"]) + ) + xa, ya, za = tile_field_periods( + surface["axis_x"], surface["axis_y"], surface["axis_z"], int(surface["nfp"]) + ) + axis_true = np.stack( + [xb_true.mean(axis=0), yb_true.mean(axis=0), zb_true.mean(axis=0)], + axis=1, + ) + axis_chart = np.stack([xa.mean(axis=0), ya.mean(axis=0), za.mean(axis=0)], axis=1) + + def close_curve(curve: np.ndarray) -> np.ndarray: + curve = np.asarray(curve) + return np.concatenate([curve, curve[:1]]) + + def close_surface(xsurf: np.ndarray, ysurf: np.ndarray, zsurf: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + return ( + np.concatenate([xsurf, xsurf[:, :1]], axis=1), + np.concatenate([ysurf, ysurf[:, :1]], axis=1), + np.concatenate([zsurf, zsurf[:, :1]], axis=1), + ) + + def sample_indices(size: int, stride: int) -> list[int]: + indices = list(range(0, size, stride)) + last = size - 1 + if last not in indices: + indices.append(last) + return indices + + xb_true_plot, yb_true_plot, zb_true_plot = close_surface(xb_true, yb_true, zb_true) + xb_plot, yb_plot, zb_plot = close_surface(xb, yb, zb) + axis_true_plot = np.column_stack( + [close_curve(axis_true[:, 0]), close_curve(axis_true[:, 1]), close_curve(axis_true[:, 2])] + ) + axis_chart_plot = np.column_stack( + [close_curve(axis_chart[:, 0]), close_curve(axis_chart[:, 1]), close_curve(axis_chart[:, 2])] + ) + + def draw_projection(ax, xsurf, ysurf, xlabel, ylabel, title, color): + for j in sample_indices(xsurf.shape[1], max(1, xsurf.shape[1] // 12)): + ax.plot(xsurf[:, j], ysurf[:, j], color=color, lw=0.8, alpha=0.55) + for i in sample_indices(xsurf.shape[0], max(1, xsurf.shape[0] // 12)): + ax.plot(xsurf[i], ysurf[i], color=color, lw=0.6, alpha=0.3) + ax.set_aspect("equal", adjustable="box") + ax.set_xlabel(xlabel) + ax.set_ylabel(ylabel) + ax.set_title(title) + + def draw_surface_3d(ax, xsurf, ysurf, zsurf, color, label, axis_xyz, linestyle="-"): + for j in sample_indices(xsurf.shape[1], max(1, xsurf.shape[1] // 12)): + ax.plot(xsurf[:, j], ysurf[:, j], zsurf[:, j], color=color, lw=0.7, alpha=0.45, ls=linestyle) + for i in sample_indices(xsurf.shape[0], max(1, xsurf.shape[0] // 12)): + ax.plot(xsurf[i], ysurf[i], zsurf[i], color=color, lw=0.5, alpha=0.22, ls=linestyle) + ax.plot(axis_xyz[:, 0], axis_xyz[:, 1], axis_xyz[:, 2], color=color, lw=1.8, ls=linestyle, label=label) + + def set_equal_3d(ax, x, y, z): + mins = np.array([np.min(x), np.min(y), np.min(z)], dtype=float) + maxs = np.array([np.max(x), np.max(y), np.max(z)], dtype=float) + center = 0.5 * (mins + maxs) + radius = 0.55 * np.max(maxs - mins) + if radius <= 0.0: + radius = 1.0 + ax.set_xlim(center[0] - radius, center[0] + radius) + ax.set_ylim(center[1] - radius, center[1] + radius) + ax.set_zlim(center[2] - radius, center[2] + radius) + ax.set_box_aspect((1.0, 1.0, 1.0)) + + fig = plt.figure(figsize=(16, 12), constrained_layout=True) + axes = np.empty((3, 3), dtype=object) + for i in range(3): + axes[0, i] = fig.add_subplot(3, 3, i + 1) + for i in range(3): + axes[1, i] = fig.add_subplot(3, 3, i + 4, projection="3d") + for i in range(3): + axes[2, i] = fig.add_subplot(3, 3, i + 7, projection="3d" if i < 2 else None) + + draw_projection(axes[0, 0], xb_true_plot, zb_true_plot, "x [m]", "z [m]", "QUASR boundary x-z", "C0") + draw_projection(axes[0, 1], yb_true_plot, zb_true_plot, "y [m]", "z [m]", "QUASR boundary y-z", "C0") + draw_projection(axes[0, 2], xb_true_plot, yb_true_plot, "x [m]", "y [m]", "QUASR boundary x-y", "C0") + axes[0, 0].plot(axis_true_plot[:, 0], axis_true_plot[:, 2], color="black", lw=1.6, label="QUASR axis") + axes[0, 0].plot(axis_chart_plot[:, 0], axis_chart_plot[:, 2], color="goldenrod", lw=1.2, ls="--", label="chartmap axis") + axes[0, 0].legend(fontsize=8) + axes[0, 1].plot(axis_true_plot[:, 1], axis_true_plot[:, 2], color="black", lw=1.6) + axes[0, 1].plot(axis_chart_plot[:, 1], axis_chart_plot[:, 2], color="goldenrod", lw=1.2, ls="--") + axes[0, 2].plot(axis_true_plot[:, 0], axis_true_plot[:, 1], color="black", lw=1.6) + axes[0, 2].plot(axis_chart_plot[:, 0], axis_chart_plot[:, 1], color="goldenrod", lw=1.2, ls="--") + + xb_all = np.concatenate([xb_true_plot.ravel(), xb_plot.ravel()]) + yb_all = np.concatenate([yb_true_plot.ravel(), yb_plot.ravel()]) + zb_all = np.concatenate([zb_true_plot.ravel(), zb_plot.ravel()]) + views = [(22.0, -62.0), (18.0, 24.0), (78.0, -90.0)] + for ax, (elev, azim) in zip(axes[1], views): + draw_surface_3d(ax, xb_true_plot, yb_true_plot, zb_true_plot, "C0", "QUASR boundary", axis_true_plot) + draw_surface_3d(ax, xb_plot, yb_plot, zb_plot, "C3", "chartmap surface", axis_chart_plot, linestyle="--") + ax.view_init(elev=elev, azim=azim) + set_equal_3d(ax, xb_all, yb_all, zb_all) + ax.set_xlabel("x [m]") + ax.set_ylabel("y [m]") + ax.set_zlabel("z [m]") + ax.set_title(f"Overlay 3D elev={elev:.0f}, azim={azim:.0f}") + axes[1, 0].legend(fontsize=8) + + draw_surface_3d(axes[2, 0], xb_true_plot, yb_true_plot, zb_true_plot, "C0", "QUASR boundary", axis_true_plot) + axes[2, 0].view_init(elev=28.0, azim=-45.0) + set_equal_3d(axes[2, 0], xb_true_plot, yb_true_plot, zb_true_plot) + axes[2, 0].set_title("QUASR boundary only") + axes[2, 0].set_xlabel("x [m]") + axes[2, 0].set_ylabel("y [m]") + axes[2, 0].set_zlabel("z [m]") + + draw_surface_3d(axes[2, 1], xb_plot, yb_plot, zb_plot, "C3", "chartmap surface", axis_chart_plot) + axes[2, 1].view_init(elev=28.0, azim=-45.0) + set_equal_3d(axes[2, 1], xb_plot, yb_plot, zb_plot) + axes[2, 1].set_title("Chartmap surface only") + axes[2, 1].set_xlabel("x [m]") + axes[2, 1].set_ylabel("y [m]") + axes[2, 1].set_zlabel("z [m]") + + bmod = np.asarray(fields["Bmod"], dtype=float) + mid = bmod.shape[0] // 2 + im = axes[2, 2].imshow(bmod[mid].T, origin="lower", aspect="auto", cmap="magma") + axes[2, 2].set_title("|B| at mid rho") + axes[2, 2].set_xlabel("zeta index") + axes[2, 2].set_ylabel("theta index") + fig.colorbar(im, ax=axes[2, 2], shrink=0.8) + + outpath = run_dir / "figure8_review.png" + fig.savefig(outpath, dpi=180) + print(f"Saved {outpath}") + + +def maybe_plot_all_cases(run_dir: Path, figure8_metrics: dict[str, dict[str, np.ndarray | float]]) -> None: + common_dir = Path.cwd() / "boozer_chartmap_e2e" + if not common_dir.exists(): + return + cases = {} + for case_dir in sorted(common_dir.iterdir()): + if not case_dir.is_dir(): + continue + series = {} + for label, subdir in [ + ("VMEC-Boozer", "vmec_run"), + ("VMEC chartmap", "vmec_chartmap_run"), + ("GVEC chartmap", "gvec_chartmap_run"), + ]: + data_file = case_dir / subdir / "confined_fraction.dat" + if data_file.exists(): + series[label] = confined_metrics(load_confined_fraction(data_file)) + if series: + cases[case_dir.name] = series + cases["Figure-8"] = figure8_metrics + plot_all_cases_total_losses(run_dir / "all_cases_losses.png", cases) + + +def main() -> None: + for path, label in [ + (SIMPLE_X, "simple.x"), + (GVEC_TO_CHARTMAP, "GVEC conversion script"), + (QUASR_TO_BOUNDARY, "QUASR boundary conversion script"), + (REFERENCE_SIGNATURE, "figure-8 golden signature"), + ]: + if not path.exists(): + raise SystemExit(f"Missing {label}: {path}") + + run_dir = Path.cwd() / "figure8_chartmap_smoke" + if run_dir.exists(): + shutil.rmtree(run_dir) + run_dir.mkdir(parents=True) + + boundary, chartmap = build_figure8_chartmap(run_dir) + write_inputs(run_dir) + run_cmd([str(SIMPLE_X)], cwd=run_dir, label="figure8 SIMPLE") + + data = load_confined_fraction(run_dir / "confined_fraction.dat") + confined = np.asarray(data[:, 1:3], dtype=float) + if not np.all(np.isfinite(confined)): + raise RuntimeError("Non-finite confined fractions in figure-8 benchmark") + if np.any(confined < -1.0e-12) or np.any(confined > 1.0 + 1.0e-12): + raise RuntimeError("Confined fractions outside [0, 1] in figure-8 benchmark") + + metrics = {"Figure-8 GVEC chartmap": confined_metrics(data)} + plot_case_loss_comparison(run_dir / "figure8_losses.png", "Figure-8", metrics) + plot_review(run_dir, boundary, chartmap) + assert_figure8_golden(chartmap, run_dir) + maybe_plot_all_cases(run_dir, metrics) + print("FIGURE-8 GVEC CHARTMAP PASS") + + +if __name__ == "__main__": + main()