Running nonadiabatic MD in Julia

A look under the hood of the NQCD packages

Alex Spears

13.06.2024

Table of contents

• The maths behind MD: Solving differential equations
• Now where does NQCDynamics come in?
• How does NQCDynamics work?
• How do I scale dynamics up for many trajectories?
• How can I contribute to NQCD?
• Questions?

The maths behind MD: Solving differential equations

What is MD?

Differential equation that is solved for given starting conditions

Let’s look at example

Formulating a problem

Classical system Hamiltonian \[ \mathrm{H=\frac{\mathbf{P}^2}{2M}+V(\mathbf{R})} \]

Starting conditions \[ \vec{v}=\begin{bmatrix}0\\\vdots\\0\end{bmatrix}, \vec{r}=\begin{bmatrix}\vdots\end{bmatrix} \]

Integration scheme

Velocity Verlet integration

• Need an equation of motion for dynamics to perform
  • Classical simplest example
• Starting conditions to propagate DE from
• Numerical integration scheme since we can’t get analytical results.

Solving Differential equations in Julia

DifferentialEquations.jl

• Consistent structure for creating problems and solving them.

• Identical I/O for different propagation methods.

• Flexible, extendable.

Unitful.jl

• Units
• Unit conversion

So how to do this in Julia? SciML packages

• DifferentialEquations.jl provides rigid structure for solving any kind of differential equation problem with a consistent structure.
  • Useful because different MD methods need different integration schemes, but base problem stays similar.
• Unitful.jl good for translating physically sensible input units into atomic units for calculations.

How do we translate formulated problem into Julia code?

Example: defining an ODEProblem

Need to build an ODEProblem, which needs

starting conditions,

differential equation to solve and

any other parameters

Equation of motion (1D Harmonic oscillator)

function hamiltonian_1d(u; omega=omega_1, m=m1)
  return (0.5*omega^2*m*u.x[2]^2) + (1/2 * u.x[1]^2 * m)
end

Starting parameters

u_initial = ComponentVector(r=1.0, v=0.1)

Resulting ODE problem:

ode_problem = ODEProblem(
    hamiltonian_1d,
    u_initial,
    (0.0,1000.0), # Time span
    dt=0.1 # Time step
)

Example: Solving the ODEProblem

sol = solve(
    dynamical_problem1, # Problem
    VelocityVerlet(), # Integrator
    reltol=1e-8, # Numerical accuracy
    dt=0.1, # Time step
    # Any other arguments
)

• Yields a solution object with
  • all initial parameters
  • propagated variables
  • SDEs: random noise values

ODEProblem then passed on to solve function.

Takes an integrator (describes how to numerically integrate)

and other parameters.

The DifferentialEquations.jl pipeline

MD workflows may require:

• Parallelisation
• Callbacks
• Reductions

• Problem type contains all inputs to the problem
• Problem, integrator method and integration parameters passed to solve command
• Built in support for
  • parallelisation
  • callbacks e.g. stopping on a condition
  • reductions e.g. ensemble averaging
• Result is a solution object holding information about simulation(s)

Now where does NQCDynamics come in?

This was just one example, but we don’t want to define EoM and integrator to use ourselves every time.

The NQCD “ecosystem”

NQCD packages function as an interface around SciML packages.

Translation of atomic structure representation we can interpret to mathematical construct.

Choice of sensible integration scheme is handled for us.

Mathematical construct of solution translated back into the outputs we need.

NQCDynamics.jl

NQCDynamics.jl provides the same kind of rigid I/O structure for nonadiabatic MD methods.

• wraps around DifferentialEquations.jl to propagate dynamics
  • unit conversions
  • equations of motion
  • consistent inputs/outputs
• Handling of ensemble simulations for statistical distributions

Atomic structure: NQCBase.jl

• Atoms
• PeriodicCell
• IO functions:¹
  • convert_from_ase_atoms
  • convert_to_ase_atoms

