Analysis functions for diatomic molecules.

Project.toml

@@ -36,6 +36,7 @@ Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"

docs/make.jl

@@ -44,6 +44,7 @@ end
             "dynamicssimulations/dynamicssimulations.md"
             find_all_files("dynamicssimulations/dynamicsmethods")
         ]
+        "Outputs and Analysis" => find_all_files("output_and_analysis")
         "Examples" => find_all_files("examples")
         "integration_algorithms.md"
         "Developer documentation" => find_all_files("devdocs")

docs/src/api/NQCDynamics/analysis.md

@@ -0,0 +1,14 @@
+```@setup logging
+@info "Expanding src/api/NQCDynamics/analysis.md..."
+start_time = time()
+```
+# Analysis
+
+```@autodocs
+Modules=[NQCDynamics.Analysis, NQCDynamics.Analysis.Diatomic]
+```
+
+```@setup logging
+runtime = round(time() - start_time; digits=2)
+@info "...done after $runtime s."
+```

docs/src/api/NQCDynamics/structure.md

@@ -0,0 +1,14 @@
+```@setup logging
+@info "Expanding src/api/NQCDynamics/structure.md..."
+start_time = time()
+```
+# Structure
+
+```@autodocs
+Modules=[NQCDynamics.Structure]
+```
+
+```@setup logging
+runtime = round(time() - start_time; digits=2)
+@info "...done after $runtime s."
+```

docs/src/getting_started.md

@@ -238,7 +238,7 @@ For classical dynamics we also provide a timestep `dt` since we're using the
 
 The final keyword argument `output` is used to specify the quantities we want
 to save during the dynamics.
-A list of the available quantities can be found [here](@ref DynamicsOutputs).
+A list of the available quantities can be found [here](@ref NQCDynamics.DynamicsOutputs).
 
 !!! tip "Output format"
 

docs/src/output_and_analysis/intro.md

@@ -0,0 +1,42 @@
+# [Simulations outputs and analysis functions](@id output_and_analysis)
+
+NQCDynamics.jl is able to output a variety of common observables when running simulations with the `run_dynamics` function. Further output functions can be implemented easily as well. 
+
+In this section you will find an overview of all available output types, as well as explanations of some common analysis functions in the realm of surface chemistry which we have implemented in the package. 
+
+## DynamicsOutputs
+
+In many examples within this documentation, you have seen calls to `run_dynamics`:
+
+```julia
+ensemble = run_dynamics(sim, tspan, distribution;selection=1:20,
+    dt=0.1u"fs", output=OutputPosition, trajectories=20, callback=terminate)
+```
+
+Within `run_dynamics`, the `output` argument specifies the desired output values for each trajectory. `output` can either be given as a single function, or as a tuple of multiple output functions, for example:
+
+```julia
+output=OutputPosition # or
+output=(OutputPosition, OutputVelocity, OutputKineticEnergy)
+
+ensemble[3][:OutputPosition] # will output the positions at all timesteps in trajectory 3
+```
+
+Every output type is a function which can use the [`DynamicsVariables`](@ref) and [`Simulation`](@ref) values of the respective trajectory, allowing you to create custom output types of any kind. See the [developer documentation] for more information on how to implement a custom output type. 
+
+You can find an overview of all available output types in the [`DynamicsOutputs`](@ref NQCDynamics.DynamicsOutputs) API. 
+
+## Analysis functions
+
+The [Analysis](@ref Analysis) submodule in NQCDynamics contains functions commonly used in the analysis of trajectories to make the analysis of existing trajectories easier. 
+Ideally, most observable quantities could be implemented with a combination of [`DynamicsOutputs`](@ref NQCDynamics.DynamicsOutputs) and `Reduction` types, however we might want to data from existing ensemble simulations where re-running the entire set of trajectories is impractical. 
+
+As a result, most functions in the `Analysis` submodule are also implemented as a `DynamicsOutput`. 
+
+### Convenient functions for periodic structures
+[`NQCDynamics.Structure`](@ref Structure) contains several useful functions for periodic structures, such as `pbc_distance, pbc_center_of_mass`. 
+
+These functions take into account periodic copies of the atoms in question, returning the respective values for the closest set of periodic copies. 
+
+### Analysis of diatomic molecules
+[`NQCDynamics.Analysis.Diatomic`](@ref Analysis)

src/Analysis/Analysis.jl

@@ -0,0 +1,18 @@
+"""
+Analysis functions common enough to be included in the main package.
+"""
+module Analysis
+using Reexport: @reexport
+
+# Diatomic analysis functions @alexsp32
+include("diatomic.jl")
+export Diatomic
+
+# Rebinding of quantise_diatomic under Analysis. 
+using NQCDynamics: InitialConditions
+function quantise_diatomic(sim, v, r; args...)
+	InitialConditions.QuantisedDiatomic.quantise_diatomic(sim, v, r; args...)
+end
+export quantise_diatomic
+
+end

