text:
using DrWatson
quickactivate(@__DIR__)
@show projectdir()
#-----------------------------------------------------------------------------
# load parameters
using ArgParse
using Dates
argtable3 = ArgParseSettings(
    description="This script samples from National.model using NUTS."
)
@add_arg_table! argtable3 begin
    "--chunksize"
        help = "chunksize to be used by ForwardDiff.jl"
        arg_type = Int
        default = 40
    "--num-samples", "-n"
        help = "number of samples"
        arg_type = Int
        default = 10
    "--num-chains", "-t"
        help = "number of chains sampled with multi-threading"
        arg_type = Int
        default = 1
    "--num-warmup", "-w"
        help = "number of samples to use for warmup/adaptation"
        arg_type = Int
        default = 10
    "--observations-end", "-o"
        help = "end of observation date as string yyyy-mm-dd. Use full data set by defaults"
        arg_type = String
        default = nothing
    "--cases-start", "-s"
        help = "start date for case observations as string yyyy-mm-dd. No case observations by default"
        arg_type = String
        default = "2020-06-01"
    "--covariates-start", "-p"
        help = "start date for covariates as string yyyy-mm-dd. starts on 2020-11-10 by default"
        arg_type = String
        default = "2020-11-10"
    "--add-covariates", "-c"
        help = "list of covariate names (default empty string = no covariates). Choose: C: contacts, R:"
        arg_type = String
        default = nothing
    "--model", "-m"
        help = "choose from: 'v1', 'gp'. Defaul: 'v1'"
        arg_type = String
        default = "hospit"
    "--semiparametric", "-S"
        arg_type = Bool
        default = true
    "--foldername", "-f"
        arg_type = String
        default = "$(today())"
    "--name-prefix", "-u"
        arg_type = String
        default = ""
    "--seed"
        help = "random seed to use"
        arg_type = Int
        default = nothing
end
parsed_args = parse_args(ARGS, argtable3)
#-----------------------------------------------------------------------------
# load packages
@info "load packages"
using CovidSurvey
using Random, Turing
using Base.Threads
using StatsPlots
using PrettyTables
using DataFrames
using BSON

using Memoization
using ReverseDiff
Turing.setrdcache(true)
setadbackend(:reversediff)

plotlyjs()
@info "number of threads available: $(nthreads())"
#----------------------------------------------------------------------------
# load data
@info "load data"
data = National.load_data(
      observationsend  = parsed_args["observations-end"]
    , predictors        = parsed_args["add-covariates"]|> CovidSurvey.parse_predictors
    , rwstep           = 1
    , epidemicstart    = 30
    , numimpute        = 6
    , deathmodel        = parsed_args["model"] == "cases" ? National.DeathInit(obs_stop="$( Date(parsed_args["cases-start"]) + Day(30))") : National.DeathInit()
    , casemodel         = National.CaseInit(obs_start=parsed_args["cases-start"])
    # , hospitmodel       = National.HospitInit()
    , seromodel         = National.SeroInit1(delay=0, std=.5)
    , link              = KLogistic(3.)
    , invlink           = KLogit(3.)
    , lockdown          = "2020-03-18"
    , covariates_kwargs = Dict(
        :semiparametric => parsed_args["semiparametric"],
        :fname => projectdir("data","raw_contact_rates.csv"),
        :shift => -1,
        :startdate => parsed_args["covariates-start"],
        :enddate => nothing))

turing_data = data.turing_data;
#----------------------------------------------------------------------------
# prepare model
ps = (
    regional = false,
    model    = parsed_args["model"],
    preds    = parsed_args["add-covariates"],
    steps    = parsed_args["num-samples"],
    cases    = parsed_args["cases-start"],
    warmup   = parsed_args["num-warmup"],
    observ   = parsed_args["observations-end"],
    chains   = parsed_args["num-chains"],
    seed     = parsed_args["seed"],
    semipar  = parsed_args["semiparametric"],
    predstart= parsed_args["covariates-start"]
)
PREFIX = isempty(parsed_args["name-prefix"]) ? "" : parsed_args["name-prefix"]*"-"

model = National.name2model[ps.model]
# model = name2model["hospit"]
!isnothing(ps.seed) && Random.seed!(ps.seed);
@info ps
m = model(turing_data, false)
#----------------------------------------------------------------------------
@time chain = let
    thinning = 10
    if ps.chains > 1
        sample(m, NUTS(ps.warmup, 0.95; max_depth=5), MCMCThreads(), ps.steps + ps.warmup, ps.chains; progress=true, thinning)
    else
        sample(m, NUTS(ps.warmup, 0.95; max_depth=5), ps.steps + ps.warmup; progress=true, thinning) #; max_depth=15
    end
end
#-----------------------------------------------------------------------------
@info "save result"
fdir = projectdir("reports/", parsed_args["foldername"])
fname = National.save_results(fdir, ps, data, chain)

#-----------------------------------------------------------------------------
@info "post processing"
National.postprocessing(fname; plot_results=false)

max_stars_repo_name: andreaskoher/CovidSurvey

max_issues_repo_name: andreaskoher/CovidSurvey

max_forks_repo_name: andreaskoher/CovidSurvey


text:
using DynamicHMCModels
using Test

scripts = [
  "../scripts/02/m2.1d.jl",
  "../scripts/04/m4.1d.jl",
  "../scripts/04/m4.2d.jl",
  "../scripts/04/m4.5d.jl",
  "../scripts/04/m4.5d1.jl",
  "../scripts/05/m5.1d.jl",
  "../scripts/05/m5.3d.jl",
  "../scripts/05/m5.6d.jl",
  "../scripts/10/m10.2d.jl",
  "../scripts/10/m10.4d.jl",
  "../scripts/12/m12.6d.jl"
]

for script in scripts
  println("\n * $script *\n")
  include(script)
  println("\n * $script completed\n")
end

ext: jl


text:
using LinearAlgebra, SparseArrays

"""
    QuadtoSOCBridge{T}

The set of points `x` satisfying the constraint
```math
\\frac{1}{2}x^T Q x + a^T x + b \\le 0
```
is a convex set if `Q` is positive semidefinite and is the union of two convex
cones if `a` and `b` are zero (i.e. *homogeneous* case) and `Q` has only one
negative eigenvalue. Currently, only the non-homogeneous transformation
is implemented, see the Note section below for more details.

## Non-homogeneous case

If `Q` is positive semidefinite, there exists `U` such that ``Q = U^T U``, the
inequality can then be rewritten as
```math
\\|U x\\|_2^2 \\le 2 (-a^T x - b)
```
which is equivalent to the membership of `(1, -a^T x - b, Ux)` to the rotated
second-order cone.

## Homogeneous case

If `Q` has only one negative eigenvalue, the set of `x` such that ``x^T Q x \\le
0`` is the union of a convex cone and its opposite. We can choose which one to
model by checking the existence of bounds on variables as shown below.

### Second-order cone

If `Q` is diagonal and has eigenvalues `(1, 1, -1)`, the inequality
``x^2 + x^2 \\le z^2`` combined with ``z \\ge 0`` defines the Lorenz cone (i.e.
the second-order cone) but when combined with ``z \\le 0``, it gives the
opposite of the second order cone. Therefore, we need to check if the variable
`z` has a lower bound 0 or an upper bound 0 in order to determine which cone is

### Rotated second-order cone

The matrix `Q` corresponding to the inequality ``x^2 \\le 2yz`` has one
eigenvalue 1 with eigenvectors `(1, 0, 0)` and `(0, 1, -1)` and one eigenvalue
`-1` corresponding to the eigenvector `(0, 1, 1)`. Hence if we intersect this
union of two convex cone with the halfspace ``x + y \\ge 0``, we get the rotated
second-order cone and if we intersect it with the halfspace ``x + y \\le 0`` we
get the opposite of the rotated second-order cone. Note that `y` and `z` have
the same sign since `yz` is nonnegative hence ``x + y \\ge 0`` is equivalent to
``x \\ge 0`` and ``y \\ge 0``.

### Note

The check for existence of bound can be implemented (but inefficiently) with the
current interface but if bound is removed or transformed (e.g. `≤ 0` transformed
into `≥ 0`) then the bridge is no longer valid. For this reason the homogeneous
version of the bridge is not implemented yet.
"""
struct QuadtoSOCBridge{T} <: AbstractBridge
    soc::CI{MOI.VectorAffineFunction{T}, MOI.RotatedSecondOrderCone}
    dimension::Int  # dimension of the SOC constraint
    less_than::Bool # whether the constraint was ≤ or ≥
    set_constant::T # the constant that was on the set
end
function bridge_constraint(::Type{QuadtoSOCBridge{T}}, model,
                           func::MOI.ScalarQuadraticFunction{T},
                           set::Union{MOI.LessThan{T},
                                      MOI.GreaterThan{T}}) where T
    set_constant = MOI.constant(set)
    less_than = set isa MOI.LessThan
    Q, index_to_variable_map = matrix_from_quadratic_terms(func.quadratic_terms)
    if !less_than
        rmul!(Q, -1)
    end
    # We have L × L' ≈ Q[p, p]
    L, p = try
        F = cholesky(Symmetric(Q))
        sparse(F.L), F.p
    catch err
        if err isa PosDefException
            error("The optimizer supports second-order cone constraints and",
                  " not quadratic constraints but you entered a quadratic",
                  " constraint of type: `$(typeof(func))`-in-`$(typeof(set))`.",
                  " A bridge attempted to transform the quadratic constraint",
                  " to a second order cone constraint but the constraint is",
                  " not strongly convex, i.e., the symmetric matrix of",
                  " quadratic coefficients is not positive definite. Convex",
                  " constraints that are not strongly convex, i.e. the matrix is",
                  " positive semidefinite but not positive definite, are not",
                  " supported yet.")
        else
            rethrow(err)
        end
    end
    Ux_terms = matrix_to_vector_affine_terms(L, p, index_to_variable_map)
    Ux = MOI.VectorAffineFunction(Ux_terms, zeros(T, size(L, 2)))
    t = MOI.ScalarAffineFunction(less_than ? MOIU.operate_terms(-, func.affine_terms) : func.affine_terms,
                                 less_than ? set_constant - func.constant : func.constant - set_constant)
    f = MOIU.operate(vcat, T, one(T), t, Ux)
    dimension = MOI.output_dimension(f)
    soc = MOI.add_constraint(model, f,
                             MOI.RotatedSecondOrderCone(dimension))
    return QuadtoSOCBridge(soc, dimension, less_than, set_constant)
end

function matrix_from_quadratic_terms(terms::Vector{MOI.ScalarQuadraticTerm{T}}) where T
    variable_to_index_map = Dict{VI, Int}()
    index_to_variable_map = Dict{Int, VI}()
    n = 0
    for term in terms
        for variable in (term.variable_index_1, term.variable_index_2)
            if !(variable in keys(variable_to_index_map))
                n += 1
                variable_to_index_map[variable] = n
                index_to_variable_map[n] = variable
            end
        end
    end
    I = Int[]
    J = Int[]
    V = T[]
    for term in terms
        i = variable_to_index_map[term.variable_index_1]
        j = variable_to_index_map[term.variable_index_2]
        push!(I, i)
        push!(J, j)
        push!(V, term.coefficient)
        if i != j
            push!(I, j)
            push!(J, i)
            push!(V, term.coefficient)
        end
    end
    # Duplicate terms are summed together in `sparse`
    Q = sparse(I, J, V, n, n)
    return Q, index_to_variable_map
end

function matrix_to_vector_affine_terms(L::SparseMatrixCSC{T},
                                       p::Vector,
                                       index_to_variable_map::Dict{Int, VI}) where T
    # We know that L × L' ≈ Q[p, p] hence (L × L')[i, :] ≈ Q[p[i], p]
    # We precompute the map to avoid having to do a dictionary lookup for every
    # term
    variable = map(i -> index_to_variable_map[p[i]], 1:size(L, 1))
    function term(i::Integer, j::Integer, v::T)
        return MOI.VectorAffineTerm(j, MOI.ScalarAffineTerm(v, variable[i]))
    end
    return map(ijv -> term(ijv...), zip(findnz(L)...))
end

function MOI.supports_constraint(::Type{QuadtoSOCBridge{T}},
                                 ::Type{MOI.ScalarQuadraticFunction{T}},
                                 ::Type{<:Union{MOI.LessThan{T},
                                                MOI.GreaterThan{T}}}) where T
    return true
