index.qmd

---
title: "Running nonadiabatic MD in Julia"

subtitle: "*A look under the hood of the NQCD packages*"
---

# The maths behind MD: *Solving differential equations*

::: {.notes}
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

::: {.notes}
- 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**](https://docs.sciml.ai/DiffEqDocs/stable/)

- Consistent structure for creating problems and solving them. 

- Identical I/O for different propagation methods. 

- Flexible, extendable. 

[**Unitful.jl**](https://docs.sciml.ai/Unitful/stable/)

- Units
- Unit conversion

::: {.notes}
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`
![](images/ODEProblem.svg)

::: {.notes}
Need to build an ODEProblem, which needs 

starting conditions, 

differential equation to solve and

any other parameters
:::

---

**Equation of motion** (1D Harmonic oscillator)

```julia
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**

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

**Resulting ODE problem:**

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

## Example: Solving the `ODEProblem`
::: {.columns}
::: {.column}
```julia
sol = solve(
    dynamical_problem1, # Problem
    VelocityVerlet(), # Integrator
    reltol=1e-8, # Numerical accuracy
    dt=0.1, # Time step
    # Any other arguments
)
```
:::
::: {.column}
- Yields a solution object with
    - all initial parameters
    - propagated variables
    - SDEs: random noise values
:::
::: 

::: {.notes}
ODEProblem then passed on to `solve` function. 

Takes an integrator (describes how to numerically integrate)

and other parameters. 
:::

---

## The `DifferentialEquations.jl` pipeline
![](images/DiffEQ.svg)

**MD workflows may require:**

- Parallelisation
- Callbacks
- Reductions

::: {.notes}
- 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?

::: {.notes}
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.svg)

::: {.notes}
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
![](images/nqcdynamics.svg)

::: {.notes}
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`{ncols=2}
- `Atoms`
- `PeriodicCell`
- IO functions:^[[Documentation](https://nqcd.github.io/NQCDynamics.jl/stable/api/NQCBase/nqcbase/)]
    - `convert_from_ase_atoms`
    - `convert_to_ase_atoms`

::: {.notes}
- NQCBase contains very basics needed to describe atomic structures
- Atoms, unit cells and structure loading. 
:::

## Potential Energy Surfaces: `NQCModels.jl`

::: {.columns}
::: {.column}
- `Model`
    - Defines a potential energy surface^[[Overview](https://nqcd.github.io/NQCDynamics.jl/stable/NQCModels/overview/)] ^[[Analytical Model Library](https://nqcd.github.io/NQCDynamics.jl/stable/NQCModels/analyticmodels/)]
    - Dimensionality: arbitrary (in principle)
    - Methods: `potential(model, R), derivative(model, R)` (and in-place versions)
- Interface to ASE for ML models / AIMD^[[Documentation](https://nqcd.github.io/NQCDynamics.jl/stable/NQCModels/ase/)]
:::
::: {.column}
```{.julia}
using NQCModels
model=AdiabaticASEModel(ase_structure)
NQCModels.potential(model, R)
NQCModels.derivative!(model, D, R)
```
:::
::: 

::: {.notes}
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
:::

---

![Model types implemented in NQCModels.jl](images/Models.svg)

## Statistical sampling: NQCDistributions.jl
![](images/nqcdistributions.svg)

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

::: {.notes}
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
```julia
atoms = Atoms([:H])
```
- Chemical symbols
- or arbitrary masses (`[:X]` atoms)

### Initial conditions
```julia
initial_conditions = DynamicsVariables(
    v = hcat([1.0]),
    r = hcat([0.0]),
)
```

### Potential energy surface
```julia
model = Harmonic(m = 0.4, ω = 2.0, r₀ = 0.1)
```

### Combining them in a `Simulation`
```julia
sim = Simulation{Classical}(atoms, model)
```

## `Simulation`s^[[NQCD Docs: Getting Started](https://nqcd.github.io/NQCDynamics.jl/stable/getting_started/)]
```julia
sim = Simulation{Method}( # <1>
    atoms::Atoms{T}, # <2>
    model::Model; # <3>
    temperature=0u"K", # <4> 
    cell::AbstractCell=InfiniteCell() # <5>
)
```
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. 

- `Simulation`s contain a `Calculator` to cache energies, forces,...
- Containers for simulation parameters


## Running an MD simulation (2/2)
```julia
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^[[NQCD Docs: Ensemble Simulations](https://nqcd.github.io/NQCDynamics.jl/stable/ensemble_simulations/)]
- Outputs are processed from solution objects

---

### Dynamics Parameters

```{.julia code-line-numbers="3"}
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

```{.julia code-line-numbers="4"}
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

```{.julia code-line-numbers="5"}
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](https://nqcd.github.io/NQCDynamics.jl/stable/api/NQCDynamics/dynamicsoutputs/)

---

### Dynamics Parameters

```{.julia code-line-numbers="6"}
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

```{.julia code-line-numbers="7-9"}
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. 
- `Reduction`s applied to outputs
    - Default: Append
    - `MeanReduction` for ensemble Mean
    - `FileReduction`

---

### Dynamics Parameters

```{.julia code-line-numbers="10"}
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

```{.julia code-line-numbers="11"}
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^[[NQCD Docs: Callbacks](https://nqcd.github.io/NQCDynamics.jl/stable/devdocs/diffeq/#devdocs-callbacks)]:**

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

# 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

```{.julia}
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

```{.julia}
some_vector = Distributed.pmap(a_function, a_vector)
```

## Parallelisation in NQCDynamics.jl through `EnsembleAlgorithm`s

```{.julia code-line-numbers="10"}
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`^[[GitHub](https://github.com/NQCD/ClusterScripts.jl)]
- Designed for simulations across multiple combinations of parameters. 
- Trivial taskfarming
- Temporary file storage, functions to concatenate back into input parameter space. 

```julia
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")
```

# How can I contribute to NQCD?
![NQCDynamics.jl and Submodules](images/NQCD_Modules.svg)

---

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

::: {.notes}
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]{style="color:grey"}

```{.julia code-summary="Simple example for a DynamicsOutput"}
OutputKineticEnergy(sol, i) = DynamicsUtils.classical_kinetic_energy.(sol.prob.p, sol.u)
export OutputKineticEnergy
```

- Functions that take DiffEq solution object
- [DifferentialEquations.jl: Solution Handling](https://docs.sciml.ai/DiffEqDocs/stable/basics/solution/)

```{.julia code-summary="A more complex DynamicsOutput that needs to store its own parameters"}
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
```

::: {.notes}
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]{style="color:grey"}

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

[src/Analysis/Analysis.jl]{style="color:grey"}

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

::: {.notes}
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

::: {.notes}
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](images/DynamicsMethods.svg)

## Something completely different?
- Let's have a chat 
- Please document 
- Sanity check type hierarchy with other devs/maintainers

[NQCD Docs](https://nqcd.github.io/NQCDynamics.jl/stable/){preview-link="true"}

::: {.notes}
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`
:::

::: {.div style="font-weight:400"}
# Questions?
::: {.columns}
::: {.column}
![](images/qr-code.svg){width=66%}

Link to this presentation
:::
::: {.column}
**Links**

[GitHub Organisation](https://github.com/NQCD)

[Issues and Suggestions](https://github.com/NQCD/NQCDynamics.jl/issues)

[NQCD Documentation](https://nqcd.github.io/NQCDynamics.jl/stable/)



:::
::: 
:::