text:
import Oceananigans: short_show

@inline zerofunction(args...) = 0
@inline onefunction(args...) = 1

T_zerofunction = typeof(zerofunction)
T_onefunction = typeof(onefunction)

short_show(::T_zerofunction) = "0"
short_show(::T_onefunction) = "1"

"""
    struct Relaxation{R, M, T}

Callable object for restoring fields to a `target` at
some `rate` and within a `mask`ed region in `x, y, z`.
"""
struct Relaxation{R, M, T}
      rate :: R
      mask :: M
    target :: T
end

"""
    Relaxation(; rate, mask=onefunction, target=zerofunction)

Returns a `Forcing` that restores a field to `target(x, y, z, t)`
at the specified `rate`, in the region `mask(x, y, z)`.

The functions `onefunction` and `zerofunction` always return 1 and 0, respectively.
Thus the default `mask` leaves the whole domain uncovered, and the default `target` is zero.

Example
=======

* Restore a field to zero on a timescale of "3600" (equal
  to one hour if the time units of the simulation are seconds).

```jldoctest relaxation
using Oceananigans

damping = Relaxation(rate = 1/3600)

# output
Relaxation{Float64, typeof(Oceananigans.Forcings.onefunction), typeof(Oceananigans.Forcings.zerofunction)}
├── rate: 0.0002777777777777778
├── mask: 1
└── target: 0
```

* Restore a field to a linear z-gradient within the bottom 1/4 of a domain
  on a timescale of "60" (equal to one minute if the time units of the simulation
  are seconds).

```jldoctest relaxation
dTdz = 0.001 # ⁰C m⁻¹, temperature gradient

T₀ = 20 # ⁰C, surface temperature at z=0

Lz = 100 # m, depth of domain

bottom_sponge_layer = Relaxation(; rate = 1/60,
                                   target = LinearTarget{:z}(intercept=T₀, gradient=dTdz),
                                   mask = GaussianMask{:z}(center=-Lz, width=Lz/4))

# output
Relaxation{Float64, GaussianMask{:z,Float64}, LinearTarget{:z,Float64}}
├── rate: 0.016666666666666666
├── mask: exp(-(z + 100.0)^2 / (2 * 25.0^2))
└── target: 20.0 + 0.001 * z
```
"""
Relaxation(; rate, mask=onefunction, target=zerofunction) = Relaxation(rate, mask, target)

""" Wrap `forcing::Relaxation` in `ContinuousForcing` and add the appropriate field dependency. """
function regularize_forcing(forcing::Relaxation, field_name, model_field_names)
    continuous_relaxation = ContinuousForcing(forcing, field_dependencies=field_name)
    return regularize_forcing(continuous_relaxation, field_name, model_field_names)
end

@inline (f::Relaxation)(x, y, z, t, field) =
    f.rate * f.mask(x, y, z) * (f.target(x, y, z, t) - field)

@inline (f::Relaxation{R, M, <:Number})(x, y, z, t, field) where {R, M} =
    f.rate * f.mask(x, y, z) * (f.target - field)

"""Show the innards of a `Relaxation` in the REPL."""
Base.show(io::IO, relaxation::Relaxation{R, M, T}) where {R, M, T} =
    print(io, "Relaxation{$R, $M, $T}", '\n',
        "├── rate: $(relaxation.rate)", '\n',
        "├── mask: $(short_show(relaxation.mask))", '\n',
        "└── target: $(short_show(relaxation.target))")

short_show(relaxation::Relaxation) =
    "Relaxation(rate=$(relaxation.rate), mask=$(short_show(relaxation.mask)), target=$(short_show(relaxation.target)))"

#####
##### Sponge layer functions
#####

"""
    GaussianMask{D}(center, width)

Callable object that returns a Gaussian masking function centered on
`center`, with `width`, and varying along direction `D`.

Examples
========

* Create a Gaussian mask centered on `z=0` with width `1` meter.

```julia
julia> mask = GaussianMask{:z}(center=0, width=1)
```
"""
struct GaussianMask{D, T}
    center :: T
     width :: T

    function GaussianMask{D}(; center, width) where D
        T = promote_type(typeof(center), typeof(width))
        return new{D, T}(center, width)
    end
end

@inline (g::GaussianMask{:x})(x, y, z) = exp(-(x - g.center)^2 / (2 * g.width^2))
@inline (g::GaussianMask{:y})(x, y, z) = exp(-(y - g.center)^2 / (2 * g.width^2))
@inline (g::GaussianMask{:z})(x, y, z) = exp(-(z - g.center)^2 / (2 * g.width^2))

show_exp_arg(D, c) = c == 0 ? "$D^2" :
                     c > 0  ? "($D - $c)^2" :
                              "($D + $(-c))^2"

short_show(g::GaussianMask{D}) where D =
    "exp(-$(show_exp_arg(D, g.center)) / (2 * $(g.width)^2))"

#####
##### Linear target functions
#####

"""
    LinearTarget{D}(intercept, gradient)

Callable object that returns a Linear target function
with `intercept` and `gradient`, and varying along direction `D`.

Examples
========

* Create a linear target function varying in `z`, equal to `0` at
  `z=0` and with gradient 10⁻⁶:

```julia
julia> target = LinearTarget{:z}(intercept=0, gradient=1e-6)
```
"""
struct LinearTarget{D, T}
    intercept :: T
     gradient :: T

    function LinearTarget{D}(; intercept, gradient) where D
        T = promote_type(typeof(gradient), typeof(intercept))
        return new{D, T}(intercept, gradient)
    end
end

@inline (p::LinearTarget{:x})(x, y, z, t) = p.intercept + p.gradient * x
@inline (p::LinearTarget{:y})(x, y, z, t) = p.intercept + p.gradient * y
@inline (p::LinearTarget{:z})(x, y, z, t) = p.intercept + p.gradient * z

short_show(l::LinearTarget{:x}) = "$(l.intercept) + $(l.gradient) * x"
short_show(l::LinearTarget{:y}) = "$(l.intercept) + $(l.gradient) * y"
short_show(l::LinearTarget{:z}) = "$(l.intercept) + $(l.gradient) * z"


text:
# Load modules
using Plots, LaTeXStrings
pyplot()
using ThinFilmsTools
using Optim

##
function bilayer_reflectance(beam, incident, emergent)
    # Create a measured spectrum of reflection
    layers2 = [ LayerTMMO(incident),
                LayerTMMO(RI.looyenga([0.70 0.3],[incident emergent]); d=300.),
                LayerTMMO(RIdb.glass(beam.λ./1e3); d=250.),
                LayerTMMO(emergent) ]
    sol = tmm_optics(beam, layers2)
    return vec(beam.p.*sol.Spectra.Rp .+ (1.0 - beam.p).*sol.Spectra.Rs)
end
##

# Wavelength range [nm]
λ = 250:900
# Angle of incidence [degrees]
θ = [5.]
# Polarisation (1.0 = p, 0.0 = s, between 0.0 and 1.0 = average)
pol = 0.5
beam = PlaneWave(λ, θ; p=pol)

# Refractive indices of incident (0) and substrate (2)
incident = RIdb.air(beam.λ)
emergent = RIdb.silicon(beam.λ)

# Define the RI model to use
layers = [
    LayerTMMO(incident),
    ModelFit(:looyenga; N=(ninc=incident, nhost=emergent)),
    ModelFit(:sellmeier),
    LayerTMMO(emergent)
]

# Create absolute reflectance spectrum to fit
Rexp = bilayer_reflectance(beam, incident, emergent)
plot(beam.λ, Rexp_norm)
gui()

options = Optim.Options(
    g_abstol=1e-8, g_reltol=1e-8, iterations=10^6, show_trace=true, store_trace=true,
);

seed = [
    [300, 0.7],
    vcat(250.0, [1.0, 0.23, 1.0, 6e-3, 2e-2, 1e-2]),
]

solOptim = fit_tmm_optics(
    Reflectance(Rexp), seed, beam, layers;
    options=options, alg=SAMIN(), lb=0.5.*seed, ub=1.5.*seed,
)

plot(FitSpectrum(), solOptim.beam.λ, solOptim.spectrumExp, solOptim.spectrumFit)
gui()

# julia> solOptim.objfunMin
# 8.487618158368491e-9
#
# julia> solOptim.optParams
# 2-element Array{Array{Float64,1},1}:
#  [300.06736895487506, 0.700027712313205]
#  [249.97204958549517, 0.5080567934627047, 0.11979359414683158, 0.6358886116061814, 0.008572122867045586, 0.010000126624594001, 0.010114195402292319]

ext: jl


text:
module CUFFT

using ..APIUtils

using ..CUDA
using ..CUDA: CUstream, cuComplex, cuDoubleComplex, libraryPropertyType
using ..CUDA: libcufft, unsafe_free!, @retry_reclaim, initialize_context

using CEnum: @cenum

using Reexport: @reexport

# core library
include("libcufft_common.jl")
include("error.jl")
include("libcufft.jl")

# low-level wrappers
include("util.jl")
include("wrappers.jl")

# high-level integrations
include("fft.jl")

end


text:
module GeoPhyInv

# include modules (note: due to dependencies, order is important!)
include("Interpolation/Interpolation.jl")
include("Utils/Utils.jl")
include("Operators.jl")
include("Smooth.jl")
include("IO.jl")
include("Models/Models.jl")
include("Acquisition/Acquisition.jl")
include("Coupling.jl")
include("Data/Data.jl")
include("Born/Born.jl")
include("Fdtd/Fdtd.jl")
include("FWI/FWI.jl")
include("Poisson/Poisson.jl")
include("Gallery/Gallery.jl")
include("Plots.jl")

# export the Expt for Seismic Forward Modelling
const SeisForwExpt=GeoPhyInv.Fdtd.Param
export SeisForwExpt
mod!(a::SeisForwExpt)=GeoPhyInv.Fdtd.mod!(a)

# export the Expt for Seismic Inversion
const SeisInvExpt=GeoPhyInv.FWI.Param
export SeisInvExpt

# export the Expt for Poisson
const PoissonExpt=GeoPhyInv.Poisson.ParamExpt
export PoissonExpt
mod!(a::PoissonExpt,b,c)=GeoPhyInv.Poisson.mod!(a,b,c)
mod!(a::PoissonExpt,b)=GeoPhyInv.Poisson.mod!(a,b)
mod!(a::PoissonExpt)=GeoPhyInv.Poisson.mod!(a)
operator_Born(a::PoissonExpt,b)=GeoPhyInv.Poisson.operator_Born(a,b)

export mod!, operator_Born

const JuMIT=GeoPhyInv
export JuMIT

const GIPh=GeoPhyInv
export GIPh

const J=GeoPhyInv
export J

const GInterp=GeoPhyInv.Interpolation
export GInterp

const JP=GeoPhyInv.Plots
export JP

const JF=GeoPhyInv.FWI
export JF

const JD=GeoPhyInv.Data
export JD

const JA=GeoPhyInv.Acquisition
export JA

const JG=GeoPhyInv.Gallery
export JG

end # module


text:
@testset "default options" begin
    res0 = [get_interval(
        [3., 2., 2.1],
        i,
        (x::Vector{Float64}) -> f_3p_1im_dep(x),
        :CICO_ONE_PASS;
        loss_crit = 9.,
        silent = true
    ) for i in 1:3]

    @test isapprox(res0[1].result[1].value, 1.0, atol=1e-2)
    @test isapprox(res0[1].result[2].value, 5.0, atol=1e-2)
    @test length(res0[1].result[1].profilePoints) > 0
    @test length(res0[1].result[2].profilePoints) > 0
    @test res0[1].result[1].status == :BORDER_FOUND_BY_SCAN_TOL
    @test res0[1].result[2].status == :BORDER_FOUND_BY_SCAN_TOL
    @test res0[1].result[1].direction == :left
    @test res0[1].result[2].direction == :right
    @test isapprox(res0[2].result[1].value, 2.0-2.0*sqrt(2.), atol=1e-2)
    @test isapprox(res0[2].result[2].value, 2.0+2.0*sqrt(2.), atol=1e-2)
    @test res0[2].result[1].status == :BORDER_FOUND_BY_SCAN_TOL
    @test res0[2].result[2].status == :BORDER_FOUND_BY_SCAN_TOL
    @test res0[3].result[1].status == :SCAN_BOUND_REACHED
    @test res0[3].result[2].status == :SCAN_BOUND_REACHED
end


text:
using SparseArrays
using ViewerGL
GL = ViewerGL
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation

cubohemioctahedron = Lar.cubohemioctahedron()

V,EV,FV,CV = Lar.struct2lar(cubohemioctahedron)
GL.VIEW([ GL.GLGrid(V,FV, GL.Point4d(1,1,1,0.1)) ,GL.GLFrame]);

