boozer_sub: VMEC<->Boozer angle transforms

src/boozer_converter.F90

@@ -16,6 +16,8 @@ module boozer_sub
     public :: get_boozer_coordinates
     public :: splint_boozer_coord
     public :: reset_boozer_batch_splines
+    public :: vmec_to_boozer, boozer_to_vmec
+    public :: delthe_delphi_BV
 
     ! Constants
     real(dp), parameter :: TWOPI = 2.0_dp*3.14159265358979_dp
@@ -37,6 +39,17 @@ module boozer_sub
     type(BatchSplineData1D), save :: bcovar_tp_batch_spline
     logical, save :: bcovar_tp_batch_spline_ready = .false.
 
+    ! Batch splines for angle transformations (VMEC <-> Boozer).
+    ! delt_delp_V holds (theta_B-theta_V, phi_B-phi_V) on the VMEC-angle grid;
+    ! delt_delp_B holds the same differences on the Boozer-angle grid.
+    type(BatchSplineData3D), save :: delt_delp_V_batch_spline
+    logical, save :: delt_delp_V_batch_spline_ready = .false.
+    real(dp), allocatable, save :: delt_delp_V_grid(:, :, :, :)
+
+    type(BatchSplineData3D), save :: delt_delp_B_batch_spline
+    logical, save :: delt_delp_B_batch_spline_ready = .false.
+    real(dp), allocatable, save :: delt_delp_B_grid(:, :, :, :)
+
 contains
 
     !> Initialize Boozer coordinates using VMEC field (backward compatibility)
@@ -110,6 +123,7 @@ subroutine get_boozer_coordinates_impl
         call build_boozer_aphi_batch_spline
         call build_boozer_bcovar_tp_batch_spline
         call build_boozer_field3d_batch_spline
+        call build_boozer_delt_delp_batch_splines
 
     end subroutine get_boozer_coordinates_impl
 
