text:
function mean_hausdorff(u, v)
    "
    This function extracts the edges of two images and then finds the average Hausdorff distance
    Both arrays u and v are required to be the same size and they must be binary or boolean images.
    "
    d = Euclidean()

    # Find the coordinates of the pixels that are edges in the images
    edges_1 = find_edges(u)
    edges_2 = find_edges(v)

    # Turn the array of CartesianIndices into an array of integers corresponding to the CartesianIndices
    edges_1 = (transpose ∘ reshape)(reinterpret(Int, edges_1), 2, :)
    edges_2 = (transpose ∘ reshape)(reinterpret(Int, edges_2), 2, :)

    min_euc_list_u = []
    min_euc_list_v = []

    # Loop through every edge point on `edge_1` and find its corresponding closest point to `edge_2`
    for i in 1:size(edges_1, 1)
        euc_list_1 = []
        for j in 1:size(edges_2, 1)
            euc = evaluate(d, edges_1[i, :], edges_2[j, :])
            append!(euc_list_1 , euc)
        end
        append!(min_euc_list_u, minimum(euc_list_1))
    end

    # Loop through every edge point on `edge_2` and find its corresponding closest point to `edge_1`
    for i in 1:size(edges_2, 1)
        euc_list_1 = []
        for j in 1:size(edges_1, 1)
            euc = evaluate(d, edges_2[i, :], edges_1[j, :])
            append!(euc_list_1 , euc)
        end
        append!(min_euc_list_v, minimum(euc_list_1))
    end

    # Take the average of each of these points to return the average Hausdorff distance
    return mean([mean(min_euc_list_u), mean(min_euc_list_v)])
end

function mean_dice(score, target)
    smooth = 1e-5
    intersect = sum(score .* target)
    y_sum = sum(target .* target)
    z_sum = sum(score .* score)
    metric = (2 * intersect + smooth) / (z_sum + y_sum + smooth)
    return metric
end


text:
getdayfields(days::Int) = (d->d<0 ? Symbol("DM",abs(d)) : Symbol("D",d)).(collect((-1*days):days))

function formeventwindows(event::DataFrame, days::Int;
  daysextended::Int = days,
  refresheventcrsp::Bool = true,
  Fret::Symbol = error("Fret is required"),
  validpermnos::Vector{Int} = unique(event[!, :permno]),
  Fdate::Symbol = error("Fdate is required"),
  retbound::NFloat64 = nothing,
  prefix::String = error("prefix is required"),
  daystostaleprice::NInt=nothing,
  yearrange::UnitRange = error("yearrange (for crsp) is required")
  )

  days>0 || error("days must be greater than 0")
  Nevent::Int = size(event,1)
  event.eid = collect(1:Nevent)

  #load the crsp df
  local crsp::DataFrame = prepcrsp(
    crspcolumns = [:date, :permno, Fret, :price, :divamt, :distcd, :numtrd, :vol],
    crsptype = :dailyexpanded,
    refreshcrsp=refresheventcrsp,
    validpermnos = validpermnos,
    prefix = "$(prefix)event",
    yearrange=yearrange) do df::DataFrame
    computenetreturns!(df)
    computevolumes!(df)
    sort!(df,[:permno,:date])
    #df.Tprice = Vector{MFloat64}(missing, size(df,1))
    #(!isnothing(daystostaleprice)) && lagprice!(df)
  end

  #println("namescrsp", names(crsp))
  #if we use the crsp event file, use the additional data source
  if DIST_TYPE[] ∈ [:primaryshort, :primaryfull]
    select!(crsp, Not([:divamt, :distcd]))
  elseif DIST_TYPE[] ∈ [:oldshort, :oldfull]
    rename!(crsp, [:divamt=>:divamt2, :distcd=>:distcd2])
  else
    @assert false #shouldn't happen due to previous checks
  end

  #make an index of possible dates
  crspdates::Vector{Date} = sort(unique(crsp.date))
  Ndates::Int = length(crspdates)

  #the following interlude only applies to the distributions dataframe
  if :exdate ∈ names(event)
    event.days2lastdiv = Vector{MInt}(missing, size(event,1))

    #looks up the date against crspdates. If it finds the date,
    #returns the index, otherwise returns missing
    @inline function crspdateormissing(d::Date)::MInt
      idx::Int = searchsortedfirst(crspdates, d)
      ((idx > Ndates) || (crspdates[idx] ≠ d)) && return missing
      return idx
    end

    #get the difference between the dividend dates in terms of trading days
    @mpar for r::DataFrameRow ∈ eachrow(event)
      ismissing(r.Lexdate) && continue

      Lexidx::MInt = crspdateormissing(r.Lexdate)
      ismissing(Lexidx) && continue

      exidx::MInt = crspdateormissing(r.exdate)
      r.days2lastdiv = exidx - Lexidx
    end
  end

  #winsorize to an fixed vlaue if desired
  deleterows!(crsp, (ismissing).(crsp[!,Fret]))
  if !isnothing(retbound)
    (retbound > 0.0) || error("retbound must be > 0")
    crsp.winsorized = (f::MFloat64->ismissing(f) ? missing : abs(f) > retbound).(crsp[!, Fret])

    crsp[!, Fret] .= (f::Float64->(max(min(f,retbound), -retbound))).(crsp[!, Fret])
  end

  sort!(event, [:permno, Fdate])

  #assume a fixed set of dates for each event
  dateranges::Dict =Dict{Date,NTuple{2, Date}}(
    crspdates[i]=> (crspdates[i-days], crspdates[i+days])
    for i ∈ (1+days):(length(crspdates)-days))

  #get the begin and end date of each range
  filter!(r->haskey(dateranges, r[Fdate]), event)
  event.begindate = (d->dateranges[d][1]).(event[!,Fdate])
  event.enddate = (d->dateranges[d][2]).(event[!,Fdate])

  #now check for overlapping dates
  event.oneex=trues(size(event,1))
  sevents::GroupedDataFrame = groupby(event, :permno)
  for sevent ∈ groupby(event, :permno)
    Nsevent::Int = size(sevent,1)
    for (i,r) ∈ enumerate(eachrow(sevent))
      t₀::Date = r[Fdate]
      (i > 1) &&  (r[:oneex] *= t₀>sevent[i-1,:enddate])
      (i < Nsevent) && (r[:oneex] *= t₀<sevent[i+1,:begindate])
    end
  end

  (:exyear ∈ names(event)) && event[event.exyear .== 2015, :] |> CSV.write(
    "output\\oneex-$(REPLICATION_TYPE[])-$(DIST_TYPE[]).csv")

  #now build the panel
  #ensure we capture the declaration day with the dividends
  lowerdays::Int = isnothing(daystostaleprice) ? days : max(daysextended, daystostaleprice)
  upperdays::Int = daysextended
  eventXday::Vector{NTuple{2,Int}} = vec(collect(Iterators.product(event.eid, -lowerdays:upperdays)))
  newevent::DataFrame = DataFrame(eid= (t->t[1]).(eventXday), day = (t->t[2]).(eventXday))
  event = join(event, newevent, on=:eid)

  dateidx::Dict = Dict(d=>i for (i,d) ∈ enumerate(crspdates))
  Ncrspdates::Int = length(crspdates)
  filter!(event) do r::DataFrameRow #only keep data within ranges
    idx = dateidx[r[Fdate]] + r.day
    return (idx ≥ 1) && (idx ≤ Ncrspdates)
  end
  event.date = (r::DataFrameRow->crspdates[dateidx[r[Fdate]] + r.day]).(eachrow(event))

  Nevent = size(event,1)
  event = join(event, crsp, on=[:permno, :date], kind=:left)
  @assert Nevent == size(event,1) #make sure we picked the correct keys

  #qc on the returns and net returns
  (s->finiteormissing!(event, s)).([Fret, :price])

  println("################# eventsize: $(size(event)) crspsize: $(size(crsp))")
  #println(describe(crsp))
  #println(describe(event))
  sort!(event, [:eid, :day])

  return event
end

function formevent(;refreshevent::Bool = true,
    eventpath::String = EVENT_PATH,
    refreshseocrsp::Bool=true,
    refreshdistcrsp::Bool=true,
    eventname::String = EVENT_NAME,
    injlsstream::Function = IN_JLS_STREAM,
    outjlsstream::Function = OUT_JLS_STREAM,
    seodays::Int = SEO_DAYS,
    seoyearrange::UnitRange = SEO_YEAR_RANGE,
    distdays::Int = DIST_DAYS,
    distdaysextended::Int = DIST_DAYS_EXTENDED,
    distyearrange::UnitRange = DIST_YEAR_RANGE
    )::NTuple{2,DataFrame}

  local bb::DataFrame, dist::DataFrame
  eventname = "$eventname-$(REPLICATION_TYPE[])-$(DIST_TYPE[])"

  if   refreshevent
    bb = readbbseo()
    #bb |> CSV.write("output\\bbevent-$(REPLICATION_TYPE[])-$(DIST_TYPE[]).csv")

    bb = formeventwindows(bb,seodays,
      refresheventcrsp=refreshseocrsp,
      validpermnos=unique(bb.permno),
      Fret = SEO_RET,
      Fdate=:date,
      prefix="seo",
      yearrange=seoyearrange
      )

    if DIST_TYPE[] ∈ [:primaryshort, :primaryfull]
      dist = primaryreaddist(prefix="", refreshdist=refreshdistcrsp)
    elseif DIST_TYPE[] ∈ [:oldshort, :oldfull]
      dist = oldreaddist()
    else
      @assert false #shouldn't happen due to previous checks
    end

    dist = preprocessdist(dist)

    (DIST_TYPE[] == :primaryshort || DIST_TYPE[]==:oldshort) && (
      dist = dist[rand(1:(size(dist,1)),5000),:])  #WARNING TEST DEBUG ONLY

    dist = formeventwindows(dist,distdays,
      daysextended=distdaysextended,
      refresheventcrsp=refreshdistcrsp,
      validpermnos=unique(dist.permno),
      Fret = DIST_RET,
      Fdate=:exdate,
      retbound = DIST_RET_BOUND,
      prefix="dist-$(DIST_TYPE[])",
      daystostaleprice=DIST_DAYS_TO_STALE_PRICE,
      yearrange=distyearrange
      )
    #println("Size of file: $(length(unique(dist.eid)))")
    #println(describe(dist))
    #error("stop")
    processdist!(dist)

    outjlsstream("$eventpath\\$eventname.jls.lz4") do s
      serialize(s, (bb,dist))
    end
  else
    injlsstream("$eventpath\\$eventname.jls.lz4") do s
      (bb,dist) = deserialize(s)
    end
  end

  #bb |> CSV.write("output\\cleanbb-$(REPLICATION_TYPE[]).csv")
  dist[dist.exyear .== 2015,:] |> CSV.write(
    "output\\cleandist-$(REPLICATION_TYPE[])-$(DIST_TYPE[])-2015.csv")
  #println(describe(dist[dist.exyear .== 2015,:]))
  return (bb,dist)
end

ext: jl

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
using PastaQ
using ITensors
using Test
using LinearAlgebra

@testset "array" begin
  N = 5
  ψ = productstate(N)
  ψvec = PastaQ.array(ψ)
  @test size(ψvec) == (1 << N,)

  ρ = MPO(productstate(N))
  ρmat = PastaQ.array(ρ)
  @test size(ρmat) == (1 << N, 1 << N)

  ρ = randomstate(N; mixed=true)
  ρmat = PastaQ.array(ρ)
  @test size(ρmat) == (1 << N, 1 << N)

  U = randomprocess(N)
  Umat = PastaQ.array(U)
  @test size(Umat) == (1 << N, 1 << N)

  N = 3
  Λ = randomprocess(N; mixed=true)
end

