text:
struct OrthographicCamera <: AbstractCamera end

"""
    layoutable(Axis3, fig_or_scene; bbox = nothing, kwargs...)

Creates an `Axis3` object in the parent `fig_or_scene` which consists of a child scene
with orthographic projection for 2D plots and axis decorations that live in the
parent.
"""
function layoutable(::Type{<:Axis3}, fig_or_scene::Union{Figure, Scene}; bbox = nothing, kwargs...)

    topscene = get_topscene(fig_or_scene)

    default_attrs = default_attributes(Axis3, topscene).attributes
    theme_attrs = subtheme(topscene, :Axis3)
    attrs = merge!(merge!(Attributes(kwargs), theme_attrs), default_attrs)

    @extract attrs (elevation, azimuth, perspectiveness, aspect, viewmode,
        xlabel, ylabel, zlabel,
    )

    decorations = Dict{Symbol, Any}()

    protrusions = lift(to_protrusions, attrs.protrusions)
    layoutobservables = LayoutObservables{Axis3}(attrs.width, attrs.height, attrs.tellwidth, attrs.tellheight, attrs.halign, attrs.valign, attrs.alignmode;
        suggestedbbox = bbox, protrusions = protrusions)

    notify(protrusions)

    finallimits = Observable(Rect3f(Vec3f(0f0, 0f0, 0f0), Vec3f(100f0, 100f0, 100f0)))

    scenearea = lift(round_to_IRect2D, layoutobservables.computedbbox)

    scene = Scene(topscene, scenearea, clear = false, backgroundcolor = attrs.backgroundcolor)
    cam = OrthographicCamera()
    cameracontrols!(scene, cam)

    matrices = lift(calculate_matrices, finallimits, scene.px_area, elevation, azimuth, perspectiveness, aspect, viewmode)

    on(matrices) do (view, proj, eyepos)
        cam = camera(scene)
        Makie.set_proj_view!(cam, proj, view)
        cam.eyeposition[] = eyepos
    end

    ticknode_1 = lift(finallimits, attrs.xticks, attrs.xtickformat) do lims, ticks, format
        tl = get_ticks(ticks, identity, format, minimum(lims)[1], maximum(lims)[1])
    end

    ticknode_2 = lift(finallimits, attrs.yticks, attrs.ytickformat) do lims, ticks, format
        tl = get_ticks(ticks, identity, format, minimum(lims)[2], maximum(lims)[2])
    end

    ticknode_3 = lift(finallimits, attrs.zticks, attrs.ztickformat) do lims, ticks, format
        tl = get_ticks(ticks, identity, format, minimum(lims)[3], maximum(lims)[3])
    end

    mi1 = @lift(!(pi/2 <= $azimuth % 2pi < 3pi/2))
    mi2 = @lift(0 <= $azimuth % 2pi < pi)
    mi3 = @lift($elevation > 0)

    add_panel!(scene, 1, 2, 3, finallimits, mi3, attrs)
    add_panel!(scene, 2, 3, 1, finallimits, mi1, attrs)
    add_panel!(scene, 1, 3, 2, finallimits, mi2, attrs)

    add_gridlines_and_frames!(topscene, scene, 1, finallimits, ticknode_1, mi1, mi2, mi3, attrs)
    add_gridlines_and_frames!(topscene, scene, 2, finallimits, ticknode_2, mi2, mi1, mi3, attrs)
    add_gridlines_and_frames!(topscene, scene, 3, finallimits, ticknode_3, mi3, mi1, mi2, attrs)

    add_ticks_and_ticklabels!(topscene, scene, 1, finallimits, ticknode_1, mi1, mi2, mi3, attrs, azimuth)
    add_ticks_and_ticklabels!(topscene, scene, 2, finallimits, ticknode_2, mi2, mi1, mi3, attrs, azimuth)
    add_ticks_and_ticklabels!(topscene, scene, 3, finallimits, ticknode_3, mi3, mi1, mi2, attrs, azimuth)

    titlepos = lift(scene.px_area, attrs.titlegap, attrs.titlealign) do a, titlegap, align

        x = if align == :center
            a.origin[1] + a.widths[1] / 2
        elseif align == :left
            a.origin[1]
        elseif align == :right
            a.origin[1] + a.widths[1]
        else
            error("Title align $align not supported.")
        end

        yoffset = top(a) + titlegap

        Point2(x, yoffset)
    end

    titlealignnode = lift(attrs.titlealign) do align
        (align, :bottom)
    end

    titlet = text!(
        topscene, attrs.title,
        position = titlepos,
        visible = attrs.titlevisible,
        textsize = attrs.titlesize,
        align = titlealignnode,
        font = attrs.titlefont,
        color = attrs.titlecolor,
        space = :data,
        inspectable = false)
    decorations[:title] = titlet

    mouseeventhandle = addmouseevents!(scene)
    scrollevents = Observable(ScrollEvent(0, 0))
    keysevents = Observable(KeysEvent(Set()))

    on(scene.events.scroll) do s
        if is_mouseinside(scene)
            scrollevents[] = ScrollEvent(s[1], s[2])
            return Consume(true)
        end
        return Consume(false)
    end

    on(scene.events.keyboardbutton) do e
        keysevents[] = KeysEvent(scene.events.keyboardstate)
        return Consume(false)
    end

    interactions = Dict{Symbol, Tuple{Bool, Any}}()

    ax = Axis3(fig_or_scene, layoutobservables, attrs, decorations, scene, finallimits,
        mouseeventhandle, scrollevents, keysevents, interactions, Cycler())

    on(attrs.limits) do lims
        reset_limits!(ax)
    end

    on(attrs.targetlimits) do lims
        # adjustlimits!(ax)
        # we have no aspect constraints here currently, so just update final limits
        ax.finallimits[] = lims
    end

    function process_event(event)
        for (active, interaction) in values(ax.interactions)
            if active
                maybe_consume = process_interaction(interaction, event, ax)
                maybe_consume == Consume(true) && return Consume(true)
            end
        end
        return Consume(false)
    end

    on(process_event, mouseeventhandle.obs)
    on(process_event, scrollevents)
    on(process_event, keysevents)

    register_interaction!(ax,
        :dragrotate,
        DragRotate())

    # in case the user set limits already
    notify(attrs.limits)

    ax
end

can_be_current_axis(ax3::Axis3) = true

function calculate_matrices(limits, px_area, elev, azim, perspectiveness, aspect,
    viewmode)
    ws = widths(limits)

    t = Makie.translationmatrix(-Float64.(limits.origin))
    s = if aspect == :equal
        scales = 2 ./ Float64.(ws)
    elseif aspect == :data
        scales = 2 ./ max.(maximum(ws), Float64.(ws))
    elseif aspect isa VecTypes{3}
        scales = 2 ./ Float64.(ws) .* Float64.(aspect) ./ maximum(aspect)
    else
        error("Invalid aspect $aspect")
    end |> Makie.scalematrix

    t2 = Makie.translationmatrix(-0.5 .* ws .* scales)
    scale_matrix = t2 * s * t

    ang_max = 90
    ang_min = 0.5

    @assert 0 <= perspectiveness <= 1

    angle = ang_min + (ang_max - ang_min) * perspectiveness

    # vFOV = 2 * Math.asin(sphereRadius / distance);
    # distance = sphere_radius / Math.sin(vFov / 2)

    # radius = sqrt(3) / tand(angle / 2)
    radius = sqrt(3) / sind(angle / 2)

    x = radius * cos(elev) * cos(azim)
    y = radius * cos(elev) * sin(azim)
    z = radius * sin(elev)

    eyepos = Vec3{Float64}(x, y, z)

    lookat_matrix = Makie.lookat(
        eyepos,
        Vec3{Float64}(0, 0, 0),
        Vec3{Float64}(0, 0, 1))

    w = width(px_area)
    h = height(px_area)

    view_matrix = lookat_matrix * scale_matrix

    projection_matrix = projectionmatrix(view_matrix, limits, eyepos, radius, azim, elev, angle, w, h, scales, viewmode)

    # for eyeposition dependent algorithms, we need to present the position as if
    # there was no scaling applied
    eyeposition = Vec3f(inv(scale_matrix) * Vec4f(eyepos..., 1))

    view_matrix, projection_matrix, eyeposition
end

function projectionmatrix(viewmatrix, limits, eyepos, radius, azim, elev, angle, width, height, scales, viewmode)
    near = 0.5 * (radius - sqrt(3))
    far = radius + 2 * sqrt(3)

    aspect_ratio = width / height

    projection_matrix = if viewmode in (:fit, :fitzoom, :stretch)
        if height > width
            angle = angle / aspect_ratio
        end

        pm = Makie.perspectiveprojection(Float64, angle, aspect_ratio, near, far)

        if viewmode in (:fitzoom, :stretch)
            points = decompose(Point3f, limits)
            # @show points
            projpoints = Ref(pm * viewmatrix) .* to_ndim.(Point4f, points, 1)

            maxx = maximum(x -> abs(x[1] / x[4]), projpoints)
            maxy = maximum(x -> abs(x[2] / x[4]), projpoints)

            ratio_x = maxx
            ratio_y = maxy

            if viewmode == :fitzoom
                if ratio_y > ratio_x
                    pm = Makie.scalematrix(Vec3(1/ratio_y, 1/ratio_y, 1)) * pm
                else
                    pm = Makie.scalematrix(Vec3(1/ratio_x, 1/ratio_x, 1)) * pm
                end
            else
                pm = Makie.scalematrix(Vec3(1/ratio_x, 1/ratio_y, 1)) * pm
            end
        end
        pm
    else
        error("Invalid viewmode $viewmode")
    end
end

function Makie.plot!(
    ax::Axis3, P::Makie.PlotFunc,
    attributes::Makie.Attributes, args...;
    kw_attributes...)

    allattrs = merge(attributes, Attributes(kw_attributes))

    cycle = get_cycle_for_plottype(allattrs, P)
    add_cycle_attributes!(allattrs, P, cycle, ax.cycler, ax.palette)

    plot = Makie.plot!(ax.scene, P, allattrs, args...)

    reset_limits!(ax)
    plot
end

function Makie.plot!(P::Makie.PlotFunc, ax::Axis3, args...; kw_attributes...)
    attributes = Makie.Attributes(kw_attributes)
    Makie.plot!(ax, P, attributes, args...)
end

function autolimits!(ax::Axis3)
    xlims = getlimits(ax, 1)
    ylims = getlimits(ax, 2)
    zlims = getlimits(ax, 3)

    ori = Vec3f(xlims[1], ylims[1], zlims[1])
    widths = Vec3f(xlims[2] - xlims[1], ylims[2] - ylims[1], zlims[2] - zlims[1])

    enlarge_factor = 0.1

    nori = ori .- (0.5 * enlarge_factor) * widths
    nwidths = widths .* (1 + enlarge_factor)

    lims = Rect3f(nori, nwidths)

    ax.finallimits[] = lims
    nothing
end

to_protrusions(x::Number) = GridLayoutBase.RectSides{Float32}(x, x, x, x)
to_protrusions(x::Tuple{Any, Any, Any, Any}) = GridLayoutBase.RectSides{Float32}(x...)

function getlimits(ax::Axis3, dim)

    plots_with_autolimits = if dim == 1
        filter(p -> !haskey(p.attributes, :xautolimits) || p.attributes.xautolimits[], ax.scene.plots)
    elseif dim == 2
        filter(p -> !haskey(p.attributes, :yautolimits) || p.attributes.yautolimits[], ax.scene.plots)
    elseif dim == 3
        filter(p -> !haskey(p.attributes, :zautolimits) || p.attributes.zautolimits[], ax.scene.plots)
    else
        error("Dimension $dim not allowed. Only 1, 2 or 3.")
    end

    visible_plots = filter(
        p -> !haskey(p.attributes, :visible) || p.attributes.visible[],
        plots_with_autolimits)

    bboxes = Makie.data_limits.(visible_plots)
    finite_bboxes = filter(Makie.isfinite_rect, bboxes)

    isempty(finite_bboxes) && return nothing

    templim = (finite_bboxes[1].origin[dim], finite_bboxes[1].origin[dim] + finite_bboxes[1].widths[dim])

    for bb in finite_bboxes[2:end]
        templim = limitunion(templim, (bb.origin[dim], bb.origin[dim] + bb.widths[dim]))
    end

    templim
end

# mutable struct LineAxis3D

# end

function dimpoint(dim, v, v1, v2)
    if dim == 1
        Point(v, v1, v2)
    elseif dim == 2
        Point(v1, v, v2)
    elseif dim == 3
        Point(v1, v2, v)
    end
end

function dim1(dim)
    if dim == 1
        2
    elseif dim == 2
        1
    elseif dim == 3
        1
    end
end

function dim2(dim)
    if dim == 1
        3
    elseif dim == 2
        3
    elseif dim == 3
        2
    end
end

