text:
# Test basic methods
@testset "Basics" begin
    # test _macro_error
    @testset "_macro_error" begin
        source = LineNumberNode(42, :bob)
        @test_throws ErrorException InfiniteOpt._macro_error(:test, (:x,), source, "hello")
    end
    # test _esc_non_constant
    @testset "_esc_non_constant" begin
        @test InfiniteOpt._esc_non_constant(42) == 42
        @test InfiniteOpt._esc_non_constant(:(quote 42 end)) == :(quote 42 end)
        @test InfiniteOpt._esc_non_constant(:(x[i])) == esc(:(x[i]))
        @test InfiniteOpt._esc_non_constant(:x) == esc(:x)
    end
    # test _get_name
    @testset "_get_name" begin
        @test InfiniteOpt._get_name(:a) == :a
        @test InfiniteOpt._get_name(nothing) == ()
        @test InfiniteOpt._get_name("bob") == "bob"
        @test InfiniteOpt._get_name(:("a$(x)")) == :("a$(x)")
        @test InfiniteOpt._get_name(:(x[1:2])) == :x
    end
    # test _name_call
    @testset "_name_call" begin
        @test InfiniteOpt._name_call("bob", []) == "bob"
        @test InfiniteOpt._name_call("", []) == ""
        expected = :(string("x", "[", string($(Expr(:escape, "2"))), ",", string($(Expr(:escape, "1"))), "]"))
        @test InfiniteOpt._name_call("x", ["2", "1"]) == expected
    end
    # test _extract_kwargs
    @testset "_extract_kwargs" begin
        # test normal
        args = (:x, :(d=42), :(container=Cat))
        @test InfiniteOpt._extract_kwargs(args)[1] == [:x]
        @test InfiniteOpt._extract_kwargs(args)[2] == [:(d=42)]
        @test InfiniteOpt._extract_kwargs(args)[3] == :(Cat)
        @test InfiniteOpt._extract_kwargs(args)[4] === nothing
        # test with semi-colon keywords
        args = (Expr(:parameters, :(d=42), :(base_name = "cat")), :x, :(container=Cat))
        @test InfiniteOpt._extract_kwargs(args)[1] == [:x]
        @test InfiniteOpt._extract_kwargs(args)[2] == [:(d=42)]
        @test InfiniteOpt._extract_kwargs(args)[3] == :(Cat)
        @test InfiniteOpt._extract_kwargs(args)[4] == esc("cat")
    end
    # test _add_kw_args
    @testset "_add_kwargs" begin
        esc_kw = esc(Expr(:kw, :d, 2))
        call = :(f(a))
        @test InfiniteOpt._add_kwargs(call, [:(d = 2)]) isa Nothing
        @test call == :(f(a, $esc_kw))
    end
     # test _has_idxvars
     @testset "_has_idxvars" begin
        vars = Any[:i, :j]
        @test !InfiniteOpt._has_idxvars(:w, vars)
        @test InfiniteOpt._has_idxvars(:i, vars)
        @test InfiniteOpt._has_idxvars(:(a[i]), vars)
        @test !InfiniteOpt._has_idxvars(:(ones(a)), vars)
    end
    # test _valid_model
    @testset "_valid_model" begin
        m = InfiniteModel()
        m2 = 2
        @test InfiniteOpt._valid_model(error, m, :m) isa Nothing
        @test_throws ErrorException InfiniteOpt._valid_model(error, m2, :m2)
    end
    # test _error_if_cannot_register
    @testset "_error_if_cannot_register" begin
        m = InfiniteModel()
        @test InfiniteOpt._error_if_cannot_register(error, m, :x) isa Nothing
        object_dictionary(m)[:x] = 2
        @test_throws ErrorException InfiniteOpt._error_if_cannot_register(error, m, :x)
        @test_throws ErrorException InfiniteOpt._error_if_cannot_register(error, m, :((x, 2)))
    end
    # test _macro_assign_and_return
    @testset "_macro_assign_and_return" begin
        m = InfiniteModel()
        esc_m = esc(:m)
        code = :(GeneralVariableRef($esc_m, 1, FiniteVariableIndex))
        @test InfiniteOpt._macro_assign_and_return(error, code, :a, esc_m) isa Expr
    end
    # test _finalize_macro
    @testset "_finalize_macro" begin
        m = InfiniteModel()
        esc_m = esc(:m)
        code = :(GeneralVariableRef($esc_m, 1, FiniteVariableIndex))
        source = LineNumberNode(42, :bob)
        @test InfiniteOpt._finalize_macro(error, esc_m, code, source) isa Expr
    end
end

# Test container methods
@testset "Container Helpers" begin
    # test _parse_idxset
    @testset "_parse_idxset" begin
        @test InfiniteOpt._parse_idxset(:(i = 1:2)) == (:i, esc(:(1:2)))
        @test InfiniteOpt._parse_idxset(:(i in 1:2)) == (:i, esc(:(1:2)))
        @test InfiniteOpt._parse_idxset(:(1:2))[1] isa Symbol
        @test InfiniteOpt._parse_idxset(:(1:2))[2] == esc(:(1:2))
        @test InfiniteOpt._parse_idxset(:a)[1] isa Symbol
        @test InfiniteOpt._parse_idxset(:a)[2] == esc(:a)
    end
    # test _explicit_oneto
    @testset "_explicit_oneto" begin
        @test InfiniteOpt._explicit_oneto(:(1:2)) == :(Base.OneTo(1:2))
        s = esc(:(1:2))
        @test InfiniteOpt._explicit_oneto(s) == :(Base.OneTo($s))
        @test InfiniteOpt._explicit_oneto(:(2:4)) == :(2:4)
    end
    # test _expr_is_splat
    @testset "_expr_is_splat" begin
        @test !InfiniteOpt._expr_is_splat(:(1:2))
        @test !InfiniteOpt._expr_is_splat(esc(:(1:2)))
        @test InfiniteOpt._expr_is_splat(:(a...))
        @test !InfiniteOpt._expr_is_splat(:a)
    end
    # test _parse_ref_sets
    @testset "_parse_ref_sets" begin
        # test errors
        bad_vcat = :([1:2; k; j])
        @test_throws ErrorException InfiniteOpt._parse_ref_sets(error, bad_vcat)
        bad_vect = :([1:2, 1:2; k; j])
        @test_throws ErrorException InfiniteOpt._parse_ref_sets(error, bad_vect)
        duplicate_ref = :(a[i in 1:2, i in 1:2])
        @test_throws ErrorException InfiniteOpt._parse_ref_sets(error, duplicate_ref)
        bad_type = :((i in 2:3))
        @test_throws ErrorException InfiniteOpt._parse_ref_sets(error, bad_type)
        # test vcat
        vcat = :([i = s; i <= 2])
        @test InfiniteOpt._parse_ref_sets(error, vcat)[1] == [:i]
        @test InfiniteOpt._parse_ref_sets(error, vcat)[2] == [esc(:s)]
        @test InfiniteOpt._parse_ref_sets(error, vcat)[3] == :(i <= 2)
        # test typed_vcat
        typed_vcat = :(a[i = 2:w; k])
        @test InfiniteOpt._parse_ref_sets(error, typed_vcat)[1] == [:i]
        @test InfiniteOpt._parse_ref_sets(error, typed_vcat)[2] == [esc(:(2:w))]
        @test InfiniteOpt._parse_ref_sets(error, typed_vcat)[3] == :k
        # test ref w/o condition
        ref = :(a[1:2, i = S])
        @test length(InfiniteOpt._parse_ref_sets(error, ref)[1]) == 2
        @test InfiniteOpt._parse_ref_sets(error, ref)[2] == [esc(:(1:2)), esc(:S)]
        @test InfiniteOpt._parse_ref_sets(error, ref)[3] == :()
        # test ref w/ condition
        ref = :(a[i in 2:4, j = [:d, :e]; i >= 3])
        @test InfiniteOpt._parse_ref_sets(error, ref)[1] == [:i, :j]
        @test InfiniteOpt._parse_ref_sets(error, ref)[2] == [esc(:(2:4)), esc(:([:d, :e]))]
        @test InfiniteOpt._parse_ref_sets(error, ref)[3] == :(i >= 3)
        # test vect w/o condition
        vect = :([s in S, 4:5])
        @test InfiniteOpt._parse_ref_sets(error, vect)[1][1] == :s
        @test InfiniteOpt._parse_ref_sets(error, vect)[1][2] isa Symbol
        @test InfiniteOpt._parse_ref_sets(error, vect)[2] == [esc(:S), esc(:(4:5))]
        @test InfiniteOpt._parse_ref_sets(error, vect)[3] == :()
        # test vect w/ condition
        vect = :([s in S, j = 4:5; s + j == 6])
        @test InfiniteOpt._parse_ref_sets(error, vect)[1] == [:s, :j]
        @test InfiniteOpt._parse_ref_sets(error, vect)[2] == [esc(:S), esc(:(4:5))]
        @test InfiniteOpt._parse_ref_sets(error, vect)[3] == :(s + j == 6)
        # test symbol
        @test InfiniteOpt._parse_ref_sets(error, :a) == (Any[], Any[], :())
    end
    # test _build_ref_sets
    @testset "_build_ref_sets" begin
        # test errors
        @test_throws ErrorException InfiniteOpt._build_ref_sets(error, :([a...]))
        bad_vcat = :([1:2; k; j])
        @test_throws ErrorException InfiniteOpt._build_ref_sets(error, bad_vcat)
        # test dependent build
        cont = JuMPC.SparseAxisArray(Dict((1, 1) => 2, (2, 1) => 3, (2, 2) => 4))
        @test @gen_container([i = 1:2, j = 1:i], i + j) == cont
        # test conditioned build
        cont = JuMPC.SparseAxisArray(Dict((1, 2) => 3, (1, 3) => 4, (2, 2) => 4))
        @test @gen_container([i = 1:2, j = 2:3; i + j <= 4], i + j) == cont
        # test easy build
        cont = JuMPC.DenseAxisArray(ones(2, 2), Base.OneTo(2), 2:3)
        @test @gen_container([1:2, j = 2:3], 1.0) == cont
    end
end


text:
const mpl = matplotlib

# const Line2D = matplotlib.lines.Line2D
# const Patch = matplotlib.patches.Patch

# const GridSpec = mpl.gridspec.GridSpec

plt.style.use("seaborn-colorblind")  # "ggplot"
palette = plt.cm.Dark2.colors

font = Dict(:family => "STIXGeneral", :size => 16)
savefig = Dict(:dpi => 600, :bbox => "tight")
lines = Dict(:linewidth => 4)
figure = Dict(:figsize => (8, 5))
axes = Dict(:prop_cycle => mpl.cycler(; color=palette))
legend = Dict(:fontsize => "x-large")  # medium for presentations, x-large for papers

mpl.rc("font"; font...)
mpl.rc("savefig"; savefig...)
mpl.rc("lines"; lines...)
mpl.rc("figure"; figure...)
mpl.rc("axes"; axes...)
mpl.rc("legend"; legend...)

function fun_plotter(fun, array; xlim=(0, 0))
    output = map(fun, eachrow(array)...)
    return lineplot(output; title="Custom Function", name="fun", ylim=extrema(output), xlim)
end

# handy terminal plots
function unicode_plotter(states, controls; only=nothing, vars=nothing, fun=nothing)
    @argcheck size(states, 2) == size(controls, 2)
    xlim = (0, size(states, 2))
    if only == :states
        if !isnothing(fun)
            return fun_plotter(fun, states; xlim)
        end
        typeof(vars) <: Vector && (states = @view states[vars, :])
        ylim = states |> extrema
        tag = isnothing(vars) ? "x1" : "x$(vars[1])"
        plt = lineplot(
            states[1, :]; title="State Evolution", name=tag, xlabel="step", ylim, xlim
        )
        for (i, s) in enumerate(eachrow(states[2:end, :]))
            tag = isnothing(vars) ? "x$(i+1)" : "x$(vars[i+1])"
            lineplot!(plt, collect(s); name=tag)
        end
    elseif only == :controls
        if !isnothing(fun)
            return fun_plotter(fun, controls; xlim)
        end
        typeof(vars) <: Vector && (controls = @view controls[vars, :])
        ylim = controls |> extrema
        tag = isnothing(vars) ? "c1" : "c$(vars[1])"
        plt = lineplot(
            controls[1, :]; title="Control Evolution", name=tag, xlabel="step", ylim, xlim
        )
        for (i, s) in enumerate(eachrow(controls[2:end, :]))
            tag = isnothing(vars) ? "c$(i+1)" : "c$(vars[i+1])"
            lineplot!(plt, collect(s); name=tag)
        end
    else
        if !isnothing(fun)
            return fun_plotter(fun, hcat(states, controls); xlim)
        end
        ylim = Iterators.flatten((states, controls)) |> extrema
        plt = lineplot(
            states[1, :];
            title="State and Control Evolution",
            name="x1",
            xlabel="step",
            ylim,
            xlim,
        )
        for (i, s) in enumerate(eachrow(states[2:end, :]))
            lineplot!(plt, collect(s); name="x$(i+1)")
        end
        for (i, c) in enumerate(eachrow(controls))
            lineplot!(plt, collect(c); name="c$i")
        end
    end
    return plt
end

function plot_simulation(
    controlODE,
    params;
    show=nothing,
    only=nothing,
    vars=nothing,
    fun=nothing,
    yrefs=nothing,
)
    !isnothing(show) && @info show

    # TODO: use times in plotting?
    times, states, controls = run_simulation(controlODE, params)
    plt = unicode_plotter(states, controls; only, vars, fun)
    if !isnothing(yrefs)
        for yref in yrefs
            lineplot!(plt, x -> yref; name="$yref")
        end
    end
    display(plt)
    return false  # if return true, then optimization stops
end

function plot_collocation(controls_collocation, interpol, tsteps)
    plt.figure()
    finer_tsteps = range(tsteps[1], tsteps[end]; length=1000)
    plt.plot(finer_tsteps, [interpol(t) for t in finer_tsteps]; label="interpolation")
    plt.plot(tsteps, controls_collocation, "xg"; label="collocation")
    plt.title("Control collocation")
    plt.xlabel("time")
    plt.legend()
    return plt.show()
end

abstract type PhasePlotMarker end

@kwdef struct IntegrationPath <: PhasePlotMarker
    points
    fmt = "m:"
    label = "Integration path"
    markersize = nothing
    linewidth = 4
end

@kwdef struct InitialState <: PhasePlotMarker
    points
    fmt = "bD"
    label = "Initial state"
    markersize = 12
    linewidth = nothing
end

@kwdef struct FinalState <: PhasePlotMarker
    points
    fmt = "r*"
    label = "Final state"
    markersize = 18
    linewidth = nothing
end

@kwdef struct ShadeConf
    indicator::Function
    cmap = "gray"
    transparency = 1
end

