text:
# Endogenous Grid Points with AR1 + IID Income
# Cash on Hand as State variable
# Includes NIT and discount factor heterogeneity
# Greg Kaplan 2017
# Translated by Tom Sweeney Dec 2020

using Random, Interpolations, NLsolve, Plots
include("discrete_normal.jl")
include("lininterp1.jl")
include("rouwenhorst.jl")
include("ginicoeff.jl")

# PARAMETERS

## preferences
risk_aver = 1
beta = 0.955

## warm glow bequests: bequest_weight = 0 is accidental
bequest_weight = 0 # 0.07
bequest_luxury = 2 # 0.01
dieprob     = 0 # 1/50

## returns
r = 0.02
R = 1+r

## income risk: AR(1) + IID in logs
nyT = 5 # transitory component (not a state variable)
sd_logyT = 0.20 # relevant if nyT>1

nyP = 11 # 9 # persistent component
sd_logyP = 0.24
rho_logyP = 0.97

## cash on hand / savings grid
nx = 50
xmax = 40
xgrid_par = 0.4 # 1 for linear, 0 for L-shaped
borrow_lim = 0

## government
labtaxlow = 0 # proportional tax
labtaxhigh = 0 # additional tax on incomes above threshold
labtaxthreshpc = 0.99 # percentile of earnings distribution where high tax rate kicks in

savtax = 0.0001 # tax rate on savings
savtaxthresh = 0 # multiple of mean gross labor income

## discount factor shocks
nb = 1 # 1 or 2
betawidth = 0.065 # beta +/- beta width
betaswitch = 1/40 # 0

## computation
max_iter = 1000
tol_iter = 1.0e-6
Nsim = 100000
Tsim = 300

## mpc options
mpcamount1 = 1.0e-10 # approximate thoeretical mpc
mpcamount2 = 0.10 # one percent of average income: approx $500

# OPTIONS
Display = 1
DoSimulate = 1
MakePlots = 1
ComputeMPC = 1

## which function to interpolation
InterpMUC = 0
InterpCon = 1 # sometimes useful to stop extrapolating to negative MUC

# DRAW RANDOM NUMBERS
Random.seed!(2020)
yPrand = rand(Nsim,Tsim)
yTrand = randn(Nsim,Tsim)
betarand = rand(Nsim,Tsim)
dierand = rand(Nsim,Tsim)

# INCOME GRIDS

## persistent income: rowenhurst
logyPgrid, yPtrans, yPdist = rouwenhorst(nyP, -0.5*sd_logyP^2, sd_logyP, rho_logyP)
yPgrid = exp.(logyPgrid)
yPcumdist = cumsum(yPdist)
yPcumtrans = cumsum(yPtrans,dims=2)

## transitory income: disretize normal distribution
if nyT>1
    width = nlsolve(x -> discrete_normal(5,0,1,x...)[1],[2.0]).zero
    temp,logyTgrid,yTdist = discrete_normal(nyT,-0.5*sd_logyT^2,sd_logyT,width...)
elseif nyT==1
    logyTgrid = 0
    yTdist = 1
end
yTgrid = exp.(logyTgrid)

## gross labor income grid
grosslabincgrid = zeros(nyP,nyT)
grosslabincdist= zeros(nyP,nyT)
for iyP = 1:nyP
    for iyT= 1:nyT
        grosslabincgrid[iyP,iyT] = yPgrid[iyP]*yTgrid[iyT]
        grosslabincdist[iyP,iyT] = yPdist[iyP]*yTdist[iyT]
    end
end

## gross labor income distribution
grosslabincgridvec = grosslabincgrid[:]
grosslabincdistvec = grosslabincdist[:]
temp = [grosslabincgrid[:] grosslabincdist[:]]
temp = temp[sortperm(temp[:,1]),:]
glincgridvec = temp[:,1]
glincdistvec = temp[:,2]
glincdistcumvec = cumsum(glincdistvec)

## labor taxes and transfers function
meangrosslabinc = sum(glincgridvec.*glincdistvec)
totgrosslabinc = meangrosslabinc
labtaxthresh = lininterp1(glincdistcumvec,glincgridvec,labtaxthreshpc)
totgrosslabinchigh = sum(max.(glincgridvec .- labtaxthresh,0).*glincdistvec)

lumptransfer = labtaxlow*totgrosslabinc + labtaxhigh*totgrosslabinchigh

## net labor income
netlabincgrid = lumptransfer .+ (1-labtaxlow).*grosslabincgrid - labtaxhigh.*max.(grosslabincgrid .- labtaxthresh,0)

# ASSET GRIDS

xgrid = zeros(nx,nyP)

## cash on hand grids: different min points for each value of (iyP)
for iyP = 1:nyP
    xgrid[:,iyP] = range(0,1,length=nx)
    xgrid[:,iyP] = xgrid[:,iyP].^(1/xgrid_par)
    xgrid[:,iyP] = borrow_lim + minimum(netlabincgrid[iyP,:]) .+ (xmax-borrow_lim).*xgrid[:,iyP]
end

## savings grid for EGP
ns = nx
sgrid = range(0,1,length=nx)
sgrid = sgrid.^(1/xgrid_par)
sgrid = borrow_lim .+ (xmax-borrow_lim).*sgrid

## discount factor grid
if  nb == 1
    betagrid = [beta]
    betadist = [1.0]
    betatrans = [1.0]
elseif nb == 2
    betagrid = [beta-betawidth; beta+betawidth]
    betadist = [0.5; 0.5]
    betatrans = [1 .-betaswitch betaswitch; betaswitch 1 .-betaswitch]; # transitions on average once every 40 years
else
    error("nb must be 1 or 2")
end

betacumdist = cumsum(betadist)
betacumtrans = cumsum(betadist,dims=2)

# UTILITY FUNCTION, BEQUEST FUNCTION

if risk_aver==1
    u(c)= log.(c)
    beq(a) = bequest_weight.*log.(a.+bequest_luxury)
else
    u(c) = (c.^(1-risk_aver)-1)./(1-risk_aver)
    beq(a) = bequest_weight.*((a.+bequest_luxury).^(1-risk_aver)-1)./(1-risk_aver)
end

u1(c) = c.^(-risk_aver)
u1inv(u) = u.^(-1/risk_aver)

beq1(a) = bequest_weight.*(a.+bequest_luxury).^(-risk_aver)

# INITIALIZE CONSUMPTION FUNCTION

conguess = zeros(nx,nyP,nb)
for ib = 1:nb
    for iy = 1:nyP
        conguess[:,iy,ib] = r.*xgrid[:,iy]
    end
end

# ITERATE ON EULER EQUATION WITH ENDOGENOUS GRID POINTS

con = copy(conguess)

iter = 0
cdiff = 1000

while iter <= max_iter && cdiff>tol_iter
    iter = iter + 1
    global sav = zeros(nx,nyP,nb)
    emuc = zeros(ns,nyP,nb)

    conlast = copy(con)
    muc = u1(conlast) # muc on grid for x'

    coninterp = Array{Any}(undef,nyP,nb)
    mucinterp = Array{Any}(undef,nyP,nb)

    ## create interpolating function for each yP'
    for ib = 1:nb
        for iy = 1:nyP
            if InterpMUC==1
                mucinterp[iy,ib] = extrapolate(interpolate((xgrid[:,iy],), muc[:,iy,ib], Gridded(Linear())), Line())
            elseif InterpCon==1
                coninterp[iy,ib] = extrapolate(interpolate((xgrid[:,iy],), conlast[:,iy,ib], Gridded(Linear())), Line())
            end
        end
    end

    ## loop over current persistent income and discount factor
    for ib = 1:nb
        for iy = 1:nyP

            muc1 = zeros(nx,nyP,nb)
            con1 = zeros(nx,nyP,nb)
            cash1 = zeros(nx,nyP,nb)

            ## loop over future income realizations    and discount factor
            for ib2 = 1:nb
                for iyP2 = 1:nyP
                    for iyT2 = 1:nyT
                        if InterpMUC==1
                            emuc[:,iy,ib] = emuc[:,iy,ib] .+ mucinterp[iyP2,ib2](R.*sgrid[:] .+ netlabincgrid[iyP2,iyT2]) * yPtrans[iy,iyP2] * yTdist[iyT2] * betatrans[ib,ib2]
                        elseif InterpCon==1
                            emuc[:,iy,ib] = emuc[:,iy,ib] .+ u1(coninterp[iyP2,ib2](R.*sgrid[:] .+ netlabincgrid[iyP2,iyT2])) * yPtrans[iy,iyP2] * yTdist[iyT2] * betatrans[ib,ib2]
                        end
                    end
                end
            end
            muc1[:,iy,ib] = (1-dieprob) .*  betagrid[ib].*R.*emuc[:,iy,ib]./(1 .+savtax.*(sgrid.>=savtaxthresh)) .+ dieprob.*beq1(sgrid)
            con1[:,iy,ib] = u1inv(muc1[:,iy,ib])
            cash1[:,iy,ib] = con1[:,iy,ib] .+ sgrid .+ savtax.*sgrid.*(sgrid.>=savtaxthresh)

            ## loop over current period cash on hand
            for ix = 1:nx
                if xgrid[ix,iy]<cash1[1,iy,ib] # borrowing constraint binds
                    sav[ix,iy,ib] = borrow_lim
                else # borrowing constraint does not bind;
                    sav[ix,iy,ib] = lininterp1(cash1[:,iy,ib],sgrid,xgrid[ix,iy])
                end
            end
            con[:,iy,ib] = xgrid[:,iy] - sav[:,iy,ib]
        end
    end

    cdiff = maximum(abs.(con[:] - conlast[:]))
    if Display>=1
        println("Iteration no. " * string(iter), " max con fn diff is " * string(cdiff))
    end
end

# SIMULATE
if DoSimulate==1

    yPindsim = zeros(Int,Nsim,Tsim)
    logyTsim = zeros(Nsim,Tsim)
    logyPsim = zeros(Nsim,Tsim)
    ygrosssim = zeros(Nsim,Tsim)
    ynetsim = zeros(Nsim,Tsim)
    xsim = zeros(Nsim,Tsim)
    ssim = zeros(Nsim,Tsim)
    betaindsim = zeros(Int,Nsim,Tsim)
    betasim = zeros(Nsim,Tsim)

    ## indicators for dying
    diesim = dierand.<dieprob

    ## create interpolating functions
    savinterp = Array{Any}(undef,nyP,nb)
    for ib = 1:nb
        for iy = 1:nyP
            savinterp[iy,ib] = extrapolate(interpolate((xgrid[:,iy],), sav[:,iy,ib], Gridded(Linear())), Line())
        end
    end

    ## initialize permanent income;
    it = 1
    yPindsim[yPrand[:,it].<=yPcumdist[1],it] .= 1
    for iy = 2:nyP
        yPindsim[(yPrand[:,it].>yPcumdist[iy-1]) .& (yPrand[:,it].<=yPcumdist[iy]),it] .= iy
    end

    ## initialize discount factor;
    it = 1
    betaindsim[betarand[:,it].<=betacumdist[1],it] .= 1
    for ib = 2:nb
        betaindsim[(betarand[:,it].>betacumdist[ib-1]) .& (betarand[:,it].<=betacumdist[ib]),it] .= ib
    end

    ## loop over time periods
    for it = 1:Tsim
        if Display>=1 && mod(it,100)==0
            println("Simulating, time period " * string(it))
        end

        ## permanent income realization: note we vectorize simulations at once because
        ## of matlab, in other languages we would loop over individuals

        ## people who dont die: use transition matrix
        if it > 1
            yPindsim[(diesim[:,it].==0) .& (yPrand[:,it].<=yPcumtrans[yPindsim[:,it-1],1]),it] .= 1
            for iy = 2:nyP
                yPindsim[(diesim[:,it].==0) .& (yPrand[:,it].>yPcumtrans[yPindsim[:,it-1],iy-1]) .& (yPrand[:,it].<=yPcumtrans[yPindsim[:,it-1],iy]),it] .= iy
            end
        end

        ## people who die: descendendents draw income from stationary distribution
        yPindsim[(diesim[:,it].==1) .& (yPrand[:,it].<=yPcumdist[1]),it] .= 1
        for iy = 2:nyP
            yPindsim[(diesim[:,it].==1) .& (yPrand[:,it].>yPcumdist[iy-1]) .& (yPrand[:,it].<=yPcumdist[iy]),it] .= iy
        end

        logyPsim[:,it] = logyPgrid[yPindsim[:,it]]

        ## transitory income realization
        if nyT>1
            logyTsim[:,it] = - 0.5*sd_logyT^2 .+ yTrand[:,it].*sd_logyT
        else
            logyTsim[:,it] .= 0
        end
        ygrosssim[:,it] = exp.(logyPsim[:,it] .+ logyTsim[:,it])
        ynetsim[:,it] = lumptransfer .+ (1-labtaxlow).*ygrosssim[:,it] .- labtaxhigh.*max.(ygrosssim[:,it].-labtaxthresh,0)

        ## discount factor realization: note we vectorize simulations at once because
        ## of matlab, in other languages we would loop over individuals
        if it > 1
            betaindsim[betarand[:,it].<=betacumtrans[betaindsim[:,it-1],1],it] .= 1
            for ib = 2:nb
                betaindsim[(betarand[:,it].>betacumtrans[betaindsim[:,it-1],ib-1]) .& (betarand[:,it].<=betacumtrans[betaindsim[:,it-1],ib]),it] .= ib
            end
        end

        ## update cash on hand
        if it>1
            xsim[:,it] = R.*ssim[:,it-1] .+ ynetsim[:,it]
        end

        ## savings choice
        for ib = 1:nb
            for iy = 1:nyP
                ssim[(yPindsim[:,it].==iy) .& (betaindsim[:,it].==ib),it] = savinterp[iy,ib](xsim[(yPindsim[:,it].==iy) .& (betaindsim[:,it].==ib),it])
            end
        end
        ssim[ssim.<borrow_lim] .= borrow_lim
    end

    ## convert to assets
    asim = (xsim .- ynetsim)./R

    ## consumption
    csim = xsim .- ssim .- savtax.*max.(ssim.-savtaxthresh,0)
