Merge pull request #3 from HPMolSim/Taylor

added cheb interp for U_series

Project.toml

@@ -7,6 +7,7 @@ version = "1.0.0-DEV"
 ChebParticleMesh = "1983ef0c-217d-4026-99b0-9163e7750d85"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ExTinyMD = "fec76197-d59f-46dd-a0ed-76a83c21f7aa"
+FastChebInterp = "cf66c380-9a80-432c-aff8-4f9c79c0bdde"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922"

src/FFCT/cheb_interpolate.jl

@@ -0,0 +1,32 @@
+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 * exp(-sl^2 * k2 / 4) * exp(-z^2 / sl^2) * (T(2) - T(4) * z^2 / sl^2 + k2 * 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

src/FFCT/nugrid_interpolate.jl

@@ -29,8 +29,7 @@ H2[i, j, k] := phase_xys[i, j, n] * exp_coef[l, k, n] * k_mat[i, j] * us_mat[i,
     k_x::Vector{T}, k_y::Vector{T}, k_mat::Array{T, 2}, 
     phase_xs::Array{Complex{T}, 2}, phase_ys::Array{Complex{T}, 2}, phase_xys::Array{Complex{T}, 3}, 
     z_coef::Array{T, 3}, exp_coef::Array{T, 3}, temp_ijlk::Array{Complex{T}, 4}, temp_ijl::Array{Complex{T}, 3},
-    r_z::Vector{T}, us_mat::Array{Complex{T}, 3}, uspara::USeriesPara{T}, M_mid::Int,  
-    size_dict::Dict{Char, Int64}) where{T}
+    r_z::Vector{T}, us_mat::Array{Complex{T}, 3}, uspara::USeriesPara{T}, M_mid::Int, size_dict::Dict{Char, Int64}) where{T}
 
     revise_phase_neg_all!(qs, poses, L, phase_xs, phase_ys, phase_xys, k_x, k_y)
     revise_z_coef!(z_coef, exp_coef, r_z, poses, uspara, M_mid)
@@ -68,6 +67,48 @@ end
     return H_r
 end
 
