text:
module ImpvolEquilibrium

export period_wrapper, coerce_parameters!, rotate_sectors, CES_price_index, array_transpose, remove_shock!

using Logging, Base.Test

include("utils.jl")
parameters = Dict{Symbol, Any}()

# for matrix conformity, store all variables in a 4-dimensional array:
# mnjt: destination, source, sector, time

# per-period random variables are stored as
# mnjs: destination, source, sector, state

function report(A)
	display(A[1,1:2,end-1:end,1])
	sleep(1)
end

function where_in_simplex(labor_share, s)
	l = labor_share[:,:,:,s]
	return (minimum(l), ind2sub(l,indmin(l)))
end

function biggest_gap(x1, x2, s)
	y = (x1 .- x2)[:,:,:,:] .^ 2
	return (maximum(y), ind2sub(y,indmax(y)))
end

function expected_value(y)
	_, _, _, S = size(y)
	if S==1
		return y[:,:,:,1]
	else
		# ignore first realization of S
		return mean(y[:,:,:,2:end], 4)
	end
end

function rotate_sectors(A, y)
	# y may be a time series variable or a random variable
	M, N, J, T = size(y)
	K = size(A,1)
	B = zeros(M,N,K,T)
	for m=1:M
		for n=1:N
			for t=1:T
				B[m,n,:,t] = A * y[m,n,:,t]
			end
		end
	end
	return B
end

function p_distance(p1, p2, theta)
	s1 = p1 .^ (-theta) ./ sum(p1 .^ (-theta), 2)
	s2 = p2 .^ (-theta) ./ sum(p2 .^ (-theta), 2)
	return share_distance(s1, s2)
end

function share_distance(p1, p2)
	return mean((p1 .- p2) .^ 2) .^ 0.5
end

function deflate_all_nominal_variables!(random_variables, parameters, t)
	compute_price_index!(random_variables, parameters, t)
	US_price_index = random_variables[:P_ns][1:1,end:end,1:1,:]
	for variable in [:R_njs, :rho_njs, :w_njs, :P_ns, :input_price_njs, :P_njs, :E_mjs]
		random_variables[variable] = random_variables[variable] ./ US_price_index
	end
end

