text:
module Trace

using Memento
using Humanize

import Memento: Attribute

global ENABLED = false

export
    @log,
    @trace,
    @debug,
    @info,
    @notice,
    @warn,
    @error,
    @critical,
    @alert,
    @emergency

"""
    Trace.enable()

Enables logging for all subsequent tracing macros and adds a "trace" (5) logging level.
"""
function enable()
    ENABLED::Bool &&  return  # exit early if enabled has already been set
    global ENABLED = true
    add_level(get_logger(), "trace", 5)
end

include(joinpath(Pkg.dir("Trace"), "src", "record.jl"))
include(joinpath(Pkg.dir("Trace"), "src", "macros.jl"))

end  # module


text:
const STB = ScientificTypesBase

# aliases not exported:
const PlainNGram{N}  = NTuple{N,<:AbstractString}
const TaggedNGram{N} = NTuple{N,<:CL.TaggedWord}

# This can be made less of a hack once ScientificTypes #155 is sorted.

type2scitype(T::Type) = STB.Scitype(T, DefaultConvention())
type2scitype(::Type{<:AbstractVector{T}}) where T =
    AbstractVector{type2scitype(T)}
type2scitype(::NTuple{N,T}) where {N,T} = NTuple{type2scitype{T}}

STB.scitype(::CL.TaggedWord, ::DefaultConvention) = Annotated{Textual}
STB.scitype(::CL.Document{<:AbstractVector{T}}, ::DefaultConvention) where T =
    Annotated{AbstractVector{type2scitype(T)}}
STB.scitype(::AbstractDict{<:AbstractString,<:Integer},
           ::DefaultConvention) = Multiset{Textual}
STB.scitype(::AbstractDict{<:CL.TaggedWord,<:Integer},
           ::DefaultConvention) = Multiset{Annotated{Textual}}
STB.scitype(::AbstractDict{<:Union{CL.TaggedWord,AbstractString},<:Integer},
           ::DefaultConvention) =
               Multiset{Union{Textual,Annotated{Textual}}}
STB.scitype(::AbstractDict{<:PlainNGram{N}}, ::DefaultConvention) where N =
    Multiset{NTuple{N,Textual}}
STB.scitype(::AbstractDict{<:TaggedNGram{N}}, ::DefaultConvention) where N =
    Multiset{NTuple{N,Annotated{Textual}}}
STB.scitype(::AbstractDict{<:PlainNGram}, ::DefaultConvention) =
    Multiset{NTuple{<:Any,Textual}}
STB.scitype(::AbstractDict{<:TaggedNGram}, ::DefaultConvention) =
    Multiset{NTuple{<:Any,Annotated{Textual}}}

STB.Scitype(::Type{<:CL.TaggedWord}, ::DefaultConvention) =
    Annotated{Textual}

ext: jl


text:
local_dir = "/home/awise/Desktop/RvSpectML"
#local_dir = "/storage/work/afw5465/RvSpectML"

using Pkg
 Pkg.activate(joinpath(local_dir,"alex_sandbox/"))
 using Query, DataFrames, DataStructures, CSV, JLD, Statistics, StatsBase, Plots, JLD2, FileIO, ArgParse
 using RvSpectMLBase, EchelleInstruments, EchelleCCFs, RvSpectML, Scalpels

 include(joinpath(local_dir,"data_paths.jl"))
 include(joinpath(local_dir,"RvLineList/inputs/param.jl"))
 include(joinpath(local_dir,"RvLineList/src/masks.jl"))

 s = ArgParseSettings()

 @add_arg_table s begin
    "--norm_type"
       help = "normalization type. Options: :raw, :continuum, :blaze"
       arg_type = Symbol
    "--EXPRES_excalibur_only"
       help = "whether or not to only include excalibur wavelengths in expres data - could be replaced with orders_to_use param here"
       arg_type = Bool
   "--mask_type"
      help = "mask shape. Options: :tophat, :gaussian, :supergaussian, :halfcos"
      arg_type = Symbol
   "--mask_scale_factor"
      help = "approximate number of pixels in the width scale of the mask"
      arg_type = Float64
   "--overlap_cutoff"
      help = "distance between lines required for them to be categorized as a blend, expressed as a fraction of the speed of light x 10^-5"
      arg_type = Int
   "--quant"
      help = "quantile for stability of fit params for empirical masks"
      arg_type = String
   "--nbin"
      help = "number of mask subdivisions to make"
      arg_type = Int
   "--bin_n"
      help = "which subdivision from nbin to use - one is first as this is a julia index"
      arg_type = Int
 end

 parsed_args = parse_args(s)

if parsed_args["norm_type"] != nothing
   println("Found command-line arg norm_type. Overwriting param file definition of this arg.")
   global norm_type = parsed_args["norm_type"]
end

if parsed_args["EXPRES_excalibur_only"] != nothing
   println("Found command-line arg EXPRES_excalibur_only. Overwriting param file definition of this arg.")
   global EXPRES_excalibur_only = parsed_args["EXPRES_excalibur_only"]
end

if parsed_args["mask_type"] != nothing
   println("Found command-line arg mask_type. Overwriting param file definition of this arg.")
   global mask_type = parsed_args["mask_type"]
end

if parsed_args["mask_scale_factor"] != nothing
   println("Found command-line arg mask_scale_factor. Overwriting param file definition of this arg.")
   global mask_scale_factor = parsed_args["mask_scale_factor"]
end

if parsed_args["overlap_cutoff"] != nothing
   println("Found command-line arg overlap_cutoff. Overwriting param file definition of this arg.")
   global overlap_cutoff = parsed_args["overlap_cutoff"]
end

if parsed_args["quant"] != nothing
   println("Found command-line arg quant. Overwriting param file definition of this arg.")
   global quant = parsed_args["quant"]
end

if parsed_args["nbin"] != nothing
   println("Found command-line arg nbin. Overwriting param file definition of this arg.")
   global nbin = parsed_args["nbin"]
end

if parsed_args["bin_n"] != nothing
   println("Found command-line arg bin_n. Overwriting param file definition of this arg.")
   global bin_n = parsed_args["bin_n"]
end

function hasOneToOneMatch(X::A1,Y::A2) where { T1<:Real, A1<:AbstractArray{T1,1}, T2<:Real, A2<:AbstractArray{T2,1}}
   xTrue = zeros(Bool,size(X,1))
   for y in 1:length(Y)
      n_matches = 0
      match_i = -1
      for x in 1:length(X)
         if (abs(X[x]-Y[y])/X[x])<1e-5
            n_matches += 1
            match_i = x
         end
      end
      if n_matches == 1
         xTrue[match_i]=true
      end
   end
   return xTrue
end

#norm_type_results = Dict()

ccfs = Dict()

v_grids = Dict()

pipeline_plans = Dict()

rvs = Dict()

times = Dict()

targs = ["101501", "26965", "34411", "10700"]

for targ in targs

   GC.gc()

   global target = targ

   maskWavelengths = string("Reiners",target) #keyword for mask wavelengths - should be e.g. Reiners or Reiners101501
   badLineFilter = string("HD",target)

   linelist_names = SortedDict(Dict())

   push!(linelist_names, "empirical_mask"=>string(target,"_good_lines_fit_quant=0.",quant,"_frac=",min_frac_converged,"_fmt=espresso.mas"))

   VALDmaskStr = string("VALD", depthPercentile ? "" : "_nonDP", "_species=", iron1Only, "_depthcutoff=", depth_cutoff, "_overlapcutoff=", overlap_cutoff, "e-05_allowBlends=", allowBlends, "_badLineFilter=", badLineFilter, "_rejectTelluricSlope=", rejectTelluricSlope, "_waves=", maskWavelengths, "_nbin=", 1, "_binParam=", binParam, "_n=", 0, ".mas")
   push!(linelist_names, "vald_mask"=>VALDmaskStr)

   push!(linelist_names, "espresso"=>"G2.espresso.mas")

   global fits_target_str = target
   verbose=true
   #path_to_EXPRES_data = joinpath("/home/awise/data/",target)
   path_to_EXPRES_data = joinpath(expres_data_path,target)
   default_paths_to_search = [local_dir]
   pipeline_plan = PipelinePlan()
   dont_make_plot!(pipeline_plan, :movie)

   reset_all_needs!(pipeline_plan)
   if need_to(pipeline_plan,:read_spectra)
      if verbose println("# Finding what data files are avaliable.")  end
      global df_files = EXPRES.make_manifest(path_to_EXPRES_data, fits_target_str)

      if verbose println("# Reading in customized parameters from param.jl.")  end
      include(joinpath(local_dir,"param.jl"))

      if verbose println("# Reading in ", size(df_files_use,1), " FITS files.")  end
      @time all_spectra = map(row->EXPRES.read_data(row,store_min_data=true, store_tellurics=true, normalization=norm_type), eachrow(df_files_use))
      #for i in 1:length(all_spectra)
      #  all_spectra[i].metadata[:bjd] = all_spectra[i].metadata[:BARYMJD]
      #end
      #RvSpectMLBase.discard_blaze(all_spectra)
      RvSpectMLBase.discard_continuum(all_spectra)
      dont_need_to!(pipeline_plan,:read_spectra)
   end

   # ccf_mid_velocity = 0.0 #since the empirical mask is at v=0, and we shifted the VALD masks to be at v=0 too, we must reset this global variable to zero.

   pixels_to_use = map(ord->min_col_default(all_spectra[1].inst,ord):max_col_default(all_spectra[1].inst,ord),orders_to_use)
   order_list_timeseries = extract_orders(all_spectra,pipeline_plan, orders_to_use=orders_to_use);
   n_spectra = length(order_list_timeseries)
   n_orders = length(orders_to_use)

   ### NEW RUN: empirical mask
   # ccf_mid_velocity = 0.0 #since the empirical mask is at v=0, and we shifted the VALD masks to be at v=0 too, we must reset this global variable to zero.

   ll = string("empirical_mask")
   #ll = "espresso"
   ll_v = string("vald_mask")

   line_width_50 = 7000.

   need_to!(pipeline_plan,:read_line_list)
   need_to!(pipeline_plan,:clean_line_list_tellurics)
   empty!(pipeline_plan.cache)
   println(linelist_names[ll])
   linelist_path = joinpath(local_dir,"RvLineList","inputs","empirical_masks",linelist_names[ll])
   mask_ll = prepare_line_list(linelist_path, all_spectra, pipeline_plan,  v_center_to_avoid_tellurics=0.0,
      Δv_to_avoid_tellurics = 1*RvSpectMLBase.max_bc+5*line_width_50+2*default_ccf_mask_v_width(EXPRES1D()),  orders_to_use=orders_to_use, verbose=true,
      convert_air_to_vacuum=false)

   need_to!(pipeline_plan,:read_line_list)
   need_to!(pipeline_plan,:clean_line_list_tellurics)
   empty!(pipeline_plan.cache)
   println(linelist_names[ll_v])
   linelist_path = joinpath(local_dir,"RvLineList","outputs","masks",linelist_names[ll_v])
   mask_ll_v = prepare_line_list(linelist_path, all_spectra, pipeline_plan,  v_center_to_avoid_tellurics=0.0,
      Δv_to_avoid_tellurics = 1*RvSpectMLBase.max_bc+5*line_width_50+2*default_ccf_mask_v_width(EXPRES1D()),  orders_to_use=orders_to_use, verbose=true,
      convert_air_to_vacuum=false)

   nf = length(mask_ll.lambda)
   # hm = zeros(Bool,nf)
   # for i in 1:nf
   #   hm[i] = hasOneMatch(mask_ll.lambda[i],mask_ll_v.lambda)
      #hm[i] = hasOneMatch(linelists[ll].lambda[i]*(1.0-5e3/3e8),linelists[ll_v].lambda)
   hm = hasOneToOneMatch(mask_ll.lambda,mask_ll_v.lambda)
   #end
   mask_empirical_total = mask_ll[hm,:]

   masks_empirical = binMask(mask_empirical_total,nbin)
   mask_empirical = masks_empirical[bin_n]

   println("# mask_scale_factor = ", mask_scale_factor, ", frac_width_to_fit =  ", fwtf, ". ")
   @time (ccfs_empirical, v_grid) = ccf_total(order_list_timeseries, mask_empirical, pipeline_plan,  mask_scale_factor=mask_scale_factor, range_no_mask_change=5*line_width_50, ccf_mid_velocity=0.0, v_step=100,
      #v_max=RvSpectMLBase.max_bc,
      v_max= 0*RvSpectMLBase.max_bc+5*line_width_50, #+2*line_width_50,
      calc_ccf_var=false, recalc=true, mask_type = mask_type)
   push!(ccfs, string("ccfs_empirical_",target) => ccfs_empirical)
   push!(v_grids, string("v_grid_empirical_",target) => v_grid)
   push!(pipeline_plans, string("pipeline_plan_empirical_",target) => pipeline_plan)
   #if nbin == 1 #if nbin > 1, rv fits sometimes fail
   alg_fit_rv =  EchelleCCFs.MeasureRvFromCCFGaussian(frac_of_width_to_fit=fwtf)
   rvs_ccf_empirical = RvSpectML.calc_rvs_from_ccf_total(ccfs_empirical, pipeline_plan, v_grid=v_grid, times = order_list_timeseries.times, recalc=true, bin_nightly=true, alg_fit_rv=alg_fit_rv)
   push!(rvs, string("rvs_empirical_",target) => rvs_ccf_empirical)
   #end

   k="espresso"
   need_to!(pipeline_plan,:read_line_list)
   need_to!(pipeline_plan,:clean_line_list_tellurics)
   empty!(pipeline_plan.cache)
   println(linelist_names[k])
   linelist_path = joinpath("alex_sandbox","masks",linelist_names[k])
   line_list_df = prepare_line_list(linelist_path, all_spectra, pipeline_plan, v_center_to_avoid_tellurics=ccf_mid_velocity,
   Δv_to_avoid_tellurics = 2*RvSpectMLBase.max_bc+3*lsf_width, verbose=true)

   masks_espresso = binMask(line_list_df,nbin)
   mask_espresso = masks_espresso[bin_n]

   println("# mask_scale_factor = ", mask_scale_factor, ", frac_width_to_fit =  ", fwtf, ". ")
   @time (ccfs_espresso, v_grid) = ccf_total(order_list_timeseries, mask_espresso, pipeline_plan,  mask_scale_factor=mask_scale_factor, range_no_mask_change=5*line_width_50, ccf_mid_velocity=ccf_mid_velocity, v_step=100,
      #v_max=RvSpectMLBase.max_bc,
      v_max= 0*RvSpectMLBase.max_bc+5*line_width_50, #+2*line_width_50,
      calc_ccf_var=false, recalc=true, mask_type = mask_type)
   push!(ccfs, string("ccfs_espresso_",target) => ccfs_espresso)
   push!(v_grids, string("v_grid_espresso_",target) => v_grid)
   push!(pipeline_plans, string("pipeline_plan_espresso_",target) => pipeline_plan)
   #if nbin == 1 #if nbin > 1, rv fits sometimes fail
   alg_fit_rv =  EchelleCCFs.MeasureRvFromCCFGaussian(frac_of_width_to_fit=fwtf)
   rvs_ccf_espresso = RvSpectML.calc_rvs_from_ccf_total(ccfs_espresso, pipeline_plan, v_grid=v_grid, times = order_list_timeseries.times, recalc=true, bin_nightly=true, alg_fit_rv=alg_fit_rv)
   push!(rvs, string("rvs_espresso_",target) => rvs_ccf_espresso)
   #end

   times[target] = order_list_timeseries.times

