text:
# Attempt to compute a reasonable default mu: at the starting
# position, the gradient of the input function should dominate the
# gradient of the barrier.
function initial_mu{T}(gfunc::Array{T}, gbarrier::Array{T}; mu0::T = convert(T, NaN), mu0factor::T = 0.001)
    if isnan(mu0)
        gbarriernorm = sum(abs, gbarrier)
        if gbarriernorm > 0
            mu = mu0factor*sum(abs, gfunc)/gbarriernorm
        else
            # Presumably, there is no barrier function
            mu = zero(T)
        end
    else
        mu = mu0
    end
    return mu
end

function barrier_box{T}(x::Array{T}, g, l::Array{T}, u::Array{T})
    n = length(x)
    calc_g = !(g === nothing)

    v = zero(T)
    for i = 1:n
        thisl = l[i]
        if isfinite(thisl)
            dx = x[i] - thisl
            if dx <= 0
                return convert(T, Inf)
            end
            v -= log(dx)
            if calc_g
                g[i] = -one(T)/dx
            end
        else
            if calc_g
                g[i] = zero(T)
            end
        end
        thisu = u[i]
        if isfinite(thisu)
            dx = thisu - x[i]
            if dx <= 0
                return convert(T, Inf)
            end
            v -= log(dx)
            if calc_g
                g[i] += one(T)/dx
            end
        end
    end
    return v
end

function function_barrier{T, F<:Function, FB<:Function}(x::Array{T}, gfunc, gbarrier, f::F, fbarrier::FB)
    vbarrier = fbarrier(x, gbarrier)
    if isfinite(vbarrier)
        vfunc = f(x, gfunc)
    else
        vfunc = vbarrier
    end
    return vfunc, vbarrier
end

function barrier_combined{T, FB<:Function}(x::Array{T}, g, gfunc, gbarrier, fb::FB, mu::T)
    calc_g = !(g === nothing)
    valfunc, valbarrier = fb(x, gfunc, gbarrier)
    if calc_g
        @simd for i = 1:length(g)
            @inbounds g[i] = gfunc[i] + mu*gbarrier[i]
        end
    end
    return convert(T, valfunc + mu*valbarrier) # FIXME make this unnecessary
end

function limits_box{T}(x::Array{T}, d::Array{T}, l::Array{T}, u::Array{T})
    alphamax = convert(T, Inf)
    for i = 1:length(x)
        if d[i] < 0
            @inbounds alphamax = min(alphamax, ((l[i]-x[i])+eps(l[i]))/d[i])
        elseif d[i] > 0
            @inbounds alphamax = min(alphamax, ((u[i]-x[i])-eps(u[i]))/d[i])
        end
    end
    epsilon = eps(max(alphamax, one(T)))
    if !isinf(alphamax) && alphamax > epsilon
        alphamax -= epsilon
    end
    return alphamax
end

# Default preconditioner for box-constrained optimization
# This creates the inverse Hessian of the barrier penalty
function precondprepbox!(P, x, l, u, mu)
    @inbounds @simd for i = 1:length(x)
        xi = x[i]
        li = l[i]
        ui = u[i]
        P.diag[i] = 1/(mu*(1/(xi-li)^2 + 1/(ui-xi)^2) + 1) # +1 like identity far from edges
    end
end

immutable Fminbox <: Optimizer end

function optimize{T<:AbstractFloat}(
        df::DifferentiableFunction,
        initial_x::Array{T},
        l::Array{T},
        u::Array{T},
        ::Fminbox;
        x_tol::T = eps(T),
        f_tol::T = sqrt(eps(T)),
        g_tol::T = sqrt(eps(T)),
        iterations::Integer = 1_000,
        store_trace::Bool = false,
        show_trace::Bool = false,
        extended_trace::Bool = false,
        callback = nothing,
        show_every = 1,
        linesearch! = nothing,
        linesearch = LineSearches.hagerzhang!,
        eta::Real = convert(T,0.4),
        mu0::T = convert(T, NaN),
        mufactor::T = convert(T, 0.001),
        precondprep! = nothing,
        precondprep = (P, x, l, u, mu) -> precondprepbox!(P, x, l, u, mu),
        optimizer = ConjugateGradient,
        optimizer_o = Options(store_trace = store_trace,
                                          show_trace = show_trace,
                                          extended_trace = extended_trace),
        nargs...)

    # remove in v0.8.0

    linesearch = get_linesearch(linesearch!, linesearch)
    precondprep = get_precondprep(precondprep!, precondprep)

    optimizer == Newton && warning("Newton is not supported as the inner optimizer. Defaulting to ConjugateGradient.")
    x = copy(initial_x)
    fbarrier = (x, gbarrier) -> barrier_box(x, gbarrier, l, u)
    fb = (x, gfunc, gbarrier) -> function_barrier(x, gfunc, gbarrier, df.fg!, fbarrier)
    gfunc = similar(x)
    gbarrier = similar(x)
    P = InverseDiagonal(Array(T, length(initial_x)))
    # to be careful about one special case that might occur commonly
    # in practice: the initial guess x is exactly in the center of the
    # box. In that case, gbarrier is zero. But since the
    # initialization only makes use of the magnitude, we can fix this
    # by using the sum of the absolute values of the contributions
    # from each edge.
    for i = 1:length(gbarrier)
        thisx = x[i]
        thisl = l[i]
        thisu = u[i]
        if thisx < thisl || thisx > thisu
            error("Initial position must be inside the box")
        end
        gbarrier[i] = (isfinite(thisl) ? one(T)/(thisx-thisl) : zero(T)) + (isfinite(thisu) ? one(T)/(thisu-thisx) : zero(T))
    end
    df.g!(x, gfunc)
    mu = isnan(mu0) ? initial_mu(gfunc, gbarrier; mu0factor=mufactor) : mu0
    if show_trace > 0
        println("######## fminbox ########")
        println("Initial mu = ", mu)
    end

    g = similar(x)
    fval_all = Array(Vector{T}, 0)

    # Count the total number of outer iterations
    iteration = 0

    xold = similar(x)
    converged = false
    local results
    first = true
    while !converged && iteration < iterations
        # Increment the number of steps we've had to perform
        iteration += 1

        copy!(xold, x)
        # Optimize with current setting of mu
        funcc = (x, g) -> barrier_combined(x, g, gfunc, gbarrier, fb, mu)
        fval0 = funcc(x, nothing)
        dfbox = DifferentiableFunction(x->funcc(x,nothing), (x,g)->(funcc(x,g); g), funcc)
        if show_trace > 0
            println("#### Calling optimizer with mu = ", mu, " ####")
        end
        pcp = (P, x) -> precondprep(P, x, l, u, mu)
        if optimizer == ConjugateGradient
            _optimizer = optimizer(eta = eta, linesearch = linesearch, P = P, precondprep = pcp)
        elseif optimizer in (LBFGS, GradientDescent)
            _optimizer = optimizer(linesearch = linesearch, P = P, precondprep = pcp)
        elseif optimizer in (NelderMead, SimulatedAnnealing)
            _optimizer = optimizer()
        elseif optimizer == Newton
            _optimizer = ConjugateGradient(eta = eta, linesearch = linesearch, P = P, precondprep = pcp)
        else
            _optimizer = optimizer(linesearch = linesearch)
        end
        resultsnew = optimize(dfbox, x, _optimizer, optimizer_o)
        if first
            results = resultsnew
            first = false
        else
            append!(results, resultsnew)
        end
        copy!(x, minimizer(results))
        if show_trace > 0
            println("x: ", x)
        end

        # Decrease mu
        mu *= mufactor

        # Test for convergence
        @simd for i = 1:length(x)
            @inbounds g[i] = gfunc[i] + mu*gbarrier[i]
        end

        results.x_converged, results.f_converged, results.g_converged, converged = assess_convergence(x, xold, minimum(results), fval0, g, x_tol, f_tol, g_tol)

    end
    results.method = "Fminbox with $(method(results))"
    results.iterations = iteration
    results.initial_x = initial_x
    results.minimum = df.f(minimizer(results))
    results
end


text:
using ITensors
using ITensors.ITensorNetworkMaps

struct InfMPS
  data::Vector{ITensor}
end
Base.length(m::InfMPS) = length(m.data)
Base.copy(m::InfMPS) = InfMPS(copy(m.data))
Base.reverse(m::InfMPS) = InfMPS(reverse(m.data))
Base.iterate(m::InfMPS, args...) = iterate(m.data, args...)
Base.getindex(m::InfMPS, args...) = getindex(m.data, args...)
Base.setindex!(m::InfMPS, args...) = setindex!(m.data, args...)

function infmps(N; χ⃗, d)
  n⃗ = 1:N
  e⃗ = [n => mod1(n + 1, N) for n in 1:N]
  linkindex(χ⃗, e) = Index(χ⃗[e], "l=$(e[1])↔$(e[2])")
  l⃗ = Dict([e .=> linkindex(χ⃗, e) for e in e⃗])
  s⃗ = [Index(d, "s=$n") for n in n⃗]
  neigbhors(n, N) = [mod1(n - 1, N) => n, n => mod1(n + 1, N)]
  return InfMPS([ITensor(getindex.(Ref(l⃗), neigbhors(n, N))..., s⃗[n]) for n in n⃗])
end

ITensors.dag(tn::InfMPS) = InfMPS(dag.(tn.data))
function ITensors.prime(::typeof(linkinds), tn::InfMPS)
  tn_p = copy(tn)
  N = length(tn)
  for i in 1:N, j in (i + 1):N
    l = commoninds(tn[i], tn[j])
    tn_p[i] = prime(tn_p[i]; inds=l)
    tn_p[j] = prime(tn_p[j]; inds=l)
  end
  return tn_p
end

interleave(x::Vector, y::Vector) = permutedims([x y])[:]

function ITensors.linkind(tn::InfMPS, e)
  return commonind(tn[e[1]], tn[e[2]])
end

function transfer_matrix(ψ::InfMPS)
  N = length(ψ)
  ψ′ = prime(linkinds, dag(ψ))
  tn = interleave(reverse(ψ.data), reverse(ψ′.data))
  right_inds = [linkind(ψ, N => 1), linkind(ψ′, N => 1)]
  left_inds = [linkind(ψ, N => 1), linkind(ψ′, N => 1)]
  T = ITensorNetworkMap(tn; input_inds=right_inds, output_inds=left_inds)
  return T
end

function transfer_matrices(ψ::InfMPS)
  N = length(ψ)
  ψ′ = prime(linkinds, dag(ψ))
  # Build from individual transfer matrices
  T⃗ = Vector{ITensorNetworkMap}(undef, N)
  for n in 1:N
    n⁺¹ = mod1(n + 1, N)
    n⁻¹ = mod1(n - 1, N)
    right_inds = [linkind(ψ, n => n⁺¹), linkind(ψ′, n => n⁺¹)]
    left_inds = [linkind(ψ, n⁻¹ => n), linkind(ψ′, n⁻¹ => n)]
    T⃗[n] = ITensorNetworkMap([ψ[n], ψ′[n]]; input_inds=right_inds, output_inds=left_inds)
  end
  return T⃗
end

ext: jl


text:
using ReachabilityAnalysis, Plots

@taylorize function duffing!(du, u, p, t)
    local α = -1.0
    local β = 1.0
    local δ = 0.3
    local γ = 0.37

    x, v = u
    f = γ * cos(ω * t)

    du[1] = u[2]
    du[2] = - α*x - δ*v - β*x^3 + f
end

ω = 1.2
T = 2*pi / ω
X0 = Singleton([1.0, 0.0]) ⊕ BallInf(zeros(2), 0.1)
prob = @ivp(x' = duffing!(x), x(0) ∈ X0, dim=2);

sol = solve(prob, tspan=(0.0, 20*T), alg=TMJets21a());

plot(sol, vars=(0, 1), xlab="t", ylab="x", lw=0.2, color=:blue)

plot(sol, vars=(0, 2), xlab="t", ylab="v", lw=0.2, color=:blue)

plot(sol, vars=(1, 2), xlab="x", ylab="v", lw=0.5, color=:red)


text:
# MSF_compute_v3.jl
# ==================

# master stability functions for the Roessler system with different coupling terms

using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
#using DynamicalSystems
using PyPlot, OrdinaryDiffEq
using StaticArrays

function Roessler_ODE(u,p,t)
    a, b, c = p
    x, y, z = u
    dx =  - y - z
    dy =  x + a*y
    dz = b + z*(x-c)
    return SVector(dx, dy, dz)
end

function Roessler_MSF_ODE(u,p,t)
    a, b, c, alpha, Jh = p
    x1, x2, x3, y1, y2, y3  = u
    dx1 =  - x2 - x3
    dx2 =  x1 + a*x2
    dx3 = b + x3*(x1-c)

   # linearized equation
    Jg = [0 -1 -1 ; 1 a 0 ; x3 0 x1-c]
    dy = (Jg .- alpha.*Jh) * u[4:6]

    return SVector(dx1, dx2, dx3, dy[1], dy[2], dy[3])
end

function MSF(uinit,a,b,c,Jh)

    nle = 20000 # number of renormalization steps
    tle = 6.  # period of time between renormalizations

    alpha_min = -2.
    alpha_max = 8.
    nalpha = 200
    alpha_vec = LinRange(alpha_min, alpha_max, nalpha)
  #  avec = zeros(nalpha)   # alpha values
    LE_vec = zeros(nalpha)  # corresponding Lyapunov exponent

   # uinit = u0[1:3]

    for ialpha = 1:nalpha

        alpha = alpha_vec[ialpha]
        cum = 0.

        p = a, b, c, alpha, Jh
        u0 = [uinit ; SVector(1., 0., 0.) ]

        for ile = 1:nle   # loop over renormalization steps

            tspan = (0., tle)
            prob = ODEProblem(Roessler_MSF_ODE, u0, tspan, p)
            sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8,tsamp = tle)

            yvec = [u[4:6] for u in sol.u]  # perturbation vector
            yend = yvec[end]
            ynorm = sqrt( sum(yend .* yend)  )
            yend = yend / ynorm      # renormalization
            cum = cum + log(ynorm)

            uvec = [u[1:3] for u in sol.u]
            uend = uvec[end]

            u0 = [uend ; yend]   # initial condition for next integration step

        end

        lambda_max = cum / (nle*tle) # Lyapunov exponent
       # println(alpha, "  ", lambda_max)

       # avec[ialpha] = alpha
        LE_vec[ialpha] = lambda_max

    end # of alpha loop
    return alpha_vec, LE_vec
    end # of function

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

