Ensemble simulations

Typically we'll be interested in computing observables based upon the statistics obtained from many trajectories. Technically it is possible to manually run many trajectories using the single trajectory procedure introduced in the Getting started section. However, by using the methods introduced on this page it is possible to run many trajectories at once, using parallelism and computing ensemble observables automatically.

The key function for performing ensemble simulations is run_dynamics.

run_dynamics(sim::AbstractSimulation, tspan, distribution;
    output,
    selection::Union{Nothing,AbstractVector}=nothing,
    reduction=AppendReduction(),
    ensemble_algorithm=SciMLBase.EnsembleSerial(),
    algorithm=DynamicsMethods.select_algorithm(sim),
    trajectories=1,
    kwargs...
    )

This is the same function used to perform single trajectory simulations, but by replacing the single initial condition with a distribution and changing the number of trajectories it is possible to run an ensemble of trajectories. The distributions are defined such that they can be sampled to provide initial conditions for each trajectory. The Storing and sampling distributions page details the format the distributions must take.

Internally, the DifferentialEquations.jl ensemble infrastructure is used to handle per trajectory parallelism. The ensemble_algorithm keyword takes one of the EnsembleAlgorithms. To use these, you must first add using DiffEqBase to your script.

Example

To demonstrate usage of run_dynamics, let's investigate different ways to calculate the time-dependent population with FSSH.

First, we set up our system using one of Tully's simple models ([1]).

using NQCDynamics

atoms = Atoms(2000)
model = TullyModelOne()
sim = Simulation{FSSH}(atoms, model)

As mentioned above, before running the ensemble, we must prepare a distribution to generate initial conditions for each trajectory. This procedure is detailed in the Storing and sampling distributions section.

using Distributions: Normal # Import the Normal distribution

k = 10 # Initial momentum = 10
v = k / atoms.masses[1] # Convert momentum to velocity
r = Normal(-8) # Create Normal distribution centred at -8 for sampling initial position
nuclear_distribution = DynamicalDistribution(v, r, (1,1)) # Combine position and velocity
electronic_distribution = PureState(2) # Create nonequilibrium electronic distribution
product_distribution = nuclear_distribution * electronic_distribution

In this case, we have used a deterministic momentum of $10$ a.u. and a gaussian position distribution with width 1 centered at $-8$ a.u.. The electronic variables will be sampled such that the initial population is confined to the second state by PureState(2).

The final step before running the dynamics is to decide how to output the results. The simplest option is to use the built-in tuple format familiar from run_dynamics.

ensemble = run_dynamics(sim, (0.0, 3000.0), product_distribution;
    trajectories=20, output=OutputDiabaticPopulation)

[ Info: Sampling randomly from provided distribution.
[ Info: Performing 20 trajectories.
[ Info: Finished after 2.810021555 seconds.

This is equivalent to performing single trajectories in a loop and manually re-sampling the initial conditions each time. However, here we have been able to do this more concisely, using internal mechanisms for sampling from the product_distribution.

The output of this function is a vector containing the output from each trajectory. Each entry is equivalent to the output from a call to run_dynamics and can be plotted by iterating through ensemble.

using Plots

p = plot(legend=false)
for traj in ensemble
    plot!(traj[:Time], [population[2] - population[1] for population in traj[:OutputDiabaticPopulation]])
end
p

This plot shows the population difference between the two states for each trajectory. To approximate the exact quantum dynamics for this model, the average over all trajectories should be computed. Instead of manually averaging the result, we can use reduction=MeanReduction() or reduction=SumReduction() which will reduce the data accordingly before outputting the result:

ensemble = run_dynamics(sim, (0.0, 3000.0), product_distribution;
    trajectories=20, output=OutputDiabaticPopulation, reduction=MeanReduction(), saveat=0.0:10.0:3000.0)
plot(ensemble, :OutputDiabaticPopulation)

This workflow can be applied for any of the quantities defined in the DynamicsOutputs submodule. If we want a more complex output, such as a scattering probability or a time-correlation function, we can provide a function to the output argument as described in the DifferentialEquations.jl documentation. The advantage of this approach is that memory can be saved by reducing the data as the trajectories accumulate, it also allows greater flexibility when modifying the output.

Inside the Ensembles submodule we define a few premade functions of this sort, but here we can demonstrate how to reformulate the previous simulation using the alternative format.

function output_function(sol, i)
    output = zeros(2,div(3000, 50) + 1)
    for (i,u) in enumerate(sol.u)
        output[:,i] .= Estimators.diabatic_population(sim, u)
    end
    return output
end

ensemble = run_dynamics(sim, (0.0, 3000.0), product_distribution;
    trajectories=20, output=output_function, reduction=MeanReduction(), saveat=50.0)

2-element Dictionaries.Dictionary{Symbol, Array{Float64}}
            :Time │ [0.0, 50.0, 100.0, 150.0, 200.0, 250.0, 300.0, 350.0, 400.0…
 :output_function │ [0.0 0.0 … 0.65 0.65; 1.0 1.0 … 0.34999999999999654 0.34999…

This function provides us the same output as above, but this gives us the flexibility to calculate any observable we want.

Throughout the documentation, ensemble simulations like this one are used to demonstrate many of the dynamics methods. Now that you have understood the contents of this page, all of the ensemble simulations will appear familiar.