text:
# WARNING: Work in progress
# E.g., probably needs to update function names/signatures
# Need way to optimize covar to be useful

using PDMats

function kernel_matern32(d::Float64; rho::Float64 = 1.0, sigmasq::Float64 = 1.0)
  x = abs(sqrt(3)*d/rho)
  sigmasq * (1+x) * exp(-x)
end

function dkerneldx_matern32(d::Float64; rho::Float64 = 1.0, sigmasq::Float64 = 1.0)
  x = abs(sqrt(3)*d/rho)
  -sigmasq * (-3*d/rho^2) * exp(-x)
end

function d2kerneldx2_matern32(d::Float64; rho::Float64 = 1.0, sigmasq::Float64 = 1.0)
  x = abs(sqrt(3)*d/rho)
  -sigmasq * (3/rho^2) * (1-x)*exp(-x)
end

function nuttall_kernel{T<:Real}(x::T; rho::T = 1.0)
   const   a0 = 0.355768
   const   a1 = 0.487396
   const   a2 = 0.144232
   const   a3 = 0.012604
   if abs(x)>rho return zero(T) end
   cpx = cospi(x/rho)
   spx = sinpi(x/rho)
   c2px = cpx*cpx-spx*spx
   return a0+a1*cpx+(a2+a3*cpx)*cpx*cpx-(a2+3*a3*cpx)*spx*spx
end

function nuttall_dkerneldx{T<:Real}(x::T; rho::T = 1.0)
   const   a0 = 0.355768
   const   a1 = 0.487396
   const   a2 = 0.144232
   const   a3 = 0.012604
   if abs(x)>rho return zero(T) end
   cpx = cospi(x/rho)
   spx = sinpi(x/rho)
   c2px = cpx*cpx-spx*spx
   s2px = 2*cpx*spx
   s3px = spx*c2px+cpx*s2px
   return -pi/rho*(a1*spx+2*a2*s2px+3*a3*s3px)
end

function nuttall_d2kerneldx2{T<:Real}(x::T; rho::T = 1.0)
   const   a0 = 0.355768
   const   a1 = 0.487396
   const   a2 = 0.144232
   const   a3 = 0.012604
   if abs(x)>rho return zero(T) end
   cpx = cospi(x/rho)
   spx = sinpi(x/rho)
   c2px = cpx*cpx-spx*spx
   c3px = cpx*(cpx*cpx-3*spx*spx)
   return -(pi/rho)^2*(a1*cpx+4*a2*c2px+9*a3*c3px)
end

function matern32_sparse_kernel{T<:Real}(x::T; rho::T = 1.0, sigmasq::T = 1.0)
  km = kernel_matern32(x,rho=rho,sigmasq=sigmasq)
  kn = nuttall_kernel(x,rho=rho*4)
  km*kn
end

function dkerneldx_matern32_sparse{T<:Real}(x::T; rho::T = 1.0, sigmasq::T = 1.0)
  km = kernel_matern32(x,rho=rho,sigmasq=sigmasq)
  kn = nuttall_kernel(x,rho=rho*4)
  dkmdx = dkerneldx_matern32(x,rho=rho,sigmasq=sigmasq)
  dkndx = nuttall_dkerneldx(x,rho=rho*4)
  km*dkndx+kn*dkmdx
end

function d2kerneldx2_matern32_sparse{T<:Real}(x::T; rho::T = 1.0, sigmasq::T = 1.0)
  km = kernel_matern32(x,rho=rho,sigmasq=sigmasq)
  kn = nuttall_kernel(x,rho=rho*4)
  dkmdx = dkerneldx_matern32(x,rho=rho,sigmasq=sigmasq)
  dkndx = nuttall_dkerneldx(x,rho=rho*4)
  d2kmdx2 = d2kerneldx2_matern32(x,rho=rho,sigmasq=sigmasq)
  d2kndx2 = nuttall_d2kerneldx2(x,rho=rho*4)
  km*d2kndx2+2*dkndx*dkmdx+kn*d2kmdx2
end

function make_kernel_data{T}(x::AbstractArray{T,1}, kernel::Function;
			sigmasq_obs::AbstractArray{T,1} = zeros(length(x)),
			sigmasq_cor::T = 1.0, rho::T = 1.0)
  @assert length(x) == length(sigmasq_obs)
  K = diagm(sigmasq_obs)
  for i in 1:size(K,1)
      for j in 1:size(K,2)
	      K[i,j] += kernel(x[i]-x[j], sigmasq=sigmasq_cor, rho=rho)
	  end
  end
  return PDMat(K)
end

function make_kernel_obs_pred{T}(xobs::AbstractArray{T,1}, xpred::AbstractArray{T,1}, kernel::Function;
			sigmasq_cor::T = 1.0, rho::T = 1.0)
  K = zeros(length(xobs),length(xpred))
  for i in 1:length(xobs)
      for j in 1:length(xpred)
	      K[i,j] += kernel(xobs[i]-xpred[j], sigmasq=sigmasq_cor, rho=rho)
	  end
  end
  return K
end

function make_kernel_matern32_data{T}(x::AbstractArray{T,1};
			sigmasq_obs::AbstractArray{T,1} = zeros(length(x)),
			sigmasq_cor::T = 1.0, rho::T = 1.0)
  make_kernel_data(x, kernel_matern32, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
end

function make_kernel_matern32_obs_pred{T}(xobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			#sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(x)),
			sigmasq_cor::T = 1.0, rho::T = 1.0)
  make_kernel_obs_pred(xobs, xpred, kernel_matern32, sigmasq_cor=sigmasq_cor, rho=rho)
end

function make_kernel_dmatern32dx_obs_pred{T}(xobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			#sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(x)),
			sigmasq_cor::T = 1.0, rho::T = 1.0)
  make_kernel_obs_pred(xobs, xpred, dkerneldx_matern32, sigmasq_cor=sigmasq_cor, rho=rho)
end

function make_kernel_d2matern32dx2_obs_pred{T}(xobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			#sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(x)),
			sigmasq_cor::T = 1.0, rho::T = 1.0)
  make_kernel_obs_pred(xobs, xpred, d2kerneldx2_matern32, sigmasq_cor=sigmasq_cor, rho=rho)
end

function make_kernel_matern32_sparse_data{T}(x::AbstractArray{T,1};
			sigmasq_obs::AbstractArray{T,1} = zeros(length(x)),
			sigmasq_cor::T = 1.0, rho::T = 1.0)
  make_kernel_data(x, matern32_sparse_kernel, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
end

function make_kernel_matern32_sparse_obs_pred{T}(xobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			#sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(x)),
			sigmasq_cor::T = 1.0, rho::T = 1.0)
  make_kernel_obs_pred(xobs, xpred, matern32_sparse_kernel, sigmasq_cor=sigmasq_cor, rho=rho)
end

function make_kernel_dmatern32dx_sparse_obs_pred{T}(xobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			#sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(x)),
			sigmasq_cor::T = 1.0, rho::T = 1.0)
  make_kernel_obs_pred(xobs, xpred, dkerneldx_matern32_sparse, sigmasq_cor=sigmasq_cor, rho=rho)
end

function make_kernel_d2matern32dx2_sparse_obs_pred{T}(xobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			#sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(x)),
			sigmasq_cor::T = 1.0, rho::T = 1.0)
  make_kernel_obs_pred(xobs, xpred, d2kerneldx2_matern32_sparse, sigmasq_cor=sigmasq_cor, rho=rho)
end

function predict_mean{T}(xobs::AbstractArray{T,1}, yobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(xobs)),	sigmasq_cor::T = 1.0, rho::T = 1.0)
  kobs = make_kernel_matern32_sparse_data(xobs, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
  kobs_pred = make_kernel_matern32_sparse_obs_pred(xobs,xpred, sigmasq_cor=sigmasq_cor, rho=rho)
  alpha = kobs \ yobs
  pred_mean = kobs_pred' * alpha
end

function predict_deriv{T}(xobs::AbstractArray{T,1}, yobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(xobs)),	sigmasq_cor::T = 1.0, rho::T = 1.0)
  kobs = make_kernel_matern32_sparse_data(xobs, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
  kobs_pred_deriv = make_kernel_dmatern32dx_sparse_obs_pred(xobs,xpred, sigmasq_cor=sigmasq_cor, rho=rho)
  alpha = kobs \ yobs
  pred_deriv = kobs_pred_deriv' * alpha
end

function predict_deriv2{T}(xobs::AbstractArray{T,1}, yobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(xobs)),	sigmasq_cor::T = 1.0, rho::T = 1.0)
  kobs = make_kernel_matern32_data(xobs, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
  kobs_pred_deriv2 = make_kernel_d2matern32dx2_sparse_obs_pred(xobs,xpred, sigmasq_cor=sigmasq_cor, rho=rho)
  alpha = kobs \ yobs
  pred_deriv = kobs_pred_deriv2' * alpha
end

function predict_mean_and_derivs{T}(xobs::AbstractArray{T,1}, yobs::AbstractArray{T,1}, xpred::AbstractArray{T,1};
			sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(xobs)),	sigmasq_cor::T = 1.0, rho::T = 1.0)
  kobs = make_kernel_matern32_data(xobs, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
  alpha = kobs \ yobs
  kobs_pred = make_kernel_matern32_sparse_obs_pred(xobs,xpred, sigmasq_cor=sigmasq_cor, rho=rho)
  pred_mean = kobs_pred' * alpha
  kobs_pred_deriv = make_kernel_dmatern32dx_sparse_obs_pred(xobs,xpred, sigmasq_cor=sigmasq_cor, rho=rho)
  pred_deriv = kobs_pred_deriv' * alpha
  kobs_pred_deriv2 = make_kernel_d2matern32dx2_sparse_obs_pred(xobs,xpred, sigmasq_cor=sigmasq_cor, rho=rho)
  pred_deriv2 = kobs_pred_deriv2' * alpha
  return (pred_mean, pred_deriv, pred_deriv2)
end

function gp_marginal{T}(xobs::AbstractArray{T,1}, yobs::AbstractArray{T,1};
			sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(xobs)),	sigmasq_cor::T = 1.0, rho::T = 1.0)
  kobs = make_kernel_matern32_sparse_data(xobs, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
  -0.5*( invquad(kobs, yobs) + logdet(kobs) + length(xobs)*log(2pi) )
end

function calc_gp_marginal_on_segments{T}(lambda::AbstractArray{T,1}, flux::AbstractArray{T,1};
                                sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(lambda)),	sigmasq_cor::T = 1.0, rho::T = 1.0,
                                half_chunck_size::Integer = 100)
  @assert length(lambda) == length(flux) == length(sigmasq_obs)
  println("# sigmasq_obs[1,1] = ", sigmasq_obs[1,1], " sigmasq_cor= ", sigmasq_cor, " rho= ", rho)
  output = 0.0
  num_seg = convert(Int64,ceil(length(lambda)/half_chunck_size)-1)
  for i in 1:num_seg
    idx_begin = 1+half_chunck_size*(i-1)
    idx_end = min(half_chunck_size*(i+1), length(lambda))
    write_idx_begin = idx_begin + div(half_chunck_size,2)
    write_idx_end = idx_end - div(half_chunck_size,2)
    if i==1 write_idx_begin=1 end
    if i==num_seg
      idx_begin = max(1,idx_end-2*half_chunck_size)
      write_idx_end=length(lambda)
    end
    #println("# i= ",i,": ", idx_begin, " - ", idx_end, " -> ", write_idx_begin, " - ", write_idx_end)
    #output[write_idx_begin:write_idx_end] = predict_gp(view(lambda,idx_begin:idx_end), view(flux,idx_begin:idx_end), view(lambda,write_idx_begin:write_idx_end), sigmasq_obs=view(sigmasq_obs,idx_begin:idx_end), sigmasq_cor=sigmasq_cor, rho=rho)
    output += gp_marginal(lambda[idx_begin:idx_end], flux[idx_begin:idx_end], sigmasq_obs=sigmasq_obs[idx_begin:idx_end], sigmasq_cor=sigmasq_cor, rho=rho)
  end
  output
end

function gp_marginal_wrapper(param::Vector)
   @assert length(param) == 2
   idx_min = 40000
   idx_max = 45000
   println("# wrapper: ", param)
   calc_gp_marginal_on_segments(lambda[idx_min:idx_max],vec(mean(obs[idx_min:idx_max],2)), sigmasq_obs = ones(length(lambda[idx_min:idx_max]))/150000,
                    sigmasq_cor = exp(param[1]), rho = exp(param[2]) )
end

function calc_doppler_component_gp{T}(lambda::AbstractArray{T,1}, flux::AbstractArray{T,1};
                                sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(lambda)),	sigmasq_cor::T = 1.0, rho::T = 1.0,
                                half_chunck_size::Integer = 100)
   lambda.*calc_gp_on_segments(predict_deriv,lambda, flux, sigmasq_obs=sigmasq_obs,	sigmasq_cor=sigmasq_cor, rho=rho, half_chunck_size=half_chunck_size)
end

function calc_gp_on_segments{T}(predict_gp::Function, lambda::AbstractArray{T,1}, flux::AbstractArray{T,1};
                                sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(lambda)),	sigmasq_cor::T = 1.0, rho::T = 1.0,
                                half_chunck_size::Integer = 100)
  @assert length(lambda) == length(flux) == length(sigmasq_obs)
  output = Array{Float64}(length(lambda))
  num_seg = convert(Int64,ceil(length(lambda)/half_chunck_size)-1)
  for i in 1:num_seg
    idx_begin = 1+half_chunck_size*(i-1)
    idx_end = min(half_chunck_size*(i+1), length(lambda))
    write_idx_begin = idx_begin + div(half_chunck_size,2)
    write_idx_end = idx_end - div(half_chunck_size,2)
    if i==1 write_idx_begin=1 end
    if i==num_seg
      idx_begin = max(1,idx_end-2*half_chunck_size)
      write_idx_end=length(lambda)
    end
    #println("# i= ",i,": ", idx_begin, " - ", idx_end, " -> ", write_idx_begin, " - ", write_idx_end)
    #output[write_idx_begin:write_idx_end] = predict_gp(view(lambda,idx_begin:idx_end), view(flux,idx_begin:idx_end), view(lambda,write_idx_begin:write_idx_end), sigmasq_obs=view(sigmasq_obs,idx_begin:idx_end), sigmasq_cor=sigmasq_cor, rho=rho)
    output[write_idx_begin:write_idx_end] = predict_gp(lambda[idx_begin:idx_end], flux[idx_begin:idx_end], lambda[write_idx_begin:write_idx_end], sigmasq_obs=sigmasq_obs[idx_begin:idx_end], sigmasq_cor=sigmasq_cor, rho=rho)
  end
  output
