text:
# This file contains methods for validated integration of ODEs using Taylor models,
# originally from TaylorModels.jl. See license below.

#==============================================================================
The TaylorModels.jl package is licensed under the MIT "Expat" License:

> Copyright (c) 2018-2020: David Sanders and Luis Benet.
>
>
> Permission is hereby granted, free of charge, to any person obtaining a copy
>
> of this software and associated documentation files (the "Software"), to deal
>
> in the Software without restriction, including without limitation the rights
>
> to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
>
> copies of the Software, and to permit persons to whom the Software is
>
> furnished to do so, subject to the following conditions:
>
>
>
> The above copyright notice and this permission notice shall be included in all
>
> copies or substantial portions of the Software.
>
>
>
> THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
>
> IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
>
> FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
>
> AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
>
> LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
>
> OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
>
> SOFTWARE.
>
>
=#
#===============================================================================#

"""
    remainder_taylorstep!(f!, t, x, dx, xI, dxI, δI, δt, params, adaptive)

Returns a remainder for the integration step for the dependent variables (`x`)
checking that the solution satisfies the criteria for existence and uniqueness.
"""
function remainder_taylorstep!(f!::Function, t::Taylor1{T},
        x::Vector{Taylor1{TaylorN{T}}}, dx::Vector{Taylor1{TaylorN{T}}},
        xI::Vector{Taylor1{IA.Interval{T}}}, dxI::Vector{Taylor1{IA.Interval{T}}},
        δI::IntervalBox{N,T}, δt::IA.Interval{T}, params, adaptive::Bool) where {N,T}

    orderT = get_order(dx[1])
    aux = δt^(orderT+1)
    Δx  = IntervalBox([xI[i][orderT+1] for i in eachindex(xI)]) * aux
    Δdx = IntervalBox([dxI[i][orderT+1] for i in eachindex(xI)]) * aux
    Δ0  = IntervalBox([dx[i][orderT](δI) for i in eachindex(x)]) * aux / (orderT+1)
    Δ = Δ0 + Δdx * δt
    Δxold = Δx

    # Checking existence and uniqueness
    iscontractive(Δ, Δx) && return (true, Δx, t[0])

    # If the check didn't work, compute new remainders. A new Δx is proposed,
    # and the corresponding Δdx is computed
    xxI  = Array{Taylor1{TaylorN{IA.Interval{T}}}}(undef, N)
    dxxI = Array{Taylor1{TaylorN{IA.Interval{T}}}}(undef, N)
    vv = Array{IA.Interval{T}}(undef, N)
    for its = 1:50
        # Remainder of Picard iteration
        Δ = picard_remainder!(f!, t, x, dx, xxI, dxxI, δI, δt, Δx, Δ0, params)

        # Checking existence and uniqueness
        iscontractive(Δ, Δx) && return (true, Δx, t[0])
        # iscontractive(Δ, Δx) && return _contract_iteration!(f!, t, x, dx, xxI, dxxI, δI, δt, Δx, Δdx, Δ0, params)

        # Expand Δx in the directions needed
        Δxold = Δx
        if Δ ⊆ Δx
            @inbounds for ind in 1:N
                # Widen the directions where ⊂ does not hold
                vv[ind] = Δx[ind]
                if Δ[ind] == Δx[ind]
                    vv[ind] = widen.(Δ[ind])
                end
            end
            Δx = IntervalBox(vv)
            continue
        end
        Δx = Δ
    end

    if adaptive
        return (false, Δx, t[0])
    else
        @format full
        error("""
        Error: cannot prove existence and unicity of the solution:
            t0 = $(t[0])
            δt = $(δt)
            Δ  = $(Δ)
            Δxo = $(Δxold)
            Δx = $(Δx)
            $(Δ .⊆ Δxold)
        """)
    end
end

"""
    iscontractive(Δ, Δx)

Checks if `Δ .⊂ Δx` is satisfied. If ``Δ ⊆ Δx` is satisfied, it returns
`true` if all cases where `==` holds corresponds to the zero `IA.Interval`.
"""
function iscontractive(Δ::IntervalBox{N,T}, Δx::IntervalBox{N,T}) where{N,T}
    zi = IA.Interval{T}(0, 0)
    @inbounds for ind in 1:N
        Δ[ind] ⊂ Δx[ind] && continue
        Δ[ind] == Δx[ind] == zi && continue
        return false
    end
    return true
end

"""
    picard_remainder!(f!, t, x, dx, xxI, dxxI, δI, δt, Δx, Δ0, params)

Return the remainder of Picard operator
"""
function picard_remainder!(f!::Function, t::Taylor1{T},
    x::Vector{Taylor1{TaylorN{T}}}, dx::Vector{Taylor1{TaylorN{T}}},
    xxI::Vector{Taylor1{TaylorN{IA.Interval{T}}}},
    dxxI::Vector{Taylor1{TaylorN{IA.Interval{T}}}},
    δI::IntervalBox{N,T}, δt::IA.Interval{T},
    Δx::IntervalBox{N,T}, Δ0::IntervalBox{N,T}, params) where {N,T}

    # Extend `x` and `dx` to have IA.Interval coefficients
    zI = zero(IA.Interval{T})
    @inbounds for ind in eachindex(x)
        xxI[ind]  = x[ind] + Δx[ind]
        dxxI[ind] = dx[ind] + zI
    end

    # Compute `dxxI` from the equations of motion
    f!(dxxI, xxI, params, t)

    # Picard iteration, considering only the bound of `f` and the last coeff of f
    Δdx = IntervalBox(TM.evaluate.( (dxxI - dx)(δt), (δI,) ) )
    Δ = Δ0 + Δdx * δt
    return Δ
end

"""
    absorb_remainder(a::TaylorModelN)

Returns a TaylorModelN, equivalent to `a`, such that the remainder
is mostly absorbed in the constant and linear coefficients. The linear shift assumes
that `a` is normalized to the `IntervalBox(-1..1, Val(N))`.

Ref: Xin Chen, Erika Abraham, and Sriram Sankaranarayanan,
"Taylor Model Flowpipe Construction for Non-linear Hybrid
Systems", in Real Time Systems Symposium (RTSS), pp. 183-192 (2012),
IEEE Press.
"""
function absorb_remainder(a::TaylorModelN{N,T,T}) where {N,T}
    Δ = remainder(a)
    orderQ = get_order(a)
    δ = IntervalBox(IA.Interval{T}(-1,1), Val(N))
    aux = diam(Δ)/(2N)
    rem = IA.Interval{T}(0, 0)

    # Linear shift
    lin_shift = mid(Δ) + sum((aux*TaylorN(i, order=orderQ) for i in 1:N))
    bpol = a.pol + lin_shift

    # Compute the new remainder
    aI = a(δ)
    bI = bpol(δ)

    if bI ⊆ aI
        rem = IA.Interval(aI.lo-bI.lo, aI.hi-bI.hi)
    elseif aI ⊆ bI
        rem = IA.Interval(bI.lo-aI.lo, bI.hi-aI.hi)
    else
        r_lo = aI.lo-bI.lo
        r_hi = aI.hi-bI.hi
        if r_lo > 0
            rem = IA.Interval(-r_lo, r_hi)
        else
            rem = IA.Interval( r_lo, -r_hi)
        end
    end

    return TaylorModelN(bpol, rem, a.x0, a.dom)
end

"""
    validated-step!
"""
function validated_step!(f!, t::Taylor1{T}, x::Vector{Taylor1{TaylorN{T}}},
        dx::Vector{Taylor1{TaylorN{T}}}, xaux::Vector{Taylor1{TaylorN{T}}},
        tI::Taylor1{T}, xI::Vector{Taylor1{IA.Interval{T}}},
        dxI::Vector{Taylor1{IA.Interval{T}}}, xauxI::Vector{Taylor1{IA.Interval{T}}},
        t0::T, tmax::T, sign_tstep::Int,
        xTMN::Vector{TaylorModelN{N,T,T}}, xv::Vector{IntervalBox{N,T}},
        rem::Vector{IA.Interval{T}}, zbox::IntervalBox{N,T}, symIbox::IntervalBox{N,T},
        nsteps::Int, orderT::Int, abstol::T, params, tmpTaylor, arrTaylor,
        tmpTaylorI, arrTaylorI, parse_eqs::Bool,
        check_property::Function=(t, x)->true, adaptive::Bool=true) where {N,T}

    # One step integration (non-validated)
    # TaylorIntegration.__jetcoeffs!(Val(parse_eqs), f!, t, x, dx, xaux, params)
    # δt = TaylorIntegration.stepsize(x, abstol)
    δt = TaylorIntegration.taylorstep!(f!, t, x, dx, xaux, abstol, params, tmpTaylor, arrTaylor, parse_eqs)
    f!(dx, x, params, t)  # Update `dx[:][orderT]`

    # One step integration for the initial box
    # TaylorIntegration.__jetcoeffs!(Val(parse_eqs), f!, tI, xI, dxI, xauxI, params)
    # δtI = TaylorIntegration.stepsize(xI, abstol)
    δtI = TaylorIntegration.taylorstep!(f!, tI, xI, dxI, xauxI, abstol, params, tmpTaylorI, arrTaylorI, parse_eqs)
    f!(dxI, xI, params, tI)  # Update `dxI[:][orderT+1]`

    # Step size
    δt = min(δt, sign_tstep*(tmax-t0))
    δt = sign_tstep * δt

    # Test if `check_property` is satisfied; if not, half the integration time.
    # If after 25 checks `check_property` is not satisfied, throw an error.
    nsteps += 1
    issatisfied = false
    rem_old = copy(rem)
    local success, _t0

    for nchecks = 1:25
        # Validate the solution: remainder consistent with Schauder thm
        δtI = sign_tstep * IA.Interval(0, sign_tstep*δt)
        (success, Δ, _t0) = remainder_taylorstep!(f!, t, x, dx, xI, dxI, symIbox, δtI, params, adaptive)
        if adaptive && !success
            break
        end

        rem .= rem_old .+ Δ

        # Create TaylorModelN to store remainders and evaluation
        @inbounds begin
            for i in eachindex(x)
                xTMN[i] = fp_rpa( TaylorModelN(x[i](δtI), rem[i], zbox, symIbox) )

                # If remainder is still too big, do it again
                # TODO: check if this can be avoided.
                j = 0
                while (j < 10) && (mag(rem[i]) > 1.0e-10)
                    j += 1
                    xTMN[i] = absorb_remainder(xTMN[i])
                    rem[i] = remainder(xTMN[i])
                end
            end
            xv[nsteps] = TM.evaluate(xTMN, symIbox) # IntervalBox

            if !check_property(t0+δt, xv[nsteps])
                δt = δt/2
                continue
            end
        end # @inbounds

        issatisfied = true
        break
    end

    if !issatisfied && !adaptive # TEMP do not return the error msg if adaptive is true
        error("""
            `check_property` is not satisfied:
            $t0 $nsteps $δt
            $(xv[nsteps])
            $(check_property(t0+δt, xv[nsteps]))""")
    end

    return (success, δt, _t0)
end

