From 223915d077cfab938356d1398f26f64f04bbf4d1 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Tue, 25 Jun 2024 12:49:50 +0100 Subject: [PATCH 01/74] no branching for cutoffs --- src/cutoffs.jl | 36 ++++++++++++------------------------ 1 file changed, 12 insertions(+), 24 deletions(-) diff --git a/src/cutoffs.jl b/src/cutoffs.jl index bade01c03..542d84bcf 100644 --- a/src/cutoffs.jl +++ b/src/cutoffs.jl @@ -150,19 +150,13 @@ function force_divr_with_cutoff(inter, r2, params, cutoff::C, coord_i::SVector{D if cutoff_points(C) == 0 return force_divr(inter, r2, inv(r2), params) elseif cutoff_points(C) == 1 - if r2 > cutoff.sqdist_cutoff - return zero(inv(oneunit(T))) * force_units - else - return force_divr_cutoff(cutoff, r2, inter, params) - end + return force_divr_cutoff(cutoff, r2, inter, params) * (r2 <= cutoff.sqdist_cutoff) elseif cutoff_points(C) == 2 - if r2 > cutoff.sqdist_cutoff - return zero(inv(oneunit(T))) * force_units - elseif r2 < cutoff.sqdist_activation - return force_divr(inter, r2, inv(r2), params) - else - return force_divr_cutoff(cutoff, r2, inter, params) - end + return ifelse( + r2 < cutoff.sqdist_activation, + force_divr(inter, r2, inv(r2), params), + force_divr_cutoff(cutoff, r2, inter, params) * (r2 <= cutoff.sqdist_cutoff), + ) end end @@ -171,18 +165,12 @@ function potential_with_cutoff(inter, r2, params, cutoff::C, coord_i::SVector{D, if cutoff_points(C) == 0 return potential(inter, r2, inv(r2), params) elseif cutoff_points(C) == 1 - if r2 > cutoff.sqdist_cutoff - return ustrip(zero(T)) * energy_units - else - return potential_cutoff(cutoff, r2, inter, params) - end + return potential_cutoff(cutoff, r2, inter, params) * (r2 <= cutoff.sqdist_cutoff) elseif cutoff_points(C) == 2 - if r2 > cutoff.sqdist_cutoff - return ustrip(zero(T)) * energy_units - elseif r2 < cutoff.sqdist_activation - return potential(inter, r2, inv(r2), params) - else - return potential_cutoff(cutoff, r2, inter, params) - end + return ifelse( + r2 < cutoff.sqdist_activation, + potential(inter, r2, inv(r2), params), + potential_cutoff(cutoff, r2, inter, params) * (r2 <= cutoff.sqdist_cutoff), + ) end end From 5b83a0162a78a5ffefb0ee9d31a45469dda7053a Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Tue, 25 Jun 2024 14:35:29 +0100 Subject: [PATCH 02/74] coord not needed for cutoff function --- src/cutoffs.jl | 6 ++---- src/interactions/buckingham.jl | 2 +- src/interactions/coulomb.jl | 4 ++-- src/interactions/lennard_jones.jl | 4 ++-- src/interactions/mie.jl | 2 +- src/interactions/soft_sphere.jl | 2 +- 6 files changed, 9 insertions(+), 11 deletions(-) diff --git a/src/cutoffs.jl b/src/cutoffs.jl index 542d84bcf..b1d2e65cf 100644 --- a/src/cutoffs.jl +++ b/src/cutoffs.jl @@ -145,8 +145,7 @@ end Base.:+(c1::T, ::T) where {T <: Union{NoCutoff, DistanceCutoff, ShiftedPotentialCutoff, ShiftedForceCutoff, CubicSplineCutoff}} = c1 -function force_divr_with_cutoff(inter, r2, params, cutoff::C, coord_i::SVector{D, T}, - force_units) where {C, D, T} +function force_divr_with_cutoff(inter, r2, params, cutoff::C, force_units) where C if cutoff_points(C) == 0 return force_divr(inter, r2, inv(r2), params) elseif cutoff_points(C) == 1 @@ -160,8 +159,7 @@ function force_divr_with_cutoff(inter, r2, params, cutoff::C, coord_i::SVector{D end end -function potential_with_cutoff(inter, r2, params, cutoff::C, coord_i::SVector{D, T}, - energy_units) where {C, D, T} +function potential_with_cutoff(inter, r2, params, cutoff::C, energy_units) where C if cutoff_points(C) == 0 return potential(inter, r2, inv(r2), params) elseif cutoff_points(C) == 1 diff --git a/src/interactions/buckingham.jl b/src/interactions/buckingham.jl index 102b215ca..af54823ef 100644 --- a/src/interactions/buckingham.jl +++ b/src/interactions/buckingham.jl @@ -64,7 +64,7 @@ use_neighbors(inter::Buckingham) = inter.use_neighbors r2 = sum(abs2, dr) params = (Aij, Bij, Cij) - f = force_divr_with_cutoff(inter, r2, params, cutoff, coord_i, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, inter.force_units) if special return f * dr * inter.weight_special else diff --git a/src/interactions/coulomb.jl b/src/interactions/coulomb.jl index 595d160e7..94ecced5c 100644 --- a/src/interactions/coulomb.jl +++ b/src/interactions/coulomb.jl @@ -73,7 +73,7 @@ end qi, qj = atom_i.charge, atom_j.charge params = (coulomb_const, qi, qj) - f = force_divr_with_cutoff(inter, r2, params, cutoff, coord_i, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, inter.force_units) if special return f * dr * inter.weight_special else @@ -205,7 +205,7 @@ end σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) params = (coulomb_const, qi, qj, σ, inter.σ6_fac) - f = force_divr_with_cutoff(inter, r2, params, cutoff, coord_i, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, inter.force_units) if special return f * dr * inter.weight_special else diff --git a/src/interactions/lennard_jones.jl b/src/interactions/lennard_jones.jl index 3dc1e83af..7dd0b01cd 100644 --- a/src/interactions/lennard_jones.jl +++ b/src/interactions/lennard_jones.jl @@ -102,7 +102,7 @@ end σ2 = σ^2 params = (σ2, ϵ) - f = force_divr_with_cutoff(inter, r2, params, cutoff, coord_i, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, inter.force_units) if special return f * dr * inter.weight_special else @@ -269,7 +269,7 @@ end σ2 = σ^2 params = (σ2, ϵ, inter.σ6_fac) - f = force_divr_with_cutoff(inter, r2, params, cutoff, coord_i, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, inter.force_units) if special return f * dr * inter.weight_special else diff --git a/src/interactions/mie.jl b/src/interactions/mie.jl index 27b58615c..2568e2e06 100644 --- a/src/interactions/mie.jl +++ b/src/interactions/mie.jl @@ -68,7 +68,7 @@ function force(inter::Mie{S, C, T}, σ_r = σ / r params = (m, n, σ_r, const_mn) - f = force_divr_with_cutoff(inter, r2, params, cutoff, coord_i, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, inter.force_units) return f * dr end diff --git a/src/interactions/soft_sphere.jl b/src/interactions/soft_sphere.jl index df3404c21..45ef44029 100644 --- a/src/interactions/soft_sphere.jl +++ b/src/interactions/soft_sphere.jl @@ -53,7 +53,7 @@ use_neighbors(inter::SoftSphere) = inter.use_neighbors σ2 = σ^2 params = (σ2, ϵ) - f = force_divr_with_cutoff(inter, r2, params, cutoff, coord_i, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, inter.force_units) return f * dr end From a4251a08c53f2ad47d36696dd8b9b39f3065f235 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 27 Jun 2024 16:50:40 +0100 Subject: [PATCH 03/74] potential cutoff args --- src/interactions/buckingham.jl | 2 +- src/interactions/coulomb.jl | 4 ++-- src/interactions/lennard_jones.jl | 4 ++-- src/interactions/mie.jl | 2 +- src/interactions/soft_sphere.jl | 2 +- 5 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/interactions/buckingham.jl b/src/interactions/buckingham.jl index af54823ef..6b74d58ec 100644 --- a/src/interactions/buckingham.jl +++ b/src/interactions/buckingham.jl @@ -98,7 +98,7 @@ end r2 = sum(abs2, dr) params = (Aij, Bij, Cij) - pe = potential_with_cutoff(inter, r2, params, cutoff, coord_i, inter.energy_units) + pe = potential_with_cutoff(inter, r2, params, cutoff, inter.energy_units) if special return pe * inter.weight_special else diff --git a/src/interactions/coulomb.jl b/src/interactions/coulomb.jl index 94ecced5c..effd1de7c 100644 --- a/src/interactions/coulomb.jl +++ b/src/interactions/coulomb.jl @@ -99,7 +99,7 @@ end qi, qj = atom_i.charge, atom_j.charge params = (coulomb_const, qi, qj) - pe = potential_with_cutoff(inter, r2, params, cutoff, coord_i, inter.energy_units) + pe = potential_with_cutoff(inter, r2, params, cutoff, inter.energy_units) if special return pe * inter.weight_special else @@ -236,7 +236,7 @@ end σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) params = (coulomb_const, qi, qj, σ, inter.σ6_fac) - pe = potential_with_cutoff(inter, r2, params, cutoff, coord_i, inter.energy_units) + pe = potential_with_cutoff(inter, r2, params, cutoff, inter.energy_units) if special return pe * inter.weight_special else diff --git a/src/interactions/lennard_jones.jl b/src/interactions/lennard_jones.jl index 7dd0b01cd..37f8da436 100644 --- a/src/interactions/lennard_jones.jl +++ b/src/interactions/lennard_jones.jl @@ -140,7 +140,7 @@ end σ2 = σ^2 params = (σ2, ϵ) - pe = potential_with_cutoff(inter, r2, params, cutoff, coord_i, inter.energy_units) + pe = potential_with_cutoff(inter, r2, params, cutoff, inter.energy_units) if special return pe * inter.weight_special else @@ -310,7 +310,7 @@ end σ2 = σ^2 params = (σ2, ϵ, inter.σ6_fac) - pe = potential_with_cutoff(inter, r2, params, cutoff, coord_i, inter.energy_units) + pe = potential_with_cutoff(inter, r2, params, cutoff, inter.energy_units) if special return pe * inter.weight_special else diff --git a/src/interactions/mie.jl b/src/interactions/mie.jl index 2568e2e06..d3fa52c9b 100644 --- a/src/interactions/mie.jl +++ b/src/interactions/mie.jl @@ -100,7 +100,7 @@ end σ_r = σ / r params = (m, n, σ_r, const_mn) - return potential_with_cutoff(inter, r2, params, cutoff, coord_i, inter.energy_units) + return potential_with_cutoff(inter, r2, params, cutoff, inter.energy_units) end function potential(::Mie, r2, invr2, (m, n, σ_r, const_mn)) diff --git a/src/interactions/soft_sphere.jl b/src/interactions/soft_sphere.jl index 45ef44029..a4b993fa4 100644 --- a/src/interactions/soft_sphere.jl +++ b/src/interactions/soft_sphere.jl @@ -82,7 +82,7 @@ function potential_energy(inter::SoftSphere{S, C}, σ2 = σ^2 params = (σ2, ϵ) - return potential_with_cutoff(inter, r2, params, cutoff, coord_i, inter.energy_units) + return potential_with_cutoff(inter, r2, params, cutoff, inter.energy_units) end function potential(::SoftSphere, r2, invr2, (σ2, ϵ)) From 28113c8fa81f8e432fcd0c14cb1012d20c09cd1f Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 28 Jun 2024 13:52:14 +0100 Subject: [PATCH 04/74] check units changes --- src/units.jl | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/src/units.jl b/src/units.jl index f16ec2f27..31934eae7 100644 --- a/src/units.jl +++ b/src/units.jl @@ -213,19 +213,13 @@ end function check_energy_units(E, energy_units) if unit(E) != energy_units - error("system energy units are ", energy_units, " but encountered energy units ", - unit(E)) + error("system energy units are ", energy_units, " but encountered energy units ", unit(E)) end end -function check_force_units(fdr::AbstractArray, sys_force_units) - return check_force_units(unit(first(fdr)), sys_force_units) -end - -function check_force_units(force_units, sys_force_units) - if force_units != sys_force_units - error("system force units are ", sys_force_units, " but encountered force units ", - force_units) +function check_force_units(F, force_units) + if unit(F) != force_units + error("system force units are ", force_units, " but encountered force units ", unit(F)) end end From e2d3a96fa55704bfe2a5e33be8f36e4a430a9dbd Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 28 Jun 2024 14:35:05 +0100 Subject: [PATCH 05/74] pairwise interaction arguments and fields --- src/interactions/buckingham.jl | 48 ++++----- src/interactions/coulomb.jl | 165 ++++++++++++------------------ src/interactions/gravity.jl | 20 ++-- src/interactions/lennard_jones.jl | 123 +++++++++------------- src/interactions/mie.jl | 42 ++++---- src/interactions/soft_sphere.jl | 42 ++++---- 6 files changed, 179 insertions(+), 261 deletions(-) diff --git a/src/interactions/buckingham.jl b/src/interactions/buckingham.jl index 6b74d58ec..3684669b7 100644 --- a/src/interactions/buckingham.jl +++ b/src/interactions/buckingham.jl @@ -1,7 +1,7 @@ export Buckingham @doc raw""" - Buckingham(; cutoff, use_neighbors, weight_special, force_units, energy_units) + Buckingham(; cutoff, use_neighbors, weight_special) The Buckingham interaction between two atoms. @@ -23,37 +23,32 @@ C_{ij} &= (C_{ii} C_{jj})^{1/2} ``` so atoms that use this interaction should have fields `A`, `B` and `C` available. """ -struct Buckingham{C, W, F, E} <: PairwiseInteraction +struct Buckingham{C, W} <: PairwiseInteraction cutoff::C use_neighbors::Bool weight_special::W - force_units::F - energy_units::E end function Buckingham(; cutoff=NoCutoff(), use_neighbors=false, - weight_special=1, - force_units=u"kJ * mol^-1 * nm^-1", - energy_units=u"kJ * mol^-1") - return Buckingham{typeof(cutoff), typeof(weight_special), typeof(force_units), typeof(energy_units)}( - cutoff, use_neighbors, weight_special, force_units, energy_units) + weight_special=1) + return Buckingham{typeof(cutoff), typeof(weight_special)}( + cutoff, use_neighbors, weight_special) end use_neighbors(inter::Buckingham) = inter.use_neighbors @inline function force(inter::Buckingham{C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) where C + dr, + atom_i, + atom_j, + force_units, + special::Bool=false, + args...) where C if (iszero_value(atom_i.A) || iszero_value(atom_j.A)) && (iszero_value(atom_i.C) || iszero_value(atom_j.C)) - return ustrip.(zero(coord_i)) * inter.force_units + return ustrip.(zero(coord_i)) * force_units end Aij = sqrt(atom_i.A * atom_j.A) @@ -64,7 +59,7 @@ use_neighbors(inter::Buckingham) = inter.use_neighbors r2 = sum(abs2, dr) params = (Aij, Bij, Cij) - f = force_divr_with_cutoff(inter, r2, params, cutoff, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, force_units) if special return f * dr * inter.weight_special else @@ -78,16 +73,15 @@ function force_divr(::Buckingham, r2, invr2, (A, B, C)) end @inline function potential_energy(inter::Buckingham{C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) where C + dr, + atom_i, + atom_j, + energy_units, + special::Bool=false, + args...) where C if (iszero_value(atom_i.A) || iszero_value(atom_j.A)) && (iszero_value(atom_i.C) || iszero_value(atom_j.C)) - return ustrip(zero(coord_i[1])) * inter.energy_units + return ustrip(zero(coord_i[1])) * energy_units end Aij = sqrt(atom_i.A * atom_j.A) @@ -98,7 +92,7 @@ end r2 = sum(abs2, dr) params = (Aij, Bij, Cij) - pe = potential_with_cutoff(inter, r2, params, cutoff, inter.energy_units) + pe = potential_with_cutoff(inter, r2, params, cutoff, energy_units) if special return pe * inter.weight_special else diff --git a/src/interactions/coulomb.jl b/src/interactions/coulomb.jl index effd1de7c..f2ab87bbb 100644 --- a/src/interactions/coulomb.jl +++ b/src/interactions/coulomb.jl @@ -4,7 +4,7 @@ export CoulombReactionField @doc raw""" - Coulomb(; cutoff, use_neighbors, weight_special, coulomb_const, force_units, energy_units) + Coulomb(; cutoff, use_neighbors, weight_special, coulomb_const) The Coulomb electrostatic interaction between two atoms. @@ -13,13 +13,11 @@ The potential energy is defined as V(r_{ij}) = \frac{q_i q_j}{4 \pi \varepsilon_0 r_{ij}} ``` """ -struct Coulomb{C, W, T, F, E} <: PairwiseInteraction +struct Coulomb{C, W, T} <: PairwiseInteraction cutoff::C use_neighbors::Bool weight_special::W coulomb_const::T - force_units::F - energy_units::E end const coulombconst = 138.93545764u"kJ * mol^-1 * nm" # 1 / 4πϵ0 @@ -29,23 +27,14 @@ function Coulomb(; use_neighbors=false, weight_special=1, coulomb_const=coulombconst, - force_units=u"kJ * mol^-1 * nm^-1", - energy_units=u"kJ * mol^-1") - return Coulomb{typeof(cutoff), typeof(weight_special), typeof(coulomb_const), typeof(force_units), typeof(energy_units)}( - cutoff, use_neighbors, weight_special, coulomb_const, force_units, energy_units) + return Coulomb{typeof(cutoff), typeof(weight_special), typeof(coulomb_const)}( + cutoff, use_neighbors, weight_special, coulomb_const) end use_neighbors(inter::Coulomb) = inter.use_neighbors -function Base.zero(coul::Coulomb{C, W, T, F, E}) where {C, W, T, F, E} - return Coulomb{C, W, T, F, E}( - coul.cutoff, - false, - zero(W), - zero(T), - coul.force_units, - coul.energy_units, - ) +function Base.zero(coul::Coulomb{C, W, T}) where {C, W, T} + return Coulomb{C, W, T}(coul.cutoff, false, zero(W), zero(T)) end function Base.:+(c1::Coulomb, c2::Coulomb) @@ -54,26 +43,23 @@ function Base.:+(c1::Coulomb, c2::Coulomb) c1.use_neighbors, c1.weight_special + c2.weight_special, c1.coulomb_const + c2.coulomb_const, - c1.force_units, - c1.energy_units, ) end @inline function force(inter::Coulomb{C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) where C + dr, + atom_i, + atom_j, + force_units, + special::Bool=false, + args...) where C r2 = sum(abs2, dr) cutoff = inter.cutoff coulomb_const = inter.coulomb_const qi, qj = atom_i.charge, atom_j.charge params = (coulomb_const, qi, qj) - f = force_divr_with_cutoff(inter, r2, params, cutoff, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, force_units) if special return f * dr * inter.weight_special else @@ -86,20 +72,19 @@ function force_divr(::Coulomb, r2, invr2, (coulomb_const, qi, qj)) end @inline function potential_energy(inter::Coulomb{C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) where C + dr, + atom_i, + atom_j, + energy_units, + special::Bool=false, + args...) where C r2 = sum(abs2, dr) cutoff = inter.cutoff coulomb_const = inter.coulomb_const qi, qj = atom_i.charge, atom_j.charge params = (coulomb_const, qi, qj) - pe = potential_with_cutoff(inter, r2, params, cutoff, inter.energy_units) + pe = potential_with_cutoff(inter, r2, params, cutoff, energy_units) if special return pe * inter.weight_special else @@ -112,8 +97,7 @@ function potential(::Coulomb, r2, invr2, (coulomb_const, qi, qj)) end @doc raw""" - CoulombSoftCore(; cutoff, α, λ, p, use_neighbors, lorentz_mixing, weight_special, - coulomb_const, force_units, energy_units) + CoulombSoftCore(; cutoff, α, λ, p, use_neighbors, lorentz_mixing, weight_special, coulomb_const) The Coulomb electrostatic interaction between two atoms with a soft core. @@ -123,7 +107,7 @@ V(r_{ij}) = \frac{q_i q_j}{4 \pi \varepsilon_0 (r_{ij}^6 + \alpha \sigma_{ij}^6 ``` If ``\alpha`` or ``\lambda`` are zero this gives the standard [`Coulomb`](@ref) potential. """ -struct CoulombSoftCore{C, A, L, P, R, W, T, F, E} <: PairwiseInteraction +struct CoulombSoftCore{C, A, L, P, R, W, T} <: PairwiseInteraction cutoff::C α::A λ::L @@ -133,8 +117,6 @@ struct CoulombSoftCore{C, A, L, P, R, W, T, F, E} <: PairwiseInteraction lorentz_mixing::Bool weight_special::W coulomb_const::T - force_units::F - energy_units::E end function CoulombSoftCore(; @@ -145,21 +127,17 @@ function CoulombSoftCore(; use_neighbors=false, lorentz_mixing=true, weight_special=1, - coulomb_const=coulombconst, - force_units=u"kJ * mol^-1 * nm^-1", - energy_units=u"kJ * mol^-1") + coulomb_const=coulombconst) σ6_fac = α * λ^p return CoulombSoftCore{typeof(cutoff), typeof(α), typeof(λ), typeof(p), typeof(σ6_fac), - typeof(weight_special), typeof(coulomb_const), typeof(force_units), - typeof(energy_units)}( - cutoff, α, λ, p, σ6_fac, use_neighbors, lorentz_mixing, weight_special, coulomb_const, - force_units, energy_units) + typeof(weight_special), typeof(coulomb_const)}( + cutoff, α, λ, p, σ6_fac, use_neighbors, lorentz_mixing, weight_special, coulomb_const) end use_neighbors(inter::CoulombSoftCore) = inter.use_neighbors -function Base.zero(coul::CoulombSoftCore{C, A, L, P, R, W, T, F, E}) where {C, A, L, P, R, W, T, F, E} - return CoulombSoftCore{C, A, L, P, R, W, T, F, E}( +function Base.zero(coul::CoulombSoftCore{C, A, L, P, R, W, T}) where {C, A, L, P, R, W, T} + return CoulombSoftCore{C, A, L, P, R, W, T}( coul.cutoff, zero(A), zero(L), @@ -169,8 +147,6 @@ function Base.zero(coul::CoulombSoftCore{C, A, L, P, R, W, T, F, E}) where {C, A false, zero(W), zero(T), - coul.force_units, - coul.energy_units, ) end @@ -185,19 +161,16 @@ function Base.:+(c1::CoulombSoftCore, c2::CoulombSoftCore) c1.lorentz_mixing, c1.weight_special + c2.weight_special, c1.coulomb_const + c2.coulomb_const, - c1.force_units, - c1.energy_units, ) end @inline function force(inter::CoulombSoftCore{C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) where C + dr, + atom_i, + atom_j, + force_units, + special::Bool=false, + args...) where C r2 = sum(abs2, dr) cutoff = inter.cutoff coulomb_const = inter.coulomb_const @@ -205,7 +178,7 @@ end σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) params = (coulomb_const, qi, qj, σ, inter.σ6_fac) - f = force_divr_with_cutoff(inter, r2, params, cutoff, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, force_units) if special return f * dr * inter.weight_special else @@ -222,13 +195,12 @@ function force_divr(::CoulombSoftCore, r2, invr2, (coulomb_const, qi, qj, σ, σ end @inline function potential_energy(inter::CoulombSoftCore{C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) where C + dr, + atom_i, + atom_j, + energy_units, + special::Bool=false, + args...) where C r2 = sum(abs2, dr) cutoff = inter.cutoff coulomb_const = inter.coulomb_const @@ -236,7 +208,7 @@ end σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) params = (coulomb_const, qi, qj, σ, inter.σ6_fac) - pe = potential_with_cutoff(inter, r2, params, cutoff, inter.energy_units) + pe = potential_with_cutoff(inter, r2, params, cutoff, energy_units) if special return pe * inter.weight_special else @@ -251,19 +223,17 @@ end """ CoulombReactionField(; dist_cutoff, solvent_dielectric, use_neighbors, weight_special, - coulomb_const, force_units, energy_units) + coulomb_const) The Coulomb electrostatic interaction modified using the reaction field approximation between two atoms. """ -struct CoulombReactionField{D, S, W, T, F, E} <: PairwiseInteraction +struct CoulombReactionField{D, S, W, T} <: PairwiseInteraction dist_cutoff::D solvent_dielectric::S use_neighbors::Bool weight_special::W coulomb_const::T - force_units::F - energy_units::E end const crf_solvent_dielectric = 78.3 @@ -273,26 +243,21 @@ function CoulombReactionField(; solvent_dielectric=crf_solvent_dielectric, use_neighbors=false, weight_special=1, - coulomb_const=coulombconst, - force_units=u"kJ * mol^-1 * nm^-1", - energy_units=u"kJ * mol^-1") - return CoulombReactionField{typeof(dist_cutoff), typeof(solvent_dielectric), typeof(weight_special), - typeof(coulomb_const), typeof(force_units), typeof(energy_units)}( - dist_cutoff, solvent_dielectric, use_neighbors, weight_special, - coulomb_const, force_units, energy_units) + coulomb_const=coulombconst) + return CoulombReactionField{typeof(dist_cutoff), typeof(solvent_dielectric), + typeof(weight_special), typeof(coulomb_const)}( + dist_cutoff, solvent_dielectric, use_neighbors, weight_special, coulomb_const) end use_neighbors(inter::CoulombReactionField) = inter.use_neighbors -function Base.zero(coul::CoulombReactionField{D, S, W, T, F, E}) where {D, S, W, T, F, E} - return CoulombReactionField{D, S, W, T, F, E}( +function Base.zero(coul::CoulombReactionField{D, S, W, T}) where {D, S, W, T} + return CoulombReactionField{D, S, W, T}( zero(D), zero(S), false, zero(W), zero(T), - coul.force_units, - coul.energy_units, ) end @@ -303,22 +268,19 @@ function Base.:+(c1::CoulombReactionField, c2::CoulombReactionField) c1.use_neighbors, c1.weight_special + c2.weight_special, c1.coulomb_const + c2.coulomb_const, - c1.force_units, - c1.energy_units, ) end @inline function force(inter::CoulombReactionField, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) + dr, + atom_i, + atom_j, + force_units, + special::Bool=false, + args...) r2 = sum(abs2, dr) if r2 > (inter.dist_cutoff ^ 2) - return ustrip.(zero(coord_i)) * inter.force_units + return ustrip.(zero(coord_i)) * force_units end coulomb_const = inter.coulomb_const @@ -343,16 +305,15 @@ end end @inline function potential_energy(inter::CoulombReactionField, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) + dr, + atom_i, + atom_j, + energy_units, + special::Bool=false, + args...) r2 = sum(abs2, dr) if r2 > (inter.dist_cutoff ^ 2) - return ustrip(zero(coord_i[1])) * inter.energy_units + return ustrip(zero(coord_i[1])) * energy_units end coulomb_const = inter.coulomb_const diff --git a/src/interactions/gravity.jl b/src/interactions/gravity.jl index 5246c7ce2..02eabe250 100644 --- a/src/interactions/gravity.jl +++ b/src/interactions/gravity.jl @@ -20,12 +20,10 @@ Gravity(; G=Unitful.G, use_neighbors=false) = Gravity{typeof(G)}(G, use_neighbor use_neighbors(inter::Gravity) = inter.use_neighbors @inline function force(inter::Gravity, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary) + dr, + atom_i, + atom_j, + args...) r2 = sum(abs2, dr) params = (inter.G, mass(atom_i), mass(atom_j)) f = force_divr(inter, r2, inv(r2), params) @@ -37,12 +35,10 @@ function force_divr(::Gravity, r2, invr2, (G, mi, mj)) end @inline function potential_energy(inter::Gravity, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary) + dr, + atom_i, + atom_j, + args...) r2 = sum(abs2, dr) params = (inter.G, mass(atom_i), mass(atom_j)) potential(inter, r2, inv(r2), params) diff --git a/src/interactions/lennard_jones.jl b/src/interactions/lennard_jones.jl index 37f8da436..b79de4c1f 100644 --- a/src/interactions/lennard_jones.jl +++ b/src/interactions/lennard_jones.jl @@ -4,7 +4,7 @@ export @doc raw""" LennardJones(; cutoff, use_neighbors, lorentz_mixing, weight_special, weight_solute_solvent, - force_units, energy_units, skip_shortcut) + skip_shortcut) The Lennard-Jones 6-12 interaction between two atoms. @@ -20,14 +20,12 @@ and the force on each atom by \end{aligned} ``` """ -struct LennardJones{S, C, W, WS, F, E} <: PairwiseInteraction +struct LennardJones{S, C, W, WS} <: PairwiseInteraction cutoff::C use_neighbors::Bool lorentz_mixing::Bool weight_special::W weight_solute_solvent::WS - force_units::F - energy_units::E end function LennardJones(; @@ -36,13 +34,10 @@ function LennardJones(; lorentz_mixing=true, weight_special=1, weight_solute_solvent=1, - force_units=u"kJ * mol^-1 * nm^-1", - energy_units=u"kJ * mol^-1", skip_shortcut=false) return LennardJones{skip_shortcut, typeof(cutoff), typeof(weight_special), - typeof(weight_solute_solvent), typeof(force_units), typeof(energy_units)}( - cutoff, use_neighbors, lorentz_mixing, weight_special, weight_solute_solvent, - force_units, energy_units) + typeof(weight_solute_solvent)}( + cutoff, use_neighbors, lorentz_mixing, weight_special, weight_solute_solvent) end use_neighbors(inter::LennardJones) = inter.use_neighbors @@ -50,42 +45,37 @@ use_neighbors(inter::LennardJones) = inter.use_neighbors is_solute(at::Atom) = at.solute is_solute(at) = false -function Base.zero(lj::LennardJones{S, C, W, WS, F, E}) where {S, C, W, WS, F, E} - return LennardJones{S, C, W, WS, F, E}( +function Base.zero(lj::LennardJones{S, C, W, WS}) where {S, C, W, WS} + return LennardJones{S, C, W, WS}( lj.cutoff, false, false, zero(W), zero(WS), - lj.force_units, - lj.energy_units, ) end -function Base.:+(l1::LennardJones{S, C, W, WS, F, E}, - l2::LennardJones{S, C, W, WS, F, E}) where {S, C, W, WS, F, E} - return LennardJones{S, C, W, WS, F, E}( +function Base.:+(l1::LennardJones{S, C, W, WS}, + l2::LennardJones{S, C, W, WS}) where {S, C, W, WS} + return LennardJones{S, C, W, WS}( l1.cutoff, l1.use_neighbors, l1.lorentz_mixing, l1.weight_special + l2.weight_special, l1.weight_solute_solvent + l2.weight_solute_solvent, - l1.force_units, - l1.energy_units, ) end @inline function force(inter::LennardJones{S, C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) where {S, C} + dr, + atom_i, + atom_j, + force_units, + special::Bool=false, + args...) where {S, C} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) - return ustrip.(zero(coord_i)) * inter.force_units + return ustrip.(zero(dr)) * force_units end # Lorentz-Berthelot mixing rules use the arithmetic average for σ @@ -102,7 +92,7 @@ end σ2 = σ^2 params = (σ2, ϵ) - f = force_divr_with_cutoff(inter, r2, params, cutoff, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, force_units) if special return f * dr * inter.weight_special else @@ -116,16 +106,15 @@ function force_divr(::LennardJones, r2, invr2, (σ2, ϵ)) end @inline function potential_energy(inter::LennardJones{S, C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) where {S, C} + dr, + atom_i, + atom_j, + energy_units, + special::Bool=false, + args...) where {S, C} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) - return ustrip(zero(coord_i[1])) * inter.energy_units + return ustrip(zero(coord_i[1])) * energy_units end σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) @@ -140,7 +129,7 @@ end σ2 = σ^2 params = (σ2, ϵ) - pe = potential_with_cutoff(inter, r2, params, cutoff, inter.energy_units) + pe = potential_with_cutoff(inter, r2, params, cutoff, energy_units) if special return pe * inter.weight_special else @@ -155,7 +144,7 @@ end @doc raw""" LennardJonesSoftCore(; cutoff, α, λ, p, use_neighbors, lorentz_mixing, weight_special, - weight_solute_solvent, force_units, energy_units, skip_shortcut) + weight_solute_solvent, skip_shortcut) The Lennard-Jones 6-12 interaction between two atoms with a soft core. @@ -173,7 +162,7 @@ r_{ij}^{\text{sc}} = \left(r_{ij}^6 + \alpha \sigma_{ij}^6 \lambda^p \right)^{1/ ``` If ``\alpha`` or ``\lambda`` are zero this gives the standard [`LennardJones`](@ref) potential. """ -struct LennardJonesSoftCore{S, C, A, L, P, R, W, WS, F, E} <: PairwiseInteraction +struct LennardJonesSoftCore{S, C, A, L, P, R, W, WS} <: PairwiseInteraction cutoff::C α::A λ::L @@ -183,8 +172,6 @@ struct LennardJonesSoftCore{S, C, A, L, P, R, W, WS, F, E} <: PairwiseInteractio lorentz_mixing::Bool weight_special::W weight_solute_solvent::WS - force_units::F - energy_units::E end function LennardJonesSoftCore(; @@ -196,21 +183,19 @@ function LennardJonesSoftCore(; lorentz_mixing=true, weight_special=1, weight_solute_solvent=1, - force_units=u"kJ * mol^-1 * nm^-1", - energy_units=u"kJ * mol^-1", skip_shortcut=false) σ6_fac = α * λ^p return LennardJonesSoftCore{skip_shortcut, typeof(cutoff), typeof(α), typeof(λ), typeof(p), - typeof(σ6_fac), typeof(weight_special), typeof(weight_solute_solvent), - typeof(force_units), typeof(energy_units)}( + typeof(σ6_fac), typeof(weight_special), + typeof(weight_solute_solvent)}( cutoff, α, λ, p, σ6_fac, use_neighbors, lorentz_mixing, weight_special, - weight_solute_solvent, force_units, energy_units) + weight_solute_solvent) end use_neighbors(inter::LennardJonesSoftCore) = inter.use_neighbors -function Base.zero(lj::LennardJonesSoftCore{S, C, A, L, P, R, W, WS, F, E}) where {S, C, A, L, P, R, W, WS, F, E} - return LennardJonesSoftCore{S, C, A, L, P, R, W, WS, F, E}( +function Base.zero(lj::LennardJonesSoftCore{S, C, A, L, P, R, W, WS}) where {S, C, A, L, P, R, W, WS} + return LennardJonesSoftCore{S, C, A, L, P, R, W, WS}( lj.cutoff, zero(A), zero(L), @@ -220,14 +205,12 @@ function Base.zero(lj::LennardJonesSoftCore{S, C, A, L, P, R, W, WS, F, E}) wher false, zero(W), zero(WS), - lj.force_units, - lj.energy_units, ) end -function Base.:+(l1::LennardJonesSoftCore{S, C, A, L, P, R, W, WS, F, E}, - l2::LennardJonesSoftCore{S, C, A, L, P, R, W, WS, F, E}) where {S, C, A, L, P, R, W, WS, F, E} - return LennardJonesSoftCore{S, C, A, L, P, R, W, WS, F, E}( +function Base.:+(l1::LennardJonesSoftCore{S, C, A, L, P, R, W, WS}, + l2::LennardJonesSoftCore{S, C, A, L, P, R, W, WS}) where {S, C, A, L, P, R, W, WS} + return LennardJonesSoftCore{S, C, A, L, P, R, W, WS}( l1.cutoff, l1.α + l2.α, l1.λ + l2.λ, @@ -237,22 +220,19 @@ function Base.:+(l1::LennardJonesSoftCore{S, C, A, L, P, R, W, WS, F, E}, l1.lorentz_mixing, l1.weight_special + l2.weight_special, l1.weight_solute_solvent + l2.weight_solute_solvent, - l1.force_units, - l1.energy_units, ) end @inline function force(inter::LennardJonesSoftCore{S, C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) where {S, C} + dr, + atom_i, + atom_j, + force_units, + special::Bool=false, + args...) where {S, C} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) - return ustrip.(zero(coord_i)) * inter.force_units + return ustrip.(zero(coord_i)) * force_units end # Lorentz-Berthelot mixing rules use the arithmetic average for σ @@ -269,7 +249,7 @@ end σ2 = σ^2 params = (σ2, ϵ, inter.σ6_fac) - f = force_divr_with_cutoff(inter, r2, params, cutoff, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, force_units) if special return f * dr * inter.weight_special else @@ -286,16 +266,15 @@ function force_divr(::LennardJonesSoftCore, r2, invr2, (σ2, ϵ, σ6_fac)) end @inline function potential_energy(inter::LennardJonesSoftCore{S, C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) where {S, C} + dr, + atom_i, + atom_j, + energy_units, + special::Bool=false, + args...) where {S, C} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || iszero_value(atom_i.σ) ||iszero_value(atom_j.σ)) - return ustrip(zero(coord_i[1])) * inter.energy_units + return ustrip(zero(coord_i[1])) * energy_units end σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) @@ -310,7 +289,7 @@ end σ2 = σ^2 params = (σ2, ϵ, inter.σ6_fac) - pe = potential_with_cutoff(inter, r2, params, cutoff, inter.energy_units) + pe = potential_with_cutoff(inter, r2, params, cutoff, energy_units) if special return pe * inter.weight_special else diff --git a/src/interactions/mie.jl b/src/interactions/mie.jl index d3fa52c9b..883669412 100644 --- a/src/interactions/mie.jl +++ b/src/interactions/mie.jl @@ -1,7 +1,7 @@ export Mie @doc raw""" - Mie(; m, n, cutoff, use_neighbors, lorentz_mixing, force_units, energy_units, skip_shortcut) + Mie(; m, n, cutoff, use_neighbors, lorentz_mixing, skip_shortcut) The Mie generalized interaction between two atoms. @@ -15,14 +15,12 @@ where C = \frac{n}{n - m} \left( \frac{n}{m} \right) ^\frac{m}{n - m} ``` """ -struct Mie{S, C, T, F, E} <: PairwiseInteraction +struct Mie{S, C, T} <: PairwiseInteraction m::T n::T cutoff::C use_neighbors::Bool lorentz_mixing::Bool - force_units::F - energy_units::E mn_fac::T end @@ -32,26 +30,23 @@ function Mie(; cutoff=NoCutoff(), use_neighbors=false, lorentz_mixing=true, - force_units=u"kJ * mol^-1 * nm^-1", - energy_units=u"kJ * mol^-1", skip_shortcut=false) m_p, n_p, mn_fac = promote(m, n, (n / (n - m)) * (n / m) ^ (m / (n - m))) - return Mie{skip_shortcut, typeof(cutoff), typeof(m_p), typeof(force_units), typeof(energy_units)}( - m_p, n_p, cutoff, use_neighbors, lorentz_mixing, force_units, energy_units, mn_fac) + return Mie{skip_shortcut, typeof(cutoff), typeof(m_p)}( + m_p, n_p, cutoff, use_neighbors, lorentz_mixing, mn_fac) end use_neighbors(inter::Mie) = inter.use_neighbors function force(inter::Mie{S, C, T}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary) where {S, C, T} + dr, + atom_i, + atom_j, + force_units, + args...) where {S, C, T} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) - return ustrip.(zero(coord_i)) * inter.force_units + return ustrip.(zero(coord_i)) * force_units end # Lorentz-Berthelot mixing rules use the arithmetic average for σ @@ -68,7 +63,7 @@ function force(inter::Mie{S, C, T}, σ_r = σ / r params = (m, n, σ_r, const_mn) - f = force_divr_with_cutoff(inter, r2, params, cutoff, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, force_units) return f * dr end @@ -77,15 +72,14 @@ function force_divr(::Mie, r2, invr2, (m, n, σ_r, const_mn)) end @inline function potential_energy(inter::Mie{S, C, T}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary) where {S, C, T} + dr, + atom_i, + atom_j, + energy_units, + args...) where {S, C, T} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) - return ustrip(zero(coord_i[1])) * inter.energy_units + return ustrip(zero(coord_i[1])) * energy_units end σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) @@ -100,7 +94,7 @@ end σ_r = σ / r params = (m, n, σ_r, const_mn) - return potential_with_cutoff(inter, r2, params, cutoff, inter.energy_units) + return potential_with_cutoff(inter, r2, params, cutoff, energy_units) end function potential(::Mie, r2, invr2, (m, n, σ_r, const_mn)) diff --git a/src/interactions/soft_sphere.jl b/src/interactions/soft_sphere.jl index a4b993fa4..cca4fa8f3 100644 --- a/src/interactions/soft_sphere.jl +++ b/src/interactions/soft_sphere.jl @@ -1,7 +1,7 @@ export SoftSphere @doc raw""" - SoftSphere(; cutoff, use_neighbors, lorentz_mixing, force_units, energy_units, skip_shortcut) + SoftSphere(; cutoff, use_neighbors, lorentz_mixing, skip_shortcut) The soft-sphere potential. @@ -10,37 +10,32 @@ The potential energy is defined as V(r_{ij}) = 4\varepsilon_{ij} \left(\frac{\sigma_{ij}}{r_{ij}}\right)^{12} ``` """ -struct SoftSphere{S, C, F, E} <: PairwiseInteraction +struct SoftSphere{S, C} <: PairwiseInteraction cutoff::C use_neighbors::Bool lorentz_mixing::Bool - force_units::F - energy_units::E end function SoftSphere(; cutoff=NoCutoff(), use_neighbors=false, lorentz_mixing=true, - force_units=u"kJ * mol^-1 * nm^-1", - energy_units=u"kJ * mol^-1", skip_shortcut=false) - return SoftSphere{skip_shortcut, typeof(cutoff), typeof(force_units), typeof(energy_units)}( - cutoff, use_neighbors, lorentz_mixing, force_units, energy_units) + return SoftSphere{skip_shortcut, typeof(cutoff)}( + cutoff, use_neighbors, lorentz_mixing) end use_neighbors(inter::SoftSphere) = inter.use_neighbors @inline function force(inter::SoftSphere{S, C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary) where {S, C} + dr, + atom_i, + atom_j, + force_units, + args...) where {S, C} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) - return ustrip.(zero(coord_i)) * inter.force_units + return ustrip.(zero(coord_i)) * force_units end # Lorentz-Berthelot mixing rules use the arithmetic average for σ @@ -53,7 +48,7 @@ use_neighbors(inter::SoftSphere) = inter.use_neighbors σ2 = σ^2 params = (σ2, ϵ) - f = force_divr_with_cutoff(inter, r2, params, cutoff, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, force_units) return f * dr end @@ -63,15 +58,14 @@ function force_divr(::SoftSphere, r2, invr2, (σ2, ϵ)) end function potential_energy(inter::SoftSphere{S, C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary) where {S, C} + dr, + atom_i, + atom_j, + energy_units, + args...) where {S, C} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) - return ustrip(zero(coord_i[1])) * inter.energy_units + return ustrip(zero(coord_i[1])) * energy_units end σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) @@ -82,7 +76,7 @@ function potential_energy(inter::SoftSphere{S, C}, σ2 = σ^2 params = (σ2, ϵ) - return potential_with_cutoff(inter, r2, params, cutoff, inter.energy_units) + return potential_with_cutoff(inter, r2, params, cutoff, energy_units) end function potential(::SoftSphere, r2, invr2, (σ2, ϵ)) From fdd51110d805250f1ab134c04a6965ef200b3db6 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 28 Jun 2024 14:36:42 +0100 Subject: [PATCH 06/74] remove Zygote functions --- src/zygote.jl | 724 -------------------------------------------------- 1 file changed, 724 deletions(-) delete mode 100644 src/zygote.jl diff --git a/src/zygote.jl b/src/zygote.jl deleted file mode 100644 index 3c8d0d9e8..000000000 --- a/src/zygote.jl +++ /dev/null @@ -1,724 +0,0 @@ -# Extend Zygote to work with static vectors and custom types on the -# fast broadcast/GPU path -# Here be dragons - -using ForwardDiff: Chunk, Dual, partials, value -using Zygote: unbroadcast - -iszero_value(x::Dual) = iszero(value(x)) -iszero_value(x) = iszero(x) - -Zygote.accum(x::AbstractArray{<:SizedVector}, ys::AbstractArray{<:SVector}...) = Zygote.accum.(convert(typeof(ys[1]), x), ys...) -Zygote.accum(x::AbstractArray{<:SVector}, ys::AbstractArray{<:SizedVector}...) = Zygote.accum.(x, convert.(typeof(x), ys)...) - -Zygote.accum(x::Vector{<:SVector} , y::CuArray{<:SVector}) = Zygote.accum(CuArray(x), y) -Zygote.accum(x::CuArray{<:SVector}, y::Vector{<:SVector} ) = Zygote.accum(x, CuArray(y)) - -Zygote.accum(x::SVector{D, T}, y::T) where {D, T} = x .+ y - -Zygote.accum(u1::T, ::T) where {T <: Unitful.FreeUnits} = u1 - -Base.:+(x::Real, y::SizedVector) = x .+ y -Base.:+(x::SizedVector, y::Real) = x .+ y - -Base.:+(x::Real, y::Zygote.OneElement) = x .+ y -Base.:+(x::Zygote.OneElement, y::Real) = x .+ y - -function Zygote.accum(x::CuArray{Atom{T, T, T, T}}, - y::Vector{NamedTuple{(:index, :charge, :mass, :σ, :ϵ, :solute)}}) where T - CuArray(Zygote.accum(Array(x), y)) -end - -function Base.:+(x::Atom{T, T, T, T}, y::NamedTuple{(:index, :charge, :mass, :σ, :ϵ, :solute), - Tuple{Int, C, M, S, E, Bool}}) where {T, C, M, S, E} - Atom{T, T, T, T}( - 0, - Zygote.accum(x.charge, y.charge), - Zygote.accum(x.mass, y.mass), - Zygote.accum(x.σ, y.σ), - Zygote.accum(x.ϵ, y.ϵ), - false, - ) -end - -function Base.:+(r::Base.RefValue{Any}, y::NamedTuple{(:atoms, :coords, :boundary, - :velocities, :atoms_data, :topology, :pairwise_inters, :specific_inter_lists, - :general_inters, :constraints, :neighbor_finder, :loggers, :df, :force_units, - :energy_units, :k, :masses, :data)}) - x = r.x - ( - atoms=Zygote.accum(x.atoms, y.atoms), - coords=Zygote.accum(x.coords, y.coords), - boundary=Zygote.accum(x.boundary, y.boundary), - velocities=Zygote.accum(x.velocities, y.velocities), - atoms_data=Zygote.accum(x.atoms_data, y.atoms_data), - topology=nothing, - pairwise_inters=Zygote.accum(x.pairwise_inters, y.pairwise_inters), - specific_inter_lists=Zygote.accum(x.specific_inter_lists, y.specific_inter_lists), - general_inters=Zygote.accum(x.general_inters, y.general_inters), - constraints=Zygote.accum(x.constraints, y.constraints), - neighbor_finder=nothing, - loggers=nothing, - df=nothing, - force_units=nothing, - energy_units=nothing, - k=Zygote.accum(x.k, y.k), - masses=Zygote.accum(x.masses, y.masses), - data=nothing, - ) -end - -function Base.:+(y::NamedTuple{(:atoms, :coords, :boundary, - :velocities, :atoms_data, :topology, :pairwise_inters, :specific_inter_lists, - :general_inters, :constraints, :neighbor_finder, :loggers, :df, :force_units, - :energy_units, :k, :masses, :data)}, r::Base.RefValue{Any}) - return r + y -end - -function Zygote.accum(x::NamedTuple{(:side_lengths,), Tuple{SVector{3, T}}}, y::SVector{3, T}) where T - CubicBoundary(x.side_lengths .+ y; check_positive=false) -end - -function Zygote.accum(x::NamedTuple{(:side_lengths,), Tuple{SVector{2, T}}}, y::SVector{2, T}) where T - RectangularBoundary(x.side_lengths .+ y; check_positive=false) -end - -function Zygote.accum(x::NamedTuple{(:side_lengths,), Tuple{SVector{3, T}}}, y::SizedVector{3, T, Vector{T}}) where T - CubicBoundary(x.side_lengths .+ y; check_positive=false) -end - -function Zygote.accum(x::NamedTuple{(:side_lengths,), Tuple{SVector{2, T}}}, y::SizedVector{2, T, Vector{T}}) where T - RectangularBoundary(x.side_lengths .+ y; check_positive=false) -end - -function Zygote.accum(x::NamedTuple{(:side_lengths,), Tuple{SizedVector{3, T, Vector{T}}}}, y::SVector{3, T}) where T - CubicBoundary(SVector{3, T}(x.side_lengths .+ y); check_positive=false) -end - -function Zygote.accum(x::NamedTuple{(:side_lengths,), Tuple{SizedVector{2, T, Vector{T}}}}, y::SVector{2, T}) where T - RectangularBoundary(SVector{2, T}(x.side_lengths .+ y); check_positive=false) -end - -function Base.:+(x::NamedTuple{(:side_lengths,), Tuple{SizedVector{3, T, Vector{T}}}}, y::CubicBoundary{T}) where T - CubicBoundary(SVector{3, T}(x.side_lengths .+ y.side_lengths); check_positive=false) -end - -function Base.:+(x::NamedTuple{(:side_lengths,), Tuple{SizedVector{2, T, Vector{T}}}}, y::RectangularBoundary{T}) where T - RectangularBoundary(SVector{2, T}(x.side_lengths .+ y.side_lengths); check_positive=false) -end - -function Base.:+(x::CubicBoundary{T}, y::NamedTuple{(:side_lengths,), Tuple{SVector{3, T}}}) where T - CubicBoundary(SVector{3, T}(x.side_lengths .+ y.side_lengths); check_positive=false) -end - -function Base.:+(x::RectangularBoundary{T}, y::NamedTuple{(:side_lengths,), Tuple{SVector{2, T}}}) where T - RectangularBoundary(SVector{2, T}(x.side_lengths .+ y.side_lengths); check_positive=false) -end - -function Base.:+(x::SVector{3, T}, y::CubicBoundary{T}) where T - CubicBoundary(x .+ y.side_lengths; check_positive=false) -end - -function Base.:+(x::SVector{2, T}, y::RectangularBoundary{T}) where T - RectangularBoundary(x .+ y.side_lengths; check_positive=false) -end - -Base.:+(x::CubicBoundary{T}, y::SVector{3, T}) where {T} = y + x -Base.:+(x::RectangularBoundary{T}, y::SVector{2, T}) where {T} = y + x - -atom_or_empty(at::Atom, T) = at -atom_or_empty(at::Nothing, T) = zero(Atom{T, T, T, T}) - -Zygote.z2d(dx::AbstractArray{Union{Nothing, Atom{T, T, T, T}}}, primal::AbstractArray{Atom{T, T, T, T}}) where {T} = atom_or_empty.(dx, T) -Zygote.z2d(dx::SVector{3, T}, primal::T) where {T} = sum(dx) - -function Zygote.unbroadcast(x::AbstractArray{<:Real}, x̄::AbstractArray{<:StaticVector}) - if length(x) == length(x̄) - Zygote._project(x, sum.(x̄)) - else - dims = ntuple(d -> size(x, d) == 1 ? d : ndims(x̄) + 1, ndims(x̄)) - Zygote._project(x, accum_sum(x̄; dims=dims)) - end -end - -Zygote._zero(xs::AbstractArray{<:StaticVector}, T) = fill!(similar(xs, T), zero(T)) - -function Zygote._zero(xs::AbstractArray{Atom{T, T, T, T}}, ::Type{Atom{T, T, T, T}}) where T - fill!(similar(xs), Atom{T, T, T, T}(0, zero(T), zero(T), zero(T), zero(T), false)) -end - -function Base.zero(::Type{Union{Nothing, SizedVector{D, T, Vector{T}}}}) where {D, T} - zero(SizedVector{D, T, Vector{T}}) -end - -# ChainRules._setindex_zero returns a union type with NoTangent via the default path -# This causes a problem on the GPU -function ChainRules.∇getindex(x::CuArray{<:StaticVector}, dy, inds...) - plain_inds = Base.to_indices(x, inds) - dx = ChainRules.∇getindex!(zero(x), dy, plain_inds...) - return ChainRules.ProjectTo(x)(dx) -end - -# Modified version of ForwardDiff.ForwardDiffStaticArraysExt.dualize -# Extend to add extra empty partials before (B) and after (A) the SVector partials -@generated function dualize(::Type{T}, x::StaticArray, ::Val{B}, ::Val{A}) where {T, B, A} - N = length(x) - dx = Expr(:tuple, [:(Dual{T}(x[$i], chunk, Val{$i + $B}())) for i in 1:N]...) - V = StaticArrays.similar_type(x, Dual{T, eltype(x), N + B + A}) - return quote - chunk = Chunk{$N + $B + $A}() - $(Expr(:meta, :inline)) - return $V($(dx)) - end -end - -@inline function sum_partials(sv::SVector{3, Dual{Nothing, T, P}}, y1, i::Integer) where {T, P} - partials(sv[1], i) * y1[1] + partials(sv[2], i) * y1[2] + partials(sv[3], i) * y1[3] -end - -sized_to_static(v::SizedVector{3, T, Vector{T}}) where {T} = SVector{3, T}(v[1], v[2], v[3]) -sized_to_static(v::SizedVector{2, T, Vector{T}}) where {T} = SVector{2, T}(v[1], v[2]) - -function modify_grad(ȳ_in::AbstractArray{SizedVector{D, T, Vector{T}}}, arg::CuArray) where {D, T} - CuArray(sized_to_static.(ȳ_in)) -end - -function modify_grad(ȳ_in::AbstractArray{SizedVector{D, T, Vector{T}}}, arg) where {D, T} - sized_to_static.(ȳ_in) -end - -modify_grad(ȳ_in, arg::CuArray) = CuArray(ȳ_in) -modify_grad(ȳ_in, arg) = ȳ_in - -# Dualize a value with extra partials -macro dualize(x, n_partials::Integer, active_partial::Integer) - ps = [i == active_partial for i in 1:n_partials] - return :(ForwardDiff.Dual($(esc(x)), $(ps...))) -end - -function dual_function_svec(f::F) where F - function (arg1) - ds1 = dualize(Nothing, arg1, Val(0), Val(0)) - return f(ds1) - end -end - -@inline function Zygote.broadcast_forward(f, arg1::AbstractArray{SVector{D, T}}) where {D, T} - out = dual_function_svec(f).(arg1) - y = map(x -> value.(x), out) - function bc_fwd_back(ȳ_in) - ȳ = modify_grad(ȳ_in, arg1) - barg1 = broadcast(ȳ, out) do y1, o1 - if length(y1) == 1 - y1 .* SVector{D, T}(partials(o1)) - else - SVector{D, T}(sum_partials(o1, y1, 1), sum_partials(o1, y1, 2), sum_partials(o1, y1, 3)) - end - end - darg1 = unbroadcast(arg1, barg1) - (nothing, nothing, darg1) - end - return y, bc_fwd_back -end - -function dual_function_svec_real(f::F) where F - function (arg1::SVector{D, T}, arg2) where {D, T} - ds1 = dualize(Nothing, arg1, Val(0), Val(1)) - # Leaving the integer type in here results in Float32 -> Float64 conversion - ds2 = Zygote.dual(arg2 isa Int ? T(arg2) : arg2, 4, Val(4)) - return f(ds1, ds2) - end -end - -@inline function Zygote.broadcast_forward(f, arg1::AbstractArray{SVector{D, T}}, arg2) where {D, T} - out = dual_function_svec_real(f).(arg1, arg2) - y = map(x -> value.(x), out) - function bc_fwd_back(ȳ_in) - ȳ = modify_grad(ȳ_in, arg1) - barg1 = broadcast(ȳ, out) do y1, o1 - if length(y1) == 1 - y1 .* SVector{D, T}(partials.((o1,), (1, 2, 3))) - else - SVector{D, T}(sum_partials(o1, y1, 1), sum_partials(o1, y1, 2), sum_partials(o1, y1, 3)) - end - end - darg1 = unbroadcast(arg1, barg1) - darg2 = unbroadcast(arg2, broadcast((y1, o1) -> y1 .* partials.(o1, 4), ȳ, out)) - (nothing, nothing, darg1, darg2) - end - return y, bc_fwd_back -end - -function dual_function_real_svec(f::F) where F - function (arg1, arg2::SVector{D, T}) where {D, T} - ds1 = Zygote.dual(arg1 isa Int ? T(arg1) : arg1, 1, Val(4)) - ds2 = dualize(Nothing, arg2, Val(1), Val(0)) - return f(ds1, ds2) - end -end - -@inline function Zygote.broadcast_forward(f, - arg1::Union{AbstractArray{R}, Tuple{R}, R}, - arg2::AbstractArray{SVector{D, T}}) where {D, T, R <: Real} - out = dual_function_real_svec(f).(arg1, arg2) - y = map(x -> value.(x), out) - function bc_fwd_back(ȳ_in) - ȳ = modify_grad(ȳ_in, arg2) - darg1 = unbroadcast(arg1, broadcast((y1, o1) -> y1 .* partials.(o1, 1), ȳ, out)) - barg2 = broadcast(ȳ, out) do y1, o1 - if length(y1) == 1 - y1 .* SVector{D, T}(partials.((o1,), (2, 3, 4))) - else - SVector{D, T}(sum_partials(o1, y1, 2), sum_partials(o1, y1, 3), sum_partials(o1, y1, 4)) - end - end - darg2 = unbroadcast(arg2, barg2) - (nothing, nothing, darg1, darg2) - end - return y, bc_fwd_back -end - -function dual_function_svec_svec(f::F) where F - function (arg1, arg2) - ds1 = dualize(Nothing, arg1, Val(0), Val(3)) - ds2 = dualize(Nothing, arg2, Val(3), Val(0)) - return f(ds1, ds2) - end -end - -@inline function Zygote.broadcast_forward(f, - arg1::AbstractArray{SVector{D, T}}, - arg2::Union{AbstractArray{SVector{D, T}}, Tuple{SVector{D, T}}}) where {D, T} - out = dual_function_svec_svec(f).(arg1, arg2) - y = map(x -> value.(x), out) - function bc_fwd_back(ȳ_in) - ȳ = modify_grad(ȳ_in, arg1) - barg1 = broadcast(ȳ, out) do y1, o1 - if length(y1) == 1 - y1 .* SVector{D, T}(partials.((o1,), (1, 2, 3))) - else - SVector{D, T}(sum_partials(o1, y1, 1), sum_partials(o1, y1, 2), sum_partials(o1, y1, 3)) - end - end - darg1 = unbroadcast(arg1, barg1) - barg2 = broadcast(ȳ, out) do y1, o1 - if length(y1) == 1 - y1 .* SVector{D, T}(partials.((o1,), (4, 5, 6))) - else - SVector{D, T}(sum_partials(o1, y1, 4), sum_partials(o1, y1, 5), sum_partials(o1, y1, 6)) - end - end - darg2 = unbroadcast(arg2, barg2) - (nothing, nothing, darg1, darg2) - end - return y, bc_fwd_back -end - -function dual_function_atom(f::F) where F - function (arg1) - c, m, σ, ϵ = arg1.charge, arg1.mass, arg1.σ, arg1.ϵ - ds1 = Atom( - arg1.index, - @dualize(c, 4, 1), - @dualize(m, 4, 2), - @dualize(σ, 4, 3), - @dualize(ϵ, 4, 4), - arg1.solute, - ) - return f(ds1) - end -end - -# For mass, charge etc. -@inline function Zygote.broadcast_forward(f, arg1::AbstractArray{<:Atom}) - out = dual_function_atom(f).(arg1) - y = map(x -> value.(x), out) - function bc_fwd_back(ȳ_in) - ȳ = modify_grad(ȳ_in, arg1) - barg1 = broadcast(ȳ, out) do y1, o1 - ps = partials(o1) - Atom(0, y1 * ps[1], y1 * ps[2], y1 * ps[3], y1 * ps[4], false) - end - darg1 = unbroadcast(arg1, barg1) - (nothing, nothing, darg1) - end - return y, bc_fwd_back -end - -function dual_function_born_radii_loop_OBC(f::F) where F - function (arg1, arg2, arg3, arg4, arg5, arg6) - ds1 = dualize(Nothing, arg1, Val(0), Val(5)) - ds2 = dualize(Nothing, arg2, Val(3), Val(2)) - ds3 = Zygote.dual(arg3, 7, Val(8)) - ds4 = Zygote.dual(arg4, 8, Val(8)) - ds5 = arg5 - ds6 = arg6 - return f(ds1, ds2, ds3, ds4, ds5, ds6) - end -end - -# For born_radii_loop_OBC -@inline function Zygote.broadcast_forward(f, - arg1::AbstractArray{SVector{D, T}}, - arg2::AbstractArray{SVector{D, T}}, - arg3::AbstractArray{T}, - arg4::AbstractArray{T}, - arg5::T, - arg6) where {D, T} - out = dual_function_born_radii_loop_OBC(f).(arg1, arg2, arg3, arg4, arg5, arg6) - y = value.(out) - function bc_fwd_back(ȳ_in) - ȳ = modify_grad(ȳ_in, arg1) - darg1 = unbroadcast(arg1, broadcast((y1, o1) -> SVector{D, T}(partials(o1, 1) * y1, - partials(o1, 2) * y1, partials(o1, 3) * y1), ȳ, out)) - darg2 = unbroadcast(arg2, broadcast((y1, o1) -> SVector{D, T}(partials(o1, 4) * y1, - partials(o1, 5) * y1, partials(o1, 6) * y1), ȳ, out)) - darg3 = unbroadcast(arg3, broadcast((y1, o1) -> partials(o1, 7) * y1, ȳ, out)) - darg4 = unbroadcast(arg4, broadcast((y1, o1) -> partials(o1, 8) * y1, ȳ, out)) - darg5 = nothing - darg6 = nothing - return (nothing, nothing, darg1, darg2, darg3, darg4, darg5, darg6) - end - return y, bc_fwd_back -end - -# For born_radii_sum -@inline function Zygote.broadcast_forward(f, - arg1::AbstractArray{T}, - arg2::T, - arg3::AbstractArray{T}, - arg4, - arg5, - arg6) where T - out = Zygote.dual_function(f).(arg1, arg2, arg3, arg4, arg5, arg6) - y = broadcast(o1 -> (value(o1[1]), value(o1[2])), out) - function bc_fwd_back(ȳ_in) - ȳ = modify_grad(ȳ_in, arg1) - darg1 = unbroadcast(arg1, broadcast((y1, o1) -> partials(o1[1], 1) * y1[1] + partials(o1[2], 1) * y1[2], ȳ, out)) - darg2 = unbroadcast(arg2, broadcast((y1, o1) -> partials(o1[1], 2) * y1[1] + partials(o1[2], 2) * y1[2], ȳ, out)) - darg3 = unbroadcast(arg3, broadcast((y1, o1) -> partials(o1[1], 3) * y1[1] + partials(o1[2], 3) * y1[2], ȳ, out)) - darg4 = unbroadcast(arg4, broadcast((y1, o1) -> partials(o1[1], 4) * y1[1] + partials(o1[2], 4) * y1[2], ȳ, out)) - darg5 = unbroadcast(arg5, broadcast((y1, o1) -> partials(o1[1], 5) * y1[1] + partials(o1[2], 5) * y1[2], ȳ, out)) - darg6 = unbroadcast(arg6, broadcast((y1, o1) -> partials(o1[1], 6) * y1[1] + partials(o1[2], 6) * y1[2], ȳ, out)) - return (nothing, nothing, darg1, darg2, darg3, darg4, darg5, darg6) - end - return y, bc_fwd_back -end - -function dual_function_born_radii_loop_GBN2(f::F) where F - function (arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10, arg11, arg12) - ds1 = dualize(Nothing, arg1, Val(0), Val(11)) - ds2 = dualize(Nothing, arg2, Val(3), Val(8)) - ds3 = Zygote.dual(arg3, 7 , Val(14)) - ds4 = Zygote.dual(arg4, 8 , Val(14)) - ds5 = Zygote.dual(arg5, 9 , Val(14)) - ds6 = arg6 - ds7 = Zygote.dual(arg7, 10, Val(14)) - ds8 = Zygote.dual(arg8, 11, Val(14)) - ds9 = Zygote.dual(arg9, 12, Val(14)) - ds10 = Zygote.dual(arg10, 13, Val(14)) - ds11 = Zygote.dual(arg11, 14, Val(14)) - ds12 = arg12 - return f(ds1, ds2, ds3, ds4, ds5, ds6, ds7, ds8, ds9, ds10, ds11, ds12) - end -end - -# For born_radii_loop_GBN2 -@inline function Zygote.broadcast_forward(f, - arg1::AbstractArray{SVector{D, T}}, - arg2::AbstractArray{SVector{D, T}}, - arg3::AbstractArray{T}, - arg4::AbstractArray{T}, - arg5::AbstractArray{T}, - arg6::T, - arg7::T, - arg8::T, - arg9::T, - arg10::AbstractArray{T}, - arg11::AbstractArray{T}, - arg12) where {D, T} - out = dual_function_born_radii_loop_GBN2(f).(arg1, arg2, arg3, arg4, arg5, arg6, - arg7, arg8, arg9, arg10, arg11, arg12) - y = broadcast(o1 -> BornRadiiGBN2LoopResult{T, T}(value(o1.I), value(o1.I_grad)), out) - function bc_fwd_back(ȳ_in) - ȳ = modify_grad(ȳ_in, arg1) - darg1 = unbroadcast(arg1, broadcast((y1, o1) -> SVector{D, T}( - partials(o1.I, 1) * y1.I + partials(o1.I_grad, 1) * y1.I_grad, - partials(o1.I, 2) * y1.I + partials(o1.I_grad, 2) * y1.I_grad, - partials(o1.I, 3) * y1.I + partials(o1.I_grad, 3) * y1.I_grad), - ȳ, out)) - darg2 = unbroadcast(arg2, broadcast((y1, o1) -> SVector{D, T}( - partials(o1.I, 4) * y1.I + partials(o1.I_grad, 4) * y1.I_grad, - partials(o1.I, 5) * y1.I + partials(o1.I_grad, 5) * y1.I_grad, - partials(o1.I, 6) * y1.I + partials(o1.I_grad, 6) * y1.I_grad), - ȳ, out)) - darg3 = unbroadcast(arg3, broadcast((y1, o1) -> partials(o1.I, 7) * y1.I + partials(o1.I_grad, 7) * y1.I_grad, ȳ, out)) - darg4 = unbroadcast(arg4, broadcast((y1, o1) -> partials(o1.I, 8) * y1.I + partials(o1.I_grad, 8) * y1.I_grad, ȳ, out)) - darg5 = unbroadcast(arg5, broadcast((y1, o1) -> partials(o1.I, 9) * y1.I + partials(o1.I_grad, 9) * y1.I_grad, ȳ, out)) - darg6 = nothing - darg7 = unbroadcast(arg7, broadcast((y1, o1) -> partials(o1.I, 10) * y1.I + partials(o1.I_grad, 10) * y1.I_grad, ȳ, out)) - darg8 = unbroadcast(arg8, broadcast((y1, o1) -> partials(o1.I, 11) * y1.I + partials(o1.I_grad, 11) * y1.I_grad, ȳ, out)) - darg9 = unbroadcast(arg9, broadcast((y1, o1) -> partials(o1.I, 12) * y1.I + partials(o1.I_grad, 12) * y1.I_grad, ȳ, out)) - darg10 = unbroadcast(arg10, broadcast((y1, o1) -> partials(o1.I, 13) * y1.I + partials(o1.I_grad, 13) * y1.I_grad, ȳ, out)) - darg11 = unbroadcast(arg11, broadcast((y1, o1) -> partials(o1.I, 14) * y1.I + partials(o1.I_grad, 14) * y1.I_grad, ȳ, out)) - darg12 = nothing - return (nothing, nothing, darg1, darg2, darg3, darg4, darg5, darg6, - darg7, darg8, darg9, darg10, darg11, darg12) - end - return y, bc_fwd_back -end - -function dual_function_gb_force_loop_1(f::F) where F - function (arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10, arg11, arg12, arg13) - ds1 = dualize(Nothing, arg1, Val(0), Val(10)) - ds2 = dualize(Nothing, arg2, Val(3), Val(7)) - ds3 = arg3 - ds4 = arg4 - ds5 = Zygote.dual(arg5 , 7 , Val(13)) - ds6 = Zygote.dual(arg6 , 8 , Val(13)) - ds7 = Zygote.dual(arg7 , 9 , Val(13)) - ds8 = Zygote.dual(arg8 , 10, Val(13)) - ds9 = arg9 - ds10 = Zygote.dual(arg10, 11, Val(13)) - ds11 = Zygote.dual(arg11, 12, Val(13)) - ds12 = Zygote.dual(arg12, 13, Val(13)) - ds13 = arg13 - return f(ds1, ds2, ds3, ds4, ds5, ds6, ds7, ds8, ds9, ds10, ds11, ds12, ds13) - end -end - -# For gb_force_loop_1 -@inline function Zygote.broadcast_forward(f, - arg1::AbstractArray{SVector{D, T}}, - arg2::AbstractArray{SVector{D, T}}, - arg3::AbstractArray{Int}, - arg4::AbstractArray{Int}, - arg5::AbstractArray{T}, - arg6::AbstractArray{T}, - arg7::AbstractArray{T}, - arg8::AbstractArray{T}, - arg9::T, - arg10::T, - arg11::T, - arg12::T, - arg13) where {D, T} - out = dual_function_gb_force_loop_1(f).(arg1, arg2, arg3, arg4, arg5, arg6, arg7, - arg8, arg9, arg10, arg11, arg12, arg13) - y = broadcast(o1 -> ForceLoopResult1{T, SVector{D, T}}(value(o1.bi), value(o1.bj), - value.(o1.fi), value.(o1.fj)), out) - function bc_fwd_back(ȳ_in) - ȳ = modify_grad(ȳ_in, arg1) - darg1 = unbroadcast(arg1, broadcast((y1, o1) -> SVector{D, T}( - sum_partials(o1.fi, y1.fi, 1) + sum_partials(o1.fj, y1.fj, 1) + partials(o1.bi, 1) * y1.bi + partials(o1.bj, 1) * y1.bj, - sum_partials(o1.fi, y1.fi, 2) + sum_partials(o1.fj, y1.fj, 2) + partials(o1.bi, 2) * y1.bi + partials(o1.bj, 2) * y1.bj, - sum_partials(o1.fi, y1.fi, 3) + sum_partials(o1.fj, y1.fj, 3) + partials(o1.bi, 3) * y1.bi + partials(o1.bj, 3) * y1.bj), - ȳ, out)) - darg2 = unbroadcast(arg2, broadcast((y1, o1) -> SVector{D, T}( - sum_partials(o1.fi, y1.fi, 4) + sum_partials(o1.fj, y1.fj, 4) + partials(o1.bi, 4) * y1.bi + partials(o1.bj, 4) * y1.bj, - sum_partials(o1.fi, y1.fi, 5) + sum_partials(o1.fj, y1.fj, 5) + partials(o1.bi, 5) * y1.bi + partials(o1.bj, 5) * y1.bj, - sum_partials(o1.fi, y1.fi, 6) + sum_partials(o1.fj, y1.fj, 6) + partials(o1.bi, 6) * y1.bi + partials(o1.bj, 6) * y1.bj), - ȳ, out)) - darg3 = nothing - darg4 = nothing - darg5 = unbroadcast(arg5, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 7) + sum_partials(o1.fj, y1.fj, 7) + partials(o1.bi, 7) * y1.bi + partials(o1.bj, 7) * y1.bj, ȳ, out)) - darg6 = unbroadcast(arg6, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 8) + sum_partials(o1.fj, y1.fj, 8) + partials(o1.bi, 8) * y1.bi + partials(o1.bj, 8) * y1.bj, ȳ, out)) - darg7 = unbroadcast(arg7, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 9) + sum_partials(o1.fj, y1.fj, 9) + partials(o1.bi, 9) * y1.bi + partials(o1.bj, 9) * y1.bj, ȳ, out)) - darg8 = unbroadcast(arg8, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 10) + sum_partials(o1.fj, y1.fj, 10) + partials(o1.bi, 10) * y1.bi + partials(o1.bj, 10) * y1.bj, ȳ, out)) - darg9 = nothing - darg10 = unbroadcast(arg10, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 11) + sum_partials(o1.fj, y1.fj, 11) + partials(o1.bi, 11) * y1.bi + partials(o1.bj, 11) * y1.bj, ȳ, out)) - darg11 = unbroadcast(arg11, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 12) + sum_partials(o1.fj, y1.fj, 12) + partials(o1.bi, 12) * y1.bi + partials(o1.bj, 12) * y1.bj, ȳ, out)) - darg12 = unbroadcast(arg12, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 13) + sum_partials(o1.fj, y1.fj, 13) + partials(o1.bi, 13) * y1.bi + partials(o1.bj, 13) * y1.bj, ȳ, out)) - darg13 = nothing - return (nothing, nothing, darg1, darg2, darg3, darg4, darg5, darg6, darg7, - darg8, darg9, darg10, darg11, darg12, darg13) - end - return y, bc_fwd_back -end - -function dual_function_gb_force_loop_2(f::F) where F - function (arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8) - ds1 = dualize(Nothing, arg1, Val(0), Val(7)) - ds2 = dualize(Nothing, arg2, Val(3), Val(4)) - ds3 = Zygote.dual(arg3, 7 , Val(10)) - ds4 = Zygote.dual(arg4, 8 , Val(10)) - ds5 = Zygote.dual(arg5, 9 , Val(10)) - ds6 = Zygote.dual(arg6, 10, Val(10)) - ds7 = arg7 - ds8 = arg8 - return f(ds1, ds2, ds3, ds4, ds5, ds6, ds7, ds8) - end -end - -# For gb_force_loop_2 -@inline function Zygote.broadcast_forward(f, - arg1::AbstractArray{SVector{D, T}}, - arg2::AbstractArray{SVector{D, T}}, - arg3::AbstractArray{T}, - arg4::AbstractArray{T}, - arg5::AbstractArray{T}, - arg6::AbstractArray{T}, - arg7::T, - arg8) where {D, T} - out = dual_function_gb_force_loop_2(f).(arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8) - y = broadcast(o1 -> ForceLoopResult2{SVector{D, T}}(value.(o1.fi), value.(o1.fj)), out) - function bc_fwd_back(ȳ_in) - ȳ = modify_grad(ȳ_in, arg1) - darg1 = unbroadcast(arg1, broadcast((y1, o1) -> SVector{D, T}( - sum_partials(o1.fi, y1.fi, 1) + sum_partials(o1.fj, y1.fj, 1), - sum_partials(o1.fi, y1.fi, 2) + sum_partials(o1.fj, y1.fj, 2), - sum_partials(o1.fi, y1.fi, 3) + sum_partials(o1.fj, y1.fj, 3)), - ȳ, out)) - darg2 = unbroadcast(arg2, broadcast((y1, o1) -> SVector{D, T}( - sum_partials(o1.fi, y1.fi, 4) + sum_partials(o1.fj, y1.fj, 4), - sum_partials(o1.fi, y1.fi, 5) + sum_partials(o1.fj, y1.fj, 5), - sum_partials(o1.fi, y1.fi, 6) + sum_partials(o1.fj, y1.fj, 6)), - ȳ, out)) - darg3 = unbroadcast(arg3, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 7) + sum_partials(o1.fj, y1.fj, 7), ȳ, out)) - darg4 = unbroadcast(arg4, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 8) + sum_partials(o1.fj, y1.fj, 8), ȳ, out)) - darg5 = unbroadcast(arg5, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 9) + sum_partials(o1.fj, y1.fj, 9), ȳ, out)) - darg6 = unbroadcast(arg6, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 10) + sum_partials(o1.fj, y1.fj, 10), ȳ, out)) - darg7 = nothing - darg8 = nothing - return (nothing, nothing, darg1, darg2, darg3, darg4, darg5, darg6, darg7, darg8) - end - return y, bc_fwd_back -end - -function dual_function_gb_energy_loop(f::F) where F - function (arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10, - arg11, arg12, arg13, arg14, arg15, arg16, arg17, arg18) - ds1 = dualize(Nothing, arg1, Val(0), Val(14)) - ds2 = dualize(Nothing, arg2, Val(3), Val(11)) - ds3 = arg3 - ds4 = arg4 - # Using Zygote.dual errors on GPU so Dual is called explicitly - ds5 = Dual(arg5 , (false, false, false, false, false, false, true , false, false, false, false, false, false, false, false, false, false)) - ds6 = Dual(arg6 , (false, false, false, false, false, false, false, true , false, false, false, false, false, false, false, false, false)) - ds7 = Dual(arg7 , (false, false, false, false, false, false, false, false, true , false, false, false, false, false, false, false, false)) - ds8 = Dual(arg8 , (false, false, false, false, false, false, false, false, false, true , false, false, false, false, false, false, false)) - ds9 = Dual(arg9 , (false, false, false, false, false, false, false, false, false, false, true , false, false, false, false, false, false)) - ds10 = arg10 - ds11 = Dual(arg11, (false, false, false, false, false, false, false, false, false, false, false, true , false, false, false, false, false)) - ds12 = Dual(arg12, (false, false, false, false, false, false, false, false, false, false, false, false, true , false, false, false, false)) - ds13 = Dual(arg13, (false, false, false, false, false, false, false, false, false, false, false, false, false, true , false, false, false)) - ds14 = Dual(arg14, (false, false, false, false, false, false, false, false, false, false, false, false, false, false, true , false, false)) - ds15 = Dual(arg15, (false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, true , false)) - ds16 = Dual(arg16, (false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, true )) - ds17 = arg17 - ds18 = arg18 - return f(ds1, ds2, ds3, ds4, ds5, ds6, ds7, ds8, ds9, ds10, - ds11, ds12, ds13, ds14, ds15, ds16, ds17, ds18) - end -end - -# For gb_energy_loop -@inline function Zygote.broadcast_forward(f, - arg1::AbstractArray{SVector{D, T}}, - arg2::AbstractArray{SVector{D, T}}, - arg3::AbstractArray{Int}, - arg4::AbstractArray{Int}, - arg5::AbstractArray{T}, - arg6::AbstractArray{T}, - arg7::AbstractArray{T}, - arg8::AbstractArray{T}, - arg9::AbstractArray{T}, - arg10::T, - arg11::T, - arg12::T, - arg13::T, - arg14::T, - arg15::T, - arg16::T, - arg17::Bool, - arg18) where {D, T} - out = dual_function_gb_energy_loop(f).(arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, - arg10, arg11, arg12, arg13, arg14, arg15, arg16, arg17, arg18) - y = value.(out) - function bc_fwd_back(ȳ_in) - ȳ = modify_grad(ȳ_in, arg1) - darg1 = unbroadcast(arg1, broadcast((y1, o1) -> SVector{D, T}(partials(o1, 1) * y1, - partials(o1, 2) * y1, partials(o1, 3) * y1), ȳ, out)) - darg2 = unbroadcast(arg2, broadcast((y1, o1) -> SVector{D, T}(partials(o1, 4) * y1, - partials(o1, 5) * y1, partials(o1, 6) * y1), ȳ, out)) - darg3 = nothing - darg4 = nothing - darg5 = unbroadcast(arg5, broadcast((y1, o1) -> partials(o1, 7) * y1, ȳ, out)) - darg6 = unbroadcast(arg6, broadcast((y1, o1) -> partials(o1, 8) * y1, ȳ, out)) - darg7 = unbroadcast(arg7, broadcast((y1, o1) -> partials(o1, 9) * y1, ȳ, out)) - darg8 = unbroadcast(arg8, broadcast((y1, o1) -> partials(o1, 10) * y1, ȳ, out)) - darg9 = unbroadcast(arg9, broadcast((y1, o1) -> partials(o1, 11) * y1, ȳ, out)) - darg10 = nothing - darg11 = unbroadcast(arg11, broadcast((y1, o1) -> partials(o1, 12) * y1, ȳ, out)) - darg12 = unbroadcast(arg12, broadcast((y1, o1) -> partials(o1, 13) * y1, ȳ, out)) - darg13 = unbroadcast(arg13, broadcast((y1, o1) -> partials(o1, 14) * y1, ȳ, out)) - darg14 = unbroadcast(arg14, broadcast((y1, o1) -> partials(o1, 15) * y1, ȳ, out)) - darg15 = unbroadcast(arg15, broadcast((y1, o1) -> partials(o1, 16) * y1, ȳ, out)) - darg16 = unbroadcast(arg16, broadcast((y1, o1) -> partials(o1, 17) * y1, ȳ, out)) - darg17 = nothing - darg18 = nothing - return (nothing, nothing, darg1, darg2, darg3, darg4, darg5, darg6, darg7, darg8, darg9, - darg10, darg11, darg12, darg13, darg14, darg15, darg16, darg17, darg18) - end - return y, bc_fwd_back -end - -@inline function Zygote.broadcast_forward(f::typeof(get_i1), arg1) - return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, broadcast(y1 -> (y1, zero(y1)), ȳ))) -end - -@inline function Zygote.broadcast_forward(f::typeof(get_i2), arg1) - return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, broadcast(y1 -> (zero(y1), y1), ȳ))) -end - -@inline function Zygote.broadcast_forward(f::typeof(get_I), - arg1::AbstractArray{<:BornRadiiGBN2LoopResult}) - return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, broadcast(y1 -> (I=y1, I_grad=zero(y1)), ȳ))) -end - -@inline function Zygote.broadcast_forward(f::typeof(get_I_grad), - arg1::AbstractArray{<:BornRadiiGBN2LoopResult}) - return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, broadcast(y1 -> (I=zero(y1), I_grad=y1), ȳ))) -end - -@inline function Zygote.broadcast_forward(f::typeof(get_bi), - arg1::AbstractArray{<:ForceLoopResult1}) - return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, - broadcast(y1 -> (bi=y1, bj=zero(y1), fi=zero(y1), fj=zero(y1)), ȳ))) -end - -@inline function Zygote.broadcast_forward(f::typeof(get_bj), - arg1::AbstractArray{<:ForceLoopResult1}) - return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, - broadcast(y1 -> (bi=zero(y1), bj=y1, fi=zero(y1), fj=zero(y1)), ȳ))) -end - -@inline function Zygote.broadcast_forward(f::typeof(get_fi), - arg1::AbstractArray{<:ForceLoopResult1}) - return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, - broadcast(y1 -> (bi=zero(y1), bj=zero(y1), fi=y1, fj=zero(y1)), ȳ))) -end - -@inline function Zygote.broadcast_forward(f::typeof(get_fj), - arg1::AbstractArray{<:ForceLoopResult1}) - return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, - broadcast(y1 -> (bi=zero(y1), bj=zero(y1), fi=zero(y1), fj=y1), ȳ))) -end - -@inline function Zygote.broadcast_forward(f::typeof(get_fi), - arg1::AbstractArray{<:ForceLoopResult2}) - return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, broadcast(y1 -> (fi=y1, fj=zero(y1)), ȳ))) -end - -@inline function Zygote.broadcast_forward(f::typeof(get_fj), - arg1::AbstractArray{<:ForceLoopResult2}) - return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, broadcast(y1 -> (fi=zero(y1), fj=y1), ȳ))) -end - -# Use fast broadcast path on CPU -for op in (:+, :-, :*, :/, :mass, :charge, :ustrip, :ustrip_vec, :wrap_coords, - :born_radii_loop_OBC, :get_i1, :get_i2, :get_I, :get_I_grad, :born_radii_loop_GBN2, - :get_bi, :get_bj, :get_fi, :get_fj, :gb_force_loop_1, :gb_force_loop_2, :gb_energy_loop) - @eval Zygote.@adjoint Broadcast.broadcasted(::Broadcast.AbstractArrayStyle, f::typeof($op), args...) = Zygote.broadcast_forward(f, args...) - # Avoid ambiguous dispatch - @eval Zygote.@adjoint Broadcast.broadcasted(::CUDA.AbstractGPUArrayStyle , f::typeof($op), args...) = Zygote.broadcast_forward(f, args...) -end From 9467d76da1294b39fc685391d2b065aaa7bd4fd7 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 28 Jun 2024 14:37:58 +0100 Subject: [PATCH 07/74] remove chain rules --- src/chain_rules.jl | 917 --------------------------------------------- 1 file changed, 917 deletions(-) delete mode 100644 src/chain_rules.jl diff --git a/src/chain_rules.jl b/src/chain_rules.jl deleted file mode 100644 index b40d05c53..000000000 --- a/src/chain_rules.jl +++ /dev/null @@ -1,917 +0,0 @@ -# Chain rules to allow differentiable simulations - -@non_differentiable CUDA.zeros(args...) -@non_differentiable n_infinite_dims(args...) -@non_differentiable random_velocities(args...) -@non_differentiable random_velocities!(args...) -@non_differentiable cuda_threads_blocks_pairwise(args...) -@non_differentiable cuda_threads_blocks_specific(args...) -@non_differentiable check_force_units(args...) -@non_differentiable atoms_bonded_to_N(args...) -@non_differentiable lookup_table(args...) -@non_differentiable cuda_threads_blocks_gbsa(args...) -@non_differentiable find_neighbors(args...) -@non_differentiable DistanceNeighborFinder(args...) -@non_differentiable run_loggers!(args...) -@non_differentiable visualize(args...) -@non_differentiable place_atoms(args...) -@non_differentiable place_diatomics(args...) -@non_differentiable MolecularForceField(T::Type, ff_files::AbstractString...) -@non_differentiable MolecularForceField(ff_files::AbstractString...) -@non_differentiable System(coord_file::AbstractString, force_field::MolecularForceField) -@non_differentiable System(T::Type, coord_file::AbstractString, top_file::AbstractString) -@non_differentiable System(coord_file::AbstractString, top_file::AbstractString) - -function ChainRulesCore.rrule(T::Type{<:Atom}, vs...) - Y = T(vs...) - function Atom_pullback(Ȳ) - return NoTangent(), Ȳ.index, Ȳ.charge, Ȳ.mass, Ȳ.σ, Ȳ.ϵ, Ȳ.solute - end - return Y, Atom_pullback -end - -function ChainRulesCore.rrule(T::Type{<:SpecificInteraction}, vs...) - Y = T(vs...) - function SpecificInteraction_pullback(Ȳ) - return NoTangent(), Ȳ... - end - return Y, SpecificInteraction_pullback -end - -function ChainRulesCore.rrule(T::Type{<:PairwiseInteraction}, vs...) - Y = T(vs...) - function PairwiseInteraction_pullback(Ȳ) - return NoTangent(), getfield.((Ȳ,), fieldnames(T))... - end - return Y, PairwiseInteraction_pullback -end - -function ChainRulesCore.rrule(T::Type{<:InteractionList1Atoms}, vs...) - Y = T(vs...) - function InteractionList1Atoms_pullback(Ȳ) - return NoTangent(), NoTangent(), Ȳ.inters, NoTangent() - end - return Y, InteractionList1Atoms_pullback -end - -function ChainRulesCore.rrule(T::Type{<:InteractionList2Atoms}, vs...) - Y = T(vs...) - function InteractionList2Atoms_pullback(Ȳ) - return NoTangent(), NoTangent(), NoTangent(), Ȳ.inters, NoTangent() - end - return Y, InteractionList2Atoms_pullback -end - -function ChainRulesCore.rrule(T::Type{<:InteractionList3Atoms}, vs...) - Y = T(vs...) - function InteractionList3Atoms_pullback(Ȳ) - return NoTangent(), NoTangent(), NoTangent(), NoTangent(), Ȳ.inters, NoTangent() - end - return Y, InteractionList3Atoms_pullback -end - -function ChainRulesCore.rrule(T::Type{<:InteractionList4Atoms}, vs...) - Y = T(vs...) - function InteractionList4Atoms_pullback(Ȳ) - return NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), Ȳ.inters, - NoTangent() - end - return Y, InteractionList4Atoms_pullback -end - -function ChainRulesCore.rrule(T::Type{<:SpecificForce1Atoms}, vs...) - Y = T(vs...) - function SpecificForce1Atoms_pullback(Ȳ) - return NoTangent(), Ȳ.f1 - end - return Y, SpecificForce1Atoms_pullback -end - -function ChainRulesCore.rrule(T::Type{<:SpecificForce2Atoms}, vs...) - Y = T(vs...) - function SpecificForce2Atoms_pullback(Ȳ) - return NoTangent(), Ȳ.f1, Ȳ.f2 - end - return Y, SpecificForce2Atoms_pullback -end - -function ChainRulesCore.rrule(T::Type{<:SpecificForce3Atoms}, vs...) - Y = T(vs...) - function SpecificForce3Atoms_pullback(Ȳ) - return NoTangent(), Ȳ.f1, Ȳ.f2, Ȳ.f3 - end - return Y, SpecificForce3Atoms_pullback -end - -function ChainRulesCore.rrule(T::Type{<:SpecificForce4Atoms}, vs...) - Y = T(vs...) - function SpecificForce4Atoms_pullback(Ȳ) - return NoTangent(), Ȳ.f1, Ȳ.f2, Ȳ.f3, Ȳ.f4 - end - return Y, SpecificForce4Atoms_pullback -end - -# Required for SVector gradients in RescaleThermostat -function ChainRulesCore.rrule(::typeof(sqrt), x::Real) - Y = sqrt(x) - function sqrt_pullback(Ȳ) - return NoTangent(), sum(Ȳ * inv(2 * Y)) - end - return Y, sqrt_pullback -end - -function ChainRulesCore.rrule(::typeof(reinterpret), - ::Type{T}, - arr::SVector{D, T}) where {D, T} - Y = reinterpret(T, arr) - function reinterpret_pullback(Ȳ::Vector{T}) - return NoTangent(), NoTangent(), SVector{D, T}(Ȳ) - end - return Y, reinterpret_pullback -end - -function ChainRulesCore.rrule(::typeof(reinterpret), - ::Type{T}, - arr::AbstractArray{SVector{D, T}}) where {D, T} - Y = reinterpret(T, arr) - function reinterpret_pullback(Ȳ::Vector{T}) - return NoTangent(), NoTangent(), reinterpret(SVector{D, T}, Ȳ) - end - return Y, reinterpret_pullback -end - -function ChainRulesCore.rrule(::typeof(reinterpret), - ::Type{SVector{D, T}}, - arr::AbstractVector{T}) where {D, T} - Y = reinterpret(SVector{D, T}, arr) - function reinterpret_pullback(Ȳ::AbstractArray{SVector{D, T}}) - return NoTangent(), NoTangent(), reinterpret(T, Ȳ) - end - return Y, reinterpret_pullback -end - -function ChainRulesCore.rrule(::typeof(sum_svec), arr::AbstractArray{SVector{D, T}}) where {D, T} - Y = sum_svec(arr) - function sum_svec_pullback(Ȳ::SVector{D, T}) - return NoTangent(), zero(arr) .+ (Ȳ,) - end - return Y, sum_svec_pullback -end - -function ChainRulesCore.rrule(::typeof(mean), arr::AbstractArray{SVector{D, T}}) where {D, T} - Y = mean(arr) - function mean_pullback(Ȳ::SVector{D, T}) - return NoTangent(), zero(arr) .+ (Ȳ ./ length(arr),) - end - return Y, mean_pullback -end - -function ChainRulesCore.rrule(T::Type{<:DistanceCutoff}, dist_cutoff) - Y = T(dist_cutoff) - function DistanceCutoff_pullback(Ȳ) - return NoTangent(), NoTangent() - end - return Y, DistanceCutoff_pullback -end - -function ChainRulesCore.rrule(T::Type{<:HarmonicBond}, vs...) - Y = T(vs...) - function HarmonicBond_pullback(Ȳ) - return NoTangent(), Ȳ.k, Ȳ.r0 - end - return Y, HarmonicBond_pullback -end - -function ChainRulesCore.rrule(T::Type{<:HarmonicAngle}, vs...) - Y = T(vs...) - function HarmonicAngle_pullback(Ȳ) - return NoTangent(), Ȳ.k, Ȳ.θ0 - end - return Y, HarmonicAngle_pullback -end - -function ChainRulesCore.rrule(T::Type{<:PeriodicTorsion}, vs...) - Y = T(vs...) - function PeriodicTorsion_pullback(Ȳ) - return NoTangent(), NoTangent(), Ȳ.phases, Ȳ.ks, NoTangent() - end - return Y, PeriodicTorsion_pullback -end - -duplicated_if_present(x, dx) = length(x) > 0 ? Duplicated(x, dx) : Const(x) -active_if_present(x) = length(x) > 0 ? Active(x) : Const(x) - -nothing_to_notangent(x) = x -nothing_to_notangent(::Nothing) = NoTangent() - -function ChainRulesCore.rrule(::typeof(forces_pair_spec), coords::AbstractArray{SVector{D, T}}, - atoms::AbstractArray{A}, pairwise_inters_nonl, pairwise_inters_nl, - sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, boundary, - force_units, neighbors, n_threads) where {D, T, A} - if force_units != NoUnits - error("taking gradients through force calculation is not compatible with units, " * - "system force units are $force_units") - end - Y = forces_pair_spec(coords, atoms, pairwise_inters_nonl, pairwise_inters_nl, sils_1_atoms, - sils_2_atoms, sils_3_atoms, sils_4_atoms, boundary, force_units, - neighbors, n_threads) - - function forces_pair_spec_pullback(d_forces) - fs = zero(coords) - z = zero(T) - d_coords = zero(coords) - d_atoms = fill(zero(A), length(coords)) - d_sils_1_atoms = zero.(sils_1_atoms) - d_sils_2_atoms = zero.(sils_2_atoms) - d_sils_3_atoms = zero.(sils_3_atoms) - d_sils_4_atoms = zero.(sils_4_atoms) - grads = autodiff( - Enzyme.Reverse, - forces_pair_spec!, - Const, - Duplicated(fs, convert(typeof(fs), d_forces)), - Duplicated(coords, d_coords), - Duplicated(atoms, d_atoms), - active_if_present(pairwise_inters_nonl), - active_if_present(pairwise_inters_nl), - duplicated_if_present(sils_1_atoms, d_sils_1_atoms), - duplicated_if_present(sils_2_atoms, d_sils_2_atoms), - duplicated_if_present(sils_3_atoms, d_sils_3_atoms), - duplicated_if_present(sils_4_atoms, d_sils_4_atoms), - Const(boundary), - Const(force_units), - Const(neighbors), - Const(n_threads), - )[1] - d_pairwise_inters_nonl = nothing_to_notangent(grads[4]) - d_pairwise_inters_nl = nothing_to_notangent(grads[5]) - d_boundary = grads[10] - return NoTangent(), d_coords, d_atoms, d_pairwise_inters_nonl, d_pairwise_inters_nl, - d_sils_1_atoms, d_sils_2_atoms, d_sils_3_atoms, d_sils_4_atoms, - d_boundary, NoTangent(), NoTangent(), NoTangent() - end - - return Y, forces_pair_spec_pullback -end - -function ChainRulesCore.rrule(::typeof(potential_energy_pair_spec), coords, atoms::AbstractArray{A}, - pairwise_inters_nonl, pairwise_inters_nl, sils_1_atoms, sils_2_atoms, - sils_3_atoms, sils_4_atoms, boundary, energy_units, neighbors, - n_threads, val_ft::Val{T}) where {A, T} - if energy_units != NoUnits - error("taking gradients through potential energy calculation is not compatible with " * - "units, system energy units are $energy_units") - end - Y = potential_energy_pair_spec(coords, atoms, pairwise_inters_nonl, pairwise_inters_nl, - sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, boundary, - energy_units, neighbors, n_threads, val_ft) - - function potential_energy_pair_spec_pullback(d_pe_vec) - pe_vec = zeros(T, 1) - z = zero(T) - d_coords = zero(coords) - d_atoms = fill(zero(A), length(coords)) - d_sils_1_atoms = zero.(sils_1_atoms) - d_sils_2_atoms = zero.(sils_2_atoms) - d_sils_3_atoms = zero.(sils_3_atoms) - d_sils_4_atoms = zero.(sils_4_atoms) - grads = autodiff( - Enzyme.Reverse, - potential_energy_pair_spec!, - Const, - Duplicated(pe_vec, [d_pe_vec]), - Duplicated(coords, d_coords), - Duplicated(atoms, d_atoms), - active_if_present(pairwise_inters_nonl), - active_if_present(pairwise_inters_nl), - duplicated_if_present(sils_1_atoms, d_sils_1_atoms), - duplicated_if_present(sils_2_atoms, d_sils_2_atoms), - duplicated_if_present(sils_3_atoms, d_sils_3_atoms), - duplicated_if_present(sils_4_atoms, d_sils_4_atoms), - Const(boundary), - Const(energy_units), - Const(neighbors), - Const(n_threads), - Const(val_ft), - )[1] - d_pairwise_inters_nonl = nothing_to_notangent(grads[4]) - d_pairwise_inters_nl = nothing_to_notangent(grads[5]) - d_boundary = grads[10] - return NoTangent(), d_coords, d_atoms, d_pairwise_inters_nonl, d_pairwise_inters_nl, - d_sils_1_atoms, d_sils_2_atoms, d_sils_3_atoms, d_sils_4_atoms, - d_boundary, NoTangent(), NoTangent(), NoTangent(), NoTangent() - end - - return Y, potential_energy_pair_spec_pullback -end - -function grad_pairwise_force_kernel!(fs_mat, d_fs_mat, coords, d_coords, atoms, d_atoms, - boundary, inters::I, grad_inters, neighbors, val_dims, - val_force_units, ::Val{N}) where {I, N} - shared_grad_inters = CuStaticSharedArray(I, N) - sync_threads() - - grads = Enzyme.autodiff_deferred( - Enzyme.Reverse, - pairwise_force_kernel_nl!, - Const, - Duplicated(fs_mat, d_fs_mat), - Duplicated(coords, d_coords), - Duplicated(atoms, d_atoms), - Const(boundary), - Active(inters), - Const(neighbors), - Const(val_dims), - Const(val_force_units), - )[1] - - tidx = threadIdx().x - shared_grad_inters[tidx] = grads[5] - sync_threads() - - if tidx == 1 - grad_inters_sum = shared_grad_inters[1] - for ti in 2:N - grad_inters_sum = map(+, grad_inters_sum, shared_grad_inters[ti]) - end - grad_inters[blockIdx().x] = grad_inters_sum - end - return nothing -end - -function ChainRulesCore.rrule(::typeof(pairwise_force_gpu), coords::AbstractArray{SVector{D, C}}, - atoms::AbstractArray{A}, boundary, pairwise_inters, nbs, force_units, - val_ft::Val{T}) where {D, C, A, T} - if force_units != NoUnits - error("taking gradients through force calculation is not compatible with units, " * - "system force units are $force_units") - end - Y = pairwise_force_gpu(coords, atoms, boundary, pairwise_inters, nbs, force_units, val_ft) - - function pairwise_force_gpu_pullback(d_fs_mat) - n_atoms = length(atoms) - z = zero(T) - fs_mat = CUDA.zeros(T, D, n_atoms) - d_coords = zero(coords) - d_atoms = CuArray(fill(zero(A), n_atoms)) - n_threads_gpu, n_blocks = cuda_threads_blocks_pairwise(length(nbs)) - grad_pairwise_inters = CuArray(fill(pairwise_inters, n_blocks)) - - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks grad_pairwise_force_kernel!(fs_mat, - d_fs_mat, coords, d_coords, atoms, d_atoms, boundary, pairwise_inters, - grad_pairwise_inters, nbs, Val(D), Val(force_units), Val(n_threads_gpu)) - - d_pairwise_inters = reduce((t1, t2) -> map(+, t1, t2), Array(grad_pairwise_inters)) - return NoTangent(), d_coords, d_atoms, NoTangent(), d_pairwise_inters, NoTangent(), - NoTangent(), NoTangent() - end - - return Y, pairwise_force_gpu_pullback -end - -function grad_pairwise_pe_kernel!(pe_vec, d_pe_vec, coords, d_coords, atoms, d_atoms, boundary, - inters::I, grad_inters, neighbors, val_energy_units, - ::Val{N}) where {I, N} - shared_grad_inters = CuStaticSharedArray(I, N) - sync_threads() - - grads = Enzyme.autodiff_deferred( - Enzyme.Reverse, - pairwise_pe_kernel!, - Const, - Duplicated(pe_vec, d_pe_vec), - Duplicated(coords, d_coords), - Duplicated(atoms, d_atoms), - Const(boundary), - Active(inters), - Const(neighbors), - Const(val_energy_units), - )[1] - - tidx = threadIdx().x - shared_grad_inters[tidx] = grads[5] - sync_threads() - - if tidx == 1 - grad_inters_sum = shared_grad_inters[1] - for ti in 2:N - grad_inters_sum = map(+, grad_inters_sum, shared_grad_inters[ti]) - end - grad_inters[blockIdx().x] = grad_inters_sum - end - return nothing -end - -function ChainRulesCore.rrule(::typeof(pairwise_pe_gpu), coords::AbstractArray{SVector{D, C}}, - atoms::AbstractArray{A}, boundary, pairwise_inters, nbs, energy_units, - val_ft::Val{T}) where {D, C, A, T} - if energy_units != NoUnits - error("taking gradients through potential energy calculation is not compatible with " * - "units, system energy units are $energy_units") - end - Y = pairwise_pe_gpu(coords, atoms, boundary, pairwise_inters, nbs, energy_units, val_ft) - - function pairwise_pe_gpu_pullback(d_pe_vec_arg) - n_atoms = length(atoms) - z = zero(T) - pe_vec = CUDA.zeros(T, 1) - d_pe_vec = CuArray([d_pe_vec_arg[1]]) - d_coords = zero(coords) - d_atoms = CuArray(fill(zero(A), n_atoms)) - n_threads_gpu, n_blocks = cuda_threads_blocks_pairwise(length(nbs)) - grad_pairwise_inters = CuArray(fill(pairwise_inters, n_blocks)) - - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks grad_pairwise_pe_kernel!(pe_vec, - d_pe_vec, coords, d_coords, atoms, d_atoms, boundary, pairwise_inters, - grad_pairwise_inters, nbs, Val(energy_units), Val(n_threads_gpu)) - - d_pairwise_inters = reduce((t1, t2) -> map(+, t1, t2), Array(grad_pairwise_inters)) - return NoTangent(), d_coords, d_atoms, NoTangent(), d_pairwise_inters, NoTangent(), - NoTangent(), NoTangent() - end - - return Y, pairwise_pe_gpu_pullback -end - -function grad_specific_force_1_atoms_kernel!(fs_mat, d_fs_mat, coords, d_coords, - boundary, is, inters, d_inters, - val_dims, val_force_units) - Enzyme.autodiff_deferred( - Enzyme.Reverse, - specific_force_1_atoms_kernel!, - Const, - Duplicated(fs_mat, d_fs_mat), - Duplicated(coords, d_coords), - Const(boundary), - Const(is), - Duplicated(inters, d_inters), - Const(val_dims), - Const(val_force_units), - ) - return nothing -end - -function grad_specific_force_2_atoms_kernel!(fs_mat, d_fs_mat, coords, d_coords, - boundary, is, js, inters, d_inters, - val_dims, val_force_units) - Enzyme.autodiff_deferred( - Enzyme.Reverse, - specific_force_2_atoms_kernel!, - Const, - Duplicated(fs_mat, d_fs_mat), - Duplicated(coords, d_coords), - Const(boundary), - Const(is), - Const(js), - Duplicated(inters, d_inters), - Const(val_dims), - Const(val_force_units), - ) - return nothing -end - -function grad_specific_force_3_atoms_kernel!(fs_mat, d_fs_mat, coords, d_coords, - boundary, is, js, ks, inters, d_inters, - val_dims, val_force_units) - Enzyme.autodiff_deferred( - Enzyme.Reverse, - specific_force_3_atoms_kernel!, - Const, - Duplicated(fs_mat, d_fs_mat), - Duplicated(coords, d_coords), - Const(boundary), - Const(is), - Const(js), - Const(ks), - Duplicated(inters, d_inters), - Const(val_dims), - Const(val_force_units), - ) - return nothing -end - -function grad_specific_force_4_atoms_kernel!(fs_mat, d_fs_mat, coords, d_coords, - boundary, is, js, ks, ls, inters, d_inters, - val_dims, val_force_units) - Enzyme.autodiff_deferred( - Enzyme.Reverse, - specific_force_4_atoms_kernel!, - Const, - Duplicated(fs_mat, d_fs_mat), - Duplicated(coords, d_coords), - Const(boundary), - Const(is), - Const(js), - Const(ks), - Const(ls), - Duplicated(inters, d_inters), - Const(val_dims), - Const(val_force_units), - ) - return nothing -end - -function ChainRulesCore.rrule(::typeof(specific_force_gpu), inter_list, - coords::AbstractArray{SVector{D, C}}, boundary, - force_units, val_ft::Val{T}) where {D, C, T} - if force_units != NoUnits - error("taking gradients through force calculation is not compatible with units, " * - "system force units are $force_units") - end - Y = specific_force_gpu(inter_list, coords, boundary, force_units, val_ft) - - function specific_force_gpu_pullback(d_fs_mat) - fs_mat = CUDA.zeros(T, D, length(coords)) - d_inter_list = zero(inter_list) - d_coords = zero(coords) - n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) - - if inter_list isa InteractionList1Atoms - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks grad_specific_force_1_atoms_kernel!( - fs_mat, d_fs_mat, coords, d_coords, boundary, - inter_list.is, - inter_list.inters, d_inter_list.inters, Val(D), Val(force_units)) - elseif inter_list isa InteractionList2Atoms - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks grad_specific_force_2_atoms_kernel!( - fs_mat, d_fs_mat, coords, d_coords, boundary, - inter_list.is, inter_list.js, - inter_list.inters, d_inter_list.inters, Val(D), Val(force_units)) - elseif inter_list isa InteractionList3Atoms - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks grad_specific_force_3_atoms_kernel!( - fs_mat, d_fs_mat, coords, d_coords, boundary, - inter_list.is, inter_list.js, inter_list.ks, - inter_list.inters, d_inter_list.inters, Val(D), Val(force_units)) - elseif inter_list isa InteractionList4Atoms - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks grad_specific_force_4_atoms_kernel!( - fs_mat, d_fs_mat, coords, d_coords, boundary, - inter_list.is, inter_list.js, inter_list.ks, inter_list.ls, - inter_list.inters, d_inter_list.inters, Val(D), Val(force_units)) - end - - return NoTangent(), d_inter_list, d_coords, NoTangent(), NoTangent(), NoTangent() - end - - return Y, specific_force_gpu_pullback -end - -function grad_specific_pe_1_atoms_kernel!(energy, d_energy, coords, d_coords, - boundary, is, inters, d_inters, - val_energy_units) - Enzyme.autodiff_deferred( - Enzyme.Reverse, - specific_pe_1_atoms_kernel!, - Const, - Duplicated(energy, d_energy), - Duplicated(coords, d_coords), - Const(boundary), - Const(is), - Duplicated(inters, d_inters), - Const(val_energy_units), - ) - return nothing -end - -function grad_specific_pe_2_atoms_kernel!(energy, d_energy, coords, d_coords, - boundary, is, js, inters, d_inters, - val_energy_units) - Enzyme.autodiff_deferred( - Enzyme.Reverse, - specific_pe_2_atoms_kernel!, - Const, - Duplicated(energy, d_energy), - Duplicated(coords, d_coords), - Const(boundary), - Const(is), - Const(js), - Duplicated(inters, d_inters), - Const(val_energy_units), - ) - return nothing -end - -function grad_specific_pe_3_atoms_kernel!(energy, d_energy, coords, d_coords, - boundary, is, js, ks, inters, d_inters, - val_energy_units) - Enzyme.autodiff_deferred( - Enzyme.Reverse, - specific_pe_3_atoms_kernel!, - Const, - Duplicated(energy, d_energy), - Duplicated(coords, d_coords), - Const(boundary), - Const(is), - Const(js), - Const(ks), - Duplicated(inters, d_inters), - Const(val_energy_units), - ) - return nothing -end - -function grad_specific_pe_4_atoms_kernel!(energy, d_energy, coords, d_coords, - boundary, is, js, ks, ls, inters, d_inters, - val_energy_units) - Enzyme.autodiff_deferred( - Enzyme.Reverse, - specific_pe_4_atoms_kernel!, - Const, - Duplicated(energy, d_energy), - Duplicated(coords, d_coords), - Const(boundary), - Const(is), - Const(js), - Const(ks), - Const(ls), - Duplicated(inters, d_inters), - Const(val_energy_units), - ) - return nothing -end - -function ChainRulesCore.rrule(::typeof(specific_pe_gpu), inter_list, - coords::AbstractArray{SVector{D, C}}, boundary, - energy_units, val_ft::Val{T}) where {D, C, T} - if energy_units != NoUnits - error("taking gradients through potential energy calculation is not compatible with " * - "units, system energy units are $energy_units") - end - Y = specific_pe_gpu(inter_list, coords, boundary, energy_units, val_ft) - - function specific_pe_gpu_pullback(d_pe_vec_arg) - pe_vec = CUDA.zeros(T, 1) - d_pe_vec = CuArray([d_pe_vec_arg[1]]) - d_inter_list = zero(inter_list) - d_coords = zero(coords) - n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) - - if inter_list isa InteractionList1Atoms - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks grad_specific_pe_1_atoms_kernel!( - pe_vec, d_pe_vec, coords, d_coords, boundary, - inter_list.is, - inter_list.inters, d_inter_list.inters, Val(energy_units)) - elseif inter_list isa InteractionList2Atoms - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks grad_specific_pe_2_atoms_kernel!( - pe_vec, d_pe_vec, coords, d_coords, boundary, - inter_list.is, inter_list.js, - inter_list.inters, d_inter_list.inters, Val(energy_units)) - elseif inter_list isa InteractionList3Atoms - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks grad_specific_pe_3_atoms_kernel!( - pe_vec, d_pe_vec, coords, d_coords, boundary, - inter_list.is, inter_list.js, inter_list.ks, - inter_list.inters, d_inter_list.inters, Val(energy_units)) - elseif inter_list isa InteractionList4Atoms - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks grad_specific_pe_4_atoms_kernel!( - pe_vec, d_pe_vec, coords, d_coords, boundary, - inter_list.is, inter_list.js, inter_list.ks, inter_list.ls, - inter_list.inters, d_inter_list.inters, Val(energy_units)) - end - - return NoTangent(), d_inter_list, d_coords, NoTangent(), NoTangent(), NoTangent() - end - - return Y, specific_pe_gpu_pullback -end - -function grad_gbsa_born_kernel!(Is, d_Is, I_grads, d_I_grads, coords, d_coords, offset_radii, - d_offset_radii, scaled_offset_radii, d_scaled_offset_radii, - dist_cutoff, offset, neck_scale::T, grad_neck_scale, neck_cut, d0s, - d_d0s, m0s, d_m0s, boundary, val_coord_units, ::Val{N}) where {T, N} - shared_grad_neck_scale = CuStaticSharedArray(T, N) - sync_threads() - - grads = Enzyme.autodiff_deferred( - Enzyme.Reverse, - gbsa_born_kernel!, - Const, - Duplicated(Is, d_Is), - Duplicated(I_grads, d_I_grads), - Duplicated(coords, d_coords), - Duplicated(offset_radii, d_offset_radii), - Duplicated(scaled_offset_radii, d_scaled_offset_radii), - Const(dist_cutoff), - Const(offset), - Active(neck_scale), - Const(neck_cut), - Duplicated(d0s, d_d0s), - Duplicated(m0s, d_m0s), - Const(boundary), - Const(val_coord_units), - )[1] - - tidx = threadIdx().x - shared_grad_neck_scale[tidx] = grads[8] - sync_threads() - - if tidx == 1 - grad_neck_scale_sum = shared_grad_neck_scale[1] - for ti in 2:N - grad_neck_scale_sum += shared_grad_neck_scale[ti] - end - if !iszero(grad_neck_scale_sum) - Atomix.@atomic grad_neck_scale[1] += grad_neck_scale_sum - end - end - return nothing -end - -function ChainRulesCore.rrule(::typeof(gbsa_born_gpu), coords::AbstractArray{SVector{D, C}}, - offset_radii, scaled_offset_radii, dist_cutoff, offset, neck_scale, - neck_cut, d0s, m0s, boundary, val_ft::Val{T}) where {D, C, T} - if unit(C) != NoUnits - error("taking gradients through force/energy calculation is not compatible with units, " * - "coordinate units are $(unit(C))") - end - Y = gbsa_born_gpu(coords, offset_radii, scaled_offset_radii, dist_cutoff, offset, neck_scale, - neck_cut, d0s, m0s, boundary, val_ft) - - function gbsa_born_gpu_pullback(d_args) - n_atoms = length(coords) - d_Is = d_args[1] == ZeroTangent() ? CUDA.zeros(T, n_atoms) : d_args[1] - d_I_grads = d_args[2] == ZeroTangent() ? CUDA.zeros(T, n_atoms, n_atoms) : d_args[2] - Is = CUDA.zeros(T, n_atoms) - I_grads = CUDA.zeros(T, n_atoms, n_atoms) - d_coords = zero(coords) - d_offset_radii = zero(offset_radii) - d_scaled_offset_radii = zero(scaled_offset_radii) - grad_neck_scale = CUDA.zeros(T, 1) - d_d0s = zero(d0s) - d_m0s = zero(m0s) - n_inters = n_atoms ^ 2 - n_threads_gpu, n_blocks = cuda_threads_blocks_gbsa(n_inters) - - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks grad_gbsa_born_kernel!( - Is, d_Is, I_grads, d_I_grads, coords, d_coords, offset_radii, - d_offset_radii, scaled_offset_radii, d_scaled_offset_radii, - dist_cutoff, offset, neck_scale, grad_neck_scale, neck_cut, d0s, - d_d0s, m0s, d_m0s, boundary, Val(C), Val(n_threads_gpu)) - - d_neck_scale = Array(grad_neck_scale)[1] - return NoTangent(), d_coords, d_offset_radii, d_scaled_offset_radii, NoTangent(), - NoTangent(), d_neck_scale, NoTangent(), d_d0s, d_m0s, NoTangent(), NoTangent() - end - - return Y, gbsa_born_gpu_pullback -end - -function grad_gbsa_force_1_kernel!(fs_mat, d_fs_mat, born_forces_mod_ustrip, - d_born_forces_mod_ustrip, coords, d_coords, boundary, - dist_cutoff, factor_solute::T, grad_factor_solute, - factor_solvent::T, grad_factor_solvent, kappa::T, grad_kappa, - Bs, d_Bs, chs, d_chs, val_dims, val_force_units, - ::Val{N}) where {T, N} - shared_grad_factor_solute = CuStaticSharedArray(T, N) - shared_grad_factor_solvent = CuStaticSharedArray(T, N) - shared_grad_kappa = CuStaticSharedArray(T, N) - sync_threads() - - grads = Enzyme.autodiff_deferred( - Enzyme.Reverse, - gbsa_force_1_kernel!, - Const, - Duplicated(fs_mat, d_fs_mat), - Duplicated(born_forces_mod_ustrip, d_born_forces_mod_ustrip), - Duplicated(coords, d_coords), - Const(boundary), - Const(dist_cutoff), - Active(factor_solute), - Active(factor_solvent), - Active(kappa), - Duplicated(Bs, d_Bs), - Duplicated(chs, d_chs), - Const(val_dims), - Const(val_force_units), - )[1] - - tidx = threadIdx().x - shared_grad_factor_solute[tidx] = grads[6] - shared_grad_factor_solvent[tidx] = grads[7] - shared_grad_kappa[tidx] = grads[8] - sync_threads() - - if tidx == 1 - grad_factor_solute_sum = shared_grad_factor_solute[1] - for ti in 2:N - grad_factor_solute_sum += shared_grad_factor_solute[ti] - end - if !iszero(grad_factor_solute_sum) - Atomix.@atomic grad_factor_solute[1] += grad_factor_solute_sum - end - elseif tidx == 2 - grad_factor_solvent_sum = shared_grad_factor_solvent[1] - for ti in 2:N - grad_factor_solvent_sum += shared_grad_factor_solvent[ti] - end - if !iszero(grad_factor_solvent_sum) - Atomix.@atomic grad_factor_solvent[1] += grad_factor_solvent_sum - end - elseif tidx == 3 - grad_kappa_sum = shared_grad_kappa[1] - for ti in 2:N - grad_kappa_sum += shared_grad_kappa[ti] - end - if !iszero(grad_kappa_sum) - Atomix.@atomic grad_kappa[1] += grad_kappa_sum - end - end - return nothing -end - -function ChainRulesCore.rrule(::typeof(gbsa_force_1_gpu), coords::AbstractArray{SVector{D, C}}, - boundary, dist_cutoff, factor_solute, factor_solvent, kappa, Bs, - chs::AbstractArray{T}, force_units) where {D, C, T} - if force_units != NoUnits - error("taking gradients through force calculation is not compatible with units, " * - "system force units are $force_units") - end - Y = gbsa_force_1_gpu(coords, boundary, dist_cutoff, factor_solute, factor_solvent, kappa, - Bs, chs, force_units) - - function gbsa_force_1_gpu_pullback(d_args) - n_atoms = length(coords) - d_fs_mat = d_args[1] == ZeroTangent() ? CUDA.zeros(T, D, n_atoms) : d_args[1] - d_born_forces_mod_ustrip = d_args[2] == ZeroTangent() ? CUDA.zeros(T, n_atoms) : d_args[2] - fs_mat = CUDA.zeros(T, D, n_atoms) - born_forces_mod_ustrip = CUDA.zeros(T, n_atoms) - d_coords = zero(coords) - grad_factor_solute = CUDA.zeros(T, 1) - grad_factor_solvent = CUDA.zeros(T, 1) - grad_kappa = CUDA.zeros(T, 1) - d_Bs = zero(Bs) - d_chs = zero(chs) - n_inters = n_atoms_to_n_pairs(n_atoms) + n_atoms - n_threads_gpu, n_blocks = cuda_threads_blocks_gbsa(n_inters) - - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks grad_gbsa_force_1_kernel!( - fs_mat, d_fs_mat, born_forces_mod_ustrip, d_born_forces_mod_ustrip, coords, - d_coords, boundary, dist_cutoff, factor_solute, grad_factor_solute, - factor_solvent, grad_factor_solvent, kappa, grad_kappa, Bs, d_Bs, chs, - d_chs, Val(D), Val(force_units), Val(n_threads_gpu)) - - d_factor_solute = Array(grad_factor_solute )[1] - d_factor_solvent = Array(grad_factor_solvent)[1] - d_kappa = Array(grad_kappa )[1] - return NoTangent(), d_coords, NoTangent(), NoTangent(), d_factor_solute, - d_factor_solvent, d_kappa, d_Bs, d_chs, NoTangent() - end - - return Y, gbsa_force_1_gpu_pullback -end - -function grad_gbsa_force_2_kernel!(fs_mat, d_fs_mat, born_forces, d_born_forces, coords, d_coords, - boundary, dist_cutoff, or, d_or, sor, d_sor, Bs, d_Bs, B_grads, - d_B_grads, I_grads, d_I_grads, val_dims, val_force_units) - Enzyme.autodiff_deferred( - Enzyme.Reverse, - gbsa_force_2_kernel!, - Const, - Duplicated(fs_mat, d_fs_mat), - Duplicated(born_forces, d_born_forces), - Duplicated(coords, d_coords), - Const(boundary), - Const(dist_cutoff), - Duplicated(or, d_or), - Duplicated(sor, d_sor), - Duplicated(Bs, d_Bs), - Duplicated(B_grads, d_B_grads), - Duplicated(I_grads, d_I_grads), - Const(val_dims), - Const(val_force_units), - ) - return nothing -end - -function ChainRulesCore.rrule(::typeof(gbsa_force_2_gpu), coords::AbstractArray{SVector{D, C}}, - boundary, dist_cutoff, Bs, B_grads, I_grads, born_forces, offset_radii, - scaled_offset_radii, force_units, val_ft::Val{T}) where {D, C, T} - if force_units != NoUnits - error("taking gradients through force calculation is not compatible with units, " * - "system force units are $force_units") - end - Y = gbsa_force_2_gpu(coords, boundary, dist_cutoff, Bs, B_grads, I_grads, born_forces, - offset_radii, scaled_offset_radii, force_units, val_ft) - - function gbsa_force_2_gpu_pullback(d_fs_mat) - n_atoms = length(coords) - fs_mat = CUDA.zeros(T, D, n_atoms) - d_coords = zero(coords) - d_born_forces = zero(born_forces) - d_offset_radii = zero(offset_radii) - d_scaled_offset_radii = zero(scaled_offset_radii) - d_Bs = zero(Bs) - d_B_grads = zero(B_grads) - d_I_grads = zero(I_grads) - n_inters = n_atoms ^ 2 - n_threads_gpu, n_blocks = cuda_threads_blocks_gbsa(n_inters) - - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks grad_gbsa_force_2_kernel!( - fs_mat, d_fs_mat, born_forces, d_born_forces, coords, d_coords, boundary, - dist_cutoff, offset_radii, d_offset_radii, scaled_offset_radii, - d_scaled_offset_radii, Bs, d_Bs, B_grads, d_B_grads, I_grads, d_I_grads, - Val(D), Val(force_units)) - - return NoTangent(), d_coords, NoTangent(), NoTangent(), d_Bs, d_B_grads, d_I_grads, - d_born_forces, d_offset_radii, d_scaled_offset_radii, NoTangent(), NoTangent() - end - - return Y, gbsa_force_2_gpu_pullback -end From f4d6c80ab0b5d0ff5fd590a5fb3c6d09a80fd339 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 28 Jun 2024 15:35:12 +0100 Subject: [PATCH 08/74] interaction arguments --- benchmark/benchmarks.jl | 8 ++--- docs/src/differentiable.md | 2 -- docs/src/examples.md | 51 ++++++++-------------------- src/interactions/buckingham.jl | 8 ++--- src/interactions/coulomb.jl | 16 ++++----- src/interactions/lennard_jones.jl | 14 ++++---- src/interactions/mie.jl | 8 ++--- src/interactions/soft_sphere.jl | 8 ++--- src/setup.jl | 18 +++------- test/agent.jl | 13 ++------ test/basic.jl | 4 +-- test/interactions.jl | 55 +++++++++++++++---------------- test/minimization.jl | 2 +- test/simulation.jl | 16 ++------- test/zygote.jl | 15 +++------ 15 files changed, 89 insertions(+), 149 deletions(-) diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index a144cf54e..a521ccb16 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -44,10 +44,10 @@ coords = [c1, c2] dr = vector(c1, c2, boundary) b1 = HarmonicBond(k=100_000.0u"kJ * mol^-1 * nm^-2", r0=0.6u"nm") -SUITE["interactions"]["LennardJones force" ] = @benchmarkable force($(LennardJones()), $(dr), $(c1), $(c2), $(a1), $(a1), $(boundary)) -SUITE["interactions"]["LennardJones energy"] = @benchmarkable potential_energy($(LennardJones()), $(dr), $(c1), $(c2), $(a1), $(a1), $(boundary)) -SUITE["interactions"]["Coulomb force" ] = @benchmarkable force($(Coulomb()), $(dr), $(c1), $(c2), $(a1), $(a1), $(boundary)) -SUITE["interactions"]["Coulomb energy" ] = @benchmarkable potential_energy($(Coulomb()), $(dr), $(c1), $(c2), $(a1), $(a1), $(boundary)) +SUITE["interactions"]["LennardJones force" ] = @benchmarkable force($(LennardJones()), $(dr), $(a1), $(a1)) +SUITE["interactions"]["LennardJones energy"] = @benchmarkable potential_energy($(LennardJones()), $(dr), $(a1), $(a1)) +SUITE["interactions"]["Coulomb force" ] = @benchmarkable force($(Coulomb()), $(dr), $(a1), $(a1)) +SUITE["interactions"]["Coulomb energy" ] = @benchmarkable potential_energy($(Coulomb()), $(dr), $(a1), $(a1)) SUITE["interactions"]["HarmonicBond force" ] = @benchmarkable force($(b1), $(c1), $(c2), $(boundary)) SUITE["interactions"]["HarmonicBond energy"] = @benchmarkable potential_energy($(b1), $(c1), $(c2), $(boundary)) diff --git a/docs/src/differentiable.md b/docs/src/differentiable.md index 777a18dad..79646ccb8 100644 --- a/docs/src/differentiable.md +++ b/docs/src/differentiable.md @@ -66,8 +66,6 @@ neighbor_finder = DistanceNeighborFinder( lj = LennardJones( cutoff=DistanceCutoff(1.5), use_neighbors=true, - force_units=NoUnits, - energy_units=NoUnits, ) pairwise_inters = (lj,) coords = place_atoms(n_atoms, boundary; min_dist=0.6) diff --git a/docs/src/examples.md b/docs/src/examples.md index 53a49afc7..7bab7cee4 100644 --- a/docs/src/examples.md +++ b/docs/src/examples.md @@ -166,11 +166,9 @@ end # Custom force function function Molly.force(inter::SIRInteraction, vec_ij, - coord_i, - coord_j, atom_i, atom_j, - boundary) + args...) if (atom_i.status == infected && atom_j.status == susceptible) || (atom_i.status == susceptible && atom_j.status == infected) # Infect close people randomly @@ -187,7 +185,7 @@ function Molly.force(inter::SIRInteraction, atom_i.status = recovered end end - return zero(coord_i) + return zero(vec_ij) end # Custom logger @@ -211,12 +209,7 @@ atoms = [Person(i, i <= n_starting ? infected : susceptible, 1.0, 0.1, 0.02) for coords = place_atoms(n_people, boundary; min_dist=0.1) velocities = [random_velocity(1.0, temp; dims=2) for i in 1:n_people] -lj = LennardJones( - cutoff=DistanceCutoff(1.6), - use_neighbors=true, - force_units=NoUnits, - energy_units=NoUnits, -) +lj = LennardJones(cutoff=DistanceCutoff(1.6), use_neighbors=true) sir = SIRInteraction(0.5, 0.06, 0.01) pairwise_inters = (LennardJones=lj, SIR=sir) neighbor_finder = DistanceNeighborFinder( @@ -640,11 +633,9 @@ Molly.use_neighbors(::BondableInteraction) = true function Molly.force(inter::BondableInteraction, dr, - coord_i, - coord_j, atom_i, atom_j, - boundary) + args...) # Break bonds randomly if atom_j.i in atom_i.partners && rand() < inter.prob_break delete!(atom_i.partners, atom_j.i) @@ -662,7 +653,7 @@ function Molly.force(inter::BondableInteraction, fdr = -c * normalize(dr) return fdr else - return zero(coord_i) + return zero(dr) end end @@ -687,12 +678,7 @@ atoms = [BondableAtom(i, 1.0, 0.1, 0.02, Set([])) for i in 1:n_atoms] coords = place_atoms(n_atoms, boundary; min_dist=0.1) velocities = [random_velocity(1.0, temp; dims=2) for i in 1:n_atoms] pairwise_inters = ( - SoftSphere( - cutoff=DistanceCutoff(2.0), - use_neighbors=true, - force_units=NoUnits, - energy_units=NoUnits, - ), + SoftSphere(cutoff=DistanceCutoff(2.0), use_neighbors=true), BondableInteraction(0.1, 0.1, 1.1, 2.0, 0.1), ) neighbor_finder = DistanceNeighborFinder( @@ -751,7 +737,7 @@ using Molly using Zygote using GLMakie -inter = LennardJones(force_units=NoUnits, energy_units=NoUnits) +inter = LennardJones() boundary = CubicBoundary(5.0) a1, a2 = Atom(σ=0.3, ϵ=0.5), Atom(σ=0.3, ϵ=0.5) @@ -759,7 +745,7 @@ function force_direct(dist) c1 = SVector(1.0, 1.0, 1.0) c2 = SVector(dist + 1.0, 1.0, 1.0) vec = vector(c1, c2, boundary) - F = force(inter, vec, c1, c2, a1, a2, boundary) + F = force(inter, vec, a1, a2, NoUnits) return F[1] end @@ -768,7 +754,7 @@ function force_grad(dist) c1 = SVector(1.0, 1.0, 1.0) c2 = SVector(dist + 1.0, 1.0, 1.0) vec = vector(c1, c2, boundary) - potential_energy(inter, vec, c1, c2, a1, a2, boundary) + potential_energy(inter, vec, a1, a2, NoUnits) end return -grad[1] end @@ -810,7 +796,7 @@ ab_sys = AtomsBase.AbstractSystem( [0.0 , 0.0 , 1.7928950]]u"Å", ) -coul = Coulomb(coulomb_const=2.307e-21u"kJ*Å", force_units=u"kJ/Å", energy_units=u"kJ") +coul = Coulomb(coulomb_const=2.307e-21u"kJ*Å") calc = MollyCalculator(pairwise_inters=(coul,), force_units=u"kJ/Å", energy_units=u"kJ") AtomsCalculators.potential_energy(ab_sys, calc) @@ -901,7 +887,7 @@ function energies(m, n) c1 = SVector(1.0, 1.0, 1.0) c2 = SVector(dist + 1.0, 1.0, 1.0) vec = vector(c1, c2, boundary) - potential_energy(inter, vec, c1, c2, a1, a2, boundary) + potential_energy(inter, vec, a1, a2, NoUnits) end end @@ -948,7 +934,7 @@ function energies(α, λ, p) c1 = SVector(1.0, 1.0, 1.0) c2 = SVector(dist + 1.0, 1.0, 1.0) vec = vector(c1, c2, boundary) - potential_energy(inter, vec, c1, c2, a1, a2, boundary) + potential_energy(inter, vec, a1, a2, NoUnits) end end @@ -1005,11 +991,7 @@ r_cut = 0.85u"nm" sys = System( fcc_crystal; velocities=velocities, - pairwise_inters=(LennardJones( - cutoff=ShiftedForceCutoff(r_cut), - energy_units=u"kJ * mol^-1", - force_units=u"kJ * mol^-1 * nm^-1", - ),), + pairwise_inters=(LennardJones(cutoff=ShiftedForceCutoff(r_cut)),), loggers=( kinetic_eng=KineticEnergyLogger(100), pot_eng=PotentialEnergyLogger(100), @@ -1060,12 +1042,7 @@ atoms = [Atom(index=i, mass=atom_mass, σ=2.8279u"Å", ϵ=0.074u"kcal* mol^-1") max_coord = 200.0u"Å" coords = [max_coord .* rand(SVector{3}) for i in 1:n_atoms_half] boundary = CubicBoundary(200.0u"Å") -lj = LennardJones( - cutoff=ShiftedPotentialCutoff(r_cut), - use_neighbors=true, - energy_units=u"kcal * mol^-1", - force_units=u"kcal * mol^-1 * Å^-1", -) +lj = LennardJones(cutoff=ShiftedPotentialCutoff(r_cut), use_neighbors=true) # Add bonded atoms bond_length = 0.74u"Å" # Hydrogen bond length diff --git a/src/interactions/buckingham.jl b/src/interactions/buckingham.jl index 3684669b7..d63a8a8f3 100644 --- a/src/interactions/buckingham.jl +++ b/src/interactions/buckingham.jl @@ -43,12 +43,12 @@ use_neighbors(inter::Buckingham) = inter.use_neighbors dr, atom_i, atom_j, - force_units, + force_units=u"kJ * mol^-1 * nm^-1", special::Bool=false, args...) where C if (iszero_value(atom_i.A) || iszero_value(atom_j.A)) && (iszero_value(atom_i.C) || iszero_value(atom_j.C)) - return ustrip.(zero(coord_i)) * force_units + return ustrip.(zero(dr)) * force_units end Aij = sqrt(atom_i.A * atom_j.A) @@ -76,12 +76,12 @@ end dr, atom_i, atom_j, - energy_units, + energy_units=u"kJ * mol^-1", special::Bool=false, args...) where C if (iszero_value(atom_i.A) || iszero_value(atom_j.A)) && (iszero_value(atom_i.C) || iszero_value(atom_j.C)) - return ustrip(zero(coord_i[1])) * energy_units + return ustrip(zero(dr[1])) * energy_units end Aij = sqrt(atom_i.A * atom_j.A) diff --git a/src/interactions/coulomb.jl b/src/interactions/coulomb.jl index f2ab87bbb..2811a6d99 100644 --- a/src/interactions/coulomb.jl +++ b/src/interactions/coulomb.jl @@ -50,7 +50,7 @@ end dr, atom_i, atom_j, - force_units, + force_units=u"kJ * mol^-1 * nm^-1", special::Bool=false, args...) where C r2 = sum(abs2, dr) @@ -75,7 +75,7 @@ end dr, atom_i, atom_j, - energy_units, + energy_units=u"kJ * mol^-1", special::Bool=false, args...) where C r2 = sum(abs2, dr) @@ -168,7 +168,7 @@ end dr, atom_i, atom_j, - force_units, + force_units=u"kJ * mol^-1 * nm^-1", special::Bool=false, args...) where C r2 = sum(abs2, dr) @@ -198,7 +198,7 @@ end dr, atom_i, atom_j, - energy_units, + energy_units=u"kJ * mol^-1", special::Bool=false, args...) where C r2 = sum(abs2, dr) @@ -275,12 +275,12 @@ end dr, atom_i, atom_j, - force_units, + force_units=u"kJ * mol^-1 * nm^-1", special::Bool=false, args...) r2 = sum(abs2, dr) if r2 > (inter.dist_cutoff ^ 2) - return ustrip.(zero(coord_i)) * force_units + return ustrip.(zero(dr)) * force_units end coulomb_const = inter.coulomb_const @@ -308,12 +308,12 @@ end dr, atom_i, atom_j, - energy_units, + energy_units=u"kJ * mol^-1", special::Bool=false, args...) r2 = sum(abs2, dr) if r2 > (inter.dist_cutoff ^ 2) - return ustrip(zero(coord_i[1])) * energy_units + return ustrip(zero(dr[1])) * energy_units end coulomb_const = inter.coulomb_const diff --git a/src/interactions/lennard_jones.jl b/src/interactions/lennard_jones.jl index b79de4c1f..664009ccf 100644 --- a/src/interactions/lennard_jones.jl +++ b/src/interactions/lennard_jones.jl @@ -70,7 +70,7 @@ end dr, atom_i, atom_j, - force_units, + force_units=u"kJ * mol^-1 * nm^-1", special::Bool=false, args...) where {S, C} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || @@ -109,12 +109,12 @@ end dr, atom_i, atom_j, - energy_units, + energy_units=u"kJ * mol^-1", special::Bool=false, args...) where {S, C} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) - return ustrip(zero(coord_i[1])) * energy_units + return ustrip(zero(dr[1])) * energy_units end σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) @@ -227,12 +227,12 @@ end dr, atom_i, atom_j, - force_units, + force_units=u"kJ * mol^-1 * nm^-1", special::Bool=false, args...) where {S, C} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) - return ustrip.(zero(coord_i)) * force_units + return ustrip.(zero(dr)) * force_units end # Lorentz-Berthelot mixing rules use the arithmetic average for σ @@ -269,12 +269,12 @@ end dr, atom_i, atom_j, - energy_units, + energy_units=u"kJ * mol^-1", special::Bool=false, args...) where {S, C} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || iszero_value(atom_i.σ) ||iszero_value(atom_j.σ)) - return ustrip(zero(coord_i[1])) * energy_units + return ustrip(zero(dr[1])) * energy_units end σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) diff --git a/src/interactions/mie.jl b/src/interactions/mie.jl index 883669412..74fd9c5a3 100644 --- a/src/interactions/mie.jl +++ b/src/interactions/mie.jl @@ -42,11 +42,11 @@ function force(inter::Mie{S, C, T}, dr, atom_i, atom_j, - force_units, + force_units=u"kJ * mol^-1 * nm^-1", args...) where {S, C, T} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) - return ustrip.(zero(coord_i)) * force_units + return ustrip.(zero(dr)) * force_units end # Lorentz-Berthelot mixing rules use the arithmetic average for σ @@ -75,11 +75,11 @@ end dr, atom_i, atom_j, - energy_units, + energy_units=u"kJ * mol^-1", args...) where {S, C, T} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) - return ustrip(zero(coord_i[1])) * energy_units + return ustrip(zero(dr[1])) * energy_units end σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) diff --git a/src/interactions/soft_sphere.jl b/src/interactions/soft_sphere.jl index cca4fa8f3..cb71fe127 100644 --- a/src/interactions/soft_sphere.jl +++ b/src/interactions/soft_sphere.jl @@ -31,11 +31,11 @@ use_neighbors(inter::SoftSphere) = inter.use_neighbors dr, atom_i, atom_j, - force_units, + force_units=u"kJ * mol^-1 * nm^-1", args...) where {S, C} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) - return ustrip.(zero(coord_i)) * force_units + return ustrip.(zero(dr)) * force_units end # Lorentz-Berthelot mixing rules use the arithmetic average for σ @@ -61,11 +61,11 @@ function potential_energy(inter::SoftSphere{S, C}, dr, atom_i, atom_j, - energy_units, + energy_units=u"kJ * mol^-1", args...) where {S, C} if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) - return ustrip(zero(coord_i[1])) * energy_units + return ustrip(zero(dr[1])) * energy_units end σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) diff --git a/src/setup.jl b/src/setup.jl index e1ffb0a0d..c1dd866cb 100644 --- a/src/setup.jl +++ b/src/setup.jl @@ -792,8 +792,6 @@ function System(coord_file::AbstractString, cutoff=DistanceCutoff(T(dist_cutoff)), use_neighbors=true, weight_special=force_field.weight_14_lj, - force_units=force_units, - energy_units=energy_units, ) if isnothing(implicit_solvent) crf = CoulombReactionField( @@ -802,8 +800,6 @@ function System(coord_file::AbstractString, use_neighbors=true, weight_special=force_field.weight_14_coulomb, coulomb_const=(units ? T(coulombconst) : T(ustrip(coulombconst))), - force_units=force_units, - energy_units=energy_units, ) else crf = Coulomb( @@ -811,8 +807,6 @@ function System(coord_file::AbstractString, use_neighbors=true, weight_special=force_field.weight_14_coulomb, coulomb_const=(units ? T(coulombconst) : T(ustrip(coulombconst))), - force_units=force_units, - energy_units=energy_units, ) end pairwise_inters = (lj, crf) @@ -954,8 +948,8 @@ function System(coord_file::AbstractString, general_inters=general_inters, neighbor_finder=neighbor_finder, loggers=loggers, - force_units=units ? u"kJ * mol^-1 * nm^-1" : NoUnits, - energy_units=units ? u"kJ * mol^-1" : NoUnits, + force_units=(units ? u"kJ * mol^-1 * nm^-1" : NoUnits), + energy_units=(units ? u"kJ * mol^-1" : NoUnits), k=k, data=data, ) @@ -1214,8 +1208,6 @@ function System(T::Type, cutoff=DistanceCutoff(T(dist_cutoff)), use_neighbors=true, weight_special=T(0.5), - force_units=force_units, - energy_units=energy_units, ) crf = CoulombReactionField( dist_cutoff=T(dist_cutoff), @@ -1223,8 +1215,6 @@ function System(T::Type, use_neighbors=true, weight_special=T(0.5), coulomb_const=(units ? T(coulombconst) : T(ustrip(coulombconst))), - force_units=force_units, - energy_units=energy_units, ) if isnothing(boundary) @@ -1322,8 +1312,8 @@ function System(T::Type, specific_inter_lists=specific_inter_lists, neighbor_finder=neighbor_finder, loggers=loggers, - force_units=units ? u"kJ * mol^-1 * nm^-1" : NoUnits, - energy_units=units ? u"kJ * mol^-1" : NoUnits, + force_units=(units ? u"kJ * mol^-1 * nm^-1" : NoUnits), + energy_units=(units ? u"kJ * mol^-1" : NoUnits), k=k, data=data, ) diff --git a/test/agent.jl b/test/agent.jl index 850d6e687..75cf6ae97 100644 --- a/test/agent.jl +++ b/test/agent.jl @@ -28,11 +28,9 @@ # Custom force function function Molly.force(inter::SIRInteraction, vec_ij, - coord_i, - coord_j, atom_i, atom_j, - boundary) + args...) if (atom_i.status == infected && atom_j.status == susceptible) || (atom_i.status == susceptible && atom_j.status == infected) # Infect close people randomly @@ -49,7 +47,7 @@ atom_i.status = recovered end end - return zero(coord_i) + return zero(vec_ij) end # Test log_property! definition rather than just using GeneralObservableLogger @@ -74,12 +72,7 @@ coords = place_atoms(n_people, boundary; min_dist=0.1) velocities = [random_velocity(1.0, temp; dims=2) for i in 1:n_people] - lj = LennardJones( - cutoff=DistanceCutoff(1.6), - use_neighbors=true, - force_units=NoUnits, - energy_units=NoUnits, - ) + lj = LennardJones(cutoff=DistanceCutoff(1.6), use_neighbors=true) sir = SIRInteraction(0.5, 0.06, 0.01) @test !use_neighbors(sir) pairwise_inters = (LennardJones=lj, SIR=sir) diff --git a/test/basic.jl b/test/basic.jl index eeb038ee6..d92da4bd6 100644 --- a/test/basic.jl +++ b/test/basic.jl @@ -455,7 +455,7 @@ end # Mis-matched energy units in interaction and system coords = place_atoms(1, b_right; min_dist=0.01u"nm") - lj = LennardJones(energy_units="kcal/mol") + lj = LennardJones() @test_throws ArgumentError System(atoms=atoms, coords=coords, boundary=b_right, pairwise_inters=(lj,)) @@ -497,7 +497,7 @@ end [0.0 , 1.4654985, 0.0 ], [0.0 , 0.0 , 1.7928950]]u"Å", ) - coul = Coulomb(coulomb_const=2.307e-21u"kJ*Å", force_units=u"kJ/Å", energy_units=u"kJ") + coul = Coulomb(coulomb_const=2.307e-21u"kJ*Å") calc = MollyCalculator(pairwise_inters=(coul,), force_units=u"kJ/Å", energy_units=u"kJ") pe = AtomsCalculators.potential_energy(ab_sys, calc) diff --git a/test/interactions.jl b/test/interactions.jl index e8f248974..f6706604f 100644 --- a/test/interactions.jl +++ b/test/interactions.jl @@ -12,22 +12,22 @@ for inter in (LennardJones(), Mie(m=6, n=12), LennardJonesSoftCore(α=1, λ=0, p=2)) @test isapprox( - force(inter, dr12, c1, c2, a1, a1, boundary), + force(inter, dr12, a1, a1), SVector(16.0, 0.0, 0.0)u"kJ * mol^-1 * nm^-1"; atol=1e-9u"kJ * mol^-1 * nm^-1", ) @test isapprox( - force(inter, dr13, c1, c3, a1, a1, boundary), + force(inter, dr13, a1, a1), SVector(-1.375509739, 0.0, 0.0)u"kJ * mol^-1 * nm^-1"; atol=1e-9u"kJ * mol^-1 * nm^-1", ) @test isapprox( - potential_energy(inter, dr12, c1, c2, a1, a1, boundary), + potential_energy(inter, dr12, a1, a1), 0.0u"kJ * mol^-1"; atol=1e-9u"kJ * mol^-1", ) @test isapprox( - potential_energy(inter, dr13, c1, c3, a1, a1, boundary), + potential_energy(inter, dr13, a1, a1), -0.1170417309u"kJ * mol^-1"; atol=1e-9u"kJ * mol^-1", ) @@ -35,44 +35,44 @@ inter = LennardJonesSoftCore(α=1, λ=0.5, p=2) @test isapprox( - force(inter, dr12, c1, c2, a1, a1, boundary), + force(inter, dr12, a1, a1), SVector(6.144, 0.0, 0.0)u"kJ * mol^-1 * nm^-1"; atol=1e-9u"kJ * mol^-1 * nm^-1", ) @test isapprox( - force(inter, dr13, c1, c3, a1, a1, boundary), + force(inter, dr13, a1, a1), SVector(-1.290499537, 0.0, 0.0)u"kJ * mol^-1 * nm^-1"; atol=1e-9u"kJ * mol^-1 * nm^-1", ) @test isapprox( - potential_energy(inter, dr12, c1, c2, a1, a1, boundary), + potential_energy(inter, dr12, a1, a1), -0.128u"kJ * mol^-1"; atol=1e-9u"kJ * mol^-1", ) @test isapprox( - potential_energy(inter, dr13, c1, c3, a1, a1, boundary), + potential_energy(inter, dr13, a1, a1), -0.1130893709u"kJ * mol^-1"; atol=1e-9u"kJ * mol^-1", ) inter = SoftSphere() @test isapprox( - force(inter, dr12, c1, c2, a1, a1, boundary), + force(inter, dr12, a1, a1), SVector(32.0, 0.0, 0.0)u"kJ * mol^-1 * nm^-1"; atol=1e-9u"kJ * mol^-1 * nm^-1", ) @test isapprox( - force(inter, dr13, c1, c3, a1, a1, boundary), + force(inter, dr13, a1, a1), SVector(0.7602324486, 0.0, 0.0)u"kJ * mol^-1 * nm^-1"; atol=1e-9u"kJ * mol^-1 * nm^-1", ) @test isapprox( - potential_energy(inter, dr12, c1, c2, a1, a1, boundary), + potential_energy(inter, dr12, a1, a1), 0.8u"kJ * mol^-1"; atol=1e-9u"kJ * mol^-1", ) @test isapprox( - potential_energy(inter, dr13, c1, c3, a1, a1, boundary), + potential_energy(inter, dr13, a1, a1), 0.0253410816u"kJ * mol^-1"; atol=1e-9u"kJ * mol^-1", ) @@ -92,46 +92,46 @@ buck_dr12 = vector(buck_c1, buck_c2, buck_boundary) buck_dr13 = vector(buck_c1, buck_c3, buck_boundary) - inter = Buckingham(force_units=u"eV * Å^-1", energy_units=u"eV") + inter = Buckingham() @test isapprox( - force(inter, buck_dr12, buck_c1, buck_c2, buck_a1, buck_a1, buck_boundary), + force(inter, buck_dr12, buck_a1, buck_a1, u"eV * Å^-1"), SVector(0.3876527503, 0.0, 0.0)u"eV * Å^-1"; atol=1e-9u"eV * Å^-1", ) @test isapprox( - force(inter, buck_dr13, buck_c1, buck_c3, buck_a1, buck_a1, buck_boundary), + force(inter, buck_dr13, buck_a1, buck_a1, u"eV * Å^-1"), SVector(-0.0123151202, 0.0, 0.0)u"eV * Å^-1"; atol=1e-9u"eV * Å^-1", ) @test isapprox( - potential_energy(inter, buck_dr12, buck_c1, buck_c2, buck_a1, buck_a1, buck_boundary), + potential_energy(inter, buck_dr12, buck_a1, buck_a1, u"eV"), 0.0679006736u"eV"; atol=1e-9u"eV", ) @test isapprox( - potential_energy(inter, buck_dr13, buck_c1, buck_c3, buck_a1, buck_a1, buck_boundary), + potential_energy(inter, buck_dr13, buck_a1, buck_a1, u"eV"), -0.0248014380u"eV"; atol=1e-9u"eV", ) for inter in (Coulomb(), CoulombSoftCore(α=1, λ=0, p=2)) @test isapprox( - force(inter, dr12, c1, c2, a1, a1, boundary), + force(inter, dr12, a1, a1), SVector(1543.727311, 0.0, 0.0)u"kJ * mol^-1 * nm^-1"; atol=1e-5u"kJ * mol^-1 * nm^-1", ) @test isapprox( - force(inter, dr13, c1, c3, a1, a1, boundary), + force(inter, dr13, a1, a1), SVector(868.3466125, 0.0, 0.0)u"kJ * mol^-1 * nm^-1"; atol=1e-5u"kJ * mol^-1 * nm^-1", ) @test isapprox( - potential_energy(inter, dr12, c1, c2, a1, a1, boundary), + potential_energy(inter, dr12, a1, a1), 463.1181933u"kJ * mol^-1"; atol=1e-5u"kJ * mol^-1", ) @test isapprox( - potential_energy(inter, dr13, c1, c3, a1, a1, boundary), + potential_energy(inter, dr13, a1, a1), 347.338645u"kJ * mol^-1"; atol=1e-5u"kJ * mol^-1", ) @@ -139,22 +139,22 @@ inter = CoulombSoftCore(α=1, λ=0.5, p=2) @test isapprox( - force(inter, dr12, c1, c2, a1, a1, boundary), + force(inter, dr12, a1, a1), SVector(1189.895726, 0.0, 0.0)u"kJ * mol^-1 * nm^-1"; atol=1e-5u"kJ * mol^-1 * nm^-1", ) @test isapprox( - force(inter, dr13, c1, c3, a1, a1, boundary), + force(inter, dr13, a1, a1), SVector(825.3456507, 0.0, 0.0)u"kJ * mol^-1 * nm^-1"; atol=1e-5u"kJ * mol^-1 * nm^-1", ) @test isapprox( - potential_energy(inter, dr12, c1, c2, a1, a1, boundary), + potential_energy(inter, dr12, a1, a1), 446.2108973u"kJ * mol^-1"; atol=1e-5u"kJ * mol^-1", ) @test isapprox( - potential_energy(inter, dr13, c1, c3, a1, a1, boundary), + potential_energy(inter, dr13, a1, a1), 344.8276396u"kJ * mol^-1"; atol=1e-5u"kJ * mol^-1", ) @@ -166,13 +166,12 @@ dr12_grav = vector(c1_grav, c2_grav, boundary_grav) inter = Gravity() @test isapprox( - force(inter, dr12_grav, c1_grav, c2_grav, a1_grav, a2_grav, boundary_grav), + force(inter, dr12_grav, a1_grav, a2_grav), SVector(-0.266972, 0.0, 0.0)u"kg * m * s^-2"; atol=1e-9u"kg * m * s^-2", ) @test isapprox( - potential_energy(inter, dr12_grav, c1_grav, c2_grav, - a1_grav, a2_grav, boundary_grav), + potential_energy(inter, dr12_grav, a1_grav, a2_grav), -1.33486u"kg * m^2 * s^-2"; atol=1e-9u"kg * m^2 * s^-2", ) diff --git a/test/minimization.jl b/test/minimization.jl index 61a83521e..83a10f0e8 100644 --- a/test/minimization.jl +++ b/test/minimization.jl @@ -29,7 +29,7 @@ atoms=[Atom(σ=0.4 / (2 ^ (1 / 6)), ϵ=1.0, mass=1.0) for i in 1:3], coords=coords, boundary=CubicBoundary(5.0), - pairwise_inters=(LennardJones(force_units=NoUnits, energy_units=NoUnits),), + pairwise_inters=(LennardJones(),), force_units=NoUnits, energy_units=NoUnits, ) diff --git a/test/simulation.jl b/test/simulation.jl index dd952a808..f7e12f0fe 100644 --- a/test/simulation.jl +++ b/test/simulation.jl @@ -391,8 +391,7 @@ end coords=ustrip_vec.(coords), boundary=CubicBoundary(ustrip.(boundary)), velocities=ustrip_vec.(u"nm/ps",velocities), - pairwise_inters=(LennardJones(use_neighbors=true, force_units=NoUnits, - energy_units=NoUnits),), + pairwise_inters=(LennardJones(use_neighbors=true),), neighbor_finder=DistanceNeighborFinder( eligible=trues(n_atoms, n_atoms), n_steps=10, @@ -478,12 +477,7 @@ end atoms = [Atom(index=j, mass=atom_mass, σ=2.8279u"Å", ϵ=0.074u"kcal* mol^-1") for j in 1:n_atoms] cons = SHAKE_RATTLE(constraints, n_atoms, 1e-8u"Å", 1e-8u"Å^2 * ps^-1") boundary = CubicBoundary(200.0u"Å") - lj = LennardJones( - cutoff=ShiftedPotentialCutoff(r_cut), - use_neighbors=true, - energy_units=u"kcal * mol^-1", - force_units=u"kcal * mol^-1 * Å^-1" - ) + lj = LennardJones(cutoff=ShiftedPotentialCutoff(r_cut), use_neighbors=true) neighbor_finder = DistanceNeighborFinder( eligible=trues(n_atoms, n_atoms), dist_cutoff=1.5*r_cut, @@ -1089,11 +1083,7 @@ end sys = System( fcc_crystal; velocities=velocities, - pairwise_inters=(LennardJones( - cutoff=ShiftedForceCutoff(r_cut), - energy_units=u"kJ * mol^-1", - force_units=u"kJ * mol^-1 * nm^-1", - ),), + pairwise_inters=(LennardJones(cutoff=ShiftedForceCutoff(r_cut)),), loggers=(tot_eng=TotalEnergyLogger(100),), energy_units=u"kJ * mol^-1", force_units=u"kJ * mol^-1 * nm^-1", diff --git a/test/zygote.jl b/test/zygote.jl index d1c261a07..04cc1e70e 100644 --- a/test/zygote.jl +++ b/test/zygote.jl @@ -1,5 +1,5 @@ @testset "Gradients" begin - inter = LennardJones(force_units=NoUnits, energy_units=NoUnits) + inter = LennardJones() boundary = CubicBoundary(5.0) a1, a2 = Atom(σ=0.3, ϵ=0.5), Atom(σ=0.3, ϵ=0.5) @@ -7,7 +7,7 @@ c1 = SVector(1.0, 1.0, 1.0) c2 = SVector(dist + 1.0, 1.0, 1.0) vec = vector(c1, c2, boundary) - F = force(inter, vec, c1, c2, a1, a2, boundary) + F = force(inter, vec, a1, a2, NoUnits) return F[1] end @@ -16,7 +16,7 @@ c1 = SVector(1.0, 1.0, 1.0) c2 = SVector(dist + 1.0, 1.0, 1.0) vec = vector(c1, c2, boundary) - potential_energy(inter, vec, c1, c2, a1, a2, boundary) + potential_energy(inter, vec, a1, a2, NoUnits) end return -grad[1] end @@ -54,19 +54,12 @@ end coords_dual = [ForwardDiff.Dual.(x, f32 ? 0.0f0 : 0.0) for x in coords] velocities_dual = [ForwardDiff.Dual.(x, f32 ? 0.0f0 : 0.0) for x in velocities] nb_cutoff = f32 ? 1.2f0 : 1.2 - lj = LennardJones( - cutoff=DistanceCutoff(nb_cutoff), - use_neighbors=true, - force_units=NoUnits, - energy_units=NoUnits, - ) + lj = LennardJones(cutoff=DistanceCutoff(nb_cutoff), use_neighbors=true) crf = CoulombReactionField( dist_cutoff=nb_cutoff, solvent_dielectric=f32 ? Float32(Molly.crf_solvent_dielectric) : Molly.crf_solvent_dielectric, use_neighbors=true, coulomb_const=f32 ? Float32(ustrip(Molly.coulombconst)) : ustrip(Molly.coulombconst), - force_units=NoUnits, - energy_units=NoUnits, ) pairwise_inters = pis ? (lj, crf) : () bond_is = gpu ? CuArray(Int32.(collect(1:(n_atoms ÷ 2)))) : Int32.(collect(1:(n_atoms ÷ 2))) From 004f38cebb70948f3641e3277153d90b213d4eae Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 28 Jun 2024 15:35:26 +0100 Subject: [PATCH 09/74] pairwise interaction docs --- docs/src/documentation.md | 47 ++++++++++++++++++++++++++++++++------- 1 file changed, 39 insertions(+), 8 deletions(-) diff --git a/docs/src/documentation.md b/docs/src/documentation.md index 642c09bb1..936ea214a 100644 --- a/docs/src/documentation.md +++ b/docs/src/documentation.md @@ -586,12 +586,17 @@ Next, you need to define a method for the [`force`](@ref) function acting betwee This has a set series of arguments: ```julia function Molly.force(inter::MyPairwiseInter, - vec_ij, - coord_i, - coord_j, - atom_i, - atom_j, - boundary) + vec_ij, + atom_i, + atom_j, + force_units, + special, + coord_i, + coord_j, + boundary, + velocity_i, + velocity_j, + step_n) # Replace this with your force calculation # A positive force causes the atoms to move apart f = 0.0 @@ -601,12 +606,17 @@ function Molly.force(inter::MyPairwiseInter, return fdr end ``` +Most of the arguments will generally not be used but are passed to allow maximum flexibility. +You can use `args...` to indicate unused further arguments, e.g. `Molly.force(inter::MyPairwiseInter, vec_ij, args...)`. `vec_ij` is the vector between the closest images of atoms `i` and `j` accounting for the periodic boundary conditions. Atom properties can be accessed, e.g. `atom_i.σ`. +`force_units` can be useful for returning a zero force under certain conditions. +`step_n` is the step number in the simulator, allowing time-dependent interactions. +Beware that this step counter starts from 1 every time [`simulate!`](@ref) is called. Typically the force function is where most computation time is spent during the simulation, so consider optimising this function if you want high performance. One nice feature of Molly is that this function will work on both the CPU and the GPU. -An optional final argument `special` is a `Bool` determining whether the atom pair interaction should be treated as special. +The argument `special` is a `Bool` determining whether the atom pair interaction should be treated as special. This is specified during neighbor finder construction. When simulating molecules, for example, non-bonded interactions for atoms in a 1-4 bonding arrangement (i-x-x-j) are often weighted by a factor such as 0.5. For interactions where this is relevant, `special` can be used to apply this weighting in the interaction. @@ -624,7 +634,28 @@ pairwise_inters = (MyPairwiseInter=MyPairwiseInter(),) For performance reasons it is best to [avoid containers with abstract type parameters](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-abstract-container-1), such as `Vector{PairwiseInteraction}`. If you wish to calculate potential energies or log the energy throughout a simulation, you will need to define a method for the [`potential_energy`](@ref) function. -This has the same arguments as [`force`](@ref) and should return a single value corresponding to the potential energy. +This has the same arguments as [`force`](@ref), except the fifth argument is the energy units not the force units, and should return a single value corresponding to the potential energy: +```julia +function Molly.potential_energy(inter::MyPairwiseInter, + vec_ij, + atom_i, + atom_j, + energy_units, + special, + coord_i, + coord_j, + boundary, + velocity_i, + velocity_j, + step_n) + # Example Lennard-Jones interaction + σ = (atom_i.σ + atom_j.σ) / 2 + ϵ = sqrt(atom_i.ϵ * atom_j.ϵ) + r = norm(vec_ij) + E = 4ϵ * ((σ/r)^6 - (σ/r)^12) + return E +end +``` ### Specific interactions From 0df6b9fe49b1881476881b75f98f1abb5d6beaff Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 28 Jun 2024 18:37:28 +0100 Subject: [PATCH 10/74] rename test groups --- .github/workflows/CI.yml | 4 ++-- docs/src/development.md | 2 +- test/{zygote.jl => gradients.jl} | 0 test/runtests.jl | 10 +++++----- 4 files changed, 8 insertions(+), 8 deletions(-) rename test/{zygote.jl => gradients.jl} (100%) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 6a3032422..34d07a863 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -26,8 +26,8 @@ jobs: arch: - x64 test: - - NotZygote - - Zygote + - NotGradients + - Gradients steps: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v2 diff --git a/docs/src/development.md b/docs/src/development.md index 54346e2d0..a42211bfd 100644 --- a/docs/src/development.md +++ b/docs/src/development.md @@ -8,7 +8,7 @@ Various environmental variables can be set to modify the tests: - `VISTESTS` determines whether to run the [GLMakie.jl](https://github.com/JuliaPlots/Makie.jl) plotting tests which will error on remote systems where a display is not available, default `VISTESTS=1`. - `GPUTESTS` determines whether to run the GPU tests, default `GPUTESTS=1`. - `DEVICE` determines which GPU to run the GPU tests on, default `DEVICE=0`. -- `GROUP` can be used to run a subset of the tests, options `All`/`Protein`/`Zygote`/`NotZygote`, default `GROUP=All`. +- `GROUP` can be used to run a subset of the tests, options `All`/`Protein`/`Gradients`/`NotGradients`, default `GROUP=All`. The CI run does not carry out all tests - for example the GPU tests are not run - and this is reflected in the code coverage. ## Benchmarks diff --git a/test/zygote.jl b/test/gradients.jl similarity index 100% rename from test/zygote.jl rename to test/gradients.jl diff --git a/test/runtests.jl b/test/runtests.jl index 5a7d3ab54..8f0e74461 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -23,7 +23,7 @@ using Test # Allow testing of particular components const GROUP = get(ENV, "GROUP", "All") -if GROUP in ("Protein", "Zygote", "NotZygote") +if GROUP in ("Protein", "Gradients", "NotGradients") @warn "Only running $GROUP tests as GROUP is set to $GROUP" end @@ -70,7 +70,7 @@ const temp_fp_viz = tempname(cleanup=true) * ".mp4" # Required for parallel gradient tests Enzyme.API.runtimeActivity!(true) -if GROUP in ("All", "NotZygote") +if GROUP in ("All", "NotGradients") # Some failures due to dependencies but there is an unbound args error Aqua.test_all( Molly; @@ -86,10 +86,10 @@ if GROUP in ("All", "NotZygote") include("agent.jl") end -if GROUP in ("All", "Protein", "NotZygote") +if GROUP in ("All", "Protein", "NotGradients") include("protein.jl") end -if GROUP in ("All", "Zygote") - include("zygote.jl") +if GROUP in ("All", "Gradients") + include("gradients.jl") end From e8a6dd5b396643e325134475419fb855a5b2a4f5 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 28 Jun 2024 18:39:54 +0100 Subject: [PATCH 11/74] remove deps --- Project.toml | 4 ---- src/Molly.jl | 4 ---- test/Project.toml | 2 -- test/runtests.jl | 2 -- 4 files changed, 12 deletions(-) diff --git a/Project.toml b/Project.toml index 53ff5c50e..9dd4e0df4 100644 --- a/Project.toml +++ b/Project.toml @@ -20,7 +20,6 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" EzXML = "8f5d6c58-4d21-5cfd-889c-e3ad7ee6a615" FLoops = "cc61a311-1640-44b5-9fba-1b764f453329" -ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" KernelDensity = "5ab0869b-81aa-558d-bb23-cbf5423bbe9b" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -36,7 +35,6 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a" UnitfulChainRules = "f31437dd-25a7-4345-875f-756556e6935d" UnsafeAtomicsLLVM = "d80eeb9a-aca5-4d75-85e5-170c8b632249" -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [weakdeps] Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" @@ -65,7 +63,6 @@ Distributions = "0.23, 0.24, 0.25" Enzyme = "0.11.15, 0.12" EzXML = "1" FLoops = "0.2" -ForwardDiff = "0.10.35" GLMakie = "0.8, 0.9, 0.10" Graphs = "1.8" KernelDensity = "0.5, 0.6" @@ -83,5 +80,4 @@ Unitful = "1" UnitfulAtomic = "1" UnitfulChainRules = "0.1.2" UnsafeAtomicsLLVM = "0.1, 0.2" -Zygote = "0.6.67" julia = "1.9" diff --git a/src/Molly.jl b/src/Molly.jl index ebe7f1d90..3b3a0eea5 100644 --- a/src/Molly.jl +++ b/src/Molly.jl @@ -20,7 +20,6 @@ using Distributions using Enzyme using EzXML using FLoops -using ForwardDiff using Graphs using KernelDensity using NearestNeighbors @@ -30,7 +29,6 @@ using Unitful using UnitfulAtomic using UnitfulChainRules using UnsafeAtomicsLLVM -using Zygote using LinearAlgebra using Random @@ -68,8 +66,6 @@ include("neighbors.jl") include("loggers.jl") include("analysis.jl") include("setup.jl") -include("chain_rules.jl") -include("zygote.jl") include("gradients.jl") end diff --git a/test/Project.toml b/test/Project.toml index 60e2d7bb8..8d6edd647 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -8,14 +8,12 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" -ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SimpleCrystals = "64031d72-e220-11ed-1a7e-43a2532b2fa8" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] Aqua = "0.8" diff --git a/test/runtests.jl b/test/runtests.jl index 8f0e74461..398175820 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -8,9 +8,7 @@ import BioStructures # Imported to avoid clashing names using CUDA import Enzyme using FiniteDifferences -using ForwardDiff import SimpleCrystals -using Zygote using DelimitedFiles using LinearAlgebra From 6cc14b304894dcc72e8f4710a81cb90faf2840ab Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 28 Jun 2024 18:41:44 +0100 Subject: [PATCH 12/74] force and energy arguments --- src/cuda.jl | 55 +++++++++++++++++------------- src/energy.jl | 67 +++++++++++++++++-------------------- src/force.jl | 63 +++++++++++++++++----------------- src/gradients.jl | 10 ++---- src/interactions/coulomb.jl | 2 +- src/spatial.jl | 19 ++++++----- src/types.jl | 2 ++ src/units.jl | 28 ++-------------- test/basic.jl | 6 ---- 9 files changed, 110 insertions(+), 142 deletions(-) diff --git a/src/cuda.jl b/src/cuda.jl index 15b2e5d74..fdc71eee2 100644 --- a/src/cuda.jl +++ b/src/cuda.jl @@ -29,32 +29,33 @@ function cuda_threads_blocks_specific(n_inters) return n_threads_gpu, n_blocks end -function pairwise_force_gpu(coords::AbstractArray{SVector{D, C}}, atoms, boundary, +function pairwise_force_gpu(coords::AbstractArray{SVector{D, C}}, velocities, atoms, boundary, pairwise_inters, nbs, force_units, ::Val{T}) where {D, C, T} fs_mat = CUDA.zeros(T, D, length(atoms)) if typeof(nbs) == NoNeighborList kernel = @cuda launch=false pairwise_force_kernel_nonl!( - fs_mat, coords, atoms, boundary, pairwise_inters, Val(D), Val(force_units)) + fs_mat, coords, velocities, atoms, boundary, pairwise_inters, Val(D), Val(force_units)) conf = launch_configuration(kernel.fun) threads_basic = parse(Int, get(ENV, "MOLLY_GPUNTHREADS_PAIRWISE", "512")) nthreads = min(length(atoms), threads_basic, conf.threads) nthreads = cld(nthreads, WARPSIZE) * WARPSIZE n_blocks_i = cld(length(atoms), WARPSIZE) n_blocks_j = cld(length(atoms), nthreads) - kernel(fs_mat, coords, atoms, boundary, pairwise_inters, Val(D), Val(force_units); + kernel(fs_mat, coords, velocities, atoms, boundary, pairwise_inters, Val(D), Val(force_units); threads=nthreads, blocks=(n_blocks_i, n_blocks_j)) else n_threads_gpu, n_blocks = cuda_threads_blocks_pairwise(length(nbs)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks pairwise_force_kernel_nl!( - fs_mat, coords, atoms, boundary, pairwise_inters, nbs, Val(D), Val(force_units)) + fs_mat, coords, velocities, atoms, boundary, pairwise_inters, nbs, Val(D), Val(force_units)) end return fs_mat end -function pairwise_force_kernel_nl!(forces, coords_var, atoms_var, boundary, inters, +function pairwise_force_kernel_nl!(forces, coords_var, velocities_var, atoms_var, boundary, inters, neighbors_var, ::Val{D}, ::Val{F}) where {D, F} coords = CUDA.Const(coords_var) + velocities = CUDA.Const(velocities_var) atoms = CUDA.Const(atoms_var) neighbors = CUDA.Const(neighbors_var) @@ -62,7 +63,8 @@ function pairwise_force_kernel_nl!(forces, coords_var, atoms_var, boundary, inte @inbounds if inter_i <= length(neighbors) i, j, special = neighbors[inter_i] - f = sum_pairwise_forces(inters, coords[i], coords[j], atoms[i], atoms[j], boundary, special, Val(F)) + f = sum_pairwise_forces(inters, atoms[i], atoms[j], Val(F), special, coords[i], coords[j], + boundary, velocities[i], velocities[j], 0) for dim in 1:D fval = ustrip(f[dim]) Atomix.@atomic :monotonic forces[dim, i] += -fval @@ -112,9 +114,10 @@ That's why the calculations are done in the following order: h | 1 2 3 4 5 6 ``` =# -function pairwise_force_kernel_nonl!(forces::AbstractArray{T}, coords_var, atoms_var, boundary, inters, +function pairwise_force_kernel_nonl!(forces::AbstractArray{T}, coords_var, velocities_var, atoms_var, boundary, inters, ::Val{D}, ::Val{F}) where {T, D, F} coords = CUDA.Const(coords_var) + velocities = CUDA.Const(velocities_var) atoms = CUDA.Const(atoms_var) n_atoms = length(atoms) @@ -133,12 +136,13 @@ function pairwise_force_kernel_nonl!(forces::AbstractArray{T}, coords_var, atoms if i_0_tile + warpsize() > n_atoms || j_0_tile + warpsize() > n_atoms @inbounds if i <= n_atoms njs = min(warpsize(), n_atoms - j_0_tile) - atom_i, coord_i = atoms[i], coords[i] + atom_i, coord_i, vel_i = atoms[i], coords[i], velocities[i] for del_j in 1:njs j = j_0_tile + del_j if i != j - atom_j, coord_j = atoms[j], coords[j] - f = sum_pairwise_forces(inters, coord_i, coord_j, atom_i, atom_j, boundary, false, Val(F)) + atom_j, coord_j, vel_j = atoms[j], coords[j], velocities[j] + f = sum_pairwise_forces(inters, atom_i, atom_j, Val(F), false, coord_i, coord_j, + boundary, vel_i, vel_j, 0) for dim in 1:D forces_shmem[dim, tidx] += -ustrip(f[dim]) end @@ -157,12 +161,13 @@ function pairwise_force_kernel_nonl!(forces::AbstractArray{T}, coords_var, atoms tilesteps -= 1 end - atom_i, coord_i = atoms[i], coords[i] - coord_j = coords[j] + atom_i, coord_i, vel_i = atoms[i], coords[i], velocities[i] + coord_j, vel_j = coords[j], velocities[j] @inbounds for _ in 1:tilesteps sync_warp() atom_j = atoms[j] - f = sum_pairwise_forces(inters, coord_i, coord_j, atom_i, atom_j, boundary, false, Val(F)) + f = sum_pairwise_forces(inters, atom_i, atom_j, Val(F), false, coord_i, coord_j, + boundary, vel_i, vel_j, 0) for dim in 1:D forces_shmem[dim, tidx] += -ustrip(f[dim]) end @@ -177,11 +182,12 @@ function pairwise_force_kernel_nonl!(forces::AbstractArray{T}, coords_var, atoms return nothing end -@inline function sum_pairwise_forces(inters, coord_i, coord_j, atom_i, atom_j, - boundary, special, ::Val{F}) where F +@inline function sum_pairwise_forces(inters, atom_i, atom_j, ::Val{F}, special, coord_i, coord_j, + boundary, vel_i, vel_j, step_n) where F dr = vector(coord_i, coord_j, boundary) f_tuple = ntuple(length(inters)) do inter_type_i - force_gpu(inters[inter_type_i], dr, coord_i, coord_j, atom_i, atom_j, boundary, special) + force_gpu(inters[inter_type_i], dr, atom_i, atom_j, F, special, coord_i, coord_j, boundary, + vel_i, vel_j, step_n) end f = sum(f_tuple) if unit(f[1]) != F @@ -328,18 +334,19 @@ function specific_force_4_atoms_kernel!(forces, coords_var, boundary, is_var, js return nothing end -function pairwise_pe_gpu(coords::AbstractArray{SVector{D, C}}, atoms, boundary, +function pairwise_pe_gpu(coords::AbstractArray{SVector{D, C}}, velocities, atoms, boundary, pairwise_inters, nbs, energy_units, ::Val{T}) where {D, C, T} pe_vec = CUDA.zeros(T, 1) n_threads_gpu, n_blocks = cuda_threads_blocks_pairwise(length(nbs)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks pairwise_pe_kernel!( - pe_vec, coords, atoms, boundary, pairwise_inters, nbs, Val(energy_units)) + pe_vec, coords, velocities, atoms, boundary, pairwise_inters, nbs, Val(energy_units)) return pe_vec end -function pairwise_pe_kernel!(energy, coords_var, atoms_var, boundary, inters, neighbors_var, +function pairwise_pe_kernel!(energy, coords_var, velocities_var, atoms_var, boundary, inters, neighbors_var, ::Val{E}) where E coords = CUDA.Const(coords_var) + velocities = CUDA.Const(velocities_var) atoms = CUDA.Const(atoms_var) neighbors = CUDA.Const(neighbors_var) @@ -347,13 +354,13 @@ function pairwise_pe_kernel!(energy, coords_var, atoms_var, boundary, inters, ne @inbounds if inter_i <= length(neighbors) i, j, special = neighbors[inter_i] - coord_i, coord_j = coords[i], coords[j] + coord_i, coord_j, vel_i, vel_j = coords[i], coords[j], velocities[i], velocities[j] dr = vector(coord_i, coord_j, boundary) - pe = potential_energy_gpu(inters[1], dr, coord_i, coord_j, atoms[i], atoms[j], - boundary, special) + pe = potential_energy_gpu(inters[1], dr, atoms[i], atoms[j], E, special, coord_i, coord_j, + boundary, vel_i, vel_j, 0) for inter in inters[2:end] - pe += potential_energy_gpu(inter, dr, coord_i, coord_j, atoms[i], atoms[j], - boundary, special) + pe += potential_energy_gpu(inter, dr, atoms[i], atoms[j], E, special, coord_i, coord_j, + boundary, vel_i, vel_j, 0) end if unit(pe) != E error("wrong energy unit returned, was expecting $E but got $(unit(pe))") diff --git a/src/energy.jl b/src/energy.jl index bc536e6ab..bed391a0f 100644 --- a/src/energy.jl +++ b/src/energy.jl @@ -53,8 +53,8 @@ end Calculate the potential energy of a system using the pairwise, specific and general interactions. - potential_energy(inter::PairwiseInteraction, vec_ij, coord_i, coord_j, - atom_i, atom_j, boundary) + potential_energy(inter::PairwiseInteraction, vec_ij, atom_i, atom_j, energy_units, special, + coord_i, coord_j, boundary, velocity_i, velocity_j, step_n) potential_energy(inter::SpecificInteraction, coords_i, coords_j, boundary) potential_energy(inter::SpecificInteraction, coords_i, coords_j, @@ -70,6 +70,13 @@ function potential_energy(sys; n_threads::Integer=Threads.nthreads()) return potential_energy(sys, find_neighbors(sys; n_threads=n_threads); n_threads=n_threads) end +# Allow GPU-specific potential energy functions to be defined if required +potential_energy_gpu(inter::PairwiseInteraction, dr, ai, aj, eu, sp, ci, cj, bnd, vi, vj, sn) = potential_energy(inter, dr, ai, aj, eu, sp, ci, cj, bnd, vi, vj, sn) +potential_energy_gpu(inter::SpecificInteraction, ci, bnd) = potential_energy(inter, ci, bnd) +potential_energy_gpu(inter::SpecificInteraction, ci, cj, bnd) = potential_energy(inter, ci, cj, bnd) +potential_energy_gpu(inter::SpecificInteraction, ci, cj, ck, bnd) = potential_energy(inter, ci, cj, ck, bnd) +potential_energy_gpu(inter::SpecificInteraction, ci, cj, ck, cl, bnd) = potential_energy(inter, ci, cj, ck, cl, bnd) + function potential_energy(sys::System{D, false}, neighbors; n_threads::Integer=Threads.nthreads()) where D pairwise_inters_nonl = filter(!use_neighbors, values(sys.pairwise_inters)) @@ -80,7 +87,7 @@ function potential_energy(sys::System{D, false}, neighbors; sils_4_atoms = filter(il -> il isa InteractionList4Atoms, values(sys.specific_inter_lists)) ft = typeof(ustrip(sys.coords[1][1])) # Allow types like those from Measurements.jl - pe = potential_energy_pair_spec(sys.coords, sys.atoms, pairwise_inters_nonl, pairwise_inters_nl, + pe = potential_energy_pair_spec(sys.coords, sys.velocities, sys.atoms, pairwise_inters_nonl, pairwise_inters_nl, sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, sys.boundary, sys.energy_units, neighbors, n_threads, Val(ft)) @@ -91,18 +98,18 @@ function potential_energy(sys::System{D, false}, neighbors; return pe end -function potential_energy_pair_spec(coords, atoms, pairwise_inters_nonl, pairwise_inters_nl, +function potential_energy_pair_spec(coords, velocities, atoms, pairwise_inters_nonl, pairwise_inters_nl, sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, boundary, energy_units, neighbors, n_threads, val_ft::Val{T}) where T pe_vec = zeros(T, 1) - potential_energy_pair_spec!(pe_vec, coords, atoms, pairwise_inters_nonl, pairwise_inters_nl, + potential_energy_pair_spec!(pe_vec, coords, velocities, atoms, pairwise_inters_nonl, pairwise_inters_nl, sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, boundary, energy_units, neighbors, n_threads, val_ft) return pe_vec[1] * energy_units end -function potential_energy_pair_spec!(pe_vec, coords, atoms, pairwise_inters_nonl, +function potential_energy_pair_spec!(pe_vec, coords, velocities, atoms, pairwise_inters_nonl, pairwise_inters_nl, sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, boundary, energy_units, neighbors, n_threads, ::Val{T}) where T pe_sum = zero(T) @@ -116,11 +123,11 @@ function potential_energy_pair_spec!(pe_vec, coords, atoms, pairwise_inters_nonl for i in chunk_i:n_threads:n_atoms for j in (i + 1):n_atoms dr = vector(coords[i], coords[j], boundary) - pe = potential_energy(pairwise_inters_nonl[1], dr, coords[i], coords[j], atoms[i], - atoms[j], boundary) + pe = potential_energy(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], energy_units, false, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) for inter in pairwise_inters_nonl[2:end] - pe += potential_energy(inter, dr, coords[i], coords[j], atoms[i], - atoms[j], boundary) + pe += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, false, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) end check_energy_units(pe, energy_units) pe_sum_chunks[chunk_i] += ustrip(pe) @@ -137,11 +144,11 @@ function potential_energy_pair_spec!(pe_vec, coords, atoms, pairwise_inters_nonl for ni in chunk_i:n_threads:length(neighbors) i, j, special = neighbors[ni] dr = vector(coords[i], coords[j], boundary) - pe = potential_energy(pairwise_inters_nl[1], dr, coords[i], coords[j], atoms[i], - atoms[j], boundary, special) + pe = potential_energy(pairwise_inters_nl[1], dr, atoms[i], atoms[j], energy_units, special, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) for inter in pairwise_inters_nl[2:end] - pe += potential_energy(inter, dr, coords[i], coords[j], atoms[i], - atoms[j], boundary, special) + pe += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, special, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) end check_energy_units(pe, energy_units) pe_sum_chunks[chunk_i] += ustrip(pe) @@ -156,11 +163,11 @@ function potential_energy_pair_spec!(pe_vec, coords, atoms, pairwise_inters_nonl for i in 1:n_atoms for j in (i + 1):n_atoms dr = vector(coords[i], coords[j], boundary) - pe = potential_energy(pairwise_inters_nonl[1], dr, coords[i], coords[j], atoms[i], - atoms[j], boundary) + pe = potential_energy(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], energy_units, false, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) for inter in pairwise_inters_nonl[2:end] - pe += potential_energy(inter, dr, coords[i], coords[j], atoms[i], - atoms[j], boundary) + pe += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, false, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) end check_energy_units(pe, energy_units) pe_sum += ustrip(pe) @@ -175,11 +182,11 @@ function potential_energy_pair_spec!(pe_vec, coords, atoms, pairwise_inters_nonl for ni in eachindex(neighbors) i, j, special = neighbors[ni] dr = vector(coords[i], coords[j], boundary) - pe = potential_energy(pairwise_inters_nl[1], dr, coords[i], coords[j], atoms[i], - atoms[j], boundary, special) + pe = potential_energy(pairwise_inters_nl[1], dr, atoms[i], atoms[j], energy_units, special, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) for inter in pairwise_inters_nl[2:end] - pe += potential_energy(inter, dr, coords[i], coords[j], atoms[i], - atoms[j], boundary, special) + pe += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, special, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) end check_energy_units(pe, energy_units) pe_sum += ustrip(pe) @@ -233,7 +240,7 @@ function potential_energy(sys::System{D, true, T}, neighbors; pairwise_inters_nonl = filter(!use_neighbors, values(sys.pairwise_inters)) if length(pairwise_inters_nonl) > 0 nbs = NoNeighborList(n_atoms) - pe_vec += pairwise_pe_gpu(sys.coords, sys.atoms, sys.boundary, pairwise_inters_nonl, + pe_vec += pairwise_pe_gpu(sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nonl, nbs, sys.energy_units, val_ft) end @@ -244,7 +251,7 @@ function potential_energy(sys::System{D, true, T}, neighbors; end if length(neighbors) > 0 nbs = @view neighbors.list[1:neighbors.n] - pe_vec += pairwise_pe_gpu(sys.coords, sys.atoms, sys.boundary, pairwise_inters_nl, + pe_vec += pairwise_pe_gpu(sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nl, nbs, sys.energy_units, val_ft) end end @@ -264,15 +271,3 @@ function potential_energy(sys::System{D, true, T}, neighbors; return pe * sys.energy_units end - -function potential_energy(inter, dr, coord_i, coord_j, atom_i, atom_j, boundary, special) - # Fallback for interactions where special interactions are not relevant - return potential_energy(inter, dr, coord_i, coord_j, atom_i, atom_j, boundary) -end - -# Allow GPU-specific potential energy functions to be defined if required -potential_energy_gpu(inter::PairwiseInteraction, dr, ci, cj, ai, aj, bnd, spec) = potential_energy(inter, dr, ci, cj, ai, aj, bnd, spec) -potential_energy_gpu(inter::SpecificInteraction, ci, bnd) = potential_energy(inter, ci, bnd) -potential_energy_gpu(inter::SpecificInteraction, ci, cj, bnd) = potential_energy(inter, ci, cj, bnd) -potential_energy_gpu(inter::SpecificInteraction, ci, cj, ck, bnd) = potential_energy(inter, ci, cj, ck, bnd) -potential_energy_gpu(inter::SpecificInteraction, ci, cj, ck, cl, bnd) = potential_energy(inter, ci, cj, ck, cl, bnd) diff --git a/src/force.jl b/src/force.jl index fe0ba8165..3d65277b9 100644 --- a/src/force.jl +++ b/src/force.jl @@ -25,10 +25,8 @@ function accelerations(sys, neighbors; n_threads::Integer=Threads.nthreads()) end """ - force(inter::PairwiseInteraction, vec_ij, coord_i, coord_j, - atom_i, atom_j, boundary) - force(inter::PairwiseInteraction, vec_ij, coord_i, coord_j, - atom_i, atom_j, boundary, special) + force(inter::PairwiseInteraction, vec_ij, atom_i, atom_j, force_units, special, + coord_i, coord_j, boundary, velocity_i, velocity_j, step_n) force(inter::SpecificInteraction, coord_i, coord_j, boundary) force(inter::SpecificInteraction, coord_i, coord_j, @@ -42,17 +40,14 @@ For [`PairwiseInteraction`](@ref)s returns a single force vector and for [`SpecificInteraction`](@ref)s returns a type such as [`SpecificForce2Atoms`](@ref). Custom pairwise and specific interaction types should implement this function. """ -function force(inter, dr, coord_i, coord_j, atom_i, atom_j, boundary, special) - # Fallback for interactions where special interactions are not relevant - return force(inter, dr, coord_i, coord_j, atom_i, atom_j, boundary) -end +function force end # Allow GPU-specific force functions to be defined if required -force_gpu(inter, dr, ci, cj, ai, aj, bnd, spec) = force(inter, dr, ci, cj, ai, aj, bnd, spec) -force_gpu(inter, ci, bnd) = force(inter, ci, bnd) -force_gpu(inter, ci, cj, bnd) = force(inter, ci, cj, bnd) -force_gpu(inter, ci, cj, ck, bnd) = force(inter, ci, cj, ck, bnd) -force_gpu(inter, ci, cj, ck, cl, bnd) = force(inter, ci, cj, ck, cl, bnd) +force_gpu(inter::PairwiseInteraction, dr, ai, aj, fu, sp, ci, cj, bnd, vi, vj, sn) = force(inter, dr, ai, aj, fu, sp, ci, cj, bnd, vi, vj, sn) +force_gpu(inter::SpecificInteraction, ci, bnd) = force(inter, ci, bnd) +force_gpu(inter::SpecificInteraction, ci, cj, bnd) = force(inter, ci, cj, bnd) +force_gpu(inter::SpecificInteraction, ci, cj, ck, bnd) = force(inter, ci, cj, ck, bnd) +force_gpu(inter::SpecificInteraction, ci, cj, ck, cl, bnd) = force(inter, ci, cj, ck, cl, bnd) """ SpecificForce1Atoms(f1) @@ -138,7 +133,7 @@ function forces(sys::System{D, false}, neighbors; sils_3_atoms = filter(il -> il isa InteractionList3Atoms, values(sys.specific_inter_lists)) sils_4_atoms = filter(il -> il isa InteractionList4Atoms, values(sys.specific_inter_lists)) - fs = forces_pair_spec(sys.coords, sys.atoms, pairwise_inters_nonl, pairwise_inters_nl, + fs = forces_pair_spec(sys.coords, sys.velocities, sys.atoms, pairwise_inters_nonl, pairwise_inters_nl, sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, sys.boundary, sys.force_units, neighbors, n_threads) @@ -149,17 +144,17 @@ function forces(sys::System{D, false}, neighbors; return fs end -function forces_pair_spec(coords, atoms, pairwise_inters_nonl, pairwise_inters_nl, +function forces_pair_spec(coords, velocities, atoms, pairwise_inters_nonl, pairwise_inters_nl, sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, boundary, force_units, neighbors, n_threads) fs = ustrip_vec.(zero(coords)) - forces_pair_spec!(fs, coords, atoms, pairwise_inters_nonl, pairwise_inters_nl, + forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, pairwise_inters_nl, sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, boundary, force_units, neighbors, n_threads) return fs * force_units end -function forces_pair_spec!(fs, coords, atoms, pairwise_inters_nonl, pairwise_inters_nl, +function forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, pairwise_inters_nl, sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, boundary, force_units, neighbors, n_threads) n_atoms = length(coords) @@ -171,10 +166,11 @@ function forces_pair_spec!(fs, coords, atoms, pairwise_inters_nonl, pairwise_int for i in chunk_i:n_threads:n_atoms for j in (i + 1):n_atoms dr = vector(coords[i], coords[j], boundary) - f = force(pairwise_inters_nonl[1], dr, coords[i], coords[j], atoms[i], - atoms[j], boundary) + f = force(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], force_units, false, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) for inter in pairwise_inters_nonl[2:end] - f += force(inter, dr, coords[i], coords[j], atoms[i], atoms[j], boundary) + f += force(inter, dr, atoms[i], atoms[j], force_units, false, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) end check_force_units(f, force_units) f_ustrip = ustrip.(f) @@ -193,11 +189,11 @@ function forces_pair_spec!(fs, coords, atoms, pairwise_inters_nonl, pairwise_int for ni in chunk_i:n_threads:length(neighbors) i, j, special = neighbors[ni] dr = vector(coords[i], coords[j], boundary) - f = force(pairwise_inters_nl[1], dr, coords[i], coords[j], atoms[i], - atoms[j], boundary, special) + f = force(pairwise_inters_nl[1], dr, atoms[i], atoms[j], force_units, special, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) for inter in pairwise_inters_nl[2:end] - f += force(inter, dr, coords[i], coords[j], atoms[i], atoms[j], boundary, - special) + f += force(inter, dr, atoms[i], atoms[j], force_units, special, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) end check_force_units(f, force_units) f_ustrip = ustrip.(f) @@ -213,10 +209,11 @@ function forces_pair_spec!(fs, coords, atoms, pairwise_inters_nonl, pairwise_int for i in 1:n_atoms for j in (i + 1):n_atoms dr = vector(coords[i], coords[j], boundary) - f = force(pairwise_inters_nonl[1], dr, coords[i], coords[j], atoms[i], - atoms[j], boundary) + f = force(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], force_units, false, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) for inter in pairwise_inters_nonl[2:end] - f += force(inter, dr, coords[i], coords[j], atoms[i], atoms[j], boundary) + f += force(inter, dr, atoms[i], atoms[j], force_units, false, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) end check_force_units(f, force_units) f_ustrip = ustrip.(f) @@ -233,11 +230,11 @@ function forces_pair_spec!(fs, coords, atoms, pairwise_inters_nonl, pairwise_int for ni in eachindex(neighbors) i, j, special = neighbors[ni] dr = vector(coords[i], coords[j], boundary) - f = force(pairwise_inters_nl[1], dr, coords[i], coords[j], atoms[i], - atoms[j], boundary, special) + f = force(pairwise_inters_nl[1], dr, atoms[i], atoms[j], force_units, special, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) for inter in pairwise_inters_nl[2:end] - f += force(inter, dr, coords[i], coords[j], atoms[i], atoms[j], boundary, - special) + f += force(inter, dr, atoms[i], atoms[j], force_units, special, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) end check_force_units(f, force_units) f_ustrip = ustrip.(f) @@ -304,7 +301,7 @@ function forces(sys::System{D, true, T}, neighbors; pairwise_inters_nonl = filter(!use_neighbors, values(sys.pairwise_inters)) if length(pairwise_inters_nonl) > 0 nbs = NoNeighborList(n_atoms) - fs_mat += pairwise_force_gpu(sys.coords, sys.atoms, sys.boundary, pairwise_inters_nonl, + fs_mat += pairwise_force_gpu(sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nonl, nbs, sys.force_units, val_ft) end @@ -315,7 +312,7 @@ function forces(sys::System{D, true, T}, neighbors; end if length(neighbors) > 0 nbs = @view neighbors.list[1:neighbors.n] - fs_mat += pairwise_force_gpu(sys.coords, sys.atoms, sys.boundary, pairwise_inters_nl, + fs_mat += pairwise_force_gpu(sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nl, nbs, sys.force_units, val_ft) end end diff --git a/src/gradients.jl b/src/gradients.jl index 7238d6bee..b621e37db 100644 --- a/src/gradients.jl +++ b/src/gradients.jl @@ -165,16 +165,14 @@ function inject_interaction_list(inter::InteractionList4Atoms, params_dic, gpu) InteractionList4Atoms(inter.is, inter.js, inter.ks, inter.ls, inters_grad, inter.types) end -function inject_interaction(inter::LennardJones{S, C, W, WS, F, E}, params_dic) where {S, C, W, WS, F, E} +function inject_interaction(inter::LennardJones{S, C, W, WS}, params_dic) where {S, C, W, WS} key_prefix = "inter_LJ_" - LennardJones{S, C, W, WS, F, E}( + LennardJones{S, C, W, WS}( inter.cutoff, inter.use_neighbors, inter.lorentz_mixing, dict_get(params_dic, key_prefix * "weight_14", inter.weight_special), dict_get(params_dic, key_prefix * "weight_solute_solvent", inter.weight_solute_solvent), - inter.force_units, - inter.energy_units, ) end @@ -185,8 +183,6 @@ function inject_interaction(inter::Coulomb, params_dic) inter.use_neighbors, dict_get(params_dic, key_prefix * "weight_14", inter.weight_special), dict_get(params_dic, key_prefix * "coulomb_const", inter.coulomb_const), - inter.force_units, - inter.energy_units, ) end @@ -198,8 +194,6 @@ function inject_interaction(inter::CoulombReactionField, params_dic) inter.use_neighbors, dict_get(params_dic, key_prefix * "weight_14", inter.weight_special), dict_get(params_dic, key_prefix * "coulomb_const", inter.coulomb_const), - inter.force_units, - inter.energy_units, ) end diff --git a/src/interactions/coulomb.jl b/src/interactions/coulomb.jl index 2811a6d99..db43fa694 100644 --- a/src/interactions/coulomb.jl +++ b/src/interactions/coulomb.jl @@ -26,7 +26,7 @@ function Coulomb(; cutoff=NoCutoff(), use_neighbors=false, weight_special=1, - coulomb_const=coulombconst, + coulomb_const=coulombconst) return Coulomb{typeof(cutoff), typeof(weight_special), typeof(coulomb_const)}( cutoff, use_neighbors, weight_special, coulomb_const) end diff --git a/src/spatial.jl b/src/spatial.jl index be43acfdc..87e2909bc 100644 --- a/src/spatial.jl +++ b/src/spatial.jl @@ -714,14 +714,14 @@ end function virial(sys::System{D, G, T}, neighbors_dev, pairwise_inters_nonl, pairwise_inters_nl) where {D, G, T} if G - coords, atoms = Array(sys.coords), Array(sys.atoms) + coords, velocities, atoms = Array(sys.coords), Array(sys.velocities), Array(sys.atoms) if isnothing(neighbors_dev) neighbors = neighbors_dev else neighbors = NeighborList(neighbors_dev.n, Array(neighbors_dev.list)) end else - coords, atoms = sys.coords, sys.atoms + coords, velocities, atoms = sys.coords, sys.velocities, sys.atoms neighbors = neighbors_dev end @@ -733,10 +733,11 @@ function virial(sys::System{D, G, T}, neighbors_dev, pairwise_inters_nonl, for i in 1:n_atoms for j in (i + 1):n_atoms dr = vector(coords[i], coords[j], boundary) - f = force(pairwise_inters_nonl[1], dr, coords[i], coords[j], atoms[i], - atoms[j], boundary) + f = force(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], sys.force_units, false, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) for inter in pairwise_inters_nonl[2:end] - f += force(inter, dr, coords[i], coords[j], atoms[i], atoms[j], boundary) + f += force(inter, dr, atoms[i], atoms[j], sys.force_units, false, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) end v += dot(f, dr) end @@ -750,11 +751,11 @@ function virial(sys::System{D, G, T}, neighbors_dev, pairwise_inters_nonl, for ni in eachindex(neighbors) i, j, special = neighbors[ni] dr = vector(coords[i], coords[j], boundary) - f = force(pairwise_inters_nl[1], dr, coords[i], coords[j], atoms[i], - atoms[j], boundary, special) + f = force(pairwise_inters_nl[1], dr, atoms[i], atoms[j], sys.force_units, special, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) for inter in pairwise_inters_nl[2:end] - f += force(inter, dr, coords[i], coords[j], atoms[i], atoms[j], boundary, - special) + f += force(inter, dr, atoms[i], atoms[j], sys.force_units, special, + coords[i], coords[j], boundary, velocities[i], velocities[j], 0) end v += dot(f, dr) end diff --git a/src/types.jl b/src/types.jl index a4a4c7b3f..2996a2670 100644 --- a/src/types.jl +++ b/src/types.jl @@ -1342,3 +1342,5 @@ struct ASECalculator{T} ase_atoms::T # T will be Py but that is not available here ase_calc::T end + +iszero_value(x) = iszero(x) diff --git a/src/units.jl b/src/units.jl index 31934eae7..ef00e6332 100644 --- a/src/units.jl +++ b/src/units.jl @@ -19,10 +19,7 @@ function check_units(atoms, coords, velocities, energy_units, force_units, p_inters, s_inters, g_inters, boundary) masses = mass.(atoms) sys_units = check_system_units(masses, coords, velocities, energy_units, force_units) - - check_interaction_units(p_inters, s_inters, g_inters, sys_units) check_other_units(atoms, boundary, sys_units) - return sys_units end @@ -63,25 +60,6 @@ function check_system_units(masses, coords, velocities, energy_units, force_unit vel_units, mass_units, energy_units, force_units)) end -function check_interaction_units(p_inters, s_inters, g_inters, sys_units::NamedTuple) - for inter_tuple in [p_inters, s_inters, g_inters] - for inter in inter_tuple - if hasproperty(inter, :energy_units) - if inter.energy_units != sys_units[:energy] - throw(ArgumentError("energy units passed to system do not match those passed in an interaction")) - end - end - - if hasproperty(inter, :force_units) - if inter.force_units != sys_units[:force] - throw(ArgumentError("force units passed to system do not match those passed in an interaction")) - end - end - end - end - -end - function check_other_units(atoms_dev, boundary, sys_units::NamedTuple) atoms = Array(atoms_dev) box_units = unit(length_type(boundary)) @@ -195,9 +173,7 @@ function convert_k_units(T, k, energy_units) # Use user-supplied unitless Boltzmann constant k_converted = T(k) else - Zygote.ignore() do - @warn "Units will be stripped from Boltzmann constant: energy_units was passed as NoUnits and units were provided on k: $(unit(k))" - end + @warn "Units will be stripped from Boltzmann constant: energy_units was passed as NoUnits and units were provided on k: $(unit(k))" k_converted = T(ustrip(k)) end elseif dimension(energy_units) in (u"𝐋^2 * 𝐌 * 𝐍^-1 * 𝐓^-2", u"𝐋^2 * 𝐌 * 𝐓^-2") @@ -223,6 +199,8 @@ function check_force_units(F, force_units) end end +check_force_units(F::SVector, force_units) = @inbounds check_force_units(F[1], force_units) + function energy_remove_mol(x) if dimension(x) == u"𝐋^2 * 𝐌 * 𝐍^-1 * 𝐓^-2" T = typeof(ustrip(x)) diff --git a/test/basic.jl b/test/basic.jl index d92da4bd6..00326ced0 100644 --- a/test/basic.jl +++ b/test/basic.jl @@ -453,12 +453,6 @@ end coords = place_atoms(1, b_right; min_dist=0.01u"nm") @test_throws ArgumentError System(atoms=atoms, coords=coords, boundary=b_wrong) - # Mis-matched energy units in interaction and system - coords = place_atoms(1, b_right; min_dist=0.01u"nm") - lj = LennardJones() - @test_throws ArgumentError System(atoms=atoms, coords=coords, boundary=b_right, - pairwise_inters=(lj,)) - # Mixed units or other invalid units bad_velo = [random_velocity(1.0u"g/mol",10u"K",Unitful.k*Unitful.Na) .* 2u"g"] @test_throws ArgumentError System(atoms=atoms, coords=coords, boundary=b_right, From 05a0a1702c8369fc0429cd93ce9ba44a57866675 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 28 Jun 2024 19:14:13 +0100 Subject: [PATCH 13/74] fix unit check --- src/force.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/force.jl b/src/force.jl index 3d65277b9..16184b468 100644 --- a/src/force.jl +++ b/src/force.jl @@ -325,7 +325,7 @@ function forces(sys::System{D, true, T}, neighbors; for inter in values(sys.general_inters) fs_gen = AtomsCalculators.forces(sys, inter; neighbors=neighbors, n_threads=n_threads) - check_force_units(unit(eltype(eltype(fs_gen))), sys.force_units) + check_force_units(zero(eltype(fs_gen)), sys.force_units) fs += ustrip_vec.(fs_gen) end From 60d373b0f90495d167b852688a1921f37d54e6c5 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Mon, 1 Jul 2024 14:27:14 +0100 Subject: [PATCH 14/74] rename coulombconst to coulomb_const --- src/interactions/coulomb.jl | 48 ++++++++++++++-------------- src/interactions/implicit_solvent.jl | 20 ++++++------ src/setup.jl | 6 ++-- test/gradients.jl | 2 +- 4 files changed, 38 insertions(+), 38 deletions(-) diff --git a/src/interactions/coulomb.jl b/src/interactions/coulomb.jl index db43fa694..00878a0f9 100644 --- a/src/interactions/coulomb.jl +++ b/src/interactions/coulomb.jl @@ -20,13 +20,13 @@ struct Coulomb{C, W, T} <: PairwiseInteraction coulomb_const::T end -const coulombconst = 138.93545764u"kJ * mol^-1 * nm" # 1 / 4πϵ0 +const coulomb_const = 138.93545764u"kJ * mol^-1 * nm" # 1 / 4πϵ0 function Coulomb(; cutoff=NoCutoff(), use_neighbors=false, weight_special=1, - coulomb_const=coulombconst) + coulomb_const=coulomb_const) return Coulomb{typeof(cutoff), typeof(weight_special), typeof(coulomb_const)}( cutoff, use_neighbors, weight_special, coulomb_const) end @@ -55,9 +55,9 @@ end args...) where C r2 = sum(abs2, dr) cutoff = inter.cutoff - coulomb_const = inter.coulomb_const + ke = inter.coulomb_const qi, qj = atom_i.charge, atom_j.charge - params = (coulomb_const, qi, qj) + params = (ke, qi, qj) f = force_divr_with_cutoff(inter, r2, params, cutoff, force_units) if special @@ -67,8 +67,8 @@ end end end -function force_divr(::Coulomb, r2, invr2, (coulomb_const, qi, qj)) - return (coulomb_const * qi * qj) / √(r2 ^ 3) +function force_divr(::Coulomb, r2, invr2, (ke, qi, qj)) + return (ke * qi * qj) / √(r2 ^ 3) end @inline function potential_energy(inter::Coulomb{C}, @@ -80,9 +80,9 @@ end args...) where C r2 = sum(abs2, dr) cutoff = inter.cutoff - coulomb_const = inter.coulomb_const + ke = inter.coulomb_const qi, qj = atom_i.charge, atom_j.charge - params = (coulomb_const, qi, qj) + params = (ke, qi, qj) pe = potential_with_cutoff(inter, r2, params, cutoff, energy_units) if special @@ -92,8 +92,8 @@ end end end -function potential(::Coulomb, r2, invr2, (coulomb_const, qi, qj)) - return (coulomb_const * qi * qj) * √invr2 +function potential(::Coulomb, r2, invr2, (ke, qi, qj)) + return (ke * qi * qj) * √invr2 end @doc raw""" @@ -127,7 +127,7 @@ function CoulombSoftCore(; use_neighbors=false, lorentz_mixing=true, weight_special=1, - coulomb_const=coulombconst) + coulomb_const=coulomb_const) σ6_fac = α * λ^p return CoulombSoftCore{typeof(cutoff), typeof(α), typeof(λ), typeof(p), typeof(σ6_fac), typeof(weight_special), typeof(coulomb_const)}( @@ -173,10 +173,10 @@ end args...) where C r2 = sum(abs2, dr) cutoff = inter.cutoff - coulomb_const = inter.coulomb_const + ke = inter.coulomb_const qi, qj = atom_i.charge, atom_j.charge σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) - params = (coulomb_const, qi, qj, σ, inter.σ6_fac) + params = (ke, qi, qj, σ, inter.σ6_fac) f = force_divr_with_cutoff(inter, r2, params, cutoff, force_units) if special @@ -186,11 +186,11 @@ end end end -function force_divr(::CoulombSoftCore, r2, invr2, (coulomb_const, qi, qj, σ, σ6_fac)) +function force_divr(::CoulombSoftCore, r2, invr2, (ke, qi, qj, σ, σ6_fac)) inv_rsc6 = inv(r2^3 + σ6_fac * σ^6) inv_rsc2 = cbrt(inv_rsc6) inv_rsc3 = sqrt(inv_rsc6) - ff = (coulomb_const * qi * qj) * inv_rsc2 * sqrt(r2)^5 * inv_rsc2 * inv_rsc3 + ff = (ke * qi * qj) * inv_rsc2 * sqrt(r2)^5 * inv_rsc2 * inv_rsc3 return ff * √invr2 end @@ -203,10 +203,10 @@ end args...) where C r2 = sum(abs2, dr) cutoff = inter.cutoff - coulomb_const = inter.coulomb_const + ke = inter.coulomb_const qi, qj = atom_i.charge, atom_j.charge σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) - params = (coulomb_const, qi, qj, σ, inter.σ6_fac) + params = (ke, qi, qj, σ, inter.σ6_fac) pe = potential_with_cutoff(inter, r2, params, cutoff, energy_units) if special @@ -216,9 +216,9 @@ end end end -function potential(::CoulombSoftCore, r2, invr2, (coulomb_const, qi, qj, σ, σ6_fac)) +function potential(::CoulombSoftCore, r2, invr2, (ke, qi, qj, σ, σ6_fac)) inv_rsc6 = inv(r2^3 + σ6_fac * σ^6) - return (coulomb_const * qi * qj) * √cbrt(inv_rsc6) + return (ke * qi * qj) * √cbrt(inv_rsc6) end """ @@ -243,7 +243,7 @@ function CoulombReactionField(; solvent_dielectric=crf_solvent_dielectric, use_neighbors=false, weight_special=1, - coulomb_const=coulombconst) + coulomb_const=coulomb_const) return CoulombReactionField{typeof(dist_cutoff), typeof(solvent_dielectric), typeof(weight_special), typeof(coulomb_const)}( dist_cutoff, solvent_dielectric, use_neighbors, weight_special, coulomb_const) @@ -283,7 +283,7 @@ end return ustrip.(zero(dr)) * force_units end - coulomb_const = inter.coulomb_const + ke = inter.coulomb_const qi, qj = atom_i.charge, atom_j.charge r = √r2 if special @@ -295,7 +295,7 @@ end (2 * inter.solvent_dielectric + 1)) end - f = (coulomb_const * qi * qj) * (inv(r) - 2 * krf * r2) * inv(r2) + f = (ke * qi * qj) * (inv(r) - 2 * krf * r2) * inv(r2) if special return f * dr * inter.weight_special @@ -316,7 +316,7 @@ end return ustrip(zero(dr[1])) * energy_units end - coulomb_const = inter.coulomb_const + ke = inter.coulomb_const qi, qj = atom_i.charge, atom_j.charge r = √r2 if special @@ -330,7 +330,7 @@ end (2 * inter.solvent_dielectric + 1)) end - pe = (coulomb_const * qi * qj) * (inv(r) + krf * r2 - crf) + pe = (ke * qi * qj) * (inv(r) + krf * r2 - crf) if special return pe * inter.weight_special diff --git a/src/interactions/implicit_solvent.jl b/src/interactions/implicit_solvent.jl index 5d90cdc4a..37f63c311 100644 --- a/src/interactions/implicit_solvent.jl +++ b/src/interactions/implicit_solvent.jl @@ -399,16 +399,16 @@ function ImplicitSolventOBC(atoms::AbstractArray{Atom{T, M, D, E}}, inds_j = hcat(1:n_atoms...) inds_i = permutedims(inds_j, (2, 1)) - coulomb_const = units ? coulombconst : ustrip(coulombconst) + coulomb_const_units = units ? coulomb_const : ustrip(coulomb_const) if !iszero_value(solute_dielectric) - factor_solute = -T(coulomb_const) / T(solute_dielectric) + factor_solute = -T(coulomb_const_units) / T(solute_dielectric) else - factor_solute = zero(T(coulomb_const)) + factor_solute = zero(T(coulomb_const_units)) end if !iszero_value(solvent_dielectric) - factor_solvent = T(coulomb_const) / T(solvent_dielectric) + factor_solvent = T(coulomb_const_units) / T(solvent_dielectric) else - factor_solvent = zero(T(coulomb_const)) + factor_solvent = zero(T(coulomb_const_units)) end if isa(atoms, CuArray) @@ -551,16 +551,16 @@ function ImplicitSolventGBN2(atoms::AbstractArray{Atom{T, M, D, E}}, table_m0 = ustrip.(table_m0_units) end - coulomb_const = units ? coulombconst : ustrip(coulombconst) + coulomb_const_units = units ? coulomb_const : ustrip(coulomb_const) if !iszero_value(solute_dielectric) - factor_solute = -T(coulomb_const) / T(solute_dielectric) + factor_solute = -T(coulomb_const_units) / T(solute_dielectric) else - factor_solute = zero(T(coulomb_const)) + factor_solute = zero(T(coulomb_const_units)) end if !iszero_value(solvent_dielectric) - factor_solvent = T(coulomb_const) / T(solvent_dielectric) + factor_solvent = T(coulomb_const_units) / T(solvent_dielectric) else - factor_solvent = zero(T(coulomb_const)) + factor_solvent = zero(T(coulomb_const_units)) end if isa(atoms, CuArray) diff --git a/src/setup.jl b/src/setup.jl index c1dd866cb..476fff683 100644 --- a/src/setup.jl +++ b/src/setup.jl @@ -799,14 +799,14 @@ function System(coord_file::AbstractString, solvent_dielectric=T(crf_solvent_dielectric), use_neighbors=true, weight_special=force_field.weight_14_coulomb, - coulomb_const=(units ? T(coulombconst) : T(ustrip(coulombconst))), + coulomb_const=(units ? T(coulomb_const) : T(ustrip(coulomb_const))), ) else crf = Coulomb( cutoff=DistanceCutoff(T(dist_cutoff)), use_neighbors=true, weight_special=force_field.weight_14_coulomb, - coulomb_const=(units ? T(coulombconst) : T(ustrip(coulombconst))), + coulomb_const=(units ? T(coulomb_const) : T(ustrip(coulomb_const))), ) end pairwise_inters = (lj, crf) @@ -1214,7 +1214,7 @@ function System(T::Type, solvent_dielectric=T(crf_solvent_dielectric), use_neighbors=true, weight_special=T(0.5), - coulomb_const=(units ? T(coulombconst) : T(ustrip(coulombconst))), + coulomb_const=(units ? T(coulomb_const) : T(ustrip(coulomb_const))), ) if isnothing(boundary) diff --git a/test/gradients.jl b/test/gradients.jl index 04cc1e70e..efd341204 100644 --- a/test/gradients.jl +++ b/test/gradients.jl @@ -59,7 +59,7 @@ end dist_cutoff=nb_cutoff, solvent_dielectric=f32 ? Float32(Molly.crf_solvent_dielectric) : Molly.crf_solvent_dielectric, use_neighbors=true, - coulomb_const=f32 ? Float32(ustrip(Molly.coulombconst)) : ustrip(Molly.coulombconst), + coulomb_const=f32 ? Float32(ustrip(Molly.coulomb_const)) : ustrip(Molly.coulomb_const), ) pairwise_inters = pis ? (lj, crf) : () bond_is = gpu ? CuArray(Int32.(collect(1:(n_atoms ÷ 2)))) : Int32.(collect(1:(n_atoms ÷ 2))) From fbc07568534a508f2c288cdbad510f088c1fefb6 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Mon, 1 Jul 2024 15:37:54 +0100 Subject: [PATCH 15/74] change specific interaction arguments --- docs/src/documentation.md | 31 +++- src/cuda.jl | 154 +++++++++++------- src/energy.jl | 48 +++--- src/force.jl | 56 ++++--- src/interactions/cosine_angle.jl | 4 +- src/interactions/fene_bond.jl | 4 +- src/interactions/harmonic_angle.jl | 4 +- src/interactions/harmonic_bond.jl | 4 +- .../harmonic_position_restraint.jl | 4 +- src/interactions/morse_bond.jl | 4 +- src/interactions/periodic_torsion.jl | 6 +- src/interactions/rb_torsion.jl | 5 +- src/spatial.jl | 9 +- 13 files changed, 205 insertions(+), 128 deletions(-) diff --git a/docs/src/documentation.md b/docs/src/documentation.md index 936ea214a..42ac5a359 100644 --- a/docs/src/documentation.md +++ b/docs/src/documentation.md @@ -669,7 +669,15 @@ Next, you need to define a method for the [`force`](@ref) function. The form of this will depend on whether the interaction involves 1, 2, 3 or 4 atoms. For example in the 2 atom case: ```julia -function Molly.force(inter::MySpecificInter, coords_i, coords_j, boundary) +function Molly.force(inter::MySpecificInter, + coord_i, + coord_j, + boundary, + atom_i, + atom_j, + force_units, + velocity_i, + velocity_j) dr = vector(coords_i, coords_j, boundary) # Replace this with your force calculation @@ -680,7 +688,8 @@ function Molly.force(inter::MySpecificInter, coords_i, coords_j, boundary) return SpecificForce2Atoms(-fdr, fdr) end ``` -The 3 atom case would define `Molly.force(inter::MySpecificInter, coords_i, coords_j, coords_k, boundary)` and return `SpecificForce3Atoms(f1, f2, f3)`. +Again, most of these arguments are rarely used and can be replaced with `args...`. +The 3 atom case would define `Molly.force(inter::MySpecificInter, coord_i, coord_j, coord_k, boundary, atom_i, atom_j, atom_k, force_units, velocity_i, velocity_j, velocity_k)` and return `SpecificForce3Atoms(f1, f2, f3)`. To use your custom interaction, add it to the specific interaction lists along with the atom indices: ```julia specific_inter_lists = ( @@ -695,7 +704,23 @@ For 3 atom interactions use [`InteractionList3Atoms`](@ref) and pass 3 sets of i If using the GPU, the inner list of indices and interactions should be moved to the GPU with `CuArray`. The number in the interaction list and the return type from [`force`](@ref) must match, e.g. [`InteractionList3Atoms`](@ref) must always return [`SpecificForce3Atoms`](@ref) from the corresponding [`force`](@ref) function. If some atoms are required in the interaction for force calculation but have no force applied to them by the interaction, give a zero force vector for those atoms. -Again a method for the [`potential_energy`](@ref) function with the same arguments can be defined. +Again a method for [`potential_energy`](@ref) with the same arguments, except the seventh argument is the energy units not the force units, can be defined: +```julia +function Molly.potential_energy(inter::MySpecificInter, + coord_i, + coord_j, + boundary, + atom_i, + atom_j, + energy_units, + velocity_i, + velocity_j) + # Example harmonic bond interaction + dr = vector(coord_i, coord_j, boundary) + r = norm(dr) + return (inter.k / 2) * (r - inter.r0) ^ 2 +end +``` ### General interactions diff --git a/src/cuda.jl b/src/cuda.jl index fdc71eee2..8a8c8ee96 100644 --- a/src/cuda.jl +++ b/src/cuda.jl @@ -30,12 +30,13 @@ function cuda_threads_blocks_specific(n_inters) end function pairwise_force_gpu(coords::AbstractArray{SVector{D, C}}, velocities, atoms, boundary, - pairwise_inters, nbs, force_units, ::Val{T}) where {D, C, T} + pairwise_inters, nbs, force_units, step_n, ::Val{T}) where {D, C, T} fs_mat = CUDA.zeros(T, D, length(atoms)) if typeof(nbs) == NoNeighborList kernel = @cuda launch=false pairwise_force_kernel_nonl!( - fs_mat, coords, velocities, atoms, boundary, pairwise_inters, Val(D), Val(force_units)) + fs_mat, coords, velocities, atoms, boundary, pairwise_inters, step_n, + Val(D), Val(force_units)) conf = launch_configuration(kernel.fun) threads_basic = parse(Int, get(ENV, "MOLLY_GPUNTHREADS_PAIRWISE", "512")) nthreads = min(length(atoms), threads_basic, conf.threads) @@ -47,13 +48,14 @@ function pairwise_force_gpu(coords::AbstractArray{SVector{D, C}}, velocities, at else n_threads_gpu, n_blocks = cuda_threads_blocks_pairwise(length(nbs)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks pairwise_force_kernel_nl!( - fs_mat, coords, velocities, atoms, boundary, pairwise_inters, nbs, Val(D), Val(force_units)) + fs_mat, coords, velocities, atoms, boundary, pairwise_inters, nbs, step_n, + Val(D), Val(force_units)) end return fs_mat end function pairwise_force_kernel_nl!(forces, coords_var, velocities_var, atoms_var, boundary, inters, - neighbors_var, ::Val{D}, ::Val{F}) where {D, F} + neighbors_var, step_n, ::Val{D}, ::Val{F}) where {D, F} coords = CUDA.Const(coords_var) velocities = CUDA.Const(velocities_var) atoms = CUDA.Const(atoms_var) @@ -64,7 +66,7 @@ function pairwise_force_kernel_nl!(forces, coords_var, velocities_var, atoms_var @inbounds if inter_i <= length(neighbors) i, j, special = neighbors[inter_i] f = sum_pairwise_forces(inters, atoms[i], atoms[j], Val(F), special, coords[i], coords[j], - boundary, velocities[i], velocities[j], 0) + boundary, velocities[i], velocities[j], step_n) for dim in 1:D fval = ustrip(f[dim]) Atomix.@atomic :monotonic forces[dim, i] += -fval @@ -114,8 +116,8 @@ That's why the calculations are done in the following order: h | 1 2 3 4 5 6 ``` =# -function pairwise_force_kernel_nonl!(forces::AbstractArray{T}, coords_var, velocities_var, atoms_var, boundary, inters, - ::Val{D}, ::Val{F}) where {T, D, F} +function pairwise_force_kernel_nonl!(forces::AbstractArray{T}, coords_var, velocities_var, + atoms_var, boundary, inters, step_n, ::Val{D}, ::Val{F}) where {T, D, F} coords = CUDA.Const(coords_var) velocities = CUDA.Const(velocities_var) atoms = CUDA.Const(atoms_var) @@ -142,7 +144,7 @@ function pairwise_force_kernel_nonl!(forces::AbstractArray{T}, coords_var, veloc if i != j atom_j, coord_j, vel_j = atoms[j], coords[j], velocities[j] f = sum_pairwise_forces(inters, atom_i, atom_j, Val(F), false, coord_i, coord_j, - boundary, vel_i, vel_j, 0) + boundary, vel_i, vel_j, step_n) for dim in 1:D forces_shmem[dim, tidx] += -ustrip(f[dim]) end @@ -167,7 +169,7 @@ function pairwise_force_kernel_nonl!(forces::AbstractArray{T}, coords_var, veloc sync_warp() atom_j = atoms[j] f = sum_pairwise_forces(inters, atom_i, atom_j, Val(F), false, coord_i, coord_j, - boundary, vel_i, vel_j, 0) + boundary, vel_i, vel_j, step_n) for dim in 1:D forces_shmem[dim, tidx] += -ustrip(f[dim]) end @@ -200,47 +202,50 @@ end end function specific_force_gpu(inter_list::InteractionList1Atoms, coords::AbstractArray{SVector{D, C}}, - boundary, force_units, ::Val{T}) where {D, C, T} + velocities, atoms, boundary, force_units, step_n, ::Val{T}) where {D, C, T} fs_mat = CUDA.zeros(T, D, length(coords)) n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_force_1_atoms_kernel!(fs_mat, - coords, boundary, inter_list.is, inter_list.inters, Val(D), Val(force_units)) + coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.inters, + Val(D), Val(force_units)) return fs_mat end function specific_force_gpu(inter_list::InteractionList2Atoms, coords::AbstractArray{SVector{D, C}}, - boundary, force_units, ::Val{T}) where {D, C, T} + velocities, atoms, boundary, force_units, step_n, ::Val{T}) where {D, C, T} fs_mat = CUDA.zeros(T, D, length(coords)) n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_force_2_atoms_kernel!(fs_mat, - coords, boundary, inter_list.is, inter_list.js, inter_list.inters, Val(D), - Val(force_units)) + coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.js, + inter_list.inters, Val(D), Val(force_units)) return fs_mat end function specific_force_gpu(inter_list::InteractionList3Atoms, coords::AbstractArray{SVector{D, C}}, - boundary, force_units, ::Val{T}) where {D, C, T} + velocities, atoms, boundary, force_units, step_n, ::Val{T}) where {D, C, T} fs_mat = CUDA.zeros(T, D, length(coords)) n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_force_3_atoms_kernel!(fs_mat, - coords, boundary, inter_list.is, inter_list.js, inter_list.ks, inter_list.inters, - Val(D), Val(force_units)) + coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.js, + inter_list.ks, inter_list.inters, Val(D), Val(force_units)) return fs_mat end function specific_force_gpu(inter_list::InteractionList4Atoms, coords::AbstractArray{SVector{D, C}}, - boundary, force_units, ::Val{T}) where {D, C, T} + velocities, atoms, boundary, force_units, step_n, ::Val{T}) where {D, C, T} fs_mat = CUDA.zeros(T, D, length(coords)) n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_force_4_atoms_kernel!(fs_mat, - coords, boundary, inter_list.is, inter_list.js, inter_list.ks, inter_list.ls, - inter_list.inters, Val(D), Val(force_units)) + coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.js, + inter_list.ks, inter_list.ls, inter_list.inters, Val(D), Val(force_units)) return fs_mat end -function specific_force_1_atoms_kernel!(forces, coords_var, boundary, is_var, - inters_var, ::Val{D}, ::Val{F}) where {D, F} +function specific_force_1_atoms_kernel!(forces, coords_var, velocities_var, atoms_var, boundary, + step_n, is_var, inters_var, ::Val{D}, ::Val{F}) where {D, F} coords = CUDA.Const(coords_var) + velocities = CUDA.Const(velocities_var) + atoms = CUDA.Const(atoms_var) is = CUDA.Const(is_var) inters = CUDA.Const(inters_var) @@ -248,7 +253,7 @@ function specific_force_1_atoms_kernel!(forces, coords_var, boundary, is_var, @inbounds if inter_i <= length(is) i = is[inter_i] - fs = force_gpu(inters[inter_i], coords[i], boundary) + fs = force_gpu(inters[inter_i], coords[i], boundary, atoms[i], F, velocities[i], step_n) if unit(fs.f1[1]) != F error("wrong force unit returned, was expecting $F") end @@ -259,9 +264,11 @@ function specific_force_1_atoms_kernel!(forces, coords_var, boundary, is_var, return nothing end -function specific_force_2_atoms_kernel!(forces, coords_var, boundary, is_var, js_var, - inters_var, ::Val{D}, ::Val{F}) where {D, F} +function specific_force_2_atoms_kernel!(forces, coords_var, velocities_var, atoms_var, boundary, + step_n, is_var, js_var, inters_var, ::Val{D}, ::Val{F}) where {D, F} coords = CUDA.Const(coords_var) + velocities = CUDA.Const(velocities_var) + atoms = CUDA.Const(atoms_var) is = CUDA.Const(is_var) js = CUDA.Const(js_var) inters = CUDA.Const(inters_var) @@ -270,7 +277,8 @@ function specific_force_2_atoms_kernel!(forces, coords_var, boundary, is_var, js @inbounds if inter_i <= length(is) i, j = is[inter_i], js[inter_i] - fs = force_gpu(inters[inter_i], coords[i], coords[j], boundary) + fs = force_gpu(inters[inter_i], coords[i], coords[j], boundary, atoms[i], atoms[j], F, + velocities[i], velocities[j], step_n) if unit(fs.f1[1]) != F || unit(fs.f2[1]) != F error("wrong force unit returned, was expecting $F") end @@ -282,9 +290,11 @@ function specific_force_2_atoms_kernel!(forces, coords_var, boundary, is_var, js return nothing end -function specific_force_3_atoms_kernel!(forces, coords_var, boundary, is_var, js_var, ks_var, - inters_var, ::Val{D}, ::Val{F}) where {D, F} +function specific_force_3_atoms_kernel!(forces, coords_var, velocities_var, atoms_var, boundary, + step_n, is_var, js_var, ks_var, inters_var, ::Val{D}, ::Val{F}) where {D, F} coords = CUDA.Const(coords_var) + velocities = CUDA.Const(velocities_var) + atoms = CUDA.Const(atoms_var) is = CUDA.Const(is_var) js = CUDA.Const(js_var) ks = CUDA.Const(ks_var) @@ -294,7 +304,8 @@ function specific_force_3_atoms_kernel!(forces, coords_var, boundary, is_var, js @inbounds if inter_i <= length(is) i, j, k = is[inter_i], js[inter_i], ks[inter_i] - fs = force_gpu(inters[inter_i], coords[i], coords[j], coords[k], boundary) + fs = force_gpu(inters[inter_i], coords[i], coords[j], coords[k], boundary, atoms[i], + atoms[j], atoms[k], F, velocities[i], velocities[j], velocities[k], step_n) if unit(fs.f1[1]) != F || unit(fs.f2[1]) != F || unit(fs.f3[1]) != F error("wrong force unit returned, was expecting $F") end @@ -307,9 +318,12 @@ function specific_force_3_atoms_kernel!(forces, coords_var, boundary, is_var, js return nothing end -function specific_force_4_atoms_kernel!(forces, coords_var, boundary, is_var, js_var, ks_var, ls_var, - inters_var, ::Val{D}, ::Val{F}) where {D, F} +function specific_force_4_atoms_kernel!(forces, coords_var, velocities_var, atoms_var, boundary, + step_n, is_var, js_var, ks_var, ls_var, inters_var, + ::Val{D}, ::Val{F}) where {D, F} coords = CUDA.Const(coords_var) + velocities = CUDA.Const(velocities_var) + atoms = CUDA.Const(atoms_var) is = CUDA.Const(is_var) js = CUDA.Const(js_var) ks = CUDA.Const(ks_var) @@ -320,7 +334,9 @@ function specific_force_4_atoms_kernel!(forces, coords_var, boundary, is_var, js @inbounds if inter_i <= length(is) i, j, k, l = is[inter_i], js[inter_i], ks[inter_i], ls[inter_i] - fs = force_gpu(inters[inter_i], coords[i], coords[j], coords[k], coords[l], boundary) + fs = force_gpu(inters[inter_i], coords[i], coords[j], coords[k], coords[l], boundary, + atoms[i], atoms[j], atoms[k], atoms[l], F, velocities[i], velocities[j], + velocities[k], velocities[l], step_n) if unit(fs.f1[1]) != F || unit(fs.f2[1]) != F || unit(fs.f3[1]) != F || unit(fs.f4[1]) != F error("wrong force unit returned, was expecting $F") end @@ -335,16 +351,17 @@ function specific_force_4_atoms_kernel!(forces, coords_var, boundary, is_var, js end function pairwise_pe_gpu(coords::AbstractArray{SVector{D, C}}, velocities, atoms, boundary, - pairwise_inters, nbs, energy_units, ::Val{T}) where {D, C, T} + pairwise_inters, nbs, energy_units, step_n, ::Val{T}) where {D, C, T} pe_vec = CUDA.zeros(T, 1) n_threads_gpu, n_blocks = cuda_threads_blocks_pairwise(length(nbs)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks pairwise_pe_kernel!( - pe_vec, coords, velocities, atoms, boundary, pairwise_inters, nbs, Val(energy_units)) + pe_vec, coords, velocities, atoms, boundary, pairwise_inters, nbs, + step_n, Val(energy_units)) return pe_vec end -function pairwise_pe_kernel!(energy, coords_var, velocities_var, atoms_var, boundary, inters, neighbors_var, - ::Val{E}) where E +function pairwise_pe_kernel!(energy, coords_var, velocities_var, atoms_var, boundary, inters, + neighbors_var, step_n, ::Val{E}) where E coords = CUDA.Const(coords_var) velocities = CUDA.Const(velocities_var) atoms = CUDA.Const(atoms_var) @@ -357,10 +374,10 @@ function pairwise_pe_kernel!(energy, coords_var, velocities_var, atoms_var, boun coord_i, coord_j, vel_i, vel_j = coords[i], coords[j], velocities[i], velocities[j] dr = vector(coord_i, coord_j, boundary) pe = potential_energy_gpu(inters[1], dr, atoms[i], atoms[j], E, special, coord_i, coord_j, - boundary, vel_i, vel_j, 0) + boundary, vel_i, vel_j, step_n) for inter in inters[2:end] pe += potential_energy_gpu(inter, dr, atoms[i], atoms[j], E, special, coord_i, coord_j, - boundary, vel_i, vel_j, 0) + boundary, vel_i, vel_j, step_n) end if unit(pe) != E error("wrong energy unit returned, was expecting $E but got $(unit(pe))") @@ -371,46 +388,50 @@ function pairwise_pe_kernel!(energy, coords_var, velocities_var, atoms_var, boun end function specific_pe_gpu(inter_list::InteractionList1Atoms, coords::AbstractArray{SVector{D, C}}, - boundary, energy_units, ::Val{T}) where {D, C, T} + velocities, atoms, boundary, energy_units, step_n, ::Val{T}) where {D, C, T} pe_vec = CUDA.zeros(T, 1) n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_1_atoms_kernel!(pe_vec, - coords, boundary, inter_list.is, inter_list.inters, Val(energy_units)) + coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.inters, + Val(energy_units)) return pe_vec end function specific_pe_gpu(inter_list::InteractionList2Atoms, coords::AbstractArray{SVector{D, C}}, - boundary, energy_units, ::Val{T}) where {D, C, T} + velocities, atoms, boundary, energy_units, step_n, ::Val{T}) where {D, C, T} pe_vec = CUDA.zeros(T, 1) n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_2_atoms_kernel!(pe_vec, - coords, boundary, inter_list.is, inter_list.js, inter_list.inters, Val(energy_units)) + coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.js, + inter_list.inters, Val(energy_units)) return pe_vec end function specific_pe_gpu(inter_list::InteractionList3Atoms, coords::AbstractArray{SVector{D, C}}, - boundary, energy_units, ::Val{T}) where {D, C, T} + velocities, atoms, boundary, energy_units, step_n, ::Val{T}) where {D, C, T} pe_vec = CUDA.zeros(T, 1) n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_3_atoms_kernel!(pe_vec, - coords, boundary, inter_list.is, inter_list.js, inter_list.ks, inter_list.inters, - Val(energy_units)) + coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.js, + inter_list.ks, inter_list.inters, Val(energy_units)) return pe_vec end function specific_pe_gpu(inter_list::InteractionList4Atoms, coords::AbstractArray{SVector{D, C}}, - boundary, energy_units, ::Val{T}) where {D, C, T} + velocities, atoms, boundary, energy_units, step_n, ::Val{T}) where {D, C, T} pe_vec = CUDA.zeros(T, 1) n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_4_atoms_kernel!(pe_vec, - coords, boundary, inter_list.is, inter_list.js, inter_list.ks, inter_list.ls, - inter_list.inters, Val(energy_units)) + coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.js, + inter_list.ks, inter_list.ls, inter_list.inters, Val(energy_units)) return pe_vec end -function specific_pe_1_atoms_kernel!(energy, coords_var, boundary, is_var, - inters_var, ::Val{E}) where E +function specific_pe_1_atoms_kernel!(energy, coords_var, velocities_var, atoms_var, boundary, + step_n, is_var, inters_var, ::Val{E}) where E coords = CUDA.Const(coords_var) + velocities = CUDA.Const(velocities_var) + atoms = CUDA.Const(atoms_var) is = CUDA.Const(is_var) inters = CUDA.Const(inters_var) @@ -418,7 +439,8 @@ function specific_pe_1_atoms_kernel!(energy, coords_var, boundary, is_var, @inbounds if inter_i <= length(is) i = is[inter_i] - pe = potential_energy_gpu(inters[inter_i], coords[i], boundary) + pe = potential_energy_gpu(inters[inter_i], coords[i], boundary, atoms[i], E, + velocities[i], step_n) if unit(pe) != E error("wrong energy unit returned, was expecting $E but got $(unit(pe))") end @@ -427,9 +449,11 @@ function specific_pe_1_atoms_kernel!(energy, coords_var, boundary, is_var, return nothing end -function specific_pe_2_atoms_kernel!(energy, coords_var, boundary, is_var, js_var, - inters_var, ::Val{E}) where E +function specific_pe_2_atoms_kernel!(energy, coords_var, velocities_var, atoms_var, boundary, + step_n, is_var, js_var, inters_var, ::Val{E}) where E coords = CUDA.Const(coords_var) + velocities = CUDA.Const(velocities_var) + atoms = CUDA.Const(atoms_var) is = CUDA.Const(is_var) js = CUDA.Const(js_var) inters = CUDA.Const(inters_var) @@ -438,7 +462,8 @@ function specific_pe_2_atoms_kernel!(energy, coords_var, boundary, is_var, js_va @inbounds if inter_i <= length(is) i, j = is[inter_i], js[inter_i] - pe = potential_energy_gpu(inters[inter_i], coords[i], coords[j], boundary) + pe = potential_energy_gpu(inters[inter_i], coords[i], coords[j], boundary, atoms[i], + atoms[j], E, velocities[i], velocities[j], step_n) if unit(pe) != E error("wrong energy unit returned, was expecting $E but got $(unit(pe))") end @@ -447,9 +472,11 @@ function specific_pe_2_atoms_kernel!(energy, coords_var, boundary, is_var, js_va return nothing end -function specific_pe_3_atoms_kernel!(energy, coords_var, boundary, is_var, js_var, ks_var, - inters_var, ::Val{E}) where E +function specific_pe_3_atoms_kernel!(energy, coords_var, velocities_var, atoms_var, boundary, + step_n, is_var, js_var, ks_var, inters_var, ::Val{E}) where E coords = CUDA.Const(coords_var) + velocities = CUDA.Const(velocities_var) + atoms = CUDA.Const(atoms_var) is = CUDA.Const(is_var) js = CUDA.Const(js_var) ks = CUDA.Const(ks_var) @@ -459,7 +486,9 @@ function specific_pe_3_atoms_kernel!(energy, coords_var, boundary, is_var, js_va @inbounds if inter_i <= length(is) i, j, k = is[inter_i], js[inter_i], ks[inter_i] - pe = potential_energy_gpu(inters[inter_i], coords[i], coords[j], coords[k], boundary) + pe = potential_energy_gpu(inters[inter_i], coords[i], coords[j], coords[k], boundary, + atoms[i], atoms[j], atoms[k], E, velocities[i], velocities[j], + velocities[k], step_n) if unit(pe) != E error("wrong energy unit returned, was expecting $E but got $(unit(pe))") end @@ -468,9 +497,11 @@ function specific_pe_3_atoms_kernel!(energy, coords_var, boundary, is_var, js_va return nothing end -function specific_pe_4_atoms_kernel!(energy, coords_var, boundary, is_var, js_var, ks_var, ls_var, - inters_var, ::Val{E}) where E +function specific_pe_4_atoms_kernel!(energy, coords_var, velocities_var, atoms_var, boundary, + step_n, is_var, js_var, ks_var, ls_var, inters_var, ::Val{E}) where E coords = CUDA.Const(coords_var) + velocities = CUDA.Const(velocities_var) + atoms = CUDA.Const(atoms_var) is = CUDA.Const(is_var) js = CUDA.Const(js_var) ks = CUDA.Const(ks_var) @@ -481,7 +512,10 @@ function specific_pe_4_atoms_kernel!(energy, coords_var, boundary, is_var, js_va @inbounds if inter_i <= length(is) i, j, k, l = is[inter_i], js[inter_i], ks[inter_i], ls[inter_i] - pe = potential_energy_gpu(inters[inter_i], coords[i], coords[j], coords[k], coords[l], boundary) + pe = potential_energy_gpu(inters[inter_i], coords[i], coords[j], coords[k], coords[l], + boundary, atoms[i], atoms[j], atoms[k], atoms[l], E, + velocities[i], velocities[j], velocities[k], velocities[l], + step_n) if unit(pe) != E error("wrong energy unit returned, was expecting $E but got $(unit(pe))") end diff --git a/src/energy.jl b/src/energy.jl index bed391a0f..987d62334 100644 --- a/src/energy.jl +++ b/src/energy.jl @@ -72,10 +72,10 @@ end # Allow GPU-specific potential energy functions to be defined if required potential_energy_gpu(inter::PairwiseInteraction, dr, ai, aj, eu, sp, ci, cj, bnd, vi, vj, sn) = potential_energy(inter, dr, ai, aj, eu, sp, ci, cj, bnd, vi, vj, sn) -potential_energy_gpu(inter::SpecificInteraction, ci, bnd) = potential_energy(inter, ci, bnd) -potential_energy_gpu(inter::SpecificInteraction, ci, cj, bnd) = potential_energy(inter, ci, cj, bnd) -potential_energy_gpu(inter::SpecificInteraction, ci, cj, ck, bnd) = potential_energy(inter, ci, cj, ck, bnd) -potential_energy_gpu(inter::SpecificInteraction, ci, cj, ck, cl, bnd) = potential_energy(inter, ci, cj, ck, cl, bnd) +potential_energy_gpu(inter::SpecificInteraction, ci, bnd, ai, eu, vi, sn) = potential_energy(inter, ci, bnd, ai, eu, vi, sn) +potential_energy_gpu(inter::SpecificInteraction, ci, cj, bnd, ai, aj, eu, vi, vj, sn) = potential_energy(inter, ci, cj, bnd, ai, aj, eu, vi, vj, sn) +potential_energy_gpu(inter::SpecificInteraction, ci, cj, ck, bnd, ai, aj, ak, eu, vi, vj, vk, sn) = potential_energy(inter, ci, cj, ck, bnd, ai, aj, ak, eu, vi, vj, vk, sn) +potential_energy_gpu(inter::SpecificInteraction, ci, cj, ck, cl, bnd, ai, aj, ak, al, eu, vi, vj, vk, vl, sn) = potential_energy(inter, ci, cj, ck, cl, bnd, ai, aj, ak, al, eu, vi, vj, vk, vl, sn) function potential_energy(sys::System{D, false}, neighbors; n_threads::Integer=Threads.nthreads()) where D @@ -112,6 +112,7 @@ end function potential_energy_pair_spec!(pe_vec, coords, velocities, atoms, pairwise_inters_nonl, pairwise_inters_nl, sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, boundary, energy_units, neighbors, n_threads, ::Val{T}) where T + step_n = 0 pe_sum = zero(T) @inbounds if n_threads > 1 @@ -124,10 +125,10 @@ function potential_energy_pair_spec!(pe_vec, coords, velocities, atoms, pairwise for j in (i + 1):n_atoms dr = vector(coords[i], coords[j], boundary) pe = potential_energy(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], energy_units, false, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) for inter in pairwise_inters_nonl[2:end] pe += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, false, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) end check_energy_units(pe, energy_units) pe_sum_chunks[chunk_i] += ustrip(pe) @@ -145,10 +146,10 @@ function potential_energy_pair_spec!(pe_vec, coords, velocities, atoms, pairwise i, j, special = neighbors[ni] dr = vector(coords[i], coords[j], boundary) pe = potential_energy(pairwise_inters_nl[1], dr, atoms[i], atoms[j], energy_units, special, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) for inter in pairwise_inters_nl[2:end] pe += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, special, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) end check_energy_units(pe, energy_units) pe_sum_chunks[chunk_i] += ustrip(pe) @@ -164,10 +165,10 @@ function potential_energy_pair_spec!(pe_vec, coords, velocities, atoms, pairwise for j in (i + 1):n_atoms dr = vector(coords[i], coords[j], boundary) pe = potential_energy(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], energy_units, false, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) for inter in pairwise_inters_nonl[2:end] pe += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, false, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) end check_energy_units(pe, energy_units) pe_sum += ustrip(pe) @@ -183,10 +184,10 @@ function potential_energy_pair_spec!(pe_vec, coords, velocities, atoms, pairwise i, j, special = neighbors[ni] dr = vector(coords[i], coords[j], boundary) pe = potential_energy(pairwise_inters_nl[1], dr, atoms[i], atoms[j], energy_units, special, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) for inter in pairwise_inters_nl[2:end] pe += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, special, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) end check_energy_units(pe, energy_units) pe_sum += ustrip(pe) @@ -196,7 +197,8 @@ function potential_energy_pair_spec!(pe_vec, coords, velocities, atoms, pairwise @inbounds for inter_list in sils_1_atoms for (i, inter) in zip(inter_list.is, inter_list.inters) - pe = potential_energy(inter, coords[i], boundary) + pe = potential_energy(inter, coords[i], boundary, atoms[i], energy_units, + velocities[i], step_n) check_energy_units(pe, energy_units) pe_sum += ustrip(pe) end @@ -204,7 +206,8 @@ function potential_energy_pair_spec!(pe_vec, coords, velocities, atoms, pairwise @inbounds for inter_list in sils_2_atoms for (i, j, inter) in zip(inter_list.is, inter_list.js, inter_list.inters) - pe = potential_energy(inter, coords[i], coords[j], boundary) + pe = potential_energy(inter, coords[i], coords[j], boundary, atoms[i], atoms[j], + energy_units, velocities[i], velocities[j], step_n) check_energy_units(pe, energy_units) pe_sum += ustrip(pe) end @@ -212,7 +215,9 @@ function potential_energy_pair_spec!(pe_vec, coords, velocities, atoms, pairwise @inbounds for inter_list in sils_3_atoms for (i, j, k, inter) in zip(inter_list.is, inter_list.js, inter_list.ks, inter_list.inters) - pe = potential_energy(inter, coords[i], coords[j], coords[k], boundary) + pe = potential_energy(inter, coords[i], coords[j], coords[k], boundary, atoms[i], + atoms[j], atoms[k], energy_units, velocities[i], velocities[j], + velocities[k], step_n) check_energy_units(pe, energy_units) pe_sum += ustrip(pe) end @@ -221,7 +226,10 @@ function potential_energy_pair_spec!(pe_vec, coords, velocities, atoms, pairwise @inbounds for inter_list in sils_4_atoms for (i, j, k, l, inter) in zip(inter_list.is, inter_list.js, inter_list.ks, inter_list.ls, inter_list.inters) - pe = potential_energy(inter, coords[i], coords[j], coords[k], coords[l], boundary) + pe = potential_energy(inter, coords[i], coords[j], coords[k], coords[l], boundary, + atoms[i], atoms[j], atoms[k], atoms[l], energy_units, + velocities[i], velocities[j], velocities[k], velocities[l], + step_n) check_energy_units(pe, energy_units) pe_sum += ustrip(pe) end @@ -233,6 +241,7 @@ end function potential_energy(sys::System{D, true, T}, neighbors; n_threads::Integer=Threads.nthreads()) where {D, T} + step_n = 0 n_atoms = length(sys) val_ft = Val(T) pe_vec = CUDA.zeros(T, 1) @@ -241,7 +250,7 @@ function potential_energy(sys::System{D, true, T}, neighbors; if length(pairwise_inters_nonl) > 0 nbs = NoNeighborList(n_atoms) pe_vec += pairwise_pe_gpu(sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nonl, - nbs, sys.energy_units, val_ft) + nbs, sys.energy_units, step_n, val_ft) end pairwise_inters_nl = filter(use_neighbors, values(sys.pairwise_inters)) @@ -252,12 +261,13 @@ function potential_energy(sys::System{D, true, T}, neighbors; if length(neighbors) > 0 nbs = @view neighbors.list[1:neighbors.n] pe_vec += pairwise_pe_gpu(sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nl, - nbs, sys.energy_units, val_ft) + nbs, sys.energy_units, step_n, val_ft) end end for inter_list in values(sys.specific_inter_lists) - pe_vec += specific_pe_gpu(inter_list, sys.coords, sys.boundary, sys.energy_units, val_ft) + pe_vec += specific_pe_gpu(inter_list, sys.coords, sys.velocities, sys.atoms, sys.boundary, + sys.energy_units, step_n, val_ft) end pe = Array(pe_vec)[1] diff --git a/src/force.jl b/src/force.jl index 16184b468..5d79feff7 100644 --- a/src/force.jl +++ b/src/force.jl @@ -44,10 +44,10 @@ function force end # Allow GPU-specific force functions to be defined if required force_gpu(inter::PairwiseInteraction, dr, ai, aj, fu, sp, ci, cj, bnd, vi, vj, sn) = force(inter, dr, ai, aj, fu, sp, ci, cj, bnd, vi, vj, sn) -force_gpu(inter::SpecificInteraction, ci, bnd) = force(inter, ci, bnd) -force_gpu(inter::SpecificInteraction, ci, cj, bnd) = force(inter, ci, cj, bnd) -force_gpu(inter::SpecificInteraction, ci, cj, ck, bnd) = force(inter, ci, cj, ck, bnd) -force_gpu(inter::SpecificInteraction, ci, cj, ck, cl, bnd) = force(inter, ci, cj, ck, cl, bnd) +force_gpu(inter::SpecificInteraction, ci, bnd, ai, fu, vi, sn) = force(inter, ci, bnd, ai, fu, vi, sn) +force_gpu(inter::SpecificInteraction, ci, cj, bnd, ai, aj, fu, vi, vj, sn) = force(inter, ci, cj, bnd, ai, aj, fu, vi, vj, sn) +force_gpu(inter::SpecificInteraction, ci, cj, ck, bnd, ai, aj, ak, fu, vi, vj, vk, sn) = force(inter, ci, cj, ck, bnd, ai, aj, ak, fu, vi, vj, vk, sn) +force_gpu(inter::SpecificInteraction, ci, cj, ck, cl, bnd, ai, aj, ak, al, fu, vi, vj, vk, vl, sn) = force(inter, ci, cj, ck, cl, bnd, ai, aj, ak, al, fu, vi, vj, vk, vl, sn) """ SpecificForce1Atoms(f1) @@ -157,6 +157,7 @@ end function forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, pairwise_inters_nl, sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, boundary, force_units, neighbors, n_threads) + step_n = 0 n_atoms = length(coords) @inbounds if n_threads > 1 fs_chunks = [zero(fs) for _ in 1:n_threads] @@ -166,11 +167,12 @@ function forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, for i in chunk_i:n_threads:n_atoms for j in (i + 1):n_atoms dr = vector(coords[i], coords[j], boundary) - f = force(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], force_units, false, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + f = force(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], force_units, + false, coords[i], coords[j], boundary, velocities[i], + velocities[j], step_n) for inter in pairwise_inters_nonl[2:end] - f += force(inter, dr, atoms[i], atoms[j], force_units, false, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + f += force(inter, dr, atoms[i], atoms[j], force_units, false, coords[i], + coords[j], boundary, velocities[i], velocities[j], step_n) end check_force_units(f, force_units) f_ustrip = ustrip.(f) @@ -190,10 +192,10 @@ function forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, i, j, special = neighbors[ni] dr = vector(coords[i], coords[j], boundary) f = force(pairwise_inters_nl[1], dr, atoms[i], atoms[j], force_units, special, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) for inter in pairwise_inters_nl[2:end] - f += force(inter, dr, atoms[i], atoms[j], force_units, special, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + f += force(inter, dr, atoms[i], atoms[j], force_units, special, coords[i], + coords[j], boundary, velocities[i], velocities[j], step_n) end check_force_units(f, force_units) f_ustrip = ustrip.(f) @@ -210,10 +212,10 @@ function forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, for j in (i + 1):n_atoms dr = vector(coords[i], coords[j], boundary) f = force(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], force_units, false, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) for inter in pairwise_inters_nonl[2:end] - f += force(inter, dr, atoms[i], atoms[j], force_units, false, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + f += force(inter, dr, atoms[i], atoms[j], force_units, false, coords[i], + coords[j], boundary, velocities[i], velocities[j], step_n) end check_force_units(f, force_units) f_ustrip = ustrip.(f) @@ -231,10 +233,10 @@ function forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, i, j, special = neighbors[ni] dr = vector(coords[i], coords[j], boundary) f = force(pairwise_inters_nl[1], dr, atoms[i], atoms[j], force_units, special, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) for inter in pairwise_inters_nl[2:end] - f += force(inter, dr, atoms[i], atoms[j], force_units, special, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + f += force(inter, dr, atoms[i], atoms[j], force_units, special, coords[i], + coords[j], boundary, velocities[i], velocities[j], step_n) end check_force_units(f, force_units) f_ustrip = ustrip.(f) @@ -246,7 +248,7 @@ function forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, @inbounds for inter_list in sils_1_atoms for (i, inter) in zip(inter_list.is, inter_list.inters) - sf = force(inter, coords[i], boundary) + sf = force(inter, coords[i], boundary, atoms[i], force_units, velocities[i], step_n) check_force_units(sf.f1, force_units) fs[i] += ustrip.(sf.f1) end @@ -254,7 +256,8 @@ function forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, @inbounds for inter_list in sils_2_atoms for (i, j, inter) in zip(inter_list.is, inter_list.js, inter_list.inters) - sf = force(inter, coords[i], coords[j], boundary) + sf = force(inter, coords[i], coords[j], boundary, atoms[i], atoms[j], force_units, + velocities[i], velocities[j], step_n) check_force_units(sf.f1, force_units) check_force_units(sf.f2, force_units) fs[i] += ustrip.(sf.f1) @@ -264,7 +267,8 @@ function forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, @inbounds for inter_list in sils_3_atoms for (i, j, k, inter) in zip(inter_list.is, inter_list.js, inter_list.ks, inter_list.inters) - sf = force(inter, coords[i], coords[j], coords[k], boundary) + sf = force(inter, coords[i], coords[j], coords[k], boundary, atoms[i], atoms[j], + atoms[k], force_units, velocities[i], velocities[j], velocities[k], step_n) check_force_units(sf.f1, force_units) check_force_units(sf.f2, force_units) check_force_units(sf.f3, force_units) @@ -277,7 +281,9 @@ function forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, @inbounds for inter_list in sils_4_atoms for (i, j, k, l, inter) in zip(inter_list.is, inter_list.js, inter_list.ks, inter_list.ls, inter_list.inters) - sf = force(inter, coords[i], coords[j], coords[k], coords[l], boundary) + sf = force(inter, coords[i], coords[j], coords[k], coords[l], boundary, atoms[i], + atoms[j], atoms[k], atoms[l], force_units, velocities[i], velocities[j], + velocities[k], velocities[l], step_n) check_force_units(sf.f1, force_units) check_force_units(sf.f2, force_units) check_force_units(sf.f3, force_units) @@ -294,6 +300,7 @@ end function forces(sys::System{D, true, T}, neighbors; n_threads::Integer=Threads.nthreads()) where {D, T} + step_n = 0 n_atoms = length(sys) val_ft = Val(T) fs_mat = CUDA.zeros(T, D, n_atoms) @@ -302,7 +309,7 @@ function forces(sys::System{D, true, T}, neighbors; if length(pairwise_inters_nonl) > 0 nbs = NoNeighborList(n_atoms) fs_mat += pairwise_force_gpu(sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nonl, - nbs, sys.force_units, val_ft) + nbs, sys.force_units, step_n, val_ft) end pairwise_inters_nl = filter(use_neighbors, values(sys.pairwise_inters)) @@ -313,12 +320,13 @@ function forces(sys::System{D, true, T}, neighbors; if length(neighbors) > 0 nbs = @view neighbors.list[1:neighbors.n] fs_mat += pairwise_force_gpu(sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nl, - nbs, sys.force_units, val_ft) + nbs, sys.force_units, step_n, val_ft) end end for inter_list in values(sys.specific_inter_lists) - fs_mat += specific_force_gpu(inter_list, sys.coords, sys.boundary, sys.force_units, val_ft) + fs_mat += specific_force_gpu(inter_list, sys.coords, sys.velocities, sys.atoms, + sys.boundary, sys.force_units, step_n, val_ft) end fs = reinterpret(SVector{D, T}, vec(fs_mat)) diff --git a/src/interactions/cosine_angle.jl b/src/interactions/cosine_angle.jl index a49c36e3e..b159ecd64 100644 --- a/src/interactions/cosine_angle.jl +++ b/src/interactions/cosine_angle.jl @@ -18,7 +18,7 @@ end CosineAngle(; k, θ0) = CosineAngle{typeof(k), typeof(θ0)}(k, θ0) -@inline function force(a::CosineAngle, coords_i, coords_j, coords_k, boundary) +@inline function force(a::CosineAngle, coords_i, coords_j, coords_k, boundary, args...) # In 2D we use then eliminate the cross product ba = vector_pad3D(coords_j, coords_i, boundary) bc = vector_pad3D(coords_j, coords_k, boundary) @@ -38,7 +38,7 @@ CosineAngle(; k, θ0) = CosineAngle{typeof(k), typeof(θ0)}(k, θ0) end @inline function potential_energy(a::CosineAngle, coords_i, coords_j, - coords_k, boundary) + coords_k, boundary, args...) θ = bond_angle(coords_i, coords_j, coords_k, boundary) return a.k * (1 + cos(θ - a.θ0)) end diff --git a/src/interactions/fene_bond.jl b/src/interactions/fene_bond.jl index 3846b377b..8020cd976 100644 --- a/src/interactions/fene_bond.jl +++ b/src/interactions/fene_bond.jl @@ -28,7 +28,7 @@ end FENEBond(; k, r0, σ, ϵ) = FENEBond{typeof(k), typeof(r0), typeof(ϵ)}(k, r0, σ, ϵ) -@inline function force(b::FENEBond, coord_i, coord_j, boundary) +@inline function force(b::FENEBond, coord_i, coord_j, boundary, args...) ab = vector(coord_i, coord_j, boundary) r = norm(ab) r2 = r^2 @@ -47,7 +47,7 @@ FENEBond(; k, r0, σ, ϵ) = FENEBond{typeof(k), typeof(r0), typeof(ϵ)}(k, r0, return SpecificForce2Atoms(-f, f) end -@inline function potential_energy(b::FENEBond, coord_i, coord_j, boundary) +@inline function potential_energy(b::FENEBond, coord_i, coord_j, boundary, args...) dr = vector(coord_i, coord_j, boundary) r = norm(dr) r2 = r^2 diff --git a/src/interactions/harmonic_angle.jl b/src/interactions/harmonic_angle.jl index fd3cc940b..318c81e56 100644 --- a/src/interactions/harmonic_angle.jl +++ b/src/interactions/harmonic_angle.jl @@ -23,7 +23,7 @@ Base.zero(::HarmonicAngle{K, D}) where {K, D} = HarmonicAngle(k=zero(K), θ0=zer Base.:+(a1::HarmonicAngle, a2::HarmonicAngle) = HarmonicAngle(k=(a1.k + a2.k), θ0=(a1.θ0 + a2.θ0)) -@inline function force(a::HarmonicAngle, coords_i, coords_j, coords_k, boundary) +@inline function force(a::HarmonicAngle, coords_i, coords_j, coords_k, boundary, args...) # In 2D we use then eliminate the cross product ba = vector_pad3D(coords_j, coords_i, boundary) bc = vector_pad3D(coords_j, coords_k, boundary) @@ -42,7 +42,7 @@ Base.:+(a1::HarmonicAngle, a2::HarmonicAngle) = HarmonicAngle(k=(a1.k + a2.k), end @inline function potential_energy(a::HarmonicAngle, coords_i, coords_j, - coords_k, boundary) + coords_k, boundary, args...) θ = bond_angle(coords_i, coords_j, coords_k, boundary) return (a.k / 2) * (θ - a.θ0) ^ 2 end diff --git a/src/interactions/harmonic_bond.jl b/src/interactions/harmonic_bond.jl index b7f837272..19b5502be 100644 --- a/src/interactions/harmonic_bond.jl +++ b/src/interactions/harmonic_bond.jl @@ -21,14 +21,14 @@ Base.zero(::HarmonicBond{K, D}) where {K, D} = HarmonicBond(k=zero(K), r0=zero(D Base.:+(b1::HarmonicBond, b2::HarmonicBond) = HarmonicBond(k=(b1.k + b2.k), r0=(b1.r0 + b2.r0)) -@inline function force(b::HarmonicBond, coord_i, coord_j, boundary) +@inline function force(b::HarmonicBond, coord_i, coord_j, boundary, args...) ab = vector(coord_i, coord_j, boundary) c = b.k * (norm(ab) - b.r0) f = c * normalize(ab) return SpecificForce2Atoms(f, -f) end -@inline function potential_energy(b::HarmonicBond, coord_i, coord_j, boundary) +@inline function potential_energy(b::HarmonicBond, coord_i, coord_j, boundary, args...) dr = vector(coord_i, coord_j, boundary) r = norm(dr) return (b.k / 2) * (r - b.r0) ^ 2 diff --git a/src/interactions/harmonic_position_restraint.jl b/src/interactions/harmonic_position_restraint.jl index 862f69f06..d23f493cb 100644 --- a/src/interactions/harmonic_position_restraint.jl +++ b/src/interactions/harmonic_position_restraint.jl @@ -17,7 +17,7 @@ end HarmonicPositionRestraint(; k, x0) = HarmonicPositionRestraint{typeof(k), typeof(x0)}(k, x0) -@inline function force(pr::HarmonicPositionRestraint, coord_i, boundary) +@inline function force(pr::HarmonicPositionRestraint, coord_i, boundary, args...) ab = vector(coord_i, pr.x0, boundary) c = pr.k * norm(ab) if iszero_value(c) @@ -28,7 +28,7 @@ HarmonicPositionRestraint(; k, x0) = HarmonicPositionRestraint{typeof(k), typeof return SpecificForce1Atoms(f) end -@inline function potential_energy(pr::HarmonicPositionRestraint, coord_i, boundary) +@inline function potential_energy(pr::HarmonicPositionRestraint, coord_i, boundary, args...) dr = vector(coord_i, pr.x0, boundary) return (pr.k / 2) * dot(dr, dr) end diff --git a/src/interactions/morse_bond.jl b/src/interactions/morse_bond.jl index 56eab3d91..9d6f32422 100644 --- a/src/interactions/morse_bond.jl +++ b/src/interactions/morse_bond.jl @@ -18,7 +18,7 @@ end MorseBond(; D, a, r0) = MorseBond{typeof(D), typeof(a), typeof(r0)}(D, a, r0) -@inline function force(b::MorseBond, coord_i, coord_j, boundary) +@inline function force(b::MorseBond, coord_i, coord_j, boundary, args...) dr = vector(coord_i, coord_j, boundary) r = norm(dr) ralp = exp(-b.a * (r - b.r0)) @@ -27,7 +27,7 @@ MorseBond(; D, a, r0) = MorseBond{typeof(D), typeof(a), typeof(r0)}(D, a, r0) return SpecificForce2Atoms(f, -f) end -@inline function potential_energy(b::MorseBond, coord_i, coord_j, boundary) +@inline function potential_energy(b::MorseBond, coord_i, coord_j, boundary, args...) dr = vector(coord_i, coord_j, boundary) r = norm(dr) ralp = exp(-b.a * (r - b.r0)) diff --git a/src/interactions/periodic_torsion.jl b/src/interactions/periodic_torsion.jl index c72aa70dc..da22d43b8 100644 --- a/src/interactions/periodic_torsion.jl +++ b/src/interactions/periodic_torsion.jl @@ -79,7 +79,7 @@ end # The summation gives different errors with Enzyme on CPU and GPU # so there are two similar implementations @inline function force(d::PeriodicTorsion, coords_i, coords_j, coords_k, - coords_l, boundary) + coords_l, boundary, args...) ab, bc, cd, cross_ab_bc, cross_bc_cd, bc_norm, θ = periodic_torsion_vectors( coords_i, coords_j, coords_k, coords_l, boundary) fs = sum(zip(d.periodicities, d.phases, d.ks)) do (periodicity, phase, k) @@ -91,7 +91,7 @@ end end @inline function force_gpu(d::PeriodicTorsion{N}, coords_i, coords_j, coords_k, - coords_l, boundary) where N + coords_l, boundary, args...) where N ab, bc, cd, cross_ab_bc, cross_bc_cd, bc_norm, θ = periodic_torsion_vectors( coords_i, coords_j, coords_k, coords_l, boundary) fi_sum, fj_sum, fk_sum, fl_sum = periodic_torsion_force(d.periodicities[1], d.phases[1], @@ -108,7 +108,7 @@ end end @inline function potential_energy(d::PeriodicTorsion{N}, coords_i, coords_j, coords_k, - coords_l, boundary) where N + coords_l, boundary, args...) where N θ = torsion_angle(coords_i, coords_j, coords_k, coords_l, boundary) k1 = d.ks[1] E = k1 + k1 * cos((d.periodicities[1] * θ) - d.phases[1]) diff --git a/src/interactions/rb_torsion.jl b/src/interactions/rb_torsion.jl index 8c2610862..42147451b 100644 --- a/src/interactions/rb_torsion.jl +++ b/src/interactions/rb_torsion.jl @@ -14,8 +14,7 @@ end RBTorsion(; f1, f2, f3, f4) = RBTorsion{typeof(f1)}(f1, f2, f3, f4) -@inline function force(d::RBTorsion, coords_i, coords_j, coords_k, - coords_l, boundary) +@inline function force(d::RBTorsion, coords_i, coords_j, coords_k, coords_l, boundary, args...) ab = vector(coords_i, coords_j, boundary) bc = vector(coords_j, coords_k, boundary) cd = vector(coords_k, coords_l, boundary) @@ -36,7 +35,7 @@ RBTorsion(; f1, f2, f3, f4) = RBTorsion{typeof(f1)}(f1, f2, f3, f4) end @inline function potential_energy(d::RBTorsion, coords_i, coords_j, coords_k, - coords_l, boundary) + coords_l, boundary, args...) θ = torsion_angle(coords_i, coords_j, coords_k, coords_l, boundary) return (d.f1 * (1 + cos(θ)) + d.f2 * (1 - cos(2θ)) + d.f3 * (1 + cos(3θ)) + d.f4) / 2 end diff --git a/src/spatial.jl b/src/spatial.jl index 87e2909bc..a205196ce 100644 --- a/src/spatial.jl +++ b/src/spatial.jl @@ -713,6 +713,7 @@ end function virial(sys::System{D, G, T}, neighbors_dev, pairwise_inters_nonl, pairwise_inters_nl) where {D, G, T} + step_n = 0 if G coords, velocities, atoms = Array(sys.coords), Array(sys.velocities), Array(sys.atoms) if isnothing(neighbors_dev) @@ -734,10 +735,10 @@ function virial(sys::System{D, G, T}, neighbors_dev, pairwise_inters_nonl, for j in (i + 1):n_atoms dr = vector(coords[i], coords[j], boundary) f = force(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], sys.force_units, false, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) for inter in pairwise_inters_nonl[2:end] f += force(inter, dr, atoms[i], atoms[j], sys.force_units, false, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) end v += dot(f, dr) end @@ -752,10 +753,10 @@ function virial(sys::System{D, G, T}, neighbors_dev, pairwise_inters_nonl, i, j, special = neighbors[ni] dr = vector(coords[i], coords[j], boundary) f = force(pairwise_inters_nl[1], dr, atoms[i], atoms[j], sys.force_units, special, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) for inter in pairwise_inters_nl[2:end] f += force(inter, dr, atoms[i], atoms[j], sys.force_units, special, - coords[i], coords[j], boundary, velocities[i], velocities[j], 0) + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) end v += dot(f, dr) end From 613468608dc84397fd4889c829f50c32ceaa91da Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Mon, 1 Jul 2024 18:34:04 +0100 Subject: [PATCH 16/74] forces reuse memory --- src/cuda.jl | 40 ++++---- src/energy.jl | 21 ++-- src/force.jl | 262 +++++++++++++++++++++++++++++--------------------- 3 files changed, 183 insertions(+), 140 deletions(-) diff --git a/src/cuda.jl b/src/cuda.jl index 8a8c8ee96..a5f31822d 100644 --- a/src/cuda.jl +++ b/src/cuda.jl @@ -29,10 +29,8 @@ function cuda_threads_blocks_specific(n_inters) return n_threads_gpu, n_blocks end -function pairwise_force_gpu(coords::AbstractArray{SVector{D, C}}, velocities, atoms, boundary, - pairwise_inters, nbs, force_units, step_n, ::Val{T}) where {D, C, T} - fs_mat = CUDA.zeros(T, D, length(atoms)) - +function pairwise_force_gpu!(fs_mat, coords::AbstractArray{SVector{D, C}}, velocities, atoms, + boundary, pairwise_inters, nbs, step_n, force_units, ::Val{T}) where {D, C, T} if typeof(nbs) == NoNeighborList kernel = @cuda launch=false pairwise_force_kernel_nonl!( fs_mat, coords, velocities, atoms, boundary, pairwise_inters, step_n, @@ -43,8 +41,8 @@ function pairwise_force_gpu(coords::AbstractArray{SVector{D, C}}, velocities, at nthreads = cld(nthreads, WARPSIZE) * WARPSIZE n_blocks_i = cld(length(atoms), WARPSIZE) n_blocks_j = cld(length(atoms), nthreads) - kernel(fs_mat, coords, velocities, atoms, boundary, pairwise_inters, Val(D), Val(force_units); - threads=nthreads, blocks=(n_blocks_i, n_blocks_j)) + kernel(fs_mat, coords, velocities, atoms, boundary, pairwise_inters, step_n, Val(D), + Val(force_units); threads=nthreads, blocks=(n_blocks_i, n_blocks_j)) else n_threads_gpu, n_blocks = cuda_threads_blocks_pairwise(length(nbs)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks pairwise_force_kernel_nl!( @@ -201,9 +199,8 @@ end return f end -function specific_force_gpu(inter_list::InteractionList1Atoms, coords::AbstractArray{SVector{D, C}}, - velocities, atoms, boundary, force_units, step_n, ::Val{T}) where {D, C, T} - fs_mat = CUDA.zeros(T, D, length(coords)) +function specific_force_gpu!(fs_mat, inter_list::InteractionList1Atoms, coords::AbstractArray{SVector{D, C}}, + velocities, atoms, boundary, step_n, force_units, ::Val{T}) where {D, C, T} n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_force_1_atoms_kernel!(fs_mat, coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.inters, @@ -211,9 +208,8 @@ function specific_force_gpu(inter_list::InteractionList1Atoms, coords::AbstractA return fs_mat end -function specific_force_gpu(inter_list::InteractionList2Atoms, coords::AbstractArray{SVector{D, C}}, - velocities, atoms, boundary, force_units, step_n, ::Val{T}) where {D, C, T} - fs_mat = CUDA.zeros(T, D, length(coords)) +function specific_force_gpu!(fs_mat, inter_list::InteractionList2Atoms, coords::AbstractArray{SVector{D, C}}, + velocities, atoms, boundary, step_n, force_units, ::Val{T}) where {D, C, T} n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_force_2_atoms_kernel!(fs_mat, coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.js, @@ -221,9 +217,8 @@ function specific_force_gpu(inter_list::InteractionList2Atoms, coords::AbstractA return fs_mat end -function specific_force_gpu(inter_list::InteractionList3Atoms, coords::AbstractArray{SVector{D, C}}, - velocities, atoms, boundary, force_units, step_n, ::Val{T}) where {D, C, T} - fs_mat = CUDA.zeros(T, D, length(coords)) +function specific_force_gpu!(fs_mat, inter_list::InteractionList3Atoms, coords::AbstractArray{SVector{D, C}}, + velocities, atoms, boundary, step_n, force_units, ::Val{T}) where {D, C, T} n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_force_3_atoms_kernel!(fs_mat, coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.js, @@ -231,9 +226,8 @@ function specific_force_gpu(inter_list::InteractionList3Atoms, coords::AbstractA return fs_mat end -function specific_force_gpu(inter_list::InteractionList4Atoms, coords::AbstractArray{SVector{D, C}}, - velocities, atoms, boundary, force_units, step_n, ::Val{T}) where {D, C, T} - fs_mat = CUDA.zeros(T, D, length(coords)) +function specific_force_gpu!(fs_mat, inter_list::InteractionList4Atoms, coords::AbstractArray{SVector{D, C}}, + velocities, atoms, boundary, step_n, force_units, ::Val{T}) where {D, C, T} n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_force_4_atoms_kernel!(fs_mat, coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.js, @@ -351,7 +345,7 @@ function specific_force_4_atoms_kernel!(forces, coords_var, velocities_var, atom end function pairwise_pe_gpu(coords::AbstractArray{SVector{D, C}}, velocities, atoms, boundary, - pairwise_inters, nbs, energy_units, step_n, ::Val{T}) where {D, C, T} + pairwise_inters, nbs, step_n, energy_units, ::Val{T}) where {D, C, T} pe_vec = CUDA.zeros(T, 1) n_threads_gpu, n_blocks = cuda_threads_blocks_pairwise(length(nbs)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks pairwise_pe_kernel!( @@ -388,7 +382,7 @@ function pairwise_pe_kernel!(energy, coords_var, velocities_var, atoms_var, boun end function specific_pe_gpu(inter_list::InteractionList1Atoms, coords::AbstractArray{SVector{D, C}}, - velocities, atoms, boundary, energy_units, step_n, ::Val{T}) where {D, C, T} + velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T} pe_vec = CUDA.zeros(T, 1) n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_1_atoms_kernel!(pe_vec, @@ -398,7 +392,7 @@ function specific_pe_gpu(inter_list::InteractionList1Atoms, coords::AbstractArra end function specific_pe_gpu(inter_list::InteractionList2Atoms, coords::AbstractArray{SVector{D, C}}, - velocities, atoms, boundary, energy_units, step_n, ::Val{T}) where {D, C, T} + velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T} pe_vec = CUDA.zeros(T, 1) n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_2_atoms_kernel!(pe_vec, @@ -408,7 +402,7 @@ function specific_pe_gpu(inter_list::InteractionList2Atoms, coords::AbstractArra end function specific_pe_gpu(inter_list::InteractionList3Atoms, coords::AbstractArray{SVector{D, C}}, - velocities, atoms, boundary, energy_units, step_n, ::Val{T}) where {D, C, T} + velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T} pe_vec = CUDA.zeros(T, 1) n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_3_atoms_kernel!(pe_vec, @@ -418,7 +412,7 @@ function specific_pe_gpu(inter_list::InteractionList3Atoms, coords::AbstractArra end function specific_pe_gpu(inter_list::InteractionList4Atoms, coords::AbstractArray{SVector{D, C}}, - velocities, atoms, boundary, energy_units, step_n, ::Val{T}) where {D, C, T} + velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T} pe_vec = CUDA.zeros(T, 1) n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_4_atoms_kernel!(pe_vec, diff --git a/src/energy.jl b/src/energy.jl index 987d62334..87e463ab1 100644 --- a/src/energy.jl +++ b/src/energy.jl @@ -55,12 +55,15 @@ general interactions. potential_energy(inter::PairwiseInteraction, vec_ij, atom_i, atom_j, energy_units, special, coord_i, coord_j, boundary, velocity_i, velocity_j, step_n) - potential_energy(inter::SpecificInteraction, coords_i, coords_j, - boundary) - potential_energy(inter::SpecificInteraction, coords_i, coords_j, - coords_k, boundary) - potential_energy(inter::SpecificInteraction, coords_i, coords_j, - coords_k, coords_l, boundary) + potential_energy(inter::SpecificInteraction, coord_i, boundary, atom_i, energy_units, + velocity_i, step_n) + potential_energy(inter::SpecificInteraction, coord_i, coord_j, boundary, atom_i, atom_j, + energy_units, velocity_i, velocity_j, step_n) + potential_energy(inter::SpecificInteraction, coord_i, coord_j, coord_k, boundary, atom_i, + atom_j, atom_k, energy_units, velocity_i, velocity_j, velocity_k, step_n) + potential_energy(inter::SpecificInteraction, coord_i, coord_j, coord_k, coord_l, boundary, + atom_i, atom_j, atom_k, atom_l, energy_units, velocity_i, velocity_j, + velocity_k, velocity_l, step_n) Calculate the potential energy due to a given interaction type. @@ -250,7 +253,7 @@ function potential_energy(sys::System{D, true, T}, neighbors; if length(pairwise_inters_nonl) > 0 nbs = NoNeighborList(n_atoms) pe_vec += pairwise_pe_gpu(sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nonl, - nbs, sys.energy_units, step_n, val_ft) + nbs, step_n, sys.energy_units, val_ft) end pairwise_inters_nl = filter(use_neighbors, values(sys.pairwise_inters)) @@ -261,13 +264,13 @@ function potential_energy(sys::System{D, true, T}, neighbors; if length(neighbors) > 0 nbs = @view neighbors.list[1:neighbors.n] pe_vec += pairwise_pe_gpu(sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nl, - nbs, sys.energy_units, step_n, val_ft) + nbs, step_n, sys.energy_units, val_ft) end end for inter_list in values(sys.specific_inter_lists) pe_vec += specific_pe_gpu(inter_list, sys.coords, sys.velocities, sys.atoms, sys.boundary, - sys.energy_units, step_n, val_ft) + step_n, sys.energy_units, val_ft) end pe = Array(pe_vec)[1] diff --git a/src/force.jl b/src/force.jl index 5d79feff7..3cf0c9499 100644 --- a/src/force.jl +++ b/src/force.jl @@ -27,12 +27,15 @@ end """ force(inter::PairwiseInteraction, vec_ij, atom_i, atom_j, force_units, special, coord_i, coord_j, boundary, velocity_i, velocity_j, step_n) - force(inter::SpecificInteraction, coord_i, coord_j, - boundary) - force(inter::SpecificInteraction, coord_i, coord_j, - coord_k, boundary) - force(inter::SpecificInteraction, coord_i, coord_j, - coord_k, coord_l, boundary) + force(inter::SpecificInteraction, coord_i, boundary, atom_i, force_units, + velocity_i, step_n) + force(inter::SpecificInteraction, coord_i, coord_j, boundary, atom_i, atom_j, + force_units, velocity_i, velocity_j, step_n) + force(inter::SpecificInteraction, coord_i, coord_j, coord_k, boundary, atom_i, + atom_j, atom_k, force_units, velocity_i, velocity_j, velocity_k, step_n) + force(inter::SpecificInteraction, coord_i, coord_j, coord_k, coord_l, boundary, + atom_i, atom_j, atom_k, atom_l, force_units, velocity_i, velocity_j, + velocity_k, velocity_l, step_n) Calculate the force between atoms due to a given interaction type. @@ -114,6 +117,18 @@ Base.:+(x::SpecificForce2Atoms, y::SpecificForce2Atoms) = SpecificForce2Atoms(x. Base.:+(x::SpecificForce3Atoms, y::SpecificForce3Atoms) = SpecificForce3Atoms(x.f1 + y.f1, x.f2 + y.f2, x.f3 + y.f3) Base.:+(x::SpecificForce4Atoms, y::SpecificForce4Atoms) = SpecificForce4Atoms(x.f1 + y.f1, x.f2 + y.f2, x.f3 + y.f3, x.f4 + y.f4) +function init_forces_buffer(forces_nounits, n_threads) + if n_threads == 1 + return nothing + else + return [similar(forces_nounits) for _ in 1:n_threads] + end +end + +function init_forces_buffer(forces_nounits::CuArray{SVector{D, T}}, n_threads) where {D, T} + return CUDA.zeros(T, D, length(forces_nounits)) +end + """ forces(system, neighbors=find_neighbors(sys); n_threads=Threads.nthreads()) @@ -124,8 +139,16 @@ function forces(sys; n_threads::Integer=Threads.nthreads()) return forces(sys, find_neighbors(sys; n_threads=n_threads); n_threads=n_threads) end -function forces(sys::System{D, false}, neighbors; - n_threads::Integer=Threads.nthreads()) where D +function forces(sys, neighbors, step_n::Integer=0; n_threads::Integer=Threads.nthreads()) + forces_nounits = ustrip_vec.(zero(sys.coords)) + forces_buffer = init_forces_buffer(forces_nounits, n_threads) + forces_nounits!(forces_nounits, sys, neighbors, forces_buffer, step_n; n_threads=n_threads) + return forces_nounits .* sys.force_units +end + +function forces_nounits!(fs_nounits, sys::System{D, false}, neighbors, fs_chunks=nothing, + step_n::Integer=0; n_threads::Integer=Threads.nthreads()) where D + n_atoms = length(sys) pairwise_inters_nonl = filter(!use_neighbors, values(sys.pairwise_inters)) pairwise_inters_nl = filter( use_neighbors, values(sys.pairwise_inters)) sils_1_atoms = filter(il -> il isa InteractionList1Atoms, values(sys.specific_inter_lists)) @@ -133,103 +156,118 @@ function forces(sys::System{D, false}, neighbors; sils_3_atoms = filter(il -> il isa InteractionList3Atoms, values(sys.specific_inter_lists)) sils_4_atoms = filter(il -> il isa InteractionList4Atoms, values(sys.specific_inter_lists)) - fs = forces_pair_spec(sys.coords, sys.velocities, sys.atoms, pairwise_inters_nonl, pairwise_inters_nl, - sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, - sys.boundary, sys.force_units, neighbors, n_threads) + if length(sys.pairwise_inters) > 0 + if n_threads > 1 + pairwise_forces_threads!(fs_nounits, fs_chunks, neighbors, sys.atoms, sys.coords, + sys.velocities, sys.boundary, sys.force_units, n_atoms, + pairwise_inters_nonl, pairwise_inters_nl, n_threads, step_n) + else + pairwise_forces!(fs_nounits, neighbors, sys.atoms, sys.coords, sys.velocities, + sys.boundary, sys.force_units, n_atoms, pairwise_inters_nonl, + pairwise_inters_nl, step_n) + end + else + fill!(fs_nounits, zero(eltype(fs_nounits))) + end + + specific_forces!(fs_nounits, sys.atoms, sys.coords, sys.velocities, sys.boundary, + sys.force_units, sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, step_n) for inter in values(sys.general_inters) - fs += AtomsCalculators.forces(sys, inter; neighbors=neighbors, n_threads=n_threads) + fs_gen = AtomsCalculators.forces(sys, inter; neighbors=neighbors, n_threads=n_threads) + check_force_units(zero(eltype(fs_gen)), sys.force_units) + fs_nounits .+= ustrip_vec.(fs_gen) end - return fs + return fs_nounits end -function forces_pair_spec(coords, velocities, atoms, pairwise_inters_nonl, pairwise_inters_nl, - sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, - boundary, force_units, neighbors, n_threads) - fs = ustrip_vec.(zero(coords)) - forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, pairwise_inters_nl, - sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, boundary, - force_units, neighbors, n_threads) - return fs * force_units -end +function pairwise_forces!(fs_nounits, neighbors, atoms, coords, velocities, boundary, force_units, n_atoms, + pairwise_inters_nonl, pairwise_inters_nl, step_n=0) + fill!(fs_nounits, zero(eltype(fs_nounits))) -function forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, pairwise_inters_nl, - sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, - boundary, force_units, neighbors, n_threads) - step_n = 0 - n_atoms = length(coords) - @inbounds if n_threads > 1 - fs_chunks = [zero(fs) for _ in 1:n_threads] - - if length(pairwise_inters_nonl) > 0 - Threads.@threads for chunk_i in 1:n_threads - for i in chunk_i:n_threads:n_atoms - for j in (i + 1):n_atoms - dr = vector(coords[i], coords[j], boundary) - f = force(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], force_units, - false, coords[i], coords[j], boundary, velocities[i], - velocities[j], step_n) - for inter in pairwise_inters_nonl[2:end] - f += force(inter, dr, atoms[i], atoms[j], force_units, false, coords[i], - coords[j], boundary, velocities[i], velocities[j], step_n) - end - check_force_units(f, force_units) - f_ustrip = ustrip.(f) - fs_chunks[chunk_i][i] -= f_ustrip - fs_chunks[chunk_i][j] += f_ustrip - end + @inbounds if length(pairwise_inters_nonl) > 0 + for i in 1:n_atoms + for j in (i + 1):n_atoms + dr = vector(coords[i], coords[j], boundary) + f = force(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], force_units, false, + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) + for inter in pairwise_inters_nonl[2:end] + f += force(inter, dr, atoms[i], atoms[j], force_units, false, coords[i], + coords[j], boundary, velocities[i], velocities[j], step_n) end + check_force_units(f, force_units) + f_ustrip = ustrip.(f) + fs_nounits[i] -= f_ustrip + fs_nounits[j] += f_ustrip end end + end - if length(pairwise_inters_nl) > 0 - if isnothing(neighbors) - error("an interaction uses the neighbor list but neighbors is nothing") - end - Threads.@threads for chunk_i in 1:n_threads - for ni in chunk_i:n_threads:length(neighbors) - i, j, special = neighbors[ni] - dr = vector(coords[i], coords[j], boundary) - f = force(pairwise_inters_nl[1], dr, atoms[i], atoms[j], force_units, special, - coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) - for inter in pairwise_inters_nl[2:end] - f += force(inter, dr, atoms[i], atoms[j], force_units, special, coords[i], - coords[j], boundary, velocities[i], velocities[j], step_n) - end - check_force_units(f, force_units) - f_ustrip = ustrip.(f) - fs_chunks[chunk_i][i] -= f_ustrip - fs_chunks[chunk_i][j] += f_ustrip - end + @inbounds if length(pairwise_inters_nl) > 0 + if isnothing(neighbors) + error("an interaction uses the neighbor list but neighbors is nothing") + end + for ni in eachindex(neighbors) + i, j, special = neighbors[ni] + dr = vector(coords[i], coords[j], boundary) + f = force(pairwise_inters_nl[1], dr, atoms[i], atoms[j], force_units, special, + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) + for inter in pairwise_inters_nl[2:end] + f += force(inter, dr, atoms[i], atoms[j], force_units, special, coords[i], + coords[j], boundary, velocities[i], velocities[j], step_n) end + check_force_units(f, force_units) + f_ustrip = ustrip.(f) + fs_nounits[i] -= f_ustrip + fs_nounits[j] += f_ustrip end + end - fs .+= sum(fs_chunks) - else - if length(pairwise_inters_nonl) > 0 - for i in 1:n_atoms + return fs_nounits +end + +function pairwise_forces_threads!(fs_nounits, fs_chunks, neighbors, atoms, coords, velocities, boundary, + force_units, n_atoms, pairwise_inters_nonl, pairwise_inters_nl, + n_threads, step_n=0) + if isnothing(fs_chunks) + throw(ArgumentError("fs_chunks is not set but n_threads is > 1")) + end + if length(fs_chunks) != n_threads + throw(ArgumentError("length of fs_chunks ($(length(fs_chunks))) does not " * + "match n_threads ($n_threads)")) + end + @inbounds for chunk_i in 1:n_threads + fill!(fs_chunks[chunk_i], zero(eltype(fs_nounits))) + end + + @inbounds if length(pairwise_inters_nonl) > 0 + Threads.@threads for chunk_i in 1:n_threads + for i in chunk_i:n_threads:n_atoms for j in (i + 1):n_atoms dr = vector(coords[i], coords[j], boundary) - f = force(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], force_units, false, - coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) + f = force(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], force_units, + false, coords[i], coords[j], boundary, velocities[i], + velocities[j], step_n) for inter in pairwise_inters_nonl[2:end] f += force(inter, dr, atoms[i], atoms[j], force_units, false, coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) end check_force_units(f, force_units) f_ustrip = ustrip.(f) - fs[i] -= f_ustrip - fs[j] += f_ustrip + fs_chunks[chunk_i][i] -= f_ustrip + fs_chunks[chunk_i][j] += f_ustrip end end end + end - if length(pairwise_inters_nl) > 0 - if isnothing(neighbors) - error("an interaction uses the neighbor list but neighbors is nothing") - end - for ni in eachindex(neighbors) + @inbounds if length(pairwise_inters_nl) > 0 + if isnothing(neighbors) + error("an interaction uses the neighbor list but neighbors is nothing") + end + Threads.@threads for chunk_i in 1:n_threads + for ni in chunk_i:n_threads:length(neighbors) i, j, special = neighbors[ni] dr = vector(coords[i], coords[j], boundary) f = force(pairwise_inters_nl[1], dr, atoms[i], atoms[j], force_units, special, @@ -240,17 +278,27 @@ function forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, end check_force_units(f, force_units) f_ustrip = ustrip.(f) - fs[i] -= f_ustrip - fs[j] += f_ustrip + fs_chunks[chunk_i][i] -= f_ustrip + fs_chunks[chunk_i][j] += f_ustrip end end end + @inbounds fs_nounits .= fs_chunks[1] + @inbounds for chunk_i in 2:n_threads + fs_nounits .+= fs_chunks[chunk_i] + end + + return fs_nounits +end + +function specific_forces!(fs_nounits, atoms, coords, velocities, boundary, force_units, sils_1_atoms, sils_2_atoms, + sils_3_atoms, sils_4_atoms, step_n=0) @inbounds for inter_list in sils_1_atoms for (i, inter) in zip(inter_list.is, inter_list.inters) sf = force(inter, coords[i], boundary, atoms[i], force_units, velocities[i], step_n) check_force_units(sf.f1, force_units) - fs[i] += ustrip.(sf.f1) + fs_nounits[i] += ustrip.(sf.f1) end end @@ -260,8 +308,8 @@ function forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, velocities[i], velocities[j], step_n) check_force_units(sf.f1, force_units) check_force_units(sf.f2, force_units) - fs[i] += ustrip.(sf.f1) - fs[j] += ustrip.(sf.f2) + fs_nounits[i] += ustrip.(sf.f1) + fs_nounits[j] += ustrip.(sf.f2) end end @@ -272,9 +320,9 @@ function forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, check_force_units(sf.f1, force_units) check_force_units(sf.f2, force_units) check_force_units(sf.f3, force_units) - fs[i] += ustrip.(sf.f1) - fs[j] += ustrip.(sf.f2) - fs[k] += ustrip.(sf.f3) + fs_nounits[i] += ustrip.(sf.f1) + fs_nounits[j] += ustrip.(sf.f2) + fs_nounits[k] += ustrip.(sf.f3) end end @@ -288,28 +336,26 @@ function forces_pair_spec!(fs, coords, velocities, atoms, pairwise_inters_nonl, check_force_units(sf.f2, force_units) check_force_units(sf.f3, force_units) check_force_units(sf.f4, force_units) - fs[i] += ustrip.(sf.f1) - fs[j] += ustrip.(sf.f2) - fs[k] += ustrip.(sf.f3) - fs[l] += ustrip.(sf.f4) + fs_nounits[i] += ustrip.(sf.f1) + fs_nounits[j] += ustrip.(sf.f2) + fs_nounits[k] += ustrip.(sf.f3) + fs_nounits[l] += ustrip.(sf.f4) end end - return nothing + return fs_nounits end -function forces(sys::System{D, true, T}, neighbors; - n_threads::Integer=Threads.nthreads()) where {D, T} - step_n = 0 - n_atoms = length(sys) +function forces_nounits!(fs_nounits, sys::System{D, true, T}, neighbors, + fs_mat=CUDA.zeros(T, D, length(sys)), step_n::Integer=0; + n_threads::Integer=Threads.nthreads()) where {D, T} + fill!(fs_mat, zero(T)) val_ft = Val(T) - fs_mat = CUDA.zeros(T, D, n_atoms) - pairwise_inters_nonl = filter(!use_neighbors, values(sys.pairwise_inters)) if length(pairwise_inters_nonl) > 0 - nbs = NoNeighborList(n_atoms) - fs_mat += pairwise_force_gpu(sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nonl, - nbs, sys.force_units, step_n, val_ft) + nbs = NoNeighborList(length(sys)) + pairwise_force_gpu!(fs_mat, sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nonl, + nbs, step_n, sys.force_units, val_ft) end pairwise_inters_nl = filter(use_neighbors, values(sys.pairwise_inters)) @@ -319,23 +365,23 @@ function forces(sys::System{D, true, T}, neighbors; end if length(neighbors) > 0 nbs = @view neighbors.list[1:neighbors.n] - fs_mat += pairwise_force_gpu(sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nl, - nbs, sys.force_units, step_n, val_ft) + pairwise_force_gpu!(fs_mat, sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nonl, + nbs, step_n, sys.force_units, val_ft) end end for inter_list in values(sys.specific_inter_lists) - fs_mat += specific_force_gpu(inter_list, sys.coords, sys.velocities, sys.atoms, - sys.boundary, sys.force_units, step_n, val_ft) + specific_force_gpu!(fs_mat, inter_list, sys.coords, sys.velocities, sys.atoms, + sys.boundary, step_n, sys.force_units, val_ft) end - fs = reinterpret(SVector{D, T}, vec(fs_mat)) + fs_nounits .= reinterpret(SVector{D, T}, vec(fs_mat)) for inter in values(sys.general_inters) fs_gen = AtomsCalculators.forces(sys, inter; neighbors=neighbors, n_threads=n_threads) check_force_units(zero(eltype(fs_gen)), sys.force_units) - fs += ustrip_vec.(fs_gen) + fs_nounits .+= ustrip_vec.(fs_gen) end - return fs * sys.force_units + return fs_nounits end From c0f47d2543efe6c211c4289097eccf6256304e0d Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Tue, 2 Jul 2024 12:59:50 +0100 Subject: [PATCH 17/74] add step_n to virial --- src/spatial.jl | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/spatial.jl b/src/spatial.jl index a205196ce..488af9fce 100644 --- a/src/spatial.jl +++ b/src/spatial.jl @@ -678,8 +678,8 @@ function remove_CM_motion!(sys) end @doc raw""" - virial(sys, neighbors=find_neighbors(sys); n_threads=Threads.nthreads()) - virial(inter, sys, neighbors; n_threads=Threads.nthreads()) + virial(sys, neighbors=find_neighbors(sys), step_n=0; n_threads=Threads.nthreads()) + virial(inter, sys, neighbors; step_n=0, n_threads=Threads.nthreads()) Calculate the virial of a system or the virial resulting from a general interaction. @@ -699,21 +699,20 @@ function virial(sys; n_threads::Integer=Threads.nthreads()) return virial(sys, find_neighbors(sys; n_threads=n_threads); n_threads=n_threads) end -function virial(sys, neighbors; n_threads::Integer=Threads.nthreads()) +function virial(sys, neighbors, step_n::Integer=0; n_threads::Integer=Threads.nthreads()) pairwise_inters_nonl = filter(!use_neighbors, values(sys.pairwise_inters)) pairwise_inters_nl = filter( use_neighbors, values(sys.pairwise_inters)) - v = virial(sys, neighbors, pairwise_inters_nonl, pairwise_inters_nl) + v = virial(sys, neighbors, step_n, pairwise_inters_nonl, pairwise_inters_nl) for inter in values(sys.general_inters) - v += virial(inter, sys, neighbors; n_threads=n_threads) + v += virial(inter, sys, neighbors; step_n=step_n, n_threads=n_threads) end return v end -function virial(sys::System{D, G, T}, neighbors_dev, pairwise_inters_nonl, +function virial(sys::System{D, G, T}, neighbors_dev, step_n, pairwise_inters_nonl, pairwise_inters_nl) where {D, G, T} - step_n = 0 if G coords, velocities, atoms = Array(sys.coords), Array(sys.velocities), Array(sys.atoms) if isnothing(neighbors_dev) From 06dc81510a4ef00db523222fb07da146393d7750 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Tue, 2 Jul 2024 13:00:07 +0100 Subject: [PATCH 18/74] refactor energy functions --- src/cuda.jl | 62 +++++++------ src/energy.jl | 236 ++++++++++++++++++++++++++------------------------ src/force.jl | 28 +++--- 3 files changed, 167 insertions(+), 159 deletions(-) diff --git a/src/cuda.jl b/src/cuda.jl index a5f31822d..65ba1ede5 100644 --- a/src/cuda.jl +++ b/src/cuda.jl @@ -344,14 +344,14 @@ function specific_force_4_atoms_kernel!(forces, coords_var, velocities_var, atom return nothing end -function pairwise_pe_gpu(coords::AbstractArray{SVector{D, C}}, velocities, atoms, boundary, - pairwise_inters, nbs, step_n, energy_units, ::Val{T}) where {D, C, T} - pe_vec = CUDA.zeros(T, 1) +function pairwise_pe_gpu!(pe_vec_nounits, coords::AbstractArray{SVector{D, C}}, velocities, atoms, + boundary, pairwise_inters, nbs, step_n, energy_units, + ::Val{T}) where {D, C, T} n_threads_gpu, n_blocks = cuda_threads_blocks_pairwise(length(nbs)) CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks pairwise_pe_kernel!( - pe_vec, coords, velocities, atoms, boundary, pairwise_inters, nbs, + pe_vec_nounits, coords, velocities, atoms, boundary, pairwise_inters, nbs, step_n, Val(energy_units)) - return pe_vec + return pe_vec_nounits end function pairwise_pe_kernel!(energy, coords_var, velocities_var, atoms_var, boundary, inters, @@ -381,44 +381,40 @@ function pairwise_pe_kernel!(energy, coords_var, velocities_var, atoms_var, boun return nothing end -function specific_pe_gpu(inter_list::InteractionList1Atoms, coords::AbstractArray{SVector{D, C}}, - velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T} - pe_vec = CUDA.zeros(T, 1) +function specific_pe_gpu!(pe_vec_nounits, inter_list::InteractionList1Atoms, coords::AbstractArray{SVector{D, C}}, + velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T} n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_1_atoms_kernel!(pe_vec, - coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.inters, - Val(energy_units)) - return pe_vec + CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_1_atoms_kernel!( + pe_vec_nounits, coords, velocities, atoms, boundary, step_n, inter_list.is, + inter_list.inters, Val(energy_units)) + return pe_vec_nounits end -function specific_pe_gpu(inter_list::InteractionList2Atoms, coords::AbstractArray{SVector{D, C}}, - velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T} - pe_vec = CUDA.zeros(T, 1) +function specific_pe_gpu!(pe_vec_nounits, inter_list::InteractionList2Atoms, coords::AbstractArray{SVector{D, C}}, + velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T} n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_2_atoms_kernel!(pe_vec, - coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.js, - inter_list.inters, Val(energy_units)) - return pe_vec + CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_2_atoms_kernel!( + pe_vec_nounits, coords, velocities, atoms, boundary, step_n, inter_list.is, + inter_list.js, inter_list.inters, Val(energy_units)) + return pe_vec_nounits end -function specific_pe_gpu(inter_list::InteractionList3Atoms, coords::AbstractArray{SVector{D, C}}, - velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T} - pe_vec = CUDA.zeros(T, 1) +function specific_pe_gpu!(pe_vec_nounits, inter_list::InteractionList3Atoms, coords::AbstractArray{SVector{D, C}}, + velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T} n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_3_atoms_kernel!(pe_vec, - coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.js, - inter_list.ks, inter_list.inters, Val(energy_units)) - return pe_vec + CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_3_atoms_kernel!( + pe_vec_nounits, coords, velocities, atoms, boundary, step_n, inter_list.is, + inter_list.js, inter_list.ks, inter_list.inters, Val(energy_units)) + return pe_vec_nounits end -function specific_pe_gpu(inter_list::InteractionList4Atoms, coords::AbstractArray{SVector{D, C}}, - velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T} - pe_vec = CUDA.zeros(T, 1) +function specific_pe_gpu!(pe_vec_nounits, inter_list::InteractionList4Atoms, coords::AbstractArray{SVector{D, C}}, + velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T} n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list)) - CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_4_atoms_kernel!(pe_vec, - coords, velocities, atoms, boundary, step_n, inter_list.is, inter_list.js, - inter_list.ks, inter_list.ls, inter_list.inters, Val(energy_units)) - return pe_vec + CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_4_atoms_kernel!( + pe_vec_nounits, coords, velocities, atoms, boundary, step_n, inter_list.is, + inter_list.js, inter_list.ks, inter_list.ls, inter_list.inters, Val(energy_units)) + return pe_vec_nounits end function specific_pe_1_atoms_kernel!(energy, coords_var, velocities_var, atoms_var, boundary, diff --git a/src/energy.jl b/src/energy.jl index 87e463ab1..5e96b302e 100644 --- a/src/energy.jl +++ b/src/energy.jl @@ -48,7 +48,7 @@ function temperature(sys) end """ - potential_energy(system, neighbors=find_neighbors(sys); n_threads=Threads.nthreads()) + potential_energy(system, neighbors=find_neighbors(sys), step_n=0; n_threads=Threads.nthreads()) Calculate the potential energy of a system using the pairwise, specific and general interactions. @@ -73,15 +73,8 @@ function potential_energy(sys; n_threads::Integer=Threads.nthreads()) return potential_energy(sys, find_neighbors(sys; n_threads=n_threads); n_threads=n_threads) end -# Allow GPU-specific potential energy functions to be defined if required -potential_energy_gpu(inter::PairwiseInteraction, dr, ai, aj, eu, sp, ci, cj, bnd, vi, vj, sn) = potential_energy(inter, dr, ai, aj, eu, sp, ci, cj, bnd, vi, vj, sn) -potential_energy_gpu(inter::SpecificInteraction, ci, bnd, ai, eu, vi, sn) = potential_energy(inter, ci, bnd, ai, eu, vi, sn) -potential_energy_gpu(inter::SpecificInteraction, ci, cj, bnd, ai, aj, eu, vi, vj, sn) = potential_energy(inter, ci, cj, bnd, ai, aj, eu, vi, vj, sn) -potential_energy_gpu(inter::SpecificInteraction, ci, cj, ck, bnd, ai, aj, ak, eu, vi, vj, vk, sn) = potential_energy(inter, ci, cj, ck, bnd, ai, aj, ak, eu, vi, vj, vk, sn) -potential_energy_gpu(inter::SpecificInteraction, ci, cj, ck, cl, bnd, ai, aj, ak, al, eu, vi, vj, vk, vl, sn) = potential_energy(inter, ci, cj, ck, cl, bnd, ai, aj, ak, al, eu, vi, vj, vk, vl, sn) - -function potential_energy(sys::System{D, false}, neighbors; - n_threads::Integer=Threads.nthreads()) where D +function potential_energy(sys::System{D, false, T}, neighbors; step_n::Integer=0, + n_threads::Integer=Threads.nthreads()) where {D, T} pairwise_inters_nonl = filter(!use_neighbors, values(sys.pairwise_inters)) pairwise_inters_nl = filter( use_neighbors, values(sys.pairwise_inters)) sils_1_atoms = filter(il -> il isa InteractionList1Atoms, values(sys.specific_inter_lists)) @@ -89,171 +82,181 @@ function potential_energy(sys::System{D, false}, neighbors; sils_3_atoms = filter(il -> il isa InteractionList3Atoms, values(sys.specific_inter_lists)) sils_4_atoms = filter(il -> il isa InteractionList4Atoms, values(sys.specific_inter_lists)) - ft = typeof(ustrip(sys.coords[1][1])) # Allow types like those from Measurements.jl - pe = potential_energy_pair_spec(sys.coords, sys.velocities, sys.atoms, pairwise_inters_nonl, pairwise_inters_nl, - sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, sys.boundary, - sys.energy_units, neighbors, n_threads, Val(ft)) + if length(sys.pairwise_inters) > 0 + if n_threads > 1 + pe = pairwise_pe_threads(sys.atoms, sys.coords, sys.velocities, sys.boundary, + neighbors, sys.energy_units, length(sys), pairwise_inters_nonl, + pairwise_inters_nl, T, n_threads, step_n) + else + pe = pairwise_pe(sys.atoms, sys.coords, sys.velocities, sys.boundary, neighbors, + sys.energy_units, length(sys), pairwise_inters_nonl, + pairwise_inters_nl, T, step_n) + end + else + pe = zero(T) * sys.energy_units + end + + if length(sys.specific_inter_lists) > 0 + pe += specific_pe(sys.atoms, sys.coords, sys.velocities, sys.boundary, sys.energy_units, + sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, T, step_n) + end for inter in values(sys.general_inters) - pe += AtomsCalculators.potential_energy(sys, inter; neighbors=neighbors, n_threads=n_threads) + pe += uconvert( + sys.energy_units, + AtomsCalculators.potential_energy(sys, inter; neighbors=neighbors, n_threads=n_threads), + ) end return pe end -function potential_energy_pair_spec(coords, velocities, atoms, pairwise_inters_nonl, pairwise_inters_nl, - sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, - boundary, energy_units, neighbors, n_threads, - val_ft::Val{T}) where T - pe_vec = zeros(T, 1) - potential_energy_pair_spec!(pe_vec, coords, velocities, atoms, pairwise_inters_nonl, pairwise_inters_nl, - sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, boundary, - energy_units, neighbors, n_threads, val_ft) - return pe_vec[1] * energy_units -end +function pairwise_pe(atoms, coords, velocities, boundary, neighbors, energy_units, + n_atoms, pairwise_inters_nonl, pairwise_inters_nl, float_type, step_n=0) + pe = zero(float_type) * energy_units -function potential_energy_pair_spec!(pe_vec, coords, velocities, atoms, pairwise_inters_nonl, - pairwise_inters_nl, sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, - boundary, energy_units, neighbors, n_threads, ::Val{T}) where T - step_n = 0 - pe_sum = zero(T) - - @inbounds if n_threads > 1 - pe_sum_chunks = [zero(T) for _ in 1:n_threads] - - if length(pairwise_inters_nonl) > 0 - n_atoms = length(coords) - Threads.@threads for chunk_i in 1:n_threads - for i in chunk_i:n_threads:n_atoms - for j in (i + 1):n_atoms - dr = vector(coords[i], coords[j], boundary) - pe = potential_energy(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], energy_units, false, - coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) - for inter in pairwise_inters_nonl[2:end] - pe += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, false, - coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) - end - check_energy_units(pe, energy_units) - pe_sum_chunks[chunk_i] += ustrip(pe) - end + @inbounds if length(pairwise_inters_nonl) > 0 + n_atoms = length(coords) + for i in 1:n_atoms + for j in (i + 1):n_atoms + dr = vector(coords[i], coords[j], boundary) + pe_sum = potential_energy(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], energy_units, false, + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) + for inter in pairwise_inters_nonl[2:end] + pe_sum += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, false, + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) end + check_energy_units(pe_sum, energy_units) + pe += pe_sum end end + end - if length(pairwise_inters_nl) > 0 - if isnothing(neighbors) - error("an interaction uses the neighbor list but neighbors is nothing") - end - Threads.@threads for chunk_i in 1:n_threads - for ni in chunk_i:n_threads:length(neighbors) - i, j, special = neighbors[ni] - dr = vector(coords[i], coords[j], boundary) - pe = potential_energy(pairwise_inters_nl[1], dr, atoms[i], atoms[j], energy_units, special, - coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) - for inter in pairwise_inters_nl[2:end] - pe += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, special, - coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) - end - check_energy_units(pe, energy_units) - pe_sum_chunks[chunk_i] += ustrip(pe) - end + @inbounds if length(pairwise_inters_nl) > 0 + if isnothing(neighbors) + error("an interaction uses the neighbor list but neighbors is nothing") + end + for ni in eachindex(neighbors) + i, j, special = neighbors[ni] + dr = vector(coords[i], coords[j], boundary) + pe_sum = potential_energy(pairwise_inters_nl[1], dr, atoms[i], atoms[j], energy_units, special, + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) + for inter in pairwise_inters_nl[2:end] + pe_sum += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, special, + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) end + check_energy_units(pe_sum, energy_units) + pe += pe_sum end + end - pe_sum += sum(pe_sum_chunks) - else - if length(pairwise_inters_nonl) > 0 - n_atoms = length(coords) - for i in 1:n_atoms + return pe +end + +function pairwise_pe_threads(atoms, coords, velocities, boundary, neighbors, energy_units, n_atoms, + pairwise_inters_nonl, pairwise_inters_nl, float_type, n_threads, + step_n=0) + pe_chunks_nounits = zeros(float_type, n_threads) + + if length(pairwise_inters_nonl) > 0 + n_atoms = length(coords) + Threads.@threads for chunk_i in 1:n_threads + for i in chunk_i:n_threads:n_atoms for j in (i + 1):n_atoms dr = vector(coords[i], coords[j], boundary) - pe = potential_energy(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], energy_units, false, + pe_sum = potential_energy(pairwise_inters_nonl[1], dr, atoms[i], atoms[j], energy_units, false, coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) for inter in pairwise_inters_nonl[2:end] - pe += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, false, - coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) + pe_sum += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, false, + coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) end - check_energy_units(pe, energy_units) - pe_sum += ustrip(pe) + check_energy_units(pe_sum, energy_units) + pe_chunks_nounits[chunk_i] += ustrip(pe_sum) end end end + end - if length(pairwise_inters_nl) > 0 - if isnothing(neighbors) - error("an interaction uses the neighbor list but neighbors is nothing") - end - for ni in eachindex(neighbors) + if length(pairwise_inters_nl) > 0 + if isnothing(neighbors) + error("an interaction uses the neighbor list but neighbors is nothing") + end + Threads.@threads for chunk_i in 1:n_threads + for ni in chunk_i:n_threads:length(neighbors) i, j, special = neighbors[ni] dr = vector(coords[i], coords[j], boundary) - pe = potential_energy(pairwise_inters_nl[1], dr, atoms[i], atoms[j], energy_units, special, + pe_sum = potential_energy(pairwise_inters_nl[1], dr, atoms[i], atoms[j], energy_units, special, coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) for inter in pairwise_inters_nl[2:end] - pe += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, special, + pe_sum += potential_energy(inter, dr, atoms[i], atoms[j], energy_units, special, coords[i], coords[j], boundary, velocities[i], velocities[j], step_n) end - check_energy_units(pe, energy_units) - pe_sum += ustrip(pe) + check_energy_units(pe_sum, energy_units) + pe_chunks_nounits[chunk_i] += ustrip(pe_sum) end end end + return sum(pe_chunks_nounits) * energy_units +end + +function specific_pe(atoms, coords, velocities, boundary, energy_units, sils_1_atoms, + sils_2_atoms, sils_3_atoms, sils_4_atoms, float_type, step_n=0) + pe = zero(float_type) * energy_units + @inbounds for inter_list in sils_1_atoms for (i, inter) in zip(inter_list.is, inter_list.inters) - pe = potential_energy(inter, coords[i], boundary, atoms[i], energy_units, + pe_inter = potential_energy(inter, coords[i], boundary, atoms[i], energy_units, velocities[i], step_n) - check_energy_units(pe, energy_units) - pe_sum += ustrip(pe) + check_energy_units(pe_inter, energy_units) + pe += pe_inter end end @inbounds for inter_list in sils_2_atoms for (i, j, inter) in zip(inter_list.is, inter_list.js, inter_list.inters) - pe = potential_energy(inter, coords[i], coords[j], boundary, atoms[i], atoms[j], + pe_inter = potential_energy(inter, coords[i], coords[j], boundary, atoms[i], atoms[j], energy_units, velocities[i], velocities[j], step_n) - check_energy_units(pe, energy_units) - pe_sum += ustrip(pe) + check_energy_units(pe_inter, energy_units) + pe += pe_inter end end @inbounds for inter_list in sils_3_atoms for (i, j, k, inter) in zip(inter_list.is, inter_list.js, inter_list.ks, inter_list.inters) - pe = potential_energy(inter, coords[i], coords[j], coords[k], boundary, atoms[i], + pe_inter = potential_energy(inter, coords[i], coords[j], coords[k], boundary, atoms[i], atoms[j], atoms[k], energy_units, velocities[i], velocities[j], velocities[k], step_n) - check_energy_units(pe, energy_units) - pe_sum += ustrip(pe) + check_energy_units(pe_inter, energy_units) + pe += pe_inter end end @inbounds for inter_list in sils_4_atoms for (i, j, k, l, inter) in zip(inter_list.is, inter_list.js, inter_list.ks, inter_list.ls, inter_list.inters) - pe = potential_energy(inter, coords[i], coords[j], coords[k], coords[l], boundary, + pe_inter = potential_energy(inter, coords[i], coords[j], coords[k], coords[l], boundary, atoms[i], atoms[j], atoms[k], atoms[l], energy_units, velocities[i], velocities[j], velocities[k], velocities[l], step_n) - check_energy_units(pe, energy_units) - pe_sum += ustrip(pe) + check_energy_units(pe_inter, energy_units) + pe += pe_inter end end - pe_vec[1] = pe_sum - return nothing + return pe end -function potential_energy(sys::System{D, true, T}, neighbors; +function potential_energy(sys::System{D, true, T}, neighbors, step_n::Integer=0; n_threads::Integer=Threads.nthreads()) where {D, T} - step_n = 0 - n_atoms = length(sys) + pe_vec_nounits = CUDA.zeros(T, 1) val_ft = Val(T) - pe_vec = CUDA.zeros(T, 1) pairwise_inters_nonl = filter(!use_neighbors, values(sys.pairwise_inters)) if length(pairwise_inters_nonl) > 0 - nbs = NoNeighborList(n_atoms) - pe_vec += pairwise_pe_gpu(sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nonl, - nbs, step_n, sys.energy_units, val_ft) + nbs = NoNeighborList(length(sys)) + pairwise_pe_gpu!(pe_vec_nounits, sys.coords, sys.velocities, sys.atoms, sys.boundary, + pairwise_inters_nonl, nbs, step_n, sys.energy_units, val_ft) end pairwise_inters_nl = filter(use_neighbors, values(sys.pairwise_inters)) @@ -263,24 +266,31 @@ function potential_energy(sys::System{D, true, T}, neighbors; end if length(neighbors) > 0 nbs = @view neighbors.list[1:neighbors.n] - pe_vec += pairwise_pe_gpu(sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nl, - nbs, step_n, sys.energy_units, val_ft) + pairwise_pe_gpu!(pe_vec_nounits, sys.coords, sys.velocities, sys.atoms, sys.boundary, + pairwise_inters_nl, nbs, step_n, sys.energy_units, val_ft) end end for inter_list in values(sys.specific_inter_lists) - pe_vec += specific_pe_gpu(inter_list, sys.coords, sys.velocities, sys.atoms, sys.boundary, - step_n, sys.energy_units, val_ft) + specific_pe_gpu!(pe_vec_nounits, inter_list, sys.coords, sys.velocities, sys.atoms, + sys.boundary, step_n, sys.energy_units, val_ft) end - pe = Array(pe_vec)[1] + pe = Array(pe_vec_nounits)[1] * sys.energy_units for inter in values(sys.general_inters) - pe += ustrip( + pe += uconvert( sys.energy_units, AtomsCalculators.potential_energy(sys, inter; neighbors=neighbors, n_threads=n_threads), ) end - return pe * sys.energy_units + return pe end + +# Allow GPU-specific potential energy functions to be defined if required +potential_energy_gpu(inter::PairwiseInteraction, dr, ai, aj, eu, sp, ci, cj, bnd, vi, vj, sn) = potential_energy(inter, dr, ai, aj, eu, sp, ci, cj, bnd, vi, vj, sn) +potential_energy_gpu(inter::SpecificInteraction, ci, bnd, ai, eu, vi, sn) = potential_energy(inter, ci, bnd, ai, eu, vi, sn) +potential_energy_gpu(inter::SpecificInteraction, ci, cj, bnd, ai, aj, eu, vi, vj, sn) = potential_energy(inter, ci, cj, bnd, ai, aj, eu, vi, vj, sn) +potential_energy_gpu(inter::SpecificInteraction, ci, cj, ck, bnd, ai, aj, ak, eu, vi, vj, vk, sn) = potential_energy(inter, ci, cj, ck, bnd, ai, aj, ak, eu, vi, vj, vk, sn) +potential_energy_gpu(inter::SpecificInteraction, ci, cj, ck, cl, bnd, ai, aj, ak, al, eu, vi, vj, vk, vl, sn) = potential_energy(inter, ci, cj, ck, cl, bnd, ai, aj, ak, al, eu, vi, vj, vk, vl, sn) diff --git a/src/force.jl b/src/force.jl index 3cf0c9499..c8e6abc5a 100644 --- a/src/force.jl +++ b/src/force.jl @@ -130,7 +130,7 @@ function init_forces_buffer(forces_nounits::CuArray{SVector{D, T}}, n_threads) w end """ - forces(system, neighbors=find_neighbors(sys); n_threads=Threads.nthreads()) + forces(system, neighbors=find_neighbors(sys), step_n=0; n_threads=Threads.nthreads()) Calculate the forces on all atoms in a system using the pairwise, specific and general interactions. @@ -148,7 +148,6 @@ end function forces_nounits!(fs_nounits, sys::System{D, false}, neighbors, fs_chunks=nothing, step_n::Integer=0; n_threads::Integer=Threads.nthreads()) where D - n_atoms = length(sys) pairwise_inters_nonl = filter(!use_neighbors, values(sys.pairwise_inters)) pairwise_inters_nl = filter( use_neighbors, values(sys.pairwise_inters)) sils_1_atoms = filter(il -> il isa InteractionList1Atoms, values(sys.specific_inter_lists)) @@ -158,20 +157,22 @@ function forces_nounits!(fs_nounits, sys::System{D, false}, neighbors, fs_chunks if length(sys.pairwise_inters) > 0 if n_threads > 1 - pairwise_forces_threads!(fs_nounits, fs_chunks, neighbors, sys.atoms, sys.coords, - sys.velocities, sys.boundary, sys.force_units, n_atoms, + pairwise_forces_threads!(fs_nounits, fs_chunks, sys.atoms, sys.coords, sys.velocities, + sys.boundary, neighbors, sys.force_units, length(sys), pairwise_inters_nonl, pairwise_inters_nl, n_threads, step_n) else - pairwise_forces!(fs_nounits, neighbors, sys.atoms, sys.coords, sys.velocities, - sys.boundary, sys.force_units, n_atoms, pairwise_inters_nonl, + pairwise_forces!(fs_nounits, sys.atoms, sys.coords, sys.velocities, sys.boundary, + neighbors, sys.force_units, length(sys), pairwise_inters_nonl, pairwise_inters_nl, step_n) end else fill!(fs_nounits, zero(eltype(fs_nounits))) end - specific_forces!(fs_nounits, sys.atoms, sys.coords, sys.velocities, sys.boundary, - sys.force_units, sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, step_n) + if length(sys.specific_inter_lists) > 0 + specific_forces!(fs_nounits, sys.atoms, sys.coords, sys.velocities, sys.boundary, + sys.force_units, sils_1_atoms, sils_2_atoms, sils_3_atoms, sils_4_atoms, step_n) + end for inter in values(sys.general_inters) fs_gen = AtomsCalculators.forces(sys, inter; neighbors=neighbors, n_threads=n_threads) @@ -182,8 +183,8 @@ function forces_nounits!(fs_nounits, sys::System{D, false}, neighbors, fs_chunks return fs_nounits end -function pairwise_forces!(fs_nounits, neighbors, atoms, coords, velocities, boundary, force_units, n_atoms, - pairwise_inters_nonl, pairwise_inters_nl, step_n=0) +function pairwise_forces!(fs_nounits, atoms, coords, velocities, boundary, neighbors, force_units, + n_atoms, pairwise_inters_nonl, pairwise_inters_nl, step_n=0) fill!(fs_nounits, zero(eltype(fs_nounits))) @inbounds if length(pairwise_inters_nonl) > 0 @@ -227,9 +228,9 @@ function pairwise_forces!(fs_nounits, neighbors, atoms, coords, velocities, boun return fs_nounits end -function pairwise_forces_threads!(fs_nounits, fs_chunks, neighbors, atoms, coords, velocities, boundary, - force_units, n_atoms, pairwise_inters_nonl, pairwise_inters_nl, - n_threads, step_n=0) +function pairwise_forces_threads!(fs_nounits, fs_chunks, atoms, coords, velocities, boundary, + neighbors, force_units, n_atoms, pairwise_inters_nonl, + pairwise_inters_nl, n_threads, step_n=0) if isnothing(fs_chunks) throw(ArgumentError("fs_chunks is not set but n_threads is > 1")) end @@ -351,6 +352,7 @@ function forces_nounits!(fs_nounits, sys::System{D, true, T}, neighbors, n_threads::Integer=Threads.nthreads()) where {D, T} fill!(fs_mat, zero(T)) val_ft = Val(T) + pairwise_inters_nonl = filter(!use_neighbors, values(sys.pairwise_inters)) if length(pairwise_inters_nonl) > 0 nbs = NoNeighborList(length(sys)) From cf148a810c1039252fbb26835747a7ba44dca214 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Tue, 2 Jul 2024 15:38:29 +0100 Subject: [PATCH 19/74] step_n for general interactions --- docs/src/documentation.md | 7 ++++--- src/energy.jl | 6 ++++-- src/force.jl | 6 ++++-- src/types.jl | 6 ++++-- 4 files changed, 16 insertions(+), 9 deletions(-) diff --git a/docs/src/documentation.md b/docs/src/documentation.md index 42ac5a359..033072b98 100644 --- a/docs/src/documentation.md +++ b/docs/src/documentation.md @@ -737,9 +737,10 @@ import AtomsCalculators function AtomsCalculators.forces(sys, inter::MyGeneralInter; neighbors=nothing, + step_n=0, n_threads=Threads.nthreads(), kwargs...) - # kwargs... is required, neighbors and n_threads can be omitted if not used + # kwargs... is required, neighbors/step_n/n_threads can be omitted if not used # Calculate the forces on all atoms using the interaction and the system # The output should have the same shape as the coordinates @@ -930,7 +931,7 @@ function Molly.simulate!(sys, for step_n in 1:n_steps # Calculate accelerations like this - accels_t = accelerations(sys, neighbors; n_threads=n_threads) + accels_t = accelerations(sys, neighbors, step_n; n_threads=n_threads) # Ensure coordinates stay within the simulation box like this sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) @@ -1126,7 +1127,7 @@ Many times, a logger will just record an observation to an `Array` containing a For this purpose, you can use the [`GeneralObservableLogger`](@ref) without defining a custom logging function. Define your observation function as ```julia -function my_observable(sys::System, neighbors; n_threads::Integer, kwargs...) +function my_observable(sys::System, neighbors, step_n; n_threads::Integer, kwargs...) # Probe the system for some desired property return observation end diff --git a/src/energy.jl b/src/energy.jl index 5e96b302e..a75bcaef0 100644 --- a/src/energy.jl +++ b/src/energy.jl @@ -104,7 +104,8 @@ function potential_energy(sys::System{D, false, T}, neighbors; step_n::Integer=0 for inter in values(sys.general_inters) pe += uconvert( sys.energy_units, - AtomsCalculators.potential_energy(sys, inter; neighbors=neighbors, n_threads=n_threads), + AtomsCalculators.potential_energy(sys, inter; neighbors=neighbors, + step_n=step_n, n_threads=n_threads), ) end @@ -281,7 +282,8 @@ function potential_energy(sys::System{D, true, T}, neighbors, step_n::Integer=0; for inter in values(sys.general_inters) pe += uconvert( sys.energy_units, - AtomsCalculators.potential_energy(sys, inter; neighbors=neighbors, n_threads=n_threads), + AtomsCalculators.potential_energy(sys, inter; neighbors=neighbors, + step_n=step_n, n_threads=n_threads), ) end diff --git a/src/force.jl b/src/force.jl index c8e6abc5a..2a99bbeb4 100644 --- a/src/force.jl +++ b/src/force.jl @@ -175,7 +175,8 @@ function forces_nounits!(fs_nounits, sys::System{D, false}, neighbors, fs_chunks end for inter in values(sys.general_inters) - fs_gen = AtomsCalculators.forces(sys, inter; neighbors=neighbors, n_threads=n_threads) + fs_gen = AtomsCalculators.forces(sys, inter; neighbors=neighbors, step_n=step_n, + n_threads=n_threads) check_force_units(zero(eltype(fs_gen)), sys.force_units) fs_nounits .+= ustrip_vec.(fs_gen) end @@ -380,7 +381,8 @@ function forces_nounits!(fs_nounits, sys::System{D, true, T}, neighbors, fs_nounits .= reinterpret(SVector{D, T}, vec(fs_mat)) for inter in values(sys.general_inters) - fs_gen = AtomsCalculators.forces(sys, inter; neighbors=neighbors, n_threads=n_threads) + fs_gen = AtomsCalculators.forces(sys, inter; neighbors=neighbors, step_n=step_n, + n_threads=n_threads) check_force_units(zero(eltype(fs_gen)), sys.force_units) fs_nounits .+= ustrip_vec.(fs_gen) end diff --git a/src/types.jl b/src/types.jl index 2996a2670..7d6291165 100644 --- a/src/types.jl +++ b/src/types.jl @@ -1273,6 +1273,7 @@ AtomsCalculators.@generate_interface function AtomsCalculators.forces( abstract_sys, calc::MollyCalculator; neighbors=nothing, + step_n::Integer=0, n_threads::Integer=Threads.nthreads(), kwargs..., ) @@ -1287,13 +1288,14 @@ AtomsCalculators.@generate_interface function AtomsCalculators.forces( k=calc.k, ) nbs = isnothing(neighbors) ? find_neighbors(sys) : neighbors - return forces(sys, nbs; n_threads=n_threads) + return forces(sys, nbs, step_n; n_threads=n_threads) end AtomsCalculators.@generate_interface function AtomsCalculators.potential_energy( abstract_sys, calc::MollyCalculator; neighbors=nothing, + step_n::Integer=0, n_threads::Integer=Threads.nthreads(), kwargs..., ) @@ -1308,7 +1310,7 @@ AtomsCalculators.@generate_interface function AtomsCalculators.potential_energy( k=calc.k, ) nbs = isnothing(neighbors) ? find_neighbors(sys) : neighbors - return potential_energy(sys, nbs; n_threads=n_threads) + return potential_energy(sys, nbs, step_n; n_threads=n_threads) end """ From 7c91f14aaee14a334d6934016711c51c18116c4f Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Tue, 2 Jul 2024 15:39:31 +0100 Subject: [PATCH 20/74] loggers use step_n --- docs/src/examples.md | 4 ++-- src/loggers.jl | 40 +++++++++++++++++++++++----------------- test/agent.jl | 3 +-- test/simulation.jl | 10 +++++----- 4 files changed, 31 insertions(+), 26 deletions(-) diff --git a/docs/src/examples.md b/docs/src/examples.md index 7bab7cee4..7537e72cc 100644 --- a/docs/src/examples.md +++ b/docs/src/examples.md @@ -189,7 +189,7 @@ function Molly.force(inter::SIRInteraction, end # Custom logger -function fracs_SIR(s::System, neighbors=nothing; n_threads::Integer=Threads.nthreads()) +function fracs_SIR(s::System, args...; kwargs...) counts_sir = [ count(p -> p.status == susceptible, s.atoms), count(p -> p.status == infected , s.atoms), @@ -657,7 +657,7 @@ function Molly.force(inter::BondableInteraction, end end -function bonds(sys::System, neighbors=nothing; n_threads::Integer=Threads.nthreads()) +function bonds(sys::System, args...; kwargs...) bonds = BitVector() for i in 1:length(sys) for j in 1:(i - 1) diff --git a/src/loggers.jl b/src/loggers.jl index 3a9a35ae8..11fa933ca 100644 --- a/src/loggers.jl +++ b/src/loggers.jl @@ -72,7 +72,8 @@ Access the stored observations in a logger. Base.values(logger::GeneralObservableLogger) = logger.history """ - log_property!(logger, system, neighbors=nothing, step_n=0; n_threads=Threads.nthreads(), kwargs...) + log_property!(logger, system, neighbors=nothing, step_n=0; + n_threads=Threads.nthreads(), kwargs...) Log a property of a system throughout a simulation. @@ -82,7 +83,7 @@ Additional keyword arguments can be passed to the logger if required. function log_property!(logger::GeneralObservableLogger, s::System, neighbors=nothing, step_n::Integer=0; kwargs...) if (step_n % logger.n_steps) == 0 - obs = logger.observable(s, neighbors; kwargs...) + obs = logger.observable(s, neighbors, step_n; kwargs...) push!(logger.history, obs) end end @@ -93,7 +94,7 @@ function Base.show(io::IO, gol::GeneralObservableLogger) gol.observable) end -temperature_wrapper(sys, neighbors; kwargs...) = temperature(sys) +temperature_wrapper(sys, args...; kwargs...) = temperature(sys) """ TemperatureLogger(n_steps) @@ -112,7 +113,7 @@ function Base.show(io::IO, tl::GeneralObservableLogger{T, typeof(temperature_wra tl.n_steps, ", ", length(values(tl)), " temperatures recorded") end -coordinates_wrapper(sys, neighbors; kwargs...) = sys.coords +coordinates_wrapper(sys, args...; kwargs...) = sys.coords """ CoordinateLogger(n_steps; dims=3) @@ -136,7 +137,7 @@ function Base.show(io::IO, cl::GeneralObservableLogger{T, typeof(coordinates_wra length(values(cl)) > 0 ? length(first(values(cl))) : "?", " atoms") end -velocities_wrapper(sys, neighbors; kwargs...) = sys.velocities +velocities_wrapper(sys, args...; kwargs...) = sys.velocities """ VelocityLogger(n_steps; dims=3) @@ -160,7 +161,7 @@ function Base.show(io::IO, vl::GeneralObservableLogger{T, typeof(velocities_wrap length(values(vl)) > 0 ? length(first(values(vl))) : "?", " atoms") end -kinetic_energy_wrapper(sys, neighbors; kwargs...) = kinetic_energy(sys) +kinetic_energy_wrapper(sys, args...; kwargs...) = kinetic_energy(sys) """ KineticEnergyLogger(n_steps) @@ -179,10 +180,10 @@ function Base.show(io::IO, el::GeneralObservableLogger{T, typeof(kinetic_energy_ el.n_steps, ", ", length(values(el)), " energies recorded") end -function potential_energy_wrapper(sys, neighbors; n_threads::Integer, +function potential_energy_wrapper(sys, neighbors, step_n::Integer; n_threads::Integer, current_potential_energy=nothing, kwargs...) if isnothing(current_potential_energy) - return potential_energy(sys, neighbors; n_threads=n_threads) + return potential_energy(sys, neighbors, step_n; n_threads=n_threads) else return current_potential_energy end @@ -205,8 +206,8 @@ function Base.show(io::IO, el::GeneralObservableLogger{T, typeof(potential_energ el.n_steps, ", ", length(values(el)), " energies recorded") end -function total_energy_wrapper(sys, neighbors; kwargs...) - return kinetic_energy(sys) + potential_energy_wrapper(sys, neighbors; kwargs...) +function total_energy_wrapper(sys, args...; kwargs...) + return kinetic_energy(sys) + potential_energy_wrapper(sys, args...; kwargs...) end """ @@ -223,9 +224,10 @@ function Base.show(io::IO, el::GeneralObservableLogger{T, typeof(total_energy_wr el.n_steps, ", ", length(values(el)), " energies recorded") end -function forces_wrapper(sys, neighbors; n_threads::Integer, current_forces=nothing, kwargs...) +function forces_wrapper(sys, neighbors, step_n::Integer; n_threads::Integer, + current_forces=nothing, kwargs...) if isnothing(current_forces) - return forces(sys, neighbors; n_threads=n_threads) + return forces(sys, neighbors, step_n; n_threads=n_threads) else return current_forces end @@ -253,7 +255,9 @@ function Base.show(io::IO, fl::GeneralObservableLogger{T, typeof(forces_wrapper) length(values(fl)) > 0 ? length(first(values(fl))) : "?", " atoms") end -virial_wrapper(sys, neighbors; n_threads, kwargs...) = virial(sys, neighbors; n_threads=n_threads) +function virial_wrapper(sys, neighbors, step_n; n_threads, kwargs...) + return virial(sys, neighbors, step_n; n_threads=n_threads) +end """ VirialLogger(n_steps) @@ -273,7 +277,9 @@ function Base.show(io::IO, vl::GeneralObservableLogger{T, typeof(virial_wrapper) vl.n_steps, ", ", length(values(vl)), " virials recorded") end -pressure_wrapper(sys, neighbors; n_threads, kwargs...) = pressure(sys, neighbors; n_threads=n_threads) +function pressure_wrapper(sys, neighbors, step_n; n_threads, kwargs...) + return pressure(sys, neighbors, step_n; n_threads=n_threads) +end """ PressureLogger(n_steps) @@ -460,9 +466,9 @@ end function log_property!(logger::TimeCorrelationLogger, s::System, neighbors=nothing, step_n::Integer=0; n_threads::Integer=Threads.nthreads(), kwargs...) - A = logger.observableA(s, neighbors; n_threads=n_threads, kwargs...) + A = logger.observableA(s, neighbors, step_n; n_threads=n_threads, kwargs...) if logger.observableA != logger.observableB - B = logger.observableB(s, neighbors; n_threads=n_threads, kwargs...) + B = logger.observableB(s, neighbors, step_n; n_threads=n_threads, kwargs...) else B = A end @@ -562,7 +568,7 @@ end function log_property!(aol::AverageObservableLogger{T}, s::System, neighbors=nothing, step_n::Integer=0; kwargs...) where T if (step_n % aol.n_steps) == 0 - obs = aol.observable(s, neighbors; kwargs...) + obs = aol.observable(s, neighbors, step_n; kwargs...) push!(aol.current_block, obs) if length(aol.current_block) == aol.current_block_size diff --git a/test/agent.jl b/test/agent.jl index 75cf6ae97..d33c54af4 100644 --- a/test/agent.jl +++ b/test/agent.jl @@ -51,8 +51,7 @@ end # Test log_property! definition rather than just using GeneralObservableLogger - function Molly.log_property!(logger::SIRLogger, sys, neighbors, step_n; - n_threads=Threads.nthreads(), kwargs...) + function Molly.log_property!(logger::SIRLogger, sys, neighbors, step_n; kwargs...) if step_n % logger.n_steps == 0 counts_sir = [ count(p -> p.status == susceptible, sys.atoms), diff --git a/test/simulation.jl b/test/simulation.jl index f7e12f0fe..de460f5f2 100644 --- a/test/simulation.jl +++ b/test/simulation.jl @@ -4,7 +4,7 @@ temp = 298.0u"K" boundary = RectangularBoundary(2.0u"nm") simulator = VelocityVerlet(dt=0.002u"ps", coupling=AndersenThermostat(temp, 10.0u"ps")) - gen_temp_wrapper(s, neighbors; kwargs...) = temperature(s) + gen_temp_wrapper(s, args...; kwargs...) = temperature(s) s = System( atoms=[Atom(charge=0.0, mass=10.0u"g/mol", σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms], @@ -64,7 +64,7 @@ end TP = typeof(0.2u"kJ * mol^-1") V(sys, args...; kwargs...) = sys.velocities - pot_obs(sys, neighbors; kwargs...) = potential_energy(sys, neighbors) + pot_obs(sys, neighbors, step_n; kwargs...) = potential_energy(sys, neighbors, step_n) kin_obs(sys, args...; kwargs...) = kinetic_energy(sys) for n_threads in n_threads_list @@ -851,7 +851,7 @@ end coords = place_atoms(n_atoms, boundary; min_dist=1.0u"nm") n_log_steps = 500 - box_size_wrapper(sys, neighbors; kwargs...) = sys.boundary.side_lengths[1] + box_size_wrapper(sys, args...; kwargs...) = sys.boundary.side_lengths[1] BoundaryLogger(n_steps) = GeneralObservableLogger(box_size_wrapper, typeof(1.0u"nm"), n_steps) sys = System( @@ -941,7 +941,7 @@ end coords = place_atoms(n_atoms, boundary; min_dist=1.0u"nm") n_log_steps = 500 - box_volume_wrapper(sys, neighbors; kwargs...) = box_volume(sys.boundary) + box_volume_wrapper(sys, args...; kwargs...) = box_volume(sys.boundary) VolumeLogger(n_steps) = GeneralObservableLogger(box_volume_wrapper, typeof(1.0u"nm^3"), n_steps) baro_f(pressure) = MonteCarloAnisotropicBarostat(pressure, temp, boundary) @@ -1013,7 +1013,7 @@ end coords = place_atoms(n_atoms, boundary; min_dist=1.0u"nm") n_log_steps = 500 - box_volume_wrapper(sys, neighbors; kwargs...) = box_volume(sys.boundary) + box_volume_wrapper(sys, args...; kwargs...) = box_volume(sys.boundary) VolumeLogger(n_steps) = GeneralObservableLogger(box_volume_wrapper, typeof(1.0u"nm^3"), n_steps) lang_f(barostat) = Langevin(dt=dt, temperature=temp, friction=friction, coupling=barostat) From 0c3da3f810956dfd9670c66b1d90c38c5c72f8bb Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Tue, 2 Jul 2024 15:40:11 +0100 Subject: [PATCH 21/74] use step_n throughout --- src/coupling.jl | 12 ++++++------ src/energy.jl | 2 +- src/force.jl | 8 ++++---- src/simulators.jl | 40 +++++++++++++++++++++------------------- src/spatial.jl | 12 ++++++------ 5 files changed, 38 insertions(+), 36 deletions(-) diff --git a/src/coupling.jl b/src/coupling.jl index 5c960a94d..804f4cc89 100644 --- a/src/coupling.jl +++ b/src/coupling.jl @@ -196,7 +196,7 @@ function apply_coupling!(sys::System{D, G, T}, barostat::MonteCarloBarostat, sim recompute_forces = false for attempt_n in 1:barostat.n_iterations - E = potential_energy(sys, neighbors; n_threads=n_threads) + E = potential_energy(sys, neighbors, step_n; n_threads=n_threads) V = box_volume(sys.boundary) dV = barostat.volume_scale * (2 * rand(T) - 1) v_scale = (V + dV) / V @@ -213,7 +213,7 @@ function apply_coupling!(sys::System{D, G, T}, barostat::MonteCarloBarostat, sim # This may not be valid for larger changes neighbors_trial = neighbors end - E_trial = potential_energy(sys, neighbors_trial; n_threads=n_threads) + E_trial = potential_energy(sys, neighbors_trial, step_n; n_threads=n_threads) dE = energy_remove_mol(E_trial - E) dW = dE + uconvert(unit(dE), barostat.pressure * dV) - n_molecules * kT * log(v_scale) if dW <= zero(dW) || rand(T) < exp(-dW / kT) @@ -349,7 +349,7 @@ function apply_coupling!(sys::System{D, G, T}, mask1[axis] = true mask2[axis] = false - E = potential_energy(sys, neighbors; n_threads=n_threads) + E = potential_energy(sys, neighbors, step_n; n_threads=n_threads) V = box_volume(sys.boundary) dV = barostat.volume_scale[axis] * (2 * rand(T) - 1) v_scale = (V + dV) / V @@ -366,7 +366,7 @@ function apply_coupling!(sys::System{D, G, T}, # This may not be valid for larger changes neighbors_trial = neighbors end - E_trial = potential_energy(sys, neighbors_trial; n_threads=n_threads) + E_trial = potential_energy(sys, neighbors_trial, step_n; n_threads=n_threads) dE = energy_remove_mol(E_trial - E) dW = dE + uconvert(unit(dE), barostat.pressure[axis] * dV) - n_molecules * kT * log(v_scale) if dW <= zero(dW) || rand(T) < exp(-dW / kT) @@ -521,7 +521,7 @@ function apply_coupling!(sys::System{D, G, T}, axis = 1 end - E = potential_energy(sys, neighbors; n_threads=n_threads) + E = potential_energy(sys, neighbors, step_n; n_threads=n_threads) V = box_volume(sys.boundary) dV = barostat.volume_scale[axis] * (2 * rand(T) - 1) v_scale = (V + dV) / V @@ -557,7 +557,7 @@ function apply_coupling!(sys::System{D, G, T}, # This may not be valid for larger changes neighbors_trial = neighbors end - E_trial = potential_energy(sys, neighbors_trial; n_threads=n_threads) + E_trial = potential_energy(sys, neighbors_trial, step_n; n_threads=n_threads) dE = energy_remove_mol(E_trial - E) PdV = uconvert(unit(dE), barostat.pressure[axis] * dV) γdA = uconvert(unit(dE), barostat.tension * dA) diff --git a/src/energy.jl b/src/energy.jl index a75bcaef0..3ee29d0d4 100644 --- a/src/energy.jl +++ b/src/energy.jl @@ -73,7 +73,7 @@ function potential_energy(sys; n_threads::Integer=Threads.nthreads()) return potential_energy(sys, find_neighbors(sys; n_threads=n_threads); n_threads=n_threads) end -function potential_energy(sys::System{D, false, T}, neighbors; step_n::Integer=0, +function potential_energy(sys::System{D, false, T}, neighbors, step_n::Integer=0; n_threads::Integer=Threads.nthreads()) where {D, T} pairwise_inters_nonl = filter(!use_neighbors, values(sys.pairwise_inters)) pairwise_inters_nl = filter( use_neighbors, values(sys.pairwise_inters)) diff --git a/src/force.jl b/src/force.jl index 2a99bbeb4..6b0042cb8 100644 --- a/src/force.jl +++ b/src/force.jl @@ -11,7 +11,7 @@ export forces """ - accelerations(system, neighbors=find_neighbors(sys); n_threads=Threads.nthreads()) + accelerations(system, neighbors=find_neighbors(sys), step_n=0; n_threads=Threads.nthreads()) Calculate the accelerations of all atoms in a system using the pairwise, specific and general interactions and Newton's second law of motion. @@ -20,8 +20,8 @@ function accelerations(sys; n_threads::Integer=Threads.nthreads()) return accelerations(sys, find_neighbors(sys; n_threads=n_threads); n_threads=n_threads) end -function accelerations(sys, neighbors; n_threads::Integer=Threads.nthreads()) - return forces(sys, neighbors; n_threads=n_threads) ./ masses(sys) +function accelerations(sys, neighbors, step_n::Integer=0; n_threads::Integer=Threads.nthreads()) + return forces(sys, neighbors, step_n; n_threads=n_threads) ./ masses(sys) end """ @@ -368,7 +368,7 @@ function forces_nounits!(fs_nounits, sys::System{D, true, T}, neighbors, end if length(neighbors) > 0 nbs = @view neighbors.list[1:neighbors.n] - pairwise_force_gpu!(fs_mat, sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nonl, + pairwise_force_gpu!(fs_mat, sys.coords, sys.velocities, sys.atoms, sys.boundary, pairwise_inters_nl, nbs, step_n, sys.force_units, val_ft) end end diff --git a/src/simulators.jl b/src/simulators.jl index 2b74e2b75..cf4c996f9 100644 --- a/src/simulators.jl +++ b/src/simulators.jl @@ -73,7 +73,7 @@ function simulate!(sys, hn = sim.step_size for step_n in 1:sim.max_steps - F = forces(sys, neighbors; n_threads=n_threads) + F = forces(sys, neighbors, step_n; n_threads=n_threads) max_force = maximum(norm.(F)) coords_copy = sys.coords @@ -84,7 +84,7 @@ function simulate!(sys, neighbors_copy = neighbors neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n; n_threads=n_threads) - E_trial = potential_energy(sys, neighbors; n_threads=n_threads) + E_trial = potential_energy(sys, neighbors, step_n; n_threads=n_threads) if E_trial < E hn = 6 * hn / 5 E = E_trial @@ -158,7 +158,7 @@ function simulate!(sys, sim.dt; n_threads=n_threads) sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) - forces_t_dt = forces(sys, neighbors; n_threads=n_threads) + forces_t_dt = forces(sys, neighbors, step_n; n_threads=n_threads) accels_t_dt = forces_t_dt ./ masses(sys) sys.velocities += ((accels_t .+ accels_t_dt) .* sim.dt / 2) @@ -173,7 +173,7 @@ function simulate!(sys, neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n, recompute_forces; n_threads=n_threads) if recompute_forces - forces_t = forces(sys, neighbors; n_threads=n_threads) + forces_t = forces(sys, neighbors, step_n; n_threads=n_threads) accels_t = forces_t ./ masses(sys) else forces_t = forces_t_dt @@ -225,7 +225,7 @@ function simulate!(sys, end for step_n in 1:n_steps - forces_t = forces(sys, neighbors; n_threads=n_threads) + forces_t = forces(sys, neighbors, step_n; n_threads=n_threads) accels_t = forces_t ./ masses(sys) sys.velocities += accels_t .* sim.dt @@ -293,7 +293,7 @@ function simulate!(sys, end for step_n in 1:n_steps - forces_t = forces(sys, neighbors; n_threads=n_threads) + forces_t = forces(sys, neighbors, step_n; n_threads=n_threads) accels_t = forces_t ./ masses(sys) coords_copy = sys.coords @@ -378,7 +378,7 @@ function simulate!(sys, end for step_n in 1:n_steps - forces_t = forces(sys, neighbors; n_threads=n_threads) + forces_t = forces(sys, neighbors, step_n; n_threads=n_threads) accels_t = forces_t ./ masses(sys) sys.velocities += accels_t .* sim.dt @@ -507,7 +507,7 @@ function simulate!(sys, for step_n in 1:n_steps for (step!, args) in step_arg_pairs - step!(args..., neighbors) + step!(args..., neighbors, step_n) end sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) @@ -523,22 +523,22 @@ function simulate!(sys, return sys end -function O_step!(sys, α_eff, σ_eff, rng, temperature, neighbors) +function O_step!(sys, α_eff, σ_eff, rng, temperature, neighbors, step_n) noise = random_velocities(sys, temperature; rng=rng) sys.velocities = α_eff .* sys.velocities + σ_eff .* noise return sys end -function A_step!(sys, dt_eff, neighbors) +function A_step!(sys, dt_eff, neighbors, step_n) sys.coords += sys.velocities * dt_eff sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) return sys end function B_step!(sys, dt_eff, acceleration_vector, compute_forces::Bool, - n_threads::Integer, neighbors) + n_threads::Integer, neighbors, step_n::Integer) if compute_forces - acceleration_vector .= accelerations(sys, neighbors; n_threads=n_threads) + acceleration_vector .= accelerations(sys, neighbors, step_n; n_threads=n_threads) end sys.velocities += dt_eff * acceleration_vector return sys @@ -586,7 +586,7 @@ function simulate!(sys, run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) for step_n in 1:n_steps - forces_t = forces(sys, neighbors; n_threads=n_threads) + forces_t = forces(sys, neighbors, step_n; n_threads=n_threads) accels_t = forces_t ./ masses(sys) noise = random_velocities(sys, sim.temperature; rng=rng) @@ -666,7 +666,7 @@ function simulate!(sys, sim::NoseHoover, n_steps::Integer; T_half = uconvert(unit(sim.temperature), 2 * KE_half / (sys.df * sys.k)) zeta = zeta_half + (sim.dt / (2 * (sim.damping^2))) * ((T_half / sim.temperature) - 1) - forces_t_dt = forces(sys, neighbors; n_threads=n_threads) + forces_t_dt = forces(sys, neighbors, step_n; n_threads=n_threads) accels_t_dt = forces_t_dt ./ masses(sys) sys.velocities = (v_half .+ accels_t_dt .* (sim.dt / 2)) ./ @@ -681,7 +681,7 @@ function simulate!(sys, sim::NoseHoover, n_steps::Integer; neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n, recompute_forces; n_threads=n_threads) if recompute_forces - forces_t = forces(sys, neighbors; n_threads=n_threads) + forces_t = forces(sys, neighbors, step_n; n_threads=n_threads) accels_t = forces_t ./ masses(sys) else forces_t = forces_t_dt @@ -906,6 +906,8 @@ end n_threads=Threads.nthreads(), run_loggers=true) Run a REMD simulation on a [`ReplicaSystem`](@ref) using a REMD simulator. + +Not currently compatible with interactions that depend on step number. """ function simulate_remd!(sys::ReplicaSystem, remd_sim, @@ -998,22 +1000,22 @@ function simulate!(sys::System{D, G, T}, run_loggers=true) where {D, G, T} neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) E_old = potential_energy(sys, neighbors; n_threads=n_threads) - for i in 1:n_steps + for step_n in 1:n_steps coords_old = copy(sys.coords) sim.trial_moves(sys; sim.trial_args...) # Changes the coordinates of the system neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) - E_new = potential_energy(sys, neighbors; n_threads=n_threads) + E_new = potential_energy(sys, neighbors, step_n; n_threads=n_threads) ΔE = E_new - E_old δ = ΔE / (sys.k * sim.temperature) if δ < 0 || (rand() < exp(-δ)) - run_loggers!(sys, neighbors, i, run_loggers; n_threads=n_threads, + run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, current_potential_energy=E_new, success=true, energy_rate=(E_new / (sys.k * sim.temperature))) E_old = E_new else sys.coords = coords_old - run_loggers!(sys, neighbors, i, run_loggers; n_threads=n_threads, + run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, current_potential_energy=E_old, success=false, energy_rate=(E_old / (sys.k * sim.temperature))) end diff --git a/src/spatial.jl b/src/spatial.jl index 488af9fce..7b183a09f 100644 --- a/src/spatial.jl +++ b/src/spatial.jl @@ -679,7 +679,7 @@ end @doc raw""" virial(sys, neighbors=find_neighbors(sys), step_n=0; n_threads=Threads.nthreads()) - virial(inter, sys, neighbors; step_n=0, n_threads=Threads.nthreads()) + virial(inter, sys, neighbors, step_n; n_threads=Threads.nthreads()) Calculate the virial of a system or the virial resulting from a general interaction. @@ -705,7 +705,7 @@ function virial(sys, neighbors, step_n::Integer=0; n_threads::Integer=Threads.nt v = virial(sys, neighbors, step_n, pairwise_inters_nonl, pairwise_inters_nl) for inter in values(sys.general_inters) - v += virial(inter, sys, neighbors; step_n=step_n, n_threads=n_threads) + v += virial(inter, sys, neighbors, step_n; n_threads=n_threads) end return v @@ -765,12 +765,12 @@ function virial(sys::System{D, G, T}, neighbors_dev, step_n, pairwise_inters_non end # Default for general interactions -function virial(inter, sys::System{D, G, T}, neighbors=nothing; kwargs...) where {D, G, T} +function virial(inter, sys::System{D, G, T}, args...; kwargs...) where {D, G, T} return zero(T) * sys.energy_units end @doc raw""" - pressure(sys, neighbors=find_neighbors(sys); n_threads=Threads.nthreads()) + pressure(sys, neighbors=find_neighbors(sys), step_n=0; n_threads=Threads.nthreads()) Calculate the pressure of a system. @@ -793,13 +793,13 @@ function pressure(sys; n_threads::Integer=Threads.nthreads()) return pressure(sys, find_neighbors(sys; n_threads=n_threads); n_threads=n_threads) end -function pressure(sys::AtomsBase.AbstractSystem{D}, neighbors; +function pressure(sys::AtomsBase.AbstractSystem{D}, neighbors, step_n::Integer=0; n_threads::Integer=Threads.nthreads()) where D if has_infinite_boundary(sys.boundary) error("pressure calculation not compatible with infinite boundaries") end NkT = energy_remove_mol(length(sys) * sys.k * temperature(sys)) - vir = energy_remove_mol(virial(sys, neighbors; n_threads=n_threads)) + vir = energy_remove_mol(virial(sys, neighbors, step_n; n_threads=n_threads)) P = (NkT - (2 * vir) / D) / box_volume(sys.boundary) if sys.energy_units == NoUnits || D != 3 # If implied energy units are (u * nm^2 * ps^-2) and everything is From 60bef871f8063388c289f5b36ed7c2abe07fba3f Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Tue, 2 Jul 2024 17:37:57 +0100 Subject: [PATCH 22/74] reuse memory in simulators --- docs/src/documentation.md | 1 + src/loggers.jl | 6 +- src/simulators.jl | 208 +++++++++++++++++++++++--------------- 3 files changed, 129 insertions(+), 86 deletions(-) diff --git a/docs/src/documentation.md b/docs/src/documentation.md index 033072b98..042893434 100644 --- a/docs/src/documentation.md +++ b/docs/src/documentation.md @@ -613,6 +613,7 @@ Atom properties can be accessed, e.g. `atom_i.σ`. `force_units` can be useful for returning a zero force under certain conditions. `step_n` is the step number in the simulator, allowing time-dependent interactions. Beware that this step counter starts from 1 every time [`simulate!`](@ref) is called. +It also doesn't work with [`simulate_remd!`](@ref). Typically the force function is where most computation time is spent during the simulation, so consider optimising this function if you want high performance. One nice feature of Molly is that this function will work on both the CPU and the GPU. diff --git a/src/loggers.jl b/src/loggers.jl index 11fa933ca..b3a278c3a 100644 --- a/src/loggers.jl +++ b/src/loggers.jl @@ -113,7 +113,7 @@ function Base.show(io::IO, tl::GeneralObservableLogger{T, typeof(temperature_wra tl.n_steps, ", ", length(values(tl)), " temperatures recorded") end -coordinates_wrapper(sys, args...; kwargs...) = sys.coords +coordinates_wrapper(sys, args...; kwargs...) = copy(sys.coords) """ CoordinateLogger(n_steps; dims=3) @@ -137,7 +137,7 @@ function Base.show(io::IO, cl::GeneralObservableLogger{T, typeof(coordinates_wra length(values(cl)) > 0 ? length(first(values(cl))) : "?", " atoms") end -velocities_wrapper(sys, args...; kwargs...) = sys.velocities +velocities_wrapper(sys, args...; kwargs...) = copy(sys.velocities) """ VelocityLogger(n_steps; dims=3) @@ -229,7 +229,7 @@ function forces_wrapper(sys, neighbors, step_n::Integer; n_threads::Integer, if isnothing(current_forces) return forces(sys, neighbors, step_n; n_threads=n_threads) else - return current_forces + return copy(current_forces) end end diff --git a/src/simulators.jl b/src/simulators.jl index cf4c996f9..8f8b0c63e 100644 --- a/src/simulators.jl +++ b/src/simulators.jl @@ -63,23 +63,28 @@ function simulate!(sys, sim::SteepestDescentMinimizer; n_threads::Integer=Threads.nthreads(), run_loggers=false) - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) E = potential_energy(sys, neighbors; n_threads=n_threads) run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads, current_potential_energy=E) using_constraints = length(sys.constraints) > 0 - println(sim.log_stream, "Step 0 - potential energy ", - E, " - max force N/A - N/A") + println(sim.log_stream, "Step 0 - potential energy ", E, " - max force N/A - N/A") hn = sim.step_size + coords_copy = similar(sys.coords) + F_nounits = ustrip_vec.(similar(sys.coords)) + F = F_nounits .* sys.force_units + forces_buffer = init_forces_buffer(F_nounits, n_threads) for step_n in 1:sim.max_steps - F = forces(sys, neighbors, step_n; n_threads=n_threads) + F_nounits .= forces_nounits!(F_nounits, sys, neighbors, forces_buffer, step_n; + n_threads=n_threads) + F .= F_nounits .* sys.force_units max_force = maximum(norm.(F)) - coords_copy = sys.coords - sys.coords += hn * F ./ max_force + coords_copy .= sys.coords + sys.coords .+= hn .* F ./ max_force using_constraints && apply_position_constraints!(sys, coords_copy; n_threads=n_threads) - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) neighbors_copy = neighbors neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n; @@ -91,7 +96,7 @@ function simulate!(sys, println(sim.log_stream, "Step ", step_n, " - potential energy ", E_trial, " - max force ", max_force, " - accepted") else - sys.coords = coords_copy + sys.coords .= coords_copy neighbors = neighbors_copy hn = hn / 5 println(sim.log_stream, "Step ", step_n, " - potential energy ", @@ -134,12 +139,14 @@ function simulate!(sys, n_steps::Integer; n_threads::Integer=Threads.nthreads(), run_loggers=true) - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) !iszero(sim.remove_CM_motion) && remove_CM_motion!(sys) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) forces_t = forces(sys, neighbors; n_threads=n_threads) - forces_t_dt = zero(forces_t) accels_t = forces_t ./ masses(sys) + forces_nounits_t_dt = ustrip_vec.(similar(sys.coords)) + forces_t_dt = forces_nounits_t_dt .* sys.force_units + forces_buffer = init_forces_buffer(forces_nounits_t_dt, n_threads) accels_t_dt = zero(accels_t) run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads, current_forces=forces_t) using_constraints = length(sys.constraints) > 0 @@ -153,15 +160,17 @@ function simulate!(sys, cons_coord_storage .= sys.coords end - sys.coords += sys.velocities .* sim.dt .+ ((accels_t .* sim.dt ^ 2) ./ 2) + sys.coords .+= sys.velocities .* sim.dt .+ ((accels_t .* sim.dt .^ 2) ./ 2) using_constraints && apply_position_constraints!(sys, cons_coord_storage, cons_vel_storage, sim.dt; n_threads=n_threads) - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) - forces_t_dt = forces(sys, neighbors, step_n; n_threads=n_threads) - accels_t_dt = forces_t_dt ./ masses(sys) + forces_nounits_t_dt .= forces_nounits!(forces_nounits_t_dt, sys, neighbors, forces_buffer, + step_n; n_threads=n_threads) + forces_t_dt .= forces_nounits_t_dt .* sys.force_units + accels_t_dt .= forces_t_dt ./ masses(sys) - sys.velocities += ((accels_t .+ accels_t_dt) .* sim.dt / 2) + sys.velocities .+= ((accels_t .+ accels_t_dt) .* sim.dt ./ 2) using_constraints && apply_velocity_constraints!(sys; n_threads=n_threads) if !iszero(sim.remove_CM_motion) && step_n % sim.remove_CM_motion == 0 @@ -173,11 +182,13 @@ function simulate!(sys, neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n, recompute_forces; n_threads=n_threads) if recompute_forces - forces_t = forces(sys, neighbors, step_n; n_threads=n_threads) - accels_t = forces_t ./ masses(sys) + forces_nounits_t_dt .= forces_nounits!(forces_nounits_t_dt, sys, neighbors, + forces_buffer, step_n; n_threads=n_threads) + forces_t .= forces_nounits_t_dt .* sys.force_units + accels_t .= forces_t ./ masses(sys) else - forces_t = forces_t_dt - accels_t = accels_t_dt + forces_t .= forces_t_dt + accels_t .= accels_t_dt end run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, @@ -215,32 +226,39 @@ function simulate!(sys, n_steps::Integer; n_threads::Integer=Threads.nthreads(), run_loggers=true) - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) !iszero(sim.remove_CM_motion) && remove_CM_motion!(sys) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) + forces_nounits_t = ustrip_vec.(similar(sys.coords)) + forces_t = forces_nounits_t .* sys.force_units + forces_buffer = init_forces_buffer(forces_nounits_t, n_threads) + accels_t = forces_t ./ masses(sys) using_constraints = length(sys.constraints) > 0 if using_constraints cons_coord_storage = similar(sys.coords) end for step_n in 1:n_steps - forces_t = forces(sys, neighbors, step_n; n_threads=n_threads) - accels_t = forces_t ./ masses(sys) + forces_nounits_t .= forces_nounits!(forces_nounits_t, sys, neighbors, forces_buffer, step_n; + n_threads=n_threads) + forces_t .= forces_nounits_t .* sys.force_units + accels_t .= forces_t ./ masses(sys) - sys.velocities += accels_t .* sim.dt + sys.velocities .+= accels_t .* sim.dt if using_constraints cons_coord_storage .= sys.coords end - sys.coords += sys.velocities .* sim.dt - using_constraints && apply_position_constraints!(sys, cons_coord_storage; n_threads=n_threads) + sys.coords .+= sys.velocities .* sim.dt + using_constraints && apply_position_constraints!(sys, cons_coord_storage; + n_threads=n_threads) if using_constraints sys.velocities .= (sys.coords .- cons_coord_storage) ./ sim.dt end - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) if !iszero(sim.remove_CM_motion) && step_n % sim.remove_CM_motion == 0 remove_CM_motion!(sys) @@ -283,44 +301,43 @@ function simulate!(sys, n_steps::Integer; n_threads::Integer=Threads.nthreads(), run_loggers=true) - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) - coords_last = sys.coords + coords_last, coords_copy = similar(sys.coords), similar(sys.coords) + forces_nounits_t = ustrip_vec.(similar(sys.coords)) + forces_t = forces_nounits_t .* sys.force_units + forces_buffer = init_forces_buffer(forces_nounits_t, n_threads) + accels_t = forces_t ./ masses(sys) using_constraints = length(sys.constraints) > 0 - if using_constraints - cons_coord_storage = similar(sys.coords) - end for step_n in 1:n_steps - forces_t = forces(sys, neighbors, step_n; n_threads=n_threads) - accels_t = forces_t ./ masses(sys) - - coords_copy = sys.coords - if using_constraints - cons_coord_storage .= sys.coords - end + forces_nounits_t .= forces_nounits!(forces_nounits_t, sys, neighbors, forces_buffer, step_n; + n_threads=n_threads) + forces_t .= forces_nounits_t .* sys.force_units + accels_t .= forces_t ./ masses(sys) + coords_copy .= sys.coords if step_n == 1 # Use the velocities at the first step since there is only one set of coordinates - sys.coords += sys.velocities .* sim.dt .+ (accels_t .* sim.dt ^ 2) ./ 2 + sys.coords .+= sys.velocities .* sim.dt .+ (accels_t .* sim.dt .^ 2) ./ 2 else - sys.coords += vector.(coords_last, sys.coords, (sys.boundary,)) .+ - accels_t .* sim.dt ^ 2 + sys.coords .+= vector.(coords_last, sys.coords, (sys.boundary,)) .+ + accels_t .* sim.dt .^ 2 end - using_constraints && apply_position_constraints!(sys, cons_coord_storage; n_threads=n_threads) + using_constraints && apply_position_constraints!(sys, coords_copy; n_threads=n_threads) - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) # This is accurate to O(dt) - sys.velocities = vector.(coords_copy, sys.coords, (sys.boundary,)) ./ sim.dt + sys.velocities .= vector.(coords_copy, sys.coords, (sys.boundary,)) ./ sim.dt recompute_forces = apply_coupling!(sys, sim.coupling, sim, neighbors, step_n; n_threads=n_threads) neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n, recompute_forces; n_threads=n_threads) - coords_last = coords_copy + coords_last .= coords_copy run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, current_forces=forces_t) @@ -367,10 +384,15 @@ function simulate!(sys, n_threads::Integer=Threads.nthreads(), run_loggers=true, rng=Random.GLOBAL_RNG) - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) !iszero(sim.remove_CM_motion) && remove_CM_motion!(sys) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) + forces_nounits_t = ustrip_vec.(similar(sys.coords)) + forces_t = forces_nounits_t .* sys.force_units + forces_buffer = init_forces_buffer(forces_nounits_t, n_threads) + accels_t = forces_t ./ masses(sys) + noise = similar(sys.velocities) using_constraints = length(sys.constraints) > 0 if using_constraints cons_coord_storage = similar(sys.coords) @@ -378,24 +400,27 @@ function simulate!(sys, end for step_n in 1:n_steps - forces_t = forces(sys, neighbors, step_n; n_threads=n_threads) - accels_t = forces_t ./ masses(sys) + forces_nounits_t .= forces_nounits!(forces_nounits_t, sys, neighbors, forces_buffer, step_n; + n_threads=n_threads) + forces_t .= forces_nounits_t .* sys.force_units + accels_t .= forces_t ./ masses(sys) - sys.velocities += accels_t .* sim.dt + sys.velocities .+= accels_t .* sim.dt apply_velocity_constraints!(sys; n_threads=n_threads) if using_constraints cons_coord_storage .= sys.coords end - sys.coords += sys.velocities .* sim.dt / 2 + sys.coords .+= sys.velocities .* sim.dt ./ 2 - noise = random_velocities(sys, sim.temperature; rng=rng) - sys.velocities = sys.velocities .* sim.vel_scale .+ noise .* sim.noise_scale + noise .= random_velocities(sys, sim.temperature; rng=rng) + sys.velocities .= sys.velocities .* sim.vel_scale .+ noise .* sim.noise_scale - sys.coords += sys.velocities .* sim.dt / 2 + sys.coords .+= sys.velocities .* sim.dt ./ 2 - using_constraints && apply_position_constraints!(sys, cons_coord_storage, cons_vel_storage, sim.dt; n_threads=n_threads) - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + using_constraints && apply_position_constraints!(sys, cons_coord_storage, cons_vel_storage, + sim.dt; n_threads=n_threads) + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) if !iszero(sim.remove_CM_motion) && step_n % sim.remove_CM_motion == 0 remove_CM_motion!(sys) @@ -468,7 +493,7 @@ function simulate!(sys, α_eff = exp.(-sim.friction * sim.dt .* M_inv / count('O', sim.splitting)) σ_eff = sqrt.((1 * unit(eltype(α_eff))) .- (α_eff .^ 2)) - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) !iszero(sim.remove_CM_motion) && remove_CM_motion!(sys) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) @@ -510,7 +535,7 @@ function simulate!(sys, step!(args..., neighbors, step_n) end - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) if !iszero(sim.remove_CM_motion) && step_n % sim.remove_CM_motion == 0 remove_CM_motion!(sys) end @@ -525,13 +550,13 @@ end function O_step!(sys, α_eff, σ_eff, rng, temperature, neighbors, step_n) noise = random_velocities(sys, temperature; rng=rng) - sys.velocities = α_eff .* sys.velocities + σ_eff .* noise + sys.velocities .= α_eff .* sys.velocities .+ σ_eff .* noise return sys end function A_step!(sys, dt_eff, neighbors, step_n) - sys.coords += sys.velocities * dt_eff - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + sys.coords .+= sys.velocities .* dt_eff + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) return sys end @@ -540,7 +565,7 @@ function B_step!(sys, dt_eff, acceleration_vector, compute_forces::Bool, if compute_forces acceleration_vector .= accelerations(sys, neighbors, step_n; n_threads=n_threads) end - sys.velocities += dt_eff * acceleration_vector + sys.velocities .+= dt_eff .* acceleration_vector return sys end @@ -580,18 +605,25 @@ function simulate!(sys, @warn "OverdampedLangevin is not currently compatible with constraints, " * "constraints will be ignored" end - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) !iszero(sim.remove_CM_motion) && remove_CM_motion!(sys) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) + forces_nounits_t = ustrip_vec.(similar(sys.coords)) + forces_t = forces_nounits_t .* sys.force_units + forces_buffer = init_forces_buffer(forces_nounits_t, n_threads) + accels_t = forces_t ./ masses(sys) + noise = similar(sys.velocities) for step_n in 1:n_steps - forces_t = forces(sys, neighbors, step_n; n_threads=n_threads) - accels_t = forces_t ./ masses(sys) + forces_nounits_t .= forces_nounits!(forces_nounits_t, sys, neighbors, forces_buffer, step_n; + n_threads=n_threads) + forces_t .= forces_nounits_t .* sys.force_units + accels_t .= forces_t ./ masses(sys) - noise = random_velocities(sys, sim.temperature; rng=rng) - sys.coords += (accels_t ./ sim.friction) .* sim.dt .+ sqrt((2 / sim.friction) * sim.dt) .* noise - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + noise .= random_velocities(sys, sim.temperature; rng=rng) + sys.coords .+= (accels_t ./ sim.friction) .* sim.dt .+ sqrt((2 / sim.friction) * sim.dt) .* noise + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) if !iszero(sim.remove_CM_motion) && step_n % sim.remove_CM_motion == 0 remove_CM_motion!(sys) @@ -644,33 +676,38 @@ function simulate!(sys, sim::NoseHoover, n_steps::Integer; @warn "NoseHoover is not currently compatible with constraints, " * "constraints will be ignored" end - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) !iszero(sim.remove_CM_motion) && remove_CM_motion!(sys) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) forces_t = forces(sys, neighbors; n_threads=n_threads) - forces_t_dt = zero(forces_t) accels_t = forces_t ./ masses(sys) + forces_nounits_t_dt = ustrip_vec.(similar(sys.coords)) + forces_t_dt = forces_nounits_t_dt .* sys.force_units + forces_buffer = init_forces_buffer(forces_nounits_t_dt, n_threads) accels_t_dt = zero(accels_t) run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads, current_forces=forces_t) v_half = zero(sys.velocities) zeta = zero(inv(sim.dt)) for step_n in 1:n_steps - v_half = sys.velocities .+ (accels_t .- (sys.velocities .* zeta)) .* (sim.dt / 2) + v_half .= sys.velocities .+ (accels_t .- (sys.velocities .* zeta)) .* (sim.dt ./ 2) - sys.coords += v_half .* sim.dt - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) + sys.coords .+= v_half .* sim.dt + sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) - zeta_half = zeta + (sim.dt / (2 * (sim.damping^2))) * ((temperature(sys) / sim.temperature) - 1) + zeta_half = zeta + (sim.dt / (2 * (sim.damping^2))) * + ((temperature(sys) / sim.temperature) - 1) KE_half = sum(masses(sys) .* sum.(abs2, v_half)) / 2 T_half = uconvert(unit(sim.temperature), 2 * KE_half / (sys.df * sys.k)) zeta = zeta_half + (sim.dt / (2 * (sim.damping^2))) * ((T_half / sim.temperature) - 1) - forces_t_dt = forces(sys, neighbors, step_n; n_threads=n_threads) - accels_t_dt = forces_t_dt ./ masses(sys) + forces_nounits_t_dt .= forces_nounits!(forces_nounits_t_dt, sys, neighbors, forces_buffer, + step_n; n_threads=n_threads) + forces_t_dt .= forces_nounits_t_dt .* sys.force_units + accels_t_dt .= forces_t_dt ./ masses(sys) - sys.velocities = (v_half .+ accels_t_dt .* (sim.dt / 2)) ./ - (1 + (zeta * sim.dt / 2)) + sys.velocities .= (v_half .+ accels_t_dt .* (sim.dt / 2)) ./ + (1 + (zeta * sim.dt / 2)) if !iszero(sim.remove_CM_motion) && step_n % sim.remove_CM_motion == 0 remove_CM_motion!(sys) @@ -681,11 +718,13 @@ function simulate!(sys, sim::NoseHoover, n_steps::Integer; neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n, recompute_forces; n_threads=n_threads) if recompute_forces - forces_t = forces(sys, neighbors, step_n; n_threads=n_threads) - accels_t = forces_t ./ masses(sys) + forces_nounits_t_dt .= forces_nounits!(forces_nounits_t_dt, sys, neighbors, + forces_buffer, step_n; n_threads=n_threads) + forces_t .= forces_nounits_t_dt .* sys.force_units + accels_t .= forces_t ./ masses(sys) else - forces_t = forces_t_dt - accels_t = accels_t_dt + forces_t .= forces_t_dt + accels_t .= accels_t_dt end run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, @@ -1000,8 +1039,10 @@ function simulate!(sys::System{D, G, T}, run_loggers=true) where {D, G, T} neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) E_old = potential_energy(sys, neighbors; n_threads=n_threads) + coords_old = similar(sys.coords) + for step_n in 1:n_steps - coords_old = copy(sys.coords) + coords_old .= sys.coords sim.trial_moves(sys; sim.trial_args...) # Changes the coordinates of the system neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) E_new = potential_energy(sys, neighbors, step_n; n_threads=n_threads) @@ -1014,12 +1055,13 @@ function simulate!(sys::System{D, G, T}, energy_rate=(E_new / (sys.k * sim.temperature))) E_old = E_new else - sys.coords = coords_old + sys.coords .= coords_old run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, current_potential_energy=E_old, success=false, energy_rate=(E_old / (sys.k * sim.temperature))) end end + return sys end From 913140fbbfdb16a21f299fd4e97a5dd118f585c3 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 5 Jul 2024 14:44:26 +0100 Subject: [PATCH 23/74] reuse memory in simulators and couplers --- src/coupling.jl | 21 ++++++++++++--------- src/simulators.jl | 38 +++++++++++++++++++++++--------------- src/spatial.jl | 23 +++++++++++++++++------ 3 files changed, 52 insertions(+), 30 deletions(-) diff --git a/src/coupling.jl b/src/coupling.jl index 804f4cc89..24962b195 100644 --- a/src/coupling.jl +++ b/src/coupling.jl @@ -77,7 +77,7 @@ function apply_coupling!(sys::System{D, true, T}, thermostat::AndersenThermostat atoms_to_bump_dev = move_array(atoms_to_bump, sys) atoms_to_leave_dev = move_array(atoms_to_leave, sys) vs = random_velocities(sys, thermostat.temperature) - sys.velocities = sys.velocities .* atoms_to_leave_dev .+ vs .* atoms_to_bump_dev + sys.velocities .= sys.velocities .* atoms_to_leave_dev .+ vs .* atoms_to_bump_dev return false end @@ -101,7 +101,7 @@ end function apply_coupling!(sys, thermostat::RescaleThermostat, sim, neighbors=nothing, step_n::Integer=0; n_threads::Integer=Threads.nthreads()) - sys.velocities *= sqrt(thermostat.temperature / temperature(sys)) + sys.velocities .*= sqrt(thermostat.temperature / temperature(sys)) return false end @@ -126,7 +126,7 @@ end function apply_coupling!(sys, thermostat::BerendsenThermostat, sim, neighbors=nothing, step_n::Integer=0; n_threads::Integer=Threads.nthreads()) λ2 = 1 + (sim.dt / thermostat.coupling_const) * ((thermostat.temperature / temperature(sys)) - 1) - sys.velocities *= sqrt(λ2) + sys.velocities .*= sqrt(λ2) return false end @@ -194,6 +194,7 @@ function apply_coupling!(sys::System{D, G, T}, barostat::MonteCarloBarostat, sim kT = energy_remove_mol(sys.k * barostat.temperature) n_molecules = isnothing(sys.topology) ? length(sys) : length(sys.topology.molecule_atom_counts) recompute_forces = false + old_coords = similar(sys.coords) for attempt_n in 1:barostat.n_iterations E = potential_energy(sys, neighbors, step_n; n_threads=n_threads) @@ -201,7 +202,7 @@ function apply_coupling!(sys::System{D, G, T}, barostat::MonteCarloBarostat, sim dV = barostat.volume_scale * (2 * rand(T) - 1) v_scale = (V + dV) / V l_scale = (D == 2 ? sqrt(v_scale) : cbrt(v_scale)) - old_coords = copy(sys.coords) + old_coords .= sys.coords old_boundary = sys.boundary scale_coords!(sys, l_scale) @@ -220,7 +221,7 @@ function apply_coupling!(sys::System{D, G, T}, barostat::MonteCarloBarostat, sim recompute_forces = true barostat.n_accepted += 1 else - sys.coords = old_coords + sys.coords .= old_coords sys.boundary = old_boundary end barostat.n_attempted += 1 @@ -337,6 +338,7 @@ function apply_coupling!(sys::System{D, G, T}, kT = energy_remove_mol(sys.k * barostat.temperature) n_molecules = isnothing(sys.topology) ? length(sys) : length(sys.topology.molecule_atom_counts) recompute_forces = false + old_coords = similar(sys.coords) for attempt_n in 1:barostat.n_iterations axis = 0 @@ -354,7 +356,7 @@ function apply_coupling!(sys::System{D, G, T}, dV = barostat.volume_scale[axis] * (2 * rand(T) - 1) v_scale = (V + dV) / V l_scale = SVector{D}(mask1 * v_scale + mask2) - old_coords = copy(sys.coords) + old_coords .= sys.coords old_boundary = sys.boundary scale_coords!(sys, l_scale) @@ -373,7 +375,7 @@ function apply_coupling!(sys::System{D, G, T}, recompute_forces = true barostat.n_accepted[axis] += 1 else - sys.coords = old_coords + sys.coords .= old_coords sys.boundary = old_boundary end barostat.n_attempted[axis] += 1 @@ -510,6 +512,7 @@ function apply_coupling!(sys::System{D, G, T}, kT = energy_remove_mol(sys.k * barostat.temperature) n_molecules = isnothing(sys.topology) ? length(sys) : length(sys.topology.molecule_atom_counts) recompute_forces = false + old_coords = similar(sys.coords) for attempt_n in 1:barostat.n_iterations axis = 0 @@ -545,7 +548,7 @@ function apply_coupling!(sys::System{D, G, T}, dA = sys.boundary[1] * sys.boundary[2] * (l_scale[1] * l_scale[2] - one(T)) - old_coords = copy(sys.coords) + old_coords .= sys.coords old_boundary = sys.boundary scale_coords!(sys, l_scale) @@ -566,7 +569,7 @@ function apply_coupling!(sys::System{D, G, T}, recompute_forces = true barostat.n_accepted[axis] += 1 else - sys.coords = old_coords + sys.coords .= old_coords sys.boundary = old_boundary end barostat.n_attempted[axis] += 1 diff --git a/src/simulators.jl b/src/simulators.jl index 8f8b0c63e..a3afc7731 100644 --- a/src/simulators.jl +++ b/src/simulators.jl @@ -413,7 +413,7 @@ function simulate!(sys, end sys.coords .+= sys.velocities .* sim.dt ./ 2 - noise .= random_velocities(sys, sim.temperature; rng=rng) + random_velocities!(noise, sys, sim.temperature; rng=rng) sys.velocities .= sys.velocities .* sim.vel_scale .+ noise .* sim.noise_scale sys.coords .+= sys.velocities .* sim.dt ./ 2 @@ -497,7 +497,11 @@ function simulate!(sys, !iszero(sim.remove_CM_motion) && remove_CM_motion!(sys) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) - accels_t = accelerations(sys, neighbors; n_threads=n_threads) + forces_nounits_t = ustrip_vec.(similar(sys.coords)) + forces_t = forces_nounits_t .* sys.force_units + forces_buffer = init_forces_buffer(forces_nounits_t, n_threads) + accels_t = forces_t ./ masses(sys) + noise = similar(sys.velocities) effective_dts = [sim.dt / count(c, sim.splitting) for c in sim.splitting] @@ -524,9 +528,10 @@ function simulate!(sys, if op == 'A' return (A_step!, (sys, effective_dts[j])) elseif op == 'B' - return (B_step!, (sys, effective_dts[j], accels_t, force_computation_steps[j], n_threads)) + return (B_step!, (sys, forces_nounits_t, forces_t, forces_buffer, accels_t, + effective_dts[j], force_computation_steps[j], n_threads)) elseif op == 'O' - return (O_step!, (sys, α_eff, σ_eff, rng, sim.temperature)) + return (O_step!, (sys, noise, α_eff, σ_eff, rng, sim.temperature)) end end @@ -548,24 +553,27 @@ function simulate!(sys, return sys end -function O_step!(sys, α_eff, σ_eff, rng, temperature, neighbors, step_n) - noise = random_velocities(sys, temperature; rng=rng) - sys.velocities .= α_eff .* sys.velocities .+ σ_eff .* noise - return sys -end - function A_step!(sys, dt_eff, neighbors, step_n) sys.coords .+= sys.velocities .* dt_eff sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) return sys end -function B_step!(sys, dt_eff, acceleration_vector, compute_forces::Bool, - n_threads::Integer, neighbors, step_n::Integer) +function B_step!(sys, forces_nounits_t, forces_t, forces_buffer, accels_t, dt_eff, + compute_forces::Bool, n_threads::Integer, neighbors, step_n::Integer) if compute_forces - acceleration_vector .= accelerations(sys, neighbors, step_n; n_threads=n_threads) + forces_nounits_t .= forces_nounits!(forces_nounits_t, sys, neighbors, forces_buffer, step_n; + n_threads=n_threads) + forces_t .= forces_nounits_t .* sys.force_units + accels_t .= forces_t ./ masses(sys) end - sys.velocities .+= dt_eff .* acceleration_vector + sys.velocities .+= dt_eff .* accels_t + return sys +end + +function O_step!(sys, noise, α_eff, σ_eff, rng, temperature, neighbors, step_n) + random_velocities!(noise, sys, temperature; rng=rng) + sys.velocities .= α_eff .* sys.velocities .+ σ_eff .* noise return sys end @@ -621,7 +629,7 @@ function simulate!(sys, forces_t .= forces_nounits_t .* sys.force_units accels_t .= forces_t ./ masses(sys) - noise .= random_velocities(sys, sim.temperature; rng=rng) + random_velocities!(noise, sys, sim.temperature; rng=rng) sys.coords .+= (accels_t ./ sim.friction) .* sim.dt .+ sqrt((2 / sim.friction) * sim.dt) .* noise sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) diff --git a/src/spatial.jl b/src/spatial.jl index 7b183a09f..023e5f55b 100644 --- a/src/spatial.jl +++ b/src/spatial.jl @@ -603,15 +603,26 @@ end """ random_velocities!(sys, temp) + random_velocities!(vels, sys, temp) -Set the velocities of a [`System`](@ref) to random velocities generated from the -Maxwell-Boltzmann distribution. +Set the velocities of a [`System`](@ref), or a vector, to random velocities +generated from the Maxwell-Boltzmann distribution. """ function random_velocities!(sys, temp; rng=Random.GLOBAL_RNG) - sys.velocities = random_velocities(sys, temp; rng=rng) + sys.velocities .= random_velocities(sys, temp; rng=rng) return sys end +function random_velocities!(vels, sys::AbstractSystem{3}, temp; rng=Random.GLOBAL_RNG) + vels .= random_velocity_3D.(masses(sys), temp, sys.k, rng) + return vels +end + +function random_velocities!(vels, sys::AbstractSystem{2}, temp; rng=Random.GLOBAL_RNG) + vels .= random_velocity_2D.(masses(sys), temp, sys.k, rng) + return vels +end + # Sometimes domain error occurs for acos if the value is > 1.0 or < -1.0 acosbound(x::Real) = acos(clamp(x, -1, 1)) @@ -673,7 +684,7 @@ function remove_CM_motion!(sys) atom_masses = masses(sys) cm_momentum = sum_svec(Array(sys.velocities .* atom_masses)) cm_velocity = cm_momentum / sum(Array(atom_masses)) - sys.velocities = sys.velocities .- (cm_velocity,) + sys.velocities .= sys.velocities .- (cm_velocity,) return sys end @@ -875,7 +886,7 @@ Not currently compatible with automatic differentiation using Zygote. function scale_coords!(sys, scale_factor; ignore_molecules=false) if ignore_molecules || isnothing(sys.topology) sys.boundary = scale_boundary(sys.boundary, scale_factor) - sys.coords = scale_vec.(sys.coords, Ref(scale_factor)) + sys.coords .= scale_vec.(sys.coords, Ref(scale_factor)) elseif sys.boundary isa TriclinicBoundary error("scaling coordinates by molecule is not compatible with a TriclinicBoundary") else @@ -903,7 +914,7 @@ function scale_coords!(sys, scale_factor; ignore_molecules=false) coords_nounits[i] = wrap_coords( coords_nounits[i] .+ shift_vecs[mi] .- center_shifts[mi], boundary_nounits) end - sys.coords = coords_nounits * coord_units + sys.coords .= move_array(coords_nounits .* coord_units, sys) end return sys end From be88ae2db36a932d4608001fc52a96594e73d1cb Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 5 Jul 2024 15:39:23 +0100 Subject: [PATCH 24/74] gradient tests use Enzyme --- test/gradients.jl | 71 +++++++++++++++++++++++------------------------ test/runtests.jl | 2 +- 2 files changed, 35 insertions(+), 38 deletions(-) diff --git a/test/gradients.jl b/test/gradients.jl index efd341204..0cc5470b0 100644 --- a/test/gradients.jl +++ b/test/gradients.jl @@ -11,14 +11,21 @@ return F[1] end + function pe(dist) + c1 = SVector(1.0, 1.0, 1.0) + c2 = SVector(dist + 1.0, 1.0, 1.0) + vec = vector(c1, c2, boundary) + potential_energy(inter, vec, a1, a2, NoUnits) + end + function force_grad(dist) - grad = gradient(dist) do dist - c1 = SVector(1.0, 1.0, 1.0) - c2 = SVector(dist + 1.0, 1.0, 1.0) - vec = vector(c1, c2, boundary) - potential_energy(inter, vec, a1, a2, NoUnits) - end - return -grad[1] + grads = autodiff( + Reverse, + pe, + Active, + Active(dist), + ) + return -grads[1][1] end dists = collect(0.2:0.01:1.2) @@ -51,8 +58,6 @@ end ) coords = place_atoms(n_atoms, boundary; min_dist=f32 ? 0.6f0 : 0.6, max_attempts=500) velocities = [random_velocity(atom_mass, temp) for i in 1:n_atoms] - coords_dual = [ForwardDiff.Dual.(x, f32 ? 0.0f0 : 0.0) for x in coords] - velocities_dual = [ForwardDiff.Dual.(x, f32 ? 0.0f0 : 0.0) for x in velocities] nb_cutoff = f32 ? 1.2f0 : 1.2 lj = LennardJones(cutoff=DistanceCutoff(nb_cutoff), use_neighbors=true) crf = CoulombReactionField( @@ -125,8 +130,8 @@ end bond_js, gpu ? CuArray(bonds_inner) : bonds_inner, ) - cs = deepcopy(forward ? coords_dual : coords) - vs = deepcopy(forward ? velocities_dual : velocities) + cs = deepcopy(coords) + vs = deepcopy(velocities) sys = System( atoms=gpu ? CuArray(atoms) : atoms, @@ -179,35 +184,28 @@ end r0 = f32 ? 1.0f0 : 1.0 f = test_simulation_grad(args...) if forward - # Run once to setup - grad_zygote = ( - gradient((σ, r0) -> Zygote.forwarddiff(σ -> f(σ, r0), σ ), σ, r0)[1], - gradient((σ, r0) -> Zygote.forwarddiff(r0 -> f(σ, r0), r0), σ, r0)[2], - ) - grad_zygote = ( - gradient((σ, r0) -> Zygote.forwarddiff(σ -> f(σ, r0), σ ), σ, r0)[1], - gradient((σ, r0) -> Zygote.forwarddiff(r0 -> f(σ, r0), r0), σ, r0)[2], + grad_enzyme = ( + autodiff(Forward, f, Duplicated, Duplicated(σ, f32 ? 1.0f0 : 1.0), Const(r0))[2], + autodiff(Forward, f, Duplicated, Const(σ), Duplicated(r0, f32 ? 1.0f0 : 1.0))[2], ) else - # Run once to setup - grad_zygote = gradient(f, σ, r0) - grad_zygote = gradient(f, σ, r0) + grad_enzyme = autodiff(Reverse, f, Active, Active(σ), Active(r0))[1] end grad_fd = ( central_fdm(6, 1)(σ -> ForwardDiff.value(f(σ, r0)), σ ), central_fdm(6, 1)(r0 -> ForwardDiff.value(f(σ, r0)), r0), ) - for (prefix, gzy, gfd, tol) in zip(("σ", "r0"), grad_zygote, grad_fd, (tol_σ, tol_r0)) + for (prefix, genz, gfd, tol) in zip(("σ", "r0"), grad_enzyme, grad_fd, (tol_σ, tol_r0)) if abs(gfd) < 1e-13 - @info "$(rpad(name, 20)) - $(rpad(prefix, 2)) - FD $gfd, Zygote $gzy" + @info "$(rpad(name, 20)) - $(rpad(prefix, 2)) - FD $gfd, Enzyme $genz" ztol = contains(name, "f32") ? 1e-8 : 1e-10 - @test isnothing(gzy) || abs(gzy) < ztol - elseif isnothing(gzy) - @info "$(rpad(name, 20)) - $(rpad(prefix, 2)) - FD $gfd, Zygote $gzy" - @test !isnothing(gzy) + @test isnothing(genz) || abs(genz) < ztol + elseif isnothing(genz) + @info "$(rpad(name, 20)) - $(rpad(prefix, 2)) - FD $gfd, Enzyme $genz" + @test !isnothing(genz) else - frac_diff = abs(gzy - gfd) / abs(gfd) - @info "$(rpad(name, 20)) - $(rpad(prefix, 2)) - FD $gfd, Zygote $gzy, fractional difference $frac_diff" + frac_diff = abs(genz - gfd) / abs(gfd) + @info "$(rpad(name, 20)) - $(rpad(prefix, 2)) - FD $gfd, Enzyme $genz, fractional difference $frac_diff" @test frac_diff < tol end end @@ -419,19 +417,18 @@ end for (test_name, test_fn, test_tol) in test_runs for (platform, args) in platform_runs f = test_fn(args...) - grads_zygote = CUDA.allowscalar() do - gradient(f, params_dic)[1] - end - @test count(!iszero, values(grads_zygote)) == 67 + grads_enzyme = Dict(k => 0.0 for k in keys(params_dic)) + autodiff(Reverse, f, Active, Duplicated(params_dic, grads_enzyme)) + @test count(!iszero, values(grads_enzyme)) == 67 for param in params_to_test - gzy = grads_zygote[param] + genz = grads_enzyme[param] gfd = central_fdm(6, 1)(params_dic[param]) do val dic = deepcopy(params_dic) dic[param] = val f(dic) end - frac_diff = abs(gzy - gfd) / abs(gfd) - @info "$(rpad(test_name, 6)) - $(rpad(platform, 12)) - $(rpad(param, 21)) - FD $gfd, Zygote $gzy, fractional difference $frac_diff" + frac_diff = abs(genz - gfd) / abs(gfd) + @info "$(rpad(test_name, 6)) - $(rpad(platform, 12)) - $(rpad(param, 21)) - FD $gfd, Enzyme $genz, fractional difference $frac_diff" @test frac_diff < test_tol end end diff --git a/test/runtests.jl b/test/runtests.jl index 398175820..dac047890 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -6,7 +6,7 @@ import AtomsCalculators using AtomsCalculators.AtomsCalculatorsTesting import BioStructures # Imported to avoid clashing names using CUDA -import Enzyme +using Enzyme using FiniteDifferences import SimpleCrystals From ce94d7011abb0489572fd555687234f397a893bd Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Mon, 8 Jul 2024 13:54:06 +0100 Subject: [PATCH 25/74] move Enzyme code to extension --- Project.toml | 3 +-- ext/MollyEnzymeExt.jl | 25 +++++++++++++++++++++++++ src/Molly.jl | 1 - src/units.jl | 12 ++++++------ 4 files changed, 32 insertions(+), 9 deletions(-) create mode 100644 ext/MollyEnzymeExt.jl diff --git a/Project.toml b/Project.toml index 9dd4e0df4..b9521e42e 100644 --- a/Project.toml +++ b/Project.toml @@ -17,7 +17,6 @@ Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" -Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" EzXML = "8f5d6c58-4d21-5cfd-889c-e3ad7ee6a615" FLoops = "cc61a311-1640-44b5-9fba-1b764f453329" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" @@ -37,11 +36,11 @@ UnitfulChainRules = "f31437dd-25a7-4345-875f-756556e6935d" UnsafeAtomicsLLVM = "d80eeb9a-aca5-4d75-85e5-170c8b632249" [weakdeps] -Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" [extensions] +MollyEnzymeExt = "Enzyme" MollyGLMakieExt = ["GLMakie", "Colors"] MollyPythonCallExt = "PythonCall" diff --git a/ext/MollyEnzymeExt.jl b/ext/MollyEnzymeExt.jl new file mode 100644 index 000000000..4ddbe7c28 --- /dev/null +++ b/ext/MollyEnzymeExt.jl @@ -0,0 +1,25 @@ +# Code for taking gradients with Enzyme +# This file is only loaded when Enzyme is imported + +module MollyEnzymeExt + +using Molly +using Enzyme + +EnzymeRules.inactive(::typeof(Molly.n_infinite_dims), args...) = nothing +EnzymeRules.inactive(::typeof(random_velocities), args...) = nothing +EnzymeRules.inactive(::typeof(random_velocities!), args...) = nothing +EnzymeRules.inactive(::typeof(Molly.cuda_threads_blocks_pairwise), args...) = nothing +EnzymeRules.inactive(::typeof(Molly.cuda_threads_blocks_specific), args...) = nothing +EnzymeRules.inactive(::typeof(Molly.check_force_units), args...) = nothing +EnzymeRules.inactive(::typeof(Molly.check_energy_units), args...) = nothing +EnzymeRules.inactive(::typeof(Molly.atoms_bonded_to_N), args...) = nothing +EnzymeRules.inactive(::typeof(Molly.lookup_table), args...) = nothing +EnzymeRules.inactive(::typeof(Molly.cuda_threads_blocks_gbsa), args...) = nothing +EnzymeRules.inactive(::typeof(find_neighbors), args...) = nothing +EnzymeRules.inactive_type(::Type{DistanceNeighborFinder}) = nothing +EnzymeRules.inactive(::typeof(visualize), args...) = nothing +EnzymeRules.inactive(::typeof(place_atoms), args...) = nothing +EnzymeRules.inactive(::typeof(place_diatomics), args...) = nothing + +end diff --git a/src/Molly.jl b/src/Molly.jl index 3b3a0eea5..5720e1d67 100644 --- a/src/Molly.jl +++ b/src/Molly.jl @@ -17,7 +17,6 @@ using CUDA using DataStructures using Distances using Distributions -using Enzyme using EzXML using FLoops using Graphs diff --git a/src/units.jl b/src/units.jl index ef00e6332..4c9f510be 100644 --- a/src/units.jl +++ b/src/units.jl @@ -187,12 +187,6 @@ function convert_k_units(T, k, energy_units) return k_converted end -function check_energy_units(E, energy_units) - if unit(E) != energy_units - error("system energy units are ", energy_units, " but encountered energy units ", unit(E)) - end -end - function check_force_units(F, force_units) if unit(F) != force_units error("system force units are ", force_units, " but encountered force units ", unit(F)) @@ -201,6 +195,12 @@ end check_force_units(F::SVector, force_units) = @inbounds check_force_units(F[1], force_units) +function check_energy_units(E, energy_units) + if unit(E) != energy_units + error("system energy units are ", energy_units, " but encountered energy units ", unit(E)) + end +end + function energy_remove_mol(x) if dimension(x) == u"𝐋^2 * 𝐌 * 𝐍^-1 * 𝐓^-2" T = typeof(ustrip(x)) From ce6c0f3e162e1c8a75da275618b84db97369d8e7 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Mon, 8 Jul 2024 13:55:48 +0100 Subject: [PATCH 26/74] gradient test clearer typing --- test/gradients.jl | 77 +++++++++++++++++++++++------------------------ 1 file changed, 37 insertions(+), 40 deletions(-) diff --git a/test/gradients.jl b/test/gradients.jl index 0cc5470b0..f7b39edd8 100644 --- a/test/gradients.jl +++ b/test/gradients.jl @@ -47,97 +47,94 @@ end function test_simulation_grad(gpu::Bool, parallel::Bool, forward::Bool, f32::Bool, pis::Bool, sis::Bool, obc2::Bool, gbn2::Bool) + T = f32 ? Float32 : Float64 + AT = gpu ? CuArray : Array n_atoms = 50 n_steps = 100 - atom_mass = f32 ? 10.0f0 : 10.0 - boundary = f32 ? CubicBoundary(3.0f0) : CubicBoundary(3.0) - temp = f32 ? 1.0f0 : 1.0 + atom_mass = T(10.0) + boundary = CubicBoundary(T(3.0)) + temp = T(1.0) simulator = VelocityVerlet( - dt=f32 ? 0.001f0 : 0.001, + dt=T(0.001), coupling=RescaleThermostat(temp), ) - coords = place_atoms(n_atoms, boundary; min_dist=f32 ? 0.6f0 : 0.6, max_attempts=500) + coords = place_atoms(n_atoms, boundary; min_dist=T(0.6), max_attempts=500) velocities = [random_velocity(atom_mass, temp) for i in 1:n_atoms] - nb_cutoff = f32 ? 1.2f0 : 1.2 + nb_cutoff = T(1.2) lj = LennardJones(cutoff=DistanceCutoff(nb_cutoff), use_neighbors=true) crf = CoulombReactionField( dist_cutoff=nb_cutoff, - solvent_dielectric=f32 ? Float32(Molly.crf_solvent_dielectric) : Molly.crf_solvent_dielectric, + solvent_dielectric=T(Molly.crf_solvent_dielectric), use_neighbors=true, - coulomb_const=f32 ? Float32(ustrip(Molly.coulomb_const)) : ustrip(Molly.coulomb_const), + coulomb_const=T(ustrip(Molly.coulomb_const)), ) pairwise_inters = pis ? (lj, crf) : () - bond_is = gpu ? CuArray(Int32.(collect(1:(n_atoms ÷ 2)))) : Int32.(collect(1:(n_atoms ÷ 2))) - bond_js = gpu ? CuArray(Int32.(collect((1 + n_atoms ÷ 2):n_atoms))) : Int32.(collect((1 + n_atoms ÷ 2):n_atoms)) + bond_is = AT(Int32.(collect(1:(n_atoms ÷ 2)))) + bond_js = AT(Int32.(collect((1 + n_atoms ÷ 2):n_atoms))) bond_dists = [norm(vector(Array(coords)[i], Array(coords)[i + n_atoms ÷ 2], boundary)) for i in 1:(n_atoms ÷ 2)] - angles_inner = [HarmonicAngle(k=f32 ? 10.0f0 : 10.0, θ0=f32 ? 2.0f0 : 2.0) for i in 1:15] + angles_inner = [HarmonicAngle(k=T(10.0), θ0=T(2.0)) for i in 1:15] angles = InteractionList3Atoms( - gpu ? CuArray(Int32.(collect( 1:15))) : Int32.(collect( 1:15)), - gpu ? CuArray(Int32.(collect(16:30))) : Int32.(collect(16:30)), - gpu ? CuArray(Int32.(collect(31:45))) : Int32.(collect(31:45)), - gpu ? CuArray(angles_inner) : angles_inner, + AT(Int32.(collect( 1:15))), + AT(Int32.(collect(16:30))), + AT(Int32.(collect(31:45))), + AT(angles_inner), ) torsions_inner = [PeriodicTorsion( periodicities=[1, 2, 3], - phases=f32 ? [1.0f0, 0.0f0, -1.0f0] : [1.0, 0.0, -1.0], - ks=f32 ? [10.0f0, 5.0f0, 8.0f0] : [10.0, 5.0, 8.0], + phases=T[1.0, 0.0, -1.0], + ks=T[10.0, 5.0, 8.0], n_terms=6, ) for i in 1:10] torsions = InteractionList4Atoms( - gpu ? CuArray(Int32.(collect( 1:10))) : Int32.(collect( 1:10)), - gpu ? CuArray(Int32.(collect(11:20))) : Int32.(collect(11:20)), - gpu ? CuArray(Int32.(collect(21:30))) : Int32.(collect(21:30)), - gpu ? CuArray(Int32.(collect(31:40))) : Int32.(collect(31:40)), - gpu ? CuArray(torsions_inner) : torsions_inner, + AT(Int32.(collect( 1:10))), + AT(Int32.(collect(11:20))), + AT(Int32.(collect(21:30))), + AT(Int32.(collect(31:40))), + AT(torsions_inner), ) - atoms_setup = [Atom(charge=f32 ? 0.0f0 : 0.0, σ=f32 ? 0.0f0 : 0.0) for i in 1:n_atoms] + atoms_setup = [Atom(charge=zero(T), σ=zero(T)) for i in 1:n_atoms] if obc2 imp_obc2 = ImplicitSolventOBC( - gpu ? CuArray(atoms_setup) : atoms_setup, + AT(atoms_setup), [AtomData(element="O") for i in 1:n_atoms], InteractionList2Atoms(bond_is, bond_js, nothing); - kappa=(f32 ? 0.7f0 : 0.7), + kappa=T(0.7), use_OBC2=true, ) general_inters = (imp_obc2,) elseif gbn2 imp_gbn2 = ImplicitSolventGBN2( - gpu ? CuArray(atoms_setup) : atoms_setup, + AT(atoms_setup), [AtomData(element="O") for i in 1:n_atoms], InteractionList2Atoms(bond_is, bond_js, nothing); - kappa=(f32 ? 0.7f0 : 0.7), + kappa=T(0.7), ) general_inters = (imp_gbn2,) else general_inters = () end neighbor_finder = DistanceNeighborFinder( - eligible=gpu ? CuArray(trues(n_atoms, n_atoms)) : trues(n_atoms, n_atoms), + eligible=AT(trues(n_atoms, n_atoms)), n_steps=10, - dist_cutoff=f32 ? 1.5f0 : 1.5, + dist_cutoff=T(1.5), ) function loss(σ, r0) - if f32 - atoms = [Atom(i, i % 2 == 0 ? -0.02f0 : 0.02f0, atom_mass, σ, 0.2f0, false) for i in 1:n_atoms] - else - atoms = [Atom(i, i % 2 == 0 ? -0.02 : 0.02, atom_mass, σ, 0.2, false) for i in 1:n_atoms] - end - - bonds_inner = [HarmonicBond(f32 ? 100.0f0 : 100.0, bond_dists[i] * r0) for i in 1:(n_atoms ÷ 2)] + atoms = [Atom(i, i % 2 == 0 ? T(-0.02) : T(0.02), atom_mass, σ, T(0.2), false) for i in 1:n_atoms] + bonds_inner = [HarmonicBond(T(100.0), bond_dists[i] * r0) for i in 1:(n_atoms ÷ 2)] bonds = InteractionList2Atoms( bond_is, bond_js, - gpu ? CuArray(bonds_inner) : bonds_inner, + AT(bonds_inner), ) cs = deepcopy(coords) vs = deepcopy(velocities) sys = System( - atoms=gpu ? CuArray(atoms) : atoms, - coords=gpu ? CuArray(cs) : cs, + atoms=AT(atoms), + coords=AT(cs), boundary=boundary, - velocities=gpu ? CuArray(vs) : vs, + velocities=AT(vs), pairwise_inters=pairwise_inters, specific_inter_lists=sis ? (bonds, angles, torsions) : (), general_inters=general_inters, From 2df85058a0288b9444e7f28f8badcbbb3a360ad8 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 18 Jul 2024 21:18:09 +0100 Subject: [PATCH 27/74] LennardJones mixing and shortcut --- src/interactions/lennard_jones.jl | 195 ++++++++++++++---------------- 1 file changed, 91 insertions(+), 104 deletions(-) diff --git a/src/interactions/lennard_jones.jl b/src/interactions/lennard_jones.jl index 664009ccf..dbd4b24a4 100644 --- a/src/interactions/lennard_jones.jl +++ b/src/interactions/lennard_jones.jl @@ -2,9 +2,27 @@ export LennardJones, LennardJonesSoftCore +function lj_zero_shortcut(atom_i, atom_j) + return iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || + iszero_value(atom_i.σ) || iszero_value(atom_j.σ) +end + +no_shortcut(atom_i, atom_j) = false + +function lorentz_σ_mixing(atom_i, atom_j) + return (atom_i.σ + atom_j.σ) / 2 +end + +function geometric_σ_mixing(atom_i, atom_j) + return sqrt(atom_i.σ * atom_j.σ) +end + +function geometric_ϵ_mixing(atom_i, atom_j) + return sqrt(atom_i.ϵ * atom_j.ϵ) +end + @doc raw""" - LennardJones(; cutoff, use_neighbors, lorentz_mixing, weight_special, weight_solute_solvent, - skip_shortcut) + LennardJones(; cutoff, use_neighbors, shortcut, σ_mixing, ϵ_mixing, weight_special) The Lennard-Jones 6-12 interaction between two atoms. @@ -20,72 +38,62 @@ and the force on each atom by \end{aligned} ``` """ -struct LennardJones{S, C, W, WS} <: PairwiseInteraction +struct LennardJones{C, W} <: PairwiseInteraction cutoff::C use_neighbors::Bool - lorentz_mixing::Bool + shortcut::Function + σ_mixing::Function + ϵ_mixing::Function weight_special::W - weight_solute_solvent::WS end function LennardJones(; cutoff=NoCutoff(), use_neighbors=false, - lorentz_mixing=true, - weight_special=1, - weight_solute_solvent=1, - skip_shortcut=false) - return LennardJones{skip_shortcut, typeof(cutoff), typeof(weight_special), - typeof(weight_solute_solvent)}( - cutoff, use_neighbors, lorentz_mixing, weight_special, weight_solute_solvent) + shortcut=lj_zero_shortcut, + σ_mixing=lorentz_σ_mixing, + ϵ_mixing=geometric_ϵ_mixing, + weight_special=1) + return LennardJones(cutoff, use_neighbors, shortcut, σ_mixing, ϵ_mixing, + weight_special) end use_neighbors(inter::LennardJones) = inter.use_neighbors -is_solute(at::Atom) = at.solute -is_solute(at) = false - -function Base.zero(lj::LennardJones{S, C, W, WS}) where {S, C, W, WS} - return LennardJones{S, C, W, WS}( +function Base.zero(lj::LennardJones{C, W}) where {C, W} + return LennardJones( lj.cutoff, - false, - false, + lj.use_neighbors, + lj.shortcut, + lj.σ_mixing, + lj.ϵ_mixing, zero(W), - zero(WS), ) end -function Base.:+(l1::LennardJones{S, C, W, WS}, - l2::LennardJones{S, C, W, WS}) where {S, C, W, WS} - return LennardJones{S, C, W, WS}( +function Base.:+(l1::LennardJones, l2::LennardJones) + return LennardJones( l1.cutoff, l1.use_neighbors, - l1.lorentz_mixing, + l1.shortcut, + l1.σ_mixing, + l1.ϵ_mixing, l1.weight_special + l2.weight_special, - l1.weight_solute_solvent + l2.weight_solute_solvent, ) end -@inline function force(inter::LennardJones{S, C}, +@inline function force(inter::LennardJones, dr, atom_i, atom_j, force_units=u"kJ * mol^-1 * nm^-1", - special::Bool=false, - args...) where {S, C} - if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || - iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) + special=false, + args...) + if inter.shortcut(atom_i, atom_j) return ustrip.(zero(dr)) * force_units end - - # Lorentz-Berthelot mixing rules use the arithmetic average for σ - # Otherwise use the geometric average - σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) - if (is_solute(atom_i) && !is_solute(atom_j)) || (is_solute(atom_j) && !is_solute(atom_i)) - ϵ = inter.weight_solute_solvent * sqrt(atom_i.ϵ * atom_j.ϵ) - else - ϵ = sqrt(atom_i.ϵ * atom_j.ϵ) - end + σ = inter.σ_mixing(atom_i, atom_j) + ϵ = inter.ϵ_mixing(atom_i, atom_j) cutoff = inter.cutoff r2 = sum(abs2, dr) @@ -105,24 +113,18 @@ function force_divr(::LennardJones, r2, invr2, (σ2, ϵ)) return (24ϵ * invr2) * (2 * six_term ^ 2 - six_term) end -@inline function potential_energy(inter::LennardJones{S, C}, +@inline function potential_energy(inter::LennardJones, dr, atom_i, atom_j, energy_units=u"kJ * mol^-1", - special::Bool=false, - args...) where {S, C} - if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || - iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) + special=false, + args...) + if inter.shortcut(atom_i, atom_j) return ustrip(zero(dr[1])) * energy_units end - - σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) - if (is_solute(atom_i) && !is_solute(atom_j)) || (is_solute(atom_j) && !is_solute(atom_i)) - ϵ = inter.weight_solute_solvent * sqrt(atom_i.ϵ * atom_j.ϵ) - else - ϵ = sqrt(atom_i.ϵ * atom_j.ϵ) - end + σ = inter.σ_mixing(atom_i, atom_j) + ϵ = inter.ϵ_mixing(atom_i, atom_j) cutoff = inter.cutoff r2 = sum(abs2, dr) @@ -143,8 +145,8 @@ function potential(::LennardJones, r2, invr2, (σ2, ϵ)) end @doc raw""" - LennardJonesSoftCore(; cutoff, α, λ, p, use_neighbors, lorentz_mixing, weight_special, - weight_solute_solvent, skip_shortcut) + LennardJonesSoftCore(; cutoff, α, λ, p, use_neighbors, shortcut, σ_mixing, ϵ_mixing, + weight_special) The Lennard-Jones 6-12 interaction between two atoms with a soft core. @@ -162,16 +164,17 @@ r_{ij}^{\text{sc}} = \left(r_{ij}^6 + \alpha \sigma_{ij}^6 \lambda^p \right)^{1/ ``` If ``\alpha`` or ``\lambda`` are zero this gives the standard [`LennardJones`](@ref) potential. """ -struct LennardJonesSoftCore{S, C, A, L, P, R, W, WS} <: PairwiseInteraction +struct LennardJonesSoftCore{C, A, L, P, W, R} <: PairwiseInteraction cutoff::C α::A λ::L p::P - σ6_fac::R use_neighbors::Bool - lorentz_mixing::Bool + shortcut::Function + σ_mixing::Function + ϵ_mixing::Function weight_special::W - weight_solute_solvent::WS + σ6_fac::R end function LennardJonesSoftCore(; @@ -180,69 +183,59 @@ function LennardJonesSoftCore(; λ=0, p=2, use_neighbors=false, - lorentz_mixing=true, - weight_special=1, - weight_solute_solvent=1, - skip_shortcut=false) + shortcut=lj_zero_shortcut, + σ_mixing=lorentz_σ_mixing, + ϵ_mixing=geometric_ϵ_mixing, + weight_special=1) σ6_fac = α * λ^p - return LennardJonesSoftCore{skip_shortcut, typeof(cutoff), typeof(α), typeof(λ), typeof(p), - typeof(σ6_fac), typeof(weight_special), - typeof(weight_solute_solvent)}( - cutoff, α, λ, p, σ6_fac, use_neighbors, lorentz_mixing, weight_special, - weight_solute_solvent) + return LennardJonesSoftCore(cutoff, α, λ, p, use_neighbors, shortcut, σ_mixing, + ϵ_mixing, weight_special, σ6_fac) end use_neighbors(inter::LennardJonesSoftCore) = inter.use_neighbors -function Base.zero(lj::LennardJonesSoftCore{S, C, A, L, P, R, W, WS}) where {S, C, A, L, P, R, W, WS} - return LennardJonesSoftCore{S, C, A, L, P, R, W, WS}( +function Base.zero(lj::LennardJonesSoftCore{C, A, L, P, W, R}) where {C, A, L, P, W, R} + return LennardJonesSoftCore( lj.cutoff, zero(A), zero(L), zero(P), - zero(R), - false, - false, + lj.use_neighbors, + lj.shortcut, + lj.σ_mixing, + lj.ϵ_mixing, zero(W), - zero(WS), + zero(R), ) end -function Base.:+(l1::LennardJonesSoftCore{S, C, A, L, P, R, W, WS}, - l2::LennardJonesSoftCore{S, C, A, L, P, R, W, WS}) where {S, C, A, L, P, R, W, WS} - return LennardJonesSoftCore{S, C, A, L, P, R, W, WS}( +function Base.:+(l1::LennardJonesSoftCore, l2::LennardJonesSoftCore) + return LennardJonesSoftCore( l1.cutoff, l1.α + l2.α, l1.λ + l2.λ, l1.p + l2.p, - l1.σ6_fac + l2.σ6_fac, l1.use_neighbors, - l1.lorentz_mixing, + l1.shortcut, + l1.σ_mixing, + l1.ϵ_mixing, l1.weight_special + l2.weight_special, - l1.weight_solute_solvent + l2.weight_solute_solvent, + l1.σ6_fac + l2.σ6_fac, ) end -@inline function force(inter::LennardJonesSoftCore{S, C}, +@inline function force(inter::LennardJonesSoftCore, dr, atom_i, atom_j, force_units=u"kJ * mol^-1 * nm^-1", - special::Bool=false, - args...) where {S, C} - if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || - iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) + special=false, + args...) + if inter.shortcut(atom_i, atom_j) return ustrip.(zero(dr)) * force_units end - - # Lorentz-Berthelot mixing rules use the arithmetic average for σ - # Otherwise use the geometric average - σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) - if (is_solute(atom_i) && !is_solute(atom_j)) || (is_solute(atom_j) && !is_solute(atom_i)) - ϵ = inter.weight_solute_solvent * sqrt(atom_i.ϵ * atom_j.ϵ) - else - ϵ = sqrt(atom_i.ϵ * atom_j.ϵ) - end + σ = inter.σ_mixing(atom_i, atom_j) + ϵ = inter.ϵ_mixing(atom_i, atom_j) cutoff = inter.cutoff r2 = sum(abs2, dr) @@ -265,24 +258,18 @@ function force_divr(::LennardJonesSoftCore, r2, invr2, (σ2, ϵ, σ6_fac)) return ff * √invr2 end -@inline function potential_energy(inter::LennardJonesSoftCore{S, C}, +@inline function potential_energy(inter::LennardJonesSoftCore, dr, atom_i, atom_j, energy_units=u"kJ * mol^-1", - special::Bool=false, - args...) where {S, C} - if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || - iszero_value(atom_i.σ) ||iszero_value(atom_j.σ)) + special=false, + args...) + if inter.shortcut(atom_i, atom_j) return ustrip(zero(dr[1])) * energy_units end - - σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) - if (is_solute(atom_i) && !is_solute(atom_j)) || (is_solute(atom_j) && !is_solute(atom_i)) - ϵ = inter.weight_solute_solvent * sqrt(atom_i.ϵ * atom_j.ϵ) - else - ϵ = sqrt(atom_i.ϵ * atom_j.ϵ) - end + σ = inter.σ_mixing(atom_i, atom_j) + ϵ = inter.ϵ_mixing(atom_i, atom_j) cutoff = inter.cutoff r2 = sum(abs2, dr) From 10e11f75fee6b7a90e09f39d99071566bf5bd61b Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 18 Jul 2024 21:18:22 +0100 Subject: [PATCH 28/74] Mie mixing and shortcut --- src/interactions/mie.jl | 60 +++++++++++++++++++++++++---------------- 1 file changed, 37 insertions(+), 23 deletions(-) diff --git a/src/interactions/mie.jl b/src/interactions/mie.jl index 74fd9c5a3..bf683b659 100644 --- a/src/interactions/mie.jl +++ b/src/interactions/mie.jl @@ -1,7 +1,7 @@ export Mie @doc raw""" - Mie(; m, n, cutoff, use_neighbors, lorentz_mixing, skip_shortcut) + Mie(; m, n, cutoff, use_neighbors, shortcut, σ_mixing, ϵ_mixing) The Mie generalized interaction between two atoms. @@ -15,12 +15,14 @@ where C = \frac{n}{n - m} \left( \frac{n}{m} \right) ^\frac{m}{n - m} ``` """ -struct Mie{S, C, T} <: PairwiseInteraction +struct Mie{C, T} <: PairwiseInteraction m::T n::T cutoff::C use_neighbors::Bool - lorentz_mixing::Bool + shortcut::Function + σ_mixing::Function + ϵ_mixing::Function mn_fac::T end @@ -29,30 +31,44 @@ function Mie(; n, cutoff=NoCutoff(), use_neighbors=false, - lorentz_mixing=true, - skip_shortcut=false) + shortcut=lj_zero_shortcut, + σ_mixing=lorentz_σ_mixing, + ϵ_mixing=geometric_ϵ_mixing) m_p, n_p, mn_fac = promote(m, n, (n / (n - m)) * (n / m) ^ (m / (n - m))) - return Mie{skip_shortcut, typeof(cutoff), typeof(m_p)}( - m_p, n_p, cutoff, use_neighbors, lorentz_mixing, mn_fac) + return Mie(m_p, n_p, cutoff, use_neighbors, shortcut, σ_mixing, ϵ_mixing, mn_fac) end use_neighbors(inter::Mie) = inter.use_neighbors -function force(inter::Mie{S, C, T}, +function Base.zero(m::Mie{C, T}) where {C, T} + return Mie(zero(T), zero(T), m.cutoff, m.use_neighbors, m.shortcut, m.σ_mixing, + m.ϵ_mixing, zero(T)) +end + +function Base.:+(m1::Mie, m2::Mie) + return Mie( + m1.m + m2.m, + m1.n + m2.n, + m1.cutoff, + m1.use_neighbors, + m1.shortcut, + m1.σ_mixing, + m1.ϵ_mixing, + m1.mn_fac + m2.mn_fac, + ) +end + +function force(inter::Mie, dr, atom_i, atom_j, force_units=u"kJ * mol^-1 * nm^-1", - args...) where {S, C, T} - if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || - iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) + args...) + if inter.shortcut(atom_i, atom_j) return ustrip.(zero(dr)) * force_units end - - # Lorentz-Berthelot mixing rules use the arithmetic average for σ - # Otherwise use the geometric average - σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) - ϵ = sqrt(atom_i.ϵ * atom_j.ϵ) + σ = inter.σ_mixing(atom_i, atom_j) + ϵ = inter.ϵ_mixing(atom_i, atom_j) cutoff = inter.cutoff r2 = sum(abs2, dr) @@ -71,19 +87,17 @@ function force_divr(::Mie, r2, invr2, (m, n, σ_r, const_mn)) return -const_mn / r2 * (m * σ_r ^ m - n * σ_r ^ n) end -@inline function potential_energy(inter::Mie{S, C, T}, +@inline function potential_energy(inter::Mie, dr, atom_i, atom_j, energy_units=u"kJ * mol^-1", - args...) where {S, C, T} - if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || - iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) + args...) + if inter.shortcut(atom_i, atom_j) return ustrip(zero(dr[1])) * energy_units end - - σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) - ϵ = sqrt(atom_i.ϵ * atom_j.ϵ) + σ = inter.σ_mixing(atom_i, atom_j) + ϵ = inter.ϵ_mixing(atom_i, atom_j) cutoff = inter.cutoff r2 = sum(abs2, dr) From ea60adcfa63dcf64cf6550675116bed5f8a3202f Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 18 Jul 2024 21:18:37 +0100 Subject: [PATCH 29/74] soft sphere mixing and shortcut --- src/interactions/soft_sphere.jl | 50 ++++++++++++++++++--------------- 1 file changed, 27 insertions(+), 23 deletions(-) diff --git a/src/interactions/soft_sphere.jl b/src/interactions/soft_sphere.jl index cb71fe127..a513e7163 100644 --- a/src/interactions/soft_sphere.jl +++ b/src/interactions/soft_sphere.jl @@ -1,7 +1,7 @@ export SoftSphere @doc raw""" - SoftSphere(; cutoff, use_neighbors, lorentz_mixing, skip_shortcut) + SoftSphere(; cutoff, use_neighbors, shortcut, σ_mixing, ϵ_mixing) The soft-sphere potential. @@ -10,38 +10,44 @@ The potential energy is defined as V(r_{ij}) = 4\varepsilon_{ij} \left(\frac{\sigma_{ij}}{r_{ij}}\right)^{12} ``` """ -struct SoftSphere{S, C} <: PairwiseInteraction +struct SoftSphere{C} <: PairwiseInteraction cutoff::C use_neighbors::Bool - lorentz_mixing::Bool + shortcut::Function + σ_mixing::Function + ϵ_mixing::Function end function SoftSphere(; cutoff=NoCutoff(), use_neighbors=false, - lorentz_mixing=true, - skip_shortcut=false) - return SoftSphere{skip_shortcut, typeof(cutoff)}( - cutoff, use_neighbors, lorentz_mixing) + shortcut=lj_zero_shortcut, + σ_mixing=lorentz_σ_mixing, + ϵ_mixing=geometric_ϵ_mixing) + return SoftSphere(cutoff, use_neighbors, shortcut, σ_mixing, ϵ_mixing) end use_neighbors(inter::SoftSphere) = inter.use_neighbors -@inline function force(inter::SoftSphere{S, C}, +function Base.zero(ss::SoftSphere) + return SoftSphere(ss.cutoff, ss.use_neighbors, ss.shortcut, ss.σ_mixing, ss.ϵ_mixing) +end + +function Base.:+(s1::SoftSphere, ::SoftSphere) + return SoftSphere(s1.cutoff, s1.use_neighbors, s1.shortcut, s1.σ_mixing, s1.ϵ_mixing) +end + +@inline function force(inter::SoftSphere, dr, atom_i, atom_j, force_units=u"kJ * mol^-1 * nm^-1", - args...) where {S, C} - if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || - iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) + args...) + if inter.shortcut(atom_i, atom_j) return ustrip.(zero(dr)) * force_units end - - # Lorentz-Berthelot mixing rules use the arithmetic average for σ - # Otherwise use the geometric average - σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) - ϵ = sqrt(atom_i.ϵ * atom_j.ϵ) + σ = inter.σ_mixing(atom_i, atom_j) + ϵ = inter.ϵ_mixing(atom_i, atom_j) cutoff = inter.cutoff r2 = sum(abs2, dr) @@ -57,19 +63,17 @@ function force_divr(::SoftSphere, r2, invr2, (σ2, ϵ)) return (24ϵ * invr2) * 2 * six_term ^ 2 end -function potential_energy(inter::SoftSphere{S, C}, +function potential_energy(inter::SoftSphere, dr, atom_i, atom_j, energy_units=u"kJ * mol^-1", - args...) where {S, C} - if !S && (iszero_value(atom_i.ϵ) || iszero_value(atom_j.ϵ) || - iszero_value(atom_i.σ) || iszero_value(atom_j.σ)) + args...) + if inter.shortcut(atom_i, atom_j) return ustrip(zero(dr[1])) * energy_units end - - σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) - ϵ = sqrt(atom_i.ϵ * atom_j.ϵ) + σ = inter.σ_mixing(atom_i, atom_j) + ϵ = inter.ϵ_mixing(atom_i, atom_j) cutoff = inter.cutoff r2 = sum(abs2, dr) From 3bb4b0411f3cc1e7c9ffffad838c4876b451d9ee Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 18 Jul 2024 21:18:58 +0100 Subject: [PATCH 30/74] Buckingham mixing and shortcut --- src/interactions/buckingham.jl | 82 ++++++++++++++++++++++++---------- 1 file changed, 59 insertions(+), 23 deletions(-) diff --git a/src/interactions/buckingham.jl b/src/interactions/buckingham.jl index d63a8a8f3..71d68f4bf 100644 --- a/src/interactions/buckingham.jl +++ b/src/interactions/buckingham.jl @@ -1,7 +1,29 @@ export Buckingham +function buckingham_zero_shortcut(atom_i, atom_j) + return (iszero_value(atom_i.A) || iszero_value(atom_j.A)) && + (iszero_value(atom_i.C) || iszero_value(atom_j.C)) +end + +function geometric_A_mixing(atom_i, atom_j) + return sqrt(atom_i.A * atom_j.A) +end + +function geometric_B_mixing(atom_i, atom_j) + return sqrt(atom_i.B * atom_j.B) +end + +function geometric_C_mixing(atom_i, atom_j) + return sqrt(atom_i.C * atom_j.C) +end + +function inverse_B_mixing(atom_i, atom_j) + return 2 / (inv(atom_i.B) + inv(atom_j.B)) +end + @doc raw""" - Buckingham(; cutoff, use_neighbors, weight_special) + Buckingham(; cutoff, use_neighbors, shortcut, A_mixing, B_mixing, + C_mixing, weight_special) The Buckingham interaction between two atoms. @@ -26,38 +48,54 @@ so atoms that use this interaction should have fields `A`, `B` and `C` available struct Buckingham{C, W} <: PairwiseInteraction cutoff::C use_neighbors::Bool + shortcut::Function + A_mixing::Function + B_mixing::Function + C_mixing::Function weight_special::W end function Buckingham(; cutoff=NoCutoff(), use_neighbors=false, + shortcut=buckingham_zero_shortcut, + A_mixing=geometric_A_mixing, + B_mixing=inverse_B_mixing, + C_mixing=geometric_C_mixing, weight_special=1) - return Buckingham{typeof(cutoff), typeof(weight_special)}( - cutoff, use_neighbors, weight_special) + return Buckingham(cutoff, use_neighbors, shortcut, A_mixing, B_mixing, + C_mixing, weight_special) end use_neighbors(inter::Buckingham) = inter.use_neighbors -@inline function force(inter::Buckingham{C}, +function Base.zero(b::Buckingham{C, W}) where {C, W} + return Buckingham(b.cutoff, b.use_neighbors, b.shortcut, b.A_mixing, + b.B_mixing, b.C_mixing, zero(W)) +end + +function Base.:+(b1::Buckingham, b2::Buckingham) + return Buckingham(b1.cutoff, b1.use_neighbors, b1.shortcut, b1.A_mixing, b1.B_mixing, + b1.C_mixing, b1.weight_special + b2.weight_special) +end + +@inline function force(inter::Buckingham, dr, atom_i, atom_j, force_units=u"kJ * mol^-1 * nm^-1", - special::Bool=false, - args...) where C - if (iszero_value(atom_i.A) || iszero_value(atom_j.A)) && - (iszero_value(atom_i.C) || iszero_value(atom_j.C)) + special=false, + args...) + if inter.shortcut(atom_i, atom_j) return ustrip.(zero(dr)) * force_units end - - Aij = sqrt(atom_i.A * atom_j.A) - Bij = 2 / (inv(atom_i.B) + inv(atom_j.B)) - Cij = sqrt(atom_i.C * atom_j.C) + A = inter.A_mixing(atom_i, atom_j) + B = inter.B_mixing(atom_i, atom_j) + C = inter.C_mixing(atom_i, atom_j) cutoff = inter.cutoff r2 = sum(abs2, dr) - params = (Aij, Bij, Cij) + params = (A, B, C) f = force_divr_with_cutoff(inter, r2, params, cutoff, force_units) if special @@ -72,25 +110,23 @@ function force_divr(::Buckingham, r2, invr2, (A, B, C)) return A * B * exp(-B * r) / r - 6 * C * invr2^4 end -@inline function potential_energy(inter::Buckingham{C}, +@inline function potential_energy(inter::Buckingham, dr, atom_i, atom_j, energy_units=u"kJ * mol^-1", - special::Bool=false, - args...) where C - if (iszero_value(atom_i.A) || iszero_value(atom_j.A)) && - (iszero_value(atom_i.C) || iszero_value(atom_j.C)) + special=false, + args...) + if inter.shortcut(atom_i, atom_j) return ustrip(zero(dr[1])) * energy_units end - - Aij = sqrt(atom_i.A * atom_j.A) - Bij = 2 / (inv(atom_i.B) + inv(atom_j.B)) - Cij = sqrt(atom_i.C * atom_j.C) + A = inter.A_mixing(atom_i, atom_j) + B = inter.B_mixing(atom_i, atom_j) + C = inter.C_mixing(atom_i, atom_j) cutoff = inter.cutoff r2 = sum(abs2, dr) - params = (Aij, Bij, Cij) + params = (A, B, C) pe = potential_with_cutoff(inter, r2, params, cutoff, energy_units) if special From 24fcc660cdc83c78c54d34f2c30a586eb9ac7f98 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 18 Jul 2024 21:19:20 +0100 Subject: [PATCH 31/74] Coulomb mixing --- src/interactions/coulomb.jl | 65 ++++++++++++++++++------------------- 1 file changed, 31 insertions(+), 34 deletions(-) diff --git a/src/interactions/coulomb.jl b/src/interactions/coulomb.jl index 00878a0f9..49f6f711b 100644 --- a/src/interactions/coulomb.jl +++ b/src/interactions/coulomb.jl @@ -27,14 +27,13 @@ function Coulomb(; use_neighbors=false, weight_special=1, coulomb_const=coulomb_const) - return Coulomb{typeof(cutoff), typeof(weight_special), typeof(coulomb_const)}( - cutoff, use_neighbors, weight_special, coulomb_const) + return Coulomb(cutoff, use_neighbors, weight_special, coulomb_const) end use_neighbors(inter::Coulomb) = inter.use_neighbors function Base.zero(coul::Coulomb{C, W, T}) where {C, W, T} - return Coulomb{C, W, T}(coul.cutoff, false, zero(W), zero(T)) + return Coulomb(coul.cutoff, coul.use_neighbors, zero(W), zero(T)) end function Base.:+(c1::Coulomb, c2::Coulomb) @@ -51,7 +50,7 @@ end atom_i, atom_j, force_units=u"kJ * mol^-1 * nm^-1", - special::Bool=false, + special=false, args...) where C r2 = sum(abs2, dr) cutoff = inter.cutoff @@ -76,7 +75,7 @@ end atom_i, atom_j, energy_units=u"kJ * mol^-1", - special::Bool=false, + special=false, args...) where C r2 = sum(abs2, dr) cutoff = inter.cutoff @@ -97,7 +96,7 @@ function potential(::Coulomb, r2, invr2, (ke, qi, qj)) end @doc raw""" - CoulombSoftCore(; cutoff, α, λ, p, use_neighbors, lorentz_mixing, weight_special, coulomb_const) + CoulombSoftCore(; cutoff, α, λ, p, use_neighbors, σ_mixing, weight_special, coulomb_const) The Coulomb electrostatic interaction between two atoms with a soft core. @@ -107,16 +106,16 @@ V(r_{ij}) = \frac{q_i q_j}{4 \pi \varepsilon_0 (r_{ij}^6 + \alpha \sigma_{ij}^6 ``` If ``\alpha`` or ``\lambda`` are zero this gives the standard [`Coulomb`](@ref) potential. """ -struct CoulombSoftCore{C, A, L, P, R, W, T} <: PairwiseInteraction +struct CoulombSoftCore{C, A, L, P, W, T, R} <: PairwiseInteraction cutoff::C α::A λ::L p::P - σ6_fac::R use_neighbors::Bool - lorentz_mixing::Bool + σ_mixing::Function weight_special::W coulomb_const::T + σ6_fac::R end function CoulombSoftCore(; @@ -125,28 +124,27 @@ function CoulombSoftCore(; λ=0, p=2, use_neighbors=false, - lorentz_mixing=true, + σ_mixing=lorentz_σ_mixing, weight_special=1, coulomb_const=coulomb_const) σ6_fac = α * λ^p - return CoulombSoftCore{typeof(cutoff), typeof(α), typeof(λ), typeof(p), typeof(σ6_fac), - typeof(weight_special), typeof(coulomb_const)}( - cutoff, α, λ, p, σ6_fac, use_neighbors, lorentz_mixing, weight_special, coulomb_const) + return CoulombSoftCore(cutoff, α, λ, p, use_neighbors, σ_mixing, weight_special, + coulomb_const, σ6_fac) end use_neighbors(inter::CoulombSoftCore) = inter.use_neighbors -function Base.zero(coul::CoulombSoftCore{C, A, L, P, R, W, T}) where {C, A, L, P, R, W, T} - return CoulombSoftCore{C, A, L, P, R, W, T}( +function Base.zero(coul::CoulombSoftCore{C, A, L, P, W, T, R}) where {C, A, L, P, W, T, R} + return CoulombSoftCore( coul.cutoff, zero(A), zero(L), zero(P), - zero(R), - false, - false, + coul.use_neighbors, + coul.σ_mixing, zero(W), zero(T), + zero(R), ) end @@ -156,26 +154,26 @@ function Base.:+(c1::CoulombSoftCore, c2::CoulombSoftCore) c1.α + c2.α, c1.λ + c2.λ, c1.p + c2.p, - c1.σ6_fac + c2.σ6_fac, c1.use_neighbors, - c1.lorentz_mixing, + c1.σ_mixing, c1.weight_special + c2.weight_special, c1.coulomb_const + c2.coulomb_const, + c1.σ6_fac + c2.σ6_fac, ) end -@inline function force(inter::CoulombSoftCore{C}, +@inline function force(inter::CoulombSoftCore, dr, atom_i, atom_j, force_units=u"kJ * mol^-1 * nm^-1", - special::Bool=false, - args...) where C + special=false, + args...) r2 = sum(abs2, dr) cutoff = inter.cutoff ke = inter.coulomb_const qi, qj = atom_i.charge, atom_j.charge - σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) + σ = inter.σ_mixing(atom_i, atom_j) params = (ke, qi, qj, σ, inter.σ6_fac) f = force_divr_with_cutoff(inter, r2, params, cutoff, force_units) @@ -194,18 +192,18 @@ function force_divr(::CoulombSoftCore, r2, invr2, (ke, qi, qj, σ, σ6_fac)) return ff * √invr2 end -@inline function potential_energy(inter::CoulombSoftCore{C}, +@inline function potential_energy(inter::CoulombSoftCore, dr, atom_i, atom_j, energy_units=u"kJ * mol^-1", - special::Bool=false, - args...) where C + special=false, + args...) r2 = sum(abs2, dr) cutoff = inter.cutoff ke = inter.coulomb_const qi, qj = atom_i.charge, atom_j.charge - σ = inter.lorentz_mixing ? (atom_i.σ + atom_j.σ) / 2 : sqrt(atom_i.σ * atom_j.σ) + σ = inter.σ_mixing(atom_i, atom_j) params = (ke, qi, qj, σ, inter.σ6_fac) pe = potential_with_cutoff(inter, r2, params, cutoff, energy_units) @@ -244,9 +242,8 @@ function CoulombReactionField(; use_neighbors=false, weight_special=1, coulomb_const=coulomb_const) - return CoulombReactionField{typeof(dist_cutoff), typeof(solvent_dielectric), - typeof(weight_special), typeof(coulomb_const)}( - dist_cutoff, solvent_dielectric, use_neighbors, weight_special, coulomb_const) + return CoulombReactionField(dist_cutoff, solvent_dielectric, use_neighbors, + weight_special, coulomb_const) end use_neighbors(inter::CoulombReactionField) = inter.use_neighbors @@ -255,7 +252,7 @@ function Base.zero(coul::CoulombReactionField{D, S, W, T}) where {D, S, W, T} return CoulombReactionField{D, S, W, T}( zero(D), zero(S), - false, + coul.use_neighbors, zero(W), zero(T), ) @@ -276,7 +273,7 @@ end atom_i, atom_j, force_units=u"kJ * mol^-1 * nm^-1", - special::Bool=false, + special=false, args...) r2 = sum(abs2, dr) if r2 > (inter.dist_cutoff ^ 2) @@ -309,7 +306,7 @@ end atom_i, atom_j, energy_units=u"kJ * mol^-1", - special::Bool=false, + special=false, args...) r2 = sum(abs2, dr) if r2 > (inter.dist_cutoff ^ 2) From dbce320493c880b698bc2a975d464258818555ce Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 18 Jul 2024 21:19:44 +0100 Subject: [PATCH 32/74] Gravity zero and + functions --- src/interactions/gravity.jl | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/interactions/gravity.jl b/src/interactions/gravity.jl index 02eabe250..3256daaa7 100644 --- a/src/interactions/gravity.jl +++ b/src/interactions/gravity.jl @@ -19,6 +19,12 @@ Gravity(; G=Unitful.G, use_neighbors=false) = Gravity{typeof(G)}(G, use_neighbor use_neighbors(inter::Gravity) = inter.use_neighbors +function Base.zero(gr::Gravity{T}) where T + return Gravity(zero(T), gr.use_neighbors) +end + +Base.:+(g1::Gravity, g2::Gravity) = Gravity(g1.G + g2.G, g1.use_neighbors) + @inline function force(inter::Gravity, dr, atom_i, From ef8b4ef4d56a7023b55a4ad2d3c52ecda3ce7dfe Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 19 Jul 2024 23:53:28 +0100 Subject: [PATCH 33/74] add atom_type to Atom, reorder fields --- benchmark/benchmarks.jl | 4 +- docs/src/differentiable.md | 6 +-- docs/src/examples.md | 5 ++- src/gradients.jl | 14 +++---- src/setup.jl | 77 +++++++++++++++++++++---------------- src/types.jl | 36 ++++++++--------- test/energy_conservation.jl | 2 +- test/gradients.jl | 2 +- test/simulation.jl | 29 +++++++------- 9 files changed, 93 insertions(+), 82 deletions(-) diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index a521ccb16..e07a86cd3 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -92,12 +92,12 @@ function test_sim(nl::Bool, parallel::Bool, f32::Bool, gpu::Bool) if gpu coords = CuArray(deepcopy(f32 ? starting_coords_f32 : starting_coords)) velocities = CuArray(deepcopy(f32 ? starting_velocities_f32 : starting_velocities)) - atoms = CuArray([Atom(charge=f32 ? 0.0f0 : 0.0, mass=atom_mass, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", + atoms = CuArray([Atom(mass=atom_mass, charge=f32 ? 0.0f0 : 0.0, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", ϵ=f32 ? 0.2f0u"kJ * mol^-1" : 0.2u"kJ * mol^-1") for i in 1:n_atoms]) else coords = deepcopy(f32 ? starting_coords_f32 : starting_coords) velocities = deepcopy(f32 ? starting_velocities_f32 : starting_velocities) - atoms = [Atom(charge=f32 ? 0.0f0 : 0.0, mass=atom_mass, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", + atoms = [Atom(mass=atom_mass, charge=f32 ? 0.0f0 : 0.0, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", ϵ=f32 ? 0.2f0u"kJ * mol^-1" : 0.2u"kJ * mol^-1") for i in 1:n_atoms] end diff --git a/docs/src/differentiable.md b/docs/src/differentiable.md index 79646ccb8..ad73ad6ef 100644 --- a/docs/src/differentiable.md +++ b/docs/src/differentiable.md @@ -76,7 +76,7 @@ simulator = VelocityVerlet( ) function loss(σ, coords, velocities) - atoms = [Atom(0, 0.0, atom_mass, σ, 0.2, false) for i in 1:n_atoms] + atoms = [Atom(0, 0, atom_mass, 0.0, σ, 0.2) for i in 1:n_atoms] loggers = (coords=CoordinateLogger(Float64, 10),) sys = System( @@ -184,7 +184,7 @@ simulator = VelocityVerlet( ) function loss(θ) - atoms = [Atom(0, 0.0, atom_mass, 0.0, 0.0, false) for i in 1:n_atoms] + atoms = [Atom(0, 0, atom_mass, 0.0, 0.0, 0.0) for i in 1:n_atoms] loggers = (coords=CoordinateLogger(Float64, 2),) specific_inter_lists = ( InteractionList2Atoms( @@ -326,7 +326,7 @@ simulator = VelocityVerlet( ) function loss() - atoms = [Atom(0, 0.0f0, mass, 0.0f0, 0.0f0, false) for i in 1:n_atoms] + atoms = [Atom(0, 0, mass, 0.0f0, 0.0f0, 0.0f0) for i in 1:n_atoms] loggers = (coords=CoordinateLogger(Float32, 10),) general_inters = (NNBonds(),) diff --git a/docs/src/examples.md b/docs/src/examples.md index 7537e72cc..e53cfe1e0 100644 --- a/docs/src/examples.md +++ b/docs/src/examples.md @@ -1008,8 +1008,9 @@ These paramaters must be added to the [`System`](@ref) manually by making use of updated_atoms = [] for i in eachindex(sys) - push!(updated_atoms, Atom(index=sys.atoms[i].index, charge=sys.atoms[i].charge, - mass=sys.atoms[i].mass, σ=σ, ϵ=ϵ, solute=sys.atoms[i].solute)) + push!(updated_atoms, Atom(index=sys.atoms[i].index, atom_type=sys.atoms[i].atom_type, + mass=sys.atoms[i].mass, charge=sys.atoms[i].charge, + σ=σ, ϵ=ϵ)) end sys = System(sys; atoms=[updated_atoms...]) diff --git a/src/gradients.jl b/src/gradients.jl index b621e37db..afad5a597 100644 --- a/src/gradients.jl +++ b/src/gradients.jl @@ -26,7 +26,6 @@ function extract_parameters(sys, ff) if inter isa LennardJones key_prefix = "inter_LJ_" params_dic[key_prefix * "weight_14"] = inter.weight_special - params_dic[key_prefix * "weight_solute_solvent"] = inter.weight_solute_solvent elseif inter isa Coulomb key_prefix = "inter_CO_" params_dic[key_prefix * "weight_14"] = inter.weight_special @@ -121,11 +120,11 @@ function inject_atom(at, at_data, params_dic) key_prefix = "atom_$(at_data.atom_type)_" Atom( at.index, - at.charge, # Residue-specific + at.atom_type, dict_get(params_dic, key_prefix * "mass" , at.mass), + at.charge, # Residue-specific dict_get(params_dic, key_prefix * "σ" , at.σ ), dict_get(params_dic, key_prefix * "ϵ" , at.ϵ ), - at.solute, ) end @@ -165,14 +164,15 @@ function inject_interaction_list(inter::InteractionList4Atoms, params_dic, gpu) InteractionList4Atoms(inter.is, inter.js, inter.ks, inter.ls, inters_grad, inter.types) end -function inject_interaction(inter::LennardJones{S, C, W, WS}, params_dic) where {S, C, W, WS} +function inject_interaction(inter::LennardJones, params_dic) key_prefix = "inter_LJ_" - LennardJones{S, C, W, WS}( + LennardJones( inter.cutoff, inter.use_neighbors, - inter.lorentz_mixing, + inter.shortcut, + inter.σ_mixing, + inter.ϵ_mixing, dict_get(params_dic, key_prefix * "weight_14", inter.weight_special), - dict_get(params_dic, key_prefix * "weight_solute_solvent", inter.weight_solute_solvent), ) end diff --git a/src/setup.jl b/src/setup.jl index 476fff683..a844aeffe 100644 --- a/src/setup.jl +++ b/src/setup.jl @@ -467,7 +467,7 @@ function System(coord_file::AbstractString, top = Chemfiles.Topology(frame) n_atoms = size(top) - atoms = Atom[] + atoms_abst = Atom[] atoms_data = AtomData[] bonds = InteractionList2Atoms(HarmonicBond) angles = InteractionList3Atoms(HarmonicAngle) @@ -553,11 +553,10 @@ function System(coord_file::AbstractString, if ismissing(ch) error("atom of type ", at.type, " has not had charge set") end - solute = res_id_to_standard[res_id] || res_name in ("ACE", "NME") if (units && at.σ < zero(T)u"nm") || (!units && at.σ < zero(T)) error("atom of type ", at.type, " has not had σ or ϵ set") end - push!(atoms, Atom(index=ai, charge=ch, mass=at.mass, σ=at.σ, ϵ=at.ϵ, solute=solute)) + push!(atoms_abst, Atom(index=ai, mass=at.mass, charge=ch, σ=at.σ, ϵ=at.ϵ)) push!(atoms_data, AtomData(atom_type=at_type, atom_name=atom_name, res_number=Chemfiles.id(res), res_name=Chemfiles.name(res), element=at.element)) eligible[ai, ai] = false @@ -886,7 +885,6 @@ function System(coord_file::AbstractString, end coords = wrap_coords.(coords, (boundary_used,)) - atoms = [atoms...] if gpu || !use_cell_list neighbor_finder = DistanceNeighborFinder( eligible=(gpu ? CuArray(eligible) : eligible), @@ -905,15 +903,18 @@ function System(coord_file::AbstractString, ) end if gpu - atoms = CuArray(atoms) - coords = CuArray(coords) + atoms = CuArray([atoms_abst...]) + coords_dev = CuArray(coords) + else + atoms = [atoms_abst...] + coords_dev = coords end if isnothing(velocities) if units - vels = zero(ustrip_vec.(coords))u"nm * ps^-1" + vels = zero(ustrip_vec.(coords_dev))u"nm * ps^-1" else - vels = zero(coords) + vels = zero(coords_dev) end else vels = velocities @@ -938,7 +939,7 @@ function System(coord_file::AbstractString, k = units ? Unitful.Na * Unitful.k : ustrip(u"kJ * K^-1 * mol^-1", Unitful.Na * Unitful.k) return System( atoms=atoms, - coords=coords, + coords=coords_dev, boundary=boundary_used, velocities=vels, atoms_data=atoms_data, @@ -976,7 +977,7 @@ function System(T::Type, atomnames = Dict{String, String}() name = "?" - atoms = Atom[] + atoms_abst = Atom[] atoms_data = AtomData[] bonds = InteractionList2Atoms(HarmonicBond) pairs = Tuple{Int, Int}[] @@ -1039,11 +1040,19 @@ function System(T::Type, # Take the first version of each atom type only if !haskey(atomtypes, atomname) if units - atomtypes[atomname] = Atom(charge=parse(T, c[5]), mass=parse(T, c[4])u"g/mol", - σ=parse(T, c[7])u"nm", ϵ=parse(T, c[8])u"kJ * mol^-1") + atomtypes[atomname] = Atom( + mass=parse(T, c[4])u"g/mol", + charge=parse(T, c[5]), + σ=parse(T, c[7])u"nm", + ϵ=parse(T, c[8])u"kJ * mol^-1", + ) else - atomtypes[atomname] = Atom(charge=parse(T, c[5]), mass=parse(T, c[4]), - σ=parse(T, c[7]), ϵ=parse(T, c[8])) + atomtypes[atomname] = Atom( + mass=parse(T, c[4]), + charge=parse(T, c[5]), + σ=parse(T, c[7]), + ϵ=parse(T, c[8]), + ) end end elseif current_field == "atoms" @@ -1054,10 +1063,9 @@ function System(T::Type, else atom_mass = parse(T, c[8]) end - solute = c[4] in standard_res_names - atom_index = length(atoms) + 1 - push!(atoms, Atom(index=atom_index, charge=ch, mass=atom_mass, σ=atomtypes[attype].σ, - ϵ=atomtypes[attype].ϵ, solute=solute)) + atom_index = length(atoms_abst) + 1 + push!(atoms_abst, Atom(index=atom_index, mass=atom_mass, charge=ch, σ=atomtypes[attype].σ, + ϵ=atomtypes[attype].ϵ)) push!(atoms_data, AtomData(atom_type=attype, atom_name=c[5], res_number=parse(Int, c[3]), res_name=c[4])) elseif current_field == "bonds" @@ -1136,26 +1144,26 @@ function System(T::Type, # Read coordinate file and add solvent atoms lines = readlines(coord_file) - coords = SArray[] + coords_abst = SArray[] for (i, l) in enumerate(lines[3:end-1]) coord = SVector(parse(T, l[21:28]), parse(T, l[29:36]), parse(T, l[37:44])) if units - push!(coords, (coord)u"nm") + push!(coords_abst, (coord)u"nm") else - push!(coords, coord) + push!(coords_abst, coord) end # Some atoms are not specified explicitly in the topology so are added here - if i > length(atoms) + if i > length(atoms_abst) atname = strip(l[11:15]) attype = replace(atname, r"\d+" => "") temp_charge = atomtypes[attype].charge if attype == "CL" # Temp hack to fix charges temp_charge = T(-1.0) end - atom_index = length(atoms) + 1 - push!(atoms, Atom(index=atom_index, charge=temp_charge, mass=atomtypes[attype].mass, - σ=atomtypes[attype].σ, ϵ=atomtypes[attype].ϵ, solute=false)) + atom_index = length(atoms_abst) + 1 + push!(atoms_abst, Atom(index=atom_index, mass=atomtypes[attype].mass, charge=temp_charge, + σ=atomtypes[attype].σ, ϵ=atomtypes[attype].ϵ)) push!(atoms_data, AtomData(atom_type=attype, atom_name=atname, res_number=parse(Int, l[1:5]), res_name=strip(l[6:10]))) @@ -1181,7 +1189,7 @@ function System(T::Type, end # Calculate matrix of pairs eligible for non-bonded interactions - n_atoms = length(coords) + n_atoms = length(coords_abst) eligible = trues(n_atoms, n_atoms) for i in 1:n_atoms eligible[i, i] = false @@ -1224,7 +1232,7 @@ function System(T::Type, else boundary_used = boundary end - coords = [coords...] + coords = [coords_abst...] if center_coords coords = coords .- (mean(coords),) .+ (box_center(boundary_used),) end @@ -1266,8 +1274,6 @@ function System(T::Type, end specific_inter_lists = tuple(specific_inter_array...) - atoms = [Atom(index=a.index, charge=a.charge, mass=a.mass, σ=a.σ, ϵ=a.ϵ, solute=a.solute) for a in atoms] - if gpu || !use_cell_list neighbor_finder = DistanceNeighborFinder( eligible=(gpu ? CuArray(eligible) : eligible), @@ -1286,15 +1292,18 @@ function System(T::Type, ) end if gpu - atoms = CuArray(atoms) - coords = CuArray(coords) + atoms = CuArray([atoms_abst...]) + coords_dev = CuArray(coords) + else + atoms = [atoms_abst...] + coords_dev = coords end if isnothing(velocities) if units - vels = zero(ustrip_vec.(coords))u"nm * ps^-1" + vels = zero(ustrip_vec.(coords_dev))u"nm * ps^-1" else - vels = zero(coords) + vels = zero(coords_dev) end else vels = velocities @@ -1303,7 +1312,7 @@ function System(T::Type, k = units ? Unitful.Na * Unitful.k : ustrip(u"kJ * K^-1 * mol^-1", Unitful.Na * Unitful.k) return System( atoms=atoms, - coords=coords, + coords=coords_dev, boundary=boundary_used, velocities=vels, atoms_data=atoms_data, diff --git a/src/types.jl b/src/types.jl index 7d6291165..fa46be671 100644 --- a/src/types.jl +++ b/src/types.jl @@ -9,8 +9,8 @@ export InteractionList3Atoms, InteractionList4Atoms, Atom, - charge, mass, + charge, AtomData, MolecularTopology, NeighborList, @@ -219,43 +219,43 @@ The types used should be bits types if the GPU is going to be used. # Arguments - `index::Int`: the index of the atom in the system. -- `charge::C=0.0`: the charge of the atom, used for electrostatic interactions. +- `atom_type::T`: the type of the atom. - `mass::M=1.0u"g/mol"`: the mass of the atom. +- `charge::C=0.0`: the charge of the atom, used for electrostatic interactions. - `σ::S=0.0u"nm"`: the Lennard-Jones finite distance at which the inter-particle potential is zero. - `ϵ::E=0.0u"kJ * mol^-1"`: the Lennard-Jones depth of the potential well. -- `solute::Bool=false`: whether the atom is part of the solute. """ -struct Atom{C, M, S, E} +struct Atom{T, M, C, S, E} index::Int - charge::C + atom_type::T mass::M + charge::C σ::S ϵ::E - solute::Bool end function Atom(; index=1, - charge=0.0, + atom_type=1, mass=1.0u"g/mol", + charge=0.0, σ=0.0u"nm", - ϵ=0.0u"kJ * mol^-1", - solute=false) - return Atom(index, charge, mass, σ, ϵ, solute) + ϵ=0.0u"kJ * mol^-1") + return Atom(index, atom_type, mass, charge, σ, ϵ) end -function Base.zero(::Type{Atom{T, T, T, T}}) where T - z = zero(T) - return Atom(0, z, z, z, z, false) +function Base.zero(::Atom{T, M, C, S, E}) where {T, M, C, S, E} + return Atom(0, zero(T), zero(M), zero(C), zero(S), zero(E)) end function Base.:+(a1::Atom, a2::Atom) - return Atom(0, a1.charge + a2.charge, a1.mass + a2.mass, a1.σ + a2.σ, a1.ϵ + a2.ϵ, false) + return Atom(a1.index, a1.atom_type, a1.mass + a2.mass, a1.charge + a2.charge, + a1.σ + a2.σ, a1.ϵ + a2.ϵ) end -function Base.:-(a1::Atom, a2::Atom) - return Atom(0, a1.charge - a2.charge, a1.mass - a2.mass, a1.σ - a2.σ, a1.ϵ - a2.ϵ, false) +function Base.getindex(atom::Atom, x::Symbol) + return hasfield(Atom, x) ? getfield(atom, x) : KeyError("no field $x in Atom") end """ @@ -286,8 +286,8 @@ function Base.getindex(at::Atom, x::Symbol) end function Base.show(io::IO, a::Atom) - print(io, "Atom with index ", a.index, ", charge=", charge(a), - ", mass=", mass(a), ", σ=", a.σ, ", ϵ=", a.ϵ) + print(io, "Atom with index=", a.index, ", atom_type=", a.atom_type, ", mass=", mass(a), + ", charge=", charge(a), ", σ=", a.σ, ", ϵ=", a.ϵ) end """ diff --git a/test/energy_conservation.jl b/test/energy_conservation.jl index 26c8f1f21..a27a10917 100644 --- a/test/energy_conservation.jl +++ b/test/energy_conservation.jl @@ -13,7 +13,7 @@ using Test boundary = CubicBoundary(50.0u"nm") simulator = VelocityVerlet(dt=0.001u"ps", remove_CM_motion=false) - atoms = [Atom(charge=0.0, mass=atom_mass, σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms] + atoms = [Atom(mass=atom_mass, charge=0.0, σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms] dist_cutoff = 3.0u"nm" cutoffs = ( DistanceCutoff(dist_cutoff), diff --git a/test/gradients.jl b/test/gradients.jl index f7b39edd8..caa554494 100644 --- a/test/gradients.jl +++ b/test/gradients.jl @@ -120,7 +120,7 @@ end ) function loss(σ, r0) - atoms = [Atom(i, i % 2 == 0 ? T(-0.02) : T(0.02), atom_mass, σ, T(0.2), false) for i in 1:n_atoms] + atoms = [Atom(i, 1, atom_mass, (i % 2 == 0 ? T(-0.02) : T(0.02)), σ, T(0.2)) for i in 1:n_atoms] bonds_inner = [HarmonicBond(T(100.0), bond_dists[i] * r0) for i in 1:(n_atoms ÷ 2)] bonds = InteractionList2Atoms( bond_is, diff --git a/test/simulation.jl b/test/simulation.jl index de460f5f2..cae0665ac 100644 --- a/test/simulation.jl +++ b/test/simulation.jl @@ -7,7 +7,7 @@ gen_temp_wrapper(s, args...; kwargs...) = temperature(s) s = System( - atoms=[Atom(charge=0.0, mass=10.0u"g/mol", σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms], + atoms=[Atom(mass=10.0u"g/mol", charge=0.0, σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms], coords=place_atoms(n_atoms, boundary; min_dist=0.3u"nm"), boundary=boundary, pairwise_inters=(LennardJones(use_neighbors=true),), @@ -69,7 +69,7 @@ end for n_threads in n_threads_list s = System( - atoms=[Atom(index=i, charge=0.0, mass=atom_mass, σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") + atoms=[Atom(index=i, mass=atom_mass, charge=0.0, σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms], coords=place_atoms(n_atoms, boundary; min_dist=0.3u"nm"), boundary=boundary, @@ -166,7 +166,7 @@ end simulator = VelocityVerlet(dt=0.002u"ps", coupling=AndersenThermostat(temp, 10.0u"ps")) s = System( - atoms=[Atom(charge=0.0, mass=10.0u"g/mol", σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms], + atoms=[Atom(mass=10.0u"g/mol", charge=0.0, σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms], coords=coords, boundary=boundary, pairwise_inters=(LennardJones(use_neighbors=true),), @@ -208,7 +208,7 @@ end ] s = System( - atoms=[Atom(charge=0.0, mass=10.0u"g/mol", σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms], + atoms=[Atom(mass=10.0u"g/mol", charge=0.0, σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms], coords=coords, boundary=boundary, pairwise_inters=(LennardJones(use_neighbors=true),), @@ -249,7 +249,7 @@ end simulator = VelocityVerlet(dt=0.002u"ps", coupling=BerendsenThermostat(temp, 1.0u"ps")) s = System( - atoms=[Atom(charge=0.0, mass=10.0u"g/mol", σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms], + atoms=[Atom(mass=10.0u"g/mol", charge=0.0, σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms], coords=coords, boundary=boundary, velocities=[random_velocity(10.0u"g/mol", temp) .* 0.01 for i in 1:n_atoms], @@ -309,7 +309,7 @@ end end s = System( - atoms=[Atom(charge=i % 2 == 0 ? -1.0 : 1.0, mass=10.0u"g/mol", σ=0.2u"nm", + atoms=[Atom(mass=10.0u"g/mol", charge=(i % 2 == 0 ? -1.0 : 1.0), σ=0.2u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms], coords=place_atoms(n_atoms, boundary; min_dist=0.2u"nm"), boundary=boundary, @@ -367,7 +367,7 @@ end V(sys::System, neighbors=nothing) = sys.velocities s = System( - atoms=[Atom(charge=0.0, mass=10.0u"g/mol", σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms], + atoms=[Atom(mass=10.0u"g/mol", charge=0.0, σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms], coords=coords, boundary=boundary, velocities=velocities, @@ -387,7 +387,7 @@ end vtype_nounits = eltype(ustrip_vec.(velocities)) s_nounits = System( - atoms=[Atom(charge=0.0, mass=10.0, σ=0.3, ϵ=0.2) for i in 1:n_atoms], + atoms=[Atom(mass=10.0, charge=0.0, σ=0.3, ϵ=0.2) for i in 1:n_atoms], coords=ustrip_vec.(coords), boundary=CubicBoundary(ustrip.(boundary)), velocities=ustrip_vec.(u"nm/ps",velocities), @@ -430,7 +430,7 @@ end n_steps = 2_000 boundary = CubicBoundary(2.0u"nm") starting_coords = place_atoms(n_atoms, boundary; min_dist=0.3u"nm") - atoms = [Atom(charge=0.0, mass=10.0u"g/mol", σ=0.2u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms] + atoms = [Atom(mass=10.0u"g/mol", charge=0.0, σ=0.2u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms] atoms_data = [AtomData(atom_type=(i <= n_atoms_res ? "A1" : "A2")) for i in 1:n_atoms] sim = Langevin(dt=0.001u"ps", temperature=300.0u"K", friction=1.0u"ps^-1") @@ -577,7 +577,7 @@ end coords = place_atoms(n_atoms, boundary; min_dist=0.3u"nm") velocities = [random_velocity(10.0u"g/mol", temp) .* 0.01 for i in 1:n_atoms] s1 = System( - atoms=[Atom(charge=0.0, mass=10.0u"g/mol", σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms], + atoms=[Atom( mass=10.0u"g/mol", charge=0.0,σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1") for i in 1:n_atoms], coords=coords, boundary=boundary, velocities=velocities, @@ -1101,8 +1101,9 @@ end updated_atoms = [] for i in eachindex(sys) - push!(updated_atoms, Atom(index=sys.atoms[i].index, charge=sys.atoms[i].charge, - mass=sys.atoms[i].mass, σ=σ, ϵ=ϵ, solute=sys.atoms[i].solute)) + push!(updated_atoms, Atom(index=sys.atoms[i].index, atom_type=sys.atoms[i].atom_type, + charge=sys.atoms[i].charge, mass=sys.atoms[i].mass, + σ=σ, ϵ=ϵ)) end sys = System(sys; atoms=[updated_atoms...]) @@ -1178,12 +1179,12 @@ end if gpu coords = CuArray(deepcopy(f32 ? starting_coords_f32 : starting_coords)) velocities = CuArray(deepcopy(f32 ? starting_velocities_f32 : starting_velocities)) - atoms = CuArray([Atom(charge=f32 ? 0.0f0 : 0.0, mass=atom_mass, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", + atoms = CuArray([Atom(mass=atom_mass, charge=f32 ? 0.0f0 : 0.0, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", ϵ=f32 ? 0.2f0u"kJ * mol^-1" : 0.2u"kJ * mol^-1") for i in 1:n_atoms]) else coords = deepcopy(f32 ? starting_coords_f32 : starting_coords) velocities = deepcopy(f32 ? starting_velocities_f32 : starting_velocities) - atoms = [Atom(charge=f32 ? 0.0f0 : 0.0, mass=atom_mass, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", + atoms = [Atom(mass=atom_mass, charge=f32 ? 0.0f0 : 0.0, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", ϵ=f32 ? 0.2f0u"kJ * mol^-1" : 0.2u"kJ * mol^-1") for i in 1:n_atoms] end From 3ac68992e17ee4367733d0379a8b5b6f8a27e9e5 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Mon, 22 Jul 2024 19:26:56 +0100 Subject: [PATCH 34/74] rename run_loggers! to apply_loggers! --- docs/src/documentation.md | 4 +-- src/loggers.jl | 10 +++---- src/simulators.jl | 58 +++++++++++++++++++-------------------- 3 files changed, 36 insertions(+), 36 deletions(-) diff --git a/docs/src/documentation.md b/docs/src/documentation.md index 042893434..fad0ab1d7 100644 --- a/docs/src/documentation.md +++ b/docs/src/documentation.md @@ -949,8 +949,8 @@ function Molly.simulate!(sys, remove_CM_motion!(sys) # Apply the loggers like this - # Computed quantities can also be given as keyword arguments to run_loggers! - run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads) + # Computed quantities can also be given as keyword arguments to apply_loggers! + apply_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads) # Find new neighbors like this neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n, recompute_forces; diff --git a/src/loggers.jl b/src/loggers.jl index b3a278c3a..b74630904 100644 --- a/src/loggers.jl +++ b/src/loggers.jl @@ -1,7 +1,7 @@ # Loggers to record properties throughout a simulation export - run_loggers!, + apply_loggers!, GeneralObservableLogger, values, log_property!, @@ -22,8 +22,8 @@ export MonteCarloLogger """ - run_loggers!(system, neighbors=nothing, step_n=0, run_loggers=true; - n_threads=Threads.nthreads(), kwargs...) + apply_loggers!(system, neighbors=nothing, step_n=0, run_loggers=true; + n_threads=Threads.nthreads(), kwargs...) Run the loggers associated with a system. @@ -32,8 +32,8 @@ are not run before the first step. Additional keyword arguments can be passed to the loggers if required. Ignored for gradient calculation during automatic differentiation. """ -function run_loggers!(sys::System, neighbors=nothing, step_n::Integer=0, run_loggers=true; - n_threads::Integer=Threads.nthreads(), kwargs...) +function apply_loggers!(sys::System, neighbors=nothing, step_n::Integer=0, run_loggers=true; + n_threads::Integer=Threads.nthreads(), kwargs...) if run_loggers == true || (run_loggers == :skipzero && step_n != 0) for logger in values(sys.loggers) log_property!(logger, sys, neighbors, step_n; n_threads=n_threads, kwargs...) diff --git a/src/simulators.jl b/src/simulators.jl index a3afc7731..456487b98 100644 --- a/src/simulators.jl +++ b/src/simulators.jl @@ -66,7 +66,7 @@ function simulate!(sys, sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) E = potential_energy(sys, neighbors; n_threads=n_threads) - run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads, current_potential_energy=E) + apply_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads, current_potential_energy=E) using_constraints = length(sys.constraints) > 0 println(sim.log_stream, "Step 0 - potential energy ", E, " - max force N/A - N/A") hn = sim.step_size @@ -103,8 +103,8 @@ function simulate!(sys, E_trial, " - max force ", max_force, " - rejected") end - run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, - current_potential_energy=E) + apply_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, + current_potential_energy=E) if max_force < sim.tol break @@ -148,7 +148,7 @@ function simulate!(sys, forces_t_dt = forces_nounits_t_dt .* sys.force_units forces_buffer = init_forces_buffer(forces_nounits_t_dt, n_threads) accels_t_dt = zero(accels_t) - run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads, current_forces=forces_t) + apply_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads, current_forces=forces_t) using_constraints = length(sys.constraints) > 0 if using_constraints cons_coord_storage = similar(sys.coords) @@ -191,8 +191,8 @@ function simulate!(sys, accels_t .= accels_t_dt end - run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, - current_forces=forces_t) + apply_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, + current_forces=forces_t) end return sys end @@ -229,7 +229,7 @@ function simulate!(sys, sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) !iszero(sim.remove_CM_motion) && remove_CM_motion!(sys) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) - run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) + apply_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) forces_nounits_t = ustrip_vec.(similar(sys.coords)) forces_t = forces_nounits_t .* sys.force_units forces_buffer = init_forces_buffer(forces_nounits_t, n_threads) @@ -269,8 +269,8 @@ function simulate!(sys, neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n, recompute_forces; n_threads=n_threads) - run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, - current_forces=forces_t) + apply_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, + current_forces=forces_t) end return sys end @@ -303,7 +303,7 @@ function simulate!(sys, run_loggers=true) sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) - run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) + apply_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) coords_last, coords_copy = similar(sys.coords), similar(sys.coords) forces_nounits_t = ustrip_vec.(similar(sys.coords)) forces_t = forces_nounits_t .* sys.force_units @@ -339,8 +339,8 @@ function simulate!(sys, n_threads=n_threads) coords_last .= coords_copy - run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, - current_forces=forces_t) + apply_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, + current_forces=forces_t) end return sys end @@ -387,7 +387,7 @@ function simulate!(sys, sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) !iszero(sim.remove_CM_motion) && remove_CM_motion!(sys) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) - run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) + apply_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) forces_nounits_t = ustrip_vec.(similar(sys.coords)) forces_t = forces_nounits_t .* sys.force_units forces_buffer = init_forces_buffer(forces_nounits_t, n_threads) @@ -432,8 +432,8 @@ function simulate!(sys, neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n, recompute_forces; n_threads=n_threads) - run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, - current_forces=forces_t) + apply_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, + current_forces=forces_t) end return sys end @@ -496,7 +496,7 @@ function simulate!(sys, sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) !iszero(sim.remove_CM_motion) && remove_CM_motion!(sys) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) - run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) + apply_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) forces_nounits_t = ustrip_vec.(similar(sys.coords)) forces_t = forces_nounits_t .* sys.force_units forces_buffer = init_forces_buffer(forces_nounits_t, n_threads) @@ -548,7 +548,7 @@ function simulate!(sys, neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n; n_threads=n_threads) - run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads) + apply_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads) end return sys end @@ -616,7 +616,7 @@ function simulate!(sys, sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) !iszero(sim.remove_CM_motion) && remove_CM_motion!(sys) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) - run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) + apply_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) forces_nounits_t = ustrip_vec.(similar(sys.coords)) forces_t = forces_nounits_t .* sys.force_units forces_buffer = init_forces_buffer(forces_nounits_t, n_threads) @@ -640,8 +640,8 @@ function simulate!(sys, neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n; n_threads=n_threads) - run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, - current_forces=forces_t) + apply_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, + current_forces=forces_t) end return sys end @@ -693,7 +693,7 @@ function simulate!(sys, sim::NoseHoover, n_steps::Integer; forces_t_dt = forces_nounits_t_dt .* sys.force_units forces_buffer = init_forces_buffer(forces_nounits_t_dt, n_threads) accels_t_dt = zero(accels_t) - run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads, current_forces=forces_t) + apply_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads, current_forces=forces_t) v_half = zero(sys.velocities) zeta = zero(inv(sim.dt)) @@ -735,8 +735,8 @@ function simulate!(sys, sim::NoseHoover, n_steps::Integer; accels_t .= accels_t_dt end - run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, - current_forces=forces_t) + apply_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, + current_forces=forces_t) end return sys end @@ -1058,15 +1058,15 @@ function simulate!(sys::System{D, G, T}, ΔE = E_new - E_old δ = ΔE / (sys.k * sim.temperature) if δ < 0 || (rand() < exp(-δ)) - run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, - current_potential_energy=E_new, success=true, - energy_rate=(E_new / (sys.k * sim.temperature))) + apply_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, + current_potential_energy=E_new, success=true, + energy_rate=(E_new / (sys.k * sim.temperature))) E_old = E_new else sys.coords .= coords_old - run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, - current_potential_energy=E_old, success=false, - energy_rate=(E_old / (sys.k * sim.temperature))) + apply_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, + current_potential_energy=E_old, success=false, + energy_rate=(E_old / (sys.k * sim.temperature))) end end From e9a362b9d96ab45c25da86d3538fbc92248aec3c Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Mon, 22 Jul 2024 19:52:12 +0100 Subject: [PATCH 35/74] Use @kwdef for interactions --- src/interactions/buckingham.jl | 28 ++++------- src/interactions/coulomb.jl | 80 ++++++++++--------------------- src/interactions/gravity.jl | 8 ++-- src/interactions/lennard_jones.jl | 62 +++++++----------------- src/interactions/soft_sphere.jl | 21 +++----- 5 files changed, 59 insertions(+), 140 deletions(-) diff --git a/src/interactions/buckingham.jl b/src/interactions/buckingham.jl index 71d68f4bf..10bb2fb02 100644 --- a/src/interactions/buckingham.jl +++ b/src/interactions/buckingham.jl @@ -45,26 +45,14 @@ C_{ij} &= (C_{ii} C_{jj})^{1/2} ``` so atoms that use this interaction should have fields `A`, `B` and `C` available. """ -struct Buckingham{C, W} <: PairwiseInteraction - cutoff::C - use_neighbors::Bool - shortcut::Function - A_mixing::Function - B_mixing::Function - C_mixing::Function - weight_special::W -end - -function Buckingham(; - cutoff=NoCutoff(), - use_neighbors=false, - shortcut=buckingham_zero_shortcut, - A_mixing=geometric_A_mixing, - B_mixing=inverse_B_mixing, - C_mixing=geometric_C_mixing, - weight_special=1) - return Buckingham(cutoff, use_neighbors, shortcut, A_mixing, B_mixing, - C_mixing, weight_special) +@kwdef struct Buckingham{C, W} <: PairwiseInteraction + cutoff::C = NoCutoff() + use_neighbors::Bool = false + shortcut::Function = buckingham_zero_shortcut + A_mixing::Function = geometric_A_mixing + B_mixing::Function = inverse_B_mixing + C_mixing::Function = geometric_C_mixing + weight_special::W = 1 end use_neighbors(inter::Buckingham) = inter.use_neighbors diff --git a/src/interactions/coulomb.jl b/src/interactions/coulomb.jl index 49f6f711b..b9691ec45 100644 --- a/src/interactions/coulomb.jl +++ b/src/interactions/coulomb.jl @@ -3,6 +3,8 @@ export CoulombSoftCore, CoulombReactionField +const coulomb_const = 138.93545764u"kJ * mol^-1 * nm" # 1 / 4πϵ0 + @doc raw""" Coulomb(; cutoff, use_neighbors, weight_special, coulomb_const) @@ -13,21 +15,11 @@ The potential energy is defined as V(r_{ij}) = \frac{q_i q_j}{4 \pi \varepsilon_0 r_{ij}} ``` """ -struct Coulomb{C, W, T} <: PairwiseInteraction - cutoff::C - use_neighbors::Bool - weight_special::W - coulomb_const::T -end - -const coulomb_const = 138.93545764u"kJ * mol^-1 * nm" # 1 / 4πϵ0 - -function Coulomb(; - cutoff=NoCutoff(), - use_neighbors=false, - weight_special=1, - coulomb_const=coulomb_const) - return Coulomb(cutoff, use_neighbors, weight_special, coulomb_const) +@kwdef struct Coulomb{C, W, T} <: PairwiseInteraction + cutoff::C = NoCutoff() + use_neighbors::Bool = false + weight_special::W = 1 + coulomb_const::T = coulomb_const end use_neighbors(inter::Coulomb) = inter.use_neighbors @@ -106,30 +98,16 @@ V(r_{ij}) = \frac{q_i q_j}{4 \pi \varepsilon_0 (r_{ij}^6 + \alpha \sigma_{ij}^6 ``` If ``\alpha`` or ``\lambda`` are zero this gives the standard [`Coulomb`](@ref) potential. """ -struct CoulombSoftCore{C, A, L, P, W, T, R} <: PairwiseInteraction - cutoff::C - α::A - λ::L - p::P - use_neighbors::Bool - σ_mixing::Function - weight_special::W - coulomb_const::T - σ6_fac::R -end - -function CoulombSoftCore(; - cutoff=NoCutoff(), - α=1, - λ=0, - p=2, - use_neighbors=false, - σ_mixing=lorentz_σ_mixing, - weight_special=1, - coulomb_const=coulomb_const) - σ6_fac = α * λ^p - return CoulombSoftCore(cutoff, α, λ, p, use_neighbors, σ_mixing, weight_special, - coulomb_const, σ6_fac) +@kwdef struct CoulombSoftCore{C, A, L, P, W, T, R} <: PairwiseInteraction + cutoff::C = NoCutoff() + α::A = 1 + λ::L = 0 + p::P = 2 + use_neighbors::Bool = false + σ_mixing::Function = lorentz_σ_mixing + weight_special::W = 1 + coulomb_const::T = coulomb_const + σ6_fac::R = α * λ^p end use_neighbors(inter::CoulombSoftCore) = inter.use_neighbors @@ -219,6 +197,8 @@ function potential(::CoulombSoftCore, r2, invr2, (ke, qi, qj, σ, σ6_fac)) return (ke * qi * qj) * √cbrt(inv_rsc6) end +const crf_solvent_dielectric = 78.3 + """ CoulombReactionField(; dist_cutoff, solvent_dielectric, use_neighbors, weight_special, coulomb_const) @@ -226,24 +206,12 @@ end The Coulomb electrostatic interaction modified using the reaction field approximation between two atoms. """ -struct CoulombReactionField{D, S, W, T} <: PairwiseInteraction +@kwdef struct CoulombReactionField{D, S, W, T} <: PairwiseInteraction dist_cutoff::D - solvent_dielectric::S - use_neighbors::Bool - weight_special::W - coulomb_const::T -end - -const crf_solvent_dielectric = 78.3 - -function CoulombReactionField(; - dist_cutoff, - solvent_dielectric=crf_solvent_dielectric, - use_neighbors=false, - weight_special=1, - coulomb_const=coulomb_const) - return CoulombReactionField(dist_cutoff, solvent_dielectric, use_neighbors, - weight_special, coulomb_const) + solvent_dielectric::S = crf_solvent_dielectric + use_neighbors::Bool = false + weight_special::W = 1 + coulomb_const::T = coulomb_const end use_neighbors(inter::CoulombReactionField) = inter.use_neighbors diff --git a/src/interactions/gravity.jl b/src/interactions/gravity.jl index 3256daaa7..1b4927d41 100644 --- a/src/interactions/gravity.jl +++ b/src/interactions/gravity.jl @@ -10,13 +10,11 @@ The potential energy is defined as V(r_{ij}) = -\frac{G m_i m_j}{r_{ij}} ``` """ -struct Gravity{T} <: PairwiseInteraction - G::T - use_neighbors::Bool +@kwdef struct Gravity{T} <: PairwiseInteraction + G::T = Unitful.G + use_neighbors::Bool = false end -Gravity(; G=Unitful.G, use_neighbors=false) = Gravity{typeof(G)}(G, use_neighbors) - use_neighbors(inter::Gravity) = inter.use_neighbors function Base.zero(gr::Gravity{T}) where T diff --git a/src/interactions/lennard_jones.jl b/src/interactions/lennard_jones.jl index dbd4b24a4..c3ca90d7f 100644 --- a/src/interactions/lennard_jones.jl +++ b/src/interactions/lennard_jones.jl @@ -38,24 +38,13 @@ and the force on each atom by \end{aligned} ``` """ -struct LennardJones{C, W} <: PairwiseInteraction - cutoff::C - use_neighbors::Bool - shortcut::Function - σ_mixing::Function - ϵ_mixing::Function - weight_special::W -end - -function LennardJones(; - cutoff=NoCutoff(), - use_neighbors=false, - shortcut=lj_zero_shortcut, - σ_mixing=lorentz_σ_mixing, - ϵ_mixing=geometric_ϵ_mixing, - weight_special=1) - return LennardJones(cutoff, use_neighbors, shortcut, σ_mixing, ϵ_mixing, - weight_special) +@kwdef struct LennardJones{C, W} <: PairwiseInteraction + cutoff::C = NoCutoff() + use_neighbors::Bool = false + shortcut::Function = lj_zero_shortcut + σ_mixing::Function = lorentz_σ_mixing + ϵ_mixing::Function = geometric_ϵ_mixing + weight_special::W = 1 end use_neighbors(inter::LennardJones) = inter.use_neighbors @@ -164,32 +153,17 @@ r_{ij}^{\text{sc}} = \left(r_{ij}^6 + \alpha \sigma_{ij}^6 \lambda^p \right)^{1/ ``` If ``\alpha`` or ``\lambda`` are zero this gives the standard [`LennardJones`](@ref) potential. """ -struct LennardJonesSoftCore{C, A, L, P, W, R} <: PairwiseInteraction - cutoff::C - α::A - λ::L - p::P - use_neighbors::Bool - shortcut::Function - σ_mixing::Function - ϵ_mixing::Function - weight_special::W - σ6_fac::R -end - -function LennardJonesSoftCore(; - cutoff=NoCutoff(), - α=1, - λ=0, - p=2, - use_neighbors=false, - shortcut=lj_zero_shortcut, - σ_mixing=lorentz_σ_mixing, - ϵ_mixing=geometric_ϵ_mixing, - weight_special=1) - σ6_fac = α * λ^p - return LennardJonesSoftCore(cutoff, α, λ, p, use_neighbors, shortcut, σ_mixing, - ϵ_mixing, weight_special, σ6_fac) +@kwdef struct LennardJonesSoftCore{C, A, L, P, W, R} <: PairwiseInteraction + cutoff::C = NoCutoff() + α::A = 1 + λ::L = 0 + p::P = 2 + use_neighbors::Bool = false + shortcut::Function = lj_zero_shortcut + σ_mixing::Function = lorentz_σ_mixing + ϵ_mixing::Function = geometric_ϵ_mixing + weight_special::W = 1 + σ6_fac::R = α * λ^p end use_neighbors(inter::LennardJonesSoftCore) = inter.use_neighbors diff --git a/src/interactions/soft_sphere.jl b/src/interactions/soft_sphere.jl index a513e7163..2e54918e4 100644 --- a/src/interactions/soft_sphere.jl +++ b/src/interactions/soft_sphere.jl @@ -10,21 +10,12 @@ The potential energy is defined as V(r_{ij}) = 4\varepsilon_{ij} \left(\frac{\sigma_{ij}}{r_{ij}}\right)^{12} ``` """ -struct SoftSphere{C} <: PairwiseInteraction - cutoff::C - use_neighbors::Bool - shortcut::Function - σ_mixing::Function - ϵ_mixing::Function -end - -function SoftSphere(; - cutoff=NoCutoff(), - use_neighbors=false, - shortcut=lj_zero_shortcut, - σ_mixing=lorentz_σ_mixing, - ϵ_mixing=geometric_ϵ_mixing) - return SoftSphere(cutoff, use_neighbors, shortcut, σ_mixing, ϵ_mixing) +@kwdef struct SoftSphere{C} <: PairwiseInteraction + cutoff::C = NoCutoff() + use_neighbors::Bool = false + shortcut::Function = lj_zero_shortcut + σ_mixing::Function = lorentz_σ_mixing + ϵ_mixing::Function = geometric_ϵ_mixing end use_neighbors(inter::SoftSphere) = inter.use_neighbors From 06d1c17ed2ee3c857f811834d7ceb2d31231dc14 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Mon, 22 Jul 2024 20:14:23 +0100 Subject: [PATCH 36/74] concrete function typing for interactions --- src/interactions/buckingham.jl | 10 +++++----- src/interactions/coulomb.jl | 6 +++--- src/interactions/lennard_jones.jl | 20 ++++++++++---------- src/interactions/mie.jl | 10 +++++----- src/interactions/soft_sphere.jl | 8 ++++---- 5 files changed, 27 insertions(+), 27 deletions(-) diff --git a/src/interactions/buckingham.jl b/src/interactions/buckingham.jl index 10bb2fb02..87eaa8831 100644 --- a/src/interactions/buckingham.jl +++ b/src/interactions/buckingham.jl @@ -45,13 +45,13 @@ C_{ij} &= (C_{ii} C_{jj})^{1/2} ``` so atoms that use this interaction should have fields `A`, `B` and `C` available. """ -@kwdef struct Buckingham{C, W} <: PairwiseInteraction +@kwdef struct Buckingham{C, S, A, B, M, W} <: PairwiseInteraction cutoff::C = NoCutoff() use_neighbors::Bool = false - shortcut::Function = buckingham_zero_shortcut - A_mixing::Function = geometric_A_mixing - B_mixing::Function = inverse_B_mixing - C_mixing::Function = geometric_C_mixing + shortcut::S = buckingham_zero_shortcut + A_mixing::A = geometric_A_mixing + B_mixing::B = inverse_B_mixing + C_mixing::M = geometric_C_mixing weight_special::W = 1 end diff --git a/src/interactions/coulomb.jl b/src/interactions/coulomb.jl index b9691ec45..b178897b0 100644 --- a/src/interactions/coulomb.jl +++ b/src/interactions/coulomb.jl @@ -98,13 +98,13 @@ V(r_{ij}) = \frac{q_i q_j}{4 \pi \varepsilon_0 (r_{ij}^6 + \alpha \sigma_{ij}^6 ``` If ``\alpha`` or ``\lambda`` are zero this gives the standard [`Coulomb`](@ref) potential. """ -@kwdef struct CoulombSoftCore{C, A, L, P, W, T, R} <: PairwiseInteraction +@kwdef struct CoulombSoftCore{C, A, L, P, S, W, T, R} <: PairwiseInteraction cutoff::C = NoCutoff() α::A = 1 λ::L = 0 p::P = 2 use_neighbors::Bool = false - σ_mixing::Function = lorentz_σ_mixing + σ_mixing::S = lorentz_σ_mixing weight_special::W = 1 coulomb_const::T = coulomb_const σ6_fac::R = α * λ^p @@ -112,7 +112,7 @@ end use_neighbors(inter::CoulombSoftCore) = inter.use_neighbors -function Base.zero(coul::CoulombSoftCore{C, A, L, P, W, T, R}) where {C, A, L, P, W, T, R} +function Base.zero(coul::CoulombSoftCore{C, A, L, P, S, W, T, R}) where {C, A, L, P, S, W, T, R} return CoulombSoftCore( coul.cutoff, zero(A), diff --git a/src/interactions/lennard_jones.jl b/src/interactions/lennard_jones.jl index c3ca90d7f..37341cd19 100644 --- a/src/interactions/lennard_jones.jl +++ b/src/interactions/lennard_jones.jl @@ -38,18 +38,18 @@ and the force on each atom by \end{aligned} ``` """ -@kwdef struct LennardJones{C, W} <: PairwiseInteraction +@kwdef struct LennardJones{C, H, S, E, W} <: PairwiseInteraction cutoff::C = NoCutoff() use_neighbors::Bool = false - shortcut::Function = lj_zero_shortcut - σ_mixing::Function = lorentz_σ_mixing - ϵ_mixing::Function = geometric_ϵ_mixing + shortcut::H = lj_zero_shortcut + σ_mixing::S = lorentz_σ_mixing + ϵ_mixing::E = geometric_ϵ_mixing weight_special::W = 1 end use_neighbors(inter::LennardJones) = inter.use_neighbors -function Base.zero(lj::LennardJones{C, W}) where {C, W} +function Base.zero(lj::LennardJones{C, H, S, E, W}) where {C, H, S, E, W} return LennardJones( lj.cutoff, lj.use_neighbors, @@ -153,22 +153,22 @@ r_{ij}^{\text{sc}} = \left(r_{ij}^6 + \alpha \sigma_{ij}^6 \lambda^p \right)^{1/ ``` If ``\alpha`` or ``\lambda`` are zero this gives the standard [`LennardJones`](@ref) potential. """ -@kwdef struct LennardJonesSoftCore{C, A, L, P, W, R} <: PairwiseInteraction +@kwdef struct LennardJonesSoftCore{C, A, L, P, H, S, E, W, R} <: PairwiseInteraction cutoff::C = NoCutoff() α::A = 1 λ::L = 0 p::P = 2 use_neighbors::Bool = false - shortcut::Function = lj_zero_shortcut - σ_mixing::Function = lorentz_σ_mixing - ϵ_mixing::Function = geometric_ϵ_mixing + shortcut::H = lj_zero_shortcut + σ_mixing::S = lorentz_σ_mixing + ϵ_mixing::E = geometric_ϵ_mixing weight_special::W = 1 σ6_fac::R = α * λ^p end use_neighbors(inter::LennardJonesSoftCore) = inter.use_neighbors -function Base.zero(lj::LennardJonesSoftCore{C, A, L, P, W, R}) where {C, A, L, P, W, R} +function Base.zero(lj::LennardJonesSoftCore{C, A, L, P, H, S, E, W, R}) where {C, A, L, P, H, S, E, W, R} return LennardJonesSoftCore( lj.cutoff, zero(A), diff --git a/src/interactions/mie.jl b/src/interactions/mie.jl index bf683b659..7e76c3551 100644 --- a/src/interactions/mie.jl +++ b/src/interactions/mie.jl @@ -15,14 +15,14 @@ where C = \frac{n}{n - m} \left( \frac{n}{m} \right) ^\frac{m}{n - m} ``` """ -struct Mie{C, T} <: PairwiseInteraction +struct Mie{T, C, H, S, E} <: PairwiseInteraction m::T n::T cutoff::C use_neighbors::Bool - shortcut::Function - σ_mixing::Function - ϵ_mixing::Function + shortcut::H + σ_mixing::S + ϵ_mixing::E mn_fac::T end @@ -40,7 +40,7 @@ end use_neighbors(inter::Mie) = inter.use_neighbors -function Base.zero(m::Mie{C, T}) where {C, T} +function Base.zero(m::Mie{T}) where T return Mie(zero(T), zero(T), m.cutoff, m.use_neighbors, m.shortcut, m.σ_mixing, m.ϵ_mixing, zero(T)) end diff --git a/src/interactions/soft_sphere.jl b/src/interactions/soft_sphere.jl index 2e54918e4..e64c5ec46 100644 --- a/src/interactions/soft_sphere.jl +++ b/src/interactions/soft_sphere.jl @@ -10,12 +10,12 @@ The potential energy is defined as V(r_{ij}) = 4\varepsilon_{ij} \left(\frac{\sigma_{ij}}{r_{ij}}\right)^{12} ``` """ -@kwdef struct SoftSphere{C} <: PairwiseInteraction +@kwdef struct SoftSphere{C, H, S, E} <: PairwiseInteraction cutoff::C = NoCutoff() use_neighbors::Bool = false - shortcut::Function = lj_zero_shortcut - σ_mixing::Function = lorentz_σ_mixing - ϵ_mixing::Function = geometric_ϵ_mixing + shortcut::H = lj_zero_shortcut + σ_mixing::S = lorentz_σ_mixing + ϵ_mixing::E = geometric_ϵ_mixing end use_neighbors(inter::SoftSphere) = inter.use_neighbors From 8beea61369213af05924d5663c6305f7e9ea5dd8 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Tue, 23 Jul 2024 12:01:23 +0100 Subject: [PATCH 37/74] fix tests --- src/interactions/implicit_solvent.jl | 8 ++++---- test/protein.jl | 2 +- test/runtests.jl | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/interactions/implicit_solvent.jl b/src/interactions/implicit_solvent.jl index 37f63c311..1cf40ab49 100644 --- a/src/interactions/implicit_solvent.jl +++ b/src/interactions/implicit_solvent.jl @@ -360,7 +360,7 @@ struct ImplicitSolventOBC{T, D, V, K, S, F, I, DI} <: AbstractGBSA srjs::DI end -function ImplicitSolventOBC(atoms::AbstractArray{Atom{T, M, D, E}}, +function ImplicitSolventOBC(atoms::AbstractArray{Atom{TY, M, T, D, E}}, atoms_data, bonds; solvent_dielectric=gb_solvent_dielectric, @@ -373,7 +373,7 @@ function ImplicitSolventOBC(atoms::AbstractArray{Atom{T, M, D, E}}, use_ACE=true, use_OBC2=false, element_to_radius=mbondi2_element_to_radius, - element_to_screen=obc_element_to_screen) where {T, M, D, E} + element_to_screen=obc_element_to_screen) where {TY, M, T, D, E} units = dimension(D) == u"𝐋" radii = mbondi2_radii(atoms_data, bonds; element_to_radius=element_to_radius) @@ -474,7 +474,7 @@ struct ImplicitSolventGBN2{T, D, VT, VD, K, S, F, I, TD, TM, DI} <: AbstractGBSA srjs::DI end -function ImplicitSolventGBN2(atoms::AbstractArray{Atom{T, M, D, E}}, +function ImplicitSolventGBN2(atoms::AbstractArray{Atom{TY, M, T, D, E}}, atoms_data, bonds; solvent_dielectric=gb_solvent_dielectric, @@ -493,7 +493,7 @@ function ImplicitSolventGBN2(atoms::AbstractArray{Atom{T, M, D, E}}, atom_params=gbn2_atom_params, atom_params_nucleic=gbn2_atom_params_nucleic, data_d0=gbn2_data_d0, - data_m0=gbn2_data_m0) where {T, M, D, E} + data_m0=gbn2_data_m0) where {TY, M, T, D, E} units = dimension(D) == u"𝐋" radii = mbondi3_radii(atoms_data, bonds; element_to_radius=element_to_radius) nucleic_acid_residues = ("A", "C", "G", "U", "DA", "DC", "DG", "DT") diff --git a/test/protein.jl b/test/protein.jl index 9bf966b58..64331880a 100644 --- a/test/protein.jl +++ b/test/protein.jl @@ -171,7 +171,7 @@ end @test maximum(maximum(abs.(v)) for v in vels_diff ) < 1e-6u"nm * ps^-1" params_dic = extract_parameters(sys_nounits, ff_nounits) - @test length(params_dic) == 639 + @test length(params_dic) == 638 atoms_grad, pis_grad, sis_grad, gis_grad = inject_gradients(sys_nounits, params_dic) @test atoms_grad == sys_nounits.atoms @test pis_grad == sys_nounits.pairwise_inters diff --git a/test/runtests.jl b/test/runtests.jl index dac047890..3383a7beb 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -33,8 +33,8 @@ end const run_visualize_tests = get(ENV, "VISTESTS", "1") != "0" if run_visualize_tests - using GLMakie -elseif get(ENV, "VISTESTS", "1") == "0" + import GLMakie +else @warn "The visualization tests will not be run as VISTESTS is set to 0" end From 67ef5dd4ecdb92f63928e40ce80c02d57e29baa5 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Tue, 23 Jul 2024 16:38:31 +0100 Subject: [PATCH 38/74] remove abstract interaction types --- docs/src/documentation.md | 14 ++++---- docs/src/examples.md | 8 ++--- src/energy.jl | 31 ++++++++--------- src/force.jl | 34 +++++++++---------- src/interactions/buckingham.jl | 2 +- src/interactions/cosine_angle.jl | 2 +- src/interactions/coulomb.jl | 6 ++-- src/interactions/fene_bond.jl | 2 +- src/interactions/gravity.jl | 2 +- src/interactions/harmonic_angle.jl | 2 +- src/interactions/harmonic_bond.jl | 2 +- .../harmonic_position_restraint.jl | 2 +- src/interactions/lennard_jones.jl | 4 +-- src/interactions/mie.jl | 2 +- src/interactions/morse_bond.jl | 2 +- src/interactions/periodic_torsion.jl | 2 +- src/interactions/rb_torsion.jl | 2 +- src/interactions/soft_sphere.jl | 2 +- src/neighbors.jl | 12 +++++++ src/types.jl | 28 --------------- test/agent.jl | 4 +-- 21 files changed, 73 insertions(+), 92 deletions(-) diff --git a/docs/src/documentation.md b/docs/src/documentation.md index fad0ab1d7..7e6a6648e 100644 --- a/docs/src/documentation.md +++ b/docs/src/documentation.md @@ -534,8 +534,8 @@ The force on each particle in the system is derived from the potential correspon ``` In Molly there are three types of interactions: -- [`PairwiseInteraction`](@ref)s are present between all or most atom pairs, and account for example for non-bonded terms in molecular mechanics force fields. -- [`SpecificInteraction`](@ref)s are present between specific atoms, and account for example for bonded terms in molecular mechanics force fields. +- Pairwise interactions are present between all or most atom pairs, and account for example for non-bonded terms in molecular mechanics force fields. +- Specific interactions are present between specific atoms, and account for example for bonded terms in molecular mechanics force fields. - General interactions are a free-form interaction type that can access the whole system and outputs forces for all atoms. This is useful for neural network potentials, implicit solvent models and other cases that require maximum flexibility. General interactions should be compatible with the [AtomsCalculators.jl](https://github.com/JuliaMolSim/AtomsCalculators.jl) interface. The available pairwise interactions are: @@ -567,9 +567,9 @@ The available general interactions are: ### Pairwise interactions -To define your own [`PairwiseInteraction`](@ref), first define the `struct`: +To define your own pairwise interaction, first define the `struct`: ```julia -struct MyPairwiseInter <: PairwiseInteraction +struct MyPairwiseInter # Any properties, e.g. constants for the interaction or cutoff parameters end ``` @@ -632,7 +632,7 @@ Note that you can also use a named tuple instead of a tuple if you want to acces ```julia pairwise_inters = (MyPairwiseInter=MyPairwiseInter(),) ``` -For performance reasons it is best to [avoid containers with abstract type parameters](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-abstract-container-1), such as `Vector{PairwiseInteraction}`. +For performance reasons it is best to [avoid containers with abstract type parameters](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-abstract-container-1), such as `Vector{Any}`. If you wish to calculate potential energies or log the energy throughout a simulation, you will need to define a method for the [`potential_energy`](@ref) function. This has the same arguments as [`force`](@ref), except the fifth argument is the energy units not the force units, and should return a single value corresponding to the potential energy: @@ -660,9 +660,9 @@ end ### Specific interactions -To define your own [`SpecificInteraction`](@ref), first define the `struct`: +To define your own specific interaction, first define the `struct`: ```julia -struct MySpecificInter <: SpecificInteraction +struct MySpecificInter # Properties, e.g. a bond distance corresponding to the energy minimum end ``` diff --git a/docs/src/examples.md b/docs/src/examples.md index e53cfe1e0..4bce04177 100644 --- a/docs/src/examples.md +++ b/docs/src/examples.md @@ -156,8 +156,8 @@ mutable struct Person ϵ::Float64 end -# Custom PairwiseInteraction -struct SIRInteraction <: PairwiseInteraction +# Custom pairwise interaction +struct SIRInteraction dist_infection::Float64 prob_infection::Float64 prob_recovery::Float64 @@ -605,7 +605,7 @@ values(sys.loggers.coords)[end] ## Making and breaking bonds There is an example of mutable atom properties in the main documentation, but what if you want to make and break bonds during the simulation? -In this case you can use a [`PairwiseInteraction`](@ref) to make, break and apply the bonds. +In this case you can use a pairwise interaction to make, break and apply the bonds. The partners of the atom can be stored in the atom type. We make a logger to record when the bonds are present, allowing us to visualize them with the `connection_frames` keyword argument to [`visualize`](@ref) (this can take a while to plot). ```julia @@ -621,7 +621,7 @@ struct BondableAtom partners::Set{Int} end -struct BondableInteraction <: PairwiseInteraction +struct BondableInteraction prob_formation::Float64 prob_break::Float64 dist_formation::Float64 diff --git a/src/energy.jl b/src/energy.jl index 3ee29d0d4..c0f10c6a7 100644 --- a/src/energy.jl +++ b/src/energy.jl @@ -53,17 +53,16 @@ end Calculate the potential energy of a system using the pairwise, specific and general interactions. - potential_energy(inter::PairwiseInteraction, vec_ij, atom_i, atom_j, energy_units, special, - coord_i, coord_j, boundary, velocity_i, velocity_j, step_n) - potential_energy(inter::SpecificInteraction, coord_i, boundary, atom_i, energy_units, - velocity_i, step_n) - potential_energy(inter::SpecificInteraction, coord_i, coord_j, boundary, atom_i, atom_j, - energy_units, velocity_i, velocity_j, step_n) - potential_energy(inter::SpecificInteraction, coord_i, coord_j, coord_k, boundary, atom_i, - atom_j, atom_k, energy_units, velocity_i, velocity_j, velocity_k, step_n) - potential_energy(inter::SpecificInteraction, coord_i, coord_j, coord_k, coord_l, boundary, - atom_i, atom_j, atom_k, atom_l, energy_units, velocity_i, velocity_j, - velocity_k, velocity_l, step_n) + potential_energy(inter, vec_ij, atom_i, atom_j, energy_units, special, coord_i, coord_j, + boundary, velocity_i, velocity_j, step_n) + potential_energy(inter, coord_i, boundary, atom_i, energy_units, velocity_i, step_n) + potential_energy(inter, coord_i, coord_j, boundary, atom_i, atom_j, energy_units, + velocity_i, velocity_j, step_n) + potential_energy(inter, coord_i, coord_j, coord_k, boundary, atom_i, atom_j, atom_k, + energy_units, velocity_i, velocity_j, velocity_k, step_n) + potential_energy(inter, coord_i, coord_j, coord_k, coord_l, boundary, atom_i, atom_j, + atom_k, atom_l, energy_units, velocity_i, velocity_j, velocity_k, + velocity_l, step_n) Calculate the potential energy due to a given interaction type. @@ -291,8 +290,8 @@ function potential_energy(sys::System{D, true, T}, neighbors, step_n::Integer=0; end # Allow GPU-specific potential energy functions to be defined if required -potential_energy_gpu(inter::PairwiseInteraction, dr, ai, aj, eu, sp, ci, cj, bnd, vi, vj, sn) = potential_energy(inter, dr, ai, aj, eu, sp, ci, cj, bnd, vi, vj, sn) -potential_energy_gpu(inter::SpecificInteraction, ci, bnd, ai, eu, vi, sn) = potential_energy(inter, ci, bnd, ai, eu, vi, sn) -potential_energy_gpu(inter::SpecificInteraction, ci, cj, bnd, ai, aj, eu, vi, vj, sn) = potential_energy(inter, ci, cj, bnd, ai, aj, eu, vi, vj, sn) -potential_energy_gpu(inter::SpecificInteraction, ci, cj, ck, bnd, ai, aj, ak, eu, vi, vj, vk, sn) = potential_energy(inter, ci, cj, ck, bnd, ai, aj, ak, eu, vi, vj, vk, sn) -potential_energy_gpu(inter::SpecificInteraction, ci, cj, ck, cl, bnd, ai, aj, ak, al, eu, vi, vj, vk, vl, sn) = potential_energy(inter, ci, cj, ck, cl, bnd, ai, aj, ak, al, eu, vi, vj, vk, vl, sn) +potential_energy_gpu(inter, dr, ai, aj, eu, sp, ci, cj, bnd, vi, vj, sn) = potential_energy(inter, dr, ai, aj, eu, sp, ci, cj, bnd, vi, vj, sn) +potential_energy_gpu(inter, ci, bnd, ai, eu, vi, sn) = potential_energy(inter, ci, bnd, ai, eu, vi, sn) +potential_energy_gpu(inter, ci, cj, bnd, ai, aj, eu, vi, vj, sn) = potential_energy(inter, ci, cj, bnd, ai, aj, eu, vi, vj, sn) +potential_energy_gpu(inter, ci, cj, ck, bnd, ai, aj, ak, eu, vi, vj, vk, sn) = potential_energy(inter, ci, cj, ck, bnd, ai, aj, ak, eu, vi, vj, vk, sn) +potential_energy_gpu(inter, ci, cj, ck, cl, bnd, ai, aj, ak, al, eu, vi, vj, vk, vl, sn) = potential_energy(inter, ci, cj, ck, cl, bnd, ai, aj, ak, al, eu, vi, vj, vk, vl, sn) diff --git a/src/force.jl b/src/force.jl index 6b0042cb8..760089a54 100644 --- a/src/force.jl +++ b/src/force.jl @@ -25,32 +25,30 @@ function accelerations(sys, neighbors, step_n::Integer=0; n_threads::Integer=Thr end """ - force(inter::PairwiseInteraction, vec_ij, atom_i, atom_j, force_units, special, - coord_i, coord_j, boundary, velocity_i, velocity_j, step_n) - force(inter::SpecificInteraction, coord_i, boundary, atom_i, force_units, - velocity_i, step_n) - force(inter::SpecificInteraction, coord_i, coord_j, boundary, atom_i, atom_j, - force_units, velocity_i, velocity_j, step_n) - force(inter::SpecificInteraction, coord_i, coord_j, coord_k, boundary, atom_i, - atom_j, atom_k, force_units, velocity_i, velocity_j, velocity_k, step_n) - force(inter::SpecificInteraction, coord_i, coord_j, coord_k, coord_l, boundary, - atom_i, atom_j, atom_k, atom_l, force_units, velocity_i, velocity_j, - velocity_k, velocity_l, step_n) + force(inter, vec_ij, atom_i, atom_j, force_units, special, coord_i, coord_j, + boundary, velocity_i, velocity_j, step_n) + force(inter, coord_i, boundary, atom_i, force_units, velocity_i, step_n) + force(inter, coord_i, coord_j, boundary, atom_i, atom_j, force_units, velocity_i, + velocity_j, step_n) + force(inter, coord_i, coord_j, coord_k, boundary, atom_i, atom_j, atom_k, + force_units, velocity_i, velocity_j, velocity_k, step_n) + force(inter, coord_i, coord_j, coord_k, coord_l, boundary, atom_i, atom_j, atom_k, + atom_l, force_units, velocity_i, velocity_j, velocity_k, velocity_l, step_n) Calculate the force between atoms due to a given interaction type. -For [`PairwiseInteraction`](@ref)s returns a single force vector and for -[`SpecificInteraction`](@ref)s returns a type such as [`SpecificForce2Atoms`](@ref). +For pairwise interactions returns a single force vector and for specific interactions +returns a type such as [`SpecificForce2Atoms`](@ref). Custom pairwise and specific interaction types should implement this function. """ function force end # Allow GPU-specific force functions to be defined if required -force_gpu(inter::PairwiseInteraction, dr, ai, aj, fu, sp, ci, cj, bnd, vi, vj, sn) = force(inter, dr, ai, aj, fu, sp, ci, cj, bnd, vi, vj, sn) -force_gpu(inter::SpecificInteraction, ci, bnd, ai, fu, vi, sn) = force(inter, ci, bnd, ai, fu, vi, sn) -force_gpu(inter::SpecificInteraction, ci, cj, bnd, ai, aj, fu, vi, vj, sn) = force(inter, ci, cj, bnd, ai, aj, fu, vi, vj, sn) -force_gpu(inter::SpecificInteraction, ci, cj, ck, bnd, ai, aj, ak, fu, vi, vj, vk, sn) = force(inter, ci, cj, ck, bnd, ai, aj, ak, fu, vi, vj, vk, sn) -force_gpu(inter::SpecificInteraction, ci, cj, ck, cl, bnd, ai, aj, ak, al, fu, vi, vj, vk, vl, sn) = force(inter, ci, cj, ck, cl, bnd, ai, aj, ak, al, fu, vi, vj, vk, vl, sn) +force_gpu(inter, dr, ai, aj, fu, sp, ci, cj, bnd, vi, vj, sn) = force(inter, dr, ai, aj, fu, sp, ci, cj, bnd, vi, vj, sn) +force_gpu(inter, ci, bnd, ai, fu, vi, sn) = force(inter, ci, bnd, ai, fu, vi, sn) +force_gpu(inter, ci, cj, bnd, ai, aj, fu, vi, vj, sn) = force(inter, ci, cj, bnd, ai, aj, fu, vi, vj, sn) +force_gpu(inter, ci, cj, ck, bnd, ai, aj, ak, fu, vi, vj, vk, sn) = force(inter, ci, cj, ck, bnd, ai, aj, ak, fu, vi, vj, vk, sn) +force_gpu(inter, ci, cj, ck, cl, bnd, ai, aj, ak, al, fu, vi, vj, vk, vl, sn) = force(inter, ci, cj, ck, cl, bnd, ai, aj, ak, al, fu, vi, vj, vk, vl, sn) """ SpecificForce1Atoms(f1) diff --git a/src/interactions/buckingham.jl b/src/interactions/buckingham.jl index 87eaa8831..cf6fb790d 100644 --- a/src/interactions/buckingham.jl +++ b/src/interactions/buckingham.jl @@ -45,7 +45,7 @@ C_{ij} &= (C_{ii} C_{jj})^{1/2} ``` so atoms that use this interaction should have fields `A`, `B` and `C` available. """ -@kwdef struct Buckingham{C, S, A, B, M, W} <: PairwiseInteraction +@kwdef struct Buckingham{C, S, A, B, M, W} cutoff::C = NoCutoff() use_neighbors::Bool = false shortcut::S = buckingham_zero_shortcut diff --git a/src/interactions/cosine_angle.jl b/src/interactions/cosine_angle.jl index b159ecd64..b2dba1c95 100644 --- a/src/interactions/cosine_angle.jl +++ b/src/interactions/cosine_angle.jl @@ -11,7 +11,7 @@ The potential energy is defined as V(\theta) = k(1 + \cos(\theta - \theta_0)) ``` """ -struct CosineAngle{K, D} <: SpecificInteraction +struct CosineAngle{K, D} k::K θ0::D end diff --git a/src/interactions/coulomb.jl b/src/interactions/coulomb.jl index b178897b0..b27a47fc4 100644 --- a/src/interactions/coulomb.jl +++ b/src/interactions/coulomb.jl @@ -15,7 +15,7 @@ The potential energy is defined as V(r_{ij}) = \frac{q_i q_j}{4 \pi \varepsilon_0 r_{ij}} ``` """ -@kwdef struct Coulomb{C, W, T} <: PairwiseInteraction +@kwdef struct Coulomb{C, W, T} cutoff::C = NoCutoff() use_neighbors::Bool = false weight_special::W = 1 @@ -98,7 +98,7 @@ V(r_{ij}) = \frac{q_i q_j}{4 \pi \varepsilon_0 (r_{ij}^6 + \alpha \sigma_{ij}^6 ``` If ``\alpha`` or ``\lambda`` are zero this gives the standard [`Coulomb`](@ref) potential. """ -@kwdef struct CoulombSoftCore{C, A, L, P, S, W, T, R} <: PairwiseInteraction +@kwdef struct CoulombSoftCore{C, A, L, P, S, W, T, R} cutoff::C = NoCutoff() α::A = 1 λ::L = 0 @@ -206,7 +206,7 @@ const crf_solvent_dielectric = 78.3 The Coulomb electrostatic interaction modified using the reaction field approximation between two atoms. """ -@kwdef struct CoulombReactionField{D, S, W, T} <: PairwiseInteraction +@kwdef struct CoulombReactionField{D, S, W, T} dist_cutoff::D solvent_dielectric::S = crf_solvent_dielectric use_neighbors::Bool = false diff --git a/src/interactions/fene_bond.jl b/src/interactions/fene_bond.jl index 8020cd976..32dafc996 100644 --- a/src/interactions/fene_bond.jl +++ b/src/interactions/fene_bond.jl @@ -19,7 +19,7 @@ V_{\text{WCA}}(r) = \end{cases} ``` """ -struct FENEBond{K, D, E} <: SpecificInteraction +struct FENEBond{K, D, E} k::K r0::D σ::D diff --git a/src/interactions/gravity.jl b/src/interactions/gravity.jl index 1b4927d41..0abed2ba1 100644 --- a/src/interactions/gravity.jl +++ b/src/interactions/gravity.jl @@ -10,7 +10,7 @@ The potential energy is defined as V(r_{ij}) = -\frac{G m_i m_j}{r_{ij}} ``` """ -@kwdef struct Gravity{T} <: PairwiseInteraction +@kwdef struct Gravity{T} G::T = Unitful.G use_neighbors::Bool = false end diff --git a/src/interactions/harmonic_angle.jl b/src/interactions/harmonic_angle.jl index 318c81e56..e28ed0954 100644 --- a/src/interactions/harmonic_angle.jl +++ b/src/interactions/harmonic_angle.jl @@ -12,7 +12,7 @@ The potential energy is defined as V(\theta) = \frac{1}{2} k (\theta - \theta_0)^2 ``` """ -struct HarmonicAngle{K, D} <: SpecificInteraction +struct HarmonicAngle{K, D} k::K θ0::D end diff --git a/src/interactions/harmonic_bond.jl b/src/interactions/harmonic_bond.jl index 19b5502be..291d6eef5 100644 --- a/src/interactions/harmonic_bond.jl +++ b/src/interactions/harmonic_bond.jl @@ -10,7 +10,7 @@ The potential energy is defined as V(r) = \frac{1}{2} k (r - r_0)^2 ``` """ -struct HarmonicBond{K, D} <: SpecificInteraction +struct HarmonicBond{K, D} k::K r0::D end diff --git a/src/interactions/harmonic_position_restraint.jl b/src/interactions/harmonic_position_restraint.jl index d23f493cb..b1464c680 100644 --- a/src/interactions/harmonic_position_restraint.jl +++ b/src/interactions/harmonic_position_restraint.jl @@ -10,7 +10,7 @@ The potential energy is defined as V(\boldsymbol{x}) = \frac{1}{2} k |\boldsymbol{x} - \boldsymbol{x}_0|^2 ``` """ -struct HarmonicPositionRestraint{K, C} <: SpecificInteraction +struct HarmonicPositionRestraint{K, C} k::K x0::C end diff --git a/src/interactions/lennard_jones.jl b/src/interactions/lennard_jones.jl index 37341cd19..7a5450ac7 100644 --- a/src/interactions/lennard_jones.jl +++ b/src/interactions/lennard_jones.jl @@ -38,7 +38,7 @@ and the force on each atom by \end{aligned} ``` """ -@kwdef struct LennardJones{C, H, S, E, W} <: PairwiseInteraction +@kwdef struct LennardJones{C, H, S, E, W} cutoff::C = NoCutoff() use_neighbors::Bool = false shortcut::H = lj_zero_shortcut @@ -153,7 +153,7 @@ r_{ij}^{\text{sc}} = \left(r_{ij}^6 + \alpha \sigma_{ij}^6 \lambda^p \right)^{1/ ``` If ``\alpha`` or ``\lambda`` are zero this gives the standard [`LennardJones`](@ref) potential. """ -@kwdef struct LennardJonesSoftCore{C, A, L, P, H, S, E, W, R} <: PairwiseInteraction +@kwdef struct LennardJonesSoftCore{C, A, L, P, H, S, E, W, R} cutoff::C = NoCutoff() α::A = 1 λ::L = 0 diff --git a/src/interactions/mie.jl b/src/interactions/mie.jl index 7e76c3551..e68d53c6c 100644 --- a/src/interactions/mie.jl +++ b/src/interactions/mie.jl @@ -15,7 +15,7 @@ where C = \frac{n}{n - m} \left( \frac{n}{m} \right) ^\frac{m}{n - m} ``` """ -struct Mie{T, C, H, S, E} <: PairwiseInteraction +struct Mie{T, C, H, S, E} m::T n::T cutoff::C diff --git a/src/interactions/morse_bond.jl b/src/interactions/morse_bond.jl index 9d6f32422..ef5159547 100644 --- a/src/interactions/morse_bond.jl +++ b/src/interactions/morse_bond.jl @@ -10,7 +10,7 @@ The potential energy is defined as V(r) = D(1 - e^{-a(r - r_0)})^2 ``` """ -struct MorseBond{T, A, R} <: SpecificInteraction +struct MorseBond{T, A, R} D::T a::A r0::R diff --git a/src/interactions/periodic_torsion.jl b/src/interactions/periodic_torsion.jl index da22d43b8..ce6db4b26 100644 --- a/src/interactions/periodic_torsion.jl +++ b/src/interactions/periodic_torsion.jl @@ -11,7 +11,7 @@ The potential energy is defined as V(\phi) = \sum_{n=1}^N k_n (1 + \cos(n \phi - \phi_{s,n})) ``` """ -struct PeriodicTorsion{N, T, E} <: SpecificInteraction +struct PeriodicTorsion{N, T, E} periodicities::NTuple{N, Int} phases::NTuple{N, T} ks::NTuple{N, E} diff --git a/src/interactions/rb_torsion.jl b/src/interactions/rb_torsion.jl index 42147451b..d6742670f 100644 --- a/src/interactions/rb_torsion.jl +++ b/src/interactions/rb_torsion.jl @@ -5,7 +5,7 @@ export RBTorsion A Ryckaert-Bellemans torsion angle between four atoms. """ -struct RBTorsion{T} <: SpecificInteraction +struct RBTorsion{T} f1::T f2::T f3::T diff --git a/src/interactions/soft_sphere.jl b/src/interactions/soft_sphere.jl index e64c5ec46..89921e9e0 100644 --- a/src/interactions/soft_sphere.jl +++ b/src/interactions/soft_sphere.jl @@ -10,7 +10,7 @@ The potential energy is defined as V(r_{ij}) = 4\varepsilon_{ij} \left(\frac{\sigma_{ij}}{r_{ij}}\right)^{12} ``` """ -@kwdef struct SoftSphere{C, H, S, E} <: PairwiseInteraction +@kwdef struct SoftSphere{C, H, S, E} cutoff::C = NoCutoff() use_neighbors::Bool = false shortcut::H = lj_zero_shortcut diff --git a/src/neighbors.jl b/src/neighbors.jl index bde6ae9fb..ba23a5b6c 100644 --- a/src/neighbors.jl +++ b/src/neighbors.jl @@ -1,12 +1,24 @@ # Neighbor finders export + use_neighbors, NoNeighborFinder, find_neighbors, DistanceNeighborFinder, TreeNeighborFinder, CellListMapNeighborFinder +""" + use_neighbors(inter) + +Whether a pairwise interaction uses the neighbor list, default `false`. + +Custom pairwise interactions can define a method for this function. +For built-in interactions such as [`LennardJones`](@ref) this function accesses +the `use_neighbors` field of the struct. +""" +use_neighbors(inter) = false + """ NoNeighborFinder() diff --git a/src/types.jl b/src/types.jl index fa46be671..260f04ac4 100644 --- a/src/types.jl +++ b/src/types.jl @@ -1,9 +1,6 @@ # Types export - PairwiseInteraction, - use_neighbors, - SpecificInteraction, InteractionList1Atoms, InteractionList2Atoms, InteractionList3Atoms, @@ -25,31 +22,6 @@ export const DefaultFloat = Float64 -""" -A pairwise interaction that will apply to all or most atom pairs. - -Custom pairwise interactions should sub-type this abstract type. -""" -abstract type PairwiseInteraction end - -""" - use_neighbors(inter) - -Whether a pairwise interaction uses the neighbor list, default `false`. - -Custom pairwise interactions can define a method for this function. -For built-in interactions such as [`LennardJones`](@ref) this function accesses -the `use_neighbors` field of the struct. -""" -use_neighbors(::PairwiseInteraction) = false - -""" -A specific interaction between sets of specific atoms, e.g. a bond angle. - -Custom specific interactions should sub-type this abstract type. -""" -abstract type SpecificInteraction end - """ InteractionList1Atoms(is, inters) InteractionList1Atoms(is, inters, types) diff --git a/test/agent.jl b/test/agent.jl index d33c54af4..740aeb74e 100644 --- a/test/agent.jl +++ b/test/agent.jl @@ -10,8 +10,8 @@ ϵ::Float64 end - # Custom PairwiseInteraction - struct SIRInteraction <: PairwiseInteraction + # Custom pairwise interaction + struct SIRInteraction dist_infection::Float64 prob_infection::Float64 prob_recovery::Float64 From befdd318d8b83d52405e43f5359a9673aa275139 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Wed, 24 Jul 2024 15:06:26 +0100 Subject: [PATCH 39/74] rng arg for coordinate setup --- src/setup.jl | 16 +++++++++------- src/spatial.jl | 15 ++++++++++----- 2 files changed, 19 insertions(+), 12 deletions(-) diff --git a/src/setup.jl b/src/setup.jl index a844aeffe..614ec4b11 100644 --- a/src/setup.jl +++ b/src/setup.jl @@ -14,7 +14,7 @@ export add_position_restraints """ - place_atoms(n_atoms, boundary; min_dist=nothing, max_attempts=100) + place_atoms(n_atoms, boundary; min_dist=nothing, max_attempts=100, rng=Random.GLOBAL_RNG) Generate random coordinates. @@ -27,7 +27,8 @@ Can not be used if one or more dimensions has infinite boundaries. function place_atoms(n_atoms::Integer, boundary; min_dist=zero(length_type(boundary)), - max_attempts::Integer=100) + max_attempts::Integer=100, + rng=Random.GLOBAL_RNG) if has_infinite_boundary(boundary) throw(ArgumentError("one or more dimension has infinite boundaries, boundary is $boundary")) end @@ -41,7 +42,7 @@ function place_atoms(n_atoms::Integer, sizehint!(coords, n_atoms) failed_attempts = 0 while length(coords) < n_atoms - new_coord = random_coord(boundary) + new_coord = random_coord(boundary; rng=rng) okay = true if min_dist > zero(min_dist) for coord in coords @@ -64,7 +65,7 @@ end """ place_diatomics(n_molecules, boundary, bond_length; min_dist=nothing, - max_attempts=100, aligned=false) + max_attempts=100, aligned=false, rng=Random.GLOBAL_RNG) Generate random diatomic molecule coordinates. @@ -82,7 +83,8 @@ function place_diatomics(n_molecules::Integer, bond_length; min_dist=zero(length_type(boundary)), max_attempts::Integer=100, - aligned::Bool=false) + aligned::Bool=false, + rng=Random.GLOBAL_RNG) if has_infinite_boundary(boundary) throw(ArgumentError("one or more dimension has infinite boundaries, boundary is $boundary")) end @@ -96,11 +98,11 @@ function place_diatomics(n_molecules::Integer, sizehint!(coords, 2 * n_molecules) failed_attempts = 0 while length(coords) < (n_molecules * 2) - new_coord_a = random_coord(boundary) + new_coord_a = random_coord(boundary; rng=rng) if aligned shift = SVector{dims}([bond_length, [zero(bond_length) for d in 1:(dims - 1)]...]) else - shift = bond_length * normalize(randn(SVector{dims, typeof(ustrip(bond_length))})) + shift = bond_length * normalize(randn(rng, SVector{dims, typeof(ustrip(bond_length))})) end new_coord_b = copy(new_coord_a) + shift okay = true diff --git a/src/spatial.jl b/src/spatial.jl index 023e5f55b..d3a019e2d 100644 --- a/src/spatial.jl +++ b/src/spatial.jl @@ -323,15 +323,20 @@ function bounding_box_lines(boundary::TriclinicBoundary, dist_unit) end """ - random_coord(boundary) + random_coord(boundary; rng=Random.GLOBAL_RNG) Generate a random coordinate uniformly distributed within a bounding box. """ -random_coord(boundary::CubicBoundary ) = rand(SVector{3, float_type(boundary)}) .* boundary -random_coord(boundary::RectangularBoundary) = rand(SVector{2, float_type(boundary)}) .* boundary +function random_coord(boundary::CubicBoundary; rng=Random.GLOBAL_RNG) + return rand(rng, SVector{3, float_type(boundary)}) .* boundary +end + +function random_coord(boundary::RectangularBoundary; rng=Random.GLOBAL_RNG) + return rand(rng, SVector{2, float_type(boundary)}) .* boundary +end -function random_coord(boundary::TriclinicBoundary{T}) where T - return sum(rand(SVector{3, T}) .* boundary.basis_vectors) +function random_coord(boundary::TriclinicBoundary{T}; rng=Random.GLOBAL_RNG) where T + return sum(rand(rng, SVector{3, T}) .* boundary.basis_vectors) end """ From 7f8efbabe66906ee2e387669f8dc3594792fe777 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Wed, 24 Jul 2024 15:17:27 +0100 Subject: [PATCH 40/74] make DMS test pass with Enzyme --- ext/MollyEnzymeExt.jl | 2 + test/gradients.jl | 198 +++++++++++++++++++++++++----------------- test/runtests.jl | 4 +- 3 files changed, 122 insertions(+), 82 deletions(-) diff --git a/ext/MollyEnzymeExt.jl b/ext/MollyEnzymeExt.jl index 4ddbe7c28..90e015390 100644 --- a/ext/MollyEnzymeExt.jl +++ b/ext/MollyEnzymeExt.jl @@ -6,7 +6,9 @@ module MollyEnzymeExt using Molly using Enzyme +EnzymeRules.inactive(::typeof(Molly.check_units), args...) = nothing EnzymeRules.inactive(::typeof(Molly.n_infinite_dims), args...) = nothing +EnzymeRules.inactive(::typeof(random_velocity), args...) = nothing EnzymeRules.inactive(::typeof(random_velocities), args...) = nothing EnzymeRules.inactive(::typeof(random_velocities!), args...) = nothing EnzymeRules.inactive(::typeof(Molly.cuda_threads_blocks_pairwise), args...) = nothing diff --git a/test/gradients.jl b/test/gradients.jl index caa554494..f897012f3 100644 --- a/test/gradients.jl +++ b/test/gradients.jl @@ -35,20 +35,80 @@ end @testset "Differentiable simulation" begin - abs2_vec(x) = abs2.(x) - - # Function is strange in order to work with gradients on the GPU - function mean_min_separation(coords, boundary) - diffs = displacements(coords, boundary) - disps = Array(sum.(abs2_vec.(diffs))) - disps_diag = disps .+ Diagonal(100 * ones(typeof(boundary[1]), length(coords))) - return mean(sqrt.(minimum(disps_diag; dims=1))) + runs = [ # gpu par fwd f32 obc2 gbn2 tol_σ tol_r0 + ("CPU" , [false, false, false, false, false, false], 1e-4, 1e-4), + ("CPU forward" , [false, false, true , false, false, false], 0.5 , 0.1 ), + ("CPU f32" , [false, false, false, true , false, false], 0.01, 5e-4), + ("CPU obc2" , [false, false, false, false, true , false], 1e-4, 1e-4), + ("CPU gbn2" , [false, false, false, false, false, true ], 1e-4, 1e-4), + ("CPU gbn2 forward", [false, false, true , false, false, true ], 0.5 , 0.1 ), + ] + if run_parallel_tests # gpu par fwd f32 obc2 gbn2 tol_σ tol_r0 + push!(runs, ("CPU parallel" , [false, true , false, false, false, false], 1e-4, 1e-4)) + push!(runs, ("CPU parallel forward", [false, true , true , false, false, false], 0.5 , 0.1 )) + push!(runs, ("CPU parallel f32" , [false, true , false, true , false, false], 0.01, 5e-4)) + end + if run_gpu_tests # gpu par fwd f32 obc2 gbn2 tol_σ tol_r0 + push!(runs, ("GPU" , [true , false, false, false, false, false], 0.25, 20.0)) + push!(runs, ("GPU forward" , [true , false, true , false, false, false], 0.25, 20.0)) + push!(runs, ("GPU f32" , [true , false, false, true , false, false], 0.5 , 50.0)) + push!(runs, ("GPU obc2" , [true , false, false, false, true , false], 0.25, 20.0)) + push!(runs, ("GPU gbn2" , [true , false, false, false, false, true ], 0.25, 20.0)) + end + + function mean_min_separation(coords, boundary, ::Val{T}) where T + min_seps = T[] + for i in eachindex(coords) + min_sq_sep = T(100.0) + for j in eachindex(coords) + if i != j + sq_dist = sum(abs2, vector(coords[i], coords[j], boundary)) + min_sq_sep = min(sq_dist, min_sq_sep) + end + end + push!(min_seps, sqrt(min_sq_sep)) + end + return mean(min_seps) end - function test_simulation_grad(gpu::Bool, parallel::Bool, forward::Bool, f32::Bool, pis::Bool, - sis::Bool, obc2::Bool, gbn2::Bool) + function loss(σ, r0, coords, velocities, boundary, pairwise_inters, general_inters, + neighbor_finder, simulator, n_steps, n_threads, n_atoms, atom_mass, bond_dists, + bond_is, bond_js, angles, torsions, ::Val{T}, ::Val{AT}) where {T, AT} + atoms = [Atom(i, 1, atom_mass, (i % 2 == 0 ? T(-0.02) : T(0.02)), σ, T(0.2)) for i in 1:n_atoms] + bonds_inner = HarmonicBond{T, T}[] + for i in 1:(n_atoms ÷ 2) + push!(bonds_inner, HarmonicBond(T(100.0), bond_dists[i] * r0)) + end + bonds = InteractionList2Atoms( + bond_is, + bond_js, + AT(bonds_inner), + ) + + sys = System( + atoms=atoms, + coords=AT(coords), + boundary=boundary, + velocities=AT(velocities), + pairwise_inters=pairwise_inters, + specific_inter_lists=(bonds, angles, torsions), + general_inters=general_inters, + neighbor_finder=neighbor_finder, + force_units=NoUnits, + energy_units=NoUnits, + ) + + simulate!(sys, simulator, n_steps; n_threads=n_threads) + + return mean_min_separation(sys.coords, boundary, Val(T)) + end + + for (name, args, tol_σ, tol_r0) in runs + gpu, parallel, forward, f32, obc2, gbn2 = args T = f32 ? Float32 : Float64 AT = gpu ? CuArray : Array + σ = T(0.4) + r0 = T(1.0) n_atoms = 50 n_steps = 100 atom_mass = T(10.0) @@ -58,8 +118,9 @@ end dt=T(0.001), coupling=RescaleThermostat(temp), ) - coords = place_atoms(n_atoms, boundary; min_dist=T(0.6), max_attempts=500) - velocities = [random_velocity(atom_mass, temp) for i in 1:n_atoms] + rng = Xoshiro(1000) # Same system every time, not required but increases stability + coords = place_atoms(n_atoms, boundary; min_dist=T(0.6), max_attempts=500, rng=rng) + velocities = [random_velocity(atom_mass, temp; rng=rng) for i in 1:n_atoms] nb_cutoff = T(1.2) lj = LennardJones(cutoff=DistanceCutoff(nb_cutoff), use_neighbors=true) crf = CoulombReactionField( @@ -68,10 +129,11 @@ end use_neighbors=true, coulomb_const=T(ustrip(Molly.coulomb_const)), ) - pairwise_inters = pis ? (lj, crf) : () + pairwise_inters = (lj, crf) bond_is = AT(Int32.(collect(1:(n_atoms ÷ 2)))) bond_js = AT(Int32.(collect((1 + n_atoms ÷ 2):n_atoms))) - bond_dists = [norm(vector(Array(coords)[i], Array(coords)[i + n_atoms ÷ 2], boundary)) for i in 1:(n_atoms ÷ 2)] + bond_dists = [norm(vector(Array(coords)[i], Array(coords)[i + n_atoms ÷ 2], boundary)) + for i in 1:(n_atoms ÷ 2)] angles_inner = [HarmonicAngle(k=T(10.0), θ0=T(2.0)) for i in 1:15] angles = InteractionList3Atoms( AT(Int32.(collect( 1:15))), @@ -118,79 +180,53 @@ end n_steps=10, dist_cutoff=T(1.5), ) - - function loss(σ, r0) - atoms = [Atom(i, 1, atom_mass, (i % 2 == 0 ? T(-0.02) : T(0.02)), σ, T(0.2)) for i in 1:n_atoms] - bonds_inner = [HarmonicBond(T(100.0), bond_dists[i] * r0) for i in 1:(n_atoms ÷ 2)] - bonds = InteractionList2Atoms( - bond_is, - bond_js, - AT(bonds_inner), - ) - cs = deepcopy(coords) - vs = deepcopy(velocities) - - sys = System( - atoms=AT(atoms), - coords=AT(cs), - boundary=boundary, - velocities=AT(vs), - pairwise_inters=pairwise_inters, - specific_inter_lists=sis ? (bonds, angles, torsions) : (), - general_inters=general_inters, - neighbor_finder=neighbor_finder, - force_units=NoUnits, - energy_units=NoUnits, - ) - - simulate!(sys, simulator, n_steps; n_threads=(parallel ? Threads.nthreads() : 1)) - - return mean_min_separation(sys.coords, boundary) - end - - return loss - end - - runs = [ # gpu par fwd f32 pis sis obc2 gbn2 tol_σ tol_r0 - ("CPU" , [false, false, false, false, true , true , false, false], 0.1 , 1.0 ), - ("CPU forward" , [false, false, true , false, true , true , false, false], 0.02, 0.1 ), - ("CPU f32" , [false, false, false, true , true , true , false, false], 0.2 , 10.0), - ("CPU nospecific" , [false, false, false, false, true , false, false, false], 0.1 , 0.0 ), - ("CPU nopairwise" , [false, false, false, false, false, true , false, false], 0.0 , 1.0 ), - ("CPU obc2" , [false, false, false, false, true , true , true , false], 0.1 , 20.0), - ("CPU gbn2" , [false, false, false, false, true , true , false, true ], 0.1 , 20.0), - ("CPU gbn2 forward", [false, false, true , false, true , true , false, true ], 0.05, 0.1 ), - ] - if run_parallel_tests # gpu par fwd f32 pis sis obc2 gbn2 tol_σ tol_r0 - push!(runs, ("CPU parallel" , [false, true , false, false, true , true , false, false], 0.1 , 1.0 )) - push!(runs, ("CPU parallel forward", [false, true , true , false, true , true , false, false], 0.01, 0.05)) - push!(runs, ("CPU parallel f32" , [false, true , false, true , true , true , false, false], 0.2 , 10.0)) - end - if run_gpu_tests # gpu par fwd f32 pis sis obc2 gbn2 tol_σ tol_r0 - push!(runs, ("GPU" , [true , false, false, false, true , true , false, false], 0.25, 20.0)) - push!(runs, ("GPU f32" , [true , false, false, true , true , true , false, false], 0.5 , 50.0)) - push!(runs, ("GPU nospecific" , [true , false, false, false, true , false, false, false], 0.25, 0.0 )) - push!(runs, ("GPU nopairwise" , [true , false, false, false, false, true , false, false], 0.0 , 10.0)) - push!(runs, ("GPU obc2" , [true , false, false, false, true , true , true , false], 0.25, 20.0)) - push!(runs, ("GPU gbn2" , [true , false, false, false, true , true , false, true ], 0.25, 20.0)) - end - - for (name, args, tol_σ, tol_r0) in runs - forward, f32 = args[3], args[4] - σ = f32 ? 0.4f0 : 0.4 - r0 = f32 ? 1.0f0 : 1.0 - f = test_simulation_grad(args...) + n_threads = parallel ? Threads.nthreads() : 1 + + const_args = [ + Const(boundary), Const(pairwise_inters), + Const(general_inters), Const(neighbor_finder), Const(simulator), + Const(n_steps), Const(n_threads), Const(n_atoms), Const(atom_mass), + Const(bond_dists), Const(bond_is), Const(bond_js), Const(angles), + Const(torsions), Const(Val(T)), Const(Val(AT)), + ] if forward grad_enzyme = ( - autodiff(Forward, f, Duplicated, Duplicated(σ, f32 ? 1.0f0 : 1.0), Const(r0))[2], - autodiff(Forward, f, Duplicated, Const(σ), Duplicated(r0, f32 ? 1.0f0 : 1.0))[2], + autodiff( + Forward, loss, Duplicated, Duplicated(σ, one(T)), Const(r0), + Duplicated(copy(coords), zero(coords)), + Duplicated(copy(velocities), zero(velocities)), const_args..., + )[2], + autodiff( + Forward, loss, Duplicated, Const(σ), Duplicated(r0, one(T)), + Duplicated(copy(coords), zero(coords)), + Duplicated(copy(velocities), zero(velocities)), const_args..., + )[2], ) else - grad_enzyme = autodiff(Reverse, f, Active, Active(σ), Active(r0))[1] + grad_enzyme = autodiff( + Reverse, loss, Active, Active(σ), Active(r0), + Duplicated(copy(coords), zero(coords)), + Duplicated(copy(velocities), zero(velocities)), const_args..., + )[1][1:2] end + grad_fd = ( - central_fdm(6, 1)(σ -> ForwardDiff.value(f(σ, r0)), σ ), - central_fdm(6, 1)(r0 -> ForwardDiff.value(f(σ, r0)), r0), + central_fdm(6, 1)( + σ -> loss( + σ, r0, copy(coords), copy(velocities), boundary, pairwise_inters, general_inters, + neighbor_finder, simulator, n_steps, n_threads, n_atoms, atom_mass, bond_dists, + bond_is, bond_js, angles, torsions, Val(T), Val(AT), + ), + σ, + ), + central_fdm(6, 1)( + r0 -> loss( + σ, r0, copy(coords), copy(velocities), boundary, pairwise_inters, general_inters, + neighbor_finder, simulator, n_steps, n_threads, n_atoms, atom_mass, bond_dists, + bond_is, bond_js, angles, torsions, Val(T), Val(AT), + ), + r0, + ), ) for (prefix, genz, gfd, tol) in zip(("σ", "r0"), grad_enzyme, grad_fd, (tol_σ, tol_r0)) if abs(gfd) < 1e-13 diff --git a/test/runtests.jl b/test/runtests.jl index 3383a7beb..1be787b08 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -54,6 +54,8 @@ const gpu_list = run_gpu_tests ? (false, true) : (false,) if run_gpu_tests device!(parse(Int, DEVICE)) @info "The GPU tests will be run on device $DEVICE" +elseif get(ENV, "GPUTESTS", "1") == "0" + @warn "The GPU tests will not be run as GPUTESTS is set to 0" else @warn "The GPU tests will not be run as a CUDA-enabled device is not available" end @@ -65,7 +67,7 @@ const openmm_dir = joinpath(data_dir, "openmm_6mrr") const temp_fp_pdb = tempname(cleanup=true) * ".pdb" const temp_fp_viz = tempname(cleanup=true) * ".mp4" -# Required for parallel gradient tests +# Required for gradient tests Enzyme.API.runtimeActivity!(true) if GROUP in ("All", "NotGradients") From 5aad84109ad20b3e97bcaed1b896ec71d218fa44 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Wed, 24 Jul 2024 17:38:27 +0100 Subject: [PATCH 41/74] broadcast gradient injection --- src/gradients.jl | 3 +-- test/gradients.jl | 37 ++++++++++++++++++------------------- 2 files changed, 19 insertions(+), 21 deletions(-) diff --git a/src/gradients.jl b/src/gradients.jl index afad5a597..90829ac9c 100644 --- a/src/gradients.jl +++ b/src/gradients.jl @@ -95,8 +95,7 @@ function inject_gradients(sys::System{D, G}, params_dic) where {D, G} atoms_grad = inject_atom.(sys.atoms, sys.atoms_data, (params_dic,)) end if length(sys.pairwise_inters) > 0 - # Broadcasting fails due to https://github.com/JuliaDiff/ChainRules.jl/issues/662 - pis_grad = Tuple(inject_interaction(inter, params_dic) for inter in sys.pairwise_inters) + pis_grad = inject_interaction.(sys.pairwise_inters, (params_dic,)) else pis_grad = sys.pairwise_inters end diff --git a/test/gradients.jl b/test/gradients.jl index f897012f3..e6ed0ecf8 100644 --- a/test/gradients.jl +++ b/test/gradients.jl @@ -35,25 +35,25 @@ end @testset "Differentiable simulation" begin - runs = [ # gpu par fwd f32 obc2 gbn2 tol_σ tol_r0 - ("CPU" , [false, false, false, false, false, false], 1e-4, 1e-4), - ("CPU forward" , [false, false, true , false, false, false], 0.5 , 0.1 ), - ("CPU f32" , [false, false, false, true , false, false], 0.01, 5e-4), - ("CPU obc2" , [false, false, false, false, true , false], 1e-4, 1e-4), - ("CPU gbn2" , [false, false, false, false, false, true ], 1e-4, 1e-4), - ("CPU gbn2 forward", [false, false, true , false, false, true ], 0.5 , 0.1 ), + runs = [ # gpu par fwd f32 obc2 gbn2 tol_σ tol_r0 + ("CPU" , false, false, false, false, false, false, 1e-4, 1e-4), + ("CPU forward" , false, false, true , false, false, false, 0.5 , 0.1 ), + ("CPU f32" , false, false, false, true , false, false, 0.01, 5e-4), + ("CPU obc2" , false, false, false, false, true , false, 1e-4, 1e-4), + ("CPU gbn2" , false, false, false, false, false, true , 1e-4, 1e-4), + ("CPU gbn2 forward", false, false, true , false, false, true , 0.5 , 0.1 ), ] - if run_parallel_tests # gpu par fwd f32 obc2 gbn2 tol_σ tol_r0 - push!(runs, ("CPU parallel" , [false, true , false, false, false, false], 1e-4, 1e-4)) - push!(runs, ("CPU parallel forward", [false, true , true , false, false, false], 0.5 , 0.1 )) - push!(runs, ("CPU parallel f32" , [false, true , false, true , false, false], 0.01, 5e-4)) + if run_parallel_tests # gpu par fwd f32 obc2 gbn2 tol_σ tol_r0 + push!(runs, ("CPU parallel" , false, true , false, false, false, false, 1e-4, 1e-4)) + push!(runs, ("CPU parallel forward", false, true , true , false, false, false, 0.5 , 0.1 )) + push!(runs, ("CPU parallel f32" , false, true , false, true , false, false, 0.01, 5e-4)) end - if run_gpu_tests # gpu par fwd f32 obc2 gbn2 tol_σ tol_r0 - push!(runs, ("GPU" , [true , false, false, false, false, false], 0.25, 20.0)) - push!(runs, ("GPU forward" , [true , false, true , false, false, false], 0.25, 20.0)) - push!(runs, ("GPU f32" , [true , false, false, true , false, false], 0.5 , 50.0)) - push!(runs, ("GPU obc2" , [true , false, false, false, true , false], 0.25, 20.0)) - push!(runs, ("GPU gbn2" , [true , false, false, false, false, true ], 0.25, 20.0)) + if run_gpu_tests # gpu par fwd f32 obc2 gbn2 tol_σ tol_r0 + push!(runs, ("GPU" , true , false, false, false, false, false, 0.25, 20.0)) + push!(runs, ("GPU forward" , true , false, true , false, false, false, 0.25, 20.0)) + push!(runs, ("GPU f32" , true , false, false, true , false, false, 0.5 , 50.0)) + push!(runs, ("GPU obc2" , true , false, false, false, true , false, 0.25, 20.0)) + push!(runs, ("GPU gbn2" , true , false, false, false, false, true , 0.25, 20.0)) end function mean_min_separation(coords, boundary, ::Val{T}) where T @@ -103,8 +103,7 @@ end return mean_min_separation(sys.coords, boundary, Val(T)) end - for (name, args, tol_σ, tol_r0) in runs - gpu, parallel, forward, f32, obc2, gbn2 = args + for (name, gpu, parallel, forward, f32, obc2, gbn2, tol_σ, tol_r0) in runs T = f32 ? Float32 : Float64 AT = gpu ? CuArray : Array σ = T(0.4) From 8f7c2ec38a08ccb56a83d638d715379c40c372ac Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Wed, 24 Jul 2024 19:39:29 +0100 Subject: [PATCH 42/74] DMS protein gradient tests --- test/gradients.jl | 158 ++++++++++++++++++++-------------------------- 1 file changed, 70 insertions(+), 88 deletions(-) diff --git a/test/gradients.jl b/test/gradients.jl index e6ed0ecf8..9bcf25ff6 100644 --- a/test/gradients.jl +++ b/test/gradients.jl @@ -257,89 +257,65 @@ end ) end - function test_energy_grad(gpu::Bool, parallel::Bool) - sys_ref = create_sys(gpu) + EnzymeRules.inactive(::typeof(create_sys), args...) = nothing - function loss(params_dic) - n_threads = parallel ? Threads.nthreads() : 1 - atoms, pairwise_inters, specific_inter_lists, general_inters = inject_gradients( - sys_ref, params_dic) - - sys = System( - atoms=atoms, - coords=sys_ref.coords, - boundary=sys_ref.boundary, - pairwise_inters=pairwise_inters, - specific_inter_lists=specific_inter_lists, - general_inters=general_inters, - neighbor_finder=sys_ref.neighbor_finder, - force_units=NoUnits, - energy_units=NoUnits, - ) - - return potential_energy(sys; n_threads=n_threads) - end - - return loss + function test_energy_grad(params_dic, sys_ref, coords, neighbor_finder, n_threads) + atoms, pis, sis, gis = inject_gradients(sys_ref, params_dic) + + sys = System( + atoms=atoms, + coords=coords, + boundary=sys_ref.boundary, + pairwise_inters=pis, + specific_inter_lists=sis, + general_inters=gis, + neighbor_finder=neighbor_finder, + force_units=NoUnits, + energy_units=NoUnits, + ) + + return potential_energy(sys; n_threads=n_threads) end + + function test_force_grad(params_dic, sys_ref, coords, neighbor_finder, n_threads) + atoms, pis, sis, gis = inject_gradients(sys_ref, params_dic) - sum_abs(x) = sum(abs, x) - - function test_force_grad(gpu::Bool, parallel::Bool) - sys_ref = create_sys(gpu) - - function loss(params_dic) - n_threads = parallel ? Threads.nthreads() : 1 - atoms, pairwise_inters, specific_inter_lists, general_inters = inject_gradients( - sys_ref, params_dic) - - sys = System( - atoms=atoms, - coords=sys_ref.coords, - boundary=sys_ref.boundary, - pairwise_inters=pairwise_inters, - specific_inter_lists=specific_inter_lists, - general_inters=general_inters, - neighbor_finder=sys_ref.neighbor_finder, - force_units=NoUnits, - energy_units=NoUnits, - ) - - fs = forces(sys; n_threads=n_threads) - return sum(sum_abs.(fs)) - end + sys = System( + atoms=atoms, + coords=coords, + boundary=sys_ref.boundary, + pairwise_inters=pis, + specific_inter_lists=sis, + general_inters=gis, + neighbor_finder=neighbor_finder, + force_units=NoUnits, + energy_units=NoUnits, + ) - return loss + fs = forces(sys; n_threads=n_threads) + return sum(sum.(abs, fs)) end - function test_sim_grad(gpu::Bool, parallel::Bool) - sys_ref = create_sys(gpu) - - function loss(params_dic) - n_threads = parallel ? Threads.nthreads() : 1 - atoms, pairwise_inters, specific_inter_lists, general_inters = inject_gradients( - sys_ref, params_dic) - - sys = System( - atoms=atoms, - coords=sys_ref.coords, - boundary=sys_ref.boundary, - pairwise_inters=pairwise_inters, - specific_inter_lists=specific_inter_lists, - general_inters=general_inters, - neighbor_finder=sys_ref.neighbor_finder, - force_units=NoUnits, - energy_units=NoUnits, - ) - - simulator = Langevin(dt=0.001, temperature=300.0, friction=1.0) - n_steps = 5 - rand_seed = 1000 - simulate!(sys, simulator, n_steps; n_threads=n_threads, rng=Xoshiro(rand_seed)) - return sum(sum_abs.(sys.coords)) - end + function test_sim_grad(params_dic, sys_ref, coords, neighbor_finder, n_threads) + atoms, pis, sis, gis = inject_gradients(sys_ref, params_dic) + + sys = System( + atoms=atoms, + coords=coords, + boundary=sys_ref.boundary, + pairwise_inters=pis, + specific_inter_lists=sis, + general_inters=gis, + neighbor_finder=neighbor_finder, + force_units=NoUnits, + energy_units=NoUnits, + ) - return loss + simulator = Langevin(dt=0.001, temperature=300.0, friction=1.0) + n_steps = 5 + rng = Xoshiro(1000) + simulate!(sys, simulator, n_steps; n_threads=n_threads, rng=rng) + return sum(sum.(abs, sys.coords)) end params_dic = Dict( @@ -424,40 +400,46 @@ end "inter_PT_N/CT/C/N_k_9" => 0.238488, ) - platform_runs = [("CPU", [false, false])] + platform_runs = [("CPU", false, false)] if run_parallel_tests - push!(platform_runs, ("CPU parallel", [false, true])) + push!(platform_runs, ("CPU parallel", false, true)) end if run_gpu_tests - push!(platform_runs, ("GPU", [true, false])) + push!(platform_runs, ("GPU", true, false)) end test_runs = [ ("Energy", test_energy_grad, 1e-8), ("Force" , test_force_grad , 1e-8), ] if !running_CI - push!(test_runs, ("Sim", test_sim_grad, 0.015)) + #push!(test_runs, ("Sim", test_sim_grad, 0.015)) end params_to_test = ( - "inter_LJ_weight_14", + #"inter_LJ_weight_14", "atom_N_ϵ", "inter_PT_C/N/CT/C_k_1", "inter_GB_screen_O", - "inter_GB_neck_scale", + #"inter_GB_neck_scale", ) for (test_name, test_fn, test_tol) in test_runs - for (platform, args) in platform_runs - f = test_fn(args...) + for (platform, gpu, parallel) in platform_runs + sys_ref = create_sys(gpu) + n_threads = parallel ? Threads.nthreads() : 1 grads_enzyme = Dict(k => 0.0 for k in keys(params_dic)) - autodiff(Reverse, f, Active, Duplicated(params_dic, grads_enzyme)) - @test count(!iszero, values(grads_enzyme)) == 67 + autodiff( + Reverse, test_fn, Active, Duplicated(params_dic, grads_enzyme), + Const(sys_ref), Duplicated(sys_ref.coords, zero(sys_ref.coords)), + Duplicated(sys_ref.neighbor_finder, sys_ref.neighbor_finder), + Const(n_threads), + ) + #@test count(!iszero, values(grads_enzyme)) == 67 for param in params_to_test genz = grads_enzyme[param] gfd = central_fdm(6, 1)(params_dic[param]) do val - dic = deepcopy(params_dic) + dic = copy(params_dic) dic[param] = val - f(dic) + test_fn(dic, sys_ref, sys_ref.coords, sys_ref.neighbor_finder, n_threads) end frac_diff = abs(genz - gfd) / abs(gfd) @info "$(rpad(test_name, 6)) - $(rpad(platform, 12)) - $(rpad(param, 21)) - FD $gfd, Enzyme $genz, fractional difference $frac_diff" From 238b763f30454844f200860df9d79424f6837921 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 25 Jul 2024 17:37:40 +0100 Subject: [PATCH 43/74] move gradient code to relevant places --- src/Molly.jl | 1 - src/gradients.jl | 265 --------------------------- src/interactions/coulomb.jl | 21 +++ src/interactions/harmonic_angle.jl | 8 + src/interactions/harmonic_bond.jl | 8 + src/interactions/implicit_solvent.jl | 38 ++++ src/interactions/lennard_jones.jl | 12 ++ src/interactions/periodic_torsion.jl | 14 ++ src/types.jl | 161 ++++++++++++++++ test/gradients.jl | 4 +- 10 files changed, 264 insertions(+), 268 deletions(-) delete mode 100644 src/gradients.jl diff --git a/src/Molly.jl b/src/Molly.jl index 5720e1d67..49e893b7c 100644 --- a/src/Molly.jl +++ b/src/Molly.jl @@ -65,6 +65,5 @@ include("neighbors.jl") include("loggers.jl") include("analysis.jl") include("setup.jl") -include("gradients.jl") end diff --git a/src/gradients.jl b/src/gradients.jl deleted file mode 100644 index 90829ac9c..000000000 --- a/src/gradients.jl +++ /dev/null @@ -1,265 +0,0 @@ -# Utilities for taking gradients through simulations - -export - extract_parameters, - inject_gradients - -""" - extract_parameters(system, force_field) - -Form a `Dict` of all parameters in a [`System`](@ref), allowing gradients to be tracked. -""" -function extract_parameters(sys, ff) - params_dic = Dict() - - for at_data in sys.atoms_data - key_prefix = "atom_$(at_data.atom_type)_" - if !haskey(params_dic, key_prefix * "mass") - at = ff.atom_types[at_data.atom_type] - params_dic[key_prefix * "mass"] = at.mass - params_dic[key_prefix * "σ" ] = at.σ - params_dic[key_prefix * "ϵ" ] = at.ϵ - end - end - - for inter in values(sys.pairwise_inters) - if inter isa LennardJones - key_prefix = "inter_LJ_" - params_dic[key_prefix * "weight_14"] = inter.weight_special - elseif inter isa Coulomb - key_prefix = "inter_CO_" - params_dic[key_prefix * "weight_14"] = inter.weight_special - params_dic[key_prefix * "coulomb_const"] = inter.coulomb_const - elseif inter isa CoulombReactionField - key_prefix = "inter_CRF_" - params_dic[key_prefix * "dist_cutoff"] = inter.dist_cutoff - params_dic[key_prefix * "solvent_dielectric"] = inter.solvent_dielectric - params_dic[key_prefix * "weight_14"] = inter.weight_special - params_dic[key_prefix * "coulomb_const"] = inter.coulomb_const - end - end - - for inter in values(sys.specific_inter_lists) - if interaction_type(inter) <: HarmonicBond - for bond_type in inter.types - key_prefix = "inter_HB_$(bond_type)_" - if !haskey(params_dic, key_prefix * "k") - bond = ff.bond_types[atom_types_to_tuple(bond_type)] - params_dic[key_prefix * "k" ] = bond.k - params_dic[key_prefix * "r0"] = bond.r0 - end - end - elseif interaction_type(inter) <: HarmonicAngle - for angle_type in inter.types - key_prefix = "inter_HA_$(angle_type)_" - if !haskey(params_dic, key_prefix * "k") - ang = ff.angle_types[atom_types_to_tuple(angle_type)] - params_dic[key_prefix * "k" ] = ang.k - params_dic[key_prefix * "θ0"] = ang.θ0 - end - end - elseif interaction_type(inter) <: PeriodicTorsion - for (torsion_type, torsion_inter) in zip(inter.types, Array(inter.inters)) - if torsion_inter.proper - key_prefix = "inter_PT_$(torsion_type)_" - else - key_prefix = "inter_IT_$(torsion_type)_" - end - if !haskey(params_dic, key_prefix * "phase_1") - torsion = ff.torsion_types[atom_types_to_tuple(torsion_type)] - for i in eachindex(torsion.phases) - params_dic[key_prefix * "phase_$i"] = torsion.phases[i] - params_dic[key_prefix * "k_$i" ] = torsion.ks[i] - end - end - end - end - end - - return params_dic -end - -""" - inject_gradients(sys, params_dic) - -Add parameters from a dictionary to a [`System`](@ref). - -Allows gradients for individual parameters to be tracked. -Returns atoms, pairwise interactions, specific interaction lists and general -interactions. -""" -function inject_gradients(sys::System{D, G}, params_dic) where {D, G} - if G - atoms_grad = CuArray(inject_atom.(Array(sys.atoms), sys.atoms_data, (params_dic,))) - else - atoms_grad = inject_atom.(sys.atoms, sys.atoms_data, (params_dic,)) - end - if length(sys.pairwise_inters) > 0 - pis_grad = inject_interaction.(sys.pairwise_inters, (params_dic,)) - else - pis_grad = sys.pairwise_inters - end - if length(sys.specific_inter_lists) > 0 - sis_grad = inject_interaction_list.(sys.specific_inter_lists, (params_dic,), G) - else - sis_grad = sys.specific_inter_lists - end - if length(sys.general_inters) > 0 - gis_grad = inject_interaction.(sys.general_inters, (params_dic,), (sys,)) - else - gis_grad = sys.general_inters - end - return atoms_grad, pis_grad, sis_grad, gis_grad -end - -# get function errors with AD -dict_get(dic, key, default) = haskey(dic, key) ? dic[key] : default - -function inject_atom(at, at_data, params_dic) - key_prefix = "atom_$(at_data.atom_type)_" - Atom( - at.index, - at.atom_type, - dict_get(params_dic, key_prefix * "mass" , at.mass), - at.charge, # Residue-specific - dict_get(params_dic, key_prefix * "σ" , at.σ ), - dict_get(params_dic, key_prefix * "ϵ" , at.ϵ ), - ) -end - -function inject_interaction_list(inter::InteractionList1Atoms, params_dic, gpu) - if gpu - inters_grad = CuArray(inject_interaction.(Array(inter.inters), inter.types, (params_dic,))) - else - inters_grad = inject_interaction.(inter.inters, inter.types, (params_dic,)) - end - InteractionList1Atoms(inter.is, inters_grad, inter.types) -end - -function inject_interaction_list(inter::InteractionList2Atoms, params_dic, gpu) - if gpu - inters_grad = CuArray(inject_interaction.(Array(inter.inters), inter.types, (params_dic,))) - else - inters_grad = inject_interaction.(inter.inters, inter.types, (params_dic,)) - end - InteractionList2Atoms(inter.is, inter.js, inters_grad, inter.types) -end - -function inject_interaction_list(inter::InteractionList3Atoms, params_dic, gpu) - if gpu - inters_grad = CuArray(inject_interaction.(Array(inter.inters), inter.types, (params_dic,))) - else - inters_grad = inject_interaction.(inter.inters, inter.types, (params_dic,)) - end - InteractionList3Atoms(inter.is, inter.js, inter.ks, inters_grad, inter.types) -end - -function inject_interaction_list(inter::InteractionList4Atoms, params_dic, gpu) - if gpu - inters_grad = CuArray(inject_interaction.(Array(inter.inters), inter.types, (params_dic,))) - else - inters_grad = inject_interaction.(inter.inters, inter.types, (params_dic,)) - end - InteractionList4Atoms(inter.is, inter.js, inter.ks, inter.ls, inters_grad, inter.types) -end - -function inject_interaction(inter::LennardJones, params_dic) - key_prefix = "inter_LJ_" - LennardJones( - inter.cutoff, - inter.use_neighbors, - inter.shortcut, - inter.σ_mixing, - inter.ϵ_mixing, - dict_get(params_dic, key_prefix * "weight_14", inter.weight_special), - ) -end - -function inject_interaction(inter::Coulomb, params_dic) - key_prefix = "inter_CO_" - Coulomb( - inter.cutoff, - inter.use_neighbors, - dict_get(params_dic, key_prefix * "weight_14", inter.weight_special), - dict_get(params_dic, key_prefix * "coulomb_const", inter.coulomb_const), - ) -end - -function inject_interaction(inter::CoulombReactionField, params_dic) - key_prefix = "inter_CRF_" - CoulombReactionField( - dict_get(params_dic, key_prefix * "dist_cutoff", inter.dist_cutoff), - dict_get(params_dic, key_prefix * "solvent_dielectric", inter.solvent_dielectric), - inter.use_neighbors, - dict_get(params_dic, key_prefix * "weight_14", inter.weight_special), - dict_get(params_dic, key_prefix * "coulomb_const", inter.coulomb_const), - ) -end - -function inject_interaction(inter::HarmonicBond, inter_type, params_dic) - key_prefix = "inter_HB_$(inter_type)_" - HarmonicBond( - dict_get(params_dic, key_prefix * "k" , inter.k ), - dict_get(params_dic, key_prefix * "r0", inter.r0), - ) -end - -function inject_interaction(inter::HarmonicAngle, inter_type, params_dic) - key_prefix = "inter_HA_$(inter_type)_" - HarmonicAngle( - dict_get(params_dic, key_prefix * "k" , inter.k ), - dict_get(params_dic, key_prefix * "θ0", inter.θ0), - ) -end - -function inject_interaction(inter::PeriodicTorsion{N, T, E}, inter_type, params_dic) where {N, T, E} - if inter.proper - key_prefix = "inter_PT_$(inter_type)_" - else - key_prefix = "inter_IT_$(inter_type)_" - end - PeriodicTorsion{N, T, E}( - inter.periodicities, - ntuple(i -> dict_get(params_dic, key_prefix * "phase_$i", inter.phases[i]), N), - ntuple(i -> dict_get(params_dic, key_prefix * "k_$i" , inter.ks[i] ), N), - inter.proper, - ) -end - -function inject_interaction(inter::ImplicitSolventGBN2, params_dic, sys) - key_prefix = "inter_GB_" - bond_index = findfirst(sil -> eltype(sil.inters) <: HarmonicBond, sys.specific_inter_lists) - - element_to_radius = Dict{String, DefaultFloat}() # Units here made the gradients vanish - for k in keys(mbondi2_element_to_radius) - element_to_radius[k] = dict_get(params_dic, key_prefix * "radius_" * k, - ustrip(mbondi2_element_to_radius[k])) - end - element_to_screen = empty(gbn2_element_to_screen) - for k in keys(gbn2_element_to_screen) - element_to_screen[k] = dict_get(params_dic, key_prefix * "screen_" * k, gbn2_element_to_screen[k]) - end - atom_params = empty(gbn2_atom_params) - for k in keys(gbn2_atom_params) - atom_params[k] = dict_get(params_dic, key_prefix * "params_" * k, gbn2_atom_params[k]) - end - - ImplicitSolventGBN2( - sys.atoms, - sys.atoms_data, - sys.specific_inter_lists[bond_index]; - solvent_dielectric=dict_get(params_dic, key_prefix * "solvent_dielectric", inter.solvent_dielectric), - solute_dielectric=dict_get(params_dic, key_prefix * "solute_dielectric", inter.solute_dielectric), - kappa=dict_get(params_dic, key_prefix * "kappa", ustrip(inter.kappa))u"nm^-1", - offset=dict_get(params_dic, key_prefix * "offset", ustrip(inter.offset))u"nm", - dist_cutoff=inter.dist_cutoff, - probe_radius=dict_get(params_dic, key_prefix * "probe_radius", ustrip(inter.probe_radius))u"nm", - sa_factor=dict_get(params_dic, key_prefix * "sa_factor", ustrip(inter.sa_factor))u"kJ * mol^-1 * nm^-2", - use_ACE=inter.use_ACE, - neck_scale=dict_get(params_dic, key_prefix * "neck_scale", inter.neck_scale), - neck_cut=dict_get(params_dic, key_prefix * "neck_cut", ustrip(inter.neck_cut))u"nm", - element_to_radius=element_to_radius, - element_to_screen=element_to_screen, - atom_params=atom_params, - ) -end diff --git a/src/interactions/coulomb.jl b/src/interactions/coulomb.jl index b27a47fc4..49d98e04c 100644 --- a/src/interactions/coulomb.jl +++ b/src/interactions/coulomb.jl @@ -37,6 +37,16 @@ function Base.:+(c1::Coulomb, c2::Coulomb) ) end +function inject_interaction(inter::Coulomb, params_dic) + key_prefix = "inter_CO_" + return Coulomb( + inter.cutoff, + inter.use_neighbors, + dict_get(params_dic, key_prefix * "weight_14", inter.weight_special), + dict_get(params_dic, key_prefix * "coulomb_const", inter.coulomb_const), + ) +end + @inline function force(inter::Coulomb{C}, dr, atom_i, @@ -236,6 +246,17 @@ function Base.:+(c1::CoulombReactionField, c2::CoulombReactionField) ) end +function inject_interaction(inter::CoulombReactionField, params_dic) + key_prefix = "inter_CRF_" + return CoulombReactionField( + dict_get(params_dic, key_prefix * "dist_cutoff", inter.dist_cutoff), + dict_get(params_dic, key_prefix * "solvent_dielectric", inter.solvent_dielectric), + inter.use_neighbors, + dict_get(params_dic, key_prefix * "weight_14", inter.weight_special), + dict_get(params_dic, key_prefix * "coulomb_const", inter.coulomb_const), + ) +end + @inline function force(inter::CoulombReactionField, dr, atom_i, diff --git a/src/interactions/harmonic_angle.jl b/src/interactions/harmonic_angle.jl index e28ed0954..72e47f066 100644 --- a/src/interactions/harmonic_angle.jl +++ b/src/interactions/harmonic_angle.jl @@ -23,6 +23,14 @@ Base.zero(::HarmonicAngle{K, D}) where {K, D} = HarmonicAngle(k=zero(K), θ0=zer Base.:+(a1::HarmonicAngle, a2::HarmonicAngle) = HarmonicAngle(k=(a1.k + a2.k), θ0=(a1.θ0 + a2.θ0)) +function inject_interaction(inter::HarmonicAngle, inter_type, params_dic) + key_prefix = "inter_HA_$(inter_type)_" + return HarmonicAngle( + dict_get(params_dic, key_prefix * "k" , inter.k ), + dict_get(params_dic, key_prefix * "θ0", inter.θ0), + ) +end + @inline function force(a::HarmonicAngle, coords_i, coords_j, coords_k, boundary, args...) # In 2D we use then eliminate the cross product ba = vector_pad3D(coords_j, coords_i, boundary) diff --git a/src/interactions/harmonic_bond.jl b/src/interactions/harmonic_bond.jl index 291d6eef5..b4ed1600e 100644 --- a/src/interactions/harmonic_bond.jl +++ b/src/interactions/harmonic_bond.jl @@ -21,6 +21,14 @@ Base.zero(::HarmonicBond{K, D}) where {K, D} = HarmonicBond(k=zero(K), r0=zero(D Base.:+(b1::HarmonicBond, b2::HarmonicBond) = HarmonicBond(k=(b1.k + b2.k), r0=(b1.r0 + b2.r0)) +function inject_interaction(inter::HarmonicBond, inter_type, params_dic) + key_prefix = "inter_HB_$(inter_type)_" + return HarmonicBond( + dict_get(params_dic, key_prefix * "k" , inter.k ), + dict_get(params_dic, key_prefix * "r0", inter.r0), + ) +end + @inline function force(b::HarmonicBond, coord_i, coord_j, boundary, args...) ab = vector(coord_i, coord_j, boundary) c = b.k * (norm(ab) - b.r0) diff --git a/src/interactions/implicit_solvent.jl b/src/interactions/implicit_solvent.jl index 1cf40ab49..1f70a4af0 100644 --- a/src/interactions/implicit_solvent.jl +++ b/src/interactions/implicit_solvent.jl @@ -597,6 +597,44 @@ function ImplicitSolventGBN2(atoms::AbstractArray{Atom{TY, M, T, D, E}}, end end +function inject_interaction(inter::ImplicitSolventGBN2, params_dic, sys) + key_prefix = "inter_GB_" + bond_index = findfirst(sil -> eltype(sil.inters) <: HarmonicBond, sys.specific_inter_lists) + + element_to_radius = Dict{String, DefaultFloat}() # Units here made the gradients vanish + for k in keys(mbondi2_element_to_radius) + element_to_radius[k] = dict_get(params_dic, key_prefix * "radius_" * k, + ustrip(mbondi2_element_to_radius[k])) + end + element_to_screen = empty(gbn2_element_to_screen) + for k in keys(gbn2_element_to_screen) + element_to_screen[k] = dict_get(params_dic, key_prefix * "screen_" * k, gbn2_element_to_screen[k]) + end + atom_params = empty(gbn2_atom_params) + for k in keys(gbn2_atom_params) + atom_params[k] = dict_get(params_dic, key_prefix * "params_" * k, gbn2_atom_params[k]) + end + + ImplicitSolventGBN2( + sys.atoms, + sys.atoms_data, + sys.specific_inter_lists[bond_index]; + solvent_dielectric=dict_get(params_dic, key_prefix * "solvent_dielectric", inter.solvent_dielectric), + solute_dielectric=dict_get(params_dic, key_prefix * "solute_dielectric", inter.solute_dielectric), + kappa=dict_get(params_dic, key_prefix * "kappa", ustrip(inter.kappa))u"nm^-1", + offset=dict_get(params_dic, key_prefix * "offset", ustrip(inter.offset))u"nm", + dist_cutoff=inter.dist_cutoff, + probe_radius=dict_get(params_dic, key_prefix * "probe_radius", ustrip(inter.probe_radius))u"nm", + sa_factor=dict_get(params_dic, key_prefix * "sa_factor", ustrip(inter.sa_factor))u"kJ * mol^-1 * nm^-2", + use_ACE=inter.use_ACE, + neck_scale=dict_get(params_dic, key_prefix * "neck_scale", inter.neck_scale), + neck_cut=dict_get(params_dic, key_prefix * "neck_cut", ustrip(inter.neck_cut))u"nm", + element_to_radius=element_to_radius, + element_to_screen=element_to_screen, + atom_params=atom_params, + ) +end + function born_radii_loop_OBC(coord_i, coord_j, ori, srj, dist_cutoff, boundary) I = zero(coord_i[1] / unit(dist_cutoff)^2) r = norm(vector(coord_i, coord_j, boundary)) diff --git a/src/interactions/lennard_jones.jl b/src/interactions/lennard_jones.jl index 7a5450ac7..adc9577d7 100644 --- a/src/interactions/lennard_jones.jl +++ b/src/interactions/lennard_jones.jl @@ -71,6 +71,18 @@ function Base.:+(l1::LennardJones, l2::LennardJones) ) end +function inject_interaction(inter::LennardJones, params_dic) + key_prefix = "inter_LJ_" + return LennardJones( + inter.cutoff, + inter.use_neighbors, + inter.shortcut, + inter.σ_mixing, + inter.ϵ_mixing, + dict_get(params_dic, key_prefix * "weight_14", inter.weight_special), + ) +end + @inline function force(inter::LennardJones, dr, atom_i, diff --git a/src/interactions/periodic_torsion.jl b/src/interactions/periodic_torsion.jl index ce6db4b26..0448a276b 100644 --- a/src/interactions/periodic_torsion.jl +++ b/src/interactions/periodic_torsion.jl @@ -51,6 +51,20 @@ function Base.:+(p1::PeriodicTorsion{N, T, E}, p2::PeriodicTorsion{N, T, E}) whe ) end +function inject_interaction(inter::PeriodicTorsion{N, T, E}, inter_type, params_dic) where {N, T, E} + if inter.proper + key_prefix = "inter_PT_$(inter_type)_" + else + key_prefix = "inter_IT_$(inter_type)_" + end + return PeriodicTorsion{N, T, E}( + inter.periodicities, + ntuple(i -> dict_get(params_dic, key_prefix * "phase_$i", inter.phases[i]), N), + ntuple(i -> dict_get(params_dic, key_prefix * "k_$i" , inter.ks[i] ), N), + inter.proper, + ) +end + function periodic_torsion_vectors(coords_i, coords_j, coords_k, coords_l, boundary) ab = vector(coords_i, coords_j, boundary) bc = vector(coords_j, coords_k, boundary) diff --git a/src/types.jl b/src/types.jl index 260f04ac4..c91cbeee7 100644 --- a/src/types.jl +++ b/src/types.jl @@ -12,6 +12,8 @@ export MolecularTopology, NeighborList, System, + inject_gradients, + extract_parameters, ReplicaSystem, is_on_gpu, float_type, @@ -180,6 +182,42 @@ function Base.:+(il1::InteractionList4Atoms{I, T}, il2::InteractionList4Atoms{I, ) end +function inject_interaction_list(inter::InteractionList1Atoms, params_dic, gpu) + if gpu + inters_grad = CuArray(inject_interaction.(Array(inter.inters), inter.types, (params_dic,))) + else + inters_grad = inject_interaction.(inter.inters, inter.types, (params_dic,)) + end + InteractionList1Atoms(inter.is, inters_grad, inter.types) +end + +function inject_interaction_list(inter::InteractionList2Atoms, params_dic, gpu) + if gpu + inters_grad = CuArray(inject_interaction.(Array(inter.inters), inter.types, (params_dic,))) + else + inters_grad = inject_interaction.(inter.inters, inter.types, (params_dic,)) + end + InteractionList2Atoms(inter.is, inter.js, inters_grad, inter.types) +end + +function inject_interaction_list(inter::InteractionList3Atoms, params_dic, gpu) + if gpu + inters_grad = CuArray(inject_interaction.(Array(inter.inters), inter.types, (params_dic,))) + else + inters_grad = inject_interaction.(inter.inters, inter.types, (params_dic,)) + end + InteractionList3Atoms(inter.is, inter.js, inter.ks, inters_grad, inter.types) +end + +function inject_interaction_list(inter::InteractionList4Atoms, params_dic, gpu) + if gpu + inters_grad = CuArray(inject_interaction.(Array(inter.inters), inter.types, (params_dic,))) + else + inters_grad = inject_interaction.(inter.inters, inter.types, (params_dic,)) + end + InteractionList4Atoms(inter.is, inter.js, inter.ks, inter.ls, inters_grad, inter.types) +end + """ Atom(; ) @@ -230,6 +268,21 @@ function Base.getindex(atom::Atom, x::Symbol) return hasfield(Atom, x) ? getfield(atom, x) : KeyError("no field $x in Atom") end +# get function errors with AD +dict_get(dic, key, default) = haskey(dic, key) ? dic[key] : default + +function inject_atom(at, at_data, params_dic) + key_prefix = "atom_$(at_data.atom_type)_" + Atom( + at.index, + at.atom_type, + dict_get(params_dic, key_prefix * "mass" , at.mass), + at.charge, # Residue-specific + dict_get(params_dic, key_prefix * "σ" , at.σ ), + dict_get(params_dic, key_prefix * "ϵ" , at.ϵ ), + ) +end + """ charge(atom) @@ -673,6 +726,114 @@ function System(crystal::Crystal{D}; ) end +""" + inject_gradients(sys, params_dic) + +Add parameters from a dictionary to a [`System`](@ref). + +Allows gradients for individual parameters to be tracked. +Returns atoms, pairwise interactions, specific interaction lists and general +interactions. +""" +function inject_gradients(sys::System{D, G}, params_dic) where {D, G} + if G + atoms_grad = CuArray(inject_atom.(Array(sys.atoms), sys.atoms_data, (params_dic,))) + else + atoms_grad = inject_atom.(sys.atoms, sys.atoms_data, (params_dic,)) + end + if length(sys.pairwise_inters) > 0 + pis_grad = inject_interaction.(sys.pairwise_inters, (params_dic,)) + else + pis_grad = sys.pairwise_inters + end + if length(sys.specific_inter_lists) > 0 + sis_grad = inject_interaction_list.(sys.specific_inter_lists, (params_dic,), G) + else + sis_grad = sys.specific_inter_lists + end + if length(sys.general_inters) > 0 + gis_grad = inject_interaction.(sys.general_inters, (params_dic,), (sys,)) + else + gis_grad = sys.general_inters + end + return atoms_grad, pis_grad, sis_grad, gis_grad +end + +""" + extract_parameters(system, force_field) + +Form a `Dict` of all parameters in a [`System`](@ref), allowing gradients to be tracked. +""" +function extract_parameters(sys, ff) + params_dic = Dict() + + for at_data in sys.atoms_data + key_prefix = "atom_$(at_data.atom_type)_" + if !haskey(params_dic, key_prefix * "mass") + at = ff.atom_types[at_data.atom_type] + params_dic[key_prefix * "mass"] = at.mass + params_dic[key_prefix * "σ" ] = at.σ + params_dic[key_prefix * "ϵ" ] = at.ϵ + end + end + + for inter in values(sys.pairwise_inters) + if inter isa LennardJones + key_prefix = "inter_LJ_" + params_dic[key_prefix * "weight_14"] = inter.weight_special + elseif inter isa Coulomb + key_prefix = "inter_CO_" + params_dic[key_prefix * "weight_14"] = inter.weight_special + params_dic[key_prefix * "coulomb_const"] = inter.coulomb_const + elseif inter isa CoulombReactionField + key_prefix = "inter_CRF_" + params_dic[key_prefix * "dist_cutoff"] = inter.dist_cutoff + params_dic[key_prefix * "solvent_dielectric"] = inter.solvent_dielectric + params_dic[key_prefix * "weight_14"] = inter.weight_special + params_dic[key_prefix * "coulomb_const"] = inter.coulomb_const + end + end + + for inter in values(sys.specific_inter_lists) + if interaction_type(inter) <: HarmonicBond + for bond_type in inter.types + key_prefix = "inter_HB_$(bond_type)_" + if !haskey(params_dic, key_prefix * "k") + bond = ff.bond_types[atom_types_to_tuple(bond_type)] + params_dic[key_prefix * "k" ] = bond.k + params_dic[key_prefix * "r0"] = bond.r0 + end + end + elseif interaction_type(inter) <: HarmonicAngle + for angle_type in inter.types + key_prefix = "inter_HA_$(angle_type)_" + if !haskey(params_dic, key_prefix * "k") + ang = ff.angle_types[atom_types_to_tuple(angle_type)] + params_dic[key_prefix * "k" ] = ang.k + params_dic[key_prefix * "θ0"] = ang.θ0 + end + end + elseif interaction_type(inter) <: PeriodicTorsion + for (torsion_type, torsion_inter) in zip(inter.types, Array(inter.inters)) + if torsion_inter.proper + key_prefix = "inter_PT_$(torsion_type)_" + else + key_prefix = "inter_IT_$(torsion_type)_" + end + if !haskey(params_dic, key_prefix * "phase_1") + torsion = ff.torsion_types[atom_types_to_tuple(torsion_type)] + for i in eachindex(torsion.phases) + params_dic[key_prefix * "phase_$i"] = torsion.phases[i] + params_dic[key_prefix * "k_$i" ] = torsion.ks[i] + end + end + end + end + end + + return params_dic +end + """ ReplicaSystem(; ) diff --git a/test/gradients.jl b/test/gradients.jl index 9bcf25ff6..d5a431be4 100644 --- a/test/gradients.jl +++ b/test/gradients.jl @@ -277,7 +277,7 @@ end return potential_energy(sys; n_threads=n_threads) end - function test_force_grad(params_dic, sys_ref, coords, neighbor_finder, n_threads) + function test_forces_grad(params_dic, sys_ref, coords, neighbor_finder, n_threads) atoms, pis, sis, gis = inject_gradients(sys_ref, params_dic) sys = System( @@ -409,7 +409,7 @@ end end test_runs = [ ("Energy", test_energy_grad, 1e-8), - ("Force" , test_force_grad , 1e-8), + ("Force" , test_forces_grad , 1e-8), ] if !running_CI #push!(test_runs, ("Sim", test_sim_grad, 0.015)) From b081caa0f48e2a69cd2d68d47fdcce57f9b49263 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 25 Jul 2024 18:29:42 +0100 Subject: [PATCH 44/74] move parameter extraction to files --- src/interactions/coulomb.jl | 16 +++++++++ src/interactions/harmonic_angle.jl | 14 ++++++++ src/interactions/harmonic_bond.jl | 14 ++++++++ src/interactions/lennard_jones.jl | 6 ++++ src/interactions/periodic_torsion.jl | 20 +++++++++++ src/types.jl | 50 ++-------------------------- 6 files changed, 72 insertions(+), 48 deletions(-) diff --git a/src/interactions/coulomb.jl b/src/interactions/coulomb.jl index 49d98e04c..46815d989 100644 --- a/src/interactions/coulomb.jl +++ b/src/interactions/coulomb.jl @@ -47,6 +47,13 @@ function inject_interaction(inter::Coulomb, params_dic) ) end +function extract_parameters!(params_dic, inter::Coulomb, ff) + key_prefix = "inter_CO_" + params_dic[key_prefix * "weight_14"] = inter.weight_special + params_dic[key_prefix * "coulomb_const"] = inter.coulomb_const + return params_dic +end + @inline function force(inter::Coulomb{C}, dr, atom_i, @@ -257,6 +264,15 @@ function inject_interaction(inter::CoulombReactionField, params_dic) ) end +function extract_parameters!(params_dic, inter::CoulombReactionField, ff) + key_prefix = "inter_CRF_" + params_dic[key_prefix * "dist_cutoff"] = inter.dist_cutoff + params_dic[key_prefix * "solvent_dielectric"] = inter.solvent_dielectric + params_dic[key_prefix * "weight_14"] = inter.weight_special + params_dic[key_prefix * "coulomb_const"] = inter.coulomb_const + return params_dic +end + @inline function force(inter::CoulombReactionField, dr, atom_i, diff --git a/src/interactions/harmonic_angle.jl b/src/interactions/harmonic_angle.jl index 72e47f066..5a2ef069b 100644 --- a/src/interactions/harmonic_angle.jl +++ b/src/interactions/harmonic_angle.jl @@ -31,6 +31,20 @@ function inject_interaction(inter::HarmonicAngle, inter_type, params_dic) ) end +function extract_parameters!(params_dic, + inter::InteractionList3Atoms{<:Any, <:AbstractVector{<:HarmonicAngle}}, + ff) + for angle_type in inter.types + key_prefix = "inter_HA_$(angle_type)_" + if !haskey(params_dic, key_prefix * "k") + ang = ff.angle_types[atom_types_to_tuple(angle_type)] + params_dic[key_prefix * "k" ] = ang.k + params_dic[key_prefix * "θ0"] = ang.θ0 + end + end + return params_dic +end + @inline function force(a::HarmonicAngle, coords_i, coords_j, coords_k, boundary, args...) # In 2D we use then eliminate the cross product ba = vector_pad3D(coords_j, coords_i, boundary) diff --git a/src/interactions/harmonic_bond.jl b/src/interactions/harmonic_bond.jl index b4ed1600e..5bf33ea21 100644 --- a/src/interactions/harmonic_bond.jl +++ b/src/interactions/harmonic_bond.jl @@ -29,6 +29,20 @@ function inject_interaction(inter::HarmonicBond, inter_type, params_dic) ) end +function extract_parameters!(params_dic, + inter::InteractionList2Atoms{<:Any, <:AbstractVector{<:HarmonicBond}}, + ff) + for bond_type in inter.types + key_prefix = "inter_HB_$(bond_type)_" + if !haskey(params_dic, key_prefix * "k") + bond = ff.bond_types[atom_types_to_tuple(bond_type)] + params_dic[key_prefix * "k" ] = bond.k + params_dic[key_prefix * "r0"] = bond.r0 + end + end + return params_dic +end + @inline function force(b::HarmonicBond, coord_i, coord_j, boundary, args...) ab = vector(coord_i, coord_j, boundary) c = b.k * (norm(ab) - b.r0) diff --git a/src/interactions/lennard_jones.jl b/src/interactions/lennard_jones.jl index adc9577d7..7cc843750 100644 --- a/src/interactions/lennard_jones.jl +++ b/src/interactions/lennard_jones.jl @@ -83,6 +83,12 @@ function inject_interaction(inter::LennardJones, params_dic) ) end +function extract_parameters!(params_dic, inter::LennardJones, ff) + key_prefix = "inter_LJ_" + params_dic[key_prefix * "weight_14"] = inter.weight_special + return params_dic +end + @inline function force(inter::LennardJones, dr, atom_i, diff --git a/src/interactions/periodic_torsion.jl b/src/interactions/periodic_torsion.jl index 0448a276b..034809e38 100644 --- a/src/interactions/periodic_torsion.jl +++ b/src/interactions/periodic_torsion.jl @@ -65,6 +65,26 @@ function inject_interaction(inter::PeriodicTorsion{N, T, E}, inter_type, params_ ) end +function extract_parameters!(params_dic, + inter::InteractionList4Atoms{<:Any, <:AbstractVector{<:PeriodicTorsion}}, + ff) + for (torsion_type, torsion_inter) in zip(inter.types, Array(inter.inters)) + if torsion_inter.proper + key_prefix = "inter_PT_$(torsion_type)_" + else + key_prefix = "inter_IT_$(torsion_type)_" + end + if !haskey(params_dic, key_prefix * "phase_1") + torsion = ff.torsion_types[atom_types_to_tuple(torsion_type)] + for i in eachindex(torsion.phases) + params_dic[key_prefix * "phase_$i"] = torsion.phases[i] + params_dic[key_prefix * "k_$i" ] = torsion.ks[i] + end + end + end + return params_dic +end + function periodic_torsion_vectors(coords_i, coords_j, coords_k, coords_l, boundary) ab = vector(coords_i, coords_j, boundary) bc = vector(coords_j, coords_k, boundary) diff --git a/src/types.jl b/src/types.jl index c91cbeee7..9be353b69 100644 --- a/src/types.jl +++ b/src/types.jl @@ -778,57 +778,11 @@ function extract_parameters(sys, ff) end for inter in values(sys.pairwise_inters) - if inter isa LennardJones - key_prefix = "inter_LJ_" - params_dic[key_prefix * "weight_14"] = inter.weight_special - elseif inter isa Coulomb - key_prefix = "inter_CO_" - params_dic[key_prefix * "weight_14"] = inter.weight_special - params_dic[key_prefix * "coulomb_const"] = inter.coulomb_const - elseif inter isa CoulombReactionField - key_prefix = "inter_CRF_" - params_dic[key_prefix * "dist_cutoff"] = inter.dist_cutoff - params_dic[key_prefix * "solvent_dielectric"] = inter.solvent_dielectric - params_dic[key_prefix * "weight_14"] = inter.weight_special - params_dic[key_prefix * "coulomb_const"] = inter.coulomb_const - end + extract_parameters!(params_dic, inter, ff) end for inter in values(sys.specific_inter_lists) - if interaction_type(inter) <: HarmonicBond - for bond_type in inter.types - key_prefix = "inter_HB_$(bond_type)_" - if !haskey(params_dic, key_prefix * "k") - bond = ff.bond_types[atom_types_to_tuple(bond_type)] - params_dic[key_prefix * "k" ] = bond.k - params_dic[key_prefix * "r0"] = bond.r0 - end - end - elseif interaction_type(inter) <: HarmonicAngle - for angle_type in inter.types - key_prefix = "inter_HA_$(angle_type)_" - if !haskey(params_dic, key_prefix * "k") - ang = ff.angle_types[atom_types_to_tuple(angle_type)] - params_dic[key_prefix * "k" ] = ang.k - params_dic[key_prefix * "θ0"] = ang.θ0 - end - end - elseif interaction_type(inter) <: PeriodicTorsion - for (torsion_type, torsion_inter) in zip(inter.types, Array(inter.inters)) - if torsion_inter.proper - key_prefix = "inter_PT_$(torsion_type)_" - else - key_prefix = "inter_IT_$(torsion_type)_" - end - if !haskey(params_dic, key_prefix * "phase_1") - torsion = ff.torsion_types[atom_types_to_tuple(torsion_type)] - for i in eachindex(torsion.phases) - params_dic[key_prefix * "phase_$i"] = torsion.phases[i] - params_dic[key_prefix * "k_$i" ] = torsion.ks[i] - end - end - end - end + extract_parameters!(params_dic, inter, ff) end return params_dic From c3c1618cc569648c5ea2ad6837ed5e045b3bc409 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 26 Jul 2024 13:01:11 +0100 Subject: [PATCH 45/74] deepcopy to copy --- benchmark/benchmarks.jl | 8 ++++---- docs/src/differentiable.md | 8 ++++---- src/setup.jl | 6 +++--- test/basic.jl | 2 +- test/gradients.jl | 4 ++-- test/protein.jl | 10 +++++----- test/simulation.jl | 20 ++++++++++---------- 7 files changed, 29 insertions(+), 29 deletions(-) diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index e07a86cd3..08e6c5b4a 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -90,13 +90,13 @@ function test_sim(nl::Bool, parallel::Bool, f32::Bool, gpu::Bool) end if gpu - coords = CuArray(deepcopy(f32 ? starting_coords_f32 : starting_coords)) - velocities = CuArray(deepcopy(f32 ? starting_velocities_f32 : starting_velocities)) + coords = CuArray(copy(f32 ? starting_coords_f32 : starting_coords)) + velocities = CuArray(copy(f32 ? starting_velocities_f32 : starting_velocities)) atoms = CuArray([Atom(mass=atom_mass, charge=f32 ? 0.0f0 : 0.0, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", ϵ=f32 ? 0.2f0u"kJ * mol^-1" : 0.2u"kJ * mol^-1") for i in 1:n_atoms]) else - coords = deepcopy(f32 ? starting_coords_f32 : starting_coords) - velocities = deepcopy(f32 ? starting_velocities_f32 : starting_velocities) + coords = copy(f32 ? starting_coords_f32 : starting_coords) + velocities = copy(f32 ? starting_velocities_f32 : starting_velocities) atoms = [Atom(mass=atom_mass, charge=f32 ? 0.0f0 : 0.0, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", ϵ=f32 ? 0.2f0u"kJ * mol^-1" : 0.2u"kJ * mol^-1") for i in 1:n_atoms] end diff --git a/docs/src/differentiable.md b/docs/src/differentiable.md index ad73ad6ef..bed2d2611 100644 --- a/docs/src/differentiable.md +++ b/docs/src/differentiable.md @@ -202,9 +202,9 @@ function loss(θ) sys = System( atoms=atoms, - coords=deepcopy(coords), + coords=copy(coords), boundary=boundary, - velocities=deepcopy(velocities), + velocities=copy(velocities), specific_inter_lists=specific_inter_lists, loggers=loggers, force_units=NoUnits, @@ -332,9 +332,9 @@ function loss() sys = System( atoms=atoms, - coords=deepcopy(coords), + coords=copy(coords), boundary=boundary, - velocities=deepcopy(velocities), + velocities=copy(velocities), general_inters=general_inters, loggers=loggers, force_units=NoUnits, diff --git a/src/setup.jl b/src/setup.jl index 614ec4b11..a33919016 100644 --- a/src/setup.jl +++ b/src/setup.jl @@ -1390,9 +1390,9 @@ function add_position_restraints(sys, sis = (sys.specific_inter_lists..., restraints) return System( atoms=deepcopy(sys.atoms), - coords=deepcopy(sys.coords), - boundary=deepcopy(sys.boundary), - velocities=deepcopy(sys.velocities), + coords=copy(sys.coords), + boundary=copy(sys.boundary), + velocities=copy(sys.velocities), atoms_data=deepcopy(sys.atoms_data), topology=deepcopy(sys.topology), pairwise_inters=deepcopy(sys.pairwise_inters), diff --git a/test/basic.jl b/test/basic.jl index 00326ced0..90a1bc70d 100644 --- a/test/basic.jl +++ b/test/basic.jl @@ -188,7 +188,7 @@ eligible=(gpu ? CuArray(no_nbs) : no_nbs), dist_cutoff=1.0u"nm", ) - coords_start = deepcopy(sys.coords) + coords_start = copy(sys.coords) pe_start = potential_energy(sys, find_neighbors(sys)) scale_factor = 1.02 n_scales = 10 diff --git a/test/gradients.jl b/test/gradients.jl index d5a431be4..9b04fa21b 100644 --- a/test/gradients.jl +++ b/test/gradients.jl @@ -429,7 +429,7 @@ end grads_enzyme = Dict(k => 0.0 for k in keys(params_dic)) autodiff( Reverse, test_fn, Active, Duplicated(params_dic, grads_enzyme), - Const(sys_ref), Duplicated(sys_ref.coords, zero(sys_ref.coords)), + Const(sys_ref), Duplicated(copy(sys_ref.coords), zero(sys_ref.coords)), Duplicated(sys_ref.neighbor_finder, sys_ref.neighbor_finder), Const(n_threads), ) @@ -439,7 +439,7 @@ end gfd = central_fdm(6, 1)(params_dic[param]) do val dic = copy(params_dic) dic[param] = val - test_fn(dic, sys_ref, sys_ref.coords, sys_ref.neighbor_finder, n_threads) + test_fn(dic, sys_ref, copy(sys_ref.coords), sys_ref.neighbor_finder, n_threads) end frac_diff = abs(genz - gfd) / abs(gfd) @info "$(rpad(test_name, 6)) - $(rpad(platform, 12)) - $(rpad(param, 21)) - FD $gfd, Enzyme $genz, fractional difference $frac_diff" diff --git a/test/protein.jl b/test/protein.jl index 64331880a..c7c9db637 100644 --- a/test/protein.jl +++ b/test/protein.jl @@ -127,7 +127,7 @@ end n_steps = 100 simulator = VelocityVerlet(dt=0.0005u"ps") velocities_start = SVector{3}.(eachrow(readdlm(joinpath(openmm_dir, "velocities_300K.txt"))))u"nm * ps^-1" - sys.velocities = deepcopy(velocities_start) + sys.velocities = copy(velocities_start) @test kinetic_energy(sys) ≈ 65521.87288132431u"kJ * mol^-1" @test temperature(sys) ≈ 329.3202932884933u"K" @@ -150,7 +150,7 @@ end sys_nounits = System( joinpath(data_dir, "6mrr_equil.pdb"), ff_nounits; - velocities=deepcopy(ustrip_vec.(velocities_start)), + velocities=copy(ustrip_vec.(velocities_start)), units=false, center_coords=false, ) @@ -181,7 +181,7 @@ end sys = System( joinpath(data_dir, "6mrr_equil.pdb"), ff; - velocities=CuArray(deepcopy(velocities_start)), + velocities=CuArray(copy(velocities_start)), gpu=true, center_coords=false, ) @@ -209,7 +209,7 @@ end sys_nounits = System( joinpath(data_dir, "6mrr_equil.pdb"), ff_nounits; - velocities=CuArray(deepcopy(ustrip_vec.(velocities_start))), + velocities=CuArray(copy(ustrip_vec.(velocities_start))), units=false, gpu=true, center_coords=false, @@ -278,7 +278,7 @@ end if solvent_model == "gbn2" sim = SteepestDescentMinimizer(tol=400.0u"kJ * mol^-1 * nm^-1") - coords_start = deepcopy(sys.coords) + coords_start = copy(sys.coords) simulate!(sys, sim) @test potential_energy(sys) < E_molly @test rmsd(coords_start, sys.coords) < 0.1u"nm" diff --git a/test/simulation.jl b/test/simulation.jl index cae0665ac..6aafc3da1 100644 --- a/test/simulation.jl +++ b/test/simulation.jl @@ -436,7 +436,7 @@ end sys = System( atoms=(gpu ? CuArray(atoms) : atoms), - coords=(gpu ? CuArray(deepcopy(starting_coords)) : deepcopy(starting_coords)), + coords=(gpu ? CuArray(copy(starting_coords)) : copy(starting_coords)), boundary=boundary, atoms_data=atoms_data, pairwise_inters=(LennardJones(),), @@ -555,7 +555,7 @@ end loggers=(coords=CoordinateLogger(10),), ) - old_coords = deepcopy(sys.coords) + old_coords = copy(sys.coords) for i in eachindex(sys.coords) sys.coords[i] += [rand()*0.01, rand()*0.01, rand()*0.01]u"nm" @@ -856,7 +856,7 @@ end sys = System( atoms=atoms, - coords=deepcopy(coords), + coords=copy(coords), boundary=boundary, pairwise_inters=(LennardJones(),), loggers=( @@ -896,7 +896,7 @@ end sys = System( atoms=AT(atoms), - coords=AT(deepcopy(coords)), + coords=AT(copy(coords)), boundary=boundary, pairwise_inters=(LennardJones(),), loggers=( @@ -963,7 +963,7 @@ end sys = System( atoms=AT(atoms), - coords=AT(deepcopy(coords)), + coords=AT(copy(coords)), boundary=boundary, pairwise_inters=(LennardJones(),), loggers=( @@ -1034,7 +1034,7 @@ end sys = System( atoms=AT(atoms), - coords=AT(deepcopy(coords)), + coords=AT(copy(coords)), boundary=boundary, pairwise_inters=(LennardJones(),), loggers=( @@ -1177,13 +1177,13 @@ end show(devnull, neighbor_finder) if gpu - coords = CuArray(deepcopy(f32 ? starting_coords_f32 : starting_coords)) - velocities = CuArray(deepcopy(f32 ? starting_velocities_f32 : starting_velocities)) + coords = CuArray(copy(f32 ? starting_coords_f32 : starting_coords)) + velocities = CuArray(copy(f32 ? starting_velocities_f32 : starting_velocities)) atoms = CuArray([Atom(mass=atom_mass, charge=f32 ? 0.0f0 : 0.0, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", ϵ=f32 ? 0.2f0u"kJ * mol^-1" : 0.2u"kJ * mol^-1") for i in 1:n_atoms]) else - coords = deepcopy(f32 ? starting_coords_f32 : starting_coords) - velocities = deepcopy(f32 ? starting_velocities_f32 : starting_velocities) + coords = copy(f32 ? starting_coords_f32 : starting_coords) + velocities = copy(f32 ? starting_velocities_f32 : starting_velocities) atoms = [Atom(mass=atom_mass, charge=f32 ? 0.0f0 : 0.0, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", ϵ=f32 ? 0.2f0u"kJ * mol^-1" : 0.2u"kJ * mol^-1") for i in 1:n_atoms] end From 65e61ef4db642f3f5acfead7a3afecafe38a0e92 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 26 Jul 2024 15:00:55 +0100 Subject: [PATCH 46/74] Add @inline to simulators for Enzyme --- src/simulators.jl | 89 ++++++++++++++++++++++++----------------------- 1 file changed, 45 insertions(+), 44 deletions(-) diff --git a/src/simulators.jl b/src/simulators.jl index 456487b98..a421fc332 100644 --- a/src/simulators.jl +++ b/src/simulators.jl @@ -59,10 +59,11 @@ are not run before the first step. it is `false`. Custom simulators should implement this function. """ -function simulate!(sys, - sim::SteepestDescentMinimizer; - n_threads::Integer=Threads.nthreads(), - run_loggers=false) +@inline function simulate!(sys, + sim::SteepestDescentMinimizer; + n_threads::Integer=Threads.nthreads(), + run_loggers=false) + # @inline needed to avoid Enzyme error sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) E = potential_energy(sys, neighbors; n_threads=n_threads) @@ -134,11 +135,11 @@ function VelocityVerlet(; dt, coupling=NoCoupling(), remove_CM_motion=1) return VelocityVerlet(dt, coupling, Int(remove_CM_motion)) end -function simulate!(sys, - sim::VelocityVerlet, - n_steps::Integer; - n_threads::Integer=Threads.nthreads(), - run_loggers=true) +@inline function simulate!(sys, + sim::VelocityVerlet, + n_steps::Integer; + n_threads::Integer=Threads.nthreads(), + run_loggers=true) sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) !iszero(sim.remove_CM_motion) && remove_CM_motion!(sys) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) @@ -221,11 +222,11 @@ function Verlet(; dt, coupling=NoCoupling(), remove_CM_motion=1) return Verlet(dt, coupling, Int(remove_CM_motion)) end -function simulate!(sys, - sim::Verlet, - n_steps::Integer; - n_threads::Integer=Threads.nthreads(), - run_loggers=true) +@inline function simulate!(sys, + sim::Verlet, + n_steps::Integer; + n_threads::Integer=Threads.nthreads(), + run_loggers=true) sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) !iszero(sim.remove_CM_motion) && remove_CM_motion!(sys) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) @@ -296,11 +297,11 @@ end StormerVerlet(; dt, coupling=NoCoupling()) = StormerVerlet(dt, coupling) -function simulate!(sys, - sim::StormerVerlet, - n_steps::Integer; - n_threads::Integer=Threads.nthreads(), - run_loggers=true) +@inline function simulate!(sys, + sim::StormerVerlet, + n_steps::Integer; + n_threads::Integer=Threads.nthreads(), + run_loggers=true) sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) apply_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) @@ -378,12 +379,12 @@ function Langevin(; dt, temperature, friction, coupling=NoCoupling(), remove_CM_ vel_scale, noise_scale) end -function simulate!(sys, - sim::Langevin, - n_steps::Integer; - n_threads::Integer=Threads.nthreads(), - run_loggers=true, - rng=Random.GLOBAL_RNG) +@inline function simulate!(sys, + sim::Langevin, + n_steps::Integer; + n_threads::Integer=Threads.nthreads(), + run_loggers=true, + rng=Random.GLOBAL_RNG) sys.coords .= wrap_coords.(sys.coords, (sys.boundary,)) !iszero(sim.remove_CM_motion) && remove_CM_motion!(sys) neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) @@ -479,12 +480,12 @@ function LangevinSplitting(; dt, temperature, friction, splitting, remove_CM_mot dt, temperature, friction, splitting, Int(remove_CM_motion)) end -function simulate!(sys, - sim::LangevinSplitting, - n_steps::Integer; - n_threads::Integer=Threads.nthreads(), - run_loggers=true, - rng=Random.GLOBAL_RNG) +@inline function simulate!(sys, + sim::LangevinSplitting, + n_steps::Integer; + n_threads::Integer=Threads.nthreads(), + run_loggers=true, + rng=Random.GLOBAL_RNG) if length(sys.constraints) > 0 @warn "LangevinSplitting is not currently compatible with constraints, " * "constraints will be ignored" @@ -603,12 +604,12 @@ function OverdampedLangevin(; dt, temperature, friction, remove_CM_motion=1) return OverdampedLangevin(dt, temperature, friction, Int(remove_CM_motion)) end -function simulate!(sys, - sim::OverdampedLangevin, - n_steps::Integer; - n_threads::Integer=Threads.nthreads(), - run_loggers=true, - rng=Random.GLOBAL_RNG) +@inline function simulate!(sys, + sim::OverdampedLangevin, + n_steps::Integer; + n_threads::Integer=Threads.nthreads(), + run_loggers=true, + rng=Random.GLOBAL_RNG) if length(sys.constraints) > 0 @warn "OverdampedLangevin is not currently compatible with constraints, " * "constraints will be ignored" @@ -678,8 +679,8 @@ function NoseHoover(; dt, temperature, damping=100*dt, coupling=NoCoupling(), re return NoseHoover(dt, temperature, damping, coupling, Int(remove_CM_motion)) end -function simulate!(sys, sim::NoseHoover, n_steps::Integer; - n_threads::Integer=Threads.nthreads(), run_loggers=true) +@inline function simulate!(sys, sim::NoseHoover, n_steps::Integer; + n_threads::Integer=Threads.nthreads(), run_loggers=true) if length(sys.constraints) > 0 @warn "NoseHoover is not currently compatible with constraints, " * "constraints will be ignored" @@ -1040,11 +1041,11 @@ function MetropolisMonteCarlo(; temperature, trial_moves, trial_args=Dict()) return MetropolisMonteCarlo(temperature, trial_moves, trial_args) end -function simulate!(sys::System{D, G, T}, - sim::MetropolisMonteCarlo, - n_steps::Integer; - n_threads::Integer=Threads.nthreads(), - run_loggers=true) where {D, G, T} +@inline function simulate!(sys::System{D, G, T}, + sim::MetropolisMonteCarlo, + n_steps::Integer; + n_threads::Integer=Threads.nthreads(), + run_loggers=true) where {D, G, T} neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads) E_old = potential_energy(sys, neighbors; n_threads=n_threads) coords_old = similar(sys.coords) From 379e97cc35503f17b5b037459748aaed17130703 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 26 Jul 2024 15:01:21 +0100 Subject: [PATCH 47/74] Enzyme loggers fix --- src/loggers.jl | 2 -- src/setup.jl | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/src/loggers.jl b/src/loggers.jl index b74630904..acaa4c5d5 100644 --- a/src/loggers.jl +++ b/src/loggers.jl @@ -38,8 +38,6 @@ function apply_loggers!(sys::System, neighbors=nothing, step_n::Integer=0, run_l for logger in values(sys.loggers) log_property!(logger, sys, neighbors, step_n; n_threads=n_threads, kwargs...) end - elseif run_loggers != false && run_loggers != :skipzero - throw(ArgumentError("run_loggers must be true, false or :skipzero")) end return sys end diff --git a/src/setup.jl b/src/setup.jl index a33919016..7fecb6e36 100644 --- a/src/setup.jl +++ b/src/setup.jl @@ -1391,7 +1391,7 @@ function add_position_restraints(sys, return System( atoms=deepcopy(sys.atoms), coords=copy(sys.coords), - boundary=copy(sys.boundary), + boundary=deepcopy(sys.boundary), velocities=copy(sys.velocities), atoms_data=deepcopy(sys.atoms_data), topology=deepcopy(sys.topology), From 7ff6cdcbc16ea2ab3c6d99e43a64f27cc57d4089 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 26 Jul 2024 15:01:31 +0100 Subject: [PATCH 48/74] Enzyme compat --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index b9521e42e..eac5420e6 100644 --- a/Project.toml +++ b/Project.toml @@ -59,7 +59,7 @@ Combinatorics = "1" DataStructures = "0.18" Distances = "0.10" Distributions = "0.23, 0.24, 0.25" -Enzyme = "0.11.15, 0.12" +Enzyme = "0.12" EzXML = "1" FLoops = "0.2" GLMakie = "0.8, 0.9, 0.10" From bc813b37d3259c9051081f72b397089362922d40 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 26 Jul 2024 15:11:20 +0100 Subject: [PATCH 49/74] test simulation gradients on CI --- src/interactions/implicit_solvent.jl | 2 +- test/gradients.jl | 13 ++++++------- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/src/interactions/implicit_solvent.jl b/src/interactions/implicit_solvent.jl index 1f70a4af0..2ad224268 100644 --- a/src/interactions/implicit_solvent.jl +++ b/src/interactions/implicit_solvent.jl @@ -601,7 +601,7 @@ function inject_interaction(inter::ImplicitSolventGBN2, params_dic, sys) key_prefix = "inter_GB_" bond_index = findfirst(sil -> eltype(sil.inters) <: HarmonicBond, sys.specific_inter_lists) - element_to_radius = Dict{String, DefaultFloat}() # Units here made the gradients vanish + element_to_radius = Dict{String, DefaultFloat}() for k in keys(mbondi2_element_to_radius) element_to_radius[k] = dict_get(params_dic, key_prefix * "radius_" * k, ustrip(mbondi2_element_to_radius[k])) diff --git a/test/gradients.jl b/test/gradients.jl index 9b04fa21b..25e3322c4 100644 --- a/test/gradients.jl +++ b/test/gradients.jl @@ -407,13 +407,11 @@ end if run_gpu_tests push!(platform_runs, ("GPU", true, false)) end - test_runs = [ + test_runs = ( ("Energy", test_energy_grad, 1e-8), - ("Force" , test_forces_grad , 1e-8), - ] - if !running_CI - #push!(test_runs, ("Sim", test_sim_grad, 0.015)) - end + ("Force" , test_forces_grad, 1e-8), + ("Sim" , test_sim_grad , 1e-2), + ) params_to_test = ( #"inter_LJ_weight_14", "atom_N_ϵ", @@ -442,7 +440,8 @@ end test_fn(dic, sys_ref, copy(sys_ref.coords), sys_ref.neighbor_finder, n_threads) end frac_diff = abs(genz - gfd) / abs(gfd) - @info "$(rpad(test_name, 6)) - $(rpad(platform, 12)) - $(rpad(param, 21)) - FD $gfd, Enzyme $genz, fractional difference $frac_diff" + @info "$(rpad(test_name, 6)) - $(rpad(platform, 12)) - $(rpad(param, 21)) - " * + "FD $gfd, Enzyme $genz, fractional difference $frac_diff" @test frac_diff < test_tol end end From 68d51eac63ea02b9f832fbc32a62cbdbeb576c30 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 26 Jul 2024 16:24:00 +0100 Subject: [PATCH 50/74] remove unused gradient dependencies --- Project.toml | 6 ------ docs/src/differentiable.md | 25 +++++++++---------------- src/Molly.jl | 3 --- src/analysis.jl | 2 -- src/coupling.jl | 5 ----- src/simulators.jl | 7 ------- src/spatial.jl | 8 -------- 7 files changed, 9 insertions(+), 47 deletions(-) diff --git a/Project.toml b/Project.toml index eac5420e6..6a7f1d79f 100644 --- a/Project.toml +++ b/Project.toml @@ -10,8 +10,6 @@ AtomsCalculators = "a3e0e189-c65a-42c1-833c-339540406eb1" BioStructures = "de9282ab-8554-53be-b2d6-f6c222edabfc" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CellListMap = "69e1c6dd-3888-40e6-b3c8-31ac5f578864" -ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2" -ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" Chemfiles = "46823bd8-5fb3-5f92-9aa0-96921f3dd015" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" @@ -32,7 +30,6 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a" -UnitfulChainRules = "f31437dd-25a7-4345-875f-756556e6935d" UnsafeAtomicsLLVM = "d80eeb9a-aca5-4d75-85e5-170c8b632249" [weakdeps] @@ -51,8 +48,6 @@ AtomsCalculators = "0.2" BioStructures = "4" CUDA = "4.2, 5" CellListMap = "0.8.11, 0.9" -ChainRules = "1.44" -ChainRulesCore = "1" Chemfiles = "0.10.3" Colors = "0.11, 0.12" Combinatorics = "1" @@ -77,6 +72,5 @@ StaticArrays = "1.8.2" Statistics = "1.9" Unitful = "1" UnitfulAtomic = "1" -UnitfulChainRules = "0.1.2" UnsafeAtomicsLLVM = "0.1, 0.2" julia = "1.9" diff --git a/docs/src/differentiable.md b/docs/src/differentiable.md index bed2d2611..ba8e88a41 100644 --- a/docs/src/differentiable.md +++ b/docs/src/differentiable.md @@ -7,8 +7,8 @@ In the last few years, the deep learning revolution has broadened to include the The concept of using automatic differentiation (AD) to obtain exact gradients through physical simulations has many interesting applications, including parameterising force fields and training neural networks to describe atomic potentials. There are some projects that explore differentiable molecular simulations - see [Related software](@ref). -However Julia provides a strong suite of AD tools, with [Zygote.jl](https://github.com/FluxML/Zygote.jl) and [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) allowing source-to-source transformations for much of the language. -With Molly you can use the power of Zygote and Enzyme to obtain gradients through molecular simulations, even in the presence of complex interactions such as implicit solvation and stochasticity such as Langevin dynamics or the Andersen thermostat. +However Julia provides a strong suite of AD tools, with [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) allowing source-to-source transformations for much of the language. +With Molly you can use the power of Enzyme to obtain gradients through molecular simulations, even in the presence of complex interactions such as implicit solvation and stochasticity such as Langevin dynamics or the Andersen thermostat. Reverse mode AD can be used on the CPU with multithreading and on the GPU; performance is typically within an order of magnitude of the primal run. Forward mode AD can also be used on the CPU. Pairwise, specific and general interactions work, along with neighbor lists, and the same abstractions for running simulations are used as in the main package. @@ -44,7 +44,6 @@ end Now we can set up and run the simulation in a similar way to that described in the [Molly documentation](@ref). The difference is that we wrap the simulation in a `loss` function. This returns a single value that we want to obtain gradients with respect to, in this case the difference between the value of the above function at the end of the simulation and a target distance. -The `Zygote.ignore()` block allows us to ignore code for the purposes of obtaining gradients; you could add the [`visualize`](@ref) function there for example. ```julia using Zygote using Format @@ -96,10 +95,8 @@ function loss(σ, coords, velocities) mms_end = mean_min_separation(Array(sys.coords), boundary) loss_val = abs(mms_end - dist_true) - Zygote.ignore() do - printfmt("σ {:6.3f} | Mean min sep expected {:6.3f} | Mean min sep end {:6.3f} | Loss {:6.3f} | ", - σ, σ * (2 ^ (1 / 6)), mms_end, loss_val) - end + printfmt("σ {:6.3f} | Mean min sep expected {:6.3f} | Mean min sep end {:6.3f} | Loss {:6.3f} | ", + σ, σ * (2 ^ (1 / 6)), mms_end, loss_val) return loss_val end @@ -218,10 +215,8 @@ function loss(θ) dist_end = 0.5 * (d1 + d2) loss_val = abs(dist_end - dist_true) - Zygote.ignore() do - printfmt("θ {:5.1f}° | Final dist {:4.2f} | Loss {:5.3f} | ", - rad2deg(θ), dist_end, loss_val) - end + printfmt("θ {:5.1f}° | Final dist {:4.2f} | Loss {:5.3f} | ", + rad2deg(θ), dist_end, loss_val) return loss_val end @@ -276,7 +271,7 @@ The plot of these shows that the gradient has the expected sign either side of t ## Neural network potentials -Since gradients can be computed with Zygote, [Flux](https://fluxml.ai) models can also be incorporated into simulations. +[Flux](https://fluxml.ai) models can also be incorporated into simulations. Here we show a neural network in the force function, though they can also be used in other parts of the simulation. This example also shows how gradients for multiple parameters can be obtained, in this case the parameters of the neural network. The jump from single to multiple parameters is important because single parameters can be optimised using finite differencing, whereas differentiable simulation is well-placed to optimise many parameters simultaneously. @@ -348,10 +343,8 @@ function loss() norm(vector(sys.coords[3], sys.coords[1], boundary))) / 3 loss_val = abs(dist_end - dist_true) - Zygote.ignore() do - printfmt("Dist end {:6.3f} | Loss {:6.3f}\n", dist_end, loss_val) - visualize(sys.loggers.coords, boundary, "sim.mp4"; show_boundary=false) - end + printfmt("Dist end {:6.3f} | Loss {:6.3f}\n", dist_end, loss_val) + visualize(sys.loggers.coords, boundary, "sim.mp4"; show_boundary=false) return loss_val end diff --git a/src/Molly.jl b/src/Molly.jl index 49e893b7c..4c3b22c23 100644 --- a/src/Molly.jl +++ b/src/Molly.jl @@ -9,8 +9,6 @@ import AtomsBase import AtomsCalculators import BioStructures # Imported to avoid clashing names using CellListMap -using ChainRules -using ChainRulesCore import Chemfiles using Combinatorics using CUDA @@ -26,7 +24,6 @@ using PeriodicTable using SimpleCrystals using Unitful using UnitfulAtomic -using UnitfulChainRules using UnsafeAtomicsLLVM using LinearAlgebra diff --git a/src/analysis.jl b/src/analysis.jl index a63d9a07c..3b67b8f45 100644 --- a/src/analysis.jl +++ b/src/analysis.jl @@ -62,7 +62,6 @@ for the periodic boundary conditions. function displacements(coords, boundary) n_atoms = length(coords) coords_rep = repeat(reshape(coords, n_atoms, 1), 1, n_atoms) - # Makes gradient work with Zygote broadcasting additions vec_2_arg(c1, c2) = vector(c1, c2, boundary) diffs = vec_2_arg.(coords_rep, permutedims(coords_rep, (2, 1))) return diffs @@ -163,7 +162,6 @@ function hydrodynamic_radius(coords::AbstractArray{SVector{D, T}}, boundary) whe n_atoms = length(coords) diag_cpu = Diagonal(ones(T, n_atoms)) diag = isa(coords, CuArray) ? CuArray(diag_cpu) : diag_cpu - # Other approaches to removing the diagonal Inf didn't work with Zygote dists = distances(coords, boundary) .+ diag sum_inv_dists = sum(inv.(dists)) - sum(inv(diag)) inv_R_hyd = sum_inv_dists / (2 * n_atoms^2) diff --git a/src/coupling.jl b/src/coupling.jl index 24962b195..c975f10cc 100644 --- a/src/coupling.jl +++ b/src/coupling.jl @@ -162,8 +162,6 @@ It should be used alongside a temperature coupling method such as the [`Langevin simulator or [`AndersenThermostat`](@ref) coupling. The neighbor list is not updated when making trial moves or after accepted moves. Note that the barostat can change the bounding box of the system. - -Not currently compatible with automatic differentiation using Zygote. """ mutable struct MonteCarloBarostat{T, P, K, V} pressure::P @@ -280,8 +278,6 @@ It should be used alongside a temperature coupling method such as the [`Langevin simulator or [`AndersenThermostat`](@ref) coupling. The neighbor list is not updated when making trial moves or after accepted moves. Note that the barostat can change the bounding box of the system. - -Not currently compatible with automatic differentiation using Zygote. """ mutable struct MonteCarloAnisotropicBarostat{D, T, P, K, V} pressure::SVector{D, P} @@ -439,7 +435,6 @@ The neighbor list is not updated when making trial moves or after accepted moves Note that the barostat can change the bounding box of the system. This barostat is only available for 3D systems. -Not currently compatible with automatic differentiation using Zygote. """ mutable struct MonteCarloMembraneBarostat{T, P, K, V, S} pressure::SVector{3, P} diff --git a/src/simulators.jl b/src/simulators.jl index a421fc332..df6743c74 100644 --- a/src/simulators.jl +++ b/src/simulators.jl @@ -23,8 +23,6 @@ export Steepest descent energy minimization. -Not currently compatible with automatic differentiation using Zygote. - # Arguments - `step_size::D=0.01u"nm"`: the initial maximum displacement. - `max_steps::Int=1000`: the maximum number of steps. @@ -454,7 +452,6 @@ For more information on the sampling properties of splitting schemes, see Not currently compatible with constraints, will print a warning and continue without applying constraints. -Not currently compatible with automatic differentiation using Zygote. # Arguments - `dt::S`: the time step of the simulation. @@ -754,8 +751,6 @@ the number of temperatures in the simulator. When calling [`simulate!`](@ref), the `assign_velocities` keyword argument determines whether to assign random velocities at the appropriate temperature for each replica. -Not currently compatible with automatic differentiation using Zygote. - # Arguments - `dt::DT`: the time step of the simulation. - `temperatures::TP`: the temperatures corresponding to the replicas. @@ -862,8 +857,6 @@ The replicas are expected to have different Hamiltonians, i.e. different interac When calling [`simulate!`](@ref), the `assign_velocities` keyword argument determines whether to assign random velocities at the appropriate temperature for each replica. -Not currently compatible with automatic differentiation using Zygote. - # Arguments - `dt::DT`: the time step of the simulation. - `temperature::T`: the temperatures of the simulation. diff --git a/src/spatial.jl b/src/spatial.jl index d3a019e2d..52a78325d 100644 --- a/src/spatial.jl +++ b/src/spatial.jl @@ -113,7 +113,6 @@ image is found, which is slower. Not currently able to simulate a cubic box, use [`CubicBoundary`](@ref) or small offsets instead. Not currently compatible with infinite boundaries. -Not currently compatible with automatic differentiation using Zygote. """ struct TriclinicBoundary{T, A, D, I} basis_vectors::SVector{3, SVector{3, D}} @@ -677,7 +676,6 @@ function torsion_angle(vec_ij, vec_jk, vec_kl) return θ end -# Used to write an rrule that can override the Zygote sum adjoint sum_svec(arr) = sum(arr) """ @@ -708,8 +706,6 @@ Custom general interaction types can implement this function. This should only be used on systems containing just pairwise interactions, or where the specific interactions, constraints and general interactions without [`virial`](@ref) defined do not contribute to the virial. -Not currently compatible with automatic differentiation using Zygote when -using pairwise interactions. """ function virial(sys; n_threads::Integer=Threads.nthreads()) return virial(sys, find_neighbors(sys; n_threads=n_threads); n_threads=n_threads) @@ -802,8 +798,6 @@ This should only be used on systems containing just pairwise interactions, or where the specific interactions, constraints and general interactions without [`virial`](@ref) defined do not contribute to the virial. Not compatible with infinite boundaries. -Not currently compatible with automatic differentiation using Zygote when -using pairwise interactions. """ function pressure(sys; n_threads::Integer=Threads.nthreads()) return pressure(sys, find_neighbors(sys; n_threads=n_threads); n_threads=n_threads) @@ -837,7 +831,6 @@ Accounts for periodic boundary conditions by using the circular mean. If `topology=nothing` then the coordinates are returned. Not currently compatible with [`TriclinicBoundary`](@ref) if the topology is set. -Not currently compatible with automatic differentiation using Zygote. """ function molecule_centers(coords::AbstractArray{SVector{D, C}}, boundary, topology) where {D, C} if isnothing(topology) @@ -886,7 +879,6 @@ moved by the same amount according to the center of coordinates of the molecule. This can be disabled with `ignore_molecules=true`. Not currently compatible with [`TriclinicBoundary`](@ref) if the topology is set. -Not currently compatible with automatic differentiation using Zygote. """ function scale_coords!(sys, scale_factor; ignore_molecules=false) if ignore_molecules || isnothing(sys.topology) From 3951042c9c89142a3b2037a5e97624317c13ec34 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 26 Jul 2024 18:08:14 +0100 Subject: [PATCH 51/74] OBC Born radii don't broadcast --- src/interactions/implicit_solvent.jl | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/src/interactions/implicit_solvent.jl b/src/interactions/implicit_solvent.jl index 2ad224268..3c023c03b 100644 --- a/src/interactions/implicit_solvent.jl +++ b/src/interactions/implicit_solvent.jl @@ -324,7 +324,7 @@ function lookup_table(full_table::AbstractArray{T}, radii) where T return table end -function lookup_table(full_table::AbstractArray{T}, radii::AbstractArray{<:AbstractFloat}) where T +function lookup_table(full_table::AbstractArray, radii::AbstractArray{<:AbstractFloat}) return lookup_table(full_table, radii * u"nm") end @@ -676,6 +676,25 @@ with respect to atomic distance. Custom GBSA methods should implement this function. """ function born_radii_and_grad(inter::ImplicitSolventOBC, coords, boundary) + Is = fill(zero(eltype(eltype(coords))) / unit(inter.dist_cutoff)^2, length(coords)) + @inbounds for i in eachindex(coords) + I = zero(eltype(Is)) + for j in eachindex(coords) + I += born_radii_loop_OBC(coords[i], coords[j], inter.oris[i], + inter.srjs[j], inter.dist_cutoff, boundary) + end + Is[i] = I + end + I_grads = zeros(eltype(Is), length(Is), length(Is)) ./ unit(inter.dist_cutoff) + + Bs_B_grads = born_radii_sum.(inter.offset_radii, inter.offset, Is, + inter.α, inter.β, inter.γ) + Bs = get_i1.(Bs_B_grads) + B_grads = get_i2.(Bs_B_grads) + return Bs, B_grads, I_grads +end + +function born_radii_and_grad(inter::ImplicitSolventOBC, coords::CuArray, boundary) coords_i = @view coords[inter.is] coords_j = @view coords[inter.js] loop_res = born_radii_loop_OBC.(coords_i, coords_j, inter.oris, inter.srjs, From ced1e7d4bfb8307c105a789afd1eacb7aa846457 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 26 Jul 2024 18:16:55 +0100 Subject: [PATCH 52/74] GBSA energy don't broadcast --- src/interactions/implicit_solvent.jl | 26 +++++++++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/src/interactions/implicit_solvent.jl b/src/interactions/implicit_solvent.jl index 3c023c03b..10ad6eb3a 100644 --- a/src/interactions/implicit_solvent.jl +++ b/src/interactions/implicit_solvent.jl @@ -675,8 +675,8 @@ with respect to atomic distance. Custom GBSA methods should implement this function. """ -function born_radii_and_grad(inter::ImplicitSolventOBC, coords, boundary) - Is = fill(zero(eltype(eltype(coords))) / unit(inter.dist_cutoff)^2, length(coords)) +function born_radii_and_grad(inter::ImplicitSolventOBC{T}, coords, boundary) where T + Is = fill(zero(T) / unit(inter.dist_cutoff)^2, length(coords)) @inbounds for i in eachindex(coords) I = zero(eltype(Is)) for j in eachindex(coords) @@ -1171,7 +1171,27 @@ function gb_energy_loop(coord_i, coord_j, i, j, charge_i, charge_j, Bi, Bj, ori, end end -function AtomsCalculators.potential_energy(sys, inter::AbstractGBSA; kwargs...) +function AtomsCalculators.potential_energy(sys::System{<:Any, false, T}, inter::AbstractGBSA; + kwargs...) where T + coords, boundary = sys.coords, sys.boundary + Bs, B_grads, I_grads = born_radii_and_grad(inter, coords, boundary) + atom_charges = charge.(sys.atoms) + + E = zero(T) * sys.energy_units + @inbounds for i in eachindex(sys) + for j in eachindex(sys) + E += gb_energy_loop( + coords[i], coords[j], i, j, atom_charges[i], atom_charges[j], Bs[i], Bs[j], + inter.oris[i], inter.dist_cutoff, inter.factor_solute, inter.factor_solvent, + inter.kappa, inter.offset, inter.probe_radius, inter.sa_factor, + inter.use_ACE, boundary, + ) + end + end + return E +end + +function AtomsCalculators.potential_energy(sys::System{<:Any, true}, inter::AbstractGBSA; kwargs...) coords, atoms, boundary = sys.coords, sys.atoms, sys.boundary Bs, B_grads, I_grads = born_radii_and_grad(inter, coords, boundary) From 37583faa39375bae336196cc7db39f5005ab863f Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 26 Jul 2024 18:39:42 +0100 Subject: [PATCH 53/74] GBn2 force don't broadcast 1 --- src/interactions/implicit_solvent.jl | 33 +++++++++++++--------------- 1 file changed, 15 insertions(+), 18 deletions(-) diff --git a/src/interactions/implicit_solvent.jl b/src/interactions/implicit_solvent.jl index 10ad6eb3a..230d0ac8a 100644 --- a/src/interactions/implicit_solvent.jl +++ b/src/interactions/implicit_solvent.jl @@ -676,7 +676,7 @@ with respect to atomic distance. Custom GBSA methods should implement this function. """ function born_radii_and_grad(inter::ImplicitSolventOBC{T}, coords, boundary) where T - Is = fill(zero(T) / unit(inter.dist_cutoff)^2, length(coords)) + Is = fill(zero(T) / unit(inter.dist_cutoff), length(coords)) @inbounds for i in eachindex(coords) I = zero(eltype(Is)) for j in eachindex(coords) @@ -868,13 +868,8 @@ end get_bi(r::ForceLoopResult1) = r.bi get_bj(r::ForceLoopResult1) = r.bj -struct ForceLoopResult2{V} - fi::V - fj::V -end - -get_fi(r::Union{ForceLoopResult1, ForceLoopResult2}) = r.fi -get_fj(r::Union{ForceLoopResult1, ForceLoopResult2}) = r.fj +get_fi(r::ForceLoopResult1) = r.fi +get_fj(r::ForceLoopResult1) = r.fj function gb_force_loop_1(coord_i, coord_j, i, j, charge_i, charge_j, Bi, Bj, dist_cutoff, factor_solute, factor_solvent, kappa, boundary) @@ -920,8 +915,7 @@ function gb_force_loop_2(coord_i, coord_j, bi, ig, ori, srj, dist_cutoff, bounda dr = vector(coord_i, coord_j, boundary) r = norm(dr) if iszero_value(r) || (!iszero_value(dist_cutoff) && r > dist_cutoff) - zero_force = zero(bi ./ coord_i .^ 2) - return ForceLoopResult2(zero_force, zero_force) + return zero(bi ./ coord_i .^ 2) end rsrj = r + srj if ori < rsrj @@ -933,10 +927,9 @@ function gb_force_loop_2(coord_i, coord_j, bi, ig, ori, srj, dist_cutoff, bounda t3 = (1 + (srj^2)*r2inv)*(L^2 - U^2)/8 + log(U/L)*r2inv/4 de = bi * (t3 - ig) * rinv fdr = dr * de - return ForceLoopResult2(-fdr, fdr) + return fdr else - zero_force = zero(bi ./ coord_i .^ 2) - return ForceLoopResult2(zero_force, zero_force) + return zero(bi ./ coord_i .^ 2) end end @@ -958,12 +951,16 @@ function forces_gbsa(sys, inter, Bs, B_grads, I_grads, born_forces, atom_charges born_forces_2 = born_forces_1 .* (Bs .^ 2) .* B_grads - bis = @view born_forces_2[inter.is] - loop_res_2 = gb_force_loop_2.(coords_i, coords_j, bis, I_grads, inter.oris, inter.srjs, - inter.dist_cutoff, (boundary,)) + @inbounds for i in eachindex(sys) + for j in eachindex(sys) + f = gb_force_loop_2(coords[i], coords[j], born_forces_2[i], I_grads[i, j], + inter.oris[i], inter.srjs[j], inter.dist_cutoff, boundary) + fs[i] = fs[i] .- f + fs[j] = fs[j] .+ f + end + end - return fs .+ dropdims(sum(get_fi.(loop_res_2); dims=2); dims=2) .+ - dropdims(sum(get_fj.(loop_res_2); dims=1); dims=1) + return fs end function forces_gbsa(sys::System{D, true, T}, inter, Bs, B_grads, I_grads, born_forces, From bd3d9c7e0ec5b28ae57be99b1b8aeec39a0572fe Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 26 Jul 2024 18:56:35 +0100 Subject: [PATCH 54/74] GBn2 force don't broadcast 2 --- src/interactions/implicit_solvent.jl | 62 +++++++++++----------------- 1 file changed, 24 insertions(+), 38 deletions(-) diff --git a/src/interactions/implicit_solvent.jl b/src/interactions/implicit_solvent.jl index 230d0ac8a..2c7ccfe94 100644 --- a/src/interactions/implicit_solvent.jl +++ b/src/interactions/implicit_solvent.jl @@ -857,31 +857,17 @@ function gbsa_born_kernel!(Is, I_grads, coords_var, offset_radii_var, scaled_off return nothing end -# Store the results of the ij broadcasts during force calculation -struct ForceLoopResult1{T, V} - bi::T - bj::T - fi::V - fj::V -end - -get_bi(r::ForceLoopResult1) = r.bi -get_bj(r::ForceLoopResult1) = r.bj - -get_fi(r::ForceLoopResult1) = r.fi -get_fj(r::ForceLoopResult1) = r.fj - function gb_force_loop_1(coord_i, coord_j, i, j, charge_i, charge_j, Bi, Bj, dist_cutoff, factor_solute, factor_solvent, kappa, boundary) if j < i zero_force = zero(factor_solute ./ coord_i .^ 2) - return ForceLoopResult1(zero_force[1], zero_force[1], zero_force, zero_force) + return zero_force[1], zero_force[1], zero_force, zero_force end dr = vector(coord_i, coord_j, boundary) r2 = sum(abs2, dr) if !iszero_value(dist_cutoff) && r2 > dist_cutoff^2 zero_force = zero(factor_solute ./ coord_i .^ 2) - return ForceLoopResult1(zero_force[1], zero_force[1], zero_force, zero_force) + return zero_force[1], zero_force[1], zero_force, zero_force end alpha2_ij = Bi * Bj D = r2 / (4 * alpha2_ij) @@ -903,11 +889,10 @@ function gb_force_loop_1(coord_i, coord_j, i, j, charge_i, charge_j, Bi, Bj, dis fdr = dr * dGpol_dr change_fs_i = fdr change_fs_j = -fdr - return ForceLoopResult1(change_born_force_i, change_born_force_j, - change_fs_i, change_fs_j) + return change_born_force_i, change_born_force_j, change_fs_i, change_fs_j else zero_force = zero(factor_solute ./ coord_i .^ 2) - return ForceLoopResult1(change_born_force_i, zero_force[1], zero_force, zero_force) + return change_born_force_i, zero_force[1], zero_force, zero_force end end @@ -935,22 +920,22 @@ end function forces_gbsa(sys, inter, Bs, B_grads, I_grads, born_forces, atom_charges) coords, boundary = sys.coords, sys.boundary - coords_i = @view coords[inter.is] - coords_j = @view coords[inter.js] - charges_i = @view atom_charges[inter.is] - charges_j = @view atom_charges[inter.js] - Bsi = @view Bs[inter.is] - Bsj = @view Bs[inter.js] - loop_res_1 = gb_force_loop_1.(coords_i, coords_j, inter.is, inter.js, charges_i, - charges_j, Bsi, Bsj, inter.dist_cutoff, inter.factor_solute, - inter.factor_solvent, inter.kappa, (boundary,)) - born_forces_1 = born_forces .+ dropdims(sum(get_bi.(loop_res_1); dims=2); dims=2) .+ - dropdims(sum(get_bj.(loop_res_1); dims=1); dims=1) - fs = dropdims(sum(get_fi.(loop_res_1); dims=2); dims=2) .+ - dropdims(sum(get_fj.(loop_res_1); dims=1); dims=1) + born_forces_1 = copy(born_forces) + fs = ustrip_vec.(zero(coords)) * sys.force_units + @inbounds for i in eachindex(sys) + for j in eachindex(sys) + bi, bj, fi, fj = gb_force_loop_1( + coords[i], coords[j], i, j, atom_charges[i], atom_charges[j], Bs[i], Bs[j], + inter.dist_cutoff, inter.factor_solute, inter.factor_solvent, inter.kappa, boundary, + ) + born_forces_1[i] += bi + born_forces_1[j] += bj + fs[i] = fs[i] .+ fi + fs[j] = fs[j] .+ fj + end + end born_forces_2 = born_forces_1 .* (Bs .^ 2) .* B_grads - @inbounds for i in eachindex(sys) for j in eachindex(sys) f = gb_force_loop_2(coords[i], coords[j], born_forces_2[i], I_grads[i, j], @@ -1199,9 +1184,10 @@ function AtomsCalculators.potential_energy(sys::System{<:Any, true}, inter::Abst charges_j = @view atom_charges[inter.js] Bsi = @view Bs[inter.is] Bsj = @view Bs[inter.js] - return sum(gb_energy_loop.(coords_i, coords_j, inter.is, inter.js, charges_i, - charges_j, Bsi, Bsj, inter.oris, inter.dist_cutoff, - inter.factor_solute, inter.factor_solvent, inter.kappa, - inter.offset, inter.probe_radius, inter.sa_factor, inter.use_ACE, - (boundary,))) + return sum(gb_energy_loop.( + coords_i, coords_j, inter.is, inter.js, charges_i, charges_j, Bsi, Bsj, + inter.oris, inter.dist_cutoff, inter.factor_solute, inter.factor_solvent, + inter.kappa, inter.offset, inter.probe_radius, inter.sa_factor, inter.use_ACE, + (boundary,), + )) end From 22c16e5eeb32aaf28ea99c83e3932ed84ea29f3b Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Wed, 2 Oct 2024 16:04:59 +0100 Subject: [PATCH 55/74] fix Colors weakdep --- Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/Project.toml b/Project.toml index 6a7f1d79f..b8784c16d 100644 --- a/Project.toml +++ b/Project.toml @@ -33,6 +33,7 @@ UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a" UnsafeAtomicsLLVM = "d80eeb9a-aca5-4d75-85e5-170c8b632249" [weakdeps] +Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" From 3b72ecf80184ce03a4dfc2e30f36f4084650ae7f Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Wed, 2 Oct 2024 16:09:50 +0100 Subject: [PATCH 56/74] fix Enzyme weakdep --- Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/Project.toml b/Project.toml index b8784c16d..860cc1857 100644 --- a/Project.toml +++ b/Project.toml @@ -34,6 +34,7 @@ UnsafeAtomicsLLVM = "d80eeb9a-aca5-4d75-85e5-170c8b632249" [weakdeps] Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" From c3325da4cdd649c3ec39c3c775e5307236177593 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 4 Oct 2024 17:44:55 +0100 Subject: [PATCH 57/74] rename box_volume to volume --- docs/src/documentation.md | 2 +- ext/MollyGLMakieExt.jl | 4 ++-- src/analysis.jl | 2 +- src/coupling.jl | 12 ++++++------ src/setup.jl | 4 ++-- src/spatial.jl | 10 +++++----- test/basic.jl | 12 ++++++------ test/simulation.jl | 18 +++++++++--------- 8 files changed, 32 insertions(+), 32 deletions(-) diff --git a/docs/src/documentation.md b/docs/src/documentation.md index 7e6a6648e..133dab014 100644 --- a/docs/src/documentation.md +++ b/docs/src/documentation.md @@ -877,7 +877,7 @@ b = TriclinicBoundary( ) # Volume of bounding box -box_volume(b) # 3.8993746318188633 nm^3 +volume(b) # 3.8993746318188633 nm^3 # Random coordinate uniformly distributed within boundary random_coord(b) # SVector(2.651062310435411, 2.1702306804433973, 0.9518105403051831)u"nm" diff --git a/ext/MollyGLMakieExt.jl b/ext/MollyGLMakieExt.jl index fce5ce72f..005ffd8a0 100644 --- a/ext/MollyGLMakieExt.jl +++ b/ext/MollyGLMakieExt.jl @@ -35,12 +35,12 @@ function Molly.visualize(coord_logger, if dims == 3 PointType = Point3f ax = Axis3(fig[1, 1], aspect=:data) - max_connection_dist = cbrt(box_volume(boundary)) / 2 + max_connection_dist = cbrt(volume(boundary)) / 2 elseif dims == 2 PointType = Point2f ax = Axis(fig[1, 1]) ax.aspect = DataAspect() - max_connection_dist = sqrt(box_volume(boundary)) / 2 + max_connection_dist = sqrt(volume(boundary)) / 2 else throw(ArgumentError("found $dims dimensions but can only visualize 2 or 3 dimensions")) end diff --git a/src/analysis.jl b/src/analysis.jl index 3b67b8f45..ce32efc91 100644 --- a/src/analysis.jl +++ b/src/analysis.jl @@ -91,7 +91,7 @@ function rdf(coords, boundary; npoints::Integer=200) dists_vec = [dists[i, j] for i in 1:n_atoms, j in 1:n_atoms if j > i] dist_unit = unit(first(dists_vec)) kd = kde(ustrip.(dists_vec); npoints=npoints) - ρ = n_atoms / box_volume(boundary) + ρ = n_atoms / volume(boundary) T = float_type(boundary) if dims == 3 normalizing_factor = 4 .* T(π) .* ρ .* step(kd.x) .* kd.x .^ 2 .* dist_unit .^ 3 diff --git a/src/coupling.jl b/src/coupling.jl index c975f10cc..15e8393de 100644 --- a/src/coupling.jl +++ b/src/coupling.jl @@ -178,7 +178,7 @@ end function MonteCarloBarostat(P, T, boundary; n_steps=30, n_iterations=1, scale_factor=0.01, scale_increment=1.1, max_volume_frac=0.3, trial_find_neighbors=false) - volume_scale = box_volume(boundary) * float_type(boundary)(scale_factor) + volume_scale = volume(boundary) * float_type(boundary)(scale_factor) return MonteCarloBarostat(P, T, n_steps, n_iterations, volume_scale, scale_increment, max_volume_frac, trial_find_neighbors, 0, 0) end @@ -196,7 +196,7 @@ function apply_coupling!(sys::System{D, G, T}, barostat::MonteCarloBarostat, sim for attempt_n in 1:barostat.n_iterations E = potential_energy(sys, neighbors, step_n; n_threads=n_threads) - V = box_volume(sys.boundary) + V = volume(sys.boundary) dV = barostat.volume_scale * (2 * rand(T) - 1) v_scale = (V + dV) / V l_scale = (D == 2 ? sqrt(v_scale) : cbrt(v_scale)) @@ -301,7 +301,7 @@ function MonteCarloAnisotropicBarostat(pressure::SVector{D}, scale_increment=1.1, max_volume_frac=0.3, trial_find_neighbors=false) where D - volume_scale_factor = box_volume(boundary) * float_type(boundary)(scale_factor) + volume_scale_factor = volume(boundary) * float_type(boundary)(scale_factor) volume_scale = fill(volume_scale_factor, D) if AtomsBase.n_dimensions(boundary) != D throw(ArgumentError("pressure vector length ($(D)) must match boundary " * @@ -348,7 +348,7 @@ function apply_coupling!(sys::System{D, G, T}, mask2[axis] = false E = potential_energy(sys, neighbors, step_n; n_threads=n_threads) - V = box_volume(sys.boundary) + V = volume(sys.boundary) dV = barostat.volume_scale[axis] * (2 * rand(T) - 1) v_scale = (V + dV) / V l_scale = SVector{D}(mask1 * v_scale + mask2) @@ -465,7 +465,7 @@ function MonteCarloMembraneBarostat(pressure, xy_isotropy=false, z_axis_fixed=false, constant_volume=false) - volume_scale_factor = box_volume(boundary) * float_type(boundary)(scale_factor) + volume_scale_factor = volume(boundary) * float_type(boundary)(scale_factor) volume_scale = fill(volume_scale_factor, 3) if AtomsBase.n_dimensions(boundary) != 3 @@ -520,7 +520,7 @@ function apply_coupling!(sys::System{D, G, T}, end E = potential_energy(sys, neighbors, step_n; n_threads=n_threads) - V = box_volume(sys.boundary) + V = volume(sys.boundary) dV = barostat.volume_scale[axis] * (2 * rand(T) - 1) v_scale = (V + dV) / V l_scale = SVector{D, T}(one(T), one(T), one(T)) diff --git a/src/setup.jl b/src/setup.jl index 7fecb6e36..ce32377a1 100644 --- a/src/setup.jl +++ b/src/setup.jl @@ -33,7 +33,7 @@ function place_atoms(n_atoms::Integer, throw(ArgumentError("one or more dimension has infinite boundaries, boundary is $boundary")) end dims = AtomsBase.n_dimensions(boundary) - max_atoms = box_volume(boundary) / (min_dist ^ dims) + max_atoms = volume(boundary) / (min_dist ^ dims) if n_atoms > max_atoms throw(ArgumentError("boundary $boundary too small for $n_atoms atoms with minimum distance $min_dist")) end @@ -89,7 +89,7 @@ function place_diatomics(n_molecules::Integer, throw(ArgumentError("one or more dimension has infinite boundaries, boundary is $boundary")) end dims = AtomsBase.n_dimensions(boundary) - max_molecules = box_volume(boundary) / ((min_dist + bond_length) ^ dims) + max_molecules = volume(boundary) / ((min_dist + bond_length) ^ dims) if n_molecules > max_molecules throw(ArgumentError("boundary $boundary too small for $n_molecules diatomics with minimum distance $min_dist")) end diff --git a/src/spatial.jl b/src/spatial.jl index 52a78325d..9ef98f233 100644 --- a/src/spatial.jl +++ b/src/spatial.jl @@ -4,7 +4,7 @@ export CubicBoundary, RectangularBoundary, TriclinicBoundary, - box_volume, + volume, box_center, scale_boundary, random_coord, @@ -244,12 +244,12 @@ n_infinite_dims(b::TriclinicBoundary) = 0 n_infinite_dims(sys::System) = n_infinite_dims(sys.boundary) """ - box_volume(boundary) + volume(boundary) Calculate the volume of a 3D bounding box or the area of a 2D bounding box. """ -box_volume(b::Union{CubicBoundary, RectangularBoundary}) = prod(b.side_lengths) -box_volume(b::TriclinicBoundary) = b[1][1] * b[2][2] * b[3][3] +volume(b::Union{CubicBoundary, RectangularBoundary}) = prod(b.side_lengths) +volume(b::TriclinicBoundary) = b[1][1] * b[2][2] * b[3][3] """ box_center(boundary) @@ -810,7 +810,7 @@ function pressure(sys::AtomsBase.AbstractSystem{D}, neighbors, step_n::Integer=0 end NkT = energy_remove_mol(length(sys) * sys.k * temperature(sys)) vir = energy_remove_mol(virial(sys, neighbors, step_n; n_threads=n_threads)) - P = (NkT - (2 * vir) / D) / box_volume(sys.boundary) + P = (NkT - (2 * vir) / D) / volume(sys.boundary) if sys.energy_units == NoUnits || D != 3 # If implied energy units are (u * nm^2 * ps^-2) and everything is # consistent then this has implied units of (u * nm^-1 * ps^-2) diff --git a/test/basic.jl b/test/basic.jl index 90a1bc70d..820062a5e 100644 --- a/test/basic.jl +++ b/test/basic.jl @@ -56,8 +56,8 @@ @test ustrip(b) == CubicBoundary(4.0, 5.0, 6.0) @test ustrip(u"Å", b) == CubicBoundary(40.0, 50.0, 60.0) @test !Molly.has_infinite_boundary(b) - @test box_volume(b) == 120.0u"nm^3" - @test box_volume(CubicBoundary(0.0u"m"; check_positive=false)) == 0.0u"m^3" + @test volume(b) == 120.0u"nm^3" + @test volume(CubicBoundary(0.0u"m"; check_positive=false)) == 0.0u"m^3" @test box_center(b) == SVector(2.0, 2.5, 3.0)u"nm" sb = scale_boundary(b, 1.1) @test sb.side_lengths ≈ SVector(4.4, 5.5, 6.6)u"nm" @@ -72,8 +72,8 @@ @test ustrip(b) == RectangularBoundary(4.0, 5.0) @test ustrip(u"km", b) == RectangularBoundary(4e-3, 5e-3) @test !Molly.has_infinite_boundary(b) - @test box_volume(b) == 20.0u"m^2" - @test box_volume(RectangularBoundary(0.0u"m"; check_positive=false)) == 0.0u"m^2" + @test volume(b) == 20.0u"m^2" + @test volume(RectangularBoundary(0.0u"m"; check_positive=false)) == 0.0u"m^2" @test box_center(b) == SVector(2.0, 2.5)u"m" sb = scale_boundary(b, 0.9) @test sb.side_lengths ≈ SVector(3.6, 4.5)u"m" @@ -92,11 +92,11 @@ @test TriclinicBoundary([b.basis_vectors[1], b.basis_vectors[2], b.basis_vectors[3]]) == b @test AtomsBase.bounding_box(b) == (b.basis_vectors[1], b.basis_vectors[2], b.basis_vectors[3]) - @test box_volume(b) ≈ 3.89937463181886u"nm^3" + @test volume(b) ≈ 3.89937463181886u"nm^3" @test isapprox(box_center(b), SVector(2.28944, 1.1359815, 0.5116602)u"nm"; atol=1e-6u"nm") sb = scale_boundary(b, 1.2) @test [sb.α, sb.β, sb.γ] ≈ [b.α, b.β, b.γ] - @test box_volume(sb) ≈ box_volume(b) * 1.2^3 + @test volume(sb) ≈ volume(b) * 1.2^3 @test isapprox( Molly.cubic_bounding_box(b), SVector(4.5788800, 2.2719630, 1.0233205)u"nm"; diff --git a/test/simulation.jl b/test/simulation.jl index 6aafc3da1..0d0e7fcfe 100644 --- a/test/simulation.jl +++ b/test/simulation.jl @@ -832,7 +832,7 @@ end distance_sum += sqrt(min_dist2) end mean_distance = distance_sum / length(sys) - wigner_seitz_radius = cbrt(3 * box_volume(sys.boundary) / (4π * length(sys))) + wigner_seitz_radius = cbrt(3 * volume(sys.boundary) / (4π * length(sys))) @test wigner_seitz_radius < mean_distance < 2 * wigner_seitz_radius end @@ -941,8 +941,8 @@ end coords = place_atoms(n_atoms, boundary; min_dist=1.0u"nm") n_log_steps = 500 - box_volume_wrapper(sys, args...; kwargs...) = box_volume(sys.boundary) - VolumeLogger(n_steps) = GeneralObservableLogger(box_volume_wrapper, typeof(1.0u"nm^3"), n_steps) + volume_wrapper(sys, args...; kwargs...) = volume(sys.boundary) + VolumeLogger(n_steps) = GeneralObservableLogger(volume_wrapper, typeof(1.0u"nm^3"), n_steps) baro_f(pressure) = MonteCarloAnisotropicBarostat(pressure, temp, boundary) lang_f(barostat) = Langevin(dt=dt, temperature=temp, friction=friction, coupling=barostat) @@ -973,7 +973,7 @@ end potential_energy=PotentialEnergyLogger(n_log_steps), virial=VirialLogger(n_log_steps), pressure=PressureLogger(n_log_steps), - box_volume=VolumeLogger(n_log_steps), + volume=VolumeLogger(n_log_steps), ), ) @@ -988,8 +988,8 @@ end @test mean(values(sys.loggers.potential_energy)) < 0.0u"kJ * mol^-1" all(!isnothing, press) && @test 0.6u"bar" < mean(values(sys.loggers.pressure)) < 1.3u"bar" any(!isnothing, press) && @test 0.1u"bar" < std(values(sys.loggers.pressure)) < 2.5u"bar" - any(!isnothing, press) && @test 800.0u"nm^3" < mean(values(sys.loggers.box_volume)) < 2000u"nm^3" - any(!isnothing, press) && @test 80.0u"nm^3" < std(values(sys.loggers.box_volume)) < 500.0u"nm^3" + any(!isnothing, press) && @test 800.0u"nm^3" < mean(values(sys.loggers.volume)) < 2000u"nm^3" + any(!isnothing, press) && @test 80.0u"nm^3" < std(values(sys.loggers.volume)) < 500.0u"nm^3" axis_is_uncoupled = isnothing.(press) axis_is_unchanged = sys.boundary .== 8.0u"nm" @test all(axis_is_uncoupled .== axis_is_unchanged) @@ -1013,8 +1013,8 @@ end coords = place_atoms(n_atoms, boundary; min_dist=1.0u"nm") n_log_steps = 500 - box_volume_wrapper(sys, args...; kwargs...) = box_volume(sys.boundary) - VolumeLogger(n_steps) = GeneralObservableLogger(box_volume_wrapper, typeof(1.0u"nm^3"), n_steps) + volume_wrapper(sys, args...; kwargs...) = volume(sys.boundary) + VolumeLogger(n_steps) = GeneralObservableLogger(volume_wrapper, typeof(1.0u"nm^3"), n_steps) lang_f(barostat) = Langevin(dt=dt, temperature=temp, friction=friction, coupling=barostat) @@ -1044,7 +1044,7 @@ end potential_energy=PotentialEnergyLogger(n_log_steps), virial=VirialLogger(n_log_steps), pressure=PressureLogger(n_log_steps), - box_volume=VolumeLogger(n_log_steps), + volume=VolumeLogger(n_log_steps), ), ) From ce979b3f0d0cd4406fd1dca7902d86d242b35035 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 4 Oct 2024 17:58:30 +0100 Subject: [PATCH 58/74] density function and logger --- src/loggers.jl | 19 +++++++++++++++++++ src/spatial.jl | 25 +++++++++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/src/loggers.jl b/src/loggers.jl index acaa4c5d5..090735c55 100644 --- a/src/loggers.jl +++ b/src/loggers.jl @@ -13,6 +13,7 @@ export PotentialEnergyLogger, ForceLogger, VirialLogger, + DensityLogger, PressureLogger, StructureWriter, TimeCorrelationLogger, @@ -275,6 +276,24 @@ function Base.show(io::IO, vl::GeneralObservableLogger{T, typeof(virial_wrapper) vl.n_steps, ", ", length(values(vl)), " virials recorded") end +density_wrapper(sys, args...; kwargs...) = density(sys) + +""" + DensityLogger(n_steps) + DensityLogger(T, n_steps) + +Log the [`density`](@ref) of a system throughout a simulation. + +Not compatible with infinite boundaries. +""" +DensityLogger(T::Type, n_steps::Integer) = GeneralObservableLogger(density_wrapper, T, n_steps) +DensityLogger(n_steps::Integer) = DensityLogger(typeof(one(DefaultFloat)u"kg * m^-3"), n_steps) + +function Base.show(io::IO, dl::GeneralObservableLogger{T, typeof(density_wrapper)}) where T + print(io, "DensityLogger{", eltype(values(dl)), "} with n_steps ", + dl.n_steps, ", ", length(values(dl)), " densities recorded") +end + function pressure_wrapper(sys, neighbors, step_n; n_threads, kwargs...) return pressure(sys, neighbors, step_n; n_threads=n_threads) end diff --git a/src/spatial.jl b/src/spatial.jl index 9ef98f233..5d3aacebd 100644 --- a/src/spatial.jl +++ b/src/spatial.jl @@ -5,6 +5,7 @@ export RectangularBoundary, TriclinicBoundary, volume, + density, box_center, scale_boundary, random_coord, @@ -247,10 +248,34 @@ n_infinite_dims(sys::System) = n_infinite_dims(sys.boundary) volume(boundary) Calculate the volume of a 3D bounding box or the area of a 2D bounding box. + +Returns infinite volume for infinite boundaries. """ volume(b::Union{CubicBoundary, RectangularBoundary}) = prod(b.side_lengths) volume(b::TriclinicBoundary) = b[1][1] * b[2][2] * b[3][3] +""" + density(sys) + +The density of a [`System`](@ref). + +Returns zero density for infinite boundaries. +""" +function density(sys) + m = sum(mass, sys.atoms) + if dimension(m) == u"𝐌 * 𝐍^-1" + m_no_mol = m / Unitful.Na + else + m_no_mol = m + end + d = m_no_mol / volume(sys) + if unit(d) == NoUnits + return d + else + return uconvert(u"kg * m^-3", d) + end +end + """ box_center(boundary) From 72a4710c1fbb0d478702efdf3e4aedd4d1d73cda Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 4 Oct 2024 18:04:23 +0100 Subject: [PATCH 59/74] add VolumeLogger --- src/loggers.jl | 19 +++++++++++++++++++ src/spatial.jl | 4 +++- test/simulation.jl | 6 ------ 3 files changed, 22 insertions(+), 7 deletions(-) diff --git a/src/loggers.jl b/src/loggers.jl index 090735c55..f5382369c 100644 --- a/src/loggers.jl +++ b/src/loggers.jl @@ -13,6 +13,7 @@ export PotentialEnergyLogger, ForceLogger, VirialLogger, + VolumeLogger, DensityLogger, PressureLogger, StructureWriter, @@ -276,6 +277,24 @@ function Base.show(io::IO, vl::GeneralObservableLogger{T, typeof(virial_wrapper) vl.n_steps, ", ", length(values(vl)), " virials recorded") end +volume_wrapper(sys, args...; kwargs...) = volume(sys) + +""" + VolumeLogger(n_steps) + VolumeLogger(T, n_steps) + +Log the [`volume`](@ref) of a system throughout a simulation. + +Not compatible with infinite boundaries. +""" +VolumeLogger(T::Type, n_steps::Integer) = GeneralObservableLogger(volume_wrapper, T, n_steps) +VolumeLogger(n_steps::Integer) = VolumeLogger(typeof(one(DefaultFloat)u"nm^3"), n_steps) + +function Base.show(io::IO, vl::GeneralObservableLogger{T, typeof(volume_wrapper)}) where T + print(io, "VolumeLogger{", eltype(values(vl)), "} with n_steps ", + vl.n_steps, ", ", length(values(vl)), " volumes recorded") +end + density_wrapper(sys, args...; kwargs...) = density(sys) """ diff --git a/src/spatial.jl b/src/spatial.jl index 5d3aacebd..452bb6405 100644 --- a/src/spatial.jl +++ b/src/spatial.jl @@ -245,12 +245,14 @@ n_infinite_dims(b::TriclinicBoundary) = 0 n_infinite_dims(sys::System) = n_infinite_dims(sys.boundary) """ + volume(sys) volume(boundary) -Calculate the volume of a 3D bounding box or the area of a 2D bounding box. +Calculate the volume (3D) or area (2D) of a [`System`](@ref) or bounding box. Returns infinite volume for infinite boundaries. """ +volume(sys) = volume(sys.boundary) volume(b::Union{CubicBoundary, RectangularBoundary}) = prod(b.side_lengths) volume(b::TriclinicBoundary) = b[1][1] * b[2][2] * b[3][3] diff --git a/test/simulation.jl b/test/simulation.jl index 0d0e7fcfe..171aed7f0 100644 --- a/test/simulation.jl +++ b/test/simulation.jl @@ -941,9 +941,6 @@ end coords = place_atoms(n_atoms, boundary; min_dist=1.0u"nm") n_log_steps = 500 - volume_wrapper(sys, args...; kwargs...) = volume(sys.boundary) - VolumeLogger(n_steps) = GeneralObservableLogger(volume_wrapper, typeof(1.0u"nm^3"), n_steps) - baro_f(pressure) = MonteCarloAnisotropicBarostat(pressure, temp, boundary) lang_f(barostat) = Langevin(dt=dt, temperature=temp, friction=friction, coupling=barostat) @@ -1013,9 +1010,6 @@ end coords = place_atoms(n_atoms, boundary; min_dist=1.0u"nm") n_log_steps = 500 - volume_wrapper(sys, args...; kwargs...) = volume(sys.boundary) - VolumeLogger(n_steps) = GeneralObservableLogger(volume_wrapper, typeof(1.0u"nm^3"), n_steps) - lang_f(barostat) = Langevin(dt=dt, temperature=temp, friction=friction, coupling=barostat) barostat_test_set = ( From 15258901df4be90314c2bc8951727c945aa6b02d Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 4 Oct 2024 18:06:01 +0100 Subject: [PATCH 60/74] dipole_moment function --- src/spatial.jl | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/spatial.jl b/src/spatial.jl index 452bb6405..91f6e07dc 100644 --- a/src/spatial.jl +++ b/src/spatial.jl @@ -23,7 +23,8 @@ export virial, pressure, molecule_centers, - scale_coords! + scale_coords!, + dipole_moment """ CubicBoundary(x, y, z) @@ -942,3 +943,12 @@ function scale_coords!(sys, scale_factor; ignore_molecules=false) end return sys end + +""" + dipole_moment(sys) + +The dipole moment μ of a system. + +Requires the charges on the atoms to be set. +""" +dipole_moment(sys) = sum(sys.coords .* charges(sys)) From 38a07c8be71dbeb0e5d5170378a7e968c832c1c1 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Fri, 4 Oct 2024 19:10:25 +0100 Subject: [PATCH 61/74] density test --- docs/src/documentation.md | 6 +++++ ext/MollyGLMakieExt.jl | 4 ++-- src/loggers.jl | 46 +++++++++++++++++++-------------------- test/protein.jl | 2 ++ 4 files changed, 33 insertions(+), 25 deletions(-) diff --git a/docs/src/documentation.md b/docs/src/documentation.md index 133dab014..70a4bbd93 100644 --- a/docs/src/documentation.md +++ b/docs/src/documentation.md @@ -90,6 +90,7 @@ forces(sys) accelerations(sys) masses(sys) +density(sys) # 207.56738339673083 kg m^-3 temperature(sys) # 96.76667184796673 K random_velocities(sys, 300.0u"K") @@ -110,6 +111,7 @@ sys.loggers # For certain systems virial(sys) pressure(sys) +dipole_moment(sys) # AtomsBase.jl interface import AtomsBase @@ -1072,6 +1074,10 @@ The available loggers are: - [`KineticEnergyLogger`](@ref) - [`PotentialEnergyLogger`](@ref) - [`ForceLogger`](@ref) +- [`VolumeLogger`](@ref) +- [`DensityLogger`](@ref) +- [`VirialLogger`](@ref) +- [`PressureLogger`](@ref) - [`StructureWriter`](@ref) - [`TimeCorrelationLogger`](@ref) - [`AutoCorrelationLogger`](@ref) diff --git a/ext/MollyGLMakieExt.jl b/ext/MollyGLMakieExt.jl index 005ffd8a0..5509dddec 100644 --- a/ext/MollyGLMakieExt.jl +++ b/ext/MollyGLMakieExt.jl @@ -35,12 +35,12 @@ function Molly.visualize(coord_logger, if dims == 3 PointType = Point3f ax = Axis3(fig[1, 1], aspect=:data) - max_connection_dist = cbrt(volume(boundary)) / 2 + max_connection_dist = cbrt(Molly.volume(boundary)) / 2 elseif dims == 2 PointType = Point2f ax = Axis(fig[1, 1]) ax.aspect = DataAspect() - max_connection_dist = sqrt(volume(boundary)) / 2 + max_connection_dist = sqrt(Molly.volume(boundary)) / 2 else throw(ArgumentError("found $dims dimensions but can only visualize 2 or 3 dimensions")) end diff --git a/src/loggers.jl b/src/loggers.jl index f5382369c..e36ae0eba 100644 --- a/src/loggers.jl +++ b/src/loggers.jl @@ -12,9 +12,9 @@ export KineticEnergyLogger, PotentialEnergyLogger, ForceLogger, - VirialLogger, VolumeLogger, DensityLogger, + VirialLogger, PressureLogger, StructureWriter, TimeCorrelationLogger, @@ -255,28 +255,6 @@ function Base.show(io::IO, fl::GeneralObservableLogger{T, typeof(forces_wrapper) length(values(fl)) > 0 ? length(first(values(fl))) : "?", " atoms") end -function virial_wrapper(sys, neighbors, step_n; n_threads, kwargs...) - return virial(sys, neighbors, step_n; n_threads=n_threads) -end - -""" - VirialLogger(n_steps) - VirialLogger(T, n_steps) - -Log the [`virial`](@ref) of a system throughout a simulation. - -This should only be used on systems containing just pairwise interactions, or -where the specific interactions, general interactions and constraints do not -contribute to the virial. -""" -VirialLogger(T::Type, n_steps::Integer) = GeneralObservableLogger(virial_wrapper, T, n_steps) -VirialLogger(n_steps::Integer) = VirialLogger(typeof(one(DefaultFloat)u"kJ * mol^-1"), n_steps) - -function Base.show(io::IO, vl::GeneralObservableLogger{T, typeof(virial_wrapper)}) where T - print(io, "VirialLogger{", eltype(values(vl)), "} with n_steps ", - vl.n_steps, ", ", length(values(vl)), " virials recorded") -end - volume_wrapper(sys, args...; kwargs...) = volume(sys) """ @@ -313,6 +291,28 @@ function Base.show(io::IO, dl::GeneralObservableLogger{T, typeof(density_wrapper dl.n_steps, ", ", length(values(dl)), " densities recorded") end +function virial_wrapper(sys, neighbors, step_n; n_threads, kwargs...) + return virial(sys, neighbors, step_n; n_threads=n_threads) +end + +""" + VirialLogger(n_steps) + VirialLogger(T, n_steps) + +Log the [`virial`](@ref) of a system throughout a simulation. + +This should only be used on systems containing just pairwise interactions, or +where the specific interactions, general interactions and constraints do not +contribute to the virial. +""" +VirialLogger(T::Type, n_steps::Integer) = GeneralObservableLogger(virial_wrapper, T, n_steps) +VirialLogger(n_steps::Integer) = VirialLogger(typeof(one(DefaultFloat)u"kJ * mol^-1"), n_steps) + +function Base.show(io::IO, vl::GeneralObservableLogger{T, typeof(virial_wrapper)}) where T + print(io, "VirialLogger{", eltype(values(vl)), "} with n_steps ", + vl.n_steps, ", ", length(values(vl)), " virials recorded") +end + function pressure_wrapper(sys, neighbors, step_n; n_threads, kwargs...) return pressure(sys, neighbors, step_n; n_threads=n_threads) end diff --git a/test/protein.jl b/test/protein.jl index c7c9db637..b642077f1 100644 --- a/test/protein.jl +++ b/test/protein.jl @@ -9,6 +9,7 @@ coords=CoordinateLogger(10), energy=TotalEnergyLogger(10), writer=StructureWriter(10, temp_fp_pdb), + density=DensityLogger(10), ), data="test_data_peptide", ) @@ -33,6 +34,7 @@ s.velocities = [random_velocity(mass(a), temp) .* 0.01 for a in s.atoms] @time simulate!(s, simulator, n_steps; n_threads=1) + @test all(isapprox(1016.0870493u"kg * m^-3"), values(s.loggers.density)) traj = read(temp_fp_pdb, BioStructures.PDBFormat) rm(temp_fp_pdb) @test BioStructures.countmodels(traj) == 11 From 002006f80c6cb5806badc3185886f6149d880d74 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Mon, 7 Oct 2024 11:27:05 +0100 Subject: [PATCH 62/74] update Enzyme version --- Project.toml | 2 +- test/gradients.jl | 21 +++++++++++---------- test/runtests.jl | 3 --- 3 files changed, 12 insertions(+), 14 deletions(-) diff --git a/Project.toml b/Project.toml index 860cc1857..432e7f5b6 100644 --- a/Project.toml +++ b/Project.toml @@ -56,7 +56,7 @@ Combinatorics = "1" DataStructures = "0.18" Distances = "0.10" Distributions = "0.23, 0.24, 0.25" -Enzyme = "0.12" +Enzyme = "0.13" EzXML = "1" FLoops = "0.2" GLMakie = "0.8, 0.9, 0.10" diff --git a/test/gradients.jl b/test/gradients.jl index 25e3322c4..f70f01e21 100644 --- a/test/gradients.jl +++ b/test/gradients.jl @@ -191,20 +191,20 @@ end if forward grad_enzyme = ( autodiff( - Forward, loss, Duplicated, Duplicated(σ, one(T)), Const(r0), - Duplicated(copy(coords), zero(coords)), + set_runtime_activity(Forward), loss, Duplicated, + Duplicated(σ, one(T)), Const(r0), Duplicated(copy(coords), zero(coords)), Duplicated(copy(velocities), zero(velocities)), const_args..., - )[2], + )[1], autodiff( - Forward, loss, Duplicated, Const(σ), Duplicated(r0, one(T)), - Duplicated(copy(coords), zero(coords)), + set_runtime_activity(Forward), loss, Duplicated, + Const(σ), Duplicated(r0, one(T)), Duplicated(copy(coords), zero(coords)), Duplicated(copy(velocities), zero(velocities)), const_args..., - )[2], + )[1], ) else grad_enzyme = autodiff( - Reverse, loss, Active, Active(σ), Active(r0), - Duplicated(copy(coords), zero(coords)), + set_runtime_activity(Reverse), loss, Active, + Active(σ), Active(r0), Duplicated(copy(coords), zero(coords)), Duplicated(copy(velocities), zero(velocities)), const_args..., )[1][1:2] end @@ -426,8 +426,9 @@ end n_threads = parallel ? Threads.nthreads() : 1 grads_enzyme = Dict(k => 0.0 for k in keys(params_dic)) autodiff( - Reverse, test_fn, Active, Duplicated(params_dic, grads_enzyme), - Const(sys_ref), Duplicated(copy(sys_ref.coords), zero(sys_ref.coords)), + set_runtime_activity(Reverse), test_fn, Active, + Duplicated(params_dic, grads_enzyme), Const(sys_ref), + Duplicated(copy(sys_ref.coords), zero(sys_ref.coords)), Duplicated(sys_ref.neighbor_finder, sys_ref.neighbor_finder), Const(n_threads), ) diff --git a/test/runtests.jl b/test/runtests.jl index 1be787b08..a5daddf5b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -67,9 +67,6 @@ const openmm_dir = joinpath(data_dir, "openmm_6mrr") const temp_fp_pdb = tempname(cleanup=true) * ".pdb" const temp_fp_viz = tempname(cleanup=true) * ".mp4" -# Required for gradient tests -Enzyme.API.runtimeActivity!(true) - if GROUP in ("All", "NotGradients") # Some failures due to dependencies but there is an unbound args error Aqua.test_all( From c585953dd4d45badab192203e0c8b9d3b1c85d97 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Wed, 9 Oct 2024 16:42:58 +0100 Subject: [PATCH 63/74] add paper to docs --- docs/src/publications.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/src/publications.md b/docs/src/publications.md index a2a22a339..09236b2bc 100644 --- a/docs/src/publications.md +++ b/docs/src/publications.md @@ -11,3 +11,4 @@ Other papers that use, contribute to or are compatible with Molly are listed bel - Martínez L. CellListMap.jl: Efficient and customizable cell list implementation for calculation of pairwise particle properties within a cutoff, [Comput Phys Commun](https://doi.org/10.1016/j.cpc.2022.108452) 279, 108452 (2022) - Blassel N and Stoltz G. Fixing the flux: A dual approach to computing transport coefficients, [arXiv](https://arxiv.org/abs/2305.08224) (2023) - Witt WC et al. ACEpotentials.jl: A Julia implementation of the atomic cluster expansion, [J Chem Phys](https://doi.org/10.1063/5.0158783) 159, 164101 (2023) +- Monmarché P, Spacek R and Stoltz G. Transient subtraction: A control variate method for computing transport coefficients, [arXiv](https://arxiv.org/abs/2410.00212) (2024) From b6e305cc9c3b85db3e3aa986326c3363a4a9f2a6 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 17 Oct 2024 14:45:14 +0100 Subject: [PATCH 64/74] remove redundant method --- src/types.jl | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/types.jl b/src/types.jl index 9be353b69..ddde31f78 100644 --- a/src/types.jl +++ b/src/types.jl @@ -264,10 +264,6 @@ function Base.:+(a1::Atom, a2::Atom) a1.σ + a2.σ, a1.ϵ + a2.ϵ) end -function Base.getindex(atom::Atom, x::Symbol) - return hasfield(Atom, x) ? getfield(atom, x) : KeyError("no field $x in Atom") -end - # get function errors with AD dict_get(dic, key, default) = haskey(dic, key) ? dic[key] : default From f1e1fcacde8779604e86c4372d371105631fae99 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 17 Oct 2024 14:45:44 +0100 Subject: [PATCH 65/74] test options --- test/runtests.jl | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/test/runtests.jl b/test/runtests.jl index a5daddf5b..102a1d71c 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -23,6 +23,8 @@ using Test const GROUP = get(ENV, "GROUP", "All") if GROUP in ("Protein", "Gradients", "NotGradients") @warn "Only running $GROUP tests as GROUP is set to $GROUP" +elseif GROUP != "All" + error("Unrecognised test group, GROUP=$GROUP") end # Some CPU gradient tests give memory errors on CI @@ -47,12 +49,12 @@ else end # Allow CUDA device to be specified -const DEVICE = get(ENV, "DEVICE", "0") +const DEVICE = parse(Int, get(ENV, "DEVICE", "0")) const run_gpu_tests = get(ENV, "GPUTESTS", "1") != "0" && CUDA.functional() -const gpu_list = run_gpu_tests ? (false, true) : (false,) +const gpu_list = (run_gpu_tests ? (false, true) : (false,)) if run_gpu_tests - device!(parse(Int, DEVICE)) + device!(DEVICE) @info "The GPU tests will be run on device $DEVICE" elseif get(ENV, "GPUTESTS", "1") == "0" @warn "The GPU tests will not be run as GPUTESTS is set to 0" From 3d7bc7d3701c256211ca171d6e38ce65aa812a1e Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 17 Oct 2024 15:00:34 +0100 Subject: [PATCH 66/74] n_dimensions fixes --- docs/src/documentation.md | 2 +- src/spatial.jl | 5 ----- 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/docs/src/documentation.md b/docs/src/documentation.md index 70a4bbd93..fd912597d 100644 --- a/docs/src/documentation.md +++ b/docs/src/documentation.md @@ -470,7 +470,7 @@ sys = System( boundary=boundary, pairwise_inters=pairwise_inters, loggers=( - coords=CoordinateLogger(n_atoms, dims=n_dimensions(boundary)), + coords=CoordinateLogger(n_atoms, dims=2), montecarlo=MonteCarloLogger(), ), ) diff --git a/src/spatial.jl b/src/spatial.jl index 91f6e07dc..913dfff72 100644 --- a/src/spatial.jl +++ b/src/spatial.jl @@ -191,11 +191,6 @@ Base.getindex(b::TriclinicBoundary, i::Integer) = b.basis_vectors[i] Base.firstindex(b::TriclinicBoundary) = 1 Base.lastindex(b::TriclinicBoundary) = 3 -""" - n_dimensions(boundary) - -Number of dimensions of a bounding box. -""" AtomsBase.n_dimensions(::CubicBoundary) = 3 AtomsBase.n_dimensions(::RectangularBoundary) = 2 AtomsBase.n_dimensions(::TriclinicBoundary) = 3 From 600582553cdb7d308b2d7be9d2b4050764491976 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 17 Oct 2024 18:36:03 +0100 Subject: [PATCH 67/74] don't do sim gradient tests on CI --- test/gradients.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/gradients.jl b/test/gradients.jl index f70f01e21..059ea966a 100644 --- a/test/gradients.jl +++ b/test/gradients.jl @@ -410,8 +410,10 @@ end test_runs = ( ("Energy", test_energy_grad, 1e-8), ("Force" , test_forces_grad, 1e-8), - ("Sim" , test_sim_grad , 1e-2), ) + if !running_CI + push!(test_runs, ("Sim", test_sim_grad, 1e-2)) + end params_to_test = ( #"inter_LJ_weight_14", "atom_N_ϵ", From ab6b225893453cff7af161122c1afc2e18bd1405 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 17 Oct 2024 18:37:57 +0100 Subject: [PATCH 68/74] rename ForceLogger to ForcesLogger --- docs/src/documentation.md | 2 +- src/loggers.jl | 12 ++++++------ test/simulation.jl | 2 +- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/src/documentation.md b/docs/src/documentation.md index fd912597d..4010ca335 100644 --- a/docs/src/documentation.md +++ b/docs/src/documentation.md @@ -1073,7 +1073,7 @@ The available loggers are: - [`TotalEnergyLogger`](@ref) - [`KineticEnergyLogger`](@ref) - [`PotentialEnergyLogger`](@ref) -- [`ForceLogger`](@ref) +- [`ForcesLogger`](@ref) - [`VolumeLogger`](@ref) - [`DensityLogger`](@ref) - [`VirialLogger`](@ref) diff --git a/src/loggers.jl b/src/loggers.jl index e36ae0eba..7c54cf140 100644 --- a/src/loggers.jl +++ b/src/loggers.jl @@ -11,7 +11,7 @@ export TotalEnergyLogger, KineticEnergyLogger, PotentialEnergyLogger, - ForceLogger, + ForcesLogger, VolumeLogger, DensityLogger, VirialLogger, @@ -234,12 +234,12 @@ function forces_wrapper(sys, neighbors, step_n::Integer; n_threads::Integer, end """ - ForceLogger(n_steps; dims=3) - ForceLogger(T, n_steps; dims=3) + ForcesLogger(n_steps; dims=3) + ForcesLogger(T, n_steps; dims=3) Log the [`forces`](@ref) throughout a simulation. """ -function ForceLogger(T, n_steps::Integer; dims::Integer=3) +function ForcesLogger(T, n_steps::Integer; dims::Integer=3) return GeneralObservableLogger( forces_wrapper, Array{SArray{Tuple{dims}, T, 1, dims}, 1}, @@ -247,10 +247,10 @@ function ForceLogger(T, n_steps::Integer; dims::Integer=3) ) end -ForceLogger(n_steps::Integer; dims::Integer=3) = ForceLogger(typeof(one(DefaultFloat)u"kJ * mol^-1 * nm^-1"), n_steps; dims=dims) +ForcesLogger(n_steps::Integer; dims::Integer=3) = ForcesLogger(typeof(one(DefaultFloat)u"kJ * mol^-1 * nm^-1"), n_steps; dims=dims) function Base.show(io::IO, fl::GeneralObservableLogger{T, typeof(forces_wrapper)}) where T - print(io, "ForceLogger{", eltype(eltype(values(fl))), "} with n_steps ", + print(io, "ForcesLogger{", eltype(eltype(values(fl))), "} with n_steps ", fl.n_steps, ", ", length(values(fl)), " frames recorded for ", length(values(fl)) > 0 ? length(first(values(fl))) : "?", " atoms") end diff --git a/test/simulation.jl b/test/simulation.jl index 171aed7f0..472cf408a 100644 --- a/test/simulation.jl +++ b/test/simulation.jl @@ -88,7 +88,7 @@ end energy=TotalEnergyLogger(100), ke=KineticEnergyLogger(100), pe=PotentialEnergyLogger(100), - force=ForceLogger(100), + force=ForcesLogger(100), writer=StructureWriter(100, temp_fp_pdb), potkin_correlation=TimeCorrelationLogger(pot_obs, kin_obs, TP, TP, 1, 100), velocity_autocorrelation=AutoCorrelationLogger(V, TV, n_atoms, 100), From e5d5249fc55daa8b25f2d00d2f0b442be777f3a5 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 17 Oct 2024 18:38:56 +0100 Subject: [PATCH 69/74] rename VelocityLogger to VelocitiesLogger --- docs/src/documentation.md | 2 +- src/loggers.jl | 12 ++++++------ test/simulation.jl | 2 +- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/src/documentation.md b/docs/src/documentation.md index 4010ca335..62391b2af 100644 --- a/docs/src/documentation.md +++ b/docs/src/documentation.md @@ -1069,7 +1069,7 @@ The available loggers are: - [`GeneralObservableLogger`](@ref) - [`TemperatureLogger`](@ref) - [`CoordinateLogger`](@ref) -- [`VelocityLogger`](@ref) +- [`VelocitiesLogger`](@ref) - [`TotalEnergyLogger`](@ref) - [`KineticEnergyLogger`](@ref) - [`PotentialEnergyLogger`](@ref) diff --git a/src/loggers.jl b/src/loggers.jl index 7c54cf140..66968c4df 100644 --- a/src/loggers.jl +++ b/src/loggers.jl @@ -7,7 +7,7 @@ export log_property!, TemperatureLogger, CoordinateLogger, - VelocityLogger, + VelocitiesLogger, TotalEnergyLogger, KineticEnergyLogger, PotentialEnergyLogger, @@ -140,12 +140,12 @@ end velocities_wrapper(sys, args...; kwargs...) = copy(sys.velocities) """ - VelocityLogger(n_steps; dims=3) - VelocityLogger(T, n_steps; dims=3) + VelocitiesLogger(n_steps; dims=3) + VelocitiesLogger(T, n_steps; dims=3) Log the velocities throughout a simulation. """ -function VelocityLogger(T, n_steps::Integer; dims::Integer=3) +function VelocitiesLogger(T, n_steps::Integer; dims::Integer=3) return GeneralObservableLogger( velocities_wrapper, Array{SArray{Tuple{dims}, T, 1, dims}, 1}, @@ -153,10 +153,10 @@ function VelocityLogger(T, n_steps::Integer; dims::Integer=3) ) end -VelocityLogger(n_steps::Integer; dims::Integer=3) = VelocityLogger(typeof(one(DefaultFloat)u"nm * ps^-1"), n_steps; dims=dims) +VelocitiesLogger(n_steps::Integer; dims::Integer=3) = VelocitiesLogger(typeof(one(DefaultFloat)u"nm * ps^-1"), n_steps; dims=dims) function Base.show(io::IO, vl::GeneralObservableLogger{T, typeof(velocities_wrapper)}) where T - print(io, "VelocityLogger{", eltype(eltype(values(vl))), "} with n_steps ", + print(io, "VelocitiesLogger{", eltype(eltype(values(vl))), "} with n_steps ", vl.n_steps, ", ", length(values(vl)), " frames recorded for ", length(values(vl)) > 0 ? length(first(values(vl))) : "?", " atoms") end diff --git a/test/simulation.jl b/test/simulation.jl index 472cf408a..0db927396 100644 --- a/test/simulation.jl +++ b/test/simulation.jl @@ -84,7 +84,7 @@ end loggers=( temp=TemperatureLogger(100), coords=CoordinateLogger(100), - vels=VelocityLogger(100), + vels=VelocitiesLogger(100), energy=TotalEnergyLogger(100), ke=KineticEnergyLogger(100), pe=PotentialEnergyLogger(100), From 3e441156ff8cb9fc2199030d4670cc7f16375650 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 17 Oct 2024 18:42:20 +0100 Subject: [PATCH 70/74] rename CoordinateLogger to CoordinatesLogger --- docs/src/differentiable.md | 6 +++--- docs/src/documentation.md | 16 ++++++++-------- docs/src/examples.md | 10 +++++----- src/loggers.jl | 12 ++++++------ test/agent.jl | 2 +- test/basic.jl | 10 +++++----- test/energy_conservation.jl | 2 +- test/protein.jl | 4 ++-- test/simulation.jl | 26 +++++++++++++------------- 9 files changed, 44 insertions(+), 44 deletions(-) diff --git a/docs/src/differentiable.md b/docs/src/differentiable.md index ba8e88a41..3625625d1 100644 --- a/docs/src/differentiable.md +++ b/docs/src/differentiable.md @@ -76,7 +76,7 @@ simulator = VelocityVerlet( function loss(σ, coords, velocities) atoms = [Atom(0, 0, atom_mass, 0.0, σ, 0.2) for i in 1:n_atoms] - loggers = (coords=CoordinateLogger(Float64, 10),) + loggers = (coords=CoordinatesLogger(Float64, 10),) sys = System( atoms=atoms, @@ -182,7 +182,7 @@ simulator = VelocityVerlet( function loss(θ) atoms = [Atom(0, 0, atom_mass, 0.0, 0.0, 0.0) for i in 1:n_atoms] - loggers = (coords=CoordinateLogger(Float64, 2),) + loggers = (coords=CoordinatesLogger(Float64, 2),) specific_inter_lists = ( InteractionList2Atoms( [1, 2, 4, 5], @@ -322,7 +322,7 @@ simulator = VelocityVerlet( function loss() atoms = [Atom(0, 0, mass, 0.0f0, 0.0f0, 0.0f0) for i in 1:n_atoms] - loggers = (coords=CoordinateLogger(Float32, 10),) + loggers = (coords=CoordinatesLogger(Float32, 10),) general_inters = (NNBonds(),) sys = System( diff --git a/docs/src/documentation.md b/docs/src/documentation.md index 62391b2af..b70e5a592 100644 --- a/docs/src/documentation.md +++ b/docs/src/documentation.md @@ -66,7 +66,7 @@ sys = System( pairwise_inters=pairwise_inters, loggers=( temp=TemperatureLogger(10), - coords=CoordinateLogger(10), + coords=CoordinatesLogger(10), ), ) @@ -161,7 +161,7 @@ sys = System( pairwise_inters=(LennardJones(),), loggers=( temp=TemperatureLogger(typeof(1.0f0u"K"), 10), - coords=CoordinateLogger(typeof(1.0f0u"nm"), 10), + coords=CoordinatesLogger(typeof(1.0f0u"nm"), 10), ), ) @@ -233,7 +233,7 @@ sys = System( neighbor_finder=neighbor_finder, loggers=( temp=TemperatureLogger(10), - coords=CoordinateLogger(10), + coords=CoordinatesLogger(10), ), ) @@ -275,7 +275,7 @@ sys = System( boundary=boundary, velocities=velocities, pairwise_inters=pairwise_inters, - loggers=(coords=CoordinateLogger(Float32, 10; dims=2),), + loggers=(coords=CoordinatesLogger(Float32, 10; dims=2),), force_units=NoUnits, energy_units=NoUnits, ) @@ -470,7 +470,7 @@ sys = System( boundary=boundary, pairwise_inters=pairwise_inters, loggers=( - coords=CoordinateLogger(n_atoms, dims=2), + coords=CoordinatesLogger(n_atoms, dims=2), montecarlo=MonteCarloLogger(), ), ) @@ -1068,7 +1068,7 @@ Loggers record properties of the simulation to allow monitoring and analysis. The available loggers are: - [`GeneralObservableLogger`](@ref) - [`TemperatureLogger`](@ref) -- [`CoordinateLogger`](@ref) +- [`CoordinatesLogger`](@ref) - [`VelocitiesLogger`](@ref) - [`TotalEnergyLogger`](@ref) - [`KineticEnergyLogger`](@ref) @@ -1085,8 +1085,8 @@ The available loggers are: - [`ReplicaExchangeLogger`](@ref) - [`MonteCarloLogger`](@ref) -Many of the loggers can be initialised with just the number of steps between recorded values, e.g. `CoordinateLogger(10)`. -An optional first argument is the type of the recorded value; the above is equivalent to `CoordinateLogger(typeof(1.0u"nm"), 10)` but if the simulation did not use units then `CoordinateLogger(Float64, 10)` would be required. +Many of the loggers can be initialised with just the number of steps between recorded values, e.g. `CoordinatesLogger(10)`. +An optional first argument is the type of the recorded value; the above is equivalent to `CoordinatesLogger(typeof(1.0u"nm"), 10)` but if the simulation did not use units then `CoordinatesLogger(Float64, 10)` would be required. If the simulation is in 2D, giving `dims=2` as a keyword argument is required for some loggers. A logger's history can be accessed with `values(my_logger)`. diff --git a/docs/src/examples.md b/docs/src/examples.md index 4bce04177..91b215a32 100644 --- a/docs/src/examples.md +++ b/docs/src/examples.md @@ -111,7 +111,7 @@ sys = System( boundary=boundary, velocities=velocities, pairwise_inters=(inter,), - loggers=(coords=CoordinateLogger(typeof(1.0u"km"), 10),), + loggers=(coords=CoordinatesLogger(typeof(1.0u"km"), 10),), force_units=u"kg * km * d^-2", energy_units=u"kg * km^2 * d^-2", ) @@ -230,7 +230,7 @@ sys = System( pairwise_inters=pairwise_inters, neighbor_finder=neighbor_finder, loggers=( - coords=CoordinateLogger(Float64, 10; dims=2), + coords=CoordinatesLogger(Float64, 10; dims=2), SIR=SIRLogger(10), ), force_units=NoUnits, @@ -361,7 +361,7 @@ sys = System( pairwise_inters=(lj,), specific_inter_lists=(bonds, angles), neighbor_finder=neighbor_finder, - loggers=(coords=CoordinateLogger(200),), + loggers=(coords=CoordinatesLogger(200),), ) sim = Langevin(dt=0.002u"ps", temperature=300.0u"K", friction=1.0u"ps^-1") @@ -579,7 +579,7 @@ atoms = fill(Atom(mass=28.0), 2) coords = [SVector(1/8, 1/8, 1/8), SVector(-1/8, -1/8, -1/8)] velocities = [randn(SVector{3, Float64}) * 0.1 for _ in 1:2] boundary = CubicBoundary(Inf) -loggers = (coords=CoordinateLogger(Float64, 1),) +loggers = (coords=CoordinatesLogger(Float64, 1),) sys = System( atoms=atoms, @@ -699,7 +699,7 @@ sys = System( pairwise_inters=pairwise_inters, neighbor_finder=neighbor_finder, loggers=( - coords=CoordinateLogger(Float64, 20; dims=2), + coords=CoordinatesLogger(Float64, 20; dims=2), bonds=BondLogger(20), ), force_units=NoUnits, diff --git a/src/loggers.jl b/src/loggers.jl index 66968c4df..62c85bc41 100644 --- a/src/loggers.jl +++ b/src/loggers.jl @@ -6,7 +6,7 @@ export values, log_property!, TemperatureLogger, - CoordinateLogger, + CoordinatesLogger, VelocitiesLogger, TotalEnergyLogger, KineticEnergyLogger, @@ -116,12 +116,12 @@ end coordinates_wrapper(sys, args...; kwargs...) = copy(sys.coords) """ - CoordinateLogger(n_steps; dims=3) - CoordinateLogger(T, n_steps; dims=3) + CoordinatesLogger(n_steps; dims=3) + CoordinatesLogger(T, n_steps; dims=3) Log the coordinates throughout a simulation. """ -function CoordinateLogger(T, n_steps::Integer; dims::Integer=3) +function CoordinatesLogger(T, n_steps::Integer; dims::Integer=3) return GeneralObservableLogger( coordinates_wrapper, Array{SArray{Tuple{dims}, T, 1, dims}, 1}, @@ -129,10 +129,10 @@ function CoordinateLogger(T, n_steps::Integer; dims::Integer=3) ) end -CoordinateLogger(n_steps::Integer; dims::Integer=3) = CoordinateLogger(typeof(one(DefaultFloat)u"nm"), n_steps; dims=dims) +CoordinatesLogger(n_steps::Integer; dims::Integer=3) = CoordinatesLogger(typeof(one(DefaultFloat)u"nm"), n_steps; dims=dims) function Base.show(io::IO, cl::GeneralObservableLogger{T, typeof(coordinates_wrapper)}) where T - print(io, "CoordinateLogger{", eltype(eltype(values(cl))), "} with n_steps ", + print(io, "CoordinatesLogger{", eltype(eltype(values(cl))), "} with n_steps ", cl.n_steps, ", ", length(values(cl)), " frames recorded for ", length(values(cl)) > 0 ? length(first(values(cl))) : "?", " atoms") end diff --git a/test/agent.jl b/test/agent.jl index 740aeb74e..acbdc27d1 100644 --- a/test/agent.jl +++ b/test/agent.jl @@ -93,7 +93,7 @@ pairwise_inters=pairwise_inters, neighbor_finder=neighbor_finder, loggers=( - coords=CoordinateLogger(Float64, 10; dims=2), + coords=CoordinatesLogger(Float64, 10; dims=2), SIR=SIRLogger(10, []), ), force_units=NoUnits, diff --git a/test/basic.jl b/test/basic.jl index 820062a5e..b0fda9131 100644 --- a/test/basic.jl +++ b/test/basic.jl @@ -62,7 +62,7 @@ sb = scale_boundary(b, 1.1) @test sb.side_lengths ≈ SVector(4.4, 5.5, 6.6)u"nm" @test Molly.cubic_bounding_box(b) == SVector(4.0, 5.0, 6.0)u"nm" - @test Molly.axis_limits(CubicBoundary(4.0, 5.0, 6.0), CoordinateLogger(1), 2) == (0.0, 5.0) + @test Molly.axis_limits(CubicBoundary(4.0, 5.0, 6.0), CoordinatesLogger(1), 2) == (0.0, 5.0) @test_throws DomainError CubicBoundary(-4.0u"nm", 5.0u"nm", 6.0u"nm") @test_throws DomainError CubicBoundary( 4.0u"nm", 0.0u"nm", 6.0u"nm") @@ -78,7 +78,7 @@ sb = scale_boundary(b, 0.9) @test sb.side_lengths ≈ SVector(3.6, 4.5)u"m" @test Molly.cubic_bounding_box(b) == SVector(4.0, 5.0)u"m" - @test Molly.axis_limits(RectangularBoundary(4.0, 5.0), CoordinateLogger(1), 2) == (0.0, 5.0) + @test Molly.axis_limits(RectangularBoundary(4.0, 5.0), CoordinatesLogger(1), 2) == (0.0, 5.0) @test_throws DomainError RectangularBoundary(-4.0u"nm", 5.0u"nm") @test_throws DomainError RectangularBoundary( 4.0u"nm", 0.0u"nm") @@ -137,7 +137,7 @@ end atoms = fill(Atom(mass=1.0u"g/mol"), n_atoms) - loggers = (coords=CoordinateLogger(10),) + loggers = (coords=CoordinatesLogger(10),) temp = 100.0u"K" dt = 0.002u"ps" sim = VelocityVerlet(dt=dt, remove_CM_motion=false) @@ -405,7 +405,7 @@ end replica_velocities=replica_velocities, pairwise_inters=pairwise_inters, neighbor_finder=neighbor_finder, - replica_loggers=[(temp=TemperatureLogger(10), coords=CoordinateLogger(10)) + replica_loggers=[(temp=TemperatureLogger(10), coords=CoordinatesLogger(10)) for i in 1:n_replicas], data="test_data_repsys", ) @@ -417,7 +417,7 @@ end velocities=nothing, pairwise_inters=pairwise_inters, neighbor_finder=neighbor_finder, - loggers=(temp=TemperatureLogger(10), coords=CoordinateLogger(10)), + loggers=(temp=TemperatureLogger(10), coords=CoordinatesLogger(10)), data="test_data_sys", ) diff --git a/test/energy_conservation.jl b/test/energy_conservation.jl index a27a10917..d64a4cf2d 100644 --- a/test/energy_conservation.jl +++ b/test/energy_conservation.jl @@ -31,7 +31,7 @@ using Test boundary=boundary, pairwise_inters=(LennardJones(cutoff=cutoff, use_neighbors=false),), loggers=( - coords=CoordinateLogger(100), + coords=CoordinatesLogger(100), energy=TotalEnergyLogger(100), ), ) diff --git a/test/protein.jl b/test/protein.jl index b642077f1..7153eb633 100644 --- a/test/protein.jl +++ b/test/protein.jl @@ -6,7 +6,7 @@ joinpath(data_dir, "5XER", "gmx_top_ff.top"); loggers=( temp=TemperatureLogger(10), - coords=CoordinateLogger(10), + coords=CoordinatesLogger(10), energy=TotalEnergyLogger(10), writer=StructureWriter(10, temp_fp_pdb), density=DensityLogger(10), @@ -51,7 +51,7 @@ end joinpath(data_dir, "5XER", "gmx_top_ff.top"); loggers=( temp=TemperatureLogger(Float32, 10), - coords=CoordinateLogger(Float32, 10), + coords=CoordinatesLogger(Float32, 10), energy=TotalEnergyLogger(Float32, 10), ), units=false, diff --git a/test/simulation.jl b/test/simulation.jl index 0db927396..085a70204 100644 --- a/test/simulation.jl +++ b/test/simulation.jl @@ -18,7 +18,7 @@ ), loggers=( temp=TemperatureLogger(100), - coords=CoordinateLogger(100; dims=2), + coords=CoordinatesLogger(100; dims=2), gen_temp=GeneralObservableLogger(gen_temp_wrapper, typeof(temp), 10), avg_temp=AverageObservableLogger(Molly.temperature_wrapper, typeof(temp), 1; n_blocks=200), @@ -83,7 +83,7 @@ end ), loggers=( temp=TemperatureLogger(100), - coords=CoordinateLogger(100), + coords=CoordinatesLogger(100), vels=VelocitiesLogger(100), energy=TotalEnergyLogger(100), ke=KineticEnergyLogger(100), @@ -175,7 +175,7 @@ end n_steps=10, dist_cutoff=2.0u"nm", ), - loggers=(coords=CoordinateLogger(100),), + loggers=(coords=CoordinatesLogger(100),), ) @test Molly.has_infinite_boundary(boundary) @@ -217,7 +217,7 @@ end n_steps=10, dist_cutoff=2.0u"nm", ), - loggers=(coords=CoordinateLogger(100),), + loggers=(coords=CoordinatesLogger(100),), ) random_velocities!(s, temp) @@ -262,7 +262,7 @@ end ), loggers=( temp=TemperatureLogger(10), - coords=CoordinateLogger(10), + coords=CoordinatesLogger(10), ), ) @@ -318,7 +318,7 @@ end neighbor_finder=neighbor_finder, loggers=( temp=TemperatureLogger(100), - coords=CoordinateLogger(100), + coords=CoordinatesLogger(100), energy=TotalEnergyLogger(100), ), ) @@ -341,7 +341,7 @@ end boundary=boundary, velocities=velocities, general_inters=(MullerBrown(),), - loggers=(coords=CoordinateLogger(100; dims=2),), + loggers=(coords=CoordinatesLogger(100; dims=2),), ) simulator = VelocityVerlet(dt=0.002u"ps") @@ -379,7 +379,7 @@ end ), loggers=( temp=TemperatureLogger(100), - coords=CoordinateLogger(100), + coords=CoordinatesLogger(100), energy=TotalEnergyLogger(100), ), ) @@ -399,7 +399,7 @@ end ), loggers=( temp=TemperatureLogger(Float64, 100), - coords=CoordinateLogger(Float64, 100), + coords=CoordinatesLogger(Float64, 100), energy=TotalEnergyLogger(Float64, 100), ), force_units=NoUnits, @@ -440,7 +440,7 @@ end boundary=boundary, atoms_data=atoms_data, pairwise_inters=(LennardJones(),), - loggers=(coords=CoordinateLogger(100),), + loggers=(coords=CoordinatesLogger(100),), ) atom_selector(at, at_data) = at_data.atom_type == "A1" @@ -552,7 +552,7 @@ end pairwise_inters=(LennardJones(use_neighbors=true),), constraints=(cons,), neighbor_finder=neighbor_finder, - loggers=(coords=CoordinateLogger(10),), + loggers=(coords=CoordinatesLogger(10),), ) old_coords = copy(sys.coords) @@ -651,7 +651,7 @@ end ) n_replicas = 4 - replica_loggers = [(temp=TemperatureLogger(10), coords=CoordinateLogger(10)) for i in 1:n_replicas] + replica_loggers = [(temp=TemperatureLogger(10), coords=CoordinatesLogger(10)) for i in 1:n_replicas] repsys = ReplicaSystem( atoms=atoms, @@ -789,7 +789,7 @@ end pairwise_inters=(Coulomb(), ), neighbor_finder=neighbor_finder, loggers=( - coords=CoordinateLogger(10), + coords=CoordinatesLogger(10), mcl=MonteCarloLogger(), avgpe=AverageObservableLogger(Molly.potential_energy_wrapper, typeof(atoms[1].ϵ), 10), ), From b5903ad64c61ce701c1062602ae884d4fe003994 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 17 Oct 2024 18:50:10 +0100 Subject: [PATCH 71/74] update API docs --- docs/src/api.md | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/docs/src/api.md b/docs/src/api.md index d76480629..74c69d50f 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -2,17 +2,22 @@ The API reference can be found here. -Molly also re-exports [StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl) and [Unitful.jl](https://github.com/PainterQubits/Unitful.jl), making the likes of `SVector` and `1.0u"nm"` available when you call `using Molly`. +Molly re-exports [StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl) and [Unitful.jl](https://github.com/PainterQubits/Unitful.jl), making the likes of `SVector` and `1.0u"nm"` available when you call `using Molly`. -The [`visualize`](@ref) function is in a package extension and is only available once you have called `using GLMakie`. -The [`ASECalculator`](@ref) code is in a package extension and is only available once you have called `using PythonCall`. +Package extensions are used in order to reduce the number of dependencies: +- To use [`visualize`](@ref), call `using GLMakie`. +- To use [`ASECalculator`](@ref), call `using PythonCall`. + +## Exported names ```@index -Order = [:module, :type, :constant, :function, :macro] +Order = [:module, :type, :constant, :function, :macro] ``` +## Docstrings + ```@autodocs Modules = [Molly] Private = false -Order = [:module, :type, :constant, :function, :macro] +Order = [:module, :type, :constant, :function, :macro] ``` From bef88affd2a9e521778c7ea02436606488ce027b Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 17 Oct 2024 19:32:34 +0100 Subject: [PATCH 72/74] KernelDensity extension --- Project.toml | 3 +- docs/src/api.md | 1 + ext/MollyKernelDensityExt.jl | 28 ++++++++ src/Molly.jl | 1 - src/analysis.jl | 134 +++++++++++++++-------------------- test/Project.toml | 1 + test/runtests.jl | 1 + 7 files changed, 92 insertions(+), 77 deletions(-) create mode 100644 ext/MollyKernelDensityExt.jl diff --git a/Project.toml b/Project.toml index 432e7f5b6..66468122e 100644 --- a/Project.toml +++ b/Project.toml @@ -18,7 +18,6 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" EzXML = "8f5d6c58-4d21-5cfd-889c-e3ad7ee6a615" FLoops = "cc61a311-1640-44b5-9fba-1b764f453329" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" -KernelDensity = "5ab0869b-81aa-558d-bb23-cbf5423bbe9b" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce" PeriodicTable = "7b2266bf-644c-5ea3-82d8-af4bbd25a884" @@ -36,11 +35,13 @@ UnsafeAtomicsLLVM = "d80eeb9a-aca5-4d75-85e5-170c8b632249" Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a" +KernelDensity = "5ab0869b-81aa-558d-bb23-cbf5423bbe9b" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" [extensions] MollyEnzymeExt = "Enzyme" MollyGLMakieExt = ["GLMakie", "Colors"] +MollyKernelDensityExt = "KernelDensity" MollyPythonCallExt = "PythonCall" [compat] diff --git a/docs/src/api.md b/docs/src/api.md index 74c69d50f..c9265fb05 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -7,6 +7,7 @@ Molly re-exports [StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.j Package extensions are used in order to reduce the number of dependencies: - To use [`visualize`](@ref), call `using GLMakie`. - To use [`ASECalculator`](@ref), call `using PythonCall`. +- To use [`rdf`](@ref), call `using KernelDensity`. ## Exported names diff --git a/ext/MollyKernelDensityExt.jl b/ext/MollyKernelDensityExt.jl new file mode 100644 index 000000000..040e2b70f --- /dev/null +++ b/ext/MollyKernelDensityExt.jl @@ -0,0 +1,28 @@ +# Radial distribution function +# This file is only loaded when KernelDensity is imported + +module MollyKernelDensityExt + +using Molly +using KernelDensity + +function Molly.rdf(coords, boundary; npoints::Integer=200) + n_atoms = length(coords) + dims = length(first(coords)) + dists = distances(coords, boundary) + dists_vec = [dists[i, j] for i in 1:n_atoms, j in 1:n_atoms if j > i] + dist_unit = unit(first(dists_vec)) + kd = kde(ustrip.(dists_vec); npoints=npoints) + ρ = n_atoms / volume(boundary) + T = float_type(boundary) + if dims == 3 + normalizing_factor = 4 .* T(π) .* ρ .* step(kd.x) .* kd.x .^ 2 .* dist_unit .^ 3 + elseif dims == 2 + normalizing_factor = 2 .* T(π) .* ρ .* step(kd.x) .* kd.x .* dist_unit .^ 2 + end + bin_centers = collect(kd.x) .* dist_unit + density_weighted = kd.density ./ normalizing_factor + return bin_centers, density_weighted +end + +end diff --git a/src/Molly.jl b/src/Molly.jl index 4c3b22c23..19664debc 100644 --- a/src/Molly.jl +++ b/src/Molly.jl @@ -18,7 +18,6 @@ using Distributions using EzXML using FLoops using Graphs -using KernelDensity using NearestNeighbors using PeriodicTable using SimpleCrystals diff --git a/src/analysis.jl b/src/analysis.jl index ce32efc91..01429b5a2 100644 --- a/src/analysis.jl +++ b/src/analysis.jl @@ -1,57 +1,13 @@ # Analysis tools export - visualize, displacements, distances, - rdf, rmsd, radius_gyration, - hydrodynamic_radius - -""" - visualize(coord_logger, boundary, out_filepath; ) - -Visualize a simulation as an animation. - -This function is only available when GLMakie is imported. -It can take a while to run, depending on the length of the simulation and the -number of atoms. - -# Arguments -- `connections=Tuple{Int, Int}[]`: pairs of atoms indices to link with bonds. -- `connection_frames`: the frames in which bonds are shown. Should be a list of - the same length as the number of frames, where each item is a list of - `Bool`s of the same length as `connections`. Defaults to always `true`. -- `trails::Integer=0`: the number of preceding frames to show as transparent - trails. -- `framerate::Integer=30`: the frame rate of the animation. -- `color=:purple`: the color of the atoms. Can be a single color or a list of - colors of the same length as the number of atoms. -- `connection_color=:orange`: the color of the bonds. Can be a single color or a - list of colors of the same length as `connections`. -- `markersize=0.05`: the size of the atom markers, in the units of the data. -- `linewidth=2.0`: the width of the bond lines. -- `transparency=true`: whether transparency is active on the plot. -- `show_boundary::Bool=true`: whether to show the bounding box as lines. -- `boundary_linewidth=2.0`: the width of the boundary lines. -- `boundary_color=:black`: the color of the boundary lines. -- `kwargs...`: other keyword arguments are passed to the point plotting - function. -""" -function visualize end - -function axis_limits(boundary_conv, coord_logger, dim) - lim = boundary_conv[dim] - if isinf(lim) - # Find coordinate limits in given dimension - low = ustrip(minimum(cs -> minimum(c -> c[dim], cs), values(coord_logger))) - high = ustrip(maximum(cs -> maximum(c -> c[dim], cs), values(coord_logger))) - return low, high - else - return 0.0, lim - end -end + hydrodynamic_radius, + visualize, + rdf """ displacements(coords, boundary) @@ -75,34 +31,6 @@ periodic boundary conditions. """ distances(coords, boundary) = norm.(displacements(coords, boundary)) -""" - rdf(coords, boundary; npoints=200) - -Calculate the radial distribution function of a set of coordinates. - -This describes how density varies as a function of distance from each atom. -Returns a list of distance bin centers and a list of the corresponding -densities. -""" -function rdf(coords, boundary; npoints::Integer=200) - n_atoms = length(coords) - dims = length(first(coords)) - dists = distances(coords, boundary) - dists_vec = [dists[i, j] for i in 1:n_atoms, j in 1:n_atoms if j > i] - dist_unit = unit(first(dists_vec)) - kd = kde(ustrip.(dists_vec); npoints=npoints) - ρ = n_atoms / volume(boundary) - T = float_type(boundary) - if dims == 3 - normalizing_factor = 4 .* T(π) .* ρ .* step(kd.x) .* kd.x .^ 2 .* dist_unit .^ 3 - elseif dims == 2 - normalizing_factor = 2 .* T(π) .* ρ .* step(kd.x) .* kd.x .* dist_unit .^ 2 - end - bin_centers = collect(kd.x) .* dist_unit - density_weighted = kd.density ./ normalizing_factor - return bin_centers, density_weighted -end - """ rmsd(coords_1, coords_2) @@ -167,3 +95,59 @@ function hydrodynamic_radius(coords::AbstractArray{SVector{D, T}}, boundary) whe inv_R_hyd = sum_inv_dists / (2 * n_atoms^2) return inv(inv_R_hyd) end + +function axis_limits(boundary_conv, coord_logger, dim) + lim = boundary_conv[dim] + if isinf(lim) + # Find coordinate limits in given dimension + low = ustrip(minimum(cs -> minimum(c -> c[dim], cs), values(coord_logger))) + high = ustrip(maximum(cs -> maximum(c -> c[dim], cs), values(coord_logger))) + return low, high + else + return 0.0, lim + end +end + +""" + visualize(coord_logger, boundary, out_filepath; ) + +Visualize a simulation as an animation. + +This function is only available when GLMakie is imported. +It can take a while to run, depending on the length of the simulation and the +number of atoms. + +# Arguments +- `connections=Tuple{Int, Int}[]`: pairs of atoms indices to link with bonds. +- `connection_frames`: the frames in which bonds are shown. Should be a list of + the same length as the number of frames, where each item is a list of + `Bool`s of the same length as `connections`. Defaults to always `true`. +- `trails::Integer=0`: the number of preceding frames to show as transparent + trails. +- `framerate::Integer=30`: the frame rate of the animation. +- `color=:purple`: the color of the atoms. Can be a single color or a list of + colors of the same length as the number of atoms. +- `connection_color=:orange`: the color of the bonds. Can be a single color or a + list of colors of the same length as `connections`. +- `markersize=0.05`: the size of the atom markers, in the units of the data. +- `linewidth=2.0`: the width of the bond lines. +- `transparency=true`: whether transparency is active on the plot. +- `show_boundary::Bool=true`: whether to show the bounding box as lines. +- `boundary_linewidth=2.0`: the width of the boundary lines. +- `boundary_color=:black`: the color of the boundary lines. +- `kwargs...`: other keyword arguments are passed to the point plotting + function. +""" +function visualize end + +""" + rdf(coords, boundary; npoints=200) + +Calculate the radial distribution function of a set of coordinates. + +This function is only available when KernelDensity is imported. +This describes how density varies as a function of distance from each atom. +Returns a list of distance bin centers and a list of the corresponding +densities. +""" +function rdf end diff --git a/test/Project.toml b/test/Project.toml index 8d6edd647..a63d2fb9a 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -9,6 +9,7 @@ DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a" +KernelDensity = "5ab0869b-81aa-558d-bb23-cbf5423bbe9b" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SimpleCrystals = "64031d72-e220-11ed-1a7e-43a2532b2fa8" diff --git a/test/runtests.jl b/test/runtests.jl index 102a1d71c..00fc6654d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -8,6 +8,7 @@ import BioStructures # Imported to avoid clashing names using CUDA using Enzyme using FiniteDifferences +using KernelDensity import SimpleCrystals using DelimitedFiles From c0c2a1557a9df85a39bea96cf98cfb877ec32173 Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Thu, 17 Oct 2024 19:33:11 +0100 Subject: [PATCH 73/74] docs cover all docstrings --- docs/make.jl | 1 + src/units.jl | 6 +++--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/make.jl b/docs/make.jl index 243be41c3..187ef2cfd 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -7,6 +7,7 @@ makedocs( prettyurls=(get(ENV, "CI", nothing) == "true"), size_threshold_ignore=["api.md"], ), + modules=[Molly], pages=[ "Home" => "index.md", "Documentation" => "documentation.md", diff --git a/src/units.jl b/src/units.jl index 4c9f510be..2852b5932 100644 --- a/src/units.jl +++ b/src/units.jl @@ -1,9 +1,9 @@ export ustrip_vec # Unit types to dispatch on -@derived_dimension MolarMass Unitful.𝐌/Unitful.𝐍 true -@derived_dimension BoltzmannConstUnits Unitful.𝐌*Unitful.𝐋^2*Unitful.𝐓^-2*Unitful.𝚯^-1 true -@derived_dimension MolarBoltzmannConstUnits Unitful.𝐌*Unitful.𝐋^2*Unitful.𝐓^-2*Unitful.𝚯^-1*Unitful.𝐍^-1 true +@derived_dimension MolarMass Unitful.𝐌/Unitful.𝐍 +@derived_dimension BoltzmannConstUnits Unitful.𝐌*Unitful.𝐋^2*Unitful.𝐓^-2*Unitful.𝚯^-1 +@derived_dimension MolarBoltzmannConstUnits Unitful.𝐌*Unitful.𝐋^2*Unitful.𝐓^-2*Unitful.𝚯^-1*Unitful.𝐍^-1 """ ustrip_vec(x) From e9ad4ed37bc50d6efd44d04c083b17ab7677597b Mon Sep 17 00:00:00 2001 From: Joe Greener Date: Wed, 30 Oct 2024 11:24:41 +0000 Subject: [PATCH 74/74] Colors compat --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 66468122e..9cc8acc5b 100644 --- a/Project.toml +++ b/Project.toml @@ -52,7 +52,7 @@ BioStructures = "4" CUDA = "4.2, 5" CellListMap = "0.8.11, 0.9" Chemfiles = "0.10.3" -Colors = "0.11, 0.12" +Colors = "0.11, 0.12, 0.13" Combinatorics = "1" DataStructures = "0.18" Distances = "0.10"