Add Analysis submodule and diatomic molecule analysis functions.

- Added new DynamicsOutput: OutputDesorptionTrajectory #317
- Added new DynamicsOutput: OutputDesorptionAngle #317
- Created an `Analysis` submodule for common analysis functions

src/Analysis/Analysis.jl

@@ -0,0 +1,15 @@
+"""
+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
+@reexport using InitialConditions: 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
+using NQCBase
+using Unitful, UnitfulAtomic
+using LinearAlgebra
+
+"""
+    get_desorption_frame(trajectory::Vector, diatomic_indices::Vector{Int}, simulation::NQCDynamics.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::NQCDynamics.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(NQCDynamics.get_positions(x)[3, symdiff(1:end, diatomic_indices)]...)
+        if abs(au_to_ang(NQCBase.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(NQCBase.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::NQCDynamics.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 -> NQCBase.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 NQCBase.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 = NQCBase.Structure.fractional_mass(sim, index1, index2)
+    config[:,index2] .+= NQCBase.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

@@ -13,7 +13,8 @@ using ComponentArrays: ComponentVector
 
 using NQCDynamics:
     Estimators,
-    DynamicsUtils
+    DynamicsUtils,
+    Analysis
 
 using NQCModels: NQCModels
 
@@ -24,6 +25,10 @@ using .DynamicsUtils:
 
 using ..InitialConditions: QuantisedDiatomic
 
+using NQCBase
+using PyCall
+using Statistics
+
 """
     OutputVelocity(sol, i)
 
@@ -91,6 +96,20 @@ Evaluate the classical kinetic energy at each timestep during the trajectory.
 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
+
 OutputSpringEnergy(sol, i) = DynamicsUtils.classical_spring_energy.(sol.prob.p, sol.u)
 export OutputSpringEnergy
 
@@ -262,4 +281,45 @@ function (output::OutputStateResolvedScattering1D)(sol, i)
 end
 export OutputStateResolvedScattering1D
 
+"""
+Outputs the desorption angle in degrees (relative to the surface normal) if a desorption event is detected.
+"""
+struct OutputDesorptionAngle
+    indices::Vector{Int}
+    surface_normal::Vector
+    surface_distance_threshold
+end
+OutputDesorptionAngle(indices; surface_normal = [0,0,1], surface_distance_threshold = 5.0u"Å")=OutputDesorptionAngle(indices, surface_normal, surface_distance_threshold)
+
+"""
+    (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 NQCDynamics.Analysis.Diatomic.get_desorption_angle(sol.u, output.indices, sol.p; surface_normal=output.surface_normal, surface_distance_threshold=output.surface_distance_threshold)
+end
+
+"""
+Like OutputDynamicsVariables, but only saves parts of the trajectory where desorption is occurring.
+
+Use `extra_frames` to save additional steps before the desorption event begins. 
+"""
+struct OutputDesorptionTrajectory
+    indices::Vector{Int}
+    surface_normal::Vector
+    surface_distance_threshold
+    extra_frames::Int
+end
+OutputDesorptionTrajectory(indices; surface_normal = [0,0,1], surface_distance_threshold = 5.0u"Å", extra_frames = 0) = OutputDesorptionTrajectory(indices, 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 = NQCDynamics.Analysis.Diatomic.get_desorption_frame(sol.u, output.indices, sol.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
+
 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

src/NQCDynamics.jl

@@ -41,6 +41,9 @@ include("Ensembles/Ensembles.jl")
 export Ensembles
 @reexport using .Ensembles
 
+include("Analysis/Analysis.jl")
+export Analysis
+
 include("DynamicsOutputs.jl")
 @reexport using .DynamicsOutputs