add energy function

Project.toml

@@ -5,7 +5,6 @@ version = "1.0.0-DEV"
 
 [deps]
 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"

src/FSSoGInteraction.jl

@@ -38,15 +38,42 @@ function FSSoG_naive(L::NTuple{3, T}, n_atoms::Int, r_c::T, k_c::T; preset::Int
 end
 
 function FSSoGInteraction(
-    L::NTuple{3, T}, n_atoms::Int, 
-    n_k::NTuple{3, Int}, h::T, win_cut::Int, win_width::T, β::T,; 
+    b::T, σ::T, ω::T, M::Int,
+    L::NTuple{3, T}, n_atoms::Int,
+    r_c::T, Q_short::Int, r_min::T,
+    N_real_mid::NTuple{3, Int}, w::NTuple{3, Int}, β::NTuple{3, T}, extra_pad_ratio::Int, cheb_order::Int, M_mid::Int, 
+    N_grid_long::NTuple{3, Int}, Q_long::Int, soepara::SoePara{Complex{T}}; 
+    ϵ::T = one(T)) where{T}
+
+    uspara = USeriesPara(b, σ, ω, M)
+
+    position = [tuple(zeros(T, 3)...) for i in 1:n_atoms]
+    charge = zeros(T, n_atoms)
+
+    boundary = Q2dBoundary(L...)
+
+    uspara_cheb = Es_USeries_Cheb(uspara, r_min, r_c, Q_short)
+    F0 = F0_cal(b, σ, ω, M)
+
+    gridinfo, gridbox, cheb_coefs, scalefactor = mid_paras_gen(N_real_mid, w, β, L, extra_pad_ratio, cheb_order, uspara, M_mid)
+
+    # some of these term are for other methods, not used
+    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_long, uspara, M_mid, n_atoms)
+    cheb_mat = ChebUSeries(k_x, k_y, L[3], uspara, M_mid, Q_long)
+
+    return FSSoGInteraction{T}(b, σ, ω, M, ϵ, L, boundary, n_atoms, uspara, position, charge, uspara_cheb, r_c, F0, gridinfo, gridbox, cheb_coefs, scalefactor, 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, cheb_mat, z, sort_z, soepara)
+end
+
+function FSSoGInteraction(
+    L::NTuple{3, T}, n_atoms::Int,
+    r_c::T, Q_short::Int, r_min::T,
+    N_real_mid::NTuple{3, Int}, w::NTuple{3, Int}, β::NTuple{3, T}, extra_pad_ratio::Int, cheb_order::Int, M_mid::Int, 
+    N_grid_long::NTuple{3, Int}, Q_long::Int, soepara::SoePara{Complex{T}}; 
     preset::Int = 1, ϵ::T = one(T)) where{T}
+
     @assert preset ≤ length(preset_parameters)
 
     b, σ, ω, M = T(preset_parameters[preset][1]), T(preset_parameters[preset][2]), T(preset_parameters[preset][3]), Int(preset_parameters[preset][4])
-    uspara = USeriesPara(b, σ, ω, M)
-
 
-
-    return 
+    return FSSoGInteraction(b, σ, ω, M, L, n_atoms, r_c, Q_short, r_min, N_real_mid, w, β, extra_pad_ratio, cheb_order, M_mid, N_grid_long, Q_long, soepara, ϵ = ϵ)
 end

src/FastSpecSoG.jl

@@ -5,12 +5,12 @@ using ExTinyMD, LinearAlgebra, SpecialFunctions, ChebParticleMesh, SumOfExpVPMR,
 import FastChebInterp: ChebPoly
 
 export USeriesPara, U_series, BSA, ChebUSeries
-export FSSoG_naive
+export FSSoG_naive, FSSoGInteraction
 export short_energy_naive, long_energy_naive, energy_naive
 
 export Es_Cheb_precompute, short_energy_Cheb
 
-export mid_paras_gen, energy_mid
+export mid_paras_gen, energy_mid, energy_long, energy_short
 export FFCT_precompute, boundaries!, inverse_mat
 export interpolate_nu_einsum!, interpolate_nu_cheb!, interpolate_nu_loop!
 export real2Cheb!, solve_eqs!, gather_nu
@@ -25,19 +25,19 @@ include("types.jl")
 include("U_series.jl")
 include("FSSoGInteraction.jl")
 
-
-include("energy/energy.jl")
-include("energy/energy_short_naive.jl")
-include("energy/energy_short.jl")
-include("energy/energy_long_naive.jl")
-include("energy/energy_mid.jl")
-include("energy/energy_long.jl")
-
 include("FFCT/precompute.jl")
 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")
 
+include("energy/energy_short.jl")
+include("energy/energy_mid.jl")
+include("energy/energy_long.jl")
+
+include("energy/energy.jl")
+include("energy/energy_short_naive.jl")
+include("energy/energy_long_naive.jl")
+
 end

src/energy/energy.jl

@@ -3,8 +3,22 @@ function energy_naive(interaction::FSSoG_naive{T}, neighbor::CellList3D{T}, info
     charge = [atoms[info.particle_info[i].id].charge for i in 1:interaction.n_atoms]
     position = [info.particle_info[i].position.coo for i in 1:interaction.n_atoms]
 
-    energy_long = long_energy_naive(interaction, position, charge)
-    energy_short = short_energy_naive(interaction, neighbor, position, charge)
+    E_long = long_energy_naive(interaction, position, charge)
+    E_short = short_energy_naive(interaction, neighbor, position, charge)
 
-    return energy_long + energy_short
+    return E_long + E_short
+end
+
+function ExTinyMD.energy(interaction::FSSoGInteraction{T}, neighbor::CellList3D{T}, info::SimulationInfo{T}, atoms::Vector{Atom{T}}) where{T}
+
+    for i in 1:interaction.n_atoms
+        interaction.charge[i] = atoms[info.particle_info[i].id].charge
+        interaction.position[i] = info.particle_info[i].position.coo
+    end
+
+    E_long = energy_long(interaction)
+    E_mid = energy_mid(interaction)
+    E_short = energy_short(interaction, neighbor)
+
+    return E_long + E_mid + E_short
 end

src/energy/energy_long.jl

@@ -94,4 +94,12 @@ function energy_long_0(
     E_0 = zeroth_order(qs, z, soepara, uspara, sort_z, L, M_mid)
 
     return E_0
+end
+
+function energy_long(interaction::FSSoGInteraction{T}) where{T}
+
+    E_k = energy_long_cheb_k(interaction.charge, interaction.position, interaction.L, interaction.M_mid, interaction.k_x, interaction.k_y, interaction.r_z, interaction.phase_x, interaction.phase_y, interaction.us_mat, interaction.b_l, interaction.b_u, interaction.rhs, interaction.sol, interaction.ivsm, interaction.H_r, interaction.H_c, interaction.H_s, interaction.uspara, interaction.cheb_mat)
+    E_0 = energy_long_0(interaction.charge, interaction.position, interaction.L, interaction.M_mid, interaction.z, interaction.sort_z, interaction.uspara, interaction.soepara)
+
+    return (E_k + E_0) / (4π * interaction.ϵ)
 end

src/energy/energy_mid.jl

@@ -38,4 +38,8 @@ function energy_mid(qs::Vector{T}, poses::Vector{NTuple{3, T}}, gridinfo::GridIn
     E = gather(qs, poses, gridinfo, gridbox, cheb_coefs) / 2
 
     return E
+end
+
+function energy_mid(interaction::FSSoGInteraction{T}) where{T}
+    return energy_mid(interaction.charge, interaction.position, interaction.gridinfo, interaction.gridbox, interaction.cheb_coefs, interaction.scalefactor) / (4π * interaction.ϵ)
 end

src/energy/energy_short.jl

@@ -54,4 +54,8 @@ function short_energy_Cheb(uspara_cheb::ChebPoly{1, T, T}, r_c::T, F0::T, bounda
     energy_short += Es_self(q, F0)
 
     return energy_short / 4π
+end
+
+function energy_short(interaction::FSSoGInteraction{T}, neighbor::CellList3D{T}) where{T}
+    return short_energy_Cheb(interaction.uspara_cheb, interaction.r_c, interaction.F0, interaction.boundary, neighbor, interaction.position, interaction.charge) / interaction.ϵ
 end

src/types.jl

@@ -23,29 +23,44 @@ mutable struct FSSoGInteraction{T} <: ExTinyMD.AbstractInteraction
     M::Int
     ϵ::T
     L::NTuple{3, T}
-    uspara::USeriesPara{T}
+    boundary::Boundary{T}
     n_atoms::Int
 
-    # para for 3D FFT
-    n_k::NTuple{3, Int}
-    h::NTuple{3, T}
-    win_cut::Int
-    win_width::NTuple{3, T} # (win_cut + 0.5) * h
-    β::T # 5 * win_cut
+    uspara::USeriesPara{T}
+    position::Vector{NTuple{3, T}}
+    charge::Vector{T}
 
-    pad_ratio::Int
-    pad_cut::Int # pad_ratio * win_cut
-    pad_width::T # pad_cut * h
+    # for short range
+    uspara_cheb::ChebPoly{1, T, T}
+    r_c::T
+    F0::T
 
-    kx::Vector{T}
-    ky::Vector{T}
-    kz::Vector{T}
-    TdK::Array{T, 3}
-    grid_extended::Array{T, 3}
-    grid_real::Array{T, 3}
-    Pkbx::Array{T, 2}
-    Pkby::Array{T, 2}
-    Pkbz::Array{T, 2}
+    # for 3D FFT
+    gridinfo::GridInfo{3, T}
+    gridbox::GridBox{3, T}
+    cheb_coefs::NTuple{3, ChebCoef{T}}
+    scalefactor::ScalingFactor{3, T}
 
+    # for FFCT (using the energy_long_cheb as defalut)
+    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}
+    cheb_mat::Array{ChebPoly{1, T, T}, 2}
 
+    # for FFCT zero order term
+    z::Vector{T}
+    sort_z::Vector{Int}
+    soepara::SoePara{Complex{T}}
 end

test/Project.toml

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

test/energy.jl

@@ -0,0 +1,48 @@
+@testset "energy" begin
+
+    @info "testing the energy"
+
+    n_atoms = 100
+    L = 100.0
+    boundary = ExTinyMD.Q2dBoundary(L, L, L)
+
+    atoms = Vector{Atom{Float64}}()
+    for i in 1:n_atoms÷2
+        push!(atoms, Atom(type = 1, mass = 1.0, charge = 2.0))
+    end
+
+    for i in n_atoms÷2 + 1 : n_atoms
+        push!(atoms, Atom(type = 2, mass = 1.0, charge = - 2.0))
+    end
+
+    info = SimulationInfo(n_atoms, atoms, (0.0, L, 0.0, L, 0.0, L), boundary; min_r = 1.0, temp = 1.0)
+
+    Ewald2D_interaction = Ewald2DInteraction(n_atoms, 5.0, 0.2, (L, L, L), ϵ = 1.0)
+    Ewald2D_neighbor = CellList3D(info, Ewald2D_interaction.r_c, boundary, 1)
+    energy_ewald = energy(Ewald2D_interaction, Ewald2D_neighbor, info, atoms)
+
+    fssog_naive = FSSoG_naive((L, L, L), n_atoms, 10.0, 2.0, preset = 2)
+    fssog_neighbor = CellList3D(info, fssog_naive.r_c, boundary, 1)
+    energy_sog_naive = energy_naive(fssog_naive, fssog_neighbor, info, atoms)
+
+    N_real = (128, 128, 128)
+    w = (16, 16, 16)
+    β = 5.0 .* w
+    extra_pad_ratio = 2
+    cheb_order = 10
+    preset = 2
+    M_mid = 3
+
+    N_grid = (16, 16, 31)
+    Q = 24
+    r_c = 10.0
+
+    fssog_interaction = FSSoGInteraction((L, L, L), n_atoms, r_c, Q, 0.5, N_real, w, β, extra_pad_ratio, cheb_order, M_mid, N_grid, Q, SoePara16(); preset = preset, ϵ = 1.0)
+
+    fssog_neighbor = CellList3D(info, fssog_interaction.r_c, fssog_interaction.boundary, 1)
+    energy_sog = ExTinyMD.energy(fssog_interaction, fssog_neighbor, info, atoms)
+
+    @test abs(energy_ewald - energy_sog_naive) < 1e-4
+    @test abs(energy_ewald - energy_sog) < 1e-4
+    @test abs(energy_sog_naive - energy_sog) < 1e-6
+end

test/runtests.jl

@@ -1,10 +1,12 @@
 using FastSpecSoG
 using Test
 using ExTinyMD, EwaldSummations, TaylorSeries
+using Random
+Random.seed!(1234)
 
 @testset "FastSpecSoG.jl" begin
     include("U_series.jl")
-    # include("energy_naive.jl")
+    include("energy.jl")
     include("energy_short.jl")
     include("energy_mid.jl")
     include("interpolate_nu.jl")