Print timings

Author: jagot

Project.toml

@@ -11,6 +11,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 

src/StrongFieldApproximation.jl

@@ -9,6 +9,7 @@ using LinearAlgebra
 using StaticArrays
 
 using ProgressMeter
+using TimerOutputs
 
 include("threading.jl")
 include("telescope_iterators.jl")

src/dyson_expansions.jl

@@ -424,77 +424,82 @@ function recombination(system::System, diagram::Diagram, 𝐩, 𝐀, i)
     end
 end
 
-function integrate_diagram(::Type{Amp}, system::System, diagram::Diagram, iref, 𝐤=nothing; memory=typemax(Int), imin=1) where Amp
-    ions, unique_momenta, momenta, indeterminate_momenta, order = analyze_diagram(system, diagram)
+function integrate_diagram(::Type{Amp}, system::System, diagram::Diagram, iref, 𝐤=nothing; memory=typemax(Int), imin=1,
+                           to=TimerOutput(), verbosity=1) where Amp
+    ions, unique_momenta, momenta, indeterminate_momenta, order = @timeit to "Analyze diagram" analyze_diagram(system, diagram)
 
-    weight = (-im*system.dt)^order
+    @timeit to "Allocations" begin
+        weight = (-im*system.dt)^order
 
-    # println()
-    # @info "Integrating diagram up to" iref system diagram ions unique_momenta momenta indeterminate_momenta order weight 𝐤
+        verbosity > 1 && @info "Integrating diagram up to" iref system diagram ions unique_momenta momenta indeterminate_momenta order weight 𝐤
 
-    # @show
-    Eᵢₒₙₛ = [system.ionization_channels[ion].E for ion in ions]
-    𝐩s = complex(zeros(momentum_type(system, 𝐤), length(unique_momenta)))
-    prefactors = ones(complex(eltype(system.t)), length(unique_momenta))
-    if !isnothing(𝐤)
-        𝐩s[1] = 𝐤
+        Eᵢₒₙₛ = [system.ionization_channels[ion].E for ion in ions]
+        𝐩s = complex(zeros(momentum_type(system, 𝐤), length(unique_momenta)))
+        prefactors = ones(complex(eltype(system.t)), length(unique_momenta))
+        if !isnothing(𝐤)
+            𝐩s[1] = 𝐤
+        end
+
+        𝐀 = system.𝐀
+        amplitude = complex(zero(Amp))
+        ctT = complex(eltype(system.t))
+
+        is = vcat(iref, zeros(Int, order))
     end
-    # @show 𝐩s
 
-    𝐀 = system.𝐀
-    amplitude = complex(zero(Amp))
-    ctT = complex(eltype(system.t))
-
-    is = vcat(iref, zeros(Int, order))
-
-    for i in TelescopeIterator(max(1,imin):iref-1, order, memory)
-        is[2:end] .= i
-        # is = vcat(iref, i)
-        # is = (iref,i...)
-        # for i in 1:iref-1, j in 1:i-1
-        #     is = (iref,i,j)
-        # println(is)
-
-        evaluate_momenta!(𝐩s, prefactors, system, unique_momenta, indeterminate_momenta, is)
-        # @show 𝐩s prefactors
-
-        Sᵢₒₙ = zero(ctT)
-        Sₑₗ = zero(ctT)
-        for j = 1:order
-            a,b = is[j], is[j+1]
-            τ = system.t[a] - system.t[b]
-            Sᵢₒₙ += Eᵢₒₙₛ[j]*τ
-            Sₑₗ += volkov_phase(𝐩s[j], system.volkov, a, b)
-        end
-        S = Sᵢₒₙ + Sₑₗ
-        aₚᵣₒₚ = prod(prefactors)*exp(-im*S)
+    @timeit to "Time loop" begin
+        for i in TelescopeIterator(max(1,imin):iref-1, order, memory)
+            is[2:end] .= i
+            # is = vcat(iref, i)
+            # is = (iref,i...)
+            # for i in 1:iref-1, j in 1:i-1
+            #     is = (iref,i,j)
+            # println(is)
+
+            @timeit to "Evaluate momenta" evaluate_momenta!(𝐩s, prefactors, system, unique_momenta, indeterminate_momenta, is)
+            verbosity > 10 && @show 𝐩s prefactors
+
+            @timeit to "Evaluate propagators" begin
+                Sᵢₒₙ = zero(ctT)
+                Sₑₗ = zero(ctT)
+                for j = 1:order
+                    a,b = is[j], is[j+1]
+                    τ = system.t[a] - system.t[b]
+                    Sᵢₒₙ += Eᵢₒₙₛ[j]*τ
+                    Sₑₗ += volkov_phase(𝐩s[j], system.volkov, a, b)
+                end
+                S = Sᵢₒₙ + Sₑₗ
+                aₚᵣₒₚ = prod(prefactors)*exp(-im*S)
+            end
 