@testset "array for MPS/MPO" begin

  N = 5
  depth = 3
  circuit = randomcircuit(N; depth = depth)
  qubits = siteinds("Qubit", N)

  # MPS wavefunction
  ψ = runcircuit(qubits, circuit)
  ψvec = PastaQ.array(ψ)

  basis = vec(reverse.(Iterators.product(fill([0,1],N)...)))

  for (k,s) in enumerate(basis)
    σ = productstate(qubits, [s...])
    ψσ = inner(σ,ψ)
    @test ψσ ≈ ψvec[k]
  end

  # MPO unitary
  U = runcircuit(qubits, circuit; process = true)
  Umat = PastaQ.array(U)

  basis = vec(reverse.(Iterators.product(fill([0,1],N)...)))

  for (k,s) in enumerate(basis)
    σ = productstate(qubits, [s...])
    for (kp,sp) in enumerate(basis)
      σp = productstate(qubits, [sp...])
      Uσσp = inner(σ, U, σp)
      @test Uσσp ≈ Umat[k,kp]
    end
  end

  # MPO density matrix
  ρ = runcircuit(qubits, circuit; noise = ("DEP",(p=0.01,)))
  ρmat = PastaQ.array(ρ)

  basis = vec(reverse.(Iterators.product(fill([0,1],N)...)))

  for (k,s) in enumerate(basis)
    σ = productstate(qubits, [s...])
    for (kp,sp) in enumerate(basis)
      σp = productstate(qubits, [sp...])
      ρσσp = inner(σ, ρ, σp)
      @test ρσσp ≈ ρmat[k,kp]
    end
  end

  # LPDO density matrix
  ρ = normalize!(randomstate(qubits; χ = 4, ξ = 2))
  ρmat = PastaQ.array(ρ)

  ρmpo = MPO(ρ)
  basis = vec(reverse.(Iterators.product(fill([0,1],N)...)))

  for (k,s) in enumerate(basis)
    σ = productstate(qubits, [s...])
    for (kp,sp) in enumerate(basis)
      σp = productstate(qubits, [sp...])
      ρσσp = inner(σ, ρmpo, σp)
      @test ρσσp ≈ ρmat[k,kp]
    end
  end

  # MPO Choi matrix
  N = 3
  depth = 3
  circuit = randomcircuit(N; depth = depth)
  qubits = siteinds("Qubit", N)

  Λ = runcircuit(qubits, circuit; process = true, noise = ("DEP",(p=0.01,)))
  Λmat = PastaQ.array(Λ)

  basis = vec(reverse.(Iterators.product(fill([0,1], 2*N)...)))
  qubits_in  = [firstind(Λ[j], tags = "Input", plev = 0) for j in 1:N]
  qubits_out = [firstind(Λ[j], tags = "Output", plev = 0) for j in 1:N]

  for (k,s) in enumerate(basis)
    s_in  = s[1:2:end]
    s_out = s[2:2:end]
    for (kp,sp) in enumerate(basis)
      sp_in  = sp[1:2:end]
      sp_out = sp[2:2:end]
      Λc = copy(Λ)
      for j in 1:N
        Λc[j] = Λc[j] * prime(state(qubits_in[j],s_in[j]+1)) * state(qubits_in[j], sp_in[j]+1)
      end
      σ  = productstate(qubits_out, [s_out...])
      σp = productstate(qubits_out, [sp_out...])
      Λel = inner(σ, Λc, σp)
      @test Λel ≈ Λmat[k,kp]
    end
  end

  # LPDO Choi matrix

  Λ = normalize!(randomprocess(qubits; χ = 4, ξ = 2); localnorm = 2)
  Λmat = PastaQ.array(Λ)

  Λmpo = MPO(Λ)
  qubits_in  = [firstind(Λ.X[j], tags = "Input", plev = 0) for j in 1:N]
  qubits_out = [firstind(Λ.X[j], tags = "Output", plev = 0) for j in 1:N]

  for (k,s) in enumerate(basis)
    s_in  = s[1:2:end]
    s_out = s[2:2:end]
    for (kp,sp) in enumerate(basis)
      sp_in  = sp[1:2:end]
      sp_out = sp[2:2:end]
      Λc = copy(Λmpo)
      for j in 1:N
        Λc[j] = Λc[j] * prime(state(qubits_in[j],s_in[j]+1)) * state(qubits_in[j], sp_in[j]+1)
      end
      σ  = productstate(qubits_out, [s_out...])
      σp = productstate(qubits_out, [sp_out...])
      Λel = inner(σ, Λc, σp)
      @test Λel ≈ Λmat[k,kp]
    end
  end
end

@testset "array for full representation" begin

  N = 5
  depth = 3
  circuit = randomcircuit(N; depth = depth)
  qubits = siteinds("Qubit", N)

  # MPS wavefunction
  ψ = runcircuit(qubits, circuit)
  ψvec = PastaQ.array(ψ)
  ψprod = prod(ψ)
  ψtest = PastaQ.array(ψprod)
  @test ψvec ≈ ψtest

  # MPO unitary
  U = runcircuit(qubits, circuit; process = true)
  Umat = PastaQ.array(U)
  Uprod = prod(U)
  Utest = PastaQ.array(Uprod)
  @test Umat ≈ Utest

  # MPO density matrix
  ρ = runcircuit(qubits, circuit; noise = ("DEP",(p=0.01,)))
  ρmat = PastaQ.array(ρ)
  ρprod = prod(ρ)
  ρtest = PastaQ.array(ρprod)
  @test ρtest ≈ ρmat

  # LPDO density matrix
  ρ = normalize!(randomstate(qubits; χ = 4, ξ = 2))
  ρmat = PastaQ.array(ρ)
  ρprod = prod(ρ)
  ρtest = PastaQ.array(ρprod)
  @test ρtest ≈ ρmat

  # MPO Choi matrix
  N = 3
  depth = 3
  circuit = randomcircuit(N; depth = depth)
  qubits = siteinds("Qubit", N)

  Λ = runcircuit(qubits, circuit; process = true, noise = ("DEP",(p=0.01,)))
  Λmat = PastaQ.array(Λ)
  Λprod = prod(Λ)
  Λtest = PastaQ.array(Λprod)
  @test Λmat ≈ Λtest

  # LPDO Choi matrix
  Λ = normalize!(randomprocess(qubits; χ = 4, ξ = 2); localnorm = 2)
  Λmat = PastaQ.array(Λ)
  Λprod = prod(Λ)
  Λtest = PastaQ.array(Λprod)
  @test Λmat ≈ Λtest
end

@testset "dense states to itensors" begin
  N = 2
  d = 1<<N
  gates = randomcircuit(N; depth = 4)
  ψ = runcircuit(N,gates)

  sites = siteinds("Qubit", N)
  ψvec = PastaQ.array(ψ)
  ϕ = PastaQ.itensor(ψvec, sites)
  @test PastaQ.array(ϕ) ≈ ψvec

  ρ = runcircuit(N, gates; noise = ("DEP",(p=0.1,)))
  ρmat = PastaQ.array(ρ)
  ϱ = PastaQ.itensor(ρmat, sites)
  @test PastaQ.array(ϱ) ≈ ρmat

end


text:
export gaussian_beam!

