text:
export train_lsq, encoding_icm

"Randomly perturbs codes"
function perturb_codes!(
  B::Union{Matrix{Int16},SharedMatrix{Int16}}, # in/out. Codes to perturb
  npert::Integer,         # in. Number of entries to perturb in each code
  h::Integer,             # in. The number of codewords in each codebook
  IDX::UnitRange{Int64},  # in. Subset of codes in B to perturb
  replace::Bool=true) # in. Whether to sample with replacement

  m, _ = size(B)
  n    = length(IDX)

  # Sample random perturbation indices (places to perturb) in B
  if replace
    # With replacements this is easy
    pertidx  = rand(1:m, npert, n)
  else
    # Each call to sample has about 8 allocations, so this can results in
    # several Million alloc calls too.
    pertidx  = Matrix{Integer}(npert, n)
    for i = 1:n
      # Sample npert unique values out of m
      Distributions.sample!(1:m, view(pertidx,:,i), replace=false, ordered=true)
    end
  end

  # Sample the values that will replace the new ones
  pertvals = rand(1:h, npert, n)
  # Perturb the solutions
  for i = 1:n
    for j = 1:npert
      B[ pertidx[j,i], IDX[i] ] = pertvals[j,i]
    end
  end

  return B
end

"Run iterated conditional modes on N problems"
function iterated_conditional_modes_cpp!(
  B::Matrix{UInt8},  # in/out. Initialization, and the place where the results are saved.
  unaries::Vector{Matrix{T}},     # in. Unary terms
  binaries::Vector{Matrix{T}},    # in. The binary terms
  binaries_t::Vector{Matrix{T}},  # in. Transposed version of the above
  cbpair2binaryidx::Matrix{Int32},
  cbi::Matrix{Int32},           # in. 2-by-ncbi. indices for inter-codebook interactions
  to_look::Union{UnitRange{Int64},Vector{Int64}},
  to_condition::Matrix{Int32},
  icmiter::Integer,             # in. number of iterations for block-icm
  npert::Integer,               # in. number of codes to perturb per iteration
  IDX::UnitRange{Int64},        # in. Index to save the result
  ub::Matrix{T}, bb::Matrix{T}, # in. Preallocated memory
  h::Integer, n::Integer, m::Integer,
  V::Bool) where T <: AbstractFloat

  cbpair2binaryidx = convert(Matrix{Int32}, cbpair2binaryidx.-1)
  to_condition     = convert(Matrix{Int32}, to_condition.-1)

  binaries   = hcat(binaries...)
  binaries_t = hcat(binaries_t...)

  @inbounds for i=1:icmiter # Do the number of passed iterations

    # for j = to_look
    for j = 1:m
      # Get the unaries that we will work on
      copyto!(ub, unaries[to_look[j]])

      ccall(("condition", encode_icm_so), Nothing,
        (Ptr{Cuchar}, Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat},
        Ptr{Cint}, Ptr{Cint}, Cint, Cint, Cint),
        B, ub, binaries, binaries_t,
        cbpair2binaryidx, to_condition[:,j], to_look[j]-1, n, m)

    end # for j = 1:m
  end # for i=1:icmiter

end

"Run iterated conditional modes on N problems"
function iterated_conditional_modes!(
  B::Union{Matrix{Int16},SharedMatrix{Int16}},  # in/out. Initialization, and the place where the results are saved.
  unaries::Vector{Matrix{T}},     # in. Unary terms
  binaries::Vector{Matrix{T}},    # in. The binary terms
  binaries_t::Vector{Matrix{T}},  # in. Transposed version of the above
  cbpair2binaryidx::Matrix{Int32},
  cbi::Matrix{Int32},           # in. 2-by-ncbi. indices for inter-codebook interactions
  to_look::Union{UnitRange{Int64},Vector{Int64}},
  to_condition::Matrix{Int32},
  icmiter::Integer,             # in. number of iterations for block-icm
  npert::Integer,               # in. number of codes to perturb per iteration
  IDX::UnitRange{Int64},        # in. Index to save the result
  ub::Matrix{T}, bb::Matrix{T}, # in. Preallocated memory
  h::Integer, n::Integer, m::Integer,
  V::Bool) where T <: AbstractFloat

  @inbounds for i=1:icmiter # Do the number of passed iterations

    # This is the codebook that we are updating (i.e., the node)
    jidx = 1;
    for j = to_look
      # Get the unaries that we will work on
      copyto!(ub, unaries[j])

      # These are all the nodes that we are conditioning on (i.e., the edges to 'absorb')
      for k = to_condition[:, jidx]

        # Determine the pairwise interactions that we'll use (for cache-friendliness)
        binariidx = cbpair2binaryidx[ j, k ]
        if j < k
          bb = binaries[ binariidx ]
        else
          bb = binaries_t[ binariidx ]
        end

        # Traverse the unaries, absorbing the appropriate binaries
        for l=1:n
          codek = B[k, IDX[l]]
          for ll = 1:h
            ub[ll, l] += bb[ ll, codek ]
          end
        end
      end #for k=to_condition

      # Once we are done conditioning, traverse the absorbed unaries to find mins
      inidx = 1
      for idx=IDX
        minv = ub[1, inidx]
        mini = 1
        for k = 2:h
          ubi = ub[k, inidx]
          if ubi < minv
            minv = ubi
            mini = k
          end
        end

        B[j, idx] = mini
        inidx = inidx .+ 1
      end

      jidx = jidx .+ 1

    end # for j=to_look
  end # for i=1:icmiter

end

