Merge branch 'gpu-constraint-docs-tests' of https://github.com/ejmeitz/Molly.jl into gpu-constraint-docs-tests

Author: ejmeitz

Project.toml

@@ -50,7 +50,7 @@ MollyKernelDensityExt = "KernelDensity"
 MollyPythonCallExt = "PythonCall"
 
 [compat]
-AcceleratedKernels = "0.4.2"
+AcceleratedKernels = "0.4.3"
 Atomix = "0.1, 1.0"
 AtomsBase = "0.5"
 AtomsCalculators = "0.2"

docs/src/documentation.md

@@ -1115,7 +1115,6 @@ The available loggers are:
 - [`GeneralObservableLogger`](@ref)
 - [`TemperatureLogger`](@ref)
 - [`CoordinatesLogger`](@ref)
-- [`DisplacementsLogger`](@ref)
 - [`VelocitiesLogger`](@ref)
 - [`TotalEnergyLogger`](@ref)
 - [`KineticEnergyLogger`](@ref)
@@ -1125,6 +1124,7 @@ The available loggers are:
 - [`DensityLogger`](@ref)
 - [`VirialLogger`](@ref)
 - [`PressureLogger`](@ref)
+- [`DisplacementsLogger`](@ref)
 - [`TrajectoryWriter`](@ref)
 - [`StructureWriter`](@ref)
 - [`TimeCorrelationLogger`](@ref)

src/Molly.jl

@@ -20,7 +20,6 @@ using FFTW
 using GPUArrays
 using Graphs
 using KernelAbstractions
-import KernelAbstractions as KA
 using NearestNeighbors
 using PeriodicTable
 using SimpleCrystals

src/constraints/constraints_kernel_helper.jl

@@ -2,10 +2,10 @@
 
 
 # No-op when backends are same
-to_backend(arr, old::T, new::T) where {T <: KA.Backend} = arr
+to_backend(arr, old::T, new::T) where {T <: Backend} = arr
 
 # Allocates and copies when backends are different
-function to_backend(arr, old::A, new::B) where {A <: KA.Backend, B <: KA.Backend}
+function to_backend(arr, old::A, new::B) where {A <: Backend, B <: Backend}
     out = allocate(new, eltype(arr), size(arr))
     copy!(out, arr)
     return out

src/constraints/shake_kernels.jl

