text:
module Materials

export material, @material

import TOML

"""
    material(pkgname, name)

Find and return a material `name` from `pkgname`.
"""
function material(pkgname, name)
    return load_material(string(pkgname), Symbol(name)).value
end

"""
    @material pkgname names...

Load materials `names` from `pkgname` into the current scope.
"""
macro material(pkgname, names...)
    vals = Expr[]
    for name in names
        m = load_material(string(pkgname), name)
        push!(vals, :($(m.name) = $(m.value); @doc $(m.description) $(m.name)))
        info("load $(m.name) as $(summary(m.value))")
    end
    esc(quote
        $(vals...)
        nothing
    end)
end

"""
    materialize(pkgname)

Build materials of `pkgname`.
"""
function materialize(pkgname)
    index, dir = load_index(string(pkgname))
    cd(dir) do
        for material in index["materials"]
            name = material["name"]
            info("materializing $(name)")
            if haskey(material, "build")
                build = material["build"]
                run(`julia -e $(build)`)
            end
        end
    end
end

struct Material
    name::Symbol
    description::String
    value::Any
end

function load_material(pkgname, name::Symbol)::Material
    index, dir = load_index(pkgname)
    cd(dir) do
        for material in index["materials"]
            if material["name"] != String(name)
                continue
            end
            description = material["description"]
            action = get(material, "action", "filepath")
            if action isa AbstractString
                action = [action]
            end
            action = first(action)
            if action == "filepath"
                return Material(name, description, joinpath(dir, material["filename"]))
            elseif action == "deserialize"
                return Material(name, description, open(deserialize, material["filename"]))
            else
                throw(ArgumentError("unknown action: '$(action)'"))
            end
        end
        throw(ArgumentError("cannot find material named '$(name)'"))
    end
end

function load_index(pkgname)
    dir = Pkg.dir(pkgname, "materials")
    if !isdir(dir)
        throw(ArgumentError("package $(pkgname) does not have materials"))
    end
    indexfile = joinpath(dir, "index.toml")
    if !isfile(indexfile)
        throw(ArgumentError("index.toml does not exist"))
    end
    return TOML.parsefile(indexfile), dir
end

end # module


text:
# ROUTINES FOR ESTIMATING THE STEP LENGTH FOR THE STOCHASTIC APPROXIMATION ALGORITHM