end
function added_constraint_types(::Type{QuadtoSOCBridge{T}}) where T
    return [(MOI.VectorAffineFunction{T}, MOI.RotatedSecondOrderCone)]
end
function concrete_bridge_type(::Type{<:QuadtoSOCBridge{T}},
                              ::Type{MOI.ScalarQuadraticFunction{T}},
                              ::Type{<:Union{MOI.LessThan{T},
                                               MOI.GreaterThan{T}}}) where T
    return QuadtoSOCBridge{T}
end

# Attributes, Bridge acting as an model
function MOI.get(::QuadtoSOCBridge{T},
                 ::MOI.NumberOfConstraints{MOI.VectorAffineFunction{T},
                                           MOI.RotatedSecondOrderCone}) where T
    return 1
end
function MOI.get(bridge::QuadtoSOCBridge{T},
                 ::MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{T},
                                               MOI.RotatedSecondOrderCone}) where T
    return [bridge.soc]
end

# References
function MOI.delete(model::MOI.ModelLike, bridge::QuadtoSOCBridge)
    MOI.delete(model, bridge.soc)
end

# Attributes, Bridge acting as a constraint
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintPrimal,
                 bridge::QuadtoSOCBridge)
    soc = MOI.get(model, attr, bridge.soc)
    output = sum(soc[i]^2 for i in 3:bridge.dimension)
    output /= 2
    output -= soc[1] * soc[2]
    if !bridge.less_than
        output = -output
    end
    output += bridge.set_constant
    return output
end

# TODO
#function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintFunction,
#                 b::QuadtoSOCBridge)
#end

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintSet,
                 b::QuadtoSOCBridge{T}) where T
    if b.less_than
        return MOI.LessThan(b.set_constant)
    else
        return MOI.GreaterThan(b.set_constant)
    end
end
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintFunction,
                 b::QuadtoSOCBridge{T}) where T
    f = MOI.get(model, attr, b.soc)
    fs = MOIU.eachscalar(f)
    q = zero(MOI.ScalarQuadraticFunction{T})
    for i in 3:MOI.output_dimension(f)
        term = MOIU.operate(*, T, fs[i], fs[i])
        term = MOIU.operate!(/, T, term, 2 * one(T))
        MOIU.operate!(+, T, q, term)
    end
    MOIU.operate!(-, T, q, fs[2])
    if !b.less_than
        MOIU.operate!(-, T, q)
    end
    q.constant += b.set_constant
    return q
end


text:
module Solvers

using NLopt, NLsolve, DiffResults, ForwardDiff, LinearAlgebra

include("tunneling.jl")

end # module


text:
struct RewardTest <: TrainingTest
    idx::Integer
    state::AbstractState
    image::AbstractImage
    action::Integer
    action_image::AbstractImage
    reward::Float32
end
get_action_image(x::RewardTest) = x.action_image

mutable struct RewardTests <: TrainingTests
    tests::Vector{RewardTest}
    images::AbstractArray{Float32, 4}
    action_images::AbstractArray{Float32, 4}
    actual_reward::AbstractArray{Float32, 2}
    period::Int
    warmup::Int
    function RewardTests(tests; period=1, warmup=0)
        new(
            tests,
            tensorize(get_image.(tests)),
            tensorize(get_action_image.(tests)),
            Flux.unsqueeze(get_reward.(tests), 1),
            period,
            warmup,
        )
    end
end

function gettestfn(tests::RewardTests)
    function fn(n::Network)
        initial_output = initialinferencefn(tests)(n)
        getrewardmse(
            tests,
            recurrentinferencefn(tests, initial_output)(n)
        )
    end
end

struct RewardTestsSVG
    tests::RewardTests
end
svg(x::RewardTests) = RewardTestsSVG(x)

function Base.show(io::IO, m::MIME"image/svg+xml", x::RewardTestsSVG)
    num_tests = length(x.tests)
    range = 1:(min(10, num_tests))
    subset = 1:num_tests
    subset_indices = subset[range]
    svgbatch = [gridsvg(x.tests[i]) for i in subset_indices]
    Base.show(io, m, GridSVG(svgbatch, main_title="Reward Tests"))
end

# struct RewardTest
#     idx::Integer
#     state::AbstractState
#     action::Integer
#     image::AbstractImage
#     action_image::AbstractImage
#     reward::Float32
# end

# get_reward(x::RewardTest) = x.reward
# get_image(x::RewardTest) = x.image
# get_action_image(x::RewardTest) = x.action_image

# struct RewardTestSVG
#     test::RewardTest
# end

# svg(x::RewardTest) = RewardTestSVG(x)
# get_aspect_ratio(::Type{RewardTestSVG}) = (2,1)
# get_scale(::Type{RewardTestSVG}) = 600

# function gridsvg(test::RewardTest)
#     # @todo provision here for multiple pieces
#     # per cell
#     s = test.state
#     a = test.actions[1]
#     s_next = next_state(s, a)
#     highlight_coord = diff_coord(s_next, s)
#     action_svg = gridsvg(
#         s_next;
#         title="test action",
#         title_offset=(-210,-195),
#         highlight=highlight_coord
#     )
#     GridSVG([
#         gridsvg(s),
#         action_svg,
#     ], dim=2)
# end

# function Base.show(io::IO, m::MIME"text/plain", x::RewardTest)
#     println()
#     println("Reward Test #$(x.idx)")
#     Base.show(io, m, x.state)
#     println("\nTo Play: $(x.state.to_play)")
#     println("Test Actions: $(join(x.actions, ","))\n")
# end

# function Base.show(io::IO, m::MIME"image/svg+xml", x::RewardTestSVG)
#     Base.show(io, m, gridsvg(x.test))
# end

# import Flux

# struct RewardTests
#     tests::Vector{RewardTest}
#     images::AbstractArray{Float32, 4}
#     action_images::AbstractArray{Float32, 4}
#     actual_reward::AbstractArray{Float32, 2}
#     function RewardTests(tests)
#         new(
#             tests,
#             tensorize(map(get_image, tests)),
#             tensorize(map(get_action_image, tests)),
#             Flux.unsqueeze(map(get_reward, tests), 1)
#         )
#     end
# end

# function gettestfn(reward_tests::RewardTests)
#     function fn(n::Network)
#         initial_output = initial_inference(n, reward_tests.images) |> NetworkOutput
#         hs = hidden_state(initial_output)
#         recurrent_output = recurrent_inference(n, hs, reward_tests.action_images) |> NetworkOutput
#         pred = reward(recurrent_output)
#         Flux.mse(pred, reward_tests.actual)
#     end
# end

# struct RewardTestsSVG
#     tests::RewardTests
# end

# using Random

# svg(x::RewardTests) = RewardTestsSVG(x)
# Base.rand(x::RewardTests) = rand(x.tests)
# Base.getindex(x::RewardTests, idx::Integer) = x.tests[idx]
# Base.getindex(x::RewardTests, idx::Vector{<:Integer}) = x.tests[idx]
# Base.length(x::RewardTests) = length(x.tests)
# Base.lastindex(x::RewardTests) = length(x)
# Base.iterate(x::RewardTests, state=1) = state > length(x) ? nothing : (x[state], state+1)

# function Base.show(io::IO, m::MIME"image/svg+xml", x::RewardTestsSVG)
#     num_tests = length(x.tests)
#     range = 1:(min(10, num_tests))
#     subset = 1:num_tests
#     subset_indices = subset[range]
#     svgbatch = [gridsvg(x.tests[i]) for i in subset_indices]
#     Base.show(io, m, GridSVG(svgbatch, main_title="Reward Tests"))
# end


text:
"""
    flatten 2D matrix into (m*n, 1) matrix

        mainly used for images to flatten images

    inputs:
        x := 3D (rgp, m, n) matrix

    outputs:
        y := 2D (rgp*m*n, 1) matrix
"""
function flatten(x)
    rgp, n, m = size(x)
    return reshape(x, (rgp*m*n, 1))
end

"""
    perform the forward propagation using

    input:
        x := (n0, m) matrix
        y := (c,  m) matrix where c is the number of classes

    return cache of A, Z, Yhat, Cost
"""
function forwardProp(X::Matrix{T},
                     Y::Matrix{T},
                     model::Model) where {T}

    W::AbstractArray{Matrix{T},1},
    B::AbstractArray{Matrix{T},1},
    layers::AbstractArray{Layer,1},
    lossFun = model.W, model.B, model.layers, model.lossFun
    regulization = model.regulization
    λ = model.λ
    m = size(X)[2]
    c = size(Y)[1]
    L = length(layers)
    A = Vector{Matrix{eltype(X)}}()
    Z = Vector{Matrix{eltype(X)}}()
    push!(Z, W[1]*X .+ B[1])
    actFun = layers[1].actFun
    push!(A, eval(:($actFun.($Z[1]))))
    for l=2:L
        push!(Z, W[l]*A[l-1] .+ B[l])
        actFun = layers[l].actFun

        if isequal(actFun, :softmax)
            a = Matrix{eltype(X)}(undef, c, 0)
            for i=1:m
                zCol = Z[l][:,i]
                a = hcat(a, eval(:($(actFun)($zCol))))
            end
            push!(A, a)
        else
            push!(A, eval(:($(actFun).($Z[$l]))))
        end
    end

    if isequal(lossFun, :categoricalCrossentropy)
        cost = sum(eval(:($lossFun.($A[$L], $Y))))/ m
    else
        cost = sum(eval(:($lossFun.($A[$L], $Y)))) / (m*c)
    end

    if regulization > 0
        cost += (λ/2m) * sum([norm(w, regulization) for w in W])
    end

    return Dict("A"=>A,
                "Z"=>Z,
                "Yhat"=>A[L],
                "Cost"=>cost)
end #forwardProp

# export forwardProp

"""
    predict Y using the model and the input X and the labels Y
    inputs:
        model := the trained model
        X := the input matrix
        Y := the input labels to compare with

    output:
        a Dict of
        "Yhat" := the predicted values
        "Yhat_bools" := the predicted labels
        "accuracy" := the accuracy of the predicted labels
"""
function predict(model::Model, X, Y)
    Ŷ = forwardProp(X, Y, model)["Yhat"]
    T = eltype(Ŷ)
    layers = model.layers
    c, m = size(Y)
    # if isbool(Y)
    acc = 0
    if isequal(layers[end].actFun, :softmax)
        Ŷ_bool = BitArray(undef, c, 0)
        p = Progress(size(Ŷ)[2], 0.01)
        for v in eachcol(Ŷ)
            Ŷ_bool = hcat(Ŷ_bool, v .== maximum(v))
            next!(p)
        end

        acc = sum([Ŷ_bool[:,i] == Y[:,i] for i=1:size(Y)[2]])/m
        println("Accuracy is = $acc")
    end

    if isequal(layers[end].actFun, :σ)
        Ŷ_bool = BitArray(undef, c, 0)
        p = Progress(size(Ŷ)[2], 0.01)
        for v in eachcol(Ŷ)
            Ŷ_bool = hcat(Ŷ_bool, v .> T(0.5))
            next!(p)
        end

        acc = sum(Ŷ_bool .== Y)/(c*m)
        println("Accuracy is = $acc")
    end
    return Dict("Yhat"=>Ŷ,
                "Yhat_bool"=>Ŷ_bool,
                "accuracy"=>acc)
end #predict