# main program
# ------------
a = 0.2
b = 0.2
c = 5.7
#Jh = [0 0 0 ; 0 1 0 ; 0 0 0 ]

# transient
# ---------
p = a, b, c
u0 = SVector(0.1, 0.1, 0.)

tspan = (0., 1000)
prob = ODEProblem(Roessler_ODE, u0, tspan, p)
sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8,tsamp = 100.)

# attractor  (just to check that the solution is o.k.)
# ---------
u = sol.u
# u0 = [u0[end] ; SVector(1., 0., 0.) ]
uinit = u[end]
println(uinit)

p = a, b, c   #, alpha, Jh

tattr = 200.
tspan = (0., tattr)
prob = ODEProblem(Roessler_ODE, uinit, tspan, p)
sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8,tsamp = 0.1)

tvec = sol.t
x1vec = [u[1] for u in sol.u]
x2vec = [u[2] for u in sol.u]
x3vec = [u[3] for u in sol.u]

# plot
# ----
fig = figure("Roessler",figsize=(0.5*figx,figy))
plot(x1vec,x2vec)
xlabel(L"$x_1$")
ylabel(L"$x_2$")

# plot
# ----
fig = figure("Roessler ts",figsize=(0.5*figx,figy))
plot(tvec[1:200],x1vec[1:200])
xlabel(L"$t$")
ylabel(L"$x_1$")

fig.tight_layout(pad=0.3)

# compute MSF
# -----------
fig = figure("MSF 2",figsize=(figx,figy))

LE_min = -0.015
LE_max = 0.02

# coupling signal x
subplot(1,3,1)

Jh = [1 0 0 ; 0 0 0 ; 0 0 0 ]   # x drives x-ODE
alpha_vec, LE_vec_xx = MSF(uinit,a,b,c,Jh)
plot(alpha_vec,LE_vec_xx, color = COLORS[1], linestyle="-")

Jh = [0 0 0 ; 1 0 0 ; 0 0 0 ]  # x drives y-ODE
alpha_vec, LE_vec_xy = MSF(uinit,a,b,c,Jh)
plot(alpha_vec,LE_vec_xy, color = COLORS[1], linestyle=":")

Jh = [0 0 0 ; 0 0 0 ; 1 0 0 ]  # x drives z-ODE
alpha_vec, LE_vec_xz = MSF(uinit,a,b,c,Jh)
plot(alpha_vec,LE_vec_xz, color = COLORS[1], linestyle="--")

# zero line
plot(alpha_vec, 0 .* LE_vec_xx, linewidth=1.0,linestyle="-",color="k" )

xlabel(L"$\alpha$")
ylabel(L"$\lambda_{max}$")

ax = gca()
ax.set_xlim([alpha_vec[1], alpha_vec[end]])
ax.set_ylim([LE_min,LE_max])
ax.legend([L"x \to x"; L"x \to y"; L"x \to z"], fontsize = 28, loc = "upper left")

# coupling signal y
subplot(1,3,2)
Jh = [0 1 0 ; 0 0 0 ; 0 0 0 ]   # y drives x-ODE
alpha_vec, LE_vec_yx = MSF(uinit,a,b,c,Jh)
plot(alpha_vec,LE_vec_yx, color = COLORS[4])

Jh = [0 0 0 ; 0 1 0 ; 0 0 0 ]  # y drives y-ODE
alpha_vec, LE_vec_yy = MSF(uinit,a,b,c,Jh)
plot(alpha_vec,LE_vec_yy, color = COLORS[4], linestyle=":")

Jh = [0 0 0 ; 0 0 0 ; 0 1 0 ]  # y drives z-ODE
alpha_vec, LE_vec_yz = MSF(uinit,a,b,c,Jh)
plot(alpha_vec,LE_vec_yz, color = COLORS[4], linestyle="--")

# zero line
plot(alpha_vec, 0 .* LE_vec_yx, linewidth=1.0,linestyle="-",color="k" )

xlabel(L"$\alpha$")

ax = gca()
ax.set_xlim([alpha_vec[1], alpha_vec[end]])
ax.set_ylim([LE_min,LE_max])
ax.legend([L"y \to x"; L"y \to y"; L"y \to z"], fontsize = 28)

# coupling signal z
subplot(1,3,3)
Jh = [0 0 1 ; 0 0 0 ; 0 0 0 ] # z drives x-ODE
alpha_vec, LE_vec_zx = MSF(uinit,a,b,c,Jh)
plot(alpha_vec,LE_vec_zx, color = COLORS[2])

Jh = [0 0 0 ; 0 0 1 ; 0 0 0 ]  # z drives y-ODE
alpha_vec, LE_vec_zy = MSF(uinit,a,b,c,Jh)
plot(alpha_vec,LE_vec_zy, color = COLORS[2], linestyle=":")

Jh = [0 0 0 ; 0 0 0 ; 0 0 1 ]  # z drives z-ODE
alpha_vec, LE_vec_zz = MSF(uinit,a,b,c,Jh)
plot(alpha_vec,LE_vec_zz, color = COLORS[2], linestyle="--")

# zero line
plot(alpha_vec, 0 .* LE_vec_zx, linewidth=1.0,linestyle="-",color="k" )

xlabel(L"$\alpha$")

ax = gca()
ax.set_xlim([alpha_vec[1], alpha_vec[end]])
ax.set_ylim([LE_min,LE_max])
ax.legend([L"z \to x"; L"z \to y"; L"z \to z"], fontsize = 28)

add_identifiers!(fig)
fig.tight_layout(pad=0.3)

# save
# ----
# dname = "/Users/parlitz/ownCloud/NonlinearDynamicsTextbook/figure_generation/10/MSF/"
# fname = "MSF_Roessler_6_20000"
# ftype = ".png"
# savefig(dname*fname*ftype)

save(datadir("master_stability_function.jld"),"alpha_vec",alpha_vec,
"LE_vec_xx",LE_vec_xx, "LE_vec_xy",LE_vec_xy,"LE_vec_xz",LE_vec_xz,
"LE_vec_yx",LE_vec_yx, "LE_vec_yy",LE_vec_yy,"LE_vec_yz",LE_vec_yz,
"LE_vec_zx",LE_vec_zx, "LE_vec_zy",LE_vec_zy,"LE_vec_zz",LE_vec_zz)

max_stars_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_issues_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_forks_repo_name: JuliaDynamics/NonlinearDynamicsTextbook


text:
#!/usr/bin/env julia

# For the files given in ARGS, create all the needed finalization jobs
# for the third phase. That is, continue them at T₁ until tw=3e8.

const param = "-m 33 -t 1e4 -T 1e8 -n 10 -L"

function get_fields(file)
    regex = r"continue_at_T2_([0-9.]+)_tw1_([0-9]+)_for_T1_([0-9.]+)_L_([0-9]+)_id_([a-z]+-R[AB]).net"
    fields = match(regex, basename(file)).captures
    length(fields) == 0 && error("Couldn't get metadata from $file")
    T2 = parse(Float64, fields[1])
    tw1 = parse(Int64, fields[2])
    T1 = parse(Float64, fields[3])
    L = parse(Int64, fields[4])
    id = fields[5]
    return T1, T2, tw1, L, id
end

# Return the string containing the finalization jobs for .net file
function get_the_job(file)
    T1, T2, tw1, L, id = get_fields(file)
    outname = "finish_T2_$(T2)_tw1_$(tw1)_from_T1_$(T1)_L_$(L)_id_$(id).net"
    needediters = 3e8 - (tw1+1e8)
    cmd = "./exe/annealer -c $(tw1+1e8) $param -p -l $L -i $needediters $T1"
    job = "cat $file | $cmd \"$(outname)\""
    cleanup = "lzma $(outname)"
    return job * "\n" * cleanup * "\n" * "echo 'All done! ($outname)'"
end

jobs = String[]
for file in ARGS
    push!(jobs,get_the_job(file))
end
info("$(length(jobs)) jobs processed.")

n = 1
for job in jobs
    header=
    """
    #!/bin/sh
    #\$ -o \$HOME/glass_finish_$(n).out -j y
    #\$ -N glass_$(n)

    cd tfg/

    """
    open("finish_$n.sh", "w") do io
        println(io,header)
        println(io,job)
    end
    n+=1
end

max_stars_repo_name: RedPointyJackson/tfg

max_issues_repo_name: RedPointyJackson/tfg

max_forks_repo_name: RedPointyJackson/tfg


text:
using Revise
using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using StatsFuns: logistic
using Random
using LinearAlgebra

function run_poisson(n :: Int64, p :: Int64)
    #set random seed
    Random.seed!(1111)

    #simulate data
    k = 10 # number of true predictors
    bernoulli_rates = 0.5rand(p) #minor allele frequencies are drawn from uniform (0, 0.5)
    x = simulate_random_snparray(n, p, bernoulli_rates)
    xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);

    #construct snpmatrix, covariate files, and true model b
    z           = ones(n, 1)                   # non-genetic covariates, just the intercept
    true_b      = zeros(p)                     # model vector
    true_b[1:k] = randn(k)                     # Initialize k non-zero entries in the true model
    shuffle!(true_b)                           # Shuffle the entries
    correct_position = findall(x -> x != 0, true_b) # keep track of what the true entries are

    #check maf
    bernoulli_rates[correct_position]

    #simulate phenotypes under different noises by: y = Xb + noise
    y_temp = xbm * true_b

    # Simulate poisson data
    λ = exp.(y_temp) #inverse log link
    y = [rand(Poisson(x)) for x in λ]
    y = Float64.(y)

    #compute poisson IHT result
    result = L0_poisson_reg(x, z, y, 1, k, glm = "poisson", debias=false, convg=false, show_info=false)

    #check result
    estimated_models = result.beta[correct_position]
    true_model = true_b[correct_position]
    compare_model = DataFrame(
        correct_position = correct_position,
        true_β           = true_model,
        estimated_β      = estimated_models)

    #display results
    @show compare_model
    println("Total iteration number was " * string(result.iter))
    println("Total time was " * string(result.time))
end

Random.seed!(2019)
for i = 1:25
    @info("running the $i th model")
    n = rand(500:2000)
    p = rand(1:10)n
    println("n, p = " * string(n) * ", " * string(p))
    run_poisson(n, p)
end

#some function that runs poisson regression on same SNP matrices, using different model sizes
function run_poisson()

	n, p = 2000, 20000

    #set random seed
    Random.seed!(1111)

    #simulate data
    bernoulli_rates = 0.5rand(p) #minor allele frequencies are drawn from uniform (0, 0.5)
    x = simulate_random_snparray(n, p, bernoulli_rates)
    xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);
    z = ones(n, 1)                   # non-genetic covariates, just the intercept

    for k in 1:30
	    true_b      = zeros(p)                     # model vector
	    true_b[1:k] = randn(k)                     # Initialize k non-zero entries in the true model
	    shuffle!(true_b)                           # Shuffle the entries
	    correct_position = findall(x -> x != 0, true_b) # keep track of what the true entries are

	    #simulate phenotypes under different noises by: y = Xb + noise
	    y_temp = xbm * true_b

	    # Simulate poisson data
	    λ = exp.(y_temp) #inverse log link
	    y = [rand(Poisson(x)) for x in λ]
	    y = Float64.(y)

	    #compute poisson IHT result
	    result = L0_poisson_reg(x, z, y, 1, k, glm = "poisson", debias=false, convg=false, show_info=false)

	    #check result
	    estimated_models = result.beta[correct_position]
	    true_model = true_b[correct_position]
	    compare_model = DataFrame(
	        correct_position = correct_position,
	        true_β           = true_model,
	        estimated_β      = estimated_models)

	    #display results
	    @show compare_model
	    println("Total iteration number was " * string(result.iter))
	    println("Total time was " * string(result.time))
	end
end

function run_logistic(n :: Int64, p :: Int64, debias :: Bool)
    #set random seed
    Random.seed!(1111)

    #simulate data
    k = 10 # number of true predictors
    bernoulli_rates = 0.5rand(p) #minor allele frequencies are drawn from uniform (0, 0.5)
    x = simulate_random_snparray(n, p, bernoulli_rates)
    xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);

    #construct snpmatrix, covariate files, and true model b
    z           = ones(n, 1)                   # non-genetic covariates, just the intercept
    true_b      = zeros(p)                     # model vector
    true_b[1:k] = randn(k)                     # Initialize k non-zero entries in the true model
    shuffle!(true_b)                           # Shuffle the entries
    correct_position = findall(x -> x != 0, true_b) # keep track of what the true entries are

    #simulate phenotypes under different noises by: y = Xb + noise
    y_temp = xbm * true_b

    # Apply inverse logit link and sample from the vector of distributions
    prob = logistic.(y_temp) #inverse logit link
    y = [rand(Bernoulli(x)) for x in prob]
    y = Float64.(y)

    #compute logistic IHT result
    result = L0_logistic_reg(x, z, y, 1, k, glm = "logistic", debias=debias, show_info=false)

    #check result
    estimated_models = result.beta[correct_position]
    true_model = true_b[correct_position]
    compare_model = DataFrame(
        correct_position = correct_position,
        true_β           = true_model,
        estimated_β      = estimated_models)

    #display results
    @show compare_model
    println("n = " * string(n) * ", and p = " * string(p))
    println("Total iteration number was " * string(result.iter))
    println("Total time was " * string(result.time))
end

for i = 1:100
    println("running the $i th model")
    n = rand(500:2000)
    p = rand(1:10)n
    debias = true
    run_logistic(n, p, debias)
end

using Revise
using Random
using LinearAlgebra
using GLMNet
using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using StatsFuns: logistic
using Random
using LinearAlgebra

function iht_lasso_poisson(n :: Int64, p :: Int64)
    #set random seed
    Random.seed!(1111)

    #define maf and true model size
    bernoulli_rates = 0.5rand(p)
    k = 10

    #construct snpmatrix, covariate files, and true model b
    x = simulate_random_snparray(n, p, bernoulli_rates)
    xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);
    z           = ones(n, 1)                   # non-genetic covariates, just the intercept
    true_b      = zeros(p)                     # model vector
    true_b[1:k] = randn(k)                     # Initialize k non-zero entries in the true model
    shuffle!(true_b)                           # Shuffle the entries
    correct_position = findall(x -> x != 0, true_b) # keep track of what the true entries are

    # Simulate poisson data
    y_temp = xbm * true_b
    λ = exp.(y_temp) #inverse log link
    y = [rand(Poisson(x)) for x in λ]
    y = Float64.(y)

    #compute poisson IHT result
    path = collect(1:15)
    num_folds = 5
    folds = rand(1:num_folds, size(x, 1))
    k_est_iht = cv_iht(x, z, y, 1, path, folds, num_folds, use_maf=false, glm="poisson", debias=false)
    iht_result = L0_poisson_reg(x, z, y, 1, k_est, glm = "poisson", debias=false, convg=false, show_info=false, true_beta=true_b)

    #compute poisson lasso result
    x_float = [convert(Matrix{Float64}, x, center=true, scale=true) z]
    cv = glmnetcv(x_float, y, Poisson(), dfmax=15, nfolds=5, folds=folds)
    best = argmin(cv.meanloss)
    lasso_result = cv.path.betas[:, best]
    k_est_lasso = length(findall(!iszero, lasso_result))

    #check result
    IHT_model = iht_result.beta[correct_position]
    lasso_model = lasso_result[correct_position]
    true_model = true_b[correct_position]
    compare_model = DataFrame(
        correct_position = correct_position,
        true_β           = true_model,
        iht_β            = IHT_model,
        lasso_β          = lasso_model)
    @show compare_model

    #compute summary statistics
    lasso_num_correct_predictors = findall(!iszero, lasso_model)
    lasso_false_positives = k_est_lasso - lasso_num_correct_predictors
    lasso_false_negatives = k - lasso_num_correct_predictors
    iht_num_correct_predictors = findall(!iszero, IHT_model)
    iht_false_positives = k_est_iht - iht_num_correct_predictors
    iht_false_negatives = iht_num_correct_predictors
    println("IHT: cv found " * "$k_est_iht predictors, with " * "$iht_false_positives false positives " * "and $iht_false_negatives false negatives")
    println("lasso: cv found " * "$k_est_lasso predictors, with " * "$lasso_false_positives false positives and " * "$lasso_false_negatives false negatives \n\n")

    return lasso_false_positives, lasso_false_negatives, iht_false_positives, iht_false_negatives
end

Random.seed!(2019)
lasso_false_positives = 0
lasso_false_negatives = 0
iht_false_positives = 0
iht_false_negatives = 0
for i = 1:25
    @info("running the $i th model")
    n = rand(100:1000)
    p = rand(1:10)n
    println("n, p = " * string(n) * ", " * string(p))
    lfp, lfn, ifp, ifn = iht_lasso_poisson(n, p)

    lasso_false_positives += lfp
    lasso_false_negatives += lfn
    iht_false_positives += ifp
    iht_false_negatives += ifn
end
println("Lasso: Total number of false positives = $lasso_false_positives")
println("Lasso: Total number of false negatives = $lasso_false_negatives")
println("IHT  : Total number of false positives = $iht_false_positives")
println("IHT  : Total number of false negatives = $iht_false_negatives")

using Revise
using GLMNet #julia wrapper for GLMNet package in R, which calls fortran
using GLM
using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using StatsFuns: logistic
using Random
using LinearAlgebra
using MendelBase

n, p = 223, 2411

#set random seed
Random.seed!(1234)

#define maf and true model size
bernoulli_rates = 0.5rand(p)
k = 5

#construct snpmatrix, covariate files, and true model b
x = simulate_random_snparray(n, p, bernoulli_rates)
xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);
z           = ones(n, 1)                   # non-genetic covariates, just the intercept
true_b      = zeros(p)                     # model vector
true_b[1:k] = randn(k)                     # Initialize k non-zero entries in the true model
shuffle!(true_b)                           # Shuffle the entries
correct_position = findall(x -> x != 0, true_b) # keep track of what the true entries are

# Simulate poisson data
y_temp = xbm * true_b
λ = exp.(y_temp) #inverse log link
y = [rand(Poisson(x)) for x in λ]
y = Float64.(y)

x_float = convert(Matrix{Float64}, x, center=true, scale=true)
x_true = [x_float[:, correct_position] z]
regular_result = glm(x_true, y, Poisson(), LogLink())
regular_result = regular_result.pp.beta0
true_model = true_b[correct_position]
compare_model = DataFrame(
    true_β  = true_model,
    regular_β = regular_result[1:end-1])

#problem diagnosis
using DelimitedFiles
using Statistics
using Plots

#doesn't work at all
sim = 28

#works well
sim = 25
sim = 1

y = readdlm("/Users/biona001/.julia/dev/MendelIHT/docs/IHT_GLM_simulations_mean/data_simulation_$sim" * ".txt")
mean(y)
var(y)
histogram(y)

function simulate_random_snparray(
    n :: Int64,
    p :: Int64,
)
    #first simulate a random {0, 1, 2} matrix with each SNP drawn from Binomial(2, r[i])
    A1 = BitArray(undef, n, p)
    A2 = BitArray(undef, n, p)
    mafs = zeros(Float64, p)
    for j in 1:p
        minor_alleles = 0
        maf = 0
        while minor_alleles <= 5
            maf = 0.5rand()
            for i in 1:n
                A1[i, j] = rand(Bernoulli(maf))
                A2[i, j] = rand(Bernoulli(maf))
            end
            minor_alleles = sum(view(A1, :, j)) + sum(view(A2, :, j))
        end
        mafs[j] = maf
    end

    #fill the SnpArray with the corresponding x_tmp entry
    return _make_snparray(A1, A2), mafs
end

function _make_snparray(A1 :: BitArray, A2 :: BitArray)
    n, p = size(A1)
    x = SnpArray(undef, n, p)
    for i in 1:(n*p)
        c = A1[i] + A2[i]
        if c == 0
            x[i] = 0x00
        elseif c == 1
            x[i] = 0x02
        elseif c == 2
            x[i] = 0x03
        else
            throw(error("matrix shouldn't have missing values!"))
        end
    end
    return x
end

using Revise
using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using BenchmarkTools
using Random
using LinearAlgebra

#run directly
function should_be_slow(n :: Int64, p :: Int64)
    #set random seed
    Random.seed!(1111)
    simulate_random_snparray(n, p)
end
@time should_be_slow(10000, 100000)

using Revise
using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using BenchmarkTools
using Random
using LinearAlgebra

#run once with super small data
function should_be_fast(n :: Int64, p :: Int64)
    #set random seed
    Random.seed!(1111)
    simulate_random_snparray(n, p)
end
should_be_fast(10, 10)
@time should_be_fast(10000, 100000)

function sum_random(x :: Vector{Float64})
   s = 0.0
   for i in x
       s += i
   end
   return s
end

#lets say I really want to sum 10 million random numbers
x = rand(3 * 10^8)
@time sum_random(x) # 2.419125 seconds (22.48 k allocations: 1.161 MiB)

#Can I first run the code on a small problem, and then my larger problem?
x = rand(3 * 10^8)
y = rand(10)
@time sum_random(y) # 0.010076 seconds (22.48 k allocations: 1.161 MiB)
@time sum_random(x) # 0.324422 seconds (5 allocations: 176 bytes)

using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using BenchmarkTools
using Random
using LinearAlgebra
using StatsFuns: logistic

function time_normal_response(
    n :: Int64,  # number of cases
    p :: Int64,  # number of predictors (SNPs)
    k :: Int64,  # number of true predictors per group
    debias :: Bool # whether to debias
)
    Random.seed!(1111)

    #construct snpmatrix, covariate files, and true model b
    x, maf = simulate_random_snparray(n, p)
    xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);
    z = ones(n, 1) # non-genetic covariates, just the intercept
    true_b = zeros(p)
    true_b[1:k] = randn(k)
    shuffle!(true_b)
    correct_position = findall(x -> x != 0, true_b)
    noise = rand(Normal(0, 0.1), n) # noise vectors from N(0, s)

    #simulate phenotypes (e.g. vector y) via: y = Xb + noise
    y = xbm * true_b + noise

    #free memory from xbm.... does this work like I think it works?
    xbm = nothing
    GC.gc()

    #time the result and return
    result = @timed L0_normal_reg(x, z, y, 1, k, debias=debias)
    iter = result[1].iter
    runtime = result[2]      # seconds
    memory = result[3] / 1e6 # MB

    return iter, runtime, memory
end

#time directly:
n = 10000
p = 30000
k = 10
debias = false
iter, runtime, memory = time_normal_response(n, p, k, debias)
#(7, 17.276706309, 119.911048)

#run small example first:
n = 10000
p = 30000
k = 10
debias = false
iter, runtime, memory = time_normal_response(100, 100, 10, debias)
iter, runtime, memory = time_normal_response(n, p, k, debias)
#(14, 2.47283539, 37.778872)
#(7, 12.801054742, 82.478688)

#run same example twice:
n = 10000
p = 30000
k = 10
debias = false
iter, runtime, memory = time_normal_response(n, p, k, debias)
iter, runtime, memory = time_normal_response(n, p, k, debias)
#(7, 15.82780056, 119.911048)
#(7, 13.610244645, 82.478688)

#no GC.gc():
n = 10000
p = 30000
k = 10
debias = false
iter, runtime, memory = time_normal_response(100, 100, 10, debias)
iter, runtime, memory = time_normal_response(n, p, k, debias)
#(14, 2.661537431, 37.778872)
#(7, 12.977118564, 82.478688)

#NORMAL CROSS VALIDATION USING MEMORY MAPPED SNPARRAY - does it fix my problem?
using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using BenchmarkTools
using Random
using LinearAlgebra
using BenchmarkTools
using Distributed

#simulat data
n = 1000
p = 10000
k = 10 # number of true predictors

#set random seed
Random.seed!(2019)

#construct snpmatrix, covariate files, and true model b
x, maf = simulate_random_snparray(n, p, "tmp.bed")
xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);
z = ones(n, 1) # non-genetic covariates, just the intercept
true_b = zeros(p)
true_b[1:k] = randn(k)
shuffle!(true_b)
correct_position = findall(x -> x != 0, true_b)
noise = rand(Normal(0, 0.1), n) # noise vectors from N(0, s)

#simulate phenotypes (e.g. vector y) via: y = Xb + noise
y = xbm * true_b + noise

#specify path and folds
path = collect(1:20)
num_folds = 4
folds = rand(1:num_folds, size(x, 1))

#compute cross validation
# mses = cv_iht(x, z, y, 1, path, folds, num_folds, use_maf = false, glm = "normal", debias=false)
# @time cv_iht_distributed(x, z, y, 1, path, folds, num_folds, use_maf = false, glm = "normal", debias=false, showinfo=false, parallel=true);

@benchmark cv_iht_distributed(x, z, y, 1, path, folds, num_folds, use_maf = false, glm = "normal", debias=false, showinfo=false, parallel=true) seconds=60

rm("tmp.bed", force=true)

########### LOGISTIC CROSS VALIDATION SIMULATION CODE##############
using Revise
using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using StatsFuns: logistic
using BenchmarkTools
using Random
using LinearAlgebra

#simulat data
n = 1000
p = 10000
k = 10    # number of true predictors

#set random seed
Random.seed!(1111)

#construct snpmatrix, covariate files, and true model b
x, maf = simulate_random_snparray(n, p, "tmp.bed")
xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);
z           = ones(n, 1)                   # non-genetic covariates, just the intercept
true_b      = zeros(p)                     # model vector
true_b[1:k] = randn(k)                     # k true response
shuffle!(true_b)                           # Shuffle the entries
correct_position = findall(x -> x != 0, true_b) # keep track of what the true entries are

#simulate bernoulli data
y_temp = xbm * true_b
prob = logistic.(y_temp) #inverse logit link
y = [rand(Bernoulli(x)) for x in prob]
y = Float64.(y)

#specify path and folds
path = collect(1:20)
num_folds = 3
folds = rand(1:num_folds, size(x, 1))

#compute cross validation
mses = cv_iht_test(x, z, y, 1, path, folds, num_folds, use_maf = false, glm = "logistic", debias=true)

K = collect(1:2:10)
errors = map(K) do k
    return k
end

using Revise
using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using BenchmarkTools
using Random
using LinearAlgebra
using BenchmarkTools
using Distributed

#simulat data
n = 1000
p = 10000
k = 10 # number of true predictors

#set random seed
Random.seed!(1111)

#construct snpmatrix, covariate files, and true model b
x, maf = simulate_random_snparray(n, p, "tmp.bed")
xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);
z = ones(n, 1) # non-genetic covariates, just the intercept
true_b = zeros(p)
true_b[1:k] = randn(k)
shuffle!(true_b)
correct_position = findall(x -> x != 0, true_b)
noise = rand(Normal(0, 0.1), n) # noise vectors from N(0, s)

#simulate phenotypes (e.g. vector y) via: y = Xb + noise
y = xbm * true_b + noise

#run k = 1,2,....,20
path = collect(1:20)

#compute IHT result for each k in path
iht_run_many_models(x, z, y, 1, path, "normal", use_maf = false, debias=false, showinfo=false, parallel=true)

using Revise
using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using BenchmarkTools
using Random
using LinearAlgebra
using BenchmarkTools
using Distributed
using DelimitedFiles

#simulat data
n = 1000
p = 10000
k = 10 # number of true predictors

#set random seed
Random.seed!(2019)

#construct snpmatrix, covariate files, and true model b
x, maf = simulate_random_snparray(n, p, "normal.bed")
xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);
z = ones(n, 1) # non-genetic covariates, just the intercept
true_b = zeros(p)
true_b[1:k] = randn(k)
shuffle!(true_b)
correct_position = findall(x -> x != 0, true_b)
noise = rand(Normal(0, 0.1), n) # noise vectors from N(0, s)

#simulate phenotypes (e.g. vector y) via: y = Xb + noise
y = xbm * true_b + noise