function gaussian_beam!(arr, L, z, w₀; λ=λ0, n=1)
    L_new = let
        if typeof(L) <: Number
            (L, L)
        else
            L
        end
    end
    # array with the frequencies in Fourier space
    x = to_gpu_or_cpu(arr, Zygote.@ignore eltype(arr).(fftpos(L_new[2], size(arr, 2))'))
    y = to_gpu_or_cpu(arr, Zygote.@ignore eltype(arr).(fftpos(L_new[1], size(arr, 1))))

    k = calc_k(λ, n)
    r2 = x .^ 2 .+  y .^ 2

    zᵣ = gaussian_zᵣ(z, w₀, λ=λ, n=n)

    wz = gaussian_w(z, zᵣ, w₀)
    Rz = gaussian_R(z, zᵣ)
    ψz = gouy_phase(z, zᵣ)

    arr_out = arr .* w₀ ./ wz .* exp.(-r2/wz.^2) .* exp.(-1im .* (k .* z .+ k .* r2 ./ (2 .* Rz) .- ψz))
    return arr_out
end

# paraxial approx
function gaussian_θ(w₀; λ=λ0, n=1)
    return λ / oftype(w₀, π) / n / w₀
    #return atan(gaussian_w(z, gaussian_zᵣ(z, w₀, λ=λ, n=n), w₀) / z)
end

function gaussian_zᵣ(z, w₀; λ=λ0, n=1)
    return oftype(z, π) * w₀^2 * n / λ
end

function gouy_phase(z, zᵣ)
    return atan(z / zᵣ)
end

function gaussian_R(z, zᵣ)
    if iszero(z)
        return oftype(z, Inf)
    else
        return z * (1 + (zᵣ / z)^2)
    end
end

function gaussian_w(z, zᵣ, w₀)
    return w₀ * √(1 + (z / zᵣ)^2)
end

max_stars_repo_name: JuliaPhysics/PhysicalOptics

max_issues_repo_name: JuliaPhysics/PhysicalOptics

max_forks_repo_name: JuliaPhysics/PhysicalOptics


text:
include("alternatingMinimization.jl")
include("convexRelaxation.jl")

struct node
    zero_indices::Array{Any, 1}
    one_indices::Array{Any, 1}
    upper_bound::Float64
    lower_bound::Float64
    relaxed_binary_matrix::Array{Float64, 2}
end

function isTerminal(one_indices, zero_indices, k_sparse, matrix_dim)
    """
    This function evaluates whether the current node in the Branch and Bound
    tree is a terminal node.

    :param zero_indices: List of 2-tuples of Int64 where each entry consists of
                         an index (i, j) such that Y_ij is constrained to take
                         value 0.
    :param one_indices: List of 2-tuples of Int64 where each entry consists of
                        an index (i, j) such that S_ij is constrained to take
                        value 1 where S is the binary matrix denoting the
                        sparsity pattern of Y.
    :param k_sparse: The maximum sparsity of the sparse matrix (Int64).
    :param matrix_dim: The row and column dimension of the data (Int64).

    :return: True if this is a terminal node, false otherwise (Bool).
    """

    num_fixed_ones = size(one_indices)[1]
    num_fixed_zeros = size(zero_indices)[1]

    if num_fixed_ones == k_sparse
        return true
    end

    if num_fixed_zeros == matrix_dim ^ 2 - k_sparse
        return true
    end

    return false

end

function SLR_BnB(U, mu, lambda, k_sparse, k_rank; epsilon=0.1,
    iter_min_improv=1e-12, symmetric_tree=false, output_to_file=false,
    output_file_name="temp.txt")
    """
    This function computes a certifiably near-optimal solution to the problem
        min ||U - X - Y||_F^2 + lambda * ||X||_F^2 + mu * ||Y||_F^2
        subject to rank(X) <= k_rank, ||Y||_0 <= k_sparse

    by executing a custom branch and bound algorithm.

    :param U: An arbitrary n-by-n matrix.
    :param mu: The regularization parameter for the sparse matrix penalty
                   (Float64).
    :param lambda: The regularization parameter for the low rank matrix penalty
                   (Float64).
    :param k_sparse: The maximum sparsity of the sparse matrix (Int64).
    :param k_rank: The maximum rank of the low rank matrix (Int64).
    :param epsilon: The Branch and Bound termination criterion. This function
                    will terminate with an epsilon optimal solution to the
                    optimization problem (Float64).
    :param iter_min_improv: The minimal fractional decrease in the objective
                            value required for the alternating minimization
                            procedure to continue iterating.
    :param symmetric_tree: If true, Branch and Bound will only explore symmetric
                           sparsity patterns (Bool).
    :param output_to_file: If true, progress output will be printed to file. If
                           false, progress output will be printed to the console
                           (Bool).
    :param output_file_name: Name of the file progress output will be printed to
                             when output_to_file=true (String).

    :return: This function returns eight values:
             1) A tuple of two n-by-n arrays that correspond to the globally
                optimal solution to the optimization problem. The first element
                in the tuple is the optimal low rank matrix and the second
                element in the tuple is the optimal sparse matrix.
             2) The final global upper bound (Float64).
             3) The final global lower bound (Float64).
             4) The number of nodes explored during the optimization process
                (Int64).
             5) A list of values of type Float64 representing the evolution of
                the upper bound during the optimization procedure.
             6) A list of values of type Float64 representing the evolution of
                the lower bound during the optimization procedure.
             7) The total time elapsed during the optimiation process
                (milliseconds).
             8) The number of terminal nodes explored during the optimization
                process (Int64).

                best feasible solution found
             by this method (the first element in the tuple is the matrix X and
             the second element is the matrix Y). The second value is the
             objective value (Float64) achieved by the returned feasible
             solution.
    """

    start_time = now()

    # Solve the root node
    init_zero_indices = []
    init_one_indices = []

    upper_bound_sol, upper_bound_obj = SLR_AM(U, mu, lambda, k_sparse, k_rank,
        random_restarts=0, zero_indices=init_zero_indices,
        one_indices=init_one_indices, min_improvement=iter_min_improv)

    lower_bound_sol, lower_bound_obj = solve_perspective_relaxation(U, mu,
        lambda, k_sparse, k_rank, zero_indices=init_zero_indices,
        one_indices=init_one_indices)

    root_optimality_gap = (upper_bound_obj - lower_bound_obj) / upper_bound_obj
    if output_to_file
        open(output_file_name, "a") do io
            write(io, "Root Node has been solved.\n")
            write(io, "Root Node upper bound is: $upper_bound_obj\n")
            write(io, "Root Node lower bound is: $lower_bound_obj\n")
            write(io, "Root Node optimality gap is: $root_optimality_gap\n\n")
        end
    else
        println("Root Node has been solved.")
        println("Root Node upper bound is: $upper_bound_obj")
        println("Root Node lower bound is: $lower_bound_obj")
        println("Root Node optimality gap is: $root_optimality_gap")
        println()
    end

    # Initialize the node list and global bounds
    master_node_list = []
    root_node = node(init_zero_indices,
                     init_one_indices,
                     upper_bound_obj,
                     lower_bound_obj,
                     lower_bound_sol[3])
    push!(master_node_list, root_node)
    global_upper_bound = upper_bound_obj
    global_lower_bound = lower_bound_obj
    global_X_sol = upper_bound_sol[1]
    global_Y_sol = upper_bound_sol[2]

    num_explored_nodes = 1
    upper_bound_hist = [global_upper_bound]
    lower_bound_hist = [global_lower_bound]

    n = size(U)[1]

    terminal_nodes = 0
    # Main branch and bound loop
    while (global_upper_bound - global_lower_bound) /
                                            global_upper_bound > epsilon

        """
        # Select the most recently added node (DFS)
        current_node = pop!(master_node_list)
        """
        current_node_index = findall(this_node->
            this_node.lower_bound == global_lower_bound, master_node_list)[1]
        current_node = master_node_list[current_node_index]
        deleteat!(master_node_list, current_node_index)

        # If the node is a terminal node, place it back on the master list
        if isTerminal(current_node.one_indices, current_node.zero_indices,
                      k_sparse, n)
            prepend!(master_node_list, current_node)
            continue
        end
        # Select entry to branch on
        index = argmin(abs.(current_node.relaxed_binary_matrix .- 0.5))
        i = index[1]
        j = index[2]

        # Construct the two child sets of indices from this parent node
        one_condition = size(current_node.one_indices)[1] <= k_sparse - 2
        zero_condition = size(current_node.zero_indices)[1] <= n^2 - k_sparse - 2
        sym_branch_feasible = one_condition && zero_condition
        if symmetric_tree && sym_branch_feasible && (i != j)
            new_index_zero = (append!(copy(current_node.zero_indices),
                              [(i, j), (j, i)]), copy(current_node.one_indices))
            new_index_one = (copy(current_node.zero_indices),
                             append!(copy(current_node.one_indices),
                             [(i, j), (j, i)]))
            new_index_list = [new_index_zero, new_index_one]
        else
            new_index_zero = (append!(copy(current_node.zero_indices), [(i, j)]),
                              copy(current_node.one_indices))
            new_index_one = (copy(current_node.zero_indices),
                             append!(copy(current_node.one_indices), [(i, j)]))
            new_index_list = [new_index_zero, new_index_one]
        end

        for (zero_indices, one_indices) in new_index_list

            upper_bound_sol, upper_bound_obj = SLR_AM(U, mu, lambda, k_sparse,
                        k_rank, random_restarts=0, zero_indices=zero_indices,
                        one_indices=one_indices,
                        min_improvement=iter_min_improv)

            if isTerminal(one_indices, zero_indices, k_sparse, n)
                lower_bound_sol = upper_bound_sol
                lower_bound_obj = upper_bound_obj
                relaxed_binary_matrix = abs.(lower_bound_sol[2]) .> 1e-10
                terminal_nodes = terminal_nodes + 1
            else
                lower_bound_sol, lower_bound_obj = solve_perspective_relaxation(
                    U, mu, lambda, k_sparse, k_rank, zero_indices=zero_indices,
                    one_indices=one_indices)
                relaxed_binary_matrix = lower_bound_sol[3]
            end

            if upper_bound_obj < global_upper_bound
                global_upper_bound = upper_bound_obj
                global_X_sol = upper_bound_sol[1]
                global_Y_sol = upper_bound_sol[2]

                current_gap = global_upper_bound - global_lower_bound
                current_gap = current_gap / global_upper_bound

                println("The new upper bound is: $global_upper_bound")
                println("The current lower bound is: $global_lower_bound")
                println("The current optimality gap is: $current_gap")
                println()

                if output_to_file
                    open(output_file_name, "a") do io
                        write(io, "The new upper bound is: $global_upper_bound\n")
                        write(io, "The current lower bound is: $global_lower_bound\n")
                        write(io, "The current optimality gap is: $current_gap\n\n")
                    end
                else
                    println("The new upper bound is: $global_upper_bound")
                    println("The current lower bound is: $global_lower_bound")
                    println("The current optimality gap is: $current_gap")
                    println()
                end

                new_master_node_list = []
                for this_node in master_node_list
                    if this_node.lower_bound < global_upper_bound
                        push!(new_master_node_list, this_node)
                    end
                end

                pre_size = size(master_node_list)[1] + 1
                post_size = size(new_master_node_list)[1] + 1

                if output_to_file
                    open(output_file_name, "a") do io
                        write(io, "Tree size before pruning: $pre_size\n")
                        write(io, "Tree size after pruning: $post_size\n\n")
                    end
                else
                    println("Tree size before pruning: $pre_size")
                    println("Tree size after pruning: $post_size")
                    println()
                end

                master_node_list = new_master_node_list
            end

            if lower_bound_obj < global_upper_bound
                new_node = node(zero_indices,
                                one_indices,
                                upper_bound_obj,
                                lower_bound_obj,
                                relaxed_binary_matrix)
                push!(master_node_list, new_node)
            end

            num_explored_nodes = num_explored_nodes + 1
            append!(upper_bound_hist, global_upper_bound)
            append!(lower_bound_hist, global_lower_bound)

            if num_explored_nodes % 200 == 0
                current_time = now() - start_time

                if output_to_file
                    open(output_file_name, "a") do io
                        write(io, "$num_explored_nodes nodes have been explored.\n")
                        write(io, "Current elapsed time: $current_time\n")
                        write(io, "The current lower bound is: $global_lower_bound\n")
                    end
                else
                    println("$num_explored_nodes nodes have been explored.")
                    println("Current elapsed time: $current_time")
                    println("The current lower bound is: $global_lower_bound")
                end
            end

        end

        if size(master_node_list)[1] == 0
            if output_to_file
                open(output_file_name, "a") do io
                    write(io, "All nodes have been explored")
                end
            else
                println("All nodes have been explored")
            end
            break
        end

        global_lower_bound = master_node_list[1].lower_bound
        for current_node in master_node_list
            if current_node.lower_bound < global_lower_bound
                global_lower_bound = current_node.lower_bound
            end
        end
    end

    total_time = now() - start_time

    if output_to_file
        open(output_file_name, "a") do io
            write(io, "An optimal solution has been found!")
        end
    else
        println("An optimal solution has been found!")
    end

    return (global_X_sol, global_Y_sol),
           global_upper_bound,
           global_lower_bound,
           num_explored_nodes,
           upper_bound_hist,
           lower_bound_hist,
           total_time,
           terminal_nodes

end;


text:
using YAXArrays, YAXArrayBase, Test, Dates

@testset "YAXArrays" begin
    data = collect(reshape(1:20, 4, 5))
    axlist = [RangeAxis("XVals", 1.0:4.0), CategoricalAxis("YVals", [1, 2, 3, 4, 5])]
    props = Dict("att1" => 5, "att2" => "Hallo")
    a = YAXArray(axlist, data, props)

    @testset "Array construction sanity checks" begin
        @test_throws ArgumentError YAXArray(axlist[1:1], data, props)
        @test_throws ArgumentError YAXArray(
            [RangeAxis("XVals", 1.0:4.0), CategoricalAxis("YVals", [1, 2, 3, 4])],
            data,
            props,
        )
    end

    @testset "Basic array Functions" begin
        @test size(a) == (4, 5)
        @test size(a, 1) == 4
        @test size(a, "XVals") == 4
        @test size(a, 2) == 5
        @test size(a, :YVals) == 5
        @test eltype(a) == Int
        @test ndims(a) == 2
        @test a.XVals == axlist[1]
        @test a.YVals == axlist[2]
        pa = permutedims(a, (2, 1))
        @test pa.axes == axlist[[2, 1]]
        @test pa.data == transpose(data)
        @test pa.XVals == axlist[1]
        @test pa.YVals == axlist[2]
        @test a[2, 3] == a.data[2, 3]
        @test read(a) isa Array
        @test read(a) == collect(reshape(1:20, 4, 5))
        @test propertynames(a) == (:XVals, :YVals, :axes, :data)
        @test propertynames(a, true) == (:XVals, :YVals, :axes, :data, :properties)
    end

    @testset "YAXArray interface functions" begin

        a2 = readcubedata(a)
        @test a2.axes == a.axes
        @test a.data == a2.data
        @test a2.data isa Array

        @test caxes(a) == axlist

        @test YAXArrays.Cubes.cubechunks(a) == (4, 5)

        cs = CartesianIndices.(collect(Iterators.product([1:2, 3:5, 6:6], [1:2, 3:4])))

        renameaxis!(a2, "XVals" => "Ax1")

        @test YAXArrayBase.dimname(a2, 1) == :Ax1

        renameaxis!(a2, "Ax1" => RangeAxis("XVals", 2:5))

        @test a2.axes[1].values == 2:5

        b = IOBuffer()
        show(b, a)
        # @test String(take!(b))=="""
        # YAXArray with the following dimensions
        # XVals               Axis with 4 Elements from 1.0 to 4.0
        # YVals               Axis with 5 elements: 1 2 3 4 5
        # Total size: 160.0 bytes
        # """
    end

    @testset "YAXArrayBase interface" begin
        using YAXArrayBase
        @test YAXArrayBase.dimvals(a, 1) == 1.0:4.0
        @test YAXArrayBase.dimvals(a, 2) == [1, 2, 3, 4, 5]
        @test YAXArrayBase.dimname(a, 1) == :XVals
        @test YAXArrayBase.dimname(a, 2) == :YVals
        @test YAXArrayBase.getattributes(a) == Dict("att1" => 5, "att2" => "Hallo")
        @test YAXArrayBase.iscontdim(a, 1) == true
        @test YAXArrayBase.iscontdim(a, 2) == false
        @test YAXArrayBase.getdata(a) === data
        a2 = YAXArrayBase.yaxcreate(
            YAXArray,
            reshape(1:27, 3, 3, 3),
            ["d1", "d2", "d3"],
            [1:3, ["One", "Two", "Three"], 0.1:0.1:0.3],
            Dict("att" => 5),
        )
        @test a2 isa YAXArray
        @test a2.axes[1] == RangeAxis("d1", 1:3)
        @test a2.axes[2] == CategoricalAxis("d2", ["One", "Two", "Three"])
        @test a2.axes[3] == RangeAxis("d3", 0.1:0.1:0.3)
        @test a2.data == reshape(1:27, 3, 3, 3)
        @test a2.properties == Dict("att" => 5)
        @test YAXArrayBase.iscompressed(a) == false
    end

    @testset "Subsets" begin
        s = YAXArrays.Cubes.subsetcube(a, X = 1.5..3.5)
        @test s.data == [2 6 10 14 18; 3 7 11 15 19]
        @test s.axes[1] == RangeAxis("XVals", 2.0:3.0)
        @test s.axes[2] == CategoricalAxis("YVals", [1, 2, 3, 4, 5])
        ax = a.axes[1]
        @test YAXArrays.Cubes.interpretsubset(CartesianIndices((1:2,)), ax) == 1:2
        @test YAXArrays.Cubes.interpretsubset(CartesianIndex((2,)), ax) == 2
        @test YAXArrays.Cubes.interpretsubset(2.1, ax) == 2
        @test YAXArrays.Cubes.interpretsubset((3.5, 1.5), ax) == 2:3
        @test YAXArrays.Cubes.interpretsubset(0.8..2.2, ax) == 1:2
        tax = RangeAxis("ADate", Date(2001):Day(1):Date(2003, 2, 28))
        @test YAXArrays.Cubes.interpretsubset((Date(2001, 1, 2), Date(2001, 1, 5)), tax) ==
              2:4
        @test YAXArrays.Cubes.interpretsubset(2001:2002, tax) == 1:730
        @test YAXArrays.Cubes.interpretsubset([1, 3, 5], a.axes[2]) == [1, 3, 5]

        s2 = a[X = 0.5..3.5, Y = [1, 5, 4]]
        @test s2.data == [1 17 13; 2 18 14; 3 19 15]
        @test s2.axes[1] == RangeAxis("XVals", 1.0:3.0)
        @test s2.axes[2] == CategoricalAxis("YVals", [1, 5, 4])
    end

end

max_stars_repo_name: meggart/YAXArrays

max_forks_repo_name: meggart/YAXArrays


text:
import Base.show

function show(io::IO, z::ArbRadixFloat{radix,precision,Tdigit}) where {radix,precision,Tdigit}
    print(io, '(')
    print(io, prettystring_significand(z.significand))
    print(io, ')')
    print(io, subscriptstring(radix))
    print(io, " × (")
    print(io, radix)
    print(io, ')')
    print(io, superscriptstring(z.exponent))
end

"This character gets used to print negative digits"
const overline = '\u305'

function prettystring_significand(significand)
    s = IOBuffer()

    print_point = true
    #Test range
    maxdigit = maximum(abs.(extrema(significand)))
    if maxdigit >= 36 #Do not attempt to pretty print
        print(s, significand)
    else# maxdigit < 36, use base36
        for digit in significand
            absdigit = abs(digit)
            if absdigit > 10
                print(s, 'a'+absdigit-10)
            else
                print(s, absdigit)
            end
            if digit < 0
                print(s, overline)
            end
            if print_point
                print(s, '.')
                print_point = false
            end
        end
    end
    String(take!(s))
end

const superscriptize = Dict(
    '1' => '¹',
    '2' => '²',
    '3' => '³',
    '4' => '⁴',
    '5' => '⁵',
    '6' => '⁶',
    '7' => '⁷',
    '8' => '⁸',
    '9' => '⁹',
    '0' => '⁰',

    '-' => '\u207b',
    '!' => 'ᵎ',

    'A' => 'ᴬ',
    'B' => 'ᴮ',
    'D' => 'ᴰ',
    'E' => 'ᴱ',
    'G' => 'ᴳ',
    'H' => 'ᴴ',
    'I' => 'ᴵ',
    'J' => 'ᴶ',
    'K' => 'ᴷ',
    'L' => 'ᴸ',
    'M' => 'ᴹ',
    'N' => 'ᴺ',
    'O' => 'ᴼ',
    'P' => 'ᴾ',
    'R' => 'ᴿ',
    'T' => 'ᵀ',
    'U' => 'ᵁ',
    'V' => 'ⱽ',
    'W' => 'ᵂ',

    'a' => 'ᵃ',
    'b' => 'ᵇ',
    'c' => 'ᶜ',
    'd' => 'ᵈ',
    'e' => 'ᵉ',
    'f' => 'ᶠ',
    'g' => 'ᵍ',
    'h' => 'ʰ',
    'i' => 'ⁱ',
    'j' => 'ʲ',
    'k' => 'ᵏ',
    'l' => 'ˡ',
    'm' => 'ᵐ',
    'n' => 'ⁿ',
    'o' => 'ᵒ',
    'p' => 'ᵖ',
    'r' => 'ʳ',
    's' => 'ˢ',
    't' => 'ᵗ',
    'u' => 'ᵘ',
    'v' => 'ᵛ',
    'w' => 'ʷ',
    'x' => 'ˣ',
    'y' => 'ʸ',
    'z' => 'ᶻ',
)
function superscriptstring(num)
    s = IOBuffer()
    for d in string(num)
        print(s, get(superscriptize, d, d))
    end
    String(take!(s))
end

#Most letters don't have a subscript counterpart...
const subscriptize = Dict(
    '1' => '₁',
    '2' => '₂',
    '3' => '₃',
    '4' => '₄',
    '5' => '₅',
    '6' => '₆',
    '7' => '₇',
    '8' => '₈',
    '9' => '₉',
    '0' => '₀',

    '+' => '₊',
    '-' => '₋',
    '=' => '₌',
    '(' => '₍',
    ')' => '₎',

    'a' => 'ₐ',
    'e' => 'ₑ',
    'h' => 'ₕ',
    'i' => 'ᵢ',
    'j' => 'ⱼ',
    'k' => 'ₖ',
    'l' => 'ₗ',
    'm' => 'ₘ',
    'n' => 'ₙ',
    'o' => 'ₒ',
    'p' => 'ₚ',
    'r' => 'ᵣ',
    's' => 'ₛ',
    't' => 'ₜ',
    'u' => 'ᵤ',
    'v' => 'ᵥ',
    'x' => 'ₓ',

    'β' => 'ᵦ',
    'γ' => 'ᵧ',
    'ρ' => 'ᵨ',
    'ϕ' => 'ᵩ',
    'φ' => 'ᵩ',
    'χ' => 'ᵪ',
)
function subscriptstring(num)
    s = IOBuffer()
    for d in string(num)
        print(s, get(subscriptize, d, d))
    end
    String(take!(s))
end


text:
# This file is a part of UpROOT.jl, licensed under the MIT License (MIT).

"""
    UpROOT.pyobj(x)::PyObject

Get the python object wrapped by Julia object `x`. `x` may be a
[`TFile`](@ref)/[`TDirectory`](@ref), [`TDirectory`](@ref) or
[`TBranch`](@ref).
"""
pyobj(x) = Base.getfield(x, :pyobj)


text:
abstract type AbstractNDT end
@with_kw struct NonDecisionTime{T<:Real} <: AbstractNDT
    lower::T
    upper::T
    function NonDecisionTime{T}(lower::T, upper::T) where T<:Real
        @assert 0 <= lower < upper
        new(lower, upper)
    end
end

NonDecisionTime() = NonDecisionTime(1, 2)


text:
mutable struct Population
    # Manages all the species
    population::Vector{MLPChromosome}
    popsize::Int64
    species::Vector{Species}
    species_history::Vector{Vector{Tuple{Int, Int}}}#Vector{<:Union{Array{Int64,1}, Tuple{Int, Int}}}
    generation::Int64
    avg_fitness::Vector{Float64}
    best_fitness::Vector{MLPChromosome}
    best_no_improvement_age::Int64
    number_species::Int
    simulation_time::Float64
    #setup::Evaluation
    #evaluate::Function # Evaluates population. Override this method in your experiments
    function Population(g::Global, checkpoint::String="")

        if isfile(string(checkpoint, ".jld2"))
            resume_checkpoint(g.cf.checkpoint_filename)
        else
            popsize = g.cf.pop_size # total population size
            population = Vector{MLPChromosome}(undef, popsize)
            #minimal_chromosome = create_minimal_chromossome(g)
            for i in 1:popsize
                g.innov_number = 0
                population[i] = create_minimal_chromossome(g)
            end
            #=setup = Evaluation(g.cf.dataset_id, g.cf.fitness_function_id, normalise=g.cf.normalise_input,
                               balance=g.cf.balance_dataset)=#

            p = new(population, popsize,
                    Species[],Array{Int64,1}[], # currently living species and species history
                    -1, # generation
                    Float64[], # avg_fitness
                    MLPChromosome[], # best_fitness
                    0,
                    0,
                    0.0#,
                    #setup
                    );
            #p.evaluate = (f::Function) -> f.(p.population); # Evaluates population. Override this method in your experiments
            p
        end
    end
    function Population(g::Global, population::Vector{MLPChromosome}, popsize::Int64,
                        species::Vector{Species}, species_history::Vector{Vector{Tuple{Int, Int}}},
                        generation::Int64, avg_fitness::Vector{Float64}, best_fitness::Vector{MLPChromosome},
                        number_species::Int, simulation_time::Float64)
        setup = Evaluation(g.cf.dataset_id, g.cf.fitness_function_id, normalise=g.cf.normalise_input)
        evaluate_function = (f::Function) -> f.(p.population)
        new(population, popsize, species, species_history, generation, avg_fitness,
            best_fitness, number_species, simulation_time, setup, evaluate_function)
    end
end

function Base.show(io::IO, p::Population)
    @printf(io,"Population size: %3d   Total species: %3d", p.popsize, length(p.species))
end

function remove(p::Population, ch::MLPChromosome)
    # Removes a chromosome from the population
    deleteat!(p.population,findfirst(p.population,ch))
    return
end

function remove_specie!(p::Population, id::Int64)
    index = findfirst((s) -> s.id == id, p.species)
    if index != nothing
        deleteat!(p.species, index)
        p.population = filter((i) -> i.species_id != id, p.population)
    else
        println("Specie with id $id not found within population.")
    end
end

function limit_species!(g::Global, p::Population)
    # If the value of limit_species is less than 1, species limitation won't be enabled
    if g.cf.limit_species > 0
        sort!(p.species, by = s -> s.age, rev = true)
        i = 1
        while length(p.species) > g.cf.limit_species && i <= length(p.species)
            sp = p.species[i]
            if !sp.hasBest
                remove_specie!(p, sp.id)
                i -= 1
            end
            i += 1
        end
    end
end

function speciate(g::Global, p::Population, report::Bool)
    # Group chromosomes into species by similarity
        # Speciate the population
    for individual in p.population
        found = false
        for s in p.species
            if distance(individual, s.representant, g.cf) < g.cf.compatibility_threshold
                add(s, individual)
                found = true
                break
            end
        end

        if !found
            push!(p.species, Species(g, individual))
            p.number_species += 1
        end
    end

    # eliminate empty species
    keep = map((s)->length(s)==0 ? false : true, p.species)
    if report
        for i = 1:length(keep)
            if !keep[i] println("Removing species $(p.species[i].id) for being empty") end
        end
    end
    p.species = p.species[keep]
    #set_compatibility_threshold(g, p)
end

function set_compatibility_threshold(g::Global, p::Population)
    # controls compatibility threshold
    if length(p.species) > g.cf.species_size
        g.cf.compatibility_threshold += g.cf.compatibility_change
    elseif length(p.species) < g.cf.species_size
        if g.cf.compatibility_threshold > g.cf.compatibility_change
            g.cf.compatibility_threshold -= g.cf.compatibility_change
        else
            println("Compatibility threshold cannot be changed (minimum value has been reached)")
        end
    end
end

# Returns the average raw fitness of population
average_fitness(p::Population) = mean(map((c) -> c.fitness::Float64, p.population))

stdeviation(p::Population) = std(map((c) -> c.fitness::Float64, p.population))

function compute_spawn_levels(g::Global, p::Population)
    #  Compute each species' spawn amount (Stanley, p. 40)

    # 1. Boost if young and penalize if old
    # TODO: does it really increase the overall performance?
    species_stats = zeros(length(p.species))
    for i = 1:length(p.species)
        s = p.species[i]
        #species_stats[i] = s.age < g.cf.youth_threshold ? average_fitness(s) * g.cf.youth_boost :
        #    s.age > g.cf.old_threshold ? average_fitness(s) * g.cf.youth_boost : average_fitness(s)
        if s.age < g.cf.youth_threshold
            species_stats[i] = average_fitness(s) * g.cf.youth_boost
        elseif s.age > g.cf.old_threshold
           species_stats[i] = average_fitness(s) * g.cf.old_penalty
        else
            species_stats[i] = average_fitness(s)
        end
    end

    # 2. Share fitness (only usefull for computing spawn amounts)
    # Sharing the fitness is only meaningful here
    # we don't really have to change each individual's raw fitness
    total_average = sum(species_stats)
    #push!(p.avg_fitness, mean(total_average))

     # 3. Compute spawn
    for i= 1:length(p.species)
        s = p.species[i]
        s.spawn_amount = Int(round((species_stats[i] * p.popsize / total_average)))
    end
end

function tournamentSelection(p::Population, k=2)
    # Tournament selection with size k (default k=2).
    # randomly select k competitors
    chs = p.population[randperm(length(p.population))[1:k]]
    best = chs[1]
    for ch in chs # choose best among randomly selected
        best = ch.fitness > best.fitness ? ch : best
    end
    return best
end

function simple_log_species(p::Population)
    # Logging species data for visualizing speciation
    specById = sort(p.species, by=s->s.id)
    spec_size = zeros(Int64,specById[end].id+1)
    map(s->spec_size[s.id]=length(s) , specById)
    push!(p.species_history,spec_size)
end

function log_species(p::Population)
    # Logging species data for visualizing speciation
    specById = sort(p.species, by = s->s.id)
    spec_size = map((s) -> (s.id, length(s)) , specById)
    push!(p.species_history, spec_size)
end

function population_diversity(p::Population)
    # Calculates the diversity of population: total average weights,
    # number of connections, nodes

    num_nodes = 0
    avg_weights = 0.0

    for ch in p.population
        num_nodes += number_nodes(ch)
        sum_weights = (n) -> sum(n.weights)
        avg_weights = map((l) -> sum(sum_weights.(l.nodes)), ch.layers) |> sum
    end

    total = length(p.population)
    return num_nodes/total, avg_weights/total
end

max_stars_repo_name: Lucas-Muniz/MNEAT

max_issues_repo_name: Lucas-Muniz/MNEAT

max_forks_repo_name: Lucas-Muniz/MNEAT


text:
# An idealized baroclinic instability test case
# For optimal results consider increasing the resolution to 16x16x8 trees per cube face.
#
# Note that this elixir can take several hours to run.
# Using 24 threads of an AMD Ryzen Threadripper 3990X (more threads don't speed it up further)
# and `check-bounds=no`, this elixirs takes about one hour to run.
# With 16x16x8 trees per cube face on the same machine, it takes about 28 hours.
#
# References:
# - Paul A. Ullrich, Thomas Melvin, Christiane Jablonowski, Andrew Staniforth (2013)
#   A proposed baroclinic wave test case for deep- and shallow-atmosphere dynamical cores
#   https://doi.org/10.1002/qj.2241

using OrdinaryDiffEq
using Trixi
using LinearAlgebra

###############################################################################
# Setup for the baroclinic instability test
gamma = 1.4
equations = CompressibleEulerEquations3D(gamma)

# Initial condition for an idealized baroclinic instability test
# https://doi.org/10.1002/qj.2241, Section 3.2 and Appendix A
function initial_condition_baroclinic_instability(x, t, equations::CompressibleEulerEquations3D)
  lon, lat, r = cartesian_to_sphere(x)
  radius_earth = 6.371229e6
  # Make sure that the r is not smaller than radius_earth
  z = max(r - radius_earth, 0.0)

  # Unperturbated basic state
  rho, u, p = basic_state_baroclinic_instability_longitudinal_velocity(lon, lat, z)

  # Stream function type perturbation
  u_perturbation, v_perturbation = perturbation_stream_function(lon, lat, z)

  u += u_perturbation
  v  = v_perturbation

  # Convert spherical velocity to Cartesian
  v1 = -sin(lon) * u - sin(lat) * cos(lon) * v
  v2 =  cos(lon) * u - sin(lat) * sin(lon) * v
  v3 =  cos(lat) * v

  return prim2cons(SVector(rho, v1, v2, v3, p), equations)
end

# Steady state for RHS correction below
function steady_state_baroclinic_instability(x, t, equations::CompressibleEulerEquations3D)
  lon, lat, r = cartesian_to_sphere(x)
  radius_earth = 6.371229e6
  # Make sure that the r is not smaller than radius_earth
  z = max(r - radius_earth, 0.0)

  # Unperturbated basic state
  rho, u, p = basic_state_baroclinic_instability_longitudinal_velocity(lon, lat, z)

  # Convert spherical velocity to Cartesian
  v1 = -sin(lon) * u
  v2 =  cos(lon) * u
  v3 = 0.0

  return prim2cons(SVector(rho, v1, v2, v3, p), equations)
end

function cartesian_to_sphere(x)
  r = norm(x)
  lambda = atan(x[2], x[1])
  if lambda < 0
    lambda += 2 * pi
  end
  phi = asin(x[3] / r)

  return lambda, phi, r
end

# Unperturbated balanced steady-state.
# Returns primitive variables with only the velocity in longitudinal direction (rho, u, p).
# The other velocity components are zero.
function basic_state_baroclinic_instability_longitudinal_velocity(lon, lat, z)
  # Parameters from Table 1 in the paper
  # Corresponding names in the paper are commented
  radius_earth                   = 6.371229e6  # a
  half_width_parameter           = 2           # b
  gravitational_acceleration     = 9.80616     # g
  k                              = 3           # k
  surface_pressure               = 1e5         # p₀
  gas_constant                   = 287         # R
  surface_equatorial_temperature = 310.0       # T₀ᴱ
  surface_polar_temperature      = 240.0       # T₀ᴾ
  lapse_rate                     = 0.005       # Γ
  angular_velocity               = 7.29212e-5  # Ω

  # Distance to the center of the Earth
  r = z + radius_earth

  # In the paper: T₀
  temperature0 = 0.5 * (surface_equatorial_temperature + surface_polar_temperature)
  # In the paper: A, B, C, H
  const_a = 1 / lapse_rate
  const_b = (temperature0 - surface_polar_temperature) /
            (temperature0 * surface_polar_temperature)
  const_c = 0.5 * (k + 2) * (surface_equatorial_temperature - surface_polar_temperature) /
                            (surface_equatorial_temperature * surface_polar_temperature)
  const_h = gas_constant * temperature0 / gravitational_acceleration

  # In the paper: (r - a) / bH
  scaled_z = z / (half_width_parameter * const_h)

  # Temporary variables
  temp1 = exp(lapse_rate/temperature0 * z)
  temp2 = exp(-scaled_z^2)

  # In the paper: ̃τ₁, ̃τ₂
  tau1 = const_a * lapse_rate / temperature0 * temp1 + const_b * (1 - 2 * scaled_z^2) * temp2
  tau2 = const_c * (1 - 2 * scaled_z^2) * temp2

  # In the paper: ∫τ₁(r') dr', ∫τ₂(r') dr'
  inttau1 = const_a * (temp1 - 1) + const_b * z * temp2
  inttau2 = const_c * z * temp2

  # Temporary variables
  temp3 = r/radius_earth * cos(lat)
  temp4 = temp3^k - k/(k + 2) * temp3^(k+2)

  # In the paper: T
  temperature = 1 / ((r/radius_earth)^2 * (tau1 - tau2 * temp4))

  # In the paper: U, u (zonal wind, first component of spherical velocity)
  big_u = gravitational_acceleration/radius_earth * k * temperature * inttau2 * (temp3^(k-1) - temp3^(k+1))
  temp5 = radius_earth * cos(lat)
  u = -angular_velocity * temp5 + sqrt(angular_velocity^2 * temp5^2 + temp5 * big_u)

  # Hydrostatic pressure
  p = surface_pressure * exp(-gravitational_acceleration/gas_constant * (inttau1 - inttau2 * temp4))

  # Density (via ideal gas law)
  rho = p / (gas_constant * temperature)

  return rho, u, p
end

# Perturbation as in Equations 25 and 26 of the paper (analytical derivative)
function perturbation_stream_function(lon, lat, z)
  # Parameters from Table 1 in the paper
  # Corresponding names in the paper are commented
  perturbation_radius      = 1/6      # d₀ / a
  perturbed_wind_amplitude = 1.0      # Vₚ
  perturbation_lon         = pi/9     # Longitude of perturbation location
  perturbation_lat         = 2 * pi/9 # Latitude of perturbation location
  pertz                    = 15000    # Perturbation height cap

  # Great circle distance (d in the paper) divided by a (radius of the Earth)
  # because we never actually need d without dividing by a
  great_circle_distance_by_a = acos(sin(perturbation_lat) * sin(lat) +
                                    cos(perturbation_lat) * cos(lat) *
                                    cos(lon - perturbation_lon))

  # In the first case, the vertical taper function is per definition zero.
  # In the second case, the stream function is per definition zero.
  if z > pertz || great_circle_distance_by_a > perturbation_radius
    return 0.0, 0.0
  end

  # Vertical tapering of stream function
  perttaper = 1.0 - 3 * z^2 / pertz^2 + 2 * z^3 / pertz^3

  # sin/cos(pi * d / (2 * d_0)) in the paper
  sin_, cos_ = sincos(0.5 * pi * great_circle_distance_by_a / perturbation_radius)

  # Common factor for both u and v
  factor = 16 / (3 * sqrt(3)) * perturbed_wind_amplitude * perttaper * cos_^3 * sin_

  u_perturbation = -factor * (-sin(perturbation_lat) * cos(lat) +
                             cos(perturbation_lat) * sin(lat) * cos(lon - perturbation_lon)
                            ) / sin(great_circle_distance_by_a)

  v_perturbation = factor * cos(perturbation_lat) * sin(lon - perturbation_lon) /
      sin(great_circle_distance_by_a)

  return u_perturbation, v_perturbation
end

@inline function source_terms_baroclinic_instability(u, x, t, equations::CompressibleEulerEquations3D)
  radius_earth               = 6.371229e6  # a
  gravitational_acceleration = 9.80616     # g
  angular_velocity           = 7.29212e-5  # Ω

  r = norm(x)
  # Make sure that r is not smaller than radius_earth
  z = max(r - radius_earth, 0.0)
  r = z + radius_earth

  du1 = zero(eltype(u))

  # Gravity term
  temp = -gravitational_acceleration * radius_earth^2 / r^3
  du2 = temp * u[1] * x[1]
  du3 = temp * u[1] * x[2]
  du4 = temp * u[1] * x[3]
  du5 = temp * (u[2] * x[1] + u[3] * x[2] + u[4] * x[3])

  # Coriolis term, -2Ω × ρv = -2 * angular_velocity * (0, 0, 1) × u[2:4]
  du2 -= -2 * angular_velocity * u[3]
  du3 -=  2 * angular_velocity * u[2]

  return SVector(du1, du2, du3, du4, du5)
end

###############################################################################
# Start of the actual elixir, semidiscretization of the problem

initial_condition = initial_condition_baroclinic_instability

boundary_conditions = Dict(
  :inside  => boundary_condition_slip_wall,
  :outside => boundary_condition_slip_wall,
)

# This is a good estimate for the speed of sound in this example.
# Other values between 300 and 400 should work as well.
surface_flux = FluxLMARS(340)
volume_flux  = flux_kennedy_gruber
solver = DGSEM(polydeg=5, surface_flux=surface_flux, volume_integral=VolumeIntegralFluxDifferencing(volume_flux))

# For optimal results, use (16, 8) here
trees_per_cube_face = (8, 4)
mesh = Trixi.P4estMeshCubedSphere(trees_per_cube_face..., 6.371229e6, 30000.0,
                                  polydeg=5, initial_refinement_level=0)

semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver,
                                    source_terms=source_terms_baroclinic_instability,
                                    boundary_conditions=boundary_conditions)