src/Analysis/diatomic.jl

@@ -0,0 +1,125 @@
+"""
+Analysis functions for surface chemistry of diatomic molecules. 
+"""
+module Diatomic
+using NQCDynamics: AbstractSimulation, Simulation, get_positions, Structure
+using NQCBase
+using Unitful, UnitfulAtomic
+using LinearAlgebra
+
+"""
+    get_desorption_frame(trajectory::Vector, diatomic_indices::Vector{Int}, simulation::AbstractSimulation;surface_normal::Vector=[0,0,1], surface_distance_threshold=5.0*u"Å")
+
+Determines the index in a trajectory where surface desorption begins. 
+
+This is evaluated using two conditions:
+
+1. In the trajectory, the diatomic must be `surface_distance_threshold` or further away from the highest other atom. (In `surface_normal` direction). 
+
+2. Desorption begins at the turning point of the centre of mass velocity component along `surface_normal`, indicating overall movement away from the surface. 
+"""
+function get_desorption_frame(trajectory::Vector, diatomic_indices::Vector{Int}, simulation::AbstractSimulation;surface_normal::Vector = [0,0,1], surface_distance_threshold = 5.0*u"Å")
+    # Two criteria: Distance between two atoms (wrt PBC) must be below `distance_threshold` and CoM velocity must be positive wrt surface normal vector. 
+    function surface_distance_condition(x)
+        highest_z = max(get_positions(x)[3, symdiff(1:end, diatomic_indices)]...)
+        if abs(au_to_ang(Structure.pbc_center_of_mass(x, diatomic_indices..., simulation)[3] - highest_z)) * u"Å" ≥ surface_distance_threshold
+            @debug "Surface distance condition evaluated true"
+            return true
+        else
+            return false
+        end
+    end
+    function com_velocity_condition(x)
+        if dot(Structure.velocity_center_of_mass(x, diatomic_indices..., simulation), normalize(surface_normal)) > 0
+            @debug "Normal velocity condition evaluated true"
+            return true
+        else
+            return false
+        end
+    end
+    desorbed_frame = findfirst(surface_distance_condition, trajectory)
+    if isa(desorbed_frame, Nothing)
+        @debug "No desorption event found. "
+        return nothing
+    else
+        @debug "Found desorption event in frame $(desorbed_frame)"
+        leaving_surface_frame = findfirst(!com_velocity_condition, reverse(trajectory[1:desorbed_frame]))
+        if isa(leaving_surface_frame, Nothing)
+            @debug "Couldn't find a sub-zero CoM velocity component relative to surface normal."
+            return nothing
+        else
+            return desorbed_frame - leaving_surface_frame
+        end
+    end
+end
+
+function get_desorption_angle(trajectory::Vector, indices::Vector{Int}, simulation::AbstractSimulation; surface_normal = [0,0,1],  surface_distance_threshold = 5.0*u"Å")
+    # First determine where the reaction occurred on the surface. 
+    desorption_frame = get_desorption_frame(trajectory, indices, simulation;surface_distance_threshold = surface_distance_threshold, surface_normal = surface_normal)
+    if isa(desorption_frame, Nothing)
+        @debug "No desorption event detected in trajectory"
+        return nothing
+    end
+    @debug "Desorption frame: $(desorption_frame)"
+    # Determine the average centre of mass velocity to decrease error due to vibration and rotation orthogonal to true translational component. 
+    com_velocities = map(x -> Structure.velocity_center_of_mass(x, indices[1], indices[2], simulation), trajectory[desorption_frame:end])
+    average_velocity = mean(cat(com_velocities...;dims=2);dims=2)
+    # Now convert into an angle by arccos((a•b)/(|a|*|b|))
+    return Structure.angle_between(vec(average_velocity), surface_normal)
+end
+
+# 6✖6 Cartesian to internal coordinate transformation. 
+function transform_to_internal_coordinates(to_transform::Matrix, config::Matrix, index1::Int, index2::Int, sim::Simulation)
+    U = transform_U(config, index1, index2, sim) # Generate transformation matrix
+    return transpose(U) * to_transform * U
+end
+
+function transform_from_internal_coordinates(to_transform::Matrix, config::Matrix, index1::Int, index2::Int, sim::Simulation)
+    U_inv = inv(transform_U(config, index1, index2, sim))
+    return transpose(U_inv) * to_transform * U_inv
+end
+
+function transform_r(config, index1::Int, index2::Int)
+    return sqrt(sum(mapslices(x->(x[2]-x[1])^2, config[:, [index1, index2]];dims=2)))
+end
+
+function transform_r1(config, index1::Int, index2::Int)
+    return transform_r(config[1:2, :], index1, index2)
+end
+
+
+"""
+    transform_U(config::Matrix, index1::Int, index2::Int, sim::Simulation)
+
+Builds diatomic Cartesian to internal coordinate transformation matrix as described in the SI of `10.1021/jacsau.0c00066`
+"""
+function transform_U(config::Matrix, index1::Int, index2::Int, sim::Simulation)
+    masses = Structure.fractional_mass(sim, index1, index2)
+    config[:,index2] .+= Structure.minimum_distance_translation(config, index1, index2, sim) # PBC wrap positions for correct transformation. 
+    r = transform_r(config, index1, index2)
+    r1 = transform_r1(config, index1, index2)
+    unity = LinearAlgebra.I(3)
+    U_i1 = vcat(
+        mapslices(x->(x[1]-x[2])/r, config[:, [index1, index2]]; dims=2),
+        mapslices(x->(x[2]-x[1])/r, config[:, [index2, index1]]; dims=2),
+    )
+    U_i2 = vcat(
+        mapslices(x->(x[2]-x[1])*(config[3, index2]-config[3, index1])/(r^2*r1), config[1:2, [index1, index2]]; dims=2),
+        -r1/r^2,
+        mapslices(x->(x[1]-x[2])*(config[3, index2]-config[3, index1])/(r^2*r1), config[1:2, [index2, index1]]; dims=2),
+        r1/r^2,
+    )
+    U_i3 = [
+        -(config[2, index1]-config[2, index2]),
+        (config[1, index2]-config[1, index1]),
+        0.0,
+        -(config[2, index2]-config[2, index1]),
+        (config[1, index1]-config[1, index2]),
+        0.0,
+    ] ./ (r1^2)
+    return hcat(U_i1, U_i2, U_i3, vcat(unity.*masses[1], unity.*masses[2]))
+end
+
+export get_desorption_frame, get_desorption_angle, transform_from_internal_coordinates, transform_to_internal_coordinates, transform_U
+
+end