#specify path and folds
path = collect(1:20)
num_folds = 5
folds = rand(1:num_folds, size(x, 1))

# convert and save floating point version of the genotype matrix for glmnet
x_float = [convert(Matrix{Float64}, x, center=true, scale=true) z]
writedlm("x_float", x_float)
writedlm("y", y)
writedlm("folds", folds)

# trying GLM's fitting function
Random.seed!(2019)
n = 1000000
p = 10
x = randn(n, p)
b = randn(p)
L = LogitLink()

#simulate bernoulli data
y_temp = x * b
prob = linkinv.(L, y_temp) #inverse logit link
y = [rand(Bernoulli(i)) for i in prob]
y = Float64.(y)

glm_result = fit(GeneralizedLinearModel, x, y, Bernoulli(), L)
hi = glm_result.pp.beta0

glm_result_old = regress(x, y, "logistic")
hii = glm_result_old[1]

[b hi hii]

#deviance residual vs y - xb

function run_once()
    n = 1000
    p = 10000
    k = 10
    d = Poisson
    l = canonicallink(d())

    #construct snpmatrix, covariate files, and true model b
    x, maf = simulate_random_snparray(n, p, "tmp.bed")
    xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);
    z = ones(n, 1) # the intercept
    true_b = zeros(p)
    d == Poisson ? true_b[1:k] = rand(Normal(0, 0.3), k) : true_b[1:k] = randn(k)
    shuffle!(true_b)
    correct_position = findall(x -> x != 0, true_b)

    #simulate phenotypes (e.g. vector y) via: y = Xb + noise
    y_temp = xbm * true_b
    prob = linkinv.(l, y_temp)
    y = [rand(d(i)) for i in prob]
    y = Float64.(y)

    #run IHT
    result = L0_reg(x, xbm, z, y, 1, k, d(), l, debias=false, init=false, show_info=false, convg=true)

    #clean up
    rm("tmp.bed", force=true)

    found = length(findall(!iszero, result.beta[correct_position]))
    runtime = result.time

    return found, runtime
end

#set random seed
Random.seed!(1111)
function run()
    total_found = 0
    total_time = 0
    for i in 1:30
        f, t = run_once()
        total_found += f
        total_time += t
        println("finished $i run")
    end
    avg_found = total_found / 30
    avg_time = total_time / 30
    println(avg_found)
    println(avg_time)
end
run()

function loglik_obs(::Normal, y, μ, wt, ϕ) #this is wrong
    return wt*logpdf(Normal(μ, sqrt(ϕ)), y)
end

function test()
    y, mu = rand(1000), rand(1000)
    ϕ = MendelIHT.deviance(Normal(), y, mu)/length(y)
    ϕ = 1.0
    logl = 0.0
    for i in eachindex(y, mu)
        logl += loglik_obs(Normal(), y[i], mu[i], 1, ϕ)
    end
    println(logl)

    println(loglikelihood(Normal(), y, mu))
end
test()

b = result.beta
μ = linkinv.(l, xbm*b)
# sum(logpdf.(Poisson.(μ), y))
loglikelihood(Poisson(), y, xbm*b)
loglikelihood_test(Poisson(), y, μ)

#debugging malloc: Incorrect checksum for freed object
using Revise
using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using BenchmarkTools
using Random
using LinearAlgebra
using GLM

#simulat data with k true predictors, from distribution d and with link l.
n = 1000
p = 10000
k = 10
d = Normal
l = canonicallink(d())

#set random seed
Random.seed!(1111)

#construct snpmatrix, covariate files, and true model b
x, maf = simulate_random_snparray(n, p, "tmp.bed")
xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);
z = ones(n, 1) # the intercept
true_b = zeros(p)
d == Poisson ? true_b[1:k] = rand(Normal(0, 0.3), k) : true_b[1:k] = randn(k)
shuffle!(true_b)
correct_position = findall(x -> x != 0, true_b)

#simulate phenotypes (e.g. vector y)
y_temp = xbm * true_b
prob = linkinv.(l, y_temp)
y = [rand(d(i)) for i in prob]
y = Float64.(y)

#run IHT
result = L0_reg(x, xbm, z, y, 1, k, d(), l, debias=false, init=false, show_info=false, convg=true)

train_idx = bitrand(n)
path = collect(1:20)

p, q = size(x, 2), size(z, 2)
nmodels = length(path)
betas = zeros(p, nmodels)
cs = zeros(q, nmodels)

x_train = SnpArray(undef, sum(train_idx), p)
copyto!(x_train, @view x[train_idx, :])
y_train = @view(y[train_idx])
z_train = @view(z[train_idx, :])
x_trainbm = SnpBitMatrix{Float64}(x_train, model=ADDITIVE_MODEL, center=true, scale=true);

k = path[1]
debias = false
showinfo = false
d = d()
result = L0_reg(x_train, x_trainbm, z_train, y_train, 1, k, d, l, debias=debias, init=false, show_info=showinfo, convg=true)

#simulate correlated columns

n = 1000
p = 10000
k = 10
d = Bernoulli
l = canonicallink(d())

#set random seed
Random.seed!(1111)

#construct snpmatrix, covariate files, and true model b
x, = simulate_random_snparray(n, p, "tmp.bed")

#ad hoc method to construct correlated columns
c = 0.9
for i in 1:size(x, 1)
    prob = rand(Bernoulli(c))
    prob == 1 && (x[i, 2] = x[i, 1]) #make 2nd column the same as first column 90% of the time
end

tmp = convert(Matrix{Float64}, @view(x[:, 1:2]), center=true, scale=true)
cor(tmp[:, 1], tmp[:, 2])

#testing if IHT finds intercept and/or nongenetic covariates, normal response
using Revise
using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using BenchmarkTools
using Random
using LinearAlgebra
using GLM

#simulat data with k true predictors, from distribution d and with link l.
n = 1000
p = 10000
k = 10
d = Normal
l = canonicallink(d())

#set random seed
Random.seed!(1111)

#construct snpmatrix, covariate files, and true model b
x, = simulate_random_snparray(n, p, "tmp.bed")
xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);
z = ones(n, 2) # the intercept
z[:, 2] .= randn(n)

#define true_b and true_c
true_b = zeros(p)
true_b[1:k-2] = randn(k-2)
shuffle!(true_b)
correct_position = findall(!iszero, true_b)
true_c = [0.1; 0.1]

#simulate phenotype
prob = linkinv.(l, xbm * true_b .+ z * true_c)
y = [rand(d(i)) for i in prob]

#run result
result = L0_reg(x, xbm, z, y, 1, k, d(), l, debias=false, init=false, use_maf=false)

#compare with correct answer
compare_model = DataFrame(
    position    = correct_position,
    true_β      = true_b[correct_position],
    estimated_β = result.beta[correct_position])

compare_model = DataFrame(
    true_c      = true_c[1:2],
    estimated_c = result.c[1:2])

using BenchmarkTools
using Random

function test(z::Matrix{Float64})
    for i in 2:size(z, 2)
        col_mean = mean(z[:, i])
        col_std  = std(z[:, i])
        z[:, i] .= (z[:, i] .- col_mean) ./ col_std
    end
end
#memory estimate:  23.23 MiB
#allocs estimate:  2997
#median time:      8.278 ms (26.30% GC)

function test2(z::Matrix{Float64})
    n, q = size(z)
    for j in 2:q
        μ = 0.0
        for i in 1:n
            μ += z[i, j]
        end
        μ /= n
        σ = 0.0
        for i in 1:n
            σ += (z[i, j] - μ)^2
        end
        σ = sqrt(σ / (n - 1))
        for i in 1:n
            z[i, j] = (z[i, j] - μ) / σ
        end
    end
end
#memory estimate:  0 bytes
#allocs estimate:  0
#median time:      5.629 ms (0.00% GC)

function test3(z::Matrix{Float64})
    n, q = size(z)
    @inbounds for j in 2:q
        μ = 0.0
        @simd for i in 1:n
            μ += z[i, j]
        end
        μ /= n
        σ = 0.0
        for i in 1:n
            σ += (z[i, j] - μ)^2
        end
        σ = sqrt(σ / (n - 1))
        for i in 1:n
            z[i, j] = (z[i, j] - μ) / σ
        end
    end
end
#memory estimate:  0 bytes
#allocs estimate:  0
#median time:      3.289 ms (0.00% GC)

@inline function test4(z::Matrix{Float64})
    n, q = size(z)
    μ = _mean(z)
    σ = _std(z, μ)

    @inbounds for j in 1:q
        @simd for i in 1:n
            z[i, j] = (z[i, j] - μ[j]) * σ[j]
        end
    end
end
#memory estimate:  15.88 KiB
#allocs estimate:  2
#median time:      1.918 ms (0.00% GC)

@inline function _mean(z::Matrix{Float64})
    n, q = size(z)
    μ = zeros(q)
    @inbounds for j in 1:q
        tmp = 0.0
        @simd for i in 1:n
            tmp += z[i, j]
        end
        μ[j] = tmp / n
    end
    return μ
end

function _std(z::Matrix{Float64}, μ::Vector{Float64})
    n, q = size(z)
    σ = zeros(q)

    @inbounds for j in 1:q
        @simd for i in 1:n
            σ[j] += (z[i, j] - μ[j])^2
        end
        σ[j] = 1.0 / sqrt(σ[j] / (n - 1))
    end
    return σ
end
@benchmark _std(z, μ) setup=(z=rand(1000, 1000), μ = _mean(z))
@benchmark _mean(z) setup=(z=rand(1000, 1000))

Random.seed!(2019)
@benchmark test(z) setup=(z=rand(1000, 1000))
@benchmark test2(z) setup=(z=rand(1000, 1000))
@benchmark test3(z) setup=(z=rand(1000, 1000))
@benchmark test4(z) setup=(z=rand(1000, 1000))
@benchmark standardize!(@view(z[:, 2:end])) setup=(z=rand(1000, 1000))

function _scale!(des1::AbstractVector{T}, v1::AbstractVector{T}, w1::AbstractVector{T},
         des2::AbstractVector{T}, v2::AbstractVector{T}, w2::AbstractVector{T}) where {T <: Float64}
    copyto!(des1, v1 .* w1)
    copyto!(des2, v2 .* w2)
end

n=1000

@benchmark _scale!(des1, v1, w1, des2, v2, w2) setup=(des1 = zeros(n),des2 = zeros(n),v1 = rand(n),v2 = rand(n),w1 = rand(n),w2 = rand(n))

full_grad = zeros(2000)
weight = rand(2000)
b = rand(1000)
c = rand(1000)

function test_scale(b, c, full_grad, weight)
    lb = length(b)
    lf = length(full_grad)
    copyto!(v.b, @view(full_grad[1:lb]) ./ @view(v.weight[1:lb]))
    copyto!(v.c, @view(full_grad[lb+1:lf]) ./ @view(v.weight[lb+1:lf]))
end

@benchmark test_scale(b, c, full_grad, weight) setup=(b=rand(100000),c = rand(100000),full_grad = zeros(200000),weight = rand(200000))

function iht_stepsize(v::IHTVariable{T}, z::AbstractMatrix{T},
                      d::UnivariateDistribution, l::Link) where {T <: Float64}

    # first store relevant components of gradient
    copyto!(v.gk, view(v.df, v.idx))
    A_mul_B!(v.xgk, v.zdf2, v.xk, view(z, :, v.idc), v.gk, view(v.df2, v.idc))

    #use zdf2 as temporary storage
    v.xgk .+= v.zdf2
    v.zdf2 .= mueta.(l, v.xb).^2 ./ glmvar.(d, v.μ)

    # now compute and return step size. Note non-genetic covariates are separated from x
    numer = sum(abs2, v.gk) + sum(abs2, @view(v.df2[v.idc]))
    denom = Transpose(v.xgk) * Diagonal(v.zdf2) * v.xgk
    return (numer / denom) :: T
end

function test(x, y, z)
    x .+= y
    y .= mueta.(IdentityLink(), z).^2 ./ glmvar.(Normal(), z)
end

function test2(x, y, z)
    @inbounds for i in eachindex(x)
        x[i] += y[i]
        y[i] = mueta(IdentityLink(), z[i])^2 / glmvar(Normal(), z[i])
    end
end

using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using DelimitedFiles
using BenchmarkTools
using Random
using LinearAlgebra
using GLM

#simulat data with k true predictors, from distribution d and with link l.
n = 1000
p = 10000
k = 10
d = Normal
l = canonicallink(d())
# l = LogLink()

#set random seed
Random.seed!(1111)

for i in 1:10
    #construct SnpArraym, snpmatrix, and non genetic covariate (intercept)
    mafs = clamp!(0.5rand(p), 0.1, 0.5)
    x = simulate_random_snparray(n, p, "test1.bed", mafs=mafs)
    xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);
    z = ones(n, 1)

    # simulate response, true model b, and the correct non-0 positions of b
    y, true_b, correct_position = simulate_random_response(x, xbm, k, d, l)

    # specify weights
    weight = ones(p + 1)
    weight[1:p] .= maf_weights(x, max_weight=2.0)

    #run IHT
    result = L0_reg(x, xbm, z, y, 1, k, d(), l, debias=true, init=false, use_maf=false)
    result_weighted = L0_reg(x, xbm, z, y, 1, k, d(), l, debias=true, init=false, use_maf=false, weight=weight)

    #check result
    compare_model = DataFrame(
        unweighted = result.beta[correct_position],
        weighted   = result_weighted.beta[correct_position])
    @show compare_model

    #clean up
    rm("test1.bed", force=true)
end