###############################################################################
# ODE solvers, callbacks etc.

tspan = (0.0, 10 * 24 * 60 * 60.0) # time in seconds for 10 days

# Save RHS of the steady state and subtract it in every RHS evaluation.
# This trick preserves the steady state exactly (to machine rounding errors, of course).
# Otherwise, this elixir produces entirely unusable results for a resolution of 8x8x4 cells
# per cube face with a polydeg of 3.
# With this trick, even the polydeg 3 simulation produces usable (although badly resolved) results,
# and most of the grid imprinting in higher polydeg simulation is eliminated.
#
# See https://github.com/trixi-framework/Trixi.jl/issues/980 for more information.
u_steady_state = compute_coefficients(steady_state_baroclinic_instability, tspan[1], semi)
# Use a `let` block for performance (otherwise du_steady_state will be a global variable)
let du_steady_state = similar(u_steady_state)
  # Save RHS of the steady state
  Trixi.rhs!(du_steady_state, u_steady_state, semi, tspan[1])

  global function corrected_rhs!(du, u, semi, t)
    # Normal RHS evaluation
    Trixi.rhs!(du, u, semi, t)
    # Correct by subtracting the steady-state RHS
    Trixi.@trixi_timeit Trixi.timer() "rhs correction" begin
      # Use Trixi.@threaded for threaded performance
      Trixi.@threaded for i in eachindex(du)
        du[i] -= du_steady_state[i]
      end
    end
  end