src/DynamicsOutputs.jl

@@ -7,13 +7,15 @@
 module DynamicsOutputs
 
 using RingPolymerArrays: get_centroid
+using Unitful: @u_str
 using UnitfulAtomic: austrip
 using LinearAlgebra: norm
 using ComponentArrays: ComponentVector
 
 using NQCDynamics:
     Estimators,
-    DynamicsUtils
+    DynamicsUtils,
+    Analysis
 
 using NQCModels: NQCModels
 
@@ -24,6 +26,9 @@
 
 using ..InitialConditions: QuantisedDiatomic
 
+using NQCBase
+using Statistics
+
 """
     OutputVelocity(sol, i)
 
@@ -91,6 +96,21 @@
 OutputKineticEnergy(sol, i) = DynamicsUtils.classical_kinetic_energy.(sol.prob.p, sol.u)
 export OutputKineticEnergy
 
+"""
+    OutputCentroidKineticEnergy(sol, i)
+
+Evaluate the classical kinetic energy of a subset of the entire system at each save step. 
+
+The subset is defined by `OutputSubsetKineticEnergy(indices)`.
+"""
+struct OutputSubsetKineticEnergy{T}
+    indices::T
+end
+function (output::OutputSubsetKineticEnergy)(sol, i)
+    return map(x->DynamicsUtils.classical_kinetic_energy(sol.prob.p.atoms.masses[output.indices], x[:, output.indices]), [DynamicsUtils.get_velocities(i) for i in sol.u])
+end
+export OutputSubsetKineticEnergy
+
 OutputSpringEnergy(sol, i) = DynamicsUtils.classical_spring_energy.(sol.prob.p, sol.u)
 export OutputSpringEnergy
 
@@ -262,4 +282,63 @@
 end
 export OutputStateResolvedScattering1D
 
+"""
+Outputs the desorption angle in degrees (relative to the surface normal) if a desorption event is detected.
+"""
+struct OutputDesorptionAngle{I<:Vector{Int},S<:Vector{Float64},D}
+    indices::I
+    surface_normal::S
+    surface_distance_threshold::D
+end
+"""
+    OutputDesorptionAngle(indices; surface_normal = [0,0,1], surface_distance_threshold = 5.0u"Å")
+
+Outputs the desorption angle in degrees (relative to the surface normal) if a desorption event is detected. 
+Use `surface_normal` to define the direction "away" from the surface. Most commonly, this would be in positive z direction.
+
+A desorption is detected if the centre of mass of the molecule defined with `indices` is above `surface_distance_threshold` from the closest surface atom. 
+This is calculated with respect to `surface_normal` and will take into account periodic boundary conditions. 
+"""
+OutputDesorptionAngle(indices; surface_normal = [0,0,1], surface_distance_threshold = 5.0u"Å") = OutputDesorptionAngle(indices, convert(Vector{Float64},surface_normal), surface_distance_threshold)
+export OutputDesorptionAngle
+
+"""
+    (output::OutputDesorptionAngle)(sol, i)
+
+Outputs the desorption angle in degrees (relative to the surface normal) if a desorption event was detected.
+"""
+function (output::OutputDesorptionAngle)(sol, i)
+    return Analysis.Diatomic.get_desorption_angle(sol.u, output.indices, sol.prob.p; surface_normal=output.surface_normal, surface_distance_threshold=output.surface_distance_threshold)
+end
+
+struct OutputDesorptionTrajectory{I<:Vector{Int}, N<:Vector{Float64}, D, F<:Int}
+    indices::I
+    surface_normal::N
+    surface_distance_threshold::D
+    extra_frames::F
+end
+"""
+    `OutputDesorptionTrajectory(indices; surface_normal = [0,0,1], surface_distance_threshold = 5.0u"Å", extra_frames = 0)`
+
+Like OutputDynamicsVariables, but only saves parts of the trajectory where desorption is occurring.
+
+Use `surface_normal` to define the direction "away" from the surface. Most commonly, this would be in positive z direction. 
+
+Use `extra_frames` to save additional steps before the desorption event begins. 
+
+A desorption is detected if the centre of mass of the molecule defined with `indices` is above `surface_distance_threshold` from the closest surface atom. 
+This is calculated with respect to `surface_normal` and will take into account periodic boundary conditions. 
+"""
+OutputDesorptionTrajectory(indices; surface_normal = [0,0,1], surface_distance_threshold = 5.0u"Å", extra_frames = 0) = OutputDesorptionTrajectory(indices, convert(Vector{Float64},surface_normal), surface_distance_threshold, extra_frames)
+"""
+    (output::OutputDesorptionTrajectory)(sol, i)
+
+Only output parts of the trajectory where desorption is occurring.
+"""
+function (output::OutputDesorptionTrajectory)(sol, i)
+    desorption_frame = Analysis.Diatomic.get_desorption_frame(sol.u, output.indices, sol.prob.p; surface_distance_threshold=output.surface_distance_threshold, surface_normal=output.surface_normal)
+    return isa(desorption_frame, Int) ? sol.u[desorption_frame-output.extra_frames:end] : nothing
+end
+export OutputDesorptionTrajectory
+
 end # module

src/InitialConditions/QuantisedDiatomic/QuantisedDiatomic.jl

@@ -43,6 +43,7 @@ const TIMER = TimerOutputs.TimerOutput()
 include("energy_evaluation.jl")
 include("binding_curve.jl")
 include("random_configuration.jl")
+include("rigid_rotator.jl")
 
 struct EffectivePotential{JType,T,B,F}
     μ::T

src/InitialConditions/QuantisedDiatomic/rigid_rotator.jl

@@ -0,0 +1,29 @@
+"""
+    classical_translational_energy(config::Any, ind1::Union{Int, CartesianIndex}, ind2::Union{Int, CartesianIndex}, sim::Simulation)
+
+Returns the classical translational energy in Hartree
+"""
+function classical_translational_energy(config::Any, ind1::Union{Int, CartesianIndex}, ind2::Union{Int, CartesianIndex}, sim::Simulation)
+    return sum(sim.atoms.masses[[ind1, ind2]])/2*norm(velocity_center_of_mass(config, ind1, ind2, sim))^2
+end
+
+"""
+    classical_rotation_energy(J::Union{Int, CartesianIndex}, config::Any, ind1::Union{Int, CartesianIndex}, ind2::Union{Int, CartesianIndex}, sim::Simulation)
+
+Classical rotation energy of a rigid diatomic rotor in Hartree
+"""
+function classical_rotation_energy(J::Union{Int, CartesianIndex}, config::Any, ind1::Union{Int, CartesianIndex}, ind2::Union{Int, CartesianIndex}, sim::Simulation)
+    I = reduced_mass(sim, ind1, ind2) * austrip(pbc_distance(config, ind1, ind2, sim))^2 # Classical moment of inertia in atomic units. 
+    return J * (J+1) / (2*I)
+end
+
+"""
+    harmonic_vibration_energy(ν::Union{Int, CartesianIndex}, k::Float, ind1::Union{Int, CartesianIndex}=1, ind2::Union{Int, CartesianIndex}=2, sim::Simulation)
+
+Vibrational energy of a harmonic oscillator with the force constant k and vibrational level ν. 
+"""
+function harmonic_vibration_energy(ν::Union{Int, CartesianIndex}, k::Float64, ind1::Union{Int, CartesianIndex}, ind2::Union{Int, CartesianIndex}, sim::Simulation)
+    return (ν + 1/2)*sqrt(k/reduced_mass(sim, ind1, ind2))
+end
+
+export classical_translational_energy, classical_rotation_energy, harmonic_vibration_energy