function add_gridlines_and_frames!(topscene, scene, dim::Int, limits, ticknode, miv, min1, min2, attrs)

    dimsym(sym) = Symbol(string((:x, :y, :z)[dim]) * string(sym))
    attr(sym) = attrs[dimsym(sym)]

    dpoint = (v, v1, v2) -> dimpoint(dim, v, v1, v2)
    d1 = dim1(dim)
    d2 = dim2(dim)

    tickvalues = @lift($ticknode[1])

    endpoints = lift(limits, tickvalues, min1, min2) do lims, ticks, min1, min2
        f1 = min1 ? minimum(lims)[d1] : maximum(lims)[d1]
        f2 = min2 ? minimum(lims)[d2] : maximum(lims)[d2]
        # from tickvalues and f1 and min2:max2
        mi = minimum(lims)
        ma = maximum(lims)
        map(filter(x -> !any(y -> x ≈ y[dim], extrema(lims)), ticks)) do t
            dpoint(t, f1, mi[d2]), dpoint(t, f1, ma[d2])
        end
    end
    linesegments!(scene, endpoints, color = attr(:gridcolor),
        linewidth = attr(:gridwidth),
        xautolimits = false, yautolimits = false, zautolimits = false, transparency = true,
        visible = attr(:gridvisible), inspectable = false)

    endpoints2 = lift(limits, tickvalues, min1, min2) do lims, ticks, min1, min2
        f1 = min1 ? minimum(lims)[d1] : maximum(lims)[d1]
        f2 = min2 ? minimum(lims)[d2] : maximum(lims)[d2]
        # from tickvalues and f1 and min2:max2
        mi = minimum(lims)
        ma = maximum(lims)
        map(filter(x -> !any(y -> x ≈ y[dim], extrema(lims)), ticks)) do t
            dpoint(t, mi[d1], f2), dpoint(t, ma[d1], f2)
        end
    end
    linesegments!(scene, endpoints2, color = attr(:gridcolor),
        linewidth = attr(:gridwidth),
        xautolimits = false, yautolimits = false, zautolimits = false, transparency = true,
        visible = attr(:gridvisible), inspectable = false)

    framepoints = lift(limits, min1, min2,
            scene.camera.projectionview, scene.px_area) do lims, mi1, mi2, pview, pxa

        o = pxa.origin

        f(mi) = mi ? minimum : maximum
        p1 = dpoint(minimum(lims)[dim], f(!mi1)(lims)[d1], f(mi2)(lims)[d2])
        p2 = dpoint(maximum(lims)[dim], f(!mi1)(lims)[d1], f(mi2)(lims)[d2])
        p3 = dpoint(minimum(lims)[dim], f(mi1)(lims)[d1], f(mi2)(lims)[d2])
        p4 = dpoint(maximum(lims)[dim], f(mi1)(lims)[d1], f(mi2)(lims)[d2])
        p5 = dpoint(minimum(lims)[dim], f(mi1)(lims)[d1], f(!mi2)(lims)[d2])
        p6 = dpoint(maximum(lims)[dim], f(mi1)(lims)[d1], f(!mi2)(lims)[d2])
        # p7 = dpoint(minimum(lims)[dim], f(!mi1)(lims)[d1], f(!mi2)(lims)[d2])
        # p8 = dpoint(maximum(lims)[dim], f(!mi1)(lims)[d1], f(!mi2)(lims)[d2])

        # we are going to transform the 3d frame points into 2d of the topscene
        # because otherwise the frame lines can
        # be cut when they lie directly on the scene boundary
        to_topscene_z_2d.([p1, p2, p3, p4, p5, p6], Ref(scene))
    end

    colors = lift(vcat, Any, attr(:spinecolor_1), attr(:spinecolor_2), attr(:spinecolor_3))
    linesegments!(topscene, framepoints, color = colors, linewidth = attr(:spinewidth),
        # transparency = true,
        visible = attr(:spinesvisible), inspectable = false)

    nothing
end

# this function projects a point from a 3d subscene into the parent space with a really
# small z value
function to_topscene_z_2d(p3d, scene)
    o = scene.px_area[].origin
    p2d = Point2f(o + Makie.project(scene, p3d))
    # -10000 is an arbitrary weird constant that in preliminary testing didn't seem
    # to clip into plot objects anymore
    Point3f(p2d..., -10000)
end

function add_ticks_and_ticklabels!(topscene, scene, dim::Int, limits, ticknode, miv, min1, min2, attrs, azimuth)

    dimsym(sym) = Symbol(string((:x, :y, :z)[dim]) * string(sym))
    attr(sym) = attrs[dimsym(sym)]

    dpoint = (v, v1, v2) -> dimpoint(dim, v, v1, v2)
    d1 = dim1(dim)
    d2 = dim2(dim)

    tickvalues = @lift($ticknode[1])
    ticklabels = @lift($ticknode[2])

    tick_segments = lift(limits, tickvalues, miv, min1, min2,
            scene.camera.projectionview, scene.px_area) do lims, ticks, miv, min1, min2,
                pview, pxa

        f1 = !min1 ? minimum(lims)[d1] : maximum(lims)[d1]
        f2 = min2 ? minimum(lims)[d2] : maximum(lims)[d2]

        f1_oppo = min1 ? minimum(lims)[d1] : maximum(lims)[d1]
        f2_oppo = !min2 ? minimum(lims)[d2] : maximum(lims)[d2]

        diff_f1 = f1 - f1_oppo
        diff_f2 = f2 - f2_oppo

        map(ticks) do t
            p1 = dpoint(t, f1, f2)
            p2 = if dim == 3
                # special case the z axis, here it depends on azimuth in which direction the ticks go
                if 45 <= (rad2deg(azimuth[]) % 180) <= 135
                    dpoint(t, f1 + 0.03 * diff_f1, f2)
                else
                    dpoint(t, f1, f2 + 0.03 * diff_f2)
                end
            else
                dpoint(t, f1 + 0.03 * diff_f1, f2)
            end

            (p1, p2)
        end
    end

    # we are going to transform the 3d tick segments into 2d of the topscene
    # because otherwise they
    # be cut when they extend beyond the scene boundary
    tick_segments_2dz = lift(tick_segments,
            scene.camera.projectionview, scene.px_area) do ts, _, _
        map(ts) do p1_p2
            to_topscene_z_2d.(p1_p2, Ref(scene))
        end
    end

    linesegments!(topscene, tick_segments_2dz,
        xautolimits = false, yautolimits = false, zautolimits = false,
        transparency = true, inspectable = false,
        color = attr(:tickcolor), linewidth = attr(:tickwidth), visible = attr(:ticksvisible))

    labels_positions = lift(scene.px_area, scene.camera.projectionview,
            tick_segments, ticklabels, attr(:ticklabelpad)) do pxa, pv, ticksegs, ticklabs, pad

        o = pxa.origin

        points = map(ticksegs) do (tstart, tend)
            tstartp = Point2f(o + Makie.project(scene, tstart))
            tendp = Point2f(o + Makie.project(scene, tend))

            offset = pad * Makie.GeometryBasics.normalize(
                Point2f(tendp - tstartp))
            tendp + offset
        end

        v = collect(zip(ticklabs, points))
        v::Vector{Tuple{String, Point2f}}
    end

    align = lift(miv, min1, min2) do mv, m1, m2
        if dim == 1
            (mv ⊻ m1 ? :right : :left, m2 ? :top : :bottom)
        elseif dim == 2
            (mv ⊻ m1 ? :left : :right, m2 ? :top : :bottom)
        elseif dim == 3
            (m1 ⊻ m2 ? :left : :right, :center)
        end
    end

    ticklabel_obj = text!(topscene, labels_positions, align = align,
        color = attr(:ticklabelcolor), textsize = attr(:ticklabelsize),
        font = attr(:ticklabelfont), visible = attr(:ticklabelsvisible), inspectable = false
    )

    translate!(ticklabel_obj, 0, 0, 1000)

    label_pos_rot_valign = lift(scene.px_area, scene.camera.projectionview,
            limits, miv, min1, min2, attr(:labeloffset),
            attr(:labelrotation)) do pxa, pv, lims, miv, min1, min2,
                labeloffset, lrotation

        o = pxa.origin

        f1 = !min1 ? minimum(lims)[d1] : maximum(lims)[d1]
        f2 = min2 ? minimum(lims)[d2] : maximum(lims)[d2]

        # get end points of axis
        p1 = dpoint(minimum(lims)[dim], f1, f2)
        p2 = dpoint(maximum(lims)[dim], f1, f2)

        # project them into screen space
        pp1 = Point2f(o + Makie.project(scene, p1))
        pp2 = Point2f(o + Makie.project(scene, p2))

        # find the midpoint
        midpoint = (pp1 + pp2) / 2

        # and the difference vector
        diff = pp2 - pp1

        diffsign = if dim == 1 || dim == 3
            !(min1 ⊻ min2) ? 1 : -1
        else
            (min1 ⊻ min2) ? 1 : -1
        end

        a = pi/2

        # get the vector pointing from the axis in the direction of the label anchor
        offset_vec = (Makie.Mat2f(cos(a), sin(a), -sin(a), cos(a)) *
            Makie.GeometryBasics.normalize(diffsign * diff))

        # calculate the label offset from the axis midpoint
        plus_offset = midpoint + labeloffset * offset_vec

        offset_ang = atan(offset_vec[2], offset_vec[1])
        offset_ang_90deg = offset_ang + pi/2
        offset_ang_90deg_alwaysup = ((offset_ang + pi/2 + pi/2) % pi) - pi/2

        # # prefer rotated left 90deg to rotated right 90deg
        slight_flip = offset_ang_90deg_alwaysup < -deg2rad(88)
        if slight_flip
            offset_ang_90deg_alwaysup += pi
        end
        offset_ang_90deg_alwaysup

        labelrotation = if lrotation == Makie.automatic
            offset_ang_90deg_alwaysup
        else
            lrotation
        end

        valign = offset_vec[2] > 0 || slight_flip ? :bottom : :top

        plus_offset, labelrotation, valign
    end

    labelalign = lift(label_pos_rot_valign, attr(:labelalign)) do (_, _, valign), lalign
        if lalign == Makie.automatic
            (:center, valign)
        else
            lalign
        end
    end
    text!(topscene, attr(:label),
        color = attr(:labelcolor),
        textsize = attr(:labelsize),
        font = attr(:labelfont),
        position = @lift($label_pos_rot_valign[1]),
        rotation = @lift($label_pos_rot_valign[2]),
        align = labelalign,
        visible = attr(:labelvisible),
        inspectable = false
    )
    return
end

function dim3point(dim1, dim2, dim3, v1, v2, v3)
    if (dim1, dim2, dim3) == (1, 2, 3)
        Point(v1, v2, v3)
    elseif (dim1, dim2, dim3) == (2, 3, 1)
        Point(v3, v1, v2)
    elseif (dim1, dim2, dim3) == (1, 3, 2)
        Point(v1, v3, v2)
    else
        error("Invalid dim order $dim1, $dim2, $dim3")
    end
end

function add_panel!(scene, dim1, dim2, dim3, limits, min3, attrs)

    dimsym(sym) = Symbol(string((:x, :y, :z)[dim1]) *
        string((:x, :y, :z)[dim2]) * string(sym))
    attr(sym) = attrs[dimsym(sym)]

    vertices = lift(limits, min3) do lims, mi3

        mi = minimum(lims)
        ma = maximum(lims)

        v3 = if mi3
            mi[dim3] + 0.005 * (mi[dim3] - ma[dim3])
        else
            ma[dim3] + 0.005 * (ma[dim3] - mi[dim3])
        end

        p1 = dim3point(dim1, dim2, dim3, mi[dim1], mi[dim2], v3)
        p2 = dim3point(dim1, dim2, dim3, mi[dim1], ma[dim2], v3)
        p3 = dim3point(dim1, dim2, dim3, ma[dim1], ma[dim2], v3)
        p4 = dim3point(dim1, dim2, dim3, ma[dim1], mi[dim2], v3)
        [p1, p2, p3, p4]
    end

    faces = [1 2 3; 3 4 1]

    mesh!(scene, vertices, faces, shading = false, inspectable = false,
        xautolimits = false, yautolimits = false, zautolimits = false,
        color = attr(:panelcolor), visible = attr(:panelvisible))
end

function hidexdecorations!(ax::Axis3;
    label = true, ticklabels = true, ticks = true, grid = true)

    if label
        ax.xlabelvisible = false
    end
    if ticklabels
        ax.xticklabelsvisible = false
    end
    if ticks
        ax.xticksvisible = false
    end
    if grid
        ax.xgridvisible = false
    end
    # if minorgrid
    #     ax.xminorgridvisible = false
    # end
    # if minorticks
    #     ax.xminorticksvisible = false
    # end

    ax
end

function hideydecorations!(ax::Axis3;
    label = true, ticklabels = true, ticks = true, grid = true)

    if label
        ax.ylabelvisible = false
    end
    if ticklabels
        ax.yticklabelsvisible = false
    end
    if ticks
        ax.yticksvisible = false
    end
    if grid
        ax.ygridvisible = false
    end
    # if minorgrid
    #     ax.yminorgridvisible = false
    # end
    # if minorticks
    #     ax.yminorticksvisible = false
    # end

    ax
end

function hidezdecorations!(ax::Axis3;
    label = true, ticklabels = true, ticks = true, grid = true)

    if label
        ax.zlabelvisible = false
    end
    if ticklabels
        ax.zticklabelsvisible = false
    end
    if ticks
        ax.zticksvisible = false
    end
    if grid
        ax.zgridvisible = false
    end
    # if minorgrid
    #     ax.zminorgridvisible = false
    # end
    # if minorticks
    #     ax.zminorticksvisible = false
    # end

    ax
end

function hidedecorations!(ax::Axis3;
    label = true, ticklabels = true, ticks = true, grid = true)

    hidexdecorations!(ax; label = label, ticklabels = ticklabels,
        ticks = ticks, grid = grid)
    hideydecorations!(ax; label = label, ticklabels = ticklabels,
        ticks = ticks, grid = grid)
    hidezdecorations!(ax; label = label, ticklabels = ticklabels,
        ticks = ticks, grid = grid)

    ax
end

function hidespines!(ax::Axis3)
    ax.xspinesvisible = false
    ax.yspinesvisible = false
    ax.zspinesvisible = false
    ax
end

# this is so users can do limits = (x1, x2, y1, y2, z1, z2)
function convert_limit_attribute(lims::Tuple{Any, Any, Any, Any, Any, Any})
    (lims[1:2], lims[3:4], lims[5:6])
end

function convert_limit_attribute(lims::Tuple{Any, Any, Any})
    lims
end

function xautolimits(ax::Axis3)
    xlims = getlimits(ax, 1)

    if isnothing(xlims)
        xlims = (ax.targetlimits[].origin[1], ax.targetlimits[].origin[1] + ax.targetlimits[].widths[1])
    else
        xlims = expandlimits(xlims,
            ax.attributes.xautolimitmargin[][1],
            ax.attributes.xautolimitmargin[][2],
            identity)
    end
    xlims
end

function yautolimits(ax::Axis3)
    ylims = getlimits(ax, 2)

    if isnothing(ylims)
        ylims = (ax.targetlimits[].origin[2], ax.targetlimits[].origin[2] + ax.targetlimits[].widths[2])
    else
        ylims = expandlimits(ylims,
            ax.attributes.yautolimitmargin[][1],
            ax.attributes.yautolimitmargin[][2],
            identity)
    end
    ylims
end

function zautolimits(ax::Axis3)
    zlims = getlimits(ax, 3)

    if isnothing(zlims)
        zlims = (ax.targetlimits[].origin[3], ax.targetlimits[].origin[3] + ax.targetlimits[].widths[3])
    else
        zlims = expandlimits(zlims,
            ax.attributes.zautolimitmargin[][1],
            ax.attributes.zautolimitmargin[][2],
            identity)
    end
    zlims