@kwdef struct ConstRef
    val::Real
    direction::Symbol
    class::Symbol
    var::Integer
    linestyle = "--"
    color = "r"
    label = "Constraint"
    linewidth = 3
    min = 0
    max = 1
    function ConstRef(
        val, direction, class, var, linestyle, color, label, linewidth, min, max
    )
        @argcheck direction in (:vertical, :horizontal)
        @argcheck class in (:state, :control)
        @argcheck 0 <= min < max <= 1
        return new(val, direction, class, var, linestyle, color, label, linewidth, min, max)
    end
end

@kwdef struct FuncRef
    fn::Function
    dom::Symbol
    class::Symbol
    var::Integer
    linestyle = "--"
    color = "r"
    label = "Constraint"
    linewidth = 3
    function FuncRef(fn, dom, class, var, linestyle, color, label, linewidth)
        @argcheck dom in (:space, :time)
        @argcheck class in (:state, :control)
        return new(fn, dom, class, var, linestyle, color, label, linewidth)
    end
end

function phase_portrait(
    controlODE,
    params,
    phase_time,
    coord_lims;  #xlims, ylims
    points_per_dim=1000,
    dimension=2,
    projection=[1, 2],
    markers=nothing,
    start_points=nothing,
    start_points_x=nothing,
    start_points_y=nothing,
    title=nothing,
    shader=nothing,
    kwargs...,
)
    @argcheck length(projection) == 2
    @argcheck all(x -> isa(x, Tuple) && length(x) == 2, coord_lims)
    @argcheck all(marker isa PhasePlotMarker for marker in markers)

    function stream_interface(coords...)
        u = zeros(Float32, dimension)
        du = zeros(Float32, dimension)
        copyto!(u, coords)
        # du = deepcopy(coords)
        controlODE.system!(du, u, params, phase_time, controlODE.controller)
        return du
    end

    # evaluate system over each combination of coords in the specified ranges

    # NOTE: float64 is relevant for the conversion to pyplot due to inner
    #       numerical checks of equidistant input in the streamplot function
    ranges = [range(Float64.(lims)...; length=points_per_dim) for lims in coord_lims]
    xpoints, ypoints = collect.(ranges[projection])

    # NOTE: the transpose is required to get f.(a',b) instead of the default f.(a, b')
    xgrid, ygrid = ndgrid(xpoints, ypoints)  # NOTE: is this switched?
    phase_array_tuples = stream_interface.(xgrid, ygrid)'
    # phase_array_tuples = stream_interface.(xpoints', ypoints)

    xphase, yphase = [getindex.(phase_array_tuples, dim) for dim in projection]

    magnitude = map((x, y) -> sqrt(sum(x^2 + y^2)), xphase, yphase)

    fig = plt.figure()
    # gs = GridSpec(nrows=3, ncols=2, height_ratios=[1, 1, 2])

    if isnothing(start_points) && !isnothing(start_points_x) && !isnothing(start_points_y)
        @argcheck size(start_points_x) == size(start_points_y)
        start_grid_x, start_grid_y = ndgrid(start_points_x, start_points_y)
        start_points = hcat(reshape(start_grid_x, :, 1), reshape(start_grid_y, :, 1))
    end

    # integration_direction = isnothing(start_points) ? "both" : "forward"
    ax = fig.add_subplot()
    strm = ax.streamplot(
        xpoints,
        ypoints,
        xphase,
        yphase;
        color=magnitude,
        linewidth=1.5,
        density=1.5,
        cmap="summer",
        kwargs...,
    )
    if !isnothing(start_points)
        start_points = start_points[:, projection]
        ax.plot(start_points[:, 1], start_points[:, 2], "kX"; markersize=12)
        strm = ax.streamplot(
            xpoints,
            ypoints,
            xphase,
            yphase;
            color="darkorchid",
            linewidth=3,
            density=20,
            start_points,
            integration_direction="forward",
            kwargs...,
        )
    end

    # displaying points (handles multiple points as horizontally concatenated)
    if !isnothing(markers)
        for plotconf in markers
            points_projected = plotconf.points[projection, :]
            ax.plot(
                points_projected[1, :],
                points_projected[2, :],
                plotconf.fmt;
                label=plotconf.label,
                markersize=plotconf.markersize,
                linewidth=plotconf.linewidth,
            )
        end
    end

    xlims, ylims = coord_lims[projection]

    if !isnothing(shader)
        mask = shader.indicator.(xgrid, ygrid)'
        ax.imshow(
            mask;
            extent=(xlims..., ylims...),
            alpha=shader.transparency,
            cmap=shader.cmap,
            aspect="auto",
        )
    end

    ax.set(; xlim=xlims .+ (-0.05, 0.05), ylim=ylims .+ (-0.05, 0.05))
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    !isnothing(title) && ax.set_title(title)

    fig.colorbar(strm.lines)
    ax.legend()

    # remove frame
    ax.spines["top"].set_visible(false)
    ax.spines["right"].set_visible(false)
    ax.spines["bottom"].set_visible(false)
    ax.spines["left"].set_visible(false)

    plt.tight_layout()

    return plt.show()
end

function transparecy_scaler_abs(noise, perturbations; top=1, low=0.2)
    highest = maximum(abs, perturbations)
    amplitude = abs(noise) / highest
    return top - (top - low) * amplitude
end

function set_state_control_subplots(
    num_states, num_controls; annotation=nothing, refs=nothing
)
    if !isnothing(refs)
        @argcheck all(ref isa ConstRef for ref in refs)
        @argcheck all(
            ifelse(ref.class == :state, ref.var in 1:num_states, ref.var in 1:num_controls)
            for ref in refs
        ) BoundsError
    end

    fig_states, axs_states = plt.subplots(
        num_states;
        sharex="col",
        squeeze=false,
        # constrained_layout=true,  # incompatible with subplots_adjust
    )
    fig_controls, axs_controls = plt.subplots(
        num_controls;
        sharex="col",
        squeeze=false,
        # constrained_layout=true,  # incompatible with subplots_adjust
    )
    axs_states[1].set_title("States")
    axs_controls[1].set_title("Controls")
    axs_states[end].set_xlabel("time")
    axs_controls[end].set_xlabel("time")

    for s in 1:num_states
        axs_states[s].set_ylabel("state[$s]")
    end
    for c in 1:num_controls
        axs_controls[c].set_ylabel("control[$c]")
    end
    if !isnothing(refs)
        for r in refs
            ax = r.class == :state ? axs_states[r.var] : axs_controls[r.var]
            if r.direction == :vertical
                ax.axvline(
                    r.val;
                    ymin=r.min,
                    ymax=r.max,
                    label=r.label,
                    linestyle=r.linestyle,
                    color=r.color,
                )
            elseif r.direction == :horizontal
                ax.axhline(
                    r.val;
                    xmin=r.min,
                    xmax=r.max,
                    label=r.label,
                    linestyle=r.linestyle,
                    color=r.color,
                )
            end
        end
    end

    if !isnothing(annotation)
        fig_states.text(0, 0, string(annotation))
        fig_controls.text(0, 0, string(annotation))
    end

    return fig_states, axs_states, fig_controls, axs_controls
end

function plot_initial_perturbations(
    controlODE, θ, specs; refs=nothing, funcs=nothing
)
    state_size = size(u0, 1)
    control_size = size(controller(u0, θ), 1)

    for spec in specs
        perturbations = local_grid(
            spec.samples, spec.percentage; scale=spec.scale, type=spec.type
        )
        cmap = ColorMap("tab10")
        fig_states, axs_states, fig_controls, axs_controls = set_state_control_subplots(
            length(u0), length(controller(u0, θ)); annotation=spec, refs
        )

        for noise in perturbations
            noise_vec = zeros(typeof(noise), length(u0))
            noise_vec[spec.variable] = noise

            perturbed_u0 = u0 + noise_vec
            controlODE = ControlODE(
                controlODE.controller,
                controlODE.system!,
                perturbed_u0,
                controlODE.tspan;
                tsteps=controlODE.tsteps,
                integrator=controlODE.integrator,
                sensealg=controlODE.sensealg
            )

            times, states, controls = run_simulation(controlODE, θ)

            for s in 1:state_size
                axs_states[s].plot(
                    times,
                    states[s, :];
                    label="u0[$(spec.variable)] + " * format(noise; precision=2),
                    alpha=transparecy_scaler_abs(noise, perturbations),
                    c=cmap(s),
                )
            end
            for c in 1:control_size
                axs_controls[c].plot(
                    times,
                    controls[c, :];
                    label="u0[$(spec.variable)] + " * format(noise; precision=2),
                    alpha=transparecy_scaler_abs(noise, perturbations),
                    c=cmap(c + size(states, 1)),
                )
            end
        end
        legend_elements = [
            matplotlib.patches.Patch(;
                facecolor="black",
                edgecolor="black",
                label=format(noise; precision=2),
                alpha=transparecy_scaler_abs(noise, perturbations),
            ) for noise in sort(perturbations)
        ]
        fig_legend_div = 0.8
        fig_states.subplots_adjust(; right=fig_legend_div)
        legend_states = fig_states.legend(;
            handles=legend_elements,
            bbox_to_anchor=(fig_legend_div, 0.5),
            loc="center left",
            # borderaxespad=0,
            title="Perturbation",
        )
        fig_controls.subplots_adjust(; right=fig_legend_div)
        legend_controls = fig_controls.legend(;
            handles=legend_elements,
            bbox_to_anchor=(fig_legend_div, 0.5),
            loc="center left",
            # borderaxespad=0,
            title="Perturbation",
        )

        fig_states.show()
        fig_controls.show()

        # tight_layout alternative that considers the legend (or other artists)
        # bbox_extra_artists must be an iterable
        # fig.savefig("states_noise", bbox_extra_artists=(legend_states,), bbox_inches="tight")
        # fig.savefig("controls_noise", bbox_extra_artists=(legend_controls,), bbox_inches="tight")
    end
end

function plot_initial_perturbations_collocation(
    controlODE,
    θ,
    specs,
    collocation::Function;
    refs=nothing,
    funcs=nothing,
    storedir=nothing,
)
    state_size = size(controlODE.u0, 1)
    control_size = size(controlODE.controller(controlODE.u0, θ), 1)

    u0 = controlODE.u0
    for spec in specs
        perturbations = local_grid(
            spec.samples, spec.percentage; scale=spec.scale, type=spec.type
        )
        for (tag, noise) in enumerate(perturbations)
            cmap = ColorMap("tab20")
            fig_states, axs_states, fig_controls, axs_controls = set_state_control_subplots(
                length(controlODE.u0), length(controlODE.controller(controlODE.u0, θ)); annotation=spec, refs
            )

            noise_vec = zeros(typeof(noise), length(controlODE.u0))
            noise_vec[spec.variable] = noise

            # perturbed prob = ODEProblem(dudt!, u0 + noise_vec, tspan, θ)
            perturbed_u0 = u0 + noise_vec
            controlODE = ControlODE(
                controlODE.controller,
                controlODE.system!,
                perturbed_u0,
                controlODE.tspan;
                tsteps=controlODE.tsteps,
                integrator=controlODE.integrator,
                sensealg=controlODE.sensealg
            )

            @info "Simulation"
            @time times, states, controls = run_simulation(controlODE, θ)
            @info "Collocation"
            @time infopt_model, times_collocation, states_collocation, controls_collocation = collocation(
                perturbed_u0; constrain_states=true
            )
            @info "Interpolation"
            @time interpol = chebyshev_interpolation(controlODE.tsteps, controls_collocation)

            for s in 1:state_size
                axs_states[s].plot(
                    times,
                    states[s, :];
                    label="policy",
                    # alpha=transparecy_scaler_abs(noise, perturbations),
                    color=cmap(2s - 2),
                )
                axs_states[s].plot(
                    times,
                    states_collocation[s, :];
                    label="collocation",
                    marker="x",
                    linestyle="None",
                    # alpha=transparecy_scaler_abs(noise, perturbations),
                    color=cmap(2s - 1),
                )
            end
            for c in 1:control_size
                axs_controls[c].plot(
                    times,
                    controls[c, :];
                    label="policy",
                    # alpha=transparecy_scaler_abs(noise, perturbations),
                    color=cmap(2c - 2 + 2size(states, 1)),
                )
                axs_controls[c].plot(
                    times,
                    controls_collocation[c, :];
                    label="collocation",
                    marker="x",
                    linestyle="None",
                    # alpha=transparecy_scaler_abs(noise, perturbations),
                    color=cmap(2c - 1 + 2size(states, 1)),
                )
            end
            title = "u0[$(spec.variable)] + " * format(noise; precision=2)
            fig_states.suptitle(title)
            fig_controls.suptitle(title)

            legend_elements = [
                matplotlib.lines.Line2D(
                    [0], [0]; color="black", linewidth=3, label="policy"
                ),
                matplotlib.lines.Line2D(
                    [0],
                    [0];
                    color="black",
                    marker="x",
                    linestyle="None",
                    linewidth=3,
                    label="collocation",
                ),
            ]

            # fig_legend_div = 0.8
            # fig_states.subplots_adjust(; right=fig_legend_div)
            legend_states = fig_states.legend(;
                handles=legend_elements,
                # bbox_to_anchor=(fig_legend_div, 0.5),
                # loc="center left",
            )
            # fig_controls.subplots_adjust(; right=fig_legend_div)
            legend_controls = fig_controls.legend(;
                handles=legend_elements,
                # bbox_to_anchor=(fig_legend_div, 0.5),
                # loc="center left",
            )

            if !isnothing(storedir)
                # tight_layout alternative that considers the legend (or other artists)
                # bbox_extra_artists must be an iterable
                filename = "u0_$(spec.variable)_" * sprintf1("%04d", tag)
                fig_states.savefig(
                    joinpath(storedir, filename * "_states.png");
                    bbox_extra_artists=(legend_states,),
                    # bbox_inches="tight",  # this adjusts filesize slighly... problematic for gifs
                )
                fig_controls.savefig(
                    joinpath(storedir, filename * "_controls.png");
                    bbox_extra_artists=(legend_controls,),
                    # bbox_inches="tight",
                )
            else
                fig_states.show()
                fig_controls.show()
            end
        end
    end
end

ext: jl


text:
addprocs(4)

function test()
  a = randn(1000)
  @parallel (+) for i=1:100000
    fetch(a[randi(end)])
  end
end

max_stars_repo_name: dingzhw/RotorTrim

max_issues_repo_name: helicopterliu/RotorTrim

max_forks_repo_name: helicopterliu/RotorTrim


text:
#=
Description
    Module for QMDP solver to generate approximate POMDP solution for
    the DoubleUAV problem.
=#

module QMDP

using POMDPs, DoubleUAVs, GridInterpolations

export qmdp

const MAX_ITER = 1000
const GAMMA = 0.95
const ALPHA_TOL = 1e-6
const STATE_DIM = 5
const TERM_STATE_VAR = 1e5
const TERM_STATE = TERM_STATE_VAR * ones(STATE_DIM)

