Added "automatic integration" via Gauss quadrature

Author: jagot

Project.toml

@@ -7,6 +7,7 @@ version = "0.1.0"
 ElectricFields = "2f84ce32-9bb1-11e8-0d9f-3dce90a4beca"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 HCubature = "19dc6840-f33b-545b-b366-655c7e3ffd49"
+IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"

src/StrongFieldApproximation.jl

@@ -3,7 +3,9 @@ module StrongFieldApproximation
 using ElectricFields
 using Unitful
 using UnitfulAtomic
+
 using FastGaussQuadrature
+using IntervalSets
 
 using LinearAlgebra
 using StaticArrays
@@ -18,6 +20,7 @@ include("momentum_grid.jl")
 
 include("units.jl")
 include("field.jl")
+include("automatic_integration.jl")
 include("volkov.jl")
 
 include("channels.jl")

src/automatic_integration.jl

@@ -0,0 +1,123 @@
+struct ComplexPlane{T,R<:AbstractRange{T}} <: AbstractMatrix{Complex{T}}
+    re::R
+    im::R
+end
+
+Base.length(p::ComplexPlane) = length(p.re)*length(p.im)
+Base.size(p::ComplexPlane) = (length(p.im), length(p.re))
+Base.eltype(p::ComplexPlane) = complex(eltype(p.re))
+Base.getindex(p::ComplexPlane, i::CartesianIndex) = p.re[i[2]] + im*p.im[i[1]]
+
+function primitive_integral(f, a, b, c, w)
+    μ = (a+b)/2
+    δ = (b-a)/2
+    v = zero(f(a))
+    for (c,w) in zip(c,w)
+        v += w*f(δ*c + μ)
+    end
+    δ*v
+end
+
+interval(x::AbstractRange, r::UnitRange) = x[r[1]]..x[r[end]]
+
+function split_range(r::UnitRange)
+    n = length(r)
+    n < 1 && error("Too small range")
+    n == 1 && return (r, r[1]:(r[1]-1))
+
+    s = n ÷ 2
+    r[1]:r[s], r[s+1]:r[end]
+end
+
+function find_closest(x::AbstractRange, x₀::Real)
+    cur = firstindex(x):lastindex(x)
+    i = 0
+    while length(cur) > 1 && i < ceil(Int, log2(length(x)))
+        left,right = split_range(cur)
+        if x₀ ∈ interval(x, left)
+            cur = left
+            continue
+        elseif x₀ ∈ interval(x, right)
+            cur = right
+            continue
+        end
+        x₀ < x[left[1]] && return left[1]
+        x₀ > x[right[end]] && return right[end]
+        return norm(x[left[end]]-x₀) < norm(x[right[1]]-x₀) ? left[end] : right[1]
+        i += 1
+    end
+    cur[1]
+end
+
+function find_closest(x::ComplexPlane, x₀::Number)
+    j = find_closest(x.re, real(x₀))
+    i = find_closest(x.im, imag(x₀))
+    CartesianIndex((i,j))
+end
+
+struct AutomaticIntegration{Func, Grid, Xt, GridIntegrals,
+                            Roots, Weights}
+    f::Func
+    x::Grid
+    x₀::Xt
+    ∫f::GridIntegrals
+    c::Roots
+    w::Weights
+end
+
+function integrate!(ai::AutomaticIntegration{<:Any,<:AbstractRange}, i)
+    a,b = firstindex(ai.x),lastindex(ai.x)
+    for j = i+1:b
+        ai.∫f[j] = ai(ai.x[j],i=j-1)
+    end
+    for j = i-1:-1:a
+        ai.∫f[j] = ai(ai.x[j],i=j+1)
+    end
+end
+
+function integrate!(ai::AutomaticIntegration{<:Any,<:ComplexPlane}, i)
+    s = size(ai.x)
+    ci = CartesianIndices(s)
+
+    north(I::CartesianIndex) = CartesianIndex(I[1]-1,I[2])
+    south(I::CartesianIndex) = CartesianIndex(I[1]+1,I[2])
+    west(I::CartesianIndex) = CartesianIndex(I[1],I[2]-1)
+    east(I::CartesianIndex) = CartesianIndex(I[1],I[2]+1)
+
+    # First, we integrate parallel to the real axis, starting from
+    # index i, then we integrate parallel to the imaginary axis,
+    # starting from the real line-parallel we just created.
+    for (r,dir) in [(reverse(ci[i[1],1:i[2]-1]), east),
+                    (ci[i[1],i[2]+1:end], west),
+                    (ci[i[1]+1:end,:], north),
+                    (ci[i[1]-1:-1:1,:], south)]
+        for I in r
+            # @show I, dir(I) ai.∫f[I], ai.∫f[dir(I)]
+            ai.∫f[I] = ai(ai.x[I],i=dir(I))
+        end
+    end
+end
+
+function AutomaticIntegration(f, x, x₀; k=3, init=zero(f(x₀)))
+    Xt = eltype(x)
+    s = size(x)
+    Ft = typeof(f(Xt(x₀)))
+    ∫f = zeros(Ft, s)
+
+    c,w = gausslegendre(k)
+    ai = AutomaticIntegration(f, x, x₀, ∫f, c, w)
+
+    # We start by integrating from the point on the grid closest to x₀
+    # to x₀, which is minus the integral value desired, since the
+    # integral originates at x₀.
+    i = find_closest(x, x₀)
+    ∫f[i] = init - ai(x₀, i=i)
+    integrate!(ai, i)
+
+    ai
+end
+
+(ai::AutomaticIntegration)(x; i=find_closest(ai.x, x)) =
+    ai.∫f[i] + primitive_integral(ai.f, ai.x[i], x, ai.c, ai.w)
+
+(ai::AutomaticIntegration)(a, b) = ai(b) - ai(a)

test/runtests.jl

@@ -113,3 +113,22 @@ end
         display_diagram(system, diagram, expected_momenta, expected_unique, expected_indeterminate)
     end
 end
+
+@testset "Automatic integration" begin
+    import StrongFieldApproximation: AutomaticIntegration, ComplexPlane
+    fine_range(t, η=10) = range(t[1], stop=t[end], length=η*length(t))
+
+    tre = range(-1, stop=1, length=201)*4π
+    tim = range(-1, stop=1, length=101)*2π
+    t = ComplexPlane(tre, tim)
+    tfine = ComplexPlane(fine_range(t.re), fine_range(t.im))
+    @testset "$(f)" for (f,∫f) in [(one,identity), (identity,z->z^2/2),
+                                   (z->z^2, z->z^3/3),
+                                   (sin, z->1-cos(z))]
+        @testset "$(name)" for (name,t,tfine) in [("Real line", tre, tfine.re),
+                                                  ("Complex plane", t, tfine)]
+            ai = AutomaticIntegration(f, t, 0)
+            @test ai.(tfine) ≈ ∫f.(tfine)
+        end
+    end
+end