• NQCBase contains very basics needed to describe atomic structures
• Atoms, unit cells and structure loading.

1. Documentation

Potential Energy Surfaces: NQCModels.jl

• Model
  • Defines a potential energy surface^{1 2}
  • Dimensionality: arbitrary (in principle)
  • Methods: potential(model, R), derivative(model, R) (and in-place versions)
• Interface to ASE for ML models / AIMD³

using NQCModels
model=AdiabaticASEModel(ase_structure)
NQCModels.potential(model, R)
NQCModels.derivative!(model, D, R)

NQCModels contains PESs and how to evaluate them. - Interface to ase for use with ML models. - Model type has energy and derivative - Dynamics coupling have friction methods

1. Overview

2. Analytical Model Library

3. Documentation

Model types implemented in NQCModels.jl

Statistical sampling: NQCDistributions.jl

• DynamicalDistribution
  • Container for initial conditions (positions, velocities)
• Various distribution functions (e.g VelocityBoltzmann)

NQCDistributions enables statistical sampling from distributions. - e.g. Boltzmann distribution of velocities - Various 1D distributions, e.g. Wigner

How does NQCDynamics work?

Running an MD simulation (1/2)

Atoms

atoms = Atoms([:H])

• Chemical symbols
• or arbitrary masses ([:X] atoms)

Initial conditions

initial_conditions = DynamicsVariables(
    v = hcat([1.0]),
    r = hcat([0.0]),
)

Potential energy surface

model = Harmonic(m = 0.4, ω = 2.0, r₀ = 0.1)

Combining them in a Simulation

sim = Simulation{Classical}(atoms, model)

Simulations¹

1sim = Simulation{Method}(
2    atoms::Atoms{T},
3    model::Model;
4    temperature=0u"K",
5    cell::AbstractCell=InfiniteCell()
)

1

Defines the MD method to use: Classical, MDEF, Ehrenfest, …

2

Atoms contained in simulation

3

PES to use

4

Additional arguments for MD method

5

Additional structural parameters.

• Simulations contain a Calculator to cache energies, forces,…
• Containers for simulation parameters

1. NQCD Docs: Getting Started

Running an MD simulation (2/2)

simulation_output = run_dynamics(
    sim,
    (0,1000),
    initial_conditions;
    output = (OutputPosition, OutputVelocity),
    dt = 0.1,
    trajectories = 100, 
    selection = 1:100, 
    reduction = MeanReduction(),
    ensemble_algorithm = SciMLBase.EnsembleDistributed(),
    callback = CellBoundaryCallback(),
)

• run_dynamics creates the O/SDE problem for DifferentialEquations.jl to solve¹
• Outputs are processed from solution objects

1. NQCD Docs: Ensemble Simulations

Dynamics Parameters

simulation_output = run_dynamics(
    sim,
    (0,1000),
    initial_conditions;
    output = (OutputPosition, OutputVelocity),
    dt = 0.1,
    trajectories = 100, 
    selection = 1:100, 
    reduction = MeanReduction(),
    ensemble_algorithm = SciMLBase.EnsembleDistributed(),
    callback = CellBoundaryCallback(),
)

• Simulation time span (in a.u.)

Dynamics Parameters

simulation_output = run_dynamics(
    sim,
    (0,1000),
    initial_conditions;
    output = (OutputPosition, OutputVelocity),
    dt = 0.1,
    trajectories = 100, 
    selection = 1:100, 
    reduction = MeanReduction(),
    ensemble_algorithm = SciMLBase.EnsembleDistributed(),
    callback = CellBoundaryCallback(),
)

• Initial conditions:
  • Either DynamicsVariables to start with a single set of configurations.
  • Or a DynamicalDistribution to sample from.

Dynamics Parameters

