Merge pull request #4 from HPMolSim/dev

added thin system energy solver by fat

Project.toml

@@ -6,10 +6,12 @@ version = "1.0.0-DEV"
 [deps]
 ChebParticleMesh = "1983ef0c-217d-4026-99b0-9163e7750d85"
 ExTinyMD = "fec76197-d59f-46dd-a0ed-76a83c21f7aa"
+FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FastChebInterp = "cf66c380-9a80-432c-aff8-4f9c79c0bdde"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922"
+Polynomials = "f27b6e38-b328-58d1-80ce-0feddd5e7a45"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 SumOfExpVPMR = "2c69873a-e0bb-44e1-90b4-d15ac3b7e936"
 

src/FFCT/Chebyshev.jl

@@ -0,0 +1,66 @@
+function USeries_direct(z::T, k_xi::T, k_yj::T, uspara::USeriesPara{T}, M_mid::Int) where{T}
+    t = zero(T)
+    k2 = k_xi^2 + k_yj^2
+
+    for l in M_mid + 1:length(uspara.sw)
+        sl, wl = uspara.sw[l]
+        t += π * wl * sl^2 * exp(-sl^2 * k2 / 4) * exp(-z^2 / sl^2)
+    end
+
+    return t
+end
+
+function ChebUSeries_k(k_xi::T, k_yj::T, L_z::T, uspara::USeriesPara{T}, M::Int, Q::Int) where{T}
+
+    f = z -> USeries_direct(z, k_xi, k_yj, uspara, M)
+    x = chebpoints(Q, zero(T), L_z)
+
+    return chebinterp(f.(x), zero(T), L_z)
+end
+
+function ChebUSeries(k_x::Vector{T}, k_y::Vector{T}, L_z::T, uspara::USeriesPara{T}, M::Int, Q::Int) where{T}
+
+    cheb_mat = Array{ChebPoly{1, T, T}, 2}(undef, (length(k_x), length(k_y)))
+
+    for i in 1:length(k_x)
+        for j in 1:length(k_y)
+            cheb_mat[i, j] = ChebUSeries_k(k_x[i], k_y[j], L_z, uspara, M, Q)
+        end
+    end
+
+    return cheb_mat
+end
+
+@inline function chebpoly(n::Int, x::T, scale::T) where{T}
+    return cos(n * acos(x / scale))
+end
+
+@inbounds function real2Cheb!(H_r::Array{Complex{T}, 3}, H_c::Array{Complex{T}, 3}, r_z::Vector{T}, L_z::T) where{T}
+
+    set_zeros!(H_c)
+    N_z = size(H_r, 3)
+
+    for k in 1:N_z, l in 1:N_z
+        cheb_temp = k == 1 ? 0.5 : chebpoly(k - 1, r_z[l] - L_z / T(2), L_z / T(2))
+        for j in 1:size(H_r, 2), i in 1:size(H_r, 1)
+            H_c[i, j, k] += 2 / N_z * H_r[i, j, l] * cheb_temp
+        end
+    end
+
+    return H_c
+end
+
+@inbounds function Cheb2real!(H_c::Array{Complex{T}, 3}, H_r::Array{Complex{T}, 3}, r_z::Vector{T}, L_z::T) where{T}
+
+    set_zeros!(H_r)
+    N_z = size(H_c, 3)
+
+    for k in 1:N_z, l in 1:N_z
+        cheb_temp = chebpoly(k - 1, r_z[l] - L_z / T(2), L_z / T(2))
+        for j in 1:size(H_c, 2), i in 1:size(H_c, 1)
+            H_r[i, j, l] += H_c[i, j, k] * cheb_temp
+        end
+    end
+
+    return H_r
+end

src/FFCT/Taylor.jl

@@ -0,0 +1,52 @@
+function Greens_Q2d_direct(z::T, k_xi::T, k_yj::T, uspara::USeriesPara{T}, M_mid::Int) where{T}
+    t = zero(T)
+    k2 = k_xi^2 + k_yj^2
+
+    for l in M_mid + 1:length(uspara.sw)
+        sl, wl = uspara.sw[l]
+        t += π * wl * sl^2 * exp(-sl^2 * k2 / 4) * exp(-z^2 / sl^2)
+    end
+
+    return t
+end
+
+
+#(x^n / n!)
+@inbounds @inline function special_powern(x::T, n::Int) where{T}
+    t = one(T)
+    for i in 1:n
+        t *= x / T(i)
+    end
+    return t
+end
+
+# the resulting Polynomial is in the form of f = x -> poly((x / L_z)^2)
+function TaylorUSeries_k(k_xi::T, k_yj::T, L_z::T, uspara::USeriesPara{T}, M_mid::Int, Q::Int) where{T}
+    coefs = zeros(T, Q)
+    k2 = k_xi^2 + k_yj^2
+    if !(k2 ≈ zero(T))
+        for l in M_mid + 1:length(uspara.sw)
+            sl, wl = uspara.sw[l]
+            for n in 0:Q - 1
+                coefs[n + 1] += (-1)^n * π * wl * sl^2 * exp(-sl^2 * k2 / 4) * special_powern(L_z^2 / sl^2, n)
+            end
+        end
+    end
+    return coefs
+end
+
+function TaylorUSeries(k_x::Vector{T}, k_y::Vector{T}, L_z::T, uspara::USeriesPara{T}, M_mid::Int, Q::Int) where{T}
+
+    taylor_mats = [zeros(T, (length(k_x), length(k_y))) for i in 1:Q]
+
+    for i in 1:length(k_x)
+        for j in 1:length(k_y)
+            coefs = TaylorUSeries_k(k_x[i], k_y[j], L_z, uspara, M_mid, Q)
+            for k in 1:Q
+                taylor_mats[k][i, j] = coefs[k]
+            end
+        end
+    end
+
+    return taylor_mats
+end

