fixed SW energy loop

ejmeitz · Nov 10, 2023 · 8bddd88 · 8bddd88
1 parent 439345d
commit 8bddd88
Show file tree

Hide file tree

Showing 5 changed files with 72 additions and 115 deletions.
diff --git a/src/autodiff_test.jl b/src/autodiff_test.jl
@@ -44,16 +44,18 @@ end
 function second_order_AD_test(sys::SuperCellSystem{D}, pot::StillingerWeberSilicon, tol; r_cut = pot.r_cut) where D
     vars = make_variables(:r, D)
     r_norm = sqrt(sum(x -> x^2, vars))
+    println("at start")
     #2nd order part 
     pot2_symbolic = pair_potential_nounits(pot, r_norm)
     H2_symbolic = hessian(pot2_symbolic, vars)
     H2_exec = make_function(H2_symbolic, vars)
 
+    println("here")
     #3rd order part(s)
     pot3_symbolic = threebody_potential_nounits(pot, r_norm)
     H3_symbolic = hessian(pot3_symbolic, vars)
     H3_exec = make_function(H3_symbolic, vars)
-
+    println("here3")
     N_atoms = n_atoms(sys)
     IFC2 = zeros(D*N_atoms,D*N_atoms)
 

diff --git a/src/finite_diff_ifc.jl b/src/finite_diff_ifc.jl
@@ -25,7 +25,7 @@ function second_order_finite_diff(sys_eq::SuperCellSystem{3}, pot::PairPotential
             posns[atom_idxs[1]][cartesian_idxs[1]] += combo[1]
             posns[atom_idxs[2]][cartesian_idxs[2]] += combo[2]
 
-            energies[c] = energy_loop(pot, posns, pot.energy_unit, r_cut, sys_eq.box_sizes_SC, N_atoms)
+            energies[c] = energy_loop(pot, posns, r_cut, sys_eq.box_sizes_SC, N_atoms)
 
             #Un-modify
             posns[atom_idxs[1]][cartesian_idxs[1]] -= combo[1]
@@ -40,7 +40,7 @@ function second_order_finite_diff(sys_eq::SuperCellSystem{3}, pot::PairPotential
         for (c,combo) in enumerate(combos)
             posns[atom_idxs[1]][cartesian_idxs[1]] += combo[1]
 
-            energies[c] = energy_loop(pot, posns, pot.energy_unit, r_cut, sys_eq.box_sizes_SC, N_atoms)
+            energies[c] = energy_loop(pot, posns, r_cut, sys_eq.box_sizes_SC, N_atoms)
 
             posns[atom_idxs[1]][cartesian_idxs[1]] -= combo[1]
         end

diff --git a/src/helper_funcs.jl b/src/helper_funcs.jl
@@ -26,9 +26,9 @@ function apply_tols!(arr, tol)
     return arr
 end
 
-function energy_loop(pot::PairPotential, posns, eng_unit, r_cut, box_sizes, N_atoms)
+function energy_loop(pot::PairPotential, posns, r_cut, box_sizes, N_atoms)
 
-    U_total = 0.0*eng_unit
+    U_total = 0.0*pot.energy_unit
 
     for i in range(1,N_atoms)
         for j in range(i+1, N_atoms)             
@@ -45,78 +45,49 @@ function energy_loop(pot::PairPotential, posns, eng_unit, r_cut, box_sizes, N_at
 
 end
 
-function energy_loop(pot::StillingerWeberSilicon, posns, eng_unit, r_cut, box_sizes, N_atoms)
+#Expects positions as Vector{Vector}
+function energy_loop(pot::StillingerWeberSilicon, posns, box_sizes, N_atoms)
 
-    U_total = 0.0*eng_unit
+    U_total = 0.0*pot.energy_unit
 
+    rᵢⱼ = similar(posns[1])
+    rᵢₖ = similar(posns[1])
+    r_cut_sq = pot.r_cut * pot.r_cut #avoid doing sqrts
     for i in range(1,N_atoms)
-        for j in range(1, N_atoms)
-
-            rᵢⱼ = posns[i] .- posns[j]
-            nearest_mirror!(rᵢⱼ, box_sizes)
-            dist_ij = norm(rᵢⱼ) 
-
-            if i < j && dist_ij < r_cut
-                U_total += pair_potential(pot, rᵢⱼ)
-            end
-
-            for k in range(j+1, N_atoms)  
-                # i central
-                rᵢₖ = posns[i] .- posns[k]
-                nearest_mirror!(rᵢₖ, box_sizes)
-                dist_ik = norm(rᵢₖ)    
-
-                if dist_ij < r_cut && dist_ik < r_cut 
-                    U_total += three_body_potential(pot, rᵢⱼ, rᵢₖ, dist_ij, dist_ik)
+        for j in range(1,N_atoms)
+
+            if i != j
+                rᵢⱼ .= posns[i] .- posns[j]
+                nearest_mirror!(rᵢⱼ, box_sizes)
+                dist_ij_sq = sum(x -> x^2, rᵢⱼ)
+
+                if dist_ij_sq < r_cut_sq
+
+                    if i < j
+                        U_total += pair_potential(pot, sqrt(dist_ij_sq))
+                    end
+
+                    for k in range(j+1, N_atoms)
+                        if i != k
+                            rᵢₖ .= posns[i] .- posns[k]
+                            nearest_mirror!(rᵢₖ, box_sizes)
+                            dist_ik_sq = sum(x -> x^2, rᵢₖ)
+
+                            if dist_ik_sq < r_cut_sq
+                                U_total += three_body_potential(pot, rᵢⱼ, rᵢₖ, sqrt(dist_ij_sq), sqrt(dist_ik_sq))
+                            end
+                        end
+                    end
                 end
             end
+
         end
     end
     return U_total
 end
 
-# function energy_loop(pot::StillingerWeberSilicon, posns, eng_unit, r_cut, box_sizes, N_atoms)
-
-#     U_total = 0.0*eng_unit
-
-#     for i in range(1,N_atoms)
-#         for j in range(i+1, N_atoms)
-#             rᵢⱼ = posns[i] .- posns[j]
-#             nearest_mirror!(rᵢⱼ, box_sizes)
-#             dist_ij = norm(rᵢⱼ) 
-#             if dist_ij < r_cut
-#                 U_total += pair_potential(pot, rᵢⱼ)
-
-#                 for k in range(j+1, N_atoms)
-
-#                     #i as central atom
-#                     rᵢₖ = posns[i] .- posns[k]
-#                     nearest_mirror!(rᵢₖ, box_sizes)
-#                     dist_ik = norm(rᵢₖ)    
-
-#                     if dist_ik < r_cut 
-#                         U_total += three_body_potential(pot, rᵢⱼ, rᵢₖ)
-#                     end
-
-#                     #j,k as central atom
-#                     rⱼᵢ = -rᵢⱼ
-#                     rⱼₖ = rᵢₖ .- rᵢⱼ #rᵢⱼ .- rᵢₖ
-#                     rₖᵢ = -rᵢₖ
-#                     rₖⱼ = -rⱼₖ
-#                     dist_jk = norm(rⱼₖ)
-#                     if dist_jk < r_cut
-#                         U_total += three_body_potential(pot, rⱼᵢ, rⱼₖ) #j central atom
-#                         if dist_ik < r_cut
-#                             U_total += three_body_potential(pot, rₖᵢ, rₖⱼ) #k central atom 
-#                         end
-#                     end
-
-#                 end
-#             end
-#         end
-#     end
-#     return U_total
-# end
+
+
 
 # Calculates force on atom j in β direction
 function force_loop_j(pot::PairPotential, posns, force_unit, r_cut, box_sizes, N_atoms, j, β)

diff --git a/src/interactions/SW.jl b/src/interactions/SW.jl
@@ -18,85 +18,69 @@ function SW_Params(ϵ, σ, a, λ, γ, cosθ₀, A, B , p, q)
                                              A, B , p, q)
 end
 
-struct StillingerWeberSilicon{R,L,E} <: ThreeBodyPotential
+struct StillingerWeberSilicon{R,L,E,GS,LE} <: ThreeBodyPotential
     r_cut::R
     length_unit::L
     energy_unit::E
     params::SW_Params
+    gamma_sigma::GS
+    lambda_epsilon::LE
 end
 
 
 function StillingerWeberSilicon()
     #From LAMMPS Si.sw file
     sws_params = SW_Params(2.1683u"eV", 2.0951u"Å", 1.80, 21.0, 1.20,
-         -1/3, 7.049556277,  0.6022245584,  4.0,  0.0)
+        -0.333333333333, 7.049556277,  0.6022245584,  4.0,  0.0)
 
     r_cut = sws_params.a*sws_params.σ
     length_unit = unit(sws_params.σ)
     energy_unit = unit(sws_params.ϵ)
+    gamma_sigma = sws_params.σ*sws_params.γ
+    lambda_epsilon = sws_params.ϵ*sws_params.λ
 
-    return StillingerWeberSilicon{typeof(r_cut),typeof(length_unit),typeof(energy_unit)}(
-                r_cut, length_unit, energy_unit, sws_params)
+    return StillingerWeberSilicon{typeof(r_cut),typeof(length_unit),typeof(energy_unit), typeof(gamma_sigma), typeof(lambda_epsilon)}(
+                r_cut, length_unit, energy_unit, sws_params, gamma_sigma, lambda_epsilon)
 end
 
 
-function pair_potential(pot::StillingerWeberSilicon, rᵢⱼ::Vector)   
-    return Φ₂(pot.params.A, pot.params.B, pot.params.ϵ, pot.params.σ, pot.params.p, pot.params.q, pot.params.a, norm(rᵢⱼ))
+function pair_potential(pot::StillingerWeberSilicon, dist_ij)   
+    return Φ₂(pot.params.A, pot.params.B, pot.params.ϵ, pot.params.σ, pot.params.p, pot.params.q, pot.params.a, dist_ij)
 end
 
-function pair_potential_nounits(pot::StillingerWeberSilicon, rᵢⱼ::Vector)   
+function pair_potential_nounits(pot::StillingerWeberSilicon, dist_ij)   
     return Φ₂(pot.params.A, pot.params.B, ustrip(pot.params.ϵ), ustrip(pot.params.σ),
-                 pot.params.p, pot.params.q, pot.params.a, norm(rᵢⱼ))
+                 pot.params.p, pot.params.q, pot.params.a, dist_ij)
 end
 