function free_trade_sector_shares!(parameters)
	N, J, T = parameters[:N], parameters[:J], parameters[:T]
	gamma_jk = parameters[:gamma_jk]
	beta_j = parameters[:beta_j]
	alpha_jt = parameters[:nu_njt] ./ sum(parameters[:nu_njt], 3)

	revenue_shares = zeros(1,1,J,T)
	for t=1:T
		revenue_shares[1,1,:,t] = eigen_share(alpha_jt[1,1,:,t]*beta_j[1,1,:,1]' + gamma_jk)
	end
	parameters[:sector_shares] = revenue_shares
end

function input_price_index!(random_variables, parameters)
	P = random_variables[:P_njs]
	random_variables[:input_price_njs] = exp.(rotate_sectors(parameters[:gamma_jk]', log.(P)))
end

function compute_price!(random_variables, parameters, t)
	theta = parameters[:theta]
	kappa = non_random_variable(parameters[:kappa_mnjt], t)
	rho_njs = random_variables[:rho_njs]

	random_variables[:P_njs] = array_transpose(sum((rho_njs ./ kappa) .^ (-theta), 2) .^ (-1/theta))
end

function CES_price_index(alpha, P_njs, sigma)
	return sum(alpha .* P_njs .^ (1-sigma), 3) .^ (1/(1-sigma))
end

function compute_price_index!(random_variables, parameters, t)
	nu = non_random_variable(parameters[:nu_njt], t)
	alpha = nu ./ sum(nu, 3)
	P_njs = random_variables[:P_njs]
	sigma = parameters[:sigma]

	# Cobb-Douglas is a special case when sigma ~ 1
	random_variables[:P_ns] = CES_price_index(alpha, P_njs, sigma)
end

function free_trade_country_shares!(random_variables, parameters)
	free_trade_sector_shares!(parameters)
	A_njs = random_variables[:A_njs]
	L_njs = random_variables[:L_njs]
	theta = parameters[:theta]
	beta_j = parameters[:beta_j]

	d_njs = A_njs .^ (theta ./ (1 + theta*beta_j)) .* L_njs .^ (theta .* beta_j ./ (1 + theta*beta_j))
	random_variables[:d_njs_free] = d_njs ./ sum(d_njs,2)
end

function free_trade_labor_shares!(random_variables, parameters, t)
	free_trade_country_shares!(random_variables, parameters)

	d_njs = random_variables[:d_njs_free]
	beta_j = parameters[:beta_j]
	E_jt = non_random_variable(parameters[:sector_shares], t)

	y_njs = beta_j .* d_njs .* E_jt

	random_variables[:L_njs_free] = y_njs ./ sum(y_njs, 3)
end

function free_trade_wages!(random_variables, parameters, t)
	free_trade_country_shares!(random_variables, parameters)

	d_njs = random_variables[:d_njs_free]
	L_njs = random_variables[:L_njs]
	beta_j = parameters[:beta_j]
	E_wt = non_random_variable(parameters[:sector_shares], t) .* non_random_variable(parameters[:nominal_world_expenditure], t)

	random_variables[:w_njs] = beta_j .* d_njs .* E_wt ./ L_njs
end

function free_trade_prices!(random_variables, parameters, t)
	free_trade_country_shares!(random_variables, parameters)
	beta_j = parameters[:beta_j]
	theta = parameters[:theta]
	d_njs = random_variables[:d_njs_free]
	w_njs = random_variables[:w_njs]
	xi = parameters[:xi]
	B_j = parameters[:B_j]
	A_njs = random_variables[:A_njs]
	gamma_jk = parameters[:gamma_jk]'

	random_variables[:P_njs] = exp.(rotate_sectors(inv(eye(gamma_jk)-gamma_jk), log.(xi * d_njs .^(1/theta) .* B_j .* (w_njs .^beta_j) ./ A_njs)))
	random_variables[:rho_njs] = random_variables[:P_njs] ./ d_njs .^(1/theta)
	input_price_index!(random_variables, parameters)
end

function compute_trade_shares!(random_variables, parameters, t)
	theta = parameters[:theta]
	kappa = non_random_variable(parameters[:kappa_mnjt], t)
	rho_njs = random_variables[:rho_njs]
	P_mjs = array_transpose(random_variables[:P_njs])

	random_variables[:d_mnjs] = (rho_njs ./ kappa ./ P_mjs) .^ (-theta)
end

function compute_wage!(random_variables, parameters)
	input_price_index!(random_variables, parameters)

	rho_njs = random_variables[:rho_njs]
	input_price_njs = random_variables[:input_price_njs]

	beta_j = parameters[:beta_j]
	B_j = parameters[:B_j]
	xi = parameters[:xi]
	A_njs = random_variables[:A_njs]

	random_variables[:w_njs] = (rho_njs .* A_njs ./ input_price_njs ./ (xi * B_j)) .^ (1 ./ beta_j)
end

function compute_revenue!(random_variables, parameters)
	w_njs = random_variables[:w_njs]
	L_njs = random_variables[:L_njs]
	beta_j = parameters[:beta_j]

	random_variables[:R_njs] = w_njs .* L_njs ./ beta_j
end

function CES_share(nu, price, sigma)
	temp = nu .* price .^ (1-sigma)
	return temp ./ sum(temp, 3)
end

function compute_expenditure_shares!(random_variables, parameters, t)
	R_nks = random_variables[:R_njs]
	beta_j = parameters[:beta_j]
	gamma_jk = parameters[:gamma_jk]
	nu = non_random_variable(parameters[:nu_njt], t)
	# encompass CES and Cobb-Douglas
	alpha_njt = CES_share(nu, random_variables[:P_njs], parameters[:sigma])
	S_nt = non_random_variable(parameters[:S_nt], t)
	expenditure = sum(R_nks, 3) .- S_nt

	wagebill_ns = rotate_sectors(beta_j[:]', R_nks)
	intermediate_njs = rotate_sectors(gamma_jk, R_nks)
	e_njs = alpha_njt .* wagebill_ns .+ intermediate_njs .- alpha_njt .* S_nt
	# trade imbalance may make it negative
	e_njs = max.(parameters[:numerical_zero], e_njs)
	random_variables[:e_mjs] = array_transpose(e_njs ./ sum(e_njs, 3))
end

function compute_real_gdp!(random_variables, parameters, t)
	compute_price_index!(random_variables, parameters, t)
	w_njs = random_variables[:w_njs]
	L_njs = random_variables[:L_njs]
	P_ns = random_variables[:P_ns]

	random_variables[:real_GDP] = w_njs .* L_njs ./ P_ns
end

function fixed_expenditure_shares!(random_variables, parameters, t)
	R_mjs = array_transpose(random_variables[:R_njs])
	expenditure = sum(R_mjs, 3) .- array_transpose(non_random_variable(parameters[:S_nt], t))
	random_variables[:E_mjs] = random_variables[:e_mjs] .* max.(parameters[:numerical_zero], expenditure)
end

function shadow_price_step(random_variables, parameters, t)
	theta = parameters[:theta]
	E_mjs = random_variables[:E_mjs]
	R_njs = random_variables[:R_njs]
	kappa_mnjt = non_random_variable(parameters[:kappa_mnjt], t)
	P_mjs = array_transpose(random_variables[:P_njs])

	return sum( (kappa_mnjt .* P_mjs) .^ theta .* E_mjs ./ R_njs, 1) .^ (1/theta)
end

function starting_values!(random_variables, parameters, t)
	free_trade_wages!(random_variables, parameters, t)
	free_trade_prices!(random_variables, parameters, t)
	free_trade_labor_shares!(random_variables, parameters, t)
	compute_revenue!(random_variables, parameters)
	compute_expenditure_shares!(random_variables, parameters, t)
	fixed_expenditure_shares!(random_variables, parameters, t)
	deflate_all_nominal_variables!(random_variables, parameters, t)
end

function inner_loop!(random_variables, parameters, t)
	debug("------ BEGIN Inner loop")
	lambda = parameters[:inner_step_size]
	dist = 999
	k = 1

	while (dist > parameters[:inner_tolerance]) && (k <= parameters[:max_iter_inner])
		new_rho = shadow_price_step(random_variables, parameters, t)
		dist = p_distance(new_rho, random_variables[:rho_njs], parameters[:theta])
		debug("-------- Inner ", k, ": ", dist)
		random_variables[:rho_njs] = new_rho
		compute_price!(random_variables, parameters, t)
		compute_wage!(random_variables, parameters)
		old_R = random_variables[:R_njs]
		compute_revenue!(random_variables, parameters)
		random_variables[:R_njs] = lambda*random_variables[:R_njs] + (1-lambda)*old_R
		fixed_expenditure_shares!(random_variables, parameters, t)
		deflate_all_nominal_variables!(random_variables, parameters, t)

		k = k+1
	end
	warn("inner: ", k-1)
	debug("------ END Inner loop")
end

function middle_loop!(random_variables, parameters, t)
	debug("---- BEGIN Middle loop")
	dist = 999
	k = 1
	old_expenditure_shares = random_variables[:e_mjs]

	while (dist > parameters[:middle_tolerance]) && (k <= parameters[:max_iter_middle])
		inner_loop!(random_variables, parameters, t)
		compute_expenditure_shares!(random_variables, parameters, t)

		dist = share_distance(random_variables[:e_mjs][:,:,1:end-1,:], old_expenditure_shares[:,:,1:end-1,:])

		random_variables[:e_mjs] = parameters[:middle_step_size]*random_variables[:e_mjs]+(1-parameters[:middle_step_size])*old_expenditure_shares
		info("------ Middle ", k, ": ", dist)

		old_expenditure_shares = random_variables[:e_mjs]
		k = k+1
	end
	debug("---- END Middle loop")
end

function adjustment_loop!(random_variables, L_nj_star, parameters, t)

	function max_step_size(x0, direction)
		nulla = parameters[:numerical_zero]
		negative_entries = min.(direction, -nulla)
		return minimum(nulla - x0 ./ negative_entries, 3)
	end

	function evaluate_utility(random_variables, L_nj_star, parameters, t)
		random_variables[:L_njs] = max.(parameters[:numerical_zero], random_variables[:L_njs])
		random_variables[:L_njs] = random_variables[:L_njs] ./ sum(random_variables[:L_njs], 3)

		middle_loop!(random_variables, parameters, t)

		w_ns = sum(random_variables[:w_njs] .* random_variables[:L_njs], 3)
		wage_gap = random_variables[:w_njs] ./ w_ns
		return sum(parameters[:one_over_rho]*log.(w_ns) .- 0.5*sum((random_variables[:L_njs] .- L_nj_star).^2, 3))
	end

	function calculate_derivative(random_variables, L_nj_star, parameters)
		w_njs = random_variables[:w_njs]
		L_njs = random_variables[:L_njs]
		w_ns = sum(w_njs .* random_variables[:L_njs], 3)
		wage_gap = w_njs ./ w_ns
		gradient = parameters[:one_over_rho]*wage_gap .- (L_njs .- L_nj_star)

		# ensure that steps sum to zero: we can only reallocate labor across sectors
		return gradient .- mean(gradient, 3)
	end

	debug("-- BEGIN Adjustment loop")
	random_variables[:L_njs] = L_nj_star
	free_trade_labor_shares!(random_variables, parameters, t)
	L_njs_free = random_variables[:L_njs_free]
	k = 1

	nulla = ones(1,1,1,1)*parameters[:numerical_zero]

	utility = evaluate_utility(random_variables, L_nj_star, parameters, t)
	gradient = calculate_derivative(random_variables, L_nj_star, parameters)
	dist = mean(gradient .^ 2) .^ 0.5
	info("Starting from $dist")

	while (dist > parameters[:adjustment_tolerance]) && (k <= parameters[:max_iter_adjustment])
		snapshot = copy(random_variables)

		previous_utility = utility
		previous_L = random_variables[:L_njs]
		step_size = min.(0.33*max_step_size(previous_L, gradient), 0.1)

		snapshot[:L_njs] = previous_L .+ gradient .* step_size

		utility = evaluate_utility(snapshot, L_nj_star, parameters, t)
		gradient = calculate_derivative(snapshot, L_nj_star, parameters)
		dist = mean(gradient .^ 2) .^ 0.5

		difference = utility - previous_utility
		proportional_increase = difference / sum(gradient .^ 2 .* step_size)
		debug("Difference: ", difference)
		debug("Proportional increase: ", proportional_increase)

		info("---- Adjustment $k: $dist")

		k = k+1
		random_variables = snapshot
	end
	debug("-- END Adjustment loop")
end

function expected_wage_share(random_variables)
	w_njs = random_variables[:w_njs]
	L_njs = random_variables[:L_njs]

	wage_bill = w_njs .* L_njs
	wage_share = wage_bill ./ sum(wage_bill, 3)

	return expected_value(wage_share)
end

function outer_loop!(random_variables, parameters, t, L_nj_star)
	N, J = parameters[:N], parameters[:J]
	lambda = parameters[:outer_step_size]
	random_variables[:L_njs] = copy(L_nj_star)

	info("BEGIN Outer loop")
	starting_values!(random_variables, parameters, t)

	dist = 999
	k = 1

	while (dist > parameters[:outer_tolerance]) && (k <= parameters[:max_iter_outer])
		old_wage_share = L_nj_star ./ sum(L_nj_star, 3)
		adjustment_loop!(random_variables, L_nj_star, parameters, t)
		wage_share = expected_wage_share(random_variables)
		@test sum(wage_share, 3) ≈ ones(1,N,1,1) atol=1e-9

		dist = share_distance(wage_share, old_wage_share)
		info("-- Outer ", k, ": ", dist)

		L_nj_star = (1-lambda)*old_wage_share .+ lambda*wage_share
		k = k+1
	end
	info("END Outer loop")
	return L_nj_star
end

function period_wrapper(parameters, t)
	N, J = parameters[:N], parameters[:J]

	info("--- Period ", t, " ---")
	A_njs = parameters[:A_njs][t]
	random_variables = Dict{Symbol, Any}()
	random_variables[:A_njs] = A_njs

	random_variables[:L_njs] = zeros(1,N,J,1)
	for n=1:N
		# start from expenditure labor weights
		random_variables[:L_njs][1,n,:,1] = parameters[:importance_weight]
	end
	free_trade_labor_shares!(random_variables, parameters, t)
	stv = random_variables[:L_njs_free]
	L_nj_star = outer_loop!(random_variables, parameters, t, stv)

	compute_trade_shares!(random_variables, parameters, t)
	compute_real_gdp!(random_variables, parameters, t)
	return random_variables
end

end


text:
using WeaponSystems
using Test

@testset "WeaponSystems.jl" begin
    # Write your tests here.
end

ext: jl


text:
"""
    loglikelihood!(m::IntervalCensoredModel)

This function uses the parameter values `m.η`, `m.expη`, `m.λ₀`, `m.Λ₀`.
"""
function loglikelihood!(
    icm      :: IntervalCensoredModel{T},
    needgrad :: Bool = false,
    needhess :: Bool = false
    ) where T <: BlasReal
    n, p = size(icm.Z)
    # log-likelihood
    logl = zero(T)
    needgrad && fill!(icm.∇λ₀, 0)
    if needhess; fill!(icm.Hλ₀λ₀, 0); fill!(icm.Hλ₀β, 0); end
    @inbounds for i in 1:n
        if icm.C[i] == :left_censored
            s = exp(- icm.Λ₀[icm.Ridx[i]] * icm.expη[i])
            logl += log1p(-s)
            if needgrad
                g = icm.expη[i] / (inv(s) - 1)
                for j in 1:icm.Ridx[i]
                    icm.∇λ₀[j] += g
                end
                icm.βres[i] = icm.Λ₀[icm.Ridx[i]] * g
            end
            if needhess
                # Hλ₀λ₀
                h = abs2(icm.expη[i]) / (2 - s - inv(s))
                for j in 1:icm.Ridx[i], k in 1:j
                    icm.Hλ₀λ₀[k, j] -= h
                end
                # Hββ
                g = icm.Λ₀[icm.Ridx[i]] * icm.expη[i] / (inv(s) - 1)
                icm.glmwt[i] = - g * (1 - g * inv(s))
                # # Hλ₀β
                # g = icm.expη[i] / (inv(s) - 1) + icm.Λ₀[icm.Ridx[i]] * h
                # for k in 1:p, j in 1:icm.Ridx[i]
                #     icm.Hλ₀β[j, k] -= g * icm.Z[i, k]
                # end
            end
        elseif icm.C[i] == :right_censored
            l     = - icm.Λ₀[icm.Lidx[i]] * icm.expη[i]
            logl += l
            if needgrad
                for j in 1:icm.Lidx[i]
                    icm.∇λ₀[j] -= icm.expη[i]
                end
                icm.βres[i] = l
            end
            if needhess
                icm.glmwt[i] = -l
                for k in 1:p, j in 1:icm.Lidx[i]
                    icm.Hλ₀β[j, k] += icm.expη[i] * icm.Z[i, k]
                end
            end
        elseif icm.C[i] == :interval_censored
            ΔΛ₀  = icm.Λ₀[icm.Ridx[i]] - icm.Λ₀[icm.Lidx[i]]
            l     = - icm.Λ₀[icm.Lidx[i]] * icm.expη[i]
            s     = exp(- ΔΛ₀ * icm.expη[i])
            logl += l + log1p(-s)
            if needgrad
                g = icm.expη[i]
                for j in 1:icm.Lidx[i]
                    icm.∇λ₀[j] -= g
                end
                g = icm.expη[i] / (inv(s) - 1)
                for j in (icm.Lidx[i] + 1):icm.Ridx[i]
                    icm.∇λ₀[j] += g
                end
                icm.βres[i] = l + ΔΛ₀ * g
            end
            if needhess
                # Hλ₀λ₀
                h = abs2(icm.expη[i]) / (2 - s - inv(s))
                for j in (icm.Lidx[i]+1):icm.Ridx[i], k in (icm.Lidx[i]+1):j
                    icm.Hλ₀λ₀[k, j] -= h
                end
                # Hββ
                g = ΔΛ₀ * icm.expη[i] / (inv(s) - 1)
                icm.glmwt[i] = icm.Λ₀[icm.Lidx[i]] * icm.expη[i] - g * (1 - g * inv(s))
                # # Hλ₀β
                # for k in 1:p
                #     for j in 1:icm.Lidx[i]
                #         icm.Hλ₀β[j, k] += icm.expη[i] * icm.Z[i, k]
                #     end
                #     g = - icm.expη[i] / (inv(s) - 1) - ΔΛ₀ * h
                #     for j in (icm.Lidx[i] + 1):icm.Ridx[i]
                #         icm.Hλ₀β[j, k] += g * icm.Z[i, k]
                #     end
                # end
            end
        elseif icm.C[i] == :exact_time
            Tidx  = icm.Lidx[i]
            logl += log(icm.λ₀[Tidx]) + icm.η[i] - icm.Λ₀[Tidx] * icm.expη[i]
            if needgrad
                g = icm.expη[i]
                for j in 1:Tidx
                    icm.∇λ₀[j] -= g
                end
                icm.∇λ₀[Tidx] += inv(icm.λ₀[Tidx])
                icm.βres[i] = 1 - icm.Λ₀[Tidx] * icm.expη[i]
            end
            if needhess
                # Hλ₀λ₀
                icm.Hλ₀λ₀[Tidx, Tidx] += abs2(inv(icm.λ₀[Tidx]))
                # Hββ
                icm.glmwt[i] = icm.Λ₀[Tidx] * icm.expη[i]
                # # Hλ₀β
                # for k in 1:p, j in 1:Tidx
                #     icm.Hλ₀β[j, k] += icm.expη[i] * icm.Z[i, k]
                # end
            end
        end
    end
    needgrad && mul!(icm.∇β, transpose(icm.Z), icm.βres)
    if needhess
        copytri!(icm.Hλ₀λ₀, 'U')
        for (i, w) in enumerate(icm.glmwt)
            icm.glmwt[i] = w > 0 ? sqrt(w) : zero(w)
        end
        mul!(icm.storage_np, Diagonal(icm.glmwt), icm.Z)
        mul!(icm.Hββ, transpose(icm.storage_np), icm.storage_np)
    end
    # return log-likelihood
    logl
end

"""
    initialize_uniform!(icm::IntervalCensoredModel)

Set `icm.λ₀`, `icm.Λ₀` and `icm.S₀` according to a discrete uniform distribution
on the observed examination time points `0 < s1 < ... < sm < ∞`. `icm.β` are
set to  zero; `icm.η` and `icm.expη` are set accordingly.
"""
function initialize_uniform!(icm::IntervalCensoredModel{T}) where T <: Real
    m    = length(icm.ts)
    invm = inv(m)
    # survival function
    icm.S₀[end] = 0
    @inbounds for i in m-1:-1:1
        icm.S₀[i] = icm.S₀[i+1] + invm
    end
    # hazard function
    icm.λ₀[1] = invm
    @inbounds for i in 2:m
        icm.λ₀[i] = invm / icm.S₀[i - 1]
    end
    # cumulative hazard
    cumsum!(icm.Λ₀, icm.λ₀)
    # β, η, expη
    fill!(icm.β, 0)
    fill!(icm.η, 0)
    fill!(icm.expη, 1)
    icm
end

"""
    npmle!(icm::IntervalCensoredModel)

Non-parametric MLE (NPMLE) of survival function `S(t)` using the MM algorithm.
`icm.λ₀` and `icm.Λ₀` are set according to NPMLE of `icm.S₀`. `icm.β` are set to
zero; `icm.η` and `icm.expη` are set accordingly.
"""
function npmle!(
    icm     :: IntervalCensoredModel{T};
    maxiter :: Integer = 100_000,
    reltol  :: Real = 1e-8,
    verbose :: Bool = false
    ) where T <: Real
    n, m = length(icm.L), length(icm.ts)
    # construct membership matrix A = (αᵢⱼ), αᵢⱼ = I(sⱼ ∈ (Lᵢ, Rⱼ])
    A = Matrix{T}([icm.L[i] < icm.ts[j] ≤ icm.R[i] for i in 1:n, j in 1:m])
    # scratch space
    storage_n = Vector{T}(undef, n)
    storage_m = Vector{T}(undef, m)
    # starting point (1/m, ..., 1/m)
    p = fill(T(1 / m), m)
    logl = sum(log, mul!(storage_n, A, p))
    verbose && println("iter=0, logl=$logl")
    for iter in 1:maxiter
        logl_prev = logl
        storage_n .= inv.(storage_n)
        mul!(storage_m, transpose(A), storage_n)
        p .*= storage_m ./ n
        logl = sum(log, mul!(storage_n, A, p))
        verbose && println("iter=$iter, logl=$logl")
        # check convergence criterion
        if abs(logl - logl_prev) < reltol * (abs(logl_prev) + 1)
            break
        elseif iter == maxiter
            @warn "maximum iterations $maxiter reached"
        end
    end
    # survival function
    icm.S₀[end] = 0
    for i in m-1:-1:1
        icm.S₀[i] = icm.S₀[i+1] + p[i+1]
    end
    # hazard function
    icm.λ₀[1] = p[1]
    for i in 2:m
        icm.λ₀[i] = p[i] / icm.S₀[i - 1]
    end
    # cumulative hazard
    cumsum!(icm.Λ₀, icm.λ₀)
    # set β to 0
    fill!(icm.β, 0)
    fill!(icm.η, 0)
    fill!(icm.expη, 1)
    # return
    icm
end

function fit!(
    icm      :: IntervalCensoredModel,
    solver    = Ipopt.IpoptSolver(print_level=0, tol=1e-4);
    init     :: IntervalCensoredModel = icm,
    verbose  :: Bool = true
    )
    # set up NLP optimization problem
    m, p = length(icm.ts), size(icm.Z, 2)
    npar = m + p
    optm = MathProgBase.NonlinearModel(solver)
    lb   = fill(-Inf, npar)
    ub   = fill( Inf, npar)
    MathProgBase.loadproblem!(optm, npar, 0, lb, ub, Float64[], Float64[], :Max, icm)
    par0 = Vector{Float64}(undef, npar)
    modelpar_to_optimpar!(par0, icm)
    MathProgBase.setwarmstart!(optm, par0)
    MathProgBase.optimize!(optm)
    optstat = MathProgBase.status(optm)
    optstat == :Optimal || @warn("Optimization unsuccesful; got $optstat")
    # refresh objective, gradient, and Hessian
    optimpar_to_modelpar!(icm, MathProgBase.getsolution(optm))
    icm.S₀ .= exp.(- icm.Λ₀)
    loglikelihood!(icm, true, true)
    # inference for β (this inference is WRONG. What's implemented here is purely for convenience.)
    Hββ_eval, Hββ_evec = eigen(Symmetric(icm.Hββ))
    for (i, e) in enumerate(Hββ_eval)
        Hββ_eval[i] = e < 0 ? 0 : inv(e)
    end
    icm.Vββ .= Hββ_evec * Diagonal(Hββ_eval) * transpose(Hββ_evec)
    icm.isfitted[1] = true
    icm
end

"""
    modelpar_to_optimpar!(par, icm)

Translate model parameters in `icm` to optimization variables in `par`.
"""
function modelpar_to_optimpar!(
    par :: Vector{T},
    icm :: IntervalCensoredModel
    ) where T <: AbstractFloat
    m, p = length(icm.ts), size(icm.Z, 2)
    @inbounds for i in 1:m
        par[i] = log(max(icm.λ₀[i], floatmin(T)))
    end
    copyto!(par, m + 1, icm.β, 1, p)
    par
end

"""
    optimpar_to_modelpar!(icm, par)

Translate optimization variables in `par` to the model parameters in `m`.
"""
function optimpar_to_modelpar!(
    icm :: IntervalCensoredModel,
    par :: Vector
    )
    m, p = length(icm.ts), size(icm.Z, 2)
    @inbounds for i in 1:m
        icm.λ₀[i] = exp(par[i])
    end
    cumsum!(icm.Λ₀, icm.λ₀)
    copyto!(icm.β, 1, par, m + 1, p)
    mul!(icm.η, icm.Z, icm.β)
    icm.expη .= exp.(icm.η)
    icm
end

function MathProgBase.initialize(
                       :: IntervalCensoredModel,
    requested_features :: Vector{Symbol}
    )
    for feat in requested_features
        if !(feat in [:Grad, :Hess])
            error("Unsupported feature $feat")
        end
    end
end

MathProgBase.features_available(::IntervalCensoredModel) = [:Grad, :Hess]

function MathProgBase.eval_f(
    icm :: IntervalCensoredModel,
    par :: Vector
    )
    optimpar_to_modelpar!(icm, par)
    # println(loglikelihood!(icm))
    loglikelihood!(icm)
end

function MathProgBase.eval_grad_f(
    icm  :: IntervalCensoredModel,
    grad :: Vector,
    par  :: Vector
    )
    m, p = length(icm.ts), size(icm.Z, 2)
    optimpar_to_modelpar!(icm, par)
    loglikelihood!(icm, true)
    # gradient wrt λ₀
    # chain rule for unconstrained parameterization
    @inbounds for i in 1:m
        grad[i] = icm.∇λ₀[i] .* icm.λ₀[i]
    end
    # gradient wrt β
    copyto!(grad, m + 1, icm.∇β, 1, p)
    nothing
end

MathProgBase.eval_g(icm::IntervalCensoredModel, g, par) = nothing
MathProgBase.jac_structure(icm::IntervalCensoredModel) = Int[], Int[]
MathProgBase.eval_jac_g(icm::IntervalCensoredModel, J, par) = nothing

function MathProgBase.hesslag_structure(icm::IntervalCensoredModel)
    # our Hessian is block diagonal with two blocks: m-by-m, p-by-p
    m, p = length(icm.ts), size(icm.Z, 2)
    arr1 = Vector{Int}(undef, ◺(m) + ◺(p))
    arr2 = Vector{Int}(undef, ◺(m) + ◺(p))
    idx  = 1
    @inbounds for j in 1:m, i in 1:j
        arr1[idx] = i
        arr2[idx] = j
        idx      += 1
    end
    @inbounds for j in 1:p, i in 1:j
        arr1[idx] = m + i
        arr2[idx] = m + j
        idx      += 1
    end
    return (arr1, arr2)
end

function MathProgBase.eval_hesslag(
    icm :: IntervalCensoredModel,
    H   :: Vector{T},
    par :: Vector{T},
    σ   :: T,
    μ   :: Vector{T}
    ) where {T}
    m, p = length(icm.ts), size(icm.Z, 2)
    # refresh obj, gradient, and hessian
    optimpar_to_modelpar!(icm, par)
    loglikelihood!(icm, true, true)
    # Hλ₀λ₀
    idx = 1
    @inbounds for j in 1:m
        for i in 1:(j-1)
            H[idx] = - icm.Hλ₀λ₀[i, j] * icm.λ₀[i] * icm.λ₀[j]
            idx   += 1
        end
        H[idx] = (- icm.Hλ₀λ₀[j, j] * icm.λ₀[j] + icm.∇λ₀[j]) * icm.λ₀[j]
        idx += 1
    end
    # Hββ
    @inbounds for j in 1:p, i in 1:j
        H[idx] = - icm.Hββ[i, j]
        idx   += 1
    end
    lmul!(σ, H)
end

"""
    ◺(n::Integer)

Triangular number `n * (n + 1) / 2`.
"""
@inline ◺(n::Integer) = (n * (n + 1)) >> 1


text:
module CommandLineEquationsOfState

__precompile__()

include("Settings.jl")
include("argparse.jl")

end # module


text:
using FourierCore, FourierCore.Grid, FourierCore.Domain
using FFTW, LinearAlgebra, BenchmarkTools, Random, JLD2
using HDF5
using GLMakie
rng = MersenneTwister(12345)
# Random.seed!(123456789)
Random.seed!(12)
# jld_name = "high_order_timestep_spatial_tracer_"
jld_name = "blocky"
include("transform.jl")
include("random_phase_kernel.jl")
# using GLMakie
using CUDA
arraytype = CuArray
Ω = S¹(2π)^2
N = 2^10 # number of gridpoints
grid = FourierGrid(N, Ω, arraytype=arraytype)
nodes, wavenumbers = grid.nodes, grid.wavenumbers

x = nodes[1]
y = nodes[2]
kˣ = wavenumbers[1]
kʸ = wavenumbers[2]
# construct filter
kxmax = maximum(kˣ)
kymax = maximum(kˣ)
filter = @. (kˣ)^2 + (kʸ)^2 ≤ ((kxmax / 2)^2 + (kymax / 2)^2)
filter = @. abs(kˣ) .+ 0 * abs(kʸ) ≤ 2 / 3 * kxmax
@. filter = filter * (0 * abs(kˣ) .+ 1 * abs(kʸ) ≤ 2 / 3 * kxmax)

# now define the random field
wavemax = 5
𝓀 = arraytype([-wavemax, 0.0, wavemax]) # arraytype(1.0 .* [-wavemax, -wavemax + 1, wavemax - 1, wavemax])# arraytype(collect(-wavemax:1:wavemax))
𝓀ˣ = reshape(𝓀, (length(𝓀), 1))
𝓀ʸ = reshape(𝓀, (1, length(𝓀)))
# A = @. 0.1 * (𝓀ˣ * 𝓀ˣ + 𝓀ʸ * 𝓀ʸ)^(-11 / 12)
A = @. 0.5 * (𝓀ˣ * 𝓀ˣ + 𝓀ʸ * 𝓀ʸ)^(0.01) # @. 1e-1 / (1 * 2 * wavemax^2) .* (𝓀ˣ * 𝓀ˣ + 𝓀ʸ * 𝓀ʸ)^(0.0) # ( 1 .+ (0 .* 𝓀ˣ) .* 𝓀ʸ)
A[A.==Inf] .= 0.0
φ = arraytype(2π * rand(size(A)...))
field = arraytype(zeros(N, N))

##
# Fields
# velocity
ψ = arraytype(zeros(ComplexF64, N, N))
u = similar(ψ)
v = similar(ψ)

# auxiliary fields
S = arraytype(zeros(ComplexF64, size(ψ)..., 2))
Ṡ = arraytype(zeros(ComplexF64, size(ψ)..., 2))

∂ˣθ = similar(ψ)
∂ʸθ = similar(ψ)
∂ˣζ = similar(ψ)
∂ʸζ = similar(ψ)

𝒟θ = similar(ψ)
𝒟ζ = similar(ψ)
θ̇ = similar(ψ)

k₁ = similar(S)
k₂ = similar(S)
k₃ = similar(S)
k₄ = similar(S)
S̃ = similar(S)

# source
sθ = similar(ψ)
sζ = similar(ψ)

# phase
φ̇ = similar(A)

# operators
∂x = im * kˣ
∂y = im * kʸ
Δ = @. ∂x^2 + ∂y^2
Δ⁻¹ = 1 ./ Δ
bools = (!).(isnan.(Δ⁻¹))
Δ⁻¹ .*= bools # hack in the fact that false * NaN = 0

# plan ffts
P = plan_fft!(ψ)
P⁻¹ = plan_ifft!(ψ)

##
Δx = x[2] - x[1]
Δt = Δx / (2π) * 1
κ = 1.0 * Δx^2

# Dissipation
𝒟 = @. κ * Δ - 1e-1 * (κ * Δ)^2 + 1e-3 * (κ * Δ)^3 - 1e-5 * (κ * Δ)^4

# take the initial condition as negative of the source
tic = Base.time()

# save some snapshots
ψ_save = typeof(real.(Array(ψ)))[]
θ_save = typeof(real.(Array(ψ)))[]

θ = view(S, :, :, 1)
ζ = view(S, :, :, 2)
@. θ = 0.0 * sin(3 * x) * sin(3 * y) + -0.1 + 0im
@. ζ = sin(3 * x) * sin(3 * y)

t = [0.0]
tend = 40 # 5000

phase_speed = 1.0

iend = ceil(Int, tend / Δt)

operators = (; P, P⁻¹, filter, Δ⁻¹, 𝒟, ∂x, ∂y)
auxiliary = (; ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, u, v, ∂ˣζ, ∂ʸζ, ∂ˣθ, ∂ʸθ, 𝒟θ, 𝒟ζ, sθ, sζ)
constants = (; τ=4.0 * Δt, e=1e-2)# (; τ = 0.01, e = 0.01)

parameters = (; auxiliary, operators, constants)

function rhs!(Ṡ, S, parameters)
    θ̇ = view(Ṡ, :, :, 1)
    ζ̇ = view(Ṡ, :, :, 2)
    θ = view(S, :, :, 1)
    ζ = view(S, :, :, 2)

    P, P⁻¹, filter, Δ⁻¹, 𝒟, ∂x, ∂y = parameters.operators
    ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, u, v, ∂ˣζ, ∂ʸζ, ∂ˣθ, ∂ʸθ, 𝒟θ, 𝒟ζ, sθ, sζ = parameters.auxiliary
    τ, e = parameters.constants
    # construct random phase forcing
    event = stream_function!(ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ)
    wait(event)

    # construct source for vorticity
    @. sζ = ψ

    # P * ψ
    P * θ # in place fft ζ
    P * ζ # in place fft
    # grab stream function from vorticity
    @. ψ = Δ⁻¹ * ζ
    # ∇ᵖψ
    @. u = filter * (∂y * ψ)
    @. v = filter * -1.0 * (∂x * ψ)
    # ∇ζ
    @. ∂ˣθ = filter * ∂x * θ
    @. ∂ʸθ = filter * ∂y * θ
    @. ∂ˣζ = filter * ∂x * ζ
    @. ∂ʸζ = filter * ∂y * ζ
    @. 𝒟ζ = (Δ⁻¹ + 𝒟) * ζ
    @. 𝒟θ = 𝒟 * θ
    # go back to real space
    P⁻¹ * u
    P⁻¹ * v

    P⁻¹ * ζ
    P⁻¹ * ∂ˣζ
    P⁻¹ * ∂ʸζ
    P⁻¹ * 𝒟ζ

    P⁻¹ * θ
    P⁻¹ * ∂ˣθ
    P⁻¹ * ∂ʸθ
    P⁻¹ * 𝒟θ

    # construct source
    @. sθ = v + e - 1 / τ * θ * (real(θ) > 0)

    # rhs
    @. ζ̇ = real(-u * ∂ˣζ - v * ∂ʸζ + 𝒟ζ + sζ)
    @. θ̇ = real(-u * ∂ˣθ - v * ∂ʸθ + 𝒟θ + sθ)
    @. S = real(S)
    @. Ṡ = real(Ṡ)

    return nothing
end

include("interpolation.jl")
filename = "higher_rez.hdf5"
# rm(filename)
fid = h5open(filename, "w")
create_group(fid, "vorticity")
create_group(fid, "moisture")
saveindex = 0

for i = 1:iend
    # fourth order runge-kutta on deterministic part
    # keep ψ frozen is the correct way to do it here

    # the below assumes that φ is just a function of time
    rhs!(k₁, S, parameters)
    @. S̃ = S + Δt * k₁ * 0.5

    φ_rhs_normal!(φ̇, φ, rng)
    @. φ += phase_speed * sqrt(Δt / 2) * φ̇ # now at t = 0.5

    rhs!(k₂, S̃, parameters)
    @. S̃ = S + Δt * k₂ * 0.5
    rhs!(k₃, S̃, parameters)
    @. S̃ = S + Δt * k₃

    φ_rhs_normal!(φ̇, φ, rng)
    @. φ += phase_speed * sqrt(Δt / 2) * φ̇ # now at t = 1.0

    rhs!(k₄, S̃, parameters)
    @. S += Δt / 6 * (k₁ + 2 * k₂ + 2 * k₃ + k₄)

    t[1] += Δt
    # save output
    if t[1] > 40
        if i % div(iend, 400) == 0
            global saveindex += 1
            fid["vorticity"][string(saveindex)] = quick_interpolation(ζ)
            fid["moisture"][string(saveindex)] = quick_interpolation(θ)
        end
    end

    if i % div(iend, 100) == 0
        println("time is t = ", t[1])
        local toc = Base.time()
        println("the time for the simulation is ", toc - tic, " seconds")
        println("extrema of supersaturation are ", extrema(real.(θ)))
        println("extrema of vorticity are ", extrema(real.(ζ)))
        # println("on wavenumber index ", index_choice)
    end

end

close(fid)

toc = Base.time()
println("the time for the simulation was ", toc - tic, " seconds")

x_A = Array(x)[:] .- 2π
θ_F = Array(real.(θ))
ω_F = Array(real.(ζ))

fig = Figure(resolution=(2048, 512))
ax1 = Axis(fig[1, 1], title="moisture t = 0")
ax2 = Axis(fig[1, 2], title="moisture t = " * string(tend))
ax3 = Axis(fig[1, 4], title="vorticity t = " * string(tend))
println("the extrema of the end field is ", extrema(θ_F))
colormap = :bone_1
# colormap = :nipy_spectral
colorrange = (-1, 0)
colorrange2 = extrema(θ_F)
colorrange3 = (-maximum(ω_F), maximum(ω_F))
heatmap!(ax1, x_A, x_A, θ_F, colormap=colormap, colorrange=colorrange2, interpolate=true)
hm = heatmap!(ax2, x_A, x_A, θ_F, colormap=colormap, colorrange=colorrange2, interpolate=true)
hm_e = heatmap!(ax3, x_A, x_A, ω_F, colormap=:balance, colorrange=colorrange3, interpolate=true)
Colorbar(fig[1, 3], hm, height=Relative(3 / 4), width=25, ticklabelsize=30, labelsize=30, ticksize=25, tickalign=1,)
Colorbar(fig[1, 5], hm_e, height=Relative(3 / 4), width=25, ticklabelsize=30, labelsize=30, ticksize=25, tickalign=1,)
display(fig)

max_stars_repo_name: sandreza/FourierCore

max_issues_repo_name: sandreza/FourierCore

max_forks_repo_name: sandreza/FourierCore


text:
using Base.MathConstants

function getHarmonicNumber(n::Integer)::Float64
    H = 0
    for i = 1:n
        H = H + (1 / i)
    end
    return H
end

function getHarmonicNumberWithOrder(n::Integer, order::Integer)::Float64
    H = 0
    for i = 1:n
        H = H + (1 / i^order)
    end
    return H
end

eulermascheroni = convert(Float64, γ)

max_stars_repo_name: EllianCarlos/number-theory-but-in-julia

max_issues_repo_name: EllianCarlos/number-theory-but-in-julia

max_forks_repo_name: EllianCarlos/number-theory-but-in-julia


text:
module Test05

import Automa
import Automa.RegExp: @re_str
using Test

@testset "Test05" begin
    re = Automa.RegExp

    keyword = re"if|else|end|while"
    ident = re.diff(re"[a-z]+", keyword)
    token = re.alt(keyword, ident)

    keyword.actions[:exit] = [:keyword]
    ident.actions[:exit] = [:ident]

    machine = Automa.compile(token)

    for generator in (:table, :inline, :goto), checkbounds in (true, false), clean in (true, false)
        ctx = Automa.CodeGenContext(generator=generator, checkbounds=checkbounds, clean=clean)
        init_code = Automa.generate_init_code(ctx, machine)
        exec_code = Automa.generate_exec_code(ctx, machine, :debug)
        validate = @eval function (data)
            logger = Symbol[]
            $(init_code)
            p_end = p_eof = lastindex(data)
            $(exec_code)
            return cs == 0, logger
        end
        @test validate(b"if") == (true, [:keyword])
        @test validate(b"else") == (true, [:keyword])
        @test validate(b"end") == (true, [:keyword])
        @test validate(b"while") == (true, [:keyword])
        @test validate(b"e") == (true, [:ident])
        @test validate(b"eif") == (true, [:ident])
        @test validate(b"i") == (true, [:ident])
        @test validate(b"iff") == (true, [:ident])
        @test validate(b"1if") == (false, [])
    end
end

end


text:
using Dates
using JSON3
using Test

using Stonks: to_dataframe
using Stonks.Models: ExchangeRate, AssetInfo
using Stonks.Parsers: tryparse_js, js_to_dict, apply_filters, snake_case, JSONParser

include("test_utils.jl")

@testset "Parsing utility functions" begin
  content_ex = """{"from":"EUR", "to": "USD", "date": "2022-03-12", "close": "1.0915"}"""
  js_ex = JSON3.read(content_ex)

  @testset "Parse a flat js object into a data type with remaps" begin
    try_data = tryparse_js(ExchangeRate, js_ex)
    @test typeof(try_data) <: Exception
    ex = tryparse_js(
      ExchangeRate, js_ex; remaps=Dict(:base => :from, :target => :to, :rate => :close)
    )
    @test isa(ex, ExchangeRate)
    @test ex.base == js_ex[:from] &&
      ex.target == js_ex[:to] &&
      ex.date == Date(js_ex[:date])
  end

  @testset "Parse a flat js object into a data type with fields of wrong type" begin
    content = """{"symbol":"AAPL", "currency": "USD", "employees": "None"}"""
    data = tryparse_js(AssetInfo, JSON3.read(content))
    @test ismissing(data.employees)
  end

  @testset "Parse a flat js object into a data type with remaps and fixed values" begin
    content = """{"date": "2022-03-12", "close": "1.0915"}"""
    ex = tryparse_js(
      ExchangeRate,
      js_ex;
      remaps=Dict(:rate => :close),
      fixed=Dict(:base => "FOO", :target => "BAR"),
    )
    @test isa(ex, ExchangeRate)
    @test ex.base == "FOO" && ex.target == "BAR" && ex.date == Date(js_ex[:date])
  end

  @testset "Conert a JSON object to a dictionary" begin
    js_dict = js_to_dict(js_ex)
    @test isa(js_dict, Dict)
    @test isempty(setdiff(keys(js_dict), keys(js_ex)))
  end

  @testset "Conert a JSON nested object to a dict" begin
    content = """{
      "endDate": {
        "raw": 1640908800,
        "fmt": "2021-12-31"
      },
      "totalRevenue": {
        "raw": 3258000000,
        "fmt": "3.26B",
        "longFmt": "3,258,000,000"
      },
      "costOfRevenue": {
        "raw": -2680000000,
        "fmt": "-2.68B",
        "longFmt": "-2,680,000,000"
      }
    }"""
    js = JSON3.read(content)
    data = js_to_dict(js)
    @test isempty(setdiff(keys(data), keys(js)))
  end

  @testset "Apply date filters" begin
    prices = fake_price_data(7, Date("2022-03-12"))
    get_minmax(data) = @chain map(x -> x.date, data) (minimum(_), maximum(_))
    min, max = get_minmax(prices)
    p_filt = apply_filters(prices, "date"; from=min + Day(1), to=max - Day(1))
    p_min, p_max = get_minmax(p_filt)
    p_filt = apply_filters(prices, "date"; from=min + Day(1))
    p_min, p_max = get_minmax(p_filt)
    @test p_min == min + Day(1) && p_max == max
    p_filt = apply_filters(prices, "date"; to=max - Day(2))
    p_min, p_max = get_minmax(p_filt)
    @test p_min == min && p_max == max - Day(2)
    p_filt = apply_filters(prices, "date"; from=min + Day(3))
    p_min, p_max = get_minmax(p_filt)
    @test p_min == min + Day(3) && p_max == max
  end

  @testset "Converting camelCase to snake_case" begin
    snake_case("timeStamp")
    snake_case("askBestVolume")
    snake_case("askBest30MWPrice")
  end
end


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "cddlib"
# upstream 0.94j mapped to semantic versioning
version = v"0.94.10"

# Collection of sources required to build cddlib
sources = [
    GitSource("https://github.com/cddlib/cddlib","6d980f2258052f6cf47acca611d6a268d657dfc7")
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/*
sed -i '/^libcddgmp_la_LDFLAGS/ s/$/  $(CDD_LDFLAGS)/' lib-src/Makefile.gmp.am
./bootstrap # needed because we are building a commit and not a release
CPPFLAGS=-I${prefix}/include ./configure --prefix=${prefix} --host=${target}
cd lib-src # Call make in the lib-src subfolder to avoid building the doc folder since pdflatex is not installed
make -j${nproc}
make install
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = supported_platforms()

# The products that we will ensure are always built
products = [
    LibraryProduct("libcdd", :libcdd)
    LibraryProduct("libcddgmp", :libcddgmp)
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency(PackageSpec(name="GMP_jll", uuid="781609d7-10c4-51f6-84f2-b8444358ff6d"))
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)


text:
""" Sets parameters within Dict. """
function set_param!(d::Dict, key::Symbol, val, checks = true)
    if haskey(d, key) && val isa typeof(d[key])
        d[key] = val
        return
    elseif checks
        throw(ErrorException("Parameter with key " * string(key) * " is invalid."))
    else
        return
    end
end

function set_params!(d::Dict; kwargs...)
    for (k,v) in kwargs
        set_param!(d, k, v)
    end
    return
end

""" Gets parameters within Dict. """
function get_param(gm::Dict, key::Symbol)
    if haskey(gm, key)
        return gm[key]
    else
        throw(ErrorException("Parameter with key " * string(key) * " is invalid."))
    end
end

"""
    $(TYPEDSIGNATURES)

Normalizes a given chunk of data by column.
"""
function normalize(X::Matrix)
    bounds = [(minimum(X[:,i]), maximum(X[:,i])) for i = 1:size(X, 2)]
    X_norm = zeros(size(X))
    for i = 1:size(X,2)
        X_norm[:,i] = (X[:,i] .- bounds[i][1]) ./ (bounds[i][2] - bounds[i][1])
    end
    return X_norm, bounds
end

function normalize(X::Vector)
    bounds = (minimum(X), maximum(X))
    X_norm = (X .- bounds[1]) ./ (bounds[2] - bounds[1])
    return X_norm, bounds
end

"""
    $(TYPEDSIGNATURES)

Returns the distance matrix of X. Should almost always be done on a normalized X matrix.
"""
function pairwise_distances(X_norm::Matrix)
    dists = zeros(size(X_norm, 1), size(X_norm, 1))
    for i = 1:size(X_norm, 1)
        for j = i:size(X_norm, 1)
            dists[i,j] = sum((X_norm[i,:] .- X_norm[j,:]).^2)
            dists[j, i] = dists[i,j]
        end
    end
    return dists
end

""" Computes the mode, i.e. the most common element of an array. """
function mode(X::Vector)
    count_dict = Dict()
    for i = 1:length(X)
        if !(X[i] in keys(count_dict))
            count_dict[X[i]] = 1
        else
            count_dict[X[i]] += 1
        end
    end
    maxval = maximum(values(count_dict))
    return [k for (k,v) in count_dict if v == maxval]
end

"""
    $(TYPEDSIGNATURES)

Splits data according to the kwargs. Potential kwargs:
- bins::Int: The number of bins that you would like the data to fit into.
- sample_proportion::Union{Array, Real}: Proportion of samples in each bin. If Real, means only two bins.
- sample_count::Union{Array}: Number of samples in each bin.
"""
function split_data(X, Y;
    bins::Union{Int, Nothing} = nothing,
    sample_proportion::Union{Array, Real, Nothing} = nothing,
    sample_count::Union{Array, Nothing} = nothing)

    (sum(!isnothing(bins) + !isnothing(sample_proportion) + !isnothing(sample_count)) <= 1) || throw(ErrorException("split_data only takes one or zero kwargs."))

    subset_idxs = [0]
    if sum(!isnothing(bins) + !isnothing(sample_proportion) + !isnothing(sample_count)) == 0
        sample_proportion = [0.5, 0.5]
    elseif !isnothing(bins)
        sample_proportion = [1/bins for i = 1:bins]
    end

    if !isnothing(sample_proportion)
        isapprox(sum(sample_proportion), 1, atol = 1e-7) || throw(ErrorException("Sum of sample-proportion in split_data must be equal to the number of samples. "))
        for i = 1:length(sample_proportion)
            append!(subset_idxs, minimum([Int(ceil(length(Y)*sum(sample_proportion[1:i]))), length(Y)])) # Minimum to make sure we get no bugs with atol = 1e-7
        end
    elseif !isnothing(sample_count)
        sum(sample_count) == size(X, 1) || throw(ErrorException("Sum of sample-count in split_data must be equal to the number of samples. "))
        for i = 1:length(sample_count)
            push!(subset_idxs, subset_idxs[end] + sample_count[i])
        end
    end
    return [(X[subset_idxs[i]+1:subset_idxs[i+1], :],
    Y[subset_idxs[i]+1:subset_idxs[i+1]]) for i = 1:length(subset_idxs)-1]
end


text:
using Test
using AtomicLevels
# ^^^ -- to make it possible to run the test file separately

module GRASPParser
using Test
using AtomicLevels
import AtomicLevels: CSF, csfs
using WignerSymbols
include("grasp/rcsfparser.jl")

compare_with_grasp(f, rcsfout) = @testset "GRASP CSFs: $(rcsfout)" begin
    grasp_csfs = parse_rcsf(joinpath(@__DIR__, "grasp", rcsfout))
    atlev_csfs = f()
    @test length(atlev_csfs) == length(grasp_csfs)
    for atlev_csf in atlev_csfs
        @test atlev_csf in grasp_csfs
    end
end

export compare_with_grasp
end

using .GRASPParser

@testset "Excited configurations" begin
    @testset "Simple" begin
        @test_throws ArgumentError excited_configurations(rc"[Kr] 5s2 5p6", max_excitations=:triples)
        @test_throws ArgumentError excited_configurations(rc"[Kr] 5s2 5p6", min_occupancy=[2,0])
        @test_throws ArgumentError excited_configurations(rc"[Kr] 5s2 5p6", max_occupancy=[2,0])

        @test_throws ArgumentError excited_configurations(rc"[Kr] 5s2 5p6", min_occupancy=[-1,0,0])
        @test_throws ArgumentError excited_configurations(rc"[Kr] 5s2 5p6", max_occupancy=[3,2,4])
    end

    @testset "Ion–continuum" begin
        @test ion_continuum(c"1s2", os"k[s-d]") == [c"1s2", c"1s ks", c"1s kp", c"1s kd"]
    end

    @testset "GRASP comparisons" begin
        @test excited_configurations(c"1s2", os"2[s-p]"...) ==
            [c"1s2", c"1s 2s", c"2s2", c"2p2"]
        @test excited_configurations(c"1s2", os"k[s-p]"...) ==
            [c"1s2", c"1s ks", c"ks2", c"kp2"]

        @test excited_configurations(rc"1s2", ros"2[s-p]"...) ==
            [rc"1s2", rc"1s 2s", rc"2s2", rc"2p-2", rc"2p- 2p", rc"2p2"]
        @test excited_configurations(rc"1s2", ro"2s") == [rc"1s2", rc"1s 2s", rc"2s2"]
        @test excited_configurations(rc"1s2", ro"2p-") == [rc"1s2", rc"2p-2"]

        #= rcsfgenerate input:
        *
        1
        2s(2,i)2p(2,i)

        2s,2p
        0 10
        0
        n
        =#
        compare_with_grasp("rcsf.out.1.txt") do
            csfs(rc"1s2 2s2" ⊗ rcs"2p2")
        end

        #= rcsfgenerate input:
        *
        0
        3s(1,i)3p(3,i)3d(4,i)

        3s,3p,3d
        0 50
        0
        n
        =#
        compare_with_grasp("rcsf.out.2.txt") do
            csfs(rc"3s1" ⊗ rcs"3p3" ⊗ rcs"3d4")
        end

        #= rcsfgenerate input:
        *
        2
        3s(1,i)3p(2,i)3d(2,i)

        3s,3p,3d
        1 51
        0
        n
        =#
        compare_with_grasp("rcsf.out.3.txt") do
            csfs(rc"[Ne]* 3s1" ⊗ rcs"3p2" ⊗ rcs"3d2")
        end

        #= rcsfgenerate input:
        *
        0
        1s(2,*)

        2s,2p
        0 20
        2
        n
        =#
        compare_with_grasp("rcsf.out.4.txt") do
            csfs(excited_configurations(rc"1s2", ro"2s", ro"2p-", ro"2p"))
        end

        #= rcsfgenerate input:
        *
        0
        1s(2,*)

        3s,3p,3d
        0 20
        2
        n
        =#
        compare_with_grasp("rcsf.out.5.txt") do
            excited_configurations(
                rc"1s2",
                ro"2s", ro"2p-", ro"2p", ro"3s", ro"3p-", ro"3p", ro"3d-", ro"3d"
            ) |> csfs
        end

        # TODO:

        #= rcsfgenerate input: (problematic alphas with semicolons etc)
        *
        0
        5f(5,i)

        5s,5p,5d,5f
        1 51
        0
        n
        =#
    end
end


text:
using SpecialFunctions, Plots, HCubature, LinearAlgebra, Random

function n_ball_volume_analytic(n_dims, r)
    V = π^(n_dims/2)*r^n_dims/gamma(n_dims/2 + 1)
    return V
end

function monte_carlo_sphere_test(n_samples, r, n_dims)
    #draw uniformly random points in n_space and see if distance to
    #origin is <= r
    Random.seed!(0)
    inside_ball = 0
    for i = 1 : n_samples
        if norm(rand(n_dims)) <= r
            inside_ball += 1
        end
    end
    ratio = inside_ball/n_samples
    volume = ratio*(2*r)^n_dims
    return volume
end

function cubature_sphere_volume(n_samples, r, n_dims)
    a = -ones(n_dims-1) .+ eps()
    b = ones(n_dims-1) .- eps()
    function f(x)
        inside_bit = r^2 - sum(x.^2)
        if inside_bit <= 0
            return 0.0
        else
            return sqrt(r^2 - sum(x.^2))
        end
    end
    I, e = hcubature(f, a, b, maxevals = n_samples)
    return 2*I
end

r = 1.0
n_dims = 2:15

V = n_ball_volume_analytic.(n_dims, r)

p1 = plot(n_dims, V, xlabel = "dimension", ylabel = "volume", label = "analytic")

n_samples = Int(1e6)

V_mc = monte_carlo_sphere_test.(n_samples, r, n_dims)

plot!(p1, n_dims, V_mc, xlabel = "dimension", ylabel = "volume", label = "mc")

V_cubature = cubature_sphere_volume.(n_samples, r, n_dims)

plot!(p1, n_dims, V_cubature, ylim = [0, 6], xlabel = "dimension", ylabel = "volume", label = "cubature")

p2 = plot(n_dims, abs.(V_mc .- V)./V, yaxis = :log, xlabel = "dimension", ylabel = "relative abs error", label = "mc")

plot!(p2, n_dims, abs.(V_cubature .- V)./V, yaxis = :log, xlabel = "dimension", ylabel = "relative abs error", label = "cubature")

plot(p1, p2)


text:
@inline function update!(
    pcs::PotentialCalculationSetup{T},
    kernel,
    ndrange;
    use_nthreads::Int = Base.Threads.nthreads(),
    depletion_handling::Val{depletion_handling_enabled} = Val{false}(),
    is_weighting_potential::Val{_is_weighting_potential} = Val{false}(),
    only2d::Val{only_2d} = Val{false}()
)::Nothing where {T, only_2d, depletion_handling_enabled, _is_weighting_potential}
    update_even_points = true
    wait(kernel(
        pcs.potential, pcs.imp_scale, pcs.point_types, pcs.volume_weights, pcs.q_eff_imp, pcs.q_eff_fix, pcs.ϵ_r,
        pcs.geom_weights, pcs.sor_const, update_even_points, depletion_handling_enabled, _is_weighting_potential, only_2d,
        ndrange = ndrange
    ))
    apply_boundary_conditions!(pcs, Val(update_even_points), only2d)
    update_even_points = false
    wait(kernel(
        pcs.potential, pcs.imp_scale, pcs.point_types, pcs.volume_weights, pcs.q_eff_imp, pcs.q_eff_fix, pcs.ϵ_r,
        pcs.geom_weights, pcs.sor_const, update_even_points, depletion_handling_enabled, _is_weighting_potential, only_2d,
        ndrange = ndrange
    ))
    apply_boundary_conditions!(pcs, Val(update_even_points), only2d)
    return nothing
end


text:
using MPI, Test
include("../testhelpers.jl")

@testset "Diagnostics" begin
  tests = [
    (2,"dycoms_sin_test.jl")
   ]

  runmpi(tests, @__FILE__)
end

max_stars_repo_name: emassoud/CLIMA

max_issues_repo_name: emassoud/CLIMA

max_forks_repo_name: emassoud/CLIMA


text:
include("../src/strip2.jl")
NN = 1024
pix_size = 0.5

# White noise: evaluation of σ
dc = 0.35
observation = 24 # months
cmb = 1.1 # K at 95 GHz
system = 8.5 # K
atmosphere = 25 # K
bandwidth  = 31.5 # GHz
Tsys = cmb + system + atmosphere
Number_of_sensors = 1000.0
time_per_pixel = (observation * 30 * 24 * 60 * 60 * dc) / (NN^2)  # sec/pix

# Ma osservo 10 volte
time_per_pixel /=  10.0

sigma = ((Tsys / sqrt(31.5*1E9 * time_per_pixel )) / sqrt(Number_of_sensors)) * 1e6

# SZ Effect
N_SZ_Clusters  = 100
A_SZ_Clusters = 150.0
SZ_beta = 0.86
SZ_Theta_core = 1.0

sz_map = strip2.SZ_source(NN, pix_size, N_SZ_Clusters,
                   A_SZ_Clusters, SZ_beta,
                   SZ_Theta_core)

# Overall map (cosmological sources and noise) convolved with the
# telescope's beam
conv_map = sz_map

dir=false

mappa = zeros(NN, NN)
hit   = zeros(NN, NN)
N_baselines = 10

dataset_baselines = Array{Main.strip2.TOD_fake_pointing, 1}(undef, N_baselines)

for i = 1:N_baselines
    dataset = strip2.observe_sky(NN, conv_map; noise=true, direction_l_r=dir)
    m, h = strip2.binned_map(NN, dataset.time_order_data, dataset.pointing)
    mappa .+= m
    hit   .+= h
    dataset_baselines[i] = dataset
    global dir =  true ⊻ dir
end

tod = Array{Float64, 1}(undef, 0)
puntamenti = Array{Int64}(undef, 0)
for i in dataset_baselines
    append!(tod, i.time_order_data)
    append!(puntamenti, i.pointing[:,1])
end

for i in dataset_baselines
    append!(puntamenti, i.pointing[:,2])
end

puntamenti = reshape(puntamenti, NN*NN*10, 2)


text:
#=
These functions and types are designed to import 3D files of various formats. Currently only neurolucida3 is supported.

This will mostly address the functionality that was available in Neuron's import3d hoc files. The Import3d had lots of GUI stuff intertwined in it, though and I'm not planning on doing anything GUI related anytime soon, but I'd like to keep it separate if I do
=#

abstract Import3D

type Section3D
    parent::Array{Section3D,1}
    children::Array{Section3D,1}
    childind::Array{Int64,1}
    mytype::Int64 #Cellbody=1,Axon=2,Dendrite=3,Apical=4
    raw::Array{Float64,2}
    d::Array{Float64,1}
end

Section3D(ID::Int64)=Section3D(Array(Section3D,0),Array(Section3D,0),Array(Int64,0),ID,Array(Float64,0,3),Array(Float64,0))

type curxyz
    x::Array{Float64,1}
    y::Array{Float64,1}
    z::Array{Float64,1}
end

curxyz()=curxyz(Array(Float64,0),Array(Float64,0),Array(Float64,0))

type nlcda3 <: Import3D
    sections::Array{Section3D,1}
    mytypes::Array{Int64,1} # 4x1 array to tally total num of each section type
    file::Array{ByteString,1}
    curxyz::curxyz
    opensecs::Array{Int64,1} #inds of last section at that depth
    depth::Int64
end

function input(morphology::ASCIIString)

    if contains(morphology,".asc")
        import3d=nlcda3(morphology)
    end

    parse_file(import3d)

    return import3d

    #should then "instantiate" to create Neuron object and return it
end

function instantiate(morphology::ASCIIString)

    import3d=input(morphology)

    instantiate(import3d)

end

function instantiate(import3d::Import3D)

    neuron=make_neuron(Prop0(),Array(Node,0),Array(Section,0),[Array(Int64,0) for i=1:4])

    somaind=Array(Int64,0)
    rootind=Array(Int64,0)
    childinds=Array(Array{Int64,1},0)
    mapping=Dict{Int64,Int64}(0=>0)
    #skip somas, we'll do them later
    for i=1:length(import3d.sections)
        if import3d.sections[i].mytype>1
            if length(import3d.sections[i].parent)==0
                push!(rootind,i)
            end
            sec=Section(import3d.sections[i])
            add_sec(neuron,sec)
            push!(childinds,import3d.sections[i].childind)
            mapping[i]=length(neuron.secs)
        elseif import3d.sections[i].mytype==1
            push!(somaind,i)
        end
    end

    connect2soma(import3d,neuron,somaind,rootind,mapping)

    for i=1:length(childinds)
        for j in childinds[i]
            push!(neuron.secs[i].child,mapping[j])
        end
    end

    find_parents!(neuron)

    neuron

end

function connect2soma(import3d::Import3D,neuron::Neuron,somaind::Array{Int64,1},rootind::Array{Int64,1},mapping::Dict{Int64,Int64})

    #We need to
    #1) determine how many somas are described in file
    #2) assign a sensical 3D shape to it
    #3) find the centroid of that 3D shape

    centroids=Array(Float64,3,0)
    somas=Array(Section,0)
    #Determine how soma is described in this file

    if reduce(&, [mostly_constantz(import3d,somaind[i]) for i=1:length(somaind)])
        #if soma contours have a constant z

        if length(somaind)==1
            #approximate as sphere with diameter taken from section

            centroids=hcat(centroids,sphere_approx(import3d, somaind[1],somas))

        else
            overlaps=stack_overlap(import3d, somaind) #find which soma sections overlap
            for i=1:length(overlaps)
                if length(overlaps[i])==1 #if a soma doesn't overlap with any other sections

                    #approximate as sphere
                    centroids=hcat(centroids,sphere_approx(import3d, overlaps[i],somas))
                else
                    #multiple outlines of cell at different z positions - "stack of pancakes"
                    #find principle axis through stack
                    #approximate as "carrot" looking thing along principle axis
                    #need to make new soma Section
                end
            end
        end

    else
        #xyzd measurements along estimated centroid - "stack of frusta"
    end

    ds=zeros(Float64,length(somas))

    for i in [mapping[k] for k in rootind]

        for j=1:length(somas)
            ds[j]=sqrt((centroids[1,j]-neuron.secs[i].pt3d[1].x)^2 + (centroids[2,j]-neuron.secs[i].pt3d[1].y)^2 + (centroids[3,j]-neuron.secs[i].pt3d[1].z)^2)
        end

        push!(somas[indmin(ds)].child,i)
    end

    append!(neuron.secs,somas)

    nothing

end

function mostly_constantz(import3d::Import3D, ind::Int64)
    if std(import3d.sections[ind].raw[:,3])/mean(import3d.sections[ind].raw[:,3])<0.1
        return true
    else
        return false
    end
end

function stack_overlap(import3d::Import3D, somaind::Array{Int64,1})
    #this is the ugliest thing i have ever created. shame on me

    n = length(somaind)
    combos=[[list[c[i]-i+1] for i=1:length(c)] for c in combinations(collect(1:(n+k-1)),2)]

    overlaps=Array(Array{Int64,1},1)
    overlaps[1]=[1]
    for i=1:length(combos)
        bb1=boundbox(import3d, combos[i][1])
        bb2=boundbox(import3d, combos[i][2])

        if max(bb1[1],bb2[1])<min(bb1[2],bb2[2])
            if max(bb1[3],bb2[3])<min(bb1[4],bb2[4]) #if bounding boxes of these slices overlap
                flag=false
                for j=1:length(overlaps) #check groups of overlapping slices from previous iterations
                    if length(intersect(overlaps[j],combos[i]))>0
                        append!(overlaps[j],setdiff(combos[i],a[j]))
                        flag=true
                    end
                end
                if flag==false #if no group has yet been determined, assign to new group
                    append!(overlaps,combos[i])
                end
            end
        end
    end

    return overlaps

end

function boundbox(import3d::Import3D, ind::Int64)

    xmin=min(import3d.sections[ind].raw[:,1])
    xmax=max(import3d.sections[ind].raw[:,1])
    ymin=min(import3d.sections[ind].raw[:,2])
    ymax=max(import3d.sections[ind].raw[:,2])

    bb=collect(xmin, xmax, ymin, ymax)

end

function sphere_approx(import3d::Import3D, ind::Int64,somas::Array{Section,1})
    #going to make 3 3D points to describe 1 constant z section through the soma as a cylinder with L=D
    #trace perimeter
    perimeter=0.0
    for i=2:size(import3d.sections[ind].raw,1)
        perimeter+=sqrt((import3d.sections[ind].raw[i,1]-import3d.sections[ind].raw[i-1,1])^2 +
                        (import3d.sections[ind].raw[i,2]-import3d.sections[ind].raw[i-1,2])^2)
    end

    perimeter+=sqrt((import3d.sections[ind].raw[1,1]-import3d.sections[ind].raw[end,1])^2 +
                    (import3d.sections[ind].raw[1,2]-import3d.sections[ind].raw[end,2])^2)

    radi=perimeter/(2*pi)

    centroid=[mean(import3d.sections[ind].raw[:,1]); mean(import3d.sections[ind].raw[:,2]); mean(import3d.sections[ind].raw[:,3])]

    mypoints=[Pt3d(centroid[1]-radi,centroid[2]-radi,centroid[3],2*radi,0.0);
    Pt3d(centroid...,2*radi,.5); Pt3d(centroid[1]+radi,centroid[2]+radi,centroid[3],2*radi,1.0)]

    sec=Section(1,1:1,Array(Int64,0),0,mypoints,2*radi,Prop0)
    push!(somas,sec)

    centroid
end

#=
Neurolucida file types
=#

function nlcda3(filename::ASCIIString)
    f = open(filename)
    a=readlines(f)
    close(f)
    nlcda3(Array(Section3D,0),zeros(Int64,4),a,curxyz(),zeros(Int64,0),0)
end

const markers=["Dot","OpenStar","FilledQuadStar","CircleArrow","OpenCircle","DoubleCircle",
         "OpenQuadStar","CircleCross","Cross","Circle1","Flower3","Plus","Circle2",
         "Pinwheel","OpenUpTriangle","Circle3","TexacoStar","OpenDownTriangle","Circle4",
         "ShadedStar","OpenSquare","Circle5","SkiBasket","Asterisk","Circle6","Clock",
         "OpenDiamond","Circle7","ThinArrow","FilledStar","Circle8","ThickArrow","FilledCircle",
         "Circle9","SquareGunSight","FilledUpTriangle","Flower2","GunSight","FilledDownTriangle",
         "SnowFlake","TriStar","FilledSquare","OpenFinial","NinjaStar","FilledDiamond",
         "FilledFinial","KnightsCross","Flower","MalteseCross","Splat"]

const nonsense=["Name", "ImageCoords","Thumbnail","Color","Sections","SSM","dZI","Normal",
                "Low","High","Generated","Incomplete","SSM2","Resolution","\"CellBody\""]

function parse_file{T<:nlcda3}(nlcda::T)
    linenum=1
    leftpar=0
    rightpart=0
    skip=0
    state=0
    while linenum < length(nlcda.file)
        leftpar=length(matchall(r"\(",nlcda.file[linenum]))
        rightpar=length(matchall(r"\)",nlcda.file[linenum]))
        if leftpar>0
            if contains(nlcda.file[linenum],"(CellBody)")
                state=1
                newsec(nlcda,state)
                newparent(nlcda)
            elseif contains(nlcda.file[linenum],"(Axon)")
                state=2
                newsec(nlcda,state)
                newparent(nlcda)
            elseif contains(nlcda.file[linenum],"(Dendrite)")
                state=3
                newsec(nlcda,state)
                newparent(nlcda)
            elseif contains(nlcda.file[linenum],"(Apical)")
                state=4
                newsec(nlcda,state)
                newparent(nlcda)
            elseif markerdetect(nlcda, linenum)
                skip+=1
            elseif nonsensedetect(nlcda,linenum)
            else
                if ((leftpar-rightpar)>0) & (state>0) #check for new section
                    nlcda.depth+=1
                    newsec(nlcda,state)
                    newchild(nlcda)
                else
                    mynums=matchall(r"[-+]?[0-9]*\.?[0-9]+", nlcda.file[linenum])
                    if (state>0) & (length(mynums)>3) & (skip<1)
                        dimadd(nlcda,mynums)
                    end
                end

            end

        elseif rightpar>0
            if skip>0
                skip-=1
            else
                closesec(nlcda)
            end
        elseif contains(nlcda.file[linenum], " |")
            closesec(nlcda)
            nlcda.depth+=1
            newsec(nlcda,state)
            newchild(nlcda)
        else #no parenthesis so just ignore

        end
            linenum+=1
    end
    nothing
end

function newsec(nlcda::nlcda3,state::Int64)
    push!(nlcda.sections,Section3D(state)) #add new section to overall list of 3D
    nlcda.mytypes[state]+=1
    nothing
end

function newparent(nlcda::nlcda3)
    nlcda.curxyz=curxyz()
    nlcda.depth+=1
    push!(nlcda.opensecs,length(nlcda.sections))
    nothing
end

function newchild(nlcda::nlcda3) #new branch off of main section
    nlcda.sections[end-1].raw=vcat(nlcda.sections[end-1].raw,[nlcda.curxyz.x nlcda.curxyz.y nlcda.curxyz.z]) #add points that have accumulated to most recent section
    #make first point equal to most recent data point (where branch was)
    nlcda.sections[end].raw=zeros(Float64,1,3)
    for i=1:3; nlcda.sections[end].raw[1,i]=nlcda.sections[nlcda.opensecs[nlcda.depth-1]].raw[end,i]; end
    nlcda.sections[end].d=[nlcda.sections[nlcda.opensecs[nlcda.depth-1]].d[end]]

    push!(nlcda.sections[nlcda.opensecs[nlcda.depth-1]].children,nlcda.sections[end]) #add as child to parent
    push!(nlcda.sections[nlcda.opensecs[nlcda.depth-1]].childind,length(nlcda.sections))

    push!(nlcda.sections[end].parent, nlcda.sections[nlcda.opensecs[nlcda.depth-1]]) #add as parent to child

    nlcda.curxyz=curxyz() #reset xyz container
    push!(nlcda.opensecs,length(nlcda.sections)) #add new section to list of open

    nothing
end

function closesec(nlcda::nlcda3)
    nlcda.sections[end].raw=vcat(nlcda.sections[end].raw,[nlcda.curxyz.x nlcda.curxyz.y nlcda.curxyz.z])
    nlcda.depth-=1
    nlcda.curxyz=curxyz()
    pop!(nlcda.opensecs)

    nothing
end

function dimadd(nlcda::nlcda3,mynums::Array{SubString{UTF8String},1})
    push!(nlcda.curxyz.x,float(mynums[1]))
    push!(nlcda.curxyz.y,float(mynums[2]))
    push!(nlcda.curxyz.z,float(mynums[3]))
    push!(nlcda.sections[end].d,float(mynums[4]))
    nothing
end

function markerdetect(nlcda::nlcda3, linenum::Int64)

    for mark in markers
        if contains(nlcda.file[linenum],mark)
            return true
        end
    end

    return false
end

function nonsensedetect(nlcda::nlcda3,linenum::Int64)
    for non in nonsense
        if contains(nlcda.file[linenum],non)
            return true
        end
    end

    return false
end

function Section(section3d::Section3D) #like new_section

    mypt3d=Array(Pt3d,size(section3d.raw,1))
    mylength=0.0
    #add 3d points from 3d
    for i=1:length(section3d.d)
        if i>1
            l=sqrt((section3d.raw[i,1]-section3d.raw[i-1,1])^2+(section3d.raw[i,2]-section3d.raw[i-1,2])^2+(section3d.raw[i,3]-section3d.raw[i-1,3])^2)
            mypt3d[i]=Pt3d(section3d.raw[i,:]...,section3d.d[i],l+mypt3d[i-1].arc)
            mylength+=l
        else
            mypt3d[i]=Pt3d(section3d.raw[i,:]...,section3d.d[i],0.0)
        end
    end

    #normalize arc length
    for i=2:size(mypt3d,1)
        mypt3d[i]=Pt3d(mypt3d[i].x,mypt3d[i].y,mypt3d[i].z,mypt3d[i].d,mypt3d[i].arc/mylength)
    end

    sec=Section(section3d.mytype,1:1,Array(Int64,0),0,mypt3d,mylength,Prop0)
end

max_stars_repo_name: paulmthompson/JNeuron

max_issues_repo_name: paulmthompson/JNeuron

max_forks_repo_name: paulmthompson/JNeuron


text:
module SSPS

using Gen
using GLMNet
using ArgParse
using JSON

export julia_main

include("dbn_preprocess.jl")
include("dbn_models.jl")
include("dbn_proposals.jl")
include("mcmc_inference.jl")
include("state_updates.jl")

@load_generated_functions()

"""
Given the results of a previous run, recover the
state of the markov chain and store it in the given
choicemap object.
"""
function recover_state!(cmap, samples)

    # Recover the lambda values
    for (j, svec) in enumerate(samples["lambda"])
        cmap[:lambda_vec => j => :lambda] = svec[end][2]
    end

    # Recover the parent sets
    for (j, pdict) in enumerate(samples["parent_sets"])
        pset = Int64[]
        for (k,v) in pdict
            if v[end][2] == 1
                append!(pset, parse(Int64, k))
            end
        end
        sort!(pset)
        cmap[:parent_sets => j => :parents] = pset
    end

end

"""
parse the script's arguments using ArgParse.
"""
function parse_script_arguments()

    s = ArgParseSettings()
    @add_arg_table! s begin
        "timeseries_filename"
            help = "name of timeseries data file"
            required = true
        "ref_graph_filename"
            help = "name of reference graph data file"
            required = true
        "output_path"
            help = "name of output JSON file"
            required = true
        "timeout"
            help = "execution timeout (in seconds). When reached, terminate and output results as they are."
            arg_type = Float64
        "--thinning"
            help = "number of proposals to make for each sample taken"
            arg_type = Int64
            default=1
        "--large-indeg"
            help = "an approximate upper bound on the number of parents for any node"
            arg_type = Float64
            default = 20.0
        "--lambda-max"
            help = "hyperparameter for lambda prior"
            default = 15.0
        "--regression-deg"
            help = "regression polynomial degree in Gaussian DBN (default: full)"
            arg_type = Int64
            default = -1
        "--lambda-prop-std"
            help = "standard deviation of lambda's proposal distribution"
            arg_type = Float64
            default = 0.25
        "--n-steps"
            help = "Terminate a markov chain after it runs this many steps."
            arg_type = Int64
            default = -1
        "--proposal"
            help = "(DEPRECATED. Future versions of SSPS will only use 'sparse'.) The graph proposal distribution to use in MCMC: 'sparse' or 'uniform'."
            arg_type = String
            default = "sparse"
            range_tester = x -> x in ("sparse", "uniform")
        "--resume-sampling"
            help = "Continue sampling the markov chain stored in the given JSON file"
            required = false
            arg_type = String
            default = ""
        "--multithread"
            help = "Parallelize each iteration with multiple threads. The number of threads is specified by the `JULIA_NUM_THREADS` environment variable. This flag overrides the `--proposal` argument, setting it to 'sparse'."
            action = :store_true
    end

    args = parse_args(s)
    arg_vec = transform_arguments(args)

    return arg_vec
end

"""
Load arguments into a vector
"""
function transform_arguments(parsed_arg_dict)
    arg_vec = []
    push!(arg_vec, parsed_arg_dict["timeseries_filename"])
    push!(arg_vec, parsed_arg_dict["ref_graph_filename"])
    push!(arg_vec, parsed_arg_dict["output_path"])
    push!(arg_vec, parsed_arg_dict["timeout"])
    push!(arg_vec, parsed_arg_dict["n-steps"])
    push!(arg_vec, parsed_arg_dict["thinning"])
    push!(arg_vec, parsed_arg_dict["large-indeg"])
    push!(arg_vec, parsed_arg_dict["lambda-max"])
    push!(arg_vec, parsed_arg_dict["regression-deg"])
    push!(arg_vec, parsed_arg_dict["lambda-prop-std"])
    push!(arg_vec, parsed_arg_dict["proposal"])
    push!(arg_vec, parsed_arg_dict["resume-sampling"])
    push!(arg_vec, parsed_arg_dict["multithread"])

    return arg_vec
end

"""
A wrapper around the MCMC inference procedure
"""
function perform_inference(timeseries_filename::String,
                           ref_graph_filename::String,
                           output_path::String,
                           timeout::Float64,
                           n_steps::Int64,
                           thinning::Int64,
                           large_indeg::Float64,
                           lambda_max::Float64,
                           regression_deg::Int64,
                           lambda_prop_std::Float64,
                           proposal::String,
			   resume_sampling::String,
			   multithread::Bool)

    clear_caches()

    ts_vec, ref_ps = load_formatted_data(timeseries_filename,
                                         ref_graph_filename)

    # Some data preprocessing
    Xminus, Xplus  = combine_X(ts_vec)
    Xplus = standardize_X(Xplus)[1]
    Xminus, Xplus = vectorize_X(Xminus, Xplus)
    V = length(Xplus)

    println("Invoking SSPS on input files:\n\t",
            timeseries_filename, "\n\t", ref_graph_filename)

    # prepare parameters for proposal distributions
    ref_parent_counts = [max(sum(values(ps)), 2.0) for ps in ref_ps]
    update_loop_args = 1.0 ./ log2.(V ./ ref_parent_counts)
    push!(update_loop_args, lambda_prop_std)

    # Choose a lower bound for the lambda variable
    lambda_min = log(max(V/large_indeg - 1.0, exp(0.5)))

    # Choose the right update loop
    update_loop_fn = vertex_lambda_update_loop
    if proposal == "uniform"
        update_loop_fn = uniform_update_loop
    end
    if multithread
        update_loop_fn = multithread_update_loop
    end

    # Load observations into a choice map
    observations = Gen.choicemap()
    for (i, Xp) in enumerate(Xplus)
        observations[:Xplus => i => :Xp] = Xp
    end

    # If necessary, reload state from a previous session
    if resume_sampling != ""
        println("Resuming Markov Chain stored at ", resume_sampling)

        f = open(resume_sampling, "r")
        str = read(f, String)
        close(f)
        samples = JSON.parse(str)

        recover_state!(observations, samples)
    end

    # Prepare the Gen model's arguments
    model_args = (ref_ps,
                  Xminus,
                  lambda_min,
                  lambda_max,
                  regression_deg)

    t_start = time()

    update_results_args = [V, true]

    chain_id = string("pid_",getpid())

    results = mcmc_inference(vertex_lambda_dbn_model,
                             model_args, observations,
                             update_loop_fn,
                             update_loop_args,
                             update_results_storediff,
                             update_results_args;
                             timeout=timeout,
                             n_steps=n_steps,
                             thinning=thinning,
			     chain_id=chain_id)

    # Some last updates to the `results` object
    results["t_elapsed"] = time() - t_start
    delete!(results, "prev_parents")

    println("Saving results to JSON file:")

    f = open(output_path, "w")
    JSON.print(f, results)
    close(f)
    println("\t", output_path)

end

function julia_main()::Cint

    arg_vec = parse_script_arguments()
    perform_inference(arg_vec...)
    return 0
end

# END MODULE
end

max_stars_repo_name: gitter-lab/ssps

max_issues_repo_name: gitter-lab/ssps

max_forks_repo_name: gitter-lab/ssps


text:
include(joinpath(pwd(), "examples/implicit_dynamics/utils.jl"))
# include_implicit_dynamics()

# """
#     floating-base quadruped model
#         rigid body with force inputs via pre-specified contact sequence
# """
# struct QuadrupedFloatingBase{T}
#     dim::Dimensions
#
#     mass::T
#     inertia
#
#     gravity
# end
#
# nq = 6
# nu = 4 * 3
# nc = 0
# nw = 0
#
# model = QuadrupedFloatingBase(Dimensions(nq, nu, nc, nw),
#     1.0, Diagonal(ones(3)), [0.0; 0.0; 9.81])
#
# # eq. 14 http://roboticexplorationlab.org/papers/planning_with_attitude.pdf
# function attitude_jacobian(q)
# 	s = q[1]
# 	v = q[2:4]
#
# 	[-transpose(v);
# 	 s * I + skew(v)]
# end
#
# function G_func(q)
# 	quat = q[4:7]
# 	[1.0 0.0 0.0 0.0 0.0 0.0;
#      0.0 1.0 0.0 0.0 0.0 0.0;
# 	 0.0 0.0 1.0 0.0 0.0 0.0;
#      zeros(4, 3) attitude_jacobian(quat)]
# end
#
# function conjugate(q)
# 	s = q[1]
# 	v = q[2:4]
#
# 	return [s; -v]
# end
#
# function L_multiply(q)
# 	s = q[1]
# 	v = q[2:4]
#
# 	SMatrix{4,4}([s -transpose(v);
# 	              v s * I + skew(v)])
# end
#
# function R_multiply(q)
# 	s = q[1]
# 	v = q[2:4]
#
# 	SMatrix{4,4}([s -transpose(v);
# 	              v s * I - skew(v)])
# end
#
# function multiply(q1, q2)
# 	L_multiply(q1) * q2
# end
#
# # eq. 16 http://roboticexplorationlab.org/papers/maximal_coordinate_dynamics.pdf
# function ω_finite_difference(q1, q2, h)
# 	2.0 * multiply(conjugate(q1), (q2 - q1) ./ h)[2:4]
# end
#
# # Cayley map
# function cayley_map(ϕ)
# 	1.0 / sqrt(1.0 + norm(ϕ)^2.0) * [1.0; ϕ]
# end

# function dynamics(model::QuadrupedFloatingBase, h, q0, q1, q2, f, τ)
#
# 	p0 = q0[1:3]
#     quat0 = cayley_map(q0[4:6])
#
# 	p1 = q1[1:3]
# 	quat1 = cayley_map(q1[4:6])
#
# 	p2 = q2[1:3]
# 	quat2 = cayley_map(q2[4:6])
#
# 	# velocities
#     vm1 = (p1 - p0) / h[1]
#     vm2 = (p2 - p1) / h[1]
#
# 	ω1 = ω_finite_difference(quat0, quat1, h)
# 	ω2 = ω_finite_difference(quat1, quat2, h)
#
# 	d_linear = model.mass * (vm1 - vm2) - h[1] * model.mass * model.gravity
# 	d_angular = -1.0 * (model.inertia * ω2 * sqrt(4.0 / h[1]^2.0 - transpose(ω2) * ω2)
# 		+ cross(ω2, model.inertia * ω2)
# 		- model.inertia * ω1 * sqrt(4.0 / h[1]^2.0 - transpose(ω1) * ω1)
# 		+ cross(ω1, model.inertia * ω1))
#
# 	return [d_linear + f; d_angular - 2.0 * τ]
# end
#
# function residual(model::QuadrupedFloatingBase, z, θ, κ)
#     q0 = θ[1:nq]
#     q1 = θ[nq .+ (1:nq)]
#     h = θ[2nq .+ (1:1)]
#     f = [θ[2nq + 1 + (i - 1) * 3 .+ (1:3)] for i = 1:4]
#     p = [θ[2nq + 1 + 3 * 4 + (i - 1) * 3 .+ (1:3)] for i = 1:4]
#
#     q2 = z[1:nq]
#     p_body = q2[1:3]
#
#     τ1 = sum([cross(p[i] - p_body, f[i]) for i = 1:4])
#     f1 = sum(f)
#
#     return dynamics(model, h, q0, q1, q2, f1, τ1);
# end
#
# # generate residual methods
# nq = model.dim.q
# nu = model.dim.u
# nc = model.dim.c
# nz = nq
# nθ = nq + nq + 1 + 3 * 4 + 3 * 4
#
# # Declare variables
# @variables z[1:nz]
# @variables θ[1:nθ]
# @variables κ[1:1]
#
# # Residual
# r = residual(model, z, θ, κ)
# r = Symbolics.simplify.(r)
# rz = Symbolics.jacobian(r, z, simplify = true)
# rθ = Symbolics.jacobian(r, θ, simplify = true)
#
# # Build function
# r_func = eval(build_function(r, z, θ, κ)[2])
# rz_func = eval(build_function(rz, z, θ)[2])
# rθ_func = eval(build_function(rθ, z, θ)[2])
#
# rz_array = similar(rz, Float64)
# rθ_array = similar(rθ, Float64)
#
# u_gravity = [0.0; 0.0; 1.0 * model.mass * 9.81 / 4.0 * h[1]]
# u_stand = vcat([u_gravity for i = 1:4]...)
# z0 = copy(q_body)
# θ0 = [q_body; q_body; h; u_stand; pf1_ref[1]; pf2_ref[1]; pf3_ref[1]; pf4_ref[1]]
#
# # test dynamics
# include_implicit_dynamics()
#
# # options
# opts = InteriorPointOptions(
#    κ_init = 0.1,
#    κ_tol = 1.0e-4,
#    r_tol = 1.0e-8,
#    diff_sol = true)
#
#
# # solver
# ip = interior_point(z0, θ0,
#    r! = r_func, rz! = rz_func,
#    rz = similar(rz, Float64),
#    rθ! = rθ_func,
#    rθ = similar(rθ, Float64),
#    # idx_ineq = idx_ineq,
#    opts = opts)
#
# ip.methods.r! = r_func
# ip.methods.rz! = rz_func
# ip.methods.rθ! = rθ_func
#
# # simulate
# Tsim = 101 - 2
# q_hist = [q_body, q_body]
#
# for t = 1:Tsim
#     ip.z .= copy([q_hist[end]; 1.0 * ones(0)])
#     ip.θ .= copy([q_hist[end-1]; q_hist[end]; h; u_stand; pf1_ref[1]; pf2_ref[1]; pf3_ref[1]; pf4_ref[1]])
#     status = interior_point_solve!(ip)
#
#     if status
#         push!(q_hist, ip.z[1:nq])
#     else
#         println("dynamics failure t = $t")
#         println("res norm = $(norm(ip.r, Inf))")
#         break
#     end
# end
#
# visualize!(vis, model, q_hist,
#     [pf1_ref[1] for t = 1:T],
#     [pf2_ref[1] for t = 1:T],
#     [pf3_ref[1] for t = 1:T],
#     [pf4_ref[1] for t = 1:T],
#     Δt = h)
#

struct QuadrupedFB{I, T} <: Model{I, T}
      n::Int
      m::Int
      d::Int

      mass          # mass
      inertia       # inertia matrix
      inertia_inv   # inertia matrix inverse
      gravity       # gravity
end

function f(model::QuadrupedFB, z, u, w)
      # states
      x = view(z,1:3)
      r = view(z,4:6)
      v = view(z,7:9)
      ω = view(z,10:12)

      contact_mode = w[12 .+ (1:4)]

      # force in world frame
      F_world = [contact_mode[i] * view(u, (i - 1) * 3 .+ (1:3)) for i = 1:4]

      # torque in body frame
      τ = transpose(MRP(r...)) * sum([cross(w[(i - 1) * 3 .+ (1:3)] - x, F_world[i]) for i = 1:4])

      SVector{12}([v;
                   0.25 * ((1.0 - r' * r) * ω - 2.0 * cross(ω, r) + 2.0*(ω' * r) * r);
                   model.gravity + (1.0 / model.mass) * sum(F_world);
                   model.inertia_inv * (τ - cross(ω, model.inertia * ω))])
end

n, m, d = 12, 12, 16

mass = 1.0
inertia = Diagonal(@SVector[1.0, 1.0, 1.0])
inertia_inv = inv(inertia)
body_height = 0.2

f_min = 0.0
f_max = 10.0
μ_friction = 0.5
# a1
# mass = 108.0 / 9.81
# inertia = 10.0 * Diagonal(@SVector[0.017, 0.057, 0.064])
# inertia_inv = inv(inertia)
# body_height = 0.24

# laikago
# mass = 200.0 / 9.81
# inertia = Diagonal(@SVector[0.07335, 0.25068, 0.25447])
# inertia_inv = inv(inertia)
# body_height = 0.45

model = QuadrupedFB{Midpoint, FixedTime}(n, m, d,
                  mass,
                  inertia,
                  inertia_inv,
                  @SVector[0.0, 0.0, -9.81])

function fd(model::QuadrupedFB{Midpoint, FixedTime}, x, u, w, h, t)
	u_proj = vcat([soc_projection(u[(i - 1) * 3 .+ (1:3)], f_min, f_max, μ_friction) for i = 1:4]...)
	rz(z) = fd(model, z, x, u_proj, w, h, t) # implicit midpoint integration
	# x⁺ = newton(rz, copy(x))
    x⁺ = levenberg_marquardt(rz, copy(x))
	return x⁺
	# return view(x, 1:model.n) + h * f(model, view(x, 1:model.n) + 0.5 * h * f(model, view(x, 1:model.n), view(u, 1:model.m), w), view(u, 1:model.m), w)
end

fd(model, ones(model.n), zeros(model.m), zeros(model.d), h, 1)
function fdx(model::QuadrupedFB{Midpoint, FixedTime}, x, u, w, h, t)
    u_proj = vcat([soc_projection(u[(i - 1) * 3 .+ (1:3)], f_min, f_max, μ_friction) for i = 1:4]...)
    rz(z) = fd(model, z, x, u_proj, w, h, t) # implicit midpoint integration
    # x⁺ = newton(rz, copy(x))
    x⁺ = levenberg_marquardt(rz, copy(x))
    ∇rz = ForwardDiff.jacobian(rz, x⁺)
    rx(z) = fd(model, x⁺, z, u_proj, w, h, t) # implicit midpoint integration
    ∇rx = ForwardDiff.jacobian(rx, x)
    return -1.0 * ∇rz \ ∇rx
end

fdx(model, zeros(model.n), zeros(model.m), zeros(model.d), h, 1)

function fdu(model::QuadrupedFB{Midpoint, FixedTime}, x, u, w, h, t)
    u_proj = vcat([soc_projection(u[(i - 1) * 3 .+ (1:3)], f_min, f_max, μ_friction) for i = 1:4]...)
	u_proj_jac = cat([soc_projection_jacobian(u[(i - 1) * 3 .+ (1:3)], f_min, f_max, μ_friction) for i = 1:4]..., dims=(1,2))

    rz(z) = fd(model, z, x, u_proj, w, h, t) # implicit midpoint integration
    # x⁺ = newton(rz, copy(x))
    x⁺ = levenberg_marquardt(rz, copy(x))
    ∇rz = ForwardDiff.jacobian(rz, x⁺)
    ru(z) = fd(model, x⁺, x, z, w, h, t) # implicit midpoint integration
    ∇ru = ForwardDiff.jacobian(ru, u_proj)
    return -1.0 * ∇rz \ (∇ru * u_proj_jac)
    # return -1.0 * ∇rz \ ∇ru

	# f(z) = fd(model, x, z, w, h, t)
	# return ForwardDiff.jacobian(f, u)
end

fdu(model, zeros(model.n), zeros(model.m), zeros(model.d), h, 1)

##
x0 = [q_body; zeros(3); zeros(3)]
u_gravity = [0.0; 0.0; model.mass * 9.81 / 4.0]
u_stand = vcat([u_gravity for i = 1:4]...)
w_stand = [pf1_ref[1]; pf2_ref[2]; pf3_ref[3]; pf4_ref[4]; 1; 1; 1; 1]

soc_projection(u_gravity)
T_sim = 101
x_hist = [x0]
for t = 1:T_sim-1
    push!(x_hist, fd(model, x_hist[end], u_stand, w_stand, h, 1))
end

include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)
visualize!(vis, model, x_hist,
    [pf1_ref[1] for t = 1:T_sim],
    [pf2_ref[1] for t = 1:T_sim],
    [pf3_ref[1] for t = 1:T_sim],
    [pf4_ref[1] for t = 1:T_sim],
    Δt = h)

max_issues_repo_name: thowell/DirectMotionPlanning