function backProp(X,Y,
                  model::Model,
                  cache::Dict{})

    layers::AbstractArray{Layer, 1} = model.layers
    lossFun = model.lossFun
    m = size(X)[2]
    L = length(layers)
    A, Z = cache["A"], cache["Z"]
    W, B, regulization, λ = model.W, model.B, model.regulization, model.λ

    D = [rand(size(A[i])...) .< layers[i].keepProb for i=1:L]

    if layers[L].keepProb < 1.0 #to save time of multiplication in case keepProb was one
        A = [A[i] .* D[i] for i=1:L]
        A = [A[i] ./ layers[i].keepProb for i=1:L]
    end
    # init all Arrays
    dA = Vector{Matrix{eltype(A[1])}}([similar(mat) for mat in A])
    dZ = Vector{Matrix{eltype(Z[1])}}([similar(mat) for mat in Z])
    dW = Vector{Matrix{eltype(W[1])}}([similar(mat) for mat in W])
    dB = Vector{Matrix{eltype(B[1])}}([similar(mat) for mat in B])

    dlossFun = Symbol("d",lossFun)
    actFun = layers[L].actFun
    if ! isequal(actFun, :softmax) && !(actFun==:σ &&
                                        lossFun==:binaryCrossentropy)
        dA[L] = eval(:($dlossFun.($A[$L], $Y))) #.* eval(:($dActFun.(Z[L])))
    end
    if layers[L].keepProb < 1.0 #to save time of multiplication in case keepProb was one
        dA[L] = dA[L] .* D[L]
        dA[L] = dA[L] ./ layers[L].keepProb
    end
    for l=L:-1:2
        actFun = layers[l].actFun
        dActFun = Symbol("d",actFun)
        if l==L && (isequal(actFun, :softmax) ||
                    (actFun==:σ && lossFun==:binaryCrossentropy))
            dZ[l] = A[l] .- Y
        else
            dZ[l] = dA[l] .* eval(:($dActFun.($Z[$l])))
        end

        dW[l] = dZ[l]*A[l-1]' ./m

        if regulization > 0
            if regulization==1
                dW[l] .+= (λ/2m)
            else
                dW[l] .+= (λ/m) .* W[l]
            end
        end
        dB[l] = 1/m .* sum(dZ[l], dims=2)
        dA[l-1] = W[l]'dZ[l]
        if layers[l-1].keepProb < 1.0 #to save time of multiplication in case keepProb was one
            dA[l-1] = dA[l-1] .* D[l-1]
            dA[l-1] = dA[l-1] ./ layers[l-1].keepProb
        end
    end

    l=1 #shortcut cause I just copied from the for loop
    actFun = layers[l].actFun
    dActFun = Symbol("d",actFun)
    dZ[l] = dA[l] .* eval(:($dActFun.($Z[$l])))

    dW[l] = 1/m .* dZ[l]*X' #where X = A[0]

    if regulization > 0
        if regulization==1
            dW[l] .+= (λ/2m)
        else
            dW[l] .+= (λ/m) .* W[l]
        end
    end
    dB[l] = 1/m .* sum(dZ[l], dims=2)

    grads = Dict("dW"=>dW,
                 "dB"=>dB)
    return grads
end #backProp

function updateParams(model::Model, grads::Dict, tMiniBatch::Integer)
    W, B, V, S, layers = model.W, model.B, model.V, model.S, model.layers
    optimizer = model.optimizer
    dW, dB = grads["dW"], grads["dB"]
    L = length(layers)
    α = model.α
    β1, β2, ϵAdam = model.β1, model.β2, model.ϵAdam

    #initialize the needed variables to hold the corrected values
    #it is being init here cause these are not needed elsewhere
    VCorrected, SCorrected = deepInitVS(W,B,optimizer)
    if optimizer==:adam || optimizer==:momentum
        for i=1:L
            V[:dw][i] .= β1 .* V[:dw][i] .+ (1-β1) .* dW[i]
            V[:db][i] .= β1 .* V[:db][i] .+ (1-β1) .* dB[i]

            ##correcting
            VCorrected[:dw][i] .= V[:dw][i] ./ (1-β1^tMiniBatch)
            VCorrected[:db][i] .= V[:db][i] ./ (1-β1^tMiniBatch)

            if optimizer==:adam
                S[:dw][i] .= β2 .* S[:dw][i] .+ (1-β2) .* (dW[i].^2)
                S[:db][i] .= β2 .* S[:db][i] .+ (1-β2) .* (dB[i].^2)

                ##correcting
                SCorrected[:dw][i] .= S[:dw][i] ./ (1-β2^tMiniBatch)
                SCorrected[:db][i] .= S[:db][i] ./ (1-β2^tMiniBatch)

                ##update parameters with adam
                W[i] .-= (α .* (VCorrected[:dw][i] ./ (sqrt.(SCorrected[:dw][i]) .+ ϵAdam)))
                B[i] .-= (α .* (VCorrected[:db][i] ./ (sqrt.(SCorrected[:db][i]) .+ ϵAdam)))
            else#if optimizer==:momentum
                W[i] .-= (α .* VCorrected[:dw][i])
                B[i] .-= (α .* VCorrected[:db][i])
            end #if optimizer==:adam
        end #for i=1:L
    else
        W .-= (α .* dW)
        B .-= (α .* dB)
    end #if optimizer==:adam || optimizer==:momentum
    return W, B
end #updateParams

function updateParams!(model::Model, grads::Dict, tMiniBatch::Integer)
    W, B, V, S, layers = model.W, model.B, model.V, model.S, model.layers
    optimizer = model.optimizer
    dW, dB = grads["dW"], grads["dB"]
    L = length(layers)
    α = model.α
    β1, β2, ϵAdam = model.β1, model.β2, model.ϵAdam

    #initialize the needed variables to hold the corrected values
    #it is being init here cause these are not needed elsewhere
    VCorrected, SCorrected = deepInitVS(W,B,optimizer)
    if optimizer==:adam || optimizer==:momentum
        for i=1:L
            V[:dw][i] .= β1 .* V[:dw][i] .+ (1-β1) .* dW[i]
            V[:db][i] .= β1 .* V[:db][i] .+ (1-β1) .* dB[i]

            ##correcting
            VCorrected[:dw][i] .= V[:dw][i] ./ (1-β1^tMiniBatch)
            VCorrected[:db][i] .= V[:db][i] ./ (1-β1^tMiniBatch)

            if optimizer==:adam
                S[:dw][i] .= β2 .* S[:dw][i] .+ (1-β2) .* (dW[i].^2)
                S[:db][i] .= β2 .* S[:db][i] .+ (1-β2) .* (dB[i].^2)

                ##correcting
                SCorrected[:dw][i] .= S[:dw][i] ./ (1-β2^tMiniBatch)
                SCorrected[:db][i] .= S[:db][i] ./ (1-β2^tMiniBatch)

                ##update parameters with adam
                W[i] .-= (α .* (VCorrected[:dw][i] ./ (sqrt.(SCorrected[:dw][i]) .+ ϵAdam)))
                B[i] .-= (α .* (VCorrected[:db][i] ./ (sqrt.(SCorrected[:db][i]) .+ ϵAdam)))
            else#if optimizer==:momentum
                W[i] .-= (α .* VCorrected[:dw][i])
                B[i] .-= (α .* VCorrected[:db][i])
            end #if optimizer==:adam
        end #for i=1:L
    else
        W .-= (α .* dW)
        B .-= (α .* dB)
    end #if optimizer==:adam || optimizer==:momentum
    model.W, model.B = W, B
    return
end #updateParams!

function updateParams!(model::Model,
                       cLayer::Layer,
                       cnt::Integer = -1;
                       tMiniBatch::Integer = 1)

    optimizer = model.optimizer
    α = model.α
    β1, β2, ϵAdam = model.β1, model.β2, model.ϵAdam

    if cLayer.updateCount >= cnt
        return nothing
    end #if cLayer.updateCount >= cnt

    cLayer.updateCount += 1
    #initialize the needed variables to hold the corrected values
    #it is being init here cause these are not needed elsewhere
    VCorrected = Dict(:dw=>similar(cLayer.dW), :db=>similar(cLayer.dB))
    SCorrected = Dict(:dw=>similar(cLayer.dW), :db=>similar(cLayer.dB))
    if optimizer==:adam || optimizer==:momentum

        cLayer.V[:dw] .= β1 .* cLayer.V[:dw] .+ (1-β1) .* cLayer.dW
        cLayer.V[:db] .= β1 .* cLayer.V[:db] .+ (1-β1) .* cLayer.dB

        ##correcting
        VCorrected[:dw] .= cLayer.V[:dw] ./ (1-β1^tMiniBatch)
        VCorrected[:db] .= cLayer.V[:db] ./ (1-β1^tMiniBatch)

        if optimizer==:adam
            cLayer.S[:dw] .= β2 .* cLayer.S[:dw] .+ (1-β2) .* (cLayer.dW.^2)
            cLayer.S[:db] .= β2 .* cLayer.S[:db] .+ (1-β2) .* (cLayer.dB.^2)

            ##correcting
            SCorrected[:dw] .= cLayer.S[:dw] ./ (1-β2^tMiniBatch)
            SCorrected[:db] .= cLayer.S[:db] ./ (1-β2^tMiniBatch)

            ##update parameters with adam
            cLayer.W .-= (α .* (VCorrected[:dw] ./ (sqrt.(SCorrected[:dw]) .+ ϵAdam)))
            cLayer.B .-= (α .* (VCorrected[:db] ./ (sqrt.(SCorrected[:db]) .+ ϵAdam)))

        else#if optimizer==:momentum

            cLayer.W .-= (α .* VCorrected[:dw])
            cLayer.B .-= (α .* VCorrected[:db])

        end #if optimizer==:adam
    else
        cLayer.W .-= (α .* cLayer.dW)
        cLayer.B .-= (α .* cLayer.dB)
    end #if optimizer==:adam || optimizer==:momentum

    return nothing
end #updateParams!