# Encode using iterated conditional modes
function encode_icm_fully!(
  oldB::Matrix{Int16},  # in/out. Initialization, and the place where the results are saved.
  X::Matrix{T},                 # in. d-by-n data to encode
  C::Vector{Matrix{T}},         # in. m-long vector with d-by-h codebooks
  binaries::Vector{Matrix{T}},  # in. The binary terms
  cbi::Matrix{Int32},           # in. 2-by-ncbi. indices for inter-codebook interactions
  ilsiter::Integer,             # in. number of ILS iterations
  icmiter::Integer,             # in. number of ICM iterations
  randord::Bool,                # in. whether to randomize search order
  npert::Integer,               # in. number of codes to perturb per iteration
  IDX::UnitRange{Int64},        # in. Index to save the result
  cpp::Bool=true,               # in. Whether to use the cpp implementation
  V::Bool=true) where T <: AbstractFloat

  # Compute unaries
  @time begin
  unaries = get_unaries(X, C, V)

  h, n = size(unaries[1])
  m, _ = size(oldB)

  if cpp && h != 256
    error("The C++ implementation of ICM encoding only supports codebooks with 256 entries")
  end

  ncbi = length(binaries)

  # Create a transposed copy of the binaries
  binaries_t = similar(binaries)
  for i = 1:ncbi
    binaries_t[i] = collect(binaries[i]')
  end

  # Create an index from codebook pairs to indices
  cbpair2binaryidx = zeros(Int32, m, m)
  for i = 1:ncbi
    cbpair2binaryidx[ cbi[1,i], cbi[2,i] ] = i
    cbpair2binaryidx[ cbi[2,i], cbi[1,i] ] = i
  end

  # Preallocate some space
  bb = zeros(T, h, h)
  ub = zeros(T, h, n)
  end

  B = zeros(Int16, m, n)

  for ilstrack = 1:ilsiter

    prevcost = veccost( X, oldB, C )
    if nworkers() == 1
      B = zeros(Int16, m, n)
    else
      B = SharedArray{Int16}(m, n)
    end
    copyto!(B, oldB)

    # For codebook i, we have to condition on these codebooks
    to_look      = 1:m
    to_condition = zeros(Int32, m-1, m)
    for i = 1:m
      tmp = collect(1:m)
      splice!( tmp, i )
      to_condition[:,i] = tmp
    end
    # Make the order random
    if randord
      to_look      = randperm(m)
      to_condition = to_condition[:, to_look]
    end

    # Perturb the codes
    B = perturb_codes!(B, npert, h, IDX)

    # Run iterated conditional modes
    if cpp
      B = convert(Matrix{UInt8}, B .- one(Int16))
      iterated_conditional_modes_cpp!(B, unaries,
        binaries, binaries_t, cbpair2binaryidx, cbi,
        to_look, to_condition, icmiter, npert, IDX, ub, bb, h, n, m, V)
      B = convert(Matrix{Int16}, B) .+ one(Int16)
    else
      iterated_conditional_modes!(B, unaries,
        binaries, binaries_t, cbpair2binaryidx, cbi,
        to_look, to_condition, icmiter, npert, IDX, ub, bb, h, n, m, V)
    end

    # Keep only the codes that improved
    newcost = veccost( X, B, C )
    areequal = newcost .== prevcost
    arebetter = newcost .< prevcost
    if V; @printf(" ILS iteration %d/%d done. ", ilstrack, ilsiter); end
    if V; @printf("%5.2f%% new codes are equal. ", 100*sum(areequal)/n ); end
    if V; @printf("%5.2f%% new codes are better.\n", 100*sum(arebetter)/n ); end

    B[:, .~arebetter] = oldB[:, .~arebetter]
    copyto!(oldB, B)
  end

  return B
end

"""
    encoding_icm(X, oldB, C, ilsiter, icmiter, randord, npert, cpp=true, V=false) -> B

Given data and chain codebooks, find codes using iterated local search with ICM.

# Arguments
- `X::Matrix{T}`: `d`-by-`n` data to quantize
- `OldB::Matrix{Int16}`: `m`-by-`n` initial set of codes
- `ilsiter::Integer`: Number of iterated local search (ILS) iterations
- `icmiter::Integer`: Number of iterated conditional modes (ICM) iterations
- `randord::Bool`: Whether to use random order
- `npert::Integer`: Number of codes to perturb
- `cpp::Bool=true`: Whether to use the c++ implementation
- `V::Bool=false`: Whehter to print progress

# Returns
- `B::Matrix{Int16}`: `m`-by-`n` matrix with the new codes
"""
function encoding_icm(
  X::Matrix{T},         # d-by-n matrix. Data to encode
  oldB::Matrix{Int16},  # m-by-n matrix. Previous encoding
  C::Vector{Matrix{T}}, # m-long vector with d-by-h codebooks
  ilsiter::Integer,     # in. number of ILS iterations
  icmiter::Integer,     # in. number of ICM iterations
  randord::Bool,        # whether to use random order
  npert::Integer,       # the number of codes to perturb
  cpp::Bool=true,        # whether to use the cpp implementation
  V::Bool=true) where T <: AbstractFloat # whether to print progress

  d, n =    size( X )
  m    =  length( C )
  _, h = size( C[1] )

  # Compute binaries between all codebook pairs
  binaries, cbi = get_binaries( C )
  _, ncbi       = size( cbi )

  B = encode_icm_fully!(oldB, X, C, binaries, cbi, ilsiter, icmiter, randord, npert, 1:n, cpp, V)

  return B
end

"""
    train_lsq(X, m, h, R, B, C, niter, ilsiter, icmiter, randord, npert, cpp=true, V=false) -> C, B, obj

Train a local-search quantizer.
This method is typically initialized by [Chain quantization (ChainQ)](@ref)

# Arguments
- `X::Matrix{T}`: `d`-by-`n` data to quantize
- `m::Integer`: Number of codebooks
- `h::Integer`: Number of entries in each codebook (typically 256)
- `R::Matrix{T}`: `d`-by-`d` rotation matrix for initialization
- `B::Matrix{Int16}`: `m`-by-`n` matrix with pre-trained codes for initialization
- `C::Vector{Matrix{T}}`: `m`-long vector with `d`-by-`h` matrices. Each matrix is a pretrained codebook of size approximately `d`-by-`h`
- `niter::Integer`: Number of iterations to use
- `ilster::Integer`: Number of iterated local search (ILS) iterations
- `icmiter::Integer`: Number of iterated conditional modes (ICM) iterations
- `randord::Bool`: Whether to visit the nodes in a random order in ICM
- `npert::Integer`: Number of codes to perturb
- `cpp::Bool`: Whether to use a c++ implementation for encoding
- `V::Bool`: Whether to print progress

# Returns
- `C::Vector{Matrix{T}}`: `m`-long vector with `d`-by-`h` matrix entries. Each matrix is a codebook of size approximately `d`-by-`h`
- `B::Matrix{Int16}`: `m`-by-`n` matrix with the codes
- `obj::Vector{T}`: `niter`-long vector with the quantization error after each iteration
"""
function train_lsq(
  X::Matrix{T},         # d-by-n matrix of data points to train on.
  m::Integer,           # number of codebooks
  h::Integer,           # number of entries per codebook
  R::Matrix{T},         # init rotation
  B::Matrix{Int16},     # init codes
  C::Vector{Matrix{T}}, # init codebooks
  niter::Integer,       # number of optimization iterations
  ilsiter::Integer,     # number of ILS iterations to use during encoding
  icmiter::Integer,     # number of iterations in local search
  randord::Bool,        # whether to use random order
  npert::Integer,       # The number of codes to perturb
  cpp::Bool=true,       # whether to use ICM's cpp implementation
  V::Bool=true) where T <: AbstractFloat

  # if V
  println("**********************************************************************************************");
  println("Training LSQ with $m codebooks, $npert perturbations, $icmiter icm iterations and random order = $randord");
  println("**********************************************************************************************");
  # end

  d, n = size( X )

  # Initialize C
  RX = R' * X
  C = update_codebooks(RX, B, h, V, "fastbin")

  # Apply the rotation to the codebooks
  for i = 1:m; C[i] = R * C[i]; end
  @printf("%3d %e \n", -2, qerror( X, B, C ))

  # Initialize B
  @time B = encoding_icm(X, B, C, ilsiter, icmiter, randord, npert, cpp, V)
  @printf("%3d %e \n", -1, qerror( X, B, C ))

  obj = zeros(T, niter)

  for iter = 1:niter
    obj[iter] = qerror(X, B, C)
    @printf("%3d %e \n", iter, obj[iter])

    # Update the codebooks C
    C = update_codebooks(X, B, h, V, "fastbin")

    # Update the codes B
    @time B = encoding_icm(X, B, C, ilsiter, icmiter, randord, npert, cpp, V)
  end

  return C, B, obj
end

function experiment_lsq_query_base(
  Xt::Matrix{T}, # d-by-n. Data to learn codebooks from
  B::Matrix{T2}, # codes
  C::Vector{Matrix{T}}, # codebooks
  R::Matrix{T}, # rotation
  Xq::Matrix{T}, # d-by-n. Queries
  gt::Vector{UInt32}, # ground truth
  m::Integer,    # number of codebooks
  h::Integer,    # number of entries per codebook
  niter::Integer=25, # Number of k-means iterations for training
  knn::Integer=1000, # Compute recall @N for this value of N
  V::Bool=false) where {T <: AbstractFloat, T2 <: Integer} # whether to print progress

  # TODO expose these parameters
  ilsiter = 8
  icmiter = 4
  randord = true
  npert   = 4
  cpp     = true

  # Train LSQ
  d, _ = size(Xt)
  if V; @printf("Running CUDA LSQ training... "); end
  C, B, train_error = Rayuela.train_lsq(Xt, m, h, R, B, C, niter, ilsiter, icmiter, randord, npert, cpp, V)
  if V; @printf("done\n"); end

  norms_B, norms_C = get_norms_codebook(B, C)
  db_norms         = vec( norms_C[ norms_B ] )

  if V; print("Querying m=$m ... "); end
  @time dists, idx = linscan_lsq(B, Xq, C, db_norms, eye(Float32, d), knn)
  if V; println("done"); end

  recall = eval_recall(gt, idx, knn)
  return C, B, R, train_error, recall
end

function experiment_lsq(
  Xt::Matrix{T}, # d-by-n. Data to learn codebooks from
  B::Matrix{T2}, # codes
  C::Vector{Matrix{T}}, # codebooks
  R::Matrix{T}, # rotation
  Xb::Matrix{T}, # d-by-n. Base set
  Xq::Matrix{T}, # d-by-n. Queries
  gt::Vector{UInt32}, # ground truth
  m::Integer,    # number of codebooks
  h::Integer,    # number of entries per codebook
  niter::Integer=25, # Number of k-means iterations for training
  knn::Integer=1000,
  V::Bool=false) where {T <: AbstractFloat, T2 <: Integer} # whether to print progress

  # TODO expose these parameters
  ilsiter_train = 8
  ilsiter_base = 32
  icmiter = 4
  randord = true
  npert   = 4
  cpp     = false

  # Train LSQ
  d, _ = size(Xt)
  C, B, train_error = Rayuela.train_lsq(Xt, m, h, R, B, C, niter, ilsiter_train, icmiter, randord, npert, cpp, V)
  norms_B, norms_C = get_norms_codebook(B, C)

  # === Encode the base set ===
  B_base = convert(Matrix{Int16}, rand(1:h, m, size(Xb,2)))
  @time B_base = Rayuela.encoding_icm(Xb, B_base, C, ilsiter_base, icmiter, randord, npert, cpp, V)
  base_error = qerror(Xb, B_base, C)
  if V; @printf("Error in base is %e\n", base_error); end

  # Compute and quantize the database norms
  B_base_norms, db_norms_X = quantize_norms( B_base, C, norms_C )
  db_norms     = vec( norms_C[ B_base_norms ] )

  if V; print("Querying m=$m ... "); end
  @time dists, idx = linscan_lsq(B_base, Xq, C, db_norms, eye(Float32, d), knn)
  if V; println("done"); end

  recall = eval_recall(gt, idx, knn)
  return C, B, R, train_error, B_base, recall
end

"Runs an lsq experiment/demo"
function experiment_lsq(
  Xt::Matrix{T}, # d-by-n. Data to learn codebooks from
  Xb::Matrix{T}, # d-by-n. Base set
  Xq::Matrix{T}, # d-by-n. Queries
  gt::Vector{UInt32}, # ground truth
  m::Integer,    # number of codebooks
  h::Integer,    # number of entries per codebook
  niter::Integer=25, # Number of k-means iterations for training
  knn::Integer=1000,
  V::Bool=false) where T <: AbstractFloat # whether to print progress

  # OPQ initialization
  C, B, R, _ = train_opq(Xt, m, h, niter, "natural", V)

  # ChainQ (second initialization)
  # C, B, R, train_error = train_chainq(Xt, m, h, R, B, C, niter, V)

  # Actual experiment
  experiment_lsq(Xt, B, C, R, Xb, Xq, gt, m, h, niter, knn, V)
end


text:
module Synapse

using CuArrays

include("delta.jl")
include("alpha.jl")
include("epsp.jl")

end

ext: jl


text:
#md# The [engine](https://github.com/guo-yong-zhi/Stuffing.jl) is designed for general purpose, so the outputs don't have to be text, and shapes are OK
using WordCloud

sc = WordCloud.randomscheme()
l = 200
wc = wordcloud(
    repeat(["placeholder"], l), repeat([1], l),
    mask=shape(box, 400, 300, color=WordCloud.randommaskcolor(sc)),
    state=identity)
#md# * `words` & `weights` are just placeholders
#md# * style arguments like `colors`, `angles` and `density` have no effect
#md#
#md# And, you should manually initialize images for the placeholders, instead of calling `initwords!`
dens = 0.6
sz = 3expm1.(rand(l)) .+ 1
sz ./= √(sum(π * (sz ./ 2).^2 ./ dens) / prod(size(wc.mask))) # set a proper size according to the density
## svg version
# shapes = [shape(ellipse, round(sz[i]), round(sz[i]), color=rand(sc)) for i in 1:l]
# setsvgimages!(wc, :, shapes)
## bitmap version
shapes = WordCloud.tobitmap.([shape(ellipse, round(sz[i]), round(sz[i]), color=rand(sc)) for i in 1:l])
setimages!(wc, :, shapes)

setstate!(wc, :initwords!) # set the state flag after manual initialization
# @record "pattern_animation" overwrite=true generate!(wc, retry=1)
generate!(wc, retry=1, optimiser=WordCloud.Momentum(η=1/8)) # turn off rescale attempts. manually set images can't be rescaled
println("results are saved to pattern.png")
paint(wc, "pattern.png")
wc
#eval# runexample(:pattern)
#md# ![](pattern.png)


text:
using Base.Broadcast: AbstractArrayStyle, broadcasted

Numeric{T<:Number} = Union{T,AbstractArray{<:T}}

@tangent Broadcast.preprocess(dest, bc) =
  Broadcast.preprocess(dest, bc), (ddest, dbc) -> dbc

@tangent broadcasted(::typeof(identity), x::Numeric) = x, (_, ẋ) -> ẋ

@tangent broadcasted(::typeof(+), xs::Numeric...) =
  broadcast(+, xs...), (_, ẋs...) -> broadcast(+, ẋs...)

@tangent function broadcasted(::typeof(tanh), x::Numeric)
  y = tanh.(x)
  y, (_, ẋ) -> ẋ .* (1 .- y.^2)
end

@tangent function broadcasted(::AbstractArrayStyle, f, args...)
  error("Generic broadcast of $f not supported yet")
end


text:
function pricer(mcProcess::BaseProcess, rfCurve::AbstractZeroRateCurve, mcConfig::MonteCarloConfiguration{<:Integer, <:Integer, <:ControlVariates, <:BaseMode}, abstractPayoff::AbstractPayoff)
    set_seed(mcConfig)
    variate_handl = mcConfig.monteCarloMethod
    variate_conf = variate_handl.conf_variate
    variate_payoff = variate_handl.variate
    @assert maturity(abstractPayoff) == variate_payoff.T
    T = maturity(abstractPayoff)
    S_var = simulate(mcProcess, rfCurve, variate_conf, T)
    Payoff_var = payoff(S_var, variate_payoff, rfCurve)
    Payoff_opt_var = payoff(S_var, abstractPayoff, rfCurve)
    #c=-cov(Payoff_var,Payoff_opt_var)/var(Payoff_var)[1];
    c = -collect(cov(Payoff_var, Payoff_opt_var) / var(Payoff_var))[1]
    price_var = mean(Payoff_var)
    mcConfig_mod = MonteCarloConfiguration(mcConfig.Nsim, mcConfig.Nstep, variate_conf.monteCarloMethod, mcConfig.parallelMode, mcConfig.seed * 2)
    #END OF VARIATE SECTION
    Prices = pricer(mcProcess, rfCurve, mcConfig_mod, [abstractPayoff, variate_payoff])
    Price = Prices[1] + c * (Prices[2] - price_var)

    return Price
end

function pricer(mcProcess::BaseProcess, rfCurve::AbstractZeroRateCurve, mcConfig::MonteCarloConfiguration{<:Integer, <:Integer, <:ControlVariates{Forward{num}}, <:BaseMode}, abstractPayoff::AbstractPayoff) where {num <: Number}
    set_seed(mcConfig)
    variate_handl = mcConfig.monteCarloMethod
    variate_conf = variate_handl.conf_variate
    variate_payoff = variate_handl.variate
    @assert maturity(abstractPayoff) == variate_payoff.T
    T = maturity(abstractPayoff)
    S_var = simulate(mcProcess, rfCurve, variate_conf, T)
    Payoff_var = payoff(S_var, variate_payoff, rfCurve)
    Payoff_opt_var = payoff(S_var, abstractPayoff, rfCurve)
    c = -collect(cov(Payoff_var, Payoff_opt_var) / var(Payoff_var))[1]
    price_var = mcProcess.underlying.S0 * exp(-integral(mcProcess.underlying.d, T))
    mcConfig_mod = MonteCarloConfiguration(mcConfig.Nsim, mcConfig.Nstep, variate_conf.monteCarloMethod, mcConfig.parallelMode, mcConfig.seed + 3)
    #END OF VARIATE SECTION
    Prices = pricer(mcProcess, rfCurve, mcConfig_mod, [abstractPayoff, variate_payoff])
    Price = Prices[1] - c * (Prices[2] - price_var)
    return Price
end


text:
######################################################## model
F(z::BorderedArray,θ::AbstractVector) = F(z.u,(θ=θ,p=z.p))
function F(u::AbstractVector,parameters::NamedTuple)
	@unpack θ,p = parameters

	f = first(u)*first(p)*first(θ)
	F = similar(u,typeof(f))

	F[1] = θ[1] + p*u[1] + θ[2]*u[1]^3
	return F
end

######################################################### targets and initial guess
X = StateSpace( 1, -5:0.01:5, [0] )
θ = SizedVector{2}(4.0,-π/2)


text:
include("../src/BlockCiphers.jl")
include("testing.jl")

importall BlockCiphers

print("testing AES ")

plain = zeros(Uint8, 16)
k1 = AES(zeros(Uint8, int32(128/8)))
c1 = encrypt(k1, plain)
p1 = decrypt(k1, c1)

@assert_equal(c1, [0x66, 0xe9, 0x4b, 0xd4, 0xef, 0x8a, 0x2c, 0x3b, 0x88, 0x4c, 0xfa, 0x59, 0xca, 0x34, 0x2b, 0x2e])

k2 = AES(zeros(Uint8, int32(128/8)))
c2 = encrypt(k2, plain)
p2 = decrypt(k2, c1)

@assert_equal(c1, c2)
@assert_equal(plain, p1)
@assert_equal(plain, p2)

# from http://www.cs.utsa.edu/~wagner/laws/AEStestRuns.html
begin # Gladman's Test Data, 128-bit key
    key = AES([0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c])

    plain = [0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37, 0x07, 0x34]
    encrypted = encrypt(key, plain)

    @assert_equal(encrypted, [0x39, 0x25, 0x84, 0x1d, 0x02, 0xdc, 0x09, 0xfb, 0xdc, 0x11, 0x85, 0x97, 0x19, 0x6a, 0x0b, 0x32])

    decrypted = decrypt(key, encrypted)
    @assert_equal(plain, decrypted)
end

begin # Gladman's Test Data, 192-bit key
    key = AES([0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c, 0x76, 0x2e, 0x71, 0x60, 0xf3, 0x8b, 0x4d, 0xa5])

    plain = [0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37, 0x07, 0x34]
    encrypted = encrypt(key, plain)

    @assert_equal(encrypted, [0xf9, 0xfb, 0x29, 0xae, 0xfc, 0x38, 0x4a, 0x25, 0x03, 0x40, 0xd8, 0x33, 0xb8, 0x7e, 0xbc, 0x00])

    decrypted = decrypt(key, encrypted)
    @assert_equal(plain, decrypted)
end

begin # Gladman's Test Data, 256-bit key
    key = AES([0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c, 0x76, 0x2e, 0x71, 0x60, 0xf3, 0x8b, 0x4d, 0xa5, 0x6a, 0x78, 0x4d, 0x90, 0x45, 0x19, 0x0c, 0xfe])

    plain = [0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37, 0x07, 0x34]
    encrypted = encrypt(key, plain)

    @assert_equal(encrypted, [0x1a, 0x6e, 0x6c, 0x2c, 0x66, 0x2e, 0x7d, 0xa6, 0x50, 0x1f, 0xfb, 0x62, 0xbc, 0x9e, 0x93, 0xf3])

    decrypted = decrypt(key, encrypted)
    @assert_equal(plain, decrypted)
end

begin # AES Specification Test Data, 128-bit key
    key = AES([0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f])

    plain = [0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff]
    encrypted = encrypt(key, plain)

    @assert_equal(encrypted, [0x69, 0xc4, 0xe0, 0xd8, 0x6a, 0x7b, 0x04, 0x30, 0xd8, 0xcd, 0xb7, 0x80, 0x70, 0xb4, 0xc5, 0x5a])

    decrypted = decrypt(key, encrypted)
    @assert_equal(plain, decrypted)
end

begin # AES Specification Test Data, 192-bit key
    key = AES([0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17])

    plain = [0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff]
    encrypted = encrypt(key, plain)

    @assert_equal(encrypted, [0xdd, 0xa9, 0x7c, 0xa4, 0x86, 0x4c, 0xdf, 0xe0, 0x6e, 0xaf, 0x70, 0xa0, 0xec, 0x0d, 0x71, 0x91])

    decrypted = decrypt(key, encrypted)
    @assert_equal(plain, decrypted)
end

begin # AES Specification Test Data, 256-bit key
    key = AES([0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f])

    plain = [0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff]
    encrypted = encrypt(key, plain)

    @assert_equal(encrypted, [0x8e, 0xa2, 0xb7, 0xca, 0x51, 0x67, 0x45, 0xbf, 0xea, 0xfc, 0x49, 0x90, 0x4b, 0x49, 0x60, 0x89])

    decrypted = decrypt(key, encrypted)
    @assert_equal(plain, decrypted)
end

println()

print("testing DES ")

begin
    key = DES([0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
    plain = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]

    encrypted = encrypt(key, plain)
    @assert_equal(encrypted, [0x95, 0xa8, 0xd7, 0x28, 0x13, 0xda, 0xa9, 0x4d])
    @assert_equal(plain, decrypt(key, encrypted))
end

begin
    key = DES([0x40, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
    plain = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]

    encrypted = encrypt(key, plain)
    @assert_equal(encrypted, [0x0e, 0xec, 0x14, 0x87, 0xdd, 0x8c, 0x26, 0xd5])
    @assert_equal(plain, decrypt(key, encrypted))
end

begin
    key = DES([0x20, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
    plain = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]

    encrypted = encrypt(key, plain)
    @assert_equal(encrypted, [0x7a, 0xd1, 0x6f, 0xfb, 0x79, 0xc4, 0x59, 0x26])
    @assert_equal(plain, decrypt(key, encrypted))
end

begin
    key = DES([0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
    plain = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]

    encrypted = encrypt(key, plain)
    @assert_equal(encrypted, [0xd3, 0x74, 0x62, 0x94, 0xca, 0x6a, 0x6c, 0xf3])
    @assert_equal(plain, decrypt(key, encrypted))
end

begin
    key = DES([0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
    plain = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]

    encrypted = encrypt(key, plain)
    @assert_equal(encrypted, [0x80, 0x9f, 0x5f, 0x87, 0x3c, 0x1f, 0xd7, 0x61])
    @assert_equal(plain, decrypt(key, encrypted))
end

begin
    key = DES([0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
    plain = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]

    encrypted = encrypt(key, plain)
    @assert_equal(encrypted, [0xc0, 0x2f, 0xaf, 0xfe, 0xc9, 0x89, 0xd1, 0xfc])
    @assert_equal(plain, decrypt(key, encrypted))
end

begin
    key = DES([0x2b, 0xd6, 0x45, 0x9f, 0x82, 0xc5, 0xb3, 0x00])
    plain = [0xea, 0x02, 0x47, 0x14, 0xad, 0x5c, 0x4d, 0x84]

    encrypted = encrypt(key, plain)
    @assert_equal(encrypted, [0x12, 0x6e, 0xfe, 0x8e, 0xd3, 0x12, 0x19, 0x0a])
    @assert_equal(plain, decrypt(key, encrypted))
end

println()


text:
module KongYiji

import Base: size, split, display, read, write, first, stat, push!, length, iterate, ==, getindex, summary, collect, values, isless, get, eachmatch, in

using Pkg
using JLD2, FileIO
using Random, DataStructures, ProgressMeter, DelimitedFiles, DataFrames

export postable, h2vtable, v2htable, h2htable, hprtable, kwictable,
        ChTreebank, ZhuXian, posidsents, wordsents, rawsents,
        UselessTable, HMM, NGramHMM

#=
export generate!, train!, LM, LMHMM
=#

const Ti = UInt32
const Tv = Float32
const oo = Tv(Inf)
const EPS = Tv(1e-9)

include("utils.jl")

include("CtbTree.jl")
include("ChTreebank.jl")

include("ZhuXian.jl")

include("UselessTable.jl")
include("compress.jl")
include("AhoCorasickAutomaton.jl")

include("HmmUtils.jl")
include("HmmScoreTable.jl")
include("IcwbScoreTable.jl")
include("HMM.jl")
include("NGramHMM.jl")

#include("LM.jl")
#include("Heap.jl")
#include("LMHMM.jl")

end # module


text:
module NotifyMe

using Requires, UUIDs

const notified_args = []
add_require(args...) = push!(notified_args, args)

function __init__()
    push!(Requires.notified_pkgs, Base.PkgId(UUID(0x545cf9b9f5754090a54a9a5287d37f74), "NotifyMe"))
end

end # module


text:
#!/usr/bin/env julia

using Pkg
using Dates
using Distributed
using ArgParse
import MIToS
@everywhere using MIToS

using MIToS.Utils.Scripts

function parse_commandline()
	mitos_version = loadedversion(MIToS)
    s = ArgParseSettings(description = "Download gzipped files from PDB.",
						 version = "MIToS $(mitos_version)",
						 add_version = true)

    @add_arg_table! s begin
        "--code", "-c"
            help = "PDB code"
        "--list", "-l"
            help = "File with a list of PDB codes (one per line)"
        "--format", "-f"
            help = "Format. It should be PDBFile (pdb) or PDBML (xml)"
            arg_type = String
            default = "PDBML"
    end

    s.epilog = """
    \n
    MIToS $(mitos_version)\n
    \n
    Bioinformatics Unit\n
    Leloir Institute Foundation\n
    Av. Patricias Argentinas 435, CP C1405BWE, Buenos Aires, Argentina
    """

    return parse_args(s)
end

const parsed = parse_commandline()

function _file_names(args)
  file = args["code"]
  list = args["list"]
  if file !== nothing && list === nothing
    return String[ file ]
  elseif list !== nothing && file === nothing
    return String[ chomp(line) for line in open(readlines, list, "r") ]
  else
    throw(ErrorException("You must use --code or --list and the filename; --help for more information."))
  end
end

const files = _file_names(parsed)

@everywhere Args = remotecall_fetch(()->parsed,1) # Parsed ARGS for each worker
@everywhere FileList = remotecall_fetch(()->files,1) # List of Files for each worker

@everywhere function main(input) # input must be a file
  try
    format = Args["format"]
    if format == "PDBML" || format == "PDBFile"
      MIToS.PDB.downloadpdb(input, format=(format == "PDBML" ? MIToS.PDB.PDBML : MIToS.PDB.PDBFile))
    else
      throw(ErrorException("--format should be PDBML or PDBFile"))
    end
  catch err
    println("ERROR: ", input)
    println(err)
  end
end

pmap(main, FileList) # Run each file in parallel (with -l)


text:
module AFK

include("Knetapi.jl");   export Parameter
include("core.jl")
include("primitive.jl"); export Multiply, Embed, BatchMul, Linear, Dense
include("compound.jl");  export MLP

end # module


text:
module GSquery

end


text:
# Sensitivity Analysis (SA)
# ===========================
# NOTE: run this script after evaluating main.jl under benchmark scenario

# # ============================================
# # SECTION: define a macro to solve a model 定义一个求解一遍模型的macro()
# macro quickproc_SolveModel()
#     # NOTE: PrintMode in ["full", "concise", "final", "silent"]
#        Guess = ( r = 0.08, L = 0.2 )
#        @time EasySearch.SteadyState!( 1, Guess, Dt, Dst, Pt, Ps, Pc, env,
#           atol = 1E-8, MaxIter = 1000, PrintMode = "final", MagicNum = 2.0, StepLen = 0.5 )
#        Guess = ( r = 0.12, L = 0.75 )
#        @time EasySearch.SteadyState!( env.T, Guess, Dt, Dst, Pt, Ps, Pc, env,
#           atol = 1E-6, MaxIter = 1000,
#           PrintMode = "final", MagicNum = 2.0, StepLen = 0.5 )
#        @time PerfLog = EasySearch.Transition!( Dt, Dst, Pt, Ps, Pc, env,
#           atol = 1E-3, MaxIter = 500,
#           PrintMode = "full", MagicNum = 2.0, StepLen = 0.7, ReturnLog = true )
#        EasySearch.ProcAfterTransition!( Dt, Dst, Pt, Ps, Pc, env )  # compute other variables
# end # end macro

# ============================================
# SECTION: decorate parameters & solving & save models out 开始敏感性测试
# NOTE: we test the relative changes away from the benchmark levels of tested parameters
#       我们测试参数偏离基准情形的相对幅度
# NOTE: and follow "Pannell D J . Sensitivity analysis: strategies, methods, concepts, examples[J]. Agric Econ, 1997."
#       to design tests, then report the results
#       然后参照(Pannel 1997)设计分析并报告结果
# NOTE: the selected objective variable is: the RELATIVE CHANGE of [LI/PoolExp, LI/PoolIn, LI/GDP, LI/TexRev] from 2010 to 2110
#       e.g. x(2110) - x(2010), where x(t) is one of the four variables
#       选定的作为敏感性分析的目标量(output variable)的是正文政策模拟采用的四个指标【在2010至2100这100年间的**相对**变化幅度】
#       the reported figure uses [the % change of tested parameter] as x-axis, and [the % change of objective variable RISE] as y-axis; when lined, approximated slopes indicates elasticities
#       报告的图形以[测试参数相对基准的百分比变化]为x轴，以[目标量的百分比变化]为Y轴，当把某个参数的测试结果连线，近似的斜率就是目标量对测试参数的弹性关系
#       can report multiple tests (tested parameters) on one objective variable in one figure
#       可以在一张图里报告基于一个目标量的对不同测试参数的敏感性分析结果
# -----------------

# relative turbulence vector on parameters (in digits)
SA_RelTurb = [ -0.1, -0.05, 0.0, 0.05, 0.1 ]
# a (generalizable) Dict to save SA results (vectors with length of vecRelTurb_SensiTest::Vector)
SA_ResSet_LI2Exp = Dict( :Turbulence => SA_RelTurb )
SA_ResSet_LI2Inc = Dict( :Turbulence => SA_RelTurb )
SA_ResSet_LI2GDP = Dict( :Turbulence => SA_RelTurb )
SA_ResSet_LI2Tax = Dict( :Turbulence => SA_RelTurb )
# tested parameter name (Symbol)  -> for CONSTANT parameters in Pc::Dict
SA_ParCandidate = [ :α, :γ, :κ, :μ, :σ, :δ ]

# you should run this loop for enough times to fill the results of diff parameters in SA_ResSet_xxx::Dict
# then these Dict will be used to do plotting ()
for tmpPar in SA_ParCandidate
   # 0. create vectors for current tested parameter to save results
   SA_ResSet_LI2Exp[Symbol(tmpPar)] = []
   SA_ResSet_LI2Inc[Symbol(tmpPar)] = []
   SA_ResSet_LI2GDP[Symbol(tmpPar)] = []
   SA_ResSet_LI2Tax[Symbol(tmpPar)] = []

   for tmpTurb in SA_RelTurb
      # ----------------
      # 1. first, refresh all data & parameters 首先刷新所有数据&参数
      include("../src/proc_VarsDeclare.jl")  # NOTE: please be careful about directory references
      include("../src/proc_InitPars.jl")

      # 2. then, manually decorate/edit the parameters to test 然后手动修饰参数
      # NOTE: please MANUALLY edit the parameters to test
      #        请手动修改到底测试哪个参数
      #        though we provide all parameters we test in our paper and comment most of them
      #        虽然我们提供了所有论文中我们涉及到的参数并且将大部分注释掉了
      # ----------------------
      # all relevant parameters:
      # \alpha, \gamma, \kappa, \mu, \sigma  (CONSTANT)
      # \beta_t, \eta_t, \theta_t, z_t, \zeta_t, \phi_t, \mathbb{A}_t, \mathbb{B}_t  (TIME SERIES, not recommended)
      # ---------------------
         # for constant parameters in Pc::Dict
            Pc[ Symbol(tmpPar) ] *= 1.0 + tmpTurb

      # -------------------
      # 3. solve the model & compute temporaty variables
      Guess = ( r = 0.08, L = 0.2 )
      EasySearch.SteadyState!( 1, Guess, Dt, Dst, Pt, Ps, Pc, env,
         atol = 1E-8, MaxIter = 1000, PrintMode = "final", MagicNum = 2.0, StepLen = 0.5 )
      Guess = ( r = 0.12, L = 0.75 )
      EasySearch.SteadyState!( env.T, Guess, Dt, Dst, Pt, Ps, Pc, env,
         atol = 1E-6, MaxIter = 1000,
         PrintMode = "final", MagicNum = 2.0, StepLen = 0.5 )
      PerfLog = EasySearch.Transition!( Dt, Dst, Pt, Ps, Pc, env,
         atol = 1E-3, MaxIter = 500,
         PrintMode = "full", MagicNum = 2.0, StepLen = 0.7, ReturnLog = true )
      EasySearch.ProcAfterTransition!( Dt, Dst, Pt, Ps, Pc, env )  # compute other variables

      # ----------------
      # 4. collect objective variable series
      tmpLoc = [2010, 2110] .- env.START_YEAR .+ 1
      tmpRes = (
         LI2Exp = Dt[:LI] ./ Dt[:AggPoolExp],
         LI2Inc = Dt[:LI] ./ Dt[:AggPoolIn],
         LI2GDP = Dt[:LI] ./ Dt[:Y],
         LI2Tax = Dt[:LI] ./ (Dt[:TRc] .+ Dt[:TRw]),
      )

      # -------------------
      # 5. collect SA result data through push!() method
      # NOTE: use esp() to avoid domain error (devided by zero)
      push!( SA_ResSet_LI2Exp[Symbol(tmpPar)] , tmpRes.LI2Exp[tmpLoc[2]] / (tmpRes.LI2Exp[tmpLoc[1]] .+ eps()) - 1 )
      push!( SA_ResSet_LI2Inc[Symbol(tmpPar)] , tmpRes.LI2Inc[tmpLoc[2]] / (tmpRes.LI2Inc[tmpLoc[1]] .+ eps()) - 1 )
      push!( SA_ResSet_LI2GDP[Symbol(tmpPar)] , tmpRes.LI2GDP[tmpLoc[2]] / (tmpRes.LI2GDP[tmpLoc[1]] .+ eps()) - 1 )
      push!( SA_ResSet_LI2Tax[Symbol(tmpPar)] , tmpRes.LI2Tax[tmpLoc[2]] / (tmpRes.LI2Tax[tmpLoc[1]] .+ eps()) - 1 )

   end
end

# ==============================
# finally, do plotting with PyPlot
# but before that, save the results! (time-costing the SA is! value it! lol)
# NOTE: CAUTION! NOT NORMALIZED ON BENCHMARK!!! 注意，输出的原始数据尚未对benchmark进行rescale！
EasyIO.writecsv( string("./output/SA_LI2Exp_",EasyIO.LogTag(),".csv"), SA_ResSet_LI2Exp )
EasyIO.writecsv( string("./output/SA_LI2Inc_",EasyIO.LogTag(),".csv"), SA_ResSet_LI2Inc )
EasyIO.writecsv( string("./output/SA_LI2GDP_",EasyIO.LogTag(),".csv"), SA_ResSet_LI2GDP )
EasyIO.writecsv( string("./output/SA_LI2Tax_",EasyIO.LogTag(),".csv"), SA_ResSet_LI2Tax )

# -----------
# NOTE: SA_ParCandidate = [ :α, :γ, :δ, :κ, :μ, :σ ]; SA_RelTurb = [ -0.1, -0.05, 0.0, 0.05, 0.1 ]
#       LineStyle: all dashes --
#       markers: circle o, diamond d, star *, upper-triangle ^, crux x, square s
#       all converted from digits to per-cent-ages
# NOTE: need to normalize/centralize SA results on parameter-change = 0 (benchmark)
#       to obtain the relative change from benchmark
#       注意，展示时需要基于benchmark进行中心化得到相对于benchmark的百分比偏离!
# -----------
PyPlot.figure( figsize = [16,9] )  # figsize measured in inches
   PyPlot.subplot(2,2,1)  # subfigure: gap / pool account expenditure
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Exp[:α] ./SA_ResSet_LI2Exp[:α][3], "--o" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Exp[:γ] ./SA_ResSet_LI2Exp[:γ][3], "--d" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Exp[:δ] ./SA_ResSet_LI2Exp[:δ][3], "--*" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Exp[:κ] ./SA_ResSet_LI2Exp[:κ][3], "--^" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Exp[:μ] ./SA_ResSet_LI2Exp[:μ][3], "--x" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Exp[:σ] ./SA_ResSet_LI2Exp[:σ][3], "--s" )
      PyPlot.xlabel("Parameter (% change from benchmark)")
      PyPlot.ylabel("Relative change from benchmark (%)")
      PyPlot.title("Rise from 2010 to 2110: Pool gap / Pool expenditure")
      PyPlot.grid(true)
      PyPlot.legend(SA_ParCandidate)
   # --------------------------
   PyPlot.subplot(2,2,2)  # subfigure: gap / pool account incomes
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Inc[:α] ./SA_ResSet_LI2Inc[:α][3], "--o" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Inc[:γ] ./SA_ResSet_LI2Inc[:γ][3], "--d" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Inc[:δ] ./SA_ResSet_LI2Inc[:δ][3], "--*" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Inc[:κ] ./SA_ResSet_LI2Inc[:κ][3], "--^" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Inc[:μ] ./SA_ResSet_LI2Inc[:μ][3], "--x" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Inc[:σ] ./SA_ResSet_LI2Inc[:σ][3], "--s" )
      PyPlot.xlabel("Parameter (% change from benchmark)")
      PyPlot.ylabel("Relative change from benchmark (%)")
      PyPlot.title("Rise from 2010 to 2110: Pool gap / Pool incomes")
      PyPlot.grid(true)
      PyPlot.legend(SA_ParCandidate)
   # --------------------------
   PyPlot.subplot(2,2,3)  # subfigure: gap / GDP
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2GDP[:α] ./SA_ResSet_LI2GDP[:α][3], "--o" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2GDP[:γ] ./SA_ResSet_LI2GDP[:γ][3], "--d" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2GDP[:δ] ./SA_ResSet_LI2GDP[:δ][3], "--*" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2GDP[:κ] ./SA_ResSet_LI2GDP[:κ][3], "--^" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2GDP[:μ] ./SA_ResSet_LI2GDP[:μ][3], "--x" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2GDP[:σ] ./SA_ResSet_LI2GDP[:σ][3], "--s" )
      PyPlot.xlabel("Parameter (% change from benchmark)")
      PyPlot.ylabel("Relative change from benchmark (%)")
      PyPlot.title("Rise from 2010 to 2110: Pool gap / GDP")
      PyPlot.grid(true)
      PyPlot.legend(SA_ParCandidate)
   # --------------------------
   PyPlot.subplot(2,2,4)  # subfigure: gap / Tax Revenues
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Tax[:α] ./SA_ResSet_LI2Tax[:α][3], "--o" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Tax[:γ] ./SA_ResSet_LI2Tax[:γ][3], "--d" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Tax[:δ] ./SA_ResSet_LI2Tax[:δ][3], "--*" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Tax[:κ] ./SA_ResSet_LI2Tax[:κ][3], "--^" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Tax[:μ] ./SA_ResSet_LI2Tax[:μ][3], "--x" )
      PyPlot.plot( 100 .* SA_RelTurb, 100 .* SA_ResSet_LI2Tax[:σ] ./SA_ResSet_LI2Tax[:σ][3], "--s" )
      PyPlot.xlabel("Parameter (% change from benchmark)")
      PyPlot.ylabel("Relative change from benchmark (%)")
      PyPlot.title("Rise from 2010 to 2110: Pool gap / Tax revenues")
      PyPlot.grid(true)
      PyPlot.legend(SA_ParCandidate)
   # --------------------------
   # layout adjustment
   PyPlot.tight_layout()

# ============================
# finally, finally, save out the figure
   PyPlot.savefig( string( "./output/SAresult_",EasyIO.LogTag(),".pdf" ), format = "pdf" )

#


text:
__precompile__(true)
module Kore

include("integer.jl")
include("array.jl")
include("float.jl")

include("function.jl")

export flatten
export chopzero, chopzero!
export funcplus, functimes
export linpath
export within

end


text:
#################### Model-Based Posterior Statistics ####################
"""
    dic(mc::ModelChains)

Compute the Deviance Information Criterion (DIC) of Spiegelhalter et al. and Gelman et al. from MCMC sampler output.

Returns a `ChainSummary` type object with DIC results from the methods of Spiegelhalter and Gelman in the first and second rows of the `value` field, and the DIC value and effective numbers of parameters in the first and second columns; where

``\\text{DIC} = -2 \\mathcal{L}(\\bar{\\Theta}) + 2 p,``

such that ``\\mathcal{L}(\\bar{\\Theta})`` is the log-likelihood of model outputs given the expected values of model parameters ``\\Theta``, and ``p`` is the effective number of parameters. The latter is defined as ``p_D = -2 \\bar{\\mathcal{L}}(\\Theta) + 2 \\mathcal{L}(\\bar{\\Theta})`` for the method of Spiegelhalter and as ``p_V = \\frac{1}{2} \\operatorname{var}(-2 \\mathcal{L}(\\Theta))`` for the method of Gelman. Results are for all chains combined.

* `mc` : sampler output from a model fit with the `mcmc()` function.
"""
function dic(mc::ModelChains)
    nodekeys = keys(mc.model, :output)

    Dhat = -2.0 * logpdf(mc, mean, nodekeys)
    D = -2.0 * logpdf(mc, nodekeys).value
    p = [mean(D) - Dhat, 0.5 * var(D)]

    ChainSummary(
        [Dhat .+ 2.0 .* p p],
        ["pD", "pV"],
        ["DIC", "Effective Parameters"],
        header(mc),
    )
end

"""
    logpdf(mc::ModelChains, f::Function, nodekeys::Vector{Symbol})

Compute the sum of log-densities at each iteration of MCMC output for stochastic nodes.

Returns a `ModelChains` object of resulting summed log-densities at each MCMC iteration of the supplied chain.

* `mc` : sampler output from a model fit with the `mcmc()`` function.

*  `nodekey/nodekeys` : stochastic model node(s) over which to sum densities (default: all).

* `f` : ??
"""
function logpdf(mc::ModelChains, f::Function, nodekeys::Vector{Symbol})
    m = mc.model

    relistkeys, updatekeys = getsimkeys(mc, nodekeys)
    relistkeys = union(relistkeys, intersect(nodekeys, keys(m, :block)))
    inds = names2inds(mc, relistkeys)
    for key in relistkeys
        isa(m[key], Stochastic{<:GeneralNode}) && throw("not possible with tree objects")
    end

    m[relistkeys] = relist(m, map(i -> f(mc.value[:, i, :]), inds), relistkeys)
    update!(m, updatekeys)
    mapreduce(key -> logpdf(m[key]), +, nodekeys)
end

logpdf(mc::ModelChains, nodekey::Symbol) = logpdf(mc, [nodekey])
"""
    logpdf(mc::ModelChains,
           nodekeys::Vector{Symbol}=keys(mc.model, :stochastic))

"""
function logpdf(mc::ModelChains, nodekeys::Vector{Symbol} = keys(mc.model, :stochastic))
    N = length(mc.range)
    K = size(mc, 3)

    relistkeys, updatekeys = getsimkeys(mc, nodekeys)
    relistkeys = union(relistkeys, intersect(nodekeys, keys(mc.model, :block)))
    inds = names2inds(mc, relistkeys)

    println("MCMC Processing of $N Iterations x $K Chain" * "s"^(K > 1) * "...")

    # set up 1 RemoteChannel & a ProgressMeter for each Chain
    channel = RemoteChannel(() -> Channel{Integer}(1))
    meters = [
        Progress(
            size(mc.value, 1);
            dt = 0.5,
            desc = "Chain $k: ",
            enabled = true,
            offset = k - 1,
            showspeed = true,
        ) for k = 1:K
    ]

    lsts = [
        Any[mc[:, :, [k]], nodekeys, relistkeys, inds, updatekeys, (channel, k)] for
        k = 1:K
    ]

    sims::Vector{Chains} = []
    @sync begin
        # check RemoteChannel for new entries and update the ProgressMeters
        finished_chains = 0
        # loop until all ProgressMeters are finished
        @async while finished_chains < K
            chain = take!(channel)
            chain > 0 ? ProgressMeter.next!(meters[chain]) : finished_chains += 1
        end # while
        @async sims = pmap2(logpdf_modelchains_worker, lsts)
    end # @sync
    println("\n")
    ModelChains(
        cat(sims..., dims = 3),
        mc.model,
        mc.stats,
        mc.stat_names,
        mc.sim_params,
        mc.conv_storage,
    )
end

function logpdf_modelchains_worker(args::Vector)
    mc::ModelChains,
    nodekeys::Vector{Symbol},
    relistkeys::Vector{Symbol},
    inds::Vector{Int},
    updatekeys::Vector{Symbol},
    channel::Tuple{RemoteChannel,Integer} = args
    m = mc.model

    sim = Chains(size(mc, 1), 1, start = first(mc), thin = step(mc), names = ["logpdf"])

    for i = 1:size(mc.value, 1)
        m[relistkeys] = relist(m, mc.value[i, inds, 1], relistkeys)
        update!(m, updatekeys)
        sim.value[i, 1, 1] = mapreduce(key -> logpdf(m[key]), +, nodekeys)
        # send update to RemoteChannel so that the ProgressMeter can be updated in logpdf
        put!(channel[1], channel[2])
    end
    # signal to the logpdf function, that this chain is finished
    put!(channel[1], -1)
    sim
end

predict(mc::ModelChains, nodekey::Symbol) = predict(mc, [nodekey])
"""
    predict(mc::ModelChains,
             nodekeys::Vector{Symbol}=keys(mc.model, :output))

Generate MCMC draws from a posterior predictive distribution.

Returns a `ModelChains` object of draws simulated at each MCMC iteration of the supplied chain. For observed data node ``y``, simulation is from the posterior predictive distribution

``p(\\tilde{y} | y) = \\int p(\\tilde{y} | \\Theta) p(\\Theta | y) d\\Theta,``

where ``\\tilde{y}`` is an unknown observation on the node, ``p(\\tilde{y} | \\Theta)`` is the data likelihood, and ``p(\\Theta | y)`` is the posterior distribution of unobserved parameters ``\\Theta``.

* `mc` : sampler output from a model fit with the `mcmc()`` function.

* `nodekey/nodekeys` : observed Stochastic model node(s) for which to generate draws from the predictive distribution (default: all observed data nodes).
"""
function predict(mc::ModelChains, nodekeys::Vector{Symbol} = keys(mc.model, :output))
    m = mc.model

    outputs = keys(m, :output)
    all(key -> key in outputs, nodekeys) || throw(
        ArgumentError(
            string(
                "nodekeys are not all observed Stochastic nodess : ",
                join(map(string, outputs), ", "),
            ),
        ),
    )

    nodenames = names(m, nodekeys)
    relistkeys, updatekeys = getsimkeys(mc, nodekeys)
    for key in relistkeys
        isa(m[key], Stochastic{<:GeneralNode}) && throw("not possible with tree objects")
    end

    inds = names2inds(mc, relistkeys)

    c = Chains(
        size(mc, 1),
        length(nodenames),
        chains = size(mc, 3),
        start = first(mc),
        thin = step(mc),
        names = nodenames,
    )

    iters, _, chains = size(c.value)
    for k = 1:chains
        for i = 1:iters
            m[relistkeys] = relist(m, mc.value[i, inds, k], relistkeys)
            update!(m, updatekeys)
            f = key -> unlist(m[key], rand(m[key]))
            c.value[i, :, k] = vcat(map(f, nodekeys)...)
        end
    end

    ModelChains(c, m, c.stats, c.stat_names)
end

#################### Auxiliary Functions ####################

function getsimkeys(
    mc::ModelChains,
    nodekeys::Vector{Symbol},
)::Tuple{Vector{Symbol},Vector{Symbol}}

    relistkeys = Symbol[]
    updatekeys = Symbol[]

    m = mc.model
    dag = ModelGraph(m)

    nodekeys = intersect(nodekeys, keys(m, :stochastic))
    blockkeys = keys(m, :block)
    dynamickeys = union(blockkeys, keys(m, :target, blockkeys))
    terminalkeys = union(keys(m, :stochastic), keys(mc, :dependent))

    for v in vertices(dag.graph)
        vkey = dag.keys[v]
        if vkey in dynamickeys
            if any(key -> key in nodekeys, gettargets(dag, v, terminalkeys))
                vkey in terminalkeys ? push!(relistkeys, vkey) : push!(updatekeys, vkey)
            end
        end
    end
    if !isempty(relistkeys)
        append!(updatekeys, nodekeys)
    end

    relistkeys, intersect(keys(m, :dependent), updatekeys)
end

max_stars_repo_name: erathorn/MCPhylo

max_issues_repo_name: erathorn/MCPhylo

max_forks_repo_name: erathorn/MCPhylo


text:
#=
  Eigenvalues
  Eigenvectors
  l1 Norm
  l2 Norm
  loo Norm
  Apply Guassian Elimination
  Apply Guassian Elimination with Pivots
  Reduced-Row-Eschelon Form
  Determine Linear Transformation Matrix
  Rank and Nullity
=#

function eigenvalues(A)
  Lambda = identityMatrix(length(A)[:,1])
  'x' * Lambda
  Lambda -= A
  #TODO:fixed point iteration

function eigenvectors(A)
  Eigenvalues(A)
  Axb(A-   , ,0)

function cond(A)
  #determines condition numebr
  v = eigenvalues(A)
  #returns vector of Eigenvalues
  return abs(loo(v) / loo(1 / v))
  # largest over smallest

function l1(a)
  #accepts vector of n space
  x = 0
  for i=1:length(a)
    x += a[i]
  end

function l2(a)
  #accepts vector of n space
  x = 0
  for i = 1:length(a)
    x += a[i] * a[i]
  end
  return sqrt(x)

function loo(a)
  x = 0
  for i = 1:length(a)
    if abs(a[i]) > x
      x = abs(a[i])
    else
      continue
    end
  end
  return x

function induced_l1(A)
  #accepts matrix
  x = l1(A[:,1])
  for i = 2:length(A)[1,:]
    if l1(A[:,i]) > x
      x = l1(A[:,i])
    else
      continue
    end
  end
  return x

function induced_l2(A)
  #accepts matrix
  return loo(condition(A))
  #TODO: check negative case

function induced_loo(A)
  #accepts matrix
  x = l1(A[1,:])
  for i = 2:length(A)[:,1]
    if l1(A[i,:]) > x
      x = l1(A[i,:])
    else
      continue
    end
  end
  return x

function guass(A)
  # accepts m by n matrix
  B = copy(A)
  x = 0
  if m > n
    x = m
  else
    x = n
  end
  for i = 1:x
    for j = i+1:m
      B[j,:]- ((B[i,j] / B[i,i]) * B[i,:])
    end
  end
  return B

function pivGuass(A) # maybe unessesary
  x = 0
  a = []
  y = 0
  for i=1:(smaller)
    for j=i:m
      x = lcm(A[i,i] , A[j,i])
      a[y] = x
      A[j,:]-
    end
  end

function ltToMatrix(A,B)
  # accepts mxn set a and mxn set LT(a)
  C = [A B]
  D = rref(C)
  return

function rref(A)
  # accepts m by n matrix
  B = copy(A)
  if length(B)[1,:] > length(B)[:,1]
    x = length(B)[:,1]
  else
    x = legnth(B)[1,:]
  end
  for i = 1:x
    B[i,:] /= B[i,i]
    for j = (1 : i-1) & (i+1 : m) #make sure logic

    end
  end
  return B

function nullity(A)
  rank = 0
  nullity = 0
  B = rref(A)
  for i=1:length(A)[:,1]
    #need to fix n length issues
    if B[i,:] = 0
      nullity += 1
    else
      continue
    end
  end
  return nullity

function rank(A)
  #via rank-nullity theorem
  return length(A)[1,:] - nullity(A)

max_stars_repo_name: K-P-S/Julia-Scientific-Computing

max_issues_repo_name: K-P-S/Julia-Scientific-Computing

max_forks_repo_name: K-P-S/Julia-Scientific-Computing


text:
using RentCalculator
using Test

@testset "RentCalculator.jl" begin
    # California example
    params = (800_000.0, # home purchase price (dollars)
                5.0, # time in residence (years)
                3.5, # mortgage rate (percent)
                20.0, # mortgage down payment (percent of purchase price)
                30.0, # mortgage length (years)
                3_600.0, # monthly rent (dollars)
                3_500.0, # rental deposit (dollars)
                0.0, # rent broker's fee (dollars)
                0.5, # rental insurance rate (percent of purchase price)
                4.0, # annual home price appreciation rate (percent)
                2.0, # annual rent rate increase (percent)
                7.0, # annual nominal return on investment (percent)
                2.0, # annual inflation rate (percent)
                1.0, # annual property tax rate (percent)
                33.3, # marginal tax rate (percent)
                24.3, # tax rate on investments (percent)
                100_000.0, # captial gains deduction (dollars)
                3.0, # cost of buying house (percent of purchase price)
                5.0, # cost of selling house (percent of selling price)
                0.5, # annual renovation costs (percent of purchase price)
                0.5, # annual maintenance costs (percent of purchase price)
                0.2, # homeowners insurance (percent of purchase price)
                200.0, # additional monthly utilities compared to renting (dollars)
                400.0, # common fee (dollars)
                20.0, # percent of common fee that is deductible (percent))
    # cost output format: (total, initial costs, recurring costs, opportunity costs, net proceeds)
    @test cost(params, option="rent") == (251_705, 3_500, 225_939, 25_766, -3_500)
    @test cost(params, option="buy") == (186_786, 184_000, 253_479, 82_085, 332_778)
end


text:
# MIT License
#
# Copyright (c) 2020 Joshua E Gentry

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# =================================================================================================
"""
    @enum IFF begin
        IFF_UNKNOWN        = 0
        IFF_BMP            = 1
        IFF_JFIF_JPEG      = 2
        IFF_PNG            = 3
        IFF_TIFF           = 4
        IFF_TIFF_PACKBITS  = 5
        IFF_TIFF_RLE       = 6
        IFF_TIFF_G3        = 7
        IFF_TIFF_G4        = 8
        IFF_TIFF_LZW       = 9
        IFF_TIFF_ZIP       = 10
        IFF_PNM            = 11
        IFF_PS             = 12
        IFF_GIF            = 13
        IFF_JP2            = 14
        IFF_WEBP           = 15
        IFF_LPDF           = 16
        IFF_TIFF_JPEG      = 17
        IFF_DEFAULT        = 18
        IFF_SPIX           = 19
    end

Various constants used by leptonica to specify image types.

__Details:__

| Value             | Description
| :---------------- | :----------
| IFF_UNKNOWN       | Unknown image type.
| IFF_BMP           | BMP image.
| IFF_JFIF_JPEG     | JPEG image.
| IFF_PNG           | PNG image.
| IFF_TIFF          | TIFF image with no compression.
| IFF_TIFF_PACKBITS | TIFF image with pack bits compression.
| IFF_TIFF_RLE      | TIFF image with RLE compression.
| IFF_TIFF_G3       | TIFF image with G3 compression.
| IFF_TIFF_G4       | TIFF image with G4 compression.
| IFF_TIFF_LZW      | TIFF image with LZW compression.
| IFF_TIFF_ZIP      | TIFF image with ZIP compression.
| IFF_PNM           | PNM image.
| IFF_PS            | PostScript file.
| IFF_GIF           | GIF image.
| IFF_JP2           | JP2K image.
| IFF_WEBP          | WEBP image.
| IFF_LPDF          | PDF file.
| IFF_TIFF_JPEG     | TIFF image with JPEG compression.
| IFF_DEFAULT       | Default image type (used in saving).
| IFF_SPIX          | SPIX image.
"""
@enum IFF begin
    IFF_UNKNOWN        = 0
    IFF_BMP            = 1
    IFF_JFIF_JPEG      = 2
    IFF_PNG            = 3
    IFF_TIFF           = 4
    IFF_TIFF_PACKBITS  = 5
    IFF_TIFF_RLE       = 6
    IFF_TIFF_G3        = 7
    IFF_TIFF_G4        = 8
    IFF_TIFF_LZW       = 9
    IFF_TIFF_ZIP       = 10
    IFF_PNM            = 11
    IFF_PS             = 12
    IFF_GIF            = 13
    IFF_JP2            = 14
    IFF_WEBP           = 15
    IFF_LPDF           = 16
    IFF_TIFF_JPEG      = 17
    IFF_DEFAULT        = 18
    IFF_SPIX           = 19
end

# =================================================================================================
"""
    const TIFF_FORMATS = Set([
        IFF_TIFF, IFF_TIFF_RLE, IFF_TIFF_PACKBITS, IFF_TIFF_G3,
        IFF_TIFF_G4, IFF_TIFF_LZW, IFF_TIFF_ZIP, IFF_TIFF_JPEG
    ])

The set of ways an image can be encoded in a TIFF image.
"""
const TIFF_FORMATS = Set([
    IFF_TIFF, IFF_TIFF_RLE, IFF_TIFF_PACKBITS, IFF_TIFF_G3,
    IFF_TIFF_G4, IFF_TIFF_LZW, IFF_TIFF_ZIP, IFF_TIFF_JPEG
])


text:
#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Functions to compute general values related to the orbit.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export @check_orbit
export angvel, angvel_to_a, dargp, draan, period

################################################################################
#                                    Macros
################################################################################

"""
    @check_orbit(a, e)

Verify if the orbit with semi-major axis `a` [m] and eccentricity `e` is valid.
This macro throws an exception if the orbit is not valid.

Return `true` is the orbit is valid, and `false` otherwise.

"""
macro check_orbit(a, e)
    quote
        # Check if the arguments are valid.
        if $(esc(a)) < 0
            false
        elseif !( 0 <= $(esc(e)) < 1 )
            false
        else
            # Compute the orbit perigee and check if it is inside Earth.
            perigee = $(esc(a))*(1-$(esc(e)))

            if perigee < R0
                false
            else
                true
            end
        end
    end
end

################################################################################
#                                  Functions
################################################################################

"""
    angvel(a::Number, e::Number, i::Number, pert::Symbol = :J2)
    angvel(orb::Orbit, pert::Symbol = :J2)

Compute the angular velocity [rad/s] of an object in an orbit with semi-major
axis `a` [m], eccentricity `e`, and inclination `i` [rad], using the
perturbation terms specified by the symbol `pert`. The orbit can also be
specified by `orb` (see [`Orbit`](@ref)).

`pert` can be:

* `:J0`: Consider a Keplerian orbit.
* `:J2`: Consider the perturbation terms up to J2.
* `:J4`: Consider the perturbation terms J2, J4, and J2².

If `pert` is omitted, then it defaults to `:J2`.

"""
@inline function angvel(a::Number, e::Number, i::Number, pert::Symbol = :J2)
    # Unperturbed orbit period.
    n0 = sqrt(m0/Float64(a)^3)

    # Perturbation computed using a Keplerian orbit.
    if pert == :J0
        return n0
    # Perturbation computed using perturbations terms up to J2.
    elseif pert == :J2
        # Auxiliary variables.
        cos_i² = cos(i)^2
        aux   = 1-e^2

        # Semi-lactum rectum.
        p = a*aux

        # Orbit period considering the perturbations (up to J2).
        return n0 + 3R0^2*J2/(4p^2)*n0*(sqrt(aux)*(3cos_i²-1) + (5cos_i²-1))

    # Perturbation computed using perturbations terms J2, J4, and J2².
    elseif pert == :J4

        # Auxiliary variables
        e²     = e^2
        e⁴     = e^4
        sin_i  = sin(i)
        sin_i² = sin_i^2
        sin_i⁴ = sin_i^4
        aux    = (1-e²)
        saux   = sqrt(aux)
        p      = (a/R0)*aux
        p²     = p^2
        p⁴     = p^4

        # Notice that:
        #            .   .
        #   n = n₀ + ω + M₀
        #
        # in which the time-derivatives are computed as in [1, p. 692].

        δω   =  3/  4*n0*J2  /p²*(4 - 5sin_i²) +
                9/384*n0*J2^2/p⁴*(     56e² + (760 -  36e²)*sin_i² - (890 +  45e²)*sin_i⁴) -
               15/128*n0*J4  /p⁴*(64 + 72e² - (248 + 252e²)*sin_i² + (196 + 189e²)*sin_i⁴)

        δM_0 =  3/  4*n0*J2  /p²*saux*(2 - 3sin_i²) +
                3/512*n0*J2^2/p⁴/saux*(320e² - 280e⁴ + (1600 - 1568e² + 328e⁴)*sin_i² + (-2096 + 1072e² +  79e⁴)*sin_i⁴) -
               45/128*n0*J4  /p⁴*saux*e²*(-8 + 40sin_i - 35sin_i²)

        return n0 + δω + δM_0
    else
        throw(ArgumentError("The perturbation parameter $pert is not defined."))
    end
end

@inline function angvel(orb::Orbit, pert::Symbol = :J2)
    # Convert first to Keplerian elements.
    k = convert(KeplerianElements, orb)
    return angvel(get_a(k), get_e(k), get_i(k), pert)
end

"""
    angvel_to_a(n::Number, e::Number, i::Number, pert::Symbol = :J2; μ::Number = m0, max_iter::Int = 20, tol::Number = 1e-10)

Compute the semi-major axis that will provide an angular velocity `n` [rad/s] in
an orbit with eccentricity `e` and inclination `i` [rad], using the perturbation
terms specified by the symbol `pert`.

Notice that the angular velocity `n` is related to the nodal period, *i.e.* the
time between two consecutive passages by the ascending node.

`pert` can be:

* `:J0`: Consider a Keplerian orbit.
* `:J2`: Consider the perturbation terms up to J2.
* `:J4`: Consider the perturbation terms J2, J4, and J2².

If `pert` is omitted, then it defaults to `:J2`.

# Keyword

* `μ`: Standard gravitational parameter for Earth [m^3/s^2].
       (**Default** = `m0`)
* `max_iter`: Maximum number of iterations allowed in the Newton-Raphson
              algorithm. (**Default** = 20)
* `tol`: Tolerance to stop the Newton-Raphson algorithm. (**Default** = 1e-10)

"""
function angvel_to_a(n::Number, e::Number, i::Number, pert::Symbol = :J2;
                     μ::Number = m0, max_iter::Int = 20, tol::Number = 1e-10)

    if pert == :J0

        a = (μ/n^2)^(1/3)
        return a

    elseif pert == :J2
        # Get the semi-major axis [m] that will provide the mean motion `n`
        # using perturbation terms up to J2.
        #
        # This can only be done using a numerical algorithm to solve the
        # following equation for `a`:
        #
        #           -                                                      -
        #          |      3        sqrt(1-e²).(3cos²(i)-1) + (5cos²(i)-1))  |
        #   n = n₀ | 1 + ---. J2 .----------------------------------------- |
        #          |      4                    a^2.(1-e^2)^2                |
        #           -                                                      -
        #
        #         sqrt(μ)
        #   n₀ = ----------
        #         a^(3/2)
        #

        # To improve algorithm stability, we will compute the normalized
        # semi-major axis, i.e. `a/R0`.

        # Auxiliary variables to solve for the semi-major axis.
        sqrt_μ = sqrt(μ/R0^3)
        cos_i  = cos(i)
        K      = 3/4*J2*( sqrt(1-e^2)*(3cos_i^2-1) + (5cos_i^2-1) )/(1-e^2)^2

        # Initial guess using a non-perturbed orbit.
        a = (μ/n^2)^(1/3)/R0

        # Newton-Raphson algorithm.
        #
        # Notice that we will allow, at most, 20 iterations.
        for k = 1:20
            # Auxiliary variables.
            ap3o2  = a^(3/2)
            ap5o2  = ap3o2*a # -> a^(5/2)
            ap7o2  = ap5o2*a # -> a^(7/2)
            ap11o2 = ap7o2*a # -> a^(11/2)

            # Compute the residue.
            res = n - sqrt_μ/ap3o2 - K*sqrt_μ/ap7o2

            # Compute the Jacobian of the function.
            df = +3/2*sqrt_μ/ap5o2 + 7/2*K*sqrt_μ/ap11o2

            # Compute the new estimate.
            a = a - res/df

            (abs(res) < tol) && break
        end

        return a*R0

    elseif pert == :J4

        # Auxiliary variables
        sqrt_μ = sqrt(μ/R0^3)
        sin_i  = sin(i)
        sin_i² = sin_i^2
        sin_i⁴ = sin_i^3
        e²     = e^2
        e⁴     = e^4
        aux    = 1-e²
        saux   = sqrt(aux)

        # Get the semi-major axis using J4 perturbation theory [er].
        #
        # This can only be done using a numerical algorithm to solve the
        # following equation for `a`:
        #
        #            .   .
        #   n = n₀ + ω + M₀ ,
        #
        #         sqrt(μm)
        #   n₀ = ---------- ,
        #          a^(3/2)
        #
        # and the time-derivatives of the argument of perigee and the initial
        # mean anomaly is compute considering the terms J2, J4, and J2².

        # To improve algorithm stability, we will compute the normalized
        # semi-major axis, i.e. `a/R0`.

        K₁ = 3/4*J2/aux^2*( (4 - 5sin_i²) + sqrt(1-e²)*(2 - 3sin_i²) )

        K₂ = 3/512*J2^2/aux^4*(224e² + (3040 - 144e²)*sin_i² - (3560 + 180e²)*sin_i⁴ +
                               1/saux*((         320e² - 280e⁴) +
                                       ( 1600 - 1568e² + 328e⁴)*sin_i² +
                                       (-2096 + 1072e² +  79e⁴)*sin_i⁴))

        K₃ = 15/128*J4/aux^4*(64 + 72e² - (248 + 252e²)*sin_i² + (196 + 189e²)*sin_i⁴ +
                              e²/saux*(-8 + 40sin_i - 35sin_i²))

        # Initial guess using a non-perturbed orbit.
        a = (μ/n^2)^(1/3)/R0

        # Newton-Raphson algorithm.
        #
        # Notice that we will allow, at most, 20 iterations.
        for k = 1:20
            # Auxiliary variables.
            ap3o2  = a^(3/2)
            ap5o2  =  ap3o2*a  # -> a^(5/2)
            ap7o2  =  ap5o2*a  # -> a^(7/2)
            ap9o2  =  ap7o2*a  # -> a^(9/2)
            ap11o2 =  ap9o2*a  # -> a^(11/2)
            ap13o2 = ap11o2*a  # -> a^(13/2)

            # Compute the residue.
            res = n - sqrt_μ*(1/ap3o2 + K₁/ap7o2 + (K₂ + K₃)/ap11o2)

            # Compute the Jacobian of the function.
            df = sqrt_μ/2*( 3ap5o2 + 7K₁/ap9o2 + 11*(K₂ + K₃)/ap13o2)

            # Compute the new estimate.
            a = a - res/df

            (abs(res) < tol) && break
        end

        return a*R0
    else
        throw(ArgumentError("The perturbation parameter $pert is not defined."))
    end
end

"""
    dargp(a::Number, e::Number, i::Number, pert::Symbol = :J2)
    dargp(orb::Orbit, pert::Symbol = :J2)

Compute the time-derivative of the argument of perigee [rad/s] of an orbit with
semi-major axis `a` [m], eccentricity `e`, and inclination `i` [rad], using the
perturbation terms specified by the symbol `pert`. The orbit can also be
specified by `orb` (see [`Orbit`](@ref)).

`pert` can be:

* `:J0`: Consider a Keplerian orbit.
* `:J2`: Consider the perturbation terms up to J2.
* `:J4`: Consider the perturbation terms J2, J4, and J2².

If `pert` is omitted, then it defaults to `:J2`.

"""
@inline function dargp(a::Number, e::Number, i::Number, pert::Symbol = :J2)
    # Perturbation computed using a Keplerian orbit.
    if pert == :J0
        return 0.0
    # Perturbation computed using perturbations terms up to J2.
    elseif pert == :J2
        # Semi-lactum rectum.
        p = a*(1-e^2)

        # Unperturbed orbit period.
        n0 = angvel(a, e, i, :J0)

        # Perturbation of the argument of perigee.
        return 3R0^2*J2/(4p^2)*n0*(5cos(i)^2-1)

    elseif pert == :J4
        # Auxiliary variables
        e²     = e^2
        e⁴     = e^4
        sin_i  = sin(i)
        sin_i² = sin_i^2
        sin_i⁴ = sin_i^4
        aux    = (1-e²)
        saux   = sqrt(aux)
        p      = (a/R0)*aux
        p²     = p^2
        p⁴     = p^4

        # Unperturbed orbit period.
        n0 = angvel(a, e, i, :J0)

        # Perturbation of the argument of perigee.
        δω =  3/  4*n0*J2  /p²*(4 - 5sin_i²) +
              9/384*n0*J2^2/p⁴*(     56e² + (760 -  36e²)*sin_i² - (890 +  45*e²)*sin_i⁴) -
             15/128*n0*J4  /p⁴*(64 + 72e² - (248 + 252e²)*sin_i² + (196 + 189*e²)*sin_i⁴)

        return δω
    else
        throw(ArgumentError("The perturbation parameter $pert is not defined."))
    end
end

function dargp(orb::Orbit, pert::Symbol = :J2)
    # Convert first to Keplerian elements.
    k = convert(KeplerianElements, orb)
    return dargp(get_a(k), get_e(k), get_i(k), pert)
end

"""
    draan(a::Number, e::Number, i::Number, pert::Symbol = :J2)
    draan(orb::Orbit, pert::Symbol = :J2)

Compute the time-derivative of the right ascension of the ascending node [rad/s]
of an orbit with semi-major axis `a` [m], eccentricity `e`, and inclination `i`
[rad], using the perturbation terms specified by the symbol `pert`. The orbit
can also be specified by `orb` (see [`Orbit`](@ref)).

`pert` can be:

* `:J0`: Consider a Keplerian orbit.
* `:J2`: Consider the perturbation terms up to J2.
* `:J4`: Consider the perturbation terms J2, J4, and J2².

If `pert` is omitted, then it defaults to `:J2`.

"""
@inline function draan(a::Number, e::Number, i::Number, pert::Symbol = :J2)
    # Perturbation computed using a Keplerian orbit.
    if pert == :J0
        return 0.0
    # Perturbation computed using perturbations terms up to J2.
    elseif pert == :J2
        # Semi-lactum rectum.
        p = a*(1-e^2)

        # Unperturbed orbit period.
        n0 = angvel(a, e, i, :J0)

        # Perturbation of the right ascension of the ascending node.
        return -3/2*R0^2/(p^2)*n0*J2*cos(i)

    # Perturbation computed using perturbation terms J2, J4, and J2².
    elseif pert == :J4
        # Auxiliary variables
        sin_i, cos_i = sincos(i)
        sin_i²       = sin_i^2
        e²           = e^2
        p            = (a/R0)*(1-e²)
        p²           = p^2
        p⁴           = p^4

        # Unperturbed orbit period.
        n0 = angvel(a, e, i, :J0)

        # Perturbation of the right ascension of the ascending node.
        δΩ = -3/2 *n0*J2  /p²*cos_i +
              3/32*n0*J2^2/p⁴*cos_i*(12 -  4e² - (80 +  5e²)*sin_i²) +
             15/32*n0*J4  /p⁴*cos_i*( 8 + 12e² - (14 + 21e²)*sin_i²)

        return δΩ
    else
        throw(ArgumentError("The perturbation parameter $pert is not defined."))
    end
end

function draan(orb::Orbit, pert::Symbol = :J2)
    # Convert first to Keplerian elements.
    k = convert(KeplerianElements, orb)
    return draan(get_a(k), get_e(k), get_i(k), pert)
end

"""
    period(a::Number, e::Number, i::Number, pert::Symbol = :J2)

Compute the period [s] of an object in an orbit with semi-major axis `a` [m],
eccentricity `e`, and inclination `i` [rad], using the perturbation terms
specified by the symbol `pert`. The orbit can also be specified by `orb` (see
[`Orbit`](@ref)).

pert` can be:

* `:J0`: Consider a Keplerian orbit.
* `:J2`: Consider the perturbation terms up to J2.
* `:J4`: Consider the perturbation terms J2, J4, and J2².

If `pert` is omitted, then it defaults to `:J2`.

"""
@inline function period(a::Number, e::Number, i::Number, pert::Symbol = :J2)
    n = angvel(a, e, i, pert)
    return 2π/n
end

function period(orb::Orbit, pert::Symbol = :J2)
    # Convert first to Keplerian elements.
    k = convert(KeplerianElements, orb)
    return period(get_a(k), get_e(k), get_i(k), pert)
end


text:
module StatisticalNonlocality

include("transition_matrices.jl")

include("spectral.jl")

# Utilities
include("utils.jl")

end # module

max_stars_repo_name: sandreza/StatisticalNonlocality

max_issues_repo_name: sandreza/StatisticalNonlocality

max_forks_repo_name: sandreza/StatisticalNonlocality


text:
@enum TerminationReason begin
    TERMINATION_REASON_UNSPECIFIED
    TERMINATION_REASON_OPTIMAL
    TERMINATION_REASON_ITERATION_LIMIT
end

struct BundleMethodParams
    eps_optimal::Float64
    iteration_limit::Int64
    contraction_factor::Float64
    verbose::Bool
    printing_frequency::Int64
    full_memory::Bool
end

struct SubgradientMethodParams
    eps_optimal::Float64
    iteration_limit::Int64
    verbose::Bool
    printing_frequency::Int64
end

struct Cut
    constant_term::Float64
    gradient::Vector{Float64}
end

struct IterationInformation
    iteration::Int64
    is_null::Bool
    objective::Float64
    approximation_gap::Float64
    residual::Float64
    step_size::Float64
end

struct BundleMethodOutput
    solution::Vector{Float64}
    iteration_stats::Vector{IterationInformation}
    termination_reason::TerminationReason
end

mutable struct SolverState
    current_objective::Float64
    current_gap::Float64
    current_residual::Float64
    latest_is_null::Bool
    current_solution::Vector{Float64}
    cuts::Vector{Cut}
end

function create_cut(point::Vector{Float64}, objective_function, subgradient_section)
    subgradient = subgradient_section(point)
    return Cut(objective_function(point) - dot(point, subgradient), subgradient)
end

function evaluate_cut(cut::Cut, point::Vector{Float64})
    return cut.constant_term + dot(cut.gradient, point)
end

function evaluate_model(cuts::Vector{Cut}, point::Vector{Float64})
    return maximum(map(cut -> evaluate_cut(cut, point), cuts))
end

function aggregate_cuts(cuts::Vector{Cut}, interpolation_coefficients::Vector{Float64})
    constant_term = cuts[1].constant_term * interpolation_coefficients[1]
    gradient = cuts[1].gradient * interpolation_coefficients[1]
    for i = 2:length(cuts)
        constant_term += cuts[i].constant_term * interpolation_coefficients[i]
        gradient += cuts[i].gradient * interpolation_coefficients[i]
    end
    return [Cut(constant_term, gradient)]
end

function check_termination(params::BundleMethodParams, iteration_info::IterationInformation)
    if params.verbose && mod(iteration_info.iteration, params.printing_frequency)
        iteration_log(params, iteration_log)
    end
    if iteration_info.iteration >= params.iteration_limit
        return TERMINATION_REASON_ITERATION_LIMIT
    end
    if iteration_info.approximation_gap <= eps_optimal
        return TERMINATION_REASON_OPTIMAL
    end
end

function iteration_log(
    params::SubgradientMethodParams,
    iteration_info::IterationInformation,
)
    if params.verbose && mod(iteration_info.iteration, params.printing_frequency) == 0
        Printf.@printf(
            " %5d obj=%12g residual=%12g step_size=%12g\n",
            iteration_info.iteration,
            iteration_info.objective,
            iteration_info.residual,
            iteration_info.step_size,
        )
    end
end

function iteration_log(params::BundleMethodParams, iteration_info::IterationInformation)
    if params.verbose && mod(iteration_info.iteration, params.printing_frequency) == 0
        Printf.@printf(
            " %s %5d obj=%12g gap=%12g residual=%12g step_size=%12g\n",
            iteration_info.is_null ? "n" : "d",
            iteration_info.iteration,
            iteration_info.objective,
            iteration_info.approximation_gap,
            iteration_info.residual,
            iteration_info.step_size,
        )
    end
end

function solve_model(cuts::Vector{Cut}, point::Vector{Float64}, stepsize::Float64)
    if length(cuts) == 1
        return point - cuts[1].gradient * stepsize, [1.0]
    elseif length(cuts) == 2

        candidate = point - cuts[1].gradient * stepsize
        if evaluate_cut(cuts[1], candidate) >= evaluate_cut(cuts[2], candidate)
            return candidate, [1.0, 0.0]
        end

        candidate = point - cuts[2].gradient * stepsize
        if evaluate_cut(cuts[2], candidate) >= evaluate_cut(cuts[1], candidate)
            return candidate, [0.0, 1.0]
        end

        grad_diff = cuts[1].gradient - cuts[2].gradient
        interpolation =
            1 / (stepsize * dot(grad_diff, grad_diff)) *
            (cuts[1].constant_term - cuts[2].constant_term + dot(grad_diff, candidate))
        return candidate - grad_diff * interpolation * stepsize,
        [interpolation, 1.0 - interpolation]
    else
        # This block of code is only called with the full-memory method.
        x = Convex.Variable(length(cuts[1].gradient))
        t = Convex.Variable(1)

        problem = minimize(t + (1 / (2 * stepsize)) * Convex.sumsquares(x - point))
        for cut in cuts
            problem.constraints += t >= cut.gradient' * x + cut.constant_term
        end
        Convex.solve!(problem, () -> SCS.Optimizer(verbose = false), verbose = false)
        # Convex.solve!(problem, Gurobi.Optimizer)
        # if problem.status != 1
        # error("The inner problem couldn't be solved to optimality.") # This should never happen.
        # end
        return Convex.evaluate(x), []
    end
end

function take_step(
    objective_function,
    subgradient_map,
    params,
    step_size,
    solver_state,
    iteration_stats,
)

    (new_point, interpolation_coefficients) =
        solve_model(solver_state.cuts, solver_state.current_solution, step_size)

    new_objective = objective_function(new_point)
    solver_state.current_residual = norm(new_point - solver_state.current_solution)
    solver_state.latest_is_null = true

    model_value = evaluate_model(solver_state.cuts, new_point)
    solver_state.current_gap = new_objective - model_value

    if params.contraction_factor * (solver_state.current_objective - model_value) <=
       solver_state.current_objective - new_objective
        solver_state.latest_is_null = false
        solver_state.current_solution = new_point
        solver_state.current_objective = new_objective
    end

    if !params.full_memory
        # This combines all the previous cuts into one, while mainteining the constrains
        # for the theory to hold.
        solver_state.cuts = aggregate_cuts(solver_state.cuts, interpolation_coefficients)
    end
    push!(solver_state.cuts, create_cut(new_point, objective_function, subgradient_map))
end

"""
Solves the problem with the subgradient method.
"""
function solve(
    objective_function::Function,
    subgradient_map::Function,
    params::SubgradientMethodParams,
    stepsize_policy::Function,
    initial_point::Vector{Float64},
)
    iteration_stats =
        [IterationInformation(0, false, objective_function(initial_point), Inf, Inf, 0.0)]
    solver_state = SolverState(
        objective_function(initial_point),
        Inf,
        Inf,
        false,
        initial_point,
        [create_cut(initial_point, objective_function, subgradient_map)],
    )
    best_solution = initial_point
    best_obj = objective_function(initial_point)
    for i = 1:params.iteration_limit
        step_size = stepsize_policy(
            solver_state.current_objective,
            solver_state.current_solution,
            i,
        )
        new_point = (
            solver_state.current_solution -
            stepsize_policy(
                solver_state.current_objective,
                solver_state.current_solution,
                i,
            ) * subgradient_map(solver_state.current_solution)
        )
        solver_state.current_residual = norm(new_point - solver_state.current_solution)

        solver_state.current_solution = new_point
        solver_state.current_objective = objective_function(new_point)

        if solver_state.current_objective < best_obj
            best_solution = solver_state.current_solution
        end

        push!(
            iteration_stats,
            IterationInformation(
                i,
                solver_state.latest_is_null,
                solver_state.current_objective,
                solver_state.current_gap,
                solver_state.current_residual,
                step_size,
            ),
        )
        iteration_log(params, last(iteration_stats))
    end

    return BundleMethodOutput(
        best_solution,
        iteration_stats,
        TERMINATION_REASON_ITERATION_LIMIT,
    )
end

"""
Solves the problem with the bundle method.
"""
function solve(
    objective_function::Function,
    subgradient_map::Function,
    params::BundleMethodParams,
    stepsize_policy::Function,
    initial_point::Vector{Float64},
)

    iteration_stats =
        [IterationInformation(0, false, objective_function(initial_point), Inf, Inf, 0.0)]
    solver_state = SolverState(
        objective_function(initial_point),
        Inf,
        Inf,
        false,
        initial_point,
        [create_cut(initial_point, objective_function, subgradient_map)],
    )
    for i = 1:params.iteration_limit
        step_size = stepsize_policy(
            solver_state.current_objective,
            solver_state.current_solution,
            i,
        )
        take_step(
            objective_function,
            subgradient_map,
            params,
            step_size,
            solver_state,
            iteration_stats,
        )
        push!(
            iteration_stats,
            IterationInformation(
                i,
                solver_state.latest_is_null,
                solver_state.current_objective,
                solver_state.current_gap,
                solver_state.current_residual,
                1 / step_size,
            ),
        )
        iteration_log(params, last(iteration_stats))
    end

    return BundleMethodOutput(
        solver_state.current_solution,
        iteration_stats,
        TERMINATION_REASON_ITERATION_LIMIT,
    )
end

function combine_parallel_models(objective_function, subgradient_map, solver_states)
    min_fun, idx = findmin([state.current_objective for state in solver_states])
    new_solution = solver_states[idx].current_solution
    for state in solver_states
        if !state.latest_is_null && min_fun < state.current_objective
            state.current_residual = norm(state.current_solution - new_solution)
            state.current_solution = new_solution
            state.current_objective = objective_function(new_solution)
            state.cuts = [create_cut(new_solution, objective_function, subgradient_map)]
        end
    end
    return idx
end

function solve_adaptive(
    objective_function,
    subgradient_map,
    params,
    step_sizes,
    initial_point,
)
    iteration_stats =
        [IterationInformation(0, false, objective_function(initial_point), Inf, Inf, 0.0)]
    iteration_info_agents = [
        [IterationInformation(0, false, objective_function(initial_point), Inf, Inf, 0.0)] for _ in step_sizes
    ]
    solver_states = [
        SolverState(
            objective_function(initial_point),
            Inf,
            Inf,
            false,
            initial_point,
            [create_cut(initial_point, objective_function, subgradient_map)],
        ) for _ in step_sizes
    ]
    idx = 0
    for i = 1:params.iteration_limit
        for j = 1:length(step_sizes)
            take_step(
                objective_function,
                subgradient_map,
                params,
                step_sizes[j],
                solver_states[j],
                iteration_stats,
            )
        end
        idx = combine_parallel_models(objective_function, subgradient_map, solver_states)
        for idx = 1:length(step_sizes)
            push!(
                iteration_info_agents[idx],
                IterationInformation(
                    i,
                    solver_states[idx].latest_is_null,
                    solver_states[idx].current_objective,
                    solver_states[idx].current_gap,
                    solver_states[idx].current_residual,
                    step_sizes[idx],
                ),
            )
        end

        push!(
            iteration_stats,
            IterationInformation(
                i,
                solver_states[idx].latest_is_null,
                solver_states[idx].current_objective,
                solver_states[idx].current_gap,
                solver_states[idx].current_residual,
                1 / step_sizes[idx],
            ),
        )
        iteration_log(params, last(iteration_stats))
    end

    return BundleMethodOutput(
        solver_states[idx].current_solution,
        iteration_stats,
        TERMINATION_REASON_ITERATION_LIMIT,
    ),
    iteration_info_agents
end


text:
# Copyright (c) 2019 Matthew Help, Mathieu Tanneau and contributors              #src
#                                                                                #src
# Permission is hereby granted, free of charge, to any person obtaining a copy   #src
# of this software and associated documentation files (the "Software"), to deal  #src
# in the Software without restriction, including without limitation the rights   #src
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell      #src
# copies of the Software, and to permit persons to whom the Software is          #src
# furnished to do so, subject to the following conditions:                       #src
#                                                                                #src
# The above copyright notice and this permission notice shall be included in all #src
# copies or substantial portions of the Software.                                #src
#                                                                                #src
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR     #src
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,       #src
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE    #src
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER         #src
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,  #src
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE  #src
# SOFTWARE.                                                                      #src

# # Geographical Clustering

# **Originally Contributed by**: Matthew Helm ([with help from Mathieu Tanneau on Julia Discourse](https://discourse.julialang.org/t/which-jump-jl-solver-for-this-problem/43350/17?u=mthelm85))

# The goal of this exercise is to cluster $n$ cities into $k$ groups, minimizing
# the total pairwise distance between cities *and* ensuring that the variance in
# the total populations of each group is relatively small.

# This tutorial uses the following packages:
using JuMP
import DataFrames
import GLPK
import LinearAlgebra

# For this example, we'll use the 20 most populous cities in the United States.

cities = DataFrames.DataFrame(
    city = ["New York, NY", "Los Angeles, CA", "Chicago, IL", "Houston, TX", "Philadelphia, PA", "Phoenix, AZ", "San Antonio, TX", "San Diego, CA", "Dallas, TX", "San Jose, CA", "Austin, TX", "Indianapolis, IN", "Jacksonville, FL", "San Francisco, CA", "Columbus, OH", "Charlotte, NC", "Fort Worth, TX", "Detroit, MI", "El Paso, TX", "Memphis, TN"],
    population = [8.405,3.884,2.718,2.195,1.553,1.513,1.409,1.355,1.257,0.998,0.885,0.843,0.842,0.837,0.822,0.792,0.792,0.688,0.674,0.653],
    lat = [40.7127,34.0522,41.8781,29.7604,39.9525,33.4483,29.4241,32.7157,32.7766,37.3382,30.2671,39.7684,30.3321,37.7749,39.9611,35.2270,32.7554,42.3314,31.7775,35.1495],
    lon = [-74.0059,-118.2436,-87.6297,-95.3698,-75.1652,-112.0740,-98.4936,-117.1610,-96.7969,-121.8863,-97.7430,-86.1580,-81.6556,-122.4194,-82.9987,-80.8431,-97.3307,-83.0457,-106.4424,-90.0489],
)

# ### Model Specifics

# We will cluster these 20 cities into 3 different groups and we will assume
# that the ideal or target population $P$ for a group is simply the total
# population of the 20 cities divided by 3:

n = size(cities,1)
k = 3
P = sum(cities.population) / k

# ### Obtaining the distances between each city

# Let's compute the pairwise Haversine distance between each of the cities in
# our data set and store the result in a variable we'll call `dm`:

"""
    haversine(lat1, long1, lat2, long2, r = 6372.8)

Compute the haversine distance between two points on a sphere of radius `r`,
where the points are given by the latitude/longitude pairs `lat1/long1` and
`lat2/long2` (in degrees).
"""
function haversine(lat1, long1, lat2, long2, r = 6372.8)
    lat1, long1 = deg2rad(lat1), deg2rad(long1)
    lat2, long2 = deg2rad(lat2), deg2rad(long2)
    hav(a, b) = sin((b - a) / 2)^2
    inner_term = hav(lat1, lat2) + cos(lat1) * cos(lat2) * hav(long1, long2)
    d = 2 * r * asin(sqrt(inner_term))
    ## Round distance to nearest kilometer.
    return round(Int, d)
end

# Our distance matrix is symmetric so we'll convert it to a `LowerTriangular`
# matrix so that we can better interpret the objective value of our model:

dm = LinearAlgebra.LowerTriangular([
    haversine(cities.lat[i], cities.lon[i], cities.lat[j], cities.lon[j])
    for i = 1:n, j = 1:n
])

# ### Build the model

# Now that we have the basics taken  care of, we can set up our model, create
# decision variables, add constraints, and then solve.

# First, we'll set up a model that leverages the Cbc solver. Next, we'll set up
# a binary variable $x_{i,k}$ that takes the value $1$ if city $i$ is in group
# $k$ and $0$ otherwise. Each city must be in a group, so we'll add the
# constraint $\sum_kx_{i,k} = 1$ for every $i$.

model = Model(GLPK.Optimizer)
set_silent(model)
#-

@variable(model, x[1:n, 1:k], Bin)

#-

@constraint(model, [i = 1:n], sum(x[i, :]) == 1)

# To reduce symmetry, we fix the first city to belong to the first group.

fix(x[1, 1], 1; force = true)

# The total population of a group $k$ is $Q_k = \sum_ix_{i,k}q_i$ where $q_i$ is
# simply the $i$th value from the `population` column in our `cities` DataFrame.
# Let's add constraints so that $\alpha \leq (Q_k - P) \leq \beta$. We'll set
# $\alpha$ equal to $-3$ million and $\beta$ equal to $3$. By adjusting
# these thresholds you'll find that there is a tradeoff between having
# relatively even populations between groups and having geographically close
# cities within each group. In other words, the larger the absolute values of
# $\alpha$ and $\beta$, the closer together the cities in a group will be but
# the variance between the group populations will be higher.

@variable(model, -3 <= population_diff[1:k] <= 3)
@constraint(model, population_diff .== x' * cities.population .- P)

# Now we need to add one last binary variable $z_{i,j}$ to our model that we'll
# use to compute the total distance between the cities in our groups, defined as
# $\sum_{i,j}d_{i,j}z_{i,j}$. Variable $z_{i,j}$ will equal $1$ if cities $i$
# and $j$ are in the same group, and $0$ if they are not in the same group.

# To ensure that $z_{i,j} = 1$ if and only if cities $i$ and $j$ are in the same
# group, we add the constraints $z_{i,j} \geq x_{i,k} + x_{j,k} - 1$ for every
# pair $i,j$ and every $k$:

@variable(model, z[i = 1:n, j = 1:i], Bin)

#-

for k in 1:k, i in 1:n, j in 1:i
    @constraint(model, z[i, j] >= x[i, k] + x[j, k] - 1)
end

# We can now add an objective to our model which will simply be to minimize the
# dot product of $z$ and our distance matrix, `dm`.

@objective(model, Min, sum(dm[i, j] * z[i, j] for i = 1:n, j = 1:i))

# We can then call `optimize!` and review the results.

optimize!(model)

# ### Reviewing the Results

# Now that we have results, we can add a column to our `cities` DataFrame for
# the group and then loop through our $x$ variable to assign each city to its
# group. Once we have that, we can look at the total population for each group
# and also look at the cities in each group to verify that they are grouped by
# geographic proximity.

cities.group = zeros(n)

for i = 1:n, j = 1:k
    if round(Int, value(x[i, j])) == 1
        cities.group[i] = j
    end
end

for group in DataFrames.groupby(cities, :group)
    @show group
    println("")
    @show sum(group.population)
    println("")
end

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode