Merge branch 'master' into PythonPlot

.github/workflows/documenter.yml

@@ -11,7 +11,7 @@ jobs:
   build:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - uses: julia-actions/setup-julia@latest
         with:
           version: ^1.7.0-rc1

.github/workflows/run_tests.yml

@@ -8,17 +8,17 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        julia-version: ['1.7', '1.8']
+        julia-version: ['1.7', '1.8', '1.9']
         julia-arch: [x64]
         os: [ubuntu-latest, windows-latest, macos-latest]
     if: github.event.pull_request.draft == false
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - uses: julia-actions/setup-julia@latest
         with:
           version: ${{ matrix.julia-version }}
           arch: ${{ matrix.julia-arch }}
       - name: Install dependencies
         env:
           PYTHON: # remove PYTHON env variable to make sure packages are installed
-        uses: julia-actions/julia-runtest@master
+        uses: julia-actions/julia-runtest@v1

Project.toml

@@ -1,7 +1,7 @@
 name = "Luna"
 uuid = "30eb0fb0-5147-11e9-3356-d75b018717ce"
 authors = ["chrisbrahms <38351086+chrisbrahms@users.noreply.github.com>"]
-version = "0.2.0"
+version = "0.3.0"
 
 [deps]
 ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
@@ -57,6 +57,7 @@ CoolProp = "0.1"
 Cubature = "1.5.1"
 DSP = "0.7.3"
 DataStructures = "0.18.10"
+DelimitedFiles = "1"
 Dierckx = "0.5.1"
 Documenter = "0.27.5"
 EllipsisNotation = "1"

docs/src/scans.md

@@ -8,7 +8,7 @@ import PythonPlot: pyplot
 
 a = 125e-6
 flength = 3
-gas = :HeJ
+gas = :He
 λ0 = 800e-9
 τfwhm = 10e-15
 λlims = (100e-9, 4e-6)