end

GC.gc()

save(joinpath(local_dir,"RvLineList","outputs",string("param_study_results","_norm_type=",norm_type,
   "_EXPRES_excalibur_only=",EXPRES_excalibur_only,"_mask_type=",mask_type,"_mask_scale_factor=",mask_scale_factor,
   "_overlap_cutoff=",overlap_cutoff,"_quant=",quant,"_nbin=",nbin,"_bin_n=",bin_n,"_orders_to_use=",join(orders_to_use),".jld")),
   "ccfs",ccfs,"v_grids",v_grids,"pipeline_plans",pipeline_plans,"rvs",rvs,"times",times,"orders_to_use",orders_to_use)

"""
sample job submission scripts:

for norm_type in raw blaze continuum; do for EXPRES_excalibur_only in true false; do for mask_type in gaussian tophat; do for mask_scale_factor in 1 2 3 4 5 6 7 8; do for overlap_cutoff in 2; do qsub -v norm_type=$norm_type,EXPRES_excalibur_only=$EXPRES_excalibur_only,mask_type=$mask_type,mask_scale_factor=$mask_scale_factor,overlap_cutoff=$overlap_cutoff expres_mask_param_study.pbs; done; done; done; done; done

for norm_type in raw continuum; do for EXPRES_excalibur_only in true; do for mask_type in gaussian; do for mask_scale_factor in 0.5 1 1.5; do for overlap_cutoff in 1 2 3 4 6; do for quant in 50 70 90; do for nbin in 1; do for bin_n in 0; do qsub -v norm_type=$norm_type,EXPRES_excalibur_only=$EXPRES_excalibur_only,mask_type=$mask_type,mask_scale_factor=$mask_scale_factor,overlap_cutoff=$overlap_cutoff,quant=$quant,nbin=$nbin,n=$n expres_mask_param_study.pbs; done; done; done; done; done; done; done; done

for norm_type in raw continuum; do for EXPRES_excalibur_only in true; do for mask_type in gaussian; do for mask_scale_factor in 1 4 8; do for overlap_cutoff in 1 6; do for quant in 90; do for nbin in 8; do for bin_n in 1 2 3 4 5 6 7 8; do qsub -v norm_type=$norm_type,EXPRES_excalibur_only=$EXPRES_excalibur_only,mask_type=$mask_type,mask_scale_factor=$mask_scale_factor,overlap_cutoff=$overlap_cutoff,quant=$quant,nbin=$nbin,n=$n expres_mask_param_study.pbs; done; done; done; done; done; done; done; done

for i in $(seq 1 50); do for norm_type in raw continuum; do for EXPRES_excalibur_only in true; do for mask_type in gaussian; do for mask_scale_factor in 1; do for overlap_cutoff in 1; do for quant in 90; do for nbin in 1; do for bin_n in 1; do qsub -v norm_type=$norm_type,EXPRES_excalibur_only=$EXPRES_excalibur_only,mask_type=$mask_type,mask_scale_factor=$mask_scale_factor,overlap_cutoff=$overlap_cutoff,quant=$quant,nbin=$nbin,n=$n expres_mask_param_study.pbs; done; done; done; done; done; done; done; done; done

"""

max_stars_repo_name: alexander-wise/RvLineList

max_issues_repo_name: alexander-wise/RvLineList

max_forks_repo_name: alexander-wise/RvLineList


text:
module LinTrigInterpMod

import FEM: AbstractInterpolator, Point2, FENode2, Vertex3, inv2x2t!, det2x2,
       dNmatrix, Jmatrix, dNdxmatrix, Nvec, get_area, mass_matrix

export LinTrigInterp

immutable LinTrigInterp <: AbstractInterpolator
    N::Vector{Float64}
    dN::Matrix{Float64}
    dNdx::Matrix{Float64}
    J::Matrix{Float64}
    ref_M::Matrix{Float64}
    M::Matrix{Float64}
end

function LinTrigInterp()
    N = Array(Float64, 3)

    # Derivatives are constant
    #       ξ     η
    dN = [[ 1.0  0.0];
          [ 0.0  1.0];
          [-1.0 -1.0]]

    dNdx = Array(Float64, 3, 2)

    J = Array(Float64, 2, 2)

    ref_M = 1/12 *
             [2.0 1.0 1.0;
              1.0 2.0 1.0;
              1.0 1.0 2.0]

    M = zeros(6,6)

    LinTrigInterp(N, dN, dNdx, J, ref_M, M)
end
get_allocated_N(i::LinTrigInterp) = i.N
get_allocated_dN(i::LinTrigInterp) = i.dN
get_allocated_dNdx(i::LinTrigInterp) = i.dNdx
get_allocated_J(i::LinTrigInterp) = i.J

# Shape functions in local coords
function Nvec(interp::LinTrigInterp, loc_coords::Point2)

    ξ = loc_coords[1]
    η = loc_coords[2]

    N = get_allocated_N(interp)

    N[1] = ξ
    N[2] = η
    N[3] = 1.0 - ξ - η

    return N
end

function dNmatrix(interp::LinTrigInterp, ::Point2)
    # dN is constant so just return it
    return get_allocated_dN(interp)
end

function Jmatrix(interp::LinTrigInterp,
                 vertices::Vertex3, nodes::Vector{FENode2},
                 ::Matrix{Float64})

    J = get_allocated_J(interp)

    x1 =  nodes[vertices[1]].coords.x
    x2 =  nodes[vertices[2]].coords.x
    x3 =  nodes[vertices[3]].coords.x

    y1 =  nodes[vertices[1]].coords.y
    y2 =  nodes[vertices[2]].coords.y
    y3 =  nodes[vertices[3]].coords.y

    # Constant Jacobian
    J[1, 1] = x1 - x3
    J[2, 1] = x2 - x3
    J[1, 2] = y1 - y3
    J[2, 2] = y2 - y3

    return J
end

function dNdxmatrix(interp::LinTrigInterp, local_coords::Point2,
                    vertices::Vertex3, nodes::Vector{FENode2})

        dN = dNmatrix(interp, local_coords)
        J = Jmatrix(interp, vertices, nodes, dN)
        J_inv = inv2x2t!(J)
        A_mul_B!(interp.dNdx, dN, J_inv)
        return interp.dNdx
end

function get_area(::LinTrigInterp, vertices::Vertex3, nodes::Vector{FENode2})

    x1 =  nodes[vertices[1]].coords.x
    x2 =  nodes[vertices[2]].coords.x
    x3 =  nodes[vertices[3]].coords.x

    y1 =  nodes[vertices[1]].coords.y
    y2 =  nodes[vertices[2]].coords.y
    y3 =  nodes[vertices[3]].coords.y

    area = 0.5 * (x1*(y2 - y3) + x2*(-y1 + y3) + x3*(y1 - y2))
    return area
end

function mass_matrix(interp::LinTrigInterp, vertices::Vertex3, nodes::Vector{FENode2})
    fill!(interp.M2, 0.0)
    area = get_area(interp, vertices, nodes)
    scale!(interp.M2, interp.ref_M2, area)
    return interp.M2
end

function mass_matrix_big(interp::LinTrigInterp, vertices::Vertex3, nodes::Vector{FENode2})
    fill!(interp.M, 0.0)
    area = get_area(interp, vertices, nodes)
    scale!(interp.M, interp.ref_M, area)
    return interp.M
end

end # module


text:
# Utility functions for plotting.

"Prepare 2D data arrays for passing to plotting functions."
function getdata(data::Data, var::AbstractString, plotrange, plotinterval; griddim=1,
   innermask=false)
   @assert data.head.ndim == 2 "data must be in 2D!"

   x, w = data.x, data.w
   ndim = data.head.ndim
   varIndex_ = findindex(data, var)

   if data.head.gencoord # Generalized coordinates
      X = @view x[:,:,1]
      Y = @view x[:,:,2]
      W = @view w[:,:,varIndex_]

      if any(abs.(plotrange) .== Inf)
         if plotrange[1] == -Inf plotrange[1] = minimum(X) end
         if plotrange[2] ==  Inf plotrange[2] = maximum(X) end
         if plotrange[3] == -Inf plotrange[3] = minimum(Y) end
         if plotrange[4] ==  Inf plotrange[4] = maximum(Y) end
      end

      # Create grid values first.
      xi = range(plotrange[1], stop=plotrange[2], step=plotinterval)
      yi = range(plotrange[3], stop=plotrange[4], step=plotinterval)
      # Perform linear interpolation of the data (x,y) on grid(xi,yi)
      triang = @views matplotlib.tri.Triangulation(X[:], Y[:])
      interpolator = @views matplotlib.tri.LinearTriInterpolator(triang, W[:])
      Xi, Yi = meshgrid(xi, yi)
      Wi = interpolator(Xi, Yi)
   else # Cartesian coordinates
      xrange = range(x[1,1,1], x[end,1,1], length=size(x,1))
      yrange = range(x[1,1,2], x[1,end,2], length=size(x,2))
      if all(isinf.(plotrange))
         xi, yi = xrange, yrange
         Xi, Yi = meshgrid(xi, yi)
         Wi = w[:,:,varIndex_]'
      else
         if plotrange[1] == -Inf plotrange[1] = minimum(xrange) end
         if plotrange[2] ==  Inf plotrange[2] = maximum(xrange) end
         if plotrange[3] == -Inf plotrange[3] = minimum(yrange) end
         if plotrange[4] ==  Inf plotrange[4] = maximum(yrange) end

         xi = range(plotrange[1], stop=plotrange[2], step=plotinterval)
         yi = range(plotrange[3], stop=plotrange[4], step=plotinterval)

         spline = @views Spline2D(xrange, yrange, w[:,:,varIndex_])
         Xi, Yi = meshgrid(xi, yi)
         wi = @views spline(Xi[:], Yi[:])
         Wi = reshape(wi, size(Xi))
      end
   end

   # Mask a circle at the inner boundary
   if innermask
      varIndex_ = findlast(x->x=="rbody", data.head.variables)
      isnothing(varIndex_) && error("rbody not found in file header parameters!")
      ParamIndex_ = varIndex_ - data.head.ndim - data.head.nw
      @inbounds for i = eachindex(Xi, Yi)
         if Xi[i]^2 + Yi[i]^2 < data.head.eqpar[ParamIndex_]^2
            Wi[i] = NaN
         end
      end
   end

   if griddim == 1
      return xi, yi, Wi
   else
      return Xi, Yi, Wi
   end
end

"Find variable index in data."
function findindex(data::Data, var::AbstractString)
   varIndex_ = findfirst(x->x==lowercase(var), lowercase.(data.head.wnames))
   isnothing(varIndex_) && error("$(var) not found in file header variables!")
   varIndex_
end

"Generating consistent 2D arrays for passing to plotting functions."
function meshgrid(x, y)
   X = [x for _ in y, x in x]
   Y = [y for y in y, _ in x]
   X, Y
end


