Update AdvancedMH.jl wrapper for Monte Carlo sampling. (#325)

1. AdvancedMH or one of its dependencies may have updated a keyword from "initial_config" to "initial_params".
This caused NQCD to not deliver initial positions to the sampling chain correctly, leading to weird results for systems not centred around 0. 

2. Updated new keywords `movement_ratio` and `stop_ratio`, where `movement_ratio=1-stop_ratio` as the old `move_ratio` was confusing. Warning for deprecated kwargs was added and documentation updated with information about new keywords

* Drop compat for AdvancedMH.jl<0.8 for #320,#321

* Version bump for #325

Author: Alexsp32

Project.toml

@@ -1,7 +1,7 @@
 name = "NQCDynamics"
 uuid = "36248dfb-79eb-4f4d-ab9c-e29ea5f33e14"
 authors = ["James <[email protected]>"]
-version = "0.13.4"
+version = "0.13.5"
 
 [deps]
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
@@ -47,7 +47,7 @@ UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 
 [compat]
 AdvancedHMC = "0.5, 0.6"
-AdvancedMH = "0.6, 0.7, 0.8"
+AdvancedMH = "0.8"
 ComponentArrays = "0.11, 0.12, 0.13, 0.14, 0.15"
 DEDataArrays = "0.2"
 Dictionaries = "0.3"

docs/src/initialconditions/metropolishastings.md

@@ -14,7 +14,17 @@ random walk starting from an initial configuration.
 These are accepted or rejected based upon the Metropolis-Hastings criteria.
 The result is a Markov chain that samples the canonical distribution.
 
-## Example
+!!! Legacy version
+
+    Prior to the use of [`AdvancedMH.jl`](https://github.com/TuringLang/AdvancedMH.jl),
+    an alternative version of the algorithm was implemented that works for both classical
+    and ring polymer systems: `MetropolisHastings.run_monte_carlo_sampling(sim, R0, Δ, passes)`
+    
+    This is currently still included in the code but should be regarded as deprecated and
+    will likely be removed/combined with the [`AdvancedMH.jl`](https://github.com/TuringLang/AdvancedMH.jl)
+    version.
+
+## Example 1
 
 We can perform the sampling by setting up a classical simulation in the usual way and
 providing an appropriate initial configuration.
@@ -33,14 +43,30 @@ steps = 1e4
 step_size = Dict(:H=>1)
 ```
 
-Now we can run the sampling. The extra keyword argument `move_ratio` is used to specify
+Now we can run the sampling. The extra keyword argument `movement_ratio` is used to specify
 the fraction of the system moved during each Monte Carlo step.
 If we attempt to move the entire system at once, we can expect a very low acceptance ratio,
 whereas is we move only a single atom, the sampling will take much longer.
 You will likely have to experiment with this parameter to achieve optimal sampling.
+
+!!! note
+
+    Keyword arguments relating to how much of the system to sample were recently changed. The existing `move_ratio` and `internal_ratio` arguments are no longer used. 
+
+    Instead, there are now two options to specify how much of the system to move:
+
+    `movement_ratio`: Defines which fraction of the system to move. 1 moves the entire system. 
+
+    `stop_ratio`: Defines which fraction of the system *not* to move. 1 stops the entire system. 
+
+    `movement_ratio_internal`: Defines which proportion of ring polymer normal modes to perturb. 1 moves the entire system. 
+
+    `stop_ratio_internal`: Defines which proportion of ring polymer normal modes not to perturb. 1 stops the entire system. 
+ 
+
 ```@example mh
 using NQCDynamics.InitialConditions: ThermalMonteCarlo
-chain = ThermalMonteCarlo.run_advancedmh_sampling(sim, r0, steps, step_size; move_ratio=0.5)
+chain = ThermalMonteCarlo.run_advancedmh_sampling(sim, r0, steps, step_size; movement_ratio=0.5)
 ```
 
 Now that our sampling is complete we can evaluate the potential energy expectation value.
@@ -53,17 +79,10 @@ Estimators.@estimate potential_energy(sim, chain)
 sim.temperature / 2 * 5
 ```
 
-## Legacy version
 
-Prior to the use of [`AdvancedMH.jl`](https://github.com/TuringLang/AdvancedMH.jl),
-an alternative version of the algorithm was implemented that works for both classical
-and ring polymer systems.
-This is currently still included in the code but should be regarded as deprecated and
-will likely be removed/combined with the [`AdvancedMH.jl`](https://github.com/TuringLang/AdvancedMH.jl)
-version.
-
-Here, we use the legacy version to obtain a thermal distribution in a simple
-model system.
+## Example 2
+Here, we obtain a thermal distribution in a simple model system with some additional tweaks to 
+try and sample a larger configuration space.
 
 ```@setup monte
 using NQCDynamics
@@ -83,40 +102,37 @@ nothing # hide
 ```
 
 Then we have to specify the parameters for the Monte Carlo simulation and perform the sampling.
-`Δ` contains the step sizes for each of the species, `R0` the initial geometry and `passes` the
-number of monte carlo passes we perform (`passes*n_atoms` steps total).
+`Δ` contains the step sizes for each of the species, `R0` the initial geometry and `samples` the
+number of configurations we want to obtain. 
+
+*AdvancedMH.jl* provides some additional options to control the sampling process. To (hopefully) include
+ a larger configuration space in the final results, we set a `thinning` of 10, meaning that we only keep 
+ every 10th proposed configuration. In addition, we discard the first 100 samples, since our initial configuration might not 
+lie in the equilibrium. This is set with `discard_initial`. 
+Further explanations of the keyword arguments can be found in the *AbstractMCMC.jl* [documentation](https://turinglang.org/AbstractMCMC.jl/dev/api/#Common-keyword-arguments).
+
 ```@example monte
-Δ = Dict([(:N, 0.1), (:O, 0.1)])
+Δ = Dict(:N => 0.1, :O => 0.1)
 R0 = [1.0 0.0; 0.0 0.0; 0.0 0.0]
-passes = 1000
-output = InitialConditions.MetropolisHastings.run_monte_carlo_sampling(sim, R0, Δ, passes)
+samples = 1000
+output = InitialConditions.ThermalMonteCarlo.run_advancedmh_sampling(sim, R0, samples, Δ; movement_ratio=0.5, thinning=10, discard_initial=100)
 nothing # hide
 ```
 
-Output has three fields: the acceptance rates for each species and the energies and geometries
-obtained during sampling.
-```@repl monte
-output.acceptance
-```
-```@example monte
-plot(output.energy)
-xlabel!("Step") # hide
-ylabel!("Energy") # hide
-```
-
 We can calculate the distance between each atom and plot the bond length throughout the sampling.
 ```@example monte
 using LinearAlgebra
-plot([norm(R[:,1] .- R[:,2]) for R in output.R])
+plot([norm(R[:,1] .- R[:,2]) for R in output])
 xlabel!("Step") # hide
 ylabel!("Bond length") # hide
 ```
 
 The result of this simulation seamlessly interfaces with the `DynamicalDistribution`
-presented in the previous section and `output.R` can be readily passed to provide
+presented in the previous section and `output` can be readily passed to provide
 the position distribution.
 The Monte Carlo sampling does not include velocities but these can be readily
 obtained from the Maxwell-Boltzmann distribution.
+
 ```@setup logging
 runtime = round(time() - start_time; digits=2)
 @info "...done after $runtime s."

src/InitialConditions/ThermalMonteCarlo.jl

@@ -27,40 +27,68 @@ using NQCDynamics:
   masses
 
 """
-    run_advancedhmc_sampling(sim, r, steps, σ; move_ratio=0.0, internal_ratio=0.0)
+    run_advancedmh_sampling(sim, r, steps, σ; movement_ratio=nothing, movement_ratio_internal=nothing, kwargs...)
 
 Sample the configuration space for the simulation `sim` starting from `r`.
 
 Total number of steps is given by `steps` and `σ` is the dictionary of
 step sizes for each species.
 
-`move_ratio` defaults to `0.0` and denotes the fraction of system moved each step.
-If `move_ratio = 0`, every degree of freedom is moved at each step.
-If `move_ratio = 1`, then nothing will happen. Experiment with this parameter to achieve
-optimal sampling.
+`movement_ratio` denotes the fraction of system moved each step.
+`internal_ratio` works as for `movement_ratio` but for the internal modes of the ring polymer.
+For `movement_ratio = 0`, every degree of freedom is moved at each step, if `movement_ratio = 1`, then nothing will happen. 
 
-`internal_ratio` works as for `move_ratio` but for the internal modes of the ring polymer.
+If neither arguments are defined, default behaviour is to move one atom (and one ring polymer normal mode) per step on average. 
+
+Further kwargs are passed to `AdvancedMH.sample` to allow for [extra functionality](https://turinglang.org/AbstractMCMC.jl/dev/api/#Common-keyword-arguments).
 """
 function run_advancedmh_sampling(
   sim::AbstractSimulation,
   r,
   steps::Real,
   σ::Dict{Symbol,<:Real};
-  move_ratio=0.0,
-  internal_ratio=0.0
+  movement_ratio=nothing,
+  stop_ratio=nothing,
+  movement_ratio_internal=nothing,
+  stop_ratio_internal=nothing,
+  move_ratio=nothing, # deprecated
+  internal_ratio=nothing, # deprecated
+  kwargs...
 )
 
+  # Give a warning if move_ratio or internal_ratio are used
+  if move_ratio !== nothing || internal_ratio !== nothing
+    @warn "move_ratio and internal_ratio kwargs are deprecated and may be removed in future. More information: https://nqcd.github.io/NQCDynamics.jl/stable/initialconditions/metropolishastings/"
+    stop_ratio=move_ratio === nothing ? nothing : move_ratio
+    stop_ratio_internal=internal_ratio === nothing ? nothing : internal_ratio
+  end
+
+  # If no kwargs for system fraction to move are given, perturb one atom and one normal mode at a time
+  if movement_ratio===nothing && stop_ratio===nothing
+    @debug "No movement restriction for atoms provided, automatically setting to move one atom per step on average."
+    movement_ratio=1/size(sim)[2]
+  end
+
+  if movement_ratio_internal===nothing && stop_ratio_internal===nothing
+    @debug "No movement restriction for ring polymer normal modes provided, automatically setting to move one mode per step on average."
+    movement_ratio_internal=length(size(sim))==3 ? 1/size(sim)[3] : 0.0
+  end
+
+  # Set atom movement ratio by using whichever keyword is defined
+  stop_ratio = stop_ratio===nothing ? 1-movement_ratio : stop_ratio
+  stop_ratio_internal = stop_ratio_internal===nothing ? 1-movement_ratio_internal : stop_ratio_internal
+
   density = get_density_function(sim)
 
   density_model = AdvancedMH.DensityModel(density)
-  proposal = get_proposal(sim, σ, move_ratio, internal_ratio)
+  proposal = get_proposal(sim, σ, stop_ratio, stop_ratio_internal)
 
   sampler = AdvancedMH.MetropolisHastings(proposal)
 
   initial_config = reshape_input(sim, copy(r))
 
   chain = AdvancedMH.sample(density_model, sampler, convert(Int, steps);
-    init_params=initial_config)
+    initial_params=initial_config, kwargs...)
 
   return reshape_output(sim, chain)
 end
@@ -116,29 +144,29 @@ end
 
 struct ClassicalProposal{P,T,S<:Simulation} <: AdvancedMH.Proposal{P}
   proposal::P
-  move_ratio::T
+  stop_ratio::T
   sim::S
 end
 
 struct RingPolymerProposal{P,T,S<:RingPolymerSimulation} <: AdvancedMH.Proposal{P}
   proposal::P
-  move_ratio::T
-  internal_ratio::T
+  stop_ratio::T
+  stop_ratio_internal::T
   sim::S
 end
 
 const MolecularProposal{P,T,S} = Union{RingPolymerProposal{P,T,S},ClassicalProposal{P,T,S}}
 
-function ClassicalProposal(sim::Simulation, σ, move_ratio)
+function ClassicalProposal(sim::Simulation, σ, stop_ratio)
   proposals = Matrix{UnivariateDistribution}(undef, size(sim))
   for (i, symbol) in enumerate(sim.atoms.types) # Position proposals
     distribution = σ[symbol] == 0 ? Dirac(0) : Normal(0, σ[symbol])
     proposals[:, i] .= distribution
   end
-  ClassicalProposal(proposals[:], move_ratio, sim)
+  ClassicalProposal(proposals[:], stop_ratio, sim)
 end
 
-function RingPolymerProposal(sim::RingPolymerSimulation, σ, move_ratio, internal_ratio)
+function RingPolymerProposal(sim::RingPolymerSimulation, σ, stop_ratio, stop_ratio_internal)
   ωₖ = RingPolymers.get_matsubara_frequencies(nbeads(sim), sim.beads.ω_n)
   proposals = Array{UnivariateDistribution}(undef, size(sim))
   for i = 1:nbeads(sim)
@@ -159,12 +187,12 @@ function RingPolymerProposal(sim::RingPolymerSimulation, σ, move_ratio, interna
       end
     end
   end
-  RingPolymerProposal(proposals[:], move_ratio, internal_ratio, sim)
+  RingPolymerProposal(proposals[:], stop_ratio, stop_ratio_internal, sim)
 end
 
 function Base.rand(rng::Random.AbstractRNG, p::ClassicalProposal)
   result = map(x -> rand(rng, x), p.proposal)
-  result[Random.randsubseq(eachindex(result), p.move_ratio)] .= 0
+  result[Random.randsubseq(findall(x -> x!=0.0, result), p.stop_ratio)] .=0
   return result
 end
 
@@ -173,12 +201,12 @@ function Base.rand(rng::Random.AbstractRNG, p::RingPolymerProposal)
   reshaped_result = reshape(result, size(p.sim))
 
   # Zero some of the centroid moves
-  sequence = Random.randsubseq(CartesianIndices(size(p.sim)[1:2]), p.move_ratio)
+  sequence = Random.randsubseq(findall(x -> x != 0, CartesianIndices(size(p.sim)[1:2])), p.stop_ratio)
   reshaped_result[sequence, 1] .= 0
 
   # Zero some of the internal mode moves
   for i = 2:nbeads(p.sim)
-    sequence = Random.randsubseq(CartesianIndices((ndofs(p.sim), natoms(p.sim))), p.internal_ratio)
+    sequence = Random.randsubseq(findall(x -> x != 0, CartesianIndices((ndofs(p.sim), natoms(p.sim)))), p.stop_ratio_internal)
     reshaped_result[sequence, i] .= 0
   end