cop_EV = Lar.coboundary_0(EV::Lar.Cells);
cop_FE = Lar.coboundary_1(V, FV::Lar.Cells, EV::Lar.Cells);
W = convert(Lar.Points, V');

V, copEV, copFE, copCF = Lar.Arrangement.spatial_arrangement(W::Lar.Points, cop_EV::Lar.ChainOp, cop_FE::Lar.ChainOp);


text:
module SPICE

    # ---- Modules required
    using Distributions
    using DataFrames
    using Sobol
    using Compat
    using StatsBase
    using Statistics
    using LinearAlgebra
    using Random
    using CSV

    # ---- Imports & includes
    import Base.show, Base.copy
    include("SystemCore/System.jl")
    include("SSA/SSA.jl")
    include("Main/Main.jl")
    include("Sim/sim.jl")

    # ---- Export types
    export CEM,
    Model,
    DataSets,
    TauOpt,
    System,
    SystemState,

    # ---- Export functions
    est,
    estimate,
    resetSystem!,
    sim
end

max_stars_repo_name: pzuliani/SPICE

max_issues_repo_name: pzuliani/SPICE

max_forks_repo_name: pzuliani/SPICE


text:
using CodeTracking, JuliaInterpreter, Test
using JuliaInterpreter: enter_call, enter_call_expr, get_return, @lookup
using Base.Meta: isexpr
include("utils.jl")

const ALL_COMMANDS = (:n, :s, :c, :finish, :nc, :se, :si, :until)

function step_through_command(fr::Frame, cmd::Symbol)
    while true
        ret = JuliaInterpreter.debug_command(JuliaInterpreter.finish_and_return!, fr, cmd)
        ret == nothing && break
        fr, pc = ret
    end
    @test fr.callee === nothing
    @test fr.caller === nothing
    return get_return(fr)
end

function step_through_frame(frame_creator)
    rets = []
    for cmd in ALL_COMMANDS
        frame = frame_creator()
        ret = step_through_command(frame, cmd)
        push!(rets, ret)
    end
    @test all(ret -> ret == rets[1], rets)
    return rets[1]
end
step_through(f, args...; kwargs...) = step_through_frame(() -> enter_call(f, args...; kwargs...))
step_through(expr::Expr) = step_through_frame(() -> enter_call_expr(expr))

@generated function generatedfoo(T)
    :(return $T)
end
callgenerated() = generatedfoo(1)
@generated function generatedparams(a::Array{T,N}) where {T,N}
    :(return ($T,$N))
end
callgeneratedparams() = generatedparams([1 2; 3 4])

macro insert_some_calls()
    esc(quote
        x = sin(b)
        y = asin(x)
        z = sin(y)
    end)
end

trivial(x) = x

struct B{T} end

# Putting this into a @testset introduces a closure that breaks the kwprep detection
function complicated_keyword_stuff(a, b=-1; x=1, y=2)
    a == a
    (a, b, :x=>x, :y=>y)
end
function complicated_keyword_stuff_splatargs(args...; x=1, y=2)
    args[1] == args[1]
    (args..., :x=>x, :y=>y)
end
function complicated_keyword_stuff_splatkws(a, b=-1; kw...)
    a == a
    (a, b, kw...)
end
function complicated_keyword_stuff_splat2(args...; kw...)
    args[1] == args[1]
    (args..., kw...)
end

# @testset "Debug" begin
    @testset "Basics" begin
        frame = enter_call(map, x->2x, 1:10)
        @test debug_command(frame, :finish) === nothing
        @test frame.caller === frame.callee === nothing
        @test get_return(frame) == map(x->2x, 1:10)

        for func in (complicated_keyword_stuff, complicated_keyword_stuff_splatargs,
                     complicated_keyword_stuff_splatkws, complicated_keyword_stuff_splat2)
            for (args, kwargs) in (((1,), ()), ((1, 2), (x=7, y=33)))
                oframe = frame = enter_call(func, args...; kwargs...)
                frame = JuliaInterpreter.maybe_step_through_kwprep!(frame, false)
                frame = JuliaInterpreter.maybe_step_through_wrapper!(frame)
                @test any(stmt->isa(stmt, Expr) && JuliaInterpreter.hasarg(isequal(QuoteNode(==)), stmt.args), frame.framecode.src.code)
                f, pc = debug_command(frame, :n)
                @test f === frame
                @test isa(pc, Int)
                @test oframe.callee !== nothing
                @test debug_command(frame, :finish) === nothing
                @test oframe.caller === oframe.callee === nothing
                @test get_return(oframe) == func(args...; kwargs...)

                @test @interpret(complicated_keyword_stuff(args...; kwargs...)) == complicated_keyword_stuff(args...; kwargs...)
            end
        end

        let f22() = string(:(a+b))
            @test step_through(f22) == "a + b"
        end
        let f22() = string(QuoteNode(:a))
            @test step_through(f22) == ":a"
        end

        frame = enter_call(trivial, 2)
        @test debug_command(frame, :s) === nothing
        @test get_return(frame) == 2

        @test step_through(trivial, 2) == 2
        @test step_through(:($(+)(1,2.5))) == 3.5
        @test step_through(:($(sin)(1))) == sin(1)
        @test step_through(:($(gcd)(10,20))) == gcd(10, 20)
    end

    @testset "until" begin
        function f_with_lines(s)
            sin(2.0)
            cos(2.0)
            for i in 1:100
                s += i
            end
            sin(2.0)
        end
        meth_def = @__LINE__() - 8

        frame = enter_call(f_with_lines, 0)
        @test whereis(frame)[2] == meth_def + 1
        debug_command(frame, :until)
        @test whereis(frame)[2] == meth_def + 2
        debug_command(frame, :until; line=(meth_def + 4))
        @test whereis(frame)[2] == meth_def + 4
        debug_command(frame, :until; line=(meth_def + 6))
        @test whereis(frame)[2] == meth_def + 6
    end

    @testset "generated" begin
        frame = enter_call_expr(:($(callgenerated)()))
        f, pc = debug_command(frame, :s)
        @test isa(pc, BreakpointRef)
        @test JuliaInterpreter.scopeof(f).name == :generatedfoo
        stmt = JuliaInterpreter.pc_expr(f)
        @test JuliaInterpreter.is_return(stmt) && JuliaInterpreter.lookup_return(frame, stmt) === Int
        @test debug_command(frame, :c) === nothing
        @test frame.callee === nothing
        @test get_return(frame) === Int
        # This time, step into the generated function itself
        frame = enter_call_expr(:($(callgenerated)()))
        f, pc = debug_command(frame, :sg)
            # Aside: generators can have `Expr(:line, ...)` in their line tables, test that this is OK
            lt = JuliaInterpreter.linetable(f, 2)
            @test isexpr(lt, :line) || isa(lt, Core.LineInfoNode)
        @test isa(pc, BreakpointRef)
        @test JuliaInterpreter.scopeof(f).name == :generatedfoo
        stmt = JuliaInterpreter.pc_expr(f)
        @test JuliaInterpreter.is_return(stmt) && JuliaInterpreter.lookup_return(f, stmt) === 1
        f2, pc = debug_command(f, :finish)
        @test JuliaInterpreter.scopeof(f2).name == :callgenerated
        # Now finish the regular function
        @test debug_command(frame, :finish) === nothing
        @test frame.callee === nothing
        @test get_return(frame) === 1

        # Parametric generated function (see #157)
        frame = fr = JuliaInterpreter.enter_call(callgeneratedparams)
        while fr.pc < JuliaInterpreter.nstatements(fr.framecode) - 1
            fr, pc = debug_command(fr, :se)
        end
        fr, pc = debug_command(fr, :sg)
        @test JuliaInterpreter.scopeof(fr).name == :generatedparams
        fr, pc = debug_command(fr, :finish)
        @test debug_command(fr, :finish) === nothing
        @test JuliaInterpreter.get_return(fr) == (Int, 2)
    end

    @testset "Optional arguments" begin
        function optional(n = sin(1))
            x = asin(n)
            cos(x)
        end
        frame = JuliaInterpreter.enter_call_expr(:($(optional)()))
        # Step through the wrapper
        f = JuliaInterpreter.maybe_step_through_wrapper!(frame)
        @test frame !== f
        # asin(n)
        f, pc = debug_command(f, :n)
        # cos(1.0)
        f, pc = debug_command(f, :n)
        # return
        @test debug_command(f, :n) === nothing
    end

    @testset "Keyword arguments" begin
        f(x; b = 1) = x+b
        g() = f(1; b = 2)
        frame = JuliaInterpreter.enter_call_expr(:($(g)()));
        fr, pc = debug_command(frame, :nc)
        fr, pc = debug_command(fr, :nc)
        fr, pc = debug_command(fr, :nc)
        fr, pc = debug_command(fr, :s)
        fr, pc = debug_command(fr, :finish)
        @test debug_command(fr, :finish) === nothing
        @test frame.callee === nothing
        @test get_return(frame) == 3

        frame = JuliaInterpreter.enter_call(f, 2; b = 4)
        fr = JuliaInterpreter.maybe_step_through_wrapper!(frame)
        fr, pc = debug_command(fr, :nc)
        debug_command(fr, :nc)
        @test get_return(frame) == 6
    end

    @testset "Optional + keyword wrappers" begin
        opkw(a, b=1; c=2, d=3) = 1
        callopkw1() = opkw(0)
        callopkw2() = opkw(0, -1)
        callopkw3() = opkw(0; c=-2)
        callopkw4() = opkw(0, -1; c=-2)
        callopkw5() = opkw(0; c=-2, d=-3)
        callopkw6() = opkw(0, -1; c=-2, d=-3)
        scopes = Method[]
        for f in (callopkw1, callopkw2, callopkw3, callopkw4, callopkw5, callopkw6)
            frame = fr = JuliaInterpreter.enter_call(f)
            pc = fr.pc
            while pc <= JuliaInterpreter.nstatements(fr.framecode) - 2
                fr, pc = debug_command(fr, :se)
            end
            fr, pc = debug_command(frame, :si)
            @test stacklength(frame) == 2
            frame = fr = JuliaInterpreter.enter_call(f)
            pc = fr.pc
            while pc <= JuliaInterpreter.nstatements(fr.framecode) - 2
                fr, pc = debug_command(fr, :se)
            end
            fr, pc = debug_command(frame, :s)
            @test stacklength(frame) > 2
            push!(scopes, JuliaInterpreter.scopeof(fr))
        end
        @test length(unique(scopes)) == 1  # all get to the body method
    end

    @testset "Macros" begin
        # Work around the fact that we can't detect macro expansions if the macro
        # is defined in the same file
        include_string(Main, """
        function test_macro()
            a = sin(5)
            b = asin(a)
            @insert_some_calls
            z
        end
        ""","file.jl")
        frame = JuliaInterpreter.enter_call_expr(:($(test_macro)()))
        f, pc = debug_command(frame, :n)        # a is set
        f, pc = debug_command(f, :n)            # b is set
        f, pc = debug_command(f, :n)            # x is set
        f, pc = debug_command(f, :n)            # y is set
        f, pc = debug_command(f, :n)            # z is set
        @test debug_command(f, :n) === nothing  # return
    end

    @testset "Quoting" begin
        # Test that symbols don't get an extra QuoteNode
        f_symbol() = :limit => true
        frame = JuliaInterpreter.enter_call(f_symbol)
        fr, pc = debug_command(frame, :s)
        fr, pc = debug_command(fr, :finish)
        @test debug_command(fr, :finish) === nothing
        @test get_return(frame) == f_symbol()
    end

    @testset "Varargs" begin
        f_va_inner(x) = x + 1
        f_va_outer(args...) = f_va_inner(args...)
        frame = fr = JuliaInterpreter.enter_call(f_va_outer, 1)
        # depending on whether this is in or out of a @testset, the first statement may differ
        stmt1 = fr.framecode.src.code[1]
        if isexpr(stmt1, :call) && @lookup(frame, stmt1.args[1]) === getfield
            fr, pc = debug_command(fr, :se)
        end
        fr, pc = debug_command(fr, :s)
        fr, pc = debug_command(fr, :n)
        @test root(fr) !== fr
        fr, pc = debug_command(fr, :finish)
        @test debug_command(fr, :finish) === nothing
        @test get_return(frame) === 2
    end

    @testset "ASTI#17" begin
        function (::B)(y)
            x = 42*y
            return x + y
        end
        B_inst = B{Int}()
        step_through(B_inst, 10) == B_inst(10)
    end

    @testset "Exceptions" begin
        # Don't break on caught exceptions
        err_caught = Any[nothing]
        function f_exc_outer()
            try
                f_exc_inner()
            catch err;
                err_caught[1] = err
            end
            x = 1 + 1
            return x
        end
        f_exc_inner() = error()
        fr = JuliaInterpreter.enter_call(f_exc_outer)
        fr, pc = debug_command(fr, :s)
        fr, pc = debug_command(fr, :n)
        fr, pc = debug_command(fr, :n)
        debug_command(fr, :finish)
        @test get_return(fr) == 2
        @test first(err_caught) isa ErrorException
        @test stacklength(fr) == 1

        err_caught = Any[nothing]
        fr = JuliaInterpreter.enter_call(f_exc_outer)
        fr, pc = debug_command(fr, :s)
        debug_command(fr, :c)
        @test get_return(root(fr)) == 2
        @test first(err_caught) isa ErrorException
        @test stacklength(root(fr)) == 1

        # Rethrow on uncaught exceptions
        f_outer() = g_inner()
        g_inner() = error()
        fr = JuliaInterpreter.enter_call(f_outer)
        @test_throws ErrorException debug_command(fr, :finish)
        @test stacklength(fr) == 3

        # Break on error
        try
            break_on(:error)
            fr = JuliaInterpreter.enter_call(f_outer)
            fr, pc = debug_command(JuliaInterpreter.finish_and_return!, fr, :finish)
            @test fr.framecode.scope.name == :error

            fundef() = undef_func()
            frame = JuliaInterpreter.enter_call(fundef)
            fr, pc = debug_command(frame, :s)
            @test isa(pc, BreakpointRef)
            @test pc.err isa UndefVarError
        finally
            break_off(:error)
        end
    end

    @testset "breakpoints" begin
        # In source breakpoints
        function f_bp(x)
            #=1=#    i = 1
            #=2=#    @label foo
            #=3=#    @bp
            #=4=#    repr("foo")
            #=5=#    i += 1
            #=6=#    i > 3 && return x
            #=7=#    @goto foo
        end
        ln = @__LINE__
        method_start = ln - 9
        fr = enter_call(f_bp, 2)
        @test JuliaInterpreter.linenumber(fr) == method_start + 1
        fr, pc =  JuliaInterpreter.debug_command(fr, :c)
        # Hit the breakpoint x1
        @test JuliaInterpreter.linenumber(fr) == method_start + 3
        @test pc isa BreakpointRef
        fr, pc =  JuliaInterpreter.debug_command(fr, :n)
        @test JuliaInterpreter.linenumber(fr) == method_start + 4
        fr, pc =  JuliaInterpreter.debug_command(fr, :c)
        # Hit the breakpoint again x2
        @test pc isa BreakpointRef
        @test JuliaInterpreter.linenumber(fr) == method_start + 3
        fr, pc =  JuliaInterpreter.debug_command(fr, :c)
        # Hit the breakpoint for the last time x3
        @test pc isa BreakpointRef
        @test JuliaInterpreter.linenumber(fr) == method_start + 3
        JuliaInterpreter.debug_command(fr, :c)
        @test get_return(fr) == 2
    end

    f_inv(x::Real) = x^2;
    f_inv(x::Integer) = 1 + invoke(f_inv, Tuple{Real}, x)
    @testset "invoke" begin
        fr = JuliaInterpreter.enter_call(f_inv, 2)
        fr, pc = JuliaInterpreter.debug_command(fr, :s) # apply_type
        frame, pc = JuliaInterpreter.debug_command(fr, :s) # step into invoke
        @test frame.framecode.scope.sig == Tuple{typeof(f_inv),Real}
        JuliaInterpreter.debug_command(frame, :c)
        frame = root(frame)
        @test get_return(frame) == f_inv(2)
    end

    f_inv_latest(x::Real) = 1 + (@static isdefined(Core, :_call_latest) ? Core._call_latest(f_inv, x) : Core._apply_latest(f_inv, x))
    @testset "invokelatest" begin
        fr = JuliaInterpreter.enter_call(f_inv_latest, 2.0)
        fr, pc = JuliaInterpreter.debug_command(fr, :nc)
        frame, pc = JuliaInterpreter.debug_command(fr, :s) # step into invokelatest
        @test frame.framecode.scope.sig == Tuple{typeof(f_inv),Real}
        JuliaInterpreter.debug_command(frame, :c)
        frame = root(frame)
        @test get_return(frame) == f_inv_latest(2.0)
    end

    @testset "Issue #178" begin
        remove()
        a = [1, 2, 3, 4]
        @breakpoint length(LinearIndices(a))
        frame, bp = @interpret sum(a)
        @test debug_command(frame, :c) === nothing
        @test get_return(frame) == sum(a)
    end

    @testset "Stepping over kwfunc preparation" begin
        stepkw! = JuliaInterpreter.maybe_step_through_kwprep! # for brevity
        a = [4, 1, 3, 2]
        reversesort(x) = sort(x; rev=true)
        frame = JuliaInterpreter.enter_call(reversesort, a)
        frame = stepkw!(frame)
        @test frame.pc == JuliaInterpreter.nstatements(frame.framecode) - 1

        scopedreversesort(x) = Base.sort(x, rev=true)  # https://github.com/JuliaDebug/Debugger.jl/issues/141
        frame = JuliaInterpreter.enter_call(scopedreversesort, a)
        frame = stepkw!(frame)
        @test frame.pc == JuliaInterpreter.nstatements(frame.framecode) - 1

        frame = JuliaInterpreter.enter_call(sort, a)
        frame = stepkw!(frame)
        @test frame.pc == JuliaInterpreter.nstatements(frame.framecode) - 1

        frame, pc = debug_command(frame, :s)
        frame, pc = debug_command(frame, :se)  # get past copymutable
        frame = stepkw!(frame)
        @test frame.pc > 4

        frame = JuliaInterpreter.enter_call(sort, a; rev=true)
        frame, pc = debug_command(frame, :se)
        frame, pc = debug_command(frame, :s)
        frame, pc = debug_command(frame, :se)  # get past copymutable
        frame = stepkw!(frame)
        @test frame.pc == JuliaInterpreter.nstatements(frame.framecode) - 1
    end

    function f(x, y)
        sin(2.0)
        g(x; y = 3)
    end
    g(x; y) = x + y
    @testset "interaction of :n with kw functions" begin
        frame = JuliaInterpreter.enter_call(f, 2, 3) # at sin
        frame, pc = debug_command(frame, :n)
        # Check that we are at the kw call to g
        @test Core.kwfunc(g) == JuliaInterpreter.@lookup frame JuliaInterpreter.pc_expr(frame).args[1]
        # Step into the inner g
        frame, pc = debug_command(frame, :s)
        # Finish the frame and make sure we step out of the wrapper
        frame, pc = debug_command(frame, :finish)
        @test frame.framecode.scope == @which f(2, 3)
    end

    h_1(x, y) = h_2(x, y)
    h_2(x, y) = h_3(x; y=y)
    h_3(x; y = 2) = x + y
    @testset "stepping through kwprep after stepping through wrapper" begin
        frame = JuliaInterpreter.enter_call(h_1, 2, 1)
        frame, pc = debug_command(frame, :s)
        # Should have skipped the kwprep in h_2 and be at call to kwfunc h_3
        @test Core.kwfunc(h_3) == JuliaInterpreter.@lookup frame JuliaInterpreter.pc_expr(frame).args[1]
    end

    @testset "si should not step through wrappers or kwprep" begin
        frame = JuliaInterpreter.enter_call(h_1, 2, 1)
        frame, pc = debug_command(frame, :si)
        @test frame.pc == 1
    end

    @testset "breakpoints hit during wrapper step through" begin
        f(x = g()) = x
        g() = 5
        @breakpoint g()
        frame = JuliaInterpreter.enter_call(f)
        JuliaInterpreter.maybe_step_through_wrapper!(frame)
        @test leaf(frame).framecode.scope == @which g()
    end

    @testset "preservation of stack when throwing to toplevel" begin
        f() = "αβ"[2]
        frame1 = JuliaInterpreter.enter_call(f);
        err = try debug_command(frame1, :c)
        catch err
            err
        end
        try
            break_on(:error)
            frame2, pc = @interpret f()
            @test leaf(frame2).framecode.scope === leaf(frame1).framecode.scope
        finally
            break_off(:error)
        end
    end

    @testset "breakpoint in next line" begin
        function f(a, b)
            a == 0 && return abs(b)
            @bp
            return b
        end

        frame = JuliaInterpreter.enter_call(f, 5, 10)
        frame, pc = JuliaInterpreter.debug_command(frame, :n)
        @test pc isa BreakpointRef
    end

    @testset "kw wrapper heuristic #435" begin
        foo() = UInt8('\t')
        frame = JuliaInterpreter.enter_call(foo)
        frame, pc = debug_command(frame, :s)
        @test pc isa BreakpointRef
    end
# end

module Foo
    using ..JuliaInterpreter
    function f(x)
        x
        @bp
        x
    end
end
@testset "interpreted methods" begin
    g(x) = Foo.f(x)

    push!(JuliaInterpreter.compiled_modules, Foo)
    frame = JuliaInterpreter.enter_call(g, 5)
    frame, pc = JuliaInterpreter.debug_command(frame, :n)
    @test !(pc isa BreakpointRef)

    push!(JuliaInterpreter.interpreted_methods, first(methods(Foo.f)))
    frame = JuliaInterpreter.enter_call(g, 5)
    frame, pc = JuliaInterpreter.debug_command(frame, :n)
    @test pc isa BreakpointRef
end


text:
module LazyStack

export stack, rstack

#===== Tuples =====#

ndims(A) = Base.ndims(A)
ndims(::Tuple) = 1
ndims(::NamedTuple) = 1

axes(A) = Base.axes(A)
axes(A, d) = Base.axes(A, d)
if VERSION < v"1.1" # because, on Julia 1.0, axes((1,2)) === Base.OneTo(2)
    axes(t::Tuple) = tuple(Base.axes(t))
end
axes(nt::NamedTuple) = tuple(Base.OneTo(length(nt)))

size(A) = Base.size(A)
size(t::Tuple) = tuple(length(t))
size(t::NamedTuple) = tuple(length(t))

similar_vector(x::AbstractArray, n::Int) = similar(x, n::Int)
similar_vector(x::Tuple, n::Int) = Vector{eltype(x)}(undef, n::Int)
similar_vector(x::NamedTuple, n::Int) = Vector{eltype(x)}(undef, n::Int)

# eltype(x::Tuple) = Base.promote_type(x...) # to match choice below
# eltype(x::NamedTuple) = Base.promote_type(x...)

#===== Slices =====#

"""
    stack([A, B, C])
    stack(A, B, C)

Creates a very simple lazy `::Stacked` view of an array of arrays,
or a tuple of arrays, provided they have a consistent `eltype`.
The dimensions of the inner arrays come first:
```
julia> stack([1,2,3], [10,20,30])
3×2 LazyStack.Stacked{Int64,2,Tuple{Array{Int64,1},Array{Int64,1}}}:
 1  10
 2  20
 3  30

julia> vecs = [[1,2,3] .+ j*im for j=2:2:8];

julia> A = stack(vecs)
3×4 Stacked{Complex{Int64},2,Array{Array{Complex{Int64},1},1}}:
 1+2im  1+4im  1+6im  1+8im
 2+2im  2+4im  2+6im  2+8im
 3+2im  3+4im  3+6im  3+8im
```
A few `Base` functions will immediately look inside:
```
julia> eachcol(A) |> typeof
Array{Array{Complex{Int64},1},1}

julia> view(A, :,1)
3-element Array{Complex{Int64},1}:
 1 + 2im
 2 + 2im
 3 + 2im
```
"""
stack(xs::AbstractArray{IT,ON}) where {IT<:AbstractArray{T,IN}} where {T,IN,ON} =
    stack_slices(xs, Val(T), Val(IN+ON))
stack(xs::Tuple{Vararg{AbstractArray{T,IN}}}) where {T,IN} =
    stack_slices(xs, Val(T), Val(IN+1))
stack(xs::AbstractArray{T}...) where {T} = stack(xs)

function stack_slices(xs::AT, ::Val{T}, ::Val{N}) where {T,N,AT}
    length(xs) >= 1 || throw(ArgumentError("stacking an empty collection is not allowed"))
    storage_type(first(xs)) <: Union{Array, AbstractRange} || return stack_iter(xs)
    s = size(first(xs))
    for x in xs
        size(x) == s || throw(DimensionMismatch(
            "slices being stacked must share a common size. Expected $s, got $(size(x))"))
    end
    Stacked{T, N, AT}(xs)
end

struct Stacked{T,N,AT} <: AbstractArray{T,N}
    slices::AT
end

Base.size(x::Stacked) = (size(first(x.slices))..., size(x.slices)...)
Base.size(x::Stacked{T,N,<:Tuple}) where {T,N} = (size(first(x.slices))..., length(x.slices))

Base.axes(x::Stacked) = (axes(first(x.slices))..., axes(x.slices)...)

Base.parent(x::Stacked) = x.slices

outer_ndims(x::Stacked) = ndims(x.slices)
outer_ndims(x::Stacked{T,N,<:Tuple}) where {T,N} = 1

inner_ndims(x::Stacked) = ndims(x) - outer_ndims(x)

@inline function Base.getindex(x::Stacked{T,N}, inds::Vararg{Integer,N}) where {T,N}
    @boundscheck checkbounds(x, inds...)
    IN, ON = inner_ndims(x), outer_ndims(x)
    outer = @inbounds getindex(x.slices, ntuple(d -> inds[d+IN], ON)...)
    @inbounds getindex(outer, ntuple(d -> inds[d], IN)...)::T
end

Base.unaliascopy(x::Stacked) = stack(map(Base.unaliascopy, x.slices))

if VERSION >= v"1.1"
    Base.eachcol(x::Stacked{T,2,<:AbstractArray{<:AbstractArray{T,1}}}) where {T} = x.slices
end

Base.collect(x::Stacked{T,2,<:AbstractArray{<:AbstractArray{T,1}}}) where {T} = reduce(hcat, x.slices)

"""
    LazyStack.ensure_dense(A)

This `collect`s the result of `stack` if it is a `Stacked` container,
but does nothing to other arrays, such as those produced by `stack_iter`.
"""
ensure_dense(x::Stacked) = collect(x)
ensure_dense(x::AbstractArray) = x

Base.view(x::Stacked{T,2,<:AbstractArray{<:AbstractArray{T,1}}}, ::Colon, i::Int) where {T} = x.slices[i]

function Base.push!(x::Stacked{T,N,<:AbstractVector}, y::AbstractArray) where {T,N}
    s = size(first(x.slices))
    isempty(y) || size(y) == s || throw(DimensionMismatch(
            "slices being stacked must share a common size. Expected $s, got $(size(y))"))
    push!(x.slices, y)
    x
end

function Base.showarg(io::IO, x::Stacked, toplevel)
    print(io, "stack(")
    Base.showarg(io, parent(x), false)
    print(io, ')')
    toplevel && print(io, " with eltype ", eltype(x))
end

#===== Iteration =====#

ITERS = [:Flatten, :Drop, :Filter]
for iter in ITERS
    @eval ndims(::Iterators.$iter) = 1
end
ndims(gen::Base.Generator) = ndims(gen.iter)
ndims(zed::Iterators.Zip) = maximum(ndims, zed.is)
if VERSION < v"1.1"
    ndims(zed::Iterators.Zip2) = max(ndims(zed.a), ndims(zed.b))
end

similar_vector(x, n::Int) = throw(ArgumentError())

"""
    stack(::Generator)
    stack(::Array{T}, ::Array{S}, ...)

This constructs a new array, calling `stack_iter`.
Can handle inconsistent eltypes, but not inconsistent sizes.

```
julia> stack([i i; i 10i] for i in 1:2:3)
2×2×2 Array{Int64,3}:
[:, :, 1] =
 1   1
 1  10

[:, :, 2] =
 3   3
 3  30

julia> stack(c<9 ? [c,2c] : [9.99, 10] for c in 1:10)
2×10 Array{Real,2}:
 1  2  3  4   5   6   7   8   9.99   9.99
 2  4  6  8  10  12  14  16  10.0   10.0

julia> stack(1:3, ones(3), zeros(3) .+ im)
3×3 Array{Number,2}:
 1  1.0  0.0+1.0im
 2  1.0  0.0+1.0im
 3  1.0  0.0+1.0im
```
"""
stack(gen::Base.Generator) = stack_iter(gen)
for iter in ITERS
    @eval stack(gen::Iterators.$iter) = stack_iter(gen)
end
stack(arr::AbstractArray{Any}) = stack_iter(arr) # e.g. from arr=[]; push!(arr, rand(3)); ...

stack(arr::AbstractArray{<:AbstractArray}) = stack_iter(arr)
stack(arr::AbstractArray{<:Tuple}) = stack_iter(arr)
stack(arr::AbstractArray{<:NamedTuple}) = stack_iter(arr)

stack(tup::AbstractArray...) = stack_iter(tup)
stack(tup::Tuple...) = stack_iter(tup)
stack(tup::NamedTuple...) = stack_iter(tup)
# stack(tup::Tuple{Vararg{AbstractArray}}) = stack_iter(tup)

function stack_iter(itr)
    if itr isa Tuple || itr isa Base.Generator{<:Tuple} || ndims(itr) == 1
        outsize = tuple(:)
    else
        Base.haslength(itr) || return stack_iter(collect(itr))
        outsize = size(itr)
    end

    w, val = vstack_plus(itr)

    z = reshape(w, size(val)..., outsize...)

    rewrap_like(z, val)
end

vstack(itr) = first(vstack_plus(itr))

function vstack_plus(itr)
    zed = iterate(itr)
    zed === nothing && throw(ArgumentError("stacking an empty collection is not allowed"))
    val, state = zed
    val isa Union{AbstractArray, Tuple, NamedTuple} || return collect(itr), val

    s = size(val)
    n = Base.haslength(itr) ? prod(s)*length(itr) : nothing

    v = similar_vector(no_wraps(val), something(n, prod(s)))
    copyto!(v, 1, no_wraps(val), 1, prod(s))

    w = stack_rest(v, 0, n, s, itr, state)#::Vector
    w, val
end

function stack_rest(v, i, n, s, itr, state)
    while true
        zed = iterate(itr, state)
        zed === nothing && return v
        val, state = zed
        s == size(val) || throw(DimensionMismatch(
            "slices being stacked must share a common size. Expected $s, got $(size(val))"))

        i += 1
        if eltype(val) <: eltype(v)
            if n isa Int
                copyto!(v, i*prod(s)+1, no_wraps(val), 1, prod(s))
            else
                append!(v, vec(no_wraps(val)))
            end
        else

            T′ = Base.promote_typejoin(eltype(v), eltype(val))
            # T′ = Base.promote_type(eltype(v), eltype(val)) # which do I want?
            v′ = similar(v, T′)
            copyto!(v′, v)

            if n isa Int
                copyto!(v′, i*prod(s)+1, no_wraps(val), 1, prod(s))
            else
                append!(v′, vec(no_wraps(val)))
            end

            return stack_rest(v′, i, n, s, itr, state)
        end

    end
end

"""
    stack(fun, iters...)
    stack(eachcol(A), eachslice(B, dims=3)) do a, b
        f(a,b)
    end

If the first argument is a function, then this is mapped over the other argumenst.
The result should be `== stack(map(fun, iters...))`, but done using `stack_iter`
to return a dense `Array` instead of a `Stacked` object.
"""
stack(fun::Function, iter) = stack(fun(arg) for arg in iter)
function stack(fun::Function, iters...)
    if all(Base.haslength, iters)
        sz = size(first(iters))
        all(a -> size(a)==sz, iters) || throw(DimensionMismatch(
            "sizes of all argumens must match, in stack(f, A, B, C). " *
            "This is slightly stricter than map(f, A, B, C), for now."))
    end
    stack(fun(args...) for args in zip(iters...))
end

#===== Offset =====#

using OffsetArrays

no_wraps(a) = a
no_wraps(a::OffsetArray) = parent(a)

rewrap_like(A, a) = A
function rewrap_like(A, a::OffsetArray)
    B = rewrap_like(A, parent(a))
    OffsetArray(B, axes(a)..., axes(A, ndims(A)))
end

#===== NamedDims =====#

using NamedDims

ensure_named(a::AbstractArray, L::Tuple) = NamedDimsArray(a, L)
ensure_named(a::NamedDimsArray, L::Tuple) = refine_names(a, L)

# array of arrays
stack(xs::NamedDimsArray{<:Any,<:AbstractArray}) =
    ensure_named(stack(parent(xs)), getnames(xs))
stack(x::AT) where {AT <: AbstractArray{<:NamedDimsArray{L,T,IN},ON}} where {T,IN,ON,L} =
    ensure_named(Stacked{T, IN+ON, AT}(x), getnames(x))

getnames(xs::AbstractArray{<:AbstractArray}) =
    (dimnames(eltype(xs))..., dimnames(xs)...)

# tuple of arrays
stack(x::AT) where {AT <: Tuple{Vararg{NamedDimsArray{L,T,IN}}}} where {T,IN,L} =
    ensure_named(stack(map(parent, x)), getnames(x))

getnames(xs::Tuple{Vararg{<:NamedDimsArray}}) =
    (dimnames(first(xs))..., :_)

function rewrap_like(A, a::NamedDimsArray{L}) where {L}
    B = rewrap_like(A, parent(a))
    ensure_named(B, (L..., ntuple(_ -> :_, ndims(A) - ndims(a))...))
end

no_wraps(a::NamedDimsArray) = no_wraps(parent(a))

"""
    stack(name, things...)

If you give one `name::Symbol` before the pieces to be stacked,
this will be the name of the last dimension of the resulting `NamedDimsArray`.
(Names attached to slices, and to containers, should also be preserved.)
"""
function LazyStack.stack(s::Symbol, args...)
    data = stack(args...)
    name_last = ntuple(d -> d==ndims(data) ? s : :_, ndims(data))
    ensure_named(data, name_last)
end

#===== Gradients =====#

using ChainRulesCore: ChainRulesCore, rrule, NoTangent

function ChainRulesCore.rrule(::typeof(stack), vec::AbstractArray{<:AbstractArray{<:Any,IN}}) where {IN}
    stack(vec), Δ -> (NoTangent(), [view(Δ, ntuple(_->(:),IN)..., Tuple(I)...) for I in eachindex(vec)],)
end

function ChainRulesCore.rrule(::typeof(stack), tup::Tuple{Vararg{<:AbstractArray{<:Any,IN}}}) where {IN}
    stack(tup), Δ -> (NoTangent(), ntuple(i -> view(Δ, ntuple(_->(:),IN)..., i), length(tup)),)
end

function ChainRulesCore.rrule(::typeof(stack), gen::Base.Generator)
    stack(gen), Δ -> error("not yet!")
end

#===== CuArrays =====#
# Send these to stack_iter, by testing  storage_type(first(xs)) <: Array

function storage_type(x::AbstractArray)
    p = parent(x)
    typeof(x) === typeof(p) ? typeof(x) : storage_type(p)
end

#===== Ragged =====#

"""
    rstack(arrays; fill=0)

Ragged `stack`, which allows slices of varying size, and fills the gaps with zero
or the given `fill`. Always returns an `Array`.

```
julia> rstack(1:n for n in 1:5)
5×5 Array{Int64,2}:
 1  1  1  1  1
 0  2  2  2  2
 0  0  3  3  3
 0  0  0  4  4
 0  0  0  0  5

julia> rstack([[1,2,3], [10,20.0]], fill=missing)
3×2 Array{Union{Missing, Float64},2}:
 1.0  10.0
 2.0  20.0
 3.0    missing

julia> using OffsetArrays

julia> rstack(1:3, OffsetArray([2.0,2.1,2.2], -1), OffsetArray([3.2,3.3,3.4], +1))
5×3 OffsetArray(::Array{Real,2}, 0:4, 1:3) with eltype Real with indices 0:4×1:3:
 0  2.0  0
 1  2.1  0
 2  2.2  3.2
 3  0    3.3
 0  0    3.4
```

"""
rstack(x::AbstractArray, ys::AbstractArray...; kw...) = rstack((x, ys...); kw...)
rstack(g::Base.Generator; kw...) = rstack(collect(g); kw...)
rstack(f::Function, ABC...; kw...) = rstack(map(f, ABC...); kw...)
rstack(list::AbstractArray{<:AbstractArray}; fill=zero(eltype(first(list)))) = rstack_iter(list; fill=fill)
rstack(list::Tuple{Vararg{<:AbstractArray}}; fill=zero(eltype(first(list)))) = rstack_iter(list; fill=fill)

function rstack_iter(list; fill)
    T = mapreduce(eltype, Base.promote_typejoin, list, init=typeof(fill))
    # T = mapreduce(eltype, Base.promote_type, list, init=typeof(fill))
    N = maximum(ndims, list)
    ax = ntuple(N) do d
        hi = maximum(x -> last(axes(x,d)), list)
        if all(x -> axes(x,d) isa Base.OneTo, list)
            Base.OneTo(hi)
        else
            lo = minimum(x -> first(axes(x,d)), list)
            lo:hi
        end
    end
    arr = Array{T}(undef, map(length, ax)..., size(list)...)
    fill!(arr, fill)
    out = if ax isa Tuple{Vararg{Base.OneTo}}
        arr
    else
        OffsetArray(arr, (ax..., axes(list)...))
    end
    z = rstack_copyto!(out, list, Val(N))

    rewrap_names(z, first(list)) # now I want to separate names & offsets again!
end

function rstack_copyto!(out, list, ::Val{N}) where {N}
    for i in tupleindices(list)
        item = list[i...]
        o = ntuple(_->1, N - ndims(item))
        out[CartesianIndices(axes(item)), o..., i...] .= item

        # https://github.com/JuliaArrays/OffsetArrays.jl/issues/100
        # view(out, axes(item)..., o..., i...) .= item

        # for I in CartesianIndices(item)
        #     out[Tuple(I)..., o..., i...] = item[I]
        # end
    end
    out
end

tupleindices(t::Tuple) = ((i,) for i in 1:length(t))
tupleindices(A::AbstractArray) = (Tuple(I) for I in CartesianIndices(A))

rewrap_names(A, a) = A
function rewrap_names(A, a::NamedDimsArray{L}) where {L}
    B = rewrap_names(A, parent(a))
    ensure_named(B, (L..., ntuple(_ -> :_, ndims(A) - ndims(a))...))
end
function rstack(s::Symbol, args...; kw...)
    data = rstack(args...; kw...)
    name_last = ntuple(d -> d==ndims(data) ? s : :_, ndims(data))
    ensure_named(data, name_last)
end

end # module


text:
"""
    intersectionlength(y::MinHashSketch, y::MinHashSketch)

Count number of hashes in both `x` and `y` more effectively than using ordinary
set operations.

# Examples

```julia-repl
julia> x = minhash(1:10000, 100); y = minhash(5000:15000, 100);

julia> intersectionlength(x, y)
49
```
"""
function intersectionlength(x::MinHashSketch, y::MinHashSketch)
    xi, yi = 1, 1
    n = 0
    @inbounds while (xi ≤ length(x.hashes)) & (yi ≤ length(y.hashes))
        xv, yv = x.hashes[xi], y.hashes[yi]
        n += xv == yv
        xi += xv ≤ yv
        yi += yv ≤ xv
    end
    return n
end

function init_counters(sketches::AbstractVector{MinHashSketch})
    heap = BinaryMinHeap{Tuple{UInt, Int}}()
    counts = counter(UInt)
    for (n, sketch) in enumerate(sketches)
        if !isempty(sketch.hashes)
            h = sketch.hashes[1]
            push!(heap, (h, n))
            inc!(counts, h)
        end
    end
    return heap, counts
end

function increment_matrix!(matrix::Matrix{<:Integer}, indices::Vector{<:Integer}, N::Integer)
    @inbounds for i in 1:N-1
        r = indices[i]
        for j in i+1:N
            c = indices[j]
            c_, r_ = minmax(c, r)
            matrix[r_, c_] += 1
        end
    end
end

# Pop smallest element from heap. If the source sketch have
# more hashes, updates counter and heap with new hash
function popheap!(heap, counts, sketches, indices)
    (smallest, sketchno) = pop!(heap)
    @inbounds hashes = sketches[sketchno].hashes
    @inbounds newindex = indices[sketchno] + 1
    @inbounds indices[sketchno] = newindex
    @inbounds if newindex ≤ length(hashes)
        h = hashes[newindex]

        # We can add a new value without fear that it will be popped
        # in the same round since the new value is guaranteed to be larger
        push!(heap, (h, sketchno))
        inc!(counts, h)
    end
    return (smallest, sketchno)
end

# This function scales better than O(N^2 * S), that's why it's so complicated
"""
    intersectionlength(v::AbstractVector{MinHashSketch})

Return a `length(v)` square `Matrix{Int}` where the `[i, j]` where `i > j`'th
cell contain the number of hashes in both `v[i]` and `v[j]`.

# Examples

```julia-repl
julia> v = [sketch((1:1000) .+ (300*i), 100) for i in 1:4];

julia> intersectionlength(v)
4×4 Array{Int64,2}:
  0   0   0  0
 64   0   0  0
 27  58   0  0
  8  34  73  0
```
"""
function intersectionlength(sketches::AbstractVector{MinHashSketch})
    heap, counts = init_counters(sketches)
    indices = fill(1, length(sketches))
    smallest_ids = Vector{Int}(undef, length(sketches))
    matrix = zeros(Int, (length(sketches), length(sketches)))
    @inbounds while !isempty(heap)
        # Get smallest value
        smallest, sketchno = popheap!(heap, counts, sketches, indices)
        smallest_ids[1] = sketchno
        # How many other values are equal to smallest value?
        N = counts[smallest]
        reset!(counts, smallest)
        # Increment overlaps for all that are equal
        if N > 1
            for i in 2:N
                smallest_ids[i] = popheap!(heap, counts, sketches, indices)[2]
            end
            increment_matrix!(matrix, smallest_ids, N)
        end
    end
    return matrix
end


text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

bat_default(::typeof(bat_sample), ::Val{:rng}, ::Any) = bat_rng()

bat_default(::typeof(bat_sample), ::Val{:algorithm}, ::AnyIIDSampleable) = IIDSampling()

bat_default(::typeof(bat_sample), ::Val{:algorithm}, ::DensitySampleVector) = OrderedResampling()

bat_default(::typeof(bat_sample), ::Val{:algorithm}, ::AbstractDensity) = MCMCSampling()


text:
function exchange_action!(w, potential, u, v, x, ϕ, R)
    z = exp(im*ϕ)
    ρ = Density(u, applied(*, R, z*x))
    copyto!(potential, ρ)
    mul!(w, potential, v.args[2], 1/z)
end

function exchange_action(potential, u, v, ϕ, R)
    n = size(R,2)
    eye = complex(Matrix(1.0I, n, n))
    out = similar(eye)
    for j = 1:size(eye,2)
        exchange_action!(view(out, :, j),
                         potential, u, v,
                         view(eye, :, j), ϕ, R)
    end
    out
end

@testset "Poisson problem Hermiticity" begin
    N = 100
    ρ = 0.2

    R = StaggeredFiniteDifferences(N, ρ)
    r = axes(R,1)
    D = Derivative(r)
    ∇² = apply(*, R', D', D, R)

    T = ∇²/-2
    V = apply(*, R', QuasiDiagonal(-1 ./ r), R)
    H = T+V

    ee = eigen(H)

    @testset "n = $(n)" for n = 1:5
        potential = CoulombRepulsionPotential(R, 0, ComplexF64)
        u = applied(*, R, complex(ee.vectors[:,n]))
        v = applied(*, R, complex(ee.vectors[:,n]))

        @testset "ϕ = $(ϕ)" for ϕ ∈ π*(0:1/8:2)
            A = exchange_action(potential, u, v, ϕ, R)
            @test norm(A-A') < 1e-16
        end
    end
end


text:
using Statistics
using StatsBase: sample
using Revise
using EvoTrees
using BenchmarkTools
using BSON: @save, @load

# prepare a dataset
features = rand(Int(1.25e6), 100)
# features = rand(100, 10)
X = features
Y = rand(size(X, 1))
𝑖 = collect(1:size(X,1))

# train-eval split
𝑖_sample = sample(𝑖, size(𝑖, 1), replace = false)
train_size = 0.8
𝑖_train = 𝑖_sample[1:floor(Int, train_size * size(𝑖, 1))]
𝑖_eval = 𝑖_sample[floor(Int, train_size * size(𝑖, 1))+1:end]

X_train, X_eval = X[𝑖_train, :], X[𝑖_eval, :]
Y_train, Y_eval = Y[𝑖_train], Y[𝑖_eval]

#############################
# CPU - linear
#############################
params1 = EvoTreeRegressor(T=Float32,
    loss=:linear, metric=:none,
    nrounds=10,
    λ = 1.0, γ=0.01, η=0.1,
    max_depth = 9, min_weight = 1.0,
    rowsample=0.5, colsample=0.5, nbins=32)

# for 100k 10 rounds: 410.477 ms (44032 allocations: 182.68 MiB)
# for 100k 100 rounds: 2.177 s (404031 allocations: 626.45 MiB)
# for 1.25e6 no eval: 6.244 s (73955 allocations: 2.18 GiB)
# for 1.25e6 mse with eval data:  6.345 s (74009 allocations: 2.18 GiB)
@time model = fit_evotree(params1, X_train, Y_train);
EvoTrees.save(model, "data/model_linear_9_2.bson")

#############################
# CPU - Gaussian
#############################
params1 = EvoTreeGaussian(T=Float32,
    loss=:gaussian, metric=:none,
    nrounds=10,
    λ = 1.0, γ=0.1, η=0.1,
    max_depth = 5, min_weight = 1.0,
    rowsample=0.5, colsample=0.5, nbins=64)

@time model = fit_evotree(params1, X_train, Y_train);
@save "data/model_gaussian_5.bson" model


text:
# Author: Mathias Louboutin, mlouboutin3@gatech.edu
# Date: July 2020

parsed_args = parse_commandline()

nlayer = parsed_args["nlayer"]
tti = parsed_args["tti"]
fs =  parsed_args["fs"]

@testset "Arithmetic test with $(nlayer) layers and tti $(tti) and freesurface $(fs)" begin

        # Test 2D find_water_bottom

        dm2D = zeros(Float32,10,10)
        dm2D[:,6:end] .= 1f0
        @test find_water_bottom(dm2D) == 6*ones(Integer,10)

        ### Model
        model, model0, dm = setup_model(parsed_args["tti"], parsed_args["nlayer"])
        wb = find_water_bottom(model.m .- maximum(model.m))
        q, srcGeometry, recGeometry, info = setup_geom(model)
        dt = srcGeometry.dt[1]
        nt = length(q.data[1])
        nrec = length(recGeometry.xloc[1])

        opt = Options(free_surface=parsed_args["fs"])
        opta = Options(free_surface=parsed_args["fs"], return_array=true)
        ftol = 5f-5

        w = judiWeights(randn(Float32, model0.n))

        # Build operators
        Pr = judiProjection(info, recGeometry)
        F = judiModeling(info, model; options=opt)
        Fa = judiModeling(info, model; options=opta)
        Ps = judiProjection(info, srcGeometry)
        Pw = judiLRWF(info, q.data[1])
        J = judiJacobian(Pr*F*adjoint(Ps), q)
        Jw = judiJacobian(Pr*F*adjoint(Pw), w)

        dobs = Pr * F * Ps' * q
        dobsa = Pr * Fa * Ps' * q
        dobs_w = Pr * F * Pw' * w
        dobs_wa = Pr * Fa * Pw' * w
        dobs_out = 0f0 .* dobs
        dobs_outa = 0f0 .* dobsa
        dobs_w_out = 0f0 .* dobs_w
        dobs_w_outa = 0f0 .* dobs_wa
        q_out = 0f0 .* q
        w_out = 0f0 .* w

        # mul!
        mul!(dobs_out, Pr * F * Ps', q)
        mul!(dobs_w_out, Pr * F * Pw', w)
        mul!(dobs_outa, Pr * Fa * Ps', q)
        mul!(dobs_w_outa, Pr * Fa * Pw', w)
        @test isapprox(dobs, dobs_out; rtol=ftol)
        @test isapprox(dobsa, dobs_outa; rtol=ftol)
        @test isapprox(dobs_w, dobs_w_out; rtol=ftol)
        @test isapprox(dobs_wa, dobs_w_outa; rtol=ftol)

        mul!(w_out, adjoint(Pr * F * Pw'), dobs_w_out)
        w_a = adjoint(Pr * Fa * Pw') * dobs_w_out
        w_outa = 0f0 .* w_a
        mul!(w_outa, adjoint(Pr * Fa * Pw'), dobs_w_out)
        @test isapprox(w_out, adjoint(Pr * F * Pw') * dobs_w_out; rtol=ftol)
        @test isapprox(w_outa, w_a; rtol=ftol)

        # jacobian
        dm2 = copy(dm)
        dmd = copy(dm2.data)
        mul!(dm2, J', dobs)
        mul!(dmd, J', dobs)
        @test isapprox(dm2, J'*dobs)
        @test isapprox(dmd, dm2)

        mul!(dobs_out, J, dm)
        mul!(dobs, J, dm.data)
        dlin = J*dm
        @test isapprox(dobs_out, dlin; rtol=ftol)
        @test isapprox(dobs_out, dobs; rtol=ftol)
        # JUDI precon make sure it runs
        F = judiFilter(recGeometry, .002, .030)
        Md = judiMarineTopmute2D(0, recGeometry)
        D = judiDepthScaling(model)
        M = judiTopmute(model.n, 20, 1)

        mul!(dobs_out, F, dobs)
        @test isapprox(dobs_out, F*dobs; rtol=ftol)
        mul!(dobs_out, F', dobs)
        @test isapprox(dobs_out, F'*dobs; rtol=ftol)
        mul!(dobs_out, Md, dobs)
        @test isapprox(dobs_out, Md*dobs; rtol=ftol)
        mul!(dobs_out, Md', dobs)
        @test isapprox(dobs_out, Md'*dobs; rtol=ftol)

        mul!(w_out, D, w)
        @test isapprox(w_out, D*w; rtol=ftol)
        @test isapprox(w_out.weights[1][end, end]/w.weights[1][end, end], sqrt((model.n[2]-1)*model.d[2]); rtol=ftol)
        mul!(w_out, D', w)
        @test isapprox(w_out, D'*w; rtol=ftol)
        @test isapprox(w_out.weights[1][end, end]/w.weights[1][end, end], sqrt((model.n[2]-1)*model.d[2]); rtol=ftol)

        mul!(w_out, M, w)
        @test isapprox(w_out, M*w; rtol=ftol)
        @test isapprox(norm(w_out.weights[1][:, 1:19]), 0f0; rtol=ftol)
        mul!(w_out, M', w)
        @test isapprox(w_out, M'*w; rtol=ftol)

        # test the output of depth scaling and topmute operators, and test if they are out-of-place

        dobs1 = deepcopy(dobs)
        for Op in [F, F', Md , Md']
                m = Op*dobs
                # Test that dobs wasn't modified
                @test isapprox(dobs,dobs1,rtol=eps())
                # Test that it did compute something
                @test m != dobs
        end

        w1 = deepcopy(w)

        for Op in [D, D']
                m = Op*w
                # Test that w wasn't modified
                @test isapprox(w,w1,rtol=eps())

                w_expect = deepcopy(w)
                for j = 1:model.n[2]
                        w_expect.weights[1][:,j] = w.weights[1][:,j] * Float32(sqrt(model.d[2]*(j-1)))
                end
                @test isapprox(w_expect,m)
        end

        for Op in [M , M']
                m = Op*w
                # Test that w wasn't modified
                @test isapprox(w,w1,rtol=eps())

                @test all(isapprox.(m.weights[1][:,1:18], 0))
                @test isapprox(m.weights[1][:,21:end],w.weights[1][:,21:end])
        end

        n = (100,100,100)
        d = (10f0,10f0,10f0)
        o = (0.,0.,0.)
        m = 0.25*ones(Float32,n)
        model3D = Model(n,d,o,m)

        D3 = judiDepthScaling(model3D)

        dm = randn(Float32,prod(n))
        dm1 = deepcopy(dm)
        for Op in [D3, D3']
                opt_out = Op*dm
                # Test that dm wasn't modified
                @test dm1 == dm

                dm_expect = zeros(Float32,model3D.n)
                for j = 1:model3D.n[3]
                        dm_expect[:,:,j] = reshape(dm,model3D.n)[:,:,j] * Float32(sqrt(model3D.d[3]*(j-1)))
                end
                @test isapprox(vec(dm_expect),opt_out)
        end

        M3 = judiTopmute(model3D.n, 20, 1)

        for Op in [M3, M3']
                opt_out = Op*dm
                # Test that dm wasn't modified
                @test dm1 == dm

                @test all(isapprox.(reshape(opt_out,model3D.n)[:,:,1:18], 0))
                @test isapprox(reshape(opt_out,model3D.n)[:,:,21:end],reshape(dm,model3D.n)[:,:,21:end])
        end

        # test find_water_bottom in 3D

        dm3D = zeros(Float32,10,10,10)
        dm3D[:,:,4:end] .= 1f0
        @test find_water_bottom(dm3D) == 4*ones(Integer,10,10)

end

max_stars_repo_name: nogueirapeterson/JUDI

max_issues_repo_name: nogueirapeterson/JUDI

max_forks_repo_name: nogueirapeterson/JUDI


text:
"""
Methods to perform zero padding and truncation.

# Arguments

* `x` : real signal with dimension nplag + nnlag + 1
	first has decreasing negative lags,
	then has zero lag at nnlags + 1,
	then has increasing positive nplags lags,
	signal contains only positive lags and zero lag if nnlag=0 and vice versa
* `xpow2` : npow2 real vector with dimension npow2
* `nplags` : number of positive lags
* `nnlags` : number of negative lags
* `npow2` : number of samples in xpow2
* `pad` means xpow2 is returned using x
* `truncate` means x is returned using xpow2
"""
function pad!(
				  x::AbstractArray{T},
				  xpow2::AbstractArray{T},
				  nplags::Integer,
				  nnlags::Integer,
				  npow2::Integer,
				  ) where {T}
	@assert size(x,1) == nplags + nnlags + 1
	@assert size(xpow2,1) == npow2

	for id in 1:size(x,2)
		xpow2[1,id] = (x[nnlags+1,id]) # zero lag
		# +ve lags
		if (nplags > 0)
			for i=1:nplags
				@inbounds xpow2[i+1,id]= (x[nnlags+1+i,id])
			end
		end
		# -ve lags
		if(nnlags != 0)
			for i=1:nnlags
				@inbounds xpow2[npow2-i+1,id] =(x[nnlags+1-i,id])
			end
		end
	end
	return nothing
end

function truncate!(
				  x::AbstractArray{T},
				  xpow2::AbstractArray{T},
				  nplags::Integer,
				  nnlags::Integer,
				  npow2::Integer,
				  ) where {T}
	(size(x,1) ≠ nplags + nnlags + 1) && error("size x")
	(size(xpow2,1) ≠ npow2) && error("size xpow2")

	for id in 1:size(x,2)
		x[nnlags+1,id] = (xpow2[1,id]); # zero lag
		if(nplags != 0)
			for i=1:nplags
				@inbounds x[nnlags+1+i,id] = (xpow2[1+i,id]);
			end
		end
		if(nnlags != 0)
			for i=1:nnlags
				@inbounds x[nnlags+1-i,id] = (xpow2[npow2-i+1,id])
			end
		end
	end
	return nothing
end


text:
using FymEnvs
using LinearAlgebra

using Plots
ENV["GKSwstype"]="nul"  # do not show plot
# using Debugger

if !isdefined(Main, :TestEnvs)
    # using Revise; includet("custom_env.jl")  # to avoid conflict
    include("custom_env.jl")  # to avoid conflict
    using .TestEnvs
end

function print_msg(test_name)
    println(">"^6*" "*test_name*" "*"<"^6)
end

function test_custom_fym()
    print_msg("custom FymEnv")
    fym = TestEnv(; dt=0.001, max_t=1.000)

    env = fym.env
    log_dir = "data/test"
    file_name = "custom.jld2"
    reset!(env)
    _sample(fym, nothing, log_dir, file_name)
    data, info = load(joinpath(log_dir, file_name),
                      with_info=true)
    @show length(data["time"])
    @show length(data["group"]["time"])
    @show size(data["group"]["state"])
    @show info
    plot(data["time"], data["state"])
    for name in ["state", "input"]
        p = plot(data["time"], data[name])
        savefig(p, joinpath(log_dir, "custom_"*name*".pdf"))
    end
end

function test_reverse_time()
    print_msg("reverse_time")
    function set_dyn(env, t)
        # corresponding to `set_dot` of the original fym
        # you can use any names in this package
        sys = env.systems["sys"]
        x = sys.state
        A = Matrix(I, 3, 3)
        sys.dot = -A * x
    end
    function step(env)
        t = time(env.clock)
        sys = env.systems["sys"]
        x = sys.state
        update!(env)
        next_obs = sys.state
        reward = zeros(1)
        done = time_over(env.clock)
        info = Dict(
                    "time" => t,
                    "state" => x,
                   )
        return next_obs, reward, done, info
    end

    x0 = collect(1:3)
    systems = Dict("sys" => BaseSystem(initial_state=x0, name="3d_sys"))
    log_dir = "data/test"
    file_name = "fym.jld2"
    logger = Logger(log_dir=log_dir, file_name=file_name)
    env = BaseEnv(max_t=0.00, dt=-0.01, initial_time=1.0, logger=logger, name="test_env",)
    systems!(env, systems)  # set systems; required
    dyn!(env, set_dyn)  # set dynamics; required
    step!(env, step)  # set step; required

    reset!(env)  # reset env; required before propagation
    obs = observe_flat(env)
    i = 0
    @time while true
        render(env)  # not mendatory; would make simulator slow
        next_obs, reward, done, info = env.step()
        obs = next_obs
        i += 1
        if done
            break
        end
    end
    close!(env)
    data = load(env.logger.path)

    show(env)
    show(size(data["state"]["sys"]))
end

function _sample(env::BaseEnv, agent, log_dir, file_name)
    logger = Logger(log_dir=log_dir, file_name=file_name)
    obs = observe_flat(env)
    i = 0
    @time while true
        if agent == nothing
            action = nothing
        end
        next_obs, reward, done, info = env.step(action=action)
        record(logger, info)
        obs = next_obs
        i += 1
        if done
            break
        end
    end
    close!(env)
    close!(logger)
end

function _sample(fym::FymEnv, agent, log_dir, file_name)
    env = fym.env
    logger = Logger(log_dir=log_dir, file_name=file_name)
    obs = observe_flat(env)
    i = 0
    @time while true
        if agent == nothing
            action = nothing
            if time(env.clock) > 0.5
                fym.controller.K = - 3*Matrix(I, 3, 3)
            end
        end
        next_obs, reward, done, info = env.step(action=action)
        record(logger, info)
        obs = next_obs
        i += 1
        if done
            config = Dict(
                          "name" => env.name,
                       )
            set_info!(logger, config)
            break
        end
    end
    close!(env)
    close!(logger)
end

function test_all()
    test_custom_fym()
    test_reverse_time()
end

test_all()


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

include("learning/tasks.jl")
include("learning/models.jl")
include("learning/traits.jl")


text:
using Documenter, NuFlux

makedocs(;
    modules = [NuFlux],
    authors = "Johannes Schumann",
    format = Documenter.HTML(
        prettyurls = get(ENV, "CI", nothing) == "true",
    ),
    pages=[
        "Introduction" => "index.md",
    ],
    repo="https://github.com/8me/NuFlux.jl/blob/{commit}{path}#L{line}",
    sitename="NuFlux.jl",
)

deploydocs(;
    repo="github.com/8me/NuFlux.jl",
    devbranch="main"
)


text:
module JLLEnvs

using Pkg
using Pkg.Artifacts
using Pkg.Artifacts: load_artifacts_toml, download_artifact, artifact_path
using Downloads
using Base.BinaryPlatforms

function __triplet(p::Platform)
    for k in keys(p.tags)
        if k != "arch" && k != "os" && k != "libc"
            delete!(p.tags, k)
        end
    end
    t = triplet(p)
    if os(p) == "macos" && arch(p) == "x86_64"
        t *= "14"
    elseif os(p) == "macos" && arch(p) == "aarch64"
        t *= "20"
    end
    return t
end

include("env.jl")

function __init__()
    return merge!(SHARDS, load_artifacts_toml(Downloads.download(SHARDS_URL)))
end

include("version.jl")
include("target.jl")
include("type.jl")
include("system.jl")
# include("library.jl")

function get_default_args(is_cxx=false, version=GCC_MIN_VER)
    env = get_default_env(; is_cxx, version)
    args = ["-isystem" * dir for dir in get_system_includes(env)]
    push!(args, "--target=$(target(env.platform))")
    return args
end

end # module


text:
using Interpolations, Distributions, ForwardDiff, LinearAlgebra, FastGaussQuadrature

const TINY = 1e-30
const _FWHM = 2 * sqrt(2 * log(2.0))
const PI4 = 4e-6 * pi
const _INTLIMIT = 3.5

# Base.float(d::ForwardDiff.Dual{T}) where T = ForwardDiff.Dual{T}(float(d.value), d.partials)
# Base.prevfloat(d::ForwardDiff.Dual{T,V,N}) where {T,V,N} = ForwardDiff.Dual{T}(prevfloat(float(d.value)), d.partials)
# Base.nextfloat(d::ForwardDiff.Dual{T,V,N}) where {T,V,N} = ForwardDiff.Dual{T}(nextfloat(float(d.value)), d.partials)

function Base.ldexp(x::T, e::Integer) where T<:ForwardDiff.Dual
    if e >=0
        x * (1<<e)
    else
        x / (1<<-e)
    end
end

"""
    abeles(q::Array{Float64, 1}, layers::Array{Float64, 2})

Performs the Abeles optical matrix calculation.

### Parameters
- `q::Array{Float64, 1}` : q-vector values.
- `layers::Array{Float64, 2}` : an Nx4 array, where N is the number of layers in the system, 1st item in a given row is the thickness, the 2nd the SLD, the 3rd the imaginary SLD, and the 4th the roughness with the layer above.

### Returns
- `::Array{Float64, 1}` : unsmeared reflectometry values for the given q-vectors.
"""
function abeles(q, w)
    nlayers = size(w, 1) - 2
    npnts = length(q)

    reflectivity = Vector{Any}(undef, (npnts))
    oneC = Complex(1.0)

    for j in eachindex(q)
        qq2 = (q[j] * q[j] / 4.0) + 0.0im
        kn = (q[j] / 2.) + 0.0im

        # variables are local to if blocks
        local MRtotal11, MRtotal12, MRtotal21, MRtotal22

        for i = 1:nlayers+1
            # wavevector in the layer
            sld_next = ((w[i+1, 2] - w[1, 2]) + ((abs(w[i+1, 3]) + TINY))im) * pi * 4.0e-6
            kn_next = sqrt(qq2 - sld_next)

            # reflectance of the interface
            rj = (kn - kn_next)/(kn + kn_next) * exp(kn * kn_next * (-2.0 * w[i+1, 4]^2))

            if i == 1
                # characteristic matrix for first interface
                MRtotal11 = oneC
                MRtotal12 = rj
                MRtotal21 = rj
                MRtotal22 = oneC
            else
                # work out the beta for the layer
                beta = exp(kn * (abs(w[i, 1]) * 1im))

                # this is the characteristic matrix of a layer
                MI11 = beta
                MI12 = rj * beta
                MI22 = oneC / beta
                MI21 = rj * MI22

                # propagate optical matrix by matmul
                p11 = MRtotal11 * MI11 + MRtotal12 * MI21
                p12 = MRtotal11 * MI12 + MRtotal12 * MI22
                p21 = MRtotal21 * MI11 + MRtotal22 * MI21
                p22 = MRtotal21 * MI12 + MRtotal22 * MI22

                MRtotal11 = p11
                MRtotal12 = p12
                MRtotal21 = p21
                MRtotal22 = p22

            end
            kn = kn_next;

        end
        reflectivity[j] = MRtotal21 / MRtotal11
    end
    return real(reflectivity .* conj(reflectivity))
end

"""
    same_convolution(a::Array{Float64, 1}, b::Array{Float64, 1})

Performs a convolution of one-dimensional arrays equivalent to the [`np.convolve`](https://numpy.org/doc/stable/reference/generated/numpy.convolve.html), where the `mode` is `'same'`.

### Parameters
- `a::Array{Float64, 1}` : first array to convolve.
- `b::Array{Float64, 1}` : second array to convolve.

### Returns
- `::Array{Float64, 1}` : discrete, linear convolution of `a` and `b`.
"""
function same_convolution(a, b)

    m = length(a)
    n = length(b)

    padleft = ceil(Int32, n/2) - 1
    padright = floor(Int32, n/2)

    zero_padleft = Array{Any}(undef, padleft)
    for i in 1:padleft
        zero_padleft[i] = 0.
    end
    zero_padright = Array{Any}(undef, padleft)
    for i in 1:padleft
        zero_padright[i] = 0.
    end
    a = append!(zero_padleft, a)
    a = append!(a, zero_padright)

    output = Array{Any}(undef, m)
    for i in 1:m
        output[i] = 0.
    end
    for i in 1:m
        for j in 1:n
            output[i] += a[i+j-1] * b[n-j+1]
        end
    end
    return output
end

"""
    gauss(x::Array{Float64, 1}, s::Float64)

A Gaussian kernel for resolution smearing.

### Parameters
- `x::Array{Float64, 1}` : the kernal positions.
- `s::Float64` : the width of the kernel.

### Returns
- `::Array{Float64, 1}`: probabilities for `x`.
"""
function gauss(x::Array{Float64, 1}, s::Float64)
    g = zeros(size(x, 1))
    for i = 1:size(x, 1)
        g[i] = 1. / s / sqrt(2 * pi) * exp(-0.5 * x[i] ^ 2 / s / s)
    end
    return g
end

"""
    pointwise_smearing(q::Array{Float64, 1}, w::Array{Any, 2}, dq::Array{Float64, 1}; quad_order::Int32=17)

Perform the reflectometry calculation with a pointwise smearing.

### Parameters
- `q::Array{Float64, 1}` : q-vector values.
- `w::Array{Any, 2}` : an Nx4 array, where N is the number of layers in the system, 1st item in a given row is the thickness, the 2nd the SLD, the 3rd the imaginary SLD, and the 4th the roughness with the layer above.
- `dq::Array{Float64, 1}` : `dq` values for each `q` value.
- `quad_order{Int32, 1}` : The order of the Gaussian quadrature integration for the resolution integration.

### Returns
- `::Array{Float64, 1}` : smeared reflectometry values for the given q-vectors, with a pointise smearing.
"""
function pointwise_smearing(q, w, dq, quad_order::Int=17)
    # not sure how one checks that q, dq have same size.
    npnts = length(q)

    # get the gauss-legendre weights and abscissae
    # TODO: is it possible to use a LRU cache?
    abscissa, weights = gausslegendre(quad_order)

    # get the normal distribution at that point
    prefactor = 1.0 / sqrt(2 * pi)

    function gauss(x)
        return exp(-0.5 * x * x)
    end

    # TODO: is it possible to use a LRU cache?
    gaussvals = prefactor * gauss.(abscissa * _INTLIMIT) .* weights

    # integration between -3.5 and 3.5 sigma
    va = q .- _INTLIMIT .* dq ./ _FWHM
    vb = q .+ _INTLIMIT .* dq ./ _FWHM

    # (quad_order, npnts)
    qvals_for_res = [(ab * (li[2] - li[1]) + li[2] + li[1])/2.0 for ab=abscissa, li=zip(va, vb)]
    smeared_rvals = abeles(qvals_for_res, w)

    # abeles flattens the q vector
    smeared_rvals = reshape(smeared_rvals, size(qvals_for_res))
    m = reshape(gaussvals, (quad_order, 1))

    smeared_rvals .*= m
    rvals = sum(smeared_rvals, dims=1) .* _INTLIMIT
    return reshape(rvals, size(q))
end

"""
    constant_smearing(q::Array{Float64, 1}, w::Array{Any, 2}, resolution::Any, scale::Any, bkg::Any)

Perform the reflectometry calculation with a constant convolutional smearing.

### Parameters
- `q::Array{Float64, 1}` : q-vector values.
- `w::Array{Any, 2}` : an Nx4 array, where N is the number of layers in the system, 1st item in a given row is the thickness, the 2nd the SLD, the 3rd the imaginary SLD, and the 4th the roughness with the layer above.
- `resolution::Any` : the percentage resolution (dq/q) to be used.
- `scale::Any` : the multiplicative scale factor assocaited with the reflectometry profile.
- `bkg::Any` : the uniform background to add to the profile.

### Returns
- `::Array{Float64, 1}` : smeared reflectometry values for the given q-vectors.
"""
function constant_smearing(q, w, resolution, scale, bkg)
    if resolution < 0.5
        return abeles(q, w)
    end

    resolution /= 100
    gaussnum = 51
    gaussgpoint = (gaussnum - 1) / 2

    lowq = minimum(q[:])
    highq = maximum(q[:])
    if lowq <= 0.
        lowq = 1e-6
    end

    start = log10(lowq) - 6 * resolution / _FWHM
    finish = log10(highq * (1 + 6 * resolution / _FWHM))
    interpnum = round(abs(1 * (abs(start - finish)) / (1.7 * resolution / _FWHM / gaussgpoint)))
    xlin = 10 .^ LinRange(start, finish, Int(interpnum))

    gauss_x = collect(LinRange(-1.7 * resolution, 1.7 * resolution, gaussnum))
    gauss_y = gauss(gauss_x, resolution / _FWHM)

    rvals = abeles(xlin, w)
    smeared_rvals = same_convolution(rvals, gauss_y)
    smeared_rvals *= gauss_x[2] - gauss_x[1]

    itp = LinearInterpolation(xlin, smeared_rvals)
    smeared_output = itp(q)
    for i = 1:size(smeared_output, 1)
        smeared_output[i] *= scale
        smeared_output[i] += bkg
    end
    return smeared_output
end


text:
using Rubin
using Tests
using Elliptic
using HypergeometricFunctions
using Polylogarithms
using SpecialFunctions

@test integrate((a+b*csch(c+d*x)^2)^4, x) == :(x*a^4+-1//7*b^4*d^-1*coth(c+d*x)^7+-1//3*b^2*d^-1*coth(c+d*x)^3*(3*b^2+6*a^2+-8*a*b)+-1//5*b^3*d^-1*coth(c+d*x)^5*(-3b+4a)+-1*b*d^-1*(-1b+2a)*(b^2+2*a^2+-2*a*b)*coth(c+d*x))
@test integrate((a+b*csch(c+d*x)^2)^3, x) == :(x*a^3+-1//5*b^3*d^-1*coth(c+d*x)^5+-1*b*d^-1*(b^2+3*a^2+-3*a*b)*coth(c+d*x)+-1//3*b^2*d^-1*coth(c+d*x)^3*(-2b+3a))
@test integrate((a+b*csch(c+d*x)^2)^2, x) == :(x*a^2+-1//3*b^2*d^-1*coth(c+d*x)^3+-1*b*d^-1*(-1b+2a)*coth(c+d*x))
@test integrate(a+b*csch(c+d*x)^2, x) == :(a*x+-1*b*d^-1*coth(c+d*x))
@test integrate((a+b*csch(c+d*x)^2)^-1, x) == :(x*a^-1+-1*a^-1*b^(1/2)*d^-1*(a+-1b)^-1//2*arctan(b^-1//2*(a+-1b)^(1/2)*tanh(c+d*x)))
@test integrate((a+b*csch(c+d*x)^2)^-2, x) == :(x*a^-2+(1/2)*b*a^-1*d^-1*(a+-1b)^-1*(a+-1b+b*coth(c+d*x)^2)^-1*coth(c+d*x)+-1//2*a^-2*b^(1/2)*d^-1*(a+-1b)^-3//2*(-2b+3a)*arctan(b^-1//2*(a+-1b)^(1/2)*tanh(c+d*x)))
@test integrate((a+b*csch(c+d*x)^2)^-3, x) == :(x*a^-3+-1//8*a^-3*b^(1/2)*d^-1*(a+-1b)^-5//2*(8*b^2+15*a^2+-20*a*b)*arctan(b^-1//2*(a+-1b)^(1/2)*tanh(c+d*x))+1//4*b*a^-1*d^-1*(a+-1b)^-1*(a+-1b+b*coth(c+d*x)^2)^-2*coth(c+d*x)+1//8*b*a^-2*d^-1*(a+-1b)^-2*(a+-1b+b*coth(c+d*x)^2)^-1*(-4b+7a)*coth(c+d*x))
@test integrate((a+b*csch(c+d*x)^2)^-4, x) == :(x*a^-4+-1//16*a^-4*b^(1/2)*d^-1*(a+-1b)^-7//2*(-16*b^3+35*a^3+-70*b*a^2+56*a*b^2)*arctan(b^-1//2*(a+-1b)^(1/2)*tanh(c+d*x))+1//6*b*a^-1*d^-1*(a+-1b)^-1*(a+-1b+b*coth(c+d*x)^2)^-3*coth(c+d*x)+1//16*b*a^-3*d^-1*(a+-1b)^-3*(a+-1b+b*coth(c+d*x)^2)^-1*(8*b^2+19*a^2+-22*a*b)*coth(c+d*x)+1//24*b*a^-2*d^-1*(a+-1b)^-2*(a+-1b+b*coth(c+d*x)^2)^-2*(-6b+11a)*coth(c+d*x))
@test integrate((a+b*csch(c+d*x)^2)^5//2, x) == :(a^5//2*d^-1*arctanh(a^(1/2)*(a+-1b+b*coth(c+d*x)^2)^-1//2*coth(c+d*x))+-1//4*b*d^-1*(a+-1b+b*coth(c+d*x)^2)^3//2*coth(c+d*x)+-1//8*b^(1/2)*d^-1*(3*b^2+15*a^2+-10*a*b)*arctanh(b^(1/2)*(a+-1b+b*coth(c+d*x)^2)^-1//2*coth(c+d*x))+-1//8*b*d^-1*(a+-1b+b*coth(c+d*x)^2)^(1/2)*(-3b+7a)*coth(c+d*x))
@test integrate((a+b*csch(c+d*x)^2)^3//2, x) == :(a^3//2*d^-1*arctanh(a^(1/2)*(a+-1b+b*coth(c+d*x)^2)^-1//2*coth(c+d*x))+-1//2*b*d^-1*(a+-1b+b*coth(c+d*x)^2)^(1/2)*coth(c+d*x)+-1//2*b^(1/2)*d^-1*(-1b+3a)*arctanh(b^(1/2)*(a+-1b+b*coth(c+d*x)^2)^-1//2*coth(c+d*x)))
@test integrate((a+b*csch(c+d*x)^2)^(1/2), x) == :(a^(1/2)*d^-1*arctanh(a^(1/2)*(a+-1b+b*coth(c+d*x)^2)^-1//2*coth(c+d*x))+-1*b^(1/2)*d^-1*arctanh(b^(1/2)*(a+-1b+b*coth(c+d*x)^2)^-1//2*coth(c+d*x)))
@test integrate((a+b*csch(c+d*x)^2)^-3//2, x) == :(a^-3//2*d^-1*arctanh(a^(1/2)*(a+-1b+b*coth(c+d*x)^2)^-1//2*coth(c+d*x))+b*a^-1*d^-1*(a+-1b)^-1*(a+-1b+b*coth(c+d*x)^2)^-1//2*coth(c+d*x))
@test integrate((a+b*csch(c+d*x)^2)^-5//2, x) == :(a^-5//2*d^-1*arctanh(a^(1/2)*(a+-1b+b*coth(c+d*x)^2)^-1//2*coth(c+d*x))+1//3*b*a^-1*d^-1*(a+-1b)^-1*(a+-1b+b*coth(c+d*x)^2)^-3//2*coth(c+d*x)+1//3*b*a^-2*d^-1*(a+-1b)^-2*(a+-1b+b*coth(c+d*x)^2)^-1//2*(-3b+5a)*coth(c+d*x))
@test integrate((a+b*csch(c+d*x)^2)^-7//2, x) == :(a^-7//2*d^-1*arctanh(a^(1/2)*(a+-1b+b*coth(c+d*x)^2)^-1//2*coth(c+d*x))+1//5*b*a^-1*d^-1*(a+-1b)^-1*(a+-1b+b*coth(c+d*x)^2)^-5//2*coth(c+d*x)+1//15*b*a^-3*d^-1*(a+-1b)^-3*(a+-1b+b*coth(c+d*x)^2)^-1//2*(15*b^2+33*a^2+-40*a*b)*coth(c+d*x)+1//15*b*a^-2*d^-1*(a+-1b)^-2*(a+-1b+b*coth(c+d*x)^2)^-3//2*(-5b+9a)*coth(c+d*x))
@test integrate((1+csch(x)^2)^3//2, x) == :(-1//2*(coth(x)^2)^3//2*tanh(x)+(coth(x)^2)^(1/2)*log(sinh(x))*tanh(x))
@test integrate((1+csch(x)^2)^(1/2), x) == :((coth(x)^2)^(1/2)*log(sinh(x))*tanh(x))
@test integrate((1+csch(x)^2)^-1//2, x) == :((coth(x)^2)^-1//2*coth(x)*log(cosh(x)))
@test integrate((1+-1*csch(x)^2)^3//2, x) == :(2*arcsin((1/2)*2^(1/2)*coth(x))+(1/2)*(2+-1*coth(x)^2)^(1/2)*coth(x)+arctanh((2+-1*coth(x)^2)^-1//2*coth(x)))
@test integrate((1+-1*csch(x)^2)^(1/2), x) == :(arcsin((1/2)*2^(1/2)*coth(x))+arctanh((2+-1*coth(x)^2)^-1//2*coth(x)))
@test integrate((1+-1*csch(x)^2)^-1//2, x) == :(arctanh((2+-1*coth(x)^2)^-1//2*coth(x)))
@test integrate((-1+csch(x)^2)^3//2, x) == :(2*arctanh((-2+coth(x)^2)^-1//2*coth(x))+-1//2*(-2+coth(x)^2)^(1/2)*coth(x)+arctan((-2+coth(x)^2)^-1//2*coth(x)))
@test integrate((-1+csch(x)^2)^(1/2), x) == :(-1*arctan((-2+coth(x)^2)^-1//2*coth(x))+-1*arctanh((-2+coth(x)^2)^-1//2*coth(x)))
@test integrate((-1+csch(x)^2)^-1//2, x) == :(arctan((-2+coth(x)^2)^-1//2*coth(x)))
@test integrate((-1+-1*csch(x)^2)^3//2, x) == :((1/2)*(-1*coth(x)^2)^(1/2)*coth(x)+-1*(-1*coth(x)^2)^(1/2)*log(sinh(x))*tanh(x))
@test integrate((-1+-1*csch(x)^2)^(1/2), x) == :((-1*coth(x)^2)^(1/2)*log(sinh(x))*tanh(x))
@test integrate((-1+-1*csch(x)^2)^-1//2, x) == :((-1*coth(x)^2)^-1//2*coth(x)*log(cosh(x)))

max_stars_repo_path: test/6 Hyperbolic functions/6.6 Hyperbolic cosecant/6.6.7 (d hyper)^m (a+b (c csch)^n)^p.jl


text:
export encode

function encode(data::AbstractVector{UInt8})::Tuple{Vector{WaterDrop}, Int}
    if length(data) % PACKET_SIZE != 0
        data = [data; zeros(UInt8, PACKET_SIZE - length(data) % PACKET_SIZE)]
    end

    n_blocks = length(data) ÷ PACKET_SIZE
    n_blocks_sent = Int(n_blocks * WATER_WASTING_FACTOR)

    soliton = SolitonRobust(n_blocks)
    degrees = generate_degrees(soliton, n_blocks_sent)

    drops = Vector{Union{Nothing, WaterDrop}}(nothing, n_blocks_sent)

    for i in 1:n_blocks_sent
        degree = degrees[i]
        seed = UInt32(i)
        indices = generate_indices(seed, n_blocks, degree)
        mixed = foldl((a, b) -> a .⊻ b,
            map(i -> get_block(data, PACKET_SIZE, i), indices);
            init=zeros(UInt8, PACKET_SIZE))
        drops[i] = WaterDrop(seed, UInt32(degree), mixed)
    end

    (drops, n_blocks)
end

max_stars_repo_name: gyk/TrivialSolutions

max_issues_repo_name: gyk/TrivialSolutions

max_forks_repo_name: gyk/TrivialSolutions


text:
include("mybicgstab.jl")

#one-sided FEAST for linear eigenvalue problems
function feast_linear(A,B,x0,nc,emid,ra,rb,eps,maxit;log=false, insideEps=1e-2, set_nin=false)
    (n,m0)=size(x0)

    integrand(z,x,lest,data,resvecs)=\(z*B-A,B*x)

    function rrsolve(Q)
        Aq=Q'*A*Q
        Bq=Q'*B*Q

        le,xe=eigen(Aq,Bq)
        return (le,Q*xe)
    end

    Tf(l,x)=(B*x*diagm( 0 => l)-A*x)

    return feast_core(Tf,integrand,rrsolve,x0,nc,emid,ra,rb,eps,maxit;log=log,insideEps=insideEps, set_nin=set_nin)

end

#two-sided FEAST for linear eigenvalue problems
function feastNS_linear(A,B,x0,y0,nc,emid,ra,rb,eps,maxit; log=false,insideEps=1e-2)
    (n,m0)=size(x0)

    integrand(z,x,lest,data,resvecs)=\(z*B-A,B*x)
    hintegrand(z,y,lest,hdata,hresvecs)=\((z*B-A)',B'*y)

    function rrsolvens(Q,R)
        Aq=R'*A*Q
        Bq=R'*B*Q
        (ye,le,xe)=nseig(Aq,Bq)
        return (R*ye,le,Q*xe)
    end

    Tf(l,x)=(B*x*spdiagm(0 => l)-A*x)
    hTf(l,y)=B'*y*spdiagm(0 => l)'-A'*y

    #use SVD to B-biorthogonalize subspaces
    function biortho(Q,R)
        Bq=R'*B*Q
        (u,s,v)=svd(Bq)
        Y=R*u*diagm( 0 => 1 ./sqrt.(s))
        X=Q*v*diagm( 0 => 1 ./sqrt.(s))
        return (X,Y)
    end

    return feastNS_core(Tf,hTf,integrand,hintegrand,rrsolvens,biortho,x0,y0,nc,emid,ra,rb,eps,maxit; log=log, insideEps=insideEps)

end

#generalized inexact FEAST algorithm with gmres
function ifeast_linear(A,B,x0,alpha,isMaxit,nc,emid,ra,rb,eps,maxit;log=false,insideEps=1e-2, verbose=1)
    (n,m0)=size(x0)

    function integrand(z,x,lest,data,resvecs)
        nc=data[:shiftIndex][z]
        M=(z*B-A)
        #b=(B*x*diagm(lest)-A*x)
        #normb=maximum(sqrt.(sum(abs.(b).^2,1)))

        int=zeros(ComplexF64,n,m0)
        maxits=0
        rhs=convert(Array{ComplexF64,2},resvecs)

        for i in 1:m0
            #(int[:,i],history)=bicgstabl(M,resvecs[:,i],1,max_mv_products=isMaxit,tol=alpha,initial_zero=true,log=true)
            #(int[:,i],history)=idrs(M,rhs[:,i];maxiter=isMaxit,tol=alpha,log=true)
            (int[:,i],history)=gmres(M,resvecs[:,i],restart=isMaxit,tol=alpha,initially_zero=true,maxiter=isMaxit,log=true)
            #(int[:,i],history)=minres(M,resvecs[:,i],tol=alpha,initially_zero=true,maxiter=isMaxit,log=true)
            nlinits=size(history[:resnorm],1)
            data[:linIts][i,nc,data[:iterations]]=nlinits
            #data[:linResiduals][i,nc,data[:iterations]]=history[:resnorm][nlinits]
            data[:linResiduals][i,nc,data[:iterations]]=norm(resvecs[:,i]-M*int[:,i])/norm(resvecs[:,i])

            if(nlinits>maxits)
                maxits=nlinits
            end
        end
        #int=\(M,resvecs)
        #println("      linits=$maxits")
        #int=zbicgstabBlock(M,resvecs,zeros(n,m0),isMaxit,alpha)
        return (x-int)*spdiagm(0 => 1 ./(z.-lest))
    end

    function rrsolve(Q)
        Aq=Q'*A*Q
        Bq=Q'*B*Q
        le,xe=eigen(Aq,Bq)
        return (le,Q*xe)
    end

    Tf(l,x)=(B*x*spdiagm(0 => l)-A*x)

    return feast_core(Tf,integrand,rrsolve,x0,nc,emid,ra,rb,eps,maxit;log=log,insideEps=insideEps, verbose=verbose)

end

#generalized ifeast with bicgstab
function ifeast_linearBicgstab(A,B,x0,alpha,isMaxit,nc,emid,ra,rb,eps,maxit;log=false,insideEps=1e-2,verbose=1)
    (n,m0)=size(x0)

    function integrand(z,x,lest,data,resvecs)
        nc=data[:shiftIndex][z]
        M1=(z*B-A)
        P=speye(n)
        M=P*M1*P
        rhs=P*resvecs
        #b=(B*x*diagm(lest)-A*x)
        #normb=maximum(sqrt.(sum(abs.(b).^2,1)))

        int=zeros(ComplexF64,n,m0)

        for i in 1:m0
            (int[:,i],history)=bicgstabl(M,resvecs[:,i],1,max_mv_products=isMaxit,tol=alpha,initial_zero=true,log=true)
            nlinits=size(history[:resnorm],1)
            data[:linIts][i,nc,data[:iterations]]=nlinits
            data[:linResiduals][i,nc,data[:iterations]]=history[:resnorm][nlinits]
        end

        return (x-int)*spdiagm(0 => 1 ./(z.-lest))
    end

    function rrsolve(Q)
        Aq=Q'*A*Q
        Bq=Q'*B*Q
        le,xe=eigen(Aq,Bq)
        return (le,Q*xe)
    end

    Tf(l,x)=(B*x*spdiagm(0 => l)-A*x)

    return feast_core(Tf,integrand,rrsolve,x0,nc,emid,ra,rb,eps,maxit;log=log,insideEps=insideEps,verbose=verbose)

end

#inexact two-sided ifeast with my own bicgstab implementation
function ifeastNS_linear(A,B,x0,y0,alpha,isMaxit,nc,emid,ra,rb,eps,maxit; log=false,insideEps=1e-2)
    (n,m0)=size(x0)

    #integrand(z,x,lest,data,resvecs)=\(z*B-A,B*x)
    #hintegrand(z,y,lest,hdata,hresvecs)=\((z*B-A)',B'*y)
    function integrand(z,x,lest,data,resvecs)
        M=z*B-A
        rhs=B*x*spdiagm(0 => lest)-A*x
        int0=zeros(ComplexF64,n,m0)
        int=zeros(ComplexF64,n,m0)
        int=zbicgstabBlock(M,resvecs,int0,isMaxit,alpha)
        for i in 1:m0
            #(int[:,i],history)=gmres(M,rhs[:,i],restart=isMaxit,tol=alpha,initially_zero=true,maxiter=isMaxit,log=true)
            #(int[:,i],history)=bicgstabl(M,rhs[:,i],2,max_mv_products=isMaxit,tol=alpha,initial_zero=true,log=true)
        end

        #int=\(M,resvecs)
        return (x-int)*spdiagm(0 => 1 ./(z.-(lest)))
    end

    function hintegrand(z,y,lest,hdata,hresvecs)
        M=(z*B-A)'
        rhs=B'*y*spdiagm(0 => lest)'-A'*y
        int0=zeros(ComplexF64,n,m0)
        int=zeros(ComplexF64,n,m0)
        int=zbicgstabBlock(M,hresvecs,int0,isMaxit,alpha)
        for i in 1:m0
            #(int[:,i],history)=gmres(M,rhs[:,i],restart=isMaxit,tol=alpha,initially_zero=true,maxiter=isMaxit,log=true)

            #(int[:,i],history)=bicgstabl(M,rhs[:,i],2,max_mv_products=isMaxit,tol=alpha,initial_zero=true,log=true)
        end
        #int=\(M,hresvecs)
        return (y-int)*spdiagm(0 => 1 ./(z'.-conj.(lest)))
    end

    function rrsolvens(Q,R)
        Aq=R'*A*Q
        Bq=R'*B*Q
        (ye,le,xe)=nseig(Aq,Bq)
        return (R*ye,le,Q*xe)
    end

    Tf(l,x)=(B*x*spdiagm(0 => l)-A*x)
    hTf(l,y)=B'*y*spdiagm(0 => l)'-A'*y

    #use SVD to B-biorthogonalize subspaces
    function biortho(Q,R)
        Bq=R'*B*Q
        (u,s,v)=svd(Bq)
        Y=R*u*diagm( 0 => 1 ./sqrt.(s))
        X=Q*v*diagm( 0 => 1 ./sqrt.(s))
        return (X,Y)
    end

    return feastNS_core(Tf,hTf,integrand,hintegrand,rrsolvens,biortho,x0,y0,nc,emid,ra,rb,eps,maxit; log=log, insideEps=insideEps)

end

max_stars_repo_name: brendanedwardgavin/feastjl

max_issues_repo_name: brendanedwardgavin/feastjl

max_forks_repo_name: brendanedwardgavin/feastjl


text:
mutable struct MomentumBasedController{N, O<:MOI.AbstractOptimizer, S<:MechanismState}
    state::S
    result::DynamicsResult{Float64, Float64}
    floatingjoint::Union{Nothing, Joint{Float64}}
    floating_joint_velocity_range::UnitRange{Int}
    centroidalframe::CartesianFrame3D
    momentum_matrix::MomentumMatrix{Matrix{Float64}}
    contacts::Dict{RigidBody{Float64}, Vector{ContactPoint{N}}}
    contactwrenches::Dict{BodyID, Wrench{Float64}}
    qpmodel::Parametron.Model{Float64, O}
    v̇::Vector{Variable}
    objective::Parametron.LazyExpression # buffer to incrementally build the objective function
    initialized::Bool

    function MomentumBasedController{N}(
            mechanism::Mechanism{Float64}, optimizer::O; floatingjoint=nothing) where {N, O<:MOI.AbstractOptimizer}
        state = MechanismState(mechanism)
        result = DynamicsResult(mechanism)
        floating_joint_velocity_range = floatingjoint == nothing ? (1 : 0) : velocity_range(state, floatingjoint)
        worldframe = root_frame(mechanism)
        centroidalframe = CartesianFrame3D("centroidal")
        nv = num_velocities(state)
        momentum_matrix = MomentumMatrix(worldframe, zeros(3, nv), zeros(3, nv))
        rootframe = root_frame(mechanism)
        contacts = Dict{RigidBody{Float64}, Vector{ContactPoint{N}}}()
        contactwrenches = Dict{BodyID, Wrench{Float64}}()
        qpmodel = Parametron.Model(optimizer)
        v̇ = [Variable(qpmodel) for _ = 1 : nv]
        objective = Parametron.LazyExpression(identity, zero(QuadraticFunction{Float64}))
        new{N, O, typeof(state)}(
            state, result, floatingjoint, floating_joint_velocity_range, centroidalframe,
            momentum_matrix, contacts, contactwrenches, qpmodel, v̇, objective, false)
    end
end

Base.show(io::IO, controller::MomentumBasedController{N, O, S}) where {N, O, S} =
    print(io, "MomentumBasedController{$N, $O, $S}(…)")

centroidal_frame(controller::MomentumBasedController) = controller.centroidalframe

function (controller::MomentumBasedController)(τ::AbstractVector, t::Number, x::Union{<:Vector, <:MechanismState})
    # initialize on first solve
    if !controller.initialized
        initialize!(controller)
        controller.initialized = true
    end

    # unpack
    qpmodel = controller.qpmodel
    state = controller.state
    result = controller.result
    contacts = controller.contacts
    contactwrenches = controller.contactwrenches
    worldframe = root_frame(state.mechanism)

    # set up and solve qp
    copyto!(state, x)
    solve!(qpmodel)
    checkstatus(qpmodel)

    # extract joint accelerations and contact wrenches
    @inbounds for i in eachindex(controller.v̇)
        result.v̇[i] = value(qpmodel, controller.v̇[i])
    end
    empty!(contactwrenches)
    for body in keys(controller.contacts)
        contactwrench = zero(Wrench{Float64}, worldframe)
        for data in controller.contacts[body]
            contactwrench += value(qpmodel, data.wrench_world)
        end
        contactwrenches[BodyID(body)] = contactwrench
    end

    # compute torques
    inverse_dynamics!(τ, result.jointwrenches, result.accelerations, state, result.v̇, contactwrenches)

    # make floating joint torques precisely zero (they are approximately zero already)
    zero_floating_joint_torques!(τ, state, controller.floating_joint_velocity_range)

    τ
end

function checkstatus(qpmodel::Parametron.Model)
    ok = terminationstatus(qpmodel) == MOI.OPTIMAL && primalstatus(qpmodel) == MOI.FEASIBLE_POINT
    if !ok
        okish = terminationstatus(qpmodel) == MOI.ALMOST_OPTIMAL && primalstatus(qpmodel) == MOI.UNKNOWN_RESULT_STATUS
        if !okish
            throw(QPSolveFailure(terminationstatus(qpmodel), primalstatus(qpmodel), dualstatus(qpmodel)))
        end
    end
end

function zero_floating_joint_torques!(τ::AbstractVector, state::MechanismState, velocity_range::UnitRange)
    for i in velocity_range
        τ[i] = 0
    end
end

function addtask!(controller::MomentumBasedController, task::AbstractMotionTask)
    model = controller.qpmodel
    err = task_error(task, model, controller.state, controller.v̇)
    taskdim = dimension(task)
    @constraint(model, err == zeros(taskdim))
    nothing
end

function addtask!(controller::MomentumBasedController, task::AbstractMotionTask, weight::Union{Number, Parameter{<:Number}})
    e = add_task_error_slack_variables!(controller, task)
    controller.objective = @expression controller.objective + weight * (e ⋅ e)
    e
end

function addtask!(controller::MomentumBasedController, task::AbstractMotionTask, weight::Union{AbstractMatrix, Parameter{<:AbstractMatrix}})
    e = add_task_error_slack_variables!(controller, task)
    controller.objective = @expression controller.objective + transpose(e) * weight * e
    e
end

function add_task_error_slack_variables!(controller::MomentumBasedController, task::AbstractMotionTask)
    model = controller.qpmodel
    state = controller.state
    v̇ = controller.v̇
    e = [Variable(model) for _ = 1 : dimension(task)]
    @constraint(model, e == task_error(task, model, state, v̇))
    e
end

function regularize!(controller::MomentumBasedController, joint::Joint, weight)
    v̇joint = controller.v̇[velocity_range(controller.state, joint)]
    controller.objective = @expression controller.objective + weight * (v̇joint ⋅ v̇joint)
end

function addcontact!(
        controller::MomentumBasedController{N}, body::RigidBody{Float64},
        point::ContactPoint{N}) where N
    push!(get!(Vector{ContactPoint{N}}, controller.contacts, body), point)
    objterm = objectiveterm(point, controller.qpmodel)
    controller.objective = @expression controller.objective + objterm # TODO: currently kind of inefficient; would be better to have a single multi-arg addition
    point
end

function addcontact!(
        controller::MomentumBasedController{N}, body::RigidBody{Float64},
        position::Union{Point3D, Parameter{<:Point3D}},
        normal::Union{FreeVector3D, Parameter{<:FreeVector3D}},
        μ::Union{Float64, Parameter{Float64}}) where N
    addcontact!(controller, body, ContactPoint{N}(position, normal, μ, controller.state, controller.qpmodel))
end

@noinline function initialize!(controller::MomentumBasedController)
    setobjective!(controller)
    mechanism = controller.state.mechanism
    if controller.floatingjoint !== nothing
        add_wrench_balance_constraint!(controller, controller.floatingjoint)
    end
end

function setobjective!(controller::MomentumBasedController)
    @objective(controller.qpmodel, Minimize, controller.objective)
end

function add_wrench_balance_constraint!(controller::MomentumBasedController{N}, joint::Joint) where N
    qpmodel = controller.qpmodel
    state = controller.state
    mechanism = state.mechanism
    fg = mass(mechanism) * mechanism.gravitational_acceleration
    v̇ = controller.v̇
    A = Parameter(A -> momentum_matrix!(A, state), controller.momentum_matrix, qpmodel)
    Ȧv = Parameter{Wrench{Float64}}(() -> momentum_rate_bias(state), qpmodel)
    Wg = Parameter{Wrench{Float64}}(() -> Wrench(center_of_mass(state) × fg, fg), qpmodel)
    S = let state = state, joint = joint, floatingbody = successor(joint, mechanism)
        Parameter(qpmodel) do
            qjoint = configuration(state, joint)
            H = transform_to_root(state, floatingbody)
            transform(motion_subspace(joint, qjoint), H)
        end
    end

    torque = @expression angular(Wg)
    force = @expression linear(Wg)
    for contactvec in values(controller.contacts)
        for contact in contactvec
            torque = @expression torque + angular(contact.wrench_world)
            force = @expression force + linear(contact.wrench_world)
        end
    end

    nv = num_velocities(joint)
    @constraint(qpmodel, transpose(angular(S)) * (angular(A) * v̇ + angular(Ȧv) - torque) + transpose(linear(S)) * (linear(A) * v̇ + linear(Ȧv) - force) == zeros(nv))

    nothing
end