end

function calc_doppler_component_gp{T}(lambda::AbstractArray{T,1}, flux::AbstractArray{T,2};
                                sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(lambda)),	sigmasq_cor::T = 1.0, rho::T = 1.0)
  doppler_basis = calc_doppler_component_gp(lambda,vec(mean(flux,2)),sigmasq_obs=sigmasq_obs,sigmasq_cor=sigmasq_cor,rho=rho)
end

function calc_doppler_quadratic_term_gp{T}(lambda::AbstractArray{T,1}, flux::AbstractArray{T,1};
                                sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(lambda)),	sigmasq_cor::T = 1.0, rho::T = 1.0,
                                half_chunck_size::Integer = 100)
   0.5*lambda.^2.*calc_gp_on_segments(predict_deriv2,lambda, flux, sigmasq_obs=sigmasq_obs,	sigmasq_cor=sigmasq_cor, rho=rho, half_chunck_size=half_chunck_size)
end

function calc_doppler_quadratic_term_gp{T}(lambda::AbstractArray{T,1}, flux::AbstractArray{T,2};
                                     sigmasq_obs::AbstractArray{T,1} = 1e-16*ones(length(lambda)),	sigmasq_cor::T = 1.0, rho::T = 1.0)
  doppler_quad_term = calc_doppler_quadratic_term_gp(lambda,vec(mean(flux,2)),sigmasq_obs=sigmasq_obs,sigmasq_cor=sigmasq_cor,rho=rho)
end

function test_gp_deriv()
  #lambda_small = collect(linspace(1000.0,1100.0,101))
  #yobs_small = sin(2pi*lambda/40)+ 0.3*randn(length(lambda_small))
  #lambda_pred = collect(linspace(1000.0,1100.0,201))
  lambda_small = lambda[10000:10100]
  yobs_small = obs[10000:10100,1]
  lambda_pred = lambda_small
  (f,df,df2) = predict_mean_and_derivs(lambda_small,yobs_small,lambda_pred, sigmasq_obs=1.0/150000.0*ones(length(lambda_small)), sigmasq_cor=0.1,rho = 0.01)
  plot(lambda_small,(yobs_small),"r.")
  plot(lambda_pred,f,"b-")
  plot(lambda_pred,df,"g-")
  #plot(lambda_pred,df2,"m-")
end

#=
test_gp_deriv()

doppler_comp = calc_doppler_component(lambda,obs, sigmasq_obs=1/150000.0*ones(length(lambda)), rho=0.1)
genpca_out = fit_gen_pca_eford(obs,doppler_comp)

doppler_quadratic_term = calc_doppler_quadratic_term(lambda,obs, sigmasq_obs=1/150000.0*ones(length(lambda)), rho=0.1)
genpca_quad_out = fit_gen_pca_eford(obs,doppler_comp,qaudratic_term=doppler_quadratic_term)
=#

max_stars_repo_name: eford/RVSpectraKitLearn

max_issues_repo_name: eford/RVSpectraKitLearn

max_forks_repo_name: eford/RVSpectraKitLearn


text:
module BenchmarkEnvironments

using DecisionMakingEnvironments
using RecipesBase, StaticArrays, Distributions, LinearAlgebra
using DiffEqBase, OrdinaryDiffEq

include("utils.jl")

# Discrete environments
include("discrete/discrete.jl")
export simple_chain, simple_chain_finitetime, hard_chain_finitetime

#continuous state environments
include("continuous/continuous.jl")
export acrobot_finitetime, cartpole_finitetime, mountaincar_finitetime
export pinball_finitetime, pinball_box, pinball_empty, pinball_easy, pinball_medium, pinball_hard
export bicycle_balance, bicycle_goal
export hiv
export pendulum_finitetime
export ballandbeam_fixedgoal, ballandbeam_randomgoal, ballandbeam_tacking
end

ext: jl


text:
using Test
using Base: ImmutableDict
using BangBang
using BangBang: ismutable
using StaticArrays: SVector

macro test_error(ex)
    quote
        let err = nothing
            try
                $(esc(ex))
            catch err
            end
            @test err isa Exception
            err
        end
    end
end

macro test_inferred(ex)
    esc(:($Test.@test (($Test.@inferred $ex); true)))
end


text:
# TODO: Unit test
function _is_time_to_terminate(state::CurrentState, config::AlgorithmConfig)
    if state.total_node_count >= config.node_limit
        return true
    elseif state.total_elapsed_time_sec >= config.time_limit_sec
        return true
    elseif _optimality_gap(state.primal_bound, state.dual_bound) <=
           config.gap_tol
        return true
    else
        return false
    end
end

function optimize!(
    form::DMIPFormulation,
    config::AlgorithmConfig,
    primal_bound::Float64 = Inf,
)::Result
    result = Result()
    # TODO: Model presolve. Must happen before initial state is built.
    # Initialize search tree with LP relaxation
    state = CurrentState(primal_bound = primal_bound)
    _log_preamble(form, primal_bound, config)
    while !isempty(state.tree)
        node = pop_node!(state.tree)
        node_result = process_node!(state, form, node, config)
        update_state!(state, form, node, node_result, config)
        _log_if_necessary(state, node_result, config)
        if _is_time_to_terminate(state, config)
            break
        end
    end
    result = Result(state, config)
    _log_postamble(result, config)
    return result
end

# TODO: Store config in CurrentState, remove as argument here.
function process_node!(
    state::CurrentState,
    form::DMIPFormulation,
    node::Node,
    config::AlgorithmConfig,
)::NodeResult
    node_result = NodeResult(node)
    # 0. Check if we can bail out early by pruning by bound
    if node.dual_bound > state.primal_bound
        node_result.status = PRUNED_BY_PARENT_BOUND
        return node_result
    end
    # 1. Build model
    populate_lp_model!(state, form, node, node_result, config)
    if node_result.status == INFEASIBLE_LP
        # When trying to formulate the LP, we in fact proved it was infeasible
        return node_result
    end
    set_basis_if_available!(state, node)

    # 2. Solve model
    model = state.gurobi_model
    MOI.optimize!(model)

    # 3. Grab solution data and bundle it into a NodeResult
    node_result.simplex_iters = MOI.get(model, MOI.SimplexIterations())
    node_result.depth = node.depth
    term_status = MOI.get(model, MOI.TerminationStatus())
    if term_status == MOI.OPTIMAL
        node_result.status = OPTIMAL_LP
        node_result.cost = MOI.get(model, MOI.ObjectiveValue())
        _update_lp_solution!(state, form)
        node_result.x = state.current_solution
        node_result.int_infeas = _num_int_infeasible(state, form, config)
    elseif term_status == MOI.INFEASIBLE
        node_result.status = INFEASIBLE_LP
        node_result.cost = Inf
    elseif term_status == MOI.DUAL_INFEASIBLE
        node_result.status = UNBOUNDED_LP
        node_result.cost = -Inf
    else
        error("Unexpected termination status $term_status at node LP.")
    end
    return node_result
end

# Only checks feasibility w.r.t. integrality constraints!
function _num_int_infeasible(
    state::CurrentState,
    form::DMIPFormulation,
    config::AlgorithmConfig,
)::Int
    cnt = 0
    for cvi in all_variables(form)
        v_set = get_variable_kind(form, cvi)
        if v_set === nothing
            continue
        end
        xi = state.current_solution[index(cvi)]
        ϵ = config.int_tol
        xi_f = _approx_floor(xi, ϵ)
        xi_c = _approx_ceil(xi, ϵ)
        if v_set isa ZO
            # Should have explicitly imposed 0/1 bounds in the formulation...
            # but assert just to be safe.
            @assert -ϵ <= xi <= 1 + ϵ
        end
        if min(abs(xi - xi_f), abs(xi_c - xi)) > ϵ
            # The variable value is more than ϵ away from both its floor and
            # ceiling, so it's fractional up to our tolerance.
            cnt += 1
        end
    end
    return cnt
end

function update_state!(
    state::CurrentState,
    form::DMIPFormulation,
    node::Node,
    node_result::NodeResult,
    config::AlgorithmConfig,
)
    state.total_node_count += 1
    state.total_simplex_iters += node_result.simplex_iters
    state.polling_state.period_node_count += 1
    state.polling_state.period_simplex_iters += node_result.simplex_iters
    state.on_a_dive = false
    # 1. Prune by infeasibility
    if node_result.status == INFEASIBLE_LP
        # Do nothing
        @assert node_result.cost == Inf
    elseif node_result.status == PRUNED_BY_PARENT_BOUND ||
           node_result.cost > state.primal_bound
        # 2. Prune by bound
        # Do nothing
    elseif node_result.status == UNBOUNDED_LP
        # 3. LP is unbounded.
        #  Implies MIP is infeasible or unbounded. Should only happen at root.
        @assert _is_root_node(node)
        @assert node_result.cost == -Inf
        state.primal_bound = -Inf
    elseif node_result.int_infeas == 0
        # 4. Prune by integrality
        # Have <= to handle case where we seed the optimal cost but
        # not the optimal solution
        @assert node_result.status == OPTIMAL_LP
        if node_result.cost <= state.primal_bound
            state.primal_bound = node_result.cost
            # TODO: Make this more efficient, keys should not change.
            copy!(state.best_solution, node_result.x)
        end
    else
        # 5. Branch!
        @assert node_result.status == OPTIMAL_LP
        children = branch(state, form, node, node_result, config)
        _store_basis_if_desired!(state, children, config)
        for child in children
            child.dual_bound = node_result.cost
            push_node!(state.tree, child)
        end
        # TODO: Can be even more clever with this and reuse the same model
        # throughout the tree. However, we currently update bounds based on a
        # diff with the root. So, after backtracking we will need to reset all
        # bounds, but can otherwise reuse the same model.
        state.on_a_dive = true
        # TODO: Add a check in this branch to ensure we don't have a "funny"
        # return status. This is a little kludgy since we don't necessarily
        # store the MOI model in node_result. Maybe need to add termination
        # status as a field...
    end
    state.rebuild_model = if state.on_a_dive
        (config.model_reuse_strategy == NO_MODEL_REUSE)
    else
        (config.model_reuse_strategy != USE_SINGLE_MODEL)
    end
    state.total_elapsed_time_sec = time() - state.starting_time
    delete!(state.warm_starts, node)
    return nothing
end

function _store_basis_if_desired!(
    state::CurrentState,
    children::NTuple{N,Node},
    config::AlgorithmConfig,
) where {N}
    @assert N == length(children) >= 2
    if config.warm_start_strategy == NO_WARM_STARTS
        # Do nothing
    else
        basis = get_basis(state)
        state.warm_starts[children[1]] = basis
        ending_index = (
            if config.warm_start_strategy == WARM_START_WHEN_BACKTRACKING
                N - 1
            else
                @assert config.warm_start_strategy == WARM_START_WHENEVER_POSSIBLE
                N
            end
        )
        for i in 2:ending_index
            state.warm_starts[children[i]] = copy(basis)
        end
    end
    return nothing
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

@doc raw"""Affinity is a group of affinity scheduling rules.

    IoK8sApiCoreV1Affinity(;
        nodeAffinity=nothing,
        podAffinity=nothing,
        podAntiAffinity=nothing,
    )

    - nodeAffinity::IoK8sApiCoreV1NodeAffinity : Describes node affinity scheduling rules for the pod.
    - podAffinity::IoK8sApiCoreV1PodAffinity : Describes pod affinity scheduling rules (e.g. co-locate this pod in the same node, zone, etc. as some other pod(s)).
    - podAntiAffinity::IoK8sApiCoreV1PodAntiAffinity : Describes pod anti-affinity scheduling rules (e.g. avoid putting this pod in the same node, zone, etc. as some other pod(s)).
"""
mutable struct IoK8sApiCoreV1Affinity <: SwaggerModel
    nodeAffinity::Any # spec type: Union{ Nothing, IoK8sApiCoreV1NodeAffinity } # spec name: nodeAffinity
    podAffinity::Any # spec type: Union{ Nothing, IoK8sApiCoreV1PodAffinity } # spec name: podAffinity
    podAntiAffinity::Any # spec type: Union{ Nothing, IoK8sApiCoreV1PodAntiAffinity } # spec name: podAntiAffinity

    function IoK8sApiCoreV1Affinity(;nodeAffinity=nothing, podAffinity=nothing, podAntiAffinity=nothing)
        o = new()
        validate_property(IoK8sApiCoreV1Affinity, Symbol("nodeAffinity"), nodeAffinity)
        setfield!(o, Symbol("nodeAffinity"), nodeAffinity)
        validate_property(IoK8sApiCoreV1Affinity, Symbol("podAffinity"), podAffinity)
        setfield!(o, Symbol("podAffinity"), podAffinity)
        validate_property(IoK8sApiCoreV1Affinity, Symbol("podAntiAffinity"), podAntiAffinity)
        setfield!(o, Symbol("podAntiAffinity"), podAntiAffinity)
        o
    end
end # type IoK8sApiCoreV1Affinity

const _property_map_IoK8sApiCoreV1Affinity = Dict{Symbol,Symbol}(Symbol("nodeAffinity")=>Symbol("nodeAffinity"), Symbol("podAffinity")=>Symbol("podAffinity"), Symbol("podAntiAffinity")=>Symbol("podAntiAffinity"))
const _property_types_IoK8sApiCoreV1Affinity = Dict{Symbol,String}(Symbol("nodeAffinity")=>"IoK8sApiCoreV1NodeAffinity", Symbol("podAffinity")=>"IoK8sApiCoreV1PodAffinity", Symbol("podAntiAffinity")=>"IoK8sApiCoreV1PodAntiAffinity")
Base.propertynames(::Type{ IoK8sApiCoreV1Affinity }) = collect(keys(_property_map_IoK8sApiCoreV1Affinity))
Swagger.property_type(::Type{ IoK8sApiCoreV1Affinity }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiCoreV1Affinity[name]))}
Swagger.field_name(::Type{ IoK8sApiCoreV1Affinity }, property_name::Symbol) =  _property_map_IoK8sApiCoreV1Affinity[property_name]

