text:
mutable struct LocalApproximationValueIterationSolver{I<:LocalFunctionApproximator, RNG<:AbstractRNG} <: Solver
    interp::I # Will be copied over by value to each policy
    max_iterations::Int64 # max number of iterations
    belres::Float64 # the Bellman Residual
    verbose::Bool # Whether to print while solving or not
    rng::RNG # Seed if req'd
    is_mdp_generative::Bool # Whether to treat underlying MDP model as generative
    n_generative_samples::Int64 # If underlying model generative, how many samples to use
end

# Default constructor
function LocalApproximationValueIterationSolver(interp::I;
                                                max_iterations::Int64=100, belres::Float64=1e-3,
                                                verbose::Bool=false, rng::RNG=Random.GLOBAL_RNG,
                                                is_mdp_generative::Bool=false, n_generative_samples::Int64=0) where {I<:LocalFunctionApproximator, RNG<:AbstractRNG}
    return LocalApproximationValueIterationSolver(interp,max_iterations, belres, verbose, rng, is_mdp_generative, n_generative_samples)
end

# Unparameterized constructor just for getting requirements
function LocalApproximationValueIterationSolver()
    throw(ArgumentError("LocalApproximationValueIterationSolver needs a LocalFunctionApproximator object for construction!"))
end

# NOTE : We work directly with the value function
# And extract actions at the end by using the interpolation object
mutable struct LocalApproximationValueIterationPolicy{I<:LocalFunctionApproximator, RNG<:AbstractRNG} <: Policy
    interp::I # General approximator to be used in VI
    action_map::Vector # Maps the action index to the concrete action type
    mdp::Union{MDP,POMDP} # Uses the model for indexing in the action function
    is_mdp_generative::Bool # (Copied from solver.is_mdp_generative)
    n_generative_samples::Int64 # (Copied from solver.n_generative_samples)
    rng::RNG # (Copied from solver.rng)
end

# The policy can be created using the MDP and solver information
# The policy's function approximation object (interp) is obtained by deep-copying over the
# solver's interp object. The other policy parameters are also obtained from the solver
function LocalApproximationValueIterationPolicy(mdp::Union{MDP,POMDP},
                                                solver::LocalApproximationValueIterationSolver)
    return LocalApproximationValueIterationPolicy(deepcopy(solver.interp),ordered_actions(mdp),mdp,
                                                  solver.is_mdp_generative,solver.n_generative_samples,solver.rng)
end

@POMDP_require solve(solver::LocalApproximationValueIterationSolver, mdp::Union{MDP,POMDP}) begin

    P = typeof(mdp)
    S = statetype(P)
    A = actiontype(P)
    @req discount(::P)
    @req n_actions(::P)
    @subreq ordered_actions(mdp)

    @req action_index(::P, ::A)
    @req actions(::P, ::S)
    as = actions(mdp)
    a = first(as)

    @req convert_s(::Type{S},::V where V <: AbstractVector{Float64},::P)
    @req convert_s(::Type{V} where V <: AbstractVector{Float64},::S,::P)

    # Have different requirements depending on whether solver MDP is generative or explicit
    if solver.is_mdp_generative
        @req generate_sr(::P, ::S, ::A, ::typeof(solver.rng))
    else
        @req transition(::P, ::S, ::A)
        pts = get_all_interpolating_points(solver.interp)
        pt = first(pts)
        ss = POMDPs.convert_s(S,pt,mdp)
        dist = transition(mdp, ss, a)
        D = typeof(dist)
        @req support(::D)
    end

end

function solve(solver::LocalApproximationValueIterationSolver, mdp::Union{MDP,POMDP})

    @warn_requirements solve(solver,mdp)

    # Ensure that generative model has a non-zero number of samples
    if solver.is_mdp_generative
        @assert solver.n_generative_samples > 0
    end

    # Solver parameters
    max_iterations = solver.max_iterations
    belres = solver.belres
    discount_factor = discount(mdp)

    # Initialize the policy
    policy = LocalApproximationValueIterationPolicy(mdp,solver)

    total_time::Float64 = 0.0
    iter_time::Float64 = 0.0

    # Get attributes of interpolator
    # Since the policy object is created by the solver, it directly
    # modifies the value of the interpolator of the created policy
    num_interps::Int = n_interpolating_points(policy.interp)
    interp_points::Vector = get_all_interpolating_points(policy.interp)
    interp_values::Vector = get_all_interpolating_values(policy.interp)

    # Obtain the vector of states by converting the corresponding
    # vector of interpolation points/samples to the state type
    # using the user-provided convert_s function
    S = statetype(typeof(mdp))
    interp_states = Vector{S}(undef, num_interps)
    for (i,pt) in enumerate(interp_points)
        interp_states[i] = POMDPs.convert_s(S, pt, mdp)
    end

    # Outer loop for Value Iteration
    for i = 1 : max_iterations
        residual::Float64 = 0.0
        iter_time = @elapsed begin

        for (istate,s) in enumerate(interp_states)
            sub_aspace = actions(mdp,s)

            if isterminal(mdp, s)
                interp_values[istate] = 0.0
            else
                old_util = interp_values[istate]
                max_util = -Inf

                for a in sub_aspace
                    iaction = action_index(mdp,a)
                    u::Float64 = 0.0

                    # Do bellman backup based on generative / explicit model
                    if solver.is_mdp_generative
                        # Generative Model
                        for j in 1:solver.n_generative_samples
                            sp, r = generate_sr(mdp, s, a, solver.rng)
                            u += r

                            # Only interpolate sp if it is non-terminal
                            if !isterminal(mdp,sp)
                                sp_point = POMDPs.convert_s(Vector{Float64}, sp, mdp)
                                u += discount_factor*compute_value(policy.interp, sp_point)
                            end
                        end
                        u = u / solver.n_generative_samples
                    else
                        # Explicit Model
                        dist = transition(mdp,s,a)
                        for (sp, p) in weighted_iterator(dist)
                            p == 0.0 ? continue : nothing
                            r = reward(mdp, s, a, sp)
                            u += p*r

                            # Only interpolate sp if it is non-terminal
                            if !isterminal(mdp,sp)
                                sp_point = POMDPs.convert_s(Vector{Float64}, sp, mdp)
                                u += p * (discount_factor*compute_value(policy.interp, sp_point))
                            end
                        end # next-states
                    end

                    max_util = (u > max_util) ? u : max_util
                end #action

                # Update this interpolant value
                interp_values[istate] = max_util
                util_diff = abs(max_util - old_util)
                util_diff > residual ? (residual = util_diff) : nothing
            end
        end #state

        end #time
        total_time += iter_time
        solver.verbose ? @printf("[Iteration %-4d] residual: %10.3G | iteration runtime: %10.3f ms, (%10.3G s total)\n", i, residual, iter_time*1000.0, total_time) : nothing
        residual < belres ? break : nothing

    end #main
    return policy
end

function value(policy::LocalApproximationValueIterationPolicy, s::S) where S

    # Call the conversion function on the state to get the corresponding vector
    # That represents the point at which to interpolate the function
    s_point = POMDPs.convert_s(Vector{Float64}, s, policy.mdp)
    val = compute_value(policy.interp, s_point)
    return val
end

# Not explicitly stored in policy - extract from value function interpolation
function action(policy::LocalApproximationValueIterationPolicy, s::S) where S

    mdp = policy.mdp
    best_a_idx = -1
    max_util = -Inf
    sub_aspace = actions(mdp,s)
    discount_factor = discount(mdp)

    for a in iterator(sub_aspace)

        iaction = action_index(mdp, a)
        u::Float64 = action_value(policy,s,a)

        if u > max_util
            max_util = u
            best_a_idx = iaction
        end
    end

    return policy.action_map[best_a_idx]
end

# Compute the action-value for some state-action pair
# This is also used in the above function
function action_value(policy::LocalApproximationValueIterationPolicy, s::S, a::A) where {S,A}

    mdp = policy.mdp
    discount_factor = discount(mdp)

    u::Float64 = 0.0

    # As in solve(), do different things based on whether
    # mdp is generative or explicit
    if policy.is_mdp_generative
        for j in 1:policy.n_generative_samples
            sp, r = generate_sr(mdp, s, a, policy.rng)
            sp_point = POMDPs.convert_s(Vector{Float64}, sp, mdp)
            u += r + discount_factor*compute_value(policy.interp, sp_point)
        end
        u = u / policy.n_generative_samples
    else
        dist = transition(mdp,s,a)
        for (sp, p) in weighted_iterator(dist)
            p == 0.0 ? continue : nothing
            r = reward(mdp, s, a, sp)
            u += p*r

            # Only interpolate sp if it is non-terminal
            if !isterminal(mdp,sp)
                sp_point = POMDPs.convert_s(Vector{Float64}, sp, mdp)
                u += p*(discount_factor*compute_value(policy.interp, sp_point))
            end
        end
    end

    return u
end


text:
"""
  function calculate_D(sMMProblem::MSMProblem, theta0::Array{Float64,1}; method::Symbol = :central)

Function to calculate the jacobian of the simulated moments.
If the simulation is long enough, this provideds a good approximation for
the expected value of the jacobian. The output is the "D" matrix in the terminology
of Gouriéroux and Monfort (1996).
"""
function calculate_D(sMMProblem::MSMProblem, theta0::Array{Float64,1})

  #Calculus.jacobian(x -> sMMProblem.simulate_empirical_moments_array(x), theta0, method)
  jacobian(central_fdm(5, 1), x -> sMMProblem.simulate_empirical_moments_array(x), theta0)[1]

end

"""
  calculate_Avar!(sMMProblem::MSMProblem, theta0::Array{Float64,1}; tau::Float64 = 1.0)

Calculate the asymptotic variance of the SMM estimator using the sandwich formula.
This function assume that your model is simulating time series.
As a result, the asymptotic variance of the SMM estimator is (1+tau) times
the asymptotic variance for the GMM estimator, where tau is the ratio of the sample size
to the size of the simulated sample. See Duffie and Singleton (1993) and Gouriéroux and Monfort (1996).
"""
function calculate_Avar!(sMMProblem::MSMProblem, theta0::Array{Float64,1}; tau::Float64 = 1.0)

  # Safety Checks
  if isempty(sMMProblem.Sigma0) == true
    error("Please initialize sMMProblem.Sigma0 using the function set_Sigma0!.")
  end

  # Jacobian matrix
  D = calculate_D(sMMProblem, theta0)

  Sigma1 = transpose(D)*sMMProblem.W*D

  Sigma2 = transpose(D)*sMMProblem.W*sMMProblem.Sigma0*sMMProblem.W*D

  inv_Sigma1 = inv(Sigma1)

  # sandwich formula, where tau captures the extra noise introduced by simulation,
  # compared to the usual GMM estimator.
  # I use the function Symmetric, because otherwise a test issymmetric()
  # on the asymptotic variance will return "false" because of numerical approximations
  sMMProblem.Avar = Symmetric((1+tau)*inv_Sigma1*Sigma2*inv_Sigma1)

end

"""
  calculate_se(sMMProblem::MSMProblem, tData::Int64, tSimulation::Int64)

Function to calculate the standard error associated to the the ith parameter,
respecting the ordering given by the ordered dictionary sMMProblem.priors
"""
function calculate_se(sMMProblem::MSMProblem, tData::Int64, i::Int64)

  # Safety Checks
  if isempty(sMMProblem.Avar) == true
    error("Please caclulate the asymptotic variance using the function calculate_Avar!.")
  end

  sqrt((1/tData)*sMMProblem.Avar[i,i])

end

"""
  calculate_t(sMMProblem::MSMProblem, tData::Int64, tSimulation::Int64)

Function to calculate the t-statistic associated to following test:
H0: theta_i = 0
H1: theta_i != 0
The ordering of parameters is the one given by the ordered dictionary sMMProblem.priors
"""
function calculate_t(sMMProblem::MSMProblem, theta0::Array{Float64,1}, tData::Int64, i::Int64)

  # Safety Checks
  if isempty(sMMProblem.Avar) == true
    error("Please caclulate the asymptotic variance using the function calculate_Avar!.")
  end

  theta0[i]/calculate_se(sMMProblem, tData, i)

end

"""
  calculate_pvalue(sMMProblem::MSMProblem, tData::Int64, tSimulation::Int64)

Function to calculate the p-value associated to following test:
H0: theta_i = 0
H1: theta_i != 0
The ordering of parameters is the one given by the ordered dictionary sMMProblem.priors
"""
function calculate_pvalue(sMMProblem::MSMProblem, theta0::Array{Float64,1}, tData::Int64, i::Int64)

  # Safety Checks
  if isempty(sMMProblem.Avar) == true
    error("Please caclulate the asymptotic variance using the function calculate_Avar!.")
  end

  t =  calculate_t(sMMProblem, theta0, tData, i)

  # asymptotically normally distributed N(0,1)
  d = Normal(0,1)

  # p-value  = 2 times area to the right of t
  pvalue = 2*(1.0 - cdf(d, t))

  return pvalue

end