# get the stochastic gradient of the first-stage objective function wrt the first-stage variables for GIVEN SCENARIOS
function getStochasticGradientForFixedSample(p0::Array{Float64},rp::Array{Float64},rm::Array{Float64},alpha::Array{Float64},
						omega_d::SubArray{Float64,2,Array{Float64,2},Tuple{Base.Slice{Base.OneTo{Int64}},UnitRange{Int64}},true},
						omega_w::SubArray{Float64,2,Array{Float64,2},Tuple{Base.Slice{Base.OneTo{Int64}},UnitRange{Int64}},true},
						tau_sat::Float64,gamma_res::Float64,tau_pos::Float64,gamma_line::Float64,delta_line::Float64)

	numSamples::Int64 = size(omega_d,2)

	stochgrad_p0 = zeros(Float64,numGen)
	stochgrad_rp = zeros(Float64,numGen)
	stochgrad_rm = zeros(Float64,numGen)
	stochgrad_alpha = zeros(Float64,numGen)

	for iter = 1:numSamples
	# first, solve the recourse system
		pt, p, slack, theta = solveRecourseProblem(p0,alpha,omega_d[:,iter],omega_w[:,iter],true,tau_sat)

	# then, compute partial derivatives of recourse variables wrt first-stage variables at above solution. ignore the ones we know to be zero
		dp_dp0 = getSensitivities(alpha,pt,omega_d[:,iter],tau_sat)

	# compute stochastic gradient contribution from this scenario
	# add on the derivatives from the penalties for positive, negative reserves for generator i
		derivPosViol = zeros(Float64,numGen)
		derivNegViol = zeros(Float64,numGen)
		for i = 1:numGen
			posResViol::Float64 = p[i]-p0[i]-rp[i]
			derivPosViol[i] = reserveSmoothPenDeriv(posResViol,posResCost[i],gamma_res,tau_pos)

			negResViol::Float64 = p0[i]-p[i]-rm[i]
			derivNegViol[i] = reserveSmoothPenDeriv(negResViol,negResCost[i],gamma_res,tau_pos)
		end

		sumDemandDeviations::Float64 = sum(omega_d[j,iter] for j = 1:numLoads) - sum(meanDemand[j] for j = 1:numLoads)
		stochgrad_p0_iter1 = (dp_dp0')*(derivPosViol-derivNegViol)
		stochgrad_p0 += stochgrad_p0_iter1 + (derivNegViol - derivPosViol)
		stochgrad_alpha += stochgrad_p0_iter1*(sumDemandDeviations + slack)
		stochgrad_rp -= derivPosViol
		stochgrad_rm -= derivNegViol

	# add on the derivatives from the line limit penalties for line k
		derivLineViol = zeros(Float64,numLines)
		for k = 1:numLines
			lineFlow::Float64 = beta[k]*(theta[lines[k][1]] - theta[lines[k][2]])
			derivLineViol[k] = linePenDeriv(lineFlow,flow_max[k],gamma_line,delta_line)
		end

		stochgrad_p0_iter2 = (dp_dp0')*(Line_mat_times_B_over_G_trans*(derivLineViol.*beta))
		stochgrad_p0 += stochgrad_p0_iter2
		stochgrad_alpha += stochgrad_p0_iter2*(sumDemandDeviations + slack)
	end

	stochgrad_p0 /= numSamples
	stochgrad_rp /= numSamples
	stochgrad_rm /= numSamples
	stochgrad_alpha /= numSamples

# add the derivatives from the first-stage generation costs
	stochgrad_p0 += genCost
	stochgrad_rp += posResCost
	stochgrad_rm += negResCost

	return stochgrad_p0, stochgrad_rp, stochgrad_rm, stochgrad_alpha
end

# estimate Lipschitz constant of the gradient of the objective function
function estimateLipschitzConstant(p0_ref::Array{Float64},rp_ref::Array{Float64},rm_ref::Array{Float64},alpha_ref::Array{Float64},
									numSamples::Int64,numGradientSamples::Int64,tau_sat::Float64,gamma_res::Float64,
									tau_pos::Float64,gamma_line::Float64,delta_line::Float64)

# generate scenarios of the random variables
	omega_d, omega_w = generateScenarios(numGradientSamples*numSamples)

# four sets of Lipschitz constant estimates for the four types of first-stage variables
	L_max = zeros(Float64,4)
	for iter = 1:numSamples
	# sample points in the first-stage feasible region that are close to the reference point
		p0_1_int, rp_1_int, rm_1_int, alpha_1_int = pickPointOnSphere(p0_ref,rp_ref,rm_ref,alpha_ref)
		p0_1, rp_1, rm_1, alpha_1, solve_status_1 = projectOntoFeasibleSet(p0_1_int,rp_1_int,rm_1_int,alpha_1_int)

		p0_2_int, rp_2_int, rm_2_int, alpha_2_int = pickPointOnSphere(p0_ref,rp_ref,rm_ref,alpha_ref)
		p0_2, rp_2, rm_2, alpha_2, solve_status_2 = projectOntoFeasibleSet(p0_2_int,rp_2_int,rm_2_int,alpha_2_int)

		if((solve_status_1 == :Optimal || solve_status_1 == :Suboptimal) &&
			(solve_status_2 == :Optimal || solve_status_2 == :Suboptimal))
			# determine which samples of the random variables will be used to approximate the gradient of the objective function
			omega_start::Int64 = (iter-1)*numGradientSamples + 1
			omega_end::Int64 = iter*numGradientSamples

			# approximate gradient of the objective function using stochastic gradients
			@views grad_p0_1, grad_rp_1, grad_rm_1, grad_alpha_1 = getStochasticGradientForFixedSample(p0_1,rp_1,rm_1,alpha_1,
																	omega_d[:,omega_start:omega_end],omega_w[:,omega_start:omega_end],
																	tau_sat,gamma_res,tau_pos,gamma_line,delta_line)
			@views grad_p0_2, grad_rp_2, grad_rm_2, grad_alpha_2 = getStochasticGradientForFixedSample(p0_2,rp_2,rm_2,alpha_2,
																	omega_d[:,omega_start:omega_end],omega_w[:,omega_start:omega_end],
																	tau_sat,gamma_res,tau_pos,gamma_line,delta_line)

			L_est = zeros(Float64,4)
			L_est[1] = norm(grad_p0_2 - grad_p0_1)/norm(p0_2 - p0_1)
			L_est[2] = norm(grad_rp_2 - grad_rp_1)/norm(rp_2 - rp_1)
			L_est[3] = norm(grad_rm_2 - grad_rm_1)/norm(rm_2 - rm_1)
			L_est[4] = norm(grad_alpha_2 - grad_alpha_1)/norm(alpha_2 - alpha_1)

			for i = 1:4
				if(L_est[i] > L_max[i])
					L_max[i] = L_est[i]
				end
			end
		end
	end

	return L_max
end

# estimate step length for the stochastic approximation method using the simple rule in Ghadimi, Lan, and Zhang
# we estimate a different step length for each of the four sets of variables p0, rp, rm, and alpha because the problem is ill-conditioned
function estimateStepLength(p0_ref::Array{Float64},rp_ref::Array{Float64},rm_ref::Array{Float64},alpha_ref::Array{Float64},
							numGradientSamples::Int64,numSamplesForLipEst::Int64,tau_sat::Float64,gamma_res::Float64,
							tau_pos::Float64,gamma_line::Float64,delta_line::Float64)

	rho_est = estimateLipschitzConstant(p0_ref,rp_ref,rm_ref,alpha_ref,numSamplesForLipEst,numGradientSamples,tau_sat,
													gamma_res,tau_pos,gamma_line,delta_line)
# step length is an array of size four
	stepLength = 1.0./(2.0*rho_est)

	return stepLength
end

# given a reference point, pick a random point within a given radius
function pickPointOnSphere(p0_orig::Array{Float64},rp_orig::Array{Float64},rm_orig::Array{Float64},alpha_orig::Array{Float64})

# factor to determine sampling neighborhood
	scaleFactor::Float64 = 10.0

# generate point on the unit sphere in the dimension of the first-stage variables
	unifRand = rand(1)
	factor1::Float64 = (unifRand[1])^(1.0/numGen)
	normRand = rand(Normal(0,1), numGen)
	factor2::Float64 = norm(normRand)
# this is a vector
	factor3 = factor1*normRand/factor2

	p0_samp = p0_orig + (abs.(p0_orig)).*factor3/scaleFactor
	rp_samp = rp_orig + (abs.(rp_orig)).*factor3/scaleFactor
	rm_samp = rm_orig + (abs.(rm_orig)).*factor3/scaleFactor
	alpha_samp = alpha_orig + (abs.(alpha_orig)).*factor3/scaleFactor

	return p0_samp, rp_samp, rm_samp, alpha_samp
end

ext: jl

max_stars_repo_name: rohitkannan/DCOPF-reserve-saturation

max_issues_repo_name: rohitkannan/DCOPF-reserve-saturation

max_forks_repo_name: rohitkannan/DCOPF-reserve-saturation


text:
quoted(x) = typof(x) === Quote || typof(x) === Quotenode
unquoted(x) = typof(x) === UnaryOpCall && typof(x[1]) === OPERATOR && kindof(x[1]) == CSTParser.Tokens.EX_OR

function get_ids(x, q = false, ids = [])
    if quoted(x)
        q = true
    end
    if q && unquoted(x)
        q = false
    end
    if isidentifier(x)
        !q && push!(ids, x)
    elseif length(x) > 0
        for i in 1:length(x)
            get_ids(x[i], q, ids)
        end
    end
    ids
end

function collect_bindings_refs(x::EXPR, bindings = [], refs = [])
    if bindingof(x) !== nothing
        push!(bindings, x)
    end
    if StaticLint.hasref(x)
        push!(refs, x)
    end
    for a in x
        collect_bindings_refs(a, bindings, refs)
    end
    return bindings, refs
end

function remove_ref(x::EXPR)
    if hasref(x) && refof(x) isa Binding && refof(x).refs isa Vector
        for ia in enumerate(refof(x).refs)
            if ia[2] == x
                deleteat!(refof(x).refs, ia[1])
                setref!(x, nothing)
                return
            end
        end
        error()
    end
end

function clear_binding(x::EXPR)
    if bindingof(x) isa Binding
        for r in bindingof(x).refs
            if r isa EXPR
                setref!(r, nothing)
            elseif r isa Binding
                clear_binding(r)
            end
        end
        x.meta.binding = nothing
    end
end
function clear_scope(x::EXPR)
    if hasmeta(x) && scopeof(x) isa Scope
        setparent!(scopeof(x), nothing)
        empty!(scopeof(x).names)
        if typof(x) === FileH && scopeof(x).modules isa Dict && scopehasmodule(scopeof(x), :Base) && scopehasmodule(scopeof(x), :Core)
            m1, m2 = getscopemodule(scopeof(x), :Base), getscopemodule(scopeof(x), :Core)
            empty!(scopeof(x).modules)
            addmoduletoscope!(scopeof(x), m1)
            addmoduletoscope!(scopeof(x), m2)
        else
            scopeof(x).modules = nothing
        end
    end
end

function clear_ref(x::EXPR)
    if refof(x) isa Binding
        if refof(x).refs isa Vector
            for i in 1:length(refof(x).refs)
                if refof(x).refs[i] == x
                    deleteat!(refof(x).refs, i)
                    break
                end
            end
        end
        setref!(x, nothing)
    elseif refof(x) !== nothing
        setref!(x, nothing)
    end
end
function clear_error(x::EXPR)
    if hasmeta(x) && x.meta.error !== nothing
        x.meta.error = nothing
    end
end
function clear_meta(x::EXPR)
    clear_binding(x)
    clear_ref(x)
    clear_scope(x)
    clear_error(x)
    for a in x
        clear_meta(a)
    end
end

function get_root_method(b, server)
    return b
end

function get_root_method(b::Binding, server, b1 = nothing, visited_bindings = Binding[])
    if b.prev === nothing || b == b.prev || !(b.prev isa Binding) || b in visited_bindings
        return b
    end
    push!(visited_bindings, b)
    if b.type == b.prev.type == CoreTypes.Function
        return get_root_method(b.prev, server, b, visited_bindings)
    elseif b.type == CoreTypes.Function && b.prev.type == CoreTypes.DataType
        return b.prev
    else
        return b
    end
end

function retrieve_delayed_scope(x)
    if (CSTParser.defines_function(x) || CSTParser.defines_macro(x)) && scopeof(x) !== nothing
        if parentof(scopeof(x)) !== nothing
            return parentof(scopeof(x))
        else
            return scopeof(x)
        end
    elseif typof(x) === Export
        return retrieve_scope(x)
    elseif parentof(x) !== nothing
        return retrieve_delayed_scope(parentof(x))
    end
    return nothing
end

function retrieve_scope(x)
    if scopeof(x) !== nothing
        return scopeof(x)
    elseif parentof(x) isa EXPR
        return retrieve_scope(parentof(x))
    end
    return
end

function find_return_statements(x::EXPR)
    rets = EXPR[]
    if CSTParser.defines_function(x)
        find_return_statements(x[3], true, rets)
    end
    return rets
end

function find_return_statements(x::EXPR, last_stmt, rets)
    if last_stmt && !(typof(x) === CSTParser.Block || typof(x) === CSTParser.If || typof(x) === CSTParser.KEYWORD)
        push!(rets, x)
        return rets, false
    end

    if typof(x) === CSTParser.Return
        push!(rets, x)
        return rets, true
    end

    for i = 1:length(x)
        _, stop_iter = find_return_statements(x[i], last_stmt && (i == length(x) || (typof(x) === CSTParser.If && typof(x[i]) === CSTParser.Block)), rets)
        stop_iter && break
    end
    return rets, false
end

function _expr_assert(x::EXPR, typ, nargs)
    typof(x) == typ && length(x) == nargs
end

function _binary_assert(x, kind)
    typof(x) === CSTParser.BinaryOpCall && length(x) == 3 && typof(x[2]) === CSTParser.OPERATOR && kindof(x[2]) === kind
end

# should only be called on Bindings to functions
function last_method(func::Binding)
    if func.next isa Binding && (func.next.type === CoreTypes.Function || (func.next.type === CoreTypes.DataType && func.type === CoreTypes.Function))
        return func.next
    else
        return func
    end
end

function prev_method(func::Binding)
    if func.prev isa Binding
        if func.prev.type === CoreTypes.Function
            return func.prev
        elseif func.prev.type === CoreTypes.DataType && func.prev.val isa EXPR && CSTParser.defines_struct(func.prev.val)
            return func.prev
        elseif (func.prev.val isa SymbolServer.FunctionStore || func.prev.val isa SymbolServer.DataTypeStore) && length(func.prev.val.methods) > 0
            return func.prev.val, 1
        end
    elseif (func.prev isa SymbolServer.FunctionStore || func.prev isa SymbolServer.DataTypeStore) && length(func.prev.methods) > 0
        return func.prev, 1
    end
    return nothing
end

function prev_method(func_iter::Tuple{T,Int}) where T <: Union{SymbolServer.FunctionStore,SymbolServer.DataTypeStore}
    func = func_iter[1]
    iter = func_iter[2]
    if iter < length(func.methods)
        return func, iter + 1
    else
        return nothing
    end
end

function find_exported_names(x::EXPR)
    exported_vars = EXPR[]
    for i in 1:length(x[3])
        expr = x[3][i]
        if typof(expr) == CSTParser.Export &&
            for j = 2:length(expr)
                if CSTParser.isidentifier(expr[j]) && hasref(expr[j])
                    push!(exported_vars, expr[j])
                end
            end
        end
    end
    return exported_vars
end

"""
    is_in_fexpr(x::EXPR, f)
Check whether `x` isa the child of an expression for which `f(parent) == true`.
"""
function is_in_fexpr(x::EXPR, f)
    if f(x)
        return true
    elseif parentof(x) isa EXPR
        return is_in_fexpr(parentof(x), f)
    else
        return false
    end
end

"""
    get_in_fexpr(x::EXPR, f)
Get the `parent` of `x` for which `f(parent) == true`. (is_in_fexpr should be called first.)
"""
function get_parent_fexpr(x::EXPR, f)
    if f(x)
        return x
    elseif parentof(x) isa EXPR
        return get_parent_fexpr(parentof(x), f)
    end
end

hasreadperm(p::String) = (uperm(p) & 0x04) == 0x04

# check whether a path is in (including subfolders) the julia base dir. Returns "" if not, and the path to the base dir if so.
function _is_in_basedir(path::String)
    i = findfirst(r".*base", path)
    i == nothing && return ""
    path1 = path[i]::String
    !hasreadperm(path1) && return ""
    !isdir(path1) && return ""
    files = readdir(path1)
    if all(f -> f in files, ["Base.jl", "coreio.jl", "essentials.jl", "exports.jl"])
        return path1
    end
    return ""
end

isexportedby(k::Symbol, m::SymbolServer.ModuleStore) = haskey(m, k) && ((m[k] isa SymbolServer.SymStore && m[k].exported) || k in m.exportednames)
isexportedby(k::String, m::SymbolServer.ModuleStore) = isexportedby(Symbol(k), m)
isexportedby(x::EXPR, m::SymbolServer.ModuleStore) = isexportedby(valof(x), m)
isexportedby(k, m::SymbolServer.ModuleStore) = false

function retrieve_toplevel_scope(x)
    if scopeof(x) !== nothing && (typof(x) === CSTParser.ModuleH || typof(x) === CSTParser.BareModule || typof(x) === CSTParser.FileH)
        return scopeof(x)
    elseif parentof(x) isa EXPR
        return retrieve_toplevel_scope(parentof(x))
    else
        @info "Tried to reach toplevel scope, no scope found. Final expression $(typof(x))"
        return nothing
    end
end

# b::SymbolServer.FunctionStore or DataTypeStore
# tls is a top-level Scope (expected to contain loaded modules)
# for a FunctionStore b, checks whether additional methods are provided by other packages
# f is a function that returns `true` if we want to break early from the loop

iterate_over_ss_methods(b, tls, server, f) = false
function iterate_over_ss_methods(b::SymbolServer.FunctionStore, tls::Scope, server, f)
    if b.extends in keys(getsymbolextendeds(server)) && tls.modules !== nothing
        # above should be modified,
        rootmod = SymbolServer._lookup(b.extends.parent, getsymbolserver(server)) # points to the module containing the initial function declaration
        if rootmod !== nothing && haskey(rootmod, b.extends.name) # check rootmod exists, and that it has the variable
            rootfunc = rootmod[b.extends.name]
            # find extensoions
            if haskey(getsymbolextendeds(server), b.extends) # method extensions listed
                for vr in getsymbolextendeds(server)[b.extends] # iterate over packages with extensions
                    !(SymbolServer.get_top_module(vr) in keys(tls.modules)) && continue
                    rootmod = SymbolServer._lookup(vr, getsymbolserver(server))
                    !(rootmod isa SymbolServer.ModuleStore) && continue
                    if haskey(rootmod.vals, b.extends.name) && (rootmod.vals[b.extends.name] isa SymbolServer.FunctionStore || rootmod.vals[b.extends.name] isa SymbolServer.DataTypeStore)# check package is available and has ref
                        for m in rootmod.vals[b.extends.name].methods #
                            ret = f(m)
                            ret && return true
                        end
                    end
                end
            end
        end
    else
        for m in b.methods
            ret = f(m)
            ret && return true
        end
    end
    return false
end

function iterate_over_ss_methods(b::SymbolServer.DataTypeStore, tls::Scope, server, f)
    if b.name isa SymbolServer.VarRef
        bname = b.name
    elseif b.name isa SymbolServer.FakeTypeName
        bname = b.name.name
    end
    if (bname in keys(getsymbolextendeds(server))) && tls.modules !== nothing
        # above should be modified,
        rootmod = SymbolServer._lookup(bname.parent, getsymbolserver(server)) # points to the module containing the initial function declaration
        if rootmod !== nothing && haskey(rootmod, bname.name) # check rootmod exists, and that it has the variable
            rootfunc = rootmod[bname.name]
            # find extensoions
            if haskey(getsymbolextendeds(server), bname) # method extensions listed
                for vr in getsymbolextendeds(server)[bname] # iterate over packages with extensions
                    !(SymbolServer.get_top_module(vr) in keys(tls.modules)) && continue
                    rootmod = SymbolServer._lookup(vr, getsymbolserver(server))
                    !(rootmod isa SymbolServer.ModuleStore) && continue
                    if haskey(rootmod.vals, bname.name) && (rootmod.vals[bname.name] isa SymbolServer.FunctionStore || rootmod.vals[bname.name] isa SymbolServer.DataTypeStore)# check package is available and has ref
                        for m in rootmod.vals[bname.name].methods #
                            ret = f(m)
                            ret && return true
                        end
                    end
                end
            end
        end
    else
        for m in b.methods
            ret = f(m)
            ret && return true
        end
    end
    return false
end

fcall_name(x::EXPR) = typof(x) === Call && length(x) > 0 && valof(x[1])

is_getfield(x) = x isa EXPR && typof(x) === BinaryOpCall && length(x) == 3 && kindof(x[2]) == CSTParser.Tokens.DOT

is_getfield_w_quotenode(x) = x isa EXPR && typof(x) === BinaryOpCall && length(x) == 3 && kindof(x[2]) == CSTParser.Tokens.DOT && typof(x[3]) === CSTParser.Quotenode && length(x[3]) > 0


text:
# Cached version
mutable struct NDMBatchedCache{T,VT,VCT} <: NNBatchedCache{NDM}
    θλ_σ::VT
    θμ_σ::VT
    θλ_σp::VT
    θμ_σp::VT

    θλ_σ_tmp::VT
    θμ_σ_tmp::VT
    θλ_σp_tmp::VT
    θμ_σp_tmp::VT

    σr::VT
    σc::VT
    ∑σ::VT
    Δσ::VT

    ∑logℒ_λ_σ::T
    ∑logℒ_μ_σ::T
    ∑logℒ_λ_σp::T
    ∑logℒ_μ_σp::T
    ∂logℒ_λ_σ::VT
    ∂logℒ_μ_σ::VT
    ∂logℒ_λ_σp::VT
    ∂logℒ_μ_σp::VT

    Γ_λ::T
    Γ_μ::T
    Π::VCT

    _Π::VCT
    _Π2::VCT
    _Π_tmp::VT
    ∂logℒ_Π::VCT

    σ_row_cache::VT
    i_σ_row_cache::Int

    valid::Bool # = false
end

cache(net::NDM, batch_sz) = begin
    n_h = length(net.h_μ)
    n_v = length(net.b_μ)
    n_a = length(net.d_λ)
    CT  = complex(eltype(net.d_λ))

    NDMBatchedCache(similar(net.h_μ, n_h, batch_sz),
              similar(net.h_μ, n_h, batch_sz),
              similar(net.h_μ, n_h, batch_sz),
              similar(net.h_μ, n_h, batch_sz),

              similar(net.h_μ, n_h, batch_sz),
              similar(net.h_μ, n_h, batch_sz),
              similar(net.h_μ, n_h, batch_sz),
              similar(net.h_μ, n_h, batch_sz),

              similar(net.b_μ, n_v, batch_sz),
              similar(net.b_μ, n_v, batch_sz),
              similar(net.b_μ, n_v, batch_sz),
              similar(net.b_μ, n_v, batch_sz),

              similar(net.b_μ, 1, batch_sz),
              similar(net.b_μ, 1, batch_sz),
              similar(net.b_μ, 1, batch_sz),
              similar(net.b_μ, 1, batch_sz),
              similar(net.h_μ, n_h, batch_sz),
              similar(net.h_μ, n_h, batch_sz),
              similar(net.h_μ, n_h, batch_sz),
              similar(net.h_μ, n_h, batch_sz),

              similar(net.b_μ, 1, batch_sz),
              similar(net.b_μ, 1, batch_sz),
              similar(net.b_μ, CT, 1, batch_sz),

              similar(net.d_λ, CT, n_a, batch_sz),
              similar(net.d_λ, CT, n_a, batch_sz),
              similar(net.d_λ,     n_a, batch_sz),
              similar(net.d_λ, CT, n_a, batch_sz),

              similar(net.b_μ, n_v, batch_sz),
              -1,

              false)
end

batch_size(c::NDMBatchedCache) = size(c.σr, 2)

function Base.show(io::IO, m::NDMBatchedCache)
    print(io, "NDMBatchedCache with batch-size = $(batch_size(m))")
end

function logψ!(out::AbstractMatrix, W::NDM, c::NDMBatchedCache, σr_r::AbstractMatrix, σc_r::AbstractMatrix)
    ∑σ      = c.∑σ
    Δσ      = c.Δσ
    θλ_σ    = c.θλ_σ
    θμ_σ    = c.θμ_σ
    θλ_σp   = c.θλ_σp
    θμ_σp   = c.θμ_σp
    _Π      = c._Π
    _Π2     = c._Π2
    _Π_tmp  = c._Π_tmp
    T       = eltype(c.θλ_σ)

    σr = c.σr; copy!(σr, σr_r)
    σc = c.σc; copy!(σc, σc_r)

    if !c.valid || c.σ_row_cache ≠ σr
        c.σ_row_cache .= σr
        c.valid = true

        # θλ_σ .= W.h_λ + W.w_λ*σr
        mul!(θλ_σ, W.w_λ, σr)
        θλ_σ .+= W.h_λ

        # θμ_σ .= W.h_μ + W.w_μ*σr
        mul!(θμ_σ, W.w_μ, σr)
        θμ_σ .+= W.h_μ

        c.θλ_σ_tmp .= W.f.(θλ_σ)
        c.∑logℒ_λ_σ .= 0.0
        Base.mapreducedim!(identity, +, c.∑logℒ_λ_σ, c.θλ_σ_tmp)

        c.θμ_σ_tmp .= W.f.(θμ_σ)
        c.∑logℒ_μ_σ .= 0.0
        Base.mapreducedim!(identity, +, c.∑logℒ_μ_σ, c.θμ_σ_tmp)

        #TODO FIX
        c.∂logℒ_λ_σ .= fwd_der.(W.f, θλ_σ)
        c.∂logℒ_μ_σ .= fwd_der.(W.f, θμ_σ)
    end

    ∑logℒ_λ_σ = c.∑logℒ_λ_σ
    ∑logℒ_μ_σ = c.∑logℒ_μ_σ
    ∂logℒ_λ_σ = c.∂logℒ_λ_σ
    ∂logℒ_μ_σ = c.∂logℒ_μ_σ

    ∑σ .= σr .+ σc
    Δσ .= σr .- σc

    #θλ_σp .= W.h_λ + W.w_λ*σc
    #θμ_σp .= W.h_μ + W.w_μ*σc
    mul!(θλ_σp, W.w_λ, σc)
    θλ_σp .+= W.h_λ
    mul!(θμ_σp, W.w_μ, σc)
    θμ_σp .+= W.h_μ

    c.θλ_σp_tmp .= W.f.(θλ_σp)
    c.∑logℒ_λ_σp .= 0.0
    Base.mapreducedim!(identity, +, c.∑logℒ_λ_σp, c.θλ_σp_tmp)

    c.θμ_σp_tmp .= W.f.(θμ_σp)
    c.∑logℒ_μ_σp .= 0.0
    Base.mapreducedim!(identity, +, c.∑logℒ_μ_σp, c.θμ_σp_tmp)

    mul!(_Π_tmp, W.u_λ, ∑σ)
    _Π .= 0.5 .* _Π_tmp .+ W.d_λ
    mul!(_Π_tmp, W.u_μ, Δσ)
    _Π .+= T(0.5)im.* _Π_tmp

    Γ_λ = c.Γ_λ .=0.0
    Γ_μ = c.Γ_μ .=0.0
    mul!(Γ_λ, transpose(W.b_λ), ∑σ)
    Γ_λ .= T(0.5) .* (Γ_λ .+ c.∑logℒ_λ_σ .+ c.∑logℒ_λ_σp)
    mul!(Γ_μ, transpose(W.b_μ), Δσ)
    Γ_μ .= T(0.5) .* (Γ_μ .+ c.∑logℒ_μ_σ .- c.∑logℒ_μ_σp)

    _Π .= W.f.(_Π)
    Π = c.Π .=0.0
    Base.mapreducedim!(identity, +, Π, _Π)

    out .= Γ_λ .+ T(1.0)im .* Γ_μ .+ Π
    return out
end

function logψ_and_∇logψ!(∇logψ, out::AbstractMatrix, W::NDM, c::NDMBatchedCache, σr_r::AbstractMatrix, σc_r::AbstractMatrix)
    ∑σ      = c.∑σ
    Δσ      = c.Δσ
    θλ_σ    = c.θλ_σ
    θμ_σ    = c.θμ_σ
    θλ_σp   = c.θλ_σp
    θμ_σp   = c.θμ_σp
    _Π      = c._Π
    _Π2     = c._Π2
    _Π_tmp  = c._Π_tmp
    T       = eltype(c.θλ_σ)

    σr = c.σr; copy!(σr, σr_r)
    σc = c.σc; copy!(σc, σc_r)

    if !c.valid || c.σ_row_cache ≠ σr
        c.σ_row_cache .= σr
        c.valid = true

        # θλ_σ .= W.h_λ + W.w_λ*σr
        mul!(θλ_σ, W.w_λ, σr)
        θλ_σ .+= W.h_λ

        # θμ_σ .= W.h_μ + W.w_μ*σr
        mul!(θμ_σ, W.w_μ, σr)
        θμ_σ .+= W.h_μ

        c.θλ_σ_tmp .= W.f.(θλ_σ)
        c.∑logℒ_λ_σ .= 0.0
        Base.mapreducedim!(identity, +, c.∑logℒ_λ_σ, c.θλ_σ_tmp)

        c.θμ_σ_tmp .= W.f.(θμ_σ)
        c.∑logℒ_μ_σ .= 0.0
        Base.mapreducedim!(identity, +, c.∑logℒ_μ_σ, c.θμ_σ_tmp)

        c.∂logℒ_λ_σ .= fwd_der.(W.f, θλ_σ)
        c.∂logℒ_μ_σ .= fwd_der.(W.f, θμ_σ)
    end

    ∑logℒ_λ_σ = c.∑logℒ_λ_σ
    ∑logℒ_μ_σ = c.∑logℒ_μ_σ
    ∂logℒ_λ_σ = c.∂logℒ_λ_σ
    ∂logℒ_μ_σ = c.∂logℒ_μ_σ

    ∑σ .= σr .+ σc
    Δσ .= σr .- σc

    #θλ_σp .= W.h_λ + W.w_λ*σc
    #θμ_σp .= W.h_μ + W.w_μ*σc
    mul!(θλ_σp, W.w_λ, σc)
    θλ_σp .+= W.h_λ
    mul!(θμ_σp, W.w_μ, σc)
    θμ_σp .+= W.h_μ

    c.θλ_σp_tmp .= W.f.(θλ_σp)
    c.∑logℒ_λ_σp .= 0.0
    Base.mapreducedim!(identity, +, c.∑logℒ_λ_σp, c.θλ_σp_tmp)

    c.θμ_σp_tmp .= W.f.(θμ_σp)
    c.∑logℒ_μ_σp .= 0.0
    Base.mapreducedim!(identity, +, c.∑logℒ_μ_σp, c.θμ_σp_tmp)

    mul!(_Π_tmp, W.u_λ, ∑σ)
    _Π .= 0.5 .* _Π_tmp .+ W.d_λ
    mul!(_Π_tmp, W.u_μ, Δσ)
    _Π .+= T(0.5)im.* _Π_tmp

    Γ_λ = c.Γ_λ .=0.0
    Γ_μ = c.Γ_μ .=0.0
    mul!(Γ_λ, transpose(W.b_λ), ∑σ)
    Γ_λ .= T(0.5) .* (Γ_λ .+ c.∑logℒ_λ_σ .+ c.∑logℒ_λ_σp)
    mul!(Γ_μ, transpose(W.b_μ), Δσ)
    Γ_μ .= T(0.5) .* (Γ_μ .+ c.∑logℒ_μ_σ .- c.∑logℒ_μ_σp)

    _Π2 .= W.f.(_Π)
    Π = c.Π .=0.0
    Base.mapreducedim!(identity, +, Π, _Π2)

    out .= Γ_λ .+ T(1.0)im .* Γ_μ .+ Π
    # --- End common terms with computation of ψ --- #
    ∂logℒ_λ_σp = c.∂logℒ_λ_σp; ∂logℒ_λ_σp .= fwd_der.(W.f, θλ_σp)
    ∂logℒ_μ_σp = c.∂logℒ_μ_σp; ∂logℒ_μ_σp .= fwd_der.(W.f, θμ_σp)
    ∂logℒ_Π    = c.∂logℒ_Π;    ∂logℒ_Π    .= fwd_der.(W.f, _Π)

    # Store the derivatives
    ∇logψ.b_λ .= T(0.5)   .* ∑σ
    ∇logψ.b_μ .= T(0.5)im .* Δσ

    ∇logψ.h_λ .= T(0.5)   .* (∂logℒ_λ_σ .+ ∂logℒ_λ_σp)
    ∇logψ.h_μ .= T(0.5)im .* (∂logℒ_μ_σ .- ∂logℒ_μ_σp)

    #∇logψ.w_λ .= T(0.5)   .* (∂logℒ_λ_σ.*transpose(σr) .+ ∂logℒ_λ_σp.*transpose(σc))
    #∇logψ.w_μ .= T(0.5)im .* (∂logℒ_μ_σ.*transpose(σr) .- ∂logℒ_μ_σp.*transpose(σc))
    _batched_outer_prod_∑!(∇logψ.w_λ, T(0.5), ∂logℒ_λ_σ, σr, ∂logℒ_λ_σp, σc)
    _batched_outer_prod_Δ!(∇logψ.w_μ, T(0.5)im, ∂logℒ_μ_σ, σr, ∂logℒ_μ_σp, σc)

    ∇logψ.d_λ .= ∂logℒ_Π
#    ∇logψ.u_λ .= T(0.5) .* ∂logℒ_Π .* transpose(∑σ)
#    ∇logψ.u_μ .= T(0.5)im .*  ∂logℒ_Π .* transpose(Δσ)
    _batched_outer_prod!(∇logψ.u_λ, T(0.5), ∂logℒ_Π, ∑σ)
    _batched_outer_prod!(∇logψ.u_μ, T(0.5)im, ∂logℒ_Π, Δσ)

    return out
end


text:
for (T, args) in (
    (Mag, ((Float64, Cdouble, :d), (UInt, Culong, :ui))),
    (
     Arf, (
     (Float64, Cdouble, :d),
     (Mag, Ref{Mag}, :mag),
     (BigFloat, Ref{BigFloat}, :mpfr),
     (Int, Clong, :si),
     (UInt, Culong, :ui)),
    ),
    (
     Arb, (
     (Arf, Ref{Arf}, :arf),
     (Float64, Cdouble, :d),
     (Int, Clong, :si),
     (UInt, Culong, :ui)),
    ),
    (
     Acb, (
     (Arb, Ref{Arb}, :arb),
     (Float64, Cdouble, :d),
     (Int, Clong, :si),
     (UInt, Culong, :ui)),
    ),
)
    arbf = Symbol(cprefix(T), :_set)
    @eval begin
        function set!(res::$T, x::$T)
            ccall(@libarb($arbf), Cvoid, (Ref{$T}, Ref{$T}), res, x)
            return res
        end
    end
    for (jlT, cT, suffix) in args
        arbf = Symbol(cprefix(T), :_set_, suffix)
        @eval begin
            function set!(res::$T, x::$jlT)
                ccall(@libarb($arbf), Cvoid, (Ref{$T}, $cT), res, x)
                return res
            end
        end
    end
end

function set!(res::Arb, str::AbstractString)
    flag = ccall(@libarb(arb_set_str), Cint, (Ref{Arb}, Cstring, Int), res, str, precision(res))
    iszero(flag) || throw(ArgumentError("arblib could not parse $str as an Arb"))
    return res
end

for (jlT, cT, suffix) in (
    (Arb, Ref{Arb}, :arb_arb),
    (Float64, Cdouble, :d_d),
    (Int, Clong, :si_si),
)
    T = Acb
    arbf = Symbol(cprefix(T), :_set_, suffix)
    @eval begin
        function set!(res::$T, re::$jlT, im::$jlT)
            ccall(@libarb($arbf), Cvoid, (Ref{$T}, $cT, $cT), res, re, im)
            return res
        end
    end
end

for ArbT in (:Arf, :Arb, :Acb, :Mag)
    @eval begin
        $ArbT(t::$ArbT; prec::Integer=precision(t)) = set!($ArbT(prec=prec), t)

        function $ArbT(si::T; prec::Integer=DEFAULT_PRECISION[]) where T <: Integer
            promote_type(T, Int64) == Int64 && return $ArbT(Int64(si), prec)
            return $ArbT(BigInt(si), prec)
        end

        function $ArbT(d::T; prec::Integer=DEFAULT_PRECISION[]) where T <: AbstractFloat
            promote_type(T, Float64) == Float64 && return $ArbT(Float64(d), prec)
            return $ArbT(BigFloat(d), prec)
        end

        $ArbT(d::Float64; prec::Integer=DEFAULT_PRECISION[]) = set!($ArbT(prec=prec), d)

        if $ArbT != Arf
            $ArbT(si::Int64; prec::Integer=DEFAULT_PRECISION[]) = set!($ArbT(prec=prec), si)
        end

        Base.zero(t::$ArbT) = $ArbT(0; prec = precision(t))
        Base.one(t::$ArbT) = $ArbT(1; prec = precision(t))
    end
end

Arf(x::BigFloat; prec::Integer=precision(x)) = set!(Arf(prec=prec), x)
Arb(x::BigFloat; prec::Integer=precision(x)) = set!(Arb(prec=prec), Arf(x, prec=prec))
Acb(x::BigFloat; prec::Integer=precision(x)) = set!(Acb(prec=prec), Arb(x, prec=prec))

# fallbacks:
Arf(x::Real; prec::Integer=DEFAULT_PRECISION[]) = set!(Arf(prec=prec), BigFloat(x))
Arb(x::Real; prec::Integer=DEFAULT_PRECISION[]) = set!(Arb(prec=prec), Arf(x, prec=prec))
Acb(x::Real; prec::Integer=DEFAULT_PRECISION[]) = set!(Acb(prec=prec), Arb(x, prec=prec))

# string input
Arb(str::AbstractString; prec::Integer=DEFAULT_PRECISION[]) = set!(Arb(prec=prec), str)

function Acb(re::Integer, im::Integer; prec::Integer=DEFAULT_PRECISION[])
    promote_type(T, Int64) == Int64 && return Acb(Int64(re), Int64(im), prec=prec)
    return Acb(BigInt(re), BigInt(im), prec=prec)
end

Acb(re::Int64, im::Int64; prec::Integer=DEFAULT_PRECISION[]) = set!(Acb(prec=prec), re, im)
Acb(re::Float64, im::Float64; prec::Integer=DEFAULT_PRECISION[]) = set!(Acb(prec=prec), re, im)
Acb(re::Arb, im::Arb; prec::Integer=min(precision(re), precision(im))) = set!(Acb(prec=prec), re, im)

function Acb(z::Complex{T}; prec::Integer=DEFAULT_PRECISION[]) where T
    if promote_type(T, Float64) == Float64
        return Acb(Float64(real(z)), Float64(imag(z)), prec=prec)
    end
    return Acb(ArbReal(real(z), prec=prec), ArbReal(imag, prec=prec), prec=prec)
end

# Irrationals
for (irr, suffix) in ((:π, "pi"), (:ℯ, "e"), (:γ, "euler"))
    arbf = Symbol("arb_const_", suffix)
    jlf = Symbol("const_$suffix", "!")
    IrrT = Irrational{irr}
    @eval begin
        function $(jlf)(res::Arb)
            ccall(@libarb($arbf), Cvoid, (Ref{Arb}, Clong,), res, precision(res))
            return res
        end
        Arb(::$IrrT; prec::Integer=DEFAULT_PRECISION[])= $jlf(Arb(prec=prec))
    end
end

Acb(::Irrational{:π}; prec::Integer=DEFAULT_PRECISION[]) = const_pi!(Acb(prec=prec))

function const_pi!(res::Acb)
    ccall(@libarb(acb_const_pi), Cvoid, (Ref{Acb}, Clong,), res, precision(res=prec))
    return res
end


text:
using OrdinaryDiffEq, Test, DiffEqDevTools
using LinearAlgebra, Random

# Linear exponential solvers
A = DiffEqArrayOperator([2.0 -1.0; -1.0 2.0])
u0 = ones(2)
prob = ODEProblem(A,u0,(0.0,1.0))
sol1 = solve(prob, LinearExponential(krylov=:off))(1.0)
sol2 = solve(prob, LinearExponential(krylov=:simple))(1.0)
sol3 = solve(prob, LinearExponential(krylov=:adaptive))(1.0)
sol_analytic = exp(1.0 * Matrix(A)) * u0

@test isapprox(sol1, sol_analytic, rtol=1e-10)
@test isapprox(sol2, sol_analytic, rtol=1e-10)
@test isapprox(sol3, sol_analytic, rtol=1e-10)

# u' = A(t)u solvers

function update_func(A,u,p,t)
    A[1,1] = cos(t)
    A[2,1] = sin(t)
    A[1,2] = cos(t)*sin(t)
    A[2,2] = sin(t)^2
end
A = DiffEqArrayOperator(ones(2,2),update_func=update_func)
prob = ODEProblem(A, ones(2), (0.0, 5.))
dts = 1 ./2 .^(10:-1:1)
sol  = solve(prob,OrdinaryDiffEq.MagnusMidpoint(),dt=1/4)

dts = 1 ./2 .^(10:-1:1)
test_setup = Dict(:alg=>Vern9(),:reltol=>1e-14,:abstol=>1e-14)
sim = analyticless_test_convergence(dts,prob,MagnusMidpoint(),test_setup)
@test sim.𝒪est[:l2] ≈ 2 atol=0.2
sim = analyticless_test_convergence(dts,prob,MagnusMidpoint(krylov=true),test_setup)
@test sim.𝒪est[:l2] ≈ 2 atol=0.2

A = DiffEqArrayOperator(ones(2,2),update_func=update_func)
prob = ODEProblem(A, ones(2), (0.0, 5.))
dts = 1 ./2 .^(10:-1:1)
sol  = solve(prob,OrdinaryDiffEq.MagnusLeapfrog(),dt=1/4)

dts = 1 ./2 .^(10:-1:1)
test_setup = Dict(:alg=>Vern9(),:reltol=>1e-14,:abstol=>1e-14)
sim = analyticless_test_convergence(dts,prob,MagnusLeapfrog(),test_setup)
@test sim.𝒪est[:l2] ≈ 2 atol=0.2
sim = analyticless_test_convergence(dts,prob,MagnusLeapfrog(krylov=true),test_setup)
@test sim.𝒪est[:l2] ≈ 2 atol=0.2

function update_func(A,u,p,t)
    A[1,1] = cos(t)*u[1]
    A[2,1] = sin(t)*u[2]
    A[1,2] = -cos(t)*sin(t)
    A[2,2] = -(sin(t)^2) *u[1] * u[2]
end
A = DiffEqArrayOperator(ones(2,2),update_func=update_func)
prob = ODEProblem(A, ones(2), (0.5, 5.))
dts = 1 ./2 .^(10:-1:1)
sol  = solve(prob,OrdinaryDiffEq.LieEuler(),dt=1/4)

dts = 1 ./2 .^(10:-1:1)
test_setup = Dict(:alg=>Vern9(),:reltol=>1e-14,:abstol=>1e-14)
sim = analyticless_test_convergence(dts,prob,LieEuler(),test_setup)
@test sim.𝒪est[:l2] ≈ 1 atol=0.2
sim = analyticless_test_convergence(dts,prob,LieEuler(krylov=true),test_setup)
@test sim.𝒪est[:l2] ≈ 1 atol=0.2


text:
# only required if PyCall is not loaded allready
# pkgid = Base.PkgId(Base.UUID(0x438e738f_606a_5dbb_bf0a_cddfbfd45ab0), "PyCall")
# PC = Base.require(pkgid)

PC = PyCall;

PC.py"""
class SymStr():
    def __init__(self, *args, **kwargs):
        self.s = str(*args, **kwargs)
    def __str__(self):
        return self.s.__str__()
    def __repr__(self):
        return f'SymStr("{self.__str__()}")'
"""

sym_str_py_type = PC.py"SymStr";

PC.PyObject( s :: Symbol ) = PC.py"SymStr($(string(s)))"o
function PC.convert( ::Type{Symbol}, po :: PC.PyObject )
    sym_str = PC.pyisinstance( po, sym_str_py_type ) ? po.s : po;
    return Symbol(PC.convert(AbstractString, sym_str))
end
PC.pytype_mapping(sym_str_py_type, Symbol);

# not a good idea in general, I think, but usefull for keyword arguments:
Base.convert( :: Type{Symbol}, s :: AbstractString ) = Symbol(s);
nothing

max_stars_repo_name: manuelbb-upb/MorbitWrapper

max_issues_repo_name: manuelbb-upb/MorbitWrapper

max_forks_repo_name: manuelbb-upb/MorbitWrapper


text:
"""
Struct for Geometric Brownian Motion

		gbmProcess=GeometricBrownianMotionVec(σ::num1,μ::num2) where {num1,num2 <: Number}

Where:\n
		σ	=	volatility of the process.
		μ	=	drift of the process.
"""
mutable struct GeometricBrownianMotionVec{num <: Number, num1 <: Number, num4 <: Number, num2 <: Number, numtype <: Number} <: ItoProcess{numtype}
    σ::num
    μ::Curve{num1, num4}
    x0::num2
    function GeometricBrownianMotionVec(σ::num, μ::Curve{num1, num4}, x0::num2) where {num <: Number, num1 <: Number, num4 <: Number, num2 <: Number}
        if σ <= 0
            error("Volatility must be positive")
        else
            zero_typed = zero(num) + zero(num1) + zero(num4) + zero(num2)
            return new{num, num1, num4, num2, typeof(zero_typed)}(σ, μ, x0)
        end
    end
end

function GeometricBrownianMotion(σ::num, μ::Curve{num1, num4}, x0::num2) where {num <: Number, num1 <: Number, num4 <: Number, num2 <: Number}
    return GeometricBrownianMotionVec(σ, μ, x0)
end

export GeometricBrownianMotionVec;

function simulate!(X, mcProcess::GeometricBrownianMotionVec, mcBaseData::AbstractMonteCarloConfiguration, T::Number)
    @assert T > 0
    σ_gbm = mcProcess.σ
    mu_gbm = mcProcess.μ
    μ_bm = mu_gbm - (σ_gbm^2 / 2)
    simulate!(X, BrownianMotion(σ_gbm, μ_bm), mcBaseData, T)
    S0 = mcProcess.x0
    f(x) = S0 * exp(x)
    broadcast!(f, X, X)
    nothing
end


text:
using Test
using HMatrices
using LinearAlgebra
using Random
using StaticArrays

using HMatrices: RkMatrix

include(joinpath(HMatrices.PROJECT_ROOT,"test","testutils.jl"))

Random.seed!(1)

m = 5000
T = Float64
X    = points_on_sphere(m)
Y    = X

K = laplace_matrix(X,X)

splitter  = CardinalitySplitter(nmax=50)
Xclt      = ClusterTree(X,splitter)
Yclt      = ClusterTree(Y,splitter)
adm = StrongAdmissibilityStd(3)
comp = PartialACA(;atol=1e-10);
H = assemble_hmat(K,Xclt,Yclt;adm,comp,threads=false,distributed=false)
H_full      = Matrix(H)

##
hlu   = lu(H;atol=1e-5)
y     = rand(m)
M     = Matrix(K)
exact = M\y
approx = hlu\y
@test norm(exact - approx,Inf) < 1e-6

max_stars_repo_name: WaveProp/HMatrices

max_forks_repo_name: WaveProp/HMatrices


text:
export combineLatest

import Base: show

"""
    combineLatest(sources...; strategy = PushEach())
    combineLatest(sources::S, strategy::G = PushEach()) where { S <: Tuple, U }

Combines multiple Observables to create an Observable whose values are calculated from the latest values of each of its input Observables.
Accept optinal update strategy object.

# Arguments
- `sources`: input sources
- `strategy`: optional update strategy for batching new values together

# Examples
```jldoctest
using Rocket

latest = combineLatest(of(1), from(2:5))

subscribe!(latest, logger())
;

# output

[LogActor] Data: (1, 2)
[LogActor] Data: (1, 3)
[LogActor] Data: (1, 4)
[LogActor] Data: (1, 5)
[LogActor] Completed
```

```
using Rocket

latest = combineLatest(of(1) |> async(0), from(2:5) |> async(0), strategy = PushNew())

subscribe!(latest, logger())
;

# output

[LogActor] Data: (1, 2)
[LogActor] Completed
```

See also: [`Subscribable`](@ref), [`subscribe!`](@ref), [`PushEach`](@ref), [`PushEachBut`](@ref), [`PushNew`](@ref), [`PushNewBut`](@ref), [`PushStrategy`](@ref)
"""
function combineLatest end

combineLatest(; strategy = PushEach())                                       = error("combineLatest operator expects at least one inner observable on input")
combineLatest(args...; strategy = PushEach())                                = combineLatest(args, strategy)
combineLatest(sources::S, strategy::G = PushEach()) where { S <: Tuple, G }  = CombineLatestObservable{combined_type(sources), S, G}(sources, strategy)
combineLatest(sources::V, strategy::G = PushEach()) where { V <: Vector, G } = CombineLatestObservable{combined_type(sources), V, G}(sources, strategy)

##

struct CombineLatestInnerActor{L, W, I} <: Actor{L}
    wrapper :: W
end

Base.show(io::IO, inner::CombineLatestInnerActor{L, W, I}) where { L, W, I } = print(io, "CombineLatestInnerActor($L, $I)")

on_next!(actor::CombineLatestInnerActor{L, W, I}, data::L) where { L, W, I } = next_received!(actor.wrapper, data, Val{I}())
on_error!(actor::CombineLatestInnerActor{L, W, I}, err)    where { L, W, I } = error_received!(actor.wrapper, err, Val{I}())
on_complete!(actor::CombineLatestInnerActor{L, W, I})      where { L, W, I } = complete_received!(actor.wrapper, Val{I}())

##

struct CombineLatestActorWrapper{S, A, G, U}
    storage       :: S
    actor         :: A
    nsize         :: Int
    strategy      :: G
    updates       :: U
    subscriptions :: Vector{Teardown}
end

function CombineLatestActorWrapper(::Type{T}, actor::A, strategy::G) where { T, A, G }
    storage       = getmstorage(T)
    updates       = getustorage(T)
    nsize         = length(storage)
    subscriptions = fill!(Vector{Teardown}(undef, nsize), voidTeardown)
    return CombineLatestActorWrapper(storage, actor, nsize, strategy, updates, subscriptions)
end

push_update!(wrapper::CombineLatestActorWrapper) = push_update!(wrapper.nsize, wrapper.updates, wrapper.strategy)

dispose(wrapper::CombineLatestActorWrapper) = begin fill_cstatus!(wrapper.updates, true); foreach(s -> unsubscribe!(s), wrapper.subscriptions) end

function next_received!(wrapper::CombineLatestActorWrapper, data, index::Val{I}) where I
    setstorage!(wrapper.storage, data, index)
    vstatus!(wrapper.updates, I, true)
    ustatus!(wrapper.updates, I, true)
    if all_vstatus(wrapper.updates) && !all_cstatus(wrapper.updates)
        push_update!(wrapper)
        next!(wrapper.actor, snapshot(wrapper.storage))
    end
end

function error_received!(wrapper::CombineLatestActorWrapper, err, index::Val{I}) where I
    if !(cstatus(wrapper.updates, I))
        dispose(wrapper)
        error!(wrapper.actor, err)
    end
end

function complete_received!(wrapper::CombineLatestActorWrapper, ::Val{I}) where I
    if !all_cstatus(wrapper.updates)
        cstatus!(wrapper.updates, I, true)
        if ustatus(wrapper.updates, I)
            vstatus!(wrapper.updates, I, true)
        end
        if all_cstatus(wrapper.updates) || (!vstatus(wrapper.updates, I))
            dispose(wrapper)
            complete!(wrapper.actor)
        end
    end
end

##

@subscribable struct CombineLatestObservable{T, S, G} <: Subscribable{T}
    sources  :: S
    strategy :: G
end

function on_subscribe!(observable::CombineLatestObservable{T, S, G}, actor::A) where { T, S, G, A }
    wrapper = CombineLatestActorWrapper(T, actor, observable.strategy)
    W       = typeof(wrapper)

    for (index, source) in enumerate(observable.sources)
        @inbounds wrapper.subscriptions[index] = subscribe!(source, CombineLatestInnerActor{eltype(source), W, index}(wrapper))
        if cstatus(wrapper.updates, index) && !vstatus(wrapper.updates, index)
            dispose(wrapper)
            break
        end
    end

    if all_cstatus(wrapper.updates)
        dispose(wrapper)
    end

    return CombineLatestSubscription(wrapper)
end

##

struct CombineLatestSubscription{W} <: Teardown
    wrapper :: W
end

as_teardown(::Type{ <: CombineLatestSubscription }) = UnsubscribableTeardownLogic()

function on_unsubscribe!(subscription::CombineLatestSubscription)
    dispose(subscription.wrapper)
    return nothing
end

Base.show(io::IO, ::CombineLatestObservable{D}) where D  = print(io, "CombineLatestObservable($D)")
Base.show(io::IO, ::CombineLatestSubscription)           = print(io, "CombineLatestSubscription()")

##


text:
"""
    ValidationOk

Placeholder type if there is no validatiion error.
"""
struct ValidationOk end

"""
    ValidationOK

Placeholder if there is no validatiion error.
"""
const ValidationOK = ValidationOk()

"""
    ValidationError(msg)

Error thrown when the validatiion failed.
"""
struct ValidationError <: Exception
    msg::String
end

"""
    Construct{T}

Construct is used for serializing and deserializing objects.

## Methods

* `deserialize(cons::Construct{T}, s::IO; contextkw...)::T`
* `serialize(cons::Construct{T}, s::IO, obj::T; contextkw...)`
* `estimatesize(cons::Construct{T}; contextkw...)` - optional
"""
abstract type Construct{T} end

"""
    Construct(type)

Get default construct for type.
"""
Construct(cons::Construct) = cons

constructtype(::Construct{T}) where {T} = T
constructtype(type::Type) = type

"""
    deserialize(cons::Construct, s::IO; contextkw...)

Deserialize a stream to an object.
"""
function deserialize end

"""
    deserialize(cons::Construct, filename::AbstractString; contextkw...)

Deserialize a file to an object.
"""
deserialize(cons::Construct, filename::AbstractString; contextkw...) = open(filename, "r") do fs
    deserialize(cons, fs; contextkw...)
end

"""
    deserialize(cons::Construct, bytes::AbstractVector{UInt8}; contextkw...)

Deserialize a byte array to an object.
"""
deserialize(cons::Construct, bytes::AbstractVector{UInt8}; contextkw...) = deserialize(cons, IOBuffer(bytes); contextkw...)

"""
    deserialize(T, s::IO; contextkw...)

Deserialize a stream to an object.
"""
deserialize(t::Type, s::IO; contextkw...) = deserialize(Construct(t), s; contextkw...)

"""
    deserialize(T, filename::AbstractString; contextkw...)

Deserialize a file to an object.
"""
deserialize(t::Type, filename::AbstractString; contextkw...) = deserialize(Construct(t), filename; contextkw...)

"""
    deserialize(T, bytes::AbstractVector{UInt8}; contextkw...)

Deserialize a byte array to an object.
"""
deserialize(t::Type, bytes::AbstractVector{UInt8}; contextkw...) = deserialize(Construct(t), IOBuffer(bytes); contextkw...)

"""
    serialize(cons::Construct, s::IO, obj; contextkw...)

Serialize an object into a stream.
"""
function serialize end

"""
    serialize(cons::Construct, filename::AbstractString, obj; contextkw...)

Serialize an object to the file.
"""
serialize(cons::Construct{T}, filename::AbstractString, obj::T; contextkw...) where {T} = open(filename, "w") do fs
    serialize(cons, fs, obj; contextkw...)
end

"""
    serialize(cons::Construct, obj; contextkw...)

Serialize an object in memory (a byte array).
"""
function serialize(cons::Construct{T}, obj::T; contextkw...) where {T}
    io = IOBuffer()
    serialize(cons, io, obj; contextkw...)
    return take!(io)
end

"""
    serialize(obj, s::IO; contextkw...)

Serialize an object into a stream.
"""
serialize(s::IO, obj; contextkw...) = serialize(Construct(typeof(obj)), s, obj; contextkw...)

"""
    serialize(filename::AbstractString, obj; contextkw...)

Serialize an object to the file.
"""
serialize(filename::AbstractString, obj; contextkw...) = serialize(Construct(typeof(obj)), filename, obj; contextkw...)

"""
    serialize(obj; contextkw...)

Serialize an object in memory (a byte array).
"""
serialize(obj; contextkw...) = serialize(Construct(typeof(obj)), obj; contextkw...)

"""
    estimatesize(cons::Construct; contextkw...)

Estimate the size of the type.
"""
estimatesize(::Construct; contextkw...) = Interval(UInt(0), nothing)

"""
    estimatesize(T; contextkw...)

Estimate the size of the type.
"""
estimatesize(t::Type; contextkw...) = estimatesize(Construct(t); contextkw...)

"""
    Wrapper{TSub, T} <: Construct{T}

Base type of wrapper of `TSub`.

## Methods

* `subcon(wrapper::Wrapper{TSub, T})::Construct{TSub}`
"""
abstract type Wrapper{TSub, T} <: Construct{T} end

"""
    subcon(wrapper::Wrapper{TSub, T})::Construct{TSub}

Get sub-construct of `wrapper`.
"""
function subcon end

"""
    Adapter{TSub, T} <: Wrapper{TSub, T}

Abstract adapter type.

## Methods

* `subcon(adapter::Adapter{TSub, T})::Construct{TSub}`
* `encode(adapter::Adapter{TSub, T}, obj::T; contextkw...)`
* `decode(adapter::Adapter{TSub, T}, obj::TSub; contextkw...)`
"""
abstract type Adapter{TSub, T} <: Wrapper{TSub, T} end

"""
    encode(adapter::Adapter{TSub, T}, obj::T; contextkw...) where {TSub, T}
"""
function encode end

"""
    decode(adapter::Adapter{TSub, T}, obj::TSub; contextkw...) where {TSub, T}
"""
function decode end

function serialize(adapter::Adapter{TSub, T}, s::IO, obj::T; contextkw...) where {TSub, T}
    objsub = encode(adapter, obj; contextkw...)
    serialize(subcon(adapter), s, objsub; contextkw...)
end

function deserialize(adapter::Adapter{TSub, T}, s::IO; contextkw...) where {TSub, T}
    obj = deserialize(subcon(adapter), s; contextkw...)
    decode(adapter, obj; contextkw...)
end

estimatesize(wrapper::Adapter) = estimatesize(subcon(wrapper))

"""
    SymmetricAdapter{T} <: Adapter{T, T}

Abstract adapter type. `encode` both for serializing and deserializing.

## Methods

* `subcon(adapter::SymmetricAdapter{T})::Construct{T}`
* `encode(adapter::SymmetricAdapter{T}, obj::T; contextkw...)`
"""
abstract type SymmetricAdapter{T} <: Adapter{T, T} end

decode(adapter::SymmetricAdapter{T}, obj::T; contextkw...) where {T} = encode(adapter, obj; contextkw...)

"""
    Validator{T} <: SymmetricAdapter{T}

Abstract validator type. Validates a condition on the encoded/decoded object..

## Methods

* `subcon(validator::Validator{T})::Construct{T}`
* `validate(validator::Validator{T}, obj::T; contextkw...)::Union{ValidationOK, ValidationError}`
"""
abstract type Validator{T} <: SymmetricAdapter{T} end

"""
    validate(validator::Validator{T}, obj::T; contextkw...)::Union{ValidationOk, ValidationError}
"""
function validate end

function encode(validator::Validator{T}, obj::T; contextkw...) where {T}
    result = validate(validator, obj; contextkw...)
    if result isa ValidationError
        throw(result)
    end
    obj
end


text:
module CredentialTests

using Test
using GCP
using JSON
using ..GCPTests

@testset "Credential" begin

  @test begin
    cred = GCP.Credential(GCPTests.CredFilename)
    cred != nothing
  end
  @test begin
    j = JSON.parsefile(GCPTests.CredFilename; dicttype=Dict{Symbol, Any})
    cred = GCP.Credential(j)
    cred != nothing
  end

end

end # module CredentialTests

max_stars_repo_name: ProfLeao/mais

max_issues_repo_name: ProfLeao/mais

max_forks_repo_name: ProfLeao/mais


text:
function update_state_feature_select!(s::StateFS, c::ConstantsFS, d::DataFS,
                                      t::TunersFS;
                                      ll::Vector{Float64},
                                      fix::Vector{Symbol},
                                      use_repulsive::Bool,
                                      Z_marg_lamgam::Bool, sb_ibp::Bool,
                                      time_updates::Bool=false,
                                      verbose::Int=0,
                                      Z_thin::Int=1)

  # NOTE: `@doIf` is defined in "../util.jl"

  # Return true if parameter (sym) is not fixed
  isRandom(sym::Symbol)::Bool = !(sym in fix)

  function update_Z_!()
    for i in 1:Z_thin
      if Z_marg_lamgam  # marginalize over lambda and gamma occasionally
        update_Z_marg_lamgam!(s.theta, c.constants, d.data,
                              sb_ibp, use_repulsive=use_repulsive)
      else  # Do regular updates
        update_Z!(s.theta, c.constants, d.data, sb_ibp,
                  use_repulsive=use_repulsive)
      end
    end
  end

  # Parameter update-order should be:
  # Z -> v -> alpha ->
  # omega -> r -> lam -> W* -> gamma -> eta -> delta -> sig2 -> y*
  gibbs_sampler = [
    (:Z, () -> update_Z_!()),
    (:v, () -> update_v!(s.theta, c.constants, d.data, t.tuners, sb_ibp)),
    (:alpha, () -> update_alpha!(s.theta, c.constants, d.data, sb_ibp)),
    (:omega, () -> update_omega!(s, c, d, t)),
    (:r, () -> update_r!(s, c, d)),
    (:lam, () -> update_lam!(s.theta, c.constants, d.data)),
    (:W_star, () -> update_W_star!(s, c, d, t)),
    (:gam, () -> update_gam!(s.theta, c.constants, d.data)),
    (:eta, () -> update_eta!(s.theta, c.constants, d.data)),
    (:delta, () -> update_delta!(s.theta, c.constants, d.data)),
    (:sig2, () -> update_sig2!(s.theta, c.constants, d.data)),
    (:y_imputed, () -> update_y_imputed!(s.theta, c.constants, d.data, t.tuners)),
    # Compute loglikelihood.
    (:marg_loglike, () -> let
       _ll = compute_marg_loglike(s.theta, c.constants,
                                  d.data, c.constants.temper)
       append!(ll, _ll)
     end)
    # (:loglike, () -> append!(ll, compute_loglike(s.theta, c.constants, d.data)))
  ]

  @doIf time_updates println()
  for (param, update_param) in gibbs_sampler
    if isRandom(param)
      tic = time()
      update_param()
      toc = time()

      if time_updates
        println("update time for $(param): $(toc - tic)")
      end
    end
  end

  if verbose > 1
    println(s.theta.W)
  end
end


text:
# get_ec_number_for_gene
function get_ec_number_for_gene(gene_array::Array{String,1})::Some

    try

        # initialize -
        ec_number_array = Array{String,1}()
        for gene in gene_array

            local_ec_number_array = get_ec_number_for_gene(gene) |> check
            if (isnothing(local_ec_number_array) == false)
                for ec_number in local_ec_number_array
                    push!(ec_number_array, ec_number)
                end
            end
        end

        return Some(ec_number_array)
    catch error

        # get the original error message -
        error_message = sprint(showerror, error, catch_backtrace())
        vl_error_obj = ErrorException(error_message)

        # Package the error -
        return Some(vl_error_obj)
    end
end

function get_ec_number_for_gene(gene_location::String)::Some

    # TODO: check gene location string -

    try

        # curl -X GET http://rest.kegg.jp/link/ec/eco:b2388  => ec for this gene
        # curl -X GET http://rest.kegg.jp/link/eco/eco00010  => genes for a pathway

        # initialize -
        ec_number_array = String[]

        # formulate the url -
        url_string = "$(_KEGG_LINKAGE_EC_URL)/$(gene_location)"

        # get the data -
        ecdata = http_get_call_with_url(url_string) |> check
        if (ecdata == "\n" || isnothing(ecdata) == true)
            return Some(nothing)
        end

        # ok, parse the response which is of the form: {gene_location}\t{ec_number}\n*
        # we can have more than 1 record -
        record_array = split(ecdata, "\n")

        # process the records -
        for record in record_array

            # split along the \t -
            if (record != "")
                ec_number = string(split(record, "\t")[2])
                push!(ec_number_array, ec_number)
            end
        end

        # return this record -
        return Some(ec_number_array)

    catch error

        # get the original error message -
        error_message = sprint(showerror, error, catch_backtrace())
        vl_error_obj = ErrorException(error_message)

        # Package the error -
        return Some(vl_error_obj)
    end
end

function get_ec_numbers_for_pathway(pathway_code_array::Array{String,1})::Some

    try

        # master ec numner list -
        master_ec_number_list = Array{String,1}()

        # process the list of pathways -
        for pathway_code in pathway_code_array

            # get ec number array -
            ec_number_array = get_ec_numbers_for_pathway(pathway_code) |> check
            if (isnothing(ec_number_array) == false)
                append!(master_ec_number_list, ec_number_array)
            end
        end

        # return -
        return Some(master_ec_number_list)
    catch error

        # get the original error message -
        error_message = sprint(showerror, error, catch_backtrace())
        vl_error_obj = ErrorException(error_message)

        # Package the error -
        return Some(vl_error_obj)
    end
end

function get_ec_numbers_for_pathway(pathway_code::String)::Some

    try

        # initialize -
        ec_number_array = Array{String,1}()

        # formulate the url -
        url_string = "$(_KEGG_LINK_URL)/ec/$(pathway_code)"

        # execute the call -
        http_body_call = http_get_call_with_url(url_string) |> check
        if (http_body_call == "\n" || isnothing(http_body_call) == true)
            return Some(nothing)
        end

        # split along the new line
        record_array = split(http_body_call,"\n")
        if (isnothing(record_array) == true || length(record_array) == 0)
            return Some(nothing)
        end
        for record in record_array

            # split again along the tab -
            record_components = split(record,"\t")

            if (length(record_components) == 2)

                # the ec number is the second element -
                ec_number_value = string(record_components[2])

                # grab -
                push!(ec_number_array,ec_number_value)
            end
        end

        # wrap and return -
        return Some(ec_number_array)

    catch error

        # get the original error message -
        error_message = sprint(showerror, error, catch_backtrace())
        vl_error_obj = ErrorException(error_message)

        # Package the error -
        return Some(vl_error_obj)
    end
end


text:
"""
Abstract type representing empirical Bayes estimation methods.
"""
abstract type EBayesMethod end

function StatsBase.fit(mom::EBayesMethod, Zs)
    StatsBase.fit(mom, Zs, skedasticity(Zs))
end

struct MethodOfMoments{O} <: EBayesMethod
    object::O
    params::Any
end

MethodOfMoments(o) = MethodOfMoments(o, nothing)

struct ParametricMLE{O} <: EBayesMethod
    object::O
    params::Any
end

ParametricMLE(o) = ParametricMLE(o, nothing)

# Some defaults for parametric distribution fitting
_default_constraints(::Normal) = TwiceDifferentiableConstraints([-Inf; 0.0], [Inf; Inf])
_default_init(::Normal) = [0.0; 1.0]

function _default_constraints(::Union{Beta,Gamma})
    TwiceDifferentiableConstraints([0.0; 0.0], [Inf; Inf])
end

function _default_init(::Union{Beta,Gamma})
    [1.0; 1.0]
end

# struct SURE <: EBayesMethod
#   object::SoftThreshold,
#   location = :grandmean
# end

# SURE(Normal(); )       <---> keep dataset              SURE(Normal())          CovariateSample{EB, X}
# FixedLocationSURE( ; )      <---> keep datasaet
# GrandmeanLocationSURE( _ ; )   <----> keep dataset
# LinearModelLocationSURE( ; ); MLE(....)         @model :μ


text:
#=
    The File contains Policies definitions -
    - they are used to specify the execution policy of parallel algorithms -
      i.e., the kinds of parallelism allowed
=#

### NOT USED FOR NOW

using Base.Threads

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

# Each other policy is a subtype of this
abstract type AbstractPolicy end

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

# The execution policy type used as a unique type to disambiguate
# parallel algorithm overloading and require that a parallel algorithm's
# execution *may not be* parallelized
# The invocations of element access functions in parallel algorithms
# invoked with this policy are indeterminately sequenced in the calling thread
struct SequencedPolicy <: AbstractPolicy

end

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

# The execution policy type used as a unique type to disambiguate
# parallel algorithm overloading and indicate that a parallel algorithm's
# execution *may be* parallelized.
# The invocations of element access functions in parallel algorithms
# invoked with this policy are permitted to execute in either the invoking thread
# or in a thread implicitly created by the library to support parallel algorithm execution
# Any such invocations executing in the same thread
# are indeterminately sequenced with respect to each other
struct ParallelPolicy <: AbstractPolicy

end

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

const seq = SequencedPolicy()
const par = ParallelPolicy()


text:
ENV["GUROBI_HOME"] = ARGS[9]
ENV["GRB_LICENSE_FILE"] = ARGS[10]

using Gurobi, JuMP, CSV, DataFrames, Random, RLEVectors, DelimitedFiles, StatsBase

Random.seed!(1033)

@time begin

const CURR_PUMA = ARGS[1]
const HH_POOL_FILE = ARGS[2]
const INDV_POOL_FILE = ARGS[3]
const SYN_HHS_FILE = ARGS[4]
const SYN_INDVS_FILE = ARGS[5]
const PUMA_TRACT_EQUIV_FILE = ARGS[6]
const MARG_DIR = ARGS[7]
const GUROBI_LOGFILE = ARGS[8]
const PUMA_MARG_FILES = ["puma_hhtype.csv"]
const TRACT_MARG_FILES = ["tract_tenur_hhinc.csv", "tract_nwork.csv"]
const BLKGP_MARG_FILES = ["blkgp_tenur_hhsiz.csv"]
const MARG_COLS = [11, 21, 4, 22]
const INDV_TRACT_MARG_FILES = ["tract_i_sex_i_age.csv"]
const INDV_BLKGP_MARG_FILES = ["blkgp_emply.csv"]
const INDV_MARG_COLS = [19, 20]

## Get tracts
puma_tract_equiv = CSV.read(PUMA_TRACT_EQUIV_FILE, DataFrame)
curr_geo = filter(row -> row.STATEFP == parse(Float64, CURR_PUMA[1:2]) && row.PUMA5CE == parse(Float64, CURR_PUMA[3:7]), puma_tract_equiv)
const TRACT_GEOIDS = [string(lpad(string(sct[1]), 2, "0"), lpad(string(sct[2]), 3, "0"), lpad(string(sct[3]), 6, "0"))
                        for sct in zip(curr_geo.STATEFP, curr_geo.COUNTYFP, curr_geo.TRACTCE)]

## Process marginals
indv_index = length(vcat(PUMA_MARG_FILES, TRACT_MARG_FILES, BLKGP_MARG_FILES))+1
marginals = []
levelss = []
for marg_file in vcat(PUMA_MARG_FILES, TRACT_MARG_FILES, BLKGP_MARG_FILES, INDV_TRACT_MARG_FILES, INDV_BLKGP_MARG_FILES)
        marg_df = CSV.read("$MARG_DIR$marg_file", DataFrame)
        if marg_file in PUMA_MARG_FILES
                marg_df = filter(row -> lpad(string(row.geoid), 7, "0") == CURR_PUMA, marg_df)
        elseif marg_file in TRACT_MARG_FILES || marg_file in INDV_TRACT_MARG_FILES
                marg_df = filter(row -> lpad(string(row.geoid), 11, "0") in TRACT_GEOIDS, marg_df)
        elseif marg_file in BLKGP_MARG_FILES || marg_file in INDV_BLKGP_MARG_FILES
                marg_df = filter(row -> lpad(string(row.geoid), 12, "0")[1:11] in TRACT_GEOIDS, marg_df)
        end
        select!(marg_df, Not(:geoid))
        push!(marginals, sum.(eachcol(marg_df)))
        push!(levelss, names(marg_df))
end

const puma_hh_pop = sum(marginals[1])

## Process household  and individual pools
hh_pool_df = CSV.read(HH_POOL_FILE, DataFrame)
hh_pool = [row[:] for row in eachrow(Matrix(hh_pool_df))]

const n = length(hh_pool)

indv_pool_df = CSV.read(INDV_POOL_FILE, DataFrame)
hh_pool_hhsizs = rle(indv_pool_df.HHID)[2]
# precompute individual variable level counts for each household; factors for
#     household-level variables can all be 1 or whatever
indv_factors = Any[]
for lvls in levelss[1:(indv_index-1)]
        push!(indv_factors, [[1 for _ in 1:n] for _ in 1:length(lvls)])
end
for (indv_marg_col, lvls) in zip(INDV_MARG_COLS, levelss[indv_index:length(levelss)])
        with_counts = combine(groupby(indv_pool_df, :HHID), indv_marg_col => hh_values -> Tuple(count(val->val==l, hh_values) for l in lvls))
        counts_mat = [[v for v in vals] for vals in with_counts[!,2]]
        counts_mat = [[row[col] for row in counts_mat] for col in 1:size(counts_mat[1])[1]]
        push!(indv_factors, counts_mat)
end

println("Household Pool Size: ", n)
println("Marginals: ", vcat(PUMA_MARG_FILES, TRACT_MARG_FILES, BLKGP_MARG_FILES, INDV_TRACT_MARG_FILES, INDV_BLKGP_MARG_FILES))

#######################################################
## Define optimization model
model = Model(Gurobi.Optimizer)
set_optimizer_attribute(model, "LogFile", GUROBI_LOGFILE)

## Variables
@variable(model, x[1:n*2]) # define variable x
@variable(model, y[1:sum(length(lvls) for lvls in levelss)]) # define absolute value helper variable y

## Objective
@objective(model, Min, sum(y))

## Constraints
@constraint(model, con_lb[i=1:n*2], 0 <= x[i])
@constraint(model, con_ub[i=1:n*2], x[i] <= 1)
@constraint(model, con_rowsum[i=1:n], x[2*(i-1)+1] + x[2*(i-1)+2] == 1)
@constraint(model, con_col1sum, sum(x[2*(i-1)+1] for i in 1:n) == puma_hh_pop)
@constraint(model, con_col2sum, sum(x[2*(i-1)+2] for i in 1:n) == n - puma_hh_pop)
@constraint(model, con_col1indvsum, sum(hh_pool_hhsizs[i] * x[2*(i-1)+1] for i in 1:n) == sum(marginals[indv_index]))

# add marginal cell constraints to enforce absolute value
marg_cell_i = 1
for (marg_i, (marg_col, marg, lvls, indv_facs)) in enumerate(zip(vcat(MARG_COLS, INDV_MARG_COLS), marginals, levelss, indv_factors))
        cell_weight = length(lvls) * 2
        for (lvl_i, (lvl, ifs)) in enumerate(zip(lvls, indv_facs))
                @constraint(model, y[marg_cell_i] >= ((sum(ifs[i] * x[2*(i-1)+1] for (i,hh) in enumerate(hh_pool) if marg_i>=indv_index || hh[marg_col]==lvl)) - marg[lvl_i])/(1+marg[lvl_i])/cell_weight)
                @constraint(model, y[marg_cell_i] >= -((sum(ifs[i] * x[2*(i-1)+1] for (i,hh) in enumerate(hh_pool) if marg_i>=indv_index || hh[marg_col]==lvl)) - marg[lvl_i])/(1+marg[lvl_i])/cell_weight)
                global marg_cell_i += 1
        end
end

## Optimize
status=optimize!(model)

println("****************************************************")
println("final objective value = ", objective_value(model))

## extract selected households, and reset IDs
sol_matrix = transpose(reshape(value.(x), (2,n)))
sorted_hhids = sort([[abs(irow[2][1] - 1), hh_pool[irow[1]][1]] for irow in enumerate(eachrow(sol_matrix))])
synpop_hhids = [diff_hhid[2] for diff_hhid in sorted_hhids[1:puma_hh_pop]]

syn_hhs = filter(row -> (row.HHID in synpop_hhids), hh_pool_df)
syn_hhs.HHID = 1:DataFrames.nrow(syn_hhs)
CSV.write(SYN_HHS_FILE, syn_hhs)

syn_indvs = filter(row -> (row.HHID in synpop_hhids), indv_pool_df)
true_hhsizs = rle(syn_indvs.HHID)[2]
syn_indvs.HHID = collect(rep(collect(1:DataFrames.nrow(syn_hhs)), each = true_hhsizs))
syn_indvs.INDVID = 1:DataFrames.nrow(syn_indvs)
CSV.write(SYN_INDVS_FILE, syn_indvs)

end


text:
# Module written by Ed Scheinerman, ers@jhu.edu
# distributed under terms of the MIT license

module ShowSet

import Base: string, show, AbstractSet

function string(A::AbstractSet)
    elements = collect(A)
    try
        sort!(elements)
    catch
    end
    return "{" * join(elements,",") * "}"
end

show(io::IO, A::Set)    = print(io,string(A))
show(io::IO, A::BitSet) = print(io,string(A))

end # module ShowsSet


text:
using Flux
using Flux: crossentropy, normalise, onecold, onehotbatch
using Statistics: mean

labels = Flux.Data.Iris.labels()
features = Flux.Data.Iris.features()

# Subract mean, divide by std dev for normed mean of 0 and std dev of 1.
normed_features = normalise(features, dims=2)

klasses = sort(unique(labels))
onehot_labels = onehotbatch(labels, klasses)

# Split into training and test sets, 2/3 for training, 1/3 for test.
train_indices = [1:3:150 ; 2:3:150]

X_train = normed_features[:, train_indices]
y_train = onehot_labels[:, train_indices]

X_test = normed_features[:, 3:3:150]
y_test = onehot_labels[:, 3:3:150]

# Declare model taking 4 features as inputs and outputting 3 probabiltiies,
# one for each species of iris.
model = Chain(
    Dense(4, 3),
    softmax
)

loss(x, y) = crossentropy(model(x), y)

# Gradient descent optimiser with learning rate 0.5.
optimiser = Descent(0.5)

# Create iterator to train model over 110 epochs.
data_iterator = Iterators.repeated((X_train, y_train), 110)

println("Starting training.")
Flux.train!(loss, params(model), data_iterator, optimiser)

# Evaluate trained model against test set.
accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))