text:
function meshgrid(x,y)

    X = repeat(x', length(y), 1)
    Y = repeat(y, 1, length(x))
    return X, Y
    end #End Function


text:
using JuMP
using MathOptInterface

export Peeker

mutable struct Peeker
    model::JuMP.Model
    lpmodel::LPModel
    lprep::LPRep

    function Peeker(model::JuMP.Model)
        lpmod = get_lpmodel(model)
        this = new(model, lpmod, LPRep(lpmod))
    end

    function update()
        this.lpmod = get_lpmodel(this.model)
        this.lprep = LPRep(this.lpmod)
    end
end

max_stars_repo_name: henriean/solver-peeker

max_issues_repo_name: henriean/solver-peeker

max_forks_repo_name: henriean/solver-peeker


text:
@testset "WordMoversDistance" begin

# Just a stress test of the more abstract interface.
wmd = WordMoversDistance()
d1 = ["obama", "speaks", "media", "illinois"]
d2 = ["president", "greets", "press", "chicago"]
d12 = evaluate(wmd, d1, d2)
d3 = ["lawyer", "tanks", "car", "africa"]
d13 = evaluate(wmd, d1, d3)
d23 = evaluate(wmd, d2, d3)
@test d12 < d13
@test d12 < d23

end


text:
using Gadfly

S = 10
T = 30
x = 0.01*rand(S,T)
P = softmax(x, dims=1)

tic = time()
r, loglikely = CTC(P,[3 4 5])
toc = time()
println("loglikely = ",loglikely/T," (timeSteps averaged)")
println("ctc_time  = ",(toc-tic)*1000," ms")

Gadfly.plot(
layer(y=r[1,:],Geom.line,Theme(default_color=colorant"red")),
layer(y=r[2,:],Geom.line,Theme(default_color=colorant"yellow")),
layer(y=r[3,:],Geom.line,Theme(default_color=colorant"blue")),
layer(y=r[4,:],Geom.line,Theme(default_color=colorant"green")),
layer(y=r[5,:],Geom.line,Theme(default_color=colorant"orange"))
)

max_stars_repo_name: sonosole/Delta

max_issues_repo_name: sonosole/Delta

max_forks_repo_name: sonosole/Delta


text:
module JGraphs

using Javis
using LightGraphs
using Colors
using Animations

import NetworkLayout

include("util.jl")
export GB2Luxor, draw_node

include("struct.jl")
export JGraph, JGraphData, JGraphContainer
export jnodes, jedges

include("mst.jl")
export jmst

include("random_walk.jl")
export jgraph_walk

include("morph.jl")
export jgraph_morph

include("modify.jl")
export jadd_vertex!

end #module


text:
module BadExpressions

using Compat, JuMP, BARON, Test

@testset "UnrecognizedExpressionException" begin
    exception = BARON.UnrecognizedExpressionException("comparison", :(sin(x[1])))
    buf = IOBuffer()
    Base.showerror(buf, exception)
    @test Compat.occursin("sin(x[1])", String(take!(buf)))
end

@testset "Trig unrecognized" begin
    model = Model(BARON.Optimizer)
    @variable model x
    @NLconstraint model sin(x) == 0
    # @test_throws BARON.UnrecognizedExpressionException optimize!(model) # FIXME: currently broken due to lack of NLPBlock support.
end

end # module


text:
import MsgPack
import UUIDs: UUID
import HTTP
import Sockets

import Base: endswith
function endswith(vec::Vector{T}, suffix::Vector{T}) where T
    local liv = lastindex(vec)
    local lis = lastindex(suffix)
    liv >= lis && (view(vec, (liv - lis + 1):liv) == suffix)
end

include("./WebSocketFix.jl")

# to fix lots of false error messages from HTTP
# https://github.com/JuliaWeb/HTTP.jl/pull/546
# we do HTTP.Stream{HTTP.Messages.Request,S} instead of just HTTP.Stream to prevent the Julia warning about incremental compilation
function HTTP.closebody(http::HTTP.Stream{HTTP.Messages.Request,S}) where S <: IO
    if http.writechunked
        http.writechunked = false
        try
            write(http.stream, "0\r\n\r\n")
        catch end
    end
end

# https://github.com/JuliaWeb/HTTP.jl/pull/609
HTTP.URIs.escapeuri(query::Union{Vector,Dict}) = isempty(query) ? HTTP.URIs.absent : join((HTTP.URIs.escapeuri(k, v) for (k, v) in query), "&")

# from https://github.com/JuliaLang/julia/pull/36425
function detectwsl()
    Sys.islinux() &&
    isfile("/proc/sys/kernel/osrelease") &&
    occursin(r"Microsoft|WSL"i, read("/proc/sys/kernel/osrelease", String))
end

function open_in_default_browser(url::AbstractString)::Bool
    try
        if Sys.isapple()
            Base.run(`open $url`)
            true
        elseif Sys.iswindows() || detectwsl()
            Base.run(`powershell.exe Start $url`)
            true
        elseif Sys.islinux()
            Base.run(`xdg-open $url`)
            true
        else
            false
        end
    catch ex
        false
    end
end

isurl(s::String) = startswith(s, "http://") || startswith(s, "https://")

swallow_exception(f, exception_type::Type{T}) where T =
    try f()
    catch e
        isa(e, T) || rethrow(e)
    end

"""
    Pluto.run()

Start Pluto!

## Keyword arguments
You can configure some of Pluto's more technical behaviour using keyword arguments, but this is mostly meant to support testing and strange setups like Docker. If you want to do something exciting with Pluto, you can probably write a creative notebook to do it!

    Pluto.run(; kwargs...)

For the full list, see the [`Pluto.Configuration`](@ref) module. Some **common parameters**:

- `launch_browser`: Optional. Whether to launch the system default browser. Disable this on SSH and such.
- `host`: Optional. The default `host` is `"127.0.0.1"`. For wild setups like Docker and heroku, you might need to change this to `"0.0.0.0"`.
- `port`: Optional. The default `port` is `1234`.

## Technobabble

This will start the static HTTP server and a WebSocket server. The server runs _synchronously_ (i.e. blocking call) on `http://[host]:[port]/`.
Pluto notebooks can be started from the main menu in the web browser.
"""
function run(; kwargs...)
    options = Configuration.from_flat_kwargs(; kwargs...)
    run(options)
end

function run(options::Configuration.Options)
    session = ServerSession(;options=options)
    run(session)
end

# Deprecation errors

function run(host::String, port::Union{Nothing,Integer}=nothing; kwargs...)
    @error "Deprecated in favor of:

        run(;host=$host, port=$port)
    "
end

function run(port::Integer; kwargs...)
    @error "Oopsie! This is the old command to launch Pluto. The new command is:

        Pluto.run()

    without the port as argument - it will choose one automatically. If you need to specify the port, use:

        Pluto.run(port=$port)
    "
end

"""
    run(session::ServerSession)

Specifiy the [`Pluto.ServerSession`](@ref) to run the web server on, which includes the configuration. Passing a session as argument allows you to start the web server with some notebooks already running. See [`SessionActions`](@ref) to learn more about manipulating a `ServerSession`.
"""
function run(session::ServerSession)
    pluto_router = http_router_for(session)
    host = session.options.server.host
    port = session.options.server.port

    hostIP = parse(Sockets.IPAddr, host)
    if port === nothing
        port, serversocket = Sockets.listenany(hostIP, UInt16(1234))
    else
        try
            serversocket = Sockets.listen(hostIP, UInt16(port))
        catch e
            @error "Port with number $port is already in use. Use Pluto.run() to automatically select an available port."
            return
        end
    end

    shutdown_server = Ref{Function}(() -> ())

    servertask = @async HTTP.serve(hostIP, UInt16(port), stream=true, server=serversocket) do http::HTTP.Stream
        # messy messy code so that we can use the websocket on the same port as the HTTP server

        if HTTP.WebSockets.is_upgrade(http.message)
            try
                @assert is_authenticated(session, http.message)

                HTTP.WebSockets.upgrade(http) do clientstream
                    if !isopen(clientstream)
                        return
                    end
                    try
                    while !eof(clientstream)
                        # This stream contains data received over the WebSocket.
                        # It is formatted and MsgPack-encoded by send(...) in PlutoConnection.js
                        local parentbody
                        try
                            message = collect(WebsocketFix.readmessage(clientstream))
                            parentbody = unpack(message)

                            process_ws_message(session, parentbody, clientstream)
                        catch ex
                            if ex isa InterruptException
                                shutdown_server[]()
                            elseif ex isa HTTP.WebSockets.WebSocketError || ex isa EOFError
                                # that's fine!
                            else
                                bt = stacktrace(catch_backtrace())
                                @warn "Reading WebSocket client stream failed for unknown reason:" parentbody exception = (ex, bt)
                            end
                        end
                    end
                    catch ex
                        if ex isa InterruptException
                            shutdown_server[]()
                        elseif ex isa HTTP.WebSockets.WebSocketError || ex isa EOFError || (ex isa Base.IOError && occursin("connection reset", ex.msg))
                            # that's fine!
                        else
                            bt = stacktrace(catch_backtrace())
                            @warn "Reading WebSocket client stream failed for unknown reason:" exception = (ex, bt)
                        end
                    end
                end
            catch ex
                if ex isa InterruptException
                    shutdown_server[]()
                elseif ex isa Base.IOError
                    # that's fine!
                elseif ex isa ArgumentError && occursin("stream is closed", ex.msg)
                    # that's fine!
                elseif ex isa AssertionError && occursin("is_authenticated", ex.msg)
                    # That's fine!
                else
                    bt = stacktrace(catch_backtrace())
                    @warn "HTTP upgrade failed for unknown reason" exception = (ex, bt)
                end
            end
        else
            request::HTTP.Request = http.message
            request.body = read(http)
            HTTP.closeread(http)

            # If a "token" url parameter is passed in from binder, then we store it to add to every URL (so that you can share the URL to collaborate).
            params = HTTP.queryparams(HTTP.URI(request.target))
            if haskey(params, "token") && session.binder_token === nothing
                session.binder_token = params["token"]
            end

            request_body = IOBuffer(HTTP.payload(request))
            if eof(request_body)
                # no request body
                response_body = HTTP.handle(pluto_router, request)
            else
                @warn "HTTP request contains a body, huh?" request_body
            end

            request.response::HTTP.Response = response_body
            request.response.request = request
            try
                HTTP.setheader(http, "Referrer-Policy" => "origin-when-cross-origin")
                HTTP.startwrite(http)
                write(http, request.response.body)
                HTTP.closewrite(http)
            catch e
                if isa(e, Base.IOError) || isa(e, ArgumentError)
                    # @warn "Attempted to write to a closed stream at $(request.target)"
                else
                    rethrow(e)
                end
            end
        end
    end

    address = pretty_address(session, hostIP, port)

    println()
    if session.options.server.launch_browser && open_in_default_browser(address)
        println("Opening $address in your default browser... ~ have fun!")
    else
        println("Go to $address in your browser to start writing ~ have fun!")
    end
    println()
    println("Press Ctrl+C in this terminal to stop Pluto")
    println()

    shutdown_server[] = () -> @sync begin
        println("\n\nClosing Pluto... Restart Julia for a fresh session. \n\nHave a nice day! 🎈")
        @async swallow_exception(() -> close(serversocket), Base.IOError)
        # TODO: HTTP has a kill signal?
        # TODO: put do_work tokens back
        for client in values(session.connected_clients)
            @async swallow_exception(() -> close(client.stream), Base.IOError)
        end
        empty!(session.connected_clients)
        for (notebook_id, ws) in WorkspaceManager.workspaces
            @async WorkspaceManager.unmake_workspace(wait(ws))
        end
    end

    try
        # create blocking call and switch the scheduler back to the server task, so that interrupts land there
        wait(servertask)
    catch e
        if e isa InterruptException
            shutdown_server[]()
        elseif e isa TaskFailedException
            # nice!
        else
            rethrow(e)
        end
    end
end

function pretty_address(session::ServerSession, hostIP, port)
    root = if session.options.server.root_url === nothing
        host_str = string(hostIP)
        host_pretty = if isa(hostIP, Sockets.IPv6)
            if host_str == "::1"
                "localhost"
            else
                "[$(host_str)]"
            end
        elseif host_str == "127.0.0.1" # Assuming the other alternative is IPv4
            "localhost"
        else
            host_str
        end
        port_pretty = Int(port)
        "http://$(host_pretty):$(port_pretty)/"
    else
        @assert endswith(session.options.server.root_url, "/")
        session.options.server.root_url
    end

    Sys.set_process_title("Pluto server - $root")

    url_params = Dict{String,String}()

    if session.options.security.require_secret_for_access
        url_params["secret"] = session.secret
    end
    fav_notebook = session.options.server.notebook
    new_root = if fav_notebook !== nothing
        key = isurl(fav_notebook) ? "url" : "path"
        url_params[key] = string(fav_notebook)
        root * "open"
    else
        root
    end
    merge(HTTP.URIs.URI(new_root), query=url_params) |> string
end

"All messages sent over the WebSocket get decoded+deserialized and end up here."
function process_ws_message(session::ServerSession, parentbody::Dict, clientstream::IO)
    client_id = Symbol(parentbody["client_id"])
    client = get!(session.connected_clients, client_id, ClientSession(client_id, clientstream))
    client.stream = clientstream # it might change when the same client reconnects

    messagetype = Symbol(parentbody["type"])
    request_id = Symbol(parentbody["request_id"])

    args = []
    if haskey(parentbody, "notebook_id")
        notebook = let
            notebook_id = UUID(parentbody["notebook_id"])
            get(session.notebooks, notebook_id, nothing)
        end

        if messagetype === :connect
            if notebook === nothing
                messagetype === :connect || @warn "Remote notebook not found locally!"
            else
                client.connected_notebook = notebook
            end
        end

        push!(args, notebook)

        if haskey(parentbody, "cell_id")
            cell_id = UUID(parentbody["cell_id"])
            index = cell_index_from_id(notebook, cell_id)
            if index === nothing
                @warn "Remote cell not found locally!"
            else
                push!(args, notebook.cells[index])
            end
        end
    end

    body = parentbody["body"]

    if haskey(responses, messagetype)
        responsefunc = responses[messagetype]
        try
            responsefunc(session, body, args..., initiator=Initiator(client.id, request_id))
        catch ex
            @warn "Response function to message of type $(messagetype) failed"
            rethrow(ex)
        end
    else
        @warn "Message of type $(messagetype) not recognised"
    end
end


text:
"""
spin17.jl - derivative robustness for the δa problem
"""

WDIR = get(ENV, "ROBUST_QOC_PATH", "../../")
include(joinpath(WDIR, "src", "spin", "spin.jl"))

using Altro
using HDF5
using LinearAlgebra
using RobotDynamics
using StaticArrays
using TrajectoryOptimization
const RD = RobotDynamics
const TO = TrajectoryOptimization

# paths
const EXPERIMENT_META = "spin"
const EXPERIMENT_NAME = "spin17"
const SAVE_PATH = joinpath(WDIR, "out", EXPERIMENT_META, EXPERIMENT_NAME)

# problem
const CONTROL_COUNT = 1
const STATE_COUNT = 2
const ASTATE_SIZE_BASE = STATE_COUNT * HDIM_ISO + 3 * CONTROL_COUNT
const ACONTROL_SIZE = CONTROL_COUNT
# state indices
const STATE1_IDX = SVector{HDIM_ISO}(1:HDIM_ISO)
const STATE2_IDX = SVector{HDIM_ISO}(STATE1_IDX[end] + 1:STATE1_IDX[end] + HDIM_ISO)
const INTCONTROLS_IDX = SVector{CONTROL_COUNT}(STATE2_IDX[end] + 1:STATE2_IDX[end] + CONTROL_COUNT)
const CONTROLS_IDX = SVector{CONTROL_COUNT}(INTCONTROLS_IDX[end]
                                            + 1:INTCONTROLS_IDX[end] + CONTROL_COUNT)
const DCONTROLS_IDX = SVector{CONTROL_COUNT}(CONTROLS_IDX[end] + 1:CONTROLS_IDX[end] + CONTROL_COUNT)
const DSTATE1_IDX = SVector{HDIM_ISO}(DCONTROLS_IDX[end] + 1:DCONTROLS_IDX[end] + HDIM_ISO)
const D2STATE1_IDX = SVector{HDIM_ISO}(DSTATE1_IDX[end] + 1:DSTATE1_IDX[end] + HDIM_ISO)
# const STATE3_IDX = DCONTROLS_IDX[end] + 1:DCONTROLS_IDX[end] + HDIM_ISO
# const STATE4_IDX = STATE3_IDX[end] + 1:STATE3_IDX[end] + HDIM_ISO
# const DSTATE1_IDX = STATE4_IDX[end] + 1:STATE4_IDX[end] + HDIM_ISO
# const DSTATE2_IDX = DSTATE1_IDX[end] + 1:DSTATE1_IDX[end] + HDIM_ISO
# const DSTATE3_IDX = DSTATE2_IDX[end] + 1:DSTATE2_IDX[end] + HDIM_ISO
# const DSTATE4_IDX = DSTATE3_IDX[end] + 1:DSTATE3_IDX[end] + HDIM_ISO
# const D2STATE1_IDX = DSTATE4_IDX[end] + 1:DSTATE4_IDX[end] + HDIM_ISO
# const D2STATE2_IDX = D2STATE1_IDX[end] + 1:D2STATE1_IDX[end] + HDIM_ISO
# const D2STATE3_IDX = D2STATE2_IDX[end] + 1:D2STATE2_IDX[end] + HDIM_ISO
# const D2STATE4_IDX = D2STATE3_IDX[end] + 1:D2STATE3_IDX[end] + HDIM_ISO
# control indices
const D2CONTROLS_IDX = SVector{CONTROL_COUNT}(1:CONTROL_COUNT)

# model
struct Model{DO} <: AbstractModel
end
@inline RD.state_dim(::Model{DO}) where {DO} = ASTATE_SIZE_BASE + DO * HDIM_ISO
@inline RD.control_dim(::Model) = ACONTROL_SIZE
# @inline RD.state_dim(::Model{DO}) where {DO} = (
#     ASTATE_SIZE_BASE + (SAMPLE_STATE_COUNT - STATE_COUNT) * HDIM_ISO + DO * SAMPLE_STATE_COUNT * HDIM_ISO
# )

# dynamics
const NEGI2_H1_ISO = 2 * NEGI_H1_ISO
function RD.discrete_dynamics(::Type{RK3}, model::Model{DO}, astate::SVector,
                              acontrol::SVector, time::Real, dt::Real) where {DO}
    negi_h = (
        FQ_NEGI_H0_ISO
        + astate[CONTROLS_IDX[1]] * NEGI_H1_ISO
    )
    h_prop = exp(negi_h * dt)
    state1_ = astate[STATE1_IDX]
    state2_ = astate[STATE2_IDX]
    state1 =  h_prop * state1_
    state2 = h_prop * state2_
    intcontrols = astate[INTCONTROLS_IDX] + astate[CONTROLS_IDX] * dt
    controls = astate[CONTROLS_IDX] + astate[DCONTROLS_IDX] * dt
    dcontrols = astate[DCONTROLS_IDX] + acontrol[D2CONTROLS_IDX] * dt

    astate_ = [
        state1; state2; intcontrols; controls; dcontrols;
    ]

    if DO >= 1
        dstate1_ = astate[DSTATE1_IDX]
        dstate1 = h_prop * (dstate1_ + dt * NEGI_H1_ISO * state1_)
        astate_ = [astate_; dstate1]
        # state3_ = astate[STATE3_IDX]
        # state3 = h_prop * state3_
        # state4_ = astate[STATE4_IDX]
        # state4 = h_prop * state4_
        # dstate1_ = astate[DSTATE1_IDX]
        # dstate1 = h_prop * (dstate1_ + dt * NEGI_H1_ISO * state1_)
        # dstate2_ = astate[DSTATE2_IDX]
        # dstate2 = h_prop * (dstate2_ + dt * NEGI_H1_ISO * state2_)
        # dstate3_ = astate[DSTATE3_IDX]
        # dstate3 = h_prop * (dstate3_ + dt * NEGI_H1_ISO * state3_)
        # dstate4_ = astate[DSTATE4_IDX]
        # dstate4 = h_prop * (dstate4_ + dt * NEGI_H1_ISO * state4_)
        # append!(astate_, [state3; state4; dstate1; dstate2; dstate3; dstate4])
    end
    if DO >= 2
        d2state1_ = astate[D2STATE1_IDX]
        d2state1 = h_prop * (d2state1_ + dt * NEGI2_H1_ISO * dstate1_)
        astate_ = [astate_; d2state1]
        # d2state1_ = astate[D2STATE1_IDX]
        # d2state1 = h_prop * (d2state1_ + dt * NEGI2_H1_ISO * dstate1_)
        # d2state2_ = astate[D2STATE2_IDX]
        # d2state2 = h_prop * (d2state2_ + dt * NEGI2_H1_ISO * dstate2_)
        # d2state3_ = astate[D2STATE3_IDX]
        # d2state3 = h_prop * (d2state3_ + dt * NEGI2_H1_ISO * dstate3_)
        # d2state4_ = astate[D2STATE4_IDX]
        # d2state4 = h_prop * (d2state4_ + dt * NEGI2_H1_ISO * dstate4_)
        # append!(astate_, [d2state1; d2state2; d2state3; d2state4])
    end

    return astate_
end

# main
function run_traj(;gate_type=xpiby2, evolution_time=56.8, solver_type=altro,
                  sqrtbp=false, derivative_order=0, integrator_type=rk3, qs=ones(7),
                  smoke_test=false, dt_inv=Int64(1e1), constraint_tol=1e-8, al_tol=1e-4,
                  pn_steps=2, max_penalty=1e11, verbose=true, save=true, max_iterations=Int64(2e5))
    # model configuration
    model = Model{derivative_order}()
    n = state_dim(model)
    m = control_dim(model)
    t0 = 0.

    # initial state, target state
    x0 = zeros(n)
    xf = zeros(n)
    gate = GT_GATE_ISO[gate_type]
    x0[STATE1_IDX] = IS1_ISO_
    x0[STATE2_IDX] = IS2_ISO_
    # x0[STATE3_IDX] = IS3_ISO_
    # x0[STATE4_IDX] = IS4_ISO_
    xf[STATE1_IDX] = gate * IS1_ISO_
    xf[STATE2_IDX] = gate * IS2_ISO_
    # xf[STATE3_IDX] = gate * IS3_ISO_
    # xf[STATE4_IDX] = gate * IS4_ISO_
    x0 = SVector{n}(x0)
    xf = SVector{n}(xf)

    # control amplitude constraint
    x_max = fill(Inf, n)
    x_max[CONTROLS_IDX] .= MAX_CONTROL_NORM_0
    x_max = SVector{n}(x_max)
    x_min = fill(-Inf, n)
    x_min[CONTROLS_IDX] .= -MAX_CONTROL_NORM_0
    x_min = SVector{n}(x_min)

    # control amplitude constraint at boundary
    x_max_boundary = fill(Inf, n)
    x_max_boundary[CONTROLS_IDX] .= 0
    x_max_boundary = SVector{n}(x_max_boundary)
    x_min_boundary = fill(-Inf, n)
    x_min_boundary[CONTROLS_IDX] .= 0
    x_min_boundary = SVector{n}(x_min_boundary)

    # initial trajectory
    dt = dt_inv^(-1)
    N = Int(floor(evolution_time * dt_inv)) + 1
    U0 = [SVector{m}(
        fill(1e-4, CONTROL_COUNT)
    ) for k = 1:N-1]
    X0 = [SVector{n}([
        fill(NaN, n);
    ]) for k = 1:N]
    Z = Traj(X0, U0, dt * ones(N))

    # cost function
    Q = Diagonal(SVector{n}([
        fill(qs[1], STATE_COUNT * HDIM_ISO); # ψ1, ψ2
        fill(qs[2], CONTROL_COUNT); # ∫a
        fill(qs[3], CONTROL_COUNT); # a
        fill(qs[4], CONTROL_COUNT); # ∂a
        fill(qs[5], eval(:($derivative_order >=1 ? HDIM_ISO : 0))); #∂ψ1
        fill(qs[6], eval(:($derivative_order >=2 ? HDIM_ISO : 0))); #∂2ψ1
        # fill(0, eval(:($derivative_order >= 1 ?
        # ($SAMPLE_COUNT - $STATE_COUNT) * $HDIM_ISO : 0))) # ψ3, ψ4
        # fill(qs[5], eval(:($derivative_order >= 1 ? $SAMPLE_COUNT * $HDIM_ISO : 0))); # ∂ψ
        # fill(qs[6], eval(:($derivative_order >= 2 ? $SAMPLE_COUNT * $HDIM_ISO : 0))); # ∂2ψ
    ]))
    Qf = Q * N
    R = Diagonal(SVector{m}([
        fill(qs[7], CONTROL_COUNT);
    ]))
    objective = LQRObjective(Q, R, Qf, xf, N)

    # must satisfy control amplitude bound
    control_bnd = BoundConstraint(n, m, x_max=x_max, x_min=x_min)
    # must statisfy conrols start and end at 0
    control_bnd_boundary = BoundConstraint(n, m, x_max=x_max_boundary, x_min=x_min_boundary)
    # must reach target state, must have zero net flux
    target_astate_constraint = GoalConstraint(xf, [STATE1_IDX; STATE2_IDX]) # ; INTCONTROLS_IDX])
    # must obey unit norm.
    norm_idxs = [STATE1_IDX, STATE2_IDX]
    norm_constraints = [NormConstraint(n, m, 1, TO.Equality(), idx) for idx in norm_idxs]

    constraints = ConstraintList(n, m, N)
    # add_constraint!(constraints, control_bnd, 2:N-2)
    add_constraint!(constraints, control_bnd_boundary, N-1:N-1)
    add_constraint!(constraints, target_astate_constraint, N:N);
    for norm_constraint in norm_constraints
        add_constraint!(constraints, norm_constraint, 2:N-1)
    end

    # solve problem
    prob = Problem{IT_RDI[integrator_type]}(model, objective, constraints,
                                            x0, xf, Z, N, t0, evolution_time)
    solver = ALTROSolver(prob)
    verbose_pn = verbose ? true : false
    verbose_ = verbose ? 2 : 0
    projected_newton = solver_type == altro ? true : false
    constraint_tolerance = solver_type == altro ? constraint_tol : al_tol
    iterations_inner = smoke_test ? 1 : 300
    iterations_outer = smoke_test ? 1 : 30
    n_steps = smoke_test ? 1 : pn_steps
    set_options!(solver, square_root=sqrtbp, constraint_tolerance=constraint_tolerance,
                 projected_newton_tolerance=al_tol, n_steps=n_steps,
                 penalty_max=max_penalty, verbose_pn=verbose_pn, verbose=verbose_,
                 projected_newton=projected_newton, iterations_inner=iterations_inner,
                 iterations_outer=iterations_outer, iterations=max_iterations)
    Altro.solve!(solver)

    # post-process
    acontrols_raw = TO.controls(solver)
    acontrols_arr = permutedims(reduce(hcat, map(Array, acontrols_raw)), [2, 1])
    astates_raw = TO.states(solver)
    astates_arr = permutedims(reduce(hcat, map(Array, astates_raw)), [2, 1])
    Q_raw = Array(Q)
    Q_arr = [Q_raw[i, i] for i in 1:size(Q_raw)[1]]
    Qf_raw = Array(Qf)
    Qf_arr = [Qf_raw[i, i] for i in 1:size(Qf_raw)[1]]
    R_raw = Array(R)
    R_arr = [R_raw[i, i] for i in 1:size(R_raw)[1]]
    cidx_arr = Array(CONTROLS_IDX)
    d2cidx_arr = Array(D2CONTROLS_IDX)
    cmax = TrajectoryOptimization.max_violation(solver)
    cmax_info = TrajectoryOptimization.findmax_violation(TO.get_constraints(solver))
    iterations_ = Altro.iterations(solver)

    result = Dict(
        "acontrols" => acontrols_arr,
        "controls_idx" => cidx_arr,
        "d2controls_dt2_idx" => d2cidx_arr,
        "evolution_time" => evolution_time,
        "astates" => astates_arr,
        "Q" => Q_arr,
        "Qf" => Qf_arr,
        "R" => R_arr,
        "cmax" => cmax,
        "cmax_info" => cmax_info,
        "dt" => dt,
        "derivative_order" => derivative_order,
        "solver_type" => Integer(solver_type),
        "sqrtbp" => Integer(sqrtbp),
        "max_penalty" => max_penalty,
        "constraint_tol" => constraint_tol,
        "al_tol" => al_tol,
        "gate_type" => Integer(gate_type),
        "save_type" => Integer(jl),
        "integrator_type" => Integer(integrator_type),
        "iterations" => iterations_,
        "max_iterations" => max_iterations,
        "pn_steps" => pn_steps,
    )

    # save
    if save
        save_file_path = generate_file_path("h5", EXPERIMENT_NAME, SAVE_PATH)
        println("Saving this optimization to $(save_file_path)")
        h5open(save_file_path, "cw") do save_file
            for key in keys(result)
                write(save_file, key, result[key])
            end
        end
        result["save_file_path"] = save_file_path
    end

    return result
end

max_stars_repo_name: SchusterLab/rbqoc

max_issues_repo_name: SchusterLab/rbqoc

max_forks_repo_name: SchusterLab/rbqoc


text:
# ---------------------------------------------------------------------------- #
#
#   quadratureLine.jl
#
#   Get Legendre quadrature points for Line
#
#   λυτέος
#   Fall 2017
#
#   Max Opgenoord
#
# ---------------------------------------------------------------------------- #

function quadratureLine( P::PG1 )

  nq = 1

  weight = Array{Float64}(nq)
  xc     = Array{Float64}(nq)

  weight[1] = 2
  xc[1]     = 0

  # Scale to [0,1]
  xc     = 0.5 * ( xc + 1 )
  weight = 0.5 * weight

  return xc, weight

end

function quadratureLine( P::PG3 )

  nq = 2

  weight = Array{Float64}(nq)
  xc     = Array{Float64}(nq)

  wght = 1
  x    = 1./sqrt(3.)
  weight[1] = wght
  weight[2] = wght
  xc[1]     = -x
  xc[2]     =  x

  # Scale to [0,1]
  xc     = 0.5 * ( xc + 1 )
  weight = 0.5 * weight

  return xc, weight

end

function quadratureLine( P::PG5 )

  nq = 3

  weight = Array{Float64}(nq)
  xc     = Array{Float64}(nq)

  wght = 5./9.
  x    = sqrt(3./5.)
  weight[1] = wght
  weight[3] = wght
  xc[1]  = -x
  xc[3]  =  x

  wght = 8./9.
  x    = 0
  weight[2] = wght
  xc[2]  = x

  # Scale to [0,1]
  xc     = 0.5 * ( xc + 1 )
  weight = 0.5 * weight

  return xc, weight

end

function quadratureLine( P::PG7 )

  nq = 4

  weight = Array{Float64}(nq)
  xc     = Array{Float64}(nq)

  wght = 0.5 - sqrt(5./6.)/6.
  x    = sqrt((3 + 2*sqrt(6./5.))/7.)

  weight[1] = wght
  weight[4] = wght
  xc[1]     = -x
  xc[4]     =  x

  wght = 0.5 + sqrt(5./6.)/6.
  x    = sqrt((3 - 2*sqrt(6./5.))/7.)
  weight[2] = wght
  weight[3] = wght
  xc[2]     = -x
  xc[3]     =  x

  # Scale to [0,1]
  xc     = 0.5 * ( xc + 1 )
  weight = 0.5 * weight

  return xc, weight

end

function quadratureLine( P::PG9 )

  nq = 5

  weight = Array{Float64}(nq)
  xc     = Array{Float64}(nq)

  wght      = (322 - 13*sqrt(70.))/900.
  x         = sqrt(5 + 2*sqrt(10./7.))/3.
  weight[1] = wght
  weight[5] = wght
  xc[1]     = -x
  xc[5]     =  x

  wght      = (322 + 13*sqrt(70.))/900.
  x         = sqrt(5 - 2*sqrt(10./7.))/3.
  weight[2] = wght
  weight[4] = wght
  xc[2]     = -x
  xc[4]     =  x

  wght      = 128./225.
  x         = 0
  weight[3] = wght
  xc[3]     = x

  # Scale to [0,1]
  xc     = 0.5 * ( xc + 1 )
  weight = 0.5 * weight

  return xc, weight

end

function quadratureLine( P::PG11 )

  nq = 6

  weight = Array{Float64}(nq)
  xc     = Array{Float64}(nq)

  xc[1]     = -0.9324695142031520278123016
  xc[2]     = -0.6612093864662645136613996
  xc[3]     = -0.2386191860831969086305017
  xc[4]     =  0.2386191860831969086305017
  xc[5]     =  0.6612093864662645136613996
  xc[6]     =  0.9324695142031520278123016

  weight[1] = 0.1713244923791703450402961
  weight[2] = 0.3607615730481386075698335
  weight[3] = 0.4679139345726910473898703
  weight[4] = 0.4679139345726910473898703
  weight[5] = 0.3607615730481386075698335
  weight[6] = 0.1713244923791703450402961

  # Scale to [0,1]
  xc     = 0.5 * ( xc + 1 )
  weight = 0.5 * weight

  return xc, weight

end

function quadratureLine( P::PG13 )

  nq = 7

  weight = Array{Float64}(nq)
  xc     = Array{Float64}(nq)

  xc[1]     = -0.9491079123427585245261897
  xc[2]     = -0.7415311855993944398638648
  xc[3]     = -0.4058451513773971669066064
  xc[4]     =  0.0000000000000000000000000
  xc[5]     =  0.4058451513773971669066064
  xc[6]     =  0.7415311855993944398638648
  xc[7]     =  0.9491079123427585245261897

  weight[1] = 0.1294849661688696932706114
  weight[2] = 0.2797053914892766679014678
  weight[3] = 0.3818300505051189449503698
  weight[4] = 0.4179591836734693877551020
  weight[5] = 0.3818300505051189449503698
  weight[6] = 0.2797053914892766679014678
  weight[7] = 0.1294849661688696932706114

  # Scale to [0,1]
  xc     = 0.5 * ( xc + 1 )
  weight = 0.5 * weight

  return xc, weight

end


text:
# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# Script for some descriptive statistics
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

using DataFrames, CSV, XLSX, LinearAlgebra, Statistics, StatsBase

dir_io = "C:/Users/u0148308/data/comext/" # location of input/output (io)

# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# Import data into Julia
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

# Notes:
#   - import raw data, initial cleaning
#   - subset and compute data for figure
#   - loop over months and add to DataFrame
#   - plot figure

# initial cleaning of the data
function initial_cleaning(year::String, month::String)

    path = dir_io * "raw/" * "full" * year * month * ".dat"
    df = CSV.read(path, DataFrame)

    # column types
    # Notes:
    #   - treat PERIOD as numeric to allow sorting
    transform!(df, [:DECLARANT, :PARTNER, :PERIOD] .=> ByRow(Int64), renamecols=false)

    # additional PRODUCT_BEC5 column from 2017 onwards
    if year in ["2015", "2016"] # why does shorthand not work?
        df.PRODUCT_BEC5 .= missing
    else
    end
    cols_string = ["DECLARANT_ISO", "PARTNER_ISO", "TRADE_TYPE", "PRODUCT_NC", "PRODUCT_SITC", "PRODUCT_CPA2002", "PRODUCT_CPA2008", "PRODUCT_CPA2_1",
                   "PRODUCT_BEC", "PRODUCT_BEC5", "PRODUCT_SECTION", "FLOW", "STAT_REGIME"]
    transform!(df,  cols_string .=> ByRow(string), renamecols=false)

    transform!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG, :SUP_QUANTITY] .=> ByRow(Float64), renamecols=false)

    # formatting
    df[:, :TRADE_TYPE] .= ifelse.(df[:, :TRADE_TYPE] .== "I", "intra", "extra")
    df[:, :FLOW] .= ifelse.(df[:, :FLOW] .== "1", "imports", "exports")
    df[:, :SUPP_UNIT] .= ifelse.(ismissing.(df[:, :SUPP_UNIT]), missing, string.(df[:, :SUPP_UNIT])) # to make sure type is of Union{Missing, String}

    # missing data
    # Notes:
    #   - if value or quanity is reported as 0 use missing instead
    transform!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG, :SUP_QUANTITY] .=> ByRow(x->ifelse(x == 0.0, missing, x)), renamecols=false)

    # aggregate over STAT_REGIME, SUPP_UNIT
    # Notes:
    #   - only lose few rows ~1%
    #   - disregard SUPP_UNIT and SUP_QUANTITY for the time being
    cols_grouping = ["DECLARANT_ISO", "PARTNER_ISO", "TRADE_TYPE", "PRODUCT_NC", "PRODUCT_SITC", "PRODUCT_CPA2002", "PRODUCT_CPA2008", "PRODUCT_CPA2_1",
                     "PRODUCT_BEC", "PRODUCT_BEC5", "PRODUCT_SECTION", "FLOW", "PERIOD"]
    gdf = groupby(df, cols_grouping)
    df = combine(gdf, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> sum, renamecols=false)

    return df