-# function three_body_potential(pot::StillingerWeberSilicon, rᵢⱼ, rᵢₖ, rⱼᵢ, rⱼₖ, rₖᵢ, rₖⱼ)
-#     dist_ij = norm(rᵢⱼ); dist_ik = norm(rᵢₖ)
-#     dist_ji = norm(rⱼᵢ); dist_jk = norm(rⱼₖ)
-#     dist_ki = norm(rₖᵢ); dist_kj = norm(rₖⱼ)
-#     cosθᵢⱼₖ = dot(rᵢⱼ, rᵢₖ) / (dist_ij * dist_ik)
-#     cosθⱼᵢₖ = dot(rⱼᵢ, rⱼₖ) / (dist_ji * dist_jk)
-#     cosθₖᵢⱼ = dot(rₖᵢ, rₖⱼ) / (dist_ki * dist_kj)
-#     h1 = Φ₃(pot.params.γ, pot.params.ϵ, cosθᵢⱼₖ, pot.params.cosθ₀, pot.params.γ, pot.params.γ,
-#                  pot.params.σ, pot.params.σ, pot.params.a, pot.params.a, dist_ij, dist_ik)
-#     h2 = Φ₃(pot.params.γ, pot.params.ϵ, cosθⱼᵢₖ, pot.params.cosθ₀, pot.params.γ, pot.params.γ,
-#         pot.params.σ, pot.params.σ, pot.params.a, pot.params.a, dist_ji, dist_jk)
-#     h3 = Φ₃(pot.params.γ, pot.params.ϵ, cosθₖᵢⱼ, pot.params.cosθ₀, pot.params.γ, pot.params.γ,
-#         pot.params.σ, pot.params.σ, pot.params.a, pot.params.a, dist_ki, dist_kj)
-#     # println("dist_ij $(dist_ij) dist_ik $(dist_ik) dist_ji $(dist_ji) dist_jk $(dist_jk) dist_ki $(dist_ki) dist_kj $(dist_kj)")
-#     # println("h1: $(ustrip(h1)) h2 $(ustrip(h2)) h3 $(ustrip(h3))")
-#     # println("cos1: $(ustrip(cosθᵢⱼₖ)) cos2 $(ustrip(cosθⱼᵢₖ)) cos3 $(ustrip(cosθₖᵢⱼ))")
-#     return h1 + h2 + h3
-# end
+
 function three_body_potential(pot::StillingerWeberSilicon, rᵢⱼ, rᵢₖ, dist_ij, dist_ik)
-
+    
     cosθᵢⱼₖ = dot(rᵢⱼ, rᵢₖ) / (dist_ij * dist_ik)
 
-    return Φ₃(pot.params.γ, pot.params.ϵ, cosθᵢⱼₖ, pot.params.cosθ₀, pot.params.γ, pot.params.γ,
-                 pot.params.σ, pot.params.σ, pot.params.a, pot.params.a, dist_ij, dist_ik)
+    return Φ₃_si(pot, cosθᵢⱼₖ, pot.r_cut, dist_ij, dist_ik)
 
 end
 
-# function three_body_potential_nounits(pot::StillingerWeberSilicon, rᵢⱼ, rᵢₖ, rⱼᵢ, rⱼₖ, rₖᵢ, rₖⱼ)
-#     dist_ij = norm(rᵢⱼ); dist_ik = norm(rᵢₖ)
-#     dist_ji = norm(rⱼᵢ); dist_jk = norm(rⱼₖ)
-#     dist_ki = norm(rₖᵢ); dist_kj = norm(rₖⱼ)
-#     cosθᵢⱼₖ = dot(rᵢⱼ, rᵢₖ) / (dist_ij * dist_ik)
-#     cosθⱼᵢₖ = dot(rⱼᵢ, rⱼₖ) / (dist_ji * dist_jk)
-#     cosθₖᵢⱼ = dot(rₖᵢ, rₖⱼ) / (dist_ki * dist_kj)
-#     h1 = Φ₃(pot.params.γ, ustrip(pot.params.ϵ), cosθᵢⱼₖ, pot.params.cosθ₀, pot.params.γ, pot.params.γ,
-#                  ustrip(pot.params.σ), ustrip(pot.params.σ), pot.params.a, pot.params.a, dist_ij, dist_ik)
-#     h2 = Φ₃(pot.params.γ, ustrip(pot.params.ϵ), cosθⱼᵢₖ, pot.params.cosθ₀, pot.params.γ, pot.params.γ,
-#         ustrip(pot.params.σ), ustrip(pot.params.σ), pot.params.a, pot.params.a, dist_ji, dist_jk)
-#     h3 = Φ₃(pot.params.γ, ustrip(pot.params.ϵ), cosθₖᵢⱼ, pot.params.cosθ₀, pot.params.γ, pot.params.γ,
-#         ustrip(pot.params.σ), ustrip(pot.params.σ), pot.params.a, pot.params.a, dist_ki, dist_kj)
-#     return h1 + h2 + h3
-# end
+function three_body_potential_nounits(pot::StillingerWeberSilicon, rᵢⱼ, rᵢₖ, dist_ij, dist_ik)
+
+    cosθᵢⱼₖ = dot(rᵢⱼ, rᵢₖ) / (dist_ij * dist_ik)
+
+    return Φ₃(pot.params.λ, ustrip(pot.params.ϵ), cosθᵢⱼₖ, pot.params.cosθ₀, pot.params.γ, pot.params.γ, 
+        ustrip(pot.params.σ), ustrip(pot.params.σ), pot.params.a, pot.params.a, rᵢⱼ, rᵢₖ)
+
+end
+
+
 
 function Φ₂(A, B, ϵ, σ, p, q, a, rᵢⱼ)
     return A*ϵ*(B*((σ/rᵢⱼ)^p) - ((σ/rᵢⱼ)^q)) * exp(σ/(rᵢⱼ-(a*σ)))
 end
 
 function Φ₃(λᵢⱼₖ, ϵᵢⱼₖ, cosθᵢⱼₖ, cosθ₀, γᵢⱼ, γᵢₖ, σᵢⱼ, σᵢₖ, aᵢⱼ, aᵢₖ, rᵢⱼ, rᵢₖ)
-    # println("term1: $(ustrip(λᵢⱼₖ*ϵᵢⱼₖ*((cosθᵢⱼₖ - cosθ₀)^2))) term2 $(ustrip(exp((γᵢⱼ*σᵢⱼ)/(rᵢⱼ -(aᵢⱼ*σᵢⱼ))))) term3 $(ustrip(exp((γᵢₖ*σᵢₖ)/(rᵢₖ-(aᵢₖ*σᵢₖ)))))")
-    return λᵢⱼₖ*ϵᵢⱼₖ*((cosθᵢⱼₖ - cosθ₀)^2)*exp((γᵢⱼ*σᵢⱼ)/(rᵢⱼ -(aᵢⱼ*σᵢⱼ)))*exp((γᵢₖ*σᵢₖ)/(rᵢₖ-(aᵢₖ*σᵢₖ)))
+    return λᵢⱼₖ*ϵᵢⱼₖ*((cosθᵢⱼₖ - cosθ₀)^2)*exp((γᵢⱼ*σᵢⱼ)/(rᵢⱼ-(aᵢⱼ*σᵢⱼ)))*exp((γᵢₖ*σᵢₖ)/(rᵢₖ-(aᵢₖ*σᵢₖ)))
+end 
+
+function Φ₃_si(pot, cosθᵢⱼₖ, r_cut, rᵢⱼ, rᵢₖ)
+    return pot.lambda_epsilon*((cosθᵢⱼₖ - pot.params.cosθ₀)^2)*exp(pot.gamma_sigma/(rᵢⱼ-r_cut))*exp(pot.gamma_sigma/(rᵢₖ-r_cut))
 end 
diff --git a/src/mcc.jl b/src/mcc.jl
@@ -2,7 +2,7 @@ export mcc3, mcc3_custom_kernel
 
 
 """
-Converts third order forces constants, `Ψ` into third order modal coupling constants (MCC).
+Converts mass weighted third order forces constants, `Ψ` into third order modal coupling constants (MCC).
 Does not divide MCC calculation into smaller chunks. This might exhaust GPU memory.
 """
 function mcc3(Ψ::CuArray{Float32, 3}, phi::CuArray{Float32, 2}, tol::Float64)
@@ -19,7 +19,7 @@ function mcc3(Ψ::CuArray{Float32, 3}, phi::CuArray{Float32, 2}, tol::Float64)
 end
 
 """
-Converts third order forces constants, `Ψ` into third order modal coupling constants (MCC). The
+Converts mass weighted third order forces constants, `Ψ` into third order modal coupling constants (MCC). The
 parameter `block_size` specifies problem size when calculating the MCC. For example, if 
 `block_size` is 100, the block MCC will be calculated in 100x100x100 blocks to save GPU memory.
 The `phi` matrix will be automatically truncated to adjust for this.