end

function Makie.xlims!(ax::Axis3, xlims::Tuple{Union{Real, Nothing}, Union{Real, Nothing}})
    if length(xlims) != 2
        error("Invalid xlims length of $(length(xlims)), must be 2.")
    elseif xlims[1] == xlims[2]
        error("Can't set x limits to the same value $(xlims[1]).")
    # elseif all(x -> x isa Real, xlims) && xlims[1] > xlims[2]
    #     xlims = reverse(xlims)
    #     ax.xreversed[] = true
    # else
    #     ax.xreversed[] = false
    end

    ax.limits.val = (xlims, ax.limits[][2], ax.limits[][3])
    reset_limits!(ax, yauto = false, zauto = false)
    nothing
end

function Makie.ylims!(ax::Axis3, ylims::Tuple{Union{Real, Nothing}, Union{Real, Nothing}})
    if length(ylims) != 2
        error("Invalid ylims length of $(length(ylims)), must be 2.")
    elseif ylims[1] == ylims[2]
        error("Can't set y limits to the same value $(ylims[1]).")
    # elseif all(x -> x isa Real, ylims) && ylims[1] > ylims[2]
    #     ylims = reverse(ylims)
    #     ax.yreversed[] = true
    # else
    #     ax.yreversed[] = false
    end

    ax.limits.val = (ax.limits[][1], ylims, ax.limits[][3])
    reset_limits!(ax, xauto = false, zauto = false)
    nothing
end

function Makie.zlims!(ax::Axis3, zlims)
    if length(zlims) != 2
        error("Invalid zlims length of $(length(zlims)), must be 2.")
    elseif zlims[1] == zlims[2]
        error("Can't set y limits to the same value $(zlims[1]).")
    # elseif all(x -> x isa Real, zlims) && zlims[1] > zlims[2]
    #     zlims = reverse(zlims)
    #     ax.zreversed[] = true
    # else
    #     ax.zreversed[] = false
    end

    ax.limits.val = (ax.limits[][1], ax.limits[][2], zlims)
    reset_limits!(ax, xauto = false, yauto = false)
    nothing
end

"""
    limits!(ax::Axis3, xlims, ylims)

Set the axis limits to `xlims` and `ylims`.
If limits are ordered high-low, this reverses the axis orientation.
"""
function limits!(ax::Axis3, xlims, ylims, zlims)
    Makie.xlims!(ax, xlims)
    Makie.ylims!(ax, ylims)
    Makie.zlims!(ax, zlims)
end

"""
    limits!(ax::Axis3, x1, x2, y1, y2, z1, z2)

Set the axis x-limits to `x1` and `x2` and the y-limits to `y1` and `y2`.
If limits are ordered high-low, this reverses the axis orientation.
"""
function limits!(ax::Axis3, x1, x2, y1, y2, z1, z2)
    Makie.xlims!(ax, x1, x2)
    Makie.ylims!(ax, y1, y2)
    Makie.zlims!(ax, z1, z2)
end

"""
    limits!(ax::Axis3, rect::Rect3)

Set the axis limits to `rect`.
If limits are ordered high-low, this reverses the axis orientation.
"""
function limits!(ax::Axis3, rect::Rect3)
    xmin, ymin, zmin = minimum(rect)
    xmax, ymax, zmax = maximum(rect)
    Makie.xlims!(ax, xmin, xmax)
    Makie.ylims!(ax, ymin, ymax)
    Makie.zlims!(ax, zmin, zmax)
end


text:
using LinearAlgebra
@testset "raytracing" begin
    c = ATools.Cylinder(1.0, 0.0, 1.0)

    # check that point (1,0,0) verifies cylinder equation
    @test ATools.cylinder_equation(c, [1.0, 0.0, 0.0]) ≈ 0.0
    @test ATools.cylinder_equation(c, [0.0, 1.0, 0.0]) ≈ 0.0
    @test isapprox(ATools.cylinder_equation(c, [1.0/sqrt(2.0), 1.0/sqrt(2.0), 0.0]), 0.0, atol=1e-9)

    # check that normal to barrel is perpendicular to z axis
    nb = ATools.normal_to_barrel(c,  [1.0, 2.0, 3.0])
    @test dot(nb,[0.0, 0.0, 1.0]) ≈ 0.0

    #check functions for unit cylinder
    @test ATools.length(c) ≈ 1.0
    @test ATools.perimeter(c) ≈ 2 * π
    @test ATools.area_barrel(c) ≈ 2 * π
    @test ATools.area_endcap(c) ≈ π
    ATools.area(c) ≈ 4 * π
    ATools.volume(c) ≈ π
end

ext: jl

max_stars_repo_name: andLaing/ATools

max_issues_repo_name: andLaing/ATools

max_forks_repo_name: andLaing/ATools


text:
using SparseArrays
using CUDA
using Random
function D2_test(idata,odata,Nx,Ny,h)
    for i = 1:Nx
        for j = 1:Ny
            global_index_x = (i-1)*Ny+j

            # idata_index_x = (i-1)*Ny + j + test_mod(i,Nx) * Nx
            # idata_index_y = (i-1)*Ny + j + test_mod(j,Ny)

            idata_index_x = (i-1)*Ny + j + (div(2*Nx-i-3,Nx-2) - 1) * Nx
            idata_index_y = (i-1)*Ny + j + div(2*Ny-j-3,Ny-2) - 1

            odata[global_index_x] = ( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])
                                    + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) ) / h^2
        end
    end
end

function D2_test_v2(idata,odata,Nx,Ny,h)
    for i = 1:Nx
        for j = 1:Ny
            global_index_x = (i-1)*Ny+j

            # idata_index_x = (i-1)*Ny + j + test_mod(i,Nx) * Nx
            # idata_index_y = (i-1)*Ny + j + test_mod(j,Ny)

            offset_x = div(2*Nx-i-3,Nx-2) - 1
            offset_y = div(2*Ny-j-3,Ny-2) - 1

            idata_index_x = (i-1)*Ny + j + (offset_x) * Nx
            idata_index_y = (i-1)*Ny + j + offset_y

            odata[global_index_x] = ( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])
                                    + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) ) / 2^(abs(offset_x) + abs(offset_y))
        end
    end
end

function matrix_free_A_v2(idata,odata,Nx,Ny,h)
    # alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    odata .= 0
    for i = 1:Nx
        for j = 1:Ny
            global_index = (i-1)*Ny+j

            # idata_index_x = (i-1)*Ny + j + test_mod(i,Nx) * Nx
            # idata_index_y = (i-1)*Ny + j + test_mod(j,Ny)

            offset_x = div(2*Nx-i-3,Nx-2) - 1
            offset_y = div(2*Ny-j-3,Ny-2) - 1

            idata_index_x = (i-1)*Ny + j + (offset_x) * Nx
            idata_index_y = (i-1)*Ny + j + offset_y

            # odata[global_index] = 0

            odata[global_index] += ( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])  + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) ) / (h^2)  # D2

            odata[global_index] += abs(offset_y) * 2 * alpha1 * ( (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y]) / h^2    # BS_y

            odata[global_index] += abs(offset_x) * ( beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha4 * 2* (idata[global_index]) / h) # BS_x'
            odata[global_index+Ny*offset_x] += abs(offset_x) * beta * 2 * (-1*idata[global_index]) / h^2 # BS_x'
            odata[global_index+2*Ny*offset_x] += abs(offset_x) * beta * (0.5*idata[global_index]) / h^2 # BS_x'
        end
    end

end

function matrix_free_A_v4(idata,odata,Nx,Ny,h)
    # one function that does D2 + boundary conditions + H_tilde
    # alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    odata .= 0
    offset_j = spzeros(Int,Ny)
    offset_j[2] = -1
    offset_j[3] = -2
    offset_j[Ny-1] = 1
    offset_j[Ny-2] = 2
    coef_j = spzeros(Ny)
    coef_j[1] = 2* 1.5
    coef_j[2] = - 2.0
    coef_j[3] = 0.5
    coef_j[Ny] = 2*1.5
    coef_j[Ny-1] = -2.0
    coef_j[Ny-2] = 0.5
    for i = 1:Nx
        for j = 1:Ny
            global_index = (i-1)*Ny+j

            # idata_index_x = (i-1)*Ny + j + test_mod(i,Nx) * Nx
            # idata_index_y = (i-1)*Ny + j + test_mod(j,Ny)

            offset_x = div(2*Nx-i-3,Nx-2) - 1
            offset_y = div(2*Ny-j-3,Ny-2) - 1

            idata_index_x = (i-1)*Ny + j + (offset_x) * Nx
            idata_index_y = (i-1)*Ny + j + offset_y

            # odata[global_index] = 0

            # odata[global_index] = (( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])  + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) ) / (h^2)
            # + abs(offset_y) * 2 * (alpha1 * (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y]) / h^2
            # + abs(offset_x) * alpha4 * 2 * (idata[global_index] / h)  + coef_j[i] * beta * (idata[global_index+offset_j[i] * Ny]) / h^2)

            # Incorporating H_tilde
            odata[global_index] = (( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])  + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) )
            + abs(offset_y) * 2 * alpha1 *( (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y])
            + abs(offset_x) * alpha4 * 2 * (idata[global_index] * h)  + coef_j[i] * beta * (idata[global_index+offset_j[i] * Ny]) ) / 2^(abs(offset_x) + abs(offset_y))

            #####
            # form boundary conditions separately

            # odata[global_index] += abs(offset_y) * 2 * (alpha1 * (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y]) / h^2 # BS_y
            # odata[global_index] += abs(offset_x) * alpha4 * 2 * (idata[global_index] / h) #
            # odata[global_index] += coef_j[i] * beta * (idata[global_index+offset_j[i] * Ny]) / h^2

            # # odata[global_index] += abs(offset_x) * ( beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha4 * 2* (idata[global_index]) / h) # BS_x'
            # # odata[global_index+Ny*offset_x] += abs(offset_x) * beta * 2 * (-1*idata[global_index]) / h^2
            # # odata[global_index+2*Ny*offset_x] += abs(offset_x) * beta * (0.5*idata[global_index]) / h^2
        end
    end

end

function matrix_free_A_v5(idata,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    # one function that does D2 + boundary conditions + H_tilde
    # alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    odata .= 0
    offset_j = spzeros(Int,Ny)
    offset_j[2] = -1
    offset_j[3] = -2
    offset_j[Ny-1] = 1
    offset_j[Ny-2] = 2
    coef_j = spzeros(Ny)
    coef_j[1] = 2* 1.5
    coef_j[2] = - 2.0
    coef_j[3] = 0.5
    coef_j[Ny] = 2*1.5
    coef_j[Ny-1] = -2.0
    coef_j[Ny-2] = 0.5
    for i = 1:Nx
        for j = 1:Ny
            global_index = (i-1)*Ny+j

            offset_x = div(2*Nx-i-3,Nx-2) - 1
            offset_y = div(2*Ny-j-3,Ny-2) - 1

            idata_index_x = (i-1)*Ny + j + (offset_x) * Nx
            idata_index_y = (i-1)*Ny + j + offset_y

            odata[global_index] = (( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])  + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) )
            + abs(offset_y) * 2 * alpha1 * ( (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y])
            + abs(offset_x) * alpha4 * 2 * (idata[global_index] * h)
             + coef_j[i] * beta * (idata[global_index+offset_j[i] * Ny])
             ) / 2^(abs(offset_x) + abs(offset_y))

        end
    end
end

function matrix_free_GPU(idata,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta,offset_j,coef_j,::Val{TILE_DIM1}, ::Val{TILE_DIM2}) where {TILE_DIM1, TILE_DIM2}
    tidx = threadIdx().x
    tidy = threadIdx().y

    # for global memory indexing
    i = (blockIdx().x - 1) * TILE_DIM1 + tidx
    j = (blockIdx().y - 1) * TILE_DIM2 + tidy

    global_index = (i-1)*Ny+j

    offset_x = div(2*Nx-i-3,Nx-2) - 1
    offset_y = div(2*Ny-j-3,Ny-2) - 1
    idata_index_x = (i-1)*Ny + j + (offset_x) * Nx
    idata_index_y = (i-1)*Ny + j + offset_y

    # tidx = (blockIdx().x - 1) * TILE_DIM + threadIdx().x
    # tidy = (blockIdx().y - 1) * TILE_DIM + threadIdx().y

    # global_index = (tidx - 1) * Ny + tidy

    alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    # odata .= 0
    # offset_j = spzeros(Int,Ny)
    # offset_j = zeros(Int,Ny)
    # offset_j[2] = -1
    # offset_j[3] = -2
    # offset_j[Ny-1] = 1
    # offset_j[Ny-2] = 2
    # # coef_j = spzeros(Ny)
    # coef_j = zeros(Ny)
    # coef_j[1] = 2* 1.5
    # coef_j[2] = - 2.0
    # coef_j[3] = 0.5
    # coef_j[Ny] = 2*1.5
    # coef_j[Ny-1] = -2.0
    # coef_j[Ny-2] = 0.5

    if 1 <= i <= Nx && 1 <= j <= Ny

        odata[global_index] = (( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])  + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) )
        + abs(offset_y) * 2 * alpha1 * ( (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y])
        + abs(offset_x) * alpha4 * 2 * (idata[global_index] * h)  + coef_j[i] * beta * (idata[global_index+offset_j[i] * Ny]) ) / 2^(abs(offset_x) + abs(offset_y))
    end
    nothing
end