src/FFCT/gather.jl

@@ -0,0 +1,54 @@
+@inbounds function gather_nu_single(q::T, pos::NTuple{3, T}, L::NTuple{3, T}, k_x::Vector{T}, k_y::Vector{T}, phase_x::Vector{Complex{T}}, phase_y::Vector{Complex{T}}, H_c::Array{Complex{T}, 3}) where{T}
+    x, y, z = pos
+    L_x, L_y, L_z = L
+    ϕ = zero(Complex{T})
+
+    revise_phase_pos!(phase_x, phase_y, k_x, k_y, x - T(L[1] / 2), y - T(L[2] / 2))
+
+    for k in 1:size(H_c, 3)
+        cheb_val = chebpoly(k - 1, z - L_z / T(2), L_z / T(2))
+        ϕ += dot(phase_x', (@view H_c[:, :, k]), phase_y)* cheb_val
+    end
+
+    return q * ϕ / (2 * L_x * L_y)
+end
+
+function gather_nu(qs::Vector{T}, poses::Vector{NTuple{3, T}}, L::NTuple{3, T}, k_x::Vector{T}, k_y::Vector{T}, phase_x::Vector{Complex{T}}, phase_y::Vector{Complex{T}}, H_c::Array{Complex{T}, 3}) where{T}
+    E = zero(Complex{T})
+    for i in 1:length(qs)
+        E += gather_nu_single(qs[i], poses[i], L, k_x, k_y, phase_x, phase_y, H_c)
+    end
+    return real(E)
+end
+
+@inbounds function gather_thin_single(q::T, pos::NTuple{3, T}, L_z::T, pad_grid::Array{Complex{T}, 3}, gridinfo::GridInfo{2, T}, cheb_value::Vector{Array{T, 1}}, cheb_coefs::NTuple{2, ChebCoef{T}}) where{T}
+
+    potential_i = zero(T)
+    idl = gridinfo.index_list
+    x, y, z = pos
+
+    near_id_image = image_grid_id((pos[1], pos[2]), gridinfo)
+    near_pos_image = image_grid_pos(near_id_image, gridinfo)
+    for i in 1:2
+        dx = pos[i] - near_pos_image[i]
+        pwcheb_eval!(dx, cheb_value[i], cheb_coefs[i])
+    end
+
+    for i in gridinfo.iter_list
+        image_id = near_id_image.id .+ i
+        for k in 1:size(pad_grid, 3)
+            cheb_val = chebpoly(k - 1, z - L_z / T(2), L_z / T(2))
+            potential_i += real(pad_grid[idl[1][image_id[1]], idl[2][image_id[2]], k]) * prod(cheb_value[j][i[j] + gridinfo.w[j] + 1] for j in 1:2) * cheb_val
+        end
+    end
+
+    return q * prod(gridinfo.h) * potential_i
+end
+
+function gather_thin(qs::Vector{T}, poses::Vector{NTuple{3, T}}, L::NTuple{3, T}, H_c::Array{Complex{T}, 3}, gridinfo::GridInfo{2, T}, cheb_value::Vector{Array{T, 1}}, cheb_coefs::NTuple{2, ChebCoef{T}}) where{T}
+    E = zero(Complex{T})
+    for i in 1:length(qs)
+        E += gather_thin_single(qs[i], poses[i], L[3], H_c, gridinfo, cheb_value, cheb_coefs)
+    end
+    return real(E) / 2
+end

src/FFCT/interpolate.jl

@@ -0,0 +1,126 @@
+function interpolate_nu_loop!(
+    H_r::Array{Complex{T}, 3},
+    qs::Vector{T}, poses::Vector{NTuple{3, T}}, 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}, uspara::USeriesPara{T}, M_mid::Int) where{T}
+
+    set_zeros!(H_r)
+
+    for n in 1:length(qs)
+        x, y, z = poses[n]
+        q = qs[n]
+        revise_phase_neg!(phase_x, phase_y, k_x, k_y, x - L[1] / 2, y - L[2] / 2)
+
+        for k in 1:size(H_r, 3)
+            r_zk = r_z[k]
+            for l in M_mid + 1:length(uspara.sw)
+                sl, wl = uspara.sw[l]
+
+                for j in 1:size(H_r, 2)
+                    k_yj = k_y[j]
+                    for i in 1:size(H_r, 1)
+                        k_xi = k_x[i]
+                        k2 = k_xi^2 + k_yj^2
+                        phase = phase_x[i] * phase_y[j]
+                        if !(k2 ≈ zero(T))
+                            H_r[i, j, k] += q * π * wl * sl^2 * phase * exp(-sl^2 * k2 / 4) * exp(-(z - r_zk)^2 / sl^2)
+                        end
+                    end
+                end
+            end
+        end
+    end
+
+    return H_r
+end
+
+function interpolate_nu_cheb!(
+    H_r::Array{Complex{T}, 3},
+    qs::Vector{T}, poses::Vector{NTuple{3, T}}, 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}, cheb_mat::Array{ChebPoly{1, T, T}, 2}) where{T}
+
+    set_zeros!(H_r)
+
+    for n in 1:length(qs)
+        x, y, z = poses[n]
+        q = qs[n]
+        revise_phase_neg!(phase_x, phase_y, k_x, k_y, x - L[1] / 2, y - L[2] / 2)
+
+        for k in 1:size(H_r, 3)
+            r_zk = r_z[k]
+            for j in 1:size(H_r, 2)
+                k_yj = k_y[j]
+                for i in 1:size(H_r, 1)
+                    k_xi = k_x[i]
+                    k2 = k_xi^2 + k_yj^2
+                    if !(k2 ≈ zero(T))
+                        phase = phase_x[i] * phase_y[j]
+                        H_r[i, j, k] += q * phase * cheb_mat[i, j](abs(z - r_zk))
+                    end
+                end
+            end
+        end
+    end
+
+    return H_r
+end
+
+
+@inbounds function interpolate_thin_single!(q::T, pos::NTuple{3, T}, pad_grid::Array{Complex{T}, 3}, gridinfo::GridInfo{2, T}, cheb_value::Vector{Array{T, 1}}, cheb_coefs::NTuple{2, ChebCoef{T}}, r_z::Vector{T}, Taylor_order::Int, L_z::T) where{T}
+
+    idl = gridinfo.index_list
+    pos_new = (pos[1], pos[2])
+
+    near_id_image = image_grid_id(pos_new, gridinfo)
+    near_pos_image = image_grid_pos(near_id_image, gridinfo)
+    for i in 1:2
+        dx = pos[i] - near_pos_image[i]
+        pwcheb_eval!(dx, cheb_value[i], cheb_coefs[i])
+    end
+
+    for i in gridinfo.iter_list
+        image_id = near_id_image.id .+ i
+        for k in 1:size(pad_grid, 3)
+            qn = q * ((pos[3] - r_z[k]) / L_z)^(2 * (Taylor_order - 1))
+            pad_grid[idl[1][image_id[1]], idl[2][image_id[2]], k] += Complex{T}(qn * prod(cheb_value[j][i[j] + gridinfo.w[j] + 1] for j in 1:2))
+        end
+    end
+
+    return nothing
+end
+
+@inbounds function interpolate_thin!(qs::Vector{T}, poses::Vector{NTuple{3, T}}, pad_grids::Vector{Array{Complex{T}, 3}}, gridinfo::GridInfo{2, T}, cheb_value::Vector{Array{T, 1}}, cheb_coefs::NTuple{2, ChebCoef{T}}, r_z::Vector{T}, L_z::T) where{T}
+
+    @assert length(qs) == length(poses)
+    set_zeros!.(pad_grids)
+
+    for Taylor_order in 1:length(pad_grids)
+        for i in 1:length(qs)
+            interpolate_thin_single!(qs[i], poses[i], pad_grids[Taylor_order], gridinfo, cheb_value, cheb_coefs, r_z, Taylor_order, L_z)
+        end
+    end
+
+    return nothing
+end
+
+function ϕkz_direct(z0::T, qs::Vector{T}, poses::Vector{NTuple{3, T}}, k_x::T, k_y::T, uspara::USeriesPara{T}, M_mid::Int) where{T}
+
+    s = zero(T)
+
+    for n in 1:length(qs)
+        x, y, z = poses[n]
+        q = qs[n]
+
+        for l in M_mid + 1:length(uspara.sw)
+            sl, wl = uspara.sw[l]
+            k2 = k_x^2 + k_y^2
+            phase = exp( - T(1)im * (k_x * x + k_y * y))
+            s += q * π * wl * sl^2 * phase * exp(-sl^2 * k2 / 4) * exp(-(z - z0)^2 / sl^2)
+        end
+    end
+
+    return s
+end