function get_max_alphas!(V::Vector{Float64}, alpha::Matrix{Float64},
                         nStates::Int64)
    nActions = size(alpha, 2)
    for istate = 1:nStates
        V[istate] = alpha[istate, 1]
        for iaction = 2:nActions
            if V[istate] < alpha[istate, iaction]
                V[istate] = alpha[istate, iaction]
            end # if
        end # for iaction
    end # for istate
end # function get_max_alphas!

function qmdp(states::Matrix, actions::Matrix, sigmas::Matrix{Float64},
              weights::Vector{Float64}, verbose::Bool, maxIter::Int64,
              gamma::Float64, alphaTol::Float64, lambda::Float64)
    nStates = size(states, 2)
    nActions = size(actions, 2)
    nSigmas = size(sigmas, 2)
    g = RectangleGrid(sort(unique(vec(states[1, 1:end - 1]))),
                      sort(unique(vec(states[2, 1:end - 1]))),
                      sort(unique(vec(states[3, 1:end - 1]))),
                      sort(unique(vec(states[4, 1:end - 1]))),
                      sort(unique(vec(states[5, 1:end - 1]))))
    alpha = zeros(nStates, nActions)
    V = zeros(nStates)
    get_reward = gen_reward(lambda)

    @printf("Running QMDP alpha vectors approximation...\n")
    println("lambda = ", lambda)
    cputime = 0
    for iter = 1:maxIter
        tic()
        residual = 0
        for istate = 1:nStates
            state = TERM_STATE
            if istate != nStates # not terminal state
                state = ind2x(g, istate)
            end # if

            for iaction = 1:nActions
                action = actions[:, iaction]
                vnext = 0

                for isigma = 1:nSigmas
                    snext = get_next_state(state, iaction, isigma)
                    if snext[1] == TERM_STATE_VAR
                        vnext = vnext + V[end] * weights[isigma]
                    else
                        vnext = vnext + interpolate(g, V, snext) * weights[isigma]
                    end # if
                end # for isigma

                alphaPrev = alpha[istate, iaction]
                alpha[istate, iaction] = get_reward(state, iaction) +
                                         gamma * vnext
                residual = residual + (alpha[istate, iaction] - alphaPrev)^2
            end # for iaction
        end # for istate

        get_max_alphas!(V, alpha, nStates)

        iterTime = toq()
        cputime = cputime + iterTime
        if verbose
            @printf("Iteration %d: residual = %.2e, cputime = %.2e\n",
                    iter, residual, iterTime)
        end # if
        if residual < alphaTol; break; end
        if iter == maxIter
            println("Warning: maximum number of iterations reached;",
                    "solution may be inaccurate")
        end # if
    end # for iter
    print("lambda = ", lambda, ": ")
    @printf("QMDP done!\ncputime = %.2e sec\n\n", cputime)
    return alpha
end # function qmdp

function qmdp(d::DoubleUAV, lambda::Float64, verbose::Bool=true,
              maxIter::Int64=MAX_ITER, gamma::Float64 = GAMMA,
              alphaTol::Float64=ALPHA_TOL)
    p = d.pomdp
    alpha = qmdp(p.states, p.actions, d.sigmas, d.weights, verbose, maxIter,
                 gamma, alphaTol, lambda)
    return alpha
end # function qmdp

end # module QMDP


text:
struct Lattice{D,T,M,U}
  site::Vector{Site{D,T}}  # sorted by ID
  coord_order::Vector{Site{D,T}} # sort by coord
  neighbors::Vector{Vector{Int}}
  types::Vector{U}
end

function Lattice(coord, types::Vector;
    nbhood = VonNeumann(),
    type_list = unique(types)
  )
  dimension = coord isa Matrix ? size(coord, 1) : length(coord[1])
  # unique_types = sort!(type_list)
  unique_types = type_list
  number_types = length(unique_types)
  number_neighbors = capacity(nbhood, dimension)

  labels = OrderedDict(unique_types[i] => i+1 for i in eachindex(unique_types))
  site = build_sites(coord, types, labels)

  site_by_coord = sort(site, by = coordinates)

  neighbors = [sizehint!(Int[], number_neighbors) for i in eachindex(site)]

  mapping = [(label => l) for (label, l) in labels]
  T = coord isa Matrix ? eltype(coord) : eltype(coord[1])
  U = eltype(mapping)

  return Lattice{dimension,T,typeof(nbhood),U}(site, site_by_coord, neighbors, mapping)
end

function build_sites(coord::Matrix, types, labels)
  site = [Site(i, State(labels[types[i]]), tuple(coord[:,i]...)) for i in eachindex(types)]
end

function build_sites(coord::Vector{NTuple{N,T}}, types, labels) where {N,T}
  site = [Site(i, State(labels[types[i]]), coord[i]) for i in eachindex(types)]
end

##### accessors

dimension(::Lattice{D}) where D = D
topology(::Lattice{D,T,M}) where {D,T,M} = M()

number_types(x::Lattice) = length(x.types)
number_sites(x::Lattice) = length(x.site)

##### IO
function Base.summary(io::IO, ::Lattice{D,T,M}) where {D,T,M}
  print(io, "$(D)-D Lattice with $(M) neighborhoods")
end

function Base.show(io::IO, lattice::Lattice)
  Base.summary(io, lattice)
  print(io, "\n  species: ")
  show(io, lattice.types)
end

function Base.show(io::IO, m::MIME"text/plain", lattice::Lattice)
  Base.summary(io, lattice)
  print(io, "\n", "species: ")
  show(io, m, lattice.types)
end

##### other Base overloads
Base.copy(lattice::Lattice) = deepcopy(lattice)

# this is a dumb hack to make Lattice compatible with SamplePath
Base.size(lattice::Lattice) = size(lattice.types)

function __simple_copy(lattice)
  coord = coordinates.(lattice.site)
  tcode = get_ptype.(lattice.site)

  idx = findall(!isequal(1), tcode)

  return coord[idx], tcode[idx]
end

##### query API

function istracked(lattice::Lattice, coord)
  idx = searchsorted(lattice.coord_order, coord, by = coordinates, lt = <)

  return !isempty(idx)
end

# get the site living at the given coordinates
function get_site(lattice::Lattice, coord)
  idx = searchsorted(lattice.coord_order, coord, by = coordinates, lt = <)

  lattice.coord_order[idx[1]]
end

# get the site with the given id
function get_site(lattice::Lattice, id::Int)
  # idx = searchsorted(lattice.site, id, by = label)
  #
  # # if idx is empty, site was not found
  #
  # # if idx has length > 1, we have a duplicate
  #
  # # if idx has length = 1, return the object
  # lattice.site[idx[1]]
  lattice.site[id]
end

function neighborhood(lattice::Lattice, x::Site)
  lattice.neighbors[label(x)]
end

##### update API

function spawn_new_site(lattice::Lattice, coord)
  id = number_sites(lattice) + 1

  return Site(id, State(), coord)
end

# this function assumes new_site is not tracked!
function add_site!(lattice::Lattice, new_site::Site)
  # add neighborhood for site
  nbmax = capacity(topology(lattice), dimension(lattice))
  push!(lattice.neighbors, sizehint!(Int[], nbmax))

  # new sites are always added to the end of the list sorted by IDs
  push!(lattice.site, new_site)

  # add site to list sorted by coord
  idx = searchsortedfirst(lattice.coord_order, new_site, by = coordinates)
  insert!(lattice.coord_order, idx[1], new_site)

  new_site
end

function add_neighbor!(lattice::Lattice{D}, x::Site{D}, y::Site{D}) where D
  nb = neighborhood(lattice, x)

  nbtype = topology(lattice)
  neighbor_count  = length(nb)
  nbhood_capacity = capacity(nbtype, D)

  if label(y) in nb
    error("$(y) is already a neighbor of $(x)")
  end

  if neighbor_count ≥ nbhood_capacity
    msg = """
    Neighborhood of $(x) at capacity ($(nbhood_capacity)).
    Failed to add $(y) as a neighbor.
    """
    nb_y = neighborhood(lattice, y)

    @debug """
    Neighbor data:

    x: $(label(x)) / $(get_ptype(x))|$(get_neighbor_class(x)) @ $(x)

    y: $(label(y)) / $(get_ptype(y))|$(get_neighbor_class(y)) @ $(y)

    Neighborhood:

    x:  $(nb)

    y:  $(nb_y)

    Sites:

    x:  $(lattice.site[nb])
    y:  $(lattice.site[nb_y])
    """
    throw(ErrorException(msg))
  end

  # store neighbors in sorted fashion
  i = label(y)
  idx = searchsortedfirst(nb, i)
  insert!(nb, idx, i)
end

function rmv_neighbor!(lattice::Lattice, x::Site, y::Site)
  i = label(x)
  j = label(y)

  nb = neighborhood(lattice, x)
  idx = findfirst(isequal(j), nb)

  if idx isa Nothing
    msg = """
    $(y) is not adjacent to $(x) under $(topology(lattice)) structure.
    Failed to remove neighbor.
    """
    throw(ErrorException(msg))
  end

  # delete item, which preserves ordering
  deleteat!(nb, idx)
end

##### recipes

@recipe function f(lattice::Lattice{D,T,VonNeumann}) where {D,T}
  site = lattice.site
  types = lattice.types

  seriestype := :scatter
  aspect_ratio --> 1
  markerstrokewidth --> 0
  markershape --> :square
  # color --> map(get_ptype, site)
  # label --> map(first, lattice.types)

  for t in types
    myseries = filter(x -> get_ptype(x) == last(t), site)

    @series begin
      label --> first(t)
      [tuple(s...) for s in myseries]
    end
  end
end

@recipe function f(lattice::Lattice{2,T,Hexagonal}) where {T}
  site = lattice.site
  types = lattice.types

  seriestype := :scatter
  aspect_ratio --> 1
  markerstrokewidth --> 0
  markershape --> :hexagon
  # color --> map(get_ptype, site)
  # label --> map(first, lattice.types)

  A = cos(pi / 3)
  B = sin(pi / 3)

  for t in types
    s = filter(x -> get_ptype(x) == last(t), site)

    @series begin
      label --> first(t)
      new_coords = Vector{NTuple{2,Float64}}(undef, length(s))

      for i in eachindex(s)
        p = s[i]

        x = p[1] + A * p[2]
        y = B * p[2]

        new_coords[i] = (x, y)
      end

      new_coords
    end
  end
end

##### building neighborhoods

function build_neighborhoods!(nbtype::VonNeumann, lattice::Lattice{1})
  sites = lattice.coord_order

  sort_by_x_1D!(sites)
  sweep_neighbors!(nbtype, lattice, sites)

  return nothing
end

function build_neighborhoods!(nbtype::VonNeumann, lattice::Lattice{2})
  sites = lattice.coord_order

  sort_by_x_2D!(sites)
  sweep_neighbors!(nbtype, lattice, sites)

  sort_by_y_2D!(sites)
  sweep_neighbors!(nbtype, lattice, sites)

  return nothing
end

function build_neighborhoods!(nbtype::VonNeumann, lattice::Lattice{3})
  sites = lattice.coord_order

  sort_by_x_3D!(sites)
  sweep_neighbors!(nbtype, lattice, sites)

  sort_by_y_3D!(sites)
  sweep_neighbors!(nbtype, lattice, sites)

  sort_by_z_3D!(sites)
  sweep_neighbors!(nbtype, lattice, sites)

  return nothing
end

function sweep_neighbors!(nbtype::VonNeumann, lattice, sites)
  for i = 2:length(sites)
    x = sites[i-1]
    y = sites[i]
    d = distance(nbtype, x, y)
    if d == 1 && !(label(y) in neighborhood(lattice, x))
      add_neighbor!(lattice, x, y)
      add_neighbor!(lattice, y, x)
    end
  end
  return sites
end

function build_neighborhoods!(nbtype::Hexagonal, lattice::Lattice{2})
  sites = lattice.coord_order

  sweep_neighbors!(nbtype, lattice, sites)
end

function sweep_neighbors!(nbtype::Hexagonal, lattice, sites)
  n = length(sites)

  for i in 1:n
    for j in 1:i-1 # j < i
      x = sites[i]
      y = sites[j]
      d = distance(nbtype, x, y)

      if d == 1 && !(label(y) in neighborhood(lattice, x))
        add_neighbor!(lattice, x, y)
        add_neighbor!(lattice, y, x)
      end
    end
  end
  return sites
end

### TODO: 3D


text:
"""
 Tst
"""
function tosiunits(s)
    if !ismissing(s)
        a = split(s, " ")
        r = sum(Quantity.(eval.(Meta.parse.(a)), u"inch")) |> u"cm"
        round(typeof(1u"cm"), r)
    else
        missing
    end
end


text:
struct VelocityParameters{T <: SSDFloat}
    mu0::T
    beta::T
    E0::T
    mun::T
end

struct CarrierParameters{T <: SSDFloat}
    axis100::VelocityParameters{T}
    axis111::VelocityParameters{T}
end

struct MassParameters{T <: SSDFloat}
    ml::T
    mt::T
end

# Temperature models
include("LinearModel.jl")
include("BoltzmannModel.jl")
include("PowerLawModel.jl")
include("VacuumModel.jl")

# Electron model parametrization from [3]
@fastmath function γj(j::Integer, phi110::T, γ0::SArray{Tuple{3,3},T,2,9})::SArray{Tuple{3,3},T,2,9} where {T <: SSDFloat}
    tmp::T = 2 / 3
    a::T = acos(sqrt(tmp))
    Rx::SArray{Tuple{3,3},T,2,9} = SMatrix{3,3,T}(1, 0, 0, 0, cos(a), sin(a), 0, -sin(a), cos(a))
    b::T = -phi110 + (j - 1) * T(π) / 2
    Rzj::SArray{Tuple{3,3},T,2,9} = SMatrix{3,3,T}(cos(b), sin(b), 0, -sin(b), cos(b), 0, 0, 0, 1)
    Rj = Rx * Rzj
    transpose(Rj) * γ0 * Rj
end

@fastmath function setup_gamma_matrices(phi110::T, masses::MassParameters{T})::SVector{4, SMatrix{3,3,T}} where {T <: SSDFloat}
    ml::T = masses.ml
    mt::T = masses.mt
    γ0 = SMatrix{3,3,T}(1. / mt, 0, 0, 0, 1. / ml, 0, 0, 0, 1. / mt)
    SVector{4, SArray{Tuple{3,3},T,2,9}}(γj(1, phi110, γ0), γj(2, phi110, γ0), γj(3, phi110, γ0), γj(4, phi110, γ0))
end

# Longitudinal drift velocity formula
@fastmath function Vl(Emag::T, params::VelocityParameters{T})::T where {T <: SSDFloat}
    params.mu0 * Emag / (1 + (Emag / params.E0)^params.beta)^(1 / params.beta) - params.mun * Emag
end