function graph(inter)
    n = 1000
    p = 10000
    k = 10
    d = Poisson
    l = canonicallink(d())
    # l = LogLink()

    #set random seed
    Random.seed!(2019)

    #construct SnpArraym, snpmatrix, and non genetic covariate (intercept)
    x = simulate_random_snparray(n, p, "test1.bed")
    xbm = SnpBitMatrix{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true);
    z = ones(n, 1)

    # simulate response, true model b, and the correct non-0 positions of b
    true_b = zeros(p)
    # true_b[1:4] .= [0.1; 0.25; 0.5; 0.8]
    true_b[1:10] .= collect(0.1:0.1:1.0)
    # true_b[1:k] = rand(Normal(0, 0.3), k)
    shuffle!(true_b)
    correct_position = findall(!iszero, true_b)

    #simulate phenotypes (e.g. vector y)
    if d == Normal || d == Poisson || d == Bernoulli
        prob = linkinv.(l, xbm * true_b .+ inter)
        clamp!(prob, -20, 20)
        y = [rand(d(i)) for i in prob]
    elseif d == NegativeBinomial
        nn = 10
        μ = linkinv.(l, xbm * true_b)
        clamp!(μ, -20, 20)
        prob = 1 ./ (1 .+ μ ./ nn)
        y = [rand(d(nn, i)) for i in prob] #number of failtures before nn success occurs
    elseif d == Gamma
        μ = linkinv.(l, xbm * true_b)
        β = 1 ./ μ # here β is the rate parameter for gamma distribution
        y = [rand(d(α, i)) for i in β] # α is the shape parameter for gamma
    end
    y = Float64.(y)
    return histogram(y, bin=50)
end

using Revise
using MendelIHT
using SnpArrays
using DataFrames
using Distributions
using DelimitedFiles
using BenchmarkTools
using Random
using LinearAlgebra
using Statistics
using GLM
using BenchmarkTools

n = 1000
p = 10000
x = simulate_correlated_snparray(n, p, undef, prob=0.75)

col1, col2 = zeros(n), zeros(n)
copyto!(col1, @view(x[:, 1]));
copyto!(col2, @view(x[:, 2]));
cor(col1, col2)
sum(col1 .== col2)

col1, col2 = zeros(n), zeros(n)
for i in 1:19
    copyto!(col1, @view(x[:, i]), center=true, scale=true)
    copyto!(col2, @view(x[:, i+1]), center=true, scale=true)
    println(cor(col1, col2))
end

@benchmark simulate_random_snparray(n, p, undef)
@benchmark simulate_correlated_snparray(n, p, undef)


text:
### A Pluto.jl notebook ###
# v0.18.0

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el)
        el
    end
end

# ╔═╡ 7aa547f8-25d4-488d-9fc3-f633f7f03f57
begin
    using Pkg
	Pkg.activate(joinpath(Pkg.devdir(), "MLCourse"))
    using MLJ, MLJLinearModels, Plots, LinearAlgebra, Zygote, Flux, Random, DataFrames, Distributions, MLCourse
    import MLCourse: poly
end

# ╔═╡ e03882f9-843e-4552-90b1-c47b6cbba19b
using PlutoUI; PlutoUI.TableOfContents()

# ╔═╡ 9f18f610-79e2-403a-a792-fe2dafd2054a
md"# Gradient Descent

## Gradient Descent for Linear Regression
"

# ╔═╡ fcbac567-36db-40c5-9436-6315c0caa40a
function gradient_descent(f, x, η, T; callback = x -> nothing)
    for t in 1:T
        x .-= η * gradient(f, x)[1] # update parameters in direction of -∇f
        callback(x) # the callback will be used to save intermediate values
    end
    x
end

# ╔═╡ 49e744dc-33b7-455a-9cf7-4e8e12f11152
X, y = make_regression(40, 1, rng = 16) # make_regression is an MLJ data generator

# ╔═╡ dcc76840-6552-4a86-8268-1291ab8ea0d3
begin
    lin_reg_loss(β₀, β₁, x, y) = mean((y .- β₀  .- β₁ * x).^2)
    lin_reg_loss(β₀, β₁) = lin_reg_loss(β₀, β₁, X.x1, y)
    lin_reg_loss(β) = lin_reg_loss(β[1], β[2])
end

# ╔═╡ 3ca1e8f3-d69f-454f-b917-4bbe2dcfce01
md"β₀⁽⁰⁾ = $(@bind b0 Slider(-3.:.1:3., show_value = true, default = 0.))

β₁⁽⁰⁾ = $(@bind b1 Slider(-3.:.1:3., show_value = true, default = -3.))

η = $(@bind η Slider(.01:.01:2, show_value = true))

step t = $(@bind t Slider(0:100, show_value = true))
"

# ╔═╡ b1fc14bb-1fd2-4739-a761-7a605fd4559b
begin
    params = [b0, b1] # initial parameters
    lin_reg_path = [copy(params)] # we will copy all parameter values to this list
    gradient_descent(lin_reg_loss, params, η, 100,
                     callback = x -> push!(lin_reg_path, copy(x)))
end

