text:
try
    ts = time_ns()
    let
        test = ByteStreamTestHarness("write-write-end-pop-pop", 15)

        execute(test, Write("cat"))

        execute(test, InputEnded(false))
        execute(test, BufferEmpty(false))
        execute(test, Eof(false))
        execute(test, BytesRead(0))
        execute(test, BytesWritten(3))
        execute(test, RemainingCapacity(12))
        execute(test, BufferSize(3))
        execute(test, Peek("cat"))

        execute(test, Write("tac"))

        execute(test, InputEnded(false))
        execute(test, BufferEmpty(false))
        execute(test, Eof(false))
        execute(test, BytesRead(0))
        execute(test, BytesWritten(6))
        execute(test, RemainingCapacity(9))
        execute(test, BufferSize(6))
        execute(test, Peek("cattac"))

        execute(test, EndInput())

        execute(test, InputEnded(true))
        execute(test, BufferEmpty(false))
        execute(test, Eof(false))
        execute(test, BytesRead(0))
        execute(test, BytesWritten(6))
        execute(test, RemainingCapacity(9))
        execute(test, BufferSize(6))
        execute(test, Peek("cattac"))

        execute(test, Pop(2))

        execute(test, InputEnded(true))
        execute(test, BufferEmpty(false))
        execute(test, Eof(false))
        execute(test, BytesRead(2))
        execute(test, BytesWritten(6))
        execute(test, RemainingCapacity(11))
        execute(test, BufferSize(4))
        execute(test, Peek("ttac"))

        execute(test, Pop(4))

        execute(test, InputEnded(true))
        execute(test, BufferEmpty(true))
        execute(test, Eof(true))
        execute(test, BytesRead(6))
        execute(test, BytesWritten(6))
        execute(test, RemainingCapacity(15))
        execute(test, BufferSize(0))
    end

    let
        test = ByteStreamTestHarness("write-pop-write-end-pop", 15)

        execute(test, Write("cat"))

        execute(test, InputEnded(false))
        execute(test, BufferEmpty(false))
        execute(test, Eof(false))
        execute(test, BytesRead(0))
        execute(test, BytesWritten(3))
        execute(test, RemainingCapacity(12))
        execute(test, BufferSize(3))
        execute(test, Peek("cat"))

        execute(test, Pop(2))

        execute(test, InputEnded(false))
        execute(test, BufferEmpty(false))
        execute(test, Eof(false))
        execute(test, BytesRead(2))
        execute(test, BytesWritten(3))
        execute(test, RemainingCapacity(14))
        execute(test, BufferSize(1))
        execute(test, Peek("t"))

        execute(test, Write("tac"))

        execute(test, InputEnded(false))
        execute(test, BufferEmpty(false))
        execute(test, Eof(false))
        execute(test, BytesRead(2))
        execute(test, BytesWritten(6))
        execute(test, RemainingCapacity(11))
        execute(test, BufferSize(4))
        execute(test, Peek("ttac"))

        execute(test, EndInput())

        execute(test, InputEnded(true))
        execute(test, BufferEmpty(false))
        execute(test, Eof(false))
        execute(test, BytesRead(2))
        execute(test, BytesWritten(6))
        execute(test, RemainingCapacity(11))
        execute(test, BufferSize(4))
        execute(test, Peek("ttac"))

        execute(test, Pop(4))

        execute(test, InputEnded(true))
        execute(test, BufferEmpty(true))
        execute(test, Eof(true))
        execute(test, BytesRead(6))
        execute(test, BytesWritten(6))
        execute(test, RemainingCapacity(15))
        execute(test, BufferSize(0))
    end
    tc = (Int(time_ns() - ts)) / 1e9
    @info "byte_stream_two_writes passed......... Time: $tc"
catch e
    println(stderr, e)
end

max_stars_repo_name: AquaIndigo/JLSponge

max_issues_repo_name: AquaIndigo/JLSponge

max_forks_repo_name: AquaIndigo/JLSponge


text:
module simo_vuquoc_examples
# Simo, J. C. and L. Vuquoc (1988). "ON THE DYNAMICS IN SPACE OF RODS
# UNDERGOING LARGE MOTIONS - A GEOMETRICALLY EXACT APPROACH."
# Computer Methods in Applied Mechanics and Engineering 66(2): 125-161.

using FinEtools
using FinEtoolsDeforLinear
using FinEtoolsFlexBeams.CrossSectionModule: CrossSectionRectangle
using FinEtoolsFlexBeams.MeshFrameMemberModule: frame_member, merge_members
using FinEtoolsFlexBeams.FEMMCorotBeamModule
using FinEtoolsFlexBeams.FEMMCorotBeamModule: FEMMCorotBeam
stiffness = FEMMCorotBeamModule.stiffness
mass = FEMMCorotBeamModule.mass
distribloads_global = FEMMCorotBeamModule.distribloads_global
restoringforce = FEMMCorotBeamModule.restoringforce
gyroscopic = FEMMCorotBeamModule.gyroscopic
using FinEtoolsFlexBeams.RotUtilModule: initial_Rfield, linear_update_rotation_field!, update_rotation_field!
using LinearAlgebra: dot
using Arpack
using LinearAlgebra
using SparseArrays
using FinEtoolsFlexBeams.VisUtilModule: plot_space_box, plot_midline, plot_solid, render, react!, default_layout_3d, save_to_json
using PlotlyJS
using JSON

Aus =[0  0
   9.4612e-01   2.5692e-02
   1.5375e+00   9.2688e-01
   2.3259e+00   2.3913e+00
   3.0355e+00   3.6680e+00
   3.8633e+00   5.0949e+00
   4.5729e+00   6.3715e+00
   5.1643e+00   7.2352e+00
   5.8739e+00   8.0237e+00
   6.5834e+00   8.3241e+00
   7.4507e+00   8.0237e+00
   8.2786e+00   6.7846e+00
   8.8305e+00   5.6957e+00
   9.2641e+00   4.4190e+00
   9.8949e+00   2.6917e+00
   1.0329e+01   1.1522e+00
   1.0723e+01  -5.3755e-01
   1.1275e+01  -3.4664e+00
   1.1787e+01  -5.5692e+00
   1.2260e+01  -7.0711e+00
   1.2970e+01  -7.9723e+00
   1.3758e+01  -8.5356e+00
   1.4428e+01  -9.2115e+00
   1.5059e+01  -9.2115e+00
   1.5966e+01  -8.4229e+00
   1.6873e+01  -6.6581e+00
   1.7740e+01  -4.3300e+00
   1.8568e+01  -1.8893e+00
   1.8962e+01  -8.6957e-02
   1.9632e+01   2.5791e+00
   2.0302e+01   4.4565e+00
   2.0894e+01   5.5079e+00
   2.1879e+01   6.5217e+00
   2.2746e+01   6.7846e+00
   2.3693e+01   6.2964e+00
   2.4481e+01   6.1462e+00
   2.5191e+01   6.6719e+00
   2.5821e+01   7.4980e+00
   2.6413e+01   6.9348e+00
   2.6965e+01   5.4704e+00
   2.7438e+01   4.1561e+00
   2.7792e+01   2.8419e+00
   2.8423e+01   1.2648e+00
   2.9409e+01   8.1423e-01
   2.9961e+01   7.3913e-01];

function simo_vuquoc_animated()
    # Parameters:
    # The  parameters given in the paper  are in terms of stiffness and mass coefficients.
    asquared = 12.0e-3
    E = 1.0e6 / asquared;
    a = sqrt(asquared)
    E * 1/12 * a^4
    nu = 0.0;# Poisson ratio
    rho = 1 / asquared;
    b = a; h = a; L = 10; # cross-sectional dimensions and length of each leg
    utol = 1e-3;
    maxit = 30
    Fmag = 50;
    dt = 0.05;
    tend = 30.0;
    ng = 1/2; nb = 1/4*(1/2+ng)^2;
    # Choose the mass formulation:
    mass_type=1;

    # Cross-sectional properties
    cs = CrossSectionRectangle(s -> b, s -> h, s -> [0.0, 0.0, 1.0])

    # Select the number of elements per leg.
    n=4;
    tolerance=L/n/100;
    members = []
    push!(members, frame_member([0 L 0;L L 0], n, cs))
    push!(members, frame_member([L L 0;L 0 0], n, cs))
    fens, fes = merge_members(members; tolerance = tolerance);
    @show E * fes.A[1], E * fes.I2[1], E/2 * fes.J[1], rho * fes.A[1], rho * fes.I2[1]

    # Material properties
    material = MatDeforElastIso(DeforModelRed3D, rho, E, nu, 0.0)

    # Construct the requisite fields, geometry and displacement
    # Initialize configuration variables
    geom0 = NodalField(fens.xyz)
    u0 = NodalField(zeros(size(fens.xyz,1), 3))
    Rfield0 = initial_Rfield(fens)
    dchi = NodalField(zeros(size(fens.xyz,1), 6))

    # Locate some nodes
    box = fill(0.0, 6)
    # The pinned end is
    initbox!(box, [0,L,0])
    pinnedn = selectnode(fens; box = box, inflate = tolerance)
    # The force is applied at the elbow
    initbox!(box, [L,L,0])
    elbown = selectnode(fens; box = box, inflate = tolerance)
    # Tip is tracked
    initbox!(box, [L,0,0])
    tipn = selectnode(fens; box = box, inflate = tolerance)

    # Apply EBC's
    for i in [1, 2, 3, 4, 5, 6]
        setebc!(dchi, pinnedn, true, i)
    end
    applyebc!(dchi)
    numberdofs!(dchi);
    dchi.dofnums

    # Additional fields
    stepdchi = deepcopy(dchi)
    u1 = deepcopy(u0)
    v0 = deepcopy(dchi) # need the numbering of the degrees of freedom
    v1 = deepcopy(dchi)
    a1 = deepcopy(dchi) # zero out the Acceleration
    a0 = deepcopy(dchi) # zero out the Acceleration
    Rfield1 = deepcopy(Rfield0)
    v0v = gathersysvec(v0)
    a0v = gathersysvec(a0)
    vpv = gathersysvec(v0);
    dchipv = gathersysvec(dchi);
    stepdchiv = gathersysvec(dchi);
    rhs = gathersysvec(dchi);
    TMPv = deepcopy(rhs)
    utol = 1e-13*dchi.nfreedofs;

    tbox = plot_space_box([[-0.5*L -0.5*L -2.0*L]; [1.5*L 1.5*L 2.0*L]])
    tshape0 = plot_solid(fens, fes; x = geom0.values, u = 0.0.*dchi.values[:, 1:3], R = Rfield0.values, facecolor = "rgb(125, 155, 125)", opacity = 0.3);
    plots = cat(tbox,  tshape0; dims = 1)
    pl = render(plots)
    sleep(0.5)

    femm = FEMMCorotBeam(IntegDomain(fes, GaussRule(1, 2)), material)
    loadbdry = FESetP1(reshape(elbown, 1, 1))
    lfemm = FEMMBase(IntegDomain(loadbdry, PointRule()))

    t = 0.0; #
    step = 0;
    while (t <= tend)
        t = t + dt;
        # println("Time $(t)"); # pause
        # Initialization
        applyebc!(dchi) # Apply boundary conditions
        u1.values[:] = u0.values[:]; # guess
        Rfield1.values[:] = Rfield0.values[:]; # guess
        stepdchi.values[:] .= 0.0;# Total increment in current step
        a1.values[:] = -(1/nb/dt)*v0.values[:] -(1/2-nb)/nb*a0.values[:];
        v1.values[:] = v0.values[:] + dt*((1-ng)*a0.values[:] + ng*a1.values[:]);
        gathersysvec!(v0, v0v)
        gathersysvec!(a0, a0v)
        dchipv = dt*v0v + (dt^2/2*(1-2*nb))*a0v
        vpv = v0v +(dt*(1-ng))*a0v;
        fi = ForceIntensity(FFlt[0, 0, (t<2.0)*((t<1.0)*t+(t>1.0)*(2.0-t))*Fmag, 0, 0, 0]);

        iter = 1;
        while true
            F = distribloads(lfemm, geom0, dchi, fi, 3);
            Fr = restoringforce(femm, geom0, u1, Rfield1, dchi);       # Internal forces
            @. rhs = F + Fr;
            K = stiffness(femm, geom0, u1, Rfield1, dchi);
            M = mass(femm, geom0, u1, Rfield1, dchi);
            G = gyroscopic(femm, geom0, u1, Rfield1, v1, dchi);
            gathersysvec!(stepdchi, stepdchiv)
            @. TMPv = ((-1/(nb*dt^2))*stepdchiv+(1/(nb*dt^2))*dchipv)
            rhs .+= M*TMPv
            @. TMPv = ((-ng/nb/dt)*stepdchiv+(ng/nb/dt)*dchipv - vpv)
            rhs .+= G*TMPv;
            dchi = scattersysvec!(dchi, (K+(ng/nb/dt)*G+(1/(nb*dt^2))*M)\rhs); # Disp. incr
            u1.values[:] += (dchi.values[:,1:3])[:];   # increment displacement
            stepdchi.values[:] += dchi.values[:]
            v1.values[:] += (ng/nb/dt)*dchi.values[:];
            a1.values[:] += (1/nb/dt^2)*dchi.values[:];
            update_rotation_field!(Rfield1, dchi)
            if maximum(abs.(dchi.values[:])) < utol # convergence check
                break;
            end
            if (iter > maxit)# bailout for failed convergence
                error("Possible failed convergence");
            end
            iter += 1;
        end
        u0.values[:] = u1.values[:];       # update the displacement
        Rfield0.values[:] = Rfield1.values[:]; # update the rotations
        v0.values[:] = v1.values[:];       # update the velocities
        a0.values[:] = a1.values[:];       # update the accelerations

        if mod(step, 4) == 0
            @show step
            tshape1 = plot_solid(fens, fes; x = geom0.values, u = u1.values, R = Rfield1.values, facecolor = "rgb(125, 15, 15)");
            plots = cat(tbox,  tshape0,  tshape1; dims = 1)
            pl.plot.layout[:title] = "t=$t"
            react!(pl, plots, pl.plot.layout)
            sleep(0.5)
        end

        step=step+1;
    end

    return true
end # simo_vuquoc_animated

function simo_vuquoc_compare()
    # Parameters:
    # The  parameters given in the paper  are in terms of stiffness and mass coefficients.
    asquared = 12.0e-3
    E = 1.0e6 / asquared;
    a = sqrt(asquared)
    E * 1/12 * a^4
    nu = 0.0;# Poisson ratio
    rho = 1 / asquared;
    b = a; h = a; L = 10; # cross-sectional dimensions and length of each leg
    utol = 1e-3;
    maxit = 30
    Fmag = 50;
    dt = 0.05;
    tend = 30.0;
    ng = 1/2; nb = 1/4*(1/2+ng)^2;
    # Choose the mass formulation:
    mass_type=1;

    # Cross-sectional properties
    cs = CrossSectionRectangle(s -> b, s -> h, s -> [0.0, 0.0, 1.0])

    # Select the number of elements per leg.
    n=4;
    tolerance=L/n/100;
    members = []
    push!(members, frame_member([0 L 0;L L 0], n, cs))
    push!(members, frame_member([L L 0;L 0 0], n, cs))
    fens, fes = merge_members(members; tolerance = tolerance);
    @show E * fes.A[1], E * fes.I2[1], E/2 * fes.J[1], rho * fes.A[1], rho * fes.I2[1]

    # Material properties
    material = MatDeforElastIso(DeforModelRed3D, rho, E, nu, 0.0)

    # Construct the requisite fields, geometry and displacement
    # Initialize configuration variables
    geom0 = NodalField(fens.xyz)
    u0 = NodalField(zeros(size(fens.xyz,1), 3))
    Rfield0 = initial_Rfield(fens)
    dchi = NodalField(zeros(size(fens.xyz,1), 6))

    # Locate some nodes
    box = fill(0.0, 6)
    # The pinned end is
    initbox!(box, [0,L,0])
    pinnedn = selectnode(fens; box = box, inflate = tolerance)
    # The force is applied at the elbow
    initbox!(box, [L,L,0])
    elbown = selectnode(fens; box = box, inflate = tolerance)
    # Tip is tracked
    initbox!(box, [L,0,0])
    tipn = selectnode(fens; box = box, inflate = tolerance)

    # Apply EBC's
    for i in [1, 2, 3, 4, 5, 6]
        setebc!(dchi, pinnedn, true, i)
    end
    applyebc!(dchi)
    numberdofs!(dchi);
    dchi.dofnums

    # Additional fields
    stepdchi = deepcopy(dchi)
    u1 = deepcopy(u0)
    v0 = deepcopy(dchi) # need the numbering of the degrees of freedom
    v1 = deepcopy(dchi)
    a1 = deepcopy(dchi) # zero out the Acceleration
    a0 = deepcopy(dchi) # zero out the Acceleration
    Rfield1 = deepcopy(Rfield0)
    v0v = gathersysvec(v0)
    a0v = gathersysvec(a0)
    vpv = gathersysvec(v0);
    dchipv = gathersysvec(dchi);
    stepdchiv = gathersysvec(dchi);
    rhs = gathersysvec(dchi);
    TMPv = deepcopy(rhs)
    utol = 1e-13*dchi.nfreedofs;

    plots = cat(
            scatter(;x=Aus[:, 1], y=Aus[:, 2], mode="markers+lines", line_color = "rgb(255, 0, 0)", name="A");
            dims = 1)
        layout = Layout(width=700, height=500)
        pl = plot(plots, layout)
        display(pl)

    femm = FEMMCorotBeam(IntegDomain(fes, GaussRule(1, 2)), material)
    loadbdry = FESetP1(reshape(elbown, 1, 1))
    lfemm = FEMMBase(IntegDomain(loadbdry, PointRule()))

    t = 0.0; #
    step = 0;
    while (t <= tend)
        t = t + dt;
        # println("Time $(t)"); # pause
        # Initialization
        applyebc!(dchi) # Apply boundary conditions
        u1.values[:] = u0.values[:]; # guess
        Rfield1.values[:] = Rfield0.values[:]; # guess
        stepdchi.values[:] .= 0.0;# Total increment in current step
        a1.values[:] = -(1/nb/dt)*v0.values[:] -(1/2-nb)/nb*a0.values[:];
        v1.values[:] = v0.values[:] + dt*((1-ng)*a0.values[:] + ng*a1.values[:]);
        gathersysvec!(v0, v0v)
        gathersysvec!(a0, a0v)
        dchipv = dt*v0v + (dt^2/2*(1-2*nb))*a0v
        vpv = v0v +(dt*(1-ng))*a0v;
        fi = ForceIntensity(FFlt[0, 0, (t<2.0)*((t<1.0)*t+(t>1.0)*(2.0-t))*Fmag, 0, 0, 0]);

        iter = 1;
        while true
            F = distribloads(lfemm, geom0, dchi, fi, 3);
            Fr = restoringforce(femm, geom0, u1, Rfield1, dchi);       # Internal forces
            @. rhs = F + Fr;
            K = stiffness(femm, geom0, u1, Rfield1, dchi);
            M = mass(femm, geom0, u1, Rfield1, dchi);
            G = gyroscopic(femm, geom0, u1, Rfield1, v1, dchi);
            gathersysvec!(stepdchi, stepdchiv)
            @. TMPv = ((-1/(nb*dt^2))*stepdchiv+(1/(nb*dt^2))*dchipv)
            rhs .+= M*TMPv
            @. TMPv = ((-ng/nb/dt)*stepdchiv+(ng/nb/dt)*dchipv - vpv)
            rhs .+= G*TMPv;
            dchi = scattersysvec!(dchi, (K+(ng/nb/dt)*G+(1/(nb*dt^2))*M)\rhs); # Disp. incr
            u1.values[:] += (dchi.values[:,1:3])[:];   # increment displacement
            stepdchi.values[:] += dchi.values[:]
            v1.values[:] += (ng/nb/dt)*dchi.values[:];
            a1.values[:] += (1/nb/dt^2)*dchi.values[:];
            update_rotation_field!(Rfield1, dchi)
            if maximum(abs.(dchi.values[:])) < utol # convergence check
                break;
            end
            if (iter > maxit)# bailout for failed convergence
                error("Possible failed convergence");
            end
            iter += 1;
        end
        u0.values[:] = u1.values[:];       # update the displacement
        Rfield0.values[:] = Rfield1.values[:]; # update the rotations
        v0.values[:] = v1.values[:];       # update the velocities
        a0.values[:] = a1.values[:];       # update the accelerations

        if (mod(step,10)==0)
            plots = cat(pl.plot.data,
                scatter(;x=[t], y=[u1.values[elbown[1], 3]], mode="markers", line_color = "rgb(0, 0, 255)", showlegend=false);
                dims = 1)
            react!(pl, plots, pl.plot.layout)
            sleep(0.1)
        end

        step=step+1;
    end

    return true
end # simo_vuquoc_compare

# function argyr_swing_compare()
#     # Parameters:
#     E = 71240.0;#MPa
#     nu = 0.31;# Poisson ratio
#     rho = 5.0e-9;
#     b=10; h=30; L=240; #cross-sectional dimensions and length of each leg in millimeters
#     utol = 1e-3;
#     maxit = 30
#     Fmag = 500;
#     dt = 0.001;
#     tend = 0.5;
#     ng = 1/2; nb = 1/4*(1/2+ng)^2;
#     # Choose the mass formulation:
#     mass_type=1;

#     # Cross-sectional properties
#     cs = CrossSectionRectangle(s -> b, s -> h, s -> [0.0, 0.0, 1.0])

#     # Select the number of elements per leg.
#     n=4;
#     tolerance=L/n/100;
#     members = []
#     push!(members, frame_member([0 L 0;L L 0], n, cs))
#     push!(members, frame_member([L L 0;L 0 0], n, cs))
#     fens, fes = merge_members(members; tolerance = tolerance);

#     # Material properties
#     material = MatDeforElastIso(DeforModelRed3D, rho, E, nu, 0.0)

#     # Construct the requisite fields, geometry and displacement
#     # Initialize configuration variables
#     geom0 = NodalField(fens.xyz)
#     u0 = NodalField(zeros(size(fens.xyz,1), 3))
#     Rfield0 = initial_Rfield(fens)
#     dchi = NodalField(zeros(size(fens.xyz,1), 6))

#     # Locate some nodes
#     box = fill(0.0, 6)
#     # The pinned end is
#     initbox!(box, [0,L,0])
#     pinnedn = selectnode(fens; box = box, inflate = tolerance)
#     # The force is applied at the elbow
#     initbox!(box, [L,L,0])
#     elbown = selectnode(fens; box = box, inflate = tolerance)
#     # Tip is tracked
#     initbox!(box, [L,0,0])
#     tipn = selectnode(fens; box = box, inflate = tolerance)

#     # Apply EBC's
#     for i in [1, 2, 3]
#         setebc!(dchi, pinnedn, true, i)
#     end
#     applyebc!(dchi)
#     numberdofs!(dchi);
#     dchi.dofnums

#     # Additional fields
#     stepdchi = deepcopy(dchi)
#     u1 = deepcopy(u0)
#     v0 = deepcopy(dchi) # need the numbering of the degrees of freedom
#     v1 = deepcopy(dchi)
#     a1 = deepcopy(dchi) # zero out the Acceleration
#     a0 = deepcopy(dchi) # zero out the Acceleration
#     Rfield1 = deepcopy(Rfield0)
#     v0v = gathersysvec(v0)
#     a0v = gathersysvec(a0)
#     vpv = gathersysvec(v0);
#     dchipv = gathersysvec(dchi);
#     stepdchiv = gathersysvec(dchi);
#     rhs = gathersysvec(dchi);
#     TMPv = deepcopy(rhs)
#     utol = 1e-13*dchi.nfreedofs;

#     reft = range(0, 0.5, length = size(Aus,1))
#     plots = cat(
#         scatter(;x=reft, y=Aus[:,1], mode="lines", line_color = "rgb(255, 0, 0)", name="Tip-x"),
#         scatter(;x=reft, y=Aus[:,2], mode="lines", line_color = "rgb(0, 255, 0)", name="Tip-y"),
#         # scatter(;x=reft, y=Aus[:,3], mode="lines", line_color = "rgb(0, 0, 255)", name="Tip-z"),
#         dims = 1)
#     layout = Layout(width=700, height=500)
#     pl = plot(plots, layout)
#     display(pl)

#     femm = FEMMCorotBeam(IntegDomain(fes, GaussRule(1, 2)), material)
#     loadbdry = FESetP1(reshape(elbown, 1, 1))
#     lfemm = FEMMBase(IntegDomain(loadbdry, PointRule()))

#     t = 0.0; #
#     step = 0;
#     while (t <= tend)
#         t = t + dt;
#         # println("Time $(t)"); # pause
#         # Initialization
#         applyebc!(dchi) # Apply boundary conditions
#         u1.values[:] = u0.values[:]; # guess
#         Rfield1.values[:] = Rfield0.values[:]; # guess
#         stepdchi.values[:] .= 0.0;# Total increment in current step
#         a1.values[:] = -(1/nb/dt)*v0.values[:] -(1/2-nb)/nb*a0.values[:];
#         v1.values[:] = v0.values[:] + dt*((1-ng)*a0.values[:] + ng*a1.values[:]);
#         gathersysvec!(v0, v0v)
#         gathersysvec!(a0, a0v)
#         dchipv = dt*v0v + (dt^2/2*(1-2*nb))*a0v
#         vpv = v0v +(dt*(1-ng))*a0v;
#         fi = ForceIntensity(FFlt[0, 0, (t<0.2)*((t<0.1)*t+(t>0.1)*(0.2-t))*Fmag, 0, 0, 0]);

#         iter = 1;
#         while true
#             F = distribloads(lfemm, geom0, dchi, fi, 3);
#             Fr = restoringforce(femm, geom0, u1, Rfield1, dchi);       # Internal forces
#             @. rhs = F + Fr;
#             K = stiffness(femm, geom0, u1, Rfield1, dchi);
#             M = mass(femm, geom0, u1, Rfield1, dchi);
#             G = gyroscopic(femm, geom0, u1, Rfield1, v1, dchi);
#             gathersysvec!(stepdchi, stepdchiv)
#             @. TMPv = ((-1/(nb*dt^2))*stepdchiv+(1/(nb*dt^2))*dchipv)
#             rhs .+= M*TMPv
#             @. TMPv = ((-ng/nb/dt)*stepdchiv+(ng/nb/dt)*dchipv - vpv)
#             rhs .+= G*TMPv;
#             dchi = scattersysvec!(dchi, (K+(ng/nb/dt)*G+(1/(nb*dt^2))*M)\rhs); # Disp. incr
#             u1.values[:] += (dchi.values[:,1:3])[:];   # increment displacement
#             stepdchi.values[:] += dchi.values[:]
#             v1.values[:] += (ng/nb/dt)*dchi.values[:];
#             a1.values[:] += (1/nb/dt^2)*dchi.values[:];
#             update_rotation_field!(Rfield1, dchi)
#             if maximum(abs.(dchi.values[:])) < utol # convergence check
#                 break;
#             end
#             if (iter > maxit)# bailout for failed convergence
#                 error("Possible failed convergence");
#             end
#             iter += 1;
#         end
#         u0.values[:] = u1.values[:];       # update the displacement
#         Rfield0.values[:] = Rfield1.values[:]; # update the rotations
#         v0.values[:] = v1.values[:];       # update the velocities
#         a0.values[:] = a1.values[:];       # update the accelerations

#         if (mod(step,10)==0)
#             plots = cat(pl.plot.data,
#                 scatter(;x=[t], y=[u1.values[tipn[1], 1]], mode="markers", line_color = "rgb(255, 0, 0)", showlegend=false),
#                 scatter(;x=[t], y=[u1.values[tipn[1], 2]], mode="markers", line_color = "rgb(0, 255, 0)", showlegend=false),
#                 scatter(;x=[t], y=[u1.values[tipn[1], 3]], mode="markers", line_color = "rgb(0, 0, 255)", showlegend=false),
#                 scatter(;x=[t], y=[u1.values[elbown[1], 1]], mode="markers", line_color = "rgb(255, 0, 0)", showlegend=false),
#                 scatter(;x=[t], y=[u1.values[elbown[1], 2]], mode="markers", line_color = "rgb(0, 255, 0)", showlegend=false),
#                 scatter(;x=[t], y=[u1.values[elbown[1], 3]], mode="markers", line_color = "rgb(0, 0, 255)", showlegend=false);
#                 dims = 1)
#             react!(pl, plots, pl.plot.layout)
#             sleep(0.1)
#         end

#         step=step+1;
#     end

#     return true
# end # argyr_swing_animated