function check_required(o::IoK8sApiCoreV1Affinity)
    true
end

function validate_property(::Type{ IoK8sApiCoreV1Affinity }, name::Symbol, val)
end


text:
mutable struct TiffFile
    """A unique ID describing this file that is embedded in the XML"""
    uuid::String

    """The relative path to this file"""
    filepath::String

    io::Union{Stream, IOStream}

    """Location of the first IFD in the file stream"""
    first_offset::Int

    """Whether this file has a different endianness than the host computer"""
    need_bswap::Bool

    function TiffFile(io::Union{Stream, IOStream})
        file = new()
        file.io = io
        seekstart(io)
        # TODO: Parsing the filename from the IO name is likely to be fragile
        file.filepath = extract_filename(io)
        file.need_bswap = check_bswap(io)
        file.first_offset = Int(do_bswap(file, read(file.io, UInt32)))
        file
    end
end

function TiffFile(uuid::String, filepath::String)
    try
        file = TiffFile(open(filepath))
        file.uuid = uuid
        file.filepath = filepath
        return file
    catch e # the file probably got renamed
        (!isa(e, SystemError)) && rethrow(e)
        throw(FileIO.LoaderError("OMETIFF", "It looks like this file was renamed, "*
        "but has internal links with the original name. Please rename to $filepath "*
        "to load. See https://github.com/tlnagy/OMETIFF.jl/issues/14 for details."))
    end
end

struct IFD
    """Pointer to the file containing this IFD"""
    file::TiffFile

    """Location(s) in `file` of the data corresponding to this IFD"""
    strip_offsets::Vector{Int}
end

"""

Load all the files and run through them all to find the offsets of each IFD.
"""
function get_ifds(orig_file::TiffFile,
                  ifd_files::Array{Union{Tuple{String, String}, Nothing}})

    # open all files referenced by the tiffdatas
    files = Dict{Tuple{String, String}, TiffFile}()
    for item in unique(ifd_files)
        (item == nothing) && continue
        filepath = joinpath(dirname(orig_file.filepath), item[2])

        # if this file is the same as the base file, then don't open it again
        if filepath == orig_file.filepath
            files[item] = orig_file
            orig_file.uuid = item[1]
        # else open a new pointer to this file
        else
            files[item] = TiffFile(item[1], filepath)
        end
    end

    ifds = Array{Union{IFD, Nothing}}(nothing, length(ifd_files))

    for (fileid, file) in files
        ifd_offsets = collect(file)

        ifd_idx_in_file = 0
        for (idx, ifd_file) in enumerate(ifd_files)
            (ifd_file == nothing) && continue
            if ifd_file == fileid
                ifds[idx] = IFD(file, ifd_offsets[ifd_idx_in_file+=1])
            end
        end
    end
    files, ifds
end

"""
    load_master_xml(file::TiffFile)

Loads the master OME-XML file from `file` or from a linked file.
"""
function load_master_xml(file::TiffFile)
    omexml = loadxml(file)
    try
        # Check if the full OME-XML metadata is stored in another file
        metadata_file = findfirst("/ns:OME/ns:BinaryOnly", omexml, ["ns"=>namespace(omexml)])
        (metadata_file == nothing) && return omexml
        uuid, filepath = metadata_file["UUID"], joinpath(dirname(file.filepath), metadata_file["MetadataFile"])

        # we have a companion metadata file
        if endswith(filepath, ".companion.ome")
            xdoc = readxml(filepath)
            omexml = root(xdoc)
        else
            metadata_file = TiffFile(uuid, filepath)
            omexml = loadxml(metadata_file)
            close(metadata_file.io) # clean up
            return omexml
        end
    catch err
        isa(err, BoundsError) && return omexml
        rethrow(err)
    end
end

Base.eltype(::Type{TiffFile}) = Vector{Int}
Base.IteratorSize(::Type{TiffFile}) = Base.SizeUnknown()

Base.iterate(file::TiffFile) = iterate(file, (read_ifd(file, file.first_offset)))

"""
    iterate(file, state)

Given a `state`, returns a tuple with list of locations of the data strips corresponding to
that IFD in `file` and the updated state.
"""
function Base.iterate(file::TiffFile, state::Tuple{Union{Vector{Int}, Nothing}, Int})
    strip_locs, ifd = state
    # if current element doesn't exist, exit
    (strip_locs == nothing) && return nothing
    (ifd <= 0) && return (strip_locs, (nothing, 0))

    next_strip_locs, next_ifd = read_ifd(file, ifd)

    return (strip_locs, (next_strip_locs, next_ifd))
end

loadxml(file::TiffFile) = read_ifd(file, file.first_offset; getxml=true)

function read_ifd(file::TiffFile, offset::Int; getxml=false)
    seek(file.io, offset)

    number_of_entries = do_bswap(file, read(file.io, UInt16))

    strip_offset_list = Int[]
    strip_offset = 0
    strip_count = 0
    width = 0
    height = 0

    tag_info = Array{UInt16}(undef, 2)
    data_info = Array{UInt32}(undef, 2)

    for i in 1:number_of_entries
        read!(file.io, tag_info)
        read!(file.io, data_info)
        tag_id, tag_type = do_bswap(file, tag_info)
        data_count, data_offset = do_bswap(file, data_info)

        curr_pos = position(file.io)
        if tag_id == 256
            width = data_offset
        elseif tag_id == 257 # height of image
            height = data_offset
        elseif tag_id == 273 # offset in stream to first strip
            strip_offset = data_offset
            strip_count = data_count
        # number of rows per strip should be equal to the height of image for now
        elseif tag_id == 278
            rows_per_strip = data_offset
            strip_num = floor(Int, (height + rows_per_strip - 1) / rows_per_strip)

            # if the data is spread across multiple strips
            if strip_num > 1
                seek(file.io, strip_offset)
                strip_info = Vector{UInt32}(undef, strip_num)
                read!(file.io, strip_info)
                strip_offsets = do_bswap(file, strip_info)
                strip_offset_list = Int.(strip_offsets)
            else
                strip_offset_list = [Int(strip_offset)]
            end
        elseif getxml && tag_id == 270 # Image Description tag
            seek(file.io, data_offset)
            # strip null values from string
            _data = Array{UInt8}(undef, data_count)
            read!(file.io, _data)
            raw_str = replace(String(_data), "\0"=>"")
            # check if is xml since ImageJ display settings are also stored in
            # ImageDescription tags
            # TODO: This should be replaced with some proper validation
            if raw_str[1:5] == "<?xml"
                xdoc = parsexml(raw_str)
                return root(xdoc)
            end
        end
        seek(file.io, curr_pos)
    end

    next_ifd = Int(do_bswap(file, read(file.io, UInt32)))
    strip_offset_list, next_ifd
end

function load_comments(file)
    seek(file.io, 24)
    comment_header = Int(do_bswap(file, read(file.io, UInt32)))
    if comment_header != 99384722
        return ""
    end
    comment_offset = Int(do_bswap(file, read(file.io, UInt32)))
    seek(file.io, comment_offset)
    comment_header = read(file.io, UInt32)
    comment_length = Int(do_bswap(file, read(file.io, UInt32)))
    metadata_bytes = Array{UInt8}(undef, comment_length)
    read!(file.io, metadata_bytes)
    metadata = JSON.parse(String(metadata_bytes))
    if !haskey(metadata, "Summary")
        return ""
    end
    metadata["Summary"]
end

function do_bswap(file::TiffFile, values::AbstractArray)
    if file.need_bswap
        for i in 1:length(values)
            values[i] = bswap(values[i])
        end
    end
    values
end

do_bswap(file::TiffFile, value) = file.need_bswap ? bswap(value) : value


text:
function toRGBFloat(img::ImgArr)::ImgFlt
	c, w, h = size(img)
	ret = ImgFlt(undef, (3, w, h))
	ret .= (Float32.(img[1:3, :, :]) ./ 255)
	return ret
end

function toARGBBytes(img::ImgFlt)::ImgArr
	c, w, h = size(img)
	ret = ImgArr(undef, (4, w, h))
	ret[4, :, :] .= 0xff
	ret[1:3, :, :] .= UInt8.(round.(img .* 255))
	return ret
end

loadFloatImg(filepath::String)::ImgFlt = toRGBFloat(loadImg(filepath))

writeImg(filepath::String, img::ImgFlt) = writeImg(filepath, toARGBBytes(img))


text:
# PARALLEL example 2
N = 200000000
nheads = @parallel (+) for i = 1:N
    Int(rand(Bool))
end

println(nheads/N)

max_stars_repo_name: milankl/misc

max_issues_repo_name: milankl/misc

max_forks_repo_name: milankl/misc


text:
module IntervalSpecialFunctions

using IntervalArithmetic
using SpecialFunctions
using IntervalRootFinding

import IntervalArithmetic: @round, big53

import SpecialFunctions:
    erf, erfc,
    erfinv, erfcinv

include("erf.jl")

end # module


text:
"""
Case 6:
This case study a 19-state virtual synchronous machine against an infinite bus located at bus 1, with VSM located at bus 2.
The perturbation increase the reference power (analogy for mechanical power) from 0.5 to 0.7.
"""

##################################################
############### LOAD DATA ########################
##################################################

############### Data Network ########################

nodes_DAIB= [Bus(1 , #number
                 "Bus 1", #Name
                 "REF" , #BusType (REF, PV, PQ)
                 0, #Angle in radians
                 1.04, #Voltage in pu
                 (min=0.94, max=1.06), #Voltage limits in pu
                 0.69),  #Base voltage in kV
                 Bus(2 , "Bus 2"  , "PV" ,  0 , 1.0 , (min=0.94, max=1.06), 0.69)]

branch_DAIB = [Line("Line1", #name
                    true, #available
                    0.0, #active power flow initial condition (from-to)
                    0.0, #reactive power flow initial condition (from-to)
                    Arc(from=nodes_DAIB[1],to=nodes_DAIB[2]), #Connection between buses
                    0.0, #resistance in pu
                    0.075, #reactance in pu
                    (from=0.0, to=0.0), #susceptance in pu
                    5.0, #rate in MW
                    1.04)]  #angle limits (-min and max)

############### Data devices ########################

inf_gen_DAIB = StaticSource(1, #number
                    :InfBus, #name
                    nodes_DAIB[1],#bus
                    1.00, #VR
                    0.0, #VI
                    0.000005) #Xth

############### Inverter Data ########################

converter = AvgCnvFixedDC(690.0, #Rated Voltage
                          2.75) #Rated MVA

dc_source = FixedDCSource(600.0) #Not in the original data, guessed.

filt = LCLFilter(0.08, #Series inductance lf in pu
                   0.003, #Series resitance rf in pu
                   0.074, #Shunt capacitance cf in pu
                   0.2, #Series ractance rg to grid connection (#Step up transformer or similar)
                   0.01) #Series resistance lg to grid connection (#Step up transformer or similar)

pll = PLL(500.0, #ω_lp: Cut-off frequency for LowPass filter of PLL filter.
          0.084, #k_p: PLL proportional gain
          4.69) #k_i: PLL integral gain

virtual_H = VirtualInertia(2.0, #Ta:: VSM inertia constant
                           400.0, #kd:: VSM damping coefficient
                           20.0, #kω:: Frequency droop gain in pu
                           2*pi*50.0) #ωb:: Rated angular frequency

Q_control = ReactivePowerDroop(0.2, #kq:: Reactive power droop gain in pu
                              1000.0) #ωf:: Reactive power cut-off low pass filter frequency

outer_control = VirtualInertiaQdroop(virtual_H, Q_control)

vsc = CombinedVIwithVZ(0.59, #kpv:: Voltage controller proportional gain
                       736.0, #kiv:: Voltage controller integral gain
                       0.0, #kffv:: Binary variable enabling the voltage feed-forward in output of current controllers
                       0.0, #rv:: Virtual resistance in pu
                       0.2, #lv: Virtual inductance in pu
                       1.27, #kpc:: Current controller proportional gain
                       14.3, #kiv:: Current controller integral gain
                       0.0, #kffi:: Binary variable enabling the current feed-forward in output of current controllers
                       50.0, #ωad:: Active damping low pass filter cut-off frequency
                       0.2) #kad:: Active damping gain

Darco_Inverter = DynInverter(1, #number
                             :DARCO, #name
                             nodes_DAIB[2], #bus location
                             1.0, #ω_ref
                             1.02, #V_ref
                             0.5, #P_ref
                             0.0, #Q_ref
                             2.75, #MVABase
                             converter, #Converter
                             outer_control, #OuterControl
                             vsc, #Voltage Source Controller
                             dc_source, #DC Source
                             pll, #Frequency Estimator
                             filt) #Output Filter

######################### Dynamical System ########################

DAIB = DynamicSystem(nodes_DAIB, branch_DAIB, [Darco_Inverter], [inf_gen_DAIB], 100.0, 50.0);

##################################################
############### SOLVE PROBLEM ####################
##################################################

#Initialize variables
dx0 = zeros(LITS.get_total_rows(DAIB))
x0 = [1.00, #V1_R
          1.0648, #V2_R
          0.0, #V1_I
          0.001, #V2_I
          0.0, #δω_vsm
          0.2, #δθ_vsm
          0.025, #qm
          0.0015, #ξ_d
          -0.07, #ξ_q
          0.05, #γ_d
        -0.001, #γ_q
         0.95, #ϕ_d
         -0.10, #ϕ_q
         1.004, #vpll_d
         0.0, #vpll_q
         0.0, #ε_pll
         0.1, #δθ_pll
         0.5, #id_cv
         0.0, #iq_cv
         0.95, #vod
         -0.1, #voq
         0.49, #iod
        -0.1] #ioq
Darco_Inverter.inner_vars[13] = 0.95 #Vd_cnv var
Darco_Inverter.inner_vars[14] = -0.1 #Vq_cnv var
u0 = [0.2] #Increase in P_ref
tspan = (0.0, 4.0);

#Find initial condition
inif! = (out,x) -> LITS.system_model!(out, dx0 ,x, (u0,DAIB), 0.0)
sys_solve = nlsolve(inif!, x0, xtol=:1e-8,ftol=:1e-8,method=:trust_region)
x0_init = sys_solve.zero