"""
  calculate_CI(sMMProblem::MSMProblem, theta0::Array{Float64,1}, tData::Int64, i::Int64)

Function to calculate an alpha confidence interval for the ith parameter.
"""
function calculate_CI(sMMProblem::MSMProblem, theta0::Array{Float64,1}, tData::Int64, i::Int64, alpha::Float64)

  # Safety Checks
  if isempty(sMMProblem.Avar) == true
    error("Please caclulate the asymptotic variance using the function calculate_Avar!.")
  end

  # standard error
  se = calculate_se(sMMProblem, tData, i)

  # asymptotically normally distributed N(0,1)
  d = Normal(0,1)

  critical_value = - quantile(d, alpha)

  #return lower and upper bound of the confidence interval
  return theta0[i] - se*critical_value, theta0[i] + se*critical_value

end

#=
function summary_table(sMMProblem::MSMProblem, theta0::Array{Float64,1}, tData::Int64, alpha::Float64)

  # Safety Checks
  if isempty(sMMProblem.Avar) == true
    error("Please caclulate the asymptotic variance using the function calculate_Avar!.")
  end

  df = DataFrame(Estimate = Float64[], StdError = Float64[], tValue = Float64[], pValue = Float64[], ConfIntervalLower = Float64[], ConfIntervalUpper = Float64[])

  for i = 1:length(theta0)

    se = calculate_se(sMMProblem, tData, i)
    t = calculate_t(sMMProblem, theta0, tData, i)
    p = calculate_pvalue(sMMProblem, theta0, tData, i)
    CI_lower, CI_upper = calculate_CI(sMMProblem, theta0, tData, i, alpha)

    push!(df, [theta0[i], se, t, p, CI_lower, CI_upper])

  end

  return df

end
=#

"""
  summary_table(sMMProblem::MSMProblem, theta0::Array{Float64,1}, tData::Int64, alpha::Float64)

Returns a table with estimates and related statistics of the model.
`alpha` is the significance level.
"""
function summary_table(sMMProblem::MSMProblem, theta0::Array{Float64,1}, tData::Int64, alpha::Float64)

  # Safety Checks
  if isempty(sMMProblem.Avar) == true
    error("Please caclulate the asymptotic variance using the function calculate_Avar!.")
  end

  se = zeros(length(theta0))
  t = zeros(length(theta0))
  p = zeros(length(theta0))
  CI_lower = zeros(length(theta0))
  CI_upper = zeros(length(theta0))

  for i = 1:length(theta0)

    se[i] = calculate_se(sMMProblem, tData, i)
    t[i] = calculate_t(sMMProblem, theta0, tData, i)
    p[i] = calculate_pvalue(sMMProblem, theta0, tData, i)
    CI_lower[i], CI_upper[i] = calculate_CI(sMMProblem, theta0, tData, i, alpha)

  end

  df = CoefTable(hcat(theta0,se,t,p,CI_lower,CI_upper),
              ["Coef.","Std. Error","t","Pr(>|t|)","CI Lower","CI Upper"],
              ["x$i" for i = 1:length(theta0)], 4, 3)

  return df

end

"""
  cov_NW(data::Matrix; l::Int64 = -1)

Function to calculate Newey–West (1987) variance-covariance matrix. `data` is the
data matrix with each row representing a time period and each column representing
a variable. `l` is the nummber of lags to include.
"""
function cov_NW(data::Matrix; l::Int64 = -1)
    # See: Newey, Whitney K; West, Kenneth D (1987)
    # data: rows = time; columns = dimension
    # l: number of lags to include
    #number of periods
    T=size(data,1);
    nlag=l;
    #Default value if user does not provide l
    if nlag == -1
        nlag=Int(min(floor(1.2*T^(1/3)),T));
    end

    #Demean each column
    data = data .- repeat(mean(data, dims=1), outer = [T, 1]);

    # Newey West weights:
    w=(nlag+1 .- collect(0:nlag))./(nlag+1);

    # Variance-covariance matrix when no serial correlation
    V=transpose(data)*data./(T-1);

    # Additional terms (non-zeros if serial correlation)
    for i=1:nlag
        Gammai=(transpose(data[(i+1):T,:])*data[1:T-i,:])./(T-1);
        V = V + w[i+1].*(Gammai + transpose(Gammai));
    end

    return V
end

"""
  J_test(sMMProblem::MSMProblem, theta0::Array{Float64,1}, tData::Int64, tSimData::Int64, alpha::Float64)

Run a J-test, also called a test for over-identifying restrictions. The null hypothesis that the model is “valid”.
The alternative hypothesis that model is “invalid”. For the test to be valid,
the weigth matrix must converge in probability to the efficient weighting matrix
(Sigma0^-1).

#Ouput:
* J: value of the J-statistic
* c: critical value associated to the J-test
"""
function J_test(sMMProblem::MSMProblem, theta0::Array{Float64,1}, tData::Int64, tSimData::Int64, alpha::Float64)

  df=length(keys(sMMProblem.empiricalMoments)) - length(theta0)
  #Chi² distribution with k-l degrees of freedom, where
  d_chi=Chi(df)
  #Calculate J
  simulatedMoments=sMMProblem.simulate_empirical_moments(theta0)
  # to store the distance between empirical and simulated moments
  arrayDistance = zeros(length(keys(sMMProblem.empiricalMoments)))
  for (indexMoment, k) in enumerate(keys(sMMProblem.empiricalMoments))
    arrayDistance[indexMoment] = (sMMProblem.empiricalMoments[k][1] - simulatedMoments[k])
  end

  # See Lee and Ingram (1991) and Ruge-Murcia (2012)
  tau = tData/tSimData
  J = tData*(1.0 + tau)*transpose(arrayDistance)*sMMProblem.W*arrayDistance
  #Critical value above which the null hypothesis is rejected
  #Look at the 1.0 - alpha percentile of Chi²(df)
  c = quantile(d_chi, 1.0 - alpha)

  return J, c

end

ext: jl


text:
module Daemon

using ...SimpleTasks.Types

import SimpleTasks.Services.Queue
import SimpleTasks.Services.Bucket
import SimpleTasks.Services.Datasource
import SimpleTasks.Tasks.DaemonTask
import SimpleTasks.Tasks.BasicTask
import JSON

export DaemonService, register!, run

type DaemonService
    task_queue::QueueService
    error_queue::QueueService
    done_queue::QueueService
    bucket::BucketService
    datasource::DatasourceService
    poll_frequency_seconds::Int64
    tasks::Dict{AbstractString, Type}
end

DaemonService(task_queue::QueueService, error_queue::QueueService,
    done_queue::QueueService, bucket::BucketService,
    datasource::DatasourceService, poll_frequency_seconds::Int64) =
        DaemonService(task_queue, error_queue, done_queue, bucket,
            datasource, poll_frequency_seconds,
            Dict{AbstractString, Module}())

########### Helper types for queue notification ###########

type DaemonHostInfo
    hostname::AbstractString
    user::AbstractString
end
DAEMON_HOST_INFO = DaemonHostInfo(gethostname(), ENV["USER"])

type ErrorReport
    host_info::DaemonHostInfo
    task::Union{AbstractString, DaemonTaskDetails}
    exception::AbstractString
end

type DoneReport
    host_info::DaemonHostInfo
    task::DaemonTaskDetails
    result::DaemonTask.Result
end

function run(daemon::DaemonService)
    while true
        message = ""
        try
            message = Queue.pop_message(daemon.task_queue)

            if isempty(message)
                println("No messages found in ", Queue.string(daemon.task_queue))
            else
                println("Message received is $(message)")

                task = nothing
                try
                    task = parse(daemon, message)

                    println("Task is $(task.basic_info.id), ",
                        task.basic_info.name)

                    result = DaemonTask.run(task, daemon.datasource)

                    notify_task_done(daemon.done_queue, task, result)
                catch task_exception
                    exception_buffer = IOBuffer()
                    showerror(exception_buffer, task_exception,
                        catch_backtrace(); backtrace = true)
                    seekstart(exception_buffer)
                    notify_task_error(daemon.error_queue,
                        task == nothing ? message : task,
                        readall(exception_buffer))
                    rethrow(task_exception)
                end
            end
        catch exception
            print(STDERR, "Daemon Exception: ")
            showerror(STDERR, exception, catch_backtrace(); backtrace = true)
            println(STDERR)
        end

        sleep(daemon.poll_frequency_seconds)
    end
end

function notify_task_done(queue::QueueService, task::DaemonTaskDetails,
        result::DaemonTask.Result)
    done_report = DoneReport(DAEMON_HOST_INFO, task, result)
    Queue.push_message(queue; message_body = JSON.json(done_report))
end

"""
    notify_task_error(queue::QueueService,
        task::Union{DaemonTaskDetails, AbstractString},
        exception_message::AbstractString)
"""
function notify_task_error(queue::QueueService,
        task::Union{DaemonTaskDetails, AbstractString},
        exception_message::AbstractString)
    try
        error_report = ErrorReport(DAEMON_HOST_INFO, task,
            exception_message)
        Queue.push_message(queue; message_body = JSON.json(error_report))
    catch notify_exception
        print(STDERR, "ERROR in notifying  error-queue ", Queue.string(queue),
            " (Keep looking below for the real Daemon exception):")
        showerror(STDERR, notify_exception, catch_backtrace(); backtrace = true)
        println(STDERR)
    end
end

"""
    register!(daemon::DaemonService, task_module::Module)

Register the task type generation for the given task_name

"""
function register!(daemon::DaemonService, task_name::AbstractString,
        task_type::Type)
    if !DaemonTask.can_execute(task_type)
        error("Can not register $task_type with name $task_name " *
            " in Daemon. Could not find a registered method to execute")
    end

    if haskey(daemon.tasks, task_name)
        warn("Daemon has already registered $task_name with " *
            "$(daemon.tasks[task_name])")
    end

    daemon.tasks[task_name] = task_type
end

"""
    parse(daemon::DaemonService, text::ASCIIString)

Parse input JSON message into a task object
"""
function parse(daemon::DaemonService, text::ASCIIString)
    text = strip(text)

    if isempty(text)
        throw(ArgumentError("Trying to parse empty string for task"))
    end

    message = JSON.parse(text)

    basic_info = BasicTask.Info(message["basic_info"])

    if !haskey(daemon.tasks, basic_info.name)
        error("Task $(basic_info.name) is not registered with the daemon")
    end

    task_type = daemon.tasks[basic_info.name]
    return task_type(basic_info, message["payload_info"])
end

end # end module Daemon

max_forks_repo_name: seung-lab/SimpleTasks


text:
module StanDataFrames

using Statistics, DataFrames
import CmdStan: convert_a3d

# package code goes here

include("utilities/convert_a3d.jl")

export
  convert_a3d

end # module


text:
# in future versions, let's use the dataframe functions and then plot from the dataframe,
# JJS 11/2/20

using Plots

"""
Compare the fluctuations in tank levels calculated from WaterModels and WNTR
"""
function compare_tank_level(wm_data,wm_solution,wntr_data,wntr_simulation, outfilepath::String, tank_id)
    num_time_step = length(wm_solution["nw"])    # number of time steps

    tank_node_id = string(wm_data["nw"]["1"]["tank"][tank_id]["node"])
    tank_name = tank_node_id

    diameter = wntr_data.nodes._data[tank_name].diameter
    level_wntr = Array{Float64,1}(undef,num_time_step)
    level_watermodels = Array{Float64,1}(undef,num_time_step)

    for t in 1:num_time_step
        level_wntr[t] = wntr_simulation.node["pressure"][tank_name].values[t]
        level_watermodels[t] = wm_solution["nw"][string(t)]["node"][tank_node_id]["p"]
    end

    p = plot(1:num_time_step,ones(num_time_step,1).*wntr_data.nodes._data[tank_name].min_level,label="Min level",linecolor=:black,linestyle=:dash)
    p = plot!(p,1:num_time_step,ones(num_time_step,1).*wntr_data.nodes._data[tank_name].max_level,label="Max level",linecolor=:black,linestyle=:dot)
    p = plot!(p,1:num_time_step,level_wntr,label="WNTR",
      title=string(tank_name),xlabel="num_time_step (h)",ylabel="Water level (m)")
    p = plot!(p,1:num_time_step,level_watermodels,label="WaterModel")
    savefig(p,outfilepath*tank_name*".png")
end


text:
# Michelle Medina
# July 7, 2020
# Tests Hierarchical Basis Functions

# Going from nodal to hierarchical grid and sparse grid
xVal = collect(range(0.0, stop=1.0, length = 5))
nodal = calcNodal(x -> exp(x), xVal)
modal = Nodal_2_H(nodal)

# Going back and forth between nodal and modal
back = H_2_Nodal(modal)
@test isapprox(nodal.values,back.values)
#
# # Tests how coefficients fall off with level
# l = 12
# xYLOPHONE = collect(range(0.0, stop=1.0, length = 2^l +1))
# nodalTest = calcNodal(x -> exp(x), xYLOPHONE)
# Values = Nodal_2_H(nodalTest)
# MaxCoeff = getMaximum(Values)
#
# using PyPlot
# semilogy(MaxCoeff, "g o")
# xlabel("Levels")
# ylabel("Maximum Coefficient")
# title("Max Coefficient vs Level")
# show()