end

df = initial_cleaning("2020", "01")

# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# figure1
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

# Notes:
#   - figure 1: evolution of total imports/exports from BE to UK vis-a-vis FR, NL, DE and EU
#   - figure 2: evolution of imports/exports from BE to UK at 2digit level
#   - figure 3: evolution of share of imports/exports from BE to UK of total imports/exports from BE
#   - all figures in values and unit values (unit prices):
#       + compute a price index as using a weighted sum
#   - include vertical lines for different Brexit events

# function to create EU as a seperate partner
# Notes:
#   - with the parameter "EU" one can specify the exact countries
function append_EU(df::DataFrame, EU::Vector{String})

    # subset
    df_EU = subset(df, :TRADE_TYPE => ByRow(x-> x == "intra"), :PARTNER_ISO => ByRow(x -> x in EU))
    cols_grouping = ["DECLARANT_ISO", "PRODUCT_NC", "PRODUCT_SITC", "PRODUCT_CPA2002", "PRODUCT_CPA2008", "PRODUCT_CPA2_1",
                    "PRODUCT_BEC", "PRODUCT_BEC5", "PRODUCT_SECTION", "FLOW", "PERIOD"]
    gdf = groupby(df_EU, cols_grouping)
    df_EU = combine(gdf, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> (x -> sum(skipmissing(x))), renamecols=false)
    df_EU.TRADE_TYPE .= "intra"
    df_EU.PARTNER_ISO .= "EU"

    df = vcat(df, df_EU)

    return df