#Define Fault using Callbacks
cb = DiffEqBase.DiscreteCallback(LITS.change_t_one, LITS.step_change_first!)

#Define Simulation Problem
sim = DynamicSimulation(DAIB, tspan, u0, cb, x0_init)

#Solve problem in equilibrium
run_simulation!(sim, IDA());

#Obtain data for angles
series = get_state_series(sim, (:DARCO, :δω_vsm));

@test sim.solution.retcode == :Success


text:
############################################
#
#    dolfin.cpp.la
# https://fenicsproject.org/docs/dolfin/2017.2.0/python/programmers-reference/cpp/la/index.html
#
#    dolfin.la
# https://fenicsproject.org/docs/dolfin/latest/python/_autogenerated/dolfin.la.html
#
############################################

# 这两个类，没有从`dolfin`中导出，所以本库也不导出。 仅用作父类而存在。
@pyclass dolfin GenericTensor          # Bases: dolfin.cpp.la.LinearAlgebraObject
@pyclass dolfin GenericMatrix GenericTensor

@pyclass dolfin Scalar GenericTensor
@pyclass dolfin GenericVector GenericTensor
@pyclass dolfin EigenVector GenericVector
@pyclass dolfin PETScVector GenericVector
@pyclass dolfin Vector GenericVector FeVector
@pyclass dolfin EigenMatrix GenericMatrix
@pyclass dolfin PETScMatrix GenericMatrix
@pyclass dolfin Matrix GenericMatrix FeMatrix
@pyclass dolfin BlockMatrix            # Bases: object
@pyclass dolfin BlockVector            # Bases: object
@pyclass dolfin IndexMap               # Bases: object
@pyclass dolfin KrylovSolver           # Bases: dolfin.cpp.la.GenericLinearSolver
@pyclass dolfin LUSolver               # Bases: dolfin.cpp.la.GenericLinearSolver
@pyclass dolfin DefaultFactory         # Bases: dolfin.cpp.la.GenericLinearAlgebraFactory
@pyclass dolfin EigenFactory           # Bases: dolfin.cpp.la.GenericLinearAlgebraFactory
@pyclass dolfin LinearOperator         # Bases: dolfin.cpp.la.LinearAlgebraObject
@pyclass dolfin TensorLayout           # Bases: dolfin.cpp.common.Variable
@pyclass dolfin VectorSpaceBasis       # Bases: object
@pyclass dolfin NewtonSolver           # Bases: dolfin.cpp.common.Variable
@pyclass dolfin NonlinearProblem       # Bases: object
@pyclass dolfin OptimisationProblem NonlinearProblem
@pyclass dolfin SLEPcEigenSolver       # Bases: dolfin.cpp.common.Variable, dolfin.cpp.la.PETScObject
@pyclass dolfin PETScLUSolver          # Bases: dolfin.cpp.la.GenericLUSolver, dolfin.cpp.la.PETScObject
@pyclass dolfin PETScKrylovSolver      # Bases: dolfin.cpp.la.GenericLinearSolver, dolfin.cpp.la.PETScObject

@pyfunc ufl as_tensor
@pyfunc ufl as_vector
@pyfunc ufl as_matrix

@pyfunc dolfin list_lu_solver_methods
@pyfunc dolfin list_linear_solver_methods
@pyfunc dolfin lu_solver_methods
@pyfunc dolfin linear_solver_methods
@pyfunc dolfin list_linear_algebra_backends
@pyfunc dolfin has_linear_algebra_backend
@pyfunc dolfin list_krylov_solver_methods
@pyfunc dolfin list_krylov_solver_preconditioners
@pyfunc dolfin has_krylov_solver_method
@pyfunc dolfin has_krylov_solver_preconditioner

@pyfunc dolfin as_backend_type
@pyfunc dolfin la_index_dtype
@pyfunc dolfin normalize
@pyfunc dolfin residual

list_lu_solver_methods, list_linear_solver_methods, lu_solver_methods, linear_solver_methods,
list_krylov_solver_methods, list_krylov_solver_preconditioners, normalize,residual

*(v::GenericTensor, c::Real) = to_fetype(v.pyobject * c)
*(c::Real, v::GenericTensor) = to_fetype(c * v.pyobject)

+(v1::GenericVector, v2::GenericVector) = to_fetype(v1.pyobject + v2.pyobject)

-(v1::GenericVector, v2::GenericVector) = to_fetype(v1.pyobject - v2.pyobject)
-(v::GenericTensor) = to_fetype( (-1) * v.pyobject)

/(v::GenericTensor, c::Real) = to_fetype(v.pyobject / c)

function Base.lastindex(o::GenericVector)
    o.pyobject.size()
end

Base.getindex(o::GenericVector, idx::Union{Array{Int},UnitRange{Int}}) = to_fetype(o.pyobject.get_local(idx .- 1))
Base.getindex(o::GenericVector, idx::Colon) = o[1:end]

function Base.setindex!(o::GenericVector, values::Union{GenericVector,Array}, idx::Union{Array{Int},UnitRange{Int}})
    if isa(values,GenericVector)
        o.pyobject.set_local(values.pyobject)
    else
        cp = o.pyobject.get_local()
        cp[idx] = values
        o.pyobject.set_local(cp)
    end
end
function Base.setindex!(o::GenericVector, values::Union{GenericVector,Array}, idx::Colon)
    o[1:end] = values
end

export Scalar, GenericVector, EigenVector, PETScVector, FeVector, EigenMatrix, PETScMatrix, FeMatrix,
       BlockMatrix, BlockVector, IndexMap, KrylovSolver, LUSolver, DefaultFactory, EigenFactory,
       LinearOperator, TensorLayout, VectorSpaceBasis, NewtonSolver, NonlinearProblem,
       OptimisationProblem, SLEPcEigenSolver, PETScLUSolver, PETScKrylovSolver, as_tensor, as_vector,
       as_matrix, list_lu_solver_methods, list_linear_solver_methods, lu_solver_methods,
       linear_solver_methods, list_linear_algebra_backends, has_linear_algebra_backend,
       has_krylov_solver_method,has_krylov_solver_preconditioner, list_krylov_solver_methods,
       as_backend_type, la_index_dtype, list_krylov_solver_preconditioners, normalize,residual


text:
# # clip-05s.jl

# Load Julia packages (libraries) needed

using StatisticalRethinking, CmdStan
#gr(size=(600,800));

# CmdStan uses a tmp directory to store the output of cmdstan

ProjDir = rel_path("..", "scripts", "03")
cd(ProjDir)

# Define the Stan language model

binomialstanmodel = "
// Inferring a Rate
data {
  int N;
  int<lower=0> k[N];
  int<lower=1> n[N];
}
parameters {
  real<lower=0,upper=1> theta;
  real<lower=0,upper=1> thetaprior;
}
model {
  // Prior Distribution for Rate Theta
  theta ~ beta(1, 1);
  thetaprior ~ beta(1, 1);

  // Observed Counts
  k ~ binomial(n, theta);
}
";

# Define the Stanmodel and set the output format to :mcmcchains.

stanmodel = Stanmodel(name="binomial", monitors = ["theta"], model=binomialstanmodel,
  output_format=:mcmcchains);

# Use 16 observations

N2 = 4^2
d = Binomial(9, 0.66)
n2 = Int.(9 * ones(Int, N2))
k2 = rand(d, N2);

# Input data for cmdstan

binomialdata = Dict("N" => length(n2), "n" => n2, "k" => k2);

# Sample using cmdstan

rc, chn, cnames = stan(stanmodel, binomialdata, ProjDir, diagnostics=false,
  CmdStanDir=CMDSTAN_HOME);

# Describe the draws

MCMCChains.describe(chn)

# Plot the 4 chains

if rc == 0
  plot(chn)
end

# End of `03/clip-05s.jl`


text:
function Distributions.logpdf(d::Dict{Symbol, Vector{Sampleable}}, x, i)
    sum([logpdf(d[k][i], x[k][i]) for k in keys(d)])
end

function Distributions.logpdf(d::Dict{Symbol, Vector{Sampleable}}, x)
    sum([logpdf(d, x, i) for i=1:length(first(x)[2])])
end

function Distributions.logpdf(d::Dict{Symbol, Tuple{Sampleable, Int64}}, x, i)
    sum([logpdf(d[k][1], x[k][i]) for k in keys(d)])
end

function Distributions.logpdf(d::Dict{Symbol, Tuple{Sampleable, Int64}}, x)
    sum([logpdf(d, x, i) for i=1:length(first(x)[2])])
end

function Base.rand(rng::AbstractRNG, d::Dict{Symbol, Vector{Sampleable}})
    Dict(k => rand.(Ref(rng), d[k]) for k in keys(d))
end

function Base.rand(rng::AbstractRNG, d::Dict{Symbol, Tuple{Sampleable, Int64}})
    Dict(k => rand(rng, d[k][1], d[k][2]) for k in keys(d))
end

function Base.rand(rng::AbstractRNG, d::Dict{Symbol, Vector{Sampleable}}, N::Int)
    [rand(rng, d) for i=1:N]
end

function Base.rand(rng::AbstractRNG, d::Dict{Symbol, Tuple{Sampleable, Int64}}, N::Int)
    [rand(rng, d) for i=1:N]
end

function Distributions.fit(d::Dict{Symbol, Vector{Sampleable}}, samples, weights; add_entropy = (x) -> x)
    N = length(samples)
    new_d = Dict{Symbol, Vector{Sampleable}}()
    for s in keys(d)
        dtype = typeof(d[s][1])
        m = length(d[s])
        new_d[s] = [add_entropy(fit(dtype, [samples[j][s][i] for j=1:N], weights)) for i=1:m]
    end
    new_d
end

function Distributions.fit(d::Dict{Symbol, Tuple{Sampleable, Int64}}, samples, weights; add_entropy = (x)->x)
    N = length(samples)
    new_d = Dict{Symbol, Tuple{Sampleable, Int64}}()
    for s in keys(d)
        dtype = typeof(d[s][1])
        m = d[s][2]
        all_samples = vcat([samples[j][s][:] for j=1:N]...)
        all_weights = vcat([fill(weights[j], length(samples[j][s][:])) for j=1:N]...)
        new_d[s] = (add_entropy(fit(dtype, all_samples, all_weights)), m)
    end
    new_d
end

# Gets a function that adds entropy h to categorical distribution c
function add_categorical_entropy(h::Vector{Float64})
    function add_entropy(c::Categorical)
        p = h
        p[1:length(c.p)] .+= c.p
        Categorical(p ./ sum(p))
    end
end

# This version uses a vector of distributions for sampling
# N is the number of samples taken
# m is the length of the vector
function cross_entropy_method(loss::Function,
                              d_in;
                              max_iter,
                              N=100,
                              elite_thresh = -0.99,
                              min_elite_samples = Int64(floor(0.1*N)),
                              max_elite_samples = typemax(Int64),
                              weight_fn = (d,x) -> 1.,
                              rng::AbstractRNG = Random.GLOBAL_RNG,
                              verbose = false,
                              show_progress = false,
                              add_entropy = (x)->x
                             )
    d = deepcopy(d_in)
    show_progress ? progress = Progress(max_iter) : nothing

    for iteration in 1:max_iter
        # Get samples -> Nxm
        samples = rand(rng, d, N)

        # sort the samples by loss and select elite number
        losses = [loss(d, s) for s in samples]
        order = sortperm(losses)
        losses = losses[order]
        N_elite = losses[end] < elite_thresh ? N : findfirst(losses .> elite_thresh) - 1
        N_elite = min(max(N_elite, min_elite_samples), max_elite_samples)

        verbose && println("iteration ", iteration, " of ", max_iter, " N_elite: ", N_elite)

        #update based on elite samples
        elite_samples = samples[order[1:N_elite]]
        weights = [weight_fn(d, s) for s in elite_samples]
        if all(weights .≈ 0.)
            println("Warning: all weights are zero")
        end
        d = fit(d, elite_samples, weights, add_entropy = add_entropy)
        show_progress && next!(progress)
    end
    d
end


text:
module Minvar
##################Dependencies####################
using Revise

if  pwd() ∉ LOAD_PATH
  push!(LOAD_PATH,pwd())
end

if "$(pwd())\\data" ∉ LOAD_PATH
  push!(LOAD_PATH,"$(pwd())\\data")
end

using DataFrames, Distributions, GZip, ForwardDiff, Dates, LinearAlgebra, Statistics,
   Measures, Formatting, StatsBase, Distributions, DataStructures, Finometrics, Distributed,
  Serialization,  CSV, Random, Distributed, Gadfly, Fontconfig, Cairo# NOTE: uncomment for graphs

#########################Exported functions##########
export testportfolio,
  testcov,
  testtrisectedportfolio,
  matlab2julia,
  mathematica2juliaEq,
  mathematica2juliaW,
  testω,
  testposteriors,
  estimateminvariance,
  initializepar, #need this to set up the parallelization
  cleanup,
  makeconvergencegraphs,
  Estimation

################Priors and initial vlaues############
###Infromative Priors
const D_θG = 0.16^2/20 #current vix converted to daily
const D_δ²G = .76/20*2 #reasonably diffuse prior from volatiltiy of the vix multiplied by 2
const D_θSG = D_θG
const D_δ²SG = D_δ²G*2.0 #more uncertain
const D_θSGP = D_θG
const D_δ²SGP = D_δ²G*4.0 #most uncertain

const D_zG = .76/20/510*2*.01
const D_zP = D_zG * 2
const D_zGP = (.76/20/510)^0.5*(D_zG/2)/2 #std dev * std dev G * std dev P /2
const D_Ψ = [D_zG D_zGP; D_zGP D_zP]
const D_ν = 3.01

const U_θG = 0.0001 #current vix converted to daily
const U_δ²G = 0.0001 #reasonably diffuse prior from volatiltiy of the vix multiplied by 2
const U_θSG = 0.0001
const U_δ²SG = 0.0001 #more uncertain

const U_zG = 0.0001
const U_zP = 0.0001
const U_zGP = 0.0000 #std dev * std dev G * std dev P /2
const U_Ψ = [U_zG U_zGP; U_zGP U_zP]
const U_ν = 3.01