accuracy_score = accuracy(X_test, y_test)

println("\nAccuracy: $accuracy_score")

# Sanity check.
@assert accuracy_score > 0.8

function confusion_matrix(X, y)
    ŷ = onehotbatch(onecold(model(X)), 1:3)
    y * ŷ'
end

println("\nConfusion Matrix:\n")
display(confusion_matrix(X_test, y_test))


text:
using AlphaZero
using Test

include("../games/tictactoe/main.jl")
include("../games/connect-four/main.jl")
include("../scripts/lib/dummy_run.jl")

include("test_game.jl")

@testset "Testing Games" begin
  test_game(Tictactoe.Game)
  test_game(ConnectFour.Game)
  @test true
end

@testset "Dummy Runs" begin
  @test dummy_run(Tictactoe)
  #@test dummy_run(ConnectFour) # Takes a bit too long for Travis
end

max_stars_repo_name: btheodorou/Bridge-NN-Research

max_issues_repo_name: btheodorou/Bridge-NN-Research

max_forks_repo_name: btheodorou/Bridge-NN-Research


text:
using DrWatson
@quickactivate "StatReth"

# %%
using CSV
using DataFrames
using Turing
using Plots

include(srcdir("quap.jl"))
include(srcdir("tools.jl"))

# %%
# Load the data, don't forget to drop everything missing.
d = DataFrame(CSV.File(datadir("exp_raw/cherry_blossoms.csv"), missingstring = "NA"))
d2 = dropmissing(d, :doy)

