From 25a918b88e34be341dc2e4c75f6e7091c36e095c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Mon, 17 Apr 2023 08:46:46 +0200
Subject: [PATCH 01/20] Pretty-printing

---
 src/channels.jl | 32 ++++++++++++++++++++++++++++++--
 1 file changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/channels.jl b/src/channels.jl
index c28e8c3..fee36ae 100644
--- a/src/channels.jl
+++ b/src/channels.jl
@@ -21,6 +21,17 @@ struct DipoleSourceTerm{Dipole<:Function,Field<:Union{ElectricFields.AbstractFie
         new{Dipole,FieldAmplitude}(d, Fv)
 end
 
+Base.show(io::IO, d::DipoleSourceTerm) =
+    write(io, "DipoleSourceTerm")
+
+function Base.show(io::IO, mime::MIME"text/plain", d::DipoleSourceTerm)
+    show(io, d)
+    print(io, "\nDipole: ")
+    show(io, mime, d.d)
+    print(io, "\nField: ")
+    show(io, d.F)
+end
+
 dipole(F::Number, d::Number) = F*d
 dipole(F::SVector{3}, d::SVector{3}) = dot(F, d)
 dipole(F::Number, d::SVector{3}) = F*d[3]
@@ -54,11 +65,17 @@ IonizationChannel(E, args...) =
 Base.show(io::IO, ic::IonizationChannel) =
     write(io, "IonizationChannel: Iₚ = $(ic.E) Ha = $(27.211ic.E) eV")
 
+function Base.show(io::IO, mime::MIME"text/plain", ic::IonizationChannel)
+    show(io, ic)
+    print(io, "\nSource term: ")
+    show(io, mime, ic.st)
+end
+
 # * Couplings
 
 abstract type AbstractCanonicalMomentumConservation end
-struct CanonicalMomentumConservation <:  AbstractCanonicalMomentumConservation end
-struct NoCanonicalMomentumConservation <:  AbstractCanonicalMomentumConservation end
+struct CanonicalMomentumConservation <: AbstractCanonicalMomentumConservation end
+struct NoCanonicalMomentumConservation <: AbstractCanonicalMomentumConservation end
 
 abstract type AbstractCoupling end
 Base.iszero(::AbstractCoupling) = false
@@ -69,6 +86,8 @@ struct NoCoupling <: AbstractCoupling end
 Base.iszero(::NoCoupling) = true
 Base.zero(::AbstractCoupling) = NoCoupling()
 
+Base.show(io::IO, ::NoCoupling) = write(io, "0")
+
 # ** Dipole couplings
 
 struct DipoleCoupling{DipoleMoment<:Union{Number,SVector{3}},Field<:ElectricFields.AbstractField} <: AbstractCoupling
@@ -79,6 +98,8 @@ end
 canonical_momentum_conservation(::DipoleCoupling) = CanonicalMomentumConservation()
 canonical_momentum_conservation(::Type{<:DipoleCoupling}) = CanonicalMomentumConservation()
 
+Base.show(io::IO, ::DipoleCoupling) = write(io, "𝐅⋅𝐫")
+
 # ** Coulomb couplings
 
 momentum_transfer(k::Number, p::Number) = k-p
@@ -90,3 +111,10 @@ struct CoulombCoupling{Coupling} <: AbstractCoupling
 end
 
 (cc::CoulombCoupling)(𝐤, 𝐩, _) = cc.coupling(𝐤, 𝐩)
+
+Base.show(io::IO, ::CoulombCoupling) = write(io, "ĝ")
+
+function Base.show(io::IO, mime::MIME"text/plain", g::CoulombCoupling)
+    write(io, "CoulombCoupling: ")
+    show(io, mime, g.coupling)
+end

From a0ff45700a0bc67a388e32706489e3d047886cb4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Mon, 17 Apr 2023 08:49:55 +0200
Subject: [PATCH 02/20] Support ionic dipole couplings and tabulated electric
 fields

---
 src/channels.jl         | 11 ++++--
 src/dyson_expansions.jl | 83 +++++++++++++++++++++++++++++------------
 2 files changed, 67 insertions(+), 27 deletions(-)

diff --git a/src/channels.jl b/src/channels.jl
index fee36ae..a57f241 100644
--- a/src/channels.jl
+++ b/src/channels.jl
@@ -36,11 +36,11 @@ dipole(F::Number, d::Number) = F*d
 dipole(F::SVector{3}, d::SVector{3}) = dot(F, d)
 dipole(F::Number, d::SVector{3}) = F*d[3]
 
-source_term(d::DipoleSourceTerm{<:Any,<:ElectricFields.AbstractField}, t, p) =
-    dipole(field_amplitude(d.F, t), d.d(p))
+_field_amplitude(F::ElectricFields.AbstractField, t) =
+    field_amplitude(F, t)
+_field_amplitude(F::AbstractVector, i::Integer) = F[i]
 
-source_term(d::DipoleSourceTerm{<:Any,<:AbstractVector}, i::Integer, p) =
-    dipole(d.F[i], d.d(p))
+source_term(d::DipoleSourceTerm, t, p) = dipole(_field_amplitude(d.F, t), d.d(p))
 
 # * Ionizations channels
 
@@ -54,6 +54,8 @@ Represents an ionization channel with energy `E` above the neutral
 struct IonizationChannel{T,SourceTerm<:AbstractSourceTerm}
     E::T
     st::SourceTerm
+    IonizationChannel(E::T, st::SourceTerm) where {T,SourceTerm} =
+        new{T,SourceTerm}(E, st)
 end
 
 source_term(ic::IonizationChannel, args...) =
@@ -85,6 +87,7 @@ canonical_momentum_conservation(::Type{<:AbstractCoupling}) = NoCanonicalMomentu
 struct NoCoupling <: AbstractCoupling end
 Base.iszero(::NoCoupling) = true
 Base.zero(::AbstractCoupling) = NoCoupling()
+Base.zero(::Type{<:AbstractCoupling}) = NoCoupling()
 
 Base.show(io::IO, ::NoCoupling) = write(io, "0")
 
diff --git a/src/dyson_expansions.jl b/src/dyson_expansions.jl
index 0c51bb4..ec7569c 100644
--- a/src/dyson_expansions.jl
+++ b/src/dyson_expansions.jl
@@ -11,11 +11,13 @@ the external field, and the interactions possible, the actual process
 is described by a [`Diagram`](@ref).
 """
 struct System{T,IonizationChannels<:AbstractVector{<:IonizationChannel{T}},
+              IonDipoleCouplings,
               Couplings<:AbstractVector{<:AbstractMatrix{<:AbstractCoupling}},
               VectorPotential,
               Times<:AbstractRange,
               Volkov<:VolkovPhases}
     ionization_channels::IonizationChannels
+    𝐫ᵢₒₙ::IonDipoleCouplings
 
     couplings::Couplings
 
@@ -26,60 +28,79 @@ struct System{T,IonizationChannels<:AbstractVector{<:IonizationChannel{T}},
     volkov::Volkov
 end
 
+IonDipoleCouplingsType = Union{AbstractMatrix{<:Number},UniformScaling,SVector{3},Nothing}
+NoCouplings = Matrix{AbstractCoupling}[]
+
 """
-    System(ionization_channels, couplings, F, ndt)
+    System(ionization_channels, 𝐫ᵢₒₙ, couplings, F, ndt)
 
 Set up a [`System`](@ref) consisting of multiple
-[`IonizationChannel`](@ref)s with possible `couplings` between them,
-driven by an external field `F`, sampled at a frequency of `fs`.
+[`IonizationChannel`](@ref)s with ionic dipole moments `𝐫ᵢₒₙ` and
+possible `couplings` between them, driven by an external field `F`,
+sampled at a frequency of `fs`.
 """
 function System(ionization_channels::AbstractVector{<:IonizationChannel},
+                𝐫ᵢₒₙ::IonDipoleCouplingsType,
                 couplings::AbstractVector,
                 F::ElectricFields.AbstractField, fs::Number)
     t = timeaxis(F, fs)
     volkov = VolkovPhases(F, t)
 
     𝐀 = vector_potential.(F, t)
-    System(ionization_channels, couplings, 𝐀, t, step(t), volkov)
+    System(ionization_channels, 𝐫ᵢₒₙ, couplings, 𝐀, t, step(t), volkov)
 end
 
 @doc raw"""
-    System(ionization_channels, couplings, Fv, Av, t)
+    System(ionization_channels, 𝐫ᵢₒₙ, couplings, Fv, Av, t)
 
 Set up a [`System`](@ref) consisting of multiple
-[`IonizationChannel`](@ref)s with possible `couplings` between them,
-driven by an external field `Fv` with corresponding vector potential
-`Av`, both resolved on the times given by `t`; it is up to the user to
-ensure that ``\vec{F} = -\partial_t\vec{A}`` holds.
+[`IonizationChannel`](@ref)s with with ionic dipole moments `𝐫ᵢₒₙ` and
+possible `couplings` between them, driven by an external field `Fv`
+with corresponding vector potential `Av`, both resolved on the times
+given by `t`; it is up to the user to ensure that ``\vec{F} =
+-\partial_t\vec{A}`` holds.
 """
 function System(ionization_channels::AbstractVector{<:IonizationChannel},
-                couplings::AbstractVector,
+                𝐫ᵢₒₙ::IonDipoleCouplingsType, couplings::AbstractVector,
                 Fv::AbstractVector, Av::AbstractVector, t::AbstractRange)
     volkov = VolkovPhases(Av, t)
 
-    System(ionization_channels, couplings, Av, t, step(t), volkov)
+    System(ionization_channels, 𝐫ᵢₒₙ, couplings, Av, t, step(t), volkov)
 end
 
-System(ionization_channels::AbstractVector{<:IonizationChannel}, args...;
-       couplings=Matrix{AbstractCoupling}[], kwargs...) =
-           System(ionization_channels, couplings, args...)
-
-System(ionization_channel::IonizationChannel, args...; kwargs...) =
-    System([ionization_channel], args...; kwargs...)
+System(ionization_channels::AbstractVector{<:IonizationChannel},
+       F::ElectricFields.AbstractField, fs::Number) =
+           System(ionization_channels, nothing, NoCouplings,
+                  F, fs)
 
 System(Iₚ::Number, args...; kwargs...) =
-    System(IonizationChannel(Iₚ, args...), args...; kwargs...)
+    System([IonizationChannel(Iₚ, args...)], args...; kwargs...)
 
-function Base.show(io::IO, system::System)
-    n = length(system.ionization_channels)
-    write(io, "$(n)-channel System")
-end
+num_channels(system::System) = length(system.ionization_channels)
+
+Base.show(io::IO, system::System) =
+    write(io, "$(num_channels(system))-channel SFA System")
 
 function Base.show(io::IO, mime::MIME"text/plain", system::System)
     println(io, system, ":")
     for (i,c) in enumerate(system.ionization_channels)
         println(io, " ", i, ". ", c)
     end
+    if !isnothing(system.𝐫ᵢₒₙ)
+        nz(A) = count(!iszero, A)
+        nz(A,i) = count(e -> !iszero(e[i]), A)
+        nz(I::UniformScaling) = iszero(I) ? 0 : 1
+        println(io, "Channel dipole couplings:")
+        if system.𝐫ᵢₒₙ isa SVector{3}
+            for (i,d) in enumerate(("x","y","z"))
+                nzd = nz(system.𝐫ᵢₒₙ[i])
+                println(io, "  - $d: $(nzd) non-zero couplings")
+            end
+        else
+            nzz = nz(system.𝐫ᵢₒₙ)
+            println(io, "  - z: $(nzz) non-zero couplings")
+        end
+    end
 end
 
 canonical_momentum_conservation(system::System, which::Integer) =
@@ -185,6 +206,14 @@ struct Diagram{Couplings}
             which == 0 ||
                 throw(ArgumentError("Non-empty diagrams must have photoionization as the first interaction"))
         end
+        ncoup = length(couplings)
+        vcoup = 0:ncoup
+        invalid_couplings = findall(∉(vcoup), last.(path))
+        if !isempty(invalid_couplings)
+            length(invalid_couplings) == 1 &&
+                throw(ArgumentError("Coupling at vertex $(first(invalid_couplings)) not in valid range $(vcoup)"))
+            throw(ArgumentError("Couplings at vertices $(invalid_couplings) not in valid range $(vcoup)"))
+        end
         for i = 2:length(path)-1
             path[i][2] == 0 &&
                 throw(ArgumentError("Only the first and last interaction in a diagram may have which = 0 (corresponding to ionization/recombination)"))
@@ -209,11 +238,19 @@ diagram corresponding to photoionization into the first ionization of
 channel of `system` will automatically be generated.
 """
 function Diagram(path::AbstractVector, system::System)
+    nc = length(system.ionization_channels)
     if isempty(path)
-        length(system.ionization_channels) > 1 &&
+        nc > 1 &&
             @warn "More than one ionization channel present in system, choosing the first"
         path = [(1,0)]
     end
+    vc = 1:nc
+    invalid_channels = findall(∉(vc), first.(path))
+    if !isempty(invalid_channels)
+        length(invalid_channels) == 1 &&
+            throw(ArgumentError("Invalid channel at vertex $(first(invalid_channels))"))
+        throw(ArgumentError("Invalid channels at vertices $(invalid_channels)"))
+    end
     Diagram(path, [typeof(first(c)) for c in system.couplings])
 end
 

From a1d731e5cfc971d69757e4a87fd14d7c5340a123 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Mon, 17 Apr 2023 08:52:22 +0200
Subject: [PATCH 03/20] Minor things

---
 src/dyson_expansions.jl    | 3 ++-
 src/momentum_grid.jl       | 8 ++++++--
 src/telescope_iterators.jl | 2 +-
 src/threading.jl           | 4 ++--
 4 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/dyson_expansions.jl b/src/dyson_expansions.jl
index ec7569c..25b166b 100644
--- a/src/dyson_expansions.jl
+++ b/src/dyson_expansions.jl
@@ -394,7 +394,8 @@ set_momentum!(𝐩s::AbstractVector{<:SVector{3}}, 𝐩ₛₜ::SVector{3}, i) =
 set_momentum!(𝐩s::AbstractVector{<:SVector{3}}, 𝐩ₛₜ::T, i) where {T<:Number} =
     setindex!(𝐩s, SVector{3,T}(zero(T), zero(T), 𝐩ₛₜ), i)
 
-function evaluate_momenta!(𝐩s, prefactors, system, unique_momenta, indeterminate_momenta, i; ϵ=1e-2*√(eps(eltype(system.t))))
+function evaluate_momenta!(𝐩s, prefactors, system, unique_momenta, indeterminate_momenta, i;
+                           ϵ=1e-2*√(eps(eltype(system.t))))
     for idm in indeterminate_momenta
         uidm = unique_momenta[idm]
         a,b = i[uidm[1]],i[uidm[2]]
diff --git a/src/momentum_grid.jl b/src/momentum_grid.jl
index 2796446..305b350 100644
--- a/src/momentum_grid.jl
+++ b/src/momentum_grid.jl
@@ -1,11 +1,15 @@
-function momentum_grid(kmin, kmax, nk, nθ; spacing=:momentum)
-    kmag = if spacing == :momentum
+function momentum_grid(kmin, kmax, nk; spacing=:momentum)
+    if spacing == :momentum
         range(kmin, stop=kmax, length=nk)
     elseif spacing == :energy
         .√(2range(kmin^2/2, stop=kmax^2/2, length=nk))
     else
         throw(ArgumentError("Unknown spacing $(spacing)"))
     end
+end
+
+function momentum_grid(kmin, kmax, nk, nθ, nϕ=1; kwargs...)
+    kmag = momentum_grid(kmin, kmax, nk; kwargs...)
     nθ == 1 && return (kmag,kmag,nothing)
 
     x = range(0, stop=1, length=nθ)
diff --git a/src/telescope_iterators.jl b/src/telescope_iterators.jl
index be1b130..7c10861 100644
--- a/src/telescope_iterators.jl
+++ b/src/telescope_iterators.jl
@@ -4,7 +4,7 @@
 Iterator that is formally equivalent to
 ```julia
 for i = 1:length(v), j = max(1,i-memory):i-1, k = max(1,j-memory):j-1, ...
-    e = v[i,j,k,...]
+    e = v[[i,j,k,...]]
 end
 ```
 for `n` loop variables.
diff --git a/src/threading.jl b/src/threading.jl
index 0ec3ba8..cb6a16f 100644
--- a/src/threading.jl
+++ b/src/threading.jl
@@ -22,8 +22,8 @@ function threaded_range_loop(fun::Function, n::Integer)
     end
 end
 
-threaded_range_loop(fun::Function, v::AbstractVector) =
+threaded_range_loop(fun::Function, v::AbstractArray) =
     threaded_range_loop(i -> fun(v[i]), length(v))
 
-threaded_enumerate_range_loop(fun::Function, v::AbstractVector) =
+threaded_enumerate_range_loop(fun::Function, v::AbstractArray) =
     threaded_range_loop(i -> fun(i,v[i]), length(v))

From b83aca441c9bdb127c278a17aa611dd48dad5bfe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Mon, 17 Apr 2023 08:55:32 +0200
Subject: [PATCH 04/20] Print timings

---
 Project.toml                    |   1 +
 src/StrongFieldApproximation.jl |   1 +
 src/dyson_expansions.jl         | 136 +++++++++++++++++---------------
 3 files changed, 76 insertions(+), 62 deletions(-)

diff --git a/Project.toml b/Project.toml
index 9ef5b49..7977428 100644
--- a/Project.toml
+++ b/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"
 
diff --git a/src/StrongFieldApproximation.jl b/src/StrongFieldApproximation.jl
index 60ef8f0..e606d26 100644
--- a/src/StrongFieldApproximation.jl
+++ b/src/StrongFieldApproximation.jl
@@ -9,6 +9,7 @@ using LinearAlgebra
 using StaticArrays
 
 using ProgressMeter
+using TimerOutputs
 
 include("threading.jl")
 include("telescope_iterators.jl")
diff --git a/src/dyson_expansions.jl b/src/dyson_expansions.jl
index 25b166b..67e8d59 100644
--- a/src/dyson_expansions.jl
+++ b/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
 

From 31c3eda0f0938f4cdfcc8b5751bf5672091b3ef1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Mon, 17 Apr 2023 09:13:51 +0200
Subject: [PATCH 05/20] File separation

---
 src/StrongFieldApproximation.jl |   4 +
 src/diagrams.jl                 | 277 ++++++++++++++++++++++
 src/dyson_expansions.jl         | 408 +-------------------------------
 src/momenta.jl                  |  25 ++
 src/system.jl                   | 109 +++++++++
 5 files changed, 416 insertions(+), 407 deletions(-)
 create mode 100644 src/diagrams.jl
 create mode 100644 src/momenta.jl
 create mode 100644 src/system.jl

diff --git a/src/StrongFieldApproximation.jl b/src/StrongFieldApproximation.jl
index e606d26..4734470 100644
--- a/src/StrongFieldApproximation.jl
+++ b/src/StrongFieldApproximation.jl
@@ -22,6 +22,10 @@ include("field.jl")
 include("volkov.jl")
 
 include("channels.jl")
+include("system.jl")
+include("diagrams.jl")
+include("momenta.jl")
+
 include("dyson_expansions.jl")
 
 include("ionization_rates.jl")
diff --git a/src/diagrams.jl b/src/diagrams.jl
new file mode 100644
index 0000000..edc7e0d
--- /dev/null
+++ b/src/diagrams.jl
@@ -0,0 +1,277 @@
+# * Diagrams
+
+"""
+    Diagram(path, couplings)
+
+Goldstone diagram describing a strong-field process, i.e. after an
+initial photoionization event, the ion and photoelectron are
+propagated separately (at the chosen level of approximation, typically
+Volkov waves are used for the photoelectrons). By interacting with the
+external field, the ion state may change, and through electron
+(re)scattering, both photoelectron and ion state may change.
+
+The diagram is specified using a `path` of pairs (in antichronological
+order), where each pair designates `(resultant ion channel,
+interaction)`, where `interaction` is an index to one of the
+`couplings`. The special value `interaction == 0` corresponds to the
+initial ionization event, and should appear once at the end of `path`
+(i.e. the first event); if it also appears as the first element of
+`path` (i.e. the last event), it corresponds to recombination to the
+initial state.
+
+# Examples
+
+```jldoctest
+julia> couplings = (StrongFieldApproximation.DipoleCoupling,
+                    StrongFieldApproximation.CoulombCoupling)
+(StrongFieldApproximation.DipoleCoupling, StrongFieldApproximation.CoulombCoupling)
+
+julia> Diagram([(1,0)], couplings)
+Goldstone Diagram:
+   |0⟩
+  ╱   ╲⇜
+ 1┃   │𝐤
+
+
+julia> Diagram([(1,0),(1,0)], couplings)
+Goldstone Diagram:
+   |0⟩
+  ╱   ╲⇜
+ 1┃   │𝐩
+  ╲   ╱⇝
+   |0⟩
+
+julia> Diagram([(2,1),(1,0)], couplings)
+Goldstone Diagram:
+   |0⟩
+  ╱   ╲⇜
+ 1┃   │𝐩
+ ⇝┃   │
+ 2┃   │𝐤
+
+
+julia> Diagram([(2,2),(1,0)], couplings)
+Goldstone Diagram:
+   |0⟩
+  ╱   ╲⇜
+ 1┃   │𝐩
+  ┠┈┈┈┤
+ 2┃   │𝐤
+
+
+julia> Diagram([(1,2),(2,2),(1,0)], couplings)
+Goldstone Diagram:
+   |0⟩
+  ╱   ╲⇜
+ 1┃   │𝐩
+  ┠┈┈┈┤
+ 2┃   │𝐪
+  ┠┈┈┈┤
+ 1┃   │𝐤
+
+
+julia> Diagram([(3,1),(1,1),(2,2),(1,0)], couplings)
+Goldstone Diagram:
+   |0⟩
+  ╱   ╲⇜
+ 1┃   │𝐩
+  ┠┈┈┈┤
+ 2┃   │𝐪
+ ⇝┃   │
+ 1┃   │
+ ⇝┃   │
+ 3┃   │𝐤
+```
+"""
+struct Diagram{Couplings}
+    path::Vector{Tuple{Int,Int}}
+    couplings::Couplings
+    function Diagram(path, couplings::Couplings) where Couplings
+        if !isempty(path)
+            α,which = last(path)
+            which == 0 ||
+                throw(ArgumentError("Non-empty diagrams must have photoionization as the first interaction"))
+        end
+        ncoup = length(couplings)
+        vcoup = 0:ncoup
+        invalid_couplings = findall(∉(vcoup), last.(path))
+        if !isempty(invalid_couplings)
+            length(invalid_couplings) == 1 &&
+                throw(ArgumentError("Coupling at vertex $(first(invalid_couplings)) not in valid range $(vcoup)"))
+            throw(ArgumentError("Couplings at vertices $(invalid_couplings) not in valid range $(vcoup)"))
+        end
+        for i = 2:length(path)-1
+            path[i][2] == 0 &&
+                throw(ArgumentError("Only the first and last interaction in a diagram may have which = 0 (corresponding to ionization/recombination)"))
+        end
+        if length(path) > 1 && path[1][2] == 0
+            path[1][1] == path[2][1] ||
+                throw(ArgumentError("Must recombine from the last channel"))
+        end
+        new{Couplings}(path, couplings)
+    end
+end
+
+Diagram(path::Tuple{Int,Int}, couplings) =
+    Diagram([path], couplings)
+
+"""
+    Diagram(path::AbstractVector{Tuple{Int,Int}}, system)
+
+Convenience constructor that sets up the [`Diagram`](@ref) given a
+`path` and the couplings present in `system`. If `path` is empty, the
+diagram corresponding to photoionization into the first ionization of
+channel of `system` will automatically be generated.
+"""
+function Diagram(path::AbstractVector, system::System)
+    nc = length(system.ionization_channels)
+    if isempty(path)
+        nc > 1 &&
+            @warn "More than one ionization channel present in system, choosing the first"
+        path = [(1,0)]
+    end
+    vc = 1:nc
+    invalid_channels = findall(∉(vc), first.(path))
+    if !isempty(invalid_channels)
+        length(invalid_channels) == 1 &&
+            throw(ArgumentError("Invalid channel at vertex $(first(invalid_channels))"))
+        throw(ArgumentError("Invalid channels at vertices $(invalid_channels)"))
+    end
+    Diagram(path, [typeof(first(c)) for c in system.couplings])
+end
+
+Diagram(system::System) = Diagram([], system)
+
+for f in [:length, :first, :firstindex, :lastindex, :isempty]
+    @eval Base.$f(d::Diagram) = $f(d.path)
+end
+Base.getindex(d::Diagram, i) = Diagram(d.path[i], d.couplings)
+
+# ** Pretty-printing
+
+function draw_ionization(io, nd)
+    println(io, lpad("|0⟩", nd+4))
+    print(io, lpad("╱ ╲", nd+4))
+    printstyled(io, "⇜", color=:light_red)
+    println(io)
+end
+
+function draw_recombination(io, nd)
+    print(io, lpad("╲ ╱", nd+4))
+    printstyled(io, "⇝", color=:light_red)
+    println(io)
+    println(io, lpad("|0⟩", nd+4))
+end
+
+draw_exciton(io, ion, nd, electron="") =
+    println(io, lpad(ion, nd)*"┃   │$(electron)")
+
+draw_coulomb_interaction(io, nd) =
+    println(io, repeat(" ",nd)*"┠┈┈┈┤")
+
+function draw_dipole_interaction(io, nd)
+    printstyled(io, lpad("⇝",nd), color=:light_red)
+    println(io, "┃   │")
+end
+
+function Base.show(io::IO, ::MIME"text/plain", d::Diagram)
+    println(io, "Goldstone Diagram:")
+    if isempty(d)
+        println(io, "|0⟩")
+        return
+    end
+    nd = maximum(iw -> length(digits(iw[1])), d.path)+1
+    electrons = ["𝐩","𝐪"]
+    cur_electron = 1
+    for (i,(ion,which)) in Iterators.reverse(enumerate(d.path))
+        if which == 0
+            if i == length(d.path)
+                draw_ionization(io, nd)
+            elseif i == 1
+                draw_recombination(io, nd)
+            end
+        else
+            c = d.couplings[which]
+            if c <: CoulombCoupling
+                draw_coulomb_interaction(io, nd)
+            elseif c <: DipoleCoupling
+                draw_dipole_interaction(io, nd)
+            end
+        end
+        electron = if i == 1
+            "𝐤"
+        elseif !iszero(which) && d.couplings[which] <: DipoleCoupling
+            ""
+        else
+            ei = mod1(cur_electron,length(electrons))
+            e = electrons[ei]*repeat("′", fld1(cur_electron,length(electrons))-1)
+            cur_electron += 1
+            e
+        end
+
+        which == 0 && i == 1 && length(d) > 1 || draw_exciton(io, ion, nd, electron)
+    end
+end
+
+# ** Diagram properties
+
+function analyze_diagram(system, diagram)
+    α,which = first(diagram)
+    ions = Int[]
+    # For photoelectron spectra, time 1 is the reference time,
+    # typically at which the laser pulse has ended; for direct
+    # photoelectron diagrams, time 2 is the time of ionization. For
+    # dipoles, time 1 is the time of recombination.
+    unique_momenta = Tuple{Int,Int}[(1,2)]
+    momenta = [1]
+    indeterminate_momenta = Int[]
+    ld = length(diagram)
+    order = ld
+
+    if ld > 1
+        if which == 0
+            push!(indeterminate_momenta, 1)
+            # Recombination does not increase the order of the diagram, since
+            # it only amounts to projecting the wavefunction on ⟨0|𝐫.
+            order -= 1
+        end
+    else
+        push!(ions, α)
+    end
+
+    i = 2
+    for (α,which) in diagram.path[(iszero(which) ? 2 : 1):end]
+        push!(ions, α)
+        a,b = unique_momenta[end]
+        unique_momenta[end] = (a,i)
+        if !(iszero(which) || canonical_momentum_conservation(system, which) == CanonicalMomentumConservation())
+            push!(unique_momenta, (i,i+1))
+            push!(indeterminate_momenta, length(unique_momenta))
+        end
+        if !iszero(which)
+            push!(momenta, length(unique_momenta))
+        end
+
+        i += 1
+    end
+
+    return ions, unique_momenta, momenta, indeterminate_momenta, order
+end
+
+# 𝐤 is nothing, we want a dipole amplitude
+function diagram_amplitude(::Type{Amp}, system::System, diagram::Diagram, iref, i, ::Nothing) where Amp
+    amplitude = complex(zero(Amp))
+
+    amplitude
+end
+
+# 𝐤 is determinate, we want a photoelectron spectrum
+function diagram_amplitude(::Type{Amp}, system::System, diagram::Diagram, iref, i, 𝐤) where Amp
+    amplitude = complex(zero(Amp))
+
+    amplitude
+end
+
+# * Exports
+
+export Diagram
diff --git a/src/dyson_expansions.jl b/src/dyson_expansions.jl
index 67e8d59..bdb6700 100644
--- a/src/dyson_expansions.jl
+++ b/src/dyson_expansions.jl
@@ -1,411 +1,5 @@
-# * System
-
-"""
-    System
-
-Describes the combined system of atom/molecule (in terms of ionization
-channels) and an external, time-dependent field. The ionization
-channels may be coupled through various interactions, which may or may
-not affect the photoelectron. `System` only describes the channels,
-the external field, and the interactions possible, the actual process
-is described by a [`Diagram`](@ref).
-"""
-struct System{T,IonizationChannels<:AbstractVector{<:IonizationChannel{T}},
-              IonDipoleCouplings,
-              Couplings<:AbstractVector{<:AbstractMatrix{<:AbstractCoupling}},
-              VectorPotential,
-              Times<:AbstractRange,
-              Volkov<:VolkovPhases}
-    ionization_channels::IonizationChannels
-    𝐫ᵢₒₙ::IonDipoleCouplings
-
-    couplings::Couplings
-
-    𝐀::VectorPotential
-    t::Times
-    dt::T
-
-    volkov::Volkov
-end
-
-IonDipoleCouplingsType = Union{AbstractMatrix{<:Number},UniformScaling,SVector{3},Nothing}
-NoCouplings = Matrix{AbstractCoupling}[]
-
-"""
-    System(ionization_channels, 𝐫ᵢₒₙ, couplings, F, ndt)
-
-Set up a [`System`](@ref) consisting of multiple
-[`IonizationChannel`](@ref)s with ionic dipole moments `𝐫ᵢₒₙ` and
-possible `couplings` between them, driven by an external field `F`,
-sampled at a frequency of `fs`.
-"""
-function System(ionization_channels::AbstractVector{<:IonizationChannel},
-                𝐫ᵢₒₙ::IonDipoleCouplingsType,
-                couplings::AbstractVector,
-                F::ElectricFields.AbstractField, fs::Number)
-    t = timeaxis(F, fs)
-    volkov = VolkovPhases(F, t)
-
-    𝐀 = vector_potential.(F, t)
-    System(ionization_channels, 𝐫ᵢₒₙ, couplings, 𝐀, t, step(t), volkov)
-end
-
-@doc raw"""
-    System(ionization_channels, 𝐫ᵢₒₙ, couplings, Fv, Av, t)
-
-Set up a [`System`](@ref) consisting of multiple
-[`IonizationChannel`](@ref)s with with ionic dipole moments `𝐫ᵢₒₙ` and
-possible `couplings` between them, driven by an external field `Fv`
-with corresponding vector potential `Av`, both resolved on the times
-given by `t`; it is up to the user to ensure that ``\vec{F} =
--\partial_t\vec{A}`` holds.
-"""
-function System(ionization_channels::AbstractVector{<:IonizationChannel},
-                𝐫ᵢₒₙ::IonDipoleCouplingsType, couplings::AbstractVector,
-                Fv::AbstractVector, Av::AbstractVector, t::AbstractRange)
-    volkov = VolkovPhases(Av, t)
-
-    System(ionization_channels, 𝐫ᵢₒₙ, couplings, Av, t, step(t), volkov)
-end
-
-System(ionization_channels::AbstractVector{<:IonizationChannel},
-       F::ElectricFields.AbstractField, fs::Number) =
-           System(ionization_channels, nothing, NoCouplings,
-                  F, fs)
-
-System(Iₚ::Number, args...; kwargs...) =
-    System([IonizationChannel(Iₚ, args...)], args...; kwargs...)
-
-num_channels(system::System) = length(system.ionization_channels)
-
-Base.show(io::IO, system::System) =
-    write(io, "$(num_channels(system))-channel SFA System")
-
-function Base.show(io::IO, mime::MIME"text/plain", system::System)
-    println(io, system, ":")
-    for (i,c) in enumerate(system.ionization_channels)
-        println(io, " ", i, ". ", c)
-    end
-    if !isnothing(system.𝐫ᵢₒₙ)
-        nz(A) = count(!iszero, A)
-        nz(A,i) = count(e -> !iszero(e[i]), A)
-        nz(I::UniformScaling) = iszero(I) ? 0 : 1
-        println(io, "Channel dipole couplings:")
-        if system.𝐫ᵢₒₙ isa SVector{3}
-            for (i,d) in enumerate(("x","y","z"))
-                nzd = nz(system.𝐫ᵢₒₙ[i])
-                println(io, "  - $d: $(nzd) non-zero couplings")
-            end
-        else
-            nzz = nz(system.𝐫ᵢₒₙ)
-            println(io, "  - z: $(nzz) non-zero couplings")
-        end
-    end
-end
-
-canonical_momentum_conservation(system::System, which::Integer) =
-    iszero(which) ?
-    CanonicalMomentumConservation() :
-    canonical_momentum_conservation(first(system.couplings[which]))
-
-kinematic_momentum(k::Number, A::Number) = k + A
-kinematic_momentum(k::SVector{3}, A::SVector{3}) = k + A
-kinematic_momentum(k::SVector{3}, A::Number) = SVector{3}(k[1], k[2], k[3]+A)
-
-# * Diagrams
-
-"""
-    Diagram(path, couplings)
-
-Goldstone diagram describing a strong-field process, i.e. after an
-initial photoionization event, the ion and photoelectron are
-propagated separately (at the chosen level of approximation, typically
-Volkov waves are used for the photoelectrons). By interacting with the
-external field, the ion state may change, and through electron
-(re)scattering, both photoelectron and ion state may change.
-
-The diagram is specified using a `path` of pairs (in antichronological
-order), where each pair designates `(resultant ion channel,
-interaction)`, where `interaction` is an index to one of the
-`couplings`. The special value `interaction == 0` corresponds to the
-initial ionization event, and should appear once at the end of `path`
-(i.e. the first event); if it also appears as the first element of
-`path` (i.e. the last event), it corresponds to recombination to the
-initial state.
-
-# Examples
-
-```jldoctest
-julia> couplings = (StrongFieldApproximation.DipoleCoupling,
-                    StrongFieldApproximation.CoulombCoupling)
-(StrongFieldApproximation.DipoleCoupling, StrongFieldApproximation.CoulombCoupling)
-
-julia> Diagram([(1,0)], couplings)
-Goldstone Diagram:
-   |0⟩
-  ╱   ╲⇜
- 1┃   │𝐤
-
-
-julia> Diagram([(1,0),(1,0)], couplings)
-Goldstone Diagram:
-   |0⟩
-  ╱   ╲⇜
- 1┃   │𝐩
-  ╲   ╱⇝
-   |0⟩
-
-julia> Diagram([(2,1),(1,0)], couplings)
-Goldstone Diagram:
-   |0⟩
-  ╱   ╲⇜
- 1┃   │𝐩
- ⇝┃   │
- 2┃   │𝐤
-
-
-julia> Diagram([(2,2),(1,0)], couplings)
-Goldstone Diagram:
-   |0⟩
-  ╱   ╲⇜
- 1┃   │𝐩
-  ┠┈┈┈┤
- 2┃   │𝐤
-
-
-julia> Diagram([(1,2),(2,2),(1,0)], couplings)
-Goldstone Diagram:
-   |0⟩
-  ╱   ╲⇜
- 1┃   │𝐩
-  ┠┈┈┈┤
- 2┃   │𝐪
-  ┠┈┈┈┤
- 1┃   │𝐤
-
-
-julia> Diagram([(3,1),(1,1),(2,2),(1,0)], couplings)
-Goldstone Diagram:
-   |0⟩
-  ╱   ╲⇜
- 1┃   │𝐩
-  ┠┈┈┈┤
- 2┃   │𝐪
- ⇝┃   │
- 1┃   │
- ⇝┃   │
- 3┃   │𝐤
-```
-"""
-struct Diagram{Couplings}
-    path::Vector{Tuple{Int,Int}}
-    couplings::Couplings
-    function Diagram(path, couplings::Couplings) where Couplings
-        if !isempty(path)
-            α,which = last(path)
-            which == 0 ||
-                throw(ArgumentError("Non-empty diagrams must have photoionization as the first interaction"))
-        end
-        ncoup = length(couplings)
-        vcoup = 0:ncoup
-        invalid_couplings = findall(∉(vcoup), last.(path))
-        if !isempty(invalid_couplings)
-            length(invalid_couplings) == 1 &&
-                throw(ArgumentError("Coupling at vertex $(first(invalid_couplings)) not in valid range $(vcoup)"))
-            throw(ArgumentError("Couplings at vertices $(invalid_couplings) not in valid range $(vcoup)"))
-        end
-        for i = 2:length(path)-1
-            path[i][2] == 0 &&
-                throw(ArgumentError("Only the first and last interaction in a diagram may have which = 0 (corresponding to ionization/recombination)"))
-        end
-        if length(path) > 1 && path[1][2] == 0
-            path[1][1] == path[2][1] ||
-                throw(ArgumentError("Must recombine from the last channel"))
-        end
-        new{Couplings}(path, couplings)
-    end
-end
-
-Diagram(path::Tuple{Int,Int}, couplings) =
-    Diagram([path], couplings)
-
-"""
-    Diagram(path::AbstractVector{Tuple{Int,Int}}, system)
-
-Convenience constructor that sets up the [`Diagram`](@ref) given a
-`path` and the couplings present in `system`. If `path` is empty, the
-diagram corresponding to photoionization into the first ionization of
-channel of `system` will automatically be generated.
-"""
-function Diagram(path::AbstractVector, system::System)
-    nc = length(system.ionization_channels)
-    if isempty(path)
-        nc > 1 &&
-            @warn "More than one ionization channel present in system, choosing the first"
-        path = [(1,0)]
-    end
-    vc = 1:nc
-    invalid_channels = findall(∉(vc), first.(path))
-    if !isempty(invalid_channels)
-        length(invalid_channels) == 1 &&
-            throw(ArgumentError("Invalid channel at vertex $(first(invalid_channels))"))
-        throw(ArgumentError("Invalid channels at vertices $(invalid_channels)"))
-    end
-    Diagram(path, [typeof(first(c)) for c in system.couplings])
-end
-
-Diagram(system::System) = Diagram([], system)
-
-for f in [:length, :first, :firstindex, :lastindex, :isempty]
-    @eval Base.$f(d::Diagram) = $f(d.path)
-end
-Base.getindex(d::Diagram, i) = Diagram(d.path[i], d.couplings)
-
-function draw_ionization(io, nd)
-    println(io, lpad("|0⟩", nd+4))
-    print(io, lpad("╱ ╲", nd+4))
-    printstyled(io, "⇜", color=:light_red)
-    println(io)
-end
-
-function draw_recombination(io, nd)
-    print(io, lpad("╲ ╱", nd+4))
-    printstyled(io, "⇝", color=:light_red)
-    println(io)
-    println(io, lpad("|0⟩", nd+4))
-end
-
-draw_exciton(io, ion, nd, electron="") =
-    println(io, lpad(ion, nd)*"┃   │$(electron)")
-
-draw_coulomb_interaction(io, nd) =
-    println(io, repeat(" ",nd)*"┠┈┈┈┤")
-
-function draw_dipole_interaction(io, nd)
-    printstyled(io, lpad("⇝",nd), color=:light_red)
-    println(io, "┃   │")
-end
-
-function Base.show(io::IO, ::MIME"text/plain", d::Diagram)
-    println(io, "Goldstone Diagram:")
-    if isempty(d)
-        println(io, "|0⟩")
-        return
-    end
-    nd = maximum(iw -> length(digits(iw[1])), d.path)+1
-    electrons = ["𝐩","𝐪"]
-    cur_electron = 1
-    for (i,(ion,which)) in Iterators.reverse(enumerate(d.path))
-        if which == 0
-            if i == length(d.path)
-                draw_ionization(io, nd)
-            elseif i == 1
-                draw_recombination(io, nd)
-            end
-        else
-            c = d.couplings[which]
-            if c <: CoulombCoupling
-                draw_coulomb_interaction(io, nd)
-            elseif c <: DipoleCoupling
-                draw_dipole_interaction(io, nd)
-            end
-        end
-        electron = if i == 1
-            "𝐤"
-        elseif !iszero(which) && d.couplings[which] <: DipoleCoupling
-            ""
-        else
-            ei = mod1(cur_electron,length(electrons))
-            e = electrons[ei]*repeat("′", fld1(cur_electron,length(electrons))-1)
-            cur_electron += 1
-            e
-        end
-
-        which == 0 && i == 1 && length(d) > 1 || draw_exciton(io, ion, nd, electron)
-    end
-end
-
 # * Integrate diagrams
 
-function analyze_diagram(system, diagram)
-    α,which = first(diagram)
-    ions = Int[]
-    # For photoelectron spectra, time 1 is the reference time,
-    # typically at which the laser pulse has ended; for direct
-    # photoelectron diagrams, time 2 is the time of ionization. For
-    # dipoles, time 1 is the time of recombination.
-    unique_momenta = Tuple{Int,Int}[(1,2)]
-    momenta = [1]
-    indeterminate_momenta = Int[]
-    ld = length(diagram)
-    order = ld
-
-    if ld > 1
-        if which == 0
-            push!(indeterminate_momenta, 1)
-            # Recombination does not increase the order of the diagram, since
-            # it only amounts to projecting the wavefunction on ⟨0|𝐫.
-            order -= 1
-        end
-    else
-        push!(ions, α)
-    end
-
-    i = 2
-    for (α,which) in diagram.path[(iszero(which) ? 2 : 1):end]
-        push!(ions, α)
-        a,b = unique_momenta[end]
-        unique_momenta[end] = (a,i)
-        if !(iszero(which) || canonical_momentum_conservation(system, which) == CanonicalMomentumConservation())
-            push!(unique_momenta, (i,i+1))
-            push!(indeterminate_momenta, length(unique_momenta))
-        end
-        if !iszero(which)
-            push!(momenta, length(unique_momenta))
-        end
-
-        i += 1
-    end
-
-    return ions, unique_momenta, momenta, indeterminate_momenta, order
-end
-
-# 𝐤 is nothing, we want a dipole amplitude
-function diagram_amplitude(::Type{Amp}, system::System, diagram::Diagram, iref, i, ::Nothing) where Amp
-    amplitude = complex(zero(Amp))
-
-    amplitude
-end
-
-# 𝐤 is determinate, we want a photoelectron spectrum
-function diagram_amplitude(::Type{Amp}, system::System, diagram::Diagram, iref, i, 𝐤) where Amp
-    amplitude = complex(zero(Amp))
-
-    amplitude
-end
-
-momentum_type(_, 𝐤) = typeof(𝐤)
-momentum_type(system, ::Nothing) = eltype(system.volkov.∫A)
-
-set_momentum!(𝐩s::AbstractVector{<:Number}, 𝐩ₛₜ::Number, i) =
-    setindex!(𝐩s, 𝐩ₛₜ, i)
-set_momentum!(𝐩s::AbstractVector{<:SVector{3}}, 𝐩ₛₜ::SVector{3}, i) =
-    setindex!(𝐩s, 𝐩ₛₜ, i)
-set_momentum!(𝐩s::AbstractVector{<:SVector{3}}, 𝐩ₛₜ::T, i) where {T<:Number} =
-    setindex!(𝐩s, SVector{3,T}(zero(T), zero(T), 𝐩ₛₜ), i)
-
-function evaluate_momenta!(𝐩s, prefactors, system, unique_momenta, indeterminate_momenta, i;
-                           ϵ=1e-2*√(eps(eltype(system.t))))
-    for idm in indeterminate_momenta
-        uidm = unique_momenta[idm]
-        a,b = i[uidm[1]],i[uidm[2]]
-        set_momentum!(𝐩s, stationary_momentum(system.volkov, a, b), idm)
-        τ = system.t[a]-system.t[b]
-        ζ = (2π/(im*τ + ϵ))^(3/2)
-        prefactors[idm] = ζ
-    end
-end
-
 function ionization(system::System, diagram::Diagram, 𝐩, 𝐀, i)
     α,which = diagram.path[end]
     @assert which == 0
@@ -564,4 +158,4 @@ end
 
 # * Exports
 
-export Diagram, photoelectron_spectrum, induced_dipole
+export photoelectron_spectrum, induced_dipole
diff --git a/src/momenta.jl b/src/momenta.jl
new file mode 100644
index 0000000..3a86bd5
--- /dev/null
+++ b/src/momenta.jl
@@ -0,0 +1,25 @@
+kinematic_momentum(k::Number, A::Number) = k + A
+kinematic_momentum(k::SVector{3}, A::SVector{3}) = k + A
+kinematic_momentum(k::SVector{3}, A::Number) = SVector{3}(k[1], k[2], k[3]+A)
+
+momentum_type(_, 𝐤) = typeof(𝐤)
+momentum_type(system, ::Nothing) = eltype(system.volkov.∫A)
+
+set_momentum!(𝐩s::AbstractVector{<:Number}, 𝐩ₛₜ::Number, i) =
+    setindex!(𝐩s, 𝐩ₛₜ, i)
+set_momentum!(𝐩s::AbstractVector{<:SVector{3}}, 𝐩ₛₜ::SVector{3}, i) =
+    setindex!(𝐩s, 𝐩ₛₜ, i)
+set_momentum!(𝐩s::AbstractVector{<:SVector{3}}, 𝐩ₛₜ::T, i) where {T<:Number} =
+    setindex!(𝐩s, SVector{3,T}(zero(T), zero(T), 𝐩ₛₜ), i)
+
+function evaluate_momenta!(𝐩s, prefactors, system, unique_momenta, indeterminate_momenta, i;
+                           ϵ=1e-2*√(eps(eltype(system.t))))
+    for idm in indeterminate_momenta
+        uidm = unique_momenta[idm]
+        a,b = i[uidm[1]],i[uidm[2]]
+        set_momentum!(𝐩s, stationary_momentum(system.volkov, a, b), idm)
+        τ = system.t[a]-system.t[b]
+        ζ = (2π/(im*τ + ϵ))^(3/2)
+        prefactors[idm] = ζ
+    end
+end
diff --git a/src/system.jl b/src/system.jl
new file mode 100644
index 0000000..846b779
--- /dev/null
+++ b/src/system.jl
@@ -0,0 +1,109 @@
+# * System
+
+"""
+    System
+
+Describes the combined system of atom/molecule (in terms of ionization
+channels) and an external, time-dependent field. The ionization
+channels may be coupled through various interactions, which may or may
+not affect the photoelectron. `System` only describes the channels,
+the external field, and the interactions possible, the actual process
+is described by a [`Diagram`](@ref).
+"""
+struct System{T,IonizationChannels<:AbstractVector{<:IonizationChannel{T}},
+              IonDipoleCouplings,
+              Couplings<:AbstractVector{<:AbstractMatrix{<:AbstractCoupling}},
+              VectorPotential,
+              Times<:AbstractRange,
+              Volkov<:VolkovPhases}
+    ionization_channels::IonizationChannels
+    𝐫ᵢₒₙ::IonDipoleCouplings
+
+    couplings::Couplings
+
+    𝐀::VectorPotential
+    t::Times
+    dt::T
+
+    volkov::Volkov
+end
+
+IonDipoleCouplingsType = Union{AbstractMatrix{<:Number},UniformScaling,SVector{3},Nothing}
+NoCouplings = Matrix{AbstractCoupling}[]
+
+"""
+    System(ionization_channels, 𝐫ᵢₒₙ, couplings, F, ndt)
+
+Set up a [`System`](@ref) consisting of multiple
+[`IonizationChannel`](@ref)s with ionic dipole moments `𝐫ᵢₒₙ` and
+possible `couplings` between them, driven by an external field `F`,
+sampled at a frequency of `fs`.
+"""
+function System(ionization_channels::AbstractVector{<:IonizationChannel},
+                𝐫ᵢₒₙ::IonDipoleCouplingsType,
+                couplings::AbstractVector,
+                F::ElectricFields.AbstractField, fs::Number)
+    t = timeaxis(F, fs)
+    volkov = VolkovPhases(F, t)
+
+    𝐀 = vector_potential.(F, t)
+    System(ionization_channels, 𝐫ᵢₒₙ, couplings, 𝐀, t, step(t), volkov)
+end
+
+@doc raw"""
+    System(ionization_channels, 𝐫ᵢₒₙ, couplings, Fv, Av, t)
+
+Set up a [`System`](@ref) consisting of multiple
+[`IonizationChannel`](@ref)s with with ionic dipole moments `𝐫ᵢₒₙ` and
+possible `couplings` between them, driven by an external field `Fv`
+with corresponding vector potential `Av`, both resolved on the times
+given by `t`; it is up to the user to ensure that ``\vec{F} =
+-\partial_t\vec{A}`` holds.
+"""
+function System(ionization_channels::AbstractVector{<:IonizationChannel},
+                𝐫ᵢₒₙ::IonDipoleCouplingsType, couplings::AbstractVector,
+                Fv::AbstractVector, Av::AbstractVector, t::AbstractRange)
+    volkov = VolkovPhases(Av, t)
+
+    System(ionization_channels, 𝐫ᵢₒₙ, couplings, Av, t, step(t), volkov)
+end
+
+System(ionization_channels::AbstractVector{<:IonizationChannel},
+       F::ElectricFields.AbstractField, fs::Number) =
+           System(ionization_channels, nothing, NoCouplings,
+                  F, fs)
+
+System(Iₚ::Number, args...; kwargs...) =
+    System([IonizationChannel(Iₚ, args...)], args...; kwargs...)
+
+num_channels(system::System) = length(system.ionization_channels)
+
+Base.show(io::IO, system::System) =
+    write(io, "$(num_channels(system))-channel SFA System")
+
+function Base.show(io::IO, mime::MIME"text/plain", system::System)
+    println(io, system, ":")
+    for (i,c) in enumerate(system.ionization_channels)
+        println(io, " ", i, ". ", c)
+    end
+    if !isnothing(system.𝐫ᵢₒₙ)
+        nz(A) = count(!iszero, A)
+        nz(A,i) = count(e -> !iszero(e[i]), A)
+        nz(I::UniformScaling) = iszero(I) ? 0 : 1
+        println(io, "Channel dipole couplings:")
+        if system.𝐫ᵢₒₙ isa SVector{3}
+            for (i,d) in enumerate(("x","y","z"))
+                nzd = nz(system.𝐫ᵢₒₙ[i])
+                println(io, "  - $d: $(nzd) non-zero couplings")
+            end
+        else
+            nzz = nz(system.𝐫ᵢₒₙ)
+            println(io, "  - z: $(nzz) non-zero couplings")
+        end
+    end
+end
+
+canonical_momentum_conservation(system::System, which::Integer) =
+    iszero(which) ?
+    CanonicalMomentumConservation() :
+    canonical_momentum_conservation(first(system.couplings[which]))

From d23a9343f16555cbdb3327adf137bd6b85fc9c29 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Mon, 17 Apr 2023 10:35:09 +0200
Subject: [PATCH 06/20] Added smooth windows to limit time integrals

---
 Project.toml                    |  2 ++
 src/StrongFieldApproximation.jl |  2 ++
 src/dyson_expansions.jl         | 24 +++++++++++++++-----
 src/system.jl                   |  4 ++--
 src/windows.jl                  | 40 +++++++++++++++++++++++++++++++++
 5 files changed, 64 insertions(+), 8 deletions(-)
 create mode 100644 src/windows.jl

diff --git a/Project.toml b/Project.toml
index 7977428..620c67c 100644
--- a/Project.toml
+++ b/Project.toml
@@ -6,6 +6,7 @@ version = "0.1.0"
 [deps]
 ElectricFields = "2f84ce32-9bb1-11e8-0d9f-3dce90a4beca"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 HCubature = "19dc6840-f33b-545b-b366-655c7e3ffd49"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
@@ -18,6 +19,7 @@ UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 [compat]
 ElectricFields = "0.1, 0.2"
 FastGaussQuadrature = "0.4.7, 0.5"
+FillArrays = "1"
 HCubature = "1.5"
 ProgressMeter = "1.5,1.6"
 SpecialFunctions = "1.3,2"
diff --git a/src/StrongFieldApproximation.jl b/src/StrongFieldApproximation.jl
index 4734470..d37ddc6 100644
--- a/src/StrongFieldApproximation.jl
+++ b/src/StrongFieldApproximation.jl
@@ -7,6 +7,7 @@ using FastGaussQuadrature
 
 using LinearAlgebra
 using StaticArrays
+using FillArrays
 
 using ProgressMeter
 using TimerOutputs
@@ -20,6 +21,7 @@ include("momentum_grid.jl")
 include("units.jl")
 include("field.jl")
 include("volkov.jl")
+include("windows.jl")
 
 include("channels.jl")
 include("system.jl")
diff --git a/src/dyson_expansions.jl b/src/dyson_expansions.jl
index bdb6700..6940791 100644
--- a/src/dyson_expansions.jl
+++ b/src/dyson_expansions.jl
@@ -18,7 +18,8 @@ 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,
+function integrate_diagram(::Type{Amp}, system::System, diagram::Diagram, iref, 𝐤=nothing;
+                           window=flat_window(), imin=1,
                            to=TimerOutput(), verbosity=1) where Amp
     ions, unique_momenta, momenta, indeterminate_momenta, order = @timeit to "Analyze diagram" analyze_diagram(system, diagram)
 
@@ -39,10 +40,18 @@ function integrate_diagram(::Type{Amp}, system::System, diagram::Diagram, iref,
         ctT = complex(eltype(system.t))
 
         is = vcat(iref, zeros(Int, order))
+
+        memory = length(window)
+        if memory < iref
+            window = vcat(window, zeros(eltype(window), iref))
+        end
     end
 
     @timeit to "Time loop" begin
         for i in TelescopeIterator(max(1,imin):iref-1, order, memory)
+            windw = @timeit "Window" prod(window[(i[1] + 1) .- i])
+            iszero(windw) && continue
+
             is[2:end] .= i
             # is = vcat(iref, i)
             # is = (iref,i...)
@@ -89,7 +98,9 @@ function integrate_diagram(::Type{Amp}, system::System, diagram::Diagram, iref,
                 end
             end
 
-            amplitude += ∂a
+            @timeit to "Accumulate amplitude" begin
+                amplitude += ∂a * windw
+            end
         end
     end
 
@@ -122,7 +133,7 @@ function photoelectron_spectrum(k::AbstractArray{T},
 end
 
 function photoelectron_spectrum(k, args...; kwargs...)
-    system = System(args...; kwargs...)
+    system = System(args...)
     photoelectron_spectrum(k, system, Diagram(system); kwargs...)
 end
 
@@ -140,10 +151,11 @@ function induced_dipole(system::System, diagram::Diagram; verbosity = 1, kwargs.
     DT = eltype(system.𝐀)
     𝐝 = zeros(DT, length(t))
 
-    memory = get(kwargs, :memory, typemax(Int))
+    memory = length(get(kwargs, :window, flat_window()))
 
     @showprogress for (i,t) in enumerate(t)
-        𝐝̃ = integrate_diagram(DT, system, diagram, i; imin=i-memory, verbosity=verbosity-1, kwargs...)
+        𝐝̃ = integrate_diagram(DT, system, diagram, i; imin=i-memory,
+                             verbosity=verbosity-1, kwargs...)
         𝐝[i] = 2real(𝐝̃)
     end
 
@@ -151,7 +163,7 @@ function induced_dipole(system::System, diagram::Diagram; verbosity = 1, kwargs.
 end
 
 function induced_dipole(args...; kwargs...)
-    system = System(args...; kwargs...)
+    system = System(args...)
     diagram = Diagram([(1,0),(1,0)], system)
     induced_dipole(system, diagram; kwargs...)
 end
diff --git a/src/system.jl b/src/system.jl
index 846b779..0fad4ce 100644
--- a/src/system.jl
+++ b/src/system.jl
@@ -73,8 +73,8 @@ System(ionization_channels::AbstractVector{<:IonizationChannel},
            System(ionization_channels, nothing, NoCouplings,
                   F, fs)
 
-System(Iₚ::Number, args...; kwargs...) =
-    System([IonizationChannel(Iₚ, args...)], args...; kwargs...)
+System(Iₚ::Number, args...) =
+    System([IonizationChannel(Iₚ, args...)], args...)
 
 num_channels(system::System) = length(system.ionization_channels)
 
diff --git a/src/windows.jl b/src/windows.jl
new file mode 100644
index 0000000..811568d
--- /dev/null
+++ b/src/windows.jl
@@ -0,0 +1,40 @@
+flat_window(n=typemax(Int)) = Trues(n)
+
+#=
+
+Smooth turn-off of scalar functions, following Eqs. (18–21) of
+
+- Becke, A. D. (1988). A Multicenter Numerical Integration Scheme for
+  Polyatomic Molecules. The Journal of Chemical Physics, 88(4),
+  2547–2553. http://dx.doi.org/10.1063/1.454033
+
+=#
+
+function becke_smoother(μ, k::Integer)
+    p = μ -> 3μ/2 - μ^3/2
+    f = μ
+    for i = 1:k
+        f = p(f)
+    end
+    (1-f)/2
+end
+
+function becke_smoother(x, a, b, k)
+    if x < a
+        one(x)
+    elseif x > b
+        zero(x)
+    else
+        becke_smoother(2*(x-a)/(b-a)-1, k)
+    end
+end
+
+function smooth_window(::Type{T}, nflat, nsmooth, k) where T
+    w = ones(T, nflat+nsmooth)
+    for i = eachindex(w)
+        w[i] = becke_smoother(i, nflat, nflat+nsmooth, k)
+    end
+    w
+end
+
+export flat_window, smooth_window

From 188eda21bff380d8e28a481ee8820032d454870f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Mon, 17 Apr 2023 11:05:38 +0200
Subject: [PATCH 07/20] Fix tests

---
 src/channels.jl  | 4 ++++
 test/runtests.jl | 4 +++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/channels.jl b/src/channels.jl
index a57f241..b39ed4d 100644
--- a/src/channels.jl
+++ b/src/channels.jl
@@ -121,3 +121,7 @@ function Base.show(io::IO, mime::MIME"text/plain", g::CoulombCoupling)
     write(io, "CoulombCoupling: ")
     show(io, mime, g.coupling)
 end
+
+# * Exports
+
+export IonizationChannel
diff --git a/test/runtests.jl b/test/runtests.jl
index 830f81b..72ccc0b 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -89,7 +89,9 @@ end
     dc = StrongFieldApproximation.DipoleCoupling(0.1, F)
     couplings=Matrix{StrongFieldApproximation.AbstractCoupling}[reshape([dc],1,1),reshape([cc],1,1)]
 
-    system = StrongFieldApproximation.System(Iₚ, F, ndt, couplings=couplings)
+    ar = (F, ndt)
+    channel = IonizationChannel(Iₚ, ar...)
+    system = StrongFieldApproximation.System(repeat([channel], 10), nothing, couplings, ar...)
 
     for (path, expected_momenta, expected_unique, expected_indeterminate) in [
         ([(1,0)], [1], [(1,2)], []),

From 74e79689e4b217ccc39be16209b2c3582107772d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Mon, 17 Apr 2023 11:59:51 +0200
Subject: [PATCH 08/20] Updated compat bounds for documentation

---
 docs/Project.toml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/Project.toml b/docs/Project.toml
index c4e9453..60ecab7 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -10,8 +10,8 @@ UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 
 [compat]
 CondaPkg = "0.2"
-Documenter = "0.26"
-ElectricFields = "0.1"
+Documenter = "0.27"
+ElectricFields = "0.1, 0.2"
 FFTW = "1.4.0"
 PythonCall = "0.9"
 PythonPlot = "1"

From 287c07446ef5dbc32514fca777c10145e5f24266 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Mon, 17 Apr 2023 12:51:05 +0200
Subject: [PATCH 09/20] Updated documentation

---
 docs/plots.jl     | 13 +++++++++----
 docs/src/index.md | 23 ++++++++++++++++++++---
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/docs/plots.jl b/docs/plots.jl
index 567ee5d..5b17456 100644
--- a/docs/plots.jl
+++ b/docs/plots.jl
@@ -67,10 +67,12 @@ function hhg_example()
 
     # d will be a vector of scalars; by limiting the "memory" of the
     # integrals, we can include only the short trajectory.
-    d = induced_dipole(Iₚ, F, ndt, memory=floor(Int, 0.65ndt));
+    d = induced_dipole(Iₚ, F, ndt, window=flat_window(floor(Int, 0.65ndt)));
+    # d_all includes all trajectories.
+    d_all = induced_dipole(Iₚ, F, ndt);
 
     # d2 will be a vector of 3d vectors
-    d2 = induced_dipole(Iₚ, F2, ndt, memory=floor(Int, 0.65ndt));
+    d2 = induced_dipole(Iₚ, F2, ndt, window=flat_window(floor(Int, 0.65ndt)));
     d2x = [e[1] for e in d2]
     d2y = [e[2] for e in d2]
     d2z = [e[3] for e in d2]
@@ -80,6 +82,7 @@ function hhg_example()
     q = fftshift(fftfreq(length(t), ndt))
     qsel = ind(q,0):ind(q,100)
     D = spectrum(d)
+    D_all = spectrum(d_all)
     D2 = spectrum(d2)
     cutoff = 3.17austrip(ponderomotive_potential(F)) + Iₚ
 
@@ -98,13 +101,15 @@ function hhg_example()
             plot(tplot, d2x, "--")
             plot(tplot, d2y, "--")
             plot(tplot, d2z, "--")
-            legend(["1d", "3d x", "3d y", "3d z"])
+            plot(tplot, d_all, ":")
+            legend(["1d", "3d x", "3d y", "3d z", "1d all traj."], loc=1)
             ylabel(L"\langle\mathbf{r}\rangle(t)")
         end
         csubplot(313) do
             semilogy(q[qsel], abs2.(D[qsel]))
             semilogy(q[qsel], abs2.(D2[qsel,:]), "--")
-            legend(["1d", "3d x", "3d y", "3d z"])
+            semilogy(q[qsel], abs2.(D_all[qsel]), ":")
+            legend(["1d", "3d x", "3d y", "3d z", "1d all traj."])
             axvline(cutoff/photon_energy(F), linestyle=":", color="black")
             xlabel(L"Harmonic order of 800 nm [$q$]")
             ylabel(L"|\mathbf{r}(q)|^2")
diff --git a/docs/src/index.md b/docs/src/index.md
index 67a9589..e277a80 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -12,7 +12,7 @@ throughout, unless otherwise specified.
 
 At the moment, two kind of observables are supported: induced dipole
 moment and photoelectron spectra. Time integrals are performed
-numerically, with recursive time integrals limited by a `memory`,
+numerically, with recursive time integrals limited by a `window`,
 i.e. how many time steps are considered (default is from the beginning
 of the pulse). Integrals over intermediate photoelectron momenta are
 performed using the saddle-point method, i.e. given two times, the
@@ -75,7 +75,7 @@ julia> Iₚ = 0.5 # Hydrogen
 
 julia> # d will be a vector of scalars; by limiting the "memory" of the
        # integrals, we can include only the short trajectory.
-       d = induced_dipole(Iₚ, F, ndt, memory=floor(Int, 0.65ndt));
+       d = induced_dipole(Iₚ, F, ndt, window=flat_window(floor(Int, 0.65ndt)));
 ┌ Info: Induced dipole calculation
 │   system =
 │    1-channel System:
@@ -91,8 +91,25 @@ julia> # d will be a vector of scalars; by limiting the "memory" of the
 └
 Progress: 100%|████████████████████████████████████████████████████████████████| Time: 0:00:00
 
+julia> # d_all includes all trajectories.
+       d_all = induced_dipole(Iₚ, F, ndt);
+┌ Info: Induced dipole calculation
+│   system =
+│    1-channel SFA System:
+│     1. IonizationChannel: Iₚ = 0.5 Ha = 13.6055 eV
+│
+│   diagram =
+│    Goldstone Diagram:
+│       |0⟩
+│       ╱ ╲⇜
+│     1┃   │𝐩
+│       ╲ ╱⇝
+│       |0⟩
+└
+Progress: 100%|████████████████████████████████████████████████████████████████| Time: 0:00:08
+
 julia> # d2 will be a vector of 3d vectors
-       d2 = induced_dipole(Iₚ, F2, ndt, memory=floor(Int, 0.65ndt));
+       d2 = induced_dipole(Iₚ, F2, ndt, window=flat_window(floor(Int, 0.65ndt)));
 ┌ Info: Induced dipole calculation
 │   system =
 │    1-channel System:

From 3e5a046e20a39ce82ff7204789cbb534ebb86d5f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Mon, 17 Apr 2023 16:20:09 +0200
Subject: [PATCH 10/20] Implemented pre-propagation of laser-dressed ion
 channels

---
 Project.toml                    |   2 +
 docs/make.jl                    |   1 +
 src/StrongFieldApproximation.jl |   6 ++
 src/dyson_expansions.jl         |  34 ++++---
 src/find_blocks.jl              |  33 +++++++
 src/ion_propagators.jl          | 155 ++++++++++++++++++++++++++++++++
 src/system.jl                   |  25 ++++--
 7 files changed, 234 insertions(+), 22 deletions(-)
 create mode 100644 src/find_blocks.jl
 create mode 100644 src/ion_propagators.jl

diff --git a/Project.toml b/Project.toml
index 620c67c..db49f6d 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,12 +4,14 @@ authors = ["Stefanos Carlström <stefanos.carlstrom@gmail.com>"]
 version = "0.1.0"
 
 [deps]
+DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 ElectricFields = "2f84ce32-9bb1-11e8-0d9f-3dce90a4beca"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 HCubature = "19dc6840-f33b-545b-b366-655c7e3ffd49"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
diff --git a/docs/make.jl b/docs/make.jl
index 984d4fb..6ec87ad 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -32,6 +32,7 @@ makedocs(;
                     :Hamiltonian => "\\operator{H}",
                     :hamiltonian => "\\operator{h}",
                     :Lagrangian => "\\operator{L}",
+                    :propU => "\\mathcal{U}",
                     :fock => "\\operator{f}",
                     :lagrange => ["\\epsilon_{#1}", 1],
                     :vary => ["\\delta_{#1}", 1],
diff --git a/src/StrongFieldApproximation.jl b/src/StrongFieldApproximation.jl
index d37ddc6..068199b 100644
--- a/src/StrongFieldApproximation.jl
+++ b/src/StrongFieldApproximation.jl
@@ -6,20 +6,26 @@ using UnitfulAtomic
 using FastGaussQuadrature
 
 using LinearAlgebra
+using SparseArrays
 using StaticArrays
 using FillArrays
 
+import DataStructures: Queue, enqueue!, dequeue!
+
 using ProgressMeter
 using TimerOutputs
 
 include("threading.jl")
 include("telescope_iterators.jl")
 
+include("find_blocks.jl")
+
 include("atom_models.jl")
 include("momentum_grid.jl")
 
 include("units.jl")
 include("field.jl")
+include("ion_propagators.jl")
 include("volkov.jl")
 include("windows.jl")
 
diff --git a/src/dyson_expansions.jl b/src/dyson_expansions.jl
index 6940791..6d0ecbf 100644
--- a/src/dyson_expansions.jl
+++ b/src/dyson_expansions.jl
@@ -1,18 +1,26 @@
 # * Integrate diagrams
 
-function ionization(system::System, diagram::Diagram, 𝐩, 𝐀, i)
+function ionization(system::System{T}, diagram::Diagram, 𝐩, 𝐀, i) where T
     α,which = diagram.path[end]
     @assert which == 0
-    source_term(system.ionization_channels[α],
-                i,
-                kinematic_momentum(𝐩, 𝐀[i]))
+    s = zero(complex(T))
+    p = kinematic_momentum(𝐩, 𝐀[i])
+    for (j,q) in non_zero_ion_mapping(system.ions, α, i)
+        s += q*source_term(system.ionization_channels[j], i, p)
+    end
+    s
 end
 
 function recombination(system::System, diagram::Diagram, 𝐩, 𝐀, i)
     α,which = first(diagram)
     if which == 0 && length(diagram) > 1
-        d = system.ionization_channels[α].st.d
-        conj(d(kinematic_momentum(𝐩, 𝐀[i])))
+        s = zero(complex(T))
+        p = kinematic_momentum(𝐩, 𝐀[i])
+        for (j,q) in non_zero_ion_mapping(system.ions, α, i)
+            d = system.ionization_channels[j].st.d
+            s += q*d(p)
+        end
+        conj(s)
     else
         true
     end
@@ -28,7 +36,6 @@ function integrate_diagram(::Type{Amp}, system::System, diagram::Diagram, iref,
 
         verbosity > 1 && @info "Integrating diagram up to" iref system diagram ions unique_momenta momenta indeterminate_momenta order weight 𝐤
 
-        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(𝐤)
@@ -63,16 +70,14 @@ function integrate_diagram(::Type{Amp}, system::System, diagram::Diagram, iref,
             verbosity > 10 && @show 𝐩s prefactors
 
             @timeit to "Evaluate propagators" begin
-                Sᵢₒₙ = zero(ctT)
+                aₚᵣₒₚ_ᵢₒₙ = one(ctT)
                 Sₑₗ = zero(ctT)
                 for j = 1:order
                     a,b = is[j], is[j+1]
-                    τ = system.t[a] - system.t[b]
-                    Sᵢₒₙ += Eᵢₒₙₛ[j]*τ
+                    aₚᵣₒₚ_ᵢₒₙ *= ion_propagation(system.ions, ions[j], a, b)
                     Sₑₗ += volkov_phase(𝐩s[j], system.volkov, a, b)
                 end
-                S = Sᵢₒₙ + Sₑₗ
-                aₚᵣₒₚ = prod(prefactors)*exp(-im*S)
+                aₚᵣₒₚ = prod(prefactors)*exp(-im*Sₑₗ)*aₚᵣₒₚ_ᵢₒₙ
             end
 
             verbosity > 10 && @show is prefactors
@@ -88,13 +93,14 @@ function integrate_diagram(::Type{Amp}, system::System, diagram::Diagram, iref,
                     α = ions[j-1]
                     β = ions[j]
                     verbosity > 20 && @show j, ion, which α,β
-                    interaction = system.couplings[which][α,β]
 
                     𝐤ᵢ = 𝐩s[momenta[j-1]]
                     𝐩ᵢ = 𝐩s[momenta[j]]
                     𝐀ᵢ = 𝐀[is[j]]
 
-                    ∂a *= @timeit to "Interaction" interaction(kinematic_momentum(𝐤ᵢ, 𝐀ᵢ), kinematic_momentum(𝐩ᵢ, 𝐀ᵢ), is[j+1])
+                    ∂a *= @timeit to "Interaction" interaction(system.ions, system.couplings[which],
+                                                               α, kinematic_momentum(𝐤ᵢ, 𝐀ᵢ),
+                                                               β, kinematic_momentum(𝐩ᵢ, 𝐀ᵢ), is[j+1])
                 end
             end
 
diff --git a/src/find_blocks.jl b/src/find_blocks.jl
new file mode 100644
index 0000000..a2a4b5a
--- /dev/null
+++ b/src/find_blocks.jl
@@ -0,0 +1,33 @@
+# https://en.wikipedia.org/wiki/Breadth-first_search
+function find_blocks(G::SparseMatrixCSC)
+    m = size(G,1)
+    visited = falses(m)
+
+    rows = rowvals(G)
+    vals = nonzeros(G)
+
+    blocks = Vector{Vector{Int}}()
+    while !all(visited)
+        s = findfirst(!, visited)
+        visited[s] = true
+        nzs = nzrange(G, s)
+        (isempty(nzs) || iszero(vals[first(nzs)])) && continue
+        block = [s]
+
+        q = Queue{Int}()
+        enqueue!(q, s)
+        while !isempty(q)
+            v = dequeue!(q)
+            for k in nzrange(G, v)
+                w = rows[k]
+                (visited[w] || iszero(vals[k])) && continue
+                enqueue!(q, w)
+                visited[w] = true
+                push!(block, w)
+            end
+        end
+        push!(blocks, sort(block))
+    end
+
+    blocks
+end
diff --git a/src/ion_propagators.jl b/src/ion_propagators.jl
new file mode 100644
index 0000000..0593548
--- /dev/null
+++ b/src/ion_propagators.jl
@@ -0,0 +1,155 @@
+abstract type IonPropagator end
+
+function interaction(ions::IonPropagator, coupling, α, k, β, p, i)
+    s = zero(eltype(ions))
+    for (j′,q′) in non_zero_ion_mapping(ions, α, i)
+        for (j,q) in non_zero_ion_mapping(ions, β, i)
+            s += conj(q′)*q*coupling[j′,j](k, p, i)
+        end
+    end
+    s
+end
+
+# * Laser-dressed ions
+
+@doc raw"""
+    LaserDressedIons
+
+Pre-propagated laser-dressed basis for the ion states, i.e. the
+time-dependent eigenbasis of the ionic Hamiltonian including the
+external electric field. In the eigenbasis, the ionic propagator is
+diagonal, and the entries are
+
+```math
+\matrixel{\eta}{\propU(a,b)}{\gamma} =
+\exp[-\im E_\eta(T-a)]
+\exp[+\im E_\eta(T-b)]
+\delta_{\eta\gamma},
+```
+
+i.e. the reference time is ``T`` (the end of the pulse), which is
+consistent with [`VolkovPhases`](@ref).
+"""
+struct LaserDressedIons{T,IonBasis} <: IonPropagator
+    expimE::Matrix{T}
+    ion_basis::IonBasis
+end
+
+function LaserDressedIons(E::Matrix{T}, Q, t) where T
+    expimE = similar(E, complex(T))
+
+    nt = length(t)
+    for i = 1:nt
+        μ = im*(t[end] - t[i])
+        expimE[i,:] = exp.(μ.*E[i,:])
+    end
+
+    LaserDressedIons(expimE, Q)
+end
+
+function LaserDressedIons(ics, 𝐫ᵢₒₙ::SparseMatrixCSC, Fv::AbstractArray, t)
+    to = TimerOutput()
+
+    @timeit to "Allocations" begin
+        m,n = size(𝐫ᵢₒₙ)
+        @assert m == n == length(ics)
+
+        nt = length(t)
+        E₀ = [ic.E for ic in ics]
+        H₀ = Diagonal(E₀)
+
+        X = .!iszero.(I + (𝐫ᵢₒₙ .≠ 0))
+        bs = find_blocks(X)
+
+        @info "Diagonalizing $(m)×$(n) ionic Hamiltonian for $(nt) times" bs
+
+        T = eltype(E₀)
+        E = zeros(T, nt, m)
+        Q = zeros(T, m, m, nt)
+    end
+
+    @timeit to "Block loop" for b in bs
+        nb = length(b)
+        if nb == 1
+            j = b[1]
+            E[:,j] .= E₀[j]
+            Q[j,j,:] .= 1
+            continue
+        end
+
+        Hsub = zeros(T, nb, nb)
+        H₀sub = Diagonal(E₀[b])
+        zsub = Matrix(𝐫ᵢₒₙ[b,b])
+        @timeit to "Time loop" begin
+            @showprogress for i = 1:nt
+                Hsub .= H₀sub .+ Fv[i] .* zsub
+                ee = eigen!(Hsub)
+                for (ij,j) in enumerate(b)
+                    E[i,j] = ee.values[ij]
+                end
+                Q[b,b,i] = ee.vectors
+            end
+        end
+    end
+
+    TimerOutputs.complement!(to)
+    print_timer(to)
+
+    LaserDressedIons(E, Q, t)
+end
+
+function LaserDressedIons(ics, ::Nothing, _::AbstractArray, t)
+    nt = length(t)
+    E₀ = [ic.E for ic in ics]
+    nc = length(E₀)
+
+    T = eltype(E₀)
+    E = zeros(T, nt, nc)
+    for j = 1:nc
+        E[:,j] .= E₀[j]
+    end
+
+    LaserDressedIons(E, nothing, t)
+end
+
+LaserDressedIons(ics, 𝐫ᵢₒₙ, F::ElectricFields.AbstractField, t) =
+    LaserDressedIons(ics, 𝐫ᵢₒₙ, field_amplitude(F, t), t)
+
+Base.eltype(::LaserDressedIons{T}) where T = T
+
+ion_mapping(ions::LaserDressedIons, α, i) =
+    view(ions.ion_basis, α, :, i)
+
+ion_mapping(ions::LaserDressedIons{T,Nothing}, α, _) where T =
+    vcat(zeros(T,α-1), one(T), zeros(T, size(ions.expimE,2)-α))
+
+non_zero_ion_mapping(ions::LaserDressedIons, α, i) =
+    filter(jq -> !iszero(last(jq)),
+           collect(enumerate(ion_mapping(ions, α, i))))
+
+non_zero_ion_mapping(ions::LaserDressedIons{T,Nothing}, α, _) where T =
+    [(α, one(T))]
+
+ion_propagation(ions::LaserDressedIons, j, a, b) =
+    ions.expimE[a,j]*conj(ions.expimE[b,j])
+
+# * Field-free ions
+
+struct FieldFreeIons{T,Times} <: IonPropagator
+    E₀::Vector{T}
+    t::Times
+end
+
+FieldFreeIons(ics, _, _, t) =
+    FieldFreeIons([ic.E for ic in ics], t)
+
+Base.eltype(::FieldFreeIons{T}) where T = T
+
+ion_mapping(::FieldFreeIons{T}, α, _) where T =
+    vcat(zeros(T,α-1), one(T), zeros(T, length(ions.E₀)-α))
+
+non_zero_ion_mapping(::FieldFreeIons{T}, α, _) where T =
+    [(α, one(T))]
+
+ion_propagation(ions::FieldFreeIons, j, a, b) =
+    exp(-im*(ions.t[a]-ions.t[b])*ions.E₀[j])
diff --git a/src/system.jl b/src/system.jl
index 0fad4ce..7d61b97 100644
--- a/src/system.jl
+++ b/src/system.jl
@@ -15,6 +15,7 @@ struct System{T,IonizationChannels<:AbstractVector{<:IonizationChannel{T}},
               Couplings<:AbstractVector{<:AbstractMatrix{<:AbstractCoupling}},
               VectorPotential,
               Times<:AbstractRange,
+              Ions<:IonPropagator,
               Volkov<:VolkovPhases}
     ionization_channels::IonizationChannels
     𝐫ᵢₒₙ::IonDipoleCouplings
@@ -25,12 +26,19 @@ struct System{T,IonizationChannels<:AbstractVector{<:IonizationChannel{T}},
     t::Times
     dt::T
 
+    ions::Ions
     volkov::Volkov
 end
 
 IonDipoleCouplingsType = Union{AbstractMatrix{<:Number},UniformScaling,SVector{3},Nothing}
 NoCouplings = Matrix{AbstractCoupling}[]
 
+function System(ionization_channels, 𝐫ᵢₒₙ, couplings, 𝐅, 𝐀, t, volkov::VolkovPhases;
+                Ions::Type{<:IonPropagator}=LaserDressedIons, kwargs...)
+    ions = Ions(ionization_channels, 𝐫ᵢₒₙ, 𝐅, t)
+    System(ionization_channels, 𝐫ᵢₒₙ, couplings, 𝐀, t, step(t), ions, volkov)
+end
+
 """
     System(ionization_channels, 𝐫ᵢₒₙ, couplings, F, ndt)
 
@@ -42,12 +50,12 @@ sampled at a frequency of `fs`.
 function System(ionization_channels::AbstractVector{<:IonizationChannel},
                 𝐫ᵢₒₙ::IonDipoleCouplingsType,
                 couplings::AbstractVector,
-                F::ElectricFields.AbstractField, fs::Number)
+                F::ElectricFields.AbstractField, fs::Number; kwargs...)
     t = timeaxis(F, fs)
     volkov = VolkovPhases(F, t)
 
     𝐀 = vector_potential.(F, t)
-    System(ionization_channels, 𝐫ᵢₒₙ, couplings, 𝐀, t, step(t), volkov)
+    System(ionization_channels, 𝐫ᵢₒₙ, couplings, F, 𝐀, t, volkov; kwargs...)
 end
 
 @doc raw"""
@@ -62,19 +70,20 @@ given by `t`; it is up to the user to ensure that ``\vec{F} =
 """
 function System(ionization_channels::AbstractVector{<:IonizationChannel},
                 𝐫ᵢₒₙ::IonDipoleCouplingsType, couplings::AbstractVector,
-                Fv::AbstractVector, Av::AbstractVector, t::AbstractRange)
+                Fv::AbstractVector, Av::AbstractVector, t::AbstractRange;
+                kwargs...)
     volkov = VolkovPhases(Av, t)
 
-    System(ionization_channels, 𝐫ᵢₒₙ, couplings, Av, t, step(t), volkov)
+    System(ionization_channels, 𝐫ᵢₒₙ, couplings, Fv, Av, t, volkov; kwargs...)
 end
 
 System(ionization_channels::AbstractVector{<:IonizationChannel},
-       F::ElectricFields.AbstractField, fs::Number) =
+       F::ElectricFields.AbstractField, fs::Number; kwargs...) =
            System(ionization_channels, nothing, NoCouplings,
-                  F, fs)
+                  F, fs; kwargs...)
 
-System(Iₚ::Number, args...) =
-    System([IonizationChannel(Iₚ, args...)], args...)
+System(Iₚ::Number, args...; kwargs...) =
+    System([IonizationChannel(Iₚ, args...)], args...; kwargs...)
 
 num_channels(system::System) = length(system.ionization_channels)
 

From b2d1ee9dc42d0ff76b2521a4fabf31f2e88757e0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Mon, 17 Apr 2023 18:57:05 +0200
Subject: [PATCH 11/20] Missing type parameter

---
 src/dyson_expansions.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/dyson_expansions.jl b/src/dyson_expansions.jl
index 6d0ecbf..a9e8746 100644
--- a/src/dyson_expansions.jl
+++ b/src/dyson_expansions.jl
@@ -11,7 +11,7 @@ function ionization(system::System{T}, diagram::Diagram, 𝐩, 𝐀, i) where T
     s
 end
 
-function recombination(system::System, diagram::Diagram, 𝐩, 𝐀, i)
+function recombination(system::System{T}, diagram::Diagram, 𝐩, 𝐀, i) where T
     α,which = first(diagram)
     if which == 0 && length(diagram) > 1
         s = zero(complex(T))

From 86057f5e996fc14e8a26d9493e83a1a18e9dd7e6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Mon, 17 Apr 2023 20:09:49 +0200
Subject: [PATCH 12/20] Fixed recombination amplitude

---
 src/dyson_expansions.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/dyson_expansions.jl b/src/dyson_expansions.jl
index a9e8746..df6878f 100644
--- a/src/dyson_expansions.jl
+++ b/src/dyson_expansions.jl
@@ -14,8 +14,8 @@ end
 function recombination(system::System{T}, diagram::Diagram, 𝐩, 𝐀, i) where T
     α,which = first(diagram)
     if which == 0 && length(diagram) > 1
-        s = zero(complex(T))
         p = kinematic_momentum(𝐩, 𝐀[i])
+        s = complex(zero(first(system.ionization_channels).st.d(p)))
         for (j,q) in non_zero_ion_mapping(system.ions, α, i)
             d = system.ionization_channels[j].st.d
             s += q*d(p)

From 90b4ae4e88466db2c9a43e3df90bf3888aea53a4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Tue, 18 Apr 2023 09:33:43 +0200
Subject: [PATCH 13/20] Correct interaction time

---
 src/dyson_expansions.jl | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/dyson_expansions.jl b/src/dyson_expansions.jl
index df6878f..09b1a05 100644
--- a/src/dyson_expansions.jl
+++ b/src/dyson_expansions.jl
@@ -94,13 +94,12 @@ function integrate_diagram(::Type{Amp}, system::System, diagram::Diagram, iref,
                     β = ions[j]
                     verbosity > 20 && @show j, ion, which α,β
 
-                    𝐤ᵢ = 𝐩s[momenta[j-1]]
-                    𝐩ᵢ = 𝐩s[momenta[j]]
                     𝐀ᵢ = 𝐀[is[j]]
+                    𝐤ᵢ = kinematic_momentum(𝐩s[momenta[j-1]], 𝐀ᵢ)
+                    𝐩ᵢ = kinematic_momentum(𝐩s[momenta[j]], 𝐀ᵢ)
 
                     ∂a *= @timeit to "Interaction" interaction(system.ions, system.couplings[which],
-                                                               α, kinematic_momentum(𝐤ᵢ, 𝐀ᵢ),
-                                                               β, kinematic_momentum(𝐩ᵢ, 𝐀ᵢ), is[j+1])
+                                                               α, 𝐤ᵢ, β, 𝐩ᵢ, is[j])
                 end
             end
 

From 61c227937939cbc00fe5352cc19bb6e09431f914 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Tue, 18 Apr 2023 09:34:37 +0200
Subject: [PATCH 14/20] Verbosity changes

---
 src/dyson_expansions.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/dyson_expansions.jl b/src/dyson_expansions.jl
index 09b1a05..0bf2256 100644
--- a/src/dyson_expansions.jl
+++ b/src/dyson_expansions.jl
@@ -118,7 +118,7 @@ function photoelectron_spectrum(k::AbstractArray{T},
                                 system::System, diagram::Diagram;
                                 iref=length(system.t),
                                 verbosity=1, kwargs...) where T
-    verbosity > 0 && @info "Photoelectrum spectrum calculation" system diagram length(k)
+    verbosity == 1 && @info "Photoelectrum spectrum calculation" system diagram length(k)
 
     cT = complex(eltype(T))
     c = similar(k, cT)
@@ -133,7 +133,7 @@ function photoelectron_spectrum(k::AbstractArray{T},
         end
     end
     TimerOutputs.complement!(to)
-    verbosity > 0 && print_timer(to)
+    0 < verbosity < 4 && print_timer(to)
     c
 end
 

From 0cbc0ea5e0a10072fbc19f22a355bb2e66af4fdd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Tue, 18 Apr 2023 09:42:00 +0200
Subject: [PATCH 15/20] Added multi-channel implementation

---
 src/StrongFieldApproximation.jl |   1 +
 src/ion_propagators.jl          |  32 ++++++-
 src/multi_channel_sfa.jl        | 152 ++++++++++++++++++++++++++++++++
 3 files changed, 182 insertions(+), 3 deletions(-)
 create mode 100644 src/multi_channel_sfa.jl

diff --git a/src/StrongFieldApproximation.jl b/src/StrongFieldApproximation.jl
index 068199b..eb01e5f 100644
--- a/src/StrongFieldApproximation.jl
+++ b/src/StrongFieldApproximation.jl
@@ -35,6 +35,7 @@ include("diagrams.jl")
 include("momenta.jl")
 
 include("dyson_expansions.jl")
+include("multi_channel_sfa.jl")
 
 include("ionization_rates.jl")
 
diff --git a/src/ion_propagators.jl b/src/ion_propagators.jl
index 0593548..b22ce12 100644
--- a/src/ion_propagators.jl
+++ b/src/ion_propagators.jl
@@ -10,6 +10,24 @@ function interaction(ions::IonPropagator, coupling, α, k, β, p, i)
     s
 end
 
+function apply_interaction!(v, coupling, u, k, p, ti)
+    nc = length(u)
+    v .= false
+    for α = 1:nc
+        for β = 1:nc
+            v[α] += coupling[α,β](k, p, ti)*u[β]
+        end
+    end
+    v
+end
+
+function apply_interaction!(v, ions::IonPropagator, coupling, u, k, p, i)
+    Q = ion_mapping(ions, i)
+    mul!(v, Q, u)
+    apply_interaction!(u, coupling, v, k, p, i)
+    mul!(v, Q, u)
+end
+
 # * Laser-dressed ions
 
 @doc raw"""
@@ -83,7 +101,8 @@ function LaserDressedIons(ics, 𝐫ᵢₒₙ::SparseMatrixCSC, Fv::AbstractArray
         @timeit to "Time loop" begin
             @showprogress for i = 1:nt
                 Hsub .= H₀sub .+ Fv[i] .* zsub
-                ee = eigen!(Hsub)
+                @assert Hsub ≈ Hsub'
+                ee = eigen!(Hermitian(Hsub))
                 for (ij,j) in enumerate(b)
                     E[i,j] = ee.values[ij]
                 end
@@ -117,6 +136,11 @@ LaserDressedIons(ics, 𝐫ᵢₒₙ, F::ElectricFields.AbstractField, t) =
 
 Base.eltype(::LaserDressedIons{T}) where T = T
 
+ion_mapping(ions::LaserDressedIons, i) =
+    view(ions.ion_basis, :, :, i)
+
+ion_mapping(ions::LaserDressedIons, _) = I
+
 ion_mapping(ions::LaserDressedIons, α, i) =
     view(ions.ion_basis, α, :, i)
 
@@ -131,7 +155,7 @@ non_zero_ion_mapping(ions::LaserDressedIons{T,Nothing}, α, _) where T =
     [(α, one(T))]
 
 ion_propagation(ions::LaserDressedIons, j, a, b) =
-    ions.expimE[a,j]*conj(ions.expimE[b,j])
+    ions.expimE[a,j] .* conj(ions.expimE[b,j])
 
 # * Field-free ions
 
@@ -145,6 +169,8 @@ FieldFreeIons(ics, _, _, t) =
 
 Base.eltype(::FieldFreeIons{T}) where T = T
 
+ion_mapping(::FieldFreeIons, _) = I
+
 ion_mapping(::FieldFreeIons{T}, α, _) where T =
     vcat(zeros(T,α-1), one(T), zeros(T, length(ions.E₀)-α))
 
@@ -152,4 +178,4 @@ non_zero_ion_mapping(::FieldFreeIons{T}, α, _) where T =
     [(α, one(T))]
 
 ion_propagation(ions::FieldFreeIons, j, a, b) =
-    exp(-im*(ions.t[a]-ions.t[b])*ions.E₀[j])
+    exp.(-im*(ions.t[a]-ions.t[b])*ions.E₀[j])
diff --git a/src/multi_channel_sfa.jl b/src/multi_channel_sfa.jl
new file mode 100644
index 0000000..5e5dea0
--- /dev/null
+++ b/src/multi_channel_sfa.jl
@@ -0,0 +1,152 @@
+function ionization!(v, system, 𝐩, 𝐀, i, tmp)
+    p = kinematic_momentum(𝐩, 𝐀[i])
+    for j in eachindex(v)
+        tmp[j] = source_term(system.ionization_channels[j], i, p)
+    end
+    mul!(v, ion_mapping(system.ions, i), tmp)
+end
+
+macro swap!(x,y,tmp)
+    quote
+        $(esc(tmp)) = $(esc(x))
+        $(esc(x)) = $(esc(y))
+        $(esc(y)) = $(esc(tmp))
+    end
+end
+
+function multi_channel_sfa!(amplitudes, system, diagram, iref, 𝐤=nothing;
+                            window=flat_window(), imin=1, verbosity=0,
+                            to=TimerOutput())
+    _, unique_momenta, momenta, indeterminate_momenta, order = @timeit to "Analyze diagram" analyze_diagram(system, diagram)
+    @timeit to "Allocations" begin
+        weight = (-im*system.dt)^order
+        if verbosity > 0
+            println()
+            @info "Integrating diagram up to" iref system diagram unique_momenta momenta indeterminate_momenta order weight 𝐤
+        end
+
+        𝐩s = complex(zeros(momentum_type(system, 𝐤), length(unique_momenta)))
+        prefactors = ones(complex(eltype(system.t)), length(unique_momenta))
+        if !isnothing(𝐤)
+            𝐩s[1] = 𝐤
+        end
+        verbosity > 2 && @show 𝐩s
+
+        𝐀 = system.𝐀
+        amplitudes .= false
+        u = similar(amplitudes)
+        v = similar(amplitudes)
+        tmp = similar(amplitudes)
+        ctT = complex(eltype(system.t))
+
+        is = vcat(iref, zeros(Int, order))
+
+        memory = length(window)
+        if memory < iref
+            window = vcat(window, zeros(eltype(window), iref))
+        end
+    end
+
+    @timeit to "Time loop" begin
+        for i in TelescopeIterator(max(1,imin):iref-1, order, memory)
+            windw = @timeit "Window" prod(window[(i[1] + 1) .- i])
+            iszero(windw) && continue
+
+            is[2:end] .= i
+
+            u .= false
+            v .= false
+
+            @timeit to "Evaluate momenta" evaluate_momenta!(𝐩s, prefactors, system, unique_momenta, indeterminate_momenta, is)
+            preprod = prod(prefactors)
+
+            @timeit to "Ionization" ionization!(u, system, 𝐩s[end], 𝐀, is[end], v)
+
+            @timeit to "Interior orders" begin
+                # We have to perform this loop in chronological order, because
+                # of the pesky time-ordering operator. Since the Goldstone
+                # diagrams are stored with their vertices in
+                # anti-chronological order, we thus have to loop backwards.
+                for j = order:-1:1
+                    a,b = is[j], is[j+1]
+
+                    # First, if there is any interaction (beyond ionization)
+                    # at this vertex, we affect this by simply acting with it
+                    # on the current amplitudes.
+
+                    which = diagram.path[j][2]
+                    if which ≠ 0
+                        @timeit to "Interaction" begin
+                            𝐀ᵢ = 𝐀[b]
+                            𝐤ᵢ = kinematic_momentum(𝐩s[momenta[j]], 𝐀ᵢ)
+                            𝐩ᵢ = kinematic_momentum(𝐩s[momenta[j+1]], 𝐀ᵢ)
+
+                            apply_interaction!(v, system.ions, system.couplings[which],
+                                u, 𝐤ᵢ, 𝐩ᵢ, b)
+                            @timeit to "Swap solutions" begin
+                                @swap!(u,v,tmp)
+                            end
+                        end
+                    end
+
+                    # In-between interactions, it is assumed that the
+                    # interaction-free ion+photoelectron propagator
+                    # factorizes:
+                    #
+                    # U⁽⁰⁾(a,b) = Uᵢₒₙ(a,b) Uₑₗ(a,b)
+                    #
+                    # and we may thus propagate the electrons independently,
+                    # and then mix their channel amplitudes by the "rotation"
+                    # accumulated over the time interval by propagating the
+                    # ions (including external field).
+
+                    # Propagate electrons from b → a
+                    @timeit to "Volkov propagation" begin
+                        Sₑₗ = @timeit to "Volkov phase" volkov_phase(𝐩s[j], system.volkov, a, b)
+                        eSₑₗ = exp(-im*Sₑₗ)
+                        u .*= eSₑₗ
+                    end
+
+                    @timeit to "Ion/channel propagation" begin
+                        # Propagate ions from b → a
+                        u .*= ion_propagation(system.ions, :, a, b)
+                    end
+                end
+            end
+
+            order > 1 && last(first(diagram)) == 0 && error("Recombination not yet implemented in multi-channel case")
+
+            amplitudes .+= u .* (windw*preprod)
+        end
+    end
+
+    @timeit to "Weighting" lmul!(weight, amplitudes)
+end
+
+function multi_channel_photoelectron_spectrum(k::AbstractArray{T},
+                                              system::System, diagram::Diagram;
+                                              iref=length(system.t),
+                                              verbosity=1, kwargs...) where T
+    verbosity == 1 && @info "Ostensibly multi-channel photoelectrum spectrum calculation" system diagram length(k)
+
+    nc = num_channels(system)
+    sk = size(k)
+
+    cT = complex(eltype(T))
+    c = zeros(cT, (nc, sk...))
+    to = TimerOutput()
+    p = Progress(length(k))
+    @timeit to "k loop" begin
+        threaded_range_loop(CartesianIndices(k)) do I
+            tok = TimerOutput()
+            multi_channel_sfa!(view(c, :, I), system, diagram, iref, k[I]; verbosity=verbosity-10, to=tok, kwargs...)
+            ProgressMeter.next!(p)
+            merge!(to, tok, tree_point=["k loop"])
+        end
+    end
+    TimerOutputs.complement!(to)
+    0 < verbosity < 4 && print_timer(to)
+    c
+end
+
+export multi_channel_photoelectron_spectrum

From e50ea245d18d818d7af940008662e04212b8d705 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Tue, 18 Apr 2023 10:56:26 +0200
Subject: [PATCH 16/20] Added test calculations

---
 Project.toml                                |   3 +-
 test/diagrams.jl                            |  59 ++++++
 test/references/krypton-ref-direct.csv      | 200 ++++++++++++++++++++
 test/references/krypton-ref-rescattered.csv | 200 ++++++++++++++++++++
 test/rescattered_ati.jl                     |  93 +++++++++
 test/runtests.jl                            | 123 ++----------
 test/telescope_iterators.jl                 |  49 +++++
 7 files changed, 623 insertions(+), 104 deletions(-)
 create mode 100644 test/diagrams.jl
 create mode 100644 test/references/krypton-ref-direct.csv
 create mode 100644 test/references/krypton-ref-rescattered.csv
 create mode 100644 test/rescattered_ati.jl
 create mode 100644 test/telescope_iterators.jl

diff --git a/Project.toml b/Project.toml
index db49f6d..f74c1cd 100644
--- a/Project.toml
+++ b/Project.toml
@@ -31,7 +31,8 @@ UnitfulAtomic = "1.0"
 julia = "1.6"
 
 [extras]
+DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["DelimitedFiles", "Test"]
diff --git a/test/diagrams.jl b/test/diagrams.jl
new file mode 100644
index 0000000..ab1f949
--- /dev/null
+++ b/test/diagrams.jl
@@ -0,0 +1,59 @@
+@testset "Diagrams" begin
+    function display_diagram(system, diagram, expected_momenta, expected_unique, expected_indeterminate)
+        display(diagram)
+        ions, unique_momenta, momenta, indeterminate_momenta, order = StrongFieldApproximation.analyze_diagram(system, diagram)
+
+        expected_ions = first.(diagram.path)
+        if first(diagram.path)[2] == 0 && length(diagram) > 1
+            expected_ions = expected_ions[2:end]
+        end
+
+        @test ions == expected_ions
+        @test momenta == expected_momenta
+        @test unique_momenta == expected_unique
+        @test indeterminate_momenta == expected_indeterminate
+    end
+
+    @field(F) do
+        λ = 800.0u"nm"
+        I₀ = 1e14u"W/cm^2"
+        cycles = 4.0
+        ϕ = π
+        env = :cos²
+    end
+
+    ndt = 238
+
+    Iₚ = 14u"eV" # "Krypton"
+
+    # Elastic scattering off a Yukawa potential
+    cc = StrongFieldApproximation.CoulombCoupling((𝐤,𝐩) -> yukawa_fourier(𝐩-𝐤, 1, 0, 1))
+    dc = StrongFieldApproximation.DipoleCoupling(0.1, F)
+    couplings=Matrix{StrongFieldApproximation.AbstractCoupling}[reshape([dc],1,1),reshape([cc],1,1)]
+
+    ar = (F, ndt)
+    channel = IonizationChannel(Iₚ, ar...)
+    system = StrongFieldApproximation.System(repeat([channel], 10), nothing, couplings, ar...)
+
+    for (path, expected_momenta, expected_unique, expected_indeterminate) in [
+        ([(1,0)], [1], [(1,2)], []),
+        ([(1,0),(1,0)], [1], [(1,2)], 1:1),
+        ([(2,1),(1,0)], [1,1], [(1,3)], []),
+        ([(2,2),(1,0)], [1,2], [(1,2),(2,3)], 2:2),
+        ([(1,2),(2,2),(1,0)], [1,2,3], [(1,2),(2,3),(3,4)], 2:3),
+        ([(1,0),(1,2),(2,2),(1,0)], [1,2,3], [(1,2),(2,3),(3,4)], 1:3),
+        ([(1,1),(2,2),(1,0)], [1,1,2], [(1,3),(3,4)], 2:2),
+        ([(1,0),(1,1),(2,2),(1,0)], [1,1,2], [(1,3),(3,4)], 1:2),
+        ([(3,1),(1,1),(2,2),(1,0)], [1,1,1,2], [(1,4),(4,5)], 2:2),
+        ([(3,0),(3,1),(1,1),(2,2),(1,0)], [1,1,1,2], [(1,4),(4,5)], 1:2),
+        ([(1,2),(3,1),(1,1),(2,2),(1,0)], [1,2,2,2,3], [(1,2),(2,5),(5,6)], 2:3),
+        ([(1,2),(1,2),(3,1),(2,2),(1,0)], [1,2,3,3,4], [(1,2),(2,3),(3,5),(5,6)], 2:4),
+        ([(1,2),(1,2),(3,1),(1,1),(2,2),(1,0)], [1,2,3,3,3,4], [(1,2),(2,3),(3,6),(6,7)], 2:4),
+        ([(1,2),(10,2),(3,2),(1,2),(2,2),(1,0)], [1,2,3,4,5,6], [(1,2),(2,3),(3,4),(4,5),(5,6),(6,7)], 2:6),
+        ([(1,0),(1,2),(10,2),(3,2),(1,2),(2,2),(1,0)], [1,2,3,4,5,6], [(1,2),(2,3),(3,4),(4,5),(5,6),(6,7)], 1:6),
+        ([(1,0),(1,1),(10,1),(3,1),(1,1),(2,1),(1,0)], [1,1,1,1,1,1], [(1,7)], 1:1)
+    ]
+        diagram = Diagram(path, system)
+        display_diagram(system, diagram, expected_momenta, expected_unique, expected_indeterminate)
+    end
+end
diff --git a/test/references/krypton-ref-direct.csv b/test/references/krypton-ref-direct.csv
new file mode 100644
index 0000000..b7fb9ef
--- /dev/null
+++ b/test/references/krypton-ref-direct.csv
@@ -0,0 +1,200 @@
+-7.543456839434224e-5 + 0.030411217910974103im,-0.026717450903560192 - 0.014526723044745107im
+0.014402924017553302 + 0.033481457661180915im,0.034061084427859524 + 0.012972846833723722im
+0.011763160954343453 - 0.006357732399812491im,-0.001197304147585215 - 0.013317626653558037im
+-0.0061857076059780745 + 0.028946550049159702im,0.006813737700848334 - 0.028805185590083212im
+-0.015200429275289477 - 0.0280517880254329im,0.012323315572670858 - 0.02942943687658876im
+0.00313342519327927 + 0.006052800923066901im,-0.0011205578770404259 - 0.006723027777749032im
+-0.004471511838151416 - 0.00851116681992856im,-0.006207933263113997 - 0.007341385656408818im
+0.017715082539488287 + 0.013436964929771436im,-0.016052420114272856 - 0.015384927181490881im
+-0.018879104894497664 + 0.007555020382198454im,-0.0063941309245120655 - 0.019303212797697852im
+0.004539554586919643 - 0.021057023275245203im,0.002473947107340369 - 0.021398256255406266im
+0.01031074419627424 + 0.009452891141951632im,0.008052155807108038 - 0.01143815473232827im
+-0.011248553650085256 + 0.002189885227054618im,0.009672997148647748 - 0.0061448094102890555im
+0.0019234058252181753 - 0.003933671218754664im,0.004303992976796319 - 0.0008055455799842776im
+-0.000739944544659966 + 0.003639798463909571im,0.0033053765290337743 - 0.0016941478059331859im
+0.0011507411615440252 - 0.003121666828478544im,0.0013008780659914492 - 0.0030621439003271135im
+0.003916769783418816 + 0.004343619109917442im,0.004743290521024111 - 0.0034218865471406035im
+-0.006070247493822477 + 0.00015526262891026092im,0.005791646906141042 - 0.0018245101355647455im
+0.003008373264839285 - 0.006654171197994939im,0.007032941802134234 + 0.0019662231919554745im
+0.003971835690727766 + 0.005237752245759651im,0.004564681514173904 + 0.004730032771159276im
+-0.006576831675118637 + 0.001052187357547507im,0.0015171305052965693 + 0.006485378026617599im
+0.0015356376010813167 - 0.0055285379053008525im,-0.0020673827167768137 + 0.005352461388423846im
+0.0037088518144974205 + 0.003591118021031054im,-0.004038261659299098 + 0.0032162327642668133im
+-0.004024464548566049 + 0.0016330715702084682im,-0.004335200759917964 + 0.00026319579645109694im
+0.00037413421785105526 - 0.003576790274601549im,-0.0031938355125409835 - 0.0016531242543423952im
+0.002545970096904877 + 0.0014785653013955884im,-0.0013781158257094124 - 0.0026017140226673786im
+-0.0019258781836560958 + 0.00121629989795367im,0.0001354391182342168 - 0.0022737740576641093im
+-0.00021831613958728395 - 0.0017926098364533264im,0.001086189990394177 - 0.0014426722660809052im
+0.0012258817315888549 + 0.00045188771706886207im,0.00121723610986241 - 0.000474683875363568im
+-0.0007334809135692967 + 0.0006555354333322436im,0.0009699023618395331 + 0.0001643482991167209im
+-0.00015830408668188536 - 0.0006424034344264202im,0.0005025695460853242 + 0.000430309432586199im
+0.0004293129991974957 + 0.00015594016626833805im,0.0001882296213645946 + 0.0004161689517237313im
+-0.00022765407256953142 + 0.00016440124544514308im,-1.7605244916628566e-6 + 0.00028080428560871436im
+5.149644893587773e-5 - 0.00018130184047310626im,-2.4815827672420603e-5 + 0.00018683258899053378im
+5.465429628450371e-5 + 0.000132967061707992im,-3.639625301754912e-5 + 0.00013907783564764156im
+-0.00013152093484598918 - 8.628507886686319e-5im,-4.848269709403252e-5 + 0.00014964056676002008im
+0.00015986949409317307 - 6.151892016883622e-5im,-0.00011275210693399997 + 0.0001289565626967607im
+-4.558874660738214e-5 + 0.0001773758850697455im,-0.00016543993033795335 + 7.855041611475593e-5im
+-0.00014195439066683852 - 0.00012720382782690142im,-0.00018954721457048725 - 2.0092692476145324e-5im
+0.00017316479020220517 - 6.351751008039763e-5im,-0.00014739130110669687 - 0.00011088878667666089im
+-1.7987146319019017e-5 + 0.00017864661578211884im,-6.15235313902281e-5 - 0.00016868018807549102im
+-0.00014232120952626653 - 8.56031834671347e-5im,3.900535032634471e-5 - 0.00016143671932435712im
+0.0001281631784429033 - 8.592895535461671e-5im,0.00011348748413973027 - 0.00010454748500614394im
+2.3612084418378504e-5 + 0.00013664125232958673im,0.00013738665848485614 - 1.8795436641114076e-5im
+-0.00011975840122570262 - 3.260608312552278e-5im,0.00011148630978735953 + 5.455303888720523e-5im
+6.917950721649097e-5 - 8.60034899699037e-5im,5.285206055316999e-5 + 9.6897183660673e-5im
+4.330551510197671e-5 + 8.498411929821834e-5im,-1.0140495585392901e-5 + 9.484112251618475e-5im
+-8.311737041426236e-5 + 4.019814918331722e-6im,-5.433888782132178e-5 + 6.302334048599717e-5im
+2.5041729273910693e-5 - 6.680332717796547e-5im,-6.93271890911724e-5 + 1.6837861485335368e-5im
+4.3247268085769466e-5 + 4.2222658156639814e-5im,-5.6241801975582954e-5 - 2.2134560511793763e-5im
+-4.7458307699743654e-5 + 1.9268105658904595e-5im,-2.8052999158008787e-5 - 4.285533926305159e-5im
+8.356848631473934e-7 - 4.323492996227042e-5im,1.8750442339436873e-6 - 4.32023349696556e-5im
+3.298277434675296e-5 + 1.4706145192123155e-5im,2.246100698584512e-5 - 2.827785838863632e-5im
+-2.219686465066255e-5 + 2.0243438726722336e-5im,2.8841993519084745e-5 - 8.404583364339023e-6im
+-8.295100642198825e-6 - 2.3871184292029407e-5im,2.3829261920886725e-5 + 8.41477334301563e-6im
+2.0767344284181513e-5 + 1.276306156977402e-6im,1.1734317378133944e-5 + 1.7181889933922335e-5im
+-7.877595032397498e-6 + 1.5080514891701442e-5im,-1.4935973764689379e-7 + 1.701341013932351e-5im
+-8.882714227546477e-6 - 1.0991292214150733e-5im,-8.29180480069675e-6 + 1.1443648449868036e-5im
+1.0945156666086541e-5 - 3.2455577244230903e-6im,-1.088168166346593e-5 + 3.4524054660605488e-6im
+-1.1520542434392214e-6 + 9.011633110757094e-6im,-8.664782959910216e-6 - 2.730988192485856e-6im
+-6.398042961711891e-6 - 3.788922504944451e-6im,-4.470460023902267e-6 - 5.941874675863705e-6im
+4.700875486145184e-6 - 3.70520459466984e-6im,8.574577626438735e-8 - 5.984932671848739e-6im
+1.2562750466198507e-6 + 4.5141369228277405e-6im,2.8216105276320594e-6 - 3.7408786620458613e-6im
+-3.826193269585282e-6 - 5.598340044897815e-7im,3.6786769229026737e-6 - 1.1918494653655095e-6im
+1.5530207644364218e-6 - 2.7949898067099165e-6im,3.000569534014113e-6 + 1.1047280086831206e-6im
+1.5245612548825362e-6 + 1.9463184409895868e-6im,1.3738807499609278e-6 + 2.0554547383763185e-6im
+-1.97903958019571e-6 + 3.934595747631806e-7im,5.943573559748471e-8 + 2.0168974900373704e-6im
+3.05905014372452e-7 - 1.6450801392348956e-6im,-9.700593157718968e-7 + 1.3633970298857924e-6im
+1.0208395511138652e-6 + 6.543067695331409e-7im,-1.1545271526365437e-6 + 3.7053716951062225e-7im
+-8.091618959904539e-7 + 4.381505589944536e-7im,-8.994005960649489e-7 - 1.9441566970574282e-7im
+-1.1981041483274578e-7 - 7.599054524868312e-7im,-4.776273402777731e-7 - 6.030613248569183e-7im
+5.33546350268322e-7 + 4.5561869423670874e-8im,6.227412477050994e-8 - 5.318547986123869e-7im
+-2.0769624337199643e-7 + 3.230803390118312e-7im,2.349932424708105e-7 - 3.03803901144237e-7im
+-2.3297990984890608e-7 - 3.102658609870946e-7im,3.768956836573617e-7 - 9.216391377542706e-8im
+2.822575758064604e-7 - 1.4405817398638368e-7im,2.586519851682761e-7 + 1.8308809236646242e-7im
+-2.5852334984874314e-8 + 2.109618281293278e-7im,1.057361750835099e-7 + 1.843721739938294e-7im
+-1.7241457658465395e-7 - 1.565542882055257e-7im,2.510689851868612e-8 + 2.3152899796469723e-7im
+1.45734338086245e-7 - 1.06455519714196e-7im,-1.1863141386978108e-7 + 1.3600685162734137e-7im
+-1.2531189418968284e-8 + 7.862941580692987e-8im,-6.50563947234263e-8 + 4.5905134723616464e-8im
+-5.7332944266989374e-8 - 7.720084621506259e-8im,-8.998047071445802e-8 + 3.3919789576373916e-8im
+3.213844385958787e-8 - 4.891350540585555e-8im,-3.072176625501048e-8 - 4.981549657238044e-8im
+9.959744830768028e-9 - 1.8875770655409374e-8im,1.795483873441452e-8 + 1.1537546207694101e-8im
+-8.798047184472035e-9 - 8.118193305367473e-9im,-2.7453801261855776e-9 - 1.1652191888000551e-8im
+-2.1930969362190874e-8 - 5.208083853861391e-8im,5.269457762764693e-8 + 2.041231380046882e-8im
+3.818910985712693e-8 - 2.8727810668632856e-8im,-3.4914900191298614e-9 + 4.766030673434586e-8im
+-2.341500765621712e-8 + 5.726028164169754e-9im,1.849855201816782e-8 + 1.5454886863798676e-8im
+-4.138333992275884e-9 - 5.787555165850619e-8im,8.869867756068554e-10 + 5.801653593555485e-8im
+2.1485511093019736e-8 - 2.517233077245855e-9im,-1.8849246711868553e-8 + 1.0614590419445874e-8im
+-2.3159572109128783e-8 - 2.4099871674866045e-8im,1.3225562588208416e-8 + 3.069615608834939e-8im
+1.7328995496056558e-8 - 2.7121531771175597e-8im,-2.4346632620683494e-8 + 2.1050252536852785e-8im
+-1.1310617588002137e-8 - 7.69430945953743e-9im,1.3122992660918033e-8 + 3.8625824417613056e-9im
+5.566549302329223e-9 - 3.294568763233037e-8im,-5.765088650937689e-9 + 3.291152697429155e-8im
+1.8491296308482996e-9 - 1.4899343594953382e-9im,-2.1104936255560367e-9 - 1.0885751440346057e-9im
+-1.0041089556994748e-8 - 2.9713516512589074e-8im,1.275990279521883e-8 + 2.8651377785025365e-8im
+1.5459627434958282e-8 - 1.0366490582163345e-8im,-1.4272324698619597e-8 + 1.1948431465139256e-8im
+-1.5457606282658126e-8 - 1.14437222480803e-8im,1.502957658578102e-8 + 1.2000340260611414e-8im
+1.0406021515886165e-8 - 2.5392694849952696e-8im,-1.0161579846460135e-8 + 2.549150019706583e-8im
+-2.852996244986391e-9 + 2.2335862291877186e-11im,2.7758191420134953e-9 + 6.594722178472231e-10im
+-4.44507007691694e-9 - 2.5719297253258618e-8im,3.803945429073364e-9 + 2.5821907426814368e-8im
+9.628612726292212e-9 - 4.625187421745887e-9im,-9.786254401266537e-9 + 4.281564928006281e-9im
+-1.177026955247262e-8 - 1.399793099808178e-8im,1.1836349879615009e-8 + 1.3942101439244922e-8im
+1.0645311507017903e-8 - 1.5722417195802102e-8im,-1.0729408158188289e-8 + 1.566514756490854e-8im
+-6.7313454166753235e-9 - 2.3640395289809704e-9im,6.783050315988381e-9 + 2.211317361433909e-9im
+1.2263636152536373e-9 - 2.1086166425680713e-8im,-1.097074476219963e-9 + 2.109328851757571e-8im
+4.221138838816121e-9 - 8.959560831632557e-10im,-4.2077341421878656e-9 + 9.569268158577502e-10im
+-8.053186760584072e-9 - 1.517512876667954e-8im,8.050323211025882e-9 + 1.5176648750601133e-8im
+9.309030024886283e-9 - 8.814040977156063e-9im,-9.289291282622378e-9 + 8.83484201121718e-9im
+-7.864761828290392e-9 - 4.696483139735074e-9im,7.84735792384156e-9 + 4.72551143508538e-9im
+4.363707331639217e-9 - 1.588694373805264e-8im,-4.385862821620852e-9 + 1.588084210744543e-8im
+5.1292482638708165e-11 - 4.189642792205499e-12im,-5.131418151658949e-11 - 3.916184589909714e-12im
+-4.0938252025840765e-9 - 1.4556400662421284e-8im,4.091209353105025e-9 + 1.4557136427160202e-8im
+6.6830350404629365e-9 - 4.060002096070076e-9im,-6.686620401500871e-9 + 4.054097848336838e-9im
+-7.218656495305655e-9 - 6.781822505121999e-9im,7.222780512772376e-9 + 6.7774270514241475e-9im
+5.710468439484452e-9 - 1.0899069776170925e-8im,-5.7076995003785616e-9 + 1.0900521093052134e-8im
+-2.734340978933851e-9 - 5.85617087139724e-10im,2.7342326203912404e-9 + 5.861209629229477e-10im
+-7.708417185438558e-10 - 1.2767845496382351e-8im,7.717788205215625e-10 + 1.2767787085278765e-8im
+3.79151332009692e-9 - 1.3054066783351853e-9im,-3.791082424263053e-9 + 1.3066576523290186e-9im
+-5.527715545027323e-9 - 8.02435189790008e-9im,5.526885633988025e-9 + 8.024922862743073e-9im
+5.5966158239831245e-9 - 6.696034030408838e-9im,-5.596897941473024e-9 + 6.695799068319367e-9im
+-4.106440102184067e-9 - 1.857440520472437e-9im,4.106359858503222e-9 + 1.8576174117821984e-9im
+1.5879800807243532e-9 - 1.0272308297922903e-8im,-1.5883611857338268e-9 + 1.0272249443657547e-8im
+1.1831174385063783e-9 - 1.4040165777783176e-10im,-1.183143451286075e-9 + 1.401693809963749e-10im
+-3.419645741187761e-9 - 8.332616554116638e-9im,3.419646733057342e-9 + 8.332616848590638e-9im
+4.540604516197052e-9 - 3.5302002024164296e-9im,-4.5406542528867575e-9 + 3.5301369491536406e-9im
+-4.317274837910226e-9 - 3.2046090823471124e-9im,4.317292633408501e-9 + 3.204582833158223e-9im
+2.9099916855878497e-9 - 7.5671263946946e-9im,-2.9099954714638834e-9 + 7.567124966417843e-9im
+-7.958322324740731e-10 - 7.638641087531535e-11im,7.958285297529995e-10 + 7.641857391036855e-11im
+-1.3870588555455655e-9 - 7.812763013819974e-9im,1.3870062181986007e-9 + 7.812773000798283e-9im
+3.027351883083148e-9 - 1.4516265879144965e-9im,-3.027352714518307e-9 + 1.45162531975587e-9im
+-3.7083876906783693e-9 - 4.240675992981228e-9im,3.7083395089789416e-9 + 4.24072021926289e-9im
+3.3092044521750295e-9 - 5.041581579507456e-9im,-3.3092507805351524e-9 + 5.04155075929149e-9im
+-2.017374873724508e-9 - 6.449677024657133e-10im,2.017370052675183e-9 + 6.449835711702838e-10im
+2.5476900053096944e-10 - 6.697242920090965e-9im,-2.5482532067306464e-10 + 6.697240760259274e-9im
+1.4555465989251393e-9 - 3.460188638899199e-10im,-1.4555506570560423e-9 + 3.460015549352028e-10im
+-2.641773480964545e-9 - 4.78101575194497e-9im,2.6417448679792627e-9 + 4.781030616813682e-9im
+3.0103141693068486e-9 - 2.9555339996238736e-9im,-3.0103332450263876e-9 + 2.955514849557273e-9im
+-2.5151632620234454e-9 - 1.4506843696732794e-9im,2.515155764405974e-9 + 1.4506974470536945e-9im
+1.3532907811975836e-9 - 5.261400144251897e-9im,-1.3533219567926281e-9 + 5.261391762118579e-9im
+1.0642251578337974e-10 - 2.0919964837983413e-12im,-1.064214271413041e-10 + 2.092761662923235e-12im
+-1.4375863045107808e-9 - 4.798651032783653e-9im,1.4375559853370423e-9 + 4.7986600583552485e-9im
+2.2783406020091634e-9 - 1.4363065001686742e-9im,-2.278348948771726e-9 + 1.4362929624985666e-9im
+-2.427709885586761e-9 - 2.2029431430035703e-9im,2.427697091103722e-9 + 2.2029579561060026e-9im
+1.89037117888506e-9 - 3.763448441746972e-9im,-1.8903937242294524e-9 + 3.763436251149018e-9im
+-8.617858526117702e-10 - 1.6938042148272864e-10im,8.617851358126258e-10 + 1.6938364581391862e-10im
+-3.390323069645237e-10 - 4.374806885436236e-9im,3.390088848008817e-10 + 4.374808501453604e-9im
+1.3661818943295403e-9 - 4.967693159325117e-10im,-1.3661842200634428e-9 + 4.967613393531705e-10im
+-1.9449707888002177e-9 - 2.7234260982125e-9im,1.944957695449126e-9 + 2.72343662169197e-9im
+1.9434468405207473e-9 - 2.407153695637499e-9im,-1.9434596498643404e-9 + 2.4071433236575473e-9im
+-1.3998928320883643e-9 - 6.047814554802864e-10im,1.3998898983498652e-9 + 6.047892856283337e-10im
+5.009420469001251e-10 - 3.6501031362737527e-9im,-5.00959746085492e-10 + 3.6501009804733358e-9im
+4.798315563817024e-10 - 6.518517838022446e-11im,-4.798313693158758e-10 + 6.51833481035398e-11im
+-1.2637620732347236e-9 - 2.9362297093403045e-9im,1.2637483499976266e-9 + 2.9362361378596283e-9im
+1.6448571946663233e-9 - 1.3242963403698828e-9im,-1.6448632461296426e-9 + 1.3242892634019976e-9im
+-1.5425779246156446e-9 - 1.1048321482938567e-9im,1.5425732624424747e-9 + 1.1048382789722342e-9im
+1.0161471941801725e-9 - 2.7827211616169275e-9im,-1.0161592425449877e-9 + 2.7827163753350398e-9im
+-2.3928376122332265e-10 - 1.887153322053682e-11im,2.39283975187791e-10 + 1.8873220958625272e-11im
+-5.555203424385621e-10 - 2.8469311614193307e-9im,5.555073425821217e-10 + 2.8469333496092767e-9im
+1.1453554893600768e-9 - 5.73355517637581e-10im,-1.1453582768015782e-9 + 5.733509916091529e-10im
+-1.378489296883785e-9 - 1.5233027812980264e-9im,1.3784813514072813e-9 + 1.5233094089876584e-9im
+1.2120797010277545e-9 - 1.91686388416693e-9im,-1.2120879915901373e-9 + 1.916858223603399e-9im
+-7.174234494412367e-10 - 2.159917774607578e-10im,7.174246312665577e-10 + 2.1599590540976218e-10im
+5.290797441973872e-11 - 2.5163532173592997e-9im,-5.291894951786131e-11 + 2.5163546860746635e-9im
+5.854519987247763e-10 - 1.4973483601283103e-10im,-5.854542329819052e-10 + 1.4973247184184223e-10im
+-1.0210505992266313e-9 - 1.7758660574013023e-9im,1.0210437877497258e-9 + 1.7758703081225684e-9im
+1.1446632066745436e-9 - 1.1627055910713753e-9im,-1.1446689039905846e-9 + 1.1627005618873854e-9im
+-9.408508720727064e-10 - 5.215480220724542e-10im,9.408485752427457e-10 + 5.215504262812164e-10im
+4.864659880771292e-10 - 2.0340391699115396e-9im,-4.864741493240324e-10 + 2.0340370975014183e-9im
+7.633827435155825e-11 - 2.800663912982929e-12im,-7.633809078005535e-11 + 2.7996506894195967e-12im
+-5.837882806630385e-10 - 1.8346325870596102e-9im,5.837816826124129e-10 + 1.834634271151909e-9im
+8.974410178587642e-10 - 5.873161571597568e-10im,-8.974425564649324e-10 + 5.8731398087289e-10im
+-9.41338085579106e-10 - 8.254447676960161e-10im,9.413350205193961e-10 + 8.254483714241754e-10im
+7.193976733696434e-10 - 1.495052131325759e-9im,-7.194027708907697e-10 + 1.4950496130959195e-9im
+-3.0957489080074833e-10 - 5.539628148318073e-11im,3.095743960776314e-10 + 5.5398235255703796e-11im
+-1.623770330388753e-10 - 1.71602401109176e-9im,1.62371367357657e-10 + 1.716024865050107e-9im
+5.608302520792024e-10 - 2.1479182863314827e-10im,-5.608302512723502e-10 + 2.1478945893635196e-10im
+-7.787776854079451e-10 - 1.0521839697531692e-9im,7.787739908758316e-10 + 1.0521854238112127e-9im
+7.659884915792344e-10 - 9.824671141273876e-10im,-7.659909353194322e-10 + 9.82465134875959e-10im
+-5.397000562166127e-10 - 2.2235145025562565e-10im,5.396994996963463e-10 + 2.2235302000280616e-10im
+1.7578568068778577e-10 - 1.4654008577632635e-9im,-1.7579019962738507e-10 + 1.4654007979941763e-9im
+2.1587345193656833e-10 - 3.307433660781962e-11im,-2.1587297206470557e-10 + 3.3073251869547354e-11im
+-5.241244513334675e-10 - 1.162900675127491e-9im,5.241208114680894e-10 + 1.162902089664132e-9im
+6.676485446274053e-10 - 5.565031327410344e-10im,-6.676504661624182e-10 + 5.565002171334925e-10im
+-6.159659974426167e-10 - 4.25514545266569e-10im,6.159649268591457e-10 + 4.255163011222203e-10im
+3.9495266008446204e-10 - 1.1414435355930052e-9im,-3.9495600786650755e-10 + 1.1414423601878245e-9im
+-7.633577341703944e-11 - 4.716657424297993e-12im,7.633600275902988e-11 + 4.7168387667840456e-12im
+-2.4506135247080094e-10 - 1.1513942851325967e-9im,2.450582790558923e-10 + 1.1513948213987964e-9im
+4.791513794075973e-10 - 2.501356400497784e-10im,-4.791530459168842e-10 + 2.501338768828202e-10im
+-5.655318943570717e-10 - 6.039786807689464e-10im,5.655302007348399e-10 + 6.039808293726041e-10im
+4.886164957127593e-10 - 8.02620515362453e-10im,-4.886181318531809e-10 + 8.02620252169611e-10im
+-2.7942297961910844e-10 - 7.896917972607164e-11im,2.794222229410126e-10 + 7.897055834582488e-11im
+4.166801587134357e-12 - 1.0362019694683812e-9im,-4.16946576915406e-12 + 1.0362020201821234e-9im
+2.5628344993654204e-10 - 7.021003806176304e-11im,-2.562837901732432e-10 + 7.02090407633712e-11im
+-4.299013960063708e-10 - 7.191293832032713e-10im,4.2989879540011754e-10 + 7.191314240581321e-10im
+4.73122734644907e-10 - 4.972122349828902e-10im,-4.7312406020862e-10 + 4.972110222485415e-10im
+-3.814138526398132e-10 - 2.03032038404878e-10im,3.814143335533278e-10 + 2.0303325361109426e-10im
+1.8825836302443822e-10 - 8.508489474493774e-10im,-1.8826039729047664e-10 + 8.508490918761383e-10im
+4.639445606117728e-11 - 2.492032477051708e-12im,-4.639368444436837e-11 + 2.4930050263032215e-12im
+-2.5445839393321117e-10 - 7.549986486365908e-10im,2.544563380367115e-10 + 7.549993843229435e-10im
+3.792266370757477e-10 - 2.5749259657636945e-10im,-3.7922785370040263e-10 + 2.574919181007451e-10im
+-3.906063595219132e-10 - 3.3094053460573853e-10im,3.906050953127644e-10 + 3.309422762344281e-10im
+2.91988213033559e-10 - 6.34184390139045e-10im,-2.919897205215089e-10 + 6.341828591434792e-10im
diff --git a/test/references/krypton-ref-rescattered.csv b/test/references/krypton-ref-rescattered.csv
new file mode 100644
index 0000000..7805fd5
--- /dev/null
+++ b/test/references/krypton-ref-rescattered.csv
@@ -0,0 +1,200 @@
+-3.6687003120887313e-5 + 0.0025902157396208896im,-0.0010612424880033792 - 0.0004568511318403792im
+0.001985602183576178 + 0.0017290817653445266im,0.0011675342556811952 + 0.0010704282386948811im
+4.17180012480989e-5 - 0.0015343936628718293im,-3.423020132959156e-5 - 0.00045809793248300376im
+3.7015907458423824e-5 + 0.0013637527393382395im,0.00012454874637973117 - 0.0006882425696890472im
+-0.0012469577310956982 - 0.0008341163484935013im,0.0009893291153856346 - 0.0019344046445962204im
+0.001111123051414318 + 0.0003665981093941261im,0.0001327256000725625 + 0.00011146654467673922im
+-0.0006099801076344832 - 0.0003428607074792995im,0.0007327184889421969 - 0.0004036361665507387im
+0.0006781200590990004 - 3.1807850603184495e-5im,-0.0009864516123152098 - 0.00013365549980063645im
+-0.0008899482765173228 + 0.0009365491766828271im,0.00033784691981752146 - 0.0007285978672300551im
+7.122948855937333e-5 - 0.0010758709881924283im,-0.00027874089661010694 - 0.0009743877087091432im
+0.0010993492574300598 + 0.0006282393499952273im,0.0010824888126975253 - 0.00041187704237668343im
+-0.0010249762709698775 + 0.00021243747827111284im,0.00046616391816261995 - 0.0002415151204031273im
+0.00015574943950028074 - 0.0008556257446528235im,0.000752571330047633 + 0.0005444043839465224im
+0.0004607071332306947 + 0.0007988723961733406im,-3.872793650846325e-5 + 0.00029900526681552765im
+-0.0005101110951618798 - 6.239536548379986e-5im,-8.891694081265224e-5 + 0.000415631312524325im
+0.0003773407408717474 - 0.0005023301363900816im,-0.0002835775486062551 - 8.905502403670288e-5im
+-3.5204280884419047e-5 + 0.000545731540742111im,-7.59366740146085e-5 - 0.00013289230821707655im
+-0.00027763546245430156 - 0.00019307895108302195im,0.00014382337131556934 - 0.0001827125446159563im
+0.0004729444980347253 - 0.00020864536454528876im,0.00021891339260279733 - 4.49927696217649e-5im
+-0.0002188975115403687 + 0.0003277295203167836im,0.000243816579139473 + 0.000246392301170371im
+-0.00025460721221081604 - 0.00025318857387841745im,7.90793155536423e-6 + 0.0001991872302093402im
+0.00042886046462662036 - 1.0665350678691879e-5im,-0.0001327640362672044 + 0.0003613115397920521im
+-0.0001808262666553838 + 0.000275194984650378im,-0.00028252702164809323 - 1.7043918561582587e-5im
+-0.00021881845187093153 - 0.0002660084702132266im,-0.00027816625055810957 + 2.815921285653043e-5im
+0.00034838832630310263 + 1.3653422918705469e-5im,-0.0001178038387114123 - 0.0003276462066210169im
+-9.755155000181798e-5 + 0.0002608417527224715im,-3.525060413255612e-5 - 0.0001582727279933934im
+-0.0001928463299152632 - 0.0002266380414778093im,0.00022041968965394655 - 0.00023860851476570827im
+0.00025515618371665174 - 4.8491532066517906e-5im,0.0001179760200427581 - 5.9811391380897225e-6im
+-2.788071469370714e-5 + 0.00023422916793315628im,0.0002703823433902883 + 7.508719964669557e-5im
+-0.00018141959417192192 - 0.00013553717994957076im,-2.9191795715836844e-5 + 0.00011947182833174507im
+0.0001635758493609913 - 0.00010244407509543268im,9.115785917402933e-5 + 0.00020881988619842653im
+3.4703137745141996e-5 + 0.00018177705152273214im,-0.00020575781665736845 - 2.1314245968031917e-6im
+-0.00015423533121796385 - 3.878340458631923e-5im,-1.6677454683418112e-6 + 0.00013591891384336546im
+8.04769722865082e-5 - 0.00011833210443088209im,-0.00018634581658330798 - 0.00019139007777916792im
+7.267918919741456e-5 + 0.00010985878704074732im,5.861175520918026e-5 + 8.486704842002747e-5im
+-0.0001065644708094092 + 2.3081678177338842e-5im,-6.143163919628062e-5 - 0.00024140955028113117im
+1.9117943168058086e-5 - 9.697779190142312e-5im,0.0001112956895226135 + 0.0001387805508501105im
+7.17945254748146e-5 + 4.766128119736359e-5im,2.7828948005460773e-6 - 0.00017766364462414995im
+-5.838499465591056e-5 + 3.890656288197639e-5im,6.957615743844536e-5 + 0.00019239998919258957im
+-4.871037697058343e-6 - 6.0168770919663004e-5im,-6.820224454261186e-6 - 0.00012122758506659925im
+4.7325165766593166e-5 + 1.9715933824039773e-5im,-2.9857420438932528e-6 + 0.0001657820962438228im
+-3.352670967740023e-5 + 2.4282890006871335e-5im,-1.2713686367377363e-5 - 0.00010250066742424057im
+3.071916281373257e-6 - 3.7571748736517864e-5im,-3.491109117480372e-5 + 9.154843578098649e-5im
+2.9126229136263452e-5 + 2.4498497928885862e-5im,1.305778898873635e-5 - 7.631328847910042e-5im
+-3.75693483973441e-5 + 9.047179399497551e-6im,-3.107498968500857e-5 + 2.7492928972143134e-5im
+1.6361240221823514e-5 - 4.1555301085292275e-5im,3.913717210526768e-5 - 2.0816831112023287e-5im
+3.3844703927975e-5 + 3.7706013109929064e-5im,-2.627361759840326e-5 - 1.5709618510365135e-5im
+-5.194974064925645e-5 + 1.3888495040521972e-5im,4.1614814017429145e-5 + 4.1793186835035545e-5im
+1.1812306655782881e-5 - 5.773791882304289e-5im,-2.9805666591081588e-5 - 5.50440755972182e-5im
+5.200065704180309e-5 + 3.4953969984530314e-5im,2.8549494937256815e-5 + 8.759240340834821e-5im
+-5.21427401968616e-5 + 3.40425792459225e-5im,-2.9828190126965434e-5 - 9.601960299749844e-5im
+-9.807864258554074e-6 - 6.16264629033799e-5im,1.6192965970331915e-5 + 0.00011440153109425644im
+6.026498519381665e-5 + 1.3227821048812498e-5im,-1.9362379286129368e-5 - 0.0001292157049335489im
+-3.157995301716882e-5 + 4.7239139967128036e-5im,9.446896610562812e-6 + 0.00013103065917528im
+-2.7806814202969058e-5 - 4.3497736528966086e-5im,-3.327947361443773e-6 - 0.00014610541747825344im
+4.588504662095776e-5 - 8.780851285049196e-6im,3.975820918886292e-6 + 0.00014273537024900242im
+-6.677339104688946e-6 + 3.785768155996981e-5im,9.725277553627003e-6 - 0.00014719729210669566im
+-2.417099267482592e-5 - 1.6958152290897316e-5im,-4.62339138284058e-6 + 0.00014814415184992506im
+1.9174039329084243e-5 - 1.0930894995029525e-5im,1.5678112469404078e-5 - 0.00013953497594469598im
+9.091465702806217e-7 + 1.2605185105726217e-5im,-1.5520387764354287e-5 + 0.00014365640213891484im
+-1.6334443254224141e-6 - 4.478010655131783e-6im,1.705099887761552e-5 - 0.00012996729771370255im
+2.7279419452844055e-6 + 8.17540214557499e-6im,-2.3594176011948666e-5 + 0.0001293594109874924im
+-1.4293082894233804e-5 - 6.6366843593085276e-6im,1.889063119754962e-5 - 0.00011987801828563588im
+1.9430811303761872e-5 - 1.5475191115564262e-5im,-2.533925222590198e-5 + 0.00011079196853524834im
+9.52331603976988e-6 + 3.037916226017769e-5im,2.2425565691806128e-5 - 0.000106559462015419im
+-3.7033469764640635e-5 - 3.842813541540624e-6im,-2.2917856827889826e-5 + 9.387177285137092e-5im
+2.0546732952077137e-5 - 3.812771819422136e-5im,2.4423232766863604e-5 - 8.946839745002731e-5im
+3.15459922872445e-5 + 3.550738485667201e-5im,-2.0648699305265655e-5 + 7.94850103001702e-5im
+-4.622497834953893e-5 + 1.7047185916843394e-5im,2.288621488104551e-5 - 7.22532434821666e-5im
+1.5180002953149224e-6 - 5.1234265896352324e-5im,-1.9283275620179594e-5 + 6.544875399311335e-5im
+4.830899343632553e-5 + 1.932089090591223e-5im,1.9917691033840402e-5 - 5.786902105454994e-5im
+-3.3823427253549405e-5 + 3.702997113452878e-5im,-1.6870737679901416e-5 + 5.1948331057487764e-5im
+-2.0836492887931195e-5 - 4.3354077479320254e-5im,1.7585266917007728e-5 - 4.538600338042517e-5im
+4.550245286234667e-5 - 4.1602706760021715e-6im,-1.3719683844681776e-5 + 4.1244710500694256e-5im
+-1.0574201575646427e-5 + 3.9560961062423485e-5im,1.463478086469438e-5 - 3.3654712948954315e-5im
+-2.8461373335420596e-5 - 2.161552709169782e-5im,-1.2237494265532484e-5 + 3.298900281658954e-5im
+2.6445067117536818e-5 - 1.6128744656432012e-5im,9.925624765794493e-6 - 2.4692774448776918e-5im
+4.739513458300591e-6 + 2.412463667623693e-5im,-1.2073611210777255e-5 + 2.4309955907353168e-5im
+-1.7130063630361912e-5 - 4.011693472193267e-6im,6.081237747937999e-6 - 2.028529112212639e-5im
+7.6267687331977826e-6 - 8.913661617674163e-6im,-9.791344007199806e-6 + 1.5118808976823555e-5im
+1.4071599257549728e-6 + 5.064842725165203e-6im,6.107998409440173e-6 - 1.7694773722842367e-5im
+1.5115243794354421e-6 - 3.7741999840697305e-6im,-4.5729525728169485e-6 + 9.273004383102707e-6im
+4.183828331960153e-6 + 8.974288996852339e-6im,8.099357787096011e-6 - 1.2647212311947853e-5im
+-1.5505005451015796e-5 - 1.271182973853892e-6im,-5.695391854997875e-7 + 8.714792639950302e-6im
+1.0650689917595885e-5 - 1.944562268914237e-5im,7.202872901144672e-6 - 5.641430186174454e-6im
+1.8275978424512387e-5 + 2.1030834951124523e-5im,-1.3207181484217209e-6 + 9.661409678617662e-6im
+-3.0570558264490246e-5 + 1.0740962133676175e-5im,2.564977412847403e-6 - 1.5356129648519836e-6im
+1.5170907203124955e-6 - 3.740615761629197e-5im,-4.448053984193372e-6 + 7.265424965725659e-6im
+3.8744134986097e-5 + 1.5819937443120535e-5im,-1.2336892312209032e-6 - 2.5614972377945033e-6im
+-3.031237144118643e-5 + 3.294364239316611e-5im,-4.756210726725406e-6 + 2.0826940580805467e-6im
+-2.1122001904468074e-5 - 4.288219342887406e-5im,-5.094418915757945e-7 - 5.021145717915204e-6im
+5.0348046855492664e-5 - 5.487589686074348e-6im,-1.2361096637502915e-6 - 1.058101530089892e-6im
+-1.2303598495204043e-5 + 5.049918755034021e-5im,2.652082493339762e-6 - 4.29109206290805e-6im
+-4.368628098019674e-5 - 3.006331305312585e-5im,2.001626147623281e-6 + 1.533122634673582e-7im
+4.4319578960326526e-5 - 3.135169840887621e-5im,3.5904007808308557e-6 - 6.891823680532185e-7im
+1.4611656100543642e-5 + 5.229917623550305e-5im,1.5910822545279058e-6 + 2.7424856584947454e-6im
+-5.349797598647989e-5 - 4.619130254133441e-6im,1.2169841614552618e-6 + 1.7190353669941458e-6im
+2.2892872516133633e-5 - 4.83978469794896e-5im,-9.9009062065149e-7 + 2.8487038876074515e-6im
+3.72747533335372e-5 + 3.712283656228503e-5im,-1.2916937536917525e-6 + 8.375056956590851e-7im
+-4.61074511085138e-5 + 2.139184768306817e-5im,-2.082514067957503e-6 + 5.114664372508876e-7im
+-3.759926201093145e-6 - 4.940335290354198e-5im,-1.2055539506313733e-6 - 1.2590005773825557e-6im
+4.634051142528236e-5 + 1.2630111864006516e-5im,-6.655366318541492e-7 - 1.2603771354534885e-6im
+-2.6245941542160258e-5 + 3.729387384369372e-5im,5.92725505878343e-7 - 1.6849504275866851e-6im
+-2.4498526812196763e-5 - 3.599520066273571e-5im,1.0830309677269584e-6 - 7.840192047983314e-7im
+4.040802043480934e-5 - 1.0440438532202839e-5im,1.5259694936117611e-6 - 2.9813055060130024e-7im
+-3.4567298186247468e-6 + 3.8973907876511844e-5im,1.1706058611675044e-6 + 6.532644385201991e-7im
+-3.301772063555424e-5 - 1.586665867388639e-5im,7.559541396406643e-7 + 8.993648589623673e-7im
+2.4857329699171084e-5 - 2.4262320397896137e-5im,4.4632631169618895e-8 + 1.0922853682765564e-6im
+1.3668544815306701e-5 + 2.9260010202598018e-5im,-3.679837159799295e-7 + 6.940056890841497e-7im
+-2.958900953015067e-5 + 2.394636860455206e-6im,-6.412653613717706e-7 + 3.260910188410887e-7im
+7.561284811477915e-6 - 2.676720068857872e-5im,-5.212488991399935e-7 - 1.9381303954865294e-7im
+2.1126270231746984e-5 + 1.4719885991648837e-5im,-2.7409383184762043e-7 - 4.293859046801776e-7im
+-1.9051457955820632e-5 + 1.3391509853712786e-5im,1.090089988355091e-7 - 5.424848199790569e-7im
+-5.3063371356379836e-6 - 2.0836191189026665e-5im,3.810332902626749e-7 - 3.7775501517518906e-7im
+1.9810926736088827e-5 + 1.6768332575367055e-6im,5.452526076116123e-7 - 1.596082834022398e-7im
+-7.301672790078927e-6 + 1.618903603350552e-5im,5.139351233582443e-7 + 1.1439886261265436e-7im
+-1.1291262775582465e-5 - 1.1465098339720254e-5im,3.821276626669417e-7 + 2.733806876199923e-7im
+1.3482243166308765e-5 - 6.312632972194135e-6im,1.8191322308420946e-7 + 3.415027111311585e-7im
+1.4734423044483348e-6 + 1.3133067566506715e-5im,2.0273628256559623e-8 + 2.75707014636713e-7im
+-1.1305064152146671e-5 - 2.9977663803252284e-6im,-7.816505269322091e-8 + 1.5900886377781664e-7im
+6.2141045890283095e-6 - 8.846346431842809e-6im,-8.061274111301497e-8 + 1.7562483151008862e-8im
+5.790230473392014e-6 + 7.684843602193653e-6im,-2.0938283941951207e-8 - 7.781699952217504e-8im
+-7.961804941047287e-6 + 2.3066381654847264e-6im,7.467048941975415e-8 - 1.20961345823369e-7im
+7.030511396608388e-7 - 7.574422685191263e-6im,1.585030014315096e-7 - 9.972458776802449e-8im
+6.331816227698637e-6 + 2.6475288082154953e-6im,2.1033077374226702e-7 - 4.559702284461422e-8im
+-3.881070536297868e-6 + 4.258924312518182e-6im,2.1651734846223726e-7 + 2.1229408696507515e-8im
+-2.1570715226770414e-6 - 4.707508611056819e-6im,1.8904917969596192e-7 + 7.020208685830888e-8im
+4.753359225467798e-6 - 5.92024857522855e-7im,1.4294885625386088e-7 + 9.293377572245067e-8im
+-7.163182525194776e-7 + 3.894073357797201e-6im,9.961934120526898e-8 + 8.55467718423489e-8im
+-2.7915015510985774e-6 - 1.951005727379913e-6im,7.102894696378749e-8 + 6.015471997942864e-8im
+2.6304896441581747e-6 - 1.9210186952573685e-6im,6.33282779547663e-8 + 2.8051791350097003e-8im
+9.964267781250871e-7 + 2.5230816738424564e-6im,7.235824260903279e-8 + 2.5977620168219707e-9im
+-2.164671403270642e-6 - 1.262940384300109e-7im,9.07325179762808e-8 - 1.0752331685299218e-8im
+9.28373949289574e-7 - 1.9275217227113023e-6im,1.0899136652186293e-7 - 1.0480881572942107e-8im
+1.4848759452084648e-6 + 1.1258393855753242e-6im,1.2101007004986137e-7 - 1.24149962799995e-9im
+-1.1662530565366675e-6 + 6.853280249240823e-7im,1.2393049703847118e-7 + 1.149410444532726e-8im
+-3.6659074782693596e-8 - 1.336972057474366e-6im,1.1913878194282454e-7 + 2.1941589455057344e-8im
+1.2693075675121334e-6 + 1.6984098185907974e-7im,1.0993550089186477e-7 + 2.7223469474149793e-8im
+-3.323512143971166e-7 + 7.889988340432508e-7im,1.0038154268247313e-7 + 2.6643943437760287e-8im
+-3.7013310237057106e-7 - 6.817354176215303e-7im,9.333801420701645e-8 + 2.1959948290265644e-8im
+8.222784584405401e-7 - 2.8468041424598943e-7im,9.011614308737097e-8 + 1.5594049037201118e-8im
+1.556028075675431e-7 + 5.605159269999485e-7im,9.028110827919117e-8 + 9.98568785032577e-9im
+-3.331211572679243e-7 - 2.1651706726769675e-7im,9.245747635788377e-8 + 6.490906026793147e-9im
+4.2677049987271346e-7 - 3.828627319988332e-7im,9.496262438938686e-8 + 5.4432394628531134e-9im
+3.388100314664092e-7 + 2.7355788873447543e-7im,9.65471360404498e-8 + 6.196008303811373e-9im
+-1.623143159081012e-7 + 1.67827782314132e-8im,9.662918418001644e-8 + 7.745056356448419e-9im
+1.8266251400419462e-7 - 3.074478169232698e-7im,9.533204356387336e-8 + 9.094800596345113e-9im
+3.393309386635946e-7 + 6.050830469076738e-8im,9.319679541484804e-8 + 9.652835252326558e-9im
+2.960950931417633e-10 + 8.139125458447397e-8im,9.089281827105408e-8 + 9.262109074640483e-9im
+7.752725632763076e-8 - 1.9170523282681732e-7im,8.893561198596376e-8 + 8.15192103838466e-9im
+2.690394657690494e-7 - 5.198926230100473e-8im,8.757130032488175e-8 + 6.715996781295215e-9im
+1.0221474010879384e-7 + 6.085949899840967e-8im,8.676965040347535e-8 + 5.344572823683608e-9im
+5.8044466461011985e-8 - 1.0116676861481462e-7im,8.632748413561294e-8 + 4.280787410459781e-9im
+1.955699272140396e-7 - 8.90447654104374e-8im,8.598811666219854e-8 + 3.5939344683726132e-9im
+1.4510993241231187e-7 + 1.669626224648967e-8im,8.554876573143181e-8 + 3.205112235038951e-9im
+7.416628622917299e-8 - 5.134671965288507e-8im,8.49092690392945e-8 + 2.9665098494533543e-9im
+1.4507719770118096e-7 - 8.664100149751614e-8im,8.407454869752818e-8 + 2.7301868134845162e-9im
+1.5067147816100532e-7 - 2.025111898873895e-8im,8.311712074702608e-8 + 2.401031474920625e-9im
+9.522659576066526e-8 - 3.32935255505667e-8im,8.213104762640301e-8 + 1.9493609662538094e-9im
+1.1874906339816006e-7 - 7.17280075876247e-8im,8.119230945368217e-8 + 1.4009093888638422e-9im
+1.4005219693764018e-7 - 4.167582000966016e-8im,8.033913210516452e-8 + 8.083363418242447e-10im
+1.0882097066971662e-7 - 3.241365242005272e-8im,7.957158922038466e-8 + 2.250141450622028e-10im
+1.0833315520289424e-7 - 5.8128289297449747e-8im,7.886429718179306e-8 - 3.1386508018300335e-10im
+1.2643777406018123e-7 - 5.036115772208361e-8im,7.818364919463715e-8 - 7.969895600231406e-10im
+1.1355070523178457e-7 - 3.7499215600040406e-8im,7.75020260359685e-8 - 1.2327182900136819e-9im
+1.05323566127433e-7 - 4.997794127177591e-8im,7.6804908883345e-8 - 1.6396094498815355e-9im
+1.1537283322795446e-7 - 5.189242386952749e-8im,7.60914749369035e-8 - 2.036640404720732e-9im
+1.122944481950227e-7 - 4.266353944149984e-8im,7.536978600190284e-8 - 2.4366181450533686e-9im
+1.0429130817506547e-7 - 4.659548678029505e-8im,7.46511300741129e-8 - 2.844183604253722e-9im
+1.0776609461891065e-7 - 5.0708649245647745e-8im,7.394482185185373e-8 - 3.2569523940162155e-9im
+1.0826632594626808e-7 - 4.599104360645485e-8im,7.325571193687809e-8 - 3.6690932416242657e-9im
+1.0285267479330202e-7 - 4.597458771519699e-8im,7.258411879885587e-8 - 4.0744080849809635e-9im
+1.0267188997921122e-7 - 4.920452309937143e-8im,7.192713647984362e-8 - 4.4687568818285745e-9im
+1.0356916577283471e-7 - 4.756973792951135e-8im,7.128089471307154e-8 - 4.8507640131244735e-9im
+1.0049420275807878e-7 - 4.6448030623326745e-8im,7.064201488523165e-8 - 5.221292183924957e-9im
+9.891417235648759e-8 - 4.821950398724807e-8im,7.000863119722855e-8 - 5.58251956430019e-9im
+9.915197335046962e-8 - 4.8081748288035995e-8im,6.938041377128397e-8 - 5.936630971649536e-9im
+9.747430348137241e-8 - 4.709233709888365e-8im,6.875800318487358e-8 - 6.285231524535307e-9im
+9.569341359308484e-8 - 4.7794386210428026e-8im,6.814250914018219e-8 - 6.629020156467615e-9im
+9.524369220794479e-8 - 4.8128234516344635e-8im,6.753477021211983e-8 - 6.967900946551634e-9im
+9.41804141410585e-8 - 4.7565493863300914e-8im,6.693523317018146e-8 - 7.301460027754976e-9im
+9.262839397552993e-8 - 4.771651855498055e-8im,6.634384556152491e-8 - 7.62913084649786e-9im
+9.176763624549627e-8 - 4.805363790318213e-8im,6.576019444752713e-8 - 7.95063111038808e-9im
+9.088325773723437e-8 - 4.7830005780028085e-8im,6.51838226871418e-8 - 8.265900719696771e-9im
+8.960501123096713e-8 - 4.7783429659782186e-8im,6.461424340303102e-8 - 8.575121743315119e-9im
+8.858573595310534e-8 - 4.799411190616609e-8im,6.40511937283912e-8 - 8.87862842927089e-9im
+8.77104658589886e-8 - 4.7950810620320644e-8im,6.349450242350593e-8 - 9.17670406218151e-9im
+8.662656111932163e-8 - 4.787606352961784e-8im,6.294412725589913e-8 - 9.469651710837955e-9im
+8.559412635333225e-8 - 4.7974151807084334e-8im,6.240007512035213e-8 - 9.757605469616618e-9im
+8.469610376122668e-8 - 4.799621284455464e-8im,6.186230883833067e-8 - 1.0040686566385029e-8im
+8.372435040831792e-8 - 4.7946376935559185e-8im,6.133079777385236e-8 - 1.0318943503021262e-8im
+8.273680939906199e-8 - 4.7978543930905586e-8im,6.080542785371636e-8 - 1.0592415048745781e-8im
+8.183231705829451e-8 - 4.800864596059935e-8im,6.028607765901158e-8 - 1.0861188806077231e-8im
+8.092181803954994e-8 - 4.7985315940650027e-8im,5.977262299155406e-8 - 1.1125320469274697e-8im
+7.99903637202136e-8 - 4.798815840319824e-8im,5.926490582085714e-8 - 1.138493241323321e-8im
+7.910059777284674e-8 - 4.800709645526361e-8im,5.8762848964042945e-8 - 1.1640133465415522e-8im
+7.822817869062405e-8 - 4.799827434621112e-8im,5.826631631743104e-8 - 1.1891023233877114e-8im
+7.734675702956818e-8 - 4.799112952419451e-8im,5.777524769770255e-8 - 1.2137732254204686e-8im
+7.64846755592741e-8 - 4.799726313884568e-8im,5.728955399687908e-8 - 1.238032349009123e-8im
+7.564274683668556e-8 - 4.7992616101070005e-8im,5.680913893062402e-8 - 1.2618909716395729e-8im
diff --git a/test/rescattered_ati.jl b/test/rescattered_ati.jl
new file mode 100644
index 0000000..bdef9e0
--- /dev/null
+++ b/test/rescattered_ati.jl
@@ -0,0 +1,93 @@
+# This rather unphysical combination of electric field and vector
+# potential (disregarding short-pulse effects, i.e. the chain rule) is
+# taken from Eqs. (21–22) of
+#
+# - D B Milo\vsevi\'c, Paulus, G. G., Bauer, D., & Becker,
+#   W. (2006). Above-Threshold Ionization By Few-Cycle Pulses. Journal
+#   of Physics B: Atomic, Molecular and Optical Physics, 39(14),
+#   203–262. http://dx.doi.org/10.1088/0953-4075/39/14/r01
+#
+# to be able to reproduce the results shown in their Fig. 14.
+function milosevic_pulse(F₀, ω, cycles, ndt, ϕ)
+    T = 2π/ω
+    steps = ceil(Int, cycles*ndt)
+    t = range(0, stop=cycles*T, length=steps)
+    Fv = zeros(steps)
+    Av = zeros(steps)
+
+    ωₚ = ω/cycles
+    ω₀ = ω
+    ω₁ = ω + ωₚ
+    ω₂ = ω - ωₚ
+
+    ϕ₁ = ϕ + π/2
+
+    ℰ₀ = F₀/2
+    ℰ₁ = -ℰ₀/2
+    ℰ₂ = -ℰ₀/2
+
+    𝒜₀ = ℰ₀/ω₀
+    𝒜₁ = ℰ₁/ω₁
+    𝒜₂ = ℰ₂/ω₂
+
+    for (i,t) in enumerate(t)
+        φ₀₁ = ω₀*t + ϕ₁
+        φ₁₁ = ω₁*t + ϕ₁
+        φ₂₁ = ω₂*t + ϕ₁
+
+        Fv[i] = ℰ₀*sin(φ₀₁) + ℰ₁*sin(φ₁₁) + ℰ₂*sin(φ₂₁)
+        Av[i] = 𝒜₀*cos(φ₀₁) + 𝒜₁*cos(φ₁₁) + 𝒜₂*cos(φ₂₁)
+    end
+
+    t, Fv, Av
+end
+
+@testset "Rescattered ATI" begin
+    kmax = 3.0
+    nk = 200
+    nθ = 2
+    spacing=[:momentum,:energy][2]
+    𝐤,kmag,θ = momentum_grid(0.1, kmax, nk, nθ,
+                            spacing=spacing)
+
+    @field(F) do
+        λ = 800.0u"nm"
+        I₀ = 1e14u"W/cm^2"
+        cycles = 4.0
+        ϕ = π
+        env = :cos²
+    end
+
+    ndt = ceil(Int, 3*(kmax^2/2)/photon_energy(F))
+
+    t, Fv, Av = milosevic_pulse(amplitude(F), photon_energy(F), F.env.cycles, ndt, 0)
+
+    Iₚ = 14u"eV" # "Krypton"
+
+    # Elastic scattering off a Yukawa potential
+    cc = StrongFieldApproximation.CoulombCoupling((𝐤,𝐩) -> yukawa_fourier(𝐩-𝐤, 1, 0, 1))
+    couplings=[[cc;;]]
+
+    ar = (Fv, Av, t)
+
+    channel = IonizationChannel(Iₚ, ar...)
+    system = StrongFieldApproximation.System([channel], nothing, couplings, ar...)
+
+    direct_diagram = Diagram(system)
+    rescattering_diagram = Diagram([(1,1),(1,0)], system)
+
+    cdirect_ref = readdlm(reffile("krypton-ref-direct.csv"), ',', ComplexF64)
+    crescattered_ref = readdlm(reffile("krypton-ref-rescattered.csv"), ',', ComplexF64)
+
+    cdirect = photoelectron_spectrum(𝐤, system, direct_diagram, verbosity=1)
+    test_approx_eq(cdirect, cdirect_ref)
+    cdirect2 = multi_channel_photoelectron_spectrum(𝐤, system, direct_diagram, verbosity=1)
+    @test size(cdirect2, 1) == 1
+    test_approx_eq(cdirect2[1,:,:], cdirect_ref)
+
+    crescattered = photoelectron_spectrum(𝐤, system, rescattering_diagram, verbosity=1)
+    test_approx_eq(crescattered, crescattered_ref)
+    crescattered2 = multi_channel_photoelectron_spectrum(𝐤, system, rescattering_diagram, verbosity=1)
+    @test size(crescattered2, 1) == 1
+    test_approx_eq(crescattered2[1,:,:], crescattered_ref)
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 72ccc0b..cfb2cce 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -2,116 +2,33 @@ using StrongFieldApproximation
 using ElectricFields
 using Unitful
 using UnitfulAtomic
+using DelimitedFiles
 using Test
 
-# write your own tests here
-@test 1 == 1
-@testset "Telescope iterators" begin
-    generate_reference_common(ref) =
-        map(I -> [Tuple(I)...], findall(ref))
+reffile(name) = joinpath(dirname(@__FILE__), "references", name)
 
-    generate_1d_reference(N, memory) = [[i] for i in 1:N]
+function test_approx_eq(a, b; on_fail::Union{Nothing,Function}=nothing, kwargs...)
+    size(a) == size(b) || throw(DimensionMismatch("Cannot compare objects of sizes $(size(a)) and $(size(b))"))
+    if !isapprox(a, b; kwargs...)
+        @error "Approximate equality failed:"
+        na = norm(a)
+        nb = norm(b)
+        Δ = norm(a-b)
+        relΔ = Δ/max(na,nb)
 
-    function generate_2d_reference(N, memory)
-        ref = [0 < i - j ≤ memory
-               for i in 1:N, j in 1:N]
-        generate_reference_common(ref)
-    end
-
-    function generate_3d_reference(N, memory)
-        ref = [i > j > k && 0 < i - j ≤ memory && 0 < j - k ≤ memory
-               for i in 1:N, j in 1:N, k in 1:N]
-        generate_reference_common(ref)
-    end
+        println("   |a|", na)
+        println("   |b|", nb)
+        println("Abs. Δ", Δ)
+        println("Rel. Δ", relΔ)
 
-    function generate_4d_reference(N, memory)
-        ref = [i > j > k > l && 0 < i - j ≤ memory && 0 < j - k ≤ memory && 0 < k - l ≤ memory
-               for i in 1:N, j in 1:N, k in 1:N, l in 1:N]
-        generate_reference_common(ref)
+        isnothing(on_fail) || on_fail()
     end
 
-    function test_telescope_iterator(N, n, memory)
-        v = 1:N
-        ti = StrongFieldApproximation.TelescopeIterator(v, n, memory)
-        # At the moment we don't know how to compute length(ti), so
-        # collect(ti) does not work.
-        data = Vector{Int}[]
-        for e in ti
-            push!(data, e)
-        end
-        ref = if n == 1
-            generate_1d_reference(N, memory)
-        elseif n == 2
-            generate_2d_reference(N, memory)
-        elseif n == 3
-            generate_3d_reference(N, memory)
-        elseif n == 4
-            generate_4d_reference(N, memory)
-        end
-
-        @test data == ref
-    end
-    @testset "N = $N, n = $n, memory = $memory" for N = 1:7, n = 1:4, memory = [0,1,3,100]
-        test_telescope_iterator(N, n, memory)
-    end
+    @test isapprox(a, b; kwargs...)
 end
 
-@testset "Diagrams" begin
-    function display_diagram(system, diagram, expected_momenta, expected_unique, expected_indeterminate)
-        display(diagram)
-        ions, unique_momenta, momenta, indeterminate_momenta, order = StrongFieldApproximation.analyze_diagram(system, diagram)
-
-        expected_ions = first.(diagram.path)
-        if first(diagram.path)[2] == 0 && length(diagram) > 1
-            expected_ions = expected_ions[2:end]
-        end
-
-        @test ions == expected_ions
-        @test momenta == expected_momenta
-        @test unique_momenta == expected_unique
-        @test indeterminate_momenta == expected_indeterminate
-    end
-
-    @field(F) do
-        λ = 800.0u"nm"
-        I₀ = 1e14u"W/cm^2"
-        cycles = 4.0
-        ϕ = π
-        env = :cos²
-    end
-
-    ndt = 238
-
-    Iₚ = 14u"eV" # "Krypton"
-
-    # Elastic scattering off a Yukawa potential
-    cc = StrongFieldApproximation.CoulombCoupling((𝐤,𝐩) -> yukawa_fourier(𝐩-𝐤, 1, 0, 1))
-    dc = StrongFieldApproximation.DipoleCoupling(0.1, F)
-    couplings=Matrix{StrongFieldApproximation.AbstractCoupling}[reshape([dc],1,1),reshape([cc],1,1)]
-
-    ar = (F, ndt)
-    channel = IonizationChannel(Iₚ, ar...)
-    system = StrongFieldApproximation.System(repeat([channel], 10), nothing, couplings, ar...)
-
-    for (path, expected_momenta, expected_unique, expected_indeterminate) in [
-        ([(1,0)], [1], [(1,2)], []),
-        ([(1,0),(1,0)], [1], [(1,2)], 1:1),
-        ([(2,1),(1,0)], [1,1], [(1,3)], []),
-        ([(2,2),(1,0)], [1,2], [(1,2),(2,3)], 2:2),
-        ([(1,2),(2,2),(1,0)], [1,2,3], [(1,2),(2,3),(3,4)], 2:3),
-        ([(1,0),(1,2),(2,2),(1,0)], [1,2,3], [(1,2),(2,3),(3,4)], 1:3),
-        ([(1,1),(2,2),(1,0)], [1,1,2], [(1,3),(3,4)], 2:2),
-        ([(1,0),(1,1),(2,2),(1,0)], [1,1,2], [(1,3),(3,4)], 1:2),
-        ([(3,1),(1,1),(2,2),(1,0)], [1,1,1,2], [(1,4),(4,5)], 2:2),
-        ([(3,0),(3,1),(1,1),(2,2),(1,0)], [1,1,1,2], [(1,4),(4,5)], 1:2),
-        ([(1,2),(3,1),(1,1),(2,2),(1,0)], [1,2,2,2,3], [(1,2),(2,5),(5,6)], 2:3),
-        ([(1,2),(1,2),(3,1),(2,2),(1,0)], [1,2,3,3,4], [(1,2),(2,3),(3,5),(5,6)], 2:4),
-        ([(1,2),(1,2),(3,1),(1,1),(2,2),(1,0)], [1,2,3,3,3,4], [(1,2),(2,3),(3,6),(6,7)], 2:4),
-        ([(1,2),(10,2),(3,2),(1,2),(2,2),(1,0)], [1,2,3,4,5,6], [(1,2),(2,3),(3,4),(4,5),(5,6),(6,7)], 2:6),
-        ([(1,0),(1,2),(10,2),(3,2),(1,2),(2,2),(1,0)], [1,2,3,4,5,6], [(1,2),(2,3),(3,4),(4,5),(5,6),(6,7)], 1:6),
-        ([(1,0),(1,1),(10,1),(3,1),(1,1),(2,1),(1,0)], [1,1,1,1,1,1], [(1,7)], 1:1)
-    ]
-        diagram = Diagram(path, system)
-        display_diagram(system, diagram, expected_momenta, expected_unique, expected_indeterminate)
-    end
+@testset "StrongFieldApproximation.jl" begin
+    include("telescope_iterators.jl")
+    include("diagrams.jl")
+    include("rescattered_ati.jl")
 end
diff --git a/test/telescope_iterators.jl b/test/telescope_iterators.jl
new file mode 100644
index 0000000..8204e99
--- /dev/null
+++ b/test/telescope_iterators.jl
@@ -0,0 +1,49 @@
+@testset "Telescope iterators" begin
+    generate_reference_common(ref) =
+        map(I -> [Tuple(I)...], findall(ref))
+
+    generate_1d_reference(N, memory) = [[i] for i in 1:N]
+
+    function generate_2d_reference(N, memory)
+        ref = [0 < i - j ≤ memory
+               for i in 1:N, j in 1:N]
+        generate_reference_common(ref)
+    end
+
+    function generate_3d_reference(N, memory)
+        ref = [i > j > k && 0 < i - j ≤ memory && 0 < j - k ≤ memory
+               for i in 1:N, j in 1:N, k in 1:N]
+        generate_reference_common(ref)
+    end
+
+    function generate_4d_reference(N, memory)
+        ref = [i > j > k > l && 0 < i - j ≤ memory && 0 < j - k ≤ memory && 0 < k - l ≤ memory
+               for i in 1:N, j in 1:N, k in 1:N, l in 1:N]
+        generate_reference_common(ref)
+    end
+
+    function test_telescope_iterator(N, n, memory)
+        v = 1:N
+        ti = StrongFieldApproximation.TelescopeIterator(v, n, memory)
+        # At the moment we don't know how to compute length(ti), so
+        # collect(ti) does not work.
+        data = Vector{Int}[]
+        for e in ti
+            push!(data, e)
+        end
+        ref = if n == 1
+            generate_1d_reference(N, memory)
+        elseif n == 2
+            generate_2d_reference(N, memory)
+        elseif n == 3
+            generate_3d_reference(N, memory)
+        elseif n == 4
+            generate_4d_reference(N, memory)
+        end
+
+        @test data == ref
+    end
+    @testset "N = $N, n = $n, memory = $memory" for N = 1:7, n = 1:4, memory = [0,1,3,100]
+        test_telescope_iterator(N, n, memory)
+    end
+end

From bb50bede3cdcac40e859907632921dde2c011980 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Tue, 18 Apr 2023 11:21:42 +0200
Subject: [PATCH 17/20] Fixed threading issue

---
 src/dyson_expansions.jl  | 2 +-
 src/multi_channel_sfa.jl | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/dyson_expansions.jl b/src/dyson_expansions.jl
index 0bf2256..6003e76 100644
--- a/src/dyson_expansions.jl
+++ b/src/dyson_expansions.jl
@@ -56,7 +56,7 @@ function integrate_diagram(::Type{Amp}, system::System, diagram::Diagram, iref,
 
     @timeit to "Time loop" begin
         for i in TelescopeIterator(max(1,imin):iref-1, order, memory)
-            windw = @timeit "Window" prod(window[(i[1] + 1) .- i])
+            windw = @timeit to "Window" prod(window[(i[1] + 1) .- i])
             iszero(windw) && continue
 
             is[2:end] .= i
diff --git a/src/multi_channel_sfa.jl b/src/multi_channel_sfa.jl
index 5e5dea0..bed0d8d 100644
--- a/src/multi_channel_sfa.jl
+++ b/src/multi_channel_sfa.jl
@@ -49,7 +49,7 @@ function multi_channel_sfa!(amplitudes, system, diagram, iref, 𝐤=nothing;
 
     @timeit to "Time loop" begin
         for i in TelescopeIterator(max(1,imin):iref-1, order, memory)
-            windw = @timeit "Window" prod(window[(i[1] + 1) .- i])
+            windw = @timeit to "Window" prod(window[(i[1] + 1) .- i])
             iszero(windw) && continue
 
             is[2:end] .= i

From a30aa55add352c1f5fc5b8c2772c29ec490942a4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Tue, 18 Apr 2023 11:30:43 +0200
Subject: [PATCH 18/20] Fixed matrix syntax for Julia < 1.7

---
 test/rescattered_ati.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/rescattered_ati.jl b/test/rescattered_ati.jl
index bdef9e0..10e575e 100644
--- a/test/rescattered_ati.jl
+++ b/test/rescattered_ati.jl
@@ -66,7 +66,7 @@ end
 
     # Elastic scattering off a Yukawa potential
     cc = StrongFieldApproximation.CoulombCoupling((𝐤,𝐩) -> yukawa_fourier(𝐩-𝐤, 1, 0, 1))
-    couplings=[[cc;;]]
+    couplings=[reshape([cc], (1,1))]
 
     ar = (Fv, Av, t)
 

From 20580b0a2b8f19e7bce3142152c64561bd68bedb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Tue, 18 Apr 2023 12:43:19 +0200
Subject: [PATCH 19/20] Fixed call signature

---
 src/ion_propagators.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/ion_propagators.jl b/src/ion_propagators.jl
index b22ce12..52ae42f 100644
--- a/src/ion_propagators.jl
+++ b/src/ion_propagators.jl
@@ -139,7 +139,7 @@ Base.eltype(::LaserDressedIons{T}) where T = T
 ion_mapping(ions::LaserDressedIons, i) =
     view(ions.ion_basis, :, :, i)
 
-ion_mapping(ions::LaserDressedIons, _) = I
+ion_mapping(ions::LaserDressedIons{<:Any,Nothing}, _) = I
 
 ion_mapping(ions::LaserDressedIons, α, i) =
     view(ions.ion_basis, α, :, i)

From ac72ba18289bbcfe807ad689928c076694915e0b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefanos=20Carlstr=C3=B6m?= <stefanos.carlstrom@gmail.com>
Date: Wed, 19 Apr 2023 08:48:23 +0200
Subject: [PATCH 20/20] Support for non-linear polarization

---
 src/channels.jl        |  2 +-
 src/ion_propagators.jl | 24 +++++++++++++++++-------
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/src/channels.jl b/src/channels.jl
index b39ed4d..887e0fa 100644
--- a/src/channels.jl
+++ b/src/channels.jl
@@ -33,7 +33,7 @@ function Base.show(io::IO, mime::MIME"text/plain", d::DipoleSourceTerm)
 end
 
 dipole(F::Number, d::Number) = F*d
-dipole(F::SVector{3}, d::SVector{3}) = dot(F, d)
+dipole(F::SVector{3}, d::SVector{3}) = F[1]*d[1] + F[2]*d[2] + F[3]*d[3]
 dipole(F::Number, d::SVector{3}) = F*d[3]
 
 _field_amplitude(F::ElectricFields.AbstractField, t) =
diff --git a/src/ion_propagators.jl b/src/ion_propagators.jl
index 52ae42f..d4db7b7 100644
--- a/src/ion_propagators.jl
+++ b/src/ion_propagators.jl
@@ -65,21 +65,31 @@ function LaserDressedIons(E::Matrix{T}, Q, t) where T
     LaserDressedIons(expimE, Q)
 end
 
-function LaserDressedIons(ics, 𝐫ᵢₒₙ::SparseMatrixCSC, Fv::AbstractArray, t)
+dme_subset(A::AbstractMatrix, sel) = A[sel,sel]
+dme_subset(A::SVector{3}, sel) =
+    SVector{3}(dme_subset(A[i], sel) for i = 1:3)
+dme_subset(I::UniformScaling, _) = I
+
+incidence_matrix(A::AbstractMatrix) = A .≠ 0
+incidence_matrix(A::SVector{3,<:AbstractMatrix}) =
+    incidence_matrix(A[1]) .|
+    incidence_matrix(A[2]) .|
+    incidence_matrix(A[3])
+
+function LaserDressedIons(ics, 𝐫ᵢₒₙ::Union{SparseMatrixCSC,SVector{3,<:SparseMatrixCSC}}, Fv::AbstractArray, t)
     to = TimerOutput()
 
     @timeit to "Allocations" begin
-        m,n = size(𝐫ᵢₒₙ)
-        @assert m == n == length(ics)
+        m = length(ics)
 
         nt = length(t)
         E₀ = [ic.E for ic in ics]
         H₀ = Diagonal(E₀)
 
-        X = .!iszero.(I + (𝐫ᵢₒₙ .≠ 0))
+        X = incidence_matrix(I + incidence_matrix(𝐫ᵢₒₙ))
         bs = find_blocks(X)
 
-        @info "Diagonalizing $(m)×$(n) ionic Hamiltonian for $(nt) times" bs
+        @info "Diagonalizing $(m)×$(m) ionic Hamiltonian for $(nt) times" bs
 
         T = eltype(E₀)
         E = zeros(T, nt, m)
@@ -97,10 +107,10 @@ function LaserDressedIons(ics, 𝐫ᵢₒₙ::SparseMatrixCSC, Fv::AbstractArray
 
         Hsub = zeros(T, nb, nb)
         H₀sub = Diagonal(E₀[b])
-        zsub = Matrix(𝐫ᵢₒₙ[b,b])
+        dsub = dme_subset(𝐫ᵢₒₙ, b)
         @timeit to "Time loop" begin
             @showprogress for i = 1:nt
-                Hsub .= H₀sub .+ Fv[i] .* zsub
+                Hsub .= H₀sub .+ dipole(Fv[i,:], dsub)
                 @assert Hsub ≈ Hsub'
                 ee = eigen!(Hermitian(Hsub))
                 for (ij,j) in enumerate(b)