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/Project.toml b/Project.toml index 53ff5c50e..9cc8acc5b 100644 --- a/Project.toml +++ b/Project.toml @@ -10,19 +10,14 @@ 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" 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" -ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" 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" @@ -34,17 +29,19 @@ 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" -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [weakdeps] 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] @@ -54,18 +51,15 @@ 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" +Colors = "0.11, 0.12, 0.13" Combinatorics = "1" DataStructures = "0.18" Distances = "0.10" Distributions = "0.23, 0.24, 0.25" -Enzyme = "0.11.15, 0.12" +Enzyme = "0.13" EzXML = "1" FLoops = "0.2" -ForwardDiff = "0.10.35" GLMakie = "0.8, 0.9, 0.10" Graphs = "1.8" KernelDensity = "0.5, 0.6" @@ -81,7 +75,5 @@ StaticArrays = "1.8.2" Statistics = "1.9" Unitful = "1" UnitfulAtomic = "1" -UnitfulChainRules = "0.1.2" UnsafeAtomicsLLVM = "0.1, 0.2" -Zygote = "0.6.67" julia = "1.9" diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index a144cf54e..08e6c5b4a 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)) @@ -90,14 +90,14 @@ 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)) - atoms = CuArray([Atom(charge=f32 ? 0.0f0 : 0.0, mass=atom_mass, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", + 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) - atoms = [Atom(charge=f32 ? 0.0f0 : 0.0, mass=atom_mass, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", + 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/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/docs/src/api.md b/docs/src/api.md index d76480629..c9265fb05 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -2,17 +2,23 @@ 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`. +- To use [`rdf`](@ref), call `using KernelDensity`. + +## 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] ``` 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/docs/src/differentiable.md b/docs/src/differentiable.md index 777a18dad..3625625d1 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 @@ -66,8 +65,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) @@ -78,8 +75,8 @@ simulator = VelocityVerlet( ) function loss(σ, coords, velocities) - atoms = [Atom(0, 0.0, atom_mass, σ, 0.2, false) for i in 1:n_atoms] - loggers = (coords=CoordinateLogger(Float64, 10),) + atoms = [Atom(0, 0, atom_mass, 0.0, σ, 0.2) for i in 1:n_atoms] + loggers = (coords=CoordinatesLogger(Float64, 10),) sys = System( atoms=atoms, @@ -98,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 @@ -186,8 +181,8 @@ simulator = VelocityVerlet( ) function loss(θ) - atoms = [Atom(0, 0.0, atom_mass, 0.0, 0.0, false) for i in 1:n_atoms] - loggers = (coords=CoordinateLogger(Float64, 2),) + atoms = [Atom(0, 0, atom_mass, 0.0, 0.0, 0.0) for i in 1:n_atoms] + loggers = (coords=CoordinatesLogger(Float64, 2),) specific_inter_lists = ( InteractionList2Atoms( [1, 2, 4, 5], @@ -204,9 +199,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, @@ -220,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 @@ -278,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. @@ -328,15 +321,15 @@ simulator = VelocityVerlet( ) function loss() - atoms = [Atom(0, 0.0f0, mass, 0.0f0, 0.0f0, false) for i in 1:n_atoms] - loggers = (coords=CoordinateLogger(Float32, 10),) + atoms = [Atom(0, 0, mass, 0.0f0, 0.0f0, 0.0f0) for i in 1:n_atoms] + loggers = (coords=CoordinatesLogger(Float32, 10),) general_inters = (NNBonds(),) 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, @@ -350,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/docs/src/documentation.md b/docs/src/documentation.md index 642c09bb1..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), ), ) @@ -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 @@ -159,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), ), ) @@ -231,7 +233,7 @@ sys = System( neighbor_finder=neighbor_finder, loggers=( temp=TemperatureLogger(10), - coords=CoordinateLogger(10), + coords=CoordinatesLogger(10), ), ) @@ -273,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, ) @@ -468,7 +470,7 @@ sys = System( boundary=boundary, pairwise_inters=pairwise_inters, loggers=( - coords=CoordinateLogger(n_atoms, dims=n_dimensions(boundary)), + coords=CoordinatesLogger(n_atoms, dims=2), montecarlo=MonteCarloLogger(), ), ) @@ -534,8 +536,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 +569,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 ``` @@ -586,12 +588,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 +608,18 @@ 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. +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. -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. @@ -621,16 +634,37 @@ 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) 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 -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 ``` @@ -638,7 +672,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 @@ -649,7 +691,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 = ( @@ -664,7 +707,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 @@ -681,9 +740,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 @@ -819,7 +879,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" @@ -874,7 +934,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,)) @@ -891,8 +951,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; @@ -1008,12 +1068,16 @@ Loggers record properties of the simulation to allow monitoring and analysis. The available loggers are: - [`GeneralObservableLogger`](@ref) - [`TemperatureLogger`](@ref) -- [`CoordinateLogger`](@ref) -- [`VelocityLogger`](@ref) +- [`CoordinatesLogger`](@ref) +- [`VelocitiesLogger`](@ref) - [`TotalEnergyLogger`](@ref) - [`KineticEnergyLogger`](@ref) - [`PotentialEnergyLogger`](@ref) -- [`ForceLogger`](@ref) +- [`ForcesLogger`](@ref) +- [`VolumeLogger`](@ref) +- [`DensityLogger`](@ref) +- [`VirialLogger`](@ref) +- [`PressureLogger`](@ref) - [`StructureWriter`](@ref) - [`TimeCorrelationLogger`](@ref) - [`AutoCorrelationLogger`](@ref) @@ -1021,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)`. @@ -1070,7 +1134,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/docs/src/examples.md b/docs/src/examples.md index 53a49afc7..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", ) @@ -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 @@ -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,11 +185,11 @@ function Molly.force(inter::SIRInteraction, atom_i.status = recovered end end - return zero(coord_i) + return zero(vec_ij) 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), @@ -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( @@ -237,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, @@ -368,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") @@ -586,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, @@ -612,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 @@ -628,7 +621,7 @@ struct BondableAtom partners::Set{Int} end -struct BondableInteraction <: PairwiseInteraction +struct BondableInteraction prob_formation::Float64 prob_break::Float64 dist_formation::Float64 @@ -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,11 +653,11 @@ function Molly.force(inter::BondableInteraction, fdr = -c * normalize(dr) return fdr else - return zero(coord_i) + return zero(dr) 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) @@ -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( @@ -713,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, @@ -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), @@ -1026,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...]) @@ -1060,12 +1043,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/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) diff --git a/ext/MollyEnzymeExt.jl b/ext/MollyEnzymeExt.jl new file mode 100644 index 000000000..90e015390 --- /dev/null +++ b/ext/MollyEnzymeExt.jl @@ -0,0 +1,27 @@ +# 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.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 +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/ext/MollyGLMakieExt.jl b/ext/MollyGLMakieExt.jl index fce5ce72f..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(box_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(box_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/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 ebe7f1d90..19664debc 100644 --- a/src/Molly.jl +++ b/src/Molly.jl @@ -9,28 +9,21 @@ import AtomsBase import AtomsCalculators import BioStructures # Imported to avoid clashing names using CellListMap -using ChainRules -using ChainRulesCore import Chemfiles using Combinatorics using CUDA using DataStructures using Distances using Distributions -using Enzyme using EzXML using FLoops -using ForwardDiff using Graphs -using KernelDensity using NearestNeighbors using PeriodicTable using SimpleCrystals using Unitful using UnitfulAtomic -using UnitfulChainRules using UnsafeAtomicsLLVM -using Zygote using LinearAlgebra using Random @@ -68,8 +61,5 @@ 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/src/analysis.jl b/src/analysis.jl index a63d9a07c..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) @@ -62,7 +18,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 @@ -76,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 / box_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) @@ -163,9 +90,64 @@ 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) 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/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 diff --git a/src/coupling.jl b/src/coupling.jl index 5c960a94d..15e8393de 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 @@ -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 @@ -180,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 @@ -194,14 +192,15 @@ 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; n_threads=n_threads) - V = box_volume(sys.boundary) + E = potential_energy(sys, neighbors, step_n; n_threads=n_threads) + 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)) - old_coords = copy(sys.coords) + old_coords .= sys.coords old_boundary = sys.boundary scale_coords!(sys, l_scale) @@ -213,14 +212,14 @@ 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) 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 @@ -279,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} @@ -304,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 " * @@ -337,6 +334,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 @@ -349,12 +347,12 @@ function apply_coupling!(sys::System{D, G, T}, mask1[axis] = true mask2[axis] = false - E = potential_energy(sys, neighbors; n_threads=n_threads) - V = box_volume(sys.boundary) + E = potential_energy(sys, neighbors, step_n; n_threads=n_threads) + 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) - old_coords = copy(sys.coords) + old_coords .= sys.coords old_boundary = sys.boundary scale_coords!(sys, l_scale) @@ -366,14 +364,14 @@ 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) 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 @@ -437,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} @@ -468,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 @@ -510,6 +507,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 @@ -521,8 +519,8 @@ function apply_coupling!(sys::System{D, G, T}, axis = 1 end - E = potential_energy(sys, neighbors; n_threads=n_threads) - V = box_volume(sys.boundary) + E = potential_energy(sys, neighbors, step_n; n_threads=n_threads) + 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)) @@ -545,7 +543,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) @@ -557,7 +555,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) @@ -566,7 +564,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/cuda.jl b/src/cuda.jl index 15b2e5d74..65ba1ede5 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, - pairwise_inters, nbs, force_units, ::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, 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) 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); - 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!( - fs_mat, coords, 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, atoms_var, boundary, inters, - neighbors_var, ::Val{D}, ::Val{F}) where {D, F} +function pairwise_force_kernel_nl!(forces, coords_var, velocities_var, atoms_var, boundary, inters, + 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) 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], step_n) 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, - ::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) 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, step_n) 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, step_n) 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 @@ -193,48 +199,47 @@ end return f end -function specific_force_gpu(inter_list::InteractionList1Atoms, coords::AbstractArray{SVector{D, C}}, - boundary, force_units, ::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, 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} - 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, 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} - 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, 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} - 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, 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) @@ -242,7 +247,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 @@ -253,9 +258,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) @@ -264,7 +271,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 @@ -276,9 +284,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) @@ -288,7 +298,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 @@ -301,9 +312,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) @@ -314,7 +328,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 @@ -328,18 +344,20 @@ 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, - pairwise_inters, nbs, 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, atoms, boundary, pairwise_inters, nbs, Val(energy_units)) - return pe_vec + pe_vec_nounits, coords, velocities, atoms, boundary, pairwise_inters, nbs, + step_n, Val(energy_units)) + return pe_vec_nounits end -function pairwise_pe_kernel!(energy, coords_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) neighbors = CUDA.Const(neighbors_var) @@ -347,13 +365,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, step_n) 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, step_n) end if unit(pe) != E error("wrong energy unit returned, was expecting $E but got $(unit(pe))") @@ -363,47 +381,47 @@ function pairwise_pe_kernel!(energy, coords_var, atoms_var, boundary, inters, ne return nothing end -function specific_pe_gpu(inter_list::InteractionList1Atoms, coords::AbstractArray{SVector{D, C}}, - boundary, 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, boundary, 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}}, - boundary, 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, boundary, 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}}, - boundary, 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, boundary, 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}}, - boundary, 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, boundary, 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, 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) @@ -411,7 +429,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 @@ -420,9 +439,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) @@ -431,7 +452,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 @@ -440,9 +462,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) @@ -452,7 +476,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 @@ -461,9 +487,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) @@ -474,7 +502,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/cutoffs.jl b/src/cutoffs.jl index bade01c03..b1d2e65cf 100644 --- a/src/cutoffs.jl +++ b/src/cutoffs.jl @@ -145,44 +145,30 @@ 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 - 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 -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 - 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 diff --git a/src/energy.jl b/src/energy.jl index bc536e6ab..c0f10c6a7 100644 --- a/src/energy.jl +++ b/src/energy.jl @@ -48,19 +48,21 @@ 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. - potential_energy(inter::PairwiseInteraction, vec_ij, coord_i, coord_j, - atom_i, atom_j, boundary) - 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, 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. @@ -70,8 +72,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 -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)) @@ -79,162 +81,182 @@ 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.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, + step_n=step_n, n_threads=n_threads), + ) end return pe end -function 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::Val{T}) where T - pe_vec = zeros(T, 1) - potential_energy_pair_spec!(pe_vec, 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) - 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, 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) - - @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, coords[i], coords[j], atoms[i], - atoms[j], boundary) - for inter in pairwise_inters_nonl[2:end] - pe += potential_energy(inter, dr, coords[i], coords[j], atoms[i], - atoms[j], boundary) - 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, coords[i], coords[j], atoms[i], - atoms[j], boundary, special) - for inter in pairwise_inters_nl[2:end] - pe += potential_energy(inter, dr, coords[i], coords[j], atoms[i], - atoms[j], boundary, special) - 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, coords[i], coords[j], atoms[i], - atoms[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 += potential_energy(inter, dr, coords[i], coords[j], atoms[i], - atoms[j], boundary) + 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, coords[i], coords[j], atoms[i], - atoms[j], boundary, 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, coords[i], coords[j], atoms[i], - atoms[j], boundary, 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) - check_energy_units(pe, energy_units) - pe_sum += ustrip(pe) + pe_inter = potential_energy(inter, coords[i], boundary, atoms[i], energy_units, + velocities[i], step_n) + 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) - check_energy_units(pe, energy_units) - pe_sum += ustrip(pe) + 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_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) - check_energy_units(pe, energy_units) - pe_sum += ustrip(pe) + 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_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) - check_energy_units(pe, energy_units) - pe_sum += ustrip(pe) + 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_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} - 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.atoms, sys.boundary, pairwise_inters_nonl, - nbs, 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)) @@ -244,35 +266,32 @@ 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, - nbs, 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.boundary, 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), + AtomsCalculators.potential_energy(sys, inter; neighbors=neighbors, + step_n=step_n, n_threads=n_threads), ) end - 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) + return pe 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) +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 fe0ba8165..760089a54 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,39 +20,35 @@ 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 """ - 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::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, 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(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, 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) @@ -119,8 +115,20 @@ 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()) + 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. @@ -129,8 +137,15 @@ 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 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)) @@ -138,174 +153,210 @@ 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, - 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, 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, 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 + + 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 += 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 - return fs + return fs_nounits end -function 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) - fs = ustrip_vec.(zero(coords)) - forces_pair_spec!(fs, 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) - return fs * force_units -end +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))) -function forces_pair_spec!(fs, 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) - 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, coords[i], coords[j], atoms[i], - atoms[j], boundary) - for inter in pairwise_inters_nonl[2:end] - f += force(inter, dr, coords[i], coords[j], atoms[i], atoms[j], boundary) - 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, coords[i], coords[j], atoms[i], - atoms[j], boundary, special) - for inter in pairwise_inters_nl[2:end] - f += force(inter, dr, coords[i], coords[j], atoms[i], atoms[j], boundary, - special) - 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, 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 + 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, 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], step_n) 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], 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, 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], step_n) 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], 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 + @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) + 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 @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) - fs[j] += ustrip.(sf.f2) + fs_nounits[i] += ustrip.(sf.f1) + fs_nounits[j] += ustrip.(sf.f2) 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) - 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) - 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 @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) 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} - 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.atoms, sys.boundary, pairwise_inters_nonl, - nbs, sys.force_units, 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)) @@ -315,22 +366,24 @@ 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, - nbs, sys.force_units, val_ft) + 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 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) + 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(unit(eltype(eltype(fs_gen))), sys.force_units) - fs += ustrip_vec.(fs_gen) + 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 - return fs * sys.force_units + return fs_nounits end diff --git a/src/gradients.jl b/src/gradients.jl deleted file mode 100644 index 7238d6bee..000000000 --- a/src/gradients.jl +++ /dev/null @@ -1,272 +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 - 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 - 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 - # 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) - 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.charge, # Residue-specific - dict_get(params_dic, key_prefix * "mass" , at.mass), - dict_get(params_dic, key_prefix * "σ" , at.σ ), - dict_get(params_dic, key_prefix * "ϵ" , at.ϵ ), - at.solute, - ) -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{S, C, W, WS, F, E}, params_dic) where {S, C, W, WS, F, E} - key_prefix = "inter_LJ_" - LennardJones{S, C, W, WS, F, E}( - 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 - -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), - inter.force_units, - inter.energy_units, - ) -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), - inter.force_units, - inter.energy_units, - ) -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/buckingham.jl b/src/interactions/buckingham.jl index 102b215ca..cf6fb790d 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, force_units, energy_units) + Buckingham(; cutoff, use_neighbors, shortcut, A_mixing, B_mixing, + C_mixing, weight_special) The Buckingham interaction between two atoms. @@ -23,48 +45,47 @@ 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 - cutoff::C - use_neighbors::Bool - weight_special::W - force_units::F - energy_units::E +@kwdef struct Buckingham{C, S, A, B, M, W} + cutoff::C = NoCutoff() + use_neighbors::Bool = false + 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 -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) +use_neighbors(inter::Buckingham) = inter.use_neighbors + +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 -use_neighbors(inter::Buckingham) = inter.use_neighbors +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{C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) 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 +@inline function force(inter::Buckingham, + dr, + atom_i, + atom_j, + force_units=u"kJ * mol^-1 * nm^-1", + 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, coord_i, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, force_units) if special return f * dr * inter.weight_special else @@ -77,28 +98,25 @@ 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}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) 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 +@inline function potential_energy(inter::Buckingham, + dr, + atom_i, + atom_j, + energy_units=u"kJ * mol^-1", + 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, coord_i, 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/cosine_angle.jl b/src/interactions/cosine_angle.jl index a49c36e3e..b2dba1c95 100644 --- a/src/interactions/cosine_angle.jl +++ b/src/interactions/cosine_angle.jl @@ -11,14 +11,14 @@ 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 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/coulomb.jl b/src/interactions/coulomb.jl index 595d160e7..46815d989 100644 --- a/src/interactions/coulomb.jl +++ b/src/interactions/coulomb.jl @@ -3,8 +3,10 @@ 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, force_units, energy_units) + Coulomb(; cutoff, use_neighbors, weight_special, coulomb_const) The Coulomb electrostatic interaction between two atoms. @@ -13,39 +15,17 @@ 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 - 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 - -function Coulomb(; - cutoff=NoCutoff(), - 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) +@kwdef struct Coulomb{C, W, T} + cutoff::C = NoCutoff() + use_neighbors::Bool = false + weight_special::W = 1 + coulomb_const::T = 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(coul.cutoff, coul.use_neighbors, zero(W), zero(T)) end function Base.:+(c1::Coulomb, c2::Coulomb) @@ -54,26 +34,40 @@ 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 +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 + +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, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) where C + dr, + atom_i, + atom_j, + force_units=u"kJ * mol^-1 * nm^-1", + special=false, + 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, coord_i, inter.force_units) + f = force_divr_with_cutoff(inter, r2, params, cutoff, force_units) if special return f * dr * inter.weight_special else @@ -81,25 +75,24 @@ 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}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) where C + dr, + atom_i, + atom_j, + energy_units=u"kJ * mol^-1", + special=false, + 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, coord_i, inter.energy_units) + pe = potential_with_cutoff(inter, r2, params, cutoff, energy_units) if special return pe * inter.weight_special else @@ -107,13 +100,12 @@ 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""" - CoulombSoftCore(; cutoff, α, λ, p, use_neighbors, lorentz_mixing, weight_special, - coulomb_const, force_units, energy_units) + CoulombSoftCore(; cutoff, α, λ, p, use_neighbors, σ_mixing, weight_special, coulomb_const) The Coulomb electrostatic interaction between two atoms with a soft core. @@ -123,54 +115,31 @@ 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 - cutoff::C - α::A - λ::L - p::P - σ6_fac::R - use_neighbors::Bool - lorentz_mixing::Bool - weight_special::W - coulomb_const::T - force_units::F - energy_units::E -end - -function CoulombSoftCore(; - cutoff=NoCutoff(), - α=1, - λ=0, - p=2, - 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") - σ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) +@kwdef struct CoulombSoftCore{C, A, L, P, S, W, T, R} + cutoff::C = NoCutoff() + α::A = 1 + λ::L = 0 + p::P = 2 + use_neighbors::Bool = false + σ_mixing::S = lorentz_σ_mixing + weight_special::W = 1 + coulomb_const::T = coulomb_const + σ6_fac::R = α * λ^p 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, S, W, T, R}) where {C, A, L, P, S, 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), - coul.force_units, - coul.energy_units, + zero(R), ) end @@ -180,32 +149,29 @@ 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.force_units, - c1.energy_units, + c1.σ6_fac + c2.σ6_fac, ) end -@inline function force(inter::CoulombSoftCore{C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) where C +@inline function force(inter::CoulombSoftCore, + dr, + atom_i, + atom_j, + force_units=u"kJ * mol^-1 * nm^-1", + special=false, + args...) 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) + σ = inter.σ_mixing(atom_i, atom_j) + params = (ke, 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, force_units) if special return f * dr * inter.weight_special else @@ -213,30 +179,29 @@ 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 -@inline function potential_energy(inter::CoulombSoftCore{C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) where C +@inline function potential_energy(inter::CoulombSoftCore, + dr, + atom_i, + atom_j, + energy_units=u"kJ * mol^-1", + special=false, + args...) 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) + σ = inter.σ_mixing(atom_i, atom_j) + params = (ke, 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, energy_units) if special return pe * inter.weight_special else @@ -244,55 +209,37 @@ 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 +const crf_solvent_dielectric = 78.3 + """ 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 +@kwdef struct CoulombReactionField{D, S, W, T} 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 - -function CoulombReactionField(; - dist_cutoff, - 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) + 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 -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, + coul.use_neighbors, zero(W), zero(T), - coul.force_units, - coul.energy_units, ) end @@ -303,25 +250,42 @@ 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 +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 + +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, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) + dr, + atom_i, + atom_j, + force_units=u"kJ * mol^-1 * nm^-1", + special=false, + args...) r2 = sum(abs2, dr) if r2 > (inter.dist_cutoff ^ 2) - return ustrip.(zero(coord_i)) * inter.force_units + 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 @@ -333,7 +297,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 @@ -343,19 +307,18 @@ 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=u"kJ * mol^-1", + special=false, + args...) r2 = sum(abs2, dr) if r2 > (inter.dist_cutoff ^ 2) - return ustrip(zero(coord_i[1])) * inter.energy_units + 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 @@ -369,7 +332,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/fene_bond.jl b/src/interactions/fene_bond.jl index 3846b377b..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 @@ -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/gravity.jl b/src/interactions/gravity.jl index 5246c7ce2..0abed2ba1 100644 --- a/src/interactions/gravity.jl +++ b/src/interactions/gravity.jl @@ -10,22 +10,24 @@ 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} + 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 + 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, - 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 +39,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/harmonic_angle.jl b/src/interactions/harmonic_angle.jl index fd3cc940b..5a2ef069b 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 @@ -23,7 +23,29 @@ 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) +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 + +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) bc = vector_pad3D(coords_j, coords_k, boundary) @@ -42,7 +64,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..5bf33ea21 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 @@ -21,14 +21,36 @@ 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) +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 + +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) 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..b1464c680 100644 --- a/src/interactions/harmonic_position_restraint.jl +++ b/src/interactions/harmonic_position_restraint.jl @@ -10,14 +10,14 @@ 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 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/implicit_solvent.jl b/src/interactions/implicit_solvent.jl index 5d90cdc4a..2c7ccfe94 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 @@ -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) @@ -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) @@ -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") @@ -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) @@ -597,6 +597,44 @@ function ImplicitSolventGBN2(atoms::AbstractArray{Atom{T, M, 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}() + 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)) @@ -637,7 +675,26 @@ with respect to atomic distance. Custom GBSA methods should implement this function. """ -function born_radii_and_grad(inter::ImplicitSolventOBC, coords, boundary) +function born_radii_and_grad(inter::ImplicitSolventOBC{T}, coords, boundary) where T + 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) + 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, @@ -800,36 +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 - -struct ForceLoopResult2{V} - fi::V - fj::V -end - -get_fi(r::Union{ForceLoopResult1, ForceLoopResult2}) = r.fi -get_fj(r::Union{ForceLoopResult1, ForceLoopResult2}) = 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) @@ -851,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 @@ -863,8 +900,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 @@ -876,37 +912,40 @@ 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 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], + inter.oris[i], inter.srjs[j], inter.dist_cutoff, boundary) + fs[i] = fs[i] .- f + fs[j] = fs[j] .+ f + end + end - 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,)) - - 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, @@ -1114,7 +1153,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) @@ -1125,9 +1184,10 @@ function AtomsCalculators.potential_energy(sys, inter::AbstractGBSA; kwargs...) 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 diff --git a/src/interactions/lennard_jones.jl b/src/interactions/lennard_jones.jl index 3dc1e83af..7cc843750 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, - force_units, energy_units, skip_shortcut) + LennardJones(; cutoff, use_neighbors, shortcut, σ_mixing, ϵ_mixing, weight_special) The Lennard-Jones 6-12 interaction between two atoms. @@ -20,89 +38,76 @@ and the force on each atom by \end{aligned} ``` """ -struct LennardJones{S, C, W, WS, F, E} <: 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(; - cutoff=NoCutoff(), - use_neighbors=false, - 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) +@kwdef struct LennardJones{C, H, S, E, W} + cutoff::C = NoCutoff() + use_neighbors::Bool = false + 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 -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{C, H, S, E, W}) where {C, H, S, E, W} + return LennardJones( lj.cutoff, - false, - false, + lj.use_neighbors, + lj.shortcut, + lj.σ_mixing, + lj.ϵ_mixing, 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, 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, - 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} - 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 - 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 - # 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.ϵ) +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, + atom_j, + force_units=u"kJ * mol^-1 * nm^-1", + special=false, + args...) + if inter.shortcut(atom_i, atom_j) + return ustrip.(zero(dr)) * force_units end + σ = inter.σ_mixing(atom_i, atom_j) + ϵ = inter.ϵ_mixing(atom_i, atom_j) cutoff = inter.cutoff r2 = sum(abs2, dr) σ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, force_units) if special return f * dr * inter.weight_special else @@ -115,32 +120,25 @@ 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}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) 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 - 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.ϵ) +@inline function potential_energy(inter::LennardJones, + dr, + atom_i, + atom_j, + energy_units=u"kJ * mol^-1", + special=false, + args...) + if inter.shortcut(atom_i, atom_j) + return ustrip(zero(dr[1])) * energy_units end + σ = inter.σ_mixing(atom_i, atom_j) + ϵ = inter.ϵ_mixing(atom_i, atom_j) cutoff = inter.cutoff r2 = sum(abs2, dr) σ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, energy_units) if special return pe * inter.weight_special else @@ -154,8 +152,8 @@ function potential(::LennardJones, r2, invr2, (σ2, ϵ)) end @doc raw""" - LennardJonesSoftCore(; cutoff, α, λ, p, use_neighbors, lorentz_mixing, weight_special, - weight_solute_solvent, force_units, energy_units, 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. @@ -173,103 +171,70 @@ 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 - cutoff::C - α::A - λ::L - p::P - σ6_fac::R - use_neighbors::Bool - lorentz_mixing::Bool - weight_special::W - weight_solute_solvent::WS - force_units::F - energy_units::E -end - -function LennardJonesSoftCore(; - cutoff=NoCutoff(), - α=1, - λ=0, - p=2, - use_neighbors=false, - 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)}( - cutoff, α, λ, p, σ6_fac, use_neighbors, lorentz_mixing, weight_special, - weight_solute_solvent, force_units, energy_units) +@kwdef struct LennardJonesSoftCore{C, A, L, P, H, S, E, W, R} + cutoff::C = NoCutoff() + α::A = 1 + λ::L = 0 + p::P = 2 + use_neighbors::Bool = false + 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{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{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), zero(L), zero(P), - zero(R), - false, - false, + lj.use_neighbors, + lj.shortcut, + lj.σ_mixing, + lj.ϵ_mixing, zero(W), - zero(WS), - lj.force_units, - lj.energy_units, + zero(R), ) 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, 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.force_units, - l1.energy_units, + l1.σ6_fac + l2.σ6_fac, ) end -@inline function force(inter::LennardJonesSoftCore{S, C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary, - special::Bool=false) 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 - 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.ϵ) +@inline function force(inter::LennardJonesSoftCore, + dr, + atom_i, + atom_j, + force_units=u"kJ * mol^-1 * nm^-1", + special=false, + args...) + if inter.shortcut(atom_i, atom_j) + return ustrip.(zero(dr)) * force_units end + σ = inter.σ_mixing(atom_i, atom_j) + ϵ = inter.ϵ_mixing(atom_i, atom_j) cutoff = inter.cutoff r2 = sum(abs2, dr) σ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, force_units) if special return f * dr * inter.weight_special else @@ -285,32 +250,25 @@ function force_divr(::LennardJonesSoftCore, r2, invr2, (σ2, ϵ, σ6_fac)) return ff * √invr2 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} - 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 - 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.ϵ) +@inline function potential_energy(inter::LennardJonesSoftCore, + dr, + atom_i, + atom_j, + energy_units=u"kJ * mol^-1", + special=false, + args...) + if inter.shortcut(atom_i, atom_j) + return ustrip(zero(dr[1])) * energy_units end + σ = inter.σ_mixing(atom_i, atom_j) + ϵ = inter.ϵ_mixing(atom_i, atom_j) cutoff = inter.cutoff r2 = sum(abs2, dr) σ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, energy_units) if special return pe * inter.weight_special else diff --git a/src/interactions/mie.jl b/src/interactions/mie.jl index 27b58615c..e68d53c6c 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, shortcut, σ_mixing, ϵ_mixing) The Mie generalized interaction between two atoms. @@ -15,14 +15,14 @@ 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{T, C, H, S, E} m::T n::T cutoff::C use_neighbors::Bool - lorentz_mixing::Bool - force_units::F - energy_units::E + shortcut::H + σ_mixing::S + ϵ_mixing::E mn_fac::T end @@ -31,33 +31,44 @@ function Mie(; n, 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) + 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), typeof(force_units), typeof(energy_units)}( - m_p, n_p, cutoff, use_neighbors, lorentz_mixing, force_units, energy_units, 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}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary) 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 - end +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 + +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 - # 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.ϵ) +function force(inter::Mie, + dr, + atom_i, + atom_j, + force_units=u"kJ * mol^-1 * nm^-1", + args...) + if inter.shortcut(atom_i, atom_j) + return ustrip.(zero(dr)) * force_units + end + σ = inter.σ_mixing(atom_i, atom_j) + ϵ = inter.ϵ_mixing(atom_i, atom_j) cutoff = inter.cutoff r2 = sum(abs2, dr) @@ -68,7 +79,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, force_units) return f * dr end @@ -76,20 +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}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary) 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 +@inline function potential_energy(inter::Mie, + dr, + atom_i, + atom_j, + energy_units=u"kJ * mol^-1", + 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) @@ -100,7 +108,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, energy_units) end function potential(::Mie, r2, invr2, (m, n, σ_r, const_mn)) diff --git a/src/interactions/morse_bond.jl b/src/interactions/morse_bond.jl index 56eab3d91..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 @@ -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..034809e38 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} @@ -51,6 +51,40 @@ 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 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) @@ -79,7 +113,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 +125,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 +142,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..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 @@ -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/interactions/soft_sphere.jl b/src/interactions/soft_sphere.jl index df3404c21..89921e9e0 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, shortcut, σ_mixing, ϵ_mixing) The soft-sphere potential. @@ -10,50 +10,42 @@ 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 - cutoff::C - use_neighbors::Bool - lorentz_mixing::Bool - force_units::F - energy_units::E +@kwdef struct SoftSphere{C, H, S, E} + cutoff::C = NoCutoff() + use_neighbors::Bool = false + shortcut::H = lj_zero_shortcut + σ_mixing::S = lorentz_σ_mixing + ϵ_mixing::E = geometric_ϵ_mixing 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) +use_neighbors(inter::SoftSphere) = inter.use_neighbors + +function Base.zero(ss::SoftSphere) + return SoftSphere(ss.cutoff, ss.use_neighbors, ss.shortcut, ss.σ_mixing, ss.ϵ_mixing) end -use_neighbors(inter::SoftSphere) = inter.use_neighbors +function Base.:+(s1::SoftSphere, ::SoftSphere) + return SoftSphere(s1.cutoff, s1.use_neighbors, s1.shortcut, s1.σ_mixing, s1.ϵ_mixing) +end -@inline function force(inter::SoftSphere{S, C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary) 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 +@inline function force(inter::SoftSphere, + dr, + atom_i, + atom_j, + force_units=u"kJ * mol^-1 * nm^-1", + 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) σ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, force_units) return f * dr end @@ -62,27 +54,24 @@ function force_divr(::SoftSphere, r2, invr2, (σ2, ϵ)) return (24ϵ * invr2) * 2 * six_term ^ 2 end -function potential_energy(inter::SoftSphere{S, C}, - dr, - coord_i, - coord_j, - atom_i, - atom_j, - boundary) 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 +function potential_energy(inter::SoftSphere, + dr, + atom_i, + atom_j, + energy_units=u"kJ * mol^-1", + 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) σ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, energy_units) end function potential(::SoftSphere, r2, invr2, (σ2, ϵ)) diff --git a/src/loggers.jl b/src/loggers.jl index 3a9a35ae8..62c85bc41 100644 --- a/src/loggers.jl +++ b/src/loggers.jl @@ -1,17 +1,19 @@ # Loggers to record properties throughout a simulation export - run_loggers!, + apply_loggers!, GeneralObservableLogger, values, log_property!, TemperatureLogger, - CoordinateLogger, - VelocityLogger, + CoordinatesLogger, + VelocitiesLogger, TotalEnergyLogger, KineticEnergyLogger, PotentialEnergyLogger, - ForceLogger, + ForcesLogger, + VolumeLogger, + DensityLogger, VirialLogger, PressureLogger, StructureWriter, @@ -22,8 +24,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,14 +34,12 @@ 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...) end - elseif run_loggers != false && run_loggers != :skipzero - throw(ArgumentError("run_loggers must be true, false or :skipzero")) end return sys end @@ -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,15 +113,15 @@ 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...) = 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}, @@ -128,23 +129,23 @@ 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 -velocities_wrapper(sys, neighbors; kwargs...) = sys.velocities +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}, @@ -152,15 +153,15 @@ 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 -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,21 +224,22 @@ 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 + return copy(current_forces) end 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}, @@ -245,15 +247,53 @@ 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 -virial_wrapper(sys, neighbors; n_threads, kwargs...) = virial(sys, neighbors; n_threads=n_threads) +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) + +""" + 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 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 +313,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 +502,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 +604,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/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/setup.jl b/src/setup.jl index e1ffb0a0d..ce32377a1 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,12 +27,13 @@ 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 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 @@ -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,12 +83,13 @@ 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 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 @@ -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 @@ -467,7 +469,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 +555,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 @@ -792,8 +793,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( @@ -801,18 +800,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))), - force_units=force_units, - energy_units=energy_units, + 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))), - force_units=force_units, - energy_units=energy_units, + coulomb_const=(units ? T(coulomb_const) : T(ustrip(coulomb_const))), ) end pairwise_inters = (lj, crf) @@ -892,7 +887,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), @@ -911,15 +905,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 @@ -944,7 +941,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, @@ -954,8 +951,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, ) @@ -982,7 +979,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}[] @@ -1045,11 +1042,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" @@ -1060,10 +1065,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" @@ -1142,26 +1146,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]))) @@ -1187,7 +1191,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 @@ -1214,17 +1218,13 @@ 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), solvent_dielectric=T(crf_solvent_dielectric), use_neighbors=true, weight_special=T(0.5), - coulomb_const=(units ? T(coulombconst) : T(ustrip(coulombconst))), - force_units=force_units, - energy_units=energy_units, + coulomb_const=(units ? T(coulomb_const) : T(ustrip(coulomb_const))), ) if isnothing(boundary) @@ -1234,7 +1234,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 @@ -1276,8 +1276,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), @@ -1296,15 +1294,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 @@ -1313,7 +1314,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, @@ -1322,8 +1323,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, ) @@ -1389,9 +1390,9 @@ function add_position_restraints(sys, sis = (sys.specific_inter_lists..., restraints) return System( atoms=deepcopy(sys.atoms), - coords=deepcopy(sys.coords), + coords=copy(sys.coords), boundary=deepcopy(sys.boundary), - velocities=deepcopy(sys.velocities), + velocities=copy(sys.velocities), atoms_data=deepcopy(sys.atoms_data), topology=deepcopy(sys.topology), pairwise_inters=deepcopy(sys.pairwise_inters), diff --git a/src/simulators.jl b/src/simulators.jl index 2b74e2b75..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. @@ -59,47 +57,53 @@ 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) - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) +@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) - 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") + 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; 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; 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 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 ", 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 @@ -129,19 +133,21 @@ 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) - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) +@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) 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) + 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) @@ -153,15 +159,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; 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,15 +181,17 @@ 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) - 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, - current_forces=forces_t) + apply_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, + current_forces=forces_t) end return sys end @@ -210,37 +220,44 @@ 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) - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) +@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) - 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) + 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; 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) @@ -251,8 +268,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 @@ -278,52 +295,51 @@ end StormerVerlet(; dt, coupling=NoCoupling()) = StormerVerlet(dt, coupling) -function simulate!(sys, - sim::StormerVerlet, - n_steps::Integer; - n_threads::Integer=Threads.nthreads(), - run_loggers=true) - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) +@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) - run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads) - coords_last = sys.coords + 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 + 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; 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) + apply_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, + current_forces=forces_t) end return sys end @@ -361,16 +377,21 @@ 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) - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) +@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) - 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) + 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 +399,27 @@ function simulate!(sys, end for step_n in 1:n_steps - forces_t = forces(sys, neighbors; 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 + 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 + 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) @@ -407,8 +431,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 @@ -428,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. @@ -454,12 +477,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" @@ -468,11 +491,15 @@ 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) - accels_t = accelerations(sys, neighbors; 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) + accels_t = forces_t ./ masses(sys) + noise = similar(sys.velocities) effective_dts = [sim.dt / count(c, sim.splitting) for c in sim.splitting] @@ -499,18 +526,19 @@ 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 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,)) + 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 @@ -518,29 +546,32 @@ 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 -function O_step!(sys, α_eff, σ_eff, rng, temperature, neighbors) - noise = random_velocities(sys, temperature; rng=rng) - sys.velocities = α_eff .* sys.velocities + σ_eff .* noise +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 A_step!(sys, dt_eff, neighbors) - sys.coords += sys.velocities * dt_eff - sys.coords = wrap_coords.(sys.coords, (sys.boundary,)) +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 + 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 .* accels_t return sys end -function B_step!(sys, dt_eff, acceleration_vector, compute_forces::Bool, - n_threads::Integer, neighbors) - if compute_forces - acceleration_vector .= accelerations(sys, neighbors; n_threads=n_threads) - end - sys.velocities += dt_eff * acceleration_vector +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 @@ -570,28 +601,35 @@ 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" 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) + 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) + accels_t = forces_t ./ masses(sys) + noise = similar(sys.velocities) for step_n in 1:n_steps - forces_t = forces(sys, neighbors; 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,)) + 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,)) if !iszero(sim.remove_CM_motion) && step_n % sim.remove_CM_motion == 0 remove_CM_motion!(sys) @@ -600,8 +638,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 @@ -638,39 +676,44 @@ 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" 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) + 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)) 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; 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,15 +724,17 @@ 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) - 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, - current_forces=forces_t) + apply_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads, + current_forces=forces_t) end return sys end @@ -706,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. @@ -814,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. @@ -906,6 +947,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, @@ -991,33 +1034,36 @@ 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) - for i in 1:n_steps - coords_old = copy(sys.coords) + coords_old = similar(sys.coords) + + for step_n in 1:n_steps + 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; 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, - 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, i, run_loggers; n_threads=n_threads, - current_potential_energy=E_old, success=false, - energy_rate=(E_old / (sys.k * sim.temperature))) + sys.coords .= coords_old + 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 + return sys end diff --git a/src/spatial.jl b/src/spatial.jl index be43acfdc..913dfff72 100644 --- a/src/spatial.jl +++ b/src/spatial.jl @@ -4,7 +4,8 @@ export CubicBoundary, RectangularBoundary, TriclinicBoundary, - box_volume, + volume, + density, box_center, scale_boundary, random_coord, @@ -22,7 +23,8 @@ export virial, pressure, molecule_centers, - scale_coords! + scale_coords!, + dipole_moment """ CubicBoundary(x, y, z) @@ -113,7 +115,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}} @@ -190,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 @@ -245,12 +241,38 @@ n_infinite_dims(b::TriclinicBoundary) = 0 n_infinite_dims(sys::System) = n_infinite_dims(sys.boundary) """ - box_volume(boundary) + volume(sys) + volume(boundary) + +Calculate the volume (3D) or area (2D) of a [`System`](@ref) or bounding box. -Calculate the volume of a 3D bounding box or the area of a 2D bounding box. +Returns infinite volume for infinite boundaries. """ -box_volume(b::Union{CubicBoundary, RectangularBoundary}) = prod(b.side_lengths) -box_volume(b::TriclinicBoundary) = b[1][1] * b[2][2] * b[3][3] +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] + +""" + 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) @@ -323,15 +345,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 """ @@ -603,15 +630,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)) @@ -661,7 +699,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) """ @@ -673,13 +710,13 @@ 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 @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; n_threads=Threads.nthreads()) Calculate the virial of a system or the virial resulting from a general interaction. @@ -692,36 +729,34 @@ 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) 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; 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} 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 +768,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], step_n) 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], step_n) end v += dot(f, dr) end @@ -750,11 +786,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], step_n) 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], step_n) end v += dot(f, dr) end @@ -764,12 +800,12 @@ function virial(sys::System{D, G, T}, neighbors_dev, pairwise_inters_nonl, 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. @@ -785,21 +821,19 @@ 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) 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)) - P = (NkT - (2 * vir) / D) / box_volume(sys.boundary) + vir = energy_remove_mol(virial(sys, neighbors, step_n; n_threads=n_threads)) + 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) @@ -820,7 +854,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) @@ -869,12 +902,11 @@ 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) 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 @@ -902,7 +934,16 @@ 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 + +""" + 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)) diff --git a/src/types.jl b/src/types.jl index a4a4c7b3f..ddde31f78 100644 --- a/src/types.jl +++ b/src/types.jl @@ -1,20 +1,19 @@ # Types export - PairwiseInteraction, - use_neighbors, - SpecificInteraction, InteractionList1Atoms, InteractionList2Atoms, InteractionList3Atoms, InteractionList4Atoms, Atom, - charge, mass, + charge, AtomData, MolecularTopology, NeighborList, System, + inject_gradients, + extract_parameters, ReplicaSystem, is_on_gpu, float_type, @@ -25,31 +24,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) @@ -208,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(; ) @@ -219,43 +229,54 @@ 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) +# 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 """ @@ -286,8 +307,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 """ @@ -701,6 +722,68 @@ 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) + extract_parameters!(params_dic, inter, ff) + end + + for inter in values(sys.specific_inter_lists) + extract_parameters!(params_dic, inter, ff) + end + + return params_dic +end + """ ReplicaSystem(; ) @@ -1273,6 +1356,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 +1371,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 +1393,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 """ @@ -1342,3 +1427,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 f16ec2f27..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) @@ -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") @@ -211,21 +187,17 @@ 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)) +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 -function check_force_units(fdr::AbstractArray, sys_force_units) - return check_force_units(unit(first(fdr)), sys_force_units) -end +check_force_units(F::SVector, force_units) = @inbounds check_force_units(F[1], force_units) -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_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 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 diff --git a/test/Project.toml b/test/Project.toml index 60e2d7bb8..a63d2fb9a 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -8,14 +8,13 @@ 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" +KernelDensity = "5ab0869b-81aa-558d-bb23-cbf5423bbe9b" 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/agent.jl b/test/agent.jl index 850d6e687..acbdc27d1 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 @@ -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,12 +47,11 @@ atom_i.status = recovered end end - return zero(coord_i) + return zero(vec_ij) 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), @@ -74,12 +71,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) @@ -101,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 eeb038ee6..b0fda9131 100644 --- a/test/basic.jl +++ b/test/basic.jl @@ -56,13 +56,13 @@ @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" @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") @@ -72,13 +72,13 @@ @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" @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") @@ -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"; @@ -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) @@ -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 @@ -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", ) @@ -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(energy_units="kcal/mol") - @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, @@ -497,7 +491,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/energy_conservation.jl b/test/energy_conservation.jl index 26c8f1f21..d64a4cf2d 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), @@ -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/gradients.jl b/test/gradients.jl new file mode 100644 index 000000000..059ea966a --- /dev/null +++ b/test/gradients.jl @@ -0,0 +1,452 @@ +@testset "Gradients" begin + inter = LennardJones() + boundary = CubicBoundary(5.0) + a1, a2 = Atom(σ=0.3, ϵ=0.5), Atom(σ=0.3, ϵ=0.5) + + 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, a1, a2, NoUnits) + 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) + grads = autodiff( + Reverse, + pe, + Active, + Active(dist), + ) + return -grads[1][1] + end + + dists = collect(0.2:0.01:1.2) + forces_direct = force_direct.(dists) + forces_grad = force_grad.(dists) + @test all(forces_direct .≈ forces_grad) +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 ), + ] + 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 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, gpu, parallel, forward, f32, obc2, gbn2, tol_σ, tol_r0) in runs + 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) + boundary = CubicBoundary(T(3.0)) + temp = T(1.0) + simulator = VelocityVerlet( + dt=T(0.001), + coupling=RescaleThermostat(temp), + ) + 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( + dist_cutoff=nb_cutoff, + solvent_dielectric=T(Molly.crf_solvent_dielectric), + use_neighbors=true, + coulomb_const=T(ustrip(Molly.coulomb_const)), + ) + 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)] + angles_inner = [HarmonicAngle(k=T(10.0), θ0=T(2.0)) for i in 1:15] + angles = InteractionList3Atoms( + 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=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( + 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=zero(T), σ=zero(T)) for i in 1:n_atoms] + if obc2 + imp_obc2 = ImplicitSolventOBC( + AT(atoms_setup), + [AtomData(element="O") for i in 1:n_atoms], + InteractionList2Atoms(bond_is, bond_js, nothing); + kappa=T(0.7), + use_OBC2=true, + ) + general_inters = (imp_obc2,) + elseif gbn2 + imp_gbn2 = ImplicitSolventGBN2( + AT(atoms_setup), + [AtomData(element="O") for i in 1:n_atoms], + InteractionList2Atoms(bond_is, bond_js, nothing); + kappa=T(0.7), + ) + general_inters = (imp_gbn2,) + else + general_inters = () + end + neighbor_finder = DistanceNeighborFinder( + eligible=AT(trues(n_atoms, n_atoms)), + n_steps=10, + dist_cutoff=T(1.5), + ) + 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( + set_runtime_activity(Forward), loss, Duplicated, + Duplicated(σ, one(T)), Const(r0), Duplicated(copy(coords), zero(coords)), + Duplicated(copy(velocities), zero(velocities)), const_args..., + )[1], + autodiff( + set_runtime_activity(Forward), loss, Duplicated, + Const(σ), Duplicated(r0, one(T)), Duplicated(copy(coords), zero(coords)), + Duplicated(copy(velocities), zero(velocities)), const_args..., + )[1], + ) + else + grad_enzyme = autodiff( + 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 + + grad_fd = ( + 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 + @info "$(rpad(name, 20)) - $(rpad(prefix, 2)) - FD $gfd, Enzyme $genz" + ztol = contains(name, "f32") ? 1e-8 : 1e-10 + @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(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 + end +end + +@testset "Differentiable protein" begin + function create_sys(gpu::Bool) + ff = MolecularForceField(joinpath.(ff_dir, ["ff99SBildn.xml", "his.xml"])...; units=false) + return System( + joinpath(data_dir, "6mrr_nowater.pdb"), + ff; + units=false, + gpu=gpu, + implicit_solvent="gbn2", + kappa=0.7, + ) + end + + EnzymeRules.inactive(::typeof(create_sys), args...) = nothing + + 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_forces_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, + ) + + fs = forces(sys; n_threads=n_threads) + return sum(sum.(abs, fs)) + 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, + ) + + 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( + "atom_C8_σ" => 0.33996695084235345, + "atom_C8_ϵ" => 0.4577296, + "atom_C9_σ" => 0.33996695084235345, + "atom_C9_ϵ" => 0.4577296, + "atom_CA_σ" => 0.33996695084235345, + "atom_CA_ϵ" => 0.359824, + "atom_CT_σ" => 0.33996695084235345, + "atom_CT_ϵ" => 0.4577296, + "atom_C_σ" => 0.33996695084235345, + "atom_C_ϵ" => 0.359824, + "atom_N3_σ" => 0.32499985237759577, + "atom_N3_ϵ" => 0.71128, + "atom_N_σ" => 0.32499985237759577, + "atom_N_ϵ" => 0.71128, + "atom_O2_σ" => 0.2959921901149463, + "atom_O2_ϵ" => 0.87864, + "atom_OH_σ" => 0.30664733878390477, + "atom_OH_ϵ" => 0.8803136, + "atom_O_σ" => 0.2959921901149463, + "atom_O_ϵ" => 0.87864, + "inter_CO_weight_14" => 0.8333, + "inter_GB_neck_cut" => 0.68, + "inter_GB_neck_scale" => 0.826836, + "inter_GB_offset" => 0.0195141, + "inter_GB_params_C_α" => 0.733756, + "inter_GB_params_C_β" => 0.506378, + "inter_GB_params_C_γ" => 0.205844, + "inter_GB_params_N_α" => 0.503364, + "inter_GB_params_N_β" => 0.316828, + "inter_GB_params_N_γ" => 0.192915, + "inter_GB_params_O_α" => 0.867814, + "inter_GB_params_O_β" => 0.876635, + "inter_GB_params_O_γ" => 0.387882, + "inter_GB_probe_radius" => 0.14, + "inter_GB_radius_C" => 0.17, + "inter_GB_radius_N" => 0.155, + "inter_GB_radius_O" => 0.15, + "inter_GB_radius_O_CAR" => 0.14, + "inter_GB_sa_factor" => 28.3919551, + "inter_GB_screen_C" => 1.058554, + "inter_GB_screen_N" => 0.733599, + "inter_GB_screen_O" => 1.061039, + "inter_LJ_weight_14" => 0.5, + "inter_PT_-/C/CT/-_k_1" => 0.0, + "inter_PT_-/C/N/-_k_1" => -10.46, + "inter_PT_-/CA/CA/-_k_1" => -15.167, + "inter_PT_-/CA/CT/-_k_1" => 0.0, + "inter_PT_-/CT/C8/-_k_1" => 0.64852, + "inter_PT_-/CT/C9/-_k_1" => 0.64852, + "inter_PT_-/CT/CT/-_k_1" => 0.6508444444444447, + "inter_PT_-/CT/N/-_k_1" => 0.0, + "inter_PT_-/CT/N3/-_k_1" => 0.6508444444444447, + "inter_PT_C/N/CT/C_k_1" => -0.142256, + "inter_PT_C/N/CT/C_k_2" => 1.40164, + "inter_PT_C/N/CT/C_k_3" => 2.276096, + "inter_PT_C/N/CT/C_k_4" => 0.33472, + "inter_PT_C/N/CT/C_k_5" => 1.6736, + "inter_PT_CT/CT/C/N_k_1" => 0.8368, + "inter_PT_CT/CT/C/N_k_2" => 0.8368, + "inter_PT_CT/CT/C/N_k_3" => 1.6736, + "inter_PT_CT/CT/N/C_k_1" => 8.368, + "inter_PT_CT/CT/N/C_k_2" => 8.368, + "inter_PT_CT/CT/N/C_k_3" => 1.6736, + "inter_PT_H/N/C/O_k_1" => 8.368, + "inter_PT_H/N/C/O_k_2" => -10.46, + "inter_PT_H1/CT/C/O_k_1" => 3.3472, + "inter_PT_H1/CT/C/O_k_2" => -0.33472, + "inter_PT_HC/CT/C4/CT_k_1" => 0.66944, + "inter_PT_N/CT/C/N_k_1" => 2.7196, + "inter_PT_N/CT/C/N_k_10" => 0.1046, + "inter_PT_N/CT/C/N_k_11" => -0.046024, + "inter_PT_N/CT/C/N_k_2" => -0.824248, + "inter_PT_N/CT/C/N_k_3" => 6.04588, + "inter_PT_N/CT/C/N_k_4" => 2.004136, + "inter_PT_N/CT/C/N_k_5" => -0.0799144, + "inter_PT_N/CT/C/N_k_6" => -0.016736, + "inter_PT_N/CT/C/N_k_7" => -1.06692, + "inter_PT_N/CT/C/N_k_8" => 0.3138, + "inter_PT_N/CT/C/N_k_9" => 0.238488, + ) + + platform_runs = [("CPU", false, false)] + if run_parallel_tests + push!(platform_runs, ("CPU parallel", false, true)) + end + if run_gpu_tests + push!(platform_runs, ("GPU", true, false)) + end + test_runs = ( + ("Energy", test_energy_grad, 1e-8), + ("Force" , test_forces_grad, 1e-8), + ) + if !running_CI + push!(test_runs, ("Sim", test_sim_grad, 1e-2)) + end + params_to_test = ( + #"inter_LJ_weight_14", + "atom_N_ϵ", + "inter_PT_C/N/CT/C_k_1", + "inter_GB_screen_O", + #"inter_GB_neck_scale", + ) + + for (test_name, test_fn, test_tol) in test_runs + 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( + 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), + ) + #@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 = copy(params_dic) + dic[param] = val + 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" + @test frac_diff < test_tol + end + end + end +end 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/protein.jl b/test/protein.jl index 9bf966b58..7153eb633 100644 --- a/test/protein.jl +++ b/test/protein.jl @@ -6,9 +6,10 @@ 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), ), 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 @@ -49,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, @@ -127,7 +129,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 +152,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, ) @@ -171,7 +173,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 @@ -181,7 +183,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 +211,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 +280,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/runtests.jl b/test/runtests.jl index 5a7d3ab54..00fc6654d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -6,11 +6,10 @@ import AtomsCalculators using AtomsCalculators.AtomsCalculatorsTesting import BioStructures # Imported to avoid clashing names using CUDA -import Enzyme +using Enzyme using FiniteDifferences -using ForwardDiff +using KernelDensity import SimpleCrystals -using Zygote using DelimitedFiles using LinearAlgebra @@ -23,8 +22,10 @@ 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" +elseif GROUP != "All" + error("Unrecognised test group, GROUP=$GROUP") end # Some CPU gradient tests give memory errors on CI @@ -35,8 +36,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 @@ -49,13 +50,15 @@ 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" else @warn "The GPU tests will not be run as a CUDA-enabled device is not available" end @@ -67,10 +70,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 -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 diff --git a/test/simulation.jl b/test/simulation.jl index dd952a808..085a70204 100644 --- a/test/simulation.jl +++ b/test/simulation.jl @@ -4,10 +4,10 @@ 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], + 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),), @@ -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), @@ -64,12 +64,12 @@ 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 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, @@ -83,12 +83,12 @@ end ), loggers=( temp=TemperatureLogger(100), - coords=CoordinateLogger(100), - vels=VelocityLogger(100), + coords=CoordinatesLogger(100), + vels=VelocitiesLogger(100), 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), @@ -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),), @@ -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) @@ -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),), @@ -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) @@ -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], @@ -262,7 +262,7 @@ end ), loggers=( temp=TemperatureLogger(10), - coords=CoordinateLogger(10), + coords=CoordinatesLogger(10), ), ) @@ -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, @@ -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") @@ -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, @@ -379,7 +379,7 @@ end ), loggers=( temp=TemperatureLogger(100), - coords=CoordinateLogger(100), + coords=CoordinatesLogger(100), energy=TotalEnergyLogger(100), ), ) @@ -387,12 +387,11 @@ 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), - 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, @@ -400,7 +399,7 @@ end ), loggers=( temp=TemperatureLogger(Float64, 100), - coords=CoordinateLogger(Float64, 100), + coords=CoordinatesLogger(Float64, 100), energy=TotalEnergyLogger(Float64, 100), ), force_units=NoUnits, @@ -431,17 +430,17 @@ 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") 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(),), - loggers=(coords=CoordinateLogger(100),), + loggers=(coords=CoordinatesLogger(100),), ) atom_selector(at, at_data) = at_data.atom_type == "A1" @@ -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, @@ -558,10 +552,10 @@ end pairwise_inters=(LennardJones(use_neighbors=true),), constraints=(cons,), neighbor_finder=neighbor_finder, - loggers=(coords=CoordinateLogger(10),), + loggers=(coords=CoordinatesLogger(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" @@ -583,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, @@ -657,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, @@ -795,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), ), @@ -838,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 @@ -857,12 +851,12 @@ 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( atoms=atoms, - coords=deepcopy(coords), + coords=copy(coords), boundary=boundary, pairwise_inters=(LennardJones(),), loggers=( @@ -902,7 +896,7 @@ end sys = System( atoms=AT(atoms), - coords=AT(deepcopy(coords)), + coords=AT(copy(coords)), boundary=boundary, pairwise_inters=(LennardJones(),), loggers=( @@ -947,9 +941,6 @@ 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) - VolumeLogger(n_steps) = GeneralObservableLogger(box_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) @@ -969,7 +960,7 @@ end sys = System( atoms=AT(atoms), - coords=AT(deepcopy(coords)), + coords=AT(copy(coords)), boundary=boundary, pairwise_inters=(LennardJones(),), loggers=( @@ -979,7 +970,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), ), ) @@ -994,8 +985,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) @@ -1019,9 +1010,6 @@ 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) - 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) barostat_test_set = ( @@ -1040,7 +1028,7 @@ end sys = System( atoms=AT(atoms), - coords=AT(deepcopy(coords)), + coords=AT(copy(coords)), boundary=boundary, pairwise_inters=(LennardJones(),), loggers=( @@ -1050,7 +1038,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), ), ) @@ -1089,11 +1077,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", @@ -1111,8 +1095,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...]) @@ -1186,14 +1171,14 @@ 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)) - atoms = CuArray([Atom(charge=f32 ? 0.0f0 : 0.0, mass=atom_mass, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", + 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) - atoms = [Atom(charge=f32 ? 0.0f0 : 0.0, mass=atom_mass, σ=f32 ? 0.2f0u"nm" : 0.2u"nm", + 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/test/zygote.jl b/test/zygote.jl deleted file mode 100644 index d1c261a07..000000000 --- a/test/zygote.jl +++ /dev/null @@ -1,446 +0,0 @@ -@testset "Gradients" begin - inter = LennardJones(force_units=NoUnits, energy_units=NoUnits) - boundary = CubicBoundary(5.0) - a1, a2 = Atom(σ=0.3, ϵ=0.5), Atom(σ=0.3, ϵ=0.5) - - 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) - return F[1] - 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, c1, c2, a1, a2, boundary) - end - return -grad[1] - end - - dists = collect(0.2:0.01:1.2) - forces_direct = force_direct.(dists) - forces_grad = force_grad.(dists) - @test all(forces_direct .≈ forces_grad) -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))) - end - - function test_simulation_grad(gpu::Bool, parallel::Bool, forward::Bool, f32::Bool, pis::Bool, - sis::Bool, obc2::Bool, gbn2::Bool) - 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 - simulator = VelocityVerlet( - dt=f32 ? 0.001f0 : 0.001, - coupling=RescaleThermostat(temp), - ) - 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, - force_units=NoUnits, - energy_units=NoUnits, - ) - 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))) - bond_js = gpu ? CuArray(Int32.(collect((1 + n_atoms ÷ 2):n_atoms))) : 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 = 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, - ) - 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], - 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, - ) - atoms_setup = [Atom(charge=f32 ? 0.0f0 : 0.0, σ=f32 ? 0.0f0 : 0.0) for i in 1:n_atoms] - if obc2 - imp_obc2 = ImplicitSolventOBC( - gpu ? CuArray(atoms_setup) : atoms_setup, - [AtomData(element="O") for i in 1:n_atoms], - InteractionList2Atoms(bond_is, bond_js, nothing); - kappa=(f32 ? 0.7f0 : 0.7), - use_OBC2=true, - ) - general_inters = (imp_obc2,) - elseif gbn2 - imp_gbn2 = ImplicitSolventGBN2( - gpu ? CuArray(atoms_setup) : atoms_setup, - [AtomData(element="O") for i in 1:n_atoms], - InteractionList2Atoms(bond_is, bond_js, nothing); - kappa=(f32 ? 0.7f0 : 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), - n_steps=10, - dist_cutoff=f32 ? 1.5f0 : 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)] - bonds = InteractionList2Atoms( - bond_is, - bond_js, - gpu ? CuArray(bonds_inner) : bonds_inner, - ) - cs = deepcopy(forward ? coords_dual : coords) - vs = deepcopy(forward ? velocities_dual : velocities) - - sys = System( - atoms=gpu ? CuArray(atoms) : atoms, - coords=gpu ? CuArray(cs) : cs, - boundary=boundary, - velocities=gpu ? CuArray(vs) : 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...) - 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], - ) - else - # Run once to setup - grad_zygote = gradient(f, σ, r0) - grad_zygote = gradient(f, σ, r0) - 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)) - if abs(gfd) < 1e-13 - @info "$(rpad(name, 20)) - $(rpad(prefix, 2)) - FD $gfd, Zygote $gzy" - 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) - else - frac_diff = abs(gzy - gfd) / abs(gfd) - @info "$(rpad(name, 20)) - $(rpad(prefix, 2)) - FD $gfd, Zygote $gzy, fractional difference $frac_diff" - @test frac_diff < tol - end - end - end -end - -@testset "Differentiable protein" begin - function create_sys(gpu::Bool) - ff = MolecularForceField(joinpath.(ff_dir, ["ff99SBildn.xml", "his.xml"])...; units=false) - return System( - joinpath(data_dir, "6mrr_nowater.pdb"), - ff; - units=false, - gpu=gpu, - implicit_solvent="gbn2", - kappa=0.7, - ) - end - - function test_energy_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, - ) - - return potential_energy(sys; n_threads=n_threads) - end - - return loss - end - - 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 - - return loss - 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 - - return loss - end - - params_dic = Dict( - "atom_C8_σ" => 0.33996695084235345, - "atom_C8_ϵ" => 0.4577296, - "atom_C9_σ" => 0.33996695084235345, - "atom_C9_ϵ" => 0.4577296, - "atom_CA_σ" => 0.33996695084235345, - "atom_CA_ϵ" => 0.359824, - "atom_CT_σ" => 0.33996695084235345, - "atom_CT_ϵ" => 0.4577296, - "atom_C_σ" => 0.33996695084235345, - "atom_C_ϵ" => 0.359824, - "atom_N3_σ" => 0.32499985237759577, - "atom_N3_ϵ" => 0.71128, - "atom_N_σ" => 0.32499985237759577, - "atom_N_ϵ" => 0.71128, - "atom_O2_σ" => 0.2959921901149463, - "atom_O2_ϵ" => 0.87864, - "atom_OH_σ" => 0.30664733878390477, - "atom_OH_ϵ" => 0.8803136, - "atom_O_σ" => 0.2959921901149463, - "atom_O_ϵ" => 0.87864, - "inter_CO_weight_14" => 0.8333, - "inter_GB_neck_cut" => 0.68, - "inter_GB_neck_scale" => 0.826836, - "inter_GB_offset" => 0.0195141, - "inter_GB_params_C_α" => 0.733756, - "inter_GB_params_C_β" => 0.506378, - "inter_GB_params_C_γ" => 0.205844, - "inter_GB_params_N_α" => 0.503364, - "inter_GB_params_N_β" => 0.316828, - "inter_GB_params_N_γ" => 0.192915, - "inter_GB_params_O_α" => 0.867814, - "inter_GB_params_O_β" => 0.876635, - "inter_GB_params_O_γ" => 0.387882, - "inter_GB_probe_radius" => 0.14, - "inter_GB_radius_C" => 0.17, - "inter_GB_radius_N" => 0.155, - "inter_GB_radius_O" => 0.15, - "inter_GB_radius_O_CAR" => 0.14, - "inter_GB_sa_factor" => 28.3919551, - "inter_GB_screen_C" => 1.058554, - "inter_GB_screen_N" => 0.733599, - "inter_GB_screen_O" => 1.061039, - "inter_LJ_weight_14" => 0.5, - "inter_PT_-/C/CT/-_k_1" => 0.0, - "inter_PT_-/C/N/-_k_1" => -10.46, - "inter_PT_-/CA/CA/-_k_1" => -15.167, - "inter_PT_-/CA/CT/-_k_1" => 0.0, - "inter_PT_-/CT/C8/-_k_1" => 0.64852, - "inter_PT_-/CT/C9/-_k_1" => 0.64852, - "inter_PT_-/CT/CT/-_k_1" => 0.6508444444444447, - "inter_PT_-/CT/N/-_k_1" => 0.0, - "inter_PT_-/CT/N3/-_k_1" => 0.6508444444444447, - "inter_PT_C/N/CT/C_k_1" => -0.142256, - "inter_PT_C/N/CT/C_k_2" => 1.40164, - "inter_PT_C/N/CT/C_k_3" => 2.276096, - "inter_PT_C/N/CT/C_k_4" => 0.33472, - "inter_PT_C/N/CT/C_k_5" => 1.6736, - "inter_PT_CT/CT/C/N_k_1" => 0.8368, - "inter_PT_CT/CT/C/N_k_2" => 0.8368, - "inter_PT_CT/CT/C/N_k_3" => 1.6736, - "inter_PT_CT/CT/N/C_k_1" => 8.368, - "inter_PT_CT/CT/N/C_k_2" => 8.368, - "inter_PT_CT/CT/N/C_k_3" => 1.6736, - "inter_PT_H/N/C/O_k_1" => 8.368, - "inter_PT_H/N/C/O_k_2" => -10.46, - "inter_PT_H1/CT/C/O_k_1" => 3.3472, - "inter_PT_H1/CT/C/O_k_2" => -0.33472, - "inter_PT_HC/CT/C4/CT_k_1" => 0.66944, - "inter_PT_N/CT/C/N_k_1" => 2.7196, - "inter_PT_N/CT/C/N_k_10" => 0.1046, - "inter_PT_N/CT/C/N_k_11" => -0.046024, - "inter_PT_N/CT/C/N_k_2" => -0.824248, - "inter_PT_N/CT/C/N_k_3" => 6.04588, - "inter_PT_N/CT/C/N_k_4" => 2.004136, - "inter_PT_N/CT/C/N_k_5" => -0.0799144, - "inter_PT_N/CT/C/N_k_6" => -0.016736, - "inter_PT_N/CT/C/N_k_7" => -1.06692, - "inter_PT_N/CT/C/N_k_8" => 0.3138, - "inter_PT_N/CT/C/N_k_9" => 0.238488, - ) - - platform_runs = [("CPU", [false, false])] - if run_parallel_tests - push!(platform_runs, ("CPU parallel", [false, true])) - end - if run_gpu_tests - 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)) - end - params_to_test = ( - "inter_LJ_weight_14", - "atom_N_ϵ", - "inter_PT_C/N/CT/C_k_1", - "inter_GB_screen_O", - "inter_GB_neck_scale", - ) - - 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 - for param in params_to_test - gzy = grads_zygote[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" - @test frac_diff < test_tol - end - end - end -end