More statistics functions

src/Luna.jl

@@ -205,7 +205,7 @@ function setup(grid::Grid.RealGrid, xygrid::Grid.FreeGrid,
     end
     xo = Array{Float64}(undef, length(grid.to), length(y), length(x))
     FTo = FFTW.plan_rfft(xo, (1, 2, 3), flags=settings["fftw_flag"])
-    transform = NonlinearRHS.TransFree(grid, FTo, length(y), length(x),
+    transform = NonlinearRHS.TransFree(grid, xygrid, FTo,
                                        responses, densityfun, normfun)
     inv(FT) # create inverse FT plans now, so wisdom is saved
     inv(FTo)
@@ -226,7 +226,7 @@ function setup(grid::Grid.EnvGrid, xygrid::Grid.FreeGrid,
     end
     xo = Array{ComplexF64}(undef, length(grid.to), length(y), length(x))
     FTo = FFTW.plan_fft(xo, (1, 2, 3), flags=settings["fftw_flag"])
-    transform = NonlinearRHS.TransFree(grid, FTo, length(y), length(x),
+    transform = NonlinearRHS.TransFree(grid, xygrid, FTo,
                                        responses, densityfun, normfun)
     inv(FT) # create inverse FT plans now, so wisdom is saved
     inv(FTo)
@@ -274,10 +274,6 @@ function shotnoise!(Eω, grid::Grid.EnvGrid; seed=nothing)
     @. Eω += FFTamp * exp(1im*φ)
 end
 
-transtype(t::NonlinearRHS.TransModeAvg) = "mode average"
-transtype(t::NonlinearRHS.TransModal) = "multimode"
-transtype(t) = "unknown"
-
 linoptype(l::AbstractArray) = "constant"
 linoptype(l) = "variable"
 
@@ -286,14 +282,13 @@ gridtype(g::Grid.EnvGrid) = "envelope"
 gridtype(g) = "unknown"
 
 simtype(g, t, l) = Dict("field" => gridtype(g),
-                        "transform" => transtype(t),
+                        "transform" => string(t),
                         "linop" => linoptype(l))
 
 function run(Eω, grid,
              linop, transform, FT, output;
              min_dz=0, max_dz=Inf, init_dz=1e-4,
              rtol=1e-6, atol=1e-10, safety=0.9, norm=RK45.weaknorm,
-             stats=true,
              status_period=1)
 
 

src/Maths.jl

@@ -90,37 +90,42 @@ function fwhm(x, y; method=:linear, baseline=false, minmax=:min)
     end
     maxidx = argmax(y)
     xmax = x[maxidx]
-    left, right = try
-        if minmax == :min
-            lefti = findlast((x .< xmax) .& (y .< val))
-            righti = findfirst((x .> xmax) .& (y .< val))
-            (method == :nearest) && return abs(x[lefti] - x[righti])
-            left = linterpx(x[lefti], x[lefti+1], y[lefti], y[lefti+1], val)
-            right = linterpx(x[righti-1], x[righti], y[righti-1], y[righti], val)
-        else
-            lefti = findfirst((x .< xmax) .& (y .> val))
-            righti = findlast((x .> xmax) .& (y .> val))
-            (method == :nearest) && return abs(x[lefti] - x[righti])
-            left = linterpx(x[lefti-1], x[lefti], y[lefti-1], y[lefti], val)
-            right = linterpx(x[righti], x[righti+1], y[righti], y[righti+1], val)
+    if method in (:nearest, :linear)
+        try
+            if minmax == :min
+                lefti = findlast((x .< xmax) .& (y .< val))
+                righti = findfirst((x .> xmax) .& (y .< val))
+                (method == :nearest) && return abs(x[lefti] - x[righti])
+                left = linterpx(x[lefti], x[lefti+1], y[lefti], y[lefti+1], val)
+                right = linterpx(x[righti-1], x[righti], y[righti-1], y[righti], val)
+            else
+                lefti = findfirst((x .< xmax) .& (y .> val))
+                righti = findlast((x .> xmax) .& (y .> val))
+                (method == :nearest) && return abs(x[lefti] - x[righti])
+                left = linterpx(x[lefti-1], x[lefti], y[lefti-1], y[lefti], val)
+                right = linterpx(x[righti], x[righti+1], y[righti], y[righti+1], val)
+            end
+            return abs(right - left)
+        catch
+            return NaN
         end
-        (method == :linear) && return abs(right - left)
-        left, right
-    catch
-        return NaN
-    end
-    #spline method
-    minmax == :max && error("Spline method only supports min FWHM")
-    try
-        spl = CSpline(x, y)
-        lb = xmax - 2*(xmax-left)
-        ub = xmax + 2*(right-xmax)
-        f(x) = spl(x) - val
-        lfine = fzero(f, lb, xmax)
-        rfine = fzero(f, xmax, ub)
-        return abs(rfine - lfine)
-    catch
-        return NaN
+    elseif method == :spline
+        #spline method
+        try
+            spl = BSpline(x, y .- val)
+            r = roots(spl)
+            rleft = r[r .< xmax]
+            rright = r[r .> xmax]
+            if minmax == :min
+                return abs(minimum(rright) - maximum(rleft))
+            else
+                return abs(maximum(rright) - minimum(rleft))
+            end
+        catch
+            return NaN
+        end
+    else
+        error("Unknown FWHM method $method")
     end
 end
 

src/NonlinearRHS.jl

@@ -172,11 +172,13 @@ end
 
 function show(io::IO, t::TransModal)
     grid = "grid type: $(typeof(t.grid))"
-    nmodes = "no. of modes: $(t.nmodes)"
+    nmodes = "no. of modes: $(t.ts.nmodes)"
+    p = t.ts.indices == 1:2 ? "x,y" : t.ts.indices == 1 ? "x" : "y"
+    pol = "polarisation: $p"
     samples = "time grid size: $(length(t.grid.t)) / $(length(t.grid.to))"
     resp = "responses: "*join([string(typeof(ri)) for ri in t.resp], "\n    ")
     full = "full: $(t.full)"
-    out = join(["TransModal", nmodes, grid, samples, full, resp], "\n  ")
+    out = join(["TransModal", nmodes, pol, grid, samples, full, resp], "\n  ")
     print(io, out)
 end
 
@@ -293,7 +295,7 @@ function show(io::IO, t::TransModeAvg)
     grid = "grid type: $(typeof(t.grid))"
     samples = "time grid size: $(length(t.grid.t)) / $(length(t.grid.to))"
     resp = "responses: "*join([string(typeof(ri)) for ri in t.resp], "\n    ")
-    out = join(["TransModal", grid, samples, resp], "\n  ")
+    out = join(["TransModeAvg", grid, samples, resp], "\n  ")
     print(io, out)
 end
 
@@ -371,6 +373,16 @@ struct TransRadial{TT, HTT, FTT, nT, rT, gT, dT, iT}
     idcs::iT # CartesianIndices for Et_to_Pt! to iterate over
 end
 
+function show(io::IO, t::TransRadial)
+    grid = "grid type: $(typeof(t.grid))"
+    samples = "time grid size: $(length(t.grid.t)) / $(length(t.grid.to))"
+    resp = "responses: "*join([string(typeof(ri)) for ri in t.resp], "\n    ")
+    nr = "radial points: $(t.QDHT.N)"
+    R = "aperture: $(t.QDHT.R)"
+    out = join(["TransRadial", grid, samples, nr, R, resp], "\n  ")
+    print(io, out)
+end
+
 function TransRadial(TT, grid, HT, FT, responses, densityfun, normfun)
     Eωo = zeros(ComplexF64, (length(grid.ωo), HT.N))
     Eto = zeros(TT, (length(grid.to), HT.N))
@@ -504,11 +516,12 @@ end
 
 Transform E(ω) -> Pₙₗ(ω) for full 3D free-space propagation
 """
-mutable struct TransFree{TT, FTT, nT, rT, gT, dT, iT}
+mutable struct TransFree{TT, FTT, nT, rT, gT, xygT, dT, iT}
     FT::FTT # 3D Fourier transform (space to k-space and time to frequency)
     normfun::nT # Function which returns normalisation factor
     resp::rT # nonlinear responses (tuple of callables)
     grid::gT # time grid
+    xygrid::xygT
     densityfun::dT # callable which returns density
     Pto::Array{TT, 3} # buffer for oversampled time-domain NL polarisation
     Eto::Array{TT, 3} # buffer for oversampled time-domain field
@@ -518,13 +531,26 @@ mutable struct TransFree{TT, FTT, nT, rT, gT, dT, iT}
     idcs::iT # iterating over these slices Eto/Pto into Vectors, one at each position
 end
 
-function TransFree(TT, scale, grid, FT, Ny, Nx, responses, densityfun, normfun)
+function show(io::IO, t::TransFree)
+    grid = "grid type: $(typeof(t.grid))"
+    samples = "time grid size: $(length(t.grid.t)) / $(length(t.grid.to))"
+    resp = "responses: "*join([string(typeof(ri)) for ri in t.resp], "\n    ")
+    y = "y grid: $(minimum(t.xygrid.y)) to $(maximum(t.xygrid.y)), N=$(length(t.xygrid.y))"
+    x = "x grid: $(minimum(t.xygrid.x)) to $(maximum(t.xygrid.x)), N=$(length(t.xygrid.x))"
+    out = join(["TransFree", grid, samples, y, x, resp], "\n  ")
+    print(io, out)
+end
+
+function TransFree(TT, scale, grid, xygrid, FT, responses, densityfun, normfun)
+    Ny = length(xygrid.y)
+    Nx = length(xygrid.x)
     Eωo = zeros(ComplexF64, (length(grid.ωo), Ny, Nx))
     Eto = zeros(TT, (length(grid.to), Ny, Nx))
     Pto = similar(Eto)
     Pωo = similar(Eωo)
     idcs = CartesianIndices((Ny, Nx))
-    TransFree(FT, normfun, responses, grid, densityfun, Pto, Eto, Eωo, Pωo, scale, idcs)
+    TransFree(FT, normfun, responses, grid, xygrid, densityfun,
+              Pto, Eto, Eωo, Pωo, scale, idcs)
 end
 
 """

src/Stats.jl

@@ -173,7 +173,7 @@ function electrondensity(grid::Grid.RealGrid, ionrate!, dfun,
     end
     tospace = Modes.ToSpace(modes, components=components)
     frac = similar(to)
-    npol = components == :xy ? 2 : 1  
+    npol = tospace.npol
     Et0 = zeros(ComplexF64, (length(to), npol))
     function addstat!(d, Eω, Et, z, dz)
         # note: oversampling returns its arguments without any work done if factor==1
@@ -194,6 +194,18 @@ function density(dfun)
     end
 end
 
+function core_radius(a::Number)
+    function addstat!(d, Eω, Et, z, dz)
+        d["core_radius"] = a
+    end
+end
+
+function core_radius(afun)
+    function addstat!(d, Eω, Et, z, dz)
+        d["core_radius"] = afun(z)
+    end
+end
+
 function zdz!(d, Eω, Et, z, dz)
     d["z"] = z
     d["dz"] = dz

test/runtests.jl

@@ -29,7 +29,7 @@ end
 end
 
 @testset "ODE Solver" begin
-    @info("================= test_rect_rk45.jl")
+    @info("================= test_rk45.jl")
     include(joinpath(testdir, "test_rk45.jl"))
 end
 

test/test_freespace.jl

@@ -46,8 +46,9 @@ function gausspulse(t)
     return prop(Et, -0.3)
 end
 
-dens0 = PhysData.density(gas, pres)
-densityfun(z) = dens0
+densityfun = let dens0=PhysData.density(gas, pres)
+    z -> dens0
+end
 
 ionpot = PhysData.ionisation_potential(gas)
 ionrate = Ionisation.ionrate_fun!_PPTcached(gas, λ0)

test/test_freespace_fourier.jl

@@ -34,8 +34,9 @@ function gausspulse(t)
     Et = @. sqrt(It)*cos(ω0*t)
 end
 
-dens0 = PhysData.density(gas, pres)
-densityfun(z) = dens0
+densityfun = let dens0=PhysData.density(gas, pres)
+    z -> dens0
+end
 
 ionpot = PhysData.ionisation_potential(gas)
 ionrate = Ionisation.ionrate_fun!_ADK(ionpot)