From e67acdbf6272e647813c65abde1d9d91c56eebbc Mon Sep 17 00:00:00 2001 From: Michael Goerz Date: Thu, 7 Dec 2023 16:44:11 -0500 Subject: [PATCH] Store `timer_data` in all propagators This makes the piecewise-constant propagators consistent with the OrdinaryDiffEq propagator. --- docs/src/benchmarks/profiling.md | 4 +-- ext/QuantumPropagatorsODEExt.jl | 4 +-- src/cheby.jl | 12 +++++-- src/cheby_propagator.jl | 22 +++++++++---- src/exp_propagator.jl | 54 ++++++++++++++++++-------------- src/expprop.jl | 21 ++++++++++--- src/newton.jl | 5 ++- src/newton_propagator.jl | 26 ++++++--------- src/propagator.jl | 8 +++++ 9 files changed, 94 insertions(+), 62 deletions(-) diff --git a/docs/src/benchmarks/profiling.md b/docs/src/benchmarks/profiling.md index c435d932..36193e0a 100644 --- a/docs/src/benchmarks/profiling.md +++ b/docs/src/benchmarks/profiling.md @@ -77,7 +77,7 @@ In any subsequent propagation, we could access the timing data in a `callback` t ```@example profiling function show_timing_data(propagator, args...) if propagator.t == tlist[end] - show(propagator.wrk.timing_data, compact=true) + show(propagator.timing_data, compact=true) end end @@ -94,7 +94,7 @@ propagator = init_prop(Ψ₀, H, tlist; method=Cheby) for step ∈ 1:(length(tlist)-1) prop_step!(propagator) end -show(propagator.wrk.timing_data, compact=true) +show(propagator.timing_data, compact=true) ``` The reported runtimes here are less important than the number of function calls and the runtime percentages. In this case, the timing data shows that the propagation is dominated by the matrix-vector products (applying the Hamiltonian to the state), as it should. The percentage would go to 100% for larger Hilbert spaces. diff --git a/ext/QuantumPropagatorsODEExt.jl b/ext/QuantumPropagatorsODEExt.jl index d7eb3852..710a31aa 100644 --- a/ext/QuantumPropagatorsODEExt.jl +++ b/ext/QuantumPropagatorsODEExt.jl @@ -168,7 +168,7 @@ init_prop(state, generator, tlist, method::Val{:DifferentialEquations}; kwargs.. mutable struct ODEContinuousPropagator{IT} <: AbstractPropagator const integrator::IT const tlist::Vector{Float64} - const parameters::AbstractDict + parameters::AbstractDict const backward::Bool const inplace::Bool const timing_data::TimerOutput @@ -181,7 +181,7 @@ mutable struct ODEPWCPropagator{IT} <: PWCPropagator # out because it's in a different place in the type hierarchy const integrator::IT const tlist::Vector{Float64} - const parameters::AbstractDict + parameters::AbstractDict const backward::Bool const inplace::Bool const timing_data::TimerOutput diff --git a/src/cheby.jl b/src/cheby.jl index 65213831..0849b3de 100644 --- a/src/cheby.jl +++ b/src/cheby.jl @@ -95,13 +95,19 @@ mutable struct ChebyWrk{ST,CFS,FT<:AbstractFloat} dt::FT limit::FT timing_data::TimerOutput - function ChebyWrk(Ψ::ST, Δ::FT, E_min::FT, dt::FT; limit::FT=1e-12) where {ST,FT} + function ChebyWrk( + Ψ::ST, + Δ::FT, + E_min::FT, + dt::FT; + limit::FT=1e-12, + _timing_data=TimerOutput() + ) where {ST,FT} v0::ST = similar(Ψ) v1::ST = similar(Ψ) v2::ST = similar(Ψ) coeffs = cheby_coeffs(Δ, dt; limit=limit) n_coeffs = length(coeffs) - timing_data = TimerOutput() new{ST,typeof(coeffs),FT}( v0, v1, @@ -112,7 +118,7 @@ mutable struct ChebyWrk{ST,CFS,FT<:AbstractFloat} E_min, dt, limit, - timing_data + _timing_data ) end end diff --git a/src/cheby_propagator.jl b/src/cheby_propagator.jl index 10eb91ff..66295848 100644 --- a/src/cheby_propagator.jl +++ b/src/cheby_propagator.jl @@ -1,16 +1,16 @@ using .Controls: get_controls, evaluate, discretize -using TimerOutputs: reset_timer!, @timeit_debug +using TimerOutputs: reset_timer!, @timeit_debug, TimerOutput """Propagator for Chebychev propagation (`method=QuantumPropagators.Cheby`). This is a [`PWCPropagator`](@ref). """ mutable struct ChebyPropagator{GT,OT,ST} <: PWCPropagator - generator::GT + const generator::GT state::ST t::Float64 # time at which current `state` is defined n::Int64 # index of next interval to propagate - tlist::Vector{Float64} + const tlist::Vector{Float64} parameters::AbstractDict controls control_ranges::AbstractDict @@ -22,6 +22,7 @@ mutable struct ChebyPropagator{GT,OT,ST} <: PWCPropagator specrange_buffer::Float64 check_normalization::Bool specrange_options::Dict{Symbol,Any} + const timing_data::TimerOutput end @@ -130,7 +131,15 @@ function init_prop( if isnothing(dt) error("Chebychev propagation only works on a uniform time grid") end - wrk = Cheby.ChebyWrk(state, Δ, E_min, dt; limit=cheby_coeffs_limit) + timing_data = TimerOutput() + wrk = Cheby.ChebyWrk( + state, + Δ, + E_min, + dt; + limit=cheby_coeffs_limit, + _timing_data=timing_data + ) n = 1 t = tlist[1] if backward @@ -157,6 +166,7 @@ function init_prop( specrange_buffer, check_normalization, specrange_kwargs, + timing_data, ) end @@ -277,7 +287,7 @@ function reinit_prop!( propagator.control_ranges = control_ranges wrk = Cheby.ChebyWrk(state, Δ, E_min, dt; limit=wrk.limit) else - reset_timer!(propagator.wrk.timing_data) + reset_timer!(propagator.timing_data) end t = float(propagator.backward ? tlist[end] : tlist[1]) _pwc_set_t!(propagator, t) @@ -332,7 +342,7 @@ end function prop_step!(propagator::ChebyPropagator) - @timeit_debug propagator.wrk.timing_data "prop_step!" begin + @timeit_debug propagator.timing_data "prop_step!" begin Ψ = propagator.state H = propagator.genop n = propagator.n diff --git a/src/exp_propagator.jl b/src/exp_propagator.jl index 80b6c17a..83863c5f 100644 --- a/src/exp_propagator.jl +++ b/src/exp_propagator.jl @@ -1,4 +1,5 @@ using .Controls: get_controls +using TimerOutputs: reset_timer!, @timeit_debug, TimerOutput """Propagator for propagation via direct exponentiation (`method=QuantumPropagators.ExpProp`) @@ -6,20 +7,21 @@ using .Controls: get_controls This is a [`PWCPropagator`](@ref). """ mutable struct ExpPropagator{GT,OT,ST,WT<:ExpProp.ExpPropWrk} <: PWCPropagator - generator::GT + const generator::GT state::ST t::Float64 # time at which current `state` is defined n::Int64 # index of next interval to propagate - tlist::Vector{Float64} + const tlist::Vector{Float64} parameters::AbstractDict controls genop::OT - wrk::WT + const wrk::WT backward::Bool inplace::Bool convert_state::Type convert_operator::Type func::Function + const timing_data::TimerOutput end @@ -101,7 +103,8 @@ function init_prop( G = _pwc_get_max_genop(generator, controls, tlist) parameters = _pwc_process_parameters(parameters, controls, tlist) - wrk = ExpProp.ExpPropWrk(convert(convert_state, state)) + timing_data = TimerOutput() + wrk = ExpProp.ExpPropWrk(convert(convert_state, state); _timing_data=timing_data) n = 1 t = tlist[1] if backward @@ -127,6 +130,7 @@ function init_prop( convert_state, convert_operator, func, + timing_data, ) end @@ -137,26 +141,28 @@ init_prop(state, generator, tlist, method::Val{:expprop}; kwargs...) = function prop_step!(propagator::ExpPropagator) - n = propagator.n - tlist = getfield(propagator, :tlist) - (0 < n < length(tlist)) || return nothing - dt = tlist[n+1] - tlist[n] - if propagator.backward - dt = -dt - end - Ψ = convert(propagator.convert_state, propagator.state) - if propagator.inplace - _pwc_set_genop!(propagator, n) - H = convert(propagator.convert_operator, propagator.genop) - ExpProp.expprop!(Ψ, H, dt, propagator.wrk; func=propagator.func) - if Ψ ≢ propagator.state # `convert` of Ψ may have been a no-op - copyto!(propagator.state, convert(typeof(propagator.state), Ψ)) + @timeit_debug propagator.timing_data "prop_step!" begin + n = propagator.n + tlist = getfield(propagator, :tlist) + (0 < n < length(tlist)) || return nothing + dt = tlist[n+1] - tlist[n] + if propagator.backward + dt = -dt + end + Ψ = convert(propagator.convert_state, propagator.state) + if propagator.inplace + _pwc_set_genop!(propagator, n) + H = convert(propagator.convert_operator, propagator.genop) + ExpProp.expprop!(Ψ, H, dt, propagator.wrk; func=propagator.func) + if Ψ ≢ propagator.state # `convert` of Ψ may have been a no-op + copyto!(propagator.state, convert(typeof(propagator.state), Ψ)) + end + else + H = convert(propagator.convert_operator, _pwc_get_genop(propagator, n)) + Ψ = ExpProp.expprop(Ψ, H, dt, propagator.wrk; func=propagator.func) + setfield!(propagator, :state, convert(typeof(propagator.state), Ψ)) end - else - H = convert(propagator.convert_operator, _pwc_get_genop(propagator, n)) - Ψ = ExpProp.expprop(Ψ, H, dt, propagator.wrk; func=propagator.func) - setfield!(propagator, :state, convert(typeof(propagator.state), Ψ)) + _pwc_advance_time!(propagator) + return propagator.state end - _pwc_advance_time!(propagator) - return propagator.state end diff --git a/src/expprop.jl b/src/expprop.jl index 6739a6ae..b6688e4e 100644 --- a/src/expprop.jl +++ b/src/expprop.jl @@ -5,6 +5,7 @@ export ExpPropWrk, expprop! using LinearAlgebra import StaticArrays +using TimerOutputs: @timeit_debug, TimerOutput """ @@ -18,8 +19,9 @@ Initializes the workspace for the propagation of a vector `v0` """ struct ExpPropWrk{T} v::T - function ExpPropWrk(v0::T) where {T} - new{T}(similar(v0)) + timing_data::TimerOutput + function ExpPropWrk(v0::T; _timing_data=TimerOutput()) where {T} + new{T}(similar(v0), _timing_data) end end @@ -39,8 +41,12 @@ Keyword arguments besides `func` are ignored. """ function expprop!(Ψ, H, dt, wrk; func=(H_dt -> exp(-1im * H_dt)), _...) copyto!(wrk.v, Ψ) - U = func(H * dt) - mul!(Ψ, U, wrk.v) + @timeit_debug wrk.timing_data "matrix exponentiation" begin + U = func(H * dt) + end + @timeit_debug wrk.timing_data "matrix-vector product" begin + mul!(Ψ, U, wrk.v) + end end @@ -53,7 +59,12 @@ evaluates `Ψ_out = func(H*dt) Ψ` as in [`expprop!`](@ref), but not acting in-place. """ function expprop(Ψ, H, dt, wrk; func=(H_dt -> exp(-1im * H_dt)), _...) - return func(H * dt) * Ψ + @timeit_debug wrk.timing_data "matrix exponentiation" begin + U = func(H * dt) + end + @timeit_debug wrk.timing_data "matrix-vector product" begin + return U * Ψ + end end end diff --git a/src/newton.jl b/src/newton.jl index fff50458..0ec7051b 100644 --- a/src/newton.jl +++ b/src/newton.jl @@ -34,7 +34,7 @@ mutable struct NewtonWrk{T} restarts::Int64 m_max::Int64 timing_data::TimerOutput - function NewtonWrk(v0::T; m_max::Int64=10) where {T} + function NewtonWrk(v0::T; m_max::Int64=10, _timing_data=TimerOutput()) where {T} if m_max <= 2 error("Newton propagation requires m_max > 2") end @@ -44,7 +44,6 @@ mutable struct NewtonWrk{T} error("Newton propagation requires state dimension > 2") end end - timing_data = TimerOutput() new{T}( T[similar(v0) for _ = 1:m_max+1], # arnoldi_vecs similar(v0), # v @@ -55,7 +54,7 @@ mutable struct NewtonWrk{T} 0, # n_leja 0, # restarts m_max, - timing_data, + _timing_data, ) end end diff --git a/src/newton_propagator.jl b/src/newton_propagator.jl index 2c1dad6c..355527e3 100644 --- a/src/newton_propagator.jl +++ b/src/newton_propagator.jl @@ -6,21 +6,22 @@ using TimerOutputs: reset_timer!, @timeit_debug This is a [`PWCPropagator`](@ref). """ mutable struct NewtonPropagator{GT,OT,ST} <: PWCPropagator - generator::GT + const generator::GT state::ST t::Float64 # time at which current `state` is defined n::Int64 # index of next interval to propagate - tlist::Vector{Float64} + const tlist::Vector{Float64} parameters::AbstractDict controls genop::OT - wrk::Newton.NewtonWrk{ST} + const wrk::Newton.NewtonWrk{ST} backward::Bool inplace::Bool func::Function norm_min::Float64 relerr::Float64 max_restarts::Int64 + const timing_data::TimerOutput end set_t!(propagator::NewtonPropagator, t) = _pwc_set_t!(propagator, t) @@ -78,7 +79,8 @@ function init_prop( G = _pwc_get_max_genop(generator, controls, tlist) parameters = _pwc_process_parameters(parameters, controls, tlist) - wrk = Newton.NewtonWrk(state; m_max=m_max) + timing_data = TimerOutput() + wrk = Newton.NewtonWrk(state; m_max=m_max, _timing_data=timing_data) n = 1 t = tlist[1] if backward @@ -104,7 +106,8 @@ function init_prop( func, norm_min, relerr, - max_restarts + max_restarts, + timing_data, ) end @@ -113,19 +116,8 @@ init_prop(state, generator, tlist, method::Val{:newton}; kwargs...) = init_prop(state, generator, tlist, Val(:Newton); kwargs...) -function reinit_prop!(propagator::NewtonPropagator, state; _...) - set_state!(propagator, state) - if propagator.backward - set_t!(propagator, propagator.tlist[end]) - else - set_t!(propagator, propagator.tlist[begin]) - end - reset_timer!(propagator.wrk.timing_data) -end - - function prop_step!(propagator::NewtonPropagator) - @timeit_debug propagator.wrk.timing_data "prop_step!" begin + @timeit_debug propagator.timing_data "prop_step!" begin Ψ = propagator.state H = propagator.genop n = propagator.n # index of interval we're going to propagate diff --git a/src/propagator.jl b/src/propagator.jl index dd651765..b9bfb280 100644 --- a/src/propagator.jl +++ b/src/propagator.jl @@ -1,4 +1,5 @@ using Printf +using TimerOutputs: reset_timer! """Abstract base type for all `Propagator` objects. @@ -299,6 +300,13 @@ function reinit_prop!(propagator, state; kwargs...) else set_t!(propagator, propagator.tlist[begin]) end + try + reset_timer!(propagator.timing_data) + catch + # All or built-in propagators have `timing_data`, but custom + # propagators might not (it's not part of the required interface). + # Thus, we fail silently. + end end