text:
module GeneralPlotUtils
using Reproduce
using FileIO
using Statistics
using ProgressMeter
using JLD2
function load_results(ic, logger_key; return_type = "tensor")
    num_results = length(ic)
    results = []
    for itm in ic.items
        if !isfile(joinpath(itm.folder_str, "results.jld2"))
            println(joinpath(itm.folder_str, "results.jld2"))
            @warn "Check your sweep. It began run but did not save final results!"
        else
            data = FileIO.load(joinpath(itm.folder_str, "results.jld2"))["results"]
            push!(results,data[logger_key])
        end
    end

    if return_type == "tensor"
        if isempty(results)
            @warn "Empty IC. Check your sweep!"
            return []
        end
        return cat(results..., dims = 3)
    elseif return_type == "array"
        return results
    end
end

function combine(dict_arr, key, vals)
    new_dict_arr = []
    for dict in dict_arr
        for val in vals
            new_dict = deepcopy(Dict(dict))
            new_dict[key] = val
            push!(new_dict_arr, new_dict)
        end
    end
    return new_dict_arr
end

function split_algo(algo_ic, swept_params)
    dicts = [Dict()]
    for k in swept_params
        if k in keys(diff(algo_ic))
            list_of_params = diff(algo_ic)[k]
            dicts = combine(dicts, k, list_of_params)
        end
    end
    return dicts
end

function get_best(ic, sweep_params, metric)
    splits = split_algo(ic,sweep_params)
    errors = ones(length(splits)) * Inf
    for (ind, split) in enumerate(splits)
        res = load_results(search(ic, split), metric)
        error = mean(res)
        errors[ind] = error
    end
    low_err, low_err_ind = findmin(errors)
    return search(ic, splits[low_err_ind])
end
function get_most_episodes(ic, sweep_params, metric)
    splits = split_algo(ic,sweep_params)
    lengths = -ones(length(splits)) * Inf
    for (ind, split) in enumerate(splits)
        candidate_best = search(ic, split)
        if length(candidate_best.items) == 0
            @warn "$(split) results in not a valid combination"
            continue
        end
        res = load_results(candidate_best, metric, return_type = "array")
        @show size(res)
        num_episodes = length.(res)
        lengths[ind] = mean(num_episodes)
    end
    mx,mx_ind = findmax(lengths)
    @show lengths
    return search(ic, splits[mx_ind])

end
function get_best_final_perf(ic, sweep_params, metric, cut_per)
    splits = split_algo(ic,sweep_params)
    errors = ones(length(splits)) * Inf
    for (ind, split) in enumerate(splits)
        candidate_best = search(ic, split)
        if length(candidate_best.items) == 0
            @warn "$(split) results in not a valid combination"
            continue
        end
        res = load_results(candidate_best, metric)
        if isempty(res)
            @warn "Selecting over empty IC. This is not good and is likely not what you want"
            continue
        end
        num_steps = size(res)[2]
        cut_ind = Int(floor(num_steps * (1-cut_per)))
        error = mean(res[:,cut_ind:end,:])
        if isfinite(error)
            errors[ind] = error
        end
    end
    low_err, low_err_ind = findmin(errors)
    return search(ic, splits[low_err_ind])
end

function get_stats(data;per_gvf=false)
    mean_per_gvf, std_per_gvf = mean(data,dims=3)[:,:,1], std(data,dims=3)[:,:,1]
    return a,b = if per_gvf == true
        mean_per_gvf, std_per_gvf
    else

        # error_tot = sum(data,dims=1)[1,:,:]
        # sum(mean_per_gvf,dims=1)[1,:], sum(std_per_gvf,dims=1)[1,:]
        # @show size(std(error_tot,dims=2))
        # @show size(std(error_tot,dims=1))
        # sum(mean_per_gvf,dims=1)[1,:], std(error_tot,dims=1)[:,1]

        # sum(mean_per_gvf,dims=1)[1,:], std(error_tot,dims=1)[:,1]
        # @show size(std(data, dims=[1,3]))
        vec(sum(mean(data,dims=3)[:,:,1], dims = 1)), vec( std(sum(data, dims=1),dims=3) )
    end
end

function get_min_length(arrs)
    return minimum([length(arr) for arr in arrs])
end

function goal_visits_per_episode(arr_of_episodes, max_length; num_gvfs = 4)
    goal_visits = zeros(num_gvfs, max_length)
    for run in 1:length(arr_of_episodes)
        goal_visits += onehot(arr_of_episodes[run],num_gvfs)[:,1:max_length]
    end
    goal_visit_percentage = goal_visits / length(arr_of_episodes)
    return goal_visit_percentage
end

function onehot(data, num_gvfs = 4)
    onehot_enc = zeros(num_gvfs, length(data))
    for gvf_i in 1:num_gvfs
        onehot_enc[gvf_i,:] += (data .== gvf_i)
    end
    return onehot_enc
end

function get_label(ic, params)
    label = ""
    for p in params
        if p in keys(ic[1].parsed_args)
            label = string(label, " ", ic[1].parsed_args[p])
        end
    end

    return label
end

function print_params(ic, algo_params, sweep_params)
    println("For IC: ")
    for p in algo_params
        if p in keys(ic[1].parsed_args)
            println(p,"  ", ic[1].parsed_args[p])
        end
    end
    for p in sweep_params
        if p in keys(ic[1].parsed_args)
            println(p,"  ", ic[1].parsed_args[p])
        end
    end
    println()
end

function smooth(data, k)
    smoothed_data = zeros(size(data))
    for i = 1:size(data, 1)
        if i < k
            smoothed_data[i] = mean(data[1:i])
        else
            smoothed_data[i] = mean(data[i - k + 1:i])
        end
    end
    return smoothed_data
end

end

module LabelUtils
using Reproduce
using FileIO
using Statistics
using ProgressMeter
using JLD2
using Plots
using LaTeXStrings

# Demon Algorithm results in different line styles!!
const val_matches = Dict(["GPI","TB","Q","TB"] => 6,
    ["GPI", "TB", "SR", "TB"] => 6,
    ["GPI", "TB", "SR", "ETB"] => 2,
    ["GPI", "TB", "SR","StateETB"] => 7,
    ["GPI", "TB", "SR", "InterestTB"] => 4,
    ["GPI", "TB", "Q", "InterestTB"] => 4,
    ["GPI", "TB", "Q", "ETB"] => 2,
    ["GPI", "TB", "Q", "TB"] => 6,
    ["Q","ESARSA", "Q", "TB"] => 1,
    ["Q", "ESARSA", "SR", "TB"] => 1,
    ["Q", "TabularRoundRobin", "SR", "TB"] => 6,
    ["RoundRobin", "TB", "SR", "TB"] => 6,
    ["RoundRobin", "Q", "SR", "TB"] => 6,
    ["RoundRobin", "TB", "Q", "TB"] => 1,
    ["RoundRobin", "Q", "Q", "TB"] => 1,
    ["RoundRobin", "Q","LSTD","TB"] => 3,
    ["Q", "TabularRoundRobin", "Q", "TB"] => 1,
    ["Q", "TabularRoundRobin", "Q", "ESARSA"] => 2,
    ["Q", "TabularRoundRobin", "SR", "ESARSA"] => 6,
    ["Q", "TabularRoundRobin","LSTD","TB"] => 3,
    ["RandomDemons", "TB","Q","TB"] => 5,
    ["RandomDemons", "TB","SR","TB"] => 5,

    )

const algo_labels = Dict(["GPI","TB","Q","TB"] => L"μ(GPI), π(TB)",
    ["Q","ESARSA", "Q", "TB"] => L"μ(Sarsa), π(TB)",
    ["GPI", "TB", "SR", "TB"] => L"μ(GPI), π(SR)",
    ["GPI", "TB", "SR", "ETB"] => L"μ(GPI), π(SR with ETB)",
    ["GPI", "TB", "SR", "InterestTB"] => L"μ(GPI), π(SR with Interest)",
    ["GPI", "TB", "SR","StateETB"] =>  L"μ(GPI), π(SR with State ETB)",
    ["GPI", "TB", "Q", "InterestTB"] => L"μ(GPI), π(TB with Interest)",
    ["GPI", "TB", "Q", "ETB"] => L"μ(GPI), π(ETB)",
    ["GPI", "TB", "Q", "TB"] => L"μ(GPI), π(TB)",
    ["Q", "ESARSA", "SR", "TB"] => L"μ(Sarsa), π(SR)",
    ["Q", "TabularRoundRobin", "SR", "TB"] => L"μ(Fixed), π(SR)",
    ["Q", "TabularRoundRobin", "Q", "TB"] => L"μ(Fixed), π(TB)",
    ["RoundRobin", "TB", "SR", "TB"] => "μ(Fixed), π(SR)",
    ["RoundRobin", "Q", "SR", "TB"] => "μ(Fixed), π(SR)",
    ["RoundRobin", "TB", "Q", "TB"] => "μ(Fixed), π(TB)",
    ["RoundRobin", "Q", "Q", "TB"] => "μ(Fixed), π(TB)",
    ["RoundRobin", "Q","LSTD","TB"] => "μ(Fixed), π(LSTD)",
    ["Q", "TabularRoundRobin","LSTD","TB"] => "μ(Fixed), π(LSTD)",
    ["Q", "TabularRoundRobin", "Q", "ESARSA"] => "μ(Fixed), π(TD)",
    ["Q", "TabularRoundRobin", "SR", "ESARSA"] => "μ(Fixed), π(SR + TD)",
    ["RandomDemons", "TB","Q","TB"] => "μ(Random), π(TB)",
    ["RandomDemons", "TB","SR","TB"] => "μ(Random), π(SR)",
    )

const algo_keys = ["behaviour_learner", "behaviour_update", "demon_learner", "demon_update"]
# TOL Muted Colour Scheme
# const color_scheme = [
#     colorant"#44AA99",
#     colorant"#332288",
#     colorant"#DDCC77",
    # colorant"#999933",
#     colorant"#CC6677",
#     colorant"#AA4499",
#     colorant"#DDDDDD",
# 	colorant"#117733",
# 	colorant"#882255",
# 	colorant"#88CCEE",
#     ]

const color_scheme = [
colorant"#0077BB",
colorant"#33BBEE",
colorant"#009988",
colorant"#EE7733",
colorant"#CC3311",
colorant"#EE3377",
colorant"#BBBBBB",
]

function _is_match(ic, keys, vals)
    tst = ([ic[1].parsed_args[keys[i]] == vals[i] for i in 1:length(keys)])
    return all(tst)
end
function _print_algo_keys(ic;algo_keys = algo_keys)
    println()
    [println(ic[1].parsed_args[k]) for k in algo_keys]
end
function _get_ic_ind(ic)
    ic_diff_keys = keys(diff(ic))
    for k in algo_keys
        if k in ic_diff_keys
            @show diff(ic)
            throw(ArgumentError("Multiple types of  $(k) in item collection. Unspecified what colour palette to use"))
        end
    end

    for algo_vals in keys(val_matches)
        if _is_match(ic, algo_keys, algo_vals)
            return val_matches[algo_vals]
        end
    end

    _print_algo_keys(ic)
    @warn string("No match found for: ", diff(ic))
end

function get_colour(ic)
    # Unique Colour given on a per {behaviour learner, behaviour update, demon learner, demon update} basis.
    ind = _get_ic_ind(ic)
    return Dict(:color => color_scheme[ind])
end

function get_linestyle(ic)
    #If we are in round robin  use solid linestyle
    if ic[1].parsed_args["behaviour_update"] == "TabularRoundRobin" || ic[1].parsed_args["behaviour_learner"] == "RoundRobin"
        return Dict(:linestyle => :solid)
    elseif ic[1].parsed_args["demon_learner"] == "SR"
        return Dict(:linestyle => :solid)
    elseif ic[1].parsed_args["demon_learner"] == "Q"
        return Dict(:linestyle => :dash)
    end
end

function get_label(ic)
    for algo_vals in keys(algo_labels)
        if _is_match(ic, algo_keys, algo_vals)
            return Dict(:label => algo_labels[algo_vals])
        end
    end
    throw(ArgumentError("No Label description!!"))
end