function validated_integ!(F, f!, X0, t0::T, tmax::T, orderQ::Int, orderT::Int,
                          abstol::T, max_steps::Int, X::LazySet, disjointness, Δt0,
                          adaptive::Bool=true, params=nothing;
                          parse_eqs::Bool=true, check_property::Function=(t, x)->true) where {T<:Real}

    # Set proper parameters for jet transport
    N = get_numvars()
    dof = N

    # Some variables
    zI = zero(IA.Interval{T})
    zbox = IntervalBox(zI, Val(N))
    symIbox = IntervalBox(IA.Interval{T}(-1, 1), Val(N))
    t   = t0 + Taylor1(orderT)
    tI  = t0 + Taylor1(orderT+1)

    # Allocation of vectors

    # Output
    tv    = Vector{T}(undef, max_steps+1)
    xv    = Vector{IntervalBox{N,T}}(undef, max_steps+1)
    xTM1v = Matrix{TaylorModel1{TaylorN{T},T}}(undef, dof, max_steps+1)
    rem = Vector{IA.Interval{T}}(undef, dof)

    # Internals: jet transport integration
    x     = Vector{Taylor1{TaylorN{T}}}(undef, dof)
    dx    = Vector{Taylor1{TaylorN{T}}}(undef, dof)
    xaux  = Vector{Taylor1{TaylorN{T}}}(undef, dof)
    xTMN  = Vector{TaylorModelN{N,T,T}}(undef, dof)

    # Internals: Taylor1{IA.Interval{T}} integration
    xI    = Vector{Taylor1{IA.Interval{T}}}(undef, dof)
    dxI   = Vector{Taylor1{IA.Interval{T}}}(undef, dof)
    xauxI = Vector{Taylor1{IA.Interval{T}}}(undef, dof)

    # Set initial conditions
    TaylorModels.initialize!(X0, orderQ, orderT, x, dx, xTMN, xI, dxI, rem, xTM1v)
    sign_tstep = copysign(1, tmax-t0)

    # Output vectors
    @inbounds tv[1] = t0
    @inbounds xv[1] = TM.evaluate(xTMN, symIbox)

    # Determine if specialized jetcoeffs! method exists (built by @taylorize)
    parse_eqs, tmpTaylor, arrTaylor = TaylorIntegration._determine_parsing!(parse_eqs, f!, t, x, dx, params)
    parse_eqsI, tmpTaylorI, arrTaylorI = TaylorIntegration._determine_parsing!(parse_eqs, f!, tI, xI, dxI, params)
    @assert parse_eqs == parse_eqsI

    local success # if true, the validation step succeeded
    local _t0 # represents how much the integration could advance until validation failed

    # Integration
    nsteps = 1
    while sign_tstep*t0 < sign_tstep*tmax

        # Validated step of the integration
        (success, δt, _t0) = validated_step!(f!, t, x, dx, xaux, tI, xI, dxI, xauxI,
                                             t0, tmax, sign_tstep, xTMN, xv, rem, zbox,
                                             symIbox, nsteps, orderT, abstol, params,
                                             tmpTaylor, arrTaylor, tmpTaylorI, arrTaylorI,
                                             parse_eqs, check_property, adaptive)

        if adaptive && !success
            break
        end

        # New initial conditions and time
        nsteps += 1
        t0 += δt
        @inbounds begin
            t[0] = t0
            tI[0] = t0
            tv[nsteps] = t0
            for i in eachindex(x)
                δtI = sign_tstep * IA.Interval{T}(0, sign_tstep*δt)
                xTM1v[i, nsteps] = TaylorModel1(deepcopy(x[i]), rem[i], zI, δtI)
                aux = x[i](δt)
                x[i]  = Taylor1( aux, orderT )
                dx[i] = Taylor1( zero(aux), orderT )
                auxI = xTMN[i](symIbox)
                xI[i] = Taylor1( auxI, orderT+1 )
                dxI[i] = xI[i]
            end
        end

        # construct the taylor model reach-set
        Ri = TaylorModelReachSet(xTM1v[:, nsteps], TimeInterval(t0-δt, t0) + Δt0)

        # check intersection with invariant
        _is_intersection_empty(Ri, X, disjointness) && break

        # update output flowpipe
        # note that F has 1 less element than xTM1v and xv (we don't store the initial set)
        push!(F, Ri)

        if nsteps > max_steps
            @warn("Maximum number of integration steps reached; exiting. Try increasing `max_steps`.")
            break
        end

    end

    return F, view(tv, 1:nsteps), view(xv, 1:nsteps), view(xTM1v, :, 1:nsteps), success, _t0
end


text:
"""
    groupby(d::DTable, col::Symbol; merge=true, chunksize=0) -> GDTable

Groups `d` by distinct values of column `col`.

The process of grouping can be affected by providing kwargs `merge` and `chunksize`.
By default all the chunks belonging to a single key will be merged into a single partition.
Providing a positive value in `chunksize` will attempt to merge the smaller partitions
into partitions not bigger than `chunksize`. Please note that partitions bigger than `chunksize`
will not be split into partitions of `chunksize`.
Merging can be disabled completely by providing `merge=false`.

# Examples
```julia
julia> d = DTable((a=shuffle(repeat('a':'d', inner=4, outer=4)),), 4)
DTable with 16 partitions
Tabletype: NamedTuple

julia> Dagger.groupby(d, :a)
GDTable with 4 partitions and 4 keys
Tabletype: NamedTuple
Grouped by: [:a]

julia> Dagger.groupby(d, :a, chunksize=3)
GDTable with 24 partitions and 4 keys
Tabletype: NamedTuple
Grouped by: [:a]

julia> Dagger.groupby(d, :a, merge=false)
GDTable with 42 partitions and 4 keys
Tabletype: NamedTuple
Grouped by: [:a]
```
"""
function groupby(d::DTable, col::Symbol; merge=true, chunksize=0)
    rowmap = (_row, _col) -> Tables.getcolumn(_row, _col)
    rowmap_w_cols = _row -> rowmap(_row, col)
    _groupby(d, rowmap_w_cols, [col], merge, chunksize)
end

"""
    groupby(d::DTable, cols::Vector{Symbol}; merge=true, chunksize=0)

Groups the `d` by distinct values of columns `cols`.
The key is constructed by creating a NamedTuple from each row based on `cols` provided.

For kwargs usage details see `groupby(d::DTable, col::Symbol)`

# Examples
```julia
julia> d = DTable((a=shuffle(repeat('a':'d', inner=4, outer=4)),b=repeat(1:4, 16)), 4)
DTable with 16 partitions
Tabletype: NamedTuple

julia> Dagger.groupby(d, [:a,:b])
GDTable with 16 partitions and 16 keys
Tabletype: NamedTuple
Grouped by: [:a, :b]

julia> Dagger.groupby(d, [:a,:b], chunksize=3)
GDTable with 27 partitions and 16 keys
Tabletype: NamedTuple
Grouped by: [:a, :b]

julia> Dagger.groupby(d, [:a,:b], merge=false)
GDTable with 64 partitions and 16 keys
Tabletype: NamedTuple
Grouped by: [:a, :b]
```
"""
function groupby(d::DTable, cols::Vector{Symbol}; merge=true, chunksize=0)
    rowmap = (_row, _cols) -> (;[c => Tables.getcolumn(_row, c) for c in _cols]...)
    rowmap_w_cols = _row -> rowmap(_row, cols)
    _groupby(d, rowmap_w_cols, cols, merge, chunksize)
end

"""
    groupby(d::DTable, f::Function; merge=true, chunksize=0)

Groups `d` by the distinct set of keys created by applying `f` to each row in `d`.

For kwargs usage details see `groupby(d::DTable, col::Symbol)`

```julia
julia> d = DTable((a=shuffle(repeat('a':'d', inner=4, outer=4)),b=repeat(1:4, 16)), 4)
DTable with 16 partitions
Tabletype: NamedTuple

julia> function group_fun(row)
           row.a + row.b
       end
group_fun (generic function with 1 method)

julia> Dagger.groupby(d, group_fun)
GDTable with 7 partitions and 7 keys
Tabletype: NamedTuple
Grouped by: group_fun

julia> Dagger.groupby(d, row -> row.a + row.b, chunksize=3)
GDTable with 25 partitions and 7 keys
Tabletype: NamedTuple
Grouped by: group_fun

julia> Dagger.groupby(d, row -> row.a + row.b, merge=false)
GDTable with 52 partitions and 7 keys
Tabletype: NamedTuple
Grouped by: group_fun
```
"""
groupby(d::DTable, f::Function; merge=true, chunksize=0) = _groupby(d, f, nothing, merge, chunksize)

"""
    _groupby(d::DTable, row_function::Function, cols::Union{Nothing, Vector{Symbol}}, merge::Bool, chunksize::Int)

Internal function for performing the groupby steps based on common arguments.
"""
function _groupby(
    d::DTable,
    row_function::Function,
    cols::Union{Nothing, Vector{Symbol}},
    merge::Bool,
    chunksize::Int)

    grouping_function = cols === nothing ? row_function : nothing

    spawner = (_dchunks, _row_function) -> Vector{EagerThunk}([Dagger.@spawn distinct_partitions(c, _row_function) for c in _dchunks])

    v = Dagger.@spawn spawner(d.chunks, row_function)
    index, chunks = fetch(Dagger.@spawn build_groupby_index(merge, chunksize, tabletype(d), v))
    GDTable(DTable(chunks, d.tabletype), cols, index, grouping_function)
end

"""
    distinct_partitions(chunk, f::Function)

Takes a partition and groups its rows according based on the key value returned by `f`.
"""
function distinct_partitions(chunk, f::Function)
    rows = Tables.rows(chunk)
    keyval = f(iterate(rows)[1])

    _distinct_partitions_iterate(chunk, f, keyval)
end

function _distinct_partitions_iterate(chunk, f, keyval::T) where T
    rows = Tables.rows(chunk)
    acc = Dict{T, Vector{eltype(rows)}}()

    for row in rows
        key = convert(T, f(row))
        v = get!(acc, key, Vector{eltype(rows)}())
        push!(v, row)
    end

    Vector{Pair{T, Chunk}}([x => Dagger.tochunk(Tables.columntable(acc[x])) for x in collect(keys(acc))])
end

rowcount(chunk) = length(Tables.rows(chunk))