# %% 4.76
# Redo the model from the text ...
using BSplines: BSplineBasis, basismatrix

num_knots = 15
knot_list = quantile(d2.year, range(0, 1, length = num_knots))

basis = BSplineBasis(4, knot_list)
B = basismatrix(basis, d2.year)

@model function spline(doy, B = B)
    α ~ Normal(100, 10)
    w ~ filldist(Normal(0, 10), size(B, 2))
    σ ~ Exponential(1)
    μ = α .+ B * w
    doy ~ MvNormal(μ, σ)
end

m4_7 = quap(spline(d2.doy))

# %%
# ... and look at the splines
w_str = ["w[$i]" for i in 1:length(basis)]
post = DataFrame(rand(m4_7.distr, 1000)', ["α"; w_str; "σ"])

w = mean.(eachcol(post[:, w_str]))

plot(legend = false, xlabel = "year", ylabel = "basis * weight")
for y in eachcol(B .* w')
    plot!(d2.year, y)
end
plot!()

mu = B * Array(post[!, w_str])'
mu_m, mu_lower, mu_upper = meanlowerupper(mu)

plot!(d2.year, mu_m, ribbon = (mu_m .- mu_lower, mu_upper .- mu_m),
                    lw = 2, fa = 0.1, c = :black)

# %%
# Now do the entire thing with more knots. The model can stay the same, just need to replace
# B
num_knots = 30
knot_list = quantile(d2.year, range(0, 1, length = num_knots))

basis2 = BSplineBasis(4, knot_list)
B2 = basismatrix(basis2, d2.year)

m4_7_2 = quap(spline(d2.doy, B2))

# %%
# ... and look at the splines
w_str = ["w[$i]" for i in 1:length(basis2)]
post = DataFrame(rand(m4_7_2.distr, 1000)', ["α"; w_str; "σ"])

w = mean.(eachcol(post[:, w_str]))

plot(legend = false, xlabel = "year", ylabel = "basis * weight")
for y in eachcol(B2 .* w')
    plot!(d2.year, y)
end
plot!()

mu = B2 * Array(post[!, w_str])'
mu_m, mu_lower, mu_upper = meanlowerupper(mu)

plot!(d2.year, mu_m, ribbon = (mu_m .- mu_lower, mu_upper .- mu_m),
                    lw = 2, fa = 0.1, c = :black)

# %%
"""
More knots make it possible for the spline to have more wiggles. Since that improves the
fit, it is taken advantage of and the spline gets wigglier. This is pretty similar to how
things go when trying to fit a polynomial while increasing its order. The fit gets better
for the cost of losing generality.
I'm not going to do the prior on the weights thing, I already did the same thing in 4H6.
The amplitude goes up or down, depending in what direction I change the prior.
"""


text:
"""
    rename_keys(d::Dict, subs::Dict)

Create a copy of `d`, replacing its keys according to mapping `subs`
"""
rename_keys(dct::Dict, subs::Dict) = Dict(get(subs, k, k) => v for (k, v) in pairs(dct))


text:
# ---------------------------------
# --- UHD Bindings
# ----------------------------------------------------
# Backend to pilot USRP with UHD lib
@reexport using UHDBindings
# --- Specific UHD related functions
# We export the UHD structure
export UHDBinding

# ----------------------------------------------------
# --- Adalm Pluto managment
# ----------------------------------------------------
# Backend for Adalm Pluto
@reexport using AdalmPluto
# --- Specific Pluto exportation
# We export the Adalm Pluto structure and the specific updateGainMode function
export AdalmPluto;
export updateGainMode!;

# ----------------------------------------------------
# --- RTL-SDR bindings
# ----------------------------------------------------
include("Backends/RTLSDR/RTLSDRBindings.jl");
@reexport using .RTLSDRBindings
export RTLSDRBinding

# ----------------------------------------------------
# --- Socket System
# ----------------------------------------------------
# --- Create and load module to pilot E310 devices
# To control this device we create a pure Socket based system
# for which the AbstractSDRs package will help to bind the utils
# Have a look on minimalTransceiver.jl for the code to be ran on E310
include("Backends/SDROverNetworks.jl");
@reexport using .SDROverNetworks
# --- Specific E310 related functions
export SDROverNetwork;

# ----------------------------------------------------
# --- Simulation Radio
# ----------------------------------------------------
# --- Create and module to emulate a radio device without any actual radio connected
include("Backends/RadioSims.jl");
@reexport using .RadioSims
# --- Specific simulation related function
export updatePacketSize!;
export updateBuffer!;
export RadioSim;

# ----------------------------------------------------
# --- Define common radio type
# ----------------------------------------------------
# We define an Union type that gathers all SDR backends
# This type will be used as default fallback methods to handle 2 things
# - In case of functions not supported in the given backend to obtain a predictible (and non error) behaviour
# - To simplify access to similar backends fields
AbstractSDR = Union{RadioSim,UHDBinding,PlutoSDR,SDROverNetwork,RTLSDRBinding}


text:
using ThermoState
using Documenter

makedocs(;
    modules=[ThermoState],
    authors="longemen3000 <longemen3000@gmail.com> and contributors",
    repo="https://github.com/longemen3000/ThermoState.jl/blob/{commit}{path}#L{line}",
    sitename="ThermoState.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://longemen3000.github.io/ThermoState.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
        "The Specification Object (`Spec`)" => "spec.md",
        "Thermodynamic State Object (`ThermodynamicState`)" => "state.md",
        "Utilities" => "utils.md",
        "Examples" => [
            "new_model.md",
        ],
        "API"=> "api.md"
    ],
)

deploydocs(;
    repo="github.com/longemen3000/ThermoState.jl",
)


text:
include("fourier_utils.jl")
include("fourier_operations.jl")
include("fourier_fields.jl")

# Define Interactions
function (∇::FourierOperator)(ϕ::FourierField)
    return FourierField(∇.op .* ϕ.data, ϕ.metadata)
end
function *(∇::FourierOperator, ϕ::FourierField)
    return FourierField(∇.op .* ϕ.data, ϕ.metadata)
end
# Define Naming
function (∇::FourierOperator{S, T})(ϕ::FourierField) where
    {S, T <: FourierOperatorMetaData}
    op_name = ∇.metadata.name
    ϕ_name = ϕ.metadata.name
    new_name = op_name * "(" * ϕ_name * ")"
    fmd = FourierMetaData(new_name, ϕ.metadata.grid, ϕ.metadata.transform)
    return FourierField(∇.op .* ϕ.data, fmd)
end

max_stars_repo_name: upiterbarg/QuasiGeostrophy

max_issues_repo_name: sandreza/QuasiGeostrophy

max_forks_repo_name: upiterbarg/QuasiGeostrophy


text:
## cllr.jl  Cost of the Log Likelihood Ratio
## (c) 2014--2015 David A. van Leeuwen
##
## Licensed under the MIT software license, see LICENSE.md

## Cllr for target and non-target scores
"""
`cllr(tar, non)` computes the cost of the log-likelihood-ratio, for log-likelihood-ratio (llr) scores
`tar` and `non` (target scores and non-target scores).  Target an non-target llr's are those for which
the numerator `H1` and denominator `H2` hypothesis are actually true in the log-likelihood-ratio

    llr(x) = p(x|H1) / p(x|H2)

where `x` is the test data for a system trying to disciminate between `H1` and `H2` given `x`.

The Cllr of a perfect system, assigning a llr of `+Inf` t otarget scores and `-Inf` to non-target scores
is 0, for a system that is indifferent, producing a `llr=0` for every input `x` the `Cllr=1`.
Please note that, for badly calibrated systems, `Cllr>1`.  The units of Cllr are measured in _bits_,
and cllr can be seen as the average amount of information per trial that is _not_ extracted from the
data.

Cllr measures the _calibration_ as well as _discrimination_ ability of a system.  Discrimination
entails the ability to produces target scores that are, typically, much higher than non-target scores.
Calibration entails that for every individual `x` an optimal Bayes decision can be made if a
cost function is known.
"""
function cllr{T<:Real}(tar::Vector{T}, non::Vector{T})
    ct = cn = zero(T)
    for x in tar
        ct += softplus(-x)
    end
    for x in non
        cn += softplus(x)
    end
    (ct / length(tar) + cn / length(non)) / 2log(2)
end

## minimum Cllr: Cllr after optimal score-to-llr transformation
"""
`mincllr(tar, non)` computes the `cllr` of target and non-target scores after an optimal
transformation of the data.  This allows for measuring the disrimination performance of
a system in the units of `cllr`, which are bits.
"""
mincllr{T<:Real}(tar::Vector{T}, non::Vector{T}) = cllr(optllr(tar, non)...)


text:
using BenchmarkTools

BenchmarkTools.DEFAULT_PARAMETERS.samples = 1000
BenchmarkTools.DEFAULT_PARAMETERS.seconds = 10
BenchmarkTools.DEFAULT_PARAMETERS.evals = 1
BenchmarkTools.DEFAULT_PARAMETERS.gctrial = true
BenchmarkTools.DEFAULT_PARAMETERS.gcsample = false
@show BenchmarkTools.DEFAULT_PARAMETERS

suite = BenchmarkGroup()

suite["BN"] = BenchmarkGroup()
suite["BN"]["BigInt(::BN)"] = @benchmarkable BigInt($(rand(BN)))

suite["BN"]["invmod(::BN, ::BN)"] = @benchmarkable invmod($(rand(BN)), $(curve_order(EP)))
suite["BN"]["invmod(::BigInt, ::BigInt)"] = @benchmarkable invmod($(BigInt(rand(BN))), $(BigInt(curve_order(EP))))
suite["BN"]["mod(::BN, ::BN)"] = @benchmarkable mod($(rand(BN)), $(curve_order(EP)))
suite["BN"]["mod(::BigInt, ::BigInt)"] = @benchmarkable mod($(BigInt(rand(BN))), $(BigInt(curve_order(EP))))
suite["BN"]["::BN * ::BN"] = @benchmarkable $(rand(BN)) * $(rand(BN))
suite["BN"]["::BigInt * ::BigInt"] = @benchmarkable $(BigInt(rand(BN))) * $(BigInt(rand(BN)))

suite["FP"] = BenchmarkGroup()
suite["FP"]["rand(::Type{FP})"] = @benchmarkable rand(FP)
suite["FP"]["FP(::Int)"] = @benchmarkable FP(rand(Int))
suite["MD"] = BenchmarkGroup()
suite["MD"]["md_hmac"] = @benchmarkable md_hmac($(rand(UInt8, MD_LEN)), $(rand(UInt8, MD_LEN)))
suite["PP"] = BenchmarkGroup()
suite["PP"]["field_final_exp"] = @benchmarkable field_final_exp($(rand(FP12)))
suite["PP"]["curve_miller"] = @benchmarkable curve_miller($(rand(EP)), $(rand(EP2)))
suite["EP"] = BenchmarkGroup()
suite["EP"]["curve_dbl(::EP)"] = @benchmarkable curve_dbl($(rand(EP)))
suite["EP2"] = BenchmarkGroup()
suite["EP2"]["curve_dbl(::EP2)"] = @benchmarkable curve_dbl($(rand(EP2)))

for T in (FP, FP2, FP6, FP12)
    name = string(T)
    name in suite || (suite[name] = BenchmarkGroup())
    suite[name]["+"] = @benchmarkable $(rand(T)) + $(rand(T))
    suite[name]["-"] = @benchmarkable $(rand(T)) - $(rand(T))
    suite[name]["*"] = @benchmarkable $(rand(T)) * $(rand(T))
    suite[name]["inv"] = @benchmarkable inv($(rand(T)))
    suite[name]["^"] = @benchmarkable $(rand(T)) ^ $(rand(BN))
    suite[name]["iszero(::$name)"] = @benchmarkable iszero($(rand(T)))
    suite[name]["isone(::$name)"] = @benchmarkable isone($(rand(T)))
    suite[name]["=="] = @benchmarkable $(rand(T)) == $(rand(T))
end

for T in (EP, EP2)
    name = string(T)
    name in suite || (suite[name] = BenchmarkGroup())
    suite[name]["=="] = @benchmarkable $(rand(T)) == $(rand(T))
    suite[name]["+"] = @benchmarkable $(rand(T)) + $(rand(T))
    suite[name]["-"] = @benchmarkable $(rand(T)) - $(rand(T))
    suite[name]["*"] = @benchmarkable $(rand(T)) * $(rand(BN))
    suite[name]["isinf(::$name)"] = @benchmarkable isinf($(rand(T)))
    suite[name]["isvalid(::$name)"] = @benchmarkable isvalid($(rand(T)))
end

function format_trial(suite, group, res)
    a = allocs(res)
    gct = BenchmarkTools.prettytime(gctime(res))
    t = BenchmarkTools.prettytime(time(res))
    m = BenchmarkTools.prettymemory(memory(res))
    return "[$suite][$group]: $t (alloc: $a, mem: $m, gc: $gct)"
end

# If a cache of tuned parameters already exists, use it, otherwise, tune and cache
# the benchmark parameters. Reusing cached parameters is faster and more reliable
# than re-tuning `suite` every time the file is included.
paramspath = joinpath(@__DIR__, "params.json")
if isfile(paramspath)
    loadparams!(suite, BenchmarkTools.load(paramspath)[1], :evals);
else
    println("First run - tuning params (please be patient) ...")
    tune!(suite)
    BenchmarkTools.save(paramspath, params(suite));
end

# print the results
results = run(suite, verbose = true)
for suiteres in results
    for groupres in suiteres.second
        msg = format_trial(suiteres.first, groupres.first, groupres.second)
        println(msg)
    end
end


text:
#  ##############################################################
#  # GUniFrac: Generalized UniFrac distances for comparing microbial
#  #						communities.
#  # Reference: Jun Chen & Hongzhe (2012). Associating microbiome
#  # composition with environmental covariates using generalized
#  # UniFrac distances.
#  ###############################################################

function GUniFrac(args...; otuTableFile::String = "",
	phy = thetree, alpha::Array{Float64,1} = Float64[], VAW::Bool = false)

    if isRooted(phy = phy) == false
    	error("isRooted:wrongtpe\n", "# Rooted phylogenetic tree required!")
    end

    # read in otu count table
    if isempty(otuTableFile)
    	error("GUniFrac:NoOTUtables\n", "# need to provide OTU count file");
    else
    	otuTable = readdlm(otuTableFile,',')
    	otuTableMatrix = otuTable[2:end, 1:end]
    	otuTableMatrix = convert(Array{Float64, 2}, otuTableMatrix);
    end

    row_sum = sum(otuTableMatrix, 2)
    #row_sum[1:10,:]=0.0
    zeroIndex = vec(row_sum[:,1]).==0.0
    otu_tab = otuTableMatrix
    otu_tab[zeroIndex.==false, :] = otuTableMatrix[zeroIndex .== false, :] ./ row_sum[zeroIndex.==false,:]
    otu_tab[zeroIndex.==true, :] = 0.0
    n = size(otu_tab)[1] # nrows

    # Construct the returning array
    # d_UW: unweighted UniFrac, d_VAW: weighted UniFrac
    function unifracArray(arg...)
    	unifracTuple = Array{Float64}(n, n)
    	return unifracTuple
    end
    if VAW == true
        unifracs = ntuple(unifracArray::Function,length(alpha)+2);
    else
        unifracs = ntuple(unifracArray::Function,length(alpha)+1);
    end

    nameList = Array{AbstractString}(length(alpha))
    for i = 1:length(alpha)
        nameList[i] = join(["d_", alpha[i]])
    end
    if VAW == true
        unifracs = Dict(zip(["d_UW","d_VAW", nameList], unifracs))
        fill!(unifracs["d_VAW"], 0.0)
    else
        unifracs = Dict(zip(["d_UW", nameList], unifracs))
    end

  	# Check OTU name consistency
    otuTableName = []
    for i=1:length(otuTable[1, 1:end])
        push!(otuTableName, AbstractString(otuTable[1,i]))
    end
    #otuNames = squeeze(convert(Array{ASCIIString,2},string(int(otuTable[1,1:end])),1)
    countIdex = 0
    otuNames = otuTableName
    for i =1 :length(otuNames)
    	if otuNames[i] in phy["tip_label"] ==false
    		countIdex = countIdex +1
    	end
    end

    if countIdex != 0
    	error("GUniFrac:moreOTU\n", "# The OTU table contains unknown OTUs! OTU names in the OTU table and the tree should match!")
    end

    # Get the subtree if tree contains more OTUs
##    tipIdex=Bool[]
##    for i = 1:length(phy["tip_label"])
##    	truefalse = phy["tip_label"][i] in otuNames
##        push!(tipIdex,truefalse)
##    end
##    countnz(tipIdex)
##    absent = phy["tip_label"][tipIdex.==false]

#    tip_label = phy["tip_label"]
#
#    absent = deleteat!(tip_label, findin(tip_label,otuNames))
#    if length(absent) != 0
#    	phy = DropTip(phy,absent)
#    	warn("The tree has more OTU than the OTU table!")
#    end

  	# Reorder the otu.tab matrix if the OTU orders are different
    tip_label = phy["tip_label"]

    reorderIdex = []
    for i = 1: length(tip_label)
     append!(reorderIdex,find(all(otuNames .== tip_label[i],2)))
    end
    otu_tab = otu_tab[:,reorderIdex]
    ntip = length(phy["tip_label"])
    nbr = size(phy["edge"])[1]
    edge = phy["edge"]
    edge2 = edge[:,2]
    br_len = phy["edge_length"]

 	#  Accumulate OTU proportions up the tree
    cum = Array{Float64}(nbr, n)    # Branch abundance matrix
    for i = 1:ntip
			inx = 1:length(edge2)
    	tip_loc = collect(1:length(edge2))[edge2 .== i][1]
    	cum[tip_loc:tip_loc, :] = cum[tip_loc:tip_loc, :] + transpose(otu_tab[:, i])
    	node = edge[tip_loc,1]      # Assume the direction of edge
			if countnz(edge2 .== node)==0
				node_loc = Int[]
			else
				node_loc = collect(1:length(edge2))[edge2 .== node][1]
			end
    	while length(node_loc)>0
    			cum[node_loc:node_loc, :] = cum[node_loc:node_loc,:] + transpose(otu_tab[:, i])
    			node = Int(edge[node_loc, 1])
					if countnz(edge2 .== node)==0
						node_loc = Int[]
					else
						node_loc = collect(1:length(edge2))[edge2 .== node][1]
					end
    	end
    end

  	# Calculate various UniFrac distances
   	cum_ct = round.((cum' .* row_sum)') 	# For VAW
    fill!(unifracs["d_UW"], 0.0)
    for k = 1:length(alpha)
        fill!(unifracs[[join(["d_",alpha[k]])]], 0.0)
    end

   	for i = 2:n
   		for j = 1:(i-1)
   			cum1 = cum[:, i]
   			cum2 = cum[:, j]
   			ind = []
            append!(ind,find(all(cum[:, i] + cum[:, j].>1e-10,2)))
            cum1 = cum1[ind]
            cum2 = cum2[ind]
            br_len2 = br_len[ind]
            mi = cum_ct[ind,i] + cum_ct[ind,j]
            mt = row_sum[i] + row_sum[j]
            diff = abs.((cum1 - cum2) ./ (cum1 +cum2))

            # Generalized UniFrac distance
            for k = 1:length(alpha)
                w = br_len2 .* (cum1 +cum2) .^ alpha[k]
            	#w = br_len2 .* (cum1 + cum2) ./ sqrt(mi .* (mt - mi))
            	unifracs[[join(["d_",alpha[k]])]][i,j] = unifracs[[join(["d_",alpha[k]])]][j,i] = sum(diff .* w) / sum(w)
            end

            # Variance Adjusted UniFrac Distance
            if VAW==true
                ind2 = []
                append!(ind2, find(all(mi .!= mt,2)))
                w = br_len2 .* (cum1 + cum2) ./ sqrt(mi .* (mt - mi))
                #w = br_len2 .* (cum1 +cum2) ^ alpha[k]
                unifracs["d_VAW"][i, j] = unifracs["d_VAW"][j,i] = sum(diff[ind2].*w[ind2])/sum(w[ind2])
            end
            # Unweighted UniFrac Distance
            ind3 = []
            ind4 = []
            append!(ind3, find(all(cum1 .>1e-10,2)))
            append!(ind4, find(all(cum2 .>1e-10,2)))
            cum1[ind3]=1
            cum2[ind4]=1
            unifracs["d_UW"][i, j] = unifracs["d_UW"][j,i] = sum(abs.(cum1-cum2)./(cum1+cum2) .* br_len2)/ sum(br_len2)
#            unifracs["d_UW"][i, i] = 0.0
        end
        unifracs["d_UW"][isnan.(unifracs["d_UW"])]=0.0
        for k = 1:length(alpha)
            unifracs[[join(["d_",alpha[k]])]][isnan.(unifracs[[join(["d_",alpha[k]])]])] = 0.0
        end
        if VAW == true
            unifracs["d_VAW"][isnan.(unifracs["d_VAW"])]=0.0
        end

    end
    return unifracs
end


text:
# Additional force fields.

abstract type AdHocForce end

"""
    HarmonicRestraint

Harmonic restraint between the centers of mass of two molecules at a single
bead.
"""
struct HarmonicRestraint <: AdHocForce
    "Force constant (kJ / nm^2 mol)."
    k::Float64
    "Separation (nm)."
    R::Float64
    "Index of first molecule."
    n1::Int
    "Index of second molecule."
    n2::Int
    "Bead index."
    j::Int
end

function eval_pot(hr::HarmonicRestraint, s::State)
    0.5 * hr.k * (com_r(hr.n1, hr.n2, hr.j, s) - hr.R)^2
end

function eval_force!(result::ForceContainer, hr::HarmonicRestraint, s::State)
    factor = -hr.k * (com_r(hr.n1, hr.n2, hr.j, s) - hr.R) # kJ / nm mol
    result.f[hr.j, :, :, :] .+= factor .* com_r_grad(hr.n1, hr.n2, hr.j, s)

    nothing
end


text:
using Statistics
using Interpolations
"""
    contrast(image; step=2)

Measure the contrast of an image, in the sense of high contrast imaging.
That is, divide the image into annuli moving outwards from the centre
(index 0,0 if offset image) and calculate the standard deviation in
each.

Returns a vector of annulus locations in pixels and a vector of standard
deviations.

*NOTE* This is the 1σ contrast. Multiply by five to get the usual confidence
value.
"""
function contrast(image; step=2)
    I = origin(image)
    if I[1] == I[2] == 1
        image = centered(image)
    end
    dx = UnitRange(axes(image,1))
    dy = UnitRange(axes(image,2))
    dr = sqrt.(
        dx.^2 .+ (dy').^2
    )

    c_img = collect(image)

    bins = 0:step:maximum(dr)
    # bins = 30:step:100
    contrast = zeros(size(bins))
    mask = falses(size(image))
    mask2 = isfinite.(c_img)
    for i in eachindex(bins)
        bin = bins[i]
        mask .= (bin.-step/2) .< dr .< (bin.+step/2)
        mask .&= mask2
        c = std(view(c_img, mask))
        contrast[i] = c
    end

    return (;separation=bins, contrast)
end

function radialbins(image; step=2)
    I = origin(image)
    if I[1] == I[2] == 1
        image = centered(image)
    end
    dx = UnitRange(axes(image,1))
    dy = UnitRange(axes(image,2))
    dr = sqrt.(
        dx.^2 .+ (dy').^2
    )

    c_img = collect(image)

    bins = 0:step:maximum(dr)
    # bins = 30:step:100
    contrast = zeros(size(bins))
    mask = falses(size(image))
    mask2 = isfinite.(c_img)
    radial_bins = map(eachindex(bins)) do i
        bin = bins[i]
        mask .= (bin.-step/2) .< dr .< (bin.+step/2)
        mask .&= mask2
        return view(c_img, mask)
    end

    return (;separation=bins, data=radial_bins)
end

"""
    out,prep = contrast!_prep(image, bins)
    contrast!(out,prep)

A fast non-allocating routine for calculating
contrast. Calculates the 1σ contrast of image
moving outwards from the centre accoring to `bins`.

After calling the `prep` function, you can alter
`image` with new data and calling `contrast!`
will still be very fast and non-allocating.
"""
function contrast!_prep(image,bins,spatialmask=trues(size(image)))
    cont = zeros(size(bins))
    # masks = falses(length(bins),size(image,1),size(image,2))
    # masks = BitMatrix[]
    # We want matrices of bool, not bitmatrix for non-allocating masking view performance
    masks = [
        fill(false, size(image))
        for _ in bins
    ]
    dx = axes(image,1) .- mean(axes(image,1))
    dy = axes(image,2) .- mean(axes(image,2))
    dr = sqrt.(
        dx.^2 .+ (dy').^2
    )
    s = step(bins)/2
    for i in eachindex(bins)
        bin = bins[i]
        masks[i] .= (bin.-s) .< dr .< (bin.+s)
        masks[i] .&= spatialmask
        # push!(masks,(bin.-s) .< dr .< (bin.+s) )
    end
    views = [view(image, mask) for mask in masks]
    return cont, views
end
function contrast!(out,views)
    for (i,v) in pairs(views)
        out[i] = std(v)
    end
end

export contrast
export contrast!_prep
export contrast!

"""
    contrast_interp(image; step=2)

Returns a linear interpolation on top of the results from `contrast`.
Extrapolated results return Inf.
"""
function contrast_interp(image; step=2)
    cont = contrast(image; step)
    return LinearInterpolation(cont.separation, cont.contrast, extrapolation_bc=Inf)
end

using Distributions

"""
    contrast_dist_interp(::Type{Distribution}, image; step=2)

Returns a linear interpolation on top of the results from `contrast`.
Extrapolated results return Inf.
"""
function contrast_dist_interp(Dist::Type{<:UnivariateDistribution}, image; step=2)
    sep, bins = radialbins(image; step)
    has_px = map(!isempty, bins)
    sep = sep[has_px]
    bins = bins[has_px]
    dists = map(bins) do bin
        fit(Dist, bin)
    end
    params = Distributions.params.(dists)
    param_arrs = [zeros(eltype(first(dists)), length(params)) for _ in 1:length(first(params))]
    for j in eachindex(param_arrs)
        for (i,param) in enumerate(params)
            param_arrs[j][i] = params[i][j]
        end
    end
    param_interps = Tuple(map(param_arrs) do parr
        LinearInterpolation(sep, parr, extrapolation_bc=Inf)
    end)
    return function (sep_px)
        params = map(f->f(sep_px), param_interps)
        dist = Dist(params...)
    end
end

export contrast_interp
export contrast_dist_interp

function profile(image; step=2)
    I = origin(image)
    if I[1] == I[2] == 1
        image = centered(image)
    end
    dx = UnitRange(axes(image,1))
    dy = UnitRange(axes(image,2))
    dr = sqrt.(
        dx.^2 .+ (dy').^2
    )

    c_img = collect(image)

    bins = 0:step:maximum(dr)
    # bins = 30:step:100
    profile = zeros(size(bins))
    mask = falses(size(image))
    mask2 = isfinite.(c_img)
    for i in eachindex(bins)
        bin = bins[i]
        mask .= (bin.-step/2) .< dr .< (bin.+step/2)
        mask .&= mask2
        c = mean(view(c_img, mask))
        profile[i] = c
    end

    return (;separation=bins, profile)
end

function snrmap(image, contrast=nothing)
    I = origin(image)
    if I[1] == I[2] == 1
        image = centered(image)
    end
    dx = UnitRange(axes(image,1))
    dy = UnitRange(axes(image,2))
    dr = sqrt.(
        dx.^2 .+ (dy').^2
    )

    if isnothing(contrast)
        contrast = contrast_interp(image)
    end
    snr = profile2image(contrast, image)

    out = similar(image)
    out .= image ./ snr
    return out
end
export snrmap

max_stars_repo_name: sefffal/SpaceImages


text:
@testset "Gauss–Chebyshev" begin

    @testset "Check error" begin
        @test_throws DomainError gausschebyshev(-1,1)
        @test_throws DomainError gausschebyshev(-1,2)
        @test_throws DomainError gausschebyshev(-1,3)
        @test_throws DomainError gausschebyshev(-1,4)
        @test_throws ArgumentError gausschebyshev(0,5)
    end

    n = 10

    @testset "x.^2" begin
        x, w = gausschebyshev(n,1)
        @test dot(x.^2,w) ≈ π/2
        x, w = gausschebyshev(n,2)
        @test dot(x.^2,w) ≈ π/8
        x, w = gausschebyshev(n,3)
        @test dot(x.^2,w) ≈ π/2
        x, w = gausschebyshev(n,4)
        @test dot(x.^2,w) ≈ π/2
    end

    @testset "x^3" begin
        x, w = gausschebyshev(n,1)
        @test abs(dot(x.^3,w))<1e-15
        x, w = gausschebyshev(n,2)
        @test abs(dot(x.^3,w))<1e-15
        x, w = gausschebyshev(n,3)
        @test dot(x.^3,w) ≈ 3*π/8
        x, w = gausschebyshev(n,4)
        @test dot(x.^3,w) ≈ -3*π/8
    end
end