function matrix_free_GPU_debug(idata,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta,offset_j,coef_j,::Val{TILE_DIM1}, ::Val{TILE_DIM2}) where {TILE_DIM1, TILE_DIM2}
    tidx = threadIdx().x
    tidy = threadIdx().y

    # for global memory indexing
    i = (blockIdx().x - 1) * TILE_DIM1 + tidx
    j = (blockIdx().y - 1) * TILE_DIM2 + tidy

    global_index = (i-1)*Ny+j

    offset_x = div(2*Nx-i-3,Nx-2) - 1
    offset_y = div(2*Ny-j-3,Ny-2) - 1
    idata_index_x = (i-1)*Ny + j + (offset_x) * Nx
    idata_index_y = (i-1)*Ny + j + offset_y

    # tidx = (blockIdx().x - 1) * TILE_DIM + threadIdx().x
    # tidy = (blockIdx().y - 1) * TILE_DIM + threadIdx().y

    # global_index = (tidx - 1) * Ny + tidy

    alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    # odata .= 0
    # offset_j = spzeros(Int,Ny)
    # offset_j = zeros(Int,Ny)
    # offset_j[2] = -1
    # offset_j[3] = -2
    # offset_j[Ny-1] = 1
    # offset_j[Ny-2] = 2
    # # coef_j = spzeros(Ny)
    # coef_j = zeros(Ny)
    # coef_j[1] = 2* 1.5
    # coef_j[2] = - 2.0
    # coef_j[3] = 0.5
    # coef_j[Ny] = 2*1.5
    # coef_j[Ny-1] = -2.0
    # coef_j[Ny-2] = 0.5

    if 1 <= i <= Nx && 1 <= j <= Ny
    # if global_index <= Nx*Ny

        odata[global_index] = 0
        odata[global_index] = (( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])  + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) )
        + abs(offset_y) * 2 * alpha1 * ( (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y])
        + abs(offset_x) * alpha4 * 2 * (idata[global_index] * h)
         + coef_j[i] * beta * idata[global_index+offset_j[i]* Ny]  # * (idata[global_index+offset_j[i] * Ny])   # this line is causing issue
         ) / 2^(abs(offset_x) + abs(offset_y))
    end
    nothing
end

function matrix_free_GPU_debug_v2(idata,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta,offset_j,coef_j,::Val{TILE_DIM1}, ::Val{TILE_DIM2}) where {TILE_DIM1, TILE_DIM2}
    tidx = threadIdx().x
    tidy = threadIdx().y

    # for global memory indexing
    i = (blockIdx().x - 1) * TILE_DIM1 + tidx
    j = (blockIdx().y - 1) * TILE_DIM2 + tidy

    global_index = (i-1)*Ny+j

    # offset_x = div(2*Nx-i-3,Nx-2) - 1
    # offset_y = div(2*Ny-j-3,Ny-2) - 1
    # idata_index_x = (i-1)*Ny + j + (offset_x) * Nx
    # idata_index_y = (i-1)*Ny + j + offset_y

    # alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    # odata .= 0
    # offset_j = spzeros(Int,Ny)
    # offset_j = zeros(Int,Ny)
    # offset_j[2] = -1
    # offset_j[3] = -2
    # offset_j[Ny-1] = 1
    # offset_j[Ny-2] = 2
    # # coef_j = spzeros(Ny)
    # coef_j = zeros(Ny)
    # coef_j[1] = 2* 1.5
    # coef_j[2] = - 2.0
    # coef_j[3] = 0.5
    # coef_j[Ny] = 2*1.5
    # coef_j[Ny-1] = -2.0
    # coef_j[Ny-2] = 0.5

    if 2 <= i <= Nx-1 && 2 <= j <= Ny-1
    # if global_index <= Nx*Ny

        odata[global_index] = 0
        odata[global_index] = (( (idata[global_index-Ny] - 2*idata[global_index] + idata[global_index + Ny])  + (idata[global_index-1] - 2*idata[global_index] + idata[global_index + 1]) )
        # + abs(offset_y) * 2 * (alpha1 * (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y])
        # + abs(offset_x) * alpha4 * 2 * (idata[global_index] * h)
        #  + coef_j[i] * beta * idata[global_index+offset_j[i]* Ny]  # * (idata[global_index+offset_j[i] * Ny])   # this line is causing issue
        ) #  ) / 2^(abs(offset_x) + abs(offset_y))
    end
    nothing
end

