More statistics functions

src/Capillary.jl

@@ -9,6 +9,7 @@ import Luna: Maths
 import Luna.PhysData: c, ε_0, μ_0, ref_index_fun, roomtemp, densityspline, sellmeier_gas
 import Luna.Modes: AbstractMode, dimlimits, neff, field, Aeff, N
 import Luna.PhysData: wlfreq
+import Base: show
 
 export MarcatilliMode, dimlimits, neff, field, N, Aeff
 
@@ -27,6 +28,16 @@ struct MarcatilliMode{Ta, Tcore, Tclad, LT} <: AbstractMode
     aeff_intg::Float64 # Pre-calculated integral fraction for effective area
 end
 
+function show(io::IO, m::MarcatilliMode)
+    a = "radius (z=0): $(m.a(0))"
+    loss = "loss: " * (m.loss == Val(true) ? "true" : "false")
+    model = "model: "*string(m.model)
+    out = join(["MarcatilliMode", "kind: "*mode_string(m), a, loss, model], "\n  ")
+    print(io, out)
+end
+
+mode_string(m::MarcatilliMode) = string(m.kind)*string(m.n)*string(m.m)
+
 function MarcatilliMode(a::Number, args...; kwargs...)
     afun(z) = a
     MarcatilliMode(afun, args...; kwargs...)

src/Ionisation.jl

@@ -28,7 +28,7 @@ function ionrate_fun!_ADK(ionpot::Float64, threshold=true)
                 (4*ω_p/(ω_t_prefac*abs(E)))^(2*nstar-1)
                 *exp(-4/3*ω_p/(ω_t_prefac*abs(E))))
             else
-                0
+                zero(E)
             end
         end
         function ionrate!(out, E)

src/Luna.jl