"""
    build_groupby_index(merge::Bool, chunksize::Int, tabletype, vs...)

Takes the intermediate result of `distinct_partitions` and builds an index.
Merges partitions if possible according to the `chunksize` provided.
It will only merge chunks if their length after merging is `<= chunksize`.
It doesn't split chunks larger than `chunksize` and small chunks
may be leftover after merging if no appropriate pair was found.
"""
function build_groupby_index(
    merge::Bool,
    chunksize::Int,
    tabletype,
    spawner_out)
    vs = fetch(spawner_out)
    v1 = fetch(vs[1])
    @assert typeof(v1) <: Vector
    @assert eltype(v1) <: Pair
    keytype = eltype(v1).types[1]
    chunks = Vector{Chunk}()
    idx = Dict{keytype, Vector{UInt}}()

    i = one(UInt)
    for v in vs
        _v = fetch(v)
        for (key, chunk) in _v
            get!(idx, key, Vector{UInt}())
            push!(idx[key], i)
            push!(chunks, chunk)
            i += 1
        end
    end

    spawner_out = nothing
    v1 = nothing

    if merge && chunksize <= 0 # merge all partitions into one
        sink = Tables.materializer(tabletype())
        merged_chunks = Vector{Union{EagerThunk, Chunk}}()
        sizehint!(merged_chunks, length(keys(idx)))

        merge_spawner = (_chunks, _partition, _idx, _sink) -> begin
            map(k -> begin
                cs = getindex.(Ref(_chunks), _idx[k])
                length(cs) == 1 && return cs[1]
                Dagger.@spawn merge_chunks(_sink, cs)
            end, _partition)
        end
        vs = Dagger.@spawn merge_spawner(chunks, collect(keys(idx)), idx, sink)
        new_idx = Dict{keytype, Vector{UInt}}()
        for (i, k) in enumerate(keys(idx))
            new_idx[k] = [i]
        end
        append!(merged_chunks, fetch(vs))
        return new_idx, merged_chunks

    elseif merge && chunksize > 0 # merge all but try to merge all the small chunks into chunks of chunksize
        sink = Tables.materializer(tabletype())
        merged_chunks = Vector{Union{EagerThunk, Chunk}}()

        all_lengths = [Dagger.@spawn rowcount(c) for c in chunks]

        for k in keys(idx)
            _indices = idx[k]
            _lengths = fetch.(getindex.(Ref(all_lengths), _indices))

            ord = sortperm(_lengths, rev=true) # sorting indices and lengths by lengths
            _indices .= _indices[ord]
            _lengths .= _lengths[ord]

            l = 1
            r = length(_indices)
            prev_r = r

            while l <= r
                if _lengths[l] >= chunksize || # chunk already bigger than minimum, so move forward
                    _lengths[l] + _lengths[r] > chunksize || # potential merge would be bigger, so move forward
                    l == r # last iteration to push last chunk and trigger the last merge

                    _chunk = chunks[_indices[l]]
                    if r < prev_r # only condition for merging
                        chunks_to_merge = getindex.(Ref(chunks), _indices[r+1:prev_r])
                        _chunk = Dagger.@spawn merge_chunks(sink, [_chunk, chunks_to_merge...])
                        prev_r = r
                    end
                    push!(merged_chunks, _chunk)
                    _indices[l] = length(merged_chunks)
                    l += 1
                elseif _lengths[l] + _lengths[r] <= chunksize # merge possible, mark r for merging
                    _lengths[l] += _lengths[r]
                    r -= 1
                end
            end
            idx[k] = _indices[1:r]
        end
        return idx, merged_chunks

    else # no merge
        return idx, chunks
    end
end

ext: jl


text:
mutable struct InputFields
    taux   :: AbstractArray{Float64, 2}
    tauy   :: AbstractArray{Float64, 2}
    nswflx :: AbstractArray{Float64, 2}
    swflx  :: AbstractArray{Float64, 2}
    sumflx :: AbstractArray{Float64, 2}
    ifrac  :: AbstractArray{Float64, 2}
    frwflx :: AbstractArray{Float64, 2}
    vsflx  :: AbstractArray{Float64, 2}
    qflx_T :: AbstractArray{Float64, 2}
    qflx_S :: AbstractArray{Float64, 2}
    Tclim  :: AbstractArray{Float64, 2}
    Sclim  :: AbstractArray{Float64, 2}
    IFRACclim  :: AbstractArray{Float64, 2}
    h_ML   :: AbstractArray{Float64, 2}
end

function InputFields(datakind::Symbol, Nx::Integer, Ny::Integer)

    return InputFields(
        allocate(datakind, Float64, Nx, Ny),
        allocate(datakind, Float64, Nx, Ny),
        allocate(datakind, Float64, Nx, Ny),
        allocate(datakind, Float64, Nx, Ny),
        allocate(datakind, Float64, Nx, Ny),
        allocate(datakind, Float64, Nx, Ny),
        allocate(datakind, Float64, Nx, Ny),
        allocate(datakind, Float64, Nx, Ny),
        allocate(datakind, Float64, Nx, Ny),
        allocate(datakind, Float64, Nx, Ny),
        allocate(datakind, Float64, Nx, Ny),
        allocate(datakind, Float64, Nx, Ny),
        allocate(datakind, Float64, Nx, Ny),
        allocate(datakind, Float64, Nx, Ny),
    )
end

function SubInputFields(
    in_flds :: InputFields,
    rngs...
)
    return InputFields(
        view( in_flds.taux,            rngs...),
        view( in_flds.tauy,            rngs...),
        view( in_flds.nswflx,          rngs...),
        view( in_flds.swflx,           rngs...),
        view( in_flds.sumflx,          rngs...),
        view( in_flds.ifrac,           rngs...),
        view( in_flds.frwflx,          rngs...),
        view( in_flds.vsflx,           rngs...),
        view( in_flds.qflx_T,          rngs...),
        view( in_flds.qflx_S,          rngs...),
        view( in_flds.Tclim,           rngs...),
        view( in_flds.Sclim,           rngs...),
        view( in_flds.IFRACclim,       rngs...),
        view( in_flds.h_ML,            rngs...),
    )
end

function copyfrom!(
    dst::InputFields,
    src::InputFields,
)

    for fld in (:taux, :tauy, :nswflx, :swflx, :sumflx, :ifrac, :frwflx, :vsflx, :qflx_T, :qflx_S, :Tclim, :Sclim, :IFRACclim, :h_ML)
        getfield(dst, fld)[:] = getfield(src, fld)
    end

    #=
    dst.taux[:] = src.taux
    dst.tauy[:] = src.tauy
    dst.nswflx[:] = src.nswflx
    dst.swflx[:] = src.swflx
    dst.sumflx[:] = src.sumflx
    dst.ifrac[:] = src.ifrac
    dst.frwflx[:] = src.frwflx
    dst.qflx[:] = src.qflx
    dst.h_ML[:] = src.h_ML
    =#
end

max_stars_repo_name: meteorologytoday/IOM

max_issues_repo_name: meteorologytoday/IOM

max_forks_repo_name: meteorologytoday/IOM


text:
using Flux
using Zygote
using CUDA

X = rand(Float32,128,128,2,10)
Y = rand(Float32,128,128,1,10)

model = Conv((3,3),2=>1,pad=(1,1),relu)

X = gpu(X)
Y = gpu(Y)
model = gpu(model)

loss(x,y) = Flux.mse(model(x),y)

train_loader = Flux.Data.DataLoader(X,Y, batchsize=2,shuffle=true)

opt = ADAM(0.001)
@time Flux.train!(loss, Flux.params(model), train_loader, opt)


text:
using StaticArrays, DomainSets, Test
import DomainSets: MappedDomain, similar_interval

# test_specific_domains.jl

const v = TypeFactory{SVector}()

const io = IOBuffer()

struct DerivedUnitHyperBall<: DerivedDomain{SVector{2,Float64}}
    superdomain :: Domain

    DerivedUnitHyperBall() = new(2UnitDisk())
end