function allrun()
    println("#####################################################")
    println("# simo_vuquoc_animated ")
    simo_vuquoc_animated()
    println("#####################################################")
    println("# simo_vuquoc_compare ")
    simo_vuquoc_compare()
    # println("#####################################################")
    # println("# argyr_swing_compare ")
    # argyr_swing_compare()
    return true
end # function allrun

end # module

ext: jl


text:
to_assignments(xs::Vector, ys::Vector) = ["$x = $y" for (x, y) in zip(xs, ys)]
to_lines(xs::Vector{String}, indent::Int) = join(xs, "\n$(repeat("    ", indent))")
to_list(xs) = join(xs, ", ")

# ------------------------------------------------------------------------------

lpar(h::Int, d = "") = h == 1 ? "($(d)" : join(["⎛$(d)"; fill("⎜$(d)", h - 2); "⎝$(d)"], "\n")
rpar(h::Int, d = "") = h == 1 ? "$(d))" : join(["$(d)⎞"; fill("$(d)⎟", h - 2); "$(d)⎠"], "\n")

space(h::Int, sp = " ") = join(fill(sp, h), "\n")

function symstr(h::Int, symbol::String)
    a = fill("   ", h)
    a[Int(ceil(h/2))] = " $(symbol) "
    return join(a, "\n")
end

function mergestr(strings::String...)
    splits = split.(strings, "\n")
    rows = length(splits[1])
    cols = length(splits)
    matr = reshape(collect(Iterators.flatten(splits)), rows, cols)
    join([join(matr[i,:]) for i in 1:size(matr, 1)], "\n")
end

# ------------------------------------------------------------------------------

function _print_recsystem(io::IO, vars, body, init)
    h = size(body, 1)
    lstr, rstr = lpar(h), rpar(h, " ")
    eq = symstr(h, "=")

    # lc = "\u2099"
    # lp, rp, plus = "", "", "\u208A"
    # zero, one = "\u2080", "\u2081"

    lc = "n"
    lp, rp, plus = "(", ")", "+"
    zero, one = "0", "1"

    vars0 = Base.replace(sprint(Base.print_matrix, map(x->string(x)*"$lp$zero$rp", vars)), "\""=>"")
    vars1 = Base.replace(sprint(Base.print_matrix, map(x->string(x)*"$lp$lc$rp", vars)), "\""=>"")
    vars2 = Base.replace(sprint(Base.print_matrix, map(x->string(x)*"$lp$lc$plus$one$rp", vars)), "\""=>"")
    body = sprint(Base.print_matrix, body)
    init = sprint(Base.print_matrix, init)

    lhs1 = (lstr, vars2, rstr)
    rhs1 = (lstr, body, rstr, lstr, vars1, rstr)

    lhs2 = (lstr, vars0, rstr)
    rhs2 = (lstr, init, rstr)
    str = mergestr(space(h, "\t"), lhs1..., eq, rhs1..., space(h, "\t"), lhs2..., eq, rhs2...)
    print(io, str)
end


text:
function PrintStats(x::ASTContext)
    @check_ptrs x
    return clang_ASTContext_PrintStats(x)
end

function getTypeSize(x::ASTContext, ty_ptr::CXType_)
    @check_ptrs x
    return clang_ASTContext_getTypeSize(x, ty_ptr)
end
getTypeSize(x::ASTContext, ty::AbstractQualType) = getTypeSize(x, get_type_ptr(ty))

function getTypeDeclType(x::ASTContext, decl::TypeDecl, prev::TypeDecl=TypeDecl(C_NULL))
    @check_ptrs x decl
    return QualType(clang_ASTContext_getTypeDeclType(x, decl, prev))
end

function getRecordType(x::ASTContext, decl::RecordDecl)
    @check_ptrs x decl
    return QualType(clang_ASTContext_getRecordType(x, decl))
end

getTypeDeclType(x::ASTContext, decl::NamedDecl) = getTypeDeclType(x, TypeDecl(decl.ptr))
getTypeDeclType(x::ASTContext, decl::AbstractTypeDecl) = getTypeDeclType(x, TypeDecl(decl.ptr))
getTypeDeclType(x::ASTContext, decl::AbstractRecordDecl) = getRecordType(x, RecordDecl(decl.ptr))

function getPointerType(x::ASTContext, ty::AbstractQualType)
    @check_ptrs x
    return QualType(clang_ASTContext_getPointerType(x, ty))
end

function getLValueReferenceType(x::ASTContext, ty::AbstractQualType)
    @check_ptrs x
    return QualType(clang_ASTContext_getLValueReferenceType(x, ty))
end

function getRValueReferenceType(x::ASTContext, ty::AbstractQualType)
    @check_ptrs x
    return QualType(clang_ASTContext_getRValueReferenceType(x, ty))
end

function getMemberPointerType(x::ASTContext, ty::AbstractQualType, class_ptr::CXType_)
    @check_ptrs x
    return QualType(clang_ASTContext_getMemberPointerType(x, ty, class_ptr))
end
getMemberPointerType(x::ASTContext, ty::AbstractQualType, class::AbstractQualType) = getMemberPointerType(x, ty, get_type_ptr(class))

function getIdents(x::ASTContext)
    @check_ptrs x
    return IdentifierTable(clang_ASTContext_getIdents(x))
end

VoidTy(ctx::ASTContext) = VoidTy(clang_ASTContext_VoidTy_getAsQualType(ctx))
BoolTy(ctx::ASTContext) = BoolTy(clang_ASTContext_BoolTy_getAsQualType(ctx))
CharTy(ctx::ASTContext) = CharTy(clang_ASTContext_CharTy_getAsQualType(ctx))
WCharTy(ctx::ASTContext) = WCharTy(clang_ASTContext_WCharTy_getAsQualType(ctx))  # [C++ 3.9.1p5].
WideCharTy(ctx::ASTContext) = WideCharTy(clang_ASTContext_WideCharTy_getAsQualType(ctx))
WIntTy(ctx::ASTContext) = WIntTy(clang_ASTContext_WIntTy_getAsQualType(ctx))  # [C99 7.24.1], integer type unchanged by default promotions.
Char8Ty(ctx::ASTContext) = Char8Ty(clang_ASTContext_Char8Ty_getAsQualType(ctx))  # [C++20 proposal]
Char16Ty(ctx::ASTContext) = Char16Ty(clang_ASTContext_Char16Ty_getAsQualType(ctx))  # [C++0x 3.9.1p5], integer type in C99.
Char32Ty(ctx::ASTContext) = Char32Ty(clang_ASTContext_Char32Ty_getAsQualType(ctx))  # [C++0x 3.9.1p5], integer type in C99.
SignedCharTy(ctx::ASTContext) = SignedCharTy(clang_ASTContext_SignedCharTy_getAsQualType(ctx))
ShortTy(ctx::ASTContext) = ShortTy(clang_ASTContext_ShortTy_getAsQualType(ctx))
IntTy(ctx::ASTContext) = IntTy(clang_ASTContext_IntTy_getAsQualType(ctx))
LongTy(ctx::ASTContext) = LongTy(clang_ASTContext_LongTy_getAsQualType(ctx))
LongLongTy(ctx::ASTContext) = LongLongTy(clang_ASTContext_LongLongTy_getAsQualType(ctx))
Int128Ty(ctx::ASTContext) = Int128Ty(clang_ASTContext_Int128Ty_getAsQualType(ctx))
UnsignedCharTy(ctx::ASTContext) = UnsignedCharTy(clang_ASTContext_UnsignedCharTy_getAsQualType(ctx))
UnsignedShortTy(ctx::ASTContext) = UnsignedShortTy(clang_ASTContext_UnsignedShortTy_getAsQualType(ctx))
UnsignedIntTy(ctx::ASTContext) = UnsignedIntTy(clang_ASTContext_UnsignedIntTy_getAsQualType(ctx))
UnsignedLongTy(ctx::ASTContext) = UnsignedLongTy(clang_ASTContext_UnsignedLongTy_getAsQualType(ctx))
UnsignedLongLongTy(ctx::ASTContext) = UnsignedLongLongTy(clang_ASTContext_UnsignedLongLongTy_getAsQualType(ctx))
UnsignedInt128Ty(ctx::ASTContext) = UnsignedInt128Ty(clang_ASTContext_UnsignedInt128Ty_getAsQualType(ctx))
FloatTy(ctx::ASTContext) = FloatTy(clang_ASTContext_FloatTy_getAsQualType(ctx))
DoubleTy(ctx::ASTContext) = DoubleTy(clang_ASTContext_DoubleTy_getAsQualType(ctx))
LongDoubleTy(ctx::ASTContext) = LongDoubleTy(clang_ASTContext_LongDoubleTy_getAsQualType(ctx))
Float128Ty(ctx::ASTContext) = Float128Ty(clang_ASTContext_Float128Ty_getAsQualType(ctx))
HalfTy(ctx::ASTContext) = HalfTy(clang_ASTContext_HalfTy_getAsQualType(ctx))  # [OpenCL 6.1.1.1], ARM NEON
BFloat16Ty(ctx::ASTContext) = BFloat16Ty(clang_ASTContext_BFloat16Ty_getAsQualType(ctx))
Float16Ty(ctx::ASTContext) = Float16Ty(clang_ASTContext_Float16Ty_getAsQualType(ctx))  # C11 extension ISO/IEC TS 18661-3
FloatComplexTy(ctx::ASTContext) = FloatComplexTy(clang_ASTContext_FloatComplexTy_getAsQualType(ctx))
DoubleComplexTy(ctx::ASTContext) = DoubleComplexTy(clang_ASTContext_DoubleComplexTy_getAsQualType(ctx))
LongDoubleComplexTy(ctx::ASTContext) = LongDoubleComplexTy(clang_ASTContext_LongDoubleComplexTy_getAsQualType(ctx))
Float128ComplexTy(ctx::ASTContext) = Float128ComplexTy(clang_ASTContext_Float128ComplexTy_getAsQualType(ctx))
VoidPtrTy(ctx::ASTContext) = VoidPtrTy(clang_ASTContext_VoidPtrTy_getAsQualType(ctx))
NullPtrTy(ctx::ASTContext) = NullPtrTy(clang_ASTContext_NullPtrTy_getAsQualType(ctx))  # C++11 nullptr


text:
using SimpleSMC
import SequentialMonteCarlo.SMCModel
using RNGPool
using SMCExamples.LinearGaussian
using Random

include("test.jl")

setRNGs(0)

model, theta, ys, ko = LinearGaussian.defaultLGModel(10)
println(ko.logZhats)

numParticles = 1024
numTrials = 2

## just run the algorithm a few times
testSMC(model, numParticles, numTrials)

testCSMC(model, numParticles, numTrials)

lM = LinearGaussian.makelM(theta)
testCSMCBS(model, lM, numParticles, numTrials)


text:
module Octopart

	export get_offers, get_offers_for_FPGAs

	###########################################################################
	# Params.

	#TODO Read it from some cfg file, so everybody have it's own API key.
	const API_KEY="48a64d25"

	###########################################################################

	using Requests
	using JSON
	using DataFrames

	function get_offers(
		search_mpns::Vector,
		needed_quantity::Integer,
		currencies_to_eur::Dict
	)

		all_offers = DataFrame(
			search_mpn = [],
			mpn = [],
			sku = [],
			seller = [],
			stock = [],
			moq = [],
			price = []
		)

		for i in 1:20:length(search_mpns)
			s = i:min(i+19, length(search_mpns))
			chunk = search_mpns[s]

			queries = map(
				(s) -> Dict("mpn" => s, "limit" => 20),
				chunk
			)

			res = Requests.get(
				"http://octopart.com/api/v3/parts/match";
				query = Dict(
					"apikey" => API_KEY,
					"pretty_print" => true,
					"queries" => JSON.json(queries)
				)
			)
			if res.status != 200
				critical_error_response = JSON.parse(String(res.data))
				error(
					"Failed to obtain offers from octopart: ",
					critical_error_response["message"]
				)
			else
				part_match_response = JSON.parse(String(res.data))

				#@show part_match_response["msec"]

				results = part_match_response["results"]

				for (search_mpn, result) in zip(chunk, results)

					for item in result["items"]
						mpn = item["mpn"]

						offers = item["offers"]
						for offer in offers
							prices = offer["prices"]

							if length(prices) == 0
								continue
							end

							if !("EUR" in keys(prices))
								currency = first(keys(prices))
							else
								currency = "EUR"
							end
							cur_mul = currencies_to_eur[currency]

							om = offer["order_multiple"]
							if om != nothing && om != 1
								if om > needed_quantity
									continue
								else
									warn(
										"One offer for ", mpn,
										" have order multiple of ", om
									)
								end
							end

							moq = offer["moq"]
							if moq != nothing
								if moq > needed_quantity
									continue
								end
							else
								moq = -1
							end

							p = nothing
							for bq_price in prices[currency]
								if bq_price[1] > needed_quantity
									break
								end

								p = bq_price[2]
							end
							if p == nothing
								continue
							end
							price = parse(Float64, p) * cur_mul

							sku = offer["sku"]
							seller = offer["seller"]["name"]
							stock = offer["in_stock_quantity"]

							push!(
								all_offers,
								[search_mpn, mpn, sku, seller,
									stock, moq, price]
							)
						end
					end
				end
			end

		end

		return all_offers
	end

	using Iterators

	function get_offers_for_FPGAs(
		FPGA_families,
		needed_quantity::Integer,
		currencies_to_eur::Dict,
		simple_speed_grades::Bool
	)
		search_mpns = []
		dev_speed_pack = []
		family_names = []
		for (fn, f) in FPGA_families
			dpc = f["dev_pack_pins"][:dev_pack]
			if simple_speed_grades
				sg = f["simple_speed_grades"]
			else
				sg = f["speed_grades"]
			end
			# Cartesian product.
			cp = product(dpc, sg)
			# Concat to MPN and collect.
			m = collect(map((t) -> t[1][1] * t[2] * t[1][2] * "*", cp))
			append!(search_mpns, m)

			a = collect(map((t) -> (t[1][1], t[2], t[1][2]), cp))
			append!(dev_speed_pack, a)

			append!(family_names, fill(fn, length(m)))
		end

		o = get_offers(
			search_mpns,
			needed_quantity,
			currencies_to_eur
		)

		# All found.
		i = [findfirst(search_mpns, sm) for sm in o[:search_mpn]]
		assert(all(i != 0))
		offers = DataFrame()
		offers[:family] = family_names[i]
		dsp = dev_speed_pack[i]
		offers[:device] = [t[1] for t in dsp]
		offers[:speed_grade] = [t[2] for t in dsp]
		offers[:package] = [t[3] for t in dsp]
		offers[:mpn] = o[:mpn]
		offers[:sku] = o[:sku]
		offers[:seller] = o[:seller]
		offers[:stock] = o[:stock]
		offers[:moq] = o[:moq]
		offers[:stock_vs_need] = o[:stock] - needed_quantity
		offers[:price] = o[:price]

		return offers
	end

	###########################################################################