simulation_output = run_dynamics(
    sim,
    (0,1000),
    initial_conditions;
    output = (OutputPosition, OutputVelocity),
    dt = 0.1,
    trajectories = 100, 
    selection = 1:100, 
    reduction = MeanReduction(),
    ensemble_algorithm = SciMLBase.EnsembleDistributed(),
    callback = CellBoundaryCallback(),
)

DynamicsOutputs

• Output functions that act on a DiffEQ solution
• e.g. positions, velocities, populations, final quantum states, …
• NQCD Docs: DynamicsOutputs

Dynamics Parameters

simulation_output = run_dynamics(
    sim,
    (0,1000),
    initial_conditions;
    output = (OutputPosition, OutputVelocity),
    dt = 0.1,
    trajectories = 100, 
    selection = 1:100, 
    reduction = MeanReduction(),
    ensemble_algorithm = SciMLBase.EnsembleDistributed(),
    callback = CellBoundaryCallback(),
)

• Simulation time step
• Snapshot every time step unless saveat= is specified
  • Can be a number for time intervals
  • or a Vector of times to save at

Dynamics Parameters

simulation_output = run_dynamics(
    sim,
    (0,1000),
    initial_conditions;
    output = (OutputPosition, OutputVelocity),
    dt = 0.1,
    trajectories = 100, 
    selection = 1:100, 
    reduction = MeanReduction(),
    ensemble_algorithm = SciMLBase.EnsembleDistributed(),
    callback = CellBoundaryCallback(),
)

• Number of trajectories to simulate
• selection=nothing for random sampling, otherwise indices of distribution to sample from.
• Reductions applied to outputs
  • Default: Append
  • MeanReduction for ensemble Mean
  • FileReduction

Dynamics Parameters

simulation_output = run_dynamics(
    sim,
    (0,1000),
    initial_conditions;
    output = (OutputPosition, OutputVelocity),
    dt = 0.1,
    trajectories = 100, 
    selection = 1:100, 
    reduction = MeanReduction(),
    ensemble_algorithm = SciMLBase.EnsembleDistributed(),
    callback = CellBoundaryCallback(),
)

Parallelisation strategy

EnsembleSerial: One trajectory at a time.

EnsembleDistributed: One trajectory per process. No shared memory.

EnsembleThreads: One trajectory per thread. Shared memory.

EnsembleSplitThreads: One trajectory per thread per process.

Dynamics Parameters

simulation_output = run_dynamics(
    sim,
    (0,1000),
    initial_conditions;
    output = (OutputPosition, OutputVelocity),
    dt = 0.1,
    trajectories = 100, 
    selection = 1:100, 
    reduction = MeanReduction(),
    ensemble_algorithm = SciMLBase.EnsembleDistributed(),
    callback = CellBoundaryCallback(),
)

Callbacks¹:

• Code which is executed on dynamics as they are propagated.
  • e.g. terminate simulations if a condition is satisfied.

1. NQCD Docs: Callbacks

How do I scale dynamics up for many trajectories?

Parallelisation approaches

Multithreading

• Shared memory (lower resource usage, IO)
• Python hates this (global interpreter lock)
• Need to avoid data races

Threads.@threads for i = 1:10
   a[i] = Threads.threadid()
end

Trivial taskfarming

• No shared memory (higher resource usage, IO)
• Fewer issues with data races

some_vector = Distributed.pmap(a_function, a_vector)

Parallelisation in NQCDynamics.jl through EnsembleAlgorithms

simulation_output = run_dynamics(
    sim,
    (0,1000),
    initial_conditions;
    output = (OutputPosition, OutputVelocity),
    dt = 0.1,
    trajectories = 100, 
    selection = 1:100, 
    reduction = MeanReduction(),
    ensemble_algorithm = SciMLBase.EnsembleDistributed(),
    callback = CellBoundaryCallback(),
)

EnsembleSerial: One trajectory at a time.

EnsembleDistributed: One trajectory per process. No shared memory.

EnsembleThreads: One trajectory per thread. Shared memory.