@testset "Specific domains" begin
    @testset "empty space" begin
        d1 = EmptySpace()
        show(io,d1)
        @test isempty(d1)
        @test String(take!(io)) == "the empty space with eltype Float64"
        @test eltype(d1) == Float64
        @test 0.5 ∉ d1
        @test d1 ∩ d1 == d1
        @test d1 ∪ d1 == d1
        d2 = 0..1
        @test d1 ∩ d2 == d1
        @test d1 ∪ d2 == d2
        @test d2 \ d1 == d2
        @test d2 \ d2 == d1

        d2 = EmptySpace(SVector{2,Float64})
        @test isempty(d2)
        @test v[0.1,0.2] ∉ d2
        @test !approx_in(v[0.1,0.2], d2)
    end

    @testset "full space" begin
        d1 = FullSpace{Float64}()
        show(io,d1)
        @test String(take!(io)) == "the full space with eltype Float64"
        @test 0.5 ∈ d1
        @test d1 ∪ d1 == d1
        @test d1 ∩ d1 == d1
        d2 = 0..1
        @test d1 ∪ d2 == d1
        @test d1 ∩ d2 == d2
        @test typeof(FullSpace(0..1)+1) <: FullSpace
        @test typeof(FullSpace(0..1)*3) <: FullSpace

        d2 = FullSpace{SVector{2,Float64}}()
        @test v[0.1,0.2] ∈ d2
        @test approx_in(v[0.1,0.2], d2)
        @test !isempty(d2)

        @test d2 == Domain(SVector{2,Float64})
        @test d2 == convert(Domain,SVector{2,Float64})
        @test d2 == convert(Domain{SVector{2,Float64}}, SVector{2,Float64})
    end

    @testset "points" begin
        d = Domain(1.0)
        @test d isa Point
        @test 1 ∈ d
        @test 1.1 ∉ d
        @test approx_in(1.1, d, 0.2)
        @test !approx_in(1.2, d, 0.1)
        @test !isempty(d)

        @test d+1 == Domain(2.0)
        @test 1+d == Domain(2.0)
        @test 1-d == Domain(0.0)
        @test d-1 == Domain(0.0)
        @test 2d  == Domain(2.0)
        @test d*2 == Domain(2.0)
        @test d/2 == Domain(0.5)
        @test 2\d == Domain(0.5)

        d1 = Domain(Set([1,2,3]))
        d2 = Point(1) ∪ Point(2) ∪ Point(3)

        @test d1 == d2

        @test convert(Domain{Float64}, Point(1)) ≡ Point(1.0)
        @test Number(Point(1)) ≡ convert(Number, Point(1)) ≡ convert(Int, Point(1)) ≡ 1
    end

    @testset "intervals" begin
        T = Float64
        @testset "ClosedInterval{$T}" begin
            d = zero(T)..one(T)
            @test approx_in(-0.1, d, 0.2)
            @test approx_in(1.1, d, 0.2)
            @test !approx_in(-0.2, d, 0.1)
            @test !approx_in(1.2, d, 0.1)

            @test iscompact(d)
            @test typeof(similar_interval(d, one(T), 2*one(T))) == typeof(d)

            @test leftendpoint(d) ∈ ∂(d)
            @test rightendpoint(d) ∈ ∂(d)
        end
        @testset "UnitInterval{$T}" begin
            d = UnitInterval{T}()
            @test leftendpoint(d) == zero(T)
            @test rightendpoint(d) == one(T)
            @test minimum(d) == infimum(d) == leftendpoint(d)
            @test maximum(d) == supremum(d) == rightendpoint(d)

            @test isclosed(d)
            @test !DomainSets.isopen(d)
            @test iscompact(d)

            @test convert(Domain, d) ≡ d
            @test convert(Domain{T}, d) ≡ d
            @test convert(AbstractInterval, d) ≡ d
            @test convert(AbstractInterval{T}, d) ≡ d
            @test convert(UnitInterval, d) ≡ d
            @test convert(UnitInterval{T}, d) ≡ d
            @test convert(Domain{Float64}, d) ≡ UnitInterval()
            @test convert(AbstractInterval{Float64}, d) ≡ UnitInterval()
            @test convert(UnitInterval{Float64}, d) ≡ UnitInterval()
        end
        @testset "ChebyshevInterval{$T}" begin
            d = ChebyshevInterval{T}()
            @test leftendpoint(d) == -one(T)
            @test rightendpoint(d) == one(T)
            @test minimum(d) == infimum(d) == leftendpoint(d)
            @test maximum(d) == supremum(d) == rightendpoint(d)

            @test isclosed(d)
            @test !DomainSets.isopen(d)
            @test iscompact(d)

            @test convert(Domain, d) ≡ d
            @test convert(Domain{T}, d) ≡ d
            @test convert(AbstractInterval, d) ≡ d
            @test convert(AbstractInterval{T}, d) ≡ d
            @test convert(ChebyshevInterval, d) ≡ d
            @test convert(ChebyshevInterval{T}, d) ≡ d
            @test convert(Domain{Float64}, d) ≡ ChebyshevInterval()
            @test convert(AbstractInterval{Float64}, d) ≡ ChebyshevInterval()
            @test convert(ChebyshevInterval{Float64}, d) ≡ ChebyshevInterval()
        end
        @testset "HalfLine{$T}" begin
            d = HalfLine{T}()
            @test leftendpoint(d) == zero(T)
            @test rightendpoint(d) == T(Inf)
            @test minimum(d) == infimum(d) == leftendpoint(d)
            @test supremum(d) == rightendpoint(d)
            @test_throws ArgumentError maximum(d)

            @test !isclosed(d)
            @test !DomainSets.isopen(d)
            @test !iscompact(d)
            @test 1. ∈ d
            @test -1. ∉ d
            @test approx_in(-0.1, d, 0.5)
            @test !approx_in(-0.5, d, 0.1)
            @test similar_interval(d, T(0), T(Inf)) == d

            @test leftendpoint(d) ∈ ∂(d)
            @test rightendpoint(d) ∉ ∂(d)
        end
        @testset "NegativeHalfLine{$T}" begin
            d = NegativeHalfLine{T}()
            @test leftendpoint(d) == -T(Inf)
            @test rightendpoint(d) == zero(T)
            @test infimum(d) == leftendpoint(d)
            @test supremum(d) == rightendpoint(d)
            @test_throws ArgumentError minimum(d)
            @test_throws ArgumentError maximum(d)

            @test !isclosed(d)
            @test DomainSets.isopen(d)
            @test !iscompact(d)
            @test -1. ∈ d
            @test 1. ∉ d
            @test approx_in(0.5, d, 1.)
            @test !approx_in(0.5, d, 0.4)
            @test similar_interval(d, T(-Inf), T(0)) == d
        end

        @testset "OpenInterval{$T}" begin
            d = OpenInterval(0,1)

            @test leftendpoint(d) ∉ ∂(d)
            @test rightendpoint(d) ∉ ∂(d)
        end

        @testset "Integer intervals" begin
            d = 0..1
            @test leftendpoint(d) ∈ ∂(d)
            @test rightendpoint(d) ∈ ∂(d)

            d = Interval{:open,:closed}(0,1)
            @test leftendpoint(d) ∉ ∂(d)
            @test rightendpoint(d) ∈ ∂(d)

            d = Interval{:closed,:open}(0,1)
            @test leftendpoint(d) ∈ ∂(d)
            @test rightendpoint(d) ∉ ∂(d)
        end

        @test typeof(UnitInterval{Float64}(0.0..1.0)) <: UnitInterval
        @test typeof(ChebyshevInterval{Float64}(-1.0..1.0)) <: ChebyshevInterval

        ## Some mappings preserve the interval structure
        # Translation
        d = zero(T)..one(T)
        @test d == +d

        d2 = d + one(T)
        @test typeof(d2) == typeof(d)
        @test leftendpoint(d2) == one(T)
        @test rightendpoint(d2) == 2*one(T)

        d2 = one(T) + d
        @test typeof(d2) == typeof(d)
        @test leftendpoint(d2) == one(T)
        @test rightendpoint(d2) == 2*one(T)

        d2 = d - one(T)
        @test typeof(d2) == typeof(d)
        @test leftendpoint(d2) == -one(T)
        @test rightendpoint(d2) == zero(T)

        d2 = -d
        @test typeof(d2) == typeof(d)
        @test leftendpoint(d2) == -one(T)
        @test rightendpoint(d2) == zero(T)

        d2 = one(T) - d
        @test d2 == d

        # translation for UnitInterval
        # Does a shifted unit interval return an interval?
        d = UnitInterval{T}()
        d2 = d + one(T)
        @test typeof(d2) <: AbstractInterval
        @test leftendpoint(d2) == one(T)
        @test rightendpoint(d2) == 2*one(T)

        d2 = one(T) + d
        @test typeof(d2) <: AbstractInterval
        @test leftendpoint(d2) == one(T)
        @test rightendpoint(d2) == 2*one(T)

        d2 = d - one(T)
        @test typeof(d2) <: AbstractInterval
        @test leftendpoint(d2) == -one(T)
        @test rightendpoint(d2) == zero(T)

        d2 = -d
        @test typeof(d2) <: AbstractInterval
        @test leftendpoint(d2) == -one(T)
        @test rightendpoint(d2) == zero(T)

        d2 = one(T) - d
        @test typeof(d2) <: AbstractInterval
        @test leftendpoint(d2) == zero(T)
        @test rightendpoint(d2) == one(T)

        # translation for ChebyshevInterval
        d = ChebyshevInterval{T}()
        d2 = d + one(T)
        @test typeof(d2) <: AbstractInterval
        @test leftendpoint(d2) == zero(T)
        @test rightendpoint(d2) == 2*one(T)

        d2 = one(T) + d
        @test typeof(d2) <: AbstractInterval
        @test leftendpoint(d2) == zero(T)
        @test rightendpoint(d2) == 2*one(T)

        d2 = d - one(T)
        @test typeof(d2) <: AbstractInterval
        @test leftendpoint(d2) == -2one(T)
        @test rightendpoint(d2) == zero(T)

        @test -d == d

        d2 = one(T) - d
        @test typeof(d2) <: AbstractInterval
        @test leftendpoint(d2) == zero(T)
        @test rightendpoint(d2) == 2one(T)

        # Scaling
        d = zero(T)..one(T)
        d3 = T(2) * d
        @test typeof(d3) == typeof(d)
        @test leftendpoint(d3) == zero(T)
        @test rightendpoint(d3) == T(2)

        d3 = d * T(2)
        @test typeof(d3) == typeof(d)
        @test leftendpoint(d3) == zero(T)
        @test rightendpoint(d3) == T(2)

        d = zero(T)..one(T)
        d4 = d / T(2)
        @test typeof(d4) == typeof(d)
        @test leftendpoint(d4) == zero(T)
        @test rightendpoint(d4) == T(1)/T(2)

        d4 = T(2) \ d
        @test typeof(d4) == typeof(d)
        @test leftendpoint(d4) == zero(T)
        @test rightendpoint(d4) == T(1)/T(2)

        # Equality
        @test ChebyshevInterval() == ClosedInterval(-1.0,1.0) == ClosedInterval(-1,1)
        @test ChebyshevInterval() ≠ Interval{:closed,:open}(-1.0,1.0)
        @test ChebyshevInterval() ≈ ClosedInterval(-1.0,1.0) ≈ Interval{:closed,:open}(-1.0,1.0)

        # Union and intersection of intervals
        i1 = zero(T)..one(T)
        i2 = one(T)/3 .. one(T)/2
        i3 = one(T)/2 .. 2*one(T)
        i4 = T(2) .. T(3)
        # - union of completely overlapping intervals
        du1 = i1 ∪ i2
        @test typeof(du1) <: AbstractInterval
        @test leftendpoint(du1) == leftendpoint(i1)
        @test rightendpoint(du1) == rightendpoint(i1)

        # - intersection of completely overlapping intervals
        du2 = i1 ∩ i2
        @test typeof(du2) <: AbstractInterval
        @test leftendpoint(du2) == leftendpoint(i2)
        @test rightendpoint(du2) == rightendpoint(i2)

        # - union of partially overlapping intervals
        du3 = i1 ∪ i3
        @test typeof(du3) <: AbstractInterval
        @test leftendpoint(du3) == leftendpoint(i1)
        @test rightendpoint(du3) == rightendpoint(i3)

        # - intersection of partially overlapping intervals
        du4 = i1 ∩ i3
        @test typeof(du4) <: AbstractInterval
        @test leftendpoint(du4) == leftendpoint(i3)
        @test rightendpoint(du4) == rightendpoint(i1)

        # - union of non-overlapping intervals
        du5 = UnionDomain(i1) ∪ UnionDomain(i4)
        @test typeof(du5) <: UnionDomain

        # - intersection of non-overlapping intervals
        du6 = i1 ∩ i4
        @test isempty(du6)

        # - setdiff of intervals
        d1 = -2one(T).. 2one(T)
        @test d1 \ (3one(T) .. 4one(T)) == d1
        @test d1 \ (zero(T) .. one(T)) == UnionDomain((-2one(T)..zero(T))) ∪ UnionDomain((one(T).. 2one(T)))
        @test d1 \ (zero(T) .. 3one(T)) == (-2one(T) .. zero(T))
        @test d1 \ (-3one(T) .. zero(T)) == (zero(T) .. 2one(T))
        @test d1 \ (-4one(T) .. -3one(T)) == d1
        @test d1 \ (-4one(T) .. 4one(T)) == EmptySpace{T}()

        d1 \ (-3one(T)) == d1
        d1 \ (-2one(T)) == Interval{:open,:closed}(-2one(T),2one(T))
        d1 \ (2one(T)) == Interval{:closed,:open}(-2one(T),2one(T))
        d1 \ zero(T) == UnionDomain(Interval{:closed,:open}(-2one(T),zero(T))) ∪ UnionDomain(Interval{:open,:closed}(zero(T),2one(T)))

        # - empty interval
        @test isempty(one(T)..zero(T))
        @test zero(T) ∉ (one(T)..zero(T))
        @test isempty(Interval{:open,:open}(zero(T),zero(T)))
        @test zero(T) ∉ Interval{:open,:open}(zero(T),zero(T))
        @test isempty(Interval{:open,:closed}(zero(T),zero(T)))
        @test zero(T) ∉ Interval{:open,:closed}(zero(T),zero(T))
        @test isempty(Interval{:closed,:open}(zero(T),zero(T)))
        @test zero(T) ∉ Interval{:closed,:open}(zero(T),zero(T))

        d = one(T) .. zero(T)
        @test_throws ArgumentError minimum(d)
        @test_throws ArgumentError maximum(d)
        @test_throws ArgumentError infimum(d)
        @test_throws ArgumentError supremum(d)

        # Subset relations of intervals
        @test issubset((zero(T)..one(T)), (zero(T).. 2*one(T)))
        @test issubset((zero(T)..one(T)), (zero(T).. one(T)))
        @test issubset(OpenInterval(zero(T),one(T)), zero(T) .. one(T))
        @test !issubset(zero(T) .. one(T), OpenInterval(zero(T), one(T)))
        @test issubset(UnitInterval{T}(), ChebyshevInterval{T}())

        ## Disable iteration tests until semantics of iteration over domains are defined
        # # - iteration over intervals
        # ## for integers
        # I = ClosedInterval{Int}(1,1)
        # @test cardinality(I) == 1
        # I = ClosedInterval{Int}(1,5)
        # @test cardinality(I) == 5
        # I = OpenInterval{Int}(1,1)
        # @test cardinality(I) == 0
        # I = OpenInterval{Int}(1,5)
        # @test cardinality(I) == 3
        # ## for floats
        # I = ClosedInterval(1.0,1.0)
        # @test cardinality(I) == 1
        # I = OpenInterval(1.0,1.0)
        # @test cardinality(I) == 0
        # I = ClosedInterval(0.0,1e-316)
        # @test cardinality(I) == 20240226
        # I = OpenInterval(0.0,1e-316)
        # @test cardinality(I) == 20240224

        # - convert
        d = zero(T).. one(T)
        @test d ≡ Interval(zero(T), one(T))
        @test d ≡ ClosedInterval(zero(T), one(T))

        @test convert(Domain, d) ≡ d
        @test convert(Domain{Float32}, d) ≡ (0f0 .. 1f0)
        @test convert(Domain{Float64}, d) ≡ (0.0 .. 1.0)
        @test convert(Domain, zero(T)..one(T)) ≡ d
        @test convert(Domain{T}, zero(T)..one(T)) ≡ d
        @test convert(AbstractInterval, zero(T)..one(T)) ≡ d
        @test convert(AbstractInterval{T}, zero(T)..one(T)) ≡ d
        @test convert(Interval, zero(T)..one(T)) ≡ d
        @test Interval(zero(T)..one(T)) ≡ d
        @test convert(ClosedInterval, zero(T)..one(T)) ≡ d
        @test ClosedInterval(zero(T)..one(T)) ≡ d
        @test convert(ClosedInterval{T}, zero(T)..one(T)) ≡ d
        @test ClosedInterval{T}(zero(T)..one(T)) ≡ d

        @testset "conversion from other types" begin
            @test convert(Domain{T}, 0..1) ≡ d
            @test convert(AbstractInterval{T}, 0..1) ≡ d
            @test convert(ClosedInterval{T}, 0..1) ≡ d
            @test ClosedInterval{T}(0..1) ≡ d
        end
    end

    @testset "unit ball" begin
        D = UnitDisk()
        @test v[1.,0.] ∈ D
        @test v[1.,1.] ∉ D
        @test !isempty(D)

        D = 2UnitDisk()
        @test v[1.4, 1.4] ∈ D
        @test v[1.5, 1.5] ∉ D
        @test typeof(1.2*D)==typeof(D*1.2)
        @test v[1.5,1.5] ∈ 1.2*D
        @test v[1.5,1.5] ∈ D*1.2
        @test !isempty(D)

        D = 2UnitDisk() + v[1.0,1.0]
        @test v[2.4, 2.4] ∈ D
        @test v[3.5, 2.5] ∉ D
        @test !isempty(D)

        B = UnitBall()
        @test v[1.,0.0,0.] ∈ B
        @test v[1.,0.1,0.] ∉ B
        @test !isempty(B)

        B = 2UnitBall()
        @test v[1.9,0.0,0.0] ∈ B
        @test v[0,-1.9,0.0] ∈ B
        @test v[0.0,0.0,-1.9] ∈ B
        @test v[1.9,1.9,0.0] ∉ B
        @test !isempty(B)

        B = 2.0UnitBall() + v[1.0,1.0,1.0]
        @test v[2.9,1.0,1.0] ∈ B
        @test v[1.0,-0.9,1.0] ∈ B
        @test v[1.0,1.0,-0.9] ∈ B
        @test v[2.9,2.9,1.0] ∉ B
        @test !isempty(B)
    end

    @testset "derived unit ball" begin
        B = DerivedUnitHyperBall()
        @test v[1.4, 1.4] ∈ B
        @test v[1.5, 1.5] ∉ B
        @test typeof(1.2*B)==typeof(B*1.2)
        @test v[1.5,1.5] ∈ 1.2*B
        @test v[1.5,1.5] ∈ B*1.2
        @test DomainSets.supereltype(B) == eltype(2UnitDisk())
    end

    @testset "cube" begin
        #Square
        D = UnitInterval()^2
        @test v[0.9, 0.9] ∈ D
        @test v[1.1, 1.1] ∉ D
        @test !isempty(D)

        @test approx_in(v[-0.1,-0.1], D, 0.1)
        @test !approx_in(v[-0.1,-0.1], D, 0.09)

        #Cube
        D = (-1.5 .. 2.2) × (0.5 .. 0.7) × (-3.0 .. -1.0)
        @test v[0.9, 0.6, -2.5] ∈ D
        @test v[0.0, 0.6, 0.0] ∉ D
    end

    @testset "sphere" begin
        C = UnitCircle()
        @test v[1.,0.] ∈ C
        @test v[1.,1.] ∉ C
        @test approx_in(v[1.,0.], C)
        @test !approx_in(v[1.,1.], C)
        @test !isempty(C)

        C = 2UnitCircle() + v[1.,1.]
        @test approx_in(v[3.,1.], C)

        C = UnitCircle() + v[1,1]
        @test approx_in(v[2,1], C)

        S = UnitSphere()
        @test v[1.,0.,0.] ∈ S
        @test v[1.,0.,1.] ∉ S
        @test approx_in(v[cos(1.),sin(1.),0.], S)
        @test !isempty(S)

        S = 2UnitSphere() + v[1.,1.,1.]
        @test approx_in(v[1. + 2*cos(1.),1. + 2*sin(1.),1.], S)
        @test !approx_in(v[4.,1.,5.], S)

        # Create an ellipse, the curve
        E = ellipse(2.0, 4.0)
        @test v[2.0,0.0] ∈ E
        @test v[0.0,4.0] ∈ E
        @test v[2.0+1e-10,0.0] ∉ E
        @test v[0.0,0.0] ∉ E
        E = ellipse(1, 2.0)
        @test eltype(E) == SVector{2,Float64}

        # Create an ellipse, the domain with volume
        E2 = ellipse_shape(2.0, 4.0)
        @test v[2.0,0.0] ∈ E2
        @test v[0.0,4.0] ∈ E2
        @test v[2.0+1e-10,0.0] ∉ E2
        @test v[0.0,0.0] ∈ E2
        @test v[1.0,1.0] ∈ E2

        E2 = ellipse_shape(1, 2.0)
        @test eltype(E) == SVector{2,Float64}

        D = UnitCircle()
        @test convert(Domain{SVector{2,BigFloat}}, D) ≡ UnitCircle{BigFloat}()
        @test SVector(1,0) in D
        @test SVector(nextfloat(1.0),0) ∉ D

        D = UnitSphere()
        @test convert(Domain{SVector{3,BigFloat}}, D) ≡ UnitSphere{BigFloat}()
        @test SVector(1,0,0) in D
        @test SVector(nextfloat(1.0),0,0) ∉ D
    end

    @testset "mapped_domain" begin
        # Test chaining of maps
        D = UnitCircle()
        D1 = 2*D
        @test typeof(D1) <: MappedDomain
        @test typeof(source(D1)) <: UnitHyperSphere
        D2 = 2*D1
        @test typeof(source(D2)) <: UnitHyperSphere

        D = UnitInterval()^2
        show(io,rotate(D,1.))
        @test String(take!(io)) == "A mapped domain based on 0.0..1.0 (Unit) x 0.0..1.0 (Unit)"

        D = rotate(UnitInterval()^2, π)
        @test v[-0.9, -0.9] ∈ D
        @test v[-1.1, -1.1] ∉ D

        D = rotate(UnitInterval()^2, π, v[-.5,-.5])
        @test v[-1.5, -1.5] ∈ D
        @test v[-0.5, -0.5] ∉ D

        D = rotate(UnitInterval()^3 + v[-.5,-.5,-.5], pi, pi, pi)
        @test v[0.4, 0.4, 0.4] ∈ D
        @test v[0.6, 0.6, 0.6] ∉ D

        D = rotate((-1.5.. 2.2) × (0.5 .. 0.7) × (-3.0 .. -1.0), π, π, π, v[.35, .65, -2.])
        @test v[0.9, 0.6, -2.5] ∈ D
        @test v[0.0, 0.6, 0.0] ∉ D
    end

    @testset "simplex" begin
        d = UnitSimplex{2}()
        # We test a point in the interior, a point on each of the boundaries and
        # all corners.
        @test v[0.2,0.2] ∈ d
        @test v[0.0,0.2] ∈ d
        @test v[0.2,0.0] ∈ d
        @test v[0.5,0.5] ∈ d
        @test v[0.0,0.0] ∈ d
        @test v[1.0,0.0] ∈ d
        @test v[0.0,1.0] ∈ d
        # And then some points outside
        @test v[0.6,0.5] ∉ d
        @test v[0.5,0.6] ∉ d
        @test v[-0.2,0.2] ∉ d
        @test v[0.2,-0.2] ∉ d

        @test approx_in(v[-0.1,-0.1], d, 0.1)
        @test !approx_in(v[-0.1,-0.1], d, 0.09)

        d3 = UnitSimplex{3,BigFloat}()
        x0 = big(0.0)
        x1 = big(1.0)
        x2 = big(0.3)
        @test v[x0,x0,x0] ∈ d3
        @test v[x1,x0,x0] ∈ d3
        @test v[x0,x1,x0] ∈ d3
        @test v[x0,x0,x1] ∈ d3
        @test v[x2,x0,x0] ∈ d3
        @test v[x0,x2,x0] ∈ d3
        @test v[x0,x0,x2] ∈ d3
        @test v[x2,x2,x2] ∈ d3
        @test v[-x2,x2,x2] ∉ d3
        @test v[x2,-x2,x2] ∉ d3
        @test v[x2,x2,-x2] ∉ d3
        @test v[x1,x1,x1] ∉ d3
    end

    @testset "arithmetics" begin
        D = UnitInterval()^3
        S = 2UnitBall()
        @testset "joint domain" begin
            DS = D ∪ S
            @test v[0.0, 0.6, 0.0] ∈ DS
            @test v[0.9, 0.6,-2.5] ∉ DS
        end

        @testset "domain intersection" begin
            DS = D ∩ S
            @test v[0.1, 0.1, 0.1] ∈ DS
            @test v[0.1, -0.1, 0.1] ∉ DS
        end
        @testset "domain difference" begin
            DS = D\S
            @test v[0.1, 0.1, 0.1] ∉ DS

            D1 = 2*D
            D2 = D*2
            D3 = D/2

            @test v[2., 2., 2.] ∈ D1
            @test v[0.9, 0.6,-2.5] ∉ D1
            @test v[2., 2., 2.] ∈ D2
            @test v[0.9, 0.6,-2.5] ∉ D2
            @test v[.5, .4, .45] ∈ D3
            @test v[.3, 0.6,-.2] ∉ D3
        end
    end

    @testset "cartesian product" begin
        @testset "ProductDomain 1" begin
            T1 = (-1.0 .. 1.0)^2
            @test v[0.5,0.5] ∈ T1
            @test v[-1.1,0.3] ∉ T1

            T1 = cartesianproduct((-1.0 .. 1.0), 2)
            @test v[0.5,0.5] ∈ T1
            @test v[-1.1,0.3] ∉ T1

            T2 = (-1.0 .. 1.0) × (-1.5 .. 2.5)
            @test v[0.5,0.5] ∈ T2
            @test v[-1.1,0.3] ∉ T2
        end
        @testset "ProductDomain 2" begin
            T1 = cartesianproduct((-1.0 .. 1.0), 2)

            T3 = ProductDomain(1.05UnitDisk(), (-1.0 .. 1.0))
            @test v[0.5,0.5,0.8] ∈ T3
            @test v[-1.1,0.3,0.1] ∉ T3

            T4 = T1×(-1.0..1.)
            @test v[0.5,0.5,0.8] ∈ T4
            @test v[-1.1,0.3,0.1] ∉ T4

            T5 = (-1.0..1.)×T1
            @test v[0.,0.5,0.5] ∈ T5
            @test v[0.,-1.1,0.3] ∉ T5

            T6 = T1×T1
            @test v[0.,0.,0.5,0.5] ∈ T6
            @test v[0.,0.,-1.1,0.3] ∉ T6

            io = IOBuffer()
            show(io,T1)
            @test String(take!(io)) == "-1.0..1.0 x -1.0..1.0"
        end
        @testset "mixed intervals" begin
            d = (0..1) × (0.0..1)
            @test v[0.1,0.2] ∈ d
            @test d isa EuclideanDomain{2}
        end
    end
    @testset "embedded" begin
        i = 0.0 .. 1.0
        e = embedding_map(Float64, Complex{Float64})
        r = restriction_map(Complex{Float64}, Float64)
        ei = DomainSets.forwardmap_domain(e, i)

        @test 0.5+1im ∉ ei
        @test 0.5+0im ∈ ei
        @test 0.5-0im ∈ ei
        @test 1.1+0im ∉ ei
        @test approx_in(0.8+1e-5im, ei, 1e-5)
        @test !approx_in(0.8+1e-5im, ei, 1e-6)
    end