end
u0 = compute_coefficients(tspan[1], semi)
ode = ODEProblem(corrected_rhs!, u0, tspan, semi)

summary_callback = SummaryCallback()

analysis_interval = 5000
analysis_callback = AnalysisCallback(semi, interval=analysis_interval)

alive_callback = AliveCallback(analysis_interval=analysis_interval)

save_solution = SaveSolutionCallback(interval=5000,
                                     save_initial_solution=true,
                                     save_final_solution=true,
                                     solution_variables=cons2prim)

callbacks = CallbackSet(summary_callback,
                        analysis_callback,
                        alive_callback,
                        save_solution)

###############################################################################
# run the simulation

# Use a Runge-Kutta method with automatic (error based) time step size control
# Enable threading of the RK method for better performance on multiple threads
sol = solve(ode, RDPK3SpFSAL49(thread=OrdinaryDiffEq.True()), abstol=1.0e-6, reltol=1.0e-6,
            save_everystep=false, callback=callbacks);

summary_callback() # print the timer summary


text:
export convexhull, convexhull!, conichull

const HAny{T} = Union{HRep{T}, HRepElement{T}}
const VAny{T} = Union{VRep{T}, VRepElement{T}}

"""
    intersect(P1::HRep, P2::HRep)

Takes the intersection of `P1` and `P2` ``\\{\\, x : x \\in P_1, x \\in P_2 \\,\\}``.
It is very efficient between two H-representations or between two polyhedron for which the H-representation has already been computed.
However, if `P1` (resp. `P2`) is a polyhedron for which the H-representation has not been computed yet, it will trigger a representation conversion which is costly.
See the [Polyhedral Computation FAQ](http://www.cs.mcgill.ca/~fukuda/soft/polyfaq/node25.html) for a discussion on this operation.

The type of the result will be chosen closer to the type of `P1`. For instance, if `P1` is a polyhedron (resp. H-representation) and `P2` is a H-representation (resp. polyhedron), `intersect(P1, P2)` will be a polyhedron (resp. H-representation).
If `P1` and `P2` are both polyhedra (resp. H-representation), the resulting polyhedron type (resp. H-representation type) will be computed according to the type of `P1`.
The coefficient type however, will be promoted as required taking both the coefficient type of `P1` and `P2` into account.
"""
function Base.intersect(p1::HRep, p2::HRep, ps::HRep...)
    p = (p1, p2, ps...)
    T = promote_coefficient_type(p)
    similar(p, hmap((i, x) -> convert(similar_type(typeof(x), T), x), FullDim(p[1]), T, p...)...)
end
Base.intersect(p::HRep, el::HRepElement) = p ∩ intersect(el)
Base.intersect(el::HRepElement, p::HRep) = p ∩ el

Base.intersect(h::HRepElement) = hrep([h])
Base.intersect(hp1::HyperPlane, hp2::HyperPlane, hps::HyperPlane...) = hrep([hp1, hp2, hps...])
Base.intersect(hs1::HalfSpace, hs2::HalfSpace, hss::HalfSpace...) = hrep([hs1, hs2, hss...])
Base.intersect(h1::HyperPlane{T}, h2::HalfSpace{T}) where {T} = hrep([h1], [h2])
function Base.intersect(h1::HyperPlane, h2::HalfSpace)
    T = promote_type(coefficient_type(h1), coefficient_type(h2))
    f(h) = convert(similar_type(typeof(h), T), h)
    intersect(f(h1), f(h2))
end
Base.intersect(h1::HalfSpace, h2::HyperPlane) = h2 ∩ h1
Base.intersect(p1::HAny, p2::HAny, ps::HAny...) = intersect(p1 ∩ p2, ps...)
#function Base.intersect(p::HAny...)
#    T = promote_type(coefficient_type.(p)...)
#    f(p) = convert(similar_type(typeof(p), T), p)
#    intersect(f.(p)...)
#end

"""
    intersect!(p::HRep, h::Union{HRepresentation, HRepElement})

Same as [`intersect`](@ref) except that `p` is modified to be equal to the intersection.
"""
function Base.intersect!(p::HRep, ::Union{HRepresentation, HRepElement})
    error("intersect! not implemented for $(typeof(p)). It probably does not support in-place modification, try `intersect` (without the `!`) instead.")