end # Octopart


text:
using Sigma
using AbstractDomains
using Base.Test
x = uniform(0,1)
y = uniform(0,1)
@test isequal(prob(x+y > - 1.0), Interval(1.0,1.0))


text:
function countuniquesorted(arr::AbstractVector{T}) where T
    d = Dict{T,Int}()
    for val in arr
        count = get(d, val, 0)
        d[val] = count + 1
    end

    N = length(d)
    vals = Array{T,1}()
    sizehint!(vals,N)
    counts = Array{Int,1}()
    sizehint!(counts,N)
    for (val,count) in sort(collect(d))
        push!(vals,val)
        push!(counts,count)
    end

    return vals,counts
end

function groupptrs!(
  headptrs::AbstractVector{Ti},
  tailptrs::AbstractVector{Ti},
  arr::AbstractVector{T} where T<:Real,
  ) where {Ti<:Integer}

  N = length(arr)
  @assert N == length(headptrs) == length(tailptrs)

  fill!(headptrs, 0)
  fill!(tailptrs, 0)

  headptr = 1
  flag = arr[1]
  headptrs[1] = headptr

  for i in 2:N
    @assert (arr[i-1] <= arr[i]) "array is not sorted at i=$i" #TODO this condition is not needed
    if flag != arr[i]
      flag = arr[i]
      headptr = i
    end
    headptrs[i] = headptr
  end

  headptr = headptrs[end]
  tailptr = N

  tailptrs[end] = tailptr
  for i in (N-1):(-1):1
    if headptr != headptrs[i]
      headptr = headptrs[i]
      tailptr = i
    end
    tailptrs[i] = tailptr
  end

  headptrs, tailptrs
end

groupptrs(arr::AbstractVector{T} where T<:Real) =
  groupptrs!(zeros(Int32, length(arr)), zeros(Int32, length(arr)), arr)


text:
module RegressionDiscontinuity

using Reexport

import Base: size, getindex, getproperty, propertynames, show

using DataFrames
@reexport using Distributions
using FastGaussQuadrature
using Feather
using Intervals
using JuMP
using LinearAlgebra
using GLM
using OffsetArrays
using QuadGK
using RecipesBase
using Roots
using Setfield
using Statistics
import Statistics: var
@reexport using StatsBase
import StatsBase: fit, weights, nobs
using StatsModels
using Tables

using UnPack

include("running_variable.jl")
include("load_example_data.jl")
include("kernels.jl")
include("imbens_kalyanaraman.jl")
include("local_linear.jl")
include("minmax_optimal.jl")
include("density_test.jl")

export RunningVariable,
    RDData,
    Treated,
    Untreated,
    load_rdd_data,
    Rectangular,
    bandwidth,
    ImbensKalyanaraman,
    linearweights,
    EickerHuberWhite,
    Homoskedastic,
    NaiveLocalLinearRD,
    ImbensWagerOptRD,
    estimate_second_deriv_bound,
    McCraryTest
end


text:
@with_kw struct FastGrid
    resolution::Float64 = 1.0
    #tight::Bool         = false
end

# This is the main function
function solve(solver::FastGrid, problem::Problem) #original
    result = true
    delta = solver.resolution

    center = problem.input.center
    radius = problem.input.radius[1]

    (W, b) = (problem.network.layers[1].weights, problem.network.layers[1].bias)

    input = forward_affine_map(solver, W, b, problem.input)

    lower, upper = low(input), high(input)
    n_hypers_per_dim = BigInt.(max.(ceil.(Int, (upper-lower) / delta), 1))

    k_1 = size(W, 1)
    k_0 = size(W, 2)
#=
    C_i = vcat( Array(Diagonal(ones(k_0))),
                Array(Diagonal(ones(k_0))),
                W,
                -W
    )
=#
    C = vcat(W, -W)
    d = zeros(2k_1)

    #count = BigInt(0)
    #println("All: " * string(prod(n_hypers_per_dim)))

    # preallocate work arrays
    local_lower, local_upper, CI = similar(lower), similar(lower), similar(lower)
    for i in 1:prod(n_hypers_per_dim)
        n = i
        for j in firstindex(CI):lastindex(CI)
            n, CI[j] = fldmod1(n, n_hypers_per_dim[j])
        end
        @. local_lower = lower + delta * (CI - 1)
        @. local_upper = min(local_lower + delta, upper)
        hyper = Hyperrectangle(low = local_lower, high = local_upper)
#=
        d_i = vcat( high(problem.input),
                    -low(problem.input),
                    local_upper - b,
                    b - local_lower
        )
=#
        d = vcat(local_upper - b, b - local_lower)

        inter = intersection(problem.input, HPolyhedron(C, d))

        if isempty(inter) == false
            reach = forward_network(solver, problem.network, hyper)
            #count += 1
            if !issubset(reach, problem.output)
                result = false
            end
        end
    end

    #println("FastGrid($(delta)): " * string(count))

    if result
        return BasicResult(:holds)
    end
    return BasicResult(:violated)
end

function forward_network(solver::FastGrid, nnet::Network, input::Hyperrectangle)
    layers = nnet.layers
    act = layers[1].activation
    reach = Hyperrectangle(low = act.(low(input)), high = act.(high(input)))

    for i in 2:length(layers)
        reach = forward_layer(solver, layers[i], reach)
    end
    return reach
end

# This function is called by forward_network
function forward_layer(solver::FastGrid, L::Layer, input::Hyperrectangle)
    (W, b, act) = (L.weights, L.bias, L.activation)
    center = zeros(size(W, 1))
    gamma  = zeros(size(W, 1))
    for j in 1:size(W, 1)
        node = Node(W[j,:], b[j], act)
        center[j], gamma[j] = forward_node(solver, node, input)
    end
    return Hyperrectangle(center, gamma)
end

function forward_affine_map(solver::FastGrid, W::Matrix, b::Vector, input::Hyperrectangle)
    center = W * input.center + b
    radius = abs.(W) * input.radius
    return Hyperrectangle(center, radius)
end

function forward_node(solver::FastGrid, node::Node, input::Hyperrectangle)
    output    = node.w' * input.center + node.b
    deviation = sum(abs.(node.w) .* input.radius)
    β    = node.act(output)  # TODO expert suggestion for variable name. beta? β? O? x?
    βmax = node.act(output + deviation)
    βmin = node.act(output - deviation)
    return ((βmax + βmin)/2, (βmax - βmin)/2)
    #if solver.tight
        #return ((βmax + βmin)/2, (βmax - βmin)/2)
    #else
        #return (β, max(abs(βmax - β), abs(βmin - β)))
    #end
end


text:
cconvert(T::Type, x::VulkanStruct) = x
cconvert(T::Type{<:Ptr}, x::AbstractVector{<:VulkanStruct{false}}) = getproperty.(x, :vks)
cconvert(T::Type{<:Ptr}, x::AbstractVector{<:VulkanStruct{true}}) = (x, getproperty.(x, :vks))
cconvert(T::Type{<:Ptr}, x::VulkanStruct{false}) = Ref(x.vks)
cconvert(T::Type{<:Ptr}, x::VulkanStruct{true}) = (x, Ref(x.vks))
cconvert(T::Type{<:Ptr}, x::Handle) = x

convert(T::Type{<:Ptr}, x::Handle) = x.vks

unsafe_convert(T::Type, x::VulkanStruct) = x.vks
unsafe_convert(T::Type, x::Tuple{<:VulkanStruct{true}, <:Ref}) = unsafe_convert(T, last(x))
unsafe_convert(T::Type, x::Tuple{<:AbstractVector{<:VulkanStruct{true}}, <:Any}) = unsafe_convert(T, last(x))

"""
    `pointer_length(val)`

Return the length `val` considering it as an array.

Differ from `Base.length` in that `pointer_length(C_NULL) == 0` and that a `RefArray`s return the length of their roots.
"""
function pointer_length end

pointer_length(arr::Ptr{Nothing}) = 0
pointer_length(arr::AbstractArray) = length(arr)
pointer_length(arr::RefArray) = length(arr.roots)
pointer_length(arr::Tuple{<:Any,<:Any}) = length(first(arr))


text:
module QuantumStatePlots

using Plots
using QuantumStateBase

export
    DEFAULT_SIZE,
    Heatmap,
    Contour,
    Surface,
    plot_wigner,
    plot_ρ,
    plot_all

DEFAULT_SIZE = (1100, 900)

abstract type PlotMethod end

"""
    Heatmap <: PlotMethod

Heatmap representation for Wigner function of quantum state.
"""
struct Heatmap <: PlotMethod end

"""
    Contour <: PlotMethod

Contour representation for Wigner function of quantum state.
"""
struct Contour <: PlotMethod end

"""
    Surface <: PlotMethod

Surface representation for Wigner function of quantum state.
"""
struct Surface <: PlotMethod end

"""
    plot_wigner(ws::WignerSurface, Heatmap; size=$DEFAULT_SIZE, file_path=nothing)

Plot Wigner function of given quantum state in `Heatmap` representation.

# Examples
```julia-repl
julia> using QuantumStateBase
julia> state = CoherentState(α(5., π/4));
julia> wf = WignerFunction(-10:0.1:10, -10:0.1:10);
julia> plot_wigner(wf(state), Surface)
```
"""
function plot_wigner(
    ws::WignerSurface, ::Type{Heatmap};
    size=DEFAULT_SIZE,
    file_path=nothing
)
    !isnothing(size) && (gr(size=size) isa Plots.GRBackend) || gr()

    lim = maximum(abs.(ws.𝐰_surface))
    p = Plots.heatmap(
        ws.x_range, ws.p_range, ws.𝐰_surface',
        title="Wigner Function",
        xlabel="X",
        ylabel="P",
        clim=(-lim, lim),
        c=:coolwarm,
    )

    isnothing(file_path) || savefig(p, file_path)

    return p
end

"""
    plot_wigner(ws::WignerSurface, Contour; size=$DEFAULT_SIZE, file_path=nothing)

Plot Wigner function of given quantum state in `Contour` representation.

# Examples
```julia-repl
julia> using QuantumStateBase
julia> state = CoherentState(α(5., π/4));
julia> wf = WignerFunction(-10:0.1:10, -10:0.1:10);
julia> plot_wigner(wf(state), Contour)
```
"""
function plot_wigner(
    ws::WignerSurface, ::Type{Contour};
    levels=20,
    size=DEFAULT_SIZE,
    file_path=nothing
)
    !isnothing(size) && (gr(size=size) isa Plots.GRBackend) || gr()

    lim = maximum(abs.(ws.𝐰_surface))
    p = Plots.contour(
        ws.x_range, ws.p_range, ws.𝐰_surface',
        title="Wigner Function",
        xlabel="X",
        ylabel="P",
        clim=(-lim, lim),
        fill=true,
        levels=levels,
        c=:coolwarm,
    )

    isnothing(file_path) || savefig(p, file_path)

    return p
end

"""
    plot_wigner(ws::WignerSurface, Surface; size=$DEFAULT_SIZE, file_path=nothing)

Plot Wigner function of given quantum state in `Surface` representation.

# Examples
```julia-repl
julia> using QuantumStateBase
julia> state = CoherentState(α(5., π/4));
julia> wf = WignerFunction(-10:0.1:10, -10:0.1:10);
julia> plot_wigner(wf(state), Surface)
```
"""
function plot_wigner(
    ws::WignerSurface, ::Type{Surface};
    size=DEFAULT_SIZE,
    file_path=nothing
)
    !isnothing(size) && (gr(size=size) isa Plots.GRBackend) || gr()

    lim = maximum(abs.(ws.𝐰_surface))
    p = Plots.surface(
        ws.x_range, ws.p_range, ws.𝐰_surface',
        title="Wigner Function",
        xlabel="X",
        ylabel="P",
        clim=(-lim, lim),
        zlim=(-lim, lim),
        c=:coolwarm,
        # fillalpha=0.99, # this cause Pluto crash, work with DisplayAs
        camera=(40, 30),
    )

    isnothing(file_path) || savefig(p, file_path)

    return p
end

"""
    plot_ρ(state::AbstractState; state_n::Integer=0, size=$DEFAULT_SIZE, file_path=nothing)

Plot density matrux of given quantum state.

* `state_n=n`: truncate to `n`th photon number

# Examples
```julia-repl
julia> using QuantumStateBase
julia> state = CoherentState(α(5., π/4));
julia> plot_ρ(state)
```
"""
function plot_ρ(
    state::AbstractState;
    state_n::Integer=0,
    size=DEFAULT_SIZE,
    file_path=nothing
)
    ρᵣ = real(𝛒(state))
    if state_n != 0
        ρᵣ = ρᵣ[1:state_n+1, 1:state_n+1]
    else
        state_n = Base.size(ρᵣ)[1] - 1
    end

    !isnothing(size) && (gr(size=size) isa Plots.GRBackend) || gr()

    lim = maximum(ρᵣ)
    p = Plots.heatmap(
        0:state_n, 0:state_n, ρᵣ,
        title="Density Matrix (Real part)",
        xlabel="m",
        ylabel="n",
        c=:coolwarm,
        clim=(-lim, lim)
    )

    isnothing(file_path) || savefig(p, file_path)

    return p
end

"""
    plot_all(
        ws::WignerSurface, state::AbstractState;
        state_n::Integer=0, size=$DEFAULT_SIZE, file_path=nothing
    )

Plot Wifner function and density matrux of given quantum state
in both `Surface` and `Contour` representation.

* `state_n=n`: truncate to `n`th photon number

# Examples
```julia-repl
julia> using QuantumStateBase
julia> state = CoherentState(α(5., π/4));
julia> wf = WignerFunction(-10:0.1:10, -10:0.1:10);
julia> plot_all(wf(state), state)
```
"""
function plot_all(
    ws::WignerSurface, state::AbstractState;
    state_n=0,
    levels=20,
    size=DEFAULT_SIZE,
    file_path=nothing
)
    !isnothing(size) && (gr(size=size) isa Plots.GRBackend) || gr()

    l = @layout [
		a{0.55h}
		grid(1, 2)
	]
    p = plot(
        plot_wigner(ws, Surface, size=nothing),
        plot_wigner(ws, Contour, size=nothing, levels=levels),
        plot_ρ(state, state_n=state_n, size=nothing),
        layout=l
    )

    isnothing(file_path) || savefig(p, file_path)

    return p