end

@testset "Set operations" begin
    @testset "union" begin
        d1 = UnitDisk()
        d2 = (-.9..0.9)^2
        d3 = (-.5 .. -.1) × (.5 .. 0.1)
        d4 = (0.0..1.5)
        d5 = [1.0,3.0]

        @test isempty(d3)
        @test convert(Domain{SVector{2,Float64}}, d3) isa Domain{SVector{2,Float64}}

        u1 = d1 ∪ d2
        u2 = u1 ∪ d3

        u3 = d3 ∪ u1
        u4 = u1 ∪ u2
        x = SVector(0.,.15)
        y = SVector(1.1,.75)
        @test x∈u3
        @test x∈u4

        @test y∉u3
        @test y∉u4

        ũ1 = UnionDomain(d1,d2)
        @test u1 == ũ1
        ũ1 = UnionDomain((d1,d2))
        @test u1 == ũ1
        ũ2 = UnionDomain([d1,d2])
        @test ũ2 == ũ2
        @test u1 == ũ2

        # Don't create a union with two identical elements
        @test (d1 ∪ d1) isa typeof(d1)

        # union with non-Domain type that implements domain interface
        u45 = UnionDomain(d4, d5)
        @test u45 isa Domain{Float64}
        @test 0.2 ∈ u45
        @test 1.2 ∈ u45
        @test -1.2 ∉ u45
        @test convert(Domain{BigFloat}, u45) isa Domain{BigFloat}

        # ordering doesn't matter
        @test UnionDomain(d1,d2) == UnionDomain(d2,d1)

        @test UnionDomain(UnionDomain(d1,d2),d3) == UnionDomain(d3,UnionDomain(d1,d2))

        @test !isempty(u1)

        show(io,u1)
        @test String(take!(io)) == "a union of 2 domains:\n\t1.\t: the 2-dimensional unit ball\n\t2.\t: -0.9..0.9 x -0.9..0.9\n"
    end

    @testset "intersection" begin
        d1 = UnitDisk()
        d2 = (-.4..0.4)^2
        d3 = (-.5 .. 0.5) × (-.1.. 0.1)
        d4 = (0.0..1.5)
        d5 = [1.0,3.0]

        # intersection of productdomains
        i1 = d2 & d3
        show(io,i1)
        @test String(take!(io)) == "-0.4..0.4 x -0.1..0.1"
        i2 = d1 & d2
        show(io,i2)
        @test String(take!(io)) == "the intersection of 2 domains:\n\t1.\t: the 2-dimensional unit ball\n\t2.\t: -0.4..0.4 x -0.4..0.4\n"

        i3 = d3 & i2
        i4 = i2 & d3
        i5 = i3 & i2

        x = SVector(0.,.05)
        y = SVector(0.,.75)
        @test x∈i3
        @test x∈i4

        @test y∉i3
        @test y∉i4

        d45 = IntersectionDomain(d4, d5)
        @test d45 isa Domain{Float64}
        @test 1.0 ∈ d45
        @test 1.1 ∉ d45
        @test convert(Domain{BigFloat}, d45) isa Domain{BigFloat}
    end

    @testset "difference" begin
        d1 = UnitDisk()
        d2 = (-.5..0.5) × (-.1..0.1)
        d3 = 0.0..3.0
        d4 = [1.0, 2.5]

        # intersection of productdomains
        d = d1\d2
        show(io,d)
        @test String(take!(io)) == "the difference of 2 domains:\n\t1.\t: the 2-dimensional unit ball\n\t2.\t: -0.5..0.5 x -0.1..0.1\n"

        x = SVector(0.,.74)
        y = SVector(0.,.25)
        @test x∈d
        @test x∈d

        d34 = DifferenceDomain(d3, d4)
        @test d34 isa Domain{Float64}
        @test d34 isa DifferenceDomain
        @test 0.99 ∈ d34
        @test 1.0 ∉ d34
        @test convert(Domain{BigFloat}, d34) isa Domain{BigFloat}
    end

    @testset "arithmetic" begin
        d1 = (0..1)
        d2 = (2..3)
        d = UnionDomain(d1) ∪ UnionDomain(d2)

        @test d+1 == UnionDomain(d1+1) ∪ (d2+1)
        @test d-1 == UnionDomain(d1-1) ∪ (d2-1)
        @test 2d  == UnionDomain(2d1)  ∪ (2d2)
        @test d*2 == UnionDomain(d1*2) ∪ (d2*2)
        @test d/2 == UnionDomain(d1/2) ∪ (d2/2)
        @test 2\d == UnionDomain(2\d1) ∪ (2\d2)

        @test infimum(d) == minimum(d) == 0
        @test supremum(d) == maximum(d) == 3
    end

    @testset "different types" begin
        d̃1 = (0..1)
        d1 = (0f0.. 1f0)
        d2 = (2..3)

        @test UnionDomain(d1) ∪ d2 == UnionDomain(d̃1) ∪ d2
    end
