diff --git a/.gitignore b/.gitignore index 2e60d25b..e2e68911 100644 --- a/.gitignore +++ b/.gitignore @@ -68,3 +68,5 @@ thirdparty/pyplot_module.F90 .claude docs/_build/ build*/ +artifacts/plots/ +artifacts/notebooks/ diff --git a/requirements.txt b/requirements.txt index 77fb5fe1..5f7d6d79 100644 --- a/requirements.txt +++ b/requirements.txt @@ -12,6 +12,7 @@ f90wrap==0.2.16 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 diff --git a/src/boozer_converter.F90 b/src/boozer_converter.F90 index ecd25c86..f38b9ea2 100644 --- a/src/boozer_converter.F90 +++ b/src/boozer_converter.F90 @@ -1132,11 +1132,8 @@ subroutine export_boozer_chartmap(filename) R, Zval, alam, dR_ds, dR_dt, dR_dp, & dZ_ds, dZ_dt, dZ_dp, dl_ds, dl_dt, dl_dp) - ! Convert to Cartesian using phi_B as the azimuthal angle. - ! This creates a pseudo-Cartesian representation consistent - ! with the Boozer toroidal grid (atan2(y,x) = phi_B). - x_arr(i_rho, i_theta, i_phi) = R * cos(phi_B) - y_arr(i_rho, i_theta, i_phi) = R * sin(phi_B) + x_arr(i_rho, i_theta, i_phi) = R * cos(phi_V) + y_arr(i_rho, i_theta, i_phi) = R * sin(phi_V) z_arr(i_rho, i_theta, i_phi) = Zval end do end do @@ -1193,7 +1190,7 @@ subroutine export_boozer_chartmap(filename) ! Global attributes call nc_assert(nf90_put_att(ncid, nf90_global, "rho_convention", "rho_tor"), & "att rho_convention") - call nc_assert(nf90_put_att(ncid, nf90_global, "zeta_convention", "cyl"), & + call nc_assert(nf90_put_att(ncid, nf90_global, "zeta_convention", "boozer"), & "att zeta_convention") call nc_assert(nf90_put_att(ncid, nf90_global, "rho_lcfs", rho_arr(ns_B)), & "att rho_lcfs") diff --git a/test/test_data/figure8.gvec.chartmap.nc b/test/test_data/figure8.gvec.chartmap.nc index c469bd5a..6a7624f9 100644 Binary files a/test/test_data/figure8.gvec.chartmap.nc and b/test/test_data/figure8.gvec.chartmap.nc differ diff --git a/test/test_data/wout_qa.gvec.chartmap.nc b/test/test_data/wout_qa.gvec.chartmap.nc index 658ef1e9..b5e1bec7 100644 Binary files a/test/test_data/wout_qa.gvec.chartmap.nc and b/test/test_data/wout_qa.gvec.chartmap.nc differ diff --git a/test/tests/CMakeLists.txt b/test/tests/CMakeLists.txt index 456a7caa..66336103 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,15 @@ 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(WOUT_QA_GVEC_CHARTMAP_SOURCE "${TEST_DATA_DIR}/wout_qa.gvec.chartmap.nc") set(WOUT_QA_GVEC_CHARTMAP_BINARY "${CMAKE_CURRENT_BINARY_DIR}/wout_qa.gvec.chartmap.nc") @@ -26,6 +35,11 @@ endif() set(FIGURE8_BOUNDARY_SOURCE "${TEST_DATA_DIR}/figure8.quasr.boundary.nc") set(FIGURE8_CHARTMAP_SOURCE "${TEST_DATA_DIR}/figure8.gvec.chartmap.nc") +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 set(CHARTMAP_FILE "${CMAKE_CURRENT_BINARY_DIR}/wout.chartmap.nc") @@ -368,7 +382,7 @@ add_test(NAME test_array_utils COMMAND test_array_utils.x) ${WOUT_FILE} ${CMAKE_CURRENT_BINARY_DIR}/roundtrip_test.nc export - /tmp/boozer_chartmap_roundtrip + ${CMAKE_CURRENT_BINARY_DIR}/boozer_chartmap_roundtrip 1.0e-4 1.0e-6) set_tests_properties(test_boozer_chartmap_roundtrip PROPERTIES @@ -381,7 +395,7 @@ add_test(NAME test_array_utils COMMAND test_array_utils.x) wout.nc wout_qa.gvec.chartmap.nc external - /tmp/boozer_chartmap_gvec_qa + ${CMAKE_CURRENT_BINARY_DIR}/boozer_chartmap_gvec_qa 2.5e-4 1.0e-6) set_tests_properties(test_boozer_chartmap_gvec_qa PROPERTIES @@ -395,19 +409,22 @@ add_test(NAME test_array_utils COMMAND test_array_utils.x) # 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) if(EXISTS "${FIGURE8_BOUNDARY_SOURCE}" AND EXISTS "${FIGURE8_CHARTMAP_SOURCE}") add_test(NAME test_figure8_boozer_chartmap - COMMAND ${Python_EXECUTABLE} + 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") + LABELS "slow;plot;python;golden_record" + TIMEOUT 7200 + DEPENDS test_e2e_boozer_chartmap) endif() if(SIMPLE_ENABLE_CGAL) 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/gvec_to_boozer_chartmap.py b/test/tests/gvec_to_boozer_chartmap.py index 2f15adab..33fd34bc 100644 --- a/test/tests/gvec_to_boozer_chartmap.py +++ b/test/tests/gvec_to_boozer_chartmap.py @@ -110,18 +110,9 @@ def main(): 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 - # Use pseudo-Cartesian: R from true geometry, phi_B as azimuthal angle - X_true = pos[0] * 1e2 # m -> cm - Y_true = pos[1] * 1e2 + X = pos[0] * 1e2 # m -> cm + Y = pos[1] * 1e2 Z_geom = pos[2] * 1e2 - R_geom = np.sqrt(X_true**2 + Y_true**2) - - zeta_3d = np.broadcast_to( - zeta_geom[np.newaxis, np.newaxis, :], - (n_rho, n_theta_geom, n_phi_geom) - ) - X = R_geom * np.cos(zeta_3d) - Y = R_geom * np.sin(zeta_3d) # Write NetCDF print(f"Writing {args.output}") @@ -153,7 +144,7 @@ def main(): v = ds.createVariable("num_field_periods", "i4"); v[:] = np.int32(nfp) ds.rho_convention = "rho_tor" - ds.zeta_convention = "cyl" + ds.zeta_convention = "boozer" ds.rho_lcfs = float(rho_grid[-1]) ds.boozer_field = np.int32(1) ds.torflux = torflux diff --git a/test/tests/plot_boozer_chartmap_roundtrip.py b/test/tests/plot_boozer_chartmap_roundtrip.py index 4031ef62..b8be6382 100644 --- a/test/tests/plot_boozer_chartmap_roundtrip.py +++ b/test/tests/plot_boozer_chartmap_roundtrip.py @@ -111,7 +111,7 @@ def plot_orbit_comparison(prefix, direct_label, chartmap_label, title): parser = argparse.ArgumentParser() parser.add_argument( "--prefix", - default="/tmp/boozer_chartmap_roundtrip", + default="boozer_chartmap_roundtrip", help="Artifact prefix used by the Fortran comparison test", ) parser.add_argument("--direct-label", default="Direct (VMEC)") diff --git a/test/tests/test_boozer_chartmap_roundtrip.f90 b/test/tests/test_boozer_chartmap_roundtrip.f90 index 455ebf58..fddd4ed8 100644 --- a/test/tests/test_boozer_chartmap_roundtrip.f90 +++ b/test/tests/test_boozer_chartmap_roundtrip.f90 @@ -52,7 +52,7 @@ program test_boozer_chartmap_roundtrip mode = 'export' field_tol = 1.0e-4_dp orbit_tol = 1.0e-6_dp - artifact_prefix = '/tmp/boozer_chartmap_roundtrip' + artifact_prefix = 'boozer_chartmap_roundtrip' if (command_argument_count() >= 1) then call get_command_argument(1, wout_file) @@ -69,7 +69,7 @@ program test_boozer_chartmap_roundtrip if (command_argument_count() >= 4) then call get_command_argument(4, artifact_prefix) if (len_trim(artifact_prefix) == 0) then - artifact_prefix = '/tmp/boozer_chartmap_roundtrip' + artifact_prefix = 'boozer_chartmap_roundtrip' end if end if if (command_argument_count() >= 5) then @@ -98,7 +98,7 @@ program test_boozer_chartmap_roundtrip isw_field_type = 2 rmu = 1.0e8_dp netcdffile = trim(wout_file) - multharm = 5 + multharm = 7 ns_A = 5 ns_s = 5 ns_tp = 5 @@ -211,11 +211,15 @@ program test_boozer_chartmap_roundtrip close(u_out) print *, ' max relative error Bmod:', max_err_bmod - if (max_err_bmod > field_tol) 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 within tolerance' + print *, 'INFO: skipping Bmod pass/fail threshold, max err=', max_err_bmod end if ! ========================================================= @@ -270,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 6033b199..1a11e72e 100644 --- a/test/tests/test_e2e_boozer_chartmap.py +++ b/test/tests/test_e2e_boozer_chartmap.py @@ -1,52 +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"), - "gvec_chartmap": os.path.join(TEST_DATA, "wout_qa.gvec.chartmap.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"), - "gvec_chartmap": None, + "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}' @@ -58,12 +95,14 @@ def simple_in_vmec(wout_name, trace_time, facE_al): """ -def simple_in_chartmap(chartmap_name, 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 = '{chartmap_name}' @@ -77,216 +116,364 @@ def simple_in_chartmap(chartmap_name, 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"] - gvec_chartmap = eq["gvec_chartmap"] - 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, "vmec_chartmap_run") - gvec_dir = os.path.join(eq_dir, "gvec_chartmap_run") - os.makedirs(vmec_dir) - os.makedirs(chartmap_dir) - if gvec_chartmap and os.path.exists(gvec_chartmap): - os.makedirs(gvec_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: VMEC chartmap-only path...") - with open(os.path.join(chartmap_dir, "simple.in"), "w") as f: - f.write(simple_in_chartmap("boozer_chartmap.nc", trace_time, facE_al)) - - if not run_cmd([SIMPLE_X], cwd=chartmap_dir, label=f"{name} VMEC 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" - - series = { - "VMEC-Boozer": np.loadtxt(cf_vmec), - "VMEC chartmap": np.loadtxt(cf_chartmap), - } - - if gvec_chartmap and os.path.exists(gvec_chartmap): - print(f" [{name}] Step 4: GVEC chartmap-only path...") - shutil.copy(gvec_chartmap, os.path.join(gvec_dir, "boozer_chartmap.nc")) - shutil.copy(start_boozer, os.path.join(gvec_dir, "start.dat")) - with open(os.path.join(gvec_dir, "simple.in"), "w") as f: - f.write(simple_in_chartmap("boozer_chartmap.nc", trace_time, facE_al)) - if not run_cmd([SIMPLE_X], cwd=gvec_dir, label=f"{name} GVEC chartmap"): - return False - cf_gvec = os.path.join(gvec_dir, "confined_fraction.dat") - assert os.path.exists(cf_gvec), f"{name}: GVEC confined_fraction.dat not produced" - series["GVEC chartmap"] = np.loadtxt(cf_gvec) - - # Step 5: Compare - min_len = min(len(data) for data in series.values()) - for key in list(series): - series[key] = series[key][:min_len] - - ref = series["VMEC-Boozer"] - t = ref[:, 0] - npart = ref[0, 3] - tol = 1.0 / npart - - totals = {} - passings = {} - trappeds = {} - for key, data in series.items(): - passings[key] = data[:, 1] - trappeds[key] = data[:, 2] - totals[key] = passings[key] + trappeds[key] - - max_diff_total = 0.0 - max_diff_pass = 0.0 - max_diff_trap = 0.0 - labels = list(series) - for i, left in enumerate(labels): - for right in labels[i + 1:]: - diff_total = np.max(np.abs(totals[left] - totals[right])) - diff_pass = np.max(np.abs(passings[left] - passings[right])) - diff_trap = np.max(np.abs(trappeds[left] - trappeds[right])) - 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_diff_total = max(max_diff_total, diff_total) - max_diff_pass = max(max_diff_pass, diff_pass) - max_diff_trap = max(max_diff_trap, diff_trap) - - plot_comparison(eq_dir, name, t, totals, passings, trappeds) - - if max_diff_total > tol or max_diff_pass > tol or max_diff_trap > tol: - print( - f" [{name}] FAIL: pairwise confined-fraction disagreement exceeds {tol:.6e}" +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", ) - return False - - print(f" [{name}] PASS") - return True - -def plot_comparison(outdir, name, t, totals, passings, trappeds): - import matplotlib - matplotlib.use("Agg") - import matplotlib.pyplot as plt + 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", + ) - fig, axes = plt.subplots(1, 3, figsize=(15, 4.5)) - styles = { - "VMEC-Boozer": ("k-", 1.8), - "VMEC chartmap": ("C0--", 1.4), - "GVEC chartmap": ("C3:", 2.0), + 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) - ax = axes[0] - for key, values in totals.items(): - style, width = styles.get(key, ("-", 1.5)) - ax.plot(t * 1e3, values, style, lw=width, label=key) - 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] - for key, values in passings.items(): - style, width = styles.get(key, ("-", 1.5)) - ax.plot(t * 1e3, values, style, lw=width, label=f"{key} passing") - for key, values in trappeds.items(): - style, width = styles.get(key, ("-", 1.2)) - ax.plot(t * 1e3, values, style, lw=max(width - 0.4, 0.8), alpha=0.6, - label=f"{key} 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] - ref_key = "VMEC-Boozer" - for key, values in totals.items(): - if key == ref_key: - continue - style, width = styles.get(key, ("-", 1.5)) - ax.semilogy(t * 1e3, np.abs(totals[ref_key] - values) + 1e-16, - style, lw=width, label=f"{key} - {ref_key}") - ax.set_xlabel("time [ms]") - ax.set_ylabel("|difference|") - ax.set_title("Confined fraction difference") - ax.legend(fontsize=8) - - fig.suptitle(f"E2E confined fractions: {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): + 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 index f6a7eb72..51d55abe 100644 --- a/test/tests/test_figure8_boozer_chartmap.py +++ b/test/tests/test_figure8_boozer_chartmap.py @@ -1,34 +1,56 @@ #!/usr/bin/env python3 -"""Smoke-test the figure-8 GVEC chartmap path and generate a review plot.""" +"""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 -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") -TEST_DATA = os.path.join(os.path.dirname(SCRIPT_DIR), "test_data") -CHARTMAP = os.path.join(TEST_DATA, "figure8.gvec.chartmap.nc") -BOUNDARY = os.path.join(TEST_DATA, "figure8.quasr.boundary.nc") +from boozer_chartmap_artifacts import ( + confined_metrics, + 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" +TEST_DATA = SCRIPT_DIR.parent / "test_data" +BOUNDARY = TEST_DATA / "figure8.quasr.boundary.nc" +REFERENCE_CHARTMAP = TEST_DATA / "figure8.gvec.chartmap.nc" -def write_inputs(run_dir): - with open(os.path.join(run_dir, "start.dat"), "w") as handle: - handle.write("0.30 0.00 0.00 1.0 0.40\n") - handle.write("0.30 2.10 0.00 1.0 0.20\n") - handle.write("0.45 4.20 0.10 1.0 -0.30\n") +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 open(os.path.join(run_dir, "simple.in"), "w") as handle: + with (run_dir / "simple.in").open("w", encoding="utf-8") as handle: handle.write( """&config multharm = 5 contr_pp = -1e10 -trace_time = 1d-4 -ntestpart = 3 +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 @@ -41,111 +63,338 @@ def write_inputs(run_dir): ) -def run_simple(run_dir): +def run_cmd(cmd: list[str], cwd: Path, label: str, timeout: int = 3600) -> None: + print(f"[{label}] {' '.join(cmd)}") result = subprocess.run( - [SIMPLE_X], - cwd=run_dir, + cmd, + cwd=cwd, capture_output=True, text=True, - timeout=600, + timeout=timeout, check=False, ) - if result.returncode != 0: - print(result.stdout[-4000:]) - print(result.stderr[-4000:]) - raise SystemExit(result.returncode) + if result.returncode == 0: + return + print(result.stdout[-4000:]) + print(result.stderr[-4000:]) + raise RuntimeError(f"{label} failed with exit code {result.returncode}") -def load_chartmap_boundary(path): - with netCDF4.Dataset(path) as ds: - return ( - ds["x"][:].transpose(2, 1, 0)[-1], - ds["y"][:].transpose(2, 1, 0)[-1], - ds["z"][:].transpose(2, 1, 0)[-1], - ds["Bmod"][:], - ) +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) -> 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) + 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 output -def load_quasr_boundary(path): +def load_quasr_boundary(path: Path) -> tuple[np.ndarray, np.ndarray, np.ndarray]: with netCDF4.Dataset(path) as ds: - pos = ds["pos"][:] + pos = np.asarray(ds["pos"][:], dtype=float) return pos[:, :, 0], pos[:, :, 1], pos[:, :, 2] -def plot_review(run_dir): +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 = figure8_signature(REFERENCE_CHARTMAP) + 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, chartmap_path: Path) -> None: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt xb_true, yb_true, zb_true = load_quasr_boundary(BOUNDARY) - xb, yb, zb, bmod = load_chartmap_boundary(CHARTMAP) - - fig = plt.figure(figsize=(13, 4)) - - ax = fig.add_subplot(1, 3, 1, projection="3d") - for j in range(0, xb_true.shape[0], 8): - ax.plot(xb_true[j], yb_true[j], zb_true[j], color="C0", lw=0.8, alpha=0.7) - for i in range(0, xb_true.shape[1], 8): - ax.plot(xb_true[:, i], yb_true[:, i], zb_true[:, i], color="C3", lw=0.6, alpha=0.45) - ax.set_title("QUASR figure-8 boundary") - ax.set_xlabel("x [m]") - ax.set_ylabel("y [m]") - ax.set_zlabel("z [m]") - ax.view_init(elev=18, azim=40) - - ax = fig.add_subplot(1, 3, 2, projection="3d") - for j in range(0, xb.shape[1], 2): - ax.plot(xb[:, j], yb[:, j], zb[:, j], color="C0", lw=0.8, alpha=0.7) - for i in range(0, xb.shape[0], 2): - ax.plot(xb[i], yb[i], zb[i], color="C3", lw=0.6, alpha=0.45) - ax.set_title("SIMPLE chartmap geometry") - ax.set_xlabel("x [cm]") - ax.set_ylabel("y [cm]") - ax.set_zlabel("z [cm]") - ax.view_init(elev=18, azim=40) - - ax = fig.add_subplot(1, 3, 3) + 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 = ax.imshow(bmod[mid, :, :].T, origin="lower", aspect="auto", cmap="magma") - ax.set_title("|B| at mid zeta") - ax.set_xlabel("zeta index") - ax.set_ylabel("theta index") - fig.colorbar(im, ax=ax, shrink=0.8) - - fig.tight_layout() - outpath = os.path.join(run_dir, "figure8_review.png") + 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 main(): - if not os.path.exists(SIMPLE_X): - raise SystemExit(f"Missing simple.x: {SIMPLE_X}") - for path in (CHARTMAP, BOUNDARY): - if not os.path.exists(path): - raise SystemExit(f"Missing fixture: {path}") +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"), + (BOUNDARY, "figure-8 boundary"), + (REFERENCE_CHARTMAP, "figure-8 golden chartmap"), + ]: + if not path.exists(): + raise SystemExit(f"Missing {label}: {path}") - run_dir = "/tmp/figure8_chartmap_smoke" - if os.path.exists(run_dir): + run_dir = Path.cwd() / "figure8_chartmap_smoke" + if run_dir.exists(): shutil.rmtree(run_dir) - os.makedirs(run_dir) + run_dir.mkdir(parents=True) - shutil.copy(CHARTMAP, os.path.join(run_dir, "figure8.gvec.chartmap.nc")) + chartmap = build_figure8_chartmap(run_dir) write_inputs(run_dir) - run_simple(run_dir) + run_cmd([str(SIMPLE_X)], cwd=run_dir, label="figure8 SIMPLE") - data = np.loadtxt(os.path.join(run_dir, "confined_fraction.dat")) - if data.ndim == 1: - data = data[np.newaxis, :] - confined = data[:, 1:3] + 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 SystemExit("Non-finite confined fractions in figure-8 smoke test") - if np.any(confined < -1e-12) or np.any(confined > 1.0 + 1e-12): - raise SystemExit("Confined fractions outside [0, 1] in figure-8 smoke test") + 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") - plot_review(run_dir) - print("FIGURE-8 CHARTMAP PASS") + metrics = {"Figure-8 GVEC chartmap": confined_metrics(data)} + plot_case_loss_comparison(run_dir / "figure8_losses.png", "Figure-8", metrics) + plot_review(run_dir, chartmap) + assert_figure8_golden(chartmap, run_dir) + maybe_plot_all_cases(run_dir, metrics) + print("FIGURE-8 GVEC CHARTMAP PASS") if __name__ == "__main__":