Analysis functions for diatomic molecules. (#329)

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

-  Add `Structure` submodule for common utility functions for atomic structure. 

* Add rudimentary tests for NQCDynamics.Structure

* Improved method definitions, possible efficiency gain by using views

* Add unit tests for diatomic analysis functions for desorption

* Add "Analysis" group to unit testing

* Added explanation for Structure module scope

.github/workflows/CI.yml

@@ -12,6 +12,7 @@ jobs:
       matrix:
         group:
           - Core
+          - Analysis
           - InitialConditions
           - dynamics_classical
           - dynamics_mdef

.gitignore

@@ -7,4 +7,6 @@ dev/
 *.dat
 !/.gitignore
 !/.github
-!/.codecov.yml
+!/.codecov.yml
+!/test/artifacts/desorption_dynamics.jld2
+!/test/artifacts/desorption_test.xyz

Project.toml

@@ -1,7 +1,9 @@
 name = "NQCDynamics"
 uuid = "36248dfb-79eb-4f4d-ab9c-e29ea5f33e14"
 authors = ["James <james.gardner1421@gmail.com>"]
-version = "0.13.6"
+
+version = "0.13.7"
+
 
 [deps]
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
@@ -36,6 +38,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"
@@ -58,7 +61,6 @@ FastLapackInterface = "1, 2"
 HDF5 = "0.15, 0.16, 0.17"
 Interpolations = "0.13, 0.14"
 MuladdMacro = "0.2"
-MKL_jll = "2023" # Pin to before v2024.X to ensure tests build successfully. See #324
 NQCBase = "0.2"
 NQCDistributions = "0.1"
 NQCModels = "0.8"
@@ -92,6 +94,7 @@ DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JuLIP = "945c410c-986d-556a-acb1-167a618e0462"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 MKL = "33e6dc65-8f57-5167-99aa-e5a354878fb2"
@@ -103,4 +106,4 @@ Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "SafeTestsets", "CSV", "DataFrames", "DiffEqNoiseProcess", "FiniteDiff", "DiffEqDevTools", "Logging", "MKL", "PyCall", "JuLIP", "Symbolics", "Statistics", "Plots"]
+test = ["Test", "SafeTestsets", "CSV", "DataFrames", "DiffEqNoiseProcess", "FiniteDiff", "DiffEqDevTools", "Logging", "MKL", "PyCall", "JuLIP", "Symbolics", "Statistics", "Plots", "JLD2"]

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,26 @@
+```@setup logging
+@info "Expanding src/api/NQCDynamics/structure.md..."
+start_time = time()
+```
+# Structure
+This submodule contains utility functions to analyse and modify atomic structure data, such as interatomic distances, centres of mass, both with and without support for periodic boundary conditions. 
+
+These functions can be used to build more sophisticated output functions, or for basic analysis of simulation results in post. 
+
+This module **doesn't contain**:
+- Basic definitions of atomic structures (e.g. `Atoms`, `PeriodicCell`, ...). These are defined in [`NQCBase`](@ref).
+- Functions to generate atomic structures. These should be added to [`NQCDynamics.InitialConditions`](@ref).
+- Analysis functions for specific systems (e.g. molecules on surfaces). These should be added to [`NQCDynamics.Analysis`](@ref). 
+
+
+## Method reference
+
+
+```@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,146 @@
+"""
+Analysis functions for surface chemistry of diatomic molecules. 
+"""
+module Diatomic
+using NQCDynamics: AbstractSimulation, Simulation, get_positions, Structure
+using NQCBase
+using Unitful, UnitfulAtomic
+using LinearAlgebra
+using Statistics
+
+# Surface normal projection
+surface_normal_height(x::AbstractVector, surface_normal::AbstractVector=[0,0,1])=norm(x.*normalize(surface_normal))
+
+"""
+    surface_distance_condition(x::AbstractArray, indices::Vector{Int}, simulation::AbstractSimulation; surface_distance_threshold=5.0*u"Å")
+
+    Checks that the diatomic molecule is at least `surface_distance_threshold` away from the highest substrate atom in the simulation.
+"""
+function surface_distance_condition(x::AbstractArray, indices::Vector{Int}, simulation::AbstractSimulation; surface_distance_threshold=5.0*u"Å")
+        molecule_position=surface_normal_height(Structure.pbc_center_of_mass(x, indices..., simulation)) # molecule position wrt surface normal
+        substrate_heights=surface_normal_height.(eachcol(get_positions(x)[:, symdiff(1:length(simulation.atoms.masses), indices)])) # substrate positions along surface normal
+        highest_z=max(substrate_heights[substrate_heights.≤molecule_position]...) # Ignore substrate above molecule. 
+        #? surface_distance_threshold must be < vacuum above surface and > distance between substrate layers. 
+        if au_to_ang(molecule_position-highest_z)*u"Å"≥surface_distance_threshold
+            @debug "Surface distance condition evaluated true."
+            return true
+        else
+            return false
+        end
+end
+
+"""
+    com_velocity_condition(x::AbstractArray, indices::Vector{Int}, simulation::AbstractSimulation; surface_normal::AbstractVector=[0,0,1])
+
+    Evaluates true if the centre of mass velocity vector of the diatomic molecule points to the surface.
+"""
+function com_velocity_condition(x::AbstractArray, indices::Vector{Int}, simulation::AbstractSimulation; surface_normal::AbstractVector=[0,0,1])
+        if dot(Structure.velocity_center_of_mass(x, indices..., simulation), normalize(surface_normal))>0
+            return false
+        else
+            return true
+        end
+end
+
+"""
+    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::AbstractVector, diatomic_indices::Vector{Int}, simulation::AbstractSimulation;surface_normal::Vector = [0,0,1], surface_distance_threshold = 5.0*u"Å")
+    desorbed_frame=findfirst(surface_distance_condition.(trajectory, Ref(diatomic_indices), Ref(simulation); surface_distance_threshold=surface_distance_threshold))
+
+    if isa(desorbed_frame, Nothing)
+        @debug "No desorption event found."
+        return nothing
+    else
+        @debug "Desorption observed in snapshot $(desorbed_frame)"
+        leaving_surface_frame=findlast(com_velocity_condition.(view(trajectory, 1:desorbed_frame), Ref(diatomic_indices), Ref(simulation); surface_normal=surface_normal)) #ToDo testing views for memory efficiency, need to test time penalty. Also need to test if running on everything and findfirst-ing the Bool array is quicker. 
+        if leaving_surface_frame<0
+            @warn "Desorption event detected, but no direction change detected."
+            return nothing
+        else
+            return leaving_surface_frame
+        end
+    end
+end
+
+function get_desorption_angle(trajectory::AbstractVector, 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 = zeros(eltype(trajectory[1]), length(surface_normal), length(trajectory)-desorption_frame)
+    for i in 1:length(trajectory)-desorption_frame
+        com_velocities[:, i] .= Structure.velocity_center_of_mass(trajectory[i+desorption_frame], indices[1], indices[2], simulation)
+    end
+    average_velocity=mean(com_velocities;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