end

end


text:
# Copyright © 2021 Alexander L. Hayes
# Apache 2.0 License

import Base.Filesystem: expanduser, joinpath, mkpath, rm

"""
    get_data_home(data_home::Union{String, Nothing} = nothing)::String

Get the data home directory.
"""
function get_data_home(data_home::Union{String, Nothing} = nothing)::String

    _path = ""

    if data_home == nothing
        try
            _path = expanduser(ENV["RELATIONAL_DATASETS"])
        catch e
            _path = expanduser(
                joinpath("~", "relational_datasets")
            )
        end
    end

    return mkpath(_path)
end

function clear_data_home(data_home::Union{String, Nothing} = nothing)
    path = get_data_home(data_home)
    rm(path, force=true, recursive=true)
    return nothing
end


text:
include("euler/euler.jl")
using .BigIntegers: BigInt
using .Calculus: fibonacci_numbers

function compute(n::Int, m::Int = 1_000_000_007)
    inverted_modulo = invmod(2, m)
    result = BigInteger(-1)
    for f ∈ fibonacci_numbers(n, BigInteger)
        a, b = f % 9 + 2, f ÷ 9
        result += ((a * (a - 1) + 10) * powermod(BigInteger(10), b, m) - 2(a + 9b + 4)) * inverted_modulo % m
    end
    return result % m
end

max_stars_repo_name: Dynortice/Project-Euler

max_issues_repo_name: Dynortice/Project-Euler

max_forks_repo_name: Dynortice/Project-Euler


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

module ExpressRouteCircuitClient

using Random
using Dates
using HTTP
using Swagger
import Swagger: field_name, property_type, hasproperty, validate_property, SwaggerApi, SwaggerModel
import Base: convert, propertynames

include("modelincludes.jl")

include("api_ExpressRouteCircuitArpTableApi.jl")
include("api_ExpressRouteCircuitAuthorizationsApi.jl")
include("api_ExpressRouteCircuitPeeringsApi.jl")
include("api_ExpressRouteCircuitRoutesTableApi.jl")
include("api_ExpressRouteCircuitRoutesTableSummaryApi.jl")
include("api_ExpressRouteCircuitStatsApi.jl")
include("api_ExpressRouteCircuitsApi.jl")
include("api_ExpressRouteServiceProvidersApi.jl")

# export models
export convert, AuthorizationListResult
export convert, AuthorizationPropertiesFormat
export convert, ExpressRouteCircuitArpTable
export convert, ExpressRouteCircuitListResult
export convert, ExpressRouteCircuitPeeringConfig
export convert, ExpressRouteCircuitPeeringListResult
export convert, ExpressRouteCircuitPeeringPropertiesFormat
export convert, ExpressRouteCircuitPropertiesFormat
export convert, ExpressRouteCircuitRoutesTable
export convert, ExpressRouteCircuitRoutesTableSummary
export convert, ExpressRouteCircuitServiceProviderProperties
export convert, ExpressRouteCircuitSku
export convert, ExpressRouteCircuitStats
export convert, ExpressRouteCircuitsArpTableListResult
export convert, ExpressRouteCircuitsRoutesTableListResult
export convert, ExpressRouteCircuitsRoutesTableSummaryListResult
export convert, ExpressRouteServiceProviderBandwidthsOffered
export convert, ExpressRouteServiceProviderListResult
export convert, ExpressRouteServiceProviderPropertiesFormat
export convert, Ipv6ExpressRouteCircuitPeeringConfig
export convert, Resource
export convert, SubResource
export convert, ExpressRouteCircuit
export convert, ExpressRouteCircuitAuthorization
export convert, ExpressRouteCircuitPeering
export convert, ExpressRouteServiceProvider
export convert, RouteFilter

# export operations
export convert, ExpressRouteCircuitArpTableApi, ExpressRouteCircuitAuthorizationsApi, ExpressRouteCircuitPeeringsApi, ExpressRouteCircuitRoutesTableApi, ExpressRouteCircuitRoutesTableSummaryApi, ExpressRouteCircuitStatsApi, ExpressRouteCircuitsApi, ExpressRouteServiceProvidersApi

export check_required, field_name, property_type, hasproperty, propertynames, validate_property, convert

end


text:
module RocketOverrideOperatorTest

using Test
using Rocket

include("../test_helpers.jl")

@testset "operator: override()" begin

    println("Testing: operator override()")

    run_proxyshowcheck("Override", override(OverrideHandler(Int, nothing)))

    global_handler1 = OverrideHandler(Int, nothing)
    global_handler2 = OverrideHandler(Int, nothing)

    run_testset([
        (
            source      = from(1:5) |> override(OverrideHandler(Char, 'a')),
            values      = @ts([ 'a', 'a', 'a', 'a', 'a', c ]),
            source_type = Union{Int, Char}
        ),
        (
            source      = from(1:5) |> tap(i -> Rocket.setvalue!(global_handler1, i ^ 2)) |> override(global_handler1),
            values      = @ts([ 1, 4, 9, 16, 25, c ]),
            source_type = Int
        ),
        (
            source      = from(1:5) |> override(global_handler2) |> tap(i -> Rocket.setvalue!(global_handler2, i)),
            values      = @ts([ 1, 1, 1, 1, 1, c ]),
            source_type = Int
        ),
        (
            source      = completed(Int) |> override(OverrideHandler(Char, 'a')),
            values      = @ts(c),
            source_type = Union{Int, Char}
        ),
        (
            source      = faulted(String, "e") |> override(OverrideHandler(Char, 'a')),
            values      = @ts(e("e")),
            source_type = Union{String, Char}
        ),
        (
            source = never() |> override(OverrideHandler(Char, 'a')),
            values = @ts()
        )
    ])

end

end


text:
using ReachabilityAnalysis
Dict([
    "opts" => Dict(:T=>20.0, :alg=>GLGM06(δ=0.005)),
    "X0"   => Hyperrectangle(low=[fill(0.0002, 10); zeros(14); -0.0001; zeros(23)],
                            high=[fill(0.00025, 10); zeros(14); 0.0001; zeros(23)]),
    "info" => Dict(
        "name" => "building",
        "dim" => 48,
        "linear" => true,
        "hybrid" => false
    )
])

max_stars_repo_name: mforets/ReachabilityLibrary

max_issues_repo_name: mforets/ReachabilityLibrary

max_forks_repo_name: mforets/ReachabilityLibrary


text:
#=

@authors: Shunsuke Hori

=#

using QuantEcon

# == Parameters == #
a0 = 10.0
a1 = 2.0
β = 0.96
γ = 12.0

# == In LQ form == #

A  = eye(3)
B1 = [0.0, 1.0, 0.0]
B2 = [0.0, 0.0, 1.0]

R1 = [      0.0   -a0 / 2.0          0.0;
      -a0 / 2.0          a1     a1 / 2.0;
            0.0    a1 / 2.0          0.0]

R2 = [      0.0          0.0      -a0 / 2.0;
            0.0          0.0       a1 / 2.0;
      -a0 / 2.0     a1 / 2.0             a1]

Q1 = Q2 = γ
S1 = S2 = W1 = W2 = M1 = M2 = 0.0
# == Solve using QE's nnash function == #
F1, F2, P1, P2 = nnash(A, B1, B2, R1, R2, Q1, Q2, S1, S2, W1, W2, M1, M2,
                          beta=β)

# == Display policies == #
println("Computed policies for firm 1 and firm 2:")
println("F1 = $F1")
println("F2 = $F2")


text:
using StochasticDiffEq, DiffEqDevTools
using DiffEqProblemLibrary.SDEProblemLibrary: importsdeproblems; importsdeproblems()
using Test, Random

const TEST_PLOT = false
const LONGER_TESTS = false

if haskey(ENV,"GROUP")
    group = ENV["GROUP"]
else
    group = "All"
end

is_APPVEYOR = ( Sys.iswindows() && haskey(ENV,"APPVEYOR") )

@time begin
  if group == "All" || group == "Interface"
    @time @testset "First Rand Tests" begin include("first_rand_test.jl") end
    @time @testset "Linear SDE Tests" begin include("sde/sde_linear_tests.jl") end
    @time @testset "Linear RODE Tests" begin include("rode_linear_tests.jl") end
    @time @testset "Two-dimensional Linear SDE Tests" begin include("sde/sde_twodimlinear_tests.jl") end
    @time @testset "Number Type Tests" begin include("sde/sde_numbertype_tests.jl") end
    @time @testset "Complex Number Tests" begin include("complex_tests.jl") end
    @time @testset "Static Array Tests" begin include("static_array_tests.jl") end
    @time @testset "Noise Type Tests" begin include("noise_type_test.jl") end
    @time @testset "Mass matrix tests" begin include("mass_matrix_tests.jl") end
    @time @testset "Outofplace Arrays Tests" begin include("outofplace_arrays.jl") end
    @time @testset "tdir Tests" begin include("tdir_tests.jl") end
    @time @testset "tstops Tests" begin include("tstops_tests.jl") end
    @time @testset "saveat Tests" begin include("saveat_tests.jl") end
    @time @testset "Oval2" begin include("oval2_test.jl") end
    @time @testset "Composite Tests" begin include("composite_algorithm_test.jl") end
    @time @testset "Events Tests" begin include("events_test.jl") end
    @time @testset "Cache Tests" begin include("cache_test.jl") end
    @time @testset "Element-wise Tolerances Tests" begin include("tolerances_tests.jl") end
    @time @testset "Zero'd Noise Tests" begin include("zerod_noise_test.jl") end
    #@time @testset "Scalar Tests" begin include("scalar_noise.jl") end # Fails because of bounds checks
    @time @testset "Stiffness Detection Test" begin include("stiffness_detection_test.jl") end
    @time @testset "Adaptive SDE Linear Tests" begin include("adaptive/sde_linearadaptive_tests.jl") end
    @time @testset "Multiple Dimension Linear Adaptive Test" begin include("adaptive/sde_twodimlinearadaptive_tests.jl") end
    @time @testset "Autostepsize Test" begin include("adaptive/sde_autostepsize_test.jl") end
    @time @testset "Additive Lorenz Attractor Test" begin include("adaptive/sde_lorenzattractor_tests.jl") end
    @time @testset "Adaptive Complex Mean Test" begin include("adaptive/sde_complex_adaptive_mean_test.jl") end
    @time @testset "Utility Tests" begin include("utility_tests.jl") end
  end

  if !is_APPVEYOR && (group == "All" || group == "AlgConvergence")
    @time @testset "Additive SDE Tests" begin include("sde/sde_additive_tests.jl") end
    @time @testset "Non-diagonal SDE Tests" begin include("nondiagonal_tests.jl") end
    @time @testset "Rossler Order Tests" begin include("sde/sde_rosslerorder_tests.jl") end
    @time @testset "Convergence Tests" begin include("sde/sde_convergence_tests.jl") end
    @time @testset "Split Tests" begin include("split_tests.jl") end
    @time @testset "Stratonovich Convergence Tests" begin include("stratonovich_convergence_tests.jl") end
    @time @testset "IIF Convergence Tests" begin include("iif_methods.jl") end
    LONGER_TESTS && @time @testset "Weak Convergence Tests" begin include("weak_convergence.jl") end
    @time @testset "Cummutative Noise Methods Tests" begin include("commutative_tests.jl") end
    @time @testset "Multivariate Geometric Tests" begin include("multivariate_geometric.jl") end
  end
end


text:
using DataFrames, Query, Compat, GLM, CSV
using PlotlyJS

function percentError(estimate,exact)
  return abs((estimate-exact)/exact)*100
end

function percentDifference(val1,val2)
  return abs((val1-val2)/((val1+val2)/2))*100
end

function readTestData(f)
  re = r"^\[(?<data>[-?\d+\.\d+\s*]+)\],\s+\[(?<label>[-?\d+\.\d+?\s*]+)\]$"
  data = Float64[]
  label = Float64[]
  # Read headerline
  header = readline(f)
  while !eof(f)
    s = readline(f)
    m = match(re,s)
    data = [data;readdlm(IOBuffer(m[:data]))]
    label = [label;readdlm(IOBuffer(m[:label]))]
  end
  data = convert(DataFrame,
                 Dict(zip(["data$i" for i in 1:size(data,2)],
                          [data[:,i] for i in 1:size(data,2)])))
  label = convert(DataFrame,
                  Dict(zip(["label$i" for i in 1:size(label,2)],
                           [label[:,i] for i in 1:size(label,2)])))
  return hcat(data,label)
end

function readLqrData(filename)
  data = Float64[]
  label = Float64[]
  df = CSV.read(filename; delim=',', header=1)
  return df
end

function getQuadAngleControl(df)
  controls = Float64[]
  controlSystem = [1 1 -1 -1;
                   1 -1 -1 1;
                   1 1 1 -1;
                   1 -1 1 1]

  for i=1:size(df,1)
    controls = [controls ;
      transpose(controlSystem \
                [df[:label1][i],df[:label2][i],df[:label3][i],df[:label4][i]])]
  end
  controls = convert(DataFrame,
                     Dict(zip(["thrustCtl" "pitchCtl" "rollCtl" "yawCtl"],
                              [controls[:,i] for i in 1:size(controls,2)])))

  return hcat(df,controls)
end

#########################################

token = "LQR"
if token == "test"
  f = open("./data/test.txt","r")

  df = readTestData(f);
  df = getQuadAngleControl(df)

  linearModel = Dict("pitch" => lm(@formula(pitchCtl ~ data10 + data11 + data12 +
                                            data4 + data5 + data6), df),
                     "roll" => lm(@formula(rollCtl ~ data13 + data14 + data15 +
                                           data1 + data2 + data3), df),
                     "yaw" => lm(@formula(yawCtl ~ data7 + data8 + data9),df),
                     "thrust" => lm(@formula(thrustCtl ~ data16 + data17 + data18),df))

  print(linearModel["roll"])
  print(linearModel["yaw"])
  print(linearModel["thrust"])

elseif token == "PID"

