Store timer_data in all propagators

This makes the piecewise-constant propagators consistent with the
OrdinaryDiffEq propagator.

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.

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

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

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

src/exp_propagator.jl

@@ -1,25 +1,27 @@
 using .Controls: get_controls
+using TimerOutputs: reset_timer!, @timeit_debug, TimerOutput
 
 """Propagator for propagation via direct exponentiation
 (`method=QuantumPropagators.ExpProp`)
 
 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

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

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

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

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