Skip to content

Commit

Permalink
simulators pass PE and forces to loggers
Browse files Browse the repository at this point in the history
  • Loading branch information
@jgreener64
jgreener64 committed Apr 23, 2024
1 parent 30bd4a2 commit 43f5a89
Show file tree
Hide file tree
Showing 2 changed files with 69 additions and 28 deletions.
25 changes: 22 additions & 3 deletions docs/src/documentation.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -889,6 +889,7 @@ 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)

# Find new neighbors like this
Expand Down Expand Up @@ -1064,14 +1065,32 @@ Running loggers can be disabled entirely with `run_loggers=false`, which is the
Loggers are currently ignored for the purposes of taking gradients, so if a logger is used in the gradient calculation the gradients will appear to be nothing.

Many times, a logger will just record an observation to an `Array` containing a record of past observations.
For this purpose, you can use the [`GeneralObservableLogger`](@ref) without defining a custom logging function. Simply define your observation function as
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)
function my_observable(sys::System, neighbors; n_threads::Integer, kwargs...)
# Probe the system for some desired property
return observation
end
```
A logger which records this property every `n_steps` can be constructed through
Keyword arguments `current_forces` and `current_potential_energy` can also be used here to avoid recomputing values that are passed from the simulator:
```julia
function my_pe_observable(sys::System, neighbors; n_threads::Integer,
current_potential_energy=nothing, kwargs...)
if isnothing(current_potential_energy)
# Compute potential energy
return potential_energy(sys, neighbors; n_threads=n_threads)
else
# Potential energy was passed from simulator, reuse
return current_potential_energy
end
end
```
These keyword arguments are also available in [`log_property!`](@ref).
Which values are passed depends on the simulator being used, for example [`SteepestDescentMinimizer`](@ref) passes `current_potential_energy` because it uses it for minimization.
Note that loggers are called after [`apply_coupling!`](@ref), so the coordinates may have changed since the potential energy or forces were computed.

A logger which records the property every `n_steps` can be constructed through
```julia
my_logger = GeneralObservableLogger(my_observable, T, n_steps)
```
Expand Down
72 changes: 47 additions & 25 deletions src/simulators.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,9 +65,9 @@ function simulate!(sys,
run_loggers=false)
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)
using_constraints = length(sys.constraints) > 0
E = potential_energy(sys, neighbors; n_threads=n_threads)
run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads, current_potential_energy=E)
using_constraints = length(sys.constraints) > 0
println(sim.log_stream, "Step 0 - potential energy ",
E, " - max force N/A - N/A")
hn = sim.step_size
Expand Down Expand Up @@ -98,7 +98,8 @@ function simulate!(sys,
E_trial, " - max force ", max_force, " - rejected")
end

run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads)
run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads,
current_potential_energy=E)

if max_force < sim.tol
break
Expand Down Expand Up @@ -136,9 +137,11 @@ function simulate!(sys,
sys.coords = wrap_coords.(sys.coords, (sys.boundary,))
!iszero(sim.remove_CM_motion) && remove_CM_motion!(sys)
neighbors = find_neighbors(sys, sys.neighbor_finder; n_threads=n_threads)
run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads)
accels_t = accelerations(sys, neighbors; 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)
accels_t_dt = zero(accels_t)
run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads, current_forces=forces_t)
using_constraints = length(sys.constraints) > 0
if using_constraints
cons_coord_storage = similar(sys.coords)
Expand All @@ -155,7 +158,8 @@ function simulate!(sys,
sim.dt; n_threads=n_threads)
sys.coords = wrap_coords.(sys.coords, (sys.boundary,))

accels_t_dt = accelerations(sys, neighbors; n_threads=n_threads)
forces_t_dt = forces(sys, neighbors; n_threads=n_threads)
accels_t_dt = forces_t_dt ./ masses(sys)

sys.velocities += ((accels_t .+ accels_t_dt) .* sim.dt / 2)
using_constraints && apply_velocity_constraints!(sys; n_threads=n_threads)
Expand All @@ -169,12 +173,15 @@ function simulate!(sys,
neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n, recompute_forces;
n_threads=n_threads)
if recompute_forces
accels_t = accelerations(sys, neighbors; n_threads=n_threads)
forces_t = forces(sys, neighbors; n_threads=n_threads)
accels_t = forces_t ./ masses(sys)
else
forces_t = forces_t_dt
accels_t = accels_t_dt
end

run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads)
run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads,
current_forces=forces_t)
end
return sys
end
Expand Down Expand Up @@ -218,7 +225,8 @@ function simulate!(sys,
end

for step_n in 1:n_steps
accels_t = accelerations(sys, neighbors; n_threads=n_threads)
forces_t = forces(sys, neighbors; n_threads=n_threads)
accels_t = forces_t ./ masses(sys)

sys.velocities += accels_t .* sim.dt

Expand All @@ -243,7 +251,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)
run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads,
current_forces=forces_t)
end
return sys
end
Expand Down Expand Up @@ -284,7 +293,8 @@ function simulate!(sys,
end

for step_n in 1:n_steps
accels_t = accelerations(sys, neighbors; n_threads=n_threads)
forces_t = forces(sys, neighbors; n_threads=n_threads)
accels_t = forces_t ./ masses(sys)

coords_copy = sys.coords
if using_constraints
Expand Down Expand Up @@ -312,7 +322,8 @@ function simulate!(sys,
n_threads=n_threads)
coords_last = coords_copy

run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads)
run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads,
current_forces=forces_t)
end
return sys
end
Expand Down Expand Up @@ -367,7 +378,8 @@ function simulate!(sys,
end

for step_n in 1:n_steps
accels_t = accelerations(sys, neighbors; n_threads=n_threads)
forces_t = forces(sys, neighbors; n_threads=n_threads)
accels_t = forces_t ./ masses(sys)

sys.velocities += accels_t .* sim.dt
apply_velocity_constraints!(sys; n_threads=n_threads)
Expand Down Expand Up @@ -395,7 +407,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)
run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads,
current_forces=forces_t)
end
return sys
end
Expand Down Expand Up @@ -573,7 +586,8 @@ function simulate!(sys,
run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads)

for step_n in 1:n_steps
accels_t = accelerations(sys, neighbors; n_threads=n_threads)
forces_t = forces(sys, neighbors; n_threads=n_threads)
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
Expand All @@ -586,7 +600,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)
run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads,
current_forces=forces_t)
end
return sys
end
Expand Down Expand Up @@ -632,9 +647,11 @@ function simulate!(sys, sim::NoseHoover, n_steps::Integer;
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)
forces_t = forces(sys, neighbors; n_threads=n_threads)
forces_t_dt = zero(forces_t)
accels_t = forces_t ./ masses(sys)
accels_t_dt = zero(accels_t)
run_loggers!(sys, neighbors, 0, run_loggers; n_threads=n_threads, current_forces=forces_t)
v_half = zero(sys.velocities)
zeta = zero(inv(sim.dt))

Expand All @@ -649,7 +666,8 @@ function simulate!(sys, sim::NoseHoover, n_steps::Integer;
T_half = uconvert(unit(sim.temperature), 2 * KE_half / (sys.df * sys.k))
zeta = zeta_half + (sim.dt / (2 * (sim.damping^2))) * ((T_half / sim.temperature) - 1)

accels_t_dt = accelerations(sys, neighbors; n_threads=n_threads)
forces_t_dt = forces(sys, neighbors; n_threads=n_threads)
accels_t_dt = forces_t_dt ./ masses(sys)

sys.velocities = (v_half .+ accels_t_dt .* (sim.dt / 2)) ./
(1 + (zeta * sim.dt / 2))
Expand All @@ -663,12 +681,15 @@ 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
accels_t = accelerations(sys, neighbors; n_threads=n_threads)
forces_t = forces(sys, neighbors; n_threads=n_threads)
accels_t = forces_t ./ masses(sys)
else
forces_t = forces_t_dt
accels_t = accels_t_dt
end

run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads)
run_loggers!(sys, neighbors, step_n, run_loggers; n_threads=n_threads,
current_forces=forces_t)
end
return sys
end
Expand Down Expand Up @@ -985,14 +1006,15 @@ function simulate!(sys::System{D, G, T},

ΔE = E_new - E_old
δ = ΔE / (sys.k * sim.temperature)
if δ < 0 || rand() < exp(-δ)

run_loggers!(sys, neighbors, i, run_loggers; n_threads=n_threads, success=true,
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)))
E_old = E_new
else
sys.coords = coords_old
run_loggers!(sys, neighbors, i, run_loggers; n_threads=n_threads, success=false,
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)))
end
end
Expand Down

0 comments on commit 43f5a89

Please sign in to comment.