elseif token == "LQR"

  # LQR gains
  gains = [0.0020 1.0000 -0.0000 -0.0418 1.3856 -0.0000 1.1039 0.0001 -0.0000 4.6346 0.0395 -0.0000;
           -1.0000 0.0020 -0.0043 -1.3583 -0.0458 -0.0015 0.0004 1.1348 -0.0028 -0.0247 4.6926 -0.0028;
           -0.0035 0.0000 0.1773 -0.0070 -0.0001 -0.0188 -0.0000 -0.0010 1.0131 -0.0002 -0.0142 0.9842;
           0.0037 -0.0000 -0.9842 0.0147 0.0001 -0.9920 0.0000 0.0044 0.2049 0.0001 0.2483 0.1773]

  f = Dict{Any,Any}()
  df = Dict{Any,Any}()
  gainDiff = Dict{Any,Any}()
  linearModels = Dict{Any,Dict{String,Any}}()
  f[1] = [0,"data/lqrNn0Loops.csv"]
  f[2] = [5,"data/lqrNn5Loops.csv"]
  f[3] = [10,"data/lqrNn10Loops.csv"]
  f[4] = [15,"data/lqrNn15Loops.csv"]
  f[5] = [20,"data/lqrNn20Loops.csv"]

  for i=1:5
    numLoops = f[i][1]
    df[numLoops] = readLqrData(f[i][2])
    linearModels[i] = Dict("del_lat" => lm(@formula(del_lat ~ x + y + z + u + v
                                                            + w + p + q + r +
                                                            phi + theta + psi),
                                           df[numLoops]),
                           "del_lon" => lm(@formula(del_lon ~ x + y + z + u + v
                                                            + w + p + q + r +
                                                            phi + theta + psi),
                                           df[numLoops]),
                           "del_yaw" => lm(@formula(del_yaw ~ x + y + z + u + v
                                                            + w + p + q + r +
                                                            phi + theta + psi),
                                           df[numLoops]),
                           "del_thrust" => lm(@formula(del_thrust ~ x + y + z + u + v
                                                                  + w + p + q + r +
                                                                  phi + theta +
                                                                  psi),df[numLoops]))

    gainDiff[numLoops] = zeros(size(gains))
    r=[]

    push!(r,[coef(linearModels[i]["del_lat"])[2:end] gains[1,:]])
    push!(r,[coef(linearModels[i]["del_lon"])[2:end] gains[2,:]])
    push!(r,[coef(linearModels[i]["del_yaw"])[2:end] gains[3,:]])
    push!(r,[coef(linearModels[i]["del_thrust"])[2:end] gains[4,:]])

    for i=1:4
      for j=1:size(r[i],1)
        gainDiff[numLoops][i,j] = percentDifference(r[i][j,1],r[i][j,2])
      end
    end
  end

  numLoopsSet = [0 5 10 15 20]'
  gainDiffChange = [mean(gainDiff[i]) for i in numLoopsSet]'

  # Plotting
  #=
  trace1 = scatter(;x=numLoopsSet,y=gainDiffChange,
            mode="markers",
            marker=attr(color="#02010f",size=10,symbol="square",
                        line=attr(color="#02010f",width=0)),
            name="Gain Percent Difference")

  layout = Layout(autosize=false, width=500, height=500,
                  margin=attr(l=0,r=0,b=0,t=65))
  p = plot([trace1],layout)
  =#

  plot(scatter(x=numLoopsSet,y=gainDiffChange,
               marker=attr(color="#b3010f",size=10,symbol="square",
                           line=attr(color="#02010f",width=0)),
               name="Gain Percent Difference"))
  display(p)

end

max_stars_repo_name: umd-agrc/SimpleControlSim

max_issues_repo_name: umd-agrc/SimpleControlSim

max_forks_repo_name: umd-agrc/SimpleControlSim


text:
export Cat

import Scruff.Utils.normalize
import Distributions.Categorical

@doc """
    mutable struct Cat{O} <: Dist{O, Vector{Real}}

`Cat` defines an sfunc that represents a set of categorical output values that are not conditioned
on any input.  Its parameters are always of type `Vector{Real}`, which is the probability of each
output value.

# Supported operations
- `support`
- `support_quality`
- `sample`
- `cpdf`
- `bounded_probs`
- `compute_pi`
- `f_expectation`

# Type parameters
- `O`: the output type of the `Cat`
"""
mutable struct Cat{O} <: Dist{O}
    range :: Vector{O}
    inversemap :: Dict{O, Int}
    params :: Vector{<:Real}
    original_range :: Vector{O}
    """
        Cat(r::Vector{O}, ps::Vector{<:Real}) where O

    `Cat`'s constructor

    # Arguments
    - `r::Vector{O}`: the set of possible output values
    - `ps::Vector{<:Real}`: the set of probabilities for each value in `r`
      (will be normalized on call to `sample`)
    """
    function Cat(range::Vector{O}, params::Vector{<:Real}) where O
        @assert length(range) == length(params)
        # Handle repeated values correctly
        d = Dict{O, Float64}()
        for (x,p) in zip(range, params)
            d[x] = get(d, x, 0) + p
        end
        r = collect(keys(d))
        ps = [d[x] for x in r]
        inversemap = Dict([x => i for (i,x) in enumerate(r)])
        new{O}(r,inversemap,ps,range)
    end

    function Cat(d::Dict{O, <:Real}) where O
        Cat(collect(keys(d)), collect(values(d)))
    end

    """
        Cat(rps::Vector{Pair{O,<:Real}}) where O
    `Cat` constructor
    # Arguments
    - rps::Vector{Pair{O,Float64}}: a set of `Pair`s `output_value=>probability_of_value`
    """
    function Cat(rps::Vector{<:Pair{O,<:Real}}) where O
        range = first.(rps)
        probs = last.(rps)
        return Cat(range, probs)
    end
end

@impl begin
    struct CatSupport end
    function support(sf::Cat{O},
                    parranges::NTuple{N,Vector},
                    size::Integer,
                    curr::Vector{<:O}) where {O,N}
        sf.range
    end
end

@impl begin
    struct CatSupportQuality end
    function support_quality(::Cat, parranges)
        :CompleteSupport
    end
end

@impl begin
    struct CatSample end
    function sample(sf::Cat{O}, i::Tuple{})::O where {O}
        i = rand(Categorical(normalize(sf.params)))
        return sf.range[i]
    end
end

@impl begin
    struct CatCpdf end
    function cpdf(sf::Cat{O}, i::Tuple{}, o::O)::AbstractFloat where {O}
        if o in sf.range
            return sf.params[sf.inversemap[o]]
        else
            return 0.0
        end
    end
end

@impl begin
    struct CatBoundedProbs end
    function bounded_probs(sf::Cat{O},
                        range::Vector{O},
                        ::NTuple{N,Vector})::Tuple{Vector{<:AbstractFloat},
                        Vector{<:AbstractFloat}} where {O,N}

        ps = [x in keys(sf.inversemap) ? sf.params[sf.inversemap[x]] : 0.0 for x in range]
        (ps, ps)
    end
end

@impl begin
    struct CatComputePi end

    function compute_pi(sf::Cat{O},
                    range::Vector{O},
                    ::NTuple{N,Vector},
                    ::Tuple)::Cat{O} where {N,O}

        ps = bounded_probs(sf, range, ())[1]
        Cat(range, ps)
    end
end

#=
@impl begin
    struct CatInitialStats end

    function initial_stats(sf::Cat)
        zeros(Float64, length(sf.original_range))
    end
end

@impl begin
    struct CatAccumulateStats end

    function accumulate_stats(sf::Cat, existing_stats, new_stats)
        existing_stats .+ new_stats
    end
end

@impl begin
    struct CatExpectedStats end

    function expected_stats(sf::Cat{O},
            range::Vector{O},
            ::NTuple{N,Vector},
            ::NTuple{M,Dist},
            lambda::Score{<:O}) where {O,N,M}
        orig = sf.original_range
        ps = zeros(Float64, length(orig))
        for (i,x) in enumerate(orig)
            if x in sf.range
                ps[i] = sf.params[sf.inversemap[x]]
            end
        end
        ls = [get_score(lambda, r) for r in orig]
        return ps .* ls
    end
end

@impl begin
    struct CatMaximizeStats end

    function maximize_stats(sf::Cat, stats)
        normalize(stats)
    end
end
=#

@impl begin
    struct CatFExpectation end
    function f_expectation(sf::Cat, ::Tuple{}, fn::Function)
        sum = 0.0
        total = 0.0
        for (i,x) in enumerate(sf.range)
            sum += fn(x) * sf.params[i]
            total += sf.params[i]
        end
        return sum / total
    end
end


text:
#=
using Revise
=#

include("HMC.jl")
using Distributions, Random
using Flux

using RCall
@rimport graphics as rgraphics
@rimport rcommon

struct State
  beta
  log_sig2
end

pretty(s::State) = "State(β => $(s.beta); σ => $(sqrt(exp(s.log_sig2[1]))))"

State(K::Int) = State(param(zeros(K)), param(ones(1)))

function sim_data(;N::Int=100, K::Int=10, sig=.3)
  X = randn(N, K)
  b = randn(K)
  y = X * b + randn(N) * sig
  return Dict(:X => X, :b => b, :y => y, :sig => sig)
end

function lpdf_std_normal(z)
  return -z ^ 2 / 2 - log(2 * pi) / 2
end

lpdf_normal(y, m, s) = lpdf_std_normal((y-m)/s) - log(s)

function loglike(X, y, state::State)
  # println(state)
  m = X * state.beta
  s = sum(sqrt.(exp.(state.log_sig2))) + 1e-6
  # s = 1.0
  ll = sum(lpdf_normal.(y, m, s))
  # println("ll: $(ll)")
  return ll
end

function logprior(state::State)
  lp_beta = sum(logpdf.(Normal(0, 10), state.beta))
  lp_sig2 = sum(logpdf.(InverseGamma(2, 1), exp.(state.log_sig2)))
  lp_sig2 += sum(state.log_sig2)  # add log abs jacobian
  lp = lp_beta + lp_sig2
  # println("lp: $(lp)")
  return lp
end

### MAIN ###
Random.seed!(1)

N, K = (1000, 20)
dat = sim_data(K=K)
log_prob(state::State) = loglike(dat[:X], dat[:y], state) + logprior(state)
state = State(K)
epsilon = 5 / N
num_leapfrog_steps = 100

# Compile
_ = HMC.hmc_update(state, log_prob, num_leapfrog_steps, epsilon)

function simulate(init; nburn=500, nsamps=1000)
  state = deepcopy(init)

  log_prob_hist = zeros(nsamps)
  samps = [state for i in 1:nsamps]
  for i in 1:(nburn + nsamps)
    print("\rProgress: $i / $(nburn + nsamps)")
    state, log_prob_curr = HMC.hmc_update(state, log_prob,
                                          num_leapfrog_steps, epsilon)
    if i > nburn
      samps[i - nburn] = state
      log_prob_hist[i - nburn] = log_prob_curr
    end
  end
  return samps, log_prob_hist
end

samps, log_prob_hist = simulate(state, nburn=500, nsamps=500)
println(pretty(samps[end]));
dat

acceptance_rate = length(unique(log_prob_hist)) / length(log_prob_hist)
println("Acceptance rate: $acceptance_rate")
rgraphics.plot(log_prob_hist, main="", xlab="iter", ylab="log prob", typ="l");