"""
    ADLChargeDriftModel{T <: SSDFloat} <: AbstractChargeDriftModel

# Fields
- `electrons::CarrierParameters{T}`
- `holes::CarrierParameters{T}`
- `masses::MassParameters{T}`
- `phi110::T
- `gammas::SVector{4, SMatrix{3,3,T}}`
- `temperaturemodel::AbstractTemperatureModel{T}`
"""
struct ADLChargeDriftModel{T <: SSDFloat} <: AbstractChargeDriftModel{T}
    electrons::CarrierParameters{T}
    holes::CarrierParameters{T}
    masses::MassParameters{T}
    phi110::T
    gammas::SVector{4, SMatrix{3,3,T}}
    temperaturemodel::AbstractTemperatureModel{T}
end

function ADLChargeDriftModel{T}(chargedriftmodel::ADLChargeDriftModel)::ADLChargeDriftModel{T} where {T <: SSDFloat}
    cdmf64 = chargedriftmodel
    e100 = VelocityParameters{T}(cdmf64.electrons.axis100.mu0, cdmf64.electrons.axis100.beta, cdmf64.electrons.axis100.E0, cdmf64.electrons.axis100.mun)
    e111 = VelocityParameters{T}(cdmf64.electrons.axis111.mu0, cdmf64.electrons.axis111.beta, cdmf64.electrons.axis111.E0, cdmf64.electrons.axis111.mun)
    h100 = VelocityParameters{T}(cdmf64.holes.axis100.mu0, cdmf64.holes.axis100.beta, cdmf64.holes.axis100.E0, cdmf64.holes.axis100.mun)
    h111 = VelocityParameters{T}(cdmf64.holes.axis111.mu0, cdmf64.holes.axis111.beta, cdmf64.holes.axis111.E0, cdmf64.holes.axis111.mun)
    electrons = CarrierParameters{T}(e100, e111)
    holes     = CarrierParameters{T}(h100, h111)
    ml::T     = cdmf64.masses.ml
    mt::T     = cdmf64.masses.mt
    masses    = MassParameters{T}(ml, mt)
    phi110::T = cdmf64.phi110
    gammas = SVector{4, SArray{Tuple{3,3},T,2,9}}( cdmf64.gammas )
    temperaturemodel::AbstractTemperatureModel{T} = cdmf64.temperaturemodel
    ADLChargeDriftModel{T}(electrons, holes, masses, phi110, gammas, temperaturemodel)
end

function ADLChargeDriftModel(configfilename::Union{Missing, AbstractString} = missing; T::Type=Float32,
                             temperature::Union{Missing, Real}= missing, phi110::Union{Missing, Real} = missing)::ADLChargeDriftModel{T}

    if ismissing(configfilename) configfilename = joinpath(@__DIR__, "drift_velocity_config.json") end
    if !ismissing(temperature) temperature = T(temperature) end  #if you give the temperature it will be used, otherwise read from config file

    config = JSON.parsefile(configfilename)

    #mu0 in m^2 / ( V * s )
    #beta dimensionless
    #E0 in V / m
    #mun in m^2 / ( V * s )

    e100mu0::T  = config["drift"]["velocity"]["parameters"]["e100"]["mu0"]
    e100beta::T = config["drift"]["velocity"]["parameters"]["e100"]["beta"]
    e100E0::T   = config["drift"]["velocity"]["parameters"]["e100"]["E0"]
    e100mun::T  = config["drift"]["velocity"]["parameters"]["e100"]["mun"]
    e111mu0::T  = config["drift"]["velocity"]["parameters"]["e111"]["mu0"]
    e111beta::T = config["drift"]["velocity"]["parameters"]["e111"]["beta"]
    e111E0::T   = config["drift"]["velocity"]["parameters"]["e111"]["E0"]
    e111mun::T  = config["drift"]["velocity"]["parameters"]["e111"]["mun"]
    h100mu0::T  = config["drift"]["velocity"]["parameters"]["h100"]["mu0"]
    h100beta::T = config["drift"]["velocity"]["parameters"]["h100"]["beta"]
    h100E0::T   = config["drift"]["velocity"]["parameters"]["h100"]["E0"]
    h111mu0::T  = config["drift"]["velocity"]["parameters"]["h111"]["mu0"]
    h111beta::T = config["drift"]["velocity"]["parameters"]["h111"]["beta"]
    h111E0::T   = config["drift"]["velocity"]["parameters"]["h111"]["E0"]

    e100 = VelocityParameters{T}(e100mu0, e100beta, e100E0, e100mun)
    e111 = VelocityParameters{T}(e111mu0, e111beta, e111E0, e111mun)
    h100 = VelocityParameters{T}(h100mu0, h100beta, h100E0, 0)
    h111 = VelocityParameters{T}(h111mu0, h111beta, h111E0, 0)
    electrons = CarrierParameters{T}(e100, e111)
    holes     = CarrierParameters{T}(h100, h111)

    if "masses" in keys(config["drift"])
        ml = T(config["drift"]["masses"]["ml"])
        mt = T(config["drift"]["masses"]["mt"])
        masses = MassParameters{T}(ml, mt)
    else
        masses = MassParameters{T}(T(1.64),T(0.0819))
    end

    ismissing(phi110) ? phi110 = T(config["phi110"]) : phi110 = T(phi110)  #if you give the angle of the 110 axis it will be used, otherwise read from config file

    gammas = setup_gamma_matrices(phi110, masses)

    if "temperature_dependence" in keys(config)
        if "model" in keys(config["temperature_dependence"])
            model::String = config["temperature_dependence"]["model"]
            if model == "Linear"
                temperaturemodel = LinearModel{T}(config, temperature = temperature)
            elseif model == "PowerLaw"
                temperaturemodel = PowerLawModel{T}(config, temperature = temperature)
            elseif model == "Boltzmann"
                temperaturemodel = BoltzmannModel{T}(config, temperature = temperature)
            else
                temperaturemodel = VacuumModel{T}(config)
                println("Config File does not suit any of the predefined temperature models. The drift velocity will not be rescaled.")
            end
        else
            temperaturemodel = VacuumModel{T}(config)
            println("No temperature model specified. The drift velocity will not be rescaled.")
        end
    else
        temperaturemodel = VacuumModel{T}(config)
        # println("No temperature dependence found in Config File. The drift velocity will not be rescaled.")
    end

    return ADLChargeDriftModel{T}(electrons, holes, masses, phi110, gammas, temperaturemodel)
end

function getVe(fv::SVector{3, T}, cdm::ADLChargeDriftModel, Emag_threshold::T = T(1e-5))::SVector{3, T} where {T <: SSDFloat}
    cdmT = ADLChargeDriftModel{T}(cdm)
    getVe(fv, cdmT, Emag_threshold)
end

@fastmath function getVe(fv::SVector{3, T}, cdm::ADLChargeDriftModel{T}, Emag_threshold::T = T(1e-5))::SVector{3, T} where {T <: SSDFloat}
    @inbounds begin
        Emag::T = norm(fv)
        Emag_inv::T = inv(Emag)

        if Emag < Emag_threshold
            return SVector{3,T}(0, 0, 0)
        end

        mli::T = 1. / cdm.masses.ml
        mti::T = 1. / cdm.masses.mt
        B::T = 1/3 * sqrt(8*mti + mli)

        Gamma0::T = sqrt(2/3*mti + 1/3*mli)
        Gamma1::T = (-4*sqrt(mli) -4/3 * B)/(sqrt(mli) - B)^2
        Gamma2::T = Gamma1*(sqrt(mli) + 3*B)/(-4)

        f::NTuple{4,T} = scale_to_given_temperature(cdm.temperaturemodel)
        f100e::T = f[1]
        f111e::T = f[2]

        V100e::T = Vl(Emag, cdm.electrons.axis100) * f100e
        V111e::T = Vl(Emag, cdm.electrons.axis111) * f111e

        AE::T = V100e / Gamma0
        RE::T = Gamma1 * V111e / AE + Gamma2

        e0 = SVector{3, T}(fv * Emag_inv)

        oneOverSqrtEgE::SVector{4, T} = [1 / sqrt( cdm.gammas[j] * e0 ⋅ e0 ) for j in eachindex(1:4)] # setup.gammas[j] * e0 ⋅ e0 -> causes allocations

        g0 = MVector{3, T}(0,0,0)
        for j in eachindex(1:4)
            NiOverNj::T = RE * (oneOverSqrtEgE[j] / sum(oneOverSqrtEgE) - T(0.25)) + T(0.25)
            g0 += cdm.gammas[j] * e0 * NiOverNj * oneOverSqrtEgE[j]
        end

        return g0 * -AE
    end
end

# Hole model parametrization from [1] equations (22)-(26)
@fastmath function k0func(vrel::T)::T where {T <: SSDFloat}
    p0::T = 9.2652
    p1::T = 26.3467
    p2::T = 29.6137
    p3::T = 12.3689
    return @fastmath p0 - p1 * vrel + p2 * vrel^2 - p3 * vrel^3
end

@fastmath function lambda(k0::T)::T where {T <: SSDFloat}
    p0::T = -0.01322
    p1::T = 0.41145
    p2::T = 0.23657
    p3::T = 0.04077
    return @fastmath p0 * k0 + p1 * k0^2 - p2 * k0^3 + p3 * k0^4
end

@fastmath function omega(k0::T)::T where {T <: SSDFloat}
    p0::T = 0.006550
    p1::T = 0.19946
    p2::T = 0.09859
    p3::T = 0.01559
    return @fastmath p0 * k0 - p1 * k0^2 + p2 * k0^3 - p3 * k0^4
end

function getVh(fv::SVector{3,T}, cdm::ADLChargeDriftModel, Emag_threshold::T = T(1e-5))::SVector{3,T} where {T <: SSDFloat}
    cdmT = ADLChargeDriftModel{T}(cdm)
    getVh(fv, cdmT, Emag_threshold)
end

@fastmath function getVh(fv::SVector{3,T}, cdm::ADLChargeDriftModel{T}, Emag_threshold::T = T(1e-5))::SVector{3,T} where {T <: SSDFloat}
    @inbounds begin
        Emag::T = norm(fv)
        Emag_inv::T = inv(Emag)

        if Emag < Emag_threshold
            return SVector{3,T}(0, 0, 0)
        end

        f::NTuple{4,T} = scale_to_given_temperature(cdm.temperaturemodel)
        f100h::T = f[3]
        f111h::T = f[4]

        V100h::T = Vl(Emag, cdm.holes.axis100) * f100h
        V111h::T = Vl(Emag, cdm.holes.axis111) * f111h

        b::T = -π / 4 - cdm.phi110
        Rz = SMatrix{3, 3, T}(cos(b), sin(b), 0, -sin(b), cos(b), 0, 0, 0, 1)
        a = Rz * fv

        theta0::T = acos(a[3] / Emag)
        phi0::T = atan(a[2], a[1])

        k0::T = k0func(V111h / V100h)

        vtmp = MVector{3, T}(0, 0, 0) ## from CITATION; The implementation here is correct, mistake in the CITATION
        vtmp[3] = V100h * ( 1 - lambda(k0) * (sin(theta0)^4 * sin(2 * phi0)^2 + sin(2 * theta0)^2) )
        vtmp[1] = V100h * omega(k0) * (2 * sin(theta0)^3 * cos(theta0) * sin(2 * phi0)^2 + sin(4 * theta0))
        vtmp[2] = V100h * omega(k0) * sin(theta0)^3 * sin(4 * phi0)

        Ry = SMatrix{3, 3, T}(cos(theta0),0,-sin(theta0), 0,1,0, sin(theta0),0,cos(theta0))
        b = phi0 + pi/4 + cdm.phi110
        Rz = SMatrix{3,3, T}(cos(b),sin(b),0, -sin(b),cos(b),0, 0,0,1)
        vtmp = Rz * (Ry * vtmp)

        return vtmp
    end
end

function print(io::IO, tm::VacuumModel{T}) where {T <: SSDFloat}
    print(io, "No temperature model defined")
end
println(io::IO, tm::VacuumModel) = print(io, tm)

function print(io::IO, tm::BoltzmannModel{T}) where {T <: SSDFloat}
    print(io, "BoltzmannModel{$T}")
end
function println(io::IO, tm::BoltzmannModel{T}) where {T <: SSDFloat}
    println("\n________BoltzmannModel________")
    println("Fit function: p1 + p2 exp(-p3/T)\n")
    println("---Temperature settings---")
    println("Crystal temperature:   \t $(tm.temperature)")
    println("Reference temperature: \t $(tm.reftemperature)\n")

    println("---Fitting parameters---")
    println("   \te100      \te111      \th100      \th111")
    println("p1 \t$(tm.p1e100)   \t$(tm.p1e111)   \t$(tm.p1h100)   \t$(tm.p1h111)")
    println("p2 \t$(tm.p2e100)   \t$(tm.p2e111)   \t$(tm.p2h100)   \t$(tm.p2h111)")
    println("p3 \t$(tm.p3e100)   \t$(tm.p3e111)   \t$(tm.p3h100)   \t$(tm.p3h111)")
end

function print(io::IO, tm::LinearModel{T}) where {T <: SSDFloat}
    print(io, "LinearModel{$T}")
end
function println(io::IO, tm::LinearModel{T}) where {T <: SSDFloat}
    println("\n________LinearModel________")
    println("Fit function: p1 + p2 * T\n")
    println("---Temperature settings---")
    println("Crystal temperature:  \t$(tm.temperature)")
    println("Reference temperature:\t$(tm.reftemperature)\n")

    println("---Fitting parameters---")
    println("   \te100      \te111      \th100      \th111")
    println("p1 \t$(tm.p1e100)   \t$(tm.p1e111)   \t$(tm.p1h100)   \t$(tm.p1h111)")
    println("p2 \t$(tm.p2e100)   \t$(tm.p2e111)   \t$(tm.p2h100)   \t$(tm.p2h111)")
end

function print(io::IO, tm::PowerLawModel{T}) where {T <: SSDFloat}
    print(io, "PowerLawModel{$T}")
end
function println(io::IO, tm::PowerLawModel{T}) where {T <: SSDFloat}
    println("\n________PowerLawModel________")
    println("Fit function: p1 * T^(3/2)\n")
    println("---Temperature settings---")
    println("Crystal temperature:   \t $(tm.temperature)")
    println("Reference temperature: \t $(tm.reftemperature)\n")

    println("---Fitting parameters---")
    println("   \te100      \te111      \th100      \th111")
    println("p1 \t$(tm.p1e100)   \t$(tm.p1e111)   \t$(tm.p1h100)   \t$(tm.p1h111)")
end

show(io::IO, tm::AbstractTemperatureModel) = print(io, tm)
show(io::IO,::MIME"text/plain", tm::AbstractTemperatureModel) = show(io, tm)