# ╔═╡ bbc0b514-4789-44d1-8d90-9fc325d9ad6b
let
    p1 = scatter(X.x1, y, xlim = (-2.5, 2.5), ylim = (-2, 2),
                 legend = :bottomright,
                 xlabel = "x", ylabel = "y", label = "training data")
    plot!(x -> lin_reg_path[t+1]' * [1, x], c = :red, w = 2, label = "current fit")
    p2 = contour(-3:.1:3, -3:.1:3, lin_reg_loss, cbar = false, aspect_ratio = 1)
    scatter!(first.(lin_reg_path), last.(lin_reg_path), markersize = 1, c = :red, label = nothing)
    scatter!([lin_reg_path[t+1][1]], [lin_reg_path[t+1][2]], label = nothing,
             markersize = 4, c = :red, xlabel = "β₀", ylabel = "β₁")
    plot(p1, p2, layout = (1, 2), size = (700, 400))
end

# ╔═╡ eb9e2b3f-dc4b-49a2-a611-b45f2918adcf
md"## Gradient Descent for Logistic Regression"

# ╔═╡ 402bb576-8945-403e-a9a6-fd5bfb8016bc
begin
    X2, y2 = make_regression(50, 1, binary = true, rng = 161)
	y2 = coerce(y2, MLJ.Continuous) .- 1
    σ(x) = 1/(1 + exp(-x))
    log_reg_loss(β) = log_reg_loss(β[1], β[2])
    function log_reg_loss(β₀, β₁)
        p = σ.(β₀  .+ β₁ * X2.x1) # probability of positive class for all inputs
        -mean(@. y2 * log(p) + (1-y2) * log(1 - p)) # negative log-likelihood
    end
end

# ╔═╡ cd5f079f-a06d-4d55-9666-e2b05ddf8989
md"β₀⁽⁰⁾ = $(@bind b02 Slider(-3.:.1:3., show_value = true, default = -2.))

β₁⁽⁰⁾ = $(@bind b12 Slider(-3.:.1:3., show_value = true, default = 2.))

η = $(@bind η2 Slider(.1:.1:2, show_value = true))

step t = $(@bind t2 Slider(0:100, show_value = true))
"

# ╔═╡ 93699313-5e42-48a7-abc6-ad146e5bdcdd
begin
    params2 = [b02, b12]
    log_reg_path = [copy(params2)]
    gradient_descent(log_reg_loss, params2, η2, 100,
                     callback = x -> push!(log_reg_path, copy(x)))
end

# ╔═╡ 33bdc912-e46a-4310-9184-733be7871768
let
    p1 = scatter(X2.x1, y2, xlim = (-3.5, 3.5), ylim = (-.1, 1.1),
                 legend = :bottomleft, marker_style = :vline,
                 xlabel = "x", ylabel = "y", label = "training data")
    plot!(x -> σ(log_reg_path[t2+1]' * [1, x]), c = :red, w = 2, label = "current fit")
    p2 = contour(-3:.1:3, -3:.1:3, log_reg_loss, cbar = false, aspect_ratio = 1)
    scatter!(first.(log_reg_path), last.(log_reg_path), markersize = 1, c = :red, label = nothing)
    scatter!([log_reg_path[t2+1][1]], [log_reg_path[t2+1][2]], label = nothing,
             markersize = 4, c = :red, xlabel = "β₀", ylabel = "β₁")
    plot(p1, p2, layout = (1, 2), size = (700, 400))
end

# ╔═╡ 979feb42-c328-44b5-b519-8e8cda474140
md"## Gradient Descent to Minimize a Complicated Function"

# ╔═╡ c448238c-f712-4af3-aedb-cec3a2c1a73e
begin
	Random.seed!(5)
    f(x, θ) = θ[1] * sin.(θ[2] * x .+ θ[3]) .+ θ[4] * sin.(θ[5] * x .+ θ[6])
    f(x) = 0.3 * sin(2x) + .7 * sin(3.5x + 1)
    x3 = randn(50)
    y3 = f.(x3)
end

# ╔═╡ 6e38a3f6-a592-4dc1-a6c4-d0f0050c9399
special_loss(θ) = mean((y3 .- f(x3, θ)).^2)

# ╔═╡ 4b3ac87f-39f6-4e6c-b7b3-536925e9b112
md"step t = $(@bind t3 Slider(1:10^4, show_value = true))

seed = $(@bind special_seed Slider(1:10, default = 3, show_value = true))"

# ╔═╡ 9a61332f-cdc2-4129-b7b5-5ab54ba387a3
begin
	Random.seed!(special_seed) # defined by the slider below
    params3 = .1 * randn(6)
    special_path = [copy(params3)]
    gradient_descent(special_loss, params3, .1, 10^4,
                     callback = x -> push!(special_path, copy(x)))
end

# ╔═╡ 01221937-b6d9-4dac-be90-1b8a1c5e9d87
let
    p1 = plot(f, label = "orginal function", xlabel = "x", ylabel = "y",
              ylim = (-1.3, 1.2), xlim = (-3, 3))
    plot!(x -> f(x, special_path[t3]), label = "current fit")
	scatter!(x3, y3, label = "training data")
    th = round.(special_path[t3], digits = 2)
    annotate!([(0, -1.1, text("f̂(x) = $(th[1]) * sin($(th[2])x + $(th[3])) + $(th[4]) * sin($(th[5])x + $(th[6]))", pointsize = 7))])
    losses = special_loss.(special_path)
    p2 = plot(0:10^4, losses, label = "learning curve", c = :black, yscale = :log10)
    scatter!([t3], [losses[t3]], label = "current loss", xlabel = "t", ylabel = "loss")
    p3 = contour(-4:.1:4, -4:.1:4, (x2, x5) -> special_loss([special_path[t3][1]; x2; special_path[t3][3:4]; x5; special_path[t3][6]]), xlabel = "θ₂", ylabel = "θ₅", title = "loss")
	scatter!([special_path[t3][2]], [special_path[t3][5]])
    plot(p1, plot(p2, p3, layout = (2, 1)), layout = (1, 2), size = (700, 400))
end

# ╔═╡ f5e27275-8751-4e80-9888-c3d22d8e80e3
md"*Optional* If you want to know more about the magic of automatic differentiation and how julia computes derivates of (almost) arbitrary code: have a look at this [didactic introduction to automatic differentiation](https://github.com/MikeInnes/diff-zoo) or this [video with code examples in python](https://www.youtube.com/watch?v=wG_nF1awSSY&t)."

# ╔═╡ ca1301fd-8978-44d6-bfae-e912f939d7a8
md"# Stochastic Gradient Descent

In each step of stochastic gradient descent (SGD) the gradient is computed only on a (stochastically selected) subset of the data. As long as the selected subset of the data is somewhat representative of the full dataset the gradient will point more or less in the same direction as the full gradient computed on the full dataset. The advantage of computing the gradient only on a subset of the data is that it takes much less time to compute the gradient on a small dataset than on a large dataset. In the figure below you see the niveau lines of the full loss in black and the niveau lines of the loss computed on a green, purple, yellow and a blue subset. The gradient computed on, say, the yellow subset is always perpendicular to the yellow niveau lines but it may not be perpendicular to the niveau lines of the full loss. When selecting in each step another subset we observe a jittered trajectory around the full gradient descent trajectory.
"

# ╔═╡ 75528011-05d9-47dc-a37b-e6bb6be52c25
md"η = $(@bind η_st Slider(.01:.01:.1, show_value = true))

step t = $(@bind t5 Slider(0:100, show_value = true))
"

# ╔═╡ e821fb15-0bd3-4fa7-93ea-692bf05097b5
begin
    lin_reg_loss_st = let batch = rand(1:4, 101), i = 0
        function(β)
            xb, yb = Zygote.ignore() do
                i += 1
                idxs = (batch[i]-1)*10 + 1:batch[i]*10
                X.x1[idxs], y[idxs]
            end
            lin_reg_loss(β[1], β[2], xb, yb)
        end
    end
    lin_reg_loss_b(i) = (β₀, β₁) -> lin_reg_loss(β₀, β₁, X.x1[(i-1)*10+1:i*10],
                                                 y[(i-1)*10+1:i*10])
    params_st = [b0, b1]
    lin_reg_stoch_path = [copy(params_st)]
    gradient_descent(lin_reg_loss_st, params_st, η_st, 100,
                     callback = x -> push!(lin_reg_stoch_path, copy(x)))
end;

# ╔═╡ 7541c203-f0dc-4445-9d2a-4cf16b7e912a
let
    b = lin_reg_loss_st.batch[t5+1]
    colors = [:green, :blue, :orange, :purple]
    ma = fill(.2, 40)
    ma[(b-1)*10 + 1:b*10] .= 1
    p1 = scatter(X.x1, y, xlim = (-2.5, 2.5), ylim = (-2, 2),
                 legend = false, ma = ma,
                 c = vcat([fill(c, 10)
                           for c in colors]...),
                 xlabel = "x", ylabel = "y", label = "training data")
    plot!(x -> lin_reg_stoch_path[t5+1]' * [1, x], c = :red, w = 2,
          label = "current fit")
    p2 = contour(-3:.1:3, -3:.1:3, lin_reg_loss, cbar = false,
                 linestyle = :dash, c = :black, aspect_ratio = 1)
    for i in 1:4
        contour!(-3:.1:3, -3:.1:3, lin_reg_loss_b(i), cbar = false, w = 2,
                 linestyle = :dash, c = colors[i], alpha = b == i ? 1 : .2)
    end
    plot!(first.(lin_reg_path), last.(lin_reg_path),
          c = :black, w = 3, label = "GD")
    plot!(first.(lin_reg_stoch_path), last.(lin_reg_stoch_path),
          c = :red, w = 1.5, label = "SGD")
    ps = lin_reg_stoch_path[t5+1]
    scatter!([ps[1]], [ps[2]], label = nothing, c = :red)
    plot(p1, p2, layout = (1, 2), size = (700, 400))
end

# ╔═╡ 913cf5ee-ca1e-4063-bd34-6cccd0cc548b
md"# Improved Versions of (S)GD

There are many tricks to improve over standard (stochastic) gradient descent.
One popular idea is to use [momentum](https://distill.pub/2017/momentum/).
We do not discuss these ideas further here, but you should know that `ADAM()` and `ADAMW()` are particularly popular (and successful) improvements of standard (S)GD. These methods usually require no (or very little) tuning of the learning rate."

# ╔═╡ eb289254-7167-4183-a4d0-52f68be66b04
begin
	Random.seed!(1234)
    params4 = .1 * randn(6)
    special_path2 = [copy(params4)]
	opt = ADAMW()
    for _ in 1:10^4
        Flux.Optimise.train!(_ -> special_loss(params4), [params4], [nothing], opt,
                             cb = () -> push!(special_path2, copy(params4)))
    end
end

# ╔═╡ 166472c5-c0f4-4261-a476-4c9b0f82abd6
let
    p1 = plot(f, label = "orginal function", xlabel = "x", ylabel = "y",
              ylim = (-1.3, 1.2))
    plot!(x -> f(x, special_path2[t3]), label = "current fit")
    th = round.(special_path2[t3], digits = 2)
    annotate!([(0, -1.1, text("f̂(x) = $(th[1]) * sin($(th[2])x + $(th[3])) + $(th[4]) * sin($(th[5])x + $(th[6]))", pointsize = 7))])
    losses2 = special_loss.(special_path2)
    losses = special_loss.(special_path)
    p2 = plot(0:10^4, losses2, label = "learning curve", c = :black, yscale = :log10)
    plot!(0:10^4, losses, label = "GD learning curve", c = :black, linestyle = :dot)
    scatter!([t3], [losses2[t3]], label = "current loss", xlabel = "t", ylabel = "loss", legend = :bottomleft)
    plot(p1, p2, layout = (1, 2), size = (700, 400))
end

# ╔═╡ 08a9418f-786e-4992-b1a5-04cf9060f8fe
md"# Early Stopping as Regularization"

# ╔═╡ dc57d700-2a82-4ab0-9bd2-6ce622cb0fa5
begin
    g(x) = .3 * sin(10x) + .7x
    function regression_data_generator(; n, seed = 3, rng = MersenneTwister(seed))
        x = range(0, 1, length = n)
        DataFrame(x = x, y = g.(x) .+ .1*randn(rng, n))
    end
    regression_data = regression_data_generator(n = 10, seed = 8)
	regression_valid = regression_data_generator(n = 50, seed = 123)
end;

# ╔═╡ c72dc506-fb1d-42ee-83cf-cc49753ecd4f
begin
    h_training = Array(select(poly(regression_data, 12), Not(:y)))
    h_valid = Array(select(poly(regression_valid, 12), Not(:y)))
    poly_regression_loss(θ, x, y) = mean((y .- θ[1] .- x * θ[2:end]).^2)
    poly_params = 1e-3 * randn(13)
    poly_path = [copy(poly_params)]
	poly_opt = ADAMW()
    for _ in 1:10^5
        Flux.Optimise.train!((x, y) -> poly_regression_loss(poly_params, x, y),
                             [poly_params], [(h_training, regression_data.y)],
                             poly_opt,
                             cb = () -> push!(poly_path, copy(poly_params)))
    end
end

# ╔═╡ 075f35cf-4271-4676-b9f3-e2bcf610c2d1
md"step t = $(@bind t4 Slider(0:10^3:10^5, show_value = true))"

# ╔═╡ 0d431c00-9eef-4ce4-9542-9571728d1501
let
    p1 = scatter(regression_data.x, regression_data.y,
                 label = "training data", ylims = (-.1, 1.1))
    plot!(g, label = "generator", c = :green, w = 2)
    grid = 0:.01:1
    θ = poly_path[t4 + 1]
    pred =  Array(poly((x = grid,), 12)) * θ[2:end] .+ θ[1]
    plot!(grid, pred,
          label = "fit", w = 3, c = :red, legend = :topleft)
    losses = poly_regression_loss.(poly_path, Ref(h_training), Ref(regression_data.y))
    losses_v = poly_regression_loss.(poly_path, Ref(h_valid), Ref(regression_valid.y))
    p2 = plot(0:length(poly_path)-1, losses, yscale = :log10,
              c = :blue, label = "training loss")
    scatter!([t4], [losses[t4 + 1]], c = :blue, label = nothing)
    plot!(0:length(poly_path)-1, losses_v, c = :red, label = "validation loss")
    scatter!([t4], [losses_v[t4 + 1]], c = :red, label = nothing)
    vmin, idx = findmin(losses_v)
    vline!([idx], c = :red, linestyle = :dash, label = nothing)
    hline!([vmin], c = :red, linestyle = :dash, label = nothing)
    plot(p1, p2, layout = (1, 2), size = (700, 400))
end

# ╔═╡ e0cc188c-9c8f-47f1-b1fe-afc2a578973d
md"# Exercise

Assume the noise in a linear regression setting comes from a Laplace
distribution, i.e. the conditional probability density of the response is given by
``p(Y = y | X = x, \beta) = \frac1{2s}\exp(-\frac{|y - x^T\beta|}{s})``.
For simplicity we assume throughout this exercise that the intercept ``\beta_0 = 0``
and does not need to be fitted.

(a) Generate a training set of 100 points with the following data generator.
Notice that the noise follows a Laplace distribution instead of a normal distribution.
For once, we do not use a `DataFrame` here but represent the input explicitly as a matrix and the full dataset as a `NamedTuple`. If you run `data = data_generator()`, you can access the input matrix as `data.x` and the output vector as `data.y`.
```julia
function data_generator(; n = 100, β = [1., 2., 3.])
    x = randn(n, 3)
    y = x * β .+ rand(Laplace(0, 0.3), n)
    (x = x, y = y)
end
```

(b) Calculate with paper and pencil the negative log-likelihood loss. Apply transformations to the negative log-likelihood function to obtain a good loss function for gradient descent based on the practical considerations in the slides.

(c) Write the code to compute the loss on the training set for a given
parameter vector. *Hint:* use matrix multiplication, e.g. `data.x * β`.

(d) Perform gradient descent on the training set. Plot the learning curve to see
whether gradient descent has converged. If you see large fluctuations at the end
of training, decrease the learning rate. If the learning curve is not flat at
the end, increase the maximal number of steps. To see well the loss towards the end of gradient descent it is advisable to use log-scale for the y-axis (`yscale = :log10`).

(e) Estimate the coefficients with the standard linear regression.
Hint: do not forget that we fit without intercept (use `fit_intercept = false` in the `LinearRegressor`).

(f) Compare which method (d) or (e) found parameters closer to the one of our data generating process `[1, 2, 3]` and explain your finding.
"

# ╔═╡ cb9f858a-f60a-11eb-3f0e-a9b68cf33921
MLCourse.list_notebooks(@__FILE__)

# ╔═╡ 8459f86e-bce7-4839-9c51-57335ac6353c
MLCourse.footer()

# ╔═╡ Cell order:
# ╟─e03882f9-843e-4552-90b1-c47b6cbba19b
# ╠═7aa547f8-25d4-488d-9fc3-f633f7f03f57
# ╟─9f18f610-79e2-403a-a792-fe2dafd2054a
# ╠═fcbac567-36db-40c5-9436-6315c0caa40a
# ╠═49e744dc-33b7-455a-9cf7-4e8e12f11152
# ╠═dcc76840-6552-4a86-8268-1291ab8ea0d3
# ╟─3ca1e8f3-d69f-454f-b917-4bbe2dcfce01
# ╠═b1fc14bb-1fd2-4739-a761-7a605fd4559b
# ╟─bbc0b514-4789-44d1-8d90-9fc325d9ad6b
# ╟─eb9e2b3f-dc4b-49a2-a611-b45f2918adcf
# ╠═402bb576-8945-403e-a9a6-fd5bfb8016bc
# ╟─cd5f079f-a06d-4d55-9666-e2b05ddf8989
# ╠═93699313-5e42-48a7-abc6-ad146e5bdcdd
# ╟─33bdc912-e46a-4310-9184-733be7871768
# ╟─979feb42-c328-44b5-b519-8e8cda474140
# ╠═c448238c-f712-4af3-aedb-cec3a2c1a73e
# ╠═6e38a3f6-a592-4dc1-a6c4-d0f0050c9399
# ╠═9a61332f-cdc2-4129-b7b5-5ab54ba387a3
# ╟─4b3ac87f-39f6-4e6c-b7b3-536925e9b112
# ╟─01221937-b6d9-4dac-be90-1b8a1c5e9d87
# ╟─f5e27275-8751-4e80-9888-c3d22d8e80e3
# ╟─ca1301fd-8978-44d6-bfae-e912f939d7a8
# ╠═e821fb15-0bd3-4fa7-93ea-692bf05097b5
# ╟─75528011-05d9-47dc-a37b-e6bb6be52c25
# ╟─7541c203-f0dc-4445-9d2a-4cf16b7e912a
# ╟─913cf5ee-ca1e-4063-bd34-6cccd0cc548b
# ╠═eb289254-7167-4183-a4d0-52f68be66b04
# ╟─166472c5-c0f4-4261-a476-4c9b0f82abd6
# ╟─08a9418f-786e-4992-b1a5-04cf9060f8fe
# ╠═dc57d700-2a82-4ab0-9bd2-6ce622cb0fa5
# ╠═c72dc506-fb1d-42ee-83cf-cc49753ecd4f
# ╟─075f35cf-4271-4676-b9f3-e2bcf610c2d1
# ╟─0d431c00-9eef-4ce4-9542-9571728d1501
# ╟─e0cc188c-9c8f-47f1-b1fe-afc2a578973d
# ╟─cb9f858a-f60a-11eb-3f0e-a9b68cf33921
# ╟─8459f86e-bce7-4839-9c51-57335ac6353c

max_stars_repo_name: baltisberger/MLCourse

max_issues_repo_name: baltisberger/MLCourse

max_forks_repo_name: baltisberger/MLCourse


text:
using LinearAlgebra # For rmul!

using MathOptInterface
const MOI = MathOptInterface
const MOIU = MOI.Utilities
const AFFEQ = MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}
const BoundsSet = Union{MOI.GreaterThan{Float64}, MOI.LessThan{Float64}, MOI.Interval{Float64}}

@enum VariableType NNEG PSD

struct VariableInfo
    variable_type::VariableType
    cone_index::Int
    index_in_cone::Int
end

mutable struct Optimizer <: MOI.AbstractOptimizer
    b::Vector{Float64}
    l::Vector{Float64}
    u::Vector{Float64}

    variable_info::Vector{VariableInfo}
    single_variable_mask::Vector{UInt8}

    # PSDCone variables
    psdc_dims::Vector{Int}
    psdc_Cvar::Vector{Vector{Int}}
    psdc_Cval::Vector{Vector{Float64}}
    psdc_Avar::Vector{Vector{Int}}
    psdc_Acon::Vector{Vector{Int}}
    psdc_Aval::Vector{Vector{Float64}}
    psdc_Bvar::Vector{Vector{Int}}
    psdc_Bcon::Vector{Vector{Int}}
    psdc_Bval::Vector{Vector{Float64}}
    psdc_L::Vector{Matrix{Float64}}
    psdc_U::Vector{Matrix{Float64}}

    # NonNEGatives variables
    num_nneg::Int
    nneg_info::Vector{Int} # Similar to `info` field of `MOI.Bridges.Variable.Map`.
    nneg_Cvar::Vector{Int}
    nneg_Cval::Vector{Float64}
    nneg_Avar::Vector{Int}
    nneg_Acon::Vector{Int}
    nneg_Aval::Vector{Float64}
    nneg_Bvar::Vector{Int}
    nneg_Bcon::Vector{Int}
    nneg_Bval::Vector{Float64}
    nneg_L::Vector{Float64}
    nneg_U::Vector{Float64}

    objective_sense::MOI.OptimizationSense
    objective_constant::Float64

    primal_objective_value::Float64
    dual_objective_value::Float64
    nneg_X::Vector{Float64}
    psdc_X::Vector{Vector{Float64}}
    s::Vector{Float64}
    y::Vector{Float64}
    y2::Vector{Float64}
    nneg_Z1::Vector{Float64}
    nneg_Z2::Vector{Float64}
    psdc_Z1::Vector{Vector{Float64}}
    psdc_Z2::Vector{Vector{Float64}}
    info::Dict{String, Any}
    runhist::Dict{String, Any}
    status::Union{Nothing, Int}
    solve_time::Float64

    silent::Bool
    options::Dict{Symbol, Any}
    function Optimizer(; kwargs...)
        optimizer = new(
            Float64[], Float64[], Float64[], VariableInfo[], UInt8[],
            Int[], Vector{Int}[], Vector{Float64}[], Vector{Int}[], Vector{Int}[], Vector{Float64}[], Vector{Int}[], Vector{Int}[], Vector{Float64}[], Matrix{Float64}[], Matrix{Float64}[],
            0, Int[], Int[], Float64[], Int[], Int[], Float64[], Int[], Int[], Float64[], Float64[], Float64[],
            MOI.FEASIBILITY_SENSE, 0.0,
            NaN, NaN,
            Float64[], Vector{Float64}[], # X
            Float64[], # s
            Float64[], Float64[], # y, y2 ≈ v
            Float64[], Float64[], Vector{Float64}[], Vector{Float64}[], # Z
            Dict{String, Any}(), Dict{String, Any}(),
            nothing, NaN,
            false, Dict{Symbol, Any}())
        for (key, value) in kwargs
            MOI.set(optimizer, MOI.RawParameter(string(key)), value)
        end
        return optimizer
    end
end

function MOI.supports(optimizer::Optimizer, param::MOI.RawParameter)
    return param.name in ALLOWED_OPTIONS
end
function MOI.set(optimizer::Optimizer, param::MOI.RawParameter, value)
    if !MOI.supports(optimizer, param)
        throw(MOI.UnsupportedAttribute(param))
    end
    optimizer.options[Symbol(param.name)] = value
end
function MOI.get(optimizer::Optimizer, param::MOI.RawParameter)
    # TODO: This gives a poor error message if the name of the parameter is invalid.
    return optimizer.options[Symbol(param.name)]
end

MOI.supports(::Optimizer, ::MOI.Silent) = true
function MOI.set(optimizer::Optimizer, ::MOI.Silent, value::Bool)
    optimizer.silent = value
end
MOI.get(optimizer::Optimizer, ::MOI.Silent) = optimizer.silent

MOI.get(::Optimizer, ::MOI.SolverName) = "SDPNAL"
function MOI.get(optimizer::Optimizer, ::MOI.SolveTime)
    return optimizer.solve_time
end

function MOI.is_empty(optimizer::Optimizer)
    return isempty(optimizer.b) && isempty(optimizer.l) && isempty(optimizer.u) &&
        isempty(optimizer.variable_info) &&
        iszero(optimizer.psdc_dims) &&
        iszero(optimizer.num_nneg) &&
        optimizer.objective_sense == MOI.FEASIBILITY_SENSE &&
        iszero(optimizer.objective_constant)
end

function MOI.empty!(optimizer::Optimizer)
    empty!(optimizer.b)
    empty!(optimizer.l)
    empty!(optimizer.u)

    empty!(optimizer.variable_info)
    empty!(optimizer.single_variable_mask)

    empty!(optimizer.psdc_dims)
    empty!(optimizer.psdc_Cvar)
    empty!(optimizer.psdc_Cval)
    empty!(optimizer.psdc_Avar)
    empty!(optimizer.psdc_Acon)
    empty!(optimizer.psdc_Aval)
    empty!(optimizer.psdc_Bvar)
    empty!(optimizer.psdc_Bcon)
    empty!(optimizer.psdc_Bval)
    empty!(optimizer.psdc_L)
    empty!(optimizer.psdc_U)

    optimizer.num_nneg = 0
    empty!(optimizer.nneg_info)
    empty!(optimizer.nneg_Cvar)
    empty!(optimizer.nneg_Cval)
    empty!(optimizer.nneg_Avar)
    empty!(optimizer.nneg_Acon)
    empty!(optimizer.nneg_Aval)
    empty!(optimizer.nneg_Bvar)
    empty!(optimizer.nneg_Bcon)
    empty!(optimizer.nneg_Bval)
    empty!(optimizer.nneg_L)
    empty!(optimizer.nneg_U)

    optimizer.objective_sense = MOI.FEASIBILITY_SENSE
    optimizer.objective_constant = 0.0

    optimizer.primal_objective_value = NaN
    optimizer.dual_objective_value = NaN
    empty!(optimizer.nneg_X)
    empty!(optimizer.psdc_X)
    empty!(optimizer.s)
    empty!(optimizer.y)
    empty!(optimizer.y2)
    empty!(optimizer.nneg_Z1)
    empty!(optimizer.nneg_Z2)
    empty!(optimizer.psdc_Z1)
    empty!(optimizer.psdc_Z2)
    empty!(optimizer.info)
    empty!(optimizer.runhist)
    optimizer.status = nothing
    optimizer.solve_time = NaN
end

function MOI.supports(
    optimizer::Optimizer,
    ::Union{MOI.ObjectiveSense,
            MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}})
    return true
end

MOI.supports_add_constrained_variables(::Optimizer, ::Type{MOI.Reals}) = false

const SupportedSets = Union{MOI.Nonnegatives, MOI.PositiveSemidefiniteConeTriangle}
MOI.supports_add_constrained_variables(::Optimizer, ::Type{<:SupportedSets}) = true
function MOI.supports_constraint(
    ::Optimizer, ::Type{MOI.ScalarAffineFunction{Float64}},
    ::Type{<:Union{MOI.EqualTo{Float64}, BoundsSet}})
    return true
end

function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike; kws...)
    return MOIU.automatic_copy_to(dest, src; kws...)
end
MOIU.supports_default_copy_to(::Optimizer, copy_names::Bool) = !copy_names

# Variables
function _add_nonneg_variable(optimizer::Optimizer)
    optimizer.num_nneg += 1
    push!(optimizer.variable_info, VariableInfo(NNEG, 1, optimizer.num_nneg))
    push!(optimizer.single_variable_mask, 0x0)
    return MOI.VariableIndex(length(optimizer.variable_info))
end
function _add_psdc_variable(optimizer::Optimizer, index_in_cone)
    push!(optimizer.variable_info, VariableInfo(PSD, length(optimizer.psdc_dims), index_in_cone))
    push!(optimizer.single_variable_mask, 0x0)
    return MOI.VariableIndex(length(optimizer.variable_info))
end
function _add_constrained_variables(optimizer::Optimizer, set::MOI.Nonnegatives)
    push!(optimizer.nneg_info, MOI.dimension(set))
    for i in 2:MOI.dimension(set)
        push!(optimizer.nneg_info, i)
    end
    for i in 1:MOI.dimension(set)
        push!(optimizer.nneg_L, -Inf)
        push!(optimizer.nneg_U, Inf)
    end
    return [_add_nonneg_variable(optimizer) for i in 1:MOI.dimension(set)]
end
function _add_constrained_variables(optimizer::Optimizer, set::MOI.PositiveSemidefiniteConeTriangle)
    d = MOI.side_dimension(set)
    push!(optimizer.psdc_dims, d)
    push!(optimizer.psdc_Cvar, Int[])
    push!(optimizer.psdc_Cval, Float64[])
    push!(optimizer.psdc_Avar, Int[])
    push!(optimizer.psdc_Acon, Int[])
    push!(optimizer.psdc_Aval, Float64[])
    push!(optimizer.psdc_Bvar, Int[])
    push!(optimizer.psdc_Bcon, Int[])
    push!(optimizer.psdc_Bval, Float64[])
    push!(optimizer.psdc_L, fill(-Inf,  d, d))
    push!(optimizer.psdc_U, fill(Inf,  d, d))
    return [_add_psdc_variable(optimizer, i) for i in 1:MOI.dimension(set)]
end
function MOI.add_constrained_variables(optimizer::Optimizer, set::SupportedSets)
    vis = _add_constrained_variables(optimizer, set)
    ci = MOI.ConstraintIndex{MOI.VectorOfVariables, typeof(set)}(first(vis).value)
    return vis, ci
end

# Variable bounds
function MOI.supports_constraint(
    optimizer::Optimizer, ::Type{MOI.SingleVariable}, ::Type{<:BoundsSet})
    return true
end
function MOI.add_constraint(
    optimizer::Optimizer, f::MOI.SingleVariable, set::BoundsSet)
    vi = f.variable
    flag = MOIU.single_variable_flag(typeof(set))
    mask = optimizer.single_variable_mask[vi.value]
    MOIU.throw_if_lower_bound_set(vi, typeof(set), mask, Float64)
    MOIU.throw_if_upper_bound_set(vi, typeof(set), mask, Float64)
    info = optimizer.variable_info[vi.value]
    if !iszero(flag & MOIU.LOWER_BOUND_MASK)
        if info.variable_type == NNEG
            L = optimizer.nneg_L
        else
            L = optimizer.psdc_L[info.cone_index]
        end
        L[info.index_in_cone] = MOIU.extract_lower_bound(set)
    end
    if !iszero(flag & MOIU.UPPER_BOUND_MASK)
        if info.variable_type == NNEG
            U = optimizer.nneg_U
        else
            U = optimizer.psdc_U[info.cone_index]
        end
        U[info.index_in_cone] = MOIU.extract_upper_bound(set)
    end
    optimizer.single_variable_mask[vi.value] = mask | flag
    return MOI.ConstraintIndex{MOI.SingleVariable, typeof(set)}(vi.value)
end

# Objective
function MOI.get(optimizer::Optimizer, ::MOI.ObjectiveSense)
    return optimizer.objective_sense
end
sense_to_sign(sense::MOI.OptimizationSense) = sense == MOI.MAX_SENSE ? -1 : 1
function MOI.set(optimizer::Optimizer, ::MOI.ObjectiveSense, sense::MOI.OptimizationSense)
    if sense != optimizer.objective_sense
        sign = sense_to_sign(sense)
        rmul!(optimizer.nneg_Cval, -1)
        for i in eachindex(optimizer.psdc_dims)
            rmul!(optimizer.psdc_Cval[i], -1)
        end
    end
    optimizer.objective_sense = sense
end
function MOI.set(optimizer::Optimizer, ::MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}},
                 func::MOI.ScalarAffineFunction{Float64})
    optimizer.objective_constant = MOI.constant(func)
    empty!(optimizer.nneg_Cvar)
    empty!(optimizer.nneg_Cval)
    for i in eachindex(optimizer.psdc_dims)
        empty!(optimizer.psdc_Cvar[i])
        empty!(optimizer.psdc_Cval[i])
    end
    sign = sense_to_sign(optimizer.objective_sense)
    for term in func.terms
        info = optimizer.variable_info[term.variable_index.value]
        if info.variable_type == NNEG
            push!(optimizer.nneg_Cvar, info.index_in_cone)
            push!(optimizer.nneg_Cval, sign * term.coefficient)
        else
            @assert info.variable_type == PSD
            push!(optimizer.psdc_Cvar[info.cone_index], info.index_in_cone)
            push!(optimizer.psdc_Cval[info.cone_index], sign * term.coefficient)
        end
    end
