diff --git a/src/boozer_converter.F90 b/src/boozer_converter.F90 index e316b9b3..ecd25c86 100644 --- a/src/boozer_converter.F90 +++ b/src/boozer_converter.F90 @@ -20,6 +20,7 @@ module boozer_sub public :: delthe_delphi_BV public :: reset_boozer_batch_splines public :: load_boozer_from_chartmap + public :: export_boozer_chartmap ! Constants real(dp), parameter :: TWOPI = 2.0_dp*3.14159265358979_dp @@ -892,6 +893,7 @@ subroutine load_boozer_from_chartmap(filename) integer :: ncid, status, dimid, varid integer :: n_rho, n_theta, n_zeta, nfp_file + integer :: n_theta_field, n_zeta_field real(dp) :: torflux_val real(dp), allocatable :: rho(:), theta(:), zeta(:) real(dp), allocatable :: A_phi_arr(:), B_theta_arr(:), B_phi_arr(:) @@ -947,9 +949,22 @@ subroutine load_boozer_from_chartmap(filename) call nc_check(nf90_inq_varid(ncid, "B_phi", varid), "var B_phi") call nc_check(nf90_get_var(ncid, varid, B_phi_arr), "get B_phi") - ! Read 3D Bmod: NetCDF dims (zeta, theta, rho) map to Fortran - ! (rho, theta, zeta) via NF90 dimension reversal - allocate (Bmod_arr(n_rho, n_theta, n_zeta)) + ! Read Bmod field dimensions (endpoint-included, may differ from geometry dims) + status = nf90_inq_dimid(ncid, "theta_field", dimid) + if (status == nf90_noerr) then + call nc_check(nf90_inquire_dimension(ncid, dimid, len=n_theta_field), & + "len theta_field") + call nc_check(nf90_inq_dimid(ncid, "zeta_field", dimid), "dim zeta_field") + call nc_check(nf90_inquire_dimension(ncid, dimid, len=n_zeta_field), & + "len zeta_field") + else + ! Fallback: field uses same dims as geometry (old format) + n_theta_field = n_theta + n_zeta_field = n_zeta + end if + + ! Read 3D Bmod on field grid + allocate (Bmod_arr(n_rho, n_theta_field, n_zeta_field)) call nc_check(nf90_inq_varid(ncid, "Bmod", varid), "var Bmod") call nc_check(nf90_get_var(ncid, varid, Bmod_arr), "get Bmod") @@ -963,11 +978,12 @@ subroutine load_boozer_from_chartmap(filename) ns_s_B = 5 ns_tp_B = 5 ns_B = n_rho - n_theta_B = n_theta - n_phi_B = n_zeta + n_theta_B = n_theta_field + n_phi_B = n_zeta_field hs_B = rho(2) - rho(1) - h_theta_B = theta(2) - theta(1) - h_phi_B = zeta(2) - zeta(1) + ! Field grid step from endpoint-included dimensions + h_theta_B = TWOPI / real(n_theta_field - 1, dp) + h_phi_B = TWOPI / real(nfp_file, dp) / real(n_zeta_field - 1, dp) use_B_r = .false. use_del_tp_B = .false. @@ -1000,12 +1016,14 @@ subroutine load_boozer_from_chartmap(filename) deallocate (y_bcovar) ! Build Bmod 3D batch spline over (rho_tor, theta_B, phi_B) + ! Uses endpoint-included field grid matching original VMEC Boozer splines order_3d = [ns_s_B, ns_tp_B, ns_tp_B] periodic_3d = [.false., .true., .true.] - x_min_3d = [rho(1), theta(1), zeta(1)] - x_max_3d = [rho(n_rho), theta(n_theta), zeta(n_zeta)] + x_min_3d = [rho(1), 0.0_dp, 0.0_dp] + x_max_3d = [rho(n_rho), h_theta_B * real(n_theta_field - 1, dp), & + h_phi_B * real(n_zeta_field - 1, dp)] - allocate (y_bmod(n_rho, n_theta, n_zeta, 1)) + allocate (y_bmod(n_rho, n_theta_field, n_zeta_field, 1)) y_bmod(:, :, :, 1) = Bmod_arr call construct_batch_splines_3d(x_min_3d, x_max_3d, y_bmod, order_3d, & periodic_3d, bmod_br_batch_spline) @@ -1014,8 +1032,8 @@ subroutine load_boozer_from_chartmap(filename) deallocate (y_bmod) print *, 'Loaded Boozer splines from chartmap: ', trim(filename) - print *, ' nfp=', nfp_file, ' ns=', n_rho, ' ntheta=', n_theta, & - ' nphi=', n_zeta + print *, ' nfp=', nfp_file, ' ns=', n_rho, ' ntheta_field=', & + n_theta_field, ' nphi_field=', n_zeta_field print *, ' torflux=', torflux_val contains @@ -1032,6 +1050,194 @@ end subroutine nc_check end subroutine load_boozer_from_chartmap + subroutine export_boozer_chartmap(filename) + !> Export Boozer coordinate data computed by get_boozer_coordinates() + !> to an extended chartmap NetCDF file. Must be called after + !> get_boozer_coordinates() and while VMEC splines are still active + !> (needed for X, Y, Z geometry evaluation). + use vector_potentail_mod, only: ns, hs, sA_phi, torflux + use new_vmec_stuff_mod, only: nper + use boozer_coordinates_mod, only: ns_B, n_theta_B, n_phi_B, & + hs_B, h_theta_B, h_phi_B, & + s_Bcovar_tp_B + use spline_vmec_sub, only: splint_vmec_data + use netcdf + + character(len=*), intent(in) :: filename + + integer :: ncid, status + integer :: dim_rho, dim_theta, dim_zeta + integer :: dim_theta_field, dim_zeta_field + integer :: var_rho, var_theta, var_zeta + integer :: var_x, var_y, var_z + integer :: var_aphi, var_btheta, var_bphi, var_bmod, var_nfp + integer :: i_rho, i_theta, i_phi + integer :: n_theta_out, n_phi_out + real(dp) :: rho_tor, s, theta_B, phi_B, theta_V, phi_V + real(dp) :: R, Zval, alam + real(dp) :: A_phi_dum, A_theta_dum, dA_phi_ds, dA_theta_ds, aiota + real(dp) :: dR_ds, dR_dt, dR_dp, dZ_ds, dZ_dt, dZ_dp + real(dp) :: dl_ds, dl_dt, dl_dp + real(dp), parameter :: rho_min = sqrt(1.0e-6_dp) + real(dp), allocatable :: rho_arr(:), theta_arr(:), zeta_arr(:) + real(dp), allocatable :: A_phi_arr(:), B_theta_arr(:), B_phi_arr(:) + real(dp), allocatable :: x_arr(:, :, :), y_arr(:, :, :), z_arr(:, :, :) + + ! Chartmap geometry requires endpoint-excluded angular grids (libneo validator). + ! Boozer field data uses endpoint-included grids (matching original splines). + n_theta_out = n_theta_B - 1 + n_phi_out = n_phi_B - 1 + + allocate (rho_arr(ns_B)) + allocate (theta_arr(n_theta_out), zeta_arr(n_phi_out)) + allocate (A_phi_arr(ns_B), B_theta_arr(ns_B), B_phi_arr(ns_B)) + allocate (x_arr(ns_B, n_theta_out, n_phi_out)) + allocate (y_arr(ns_B, n_theta_out, n_phi_out)) + allocate (z_arr(ns_B, n_theta_out, n_phi_out)) + + ! Radial grid + do i_rho = 1, ns_B + rho_arr(i_rho) = max(real(i_rho - 1, dp) * hs_B, rho_min) + end do + ! Angular grids (endpoint excluded, for chartmap geometry) + do i_theta = 1, n_theta_out + theta_arr(i_theta) = real(i_theta - 1, dp) * h_theta_B + end do + do i_phi = 1, n_phi_out + zeta_arr(i_phi) = real(i_phi - 1, dp) * h_phi_B + end do + + ! Extract 1D profiles + do i_rho = 1, ns_B + A_phi_arr(i_rho) = sA_phi(1, i_rho) ! zeroth spline coeff = value + B_theta_arr(i_rho) = s_Bcovar_tp_B(1, 1, i_rho) + B_phi_arr(i_rho) = s_Bcovar_tp_B(2, 1, i_rho) + end do + + ! Compute X, Y, Z geometry on the Boozer grid (endpoint-excluded) + do i_phi = 1, n_phi_out + do i_theta = 1, n_theta_out + do i_rho = 1, ns_B + rho_tor = rho_arr(i_rho) + s = rho_tor**2 + theta_B = theta_arr(i_theta) + phi_B = zeta_arr(i_phi) + + ! Convert Boozer angles to VMEC angles + call boozer_to_vmec(s, theta_B, phi_B, theta_V, phi_V) + + ! Evaluate VMEC geometry at (s, theta_V, phi_V) + call splint_vmec_data(s, theta_V, phi_V, & + A_phi_dum, A_theta_dum, dA_phi_ds, dA_theta_ds, aiota, & + 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) + z_arr(i_rho, i_theta, i_phi) = Zval + end do + end do + end do + + ! Write NetCDF file + status = nf90_create(trim(filename), nf90_clobber, ncid) + call nc_assert(status, "create " // trim(filename)) + + ! Dimensions: geometry uses endpoint-excluded grids (chartmap validator), + ! field data uses endpoint-included grids (exact spline reproduction) + call nc_assert(nf90_def_dim(ncid, "rho", ns_B, dim_rho), "def_dim rho") + call nc_assert(nf90_def_dim(ncid, "theta", n_theta_out, dim_theta), & + "def_dim theta") + call nc_assert(nf90_def_dim(ncid, "zeta", n_phi_out, dim_zeta), & + "def_dim zeta") + call nc_assert(nf90_def_dim(ncid, "theta_field", n_theta_B, dim_theta_field), & + "def_dim theta_field") + call nc_assert(nf90_def_dim(ncid, "zeta_field", n_phi_B, dim_zeta_field), & + "def_dim zeta_field") + + ! Coordinate variables + call nc_assert(nf90_def_var(ncid, "rho", nf90_double, [dim_rho], var_rho), & + "def_var rho") + call nc_assert(nf90_def_var(ncid, "theta", nf90_double, [dim_theta], & + var_theta), "def_var theta") + call nc_assert(nf90_def_var(ncid, "zeta", nf90_double, [dim_zeta], & + var_zeta), "def_var zeta") + + ! Geometry (NF90 reverses dims: Fortran (rho,theta,zeta) -> NetCDF (zeta,theta,rho)) + call nc_assert(nf90_def_var(ncid, "x", nf90_double, & + [dim_rho, dim_theta, dim_zeta], var_x), "def_var x") + call nc_assert(nf90_put_att(ncid, var_x, "units", "cm"), "att x units") + call nc_assert(nf90_def_var(ncid, "y", nf90_double, & + [dim_rho, dim_theta, dim_zeta], var_y), "def_var y") + call nc_assert(nf90_put_att(ncid, var_y, "units", "cm"), "att y units") + call nc_assert(nf90_def_var(ncid, "z", nf90_double, & + [dim_rho, dim_theta, dim_zeta], var_z), "def_var z") + call nc_assert(nf90_put_att(ncid, var_z, "units", "cm"), "att z units") + + ! Boozer field data + call nc_assert(nf90_def_var(ncid, "A_phi", nf90_double, [dim_rho], var_aphi), & + "def_var A_phi") + call nc_assert(nf90_def_var(ncid, "B_theta", nf90_double, [dim_rho], & + var_btheta), "def_var B_theta") + call nc_assert(nf90_def_var(ncid, "B_phi", nf90_double, [dim_rho], & + var_bphi), "def_var B_phi") + call nc_assert(nf90_def_var(ncid, "Bmod", nf90_double, & + [dim_rho, dim_theta_field, dim_zeta_field], var_bmod), & + "def_var Bmod") + call nc_assert(nf90_def_var(ncid, "num_field_periods", nf90_int, var_nfp), & + "def_var nfp") + + ! 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"), & + "att zeta_convention") + call nc_assert(nf90_put_att(ncid, nf90_global, "rho_lcfs", rho_arr(ns_B)), & + "att rho_lcfs") + call nc_assert(nf90_put_att(ncid, nf90_global, "boozer_field", 1), & + "att boozer_field") + call nc_assert(nf90_put_att(ncid, nf90_global, "torflux", torflux), & + "att torflux") + + call nc_assert(nf90_enddef(ncid), "enddef") + + ! Write data + call nc_assert(nf90_put_var(ncid, var_rho, rho_arr), "put rho") + call nc_assert(nf90_put_var(ncid, var_theta, theta_arr), "put theta") + call nc_assert(nf90_put_var(ncid, var_zeta, zeta_arr), "put zeta") + call nc_assert(nf90_put_var(ncid, var_x, x_arr), "put x") + call nc_assert(nf90_put_var(ncid, var_y, y_arr), "put y") + call nc_assert(nf90_put_var(ncid, var_z, z_arr), "put z") + call nc_assert(nf90_put_var(ncid, var_aphi, A_phi_arr), "put A_phi") + call nc_assert(nf90_put_var(ncid, var_btheta, B_theta_arr), "put B_theta") + call nc_assert(nf90_put_var(ncid, var_bphi, B_phi_arr), "put B_phi") + call nc_assert(nf90_put_var(ncid, var_bmod, bmod_grid), "put Bmod") + call nc_assert(nf90_put_var(ncid, var_nfp, nper), "put nfp") + + call nc_assert(nf90_close(ncid), "close") + + print *, 'Exported Boozer chartmap to ', trim(filename) + print *, ' nfp=', nper, ' ns=', ns_B, ' ntheta=', n_theta_out, & + ' nphi=', n_phi_out + print *, ' torflux=', torflux + + contains + + subroutine nc_assert(stat, loc) + integer, intent(in) :: stat + character(len=*), intent(in) :: loc + if (stat /= nf90_noerr) then + print *, "export_boozer_chartmap: NetCDF error at ", trim(loc), & + ": ", trim(nf90_strerror(stat)) + error stop + end if + end subroutine nc_assert + + end subroutine export_boozer_chartmap + subroutine build_boozer_aphi_batch_spline use vector_potentail_mod, only: ns, hs, sA_phi use new_vmec_stuff_mod, only: ns_A diff --git a/test/tests/CMakeLists.txt b/test/tests/CMakeLists.txt index 46af006e..15ae3c15 100644 --- a/test/tests/CMakeLists.txt +++ b/test/tests/CMakeLists.txt @@ -351,6 +351,27 @@ add_test(NAME test_array_utils COMMAND test_array_utils.x) WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} LABELS "unit") + # Boozer chartmap roundtrip test: requires wout.nc + 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) + set_tests_properties(test_boozer_chartmap_roundtrip PROPERTIES + WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} + LABELS "slow") + + # 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} + ${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") + if(SIMPLE_ENABLE_CGAL) add_executable(test_stl_wall_intersection.x test_stl_wall_intersection.f90) target_link_libraries(test_stl_wall_intersection.x simple) diff --git a/test/tests/export_boozer_chartmap_tool.f90 b/test/tests/export_boozer_chartmap_tool.f90 new file mode 100644 index 00000000..f7b805af --- /dev/null +++ b/test/tests/export_boozer_chartmap_tool.f90 @@ -0,0 +1,95 @@ +program export_boozer_chartmap_tool + !> Tool to export a Boozer chartmap from VMEC and convert start.dat + !> from VMEC reference coordinates to Boozer/chartmap reference coordinates. + !> + !> Usage: export_boozer_chartmap_tool.x + use, intrinsic :: iso_fortran_env, only: dp => real64 + use new_vmec_stuff_mod, only: netcdffile, multharm, ns_A, ns_s, ns_tp, nper + use parmot_mod, only: rmu + use velo_mod, only: isw_field_type + use boozer_coordinates_mod, only: use_B_r + use boozer_sub, only: get_boozer_coordinates, vmec_to_boozer, & + export_boozer_chartmap + use spline_vmec_sub, only: spline_vmec_data + use vmecin_sub, only: stevvo + + implicit none + + real(dp), parameter :: twopi = 2.0_dp * 3.14159265358979_dp + character(len=1024) :: wout_file, chartmap_file, start_vmec, start_boozer + character(len=1024) :: line + integer :: nargs, ipart, npart, ios + real(dp) :: s, theta_v, phi_v, v, lam, theta_b, phi_b + real(dp) :: RT0, R0i, cbfi, bz0i, bf0, fper, rho + integer :: L1i + + nargs = command_argument_count() + if (nargs /= 4) then + print *, 'Usage: export_boozer_chartmap_tool.x ', & + ' ' + error stop + end if + + call get_command_argument(1, wout_file) + call get_command_argument(2, chartmap_file) + call get_command_argument(3, start_vmec) + call get_command_argument(4, start_boozer) + + ! Initialize VMEC + netcdffile = wout_file + multharm = 3 + ns_A = 5 + ns_s = 5 + ns_tp = 5 + isw_field_type = 2 + rmu = 1.0e8_dp + + call spline_vmec_data + call stevvo(RT0, R0i, L1i, cbfi, bz0i, bf0) + fper = twopi / real(L1i, dp) + + ! Compute Boozer coordinates + use_B_r = .false. + call get_boozer_coordinates + + ! Export chartmap + call export_boozer_chartmap(chartmap_file) + + ! Count particles in start_vmec + npart = 0 + open(unit=10, 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) + if (ios /= 0) exit + npart = npart + 1 + end do + close(10) + + 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) + + do ipart = 1, npart + read(10, *) 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) + + ! 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 + end do + + close(10) + close(11) + + print *, 'Written ', trim(start_boozer) + +end program export_boozer_chartmap_tool diff --git a/test/tests/plot_boozer_chartmap_roundtrip.py b/test/tests/plot_boozer_chartmap_roundtrip.py new file mode 100644 index 00000000..1fd21966 --- /dev/null +++ b/test/tests/plot_boozer_chartmap_roundtrip.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python3 +"""Plot Boozer chartmap roundtrip comparison: fields and orbits side by side.""" + +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") + s = data[:, 0] + theta = data[:, 1] + Bmod_ref = data[:, 3] + Bmod_new = data[:, 4] + rel_err = data[:, 5] + + 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.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.set_xlabel("theta_B") + ax.set_ylabel("|B| [G]") + ax.set_title("Bmod vs theta") + ax.legend() + + ax = axes[2] + ax.semilogy(s, rel_err, "ko", ms=4) + ax.set_xlabel("s") + ax.set_ylabel("Relative error") + ax.set_title("Bmod roundtrip error") + ax.axhline(1e-5, color="r", ls="--", label="tolerance") + ax.legend() + + fig.suptitle("Boozer Chartmap Field Roundtrip", fontsize=14) + fig.tight_layout() + fig.savefig("/tmp/boozer_field_comparison.png", dpi=150) + print("Saved /tmp/boozer_field_comparison.png") + + +def plot_orbit_comparison(): + try: + direct = np.loadtxt("/tmp/orbit_direct.dat") + chartmap = np.loadtxt("/tmp/orbit_chartmap.dat") + except Exception as e: + print(f"Could not load orbit data: {e}") + return + + n = min(len(direct), len(chartmap)) + direct = direct[:n] + chartmap = chartmap[:n] + + 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] + + 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.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.set_xlabel("time") + ax.set_ylabel("theta_B") + ax.set_title("Poloidal angle") + 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.set_xlabel("time") + ax.set_ylabel("phi_B") + ax.set_title("Toroidal angle") + ax.legend() + + ax = axes[1, 1] + ax.semilogy(t_d, np.abs(s_d - s_c), "k-", lw=0.5) + 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() + + fig.suptitle("Boozer Chartmap Orbit Roundtrip", fontsize=14) + fig.tight_layout() + fig.savefig("/tmp/boozer_orbit_comparison.png", dpi=150) + print("Saved /tmp/boozer_orbit_comparison.png") + + +if __name__ == "__main__": + plot_field_comparison() + plot_orbit_comparison() diff --git a/test/tests/test_boozer_chartmap_roundtrip.f90 b/test/tests/test_boozer_chartmap_roundtrip.f90 new file mode 100644 index 00000000..fb7003d9 --- /dev/null +++ b/test/tests/test_boozer_chartmap_roundtrip.f90 @@ -0,0 +1,276 @@ +program test_boozer_chartmap_roundtrip + !> Roundtrip validation: VMEC -> Boozer -> export chartmap -> reimport -> + !> compare field values and symplectic orbits between direct and file paths. + use, intrinsic :: iso_fortran_env, only: dp => real64 + use new_vmec_stuff_mod, only: nper, rmajor, netcdffile, multharm, ns_A, ns_s, ns_tp + use parmot_mod, only: rmu, ro0 + use velo_mod, only: isw_field_type + use boozer_coordinates_mod, only: use_B_r + use boozer_sub, only: splint_boozer_coord, get_boozer_coordinates, & + vmec_to_boozer, export_boozer_chartmap, load_boozer_from_chartmap, & + reset_boozer_batch_splines + use spline_vmec_sub, only: spline_vmec_data + use vmecin_sub, only: stevvo + use field_can_mod, only: field_can_from_name, field_can_init, & + eval_field => evaluate, field_can_t, get_val + use orbit_symplectic, only: orbit_sympl_init, orbit_timestep_sympl, & + symplectic_integrator_t + + implicit none + + real(dp), parameter :: pi = 3.14159265358979_dp + real(dp), parameter :: twopi = 2.0_dp * pi + + ! Field comparison + integer, parameter :: n_test = 50 + real(dp) :: s_test(n_test), th_test(n_test), ph_test(n_test) + real(dp) :: Bmod_ref(n_test), Bmod_new(n_test) + + ! Orbit comparison + integer, parameter :: n_orbit = 500 + real(dp) :: orbit_direct(n_orbit, 5), orbit_chartmap(n_orbit, 5) + + ! Work variables + real(dp) :: fper, dtau, RT0 + real(dp) :: A_theta, A_phi_val, dA_theta_dr, dA_phi_dr + real(dp) :: d2A_phi_dr2, d3A_phi_dr3 + real(dp) :: Bth, dBth, d2Bth, Bph, dBph, d2Bph + 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 + type(symplectic_integrator_t) :: si + type(field_can_t) :: f + + nfail = 0 + chartmap_file = 'roundtrip_test.nc' + + print *, 'Starting roundtrip test...' + + ! ========================================================= + ! Step 1: Initialize VMEC and compute Boozer coordinates + ! ========================================================= + isw_field_type = 2 + rmu = 1.0e8_dp + netcdffile = 'wout.nc' + multharm = 5 + ns_A = 5 + ns_s = 5 + ns_tp = 5 + + call spline_vmec_data + block + integer :: L1i + real(dp) :: R0i, cbfi, bz0i, bf0 + call stevvo(RT0, R0i, L1i, cbfi, bz0i, bf0) + fper = twopi / real(L1i, dp) + end block + phi_period = fper + RT0 = rmajor + + use_B_r = .false. + call get_boozer_coordinates + call field_can_from_name("boozer") + + ! ro0 for orbit integration: rlarm * bmod00 + ro0 = 2.23e-2_dp * 2.81e5_dp + + print *, '=== Step 1: VMEC + Boozer initialized ===' + print *, ' nper=', nper, ' R0=', RT0, ' fper=', fper + + ! ========================================================= + ! Step 2: Evaluate fields at test points (direct path) + ! ========================================================= + call init_test_points(phi_period) + + do i = 1, n_test + call splint_boozer_coord(s_test(i), th_test(i), ph_test(i), 0, & + A_theta, A_phi_val, dA_theta_dr, dA_phi_dr, d2A_phi_dr2, & + d3A_phi_dr3, Bth, dBth, d2Bth, Bph, dBph, d2Bph, & + Bmod, dBmod, d2Bmod, Br, dBr, d2Br) + Bmod_ref(i) = Bmod + end do + + print *, '=== Step 2: Direct field evaluation done ===' + + ! ========================================================= + ! Step 3: Export to Boozer chartmap NetCDF + ! ========================================================= + call export_boozer_chartmap(chartmap_file) + print *, '=== Step 3: Exported chartmap ===' + + ! ========================================================= + ! Step 4: Reimport from chartmap and evaluate fields + ! ========================================================= + call reset_boozer_batch_splines + call load_boozer_from_chartmap(chartmap_file) + + do i = 1, n_test + call splint_boozer_coord(s_test(i), th_test(i), ph_test(i), 0, & + A_theta, A_phi_val, dA_theta_dr, dA_phi_dr, d2A_phi_dr2, & + d3A_phi_dr3, Bth, dBth, d2Bth, Bph, dBph, d2Bph, & + Bmod, dBmod, d2Bmod, Br, dBr, d2Br) + Bmod_new(i) = Bmod + end do + + print *, '=== Step 4: Chartmap field evaluation done ===' + + ! ========================================================= + ! Step 5: 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' + + do i = 1, n_test + if (abs(Bmod_ref(i)) > 0.0_dp) then + rel_err = abs(Bmod_new(i) - Bmod_ref(i)) / abs(Bmod_ref(i)) + else + 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), & + Bmod_ref(i), Bmod_new(i), rel_err + end do + close(20) + + 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 + else + print *, 'PASS: Bmod roundtrip error < 1e-4' + 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' + do i = 1, n_steps_done + write(22, '(5es18.10)') orbit_chartmap(i, :) + end do + close(22) + + print *, '=== Step 7: Chartmap orbit done, steps=', n_steps_done, ' ===' + + ! ========================================================= + ! Step 8: Compare orbits + ! ========================================================= + max_err_orbit = 0.0_dp + do i = 1, n_steps_done + rel_err = abs(orbit_direct(i, 2) - orbit_chartmap(i, 2)) + max_err_orbit = max(max_err_orbit, rel_err) + end do + + print *, ' max |s_direct - s_chartmap|:', max_err_orbit + if (max_err_orbit > 1.0e-6_dp) then + print *, 'FAIL: orbit roundtrip error too large:', max_err_orbit + nfail = nfail + 1 + else + print *, 'PASS: orbit roundtrip error < 1e-6' + end if + + ! ========================================================= + ! Summary + ! ========================================================= + print *, '' + if (nfail == 0) then + print *, 'All roundtrip tests passed.' + else + print *, nfail, ' roundtrip tests failed.' + error stop 'Roundtrip test failures' + end if + +contains + + subroutine init_test_points(phi_per) + real(dp), intent(in) :: phi_per + integer :: j + real(dp) :: frac + + do j = 1, n_test + frac = real(j - 1, 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) + end do + end subroutine init_test_points + + subroutine run_symplectic_orbit(z0_in, dt, nsteps, orbit_out, steps_done) + real(dp), intent(in) :: z0_in(5), dt + integer, intent(in) :: nsteps + real(dp), intent(out) :: orbit_out(nsteps, 5) + integer, intent(out) :: steps_done + + type(symplectic_integrator_t) :: si_loc + type(field_can_t) :: f_loc + real(dp) :: z_loc(4), pphi + integer :: j, ierr_loc + + ! Initialize field_can_t + call eval_field(f_loc, z0_in(1), z0_in(2), z0_in(3), 0) + + f_loc%mu = 0.5_dp * z0_in(4)**2 * (1.0_dp - z0_in(5)**2) / & + f_loc%Bmod * 2.0_dp + f_loc%ro0 = ro0 / sqrt(2.0_dp) + f_loc%vpar = z0_in(4) * z0_in(5) * sqrt(2.0_dp) + + z_loc(1:3) = z0_in(1:3) + pphi = f_loc%vpar * f_loc%hph + f_loc%Aph / f_loc%ro0 + z_loc(4) = pphi + + ! Midpoint integrator (mode=3), single step per call + call orbit_sympl_init(si_loc, f_loc, z_loc, & + dt / sqrt(2.0_dp), 1, 1.0e-10_dp, 3) + + steps_done = 0 + do j = 1, nsteps + call orbit_timestep_sympl(si_loc, f_loc, ierr_loc) + if (ierr_loc /= 0) exit + steps_done = j + orbit_out(j, 1) = dt * real(j, dp) + orbit_out(j, 2) = si_loc%z(1) ! s + orbit_out(j, 3) = si_loc%z(2) ! theta + orbit_out(j, 4) = si_loc%z(3) ! phi + orbit_out(j, 5) = si_loc%z(4) ! pphi + end do + end subroutine run_symplectic_orbit + +end program test_boozer_chartmap_roundtrip diff --git a/test/tests/test_e2e_boozer_chartmap.py b/test/tests/test_e2e_boozer_chartmap.py new file mode 100644 index 00000000..1a0f2d8e --- /dev/null +++ b/test/tests/test_e2e_boozer_chartmap.py @@ -0,0 +1,243 @@ +#!/usr/bin/env python3 +"""End-to-end test: VMEC-Boozer vs Boozer chartmap confined fractions. + +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. +""" + +import os +import sys +import shutil +import subprocess +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") + +EQUILIBRIA = [ + { + "name": "QA", + "wout": os.path.join(TEST_DATA, "wout.nc"), + "trace_time": "1d-2", + "facE_al": "1.0d0", + }, + { + "name": "NCSX", + "wout": os.path.join(TEST_DATA, "wout_ncsx.nc"), + "trace_time": "1d-2", + "facE_al": "100.0d0", + }, +] + + +def simple_in_vmec(wout_name, trace_time, facE_al): + return f"""\ +&config +multharm = 3 +contr_pp = -1e10 +trace_time = {trace_time} +sbeg = 0.3d0 +ntestpart = 32 +netcdffile = '{wout_name}' +isw_field_type = 2 +deterministic = .True. +integmode = 1 +facE_al = {facE_al} +/ +""" + + +def simple_in_chartmap(trace_time, facE_al): + return f"""\ +&config +multharm = 3 +contr_pp = -1e10 +trace_time = {trace_time} +sbeg = 0.3d0 +ntestpart = 32 +field_input = 'boozer_chartmap.nc' +coord_input = 'boozer_chartmap.nc' +isw_field_type = 2 +deterministic = .True. +integmode = 1 +startmode = 2 +facE_al = {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): + 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 + + print() + if all_pass: + print("ALL EQUILIBRIA PASSED") + else: + print("SOME EQUILIBRIA FAILED") + sys.exit(1) + + +if __name__ == "__main__": + main()