end

# EU27
ctrys_EU = ["AT", "BE", "BG", "CY", "CZ", "DE", "DK", "EE", "ES", "FI", "FR", "GR", "HR", "HU",
             "IE", "IT", "LT", "LU", "LV", "MT", "NL", "PL", "PT", "RO", "SE", "SI", "SK"]

df = append_EU(df, ctrys_EU)

# function to create data for figure 1
function data_fig1(df::DataFrame, declarants::Vector{String}, partners::Vector{String})

    # subset dataframe
    df = subset(df, :DECLARANT_ISO => ByRow(x -> x in declarants), :PARTNER_ISO => ByRow(x -> x in partners))
    subset!(df, :PRODUCT_NC => ByRow(x -> x != "TOTAL")) # take out total
    subset!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> ByRow(x -> !ismissing(x))) # take out missing values
    subset!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> ByRow(x -> x != 0))

    # compute UNIT_PRICE
    transform!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] => ByRow((v,q) -> v/q) => :UNIT_PRICE)

    # compute weights for price index
    cols_grouping = ["DECLARANT_ISO", "PARTNER_ISO", "FLOW", "PERIOD"]
    gdf = groupby(df, cols_grouping)
    df_total = combine(gdf, :VALUE_IN_EUROS => sum => :VALUE_TOTAL)
    df_join = leftjoin(df, df_total, on=cols_grouping)
    transform!(df_join, [:VALUE_IN_EUROS, :VALUE_TOTAL] => ByRow((v,s) -> v/s) => :VALUE_WEIGHTS)
    transform!(df_join, [:UNIT_PRICE, :VALUE_WEIGHTS] => ByRow((p, w) -> p * w) => :PRICE_INDEX)

    # aggregate over products
    gdf = groupby(df_join, cols_grouping)
    df = combine(gdf, [:VALUE_IN_EUROS, :QUANTITY_IN_KG, :PRICE_INDEX] .=> sum, renamecols=false)

    return df
end

df_fig1 = DataFrame(DECLARANT_ISO=String[], PARTNER_ISO=String[], FLOW=String[], PERIOD=Int64[],
                    VALUE_IN_EUROS=Float64[], QUANTITY_IN_KG=Float64[], PRICE_INDEX=Float64[])

a = data_fig1(df, ["BE"], ["GB", "DE", "FR", "NL", "EU"])

# function to create data for figure 2
# Notes:
#   - essentially the same function as for figure 1 but with additional grouping at CN digit level
function data_fig2(df::DataFrame, declarants::Vector{String}, partners::Vector{String}, digits::Int64)

    # subset dataframe
    df = subset(df, :DECLARANT_ISO => ByRow(x -> x in declarants), :PARTNER_ISO => ByRow(x -> x in partners))
    subset!(df, :PRODUCT_NC => ByRow(x -> x != "TOTAL")) # take out total
    subset!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> ByRow(x -> !ismissing(x))) # take out missing values

    # compute UNIT_PRICE
    transform!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] => ByRow((v,q) -> v/q) => :UNIT_PRICE)

    # add 2 digit CN classification
    # Notes:
    #   - is simply the first digits
    g(x) = x[1:digits]
    df.PRODUCT_NC_digits = g.(df.PRODUCT_NC)

    # compute weights for price index
    cols_grouping = ["DECLARANT_ISO", "PARTNER_ISO", "FLOW", "PERIOD", "PRODUCT_NC_digits"]
    gdf = groupby(df, cols_grouping)
    df_total = combine(gdf, :VALUE_IN_EUROS => sum => :VALUE_TOTAL)
    df_join = leftjoin(df, df_total, on=cols_grouping)
    transform!(df_join, [:VALUE_IN_EUROS, :VALUE_TOTAL] => ByRow((v,s) -> v/s) => :VALUE_WEIGHTS)
    transform!(df_join, [:UNIT_PRICE, :VALUE_WEIGHTS] => ByRow((p, w) -> p * w) => :PRICE_INDEX)

    # aggregate over products
    gdf = groupby(df_join, cols_grouping)
    df = combine(gdf, [:VALUE_IN_EUROS, :QUANTITY_IN_KG, :PRICE_INDEX] .=> sum, renamecols=false)

    return df
end

df_fig2 = DataFrame(DECLARANT_ISO=String[], PARTNER_ISO=String[], FLOW=String[], PERIOD=Int64[], PRODUCT_NC_digits=String[],
                    VALUE_IN_EUROS=Float64[], QUANTITY_IN_KG=Float64[], PRICE_INDEX=Float64[])

a = data_fig2(df, ["BE"], ["GB"], 2)

# function to create data for figure 3
function data_fig3(df::DataFrame, declarants::Vector{String}, partners::Vector{String})

    # clean df
    subset!(df, :PRODUCT_NC => ByRow(x -> x != "TOTAL")) # take out TOTAL
    subset!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> ByRow(x -> !ismissing(x))) # take out missing values

    # compute total exports/imports
    df_total = subset(df, :DECLARANT_ISO => ByRow(x -> x in ["BE"]))
    cols_grouping = ["DECLARANT_ISO", "FLOW", "PERIOD"]
    gdf = groupby(df_total, cols_grouping)
    df_total = combine(gdf, :VALUE_IN_EUROS => sum => :VALUE_TOTAL, :QUANTITY_IN_KG => sum => :QUANTITY_TOTAL)

    # compute total exports/imports for each partner
    df_total_partner = subset(df, :DECLARANT_ISO => ByRow(x -> x in ["BE"]), :PARTNER_ISO => ByRow(x -> x in ["GB"]))
    cols_grouping = ["DECLARANT_ISO", "PARTNER_ISO", "FLOW", "PERIOD"]
    gdf = groupby(df_total_partner, cols_grouping)
    df_total_partner = combine(gdf, :VALUE_IN_EUROS => sum => :VALUE_TOTAL_PARTNER, :QUANTITY_IN_KG => sum => :QUANTITY_TOTAL_PARTNER)

    # compute export/import shares
    # Notes:
    #   - shares are in percentages
    df_join = leftjoin(df_total_partner, df_total, on=["DECLARANT_ISO", "FLOW", "PERIOD"])
    transform!(df_join, [:VALUE_TOTAL_PARTNER, :VALUE_TOTAL] => ByRow((v,s) -> v/s*100) => :VALUE_SHARE)
    transform!(df_join, [:QUANTITY_TOTAL_PARTNER, :QUANTITY_TOTAL] => ByRow((v,s) -> v/s*100) => :QUANTITY_SHARE)

    df = df_join[:,Not([:VALUE_TOTAL_PARTNER, :VALUE_TOTAL, :QUANTITY_TOTAL_PARTNER, :QUANTITY_TOTAL])]

    return df
