Refactor and update output handling (#286)

CHANGELOG.md

@@ -1,12 +1,19 @@
 # NQCDynamics.jl changelog
 
+## Version `v0.12.0`
+
+- ![BREAKING][badge-breaking] Merged `run_trajectory` and `run_ensemble` into `run_dynamics` [#286][github-286]
+- ![BREAKING][badge-enhancement] Output from `run_dynamics` is a `Dictionary` from `Dictionaries.jl` [#286][github-286]
+- ![enhancement][badge-enhancement] `FileReduction(filename)` allows for output to be stored in HDF5 format [#286][github-286]
+
 ## Version `v0.11.0`
 
 - ![Maintenance][badge-maintenance] Moved `NonadiabaticDistributions` to [NQCDistributions.jl]() [#275][github-275]
 - ![BREAKING][badge-breaking] `BoltzmannVelocityDistribution` renamed to VelocityBoltzmann [#275][github-275]
 - ![BREAKING][badge-breaking] `SingleState`, `ElectronicPopulation` renamed to `PureState`, `MixedState` [#275][github-275]
 - ![enhancement][badge-enhancement] Created changelog!
 
+[github-286]: ]https://github.com/NQCD/NQCDynamics.jl/pull/286
 [github-275]: https://github.com/NQCD/NQCDynamics.jl/pull/275
 
 [badge-breaking]: https://img.shields.io/badge/BREAKING-red.svg

Project.toml

@@ -1,15 +1,15 @@
 name = "NQCDynamics"
 uuid = "36248dfb-79eb-4f4d-ab9c-e29ea5f33e14"
 authors = ["James <james.gardner1421@gmail.com>"]
-version = "0.11.1"
+version = "0.12.0"
 
 [deps]
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 AdvancedMH = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DEDataArrays = "754358af-613d-5f8d-9788-280bf1605d4c"
+Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
-DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
@@ -38,20 +38,18 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
-TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
-UnitfulRecipes = "42071c24-d89e-48dd-8a24-8a12d9b8861f"
 
 [compat]
 AdvancedHMC = "0.3"
 AdvancedMH = "0.6"
 ComponentArrays = "0.11, 0.12, 0.13"
 DEDataArrays = "0.2"
+Dictionaries = "0.3"
 DiffEqBase = "6"
-DiffEqCallbacks = "2.16"
 Distributions = "0.25"
 FastBroadcast = "0.1, 0.2"
 HDF5 = "0.15, 0.16"
@@ -78,12 +76,10 @@ StatsBase = "0.33"
 StochasticDiffEq = "6"
 StructArrays = "0.6"
 TimerOutputs = "0.5"
-TypedTables = "1"
 UnPack = "1"
 UnicodePlots = "2, 3"
 Unitful = "1"
 UnitfulAtomic = "1"
-UnitfulRecipes = "1"
 julia = "1.7"
 
 [extras]

docs/make.jl

@@ -20,7 +20,7 @@ end
     bib,
     sitename = "NQCDynamics.jl",
     modules = [NQCDynamics, NQCDistributions, NQCModels, NQCBase, CubeLDFAModel],
-    strict = false,
+    strict = true,
     format = Documenter.HTML(
         prettyurls = get(ENV, "CI", nothing) == "true",
         canonical = "https://nqcd.github.io/NQCDynamics.jl/stable/",

docs/src/NQCDistributions/overview.md

@@ -23,10 +23,6 @@ and demonstrates how they can be combined into a product distribution.
 
 ### DynamicalDistribution
 
-```@docs
-DynamicalDistribution
-```
-
 When handling distributions for the nuclear degrees of freedom,
 the [`DynamicalDistribution`](@ref) type can be used to store initial velocities and positions:
 ```@setup distribution

docs/src/api/NQCDynamics/dynamicsoutputs.md

@@ -2,7 +2,7 @@
 # DynamicsOutputs
 
 Here are all the functions that you can specify in the output tuple when using
-`run_trajectory`.
+`run_dynamics`.
 To add more, simply add a new function in the `DynamicsOutputs` module. 
 ```@autodocs
 Modules=[NQCDynamics.DynamicsOutputs]

docs/src/devdocs/diffeq.md

@@ -14,7 +14,7 @@ This page details some of the features from DifferentialEquations.jl that we hav
 us to introduce extra code during the dynamics without needing to meddle with the
 integration code directly.
 On the developer side, [Callbacks] is the mechanism used for the saving in the
-[`run_trajectory`](@ref) function and the surface hopping procedure during FSSH.
+[`run_dynamics`](@ref) function and the surface hopping procedure during FSSH.
 The user can also write their own callbacks and give these to any of the dynamics functions
 to manipulate the progress of the dynamics or introduce their own saving mechanism.
 
@@ -36,20 +36,20 @@ sim = Simulation(atoms, model; cell=cell)
 
 z = DynamicsVariables(sim, hcat(1.0), zeros(1,1))
 
-solution = run_trajectory(z, (0.0, 300), sim; dt=0.1, output=:position)
-plot(solution, :position, label="No callbacks", legend=true)
+solution = run_dynamics(sim, (0.0, 300), z; dt=0.1, output=OutputPosition)
+plot(solution, :OutputPosition, label="No callbacks", legend=true)
 ```
 
 Now we can introduce callbacks and observe the difference:
 ```@example callbacks
-solution = run_trajectory(z, (0.0, 300), sim; callback=DynamicsUtils.CellBoundaryCallback(), dt=0.1, output=:position)
-plot!(solution, :position, label="Cell boundary" )
+solution = run_dynamics(sim, (0.0, 300), z; callback=DynamicsUtils.CellBoundaryCallback(), dt=0.1, output=OutputPosition)
+plot!(solution, :OutputPosition, label="Cell boundary" )
 
 using DiffEqBase: CallbackSet
 terminate(u, t, integrator) = t > 100
 callbacks = CallbackSet(DynamicsUtils.CellBoundaryCallback(), DynamicsUtils.TerminatingCallback(terminate))
-solution = run_trajectory(z, (0.0, 300), sim; callback=callbacks, dt=0.1, output=:position)
-plot!(solution, :position, label="Cell + termination")
+solution = run_dynamics(sim, (0.0, 300), z; callback=callbacks, dt=0.1, output=OutputPosition)
+plot!(solution, :OutputPosition, label="Cell + termination")
 ```
 See how the callbacks have altered the dynamics? The atom no longer leaves
 the simulation cell, and the termination caused the simulation to exit early. 

docs/src/devdocs/new_methods.md

@@ -118,7 +118,7 @@ value of `x` has been set to 0.5.
 sim = Simulation(Atoms(1), Free(), MyMethod(2.0))
 u = DynamicsVariables(sim, rand(1,1), rand(1,1), [0.5])
 
-sol = run_trajectory(u, (0.0, 10.0), sim, output=(:position, :velocity, :u))
+sol = run_dynamics(sim, (0.0, 10.0), u, output=(OutputPosition, OutputVelocity, OutputDynamicsVariables))
 ```
 
 !!! note
@@ -133,19 +133,19 @@ To visualise the result we can plot each of the quantities from the output table
 ```@example mymethod
 using Plots
 
-plot(sol, :position, label="Position")
-plot!(sol, :velocity, label="Velocity")
-plot!(sol, :u, label="u", legend=true)
+plot(sol, :OutputPosition, label="Position")
+plot!(sol, :OutputVelocity, label="Velocity")
+plot!(sol, :OutputDynamicsVariables, label="Dynamics variables", legend=true)
 ylabel!("Value(t)")
 ```
 
 !!! note 
 
     The additional `x` parameter that we created cannot be accessed in the output tuple
-    by name as with `position` and `velocity` since it is not a standard quantity.
-    Instead, we request `u` which contains all of the dynamical variables.
-    In the plot, two of the lines labelled `u` overlap the position and velocity result.
-    The unique line labelled `u` is the `x` variable.
+    with a pre-existing function as with `position` and `velocity` since it is not a standard quantity.
+    Instead, we request `OutputDynamicsVariables` which contains all of the dynamical variables.
+    In the plot, two of the lines labelled `DynamicsVariables` overlap the position and velocity result.
+    The unique line labelled `Dynamics variables` is the `x` variable.
     When implementing your method, if you want to add new output quantities you should do
     this inside the [`DynamicsOutputs`](@ref NQCDynamics.DynamicsOutputs) submodule.
 

docs/src/dynamicssimulations/dynamicsmethods/classical.md

@@ -29,7 +29,7 @@ sim = Simulation(Atoms([1, 1]), DiatomicHarmonic())
 v = rand(3, 2)
 u0 = DynamicsVariables(sim, zeros(3, 2), hcat(randn(3), randn(3).+1))
 
-traj = run_trajectory(u0, (0.0, 1e2), sim; dt=0.1, output=(:position))
+traj = run_dynamics(sim, (0.0, 1e2), u0; dt=0.1, output=OutputPosition)
 
-plot(traj, :position)
+plot(traj, :OutputPosition)
 ```

docs/src/dynamicssimulations/dynamicsmethods/ehrenfest.md

@@ -2,7 +2,7 @@
 
 The Ehrenfest method is a mixed quantum-classical dynamics method in which the total wavefunction is factorized into slow (nuclear) variables, which are treated classically, and fast ones (electrons) which remain quantum-mechanical. In the Ehrenfest method, nuclei move according to classical mechanics on a potential energy surface given by the expectation value of the electronic Hamiltonian. 
 
-The time depedence of the electronic wavefunction is expanded into an adiabatic basis and follows the time-dependent Schr\"odinger equation.
+The time dependence of the electronic wavefunction is expanded into an adiabatic basis and follows the time-dependent Schr\"odinger equation.
 ```math
 i\hbar \dot{c}_i(t) = V_i(\mathbf{R}) c_i (t)
 - i\hbar \sum_j \dot{\mathbf{R}} \cdot \mathbf{d}_{ij}(\mathbf{R})c_j(t)
@@ -29,13 +29,14 @@ v = k / atoms.masses[1]
 distribution = DynamicalDistribution(v, r, size(sim))* PureState(1, Adiabatic())
 ```
 
-To run an ensemble simulation we additionally choose number of trajectories `n_traj` and timespan `tspan` and we pass all the established settings to the `run_ensemble` function. In this example we output velocities by specifying `output=:velocity` and store the final values in the `final_velocities` array. Following that, we calculate final momenta.
+To run an ensemble simulation we additionally choose number of trajectories `n_traj` and timespan `tspan` and we pass all the established settings to the `run_dynamics` function.
+In this example we output velocities by specifying `output=OutputVelocity` and store the final values in the `final_velocities` array. Following that, we calculate final momenta.
 ```@example ehrenfest
 n_traj = 5000
 tspan = (0.0, 3000.0)
-solution = run_ensemble(sim, tspan, distribution; 
-    trajectories=n_traj, output=:velocity)
-final_velocities = [r.velocity[end] for r in solution]
+solution = run_dynamics(sim, tspan, distribution; 
+    trajectories=n_traj, output=OutputVelocity)
+final_velocities = [r[:OutputVelocity][end] for r in solution]
 momenta = reduce(vcat, final_velocities*atoms.masses[1])
 ```
 

docs/src/dynamicssimulations/dynamicsmethods/fssh.md

@@ -91,31 +91,31 @@ whether we want this state to be `Adiabatic()` or `Diabatic()`.
 The type of state can be important when considering the ordering of the states.
 The adiabatic states are always arranged from lowest to highest energy, whereas the diabatic
 states will be ordered as defined in the model.
-You can inspect the fields of `u` to ensure the initilisation has proceeded as you intend.
+You can inspect the fields of `u` to ensure the initialisation has proceeded as you intend.
 ```@example fssh
 u = DynamicsVariables(sim, [20/2000;;], [-10.;;], PureState(1, Adiabatic()))
 ```
 
 Finally, the trajectory can be run by passing all the parameters we have set up so far.
-Here, we request both the discrete `:state` output which is equal to ``s(t)`` and 
-`:population`, which gives us the population of each diabatic state along the trajectory.
+Here, we request both the `OutputDiscreteState` output which is equal to ``s(t)`` and 
+`OutputDiabaticPopulation`, which gives us the population of each diabatic state along the trajectory.
 ```@example fssh
-traj = run_trajectory(u, (0.0, 2000.0), sim, output=(:state, :population))
+traj = run_dynamics(sim, (0.0, 2000.0), u, output=(OutputDiscreteState, OutputDiabaticPopulation))
 ```
 
 Now we can plot ``s(t)`` throughout the trajectory. The FSSH algorithm attempts to minimise
 the total number of hops; in the limit of infinite hops the result would tend to the
 mean-field (Ehrenfest) result, which is what FSSH attempts to avoid.
 ```@example fssh
 using Plots
-plot(traj, :state)
+plot(traj, :OutputDiscreteState)
 ```
 
 Similarly, we can plot the diabatic populations. Since FSSH is performed in the adiabatic
 representation, even in the case of few hops, the diabatic populations can look dramatically
 different depending on the complexity of the model Hamiltonian. 
 ```@example fssh
-plot(traj, :population)
+plot(traj, :OutputDiabaticPopulation)
 ```
 
 [Another example is available](@ref examples-tully-model-two) where we use FSSH and other

docs/src/dynamicssimulations/dynamicsmethods/langevin.md

@@ -70,19 +70,19 @@ outputted and have selected a timestep `dt`.
 Since the default algorithm is a fixed timestep algorithm an error will be thrown if a
 timestep is not provided.
 ```@example langevin
-traj = run_trajectory(u, (0.0, 2000.0), sim; output=(:position, :velocity), dt=0.1)
+traj = run_dynamics(sim, (0.0, 2000.0), u; output=(OutputPosition, OutputVelocity), dt=0.1)
 ```
 
 Here, we plot the positions of our two atoms throughout the simulation.
 ```@example langevin
 using Plots
-plot(traj, :position, label=["Hydrogen" "Carbon"], legend=true)
+plot(traj, :OutputPosition, label=["Hydrogen" "Carbon"], legend=true)
 ```
 
 We next plot the velocities. Notice how the carbon atom with its heavier mass has a smaller
 magnitude throughout.
 ```@example langevin
-plot(traj, :velocity, label=["Hydrogen" "Carbon"], legend=true)
+plot(traj, :OutputVelocity, label=["Hydrogen" "Carbon"], legend=true)
 ```
 
 Using the configurations from the Langevin simulation we can obtain expectation values along
@@ -98,13 +98,13 @@ as simple as possible.
 Let's find the expectation for the potential energy during our simulation.
 This is the potential energy of the final configuration in the simulation:
 ```@repl langevin
-Estimators.potential_energy(sim, traj.position[end])
+Estimators.potential_energy(sim, traj[:OutputPosition][end])
 ```
 We could evaluate this for every configuration and average it manually.
 Fortunately however, we have the [`@estimate`](@ref Estimators.@estimate) macro that
 will do this for us:
 ```@repl langevin
-Estimators.@estimate potential_energy(sim, traj.position)
+Estimators.@estimate potential_energy(sim, traj[:OutputPosition])
 ```
 
 !!! tip
@@ -118,7 +118,7 @@ Estimators.@estimate potential_energy(sim, traj.position)
 
 Similarly, we can evaluate the kinetic energy expectation with:
 ```@repl langevin
-Estimators.@estimate kinetic_energy(sim, traj.velocity)
+Estimators.@estimate kinetic_energy(sim, traj[:OutputVelocity])
 ```
 Again, this takes a similar value since the total energy is evenly split between the kinetic
 and potential for a classical harmonic system.

docs/src/dynamicssimulations/dynamicsmethods/mdef.md

@@ -51,8 +51,8 @@ a function of time.
 ```@example mdef
 using Plots
 
-solution = run_trajectory(z, (0.0, 100u"fs"), sim, dt=0.1u"fs", output=(:hamiltonian))
-plot(solution, :hamiltonian)
+solution = run_dynamics(sim, (0.0, 100u"fs"), z, dt=0.1u"fs", output=OutputTotalEnergy)
+plot(solution, :OutputTotalEnergy)
 ```
 
 !!! note
@@ -81,9 +81,9 @@ Now we can re-simulate, replacing the fixed temperature with the function we hav
 
 ```@example mdef
 sim = Simulation{MDEF}(atoms, model; temperature=temperature_function)
-solution = run_trajectory(z, (0.0, 100u"fs"), sim, dt=0.1u"fs",
-    output=(:hamiltonian, :position, :velocity))
-plot(solution, :hamiltonian)
+solution = run_dynamics(sim, (0.0, 100u"fs"), z, dt=0.1u"fs",
+    output=OutputTotalEnergy)
+plot(solution, :OutputTotalEnergy)
 ```
 
 This time we see a peak in the energy in the middle of the simulation which coincides

docs/src/dynamicssimulations/dynamicsmethods/nrpmd.md

@@ -172,11 +172,10 @@ The resulting plot shows the time dependent population difference and closely ma
 the figure from the paper we were attempting to reproduce. Nice!
 
 ```@example nrpmd
-ensemble = run_ensemble(sim, (0.0, 30.0), distribution;
-    trajectories=1000, output, reduction=:mean, dt=0.1,
-    u_init=[zeros(2,2) for i=1:length(0:0.1:30.0)])
+ensemble = run_dynamics(sim, (0.0, 30.0), distribution;
+    trajectories=1000, output, reduction=MeanReduction(), dt=0.1)
 
-plt = lines(0:0.1:30, [p[1,1]-p[2,1] for p in ensemble])
+plt = lines(0:0.1:30, [p[1,1]-p[2,1] for p in ensemble[:PopulationCorrelationFunction]])
 plt.axis.xlabel = "Time"
 plt.axis.ylabel = "Population difference"
 plt

docs/src/dynamicssimulations/dynamicsmethods/rpmd.md

@@ -74,7 +74,7 @@ Now we can run the simulation, for which we use the time interval 0.0 to 500.0 a
 step of `dt = 2.5`:
 ```@example rpmd
 dt = 2.5
-traj = run_trajectory(u, (0.0, 500.0), sim; output=(:position), dt=dt)
+traj = run_dynamics(sim, (0.0, 500.0), u; output=OutputPosition, dt=dt)
 nothing # hide
 ```
 
@@ -87,7 +87,7 @@ to its two neighbours.
 ```@example rpmd
 using CairoMakie
 
-rs = traj.position
+rs = traj[:OutputPosition]
 
 index = Observable(1)
 xs = @lift(rs[$index][1,1,:])
@@ -100,7 +100,7 @@ lines!(close_loop_x, close_loop_y)
 xlims!(-3, 3)
 ylims!(-3, 3)
 
-timestamps = 1:length(traj.position)
+timestamps = 1:length(traj[:OutputPosition])
 filepath = "../../assets/figures/rpmd.mp4" # hide
 record(fig, filepath, timestamps;
         framerate = 30) do i

docs/src/dynamicssimulations/dynamicsmethods/rpsh.md

@@ -58,10 +58,10 @@ Now let's run an ensemble of trajectories that sample from this distribution.
 For the output we will receive the diabatic population at intervals of `t=50`
 and it will be averaged over all trajectories by the `:mean` keyword.
 ```@example rpsh
-solution = run_ensemble(sim, (0.0, 3000.0), distribution;
+solution = run_dynamics(sim, (0.0, 3000.0), distribution;
     saveat=50, trajectories=5e2, dt=1,
     output=TimeCorrelationFunctions.PopulationCorrelationFunction(sim, Diabatic()),
-    reduction=:mean, u_init=[zeros(3,3) for i=1:length(0:50:3000)])
+    reduction=MeanReduction())
 ```
 
 !!! note
@@ -77,9 +77,9 @@ A discussion on this topic is available from [landry2013](@cite).
 ```@example rpsh
 using Plots
 
-plot(0:50:3000, [p[1,1] for p in solution], label="State 1")
-plot!(0:50:3000, [p[1,2] for p in solution], label="State 2")
-plot!(0:50:3000, [p[1,3] for p in solution], label="State 3")
+plot(0:50:3000, [p[1,1] for p in solution[:PopulationCorrelationFunction]], label="State 1")
+plot!(0:50:3000, [p[1,2] for p in solution[:PopulationCorrelationFunction]], label="State 2")
+plot!(0:50:3000, [p[1,3] for p in solution[:PopulationCorrelationFunction]], label="State 3")
 xlabel!("Time /a.u.")
 ylabel!("Population")
 ```