# Testing analytical and numeric calculations (up to a certain level)
# 9 points --> exp(x) [Nodal Basis] --> Hierarchical basis (level 3)
# 100 points between [0, 1]
# y_analytic = exp(x) # where x has 100 points
# y_numeric_upto_level_l = evaluate(HierarchichalBasis(up to level l)) on these 100 points
# error_upto_level_l = y_numeric_upto_level_l - y_analytic
# plot(upto_level_l, error_upto_level_l)

# Testing interpolation method
xValue = collect(range(0.0, stop=1.0, length = 5))
nodalTest = calcNodal(x -> exp(x), xValue)
newVal = Nodal_2_H(nodalTest)
newVal2 = Nodal_2_H_new(nodalTest)
@test isapprox(newVal.levels[2].coefficients, newVal2.levels[2].coefficients, atol=1e-14)

# Testing error of going up to a certain level
# using LinearAlgebra
# xNew = collect(range(0.0, stop=1.0, length = 2000))
# y2 = exp.(xNew)
# E = []
# for l in 0:10
# 	y1 = evaluate_upto(newVal, l, xNew)
# 	error_upto_level_l = norm(y1 - y2)
# 	append!(E, error_upto_level_l)
# end
#
#
# using PyPlot
# semilogy(E)
# xlabel("Levels")
# ylabel("Interpolation Error")
# title("Interpolation Error vs Levels")
# show()


text:
@testset "gmesh2fem.jl" begin

end


text:
# this file create julia functions from the python modules

empty!(type_map)  # for behaving nicely in system image

for (ee_expr, julia_type) in pre_type_map
    type_map[ee_expr()] = julia_type
end

# first loop over the ee types and set functions that take type
for key in collect(keys(type_map))
    # get list of methods
    methods = py"dir"(key)
    # loop over all methods and create julia methods
    for method in methods
        if ~startswith(method, "_")
            pyattr_set([type_map[key]], Symbol(method))
            if ~(method in ee_exports)
                push!(ee_exports, method)
            end
        end
    end
end

# next get methods from the individual modules
# outer loop is to get the methods for each module
# inner loop is to map the methods from module to julia funcs

modules = [
    :Array,
    :Classifier,
    :Clusterer,
    :ConfusionMatrix,
    :Date,
    :DateRange,
    :Dictionary,
    :ErrorMargin,
    :Feature,
    :FeatureCollection,
    :Filter,
    :Geometry,
    :Image,
    :ImageCollection,
    :Join,
    :Kernel,
    :List,
    :Model,
    :Number,
    :PixelType,
    :Projection,
    :Reducer,
    :String,
    :Terrain,
    :data,
]

mod_methods = OrderedDict()

excludefuncs = ["product","cat"]

# loop over the modules
for mod in modules
    # get the module functions and add to dict
    @eval submethods = collect(py"dir"(ee.$(string(mod))))

    # loop over the module methods
    for submethod in submethods
        submethod = convert(AbstractString, submethod)
        # check if submethod is not private
        if ~startswith(submethod, "_") & ~(submethod in excludefuncs)
            m = Symbol(submethod)
            # create a julia function of the public methods
            @eval begin
                function $m(args...; kwargs...)
                    method = ee.$(string(mod)).$(submethod)
                    result = pycall(method, PyObject, args...; kwargs...)
                    ee_wrap(result)
                end
            end
            # add to list of methods for module export
            if ~(submethod in ee_exports)
                push!(ee_exports, submethod)
            end
        end
    end
end

# for some reason the ee.Algorithms module is a Dict....why???
# function to recusively search though the ee.Algorithms module and wrap all python functions
function wrap_eealgorithms(dict)
    for (k, v) in dict
        if typeof(v) == PyCall.PyObject
            m = Symbol(k)
            @eval begin
                function $m(args...; kwargs...)
                    method = $v
                    result = pycall(method, PyObject, args...; kwargs...)
                    ee_wrap(result)
                end
            end
            # add to list of methods for module export
            if ~(k in ee_exports)
                push!(ee_exports, k)
            end
        else
            wrap_eealgorithms(v)
        end
    end
end

# get the dict of ee algorithms
algorithms_dict = ee.Algorithms
# remove the GeometryConstructors functions that clobber with others
delete!(algorithms_dict, "GeometryConstructors")

# apply recusive wrapping
wrap_eealgorithms(algorithms_dict)