sig2_post = [exp(Tracker.data(s.log_sig2[1])) for s in samps]
b_first_2_post = cat([Tracker.data(s.beta[1:2])' for s in samps]..., dims=1)
rgraphics.plot(b_first_2_post, xlab="", ylab="", main="", typ="l");

rcommon.plotPosts(b_first_2_post);
rcommon.plotPost(sqrt.(sig2_post));
rgraphics.abline(v=dat[:sig], lwd=3, col="orange");


text:
using Revise # reduce recompilation time
using Plots
# using Documenter
using LinearAlgebra
using SparseArrays
using BenchmarkTools
using UnPack

push!(LOAD_PATH, "./src")
using CommonUtils
using Basis1D

using SetupDG

push!(LOAD_PATH, "./examples/EntropyStableEuler.jl/src")
using EntropyStableEuler
using EntropyStableEuler.Fluxes1D
const γ = 1.4#5/3# 1.4

function wavespeed_1D(rho,rhou,E)
    p = pfun_nd(rho,(rhou,),E)
    cvel = @. sqrt(γ*p/rho)
    return @. abs(rhou/rho) + cvel
end
unorm(U) = sum(map((x->x.^2),U))
function pfun_nd(rho, rhoU, E)
    rhoUnorm2 = unorm(rhoU)./rho
    return @. (γ-1)*(E - .5*rhoUnorm2)
end

function primitive_to_conservative_hardcode(rho,U,p)
    rhoU = rho.*U
    Unorm = unorm(U)
    E = @. p/(γ-1) + .5*rho*Unorm
    return (rho,rhoU,E)
end

"Approximation parameters"
N = 1 # The order of approximation
K = 50
T = 0.2 # endtime
#T = 6.0
#T = 0.0001

# Sod shocktube
const Bl = -0.5
const Br = 0.5
const rhoL = 1.0
const rhoR = 0.125
const pL = 1.0
const pR = 0.1
const xC = 0.0

# # Leblanc shocktube
# const Bl = 0.0
# const Br = 9.0
# const rhoL = 1.0
# const rhoR = 0.001
# const pL = 0.1
# const pR = 1e-7
# const xC = 3.0

"Mesh related variables"
#VX = LinRange(-0.5,0.5,K+1)
VX = LinRange(Bl,Br,K+1)
EToV = transpose(reshape(sort([1:K; 2:K+1]),2,K))

"Initialize reference element"
r,_ = gauss_lobatto_quad(0,0,N)   # Reference nodes
VDM = vandermonde_1D(N,r)         # modal to nodal
Dr = grad_vandermonde_1D(N,r)/VDM # nodal differentiation
V1 = vandermonde_1D(1,r)/vandermonde_1D(1,[-1;1]) # nodal linear interpolation

Nq = N+1
rq,wq = gauss_quad(0,0,Nq)
Vq = vandermonde_1D(N,rq)/VDM
M = Vq'*diagm(wq)*Vq
Mlump = zeros(size(M))
Mlump_inv = zeros(size(M))
for i = 1:N+1
    Mlump[i,i] = sum(M[i,:])
    Mlump_inv[i,i] = 1.0/Mlump[i,i]
end
S = M*Dr
L = zeros(N+1,2)
L[1,1] = -1.0/2.0
L[end,end] = 1.0/2.0

rf = [-1.0;1.0]
nrJ = [-1.0;1.0]
Vf = vandermonde_1D(N,rf)/VDM

"""Connectivity map"""
x = V1*VX[transpose(EToV)]
xf = Vf*x
mapM = reshape(1:2*K,2,K)
mapP = copy(mapM)
mapP[1,2:end] .= mapM[2,1:end-1]
mapP[2,1:end-1] .= mapM[1,2:end]

# """Periodic"""
# mapP[1] = mapM[end]
# mapP[end] = mapP[1]

"""Geometric factors"""
J = repeat(transpose(diff(VX)/2),N+1,1)
nxJ = repeat([-1;1],1,K)
rxJ = 1.0

"""Initial condition"""
# rho_x(x) = (x <= 0.0) ? 1.0 : 0.125
# u_x(x) = 0.0
# p_x(x) = (x <= 0.0) ? 1.0 : 0.1
# rho_x(x) = (x <= 3.0) ? 1.0 : 0.001
# u_x(x) = 0.0
# p_x(x) = (x <= 3.0) ? 0.1 : 1e-7
rho_x(x) = (x <= xC) ? rhoL : rhoR
u_x(x) = 0.0
p_x(x) = (x <= xC) ? pL : pR
# rho_x(x) = (x <= 0.0) ? 1.0 : 0.8
# u_x(x) = 0.0
# p_x(x) = (x <= 0.0) ? 1.0 : 0.8
# rho_x(x) = (x <= 0.0) ? 1.0 : 0.125
# u_x(x) = 1.0
# p_x(x) = (x <= 0.0) ? 1.0 : 1.0

rho = @. rho_x(x)
u = @. u_x(x)
p = @. p_x(x)
Q = primitive_to_conservative_hardcode(rho,u,p)

# hardcoded first order IDP
function rhs_IDPlow(Q,K,N)
    p = pfun_nd.(Q[1],Q[2],Q[3])
    flux = zero.(Q)
    @. flux[1] = Q[2]
    @. flux[2] = Q[2]^2/Q[1]+p
    @. flux[3] = Q[3]*Q[2]/Q[1]+p*Q[2]/Q[1]

    """Periodic case"""
    # J = 1/K/2 # assume uniform interval
    # dfdx = (zeros(N+1,K),zeros(N+1,K),zeros(N+1,K))
    # for i = 1:K*(N+1)
    #     for c = 1:3
    #         dfdx[c][i] = 1/2*(flux[c][mod1(i+1,K*(N+1))] - flux[c][mod1(i-1,K*(N+1))])
    #     end
    # end

    # visc = (zeros(N+1,K),zeros(N+1,K),zeros(N+1,K))
    # for i = 1:K*(N+1)
    #     wavespd_curr = abs(wavespeed_1D(Q[1][i],Q[2][i],Q[3][i]))
    #     wavespd_R = abs(wavespeed_1D(Q[1][mod1(i+1,K*(N+1))],Q[2][mod1(i+1,K*(N+1))],Q[3][mod1(i+1,K*(N+1))]))
    #     wavespd_L = abs(wavespeed_1D(Q[1][mod1(i-1,K*(N+1))],Q[2][mod1(i-1,K*(N+1))],Q[3][mod1(i-1,K*(N+1))]))
    #     wavespd_curr_m = abs(wavespeed_1D(Q[1][i],-Q[2][i],Q[3][i]))
    #     wavespd_R_m = abs(wavespeed_1D(Q[1][mod1(i+1,K*(N+1))],-Q[2][mod1(i+1,K*(N+1))],Q[3][mod1(i+1,K*(N+1))]))
    #     wavespd_L_m = abs(wavespeed_1D(Q[1][mod1(i-1,K*(N+1))],-Q[2][mod1(i-1,K*(N+1))],Q[3][mod1(i-1,K*(N+1))]))
    #     dL = 1/2*max(wavespd_curr,wavespd_L)
    #     dR = 1/2*max(wavespd_curr,wavespd_R)
    #     for c = 1:3
    #         visc[c][i] = dL*(Q[c][mod1(i-1,K*(N+1))]-Q[c][i]) + dR*(Q[c][mod1(i+1,K*(N+1))]-Q[c][i])
    #     end
    # end

    # rhsQ = @. 1/J*(-dfdx+visc)
    # return rhsQ

    """shocktube"""
    J = (Br-Bl)/K/2 # assume uniform interval
    dfdx = (zeros(N+1,K),zeros(N+1,K),zeros(N+1,K))
    for i = 2:K*(N+1)-1
        for c = 1:3
            dfdx[c][i] = 1/2*(flux[c][mod1(i+1,K*(N+1))] - flux[c][mod1(i-1,K*(N+1))])
        end
    end
    dfdx[1][1] = 1/2*(flux[1][2] - 0.0)
    dfdx[2][1] = 1/2*(flux[2][2] - pL)
    dfdx[3][1] = 1/2*(flux[3][2] - 0.0)
    dfdx[1][end] = 1/2*(0.0 - flux[1][end-1])
    dfdx[2][end] = 1/2*(pR - flux[2][end-1])
    dfdx[3][end] = 1/2*(0.0 - flux[3][end-1])

    visc = (zeros(N+1,K),zeros(N+1,K),zeros(N+1,K))
    for i = 2:K*(N+1)-1
        wavespd_curr = wavespeed_1D(Q[1][i],Q[2][i],Q[3][i])
        wavespd_R = wavespeed_1D(Q[1][mod1(i+1,K*(N+1))],Q[2][mod1(i+1,K*(N+1))],Q[3][mod1(i+1,K*(N+1))])
        wavespd_L = wavespeed_1D(Q[1][mod1(i-1,K*(N+1))],Q[2][mod1(i-1,K*(N+1))],Q[3][mod1(i-1,K*(N+1))])
        dL = 1/2*max(wavespd_curr,wavespd_L)
        dR = 1/2*max(wavespd_curr,wavespd_R)
        for c = 1:3
            visc[c][i] = dL*(Q[c][mod1(i-1,K*(N+1))]-Q[c][i]) + dR*(Q[c][mod1(i+1,K*(N+1))]-Q[c][i])
        end
    end

    # i = 1, (rho, rhoU, E) = (1.0, 0.0, 0.1/(γ-1))
    wavespd_curr = wavespeed_1D(Q[1][1],Q[2][1],Q[3][1])
    wavespd_R = wavespeed_1D(Q[1][2],Q[2][2],Q[3][2])
    wavespd_L = wavespeed_1D(rhoL,0.0,pL/(γ-1))
    dL = 1/2*max(wavespd_curr,wavespd_L)
    dR = 1/2*max(wavespd_curr,wavespd_R)
    visc[1][1] = dL*(rhoL-Q[1][1]) + dR*(Q[1][2]-Q[1][1])
    visc[2][1] = dL*(0.0-Q[2][1]) + dR*(Q[2][2]-Q[2][1])
    visc[3][1] = dL*(pL/(γ-1)-Q[3][1]) + dR*(Q[3][2]-Q[3][1])

    # i = end, (rho, rhoU, E) = (1e-3, 0.0, 1e-7/(γ-1))
    wavespd_curr = wavespeed_1D(Q[1][end],Q[2][end],Q[3][end])
    wavespd_R = wavespeed_1D(rhoR,0.0,pR/(γ-1))
    wavespd_L = wavespeed_1D(Q[1][end-1],Q[2][end-1],Q[3][end-1])
    dL = 1/2*max(wavespd_curr,wavespd_L)
    dR = 1/2*max(wavespd_curr,wavespd_R)
    visc[1][end] = dL*(Q[1][end-1]-Q[1][end]) + dR*(rhoR-Q[1][end])
    visc[2][end] = dL*(Q[2][end-1]-Q[2][end]) + dR*(0.0-Q[2][end])
    visc[3][end] = dL*(Q[3][end-1]-Q[3][end]) + dR*(pR/(γ-1)-Q[3][end])

    rhsQ = @. 1/J*(-dfdx+visc)
    return rhsQ
end

function rhs_nodal(Q,Vf,mapP,rxJ,S,L,N,K,J,Mlump,Mlump_inv)
    p = pfun_nd(Q...)
    flux = zero.(Q)
    @. flux[1] = Q[2]
    @. flux[2] = Q[2]^2/Q[1]+p
    @. flux[3] = Q[3]*Q[2]/Q[1]+p*Q[2]/Q[1]

    # flux derivative
    dfdx = (x->rxJ*S*x).(flux)

    # jump in flux
    flux_f = (x->Vf*x).(flux)
    df = (x->x[mapP]-x).(flux_f)
    # Boundary condition (Sod tube)
    df[1][1] = -flux_f[1][1]
    df[2][1] = 1.0-flux_f[2][1]
    df[3][1] = -flux_f[3][1]
    df[1][end] = -flux_f[1][end]
    df[2][end] = 0.1-flux_f[2][end]
    df[3][end] = -flux_f[3][end]
    df = (x->L*x).(df)

    # Artificial Viscosity
    visc = (zeros(N+1,K),zeros(N+1,K),zeros(N+1,K))
    D = zeros(N+1,N+1)
    T = zeros(N+1,N+1)
    D_arr = zeros(N+1,K) # Array of d_ii
    for k = 1:K
        # Construct graph viscosity coefficients
        for i = 1:N+1
            for j = 1:N+1
                wavespd = max(wavespeed_1D(Q[1][i,k],Q[2][i,k],Q[3][i,k]),
                              wavespeed_1D(Q[1][j,k],Q[2][j,k],Q[3][j,k]))
                # Boundary nodes (hard coded)
                if i == 1 && j == 1
                    #T[i,j] = max(wavespd*abs(S[i,j]-1/2),wavespd*abs(S[i,j]-1/2))
                elseif i == N+1 && j == N+1
                    #T[i,j] = max(wavespd*abs(S[i,j]+1/2),wavespd*abs(S[i,j]+1/2))
                else
                    T[i,j] = max(wavespd*abs(S[i,j]),wavespd*abs(S[i,j]))
                    if i != j
                        D_arr[i,k] -= T[i,j]
                    end
                end
            end
        end

        for c = 1:3
            # Difference matrix
            for i = 1:N+1
                for j = 1:N+1
                    D[i,j] = Q[c][j,k]-Q[c][i,k]
                end
            end
            visc[c][:,k] = (T.*D)*ones(N+1)

            # Exterior nodes TODO: hardcoded
            wavespd = max(wavespeed_1D(Q[1][end,mod1(k-1,K)],Q[2][end,mod1(k-1,K)],Q[3][end,mod1(k-1,K)]),
                        wavespeed_1D(Q[1][1,k],Q[2][1,k],Q[3][1,k]))
            visc[c][1,k] += wavespd/2*(Q[c][end,mod1(k-1,K)]-Q[c][1,k])
            D_arr[1,k] -= wavespd/2

            wavespd = max(wavespeed_1D(Q[1][end,k],Q[2][end,k],Q[3][end,k]),
                        wavespeed_1D(Q[1][1,mod1(k+1,K)],Q[2][1,mod1(k+1,K)],Q[3][1,mod1(k+1,K)]))
            visc[c][end,k] += wavespd/2*(Q[c][1,mod1(k+1,K)]-Q[c][end,k])
            D_arr[end,k] -= wavespd/2

            # if k <= K-1
            #     wavespd = max(wavespeed_1D(Q[1][end,k],Q[2][end,k],Q[3][end,k]),
            #                 wavespeed_1D(Q[1][1,k+1],Q[2][1,k+1],Q[3][1,k+1]))
            #     visc[c][end,k] += wavespd/2*(Q[c][1,k+1]-Q[c][end,k])
            #     D_arr[end,k] -= wavespd/2
            # else
            #     # wavespd = max(wavespeed_1D(Q[1][end,k],Q[2][end,k],Q[3][end,k]),
            #     #             wavespeed_1D(0.125,0.0,0.1/(γ-1)))
            #     # if c == 1
            #     #     visc[1][end,k] += wavespd/2*(0.125-Q[1][end,k])
            #     # elseif c == 2
            #     #     visc[2][end,k] += wavespd/2*(0.0-Q[2][end,k])
            #     # else
            #     #     visc[3][end,k] += wavespd/2*(0.1/(γ-1)-Q[3][end,k])
            #     # end
            #     # D_arr[end,k] -= wavespd/2
            # end
        end
    end

    dt = Inf
    for k = 1:K
        for i = 1:N+1
            dt = min(dt,-J[1]*Mlump[i,i]/(4*D_arr[i,k]))
        end
    end

    rhsQ = dfdx .+ df .- visc
    rhsQ = (x->-Mlump_inv*x./J).(rhsQ)
    return rhsQ, dt
end

# Time stepping
"Time integration"
t = 0.0
Q = collect(Q)
resQ = [zeros(size(x)),zeros(size(x)),zeros(size(x))]

# Forward Euler
while t < T
    #rhsQ,dt = rhs_nodal(Q,Vf,mapP,rxJ,S,L,N,K,J,Mlump,Mlump_inv)
    rhsQ = rhs_IDPlow(Q,K,N)
    dt = 0.0001
    @. Q = Q + dt*rhsQ
    global t = t + dt
    println("Current time $t with time step size $dt, and final time $T")
end

# rk4a,rk4b,rk4c = rk45_coeffs()
# Nsteps = Int(T/0.0001)
# for i = 1:1000#Nsteps
#     for INTRK = 1:5
#         rhsQ,dt = rhs_nodal(Q,Vf,mapP,rxJ,S,L,N,K,J,Mlump,Mlump_inv)
#         dt = 0.0001
#         @. resQ = rk4a[INTRK]*resQ + dt*rhsQ
#         @. Q   += rk4b[INTRK]*resQ
#     end
#     if i%10 == 0 || i == Nsteps
#         println("Number of time steps $i out of $Nsteps")
#     end
# end
#=
# SSPRK(3,3)
for i = 1:5
    rhsQ =  rhs_nodal(Q,Vf,mapP,rxJ,S,L,N,K,J,Mlump,Mlump_inv)
    w1 = @. Q + dt*rhsQ
    rhsQ = rhs_nodal(w1,Vf,mapP,rxJ,S,L,N,K,J)
    z1 = @. w1 + dt*rhsQ
    w2 = @. 3/4*Q + 1/3*z1
    rhsQ = rhs_nodal(w2,Vf,mapP,rxJ,S,L,N,K,J)
    z2 = @. w2 + dt*rhsQ
    @. Q = 1/3*Q + 2/3*z2
    if i%10 == 0 || i == Nsteps
        println("Number of time steps $i out of $Nsteps")
    end
end
=#

Vp = vandermonde_1D(N,LinRange(-1,1,10))/VDM
gr(size=(300,300),ylims=(0,1.2),legend=false,markerstrokewidth=1,markersize=2)
plt = plot(Vp*x,Vp*Q[1])

max_stars_repo_name: yiminllin/ESDG-PosLimit

max_issues_repo_name: yiminllin/ESDG-PosLimit

max_forks_repo_name: yiminllin/ESDG-PosLimit


text:
export admm

mutable struct ADMM{rT,matT,opT,ropT,vecT,rvecT,preconT} <: AbstractLinearSolver where {vecT <: AbstractVector{Union{rT, Complex{rT}}}, rvecT <: AbstractVector{rT}}
  # oerators and regularization
  A::matT
  reg::Vector{Regularization}
  regTrafo::Vector{ropT}
  # fields and operators for x update
  AᴴA::opT
  β::vecT
  β_y::vecT
  # fields for primal & dual variables
  x::vecT
  z::Vector{vecT}
  zᵒˡᵈ::Vector{vecT}
  u::Vector{vecT}
  # other parameters
  precon::preconT
  ρ::rvecT # TODO: Switch all these vectors to Tuple
  iterations::Int64
  iterationsInner::Int64
  # state variables for CG
  cgStateVars::CGStateVariables
  # convergence parameters
  rᵏ::rvecT
  sᵏ::rvecT
  ɛᵖʳⁱ::rvecT
  ɛᵈᵘᵃ::rvecT
  σᵃᵇˢ::rT
  absTol::rT
  relTol::rT
  tolInner::rT
  normalizeReg::Bool
  regFac::rT
  vary_ρ::Bool
  verbose::Bool
end

"""
    ADMM(A, x::vecT=zeros(eltype(A),size(A,2))
          ; reg=nothing, regName=["L1"], λ=[0.0], kargs...)

creates an `ADMM` object for the system matrix `A`.

# Arguments
* `A`                           - system matrix
* `x::vecT`                     - Array with the same type and size as the solution
* (`reg=nothing`)               - Regularization object
* (`regName=["L1"]`)            - name of the Regularization to use (if reg==nothing)
* (`λ=[0.0]`)                   - Regularization paramter
* (`regTrafo=nothing`)          - transformations applied inside each regularizer
* (`precon=Identity()`)         - preconditionner for the internal CG algorithm
* (`ρ::Real=1.e-2`)          - penalty of the augmented lagrangian
* (`adaptRho::Bool=false`)      - adapt rho to balance primal and dual feasibility
* (`iterations::Int64=50`)      - max number of ADMM iterations
* (`iterationsInner::Int64=10`) - max number of internal CG iterations
* (`absTol::Real=eps()`)     - abs tolerance for stopping criterion
* (`relTol::Real=eps()`)     - rel tolerance for stopping criterion
* (`tolInner::Real=1.e-5`)   - rel tolerance for CG stopping criterion
"""
function ADMM(A::matT, x::Vector{T}=zeros(eltype(A),size(A,2)); reg=nothing, regName=["L1"]
            , λ=[0.0]
            , regTrafo=nothing
            , AᴴA::opT=nothing
            , precon=Identity()
            , ρ=1e-1
            , iterations::Integer=50
            , iterationsInner::Integer=10
            , absTol::Real=eps(real(T))
            , relTol::Real=eps(real(T))
            , tolInner::Real=1e-5
            , normalizeReg::Bool=false
            , vary_ρ::Bool=false
            , verbose::Bool=false
            , kargs...) where {T,matT,opT}
# TODO: The constructor is not type stable

  # unify Floating types
  if typeof(ρ) <: Number
    ρ_vec = [real(T).(ρ)]
  else
    ρ_vec = real(T).(ρ)
  end
  λ = real(T).(λ)
  absTol = real(T)(absTol)
  relTol = real(T)(relTol)
  tolInner = real(T)(tolInner)

  if reg == nothing
    reg = [Regularization(regName, λ, kargs...)]
  else
    reg = vec(reg) # using a custom method of vec(.)
  end

  if regTrafo == nothing
    regTrafo = [opEye(eltype(x),size(A,2)) for _=1:length(reg)]
  end

  # fields for primal & dual variables
  z = [similar(x, size(regTrafo[i],1)) for i=1:length(reg)]
  zᵒˡᵈ = [similar(z[i]) for i=1:length(reg)]
  u = [similar(z[i]) for i=1:length(reg)]

  # operator and fields for the update of x
  if AᴴA == nothing
    AᴴA = A'*A
  end
  β = similar(x)
  β_y = similar(x)

  # statevariables for CG
  # we store them here to prevent CG from allocating new fields at each call
  statevars = CGStateVariables(zero(x),similar(x),similar(x))

  # convergence parameters
  rᵏ   = Array{real(T)}(undef, length(reg))
  sᵏ   = similar(rᵏ)
  ɛᵖʳⁱ = similar(rᵏ)
  ɛᵈᵘᵃ = similar(rᵏ)

  return ADMM(A,reg,regTrafo,AᴴA,β,β_y,x,z,zᵒˡᵈ,u,precon,ρ_vec,iterations
              ,iterationsInner,statevars, rᵏ,sᵏ,ɛᵖʳⁱ,ɛᵈᵘᵃ,zero(real(T)),absTol,relTol,tolInner
              ,normalizeReg,one(real(T)), vary_ρ, verbose)
end

"""
  init!(solver::ADMM{matT,opT,vecT,rvecT,preconT}, b::vecT
              ; A::matT=solver.A
              , AᴴA::opT=solver.AᴴA
              , x::vecT=similar(b,0)
              , kargs...) where {matT,opT,vecT,rvecT,preconT}

(re-) initializes the ADMM iterator
"""
function init!(solver::ADMM{rT,matT,opT,ropT,vecT,rvecT,preconT}, b::vecT
              ; A::matT=solver.A
              , AᴴA::opT=solver.AᴴA
              , x::vecT=similar(b,0)
              , kargs...) where {rT,matT,opT,ropT,vecT,rvecT,preconT}

  # operators
  if A != solver.A
    solver.A = A
    solver.AᴴA = ( AᴴA!=nothing ? AᴴA : A'*A )
  end

  # start vector
  if isempty(x)
    solver.x .= 0
 else
    solver.x[:] .= x
 end

  # primal and dual variables
  for i=1:length(solver.reg)
    solver.z[i][:] .= solver.regTrafo[i]*solver.x
    solver.zᵒˡᵈ[i][:] .= 0
    solver.u[i] .= 0
  end

  # right hand side for the x-update
  solver.β_y[:] .= adjoint(A) * b

  # convergence parameter
  solver.rᵏ .= Inf
  solver.sᵏ .= Inf
  solver.ɛᵖʳⁱ .= 0
  solver.ɛᵈᵘᵃ .= 0
  solver.σᵃᵇˢ = sqrt(length(b))*solver.absTol

  # normalization of regularization parameters
  if solver.normalizeReg
    solver.regFac = norm(b,1)/length(b)
  else
    solver.regFac = 1
  end
end

"""
    solve(solver::ADMM, b::vecT
          ; A::matT=solver.A
          , startVector::vecT=similar(b,0)
          , solverInfo=nothing
          , kargs...) where {matT,vecT}

solves an inverse problem using ADMM.

# Arguments
* `solver::ADMM`                  - the solver containing both system matrix and regularizer
* `b::Vector`                     - data vector
* (`A::matT=solver.A`)            - operator for the data-term of the problem
* (`startVector::Vector{T}=T[]`)  - initial guess for the solution
* (`solverInfo=nothing`)          - solverInfo for logging

when a `SolverInfo` objects is passed, the primal residuals `solver.rᵏ`
and the dual residual `norm(solver.sᵏ)` are stored in `solverInfo.convMeas`.
"""
function solve(solver::ADMM{rT,matT,opT,ropT,vecT,rvecT,preconT}, b::vecT; A=solver.A, AᴴA=solver.AᴴA, startVector::vecT=similar(b,0), solverInfo=nothing, kargs...) where {rT,matT,opT,ropT,vecT,rvecT,preconT}
  # initialize solver parameters
  init!(solver, b; A=A, AᴴA=AᴴA, x=startVector)

  # log solver information
  solverInfo != nothing && storeInfo(solverInfo,solver.z,solver.rᵏ...,solver.sᵏ...)

  # perform ADMM iterations
  for (iteration, item) = enumerate(solver)
    solverInfo != nothing && storeInfo(solverInfo,solver.z,solver.rᵏ...,solver.sᵏ...)
  end

  return solver.x
end

"""
  iterate(it::ADMM, iteration::Int=0)

performs one ADMM iteration.
"""
function iterate(solver::ADMM, iteration::Integer=0)
  if done(solver, iteration) return nothing end

  # 1. solve arg min_x 1/2|| Ax-b ||² + ρ/2 Σ_i||Φi*x+ui-zi||²
  # <=> (A'A+ρ Σ_i Φi'Φi)*x = A'b+ρΣ_i Φi'(zi-ui)
  copyto!(solver.β, solver.β_y)
  AᴴA = solver.AᴴA
  for i=1:length(solver.reg)
    solver.β[:] .+= solver.ρ[i]*adjoint(solver.regTrafo[i])*(solver.z[i].-solver.u[i])
    AᴴA += solver.ρ[i] * adjoint(solver.regTrafo[i]) * solver.regTrafo[i]
  end
  solver.verbose && println("conjugated gardients: ")
  cg!(solver.x, AᴴA, solver.β, Pl=solver.precon
      , maxiter=solver.iterationsInner, reltol=solver.tolInner, statevars=solver.cgStateVars, verbose = solver.verbose)

  for i=1:length(solver.reg)
    # 2. update z using the proximal map of 1/ρ*g(x)
    copyto!(solver.zᵒˡᵈ[i], solver.z[i])
    solver.z[i] .= solver.regTrafo[i]*solver.x .+ solver.u[i]
    if solver.ρ[i] != 0
      solver.reg[i].prox!(solver.z[i], solver.regFac*solver.reg[i].λ/solver.ρ[i]; solver.reg[i].params...)
    end

    # 3. update u
    solver.u[i] .+= solver.regTrafo[i]*solver.x .- solver.z[i]

    # update convergence measures (one for each constraint)
    solver.rᵏ[i] = norm(solver.regTrafo[i]*solver.x-solver.z[i])  # primal residual (x-z)
    solver.sᵏ[i] = norm(solver.ρ[i] * adjoint(solver.regTrafo[i]) * (solver.z[i] .- solver.zᵒˡᵈ[i])) # dual residual (concerning f(x))

    solver.ɛᵖʳⁱ[i] = max(norm(solver.regTrafo[i]*solver.x), norm(solver.z[i]))
    solver.ɛᵈᵘᵃ[i] = norm(solver.ρ[i] * adjoint(solver.regTrafo[i]) * solver.u[i])

    if solver.verbose
      println("rᵏ[$i] = $(solver.rᵏ[i])")
      println("sᵏ[$i] = $(solver.sᵏ[i])")
    end

    # adapt ρ according to Boyd et al.
    if solver.vary_ρ && solver.rᵏ[i] > 10solver.sᵏ[i]
      solver.ρ[i] *= 2
      solver.u[i] ./= 2
      solver.verbose && println("updated ρ[$i] = $(solver.ρ[i])")
    elseif solver.vary_ρ && solver.sᵏ[i] > 10solver.rᵏ[i]
      solver.ρ[i] /= 2
      solver.u[i] .*= 2
      solver.verbose && println("updated ρ[$i] = $(solver.ρ[i])")
    end
  end

  # return the primal feasibilty measure as item and iteration number as state
  return solver.rᵏ, iteration+1
end

function converged(solver::ADMM)
  for i=1:length(solver.reg)
    (solver.rᵏ[i] >= solver.σᵃᵇˢ + solver.relTol * solver.ɛᵖʳⁱ[i]) && return false
    (solver.sᵏ[i] >= solver.σᵃᵇˢ + solver.relTol * solver.ɛᵈᵘᵃ[i]) && return false
  end
  return true
end

@inline done(solver::ADMM,iteration::Int) = converged(solver) || iteration>=solver.iterations


text:
using
    CUDA,
    Flux,
    MLDatasets,
    Statistics

using Flux: onehotbatch, onecold, logitcrossentropy, logitcrossentropy

using DPSGD

include("Data.jl")
using .Data

if has_cuda()
    @info "Training on CUDA GPU"
    CUDA.allowscalar(false)
end

function mnist(; batchsize = 32)
    xtrain, ytrain = MLDatasets.MNIST.traindata(Float32)
    xtest, ytest = MLDatasets.MNIST.testdata(Float32)

    # Flatten each image into a linear array.
    xtrain = Flux.flatten(xtrain)
    xtest = Flux.flatten(xtest)

    ytrain = onehotbatch(ytrain, 0:9)
    ytest = onehotbatch(ytest, 0:9)

    train_data = DataLoader((xtrain, ytrain), batchsize = batchsize, shuffle = true,
        datasampler = VariableBatchSizeUniformSampleWithReplacement)
    test_data = DataLoader((xtest, ytest), batchsize = batchsize)

    return train_data, test_data
end

function make_model(; imgsize = (28, 28, 1), nclasses = 10)
    return Chain(
        Dense(prod(imgsize), 64, relu),
        Dense(64, nclasses))
end

function accuracy(model, data_loader)
    acc = 0
    num = 0
    for (x, y) in data_loader
        x, y = gpu(x), y
        acc += sum(onecold(cpu(model(x))) .== onecold(y))
        num += size(x, 2)
    end
    return acc / num
end

function train(; η = 3e-4, nepochs = 1000)
    train_data, test_data = mnist()

    mlp = make_model() |> gpu

    # Use `agg=identity` to not reduce losses over examples
    loss(x, y) = logitcrossentropy(mlp(x), y; agg = identity)

    target_delta = 1e-5
    nexamples = train_data.nobs
    privacy_opt = DifferentialPrivacy(nexamples, lotsize = train_data.batchsize)
    opt = Flux.Optimiser(privacy_opt, Descent(η))

    for epoch in 1:nepochs
        Flux.train!(loss, params(mlp), ((gpu(x), gpu(y)) for (x, y) in train_data), opt)
        @show epoch
        @show privacy_opt.niterations
        @show accuracy(mlp, test_data)
        @show privacy_spent(privacy_opt, target_delta)
    end
end

if abspath(PROGRAM_FILE) == @__FILE__
    train()
end