#Initial Parameter Values
const D_σ²G= 0.0001
const D_ζ²G= 0.0001
const D_ζ²P = 0.0001
const D_ζGP = 0.0000
const D_SG = 0.0001
const D_SGP = .0001

const SMALL = 1.0e-20

###Legacy, not used right now
const U_θSGP = 0.0001
const U_δ²SGP = 0.0001 #most uncertain

##############Other constants and methedological parameters############
#const TEX_PATH = "C:\\Users\\Clinton\\Dropbox\\Apps\\Overleaf\\Endowment Project"
#const GRAPH_PATH = TEX_PATH * "\\sub\\fig"
#const WORKING_PATH = pwd() * "\\working"

const OUT_NAME = "test"
const OUT_PATH = "output"
#const OUT_GRAPH_PATH = "C:\\Users\\Clinton Tepper\\Dropbox\\Apps\\Overleaf\\minvar\\sub"
const OUT_GRAPH_PATH = "output"
const OUT_ESTIMATE_NAME = "lastestimate"
const PlotContainer = Union{Plot,Gadfly.Compose.Context,Vector{Plot}}

const N_BURN = 100
const N_SAMPLES = 100

const DATA_NAME = "sampledata"
const DATA_PATH = pwd() * "\\data"

const DATE_BEGIN = Date(2015,1,1)
const DATE_END =  Date(2016,12,31)
const DATE_FORMAT = "yyyymmdd"

const MAX_WORKERS = Base.Sys.CPU_THREADS ÷ 2 # gets num of physical cores

const NUM_TEST_ASSETS = 100 #this shouldn't make a big difference in the results
const MIN_ADJUSTMENT = 0.75 #smaller values->more extreme weights allowed for test portfolio
const MIN_BURN = 1_000 #max used in convergence graphs
const CONVERGENCE_WINDOWS = [1000,10_000,100_000,1_000_000]
#without this throws an error on each run if already defined
if !(@isdefined RANDOM_WEIGHT_DIST)
  const RANDOM_WEIGHT_DIST = TriangularDist(-1.0,1.5,1/NUM_TEST_ASSETS)
end

#############File linkage###############################
include("code\\MinvarPar.jl")
include("code\\Universe.jl")
include("code\\Portfolio.jl")
include("code\\Parameters.jl")
include("code\\Posteriors.jl")
include("code\\TrisectedPortfolio.jl")
include("code\\X2Julia.jl")
include("code\\Estimation.jl")
include("code\\Output.jl")
include("code\\Testsandoneoffs.jl")

end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
immutable MMapIO
	# buf::Vector{UInt8}
	# pos::Int  # NOTE: zero-indexed!
	io::IOBuffer
end

# could pass in a filename or IOStream
MMapIO(arg::Union{AbstractString,IOStream}) = MMapIO(IOBuffer(Mmap.mmap(arg)))

Base.show(io::IO, mmio::MMapIO) = "MMapIO{$(length(mmio.buf)) @ $(mmio.pos)}"

Base.eof(mmio::MMapIO) = eof(mmio.io) #mmio.pos >= length(mmio.buf)
Base.read{T}(mmio::MMapIO, ::Type{T}) = read(mmio.io, T)
Base.read(mmio::MMapIO, x) = read(mmio.io, x)

# function Base.read!{T}(mmio::MMapIO, ::Type{T})
# 	struct_size = getPackedStructSize(T)
# 	if mmio.pos + struct_size > length(mmio.buf)
# 		@show mmio.pos struct_size length(mmio.buf)
# 		throw(EOFError())
# 	end
# 	obj = reinterpret(T, mmio.buf, mmio.pos)
# 	mmio.pos += struct_size
# 	obj
# end


text:
## Reference implementations of algorithms from
## D.Cox, J.Little, D.O'Shea: Ideals, Varieties, and Algorithms (3rd ed)

lexless{K<:Field,N}(t1::Term{K,N}, t2::Term{K,N}) = lexless(t1.index, t2.index)

function grlexless{N}(i1::Index{N}, i2::Index{N})
    d1 = deg(i1)
    d2 = deg(i2)
    (d1<d2) || ((d1==d2) && lexless(i1, i2))
end
grlexless{K<:Field,N}(t1::Term{K,N}, t2::Term{K,N}) = grlexless(t1.index, t2.index)

function grevlexless{N}(i1::Index{N}, i2::Index{N})
    d1 = deg(i1)
    d2 = deg(i2)
    (d1<d2) || ((d1==d2) && lexless(reverse(i2), reverse(i1)))
end
grevlexless{K<:Field,N}(t1::Term{K,N}, t2::Term{K,N}) = grevlex(t1.index, t2.index)

_order = lexless

function set_order_lex()
    global _order = lexless
end

function set_order_grlex()
    global _order = grlexless
end

function set_order_grevlex()
    global _order = grevlexless
end

_grevlexelim_l = 0
function grevlexelimless{N}(i1::Index{N}, i2::Index{N})
    d1 = sum(i1[1:_grevlexelim_l])
    d2 = sum(i2[1:_grevlexelim_l])
    (d1<d2) || ((d1==d2) && grevlexless(i1, i2))
end
grevlexelimless{K<:Field,N}(t1::Term{K,N}, t2::Term{K,N}) = grevlexelimless(t1.index, t2.index)

function set_order_grevlexelimless(l)
    global _order = grevlexelimless
    global _grevlexelim_l = l
end

isless{K<:Field,N}(t1::Term{K,N}, t2::Term{K,N}) = _order(t1.index, t2.index)
<{K<:Field,N}(t1::Term{K,N}, t2::Term{K,N}) = _order(t1.index, t2.index)
>{K<:Field,N}(t1::Term{K,N}, t2::Term{K,N}) = _order(t2.index, t1.index)

LT{K<:Field,N}(f::Polynomial{K,N}) = maximum(terms(f))
LT{K<:Field,N}(tt::PolyAsArray{K,N}) = maximum(tt)
LC{K<:Field,N}(f::Polynomial{K,N}) = LT(f).coefficient
LC{K<:Field,N}(tt::PolyAsArray{K,N}) = LT(tt).coefficient
LM{K<:Field,N}(f::Polynomial{K,N}) = Term(LT(f).index, one(K))
LM{K<:Field,N}(tt::PolyAsArray{K,N}) = Term(LT(tt).index, one(K))
multidegree{K<:Field,N}(f::Polynomial{K,N}) = index(LT(f))
multidegree{K<:Field,N}(tt::PolyAsArray{K,N}) = index(LT(tt))

function divrem{K<:Field,N}(f::Polynomial{K,N}, ff::Polynomial{K,N} ...)
    # Algorithm from Theorem 2.3 in
    # D.Cox, J.Little, D.O'Shea: Ideals, Varieties, and Algorithms (3rd ed)
    s = length(ff)
    zero = Polynomial{K,N}()
    r = copy(zero)
    a = [copy(zero) for i=1:s]
    p = copy(f)
    while p!=zero
        i = 1
        divisionoccurred = false
        while i<=s && !divisionoccurred
            ltf = LT(ff[i])
            ltp = LT(p)
            if divides(ltf, ltp)
                a[i] = a[i] + ltp/ltf
                p = p - (ltp/ltf)*ff[i]
                divisionoccurred = true
            else
                i = i + 1
            end
        end
        if !divisionoccurred
            ltp = LT(p)
            r = r + ltp
            p = p - ltp
        end
    end
    a, r
end

function rem{K<:Field,N}(f::Polynomial{K,N}, ff::Polynomial{K,N} ...)
    #The same as a divrem without computing a
    s = length(ff)
    zero = Polynomial{K,N}()
    r = copy(zero)
    p = copy(f)
    while p!=zero
        i = 1
        divisionoccurred = false
        while i<=s && !divisionoccurred
            ltf = LT(ff[i])
            ltp = LT(p)
            if divides(ltf, ltp)
                p = p - (ltp/ltf)*ff[i]
                divisionoccurred = true
            else
                i = i + 1
            end
        end
        if !divisionoccurred
            ltp = LT(p)
            r = r + ltp
            p = p - ltp
        end
    end
    r
end

LCM{K<:Field,N}(t1::Term{K,N}, t2::Term{K,N}) = Term{K,N}(Index{N}(([max(t1.index[k], t2.index[k]) for k=1:N]...)), one(K))

function S_polynomial{K<:Field,N}(f::Polynomial{K,N}, g::Polynomial{K,N})
    ltf = LT(f)
    ltg = LT(g)
    lcm = LCM(ltf, ltg)
    (lcm/ltf)*f - (lcm/ltg)*g
end

function groebner{K<:Field, N}(ff::Polynomial{K,N} ...)
    # Algorithm from Theorem 2.11 in
    # D.Cox, J.Little, D.O'Shea: Ideals, Varieties, and Algorithms (3rd ed)
    zero = Polynomial{K,N}()
    s = length(ff)
    B = Set(vcat([[(i,j)  for i=1:j-1] for j=1:s]...))
    G = [ff[k]/LC(ff[k]) for k=1:s]
    sort(G, lt=(x,y)->LT(x)<LT(y)) # Sort G[i] such that their leading terms are increasing
    t = s
    lt = [LT(G[k]) for k=1:s]
    while !isempty(B)
        first = true
        (i,j) = (0,0)
        lcm = 0
        for (i1,j1) in B # Select (i,j)∈B of such that LCM(lt[i], lt[j]) is as small as possible.
            if first
                (i,j) = (i1,j1)
                lcm = LCM(lt[i1], lt[j1])
                first = false
            else
                lcm1 = LCM(lt[i1], lt[j1])
                if lcm1<lcm
                    (i,j) = (i1,j1)
                    lcm = lcm1
                end
            end
        end
        #println((i,j), " selected")
        if lcm.index != (lt[i]*lt[j]j).index
            criterion = false
            for k = 1:t
                if k==i || k==j
                    continue
                end
                if !((min(i,k),max(i,k)) in B || (min(j,k),max(j,k)) in B) && divides(lt[k],lcm)
                    criterion = true
                    break
                end
            end
            if !criterion
                S = rem(S_polynomial(G[i],G[j]), G...)
                if S!=zero
                    t = t+1
                    push!(G, S/LC(S))
                    push!(lt, LT(G[t]))
                    union!(B, [(k,t) for k=1:t-1])
                    #println("\t",S)
                end
            end
        end
        delete!(B, (i,j))
    end

    # generate minimal Groebner basis
    to_be_deleted = Int[]
    for i=1:length(G)
        if i in to_be_deleted
            continue
        end
        for j=1:length(G)
            if j==i || j in to_be_deleted || i in to_be_deleted
                continue
            end
            if divides(LT(G[j]), LT(G[i]))
                push!(to_be_deleted, i)
            end
         end
    end
    deleteat!(G, to_be_deleted)
    for i=1:length(G)
        G[i] = G[i]/LC(G[i])
    end

    # generate reduced Groebner basis
    t = length(G)
    for i=1:t
        G[i] = rem(G[i], vcat(G[1:i-1],G[i+1:t])...)
    end

    G
end


text:
struct RLBaseEnv{E} <: RLBase.AbstractEnv
    env::E
end

Base.show(io::IO, mime::MIME"text/plain", env::RLBaseEnv{E}) where {E <: AbstractGame} = show(io, mime, env.env)

play!(env::RLBaseEnv{E}) where {E <: AbstractGame} = play!(env.env)


text:
# algorithmOutput.jl
# Use manhattanVisualizations to see step-by-step progression of algorithm
# Authored by Arthur J Delarue on 6/16/15

using SFML, LightGraphs, JLD
using Base.Collections

include("../src/LP_tools/definitions.jl")
include("../src/LP_tools/manhattan.jl")

include("manhattanVisualization.jl")

DIRECTORY = "new_r140_minr1_i5_wd_1214_avgtr_clust50_rides50000"
MAX_ROUNDS = 5

function showAlgorithmEvolution(dir=DIRECTORY, firstRound::Int=1, secondRound::Int=2)
	man = Manhattan()
	oldTimes = load("/Users/arthurdelarue/Desktop/VRP-remote/Outputs/$(dir)/manhattan-times-$firstRound.jld", "times")
	newTimes = load("/Users/arthurdelarue/Desktop/VRP-remote/Outputs/$(dir)/manhattan-times-$secondRound.jld", "times")
	man.roadTime = newTimes - oldTimes
	drawManhattan(man, colorScheme = 1)
end

function showAlgorithmResults(dir=DIRECTORY, rounds=1:MAX_ROUNDS)
	man = Manhattan()
	println("**** Original Manhattan ****")
	for i = 1:nv(man.network), j=1:nv(man.network)
		if man.roadTime[i,j] > 0
			man.roadTime[i,j] = 0.038 * man.distances[i,j]
		end
	end
	drawManhattan(man, colorScheme = 0)

	for i in rounds
		println("**** Iteration $i ****")
		man.roadTime = load("/Users/arthurdelarue/Desktop/VRP-remote/Outputs/$(dir)/manhattan-times-$i.jld", "times")
		avg_speed = 0
		num_streets = 0
		for i = 1:nv(man.network), j = out_neighbors(man.network, i)
			avg_speed = (avg_speed * num_streets + man.distances[i,j]/man.roadTime[i,j])/(num_streets + 1)
			num_streets += 1
		end
		avg_speed = 3.6 * avg_speed
		println("Average speed: $avg_speed")
		drawManhattan(man, colorScheme = 0)
	end
end

println("")

max_stars_repo_name: sebmart/NYC-timing

max_issues_repo_name: sebmart/NYC-timing

max_forks_repo_name: sebmart/NYC-timing


