From 92bf6eeb7dc9d11200cf618004493b7880303394 Mon Sep 17 00:00:00 2001 From: xzgao Date: Tue, 19 Mar 2024 20:50:01 +0800 Subject: [PATCH] add energy function --- Project.toml | 1 - src/FSSoGInteraction.jl | 39 +++++++++++++++++++++++----- src/FastSpecSoG.jl | 20 +++++++------- src/energy/energy.jl | 20 +++++++++++--- src/energy/energy_long.jl | 8 ++++++ src/energy/energy_mid.jl | 4 +++ src/energy/energy_short.jl | 4 +++ src/types.jl | 53 ++++++++++++++++++++++++-------------- test/Project.toml | 1 + test/energy.jl | 48 ++++++++++++++++++++++++++++++++++ test/energy_naive.jl | 30 --------------------- test/runtests.jl | 4 ++- 12 files changed, 162 insertions(+), 70 deletions(-) create mode 100644 test/energy.jl delete mode 100644 test/energy_naive.jl diff --git a/Project.toml b/Project.toml index c1d0e27..975a02f 100644 --- a/Project.toml +++ b/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" diff --git a/src/FSSoGInteraction.jl b/src/FSSoGInteraction.jl index dc41e63..f592120 100644 --- a/src/FSSoGInteraction.jl +++ b/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 \ No newline at end of file diff --git a/src/FastSpecSoG.jl b/src/FastSpecSoG.jl index af029d9..785a008 100644 --- a/src/FastSpecSoG.jl +++ b/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,14 +25,6 @@ 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") @@ -40,4 +32,12 @@ 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 diff --git a/src/energy/energy.jl b/src/energy/energy.jl index f73fc0f..5ee5a7b 100644 --- a/src/energy/energy.jl +++ b/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 \ No newline at end of file diff --git a/src/energy/energy_long.jl b/src/energy/energy_long.jl index b05a020..4a7b1ee 100644 --- a/src/energy/energy_long.jl +++ b/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 \ No newline at end of file diff --git a/src/energy/energy_mid.jl b/src/energy/energy_mid.jl index 112a70a..be9a23c 100644 --- a/src/energy/energy_mid.jl +++ b/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 \ No newline at end of file diff --git a/src/energy/energy_short.jl b/src/energy/energy_short.jl index d96ca52..7b43d1a 100644 --- a/src/energy/energy_short.jl +++ b/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 \ No newline at end of file diff --git a/src/types.jl b/src/types.jl index 14e9308..1c5ac58 100644 --- a/src/types.jl +++ b/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 \ No newline at end of file diff --git a/test/Project.toml b/test/Project.toml index c012f9c..8a8c659 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/energy.jl b/test/energy.jl new file mode 100644 index 0000000..ff52ae4 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/test/energy_naive.jl b/test/energy_naive.jl deleted file mode 100644 index 2d4dd20..0000000 --- a/test/energy_naive.jl +++ /dev/null @@ -1,30 +0,0 @@ -@testset "test energy naive" begin - - @info "testing the energy naive" - - 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) - - FastSpecSoG_interaction = FSSoG_naive((L, L, L), n_atoms, 10.0, 3.0, preset = 5) - FastSpecSoG_neighbor = CellList3D(info, FastSpecSoG_interaction.r_c, boundary, 1) - energy_sog = energy_naive(FastSpecSoG_interaction, FastSpecSoG_neighbor, info, atoms) - - @test abs(energy_ewald - energy_sog) < 1e-4 -end \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index 9363596..5c841fa 100644 --- a/test/runtests.jl +++ b/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")