Merge pull request #1 from HPMolSim/tensor

revise interpolate_nu_single!
HPMolSim · Feb 17, 2024 · 05d8176 · 05d8176
2 parents a8a7026 + 323e82b
commit 05d8176
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 11 deletions.
diff --git a/src/FFCT/nugrid_interpolate.jl b/src/FFCT/nugrid_interpolate.jl
@@ -1,4 +1,8 @@
-@inbounds function interpolate_nu_single!(q::T, pos::NTuple{3, T}, k_x::Vector{T}, k_y::Vector{T}, phase_x::Vector{Complex{T}}, phase_y::Vector{Complex{T}}, r_z::Vector{T}, us_mat::Array{T, 3}, H_r::Array{Complex{T}, 3}, uspara::USeriesPara{T}, M_mid::Int) where{T}
+@inline function Fourier_k(i::Int, N::Int, L::T) where{T}
+    return T(2π * (i - 1 - N) / L)
+end
+
+@inbounds function interpolate_nu_single!(q::T, pos::NTuple{3, T}, N::NTuple{3, Int}, L::NTuple{3, T}, k_x::Vector{T}, k_y::Vector{T}, phase_x::Vector{Complex{T}}, phase_y::Vector{Complex{T}}, r_z::Vector{T}, us_mat::Array{T, 3}, H_r::Array{Complex{T}, 3}, uspara::USeriesPara{T}, M_mid::Int) where{T}
 
     x, y, z = pos
     revise_phase_neg!(phase_x, phase_y, k_x, k_y, x, y)
@@ -11,12 +15,12 @@
             z_temp = T(2) - T(4) * (z - r_zk)^2 / sl^2
 
             for j in 1:size(H_r, 2)
-                k_yj = k_y[j]
-                phase_yj = phase_y[j]
+                k_yj = Fourier_k(j, N[2], L[2])
                 for i in 1:size(H_r, 1)
-                    k_xi = k_x[i]
-                    phase = phase_x[i] * phase_yj
-                    H_r[i, j, k] += q * π * phase * (z_temp + (k_xi^2 + k_yj^2) * sl^2) * exp_temp * us_mat[i, j, l - M_mid]
+                    k_xi = Fourier_k(i, N[1], L[1])
+                    phase = phase_x[i] * phase_y[j]
+                    us = us_mat[i, j, l - M_mid]
+                    H_r[i, j, k] += q * π * phase * (z_temp + (k_xi^2 + k_yj^2) * sl^2) * exp_temp * us
                 end
             end
         end
@@ -25,11 +29,11 @@
     return H_r
 end
 
-function interpolate_nu!(qs::Vector{T}, poses::Vector{NTuple{3, T}}, k_x::Vector{T}, k_y::Vector{T}, phase_x::Vector{Complex{T}}, phase_y::Vector{Complex{T}}, r_z::Vector{T}, us_mat::Array{T, 3}, H_r::Array{Complex{T}, 3}, uspara::USeriesPara{T}, M_mid::Int) where{T}
+function interpolate_nu!(qs::Vector{T}, poses::Vector{NTuple{3, T}}, N::NTuple{3, Int}, L::NTuple{3, T}, k_x::Vector{T}, k_y::Vector{T}, phase_x::Vector{Complex{T}}, phase_y::Vector{Complex{T}}, r_z::Vector{T}, us_mat::Array{T, 3}, H_r::Array{Complex{T}, 3}, uspara::USeriesPara{T}, M_mid::Int) where{T}
 
     set_zeros!(H_r)
     for i in 1:length(qs)
-        interpolate_nu_single!(qs[i], poses[i], k_x, k_y, phase_x, phase_y, r_z, us_mat, H_r, uspara, M_mid)
+        interpolate_nu_single!(qs[i], poses[i], N, L, k_x, k_y, phase_x, phase_y, r_z, us_mat, H_r, uspara, M_mid)
     end
 
     return H_r

diff --git a/src/energy/energy_long.jl b/src/energy/energy_long.jl
@@ -11,7 +11,7 @@
 # 5. gather the energy
 
 function energy_long(
-    qs::Vector{T}, poses::Vector{NTuple{3, T}}, L::NTuple{3, T}, M_mid::Int, 
+    qs::Vector{T}, poses::Vector{NTuple{3, T}}, L::NTuple{3, T}, N_grid::NTuple{3, Int}, M_mid::Int, 
     k_x::Vector{T}, k_y::Vector{T}, r_z::Vector{T}, 
     phase_x::Vector{Complex{T}}, phase_y::Vector{Complex{T}},
     z::Vector{T}, sort_z::Vector{Int}, 
@@ -24,7 +24,7 @@ function energy_long(
     @assert M_mid ≤ length(uspara.sw)
 
     b_l, b_u = boundaries!(qs, poses, b_l, b_u, k_x, k_y, phase_x, phase_y, L[3], uspara, M_mid)
-    H_r = interpolate_nu!(qs, poses, k_x, k_y, phase_x, phase_y, r_z, us_mat, H_r, uspara, M_mid)
+    H_r = interpolate_nu!(qs, poses, N_grid, L, k_x, k_y, phase_x, phase_y, r_z, us_mat, H_r, uspara, M_mid)
     H_c = real2Cheb!(H_r, H_c, r_z, L[3])
     H_s = solve_eqs!(rhs, sol, H_c, H_s, b_l, b_u, ivsm, L[3])
     E_k = gather_nu(qs, poses, L, k_x, k_y, phase_x, phase_y, H_s)

diff --git a/test/energy_long.jl b/test/energy_long.jl
@@ -16,7 +16,7 @@
     k_x, k_y, r_z, us_mat, H_r, H_c, H_s, ivsm, b_l, b_u, phase_x, phase_y, rhs, sol, sort_z, z = FFCT_precompute(L, N_grid, USeriesPara(2), M_mid, n_atoms)
 
     @info "running the FFCT for the long range part of the energy"
-    E_FFCT = energy_long(qs, poses, L, M_mid, k_x, k_y, r_z, phase_x, phase_y, z, sort_z, us_mat, b_l, b_u, rhs, sol, ivsm, H_r, H_c, H_s, uspara, soepara)
+    E_FFCT = energy_long(qs, poses, L, N_grid, M_mid, k_x, k_y, r_z, phase_x, phase_y, z, sort_z, us_mat, b_l, b_u, rhs, sol, ivsm, H_r, H_c, H_s, uspara, soepara)
 
     @info "running the direct summation for the long range part of the energy"
     # using the direct summation