function get_params(ic)
    if ic.items[1].parsed_args["exploration_param"] == 1.0
        label = Dict(:label => "μ(Random), π(SR)")
        color = Dict(:color => colorant"#BBBBBB")
        ls = get_linestyle(ic)
        return  merge(color,ls,label)
    end
    color = get_colour(ic)
    ls = get_linestyle(ic)
    label = get_label(ic)
    # label[:label] = string(label[:label], " ", ic.items[1].parsed_args["exploration_param"])

    #Hacking thing to do for ER
    if "batch_size" in keys(ic.items[1].parsed_args)
        label[:label] =  string(label[:label], " ", ic.items[1].parsed_args["batch_size"])
        ls[:linestyle] = :dot
    end
    return merge(color,ls,label)
end

end

max_stars_repo_name: MatthewMcLeod/curiosity

max_issues_repo_name: MatthewMcLeod/curiosity

max_forks_repo_name: MatthewMcLeod/curiosity


text:
# keyboard & mouse callbacks for controls
const mov_forward = Ref(false)
const mov_left = Ref(false)
const mov_right = Ref(false)
const mov_back = Ref(false)
const mov_up = Ref(false)
const mov_down = Ref(false)
const mov_boost = Ref(false)

function key_cb_fun(window, key, _, action)
    if key == MiniFB.KB_KEY_W # up
        mov_forward[] = action
    elseif key == MiniFB.KB_KEY_A # left
        mov_left[] = action
    elseif key == MiniFB.KB_KEY_D # right
        mov_right[] = action
    elseif key == MiniFB.KB_KEY_S # down
        mov_back[] = action
    elseif key == MiniFB.KB_KEY_SPACE # space
        mov_up[] = action
    elseif key == MiniFB.KB_KEY_LEFT_CONTROL # l-ctrl
        mov_down[] = action
    elseif key == MiniFB.KB_KEY_LEFT_SHIFT # l-shift
        mov_boost[] = action
    elseif key == MiniFB.KB_KEY_ESCAPE
        mfb_close(window)
    end
    return nothing
end

const mouse_pos = Int32[0,0]
const mouse_mov = Int32[0,0]

function mouse_cb_fun(window, x, y)
    mouse_mov .= (x,y) .- mouse_pos
    mouse_pos .= (x,y)
    return nothing
end

function compute_movement!(θ)
    # mouse
    dθ = -mouse_mov[1]*1f-2
    fill!(mouse_mov, zero(Int32))
    # keyboard; signs empirically adjusted...
    sθ, cθ = sincos(θ)
    dx = (mov_right[]-mov_left[])*cθ - (mov_forward[]-mov_back[])*sθ
    dy = -(mov_right[]-mov_left[])*sθ - (mov_forward[]-mov_back[])*cθ
    dz = (mov_up[]-mov_down[])
    n = dx>0 ? 1f0/hypot(dx,dy) : 1f0
    return dx*n*(1 + mov_boost[]), dy*n*(1 + mov_boost[]), dz, dθ
end

ext: jl

max_stars_repo_name: jebej/VoxelSpace

max_issues_repo_name: jebej/VoxelSpace

max_forks_repo_name: jebej/VoxelSpace