-        # @show is prefactors
+            verbosity > 10 && @show is prefactors
 
-        ∂a = (ionization(system, diagram, 𝐩s[end], 𝐀, is[end]) *
-              aₚᵣₒₚ *
-              recombination(system, diagram, 𝐩s[1], 𝐀, iref))
+            ∂a = @timeit to "Prefactor" (ionization(system, diagram, 𝐩s[end], 𝐀, is[end]) *
+                                         aₚᵣₒₚ *
+                                         recombination(system, diagram, 𝐩s[1], 𝐀, iref))
 
-        # Loop over "interior" interactions
-        for j = (order>1 && first(diagram)[2]==0 ? 3 : 2):order
-            ion,which = diagram.path[j-1]
-            α = ions[j-1]
-            β = ions[j]
-            # @show j, ion, which α,β
-            interaction = system.couplings[which][α,β]
+            @timeit to "Interior interactions" begin
+                # Loop over "interior" interactions
+                for j = (order>1 && first(diagram)[2]==0 ? 3 : 2):order
+                    ion,which = diagram.path[j-1]
+                    α = ions[j-1]
+                    β = ions[j]
+                    verbosity > 20 && @show j, ion, which α,β
+                    interaction = system.couplings[which][α,β]
 
-            𝐤ᵢ = 𝐩s[momenta[j-1]]
-            𝐩ᵢ = 𝐩s[momenta[j]]
-            𝐀ᵢ = 𝐀[is[j]]
+                    𝐤ᵢ = 𝐩s[momenta[j-1]]
+                    𝐩ᵢ = 𝐩s[momenta[j]]
+                    𝐀ᵢ = 𝐀[is[j]]
 
-            ∂a *= interaction(kinematic_momentum(𝐤ᵢ, 𝐀ᵢ), kinematic_momentum(𝐩ᵢ, 𝐀ᵢ), is[j+1])
-        end
+                    ∂a *= @timeit to "Interaction" interaction(kinematic_momentum(𝐤ᵢ, 𝐀ᵢ), kinematic_momentum(𝐩ᵢ, 𝐀ᵢ), is[j+1])
+                end
+            end
 
-        # ∂a = prod(prefactors)
-        amplitude += ∂a
+            amplitude += ∂a
+        end
     end
 
-    weight*amplitude
+    @timeit to "Weighting" weight*amplitude
 end
 
 # * High-level interface
@@ -507,11 +512,18 @@ function photoelectron_spectrum(k::AbstractArray{T},
 
     cT = complex(eltype(T))
     c = similar(k, cT)
+    to = TimerOutput()
     p = Progress(length(k))
-    threaded_range_loop(eachindex(k)) do i
-        c[i] = integrate_diagram(cT, system, diagram, iref, k[i]; kwargs...)
-        ProgressMeter.next!(p)
+    @timeit to "k loop" begin
+        threaded_range_loop(eachindex(k)) do i
+            tok = TimerOutput()
+            c[i] = integrate_diagram(cT, system, diagram, iref, k[i]; to=tok, verbosity=verbosity-1, kwargs...)
+            ProgressMeter.next!(p)
+            merge!(to, tok, tree_point=["k loop"])
+        end
     end
+    TimerOutputs.complement!(to)
+    verbosity > 0 && print_timer(to)
     c
 end
 
@@ -537,7 +549,7 @@ function induced_dipole(system::System, diagram::Diagram; verbosity = 1, kwargs.
     memory = get(kwargs, :memory, typemax(Int))
 
     @showprogress for (i,t) in enumerate(t)
-        𝐝̃ = integrate_diagram(DT, system, diagram, i; imin=i-memory, kwargs...)
+        𝐝̃ = integrate_diagram(DT, system, diagram, i; imin=i-memory, verbosity=verbosity-1, kwargs...)
         𝐝[i] = 2real(𝐝̃)
     end