# define dispatch signatures for Mathematical Operations
# from https://docs.julialang.org/en/v1/manual/mathematical-operations/#Bitwise-Operators
# arithmetic operators
Base.:+(x::T,y::T) where {T <: AbstractEEObject} = add(x,y)
Base.:+(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = add(x,y)
Base.:+(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = add(typeof(y)(x),y)
Base.:-(x::T,y::T) where {T <: AbstractEEObject} = subtract(x,y)
Base.:-(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = subtract(x,y)
Base.:-(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = subtract(typeof(y)(x),y)
Base.:*(x::T,y::T) where {T <: AbstractEEObject} = multiply(x,y)
Base.:*(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = multiply(x,y)
Base.:*(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = multiply(typeof(y)(x),y)
Base.:/(x::T,y::T) where {T <: AbstractEEObject} = divide(x,y)
Base.:/(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = divide(x,y)
Base.:/(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = divide(typeof(y)(x),y)
Base.:÷(x::T,y::T) where {T <: AbstractEEObject} = toInt(divide(x,y))
Base.:÷(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = toInt(divide(x,y))
Base.:÷(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = toInt(divide(typeof(y)(x),y))
Base.:\(x::T,y::T) where {T <: AbstractEEObject} = divide(y,x)
Base.:\(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = divide(y,x)
Base.:\(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = divide(y,typeof(y)(x))
Base.:^(x::T,y::T) where {T <: AbstractEEObject} = pow(x,y)
Base.:^(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = pow(x,y)
Base.:^(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = pow(typeof(y)(x),y)
Base.:%(x::T,y::T) where {T <: AbstractEEObject} = mod(x,y)
Base.:%(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = mod(x,y)
Base.:%(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = mod(typeof(y)(x),y)

# bitwise pperators
Base.:~(x::T) where {T <: AbstractEEObject} = bitwiseNot(x)
Base.:&(x::T,y::T) where {T <: AbstractEEObject} = bitwiseAnd(x,y)
Base.:&(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = bitwiseAnd(x,y)
Base.:&(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = bitwiseAnd(typeof(y)(x),y)
Base.:|(x::T,y::T) where {T <: AbstractEEObject} = bitwiseOr(x,y)
Base.:|(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = bitwiseOr(x,y)
Base.:|(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = bitwiseOr(typeof(y)(x),y)
Base.:⊻(x::T,y::T) where {T <: AbstractEEObject} = bitwiseXor(x,y)
Base.:⊻(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = bitwiseXor(x,y)
Base.:⊻(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = bitwiseXor(typeof(y)(x),y)
Base.:>>(x::T,y::T) where {T <: AbstractEEObject} = rightShift(x,y)
Base.:>>(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = rightShift(x,y)
Base.:>>(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = rightShift(typeof(y)(x),y)
Base.:<<(x::T,y::T) where {T <: AbstractEEObject} = leftShift(x,y)
Base.:<<(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = leftShift(x,y)
Base.:<<(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = leftShift(typeof(y)(x),y)

# numeric comparisons
isequal(x::T,y::T) where {T <: AbstractEEObject} = eq(x,y)
isequal(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = eq(x,y)
isequal(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = eq(typeof(y)(x),y)
isless(x::T,y::T) where {T <: AbstractEEObject} = lt(x,y)
isless(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = lt(x,y)
isless(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = lt(typeof(y)(x),y)
==(x::T,y::T) where {T <: AbstractEEObject} = eq(x,y)
==(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = eq(x,y)
==(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = eq(typeof(y)(x),y)
!=(x::T,y::T) where {T <: AbstractEEObject} = neq(x,y)
!=(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = neq(x,y)
!=(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = neq(typeof(y)(x),y)
≠(x::T,y::T) where {T <: AbstractEEObject} = neq(x,y)
≠(x::T,y::S) where {T <: AbstractEEObject, S <: Real} = neq(x,y)
≠(x::S,y::T) where {T <: AbstractEEObject, S <: Real} = neq(typeof(y)(x),y)


text:
### A Pluto.jl notebook ###
# v0.19.0

using Markdown
using InteractiveUtils

# ╔═╡ 5a35904c-cb93-11ec-2ee1-11009467e844
begin
	using PyCall
end

# ╔═╡ 45e2dd12-0e51-432b-9b1b-7e54282b5b89
begin
	px = pyimport("plotly.express")
	df = px.data.iris()
	fig = px.scatter(
	    df,
	    x="sepal_width",
	    y="sepal_length",
	    color="species",
	    size="petal_length",
	    hover_data=["petal_width"],
	)
	fig
end

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"

[compat]
PyCall = "~1.93.1"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

julia_version = "1.7.2"
manifest_format = "2.0"

[[deps.ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[deps.Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[deps.Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[deps.CompilerSupportLibraries_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"

[[deps.Conda]]
deps = ["Downloads", "JSON", "VersionParsing"]
git-tree-sha1 = "6e47d11ea2776bc5627421d59cdcc1296c058071"
uuid = "8f4d0f93-b110-5947-807f-2305c1781a2d"
version = "1.7.0"

[[deps.Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[deps.Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[deps.JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "3c837543ddb02250ef42f4738347454f95079d4e"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.3"

[[deps.LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[deps.LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[deps.LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[deps.Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[deps.LinearAlgebra]]
deps = ["Libdl", "libblastrampoline_jll"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[deps.MacroTools]]
deps = ["Markdown", "Random"]
git-tree-sha1 = "3d3e902b31198a27340d0bf00d6ac452866021cf"
uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
version = "0.5.9"

[[deps.Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[deps.MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

[[deps.Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[deps.MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[deps.NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[deps.OpenBLAS_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"

[[deps.Parsers]]
deps = ["Dates"]
git-tree-sha1 = "1285416549ccfcdf0c50d4997a94331e88d68413"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.3.1"

[[deps.Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[deps.PyCall]]
deps = ["Conda", "Dates", "Libdl", "LinearAlgebra", "MacroTools", "Serialization", "VersionParsing"]
git-tree-sha1 = "1fc929f47d7c151c839c5fc1375929766fb8edcc"
uuid = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
version = "1.93.1"

[[deps.Random]]
deps = ["SHA", "Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[deps.SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[deps.Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[deps.Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

[[deps.VersionParsing]]
git-tree-sha1 = "58d6e80b4ee071f5efd07fda82cb9fbe17200868"
uuid = "81def892-9a0e-5fdd-b105-ffc91e053289"
version = "1.3.0"

[[deps.Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[deps.libblastrampoline_jll]]
deps = ["Artifacts", "Libdl", "OpenBLAS_jll"]
uuid = "8e850b90-86db-534c-a0d3-1478176c7d93"

[[deps.nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"
"""

# ╔═╡ Cell order:
# ╠═5a35904c-cb93-11ec-2ee1-11009467e844
# ╠═45e2dd12-0e51-432b-9b1b-7e54282b5b89
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


text:
# Generate the first n numbers in Halton's sequence with base b.
#
#   seq = haltonsequence(b, n)
#
# Output:
#  seq - sequence of n Halton numbers
#
# Inputs:
#  b - base of the sequence
#  n - length of vector to be generated
#
function haltonsequence(b, n)
  map((i) -> haltonnumber(b, i), 1:n)
end

# Generate the nth Halton number in the sequence with base b.
#
# Notes:
#  Implementation is based on the psudo code in:
#    http://en.wikipedia.org/wiki/Halton_sequence
#
function haltonnumber(base, index)
  res = 0
  f = 1 / base
  i = index

  while (i > 0)
    res += f * (i % base)
    i = ifloor(i / base)
    f = f / base
  end

  return res
end


text:
module NCTSSOS

using DynamicPolynomials
using MultivariatePolynomials
using JuMP
using MosekTools
using LightGraphs
using ChordalGraph
using MetaGraphs
using LinearAlgebra
using SparseArrays

export nctssos_first, nctssos_higher!, cs_nctssos_first, cs_nctssos_higher!, ptraceopt_first, ptraceopt_higher!
export Werner_witness_first, Werner_witness_higher!

include("clique_merge.jl")
include("ncupop.jl")
include("nccpop.jl")
include("mixncpop.jl")
include("trace.jl")

end


text:
"""
    ADMSynthControlModel

Synthetic control model as described in [Abadie, Diamond, Hainmueller
(2010)](https://economics.mit.edu/files/11859)
The implementation follows the exposition in [Abadie (2021)](https://economics.mit.edu/files/17847),
with the weight matrix V governing weights of covariates X in the fitting procedure chosen by
cross-validation based on a train/test split of the pre-treatment outcomes

Takes in a `TreatmentPanel` and holds the results of calling the `fit!` method, optimal
weights for comparator units and the resulting predicted outcome of the synthetic control
unit as well as the estimated treatment effect from comparing synthetic control and observed
outcome of the treated unit.

"""

mutable struct ADMSynthControlModel{T1} <: SCM
    treatment_panel::T1
    w::Vector{Float64} # Vector of weights for each non-treated observation, length N-1
    v::Matrix{Float64} # Vector of importance weights for predictors X₁, ..., Xₖ, length k
    ŷ₁::Vector{Float64} # Vector of post-treatment predicted outcomes for synthetic unit
    τ̂::Vector{Float64} # Vector of estimated treatmet impacts
    p_test_res::Array{Float64, 2} # Matrix of outcomes for untreated under placebo
end

function ADMSynthControlModel(tp::BalancedPanel{TreatmentPanels.SingleUnitTreatment, S}) where S
    @unpack Y, N, T, W, X = tp

    ŷ₁ = zeros(T)
    w = zeros(N - 1)
    τ̂ = zeros(length_T₁(tp))
    p_test_res = zeros(N - 1, T)

    SynthControlModel{typeof(tp)}(tp, w, ŷ₁, τ̂, p_test_res)
end

function fit!(x::ADMSynthControlModel;
    train_share = 0.7)

    k = length(X)

    # Length of train/test period in pre-intervention data
    T₀_train = round(Int, train_share * length_T₀(x))
    T₀_test = round(Int, (1 - train_share) * length_T₀(x))

    # Weight matrix v is size (k × k), simplest choice for startes is inverse variance of X
    v = zeros(k, k)
    v[diagind(v)] .= [1/var(xᵢ) for x ∈ eachcol(x)]

    # Calculate weights on training data
    #!# Closure that captures v?
    obj(w) = (x₁ .- x₀*w)' * v * (x₁ .- x₀*w) + 1e6*(1.0 - sum(w))^2

    # Target is minimizing the distance between predicted and actual outcome for treated unit
    # in test period
    function target(v_diag)

        v[diagind(v)] .= v_diag

        # Get weights given proposal for v
        w′ = optimize(obj, initial, upper, lower).minimizer

        # Calculate loss in test period
        (y₁_test .- y₀_test * w′)*(y₁_test .- y₀_test * w′)
    end

    v = optimize(target, initial, lower, upper)
end

####################################################################################################
##################################     MANUAL ATTEMPT AT ADH     ###################################
####################################################################################################

using SynthControl
using Dates, DataFrames, Statistics, LinearAlgebra

df = SynthControl.load_germany()
sort!(df, [:country, :year])

transform!(groupby(df, :country),
    [:year, :gdp] => ((y, g) -> mean(skipmissing(g[1981 .<= y .<= 1990]))) => :gdp_81_90,
    [:year, :trade] => ((y, i) -> mean(skipmissing(i[1981 .<= y .<= 1990]))) => :trade_81_90,
    [:year, :infrate] => ((y, i) -> mean(skipmissing(i[1981 .<= y .<= 1990]))) => :infrate_81_90,
    [:year, :industry] => ((y, i) -> mean(skipmissing(i[1981 .<= y .<= 1990]))) => :industry_81_90,
    [:year, :invest80] => ((y, i) -> mean(skipmissing(i[1980 .<= y .<= 1985]))) => :invest80_80_85,
    [:year, :schooling] => ((y, i) -> mean(skipmissing(i[1980 .<= y .<= 1985]))) => :schooling_80_85)

covar_list = [:gdp_81_90,  :trade_81_90, :infrate_81_90, :industry_81_90, :schooling_80_85, :invest80_80_85]

X₁_df = combine(first, groupby(df[(df.country .== "West Germany"), [covar_list; :country]], :country))
X₁ = [x for x ∈ X₁_df[1, 2:end]]

X₀_df = combine(first, groupby(df[(df.country .!= "West Germany"), [covar_list; :country]], :country, sort = true))
X₀ = Matrix(X₀_df[:, 2:end])'

w_equal = [1/size(X₀, 2) for _ ∈ 1:size(X₀, 2)]

# Weights shown in ADH ()
w_adh = [0, 0.42, 0, 0, 0, 0, 0, 0.16, 0.09, 0, 0, 0, 0, 0.11, 0, 0.22]

sqrt((X₁ .- X₀*w_adh)'*(X₁ .- X₀*w_adh))
norm(X₁ .- X₀*w_adh)
norm(X₁ .- X₀*w_equal)

V = zeros(6, 6)
V[diagind(V)] .= 1.0

sqrt((X₁ .- X₀*w_equal)'*V*(X₁ .- X₀*w_equal))
sqrt((X₁ .- X₀*w_adh)'*V*(X₁ .- X₀*w_adh))

obj1(w) = (X₁ .- X₀*w)'*(X₁ .- X₀*w)
obj2(w) = (X₁ .- X₀*w)'*(X₁ .- X₀*w) + 1e6 * (1.0 - sum(w))^2
obj3(w) = (X₁ .- X₀*w)'*(X₁ .- X₀*w) + 1e6 * (1.0 - sum(w))^2 + 1_000 * norm(w)

using Optim

initial = w_equal
upper = fill(1.0, length(w_equal))
lower = fill(0.0, length(w_equal))

w_simple = optimize(obj1, lower, upper, initial).minimizer
w_sum_constraint = optimize(obj2, lower, upper, initial).minimizer
w_sum_ridge_constraint = optimize(obj3, lower, upper, initial).minimizer

summary = DataFrame(Country = sort(unique(df[df.country .!= "West Germany", :].country)),
            ADH_weights = w_adh,
            w_simple = round.(w_simple, digits = 2),
            w_sum_constraint = round.(w_sum_constraint, digits = 2),
            w_sum_ridge_constraint = round.(w_sum_ridge_constraint, digits = 2))

V_adh = zeros(6, 6)
covar_list = [:gdp_81_90,  :trade_81_90, :infrate_81_90, :industry_81_90, :schooling_80_85, :invest80_80_85]
V_adh[diagind(V_adh)] .= [0.442, 0.134, 0.072, 0.001, 0.107, 0.245]

sqrt((X₁.- X₀*w_adh)'*V_adh*(X₁ .- X₀*w_adh))

max_stars_repo_name: nilshg/SynthControl

max_issues_repo_name: nilshg/SynthControl

max_forks_repo_name: nilshg/SynthControl


text:
### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Statistics, DataFrames, CSV
	using Random, Distributions, Plots
	using HypothesisTests
	Random.seed!(0)
end;

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing6.03"

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
begin
	data1 = CSV.read(datadir("machine1.csv"), DataFrame; delim=",")[:, 1]
	data2 = CSV.read(datadir("machine2.csv"), DataFrame; delim=",")[:, 1]
end;

# ╔═╡ 1aa41898-3afa-11eb-2f49-570853ce6435
begin
	xBar1, xBar2 = mean(data1), mean(data2)
	n1, n2 = length(data1), length(data2)
	alpha = 0.05
	t = quantile(TDist(n1+n2-2),1-alpha/2)

	s1, s2 = std(data1), std(data2)
	sP = sqrt(((n1-1)*s1^2 + (n2-1)*s2^2) / (n1+n2-2))

	Text("Calculating formula: $((xBar1 - xBar2 - t*sP* sqrt(1/n1 + 1/n2),
		xBar1 - xBar2 + t*sP* sqrt(1/n1 + 1/n2)))")
end

# ╔═╡ 1b68b272-3b0c-11eb-013c-651a9f49938b
Text("Using confint(): $(confint(EqualVarianceTTest(data1,data2),alpha))")

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing6.03"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═1aa41898-3afa-11eb-2f49-570853ce6435
# ╠═1b68b272-3b0c-11eb-013c-651a9f49938b
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


text:
function acos(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "acos", Dfp, (Dfp,), arg0)
end

function asin(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "asin", Dfp, (Dfp,), arg0)
end

function atan(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "atan", Dfp, (Dfp,), arg0)
end

function cos(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "cos", Dfp, (Dfp,), arg0)
end

function equals(obj::Object, arg0::Object)
    return jcall(obj, "equals", jboolean, (Object,), arg0)
end

function exp(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "exp", Dfp, (Dfp,), arg0)
end

function hash_code(obj::Object)
    return jcall(obj, "hashCode", jint, ())
end

function log(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "log", Dfp, (Dfp,), arg0)
end

function pow(::Type{DfpMath}, arg0::Dfp, arg1::Dfp)
    return jcall(DfpMath, "pow", Dfp, (Dfp, Dfp), arg0, arg1)
end

function pow(::Type{DfpMath}, arg0::Dfp, arg1::jint)
    return jcall(DfpMath, "pow", Dfp, (Dfp, jint), arg0, arg1)
end

function sin(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "sin", Dfp, (Dfp,), arg0)
end

function tan(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "tan", Dfp, (Dfp,), arg0)
end


text:
using Tables
using MLJ
import EvoTrees: EvoTreeRegressor, predict
using EvoTrees: logit, sigmoid

features = rand(10_000) .* 5 .- 2
X = reshape(features, (size(features)[1], 1))
Y = sin.(features) .* 0.5 .+ 0.5
Y = logit(Y) + randn(size(Y))
Y = sigmoid(Y)
y = Y
X = Tables.table(X)

@load EvoTreeRegressor
tree_model = EvoTreeRegressor(max_depth=5, η=0.01, nrounds=200)
tree = machine(tree_model, X, y)
train, test = partition(eachindex(y), 0.7, shuffle=true); # 70:30 split
fit!(tree, rows=train)

yhat = MLJ.predict(tree, MLJ.selectrows(X,test))
mean(abs.(yhat - MLJ.selectrows(Y,test)))

fit!(tree, rows=train)
yhat = MLJ.predict(tree, MLJ.selectrows(X,test))
mean(abs.(yhat - MLJ.selectrows(Y,test)))


text:
"""
    module SaturableCPA
"""
module SaturableCPA

export SCPA
export maxwell_scpa

using ..Common
using ..Lasing
using NLsolve

function SCPA(sim::Simulation, ωs_init::Array{Float64}, ψs_init::Array; kwargs...)
    csim = conj(sim)
    ωs,ψs = SALT(sim,ωs_init,conj(ψs_init);kwargs...)
    return ωs,conj(ψs)
end

struct Maxwell_SCPA{TMS}
    MS::TMS
end
maxwell_scpa(sim::Simulation,m::Int) = Maxwell_SCPA(maxwell_salt(conj(sim),m))

fnames = (:nlsolve,:fixedpoint)
for fn ∈ fnames
    @eval begin function NLsolve.$(fn)(ms::Maxwell_SCPA, ωs_init::Array, ψs_init::Array; kwargs...)
            return nlsolve(ms.MS,ωs_init,conj(ψs_init);kwargs...)
        end
    end
end

function Base.getproperty(ms::Maxwell_SCPA,sym::Symbol)
    if sym == :ωs
        getfield(ms.MS,:ωs)
    elseif sym == :ψs
        getfield(ms.MS,:ψs)
    else
        return getfield(ms,sym)
    end
end

function Base.show(io::IO,mscpa::Maxwell_SCPA)
    ms = mscpa.MS
    print(io,"Maxwell_SCPA with $(ms.m) mode")
    ms.m>1 ? print(io,"s") : nothing
    if !ms.solved[]
        print(io," (")
        printstyled(io,"unsolved",color=:light_yellow)
        print(io,", for use in ")
        printstyled(io,"nlsolve",color=:cyan)
        print(io,")")
    elseif ms.converged[]
        print(io," (")
        printstyled(io,"solved",color=:light_green)
        print(io,", solution in fields ")
        printstyled(io,"ωs",color=:cyan)
        print(io,", ")
        printstyled(io,"ψs",color=:cyan)
        print(io,")")
    else
        print(io," (")
        printstyled(io,"solve attempted",color=:light_red)
        print(io,", ")
        printstyled(io,"no convergence",color=:light_red)
        print(io,")")
    end
end

end # module


text:
##############################################################################
##
## References
##
##############################################################################

# - https://en.bitcoin.it/wiki/Technical_background_of_version_1_Bitcoin_addresses
# - https://en.bitcoin.it/wiki/List_of_address_prefixes

##############################################################################
##
## Definitions
##
##############################################################################

type Keys
  priv_key::Array{Uint8}
  pub_key::Array{Uint8}
end

function generate_keys(network_id = 0x00)
  priv_key = Crypto.random(256)
  pub_key = get_pub_key(priv_key, network_id = network_id, version = version)

  return (Crypto.oct2hex(priv_key), pub_key)
end

function get_pub_key(priv_key::String; network_id = 0x00)
  get_pub_key(Crypto.hex2oct(priv_key),
              network_id = network_id)
end

function get_pub_key(priv_key::Array{Uint8}; network_id = 0x00)

  # Generate corresponding public key generated with against ECDSA secp256k1
  # (65 bytes, 1 byte 0x04, 32 bytes corresponding to X coordinate, 32 bytes
  # corresponding to Y coordinate)
  pub_key = Crypto.ec_pub_key(priv_key)

  # Perform SHA-256 hashing on the public key
  pub_key = Crypto.digest("SHA256", pub_key)

  # Perform RIPEMD-160 hashing on the result of SHA-256
  pub_key = Crypto.digest("RIPEMD160", pub_key)

  # Add version byte in front of RIPEMD-160 hash (0x00 for Main Network)
  # Reference: https://bitcoin.org/en/developer-reference#address-conversion
  pub_key = [Crypto.int2oct(network_id), pub_key]

  # Get checksum by performing SHA256 hash twice and taking first 4 bytes
  checksum = get_checksum(pub_key)

  # Add the 4 checksum bytes from stage 7 at the end of extended RIPEMD-160
  # hash from stage 4. This is the 25-byte binary Bitcoin Address.
  append!(pub_key, checksum)

  # Convert the result from a byte string into a base58 string using
  # Base58Check encoding. This is the most commonly used Bitcoin Address format
  # Reference: https://en.bitcoin.it/wiki/Base58Check_encoding
  # TODO: array to string to BigInt is really round-about
  pub_key = encode58(pub_key)

  return pub_key
end


text:
function parse_nwchem_output_file(output_file::String, observables::String...)

end

function get_matching_key(line::String, search_terms::Dict{String, String})
    for key in keys(search_terms)
        if occursin(search_terms[key], line)
            return key
        end
    end
    return ""
end

function get_energies(output_file_lines::Vector{String})
    search_terms = (Dict{String, String}(
        "scf"     => "Total SCF energy",
        "mp2"     => "Total MP2 energy",
        "ccsd(t)" => "Total CCSD(T) energy"
    ))
    energies = Dict{String, Vector{Float64}}()
    for (i, line) in enumerate(output_file_lines)
        possible_key = get_matching_key(line, search_terms)
        if possible_key != ""
            energy = parse(Float64, split(line)[end])

            if haskey(energies, possible_key)
                push!(energies[possible_key], energy)
            else
                energies[possible_key] = [energy]
            end
        end
    end
    return energies
end

function get_gradients(output_file_lines::Vector{String})
    search_terms = (Dict{String, String}(
        "scf"     => "RHF ENERGY GRADIENTS",
        "mp2"     => "mp2 ENERGY GRADIENTS",
        "ccsd(t)" => "CCSD ENERGY GRADIENTS"
    ))
    gradients = Dict{String, Vector{Matrix{Float64}}}()
    for (i, line) in enumerate(output_file_lines)
        possible_key = get_matching_key(line, search_terms)
        if possible_key != ""
            num_atoms::Int = 0
            j::Int = 0
            while output_file_lines[i+j+4] != ""
                num_atoms += 1
                j += 1
            end
            grads = zeros(3, num_atoms)
            j = 0
            while output_file_lines[i+j+4] != ""
                grads[:, j+1] = parse.(Float64, split(nwchem_out[i+j+4])[end-2:end])
                j += 1
            end

            if haskey(gradients, possible_key)
                push!(gradients[possible_key], grads)
            else
                gradients[possible_key] = [grads]
            end
        end
    end
    return gradients
end

max_stars_repo_name: heindelj/NWChemGen

max_issues_repo_name: heindelj/NWChemGen

max_forks_repo_name: heindelj/NWChemGen


text:
kolmogorov_smirnov(F1::Function,F2::Function,grid) =
    maximum(abs.(F1.(grid)-F2.(grid)))

function kolmogorov_smirnov(d::SFMDistribution,F2::Function)
    grid = range(d.dq.mesh.nodes[1],d.dq.mesh.nodes[end],length=2001)
    F1 = cdf(d)
    M = 0.0
    for i in 1:n_phases(d.dq)
        m = maximum(abs.(F1.(grid,i)-F2.(grid,i)))
        (m>M)&&(M=m)
    end
    return M
end
kolmogorov_smirnov(d::SFMDistribution,s::Simulation) = kolmogorov_smirnov(d,cdf(s))
kolmogorov_smirnov(d::SFMDistribution,s::SFMDistribution) = kolmogorov_smirnov(d,cdf(s))

function Lp_cell_probs(d::SFMDistribution,F2::Function, p=1)
    grid = (d.dq.mesh.nodes[1:end-1]+d.dq.mesh.nodes[2:end])./2.0
    F1 = cell_probs(d)
    M = 0.0
    for i in 1:n_phases(d.dq)
        m = sum(abs.(F1.(grid,i)-F2.(grid,i)).^p)
        M += m
    end
    return M
end
Lp_cell_probs(d::SFMDistribution,s::Simulation,p=1) = Lp_cell_probs(d,cell_probs(s),p)
Lp_cell_probs(d::SFMDistribution,s::SFMDistribution,p=1) = Lp_cell_probs(d,cell_probs(s),p)

function Lp_pdf(d::SFMDistribution,f2::Function, p=1)
    grid = range(d.dq.mesh.nodes[1],d.dq.mesh.nodes[end],length=2001)
    f1 = pdf(d)
    M = 0.0
    for i in 1:n_phases(d.dq)
        fn_evals = abs.(f1.(grid,i)-f2.(grid,i)).^p
        m = sum((fn_evals[1:end-1]+fn_evals[2:end])./2.0 .* diff(grid))
        M += m
    end
    return M
end
Lp_pdf(d::SFMDistribution,s::SFMDistribution,p=1) = Lp_pdf(d,pdf(s),p)

function Lp(f1::Function, f2::Function, grid, p=1)
    fn_evals = abs.(f1.(grid)-f2.(grid)).^p
    m = sum((fn_evals[1:end-1]+fn_evals[2:end])./2.0 .* diff(grid))
    return m
end

max_stars_repo_name: angus-lewis/StochasticFluidQueues

max_issues_repo_name: angus-lewis/StochasticFluidQueues

max_forks_repo_name: angus-lewis/StochasticFluidQueues


text:
#=
  Utils for Julia

  Here are utils for my Julia programming.

  This Julia program was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

using Printf
# include("Euler.jl")

#
# Return a Dict of the elements and their occurences
# in the collection a.
#
function make_hash(a)
    d = Dict()
    for e in collect(a)
        # get!(d,e,0)
        if !haskey(d,e)
            d[e] = 0
        end
        d[e] += 1
    end
    return d
end

# This is to simplify the conversion from WebPPL
function flip(p=0.5)
    return Bernoulli(p)
end

# Return a normalized vector, i.e.
# where the sum is 1
function simplex(v)
    return v./sum(v)
end

#
# Show a sorted Dict with keys, values and percentages
# Example
# julia> show_var_dist(chains,:d1)
#
# Distributions of variable d1
#  1 =>   13654  0.341350
#  2 =>   13417  0.335425
#  3 =>   12929  0.323225
#
# Note: I'm not sure how to get more than one variables
# E.g. this don't work now:
#  julia> show_var_dist(chains,[:d1,:d2])
function show_var_dist(chains, var)

    if var in chains.name_map.parameters
        println("Distributions of variable $var")
        len = length(vcat(chains[var]...)) # handle multiple chains
        for kv in sort(make_hash(chains[var]))
            @printf "%-3.5f => % 7d  %2.6f\n" kv[1] kv[2] kv[2]/len
        end
    else
        # println("Variable $var is not in chains")
    end
end

# Show distribution of a variable in a MCMCChain
# Sort the dictionary in order of decreasing occurrence (percentage)
# Examples:
#  - show_var_dist_pct(chains, :n)       show all entries
#  - show_var_dist_pct(chains, :n, 10)  show first 10 entries (e.g. for large tables)
function show_var_dist_pct(chains::Chains, var, num=0)

    if var in chains.name_map.parameters
        println("Distributions of variable $var (num:$num)")
        len = length(vcat(chains[var]...)) # handle multiple chains
        c = 0
        for kv in sort(collect(make_hash(chains[var])),by=x->x[2],rev=true)
            c += 1
            if (num == 0) || (num > 0 && c <= num)
                @printf "%-3.5f => % 7d  (%2.6f)\n" kv[1] kv[2] kv[2]/len
            end
        end
    else
        # println("Variable $var is not in chains")
    end
end

# Show distribution of the elements in a (simple) Array.
# Sort the dictionary in order of decreasing occurrence (percentage)
# Examples:
#  - show_var_dist_pct(a)       show all entries
#  - show_var_dist_pct(a, 10)  show first 10 entries (e.g. for large tables)
# Full example:
# julia> show_var_dist_pct(rand(Binomial(100,0.1),100000),10)
#
# Distributions of variable (num:10)
# 10.00000 =>   13143  0.131430
# 9.00000 =>   13094  0.130940
# 11.00000 =>   12002  0.120020
# 8.00000 =>   11403  0.114030
# 12.00000 =>    9989  0.099890
# 7.00000 =>    8765  0.087650
# 13.00000 =>    7475  0.074750
# 6.00000 =>    5968  0.059680
# 14.00000 =>    5117  0.051170
# 5.00000 =>    3450  0.034500
#
function show_var_dist_pct(a::Array, num=0)
    println("Distributions of variable (num:$num)")
    c = 0
    len = length(a)
    isArray = false
    if length(a[1]) > 1 # typeof(a[1]) <: Array
        isArray = true
    end
    for kv in sort(collect(make_hash(a)),by=x->x[2],rev=true)
        c += 1
        if (num == 0) || (num > 0 && c <= num)
            if isArray
                # This is an array, vector, tuple of elements
                # len = length(kv[1])
                # format = "%-$(len)s => % 7d  %2.6f\n"
                # @printf format kv[1] kv[2] kv[2]/len
                # @printf "%-35s => % 7d  %2.6f\n" kv[1] kv[2] kv[2]/len
                println("$(kv[1])\t=>\t$(kv[2]) ($(kv[2]/len))")
            else
                @printf "%-3.5f => % 7d  (%2.6f)\n" kv[1] kv[2] kv[2]/len
            end
        end
    end
end

# This works. but one have to call with _varinfo
#  Example
#  # ...
#  observe(_varinfo, d1+d2 == ss)
# Note: We return !cond so it can be called with
#     observe(_varinfo, d1+d2 == ss) && return
# function observe(_varinfo,cond)
#    cond || Turing.@addlogprob! -Inf
#    return !cond
# end

#  This works. but one have to call with _sampler and  _varinfo
#  Example
#  # ...
#  observe(_sampler, _varinfo, d1+d2 == ss)
#
# See Dirac() for a better solution
#
#=
function obs!(_sampler,_varinfo,cond)
    if !cond
        if _sampler isa Turing.Sampler{<:Union{PG,SMC}}
            produce(-Inf)
        else
            Turing.@addlogprob! -Inf
        end
    else
        if _sampler isa Turing.Sampler{<:Union{PG,SMC}}
            produce(0.0)
        end
    end
    return !cond
end
=#

#
# Ensure ("observe") than an condition is met.
# Usage:
#   @model function two_dice(ss)
#    d1 ~ DiscreteUniform(1, 6)
#    d2 ~ DiscreteUniform(1, 6)
#
#    ss ~ Dirac(d1 + d2)
#    true ~ Dirac(d > 2) # note true ~ ...
# and
#
#
# From David Widdman in https://github.com/TuringLang/Turing.jl/issues/1471
#
# Note: This is now integrated in Distributions.jl
# struct Dirac{T} <: ContinuousUnivariateDistribution
#     x::T
# end
# Distributions.logpdf(d::Dirac, x::Real) = x == d.x ? 0.0 : -Inf

max_stars_repo_name: Wikunia/hakank

max_issues_repo_name: Wikunia/hakank

max_forks_repo_name: Wikunia/hakank


text:
# ---------------------------------------------------------------------------- #
#
#   master3D.jl
#
#   Type for 3D master element
#   Inherits from the abstract "master" type
#
#   λυτέος
#   Fall 2017
#
#   Max Opgenoord
#
# ---------------------------------------------------------------------------- #

"""
    Master3D

Master3D type:
Holds basis functions and quadrature points for tetrahedron master element.
"""
struct Master3D <: Master

  dim::Int64

  porder::Porder            # Polynomial order type
  p::Int64                  # Polynomial order of mesh
  pgauss::Int64             # Polynomial order to be integrated exactly
  nodfac::Int64            # Number of nodes on face

  gpts::Array{Float64,2}    # Gauss points  3D
  gwts::Array{Float64,1}    # Gauss weights 3D

  gpts2D::Array{Float64,2}  # Gauss points  2D
  gwts2D::Array{Float64,1}  # Gauss weights 2D

  gpts1D::Array{Float64,1}  # Gauss points  1D
  gwts1D::Array{Float64,1}  # Gauss weights 1D

  ϕ::Array{Float64,2}       # Shape functions in 3D
  ∇ϕ::Array{Float64,3}      # Derivatives of shape functions in 3D

  ϕ2D::Array{Float64,2}     # Shape functions in 2D
  ∇ϕ2D::Array{Float64,3}    # Derivatives of shape functions in 2D

  ϕ1D::Array{Float64,2}     # Shape functions in 1D
  ∇ϕ1D::Array{Float64,2}    # Derivatives of shape functions in 1D

  perm::Array{Int64,3}      # Node numbers on faces

end

"""
    Master3D( porder::Porder; pgauss::PGauss = PGdef( porder ) )

Constructor for Tetrahedron master element for order `porder`.
"""
function Master3D( porder::Porder; pgauss::PGauss = PGdef( porder ) )

  p  = porder.p
  pg = pgauss.p

  if p > 3
    error(" P > 3 not supported for 3D ")
  end

  (go1D, go2D, go3D) = comporder( pgauss )

  (gpts_,   gwts_)   = quadratureTet( go3D )
  (gpts2D_, gwts2D_) = quadratureTriangle( go2D )
  (gpts1D_, gwts1D_) = quadratureLine( go1D )

  (ϕ_,   ∇ϕ_)   = basisFuncTetLag( porder, gpts_[:,1], gpts_[:,2],  gpts_[:,3] )
  (ϕ2D_, ∇ϕ2D_) = basisFuncTriangleLag( porder, gpts2D_[:,1], gpts2D_[:,2] )
  (ϕ1D_, ∇ϕ1D_) = basisFuncLineLag( porder, gpts1D_ )

  perm_ = findPerm3D( p )

  nodfac_ = size(ϕ2D_,1)

  Master3D( 3, porder, p, pg, nodfac_, gpts_, gwts_, gpts2D_, gwts2D_, gpts1D_, gwts1D_,
    ϕ_, ∇ϕ_, ϕ2D_, ∇ϕ2D_, ϕ1D_, ∇ϕ1D_, perm_ )

end

# """
#     compOrder3D( orderRq::Int64 )
#
# Computes the polynomial order to be integrated exactly.
# """
# function compOrder3D( orderRq::Int64 )
#
#   # 1D
#
#   if orderRq <= 1
#     order1D = 1
#   elseif orderRq <= 3
#     order1D = 3
#   elseif orderRq <= 5
#     order1D = 5
#   elseif orderRq <= 7
#     order1D = 7
#   elseif orderRq <= 9
#     order1D = 9
#   elseif orderRq <= 11
#     order1D = 11
#   elseif orderRq <= 13
#     order1D = 13
#   elseif orderRq <= 15
#     order1D = 15
#   elseif orderRq <= 17
#     order1D = 17
#   elseif orderRq <= 19
#     order1D = 19
#   elseif orderRq <= 21
#     order1D = 21
#   elseif orderRq <= 23
#     order1D = 23
#   end
#
#   # 2D
#
#   if orderRq <= 2
#     order2D = 1
#   elseif orderRq <= 3
#     order2D = 2
#   elseif orderRq <= 4
#     order2D = 3
#   elseif orderRq <= 5
#     order2D = 4
#   elseif orderRq <= 8
#     order2D = 5
#   # elseif orderRq <= 10
#   #   order2D = 6
#   # elseif orderRq <= 11
#   #   order2D = 7
#   elseif orderRq <= 12
#     order2D = 8
#   elseif orderRq <= 13
#     order2D = 9
#   elseif orderRq <= 14
#     order2D = 10
#   end
#
#   # 3D
#
#   if orderRq <= 2
#     order3D = 1
#   elseif orderRq <= 3
#     order3D = 2
#   elseif orderRq <= 5
#     order3D = 3
#   elseif orderRq <= 6
#     order3D = 4
#   elseif orderRq <= 7
#     order3D = 5
#   elseif orderRq <= 8
#     order3D = 6
#   elseif orderRq <= 9
#     order3D = 7
#   elseif orderRq <= 10
#     order3D = 8
#   elseif orderRq <= 11
#     order3D = 9
#   elseif orderRq <= 12
#     order3D = 10
#   elseif orderRq <= 13
#     order3D = 11
#   elseif orderRq <= 14
#     order3D = 12
#   end
#
#   return order1D, order2D, order3D
#
# end

"""
    findPerm3D( p::Int64 )

Finds which nodes belong to which faces for tetrahedron.
"""
function findPerm3D( p::Int64 )

  # 1:  1 2 3
  # 2:  1 3 2
  # 3:  2 1 3
  # 4:  2 3 1
  # 5:  3 1 2
  # 6:  3 2 1

  orderind = [1 2 3;
              1 3 2;
              2 1 3;
              2 3 1;
              3 1 2;
              3 2 1]
  orderflip = [1, -1, -1, 1, 1, -1]

  n2d = Int( round( 0.5*(p+1)*(p+2) ) )

  indfaces = fill( 0::Int64, 4, n2d )

  perm = fill( 0::Int64, n2d, 4, 6 )

  if p == 1

    # 2 3 4 - 1 4 3  - 1 2 4 - 1 3 2 : First ind is the vertices of face which does not included current vertices

    indfaces = [2 3 4;
                1 4 3;
                1 2 4;
                1 3 2]

    for jj in 1:4, ii in 1:6
      perm[:,jj,ii] = indfaces[ jj, orderind[ii,:] ]
    end

  elseif p == 2
    # [2,3,4], [1,4,3], [1,2,4], [1,3,2]
    indfaces = [2  3  4  5  6  7
                1  4  3  5  8  9 #3  1  4  9  5  8
                1  2  4  6  9  10
                1  3  2  7  10 8]#2  1  3  8  7  10]

    for jj in 1:4
      ifac = indfaces[jj,:]
      for ii in 1:6
        # corners
        ifac[1:3] = indfaces[ jj,     orderind[ii,:] ]
        # edges
        ifac[4:6] = indfaces[ jj, 3 + orderind[ii,:] ]

        perm[:,jj,ii] = ifac
      end
    end

  elseif p == 3

    indfaces = [2  3  4   5  6  7  8  9 10 17
                1  4  3   6  5 11 12 13 14 18#3  1  4  13 14  6  5 11 12 18
                1  2  4   8  7 14 13 15 16 19
                1  3  2  10  9 16 15 12 11 20]#2  1  3  12 11 10  9 16 15 20]
    for jj in 1:4
      ifac = indfaces[jj,:]
      perm[:,jj,1] = ifac
      for ii in [4,5]
        ## not flipped direction
        # corners
        ifac[1:3] = indfaces[ jj, orderind[ii,:] ]
        # edges
        oind = [ orderind[ii,:]'; orderind[ii,:]' ]
        oind = 2*oind
        oind[1,:] -= 1
        oindv = oind[:]
        ifac[4:9] = indfaces[ jj, 3 + oindv ]
        # middle point unchanged

        perm[:,jj,ii] = ifac
      end
      for ii in [2,3,6]
        # flipped direction
        ifac[1:3] = indfaces[ jj, orderind[ii,:] ]
        # edges
        oind = [ orderind[ii,:]'; orderind[ii,:]' ]
        oind = 2*oind
        oind[2,:] -= 1
        oindv = oind[:]
        ifac[4:9] = indfaces[ jj, 3 + oindv ]
        # middle point unchanged

        perm[:,jj,ii] = ifac
      end

    end

  end

  return perm

end


text:
module TestTreasuryRates

using Base.Test
using EconDatasets
using Dates

@test convertFedDate(Date(2014, 10,23)) == "10/23/2014"

end


text:
using Double64s, ErrorfreeArithmetic, Test

include("constructors.jl")


text:
#This file has the definitions for the gray Flux-Limited Diffusion Solver
#The code uses units of
# cm for length
# sh for time (1 sh = 10^-8 s)
# GJ (gigajoules) for energy; 1 GJ is also known as a jerk
# keV for temperature (1 keV is 1.1605 * 10^7 K)

using PyCall
using JLD, HDF5, SparseArrays

#radiation constant in GJ/keV^4/cm^3
const a = 0.01372
#speed of light in cm/sh; 1 sh is 10^-8 s
const c = 299.98
const Planck_int_const = 15*a*c/(4*pi^5)

"This function maps an Nx*Ny 1-D array to a i,j on an Nx by Ny grid"
function vect_to_xy(l,Nx,Ny)
    j = div(l,Nx) + 1
    i = mod(l-j,Nx) + 1
    (i,j)
end

"This function maps i,j on an Nx by Ny grid to an Nx*Ny 1-D array"
function xy_to_vect(i,j,Nx,Ny)
l = Nx*(j-1) + i
l
end
function nm(f, fp, x; tol=sqrt(eps()))

    ctr, max_steps = 0, 100

    while (abs(f(x)) > tol) && ctr < max_steps
        x = x - f(x) / fp(x)
        ctr = ctr + 1
    end

    ctr >= max_steps ? error("Method did not converge.") : return (x, ctr)

end

"""
    create_A(D,sigma,Nr,Nz,Lr,Lz; lower_z = "refl",upper_z="refl",upper_r="refl")

Compute the matrix and RHS for an RZ diffusion problem
# Arguments
- `D::Float64(Nr,Nz)`: matrix of diffusion coefficients
- `sigma::Float64(Nr,Nz)`: matrix of absorption coefficients
- `Nr::Integer`: number of cells in r
- `Nz::Integer`: number of cells in z
- `Lr::Float64`: size of domain in r
- `Lz::Float64`: size of domain in z
...
"""
function create_A(Dleft,Dright, Dtop,Dbottom, sigma,Nr,Nz,Lr,Lz; lower_z = "refl",upper_z="refl",upper_r="refl")

    redges = range(0,stop=Lr,length=Nr+1)
    zedges = range(0,stop=Lz,length=Nz+1)

    A = spzeros(Nr*Nz,Nr*Nz)
    b = zeros(Nr*Nz)
    dz = Lz/Nz
    dr = Lr/Nr
    idz = 1/dz
    idz2 = idz*idz
    idr = 1/dr

    #loop over cells
    for i in 1:Nr
        for j in 1:Nz
            here = xy_to_vect(i,j,Nr,Nz)
            Vij = pi*(redges[i+1]^2 - redges[i]^2)
            iVij = 1/Vij
            Siplus = 2*pi*redges[i+1]
            Siminus = 2*pi*redges[i]
            #add in diag
            A[here,here] = sigma[i,j]
            #add in i+1
            if (i< Nr)
                Diplus = 2*Dright[i,j]*Dleft[i+1,j]/(Dright[i,j] + Dleft[i+1,j])
                tmp_zone = xy_to_vect(i+1,j,Nr,Nz)
                A[here,tmp_zone] = -Siplus*iVij*idr*Diplus
                A[here,here] += Siplus*iVij*idr*Diplus
            elseif (upper_r == "vacuum")
                A[here,here] += Siplus*iVij*idr*Dright[i,j] #*idr*D[i,j]
            elseif (isa(upper_r,Float64))

                A[here,here] += Siplus*iVij*idr*Dright[i,j] #Siplus*iVij*0.5*c #*idr*D[i,j]
                b[here] += Siplus*iVij*idr*Dright[i,j]*upper_r
            end
            #add in i-1
            if (i>1)
                Diminus = 2*Dleft[i,j]*Dright[i-1,j]/(Dleft[i,j] + Dright[i-1,j])
                tmp_zone = xy_to_vect(i-1,j,Nr,Nz)
                A[here,tmp_zone] = -Siminus*iVij*idr*Diminus
                A[here,here] += Siminus*iVij*idr*Diminus
            end
            #add in j+1
            if (j < Nz)
                Djplus = 2*Dtop[i,j]*Dbottom[i,j+1]/(Dtop[i,j] + Dbottom[i,j+1])
                tmp_zone = xy_to_vect(i,j+1,Nr,Nz)
                #println("l = $(tmp_zone) i = $(i), j = $(j+1)")
                A[here,tmp_zone] = -idz2*Djplus
                A[here,here] += idz2*Djplus
            elseif (upper_z == "vacuum")
                A[here,here] += idz2*Dtop[i,j]   #idz2*D[i,j]
            elseif (isa(upper_z,Float64))
                A[here,here] += idz2*Dtop[i,j] #idz*0.5*c #idz2*D[i,j] #
                b[here] += idz2*Dtop[i,j]*upper_z # idz*upper_z*0.5*c #
            end
            #add in j-1
            if (j>1)
                Djminus = 2*Dbottom[i,j]*Dtop[i,j-1]/(Dbottom[i,j] + Dtop[i,j-1])
                tmp_zone = xy_to_vect(i,j-1,Nr,Nz)
                A[here,tmp_zone] = -idz2*Djminus
                A[here,here] += idz2*Djminus
            elseif (lower_z == "vacuum")
                A[here,here] += idz2*Dbottom[i,j] #idz*0.5*c #idz2*D[i,j]
            elseif (isa(lower_z,Float64))
                A[here,here] += idz2*Dbottom[i,j] #idz*c*0.5 #
                b[here] += idz2*Dbottom[i,j]*lower_z #idz*lower_z*c #
            end
        end #for j
    end #for i
    A,b
end #function

function time_dep_RT(Tfinal,delta_t,T,Er,D_func,sigma_func,Q_func,Cv_func,EOS,invEOS,
Nr,Nz,Lr,Lz;lower_z = "refl",upper_z="refl",upper_r="refl", LOUD=0, fname="tmp.jld")
    #sigma,Q,Cv,EOS,invEOS are functions of t,T,Nr,Nz,Lr,Lz
    #D is function of t,T,Nr,Nz,,Lr,Lz,Er,sigma
    done = false
    time = 0
    dz = Lz/Nz
    dr = Lr/Nr
    times = [0.]
    steps = Int64(ceil(Tfinal/delta_t+1))
    println(steps)
    Er_t = zeros(Nz*Nr,steps+1)
    Er_t[:,1] = Er
    T_t = zeros(Nz*Nr,steps+1)
    T_t[:,1] = T
    step = 1
    while !(done)
        dt = min(Tfinal-time, delta_t)
        t = time+dt
        sigma = sigma_func(t,T,Nr,Nz,Lr,Lz)
        Cv = Cv_func(t,T,Nr,Nz,Lr,Lz)
        Q = Q_func(t,T,Nr,Nz,Lr,Lz)
        beta = 4*a*c*reshape(T,Nr,Nz).^3 ./ Cv
        f = 1 ./(1 .+ beta.*sigma*dt)
        sigma_a = sigma.*f
        sigma_star = c*sigma_a .+ 1.0/(dt)
        Trect = reshape(T,Nr,Nz)

        #define Tleft
        tmp = ones(Nr,Nz)
        tmp[2:Nr,:] = 0.5*(Trect[1:(Nr-1),:] + Trect[2:Nr,:])
        sigma_left = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
        Dleft = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_left)
        #define Tright
        tmp = ones(Nr,Nz)
        tmp[1:(Nr-1),:] = 0.5*(Trect[1:(Nr-1),:] + Trect[2:Nr,:])
        sigma_right = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
        Dright = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_right)
        #define Ttop
        tmp = ones(Nr,Nz)
        tmp[:,1:(Nz-1)] = 0.5*(Trect[:,1:(Nz-1)] + Trect[:,2:Nz])
        sigma_top = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
        Dtop = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_top)
        #define Tbottom
        tmp = ones(Nr,Nz)
        tmp[:,2:Nz] = 0.5*(Trect[:,1:(Nz-1)] + Trect[:,2:Nz])
        sigma_bottom = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
        Dbottom = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_bottom)

        A,b = create_A(c*Dleft,c*Dright, c*Dtop,c*Dbottom, sigma_star, Nr, Nz, Lr, Lz, lower_z=lower_z,upper_z=upper_z, upper_r=upper_r);
        b += Er/(dt) + reshape(Q,(Nr*Nz)) + c*a*T.^4 .*reshape(sigma_a,Nr*Nz)
        Er = A\b

        dE = EOS(t,T,Nr,Nz,Lr,Lz) + c*dt*reshape(sigma_a,Nr*Nz).*(Er - a*T.^4)
        T = copy(invEOS(t,dE,Nr,Nz,Lr,Lz))

        time += dt
        done = time >= Tfinal
        times = push!(times,time)
        if (LOUD >0)
            println("Step $(step), t = $(time)")
        end
        step += 1
        Er_t[:,step] = copy(Er)
        T_t[:,step] = copy(T)

        #save(fname, "Nr", Nr, "Nz", Nz, "dr", dr, "dz", dz, "times", times, "T", T_t, "Er", Er_t)
    end #timestep loop
    if (LOUD == -1)
        println("Step $(step-1), t = $(time)")
    end
    times,Er_t,T_t
end #function

function time_dep_3T(Tfinal,delta_t,Te,Ti,Er,D_func,sigma_func,Q_func,Qe_func,Qi_func,Cve_func,EOSe,invEOSe,Cvi_func,EOSi,invEOSi,gamma_func,
Nr,Nz,Lr,Lz;lower_z = "refl",upper_z="refl",upper_r="refl", LOUD=0, fname="tmp.jld")
    #sigma,Q,Cv,EOS,invEOS are functions of t,T,Nr,Nz,Lr,Lz
    #D is function of t,T,Nr,Nz,,Lr,Lz,Er,sigma
    done = false
    time = 0
    dz = Lz/Nz
    dr = Lr/Nr
    times = [0.]
    steps = Int64(ceil(Tfinal/delta_t+1))
    println(steps)
    Er_t = zeros(Nz*Nr,steps+1)
    Er_t[:,1] = Er
    Te_t = zeros(Nz*Nr,steps+1)
    Te_t[:,1] = Te
    Ti_t = zeros(Nz*Nr,steps+1)
    Ti_t[:,1] = Ti
    step = 1
    while !(done)
        dt = min(Tfinal-time, delta_t)
        t = time+dt
        sigma = sigma_func(t,Te,Nr,Nz,Lr,Lz)
        Cve = Cve_func(t,Te,Nr,Nz,Lr,Lz)
        Cvi = Cvi_func(t,Ti,Nr,Nz,Lr,Lz)
        Q = Q_func(t,Te,Nr,Nz,Lr,Lz)
        Qe = Qe_func(t,Te,Nr,Nz,Lr,Lz)
        Qi = Qi_func(t,Ti,Nr,Nz,Lr,Lz)
        gamma = gamma_func(t,Te,Ti,Nr,Nz,Lr,Lz)
        g_res = copy(reshape(gamma,Nr,Nz))

        Ti_rect = reshape(Ti,Nr,Nz)
        Trect = reshape(Te,Nr,Nz)
        g_EvDens = Cvi./dt .+ g_res
        Si = g_res./g_EvDens.*(Cvi./dt.*Ti_rect + Qi)
        alpha = g_res.*(g_res./g_EvDens.-1)
        Sn = alpha.*Trect + Si + Qe

        beta = 4*a*c*reshape(Te,Nr,Nz).^3 ./ Cve
        f = (Cve-alpha*dt)./((1 .+ beta.*sigma*dt).*Cve-alpha*dt) #1 ./(1 .+ beta.*sigma*dt)
        sigma_a = sigma.*f
        sigma_star = c*sigma_a .+ 1.0/(dt)

        #define Tleft
        tmp = ones(Nr,Nz)
        tmp[2:Nr,:] = 0.5*(Trect[1:(Nr-1),:] + Trect[2:Nr,:])
        sigma_left = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
        Dleft = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_left)
        #define Tright
        tmp = ones(Nr,Nz)
        tmp[1:(Nr-1),:] = 0.5*(Trect[1:(Nr-1),:] + Trect[2:Nr,:])
        sigma_right = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
        Dright = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_right)
        #define Ttop
        tmp = ones(Nr,Nz)
        tmp[:,1:(Nz-1)] = 0.5*(Trect[:,1:(Nz-1)] + Trect[:,2:Nz])
        sigma_top = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
        Dtop = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_top)
        #define Tbottom
        tmp = ones(Nr,Nz)
        tmp[:,2:Nz] = 0.5*(Trect[:,1:(Nz-1)] + Trect[:,2:Nz])
        sigma_bottom = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
        Dbottom = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_bottom)

        A,b = create_A(c*Dleft,c*Dright, c*Dtop,c*Dbottom, sigma_star, Nr, Nz, Lr, Lz, lower_z=lower_z,upper_z=upper_z, upper_r=upper_r);
        b += Er/(dt) + reshape(Q,(Nr*Nz)) + c*a*Te.^4 .*reshape(sigma_a,Nr*Nz) + reshape((1.0 .-f).*Sn,Nr*Nz)
        Er = A\b

        dE = EOSe(t,Te,Nr,Nz,Lr,Lz) + (reshape(f.*Sn,Nr*Nz)+c*reshape(sigma_a,Nr*Nz).*(Er - a*Te.^4)).*reshape((Cve./(Cve-alpha*dt)),Nr*Nz)*dt

        dEi = EOSi(t,Ti,Nr,Nz,Lr,Lz) - dt*reshape(Sn,Nr*Nz) - dt*(reshape(f.*Sn,Nr*Nz)+c*reshape(sigma_a,Nr*Nz).*(Er - a*Te.^4)).*reshape((alpha*dt./(Cve-alpha*dt)),Nr*Nz) + dt*reshape(Qe+Qi,Nr*Nz)
        Te = copy(invEOSe(t,dE,Nr,Nz,Lr,Lz))
        Ti = copy(invEOSi(t,dEi,Nr,Nz,Lr,Lz))

        time += dt
        done = time >= Tfinal
        times = push!(times,time)
        if (LOUD >0)
            println("Step $(step), t = $(time)")
        end
        step += 1
        Er_t[:,step] = copy(Er)
        Te_t[:,step] = copy(Te)
        Ti_t[:,step] = copy(Ti)

        #save(fname, "Nr", Nr, "Nz", Nz, "dr", dr, "dz", dz, "times", times, "Te", Te_t, "Ti", Ti_t, "Er", Er_t)
    end #timestep loop
    if (LOUD == -1)
        println("Step $(step-1), t = $(time)")
    end
    times,Er_t,Te_t, Ti_t
end #function

function time_dep_RT_one_it(Tfinal,delta_t,T,Er,D_func,sigma_func,Q_func,Cv_func,EOS,invEOS,
Nr,Nz,Lr,Lz;lower_z = "refl",upper_z="refl",upper_r="refl", LOUD=0, fname="tmp.jld")
    #sigma,Q,Cv,EOS,invEOS are functions of t,T,Nr,Nz,Lr,Lz
    #D is function of t,T,Nr,Nz,,Lr,Lz,Er,sigma
    done = false
    time = 0
    dz = Lz/Nz
    dr = Lr/Nr
    times = [0.]
    steps = Int64(ceil(Tfinal/delta_t+1))
    println(steps)
    Er_t = zeros(Nz*Nr,steps+1)
    Er_t[:,1] = Er
    T_t = zeros(Nz*Nr,steps+1)
    T_t[:,1] = T
    step = 1
    while !(done)
        dt = min(Tfinal-time, delta_t)
        t = time+dt
        sigma = sigma_func(t,T,Nr,Nz,Lr,Lz)
        Cv = Cv_func(t,T,Nr,Nz,Lr,Lz)
        Q = Q_func(t,T,Nr,Nz,Lr,Lz)
        beta = 4*a*c*reshape(T,Nr,Nz).^3 ./Cv
        f = 1. /(1+beta.*sigma*dt)
        sigma_a = sigma.*f
        sigma_star = c*sigma_a + 1.0/(dt)
        Trect = reshape(T,Nr,Nz)

        #define Tleft
        tmp = ones(Nr,Nz)
        tmp[2:Nr,:] = 0.5*(Trect[1:(Nr-1),:] + Trect[2:Nr,:])
        sigma_left = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
        Dleft = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_left)
        #define Tright
        tmp = ones(Nr,Nz)
        tmp[1:(Nr-1),:] = 0.5*(Trect[1:(Nr-1),:] + Trect[2:Nr,:])
        sigma_right = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
        Dright = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_right)
        #define Ttop
        tmp = ones(Nr,Nz)
        tmp[:,1:(Nz-1)] = 0.5*(Trect[:,1:(Nz-1)] + Trect[:,2:Nz])
        sigma_top = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
        Dtop = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_top)
        #define Tbottom
        tmp = ones(Nr,Nz)
        tmp[:,2:Nz] = 0.5*(Trect[:,1:(Nz-1)] + Trect[:,2:Nz])
        sigma_bottom = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
        Dbottom = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_bottom)

        A,b = create_A(c*Dleft,c*Dright, c*Dtop,c*Dbottom, sigma_star, Nr, Nz, Lr, Lz, lower_z=lower_z,upper_z=upper_z, upper_r=upper_r);
        b += Er/(dt) + reshape(Q,(Nr*Nz)) + c*a*T.^4 .*reshape(sigma_a,Nr*Nz)
        Er = A\b

        dE = EOS(t,T,Nr,Nz,Lr,Lz) + c*dt*reshape(sigma_a,Nr*Nz).*(Er - a*T.^4)
        T = copy(invEOS(t,dE,Nr,Nz,Lr,Lz))

        time += dt
        done = time >= Tfinal
        times = push!(times,time)
        if (LOUD >0)
            println("Step $(step), t = $(time)")
        end
        step += 1
        Er_t[:,step] = copy(Er)
        T_t[:,step] = copy(T)

        save(fname, "Nr", Nr, "Nz", Nz, "dr", dr, "dz", dz, "times", times, "T", T_t, "Er", Er_t)
    end #timestep loop
    if (LOUD == -1)
        println("Step $(step-1), t = $(time)")
    end
    times,Er_t,T_t
end #function

function gray_update(t,delta_t,T,Er,Er_prev_it,D_func,sigma_func,Q_func,Cv_func,EOS,invEOS,rho,
Nr,Nz,Lr,Lz;lower_z = "refl",upper_z="refl",upper_r="refl", LOUD=0, fname="tmp.jld", use_fact = 1.0)
    sigma = sigma_func(t,T,Nr,Nz,Lr,Lz)
    Cv = Cv_func(t,T,Nr,Nz,Lr,Lz)
    Q = Q_func(t,T,Nr,Nz,Lr,Lz)
    JEe = -c./Cv.*sigma*4*a.*reshape(T,Nr,Nz).^3
    Jee = rho/delta_t + c./Cv.*sigma*4*a.*reshape(T,Nr,Nz).^3
    sigma_a = sigma.*(1+use_fact*JEe./Jee)
    println((JEe./Jee))
    sigma_star = c*sigma_a + 1.0/(delta_t)
    Trect = reshape(T,Nr,Nz)

    #define Tleft
    tmp = ones(Nr,Nz)
    tmp[2:Nr,:] = 0.5*(Trect[1:(Nr-1),:] + Trect[2:Nr,:])
    sigma_left = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
    Dleft = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_left)
    #define Tright
    tmp = ones(Nr,Nz)
    tmp[1:(Nr-1),:] = 0.5*(Trect[1:(Nr-1),:] + Trect[2:Nr,:])
    sigma_right = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
    Dright = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_right)
    #define Ttop
    tmp = ones(Nr,Nz)
    tmp[:,1:(Nz-1)] = 0.5*(Trect[:,1:(Nz-1)] + Trect[:,2:Nz])
    sigma_top = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
    Dtop = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_top)
    #define Tbottom
    tmp = ones(Nr,Nz)
    tmp[:,2:Nz] = 0.5*(Trect[:,1:(Nz-1)] + Trect[:,2:Nz])
    sigma_bottom = sigma_func(t,tmp,Nr,Nz,Lr,Lz)
    Dbottom = D_func(t,tmp,Nr,Nz,Lr,Lz,Er,sigma_bottom)

    A,b = create_A(c*Dleft,c*Dright, c*Dtop,c*Dbottom, sigma_star, Nr, Nz, Lr, Lz, lower_z=lower_z,upper_z=upper_z, upper_r=upper_r);
    b += Er/(dt) + reshape(Q,(Nr*Nz)) + c*a*reshape(sigma_a,Nr*Nz).*(T.^4 + use_fact*reshape(JEe./Jee,Nr*Nz).*Er_prev_it)
    #println(b,T,sigma,Q,Er/dt)
    Er = A\b

    return Er
    end #function

function gray_GARNET(Tfinal,delta_t,T,Er,D_func,sigma_func,Q_func,Cv_func,EOS,invEOS,single_invEOS,rho,
Nr,Nz,Lr,Lz;lower_z = "refl",upper_z="refl",upper_r="refl", LOUD=0, fname="tmp.jld", tol_e = 1e-6, tol_Er = 1e-6, max_its = 20, conserve = true)
    #sigma,Q,Cv,EOS,invEOS are functions of t,T,Nr,Nz,Lr,Lz
    #D is function of t,T,Nr,Nz,,Lr,Lz,Er,sigma
    done = false
    time = 0
    dz = Lz/Nz
    dr = Lr/Nr
    times = [0.]
    steps = Int64(round(Tfinal/delta_t+1))
    println(steps)
    Er_t = zeros(Nz*Nr,steps)
    Er_t[:,1] = Er
    T_t = zeros(Nz*Nr,steps)
    diff_solves = zeros(steps)
    outer_its = zeros(steps)
    T_t[:,1] = T
    step = 1
    eta = zeros(Nr,Nz)
    del_i = 1e-6*1im
    it = 1
    while !(done)
        dt = min(Tfinal-time, delta_t)
        t = time+dt
        e_old = EOS(t,T_t[:,step],Nr,Nz,Lr,Lz)
        e = copy(e_old)
        sigma = sigma_func(t,T,Nr,Nz,Lr,Lz)
        #initialize temperature
        for i in 1:Nr
            for j in 1:Nz
                l = xy_to_vect(i,j,Nr,Nz)
                f(e) = rho[i,j]/dt*(e-e_old[l]) + c*sigma[i,j]*(a*single_invEOS(t,e,i,j,Nr,Nz,Lr,Lz)^4 - Er[l])
                fp(e) = imag(f(e+del_i)/imag(del_i))
                e[l],its = nm(f, fp, e[l]; tol=sqrt(eps()))

                ft(T) = a*T^4 + EOS(t,T,i,j,Nr,Nz,Lr,Lz) - (e_old[l] + Er[l])
                ftp(T) = imag(ft(T+del_i)/imag(del_i))
                Teq, its = nm(ft, ftp, 0.5*(T[l]+ (Er[l]/a)^.25); tol=sqrt(eps()))
                e_eq = EOS(t,Teq,i,j,Nr,Nz,Lr,Lz)
                Er_eq = a*Teq^4

                println("Er_prev = ",Er[l], " Er_eq = ", Er_eq, " T_prev = ", T_t[l,step], " T_init = ",
                    single_invEOS(t,e[l],i,j,Nr,Nz,Lr,Lz), " T_eq = ", Teq)
                omega = 1.0/(1+c*sigma[i,j]*dt)
                e[l] = e[l] * omega + e_eq*(1-omega)
                Er[l] = Er[l]*omega + Er_eq*(1-omega)

            end #for j
        end #for i
        T = copy(invEOS(t,e,Nr,Nz,Lr,Lz))
        e_oi = copy(e)
        #println("e = ",e)
        #println("T = ",T)
        converged = false
        Er_o = copy(Er_t[:,step])
        Er_oi = copy(Er_o)
        while !(converged)
            conv_E = false
            E_its = 1
            sigma = sigma_func(t,T,Nr,Nz,Lr,Lz)
            Cv = Cv_func(t,T,Nr,Nz,Lr,Lz)
            Er = gray_update(t,dt,T,Er_t[:,step],Er_oi,D_func,sigma_func,Q_func,Cv_func,EOS,invEOS,rho,
                             Nr,Nz,Lr,Lz;lower_z=lower_z,upper_z=upper_z,upper_r=upper_r, LOUD=LOUD, fname=fname)
            diff_solves[step] += 1
            #solve for new temperature in each zone
            for i in 1:Nr
                for j in 1:Nz
                    l = xy_to_vect(i,j,Nr,Nz)
                    f(e) = rho[i,j]/dt*(e-e_old[l]) + c*sigma[i,j]*(a*single_invEOS(t,e,i,j,Nr,Nz,Lr,Lz)^4 - Er[l])
                    fp(e) = imag(f(e+del_i)/imag(del_i))
                    #println(fp(e[l])," ",(f(e[l]+1e-6) - f(e[l]))/1.0e-6)
                    e[l],its = nm(f, fp, e[l]; tol=sqrt(eps()))
                end #for j
            end #for i
            T = copy(invEOS(t,e,Nr,Nz,Lr,Lz))

            #println("e = ",e)
            #println("Er = ",Er)
            #println("T = ",T)
            converged = ( (norm(e - e_oi) < tol_e*norm(e)) & (norm(Er - Er_oi) < tol_Er*norm(Er))) | (it >= max_its)
            #println(norm(e - e_oi)/norm(e) , norm(Er - Er_oi)/norm(Er))
            Er_oi = copy(Er)
            Er_o = copy(Er)
            e_oi = copy(e)
            it += 1
            outer_its[step] += 1
        end

        if (conserve == true)
        #do one more solve and energy update
            Er = gray_update(t,dt,T,Er_t[:,step],0*Er,D_func,sigma_func,Q_func,Cv_func,EOS,invEOS,rho,
                                 Nr,Nz,Lr,Lz;lower_z=lower_z,upper_z=upper_z,upper_r=upper_r, LOUD=LOUD, fname=fname, use_fact = 0.0)
            diff_solves[step] += 1
            #update e in each zone
            for i in 1:Nr
                for j in 1:Nz
                    l = xy_to_vect(i,j,Nr,Nz)
                    e[l]= e_old[l] - dt/rho[i,j]*(c*sigma[i,j]*(a*T[l]^4 - Er[l]))
                end #for j
            end #for i
            T = copy(invEOS(t,e,Nr,Nz,Lr,Lz))

        end

        time += dt
        done = time >= Tfinal
        times = push!(times,time)
        if (LOUD >0)
            println("Step $(step), t = $(time)")
            println("Energy = ", sum(e+Er))
        end
        step += 1
        Er_t[:,step] = copy(Er)
        T_t[:,step] = copy(T)

        save(fname, "Nr", Nr, "Nz", Nz, "dr", dr, "dz", dz, "times", times, "T", T_t, "Er", Er_t)
    end #timestep loop
    if (LOUD == -1)
        println("Step $(step-1), t = $(time)")
    end
    times,Er_t,T_t, diff_solves, outer_its
end #function

function meshgrid(x,y)
    Nx = length(x)
    Ny = length(y)
    X = zeros(Nx,Ny)
    Y = zeros(Nx,Ny)
    for i in 1:Nx
        for j in 1:Ny
            X[i,j] = x[i]
            Y[i,j] = y[j]
        end
    end
    X,Y
end

max_stars_repo_name: DrRyanMc/HEDPRadDiffusion

max_issues_repo_name: DrRyanMc/HEDPRadDiffusion

max_forks_repo_name: DrRyanMc/HEDPRadDiffusion


text:
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:light
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.2.4
#   kernelspec:
#     display_name: Julia 1.1.0
#     language: julia
#     name: julia-1.1
# ---

# # This notebook
#
# A simplified version of `ex2_more.jl` used in `runtests.jl` (avoids downloading binary data).
#
# _Notes:_ For documentation see <https://gaelforget.github.io/MeshArrays.jl/stable/>, <https://docs.juliadiffeq.org/latest/solvers/ode_solve.html> and <https://en.wikipedia.org/wiki/Displacement_(vector)>

# ## 1. import software and `pkg/flt` trajectories

using IndividualDisplacements, MeshArrays, Plots, DifferentialEquations

# ## 2. Define gridded variables as `MeshArray`s

# Put grid variables in a dictionary:

# +
mygrid=gcmgrid("flt_example/","ll",1,[(80,42)], [80 42], Float32, read, write)
nr=8

XC=MeshArray(mygrid,Float32); XC[1]=vec(2500.:5000.:397500.0)*ones(1,42);
XG=MeshArray(mygrid,Float32); XG[1]=vec(0.:5000.:395000.0)*ones(1,42);
YC=MeshArray(mygrid,Float32); YC[1]=ones(80,1)*transpose(vec(2500.:5000.:207500.0));
YG=MeshArray(mygrid,Float32); YG[1]=ones(80,1)*transpose(vec(0.:5000.:205000.0));

GridVariables=Dict("XC" => XC,"YC" => YC,"XG" => XG,"YG" => YG,"dx" => 5000.0);
# -

# Put velocity fields in a dictionary:

# +
t0=0.0 #approximation / simplification
t1=18001.0*3600.0

u0=-(YG.-YC[1][40,21])/2000000.; u1=u0
v0=(XG.-XC[1][40,21])/2000000.; v1=v0

u0=u0./GridVariables["dx"]
u1=u1./GridVariables["dx"]
v0=v0./GridVariables["dx"]
v1=v1./GridVariables["dx"]

uvt = Dict("u0" => u0, "u1" => u1, "v0" => v0, "v1" => v1, "t0" => t0, "t1" => t1)
uvetc=merge(uvt,GridVariables);
# -

# ## 3. Define initial condition and time period

uInit=[200000.0;0.0]./uvetc["dx"]
nSteps=3000-2
du=fill(0.0,2);

# ## 4. solve through time using DifferentialEquations.jl

using DifferentialEquations
tspan = (0.0,nSteps*3600.0)
prob = ODEProblem(IndividualDisplacements.VelComp,uInit,tspan,uvetc)
sol = solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)
sol[1:4]

# ## 5. visualize trajectory

using Plots
Plots.plot(sol[1,:],sol[2,:],linewidth=5,title="Solid body rotation example",
     xaxis="lon",yaxis="lat",label="Julia Solution") # legend=false

# ## 6. check against reference result

sol[1,end],sol[2,end]
isapprox(sol[1,end],117237.0./uvetc["dx"]; atol=1.)
isapprox(sol[2,end],40448.0./uvetc["dx"]; atol=1.)


text:
using ONNX, Flux, ProtoBuf
include("ops_tests.jl")

# Cast double to float
main_test("$ONNX_TEST_PATH/test_cast_DOUBLE_to_FLOAT",
            read_output("$ONNX_TEST_PATH/test_cast_DOUBLE_to_FLOAT"),
                read_input("$ONNX_TEST_PATH/test_cast_DOUBLE_to_FLOAT")[1])

# Cast double to float16
main_test("$ONNX_TEST_PATH/test_cast_DOUBLE_to_FLOAT16",
            read_output("$ONNX_TEST_PATH/test_cast_DOUBLE_to_FLOAT16"),
                read_input("$ONNX_TEST_PATH/test_cast_DOUBLE_to_FLOAT16")[1])

# Cast Float16 to double
main_test("$ONNX_TEST_PATH/test_cast_FLOAT16_to_DOUBLE",
            read_output("$ONNX_TEST_PATH/test_cast_FLOAT16_to_DOUBLE"),
                read_input("$ONNX_TEST_PATH/test_cast_FLOAT16_to_DOUBLE")[1])

# Cast Float16 to Float
main_test("$ONNX_TEST_PATH/test_cast_FLOAT16_to_FLOAT",
            read_output("$ONNX_TEST_PATH/test_cast_FLOAT16_to_FLOAT"),
                read_input("$ONNX_TEST_PATH/test_cast_FLOAT16_to_FLOAT")[1])

# Cast Float to Double
main_test("$ONNX_TEST_PATH/test_cast_FLOAT_to_DOUBLE",
            read_output("$ONNX_TEST_PATH/test_cast_FLOAT_to_DOUBLE"),
                read_input("$ONNX_TEST_PATH/test_cast_FLOAT_to_DOUBLE")[1])

# Cast Float to Float16
main_test("$ONNX_TEST_PATH/test_cast_FLOAT_to_FLOAT16",
            read_output("$ONNX_TEST_PATH/test_cast_FLOAT_to_FLOAT16"),
                read_input("$ONNX_TEST_PATH/test_cast_FLOAT_to_FLOAT16")[1])