end


text:
using Metatheory
using Test
using Metatheory.Library
using Metatheory.EGraphs
using Metatheory.Util
using Metatheory.EGraphs.Schedulers

@metatheory_init

# ================= TEST PATEQUIV =====================

# :foo => :zoo ⟺ :foo in same class as :bar

lhs = PatEquiv(@pat(:foo), @pat(:bar))
rhs = @pat :zoo
q = RewriteRule(lhs, rhs)

g = EGraph()
fooclass, _ = addexpr!(g, :foo)
barclass, _ = addexpr!(g, :bar)
zooclass, _ = addexpr!(g, :zoo)

# display(g.classes); println()
@test !(in_same_class(g, fooclass, zooclass))

saturate!(g, [q])

# display(g.classes); println()
@test !(in_same_class(g, fooclass, zooclass))

merge!(g, fooclass.id, zooclass.id)

# display(g.classes); println()
@test in_same_class(g, fooclass, zooclass)


text:
"""
solve the 1D linear advection equation
```math
∂q/∂t=−u∂q/∂x
```
in a periodic domain, where ``q`` is the quantity being advected,
``t`` is time, ``x`` is the spatial coordinate and ``u`` is the velocity,
which is constant with ``x``.
"""
function lax_wendroff!(nt::Int, c, q_init::AbstractVector{T}, q::AbstractVector{T}) where T
    nx = length(q)
    flux = zeros(T, nx-1)   # Fluxes between boxes
    @inbounds for i=1:nx
        q[i] = q_init[i] # Initialize q
    end
    @inbounds for j=1:nt  # Main loop in time
        for i=1:nx-1
            flux[i] = 0.5*c*(q[i]+q[i+1]+c*(q[i]-q[i+1]))
        end
        for i=2:nx-1
            q[i] += flux[i-1]-flux[i]
        end
        q[1] = q[nx-1]; q[nx] = q[2] # Treat boundary conditions
    end
    return q
end

using Random
Random.seed!(2)
q_init = randn(100)
q = zeros(100)
@show lax_wendroff!(2000, 1.0, q_init, zero(q_init))
using BenchmarkTools
@benchmark lax_wendroff!(2000, 1.0, $q_init, x) setup=(x=zero(q_init))
@time lax_wendroff!(2000, 1.0, q_init, q)

using NiLang
@i function i_lax_wendroff!(nt::Int, c, q_init::AbstractVector{T}, q::AbstractVector{T},
        cache::AbstractMatrix{T}) where T
    nx ← length(q)
    @inbounds for i=1:nx
        q[i] += q_init[i] # Initialize q
    end
    @inbounds for j=1:nt  # Main loop in time
        for i=1:nx-1
            @routine begin
                @zeros T anc1 anc2 anc3
                anc1 += 0.5 * c
                anc2 += q[i] - q[i+1]
                anc3 += q[i] + q[i+1]
                anc3 += c * anc2
            end
            cache[i,j] += anc1 * anc3
            ~@routine
        end
        for i=2:nx-1
            q[i] += cache[i-1,j]-cache[i,j]
        end
        # Treat boundary conditions
        cache[nx,j] += q[nx-1]
        SWAP(q[1], cache[nx,j])
        cache[nx+1,j] += q[2]
        SWAP(q[nx], cache[nx+1,j])
    end
end
nt = 2000
i_lax_wendroff!(nt, 1.0, q_init, zero(q_init), zeros(length(q_init)+1,nt))


text:
using JWebImageDemo
using FileIO
using Test

const TEST_IMAGE = joinpath(@__DIR__, "resources/green_python.jpg")
NUM = 5

@testset "JWebImageDemo.jl" begin
    img = FileIO.load(TEST_IMAGE)
    @time img_out = JWebImageDemo.permute_image(img, NUM)
    #save(Stream(format"PNG",io), data...)
    @test size(img) != size(img_out) ||
          any(((i, o),) -> i != o, zip(iterate(img), iterate(img_out)))
    save("out.jpg", img_out)
end


text:
export σmap

"""
    σmap(s)

Map from S³ into ℝ³ using stereographic projection with the given point `s`.
"""
σmap(q::Quaternion) = ℝ³(vec(q)[1:3] ./ (1 - vec(q)[4]))
σmap(s::S³) = σmap(Quaternion(s))


text:
#!/usr/bin/env julia

using Pkg

Pkg.generate("ICI_simulateur")
Pkg.activate("ICI_simulateur")

Pkg.add(Pkg.PackageSpec(;name="CSV", version="0.9.9"))
Pkg.add(Pkg.PackageSpec(;name="JSON", version="0.21.2"))
Pkg.add(Pkg.PackageSpec(;name="DataFrames", version="1.2.2"))
Pkg.add(Pkg.PackageSpec(;name="Distributions", version="0.25.22"))
Pkg.add(Pkg.PackageSpec(;name="GraphIO", version="0.6.0"))
Pkg.add(Pkg.PackageSpec(;name="LightGraphs", version="1.3.5"))
Pkg.add(Pkg.PackageSpec(;name="MetaGraphs", version="0.6.6"))
Pkg.add(url="https://github.com/zhangliye/GraphMLReader.jl")
Pkg.add(Pkg.PackageSpec(;name="MPI", version="0.19.1"))
Pkg.add(Pkg.PackageSpec(;name="SpecialFunctions", version="1.8.1"))
Pkg.add(Pkg.PackageSpec(;name="StatsBase", version="0.33.12"))
Pkg.add(Pkg.PackageSpec(;name="PackageCompiler", version="1.4.1"))
Pkg.add(Pkg.PackageSpec(;name="Statistics"))
Pkg.add(Pkg.PackageSpec(;name="TimerOutputs"))
Pkg.add(Pkg.PackageSpec(;name="Random"))
Pkg.add(Pkg.PackageSpec(;name="Dates"))
Pkg.add(Pkg.PackageSpec(;name="DelimitedFiles"))