end
function Base.intersect!(p::Polyhedron, h::Union{HRepresentation, HRepElement})
    resethrep!(p, hrep(p) ∩ h)
end

"""
    convexhull(P1::VRep, P2::VRep)

Takes the convex hull of `P1` and `P2` ``\\{\\, \\lambda x + (1-\\lambda) y : x \\in P_1, y \\in P_2 \\,\\}``.
It is very efficient between two V-representations or between two polyhedron for which the V-representation has already been computed.
However, if `P1` (resp. `P2`) is a polyhedron for which the V-representation has not been computed yet, it will trigger a representation conversion which is costly.

The type of the result will be chosen closer to the type of `P1`. For instance, if `P1` is a polyhedron (resp. V-representation) and `P2` is a V-representation (resp. polyhedron), `convexhull(P1, P2)` will be a polyhedron (resp. V-representation).
If `P1` and `P2` are both polyhedra (resp. V-representation), the resulting polyhedron type (resp. V-representation type) will be computed according to the type of `P1`.
The coefficient type however, will be promoted as required taking both the coefficient type of `P1` and `P2` into account.
"""
function convexhull(p1::VRep, p2::VRep, ps::VRep...)
    p = (p1, p2, ps...)
    T = promote_coefficient_type(p)
    similar(p, vmap((i, x) -> convert(similar_type(typeof(x), T), x), FullDim(p[1]), T, p...)...)
end
convexhull(p::VRep, el::VRepElement) = convexhull(p, convexhull(el))
convexhull(el::VRepElement, p::VRep) = convexhull(p, el)

convexhull(v::VRepElement) = vrep([v])
convexhull(p1::AbstractVector, p2::AbstractVector, ps::AbstractVector...) = vrep([p1, p2, ps...])
convexhull(l1::Line, l2::Line, ls::Line...) = vrep([l1, l2, ls...])
convexhull(r1::Ray, r2::Ray, rs::Ray...) = vrep([r1, r2, rs...])
convexhull(p::AbstractVector{T}, r::Union{Line{T}, Ray{T}}) where {T} = vrep([p], [r])
function convexhull(p::AbstractVector, r::Union{Line, Ray})
    T = promote_type(coefficient_type(p), coefficient_type(r))
    f(x) = convert(similar_type(typeof(x), T), x)
    return convexhull(f(p), f(r))
end
convexhull(r::Union{Line, Ray}, p::AbstractVector) = convexhull(p, r)
convexhull(l::Line, r::Ray) = vrep([l], [r])
convexhull(r::Ray, l::Line) = convexhull(l, r)
convexhull(p1::VAny, p2::VAny, ps::VAny...) = convexhull(convexhull(p1, p2), ps...)
#function convexhull(p::VAny...)
#    T = promote_type(coefficient_type.(p)...)
#    f(p) = convert(similar_type(typeof(p), T), p)
#    convexhull(f.(p)...)
#end

"""
    convexhull!(p1::VRep, p2::VRep)

Same as [`convexhull`](@ref) except that `p1` is modified to be equal to the convex hull.
"""
function convexhull!(p::VRep, ine::VRepresentation)
    error("convexhull! not implemented for $(typeof(p)). It probably does not support in-place modification, try `convexhull` (without the `!`) instead.")
end
function convexhull!(p::Polyhedron, v::VRepresentation)
    resetvrep!(p, convexhull(vrep(p), v))
end

# conify: same than conichull except that conify(::VRepElement) returns a VRepElement and not a V-representation
conify(v::VRep) = vrep(lines(v), [collect(rays(v)); Ray.(collect(points(v)))])
conify(v::VCone) = v
conify(p::AbstractVector) = Ray(p)
conify(r::Union{Line, Ray}) = r

conichull(p...) = convexhull(conify.(p)...)

function sumpoints(::FullDim, ::Type{T}, p1, p2) where {T}
    _tout(p) = convert(similar_type(typeof(p), T), p)
    ps = [_tout(po1 + po2) for po1 in points(p1) for po2 in points(p2)]
    tuple(ps)
end
sumpoints(::FullDim, ::Type{T}, p1::Rep, p2::VCone) where {T} = change_coefficient_type.(preps(p1), T)
sumpoints(::FullDim, ::Type{T}, p1::VCone, p2::Rep) where {T} = change_coefficient_type.(preps(p2), T)

"""
    +(p1::VRep, p2::VRep)

Minkowski sum between `p1` and `p2` using the V-representation.
If the V-representation is not computed for `p1` or `p2`, it is computed.

    +(p::Rep, el::Union{Line, Ray})
    +(el::Union{Line, Ray}, p::Rep)

Same as `p + vrep([el])`.
"""
function Base.:+(p1::VRep{T1}, p2::VRep{T2}) where {T1, T2}
    T = typeof(zero(T1) + zero(T2))
    similar((p1, p2), FullDim(p1), T, sumpoints(FullDim(p1), T, p1, p2)..., change_coefficient_type.(rreps(p1, p2), T)...)
end
Base.:+(p::Rep, el::Union{Line, Ray}) = p + vrep([el])
Base.:+(el::Union{Line, Ray}, p::Rep) = p + el

# p1 has priority
function usehrep(p1::Polyhedron, p2::Polyhedron)
    hrepiscomputed(p1) && (!vrepiscomputed(p1) || hrepiscomputed(p2))
end

function hcartesianproduct(p1::HRep, p2::HRep)
    d = sum_fulldim(FullDim(p1), FullDim(p2))
    T = promote_coefficient_type((p1, p2))
    f = (i, x) -> zeropad(x, i == 1 ? FullDim(p2) : neg_fulldim(FullDim(p1)))
    function dimension_map(i)
        if i <= fulldim(p1)
            return (1, i)
        else
            return (2, i - fulldim(p1))
        end
    end
    similar((p1, p2), d, T, hmap(f, d, T, p1, p2)...;
            dimension_map = dimension_map)
end
function vcartesianproduct(p1::VRep, p2::VRep)
    d = sum_fulldim(FullDim(p1), FullDim(p2))
    T = promote_coefficient_type((p1, p2))
    # Always type of first arg
    f1 = (i, x) -> zeropad(x, FullDim(p2))
    f2 = (i, x) -> zeropad(x, neg_fulldim(FullDim(p1)))
    q1 = similar(p1, d, T, vmap(f1, d, T, p1)...)
    q2 = similar(p2, d, T, vmap(f2, d, T, p2)...)
    q1 + q2
end
cartesianproduct(p1::HRep, p2::HRep) = hcartesianproduct(p1, p2)
cartesianproduct(p1::VRep, p2::VRep) = vcartesianproduct(p1, p2)

function cartesianproduct(p1::Polyhedron, p2::Polyhedron)
    if usehrep(p1, p2)
        hcartesianproduct(p1, p2)
    else
        vcartesianproduct(p1, p2)
    end
end

"""
    *(p1::Rep, p2::Rep)

Cartesian product between the polyhedra `p1` and `p2`.
"""
Base.:(*)(p1::Rep, p2::Rep) = cartesianproduct(p1, p2)

"""
  \\(P::Union{AbstractMatrix, UniformScaling}, p::HRep)

Transform the polyhedron represented by ``p`` into ``P^{-1} p`` by transforming each halfspace ``\\langle a, x \\rangle \\le \\beta`` into ``\\langle P^\\top a, x \\rangle \\le \\beta`` and each hyperplane ``\\langle a, x \\rangle = \\beta`` into ``\\langle P^\\top a, x \\rangle = \\beta``.
"""
Base.:(\)(P::Union{AbstractMatrix, UniformScaling}, rep::HRep) = rep / P'

function linear_preimage_transpose(P, p::HRep{Tin}, d) where Tin
    f = (i, h) -> h / P
    T = _promote_type(Tin, eltype(P))
    return similar(p, d, T, hmap(f, d, T, p)...,
                   dimension_map = i -> nothing)
end

"""
    /(p::HRep, P::Union{AbstractMatrix, UniformScaling})

Transform the polyhedron represented by ``p`` into ``P^{-T} p`` by transforming each halfspace ``\\langle a, x \\rangle \\le \\beta`` into ``\\langle P a, x \\rangle \\le \\beta`` and each hyperplane ``\\langle a, x \\rangle = \\beta`` into ``\\langle P a, x \\rangle = \\beta``.
"""
function Base.:(/)(p::HRep, P::AbstractMatrix)
    if size(P, 2) != fulldim(p)
        throw(DimensionMismatch("The number of rows of P must match the dimension of the H-representation"))
    end
    # FIXME For a matrix P of StaticArrays, `size(P, 1)` should be type stable
    return linear_preimage_transpose(P, p, size(P, 1))
end
function Base.:(/)(p::HRep, P::UniformScaling)
    return linear_preimage_transpose(P, p, FullDim(p))
end

function linear_image(P, p::VRep{Tin}, d) where Tin
    f = (i, v) -> P * v
    T = _promote_type(Tin, eltype(P))
    return similar(p, d, T, vmap(f, d, T, p)...)
end

"""
    *(P::Union{AbstractMatrix, UniformScaling}, p::VRep)

Transform the polyhedron represented by ``p`` into ``P p`` by transforming each element of the V-representation (points, symmetric points, rays and lines) `x` into ``P x``.
"""
function Base.:(*)(P::AbstractMatrix, p::VRep)
    if size(P, 2) != fulldim(p)
        throw(DimensionMismatch("The number of rows of P must match the dimension of the V-representation"))
    end
    # FIXME For a matrix P of StaticArrays, `size(P, 1)` should be type stable
    return linear_image(P, p, size(P, 1))
end
function Base.:(*)(P::UniformScaling, p::VRep)
    return linear_image(P, p, FullDim(p))
end

"""
    *(α::Number, p::Rep)

Transform the polyhedron represented by ``p`` into ``\\alpha p`` by transforming
each element of the V-representation (points, symmetric points, rays and lines)
`x` into ``\\alpha x``.
"""
function Base.:(*)(α::Number, p::Polyhedron)
    if vrepiscomputed(p) || iszero(α)
        return (α * LinearAlgebra.I) * p
    else
        return p / (inv(α) * LinearAlgebra.I)
    end
end
function Base.:(*)(α::Number, p::VRepresentation)
    return (α * LinearAlgebra.I) * p
end
function Base.:(*)(α::Number, p::HRepresentation)
    return p / (inv(α) * LinearAlgebra.I)
end


text:
"""
    frule([::RuleConfig,] (Δf, Δx...), f, x...)

Expressing the output of `f(x...)` as `Ω`, return the tuple:

    (Ω, ΔΩ)

The second return value is the tangent w.r.t. the output.

If no method matching `frule((Δf, Δx...), f, x...)` has been defined, then return `nothing`.

Examples:

unary input, unary output scalar function:

```jldoctest frule
julia> dself = NoTangent();

julia> x = rand()
0.8236475079774124

julia> sinx, Δsinx = frule((dself, 1), sin, x)
(0.7336293678134624, 0.6795498147167869)

julia> sinx == sin(x)
true

julia> Δsinx == cos(x)
true
```

Unary input, binary output scalar function:

```jldoctest frule
julia> sincosx, Δsincosx = frule((dself, 1), sincos, x);

julia> sincosx == sincos(x)
true

julia> Δsincosx[1] == cos(x)
true

julia> Δsincosx[2] == -sin(x)
true
```

Note that techically speaking julia does not have multiple output functions, just functions
that return a single output that is iterable, like a `Tuple`.
So this is actually a [`Tangent`](@ref):
```jldoctest frule
julia> Δsincosx
Tangent{Tuple{Float64, Float64}}(0.6795498147167869, -0.7336293678134624)
```

The optional [`RuleConfig`](@ref) option allows specifying frules only for AD systems that
support given features. If not needed, then it can be omitted and the `frule` without it
will be hit as a fallback. This is the case for most rules.

See also: [`rrule`](@ref), [`@scalar_rule`](@ref), [`RuleConfig`](@ref)
"""
frule(ȧrgs, f, ::Vararg{Any}) = nothing

# if no config is present then fallback to config-less rules
frule(::RuleConfig, args...) = frule(args...)

# Manual fallback for keyword arguments. Usually this would be generated by
#
#   frule(::Any, ::Vararg{Any}; kwargs...) = nothing
#
# However - the fallback method is so hot that we want to avoid any extra code
# that would be required to have the automatically generated method package up
# the keyword arguments (which the optimizer will throw away, but the compiler
# still has to manually analyze). Manually declare this method with an
# explicitly empty body to save the compiler that work.
const frule_kwfunc = Core.kwftype(typeof(frule)).instance
(::typeof(frule_kwfunc))(::Any, ::typeof(frule), ȧrgs, f, ::Vararg{Any}) = nothing
function (::typeof(frule_kwfunc))(kws::Any, ::typeof(frule), ::RuleConfig, args...)
    return frule_kwfunc(kws, frule, args...)
end