text:
function SeisToSegy(in,out;su=true)

    if (su==true)
	file_hsize = 0
    else
	file_hsize = 900
	# add commands here to read text and binary headers out.thead and
        # out.bhead and write them to out
    end

    filename_headers = ParseHeaderName(in)
    extent = ReadTextHeader(in)
    n1 = extent.n1
    nx = extent.n2*extent.n3*extent.n4*extent.n5

    stream = open(out,"w")
    total = 60 + n1
    h_segy = Array(SegyHeader,1)
    h_seis = Array(Header,1)
    h1 = Header[]
    push!(h1,Seismic.InitSeisHeader())
    h1[1].o1 = extent.o1
    h1[1].d1 = extent.d1
    h1[1].n1 = extent.n1
    for j = 1 : nx
	if filename_headers != "NULL"
	    d,h1,e = SeisRead(in,itrace=j,ntrace=1,group="some")
	else
	    println("j=",j)
	    d,e = SeisRead(in,itrace=j,ntrace=1,group="some")
	    println("size(d)=",size(d))
	    h1[1].tracenum = j
	end
	h_segy[1] = MapHeaders(h1,j,"SeisToSegy")
	position = file_hsize + total*(j-1)*4 + segy_count["trace"]
	seek(stream,position)
	write(stream,convert(Array{Float32,1},d[:]))
	PutSegyHeader(stream,h_segy[1],n1,file_hsize,j)
    end
    close(stream)

end

max_stars_repo_name: fercarozzi/myseismicjulia

max_issues_repo_name: fercarozzi/myseismicjulia

max_forks_repo_name: fercarozzi/myseismicjulia


text:
using ApproxFun, Plots
import ApproxFun: setdomain

#####
# This solves
# u'(x) + u(x-1) = 0
# u(0) = 2.3
# on 0..3 by breaking the problem into three pieces
#####

S = Chebyshev(0..1)
D = Derivative(S)
L = [D            0I      0I;
     I             D      0I;
     0I           I        D]

B = [ldirichlet(S)   0                 0;
     rdirichlet(S)  -ldirichlet(S)     0;
     0              rdirichlet(S)      -ldirichlet(S)]

v = [B; L] \ [2.3; zeros(2); zeros(3)]

u = v[1] + setdomain(v[2], Domain(1..2)) + setdomain(v[3], Domain(2..3))

plot(u)

u'(1.3) + u(1.3-1)  # ≈ 0

max_stars_repo_path: Extras/Delay Equation.jl

max_issues_repo_path: Extras/Delay Equation.jl

max_forks_repo_path: Extras/Delay Equation.jl


text:
#import Redpitaya: measureTransferFunction

#=
Measure the transfer function of all receive chains

We use the last frequency of the send channels
=#
function measureTransferFunction(daq::DAQRedPitayaScpiNew)

  modulus = 4800 # TODO adapt me
  numFreq = div(modulus,daq.params.decimation*2)-1

  tf = zeros(ComplexF64, numFreq, numTxChannels(daq), 3)
  uall = zeros(ComplexF64, numFreq, numFreq+2, numTxChannels(daq), 2)
  ualltime = zeros(Float64, numFreq, (numFreq+1)*2, numTxChannels(daq), 2)
  freqs = [ daq.params.dfBaseFrequency/modulus*k for k=1:numFreq ]

  startTx(daq)

  for (k,freq) in enumerate(freqs)

    @info "Frequency [Hz]" freq

    # start sending
    for d=1:numTxChannels(daq)
      setTxParamsAll(daq,d,Float32[0.0,0.0,0.0,0.5],
                         Float32[0.0,0.0,0.0,0.0],
                         UInt32[1,1,1,k])
    end

    sleep(0.1)
    # receive data
    u, uref = readData(daq,1,currentFrame(daq))

    numPeriods = k
    for d=1:numTxChannels(daq)
      ualltime[k,:,d,1] = u
      ualltime[k,:,d,2] = uref
      uall[k,:,d,1] = rfft(u,1)
      uall[k,:,d,2] = rfft(uref,1)
      tf[k,d,1] = uall[k,numPeriods+1,d,1]
      tf[k,d,2] = uall[k,numPeriods+1,d,2]
      tf[k,d,3] = tf[k,d,1] / tf[k,d,2]
    end
  end

  stopTx(daq)
  #disconnect(daq)

  return freqs, tf, uall, ualltime
end


text:
module KPMPreconditioners

using LinearAlgebra
using Random
using FFTW
using Logging
using Printf

import LinearAlgebra: ldiv!, mul!, transpose

using ..Models: HolsteinModel, SSHModel, AbstractModel, Continuous, update_model!
using ..Checkerboard: checkerboard_mul!, checkerboard_transpose_mul!, checkerboard_inverse_mul!
using ..TimeFreqFFTs: TimeFreqFFT, τ_to_ω!, ω_to_τ!
using ..Utilities: get_index, reshaped

export LeftRightKPMPreconditioner, LeftKPMPreconditioner, RightKPMPreconditioner, SymmetricKPMPreconditioner, setup!, construct_Bbar

"""
Object to represent Kenerl Polynomial Expansion.
"""
mutable struct KPMExpansion{T1<:AbstractFloat,T2<:Continuous,T3<:AbstractModel}

    "If true apply preconditioner, else default to identity operator."
    active::Bool

    "Current frequency."
    ω::Int

    "Dimension of Krylov subspace used to approximate eigenvalues"
    n::Int

    "Used by Arnoldi method, columns form orthonormal basis of Krylov subspace"
    Q::Matrix{T1}

    "Used by Arnoldi method, A on Q basis, and t is upper Hessenberg"
    h::Matrix{T1}

    "amount by which to buffer min/max eigenvalues to get λ_lo/λ_hi"
    buf::T1

    "Min egeinvalues of A"
    λ_lo::T1

    "Max egeinvalues of A"
    λ_hi::T1

    "λ_avg = (λ_hi+λ_lo)/2"
    λ_avg::T1

    "λ_mag = (λ_hi-λ_lo)/2"
    λ_mag::T1

    "order = c1/phi + c2"
    c1::T1

    "order = c1/phi + c2"
    c2::T1

    "model."
    model::T3

    "TimeFreqFFT object for mapping between τ ⟷ ω"
    timefreqfft::TimeFreqFFT{T1}

    "exp{-Δτ⋅V̄} = (1/L)∑exp{-Δτ⋅V(τ)}"
    expnΔτV̄::Vector{T1}

    "For checkerboard representation of exp{-Δτ⋅K̄}=(1/L)∑exp{-Δτ⋅K(τ)}"
    cosht̄::Vector{T2}

    "For checkerboard representation of exp{-Δτ⋅K̄}=(1/L)∑exp{-Δτ⋅K(τ)}"
    sinht̄::Vector{T2}

    "ϕ=2π/L⋅(ω+1/2)"
    ϕs::Vector{T1}

    "Chebyshev Coefficients."
    coeff::Vector{Vector{Complex{T1}}}

    "Polynomial order of Chebyshev expansion for each ω value."
    order::Vector{Int}

    "Temporary Vector"
    v1::Vector{Complex{T1}}

    "Temporary Vector"
    v2::Vector{Complex{T1}}

    "Temporary Vector"
    v3::Vector{Complex{T1}}

    "Temporary Vector"
    v4::Vector{Complex{T1}}

    "Temporary Vector"
    v5::Vector{Complex{T1}}

    "Count total checkerboard multiplies."
    checkerboard_count::Int

    function KPMExpansion(model::AbstractModel{T1,T2}, n::Int, buf::T1, c1::T1, c2::T1) where {T1,T2}

        N   = model.Nsites
        L   = model.Lτ
        NL  = model.Ndim
        Lo2 = cld(L,2)

        timefreqfft = TimeFreqFFT(model.lattice,L)

        λ_lo      = 0.0
        λ_hi      = 2.0
        λ_avg     = (λ_hi+λ_lo)/2
        λ_mag     = (λ_hi-λ_lo)/2
        expnΔτV̄   = zeros(T1,N)
        cosht̄     = zeros(T2,model.Nbonds)
        sinht̄     = zeros(T2,model.Nbonds)
        ϕs        = [2*π/L*(ω+1/2) for ω in 0:Lo2-1]
        order     = ones(Int,Lo2)
        v1        = zeros(Complex{T1},NL)
        v2        = zeros(Complex{T1},NL)
        v3        = zeros(Complex{T1},N)
        v4        = zeros(Complex{T1},N)
        v5        = zeros(Complex{T1},N)

        if typeof(model) <: HolsteinModel
            cosht̄ .= model.cosht
            sinht̄ .= model.sinht
        elseif typeof(model) <: SSHModel
            expnΔτV̄ .= model.expΔτμ
        else
            throw(TypeError())
        end

        # size of krylov subspace
        n = min(n,N)

        # matrices for arnolid method
        Q = zeros(T1,N,n+1)
        h = zeros(T1,n+1,n)

        # construct expansion of the function f(x)=1.0-exp{i⋅ϕ⋅x} for each ϕ value
        coeff = [zeros(Complex{T1},1) for ω in 1:Lo2]

        return new{T1,T2,typeof(model)}(true,1,n,Q,h,buf,λ_lo,λ_hi,λ_avg,λ_mag,c1,c2,model,timefreqfft,expnΔτV̄,cosht̄,sinht̄,ϕs,coeff,order,v1,v2,v3,v4,v5,0)
    end
end

"""
Abstract type to reprepresent preconditioners based the Kernel Polynomial Method
that uses Chebyshev Polynomials to approximate M⁻¹[ω,ω].
"""
abstract type KPMPreconditioner{T1<:AbstractFloat,T2<:Continuous,T3<:AbstractModel} end

"""
For preconditioning M⋅x=b
"""
mutable struct LeftKPMPreconditioner{T1,T2,T3} <: KPMPreconditioner{T1,T2,T3}

    expansion::KPMExpansion{T1,T2,T3}

    function LeftKPMPreconditioner(model::AbstractModel{T1,T2}, n::Int, buf::T1, c1::T1, c2::T1) where {T1,T2}
        expansion = KPMExpansion(model,n,buf,c1,c2)
        return new{T1,T2,typeof(model)}(expansion)
    end

    function LeftKPMPreconditioner(expansion::KPMExpansion{T1,T2,T3}) where {T1,T2,T3}

        return new{T1,T2,T3}(expansion)
    end
end

"""
For preconditioning Mᵀ⋅x=b
"""
mutable struct RightKPMPreconditioner{T1,T2,T3} <: KPMPreconditioner{T1,T2,T3}

    expansion::KPMExpansion{T1,T2,T3}

    function RightKPMPreconditioner(model::AbstractModel{T1,T2}, n::Int, buf::T1, c1::T1, c2::T1) where {T1,T2}

        expansion = KPMExpansion(model,n,buf,c1,c2)
        T3        = typeof(model)
        return new{T1,T2,T3}(expansion)
    end

    function RightKPMPreconditioner(expansion::KPMExpansion{T1,T2,T3}) where {T1,T2,T3}

        return new{T1,T2,T3}(expansion)
    end
end

"""
For preconditioning M⋅x=b or Mᵀ⋅x=b
"""
mutable struct LeftRightKPMPreconditioner{T1,T2,T3} <: KPMPreconditioner{T1,T2,T3}

    lkpm::LeftKPMPreconditioner{T1,T2,T3}
    rkpm::RightKPMPreconditioner{T1,T2,T3}
    expansion::KPMExpansion{T1,T2,T3}

    function LeftRightKPMPreconditioner(model::AbstractModel{T1,T2}, n::Int, buf::T1, c1::T1, c2::T1) where {T1<:AbstractFloat,T2<:Number}

        lkpm = LeftKPMPreconditioner(model,n,buf,c1,c2)
        rkpm = transpose(lkpm)
        expansion = lkpm.expansion

        return new{T1,T2,typeof(model)}(lkpm,rkpm,expansion)
    end
end

"""
For preconditioning MᵀM⋅x=b
"""
mutable struct SymmetricKPMPreconditioner{T1,T2,T3} <: KPMPreconditioner{T1,T2,T3}

    expansion::KPMExpansion{T1,T2,T3}
    transposed::Bool

    function SymmetricKPMPreconditioner(model::AbstractModel{T1,T2}, n::Int, buf::T1, c1::T1, c2::T1) where {T1,T2}

        expansion = KPMExpansion(model,n,buf,c1,c2)
        T3        = typeof(model)
        return new{T1,T2,T3}(expansion,false)
    end

    function SymmetricKPMPreconditioner(expansion::KPMExpansion{T1,T2,T3}) where {T1,T2,T3}

        return new{T1,T2,T3}(expansion,false)
    end
end

"""
Return RightKPMPreconditioner given LeftKPMPreconditioner.
"""
function transpose(P::LeftKPMPreconditioner)

    return RightKPMPreconditioner(P.expansion)
end

"""
Return LeftKPMPreconditioner given RightKPMPreconditioner.
"""
function transpose(P::RightKPMPreconditioner)

    return LeftKPMPreconditioner(P.expansion)
end

"""
Set up the KPMPreconditioner i.e. calculate exp{-Δτ⋅V̄}.
"""
function setup!(op::KPMPreconditioner)

    setup!(op.expansion)
    return nothing
end

"""
Update the preconditioner.
"""
function setup!(op::KPMExpansion{T1,T2,T3}) where {T1,T2,T3}

    # update exp{-Δτ⋅V̄} and exp{-Δτ⋅K̄}
    update_A!(op)

    # approximate min/max eigenvalue of A = exp{-Δτ⋅V̄}⋅exp{-Δτ⋅K̄}
    e_min, e_max = arnoldi_eigenvalue_bounds!(op, op.Q, op.h, op.v3, op.v4, op.model.rng)

    # compute λ_lo and λ_hi
    λ_lo = max(0.0 , (1-2*op.buf)*e_min)
    λ_hi = (1+2*op.buf)*e_max

    # if λ_lo or λ_hi has changed by a factor of more than op.buf,
    # recompute expansion coefficients
    if !isapprox(λ_lo, op.λ_lo, rtol=op.buf) || !isapprox(λ_hi, op.λ_hi, rtol=op.buf)

        op.λ_lo  = λ_lo
        op.λ_hi  = λ_hi
        op.λ_avg = (op.λ_hi+op.λ_lo)/2
        op.λ_mag = (op.λ_hi-op.λ_lo)/2

        # update expansions
        for ω in 1:length(op.ϕs)
            # calculate order of expansion
            coeff       = op.coeff[ω]
            ϕ           = op.ϕs[ω]
            # order       = round(Int, op.c1/ϕ + op.c2)
            order       = round(Int, (op.λ_hi-op.λ_lo)*(op.c1/ϕ + op.c2))
            order       = max(1,order)
            op.order[ω] = order
            # resize vector containing expansion coefficients
            resize!(coeff,order)
            # calculate expansion coefficients
            kpm_coefficients!(coeff, order, op.λ_lo, op.λ_hi, ϕ)
        end
    end

    # set preconditioner to being active
    op.active = true

    return nothing
end

function setup!(op)

    return nothing
end