end

# Constraints
function is_diagonal_index(k)
    # See https://www.juliaopt.org/MathOptInterface.jl/v0.9.3/apireference/#MathOptInterface.AbstractSymmetricMatrixSetTriangle
    i = div(1 + isqrt(8k - 7), 2)
    j = k - div((i - 1) * i, 2)
    return i == j
end
function MOI.add_constraint(optimizer::Optimizer, func::MOI.ScalarAffineFunction{Float64}, set::MOI.EqualTo{Float64})
    if !iszero(MOI.constant(func))
        # We use the fact that the initial function constant was zero to
        # implement getters for `MOI.ConstraintPrimal`.
        throw(MOI.ScalarFunctionConstantNotZero{
             Float64, typeof(func), typeof(set)}(MOI.constant(func)))
    end
    push!(optimizer.b, MOI.constant(set))
    con = length(optimizer.b)
    for term in func.terms
        info = optimizer.variable_info[term.variable_index.value]
        if info.variable_type == NNEG
            push!(optimizer.nneg_Avar, info.index_in_cone)
            push!(optimizer.nneg_Acon, con)
            push!(optimizer.nneg_Aval, term.coefficient)
        else
            @assert info.variable_type == PSD
            push!(optimizer.psdc_Avar[info.cone_index], info.index_in_cone)
            push!(optimizer.psdc_Acon[info.cone_index], con)
            coef = is_diagonal_index(info.index_in_cone) ? term.coefficient : term.coefficient / √2
            push!(optimizer.psdc_Aval[info.cone_index], coef)
        end
    end
    return AFFEQ(con)
end
function MOI.add_constraint(optimizer::Optimizer, func::MOI.ScalarAffineFunction{Float64}, set::BoundsSet)
    if !iszero(MOI.constant(func))
        # We use the fact that the initial function constant was zero to
        # implement getters for `MOI.ConstraintPrimal`.
        throw(MOI.ScalarFunctionConstantNotZero{
             Float64, typeof(func), typeof(set)}(MOI.constant(func)))
    end
    flag = MOIU.single_variable_flag(typeof(set))
    push!(optimizer.l, iszero(flag & MOIU.LOWER_BOUND_MASK) ? -Inf : MOIU.extract_lower_bound(set))
    push!(optimizer.u, iszero(flag & MOIU.UPPER_BOUND_MASK) ? Inf : MOIU.extract_upper_bound(set))
    con = length(optimizer.l)
    for term in func.terms
        info = optimizer.variable_info[term.variable_index.value]
        if info.variable_type == NNEG
            push!(optimizer.nneg_Bvar, info.index_in_cone)
            push!(optimizer.nneg_Bcon, con)
            push!(optimizer.nneg_Bval, term.coefficient)
        else
            @assert info.variable_type == PSD
            push!(optimizer.psdc_Bvar[info.cone_index], info.index_in_cone)
            push!(optimizer.psdc_Bcon[info.cone_index], con)
            coef = is_diagonal_index(info.index_in_cone) ? term.coefficient : term.coefficient / √2
            push!(optimizer.psdc_Bval[info.cone_index], coef)
        end
    end
    return MOI.ConstraintIndex{typeof(func), typeof(set)}(con)
end

# TODO could do something more efficient here
#      `SparseMatrixCSC` is returned in SumOfSquares.jl test `sos_horn`
symvec(Q::SparseMatrixCSC) = symvec(Matrix(Q))
function symvec(Q::Matrix)
    n = LinearAlgebra.checksquare(Q)
    vec_dim = MOI.dimension(MOI.PositiveSemidefiniteConeTriangle(n))
    q = Vector{eltype(Q)}(undef, vec_dim)
    k = 0
    for j in 1:n
        for i in 1:j
            k += 1
            q[k] = Q[i, j]
        end
    end
    @assert k == length(q)
    return q
end
function symvec(Q::Vector)
    if length(Q) > 1
        error("Expected square matrix from SDPNAL but got a vector of length $(length(Q))")
    end
    return Q
end
function _vec(x::SparseMatrixCSC)
    if size(x, 2) > 1
        error("Expected a vector from SDPNAL but got a sparse matrix of size $(size(x))")
    end
    return Vector(x[:, 1])
end
_vec(x::Float64) = [x]
_vec(x::Vector) = x