+@inbounds function interpolate_nu_cheb_single!(
+    H_r::Array{Complex{T}, 3},
+    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}}, 
+    r_z::Vector{T}, cheb_mat::Array{ChebPoly{1, T, T}, 2}) where{T}
+
+    x, y, z = pos
+    revise_phase_neg!(phase_x, phase_y, k_x, k_y, x - L[1] / 2, y - L[2] / 2)
+
+
+    for j in 1:size(H_r, 2)
+        phase_yj = phase_y[j]
+        for i in 1:size(H_r, 1)
+            phase_xi = phase_x[i]
+            phase = phase_xi * phase_yj
+            f_cheb = cheb_mat[i, j]
+            for k in 1:size(H_r, 3)
+                r_zk = r_z[k]
+                H_r[i, j, k] += q * phase * f_cheb(abs(z - r_zk))
+            end
+        end
+    end
+
+    return H_r
+end
+
+@inbounds 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 i in 1:length(qs)
+        interpolate_nu_cheb_single!(H_r, qs[i], poses[i], L, k_x, k_y, phase_x, phase_y, r_z, cheb_mat)
+    end
+
+    return H_r
+end
+
 @inbounds function interpolate_nu_loop_single!(
     H_r::Array{Complex{T}, 3},
     q::T, pos::NTuple{3, T}, 

src/FastSpecSoG.jl

@@ -1,19 +1,21 @@
 module FastSpecSoG
 
-using ExTinyMD, LinearAlgebra, SpecialFunctions, ChebParticleMesh, SumOfExpVPMR, LoopVectorization, OMEinsum
+using ExTinyMD, LinearAlgebra, SpecialFunctions, ChebParticleMesh, SumOfExpVPMR, LoopVectorization, OMEinsum, FastChebInterp
 
-export USeriesPara, U_series, BSA
+import FastChebInterp: ChebPoly
+
+export USeriesPara, U_series, BSA, ChebUSeries
 export FSSoG_naive
 export short_energy_naive, long_energy_naive, energy_naive
 export mid_paras_gen, energy_mid
+
 export FFCT_precompute, boundaries!, inverse_mat
-export interpolate_nu_loop_single!, interpolate_nu_loop!, interpolate_nu_einsum!, interpolate_nu_einsum_non_inplace!
-export real2Cheb!, solve_eqs!
-export ather_nu, gather_nu_single
-export zeroth_order
-export energy_long_loop_k, energy_long_einsum_k
+export interpolate_nu_einsum!, interpolate_nu_cheb!, interpolate_nu_loop!
+export real2Cheb!, solve_eqs!, gather_nu
+
+export energy_long_loop_k, energy_long_einsum_k, energy_long_cheb_k, energy_long_0
+
 export long_energy_us_k, long_energy_us_0, long_energy_sw_0, long_energy_sw_k
-export energy_long_0
 export SoePara, SoePara4, SoePara8, SoePara16
 
 include("types.jl")
@@ -33,5 +35,6 @@ include("FFCT/nugrid_interpolate.jl")
 include("FFCT/linear_eqs.jl")
 include("FFCT/nugrid_gather.jl")
 include("FFCT/zeroth_order.jl")
+include("FFCT/cheb_interpolate.jl")
 
 end

src/energy/energy_long.jl

@@ -58,6 +58,30 @@ function energy_long_einsum_k(
     return E_k
 end
 
+function energy_long_cheb_k(
+    qs::Vector{T}, poses::Vector{NTuple{3, T}}, L::NTuple{3, T}, 
+    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}},
+    us_mat::Array{Complex{T}, 3}, b_l::Array{Complex{T}, 2}, b_u::Array{Complex{T}, 2}, 
+    rhs::Vector{Complex{T}}, sol::Vector{Complex{T}}, ivsm::Array{T, 4},
+    H_r::Array{Complex{T}, 3}, H_c::Array{Complex{T}, 3}, H_s::Array{Complex{T}, 3}, 
+    uspara::USeriesPara{T}, cheb_mat::Array{ChebPoly{1, T, T}, 2}
+    ) where{T}
+
+    @assert M_mid ≤ length(uspara.sw)
+
+    b_l, b_u = boundaries!(qs, poses, L, b_l, b_u, k_x, k_y, us_mat, phase_x, phase_y, L[3], uspara, M_mid)
+    H_r = interpolate_nu_cheb!(H_r,qs, poses, L, k_x, k_y, phase_x, phase_y, r_z, cheb_mat)
+    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)
+
+    @debug "long range energy, cheb, FFCT" E_k
+
+    return E_k
+end
+
 function energy_long_0(
     qs::Vector{T}, poses::Vector{NTuple{3, T}}, L::NTuple{3, T}, M_mid::Int, 
     z::Vector{T}, sort_z::Vector{Int}, 

test/Project.toml

@@ -1,4 +1,5 @@
 [deps]
 EwaldSummations = "329efeb5-5dc7-45f3-8303-d0bcfeef1c8a"
 ExTinyMD = "fec76197-d59f-46dd-a0ed-76a83c21f7aa"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+TaylorSeries = "6aa5eb33-94cf-58f4-a9d0-e4b2c4fc25ea"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/cheb_interpolate.jl

@@ -0,0 +1,10 @@
+@testset "Cheb expansion of U_series" begin
+    @info "testing the Cheb expansion"
+    k_xi = 0.1
+    k_yj = 0.2
+    L_z = 10.0
+    z = [0.01:0.01:L_z...]
+    f_cheb = FastSpecSoG.ChebUSeries_k(k_xi, k_yj, L_z, USeriesPara(2), 3, 16)
+    direct_U = map(x -> FastSpecSoG.USeries_direct(x, k_xi, k_yj, USeriesPara(2), 3), z)
+    @test isapprox(f_cheb.(z), direct_U)
+end

test/energy_long.jl

@@ -20,16 +20,18 @@
 
     E_FFCT_einsum_k = energy_long_einsum_k(qs, poses, L, M_mid, k_x, k_y, k_mat, r_z, phase_x, phase_y, phase_xs, phase_ys, phase_xys, temp_ijlk, temp_ijl, size_dict, z_coef, exp_coef, us_mat, b_l, b_u, rhs, sol, ivsm, H_r, H_c, H_s, uspara)
 
+    E_FFCT_cheb_k = energy_long_cheb_k(qs, poses, L, M_mid, k_x, k_y, r_z, phase_x, phase_y, us_mat, b_l, b_u, rhs, sol, ivsm, H_r, H_c, H_s, uspara, ChebUSeries(k_x, k_y, L[3], uspara, M_mid, 64))
+
     E_FFCT_0 = energy_long_0(qs, poses, L, M_mid, z, sort_z, uspara, soepara)
 
     @info "running the direct summation for the long range part of the energy"
     # using the direct summation
     E_direct_k = long_energy_us_k(qs, poses, 30, L, uspara, M_mid + 1, length(uspara.sw))
     E_direct_0 = long_energy_us_0(qs, poses, L, uspara, M_mid + 1, length(uspara.sw))
 
-    @show E_FFCT_loop_k, E_FFCT_einsum_k, E_direct_k
     @test E_FFCT_0 ≈ E_direct_0
     @test E_FFCT_einsum_k ≈ E_FFCT_loop_k
     @test isapprox(E_FFCT_loop_k, E_direct_k, atol=1e-8)
     @test isapprox(E_FFCT_einsum_k, E_direct_k, atol=1e-8)
+    @test isapprox(E_FFCT_cheb_k, E_direct_k, atol=1e-8)
 end

test/energy_mid.jl

@@ -25,6 +25,5 @@
     # using the direct summation
     E_direct = long_energy_us_k(qs, poses, 30, L, uspara, 1, M_mid) + long_energy_us_0(qs, poses, L, uspara, 1, M_mid)
 
-    @show E_3DFFT, E_direct
     @test isapprox(E_3DFFT, E_direct, atol=1e-4)
 end

test/interpolate_nu.jl

@@ -1,29 +1,36 @@
 @testset "interpolate on non-uniform grid" begin
+    @info "testing the interpolation on non-uniform grid"
     n_atoms = 8
     L = (100.0, 100.0, 100.0)
 
     qs = rand(n_atoms)
     qs .-= sum(qs) ./ n_atoms
     poses = [tuple(L .* rand(3)...) for i in 1:n_atoms]
 
-    # using the 3DFFT
+    # using the FFCT
     N_grid = (16, 16, 16)
     uspara = USeriesPara(2)
+    Q = 64
     soepara = SoePara16()
-    
-    for M_mid in 0:length(uspara.sw) - 1
+
+    for M_mid in 2:length(uspara.sw) - 1
         @testset "M_mid = $M_mid" begin
             k_x, k_y, k_mat, r_z, us_mat, H_r, H_c, H_s, ivsm, b_l, b_u, phase_x, phase_y, phase_xs, phase_ys, phase_xys, rhs, sol, sort_z, z, size_dict, temp_ijlk, temp_ijl, z_coef, exp_coef = FFCT_precompute(L, N_grid, USeriesPara(2), M_mid, n_atoms)
 
+            cheb_mat = ChebUSeries(k_x, k_y, L[3], uspara, M_mid, Q)
+
             H_r_einsum = copy(interpolate_nu_einsum!(copy(H_r), qs, poses, L, k_x, k_y, k_mat, phase_xs, phase_ys, phase_xys, z_coef, exp_coef, temp_ijlk, temp_ijl, r_z, us_mat, uspara, M_mid, size_dict))
 
-            H_r_einsum_direct = copy(interpolate_nu_einsum_non_inplace!(copy(H_r), qs, poses, L, k_x, k_y, k_mat, phase_xs, phase_ys, phase_xys, z_coef, exp_coef, r_z, us_mat, uspara, M_mid))
+            H_r_einsum_direct = copy(FastSpecSoG.interpolate_nu_einsum_non_inplace!(copy(H_r), qs, poses, L, k_x, k_y, k_mat, phase_xs, phase_ys, phase_xys, z_coef, exp_coef, r_z, us_mat, uspara, M_mid))
 
             H_r_loop = copy(interpolate_nu_loop!(copy(H_r), qs, poses, N_grid, L, k_x, k_y, phase_x, phase_y, r_z, us_mat, uspara, M_mid))
 
+            H_r_cheb = copy(interpolate_nu_cheb!(copy(H_r), qs, poses, L, k_x, k_y, phase_x, phase_y, r_z, cheb_mat))
+
             @test H_r_einsum ≈ H_r_einsum_direct
             @test H_r_einsum_direct ≈ H_r_loop
             @test H_r_einsum ≈ H_r_loop
+            @test isapprox(H_r_einsum, H_r_cheb; atol=1e-8)
         end
     end
 end

test/runtests.jl

@@ -1,11 +1,12 @@
 using FastSpecSoG
 using Test
-using ExTinyMD, EwaldSummations
+using ExTinyMD, EwaldSummations, TaylorSeries
 
 @testset "FastSpecSoG.jl" begin
     include("U_series.jl")
     # include("energy_naive.jl")
     include("energy_mid.jl")
     include("interpolate_nu.jl")
     include("energy_long.jl")
+    include("cheb_interpolate.jl")
 end