text:
using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:5] <= 1)

                     @variable(m, -9.734460520574642 <= q <= 2.116911729869532)

                     add_NL_constraint(m, :(softplus(-0.8046949419530325 + -0.2278494590955038*softplus(0.1770285029406451 + -0.5067338360802731*softplus(-0.7998422425711689 + 0.5883237869246916*softplus(-0.6653700708895021 + 0.6185402094345611*$(x[1]) + 0.05069108137655398*$(x[2]) + 0.9466113725109135*$(x[3]) + 0.06805104680691931*$(x[4]) + -0.30203159362192356*$(x[5])) + 0.47235910506966317*softplus(-0.04665405044600712 + -0.8840756438370558*$(x[1]) + -0.20571795637475754*$(x[2]) + 0.8733859112268667*$(x[3]) + -0.5065045210605232*$(x[4]) + -0.7647800565294718*$(x[5]))) + -0.09889031570146667*softplus(-0.7622628423532047 + 0.6518912481623378*softplus(-0.6653700708895021 + 0.6185402094345611*$(x[1]) + 0.05069108137655398*$(x[2]) + 0.9466113725109135*$(x[3]) + 0.06805104680691931*$(x[4]) + -0.30203159362192356*$(x[5])) + 0.4006903724184485*softplus(-0.04665405044600712 + -0.8840756438370558*$(x[1]) + -0.20571795637475754*$(x[2]) + 0.8733859112268667*$(x[3]) + -0.5065045210605232*$(x[4]) + -0.7647800565294718*$(x[5])))) + -0.18111143847071443*softplus(-0.12446319977672715 + 0.7018706953017659*softplus(-0.7998422425711689 + 0.5883237869246916*softplus(-0.6653700708895021 + 0.6185402094345611*$(x[1]) + 0.05069108137655398*$(x[2]) + 0.9466113725109135*$(x[3]) + 0.06805104680691931*$(x[4]) + -0.30203159362192356*$(x[5])) + 0.47235910506966317*softplus(-0.04665405044600712 + -0.8840756438370558*$(x[1]) + -0.20571795637475754*$(x[2]) + 0.8733859112268667*$(x[3]) + -0.5065045210605232*$(x[4]) + -0.7647800565294718*$(x[5]))) + 0.5275006439747312*softplus(-0.7622628423532047 + 0.6518912481623378*softplus(-0.6653700708895021 + 0.6185402094345611*$(x[1]) + 0.05069108137655398*$(x[2]) + 0.9466113725109135*$(x[3]) + 0.06805104680691931*$(x[4]) + -0.30203159362192356*$(x[5])) + 0.4006903724184485*softplus(-0.04665405044600712 + -0.8840756438370558*$(x[1]) + -0.20571795637475754*$(x[2]) + 0.8733859112268667*$(x[3]) + -0.5065045210605232*$(x[4]) + -0.7647800565294718*$(x[5]))))) + softplus(-0.4025778812055263 + 0.08853333939955377*softplus(0.1770285029406451 + -0.5067338360802731*softplus(-0.7998422425711689 + 0.5883237869246916*softplus(-0.6653700708895021 + 0.6185402094345611*$(x[1]) + 0.05069108137655398*$(x[2]) + 0.9466113725109135*$(x[3]) + 0.06805104680691931*$(x[4]) + -0.30203159362192356*$(x[5])) + 0.47235910506966317*softplus(-0.04665405044600712 + -0.8840756438370558*$(x[1]) + -0.20571795637475754*$(x[2]) + 0.8733859112268667*$(x[3]) + -0.5065045210605232*$(x[4]) + -0.7647800565294718*$(x[5]))) + -0.09889031570146667*softplus(-0.7622628423532047 + 0.6518912481623378*softplus(-0.6653700708895021 + 0.6185402094345611*$(x[1]) + 0.05069108137655398*$(x[2]) + 0.9466113725109135*$(x[3]) + 0.06805104680691931*$(x[4]) + -0.30203159362192356*$(x[5])) + 0.4006903724184485*softplus(-0.04665405044600712 + -0.8840756438370558*$(x[1]) + -0.20571795637475754*$(x[2]) + 0.8733859112268667*$(x[3]) + -0.5065045210605232*$(x[4]) + -0.7647800565294718*$(x[5])))) + -0.4269192713141794*softplus(-0.12446319977672715 + 0.7018706953017659*softplus(-0.7998422425711689 + 0.5883237869246916*softplus(-0.6653700708895021 + 0.6185402094345611*$(x[1]) + 0.05069108137655398*$(x[2]) + 0.9466113725109135*$(x[3]) + 0.06805104680691931*$(x[4]) + -0.30203159362192356*$(x[5])) + 0.47235910506966317*softplus(-0.04665405044600712 + -0.8840756438370558*$(x[1]) + -0.20571795637475754*$(x[2]) + 0.8733859112268667*$(x[3]) + -0.5065045210605232*$(x[4]) + -0.7647800565294718*$(x[5]))) + 0.5275006439747312*softplus(-0.7622628423532047 + 0.6518912481623378*softplus(-0.6653700708895021 + 0.6185402094345611*$(x[1]) + 0.05069108137655398*$(x[2]) + 0.9466113725109135*$(x[3]) + 0.06805104680691931*$(x[4]) + -0.30203159362192356*$(x[5])) + 0.4006903724184485*softplus(-0.04665405044600712 + -0.8840756438370558*$(x[1]) + -0.20571795637475754*$(x[2]) + 0.8733859112268667*$(x[3]) + -0.5065045210605232*$(x[4]) + -0.7647800565294718*$(x[5]))))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
using Distributed;

file = "../machines.csv"

machines = readlines(file)

for entry in machines
    local count, host

    v = split(entry, ",")

    host = v[1]

    if length(v) > 1
        count = parse(Int, v[2])
    else
        count = 1
    end

    println("host: $host, count: $count")

    if host == "localhost"
        println("adding $count local worker(s)")
        addprocs(count)
    else
        println("adding $count remote worker(s) on $host")
        addprocs([(host, count)])
    end
end

println(workers())


text:
# Open input
cd(raw"/Users/wapisani/Documents/Programming/Julia/Advent_of_Code/")
open("AoC_2015_3.input","r") do f
    global directions_list = read(f,String)
end

# directions_list = ">"
xy_coords = Set([(0,0)])
x = 0
y = 0

houses_at_least_1_present = 1
for direction in directions_list
    if direction == '^'
        y += 1
    elseif direction == 'v'
        y -= 1
    elseif direction == '<'
        x -= 1
    elseif direction == '>'
        x += 1
    end

    push!(xy_coords,(x,y))
end

houses_at_least_1_present = length(xy_coords)
println("Santa delivered at least 1 present to $houses_at_least_1_present houses ")

# Part 2
xy_coords_part2 = Set([(0,0)])
x_santa = 0
y_santa = 0
x_robo = 0
y_robo = 0

for (index,direction) in enumerate(directions_list)

    if index % 2 == 1 # Santa moves
        if direction == '^'
            y_santa += 1
        elseif direction == 'v'
            y_santa -= 1
        elseif direction == '<'
            x_santa -= 1
        elseif direction == '>'
            x_santa += 1
        end

        push!(xy_coords_part2,(x_santa,y_santa))

    elseif index % 2 == 0 # Robo-Santa moves
        if direction == '^'
            y_robo += 1
        elseif direction == 'v'
            y_robo -= 1
        elseif direction == '<'
            x_robo -= 1
        elseif direction == '>'
            x_robo += 1
        end

        push!(xy_coords_part2,(x_robo,y_robo))
    end
end

houses_part_2 = length(xy_coords_part2)

println("Santa and Robo-Santa delivered at least 1 present to $houses_part_2 houses ")


text:
# Adds functionality to deal with transforms from physical to lattice units.
# Mostly, defines two functions:
#
# - `dimensionless` goes from physical to lattice units
# - `physical` goes the other way
#
# Until www.github.com/Keno/SIUnits.jl#19 is resolved, we cannot use SIUnits.jl.
# However, a working implementation does exist in the tests and in the git history.
module Units
export LBunits, dimensionless, physical
# Holds all and everything about units that we know of
immutable LBUnits{T <: Number}
    δt::T
    δx::T
    δm::T
end
function LBUnits{T <: Number}(::Type{T}, δt::Number, δx::Number, δm::Number)
    LBUnits{T}(δt, δx, δm)
end

const DIMENSIONALITY = Dict{Symbol, NTuple{3, Int64}}(
    :time => (1, 0, 0), :length => (0, 1, 0), :weight => (0, 0, 1),
    :velocity => (-1, 1, 0), :pressure => (-2, -1, 1), :viscosity => (-1, -1, 1),
    :density => (0, -3, 1)
)

function dimensionless(lb::LBUnits, dimension::Symbol, x)
    if !haskey(DIMENSIONALITY, dimension)
        error("Don't know how to convert $dimensions to LB units")
    end
    const dt, dx, dm = DIMENSIONALITY[dimension]
    x * (
        if dt != 0 lb.δt^(-dt) else 1 end *
        if dx != 0 lb.δx^(-dx) else 1 end *
        if dm != 0 lb.δm^(-dm) else 1 end
    )
end
function physical(lb::LBUnits, dimension::Symbol, x)
    const dt, dx, dm = DIMENSIONALITY[dimension]
    x * (
        if dt != 0 lb.δt^dt else 1 end *
        if dx != 0 lb.δx^dx else 1 end *
        if dm != 0 lb.δm^dm else 1 end
    )
end
end


text:
# struct BarHTML end
# function Base.show(io::IO, ::MIME"juliasvscode/html", ::BarHTML)
#   print(io, "
Hello there
")
# end

# function Base.show(io::IO, m::MIME"juliavscode/html", p::Gadfly.Plot)
#     show(io, "text/html", p)
# end

# plot(z=(x, y) -> x * exp(-(x - round(Int, x))^2 - y^2), xmin=[-10], xmax=[10], ymin=[-2], ymax=[2], Geom.contour)

# BarHTML()

# println("ahhhy")

# println("hello")
# imshow(rand(Gray, 10, 10))
# println("world")


text:
## Electric field strength plotting

@recipe function f( ef::ElectricField{T, 3, Cylindrical};
                    r = missing,
                    φ = missing,
                    z = missing,
                    contours_equal_potential=false,
                    full_det = false ) where {T}

    grid::CylindricalGrid{T} = ef.grid
    cross_section::Symbol, idx::Int, value::T, units::Unitful.Units = get_crosssection_idx_and_value(grid, r, φ, z)

    title --> "Electric Field (Magn.) @ $(cross_section) = $(round(units, Float64(uconvert(units,value*(cross_section == :φ ? u"°" : internal_length_unit))), sigdigits=2))"
    colorbar_title --> "Electric Field Strength"
    seriescolor --> :inferno
    unitformat --> :slash

    ef_magn = norm.(ef)
    ElectricPotential(ef_magn, grid),  cross_section, idx, value, internal_efield_unit, contours_equal_potential, full_det
end

@recipe function f(ef::ElectricField{T, 3, Cartesian};
                    x = missing,
                    y = missing,
                    z = missing,
                    contours_equal_potential = false) where {T <: SSDFloat}

    grid::CartesianGrid3D{T} = ef.grid
    cross_section::Symbol, idx::Int, value::T, units::Unitful.Units = get_crosssection_idx_and_value(grid, x, y, z)

    title --> "Electric Field (Magn.) @ $(cross_section) = $(round(units, Float64(uconvert(units,value*internal_length_unit)), sigdigits=2))"
    colorbar_title --> "Electric Field Strength"
    seriescolor --> :inferno
    unitformat --> :slash

    ef_magn = norm.(ef)
    ElectricPotential(ef_magn, grid), cross_section, idx, value, internal_efield_unit, contours_equal_potential
end

## Electric field line plotting

function get_sample_lines(dim_symbol::Symbol, v::T, grid::CartesianGrid3D{T}, sampling::T)::Vector{ConstructiveSolidGeometry.Line} where {T}

    xrange = range(round.(endpoints(grid.x.interval)./sampling, (RoundDown, RoundUp)).*sampling..., step = sampling)
    yrange = range(round.(endpoints(grid.y.interval)./sampling, (RoundDown, RoundUp)).*sampling..., step = sampling)
    zrange = range(round.(endpoints(grid.z.interval)./sampling, (RoundDown, RoundUp)).*sampling..., step = sampling)

    return if dim_symbol == :x
        vcat(
            [ConstructiveSolidGeometry.Line(CartesianPoint{T}(v,y,0), CartesianVector{T}(0,0,1)) for y in yrange],
            [ConstructiveSolidGeometry.Line(CartesianPoint{T}(v,0,z), CartesianVector{T}(0,1,0)) for z in zrange]
        )
    elseif dim_symbol == :y
        vcat(
            [ConstructiveSolidGeometry.Line(CartesianPoint{T}(x,v,0), CartesianVector{T}(0,0,1)) for x in xrange],
            [ConstructiveSolidGeometry.Line(CartesianPoint{T}(0,v,z), CartesianVector{T}(1,0,0)) for z in zrange]
        )

    else # dim_symbol == :z
        vcat(
            [ConstructiveSolidGeometry.Line(CartesianPoint{T}(x,0,v), CartesianVector{T}(0,1,0)) for x in xrange],
            [ConstructiveSolidGeometry.Line(CartesianPoint{T}(0,y,v), CartesianVector{T}(1,0,0)) for y in yrange]
        )
    end
end

function get_sample_lines(dim_symbol::Symbol, v::T, grid::CylindricalGrid{T}, sampling::T)::Vector{ConstructiveSolidGeometry.Line} where {T}

    φsampling = 2π * sampling / (dim_symbol == :r ? v : grid.r.interval.right) # or use sampling together with the unit ?
    rrange = range(round.((-grid.r.interval.right,grid.r.interval.right)./sampling, (RoundDown, RoundUp)).*sampling..., step = sampling)
    φrange = range(round.(endpoints(grid.φ.interval)./φsampling, (RoundDown, RoundUp)).*φsampling..., step = φsampling)
    zrange = range(round.(endpoints(grid.z.interval)./sampling, (RoundDown, RoundUp)).*sampling..., step = sampling)

    return if dim_symbol == :φ
        vcat(
            [ConstructiveSolidGeometry.Line(CartesianPoint{T}(r*cos(v),r*sin(v),0), CartesianVector{T}(0,0,1)) for r in rrange],
            [ConstructiveSolidGeometry.Line(CartesianPoint{T}(0,0,z), CartesianVector{T}(cos(v),sin(v),0)) for z in zrange]
        )
    elseif dim_symbol == :z
        [ConstructiveSolidGeometry.Line(CartesianPoint{T}(0,0,v), CartesianVector{T}(cos(φ), sin(φ),0)) for φ in φrange]
    else # dim_symbol == :r
        [ConstructiveSolidGeometry.Line(CartesianPoint{T}(v*cos(φ),v*sin(φ),0), CartesianVector{T}(0,0,1)) for φ in φrange]
    end
end

@userplot Plot_electric_fieldlines
@recipe function f(gdd::Plot_electric_fieldlines; φ = missing, r = missing, x = missing, y = missing, z = missing,
                    max_nsteps = 5000,
                    sampling = 2u"mm", # Specifies in what density the Contacts are sampled to generate equally spaced surface charges. Also see spacing.
                    offset = 0.5u"mm", # should be at least as big as sampling. In doubt sampling can be reduced and the spacing keyword can be used to thin out the lines.
                    spacing = 1, # If, due to fine sampling, too many lines would clutter the plot, the spacing keyword allows to skip some fieldlines. Spacing = 2 means plot every second line. Spacing = 3 every third.
                    full_det = false,
                    skip_contact = 1) # Usually the "core" contact is skipped, and the other contacts are equally sampled for charges to drift
    sim = gdd.args[1]
    S = get_coordinate_system(sim.electric_field.grid)
    T = get_precision_type(sim.detector)
    dim_array::Vector{Union{Missing, RealQuantity}} = (S == Cartesian ? [x, y, z] : [r, φ, z])
    dim_symbols_array = (S == Cartesian ? [:x, :y, :z] : [:r, :φ, :z])

    if isempty(skipmissing(dim_array))
        if S == Cartesian
            dim_number = 1
            dim_symbol = :x
            v::T = sim.electric_field.grid[dim_number][ div(length(sim.electric_field.grid[dim_number]), 2) ]
            x = v
            units::Unitful.Units = u"m"
        else # S == Cylindrical
            v = 0
            φ = 0
            dim_number = 2
            dim_symbol = :φ
            units = u"°"
        end
        dim_array = (S == Cartesian ? [x, y, z] : [r, φ, z])
    elseif sum(ismissing.(dim_array)) == 2
        dim_number = findfirst(x -> !ismissing(x), dim_array)
        dim_symbol = dim_symbols_array[dim_number]
        vtmp = dim_array[dim_number]
        units = vtmp isa Real ? (dim_symbol == :φ ? u"°" : internal_length_unit) : unit(vtmp)
        v = T(to_internal_units(dim_symbol == :φ && φ isa Real ? deg2rad(vtmp) : vtmp))
    else
        throw(ArgumentError("Only one keyword for a certain dimension is allowed. Please choose one of "*
            (S == Cartesian ? "'x', 'y', 'z'." : "'r', 'φ', 'z'.")))
    end

    show_full_det = full_det && dim_symbol == :φ # The full_det keyword only makes sense for crossections in the xz plane in cylindrical grids

    (dim_symbol != :r && !(dim_symbol == :z && S == Cylindrical) ) ? aspect_ratio --> 1 : nothing
    title --> "Electric Field Lines @ $(dim_symbol) = $(round(units, Float64(uconvert(units, v*(dim_symbol == :φ ? u"rad" : internal_length_unit))), sigdigits=3))" * (show_full_det ? "\n(= $(round(units, Float64(uconvert(units, (v+π)%2π * u"rad")), sigdigits=3)) on left side)" : "")
    xguide --> (S == Cylindrical ? (dim_symbol == :r ? "φ" : "r") : (dim_symbol == :x ? "y" : "x"))
    yguide --> "z"
    unitformat --> :slash
    (S == Cylindrical && dim_symbol == :z) ? xguide :=  "" : nothing
    (S == Cylindrical && dim_symbol == :z) ? yguide :=  "" : nothing

    contacts_to_spawn_charges_for = Contact{T}[c for c in sim.detector.contacts if c.id != skip_contact]
    spawn_positions = CartesianPoint{T}[]
    grid = sim.electric_field.grid

    sample_lines = get_sample_lines(dim_symbol, v, grid, T(ustrip(to_internal_units(sampling))))

    for c in contacts_to_spawn_charges_for[:]
        surfs = ConstructiveSolidGeometry.surfaces(c.geometry)
        for l in sample_lines
            for surf in surfs
                pts = ConstructiveSolidGeometry.intersection(surf,l)
                for pt in pts
                    if pt in c
                        pt_in = pt + T(ustrip(to_internal_units(offset))) * normalize(ConstructiveSolidGeometry.normal(surf, pt))
                        if pt_in in sim.detector && !(pt_in in sim.detector.contacts)
                            push!(spawn_positions, pt_in)
                        end
                    end
                end
            end
        end
    end

    el_field_itp     = interpolated_vectorfield(sim.electric_field.data      , sim.electric_field.grid) #Interpolate the Electric fields, in which the charges will drift. Is passed to the drift function.
    el_field_itp_inv = interpolated_vectorfield(sim.electric_field.data .* -1, sim.electric_field.grid)
    spawn_positions = spawn_positions[findall(ipos -> ((spacing-1)+ipos)%spacing == 0, 1:length(spawn_positions))]

    for el_field in (el_field_itp, el_field_itp_inv)
        paths::Array{CartesianPoint{T}, 2} = fill(zero(CartesianVector{T}), length(spawn_positions), max_nsteps)
        last_step::Int = _drift_charge!(paths, Vector{T}(undef, max_nsteps), sim.detector, sim.point_types, sim.electric_potential.grid, CartesianPoint.(spawn_positions), ones(T, length(spawn_positions)), T(1e-9), el_field, Electron, verbose = false )
        for i in 1:size(paths, 1)
            path = @view paths[i, 1:last_step]
            @series begin
                seriescolor --> :white
                label --> ""
                x, y = if dim_symbol == :φ
                    broadcast(x -> x[1] * cos(v) + x[2] * sin(v), path)*internal_length_unit, #project to the φ-plane
                    broadcast(x -> x[3], path)*internal_length_unit
                elseif dim_symbol == :x
                    broadcast(x -> x[2], path)*internal_length_unit,
                    broadcast(x -> x[3], path)*internal_length_unit
                elseif dim_symbol == :y
                    broadcast(x -> x[1], path)*internal_length_unit,
                    broadcast(x -> x[3], path)*internal_length_unit
                elseif dim_symbol == :z
                    if S == Cylindrical
                        projection --> :polar
                        path = CylindricalPoint.(path)
                        broadcast(x -> x[2], path)*internal_angle_unit,
                        broadcast(x -> x[1], path)*internal_length_unit
                    else
                        broadcast(x -> x[1], path)*internal_length_unit,
                        broadcast(x -> x[2], path)*internal_length_unit
                    end
                elseif dim_symbol == :r
                    path = CylindricalPoint.(path)
                    broadcast(x -> x[2], path)*internal_angle_unit,
                    broadcast(x -> x[3], path)*internal_length_unit
                end
                x, y
            end
        end
    end
end


text:
"""
Single (GLM) trait IHT object, containing intermediate variables and temporary arrays
"""
mutable struct IHTVariable{T <: Float, M <: AbstractMatrix}
    # data and pre-specified parameters for fitting
    x      :: M                      # design matrix (genotypes) of subtype AbstractMatrix{T}
    y      :: Vector{T}              # response vector (phenotypes)
    z      :: VecOrMat{T}            # other non-genetic covariates
    k      :: Int                    # sparsity parameter (this is set to 0 if doubly sparse group IHT requires different k for each group)
    J      :: Int                    # Number of non-zero groups
    ks     :: Vector{Int}            # Sparsity level for each group. This is a empty vector if each group has same sparsity.
    d      :: UnivariateDistribution # Distribution for phenotypes
    l      :: Link                   # link function linking the mean μ with linear predictors xb via l(μ) = xb
    est_r  :: Symbol                 # Can be `:MM`, `:Newton`, or `:None`, specifying method used to estimate nuisance parameter in neg bin regression
    # internal IHT variables
    b      :: Vector{T}     # the statistical model for the genotype matrix, most will be 0
    b0     :: Vector{T}     # estimated model for genotype matrix in the previous iteration
    best_b :: Vector{T}     # best estimated genotype model in terms of loglikelihood
    xb     :: Vector{T}     # vector that holds x*b
    xk     :: Matrix{T}     # the n by k subset of the design matrix x corresponding to non-0 elements of b
    gk     :: Vector{T}     # numerator of step size. gk = df[idx].
    xgk    :: Vector{T}     # xk * gk, denominator of step size
    idx    :: BitVector     # idx[i] = 0 if b[i] = 0 and idx[i] = 1 if b[i] is not 0
    idx0   :: BitVector     # previous iterate of idx
    idc    :: BitVector     # idx[i] = 0 if c[i] = 0 and idx[i] = 1 if c[i] is not 0
    idc0   :: BitVector     # previous iterate of idc
    r      :: Vector{T}     # The difference between the observed and predicted response. For linear model this is the residual
    df     :: Vector{T}     # genotype portion of the score
    df2    :: Vector{T}     # non-genetic covariates portion of the score
    c      :: Vector{T}     # estimated model for non-genetic variates (first entry = intercept)
    c0     :: Vector{T}     # estimated model for non-genetic variates in the previous iteration
    best_c :: Vector{T}     # best estimated model for non-genetic variates in terms of loglikelihood
    zc     :: Vector{T}     # z * c (covariate matrix times c)
    zdf2   :: Vector{T}     # z * df2 needed to calculate non-genetic covariate contribution for denomicator of step size
    group  :: Vector{Int}   # vector denoting group membership
    weight :: Vector{T}     # weights (typically minor allele freq) that will scale b prior to projection
    μ      :: Vector{T}     # mean of the current model: μ = l^{-1}(xb)
    cv_wts :: Vector{T}     # weights for cross validation. cv_wts[i] = 0 means sample i should not be included in fitting.
    zkeep  :: BitVector     # tracks index of non-genetic covariates not subject to projection. zkeep[i] = true means `i` will not be projected.
    zkeepn :: Int           # Total number of covariates that aren't subject to projection
    full_b :: Vector{T}     # storage for full beta and full gradient
end

nsamples(v::IHTVariable) = count(!iszero, v.cv_wts)
nsnps(v::IHTVariable) = size(v.x, 2)
ncovariates(v::IHTVariable) = size(v.z, 2) # number of nongenetic covariates
ntraits(v::IHTVariable) = 1

function IHTVariable(x::M, z::AbstractVecOrMat{T}, y::AbstractVector{T},
    J::Int, k::Union{Int, Vector{Int}}, d::UnivariateDistribution, l::Link,
    group::AbstractVector{Int}, weight::AbstractVector{T}, est_r::Symbol,
    zkeep::BitVector) where {T <: Float, M <: AbstractMatrix}

    n = size(x, 1)
    p = size(x, 2)
    m = size(z, 1)
    q = size(z, 2)
    ly = length(y)
    lg = length(group)
    lw = length(weight)

    if !(ly == n == m)
        throw(DimensionMismatch("row dimension of y, x, and z ($ly, $n, $m) are not equal"))
    end

    if lg != p && lg != 0
        throw(DimensionMismatch("group must have length $p but was $lg"))
    end

    if lw != p && lw != 0
        throw(DimensionMismatch("weight must have length $p but was $lw"))
    end

    if typeof(k) == Int64
        columns = k
        ks = Int[]
    else
        columns = sum(k)
        ks, k = k, 0
    end

    if length(zkeep) != q
        throw(DimensionMismatch("zkeep must have length $q but was $(length(zkeep))"))
    end

    b      = Vector{T}(undef, p)
    b0     = Vector{T}(undef, p)
    best_b = Vector{T}(undef, p)
    xb     = Vector{T}(undef, n)
    xk     = Matrix{T}(undef, n, J * columns - 1) # subtracting 1 because the intercept will likely be selected in the first iter
    gk     = Vector{T}(undef, J * columns - 1)    # subtracting 1 because the intercept will likely be selected in the first iter
    xgk    = Vector{T}(undef, n)
    idx    = BitArray(undef, p)
    idx0   = BitArray(undef, p)
    idc    = BitArray(undef, q)
    idc0   = BitArray(undef, q)
    r      = Vector{T}(undef, n)
    df     = Vector{T}(undef, p)
    df2    = Vector{T}(undef, q)
    c      = Vector{T}(undef, q)
    best_c = Vector{T}(undef, q)
    c0     = Vector{T}(undef, q)
    zc     = Vector{T}(undef, n)
    zdf2   = Vector{T}(undef, n)
    μ      = Vector{T}(undef, n)
    cv_wts = Vector{T}(undef, n)
    storage = Vector{T}(undef, p + q)

    return IHTVariable{T, M}(
        x, y, z, k, J, ks, d, l, est_r,
        b, b0, best_b, xb, xk, gk, xgk, idx, idx0, idc, idc0, r, df, df2,
        c, c0, best_c, zc, zdf2, group, weight, μ, cv_wts, zkeep, sum(zkeep), storage)
end

function initialize(x::M, z::AbstractVecOrMat{T}, y::AbstractVecOrMat{T},
    J::Int, k::Union{Int, Vector{Int}}, d::Distribution, l::Link,
    group::AbstractVector{Int}, weight::AbstractVector{T}, est_r::Symbol,
    initialize_beta::Bool, zkeep::BitVector;
    cv_train_idx=trues(is_multivariate(y) ? size(x, 2) : size(x, 1)),
    verbose::Bool=false
    ) where {T <: Float, M <: AbstractMatrix}

    if is_multivariate(y)
        v = mIHTVariable(x, z, y, k, zkeep)
    else
        v = IHTVariable(x, z, y, J, k, d, l, group, weight, est_r, zkeep)
    end

    # initialize non-zero indices
    MendelIHT.init_iht_indices!(v, initialize_beta, cv_train_idx, verbose)

    return v
end

"""
Multivaraite Gaussian IHT object, containing intermediate variables and temporary arrays
"""
mutable struct mIHTVariable{T <: Float, M <: AbstractMatrix}
    # data and pre-specified parameters for fitting
    X      :: M             # p × n design matrix (genotypes) of subtype AbstractMatrix{T}
    Y      :: Matrix{T}     # r × n response matrix (phenotypes)
    Z      :: VecOrMat{T}   # q × n other non-genetic covariates (intercept goes here)
    k      :: Int           # sparsity parameter
    # internal IHT variables
    B      :: Matrix{T}     # r × p matrix that holds the statistical model for the genotype matrix, most will be 0
    B0     :: Matrix{T}     # estimated model for genotype matrix in the previous iteration
    best_B :: Matrix{T}     # best statistical model for the genotype matrix in terms of loglikelihood
    BX     :: Matrix{T}     # r × n matrix that holds B*X
    Xk     :: Matrix{T}     # the k × n subset of the design matrix x corresponding to non-0 elements of b
    idx    :: BitVector     # length p vector where idx[i] = 0 if i-th column of B is all zero, and idx[i] = 1 otherwise
    idx0   :: BitVector     # previous iterate of idx
    idc    :: BitVector     # length q vector where idc[i] = 0 i-th column of C is all zero, and idc[i] = 1 otherwise
    idc0   :: BitVector     # previous iterate of idc
    resid  :: Matrix{T}     # r × n matrix holding the residuals (Y - BX)
    df     :: Matrix{T}     # r × p genotype portion of the score = Γ(Y - XB)X'
    df2    :: Matrix{T}     # r × q non-genetic portion of the score
    dfidx  :: Matrix{T}     # r × k matrix storing df[:, idx], needed in stepsize calculation
    C      :: Matrix{T}     # r × q mtrix that holds the estimated model for non-genetic variates (first entry = intercept)
    C0     :: Matrix{T}     # estimated model for non-genetic variates in the previous iteration
    best_C :: Matrix{T}     # best statistical model for the non-genetic variates in terms of loglikelihood
    CZ     :: Matrix{T}     # r × n matrix holding C * Z (C times nongenetic covariates)
    μ      :: Matrix{T}     # mean of the current model: μ = BX + CZ
    Γ      :: Matrix{T}     # estimated inverse covariance matrix (TODO: try StaticArrays.jl here)
    Γ0     :: Matrix{T}     # Γ in previous iterate (TODO: try StaticArrays here)
    cv_wts :: Vector{T}     # weights for cross validation. cv_wts[i] = 0 means sample i should not be included in fitting.
    zkeep  :: BitVector     # tracks index of non-genetic covariates not subject to projection. zkeep[i] = true means `i` will not be projected.
    zkeepn :: Int           # Total number of covariates that aren't subject to projection
    # storage variables
    full_b :: Vector{T}     # storage for vectorized form of full beta [vec(B); vec(Z)]
    r_by_r1 :: Matrix{T}    # an r × r storage (needed in loglikelihood)
    r_by_r2 :: Matrix{T}    # another r × r storage (needed in loglikelihood)
    r_by_n1 :: Matrix{T}    # an r × n storage (needed in score! function)
    r_by_n2 :: Matrix{T}    # an r × n storage (needed in stepsize calculation)
    n_by_r  :: Matrix{T}    # an n × r storage (needed to efficiently compute gradient)
    p_by_r  :: Matrix{T}    # an p × r storage (needed to efficiently compute gradient)
    k_by_r  :: Matrix{T}    # an k × r storage (needed for debiasing)
    k_by_k  :: Matrix{T}    # an k × k storage (needed for debiasing)
end

function mIHTVariable(x::M, z::AbstractVecOrMat{T}, y::AbstractMatrix{T},
    k::Int, zkeep::BitVector) where {T <: Float, M <: AbstractMatrix}

    n = size(x, 2) # number of samples
    p = size(x, 1) # number of SNPs
    q = size(z, 1) # number of non-genetic covariates
    r = size(y, 1) # number of traits

    if !(n == size(y, 2) == size(z, 2))
        throw(DimensionMismatch("number of samples in y, x, and z = $(size(y, 2)), $n, $(size(z, 2)) are not equal"))
    end

    if length(zkeep) != q
        throw(DimensionMismatch("zkeep must have length $q but was $(length(zkeep))"))
    end

    B      = Matrix{T}(undef, r, p)
    B0     = Matrix{T}(undef, r, p)
    best_B = Matrix{T}(undef, r, p)
    BX     = Matrix{T}(undef, r, n)
    Xk     = Matrix{T}(undef, k - 1, n) # subtracting 1 because the intercept will likely be selected in the first iter
    idx    = BitArray(undef, p)
    idx0   = BitArray(undef, p)
    idc    = BitArray(undef, q)
    idc0   = BitArray(undef, q)
    resid  = Matrix{T}(undef, r, n)
    df     = Matrix{T}(undef, r, p)
    df2    = Matrix{T}(undef, r, q)
    dfidx  = Matrix{T}(undef, r, k - 1)
    C      = Matrix{T}(undef, r, q)
    C0     = Matrix{T}(undef, r, q)
    best_C = Matrix{T}(undef, r, q)
    CZ     = Matrix{T}(undef, r, n)
    μ      = Matrix{T}(undef, r, n)
    Γ      = Matrix{T}(I, r, r)
    Γ0     = Matrix{T}(I, r, r)
    cv_wts = Vector{T}(undef, n)
    full_b = Vector{T}(undef, r * (p + q))
    r_by_r1 = Matrix{T}(undef, r, r)
    r_by_r2 = Matrix{T}(undef, r, r)
    r_by_n1 = Matrix{T}(undef, r, n)
    r_by_n2 = Matrix{T}(undef, r, n)
    n_by_r = Matrix{T}(undef, n, r)
    p_by_r = Matrix{T}(undef, p, r)
    k_by_r = Matrix{T}(undef, k, r)
    k_by_k = Matrix{T}(undef, k, k)

    return mIHTVariable{T, M}(
        x, y, z, k,
        B, B0, best_B, BX, Xk, idx, idx0, idc, idc0, resid, df, df2, dfidx, C,
        C0, best_C, CZ, μ, Γ, Γ0, cv_wts, zkeep, r*sum(zkeep), full_b, r_by_r1,
        r_by_r2, r_by_n1, r_by_n2, n_by_r, p_by_r, k_by_r, k_by_k)
end

nsamples(v::mIHTVariable) = count(!iszero, v.cv_wts)
nsnps(v::mIHTVariable) = size(v.X, 1)
ncovariates(v::mIHTVariable) = size(v.Z, 1) # number of nongenetic covariates
ntraits(v::mIHTVariable) = size(v.Y, 1)

"""
Immutable object that houses results returned from a single-trait IHT run.
"""
struct IHTResult{T <: Float}
    time  :: Float64                    # total compute time
    logl  :: T                          # best loglikelihood achieved
    iter  :: Int64                      # number of iterations until convergence
    beta  :: Vector{T}                  # estimated beta for genetic predictors
    c     :: Vector{T}                  # estimated beta for nongenetic predictors
    J     :: Int64                      # maximum number of groups (1 for multivariate analysis)
    k     :: Union{Int64, Vector{Int}}  # maximum number of predictors (vector if group IHT have differently sized groups)
    group :: Vector{Int64}              # group membership
    d     :: Distribution               # distribution of phenotype
    σg    :: T                          # Estimated proportion of variance explained in phenotype
end
IHTResult(time, logl, iter, σg, v::IHTVariable) = IHTResult(time, logl, iter,
    v.best_b, v.best_c, v.J, v.k, v.group, v.d, σg)

"""
Immutable object that houses results returned from a multivariate Gaussian IHT run.
"""
struct mIHTResult{T <: Float}
    time   :: Float64                    # total compute time
    logl   :: T                          # best loglikelihood achieved
    iter   :: Int64                      # number of iterations until convergence
    beta   :: VecOrMat{T}                # estimated beta for genetic predictors
    c      :: VecOrMat{T}                # estimated beta for nongenetic predictors
    k      :: Int64                      # maximum number of predictors
    traits :: Int64                      # number of traits analyzed jointly
    Σ      :: Matrix{T}                  # estimated covariance matrix for multivariate analysis
    σg     :: Vector{T}                  # Estimated proportion of variance explained in phenotype
end
IHTResult(time, logl, iter, σg, v::mIHTVariable) = mIHTResult(time, logl, iter,
    v.best_B, v.best_C, v.k, ntraits(v), inv(v.Γ), σg)

"""
Displays IHTResults object
"""
function Base.show(io::IO, x::IHTResult)
    snp_position = findall(x -> x != 0, x.beta)
    nongenetic_position = findall(x -> x != 0, x.c)

    println(io, "\nIHT estimated ", count(!iszero, x.beta), " nonzero SNP predictors and ", count(!iszero, x.c), " non-genetic predictors.")
    println(io, "\nCompute time (sec):     ", x.time)
    println(io, "Final loglikelihood:    ", x.logl)
    println(io, "SNP PVE:                ", x.σg)
    println(io, "Iterations:             ", x.iter)
    println(io, "\nSelected genetic predictors:")
    print(io, DataFrame(Position=snp_position, Estimated_β=x.beta[snp_position]))
    println(io, "\n\nSelected nongenetic predictors:")
    print(io, DataFrame(Position=nongenetic_position, Estimated_β=x.c[nongenetic_position]))
end
# default IO for show is stdout
show(x::IHTResult) = show(stdout, x)

"""
Displays mIHTResult object
"""
function Base.show(io::IO, x::mIHTResult)
    println(io, "\nCompute time (sec):     ", x.time)
    println(io, "Final loglikelihood:    ", x.logl)
    println(io, "Iterations:             ", x.iter)
    for r in 1:x.traits
        println(io, "Trait $r's SNP PVE:      ", x.σg[r])
    end
    println(io, "\nEstimated trait covariance:")
    println(io, DataFrame(x.Σ, ["trait$i" for i in 1:x.traits]))

    for r in 1:x.traits
        β1 = @view(x.beta[r, :])
        C1 = @view(x.c[r, :])
        snp_position = findall(x -> x != 0, β1)
        nongenetic_position = findall(x -> x != 0, C1)
        println(io, "\nTrait $r: IHT estimated ", count(!iszero, β1),
            " nonzero SNP predictors")
        println(io, DataFrame(Position=snp_position, Estimated_β=β1[snp_position]))
        println(io, "\nTrait $r: IHT estimated ", count(!iszero, C1),
            " non-genetic predictors")
        println(io, DataFrame(Position=nongenetic_position, Estimated_β=C1[nongenetic_position]))
    end
end

"""
verbose printing of cv results
"""
function print_cv_results(io::IO, errors::Vector{T},
    path::AbstractVector{<:Integer}, k::Int) where {T <: Float}
    println(io, "\n\nCrossvalidation Results:")
    println(io, "\tk\tMSE")
    for i = 1:length(errors)
        println(io, "\t", path[i], "\t", errors[i])
    end
    println(io, "\nBest k = $k\n")
end
# default IO for print_cv_results is stdout
print_cv_results(errors::Vector{T}, path::AbstractVector{<:Integer}, k::Int
    ) where {T <: Float} = print_cv_results(stdout, errors, path, k)

"""
verbose printing of running `iht_run_many_models` with a bunch of models
"""
function print_a_bunch_of_path_results(io::IO, loglikelihoods::AbstractVector{T},
    path::AbstractVector{<:Integer}) where {T <: Float}
    println(io, "\n\nResults of running all the model sizes specified in `path`:")
    println(io, "\tk\tloglikelihoods")
    for i = 1:length(loglikelihoods)
        println(io, "\t", path[i], "\t", loglikelihoods[i])
    end
    println(io, "\nWe recommend running cross validation through `cv_iht` on " *
    "appropriate model sizes, which is roughly the values of k where the " *
    "loglikelihood stop increasing significantly.")
end
# default IO for print_a_bunch_of_path_results is stdout
print_a_bunch_of_path_results(loglikelihoods::AbstractVector{T},
    path::AbstractVector{<:Integer}) where {T <: Float} =
    print_a_bunch_of_path_results(stdout, loglikelihoods, path)


text: Dict{String, Any}("1" => Dict{String, Any}("2" => Dict{String, Any}("value" => "3", "type" => "integer")))


text:
# %% Preloads
using Statistics
using FFTW
using Plots
using BenchmarkTools
using Profile
using LinearAlgebra
using Measures
using HDF5

#%%
function err_extract(iso_out,J,N)
    pdf_err = zeros(J,J)
    temp = zeros(N)
    for j2=1:J
        for j1=1:J
            temp = (iso_out[j1,j2,:].-mean(iso_out[j1,j2,:]))./mean(iso_out[j1,j2,:])
            pdf_err[j1,j2] = maximum(temp)-minimum(temp)
         end
     end
     return maximum(pdf_err[2:7,2:7])
end

function S20_iso(data,J,L,N)
    S20 = reshape(data[2+J*L+1:2+J*L+J*L*J*L,:],J,L,J,L,N)
    iso_mat = zeros(J,J,N,L)
    for l=1:L
        intermed = zeros(J,J,N,L)
        for d=1:L
            intermed[:,:,:,d].=S20[:,l,:,mod(l+d,L)+1,:]
        end
        iso_mat[:,:,:,l] .= dropdims(sum(intermed,dims=4),dims=4)
    end
    pow_mat = zeros(J,J,N)
    pow_mat = dropdims(sum(iso_mat,dims=4),dims=4)
    return pow_mat
end

# %% Plat with some data
out0 = h5read("./Data/abs2j8l16w1p2.h5", "main/data")
angle_iso = S20_iso(out0,8,16,360)
err_extract(angle_iso,8,360)

# %%
J = 8
L= 16

plot(angle_iso[3,4,:])

out0[2+J*L+1:2+J*L+J*L*J*L,:]
test = reshape(out0[2+J*L+1:2+J*L+J*L*J*L,:],J,L,J,L,360)

plot(1:360,test[3,2,3,2,:])
plot!(1:360,test[3,4,3,4,:])

plot(1:360,test[3,1,3,16,:])
plot!(1:360,test[3,1,3,2,:])

plot!(1:360,test[3,2,3,1,:])
plot!(1:360,test[3,2,3,3,:])

plot!(1:360,test[3,3,3,2,:])
plot!(1:360,test[3,3,3,4,:])

plot!(1:360,test[3,4,3,3,:])
plot!(1:360,test[3,4,3,5,:])

fig = plot()
for i=1:L
    plot!(fig,1:360,test[3,i,3,mod1(i-1,L),:],show=true)
    plot!(fig,1:360,test[3,i,3,mod1(i+1,L),:],show=true)
end
fig

# %% Retry understanding mod... but actually, do not need to for simplest case.

function S20_iso(data,J,L,N)
    S20 = reshape(data[2+J*L+1:2+J*L+J*L*J*L,:],J,L,J,L,N)
    iso_mat0 = zeros(J,L,J,N)
    iso_mat1 = zeros(J,J,N)

    iso_mat0 .= dropdims(sum(S20,dims=4),dims=4)
    iso_mat1 .= dropdims(sum(iso_mat0,dims=2),dims=2)

    return iso_mat1
end

# %% Plot with some data
out0 = h5read("./Data/abs2j8l16w1p2.h5", "main/data")
angle_iso = S20_iso(out0,8,16,360)
err_extract(angle_iso,8,360)

plotlyjs()

plot(angle_iso[4,4,:])

test = reshape(out0[2+J*L+1:2+J*L+J*L*J*L,:],J,L,J,L,360)

plot(test[4,4,4,4,:])

plot(angle_iso[3,:,:])

# %%
function S20_iso(data,J,L,N)
    S20 = reshape(data[2+J*L+1:2+J*L+J*L*J*L,:],J,L,J,L,N)
    iso_mat = zeros(J,J,N,L)
    for l=1:L
        intermed = zeros(J,J,N,L)
        for d=1:L
            intermed[:,:,:,d].=S20[:,l,:,mod1(l+d,L)+1,:]
        end
        iso_mat[:,:,:,l] .= dropdims(sum(intermed,dims=4),dims=4)
    end
    pow_mat = zeros(J,J,N)
    pow_mat = dropdims(sum(iso_mat,dims=4),dims=4)
    return pow_mat
end

fig = plot()
for i=3
    plot!(fig,1:360,test[3,i,3,mod1(i-1,L),:],show=true)
    plot!(fig,1:360,test[3,i,3,mod1(i+1,L),:],show=true)

    plot!(fig,1:360,test[3,i-1,3,mod1(i,L),:],show=true)
    plot!(fig,1:360,test[3,i+1,3,mod1(i,L),:],show=true)
end
fig

L=8
J=8
out0 = h5read("./Data/abs2j8l8w1p1.h5", "main/data")
out0[2+J*L+1:2+J*L+J*L*J*L,:]
test = reshape(out0[2+J*L+1:2+J*L+J*L*J*L,:],J,L,J,L,360)

h5write("./Data/abs2j8l8w1p1_forpy.h5", "main/data", test)

L=16
J=8
out0 = h5read("./Data/abs2j8l16w1p2.h5", "main/data")
out0[2+J*L+1:2+J*L+J*L*J*L,:]
test = reshape(out0[2+J*L+1:2+J*L+J*L*J*L,:],J,L,J,L,360)

h5write("./Data/abs2j8l16w1p2_forpy.h5", "main/data", test)

##

UpperTriangular(test[3,:,3,:,1])

function upper_tri(A,n)
    v = zeros(Int(n*(n+1)/2),360)
    k = 0
    for i in 1:n
        for j in 1:i
            v[k+j,:] = A[j,i,:]
        end
        k+=i
    end
    v
end

temp_v = upper_tri(test[3,:,3,:,:],16)

u = 2 .*mod.(1:136,2).-1

dot_prod = dropdims(sum(temp_v .* u,dims=1),dims=1)

plot(1:360,dot_prod)

u2 = 2 .*mod.(1:16,2).-1

function upper_tri_1d(A,n)
    v = zeros(Int(n*(n+1)/2))
    k = 0
    for i in 1:n
        for j in 1:i
            v[k+j] = A[j,i]
        end
        k+=i
    end
    v
end

u3 = upper_tri_1d(ones(16,16).*transpose(u2),16)

dot_prod = dropdims(sum(temp_v .* u3,dims=1),dims=1)

plot(1:360,dot_prod)

u4 = upper_tri_1d(u2.*transpose(u2),16)

dot_prod = dropdims(sum(temp_v .* u4,dims=1),dims=1)

plot(1:360,dot_prod)

maximum(dot_prod)
minimum(dot_prod)

# %% Redoing without the S1 contribution

function S20_iso(data,J,L,N)
    S20 = reshape(data[2+J*L+1:2+J*L+J*L*J*L,:],J,L,J,L,N)
    iso_mat = zeros(Int(L*(L+1)/2),J,J,N)
    k = 0
    @inbounds for i in 1:L
        for j in 1:i
            iso_mat[k+j,:,:,:] = S20[:,j,:,i,:]
        end
        k+=i
    end
    pow_mat = dropdims(sum(iso_mat,dims=1),dims=1)
    return pow_mat
end

function err_extract(iso_out,J,N)
    pdf_err = zeros(J,J)
    temp = zeros(N)
    for j2=1:J
        for j1=1:J
            temp = (iso_out[j1,j2,:].-mean(iso_out[j1,j2,:]))./mean(iso_out[j1,j2,:])
            pdf_err[j1,j2] = maximum(temp)-minimum(temp)
         end
     end
     return maximum(pdf_err[2:7,2:7])
end

out0 = h5read("./Data/abs2j8l16w1p1.h5", "main/data")
angle_iso = S20_iso(out0,8,16,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l16w2p1.h5", "main/data")
angle_iso = S20_iso(out0,8,16,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l16w3p1.h5", "main/data")
angle_iso = S20_iso(out0,8,16,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l16w1p2.h5", "main/data")
angle_iso = S20_iso(out0,8,16,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l16w2p2.h5", "main/data")
angle_iso = S20_iso(out0,8,16,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l16w3p2.h5", "main/data")
angle_iso = S20_iso(out0,8,16,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l8w1p1.h5", "main/data")
angle_iso = S20_iso(out0,8,8,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l8w2p1.h5", "main/data")
angle_iso = S20_iso(out0,8,8,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l8w3p1.h5", "main/data")
angle_iso = S20_iso(out0,8,8,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l8w1p2.h5", "main/data")
angle_iso = S20_iso(out0,8,8,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l8w2p2.h5", "main/data")
angle_iso = S20_iso(out0,8,8,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l8w3p2.h5", "main/data")
angle_iso = S20_iso(out0,8,8,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l32w1p1.h5", "main/data")
angle_iso = S20_iso(out0,8,32,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l32w2p1.h5", "main/data")
angle_iso = S20_iso(out0,8,32,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l32w3p1.h5", "main/data")
angle_iso = S20_iso(out0,8,32,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l16w1p1.h5", "main/data")
angle_iso = S20_iso(out0,8,16,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l16w2p1.h5", "main/data")
angle_iso = S20_iso(out0,8,16,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l16w3p1.h5", "main/data")
angle_iso = S20_iso(out0,8,16,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l16w1p2.h5", "main/data")
angle_iso = S20_iso(out0,8,16,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l16w2p2.h5", "main/data")
angle_iso = S20_iso(out0,8,16,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l16w3p2.h5", "main/data")
angle_iso = S20_iso(out0,8,16,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l8w1p1.h5", "main/data")
angle_iso = S20_iso(out0,8,8,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l8w2p1.h5", "main/data")
angle_iso = S20_iso(out0,8,8,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l8w3p1.h5", "main/data")
angle_iso = S20_iso(out0,8,8,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l8w1p2.h5", "main/data")
angle_iso = S20_iso(out0,8,8,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l8w2p2.h5", "main/data")
angle_iso = S20_iso(out0,8,8,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l8w3p2.h5", "main/data")
angle_iso = S20_iso(out0,8,8,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l32w1p1.h5", "main/data")
angle_iso = S20_iso(out0,8,32,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l32w2p1.h5", "main/data")
angle_iso = S20_iso(out0,8,32,360)
err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l32w3p1.h5", "main/data")
angle_iso = S20_iso(out0,8,32,360)
err_extract(angle_iso,8,360)

# %%
function S12_iso(data,J,L,N)
    S20 = reshape(data[2+J*L+J*L*J*L+1:2+J*L+J*L*J*L+J*L*J*L,:],J,L,J,L,N)
    iso_mat = zeros(Int(L*(L+1)/2),J,J,N)
    k = 0
    @inbounds for i in 1:L
        for j in 1:i
            iso_mat[k+j,:,:,:] = S20[:,j,:,i,:]
        end
        k+=i
    end
    pow_mat = dropdims(sum(iso_mat,dims=1),dims=1)
    return pow_mat
end

function err_extract(iso_out,J,N)
     pdf_err = zeros(J,J)
     temp = zeros(N)
     for j2=1:J
         for j1=1:J
             temp = (iso_out[j1,j2,:].-mean(iso_out[j1,j2,:]))./mean(iso_out[j1,j2,:])
             pdf_err[j1,j2] = maximum(temp)-minimum(temp)
          end
      end
      pdf_err[isnan.(pdf_err)] .= 0
      return maximum(pdf_err[2:7,2:7])
 end

out0 = h5read("./Data/j8l8w1p2.h5", "main/data")
angle_iso = S12_iso(out0,8,8,360)
val = err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l8w2p2.h5", "main/data")
angle_iso = S12_iso(out0,8,8,360)
val = err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l8w3p2.h5", "main/data")
angle_iso = S12_iso(out0,8,8,360)
val = err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l16w1p2.h5", "main/data")
angle_iso = S12_iso(out0,8,16,360)
val = err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l16w2p2.h5", "main/data")
angle_iso = S12_iso(out0,8,16,360)
val = err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l16w3p2.h5", "main/data")
angle_iso = S12_iso(out0,8,16,360)
val = err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l8w1p2.h5", "main/data")
angle_iso = S12_iso(out0,8,8,360)
val = err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l8w2p2.h5", "main/data")
angle_iso = S12_iso(out0,8,8,360)
val = err_extract(angle_iso,8,360)

out0 = h5read("./Data/j8l8w3p2.h5", "main/data")
angle_iso = S12_iso(out0,8,8,360)
val = err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l16w1p2.h5", "main/data")
angle_iso = S12_iso(out0,8,16,360)
val = err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l16w2p2.h5", "main/data")
angle_iso = S12_iso(out0,8,16,360)
val = err_extract(angle_iso,8,360)

out0 = h5read("./Data/abs2j8l16w3p2.h5", "main/data")
angle_iso = S12_iso(out0,8,16,360)
val = err_extract(angle_iso,8,360)

# %%
x = 1:360
plot(x,sin.(π/180*x) .+ cos.(π/180*x))
plot!(x,abs.(sin.(π/180*x)) .+ abs.(cos.(π/180*x)))
plot!(x,abs.(sin.(π/180*x)).^2 .+ abs.(cos.(π/180*x)).^2)
plot!(x,abs.(sin.(π/180*x)).^4 .+ abs.(cos.(π/180*x)).^4)
plot!(x,abs.(sin.(π/180*x)).^6 .+ abs.(cos.(π/180*x)).^6)
plot!(x,abs.(sin.(π/180*x)).^(.5) .+ abs.(cos.(π/180*x)).^(.5))

out0 = h5read("./Data/abs2j8l16w1p2.h5", "main/data")
test = reshape(out0[2+J*L+1:2+J*L+J*L*J*L,:],J,L,J,L,360)
angle_iso = S20_iso(out0,8,16,360)

plot(angle_iso[4,4,:])

out0 = h5read("./Data/j8l16w1p2.h5", "main/data")
test = reshape(out0[2+J*L+1:2+J*L+J*L*J*L,:],J,L,J,L,360)
angle_iso = S20_iso(out0,8,16,360)

plot!(angle_iso[4,4,:])

function S20_iso(data,J,L,N)
    S20 = reshape(data[2+J*L+1:2+J*L+J*L*J*L,:],J,L,J,L,N)
    iso_mat = zeros(Int(L*(L+1)/2),J,J,N)
    k = 0
    @inbounds for i in 1:L
        for j in 1:i
            iso_mat[k+j,:,:,:] = sqrt.(S20[:,j,:,i,:])
        end
        k+=i
    end
    pow_mat = dropdims(sum(iso_mat,dims=1),dims=1)
    return pow_mat
end

out0 = h5read("./Data/abs2j8l16w1p2.h5", "main/data")
test = reshape(out0[2+J*L+1:2+J*L+J*L*J*L,:],J,L,J,L,360)
angle_iso = S20_iso(out0,8,16,360)

plot(angle_iso[4,4,:])

function S20_iso(data,J,L,N)
    S20 = reshape(data[2+J*L+1:2+J*L+J*L*J*L,:],J,L,J,L,N)
    iso_mat = zeros(Int(L*(L+1)/2),J,J,N)
    k = 0
    @inbounds for i in 1:L
        for j in 1:i
            iso_mat[k+j,:,:,:] = S20[:,j,:,i,:]
        end
        k+=i
    end
    pow_mat = dropdims(sum(iso_mat.*iso_mat,dims=1),dims=1)
    return pow_mat
end

out0 = h5read("./Data/abs2j8l16w1p2.h5", "main/data")
test = reshape(out0[2+J*L+1:2+J*L+J*L*J*L,:],J,L,J,L,360)
angle_iso = S20_iso(out0,8,16,360)

plot(angle_iso[4,4,:])

x = 1:360
plot(x,sin.(π/180*x) .+ 10)
plot(x,(sin.(π/180*x) .+ 10).^2)
plot!(x,(cos.(π/180*x) .+ 10).^2)
plot(x,(sin.(π/180*x) .+ 10).^2+(cos.(π/180*x) .+ 10).^2)

max_stars_repo_name: andrew-saydjari/DHC

max_issues_repo_name: andrew-saydjari/DHC

max_forks_repo_name: andrew-saydjari/DHC


text:
include("./imports.jl")

@testset ExtendedTestSet "`protect`" begin
    mutable struct C
        x::Int
        y::Int
    end

    mutable struct B
        c::C
    end

    mutable struct A
        b1::B
        b2::B
    end

    makea() = A(
        B(C(1, 2)),
        B(C(3, 4)),
    )

    @testset ExtendedTestSet "fully protected" begin
        a_p = protect(makea(), NamedTuple())  # protect all child structs of `a`
        @test_throws ProtectedException a_p.b1 isa Protected
    end

    @testset ExtendedTestSet "read direct child" begin
        a_p2 = protect(makea(), (b1=Read(),))  # protect everything but `b1`
        @test_nowarn a_p2.b1
        @test_throws ProtectedException a_p2.b1 = C(5, 5)
    end

    @testset ExtendedTestSet "write direct child" begin
        a_p2 = protect(makea(), (b1=Write(),))  # protect everything but `b1`
        @test_nowarn a_p2.b1
        @test_nowarn a_p2.b1 = B(C(5, 5))
    end
    @testset ExtendedTestSet "write nested" begin
        a_p3 = protect(makea(), (b1 = (c = (x=Write(),),),))  # allow mutating only a.b1.c.x
        @test_nowarn a_p3.b1.c.x = 2
        @test_nowarn a_p3.b1.c.x
        @test a_p3.b1.c isa Protected
        @test_throws ProtectedException a_p3.b1.c.y = 2
    end
end


text:
using StatisticalNonlocality, LinearAlgebra, FFTW, SparseArrays
import StatisticalNonlocality: chebyshev, fourier_nodes, fourier_wavenumbers
import StatisticalNonlocality: droprelativezeros!

# overall velocity scale
U₀ = 1e0
# transition rate
γ = 1e2  # 1e0 # 1e2
# diffusivity
κ = 1e-2 # 1e0 # 1e-2
# phase
ω = γ
# boundary condition, dirichlet = false is neumann
dirichlet = false
# output file name
filename = "local.jld2"

# minimal configuration is M = 64, N = 8
N = 8
M = 8 * 8

# The first step is to build the discretization and the block operators

Dz, z = chebyshev(M)
a, b = 0, 2π
k = fourier_wavenumbers(N, L = b - a)
x = fourier_nodes(N, a = a, b = b)

ℱ = fft(I + zeros(N, N), 1)
ℱ⁻¹ = ifft(I + zeros(N, N), 1)
Dx = real.(ℱ⁻¹ * Diagonal(im .* k) * ℱ)
x = reshape(x, (N, 1))
z = reshape(z, (1, M + 1))

ψ¹ = U₀ * sin.(x) .* cos.(π / 2 * z)
ψ² = U₀ * cos.(x) .* cos.(π / 2 * z)
ψ³ = -U₀ * sin.(x) .* cos.(π / 2 * z)
ψ⁴ = -U₀ * cos.(x) .* cos.(π / 2 * z)

∂z = kron(Dz, I + zeros(N, N))
∂x = kron(I + zeros(M + 1, M + 1), Dx)

e1 = norm(∂x * ψ¹[:] - ψ²[:])
println("The error in take in the partial with respect to x is ", e1)
e2 = norm(∂z * (ψ¹[:]) - (sin.(x).*cos.(z))[:])
println("The error in taking the partial with respect to y is ", e2)

zlifted = sparse(kron(Diagonal(z[:]), I + zeros(N, N)))
zlifted2 = sparse(kron(I + zeros(N, N), Diagonal(z[:])))

##
# u⃗⋅∇
advection_operator(ψ, ∂x, ∂z) =
    Diagonal(∂z * ψ[:]) * ∂x + Diagonal(-∂x * ψ[:]) * ∂z

A¹ = advection_operator(ψ¹, ∂x, ∂z)
A² = advection_operator(ψ², ∂x, ∂z)
A³ = advection_operator(ψ³, ∂x, ∂z)
A⁴ = advection_operator(ψ⁴, ∂x, ∂z)

Δ = ∂x^2 + ∂z^2

tmp = sparse(A⁴)
println(length(tmp.nzval))
droprelativezeros!(tmp)
println(length(tmp.nzval))

# Create the block linear operator
L = [(γ*I-κ*Δ) -ω*I 0.5*A¹; ω*I (γ*I-κ*Δ) -0.5*A²; A¹ -A² (2*γ*I-κ*Δ)]

# Grab Boundary Indices
zlifted = sparse(kron(Diagonal(z[:]), I + zeros(N, N)))
zlifted = [zlifted 0*I 0*I; 0*I zlifted 0*I; 0*I 0*I zlifted]
∂Ω¹ = zlifted.rowval[zlifted.nzval.==z[1]]
for (i, ∂i) in enumerate(∂Ω¹)
    if dirichlet
        L[∂i, :] .= 0.0
        L[∂i, ∂i] = 1.0
    else
        L[∂i, :] .= 0.0
        blockindex = div((i - 1), N)
        inds = (blockindex*N*(M+1)+1):((blockindex+1)*N*(M+1))
        L[∂i, inds] .= ∂z[∂Ω¹[(i-1)%N+1], :]
    end
end
∂Ω² = zlifted.rowval[zlifted.nzval.==z[end]]
for (i, ∂i) in enumerate(∂Ω²)
    if dirichlet
        L[∂i, :] .= 0.0
        L[∂i, ∂i] = 1.0
    else
        L[∂i, :] .= 0.0
        blockindex = div((i - 1), N)
        inds = (blockindex*N*(M+1)+1):((blockindex+1)*N*(M+1))
        L[∂i, inds] .= ∂z[∂Ω²[(i-1)%N+1], :]
    end
end
∂Ω = vcat(∂Ω¹, ∂Ω²)

Q, R = qr(L)

G = qr(L) \ I
u¹ = Diagonal(∂z * ψ¹[:])
u² = Diagonal(∂z * ψ²[:])
v¹ = Diagonal(-∂x * ψ¹[:])
v² = Diagonal(-∂x * ψ²[:])
# Account for Boundary conditions in U and V
# Can be thought of as an operator that masks the values on boundaries
# Technically the mask is applied to the entire rhs, but since we want the operator
# we just applied it here.
# This is equivalent since diagonal matrices commute and masking is idempotent
U = -0.5 * [u¹; u²; 0 * I]
Uᵀ = [u¹ u² 0 * I]
V = -0.5 * [v¹; v²; 0 * I]
Vᵀ = [v¹ v² 0 * I]
for ∂i in ∂Ω
    U[∂i, :] .= 0.0
    V[∂i, :] .= 0.0
end

EF¹¹ = Uᵀ * G * U # flux in x due to gradients in x
EF¹² = Uᵀ * G * V # flux in x due to gradients in z
EF²¹ = Vᵀ * G * U # flux in z due to gradients in x
EF²² = Vᵀ * G * V # flux in z due to gradients in z

E = [EF¹¹ EF¹²; EF²¹ EF²²]
# Check that all the eigenvalues are negative (we will multiply by -1 later)
# Note that there will be spurious eigenvalues due to the use of chebyshev
λE = eigvals(E)
bools = real.(λE) .≤ eps(1e3 * maximum(abs.(λE)))
println(sum(bools) == length(λE))

maximum(abs.(EF²¹ + EF¹²))
EF¹¹ *= -1.0
EF¹² *= -1.0
EF²¹ *= -1.0
EF²² *= -1.0

makelocal(E) = Array(Diagonal(sum(E, dims = 2)[:]))

function localdiffusivity(E, N, M)
    local localdiff = makelocal(E)
    localdiff = [localdiff[i, i] for i = 1:size(E)[1]]
    local localdiff = reshape(localdiff, (N, M + 1))
    return localdiff
end

function avglocaldiffusivity(E, N, M)
    local localdiff = localdiffusivity(E, N, M)
    local localdiff = sum(localdiff, dims = 1)[:] ./ N # zero'th mode
    return localdiff
end

# the local diffusivity is proportional to 1/γ
function grabdiagonal(A)
    MM = minimum(size(A))
    diagA = zeros(MM)
    for i in 1:MM
        diagA[i] = A[i, i]
    end
    return diagA
end

##
using JLD2
# always overwrite
data_directory = "data"
mkpath(data_directory)
filepath = data_directory * "/" * filename
if isfile(filepath)
    rm(filepath)
end
##
file = jldopen(filepath, "a+")
# diffusivity
groupname = "diffusivity"
JLD2.Group(file, groupname)
file[groupname]["K11"] = EF¹¹
file[groupname]["K12"] = EF¹²
file[groupname]["K21"] = EF²¹
file[groupname]["K22"] = EF²²
# parameters
groupname = "parameters"
JLD2.Group(file, groupname)
file[groupname]["γ"] = γ
file[groupname]["ω"] = ω
file[groupname]["κ"] = κ
# grid
groupname = "grid"
JLD2.Group(file, groupname)
file[groupname]["x"] = x
file[groupname]["z"] = z
# transition matrix
groupname = "transition"
JLD2.Group(file, groupname)

T = [-1 1 0 0; 0 -1 1 0; 0 0 -1 1; 1 0 0 -1]
sT = (T + T') / 2
aT = (T - T') / 2

file[groupname]["T"] = γ * sT + ω * aT
# streamfunction
groupname = "streamfunction"
JLD2.Group(file, groupname)
file[groupname]["ψ¹"] = ψ¹
file[groupname]["ψ²"] = ψ²
file[groupname]["ψ³"] = ψ³
file[groupname]["ψ⁴"] = ψ⁴
# velocities
groupname = "velocities"
JLD2.Group(file, groupname)
file[groupname]["u¹"] = reshape(grabdiagonal(u¹), size(ψ¹))
file[groupname]["u²"] = reshape(grabdiagonal(u²), size(ψ¹))
file[groupname]["v¹"] = reshape(grabdiagonal(v¹), size(ψ¹))
file[groupname]["v²"] = reshape(grabdiagonal(v²), size(ψ¹))

# local diffusivities
groupname = "localdiffusivity"
JLD2.Group(file, groupname)
file[groupname]["κ11"] = localdiffusivity(EF¹¹, N, M)
file[groupname]["κ12"] = localdiffusivity(EF¹², N, M)
file[groupname]["κ21"] = localdiffusivity(EF²¹, N, M)
file[groupname]["κ22"] = localdiffusivity(EF²², N, M)

# zonally averaged local diffusivities
groupname = "averagelocaldiffusivity"
JLD2.Group(file, groupname)
file[groupname]["κ11"] = avglocaldiffusivity(EF¹¹, N, M)
file[groupname]["κ12"] = avglocaldiffusivity(EF¹², N, M)
file[groupname]["κ21"] = avglocaldiffusivity(EF²¹, N, M)
file[groupname]["κ22"] = avglocaldiffusivity(EF²², N, M)

close(file)

max_stars_repo_name: sandreza/StatisticalNonlocality

max_issues_repo_name: sandreza/StatisticalNonlocality

max_forks_repo_name: sandreza/StatisticalNonlocality


text:
function norm_matrix_induced(A::Array{Float64,2})
    N = size(A)[1]
    # FIXME no check for matrix form

    # unit-norm vector
    d = 1/sqrt(N)
    v1 = ones(N)*d
    #
    v = A*v1
    return norm(v)
end


text:
using Random, Test, CollegeStratBase, CollegeStratCollege, ModelObjectsLH

cs = CollegeStratCollege;

## One college
function gradrule_main_test()
	@testset "GradRule Main" begin
		symTb = symbol_table();
		# test_header("GradRule Main")
		rng = MersenneTwister(10);
		for grSwitch in cs.grad_rule_test_switches()
			# Not all grad rules have non-empty `settings_list`
			sList = settings_list(grSwitch, symTb);

			iCollege = 3;
			g = cs.make_test_grad_rule(grSwitch, iCollege);

			# Access functions
			nMin = min_courses_to_grad(g, 4);
			@test nMin > 5
			tFirstGrad = t_first_grad(g);
			@test tFirstGrad > 2

			# One state
			h0 = 1.2;
			gProb = grad_prob(g, tFirstGrad - 1,
				100.0, 2 * nMin, h0);
			@test gProb ≈ 0.0
			gProb = grad_prob(g, tFirstGrad,
				100.0, nMin, h0);
			@test gProb > 0.0

			# Check that gradProb has the right values at the bounds
			if isa(g, GradRuleLinear)
				if use_hc_level(g)
					hLow = h_min(g);
					hHigh = h_max(g);
				else
					hLow = h_min(g) + h0;
					hHigh = h_max(g) + h0;
				end
				gProb = grad_prob(g, tFirstGrad+1, hLow, nMin, h0);
				@test isapprox(gProb, grad_prob_hmin(g))
				gProb = grad_prob(g, tFirstGrad+1, hHigh, nMin, h0);
				@test isapprox(gProb, grad_prob_max(g))
			end

			# Array inputs
			t = tFirstGrad;
			hM = 3.0 .+ rand(rng, 8,3) .- 0.5;
			hM[1,1] = 55.0;
			hM[2,1] = 0.1;
			h0M = 0.5 .* hM;
			nM = nMin .+ rand(rng, [-1,0,1], 8, 3);
			gProbM = grad_prob(g, t, hM, nM, h0M);
			@test size(gProbM) == size(hM)
			# Check that random draws produced grad and no grad outcomes
			@test (minimum(gProbM) < 0.1)  &&  (maximum(gProbM) > 0.5)

			gProb1M = grad_prob(g, t, hM, nM, h0);
			@test size(gProb1M) == size(hM)

			gProb2M = similar(gProbM);
			n1, n2 = size(hM);
			for i1 = 1 : n1
				for i2 = 1 : n2
					gProb2M[i1, i2] = grad_prob(g, t, hM[i1, i2],
						nM[i1, i2], h0M[i1, i2]);
				end
			end
			@test gProb2M ≈ gProbM
		end
	end
end

function gradrule_grid_test()
	@testset "Grid" begin
		# test_header("GradRule grid")
		for grSwitch in cs.grad_rule_test_switches()
			iCollege = 3;
			g = cs.make_test_grad_rule(grSwitch, iCollege);

			t = t_first_grad(g) + 1;
			hV = collect(range(0.1, 15.0, length = 14));
			nV = collect(0 : 4 : 22);
			h0 = 1.2;
			prob_hnM = grad_prob_grid(g, t, hV, nV, h0);
			@test size(prob_hnM) == (length(hV), length(nV))

			valid = true;
			for (i_h, h) in enumerate(hV)
				for (i_n, n) in enumerate(nV)
					isValid = (prob_hnM[i_h, i_n] ≈
						grad_prob(g, t, h, n, h0));
					valid = valid && isValid;
				end
			end
		end
	end
end

## All colleges
function gradrule_set_test()
	@testset "Set" begin
		# test_header("GradRuleSet")
		for grSwitch in cs.grad_rule_test_switches()
			gs = make_grad_rule_set(ObjectId(:gradRule), grSwitch);
			settings_table(gs)
			nc = n_colleges(gs);
			g = cs.make_grad_rule(gs, nc-1);
			@test isa(g, GradRule);

			@test is_two_year(gs, 1)
			@test !is_two_year(gs, nc)

			@test grad_prob_min(gs) < 0.5
			@test grad_prob_max(gs) > 0.5
			@test min_courses_to_grad(gs, 1) > 50
		end
    end
end

@testset "GradRule" begin
	# for switches in p_hbase_switches()
	# 	prob_h_base_test(switches);
	# end
    gradrule_main_test()
    gradrule_grid_test()
    gradrule_set_test()
end

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


text:
@testset "edgevals_type($G{$V, $E_VALS})" for
    G in (ValGraph, ValOutDiGraph, ValDiGraph),
    V in TEST_VERTEX_TYPES_SMALL,
    E_VALS in TEST_EDGEVAL_TYPES_SMALL

    g = G{V, E_VALS}(0)
    @assert g isa G{V, E_VALS}

    @test edgevals_type(g) == E_VALS
    @test edgevals_type(typeof(g)) == E_VALS
    @test edgevals_type(G{V, E_VALS}) == E_VALS

end

@testset "edgetype($G{$G, $E_VALS})" for
    G in (ValGraph, ValOutDiGraph, ValDiGraph),
    V in TEST_VERTEX_TYPES_SMALL,
    E_VALS in TEST_EDGEVAL_TYPES_SMALL

    g = G{V, E_VALS}(0)
    @assert g isa G{V, E_VALS}

    ET_should_be = (is_directed(g) ? ValDiEdge : ValEdge){V, E_VALS}

    @test edgetype(g) == ET_should_be
    @test edgetype(typeof(g)) == ET_should_be
    @test edgetype(G{V, E_VALS}) == ET_should_be

end

@testset "vertices" begin
    @testset "vertices($G{$G, $E_VALS}($n))" for
        G in (ValGraph, ValOutDiGraph, ValDiGraph),
        V in TEST_VERTEX_TYPES_SMALL,
        E_VALS in TEST_EDGEVAL_TYPES_SMALL,
        n in [0, 1, 2, 3, 10, 11]

        g = G{V, E_VALS}(n)
        @assert g isa G{V, E_VALS}

        @test vertices(g) == Base.OneTo(n)
        @test eltype(vertices(g)) == V
    end

    @testset "has_vertex($G{$G, $E_VALS}($n), $u))" for
        G in (ValGraph, ValOutDiGraph, ValDiGraph),
        V in TEST_VERTEX_TYPES_SMALL,
        E_VALS in TEST_EDGEVAL_TYPES_SMALL,
        n in [0, 1, 2, 3, 10, 11],
        u in [-1, 0, 1, 2, 4, 12, 127]

        g = G{V, E_VALS}(n)
        @assert g isa G{V, E_VALS}

        @test has_vertex(g, u) == ((u >= 1) && (u <= n))
    end
end

function xshow(io::IO, g::AbstractValGraph{V, E_VALS}) where {V, E_VALS}
    nvg = Int(nv(g))
    neg = Int(nv(g))
    dir = is_directed(g) ? "directed" : "undirected"
    name = string(nameof(typeof(g)))

    types = tuple_of_types(E_VALS)

    edgevalues_string = if g isa ZeroValGraph
        "with no edge values"
    elseif g isa OneValGraph
        if has_named_edgevals(g)
            "with named edge values of type $types"
        else
            "with edge values of type $types"
        end
    else
        if has_named_edgevals(g)
            "with multiple named edge values of types $types"
        else
            "with multiple edge values of types $types"
        end

    end

    println(io, "{$nvg, $neg} $dir $name{$V} graph $edgevalues_string.")
end

function test_show(x, expected::String)
    io = IOBuffer()
    show(io, x)
    @test String(take!(io)) == expected
end

# TODO use testset with different parameters
@testset "show" begin

    test_show(ValGraph{Int8}(10; edgeval_types=()), "{10, 0} undirected ValGraph{Int8} graph with no edge values.\n")
    test_show(ValDiGraph{Int64}(0; edgeval_types=NamedTuple()), "{0, 0} directed ValDiGraph{Int64} graph with no edge values.\n")
    test_show(ValOutDiGraph{UInt32}(0; edgeval_types=NamedTuple()), "{0, 0} directed ValOutDiGraph{UInt32} graph with no edge values.\n")

    test_show(ValGraph(undef, cycle_graph(Int16(3)), edgeval_types=(Float32,)),
              "{3, 3} undirected ValGraph{Int16} graph with edge values of type (Float32,).\n")

    test_show(ValOutDiGraph(undef, path_digraph(Int32(5)), edgeval_types=(label = String,)),
              "{5, 4} directed ValOutDiGraph{Int32} graph with named edge values of type (label = String,).\n")

    test_show(ValGraph(undef, complete_graph(UInt64(6)), edgeval_types=(String, Char)),
              "{6, 15} undirected ValGraph{UInt64} graph with multiple edge values of types (String, Char).\n")

    test_show(ValDiGraph(undef, complete_digraph(UInt64(6)), edgeval_types=(a=Bool, b=Tuple{Bool, Bool}, c=Char)),
              "{6, 30} directed ValDiGraph{UInt64} graph with multiple named edge values of types (a = Bool, b = $(Tuple{Bool,Bool}), c = Char).\n")

end


text:
@testset "Regressions" begin

# bug 137
@test isempty(lintstr("""
immutable Test137
   Test137(x) = x ? new() : error("constructor must be true")
end
"""))

# bug 88
@test isempty(lintstr("""
F = Float64
Complex{F}

type MyType{T}
a::T
end

f{A}(z::Complex{A}) = z
g{B}(z::MyType{B}) = z
"""))

# bug 122
@test isempty(lintstr("""
foo{S<:AbstractString}(args::Array{S}=ARGS) = typeof(args)
"""))

# bug 56, 84
@test isempty(lintstr("""
foo!(args::AbstractArray) = (args.x = 1)
"""))

@test messageset(lintstr("""
foo!(args::Complex) = (args.re = 1)
""")) == Set([:E525])

# bug 171
@test isempty(lintstr("""
using Base.iteratorsize
"""))

# bug 135
@test isempty(lintstr("""
import Base: ==, hash, length, size, -, +, ./, norm, dot, angle, cross, vec,
             any, print, show, parse
"""))

# bug 81
@test isempty(lintstr("""
A = [[1, 2], [3, 4]]
println(A[1][1])
"""))

# bug 180
@test messageset(lintstr("""
x, y = error()
""")) == Set([:E539])

# bug 166
@test isempty(lintstr("""
let x = :(type X end); x; end
"""))

# bug 164
@test messageset(lintstr("""
undefined()
""")) == Set([:E321])

# bug 215
@test isempty(lintstr("""
x = (rand(Bool) ? (1,) : (1.0,),)
for (i,) in x
    @show i
end
"""))

# bug 192
@test isempty(lintstr("""
type MyType end
@recipe function f(t::MyType)
    markersize --> 10
    seriestype := :scatter
    randn(10)
end
"""))

# bug 209
@test isempty(lintstr("""
function f{N}(::Array{Int,N})
    s = 0
    for j in 1:N
        s += j
    end
    s
end
"""))

# bug 187
@test isempty(lintstr("""
module Test187
import Compat.Iterators: flatten
export flatten
end
"""))

# bug 221
@test isempty(lintstr("""
@show quote
    include(string(VERSION))
end
"""))

# bug 219
@test isempty(lintstr("""
macro foo(x, y=100)
    quote
        \$(esc(x)) + \$(esc(y))
    end
end
"""))

end


text:
# Test other equivalences of prox operations which are not covered by calculus rules

using Random
using SparseArrays

Random.seed!(0)

################################################################################
### testing consistency of simplex/L1 ball
################################################################################
# Inspired by Condat, "Fast projection onto the simplex and the l1 ball", Mathematical Programming, 158:575–585, 2016.
# See Prop. 2.1 there and following remarks.

n = 20
N = 10

# projecting onto the L1 ball
for i = 1:N
  x = randn(n)
  r = 5*rand()
  f = IndSimplex(r)
  g = IndBallL1(r)

  y1, fy1 = prox(f, abs.(x))
  y1 = sign.(x).*y1
  y2, gy2 = prox(g, x)

  @test y1 ≈ y2
end

# projecting onto the simplex
for i = 1:N
  x = randn(n)
  r = 5*rand()
  f = IndSimplex(r)
  g = IndBallL1(r)

  y1, fy1 = prox(f, x)
  y2, gy2 = prox(g, x .- minimum(x) .+ r./n)

  @test y1 ≈ y2
end

################################################################################
### testing consistency of hinge loss/box indicator
################################################################################

n = 20
N = 10

# test using Moreau identity: prox(f, x, gamma) = x - gamma*prox(f*, x/gamma, 1/gamma)
# and a couple of other calculus rules in the case b = ±ones(n)
#
# f(x) = max(0, x) is conjugate to h = IndBox(0,1)
# g(x) = HingeLoss(b, mu)(x) = mu*f(1-b.*x)
# prox(g, x, gamma) = (prox(f, 1-b.*x, mu*gamma) - 1)./(-b)
#   = x + mu*gamma*prox(h, (1-b.*x)/(mu*gamma), 1/(mu*gamma))./b

b = sign.(randn(n))
mu = 0.1+rand()
g = HingeLoss(b, mu)
h = IndBox(0, 1)

for i = 1:N
  x = randn(n)
  gamma = 0.1 + rand()
  y1, ~ = prox(g, x, gamma)
  z, ~ = prox(h, (1 .- b.*x)./(mu*gamma), 1/(mu*gamma))
  y2 = mu*gamma*(z./b) + x
  @test y1 ≈ y2
end

################################################################################
### testing regularize
################################################################################

lambda = rand()
rho = rand()
g = Regularize(NormL1(lambda),rho)

x = randn(10)
y,f = prox(g,x)
y2,f2 = prox(ElasticNet(lambda,rho),x)

@test f ≈ f2
@test y ≈ y2

################################################################################
### testing IndAffine
################################################################################

A = sprand(50,100, 0.1)
b = randn(50)

g1 = IndAffine(A, b)
g2 = IndAffine(Matrix(A), b)

x = randn(100)
y1, f1 = prox(g1, x)
y2, f2 = prox(g2, x)

@test f1 ≈ f2
@test y1 ≈ y2