text:
include("matrices.jl")
include("../utils.jl")
function NBMembedding(g; matrix = ihara_matrix)
    edgeidmap, m, aristas = mapa(g)
    NBM, edgeidmap = matrix(g)
    vect = eigvecs(NBM)
    the_values = eigvals(NBM)
    the_real = real.(the_values)
    the_imag = imag.(the_values)
    the_norm = sqrt.((the_real.^2)+(the_imag.^2))
    both = hcat(the_imag,the_norm)
    sorted = sortslices(both, dims=1,by=x->x[2],rev=true)
    threshold = sorted[findall(x->x≠0,sorted[:,1]),2][1]
    hm, index = num_com(the_values,threshold)
    vectores = eigvecs(NBM)
    index = index[1:length(index)]
    matrix_embedded = real(vectores[:,index])
    cont = contraction(g,index,matrix_embedded,edgeidmap)
    M = fit(PCA, cont'; maxoutdim=2)
    contracted = transform(M, cont')
    contracted
end


text:
function longest_consecutive(avail)
    # for [1 1 1 1 0 0 1 1]
    # return 4

    longest = 0
    current = 0

    for i in 1:length(avail)
        if avail[i] == 0
            current = 0
        else
            current+=1
        end

        if current > longest
            longest = current
        end
    end

    return longest
end

function anneal_availability(avail, min_length)
    # for [1 1 1 1 0 0 1 1]
    # if min_shift_length == 3
    # return [1 1 1 1 0 0 0 0]

    start_index = 0
    current_index = 0

    for i in 1:length(avail)
        if avail[i] == 1 && start_index == 0
            start_index = i
            current_index = i
        elseif (
            (avail[i] == 0 && start_index > 0)
            || (avail[i] == 1 && i == length(avail))
        ) # also note that this must come before other avail[i] == 1
            # if it's last element  - manually correct end index
            if i == length(avail)
                current_index = i
            end

            patch_length = current_index - start_index + 1
            # plus 1 because inclusive->inclusive

            if patch_length < min_length
                for j in start_index:current_index
                    avail[j] = 0
                end
            end
            start_index = 0
            current_index = 0
        elseif avail[i] == 1 && start_index > 0
            current_index += 1
        #elseif avail[i] == 0 && start_index == 0
        end
    end
    return avail
end

function get_day_bounds(longest_availability, days_assigned, hours_scheduled, week_min, week_max, shift_min)
    # note - longest_availability should only be for days remaining

    remaining_hours_max = week_max - hours_scheduled
    remaining_hours_min = week_min - hours_scheduled

    ## Day Min ##
    # every day but today that they're working
    future_capacity = sum(dot(longest_availability[2:end], days_assigned[2:end]))

    # Calculate day min based on future schedules
    day_min = remaining_hours_min - future_capacity

    # but ceiling it at shift_min from environment
    if day_min < shift_min
        day_min = shift_min
    end

    ## Day Max
    # Do not include today
    future_min_capacity = sum(days_assigned[2:end]) * shift_min
    day_max = remaining_hours_max - future_min_capacity

    # Floor at today's longest_availability
    if day_max > longest_availability[1]
        day_max = longest_availability[1]
    end

    return day_min, day_max
end

function get_longest_availability(employee, day)
        # find longest shift on that day
        longest = longest_consecutive(employee["availability"][day])
        if longest < employee["shift_time_min"]
            # Effectively unable to work that day at all
            longest = 0
        end
        if longest > employee["shift_time_max"]
            # Effectively unable to work that day at all
            longest = employee["shift_time_max"]
        end

        return longest
end

function sub_array(array, start_index, end_index)
    # Difference: start index can be > end_index
    if start_index > end_index
        return [array[start_index:end], array[1:end_index]]
    end
    return array[start_index:end_index]
end

function days_available_per_week(employee)

    count = 0
    for t in 1:length(employee["availability"])
        employee["availability"][t] = anneal_availability(employee["availability"][t], employee["shift_time_min"])
        # Already annealed
        # Assume preprocessing sets availablility to 0 when coverage 0
        if sum(employee["availability"][t]) > 0
            count += 1
        end
    end
    return count
end

function day_sum_coverage(env)
    return sum(env["coverage"])
end

function week_sum_coverage(env)
    return sum(sum(env["coverage"]))
end

function day_hours_scheduled(schedule)
    sum = 0
    for e in keys(schedule)
        sum += schedule[e]["length"]
    end
    return sum

end

function week_hours_scheduled(weekly_schedule)
    sum_hours = 0
    for e in keys(weekly_schedule)
        num_shifts = length(weekly_schedule[e])
        for i in 1:num_shifts
            shift_length = weekly_schedule[e][i]["length"]
            sum_hours += shift_length
        end
    end
    return sum_hours
end

max_stars_repo_name: shdann/autoshift

max_issues_repo_name: shdann/autoshift

max_forks_repo_name: shdann/autoshift


text:
module Bazinga

const Maybe{T} = Union{T,Nothing}

# compatibility
include("utilities/compat.jl")

# utilities
include("utilities/fbetools.jl")
include("utilities/lbfgs.jl")
include("utilities/restore.jl")
include("utilities/scaling.jl")

# algorithms
include("algorithms/zerofpr.jl")
include("algorithms/panoc.jl")
include("algorithms/alpx.jl")

# TODO pxal : proximal augmented lagrangian
#                     optimize   f(x) + g(x)   s.t.   c(x) ∈ S
#             by considering the problem
#                     optimize   f(x) + h(x)
#             where
#                     h(x) = g(x) + Ind[S](c(x))
#             and whose proximal operator
#                     prox_{a h}(u) = argmin_x { 2 a h(x) + ||x-u||^2 }
#             corresponds to
#                     optimize   q(x) + g(x)   s.t.   c(x) ∈ S
#             where
#                     q(x) = ||x-u||^2 / (2 a)
#             Thus, it can be evaluated by calling, e.g., `alpx`.

end # module


text:
mutable struct DirectMethod{S<:RatesCacheTrait} <: DecoupledSearchAlgorithm
  rates::Vector{Float64}
  total_rate::Float64
end

RatesCache(::DirectMethod{S}) where S = S()

@inline search_jumps(algorithm::DirectMethod) = search_jump_rates(RatesCache(algorithm), algorithm.rates, algorithm.total_rate)

@inline update!(algorithm::DirectMethod, state, model, j) = update_jump_rates!(algorithm, state, model)


text:
# Autogenerated wrapper script for squashfs_tools_jll for i686-linux-gnu
export mksquashfs, unsquashfs

using Zlib_jll
using XZ_jll
using LZO_jll
using Lz4_jll
using Zstd_jll
JLLWrappers.@generate_wrapper_header("squashfs_tools")
JLLWrappers.@declare_executable_product(mksquashfs)
JLLWrappers.@declare_executable_product(unsquashfs)
function __init__()
    JLLWrappers.@generate_init_header(Zlib_jll, XZ_jll, LZO_jll, Lz4_jll, Zstd_jll)
    JLLWrappers.@init_executable_product(
        mksquashfs,
        "bin/mksquashfs",
    )

    JLLWrappers.@init_executable_product(
        unsquashfs,
        "bin/unsquashfs",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
using Sparrow, DataFrames

d = DataFrame(x = [0:0.1:30], y = sin([0:0.1:30]), rc = sin([0:0.1:30])./2 .+ 0.5, gc = zeros(length([0:0.1:30])), bc = zeros(length([0:0.1:30])))
g = Sparrow.DataGroup(d)
l = Sparrow.LinePlot(g)
p = Sparrow.PlotFrame()
Sparrow.plot(p, l)
Sparrow.print(p, l, "line.pdf")

d2 = DataFrame(x = randn(100), y = randn(100), rc = rand(100), gc = rand(100), bc = rand(100))
g2 = Sparrow.DataGroup(d2)
l2 = Sparrow.ScatterPlot(g2)
p2 = Sparrow.PlotFrame()
Sparrow.plot(p2, l2)
Sparrow.print(p2, l2, "scatter.pdf")

p3 = Sparrow.PlotFrame()
Sparrow.plot(p3, [l, l2])
Sparrow.print(p3, [l, l2], "both.pdf")

readline()


text:
# Autogenerated wrapper script for libunistring_jll for x86_64-w64-mingw32
export libunistring

using Libiconv_jll
JLLWrappers.@generate_wrapper_header("libunistring")
JLLWrappers.@declare_library_product(libunistring, "libunistring-2.dll")
function __init__()
    JLLWrappers.@generate_init_header(Libiconv_jll)
    JLLWrappers.@init_library_product(
        libunistring,
        "bin\\libunistring-2.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
using Klimakoffer

NT = 48 # Number of time-steps per year

mesh = Mesh()

co2_concentration = 315.0 # in [ppm] from year 1950
model = Model(mesh, NT; co2_concentration = co2_concentration)

discretization = Discretization(mesh, model, NT)

GlobTemp = compute_equilibrium!(discretization)


text:
module FilePaths

import Base: ==
export
    # Types
    AbstractPath,
    Path,
    PosixPath,
    WindowsPath,
    Mode,
    Status,

    # Methods
    cwd,
    home,
    parts,
    root,
    hasparent,
    parents,
    filename,
    extension,
    extensions,
    exists,
    isabs,
    mode,
    created,
    modified,
    relative,
    glob,
    uri,
    move,
    remove,
    tmpname,
    tmpdir,
    mktmp,
    mktmpdir,
    chown,
    executable,
    readable,
    writable,
    raw,

    # Macros
    @p_str,

    # Constants
    READ,
    WRITE,
    EXEC

abstract AbstractPath

# The following should be implemented in the concrete types
Base.string(path::AbstractPath) = error("`string not implemented")
parts(path::AbstractPath) = error("`parts` not implemented.")
root(path::AbstractPath) = error("`root` not implemented.")

include("constants.jl")
include("libc.jl")
include("mode.jl")
include("status.jl")
include("posix.jl")
include("windows.jl")
include("path.jl")

end # end of module


text:
"""
    OneToRange

Supertype for `OneToSRange` and `OneToMRange`. It's subtypes should behave
identically to `Base.OneTo`.
"""
abstract type OneToRange{T<:Integer} <: AbstractUnitRange{T} end

Base.firstindex(::OneToRange) = 1

Base.lastindex(r::OneToRange) = Int(last(r))

Base.issubset(r::OneToRange, s::OneTo) = last(r) <= last(s)
Base.issubset(r::OneToRange, s::OneToRange) = last(r) <= last(s)
Base.issubset(r::OneTo, s::OneToRange) = last(r) <= last(s)

function Base.getindex(v::OneToRange{T}, i::Integer) where T
    Base.@_inline_meta
    @boundscheck ((i > 0) & (i <= last(v))) || throw(BoundsError(v, i))
    return convert(T, i)
end

function Base.getindex(r::OneToRange, s::Union{OneToRange,OneTo}) where T
    Base.@_inline_meta
    @boundscheck checkbounds(r, s)
    return typeof(r)(T(last(s)))
end

Base.mod(i::Integer, r::OneToRange) = Base.mod1(i, last(r))

"""
    OneToSRange

A static range that parallels `OneTo` in behavior.
"""
struct OneToSRange{T<:Integer,E} <: OneToRange{T}
    OneToSRange{T}(stop::T) where {T} = new{T,max(zero(T), stop)}()
end

function OneToSRange{T}(r::AbstractRange) where {T<:Integer}
    first(r) == 1 || (Base.@_noinline_meta; throw(ArgumentError("first element must be 1, got $(first(r))")))
    step(r)  == 1 || (Base.@_noinline_meta; throw(ArgumentError("step must be 1, got $(step(r))")))
    return OneToSRange(last(r))
end
OneToSRange{T,<:Any}(r::OneToSRange{T}) where {T<:Integer} = r
OneToSRange{T}(r::OneToRange) where {T<:Integer} = OneToSRange{T}(last(r))
OneToSRange(stop::T) where {T<:Integer} = OneToSRange{T}(stop)
OneToSRange{T}(stop) where {T} = OneToSRange{T}(T(stop))
OneToSRange(r::AbstractRange{T}) where {T<:Integer} = OneToSRange{T}(r)

function Base.getproperty(r::OneToSRange, s::Symbol)
    if s === :stop
        return last(r)
    else
        error("type $(typeof(r)) has no property $s")
    end
end

"""
    OneToMRange

A mutable range that parallels `OneTo` in behavior.
"""
mutable struct OneToMRange{T<:Integer} <: OneToRange{T}
    stop::T

    OneToMRange{T}(stop) where {T<:Integer} = new(max(zero(T), stop))
end

function OneToMRange{T}(r::AbstractRange) where {T<:Integer}
    first(r) == 1 || (Base.@_noinline_meta; throw(ArgumentError("first element must be 1, got $(first(r))")))
    step(r)  == 1 || (Base.@_noinline_meta; throw(ArgumentError("step must be 1, got $(step(r))")))
    return OneToMRange(last(r))
end
OneToMRange{T}(r::OneToMRange{T}) where {T<:Integer} = r
OneToMRange{T}(r::OneToRange) where {T<:Integer} = OneToMRange{T}(last(r))
OneToMRange(stop::T) where {T<:Integer} = OneToMRange{T}(stop)
OneToMRange(r::AbstractRange{T}) where {T<:Integer} = OneToMRange{T}(r)
#OneToMRange{T}(r::OneTo) where {T} = OneToMRange{T}(T(last(r)))

Base.AbstractUnitRange{T}(r::OneToSRange) where {T} = OneToSRange{T}(r)
Base.AbstractUnitRange{T}(r::OneToMRange) where {T} = OneToMRange{T}(r)


text:
include("defs.jl")

bd = Billiard(Antidot([0.0, 0.0], 1.0, false))

js = [(0.2, sind(30)), (1.5, sind(30 + sqrt(2)/10))]
colors = [to_color(COLORS[i+1]) for i in 1:length(js)]

ps = Particle{Float64}[]
for (ξ, s) in js
    x, y = from_bcoords(ξ, s, bd)
    push!(ps, Particle(x..., atan(y[2], y[1])))
end

fig, = interactive_billiard(
    bd, ps;
    colors, tail = 20000, backgroundcolor = RGBf0(1,1,1),
)

record_interaction(fig, joinpath(@__DIR__, "circle_billiard.mp4"))

max_stars_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_issues_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_forks_repo_name: JuliaDynamics/NonlinearDynamicsTextbook


text:
@testset "AssAlg" begin
  include("AssAlg/CSAMaxOrd.jl")
end


text:
#==============================================================================
    Code for solving the Hamiltonian Jacboi Bellman for
	   a problem with non-convexities as in Skiba 1978

	Translated Julia code from Matlab code by Ben Moll:
        http://www.princeton.edu/~moll/HACTproject.htm
==============================================================================#

using Plots, SparseArrays, LinearAlgebra

γ= 2.0 #gamma parameter for CRRA utility
ρ = 0.05 #the discount rate
α = 0.3 # the curvature of the production function (cobb-douglas)
δ = 0.05 # the depreciation rate
AH = 0.6 #high prodictivity, costs κ
AL = 0.4 #low prodictivity, free
κ = 2 #fixed cost

kssH = (α*AH/(ρ+δ))^(1/(1-α))+κ
k_st = κ./(1-(AL/AH).^(1/α))

# create the grid for k
H = 1000 #number of points on grid
k_min = 0.001*kssH # min value
k_max = 1.3*kssH # max value
k = LinRange(k_min, k_max, H)
dk = (k_max-k_min)/(H-1)

#Create the production function
yH = AH*max.(k.-κ,0).^α #high productivity
yL = AL*k.^α #low prodictivity
y = max.(yH,yL)

plot(k,y, ylabel="\$f(k)\$", xlabel="\$k\$",
		xlims=(k_min,k_max), ylims=(0,0.9),legend=false,color=:black)
plot!(k,yH, line=:dash, color=:red)
plot!(k,yL, line=:dashdot, color=:orange)

# use Ito's lemma to find the drift and variance of our optimization equation

maxit = 150 #needed to increase this for same results in julia
ε = 10^(-6)
Δ = 1000
# set up all of these empty matrices
dVf, dVb, dV0, c, If, Ib, I0, V= [zeros(H,1) for i in 1:8]

# Now it's time to solve the model, first put in a guess for the value function
v0 = (k.^α).^(1-γ)/(1-γ)/ρ
v=v0

dist = [] # set up empty array for the convergence criteria

for n = 1:maxit
    global V=v

    #Now set up the forward difference

    dVf[1:H-1] = (V[2:H] - V[1:H-1])/dk
    dVf[H] = (y[H]-δ.*k_max).^(-γ) # imposes a constraint

    #backward difference
    dVb[2:H] = (V[2:H] - V[1:H-1])/dk
    dVb[1] = (y[1]-δ.*k_min).^(-γ)

    #I_concave = Vab .> Vaf # indicator for whether the value function is concave

    # Consumption and savings functions
    cf = max.(dVf,10^(-10)).^(-1/γ)
	μf = y -δ.*k -cf
    Hf = cf.^(1-γ)/(1-γ) + dVf.*μf

    # consumption and saving backwards difference

	cb = max.(dVb,10^(-10)).^(-1/γ)
	μb = y -δ.*k -cb
    Hb = cb.^(1-γ)/(1-γ) + dVb.*μb

    #consumption and derivative of the value function at the steady state

    c0 = y- δ.*k
    dV0 = max.(c0, 10^(-10)).^(-γ)
	H0 = c0.^(1-γ)/(1-γ)

    # Upwind scheme chooses between the forward or backward difference
	Ieither = (1 .- (μf.>0)) .* (1 .- (μb.<0))
	Iunique = (μb.<0).*(1 .- (μf.>0)) + (1 .- (μb.<0)).*(μf.>0)
	Iboth = (μb.<0).*(μf.>0)
    If= Iunique.*(μf .> 0) + Iboth.*(Hf.>=Hb)  #positive drift → forward difference
    Ib = Iunique.*(μb .< 0) + Iboth.*(Hb.>=Hf) #negative drift → backward difference
    I0 = Ieither  #at steady state

    global c = cf.*If + cb.*Ib + c0.*I0
    u = (c.^(1-γ))/(1-γ)

	# CONSTRUCT MATRIX
    global X = -μb.*Ib/dk
    global Y = - μf.*If/dk  + μb.*Ib/dk
    global Z = μf.*If/dk

	global A = spdiagm(0=>Y[:], -1=>X[2:end], 1=>Z[1:end-1])

  	B = (1/Δ + ρ)*sparse(I, H, H) - A

  	u_stacked= reshape(u, H, 1)
  	V_stacked = reshape(V,H, 1)

  	b = u_stacked + (V_stacked/Δ)

  	V_stacked = B\b

  	global V = reshape(V_stacked, H, 1)

  	V_change = V-v

  	global v= V

  	# need push function to add to an already existing array
  	push!(dist, findmax(abs.(V_change))[1])
  	if dist[n].< ε
      	println("Value Function Converged Iteration=")
      	println(n)
      	break
  	end

end

# calculate the savings for kk
kdot = y -δ.*k -c
dV_upwind = dVf.*If + dVb.*Ib+dV0.*I0
Verr = c.^(1-γ)/(1-γ) + dV_upwind.*kdot - ρ.*V

#plot for consumption
plot(k, c, ylabel="\$c(k)\$", xlabel="\$k\$", label="Consumption")
plot!(k,y-δ.*k, label="Production net of depreciation", legend=:bottomright)

#plot for savings
plot(k, kdot, ylabel="\$s(k)\$", xlabel="\$k\$", legend=false)
plot!(k,zeros(H,1), line=:dash)

#Plot the value function
plot(k, V,ylabel="\$V(k)\$", xlabel="\$k\$", legend=false)

#Plot the error in HJB equation
plot(k,Verr,ylabel="Value Function error", xlabel="\$k\$", legend=false)


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

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 920fb510-0ff1-11eb-0397-7524b4e260df
md"""
## Arrays Basics
-https://www.youtube.com/watch?v=CwDI-YOjWhc
"""

# ╔═╡ a2689530-0ff1-11eb-237b-2191e0d0e553
begin
	import Pkg
	Pkg.activate(mktempdir())
end

# ╔═╡ a23ddbb0-0ff1-11eb-13ed-39ac3383e3f7
begin
	Pkg.add(["Images", "ImageIO", "ImageMagick", "TestImages"])
	using Images
end

# ╔═╡ a22300b0-0ff1-11eb-0d31-7d1ee193c63c
	md"""Vector and array
- Defining, Type of and Size"""

# ╔═╡ a5f899c0-0ff1-11eb-33a9-c136cfa99b88
v = [1,2,3,4,5,6] # this is vector

# ╔═╡ 8b4e34c0-0ff3-11eb-1bcb-532ae49a59ed
typeof(v)

# ╔═╡ ac6b39f0-0ff3-11eb-19f3-dd8175b61e78
size(v)

# ╔═╡ 9369dbf2-0ff3-11eb-332f-cf47ac906f54
A = [4 5 6
	 9 8 3
     9 2 6]

# ╔═╡ 932e3280-0ff3-11eb-0ca8-596fcd6efe47
typeof(A)

# ╔═╡ cc5cb770-0ff3-11eb-3e5f-67172078918e
size(A)

# ╔═╡ a5e38b1e-0ff1-11eb-2d07-594e1017073f
md"""Below demonstrating selecint an emlement, range, row or column"""

# ╔═╡ a5cecaa0-0ff1-11eb-3b69-c9888018ec53
A[5] # picks element on a position

# ╔═╡ e7bdbae0-0ff4-11eb-2642-e193e1676378
A[:,2] # second column, all element

# ╔═╡ e75d9980-0ff4-11eb-0db5-29f2c9f220a4
A[3,:] # third row, all element

# ╔═╡ e70feeb0-0ff4-11eb-201a-1bce59f47b9b
A[2:3, :] # small matrix

# ╔═╡ 80066900-0ff5-11eb-1320-89d3e6476de4
A[:, 2:3] # small matrix

# ╔═╡ a5b96de2-0ff1-11eb-0e90-51478ea0c709
md"""
## Colors"""

# ╔═╡ a590b030-0ff1-11eb-1bc9-89f29eac464c
md"""Distinguishable Colors - provide the number to increase the radom"""

# ╔═╡ a57c64e0-0ff1-11eb-0d67-7390aed58a41
C1 = distinguishable_colors(3)

# ╔═╡ 2e4530f0-0ff6-11eb-283a-5fdcd7e53ef8
C2 = distinguishable_colors(7)

# ╔═╡ 2dfbcbe2-0ff6-11eb-37fe-fbf8dc60a211
C3 = distinguishable_colors(7)

# ╔═╡ a5672f30-0ff1-11eb-135b-4d6522f840ed
md"""Creating a Matrix of random Colors"""

# ╔═╡ a551ab60-0ff1-11eb-1f74-c36932fccd40
Z1 = rand(C2, 5,5)

# ╔═╡ 6c0ffa50-0ff6-11eb-0268-97c6f27d87f3
Z2 = rand(C3,5,5)

# ╔═╡ a5282a60-0ff1-11eb-3c52-97f5dc063a5f
md"""bind Slider"""

# ╔═╡ a51369e0-0ff1-11eb-35da-e5462d420967
@bind x html"<input type=range min=1 max=150>"

# ╔═╡ 50ace920-0ff7-11eb-205b-118b6c1d7e79
Z3 = rand(C3,x,x)

# ╔═╡ a5014170-0ff1-11eb-2bb3-a71e24fc2432
md"""Creating a Comprehension Table Matrix left, creating I and J loops, to create a matrix, double llop, watch the video in description at 3:35"""

# ╔═╡ a4d79960-0ff1-11eb-1016-f5306d2483e1
md"""Square of Matrix and Its Element Square"""

# ╔═╡ ca2232f0-1002-11eb-0969-eb0f53652baf
B = A

# ╔═╡ c9d09080-1002-11eb-387b-030058975f47
B^2

# ╔═╡ eec18110-1002-11eb-0713-45098e649683
md"""When ypu put a dot then its a element by element square as below:"""

# ╔═╡ e79c95f0-1002-11eb-2b73-dbbc15682a95
B.^2

# ╔═╡ a4953930-0ff1-11eb-05c5-7deb20dc7af1
md"""Concatenate matrix, vertical and horizontal as well, use semi column, thats the trick"""

# ╔═╡ b4889100-1002-11eb-272c-e721ee0984eb
 H1 = [Z1 ; Z2]

# ╔═╡ 7ac77390-1003-11eb-38ab-974ac8710b4f
H2 = [Z2 ; Z1]

# ╔═╡ c36d0420-1003-11eb-02b4-5778956fcd2c
Notice below and above, Verical concatenate happens with semi column and horizontal happens without semi columns

# ╔═╡ 8bfb2350-1003-11eb-05a1-8335c973c679
[H2  H1]

# ╔═╡ b4764180-1002-11eb-1512-0f80b099d65e
md"""--"""

# ╔═╡ b46a3390-1002-11eb-30b6-eb20022bcaa6

# ╔═╡ b4480590-1002-11eb-2d69-67d2be3a2f59
md"""--"""

# ╔═╡ b431975e-1002-11eb-2e91-3f5c9069231b

# ╔═╡ b3fdb620-1002-11eb-23d7-5d5220741f9d
md"""--"""

# ╔═╡ b3cfef60-1002-11eb-0962-6fa39009633f

# ╔═╡ Cell order:
# ╟─920fb510-0ff1-11eb-0397-7524b4e260df
# ╠═a2689530-0ff1-11eb-237b-2191e0d0e553
# ╠═a23ddbb0-0ff1-11eb-13ed-39ac3383e3f7
# ╟─a22300b0-0ff1-11eb-0d31-7d1ee193c63c
# ╠═a5f899c0-0ff1-11eb-33a9-c136cfa99b88
# ╠═8b4e34c0-0ff3-11eb-1bcb-532ae49a59ed
# ╠═ac6b39f0-0ff3-11eb-19f3-dd8175b61e78
# ╠═9369dbf2-0ff3-11eb-332f-cf47ac906f54
# ╠═932e3280-0ff3-11eb-0ca8-596fcd6efe47
# ╠═cc5cb770-0ff3-11eb-3e5f-67172078918e
# ╟─a5e38b1e-0ff1-11eb-2d07-594e1017073f
# ╠═a5cecaa0-0ff1-11eb-3b69-c9888018ec53
# ╠═e7bdbae0-0ff4-11eb-2642-e193e1676378
# ╠═e75d9980-0ff4-11eb-0db5-29f2c9f220a4
# ╠═e70feeb0-0ff4-11eb-201a-1bce59f47b9b
# ╠═80066900-0ff5-11eb-1320-89d3e6476de4
# ╟─a5b96de2-0ff1-11eb-0e90-51478ea0c709
# ╠═a590b030-0ff1-11eb-1bc9-89f29eac464c
# ╠═a57c64e0-0ff1-11eb-0d67-7390aed58a41
# ╠═2e4530f0-0ff6-11eb-283a-5fdcd7e53ef8
# ╠═2dfbcbe2-0ff6-11eb-37fe-fbf8dc60a211
# ╟─a5672f30-0ff1-11eb-135b-4d6522f840ed
# ╠═a551ab60-0ff1-11eb-1f74-c36932fccd40
# ╠═6c0ffa50-0ff6-11eb-0268-97c6f27d87f3
# ╟─a5282a60-0ff1-11eb-3c52-97f5dc063a5f
# ╠═a51369e0-0ff1-11eb-35da-e5462d420967
# ╠═50ace920-0ff7-11eb-205b-118b6c1d7e79
# ╟─a5014170-0ff1-11eb-2bb3-a71e24fc2432
# ╠═a4d79960-0ff1-11eb-1016-f5306d2483e1
# ╠═ca2232f0-1002-11eb-0969-eb0f53652baf
# ╠═c9d09080-1002-11eb-387b-030058975f47
# ╟─eec18110-1002-11eb-0713-45098e649683
# ╠═e79c95f0-1002-11eb-2b73-dbbc15682a95
# ╟─a4953930-0ff1-11eb-05c5-7deb20dc7af1
# ╠═b4889100-1002-11eb-272c-e721ee0984eb
# ╠═7ac77390-1003-11eb-38ab-974ac8710b4f
# ╠═c36d0420-1003-11eb-02b4-5778956fcd2c
# ╠═8bfb2350-1003-11eb-05a1-8335c973c679
# ╠═b4764180-1002-11eb-1512-0f80b099d65e
# ╠═b46a3390-1002-11eb-30b6-eb20022bcaa6
# ╠═b4480590-1002-11eb-2d69-67d2be3a2f59
# ╠═b431975e-1002-11eb-2e91-3f5c9069231b
# ╠═b3fdb620-1002-11eb-23d7-5d5220741f9d
# ╠═b3cfef60-1002-11eb-0962-6fa39009633f


text:
#testing
#simulates an artificial portfolio and reweights it iter times
function testPortfolio(;N=560, M=5000, iter=100)
  cols::Vector{Symbol} = (n->Symbol("X",n)).(1:M)
  crosscor::MFloat64 = 0.25
  stdev::MFloat64 = 0.1
  allmean = .02

  ρ::Matrix{MFloat64} = ones(M,M) .* crosscor
  ρ[diagind(ρ)] .= 1.0 #fix the diagonals
  σ::Vector{MFloat64} = (n->stdev).(1:M)
  Σ::Matrix{MFloat64} = ρ .* stdev^2
  L = cholesky(Σ).L

  #make the test matrix
  testmatrix::Matrix{MFloat64} = rand(Normal(), N, M)

  for r ∈ 1:N
    testmatrix[r,:] .= L * testmatrix[r,:] .+ allmean
  end

  #make missing values
  for c ∈ 1:M
    for r ∈ 1:N
      rand([true, false]) && (testmatrix[r,c] = missing)
    end
  end

  #make the dataframe
  df = DataFrame(testmatrix)
  names!(df, cols)
  df[:date] = Vector{MDate}(missing, N)

  datebegin::Date = Date(2011,11,11)

  df[:date] .= datebegin .+ (Day).(1:N)

  U = Universe(df)
  wₜ = ones(M) ./ M

  P = Portfolio(U,wₜ)

  μs::Vector{Float64} = Vector{Float64}(undef, iter)

  @time for i ∈ 1:iter
    wₜ .= rand(M)
    Portfolio!(P, wₜ)
    μs[i] = P.pμ[]
  end

  println("Sum of means: $(sum(μs))")

  return nothing
end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
using ScientificTest, Test

include("z.jl")


text:
struct Bind{T,S} <: DynamicIterator
    Y::S
    P::T
end
bind(Y,P::DynamicIterator) = Bind(Y,P)

function dyniterate(M::Bind, start::Start)
    v, q = @returnnothing dyniterate(M.Y, nothing)
    u, p = @returnnothing dyniterate(M.P, start, v)
    u, (q, p)
end
function dyniterate(M::Bind, ::Nothing)
    v, q = @returnnothing dyniterate(M.Y, nothing)
    u, p = @returnnothing dyniterate(M.P, nothing, v)
    u, (q, p)
end
function dyniterate(M::Bind, (q, p)::Tuple)
    v, q = @returnnothing iterate(M.Y, q)
    u, p = @returnnothing dyniterate(M.P, p, v)
    u, (q, p)
end

struct BindAfter{T,S} <: DynamicIterator
    Y::S
    P::T
end
bindafter(Y,P::DynamicIterator) = BindAfter(Y,P)

function dyniterate(M::BindAfter, start::Start)
    v, q = @returnnothing dyniterate(M.Y, nothing)
    u, p = @returnnothing dyniterate(M.P, start)
    u, (v, q, p)
end
function dyniterate(M::BindAfter, ::Nothing)
    v, q = @returnnothing dyniterate(M.Y, nothing)
    u, p = @returnnothing dyniterate(M.P, nothing)
    u, (q, p)
end
function dyniterate(M::BindAfter, (v, q, p)::Tuple)
    w, q = @returnnothing iterate(M.Y, q)
    u, p = @returnnothing dyniterate(M.P, p, v)
    u, (w, q, p)
end

"""

    mix(f, P, Q)

Mix two dynamic iterators by applying the mixing function `f`
to their states:

    x, y = f(x, y)

## Example
```
collectfrom(Mix((x,y) -> (x+y, y), 1:0, 1:100), (1,1)))
# last value 100*101/2 + 100
```
"""
struct Mix{F,T,S} <: DynamicIterator
    f::F
    P::T
    Q::S
end

dyniterate(M::Mix{<:Any, <:GEvolution, <:GEvolution}, u) = dub(evolve(M, u))
dyniterate(M::Mix{<:Any, <:GEvolution, <:GEvolution}, u::Value) = dub(evolve(M, u.value))
dyniterate(M::Mix{<:Any, <:GEvolution, <:GEvolution}, u::Start) = dub(evolve(M, u.value))
evolve(M::Mix{<:Any, <:GEvolution, <:GEvolution}, (i, pq)::Pair) = i+1 => evolve(M, pq)

function evolve(M::Mix{<:Any, <:GEvolution, <:GEvolution}, (p, q)::Tuple)
    p = evolve(M.P, p)
    p === nothing && return nothing
    q = evolve(M.Q, q)
    q === nothing && return nothing
    M.f(p, q)
end

function dyniterate(M::Mix, (value,)::Start)
    x, y = value
    ϕ = dyniterate(M.P, Start(x))
    ϕ === nothing && return nothing
    x, p = ϕ
    ψ = dyniterate(M.Q, Start(y))
    ψ === nothing && return nothing
    y, q = ψ
    x, y = M.f(x, y)
    (x, y), (p, q)
end
function dyniterate(M::Mix, (value, u)::Value)
    p, q = u
    x, y = value
    ϕ = dyniterate(M.P, Value(x, p))
    ϕ === nothing && return nothing
    x, p = ϕ
    ψ = dyniterate(M.Q, Value(y, q))
    ψ === nothing && return nothing
    y, q = ψ
    x, y = M.f(x, y)
    (x, y), (p, q)
end

mix(f, P, Q) = Mix(f, P, Q)

"""

    mixture(I, Ps)

    evolve(M::Mixture, (i, x))
Choose evolution in `Ps[i]` for `x` using iterate `i` of `I`.
"""
struct Mixture{S,T} <: Evolution
    I::S
    Ps::T
end

evolve(M::Mixture, (i, ix)::Pair) = i+1 => evolve(M, ix)

function evolve(M::Mixture, (i, x)::Tuple)
    i = evolve(M.I, i)
    x = evolve(M.Ps[i], x)
    (i, x)
end

mixture(I, args) = Mixture(I, args)

struct Synchronize{T} <: Evolution
    Ps::T
end

"""
    synchronize
"""
synchronize(args...) = Synchronize(args)

state((t, xs), M::Synchronize) = (t => xs, Tuple(evolve.(M.Ps, Pair.(t, xs))))

function evolve(M::Synchronize, ((t, x), next)::T) where {T}
    all(u === nothing for u in next) && return nothing
    tᵒ = minimum(first.(Iterators.filter(x -> !(t===nothing), next))) #use isnothing
    xᵒ = Any[]
    nextᵒ = Any[]
    for  (P, xᵢ, u) in zip(M.Ps, x, next)
        if !(u === nothing) && u[1] == tᵒ
            xᵢ = u[2]
            u = evolve(P, u)
        end
        push!(xᵒ, xᵢ)
        push!(nextᵒ, u)
    end
    (tᵒ => Tuple(xᵒ), Tuple(nextᵒ))::T
end


text:
using Test
using CLIMA.MoistThermodynamics
using NCDatasets
using Random
MT = MoistThermodynamics
using CLIMA.PlanetParameters
using LinearAlgebra

float_types = [Float32, Float64]

include("testdata.jl")

@testset "moist thermodynamics - isentropic processes" begin
  for FT in [Float64]
  # for FT in float_types
    e_int, ρ, q_tot, q_pt, T, p, θ_liq_ice = MT.tested_convergence_range(FT, 50)
    Φ = FT(1)
    Random.seed!(15);
    perturbation = FT(0.1)*rand(length(T))

    # TODO: Use reasonable values for ambient temperature/pressure
    T∞, p∞ = T .* perturbation, p .* perturbation
    @test air_temperature.(p, θ_liq_ice, Ref(DryAdiabaticProcess())) ≈ (p ./ FT(MSLP)) .^ (FT(R_d) / FT(cp_d)) .* θ_liq_ice
    @test air_pressure_given_θ.(θ_liq_ice, Φ, Ref(DryAdiabaticProcess())) ≈ FT(MSLP) .* (1 .- Φ ./ (θ_liq_ice .* FT(cp_d))) .^ (FT(cp_d) / FT(R_d))
    @test air_pressure.(T, T∞, p∞, Ref(DryAdiabaticProcess())) ≈ p∞ .* (T ./ T∞) .^ (FT(1) / FT(kappa_d))
  end
end

@testset "moist thermodynamics - correctness" begin
  FT = Float64
  # ideal gas law
  @test air_pressure(FT(1), FT(1), PhasePartition(FT(1))) === FT(R_v)
  @test air_pressure(FT(1), FT(2), PhasePartition(FT(1), FT(0.5), FT(0))) === FT(R_v)
  @test air_pressure(FT(1), FT(1)) === FT(R_d)
  @test air_pressure(FT(1), FT(2)) === 2*FT(R_d)
  @test air_density(FT(1), FT(1)) === 1/FT(R_d)
  @test air_density(FT(1), FT(2)) === 2/FT(R_d)

  # gas constants and heat capacities
  @test gas_constant_air(PhasePartition(FT(0))) === FT(R_d)
  @test gas_constant_air(PhasePartition(FT(1))) === FT(R_v)
  @test gas_constant_air(PhasePartition(FT(0.5), FT(0.5))) ≈ FT(R_d)/2
  @test gas_constant_air(FT) == FT(R_d)

  @test cp_m(PhasePartition(FT(0))) === FT(cp_d)
  @test cp_m(PhasePartition(FT(1))) === FT(cp_v)
  @test cp_m(PhasePartition(FT(1), FT(1))) === FT(cp_l)
  @test cp_m(PhasePartition(FT(1), FT(0), FT(1))) === FT(cp_i)
  @test cp_m(FT) == FT(cp_d)

  @test cv_m(PhasePartition(FT(0))) === FT(cp_d - R_d)
  @test cv_m(PhasePartition(FT(1))) === FT(cp_v - R_v)
  @test cv_m(PhasePartition(FT(1), FT(1))) === FT(cv_l)
  @test cv_m(PhasePartition(FT(1), FT(0), FT(1))) === FT(cv_i)
  @test cv_m(FT) == FT(cv_d)

  # speed of sound
  @test soundspeed_air(T_0 + 20, PhasePartition(FT(0))) == sqrt(cp_d/cv_d * R_d * (T_0 + 20))
  @test soundspeed_air(T_0 + 100, PhasePartition(FT(1))) == sqrt(cp_v/cv_v * R_v * (T_0 + 100))

  # specific latent heats
  @test latent_heat_vapor(FT(T_0))  ≈ LH_v0
  @test latent_heat_fusion(FT(T_0)) ≈ LH_f0
  @test latent_heat_sublim(FT(T_0)) ≈ LH_s0

  # saturation vapor pressure and specific humidity
  p=FT(1.e5); q_tot=FT(0.23); ρ=FT(1.);
  ρ_v_triple = press_triple / R_v / T_triple;
  @test saturation_vapor_pressure(FT(T_triple), Liquid()) ≈ press_triple
  @test saturation_vapor_pressure(FT(T_triple), Ice()) ≈ press_triple

  @test q_vap_saturation(FT(T_triple), ρ, PhasePartition(FT(0))) == ρ_v_triple / ρ
  @test q_vap_saturation(FT(T_triple), ρ, PhasePartition(q_tot,q_tot)) == ρ_v_triple / ρ

  @test q_vap_saturation_generic(FT(T_triple), ρ; phase=Liquid()) == ρ_v_triple / ρ
  @test q_vap_saturation_generic(FT(T_triple), ρ; phase=Ice()) == ρ_v_triple / ρ
  @test q_vap_saturation_generic.(FT(T_triple-20), ρ; phase=Liquid()) >=
        q_vap_saturation_generic.(FT(T_triple-20), ρ; phase=Ice())

  @test saturation_excess(FT(T_triple), ρ, PhasePartition(q_tot)) ==  q_tot - ρ_v_triple/ρ
  @test saturation_excess(FT(T_triple), ρ, PhasePartition(q_tot/1000)) == 0.0

  # energy functions and inverse (temperature)
  T=FT(300); e_kin=FT(11); e_pot=FT(13);
  @test air_temperature(FT(cv_d*(T-T_0))) === FT(T)
  @test air_temperature(FT(cv_d*(T-T_0)), PhasePartition(FT(0))) === FT(T)

  @test air_temperature(cv_m(PhasePartition(FT(0)))*(T-T_0), PhasePartition(FT(0))) === FT(T)
  @test air_temperature(cv_m(PhasePartition(FT(q_tot)))*(T-T_0) + q_tot*e_int_v0, PhasePartition(q_tot)) ≈ FT(T)

  @test total_energy(FT(e_kin),FT(e_pot), FT(T_0)) === FT(e_kin) + FT(e_pot)
  @test total_energy(FT(e_kin),FT(e_pot), FT(T)) ≈ FT(e_kin) + FT(e_pot) + cv_d*(T-T_0)
  @test total_energy(FT(0),FT(0), FT(T_0), PhasePartition(q_tot)) ≈ q_tot*e_int_v0

  # phase partitioning in equilibrium
  q_liq = FT(0.1);
  T = FT(T_icenuc-10); ρ = FT(1.0); q_tot = FT(0.21)
  @test liquid_fraction(T) === FT(0)
  @test liquid_fraction(T, PhasePartition(q_tot,q_liq,q_liq)) === FT(0.5)
  q = PhasePartition_equil(T, ρ, q_tot)
  @test q.liq ≈ FT(0)
  @test 0 < q.ice <= q_tot

  T = FT(T_freeze+10); ρ = FT(0.1); q_tot = FT(0.60)
  @test liquid_fraction(T) === FT(1)
  @test liquid_fraction(T, PhasePartition(q_tot,q_liq,q_liq/2)) === FT(2/3)
  q = PhasePartition_equil(T, ρ, q_tot)
  @test 0 < q.liq <= q_tot
  @test q.ice ≈ 0

  # saturation adjustment in equilibrium (i.e., given the thermodynamic
  # variables E_int, p, q_tot, compute the temperature and partitioning of the phases
  tol_T = 1e-1
  q_tot = FT(0); ρ = FT(1)
  @test MT.saturation_adjustment_SecantMethod(internal_energy_sat(300.0, ρ, q_tot), ρ, q_tot, 1e-2, 10) ≈ 300.0
  @test abs(MT.saturation_adjustment(internal_energy_sat(300.0, ρ, q_tot), ρ, q_tot, 1e-2, 10) - 300.0) < tol_T

  q_tot = FT(0.21); ρ = FT(0.1)
  @test MT.saturation_adjustment_SecantMethod(internal_energy_sat(200.0, ρ, q_tot), ρ, q_tot, 1e-2, 10) ≈ 200.0
  @test abs(MT.saturation_adjustment(internal_energy_sat(200.0, ρ, q_tot), ρ, q_tot, 1e-2, 10) - 200.0) < tol_T
  q = PhasePartition_equil(T, ρ, q_tot)
  @test q.tot - q.liq - q.ice ≈ vapor_specific_humidity(q) ≈ q_vap_saturation(T, ρ)

  ρ = FT(1)
  ρu = FT[1,2,3]
  ρe = FT(1100)
  e_pot = FT(93)
  @test internal_energy(ρ, ρe, ρu, e_pot) ≈ 1000.0

  # potential temperatures
  T = FT(300)
  @test liquid_ice_pottemp_given_pressure(T, FT(MSLP)) === T
  @test liquid_ice_pottemp_given_pressure(T, FT(MSLP)/10) ≈ T * 10^(R_d/cp_d)
  @test liquid_ice_pottemp_given_pressure(T, FT(MSLP)/10, PhasePartition(FT(1))) ≈ T * 10^(R_v/cp_v)

  # dry potential temperatures. FIXME: add correctness tests
  T = FT(300); p=FT(1.e5); q_tot=FT(0.23)
  @test dry_pottemp_given_pressure(T, p, PhasePartition(q_tot)) isa typeof(p)
  @test air_temperature_from_liquid_ice_pottemp_given_pressure(
    dry_pottemp_given_pressure(T, p, PhasePartition(q_tot)), p, PhasePartition(q_tot)) ≈ T

  # Exner function. FIXME: add correctness tests
  p=FT(1.e5); q_tot=FT(0.23)
  @test exner_given_pressure(p, PhasePartition(q_tot)) isa typeof(p)

  e_int, ρ, q_tot, q_pt, T, p, θ_liq_ice = MT.tested_convergence_range(FT, 50)
  ts = PhaseEquil.(e_int, ρ, q_tot)

  data_folder = data_folder_moist_thermo()
  ds_PhaseEquil = Dataset(joinpath(data_folder, "test_data_PhaseEquil.nc"), "r")
  e_int = Array{FT}(ds_PhaseEquil["e_int"][:])
  ρ     = Array{FT}(ds_PhaseEquil["ρ"][:])
  q_tot = Array{FT}(ds_PhaseEquil["q_tot"][:])

  # ts = PhaseEquil.(e_int, ρ, q_tot) # Fails
end

@testset "moist thermodynamics - default behavior accuracy" begin
  # Input arguments should be accurate within machine precision
  # Temperature is approximated via saturation adjustment, and should be within a physical tolerance

  for FT in float_types
    rtol = FT(1e-2)
    e_int, ρ, q_tot, q_pt, T, p, θ_liq_ice = MT.tested_convergence_range(FT, 50)

    # PhaseEquil
    ts_exact = PhaseEquil.(e_int, ρ, q_tot, 100, FT(1e-4))
    ts = PhaseEquil.(e_int, ρ, q_tot)
    # Should be machine accurate (because ts contains `e_int`,`ρ`,`q_tot`):
    @test all(getproperty.(PhasePartition.(ts),:tot) .≈ getproperty.(PhasePartition.(ts_exact),:tot))
    @test all(internal_energy.(ts) .≈ internal_energy.(ts_exact))
    @test all(air_density.(ts) .≈ air_density.(ts_exact))
    # Approximate (temperature must be computed via saturation adjustment):
    @test all(isapprox.(air_pressure.(ts), air_pressure.(ts_exact), rtol=rtol))
    @test all(isapprox.(air_temperature.(ts), air_temperature.(ts_exact), rtol=rtol))

    # PhaseEquil
    ts_exact = PhaseEquil.(e_int, ρ, q_tot, 100, FT(1e-4), MT.saturation_adjustment_SecantMethod)
    ts = PhaseEquil.(e_int, ρ, q_tot, 30, FT(1e-1), MT.saturation_adjustment_SecantMethod) # Needs to be in sync with default
    # Should be machine accurate (because ts contains `e_int`,`ρ`,`q_tot`):
    @test all(getproperty.(PhasePartition.(ts),:tot) .≈ getproperty.(PhasePartition.(ts_exact),:tot))
    @test all(internal_energy.(ts) .≈ internal_energy.(ts_exact))
    @test all(air_density.(ts) .≈ air_density.(ts_exact))
    # Approximate (temperature must be computed via saturation adjustment):
    @test all(isapprox.(air_pressure.(ts), air_pressure.(ts_exact), rtol=rtol))
    @test all(isapprox.(air_temperature.(ts), air_temperature.(ts_exact), rtol=rtol))

    # LiquidIcePotTempSHumEquil
    ts_exact = LiquidIcePotTempSHumEquil.(θ_liq_ice, ρ, q_tot, 40, FT(1e-3))
    ts = LiquidIcePotTempSHumEquil.(θ_liq_ice, ρ, q_tot)
    # Should be machine accurate:
    @test all(air_density.(ts) .≈ air_density.(ts_exact))
    @test all(getproperty.(PhasePartition.(ts),:tot) .≈ getproperty.(PhasePartition.(ts_exact),:tot))
    # Approximate (temperature must be computed via saturation adjustment):
    @test all(isapprox.(internal_energy.(ts), internal_energy.(ts_exact), rtol=rtol))
    @test all(isapprox.(liquid_ice_pottemp.(ts), liquid_ice_pottemp.(ts_exact), rtol=rtol))
    @test all(isapprox.(air_temperature.(ts), air_temperature.(ts_exact), rtol=rtol))

    # LiquidIcePotTempSHumEquil_given_pressure
    ts_exact = LiquidIcePotTempSHumEquil_given_pressure.(θ_liq_ice, p, q_tot, 40, FT(1e-3))
    ts = LiquidIcePotTempSHumEquil_given_pressure.(θ_liq_ice, p, q_tot)
    # Should be machine accurate:
    @test all(getproperty.(PhasePartition.(ts),:tot) .≈ getproperty.(PhasePartition.(ts_exact),:tot))
    # Approximate (temperature must be computed via saturation adjustment):
    @test all(isapprox.(air_density.(ts), air_density.(ts_exact), rtol=rtol))
    @test all(isapprox.(internal_energy.(ts), internal_energy.(ts_exact), rtol=rtol))
    @test all(isapprox.(liquid_ice_pottemp.(ts), liquid_ice_pottemp.(ts_exact), rtol=rtol))
    @test all(isapprox.(air_temperature.(ts), air_temperature.(ts_exact), rtol=rtol))

    # LiquidIcePotTempSHumNonEquil
    ts_exact = LiquidIcePotTempSHumNonEquil.(θ_liq_ice, ρ, q_pt, 40, FT(1e-3))
    ts = LiquidIcePotTempSHumNonEquil.(θ_liq_ice, ρ, q_pt)
    # Should be machine accurate:
    @test all(getproperty.(PhasePartition.(ts),:tot) .≈ getproperty.(PhasePartition.(ts_exact),:tot))
    @test all(getproperty.(PhasePartition.(ts),:liq) .≈ getproperty.(PhasePartition.(ts_exact),:liq))
    @test all(getproperty.(PhasePartition.(ts),:ice) .≈ getproperty.(PhasePartition.(ts_exact),:ice))
    @test all(air_density.(ts) .≈ air_density.(ts_exact))
    # Approximate (temperature must be computed via non-linear solve):
    @test all(isapprox.(internal_energy.(ts), internal_energy.(ts_exact), rtol=rtol))
    @test all(isapprox.(liquid_ice_pottemp.(ts), liquid_ice_pottemp.(ts_exact), rtol=rtol))
    @test all(isapprox.(air_temperature.(ts), air_temperature.(ts_exact), rtol=rtol))

  end

end

@testset "moist thermodynamics - constructor consistency" begin

  # Make sure `ThermodynamicState` arguments are returned unchanged

  for FT in float_types
    rtol = FT(1e-2)
    e_int, ρ, q_tot, q_pt, T, p, θ_liq_ice = MT.tested_convergence_range(FT, 50)

    # PhaseDry
    ts = PhaseDry.(e_int, ρ)
    @test all(internal_energy.(ts) .≈ e_int)
    @test all(air_density.(ts) .≈ ρ)

    # PhaseEquil
    ts = PhaseEquil.(e_int, ρ, q_tot, 30, FT(1e-1), Ref(MT.saturation_adjustment_SecantMethod))
    @test all(internal_energy.(ts) .≈ e_int)
    @test all(getproperty.(PhasePartition.(ts),:tot) .≈ q_tot)
    @test all(air_density.(ts) .≈ ρ)

    ts = PhaseEquil.(e_int, ρ, q_tot)
    @test all(internal_energy.(ts) .≈ e_int)
    @test all(getproperty.(PhasePartition.(ts),:tot) .≈ q_tot)
    @test all(air_density.(ts) .≈ ρ)

    # PhaseNonEquil
    ts = PhaseNonEquil.(e_int, ρ, q_pt)
    @test all(internal_energy.(ts) .≈ e_int)
    @test all(getproperty.(PhasePartition.(ts),:tot) .≈ getproperty.(q_pt, :tot))
    @test all(getproperty.(PhasePartition.(ts),:liq) .≈ getproperty.(q_pt, :liq))
    @test all(getproperty.(PhasePartition.(ts),:ice) .≈ getproperty.(q_pt, :ice))
    @test all(air_density.(ts) .≈ ρ)

    # air_temperature_from_liquid_ice_pottemp_given_pressure-liquid_ice_pottemp inverse
    θ_liq_ice_ = liquid_ice_pottemp_given_pressure.(T, p, q_pt)
    @test all(air_temperature_from_liquid_ice_pottemp_given_pressure.(θ_liq_ice_, p, q_pt) .≈ T)

    # liquid_ice_pottemp-air_temperature_from_liquid_ice_pottemp_given_pressure inverse
    T = air_temperature_from_liquid_ice_pottemp_given_pressure.(θ_liq_ice, p, q_pt)
    @test all(liquid_ice_pottemp_given_pressure.(T, p, q_pt) .≈ θ_liq_ice)

    # Accurate but expensive `LiquidIcePotTempSHumNonEquil` constructor (Non-linear temperature from θ_liq_ice)
    T_non_linear = air_temperature_from_liquid_ice_pottemp_non_linear.(θ_liq_ice, ρ, FT(1e-3), 10, q_pt)
    T_expansion = air_temperature_from_liquid_ice_pottemp.(θ_liq_ice, ρ, q_pt)
    @test all(isapprox.(T_non_linear, T_expansion, rtol=rtol))
    e_int_ = internal_energy.(T_non_linear, q_pt)
    ts = PhaseNonEquil.(e_int_, ρ, q_pt)
    @test all(T_non_linear .≈ air_temperature.(ts))
    @test all(θ_liq_ice .≈ liquid_ice_pottemp.(ts))

    # LiquidIcePotTempSHumEquil
    ts = LiquidIcePotTempSHumEquil.(θ_liq_ice, ρ, q_tot, 40, FT(1e-3))
    @test all(isapprox.(liquid_ice_pottemp.(ts), θ_liq_ice, atol=1e-1))
    @test all(isapprox.(air_density.(ts), ρ, rtol=rtol))
    @test all(getproperty.(PhasePartition.(ts),:tot) .≈ q_tot)

    # The LiquidIcePotTempSHumEquil_given_pressure constructor
    # passes the consistency test within sufficient physical precision,
    # however, it fails to satisfy the consistency test within machine
    # precision for the input pressure.

    # LiquidIcePotTempSHumEquil_given_pressure
    ts = LiquidIcePotTempSHumEquil_given_pressure.(θ_liq_ice, p, q_tot, 40, FT(1e-3))
    @test all(isapprox.(liquid_ice_pottemp.(ts), θ_liq_ice, atol=1e-1))
    @test all(getproperty.(PhasePartition.(ts),:tot) .≈ getproperty.(q_pt, :tot))
    @test all(isapprox.(air_pressure.(ts), p, atol=FT(MSLP)*2e-2))

    # LiquidIcePotTempSHumNonEquil_given_pressure
    ts = LiquidIcePotTempSHumNonEquil_given_pressure.(θ_liq_ice, p, q_pt)
    @test all(liquid_ice_pottemp.(ts) .≈ θ_liq_ice)
    @test all(air_pressure.(ts) .≈ p)
    @test all(getproperty.(PhasePartition.(ts),:tot) .≈ getproperty.(q_pt,:tot))
    @test all(getproperty.(PhasePartition.(ts),:liq) .≈ getproperty.(q_pt,:liq))
    @test all(getproperty.(PhasePartition.(ts),:ice) .≈ getproperty.(q_pt,:ice))

    # LiquidIcePotTempSHumNonEquil
    ts = LiquidIcePotTempSHumNonEquil.(θ_liq_ice, ρ, q_pt, 5, FT(1e-3))
    @test all(θ_liq_ice .≈ liquid_ice_pottemp.(ts))
    @test all(air_density.(ts) .≈ ρ)
    @test all(getproperty.(PhasePartition.(ts),:tot) .≈ getproperty.(q_pt,:tot))
    @test all(getproperty.(PhasePartition.(ts),:liq) .≈ getproperty.(q_pt,:liq))
    @test all(getproperty.(PhasePartition.(ts),:ice) .≈ getproperty.(q_pt,:ice))
  end

end

@testset "moist thermodynamics - type-stability" begin

  # NOTE: `Float32` saturation adjustment tends to have more difficulty
  # with converging to the same tolerances as `Float64`, so they're relaxed here.
  FT = Float32
  e_int, ρ, q_tot, q_pt, T, p, θ_liq_ice = MT.tested_convergence_range(FT, 50)

  ρu = FT[1.0, 2.0, 3.0]
  e_pot = FT(100.0)
  @test typeof.(internal_energy.(ρ, ρ.*e_int, Ref(ρu), Ref(e_pot))) == typeof.(e_int)

  ts_dry             = PhaseDry.(e_int, ρ)
  ts_eq              = PhaseEquil.(e_int, ρ, q_tot, 15, FT(1e-1))
  ts_T               = TemperatureSHumEquil.(air_temperature.(ts_dry), air_pressure.(ts_dry), q_tot)
  ts_neq             = PhaseNonEquil.(e_int, ρ, q_pt)
  ts_θ_liq_ice_eq    = LiquidIcePotTempSHumEquil.(θ_liq_ice, ρ, q_tot, 40, FT(1e-3))
  ts_θ_liq_ice_eq_p  = LiquidIcePotTempSHumEquil_given_pressure.(θ_liq_ice, p, q_tot, 40, FT(1e-3))
  ts_θ_liq_ice_neq   = LiquidIcePotTempSHumNonEquil.(θ_liq_ice, ρ, q_pt)
  ts_θ_liq_ice_neq_p = LiquidIcePotTempSHumNonEquil_given_pressure.(θ_liq_ice, p, q_pt)

  for ts in (
    ts_eq,
    ts_dry,
    ts_T,
    ts_neq,
    ts_θ_liq_ice_eq,
    ts_θ_liq_ice_eq_p,
    ts_θ_liq_ice_neq,
    ts_θ_liq_ice_neq_p
    )
    @test typeof.(soundspeed_air.(ts))          == typeof.(e_int)
    @test typeof.(gas_constant_air.(ts))        == typeof.(e_int)
    @test typeof.(vapor_specific_humidity.(ts)) == typeof.(e_int)
    @test typeof.(relative_humidity.(ts))       == typeof.(e_int)
    @test typeof.(air_pressure.(ts))            == typeof.(e_int)
    @test typeof.(air_density.(ts))             == typeof.(e_int)
    @test typeof.(total_specific_humidity.(ts)) == typeof.(e_int)
    @test typeof.(cp_m.(ts))                    == typeof.(e_int)
    @test typeof.(cv_m.(ts))                    == typeof.(e_int)
    @test typeof.(air_temperature.(ts))         == typeof.(e_int)
    @test typeof.(internal_energy_sat.(ts))     == typeof.(e_int)
    @test typeof.(internal_energy.(ts))         == typeof.(e_int)
    @test typeof.(latent_heat_vapor.(ts))       == typeof.(e_int)
    @test typeof.(latent_heat_sublim.(ts))      == typeof.(e_int)
    @test typeof.(latent_heat_fusion.(ts))      == typeof.(e_int)
    @test typeof.(q_vap_saturation.(ts))        == typeof.(e_int)
    @test typeof.(saturation_excess.(ts))       == typeof.(e_int)
    @test typeof.(liquid_fraction.(ts))         == typeof.(e_int)
    @test typeof.(liquid_ice_pottemp.(ts))      == typeof.(e_int)
    @test typeof.(dry_pottemp.(ts))             == typeof.(e_int)
    @test typeof.(exner.(ts))                   == typeof.(e_int)
    @test typeof.(liquid_ice_pottemp_sat.(ts))  == typeof.(e_int)
    @test typeof.(specific_volume.(ts))         == typeof.(e_int)
    @test typeof.(virtual_pottemp.(ts))         == typeof.(e_int)
    @test eltype.(gas_constants.(ts))           == typeof.(e_int)
    @test typeof.(getproperty.(PhasePartition.(ts),:tot)) == typeof.(e_int)
  end

end

@testset "moist thermodynamics - dry limit" begin

  FT = Float64
  e_int, ρ, q_tot, q_pt, T, p, θ_liq_ice = MT.tested_convergence_range(FT, 50)

  # PhasePartition test is noisy, so do this only once:
  ts_dry = PhaseDry(first(e_int), first(ρ))
  ts_eq  = PhaseEquil(first(e_int), first(ρ), typeof(first(ρ))(0))
  @test PhasePartition(ts_eq).tot ≈ PhasePartition(ts_dry).tot
  @test PhasePartition(ts_eq).liq ≈ PhasePartition(ts_dry).liq
  @test PhasePartition(ts_eq).ice ≈ PhasePartition(ts_dry).ice

  ts_dry = PhaseDry.(e_int, ρ)
  ts_eq  = PhaseEquil.(e_int, ρ, q_tot.*0)

  @test all(gas_constant_air.(ts_eq)                       .≈ gas_constant_air.(ts_dry))
  @test all(relative_humidity.(ts_eq)                      .≈ relative_humidity.(ts_dry))
  @test all(air_pressure.(ts_eq)                           .≈ air_pressure.(ts_dry))
  @test all(air_density.(ts_eq)                            .≈ air_density.(ts_dry))
  @test all(specific_volume.(ts_eq)                        .≈ specific_volume.(ts_dry))
  @test all(total_specific_humidity.(ts_eq)                .≈ total_specific_humidity.(ts_dry))
  @test all(cp_m.(ts_eq)                                   .≈ cp_m.(ts_dry))
  @test all(cv_m.(ts_eq)                                   .≈ cv_m.(ts_dry))
  @test all(air_temperature.(ts_eq)                        .≈ air_temperature.(ts_dry))
  @test all(internal_energy.(ts_eq)                        .≈ internal_energy.(ts_dry))
  @test all(internal_energy_sat.(ts_eq)                    .≈ internal_energy_sat.(ts_dry))
  @test all(soundspeed_air.(ts_eq)                         .≈ soundspeed_air.(ts_dry))
  @test all(latent_heat_vapor.(ts_eq)                      .≈ latent_heat_vapor.(ts_dry))
  @test all(latent_heat_sublim.(ts_eq)                     .≈ latent_heat_sublim.(ts_dry))
  @test all(latent_heat_fusion.(ts_eq)                     .≈ latent_heat_fusion.(ts_dry))
  @test all(q_vap_saturation.(ts_eq)                       .≈ q_vap_saturation.(ts_dry))
  @test all(saturation_excess.(ts_eq)                      .≈ saturation_excess.(ts_dry))
  @test all(liquid_fraction.(ts_eq)                        .≈ liquid_fraction.(ts_dry))
  @test all(liquid_ice_pottemp.(ts_eq)                     .≈ liquid_ice_pottemp.(ts_dry))
  @test all(dry_pottemp.(ts_eq)                            .≈ dry_pottemp.(ts_dry))
  @test all(virtual_pottemp.(ts_eq)                        .≈ virtual_pottemp.(ts_dry))
  @test all(liquid_ice_pottemp_sat.(ts_eq)                 .≈ liquid_ice_pottemp_sat.(ts_dry))
  @test all(exner.(ts_eq)                                  .≈ exner.(ts_dry))

  @test all(saturation_vapor_pressure.(ts_eq, Ref(Ice()))      .≈ saturation_vapor_pressure.(ts_dry, Ref(Ice())))
  @test all(saturation_vapor_pressure.(ts_eq, Ref(Liquid()))   .≈ saturation_vapor_pressure.(ts_dry, Ref(Liquid())))
  @test all(first.(gas_constants.(ts_eq))  ≈ first.(gas_constants.(ts_dry)))
  @test all(last.(gas_constants.(ts_eq))  ≈ last.(gas_constants.(ts_dry)))

end

max_stars_repo_name: favba/CLIMA

max_issues_repo_name: favba/CLIMA

max_forks_repo_name: favba/CLIMA


text:
include("../imports.jl")

@testset ExtendedTestSet "ImageKeypointRegression" begin
    method = ImageKeypointRegression((16, 16), 10)
    DLPipelines.checkmethod_core(method)
    @testset "Show backends" begin
        @testset "ShowText" begin
            FastAI.test_method_show(method, ShowText(Base.DevNull()))
        end
        @testset "ShowMakie" begin
            FastAI.test_method_show(method, ShowMakie())
        end
    end
end