"""
Calculate exp{-Δτ⋅V̄} and exp{-Δτ⋅K̄}
"""
function update_A!(op::KPMExpansion{T1,T2,T3}) where {T1,T2,T3<:HolsteinModel}

    N  = op.model.Nsites::Int
    L  = op.model.Lτ::Int
    Δτ = op.model.Δτ

    # calulcate diagonal matrix exp{-Δτ⋅V̄}
    expnΔτV = op.model.expnΔτV::Vector{T2}
    @fastmath @inbounds for i in 1:N
        op.expnΔτV̄[i] = 0.0
        for τ in 1:L
            op.expnΔτV̄[i] += expnΔτV[get_index(τ,i,L)]
        end
        op.expnΔτV̄[i] /= L
    end

    return nothing
end

"""
Update the matrix A=exp{-Δτ⋅V̄}⋅exp{-Δτ⋅K̄}
"""
function update_A!(op::KPMExpansion{T1,T2,T3}) where {T1,T2,T3<:SSHModel}

    N  = op.model.Nbonds::Int
    L  = op.model.Lτ::Int
    Δτ = op.model.Δτ

    # calulcate checkerboard representation of matrix exp{-Δτ⋅K̄}
    cosht       = op.model.cosht::Matrix{T2}
    sinht       = op.model.sinht::Matrix{T2}
    t′          = op.model.t′::Matrix{T2}
    chkbrd_perm = op.model.checkerboard_perm::Vector{Int}
    @fastmath @inbounds for i in 1:N
        op.cosht̄[i] = 0.0
        op.sinht̄[i] = 0.0
        for τ in 1:L
            op.cosht̄[i] += cosht[τ,i]
            op.sinht̄[i] += sinht[τ,i]
        end
        op.cosht̄[i] /= L
        op.sinht̄[i] /= L
    end

    # calulcate diagonal matrix exp{-Δτ⋅V̄}
    copyto!(op.expnΔτV̄,op.model.expΔτμ)

    return nothing
end

"""
Perform A⋅v where A=exp{-Δτ⋅K̄}⋅exp{-Δτ⋅V̄}
"""
function mul!(v′::AbstractVector{T4},op::KPMExpansion{T1,T2,T3},v::AbstractVector{T4}) where {T1,T2,T3,T4<:Continuous}

    expnΔτV̄ = op.expnΔτV̄::Vector{T1}
    neighbor_table = op.model.neighbor_table::Matrix{Int}
    cosht̄ = op.cosht̄::Vector{T2}
    sinht̄ = op.sinht̄::Vector{T2}

    @. v′ = expnΔτV̄ * v
    checkerboard_mul!(v′,neighbor_table,cosht̄,sinht̄)

    op.checkerboard_count += 1

    return nothing
end

"""
Perform A⁻¹⋅v where A = exp{-Δτ⋅K̄}⋅exp{-Δτ⋅V̄}
"""
function ldiv!(v′::AbstractVector{T4},op::KPMExpansion{T1,T2,T3},v::AbstractVector{T4}) where {T1,T2,T3,T4<:Continuous}

    expnΔτV̄ = op.expnΔτV̄::Vector{T1}
    neighbor_table = op.model.neighbor_table::Matrix{Int}
    cosht̄ = op.cosht̄::Vector{T2}
    sinht̄ = op.sinht̄::Vector{T2}

    copyto!(v′,v)
    checkerboard_inverse_mul!(v′,neighbor_table,cosht̄,sinht̄)
    @. v′ /= expnΔτV̄

    op.checkerboard_count += 1

    return nothing
end

"""
Apply Preconditioner.
"""
function ldiv!(vout::AbstractVector{T},P::KPMPreconditioner,vin::AbstractVector{T}) where {T<:AbstractFloat}

    op = P.expansion::KPMExpansion
    N  = op.model.Nsites::Int
    L  = op.model.Lτ::Int
    v1 = op.v1::Vector{Complex{T}}
    v2 = op.v2::Vector{Complex{T}}
    op.checkerboard_count = 0

    if op.active # apply preconditioner if active

        # 1. apply phase factor to go from (anit-periodic)⟶(periodic) in τ
        # 2. FFT from τ ⟶ ω
        τ_to_ω!(v2,op.timefreqfft,vin)

        a1  = reshaped(v1,(L,N))
        a1T = reshaped(v1,(N,L))
        a2  = reshaped(v2,(L,N))
        a2T = reshaped(v2,(N,L))

        transpose!(a1T,a2)

        # iterating over half the range of frequencies
        @fastmath @inbounds for ω in 1:cld(L,2)

            # input vector
            u1 = @view a1T[:,ω]

            # output vector
            u2 = @view a2T[:,ω]

            # set frequency
            op.ω = ω

            # multiply by KPM approximation to M⁻¹[ω,ω]
            mul!(u2,P,u1)

            # accounting for symmetry
            for i in 1:N
                a2T[i,L-ω+1] = conj(a2T[i,ω])
            end
        end

        transpose!(a1,a2T)

        # 1. iFFT from ω ⟶ τ
        # 2. apply inverse phase factor to go from (periodic)⟶(anti-periodic) in τ
        ω_to_τ!(vout,op.timefreqfft,v1)

    else # if preconditioner inactive then behave like identity matrix

        copyto!(vout,vin)
    end

    return nothing
end

function ldiv!(op::KPMPreconditioner,v::AbstractVector)

    ldiv!(v,op,v)
    return nothing
end

"""
Multiply by KPM approximation for M⁻¹[ω,ω] or M⁻ᵀ[ω,ω]
"""
function mul!(v′::AbstractVector{Complex{T1}},P::LeftRightKPMPreconditioner{T1,T2,T3},v::AbstractVector{Complex{T1}}) where {T1,T2,T3}

    op    = P.expansion::KPMExpansion{T1,T2,T3}
    model = op.model::T3
    lkpm  = P.lkpm::LeftKPMPreconditioner{T1,T2,T3}
    rkpm  = P.rkpm::RightKPMPreconditioner{T1,T2,T3}

    if model.transposed
        mul!(v′,rkpm,v)
    else
        mul!(v′,lkpm,v)
    end

    return nothing
end

"""
Multiply by KPM approximation for M⁻¹[ω,ω]
"""
function mul!(v′::AbstractVector{Complex{T1}},P::LeftKPMPreconditioner{T1,T2,T3},v::AbstractVector{Complex{T1}}) where {T1,T2,T3}

    op    = P.expansion::KPMExpansion{T1,T2,T3}
    ω     = op.ω # current frequency
    order = op.order[ω] # order of expansion
    coeff = op.coeff::Vector{Vector{Complex{T1}}}
    c     = coeff[ω]::Vector{Complex{T1}}

    # Recursively build `uₙ = Tₙ(A)⋅v`.
    uₙ₋₁ = op.v3
    uₙ   = op.v4
    uₙ₊₁ = op.v5

    # v′ = c₁⋅u₁ = c₁⋅T₁(A)⋅v
    @. v′ = c[1] * v
    if order>1
        n = 1 # current order
        copyto!(uₙ,v)     # u₁ = v
        mulA′!(uₙ₊₁,P,uₙ) # u₂ = A′⋅u₁
        @fastmath @inbounds while true
            n += 1 # increment order counter
            # uₙ₋₁ = uₙ
            # uₙ   = uₙ₊₁
            temp = uₙ₋₁
            uₙ₋₁ = uₙ
            uₙ   = uₙ₊₁
            uₙ₊₁ = temp
            # v′ = v′ + cₙ⋅uₙ = v′ + cₙ⋅Tₙ(A)⋅v
            @. v′ += c[n] * uₙ
            if n==order
                break
            end
            # uₙ₊₁ = A′⋅uₙ
            mulA′!(uₙ₊₁,P,uₙ)
            # uₙ₊₁ = 2⋅A′⋅uₙ - uₙ₋₁
            @. uₙ₊₁ = 2 * uₙ₊₁ - uₙ₋₁
        end
    end

    return nothing
end

"""
Multiply by KPM approximation for M⁻ᵀ[ω,ω]
"""
function mul!(v′::AbstractVector{Complex{T1}},P::RightKPMPreconditioner{T1,T2},v::AbstractVector{Complex{T1}}) where {T1<:AbstractFloat,T2<:Number}

    op    = P.expansion::KPMExpansion{T1,T2}
    ω     = op.ω # current frequency
    order = op.order[ω] # order of expansion
    coeff = op.coeff::Vector{Vector{Complex{T1}}}
    c     = coeff[ω]::Vector{Complex{T1}}

    # Recursively build `uₙ = Tₙ(A)⋅v`.
    uₙ₋₁ = op.v3
    uₙ   = op.v4
    uₙ₊₁ = op.v5

    # v′ = c₁⋅u₁ = c₁⋅T₁(A)⋅v
    @. v′ = conj(c[1]) * v
    if order>1
        n = 1 # current order
        copyto!(uₙ,v)     # u₁ = v
        mulA′!(uₙ₊₁,P,uₙ) # u₂ = A′⋅u₁
        @fastmath @inbounds while true
            n += 1 # increment order counter
            # uₙ₋₁ = uₙ
            # uₙ   = uₙ₊₁
            temp = uₙ₋₁
            uₙ₋₁ = uₙ
            uₙ   = uₙ₊₁
            uₙ₊₁ = temp
            # v′ = v′ + cₙ⋅uₙ = v′ + cₙ⋅Tₙ(A)⋅v
            @. v′ += conj(c[n]) * uₙ
            if n==order
                break
            end
            # uₙ₊₁ = A′⋅uₙ
            mulA′!(uₙ₊₁,P,uₙ)
            # uₙ₊₁ = 2⋅A′⋅uₙ - uₙ₋₁
            @. uₙ₊₁ = 2 * uₙ₊₁ - uₙ₋₁
        end
    end

    return nothing
end

"""
Multiply by KPM approximation for M⁻¹[ω,ω]⋅M⁻ᵀ[ω,ω]
"""
function mul!(v′::AbstractVector{Complex{T1}},P::SymmetricKPMPreconditioner{T1,T2},v::AbstractVector{Complex{T1}}) where {T1<:AbstractFloat,T2<:Number}

    op    = P.expansion::KPMExpansion{T1,T2}
    ω     = op.ω # current frequency
    order = op.order[ω] # order of expansion
    coeff = op.coeff::Vector{Vector{Complex{T1}}}
    c     = coeff[ω]::Vector{Complex{T1}}

    # Recursively build `uₙ = Tₙ(A)⋅v`.
    uₙ₋₁ = op.v3
    uₙ   = op.v4
    uₙ₊₁ = op.v5

    # multiply by M⁻ᵀ[ω,ω]
    P.transposed = true

    # v′ = c₁⋅u₁ = c₁⋅T₁(A)⋅v
    @. v′ = conj(c[1]) * v
    if order>1
        n = 1 # current order
        copyto!(uₙ,v)     # u₁ = v
        mulA′!(uₙ₊₁,P,uₙ) # u₂ = A′⋅u₁
        @fastmath @inbounds while true
            n += 1 # increment order counter
            # uₙ₋₁ = uₙ
            # uₙ   = uₙ₊₁
            temp = uₙ₋₁
            uₙ₋₁ = uₙ
            uₙ   = uₙ₊₁
            uₙ₊₁ = temp
            # v′ = v′ + cₙ⋅uₙ = v′ + cₙ⋅Tₙ(A)⋅v
            @. v′ += conj(c[n]) * uₙ
            if n==order
                break
            end
            # uₙ₊₁ = A′⋅uₙ
            mulA′!(uₙ₊₁,P,uₙ)
            # uₙ₊₁ = 2⋅A′⋅uₙ - uₙ₋₁
            @. uₙ₊₁ = 2 * uₙ₊₁ - uₙ₋₁
        end
    end

    # Multiply by M⁻¹[ω,ω]
    P.transposed = false

    if order==1
        @. v′ *= c[1] # v′ = c₁⋅u₁ = c₁⋅T₁(A)⋅v
    else
        copyto!(uₙ,v′)    # u₁ = v
        @. v′ = c[1] * v′ # v′ = c₁⋅u₁ = c₁⋅T₁(A)⋅v
        n = 1 # current order
        mulA′!(uₙ₊₁,P,uₙ) # u₂ = A′⋅u₁
        @fastmath @inbounds while true
            n += 1 # increment order counter
            # uₙ₋₁ = uₙ
            # uₙ   = uₙ₊₁
            temp = uₙ₋₁
            uₙ₋₁ = uₙ
            uₙ   = uₙ₊₁
            uₙ₊₁ = temp
            # v′ = v′ + cₙ⋅uₙ = v′ + cₙ⋅Tₙ(A)⋅v
            @. v′ += c[n] * uₙ
            if n==order
                break
            end
            # uₙ₊₁ = A′⋅uₙ
            mulA′!(uₙ₊₁,P,uₙ)
            # uₙ₊₁ = 2⋅A′⋅uₙ - uₙ₋₁
            @. uₙ₊₁ = 2 * uₙ₊₁ - uₙ₋₁
        end
    end

    return nothing
end

"""
Perform A′⋅v where A′= 2⋅(A-λ_lo)/(λ_hi-λ_lo)-I
"""
function mulA′!(v′::AbstractVector{Complex{T}},P::KPMPreconditioner,v::AbstractVector{Complex{T}}) where {T<:AbstractFloat}

    λ_avg = P.expansion.λ_avg::T
    λ_mag = P.expansion.λ_mag::T
    mulA!(v′,P,v)
    @. v′ = (1/λ_mag)*v′ - (λ_avg/λ_mag)*v

    return nothing
end

"""
Perform A⋅v where A=exp{-Δτ⋅V̄}⋅exp{-Δτ⋅K̄} or A=exp{-Δτ⋅K}⋅exp{-Δτ⋅V̄}
"""
function mulA!(v′::AbstractVector{Complex{T1}},P::LeftRightKPMPreconditioner{T1,T2},v::AbstractVector{Complex{T1}}) where {T1<:AbstractFloat,T2<:Number}

    op    = P.expansion::KPMExpansion{T1,T2}
    model = op.model::T3
    lkpm  = P.lkpm::LeftKPMPreconditioner{T1,T2}
    rkpm  = P.rkpm::RightKPMPreconditioner{T1,T2}

    if model.transposed
        mul!(v′,rkpm,v)
    else
        mul!(v′,lkpm,v)
    end

    return nothing
end

"""
Perform A⋅v where A=exp{-Δτ⋅V̄}⋅exp{-Δτ⋅K̄}
"""
function mulA!(v′::AbstractVector{Complex{T1}},P::LeftKPMPreconditioner{T1,T2,T3},v::AbstractVector{Complex{T1}}) where {T1,T2,T3}

    op      = P.expansion::KPMExpansion{T1,T2,T3}
    expnΔτV̄ = op.expnΔτV̄::Vector{T1}
    neighbor_table = op.model.neighbor_table::Matrix{Int}
    cosht̄ = op.cosht̄::Vector{T2}
    sinht̄ = op.sinht̄::Vector{T2}

    @. v′ = expnΔτV̄ * v
    checkerboard_mul!(v′,neighbor_table,cosht̄,sinht̄)

    op.checkerboard_count += 1

    return nothing
