diff --git a/CHANGELOG.md b/CHANGELOG.md index a3e9469b5..1aa419adc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # 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] @@ -7,6 +13,7 @@ - ![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 diff --git a/Project.toml b/Project.toml index 76a4c8840..9eb1dd3bc 100644 --- a/Project.toml +++ b/Project.toml @@ -1,15 +1,15 @@ name = "NQCDynamics" uuid = "36248dfb-79eb-4f4d-ab9c-e29ea5f33e14" authors = ["James "] -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] diff --git a/docs/make.jl b/docs/make.jl index 94e3a85d6..b97888de7 100644 --- a/docs/make.jl +++ b/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/", diff --git a/docs/src/NQCDistributions/overview.md b/docs/src/NQCDistributions/overview.md index 3a8e2989c..141dcea4a 100644 --- a/docs/src/NQCDistributions/overview.md +++ b/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 diff --git a/docs/src/api/NQCDynamics/dynamicsoutputs.md b/docs/src/api/NQCDynamics/dynamicsoutputs.md index 0384f3a3f..d998122f9 100644 --- a/docs/src/api/NQCDynamics/dynamicsoutputs.md +++ b/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] diff --git a/docs/src/devdocs/diffeq.md b/docs/src/devdocs/diffeq.md index cece78d6d..cd50d956d 100644 --- a/docs/src/devdocs/diffeq.md +++ b/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. diff --git a/docs/src/devdocs/new_methods.md b/docs/src/devdocs/new_methods.md index bbb7d6ee1..f80338246 100644 --- a/docs/src/devdocs/new_methods.md +++ b/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. diff --git a/docs/src/dynamicssimulations/dynamicsmethods/classical.md b/docs/src/dynamicssimulations/dynamicsmethods/classical.md index 34585dca0..62b2b941b 100644 --- a/docs/src/dynamicssimulations/dynamicsmethods/classical.md +++ b/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) ``` diff --git a/docs/src/dynamicssimulations/dynamicsmethods/ehrenfest.md b/docs/src/dynamicssimulations/dynamicsmethods/ehrenfest.md index 223875dab..ba9820c8c 100644 --- a/docs/src/dynamicssimulations/dynamicsmethods/ehrenfest.md +++ b/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]) ``` diff --git a/docs/src/dynamicssimulations/dynamicsmethods/fssh.md b/docs/src/dynamicssimulations/dynamicsmethods/fssh.md index 2426df613..61a37259e 100644 --- a/docs/src/dynamicssimulations/dynamicsmethods/fssh.md +++ b/docs/src/dynamicssimulations/dynamicsmethods/fssh.md @@ -91,16 +91,16 @@ 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 @@ -108,14 +108,14 @@ the total number of hops; in the limit of infinite hops the result would tend to 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 diff --git a/docs/src/dynamicssimulations/dynamicsmethods/langevin.md b/docs/src/dynamicssimulations/dynamicsmethods/langevin.md index 734e45bd2..47b36eaa4 100644 --- a/docs/src/dynamicssimulations/dynamicsmethods/langevin.md +++ b/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. diff --git a/docs/src/dynamicssimulations/dynamicsmethods/mdef.md b/docs/src/dynamicssimulations/dynamicsmethods/mdef.md index edc3edf80..76b157d82 100644 --- a/docs/src/dynamicssimulations/dynamicsmethods/mdef.md +++ b/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 diff --git a/docs/src/dynamicssimulations/dynamicsmethods/nrpmd.md b/docs/src/dynamicssimulations/dynamicsmethods/nrpmd.md index b5b643434..209b6331a 100644 --- a/docs/src/dynamicssimulations/dynamicsmethods/nrpmd.md +++ b/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 diff --git a/docs/src/dynamicssimulations/dynamicsmethods/rpmd.md b/docs/src/dynamicssimulations/dynamicsmethods/rpmd.md index 35f6e246c..2de98f9f1 100644 --- a/docs/src/dynamicssimulations/dynamicsmethods/rpmd.md +++ b/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 diff --git a/docs/src/dynamicssimulations/dynamicsmethods/rpsh.md b/docs/src/dynamicssimulations/dynamicsmethods/rpsh.md index d17c24ad2..aef6ab073 100644 --- a/docs/src/dynamicssimulations/dynamicsmethods/rpsh.md +++ b/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") ``` diff --git a/docs/src/dynamicssimulations/dynamicssimulations.md b/docs/src/dynamicssimulations/dynamicssimulations.md index 2880074f7..f596f868a 100644 --- a/docs/src/dynamicssimulations/dynamicssimulations.md +++ b/docs/src/dynamicssimulations/dynamicssimulations.md @@ -54,35 +54,35 @@ The output of this function takes the place of the `u` argument seen throughout [DifferentialEquations](https://diffeq.sciml.ai/stable/). With both the [`Simulation`](@ref) and [`DynamicsVariables`](@ref) in hand, -the central function is [`run_trajectory`](@ref) which allows us to perform a single dynamics trajectory. -[`run_trajectory`](@ref) takes the simulation parameters `sim` and the initial conditions `u0`, along with a time span `tspan` +the central function is [`run_dynamics`](@ref) which allows us to perform a single dynamics trajectory. +[`run_dynamics`](@ref) takes the simulation parameters `sim` and the initial conditions `u0`, along with a time span `tspan` that the trajectory will cover. ```@example dynamics tspan = (0.0, 2000.0) -run_trajectory(u0, tspan, sim) +run_dynamics(sim, tspan, u0; output=OutputDynamicsVariables) ``` -By default, we can see that the output contains both `t` and `u`. These are the time and dynamics variables -respectively. -By passing a `Tuple` to the `output` keyword argument we can ask for specific quantities. +The output is a dictionary containing entries for `:Time` and our requested output quantity. +Output is a required keyword and the code will error unless at least one quantity is specified. +By passing a `Tuple` to the `output` keyword argument we can ask for multiple quantities. ```@example dynamics -out = run_trajectory(u0, tspan, sim; output=(:position, :adiabatic_population)) +out = run_dynamics(sim, tspan, u0; output=(OutputPosition, OutputAdiabaticPopulation)) ``` The quantities that are available are listed [here](@ref NQCDynamics). -Adding more quantities requires that a new function is defined inside the [`DynamicsOutputs`](@ref NQCDynamics.DynamicsOutputs) module -that has the same set of parameters as the existing quantities. +More quantities can be added by defining new functions with the signature `f(sol, i)`. +The first argument is the `DifferentialEquations.jl` solution object and the second is the trajectory index. -This time we can see that the output contains only the quantities that we asked for. +This time we can see that the output contains the two quantities that we asked for. ```@example dynamics using Plots -plot(out, :position) +plot(out, :OutputPosition) ``` ```@example dynamics -plot(out, :adiabatic_population) +plot(out, :OutputAdiabaticPopulation) ``` !!! note diff --git a/docs/src/ensemble_simulations.md b/docs/src/ensemble_simulations.md index 425eb32bc..91223db88 100644 --- a/docs/src/ensemble_simulations.md +++ b/docs/src/ensemble_simulations.md @@ -7,23 +7,16 @@ procedure introduced in the [Getting started](@ref) 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_ensemble`](@ref). +The key function for performing ensemble simulations is [`run_dynamics`](@ref). ```@docs -run_ensemble +run_dynamics ``` -From the function signature displayed above it should be possible to identify the similarities to the [`run_trajectory`](@ref) function. -The `sim` and `tspan` positional arguments are the same, but the initial [`DynamicsVariables`](@ref) have been replaced by a distribution. -These distributions are defined such that they can be sampled to provide initial conditions for each trajectory. -The [Storing and sampling distributions](@ref) page details the format this distribution must take. - -The output may take two distinct forms: the tuple structure familiar from [`run_trajectory`](@ref) `output=(:position, :velocity, :hamiltonian)` -or a function as described in the [DifferentialEquations.jl](https://diffeq.sciml.ai/stable/features/ensemble/#Performing-an-Ensemble-Simulation) documentation. -These options allow access to the common output quantities defined in the [`DynamicsOutputs`](@ref NQCDynamics.DynamicsOutputs) module -along with specific customised output. -Already implemented in the code are a small library of existing functions of this type, but it is possible to use any Julia function in its place. -The existing ensemble outputs can be found [here](@ref Ensembles). +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](@ref) page details the format the distributions must take. Internally, the [DifferentialEquations.jl](https://diffeq.sciml.ai/stable/features/ensemble/#Performing-an-Ensemble-Simulation) ensemble infrastructure is used to handle per trajectory parallelism. @@ -32,7 +25,7 @@ To use these, you must first add `using DiffEqBase` to your script. ## Example -To demonstrate usage of [`run_ensemble`](@ref), let's investigate different ways to calculate the time-dependent population +To demonstrate usage of [`run_dynamics`](@ref), let's investigate different ways to calculate the time-dependent population with [FSSH](@ref fssh-dynamics). First, we set up our system using one of Tully's simple models ([Tully1990](@cite)). @@ -65,35 +58,35 @@ The electronic variables will be sampled such that the initial population is con 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_trajectory`](@ref). +The simplest option is to use the built-in tuple format familiar from [`run_dynamics`](@ref). ```@example ensemble -ensemble = run_ensemble(sim, (0.0, 3000.0), product_distribution; - trajectories=20, output=:population) +ensemble = run_dynamics(sim, (0.0, 3000.0), product_distribution; + trajectories=20, output=OutputDiabaticPopulation) nothing # hide ``` 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_trajectory`](@ref) and +Each entry is equivalent to the output from a call to [`run_dynamics`](@ref) and can be plotted by iterating through `ensemble`. ```@example ensemble using Plots p = plot(legend=false) for traj in ensemble - plot!(traj.t, [population[2] - population[1] for population in traj.population]) + 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=:mean` or `reduction=:sum` +Instead of manually averaging the result, we can use `reduction=MeanReduction()` or `reduction=SumReduction()` which will reduce the data accordingly before outputting the result: ```@example ensemble -ensemble = run_ensemble(sim, (0.0, 3000.0), product_distribution; - trajectories=20, output=:population, reduction=:mean, saveat=0.0:10.0:3000.0) -plot(ensemble, :population) +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) ``` !!! note @@ -116,22 +109,15 @@ we can demonstrate how to reformulate the previous simulation using the alternat 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) + output[:,i] .= Estimators.diabatic_population(sim, u) end - return (output, false) + return output end -ensemble = run_ensemble(sim, (0.0, 3000.0), product_distribution; - trajectories=20, output=output_function, reduction=:mean, u_init=zeros(2,div(3000, 50)+1), saveat=50.0) +ensemble = run_dynamics(sim, (0.0, 3000.0), product_distribution; + trajectories=20, output=output_function, reduction=MeanReduction(), saveat=50.0) ``` -This function provides us the same output as above, but here we have defined it -in a way compatible with the [DifferentialEquations.jl format](https://diffeq.sciml.ai/stable/features/ensemble/#Building-a-Problem). -To use this format, we have additionally specified `u_init` and `saveat`. -These keywords come from [DifferentialEquations.jl](https://diffeq.sciml.ai/stable/features/ensemble/#Building-a-Problem), -and specify the shape of the output and at which timesteps we want to evaluate `output_function` and save the values. -In this example we have saved the output at intervals of 50 from 0 to 3000, our `tspan`. -The advantage of this format is that greater flexibility is available -since `output_function` can be modified to output any quantity. +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. diff --git a/docs/src/examples/reactive_scattering.md b/docs/src/examples/reactive_scattering.md index ad6f8180b..92871f7aa 100644 --- a/docs/src/examples/reactive_scattering.md +++ b/docs/src/examples/reactive_scattering.md @@ -119,12 +119,12 @@ model = LDFAModel(model, "../assets/friction/test.cube", atoms, friction_atoms=[ ``` Now we can pass all the variables defined so far to the `Simulation` and run multiple -trajectories using [`run_ensemble`](@ref). +trajectories using [`run_dynamics`](@ref). ```@example h2scatter sim = Simulation{MDEF}(atoms, model, cell=cell, temperature=300u"K") -ensemble = run_ensemble(sim, tspan, distribution;selection=1:20, - dt=0.1u"fs", output=:position, trajectories=20, callback=terminate) +ensemble = run_dynamics(sim, tspan, distribution;selection=1:20, + dt=0.1u"fs", output=OutputPosition, trajectories=20, callback=terminate) ``` ## MDEF with neural network friction @@ -138,8 +138,8 @@ This can be used by simply using the model directly, without wrapping it with th ```@example h2scatter model = H2AgModel() sim = Simulation{MDEF}(atoms, model, cell=cell, temperature=300u"K") -ensemble = run_ensemble(sim, tspan, distribution;selection=1:20, - dt=0.1u"fs", output=:position, trajectories=20, callback=terminate) +ensemble = run_dynamics(sim, tspan, distribution;selection=1:20, + dt=0.1u"fs", output=OutputPosition, trajectories=20, callback=terminate) ``` ## Visualisation diff --git a/docs/src/examples/spinboson.md b/docs/src/examples/spinboson.md index 51ad073c8..28fee937b 100644 --- a/docs/src/examples/spinboson.md +++ b/docs/src/examples/spinboson.md @@ -54,20 +54,19 @@ ehrenfest = Simulation{Ehrenfest}(atoms, model) saveat = 0:0.1:20 output = TimeCorrelationFunctions.PopulationCorrelationFunction(fssh, Diabatic()) -u_init = [zeros(2,2) for i=1:length(saveat)] -ensemble_fssh = run_ensemble(fssh, (0.0, 20.0), distribution; - saveat=saveat, trajectories=100, output, reduction=:mean, u_init=copy(u_init)) +ensemble_fssh = run_dynamics(fssh, (0.0, 20.0), distribution; + saveat=saveat, trajectories=100, output, reduction=MeanReduction()) output = TimeCorrelationFunctions.PopulationCorrelationFunction(ehrenfest, Diabatic()) -ensemble_ehrenfest = run_ensemble(ehrenfest, (0.0, 20.0), distribution; - saveat=saveat, trajectories=100, output, reduction=:mean, u_init=copy(u_init)) +ensemble_ehrenfest = run_dynamics(ehrenfest, (0.0, 20.0), distribution; + saveat=saveat, trajectories=100, output, reduction=MeanReduction()) nothing # hide ``` Here, we can see the population difference between the two states. ```@example spinboson using Plots -plot(saveat, [p[1,1] - p[1,2] for p in ensemble_fssh], label="FSSH") -plot!(saveat, [p[1,1] - p[1,2] for p in ensemble_ehrenfest], label="Ehrenfest") +plot(saveat, [p[1,1] - p[1,2] for p in ensemble_fssh[:PopulationCorrelationFunction]], label="FSSH") +plot!(saveat, [p[1,1] - p[1,2] for p in ensemble_ehrenfest[:PopulationCorrelationFunction]], label="Ehrenfest") xlabel!("Time /a.u.") ylabel!("Population difference") ``` diff --git a/docs/src/examples/threestatemorse.md b/docs/src/examples/threestatemorse.md index 682aa52f8..f85365f55 100644 --- a/docs/src/examples/threestatemorse.md +++ b/docs/src/examples/threestatemorse.md @@ -69,27 +69,27 @@ the intial population with the final population at each timestep. ```@example threestatemorse sim = Simulation{FSSH}(atoms, model) -fssh_result = run_ensemble(sim, (0.0, 3000.0), distribution; +fssh_result = run_dynamics(sim, (0.0, 3000.0), distribution; saveat=10, trajectories=1e3, output=TimeCorrelationFunctions.PopulationCorrelationFunction(sim, Diabatic()), - reduction=:mean, dt=1.0, u_init=[zeros(3,3) for i=1:length(0:10:3000)]) + reduction=MeanReduction(), dt=1.0) sim = Simulation{Ehrenfest}(atoms, model) -ehrenfest_result = run_ensemble(sim, (0.0, 3000.0), distribution; +ehrenfest_result = run_dynamics(sim, (0.0, 3000.0), distribution; saveat=10, trajectories=1e3, output=TimeCorrelationFunctions.PopulationCorrelationFunction(sim, Diabatic()), - reduction=:mean, dt=1.0, u_init=[zeros(3,3) for i=1:length(0:10:3000)]) + reduction=MeanReduction(), dt=1.0) fig = Figure() ax = Axis(fig[1,1], xlabel="Time /a.u.", ylabel="Population") x = 0:10:3000 -lines!(ax, x, [p[1,1] for p in fssh_result], label="FSSH State 1", color=:red) -lines!(ax, x, [p[1,2] for p in fssh_result], label="FSSH State 2", color=:green) -lines!(ax, x, [p[1,3] for p in fssh_result], label="FSSH State 3", color=:blue) +lines!(ax, x, [p[1,1] for p in fssh_result[:PopulationCorrelationFunction]], label="FSSH State 1", color=:red) +lines!(ax, x, [p[1,2] for p in fssh_result[:PopulationCorrelationFunction]], label="FSSH State 2", color=:green) +lines!(ax, x, [p[1,3] for p in fssh_result[:PopulationCorrelationFunction]], label="FSSH State 3", color=:blue) -lines!(ax, x, [p[1,1] for p in ehrenfest_result], label="Ehrenfest State 1", color=:red, linestyle=:dash) -lines!(ax, x, [p[1,2] for p in ehrenfest_result], label="Ehrenfest State 2", color=:green, linestyle=:dash) -lines!(ax, x, [p[1,3] for p in ehrenfest_result], label="Ehrenfest State 3", color=:blue, linestyle=:dash) +lines!(ax, x, [p[1,1] for p in ehrenfest_result[:PopulationCorrelationFunction]], label="Ehrenfest State 1", color=:red, linestyle=:dash) +lines!(ax, x, [p[1,2] for p in ehrenfest_result[:PopulationCorrelationFunction]], label="Ehrenfest State 2", color=:green, linestyle=:dash) +lines!(ax, x, [p[1,3] for p in ehrenfest_result[:PopulationCorrelationFunction]], label="Ehrenfest State 3", color=:blue, linestyle=:dash) axislegend(ax) fig diff --git a/docs/src/examples/tully_scattering.md b/docs/src/examples/tully_scattering.md index ab691d7ef..2cb6d4f4a 100644 --- a/docs/src/examples/tully_scattering.md +++ b/docs/src/examples/tully_scattering.md @@ -23,10 +23,10 @@ Firstly, we can prepare the parts that will be the same for every ensemble: ```@example tullymodeltwo using ComponentArrays: ComponentVector -output = Ensembles.OutputStateResolvedScattering1D(sim, :adiabatic) +output = OutputStateResolvedScattering1D(sim, :adiabatic) ``` Here, we are using the -[`OutputStateResolvedScattering1D`](@ref Ensembles.OutputStateResolvedScattering1D) +[`OutputStateResolvedScattering1D`](@ref OutputStateResolvedScattering1D) along with the [`MeanReduction`](@ref Ensembles.MeanReduction) which will give us the average scattering outcome from the entire ensemble. Each trajectory outputs the scattering outcome along with its final adiabatic state, and the reduction @@ -35,7 +35,7 @@ computes the average over all trajectories. Next, we can choose how many trajectories we want to perform for each ensemble, and choose the range of momentum values: ```@example tullymodeltwo -ntraj = 500 +trajectories = 500 momenta = 9:2:50 ``` @@ -60,12 +60,11 @@ for k in momenta # Iterate through each momentum value tspan = (0, 2abs(r)/v) distribution = DynamicalDistribution(v, -5, size(sim)) * PureState(1, Adiabatic()) - out = run_ensemble(sim, tspan, distribution; - trajectories=ntraj, output=output, reduction=:mean, - u_init=ComponentVector(reflection=zeros(2), transmission=zeros(2)) + out = run_dynamics(sim, tspan, distribution; + saveat=tspan[end], trajectories, output, reduction=MeanReduction() ) - push!(result, out) + push!(result, out[:OutputStateResolvedScattering1D]) end result diff --git a/docs/src/getting_started.md b/docs/src/getting_started.md index ac9811009..124ad02a4 100644 --- a/docs/src/getting_started.md +++ b/docs/src/getting_started.md @@ -219,16 +219,16 @@ z = DynamicsVariables(sim, randn(size(sim)), randn(size(sim))) nothing # hide ``` -Now, we can finally run the trajectory using the `run_trajectory` function. -This takes three positional arguments: the dynamics variables `z`, the time span -we want to solve for `tspan`, and the simulation parameters `sim`. +Now, we can finally run the trajectory using the `run_dynamics` function. +This takes three positional arguments: the simulation parameters `sim, the time span +we want to solve for `tspan`, and the dynamics variables `z`. For classical dynamics we also provide a timestep `dt` since we're using the `VelocityVerlet` algorithm by default. !!! note "Integration algorithms" Each method will default to an appropriate integration algorithm though it is possible - to specify via a keyword argument to [`run_trajectory`](@ref) if an alternative + to specify via a keyword argument to [`run_dynamics`](@ref) if an alternative algorithm is preferred. Refer to the [dynamics documentation](@ref dynamicssimulations) for more information. @@ -238,43 +238,39 @@ A list of the available quantities can be found [here](@ref DynamicsOutputs). !!! tip "Output format" - `run_trajectory` returns a `Table` from `TypedTables.jl` that has columns containing - the time and the output quantities saved at each time. - By default, it outputs the value of the dynamics variables into the field `u`. + `run_dynamics` returns a `Dictionary` from `Dictionaries.jl` that has entries containing + the time and the output quantities saved at each time step. ```@example classical tspan = (0.0, 50.0) -solution = run_trajectory(z, (0.0, 50.0), sim; - dt=0.1, output=(:position, :velocity)) +solution = run_dynamics(sim, (0.0, 50.0), z; + dt=0.1, output=(OutputPosition, OutputVelocity)) ``` -Here you can see the output table with columns for the time and the output quantities -we specified. +Here you can see the output containing the time steps and the output quantities we specified. These can be accessed directly as shown here: ```@repl classical -solution.t -solution.position +solution[:Time] +solution[:OutputPosition] ``` As with the models, we provide custom plotting recipes to quickly visualise the results before performing further analysis by manually accessing the fields of the solution table. To use these recipes, simply provide the solution to the `plot` function from `Plots.jl` and give the name of the output quantity as the second argument. -This will only work if this quantity was specified in `run_trajectory`. +This will only work if this quantity was specified in `run_dynamics`. ```@example classical using Plots # hide -plot(solution, :position) -plot!(solution, :velocity) +plot(solution, :OutputPosition) +plot!(solution, :OutputVelocity) ``` ### Ensemble simulations -So we have solved a single trajectory? That's pretty cool but wouldn't it be great -if we could do a whole bunch at once? Well, fortunately we can thanks to the -[`Ensembles`](@ref) module. -This provides [`run_ensemble`](@ref) which can be used to perform many trajectories in parallel. -To learn more and see some examples, refer to the -[Ensemble simulations](@ref ensembles) section. +We have shown how to perform a single trajectory, but usually we are interested in performing many +and calculating observables using statistical methods. +Running more trajectories is as simple as providing the `trajectories` keyword to `run_dynamics`, +but we'll go through this in more detail in the [Ensemble simulations section.](@ref ensembles) ### What's next? diff --git a/src/DynamicsMethods/DynamicsMethods.jl b/src/DynamicsMethods/DynamicsMethods.jl index 665a74fa2..ff4d158e0 100644 --- a/src/DynamicsMethods/DynamicsMethods.jl +++ b/src/DynamicsMethods/DynamicsMethods.jl @@ -17,14 +17,12 @@ provided by `DifferentialEquations.jl`. module DynamicsMethods using DiffEqBase: DiffEqBase -using TypedTables: TypedTables using UnitfulAtomic: austrip using Reexport: @reexport using OrdinaryDiffEq: OrdinaryDiffEq using ComponentArrays: ComponentVector -using NQCDynamics: AbstractSimulation, Simulation, RingPolymerSimulation, - DynamicsOutputs +using NQCDynamics: AbstractSimulation, Simulation, RingPolymerSimulation using NQCBase: austrip_kwargs export DynamicsVariables @@ -71,51 +69,12 @@ classical dynamics. """ DynamicsVariables(::AbstractSimulation, v, r) = ComponentVector(v=v, r=r) +run_trajectory(args...; kwargs...) = throw(error(""" +`run_trajectory` has been replaced, use `run_dynamics` instead. +`run_dynamics` unifies single trajectory and ensemble simulations into one function. +Refer to the `run_dynamics` docstring for more information. """ - run_trajectory(u0, tspan::Tuple, sim::AbstractSimulation; - output=(:u,), saveat=[], algorithm=select_algorithm(sim), kwargs...) - -Solve a single trajectory starting from `u0` with a timespan `tspan` -for the simulation `sim`. - -# Keyword arguments - -`output` specifies the quantities that should be saved during the dynamics simulation. -The options for this keyword are any of the functions found in -`src/DynamicsMethods/output.jl`. - -The rest of the keywords are the usual arguments for the `solve` function from -`DifferentialEquations.jl`. -It is possible to use `Unitful` quantities for any of the arguments since these are -automatically converted to atomic units internally. - -# Output - -The function returns a `Table` from `TypedTables.jl` with columns for time `t` and -every quantity specified in the `output` tuple. -""" -function run_trajectory(u0, tspan::Tuple, sim::AbstractSimulation; - output=(:u,), saveat=[], callback=nothing, algorithm=select_algorithm(sim), - kwargs...) - - if !(output isa Tuple) - output = (output,) - end - - problem = create_problem(u0, austrip.(tspan), sim) - - saving_callback, vals = DynamicsOutputs.create_saving_callback(output; saveat=austrip.(saveat)) - callback_set = DiffEqBase.CallbackSet(callback, saving_callback, get_callbacks(sim)) - problem = DiffEqBase.remake(problem, callback=callback_set) - - stripped_kwargs = austrip_kwargs(;kwargs...) - @info "Performing a single trajectory." - stats = @timed DiffEqBase.solve(problem, algorithm; stripped_kwargs...) - @info "Finished after $(stats.time) seconds." - @debug "Simulation details" algorithm=stats.value.alg defunction=problem.f - out = [(;zip(output, val)...) for val in vals.saveval] - TypedTables.Table(t=vals.t, out) -end +)) include("ClassicalMethods/ClassicalMethods.jl") @reexport using .ClassicalMethods diff --git a/src/DynamicsMethods/SurfaceHoppingMethods/fssh.jl b/src/DynamicsMethods/SurfaceHoppingMethods/fssh.jl index db636ab60..69e6d3d73 100644 --- a/src/DynamicsMethods/SurfaceHoppingMethods/fssh.jl +++ b/src/DynamicsMethods/SurfaceHoppingMethods/fssh.jl @@ -227,6 +227,6 @@ end function DynamicsUtils.classical_hamiltonian(sim::Simulation{<:FSSH}, u) kinetic = DynamicsUtils.classical_kinetic_energy(sim, DynamicsUtils.get_velocities(u)) eigs = Calculators.get_eigen(sim.calculator, DynamicsUtils.get_positions(u)) - potential = eigs.values[sim.method.state] + potential = eigs.values[u.state] return kinetic + potential end diff --git a/src/DynamicsOutputs.jl b/src/DynamicsOutputs.jl index 1c20ab3ca..c4ef28e1f 100644 --- a/src/DynamicsOutputs.jl +++ b/src/DynamicsOutputs.jl @@ -2,177 +2,236 @@ """ DynamicsOutputs -Infrastructure for saving quantities during trajectories. - -Defines all available options that can be used within the `output` tuple along with -the functions that perform the saving operation. +Defines a set of functions that can be used to calculate outputs for dynamics simulations. """ module DynamicsOutputs -using RecursiveArrayTools: ArrayPartition -using DiffEqCallbacks: SavedValues, SavingCallback -using ComponentArrays: ComponentVector using RingPolymerArrays: get_centroid +using UnitfulAtomic: austrip +using LinearAlgebra: norm +using ComponentArrays: ComponentVector -using NQCModels: NQCModels using NQCDynamics: Estimators, - DynamicsUtils, - AbstractSimulation + DynamicsUtils -using ..DynamicsUtils: +using NQCModels: NQCModels + +using .DynamicsUtils: get_positions, get_velocities, get_quantum_subsystem -struct EnsembleSaver{N,F,S<:AbstractSimulation} - function_names::NTuple{N, Symbol} - output_functions::F - sim::S -end +using ..InitialConditions: QuantisedDiatomic -function EnsembleSaver(function_names::NTuple{N, Symbol}, sim) where {N} - output_functions = tuple([getfield(DynamicsOutputs, f) for f in function_names]...) - EnsembleSaver(function_names, output_functions, sim) -end +""" + OutputVelocity(sol, i) -function output_template(output::EnsembleSaver, savepoints, u0) - sol = (t=savepoints, u=[u0 for _ in savepoints]) - out = output(sol, 0)[1] - for i in eachindex(out) - for j in eachindex(out[i]) - out[i][j] = zero(out[i][j]) - end - end - return out -end +Output the velocity at each timestep during the trajectory. +""" +OutputVelocity(sol, i) = [copy(get_velocities(u)) for u in sol.u] +export OutputVelocity -function (output::EnsembleSaver)(sol, _) - out = [Any[] for _ in 1:length(sol.t)] - for (i, (t, u)) in enumerate(zip(sol.t, sol.u)) - sizehint!(out[i], length(output.function_names)+1) - push!(out[i], t) - evaluate_output!(out[i], u, t, output.sim, output.output_functions...) - end - return (out, false) -end +""" + OutputCentroidVelocity(sol, i) +Output the velocity of the ring polymer centroid at each timestep during the trajectory. """ - create_saving_callback(quantities::NTuple{N, Symbol}; saveat=[]) where {N} +OutputCentroidVelocity(sol, i) = [get_centroid(get_velocities(u)) for u in sol.u] +export OutputCentroidVelocity -Get the `SavingCallback` that will populate `saved_values` with the result obtained -by evaluating the functions provided in `function_names`. """ -function create_saving_callback(function_names::NTuple{N, Symbol}; saveat=[]) where {N} - saved_values = SavedValues(Float64, Vector{Any}) - saving_function = OutputSaver(function_names) - SavingCallback(saving_function, saved_values; saveat=saveat), saved_values -end + OutputPosition(sol, i) -struct OutputSaver{N,F} - function_names::NTuple{N, Symbol} - output_functions::F -end +Output the position at each timestep during the trajectory. +""" +OutputPosition(sol, i) = [copy(get_positions(u)) for u in sol.u] +export OutputPosition """ - OutputSaver(function_names::NTuple{N, Symbol}) where {N} + OutputCentroidPosition(sol, i) -Used to obtain the outputs for all functions given in `function_names`. +Output the position of the ring polymer centroid at each timestep during the trajectory. """ -function OutputSaver(function_names::NTuple{N, Symbol}) where {N} - output_functions = tuple([getfield(DynamicsOutputs, f) for f in function_names]...) - OutputSaver(function_names, output_functions) -end +OutputCentroidPosition(sol, i) = [get_centroid(get_position(u)) for u in sol.u] +export OutputCentroidPosition """ - (output::OutputSaver)(u, t, integrator) + OutputPotentialEnergy(sol, i) -Evaluates every function listed in `output.function_names` and returns all the results. +Output the adiabatic potential energy at each timestep during the trajectory. """ -function (output::OutputSaver)(u, t, integrator) - out = Any[] - sizehint!(out, length(output.function_names)) - evaluate_output!(out, u, t, integrator.p, output.output_functions...) - return out -end +OutputPotentialEnergy(sol, i) = DynamicsUtils.classical_potential_energy.(sol.prob.p, sol.u) +export OutputPotentialEnergy """ -Used to recursively evaluate every function for the [`OutputSaver`](@ref). + OutputTotalEnergy(sol, i) -See here for a description of why it is written like this: -https://stackoverflow.com/questions/55840333/type-stability-for-lists-of-closures +Evaluate the classical Hamiltonian at each timestep during the trajectory. """ -function evaluate_output!(out, u, t, sim, f::F, output_functions...) where {F} - push!(out, f(u, t, sim)) - return evaluate_output!(out, u, t, sim, output_functions...) -end -evaluate_output!(out, u, t, sim) = out +OutputTotalEnergy(sol, i) = DynamicsUtils.classical_hamiltonian.(sol.prob.p, sol.u) +export OutputTotalEnergy + +""" + OutputKineticEnergy(sol, i) + +Evaluate the classical kinetic energy at each timestep during the trajectory. +""" +OutputKineticEnergy(sol, i) = DynamicsUtils.classical_kinetic_energy.(sol.prob.p, sol.u) +export OutputKineticEnergy + +""" + OutputDynamicsVariables(sol, i) + +Output all of the dynamics variables at each timestep during the trajectory. +""" +OutputDynamicsVariables(sol, i) = copy.(sol.u) +export OutputDynamicsVariables + +""" + OutputQuantumSubsystem(sol, i) + +Output the quantum subsystem at each timestep during the trajectory. +Usually this will refer to a wavefunction or density matrix but will depend on the particular dynamics method. +""" +OutputQuantumSubsystem(sol, i) = [copy(get_quantum_subsystem(u)) for u in sol.u] +export OutputQuantumSubsystem + +""" + OutputMappingPosition(sol, i) -export force -export velocity -export position -export centroid_position -export potential -export hamiltonian -export kinetic -export u -export quantum_subsystem -export state -export noise -export population -export adiabatic_population -export friction +Output the position mapping variables at each timestep during the trajectory. +""" +OutputMappingPosition(sol, i) = [copy(DynamicsUtils.get_mapping_positions(u)) for u in sol.u] +export OutputMappingPosition -"Get the force" -force(u, t, sim) = -copy(sim.calculator.derivative) +""" + OutputMappingMomentum(sol, i) -"Get the velocity" -velocity(u, t, sim) = copy(get_velocities(u)) +Output the momentum mapping variable at each timestep during the trajectory. +""" +OutputMappingMomentum(sol, i) = [copy(DynamicsUtils.get_mapping_momenta(u)) for u in sol.u] +export OutputMappingMomentum -"Get the velocity of the ring polymer centroid" -centroid_velocity(u, t, sim) = get_centroid(get_velocities(u)) +""" + OutputDiscreteState(sol, i) -"Get the position" -position(u, t, sim) = copy(get_positions(u)) +Output the discrete state variable at each timestep during the trajectory. +This is used for surface hopping simulations and returns the variable that determines the currently occupied adiabatic state. -"Get the position of the ring polymer centroid" -centroid_position(u, t, sim) = get_centroid(get_positions(u)) +Requires that the dynamics variable has a field `state`. -"Evaluate the potential from the model" -potential(u, t, sim) = DynamicsUtils.classical_potential_energy(sim, u) +Use [`OutputDiabaticPopulation`](@ref) or [`OutputAdiabaticPopulation`](@ref) to get the population estimators. +""" +OutputDiscreteState(sol, i) = [copy(u.state) for u in sol.u] +export OutputDiscreteState -"Evaluate the classical Hamiltonian" -hamiltonian(u, t, sim) = DynamicsUtils.classical_hamiltonian(sim, u) +""" + OutputDiabaticPopulation(sol, i) -"Evaluate the classical kinetic energy" -kinetic(u, t, sim) = DynamicsUtils.classical_kinetic_energy(sim, get_velocities(u)) +Output the diabatic population at each timestep during the trajectory. +""" +OutputDiabaticPopulation(sol, i) = Estimators.diabatic_population.(sol.prob.p, sol.u) +export OutputDiabaticPopulation -"Get all the dynamics variables. This is the default" -u(u, t, sim) = copy(u) -u(u::ArrayPartition, t, sim) = ComponentVector(v=copy(get_velocities(u)), r=copy(get_positions(u))) +""" + OutputTotalDiabaticPopulation(sol, i) +Output the total diabatic population at eah timestep during the trajectory. """ -Get the quantum subsystem of the dynamics variables. -Requires that `DynamicsUtils.get_quantum_subsystem` is implemented for your chosen method. +OutputTotalDiabaticPopulation(sol, i) = sum.(Estimators.diabatic_population(sol.prob.p, sol.u)) +export OutputTotalDiabaticPopulation + """ -quantum_subsystem(u, t, sim) = copy(get_quantum_subsystem(u)) + OutputAdiabaticPopulation(sol, i) -mapping_position(u, t, sim) = copy(DynamicsUtils.get_mapping_positions(u)) -mapping_momentum(u, t, sim) = copy(DynamicsUtils.get_mapping_momenta(u)) +Output the adiabatic population at each timestep during the trajectory. +""" +OutputAdiabaticPopulation(sol, i) = Estimators.adiabatic_population.(sol.prob.p, sol.u) +export OutputAdiabaticPopulation """ -Get the currently occupied state from the dynamics variables. -Requires that the dynamics variable has a field `state`. -Currently this is for surface hopping methods only. -Use [`population`](@ref) or [`adiabatic_population`](@ref) for most other methods. + OutputTotalAdiabaticPopulation(sol, i) + +Output the total adiabatic population at each timestep during the trajectory. """ -state(u, t, sim) = copy(u.state) +OutputTotalAdiabaticPopulation(sol, i) = sum.(Estimators.adiabatic_population.(sol.prob.p, sol.u)) +export OutputTotalAdiabaticPopulation -"Evaluate the diabatic population" -population(u, t, sim) = Estimators.diabatic_population(sim, u) -total_population(u, t, sim) = sum(Estimators.diabatic_population(sim, u)) +""" +Output the end point of each trajectory. +""" +OutputFinal(sol, i) = last(sol.u) +export OutputFinal -"Evaluate the adiabatic population" -adiabatic_population(u, t, sim) = Estimators.adiabatic_population(sim, u) +""" +Output a 1 if the molecule has dissociated, 0 otherwise. +""" +struct OutputDissociation{T} + "The maximum distance at which the two atoms can be considered bonded." + distance::T + "The indices of the two atoms in the molecule of interest." + atom_indices::Tuple{Int,Int} + OutputDissociation(distance, atom_indices) = new{typeof(distance)}(austrip(distance), atom_indices) +end +export OutputDissociation + +function (output::OutputDissociation)(sol, i) + R = DynamicsUtils.get_positions(last(sol.u)) + dissociated = norm(R[:,output.atom_indices[1]] .- R[:,output.atom_indices[2]]) > output.distance + return dissociated ? 1 : 0 +end + +""" +Output the vibrational and rotational quantum numbers of the final image. +""" +struct OutputQuantisedDiatomic{S,H,V} + sim::S + height::H + normal_vector::V +end +OutputQuantisedDiatomic(sim; height=10, normal_vector=[0, 0, 1]) = OutputQuantisedDiatomic(sim, height, normal_vector) +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 + +""" +Output a `ComponentVector` with fields `reflection` and `transmission` containing +the probability of the outcome. +Each index in the arrays refers to the adiabatic state. +""" +struct OutputStateResolvedScattering1D{S} + sim::S + type::Symbol +end +function (output::OutputStateResolvedScattering1D)(sol, i) + final = last(sol.u) # get final configuration from trajectory + if output.type == :adiabatic + populations = Estimators.adiabatic_population(output.sim, final) + elseif output.type == :diabatic + populations = Estimators.diabatic_population(output.sim, final) + else + throw(ArgumentError("$(output.type) not recognised. + Only `:diabatic` or `:adiabatic` accepted.")) + end + output = ComponentVector( + reflection=zeros(NQCModels.nstates(output.sim)), + transmission=zeros(NQCModels.nstates(output.sim)) + ) + x = DynamicsUtils.get_positions(final)[1] + if x > 0 # If final position past 0 then we count as transmission + output.transmission .= populations + else # If final position left of 0 then we count as reflection + output.reflection .= populations + end + return output +end +export OutputStateResolvedScattering1D end # module diff --git a/src/DynamicsUtils/plot.jl b/src/DynamicsUtils/plot.jl index f0a55e132..3cdddde0e 100644 --- a/src/DynamicsUtils/plot.jl +++ b/src/DynamicsUtils/plot.jl @@ -1,14 +1,14 @@ -using RecipesBase -using TypedTables: Table +using RecipesBase: RecipesBase +using Dictionaries: Dictionary -@recipe function f(table::Table, quantity::Symbol) +RecipesBase.@recipe function f(dict::Dictionary, quantity::Symbol) xguide --> "t" yguide --> String(quantity) - for i in eachindex(eval(:($table.$quantity[1]))) - @series begin + for i in eachindex(dict[quantity][1]) + RecipesBase.@series begin legend --> :false - table.t, eval(:([value[$i] for value in $table.$quantity])) + dict[:Time], [value[i] for value in dict[quantity]] end end end diff --git a/src/Ensembles/Ensembles.jl b/src/Ensembles/Ensembles.jl index 88b315129..6b708706b 100644 --- a/src/Ensembles/Ensembles.jl +++ b/src/Ensembles/Ensembles.jl @@ -1,141 +1,130 @@ """ Ensembles -This module provides the main function [`run_ensemble`](@ref Ensembles.run_ensemble). +This module provides the main function [`run_dynamics`](@ref run_dynamics). This serves to run multiple trajectories for a given simulation type, sampling from an initial distribution. """ module Ensembles -using SciMLBase: SciMLBase, EnsembleProblem, solve -using RecursiveArrayTools: ArrayPartition -using TypedTables: TypedTables +using SciMLBase: SciMLBase +using Dictionaries: Dictionary +using UnitfulAtomic: austrip using NQCBase: NQCBase using NQCDynamics: AbstractSimulation, - Simulation, - RingPolymerSimulation, - DynamicsUtils, - DynamicsMethods, - DynamicsOutputs -using RingPolymerArrays: RingPolymerArray + DynamicsMethods -export run_ensemble +export run_dynamics include("selections.jl") -include("outputs.jl") + include("reductions.jl") +export SumReduction +export MeanReduction +export AppendReduction +export FileReduction + +""" + EnsembleSaver{F<:Tuple} +Store a tuple of functions with the signature `f(sol)` where `sol` is a DiffEq solution object. +`EnsembleSaver` will evaluate each of these functions and return the result in a `Dictionary`. """ - run_ensemble(sim::AbstractSimulation, tspan, distribution; +struct EnsembleSaver{F<:Tuple} + functions::F +end + +function (output::EnsembleSaver)(sol, i) + out = Dictionary{Symbol,Any}([:Time], [sol.t]) + return (evaluate_output_functions!(out, sol, i, output.functions...), false) +end + +function evaluate_output_functions!(out::Dictionary, sol, i, f::F, fs...) where {F} + if f isa Function + name = nameof(f) + else + name = nameof(typeof(f)) + end + + insert!(out, name, f(sol, i)) + return evaluate_output_functions!(out, sol, i, fs...) +end +evaluate_output_functions!(out::Dictionary, sol, i) = out + +""" + run_dynamics(sim::AbstractSimulation, tspan, distribution; + output, selection::Union{Nothing,AbstractVector}=nothing, - output=(sol,i)->(sol,false), - reduction=:append - reduction::Symbol=:append, - ensemble_algorithm=SciMLBase.EnsembleThreads(), + reduction=AppendReduction(), + ensemble_algorithm=SciMLBase.EnsembleSerial(), algorithm=DynamicsMethods.select_algorithm(sim), - saveat=[], trajectories=1, kwargs... ) -Run multiple trajectories for timespan `tspan` sampling from `distribution`. -The DifferentialEquations ensemble interface is used which allows us to specify -functions to modify the output and how it is reduced across trajectories. +Run trajectories for timespan `tspan` sampling from `distribution`. # Keywords +* `output` either a single function or a Tuple of functions with the signature `f(sol, i)` that takes the DifferentialEquations solution and returns the desired output quantity. * `selection` should be an `AbstractVector` containing the indices to sample from the `distribution`. By default, `nothing` leads to random sampling. -* `output` can be a function that transforms the DiffEq solution to an output, or a tuple of output quantities as for `run_trajectory`. -* `reduction` defines how the data is reduced across trajectories. Options are `:append`, `:mean` or `:sum`. +* `reduction` defines how the data is reduced across trajectories. Options are `AppendReduction()`, `MeanReduction()`, `SumReduction` and `FileReduction(filename)`. * `ensemble_algorithm` is the algorithm from DifferentialEquations which determines which form of parallelism is used. * `algorithm` is the algorithm used to integrate the equations of motion. -* `saveat` mirrors the DiffEq keyword and labels the time points to save the output. * `trajectories` is the number of trajectories to perform. - -This function wraps the EnsembleProblem from DifferentialEquations and passes the `kwargs` -to the `solve` function. +* `kwargs...` any additional keywords are passed to DifferentialEquations `solve``. """ -function run_ensemble( +function run_dynamics( sim::AbstractSimulation, tspan, distribution; + output, selection::Union{Nothing,AbstractVector}=nothing, - output=SciMLBase.DEFAULT_OUTPUT_FUNC, - reduction::Symbol=:append, - ensemble_algorithm=SciMLBase.EnsembleThreads(), + reduction=AppendReduction(), + ensemble_algorithm=SciMLBase.EnsembleSerial(), algorithm=DynamicsMethods.select_algorithm(sim), - saveat=[], trajectories=1, - u_init=nothing, kwargs... ) + if !(output isa Tuple) + output = (output,) + end + kwargs = NQCBase.austrip_kwargs(;kwargs...) trajectories = convert(Int, trajectories) - saveat = austrip.(saveat) tspan = austrip.(tspan) - (output isa Symbol) && (output = (output,)) - reduction = Reduction(reduction) - - ensemble_problem = EnsembleProblem(sim, tspan, distribution, selection, output, reduction, u_init, saveat, trajectories, kwargs) - @info "Performing $trajectories trajectories." - stats = @timed solve(ensemble_problem, algorithm, ensemble_algorithm; saveat, trajectories, kwargs...) - @info "Finished after $(stats.time) seconds." + prob_func = Selection(distribution, selection, trajectories) - return extract_output(stats, output, reduction, trajectories) -end - -function SciMLBase.EnsembleProblem(sim::AbstractSimulation, tspan, distribution, selection, output, reduction, u_init, saveat, trajectories, kwargs) - - prob_func = Selection(distribution, selection) - - problem = create_problem(sim, distribution, tspan) - output_func = Output(output, sim) - if output_func isa DynamicsOutputs.EnsembleSaver - u_init = get_u_init(reduction, saveat, kwargs, tspan, problem.u0, output_func) - end - - return EnsembleProblem(problem; prob_func, output_func, reduction, u_init) -end - -Output(output::Tuple, sim) = DynamicsOutputs.EnsembleSaver(output, sim) -Output(output, _) = output - -function create_problem(sim::AbstractSimulation, distribution, tspan) u0 = sample_distribution(sim, distribution) - return DynamicsMethods.create_problem(u0, tspan, sim) -end + problem = DynamicsMethods.create_problem(u0, tspan, sim) -function extract_output(stats, output, reduction, trajectories) - sol = stats.value + output_func = EnsembleSaver(output) - if output isa Tuple - process_output(reduction, sol, output, trajectories) + ensemble_problem = SciMLBase.EnsembleProblem(problem; prob_func, output_func, reduction) + + if trajectories == 1 + @info "Performing 1 trajectory." else - return sol.u + @info "Performing $trajectories trajectories." end -end -function process_output(::AppendReduction, sol, output, trajectories) - return [generate_table(output, s) for s in sol.u] -end + stats = @timed SciMLBase.solve(ensemble_problem, algorithm, ensemble_algorithm; trajectories, kwargs...) + @info "Finished after $(stats.time) seconds." -function process_output(::SumReduction, sol, output, trajectories) - for i in eachindex(sol.u) - sol.u[i][1] /= trajectories + if trajectories == 1 + return stats.value.u[1] + else + return stats.value.u end - return generate_table(output, sol.u) end -function process_output(::MeanReduction, sol, output, trajectories) - return generate_table(output, sol.u) -end - -function generate_table(output, u) - output = (:t, output...) - out = [(;zip(output, quantity)...) for quantity in u] - return TypedTables.Table(out) -end +run_ensemble(args...; kwargs...) = throw(error(""" +`run_ensemble` has been replaced, use `run_dynamics` instead. +`run_dynamics` unifies single trajectory and ensemble simulations into one function. +Refer to the `run_dynamics` docstring for more information. +""" +)) end # module diff --git a/src/Ensembles/outputs.jl b/src/Ensembles/outputs.jl deleted file mode 100644 index aecd7a103..000000000 --- a/src/Ensembles/outputs.jl +++ /dev/null @@ -1,88 +0,0 @@ - -using UnitfulAtomic: austrip -using LinearAlgebra: norm -using ComponentArrays: ComponentVector - -using ..InitialConditions: QuantisedDiatomic -using NQCDynamics: Estimators -using NQCModels: nstates -using NQCDistributions: Diabatic, Adiabatic - -function output_template(output, u0) - return zero(output(ComponentVector(u=[u0]), 1)[1]) -end - -""" -Output the end point of each trajectory. -""" -struct OutputFinal end - -(::OutputFinal)(sol, i) = (last(sol.u), false) - -""" -Output a 1 if the molecule has dissociated, 0 otherwise. -""" -struct OutputDissociation{T} - "The maximum distance at which the two atoms can be considered bonded." - distance::T - "The indices of the two atoms in the molecule of interest." - atom_indices::Tuple{Int,Int} - OutputDissociation(distance, atom_indices) = new{typeof(distance)}(austrip(distance), atom_indices) -end - -function (output::OutputDissociation)(sol, i) - R = DynamicsUtils.get_positions(last(sol.u)) - dissociated = norm(R[:,output.atom_indices[1]] .- R[:,output.atom_indices[2]]) > output.distance - return dissociated ? (1, false) : (0, false) -end - -""" -Output the vibrational and rotational quantum numbers of the final image. -""" -struct OutputQuantisedDiatomic{S,H,V} - sim::S - height::H - normal_vector::V -end -OutputQuantisedDiatomic(sim; height=10, normal_vector=[0, 0, 1]) = OutputQuantisedDiatomic(sim, height, normal_vector) - -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), false) -end -output_template(::OutputQuantisedDiatomic, u0) = (0.0, 0.0) - -""" -Output a `ComponentVector` with fields `reflection` and `transmission` containing -the probability of the outcome. -Each index in the arrays refers to the adiabatic state. -""" -struct OutputStateResolvedScattering1D{S} - sim::S - type::Symbol -end -function (output::OutputStateResolvedScattering1D)(sol, i) - final = last(sol.u) # get final configuration from trajectory - if output.type == :adiabatic - populations = Estimators.adiabatic_population(output.sim, final) - elseif output.type == :diabatic - populations = Estimators.diabatic_population(output.sim, final) - else - throw(ArgumentError("$(output.type) not recognised. - Only `:diabatic` or `:adiabatic` accepted.")) - end - output = ComponentVector( - reflection=zeros(nstates(output.sim)), - transmission=zeros(nstates(output.sim)) - ) - x = DynamicsUtils.get_positions(final)[1] - if x > 0 # If final position past 0 then we count as transmission - output.transmission .= populations - else # If final position left of 0 then we count as reflection - output.reflection .= populations - end - return (output, false) -end diff --git a/src/Ensembles/reductions.jl b/src/Ensembles/reductions.jl index 8fbb6511e..ed81fe627 100644 --- a/src/Ensembles/reductions.jl +++ b/src/Ensembles/reductions.jl @@ -1,24 +1,48 @@ using StatsBase: mean -using NQCDynamics: TimeCorrelationFunctions, DynamicsOutputs -using .TimeCorrelationFunctions: TimeCorrelationFunction +using HDF5: HDF5 +using Dictionaries: Dictionary """ Sum the outputs from each trajectory. """ -struct SumReduction end +mutable struct SumReduction + initialised::Bool + SumReduction() = new(false) +end function (reduction::SumReduction)(u, batch, I) + if !reduction.initialised + u = get_initial_u(batch) + reduction.initialised = true + end sum_outputs!(u, batch) return (u, false) end +function get_initial_u(batch) + template = first(batch) + return Dictionary(keys(template), [zero.(t) for t in template]) +end + +function sum_outputs!(u, batch) + for b in batch + u .+= b + end +end + """ Average the outputs over all trajectories. """ mutable struct MeanReduction counter::Int + initialised::Bool + MeanReduction() = new(0, false) end function (reduction::MeanReduction)(u, batch, I) + if !reduction.initialised + u = get_initial_u(batch) + reduction.initialised = true + end u .*= reduction.counter reduction.counter += length(batch) sum_outputs!(u, batch) @@ -26,71 +50,41 @@ function (reduction::MeanReduction)(u, batch, I) return (u, false) end -function sum_outputs!(u, batch) - for i=1:length(batch) - if size(batch[i]) != size(u) - @warn "DimensionMismatch encountered when reducing outputs. \ - Discarding trajectory. Check `u_init` matches the output size." - else - u .+= batch[i] - end - end -end - struct AppendReduction end (::AppendReduction)(u,data,I) = (append!(u,data), false) -""" - Reduction(reduction::Symbol) - -Converts reduction keyword into function that performs reduction. +struct FileReduction + filename::String -* `:discard` is used internally when using callbacks to save outputs instead. -""" -function Reduction(reduction::Symbol) - if reduction === :mean - return MeanReduction(0) - elseif reduction === :sum - return SumReduction() - elseif reduction === :append - return AppendReduction() - else - throw(ArgumentError("`reduction` $reduction not recognised.")) + function FileReduction(filename::String) + splitname = splitext(filename) + ext = splitname[2] + if (ext == ".h5") || (ext == ".hdf5") + return new(filename) + else + filename = splitname[1] * ".h5" + return new(filename) + end end end -function get_u_init(::Union{MeanReduction, SumReduction}, saveat, stripped_kwargs, tspan, u0, output::TimeCorrelationFunction) - savepoints = get_savepoints(saveat, stripped_kwargs, tspan) - u_init = [TimeCorrelationFunctions.correlation_template(output) for _ ∈ savepoints] - return u_init -end - -function get_u_init(::Union{MeanReduction, SumReduction}, saveat, stripped_kwargs, tspan, u0, output::DynamicsOutputs.EnsembleSaver) - savepoints = get_savepoints(saveat, stripped_kwargs, tspan) - u_init = DynamicsOutputs.output_template(output, savepoints, u0) - return u_init -end - -function get_u_init(::Union{MeanReduction, SumReduction}, saveat, stripped_kwargs, tspan, u0, output) - return output_template(output, u0) -end - -function get_u_init(::AppendReduction, saveat, stripped_kwargs, tspan, u0, output) - return nothing -end - -function get_savepoints(saveat, stripped_kwargs, tspan) - if saveat != [] - if saveat isa Number - savepoints = tspan[1]:saveat:tspan[2] - else - savepoints = saveat +function (reduction::FileReduction)(u, batch, I) + HDF5.h5open(reduction.filename, "w") do file + for (i, trajectory_id) in enumerate(I) + trajectory_group = HDF5.create_group(file, "trajectory_$(trajectory_id)") + trajectory = batch[i] + for k in keys(trajectory) + value = trajectory[k] + if (value isa Array{<:Number}) || (value isa Number) + trajectory_group[string(k)] = trajectory[k] + elseif (value isa Vector{<:Array}) + output = cat(value...; dims=3) + trajectory_group[string(k)] = output + else + throw(error("Cannot convert output type to HDF5 format")) + end + end end - elseif haskey(stripped_kwargs, :dt) - dt = stripped_kwargs[:dt] - savepoints = tspan[1]:dt:tspan[2] - else - throw(ArgumentError("Make sure to pass either `saveat` or `dt` as keyword arguments - to compute average/sum over many trajectories.")) end + return ("Output written to $(reduction.filename).", false) end diff --git a/src/Ensembles/selections.jl b/src/Ensembles/selections.jl index a4866cb73..feb7db307 100644 --- a/src/Ensembles/selections.jl +++ b/src/Ensembles/selections.jl @@ -1,6 +1,4 @@ -using DiffEqBase: CallbackSet -using NQCDynamics: InitialConditions, DynamicsOutputs using NQCDistributions: DynamicalDistribution, ProductDistribution abstract type AbstractSelection end @@ -22,13 +20,17 @@ struct RandomSelection{D} <: AbstractSelection distribution::D end -function Selection(distribution, selection::AbstractVector) - @info "Sampling the provided distribution in the range $selection." +function Selection(distribution, selection::AbstractVector, trajectories) + if trajectories > 1 + @info "Sampling the provided distribution in the range $selection." + end OrderedSelection(distribution, selection) end -function Selection(distribution, ::Any) - @info "Sampling randomly from provided distribution." +function Selection(distribution, ::Any, trajectories) + if trajectories > 1 + @info "Sampling randomly from provided distribution." + end RandomSelection(distribution) end @@ -63,3 +65,5 @@ function sample_distribution(sim::AbstractSimulation, distribution::ProductDistr DynamicsMethods.DynamicsVariables(sim, u.v, u.r, distribution.electronic) end +sample_distribution(::AbstractSimulation, distribution) = copy(distribution) + diff --git a/src/NQCDynamics.jl b/src/NQCDynamics.jl index d37031f8f..cda72f802 100644 --- a/src/NQCDynamics.jl +++ b/src/NQCDynamics.jl @@ -26,8 +26,6 @@ export DynamicsUtils include("Estimators.jl") export Estimators -include("DynamicsOutputs.jl") - include("TimeCorrelationFunctions.jl") export TimeCorrelationFunctions @@ -41,4 +39,7 @@ include("Ensembles/Ensembles.jl") export Ensembles @reexport using .Ensembles +include("DynamicsOutputs.jl") +@reexport using .DynamicsOutputs + end # module diff --git a/src/TimeCorrelationFunctions.jl b/src/TimeCorrelationFunctions.jl index 348a8d223..7e7179fd3 100644 --- a/src/TimeCorrelationFunctions.jl +++ b/src/TimeCorrelationFunctions.jl @@ -36,7 +36,7 @@ function (correlation::TimeCorrelationFunction)(sol, _) correlate!(out[i], normalisation, initial_value, final_value, correlation) end - return (out, false) + return out end evaluate_normalisation(sim::AbstractSimulation, correlation::TimeCorrelationFunction) = 1 diff --git a/test/Dynamics/bcbwithtsit5.jl b/test/Dynamics/bcbwithtsit5.jl index b3bca6591..91dbe14ab 100644 --- a/test/Dynamics/bcbwithtsit5.jl +++ b/test/Dynamics/bcbwithtsit5.jl @@ -8,10 +8,10 @@ using Random: seed! u = DynamicsVariables(sim, fill(10/2000, size(sim)), fill(-5, size(sim)) .+ randn(size(sim)), PureState(1)) dt = 0.1 - sol = run_trajectory(u, (0, 2000.0), sim; algorithm=DynamicsMethods.IntegrationAlgorithms.BCBwithTsit5(), saveat=0:10:2000, dt=dt) - sol1 = run_trajectory(u, (0, 2000.0), sim; algorithm=Tsit5(), saveat=0:10:2000, abstol=1e-6, reltol=1e-6) + sol = run_dynamics(sim, (0, 2000.0), u; output=OutputDynamicsVariables, algorithm=DynamicsMethods.IntegrationAlgorithms.BCBwithTsit5(), saveat=0:10:2000, dt=dt) + sol1 = run_dynamics(sim, (0, 2000.0), u; output=OutputDynamicsVariables, algorithm=Tsit5(), saveat=0:10:2000, abstol=1e-6, reltol=1e-6) - @test sol.u ≈ sol1.u rtol=1e-3 + @test sol[:OutputDynamicsVariables] ≈ sol1[:OutputDynamicsVariables] rtol=1e-3 end @testset "FSSH" begin @@ -20,9 +20,9 @@ end dt = 0.1 seed!(1) - sol = run_trajectory(u, (0, 2000.0), sim; algorithm=DynamicsMethods.IntegrationAlgorithms.BCBwithTsit5(), saveat=0:10:2000, dt=dt) + sol = run_dynamics(sim, (0, 2000.0), u; output=OutputDynamicsVariables, algorithm=DynamicsMethods.IntegrationAlgorithms.BCBwithTsit5(), saveat=0:10:2000, dt=dt) seed!(1) - sol1 = run_trajectory(u, (0, 2000.0), sim; algorithm=Tsit5(), saveat=0:10:2000, dt=dt, adaptive=false) + sol1 = run_dynamics(sim, (0, 2000.0), u; output=OutputDynamicsVariables, algorithm=Tsit5(), saveat=0:10:2000, dt=dt, adaptive=false) - @test sol.u ≈ sol1.u rtol=1e-3 + @test sol[:OutputDynamicsVariables] ≈ sol1[:OutputDynamicsVariables] rtol=1e-3 end diff --git a/test/Dynamics/cell_boundary_callback.jl b/test/Dynamics/cell_boundary_callback.jl index da8c2e8d3..cb5b2dde9 100644 --- a/test/Dynamics/cell_boundary_callback.jl +++ b/test/Dynamics/cell_boundary_callback.jl @@ -9,10 +9,10 @@ sim = Simulation(atoms, model; cell=cell) z = ComponentVector(v=fill(1.0, 1, length(atoms)), r=zeros(1, length(atoms))) -solution = run_trajectory(z, (0.0, 300), sim; dt=0.1) -Rs = vcat(get_positions.(solution.u)...) +solution = run_dynamics(sim, (0.0, 300), z; output=OutputPosition, dt=0.1) +Rs = reduce(vcat, solution[:OutputPosition]) @test !isapprox(minimum(Rs), 0, atol=1e-4) | (minimum(Rs) >= 0) # check atom leaves cell -solution = run_trajectory(z, (0.0, 300), sim; callback=DynamicsUtils.CellBoundaryCallback(), dt=0.1) -Rs = vcat(get_positions.(solution.u)...) +solution = run_dynamics(sim, (0.0, 300), z; output=OutputPosition, callback=DynamicsUtils.CellBoundaryCallback(), dt=0.1) +Rs = reduce(vcat, solution[:OutputPosition]) @test isapprox(minimum(Rs), 0, atol=1e-4) | (minimum(Rs) >= 0) # Check atom remains inside diff --git a/test/Dynamics/classical.jl b/test/Dynamics/classical.jl index 92571e62c..3dd18024b 100644 --- a/test/Dynamics/classical.jl +++ b/test/Dynamics/classical.jl @@ -19,8 +19,8 @@ model = NQCModels.Harmonic() r = get_blank(sim) u0 = ComponentVector(v=v, r=r) - sol = run_trajectory(u0, (0.0, 1000.0), sim; dt=0.1, output=(:hamiltonian)) - @test sol.hamiltonian[1] ≈ sol.hamiltonian[end] rtol=1e-2 + sol = run_dynamics(sim, (0.0, 1000.0), u0; dt=0.1, output=(OutputTotalEnergy)) + @test sol[:OutputTotalEnergy][1] ≈ sol[:OutputTotalEnergy][end] rtol=1e-2 end @testset "Ring polymer classical" begin @@ -34,8 +34,8 @@ end r = get_blank(sim) u0 = ComponentVector(v=v, r=r) - sol = run_trajectory(u0, (0.0, 1000.0), sim; dt=0.1, output=(:hamiltonian)) - @test sol.hamiltonian[1] ≈ sol.hamiltonian[end] rtol=1e-2 + sol = run_dynamics(sim, (0.0, 1000.0), u0; dt=0.1, output=(OutputTotalEnergy)) + @test sol[:OutputTotalEnergy][1] ≈ sol[:OutputTotalEnergy][end] rtol=1e-2 end @testset "Fermion model classical adiabatic dynamics" begin @@ -44,8 +44,8 @@ end v = rand(VelocityBoltzmann(300u"K", atoms.masses, size(sim))) r = [model.model.morse.x₀;;] u0 = DynamicsVariables(sim, v, r) - sol = run_trajectory(u0, (0.0, 900.0u"fs"), sim; dt=1u"fs", output=(:hamiltonian, :position)) - @test sol.hamiltonian[1] ≈ sol.hamiltonian[end] rtol=1e-2 + sol = run_dynamics(sim, (0.0, 900.0u"fs"), u0; dt=1u"fs", output=(OutputTotalEnergy)) + @test sol[:OutputTotalEnergy][1] ≈ sol[:OutputTotalEnergy][end] rtol=1e-2 end @testset "Fermion model ring polymer adiabatic dynamics" begin @@ -56,6 +56,6 @@ end r = model.model.morse.x₀ d = DynamicalDistribution(v, r, size(sim)) u0 = rand(d) - sol = run_trajectory(u0, (0.0, 900.0u"fs"), sim; dt=1u"fs", output=(:hamiltonian, :position)) - @test sol.hamiltonian[1] ≈ sol.hamiltonian[end] rtol=1e-2 + sol = run_dynamics(sim, (0.0, 900.0u"fs"), u0; dt=1u"fs", output=(OutputTotalEnergy)) + @test sol[:OutputTotalEnergy][1] ≈ sol[:OutputTotalEnergy][end] rtol=1e-2 end diff --git a/test/Dynamics/cmm.jl b/test/Dynamics/cmm.jl index 47d11d87c..e44bcd483 100644 --- a/test/Dynamics/cmm.jl +++ b/test/Dynamics/cmm.jl @@ -29,16 +29,18 @@ u1 = DynamicsVariables(sim1, v, r, PureState(1)) test_motion!(sim, u) test_motion!(sim1, u1) -sol = run_trajectory(u, (0, 100.0), sim; output=(:hamiltonian, :position, :u), reltol=1e-10, abstol=1e-10) -@test sol.hamiltonian[1] ≈ sol.hamiltonian[end] rtol=1e-3 -qmap = [u.qmap for u in sol.u] -pmap = [u.pmap for u in sol.u] +sol = run_dynamics(sim, (0, 100.0), u; output=(OutputTotalEnergy, OutputPosition, OutputDynamicsVariables), reltol=1e-10, abstol=1e-10) +@test sol[:OutputTotalEnergy][1] ≈ sol[:OutputTotalEnergy][end] rtol=1e-3 +qmap = [u.qmap for u in sol[:OutputDynamicsVariables]] +pmap = [u.pmap for u in sol[:OutputDynamicsVariables]] total_population = sum.(DynamicsMethods.MappingVariableMethods.mapping_kernel.(qmap, pmap, sim.method.γ)) @test all(isapprox.(total_population, 1, rtol=1e-3)) -sol = run_trajectory(u1, (0, 100.0), sim1; output=(:hamiltonian, :position, :u), reltol=1e-10, abstol=1e-10) -@test sol.hamiltonian[1] ≈ sol.hamiltonian[end] rtol=1e-3 -total_population = sum.(DynamicsMethods.MappingVariableMethods.mapping_kernel.(qmap, pmap, sim.method.γ)) +sol = run_dynamics(sim1, (0, 100.0), u1; output=(OutputTotalEnergy, OutputPosition, OutputDynamicsVariables), reltol=1e-10, abstol=1e-10) +@test sol[:OutputTotalEnergy][1] ≈ sol[:OutputTotalEnergy][end] rtol=1e-3 +qmap = [u.qmap for u in sol[:OutputDynamicsVariables]] +pmap = [u.pmap for u in sol[:OutputDynamicsVariables]] +total_population = sum.(DynamicsMethods.MappingVariableMethods.mapping_kernel.(qmap, pmap, sim1.method.γ)) @test all(isapprox.(total_population, 1, rtol=1e-3)) @testset "generate_random_points_on_nsphere" begin diff --git a/test/Dynamics/diabatic_mdef.jl b/test/Dynamics/diabatic_mdef.jl index a70309ab4..dcf28d06a 100644 --- a/test/Dynamics/diabatic_mdef.jl +++ b/test/Dynamics/diabatic_mdef.jl @@ -38,12 +38,12 @@ end @testset "Simulation{DiabaticMDEF}" begin sim = Simulation{DiabaticMDEF}(atoms, model; temperature=T, friction_method=ClassicalMethods.WideBandExact(model.ρ)) u = DynamicsVariables(sim, rand(1,1), rand(1,1)) - run_trajectory(u, (0.0, 1.0), sim; dt=0.1) + run_dynamics(sim, (0.0, 1.0), u; dt=0.1, output=OutputDynamicsVariables) end @testset "RingPolymerSimulation{DiabaticMDEF}" begin n_beads = 10 sim = RingPolymerSimulation{DiabaticMDEF}(atoms, model, n_beads; temperature=T, friction_method=ClassicalMethods.WideBandExact(model.ρ)) u = DynamicsVariables(sim, rand(1,1,n_beads), rand(1,1,n_beads)) - run_trajectory(u, (0.0, 1.0), sim; dt=0.1) + run_dynamics(sim, (0.0, 1.0), u; dt=0.1, output=OutputDynamicsVariables) end diff --git a/test/Dynamics/ehrenfest.jl b/test/Dynamics/ehrenfest.jl index f2f779a53..7b7b4948f 100644 --- a/test/Dynamics/ehrenfest.jl +++ b/test/Dynamics/ehrenfest.jl @@ -39,16 +39,16 @@ atoms = Atoms(:H) r = fill(-5.0, size(sim)) v1 = fill(8.9, size(sim)) ./ sim1.atoms.masses[1] z1 = DynamicsVariables(sim, v1, r, PureState(1, Adiabatic())) - solution1 = run_trajectory(z1, (0.0, 2500.0), sim1; output=(:population), reltol=1e-6) + solution1 = run_dynamics(sim1, (0.0, 2500.0), z1; output=OutputDiabaticPopulation, reltol=1e-6) - @testset "run_trajectory" begin + @testset "run_dynamics" begin atoms = Atoms(2000) sim = Simulation{Ehrenfest}(atoms, NQCModels.TullyModelTwo()) v = hcat(100 / 2000) r = hcat(-10.0) u = DynamicsVariables(sim, v, r, PureState(1, Adiabatic())) - solution = run_trajectory(u, (0.0, 500.0), sim, output=(:hamiltonian), reltol=1e-6) - @test solution.hamiltonian[1] ≈ solution.hamiltonian[end] rtol=1e-2 + solution = run_dynamics(sim, (0.0, 500.0), u, output=OutputTotalEnergy, reltol=1e-6) + @test solution[:OutputTotalEnergy][1] ≈ solution[:OutputTotalEnergy][end] rtol=1e-2 end end @@ -58,6 +58,6 @@ end v = fill(100 / 2000, 1, 1, 10) r = fill(-10.0, 1, 1, 10) u = DynamicsVariables(sim, v, r, PureState(1, Adiabatic())) - solution = run_trajectory(u, (0.0, 500.0), sim, output=(:hamiltonian), dt=1) - @test solution.hamiltonian[1] ≈ solution.hamiltonian[end] rtol=1e-2 + solution = run_dynamics(sim, (0.0, 500.0), u, output=OutputTotalEnergy, dt=1) + @test solution[:OutputTotalEnergy][1] ≈ solution[:OutputTotalEnergy][end] rtol=1e-2 end diff --git a/test/Dynamics/fssh.jl b/test/Dynamics/fssh.jl index 41beb16ee..0956247c9 100644 --- a/test/Dynamics/fssh.jl +++ b/test/Dynamics/fssh.jl @@ -97,14 +97,14 @@ atoms = Atoms(2) @test ΔKE ≈ -ΔE rtol=1e-3 # Test for energy conservation end - @testset "run_trajectory" begin + @testset "run_dynamics" begin atoms = Atoms(2000) sim = Simulation{FSSH}(atoms, NQCModels.TullyModelTwo()) v = hcat(100 / 2000) r = hcat(-10.0) u = DynamicsVariables(sim, v, r, PureState(1, Adiabatic())) - solution = run_trajectory(u, (0.0, 500.0), sim, output=(:hamiltonian, :state), reltol=1e-6) - @test solution.hamiltonian[1] ≈ solution.hamiltonian[end] rtol=1e-2 + solution = run_dynamics(sim, (0.0, 500.0), u, output=OutputTotalEnergy, reltol=1e-6) + @test solution[:OutputTotalEnergy][1] ≈ solution[:OutputTotalEnergy][end] rtol=1e-2 end end @@ -163,18 +163,18 @@ end @test ΔKE ≈ -sum(ΔE) rtol=1e-3 # Test for energy conservation end - @testset "run_trajectory" begin + @testset "run_dynamics" begin atoms = Atoms(2000) sim = RingPolymerSimulation{FSSH}(atoms, TullyModelTwo(), 5; temperature=0.01) v = fill(100 / 2000, size(sim)) r = fill(-10.0, size(sim)) .+ randn(1,1,5) u = DynamicsVariables(sim, v, r, PureState(1, Adiabatic())) seed!(1) - solution = run_trajectory(u, (0.0, 1000.0), sim, output=(:hamiltonian), dt=0.1) + solution = run_dynamics(sim, (0.0, 1000.0), u, output=OutputTotalEnergy, dt=0.1) seed!(1) - solution2 = run_trajectory(u, (0.0, 1000.0), sim, output=(:hamiltonian), algorithm=Tsit5(), abstol=1e-10, reltol=1e-10, saveat=0:0.1:1000.0) + solution2 = run_dynamics(sim, (0.0, 1000.0), u, output=OutputTotalEnergy, algorithm=Tsit5(), abstol=1e-10, reltol=1e-10, saveat=0:0.1:1000.0) # Ring polymer Hamiltonian is not strictly conserved during hoppping - @test solution.hamiltonian[1] ≈ solution.hamiltonian[end] rtol=1e-2 + @test solution[:OutputTotalEnergy][1] ≈ solution[:OutputTotalEnergy][end] rtol=1e-2 end end diff --git a/test/Dynamics/langevin.jl b/test/Dynamics/langevin.jl index 5f5f2a96a..5fb2c2776 100644 --- a/test/Dynamics/langevin.jl +++ b/test/Dynamics/langevin.jl @@ -8,4 +8,4 @@ sim = Simulation{Langevin}(atoms, model; temperature=1, γ=1) z = DynamicsVariables(sim, randn(1,1), randn(1,1)) -solution = run_trajectory(z, (0.0, 500.0), sim; dt=1) +solution = run_dynamics(sim, (0.0, 500.0), z; output=OutputDynamicsVariables, dt=1) diff --git a/test/Dynamics/mdef.jl b/test/Dynamics/mdef.jl index 75f868774..f08bcd584 100644 --- a/test/Dynamics/mdef.jl +++ b/test/Dynamics/mdef.jl @@ -20,11 +20,11 @@ du = zero(u) @test all(diag(gtmp) .≈ 1.0) end -sol = run_trajectory(u, (0.0, 100.0), sim; dt=1) -@test sol.u[1] ≈ u +sol = run_dynamics(sim, (0.0, 100.0), u; output=OutputDynamicsVariables, dt=1) +@test sol[:OutputDynamicsVariables][1] ≈ u f(t) = 100u"K"*exp(-ustrip(t)) model = CompositeFrictionModel(Harmonic(), RandomFriction(1)) sim = Simulation{MDEF}(atoms, model; temperature=f) -sol = run_trajectory(u, (0.0, 100.0), sim; dt=1) +sol = run_dynamics(sim, (0.0, 100.0), u; output=OutputDynamicsVariables, dt=1) \ No newline at end of file diff --git a/test/Dynamics/nrpmd.jl b/test/Dynamics/nrpmd.jl index 37bec2f5d..da79133d4 100644 --- a/test/Dynamics/nrpmd.jl +++ b/test/Dynamics/nrpmd.jl @@ -52,14 +52,14 @@ end algs = (DynamicsMethods.IntegrationAlgorithms.RingPolymerMInt(), Tsit5()) @testset "Energy conservation $alg" for alg in algs - sol = run_trajectory(u, (0, 10.0), sim; output=(:hamiltonian, :population), dt=1e-2, algorithm=alg, abstol=1e-8, reltol=1e-8) - @test sol.hamiltonian[1] ≈ sol.hamiltonian[end] rtol=1e-2 + sol = run_dynamics(sim, (0, 10.0), u; output=OutputTotalEnergy, dt=1e-2, algorithm=alg, abstol=1e-8, reltol=1e-8) + @test sol[:OutputTotalEnergy][1] ≈ sol[:OutputTotalEnergy][end] rtol=1e-2 end @testset "Algorithm comparison" begin - sol = run_trajectory(u, (0, 10.0), sim; dt=1e-2, algorithm=DynamicsMethods.IntegrationAlgorithms.RingPolymerMInt()) - sol1 = run_trajectory(u, (0, 10.0), sim; algorithm=Tsit5(), reltol=1e-10, abstol=1e-10, saveat=sol.t) - @test sol.u ≈ sol1.u rtol=1e-2 + sol = run_dynamics(sim, (0, 10.0), u; output=OutputDynamicsVariables, dt=1e-2, algorithm=DynamicsMethods.IntegrationAlgorithms.RingPolymerMInt()) + sol1 = run_dynamics(sim, (0, 10.0), u; output=OutputDynamicsVariables, algorithm=Tsit5(), reltol=1e-10, abstol=1e-10, saveat=sol[:Time]) + @test sol[:OutputDynamicsVariables] ≈ sol1[:OutputDynamicsVariables] rtol=1e-2 end @testset "MInt algorithm convergence" begin diff --git a/test/Dynamics/rpecmm.jl b/test/Dynamics/rpecmm.jl index 26bc35d16..c4611ba6e 100644 --- a/test/Dynamics/rpecmm.jl +++ b/test/Dynamics/rpecmm.jl @@ -31,16 +31,18 @@ u1 = DynamicsVariables(sim1, v, r, PureState(1)) test_motion!(sim, u) test_motion!(sim1, u1) -sol = run_trajectory(u, (0, 100.0), sim; output=(:hamiltonian, :position, :u), dt=0.01) -@test sol.hamiltonian[1] ≈ sol.hamiltonian[end] rtol=1e-3 -qmap = [u.qmap for u in sol.u] -pmap = [u.pmap for u in sol.u] +sol = run_dynamics(sim, (0, 100.0), u; output=(OutputTotalEnergy, OutputMappingPosition, OutputMappingMomentum), dt=0.01) +@test sol[:OutputTotalEnergy][1] ≈ sol[:OutputTotalEnergy][end] rtol=1e-3 +qmap = sol[:OutputMappingPosition] +pmap = sol[:OutputMappingMomentum] total_population = sum.(DynamicsMethods.MappingVariableMethods.mapping_kernel.(qmap, pmap, sim.method.γ)) @test all(isapprox.(total_population, 1, rtol=1e-2)) -sol = run_trajectory(u1, (0, 100.0), sim1; output=(:hamiltonian, :position, :u), dt=0.01) -@test sol.hamiltonian[1] ≈ sol.hamiltonian[end] rtol=1e-3 -total_population = sum.(DynamicsMethods.MappingVariableMethods.mapping_kernel.(qmap, pmap, sim.method.γ)) +sol = run_dynamics(sim1, (0, 100.0), u1; output=(OutputTotalEnergy, OutputMappingPosition, OutputMappingMomentum), dt=0.01) +@test sol[:OutputTotalEnergy][1] ≈ sol[:OutputTotalEnergy][end] rtol=1e-3 +qmap = sol[:OutputMappingPosition] +pmap = sol[:OutputMappingMomentum] +total_population = sum.(DynamicsMethods.MappingVariableMethods.mapping_kernel.(qmap, pmap, sim1.method.γ)) @test all(isapprox.(total_population, 1, rtol=1e-2)) @testset "Population correlation" begin @@ -64,10 +66,10 @@ total_population = sum.(DynamicsMethods.MappingVariableMethods.mapping_kernel.(q end @testset "Algorithm comparison" begin - sol = run_trajectory(u, (0, 10.0), sim; dt=1e-2, algorithm=DynamicsMethods.IntegrationAlgorithms.RingPolymerMInt()) - sol1 = run_trajectory(u, (0, 10.0), sim; algorithm=Tsit5(), reltol=1e-10, abstol=1e-10, saveat=sol.t) - @test sol.u ≈ sol1.u rtol=1e-2 - sol = run_trajectory(u, (0, 10.0), sim1; dt=1e-2, algorithm=DynamicsMethods.IntegrationAlgorithms.RingPolymerMInt()) - sol1 = run_trajectory(u, (0, 10.0), sim1; algorithm=Tsit5(), reltol=1e-10, abstol=1e-10, saveat=sol.t) - @test sol.u ≈ sol1.u rtol=1e-2 + sol = run_dynamics(sim, (0, 10.0), u; output=OutputDynamicsVariables, dt=1e-2, algorithm=DynamicsMethods.IntegrationAlgorithms.RingPolymerMInt()) + sol1 = run_dynamics(sim, (0, 10.0), u; output=OutputDynamicsVariables, algorithm=Tsit5(), reltol=1e-10, abstol=1e-10, saveat=sol[:Time]) + @test sol[:OutputDynamicsVariables] ≈ sol1[:OutputDynamicsVariables] rtol=1e-2 + sol = run_dynamics(sim1, (0, 10.0), u; output=OutputDynamicsVariables, dt=1e-2, algorithm=DynamicsMethods.IntegrationAlgorithms.RingPolymerMInt()) + sol1 = run_dynamics(sim1, (0, 10.0), u; output=OutputDynamicsVariables, algorithm=Tsit5(), reltol=1e-10, abstol=1e-10, saveat=sol[:Time]) + @test sol[:OutputDynamicsVariables] ≈ sol1[:OutputDynamicsVariables] rtol=1e-2 end diff --git a/test/Dynamics/rpmdef.jl b/test/Dynamics/rpmdef.jl index 54f7d09f7..7f517efd4 100644 --- a/test/Dynamics/rpmdef.jl +++ b/test/Dynamics/rpmdef.jl @@ -57,7 +57,7 @@ end v = RingPolymerArray(zeros(size(sim))) r = RingPolymerArray(zeros(size(sim))) - sol = run_trajectory(ArrayPartition(v,r), (0.0, 1e4), sim; dt=1, output=(:kinetic)) + sol = run_dynamics(sim, (0.0, 1e4), ArrayPartition(v,r); dt=1, output=OutputKineticEnergy) # @test mean(sol.kinetic) ≈ austrip(100u"K") * length(sim.beads) rtol=5e-1 end @@ -69,5 +69,5 @@ end v = RingPolymerArray(zeros(size(sim))) r = RingPolymerArray(zeros(size(sim))) - sol = run_trajectory(ArrayPartition(v,r), (0.0, 1e2), sim; dt=1, output=(:kinetic)) + sol = run_dynamics(sim, (0.0, 1e2), ArrayPartition(v,r); dt=1, output=OutputKineticEnergy) end diff --git a/test/Dynamics/spin_mapping.jl b/test/Dynamics/spin_mapping.jl index e86da4fcb..0c270caec 100644 --- a/test/Dynamics/spin_mapping.jl +++ b/test/Dynamics/spin_mapping.jl @@ -34,14 +34,14 @@ u = DynamicsVariables(sim, v, r, PureState(2)) end @testset "Energy conservation" begin - sol = run_trajectory(u, (0, 10.0), sim; output=:hamiltonian, dt=1e-2) - @test sol.hamiltonian[1] ≈ sol.hamiltonian[end] rtol=1e-2 + sol = run_dynamics(sim, (0, 10.0), u; output=OutputTotalEnergy, dt=1e-2) + @test sol[:OutputTotalEnergy][1] ≈ sol[:OutputTotalEnergy][end] rtol=1e-2 end @testset "Algorithm comparison" begin - sol = run_trajectory(u, (0, 3.0), sim; dt=1e-2, algorithm=DynamicsMethods.IntegrationAlgorithms.MInt()) - sol1 = run_trajectory(u, (0, 3.0), sim; algorithm=Tsit5(), reltol=1e-10, abstol=1e-10, saveat=sol.t) - @test sol.u ≈ sol1.u rtol=1e-2 + sol = run_dynamics(sim, (0, 3.0), u; output=OutputDynamicsVariables, dt=1e-2, algorithm=DynamicsMethods.IntegrationAlgorithms.MInt()) + sol1 = run_dynamics(sim, (0, 3.0), u; output=OutputDynamicsVariables, algorithm=Tsit5(), reltol=1e-10, abstol=1e-10, saveat=sol[:Time]) + @test sol[:OutputDynamicsVariables] ≈ sol1[:OutputDynamicsVariables] rtol=1e-2 end @testset "MInt algorithm convergence" begin diff --git a/test/Ensembles/ensembles.jl b/test/Ensembles/ensembles.jl index 1ac5d5e3e..325986ccc 100644 --- a/test/Ensembles/ensembles.jl +++ b/test/Ensembles/ensembles.jl @@ -1,7 +1,6 @@ using NQCDynamics using Test using RecursiveArrayTools: ArrayPartition -using ComponentArrays: ComponentVector atoms = Atoms([:H]) model = NQCModels.Harmonic() @@ -13,30 +12,34 @@ distribution = DynamicalDistribution(positions, velocities, (1, 1)) u0 = rand(distribution) tspan = (0.0, 10.0) -@testset "run_ensemble" for reduction in (:append, :sum, :mean) - out = run_ensemble(sim, tspan, distribution; reduction, - output=(:position, :velocity, :hamiltonian), dt=1, trajectories=10) - if (reduction === :sum) || (reduction === :mean) - @test out.t == 0.0:10.0 - @test out.position isa Vector{<:AbstractMatrix} - @test out.velocity isa Vector{<:AbstractMatrix} - @test out.hamiltonian isa Vector{<:Number} - else - @test all(x -> x == 0.0:10.0, (traj.t for traj in out)) - @test all(x -> x isa Vector{<:AbstractMatrix}, (traj.position for traj in out)) - @test all(x -> x isa Vector{<:AbstractMatrix}, (traj.velocity for traj in out)) - @test all(x -> x isa Vector{<:Number}, (traj.hamiltonian for traj in out)) +@testset "run_dynamics" for reduction in ( + AppendReduction(), MeanReduction(), SumReduction(), FileReduction("test.h5") +) + out = run_dynamics(sim, tspan, distribution; reduction, + output=(OutputPosition, OutputVelocity, OutputTotalEnergy), dt=1, trajectories=10) + if (reduction isa MeanReduction) || (reduction isa SumReduction) + reduction isa SumReduction && @test out[:Time] == 0.0:10.0:100.0 + reduction isa MeanReduction && @test out[:Time] == 0.0:10.0 + @test out[:OutputPosition] isa Vector{<:AbstractMatrix} + @test out[:OutputVelocity] isa Vector{<:AbstractMatrix} + @test out[:OutputTotalEnergy] isa Vector{<:Number} + elseif (reduction isa AppendReduction) + @test all(x -> x == 0.0:10.0, (traj[:Time] for traj in out)) + @test all(x -> x isa Vector{<:AbstractMatrix}, (traj[:OutputPosition] for traj in out)) + @test all(x -> x isa Vector{<:AbstractMatrix}, (traj[:OutputVelocity] for traj in out)) + @test all(x -> x isa Vector{<:Number}, (traj[:OutputTotalEnergy] for traj in out)) end end -@testset "run_ensemble" begin - out = run_ensemble(sim, tspan, distribution; - output=Ensembles.OutputFinal(), dt=1, trajectories=10, reduction=:append) - @test out isa Vector{<:ArrayPartition} +@testset "run_dynamics" begin + out = run_dynamics(sim, tspan, distribution; + output=OutputFinal, dt=1, trajectories=10, reduction=AppendReduction()) + final = [o[:OutputFinal] for o in out] + @test final isa Vector{<:ArrayPartition} end -@testset "run_ensemble, reduction=$reduction" for reduction ∈ (:sum, :mean) - out = run_ensemble(sim, tspan, distribution; - output=Ensembles.OutputFinal(), dt=1, trajectories=10, reduction=reduction, u_init=u0) - @test out isa ComponentVector +@testset "run_dynamics, reduction=$reduction" for reduction ∈ (SumReduction(), MeanReduction()) + out = run_dynamics(sim, tspan, distribution; + output=OutputFinal, dt=1, trajectories=10, reduction=reduction) + @test out[:OutputFinal] isa ArrayPartition end diff --git a/test/Ensembles/outputs.jl b/test/Ensembles/outputs.jl index 04a54a2be..3744ffa99 100644 --- a/test/Ensembles/outputs.jl +++ b/test/Ensembles/outputs.jl @@ -3,31 +3,29 @@ using Test using ComponentArrays: ComponentVector @testset "OutputFinal" begin - output = Ensembles.OutputFinal() + output = OutputFinal sol = ComponentVector(u=1:10) - @test output(sol, 1) == (10, false) + @test output(sol, 1) == 10 end @testset "OutputDissociation" begin - output = Ensembles.OutputDissociation(1.0, (1, 2)) + output = OutputDissociation(1.0, (1, 2)) r = [1 0; 0 0; 0 0] v = zeros(3, 2) u = ComponentVector(v=v, r=r) sol = ComponentVector(u=[u]) - @test output(sol, 1) == (0, false) + @test output(sol, 1) == 0 r = [2 0; 0 0; 0 0] u = ComponentVector(v=v, r=r) sol = ComponentVector(u=[u]) - @test output(sol, 1) == (1, false) - @test Ensembles.output_template(output, u) == 0 + @test output(sol, 1) == 1 end @testset "OutputQuantisedDiatomic" begin sim = Simulation(Atoms(1), Harmonic()) - output = Ensembles.OutputQuantisedDiatomic(sim; height=1, normal_vector = [1, 1, 1]) - output = Ensembles.OutputQuantisedDiatomic(sim) + output = OutputQuantisedDiatomic(sim; height=1, normal_vector = [1, 1, 1]) + output = OutputQuantisedDiatomic(sim) u = DynamicsVariables(sim, hcat(1), hcat(1)) - @test Ensembles.output_template(output, u) == (0, 0) end @testset "PopulationCorrelationFunction" begin @@ -36,8 +34,8 @@ end sim = Simulation{Ehrenfest}(Atoms(2000), TullyModelTwo()) output = TimeCorrelationFunctions.PopulationCorrelationFunction(sim, Adiabatic()) u = DynamicsVariables(sim, hcat(20/2000), hcat(-5), PureState(1, Adiabatic())) - sol = run_trajectory(u, (0, 1000.0), sim) - out, cont = output(sol, 1) + sol = run_dynamics(sim, (0, 1000.0), u; output) + out = sol[:PopulationCorrelationFunction] @test out[1][1,1] ≈ 1 @test out[1][2,2] ≈ 0 @test [o[1,2] for o in out] ≈ [1 - o[1,1] for o in out] @@ -47,8 +45,8 @@ end sim = Simulation{FSSH}(Atoms(2000), TullyModelTwo()) output = TimeCorrelationFunctions.PopulationCorrelationFunction(sim, Adiabatic()) u = DynamicsVariables(sim, hcat(20/2000), hcat(-5), PureState(1, Adiabatic())) - sol = run_trajectory(u, (0, 1000.0), sim) - out, cont = output(sol, 1) + sol = run_dynamics(sim, (0, 1000.0), u; output) + out = sol[:PopulationCorrelationFunction] @test out[1][1,1] ≈ 1 @test out[1][2,2] ≈ 0 @test [o[1,2] for o in out] ≈ [1 - o[1,1] for o in out] @@ -58,8 +56,8 @@ end sim = Simulation{eCMM}(Atoms(2000), TullyModelTwo()) output = TimeCorrelationFunctions.PopulationCorrelationFunction(sim, Diabatic()) u = DynamicsVariables(sim, hcat(20/2000), hcat(-5), PureState(1, Diabatic())) - sol = run_trajectory(u, (0, 1000.0), sim) - out, cont = output(sol, 1) + sol = run_dynamics(sim, (0, 1000.0), u; output) + out = sol[:PopulationCorrelationFunction] end end diff --git a/test/Ensembles/reductions.jl b/test/Ensembles/reductions.jl index eba8c5308..2cad52a65 100644 --- a/test/Ensembles/reductions.jl +++ b/test/Ensembles/reductions.jl @@ -3,55 +3,15 @@ using NQCDynamics using Test using OrdinaryDiffEq -@testset "Reduction" begin - @test Ensembles.Reduction(:mean) isa Ensembles.MeanReduction - @test Ensembles.Reduction(:sum) isa Ensembles.SumReduction - @test Ensembles.Reduction(:append) isa Ensembles.AppendReduction -end - -@testset "get_u_init" begin - - @testset "Output, reduction=$reduction" for reduction in (:mean, :sum) - reduction = Ensembles.Reduction(reduction) - tspan = (0.0, 1.0) - u0 = 1.0 - saveat = 1 - output = Ensembles.OutputFinal() - u_init = Ensembles.get_u_init(reduction, saveat, Dict(), tspan, u0, output) - @test u_init ≈ zero(u0) - end - - @testset "reduction=:append" begin - reduction = Ensembles.Reduction(:append) - output = Ensembles.OutputFinal() - u0 = 1 - saveat = 1 - u_init = Ensembles.get_u_init(reduction, saveat, Dict(), (0, 1), u0, output) - @test u_init === nothing - end - - @testset "TimeCorrelationFunction, reduction=$reduction" for reduction in (:mean, :sum) - reduction = Ensembles.Reduction(reduction) - sim = Simulation{FSSH}(Atoms(1), DoubleWell()) - saveat = 0.1 - tspan = (0.0, 1.0) - u0 = DynamicsVariables(sim, 1, 0, PureState(1, Adiabatic())) - output = TimeCorrelationFunctions.PopulationCorrelationFunction(sim, Diabatic()) - u_init = Ensembles.get_u_init(reduction, saveat, Dict(), tspan, u0, output) - @test u_init ≈ [zeros(2, 2) for _ ∈ tspan[1]:0.1:tspan[2]] - end - -end - -@testset "$reduction" for reduction in (Ensembles.MeanReduction(0), Ensembles.SumReduction()) +@testset "$reduction" for reduction in (MeanReduction(), SumReduction()) prob = ODEProblem((u,p,t) -> 1.01u, [0.5], (0.0,1.0)) prob_func(prob, i, repeat) = remake(prob,u0=[rand()]) - output = Ensembles.OutputFinal() - ensemble = EnsembleProblem(prob, prob_func=prob_func, output_func=output, reduction=reduction, u_init=[0.0]) + output = Ensembles.EnsembleSaver((OutputFinal,)) + ensemble = EnsembleProblem(prob, prob_func=prob_func, output_func=output, reduction=reduction) sol = solve(ensemble, Tsit5(), trajectories=1e3, batch_size=20) if reduction isa Ensembles.SumReduction - @test sol.u[1] / 1e3 ≈ 0.5 * exp(1.01) rtol=1e-1 + @test sol.u[:OutputFinal][1] / 1e3 ≈ 0.5 * exp(1.01) rtol=1e-1 else - @test sol.u[1] ≈ 0.5 * exp(1.01) rtol=1e-1 + @test sol.u[:OutputFinal][1] ≈ 0.5 * exp(1.01) rtol=1e-1 end end