end

# MAKE PLOTS
if MakePlots==1

    ## consumption policy function
    p1 = plot([xgrid[:,1] xgrid[:,nyP]], [con[:,1,1] con[:,nyP,1]], xlims=(borrow_lim,xmax), title="Consumption Policy Function", color=[:blue :red], label=["Lowest income state" "Highest income state"])
    display(p1)

    ## savings policy function
    p2 = plot([xgrid[:,1] xgrid[:,nyP]], [sav[:,1,1]./xgrid[:,1] sav[:,nyP,1]./xgrid[:,nyP]], xlims=(borrow_lim,xmax), title="Savings Policy Function s/x", color=[:blue :red], label=["Lowest income state" "Highest income state"], legend=false)
    plot!(sgrid, zeros(nx,1), color=:black, lw=0.5)
    display(p2)

    ## nice zoom
    xlimits = (0,4)
    xlimind = trues(nx,nyP)
    for iy = 1:nyP
        if minimum(xgrid[:,iy]) < xlimits[1]
            xlimind[:,iy] = xlimind[:,iy] .& (xgrid[:,iy].>=maximum(xgrid[xgrid[:,iy].<xlimits[1],iy]))
        elseif minimum(xgrid[:,iy]) > xlimits[2]
            xlimind[:,iy] .= 0
        end
        if maximum(xgrid[:,iy]) > xlimits[2]
            xlimind[:,iy] = xlimind[:,iy] .& (xgrid[:,iy].<=minimum(xgrid[xgrid[:,iy].>xlimits[2],iy]))
        elseif maximum(xgrid[:,iy]) < xlimits[1]
            xlimind[:,iy] .= 0
        end
    end

    ## consumption policy function: zoomed in
    p3 = plot(xgrid[xlimind[:,1],1],con[xlimind[:,1],1,1], xlims=xlimits, title="Consumption: Zoomed", marker=:circle, color=:blue, linewidth=2, legend=false)
    plot!(xgrid[xlimind[:,nyP],nyP], con[xlimind[:,nyP],nyP,1], color=:red, linewidth=2)
    display(p3)

    ## savings policy function: zoomed in
    p4 = plot(xgrid[xlimind[:,1],1],sav[xlimind[:,1],1,1]./xgrid[xlimind[:,1],1], xlims=xlimits, title="Consumption: Zoomed", marker=:circle, color="blue", linewidth=2, legend=false)
    plot!(xgrid[xlimind[:,nyP],nyP], sav[xlimind[:,nyP],nyP,1]./xgrid[xlimind[:,nyP],nyP], color=:red, linewidth=2)
    plot!(sgrid, zeros(nx,1), color=:black, lw=0.5)
    display(p4)

    ## income distribution
    p5 = histogram(ygrosssim[:,Tsim], nbins=50, title="Income distribution", color=RGB(0,0.5,0.5), linecolor=:blue, legend=false)
    display(p5)

    ## asset distribution
    p6 = histogram(asim[:,Tsim], nbins=100, title="Asset distribution", color=RGB(.7,.7,.7), linecolor=:black, legend=false)
    display(p6)

    ## convergence check
    p7 = plot(1:Tsim, mean(asim,dims=1)', title="Mean Asset Convergence", xlabel="Time Period", color=:black, linewidth=1.5, legend=false)
    display(p7)

    ## asset distribution statistics
    aysim = asim[:,Tsim]./meangrosslabinc
    println("Mean assets (relative to mean income): " * string(mean(aysim)))
    println("Fraction borrowing constrained: " * string(sum(aysim.==borrow_lim)/Nsim * 100) * '%')
    println("10th Percentile: " * string(quantile(aysim,.1)))
    println("50th Percentile: " * string(quantile(aysim,.5)))
    println("90th Percentile: " * string(quantile(aysim,.9)))
    println("99th Percentile: " * string(quantile(aysim,.99)))

    ## table of parameters
    tabparam  = zeros(4,1)
    tabparam[1] = beta # discount factor
    tabparam[2] = risk_aver # risk aversion
    tabparam[3] = labtaxlow # labor tax rate
    tabparam[4] = lumptransfer/meangrosslabinc # relative to mean gross lab inc

    ## table of statistics
    tabstat = zeros(16,1)
    tabstat[1] = var(log.(ygrosssim[:,Tsim-1])) # variance log gross earnings
    tabstat[2] = ginicoeff(ygrosssim[:,Tsim-1]) # gini gross earnings
    tabstat[3] = var(log.(ynetsim[:,Tsim-1])) # variance log net earnings
    tabstat[4] = ginicoeff(ynetsim[:,Tsim-1]) # gini net earnings
    tabstat[5] = var(log.(csim[:,Tsim-1])) # variance log consumption
    tabstat[6] = ginicoeff(csim[:,Tsim-1]) # gini consumption
    tabstat[7] = mean(aysim) # mean wealth (relative to mean earnings)
    tabstat[8] = quantile(aysim,.5) # median wealth (relative to mean earnings)
    tabstat[9] = ginicoeff(aysim) # wealth gini
    tabstat[10] = quantile(aysim,.9) / quantile(aysim,.5) # 90-10 wealth distribution
    tabstat[11] = quantile(aysim,.99) / quantile(aysim,.5) # 99-50 wealth distribution
    tabstat[12] = sum(aysim.<=0)/Nsim # fraction  less than equal to zero
    tabstat[13] = sum(aysim.<=0.05)/Nsim # fraction  less than equal to 5% av. annual gross earnings
    tabstat[14] = sum(aysim[aysim.>=quantile(aysim,.9)]) / sum(aysim) # top 10% wealth share
    tabstat[15] = sum(aysim[aysim.>=quantile(aysim,.99)]) / sum(aysim) # top 1% wealth share
    tabstat[16] = sum(aysim[aysim.>=quantile(aysim,.999)]) / sum(aysim) # top 0.1% wealth share

    taboutput = [tabparam; tabstat]
end

# COMPUTE MPCs
if ComputeMPC==1

    ## theoretical mpc lower bound
    mpclim = R*((beta*R)^(-1/risk_aver))-1

    coninterp = Array{Any}(undef,nyP)
    mpc1 = zeros(nx,nyP)
    mpc2 = zeros(nx,nyP)

    for iy = 1:nyP
        ## create interpolating function
        coninterp[iy] = extrapolate(interpolate((xgrid[:,iy],), con[:,iy], Gridded(Linear())), Line())

        mpc1[:,iy] = ( coninterp[iy](xgrid[:,iy].+mpcamount1) - con[:,iy] ) ./ mpcamount1
        mpc2[:,iy] = ( coninterp[iy](xgrid[:,iy].+mpcamount2) - con[:,iy] ) ./ mpcamount2

    end

    ## mpc functions
    p8 = plot([xgrid xgrid],[mpc1[:,Int((nyP+1)/2)] mpc2[:,Int((nyP+1)/2)]], xlims=(0,10), linestyle=[:solid :dash], color=:blue, label=["Amount 1" fill("",(1,nyP-1)) "Amount 2" fill("",(1,nyP-1))])
    plot!(xgrid[:,1],mpclim.*ones(size(xgrid[:,1])),color=:black,linestyle=:dot,linewidth=2,label="Theoretical MPC limit = " * string(mpclim))
    display(p8)

    ## mpc distribution
    mpc1sim = zeros(Nsim,1)
    mpc2sim = zeros(Nsim,1)
    for iy = 1:nyP
        mpc1sim[yPindsim[:,Tsim].==iy] = ( coninterp[iy](xsim[yPindsim[:,Tsim].==iy,Tsim].+mpcamount1) - coninterp[iy](xsim[yPindsim[:,Tsim].==iy,Tsim]) ) ./ mpcamount1
        mpc2sim[yPindsim[:,Tsim].==iy] = ( coninterp[iy](xsim[yPindsim[:,Tsim].==iy,Tsim].+mpcamount2) - coninterp[iy](xsim[yPindsim[:,Tsim].==iy,Tsim]) ) ./ mpcamount2
    end

    p9 = histogram(mpc1sim,bins=0:0.02:1.5, title="MPC distribution",color=RGB(.7,.7,.7),legend=false)
    display(p9)

    # mpc distribution statistics
    println("Mean MPC amount 1: " * string(mean(mpc1sim)))
    println("Mean MPC amount 2: " * string(mean(mpc2sim)))
end


text:
using DifferentiableStateSpaceModels, Symbolics, LinearAlgebra, Test
using DifferentiableStateSpaceModels.Examples

@testset "SGU First Order" begin
    m = @include_example_module(Examples.sgu)
    p_d = (γ = 2.0, ω = 1.455, ρ = 0.42, σe = 0.0129, δ = 0.1, ψ = 0.000742, α = 0.32,
           ϕ = 0.028, β = 1.0 / (1.0 + 0.04), r_w = 0.04, d_bar = 0.7442)
    p_f = nothing

    c = SolverCache(m, Val(1), p_d)
    sol = generate_perturbation(m, p_d, p_f; cache = c)
    generate_perturbation_derivatives!(m, p_d, p_f, c)

    # TODO check that these are actually the correct results, currently only a regression test
    # Also, does not check the derivative details in the cache

    @test sol.retcode == :Success
    @test sol.A ≈ [0.42 1.7519543732312002e-17 -4.1734466435163804e-17 9.83532499523268e-19;
                   0.9006855710385308 0.9738235409647132 -1.685401161087487 0.027873826122536263;
                   0.6721065922179434 -0.006588281444814813 0.5003113265680597 -0.0001885768865857994;
                   0.016707717342759755 0.01806442668489436 -0.031264191538169755 0.000517059474573019]
    @test sol.C ≈
          [0.9006855710385274 0.9738235409647122 -1.6854011610874828 0.02787382612253622;
           1.2604305985051094 -0.03920634002915657 0.50643088763379 -0.001122206086526858;
           1.2903225806451621 -1.1817797818825905e-16 0.4129032258064517 -1.1862432475956954e-18;
           1.8774193548387106 4.1726226271227295e-17 0.600774193548387 6.519722854879265e-19;
           6.72106592217944 -0.06588281444815657 -3.996886734319368 -0.0018857688658582676;
           0.6721065922179433 -0.006588281444815545 0.5003113265680625 -0.00018857688658582318;
           -0.4905618464963626 0.207427359629203 -0.7024101653193269 0.005937209270617432;
           -0.643512113444817 0.04451868414563279 1.1217837540922557 0.0012742617208145878;
           -0.6059154121590582 0.017609608142098897 1.1338146985036983 -0.018751472641015628;
           0.0006683086937103991 0.0007225770673958052 -0.0012505676615268426 2.06823789829215e-5;
           0.01670771734276043 0.018064426684894292 -0.03126419153816984 0.0005170594745730165;
           1.0 0.0 0.0 0.0; 0.0 1.0 0.0 0.0; 0.0 0.0 1.0 0.0; 0.0 0.0 0.0 1.0]
    @test sol.Q == I
    @test sol.y ≈ [0.7442, 0.11060245636938473, 0.007390615600776165, 0.3964158265111139,
                   -1.0794906932984638, 1.2230943996955819, 1.7243861964374876, -3.5143212879378156,
                   0.0, 0.0, -3.2188758248681983]
    @test sol.x ≈ [0.0, 0.7442, 1.2230943996955819, -3.2188758248681983]
    @test sol.B ≈ [-0.0129; 0.0; 0.0; 0.0]
    @test sol.D === nothing
    @test sol.g_x ≈
          [0.9006855710385274 0.9738235409647122 -1.6854011610874828 0.02787382612253622;
           1.2604305985051094 -0.03920634002915657 0.50643088763379 -0.001122206086526858;
           1.2903225806451621 -1.1817797818825905e-16 0.4129032258064517 -1.1862432475956954e-18;
           1.8774193548387106 4.1726226271227295e-17 0.600774193548387 6.519722854879265e-19;
           6.72106592217944 -0.06588281444815657 -3.996886734319368 -0.0018857688658582676;
           0.6721065922179433 -0.006588281444815545 0.5003113265680625 -0.00018857688658582318;
           -0.4905618464963626 0.207427359629203 -0.7024101653193269 0.005937209270617432;
           -0.643512113444817 0.04451868414563279 1.1217837540922557 0.0012742617208145878;
           -0.6059154121590582 0.017609608142098897 1.1338146985036983 -0.018751472641015628;
           0.0006683086937103991 0.0007225770673958052 -0.0012505676615268426 2.06823789829215e-5;
           0.01670771734276043 0.018064426684894292 -0.03126419153816984 0.0005170594745730165]
    @test sol.η ≈ [-1; 0; 0; 0]
    @test sol.Γ ≈ [0.0129]
end

@testset "SGU Second Order" begin
    m = @include_example_module(Examples.sgu)
    p_d = (γ = 2.0, ω = 1.455, ρ = 0.42, σe = 0.0129, δ = 0.1, ψ = 0.000742, α = 0.32,
           ϕ = 0.028, β = 1.0 / (1.0 + 0.04), r_w = 0.04, d_bar = 0.7442)
    p_f = nothing

    c = SolverCache(m, Val(2), p_d)
    sol = generate_perturbation(m, p_d, p_f, Val(2); cache = c)
    generate_perturbation_derivatives!(m, p_d, p_f, c)  # Solves and fills the cache

    # TODO check that these are actually the correct results, currently only a regression test
    # Also, does not check the derivative details in the cache
    @test sol.retcode == :Success
    @test sol.Q == I
    @test sol.y ≈ [0.7442, 0.11060245636938473, 0.007390615600776165, 0.3964158265111139,
                   -1.0794906932984638, 1.2230943996955819, 1.7243861964374876, -3.5143212879378156,
                   0.0, 0.0, -3.2188758248681983]
    @test sol.x ≈ [0.0, 0.7442, 1.2230943996955819, -3.2188758248681983]
    @test sol.B ≈ [-0.0129; 0.0; 0.0; 0.0]
    @test sol.D === nothing
    @test sol.g_x ≈
          [0.9006855710385274 0.9738235409647122 -1.6854011610874828 0.02787382612253622;
           1.2604305985051094 -0.03920634002915657 0.50643088763379 -0.001122206086526858;
           1.2903225806451621 -1.1817797818825905e-16 0.4129032258064517 -1.1862432475956954e-18;
           1.8774193548387106 4.1726226271227295e-17 0.600774193548387 6.519722854879265e-19;
           6.72106592217944 -0.06588281444815657 -3.996886734319368 -0.0018857688658582676;
           0.6721065922179433 -0.006588281444815545 0.5003113265680625 -0.00018857688658582318;
           -0.4905618464963626 0.207427359629203 -0.7024101653193269 0.005937209270617432;
           -0.643512113444817 0.04451868414563279 1.1217837540922557 0.0012742617208145878;
           -0.6059154121590582 0.017609608142098897 1.1338146985036983 -0.018751472641015628;
           0.0006683086937103991 0.0007225770673958052 -0.0012505676615268426 2.06823789829215e-5;
           0.01670771734276043 0.018064426684894292 -0.03126419153816984 0.0005170594745730165]
    @test sol.η ≈ [-1; 0; 0; 0]
    @test sol.Γ ≈ [0.0129]

    @test sol.A_0 ≈ [-3.9402619560802135e-20, 0.0006391584993811839, 0.00015405791020935275,
                     1.185639016352093e-5]
    @test sol.A_1 ≈
          [0.42 1.7519543732312002e-17 -4.1734466435163804e-17 9.83532499523268e-19;
           0.9006855710385308 0.9738235409647132 -1.685401161087487 0.027873826122536263;
           0.6721065922179434 -0.006588281444814813 0.5003113265680597 -0.0001885768865857994;
           0.016707717342759755 0.01806442668489436 -0.031264191538169755 0.000517059474573019]
    @test sol.A_2[:, :, 1] ≈ [0.0 0.0 0.0 0.0;
                              -0.6058913712160838 -0.0031020337841137775 -0.40284701288926217 -8.878975162048575e-5;
                              -0.03550021179812225 0.0004999494367584898 -0.14268843831504846 1.4310091185952039e-5;
                              -0.0038546588776544625 0.007926733837883587 -0.021291238555819474 0.0002268875123904971]
    @test sol.A_2[:, :, 2] ≈ [0.0 0.0 0.0 0.0;
                              -0.0031020337841140833 -0.012114331571761804 0.026596726056904527 0.01838062634333291;
                              0.0004999494367584289 -0.0030632500289434696 0.007620095670675804 -0.00021437736130562276;
                              0.007926733837883589 0.008407899373508049 -0.014447148738243992 0.0005880527713927997] # add more slices
    @test sol.C_0 ≈ [0.0006391584993811839, 0.0001036018842634584, -6.50134529592093e-21,
                     -4.420914801226232e-21, 0.0015405791020935281, 0.0001540579102093528,
                     -0.0005481221988433036, -0.00042997911595382926, -0.0004299791159538292,
                     4.7425560654083843e-7, 1.185639016352093e-5, 0.0, 0.0, 0.0, 0.0]
    @test sol.C_1 ≈
          [0.9006855710385274 0.9738235409647122 -1.6854011610874828 0.02787382612253622;
           1.2604305985051094 -0.03920634002915657 0.50643088763379 -0.001122206086526858;
           1.2903225806451621 -1.1817797818825905e-16 0.4129032258064517 -1.1862432475956954e-18;
           1.8774193548387106 4.1726226271227295e-17 0.600774193548387 6.519722854879265e-19;
           6.72106592217944 -0.06588281444815657 -3.996886734319368 -0.0018857688658582676;
           0.6721065922179433 -0.006588281444815545 0.5003113265680625 -0.00018857688658582318;
           -0.4905618464963626 0.207427359629203 -0.7024101653193269 0.005937209270617432;
           -0.643512113444817 0.04451868414563279 1.1217837540922557 0.0012742617208145878;
           -0.6059154121590582 0.017609608142098897 1.1338146985036983 -0.018751472641015628;
           0.0006683086937103991 0.0007225770673958052 -0.0012505676615268426 2.06823789829215e-5;
           0.01670771734276043 0.018064426684894292 -0.03126419153816984 0.0005170594745730165;
           1.0 0.0 0.0 0.0; 0.0 1.0 0.0 0.0; 0.0 0.0 1.0 0.0; 0.0 0.0 0.0 1.0]
    @test sol.C_2[:, :, 1] ≈
          [-0.6058913712160838 -0.003102033784113778 -0.40284701288926217 -8.878975162048576e-5;
           0.3641856225490242 0.02778598688564625 0.041906068536917314 0.0007953204400113982;
           -9.250270498653839e-15 -1.6521083580989144e-16 2.6335403233441523e-15 -2.259643246988632e-18;
           -6.2196715024389655e-15 -2.032077242658297e-16 2.0907564007853216e-15 -2.070244672048874e-18;
           -20.682729326608023 0.20426072693774494 13.686097933010494 0.005846570499509324;
           -0.035500211798123044 0.0004999494367584962 -0.14268843831504868 1.431009118595195e-5;
           -0.2699169932252056 -0.04172126179447118 0.03163930042536679 -0.0011941908856709196;
           1.5804491334669595 -0.039703298061926455 -0.6000140721981689 -0.0011364305545266838;
           1.545156746130972 -0.014443488006995794 -0.6113076361456838 0.017661920088352694;
           -0.0001486033987300181 0.00032310566021555423 -0.0008620966077361734 9.248278166631295e-6;
           -0.0038546588776544816 0.00792673383788359 -0.02129123855581945 0.00022688751239049705;
           0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0] # TODO, add more slices
    @test sol.g_xx[:, :, 1] ≈
          [-1.2117827424321677 -0.006204067568227556 -0.8056940257785243 -0.00017757950324097152;
           0.7283712450980484 0.0555719737712925 0.08381213707383463 0.0015906408800227963;
           -1.8500540997307678e-14 -3.304216716197829e-16 5.2670806466883046e-15 -4.519286493977264e-18;
           -1.2439343004877931e-14 -4.064154485316594e-16 4.181512801570643e-15 -4.140489344097748e-18;
           -41.365458653216045 0.4085214538754899 27.372195866020988 0.011693140999018648;
           -0.07100042359624609 0.0009998988735169924 -0.28537687663009736 2.86201823719039e-5;
           -0.5398339864504113 -0.08344252358894236 0.06327860085073359 -0.002388381771341839;
           3.160898266933919 -0.07940659612385291 -1.2000281443963379 -0.0022728611090533676;
           3.090313492261944 -0.028886976013991587 -1.2226152722913677 0.03532384017670539;
           -0.0002972067974600362 0.0006462113204311085 -0.0017241932154723467 1.849655633326259e-5;
           -0.007709317755308963 0.01585346767576718 -0.0425824771116389 0.0004537750247809941] # TODO add more slices
    @test sol.g_σσ ≈ [0.0012783169987623678, 0.0002072037685269168, -1.300269059184186e-20,
                      -8.841829602452464e-21, 0.0030811582041870562, 0.0003081158204187056,
                      -0.0010962443976866073, -0.0008599582319076585, -0.0008599582319076584,
                      9.485112130816769e-7, 2.371278032704186e-5]
end

ext: jl


text:
# Pairwise Gap Percentage
# =======================

"It calculates the gap intersection as percentage from a table of `Counts`."
function gap_intersection_percentage(nxy::Counts{T,2,GappedAlphabet}) where T
    T(100.0) * gettablearray(nxy)[21,21] / gettotal(nxy)
end

"It calculates the gap union as percentage from a table of `Counts`."
function gap_union_percentage(nxy::Counts{T,2,GappedAlphabet}) where T
    marginals = getmarginalsarray(nxy)
    T(100.0) * (marginals[21,1] + marginals[21,2] - gettablearray(nxy)[21,21]) / gettotal(nxy)
end

# MIToS Pairwise Gap Percentage
# =============================

"""
It takes a MSA or a file and a `FileFormat` as first arguments. It calculates the percentage
of gaps on columns pairs (union and intersection) using sequence clustering (Hobohm I).

Argument, type, default value and descriptions:

```
    - clustering  Bool      true    Sequence clustering (Hobohm I)
    - threshold             62      Percent identity threshold for sequence clustering (Hobohm I)
```

This function returns:

```
    - pairwise gap union as percentage
    - pairwise gap intersection as percentage
```
"""
function pairwisegapfraction(aln::AbstractMatrix{Residue}; clustering::Bool=true, threshold=62)
    clusters = clustering ? hobohmI(aln, threshold) : NoClustering()
    table = Counts(ContingencyTable(Float64,Val{2},GappedAlphabet()))
    gu = mapcolpairfreq!(gap_union_percentage, aln, table, Val{true}, weights=clusters)
    gi = mapcolpairfreq!(gap_intersection_percentage, aln, table, Val{true}, weights=clusters)
    gu, gi
end

function pairwisegapfraction(filename::String, format::Type{T}; kargs...) where T <: FileFormat
    aln = read(filename, T, AnnotatedMultipleSequenceAlignment, generatemapping=true)
    pairwisegapfraction(aln; kargs...)
end


text:
# API demo
# This is a stripped-down version for the layout here: http://postimage.org/image/s6pskvknf/
# See scatter_image.jl for a more comprehensive implementation. This one
# is here to provide a relatively simple test-case for early development.
# It cuts several corners:
#    hard-codes the minimum space needed for the axis labels
#    uses pixels for all measurements

require("Layout")
using Layout

# Create a fake window
win = Window(0, 0, 800, 600)

# Create the overall layout inside win. This is specified in terms of named tabstops.
xtabs = [:scatterL, :scatterR, :margeEL, :margeER, :imageL, :imageR]
ytabs = [:margeNT, :margeNB, :scatterT, :scatterB, :buttonsT]
l = LayoutLP(win, xtabs, ytabs)

# Create a sub-layout for the buttons
# After the parent, the next two arguments refer to parent coordinates.
# The last two arguments introduce tab stops in the child
# lbuttons = LayoutLP(l, [:W, :E], [:buttonsT, :S], [:cancelL, :cancelR, :doneL], [])

# Specify image dimensions
imsz = [384,512]
imheight = () -> imsz[1]
imwidth = () -> imsz[2]
# These are only nominally hard-coding the image size; you could change
# imsz[1] or imsz[2] and it would update to the new dimensions.

# Now specify the geometry. This is done by adding constraints and penalties. Note that a "subplot" function would take care of many of these for you, but here we're going low-level because this GUI has a lot of custom requirements.

# First we require that the tabstops are in increasing order. One can pick out
# a subset of the tab stops and do something different for the rest
# (partial ordering)
increasing(l, xtabs)
increasing(l, ytabs)

# Now add particular constraints
addconstraints(l,
        # Ensure the proper aspect ratio for the image
        :($imheight()*(imageR-imageL) == $imwidth()*(scatterB-scatterT)),
        # Leave enough room for scatterplot y labels
        :(scatterL-W > 40),
        # Scatterplot bottom label
        :(buttonsT-scatterB > 40),
        # Ensure a small gap between east marginal axis and the image axis
        :(imageL > margeER + 10)
    )
# addconstraints(lbuttons, :(S-N > 25))
addconstraints(l, :(S - buttonsT > 25))

# Yet more constraints
addconstraints(l,
        # we want lots of space for the primary displays
        :(soft(imageR-imageL>0.3*(E-W), 1.0)),
        :(soft(scatterR-scatterL>0.4*(E-W), 1.0)),
        # we'd like the marginal axes to be substantially thinner than the scatterplot
        :(soft(margeER-margeEL == (scatterR-scatterL)/3, 1.0)),
        :(soft(margeNB-margeNT == (scatterB-scatterT)/3, 1.0)),
        # small gaps for the marginal axes
        :(margeEL-scatterR>5),
        :(scatterT-margeNB>5)
    )

# Reward leaving as much vertical space as possible
# addobjective(l, :(100*(scatterB-scatterT)))

lpp, stuffer = lpparse(l)
lpsolve(lpp, stuffer)

@show l[xtabs]
@show l[ytabs]


text:
using ImaginePlots
using ImagineInterface, Unitful, Plots
using Test

pyplot()

sigs = rigtemplate("ocpi-2")
p = getpositioners(sigs)[1]
append!(p, "test", Unitful.V*rand(100))
@test_nowarn plot(p)
@test_nowarn plot([p;p])


text:
#!/usr/bin/env julia
#-*- mode: julia -*-
#-*- code: utf-8 -*-

const rFlag = r"^##" :: Regex
const rName = r"^(.+)\.md$" :: Regex

function splitF(f :: IOStream)
    tmp = []
    out = []

    while !eof(f)
        x = readline(f)
        if ismatch(rFlag, x)
            push!(out, tmp)
            tmp = []
        end
        push!(tmp, x)
    end

    push!(out, tmp)
end

function writeF(fn :: AbstractString)
    return function(list :: Array{Any})
        rm(fn, force=true, recursive=true)
        mkdir(fn)

        @sync @parallel for i=1:length(list)
            write(string(fn, "/", i), list[i])
        end
    end
end

function main(f :: AbstractString)
    if !ismatch(rName, f)
        error("Input is not a markdown file.")
    end
    fn = match(rName, f)[1]
    open(f) |> splitF |> writeF(fn)
end

ARGS[1] |> main

max_stars_repo_name: yuffiy/me

max_issues_repo_name: yuffiy/me

max_forks_repo_name: yuffiy/me


text:
include("CalEqOdds.jl")
include("DisparateImpact.jl")
include("adversarial.jl")


text:
using Distributed
addprocs(2)

@everywhere using CharibdeOptim
@everywhere using ModelingToolkit
@everywhere using IntervalArithmetic

# --- Michalewicz Problem ---

@everywhere function michalewicz(X)
      sum = 0
      for i in 1:length(X)
            sum = sum + sin(X[i])*(sin(i*(X[i]^2)/π))^20
      end
      return -sum
end

(global_min, minimisers, information) = charibde_min(michalewicz, IntervalBox(0..π, 50))

# --- Sine Envolope Problem ---

@everywhere function sine_envolope(X)
      sum = 0
      for i in 1:length(X)-1
            temp = X[i+1]^2 + X[i]^2
            sum = sum + 0.5 + (sin(√temp -0.5)/(0.001*temp + 1))^2
      end
      return -sum
end

(global_min, minimisers, information) = charibde_min(sine_envolope, IntervalBox(-100..100, 5))

# --- Shekel Problem ---

@everywhere function shekel(X)
      sum = 0
      for i in 1:30
            s = 0
            for j in 1:length((X))
                  s = s + (X[j] - i*j)^2
            end
            sum = sum + (1 / (i + s))
      end
      return -sum
end

(global_min, minimisers, information) = charibde_min(shekel, IntervalBox(0..10, 5))

# --- Egg-Holder Problem ---

@everywhere function egg_holder(X)
      sum = 0
      for i in 1:length(X)-1
            sum = sum + ((X[i+1] + 47)sin(√abs(X[i+1] + 47 + X[i]/2)) + X[i]sin(√abs(X[i] - X[i+1] -47)))
      end
      return -sum
end

(global_min, minimisers, information) = charibde_min(egg_holder, IntervalBox(-512..512, 5))

# --- Rana Problem ---

@everywhere function rana(X)
      sum = 0
      for i in 1:length(X)-1
            temp1 = √abs(X[i+1] + X[i] + 1)
            temp2 = √abs(X[i+1] - X[i] + 1)
            sum = sum + (X[i]*cos(temp1) * sin(temp2) + (1+ X[i+1]) * sin(temp1) * cos(temp2))
      end
      return -sum
end

(global_min, minimisers, information) = charibde_min(rana , IntervalBox(-512..512, 5))

# --- Keane Problem ---

@everywhere function keane(X)
      temp_sum = 0
      temp_pro = 1
      temp_sum1 = 0
      for i in 1:length(X)
            temp_sum = temp_sum + cos(X[i])^4
            temp_pro = temp_pro * cos(X[i])^2
            temp_sum1 = temp_sum1 + i * X[i]^2
      end
      sum = abs(temp_sum - 2 * temp_pro) / √temp_sum1
      return -sum
end

@everywhere vars = ModelingToolkit.@variables x y w z
@everywhere C1 = constraint(vars , x*y*w*z, Interval(0.75, Inf))
@everywhere C2 = constraint(vars , x+y+z+w, Interval(-Inf, 30))

(global_min, minimisers, information) = charibde_min(keane, IntervalBox(0..10, 4), [C1, C2])


text:
"""
$(TYPEDEF)
"""
struct LinearSolution{T,N,uType,R,P,A} <: AbstractLinearSolution{T,N}
  u::uType
  resid::R
  prob::P
  alg::A
  retcode::Symbol
end

function build_solution(prob::AbstractLinearProblem,
                        alg,u,resid;calculate_error = true,
                        retcode = :Default, kwargs...)

  T = eltype(eltype(u))
  N = length((size(u)...,))

  LinearSolution{T,N,typeof(u),typeof(resid),typeof(prob),typeof(alg)}(u,resid,prob,alg,retcode)
end

struct NonlinearSolution{T,N,uType,R,P,A} <: AbstractNonlinearSolution{T,N}
  u::uType
  resid::R
  prob::P
  alg::A
  retcode::Symbol
end

function build_solution(prob::AbstractNonlinearProblem,
                        alg,u,resid;calculate_error = true,
                        retcode = :Default, kwargs...)

  T = eltype(eltype(u))
  N = length((size(u)...,))

  NonlinearSolution{T,N,typeof(u),typeof(resid),typeof(prob),typeof(alg)}(u,resid,prob,alg,retcode)
end

struct QuadratureSolution{T,N,uType,R,P,A,C} <: AbstractQuadratureSolution{T,N}
  u::uType
  resid::R
  prob::P
  alg::A
  retcode::Symbol
  chi::C
end

function build_solution(prob::AbstractQuadratureProblem,
                        alg,u,resid;calculate_error = true,
                        chi = nothing,
                        retcode = :Default, kwargs...)

  T = eltype(eltype(u))
  N = length((size(u)...,))

  QuadratureSolution{T,N,typeof(u),typeof(resid),
                     typeof(prob),typeof(alg),typeof(chi)}(
                     u,resid,prob,alg,retcode,chi)
end


text:
module Day14

export get_inputs, get_solution1, get_solution2

using ..Advent2020.Utils
using SparseArrays

## Input getting
function get_inputs()
    test_input1 = read_input(IOBuffer(
        """
        mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X
        mem[8] = 11
        mem[7] = 101
        mem[8] = 0
        """
    ))
    test_output1 = 165
    test_input2 = read_input(IOBuffer(
        """
        mask = 000000000000000000000000000000X1001X
        mem[42] = 100
        mask = 00000000000000000000000000000000X0XX
        mem[26] = 1
        """
    ))
    test_output2 = 208
    data = read_input(joinpath(@__DIR__, "input.txt"))

    return (; test_input1, test_input2, test_output1, test_output2, data)
end

function read_input(io)
    data = read(io, String)
    chunks = split(data, "mask = ", keepempty=false)
    return parse_chunk.(chunks)
end

function parse_chunk(chunk)
    lines = split(chunk, r"mask = |\n|\r\n", keepempty=false)
    mask = lines[1]
    mem = parse_mem.(lines[2:end])
    return (; mask, mem)
end

function parse_mem(mem)
    address, value = parse.(Int, split(mem, r"mem\[|\] = ", keepempty=false))
    return address => value
end

## Solution functions
initialize_memory(data) = spzeros(Int, typemax(Int))

mask_replace(mask, x='0') = parse(Int, replace(mask, 'X'=>x), base=2)

apply_mask(value, mask) = (value | mask_replace(mask, '0')) & mask_replace(mask, '1')

# Part 1
function get_solution1(data)
    memory = initialize_memory(data)

    for chunk in data, (address, value) in chunk.mem
        memory[address] = apply_mask(value, chunk.mask)
    end

    return sum(memory)
end

# Part 2
function get_solution2(data)
    memory = initialize_memory(data)

    for line in data
        mask = zero_based(replace(collect(line.mask), 'X'=>'C', '0'=>'X'))
        C_positions = findall(==('C'), mask)
        num_C = length(C_positions)

        for short_mask in 0:(2^num_C-1)
            short_mask = last(bitstring(short_mask), num_C)
            mask[C_positions] = zero_based(collect(short_mask))
            mask_string = String(mask)

            for (encoded_address, value) in line.mem
                address = apply_mask(encoded_address, mask_string)
                memory[address] = value
            end
        end
    end

    return sum(memory)
end

end


text:
module Transpositions

import LinearAlgebra: transpose!

using TimerOutputs
import MPI

using ..PencilArrays
using ..Pencils: ArrayRegion
using StaticPermutations

# Declare transposition approaches.
abstract type AbstractTransposeMethod end

struct PointToPoint <: AbstractTransposeMethod end
struct Alltoallv <: AbstractTransposeMethod end

function Base.show(io::IO, ::T) where {T<:AbstractTransposeMethod}
    print(io, nameof(T))
end

"""
    Transposition

Holds data for transposition between two pencil configurations.

---

    Transposition(dest::PencilArray{T,N}, src::PencilArray{T,N};
                  method = Transpositions.PointToPoint())

Prepare transposition of arrays from one pencil configuration to the other.

The two pencil configurations must be compatible for transposition:

- they must share the same MPI Cartesian topology,

- they must have the same global data size,

- when written as a sorted tuple, the decomposed dimensions must be almost the
  same, with at most one difference. For instance, if the input of a 3D dataset
  is decomposed in `(2, 3)`, then the output may be decomposed in `(1, 3)`, but
  not in `(1, 2)`. If the decomposed dimensions are the same, then no
  transposition is performed, and data is just copied if needed.

The `src` and `dest` arrays may be aliased (they can share memory space).

# Performance tuning

The `method` argument allows to choose between transposition implementations.
This can be useful to tune performance of MPI data transfers.
Two values are currently accepted:

- `Transpositions.PointToPoint()` uses non-blocking point-to-point data transfers
  (`MPI_Isend` and `MPI_Irecv`).
  This may be more performant since data transfers are interleaved with local
  data transpositions (index permutation of received data).
  This is the default.

- `Transpositions.Alltoallv()` uses collective `MPI_Alltoallv` for global data
  transpositions.
"""
struct Transposition{T, N,
                     PencilIn  <: Pencil,
                     PencilOut <: Pencil,
                     ArrayIn   <: PencilArray{T,N},
                     ArrayOut  <: PencilArray{T,N},
                     M  <: AbstractTransposeMethod,
                    }
    Pi :: PencilIn
    Po :: PencilOut
    Ai :: ArrayIn
    Ao :: ArrayOut
    method :: M
    dim :: Union{Nothing,Int}  # dimension along which transposition is performed
    send_requests :: Vector{MPI.Request}

    function Transposition(Ao::PencilArray{T,N}, Ai::PencilArray{T,N};
                           method = PointToPoint()) where {T,N}
        Pi = pencil(Ai)
        Po = pencil(Ao)

        # Verifications
        if extra_dims(Ai) !== extra_dims(Ao)
            throw(ArgumentError(
                "incompatible number of extra dimensions of PencilArrays: " *
                "$(extra_dims(Ai)) != $(extra_dims(Ao))"))
        end

        assert_compatible(Pi, Po)

        # The `decomp_dims` tuples of both pencils must differ by at most one
        # value (as just checked by `assert_compatible`). The transposition
        # is performed along the dimension R where that difference happens.
        dim = findfirst(decomposition(Pi) .!= decomposition(Po))

        reqs = MPI.Request[]

        new{T, N, typeof(Pi), typeof(Po), typeof(Ai), typeof(Ao),
            typeof(method)}(Pi, Po, Ai, Ao, method, dim, reqs)
    end
end

"""
    MPI.Waitall!(t::Transposition)

Wait for completion of all unfinished MPI communications related to the
transposition.
"""
MPI.Waitall!(t::Transposition) =
    isempty(t.send_requests) || MPI.Waitall!(t.send_requests)

"""
    transpose!(t::Transposition; waitall=true)
    transpose!(dest::PencilArray{T,N}, src::PencilArray{T,N};
               method = Transpositions.PointToPoint())

Transpose data from one pencil configuration to the other.

The first variant allows to optionally delay the wait for MPI send operations to
complete.
This is useful if the caller wants to perform other operations with the already received data.
To do this, the caller should pass `waitall=false`, and manually invoke
[`MPI.Waitall!`](@ref) on the `Transposition` object once the operations are
done.
Note that this option only has an effect when the transposition method is
`PointToPoint`.

See [`Transposition`](@ref) for details.
"""
function transpose! end

function transpose!(
        dest::PencilArray, src::PencilArray;
        method::AbstractTransposeMethod = PointToPoint(),
    )
    dest === src && return dest  # same pencil & same data
    t = Transposition(dest, src, method=method)
    transpose!(t, waitall=true)
    dest
end

function transpose!(t::Transposition; waitall=true)
    timer = Pencils.timer(t.Pi)
    @timeit_debug timer "transpose!" begin
        transpose_impl!(t.dim, t)
        if waitall
            @timeit_debug timer "wait send" MPI.Waitall!(t)
        end
    end
    t
end

function assert_compatible(p::Pencil, q::Pencil)
    if p.topology !== q.topology
        throw(ArgumentError("pencil topologies must be the same."))
    end
    if p.size_global !== q.size_global
        throw(ArgumentError(
            "global data sizes must be the same between different pencil " *
            " configurations. Got $(p.size_global) ≠ $(q.size_global)."))
    end
    # Check that decomp_dims differ on at most one value.
    # Both are expected to be sorted.
    dp, dq = map(decomposition, (p, q))
    @assert all(map(issorted, (dp, dq)))
    if sum(dp .!= dq) > 1
        throw(ArgumentError(
            "pencil decompositions must differ in at most one dimension. " *
            "Got decomposed dimensions $dp and $dq."))
    end
    nothing
end

# Reinterpret UInt8 array as a different type of array.
# The input array should have enough space for the reinterpreted array with the
# given dimensions.
# This is a workaround to the performance issues when using `reinterpret`.
# See for instance:
# - https://discourse.julialang.org/t/big-overhead-with-the-new-lazy-reshape-reinterpret/7635
# - https://github.com/JuliaLang/julia/issues/28980
function unsafe_as_array(::Type{T}, x::Vector{UInt8}, dims) where T
    p = Ptr{T}(pointer(x))
    A = unsafe_wrap(Array, p, dims, own=false)
    @assert sizeof(A) <= sizeof(x)
    A
end

# Only local transposition.
function transpose_impl!(::Nothing, t::Transposition)
    Pi = t.Pi
    Po = t.Po
    Ai = t.Ai
    Ao = t.Ao
    timer = Pencils.timer(Pi)

    # Both pencil configurations are identical, so we just copy the data,
    # permuting dimensions if needed.
    @assert size(Ai) === size(Ao)
    ui = parent(Ai)
    uo = parent(Ao)

    if permutation(Pi) == permutation(Po)
        @timeit_debug timer "copy!" copy!(uo, ui)
    else
        @timeit_debug timer "permute_local!" permute_local!(Ao, Ai)
    end

    t
end

function permute_local!(Ao::PencilArray{T,N},
                        Ai::PencilArray{T,N}) where {T, N}
    Pi = pencil(Ai)
    Po = pencil(Ao)

    perm = let
        perm_base = permutation(Po) / permutation(Pi)  # relative permutation
        p = append(perm_base, Val(ndims_extra(Ai)))
        Tuple(p)
    end

    ui = parent(Ai)
    uo = parent(Ao)

    inplace = Base.mightalias(ui, uo)

    if inplace
        # TODO optimise in-place version?
        # For now we permute into a temporary buffer, and then we copy to `Ao`.
        # We reuse `recv_buf` used for MPI transposes.
        buf = let x = Pi.recv_buf
            n = length(uo)
            dims = size(uo)
            resize!(x, sizeof(T) * n)
            vec = unsafe_as_array(T, x, n)
            reshape(vec, dims)
        end
        permutedims!(buf, ui, perm)
        copy!(uo, buf)
    else
        # Permute directly onto the output.
        permutedims!(uo, ui, perm)
    end

    Ao
end

mpi_buffer(p::Ptr{T}, count) where {T} =
    MPI.Buffer(p, Cint(count), MPI.Datatype(T))

# Transposition among MPI processes in a subcommunicator.
# R: index of MPI subgroup (dimension of MPI Cartesian topology) along which the
# transposition is performed.
function transpose_impl!(R::Int, t::Transposition{T}) where {T}
    @assert t.dim === R
    Pi = t.Pi
    Po = t.Po
    Ai = t.Ai
    Ao = t.Ao
    method = t.method
    timer = Pencils.timer(Pi)

    @assert Pi.topology === Po.topology
    @assert extra_dims(Ai) === extra_dims(Ao)

    topology = Pi.topology
    comm = topology.subcomms[R]  # exchange among the subgroup R
    Nproc = topology.dims[R]
    subcomm_ranks = topology.subcomm_ranks[R]
    myrank = subcomm_ranks[topology.coords_local[R]]  # rank in subgroup

    remote_inds = get_remote_indices(R, topology.coords_local, Nproc)

    # Length of data that I will "send" to myself.
    length_self = let
        range_intersect = map(intersect, Pi.axes_local, Po.axes_local)
        prod(map(length, range_intersect)) * prod(extra_dims(Ai))
    end

    # Total data to be sent / received.
    length_send = length(Ai) - length_self
    length_recv_total = length(Ao)  # includes local exchange with myself

    resize!(Po.send_buf, sizeof(T) * length_send)
    send_buf = unsafe_as_array(T, Po.send_buf, length_send)

    resize!(Po.recv_buf, sizeof(T) * length_recv_total)
    recv_buf = unsafe_as_array(T, Po.recv_buf, length_recv_total)
    recv_offsets = Vector{Int}(undef, Nproc)  # all offsets in recv_buf

    req_length = method === Alltoallv() ? 0 : Nproc
    send_req = t.send_requests
    resize!(send_req, req_length)
    recv_req = similar(send_req)

    buffers = (send_buf, recv_buf)
    requests = (send_req, recv_req)

    # 1. Pack and send data.
    @timeit_debug timer "pack data" index_local_req = transpose_send!(
        buffers, recv_offsets, requests, length_self, remote_inds,
        (comm, subcomm_ranks, myrank),
        Ao, Ai, method, timer,
    )

    # 2. Unpack data and perform local transposition.
    @timeit_debug timer "unpack data" transpose_recv!(
        recv_buf, recv_offsets, recv_req,
        remote_inds, index_local_req,
        Ao, Ai, method, timer,
    )

    t
end

function transpose_send!(
        (send_buf, recv_buf),
        recv_offsets, requests,
        length_self, remote_inds,
        (comm, subcomm_ranks, myrank),
        Ao::PencilArray{T}, Ai::PencilArray{T},
        method::AbstractTransposeMethod,
        timer::TimerOutput,
    ) where {T}
    Pi = pencil(Ai)  # input (sent data)
    Po = pencil(Ao)  # output (received data)

    idims_local = Pi.axes_local
    odims_local = Po.axes_local

    idims = Pi.axes_all
    odims = Po.axes_all

    exdims = extra_dims(Ai)
    prod_extra_dims = prod(exdims)

    isend = 0  # current index in send_buf
    irecv = 0  # current index in recv_buf

    index_local_req = -1  # request index associated to local exchange

    # Data received from other processes.
    length_recv = length(Ao) - length_self

    Nproc = length(subcomm_ranks)
    @assert Nproc == MPI.Comm_size(comm)
    @assert myrank == MPI.Comm_rank(comm)

    buf_info = make_buffer_info(method, (send_buf, recv_buf), Nproc)

    for (n, ind) in enumerate(remote_inds)
        # Global data range that I need to send to process n.
        srange = map(intersect, idims_local, odims[ind])
        length_send_n = prod(map(length, srange)) * prod_extra_dims
        local_send_range = to_local(Pi, srange, MemoryOrder())

        # Determine amount of data to be received.
        rrange = map(intersect, odims_local, idims[ind])
        length_recv_n = prod(map(length, rrange)) * prod_extra_dims
        recv_offsets[n] = irecv

        rank = subcomm_ranks[n]  # actual rank of the other process

        if rank == myrank
            # Copy directly from `Ai` to `recv_buf`.
            # For convenience, data is put at the end of `recv_buf`.
            # This makes it easier to implement an alternative based on MPI_Alltoallv.
            @assert length_recv_n == length_self
            recv_offsets[n] = length_recv
            @timeit_debug timer "copy_range!" copy_range!(
                recv_buf, length_recv, Ai, local_send_range, exdims, timer)
            transpose_send_self!(method, n, requests, buf_info)
            index_local_req = n
        else
            # Copy data into contiguous buffer, then send the buffer.
            @timeit_debug timer "copy_range!" copy_range!(
                send_buf, isend, Ai, local_send_range, exdims, timer)
            transpose_send_other!(
                method, buf_info, (length_send_n, length_recv_n), n,
                requests, (rank, comm), eltype(Ai),
            )
            irecv += length_recv_n
            isend += length_send_n
        end
    end

    if method === Alltoallv()
        # This @view is needed because the Alltoallv wrapper checks that the
        # length of the buffer is consistent with recv_counts.
        recv_buf_view = @view recv_buf[1:length_recv]
        @timeit_debug timer "MPI.Alltoallv!" MPI.Alltoallv!(
            MPI.VBuffer(send_buf, buf_info.send_counts),
            MPI.VBuffer(recv_buf_view, buf_info.recv_counts),
            comm,
        )
    end

    index_local_req
end

function make_buffer_info(::PointToPoint, (send_buf, recv_buf), Nproc)
    (
        send_ptr = Ref(pointer(send_buf)),
        recv_ptr = Ref(pointer(recv_buf)),
    )
end

function make_buffer_info(::Alltoallv, bufs, Nproc)
    counts = Vector{Cint}(undef, Nproc)
    (
        send_counts = counts,
        recv_counts = similar(counts),
    )
end

function transpose_send_self!(::PointToPoint, n, (send_req, recv_req), etc...)
    send_req[n] = recv_req[n] = MPI.REQUEST_NULL
    nothing
end

function transpose_send_self!(::Alltoallv, n, reqs, buf_info)
    # Don't send data to myself via Alltoallv.
    buf_info.send_counts[n] = buf_info.recv_counts[n] = zero(Cint)
    nothing
end

function transpose_send_other!(
        ::PointToPoint, buf_info, (length_send_n, length_recv_n),
        n, (send_req, recv_req), (rank, comm), ::Type{T}
    ) where {T}
    # Exchange data with the other process (non-blocking operations).
    # Note: data is sent and received with the permutation associated to Pi.
    tag = 42
    send_req[n] = MPI.Isend(
        mpi_buffer(buf_info.send_ptr[], length_send_n),
        rank, tag, comm
    )
    recv_req[n] = MPI.Irecv!(
        mpi_buffer(buf_info.recv_ptr[], length_recv_n),
        rank, tag, comm
    )
    buf_info.send_ptr[] += length_send_n * sizeof(T)
    buf_info.recv_ptr[] += length_recv_n * sizeof(T)
    nothing
end

function transpose_send_other!(
        ::Alltoallv, buf_info, (length_send_n, length_recv_n), n, args...
    )
    buf_info.send_counts[n] = length_send_n
    buf_info.recv_counts[n] = length_recv_n
    nothing
end

function transpose_recv!(
        recv_buf, recv_offsets, recv_req,
        remote_inds, index_local_req,
        Ao::PencilArray, Ai::PencilArray,
        method::AbstractTransposeMethod,
        timer::TimerOutput,
    )
    Pi = pencil(Ai)  # input (sent data)
    Po = pencil(Ao)  # output (received data)

    odims_local = Po.axes_local
    idims = Pi.axes_all

    exdims = extra_dims(Ao)
    prod_extra_dims = prod(exdims)

    # Relative index permutation to go from Pi ordering to Po ordering.
    perm = permutation(Po) / permutation(Pi)

    Nproc = length(remote_inds)

    for m = 1:Nproc
        if method === Alltoallv()
            n = m
        elseif m == 1
            n = index_local_req  # copy local data first
        else
            @timeit_debug timer "wait receive" n, status =
                MPI.Waitany!(recv_req)
        end

        # Non-permuted global indices of received data.
        ind = remote_inds[n]
        g_range = map(intersect, odims_local, idims[ind])

        length_recv_n = prod(map(length, g_range)) * prod_extra_dims
        off = recv_offsets[n]

        # Local output data range in the **input** permutation.
        o_range_iperm = permutation(Pi) * to_local(Po, g_range, LogicalOrder())

        # Copy data to `Ao`, permuting dimensions if required.
        @timeit_debug timer "copy_permuted!" copy_permuted!(
            Ao, o_range_iperm, recv_buf, off, perm, exdims)
    end

    Ao
end

# Cartesian indices of the remote MPI processes included in the subgroup of
# index `R`.
# Example: if coords_local = (2, 3, 5) and R = 1, then this function returns the
# indices corresponding to (:, 3, 5).
function get_remote_indices(R::Int, coords_local::Dims{M}, Nproc::Int) where M
    t = ntuple(Val(M)) do i
        if i == R
            1:Nproc
        else
            c = coords_local[i]
            c:c
        end
    end
    CartesianIndices(t)
end

function copy_range!(dest::Vector{T}, dest_offset::Int, src::PencilArray{T,N},
                     src_range::ArrayRegion{P}, extra_dims::Dims{E}, timer,
                    ) where {T,N,P,E}
    @assert P + E == N

    n = dest_offset
    src_p = parent(src)  # array with non-permuted indices
    for K in CartesianIndices(extra_dims)
        for I in CartesianIndices(src_range)
            @inbounds dest[n += 1] = src_p[I, K]
        end
    end

    dest
end

function copy_permuted!(dest::PencilArray{T,N}, o_range_iperm::ArrayRegion{P},
                        src::Vector{T}, src_offset::Int,
                        perm::AbstractPermutation, extra_dims::Dims{E}) where {T,N,P,E}
    @assert P + E == N

    src_view = let src_dims = (map(length, o_range_iperm)..., extra_dims...)
        Ndata = prod(src_dims)
        n = src_offset
        v = view(src, (n + 1):(n + Ndata))
        reshape(v, src_dims)
    end

    dest_view = let dest_p = parent(dest)  # array with non-permuted indices
        indices = perm * o_range_iperm
        v = view(dest_p, indices..., map(Base.OneTo, extra_dims)...)
        if isidentity(perm)
            v
        else
            pperm = append(perm, Val(E))
            # Use fully inferred constructor defined in StaticPermutations
            PermutedDimsArray(v, inv(pperm))
        end
    end

    copyto!(dest_view, src_view)

    dest
end

end  # module Transpositions


text:
farmer = @stochastic_model begin
    @stage 1 begin
        @parameters begin
            Crops = [:wheat, :corn, :beets]
            Cost = Dict(:wheat=>150, :corn=>230, :beets=>260)
            Budget = 500
        end
        @decision(model, x[c in Crops] >= 0)
        @objective(model, Min, sum(Cost[c]*x[c] for c in Crops))
        @constraint(model, sum(x[c] for c in Crops) <= Budget)
    end
    @stage 2 begin
        @parameters begin
            Crops = [:wheat, :corn, :beets]
            Required = Dict(:wheat=>200, :corn=>240, :beets=>0)
            PurchasePrice = Dict(:wheat=>238, :corn=>210)
            SellPrice = Dict(:wheat=>170, :corn=>150, :beets=>36, :extra_beets=>10)
        end
        @uncertain ξ[c in Crops]
        @variable(model, y[p in setdiff(Crops, [:beets])] >= 0)
        @variable(model, w[s in Crops ∪ [:extra_beets]] >= 0)
        @objective(model, Min, sum(PurchasePrice[p] * y[p] for p in setdiff(Crops, [:beets]))
                   - sum(SellPrice[s] * w[s] for s in Crops ∪ [:extra_beets]))
        @constraint(model, minimum_requirement[p in setdiff(Crops, [:beets])],
            ξ[p] * x[p] + y[p] - w[p] >= Required[p])
        @constraint(model, minimum_requirement_beets,
            ξ[:beets] * x[:beets] - w[:beets] - w[:extra_beets] >= Required[:beets])
        @constraint(model, beets_quota, w[:beets] <= 6000)
    end
end

ξ₁ = Scenario(wheat = 3.0, corn = 3.6, beets = 24.0, probability = 1/3)
ξ₂ = Scenario(wheat = 2.5, corn = 3.0, beets = 20.0, probability = 1/3)
ξ₃ = Scenario(wheat = 2.0, corn = 2.4, beets = 16.0, probability = 1/3)

farmer_res = SPResult([170,80,250], -108390, -115405.56, 7015.56, 1150, -118600, -107240)
push!(problems, (farmer, [ξ₁,ξ₂,ξ₃], farmer_res,"Farmer"))


text:
@testset "104.maximum-depth-of-binary-tree.jl" begin
    @test maximum_depth_of_btree(TreeNode{Int}([3, 9, 20, nothing, nothing, 15, 7])) == 3
    @test maximum_depth_of_btree(TreeNode{Int}([1, nothing, 2])) == 2
    @test maximum_depth_of_btree(nothing) == 0
    @test maximum_depth_of_btree(TreeNode{Int}([5])) == 1
end


text:
function correct(string::AbstractString)::AbstractString
    function correct(c::Char)::Char
        c == '0' && return 'O'
        c == '1' && return 'I'
        c == '5' && return 'S'
                    return c
    end
    map(correct,string)
end

max_stars_repo_name: UlrichBerntien/Codewars-Katas

max_issues_repo_name: UlrichBerntien/Codewars-Katas

max_forks_repo_name: UlrichBerntien/Codewars-Katas


text:
include("fuzzlib.jl")

using Random: seed!

seed!(6174)
@testset "Fuzz test" begin
    @testset "polynormalize fuzz" begin
        for i=1:500
            fuzz_test(5, num_spec, SymbolicUtils.polynormalize; min_depth=3)
        end
    end
    @testset "num fuzz" begin
        for i=1:1500
            fuzz_test(5, num_spec)
        end
    end
    @testset "bool fuzz" begin
        for i=1:500
            seed!(i)
            fuzz_test(5, bool_spec)
        end
    end
end


text:
@testset "matmul" begin
    using NeuralAttentionlib.Matmul

    function matmul_test(x, y, s)
        cx = x isa CollapsedDimArray ? collapseddim(x) : x
        cy = y isa CollapsedDimArray ? collapseddim(y) : y
        return matmul(x, y, s) ≈ batched_mul(cx, cy) .* s
    end
    uwcs(x) = size(unwrap_collapse(x))

    @testset "gemm_strided Float64" begin
        a = randn(7, 6, 5, 4, 3, 2)
        b = randn(42, 20, 6)
        bt = randn(20, 42, 6)
        c = randn(42, 60, 2)
        ct = randn(60, 42, 2)
        s = rand() + 1

        ca1 = CollapsedDimArray(a, 3, 5)
        ca2 = CollapsedDimArray(a, 3, 6)

        @test matmul_test(ca1, bt, s)
        @test matmul_test(ca1, batched_transpose(b), s)
        @test matmul_test(bt, ca1, s)
        @test matmul_test(batched_transpose(b), ca1, s)
        @test matmul_test(batched_transpose(ca1), b, s)
        @test matmul_test(batched_transpose(ca1), batched_transpose(bt), s)
        @test matmul_test(b, batched_transpose(ca1), s)
        @test matmul_test(batched_transpose(bt), batched_transpose(ca1), s)
        @test matmul_test(ca1, batched_adjoint(b), s)
        @test matmul_test(batched_adjoint(b), ca1, s)
        @test matmul_test(batched_adjoint(ca1), b, s)
        @test matmul_test(batched_adjoint(ca1), batched_adjoint(bt), s)
        @test matmul_test(b, batched_adjoint(ca1), s)
        @test matmul_test(batched_adjoint(bt), batched_adjoint(ca1), s)

        @test uwcs(matmul(ca1, bt)) == (7, 6, 42, 6)
        @test uwcs(matmul(bt, ca1)) == (20, 5, 4, 3, 2)
        @test uwcs(matmul(batched_transpose(ca1), b)) == (5, 4, 20, 6)
        @test uwcs(matmul(batched_transpose(b), ca1)) == (20, 5, 4, 3, 2)

        @test matmul_test(ca2, ct, s)
        @test matmul_test(ca2, batched_transpose(c), s)
        @test matmul_test(ct, ca2, s)
        @test matmul_test(batched_transpose(c), ca2, s)
        @test matmul_test(batched_transpose(ca2), c, s)
        @test matmul_test(batched_transpose(ca2), batched_transpose(ct), s)
        @test matmul_test(c, batched_transpose(ca2), s)
        @test matmul_test(batched_transpose(ct), batched_transpose(ca2), s)

        @test uwcs(matmul(ca2, ct)) == (7, 6, 42, 2)
        @test uwcs(matmul(ct, ca2)) == (60, 5, 4, 3, 2)
        @test uwcs(matmul(batched_transpose(ca2), c)) == (5, 4, 3, 60, 2)
        @test uwcs(matmul(batched_transpose(c), ca2)) == (60, 5, 4, 3, 2)

    end

    @testset "gemm_strided Float32" begin
        a = randn(Float32, 7, 6, 5, 4, 3, 2)
        b = randn(Float32, 42, 20, 6)
        bt = randn(Float32, 20, 42, 6)
        c = randn(Float32, 42, 60, 2)
        ct = randn(Float32, 60, 42, 2)
        s = rand(Float32) + 1

        ca1 = CollapsedDimArray(a, 3, 5)
        ca2 = CollapsedDimArray(a, 3, 6)

        @test matmul_test(ca1, bt, s)
        @test matmul_test(ca1, batched_transpose(b), s)
        @test matmul_test(bt, ca1, s)
        @test matmul_test(batched_transpose(b), ca1, s)
        @test matmul_test(batched_transpose(ca1), b, s)
        @test matmul_test(batched_transpose(ca1), batched_transpose(bt), s)
        @test matmul_test(b, batched_transpose(ca1), s)
        @test matmul_test(batched_transpose(bt), batched_transpose(ca1), s)
        @test matmul_test(ca1, batched_adjoint(b), s)
        @test matmul_test(batched_adjoint(b), ca1, s)
        @test matmul_test(batched_adjoint(ca1), b, s)
        @test matmul_test(batched_adjoint(ca1), batched_adjoint(bt), s)
        @test matmul_test(b, batched_adjoint(ca1), s)
        @test matmul_test(batched_adjoint(bt), batched_adjoint(ca1), s)

        @test matmul_test(ca2, ct, s)
        @test matmul_test(ca2, batched_transpose(c), s)
        @test matmul_test(ct, ca2, s)
        @test matmul_test(batched_transpose(c), ca2, s)
        @test matmul_test(batched_transpose(ca2), c, s)
        @test matmul_test(batched_transpose(ca2), batched_transpose(ct), s)
        @test matmul_test(c, batched_transpose(ca2), s)
        @test matmul_test(batched_transpose(ct), batched_transpose(ca2), s)

    end

    @testset "gemm_strided ComplexF64" begin
        a = randn(ComplexF64, 7, 6, 5, 4, 3, 2)
        b = randn(ComplexF64, 42, 20, 6)
        bt = randn(ComplexF64, 20, 42, 6)
        c = randn(ComplexF64, 42, 60, 2)
        ct = randn(ComplexF64, 60, 42, 2)
        s = rand(ComplexF64) + 1

        ca1 = CollapsedDimArray(a, 3, 5)
        ca2 = CollapsedDimArray(a, 3, 6)

        @test matmul_test(ca1, bt, s)
        @test matmul_test(ca1, batched_transpose(b), s)
        @test matmul_test(bt, ca1, s)
        @test matmul_test(batched_transpose(b), ca1, s)
        @test matmul_test(batched_transpose(ca1), b, s)
        @test matmul_test(batched_transpose(ca1), batched_transpose(bt), s)
        @test matmul_test(b, batched_transpose(ca1), s)
        @test matmul_test(batched_transpose(bt), batched_transpose(ca1), s)
        @test matmul_test(ca1, batched_adjoint(b), s)
        @test matmul_test(batched_adjoint(b), ca1, s)
        @test matmul_test(batched_adjoint(ca1), b, s)
        @test matmul_test(batched_adjoint(ca1), batched_adjoint(bt), s)
        @test matmul_test(b, batched_adjoint(ca1), s)
        @test matmul_test(batched_adjoint(bt), batched_adjoint(ca1), s)

        @test matmul_test(ca2, ct, s)
        @test matmul_test(ca2, batched_transpose(c), s)
        @test matmul_test(ct, ca2, s)
        @test matmul_test(batched_transpose(c), ca2, s)
        @test matmul_test(batched_transpose(ca2), c, s)
        @test matmul_test(batched_transpose(ca2), batched_transpose(ct), s)
        @test matmul_test(c, batched_transpose(ca2), s)
        @test matmul_test(batched_transpose(ct), batched_transpose(ca2), s)

    end

    @testset "gemm_strided ComplexF32" begin
        a = randn(ComplexF32, 7, 6, 5, 4, 3, 2)
        b = randn(ComplexF32, 42, 20, 6)
        bt = randn(ComplexF32, 20, 42, 6)
        c = randn(ComplexF32, 42, 60, 2)
        ct = randn(ComplexF32, 60, 42, 2)
        s = rand(ComplexF32) + 1

        ca1 = CollapsedDimArray(a, 3, 5)
        ca2 = CollapsedDimArray(a, 3, 6)

        @test matmul_test(ca1, bt, s)
        @test matmul_test(ca1, batched_transpose(b), s)
        @test matmul_test(bt, ca1, s)
        @test matmul_test(batched_transpose(b), ca1, s)
        @test matmul_test(batched_transpose(ca1), b, s)
        @test matmul_test(batched_transpose(ca1), batched_transpose(bt), s)
        @test matmul_test(b, batched_transpose(ca1), s)
        @test matmul_test(batched_transpose(bt), batched_transpose(ca1), s)
        @test matmul_test(ca1, batched_adjoint(b), s)
        @test matmul_test(batched_adjoint(b), ca1, s)
        @test matmul_test(batched_adjoint(ca1), b, s)
        @test matmul_test(batched_adjoint(ca1), batched_adjoint(bt), s)
        @test matmul_test(b, batched_adjoint(ca1), s)
        @test matmul_test(batched_adjoint(bt), batched_adjoint(ca1), s)

        @test matmul_test(ca2, ct, s)
        @test matmul_test(ca2, batched_transpose(c), s)
        @test matmul_test(ct, ca2, s)
        @test matmul_test(batched_transpose(c), ca2, s)
        @test matmul_test(batched_transpose(ca2), c, s)
        @test matmul_test(batched_transpose(ca2), batched_transpose(ct), s)
        @test matmul_test(c, batched_transpose(ca2), s)
        @test matmul_test(batched_transpose(ct), batched_transpose(ca2), s)

    end

    @testset "AD" begin
        test_rrule(matmul, randn(7,6,5), randn(6, 2), randn())
        test_rrule(matmul, randn(7,6,5,4), randn(6), randn())
        test_rrule(matmul, CollapsedDimArray(randn(2,2,2,2,2,3), 4, 6), batched_transpose(randn(5,4,3)), randn(); check_inferred=false)
    end

end


text:
__precompile__()

module Sim

import GLFW
using MuJoCo
using StaticArrays

##################################################### globals
const fontscale = mj.FONTSCALE_150 # can be 100, 150, 200

mutable struct Simulation
   # visual interaction controls
   lastx::Float64
   lasty::Float64
   button_left::Bool
   button_middle::Bool
   button_right::Bool

   lastbutton::Int32
   lastclicktm::Float64

   refreshrate::Int

   showhelp::Int
   showoption::Bool
   showdepth::Bool
   showfullscreen::Bool
   showsensor::Bool
   slowmotion::Bool

   showinfo::Bool
   paused::Bool
   keyreset::Int

   framenum::Int
   lastframenum::Int

   # MuJoCo things
   scn::Ptr{mj.mjvScene}
   cam::Ptr{mj.mjvCamera}
   vopt::Ptr{mj.mjvOption}
   pert::Ptr{mj.mjvPerturb}
   con::Ptr{mj.mjrContext}
   figsensor::Ptr{mj.mjvFigure}
   m::mj.jlModel
   d::mj.jlData

   # GLFW handle
   window::GLFW.Window

   function Simulation(m::mj.jlModel, d::mj.jlData, window::GLFW.Window)
      new(0.0, 0.0, false, false, false, 0, 0.0,
          0, 0, false, false, false, false, false, true, true, 0, 0, 0,
          Ptr{mj.mjvScene}(Base.Libc.malloc(sizeof(mj.mjvScene))),
          Ptr{mj.mjvCamera}(Base.Libc.malloc(sizeof(mj.mjvCamera))),
          Ptr{mj.mjvOption}(Base.Libc.malloc(sizeof(mj.mjvOption))),
          Ptr{mj.mjvPerturb}(Base.Libc.malloc(sizeof(mj.mjvPerturb))),
          Ptr{mj.mjrContext}(Base.Libc.malloc(sizeof(mj.mjrContext))),
          Ptr{mj.mjvFigure}(Base.Libc.malloc(sizeof(mj.mjvFigure))),
          m, d, window)
   end
end

export Simulation

const keycmds = Dict{GLFW.Key, Function}(GLFW.KEY_F1=>(s)->begin  # help
                                         s.showhelp += 1
                                         if s.showhelp > 2 s.showhelp = 0 end
                                      end,
                                      GLFW.KEY_F2=>(s)->begin  # option
                                         s.showoption = !s.showoption;
                                      end,
                                      GLFW.KEY_F3=>(s)->begin  # info
                                         s.showinfo = !s.showinfo;
                                      end,
                                      GLFW.KEY_F4=>(s)->begin  # depth
                                         s.showdepth = !s.showdepth;
                                      end,
                                      GLFW.KEY_F5=>(s)->begin  # toggle full screen
                                         s.showfullscreen = !s.showfullscreen;
                                         s.showfullscreen ? GLFW.MaximizeWindow(s.window):GLFW.RestoreWindow(s.window)
                                      end,
                                      #GLFW.KEY_F6=>(s)->begin  # stereo
                                      #   s.stereo = mj.get(s.scn, :stereo) == mj.mjSTEREO_NONE ? mjSTEREO_QUADBUFFERED : mj.mjSTEREO_NONE
                                      #   mj.set(s.scn, :stereo)
                                      #end,
                                      GLFW.KEY_F7=>(s)->begin  # sensor figure
                                         s.showsensor = !s.showsensor;
                                      end,
                                      GLFW.KEY_F8=>(s)->begin  # profiler
                                         s.showprofiler = !s.showprofiler;
                                      end,
                                      GLFW.KEY_ENTER=>(s)->begin  # slow motion
                                         s.slowmotion = !s.slowmotion;
                                         s.slowmotion ? println("Slow Motion Mode!"):println("Normal Speed Mode!")
                                      end,
                                      GLFW.KEY_SPACE=>(s)->begin  # pause
                                         s.paused = !s.paused
                                         s.paused ? println("Paused"):println("Running")
                                      end,
                                      GLFW.KEY_PAGE_UP=>(s)->begin    # previous keyreset
                                         s.keyreset = min(mj.get(s.m.m, :nkey) - 1, s.keyreset + 1)
                                      end,
                                      GLFW.KEY_PAGE_DOWN=>(s)->begin  # next keyreset
                                         s.keyreset = max(-1, s.keyreset - 1)
                                      end,
                                      # continue with reset
                                      GLFW.KEY_BACKSPACE=>(s)->begin  # reset
                                         mj.resetData(s.m.m, s.d.d)
                                         if s.keyreset >= 0 && s.keyreset < mj.get(s.m.m, :nkey)
                                            mj.set(s.d, :time, s.m.key_time[s.keyreset+1])
                                            s.d.qpos[:] = s.m.key_qpos[:,s.keyreset+1]
                                            s.d.qvel[:] = s.m.key_qvel[:,s.keyreset+1]
                                            s.d.act[:]  = s.m.key_act[:,s.keyreset+1]
                                         end
                                         mj.forward(s.m, s.d)
                                         #profilerupdate()
                                         sensorupdate(s)
                                      end,
                                      GLFW.KEY_RIGHT=>(s)->begin  # step forward
                                         if s.paused
                                            mj.step(s.m, s.d)
                                            #profilerupdate()
                                            sensorupdate(s)
                                         end
                                      end,
                                      GLFW.KEY_LEFT=>(s)->begin  # step back
                                         if s.paused
                                            dt = mj.get(s.m, :opt, :timestep)
                                            mj.set(s.m, :opt, :timestep, -dt)
                                            #cleartimers(s.d);
                                            mj.step(s.m, s.d);
                                            mj.set(s.m, :opt, :timestep, dt)
                                            #profilerupdate()
                                            sensorupdate(s)
                                         end
                                      end,
                                      GLFW.KEY_DOWN=>(s)->begin  # step forward 100
                                         if s.paused
                                            #cleartimers(d);
                                            for n=1:100 mj.step(s.m, s.d) end
                                            #profilerupdate();
                                            sensorupdate(s)
                                         end
                                      end,
                                      GLFW.KEY_UP=>(s)->begin  # step back 100
                                         if s.paused
                                            dt = mj.get(s.m, :opt, :timestep)
                                            mj.set(s.m, :opt, :timestep, -dt)
                                            #cleartimers(d)
                                            for n=1:100 mj.step(s.m, s.d) end
                                            mj.set(s.m, :opt, :timestep, dt)
                                            #profilerupdate();
                                            sensorupdate(s)
                                         end
                                      end,
                                      GLFW.KEY_ESCAPE=>(s)->begin  # free camera
                                         mj.set(s.cam, :type, Int(mj.CAMERA_FREE))
                                      end,
                                      GLFW.KEY_EQUAL=>(s)->begin  # bigger font
                                         if fontscale < 200
                                            fontscale += 50
                                            mj.mjr_makeContext(s.m.m, s.con, fontscale)
                                         end
                                      end,
                                      GLFW.KEY_MINUS=>(s)->begin  # smaller font
                                         if fontscale > 100
                                            fontscale -= 50;
                                            mj.mjr_makeContext(s.m.m, s.con, fontscale);
                                         end
                                      end,
                                      GLFW.KEY_LEFT_BRACKET=>(s)->begin  # '[' previous fixed camera or free
                                           fixedcam = mj.get(s.cam, :type)
                                           if mj.get(s.m.m, :ncam) > 0 && fixedcam == Int(mj.CAMERA_FIXED)
                                              fixedcamid = mj.get(s.cam, :fixedcamid)
                                              if (fixedcamid  > 0)
                                                 mj.set(s.cam, :fixedcamid, fixedcamid-1)
                                              else
                                                 mj.set(s.cam, :type, Int(mj.CAMERA_FREE))
                                              end
                                           end
                                        end,
                                        GLFW.KEY_RIGHT_BRACKET=>(s)->begin  # ']' next fixed camera
                                           if mj.get(s.m.m, :ncam) > 0
                                              fixedcam = mj.get(s.cam, :type)
                                              fixedcamid = mj.get(s.cam, :fixedcamid)
                                              if fixedcam != Int(mj.CAMERA_FIXED)
                                                 mj.set(s.cam, :type, Int(mj.CAMERA_FIXED))
                                              elseif fixedcamid < mj.get(s.m.m, :ncam) - 1
                                                 mj.set(s.cam, :fixedcamid, fixedcamid+1)
                                              end
                                           end
                                        end,
                                        GLFW.KEY_SEMICOLON=>(s)->begin  # cycle over frame rendering modes
                                           frame = mj.get(s.vopt, :frame)
                                           mj.set(s.vopt, :frame, max(0, frame - 1))
                                        end,
                                        GLFW.KEY_APOSTROPHE=>(s)->begin  # cycle over frame rendering modes
                                           frame = mj.get(s.vopt, :frame)
                                           mj.set(s.vopt, :frame,
                                                  min(Int(mj.NFRAME)-1, frame+1))
                                        end,
                                        GLFW.KEY_PERIOD=>(s)->begin  # cycle over label rendering modes
                                           label = mj.get(s.vopt, :label)
                                           mj.set(s.vopt, :label, max(0, label-1))
                                        end,
                                        GLFW.KEY_SLASH=>(s)->begin  # cycle over label rendering modes
                                           label = mj.get(s.vopt, :label)
                                           mj.set(s.vopt, :label,
                                                  min(Int(mj.NLABEL)-1, label+1))
                                        end)

##################################################### functions
function autoscale(s::Simulation)
   # autoscale
   center = mj.get(s.m, :stat, :center)
   mj.set(s.cam, :lookat, center[1], 1)
   mj.set(s.cam, :lookat, center[2], 2)
   mj.set(s.cam, :lookat, center[3], 3)
   mj.set(s.cam, :distance, 1.5*mj.get(s.m, :stat, :extent))

   # set to free camera
   mj.set(s.cam, :_type, Int(mj.CAMERA_FREE))
end

# init sensor figure
function sensorinit(s::Simulation)
   # set figure to default
   mj.mjv_defaultFigure(s.figsensor)

   # set flags
   mj.set(s.figsensor, :flg_extend, Cint(1))
   mj.set(s.figsensor, :flg_barplot, Cint(1))

   # title
   title = Vector{UInt8}("Sensor data")
   for i=1:length(title)
      mj.set(s.figsensor, :title, title[i], i)
   end

   # y-tick nubmer format
   format = "%.0f"
   for i=1:length(format)
      mj.set(s.figsensor, :yformat, format[i], i)
   end

   # grid size
   mj.set(s.figsensor, :gridsize, 2, 1)
   mj.set(s.figsensor, :gridsize, 3, 2)

   # minimum range
   mj.set(s.figsensor, :range,  0, 1, 1)
   mj.set(s.figsensor, :range,  0, 1, 2)
   mj.set(s.figsensor, :range, -1, 2, 1)
   mj.set(s.figsensor, :range,  1, 2, 2)
end

# update sensor figure
function sensorupdate(s::Simulation)
   const maxline = 10

   for i=1:maxline # clear linepnt
      mj.set(s.figsensor, :linepnt, Cint(0), i)
   end

   lineid = 1 # start with line 0
   m = s.m
   d = s.d

   # loop over sensors
   for n=1:mj.get(m, :nsensor)
      # go to next line if type is different
      if (n > 1 && m.sensor_type[n] != m.sensor_type[n - 1])
         lineid = min(lineid+1, maxline)
      end

      # get info about this sensor
      cutoff = m.sensor_cutoff[n] > 0 ? m.sensor_cutoff[n] : 1.0
      adr = m.sensor_adr[n]
      dim = m.sensor_dim[n]

      # data pointer in line
      p = mj.get(s.figsensor, :linepnt, lineid)

      # fill in data for this sensor
      for i=0:(dim-1)
         # check size
         if ((p + 2i) >= Int(mj.MAXLINEPNT) / 2) break end

         x1 = 2p + 4i + 1
         x2 = 2p + 4i + 3
         mj.set(s.figsensor, :linedata, adr+i, lineid, x1)
         mj.set(s.figsensor, :linedata, adr+i, lineid, x2)

         y1 = 2p + 4i + 2
         y2 = 2p + 4i + 4
         se = d.sensordata[adr+i+1]/cutoff
         mj.set(s.figsensor, :linedata,  0, lineid, y1)
         mj.set(s.figsensor, :linedata, se, lineid, y2)
      end

      # update linepnt
      mj.set(s.figsensor, :linepnt,
             min(Int(mj.MAXLINEPNT)-1, p+2dim),
             lineid)
   end
end

# show sensor figure
function sensorshow(s::Simulation, rect::mj.mjrRect)
   # render figure on the right
   viewport = mj.mjrRect(rect.width - rect.width / 4,
                         rect.bottom,
                         rect.width / 4,
                         rect.height / 3)
   mj.mjr_figure(viewport, s.figsensor, s.con)
end

##################################################### callbacks
function mykeyboard(s::Simulation, window::GLFW.Window,
                    key::GLFW.Key, scancode::Int32, act::GLFW.Action, mods::Int32)
   # do not act on release
   if act == GLFW.RELEASE return end

   try
      keycmds[key](s) # call anon function in Dict with s struct passed in
   catch
      # control keys
      if mods & GLFW.MOD_CONTROL > 0
         if key == GLFW.KEY_A
            autoscale(s)
         elseif key == GLFW.KEY_P
            println(s.d.qpos)
         #elseif key == GLFW.KEY_L && lastfile[0]
         #   loadmodel(window, s.)
         elseif key == GLFW.KEY_Q
            GLFW.SetWindowShouldClose(window, true)
         end
      end

      # toggle visualiztion flag
      flags = mj.get(s.vopt, :flags)
      for i=1:Int(mj.NVISFLAG)
         if (key == Int(mj.VISSTRING[i,3][1]))
            mj.set(s.vopt, :flags, flags[i]==0x00?0x01:0x00, i)
         end
      end
      # toggle rendering flag
      flags = mj.get(s.scn, :flags)
      for i=1:Int(mj.NRNDFLAG)
         if (key == Int(mj.RNDSTRING[i,3][1]))
            mj.set(s.scn, :flags, flags[i]==0x00?0x01:0x00, i)
         end
      end
      # toggle geom/site group
      for i=1:Int(mj.NGROUP)
         if key == i + Int('0')
            sitegroup = mj.get(s.vopt, :sitegroup)
            geomgroup = mj.get(s.vopt, :geomgroup)
            if mods & GLFW.MOD_SHIFT == true
               mj.set(s.vopt, :sitegroup, sitegroup[i]>0?0:1, i)
            else
               mj.set(s.vopt, :geomgroup, geomgroup[i]>0?0:1, i)
            end
         end
      end
   end
end

function mouse_move(s::Simulation, window::GLFW.Window,
                    xpos::Float64, ypos::Float64)
   # no buttons down: nothing to do
   if !s.button_left && !s.button_middle && !s.button_right
      return
   end

   # compute mouse displacement, save
   dx = xpos - s.lastx
   dy = ypos - s.lasty
   s.lastx = xpos
   s.lasty = ypos

   width, height = GLFW.GetWindowSize(window)

   mod_shift = GLFW.GetKey(window, GLFW.KEY_LEFT_SHIFT) || GLFW.GetKey(window, GLFW.KEY_RIGHT_SHIFT)

   # determine action based on mouse button
   if s.button_right
      action = mod_shift ? Int(mj.MOUSE_MOVE_H) : Int(mj.MOUSE_MOVE_V)
   elseif s.button_left
      action = mod_shift ? Int(mj.MOUSE_ROTATE_H) : Int(mj.MOUSE_ROTATE_V)
   else
      action = Int(mj.MOUSE_ZOOM)
   end

   # move perturb or camera
   pert = unsafe_load(s.pert)
   if pert.active != 0
      mj.mjv_movePerturb(s.m.m, s.d.d, action,
                         dx / height, dy / height,
                         s.scn, s.pert);
   else
      mj.mjv_moveCamera(s.m.m, action,
                        dx / height, dy / height,
                        s.scn, s.cam)
   end
end

# past data for double-click detection
function mouse_button(s::Simulation, window::GLFW.Window,
                      button::Int32, act::GLFW.Action, mods::Int32)
   # update button state
   s.button_left = GLFW.GetMouseButton(window, GLFW.MOUSE_BUTTON_LEFT)
   s.button_middle = GLFW.GetMouseButton(window, GLFW.MOUSE_BUTTON_MIDDLE)
   s.button_right = GLFW.GetMouseButton(window, GLFW.MOUSE_BUTTON_RIGHT)

   # Alt: swap left and right
   if mods == GLFW.MOD_ALT
      tmp = s.button_left
      s.button_left = s.button_right
      s.button_right = tmp

      if button == GLFW.MOUSE_BUTTON_LEFT
         button = GLFW.MOUSE_BUTTON_RIGHT;
      elseif button == GLFW.MOUSE_BUTTON_RIGHT
         button = GLFW.MOUSE_BUTTON_LEFT;
      end
   end

   # update mouse position
   x, y = GLFW.GetCursorPos(window)
   s.lastx = x
   s.lasty = y

   # set perturbation
   newperturb = 0;
   pert = unsafe_load(s.pert)
   cam = unsafe_load(s.cam)
   if act == GLFW.PRESS && mods == GLFW.MOD_CONTROL && pert.select > 0
      # right: translate;  left: rotate
      if s.button_right
         newperturb = Int(mj.PERT_TRANSLATE)
      elseif s.button_left
         newperturb = Int(mj.PERT_ROTATE)
      end
      # perturbation onset: reset reference
      if newperturb>0 && pert.active==0
         mj.mjv_initPerturb(s.m.m, s.d.d, s.scn, s.pert)
      end
   end
   mj.set(s.pert, :active, newperturb)

   # detect double-click (250 msec)
   if act == GLFW.PRESS && (time() - s.lastclicktm < 0.25) && (button == s.lastbutton)
      # determine selection mode
      if button == GLFW.MOUSE_BUTTON_LEFT
         selmode = 1;
      elseif mods == GLFW.MOD_CONTROL
         selmode = 3;
      else
         selmode = 2;
      end
      # get current window size
      width, height = GLFW.GetWindowSize(window)

      # find geom and 3D click point, get corresponding body
      selpnt = zeros(3)
      selgeom = mj.mjv_select(s.m.m, s.d.d, s.vopt,
                              width / height, x / width,
                              (height - y) / height,
                              s.scn, selpnt)
      selbody = (selgeom >= 0 ? s.m.geom_bodyid[selgeom+1] : 0) # 0 indexed

      # set lookat point, start tracking is requested
      if selmode == 2 || selmode == 3
         # copy selpnt if geom clicked
         # TODO hacks
         if selgeom >= 0
            mj.mju_copy3(cam.lookat, @SVector [selpnt[1], selpnt[2], selpnt[3]])
         end

         # switch to tracking camera
         if selmode == 3 && selbody >= 0
            #cam.trackbodyid = selbody
            #cam.fixedcamid = -1
            mj.set(s.cam, :_type, Int(mj.CAMERA_TRACKING))
            mj.set(s.cam, :trackbodyid, selbody)
            mj.set(s.cam, :fixedcamid, -1)
         end
      else # set body selection
         if selbody >= 0
            # compute localpos
            tmp = SVector{3, Float64}(0.0, 0.0, 0.0)
            #mju_sub3(tmp, selpnt, d->xpos + 3 * pert.select);
            tmp = selpnt - s.d.xpos[:,selbody+1]
            #mju_mulMatTVec(pert.localpos, d->xmat + 9 * pert.select, tmp, 3, 3);
            res = reshape(s.d.xmat[:,selbody+1], 3, 3)' * tmp
            mj.set(s.pert, :localpos, SVector{3}(res))

            # record selection
            mj.set(s.pert, :select, selbody)
         else
            mj.set(s.pert, :select, 0)
         end
      end

      # stop perturbation on select
      mj.set(s.pert, :active, 0)
   end
   # save info
   if act == GLFW.PRESS
      s.lastbutton = button
      s.lastclicktm = time()
   end
end

function scroll(s::Simulation, window::GLFW.Window,
                xoffset::Float64, yoffset::Float64)
   # scroll: emulate vertical mouse motion = 5% of window height
   mj.mjv_moveCamera(s.m.m, Int(mj.MOUSE_ZOOM),
                     0.0, -0.05 * yoffset, s.scn, s.cam);

end

function drop(window::GLFW.Window,
              count::Int, paths::String)
end

function start(mm::mj.jlModel, dd::mj.jlData,
               width=1200, height=900) # TODO named args for callbacks
   # TODO make simulation
   s = Simulation(mm, dd, GLFW.CreateWindow(width, height, "Simulate"))

   # Make the window's context current
   GLFW.MakeContextCurrent(s.window)
   GLFW.SwapInterval(1)

   s.refreshrate = GLFW.GetVideoMode(GLFW.GetPrimaryMonitor()).refreshrate
   println("Refresh Rate: $(s.refreshrate)")

   # mujoco setup
   mj.mjv_makeScene(s.scn, 1000)
   mj.mjv_defaultPerturb(s.pert)
   mj.mjv_defaultCamera(s.cam)
   mj.mjv_defaultOption(s.vopt)
   mj.mjr_defaultContext(s.con)
   mj.mjr_makeContext(s.m.m, s.con, Int(fontscale)) # model specific setup

   #profilerinit();
   sensorinit(s)
   #GLFW.SetKeyCallback(s.window, (w,k,sc,a,m)->mykeyboard(s,w,k,sc,a,m))

   GLFW.SetCursorPosCallback(s.window, (w,x,y)->mouse_move(s,w,x,y))
   GLFW.SetMouseButtonCallback(s.window, (w,b,a,m)->mouse_button(s,w,b,a,m))
   GLFW.SetScrollCallback(s.window, (w,x,y)->scroll(s,w,x,y))
   ##GLFW.SetDropCallback(s.window, drop)

   return s
end

#function __init__()
#   atexit(cleanup)
#end
#function cleanup()
#   GLFW.DestroyWindow(window)
#end

end


text:
A = zeros(UInt8, 10, 10, 10)
for i in 1:8
    A[i, 5, i:i+2] = 0x01
end
A[9:10, 5, 8:10] = 0x01
objects = Vector{Object}()

blob_series = form_blobs(A)
for blobs in blob_series
    match!(objects, blobs, 1; min_area=2)
    @test length(objects) == 1
end


text:
"""
    assortativity(g)

Return the [assortativity coefficient](https://en.wikipedia.org/wiki/Assortativity)
of graph `g`, defined as the Pearson correlation of excess degree between
the end vertices of all the edges of the graph.

The excess degree is equal to the degree of linked vertices minus one,
i.e. discounting the edge that links the pair.
In directed graphs, the paired values are the out-degree of source vertices
and the in-degree of destination vertices.

# Examples
```jldoctest
julia> using LightGraphs

julia> assortativity(star_graph(4))
-1.0
```
"""
function assortativity(g::AbstractGraph{T}) where T
    P = promote_type(Int64, T) # at least Int64 to reduce risk of overflow
    nue  = ne(g)
    sjk = sj = sk = sjs = sks = zero(P)
    for d in edges(g)
        j = P(outdegree(g, src(d)) - 1)
        k = P(indegree(g, dst(d)) - 1)
        sjk += j*k
        sj  += j
        sk  += k
        sjs += j^2
        sks += k^2
    end
    return assortativity_coefficient(g, sjk, sj, sk, sjs, sks, nue)
end

#=
assortativity coefficient for directed graphs:
see equation (21) in M. E. J. Newman: Mixing patterns in networks, Phys. Rev. E 67, 026126 (2003),
http://arxiv.org/abs/cond-mat/0209450
=#
@traitfn function assortativity_coefficient(g::::IsDirected, sjk, sj, sk, sjs, sks, nue)
    return (sjk - sj*sk/nue) / sqrt((sjs - sj^2/nue)*(sks - sk^2/nue))
end

#=
assortativity coefficient for undirected graphs:
see equation (4) in M. E. J. Newman: Assortative mixing in networks, Phys. Rev. Lett. 89, 208701 (2002),
http://arxiv.org/abs/cond-mat/0205405/
=#
@traitfn function assortativity_coefficient(g::::(!IsDirected), sjk, sj, sk, sjs, sks, nue)
    return (sjk/nue - ((sj + sk)/(2*nue))^2) / ((sjs + sks)/(2*nue) - ((sj + sk)/(2*nue))^2)
end


text:
#=
test_module:
- Julia version: 1.5.0
- Author: shisa
- Date: 2020-08-06
=#

module TestModule
    # packages
    using Test

    # extarnel modules
    include("simple_module.jl")
    include("nest_module.jl")

    # methods
    function test()
        @testset "Module" begin
            @testset "SimpleModule" begin
                @test_nowarn SimpleModule.main()
            end

            @testset "NestModule" begin
                @test_nowarn NestModule.main()
            end
        end
    end
end

if abspath(PROGRAM_FILE) == @__FILE__
    using. TestModule
    TestModule.test()
end

max_stars_repo_name: ShisatoYano/JuliaPractice

max_issues_repo_name: ShisatoYano/JuliaPractice

max_forks_repo_name: ShisatoYano/JuliaPractice


text:
function extract_nodes_branches(
  filename  # name of the iSDF file which contains the information of the circuit to be simulated.
  )  #=

reads the circuit to be simulated from a file whose file name is designated by the argument "filename" and then extracts the information of the nodes and branches in the circuit. The following variables are returned as the elements of the tuple "nbset".
- nodes: array of Strings containing independent node names.
- num_nodes: number of nodes.
- branches: array of struct Branch.
- num_branches: sum of branch phases.
- num_dyns: number of dynamic branches.

Author: Taku Noda
Started on Jan. 6, 2021
Update on June 29, 2021

Copyright 2021 Central Research Institute of Electric Power Industry

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

=#

sys = isdf_pp1( filename )
subsyss = elements( sys.root )
subsys = subsyss[ 1 ]

nodes = collect_nodes( subsys )
num_nodes = length( nodes )
branches, num_branches, num_dyns = collect_branches( subsys, nodes )
nbset = ( nodes, num_nodes, branches, num_branches, num_dyns )

return nbset

end

max_stars_repo_name: takunoda/omegasim

max_issues_repo_name: takunoda/omegasim

max_forks_repo_name: takunoda/omegasim


text:
# SPD-License-Identifier: MIT

using Knet
using CUDA

include("layers.jl")

#======================================================================
Model definition.
======================================================================#

function buildmodel(config, indim, outdim)
    dropout = config["dropout"]
    dilations = config["dilations"]
    hdims = config["hidden-dims"]
    kernelsizes = config["kernel-sizes"]
    strides = config["strides"]

    @assert (length(kernelsizes) == length(dilations)
             && length(dilations) == length(strides)
             && length(dilations) == length(hdims))

    layers = Any[]
    push!(layers, PermuteDims([2, 1, 3]))
    push!(layers, AddExtraDim()) # Necessary as convolution is 2D.
    for i in 1:length(hdims)
        layer = Conv(
            (kernelsizes[i], 1),
            indim => hdims[i],
            stride = (strides[i], 1),
            dilation = (dilations[i], 1),
            padding = ((kernelsizes[i] - 1 ) ÷ 2, 0)
           )
        push!(layers, layer)
        push!(layers, BatchNorm(hdims[i], relu))
        push!(layers, Dropout(dropout))
        indim = hdims[i]
    end

    push!(layers, RemoveExtraDim())
    push!(layers, PermuteDims([2, 1, 3]))

    # Final layer is just an affine transform.
    push!(layers, Dense(hdims[end], outdim))

    network = Chain(layers...)
end

function getlengths(config, inlengths)
    dilations = config["dilations"]
    hdims = config["hidden-dims"]
    kernelsizes = config["kernel-sizes"]
    strides = config["strides"]

    seqlengths = inlengths
    for i in 1:length(dilations)
        ksize = kernelsizes[i]
        pad = (ksize - 1 ) ÷ 2
        newlengths = seqlengths .+ 2*pad .- dilations[i]*(ksize-1) .+ strides[i] .- 1
        seqlengths = newlengths .÷ strides[i]
    end
    seqlengths
end

#======================================================================
Training utilities.
======================================================================#

mutable struct PlateauScheduler
    factor
    patience
    threshold
    best_loss
    nsteps
end

function getscheduler(trainconfig)
    sconf = trainconfig["scheduler"]
    PlateauScheduler(sconf["factor"], sconf["patience"], 1e-4, -Inf, 0)
end

function update_scheduler!(s::PlateauScheduler, opts, loss)
    if loss < s.best_loss * (1 - s.threshold)
        s.nsteps = 0
        s.best_loss = loss
    elseif s.nsteps > s.patience
        @debug "$(s.nsteps) epoch(s) without improvement, setting learning rate to: $(opts[1].lr)"
        for opt in opts opt.lr = opt.lr * s.factor end
        s.nsteps = 0
    end
    s.nsteps += 1
end

function getoptimizer(trainconfig)
    Adam(lr = trainconfig["learning_rate"], beta1 = 0.9, beta2 = 0.999)
end

max_stars_repo_name: MartinKocour/SpeechLab

max_issues_repo_name: MartinKocour/SpeechLab

max_forks_repo_name: MartinKocour/SpeechLab