end

df_fig3 = DataFrame(DECLARANT_ISO=String[], PARTNER_ISO=String[], FLOW=String[], PERIOD=Int64[],
                    VALUE_SHARE=Float64[], QUANTITY_SHARE=Float64[])

a = data_fig3(df, ["BE"], ["GB"])


text:
#=
Two pooling layers - MeanPool and MaxPool
=#

########################################
# MeanPool
########################################

#=
Mean pooling function
 2D input -> 2D pooled output for each patch of each image
 Input:
 A - inputs from previous layer - dim x dim x num_patches x num_obs
 Output (updated in-place):
 Z - pooled features - pool_dim x pool_dim x num_patches x num_obs
=#
function meanpool!{F<:AbstractFloat}(Z::Array{F,4}, A::Array{F,4})

    stride = div(size(A, 1), size(Z, 1))
    t = stride * stride

    for img in indices(Z, 4)
        for patch in indices(Z, 3)
            col_start = 1
            col_end = stride
            for pcol in indices(Z, 2)
                row_start = 1
                row_end = stride
                for prow in indices(Z, 1)
                    s = F(0)
                    for incol = col_start:col_end
                        for inrow = row_start:row_end
                            @inbounds s += A[inrow, incol, patch, img]
                        end
                    end
                    @inbounds Z[prow, pcol, patch, img] = s / t
                    row_start += stride
                    row_end += stride
                end
                col_start += stride
                col_end += stride
            end
        end
    end
end

#=
Upsample function for mean pooling - uniformly distributes the error for a
single pooling unit among the units which feed into it in the previous layer
 Input:
 Δpool - delta of pool layer
       - pool_dim x pool_dim x num_patches x num_obs
 Output (updated in-place):
 Δconv - delta of previous layer (usually a convolution layer)
       - conv_dim x conv_dim x num_patches x num_obs
=#
function upsample_mean!{F<:AbstractFloat}(Δconv::Array{F,4}, Δpool::Array{F,4})

    stride = div(size(Δconv, 1), size(Δpool, 1))
    m = inv(stride^2)

    for img in indices(Δpool, 4)
        for patch in indices(Δpool, 3)
            start_col = 1
            stop_col = stride
            for pcol in indices(Δpool, 2)
                start_row = 1
                stop_row = stride
                for prow in indices(Δpool, 1)
                    @inbounds tmp = Δpool[prow, pcol, patch, img] * m
                    for col in start_col:stop_col
                        for row in start_row:stop_row
                            @inbounds Δconv[row, col, patch, img] = tmp
                        end
                    end
                    start_row += stride
                    stop_row += stride
                end
                start_col += stride
                stop_col += stride
            end
        end
    end
end

########################################
# MaxPool
########################################

#=
Max pooling function
 2D input -> 2D pooled output for each patch of each image
 Inputs:
 A - inputs from previous layer - dim x dim x num_patches x num_obs
 Output (updated in-place):
 Z - pooled features - pool_dim x pool_dim x num_patches x num_obs
 maxinds - linear index of the max value from each patch
=#
function maxpool!{F<:AbstractFloat, T<:Integer}(Z::Array{F,4}, A::Array{F,4},
                                                maxinds::Array{T,4})

    dims = size(A)
    stride = div(size(A, 1), size(Z, 1))

    for img in indices(Z, 4)
        for patch in indices(Z, 3)
            col_start = 1
            col_end = stride
            for pcol in indices(Z, 2)
                row_start = 1
                row_end = stride
                for prow in indices(Z, 1)
                    maxval = F(-Inf)
                    maxind = T(0)
                    for incol = col_start:col_end
                        for inrow = row_start:row_end
                            @inbounds val = A[inrow, incol, patch, img]
                            if val > maxval
                                maxval = val
                                maxind = sub2ind(dims, inrow, incol, patch, img)
                            end
                        end
                    end
                    @inbounds Z[prow, pcol, patch, img] = maxval
                    @inbounds maxinds[prow, pcol, patch, img] = maxind
                    row_start += stride
                    row_end += stride
                end
                col_start += stride
                col_end += stride
            end
        end
    end
end

#=
Upsample function for max pooling - the unit which was chosen as the max
receives all the error
 Input:
 Δpool - delta of pool layer
       - pool_dim x pool_dim x num_patches x num_obs
 maxinds - linear index of the max value from each patch
 Output (updated in-place):
 Δconv - delta of previous layer (usually a convolution layer)
       - conv_dim x conv_dim x num_patches x num_obs
=#
function upsample_max!(Δconv, Δpool, maxinds)
    fill!(Δconv, 0)
    for img in indices(Δpool, 4)
        for patch in indices(Δpool, 3)
            for pcol in indices(Δpool, 2)
                for prow in indices(Δpool, 1)
                    @inbounds ind = maxinds[prow, pcol, patch, img]
                    @inbounds Δconv[ind] = Δpool[prow, pcol, patch, img]
                end
            end
        end
    end
end


text:
include("download_soap_example.jl")
# include("download_soap_example_from_gilbertson_etal_2020.jl")   # Commented out since too large

max_stars_repo_name: RvSpectML/EchelleCCFs

max_issues_repo_name: RvSpectML/EchelleCCFs

max_forks_repo_name: RvSpectML/EchelleCCFs


text:
module Quadrature

include("basic.jl")
export trapz, simps

end # module

max_stars_repo_name: gabrielgellner/Quadrature

max_issues_repo_name: gabrielgellner/Quadrature

max_forks_repo_name: gabrielgellner/Quadrature


text:
using Base.Cartesian

"""
    Sequence{N,T}(::AbstractArray)

A `Sequence` is a thin wrapper around a multi-dimensional array that
allows you to index it like a vector. For example, the following
code shows a 10-dimensional time series, sampled at 100 time points:

```
data = randn(10,100)
seq = Sequence(data)
seq[1] == data[:,1] # true
```

This generalizes to higher-order arrays. Consider a time series
where a 10x10 matrix of data is collected at 100 time points:

```
data = randn(10,10,100)
seq = Sequence(data)
seq[1] == data[:,:,1] # true
```
"""
struct Sequence{N,T} <: AbstractArray{T,1}
    val::AbstractArray{T,N}
end

Base.size(x::Sequence{N}) where N = (size(x.val,N),)
Base.eltype(x::Sequence{N,T}) where {N, T} = T

@generated function Base.getindex(x::Sequence{N}, i) where N
    :( x.val[@ntuple($N, (n-> n==$N ? i : Colon()))...] )
end

@generated function Base.setindex!(x::Sequence{N}, val, i) where N
    :( x.val[@ntuple($N, (n-> n==$N ? i : Colon()))...] = val )
end

# convert a sequence back into an array
# in case sequences have different lenght, throw error
function seq_to_array(seq::AbstractVector{T}) where T <: Sequence
  len_seq = length(seq[1])
  arr = zeros(typeof(seq[1][1]),len_seq,length(seq))
  for i=1:length(seq)
    length(seq[i]) != len_seq ? error("Sequences do not have the same length, cannot construct array") : nothing
    arr[:,i] = seq[i][:]
  end
  return arr
end


text:
abstract type Scaler end

export ScalePath
include("./ScalePath.jl")

export SwitchPath
include("./SwitchPath.jl")

export ScaleChannels
include("./ScaleChannels.jl")


text:
#=
P27 (**) Group the elements of a set into disjoint subsets.
a) In how many ways can a group of 9 people work in 3 disjoint subgroups of 2, 3 and 4 persons? Write a function that generates all the possibilities and returns them in a list.

Example:
* (group3 '(aldo beat carla david evi flip gary hugo ida))
( ( (ALDO BEAT) (CARLA DAVID EVI) (FLIP GARY HUGO IDA) )
... )

b) Generalize the above predicate in a way that we can specify a list of group sizes and the predicate will return a list of groups.

Example:
* (group '(aldo beat carla david evi flip gary hugo ida) '(2 2 5))
( ( (ALDO BEAT) (CARLA DAVID) (EVI FLIP GARY HUGO IDA) )
... )

Note that we do not want permutations of the group members; i.e. ((ALDO BEAT) ...) is the same solution as ((BEAT ALDO) ...). However, we make a difference between ((ALDO BEAT) (CARLA DAVID) ...) and ((CARLA DAVID) (ALDO BEAT) ...).

You may find more about this combinatorial problem in a good book on discrete mathematics under the term "multinomial coefficients".
=#


text:
function mincut(graph, source, sink, interdiction, ::Val{:raw})
    m = model(; kind=:cut)
    n = size(graph, 1)

    d = domain(0:n)

    separator = n + 1 # value that separate the two sides of the cut

    # Add variables:
    foreach(_ -> variable!(m, d), 0:n)

    # Extract error function from usual_constraint
    e1 = (x; param=nothing, dom_size=n + 1) -> error_f(
        usual_constraints[:ordered])(x; param, dom_size
    )
    e2 = (x; param=nothing, dom_size=n + 1) -> error_f(
        usual_constraints[:all_different])(x; param, dom_size
    )

    # Add constraint
    constraint!(m, e1, [source, separator, sink])
    constraint!(m, e2, 1:(n + 1))

    # Add objective
    objective!(m, (x...) -> o_mincut(graph, x...; interdiction))

    return m
end

function mincut(graph, source, sink, interdiction, ::Val{:JuMP})
    m = JuMP.Model(CBLS.Optimizer)
    n = size(graph, 1)
    separator = n + 1

    @variable(m, 0 ≤ X[1:separator] ≤ n, Int)

    @constraint(m, [X[source], X[separator], X[sink]] in Ordered())
    @constraint(m, X in AllDifferent())

    obj(x...) = o_mincut(graph, x...; interdiction)
    @objective(m, Min, ScalarFunction(obj))

    return m, X
end

"""
    mincut(graph; source, sink, interdiction =0, modeler = :JuMP)

Compute the minimum cut of a graph.

# Arguments:
- `graph`: Any matrix <: AbstractMatrix that describes the capacities of the graph
- `source`: Id of the source node; must be set
- `sink`: Id of the sink node; must be set
- `interdiction`: indicates the number of forbidden links
- `modeler`: Default to `:JuMP`.
"""
function mincut(graph; source, sink, interdiction =0, modeler = :JuMP)
    return mincut(graph, source, sink, interdiction, Val(modeler))
end


text:
using Images

import Pkg
Pkg.activate("../../WCvsLHE")
using WCvsLHE

normalize(x) = typeof(x)( (x-minimum(x).*ones(size(x)))*0.7./(maximum(x)-minimum(x)) + 0.15*(ones(size(x))) )

function plot_comparison_middle(A, B,C;args...)
    mA = floor(size(A,1)/2) |> Int64
    mB = floor(size(B,1)/2) |> Int64
    mC = floor(size(C,1)/2)|> Int64
    Plots.plot([Float64.(A[mA,:]),Float64.(B[mB,:]), Float64.(C[mC,:])]; args...)
end

function test(name, params; lhe3d = false, dest = name)
	println(string("Processing: ", name))
	img = Gray.(load(string("originals/",name,".png")))
	img = Gray.(normalize(Float64.(img)))
	print(string("\tContrast enhancement: WC2D "))
	wc2d = wc(img, params[:WC2D], algo_type = :planar) |> normalize
	print(string("LHE2D "))
	lhe2d = lhe(img, params[:LHE2D], algo_type = :planar) |> normalize
	print(string("WC3D "))
	wc3d = wc(img, params[:WC3D], algo_type = :cortical) |> normalize

	if lhe3d
		print(string("LHE3D "))
		lhe3d = lhe(img, params[:LHE3D], algo_type = :cortical) |> normalize
		save(string("results/",dest,"_lhe3d.png"), map(clamp01nan,lhe3d))
	end

	println("\n\tSaving.")

	save(string("results/",dest,".png"), img)
	save(string("results/",dest,"_wc2d.png"), map(clamp01nan, wc2d))
	save(string("results/",dest,"_lhe2d.png"), map(clamp01nan, lhe2d))
	save(string("results/",dest,"_wc3d.png"), map(clamp01nan, wc3d))
end

# Illusion parameters
illusions = Dict(
"white" => Dict(:WC2D => Params(10,20,.7,1.4),
				:LHE2D => Params(10,50,.7,1),
				:WC3D => Params(20,30,.7,1.4),
				:LHE3D => Params(2, 50,.7,1))

, "brightness" => Dict(:WC2D => Params(2,10,.7,1.4),
				:LHE2D => Params(2,10,.7,1),
				:WC3D => Params(2,10,.7,1.4),
				:LHE3D => Params(2, 10,.7,1))

, "checkerboard" => Dict(:WC2D => Params(10,70,.7,1.4),
				:LHE2D => Params(10,70,.7,1),
				:WC3D => Params(10,70,.7,1.4),
				:LHE3D => Params(10,70,.7,1))

, "chevreul" => Dict(:WC2D => Params(2, 5,.7,1),
				:LHE2D => Params(2,10,.7,1),
				:WC3D => Params(2, 40,.5,1),
				:LHE3D => Params(5,7,.7,1))

, "chevreulCanc" => Dict(:WC2D => Params(5, 3,.9,1),
				:LHE2D => Params(5, 3,.9,1),
				:WC3D => Params(2, 20,.5,1.4),
				:LHE3D => Params(5, 3,.9,1))

, "dungeon" => Dict(:WC2D => Params(6,10,.7,1.4),
				:LHE2D => Params(5,40,.7,1),
				:WC3D => Params(2,50,.7,1.4),
				:LHE3D => Params(5,50,.7,1))

, "gratings" => Dict(:WC2D => Params(2,6,.7,1),
				:LHE2D => Params(2,6,.7,1),
				:WC3D => Params(2,6,.7,1),
				:LHE3D => Params(2, 6,.7,1))

, "hong_shevell" => Dict(:WC2D => Params(5,20,.7,1),
				:LHE2D => Params(5,.5,.7,1),
				:WC3D => Params(10,30,.7,1),
				:LHE3D => Params(10,30,.7,1))

, "luminance" => Dict(:WC2D => Params(2,6,.7,1),
				:LHE2D => Params(2,6,.7,1),
				:WC3D => Params(2,6,.7,1),
				:LHE3D => Params(2,6,.7,1))

, "poggendorff" => Dict(:WC2D => Params(3,10,.5,1),
				:LHE2D => Params(3,10,.5,1),
				:WC3D => Params(3,10,.5,1),
				:LHE3D => Params(3,10,.5,1))

, "tilt" => Dict(:WC2D => Params(15,20,.7,1),
				:LHE2D => Params(15,20,.7,1),
				:WC3D => Params(15,20,.7,1),
				:LHE3D => Params(15,20,.7,1))
)