text:
"""
    readbed(bedfile::String; pad=0, flanking=false, delim='\t')

Read and optionally pad coordinates of a 3+ column bedfile.
Return an array of positions in the chr:start-stop format.
Bedfile may be commented with #.

Examples:
Simply reformat sites to UCSC input: readbed(bedfile)
Expand the region by 1kb: readbed(bedfile, pad=1000);
Omit target and get the flanking regions of size "pad";
(leftsites, rightsites) = readbed(bedfile; pad=100, flanking=true).
"""
function readbed(bedfile::String; pad=0, flanking=false, delim='\t')

    w = pad
    m = readdlm(bedfile, Char(delim))

    chrs = string.(m[:,1])
    chrs = replace.(chrs, r"\.0" => "")
    o1 = BigInt.(m[:,2])
    o2 = BigInt.(m[:,3])
    p1 = BigInt.(m[:,2] .- w)
    p1 = [if i < 0; 1 else i end for i in p1]
    p2 = BigInt.(m[:,3] .+ w)
    p2 = [if i < 0; 1 else i end for i in p2]

    if pad == 0 && flanking == true
        error("Pad sequence must be set and >0 if excluding region.")
    end

    if sum(collect(p1 .> p2)) > 0
        error("Second coordinate must be less than first.")
    end

    if flanking == true

        nb1 = hcat(string.(chrs), string.(p1 .- 1), string.(o1 .- 1))
        nb2 = hcat(string.(chrs), string.(o2 .+ 1), string.(p2 .+ 1))

        newpos1 = [join(nb1[i,:], ":","-") for i in 1:size(nb1,1)]
        newpos2 = [join(nb2[i,:], ":","-") for i in 1:size(nb2,1)]
        return newpos1, newpos2

    else

        nb = hcat(string.(chrs), string.(p1), string.(p2))
        newpos = [join(nb[i,:], ":","-") for i in 1:size(nb,1)]
        return newpos

    end

end


text:
### Joint recovery
## author: Ziyi Yin (ziyi.yin@gatech.edu)

using DrWatson
@quickactivate "CompassTimeLapseCCS"
using TimeProbeSeismic
using JLD2, FFTW, DSP, PyPlot, Images

import Base.*

Random.seed!(1234);

# Needed for A*z
*(A::joLinearFunction, v::Array{judiVector{Float32,Array{Float32,2}},1}) = A.fop(v)
*(A::joLinearFunction, v::Array{judiVector{Float32},1}) = A.fop(v)
# adjoint(A)*b
*(A::joLinearFunction, v::Array{Array{Float32,1},1}) = A.fop(v)

JLD2.@load "../model/CompassTimeLapse.jld2"
JLD2.@load "../data/d_lin_cut.jld2" d_lin_cut q

model0_stack = [Model(n,d,o,m0_stack[i]; nb=100) for i = 1:num_vintage]

srcGeometry = q.geometry
recGeometry_stack = [d_lin_cut[i].geometry for i = 1:num_vintage]

ntComp = get_computational_nt(srcGeometry, recGeometry_stack[1], model0_stack[1])
info = Info(prod(n), d_lin_cut[1].nsrc, ntComp)

opt = Options(isic=true, limit_m=true)

F0_stack = [judiProjection(info, recGeometry_stack[i])*judiModeling(info, model0_stack[i]; options=opt)*adjoint(judiProjection(info, srcGeometry)) for i = 1:num_vintage]
J_stack = [judiJacobian(F0_stack[i],q) for i = 1:num_vintage]
data_ = d_lin_cut
###

L = num_vintage

function Setup4D_forward(As, z)
	# Sum z_i according to identity matrix
	vec = As[end] * z
	# Apply diag(ops) to summed z_i
	return [As[i]*vec[i] for i=1:length(As)-1]
end

function Setup4D_adjoint(As, b)
	# Apply diag(adj(ops))  * dlin_i then apply IMatrix
    z = adjoint(As[end]) * [adjoint(As[i]) * b[i] for i=1:length(b)]
    return z
end

# Setup JRM
ps = 32

niter = 23
batchsize = 8 * ones(Int,L)

soft_thresholding(x::Array{Float64}, lambda) = sign.(x) .* max.(abs.(x) .- convert(Float64, lambda), 0.0)
soft_thresholding(x::Array{Float32}, lambda) = sign.(x) .* max.(abs.(x) .- convert(Float32, lambda), 0f0)

n = model0_stack[1].n
C0 = joCurvelet2D(n[1], 2*n[2]; zero_finest = false, DDT = Float32, RDT = Float64)

function C_fwd(im, C, n)
	im = hcat(reshape(im, n), reshape(im, n)[:, end:-1:1])
	coeffs = C*vec(im)/sqrt(2f0)
	return coeffs
end

