Improve bond potentials + adapt (#49)

Author: romainljsimon

src/IO/IO.jl

@@ -124,15 +124,29 @@ function construct_bonds_array(io, N_bonds, N)
     return bond
 end
 
+filter_kwargs(pairs...) = (; (k => v for (k, v) in pairs if v !== nothing)...)
+
 function get_model(data, i::Int, j::Int)
     key = i <= j ? "$i-$j" : "$j-$i"
     m = data[key]
     if m["name"] == "GeneralKG"
-        return GeneralKG(m["epsilon"], m["sigma"], m["k"], m["r0"]; rcut = m["rcut"])
+        rcut = get(m, "rcut", nothing)
+        
+        return GeneralKG(m["epsilon"], m["sigma"], m["k"], m["r0"];
+                        filter_kwargs(
+                            :rcut => get(m, "rcut", nothing),
+                            :ϵbond => get(m, "epsilonbond", nothing),
+                            :σbond => get(m, "sigmabond", nothing),
+                            :rcutbond => get(m, "rcutbond", nothing),
+                            )...)
     elseif m["name"] == "SmoothLennardJones"
-        return SmoothLennardJones(m["epsilon"], m["sigma"]; rcut = m["rcut"])
+        return SmoothLennardJones(m["epsilon"], m["sigma"];
+                                filter_kwargs(
+                                :rcut => get(m, "rcut", nothing))...)
     elseif m["name"] == "LennardJones"
-        return LennardJones(m["epsilon"], m["sigma"]; rcut = m["rcut"])
+        return LennardJones(m["epsilon"], m["sigma"];
+                                filter_kwargs(
+                                    :rcut => get(m, "rcut", nothing))...)            
     else
         error("Model $(m["name"]) is not implemented")
         return nothing

src/models.jl

@@ -180,32 +180,45 @@ struct GeneralKG{T<:AbstractFloat} <: DiscreteModel
     name::String
     ϵ::T
     ϵ4::T
+    ϵbond::T
+    ϵ4bond::T
     σ2::T
+    σ2bond::T
     shift::T
+    shiftbond::T
     k::T
     kr02::T
     r02::T
     rcut::T
     rcut2::T
-
+    rcutbond::T
+    rcut2bond::T
 end
 
-function GeneralKG(ϵ, σ, k, r0; rcut=2^(1 / 6) * σ)
+function GeneralKG(ϵ, σ, k, r0; rcut=2^(1 / 6) * σ, ϵbond=ϵ, σbond=σ, rcutbond=rcut)
     name = "GeneralKG"
     r02 = r0 ^ 2
     rcut2 = rcut ^ 2
+    rcut2bond = rcutbond ^2
     kr02 = - k * r02 / 2
     σ2 = σ ^ 2
+    σ2bond = σbond ^ 2
     shift = lennard_jones(rcut2, 4ϵ, σ2)
-    return GeneralKG(name, ϵ, 4ϵ, σ2, shift, k, kr02, r02, rcut, rcut2)
+    shiftbond = lennard_jones(rcut2bond, 4ϵbond, σ2bond)
+    return GeneralKG(name, ϵ, 4ϵ, ϵbond, 4ϵbond, σ2, σ2bond, shift, shiftbond, k, kr02, r02, rcut, rcut2, rcutbond, rcut2bond)
 end
 
 @inline function potential(r2::T, model::GeneralKG) where T<:AbstractFloat
     return lennard_jones(r2, model.ϵ4, model.σ2) - model.shift
 end
 
 @inline function bond_potential(r2::T, model::GeneralKG) where T<:AbstractFloat
-    return r2 ≤ model.r02 ? fene(r2, model.kr02, model.r02) : Inf
+    u_fene = r2 ≤ model.r02 ? fene(r2, model.kr02, model.r02) : Inf
+    u_lj = 0.0
+    if r2 ≤ cutoff2(model)
+       u_lj += lennard_jones(r2, model.ϵ4bond, model.σ2bond) - model.shiftbond
+    end
+    return u_fene + u_lj
 end
 
 cutoff(model::DiscreteModel) = model.rcut

src/molecules.jl

@@ -105,14 +105,9 @@ end
 function compute_energy_ij(system::Molecules, position_i, position_j, model_ij::Model, ::Bonded)
     box = get_box(system)
     r2 = nearest_image_distance_squared(position_i, position_j, box)
-    energy_ij = bond_potential(r2, model_ij)
-    cutoff2_val = cutoff2(model_ij)
-    if r2 ≤ cutoff2_val
-        energy_ij += potential(r2, model_ij)
-    end
-
-    return energy_ij
+    return bond_potential(r2, model_ij)
 end
+
 function compute_energy_ij(system::Molecules, position_i, position_j, model_ij::Model, ::NonBonded)
     r2 = nearest_image_distance_squared(position_i, position_j, get_box(system))
     cutoff2_val = cutoff2(model_ij)