function batch_test(dict; args...)
	for k in keys(dict)
		if k ∉ ["hong_shevell", "tilt", "dungeon"]
			test(k, dict[k]; args...)
		end
	end

	test("hs1", dict["hong_shevell"]; args...)
	test("hs2", dict["hong_shevell"]; args...)

	test("dungeon1", dict["dungeon"]; args...)
	test("dungeon2", dict["dungeon"]; args...)

	test("tilt1", dict["tilt"]; args...)
	test("tilt2", dict["tilt"]; args...)
end

mkpath("results")
batch_test(illusions, lhe3d = true)

max_stars_repo_name: dprn/WCvsLHE

max_issues_repo_name: dprn/WCvsLHE

max_forks_repo_name: dprn/WCvsLHE


text:
module Expt8001
    vec(eltype::Type, len::Int64) = Vector{eltype}(len)

    # Return an empty matrix of the same type as `A`, with element `eltype`
    # and of size `sz`
    empty(A::Diagonal, eltype::Type, sz) = Diagonal(vec(eltype, sz[1]))
    empty(A::Bidiagonal, eltype::Type, sz) = Bidiagonal(vec(eltype, sz[1]), vec(eltype, sz[1]-1), A.isupper)
    empty(A::Tridiagonal, eltype::Type, sz) = empty(Tridiagonal, eltype, sz)

    empty(::Type{Diagonal}, eltype::Type, sz) = Diagonal(vec(eltype, sz[1]))
    empty(::Type{Bidiagonal}, eltype::Type, sz, isupper::Bool) = Bidiagonal(vec(eltype, sz[1]), vec(eltype, sz[1]-1), isupper)
    empty(::Type{UpperTriangular}, eltype::Type, sz) = UpperTriangular(zeros(Float64, sz))
    empty(::Type{LowerTriangular}, eltype::Type, sz) = LowerTriangular(zeros(Float64, sz))
    empty(::Type{Tridiagonal}, eltype::Type, sz) = Tridiagonal(vec(eltype, sz[1]-1), vec(eltype, sz[1]), vec(eltype, sz[1]-1))
    empty(::Type{Matrix}, eltype::Type, sz) = zeros(eltype, sz...)

    # Return an empty matrix which is the tightest type possible to hold
    # the result of performing * on `A` and `B`
    compat(*, A::Diagonal, B::Diagonal) = empty(A, Float64, (size(A,1), size(B,2)) )
    compat(*, A::Diagonal, B::AbstractMatrix) = empty(B, Float64, (size(A,1), size(B,2)) )
    compat(*, A::AbstractMatrix, B::Diagonal) = empty(A, Float64, (size(A,1), size(B,2)) )
    compat(*, A::Bidiagonal, B::Bidiagonal) =
        if A.isupper && B.isupper
            empty(UpperTriangular, Float64, (size(A,1), size(A,1)))
        elseif !A.isupper && !B.isupper
            empty(LowerTriangular, Float64, (size(A,1), size(A,1)))
        else
            empty(Tridiagonal, Float64, (size(A,1), size(A,1)) )
        end
    compat(*, A::Bidiagonal, B::Tridiagonal) = empty(Matrix, Float64, (size(A,1), size(A,1)) )
    compat(*, A::Tridiagonal, B::Bidiagonal) = empty(Matrix, Float64, (size(A,1), size(A,1)) )
    compat(*, A::Tridiagonal, B::Tridiagonal) = empty(Matrix, Float64, (size(A,1), size(A,1)) )

    # Indexing
    struct Index
        i::Int64
        j::Int64
    end

    # Represents the shape/structure/indexes of a banded matrix
    abstract type AbstractShape end
    struct DiagonalShape <: AbstractShape
        # An optimisation to keep iteration fast for diagonals
        m::Int64
        n::Int64
    end
    struct BandedShape{B, A} <: AbstractShape
        # B = number diagonals below main, A = number of diagonals above
        # B, A at type level in failed attempt to simplify compat
        m::Int64
        n::Int64
    end

    compose(*, A::DiagonalShape, B::DiagonalShape) = A
    compose(*, A::DiagonalShape, B::AbstractShape) = B
    compose(*, A::AbstractShape, B::DiagonalShape) = A
    # Type inference doesn't seem to be able to see through the
    # type-level addition here, so can't inline choice of
    # matrixfromshape()
    compose{b1,a1,b2,a2}(*, A::BandedShape{b1,a1}, B::BandedShape{b2,a2}) =
        BandedShape{b1+b2,a1+a2}(A.m, B.n)

    shapefrommatrix(D::Diagonal) = DiagonalShape(size(D,1), size(D,2))
    shapefrommatrix(B::Bidiagonal)::BandedShape =
        if B.isupper
            BandedShape{0,1}(size(B,1), size(B,2))
        else
            BandedShape{1,0}(size(B,1), size(B,2))
        end
    shapefrommatrix(T::Tridiagonal)::BandedShape = BandedShape{1,1}(size(T,1), size(T,2))

    iterateshape(D::DiagonalShape) = (Index(i,i) for i in 1:D.m)

    # FIXME this returns some indices more than once
    banded(n::Int64, m::Int64, below::Int64, above::Int64) =
        ( Index(i, clamp(i+j, 1, m))  for i in 1:n, j in -below:above )
    iterateshape{b,a}(B::BandedShape{b,a}) = banded(B.n, B.m, b, a)

    # matrixfromshape(D::DiagonalShape) = empty(Diagonal, Float64, (D.m, D.n))

    # matrixfromshape(B::BandedShape{0,0}) = empty(Diagonal, Float64, (B.m, B.n))
    # matrixfromshape(B::BandedShape{0,1}) = empty(Bidiagonal, Float64, (B.m, B.n), true)
    # matrixfromshape(B::BandedShape{1,0}) = empty(Bidiagonal, Float64, (B.m, B.n), false)
    # matrixfromshape(B::BandedShape{1,1}) = empty(Tridiagonal, Float64, (B.m, B.n))
    # matrixfromshape{a}(B::BandedShape{0,a}) = empty(UpperTriangular, Float64, (B.m, B.n))
    # matrixfromshape{b}(B::BandedShape{b,0}) = empty(LowerTriangular, Float64, (B.m, B.n))
    # matrixfromshape(B::BandedShape) = empty(Matrix, Float64, (B.m, B.n))

    indexes(*, A::AbstractMatrix, B::AbstractMatrix) =
        iterateshape(compose(*, shapefrommatrix(A), shapefrommatrix(B)))

    # Dot product of the `i`th row of `A` with the `j`th column of `B`
    dot(A::Diagonal,  i::Int64, B::AbstractMatrix, j::Int64) = @inbounds return A.diag[i] * B[i, j]

    @inline function dot(A::Bidiagonal, i::Int64, B::AbstractMatrix, j::Int64)
        (lim, evoff, Boff) = if A.isupper
            (size(A,2), 0, 1)
        else
            (1, -1, -1)
        end
        @inbounds return if i == lim
            A.dv[i] * B[i, j]
        else
            A.dv[i] * B[i, j]  + A.ev[i + evoff] * B[i + Boff, j]
        end
    end

    @inline function dot(A::Tridiagonal, i::Int64, B::AbstractMatrix, j::Int64)
        # FIXME Try copying diagonals of Tridiagonal to banded form and working with them instead
        # With zeros in place can remove branching at edges
        # FIXME out of bounds if A is 2x2?
        @inbounds return if i == 1
            A.d[i] * B[i, j] + A.du[i] * B[i+1, j]
        elseif i == size(A, 2)
            A.d[i] * B[i, j] + A.dl[end] * B[i-1, j]
        else
            A.d[i] * B[i, j] + A.dl[i-1] * B[i-1, j] + A.du[i] * B[i+1, j]
        end
    end

    # Actual multiplication
    function A_mul_B(A::AbstractMatrix, B::AbstractMatrix)
        T = compat(*, A, B)

        # Attempt to use shape composition to determine
        # target matrix type seems to defeat inliner
        # Rshape = compose(*, shapefrommatrix(A), shapefrommatrix(B))
        # T = matrixfromshape(Rshape)

        for ndx in indexes(*, A, B)
            @inbounds T[ndx.i, ndx.j] = dot(A, ndx.i, B, ndx.j)
        end
        T
    end

    # Performances tests to see where this approach looses time
    function a(A::AbstractMatrix, B::AbstractMatrix)
        # Cost of indexes()
        sum = 0
        for ndx in indexes(*, A, B)
            sum += ndx.i
        end
        sum
    end

    function b(A::AbstractMatrix, B::AbstractMatrix)
        # Cost of dot()
        sum = 0.0
        for i in 1:100, j in 1:4
            sum += dot(A, i, B, j)
        end
        sum
    end

    function c(A::AbstractMatrix, B::AbstractMatrix)
        # Cost of indexes() and dot()
        sum = 0.0
        for ndx in indexes(*, A, B)
            @inbounds sum += dot(A, ndx.i, B, ndx.j)
        end
        sum
    end

end

ops = [(Expt8001.A_mul_B, A_mul_B!)]

mats = [n -> Diagonal(randn(n)),
        n -> Bidiagonal(randn(n), randn(n-1), true),
        n -> Bidiagonal(randn(n), randn(n-1), false),
        n -> Tridiagonal(randn(n-1), randn(n), randn(n-1))
        ]

for (opnew, opold) in ops, m1f in mats, m2f in mats
    n = 10
    m1, m2 = m1f(n), m2f(n)

    fullres = Matrix{Float64}(n, n)
    try
        opold(fullres, full(m1), full(m2))
    catch err
        print("FULL OP FAILED ", typeof(err), "\t: ", typeof(m1), " ", opold, " ", typeof(m2), "\n")
        continue
    end
    try
        if !(full(opnew(m1, m2)) ≈ fullres)
            print("NOT EQUAL TO FULL OP ", typeof(m1), " ", opnew, " ", typeof(m2), " ", "\n")
        end
    catch err
        println("\n>>> ", typeof(err), "\t: ", typeof(m1), " ", opnew, " ", typeof(m2))
        showerror(STDOUT, err) #, catch_backtrace())
        println()
    end
end


text:
# ---
# title: 793. Preimage Size of Factorial Zeroes Function
# id: problem793
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Binary Search
# link: <https://leetcode.com/problems/preimage-size-of-factorial-zeroes-function/description/>
# hidden: true
# ---
#
# Let `f(x)` be the number of zeroes at the end of `x!`. (Recall that `x! = 1 *
# 2 * 3 * ... * x`, and by convention, `0! = 1`.)
#
# For example, `f(3) = 0` because 3! = 6 has no zeroes at the end, while `f(11)
# = 2` because 11! = 39916800 has 2 zeroes at the end. Given `K`, find how many
# non-negative integers `x` have the property that `f(x) = K`.
#
#
#
#     **Example 1:**
#     Input: K = 0
#     Output: 5
#     Explanation: 0!, 1!, 2!, 3!, and 4! end with K = 0 zeroes.
#
#     **Example 2:**
#     Input: K = 5
#     Output: 0
#     Explanation: There is no x such that x! ends in K = 5 zeroes.
#
#
# **Note:**
#
#   * `K` will be an integer in the range `[0, 10^9]`.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
using Test, AdmittanceModels, LinearAlgebra
import .CircuitExample, .PSOModelExample, .CircuitCascadeExample, .LRCExample
const ce = CircuitExample
const ae = PSOModelExample
const cce = CircuitCascadeExample
const lrce = LRCExample
using SparseArrays: SparseMatrixCSC

@testset "PSOModel" begin
    @test length(ae.pso.ports) == 2
    for m in get_Y(ae.pso)
        @test size(m) == (4,4)
        @test issymmetric(m)
    end
    @test size(get_P(ae.pso)) == (4, 2)
    @test get_Q(ae.pso) == get_P(ae.pso)
end

@testset "coordinate_transform" begin
    coord_matrix_from = coordinate_matrix(ce.circuit)
    coord_matrix_to = coordinate_matrix(ce.circuit, ce.tree)
    transform = AdmittanceModels.coordinate_transform(coord_matrix_from, coord_matrix_to)
    @test AdmittanceModels.coordinate_transform(coord_matrix_from) == I
    @test AdmittanceModels.coordinate_transform(coord_matrix_to) != I
    @test transform * AdmittanceModels.coordinate_transform(coord_matrix_from) ==
        AdmittanceModels.coordinate_transform(coord_matrix_to)
end

@testset "Circuit to PSOModel" begin
    converted_pso = PSOModel(ce.circuit, [(ce.Q0_SIGNAL, ce.Q0_GROUND),
                                          (ce.Q1_SIGNAL, ce.Q1_GROUND)],
                             get_ports(ae.pso), ce.tree)
    @test ae.pso == converted_pso
end