end

"""
Perform A⋅v where A=exp{-Δτ⋅K̄}⋅exp{-Δτ⋅V̄}
"""
function mulA!(v′::AbstractVector{Complex{T1}},P::RightKPMPreconditioner{T1,T2,T3},v::AbstractVector{Complex{T1}}) where {T1,T2,T3}

    op      = P.expansion::KPMExpansion{T1,T2,T3}
    expnΔτV̄ = op.expnΔτV̄::Vector{T1}
    neighbor_table = op.model.neighbor_table::Matrix{Int}
    cosht̄ = op.cosht̄::Vector{T2}
    sinht̄ = op.sinht̄::Vector{T2}

    copyto!(v′,v)
    checkerboard_transpose_mul!(v′,neighbor_table,cosht̄,sinht̄)
    @. v′ *= expnΔτV̄

    op.checkerboard_count += 1

    return nothing
end

"""
Perform A⋅v where A=exp{-Δτ⋅V̄}⋅exp{-Δτ⋅K̄} or A=exp{-Δτ⋅K̄}⋅exp{-Δτ⋅V̄}
"""
function mulA!(v′::AbstractVector{Complex{T1}},P::SymmetricKPMPreconditioner{T1,T2,T3},v::AbstractVector{Complex{T1}}) where {T1,T2,T3}

    op      = P.expansion::KPMExpansion{T1,T2,T3}
    expnΔτV̄ = op.expnΔτV̄::Vector{T1}
    neighbor_table = op.model.neighbor_table::Matrix{Int}
    cosht̄ = op.cosht̄::Vector{T2}
    sinht̄ = op.sinht̄::Vector{T2}

    if P.transposed
        copyto!(v′,v)
        checkerboard_transpose_mul!(v′,neighbor_table,cosht̄,sinht̄)
        @. v′ *= expnΔτV̄
    else
        @. v′ = expnΔτV̄ * v
        checkerboard_mul!(v′,neighbor_table,cosht̄,sinht̄)
    end

    op.checkerboard_count += 1

    return nothing
end

"""
Calculate coefficients c_m of a polynomial approximation,
    M⁻¹[ω,ω] ~ sum_m c_m T_m(x)
valid for all x in the range
    λ_lo < x < λ_hi.
A Chebyshev approximation naturally lies in the range -1 < x < 1, so some rescaling
factors are necessary.
"""
function kpm_coefficients!(c::AbstractVector{T}, order::Int, λ_lo, λ_hi, ϕ) where {T<:Complex}

    M     = order
    N_M   = 2*M
    λ_avg = (λ_hi+λ_lo)/2
    λ_mag = (λ_hi-λ_lo)/2

    # initialize coefficients to zero
    fill!(c,0)

    # declare temporary array
    c′ = zeros(T,N_M)

    #####################################
    ## Real Part Exapnsion Coefficient ##
    #####################################

    for n in 0:(N_M-1)
        c′[n+1] = real(scalar_invM(λ_mag*cos(π*(n+0.5)/N_M)+λ_avg,ϕ))
    end
    FFTW.dct!(c′)

    # FFTW uses the "unitary" normalization. Undo that.
    c′    *= sqrt(2*N_M)/2
    c′[1] *= sqrt(2)

    for m in 0:(M-1)
        q_m = π / (m == 0 ? 1 : 2)
        c[m+1] += (π * c′[m+1]) / (N_M * q_m)
    end

    ##########################################
    ## Imaginary Part Exapnsion Coefficient ##
    ##########################################

    for n in 0:(N_M-1)
        c′[n+1] = imag(scalar_invM(λ_mag*cos(π*(n+0.5)/N_M)+λ_avg,ϕ))
    end
    FFTW.dct!(c′)

    # FFTW uses the "unitary" normalization. Undo that.
    c′    *= sqrt(2*N_M)/2
    c′[1] *= sqrt(2)

    for m in 0:(M-1)
        q_m = π / (m == 0 ? 1 : 2)
        c[m+1] += im * (π * c′[m+1]) / (N_M * q_m)
    end

    return nothing
end

"""
Computes a basis of the (n + 1)-Krylov subspace of A: the space spanned by {b, Ab, ..., A^n b}.
Then use this to approximate the min and max eigenvalues of A.
"""
function arnoldi_eigenvalue_bounds!(A, Q::AbstractMatrix{T1}, h::AbstractMatrix{T1}, b::AbstractVector{T2}, v::AbstractVector{T2}, rng::AbstractRNG) where {T1<:AbstractFloat,T2<:Continuous}

    # dimension of Krylov subspace, must be >= 1
    n = size(h,2)

    # dimension of A matrix
    m = size(Q,1)

    ################################
    ## Arnoldi for Max Eigenvalue ##
    ################################

    # randomize input vector using for loop because
    # b vector be of complex valued data type
    for i in 1:m
        b[i] = randn(rng,T1)
    end

    # normalize input vector
    normalize!(b)

    # Use it as the first Krylov vector
    @. Q[:,1] = real(b)

    l = n
    for k in 1:n
        mul!(v,A,b) # Generate a new candidate vector
        for j in 1:k # Subtract the projections on previous vectors
            Qj      = @view Q[:, j]
            h[j, k] = real(dot(Qj, v))
            @. v   -= h[j, k] * Qj
        end
        h[k+1, k] = norm(v)
        # Add the produced vector to the list, unless the zero vector is produced
        if h[k+1, k] > 1e-12
            @. b = v / h[k + 1, k]
            @. Q[:,k+1] = real(b)
        else  # If that happens, stop iterating.
            l = k
            break
        end
    end

    # calulcate min and max eigenvalues
    h′       = @view h[1:l,1:l]
    eigvs    = eigvals!(h′)
    e_max    = maximum(real, eigvs)

    ################################
    ## Arnoldi for Min Eigenvalue ##
    ################################

    # randomize input vector
    for i in 1:m
        b[i] = randn(rng,T1)
    end

    # normalize input vector
    normalize!(b)

    # Use it as the first Krylov vector
    Q1 = @view Q[:,1]
    @. Q1 = real(b)

    l = n
    for k in 1:n
        ldiv!(v,A,b) # Generate a new candidate vector
        for j in 1:k # Subtract the projections on previous vectors
            Qj      = @view Q[:, j]
            h[j, k] = real(dot(Qj, v))
            @. v   -= h[j, k] * Qj
        end
        h[k+1, k] = norm(v)
        # Add the produced vector to the list, unless the zero vector is produced
        if h[k+1, k] > 1e-12
            @. b = v / h[k + 1, k]
            @. Q[:,k+1] = real(b)
        else  # If that happens, stop iterating.
            l = k
            break
        end
    end

    # calulcate min and max eigenvalues
    h′     = @view h[1:l,1:l]
    eigvs  = eigvals!(h′)
    e_min = 1/maximum(real, eigvs)

    return e_min, e_max
end

"""
Scalar function of x where x has replaced A=exp{-Δτ⋅V̄}⋅exp{-Δτ⋅K} in
the expresion M⁻¹[ω,ω] = (I - exp{-i⋅ϕ(ω)}⋅A)⁻¹
"""
function scalar_invM(x,ϕ)

    return inv(1.0 - exp(-im * ϕ) * x)
end

"""
Function to construct Bbar (also called A) matrix.
"""
function construct_Bbar(op::KPMPreconditioner{T1,T2,T3};threshold::T1=1e-10) where {T1,T2,T3}

    expansion = op.expansion
    N         = expansion.model.Nsites

    # store matrix elements
    rows     = Int[]
    cols     = Int[]
    elements = T1[]

    # stores columns vector
    v  = zeros(T1,N)
    Av = zeros(T1,N)

    # iterating over rows
    for col in 1:N
        # initialize column vector as unit vector
        v      .= 0.0
        v[col]  = 1.0
        # doing matrix vector multiply
        mul!(Av,expansion,v)
        # iterate of column vecto
        for row in 1:N
            # if nonzero
            if abs(Av[row])>threshold
                # save matrix element
                append!(rows,row)
                append!(cols,col)
                append!(elements,Av[row])
            end
        end
    end
    return rows, cols, elements

    return nothing
end

end

max_stars_repo_name: cohensbw/Langevin

max_issues_repo_name: cohensbw/Langevin

max_forks_repo_name: cohensbw/Langevin


text:
# This file was generated, do not modify it. # hide
evaluate(
    pipe,
    X,
    height,
    resampling=Holdout(),
    measure=rms
) |> pprint

max_issues_repo_name: alan-turing-institute/MLJTutorials

max_forks_repo_name: alan-turing-institute/MLJTutorials


text:
function acqr_simula(m::acqr, nt::Int, seed::Int)
	#
	# simulated model to test the functions
	#
	# nt: number of time instants
	# seed: seed for the random numbers

	# random seed
	# ------------------------------
	Random.seed!(seed)

	# state equation, no noise
	# ---------------------------------
	w = sqrt(m.q)*randn(nt) # w ~ N(0,q)
	x = zeros(nt)
	x[1] = m.m1 + sqrt(m.P1)*randn() # x1 ~ N(m1,P1)
	for t=1:nt-1
		x[t+1] = m.a*x[t] + w[t]
	end

	# observation equation
	# -------------------------------------
	v = sqrt(m.r)*randn(nt) # v ~ N(0,r)
	y = m.c*x + v

	salida = Dict(
	"y"=>y,"x"=>x,
	"w"=>w,"v"=>v)

	return salida
end


text:
using Apriori
using Test
using Random

@testset "Apriori.jl" begin

    #testy uruchamia się z konsoli, wpisujesz ]test
    # ten pierwszy znak ] przejdzie do trybu obsługi paczek i wtedy na niebieski kolor zmieny prompt
    # jak zmieni kolor to dopisujesz test i enter
    # tak powinien prompt wyglądać przed enter
    # (Apriori) pkg> test

    let  #IZOLOWANY SCOPE, przydatny jak chcesz tych samych zmiennych używać i mieć pewność że nic nie wycieka poza test
        df = dummy_dataset_biased(20,100, 0.3, ['a'], ['b'], 1)
        result = apriori(df, 0.9, 0.9)
        rules = map(x->x[1], result)
        @test ([:b], [:a]) in rules
        @test ([:a], [:b]) in rules

        df = dummy_dataset_biased(20,100, 0.2, ['a', 'b'], ['c','d'], 1, 0.8)
        result = apriori(df, 0.9, 0.9)
        rules = map(x->x[1], result)
        @test !(([:a, :b], [:c, :d]) in rules)
        result = apriori(df, 0.7, 0.7)
        rules = map(x->x[1], result)
        @test ([:a, :b], [:c, :d]) in rules

        #CONSTRAINTS
        X = [:a]
        df = dummy_dataset_biased(20,100, 0.3, ['a', 'b'], ['c','d'], 1)
        result = apriori(df, 0.9, 0.9, X_in_antecedent=X)
        rules = map(x->x[1], result)
        for rule in rules
            for item in X
                @test item in rule[1]
            end
        end

        X = [:a]
        df = dummy_dataset_biased(20,100, 0.3, ['a', 'b'], ['c','d'], 1)
        result = apriori(df, 0.9, 0.9,
                        X_in_antecedent=Vector{Symbol}()::Vector{Symbol},
                        antecedent_in_X=X)
        rules = map(x->x[1], result)
        for rule in rules
            for item in rule[1]
                @test item in X
            end
        end

        X = [:a]
        df = dummy_dataset_biased(20,100, 0.3, ['a', 'b'], ['c','d'], 1)
        result = apriori(df, 0.9, 0.9,
                        X_in_antecedent=Vector{Symbol}()::Vector{Symbol},
                        antecedent_in_X=Vector{Symbol}()::Vector{Symbol},
                        X_in_consequent=X)
        rules = map(x->x[1], result)
        for rule in rules
            for item in X
                @test item in rule[2]
            end
        end

        df = dummy_dataset_biased(20,100, 0.3, ['a', 'b'], ['c','d'], 1)
        result = apriori(df, 0.9, 0.9,
                        X_in_antecedent=Vector{Symbol}()::Vector{Symbol},
                        antecedent_in_X=Vector{Symbol}()::Vector{Symbol},
                        X_in_consequent=Vector{Symbol}()::Vector{Symbol},
                        consequent_in_X=X)
        rules = map(x->x[1], result)
        for rule in rules
            for item in rule[2]
                @test item in X
            end
        end
    end

    # ale nie musisz w tych blokach let/end pisać, możesz zrobić jakiś tu globalny d i sprawdzać wieloma @test

end


text:
# Load X and y variable
using JLD
data = load("outliersData.jld")
(X,y,Xtest,ytest) = (data["X"],data["y"],data["Xtest"],data["ytest"])

# Fit a least squares model
include("leastMax.jl")
model = leastMax(X,y)

# Evaluate training error
yhat = model.predict(X)
trainError = mean((yhat - y).^2)
@printf("Squared train Error with least squares: %.3f\n",trainError)

# Evaluate test error
yhat = model.predict(Xtest)
testError = mean((yhat - ytest).^2)
@printf("Squared test Error with least squares: %.3f\n",testError)

# Plot model
using PyPlot
figure()
plot(X,y,"b.")
Xhat = minimum(X):.01:maximum(X)
yhat = model.predict(Xhat)
plot(Xhat,yhat,"g")


text:
function request_body(url; kwargs...)
    resp = nothing
    body = sprint() do output
        resp = request(url; output=output, kwargs...)
    end
    return resp, body
end

function status_error(resp, log=nothing)
    logs = !isnothing(log) ? ": $log" : ""
    error("request status $(resp.message)$logs")
end

# HfApi

"""
    login(user::AbstractString = readline())

Login to huggingface hub and get/save the account token for operations that need authentication.
 Re-run this function with username (i.e. `login(username)`) if token expired.
"""
login(user::AbstractString) = login(maybe_login_prompt(user))
function login(userpass::LibGit2.UserPasswordCredential = login_prompt())
    endpoint = get_endpoint()
    resp, body = withbuf(LoginBuffer(userpass)) do input
        request_body("$(endpoint)/api/login"; method="POST", input=input,
                     headers = ["content-type" => "application/json"])
    end

    if resp.status == 200
        resp_body = JSON3.read(body)
        write_cred(userpass)
        Base.shred!(userpass)
        token = resp_body.token
        save_token(token)
        return token
    else
        Base.shred!(userpass)
        status_error(resp, "$body")
    end
end

function whoami(token = get_token())
    ensure_token(token)
    endpoint = get_endpoint()
    resp, body = request_body("$(endpoint)/api/whoami-v2"; method="GET", headers = auth_header(token))

    if resp.status >= 400
        status_error(resp, "Invalid user token. Make sure you login or pass a correct token.")
    else
        return JSON3.read(body)
    end
end