EnsembleSplitThreads: One trajectory per thread per process.

Parallelisation with ClusterScripts.jl¹

• Designed for simulations across multiple combinations of parameters.
• Trivial taskfarming
• Temporary file storage, functions to concatenate back into input parameter space.

fixed_parameters=Dict(
    "task" => "mdef+2tm",
    "trajectories" => 10000, 
    "runtime" => 4.7u"ps", 
    "timestep" => 0.1u"fs",
    "ensemble_algorithm" => EnsembleSerial(),
    "outputs" => (OutputInitial, OutputFinal),
    "friction_atoms" => friction_atoms,
)
variables=Dict(
    "starting_temperature" => [100, 150, 200, 250, 300],
    "fluence" => [10,20,40,60,80,100,120],
)
job_queue=build_job_queue(fixed_parameters, variables, postprocess_queue)
serialise_queue!(job_queue; filename="simulation_parameters.jld2")

1. GitHub

How can I contribute to NQCD?

NQCDynamics.jl and Submodules

• Different packages to compartmentalise categories of functionality
  • Atomic structure
  • PESs
  • …
• Submodules to isolate functional components
  • Simulations
  • Calculators
  • DynamicsOutputs
  • …

NQCD package and submodule structure a bit complicated when making changes for first time.

Feel free to talk to me since I’ve touched many parts of the codebase recently.

Use #julia_jl channel on Slack to talk about things.

Will briefly talk about most likely changes / additions.

A new type of output

DynamicsOutputs

src/DynamicsOutputs.jl

OutputKineticEnergy(sol, i) = DynamicsUtils.classical_kinetic_energy.(sol.prob.p, sol.u)
export OutputKineticEnergy

• Functions that take DiffEq solution object
• DifferentialEquations.jl: Solution Handling

struct OutputQuantisedDiatomic{S,H,V}
    sim::S
    height::H
    normal_vector::V
end
export OutputQuantisedDiatomic

function (output::OutputQuantisedDiatomic)(sol, i)
    final = last(sol.u) 
    ν, J = QuantisedDiatomic.quantise_diatomic(output.sim,
        DynamicsUtils.get_velocities(final), DynamicsUtils.get_positions(final);
        height=output.height, normal_vector=output.normal_vector)
    return (ν, J)
end

DynamicsOutputs act on a DiffEq solution object

• Can be either functions
• or Structs with functors

I’m adding an OutputSolution output to help develop new DynamicsOutputs.

Analysis methods for post-processing

src/structure.jl

• Basic functions for atomic structure, e.g. distance across PBC, centre of mass

src/Analysis/Analysis.jl

• Submodules for analysis type, e.g. functions for diatomic molecules in diatomic.jl.

Recent addition to NQCDynamics.jl:

• Basic structure functions to add more interesting stuff to.
• Analysis submodule to encompass anything needed for postprocessing.
  • e.g. Diatomic final state quantisation.

A new model?

A model in NQCModels.jl needs to implement:

• ndofs(m::Model)
• potential(m::Model, R::AbstractArray)
• derivative(m::Model, R::AbstractArray)

Adiabatic models: - Single electronic state

Diabatic models: \[ \mathbf{V(R)}=\begin{bmatrix}V_1 & \Lambda_{12}\\\Lambda_{21} & V_2\end{bmatrix} \]

• same for derivative

e.g. I’m currently working on a MACE interface which might speed up ring polymer dynamics a bit.

Feel free to add more.

A new dynamics method?

DynamicsMethods.Method types

Something completely different?

• Let’s have a chat
• Please document
• Sanity check type hierarchy with other devs/maintainers

NQCD Docs

Lots of stuff could be done.

Add issues on Github with ideas so we can coordinate.

Explain your changes to someone else - helps make code more understandable, avoids hacks like model.model

Questions?

Link to this presentation

Links

GitHub Organisation

Issues and Suggestions

NQCD Documentation