@@ -274,33 +274,43 @@ 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"
+
+gridtype(g::Grid.RealGrid) = "field-resolved"
+gridtype(g::Grid.EnvGrid) = "envelope"
+gridtype(g) = "unknown"
+
+simtype(g, t, l) = Dict("field" => gridtype(g),
+                        "transform" => transtype(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)
 
 
     Et = FT \ Eω
 
     z = 0.0
 
-    window! = let window=grid.ωwin, twindow=grid.twin, FT=FT, Et=Et
-        function window!(Eω)
-            Eω .*= window
-            ldiv!(Et, FT, Eω)
-            Et .*= twindow
-            mul!(Eω, FT, Et)
-        end
-    end
-
     function stepfun(Eω, z, dz, interpolant)
-        window!(Eω)
+        Eω .*= grid.ωwin
+        ldiv!(Et, FT, Eω)
+        Et .*= grid.twin
+        mul!(Eω, FT, Et)
         output(Eω, z, dz, interpolant)
     end
 
     output(Grid.to_dict(grid), group="grid")
+    output(simtype(grid, transform, linop), group="simulation_type")
 
     RK45.solve_precon(
         transform, linop, Eω, z, init_dz, grid.zmax, stepfun=stepfun,

src/Maths.jl

@@ -7,6 +7,7 @@ import Random: AbstractRNG, randn, MersenneTwister
 import FFTW
 import Luna
 import Luna.Utils: saveFFTwisdom, loadFFTwisdom
+import Roots: fzero
 import Dierckx
 
 "Calculate derivative of function f(x) at value x using finite differences"
@@ -64,18 +65,123 @@ function rms_width(x::Vector, y; dim = 1)
 end
 
 """
-Trapezoidal integration for multi-dimensional arrays, in-place or with output array.
-In all of these functions, x can be an array (the x axis) or a number (the x axis spacing)
+    fwhm(x, y; method=:linear, baseline=false, minmax=:min)
 
-In-place integration for multi-dimensional arrays
+Calculate the full width at half maximum (FWHM) of `y` on the axis `x`
+
+`method` can be `:spline` or `:nearest`. `:spline` uses a [`CSpline`](@ref), whereas
+`:nearest` finds the closest values either side of the crossing point and interpolates linearly.
+
+If `baseline` is true, the width is not taken at
+half the global maximum, but at half of the span of `y`.
+
+`minmax` determines whether the FWHM is taken at the narrowest (`:min`) or the widest (`:max`)
+point of y.
+"""
+function fwhm(x, y; method=:linear, baseline=false, minmax=:min)
+    minmax in (:min, :max) || error("minmax has to be :min or :max")
+    if baseline
+        val = minimum(y) + 0.5*(maximum(y) - minimum(y))
+    else
+        val = 0.5*maximum(y)
+    end
+    if !(method in (:spline, :linear, :nearest))
+        error("Unknown FWHM method $method")
+    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)
+        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
+    end
+end
+
+"""
+    linterpx(x1, x2, y1, y2, val)
+
+Given two points on a straight line, `(x1, y1)` and `(x2, y2)`, calculate the value of `x` 
+at which this straight line intercepts with `val`.
+
+# Examples
+```jldoctest
+julia> x1 = 0; x2 = 1; y1 = 0; y2 = 2; # y = 2x
+julia> linterpx(x1, x2, y1, y2, 0.5)
+0.25
+```
+"""
+function linterpx(x1, x2, y1, y2, val)
+    slope = (y2-y1)/(x2-x1)
+    icpt = y1 - slope*x1
+    (val-icpt)/slope
+end
+
+"""
+    hwhwm(f, x0=0; direction=:fwd)
+
+Find the value `x` where the function `f(x)` drops to half of its maximum, which is located
+at `x0`. For `direction==:fwd`, search in the region `x > x0`, for :bwd, search in `x < x0`.
 """
+function hwhm(f, x0=0; direction=:fwd)
+    m = f(x0)
+    fhalf(x) = f(x) - 0.5*m
+    if direction == :fwd
+        xhw = fzero(fhalf, x0, Inf)
+    elseif direction == :bwd
+        xhw = fzero(fhalf, -Inf, x0)
+    end
+    return abs(xhw - x0)
+end
+
+"""
+    cumtrapz!([out, ] y, x; dim=1)
+
+Trapezoidal integration for multi-dimensional arrays or vectors, in-place or with output array.
+
+If `out` is omitted, `y` is integrated in place. Otherwise the result is placed into `out`.
+
+`x` can be an array (the x axis) or a number (the x axis spacing).
+"""
+function cumtrapz! end
+
 function cumtrapz!(y, x; dim=1)
     idxlo = CartesianIndices(size(y)[1:dim-1])
     idxhi = CartesianIndices(size(y)[dim+1:end])
     _cumtrapz!(y, x, idxlo, idxhi)
 end
 
-"Inner function for multi-dimensional arrays - uses 1-D routine internally"
+"""
+    _cumtrapz!([out, ] y, x, idxlo, idxhi)
+
+Inner function for multi-dimensional `cumtrapz!` - uses 1-D routine internally
+"""
 function _cumtrapz!(y, x, idxlo, idxhi)
     for lo in idxlo
         for hi in idxhi
@@ -84,7 +190,6 @@ function _cumtrapz!(y, x, idxlo, idxhi)
     end
 end
 
-"In-place integration for 1-D arrays"
 function cumtrapz!(y::T, x) where T <: Union{SubArray, Vector}
     tmp = y[1]
     y[1] = 0
@@ -95,14 +200,12 @@ function cumtrapz!(y::T, x) where T <: Union{SubArray, Vector}
     end
 end
 
-"Integration into output array for multi-dimensional arrays"
 function cumtrapz!(out, y, x; dim=1)
     idxlo = CartesianIndices(size(y)[1:dim-1])
     idxhi = CartesianIndices(size(y)[dim+1:end])
     _cumtrapz!(out, y, x, idxlo, idxhi)
 end
 
-"Inner function for multi-dimensional arrays - uses 1-D routine internally"
 function _cumtrapz!(out, y, x, idxlo, idxhi)
     for lo in idxlo
         for hi in idxhi
@@ -111,24 +214,30 @@ function _cumtrapz!(out, y, x, idxlo, idxhi)
     end
 end
 
-"Integration into output array for 1-D array"
 function cumtrapz!(out, y::Union{SubArray, Vector}, x)
     out[1] = 0
     for i in 2:length(y)
         out[i] = out[i-1]+ 1//2*(y[i-1] + y[i])*_dx(x, i)
     end
 end
 
-"x axis spacing if x is given as an array"
-function _dx(x, i)
-    x[i] - x[i-1]
-end
+"""
+    _dx(x, i)
 
-"x axis spacing if x is given as a number (i.e. dx)"
-function _dx(x::Number, i)
-    x
-end
+Calculate the axis spacing at index `i` given an axis `x`. If `x` is a number, interpret this
+as `δx` directly
+"""
+_dx(x, i) = x[i] - x[i-1]
+_dx(δx::Number, i) = δx
 
+
+"""
+    cumtrapz(y, x; dim=1)
+
+Calculate the cumulative trapezoidal integral of `y`.
+
+`x` can be an array (the x axis) or a number (the x axis spacing).
+"""
 function cumtrapz(y, x; dim=1)
     out = similar(y)
     cumtrapz!(out, y, x; dim=dim) 

src/NonlinearRHS.jl

@@ -16,6 +16,7 @@ module NonlinearRHS
 import FFTW
 import Hankel
 import Cubature
+import Base: show
 import LinearAlgebra: mul!, ldiv!
 import NumericalIntegration: integrate, SimpsonEven
 import Luna: PhysData, Modes, Maths, Grid
@@ -169,6 +170,16 @@ mutable struct TransModal{tsT, lT, TT, FTT, rT, gT, dT, nT}
     mfcn::Int
 end
 
+function show(io::IO, t::TransModal)
+    grid = "grid type: $(typeof(t.grid))"
+    nmodes = "no. of modes: $(t.nmodes)"
+    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  ")
+    print(io, out)
+end
+
 "Transform E(ω) -> Pₙₗ(ω) for modal field."
 # FT - forward FFT for the grid
 # resp - tuple of nonlinear responses
@@ -195,8 +206,6 @@ function TransModal(grid::Grid.EnvGrid, args...; kwargs...)
     TransModal(ComplexF64, grid, args...; kwargs...)
 end
 
-show(io::IO, t::TransModal) = print(io, "TransModal{$(t.ts.nmodes) modes}")
-
 @noinline function reset!(t::TransModal, Emω::Array{ComplexF64,2}, z::Float64)
     t.Emω .= Emω
     t.ncalls = 0
@@ -280,6 +289,14 @@ struct TransModeAvg{TT, FTT, rT, gT, dT, nT, aT}
     aeff::aT # function which returns effective area
 end
 
+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  ")
+    print(io, out)
+end
+
 function TransModeAvg(TT, grid, FT, resp, densityfun, normfun, aeff)
     Eωo = zeros(ComplexF64, length(grid.ωo))
     Eto = zeros(TT, length(grid.to))
@@ -308,10 +325,32 @@ function (t::TransModeAvg)(nl, Eω, z)
 end
 
 "Calculate energy from modal field E(t)"
-energy_modal() = _energy_modal
+function energy_modal(grid::Grid.RealGrid)
+    function energy_t(t, Et)
+        Eta = Maths.hilbert(Et)
+        return integrate(grid.t, abs2.(Eta), SimpsonEven())
+    end
+
+    prefac = 2π/(grid.ω[end]^2)
+    function energy_ω(ω, Eω)
+        prefac*integrate(ω, abs2.(Eω), SimpsonEven())
+    end
+    return energy_t, energy_ω
+end
+
+function energy_modal(grid::Grid.EnvGrid)
+    function energy_t(t, Et)
+        return integrate(grid.t, abs2.(Et), SimpsonEven())
+    end
 
-_energy_modal(t, Et::Array{T, N}) where T <: Real where N = _energy_modal(t, Maths.hilbert(Et))
-_energy_modal(t, Et::Array{T, N}) where T <: Complex where N = abs(integrate(t, abs2.(Et), SimpsonEven()))
+    δω = grid.ω[2] - grid.ω[1]
+    Δω = length(grid.ω)*δω
+    prefac = 2π*δω/(Δω^2)
+    function energy_ω(ω, Eω)
+        prefac*sum(abs2.(Eω))
+    end
+    return energy_t, energy_ω
+end
 
 """
     TransRadial