function C_adj(coeffs, C, n)
	im = reshape(C'*coeffs, n[1], 2*n[2])
	return vec(im[:, 1:n[2]] .+ im[:, end:-1:n[2]+1])/sqrt(2f0)
end

C = joLinearFunctionFwd_T(size(C0, 1), n[1]*n[2],
                          x -> C_fwd(x, C0, n),
                          b -> C_adj(b, C0, n),
                          Float32,Float64, name="Cmirrorext")

x = [zeros(Float32, info.n) for i=1:L+1];
z = [zeros(Float32, info.n) for i=1:L+1];

src_list = [collect(1:d_lin_cut[1].nsrc) for i = 1:L]

γ   = 1f0 # hyperparameter to tune
Iz  = Array{Any}(undef, L, L+1)

# Preconditioner

function dip(x,n;k=20)
    image = reshape(x,n)
    image_ext = zeros(Float32,n[1],2*n[2])
    image_ext[:,1:n[2]] = image
    image_ext[:,n[2]+1:end] = reverse(image,dims=2)
    image_f = fftshift(fft(image_ext))
    mask = ones(Float32,n[1],2*n[2])
    for i = 1:n[1]
        for j = 1:2*n[2]
            if (i-(n[1]+1)/2-k*j+k*(2*n[2]+1)/2)*(i-(n[1]+1)/2+k*j-k*(2*n[2]+1)/2)>0
                mask[i,j] = 0f0
            end
        end
    end
    mask = convert(Array{Float32},imfilter(mask,Kernel.gaussian(20)))
    image_f1 = mask .* image_f
    image_out = (vec(real.(ifft(ifftshift(image_f1)))[:,1:n[2]])+vec(real.(ifft(ifftshift(image_f1)))[:,end:-1:n[2]+1]))/2f0
    return image_out
end

D = joLinearFunctionFwd_T(prod(n), prod(n),
                                 v -> dip(v,n;k=10),
                                 w -> dip(w,n;k=10),
                                 Float32,Float32,name="dip filter")
Tm = judiTopmute(model0_stack[1].n, idx_wb, 1)  # Mute water column
S = judiDepthScaling(model0_stack[1])  # depth scaling
Mr = S*Tm*D

for i=1:L

	for j=1:L+1
		Iz[i,j] = 0
	end

	Iz[i,1]  = 1f0/γ*joEye(prod(info.n); RDT=Float32, DDT=Float32)

end

collect(Iz[i,i+1] = joEye(prod(info.n); RDT=Float32, DDT=Float32) for i=1:L)
J  = [[undef for i=1:L]; [Iz]]

b = Array{judiVector{Float32,Array{Float32,2}},1}(undef, L)
lambda = zeros(Float64,L+1)

for  j=1:niter

	# Main loop
	@printf("Iteration: %d \n", j)
    flush(Base.stdout)

	for i=1:L
		length(src_list[i]) < batchsize[i] && (global src_list[i] = collect(1:data_[i].nsrc))
		src_list[i] = src_list[i][randperm(MersenneTwister(i+2000*j),length(src_list[i]))]
		if j == 1
			global inds = [pop!(src_list[i]) for b=1:2*batchsize[i]]
		else
			global inds = [pop!(src_list[i]) for b=1:batchsize[i]]
		end
		global b[i]  = data_[i][inds]
		global J[i]  = judiJacobian(F0_stack[i][inds], q[inds],ps,data_[i][inds])*Mr
		println("Vintage $(i) Imaging source $(inds)")

  	end

	A = joLinearFunctionFwd_T(L, L+1,
		    z -> Setup4D_forward(J,z),
		    b -> Setup4D_adjoint(J,b),
		    Float32,Float32, name="JRM")

    r = A*x - b
    g = adjoint(A)*r

	phi = .5*norm(r)^2

	# Step size and update variable

	t = Float32.(2*phi/norm(g)^2)

	@printf("At iteration %d function value is %2.2e \n", j, phi)
	flush(Base.stdout)

	global z -= t*g
	C_z = [C*z[i] for i = 1:L+1]
	if j == 1
		for k = 1:L+1
			global lambda[k] = quantile(abs.(vec(C_z[k])), .9) # estimate thresholding parameter at 1st iteration
		end
	end
	# Update variables and save snapshot
	global x = [adjoint(C)*soft_thresholding(C_z[i], lambda[i]) for i = 1:L+1]

    JLD2.@save "../results/JointRecovIter$(j).jld2" x z g lambda phi t
end


text:
using LAP_julia, TimerOutputs

for _ in 1:3
    img, imgw, flow = gen_init(:lena, :tiled, flow_args=[15, 30])

    println("ONCE")
    timer=TimerOutput("pf lap");
    method_kwargs =Dict(:timer => timer, :display => false, :max_repeats => 1)
    timer = @time flow_est, source_reg, timer, results = test_registration_alg(pflap, img, imgw, flow, [], method_kwargs, timer=timer, display=true);
end


text:
# ---
# title: 554. Brick Wall
# id: problem554
# author: Indigo
# date: 2021-05-13
# difficulty: Medium
# categories: Hash Table
# link: <https://leetcode.com/problems/brick-wall/description/>
# hidden: true
# ---
#
# There is a brick wall in front of you. The wall is rectangular and has several
# rows of bricks. The bricks have the same height but different width. You want
# to draw a vertical line from the **top** to the **bottom** and cross the
# **least** bricks.
#
# The brick wall is represented by a list of rows. Each row is a list of
# integers representing the width of each brick in this row from left to right.
#
# If your line go through the edge of a brick, then the brick is not considered
# as crossed. You need to find out how to draw the line to cross the least
# bricks and return the number of crossed bricks.
#
# **You cannot draw a line just along one of the two vertical edges of the wall,
# in which case the line will obviously cross no bricks.**
#
#
#
# **Example:**
#
#
#
#     Input: [[1,2,2,1],
#             [3,1,2],
#             [1,3,2],
#             [2,4],
#             [3,1,2],
#             [1,3,1,1]]
#
#     Output: 2
#
#     Explanation:
#     ![](https://assets.leetcode.com/uploads/2018/10/12/brick_wall.png)
#
#
#
#
# **Note:**
#
#   1. The width sum of bricks in different rows are the same and won't exceed INT_MAX.
#   2. The number of bricks in each row is in range [1,10,000]. The height of wall is in range [1,10,000]. Total number of bricks of the wall won't exceed 20,000.
#
#
## @lc code=start
using LeetCode

function least_bricks(wall::Vector{Vector{Int}})
    tb = Dict{Int,Int}()
    for row in wall
        acc = 0
        for i in @view(row[1:(end - 1)])
            acc += i
            haskey(tb, acc) ? tb[acc] += 1 : tb[acc] = 1
        end
    end
    isempty(tb) && return length(wall)
    return length(wall) - maximum(values(tb))
end
## @lc code=end


text:
@Maple.fieldproxy struct LangData <: AbstractDict{Symbol, Any}
    data::Dict{Symbol, Any}

    LangData(d::Dict{String, Any}) = new(Dict{Symbol, Any}(Symbol(k) => isa(v, Dict) ? LangData(v) : v for (k, v) in d))
    LangData(d::Dict{Symbol, Any}) = new(d)
    LangData() = new(Dict{Symbol, Any}())
end

Base.iterate(ld::LangData) = iterate(getfield(ld, :data))
Base.iterate(ld::LangData, i) = iterate(getfield(ld, :data), i)
Base.merge(ld::LangData, others::LangData...) = LangData(merge(getfield(ld, :data), getfield.(others, Ref(:data))...))

function Base.get(ld::LangData, parts::Array{Symbol, 1}, default::Any=LangData())
    target = ld

    for part in parts
        if !isa(target, LangData) || !haskey(getfield(target, :data), part)
            return default
        end

        target = getfield(target, :data)[part]
    end

    return target
end

Base.get(ld::LangData, parts::Array{String, 1}, default::Any=LangData()) = get(ld, Symbol.(parts), default)
Base.get(ld::LangData, key::Symbol) = get(ld, key, LangData())

function setDictPathValue!(d::Dict{String, Any}, parts::Array{String, 1}=String[], value::Any=nothing; allowOverwrite::Bool=false)
    target = d

    for part in parts[1:end - 1]
        if !haskey(target, part)
            target[part] = Dict{String, Any}()
        end

        target = target[part]
    end

    if !allowOverwrite || (allowOverwrite && haskey(target, parts[end]))
        target[parts[end]] = value
    end
end

function parseLangfile(s::String; init::Dict{String, Any}=Dict{String, Any}(), commentPrefix::String="#", assignment::String="=", seperator::String=".", allowOverwrite::Bool=false, stripValues::Bool=true)
    res = init

    replace(s, "\r\n" => "\n")

    for line in split(s, "\n")
        if !startswith(line, commentPrefix)
            parts = split(line, assignment)
            key, value = parts[1], join(parts[2:end], assignment)

            if stripValues
                value = convert(String, strip(value))
            end

            if !isempty(key) && !isempty(value)
                fieldParts = String.(split(key, seperator))
                setDictPathValue!(res, fieldParts, value, allowOverwrite=allowOverwrite)
            end
        end
    end

    return res
end

function expandTooltipText(s::String)
    # Replace \n with actuall newline
    # Strip leading/trailing whitespace from all lines

    s = replace(s, "\\n" => "\n")
    s = join(strip.(split(s, "\n")), "\n")

    return s
end

function loadLangfile()
    global langdata = LangData(loadLangdata())
end


text:
module TrajectoryOptimization

using SparseArrays
using StaticArrays
using LinearAlgebra
using ForwardDiff
using Ipopt
using Requires
using Symbolics: jacobian_sparsity

# Utilities
include("Utilities/typeFlags.jl")
include("Utilities/sparseMatrixUtils.jl")
export Dynamics, Cost, Algebraic, AllFunctions # Function type flags

# User function utilities
include("UserFunctionUtilities/PathFunction.jl")
include("UserFunctionUtilities/ADPathFunction.jl")
include("UserFunctionUtilities/AnalyticPathFunction.jl")
export PathFunction

# Conditionally load Snopt
function __init__()
    @require Snopt="0e9dc826-d618-11e8-1f57-c34e87fde2c0" include("NLPUtilities/snopt.jl")
end

end


text:
# Options structure
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: May 2017
#

export Options, subsample

# Object for velocity/slowness models
mutable struct Options
    space_order::Integer
    free_surface::Bool
    limit_m::Bool
    buffer_size::Real
    save_data_to_disk::Bool
    save_wavefield_to_disk::Bool
    file_path::String
    file_name::String
    sum_padding::Bool
    optimal_checkpointing::Bool
    num_checkpoints::Union{Integer, Nothing}
    checkpoints_maxmem::Union{Real, Nothing}
    frequencies::Array
    dft_subsampling_factor::Integer
    isic::Bool
end

"""
    Options
        space_order::Integer
        free_surface::Bool
        limit_m::Bool
        buffer_size::Real
        save_rate::Real
        save_data_to_disk::Bool
        save_wavefield_to_disk::Bool
        file_path::String
        file_name::String
        sum_padding::Bool
        optimal_checkpointing::Bool
        num_checkpoints::Integer
        checkpoints_maxmem::Real
        frequencies::Array
        dft_subsampling_factor::Integer
        isic::Bool

Options structure for seismic modeling.

`space_order`: finite difference space order for wave equation (default is 8, needs to be multiple of 4)

`free_surface`: set to `true` to enable a free surface boundary condition.

`limit_m`: for 3D modeling, limit modeling domain to area with receivers (saves memory)

`buffer_size`: if `limit_m=true`, define buffer area on each side of modeling domain (in meters)

`save_data_to_disk`: if `true`, saves shot records as separate SEG-Y files

`save_wavefield_to_disk`: If wavefield is return value, save wavefield to disk as pickle file

`file_path`: path to directory where data is saved

`file_name`: shot records will be saved as specified file name plus its source coordinates

`sum_padding`: when removing the padding area of the gradient, sum into boundary rows/columns for true adjoints

`optimal_checkpointing`: instead of saving the forward wavefield, recompute it using optimal checkpointing

`num_checkpoints`: number of checkpoints. If not supplied, is set to log(num_timesteps).

`checkpoints_maxmem`: maximum amount of memory that can be allocated for checkpoints (MB)

`frequencies`: calculate the FWI/LS-RTM gradient in the frequency domain for a given set of frequencies

`dft_subsampling_factor`: compute on-the-fly DFTs on a time axis that is reduced by a given factor (default is 1)

`isic`: use linearized inverse scattering imaging condition

Constructor
==========

All arguments are optional keyword arguments with the following default values:

    Options(; space_order, free_surface, limit_m=false, buffer_size=1e3, save_data_to_disk=false, save_wavefield_to_disk=false, file_path=pwd(), file_name="shot",
        sum_padding=false, optimal_checkpointing=false, num_checkpoints=log(nt), checkpoints_maxmem=[], frequencies=[], dft_subsampling_factor=[], isic=false)

"""
Options(; space_order=8, free_surface=false, limit_m=false, buffer_size=1e3, save_data_to_disk=false, save_wavefield_to_disk=false, file_path="", file_name="shot", sum_padding=false,
    optimal_checkpointing=false, num_checkpoints=nothing, checkpoints_maxmem=nothing, frequencies=[], dft_subsampling_factor=1, isic=false) =
    Options(space_order, free_surface, limit_m, buffer_size, save_data_to_disk, save_wavefield_to_disk, file_path, file_name,
    sum_padding, optimal_checkpointing, num_checkpoints, checkpoints_maxmem, frequencies, dft_subsampling_factor, isic)

function subsample(options::Options, srcnum)
    if isempty(options.frequencies)
        return options
    else
        opt_out = deepcopy(options)
        opt_out.frequencies = Array{Any}(undef, 1)
        opt_out.frequencies[1] = options.frequencies[srcnum]
        return opt_out
    end
end


text:
using ReinforcementLearning
using OpenSpiel

function RL.Experiment(::Val{:JuliaRL}, ::Val{:Minimax}, ::Val{:OpenSpiel}, game;)
    env = OpenSpielEnv(string(game))
    agents = MultiAgentManager(
        NamedPolicy(0 => MinimaxPolicy()),
        NamedPolicy(1 => MinimaxPolicy()),
    )
    hooks = MultiAgentHook(0 => TotalRewardPerEpisode(), 1 => TotalRewardPerEpisode())
    description = "# Play `$game` in OpenSpiel with Minimax"
    Experiment(agents, env, StopAfterEpisode(1, is_show_progress=!haskey(ENV, "CI")), hooks, description)
end

using Plots
ex = E`JuliaRL_Minimax_OpenSpiel(tic_tac_toe)`
run(ex)

ex.hook

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


text:
# make your scripts automatically re-activate your project
cd(@__DIR__)
using Pkg; Pkg.activate("../../..")

# load packages
using ml_for_pt
using DelimitedFiles
using LaTeXStrings
using Plots
using Random
using Flux
ENV["GKSwstype"]="nul"

# set path to save folder
save_folder = "./results/"

# set critical value of tuning parameter
p_crit = 1.0f0

# define parameter ranges
p_min = 0.1f0
p_max = 3.0f0
dp = 0.05f0
p_range = collect(p_min:dp:p_max)
p_range_LBC = collect(p_min-dp/2:dp:p_max+dp/2)

# construct probability distribution
distr,samples = MLP.jump_distr(p_crit=p_crit)

# optionally run methods using the analytical expressions to compare with neural network results

# supervised learning using analytical expression
# returns optimal predictions, indicator, and loss
pred_opt_SL, indicator_opt_SL, loss_opt_SL = MLP.get_indicators_SL_analytical(samples, distr, p_range, dp, p_min, p_max)

# prediction-based method using analytical expression
# returns optimal predictions, indicator, and loss
pred_opt_PBM, indicator_opt_PBM, loss_opt_PBM = MLP.get_indicators_PBM_analytical(samples, distr, p_range, dp)

# learning by confusion using analytical expression
# returns optimal indicator and loss
indicator_opt_LBC, loss_opt_LBC = MLP.get_indicators_LBC_analytical(samples, distr, p_range, p_range_LBC)

# construct dataset for training neural networks
dataset = MLP.get_dataset_proto(p_range, distr, samples)
dataset_train_SL = hcat(dataset[:, 1:length(samples)], dataset[:, end-length(samples)+1:end])

# standardize inputs
inputs_one_hot = reshape(samples, 1, length(samples))
mean_train, std_train = MLP.get_dataset_stats(dataset,inputs_one_hot, length(p_range))
inputs_one_hot_stand = (inputs_one_hot.-mean_train)./std_train
mean_train_SL, std_train_SL = MLP.get_dataset_stats(dataset_train_SL, inputs_one_hot, length(p_range))
inputs_one_hot_stand_SL = (inputs_one_hot.-mean_train_SL)./std_train_SL

# supervised learning using neural networks

# initialize neural network
n_nodes_SL = [size(inputs_one_hot)[1], 64, 2]
NN_SL = Chain(
  Dense(n_nodes_SL[1], n_nodes_SL[2], relu),
  Dense(n_nodes_SL[2], n_nodes_SL[end], identity))
pnn_SL, re_SL = Flux.destructure(NN_SL)

# set hyperparameters
lr_SL = 0.001f0
epochs_SL = 1000
saveat_SL = 100
opt_SL = ADAM(lr_SL)
verbose = true

# train neural network
# returns predictions, indicators, and loss saved at epochs specified by saveat_SL variable
pred_logger_SL, losses_SL, NN_logger_SL = MLP.get_indicators_SL_numerical(pnn_SL,  re_SL, dataset_train_SL, dataset, epochs_SL, p_range, dp, p_min, p_max, opt_SL, inputs_one_hot_stand, verbose=verbose, saveat=saveat_SL)

# extract results at last save point
pred_NN_SL = pred_logger_SL[end][1]
indicator_NN_SL = pred_logger_SL[end][2]
loss_NN_SL = pred_logger_SL[end][3]

# save results of supervised learning in save_folder
open(save_folder*"hyperparams_SL.txt", "w") do io
  writedlm(io, [saveat_SL, lr_SL, epochs_SL])
end

open(save_folder*"NN_logger_SL.txt", "w") do io
  writedlm(io, NN_logger_SL)
end

for i in 1:length(pred_logger_SL)
  open(save_folder*"pred_logger_SL_"*string(i)*".txt", "w") do io
    writedlm(io, pred_logger_SL[i][1])
  end

  open(save_folder*"indicator_logger_SL_"*string(i)*".txt", "w") do io
    writedlm(io, pred_logger_SL[i][2])
  end

  open(save_folder*"loss_logger_SL_"*string(i)*".txt", "w") do io
    writedlm(io, pred_logger_SL[i][3])
  end

end

open(save_folder*"pred_NN_SL.txt", "w") do io
  writedlm(io, pred_NN_SL)
end

open(save_folder*"indicator_NN_SL.txt", "w") do io
  writedlm(io, indicator_NN_SL)
end

open(save_folder*"loss_NN_SL.txt", "w") do io
  writedlm(io, loss_NN_SL)
end

open(save_folder*"NN_structure_SL.txt", "w") do io
  writedlm(io, n_nodes_SL)
end

open(save_folder*"losses_SL.txt", "w") do io
  writedlm(io, losses_SL)
end

# plot and save results of supervised learning in save_folder
plt = plot(p_range, pred_NN_SL, dpi=300, c="black", label="NN", ylims=(0.0, 1.0))
plot!(p_range, pred_opt_SL, dpi=300, c="blue", label="analytical", ylims=(0.0, 1.0))
vline!([p_crit], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$\hat{y}_{\mathrm{SL}}$")
savefig(save_folder*"pred_NN_SL.png")

plt = plot(p_range[2:end-1], indicator_NN_SL, dpi=300, c="black", label="NN")
plot!(p_range[2:end-1], indicator_opt_SL, dpi=300, c="blue", label="analytical")
vline!([p_crit], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$I_{\mathrm{SL}}$")
savefig(save_folder*"indicator_NN_SL.png")

plt = plot(collect(1:length(losses_SL)), losses_SL, dpi=300, c="black", label="NN")
hline!([loss_opt_SL], label="analytical", c="blue")
xlabel!("Epochs")
ylabel!(L"$\mathcal{L}_{\mathrm{SL}}$")
savefig(save_folder*"loss_NN_SL.png")

# prediction-based method using neural networks

# initialize neural network
n_nodes_PBM = [size(inputs_one_hot)[1], 64, 1]
NN_PBM = Chain(
  Dense(n_nodes_PBM[1], n_nodes_PBM[2], relu),
  Dense(n_nodes_PBM[2], n_nodes_PBM[end], identity))
pnn_PBM, re_PBM = Flux.destructure(NN_PBM)

# set hyperparameters
lr_PBM = 0.001f0
epochs_PBM = 1000
saveat_PBM = 100
opt_PBM = ADAM(lr_PBM)
verbose = true

# train neural network
# returns predictions, indicators, and loss saved at epochs specified by saveat_PBM variable
pred_logger_PBM, losses_PBM, NN_logger_PBM = MLP.get_indicators_PBM_numerical(pnn_PBM, re_PBM, dataset, epochs_PBM, p_range, dp, opt_PBM, inputs_one_hot_stand, verbose=verbose, saveat=saveat_PBM)

# extract results at last save point
pred_NN_PBM = pred_logger_PBM[end][1]
indicator_NN_PBM = pred_logger_PBM[end][2]
loss_NN_PBM = pred_logger_PBM[end][3]

# save results of prediction-based method in save_folder
open(save_folder*"hyperparams_PBM.txt", "w") do io
  writedlm(io, [saveat_PBM, lr_PBM, epochs_PBM])
end

open(save_folder*"NN_logger_PBM.txt", "w") do io
  writedlm(io, NN_logger_PBM)
end

for i in 1:length(pred_logger_PBM)
  open(save_folder*"pred_logger_PBM_"*string(i)*".txt", "w") do io
    writedlm(io, pred_logger_PBM[i][1])
  end

  open(save_folder*"indicator_logger_PBM_"*string(i)*".txt", "w") do io
    writedlm(io, pred_logger_PBM[i][2])
  end

  open(save_folder*"loss_logger_PBM_"*string(i)*".txt", "w") do io
    writedlm(io, pred_logger_PBM[i][3])
  end

end

open(save_folder*"pred_NN_PBM.txt", "w") do io
  writedlm(io, pred_NN_PBM)
end

open(save_folder*"indicator_NN_PBM.txt", "w") do io
  writedlm(io, indicator_NN_PBM)
end

open(save_folder*"loss_NN_PBM.txt", "w") do io
  writedlm(io, loss_NN_PBM)
end

open(save_folder*"NN_structure_PBM.txt", "w") do io
  writedlm(io, n_nodes_PBM)
end

open(save_folder*"losses_PBM.txt", "w") do io
  writedlm(io, losses_PBM)
end

# plot and save results of prediction-based method in save_folder
plt = plot(p_range, pred_NN_PBM, dpi=300, c="black", label="NN", ylims=(p_range[1], p_range[end]))
plot!(p_range, pred_opt_PBM, dpi=300, c="blue", label="analytical", ylims=(p_range[1],p_range[end]))
vline!([p_crit], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$\hat{y}_{\mathrm{PBM}}$")
savefig(save_folder*"pred_NN_PBM.png")

plt = plot(p_range[2:end-1], indicator_NN_PBM, dpi=300, c="black", label="NN")
plot!(p_range[2:end-1], indicator_opt_PBM, dpi=300, c="blue", label="analytical")
vline!([p_crit], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$I^{\mathrm{NN}}_{\mathrm{PBM}}$")
savefig(save_folder*"indicator_NN_PBM.png")

plt = plot(collect(1:length(losses_PBM)), losses_PBM, dpi=300, c="black", label="NN")
hline!([loss_opt_PBM], label="analytical", c="blue")
xlabel!("Epochs")
ylabel!(L"$\mathcal{L}_{\mathrm{PBM}}$")
savefig(save_folder*"loss_NN_PBM.png")

# learning by confusion using neural networks

# initialize neural network
n_nodes_LBC = [size(inputs_one_hot)[1], 64, 2]
NN_LBC = Chain(
  Dense(n_nodes_LBC[1], n_nodes_LBC[2], relu),
  Dense(n_nodes_LBC[2], n_nodes_LBC[end], identity))
pnn_LBC, re_LBC = Flux.destructure(NN_LBC)

# set hyperparameters
lr_LBC = 0.001f0
epochs_LBC = 1000
saveat_LBC = epochs_LBC
opt_LBC = ADAM(lr_LBC)
verbose = true

indicator_NN_LBC = []
loss_NN_LBC = []
for indx in collect(1:length(p_range_LBC))
  pred_logger_LBC, losses_LBC, NN_logger_LBC = MLP.get_indicators_LBC_numerical_fixed_p(deepcopy(pnn_LBC), re_LBC, dataset, epochs_LBC, p_range, dp, opt_LBC, p_range_LBC, indx, inputs_one_hot_stand, saveat=saveat_LBC, verbose=verbose)

  indicator_NN_LBC_p = pred_logger_LBC[end][1][1]
  push!(indicator_NN_LBC, indicator_NN_LBC_p)
  loss_NN_LBC_p = pred_logger_LBC[end][2][1]
  push!(loss_NN_LBC, loss_NN_LBC_p)
end

# save results of learning by confusion in save_folder
open(save_folder*"hyperparams_LBC.txt", "w") do io
  writedlm(io, [saveat_LBC, lr_LBC, epochs_LBC])
end

open(save_folder*"indicator_NN_LBC.txt", "w") do io
  writedlm(io, indicator_NN_LBC)
end

open(save_folder*"loss_NN_LBC.txt", "w") do io
  writedlm(io, loss_NN_LBC)
end

open(save_folder*"NN_structure_LBC.txt", "w") do io
  writedlm(io, n_nodes_LBC)
end

# plot and save results of learning by confusion in save_folder
plt = plot(p_range_LBC, indicator_NN_LBC, dpi=300, c="black", label="NN")
plot!(p_range_LBC, indicator_opt_LBC, dpi=300, c="blue", label="analytical")
vline!([p_crit], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$I_{\mathrm{LBC}}$")
savefig(save_folder*"indicator_NN_LBC.png")

plt = plot(p_range_LBC, loss_NN_LBC, dpi=300, c="black", label="NN")
plot!(p_range_LBC, loss_opt_LBC, dpi=300, c="blue", label="analytical")
vline!([p_crit], label=L"p_{c}",c="red")
xlabel!(L"$p$")
ylabel!(L"$\mathcal{L}_{\mathrm{LBC}}$")
savefig(save_folder*"loss_NN_LBC.png")


text:
module OldDynamics
using DifferentialEquations
using ForwardDiff
using DiffResults
using StaticArrays
using ..ProbInfo, ..LinPoint

const mass_idx = 1
const r_idx = SVector(2,3,4)
const v_idx = SVector(5,6,7)
const qbi_idx = SVector(8,9,10,11)
const omb_idx = SVector(12,13,14)
const state_dim = 14
const control_dim = 3

struct IntegratorParameters{T}
    dt :: Float64
    uk :: SArray{Tuple{3}, T, 1, 3}
    up :: SArray{Tuple{3}, T, 1, 3}
end

struct LinRes
    Ak::Array{Float64,2}
    Bk::Array{Float64,2}
    Ck::Array{Float64,2}
    Sigk::Array{Float64,1}
    Zk::Array{Float64,1}
    iter_end::Array{Float64, 1} #dunno
end

function DCM(quat::SArray{Tuple{4}, T, 1, 4} where T)
    q0 = quat[1]
    q1 = quat[2]
    q2 = quat[3]
    q3 = quat[4]

    return transpose(@SMatrix [(1-2*(q2^2 + q3^2))         (2*(q1 * q2 + q0 * q3))  (2*(q1*q3 - q0 * q2));
                    (2*(q1 * q2 - q0 * q3))     (1-2*(q1^2 + q3^2))      (2*(q2*q3 + q0 * q1));
                    (2*(q1 * q3 + q0 * q2))     (2*(q2*q3 - q0 * q1))    (1-2*(q1^2 + q2^2)) ])
end

function Omega{T}(omegab::SArray{Tuple{3}, T, 1, 3})
    z = zero(T)
    return @SMatrix [z        -omegab[1] -omegab[2] -omegab[3];
                    omegab[1] z           omegab[3] -omegab[2];
                    omegab[2] -omegab[3]  z          omegab[1];
                    omegab[3]  omegab[2] -omegab[1]  z         ]
end

function dx(state::SArray{Tuple{14}, T, 1, 14} where T, u::MArray{Tuple{3}, T, 1, 3} where T, info::ProbInfo)
    qbi = state[qbi_idx]
    omb = state[omb_idx]
    thr_acc = DCM(qbi) * u/state[mass_idx]
    grv_acc = @SVector [-info.g0,0,0]
    acc = thr_acc + grv_acc
    rot_vel = 0.5*Omega(omb)*qbi
    rot_acc = info.jBi*(cross(info.rTB,u) - cross(omb,info.jB*omb))
    return @SVector [-info.a*norm(u),
            state[5],state[6],state[7],
            acc[1], acc[2], acc[3],
            rot_vel[1], rot_vel[2], rot_vel[3], rot_vel[4],
            rot_acc[1], rot_acc[2], rot_acc[3]]
end

function compute_A(state::SArray{Tuple{14}, T, 1, 14} where T, u::MArray{Tuple{3}, T, 1, 3} where T, target::MArray{Tuple{14,14}, T, 2, 196} where T, info::ProbInfo)
    ForwardDiff.jacobian!(target, istate -> dx(istate, u, info), state)
end

function compute_B(state::SArray{Tuple{14}, T, 1, 14} where T, u::MArray{Tuple{3}, T, 1, 3} where T, target::MArray{Tuple{14,3}, T, 2, 42} where T, info::ProbInfo)
    ForwardDiff.jacobian!(target, iu -> dx(state, iu, info), u)
end

function linearize_dynamics(states::Array{LinPoint,1}, sigma_lin::Float64, dt::Float64, info::ProbInfo)
    #preallocate memory
    const Aval = @MArray zeros(state_dim, state_dim)
    const Bval = @MArray zeros(state_dim, control_dim)
    const clin = @MVector zeros(control_dim)
    const cneg = @MVector zeros(control_dim)
    const cpls = @MVector zeros(control_dim)

    #compute indices into the combined accumulator matrix
    const state_idx = 1
    const A_start = state_idx + 1
    const A_end = A_start + state_dim - 1
    const B_start = A_end + 1
    const B_end = B_start + control_dim - 1
    const C_start = B_end + 1
    const C_end = C_start + control_dim - 1
    const Sig_idx = C_end + 1
    const Z_idx = Sig_idx + 1

    #integrand worker
    #the accumulator is laid out as
    # state:1 |     Ak:state_dim     | Bkm:control_dim | Bkp:control_dim | sigK:1 | zk:1
    # all are state_dim high. Dims are thus state_dim+control_dim*2+3 by state_dim.
    function integrand(du,
                       u,#::MArray{Tuple{acc_height,acc_width}, Float64, 2, acc_height*acc_width},
                       p::IntegratorParameters,
                       t::Float64)
        lkm = (p.dt-t)/p.dt
        lkp = t/p.dt
        broadcast!(*, cneg, lkm, p.uk)
        broadcast!(*, cpls, lkp, p.up)
        broadcast!(+, clin, cneg, cpls)

        sv = SVector{14}(u[:,state_idx])
        du[:,1] = dx(sv, clin, info)
        compute_A(sv, clin, Aval, info)
        compute_B(sv, clin, Bval, info)
        broadcast!(*, Aval, sigma_lin, Aval)
        broadcast!(*, Bval, sigma_lin, Bval)
        phi = u[:,A_start:A_end]
        iPhi = inv(phi)

        du[:,A_start:A_end] = Aval*phi
        du[:,B_start:B_end] = lkm*(iPhi * Bval)
        du[:,C_start:C_end] = lkp*(iPhi * Bval)
        du[:,Sig_idx] = iPhi*(du[:,1])
        du[:,Z_idx] = iPhi*(-Aval*u[:,1] - Bval * clin)
    end

    results = Array{LinRes,1}(length(states)-1)
    for i=1:length(states)-1
        uk = states[i].control
        up = states[i+1].control
        int_iv = hcat(states[i].state, eye(Float64, state_dim), zeros(Float64, state_dim, control_dim*2 + 2))
        prob = ODEProblem(integrand, int_iv, (0.0,dt), IntegratorParameters(dt, uk, up))
        sol = DifferentialEquations.solve(prob, DP5(); dtmin=0.001, force_dtmin=true)
        result = sol.u[end]

        Ak::Array{Float64,2} = result[:,A_start:A_end]
        Bkm::Array{Float64,2} = Ak*result[:,B_start:B_end]
        Bkp::Array{Float64,2} = Ak*result[:,C_start:C_end]
        Sigk::Array{Float64,1} = Ak*result[:,Sig_idx]
        zk::Array{Float64,1} = Ak*result[:,Z_idx]
        results[i] = LinRes(Ak,
            Bkm,
            Bkp,
            Sigk,
            zk,
            result[:,1])
    end
    return results
end

function next_step(dynam, ab, abn, state, control_k, control_kp, sigma, sigHat, relax)
    dk = dynam
    return dk.Ak * state + dk.Bk * control_k + dk.Ck * control_kp + dk.Sigk*sigma + dk.Zk + relax
end

export linearize_dynamics, LinRes
end

max_stars_repo_name: BenChung/SuccessiveConvexification

max_issues_repo_name: BenChung/SuccessiveConvexification

max_forks_repo_name: BenChung/SuccessiveConvexification


text:
export PolytopeVertexMesh
export name, vertices, polytopes, materials, material_names, groups, vtk_type

struct PolytopeVertexMesh{Dim,T,P<:Polytope} <: AbstractMesh
    vertices::Vector{Point{Dim,T}}
    polytopes::Vector{P}
    materials::Vector{UInt8}
    material_names::Vector{String}
    name::String
    groups::Dict{String,BitSet}
end

name(mesh::PolytopeVertexMesh) = mesh.name
points(mesh::PolytopeVertexMesh) = mesh.vertices
vertices(mesh::PolytopeVertexMesh) = mesh.vertices
polytopes(mesh::PolytopeVertexMesh) = mesh.polytopes
materials(mesh::PolytopeVertexMesh) = mesh.materials
material_names(mesh::PolytopeVertexMesh) = mesh.material_names
groups(mesh::PolytopeVertexMesh) = mesh.groups

# constructors
function PolytopeVertexMesh(vertices::Vector{Point{Dim,T}},
                            polytopes::Vector{P},
                            name::String
                           ) where {Dim,T,P}
    return PolytopeVertexMesh(vertices, polytopes, name, Dict{String,BitSet}())
end

function PolytopeVertexMesh(vertices::Vector{Point{Dim,T}},
                            polytopes::Vector{P}
                           ) where {Dim,T,P}
    return PolytopeVertexMesh(vertices, polytopes, "", Dict{String,BitSet}())
end

function PolytopeVertexMesh(mesh::VolumeMesh{Dim,T,U}) where {Dim,T,U}
    return PolytopeVertexMesh(
             mesh.points,
             map(i->_materialize_face_connectivity(i, mesh), eachindex(mesh.types)),
             mesh.materials,
             mesh.material_names,
             mesh.name,
             mesh.groups
            )
end

vtk_type(::Type{<:Triangle})               = VTK_TRIANGLE
vtk_type(::Type{<:Quadrilateral})          = VTK_QUAD
vtk_type(::Type{<:QuadraticTriangle})      = VTK_QUADRATIC_TRIANGLE
vtk_type(::Type{<:QuadraticQuadrilateral}) = VTK_QUADRATIC_QUAD
vtk_type(::Type{<:Tetrahedron})            = VTK_TETRA
vtk_type(::Type{<:Hexahedron})             = VTK_HEXAHEDRON
vtk_type(::Type{<:QuadraticTetrahedron})   = VTK_QUADRATIC_TETRA
vtk_type(::Type{<:QuadraticHexahedron})    = VTK_QUADRATIC_HEXAHEDRON

function Base.show(io::IO, mesh::PolytopeVertexMesh{Dim,T,P}) where {Dim,T,P}
    print(io, "PolytopeVertexMesh{",Dim,", ",T,", ")
    if isconcretetype(P)
        println(io, alias_string(P),"{",vertex_type(P),"}}")
    else
        println(io, P,"}")
    end
    println(io, "  ├─ Name      : ", mesh.name)
    size_B = Base.summarysize(mesh)
    if size_B < 1e6
        println(io, "  ├─ Size (KB) : ", string(@sprintf("%.3f", size_B/1000)))
    else
        println(io, "  ├─ Size (MB) : ", string(@sprintf("%.3f",size_B/1e6)))
    end
    println(io, "  ├─ Vertices  : ", length(mesh.vertices))
    println(io, "  ├─ Polytopes : ", length(mesh.polytopes))
    poly_types = unique!(map(typeof, mesh.polytopes))
    npoly_types = length(poly_types)
    for i = 1:npoly_types
        poly_type = poly_types[i]
        npoly = count(x->x isa poly_type,  mesh.polytopes)
        if i === npoly_types
            println(io, "  │  └─ ", rpad(alias_string(poly_type), 22), ": ", npoly)
        else
            println(io, "  │  ├─ ", rpad(alias_string(poly_type), 22), ": ", npoly)
        end
    end
    println(io, "  ├─ Materials : ", length(mesh.material_names))
    ngroups = length(mesh.groups)
    println(io, "  └─ Groups    : ", ngroups)
    if 0 < ngroups ≤ 5
        group_keys = sort(collect(keys(mesh.groups)))
        for i = 1:ngroups
            if i === ngroups
                println(io, "     └─ ", group_keys[i])
            else
                println(io, "     ├─ ", group_keys[i])
            end
        end
    end
end

max_stars_repo_name: khurrumsaleem/MOCNeutronTransport

max_issues_repo_name: khurrumsaleem/MOCNeutronTransport

max_forks_repo_name: khurrumsaleem/MOCNeutronTransport


text:
julia> s = "abcdefg"
"abcdefg"

julia> n = 3
3

julia> s[n:end]
"cdefg"

julia> m=2
2

julia> s[n:n+m]
"cde"

julia> s[1:end-1]
"abcdef"

julia> s[search(s,'c')]
'c'

julia> s[search(s,'c'):search(s,'c')+m]
"cde"

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode