QuantiseDiatomic updates (#300)

* Correct for periodicity in quantise_diatomic

* Split velocity array into slab & molecule as well

* Checking for PeriodicCell now type invariant

* Partial fix for #292

* Add debug info

* More debug info

* Possibly translating in the wrong direction

* More debug outputs

* Fix unrelated typo

* Improve accuracy of equilibrium bond length

* Fix missing import

* Add vectorised version of `quantise_diatomic`

`quantise_diatomic` is now able to process a single configuration
or a vector of multiple configurations.
This improves performance for multiple configurations,
as only one potential-specific binding curve is
generated to analyse all configurations.

* Delay rounding of J values until return

* Fix type definition

---------

Co-authored-by: Alexander Spears <alexander.spears@warwick.ac.uk>
Co-authored-by: James Gardner <james.gardner1421@gmail.com>

src/DynamicsOutputs.jl

@@ -64,7 +64,7 @@ export OutputPosition
 
 Output the position of the ring polymer centroid at each timestep during the trajectory.
 """
-OutputCentroidPosition(sol, i) = [get_centroid(get_position(u)) for u in sol.u]
+OutputCentroidPosition(sol, i) = [get_centroid(get_positions(u)) for u in sol.u]
 export OutputCentroidPosition
 
 """

src/InitialConditions/QuantisedDiatomic/QuantisedDiatomic.jl

@@ -27,7 +27,8 @@ using UnicodePlots: lineplot, lineplot!, DotCanvas, histogram
 using Optim: Optim
 using Roots: Roots
 using TimerOutputs: TimerOutputs, @timeit
-using Interpolations: interpolate, BSpline, Cubic, Line, OnGrid, scale, hessian
+using ProgressMeter
+using Interpolations: interpolate, BSpline, Cubic, Line, OnGrid, scale, hessian, knots
 
 using NQCDynamics: Simulation, Calculators, DynamicsUtils, masses
 using NQCBase: Atoms, PeriodicCell, InfiniteCell
@@ -43,9 +44,9 @@ include("energy_evaluation.jl")
 include("binding_curve.jl")
 include("random_configuration.jl")
 
-struct EffectivePotential{T,B,F}
+struct EffectivePotential{JType,T,B,F}
     μ::T
-    J::Int
+    J::JType
     binding_curve::BindingCurve{T,B,F}
 end
 
@@ -57,9 +58,9 @@ function (effective_potential::EffectivePotential)(r)
     return rotational + potential
 end
 
-struct RadialMomentum{T,B,F}
+struct RadialMomentum{JType,T,B,F}
     total_energy::T
-    V::EffectivePotential{T,B,F}
+    V::EffectivePotential{JType,T,B,F}
 end
 
 function (radial_momentum::RadialMomentum)(r)
@@ -251,42 +252,174 @@ function find_integral_bounds(total_energy::Real, V::EffectivePotential)
     return r₁, r₂
 end
 
+function binding_curve_from_structure(sim::Simulation, v::Matrix, r::Matrix; 
+    bond_lengths=0.5:0.01:5.0, 
+    height=10.0, 
+    surface_normal=[0, 0, 1.0], 
+    atom_indices=[1, 2], 
+    args...
+)
+    # Separate slab from molecule if necessary. 
+    r, slab = separate_slab_and_molecule(atom_indices, r)
+    v, slab_v = separate_slab_and_molecule(atom_indices, v)
+    environment = EvaluationEnvironment(atom_indices, size(sim), slab, austrip(height), surface_normal)
+
+    binding_curve = calculate_binding_curve(bond_lengths, sim.calculator.model, environment)
+    @debug "Finished generating binding curve for the given potential."
+
+    return binding_curve
+end
+
+"""
+    quantise_diatomic(sim::Simulation, v::Matrix, r::Matrix; 
+    bond_lengths=0.5:0.01:5.0, 
+    height=10.0, 
+    surface_normal=[0, 0, 1.0], 
+    atom_indices=[1, 2], 
+    max_translation=1,
+    show_timer=false, reset_timer=false
+    )
+
+Quantise the vibrational and rotational degrees of freedom for the specified
+positions and velocities.
+
+If the potential can be evaluated for the diatomic only, independent of position, 
+supplying a `Simulation` for just the diatomic will speed up evaluation. 
+
+When evaluating the potential, the molecule is moved to `height` in direction `normal_vector`.
+If the potential is independent of centre of mass position, this has no effect.
+Otherwise, be sure to modify these parameters to give the intended behaviour.
+
+If a `Simulation` with a `PeriodicCell` is supplied, periodic copies of the diatomic atoms 
+will be used if positions are close to cell boundaries. 
+Set `max_translation` to the radius of surrounding unit cells to search. 
+(e.g. 1 if positions are already wrapped around cell boundaries)
+"""
+function quantise_diatomic(sim::Simulation, v::Matrix, r::Matrix; 
+    bond_lengths=0.5:0.01:5.0, 
+    height=10.0, 
+    surface_normal=[0, 0, 1.0], 
+    atom_indices=[1, 2], 
+    args...
+)
+    binding_curve=binding_curve_from_structure(sim, v, r; bond_lengths=bond_lengths, height=height, surface_normal=surface_normal, atom_indices=atom_indices)
+    return quantise_diatomic(sim, v, r, binding_curve; height=height, surface_normal=surface_normal, atom_indices=atom_indices, args...)
+end
+
+"""
+    quantise_diatomic(sim::Simulation, v::Vector{Matrix}, r::Vector{Matrix};
+    bond_lengths=0.5:0.01:5.0, 
+    height=10.0, 
+    surface_normal=[0, 0, 1.0], 
+    atom_indices=[1, 2], 
+    show_timer=false, reset_timer=false
+    )
+
+Quantise the vibrational and rotational degrees of freedom of multiple atomic configurations 
+given as a vector of velocity matrices and a vector of position matrices. 
+
+If the potential can be evaluated for the diatomic only, independent of position, 
+supplying a `Simulation` for just the diatomic will speed up evaluation. 
+
+When evaluating the potential, the molecule is moved to `height` in direction `normal_vector`.
+If the potential is independent of centre of mass position, this has no effect.
+Otherwise, be sure to modify these parameters to give the intended behaviour.
+
+If a `Simulation` with a `PeriodicCell` is supplied, periodic copies of the diatomic atoms 
+will be used if positions are close to cell boundaries. 
+Set `max_translation` to the radius of surrounding unit cells to search. 
+(e.g. 1 if positions are already wrapped around cell boundaries)
+
+Specify `show_timer=true` for performance timings of the EBK quantisation process and
+`reset_timer=true` to see timings for each individual quantisation. 
+"""
+function quantise_diatomic(sim::Simulation, v::Vector{<: Matrix{<: Any}}, r::Vector{<: Matrix{<: Any}};
+    bond_lengths=0.5:0.01:5.0, 
+    height=10.0, 
+    surface_normal=[0, 0, 1.0], 
+    atom_indices=[1, 2], 
+    args...
+)
+    # Generate binding curve for all structures
+    binding_curve=binding_curve_from_structure(sim, v[1], r[1]; bond_lengths=bond_lengths, height=height, surface_normal=surface_normal, atom_indices=atom_indices)
+    ν, J=[],[]
+
+    # Quantise all configurations in the given vectors. 
+    @showprogress for index in eachindex(v)
+        try
+            result=quantise_diatomic(sim, v[index], r[index], binding_curve; height=height, surface_normal=surface_normal, atom_indices=atom_indices, args...)
+            push!(ν, result[1])
+            push!(J, result[2])
+        catch e
+            @warn "Quantisation was unsuccessful for configuration at $(index).\nThe final results vectors will show ν=-1 and J=-1 for this configuration.\n The error was: $(e)"
+            push!(ν, -1)
+            push!(J, -1)
+        end
+    end
+    return ν, J
+end
+
 """
-    quantise_diatomic(sim::Simulation, v::Matrix, r::Matrix;
-        height=10, normal_vector=[0, 0, 1])
+    quantise_diatomic(sim::Simulation, v::Matrix, r::Matrix, binding_curve::BindingCurve;
+    show_timer=false, reset_timer=false,
+    height=10, normal_vector=[0, 0, 1], atom_indices=[1,2], max_translation=1) 
+    )
 
 Quantise the vibrational and rotational degrees of freedom for the specified
-positions and velocities
+positions and velocities using the `BindingCurve` specified. 
+A binding curve will be automatically generated if you do not supply one. 
+
+If the potential can be evaluated for the diatomic only, independent of position, 
+supplying a `Simulation` for just the diatomic will speed up evaluation. 
 
 When evaluating the potential, the molecule is moved to `height` in direction `normal_vector`.
 If the potential is independent of centre of mass position, this has no effect.
 Otherwise, be sure to modify these parameters to give the intended behaviour.
+
+If a `Simulation` with a `PeriodicCell` is supplied, periodic copies of the diatomic atoms 
+will be used if positions are close to cell boundaries. 
+Set `max_translation` to the radius of surrounding unit cells to search. 
+(e.g. 1 if positions are already wrapped around cell boundaries)
 """
-function quantise_diatomic(sim::Simulation, v::Matrix, r::Matrix;
-    height=10.0, surface_normal=[0, 0, 1.0], atom_indices=[1, 2],
+function quantise_diatomic(sim::Simulation, v::Matrix, r::Matrix, binding_curve::BindingCurve;
     show_timer=false, reset_timer=false,
-    bond_lengths=0.5:0.01:5.0
+    max_translation=1, height=10.0, 
+    surface_normal=[0, 0, 1.0], atom_indices=[1, 2], 
 )
 
     reset_timer && TimerOutputs.reset_timer!(TIMER)
 
+    if isa(sim.cell,PeriodicCell)
+        # If the simulation used a `PeriodicCell`, translate `atom_indices` so they are at their minimum distance. (This is necessary if atoms were translated back into the original unit cell)
+        translations=[[i,j,k] for i in -max_translation:max_translation for j in -max_translation:max_translation for k in -max_translation:max_translation]
+        which_translation=argmin([norm(abs.(r[:,atom_indices[2]]-r[:,atom_indices[1]]+sim.cell.vectors*operation)) for operation in translations])
+        # Translate one atom for minimal distance. 
+        if translations[which_translation]!=[0,0,0]
+            r[:,atom_indices[2]].=r[:,atom_indices[2]]+sim.cell.vectors*translations[which_translation]
+            @debug "Using a periodic copy of atom "*string(atom_indices[end])*"  to bring it closer to atom "*string(atom_indices[begin])
+        end
+    end
+
     r, slab = separate_slab_and_molecule(atom_indices, r)
+    v, slab_v = separate_slab_and_molecule(atom_indices, v)
     environment = EvaluationEnvironment(atom_indices, size(sim), slab, austrip(height), surface_normal)
 
+    @debug "After PBC check and separation from slab, diatomic positions are:\n$(r)"
+
     r_com = subtract_centre_of_mass(r, masses(sim)[atom_indices])
     v_com = subtract_centre_of_mass(v, masses(sim)[atom_indices])
     p_com = v_com .* masses(sim)[atom_indices]'
 
-    k = DynamicsUtils.classical_kinetic_energy(sim, v_com)
+    k = DynamicsUtils.classical_kinetic_energy(masses(sim)[atom_indices], v_com)
     p = calculate_diatomic_energy(bond_length(r_com), sim.calculator.model, environment)
     E = k + p
 
     L = total_angular_momentum(r_com, p_com)
-    J = round(Int, (sqrt(1 + 4 * L^2) - 1) / 2) # L^2 = J(J+1)ħ^2
+    J = (sqrt(1+4*L^2) - 1) / 2 # L^2 = J(J+1)ħ^2
 
     μ = reduced_mass(masses(sim)[atom_indices])
 
-    binding_curve = calculate_binding_curve(bond_lengths, sim.calculator.model, environment)
+    @debug "k=$k\np=$p\nE=$E\nL=$L\nJ=$J\n"
 
     V = EffectivePotential(μ, J, binding_curve)
 
@@ -303,7 +436,8 @@ function quantise_diatomic(sim::Simulation, v::Matrix, r::Matrix;
     TimerOutputs.complement!(TIMER)
     show_timer && show(TIMER)
 
-    return round(Int, ν), J
+    @debug "Found ν=$ν"
+    return round(Int, ν), round(Int, J)
 end
 
 function quantise_1D_vibration(model::AdiabaticModel, μ::Real, r::Real, v::Real;

src/InitialConditions/QuantisedDiatomic/binding_curve.jl

@@ -1,3 +1,4 @@
+using Optim: Optim
 
 struct BindingCurve{T,B,F}
     bond_lengths::B
@@ -11,9 +12,8 @@ function calculate_binding_curve(
     bond_lengths::AbstractVector, model::AdiabaticModel, environment::EvaluationEnvironment
 )
     potential = calculate_diatomic_energy.(bond_lengths, model, environment) # Calculate binding curve
-    potential_minimum, index = findmin(potential)
-    equilibrium_bond_length = bond_lengths[index]
     fit = fit_binding_curve(bond_lengths, potential)
+    equilibrium_bond_length, potential_minimum = find_minimum(fit)
     return BindingCurve(bond_lengths, potential, equilibrium_bond_length, potential_minimum, fit)
 end
 
@@ -23,6 +23,12 @@ function fit_binding_curve(bond_lengths, binding_curve)
     return sitp
 end
 
+function find_minimum(fit)
+    grid = knots(fit)
+    results = Optim.optimize(fit, minimum(grid), maximum(grid), Optim.Brent())
+    return Optim.minimizer(results), Optim.minimum(results)
+end
+
 function calculate_force_constant(binding_curve)
     return hessian(binding_curve.fit, binding_curve.equilibrium_bond_length)[1]
 end