@@ -37,7 +37,7 @@ runscan(scan) do scanidx, energy, pressure
                    λlims, trange, scan, scanidx, filepath=outputdir)
 end
 ```
-Here, `runscan` uses the [do-block syntax](https://docs.julialang.org/en/v1/manual/functions/#Do-Block-Syntax-for-Function-Arguments), which wraps its body in a function. In our example this function takes three arguments: `scanidx, energy, pressure`. Generally, `scanidx` is always present and uniquely identifies a specific simulation in the scan (it simply runs from 1 to the number of simulations in the scan, `length(scan)`). The number of subsequent arguments is equal to the number of scan variables you have added to the `scan`. **Note** that the order in which you add the variables to the scan matters, and it must match the arguments in the `do` block. Alternatively, we could also have wrapped our scan in a function ourselves:
+Here, `runscan` uses the [do-block syntax](https://docs.julialang.org/en/v1/manual/functions/#Do-Block-Syntax-for-Function-Arguments), which wraps its body in a function. In our example this function takes three arguments: `scanidx, energy, pressure`. Generally, `scanidx` is always present and uniquely identifies a specific simulation in the scan (it simply runs from 1 to the number of simulations in the scan, `length(scan)`). The number of subsequent arguments is equal to the number of scan variables you have added to the `scan`, **and their order must match the order in which variables were added to the scan**. Alternatively, we could also have wrapped our scan in a function ourselves:
 ```julia
 function runone(scanidx, energy, pressure)
    prop_capillary(a, flength, gas, pressure; λ0, τfwhm, energy,
@@ -47,6 +47,9 @@ runscan(runone, scan)
 ```
 but the `do` block syntax is exactly equivalent to this and usually easier.
 
+!!! warning
+    The order in which you add the variables to the scan is important, and it **must** match the arguments in the `do` block. `scan = Scan("pressure_energy_example"; energy=energies, pressure=pressures)` is **not** equivalent to `scan = Scan("pressure_energy_example"; pressure=pressures, energy=energies)`.
+
 Running this script will simply run all the simulations, one after the other, in the current Julia session. To alter this behaviour, `Scan` also takes another argument, the execution mode, which has to be a subtype of `Scans.AbstractExec`. For example, to run only the first 10 items in the scan, we can use a `RangeExec`:
 ```julia
 scan = Scan("pressure_energy_example", Scans.RangeExec(1:10); energy=energies) # note the second argument here

examples/simple_interface/scan.jl

@@ -4,7 +4,7 @@ import PythonPlot: pyplot
 # Fixed parameters:
 a = 125e-6
 flength = 3
-gas = :HeJ
+gas = :He
 
 λ0 = 800e-9
 τfwhm = 10e-15

src/Fields.jl

@@ -392,14 +392,23 @@ function (s::SpatioTemporalField)(grid, spacegrid, FT)
     Eωk
 end
 
-# TODO: this for FreeGrid
 function prop!(Eωk, z, grid, q::Hankel.QDHT)
     kzsq = @. (grid.ω/PhysData.c)^2 - (q.k^2)'
     kzsq[kzsq .< 0] .= 0
     kz = sqrt.(kzsq)
     @. Eωk *= exp(-1im * z * (kz - grid.ω/PhysData.c))
 end
 
+function prop!(Eωk, z, grid, xygrid)
+    kzsq = ((grid.ω ./ PhysData.c).^2
+            .- reshape(xygrid.ky.^2, (1, length(xygrid.ky), 1))
+            .- reshape(xygrid.kx.^2, (1, 1, length(xygrid.kx)))
+    )
+    kzsq[kzsq.<0] .= 0
+    kz = sqrt.(kzsq)
+    @. Eωk *= exp(-1im * z * (kz - grid.ω / PhysData.c))
+end
+
 """
     gauss_beam(k, ω0; z=0.0, pol=:y)
 

src/Interface.jl

@@ -36,7 +36,7 @@ peak power specified.
 - `mode::Symbol`: Mode in which this input should be coupled. Can be `:lowest` for the
                   lowest-order mode in the simulation, or a mode designation
                   (e.g. `:HE11`, `:HE12`, `:TM01`, etc.). Defaults to `:lowest`.
-- `polarisation`: Can be `:linear`, `:circular`, or an ellipticity number -1 ≤ ε ≤ 1,
+- `polarisation`: Can be `:linear`, `:x`, `:y`, `:circular`, or an ellipticity number -1 ≤ ε ≤ 1,
                   where ε=-1 corresponds to left-hand circular, ε=1 to right-hand circular,
                   and ε=0 to linear polarisation.
 - `propagator`: A function `propagator!(Eω, grid)` which **mutates** its first argument to
@@ -71,7 +71,7 @@ specified.
 - `mode::Symbol`: Mode in which this input should be coupled. Can be `:lowest` for the
                   lowest-order mode in the simulation, or a mode designation
                   (e.g. `:HE11`, `:HE12`, `:TM01`, etc.). Defaults to `:lowest`.
-- `polarisation`: Can be `:linear`, `:circular`, or an ellipticity number -1 ≤ ε ≤ 1,
+- `polarisation`: Can be `:linear`, `:x`, `:y`, `:circular`, or an ellipticity number -1 ≤ ε ≤ 1,
                   where ε=-1 corresponds to left-hand circular, ε=1 to right-hand circular,
                   and ε=0 to linear polarisation.
 - `propagator`: A function `propagator!(Eω, grid)` which **mutates** its first argument to
@@ -105,7 +105,7 @@ specified, and duration given either as `τfwhm` or `τw`.
 - `mode::Symbol`: Mode in which this input should be coupled. Can be `:lowest` for the
                   lowest-order mode in the simulation, or a mode designation
                   (e.g. `:HE11`, `:HE12`, `:TM01`, etc.). Defaults to `:lowest`.
-- `polarisation`: Can be `:linear`, `:circular`, or an ellipticity number -1 ≤ ε ≤ 1,
+- `polarisation`: Can be `:linear`, `:x`, `:y`, `:circular`, or an ellipticity number -1 ≤ ε ≤ 1,
                   where ε=-1 corresponds to left-hand circular, ε=1 to right-hand circular,
                   and ε=0 to linear polarisation.
 - `propagator`: A function `propagator!(Eω, grid)` which **mutates** its first argument to
@@ -149,7 +149,7 @@ A custom pulse defined by tabulated data to be used with `prop_capillary`.
 - `mode::Symbol`: Mode in which this input should be coupled. Can be `:lowest` for the
                   lowest-order mode in the simulation, or a mode designation
                   (e.g. `:HE11`, `:HE12`, `:TM01`, etc.). Defaults to `:lowest`.
-- `polarisation`: Can be `:linear`, `:circular`, or an ellipticity number -1 ≤ ε ≤ 1,
+- `polarisation`: Can be `:linear`, `:x`, `:y`, `:circular`, or an ellipticity number -1 ≤ ε ≤ 1,
                   where ε=-1 corresponds to left-hand circular, ε=1 to right-hand circular,
                   and ε=0 to linear polarisation.
 - `propagator`: A function `propagator!(Eω, grid)` which **mutates** its first argument to
@@ -280,8 +280,8 @@ In this case, all keyword arguments except for `λ0` are ignored.
 - `power`: Peak power **after any spectral phases are added**.
 - `pulseshape`: Shape of the transform-limited pulse. Can be `:gauss` for a Gaussian pulse
     or `:sech` for a sech² pulse.
-- `polarisation`: Polarisation of the input pulse. Can be `:linear` (default), `:circular`,
-    or an ellipticity number -1 ≤ ε ≤ 1, where ε=-1 corresponds to left-hand circular,
+- `polarisation`: Polarisation of the input pulse. Can be `:linear` (default), `:x`, `:y`,
+    `:circular`, or an ellipticity number -1 ≤ ε ≤ 1, where ε=-1 corresponds to left-hand circular,
     ε=1 to right-hand circular, and ε=0 to linear polarisation. The major axis for
     elliptical polarisation is always the y-axis.
 - `propagator`: A function `propagator!(Eω, grid)` which **mutates** its first argument to
@@ -364,8 +364,6 @@ function prop_capillary_args(radius, flength, gas, pressure;
     plasma = isnothing(plasma) ? !envelope : plasma
     thg = isnothing(thg) ? !envelope : thg
 
-    gas = (gas == :He) ? :HeJ : gas
-
     grid = makegrid(flength, λ0, λlims, trange, envelope, thg, δt)
     mode_s = makemode_s(modes, flength, radius, gas, pressure, model, loss, pol)
     check_orth(mode_s)
@@ -407,12 +405,13 @@ end
 function needpol(pol)
     if pol == :linear
         return false
-    elseif pol == :circular
+    elseif pol in (:circular, :x, :y)
         return true
     else
-        error("Polarisation must be :linear, :circular, or an ellipticity, not $pol")
+        error("Polarisation must be :linear, :circular, :x/:y, or an ellipticity, not $pol")
     end
 end
+
 needpol(pol::Number) = true
 needpol(pulse::Pulses.AbstractPulse) = needpol(pulse.polarisation)
 needpol(pulses::Vector{<:Pulses.AbstractPulse}) = any(needpol, pulses)
@@ -661,8 +660,10 @@ _findmode(mode_s, md) = _findmode([mode_s], md)
 function makeinputs(mode_s, λ0, pulse::Pulses.AbstractPulse)
     idcs = findmode(mode_s, pulse)
     (length(idcs) > 0) || error("Mode $(pulse.mode) not found in mode list: $mode_s")
-    if pulse.polarisation == :linear
+    if pulse.polarisation == :linear || pulse.polarisation == :x
         ((mode=idcs[1], fields=(pulse.field,)),)
+    elseif pulse.polarisation == :y
+        ((mode=idcs[2], fields=(pulse.field,)),)
     else
         (length(idcs) == 2) || error("Modes not set up for circular/elliptical polarisation")
         f1, f2 = ellfields(pulse)

src/Luna.jl

@@ -371,4 +371,31 @@ function run(Eω, grid,
         status_period=status_period)
 end
 
+# run some code for precompilation
+Logging.with_logger(Logging.NullLogger()) do
+    prop_capillary(125e-6, 0.3, :He, 1.0; λ0=800e-9, energy=1e-9,
+                    τfwhm=10e-15, λlims=(150e-9, 4e-6), trange=1e-12, saveN=11)
+    prop_capillary(125e-6, 0.3, :He, (1.0, 0); λ0=800e-9, energy=1e-9,
+                    τfwhm=10e-15, λlims=(150e-9, 4e-6), trange=1e-12, saveN=11)
+    prop_capillary(125e-6, 0.3, :He, 1.0; λ0=800e-9, energy=1e-9,
+                    τfwhm=10e-15, λlims=(150e-9, 4e-6), trange=1e-12, saveN=11,
+                    modes=4)
+    p = Tools.capillary_params(120e-6, 10e-15, 800e-9, 125e-6, :He, P=1.0)
+
+    # gnlse_sol.jl example but with N=1 and 100th of the fibre length
+    γ = 0.1
+    β2 = -1e-26
+    τ0 = 280e-15
+    τfwhm = (2*log(1 + sqrt(2)))*τ0
+    fr = 0.18
+    P0 = abs(β2)/((1-fr)*γ*τ0^2)
+    flength = π*τ0^2/abs(β2)/100
+    βs =  [0.0, 0.0, β2]
+    λ0 = 835e-9
+    λlims = [450e-9, 8000e-9]
+    trange = 4e-12
+    output = prop_gnlse(γ, flength, βs; λ0, τfwhm, power=P0, pulseshape=:sech, λlims, trange,
+                        raman=true, shock=true, fr, shotnoise=true, saveN=11)
+end
+
 end # module

src/Maths.jl

@@ -7,6 +7,7 @@ import Random: AbstractRNG, randn, GLOBAL_RNG
 import FFTW
 import Luna
 import Luna.Utils: saveFFTwisdom, loadFFTwisdom
+import Luna: settings
 import Roots: fzero
 import Dierckx
 import Peaks
@@ -475,6 +476,82 @@ end
 _gaborFT(x::Array{T, 2}) where T <: Real = FFTW.rfft(x, 1)
 _gaborFT(x::Array{T, 2}) where T <: Complex = FFTW.fft(x, 1)
 
+
+"""
+    wigner(t, A; downsample=1, crop=1)
+
+Compute the wigner distribution function for the field `A` sampled on time axis `t`.
+The size of `t` and `A` should be a power of 2. The size of the transform can be reduced
+(and the computation sped up) by either `downsample`, which reduces the time resolution (frequency range),
+or by `crop` which reduces the time range (frequency resolution). If given, `downsample` and `crop` must
+also be a power of 2.
+
+See also https://en.wikipedia.org/wiki/Wigner_distribution_function
+"""
+function wigner(t, A::Vector{<:Complex}; downsample=1, crop=1)
+    # crop and/or downsample
+    log2(downsample) % 1 ≠ 0 && error("downsample factor must be a power of 2")
+    if crop > 1
+        log2(crop) % 1 ≠ 0 && error("cropping factor must be a power of 2")
+        ncrop = (length(t) - length(t) ÷ crop) ÷ 2
+        startidx = ncrop
+        endidx = length(t) - ncrop - 1
+    else
+        startidx = 1
+        endidx = length(t)
+    end
+
+    A = A[startidx:downsample:endidx]
+    t = t[startidx:downsample:endidx]
+
+    # pad with zeros in the time domain to allow for shifting
+    l = length(t)
+    n = l ÷ 2
+    Ao = vcat(zeros(n), A, zeros(n))
+
+    # make frequency axis for expanded time axis
+    δt = t[2] - t[1]
+    Nt = collect(range(0, length=2l))
+    to = (Nt .- l) .* δt
+    ωo = fftfreq(to)
+    ωos = FFTW.fftshift(ωo)
+
+    # plan FFT
+    FT = FFTW.plan_fft(copy(Ao), 1; flags=settings["fftw_flag"])
+
+    Af = FT * Ao
+    Afs = similar(Af)
+
+    function τshift!(x, Af, τ, cc)
+        @. Afs = Af * exp.(-1im * ωos * τ)
+        ldiv!(x, FT, Afs)
+        cc && conj!(x)
+    end
+
+    Wt = zeros(ComplexF64, (length(t), length(ωo)))
+    Ats = similar(Ao)
+    Atc = similar(Ao)
+    for (idx, τi) in enumerate(t) # iterate over time which here turns into a delay
+        τshift!(Ats, Af, -τi/2, false) # A(t + τ/2)
+        τshift!(Atc, Af, τi/2, true) # A*(t - τ/2)
+        Wt[idx, :] .= Ats .* Atc
+    end
+
+    # middle of the time window
+    τgrid = l*δt/2
+
+    ω = fftfreq(t)
+    # Fourier transform along the delay axis to get the Wigner distribution
+    Wf = FFTW.fftshift(
+        FFTW.fft(Wt, 1) .* exp.(1im .* FFTW.fftshift(ω) .* τgrid),
+        1
+    )[:, n:end-n-1] # crop to remove the time domain padding we added earlier
+
+    t, ω, real(Wf)
+end
+
+wigner(t, A::Vector{<:Real}; kwargs...) = wigner(t, hilbert(A); kwargs...)
+
 """
     hilbert(x; dim=1)
 

src/PhysData.jl

@@ -45,9 +45,10 @@ const amg = atm/(k_B*273.15)
 "Atomic mass unit"
 const m_u = ustrip(CODATA2014.m_u)
 
-const gas = (:Air, :He, :HeJ, :Ne, :Ar, :Kr, :Xe, :N2, :H2, :O2, :CH4, :SF6, :N2O, :D2)
+const gas = (:Air, :He, :HeJ, :HeB, :Ne, :Ar, :Kr, :Xe, :N2, :H2, :O2, :CH4, :SF6, :N2O, :D2)
 const gas_str = Dict(
     :He => "He",
+    :HeB => "He",
     :HeJ => "He",
     :Ar => "Ar",
     :Ne => "Neon",
@@ -142,13 +143,13 @@ calculated from Sellmeier expansions.
 """
 function sellmeier_gas(material::Symbol)
     dens = dens_1bar_0degC[material]
-    if material == :He
+    if material == :HeB
         B1 = 4977.77e-8
         C1 = 28.54e-6
         B2 = 1856.94e-8
         C2 = 7.76e-3
         return γ_Börzsönyi(B1/dens, C1, B2/dens, C2)
-    elseif material == :HeJ
+    elseif material == :He || material == :HeJ
         B1 = 2.16463842e-05
         C1 = -6.80769781e-04
         B2 = 2.10561127e-07
@@ -570,7 +571,7 @@ References:
 function γ3_gas(material::Symbol; source=nothing)
     # TODO: More Bishop/Shelton; Wahlstrand updated values.
     if source === nothing
-        if material in (:He, :HeJ, :Ne, :Ar, :Kr, :Xe, :N2)
+        if material in (:He, :HeB, :HeJ, :Ne, :Ar, :Kr, :Xe, :N2)
             source = :Lehmeier
         elseif material in (:H2, :CH4, :SF6, :D2)
             source = :Shelton
@@ -585,7 +586,7 @@ function γ3_gas(material::Symbol; source=nothing)
     if source == :Lehmeier
         dens = dens_1atm_0degC[material]
         # Table 1 in [3]
-        if material in (:He, :HeJ)
+        if material in (:He, :HeB, :HeJ)
             fac = 1
         elseif material == :Ne
             fac = 1.8
@@ -708,7 +709,7 @@ Return the first ionisation potential of the `material` in a specific unit (defa
 Possible units are `:SI`, `:atomic` and `:eV`.
 """
 function ionisation_potential(material; unit=:SI)
-    if material in (:He, :HeJ)
+    if material in (:He, :HeB, :HeJ)
         Ip = 0.9036
     elseif material == :Ne
         Ip = 0.7925
@@ -764,7 +765,7 @@ function quantum_numbers(material)
         return 4, 1, 1
     elseif material == :Xe
         return 5, 1, 1
-    elseif material in (:He, :HeJ)
+    elseif material in (:He, :HeB, :HeJ)
         return 1, 0, 1
     elseif material == :O2
         return 2, 0, 0.53 # https://doi.org/10.1016/S0030-4018(99)00113-3

src/Plotting.jl

@@ -672,7 +672,7 @@ function cornertext(ax, text; corner="ul", pad=0.02, xpad=nothing, ypad=nothing,
     if corner[2] == 'l'
         hal = "left"
         x = xpad
-    elseif [2] == 'r'
+    elseif corner[2] == 'r'
         hal = "right"
         x = 1 - xpad
     else