text:
module Modes
import Roots: find_zero, Order2
import Cubature: hcubature, hquadrature
import LinearAlgebra: dot, norm
import NumericalIntegration: integrate, Trapezoidal
import Luna: Maths, Grid
import Luna.PhysData: c, ε_0, μ_0
import Memoize: @memoize
import LinearAlgebra: mul!
import DSP: unwrap

export dimlimits, neff, β, α, losslength, transmission, dB_per_m, dispersion, zdw, field, Exy, Aeff, @delegated, @arbitrary, chkzkwarg

"""
    AbstractMode

Abstract type representing a single mode of a waveguide.
"""
abstract type AbstractMode end

ModeCollection = Union{Tuple{Vararg{T} where T <: Modes.AbstractMode},
                       AbstractArray{T} where T <: Modes.AbstractMode}

# make modes broadcast like a scalar
Broadcast.broadcastable(m::AbstractMode) = Ref(m)

"""
    dimlimits(m::AbstractMode; z=0.0)

Maximum dimensional limits of validity for mode `m` at position `z`.
"""
function dimlimits end

"""
    field(m::AbstractMode, xs; z=0.0)

Get the field components `(Ex, Ey)`` at position `xs`, `z`
"""
function field end

"""
    field(m::AbstractMode; z=0)

Create function of coords that returns (xs) -> (Ex, Ey) for the mode `m`.
"""
field(m::AbstractMode; z=0) =  (xs) -> field(m, xs, z=z)

"""
    N(m::AbstractMode; z=0.0)

Get mode normalization constant of mode `m` at longitudinal position `z`.
"""
@memoize function N(m::AbstractMode; z=0.0)
    # we memoize this so it is only called once for each mode and z position
    f = field(m, z=z)
    dl = dimlimits(m, z=z)
    function Nfunc(xs)
        E = f(xs)
        ret = sqrt(ε_0/μ_0)*dot(E, E)
        dl[1] == :polar ? xs[1]*ret : ret
    end
    val, err = hcubature(Nfunc, dl[2], dl[3])
    0.5*abs(val)
end

"""
    Exy(m::AbstractMode, xs; z=0.0)

Get the normalised field components at position `xs`, `z`.
"""
Exy(m::AbstractMode, xs; z=0.0) = field(m, xs, z=z) ./ sqrt(N(m, z=z))

"""
    Exy(m::AbstractMode; z=0.0)

Create function that returns normalised (xs) -> (Ex, Ey)"""
Exy(m::AbstractMode; z=0.0) = (xs) -> Exy(m, xs, z=z)

"""
    absE(m::AbstractMode, xs; z=0.0)

Get the field norm ``|E|`` at position `xs`, `z`
"""
absE(m::AbstractMode, xs; z=0.0) = norm(Exy(m, xs, z=z))

"""
    absE(m::AbstractMode; z=0.0)

Create function that returns normalised (xs) -> ``|E|``.
"""
absE(m::AbstractMode; z=0.0) = (xs) -> absE(m, xs, z=z)

"""
    Aeff(m::AbstractMode; z=0.0)

Get effective area of mode `m` and longitudinal position `z`.

This is the brute-force implementation valid for any mode.
"""
@memoize function Aeff(m::AbstractMode; z=0.0)
    # we memoize this so it is only called once for each mode and z position
    em = absE(m, z=z)
    dl = dimlimits(m, z=z)
    # Numerator
    function Aeff_num(xs)
        e = em(xs)
        dl[1] == :polar ? xs[1]*e^2 : e^2
    end
    val, err = hcubature(Aeff_num, dl[2], dl[3])
    num = val^2
    # Denominator
    function Aeff_den(xs)
        e = em(xs)
        dl[1] == :polar ? xs[1]*e^4 : e^4
    end
    den, err = hcubature(Aeff_den, dl[2], dl[3])
    return num / den
end

"""
    neff(m::AbstractMode, ω; z=0.0)

Get the full complex refractive index of mode `m` at frequency `ω`
and longitudinal position `z`.
"""
function neff end

"""
    β(m::AbstractMode, ω; z=0.0)

Calculate the propagation constant ``β`` for the mode `m` at frequency `ω` and
longitudinal position `z`.
"""
function β(m::AbstractMode, ω; z=0.0)
    return ω/c*real(neff(m, ω, z=z))
end

"""
    α(m::AbstractMode, ω; z=0.0)

Calculate the attenuation constant ``α`` for the mode `m` at frequency `ω` and
longitudinal position `z`.
"""
function α(m::AbstractMode, ω; z=0.0)
    return 2*ω/c*imag(neff(m, ω, z=z))
end

"""
    losslength(m::AbstractMode, ω; z=0.0)

Calculate the losslength (``1/α``) for the mode `m` at frequency `ω` and
longitudinal position `z`.
"""
function losslength(m::AbstractMode, ω; z=0.0)
    return 1/α(m, ω, z=z)
end

"""
    dB_per_m(m::AbstractMode, ω; z=0.0)

Calculate the attenuation in dB/m for the mode `m` at frequency `ω` and
longitudinal position `z`.
"""
function dB_per_m(m::AbstractMode, ω; z=0.0)
    return 10/log(10).*α(m, ω; z=z)
end

"""
    transmission(m::AbstractMode, ω, L)

Calculate the power transmission after propagation through length `L` in the mode `m` for
radiation at the frequency `ω`.
"""
function transmission(m::AbstractMode, ω, L)
    αint, err = hquadrature(0, L) do z
        -α(m, ω; z)
    end
    return exp(αint)
end

"""
    dispersion_func(m::AbstractMode, order; z=0.0)

Get a function `βn(ω)` which returns the dispersion of a given `order` at frequency `ω`.
"""
function dispersion_func(m::AbstractMode, order; z=0.0)
    βn(ω) = Maths.derivative(ω -> β(m, ω, z=z), ω, order)
    return βn
end

"""
    dispersion_func(m::AbstractMode, order; z=0.0)

Calculate the dispersion of a given `order` at frequency `ω`.
"""
function dispersion(m::AbstractMode, order, ω; z=0.0)
    return dispersion_func(m, order, z=z).(ω)
end

"""
    zdw(m::AbstractMode; λmin=100e-9, λmax=3000e-9, z=0.0)

Calculate the zero-dispersion wavelength (ZDW) of mode `m` within the range `ub` to `lb`.
"""
function zdw(m::AbstractMode; λmin=100e-9, λmax=3000e-9, z=0.0)
    ubω = 2π*c/λmin
    lbω = 2π*c/λmax
    ω0 = missing
    try
        ω0 = find_zero(dispersion_func(m, 2, z=z), (lbω, ubω))
    catch
    end
    return 2π*c/ω0
end

"""
    zdw(m::AbstractMode, λ0; z=0.0, rtol=1e-4)

Calculate the zero-dispersion wavelength (ZDW) of mode `m` with an initial guess of `λ0`.

This method is faster than the bounded version if the ZDW is known to be close to `λ0`.
"""
function zdw(m::AbstractMode, λ0; z=0.0, rtol=1e-6)
    ωguess = 2π*c/λ0 * 1e-15 # convert everything find_zero sees to units close to 1.0
    dfun = dispersion_func(m, 2, z=z)
    β2(ω) = 1e30 * dfun(ω*1e15)
    ω0 = missing
    try
        ω0 = 1e15*find_zero(β2, ωguess, Order2(), rtol=rtol)
    catch
    end
    return 2π*c/ω0
end

"""
    β_ret(m::AbstractMode, ω; z=0, λ0)

Calculate the propagation constant ``β`` for mode `m` transformed into a frame which moves with
the group and phase velocity of `m` at wavelength `λ0`.
"""
function β_ret(m::AbstractMode, ω; z=0, λ0)
    ω0 = 2π*c/λ0
    β.(m, ω; z=z) .- dispersion(m, 1, ω0; z=z)*(ω.-ω0) .- β(m, ω0; z=z)
end

"""
    chkzkwarg(func)

Check that function accepts `z` keyword argument and add it if necessary.
"""
function chkzkwarg(func)
    try
        func(2.5e15, z=0.0)
        return func
    catch e
        if isa(e, ErrorException)
            f = (ω; z) -> func(ω)
            return f
        else
            throw(e)
        end
    end
end

"""
    overlap(m::AbstractMode, r, E; dim)

Calculate mode overlap between radially symmetric field and radially symmetric mode.

# Examples
```jldoctest
julia> a = 100e-6;
julia> m = Capillary.MarcatiliMode(a, :He, 1.0);
julia> unm = approx_besselroots(0, 1)[end]
julia> r = collect(range(0, a, length=512));
julia> Er = besselj.(0, unm*r/a);

julia> η = Modes.overlap(m, r, Er; dim=1);
julia> abs2(η[1]) ≈ 1
true
```
"""
function overlap(m::AbstractMode, r, E; dim, norm=true)
    dl = dimlimits(m) # integration limits
    # sample the modal field at the same coords as E - select y polarisation component
    Er = [Exy(m, (ri, 0))[2] for ri in r]  #field [Ex(r, θ), Ey(r, θ)] of the mode
    Er[r .> dl[3][1]] .= 0
    #= normalisation factor - we want the integral of the modal intensity over the waveguide
        to be 1, but the  fields Exy(...) are normalised to produce modal power,
        so they include the factor of cε₀/2 =#
    normEr = 1/sqrt(c*ε_0/2)

    # Generate output array: same shape as input, except length in space is 1
    shape = collect(size(E))
    shape[dim] = 1
    integral = zeros(eltype(E), Tuple(shape)) # make output array

    # Indices to iterate over all other dimensions (e.g. polarisation, frequency)
    idxlo = CartesianIndices(size(E)[1:dim-1])
    idxhi = CartesianIndices(size(E)[dim+1:end])
    for hi in idxhi
        for lo in idxlo
                # normalisation factor for the other field
                normE = norm ? sqrt(2π*integrate(r, r.*abs2.(E[lo, :, hi]), Trapezoidal())) :
                               1/sqrt(c*ε_0/2)
                # E[lo, :, hi] is a vector
                integrand = 2π .* E[lo, :, hi] .* Er.*r./(normE*normEr)
                integral[lo, 1, hi] = integrate(r, integrand, Trapezoidal())
        end
    end
    return integral
end

"""
    overlap(m::AbstractMode, E)

Calculate mode overlap between (analytic) 2D field `E` and mode `m`.
The field function `E(xs)` should return the normalised cartesian vector
components of the field `(Ex, Ey)` as an `SVector` as a function of polar
coordinates `xs = (r,θ)`.
```
"""
function overlap(m::AbstractMode, E)
    dl = dimlimits(m)
    function f(xs)
        0.5*sqrt(ε_0/μ_0)*dot(conj(Exy(m, xs)), E(xs))*xs[1]
    end
    val, err = hcubature(f, dl[2], dl[3]; maxevals=1000)
    abs(val)
end