function Boundary_Conditions(idata,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    odata .= 0
    for i = 1:Nx
        for j = 1:Ny
            global_index = (i-1)*Ny + j
            if j == Ny # N
                odata[global_index] += alpha1 * 2 * (1.5*idata[global_index] - 2*idata[global_index-1] + 0.5*idata[global_index-2]) / h^2
                # odata[global_index] = alpha1 * 2 * (1.5*idata[global_index] - 2*idata[global_index-1] + 0.5*idata[global_index-2]) / h^2
            end

            if j == 1 # S
                odata[global_index] += alpha2 * 2 *(1.5*idata[global_index] - 2*idata[global_index+1] + 0.5*idata[global_index+2]) / h^2
            end

            if i == Nx # E
                odata[global_index] += beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha4 * 2* (idata[global_index]) / h
                odata[global_index-Ny] += beta * 2 * (-1*idata[global_index]) / h^2
                odata[global_index-2*Ny] += beta * (0.5*idata[global_index]) / h^2
            end

            if i == 1 # W
                odata[global_index] += beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha3 * 2* (idata[global_index]) / h
                odata[global_index+Ny] += beta * 2 * (-1*idata[global_index]) / h^2
                odata[global_index+2*Ny] += beta * (0.5*idata[global_index]) / h^2
            end
        end
    end
end

function Boundary_Conditions_v2(idata,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    odata .= 0
    for i = 1:Nx
        for j = 1:Ny
            global_index = (i-1)*Ny + j
            offset_x = div(2*Nx-i-3,Nx-2) - 1
            offset_y = div(2*Ny-j-3,Ny-2) - 1
            # if j == Ny # N
            #     odata[global_index] += alpha1 * 2 * (1.5*idata[global_index] - 2*idata[global_index-1] + 0.5*idata[global_index-2]) / (2^(abs(offset_x)+abs(offset_y)))
            #     # odata[global_index] = alpha1 * 2 * (1.5*idata[global_index] - 2*idata[global_index-1] + 0.5*idata[global_index-2]) / h^2
            # end

            # if j == 1 # S
            #     odata[global_index] += alpha2 * 2 *(1.5*idata[global_index] - 2*idata[global_index+1] + 0.5*idata[global_index+2]) / h^2
            # end
            odata[global_index] += abs(offset_y) * 2 * alpha1 * ( (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y]) / h^2# + abs(offset_x) * ( beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha4 * 2* (idata[global_index]) / h)
                                # +  abs(offset_x) * ( beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha4 * 2* (idata[global_index]) / h)

            # if i == Nx # E
            #     odata[global_index] += beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha4 * 2* (idata[global_index]) / h
            #     odata[global_index-Ny] += beta * 2 * (-1*idata[global_index]) / h^2
            #     odata[global_index-2*Ny] += beta * (0.5*idata[global_index]) / h^2
            # end

            # if i == 1 # W
            #     odata[global_index] += beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha3 * 2* (idata[global_index]) / h
            #     odata[global_index+Ny] += beta * 2 * (-1*idata[global_index]) / h^2
            #     odata[global_index+2*Ny] += beta * (0.5*idata[global_index]) / h^2
            # end
            odata[global_index] += abs(offset_x) * ( beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha4 * 2* (idata[global_index]) / h)
            odata[global_index+Ny*offset_x] += abs(offset_x) * beta * 2 * (-1*idata[global_index]) / h^2
            odata[global_index+2*Ny*offset_x] += abs(offset_x) * beta * (0.5*idata[global_index]) / h^2
        end
    end
end

function H_tilde_matrix_free(idata,odata,Nx,Ny,h)
    for i in 1:Nx
        for j in 1:Ny
            global_index = (i-1)*Ny+j
            odata[global_index] = idata[global_index] * h^2
            if i == 1
                odata[global_index] /= 2
            end
            if j == 1
                odata[global_index] /= 2
            end
            if i == Nx
                odata[global_index] /= 2
            end
            if j == Ny
                odata[global_index] /= 2
            end
        end
    end
end

function H_tilde_matrix_free_v2(idata,odata,Nx,Ny,h)
    for i in 1:Nx
        for j in 1:Ny
            global_index = (i-1)*Ny + j
            var = abs(div(2*Nx-i-3,Nx-2) - 1) + abs(div(2*Ny-j-3,Ny-2)-1)
            odata[global_index] = idata[global_index] * h^2 / (2^var)
            # odata[global_index] = var
        end
    end
end

function test_H_tilde(Nx,Ny)
    h = 1/(Nx-1)
    # idata = randn(Nx*Ny)
    idata = ones(Nx*Ny)
    odata1 = spzeros(Nx*Ny)
    odata2 = spzeros(Nx*Ny)
    H_tilde_matrix_free(idata,odata1,Nx,Ny,h)
    H_tilde_matrix_free_v2(idata,odata2,Nx,Ny,h)
    # @show odata1
    # @show odata2
    @show odata1 - odata2
end

function test_D2(Nx,Ny)
    h = 1/(Nx-1)
    # idata = randn(Nx*Ny)
    idata = ones(Nx*Ny)
    odata1 = spzeros(Nx*Ny)
    odata2 = spzeros(Nx*Ny)
    odata3 = spzeros(Nx*Ny)
    D2_test(idata,odata1,Nx,Ny,h)
    H_tilde_matrix_free(odata1,odata2,Nx,Ny,h)
    D2_test_v2(idata,odata3,Nx,Ny,h)
    # @show odata1
    # @show odata2
    @show odata2 - odata3
end

function test_Boundary_Conditions(Nx,Ny)
    h = 1/(Nx-1)
    alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    # idata = randn(Nx*Ny)
    idata = randn(Nx*Ny)
    odata1 = spzeros(Nx*Ny)
    odata2 = spzeros(Nx*Ny)
    Boundary_Conditions(idata,odata1,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    Boundary_Conditions_v2(idata,odata2,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    # @show odata1
    # @show odata2
    @show odata1 - odata2
end

function matrix_free_A(idata,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    odata1 = spzeros(Nx*Ny)
    odata2 = spzeros(Nx*Ny)
    D2_test(idata,odata1,Nx,Ny,h)
    Boundary_Conditions(idata,odata2,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    # odata3 = odata1 .+ odata2
    # return odata
    odata .= odata1 .+ odata2
    H_tilde_matrix_free(odata,odata,Nx,Ny,h)
    # return odata3
end

function test_matrix_free(Nx,Ny)
    h = 1/(Nx-1)
    alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    alpha1 = alpha2 = -13/h
    idata = randn(Nx*Ny)
    odata1 = spzeros(Nx*Ny)
    odata2 = spzeros(Nx*Ny)
    odata3 = spzeros(Nx*Ny)
    odata4 = spzeros(Nx*Ny)
    odata5 = spzeros(Nx*Ny)
    odata6 = spzeros(Nx*Ny)
    D2_test(idata,odata1,Nx,Ny,h)
    Boundary_Conditions(idata,odata2,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    # D2_test_v2(idata,odata3,Nx,Ny,h)
    Boundary_Conditions_v2(idata,odata4,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    matrix_free_A(idata,odata3,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    matrix_free_A_v2(idata,odata5,Nx,Ny,h) # without incorporating H_tilde
    matrix_free_A_v4(idata,odata6,Nx,Ny,h) # including H_tilde
    @show reshape(odata1 + odata2 - odata5,Nx,Ny)
    @show reshape(odata3 - odata6,Nx,Ny)
    @assert odata1 + odata2 ≈ odata5
    @assert odata3 ≈ odata6
end

function test_GPU_kernel(level)
    alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    Nx = Ny = 2^level + 1
    h = 1/(Nx-1)
    Random.seed!(0)
    u = randn(Nx*Ny)
    odata = randn(Nx*Ny)
    d_u = CuArray(u)
    d_odata = CuArray(odata)
    TILE_DIM_1 = 16
    TILE_DIM_2 = 16
    griddim = (div(Nx,TILE_DIM_1) + 1, div(Ny,TILE_DIM_2) + 1)
	blockdim = (TILE_DIM_1,TILE_DIM_2)

    offset_j = CuArray(zeros(Int,Ny))
    # offset_j = zeros(Int,Ny)
    offset_j[2] = -1
    offset_j[3] = -2
    offset_j[Ny-1] = 1
    offset_j[Ny-2] = 2
    # coef_j = spzeros(Ny)
    coef_j = CuArray(zeros(Float64,Ny))
    coef_j[1] = 2* 1.5
    coef_j[2] = - 2.0
    coef_j[3] = 0.5
    coef_j[Ny] = 2*1.5
    coef_j[Ny-1] = -2.0
    coef_j[Ny-2] = 0.5
    matrix_free_A_v5(u,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    @cuda threads=blockdim blocks=griddim matrix_free_GPU_debug(d_u,d_odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta,offset_j,coef_j,Val(TILE_DIM_1), Val(TILE_DIM_2))
    # @show reshape(odata,Nx,Ny)
    # @show reshape(d_odata,Nx,Ny)

    @assert odata ≈ Array(d_odata)

    iter_times = 100

    t0 = time()
    for i in 1:iter_times
        matrix_free_A_v5(u,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    end
    t_cpu = time() - t0

    t1 = time()
    for i in 1:iter_times
        @cuda threads=blockdim blocks=griddim matrix_free_GPU_debug(d_u,d_odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta,offset_j,coef_j,Val(TILE_DIM_1), Val(TILE_DIM_2))
    end
    synchronize()
    t_gpu = time() - t1

    t_D2 = time()
    for i in 1:iter_times
        @cuda threads=blockdim blocks=griddim matrix_free_GPU_debug_v2(d_u,d_odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta,offset_j,coef_j,Val(TILE_DIM_1), Val(TILE_DIM_2))
    end
    synchronize()
    t_D2 = time() - t_D2

    @show t_cpu
    @show t_gpu
    @show t_D2
end

function test_sparse_matrix(A,b)
    iter_times = 1000
    t3 = time()
    for i in 1:iter_times
        out = A*b
    end
    cpu_time = time() - t3
    @show cpu_time
end


text:
module Models

export Model, ChannelModel, NonDimensionalModel, Clock

using Oceananigans.Architectures
using Oceananigans.Fields
using Oceananigans.Coriolis
using Oceananigans.Buoyancy
using Oceananigans.TurbulenceClosures
using Oceananigans.BoundaryConditions
using Oceananigans.Solvers
using Oceananigans.Forcing
using Oceananigans.Utils

"""
    AbstractModel

Abstract supertype for models.
"""
abstract type AbstractModel end

include("clock.jl")
include("model.jl")
include("channel_model.jl")
include("non_dimensional_model.jl")
include("show_models.jl")

end


text:
function load_lattice_ase(T, pyobj::PyObject)
    ase = pyimport("ase")

    if pyisinstance(pyobj, ase.Atoms)
        if !all(pyobj.pbc)
            error("DFTK only supports calculations with periodic boundary conditions.")
        end
        load_lattice_ase(T, pyobj.cell)
    elseif pyisinstance(pyobj, ase.cell.Cell)
        lattice = zeros(3, 3)
        cell_julia = convert(Array, pyobj)  # Array of arrays
        for i = 1:3, j = 1:3
            lattice[i, j] = austrip(cell_julia[j][i] * u"Å")
        end
        Mat3{T}(lattice)
    else
        error("load_lattice_ase not implemented for python type $pyobj")
    end
end

function load_atoms_ase(T, pyobj::PyObject)
    @assert pyisinstance(pyobj, pyimport("ase").Atoms)
    # TODO Be smarter and look at the calculator to determine the psps

    frac_coords = pyobj.get_scaled_positions()
    map(ase_atoms_translation_map(pyobj)) do (symbol, atom_indices)
        coords = Vec3{T}[]
        for idx in atom_indices
            push!(coords, frac_coords[idx, :])
        end
        ElementCoulomb(symbol) => coords
    end
end

function load_magnetic_moments_ase(pyobj::PyObject)
    @assert pyisinstance(pyobj, pyimport("ase").Atoms)
    magmoms = pyobj.get_initial_magnetic_moments()
    map(ase_atoms_translation_map(pyobj)) do (symbol, atom_indices)
        ElementCoulomb(symbol) => [normalize_magnetic_moment(magmoms[i])
                                   for i in atom_indices]
    end
end

"""
Map translating between the `atoms` datastructure used in DFTK
and the indices of the respective atoms in the ASE.Atoms object
in `pyobj`.
"""
function ase_atoms_translation_map(pyobj::PyObject)
    @assert pyisinstance(pyobj, pyimport("ase").Atoms)
    numbers = pyobj.get_atomic_numbers()
    map(unique(numbers)) do atnum
        ElementCoulomb(atnum).symbol => findall(isequal(atnum), numbers)
    end
end

#                                                             convert A.U. -> Å
ase_cell(lattice) = pyimport("ase").cell.Cell(Array(lattice)' / austrip(1u"Å"))
ase_cell(model::Model) = ase_cell(model.lattice)

function ase_atoms(lattice_or_model, atoms, magnetic_moments=[])
    cell = ase_cell(lattice_or_model)
    symbols = String[]
    for (elem, pos) in atoms
        append!(symbols, fill(string(atomic_symbol(elem)), length(pos)))
    end
    scaled_positions = vcat([pos for (elem, pos) in atoms]...)

    magmoms = nothing
    if !isempty(magnetic_moments)
        @assert length(magnetic_moments) == length(atoms)
        get_z_magmom(m) = normalize_magnetic_moment(m)[3]
        magmoms = vcat([get_z_magmom.(magmom)
                        for (elem, magmom) in magnetic_moments]...)
        @assert length(magmoms) == length(scaled_positions)
    end

    pyimport("ase").Atoms(symbols=symbols, cell=cell, pbc=true,
                          scaled_positions=hcat(scaled_positions...)',
                          magmoms=magmoms)
end
function ase_atoms(model::Model, magnetic_moments=nothing)
    ase_atoms(model.lattice, model.atoms, magnetic_moments)
end


text:
function WeightedObservedPoints()
    return WeightedObservedPoints(())
end

function add(obj::WeightedObservedPoints, arg0::WeightedObservedPoint)
    return jcall(obj, "add", void, (WeightedObservedPoint,), arg0)
end

function add(obj::WeightedObservedPoints, arg0::jdouble, arg1::jdouble)
    return jcall(obj, "add", void, (jdouble, jdouble), arg0, arg1)
end

function add(obj::WeightedObservedPoints, arg0::jdouble, arg1::jdouble, arg2::jdouble)
    return jcall(obj, "add", void, (jdouble, jdouble, jdouble), arg0, arg1, arg2)
end

function clear(obj::WeightedObservedPoints)
    return jcall(obj, "clear", void, ())
end

function to_list(obj::WeightedObservedPoints)
    return jcall(obj, "toList", List, ())
end


text:
@enum Cell Dead Alive

function initBoard(strarr, h, w)
    board = [[Dead for _ in 1:w] for _ in 1:h]
    for i in 1:length(strarr) for j in 1:length(strarr[i])
            (strarr[i][j] == '#') && (board[i][j] = Alive)
    end end
    return board
end

function showBoard(board, aliveC, deadC)
    for row in board
        for cell in row print( ifelse(cell==Alive, aliveC, deadC)) end
        println()
    end
end

function countNeighbors(board, row, col)
    count  = ifelse(board[row][col] == Alive, -1, 0)
    height = length(board)
    width  = length(board[1])
    for dr in -1:1 for dc in -1:1
        r = 1 + (row + dr + height - 1) % height
        c = 1 + (col + dc + width - 1) % width
        (board[r][c] == Alive) && (count += 1)
    end end
    return count
end

function nextStage(board)
    height = length(board)
    width  = length(board[1])
    next   = initBoard([[]], height, width)
    for i in 1:height for j in 1:width
        n = countNeighbors(board, i, j)
        ( (board[i][j]==Alive && 2 <= n <= 3) || n == 3 ) && ( next[i][j] = Alive )
    end end
    return next
end

inputFile = open("input.txt", "r")

height, width = [parse(Int,x) for x in split(readline(inputFile), " ")]
initstr = []
while ! eof(inputFile) push!(initstr, readline(inputFile)) end
close(inputFile)

world = initBoard(initstr, height, width)
showBoard(world, '#', '.')
while true
    global world
    world = nextStage(world)
    print("\033[",height,"A")
    showBoard(world, '#', '.')
    sleep(0.25)
end

max_stars_repo_name: ray-pH/gameo-life

max_issues_repo_name: ray-pH/gameo-life

max_forks_repo_name: ray-pH/gameo-life


text: Zygote.@nograd Bijectors._debug


text:
module MPT

# Write your package code here.
# Mabye we should not have to build in plots?
# The user perhaps can plot themselves.
using Statistics, XLSX, DataFrames, JuMP, Dates
import Ipopt

export
    equal_weights,
    walk_forward,
    long_only_minimum_variance,
    random_weights,
    mdp

include("import.jl")
include("portfolios.jl")
include("tools.jl")
include("walkforward.jl")

end


text:
"""
Voxel corner and edge indexing conventions

        Z
        |

        5------5------6              Extra edges not drawn
       /|            /|              -----------
      8 |           6 |              - face diagonals
     /  9          /  10                - 13: 1 to 3
    8------7------7   |                 - 14: 1 to 8
    |   |         |   |                 - 15: 1 to 6
    |   1------1--|---2  -- Y           - 16: 5 to 7
    12 /          11 /                  - 17: 2 to 7
    | 4           | 2                   - 18: 4 to 7
    |/            |/                 - body diagonal
    4------3------3                     - 19: 1 to 7

  /
 X
"""
struct VoxelIndices{T <: Integer}
    voxCrnrPos::NTuple{8,NTuple{3,T}}
    voxEdgeCrnrs::NTuple{19, NTuple{2,T}}
    voxEdgeDir::NTuple{19,T}
    voxEdgeIx::NTuple{8,NTuple{8,T}}
    subTets::NTuple{6,NTuple{4,T}}
    tetEdgeCrnrs::NTuple{6,NTuple{2,T}}
    tetTri::NTuple{16,NTuple{6,T}}

    function VoxelIndices{T}() where {T <: Integer}
        voxCrnrPos = ((0, 0, 0),
                      (0, 1, 0),
                      (1, 1, 0),
                      (1, 0, 0),
                      (0, 0, 1),
                      (0, 1, 1),
                      (1, 1, 1),
                      (1, 0, 1))
        # the voxel IDs at either end of the tetrahedra edges, by edge ID
        voxEdgeCrnrs = ((1, 2),
                        (2, 3),
                        (4, 3),
                        (1, 4),
                        (5, 6),
                        (6, 7),
                        (8, 7),
                        (5, 8),
                        (1, 5),
                        (2, 6),
                        (3, 7),
                        (4, 8),
                        (1, 3),
                        (1, 8),
                        (1, 6),
                        (5, 7),
                        (2, 7),
                        (4, 7),
                        (1, 7))

        # direction codes:
        # 0 => +x, 1 => +y, 2 => +z,
        # 3 => +xy, 4 => +xz, 5 => +yz, 6 => +xyz
        voxEdgeDir = (1,0,1,0,1,0,1,0,2,2,2,2,3,4,5,3,4,5,6)

        # For a pair of corner IDs, the edge ID joining them
        # 0 denotes a pair with no edge
        voxEdgeIx = ((0,1,13,4,9,15,19,14),
                     (1,0,2,0,0,10,17,0),
                     (13,2,0,3,0,0,11,0),
                     (4,0,3,0,0,0,18,12),
                     (9,0,0,0,0,5,16,8),
                     (15,10,0,0,5,0,6,0),
                     (19,17,11,18,16,6,0,7),
                     (14,0,0,12,8,0,7,0))

        # voxel corners that comprise each of the six tetrahedra
        subTets = ((1,3,2,7),
                      (1,8,4,7),
                      (1,4,3,7),
                      (1,2,6,7),
                      (1,5,8,7),
                      (1,6,5,7))
        # tetrahedron corners for each edge (indices 1-4)
        tetEdgeCrnrs = ((1,2),
                        (2,3),
                        (1,3),
                        (1,4),
                        (2,4),
                        (3,4))

        # triangle cases for a given tetrahedron edge code
        tetTri = ((0,0,0,0,0,0),
                    (1,3,4,0,0,0),
                    (1,5,2,0,0,0),
                    (3,5,2,3,4,5),
                    (2,6,3,0,0,0),
                    (1,6,4,1,2,6),
                    (1,5,6,1,6,3),
                    (4,5,6,0,0,0),
                    (4,6,5,0,0,0),
                    (1,6,5,1,3,6),
                    (1,4,6,1,6,2),
                    (2,3,6,0,0,0),
                    (3,2,5,3,5,4),
                    (1,2,5,0,0,0),
                    (1,4,3,0,0,0),
                    (0,0,0,0,0,0))

        new{T}(voxCrnrPos,
            voxEdgeCrnrs,
            voxEdgeDir,
            voxEdgeIx,
            subTets,
            tetEdgeCrnrs,
            tetTri)
    end
end


text:
using DataFrames

#Arrays and Dataframe example

dfGenPath = "C://Users//rapiduser//github//GRACE-ARPAE//examples_Julia//UC_DukeEnergy_Sample//inputs//data_generators.csv"
dfGenerator = CSV.read(dfGenPath, DataFrame)

myArr = [0.0, 0.0, 0.0, -0.0, -0.0, -0.0, 0.0, 0.0, 0.0, 1.0, 1.0, -0.0, -0.0, 1.0]
arrSalary = [3000.50,450_000.70,80_000.1,99_000.9,77_777.85]

m = Array{Bool}(undef, 14)

df = DataFrame(Name = ["Jon","Bill","Maria","Julia","Mark"],
               Age = [22,43,81,52,27],
               Salary = [3000.50,45000.70,60000.1,50000.9,55000.85]
               )

# These commands work
#df = DataFrame(Age = round.(Int, rand(5)), Salary = arrSalary )
#df.Salary = round.(arrSalary; digits =3)

df.Salary = arrSalary
#replace!(Salary,)

m = myArr
println("ArraySize")
println(size(m))
println(size(myArr))
println(df)

# MODIFYING DATAFRAME
println("dfGenerator.Uini")
println(dfGenerator.Uini)

arrayUini = [1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1]
dfGenerator.Uini = arrayUini

println("dfGenerator.Uini (MODIFIED)")
println(dfGenerator.Uini)
#@. df.Salary = arrSalary

max_stars_repo_name: power-systems-op/grace-arpae

max_issues_repo_name: power-systems-op/grace-arpae

max_forks_repo_name: power-systems-op/grace-arpae


text:
using MULTEM
using TOML
using Test

@testset "MULTEM.jl" begin
    test_defaults = from_dict(MULTEMInput, TOML.parsefile("default.toml"))
    isfile("default_generated.toml") ? rm("default_generated.toml") : nothing
    to_toml("default_generated.toml", test_defaults; include_defaults=true)
    @test from_dict(MULTEMInput, TOML.parsefile("default_generated.toml")) == from_dict(MULTEMInput, TOML.parsefile("default.toml"))
    @test TOML.parsefile("default_generated.toml") == TOML.parsefile("default.toml")
end


text:
using Documenter, StackOverflow

makedocs(sitename="StackOverflow.jl")

deploydocs(
    repo = "github.com/logankilpatrick/StackOverflow.jl.git",
)


text:
function version()
    ver = Ref{Cint}()
    ncclGetVersion(ver)
    major, ver = divrem(ver[], 1000)
    minor, patch = divrem(ver, 100)

    VersionNumber(major, minor, patch)
end

function ncclDataType(T::DataType)
    if T == Float32
        return ncclFloat32
    elseif T == Float16
        return ncclFloat16
    elseif T == Float64
        return ncclFloat64
    elseif T == Int8
        return ncclInt8
    elseif T == Char
        return ncclInt8
    elseif T == Int32
        return ncclInt32
    elseif T == UInt32
        return ncclUint32
    elseif T == Int64
        return ncclInt64
    elseif T == UInt64
        return ncclUint64
    else
        throw(ArgumentError("ncclDataType equivalent for input type $T does not exist!"))
    end
end


text:
using AutomotiveDrivingModels
using POMDPs
using POMDPSimulators
using POMDPPolicies
using BeliefUpdaters
using RLInterface
using Random

@testset  "Urban POMDP" begin
    rng = MersenneTwister(1)

    params = UrbanParams(nlanes_main=1,
                     crosswalk_pos =[VecSE2(6, 0., pi/2), VecSE2(-6, 0., pi/2), VecSE2(0., -5., 0.)],
                     crosswalk_length =  [14.0, 14., 14.0],
                     crosswalk_width = [4.0, 4.0, 3.1],
                     stop_line = 22.0)
    env = UrbanEnv(params=params)
    pomdp = UrbanPOMDP(env=env,
                   ego_goal = LaneTag(2, 1),
                   max_cars=5,
                   max_peds=5,
                   car_birth=0.5,
                   ped_birth=0.2,
                   max_obstacles=0, # no fixed obstacles
                   ego_start=20,
                   ΔT=0.1);

    up = KMarkovUpdater(4)

    policy = RandomPolicy(pomdp, rng=rng)

    hr = HistoryRecorder(rng=rng, max_steps = 100)
    s0 = rand(rng, initialstate(pomdp))
    initial_observation = rand(rng, initialobs(pomdp, s0))
    initial_obs_vec = fill(initial_observation, 4)
    hist = POMDPs.simulate(hr, pomdp, policy, up, initial_obs_vec, s0)

    @test n_steps(hist) > 1
end


text:
using Pkg;
Pkg.activate("././.")
using Revise
using ITensors

## FUNCTIONS ##

struct Parameters
    # Model parameters
    N::Int
    t::Real
    U::Real
    ω::Real
    g0::Real
    g1::Real
    doping::Real

    # DMRG parameters
    sweeps::Sweeps

    # TEBD parameters
    mid::Int
    T::Int
    τ::Real
    TEBD_cutoff
    TEBD_maxdim
end

struct HubbardHolsteinModel
    # Lattice
    sites
    # Hamiltonian (in two forms)
    mpo::MPO
    gates
    ampo
end

struct DMRGResults
    ground_state
    ground_state_energy
    entropy
end

struct TEBDResults
    entropy
    self_overlap
    phonon_flux
    corrs
end

function parameters(;N::Int, t::Real, U::Real=nothing, ω::Real=nothing,
                    g0::Real=nothing, g1::Real=nothing, doping::Real=0,
                    DMRG_numsweeps::Int=20, DMRG_noise=nothing, DMRG_maxdim=nothing, DMRG_cutoff=nothing,
                    T::Int=25, τ::Real=0.1, TEBD_cutoff=1E-14, TEBD_maxdim=400)

    if isnothing(U)
        U = 8*t
    end
    if isnothing(ω)
        ω = 0.5*t
    end
    if isnothing(g0)
        g0 = 0.05*t
    end
    if isnothing(g1)
        g1 = 0*g0
    end

    sweeps = Sweeps(DMRG_numsweeps)
    if isnothing(DMRG_noise)
        DMRG_noise = [1E-6,1E-6,1E-8,0]
    end
    if isnothing(DMRG_maxdim)
        DMRG_maxdim = [20,40,100,200,400]
    end

    if isnothing(DMRG_cutoff)
        DMRG_cutoff = 1E-14
    end
    setnoise!(sweeps, DMRG_noise...) # Very important to use noise for this model
    setmaxdim!(sweeps, DMRG_maxdim...)
    setcutoff!(sweeps, DMRG_cutoff...)

    Parameters(N,t,U,ω,g0,g1,doping,sweeps,ceil(Int,N/2),T,τ,TEBD_cutoff,TEBD_maxdim)
end

function HubbardHolsteinModel(p::Parameters)
    N, t, U, ω, g0, g1, τ = p.N, p.t, p.U, p.ω, p.g0, p.g1, p.τ

    # make the sites
    sites = siteinds("HubHolst", N)

    # make the hamiltonian
    ampo = Sum{Op}()
    for j=1:N-1
        # ∑_j,σ t * (c^†_jσ c_{j+1}σ + h.c.)
        ampo += -t,"Cdagup",j,"Cup",j+1
        ampo += -t,"Cdagup",j+1,"Cup",j
        ampo += -t,"Cdagdn",j,"Cdn",j+1
        ampo += -t,"Cdagdn",j+1,"Cdn",j

        # ∑_j U * n_j↑ n_j↓
        ampo += U,"Nupdn",j,"I",j

        # ∑_j ω * nb_j
        ampo += ω,"Nb",j

        # # ∑_j g0 * nf_j (b^†_j + b_j)
        ampo += g0,"Ntot(Bd+B)",j

        # # ∑_⟨ij⟩ g1 * nf_j (b^†_i + b_i)
        ampo += g1,"Ntot",j,"Bdag+B",j+1
        ampo += g1,"Ntot",j+1,"Bdag+B",j
    end
    # Edge site
    ampo += U,"Nupdn",N
    ampo += ω,"Nb",N
    ampo += g0,"Ntot(Bd+B)",N

    H = MPO(ampo,sites)

    # make the trotter gates
    gates = ITensor[]
    hgates = ITensor[]
    for j=1:N-1
        s1 = sites[j] # site j
        s2 = sites[j+1] # site j+1

        hj = -t*(op("Cdagup*F",s1) * op("Cup",s2)  # t * (c^†_jσ c_{j+1}σ + h.c.)
                 -op("Cup*F",s1) * op("Cdagup",s2)
                 +op("Cdagdn*F",s1) * op("Cdn",s2)
                 -op("Cdn*F",s1) * op("Cdagdn",s2))

            + U*(op("Nupdn",s1) * op("I",s2))

            + ω*(op("Nb",s1) * op("I",s2))

            + g0*(op("Ntot(Bd+B)",s1) * op("I",s2))

            + g1*(op("Ntot",s1) * op("Bdag+B",s2))
            + g1*(op("Bdag+B",s1) * op("Ntot",s2))

        Gj = exp(-1.0im * τ/2 * hj)
        push!(gates,Gj)
        push!(hgates, hj)
    end

    # End site
    hn = U*op("Nupdn",sites[N]) + ω*op("Nb",sites[N]) +
            g0*op("Ntot(Bd+B)",sites[N])
    Gn = exp(-1.0im * τ/2 * hn)
    push!(gates,Gn)
    push!(hgates, hn)

    append!(gates,reverse(gates))
    HubbardHolsteinModel(sites, H, gates, ampo)
end

function initialize_wavefcn(HH::HubbardHolsteinModel, p::Parameters)
    state = [isodd(n) ? "Up,0" : "Dn,0" for n=1:p.N] # NOTE: the QN of this state is preserved through DMRG

    # Account for doping
    if p.doping > 0
        spacing = floor(1/p.doping)
        state[1:spacing:end] .= "Emp,0"
    end

    productMPS(HH.sites,state)
end

function run_DMRG(HH::HubbardHolsteinModel, p::Parameters)
    ϕ0 = initialize_wavefcn(HH,p)
    @show flux(ϕ0)
    energy, ϕ = dmrg(HH.mpo, ϕ0, p.sweeps)
    entropy = compute_entropy(ϕ, p.mid)
    return DMRGResults(ϕ, energy, entropy)
end

function apply_onesite_operator(ϕ::MPS, opname::String, sites, siteidx::Int)
    ψ = copy(ϕ) # Make a copy of the original state
    new_ψj = op(opname,sites[siteidx]) * ψ[siteidx] # Apply the local operator
    noprime!(new_ψj)
    ψ[siteidx] = new_ψj
    return ψ
end

function compute_entropy(ψ::MPS, s)
    orthogonalize!(copy(ψ), s)
    _,S,_ = svd(ψ[s], (linkind(ψ, s-1), siteind(ψ,s)))
    S = S ./ sum(S.^2) # Normalize so that the sum squared of eigenvalues is 1
    SvN = 0.0
    for n=1:dim(S, 1)
        p = S[n,n]
        SvN -= p * log(p)
    end
    return SvN
end

function compute_overlap(ψ1::MPS, ψ2::MPS)
    LinearAlgebra.norm(inner(ψ1, ψ2))
end

function compute_phonon_number(ψ::MPS)
    expect(ψ,"Nb")
end

function time_evolve(dmrg_results::DMRGResults, HH::HubbardHolsteinModel, p::Parameters; order=4)
    ℋ = HH.ampo
    s = HH.sites
    ψ₀ = dmrg_results.ground_state
    nsteps = 10#floor(p.T/p.τ)
    t = 0.1
    println("Making 𝒰")
    𝒰 = exp(im * t * ℋ; alg=Trotter{order}(nsteps))
    println("Making U")
    U = Prod{ITensor}(𝒰, s)
    println("Applying the gate")
    ψ = apply(U,ψ₀)
    @show inner(ψ₀, ψ)

    #H = ITensor(ℋ, s)
    #𝒰ʳᵉᶠ = exp(im * t * ℋ)
    #Uʳᵉᶠ = ITensor(𝒰ʳᵉᶠ, s)
    #Uʳᵉᶠψ₀ = replaceprime(Uʳᵉᶠ * prod(ψ₀), 1 => 0)
end

function make_spectral_fcn(corrs, p::Parameters)
    # now has dimensions time x space
    sqw = fft.fftshift(fft.fft2(corrs * p.τ, norm="ortho"),axes=0) # TODO: Check the axis
    qs = range(0, stop=2*π, length=p.N+2)[2:end-1]
    ωs = 2 * π * fft.fftshift(fft.fftfreq(Int(p.T/p.τ), p.τ))
    for i in 1:p.N
        sqw[:,i] = imag.(exp(1im * qs[i] * p.mid) * sqw[:,i])
    end
    return real(sqw)/π, ωs, qs
end

function plot_correlation_function(tebd_results::TEBDResults)
    heatmap(LinearAlgebra.norm.(tebd_results.corrs'))
end

function plot_spectral_function(tebd_results::TEBDResults, p::Parameters; lims=nothing)
    ff, ωs, qs = make_spectral_fcn(tebd_results.corrs', p)
    if !isnothing(lims)
        ff = ff[lims[1]:lims[2],:]
        ωs = ωs[lims[1]:lims[2]]
    end
    ff = circshift(ff', p.mid-1)'
    qs = range(-π, stop=π, length=p.N+2)[2:end-1]
    maxval = maximum(abs.(ff))
    heatmap(qs, ωs, abs.(ff), c=:bwr, clims=(-maxval, maxval))
end

function plot_entropy(tebd_results::TEBDResults)
    niters = length(tebd_results.entropy)
    ϕ_entropy = [tebd_results.entropy[n][1] for n in 1:niters]
    ψ_entropy = [tebd_results.entropy[n][2] for n in 1:niters]
    plot(1:niters, ϕ_entropy, label="ϕ(t)")
    plot!(1:niters, ψ_entropy, label="ψ(t)")
    title!("Von Neumann Entropy")
    xlabel!("Iteration")
end

function plot_phonon_flux(tebd_results::TEBDResults)
    ## TO DO
end

function plot_overlap(tebd_results::TEBDResults)
    plot(1:length(tebd_results.self_overlap), tebd_results.self_overlap)
end

## CODE ##

## PARAMETERS ##

# Model
N = 8
t = 1 ## THIS TERM IS FINE
U = 8
ω = 0*t ## THIS TERM IS FINE (by itself)
g0 = 0*t ## THIS TERM IS FINE (by itself)
g1 = 0*g0 ## THIS TERM IS FINE (by itself)

# Simulation
T = 5
τ = 0.01
DMRG_numsweeps = 20
TEBD_maxdim = 800
TEBD_cutoff = 1E-14

## TODO: Modify number of phonons on each site from script
## TODO: Confirm that it's okay to overwrite the operators in the hubbardholstein.jl sites file
## TODO: How many phonons should we initially add to each site?

# Specify operators of interest
A_t0 = "Cup"
A_t = "Cup"

# Initialize
println("Initializing...")
params = parameters(N=N, t=t, U=U, ω=ω, g0=g0, g1=g1,
                    DMRG_numsweeps=DMRG_numsweeps,
                    T=T, τ=τ, TEBD_cutoff=TEBD_cutoff)
hubbholst = HubbardHolsteinModel(params)

# Run DMRG
println("Finding ground state...")
dmrg_results = run_DMRG(hubbholst, params)
@show compute_phonon_number(dmrg_results.ground_state)

# Time evolve
println("Time evolving...")
tebd_results = time_evolve(dmrg_results, hubbholst, params)


text:
function TimeMarching.r₁(u::Tuple{Nodes{Dual,NX,NY},Vector{Float64}},t,sys::NavierStokes{NX,NY},
                            motion::RigidBodyMotions.RigidBodyMotion) where {NX,NY}
    _,ċ,_,_,α̇,_ = motion(t)
    return TimeMarching.r₁(u[1],t,sys), [real(ċ),imag(ċ),α̇]
end

function TimeMarching.r₂(u::Tuple{Nodes{Dual,NX,NY},Vector{Float64}},t,sys::NavierStokes{NX,NY,N,false},
                            motion::RigidBodyMotions.RigidBodyMotion) where {NX,NY,N}

  # for now, just assume that there is only one body. will fix later.
  xc, yc, α = u[2]
  T = Bodies.RigidTransform((xc,yc),α)
  x, y = T(sys.X̃.u,sys.X̃.v)

  ΔV = VectorData(sys.X̃)
  _,ċ,_,_,α̇,_ = motion(t)
  for i = 1:N
      Δz = (x[i]-xc)+im*(y[i]-yc)
      ċi = ċ + im*α̇*Δz
      ΔV.u[i] = real(ċi)
      ΔV.v[i] = imag(ċi)
  end
  ΔV.u .-= sys.U∞[1]
  ΔV.v .-= sys.U∞[2]
  return ΔV, Vector{Float64}()
end

function TimeMarching.plan_constraints(u::Tuple{Nodes{Dual,NX,NY},Vector{Float64}},t,sys::NavierStokes{NX,NY,N,false}) where {NX,NY,N}

  # for now, just assume that there is only one body. will fix later.

  xc, yc, α = u[2]
  T = Bodies.RigidTransform((xc,yc),α)
  # should be able to save some time and memory allocation here...
  x, y = T(sys.X̃.u,sys.X̃.v)
  X = VectorData(x,y)
  regop = Regularize(X,cellsize(sys);issymmetric=true,I0=origin(sys))
  if sys._isstore
    Hmat, Emat = RegularizationMatrix(regop,VectorData{N}(),Edges{Primal,NX,NY}())
    sys.Hmat = Hmat
    sys.Emat = Emat
    return (f->TimeMarching.B₁ᵀ(f,sys), f->zeros(Float64,size(u[2]))),
           (w->TimeMarching.B₂(w,sys), u->Vector{Float64}())
  else
    return (f->TimeMarching.B₁ᵀ(f,regop,sys), f->zeros(Float64,size(u[2]))),
           (w->TimeMarching.B₂(w,regop,sys), u->Vector{Float64}())
  end

end


text:
function number_horizontal!(r)

    function change(side, kol)
        return HorizonSide((Int(side) + kol) % 4)
    end

    function move_and_kol(side, kol)
        flag = false
        count = 0
        while !isborder(r, side)
            if (isborder(r, change(side, kol)) && !flag)
                flag = true
                count += 1
            elseif !isborder(r, change(side, kol))
                flag = false
            end
            move!(r, side)
        end
        return count
    end

    function solve()
        side = Ost
        n = 1
        count = 0
        while !isborder(r, Nord)
            count += move_and_kol(side, n)
            side = change(side, 2)
            n = (n + 2) % 4
            move!(r, Nord)
        end
        println(count)
    end

    solve()
end

max_stars_repo_name: dmitrii-a-piskunov/mirea-progs

max_issues_repo_name: dmitrii-a-piskunov/mirea-progs

max_forks_repo_name: dmitrii-a-piskunov/mirea-progs


text:
using JLD2, Attitude, MATLAB, SatelliteDynamics, Dierckx
# include ksfunctions.jl
include(joinpath(dirname(dirname(@__FILE__)),"ks_functions.jl"))
@load "rate_control_transfers/60_day_v2_transfer.jld2" X U t_hist
dscale = 1e7 # m
tscale = 20000 # s
uscale = 10000.0
dt = 3e-2

r_eci_hist = dscale*mat_from_vec([x_from_u(X[i][1:4]) for i = 1:length(X)])
v_eci_hist = (dscale/tscale)*mat_from_vec([xdot_from_u(X[i][1:4],X[i][5:8]) for i = 1:length(X)])

# TODO: Spline this in a data efficient way
mat"
figure
hold on
plot3($r_eci_hist(1,:),$r_eci_hist(2,:),$r_eci_hist(3,:))
axis equal
hold off
"

Xm = mat_from_vec(X)
U_real = Xm[12:14,:]

mat"
figure
hold on
title('Thrust Vector')
plot($t_hist,$U_real')
ylabel('Thrust Acceleration (0.0001m/s^2)')
xlabel('Time (days)')
hold off
xlim([0,$t_hist(end)])
legend('u_x','u_y','u_z')
%saveas(gcf,'long_traj.png')
"

U_real_vec = [X[i][12:14] for i = 1:length(X)]
Unorm = [norm(U_real_vec[i]) for i = 1:length(X)]

mat"
figure
hold on
plot($t_hist,$Unorm/$uscale)
xlim([$t_hist(4),$t_hist(end)])
xlabel('Time (days)')
ylabel('Thrust Magnitude (m/s^2)')
hold off
"

println("Time in days:  $t_days")

# now let's get orbital elements
const SD = SatelliteDynamics
oe_hist = [SD.sCARTtoOSC([r_eci_hist[:,i];v_eci_hist[:,i]]) for i = 1:size(r_eci_hist,2)]

a_hist = [oe_hist[i][1] for i = 1:length(oe_hist)]
e_hist = [oe_hist[i][2] for i = 1:length(oe_hist)]
i_hist = [rad2deg(oe_hist[i][3]) for i = 1:length(oe_hist)]

mat"
figure
hold on
subplot(3,1,1)
hold on
plot($t_hist,$a_hist/1000)
h = plot($t_hist, 42e3*ones(length($t_hist),1),'r');
legend([h(1)],'GEO SMA','Location','SouthEast')
ylabel('Semi-Major Axis (km)')
xlabel('Time (days)')
xlim([0,$t_hist(end)])
ylim([20e3,46e3])
subplot(3,1,2)
hold on
plot($t_hist,$e_hist)
ylabel('Eccentricity')
xlabel('Time (days)')
xlim([0,$t_hist(end)])
subplot(3,1,3)
hold on
plot($t_hist,$i_hist)
ylabel('Inclination (deg)')
xlabel('Time (days)')
xlim([0,$t_hist(end)])
hold off
"

angle_hist = get_rtn_in_out(U_real,r_eci_hist,v_eci_hist)
angles = mat_from_vec(angle_hist)

sixty = (t_hist = t_hist, a_hist = a_hist, e_hist = e_hist, i_hist = i_hist, Unorm = Unorm,angles = angles)

using JLD2
@save "plotting_data/plotting_data_60.jld2" sixty

max_stars_repo_name: RoboticExplorationLab/KSLowThrust

max_issues_repo_name: RoboticExplorationLab/KSLowThrust

max_forks_repo_name: RoboticExplorationLab/KSLowThrust


text:
#############################################################################
# add_subtract.jl
# Handles unary negation, addition and subtraction of variables, constants
# and expressions.
# All expressions and atoms are subtpyes of AbstractExpr.
# Please read expressions.jl first.
#############################################################################

export +, -
export sign, curvature, monotonicity, evaluate

### Unary Negation

struct NegateAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr}
    size::Tuple{Int, Int}

    function NegateAtom(x::AbstractExpr)
        children = (x,)
        return new(:-, hash(children), children, x.size)
    end
end

function sign(x::NegateAtom)
    return -sign(x.children[1])
end

function monotonicity(x::NegateAtom)
    return (Nonincreasing(),)
end

function curvature(x::NegateAtom)
    return ConstVexity()
end

function evaluate(x::NegateAtom)
    return -evaluate(x.children[1])
end

-(x::AbstractExpr) = NegateAtom(x)

function conic_form!(x::NegateAtom, unique_conic_forms::UniqueConicForms)
    if !has_conic_form(unique_conic_forms, x)
        objective = conic_form!(x.children[1], unique_conic_forms)
        objective = -objective
        cache_conic_form!(unique_conic_forms, x, objective)
    end
    return get_conic_form(unique_conic_forms, x)
end

### Addition
struct AdditionAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Array{AbstractExpr, 1}
    size::Tuple{Int, Int}

    function AdditionAtom(x::AbstractExpr, y::AbstractExpr)
        # find the size of the expression = max of size of x and size of y
        if x.size == y.size || y.size == (1, 1)
            sz = x.size
        elseif x.size == (1, 1)
            sz = y.size
        else
            error("Cannot add expressions of sizes $(x.size) and $(y.size)")
        end
        # see if we're forming a sum of more than two terms and condense them
        children = AbstractExpr[]
        if isa(x, AdditionAtom)
            append!(children, x.children)
        else
            push!(children, x)
        end
        if isa(y, AdditionAtom)
            append!(children, y.children)
        else
            push!(children, y)
        end
        return new(:+, hash(children), children, sz)
    end
end

function sign(x::AdditionAtom)
    return sum(Sign[sign(child) for child in x.children])
    # Creating an array of type Sign and adding all the sign of xhildren of x so if anyone is complex the resultant sign would be complex.
end

function monotonicity(x::AdditionAtom)
    return Monotonicity[Nondecreasing() for child in x.children]
end

function curvature(x::AdditionAtom)
    return ConstVexity()
end

function evaluate(x::AdditionAtom)
    # broadcast is used in reduction instead of using sum directly to support addition
    # between scalars and arrays
    return mapreduce(evaluate, (a, b) -> a .+ b, x.children)
end

function conic_form!(x::AdditionAtom, unique_conic_forms::UniqueConicForms)
    if !has_conic_form(unique_conic_forms, x)
        objective = ConicObj()
        for child in x.children
            child_objective = conic_form!(child, unique_conic_forms)
            if x.size != child.size
                child_objective = promote_size(child_objective, length(x))
            end
            objective += child_objective
        end
        cache_conic_form!(unique_conic_forms, x, objective)
    end
    return get_conic_form(unique_conic_forms, x)
end

+(x::AbstractExpr, y::AbstractExpr) = AdditionAtom(x, y)
+(x::Value, y::AbstractExpr) = AdditionAtom(Constant(x), y)
+(x::AbstractExpr, y::Value) = AdditionAtom(x, Constant(y))
-(x::AbstractExpr, y::AbstractExpr) = x + (-y)
-(x::Value, y::AbstractExpr) = Constant(x) + (-y)
-(x::AbstractExpr, y::Value) = x + Constant(-y)


text:
"""
    hsvd(fid, rank)

Use the Hankel singular value decomposition method for fitting a sum of
exponential signals of the form

    xₙ = Σₖ Aₖ zₖⁿ

to compute the ampltudes Aₖ and complex-valued poles zₖ. Any poles with |zₖ| >
1 are discarded to avoid numerical instability during the fit. A tuple of poles,
amplitudes and the exponential basis (z,A,basis) is returned.

The discrete-time and continuous time signals are related by

    sₙ =  Aₖ zₖⁿ  =  Aₖ exp((-bₖ + iωₖ) * tₙ)

or

    zₖ  =  exp((-bₖ + iωₖ) * Δt)

where ωₖ is an angular frequency and bₖ the decay rate in Hz.

Ref. to NMR literature: Barkhuijsen et al., "Improved Algorithm for Noniterative
Time-Domain Model Fitting to Exponentially Damped Magnetic Resonance Signals",
J. Mag. Res., (1987) https://doi.org/10.1016/0022-2364(87)90023-0

More broadly these techniques seem to come under the title "Singular Spectrum
analysis".
"""
function hsvd(fid, rank)
    # In the code below we attempt to use the same notation as the paper.
    # Construct the Hankel matrix
    N = length(fid)
    L = N÷2
    M = N-L
    X = zeros(ComplexF64, L, M)
    for i = 1:M
        X[:, i] .= fid[i:i+L-1]
    end

    # Form a low-rank approximation to the signal basis by truncating the SVD U
    # factor to the given rank.
    hsvd = svd(X)
    Xl = hsvd.U[:, 1:rank]

    # As described in the paper, each row of Xl is the row below multiplied by
    # some matrix Zp, so we can compute Zp with a linear by solving the
    # overdetermined system `Ub*Zp = Ut` as a least squares problem:
    Ub = Xl[1:end-1, :]
    Ut = Xl[2:end, :]
    Zp = Ub \ Ut

    # The eigen decomposition of Zp now gives the eigenvalues as the poles
    Zeig = eigen(Zp)
    Zeig.values
end

"""
Use HSVD to suppress a spectral region.

This is rather a WIP implementation. TODO
* Make all this work with AxisArrays
* Proper testing...
"""
function hsvd_water_suppression(fid; dt = 0.0008, water_center_freq=0.0,
                                bandwidth=100.0, hsvd_rank=50)
    zs0 = hsvd(fid, hsvd_rank)

    # Filter out any exponentially growing components. These can occur if a
    # naive truncation-based frequency domain downsampling of the time domain
    # signal is used. (For example this seems to be the default on Siemens
    # scanners when using the "remove oversampling" option.)
    zs = zs0[abs.(zs0) .< 1]

    # Generate a basis from the components and fit to it to determine the
    # amplitudes of the various signal components.
    basis = [zs[i]^n for n=0:length(fid)-1, i=1:length(zs)]
    A = basis \ fid

    # Compute MR frequencies of components
    frequencies = imag(log.(zs))/dt

    # Assume all signal within some spectral bandwidth of water centre frequency
    water_components = abs.(frequencies .- water_center_freq) .< bandwidth

    water = basis[:,water_components] * A[water_components]

    fid .- water
end

max_stars_repo_name: TRIImaging/MagneticResonanceSignals

max_issues_repo_name: TRIImaging/MagneticResonanceSignals

max_forks_repo_name: TRIImaging/MagneticResonanceSignals


text:
export RainbowLearner, update!

using Flux, Zygote
using StatsBase

"""
    RainbowLearner(;kwargs...)

See paper: [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://arxiv.org/abs/1710.02298)

# Keywords

- `approximator`::[`AbstractQApproximator`](@ref): used to get Q-values of a state.
- `target_approximator`::[`AbstractQApproximator`](@ref): similar to `approximator`, but used to estimate the target (the next state).
- `loss_fun`: the loss function.
- `Vₘₐₓ::Float32`: the maximum value of distribution.
- `Vₘᵢₙ::Float32`: the minimum value of distribution.
- `n_actions::Int`: number of possible actions.
- `γ::Float32=0.99f0`: discount rate.
- `batch_size::Int=32`
- `update_horizon::Int=1`: length of update ('n' in n-step update).
- `min_replay_history::Int=32`: number of transitions that should be experienced before updating the `approximator`.
- `update_freq::Int=4`: the frequency of updating the `approximator`.
- `target_update_freq::Int=500`: the frequency of syncing `target_approximator`.
- `stack_size::Union{Int, Nothing}=4`: use the recent `stack_size` frames to form a stacked state.
- `default_priority::Float64=100.`: the default priority for newly added transitions.
- `n_atoms::Int=51`: the number of buckets of the value function distribution.
- `stack_size::Union{Int, Nothing}=4`: use the recent `stack_size` frames to form a stacked state.
- `default_priority::Float64=100.`: the default priority for newly added transitions.

!!! warning
    The `Zygote_gpu` version is slow due to that the `argmax(A, dims=1)` falls back to the CPU version in CuArrays.
"""
Base.@kwdef mutable struct RainbowLearner{Tq<:AbstractQApproximator, Tt<:AbstractQApproximator,Tf,Ts,Tss<:Union{Int, Nothing}} <: AbstractLearner
    approximator::Tq
    target_approximator::Tq
    loss_fun::Tf
    Vₘₐₓ::Float32
    Vₘᵢₙ::Float32
    n_actions::Int
    n_atoms::Int = 51
    support::Ts = collect(range(Float32(-Vₘₐₓ), Float32(Vₘₐₓ), length = n_atoms))
    stack_size::Tss = 4
    delta_z::Float32 = Float32(support[2] - support[1])
    γ::Float32 = 0.99
    batch_size::Int = 32
    update_horizon::Int = 1
    min_replay_history::Int = 32
    update_freq::Int = 1
    target_update_freq::Int = 500
    update_step::Int = 0
    default_priority::Float64 = 100.0

    function RainbowLearner(approximator::Tq, target_approximator::Tt, loss_fun::Tf, Vₘₐₓ::Float32, Vₘᵢₙ::Float32, n_actions::Int, n_atoms::Int, support::Ts, stack_size::Tss, args...) where {Tq,Tt, Tf, Ts, Tss}
        copyto!(approximator, target_approximator)  # force sync
        support = to_device(approximator, support)
        new{Tq,Tt,Tf, typeof(support), Tss}(approximator, target_approximator, loss_fun, Vₘₐₓ, Vₘᵢₙ, n_actions, n_atoms, support, stack_size, args...)
    end
end

function (learner::RainbowLearner)(obs::Observation)
    state = get_state(obs)
    logits = batch_estimate(learner.approximator, to_device(learner.approximator, state))
    q = learner.support .* softmax(reshape(logits, :, learner.n_actions))
    # probs = vec(sum(q, dims=1)) .+ legal_action
    vec(sum(q, dims = 1)) |> to_host
end

function update!(learner::RainbowLearner, batch)
    learner.update_step += 1
    learner.update_step % learner.update_freq == 0 || return nothing

    Q, Qₜ, γ, loss_fun, n_atoms, n_actions, support, delta_z, update_horizon, batch_size = learner.approximator,
        learner.target_approximator,
        learner.γ,
        learner.loss_fun,
        learner.n_atoms,
        learner.n_actions,
        learner.support,
        learner.delta_z,
        learner.update_horizon,
        learner.batch_size

    states, rewards, terminals, next_states = map(x->to_device(Q, x), (batch.states, batch.rewards, batch.terminals, batch.next_states))
    actions = CartesianIndex.(batch.actions, 1:batch_size)

    target_support = reshape(rewards, 1, :) .+
                     (reshape(support, :, 1) *
                      reshape((γ^update_horizon) .* (1 .- terminals), 1, :))

    updated_priorities = Vector{Float32}()

    gs = gradient(Q) do
        logits = reshape(batch_estimate(Q, states), n_atoms, n_actions, :)
        select_logits = logits[:, actions]
        next_logits = batch_estimate(Qₜ, next_states)
        next_probs = reshape(softmax(reshape(next_logits, n_atoms, :)), n_atoms, n_actions, :)
        next_q = reshape(sum(support .* next_probs, dims = 1), n_actions, :)
        # next_q_argmax = argmax(cpu(next_q .+ next_legal_actions), dims=1)
        next_q_argmax = argmax(next_q, dims = 1)
        next_prob_select = reshape(next_probs[:, next_q_argmax], n_atoms, :)

        target_distribution = project_distribution(
            target_support,
            next_prob_select,
            support,
            delta_z,
            learner.Vₘᵢₙ,
            learner.Vₘₐₓ,
        )

        batch_losses = loss_fun(select_logits, Zygote.dropgrad(target_distribution))
        updated_priorities = vec(clamp.(sqrt.(batch_losses .+ 1f-10), 1.f0, 1.f2))

        target_priorities = 1.0f0 ./ sqrt.(updated_priorities .+ 1f-10)
        target_priorities = target_priorities ./ maximum(target_priorities)

        mean(Zygote.dropgrad(target_priorities) .* batch_losses)
    end

    update!(Q, gs)

    if learner.update_step % learner.target_update_freq == 0
        copyto!(Qₜ, Q)
    end

    updated_priorities |> to_host
end

function project_distribution(supports, weights, target_support, delta_z, vmin, vmax)
    batch_size, n_atoms = size(supports, 2), length(target_support)
    clampped_support = clamp.(supports, vmin, vmax)
    tiled_support = reshape(repeat(clampped_support; outer=(n_atoms, 1)), n_atoms, n_atoms, batch_size)

    projection = clamp.(
        1 .- abs.(tiled_support .- reshape(target_support, 1, :)) ./ delta_z,
        0,
        1,
    ) .* reshape(weights, n_atoms, 1, batch_size)
    reshape(sum(projection, dims = 1), n_atoms, batch_size)
end

function extract_transitions(buffer::CircularTurnBuffer{PRTSA}, learner::RainbowLearner)
    if length(buffer) > learner.min_replay_history
        inds, consecutive_batch = sample(buffer, learner.batch_size, learner.update_horizon, learner.stack_size)
        inds, extract_SARTS(consecutive_batch, learner.γ)
    else
        nothing
    end
end


text:
module MyPkg
# this is a comment. V4
greet() = print("Hello World!")

# Another comment
# Comment2
end # module


text:
module CairoMakie

using Makie, LinearAlgebra
using Colors, GeometryBasics, FileIO, StaticArrays
import SHA
import Base64
import Cairo

using Makie: Scene, Lines, Text, Image, Heatmap, Scatter, @key_str, broadcast_foreach
using Makie: convert_attribute, @extractvalue, LineSegments, to_ndim, NativeFont
using Makie: @info, @get_attribute, Combined
using Makie: to_value, to_colormap, extrema_nan
using Makie: inline!

const OneOrVec{T} = Union{
    T,
    Vec{N1, T} where N1,
    NTuple{N2, T} where N2,
}

# re-export Makie, including deprecated names
for name in names(Makie, all=true)
    if Base.isexported(Makie, name)
        @eval using Makie: $(name)
        @eval export $(name)
    end
end
export inline!

include("infrastructure.jl")
include("utils.jl")
include("fonts.jl")
include("primitives.jl")
include("overrides.jl")

function __init__()
    activate!()
    Makie.register_backend!(Makie.current_backend[])
end

function display_path(type::String)
    if !(type in ("svg", "png", "pdf", "eps"))
        error("Only \"svg\", \"png\", \"eps\" and \"pdf\" are allowed for `type`. Found: $(type)")
    end
    return joinpath(@__DIR__, "display." * type)
end

const _last_inline = Ref(true)
const _last_type = Ref("png")
const _last_px_per_unit = Ref(1.0)
const _last_pt_per_unit = Ref(0.75)

function activate!(; inline = _last_inline[], type = _last_type[], px_per_unit=_last_px_per_unit[], pt_per_unit=_last_pt_per_unit[])
    backend = CairoBackend(display_path(type); px_per_unit=px_per_unit, pt_per_unit=pt_per_unit)
    Makie.current_backend[] = backend
    Makie.use_display[] = !inline
    _last_inline[] = inline
    _last_type[] = type
    _last_px_per_unit[] = px_per_unit
    _last_pt_per_unit[] = pt_per_unit
    return
end

if Base.VERSION >= v"1.4.2"
    include("precompiles.jl")
    _precompile_()
end

end


text:
using GXBeam, LinearAlgebra, BenchmarkTools

function steady_benchmark()

    # straight section of the beam
    L_b1 = 31.5 # inch
    r_b1 = [0, 0, 0]
    nelem_b1 = 20
    lengths_b1, xp_b1, xm_b1, Cab_b1 = discretize_beam(L_b1, r_b1, nelem_b1)

    # swept section of the beam
    L_b2 = 6 # inch
    r_b2 = [31.5, 0, 0]
    nelem_b2 = 4
    frame_b2 = [
        0.707106781186548 0.707106781186548 0.0;
        -0.707106781186548 0.707106781186548 0.0;
        0.0               0.0               1.0]
    lengths_b2, xp_b2, xm_b2, Cab_b2 = discretize_beam(L_b2, r_b2, nelem_b2;
        frame = frame_b2)

    # combine elements and points into one array
    nelem = nelem_b1 + nelem_b2
    points = vcat(xp_b1, xp_b2[2:end])
    start = 1:nelem_b1 + nelem_b2
    stop = 2:nelem_b1 + nelem_b2 + 1
    lengths = vcat(lengths_b1, lengths_b2)
    midpoints = vcat(xm_b1, xm_b2)
    Cab = vcat(Cab_b1, Cab_b2)

    compliance = fill(Diagonal([
        1.4974543276E-06,
        4.7619054919E-06,
        5.8036221902E-05,
        3.1234387938E-03,
        4.5274507258E-03,
        1.7969451932E-05]), nelem)

    mass = fill(Diagonal([
        1.5813000000E-05,
        1.5813000000E-05,
        1.5813000000E-05,
        1.3229801498E-06,
        5.2301497500E-09,
        1.3177500000E-06]), nelem)

    # create assembly
    assembly = Assembly(points, start, stop;
        compliance = compliance,
        mass = mass,
        frames = Cab,
        lengths = lengths,
        midpoints = midpoints)

    # create dictionary of prescribed conditions
    prescribed_conditions = Dict(
        # root section is fixed
        1 => PrescribedConditions(ux=0, uy=0, uz=0, theta_x=0, theta_y=0, theta_z=0)
        )

    # perform nonlinear steady state analysis
    system, converged = steady_state_analysis(assembly,
        linear_velocity = [0, 196.349540849362, 0],
        angular_velocity = [0, 0, 78.5398163397448],
        prescribed_conditions = prescribed_conditions)

    # test convergence
    @assert converged

    result = AssemblyState(system, assembly;
        prescribed_conditions = prescribed_conditions)

    return result
end

function eigen_benchmark()

    # straight section of the beam
    L_b1 = 31.5 # inch
    r_b1 = [0, 0, 0]
    nelem_b1 = 20
    lengths_b1, xp_b1, xm_b1, Cab_b1 = discretize_beam(L_b1, r_b1, nelem_b1)

    # swept section of the beam
    L_b2 = 6 # inch
    r_b2 = [31.5, 0, 0]
    nelem_b2 = 4
    frame_b2 = [
        0.707106781186548 0.707106781186548 0.0;
        -0.707106781186548 0.707106781186548 0.0;
        0.0               0.0               1.0]
    lengths_b2, xp_b2, xm_b2, Cab_b2 = discretize_beam(L_b2, r_b2, nelem_b2;
        frame = frame_b2)

    # combine elements and points into one array
    nelem = nelem_b1 + nelem_b2
    points = vcat(xp_b1, xp_b2[2:end])
    start = 1:nelem_b1 + nelem_b2
    stop = 2:nelem_b1 + nelem_b2 + 1
    lengths = vcat(lengths_b1, lengths_b2)
    midpoints = vcat(xm_b1, xm_b2)
    Cab = vcat(Cab_b1, Cab_b2)

    compliance = fill(Diagonal([
        1.4974543276E-06,
        4.7619054919E-06,
        5.8036221902E-05,
        3.1234387938E-03,
        4.5274507258E-03,
        1.7969451932E-05]), nelem)

    mass = fill(Diagonal([
        1.5813000000E-05,
        1.5813000000E-05,
        1.5813000000E-05,
        1.3229801498E-06,
        5.2301497500E-09,
        1.3177500000E-06]), nelem)

    # create assembly
    assembly = Assembly(points, start, stop;
        compliance = compliance,
        mass = mass,
        frames = Cab,
        lengths = lengths,
        midpoints = midpoints)

    # create dictionary of prescribed conditions
    prescribed_conditions = Dict(
        # root section is fixed
        1 => PrescribedConditions(ux=0, uy=0, uz=0, theta_x=0, theta_y=0, theta_z=0)
        )

    # perform eigenvalue analysis
    system, λ, V, converged = eigenvalue_analysis(assembly; nev = 20,
        linear_velocity = [0, 196.349540849362, 0],
        angular_velocity = [0, 0, 78.5398163397448],
        prescribed_conditions = prescribed_conditions)

    # test convergence
    @assert converged

    return λ, V
end

function unsteady_benchmark()

    tvec = range(0, 0.5, length=2001)

    # straight section of the beam
    L_b1 = 31.5 # inch
    r_b1 = [0, 0, 0]
    nelem_b1 = 20
    lengths_b1, xp_b1, xm_b1, Cab_b1 = discretize_beam(L_b1, r_b1, nelem_b1)

    # swept section of the beam
    L_b2 = 6 # inch
    r_b2 = [31.5, 0, 0]
    nelem_b2 = 4
    frame_b2 = [
        0.707106781186548 0.707106781186548 0.0;
        -0.707106781186548 0.707106781186548 0.0;
        0.0               0.0               1.0]
    lengths_b2, xp_b2, xm_b2, Cab_b2 = discretize_beam(L_b2, r_b2, nelem_b2;
        frame = frame_b2)

    # combine elements and points into one array
    nelem = nelem_b1 + nelem_b2
    points = vcat(xp_b1, xp_b2[2:end])
    start = 1:nelem_b1 + nelem_b2
    stop = 2:nelem_b1 + nelem_b2 + 1
    lengths = vcat(lengths_b1, lengths_b2)
    midpoints = vcat(xm_b1, xm_b2)
    Cab = vcat(Cab_b1, Cab_b2)

    compliance = fill(Diagonal([
        1.4974543276E-06,
        4.7619054919E-06,
        5.8036221902E-05,
        3.1234387938E-03,
        4.5274507258E-03,
        1.7969451932E-05]), nelem)

    mass = fill(Diagonal([
        1.5813000000E-05,
        1.5813000000E-05,
        1.5813000000E-05,
        1.3229801498E-06,
        5.2301497500E-09,
        1.3177500000E-06]), nelem)

    # create assembly
    assembly = Assembly(points, start, stop;
        compliance = compliance,
        mass = mass,
        frames = Cab,
        lengths = lengths,
        midpoints = midpoints)

    # time function
    tf1 = (t) -> t < 0.2 ? 0.1*(1 - sin(2*pi*(t/0.4 + 0.25))) : zero(t)

    # prescribed condition
    prescribed_conditions = (t) -> Dict(
        1 => PrescribedConditions(ux=0, uy=0, uz=0, theta_x=0, theta_y=0, theta_z=0),
        nelem_b1 + 1 => PrescribedConditions(Fx=0, Fy=100*tf1(t), Fz=0, Mx=0, My=0, Mz=0)
    )

    # perform time marching analysis
    system, history, converged = time_domain_analysis(assembly, tvec;
        linear_velocity = [0, 196.349540849362, 0],
        angular_velocity = [0, 0, 78.5398163397448],
        prescribed_conditions = prescribed_conditions)

    # test convergence
    @assert converged

    return history
end

println("Steady State Analysis: ")
@btime steady_benchmark();

# Results on a single i7-5500U CPU @ 2.40GHz core.
# GEBT: 13.722 ms
# GXBeam: 4.716 ms (2207 allocations: 2.79 MiB)

println("Eigenvalue Analysis: ")
@btime eigenvalue_benchmark();

# Results on a single i7-5500U CPU @ 2.40GHz core.
# GEBT: 33.712 ms
# GXBeam: 18.478 ms (2568 allocations: 5.16 MiB)

println("Time Marching Analysis: ")
@btime unsteady_benchmark();

# Results on a single i7-5500U CPU @ 2.40GHz core.
# GEBT: 26.870 s
# GXBeam: 9.019 s (851711 allocations: 4.30 GiB)