"""
    rrule([::RuleConfig,] f, x...)

Expressing `x` as the tuple `(x₁, x₂, ...)` and the output tuple of `f(x...)`
as `Ω`, return the tuple:

    (Ω, (Ω̄₁, Ω̄₂, ...) -> (s̄elf, x̄₁, x̄₂, ...))

Where the second return value is the the propagation rule or pullback.
It takes in cotangents corresponding to the outputs (`x̄₁, x̄₂, ...`),
and `s̄elf`, the internal values of the function itself (for closures)

If no method matching `rrule(f, xs...)` has been defined, then return `nothing`.

Examples:

unary input, unary output scalar function:

```jldoctest
julia> x = rand();

julia> sinx, sin_pullback = rrule(sin, x);

julia> sinx == sin(x)
true

julia> sin_pullback(1) == (NoTangent(), cos(x))
true
```

binary input, unary output scalar function:

```jldoctest
julia> x, y = rand(2);

julia> hypotxy, hypot_pullback = rrule(hypot, x, y);

julia> hypotxy == hypot(x, y)
true

julia> hypot_pullback(1) == (NoTangent(), (x / hypot(x, y)), (y / hypot(x, y)))
true
```

The optional [`RuleConfig`](@ref) option allows specifying rrules only for AD systems that
support given features. If not needed, then it can be omitted and the `rrule` without it
will be hit as a fallback. This is the case for most rules.

See also: [`frule`](@ref), [`@scalar_rule`](@ref), [`RuleConfig`](@ref)
"""
rrule(::Any, ::Vararg{Any}) = nothing

# if no config is present then fallback to config-less rules
rrule(::RuleConfig, args...) = rrule(args...)

# Manual fallback for keyword arguments. See above
const rrule_kwfunc = Core.kwftype(typeof(rrule)).instance
(::typeof(rrule_kwfunc))(::Any, ::typeof(rrule), ::Any, ::Vararg{Any}) = nothing
function (::typeof(rrule_kwfunc))(kws::Any, ::typeof(rrule), ::RuleConfig, args...)
    return rrule_kwfunc(kws, rrule, args...)
end

##############################################################
### Opt out functionality

const NO_RRULE_DOC = """
    no_rrule

This is an piece of infastructure supporting opting out of [`rrule`](@ref).
It follows the signature for `rrule` exactly.
A collection of type-tuples is stored in its method-table.
If something has this defined, it means that it must having a must also have a `rrule`,
defined that returns `nothing`.

!!! warning "Do not overload no_rrule directly"
    It is fine and intended to query the method table of `no_rrule`.
    It is not safe to add to that directly, as corresponding changes also need to be made to
    `rrule`.
    The [`@opt_out`](@ref) macro does both these things, and so should almost always be used
    rather than defining a method of `no_rrule` directly.

### Mechanics
note: when the text below says methods `==` it actually means:
`parameters(m.sig)[2:end]` (i.e. the signature type tuple) rather than the method object `m` itself.

To decide if should opt-out using this mechanism.
 - find the most specific method of `rrule` and `no_rule` e.g with `Base.which`
  - if the method of `no_rrule` `==` the method of `rrule`, then should opt-out

To just ignore the fact that rules can be opted-out from, and that some rules thus return
`nothing`, then filter the list of methods of `rrule` to remove those that are `==` to ones
that occur in the method table of `no_rrule`.

Note also when doing this you must still also handle falling back from rule with config, to
rule without config.

On the other-hand if your AD can work with `rrule`s that return `nothing`, then it is
simpler to just use that mechanism for opting out; and you don't need to worry about this
at all.

For more information see the [documentation on opting out of rules](@ref opt_out)
"""

"""
$NO_RRULE_DOC

See also [`ChainRulesCore.no_frule`](@ref).
"""
function no_rrule end
no_rrule(::Any, ::Vararg{Any}) = nothing

"""
$(replace(NO_RRULE_DOC, "rrule"=>"frule"))

See also [`ChainRulesCore.no_rrule`](@ref).
"""
function no_frule end
no_frule(ȧrgs, f, ::Vararg{Any}) = nothing


text:
module SimpleNeuralNet

using LinearAlgebra
using Printf
using HDF5
using StatsBase
using Statistics

export neuralNet;
export nnForward;
export nnBackPropagate;
export nnTrain;
export nnComputeCosts;
export nnComputeOutputs;
export nnLoad;
export nnSave;

#useful functions
export sigmoid;
export dsigmoid;

#function stubs
export nnAct;
export nnDAct;
export nnCost;
export nnDCost;

include("neuralNet.jl");
include("nnForward.jl");
include("nnBackPropagate.jl");
include("nnTrain.jl");
include("nnComputeCosts.jl");
include("nnComputeOutputs.jl");
include("nnLoad.jl");
include("nnSave.jl");
include("nnStubs.jl");
include("nnUtils.jl");

end #end module

max_stars_repo_name: krometis/SimpleNeuralNet

max_issues_repo_name: krometis/SimpleNeuralNet

max_forks_repo_name: krometis/SimpleNeuralNet


text:
#########################################################################
##Define the function that extracts the first r factors of
##the data in the X matrix from time 1 to time t
#########################################################################
function get_factors(vec_î, X_lag, r)
    x_j = X_lag[1:vec_î, :]
    x_j_c = x_j .- mean(x_j, 1)
    # computing the weight matrix
    covx = cov(x_j_c)
    # W is the weight matrix
    F = eig(covx)[2]
    W = flipdim(F, 2)::Array{Float64, 2}
    res = (x_j_c * W)[:, 1:r]
    return res
end

#########################################################################
#Forecast using a dynamic factor model with lagged aggregate variables
#and lagged values of first "r" principal components (factors).
#########################################################################
function dfm_oos_forc(J, h, h_lag, Y, Y_lag, X_lag, r)
    ## Vector Index for oos forecast goes from J - (h_lag + h) to length(Y)
    ŷ = zeros(Float64, size(J - (h_lag + h):length(Y)))
    for (n, vec_î) in enumerate(J - (h_lag + h):length(Y))
        LHS = Y[1:(vec_î - h)]
        sample_size = length(LHS)
        factors = get_factors(vec_î, X_lag, r)
        RHS = [ones(sample_size) Y_lag[1:sample_size] factors[1:sample_size, :]]
        β̂ = OLSestimator(LHS, RHS)
        # Compute Forecasts
        RHS_crnt = [1; Y_lag[vec_î]; factors[vec_î, :]]
        ŷ[n] = dot(β̂, RHS_crnt)
    end
    RMSE = sqrt(mean((Y[J - (h_lag + h):length(Y)] - ŷ) .^ 2))
    return ŷ, RMSE
end

#########################################################################
#Forecast using a dynamic factor model with lagged aggregate variables
#and lagged values of first 1 principal components from
#1) the first order terms and
#2) the second order interactions terms.
#########################################################################
function dfm_oos_forc(J, h, h_lag, Y, Y_lag, X_lag, X_lag2, r1, r2)
    ## Vector Index for oos forecast goes from J - (h_lag + h) to length(Y)
    ŷ = zeros(Float64, size(J - (h_lag + h):length(Y)))
    for (n, vec_î) in enumerate(J - (h_lag + h):length(Y))
        println(n)
        LHS = Y[1:(vec_î - h)]
        sample_size = length(LHS)
        factors1 = get_factors(vec_î, X_lag, r1)
        factors2 = get_factors(vec_î, X_lag2, r2)
        RHS = [ ones(sample_size) Y_lag[1:sample_size] factors1[1:sample_size, :] factors2[1:sample_size, :] ]
        β̂ = OLSestimator(LHS, RHS)
        # Compute Forecasts
        RHS_crnt = [1; Y_lag[vec_î]; factors1[vec_î, :]; factors2[vec_î, :]]
        ŷ[n] = dot(β̂, RHS_crnt)
    end
    RMSE = sqrt(mean((Y[J - (h_lag + h):length(Y)] - ŷ) .^ 2))
    return ŷ, RMSE
end

max_stars_repo_name: jungangl/forecast-machine-learning

max_issues_repo_name: jungangl/forecast-machine-learning

max_forks_repo_name: jungangl/forecast-machine-learning


text:
module ConditionalMean

using Base: indices1, tail
export condmean, condvar, condstd
export nanmean, nanstd

"""
`M = condmean(A, cond, region)` calculates the mean value of A along the dimensions listed in `region`, ignoring values when cond is not true, or values equal to cond if cond is a number.
`M = condmean(f, A, cond, region)` calculates the mean of f(A).
"""
#condmean{T<:Real}(A::AbstractArray{T}, cond::Function, region=1) = _condmean(A, T, region)
# vectorizing mean(A[:]) is trivial, so keep default region definition of 1
condmean{T<:Number}(             A::AbstractArray{T}, cond::Function=(x->true), region=1; Km1::Bool=false) = _condmean(identity, A, T, cond      , region, Km1)
condmean{T<:Number}(             A::AbstractArray{T}, cond::Number            , region=1; Km1::Bool=false) = _condmean(identity, A, T, x->x!=cond, region, Km1)
condmean{T<:Number}(f::Function, A::AbstractArray{T}, cond::Function=(x->true), region=1; Km1::Bool=false) = _condmean(f       , A, T, cond      , region, Km1)
condmean{T<:Number}(f::Function, A::AbstractArray{T}, cond::Number            , region=1; Km1::Bool=false) = _condmean(f       , A, T, x->x!=cond, region, Km1)

# redirect _condmean to _condmeanoffset for Array{Float32} to improve precision
#_condmean{T<:Float32}(f::Function, A::AbstractArray, ::Type{T}, cond, region, Km1::Bool=false) = _condmeanoffset(f, A, T, cond, region, Km1)

# # general condition cond generalizes to all Number types, but behavior maybe unpredictable for some
# function _condmean{T<:Number}(f::Function, A::AbstractArray, ::Type{T}, cond, region, Km1::Bool=false)
#     inds = Base.reduced_indices(A, region)
#     K = similar(Array{Int}, inds)
#     S = similar(array{eltype(A)}, inds)
#     fill!(K,0)
#     fill!(S, zero(eltype(A)))
#     condsum!(S, K, A, f, cond)
#     if Km1
#        S./(K-1)
#     else
#        S./K
#     end
# end

# add offset to improve precision for floating point numerics, behavior unpredictable for other types
function _condmeanoffset{T<:Number}(f::Function, A::AbstractArray, ::Type{T}, cond, region, Km1::Bool=false)
    inds = Base.reduced_indices(A, region)
    K = similar(Array{Int}, inds)
    S = similar(array{eltype(A)}, inds)
    P = similar(array{eltype(A)}, inds)
    fill!(K,0)
    fill!(S, zero(eltype(A)))
    fill!(P, zero(eltype(A)))
    condsumoffset!(S, P, K, A, f, cond)
    if Km1
        (S + P.*K) ./ (K-1)
    else
        S./K + P
    end
end

# always redirect _condmean to _condmeanoffset to improve precision
_condmean=_condmeanoffset

using Base: check_reducedims, reducedim1, safe_tail
using Base.Broadcast: newindex

"""
    condsum!(S, K, A, f, cond)

Compute the sum `S` and number of contributing values `K` for
reductions of the array `A` over dimensions. `S` and `K` must have
identical indices, and either match `A` or have singleton-dimensions for
the dimensions that are being summed over. Only values x with cond(x)==true
are included in the tallies of `S` and `K`.

Note that the mean is just S./K.
"""
function condsum!{T,N}(S, K, A::AbstractArray{T,N}, f, cond)
    check_reducedims(S, A)
    isempty(A) && return S, K
    indices(S) == indices(K) || throw(DimensionMismatch("S and K must have identical indices"))

    indsAt, indsSt = safe_tail(indices(A)), safe_tail(indices(S))
    keep, Idefault = _newindexer(indsAt, indsSt)
    if reducedim1(S, A)
        # keep the accumulators as a local variable when reducing along the first dimension
        i1 = first(indices1(S))
        @inbounds for IA in CartesianRange(indsAt)
            IS = newindex(IA, keep, Idefault)
            s, k = S[i1,IS], K[i1,IS]
            for i in indices(A, 1)
                tmp = A[i, IA]
                if cond(tmp)
                    s += f(tmp)
                    k += 1
                end
            end
            S[i1,IS], K[i1,IS] = s, k
        end
    else
        @inbounds for IA in CartesianRange(indsAt)
            IS = newindex(IA, keep, Idefault)
            for i in indices(A, 1)
                tmp = A[i, IA]
                if cond(tmp)
                    S[i, IS] += f(tmp)
                    K[i, IS] += 1
                end
            end
        end
    end
    S, K
end
if VERSION < v"0.6.0-dev.693"
    _newindexer(shape, inds) = Base.Broadcast.newindexer(shape, inds)
else
    _newindexer(shape, inds) = Base.Broadcast.shapeindexer(shape, inds)
end