"""
    overlap(modes::ModeCollection, newgrid, oldgrid, r, Eωr)

Decompose the spatio-spectral field `Eωr`, sampled on radial coordinate `r` and time-grid
`oldgrid`, into the given `modes` and resample onto `newgrid` via cubic interpolation.
"""
function overlap(modes::ModeCollection, newgrid::Grid.RealGrid, oldgrid::Grid.RealGrid, r, Eωr)
    Egm = zeros(ComplexF64, (length(newgrid.ω), length(modes))) # output array
    # If old and new grids have different number of samples, we need to scale by δt1/δt2
    scale = (oldgrid.t[2] - oldgrid.t[1]) / (newgrid.t[2] - newgrid.t[1])
    #= Pulses centred on t=0 have large linear spectral phase components which can confuse
        phase unwrapping and lead to oscillations in the spectral phase after interpolation.
        Here we shift the pulse to t=-t_max of the old grid to remove this, then do the
        interpolation, then shift back to t=0 on the new grid. Note that this preserves any
        actual time shifts on the original pulse =#
    τold = length(oldgrid.t) * (oldgrid.t[2] - oldgrid.t[1])/2
    τnew = length(newgrid.t) * (newgrid.t[2] - newgrid.t[1])/2
    for (midx, mode) in enumerate(modes)
        Eωm = overlap(mode, r, Eωr; norm=false, dim=2)[:, 1]
        Eωm .*= exp.(1im.*oldgrid.ω.*τold) # shift to -t_max before unwrapping
        Aω = abs.(Eωm) # spectral amplitude
        ϕω = unwrap(angle.(Eωm)) # spectral phase
        Ag = Maths.BSpline(oldgrid.ω, Aω).(newgrid.ω)
        ϕg = Maths.BSpline(oldgrid.ω, ϕω).(newgrid.ω) .- newgrid.ω*τnew # shift to t=0
        Egm[:, midx] = scale * Ag .* exp.(1im*ϕg)
    end
    Egm
end

function _overlap(mode1::AbstractMode, mode2::AbstractMode)
    dl = dimlimits(mode1)
    if dl != dimlimits(mode2)
        error("Integration limits for both modes must be the same")
    end
    function f(xs)
        ret = 1/2*sqrt(ε_0/μ_0)*dot(conj(Exy(mode1, xs)), Exy(mode2, xs))
        dl[1] == :polar ? xs[1]*ret : ret
    end
    val, err = hcubature(f, dl[2], dl[3]; maxevals=5000)
end

"""
    overlap(mode1, mode2)

Calculate the normalised overlap between two `AbstractMode`s.
"""
function overlap(mode1::AbstractMode, mode2::AbstractMode)
    val, err = _overlap(mode1, mode2)
    val
end

"""
    orthogonal(mode1, mode2)

Test whether the `AbstractMode`s `mode1` and `mode2` are orthogonal to each other.
"""
function orthogonal(mode1::AbstractMode, mode2::AbstractMode)
    val, err = _overlap(mode1, mode2)
    isapprox(val, 0; atol=max(err, eps(typeof(val))))
end

"""
    orthonormal(modes)

Test whether the `modes` form an orthonormal set.
"""
function orthonormal(modes)
    out = true
    for (ii1, ii2) in Iterators.product(eachindex(modes), eachindex(modes))
        val, err = _overlap(modes[ii1], modes[ii2])
        atol = max(err, eps(typeof(val)))
        this = ii1 == ii2 ? isapprox(val, 1; atol) : isapprox(val, 0; atol)
        out = out && this
    end
    out
end

struct ToSpace{mT,iT}
    ms::mT
    indices::iT
    nmodes::Int
    npol::Int
    Ems::Array{Float64,2}
end

"""
    ToSpace(ms; components=:xy)

Construct a `ToSpace` for high performance conversion between modal fields and real space.

# Arguments
- `ms::Tuple`: a tuple of modes
- `components::Symbol`: which polarisation components to return: :x, :y, :xy
"""
function ToSpace(ms; components=:xy)
    if components == :xy
        indices = 1:2
    elseif components == :x
        indices = 1
    elseif components == :y
        indices = 2
    else
        error("components $components not recognised")
    end
    nmodes = length(ms)
    npol = length(indices)
    Ems = Array{Float64,2}(undef, nmodes, npol)
    ToSpace(ms, indices, nmodes, npol, Ems)
end

"""
    to_space!(Erω, Emω, xs, ts::ToSpace; z=0.0)

Convert from modal fields to real space using provided `ToSpace` struct.

# Arguments
- `Erω::Array{ComplexF64}`: a dimension nω x npol array where the real space frequency domain
                            field will be written to
- `Emω::Array{ComplexF64}`: a dimension nω x nmodes array containing the frequency domain
                            modal fields
- `xs:Tuple`: the transverse coordinates, `x,y` for cartesian, `r,θ` for polar
- `ts::ToSpace`: the corresponding `ToSpace` struct
- `z::Real`: the axial position
"""
@noinline function to_space!(Erω, Emω, xs, ts::ToSpace; z=0.0)
    if ts.nmodes != size(Emω,2)
        error("the number of modes must match the number of modal fields")
    end
    if ts.npol != size(Erω,2)
        error("the number of output fields must match the number of polarisation components")
    end
    # we assume all dimlimits are the same
    dimlims = dimlimits(ts.ms[1], z=z)
    # handle limits
    if dimlims[1] == :cartesian
        # for the cartesian case
        # if either coordinate is outside dimlimits we return 0
        if xs[1] <= dimlims[2][1] || xs[1] >= dimlims[3][1]
            fill!(Erω, 0.0)
            return
        elseif xs[2] <= dimlims[2][2] || xs[2] >= dimlims[3][2]
            fill!(Erω, 0.0)
            return
        end
    elseif dimlims[1] == :polar
        # for the polar case
        # if the r coordinate is negative we error
        if xs[1] < 0.0
            error("polar coordinate r cannot be smaller than 0")
        # if r is greater or equal to the boundary we return 0
        elseif xs[1] >= dimlims[3][1]
            fill!(Erω, 0.0)
            return
        end
    end
    # get the field at x1, x2
    for i = 1:ts.nmodes
        ts.Ems[i,:] .= Exy(ts.ms[i], xs, z=z)[ts.indices] # field matrix (nmodes x npol)
    end
    mul!(Erω, Emω, ts.Ems) # matrix product (nω x nmodes) * (nmodes x npol) -> (nω x npol)
end

"""
    to_space(Emω, xs, ts::ToSpace; z=0.0)

Convert from modal fields to real space using provided `ToSpace` struct.

# Arguments
- `Emω::Array{ComplexF64}`: a dimension nω x nmodes array containing the frequency domain
                            modal fields
- `xs:Tuple`: the transverse coordinates, `x,y` for cartesian, `r,θ` for polar
- `ts::ToSpace`: the corresponding `ToSpace` struct
- `z::Real`: the axial position
"""
function to_space(Emω, xs, ts::ToSpace; z=0.0)
    Erω = Array{ComplexF64,2}(undef, size(Emω,1), ts.npol)
    to_space!(Erω, Emω, xs, ts, z=z)
    Erω
end

"""
    to_space(Emω, xs, ms; components=:xy, z=0.0)

Convert from modal fields to real space.

# Arguments
- `Emω::Array{ComplexF64}`: a dimension nω x nmodes array containing the frequency domain
                            modal fields
- `xs:Tuple`: the transverse coordinates, `x,y` for cartesian, `r,θ` for polar
- `ms::Tuple`: a tuple of modes
- `components::Symbol`: which polarisation components to return: :x, :y, :xy
- `z::Real`: the axial position
"""
function to_space(Emω, xs, ms; components=:xy, z=0.0)
    ts = ToSpace(ms, components=components)
    Erω = Array{ComplexF64,2}(undef, size(Emω,1), ts.npol)
    to_space!(Erω, Emω, xs, ts, z=z)
    Erω
end

struct DelegatedMode{mT, idT} <: AbstractMode
    mode::mT # wrapped mode
    id::idT # unique identifier type to distinguish different DelegatedModes
end

function DelegatedMode(mode)
    DelegatedMode(mode, Val(gensym()))
end

"""
    delegated(mode, kwargs...)

Create a delegated mode, which takes its methods from an existing mode except
for those which are overwritten
# Arguments:
- `mode::AbstractMode`: The wrapped mode to which non-specified methods are delegated
- `kwargs`: functions that override: `neff`, `field`, `dimlimits`, `Aeff`, or `N`.

The functions given should have the **same signature** as the mode methods, i.e. take
an `AbstractMode` as their first argument, **even if** they do not do anything with it. This
is to ensure that the delegated functions can access the data of the wrapped mode if necessary.

To override `neff` with functions for `α` and `β`, create the `neff` function using
[`neff_from_αβ`](@ref).
"""
function delegated(mode; kwargs...)
    dmode = DelegatedMode(mode)
    mT = typeof(dmode) # this is unique for each instance by virtue of the id field
    kw = Dict(kwargs)
    for mfun in (:neff, :field, :dimlimits)
        if haskey(kw, mfun)
            dfun = kw[mfun]
            @eval $mfun(dm::$mT, args...; kwargs...) = $dfun(dm.mode, args...; kwargs...)
        else
            @eval $mfun(dm::$mT, args...; kwargs...) = $mfun(dm.mode, args...; kwargs...)
        end
    end
    for mfun in (:Aeff, :N)
        if haskey(kw, mfun)
            dfun = kw[mfun]
            @eval $mfun(dm::$mT, args...; kwargs...) = $dfun(dm.mode, args...; kwargs...)
        elseif (!haskey(kw, :dimlimits)) && (!haskey(kw, :field))
            #= if dimlimits and field have not been changed, Aeff and N are the same too,
                so we can safely delegate them to the wrapped mode (likely faster) =#
            @eval $mfun(dm::$mT, args...; kwargs...) = $mfun(dm.mode, args...; kwargs...)
        end
    end
    dmode
end

"""
    arbitrary(kwargs...)

Create an arbitrary mode, which takes its methods from given functions.

The functions given should have the same signature as defined `Luna.Modes`, **except** that
the first argument (the `AbstractMode`) is omitted, e.g. for `neff` the function should be
of the form `n(ω; z) = ...`

To define `neff` with functions for `α` and `β`, create the `neff` function using
[`neff_from_αβ`](@ref).
"""
function arbitrary(;kwargs...)
    dmode = DelegatedMode(nothing)
    mT = typeof(dmode)
    kw = Dict(kwargs)
    for mfun in (:neff, :field, :dimlimits)
        if haskey(kw, mfun)
            dfun = kw[mfun]
            @eval $mfun(dm::$mT, args...; kwargs...) = $dfun(args...; kwargs...)
        else
            err = "method $mfun not defined for this mode"
            @eval $mfun(dm::$mT, args...; kwargs...) = error($err)
        end
    end
    for mfun in (:Aeff, :N)
        if haskey(kw, mfun)
            dfun = kw[mfun]
            @eval $mfun(dm::$mT, args...; kwargs...) = $dfun(args...; kwargs...)
        end
    end
    dmode
end

"""
    neff_from_αβ(α, β)

Create a closure converting the functions `α(ω; z)` and `β(ω; z)` into an effective index.
"""
neff_from_αβ(α, β) = (ω; z=0) -> c/ω * (β(ω; z=z) + 0.5im*α(ω; z=z))

end