function MOI.optimize!(optimizer::Optimizer)
    options = optimizer.options
    if optimizer.silent
        options = copy(options)
        options[:printlevel] = 0
    end

    blk1 = Any[]
    blk2 = Any[]
    m = length(optimizer.b)
    p = length(optimizer.l)
    At = SparseArrays.SparseMatrixCSC{Float64,Int}[]
    Bt = SparseArrays.SparseMatrixCSC{Float64,Int}[]
    # FIXME I get a strange failure with sparse vectors, need to investigate
    C = Union{Matrix{Float64}, Vector{Float64}}[]
    L = Vector{Union{Matrix{Float64}, Vector{Float64}}}(undef, 0)
    U = Vector{Union{Matrix{Float64}, Vector{Float64}}}(undef, 0)
    if !isempty(optimizer.psdc_dims)
        for (i, dim) in enumerate(optimizer.psdc_dims)
            vec_dim = MOI.dimension(MOI.PositiveSemidefiniteConeTriangle(dim))
            push!(At, sparse(optimizer.psdc_Avar[i], optimizer.psdc_Acon[i], optimizer.psdc_Aval[i], vec_dim, m))
            push!(Bt, sparse(optimizer.psdc_Bvar[i], optimizer.psdc_Bcon[i], optimizer.psdc_Bval[i], vec_dim, p))
            c = Vector(sparsevec(optimizer.psdc_Cvar[i], optimizer.psdc_Cval[i], vec_dim))
            Ci = zeros(dim, dim)
            Li_vec = optimizer.psdc_L[i]
            Li = zeros(dim, dim)
            Ui_vec = optimizer.psdc_U[i]
            Ui = zeros(dim, dim)
            k = 0
            for col in 1:dim
                for row in 1:(col - 1)
                    k += 1
                    Ci[row, col] = c[k] / 2
                    Ci[col, row] = c[k] / 2
                    Li[row, col] = Li_vec[k]
                    Li[col, row] = Li_vec[k]
                    Ui[row, col] = Ui_vec[k]
                    Ui[col, row] = Ui_vec[k]
                end
                k += 1
                Ci[col, col] = c[k]
                Li[col, col] = Li_vec[k]
                Ui[col, col] = Ui_vec[k]
            end
            push!(C, Ci)
            push!(L, Li)
            push!(U, Ui)
            push!(blk1, "s")
            push!(blk2, Float64(dim))
        end
    end
    if !iszero(optimizer.num_nneg)
        push!(At, sparse(optimizer.nneg_Avar, optimizer.nneg_Acon, optimizer.nneg_Aval, optimizer.num_nneg, m))
        push!(Bt, sparse(optimizer.nneg_Bvar, optimizer.nneg_Bcon, optimizer.nneg_Bval, optimizer.num_nneg, p))
        push!(C, Vector(sparsevec(optimizer.nneg_Cvar, optimizer.nneg_Cval, optimizer.num_nneg)))
        push!(blk1, "l")
        push!(blk2, Float64(optimizer.num_nneg))
        push!(L, optimizer.nneg_L)
        push!(U, optimizer.nneg_U)
    end
    blk = [blk1 blk2]

    options = optimizer.options
    if optimizer.silent
        options = copy(options)
        options[:printlevel] = 0
    end

    obj, X, s, y, Z1, Z2, y2, v, optimizer.info, optimizer.runhist = sdpnalplus(
        blk, At, C, optimizer.b, L, U, Bt, optimizer.l, optimizer.u; options...)

    optimizer.primal_objective_value, optimizer.dual_objective_value = obj
    k = 0
    optimizer.psdc_X = symvec.(X[k .+ eachindex(optimizer.psdc_dims)])
    optimizer.psdc_Z1 = symvec.(Z1[k .+ eachindex(optimizer.psdc_dims)])
    optimizer.psdc_Z2 = symvec.(Z2[k .+ eachindex(optimizer.psdc_dims)])
    k += length(optimizer.psdc_dims)
    if iszero(optimizer.num_nneg)
        empty!(optimizer.nneg_X)
        empty!(optimizer.nneg_Z1)
        empty!(optimizer.nneg_Z2)
    else
        k += 1
        optimizer.nneg_X = X[k]
        optimizer.nneg_Z1 = Z1[k]
        optimizer.nneg_Z2 = _vec(Z2[k])
    end
    optimizer.s = s
    optimizer.y = y
    optimizer.y2 = _vec(y2) # Should we used `v` instead ? They should be approximately equal
    optimizer.status = optimizer.info["termcode"]
    optimizer.solve_time = optimizer.info["totaltime"]
end

function MOI.get(optimizer::Optimizer, ::MOI.RawStatusString)
    if optimizer.status === nothing
        throw(MOI.OptimizeNotCalled())
    else
        # See `solver_main_default/sdpnalplus.m`
        if optimizer.status == 0
            return "SDP is solved to the required accuracy"
        elseif optimizer.status == 1
            return "SDP is not solved successfully"
        else
            if !(-1 <= optimizer.status <= -3)
                error("Unrecognized SDPNAL termcode `$(optimizer.status)`.")
            end
            stoptol = get(optimizer.options, "stoptol", 1e-6)
            prim = optimizer.status == -2 ? 1.05stoptol : 10stoptol
            dual = optimizer.status == -1 ? 1.05stoptol : (optimizer.status == -3 ? 10stoptol : stoptol)
            return "SDP is partially solved successfully: primfeasorg < $prim, dualfeasorg < $dual"
        end
    end
end

const TERMINATION_STATUS = [
    MOI.NUMERICAL_ERROR,
    MOI.NUMERICAL_ERROR,
    MOI.NUMERICAL_ERROR,
    MOI.OPTIMAL,
    MOI.NUMERICAL_ERROR
]
function MOI.get(optimizer::Optimizer, ::MOI.TerminationStatus)
    if optimizer.status === nothing
        return MOI.OPTIMIZE_NOT_CALLED
    else
        return TERMINATION_STATUS[4 + optimizer.status]
    end
end

const PRIMAL_STATUS = [
    MOI.NEARLY_FEASIBLE_POINT,
    MOI.NEARLY_FEASIBLE_POINT,
    MOI.NEARLY_FEASIBLE_POINT,
    MOI.FEASIBLE_POINT,
    MOI.NO_SOLUTION
]
function MOI.get(optimizer::Optimizer, attr::MOI.PrimalStatus)
    if attr.N > MOI.get(optimizer, MOI.ResultCount())
        return MOI.NO_SOLUTION
    end
    return PRIMAL_STATUS[4 + optimizer.status]
end

function MOI.get(optimizer::Optimizer, attr::MOI.DualStatus)
    if attr.N > MOI.get(optimizer, MOI.ResultCount())
        return MOI.NO_SOLUTION
    end
    return PRIMAL_STATUS[4 + optimizer.status]
end

MOI.get(::Optimizer, ::MOI.ResultCount) = 1
function MOI.get(optimizer::Optimizer, attr::MOI.ObjectiveValue)
    MOI.check_result_index_bounds(optimizer, attr)
    sign = sense_to_sign(optimizer.objective_sense)
    return sign * optimizer.primal_objective_value + optimizer.objective_constant
end
function MOI.get(optimizer::Optimizer, attr::MOI.DualObjectiveValue)
    MOI.check_result_index_bounds(optimizer, attr)
    sign = sense_to_sign(optimizer.objective_sense)
    return sign * optimizer.dual_objective_value + optimizer.objective_constant
end

function MOI.get(optimizer::Optimizer, attr::MOI.VariablePrimal, vi::MOI.VariableIndex)
    MOI.check_result_index_bounds(optimizer, attr)
    info = optimizer.variable_info[vi.value]
    if info.variable_type == NNEG
        return optimizer.nneg_X[info.index_in_cone]
    else
        @assert info.variable_type == PSD
        return optimizer.psdc_X[info.cone_index][info.index_in_cone]
    end
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintPrimal,
                 ci::MOI.ConstraintIndex{MOI.SingleVariable, <:BoundsSet})
    MOI.check_result_index_bounds(optimizer, attr)
    return MOI.get(optimizer, MOI.VariablePrimal(attr.N),
                   MOI.VariableIndex(ci.value))
end

function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintPrimal,
                 ci::MOI.ConstraintIndex{MOI.VectorOfVariables, S}) where S<:SupportedSets
    MOI.check_result_index_bounds(optimizer, attr)
    info = optimizer.variable_info[ci.value]
    if info.variable_type == NNEG
        dim = optimizer.nneg_info[info.index_in_cone]
        return optimizer.nneg_X[(info.index_in_cone - 1) .+ (1:dim)]
    else
        @assert info.variable_type == PSD
        return optimizer.psdc_X[info.cone_index]
    end
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintPrimal, ci::AFFEQ)
    MOI.check_result_index_bounds(optimizer, attr)
    return optimizer.b[ci.value]
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintPrimal,
                 ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64},
                                         <:BoundsSet})
    MOI.check_result_index_bounds(optimizer, attr)
    return optimizer.s[ci.value]
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintDual,
                 ci::MOI.ConstraintIndex{MOI.VectorOfVariables, <:SupportedSets})
    MOI.check_result_index_bounds(optimizer, attr)
    info = optimizer.variable_info[ci.value]
    if info.variable_type == NNEG
        dim = optimizer.nneg_info[info.index_in_cone]
        return optimizer.nneg_Z1[(info.index_in_cone - 1) .+ (1:dim)]
    else
        @assert info.variable_type == PSD
        return optimizer.psdc_Z1[info.cone_index]
    end
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintDual,
                 ci::MOI.ConstraintIndex{MOI.SingleVariable, <:BoundsSet})
    MOI.check_result_index_bounds(optimizer, attr)
    info = optimizer.variable_info[ci.value]
    if info.variable_type == NNEG
        Z2 = optimizer.nneg_Z2
    else
        @assert info.variable_type == PSD
        Z2 = optimizer.psdc_Z2[info.cone_index]
    end
    return Z2[info.index_in_cone]
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintDual, ci::AFFEQ)
    MOI.check_result_index_bounds(optimizer, attr)
    return optimizer.y[ci.value]
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintDual,
                 ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64},
                                         <:BoundsSet})
    MOI.check_result_index_bounds(optimizer, attr)
    return optimizer.y2[ci.value]
end


text:
module IncrementalInference

blas_set_num_threads(2)

using
  Graphs,
  GraphViz,
  Gadfly,
  Colors,
  NLsolve,
  Distributions,
  KernelDensityEstimate,
  TransformUtils#,
  #CloudGraphs

export
  # actual CloudGraphs integration experimental code
  setCloudDataLayerAPI!,

  # using either dictionary or cloudgraphs
  VariableNodeData,
  PackedVariableNodeData,
  VNDencoder,
  VNDdecoder,
  FNDencode,
  FNDdecode,
  FunctionNodeData,
  FactorGraph,
  addNode!,
  addFactor!,
  resetData!,
  getData,
  getVarNode,
  getVal,
  setVal!,
  getBWVal,
  setBW!,
  setValKDE!,
  updateVertData!,
  BayesTree,
  EasyMessage,
  NBPMessage,
  ExploreTreeType,
  emptyFactorGraph,
  subgraphShortestPath,
  subgraphFromVerts,
  getEliminationOrder,
  prepBatchTree!,
  wipeBuildNewTree!,
  whichCliq,
  getKDE,
  getVertKDE,

  #Visualization
  investigateMultidimKDE,
  writeGraphPdf,
  ls,
  drawHorDens,
  drawHorBeliefsList,
  # vstackedDensities, # global scope naming issue with msgPlots

  # Tree stuff
  spyCliqMat,
  evalPotential,
  evalFactor2,

  # solve inference
  inferOverTree!,
  inferOverTreeR!,
    #development interface
    upMsgPassingRecursive,

  GenericMarginal,
  #Robot stuff
  PriorPose2,
  PackedPriorPose2,
  Pose2Pose2,
  PackedPose2Pose2,
  addPose2Pose2,
  Pose2DPoint2DBearingRange,
  PackedPose2DPoint2DBearingRange,
  solveLandm,
  solvePose2,
  solveSetSeps,
  addPose2Pose2!,
  uppA,
  convert, # for magic protobuf stuff
  compare,

  # Going to move to RoME.jl in future
  # For 1D example
  Odo,
  odoAdd,
  PackedOdo,
  Obsv2,
  PackedObsv2,
  Ranged,

  PriorPose3,
  Pose3Pose3,
  projectParticles,

  # dev exports
  addGraphsVert!

include("FactorGraphTypes.jl")
# include("CloudGraphIntegration.jl") # Work in progress code
include("DataLayerAPI.jl")
include("FactorGraph01.jl")
include("JunctionTree.jl")
include("GraphConstraintTypes.jl")
include("TreePotentials01.jl")
include("TreePotentials02.jl")
include("TreePotentials03.jl")
include("SolveTree01.jl")
include("SolverVisualization.jl")

end


text:
export CGPFunctions

module CGPFunctions

global arity = Dict()

SorX = Union{Symbol, Expr}

function scaled(x::Float64)
    if isnan(x)
        return 0.0
    end
    min(max(x, -1.0), 1.0)
end

function fgen(name::Symbol, ar::Int, s1::SorX; safe::Bool=false)
    if safe
        @eval function $name(x::Float64, y::Float64)::Float64
            try
                return $s1
            catch
                return x
            end
        end
    else
        @eval function $name(x::Float64, y::Float64)::Float64
            $s1
        end
    end
    arity[String(name)] = ar
end

# Mathematical
fgen(:f_add, 2, :((x + y) / 2.0))
fgen(:f_subtract, 2, :(abs(x - y) / 2.0))
fgen(:f_mult, 2, :(x * y))
fgen(:f_div, 2, :(scaled(x / y)))
fgen(:f_abs, 1, :(abs(x)))
fgen(:f_sqrt, 1, :(sqrt(abs(x))))
fgen(:f_pow, 2, :(abs(x) ^ abs(y)))
fgen(:f_exp, 1, :((2 * (exp(x+1)-1.0))/(exp(2.0)-1.0) -1))
fgen(:f_sin, 1, :(sin(x)))
fgen(:f_cos, 1, :(cos(x)))
fgen(:f_tanh, 1, :(tanh(x)))
fgen(:f_sqrt_xy, 2, :(sqrt(x^2 + y^2) / sqrt(2.0)))
fgen(:f_lt, 2, :(Float64(x < y)))
fgen(:f_gt, 2, :(Float64(x > y)))

# Logical
fgen(:f_and, 2, :(Float64((&)(Int(round(x)), Int(round(y))))))
fgen(:f_or, 2, :(Float64((|)(Int(round(x)), Int(round(y))))))
fgen(:f_xor, 2, :(Float64(xor(Int(abs(round(x))), Int(abs(round(y)))))))
fgen(:f_not, 1, :(1 - abs(round(x))))

end


text: if true 1 elseif true 2 end

max_stars_repo_name: gabriel-doriath-dohler/pjulia

max_issues_repo_name: gabriel-doriath-dohler/pjulia

max_forks_repo_name: gabriel-doriath-dohler/pjulia