# add offset to improve precision for floating point numerics, behavior unpredictable for other types
"""
    condsumoffset!(S, P, K, A, f, cond)

Compute the adjusted sum `S` and number of contributing values `K` for
reductions of `f` mapped on array `A` over dimensions.
`S`, `P`, and `K` must have identical indices, and either match `A`
or have singleton-dimensions for the dimensions that are being summed
over. Only values of A with cond(A[i])==true are included in the tallies
of `S` and `K`. `P` is the first good value in each A set of `A` values
to be summed.

By subtracting off the first good value `P` accumulated to each sum,
condsumoffset! improves precision in case the values of A have a nearly
constant offset. The improved precision is beneficial if `K`×`P` is so large
that interesting variations of A are truncated in the sum S. However, precision
of this implementation will be worse than condsum! if the first good value is an
outlier farther from zero than the mean of the values of `A` summed.

The value mean computed in `condmeanprecise` is S./K+P.

See also `condsum!` and `condmeanoffset`.
"""
function condsumoffset!{T,N}(S, P, K, A::AbstractArray{T,N}, f, cond)
  #loops over dimensions of A _then_ determines dimensions of reduced S
    check_reducedims(S, A)
    isempty(A) && return S, K
    indices(S) == indices(K) || throw(DimensionMismatch("S and K must have identical indices"))

    indsAt, indsSt = safe_tail(indices(A)), safe_tail(indices(S))
    keep, Idefault = _newindexer(indsAt, indsSt)
    if reducedim1(S, A)
        # keep the accumulators as a local variable when reducing along the first dimension
        i1 = first(indices1(S))
        @inbounds for IA in CartesianRange(indsAt)
            IS = newindex(IA, keep, Idefault)
            s, p, k = S[i1,IS], P[i1,IS], K[i1,IS]
            #isfirst = true
            for i in indices(A, 1)
                tmp = A[i, IA]
                if cond(tmp)
                    if k==0 #isfirst
                        p = f(tmp)
                        #isfirst = false
                    else
                        s += f(tmp)-p
                    end
                    k += 1
                end
            end
            S[i1,IS], P[i1,IS], K[i1,IS] = s, p, k
        end
    else
        @inbounds for IA in CartesianRange(indsAt)
            IS = newindex(IA, keep, Idefault)
            #isfirst = true
            for i in indices(A, 1)
                tmp = A[i, IA]
                if cond(tmp)
                    if K[i, IS]==0 #isfirst
                        P[i, IS] = f(tmp)
                        #isfirst = false
                    else
                        S[i, IS] += f(tmp)-P[i, IS]
                    end
                    K[i, IS] += 1
                end
            end
        end
    end
    S, P, K
end
#=
if testing for k==0 is slow, then an alternative that works in many cases for
nearly constant offsets across the whole of A would be to just subtract off the
first good value in all of A. This would be good for surface pressure, but not
good for a significant gradient of A, like atmospheric pressure from the surface
to the stratosphere, where p varies by an order of magnitude.
Another solution, that I don't think will be faster, is to make a boolean array
`isfirst` like `K` but private to condsumoffset that stores whether the offset
has been recorded in `P`.
=#

# NOT TESTED
# precise, but requires two passes of condmean through A
# instead could write a version of condsum that accumulates multiple outputs of f, e.g. f(x)=[x, x*x]
"""
`condvar(A, cond, region; m=knownmean, corrected=true)` gives the unbiased
estimate of the variance of A (normalized by N-1) along dims in region.
`corrected=false` gives the variance normalized by N.
"""
condvar(A,cond,region; m=condmean(A,cond,region), corrected::Bool=true) =       condmean(x->abs(x)*abs(x), broadcast(-,A,m), cond, region, Km1=corrected)
condstd(A,cond,region; m=condmean(A,cond,region), corrected::Bool=true) = sqrt.(condmean(x->abs(x)*abs(x), broadcast(-,A,m), cond, region, Km1=corrected))
# corrected=true normalizes variance by K-1

"nanmean(X,r) mean of X long dimension(s) r, ignoring NaNs."
nanmean(X,r)=condmean(X,!(isnan),r)
"nanstd(X,r) standard deviation of X along dimension(s) r, ignoring NaNs."
nanstd(X,r)=sqrt.(condmean(x->abs(x)*abs(x),X.-nanmean(X,r),!(isnan),r))
end # module


text:
resultfile, target, = ARGS
using PkgBenchmark
mkpath(dirname(resultfile))
PkgBenchmark.benchmarkpkg(dirname(@__DIR__), target; resultfile = resultfile)


text:
# This file was generated, do not modify it. # hide
misclassification_rate(mode.(ŷ), ytrain)


text:
@testset verbose=true "Alignment" begin
    Random.seed!(1234)
    @testset verbose=true "Align QuatVec{$T}" for T in [Float16, Float32, Float64]
        for N in [1, 2, 3, 4, 5, 10, 20]
            a⃗ = randn(QuatVec{T}, N)
            R = randn(Rotor{T})

            # Test the exact result
            b⃗ = R .* a⃗ .* conj(R)
            R′ = align(a⃗, b⃗)
            if N > 1
                @test distance(R, conj(R′)) < 25eps(T)
            end
            @test maximum(abs, a⃗ - R′ .* b⃗ .* conj(R′)) < 40eps(T)
            @test_throws DimensionMismatch align(a⃗, b⃗[2:end])

            # Uniform weights
            w = 17ones(T, size(a⃗)) / T(3)
            R′′ = align(a⃗, b⃗, w)
            if N > 1
                @test distance(R, conj(R′′)) < 25eps(T)
            end
            @test maximum(abs, a⃗ - R′′ .* b⃗ .* conj(R′′)) < 40eps(T)
            @test_throws DimensionMismatch align(a⃗, b⃗, w[2:end])

            # Perturb b⃗ slightly
            δ = √eps(T)
            b⃗′′′ = [b⃗i + QuatVec((2*(rand(T, 3) .- 1/T(2)) * δ/√T(3))...) for b⃗i in b⃗]
            R′′′ = align(a⃗, b⃗′′′)
            if N > 1
                @test distance(R, conj(R′′′)) < 25δ
            end
            @test maximum(abs, a⃗ - R′′′ .* b⃗′′′ .* conj(R′′′)) < 40δ

            # Change first third, but use weights to ignore
            if N > 3
                N′ = N ÷ 3
                b⃗′ = copy(b⃗)
                w′ = copy(w)
                b⃗′[1:N′] = randn(QuatVec{T}, N′)
                w′[1:N′] .= 0
                R1 = align(a⃗, b⃗, w′)
                R2 = align(a⃗, b⃗′, w′)
                @test distance(R, conj(R1)) < 25eps(T)
                @test distance(R, conj(R2)) < 25eps(T)
                @test maximum(abs, (a⃗ - R1 .* b⃗ .* conj(R1))[N′+1:end]) < 40eps(T)
                @test maximum(abs, (a⃗ - R2 .* b⃗ .* conj(R2))[N′+1:end]) < 40eps(T)
            end
        end
    end

    @testset verbose=true "Align Rotor{$T}" for T in [Float16, Float32, Float64]
        for N in [1, 2, 3, 4, 5, 10, 20]
            A = randn(Rotor{T}, N)
            R = randn(Rotor{T})

            # Test the exact result
            B = R .* A
            R′ = align(A, B)
            if N > 1
                @test distance(R, conj(R′)) < 25eps(T)
            end
            @test maximum(abs, A - R′ .* B) < 40eps(T)
            @test_throws DimensionMismatch align(A, B[2:end])

            # Uniform weights
            w = 17ones(T, size(A)) / T(3)
            R′′ = align(A, B, w)
            if N > 1
                @test distance(R, conj(R′′)) < 25eps(T)
            end
            @test maximum(abs, A - R′′ .* B) < 40eps(T)
            @test_throws DimensionMismatch align(A, B, w[2:end])

            # Perturb B slightly
            δ = √eps(T)
            B′′′ = [Rotor(Bi + Quaternion((2*(rand(T, 4) .- 1/T(2)) * δ/√T(3))...)) for Bi in B]
            R′′′ = align(A, B′′′)
            if N > 1
                @test distance(R, conj(R′′′)) < 25δ
            end
            @test maximum(abs, A - R′′′ .* B′′′) < 40δ

            # Change first third, but use weights to ignore
            if N > 3
                N′ = N ÷ 3
                B′ = copy(B)
                w′ = copy(w)
                B′[1:N′] = randn(Rotor{T}, N′)
                w′[1:N′] .= 0
                R1 = align(A, B, w′)
                R2 = align(A, B′, w′)
                @test distance(R, conj(R1)) < 25eps(T)
                @test distance(R, conj(R2)) < 25eps(T)
                @test maximum(abs, (A - R1 .* B)[N′+1:end]) < 40eps(T)
                @test maximum(abs, (A - R2 .* B)[N′+1:end]) < 40eps(T)
            end
        end
    end
end


text:
using BandwidthBenchmark

bwbench(; verbose=true, write_allocate=true)

max_stars_repo_name: carstenbauer/BandwidthBenchmark

max_issues_repo_name: carstenbauer/BandwidthBenchmark

max_forks_repo_name: carstenbauer/BandwidthBenchmark


text:
### load a few 700 MHz BMRB experiments, each with a known concentration.
# combine them into a single FID.

# Given multiple experiments, this package also additively (in a weighted manner for each experiment) combines the experiments to get a single experiment. It is up to the user to ensure that the input experiments have the same `fs` and `SW`, otherwise the output won't be a realistic NMR spectrum. This combination is referred to as weighted mixing.

import NMRData

#include("../src/NMRDataSetup.jl")
import NMRDataSetup

using FFTW
import PyPlot
import BSON

PyPlot.close("all")
fig_num = 1

PyPlot.matplotlib["rcParams"][:update](["font.size" => 22, "font.family" => "serif"])

# user inputs.
BMRB_setup_label = "BMRB-700" # for use with the NMRData package. See that package's readme for details.

save_dir = "/home/roy/MEGAsync/outputs/NMR/combined/"
project_name = "cal01"

target_names = ["D-(+)-Glucose"; # compounds you wish to add to the mixture.
"L-Alanine"]

weights = ones(Float64, length(target_names)) # mixing weights for each experiment.

solvent_ppm_guess = 4.7 # guess where the solvent ppm is for each BMRB experiment.
solvent_window_ppm = 0.1 # allowed solvent shift, in ppm.

### end inputs.

## get BMRB experiment paths from the NMRData package. You can specify your own array of experiments to mix.
paths = Vector{String}(undef, 0)
for i = 1:length(target_names)
    push!( paths, NMRData.getBMRBexperimentpath(BMRB_setup_label, target_names[i])[1] )
end
@assert all(isdir.(paths)) # check for bad paths.

## mix.
s_t, ν_0ppm, fs, SW,
s_t_set, ν_0ppm_set, fs_set, SW_set = NMRDataSetup.assemblemixture(paths; weights = weights)

## store.
save_folder_path = joinpath(save_dir, project_name)
isdir(save_folder_path) || mkdir(save_folder_path); # make save folder if it doesn't exist.

save_path = joinpath(save_folder_path, "experiment.bson")
BSON.bson(save_path,
s_t = s_t,
molecule_names = target_names,
fs = fs,
SW = SW,
ν_0ppm = ν_0ppm,
s_t_set = s_t_set,
ν_0ppm_set = ν_0ppm_set,
fs_set = fs_set,
SW_set = SW_set)

#### plot.

function getDFTfreqrange(N::Int, fs::T)::LinRange{T} where T
    a = zero(T)
    b = fs-fs/N

    return LinRange(a, b, N)
end

# Assumes U0 is sorted in ascending order.
function wrapfreqrange(U0, ν_begin::T, fs::T) where T <: Real

    N = length(U0)
    ind = findfirst(xx->(xx>ν_begin), U0)

    # TODO handle these exceptions with more grace.
    @assert typeof(ind) == Int
    @assert ind <= N

    out = Vector{T}(undef, N)
    #out[1:ind] = U0[1:ind] .+ fs
    #out[ind+1:end] = U0[ind+1:end]

    M = N-ind
    out[1:M] = U0[ind+1:end]
    out[M+1:end] = U0[1:ind] .+ fs

    inds = collect(1:N)
    inds[1:M] = collect(ind+1:N)
    inds[M+1:end] = collect(1:ind)

    return out, inds
end

function getwraparoundDFTfreqs(N::Int, fs::T, ν_begin::T) where T

    U0 = getDFTfreqrange(N, fs)
    out, inds = wrapfreqrange(U0, ν_begin, fs)

    return U0, out, inds
end

hz2ppmfunc = uu->(uu - ν_0ppm)*SW/fs
ppm2hzfunc = pp->(ν_0ppm + pp*fs/SW)

offset_Hz = ν_0ppm - (ppm2hzfunc(0.3)-ppm2hzfunc(0.0))

N = length(s_t)
DFT_s = fft(s_t)
U_DFT, U, U_inds = getwraparoundDFTfreqs(N, fs, offset_Hz)
y = DFT_s[U_inds]

##############################

S_U = DFT_s[U_inds]
P = hz2ppmfunc.(U)

S_set_U = collect( (fft(s_t_set[n]) ./fs)[U_inds] for n = 1:length(s_t_set))
#P_cost = hz2ppmfunc.(U_cost)

PyPlot.figure(fig_num)
fig_num += 1

PyPlot.plot(P, real.(S_U), label = "combined spectrum")

PyPlot.legend()
PyPlot.xlabel("ppm")
PyPlot.ylabel("real")
PyPlot.title("simulated spectra")