Particle Mesh Ewald

src/long_range_interactions/SPME/PME_TODO.txt

@@ -0,0 +1,11 @@
+- dQdR calculation in charge interpolation never worked (GPU & CPU)
+    - required for force calculation
+- Parallelize charge interpolation on CPU
+- remove repeated allocations for things in called force loop
+- add LongRangeInteraction type?, SPME is a pair interaction though
+    - figure out how to adapt run! function to Molly potential interface
+- write kernel to calculate the scaled fractional coords for charge interpolation on GPU
+
+- bundle this into a package extension? would separate FFTW and potentially Adapt.jl
+
+- do particle-particle part with Molly existing infastructure

src/long_range_interactions/SPME/SPME.jl

@@ -0,0 +1,96 @@
+function vol(lat_vecs::Vector{Vector{T}}) where T
+    return dot(lat_vecs[1], cross(lat_vecs[2], lat_vecs[3]))
+end
+
+function reciprocal_vecs(lat_vecs::Vector{Vector{T}}) where T
+    V = vol(lat_vecs)
+    m1 = cross(lat_vecs[2], lat_vecs[3])/V
+    m2 = cross(lat_vecs[3], lat_vecs[1])/V
+    m3 = cross(lat_vecs[1], lat_vecs[2])/V
+    return [m1,m2,m3]
+end
+
+function reciprocal_vecs_twopi(lat_vecs)
+    V = vol(lat_vecs)
+    m1 = 2*π*cross(lat_vecs[2], lat_vecs[3])/V
+    m2 = 2*π*cross(lat_vecs[3], lat_vecs[1])/V
+    m3 = 2*π*cross(lat_vecs[1], lat_vecs[2])/V
+    return [m1,m2,m3]
+end
+
+
+struct SPME{T, R, B, E, C, L}
+    sys::System
+    error_tol::T
+    r_cut_dir::R
+    β::B
+    self_energy::E
+    BC::Array{C,3}
+    K::SVector{3,Integer}
+    spline_order::Integer
+    recip_lat::Vector{SVector{3,L}}
+end
+
+function SPME(sys, error_tol, r_cut_dir, spline_order)
+
+    β = sqrt(-log(2*error_tol))/r_cut_dir
+    box_sizes = norm.(lattice_vec(sys))
+    K = ceil.(Int, 2*β.*box_sizes./(3*(error_tol ^ 0.2))) 
+
+    recip_lat = reciprocal_vecs(sys.lattice_vec)
+
+    # Calculate self-energy
+    self_energy = -(β/sqrt(π))*sum(x -> x*x, charges(sys))
+    # Calcualte pre-factors for mesh component of energy/force
+    C_type = Complex{float_type(sys)}
+    BC = zeros(C_type, K)
+    calc_BC!(BC, spme)
+
+    return SPME{typeof(error_tol), typeof(r_cut_dir), typeof(β), typeof(self_energy), C_type, eltype(recip_lat[1])}(
+                sys, error_tol, r_cut_dir, β, self_energy, K,
+                spline_order, recip_lat)
+end
+
+reciprocal_lattice(spme::SPME) = spme.recip_lat
+self_energy(spme::SPME) = spme.self_energy
+n_mesh(spme::SPME) = spme.K
+spline_order(spme::SPME) = spme.spline_order
+
+function run!(spme::SPME)
+
+    #& could write to reuse storage for F
+    E_dir, F_dir = particle_particle(spme) # Uses Molly pair-pair
+    E_rec, F_rec = particle_mesh(spme)
+    E_self = self_energy(spme) #calculated on construction, re-use storage
+
+    #* Need to convert units here, try not to hard code to specific unit system
+    #* Probably best to just use Unitful and multiply E by 1/4πϵ₀
+
+    E_SPME = E_dir + E_rec + E_self
+    F_SPME = F_dir .+ F_rec
+
+    return E_SPME, F_SPME
+
+end
+
+
+# #& GPU version of Q
+# u2 = [Vector{Float64}(undef, (length(n_mesh(spme)), )) for _ in eachindex(positions)];
+# u2 = scaled_fractional_coords!(u2, spme.sys.positions, n_mesh(spme), spme.recip_lat);
+# M0, M1, M2, _, _, _ = calc_spline_values(u2, n, N_atoms);
+# cuQ = CUDA.zeros(Float32,n_mesh(spme)...);
+
+# cuM0 = CuArray{Float32}(M0);
+# cuM1 = CuArray{Float32}(M1);
+# cuM2 = CuArray{Float32}(M2);
+# n_half = ceil(Int64,n/2);
+# cu_u = CuArray{Float32}(reduce(hcat, u2)'); #try transposing
+# cuCharges = CuArray{Int32}(spme.sys.atoms.charge);
+# BC_cuda = CuArray{Float32}(BC)
+
+# thread_per_block = 64
+# N_blocks = ceil(Int64, N_atoms/thread_per_block)
+
+
+# @cuda threads=thread_per_block blocks=N_blocks interpolate_charge_kernel!(cu_u, cuM0, cuM1, cuM2, cuQ, 
+#     n_half,cuCharges, n_mesh(spme)..., n, N_atoms)