@@ -358,12 +372,139 @@ subroutine splint_boozer_coord(r, vartheta_B, varphi_B, mode_secders, &
 
     end subroutine splint_boozer_coord
 
+    !> Computes deltheta_BV = vartheta_B - theta_V and delphi_BV = varphi_B - varphi_V
+    !> together with their first derivatives over the angles.
+    !> isw=0: arguments are VMEC angles (r, theta, varphi)
+    !> isw=1: arguments are Boozer angles (r, vartheta_B, varphi_B)
+    subroutine delthe_delphi_BV(isw, r, vartheta, varphi, deltheta_BV, delphi_BV, &
+                                ddeltheta_BV, ddelphi_BV)
+        use boozer_coordinates_mod, only: use_del_tp_B
+        use chamb_mod, only: rnegflag
+
+        integer, intent(in) :: isw
+        real(dp), intent(in) :: r, vartheta, varphi
+        real(dp), intent(out) :: deltheta_BV, delphi_BV
+        real(dp), dimension(2), intent(out) :: ddeltheta_BV, ddelphi_BV
+
+        real(dp) :: rho_tor, x_eval(3), y_eval(2), dy_eval(3, 2)
+        real(dp) :: r_local
+
+        r_local = r
+        if (r_local <= 0.0_dp) then
+            rnegflag = .true.
+            r_local = abs(r_local)
+        end if
+
+        rho_tor = sqrt(r_local)
+        x_eval(1) = rho_tor
+        x_eval(2) = vartheta
+        x_eval(3) = varphi
+
+        if (isw .eq. 0) then
+            if (.not. delt_delp_V_batch_spline_ready) then
+                error stop "delthe_delphi_BV: V batch spline not initialized"
+            end if
+            call evaluate_batch_splines_3d_der(delt_delp_V_batch_spline, x_eval, &
+                                               y_eval, dy_eval)
+        elseif (isw .eq. 1) then
+            if (.not. use_del_tp_B) then
+                print *, 'delthe_delphi_BV : Boozer data is not loaded'
+                return
+            end if
+            if (.not. delt_delp_B_batch_spline_ready) then
+                error stop "delthe_delphi_BV: B batch spline not initialized"
+            end if
+            call evaluate_batch_splines_3d_der(delt_delp_B_batch_spline, x_eval, &
+                                               y_eval, dy_eval)
+        else
+            print *, 'delthe_delphi_BV : unknown value of switch isw'
+            return
+        end if
+
+        deltheta_BV = y_eval(1)
+        delphi_BV = y_eval(2)
+
+        ddeltheta_BV(1) = dy_eval(2, 1)
+        ddelphi_BV(1) = dy_eval(2, 2)
+
+        ddeltheta_BV(2) = dy_eval(3, 1)
+        ddelphi_BV(2) = dy_eval(3, 2)
+
+    end subroutine delthe_delphi_BV
+
+    !> Convert VMEC angles (r, theta, varphi) to Boozer angles (vartheta_B, varphi_B)
+    subroutine vmec_to_boozer(r, theta, varphi, vartheta_B, varphi_B)
+        use new_vmec_stuff_mod, only: nper
+
+        real(dp), intent(in) :: r, theta, varphi
+        real(dp), intent(out) :: vartheta_B, varphi_B
+
+        real(dp) :: deltheta_BV, delphi_BV
+        real(dp), dimension(2) :: ddeltheta_BV, ddelphi_BV
+
+        call delthe_delphi_BV(0, r, theta, varphi, deltheta_BV, delphi_BV, &
+                              ddeltheta_BV, ddelphi_BV)
+
+        vartheta_B = modulo(theta + deltheta_BV, TWOPI)
+        varphi_B = modulo(varphi + delphi_BV, TWOPI/real(nper, dp))
+
+    end subroutine vmec_to_boozer
+
+    !> Convert Boozer angles (r, vartheta_B, varphi_B) to VMEC angles (theta, varphi)
+    !> by Newton iteration on the VMEC-angle transform.
+    subroutine boozer_to_vmec(r, vartheta_B, varphi_B, theta, varphi)
+        use boozer_coordinates_mod, only: use_del_tp_B
+
+        real(dp), intent(in) :: r, vartheta_B, varphi_B
+        real(dp), intent(out) :: theta, varphi
+
+        real(dp), parameter :: epserr = 1.0e-14_dp
+        integer, parameter :: niter = 100
+
+        integer :: iter
+        real(dp) :: deltheta_BV, delphi_BV
+        real(dp) :: f1, f2, f11, f12, f21, f22, delthe, delphi, det
+        real(dp), dimension(2) :: ddeltheta_BV, ddelphi_BV
+
+        if (use_del_tp_B) then
+            call delthe_delphi_BV(1, r, vartheta_B, varphi_B, deltheta_BV, delphi_BV, &
+                                  ddeltheta_BV, ddelphi_BV)
+
+            theta = vartheta_B - deltheta_BV
+            varphi = varphi_B - delphi_BV
+        else
+            theta = vartheta_B
+            varphi = varphi_B
+        end if
+
+        do iter = 1, niter
+            call delthe_delphi_BV(0, r, theta, varphi, deltheta_BV, delphi_BV, &
+                                  ddeltheta_BV, ddelphi_BV)
+
+            f1 = theta + deltheta_BV - vartheta_B
+            f2 = varphi + delphi_BV - varphi_B
+            f11 = 1.0_dp + ddeltheta_BV(1)
+            f12 = ddeltheta_BV(2)
+            f21 = ddelphi_BV(1)
+            f22 = 1.0_dp + ddelphi_BV(2)
+
+            det = f11*f22 - f12*f21
+            delthe = (f2*f12 - f1*f22)/det
+            delphi = (f1*f21 - f2*f11)/det
+
+            theta = theta + delthe
+            varphi = varphi + delphi
+            if (abs(delthe) + abs(delphi) .lt. epserr) exit
+        end do
+
+    end subroutine boozer_to_vmec
+
     subroutine compute_boozer_data
         ! Computes Boozer coordinate transformations and magnetic field data
         use boozer_coordinates_mod, only: ns_s_B, ns_tp_B, ns_B, n_theta_B, n_phi_B, &
                                           hs_B, h_theta_B, h_phi_B, &
                                           s_Bcovar_tp_B, &
-                                          use_B_r
+                                          use_B_r, use_del_tp_B
         use binsrc_sub, only: binsrc
         use plag_coeff_sub, only: plag_coeff
         use spline_vmec_sub
@@ -450,6 +591,9 @@ subroutine compute_boozer_data
         call ensure_grid_3d(bmod_grid, ns_B, n_theta_B, n_phi_B)
         call ensure_grid_3d(sqrt_g_ss_grid, ns_B, n_theta_B, n_phi_B)
         if (use_B_r) call ensure_grid_3d(br_grid, ns_B, n_theta_B, n_phi_B)
+        call ensure_grid_4d(delt_delp_V_grid, ns_B, n_theta_B, n_phi_B, 2)
+        if (use_del_tp_B) call ensure_grid_4d(delt_delp_B_grid, ns_B, n_theta_B, &
+                                              n_phi_B, 2)
 
         do i = 0, ns_tp_B
             wint_t(i) = h_theta_B**(i + 1)/real(i + 1, dp)
@@ -561,6 +705,9 @@ subroutine compute_boozer_data
 ! $\Delta \vartheta_{BV}=\vartheta_B-\theta$:
             deltheta_BV_Vg = aiota*delphi_BV_Vg + alam_2D
 
+            delt_delp_V_grid(i_rho, :, :, 1) = deltheta_BV_Vg
+            delt_delp_V_grid(i_rho, :, :, 2) = delphi_BV_Vg
+
 ! At this point, all quantities are specified on
 ! equidistant grid in VMEC angles $(\theta,\varphi)$
 
@@ -616,6 +763,10 @@ subroutine compute_boozer_data
                 end do
             end do
 
+            if (use_del_tp_B) then
+                delt_delp_B_grid(i_rho, :, :, 1) = perqua_2D(1, :, :)
+                delt_delp_B_grid(i_rho, :, :, 2) = perqua_2D(2, :, :)
+            end if
             bmod_grid(i_rho, :, :) = perqua_2D(3, :, :)
             sqrt_g_ss_grid(i_rho, :, :) = perqua_2D(7, :, :)
 
@@ -754,6 +905,19 @@ subroutine ensure_grid_3d(grid, n1, n2, n3)
         end if
     end subroutine ensure_grid_3d
 
+    !> Ensure 4D grid is allocated with correct dimensions
+    subroutine ensure_grid_4d(grid, n1, n2, n3, n4)
+        real(dp), allocatable, intent(inout) :: grid(:, :, :, :)
+        integer, intent(in) :: n1, n2, n3, n4
+
+        if (.not. allocated(grid)) then
+            allocate (grid(n1, n2, n3, n4))
+        else if (any(shape(grid) /= [n1, n2, n3, n4])) then
+            deallocate (grid)
+            allocate (grid(n1, n2, n3, n4))
+        end if
+    end subroutine ensure_grid_4d
+
     subroutine reset_boozer_batch_splines
         if (aphi_batch_spline_ready) then
             call destroy_batch_splines_1d(aphi_batch_spline)
@@ -771,6 +935,16 @@ subroutine reset_boozer_batch_splines
         if (allocated(bmod_grid)) deallocate (bmod_grid)
         if (allocated(sqrt_g_ss_grid)) deallocate (sqrt_g_ss_grid)
         if (allocated(br_grid)) deallocate (br_grid)
+        if (delt_delp_V_batch_spline_ready) then
+            call destroy_batch_splines_3d(delt_delp_V_batch_spline)
+            delt_delp_V_batch_spline_ready = .false.
+        end if
+        if (allocated(delt_delp_V_grid)) deallocate (delt_delp_V_grid)
+        if (delt_delp_B_batch_spline_ready) then
+            call destroy_batch_splines_3d(delt_delp_B_batch_spline)
+            delt_delp_B_batch_spline_ready = .false.
+        end if
+        if (allocated(delt_delp_B_grid)) deallocate (delt_delp_B_grid)
     end subroutine reset_boozer_batch_splines
 
     subroutine build_boozer_aphi_batch_spline
@@ -890,4 +1064,68 @@ subroutine build_boozer_field3d_batch_spline
         deallocate (y_batch)
     end subroutine build_boozer_field3d_batch_spline
 
+    subroutine build_boozer_delt_delp_batch_splines
+        ! Build angle-transform splines from the 4D grids filled in
+        ! compute_boozer_data. The V grid is always built; the Boozer-angle B
+        ! grid is built only when use_del_tp_B is set.
+        use boozer_coordinates_mod, only: ns_s_B, ns_tp_B, ns_B, n_theta_B, n_phi_B, &
+                                          hs_B, h_theta_B, h_phi_B, use_del_tp_B
+
+        integer :: order(3)
+        real(dp) :: x_min(3), x_max(3)
+        logical :: periodic(3)
+        real(dp), allocatable :: y_batch(:, :, :, :)
+
+        if (.not. allocated(delt_delp_V_grid)) then
+            error stop &
+                "build_boozer_delt_delp_batch_splines: delt_delp_V_grid not allocated"
+        end if
+
+        if (delt_delp_V_batch_spline_ready) then
+            call destroy_batch_splines_3d(delt_delp_V_batch_spline)
+            delt_delp_V_batch_spline_ready = .false.
+        end if
+
+        order = [ns_s_B, ns_tp_B, ns_tp_B]
+        if (any(order < 3) .or. any(order > 5)) then
+            error stop "build_boozer_delt_delp_batch_splines: order must be 3..5"
+        end if
+
+        x_min = [0.0_dp, 0.0_dp, 0.0_dp]
+        x_max(1) = hs_B*real(ns_B - 1, dp)
+        x_max(2) = h_theta_B*real(n_theta_B - 1, dp)
+        x_max(3) = h_phi_B*real(n_phi_B - 1, dp)
+
+        periodic = [.false., .true., .true.]
+
+        allocate (y_batch(ns_B, n_theta_B, n_phi_B, 2))
+        y_batch(:, :, :, 1) = delt_delp_V_grid(:, :, :, 1)
+        y_batch(:, :, :, 2) = delt_delp_V_grid(:, :, :, 2)
+
+        call construct_batch_splines_3d(x_min, x_max, y_batch, order, periodic, &
+                                        delt_delp_V_batch_spline)
+        delt_delp_V_batch_spline_ready = .true.
+
+        if (use_del_tp_B) then
+            if (.not. allocated(delt_delp_B_grid)) then
+                error stop &
+            "build_boozer_delt_delp_batch_splines: delt_delp_B_grid not allocated"
+            end if
+
+            if (delt_delp_B_batch_spline_ready) then
+                call destroy_batch_splines_3d(delt_delp_B_batch_spline)
+                delt_delp_B_batch_spline_ready = .false.
+            end if
+
+            y_batch(:, :, :, 1) = delt_delp_B_grid(:, :, :, 1)
+            y_batch(:, :, :, 2) = delt_delp_B_grid(:, :, :, 2)
+
+            call construct_batch_splines_3d(x_min, x_max, y_batch, order, periodic, &
+                                            delt_delp_B_batch_spline)
+            delt_delp_B_batch_spline_ready = .true.
+        end if
+
+        deallocate (y_batch)
+    end subroutine build_boozer_delt_delp_batch_splines
+
 end module boozer_sub

test/CMakeLists.txt

@@ -493,6 +493,15 @@ set_tests_properties(test_boozer_converter_vs_simple PROPERTIES
   WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}
   DEPENDS setup_vmec_wout)
 
+add_executable(test_boozer_angle_roundtrip.x source/test_boozer_angle_roundtrip.f90)
+target_link_libraries(test_boozer_angle_roundtrip.x neo)
+add_test(NAME test_boozer_angle_roundtrip
+  COMMAND test_boozer_angle_roundtrip.x)
+set_tests_properties(test_boozer_angle_roundtrip PROPERTIES
+  FAIL_REGULAR_EXPRESSION "STOP"
+  WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}
+  DEPENDS setup_vmec_wout)
+
 add_executable(test_chartmap_vmec_mapping.x source/test_chartmap_vmec_mapping.f90)
 target_link_libraries(test_chartmap_vmec_mapping.x neo)
 add_test(NAME test_chartmap_vmec_mapping