@@ -369,7 +369,7 @@ function shake_gpu!(
     active_idxs .= 1:N_active_clusters
     # Doesnt need to be initialized, kernel will do that
     still_active = allocate(backend, Bool, N_active_clusters)
-    KA.pagelock!(backend, still_active)
+    KernelAbstractions.pagelock!(backend, still_active)
 
     iter = 1
     while iter <= max_iters

src/loggers.jl

@@ -16,15 +16,15 @@ export
     DensityLogger,
     VirialLogger,
     PressureLogger,
+    DisplacementsLogger,
     write_structure,
     TrajectoryWriter,
     StructureWriter,
     TimeCorrelationLogger,
     AutoCorrelationLogger,
     AverageObservableLogger,
     ReplicaExchangeLogger,
-    MonteCarloLogger,
-    DisplacementsLogger
+    MonteCarloLogger
 
 """
     apply_loggers!(system, neighbors=nothing, step_n=0, run_loggers=true;
@@ -301,7 +301,6 @@ function virial_wrapper(sys, neighbors, step_n; n_threads, kwargs...)
     return virial(sys, neighbors, step_n; n_threads=n_threads)
 end
 
-
 """
     VirialLogger(n_steps)
     VirialLogger(T, n_steps)
@@ -342,6 +341,58 @@ function Base.show(io::IO, pl::GeneralObservableLogger{T, typeof(pressure_wrappe
             pl.n_steps, ", ", length(values(pl)), " pressures recorded")
 end
 
+"""
+    DisplacementsLogger(n_steps, coords_start; n_steps_update::Integer=10)
+
+Log the displacements of atoms in a system throughout a simulation, useful for calculating
+properties like mean square displacement in periodic systems.
+
+Displacements are updated every `n_steps_update` steps and a copy is saved every
+`n_steps` steps.
+`coords_start` are the initial reference positions and should match the coordinate type
+in the system.
+
+It is assumed that a particle does not cross half the box size in `n_steps_update` steps. 
+By default `n_steps_update` is set to 10 to mitigate this assumption, but it can be 
+set to a higher value to reduce cost.
+`n_steps` must be a multiple of `n_steps_update`. 
+"""
+mutable struct DisplacementsLogger{D, C}
+    displacements::D
+    coords_ref::C
+    last_displacements::C
+    n_steps::Int
+    n_steps_update::Int
+end
+
+function DisplacementsLogger(n_steps::Integer, coords_start; n_steps_update::Integer=10)
+    if n_steps % n_steps_update != 0
+        throw(ArgumentError("DisplacementsLogger n_steps ($n_steps) must be a multiple of " *
+                            "n_steps_update ($n_steps_update)"))
+    end
+    displacements = typeof(from_device(coords_start))[]
+    return DisplacementsLogger(displacements, copy(coords_start), zero(coords_start),
+                               n_steps, n_steps_update)
+end
+
+Base.values(dl::DisplacementsLogger) = dl.displacements
+
+function log_property!(dl::DisplacementsLogger, sys::System, neighbors=nothing,
+                       step_n::Integer=0; kwargs...)     
+    if (step_n % dl.n_steps_update) == 0
+        dl.last_displacements .+= vector.(dl.coords_ref, sys.coords, (sys.boundary,))
+        dl.coords_ref .= sys.coords
+        if (step_n % dl.n_steps) == 0
+            push!(dl.displacements, Array(dl.last_displacements))
+        end
+    end
+end
+
+function Base.show(io::IO, dl::DisplacementsLogger)
+    print(io, "DisplacementsLogger with update n_steps ", dl.n_steps_update, ", saving n_steps ",
+            dl.n_steps, ", ", length(dl.displacements), " displacements recorded")
+end
+
 pdb_cryst1_length(l_Å) = lpad(round(l_Å; digits=3), 9)
 pdb_cryst1_angle(θ_rad) = lpad(round(rad2deg(θ_rad); digits=2), 7)
 
@@ -958,54 +1009,3 @@ function log_property!(mcl::MonteCarloLogger{T},
     push!(mcl.state_changed, success)
     push!(mcl.energy_rates, energy_rate)
 end
-
-
-"""
-    DisplacementsLogger(n_steps, r0; n_update::Integer=1, dims::Integer=3)
-
-Log the displacements of atoms in a system throughout a simulation. Displacements are
-updated every `n_update` steps and saved every `n_steps` steps. `r0` are the
-intitial refernce positions and should match the coords type in your `System` object.
-
-The logger assumes a particle does not cross 2 periodic boxes in `n_update` steps. 
-By default `n_update` is set to one to mitigate this assumption, but it can be 
-set to a higher value to reduce cost. `n_steps` must be a multiple of `n_update`. 
-"""
-mutable struct DisplacementsLogger{A, B}
-    displacements::Vector{A}
-    reference::Vector{B}
-    last_displacements::Vector{B}
-    n_steps::Int
-    n_update::Int
-end
-
-
-function DisplacementsLogger(n_steps::Integer, r0;  n_update::Integer = 1, dims::Integer = 3)
-    T = eltype(first(r0))
-    B = SArray{Tuple{dims}, T, 1, dims}
-    A = Array{B, 1}
-    if n_steps % n_update != 0
-        throw(ArgumentError("DisplacementsLogger: n_steps ($n_steps) must be a multiple n_update ($(n_update))"))
-    end
-    return DisplacementsLogger{A, B}(A[], copy(r0), zero(r0), n_steps, n_update)
-end
-
-Base.values(dl::DisplacementsLogger) = dl.displacements
-
-function log_property!(dl::DisplacementsLogger, s::System, neighbors=nothing,
-                        step_n::Integer=0; kwargs...)
-                        
-    if (step_n % dl.n_update) == 0
-        dl.last_displacements .+= vector.(dl.reference, s.coords, Ref(s.boundary))
-        dl.reference .= s.coords
-        if (step_n % dl.n_steps) == 0
-            push!(dl.displacements, copy(dl.last_displacements))
-        end
-    end
-end
-
-function Base.show(io::IO, dl::DisplacementsLogger)
-    print(io, "DisplacementsLogger with updating every ", dl.n_update, " steps, saving every ",
-            dl.n_steps, " steps with ", length(dl.displacements), " displacements in storage.")
-end
-

test/simulation.jl

@@ -370,12 +370,12 @@ end
         ),
         loggers=(
             coords=CoordinatesLogger(100),
-            disp=DisplacementsLogger(100, coords)
+            disp=DisplacementsLogger(100, coords),
         ),
 
     )
 
-    @test_throws ArgumentError DisplacementsLogger(100, coords; n_update = 17)
+    @test_throws ArgumentError DisplacementsLogger(100, coords; n_steps_update=17)
 
     if run_cuda_tests
         sys_gpu = System(
@@ -390,17 +390,19 @@ end
             ),
             loggers=(
                 coords=CoordinatesLogger(100),
-                disp=DisplacementsLogger(100, CuArray(coords))
+                disp=DisplacementsLogger(100, CuArray(coords)),
             ),
         )
     end
 
     for simulator in simulators
         @time simulate!(sys, simulator, n_steps; n_threads=1)
-        @test sum(sum(first(values(sys.loggers.disp)))) == 0.0u"nm"
+        @test all(isequal(0.0u"nm"), norm.(first(values(sys.loggers.disp))))
+        @test mean(norm.(sys.loggers.disp.displacements[end])) > 0.005u"nm"
         if run_cuda_tests
             @time simulate!(sys_gpu, simulator, n_steps; n_threads=1)
-            @test sum(first(values(sys_gpu.loggers.disp))) == 0.0u"nm"
+            @test all(isequal(0.0u"nm"), norm.(first(values(sys_gpu.loggers.disp))))
+            @test mean(norm.(sys.loggers.disp.displacements[end])) > 0.005u"nm"
             coord_diff = sys.coords .- from_device(sys_gpu.coords)
             coord_diff_size = sum(sum(map(x -> abs.(x), coord_diff))) / (3 * n_atoms)
             E_diff = abs(potential_energy(sys) - potential_energy(sys_gpu))