src/long_range_interactions/SPME/bsplines.jl

@@ -0,0 +1,84 @@
+export calc_BC
+
+function M(u, n)
+    if n > 2
+        return (u/(n-1))*M(u,n-1) + ((n-u)/(n-1))*M(u-1,n-1)
+    elseif n == 2
+        if u >= 0 && u <= 2
+            return 1 - abs(u-1)
+        else
+            return 0
+        end
+    else
+        println("Shouldn't be here")
+    end
+end
+
+function dMdu(u,n)
+    return M(u, n-1) - M(u-1, n-1)
+end
+
+function b(mⱼ,n,Kⱼ)
+    if iseven(n) && (2*abs(mⱼ) == Kⱼ) #interp fails in this case
+        return 0.0
+    end
+
+    m_K = 2*π*mⱼ/Kⱼ
+    v = m_K*(n-1)
+    num = cos(v) + 1im*sin(v)
+    denom = 0.0 + 0.0im
+    for k in 0:n-2
+        v2 = m_K*k
+        denom += M(k+1, n)*(cos(v2) + 1im*sin(v2))
+    end
+    return  num/denom 
+end
+
+function calc_C(β, V, ms, recip_lat)
+    m_star = ms[1].*recip_lat[1] .+ ms[2].*recip_lat[2] .+ ms[3].*recip_lat[3]
+    m_sq = dot(m_star,m_star)
+    return (1/(π*V))*(exp(-(π^2)*m_sq/(β^2))/m_sq)
+end
+
+
+function calc_BC!(BC::Array{T,3}, spme::SPME) where T
+    V = vol(spme.sys)
+    K1, K2, K3 = n_mesh(spme)
+    recip_lat = reciprocal_lattice(spme) #SPME uses the 1 normalized version
+    n = spline_order(spme)
+
+    #Gather h indices 
+    h1s = collect(0:(K1-1))
+    h1s[h1s .> K1/2] .-= K1 #* this has allocation
+
+    h2s = collect(0:(K2-1))
+    h2s[h2s .> K2/2] .-= K2
+
+    h3s = collect(0:(K3-1))
+    h3s[h3s .> K3/2] .-= K3
+
+    B1 = abs2.(b.(0:K1-1, n, K1))
+    B2 = abs2.(b.(0:K2-1, n, K2))
+    B3 = abs2.(b.(0:K3-1, n, K3))
+
+    hs = [0,0,0]
+    for m1 in range(0,K1-1)
+        hs[1] = h1s[m1+1]
+        for m2 in range(0,K2-1)
+            hs[2] = h2s[m2+1]
+            for m3 in range(0,K3-1) 
+                hs[3] = h3s[m3+1]        
+                C = calc_C(spme.β, V, hs, recip_lat)
+                BC[m1+1,m2+1,m3+1] = B1[m1+1]*B2[m2+1]*B3[m3+1]*C
+            end
+        end
+    end
+
+    BC[1,1,1] = T(0.0)
+
+    return BC
+end
+
+#equivalent, time to see whats faster
+# def M(u, n):
+#     return (1/math.factorial(n-1)) * np.sum([((-1)**k)*math.comb(n,k)*np.power(max(u-k, 0), n-1) for k in range(n+1)])

src/long_range_interactions/SPME/charge_interpolation.jl

@@ -0,0 +1,164 @@
+export interpolate_charge!, interpolate_charge2!, interpolate_charge_kernel!,calc_spline_values
+
+#Assumes Vector of Vectors for r and recip_vectors
+function scaled_fractional_coords!(u, r, n_mesh::AbstractVector, recip_vectors)
+
+    Threads.@threads for i in eachindex(r)
+        for dim in eachindex(n_mesh)
+            u[i][dim] = n_mesh[dim]*dot(recip_vectors[dim], r[i])
+        end
+    end
+    return u
+end
+
+function calc_spline_values(u, n, N_atoms)
+
+    #For each atom & dir there are n spline values
+    M0 = zeros(N_atoms, n+1)
+    dM0 = zeros(N_atoms, n+1)
+    M1 = zeros(N_atoms, n+1)
+    dM1 = zeros(N_atoms, n+1)
+    M2 = zeros(N_atoms, n+1)
+    dM2 = zeros(N_atoms, n+1)
+    Threads.@threads for i in range(1,N_atoms)
+        for c in range(0,n) #&either the zero or n case is not used???
+            l0 = round(u[i][1]) - c
+            l1 = round(u[i][2]) - c
+            l2 = round(u[i][3]) - c
+
+            M0[i,c+1] = M(u[i][1] - l0, n)
+            dM0[i,c+1] = dMdu(u[i][1] - l0, n)
+            M1[i,c+1] = M(u[i][2] - l1, n)
+            dM1[i,c+1] = dMdu(u[i][2] - l1, n)
+            M2[i,c+1] = M(u[i][3] - l2, n)
+            dM2[i,c+1] = dMdu(u[i][3] - l2, n)
+        end
+    end
+    return M0, M1, M2, dM0, dM1, dM2
+end
+
+
+function interpolate_charge!(u, Q, dQdr, spme::SPME{SingleThread})
+    K1,K2,K3 = n_mesh(spme)
+    recip_lat = reciprocal_lattice(spme)
+    q_arr = charges(spme.sys)
+    N_atoms = length(q_arr)
+    n = spme.spline_order
+
+    u = scaled_fractional_coords!(u, positions(spme.sys), n_mesh(spme), recip_lat)
+
+
+    for i in 1:N_atoms
+        for c0 in 0:n
+            l0 = round(Int64,u[i][1]) - c0 # Grid point to interpolate onto
+
+            M0 = M(u[i][1] - l0, n)
+            q_n_0 = q_arr[i]*M0 #if 0 <= u_i0 - l0 <= n will be non-zero
+            dM0 = dMdu(u[i][1] - l0,n)
+
+            l0 += ceil(Int64,n/2) # Shift
+            if l0 < 0 # Apply PBC
+                l0 += K1
+            elseif l0 >= K1
+                l0 -= K1
+            end
+
+            for c1 in 0:n
+                l1 = round(Int64,u[i][2]) - c1 # Grid point to interpolate onto
+
+                M1 = M(u[i][2] - l1, n)
+                q_n_1 = q_n_0*M1 #if 0 <= u_i1 - l1 <= n will be non-zero
+                dM1 = dMdu(u[i][2] - l1,n)
+
+
+                l1 += ceil(Int64,n/2) # Shift
+                if l1 < 0 # Apply PBC
+                    l1 += K2
+                elseif l1 >= K2
+                    l1 -= K2
+                end
+
+                for c2 in 0:n
+                    l2 = round(Int64,u[i][3]) - c2 # Grid point to interpolate onto
+
+                    M2 = M(u[i][3] - l2, n)
+                    q_n_2 = q_n_1*M2 #if 0 <= u_i1 - l1 <= n will be non-zero
+                    dM2 = dMdu(u[i][3] - l2,n)
+
+                    l2 += ceil(Int64,n/2) # Shift
+                    if l2 < 0 # Apply PBC
+                        l2 += K3
+                    elseif l2 >= K3
+                        l2 -= K3
+                    end
+
+                    Q[l0+1,l1+1,l2+1] += q_arr[i]*M0*M1*M2
+
+                    #*Does it matter that l0,l1,l2 is also a function of r_ia
+                    #*This looks like its probably equivalent to some matrix multiply
+                    # dQdr[i, 1, l0+1, l1+1, l2+1] = q_arr[i]*(K1*recip_lat[1][1]*dM0*M1*M2 + K2*recip_lat[2][1]*dM1*M0*M2 + K3*recip_lat[3][1]*dM2*M0*M1)
+                    # dQdr[i, 2, l0+1, l1+1, l2+1] = q_arr[i]*(K1*recip_lat[1][2]*dM0*M1*M2 + K2*recip_lat[2][2]*dM1*M0*M2 + K3*recip_lat[3][2]*dM2*M0*M1)
+                    # dQdr[i, 3, l0+1, l1+1, l2+1] = q_arr[i]*(K1*recip_lat[1][3]*dM0*M1*M2 + K2*recip_lat[2][3]*dM1*M0*M2 + K3*recip_lat[3][3]*dM2*M0*M1)
+
+                end
+            end
+        end
+    end
+
+    return Q, dQdr
+end
+
+function interpolate_charge_kernel!(u, Mx_vals, My_vals, Mz_vals, Q, 
+    n_half, charges, K1, K2, K3, n, N_atoms)
+
+
+    i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x
+
+
+    if i <= N_atoms
+        #Load atom data into thread registers
+        u_ix = round(Int32,u[i,1])
+        u_iy = round(Int32,u[i,2])
+        u_iz = round(Int32,u[i,3])
+        q = charges[i]
+        l0 = 0i32; l1 = 0i32; l2 = 0i32;
+        for c0 in 0:n
+            l0 = u_ix - c0 # Grid point to interpolate onto
+
+
+            l0 += n_half # Shift
+            if l0 < 0 # Apply PBC
+                l0 += K1
+            elseif l0 >= K1
+                l0 -= K1
+            end
+
+            for c1 in 0:n
+                l1 = u_iy - c1 # Grid point to interpolate onto
+
+                l1 += n_half # Shift
+                if l1 < 0 # Apply PBC
+                    l1 += K2
+                elseif l1 >= K2
+                    l1 -= K2
+                end
+
+                for c2 in 0:n
+                    l2 = u_iz - c2 # Grid point to interpolate onto
+
+                    l2 += n_half # Shift
+                    if l2 < 0 # Apply PBC
+                        l2 += K3
+                    elseif l2 >= K3
+                        l2 -= K3
+                    end
+
+                    v = q*Mx_vals[i,c0+1]*My_vals[i,c1+1]*Mz_vals[i,c2+1]
+                    CUDA.@atomic Q[l0+1,l1+1,l2+1] += v
+                end
+            end
+        end
+    end
+
+end
+

src/long_range_interactions/SPME/particle_mesh.jl

@@ -0,0 +1,10 @@
+function particle_mesh(spme::SPME, cuQ)
+    # Q_inv = fft(cuQ)
+    # Q_inv .*= BC_cuda
+    # Q_inv = ifft!(Q_inv) # Q_conv_theta, but do in place
+
+    # A = 332.0637132991921
+    # E = 0.5 * A* sum(real(Q_inv) .* real(cuQ))
+    # F = None
+    # return E, F
+end