@testset "cascade, unite_ports, open_ports, short_ports" begin
    # cascade and unite_ports
    circ = cascade(cce.circuit0, cce.circuit1, cce.circuit2)
    united_circ = unite_vertices(circ, 0, 4, 8)
    united_circ = unite_vertices(united_circ, 1, 5, 9)
    port_edges = [(united_circ.vertices[1], v) for v in united_circ.vertices[2:end]]
    port_names = united_circ.vertices[2:end]
    pso_from_circuit = PSOModel(united_circ, port_edges, port_names)
    pso0 = PSOModel(cce.circuit0, [(0,1), (0,2), (0, 3)], [1, 2, 3])
    pso1 = PSOModel(cce.circuit1, [(4,5), (4,6), (4,7)], [5, 6, 7])
    pso2 = PSOModel(cce.circuit2, [(8,9), (8,10), (8,11)], [9, 10, 11])
    pso = cascade(pso0, pso1, pso2)
    united_pso = unite_ports(pso, 1, 5, 9)
    transform = transpose(get_P(united_pso)/get_P(pso_from_circuit))
    pso_transform = apply_transform(pso_from_circuit, transform)
    @test pso_transform ≈ united_pso
    # open_ports
    port_edges = port_edges[1:2]
    port_names = port_names[1:2]
    pso_from_circuit = PSOModel(united_circ, port_edges, port_names)
    pso_transform = apply_transform(pso_from_circuit, transform)
    united_pso = open_ports(united_pso, get_ports(united_pso)[3:end])
    @test pso_transform ≈ united_pso
    # short_ports
    united_circ = unite_vertices(united_circ, 0, 2, 6)
    port_edges = [(united_circ.vertices[1], v) for v in united_circ.vertices[2:end]]
    port_names = united_circ.vertices[2:end]
    pso_from_circuit = PSOModel(united_circ, port_edges, port_names)
    united_pso = unite_ports(pso, 1, 5, 9)
    united_pso = short_ports(united_pso, 2, 6)
    transform = transpose(get_P(united_pso)/get_P(pso_from_circuit))
    pso_transform = apply_transform(pso_from_circuit, transform)
    @test pso_transform ≈ united_pso
end

@testset "lossy_modes_dense and lossless_modes_dense" begin
    eigenvalues, eigenvectors = lossy_modes_dense(lrce.pso)
    @test length(eigenvalues) == 1
    @test size(eigenvectors) == (1,1)
    decay_rate = -2 * real(eigenvalues[1])
    frequency = imag(eigenvalues[1])/(2π)
    @test frequency ≈ lrce.ω/(2π)
    @test decay_rate ≈ lrce.κ
    eigenvalues, eigenvectors = lossless_modes_dense(lrce.pso)
    @test length(eigenvalues) == 1
    @test size(eigenvectors) == (1,1)
    decay_rate = -2 * real(eigenvalues[1])
    frequency = imag(eigenvalues[1])/(2π)
    @test frequency ≈ 1/(sqrt(lrce.l * lrce.c) * 2π)
    @test decay_rate == 0
end

@testset "lossless_modes_sparse and tline resonator modes" begin
    ν, Z0 = 1.2e8, 50.0
    L = 5e-3
    δ = 50e-6
    resonator = TransmissionLine(["0", "1"], ν, Z0, L, δ=δ)
    pso = PSOModel(resonator)
    pso = short_ports(pso, pso.ports)
    for Y in get_Y(pso) # pso model is sparse
        @test count(!iszero, Y)/(size(Y,1) * size(Y, 2)) < .05
    end
    num_modes = 5
    values_s, vectors_s = lossless_modes_sparse(pso, num_modes=num_modes)
    values_d, vectors_d = lossless_modes_dense(pso)
    values_d, vectors_d = values_d[1:num_modes], vectors_d[:, 1:num_modes]
    @test isapprox(values_s, values_d, rtol=1e-6)
    for i in 1:5
        v_s, v_d = vectors_s[:, i], vectors_d[:, i]
        j = argmax(abs.(v_s))
        v_s /= v_s[j]
        v_d /= v_d[j]
        @test isapprox(v_s, v_d, rtol=1e-2)
    end
    @test isapprox(ν/(2*L) * (1:num_modes), imag.(values_s)/(2π), rtol=1e-3)
    @test isapprox(ν/(2*L) * (1:num_modes), imag.(values_d)/(2π), rtol=1e-3)
end

@testset "canonical_gauge" begin
    circ = cce.circuit0
    port_edges = [(1,2), (2,3)]
    port_names = [1, 2]
    pso = PSOModel(circ, port_edges, port_names)
    P = canonical_gauge(pso).P
    @test typeof(pso.P) <: SparseMatrixCSC
    @test typeof(P) <: Matrix
    m, n = size(P)
    @test P ≈ [Matrix(I, n, n); zeros(m-n, n)]
end


text:
using SSE
using Test

@testset "lattice" begin
    lt = Chain(20)
    @test neighbors(lt, 3) == [2, 4]
    @test neighbors(lt, 3) == [2, 4]
    @test neighbors(lt, 20) == [19, 1]
    @test neighbors(lt, 1) == [20, 2]
end


text:
function RinexLoader(arg0::InputStream, arg1::JString)
    return RinexLoader((InputStream, JString), arg0, arg1)
end

function RinexLoader(arg0::InputStream, arg1::JString, arg2::TimeScales)
    return RinexLoader((InputStream, JString, TimeScales), arg0, arg1, arg2)
end

function RinexLoader(arg0::JString)
    return RinexLoader((JString,), arg0)
end

function RinexLoader(arg0::JString, arg1::DataProvidersManager, arg2::TimeScales)
    return RinexLoader((JString, DataProvidersManager, TimeScales), arg0, arg1, arg2)
end

function get_observation_data_sets(obj::RinexLoader)
    return jcall(obj, "getObservationDataSets", List, ())
end


text:
"""
    Dogleg()

A trust region sub-problem solver that assumes positive definite hessians (exact
or quasi-Newton approximations such as BFGS or variants).
"""
struct Dogleg{T} <: TRSPSolver
    γ::T # unused, for double-dogleg
end
Dogleg() = Dogleg(nothing)

function (dogleg::Dogleg)(∇f, H, Δ, p, scheme; abstol=1e-10, maxiter=50)
    T = eltype(p)
    n = length(∇f)

    # find the Cauchy point; assumes ∇f is not ≈ 0
    d_cauchy = - ∇f*norm(∇f)^2/(∇f'*H*∇f)

    # If it lies outside of the trust region, accept the Cauchy point and
    # move on
    norm_d_cauchy = norm(d_cauchy)
    if norm_d_cauchy ≥ Δ
        shrink = Δ/norm_d_cauchy # inv(Δ/norm_d_cauchy) puts it on the border
        p .= d_cauchy.*shrink
        interior = false
    else
        # Else, calculate (Quasi-)Newton step. If this is interior, then take the
        # step. Otherwise find where the dog-leg path crosses the trust region

        # find the (quasi-)Newton step
        d_qn = copy(p)# -H\∇f # use find step and scheme
        find_direction!(d_qn, H, ∇f, scheme)
        norm_d_qn = norm(d_qn)
        if norm_d_qn ≤ Δ
            p .= d_qn
            if norm_d_qn < Δ
                interior = true
            else
                interior = false
            end
        else
            # we now solve a quadratic to find the step size t from d_cauchy
            # towards d_qn. We use a numerically stable way of doing this (see
            # any numerical analysis text book) See [NW, p. 75] for the expression
            # giving rise to the quadratic equation
            dot_cachy_qn = dot(d_cauchy, d_qn)
            # a is ||d_c - d_n||^2 expanded into scalar operations
            a = norm_d_cauchy^2 + norm_d_qn^2 - 2*dot_cachy_qn
            b = -dot_cachy_qn * norm_d_cauchy^2
            c = norm_d_cauchy^2 - Δ^2 # move the rhs over
            q = -(b + sign(b)*√(b^2-4*a*c))/2
            # since we know that c is necessarily negative (since d_cauchy was
            # not at the border) the discriminant is positive, and there are two
            # roots - pick the positive one. There has to be one positive and one
            # negative.
            if b > 0
                # if b is positive, q is negative. Then if c is negative, we must
                # have that c / q is the positive root.
                t = c / q
            elseif b < 0
                # else, the other root must be positive
                t = q / a
            else
                t = T(0)
            end
            p .= d_cauchy .+ t.*d_qn
            interior = false
        end
    end
    m = dot(∇f, p) + dot(p, H * p)/2

    return (p=p, mz=m, interior=interior, λ=nothing, hard_case=nothing, solved=true, Δ=Δ)
end


text:
#=
def rounds(m, T, scale):      # generate a price schedule vector
    mask = []
    for i in range(1, m):
    	print(i)
    	print(int(scale * math.ceil(logx(T, m - i))))
    	print(np.full(int(scale * math.ceil(logx(T, m - i))), i - 1))
        mask.extend( np.full(int(scale * math.ceil(logx(T, m - i))), i - 1) )
        print(mask)
    return np.append(mask, np.full(T - len(mask), m - 1))

h_vec = list(demand_hypotheses(np.linspace(start_a, end_a, 4),
	np.linspace(start_b, end_b, 4)))
=#

function intervals(m, T, scale)  # generate a price schedule vector
    mask = []
    for i in 1:m-1
        append!(mask, repeat([i - 1], Int(scale * ceil(logx(T, m - i)))) )
    end
    return append!(mask, repeat([m - 1], T - length(mask)))
end


text:
using Gridap

# Geometry
domain = (0,1,0,1); cells = (10,10)
model = CartesianDiscreteModel(domain,cells)

# Integration
degree = 2
Ω = Triangulation(model)
Γ = BoundaryTriangulation(model)
dΩ = LebesgueMeasure(Ω,degree)
dΓ = LebesgueMeasure(Γ,degree)

# Manufactured solution
u(x) = x[1]+x[2]

# Parameter-dependent weak form
a(p,u,v) = ∫( (p*p)*∇(v)⋅∇(u) )*dΩ
l(v) = ∫( -v*Δ(u) )*dΩ

# Objective function in terms of primal solution
# (squared distance wrt exact solution u on Ω and Γ
# as an example involving different domains)
# Note: the following function returns cell contributions.
# The final objective is computed as sum(j(v))
j(v) = ∫( abs2(u-v) )*dΩ + ∫( abs2(u-v) )*dΓ

# FE spaces: P for params, V, U for primal solution
P = FESpace(model,ReferenceFE(:Lagrangian,Float64,0),conformity=:L2)
V = TestFESpace(model,ReferenceFE(:Lagrangian,Float64,1),dirichlet_tags="boundary")
U = TrialFESpace(V,u)

using Gridap.FESpaces

# Objective function (scalar), plus its gradient (vector) in terms of p
# fp, ∇fp = f_∇f(p)
function f_∇f(p)
  # Direct solve
  a_at_p(u,v) = a(p,u,v)
  op = AffineFEOperator(a_at_p,l,U,V)
  u = solve(op)
  # Adjoint solve
  A = get_matrix(op)
  dj_du = ∇(j)(u) # autodiff
  b = assemble_vector(dj_du,V)
  λ = FEFunction(V,A\b) # assuming A self-adjoint
  # we want to have something like this
  # op_t = AdjointFEOperator(op, ∇(j), selfadjoint=true )
  # u, lambda = solve(op,op_t)
  # for the moment, we can do more low level things at the driver level, e.g., extracting matrices, etc for the adjoint
  # in any case, it is not needed for self-adjoint problems
  # Objective
  fp = sum(j(u))
  # Its gradient
  ∂a_∂p = ∇(x -> a(x,u,λ))(p) # Autodiff
  ∇fp = -1*assemble_vector(∂a_∂p,P)
  fp, ∇fp
end

# run for a given realization of p
p = FEFunction(P,rand(num_free_dofs(P)).+1)
fp, ∇fp = f_∇f(p)
# Print results
@show fp
writevtk(Ω,"trian",cellfields=["∇fp"=>FEFunction(P,∇fp)])

function df_dp(uvec,pvec,λvec)
  da(p,u,v,dp) = (2*p*dp)*∇(v)⋅∇(u)
  uh = FEFunction(U,uvec)
  ph = FEFunction(P,pvec)
  λh = FEFunction(V,λvec)
  t = FESource((dp)->da(ph,uh,λh,dp),trian,quad)
  op = AffineFEOperator(P,Q,t)
  get_vector(op)
end

max_stars_repo_name: santiagobadia/GridapGeosciences

max_issues_repo_name: santiagobadia/GridapGeosciences

max_forks_repo_name: santiagobadia/GridapGeosciences


text:
using OMEinsum
using LightGraphs
using Test, Random
using SparseArrays
using OMEinsumContractionOrders

@testset "kahypar" begin
    Random.seed!(2)
    function random_regular_eincode(n, k)
        g = LightGraphs.random_regular_graph(n, k)
        ixs = [minmax(e.src,e.dst) for e in LightGraphs.edges(g)]
        return EinCode((ixs..., [(i,) for i in LightGraphs.vertices(g)]...), ())
    end

    g = random_regular_graph(220, 3)
    rows = Int[]
    cols = Int[]
    for (i,edge) in enumerate(edges(g))
        push!(rows, edge.src, edge.dst)
        push!(cols, i, i)
    end
    graph = sparse(rows, cols, ones(Int, length(rows)))
    sc_target = 28.0
    group1, group2 = OMEinsumContractionOrders.kahypar_partitions_sc(graph, collect(1:size(graph, 1)); log2_sizes=fill(1, size(graph, 2)), sc_target=sc_target, imbalances=[0.0:0.02:0.8...])
    @test OMEinsumContractionOrders.group_sc(graph, group1) <= sc_target
    @test OMEinsumContractionOrders.group_sc(graph, group2) <= sc_target
    sc_target = 27.0
    group11, group12 = OMEinsumContractionOrders.kahypar_partitions_sc(graph, group1; log2_sizes=fill(1, size(graph, 2)), sc_target=37.0, imbalances=[0.0:0.02:1.0...])
    @test OMEinsumContractionOrders.group_sc(graph, group11) <= sc_target
    @test OMEinsumContractionOrders.group_sc(graph, group12) <= sc_target

    code = random_regular_eincode(220, 3)
    res = optimize_kahypar(code,uniformsize(code, 2); max_group_size=50, sc_target=30)
    tc, sc = OMEinsum.timespace_complexity(res, uniformsize(code, 1))
    @test sc <= 30

    # contraction test
    code = random_regular_eincode(50, 3)
    codeg = optimize_kahypar(code, uniformsize(code, 2); max_group_size=10, sc_target=12)
    codek = optimize_greedy(code, uniformsize(code, 2))
    tc, sc = OMEinsum.timespace_complexity(codek, uniformsize(code, 1))
    @test sc <= 12
    xs = [[2*randn(2, 2) for i=1:75]..., [randn(2) for i=1:50]...]
    resg = codeg(xs...)
    resk = codek(xs...)
    @test resg ≈ resk
end