"""
    logout()

Login from huggingface hub and remove all authentication cache.
"""
function logout(token = get_token())
    ensure_token(token)
    username = whoami(token).name
    erase_cred(username)

    endpoint = get_endpoint()
    resp = request("$(endpoint)/api/logout"; method="POST", headers = auth_header(token))
    if resp.status >= 400
        status_error(resp)
    else
        delete_token()
    end
end


text:
include("intcode.jl")
using Images: save, Gray
using LightGraphs

function render(plan, file)
    mincol, maxcol = extrema(first.(keys(plan)))
    minrow, maxrow = extrema(last.(keys(plan)))
    screen = zeros(Int32, maxrow - minrow + 1, maxcol - mincol + 1)
    for (coord, val) in plan
        screen[last(coord) - minrow + 1, first(coord) - mincol + 1] = val + 1
    end
    save(file, Gray.(screen / 4))
end

"""Explore the map to get a plan. Return the plan when all tiles are explored."""
function explore(input, output)
    # Plan: 0 means known space, 1 means wall, 2 means oxygen system, 3 (internal only) means fully explored
    pos = [0 0]
    plan = Dict(pos => 0)
    # Directions North, South, West, East correspond to commands 1 to 4
    cardinal_directions = [[0 -1], [0 1], [-1 0], [1 0]]
    oxygen_system_pos = nothing

    while true
        moved = false
        for (cmd, vec) in enumerate(cardinal_directions)
            # Have we already taken a look at this one? If so, skip it
            new_pos = pos + vec
            if new_pos ∈ keys(plan) continue end

            # Can we move in this direction?
            put!(input, cmd)
            result = take!(output)
            if result == 0
                # It's a wall, mark it as such
                plan[new_pos] = 1
            elseif result ∈ [1, 2]
                # We have moved to a previously unknown space, or the oxygen system space
                plan[new_pos] = 0
                pos = new_pos
                moved = true
                if result == 2 oxygen_system_pos = new_pos end
                break
            end
        end

        if !moved
            # We've looked at all directions from here now. Mark it as complete and move elsewhere
            plan[pos] = 3
            for (cmd, vec) in enumerate(cardinal_directions)
                if plan[pos + vec] == 0
                    moved = true
                    put!(input, cmd)
                    take!(output)
                    pos += vec
                    break
                end
            end
        end

        if !moved break end
    end

    plan[oxygen_system_pos] = 2
    return Dict(k => plan[k] == 3 ? 0 : plan[k] for k in keys(plan))
end

"""Create graph from plan, returning graph, index of start node, and index of the oxygen system node"""
function make_graph(plan)
    space_nodes = [k for k in keys(plan) if plan[k] ∈ [0, 2]]
    index_lookup = Dict(k => i for (i, k) in enumerate(space_nodes))
    graph = SimpleGraph(length(space_nodes))
    start_node = -1
    oxygen_system_node = -1
    for (i, k) in enumerate(space_nodes)
        # Each node can add its E and S neighbours if they exist (N and W will be added by another node)
        for neighbour in [k + [1 0], k + [0 1]]
            if neighbour ∈ keys(index_lookup) add_edge!(graph, i, index_lookup[neighbour]) end
        end
        if k == [0 0] start_node = i end
        if plan[k] == 2 oxygen_system_node = i end
    end
    return graph, start_node, oxygen_system_node
end

function main()
    # Part 1: find the oxygen system and its min-path distance from start point
    prog = open(f -> parse.(Int64, split(readline(f), ",")), "day15_input.txt")
    input = Channel(1)
    output = Channel(1)
    @async run_intcode!(copy(prog), input, output)
    plan = explore(input, output)
    render(plan, "day15_result.png")
    graph, start_node, oxygen_system_node = make_graph(plan)
    dijkstra_state = dijkstra_shortest_paths(graph, oxygen_system_node)
    println("Shortest path length: $(dijkstra_state.dists[start_node])")

    # Part 2: max distance from oxygen system
    println("Time for oxygen to fill all spaces: $(maximum(dijkstra_state.dists))")
end
main()

max_stars_repo_name: ericgreveson/adventofcode

max_issues_repo_name: ericgreveson/adventofcode

max_forks_repo_name: ericgreveson/adventofcode


text:
"""
Notes:

This file defines grid lengths, areas, and volumes for staggered structured grids.

Each "reference cell" is associated with an index i, j, k.
The "location" of each reference cell is roughly the geometric centroid of the reference cell.

On the staggered grid, there are 7 cells additional to the "reference cell"
that are staggered with respect to the reference cell in x, y, and/or z.
The staggering is denoted by the locations "Center" and "Face":

    - "Center" is shared with the reference cell;
    - "Face" lies between reference cell centers, roughly at the interface between
      reference cells.

The three-dimensional location of an object is defined by a 3-tuple of locations, and
denoted by a triplet of superscripts. For example, an object `ϕ` whose cell is located at
(Center, Center, Face) is denoted `ϕᶜᶜᶠ`. `ᶜᶜᶠ` is Centered in `x`, `Centered` in `y`, and on
reference cell interfaces in `z` (this is where the vertical velocity is located, for example).

The super script `ᵃ` denotes "any" location.

The operators in this file fall into three categories:

1. Operators needed for an algorithm valid on rectilinear grids with
   at most a stretched vertical dimension and regular horizontal dimensions.

2. Operators needed for an algorithm on a grid that is curvilinear in the horizontal
   at rectilinear (possibly stretched) in the vertical.
"""

#####
##### Grid lengths for horizontally-regular algorithms
#####

@inline Δx(i, j, k, grid::ARG) = grid.Δx
@inline Δy(i, j, k, grid::ARG) = grid.Δy

@inline ΔzC(i, j, k, grid::RegularRectilinearGrid) = grid.Δz
@inline ΔzC(i, j, k, grid::VerticallyStretchedRectilinearGrid) = @inbounds grid.Δzᵃᵃᶠ[k]

@inline ΔzF(i, j, k, grid::RegularRectilinearGrid) = grid.Δz
@inline ΔzF(i, j, k, grid::VerticallyStretchedRectilinearGrid) = @inbounds grid.Δzᵃᵃᶜ[k]

@inline Δzᵃᵃᶠ(i, j, k, grid::RegularRectilinearGrid) = grid.Δz
@inline Δzᵃᵃᶠ(i, j, k, grid::VerticallyStretchedRectilinearGrid) = @inbounds grid.Δzᵃᵃᶠ[k]

@inline Δzᵃᵃᶜ(i, j, k, grid::RegularRectilinearGrid) = grid.Δz
@inline Δzᵃᵃᶜ(i, j, k, grid::VerticallyStretchedRectilinearGrid) = @inbounds grid.Δzᵃᵃᶜ[k]

#####
##### Areas for horizontally-regular algorithms
#####

@inline Axᵃᵃᶜ(i, j, k, grid) = Δy(i, j, k, grid) * ΔzF(i, j, k, grid)
@inline Axᵃᵃᶠ(i, j, k, grid) = Δy(i, j, k, grid) * ΔzC(i, j, k, grid)

@inline Ayᵃᵃᶜ(i, j, k, grid) = Δx(i, j, k, grid) * ΔzF(i, j, k, grid)
@inline Ayᵃᵃᶠ(i, j, k, grid) = Δx(i, j, k, grid) * ΔzC(i, j, k, grid)

@inline Azᵃᵃᵃ(i, j, k, grid) = Δx(i, j, k, grid) * Δy(i, j, k, grid)

#####
##### Volumes for horizontally-regular algorithms
#####

@inline Vᵃᵃᶜ(i, j, k, grid) = Δx(i, j, k, grid) * Δy(i, j, k, grid) * ΔzF(i, j, k, grid)
@inline Vᵃᵃᶠ(i, j, k, grid) = Δx(i, j, k, grid) * Δy(i, j, k, grid) * ΔzC(i, j, k, grid)

#####
##### Grid lengths for horizontally-curvilinear, vertically-rectilinear algorithms
#####

@inline Δxᶜᶜᵃ(i, j, k, grid::ARG) = grid.Δx
@inline Δxᶜᶠᵃ(i, j, k, grid::ARG) = grid.Δx
@inline Δxᶠᶠᵃ(i, j, k, grid::ARG) = grid.Δx
@inline Δxᶠᶜᵃ(i, j, k, grid::ARG) = grid.Δx

@inline Δyᶜᶜᵃ(i, j, k, grid::ARG) = grid.Δy
@inline Δyᶠᶜᵃ(i, j, k, grid::ARG) = grid.Δy
@inline Δyᶜᶠᵃ(i, j, k, grid::ARG) = grid.Δy
@inline Δyᶠᶠᵃ(i, j, k, grid::ARG) = grid.Δy

#####
##### Areas for horizontally-curvilinear, vertically-rectilinear algorithms
#####

@inline Azᶜᶜᵃ(i, j, k, grid::ARG) = Δx(i, j, k, grid) * Δy(i, j, k, grid)
@inline Azᶠᶠᵃ(i, j, k, grid::ARG) = Δx(i, j, k, grid) * Δy(i, j, k, grid)
@inline Azᶜᶠᵃ(i, j, k, grid::ARG) = Δx(i, j, k, grid) * Δy(i, j, k, grid)
@inline Azᶠᶜᵃ(i, j, k, grid::ARG) = Δx(i, j, k, grid) * Δy(i, j, k, grid)

#####
##### Areas for three-dimensionally curvilinear algorithms
#####

@inline Axᶜᶜᶜ(i, j, k, grid::Union{ARG, AHCG}) = Δyᶜᶜᵃ(i, j, k, grid) * Δzᵃᵃᶜ(i, j, k, grid) # c
@inline Axᶠᶜᶜ(i, j, k, grid::Union{ARG, AHCG}) = Δyᶠᶜᵃ(i, j, k, grid) * Δzᵃᵃᶜ(i, j, k, grid) # u
@inline Axᶠᶠᶜ(i, j, k, grid::Union{ARG, AHCG}) = Δyᶠᶠᵃ(i, j, k, grid) * Δzᵃᵃᶜ(i, j, k, grid) # ζ
@inline Axᶠᶜᶠ(i, j, k, grid::Union{ARG, AHCG}) = Δyᶠᶜᵃ(i, j, k, grid) * Δzᵃᵃᶠ(i, j, k, grid) # η

@inline Ayᶜᶜᶜ(i, j, k, grid::Union{ARG, AHCG}) = Δxᶜᶜᵃ(i, j, k, grid) * Δzᵃᵃᶜ(i, j, k, grid) # c
@inline Ayᶜᶠᶜ(i, j, k, grid::Union{ARG, AHCG}) = Δxᶜᶠᵃ(i, j, k, grid) * Δzᵃᵃᶜ(i, j, k, grid) # v
@inline Ayᶠᶠᶜ(i, j, k, grid::Union{ARG, AHCG}) = Δxᶠᶠᵃ(i, j, k, grid) * Δzᵃᵃᶜ(i, j, k, grid) # ζ
@inline Ayᶜᶠᶠ(i, j, k, grid::Union{ARG, AHCG}) = Δxᶜᶠᵃ(i, j, k, grid) * Δzᵃᵃᶠ(i, j, k, grid) # ξ

#####
##### Volumes for three-dimensionally curvilinear algorithms
#####

@inline Vᶜᶜᶜ(i, j, k, grid::Union{ARG, AHCG}) = Azᶜᶜᵃ(i, j, k, grid) * Δzᵃᵃᶜ(i, j, k, grid)
@inline Vᶠᶜᶜ(i, j, k, grid::Union{ARG, AHCG}) = Azᶠᶜᵃ(i, j, k, grid) * Δzᵃᵃᶜ(i, j, k, grid)
@inline Vᶜᶠᶜ(i, j, k, grid::Union{ARG, AHCG}) = Azᶜᶠᵃ(i, j, k, grid) * Δzᵃᵃᶜ(i, j, k, grid)
@inline Vᶜᶜᶠ(i, j, k, grid::Union{ARG, AHCG}) = Azᶜᶜᵃ(i, j, k, grid) * Δzᵃᵃᶠ(i, j, k, grid)

#####
##### Temporary place for grid spacings and areas for RegularLatitudeLongitudeGrid
#####

@inline hack_cosd(ϕ) = cos(π * ϕ / 180)
@inline hack_sind(ϕ) = sin(π * ϕ / 180)

@inline Δxᶜᶠᵃ(i, j, k, grid::RegularLatitudeLongitudeGrid) = @inbounds grid.radius * hack_cosd(grid.ϕᵃᶠᵃ[j]) * deg2rad(grid.Δλ)
@inline Δxᶠᶜᵃ(i, j, k, grid::RegularLatitudeLongitudeGrid) = @inbounds grid.radius * hack_cosd(grid.ϕᵃᶜᵃ[j]) * deg2rad(grid.Δλ)
@inline Δxᶜᶜᵃ(i, j, k, grid::RegularLatitudeLongitudeGrid) = Δxᶠᶜᵃ(i, j, k, grid)
@inline Δxᶠᶠᵃ(i, j, k, grid::RegularLatitudeLongitudeGrid) = Δxᶜᶠᵃ(i, j, k, grid)

@inline Δyᶜᶠᵃ(i, j, k, grid::RegularLatitudeLongitudeGrid) = @inbounds grid.radius * deg2rad(grid.Δϕ)
@inline Δyᶠᶜᵃ(i, j, k, grid::RegularLatitudeLongitudeGrid) = Δyᶜᶠᵃ(i, j, k, grid)
@inline Δyᶜᶜᵃ(i, j, k, grid::RegularLatitudeLongitudeGrid) = Δyᶜᶠᵃ(i, j, k, grid)
@inline Δyᶠᶠᵃ(i, j, k, grid::RegularLatitudeLongitudeGrid) = Δyᶜᶠᵃ(i, j, k, grid)

@inline Δzᵃᵃᶜ(i, j, k, grid::RegularLatitudeLongitudeGrid) = grid.Δz
@inline Δzᵃᵃᶠ(i, j, k, grid::RegularLatitudeLongitudeGrid) = grid.Δz

@inline Azᶜᶜᵃ(i, j, k, grid::RegularLatitudeLongitudeGrid) = @inbounds grid.radius^2 * deg2rad(grid.Δλ) * (hack_sind(grid.ϕᵃᶠᵃ[j+1]) - hack_sind(grid.ϕᵃᶠᵃ[j]))
@inline Azᶠᶠᵃ(i, j, k, grid::RegularLatitudeLongitudeGrid) = @inbounds grid.radius^2 * deg2rad(grid.Δλ) * (hack_sind(grid.ϕᵃᶜᵃ[j])   - hack_sind(grid.ϕᵃᶜᵃ[j-1]))
@inline Azᶠᶜᵃ(i, j, k, grid::RegularLatitudeLongitudeGrid) = Azᶜᶜᵃ(i, j, k, grid)
@inline Azᶜᶠᵃ(i, j, k, grid::RegularLatitudeLongitudeGrid) = Azᶠᶠᵃ(i, j, k, grid)