text:
# Copyright (c) 2020: Matthew Wilhelm & Matthew Stuber.
# This work is licensed under the Creative Commons Attribution-NonCommercial-
# ShareAlike 4.0 International License. To view a copy of this license, visit
# http://creativecommons.org/licenses/by-nc-sa/4.0/ or send a letter to Creative
# Commons, PO Box 1866, Mountain View, CA 94042, USA.
#############################################################################
# DynamicBoundspODEIneq.jl
# See https://github.com/PSORLab/DynamicBoundspODEIneq.jl
#############################################################################
# src/DifferentialInequality.jl
# Defines storage structures, the main relaxation routine, and access functions.
#############################################################################

"""
DifferentialInequalityCond <: Function

A functor (`<: Function`) used to check for an event were a relaxation ODE crosses the an interval bound ODE.
The functor is called using `(d::DifferentialInequalityCond)(g, x, t, integrator)`. The
`p` field of the `integrator` hold parameter values in components 1, ..., np. Components
np + 1, ..., 2*np are a flag is indicating a positive crossing. Components
2*np + 1, ..., 3*np are a flag is indicating a negative crossing.
"""
struct DifferentialInequalityCond <: Function
    nx::Int
    np::Int
    etol::Float64
end

function (d::DifferentialInequalityCond)(g, x::Vector{Float64}, t::Float64, integrator)::Nothing
    for i = 1:d.nx
        g[i] = x[2*d.nx + i] - x[i] - integrator.p[d.np + i]*d.etol
        g[d.nx + i] = x[d.nx + i] - x[3*d.nx + i]  - integrator.p[d.np + d.nx + i]*d.etol
        integrator.p[d.np + 2*d.nx + i] = x[2*d.nx + i] - x[i]          # positive if cv > lo
        integrator.p[d.np + 3*d.nx + i] = x[d.nx + i] - x[3*d.nx + i]   # positive if hi > cc
    end
    return
end

"""
DifferentialInequalityAffect <: Function

A functor (`<: Function`) that sets an indicator parameter to zero if a positive event crossing occurs.
"""
struct DifferentialInequalityAffect <: Function
    np::Int
    nx::Int
end

function (d::DifferentialInequalityAffect)(integrator, idx::Int)::Nothing
    @inbounds integrator.p[d.np + idx] = 0.0
    return
end

"""
DifferentialInequalityAffectNeg <: Function

A functor (`<: Function`) that sets an indicator parameter to zero if a negative event crossing occurs.
"""
struct DifferentialInequalityAffectNeg <: Function
    np::Int
    nx::Int
end

function (d::DifferentialInequalityAffectNeg)(integrator, idx::Int)::Nothing
    @inbounds integrator.p[d.np + idx] = 1.0
    return
end

"""
DifferentialInequalityf{Z, F} <: Function

A functor (`<: Function`) used to evaluate the r.h.s of the differential inequality. The
following constructor is used to initialize it:
`DifferentialInequalityf(f!, Z, nx, np, P, relax, subgrad, polyhedral_constraint, Xapriori)`.

$(TYPEDFIELDS)
"""
mutable struct DifferentialInequalityf{Z, F} <: Function
    "Right-hand side function"
    f!::F
    "Number of state variables"
    nx::Int
    "Dimensionality of polyhedral constraints."
    nm::Int
    "Number of decision variables"
    np::Int
    "Decision variable relaxation storage"
    p_mc::Vector{Z}
    "Decision variable interval bounds"
    P::Vector{Interval{Float64}}
    "Constraint state variable interval bounds"
    X::Vector{Interval{Float64}}
    "Input State Variable Temporary Storage"
    x_mc::Vector{Z}
    "Output State Variable Temporary Storage"
    xout_mc::Vector{Z}
    "Temporary Storage for Lower Beta"
    BetaL::Vector{Interval{Float64}}
    "Temporary Storage for Upper Beta"
    BetaU::Vector{Interval{Float64}}
    "Temporary Storage"
    xout_intv1::Vector{Interval{Float64}}
    "Temporary Storage"
    xout_intv2::Vector{Interval{Float64}}
    "Indicates that relaxations should be computed."
    calculate_relax::Bool
    "Indicates that subgradients should be computed (`calculate_relax` must be `true`)."
    calculate_subgradient::Bool
    prng::UnitRange{Int64}
    xrng::UnitRange{Int64}
    "Polyhedral constraint used"
    polyhedral_constraint::Union{PolyhedralConstraint, Nothing}
    has_apriori::Bool
    Xapriori::Vector{Interval{Float64}}
    "Constant value parameters"
    params::Vector{Float64}
    has_params::Bool
end

size(x::PolyhedralConstraint) = size(x.A)
size(x::Nothing) = 0,0

function set_parameters!(cb::DifferentialInequalityf, values::Vector{Float64})
    cb.params = values
    nothing
end

function DifferentialInequalityf(f!, Z, nx::Int, np::Int, P, relax::Bool, subgrad::Bool,
                                 polyhedral_constraint, Xapriori, params)
     has_params = !isempty(params)
     np = length(P)
     s1, s2 = size(polyhedral_constraint)
     DifferentialInequalityf{Z,typeof(f!)}(f!, nx, s1, np, zeros(Z, np), P,
                              zeros(Interval{Float64}, nx), zeros(Z, nx), zeros(Z, nx),
                              zeros(Interval{Float64}, nx), zeros(Interval{Float64}, nx),
                              zeros(Interval{Float64}, nx), zeros(Interval{Float64}, nx),
                              relax, subgrad, 1:np, 1:nx, polyhedral_constraint, false, Xapriori,
                              params, has_params)
end

function (d::DifferentialInequalityf{MC{N,T},F})(dx::Vector{Float64}, x::Vector{Float64},
                                     p::Vector{Float64}, t::Float64)::Nothing where {T<:RelaxTag, N, F}

    np = d.np
    nm = d.nm
    nx = d.nx
    has_params = d.has_params

    if d.calculate_relax
        d.p_mc.= MC{N,T}.(p[1:np], d.P, d.prng)
    end

    for i = d.xrng
        d.X[i] = Interval{Float64}(x[i], x[nx + i])
        if d.calculate_relax
            is_cvmid::Bool = (x[i] <= x[2*nx + i] <= x[nx + i]) && d.calculate_subgradient
            is_ccmid::Bool = (x[i] <= x[3*nx + i] <= x[nx + i]) && d.calculate_subgradient
            cvmid::Float64, _ = mid3(x[i], x[nx + i], x[2*nx + i])
            ccmid::Float64, _ = mid3(x[i], x[nx + i], x[3*nx + i])
            if is_cvmid
                cvsgmid::SVector{N,Float64} = SVector{N}([x[4*nx + N*(i-1) + j] for j = d.prng])
            else
                cvsgmid = zero(SVector{N,Float64})
            end
            if is_ccmid
                ccsgmid::SVector{N,Float64} = SVector{N}([x[4*nx + N*(nx+i-1) + j] for j = d.prng])
            else
                ccsgmid = zero(SVector{N,Float64})
            end
            d.x_mc[i] = MC{N,T}(cvmid,  ccmid, d.X[i], cvsgmid, ccsgmid, (~is_cvmid || ~is_ccmid))
        end
    end
    d.BetaL .= d.X
    d.BetaU .= d.X

    if d.calculate_relax
        if d.has_params
            d.f!(d.xout_mc, d.x_mc, d.p_mc, t, d.params)
        else
            d.f!(d.xout_mc, d.x_mc, d.p_mc, t)
        end
    end

    for i = d.xrng
        d.BetaL[i] = @interval(x[i])
        polyhedral_contact!(d.polyhedral_constraint, d.BetaL, d.Xapriori, nx, nm)
        if has_params
            d.f!(d.xout_intv1, d.BetaL, d.P, @interval(t), d.params)
        else
            d.f!(d.xout_intv1, d.BetaL, d.P, @interval(t))
        end
        dx[i] = d.xout_intv1[i].lo

        d.BetaU[i] = @interval(x[nx+i])
        polyhedral_contact!(d.polyhedral_constraint, d.BetaU, d.Xapriori, nx, nm)
        d.f!(d.xout_intv2, d.BetaU, d.P, @interval(t))
        if has_params
            d.f!(d.xout_intv2, d.BetaU, d.P, @interval(t), d.params)
        else
            d.f!(d.xout_intv2, d.BetaU, d.P, @interval(t))
        end
        dx[nx+i] = d.xout_intv2[i].hi

        if d.calculate_relax
            if (dx[i] > d.xout_mc[i].cv) && (p[np + i] == 1.0)
                dx[2*nx+i] = dx[i]
                if d.calculate_subgradient
                    for j=d.prng
                        dx[4*nx+(i-1)*np+j] = 0.0
                    end
                end
            else
                dx[2*nx+i] = d.xout_mc[i].cv
                if d.calculate_subgradient
                    for j=d.prng
                        dx[4*nx+(i-1)*np+j] = d.xout_mc[i].cv_grad[j]
                    end
                end
            end
            if (dx[nx+i] < d.xout_mc[i].cc) && (p[np + nx + i] == 1.0)
                dx[3*nx+i] = dx[nx+i]
                if d.calculate_subgradient
                    for j=d.prng
                        dx[(4+np)*nx+(i-1)*np+j] = 0.0
                    end
                end
            else
                dx[3*nx+i] = d.xout_mc[i].cc
                if d.calculate_subgradient
                    for j=d.prng
                        dx[(4+np)*nx+(i-1)*np+j] = d.xout_mc[i].cc_grad[j]
                    end
                end
            end
        end
        d.BetaL[i] = d.X[i]
        d.BetaU[i] = d.X[i]
    end
    return
end

"""
$(TYPEDEF)

The DifferentialInequality type integrator represents the relaxed pODE problem as a
2*nx dimension ODE problem if `calculate_relax` is `false`, a 4*nx dimension ODE problem
if `calculate_relax` is `true` and `4*nx + 4*np*nx` if `calculate_subgradients` is also
set to `true`.
- `x[1:nx]` are the lower interval bounds
- `x[(1+nx):2*nx]` are upper interval bounds
- `x[(1+2*nx):3*nx]` are convex relaxations (computed if calculate_relax = true)
- `x[(1+3*nx):4*nx]` are concave relaxations (computed if calculate_relax = true)
- `x[(1+4*nx):(4+np)*nx]` are subgradients of the convex relaxations (computed if calculate_subgradients = true)
- `x[(1+(4+np)*nx):(4+2*np)*nx]` are subgradients of the concave relaxations (computed if calculate_subgradients = true)
The first `np` parameter values correspond to the parameter values in the original
problem. If (`calculate_relax == true`) then `(np+1):(np+nx)` parameter values correspond
to the `b_i^c` variables used to detect a relaxation crossing a lower state bound.
The `(np+nx+1):(np+nx)` parameter values correspond to the `b_i^C` variables used
to detect a relaxation crossing a upper state bound. The variables are floats but
are valued 0.0 and 1.0 and can be intepreted as the corresponding 0-1 Boolean values.
Otherwise, only `np` parameter values are used.

$(TYPEDFIELDS)
"""
mutable struct DifferentialInequality{F, N, T<:RelaxTag, PRB1<:AbstractODEProblem,
                 INT1, CB<:AbstractContinuousCallback} <: AbstractODERelaxIntegrator
    calculate_relax::Bool
    calculate_subgradient::Bool
    calculate_local_sensitivity::Bool
    differentiable::Bool
    event_soft_tol::Float64
    params::Vector{Float64}
    p::Vector{Float64}
    pL::Vector{Float64}
    pU::Vector{Float64}
    p_mc::Vector{MC{N,T}}
    x0f::F
    x0::Vector{Float64}
    x0_mc::Vector{MC{N,T}}
    xL::ElasticArray{Float64,2,1}
    xU::ElasticArray{Float64,2,1}
    relax_ode_prob::PRB1
    relax_ode_integrator::INT1
    relax_t::Vector{Float64}
    relax_lo::ElasticArray{Float64,2,1}
    relax_hi::ElasticArray{Float64,2,1}
    relax_cv::ElasticArray{Float64,2,1}
    relax_cc::ElasticArray{Float64,2,1}
    relax_cv_grad::ElasticArray{SVector{N,Float64},2,1}
    relax_cc_grad::ElasticArray{SVector{N,Float64},2,1}
    relax_mc::ElasticArray{MC{N,T},2,1}
    vector_callback::CB
    integrator_state::IntegratorStates
    local_problem_storage
    np::Int
    nx::Int
    nt::Int
    relax_t_dict_flt::Dict{Float64,Int64}
    relax_t_dict_indx::Dict{Int64,Int64}
    polyhedral_constraint::Union{PolyhedralConstraint,Nothing}
    has_params::Bool
    prob::ODERelaxProb
end

function DifferentialInequality(d::ODERelaxProb; calculate_relax::Bool = true,
                   calculate_subgradient::Bool = true,
                   differentiable::Bool = false,
                   calculate_local_sensitivity::Bool = true,
                   event_soft_tol = 1E-4,
                   relax_ode_integrator = CVODE_Adams(),
                   local_ode_integrator = CVODE_Adams())

    @assert ~calculate_subgradient || calculate_relax "Relaxations must be computed in order to compute subgradients"

    np = length(d.p)
    Z = differentiable ? MC{np,Diff} : MC{np,NS}
    xdim = calculate_relax ? 4*d.nx : 2*d.nx
    xdim += calculate_subgradient ? 2*d.nx*d.np : 0
    pdim = d.np + (calculate_relax ? 4*d.nx : 0)
    pL = d.pL
    pU = d.pU
    p = zeros(Float64, pdim)
    p[1:d.np] = 0.5*(d.pL + d.pU)

    utemp = zeros(Float64, xdim)
    x0temp = zeros(Float64, xdim)
    x0mctemp = zeros(Z, xdim)
    differentialInequalitycond! = DifferentialInequalityCond(d.nx, np, event_soft_tol)
    differentialInequalityaffect! = DifferentialInequalityAffect(np, d.nx)
    differentialInequalityaffectneg! = DifferentialInequalityAffectNeg(np, d.nx)
    vector_callback = VectorContinuousCallback(differentialInequalitycond!, differentialInequalityaffect!,
                                               2*d.nx; affect_neg! = differentialInequalityaffectneg!,
                                               rootfind=true, save_positions=(false,false),
                                               interp_points=20)

    has_params = !isempty(d.params)
    params = has_params ? d.params : Float64[]
    P = Interval{Float64}.(pL, pU)
    p_mc = zeros(Z, np)
    const_bnds = d.constant_state_bounds # const_bnds = get(d, ConstantBounds())
    polyhedral_constraint = d.polyhedral_constraint #get(d, DenseLinearInvariant())
    if polyhedral_constraint !== nothing
        X_natural_box = Interval{Float64}.(const_bnds.xL, const_bnds.xU)
    else
        X_natural_box = fill(Interval{Float64}(-Inf,Inf), d.nx)
    end
    f = DifferentialInequalityf(d.f, Z, d.nx, np, P, calculate_relax, calculate_subgradient,
                                polyhedral_constraint, X_natural_box, params)

    relax_ode_prob = ODEProblem(f, utemp, d.tspan, p)
    local_problem_storage = ODELocalIntegrator(d, local_ode_integrator)

    relax_t = Float64[]
    relax_lo = ElasticArray(zeros(Float64,d.nx,2))
    relax_hi = ElasticArray(zeros(Float64,d.nx,2))
    relax_cv = ElasticArray(zeros(Float64,d.nx,2))
    relax_cc = ElasticArray(zeros(Float64,d.nx,2))
    relax_cv_grad = ElasticArray(zeros(SVector{np,Float64},d.nx,2))
    relax_cc_grad = ElasticArray(zeros(SVector{np,Float64},d.nx,2))
    relax_mc = ElasticArray(zeros(Z,d.nx,2))

    xL = ElasticArray(zeros(Float64,d.nx,2))
    xU = ElasticArray(zeros(Float64,d.nx,2))

    support_set = DBB.get(d, DBB.SupportSet())
    nt = length(support_set.s)
    relax_t_dict_flt = Dict{Float64,Int64}()
    relax_t_dict_indx = Dict{Int64,Int64}()
    for (i,s) in enumerate(support_set.s)
        relax_t_dict_flt[s] = i
        relax_t_dict_indx[i] = i
    end

    DifferentialInequality(calculate_relax, calculate_subgradient, calculate_local_sensitivity,
                           differentiable,
                          event_soft_tol, params, p, pL, pU, p_mc, d.x0, x0temp,
                          x0mctemp, xL, xU, relax_ode_prob, relax_ode_integrator, relax_t,
                          relax_lo, relax_hi, relax_cv, relax_cc, relax_cv_grad, relax_cc_grad,
                          relax_mc, vector_callback, IntegratorStates(),
                          local_problem_storage, np, d.nx, nt, relax_t_dict_flt,
                          relax_t_dict_indx, d.polyhedral_constraint, has_params, d)
end

function relax!(d::DifferentialInequality{F, N, T, PRB1, INT1, CB}) where {F, N, T<:RelaxTag,
                                                                    PRB1<:AbstractODEProblem, INT1,
                                                                    CB<:AbstractContinuousCallback}

    empty!(d.relax_t_dict_flt)
    empty!(d.relax_t_dict_indx)

    # load functors
    if d.integrator_state.new_decision_pnt
        for i=1:d.np
            d.relax_ode_prob.f.f.P[i] = Interval{Float64}(d.pL[i], d.pU[i])
            d.p_mc[i] = MC{N, NS}(d.p[i], d.relax_ode_prob.f.f.P[i], i)
        end
        # populate initial condition
        if d.has_params
            d.x0_mc = d.x0f(d.p_mc, d.params)
        else
            d.x0_mc = d.x0f(d.p_mc)
        end
        for i=1:d.nx
            d.x0[i] = d.x0_mc[i].Intv.lo
            d.x0[d.nx + i] = d.x0_mc[i].Intv.hi
            if d.calculate_relax
                d.x0[2*d.nx + i] = d.x0_mc[i].cv
                d.x0[3*d.nx + i] = d.x0_mc[i].cc
                if d.calculate_subgradient
                    for j = 1:d.np
                        d.x0[4*d.nx + j + d.np*(i - 1)] = d.x0_mc[i].cv_grad[j]
                        d.x0[(4 + d.np)*d.nx + j + d.np*(i - 1)] = d.x0_mc[i].cc_grad[j]
                    end
                end
            end
        end
        d.relax_ode_prob.f.f.params .= d.params
        d.relax_ode_prob = deepcopy(remake(d.relax_ode_prob; u0=d.x0))
        if d.calculate_relax
            relax_ode_sol = solve(d.relax_ode_prob, d.relax_ode_integrator,
                                  callback = d.vector_callback, abstol = 1E-6,
                                  reltol = 1E-3,
                                  saveat = d.local_problem_storage.user_t)
        else
            relax_ode_sol = solve(d.relax_ode_prob, d.relax_ode_integrator,
                                  abstol = 1E-6, reltol = 1E-3,
                                  saveat = d.local_problem_storage.user_t)
        end
        relax_ode_sol_t = relax_ode_sol.t::Vector{Float64}
        new_length::Int = length(relax_ode_sol_t)
        resize!(d.relax_t, new_length)
        d.relax_t .= relax_ode_sol_t

        for (tindx, t) in enumerate(relax_ode_sol_t)
            d.relax_t_dict_flt[t] = tindx
        end

        if !isempty(d.local_problem_storage.user_t)
            next_support_time = d.local_problem_storage.user_t[1]
            supports_left = length(d.local_problem_storage.user_t)
            loc_count = 1
            for (tindx, t) in enumerate(relax_ode_sol_t)
                if t == next_support_time
                    d.relax_t_dict_indx[loc_count] = tindx
                    loc_count += 1
                    supports_left -= 1
                    if supports_left > 0
                        next_support_time = d.local_problem_storage.user_t[loc_count]
                    end
                end
            end
        end

        resize!(d.relax_lo, d.nx, new_length)
        resize!(d.relax_hi, d.nx, new_length)
        resize!(d.relax_cv, d.nx, new_length)
        resize!(d.relax_cc, d.nx, new_length)
        resize!(d.relax_cv_grad, d.nx, new_length)
        resize!(d.relax_cc_grad, d.nx, new_length)
        resize!(d.relax_mc, d.nx, new_length)

        for i in 1:new_length
            time_step_sol::Vector{Float64} = relax_ode_sol[i]
            for j in 1:d.nx
                step_lo::Float64 = time_step_sol[j]
                step_hi::Float64 = time_step_sol[d.nx + j]
                d.relax_lo[j,i] = step_lo
                d.relax_hi[j,i] = step_hi
                if d.calculate_relax
                    step_cv::Float64 = time_step_sol[2*d.nx + j]
                    step_cc::Float64 = time_step_sol[3*d.nx + j]
                    d.relax_cv[j, i] = step_cv
                    d.relax_cc[j, i] = step_cc
                    if d.calculate_subgradient
                        step_cv_grad::SVector{N,Float64} = SVector{N}(Float64[time_step_sol[4*d.nx + d.np*(j - 1) + k] for k = 1:d.np])
                        step_cc_grad::SVector{N,Float64} = SVector{N}(Float64[time_step_sol[4*d.nx + d.np*(d.nx + j - 1) + k] for k = 1:d.np])
                        d.relax_cv_grad[j, i] = step_cv_grad
                        d.relax_cc_grad[j, i] = step_cc_grad
                    else
                        step_cv_grad = zero(SVector{d.np,Float64})::SVector{N,Float64}
                        step_cc_grad = zero(SVector{d.np,Float64})::SVector{N,Float64}
                    end
                    d.relax_mc[j,i] = MC{N,NS}(step_cv, step_cc, Interval{Float64}(step_lo, step_hi),
                                               step_cv_grad, step_cc_grad, false)
                end
            end
        end
    end
    return
end

DBB.supports(::DifferentialInequality, ::DBB.IntegratorName) = true
function DBB.supports(t::DifferentialInequality, ::DBB.Gradient{T}) where {T <: AbstractBoundLoc}
    t.calculate_subgradient
end
function DBB.supports(t::DifferentialInequality, ::DBB.Subgradient{T}) where {T <: AbstractBoundLoc}
    t.calculate_subgradient
end
DBB.supports(::DifferentialInequality, ::DBB.Bound{Lower}) = true
DBB.supports(::DifferentialInequality, ::DBB.Bound{Upper}) = true
DBB.supports(t::DifferentialInequality, ::DBB.Relaxation{Lower}) = t.calculate_relax
DBB.supports(t::DifferentialInequality, ::DBB.Relaxation{Upper}) = t.calculate_relax
DBB.supports(::DifferentialInequality, ::DBB.IsNumeric) = true
DBB.supports(::DifferentialInequality, ::DBB.IsSolutionSet) = true
DBB.supports(::DifferentialInequality, ::DBB.TerminationStatus) = true
DBB.supports(::DifferentialInequality, ::DBB.Value) = true
DBB.supports(::DifferentialInequality, ::DBB.ParameterBound{Lower}) = true
DBB.supports(::DifferentialInequality, ::DBB.ParameterBound{Upper}) = true
DBB.supports(::DifferentialInequality, ::DBB.ParameterValue) = true
DBB.supports(::DifferentialInequality, ::DBB.SupportSet) = true
DBB.supports(::DifferentialInequality, ::DBB.ParameterNumber) = true
DBB.supports(::DifferentialInequality, ::DBB.StateNumber) = true
DBB.supports(::DifferentialInequality, ::DBB.SupportNumber) = true
DBB.supports(t::DifferentialInequality, ::DBB.LocalSensitivityOn) = t.calculate_local_sensitivity

DBB.get(t::DifferentialInequality, v::DBB.IntegratorName) = "DifferentialInequality Integrator"
DBB.get(t::DifferentialInequality, v::DBB.IsNumeric) = false
DBB.get(t::DifferentialInequality, v::DBB.IsSolutionSet) = true
DBB.get(t::DifferentialInequality, v::DBB.TerminationStatus) = t.integrator_state.termination_status
DBB.get(t::DifferentialInequality, v::DBB.SupportSet) = t.prob.support_set
DBB.get(t::DifferentialInequality, v::DBB.ParameterNumber) = t.np
DBB.get(t::DifferentialInequality, v::DBB.StateNumber) = t.nx
DBB.get(t::DifferentialInequality, v::DBB.SupportNumber) = t.nt
DBB.get(t::DifferentialInequality, v::DBB.LocalSensitivityOn) = t.calculate_local_sensitivity

function DBB.get(t::DifferentialInequality, v::DBB.LocalIntegrator)
    return t.local_problem_storage
end

DBB.get(t::DifferentialInequality, v::DBB.AttachedProblem) = t.prob

function DBB.set!(t::DifferentialInequality, v::DBB.LocalSensitivityOn, b::Bool)
    t.calculate_local_sensitivity = b
    return
end

function get_val_loc(t::DifferentialInequality, index::Int64, time::Float64)
    (index <= 0 && time == -Inf) && error("Must set either index or time.")
    if index > 0
        return t.relax_t_dict_indx[index]
    end
    t.relax_t_dict_flt[time]
end

function DBB.get(out::Vector{Float64}, t::DifferentialInequality, v::DBB.Bound{Lower})
     val_loc = get_val_loc(t, v.index, v.time)
     out .= t.relax_lo[:, val_loc]
     return
end

function DBB.get(out::Vector{Float64}, t::DifferentialInequality, v::DBB.Bound{Upper})
    val_loc = get_val_loc(t, v.index, v.time)
    out .= t.relax_hi[:, val_loc]
    return
end

function DBB.get(out::Vector{Float64}, t::DifferentialInequality, v::DBB.Relaxation{Lower})
    val_loc = get_val_loc(t, v.index, v.time)
    out .= t.relax_cv[:, val_loc]
    return
end

function DBB.get(out::Vector{Float64}, t::DifferentialInequality, v::DBB.Relaxation{Upper})
    val_loc = get_val_loc(t, v.index, v.time)
    out .= t.relax_cc[:, val_loc]
    return
end

function DBB.get(out::Matrix{Float64}, t::DifferentialInequality, v::DBB.Subgradient{Lower})
    val_loc = get_val_loc(t, v.index, v.time)

    for i = 1:t.np
        if !t.calculate_relax
            fill!(out, 0.0)
        else
            for j = 1:t.nx
                out[j,i] = t.relax_cv_grad[j,val_loc][i]
            end
        end
    end
    return
end

function DBB.get(out::Matrix{Float64}, t::DifferentialInequality, v::DBB.Subgradient{Upper})
    val_loc = get_val_loc(t, v.index, v.time)
    for i = 1:t.np
        if !t.calculate_relax
            fill!(out, 0.0)
        else
            for j = 1:t.nx
                out[j,i] = t.relax_cc_grad[j,val_loc][i]
            end
        end
    end
    return
end

function DBB.getall!(out::Vector{Array{Float64,2}}, t::DifferentialInequality, g::DBB.Subgradient{Lower})
    if !t.calculate_relax
        for i = 1:t.np
            fill!(out[i], 0.0)
        end
    else
        for i = 1:length(t.relax_cv_grad[1])
            for j = 1:t.np
                for k = 1:t.nx
                    out[j][k,i]= t.relax_cv_grad[k,i][j]
                end
            end
        end
    end
    return
end
function DBB.getall!(out::Vector{Array{Float64,2}}, t::DifferentialInequality, g::DBB.Subgradient{Upper})
    if !t.calculate_relax
        for i = 1:t.np
            fill!(out[i], 0.0)
        end
    else
        for i = 1:length(t.relax_cc_grad[1])
            for j = 1:t.np
                for k = 1:t.nx
                    out[j][k,i]= t.relax_cc_grad[k,i][j]
                end
            end
        end
    end
    return
end

function DBB.getall!(out::Vector{Array{Float64,2}}, t::DifferentialInequality, g::DBB.Gradient{Lower})
    if ~t.differentiable
        error("Integrator does not generate differential relaxations. Create integrator with
               differentiable field to set to true and reintegrate.")
    end
    DBB.getall!(out, t, DBB.Subgradient{Lower}())
    return
end
function DBB.getall!(out::Vector{Array{Float64,2}}, t::DifferentialInequality, g::DBB.Gradient{Upper})
    if ~t.differentiable
        error("Integrator does not generate differential relaxations. Create integrator with
               differentiable field to set to true and reintegrate.")
    end
    DBB.getall!(out, t, DBB.Subgradient{Upper}())
    return
end

function DBB.getall!(out::Array{Float64,2}, t::DifferentialInequality, v::DBB.Bound{Lower})
    out .= t.relax_lo
    return
end

function DBB.getall!(out::Vector{Float64}, t::DifferentialInequality, v::DBB.Bound{Lower})
    @inbounds for i in eachindex(out)
        out[i] = t.relax_lo[1,i]
    end
    return
end

function DBB.getall!(out::Array{Float64,2}, t::DifferentialInequality, v::DBB.Bound{Upper})
    out .= t.relax_hi
    return
end

function DBB.getall!(out::Vector{Float64}, t::DifferentialInequality, v::DBB.Bound{Upper})
    @inbounds for i in eachindex(out)
        out[i] = t.relax_hi[1,i]
    end
    return
end

function DBB.getall!(out::Array{Float64,2}, t::DifferentialInequality, v::DBB.Relaxation{Lower})
    if !t.calculate_relax
        @inbounds for i in eachindex(out)
            out[i] = t.relax_lo[i]
        end
    else
        @inbounds for i in eachindex(out)
            out[i] = t.relax_cv[i]
        end
    end
    return
end

function DBB.getall!(out::Vector{Float64}, t::DifferentialInequality, v::DBB.Relaxation{Lower})
    if !t.calculate_relax
        @inbounds for i in eachindex(out)
            out[i] = t.X[1,i].lo
        end
    else
        for i in eachindex(out)
            out[i] = t.relax_cv[1,i]
        end
    end
    return
end

function DBB.getall!(out::Array{Float64,2}, t::DifferentialInequality, v::DBB.Relaxation{Upper})
    if !t.calculate_relax
        @inbounds for i in eachindex(out)
            out[i] = t.X[i].hi
        end
    else
        @inbounds for i in eachindex(out)
            out[i] = t.relax_cc[i]
        end
    end
    return
end

function DBB.getall!(out::Vector{Float64}, t::DifferentialInequality, v::DBB.Relaxation{Upper})
    if !t.calculate_relax
        @inbounds for i in eachindex(out)
            out[i] = t.X[1,i].hi
        end
    else
        for i in eachindex(out)
            out[i] = t.relax_cc[1,i]
        end
    end
    return
end

function DBB.setall!(t::DifferentialInequality, v::DBB.ParameterBound{Lower}, value::Vector{Float64})
    t.integrator_state.new_decision_box = true
    @inbounds for i = 1:t.np
        t.pL[i] = value[i]
    end
    return
end

function DBB.setall!(t::DifferentialInequality, v::DBB.ParameterBound{Upper}, value::Vector{Float64})
    t.integrator_state.new_decision_box = true
    @inbounds for i = 1:t.np
        t.pU[i] = value[i]
    end
    return
end

function DBB.setall!(t::DifferentialInequality, v::DBB.ParameterValue, value::Vector{Float64})
    t.integrator_state.new_decision_pnt = true
    @inbounds for i = 1:t.np
        t.p[i] = value[i]
    end
    return
end

function DBB.set!(t::DifferentialInequality, v::DBB.ConstantParameterValue, value)
    t.integrator_state.new_decision_pnt = true
    t.params[v.i] = value
    return
end
function DBB.setall!(t::DifferentialInequality, v::DBB.ConstantParameterValue, value::Vector{Float64})
    t.integrator_state.new_decision_pnt = true
    @inbounds for i = 1:t.np
        t.params[i] = value[i]
    end
    return
end

function DBB.get(t::DifferentialInequality, v::DBB.ConstantParameterValue)
    return t.params[v.i]
end
function DBB.getall(t::DifferentialInequality, v::DBB.ConstantParameterValue)
    return t.params
end
function DBB.getall(out, t::DifferentialInequality, v::DBB.ConstantParameterValue)
    out .= t.params
    return
end

function DBB.set!(t::DifferentialInequality, v::DBB.ParameterBound{Lower}, value::T) where T <: Union{Integer, AbstractFloat}
    t.integrator_state.new_decision_box = true
    @inbounds t.pL[v.i] = value
    return
end

function DBB.set!(t::DifferentialInequality, v::DBB.ParameterBound{Upper}, value::T) where T <: Union{Integer, AbstractFloat}
    t.integrator_state.new_decision_box = true
    @inbounds t.pU[v.i] = value
    return
end

DBB.getall(t::DifferentialInequality, v::DBB.ParameterBound{Lower}) = t.pL
DBB.getall(t::DifferentialInequality, v::DBB.ParameterBound{Upper}) = t.pU

function DBB.getall!(out::Vector{Float64}, t::DifferentialInequality, v::DBB.ParameterBound{Lower})
    out .= t.pL
    return
end
function DBB.getall!(out::Vector{Float64}, t::DifferentialInequality, v::DBB.ParameterBound{Upper})
    out .= t.pU
    return
end

function DBB.getall!(out::Vector{Float64}, t::DifferentialInequality, v::DBB.ParameterValue)
    out .= t.p[1:t.np]
    return
end

function DBB.set!(t::DifferentialInequality, v::DBB.ParameterValue, value::T) where T <: Union{Integer, AbstractFloat}
    t.integrator_state.new_decision_pnt = true
    @inbounds t.p[v.i] = value
    return
end

function DBB.setall!(t::DifferentialInequality, v::DBB.ConstantStateBounds)
    if t.integrator_state.new_decision_box
        t.integrator_state.set_lower_state = true
        t.integrator_state.set_upper_state = true
    end
    relax_ode_prob.f!.Xapriori = Interval.(v.xL, v.xU)
    return
end


text:
###################getHomosked!#######################
#Gets the covariance matrix under Homoskedastic assumptions σ^2[X'X]^-1

#main method to get the covariance matrix
#IN: The QR decomposition from the regression, a vector of residuals,
#optionally memory for the covariance matrix
#OUT: Writes and returns the covariance matrix
getHomoskedΣ!(args...; keyargs...) = error("use homoskedasticΣ instead")

function homoskedasticΣ!(xqr::FMQR{M}, ε::V, Σ::M = M(undef, lin.K, lin.K);
  dofcorrect::Float64 = 1.0)::M where  {M<:AbstractMatrix, V<:AbstractVector}

  BLAS.gemm!('N','T',(ε⋅ε)/xqr.N,xqr.Rinv,xqr.Rinv,0.0,Σ)

  Σ .*= dofcorrect

  return Σ #[X'X]^-1*σ2
end

#=helper method for the above function which extracts the QR decomposition
from a linear model
IN: A linear model, optionally memory for the covariance matrix
OUT: WRites and returns the covariance matrix=#
function homoskedasticΣ!(lin::FMLM{M},
  Σ::M = M(undef, lin.K, lin.K))::M where M<:AbstractMatrix

  return homoskedasticΣ!(lin.xqr, lin.ε, Σ, dofcorrect = lin.N / lin.dof)
end

#=Simple test function to verify the QR-related optimizations
Output should be identical to the standard methods
IN: Independent variable X matrix
OUT: Returns the covariance matrix=#
getHomoskedΣSlow(args...; keyargs...) = error("use homoskedasticΣslow instead")
function homoskedasticΣslow(X::M, ε::V; dofcorrect::Float64=1.0)::M where {M<:AbstractMatrix, V<:AbstractVector}
  T::Type = eltype(V)

  out = Matrix(X' * X)\I .* (ε⋅ε)/size(X,1)
  return out * dofcorrect
end

#Convenience method for above: IN: A linear model #OUT: A covariance matrix
function homoskedasticΣslow(lin::FMLM{M, V}) where {M<:AbstractMatrix, V<:AbstractVector}
  return homoskedasticΣslow(lin.X, lin.Y.-lin.X * lin.β, dofcorrect = lin.N / lin.dof)::M
end

###################getWhiteERRORs ###OLS only
#checked against vcovHC in R (which uses the HC3 algorithm)
getWhiteΣ!(args...; keyargs...) = error("use whiteΣ! instead")
function whiteΣ!(xqr::FMQR{M}, ε::V,
  Σ::M = M(undef, xqr.K, xqr.K),
  dofcorrect::Float64=1.0)::M where {M<:AbstractMatrix, V<:AbstractVector}

  Λ::V = ε.^2.0

  QRtinv::M = BLAS.gemm('N','T',xqr.Q,xqr.Rinv) #Q(R')^-1
  RinvQΛ::M = QRtinv'
  RinvQΛ *= Diagonal(Λ)

  BLAS.gemm!('N','N',1.0,RinvQΛ, QRtinv, 0.0, Σ)

  Σ .*= dofcorrect

  #final multiplicaiton and assignment
  return Σ #R^-1Q'ΛQ(R')^-1
end

function whiteΣ!(lin::FMLM{M},
  Σ::M = M(undef, lin.K, lin.K))::M where M<:AbstractMatrix

  return whiteΣ!(lin.xqr, lin.ε, Σ, lin.N/lin.dof)
end

#for testing purposes only
whiteΣslow(args...; keyargs...) = error("use whiteΣslow instead")
function whiteΣslow(X::M, ε::V, dofcorrect::Float64=1.0)  where {M<:AbstractMatrix, V<:AbstractVector}
  XXinv::M = (X' * X)\I

  return XXinv * X' * diagm(ε) .^ 2.0 * X * XXinv .*dofcorrect
end

#Convenience method for above: IN: A linear model #OUT: A covariance matrix
whiteΣslow(lin::FMLM) = whiteΣslow(lin.X, lin.Y.-lin.X * lin.β, lin.N/lin.dof)

###########################getClustered SEs
#method for getting clustered errors
#from http://cameron.econ.ucdavis.edu/research/Cameron_Miller_JHR_2015_February.pdf
#also provides methedology for multi-way clsutering
#forumala: [X'X]^-1*B*[X'X]^-1 where B= sum over G (Xg'*εg*εg'*Xg) and g is indexed for G clusters
#checked against felm's [modelvar].cse with exactDOF=TRUE in R
getClustered!(args...; keyargs...) = error("use clusteredΣ! instead")

function clusteredΣ!(X::M, xqr::FMQR{M}, ε::V, clusters::C,
    Σ::M = M(undef, xqr.K, xqr.K),
    dof::Int = xqr.N-xqr.K)::M where {M<:AbstractMatrix, V<:AbstractVector, C<:AbstractVector}

  #@info "xqr.K: $(xqr.K)"
  #@info "Σ: $(size(Σ))"

  #convenience vectors and values
  inddf::DataFrame = DataFrame(clusters=clusters, idx = 1:xqr.N)
  sindfs::GroupedDataFrame = groupby(inddf, :clusters)
  Nclusters::Int = length(sindfs)

  #ctable::Dict = Dict(cvars[i] => i for i::Int ∈ 1:Nclusters)
  #clustercode::Vector{Int} = ((x)->ctable[x]).(clusters)

  #iterate through the groups
  B::M = zeros(xqr.K, xqr.K)
  @fastmath for sindf::SubDataFrame ∈ sindfs
    Xg::SubArray = view(X, sindf.idx, :)
    ug::SubArray = view(ε, sindf.idx)
    B .+= Xg'*ug*ug'*Xg
  end

  # Calc [X'X]^-1*B*[X'X]^-1
  RRinv::M =  BLAS.gemm('N','T',xqr.Rinv,xqr.Rinv)
  RRinvB::M = BLAS.gemm('N','N',RRinv,B)

  dofcorrect::Float64 = Nclusters/(Nclusters-1.)*(xqr.N-1.)/dof #small sample correction

  if M<:Matrix{Float64}
    ##try
      BLAS.gemm!('N','N',dofcorrect,RRinvB,RRinv,0.0,Σ)
    #=catch err
      @info "RRinvB size: $(size(RRinvB))"
      @info "RRinv size: $(size(RRinv))"
      @info "Σ size: $(size(Σ))"
      @info "B size: $(size(B))"
      @info "xqr.K: $(xqr.K)"
      error("err: $err")
    end=#
  else
    Σ = Matrix(BLAS.gemm('N','N',dofcorrect,RRinvB,RRinv))
  end

  return Σ
end

#handles 2-way clustering
#checked against felm's [modelvar].cse with exactDOF=TRUE in R
function clusteredΣ!(lin::FMLM{M, V}, Σ₁::M = M(undef, lin.K, lin.K);
    clusters::Vector{<:Vector}#=Union{FMClusters, Vector{<:FMData}}=# = lin.clusters, #allows for an override on the standard errors
    testequivelance::Bool = false)::M where {M<:AbstractMatrix, V<:AbstractVector}

  @assert lin.K == size(Σ₁,1) == size(Σ₁,2)
  length(clusters) ∈ [1,2] || error("Illegal number of clusters = $(length(clusters))")
  clusteredΣ!(lin.X, lin.xqr, lin.ε, clusters[1], Σ₁, lin.dof)
  if length(clusters) == 2
    Σ₂::M = similar(Σ₁)
    Σ₃::M = similar(Σ₁)

    #create the intersection of the clusters

    cluster⋂::Vector{Symbol} = ((s1, s2)->
      Symbol(string(s1),"_∩_",string(s2))).(clusters[1],clusters[2])

    #get the remaining clusters
    clusteredΣ!(lin.X, lin.xqr, lin.ε, clusters[2], Σ₂, lin.dof)

    #may verify the equivelance of white Σ and the cluster intersection Σ in the case where
    #cluster 1 within cluster 2 accounts for all variation (e.g. combination of clusters is unique)
    if testequivelance
      clusteredΣ!(lin.X, lin.xqr, lin.ε, cluster⋂, Σ₃, lin.dof)
      if length(unique(cluster⋂)) ≠ length(cluster⋂)
        @warn "Note: cluster intersection not unique. White results will not reconcile."
      end
      println("Clustered version:")
      display(Σ₃)

      whiteΣ!(lin.xqr, lin.ε, Σ₃, lin.N/lin.dof)
      println("White version:")
      display(Σ₃)
    else #standard case
      if length(unique(cluster⋂)) == length(cluster⋂) #special case allowing for performance optimization
        whiteΣ!(lin.xqr, lin.ε, Σ₃, lin.N/lin.dof)
      else
        clusteredΣ!(lin.X, lin.xqr, lin.ε, cluster⋂, Σ₃, lin.dof)
      end
    end

    Σ₁ .= Σ₁ .+ Σ₂ .- Σ₃

  end

  return Σ₁
end

###################getMWErrors!##############################
#Main method to get the modified white SEs Does #[X'X]^-1X'ΛX[X'X]^-1

#=Main method to calculate the modified white SEs
IN: A FMQR decomposition object, a residuals vector,
optional memory for a covariance matrix
OUT: Writes and returns the covariance matrix=#
getModWhiteΣ!(args...; keyargs...) = error("Use modifiedwhiteΣ! instead")
function modifiedwhiteΣ!(xqr::FMQR{M}, ε::V,
  Σ::M = M(undef, xqr.K, xqr.K), dofcorrect::Float64 = 1.0
  )::M where {M<:AbstractMatrix, V<:AbstractVector}

  Λ::V = similar(ε)

  project!(xqr,Λ) #get the projection diagonal
  Λ .= (ε./(1.0.-Λ)).^2.0

  QRtinv::M = BLAS.gemm('N','T',xqr.Q,xqr.Rinv) #Q(R')^-1
  RinvQΛ::M = QRtinv'

  #loop to scale the matrix by Λ (modified part of modified white)
  #=@fastmath for j ∈ 1:xqr.N, i∈1:xqr.K
    RinvQΛ[i,j] *= Λ[j]
  end=#
  RinvQΛ *= Diagonal(Λ) #hopefully this is fast
  T::Type = eltype(M)

  BLAS.gemm!('N','N',dofcorrect,RinvQΛ, QRtinv,0.0,Σ)
  #final multiplicaiton and assignment
  return  Σ
end

#=helper method which extracts the required components from a linear model
IN: A linear model, optionally memory for the covariance matrix
OUT: Writes and allocates the covariance matrix=#
function modifiedwhiteΣ!(lin::FMLM{M},
  Σ::M = M(undef, lin.K, lin.K))::M where M <: AbstractMatrix

  return modifiedwhiteΣ!(lin.xqr, lin.ε, Σ, lin.N/lin.dof)
end

#=Simple test function to verify the QR-related optimizations
Output should be identical to the standard methods
IN: Independent variable X matrix
OUT: Returns the covariance matrix=#
getModWhiteΣSlow(args...; keyargs...) = error("Use modifiedwhiteΣslowinstead")
function modifiedwhiteΣslow(X::M, ε::V, dofcorrect::Float64) where {M<:AbstractMatrix, V<:AbstractVector}

  #do this since we only care about clarity (and not performance)
  #cu hack diabled
  #(M<:CuArray) && CuArrays.allowscalar(true)

  XXinv::M = Matrix(X' * X)\I

  P::M = X * (XXinv) * X'
  Σ::M = XXinv * X' * diagm(ε  ./ (1.0 .- diag(P))) .^ 2.0 * X * XXinv

  #(M<:CuArray) && CuArrays.allowscalar(false)

  return Σ .* dofcorrect

end

#Convenience method for above: IN: A linear model #OUT: A covariance matrix
modifiedwhiteΣslow(lin::FMLM, dofcorrect::Float64 = lin.N/lin.dof) = modifiedwhiteΣslow(
  lin.X, lin.Y.-lin.X * lin.β, dofcorrect)

  #########################getNeweyWest ##########OLS Only
  #NOTE: To match Sandwich NeweyWest in R, set prewhite = FALSE, adjust = TRUE
  getNeweyWest!(args...; keyargs...) = error("use neweywestΣ! instead")

  function neweywestΣ!(X::M, xqr::FMQR{M}, ε::V, lag::Int,
    Σ::M = M(undef, xqr.K, xqr.K),
    dofcorrect::Float64 = 1.0)::M where {M<:AbstractMatrix, V<:AbstractVector}

    #pre-allocate for the spectral matrix

    #NOTE: we need to switch away from gpu arrays due to the lagging
    #that is, views don't seem to work properly on gpu matrices
    T::Type = eltype(M)
    Rv::Matrix{T} = Matrix{T}(undef, xqr.K, xqr.K) #pre-allocate working matrix
    RRinv::Matrix{T} = BLAS.gemm('N', 'T', xqr.Rinv, xqr.Rinv) #this is equivelent to [X'X]^-1

    #need to multiply through by the error
    Xe::Matrix{T} = X .* ε
    ST::Matrix{T} = BLAS.gemm('T','N',T(1.0/xqr.N), Xe, Xe)
    for v::Int ∈ 1:lag
      #overwrites Rv with (1/N)R'R
      BLAS.gemm!('T', 'N', T(1.0/xqr.N), view(Xe, (v+1):(xqr.N), :),view(Xe, 1:(xqr.N-v), :), T(0.0), Rv)
      #Rv .= view(Xe, (v+1):(xqr.N), :)' * view(Xe, 1:(xqr.N-v), :) .* (1.0/xqr.N)
      ST += (lag + 1 - v)/(lag+1.) .* (Rv .+ Rv')
    end

    #this is [X'X]^-1S=[R'R]^-1S
    RRinvS::Matrix{T} = BLAS.gemm('N', 'N', RRinv, ST)

    #finally we have T[X'X]^-1S[X'X]^-1
    if M<:Matrix{Float64}
      BLAS.gemm!('N','N',Float64(xqr.N), RRinvS, RRinv, 0.0, Σ)
    else
      Σ = BLAS.gemm('N','N',T(xqr.N), RRinvS, RRinv)
    end

    Σ .*= dofcorrect
    #println("$(diag(Σ .* dofcorrect))")sasa
    return Σ
  end

  function neweywestΣ!(lin::FMLM{M,V}, lag::Int,
    Σ::M = M(undef, lin.K, lin.K))::M where {M<: AbstractMatrix, V<:AbstractVector}

    return neweywestΣ!(lin.X, lin.xqr, lin.ε, lag, Σ, lin.N/lin.dof)
  end

  #helper function in case the input requires a single argument function
  getNeweyWestFunc(args...; keyargs...) = error("use neweywestΣfunc instead")
  neweywestΣfunc(lag::Int) = (lin::FMLM)-> neweywestΣ!(lin, lag)

  #testing for NeweyWest
  #NOTE: assumes dofcorrect value
  getNeweyWestSlow(args...; keyargs...) = error("use neweywestΣslow instead")
  function neweywestΣslow(X::M, ε::V, lag::Int, dofcorrect::Float64
      )  where {M<:AbstractMatrix, V<:AbstractVector}

    N::Int, K::Int = size(X)
    #cu hack diabled
    #(M<:CuArray) && CuArrays.allowscalar(true) #don't care about performance here

    #first form the spectral matrix

    #initial value for spectral matrix
    xxt::Matrix = zeros(K,K)
    for t::Int ∈ 1:N
      xxt .+= (X[t, :] * X[t, :]') .* ε[t]^2
    end
    ST::Matrix =  xxt ./ N

    #make the rest of the spectral matrix
    for v::Int ∈ 1:lag
      xxt .= zeros(K,K)
      for t::Int ∈ (1+v):(N)
        xxt .+= (X[t, :] * X[t-v, :]') .* ε[t]*ε[t-v]
      end
      xxt ./= N
      ST .+= (lag + 1. - v)/(lag + 1.) .* (xxt .+ xxt')
    end

    XXinv::Matrix = Matrix(X' * X) \ I

    ret::M = N*XXinv*ST* XXinv
    #(M<:CuArray) && CuArrays.allowscalar(false) #don't care about performance here
    return ret * dofcorrect
  end

  neweywestΣslow(lin::FMLM, lag::Int) = neweywestΣslow(lin.X, lin.Y .- lin.X * lin.β, lag, lin.N/lin.dof)

  getNeweyWestFuncSlow(args...; keyargs...) = error("use neweywestΣslow instead")
  neweywestΣfuncslow(lag::Int)::Function = (lin::FMLM)-> neweywestΣslow(lin, lag)

  ######################Panel Newy-West#######################

function neweywestpanelΣ!(X::M, xqr::FMQR{M}, ε::V, clusters::FMClusters,
    lag::Int, Σ::M = M(undef, xqr.K, xqr.K))::M where {
    M<:AbstractMatrix, V<:AbstractVector}

  if length(clusters)==0
    error("No clusters detected in neweywestpanelΣ! - use neweywestΣ! instead")
  elseif length(clusters)>1
    error("Multi-way clusters not supported in neweywestpanelΣ!")
  end

  N::Int = length(ε)
  E::Type = eltype(M)

  K::Int = size(X,2)
  Sₜ::M = zeros(K,K) #holds teh central matrix

  temp::M = M(undef, K, K) #pre-allocate working matrix
  RRinv::M = BLAS.gemm('N', 'T', xqr.Rinv, xqr.Rinv) #this is equivelent to [X'X]^-1

  inddf::DataFrame = DataFrame(clusters=clusters[1], idx = 1:xqr.N)
  sindfs::GroupedDataFrame = groupby(inddf, :clusters)
  Nclusters::Int = length(sindfs)

  #iterate over all stocks
  for sindf::SubDataFrame ∈ sindfs
    Xₙ::SubArray = view(X, sindf.idx, :)
    εₙ::SubArray = view(ε, sindf.idx)

    #need to multiply through by the error
    T::Int = length(εₙ)
    Xₑ::M = Xₙ .* εₙ

    Sₜ .+= BLAS.gemm('T','N',  E(1.0/N), Xₑ, Xₑ)

    for v::Int ∈ 1:lag
      #overwrites temp with (1/N)R'R
      BLAS.gemm!('T', 'N',   E(1.0/N), view(Xₑ, (v+1):T, :),view(Xₑ, 1:(T-v), :),   E(0.0), temp)
      Sₜ .+= (lag + 1 - v)/(lag+1.) .* (temp .+ temp')
    end

  end
  #Sₜ = Matrix{Float64}(I,K,K)
  #this is [X'X]^-1S=[R'R]^-1S
  RRinvS::M = BLAS.gemm('N', 'N', RRinv, Sₜ)

  #display(RRinv)

  #finally we have T[X'X]^-1S[X'X]^-1
  BLAS.gemm!('N','N', E(N), RRinvS, RRinv, E(0.0), Σ)
  #println("$(diag(Σ .* dofCorrect))")sasa
  return Σ
end

function neweywestpanelΣ!(lin::FMLM{M,V}, lag::Int,
  Σ::M = M(undef, lin.K, lin.K))::M where {M<: AbstractMatrix, V<:AbstractVector}

  return neweywestpanelΣ!(lin.X, lin.xqr, lin.ε, lin.clusters, lag, Σ)
end

neweywestpanelΣfunc(lag::Int) = (lin::FMLM)-> neweywestpanelΣ!(lin, lag)

function neweywestpanelΣslow!(X::M, xqr::FMQR{M}, ε::V, clusters::FMClusters,
    lag::Int, Σ::M = M(undef, xqr.K, xqr.K))::M where {
    M<:AbstractMatrix, V<:AbstractVector}

  N::Int = size(X,1)
  K::Int = size(X,2)
  T::Type = eltype(M)

  if length(clusters)==0
    error("No clusters detected in neweywestpanelΣslow! - use neweywestΣslow! instead")
  elseif length(clusters)>1
    error("Multi-way clusters not supported in neweywestpanelΣslow!")
  end

  #pre-allocate
  Sₜ::M = zeros(K, K)

  #this method requires a sorted array
  #will do this via dataframes
  df::DataFrame = DataFrame(deepcopy(X), :auto)
  xnames = names(df)
  df.clusters = deepcopy(clusters[1])
  df.ε = deepcopy(ε)

  sort!(df, :clusters)
  X = M(df[!, xnames])
  ε = V(df.ε)
  clusters = df.clusters

  #helper function using the Lars method
  for t ∈ 1:N
    Sₜ .+= X[t,:] * X[t,:]' * ε[t]^2
  end

  for v::Int ∈ 1:lag
    for t ∈ (v+1):N
      (clusters[t] == clusters[t-v]) && (
        Sₜ .+= (lag + 1 - v)/(lag+1.) .* (X[t,:] * X[t-v,:]' .+ X[t-v,:] * X[t,:]') * ε[t] * ε[t-v])
    end
  end
  Sₜ ./= N
  #Sₜ = Matrix{Float64}(I,K,K)

  #the final step is to multiply out the var-covar sandwhich
  RRinv::M = BLAS.gemm('N', 'T', xqr.Rinv, xqr.Rinv) #this is equivelent to [X'X]^-1

  #this is [X'X]^-1S=[R'R]^-1S
  RRinvS::M= BLAS.gemm('N', 'N', RRinv, Sₜ)

  #finally we have T[X'X]^-1S[X'X]^-1
  BLAS.gemm!('N','N',T(N), RRinvS, RRinv, T(0.0), Σ)

  return Σ
end

function neweywestpanelΣslow!(lin::FMLM{M,V}, lag::Int,
  Σ::M = M(undef, lin.K, lin.K))::M where {M<: AbstractMatrix, V<:AbstractVector}

  return neweywestpanelΣslow!(lin.X, lin.xqr, lin.ε, lin.clusters, lag, Σ)
end

neweywestΣslowfunc(lag::Int) = (lin::FMLM)-> neweywestpanelΣslow!(lin, lag)

ext: jl

max_stars_repo_name: clintonTE/Finometrics

max_issues_repo_name: clintonTE/Finometrics

max_forks_repo_name: clintonTE/Finometrics


text:
#
# Build script
#
# author: Atsushi Sakai
#
println("build")

using Pkg
Pkg.add("PyPlot")
Pkg.add("DataFrames")
Pkg.add("NearestNeighbors")
Pkg.add("DataStructures")

max_stars_repo_name: yinflight/HybridAStarTrailer

max_issues_repo_name: GeoffGao/HybridAStarTrailer

max_forks_repo_name: GeoffGao/HybridAStarTrailer


text:
import Adapt

@inline fractional_index(x, xs, Δx) = @inbounds (x - xs[1]) / Δx

# Linear Lagrange polynomials
@inline ϕ₀(ξ) = 1 - ξ
@inline ϕ₁(ξ) = ξ

#####
##### Linear interpolation for regularly spaced profiles e.g. T(z).
#####

struct InterpolatedProfile{D, Z, Δ, T}
     data :: D
        z :: Z
       Δz :: Δ
    z_max :: T
end

@inline function _interpolate(profile::InterpolatedProfile, zᵢ)
    k = 1 + fractional_index(zᵢ, profile.z, profile.Δz)
    ξ, k = mod(k, 1), Base.unsafe_trunc(Int, k)

    # Ensure we don't go out of bounds.
    if zᵢ >= profile.z_max
        return @inbounds profile.data[k]
    else
        return @inbounds ϕ₀(ξ) * profile.data[k] + ϕ₁(ξ) * profile.data[k+1]
    end
end

@inline (profile::InterpolatedProfile)(z) = _interpolate(profile, z)

@inline function ∂z(profile::InterpolatedProfile, zᵢ)
    k = 1 + fractional_index(zᵢ, profile.z, profile.Δz)
    k = Base.unsafe_trunc(Int, k)
    zᵢ >= profile.z_max && (k = k - 1) # Ensure we don't go out of bounds.
    return @inbounds (profile.data[k+1] - profile.data[k]) / profile.Δz
end

Adapt.adapt_structure(to, profile::InterpolatedProfile) =
    InterpolatedProfile(Adapt.adapt(to, profile.data),
                        Adapt.adapt(to, profile.z),
                        Adapt.adapt(to, profile.Δz),
                        Adapt.adapt(to, profile.z_max))

#####
##### Linear interpolated for regularly spaced profile time series, e.g. T(z, t).
#####

struct InterpolatedProfileTimeSeries{D, Z, T, ΔZ, ΔT, FT}
     data :: D
        z :: Z
        t :: T
       Δz :: ΔZ
       Δt :: ΔT
    z_max :: FT
    t_max :: FT
end

@inline function _interpolate(profile::InterpolatedProfileTimeSeries, zᵢ, tᵢ)
    k = 1 + fractional_index(zᵢ, profile.z, profile.Δz)
    n = 1 + fractional_index(tᵢ, profile.t, profile.Δt)

    ξ, k = mod(k, 1), Base.unsafe_trunc(Int, k)
    η, n = mod(n, 1), Base.unsafe_trunc(Int, n)

    # Ensure we don't go out of bounds.
    if zᵢ >= profile.z_max && tᵢ >= profile.t_max
        return @inbounds profile.data[k, n]
    elseif zᵢ >= profile.z_max
        return @inbounds ϕ₀(η) * profile.data[k, n] + ϕ₁(η) * profile.data[k, n+1]
    elseif tᵢ >= profile.t_max
        return @inbounds ϕ₀(ξ) * profile.data[k, n] + ϕ₁(ξ) * profile.data[k+1, n]
    else
        return @inbounds (  ϕ₀(ξ) * ϕ₀(η) * profile.data[k,   n  ]
                          + ϕ₀(ξ) * ϕ₁(η) * profile.data[k,   n+1]
                          + ϕ₁(ξ) * ϕ₀(η) * profile.data[k+1, n  ]
                          + ϕ₁(ξ) * ϕ₁(η) * profile.data[k+1, n+1])
    end
end

@inline (profile::InterpolatedProfileTimeSeries)(z, t) = _interpolate(profile, z, t)

@inline function ∂z(profile::InterpolatedProfileTimeSeries, zᵢ, tᵢ)
    k = 1 + fractional_index(zᵢ, profile.z, profile.Δz)
    n = 1 + fractional_index(tᵢ, profile.t, profile.Δt)

    ξ, k = mod(k, 1), Base.unsafe_trunc(Int, k)
    η, n = mod(n, 1), Base.unsafe_trunc(Int, n)

    # Ensure we don't go out of bounds.
    zᵢ >= profile.z_max && (k = k - 1)
    tᵢ >= profile.t_max && (n = n - 1)

    dataᵏ   = @inbounds ϕ₀(η) * profile.data[k,   n] + ϕ₁(η) * profile.data[k,   n+1]
    dataᵏ⁺¹ = @inbounds ϕ₀(η) * profile.data[k+1, n] + ϕ₁(η) * profile.data[k+1, n+1]

    return (dataᵏ⁺¹ - dataᵏ) / profile.Δz
end

@inline function ∂t(profile::InterpolatedProfileTimeSeries, zᵢ, tᵢ)
    k = 1 + fractional_index(zᵢ, profile.z, profile.Δz)
    n = 1 + fractional_index(tᵢ, profile.t, profile.Δt)

    ξ, k = mod(k, 1), Base.unsafe_trunc(Int, k)
    η, n = mod(n, 1), Base.unsafe_trunc(Int, n)

    # Ensure we don't go out of bounds.
    zᵢ >= profile.z_max && (k = k - 1)
    tᵢ >= profile.t_max && (n = n - 1)

    dataⁿ   = @inbounds ϕ₀(ξ) * profile.data[k,   n] + ϕ₁(ξ) * profile.data[k+1,   n]
    dataⁿ⁺¹ = @inbounds ϕ₀(ξ) * profile.data[k, n+1] + ϕ₁(ξ) * profile.data[k+1, n+1]

    return (dataⁿ⁺¹ - dataⁿ) / profile.Δt
end

Adapt.adapt_structure(to, profile::InterpolatedProfileTimeSeries) =
    InterpolatedProfileTimeSeries(Adapt.adapt(to, profile.data),
                                  Adapt.adapt(to, profile.z),
                                  Adapt.adapt(to, profile.t),
                                  Adapt.adapt(to, profile.Δz),
                                  Adapt.adapt(to, profile.Δt),
                                  Adapt.adapt(to, profile.z_max),
                                  Adapt.adapt(to, profile.t_max))


text:
using BivarNormalProb
using Base.Test

# test orthants probability. See section 2.2.1 of
# Computation of Multivariate Normal and t Probabilities
# Alan Genz • Frank Bretz
# Springer-Verlag,  2009
@test bvnu(0,0,1/2) == (1/4)+asin(1/2)/(2pi)


text:
## logistic regression

struct analysis_logistic_Binom_NIP_Optim
    formula::FormulaTerm
    modelClass::String
    LikelihoodMod::String
    PriorMod::String
    Link::String
    ComputeMethod::String
    fit
    beta
    LogLike::Float64
    LogPost::Float64
    AIC::Float64
    BIC::Float64
    #R_sqr::Float64
    #Adjusted_R_sqr::Float64
    #sigma::Float64
    #Cooks_distance
end

function logistic_reg(formula::FormulaTerm,data,Link::String="LogitLink")

    modelClass = "LogisticReg";
    LikelihoodMod="Binomial";
    PriorMod="NIP";
    ComputeMethod="Optimization";
    formula = apply_schema(formula, schema(formula, data));
    y, X = modelcols(formula, data);
    fm_frame=ModelFrame(formula,data);
    X=modelmatrix(fm_frame);
    p = size(X, 2);
    n = size(X, 1);

    if (Link=="LogitLink")
        res = glm(formula,data,Binomial(), LogitLink());

    elseif (Link=="ProbitLink")
        res = glm(formula,data,Binomial(), ProbitLink());

    elseif (Link=="CauchitLink")
        res = glm(formula,data,Binomial(), CauchitLink());

    elseif (Link=="CloglogLink")
        res = glm(formula,data,Binomial(), CloglogLink());

    else
        println("This link function is not part of logistic regression family.")
        println("-------------------------------------------------------------")
    end

    fit = coeftable(res)
    beta_hat = coef(res)
    logLike = GLM.loglikelihood(res)
    LogPost = logLike
    npar = p;
    AIC = 2*npar - 2*logLike
    BIC = log(n)*npar - 2*logLike

    ans = analysis_logistic_Binom_NIP_Optim(formula,modelClass,LikelihoodMod
            ,PriorMod,Link,ComputeMethod,fit,beta_hat
            ,logLike,LogPost,AIC,BIC)
    ans

end

function logistic_reg_predicts(obj,newdata::DataFrame)

    formula = obj.formula;
    fm_frame=ModelFrame(formula,newdata);
    X=modelmatrix(fm_frame);
    beta = obj.beta
    z = X*beta;

    if (obj.Link=="LogitLink")
      p = exp.(z) ./ (1 .+ exp.(z));

    elseif (obj.Link=="ProbitLink")
      p = Probit_Link.(z);

    elseif (obj.Link=="CauchitLink")
      p = Cauchit_Link.(z);

    elseif (obj.Link=="CloglogLink")
      p = Cloglog_Link.(z);

    else
        println("This link function is not part of logistic regression family.")
        println("-------------------------------------------------------------")
    end
    p
end

## logistic regression with Ridge Prior

function logistic_reg_internal(formula::FormulaTerm, data::DataFrame, link_function, PriorMod::Prior_Ridge, h::Float64=0.5,sim_size::Int64=10000)

    formula = apply_schema(formula, schema(formula, data));
    y, X = modelcols(formula, data);
    fm_frame=ModelFrame(formula,data);
    X=modelmatrix(fm_frame);

    @model LogisticReg(X, y) = begin
        p = size(X, 2);
        n = size(X, 1);
        #priors
        λ~InverseGamma(h,h)
        β ~ filldist(Normal(0,λ), p)

        z = X*β

        ## Link Function

        prob = link_function.(z)

        #likelihood
        for i = 1:n
            y[i] ~ Bernoulli(prob[i])
        end

    end;

    LogisticReg_model=LogisticReg(X,y);

    chain = sample(CRRao_rng, LogisticReg_model, NUTS(), sim_size);

    summaries, quantiles = describe(chain);

    ans = MCMC_chain(chain,summaries,quantiles)

    ans

end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Logit, PriorMod::Prior_Ridge, h::Float64=0.5,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Logit_Link, PriorMod, h, sim_size)
end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Probit, PriorMod::Prior_Ridge, h::Float64=0.5,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Probit_Link, PriorMod, h, sim_size)
end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Cloglog, PriorMod::Prior_Ridge, h::Float64=0.5,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Cloglog_Link, PriorMod, h, sim_size)
end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Cauchit, PriorMod::Prior_Ridge, h::Float64=0.5,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Cauchit_Link, PriorMod, h, sim_size)
end

## logistic regression with Laplace Prior

function logistic_reg_internal(formula::FormulaTerm,data::DataFrame, link_function,PriorMod::Prior_Laplace,h::Float64=0.5,sim_size::Int64=10000)

    formula = apply_schema(formula, schema(formula, data));
    y, X = modelcols(formula, data);
    fm_frame=ModelFrame(formula,data);
    X=modelmatrix(fm_frame);

    @model LogisticReg(X, y) = begin
        p = size(X, 2);
        n = size(X, 1);
        #priors
        λ~InverseGamma(h,h)
        β ~ filldist(Laplace(0,λ), p)

        z = X*β

        ## Link Function

        prob = Logit_Link.(z)

        #likelihood
        for i = 1:n
            y[i] ~ Bernoulli(prob[i])
        end

    end;

    LogisticReg_model=LogisticReg(X,y);

    chain = sample(CRRao_rng, LogisticReg_model, NUTS(), sim_size);

    summaries, quantiles = describe(chain);

    ans = MCMC_chain(chain,summaries,quantiles)

    return ans

end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Logit, PriorMod::Prior_Laplace, h::Float64=0.5,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Logit_Link, PriorMod, h, sim_size)
end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Probit, PriorMod::Prior_Laplace, h::Float64=0.5,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Probit_Link, PriorMod, h, sim_size)
end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Cloglog, PriorMod::Prior_Laplace, h::Float64=0.5,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Cloglog_Link, PriorMod, h, sim_size)
end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Cauchit, PriorMod::Prior_Laplace, h::Float64=0.5,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Cauchit_Link, PriorMod, h, sim_size)
end

## logistic regression with Cauchy Prior

function logistic_reg_internal(formula::FormulaTerm,data,link_function,PriorMod::Prior_Cauchy,h::Float64,sim_size::Int64)

    formula = apply_schema(formula, schema(formula, data));
    y, X = modelcols(formula, data);
    fm_frame=ModelFrame(formula,data);
    X=modelmatrix(fm_frame);

    @model LogisticReg(X, y) = begin
        p = size(X, 2);
        n = size(X, 1);
        #priors
        λ~ Truncated(TDist(1),0,Inf)
        β ~ filldist(TDist(1)*λ, p)

        z = X*β

        ## Link Function

        prob = Logit_Link.(z)

        #likelihood
        for i = 1:n
            y[i] ~ Bernoulli(prob[i])
        end

    end;

    LogisticReg_model=LogisticReg(X,y);

    chain = sample(CRRao_rng, LogisticReg_model, NUTS(), sim_size);

    summaries, quantiles = describe(chain);

    ans = MCMC_chain(chain,summaries,quantiles)

    return ans

end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Logit, PriorMod::Prior_Cauchy, h::Float64=1.0,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Logit_Link, PriorMod, h, sim_size)
end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Probit, PriorMod::Prior_Cauchy, h::Float64=1.0,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Probit_Link, PriorMod, h, sim_size)
end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Cloglog, PriorMod::Prior_Cauchy, h::Float64=1.0,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Cloglog_Link, PriorMod, h, sim_size)
end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Cauchit, PriorMod::Prior_Cauchy, h::Float64=1.0,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Cauchit_Link, PriorMod, h, sim_size)
end

## logistic regression with TDist Prior

function logistic_reg_internal(formula::FormulaTerm,data,link_function,PriorMod::Prior_TDist,h::Float64,sim_size::Int64)

    formula = apply_schema(formula, schema(formula, data));
    y, X = modelcols(formula, data);
    fm_frame=ModelFrame(formula,data);
    X=modelmatrix(fm_frame);

    @model LogisticReg(X, y) = begin
        p = size(X, 2);
        n = size(X, 1);
        #priors
        λ ~ InverseGamma(h,h)
        ν ~ InverseGamma(h,h)
        β ~ filldist(TDist(ν)*λ, p)

        z = X*β

        ## Link Function

        prob = Logit_Link.(z)

        #likelihood
        for i = 1:n
            y[i] ~ Bernoulli(prob[i])
        end

    end;

    LogisticReg_model=LogisticReg(X,y);

    chain = sample(CRRao_rng, LogisticReg_model, NUTS(), sim_size);

    summaries, quantiles = describe(chain);

    ans = MCMC_chain(chain,summaries,quantiles)

    return ans

end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Logit, PriorMod::Prior_TDist, h::Float64=1.0,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Logit_Link, PriorMod, h, sim_size)
end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Probit, PriorMod::Prior_TDist, h::Float64=1.0,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Probit_Link, PriorMod, h, sim_size)
end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Cloglog, PriorMod::Prior_TDist, h::Float64=1.0,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Cloglog_Link, PriorMod, h, sim_size)
end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Cauchit, PriorMod::Prior_TDist, h::Float64=1.0,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Cauchit_Link, PriorMod, h, sim_size)
end

## logistic regression with Uniform Prior

function logistic_reg_internal(formula::FormulaTerm,data,link_function,PriorMod::Prior_Uniform,h::Float64=0.5,sim_size::Int64=10000)

    formula = apply_schema(formula, schema(formula, data));
    y, X = modelcols(formula, data);
    fm_frame=ModelFrame(formula,data);
    X=modelmatrix(fm_frame);

    @model LogisticReg(X, y) = begin
        p = size(X, 2);
        n = size(X, 1);
        #priors
        v ~ InverseGamma(h,h)
        β ~ filldist(Uniform(-v,v), p)

        z = X*β

        ## Link Function

        prob = Logit_Link.(z)

        #likelihood
        for i = 1:n
            y[i] ~ Bernoulli(prob[i])
        end

    end;

    LogisticReg_model=LogisticReg(X,y);

    chain = sample(CRRao_rng, LogisticReg_model, NUTS(), sim_size);

    summaries, quantiles = describe(chain);

    ans = MCMC_chain(chain,summaries,quantiles)

    return ans

end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Logit, PriorMod::Prior_Uniform, h::Float64=1.0,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Logit_Link, PriorMod, h, sim_size)
end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Probit, PriorMod::Prior_Uniform, h::Float64=1.0,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Probit_Link, PriorMod, h, sim_size)
end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Cauchit, PriorMod::Prior_Uniform, h::Float64=1.0,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Cauchit_Link, PriorMod, h, sim_size)
end

function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::Cloglog, PriorMod::Prior_Uniform, h::Float64=1.0,sim_size::Int64=10000)
    logistic_reg_internal(formula, data, Cloglog_Link, PriorMod, h, sim_size)
end


text:
using Test, Pkg

ENV["DATADEPS_ALWAYS_ACCEPT"] = true
ENV["JULIA_LOG_LEVEL"] = "WARN"

function include_test(_module)
    println("Starting tests for $_module")
    t = @elapsed include(joinpath(_module, "runtests.jl"))
    println("Completed tests for $_module, $(round(Int, t)) seconds elapsed")
    return nothing
end

@testset "CLIMA" begin
    all_tests = isempty(ARGS) || "all" in ARGS ? true : false

    for submodule in [
        "InputOutput",
        "Utilities",
        "Common",
        "Arrays",
        "Atmos",
        "Numerics",
        "Diagnostics",
        "Ocean",
        "Driver",
    ]
        if all_tests || "$submodule" in ARGS || "CLIMA" in ARGS
            include_test(submodule)
        end
    end

end


text:
if !isdefined(Main, :CustomStencil)
    include("../module/CustomStencil.jl")
    using Main.CustomStencil
end
include("misc.jl")
include("cpu_stencils.jl")

## Stencil Definition
coefs = Float64.([1; 0.5; 0.25; 0.125; 0.06])

nv_star_stencil = @def_stencil_expression -c[0]D[x,y,z] + @sum(i, 1, 4, c[i]*(
            D[x+i,y,z] + D[x,y+i,z] + D[x,y,z+i] +
            D[x-i,y,z] + D[x,y-i,z] + D[x,y,z-i]))

star_stencil = @def_stencil_expression -v*c[0]D[x,y,z] + v*@sum(i, 1, 4, c[i]*(
            D[x+i,y,z] + D[x,y+i,z] + D[x,y,z+i] +
            D[x-i,y,z] + D[x,y-i,z] + D[x,y,z-i]))

v_def = CreateStencilDefinition(star_stencil, coefs)
nv_def = CreateStencilDefinition(nv_star_stencil, coefs)
st_inst = NewStencilInstance(v_def, nv_def)

## Input Data size Definition
radius = 4
nx = 5
ny = 5
nz = 5
dx = 1<<(nx)
dy = 1<<(ny)
dz = 1<<(nz)

t_steps = 4
## GPU Implementation
original_data = CreateData(dx,dy,dz)
data = PadData(radius, original_data)
vsq = ones(Float32, dx, dy, dz)
vsq[10:15,:,:] .= 0.1
dout = ApplyStencil(st_inst, data, t_steps,vsq=vsq)
gpu_out = Array(dout)

## CPU Implementation
cpu_out =  cpu_star_stencil_v_nv(original_data, vsq, radius, coefs, t_steps)

## Check Result
mean_er = sum(abs.(cpu_out .- gpu_out))./(reduce(*, size(gpu_out)))
@assert(sum(abs.(cpu_out .- gpu_out))./(reduce(*, size(gpu_out))) < 0.001)
## Visual Check
# using Plots
# shift = -10
# Int(round(dz/2))+shift
# heatmap(cout[:,:,Int(round(dz/2))+shift])
# heatmap(o[:,:,Int(round(dz/2))+shift])


text:
# Julia script to run set of models in tuple modeltypes on region
t0 = time()

@everywhere using DelimitedFiles
@everywhere using StatsBase
@everywhere using Dates
@everywhere using LinearAlgebra
@everywhere using StatsBase

const txtstr = ".txt"

@everywhere include("SICRnew.jl")

function getdatafiles(datafolder)
    datafiles = readdir("data")
    datafiles = datafiles[occursin.("covid_timeseries",datafiles)]
    regions = Array{String,1}(undef,0)
    for datafile in datafiles
        push!(regions,replace(replace(datafile,".csv"=>""),"covid_timeseries_"=>""))
    end
    return datafiles,regions
end

function getprior(modelname)
    if modelname == "r"
        return [.25,.1,.1,.1,.1]
    end
    if modelname == "rq" || modelname == "rf"
        return [.25,.1,.1,.1,.1,.1]
    end
    if modelname == "rm"
        return [.25,.1,.1,.1,.1,5.,.1]
    end
    if modelname == "rqm" || modelname == "rfm"
        return [.25,.1,.1,.1,.1,.1,5.,.1]
    end
    if modelname == "rqfm"
        return [.25,.1,.1,.1,.1,.1,.1,5.,.1]
    end
    if modelname == "u"
        return [.25,.1,.1,.1]
    end
    if modelname == "uq" || modelname == "uf"
        return [.25,.1,.1,.1,.1]
    end
    if modelname == "um"
        return [.25,.1,.1,.1,5.,.1]
    end
    if modelname == "uqm" || modelname == "ufm"
        return [.25,.1,.1,.1,.1,5.,.1]
    end
    if modelname == "uqfm"
        return [.25,.1,.1,.1,.1,.1,5.,.1]
    end
end

name = Dict(sicr! => "r", sicrq! => "rq", sicrf! => "rf", sicrqf! => "rqf", sicrm! => "rm", sicrqm! => "rqm", sicrfm! => "rfm", sicrqfm! => "rqfm",sicu! => "u", sicuq! => "uq", sicuf! => "uf", sicuq! => "uq", sicuqf! => "uqf", sicum! => "um", sicuqm! => "uqm", sicufm! => "ufm", sicuqfm! => "uqfm")

datestrin = "2020-04-04"
datestrout = "2020-04-04"
println(datestrin,"->",datestrout)
if ~isdir("Results")
    mkdir("Results")
end
pathin = "Results/$datestrin/"
pathout = "Results/$datestrout/"
if ~isdir(pathout)
    mkdir(pathout)
end

datafiles,regions = getdatafiles("data")

colsr = [5,3,4]
colsu = [5,6]

N = 1e7
total = Int(1e7)
temp = 1.

modeltype = (sicr!,sicrq!,sicrf!,sicrm!,sicu!,sicuq!,sicuf!,sicum!)

out = Array{Future,2}(undef,length(modeltype),length(datafiles))

for j in eachindex(datafiles)
    datafile = datafiles[j]
    df = DataFrame(CSV.File("data/" * datafile))
    region = regions[j]
    println(region)

    for i in eachindex(modeltype)

        modrun = modeltype[i]

        if occursin("u",name[modeltype[i]])
            data = sicuprep(df)
            cols = colsu
        else
            data = df
            cols = colsr
        end

        infile = pathin * "maxll" * "_" * region * "_" * name[modrun] * txtstr

        if isfile(infile) && ~isempty(read(infile))
            p = readdlm(infile)[2:end]
        else
            p = getprior(name[modrun])
            println(name[modrun],": No prior")
        end

        out[i,j] = @time @spawn mcmc(data,cols,p,N,total,modrun,temp)

        println("maxll: ",out[i][4],", Accept: ",out[i][3])

    end
end

for j in eachindex(datafiles), i in eachindex(modeltype)

    outtuple = fetch(out[i,j])
    region = regions[j]
    modrun = modeltype[i]

    mlstr = pathout * "maxll" * "_" * region * "_" * name[modrun] * txtstr
    llstr = pathout * "llsamples" * "_" * region * "_" * name[modrun] * txtstr
    pstr = pathout * "pl" * "_" * region * "_" * name[modrun] * txtstr

    fml=open(mlstr,"w")
    fml=open(mlstr,"a")
    writedlm(fml,[outtuple[4];outtuple[5]])
    close(fml)

    fchil=open(llstr,"w")
    fchil=open(llstr,"a")
    writedlm(fchil,outtuple[1])
    close(fchil)

    fpl=open(pstr,"w")
    fpll=open(pstr,"a")
    writedlm(fpl,outtuple[1])
    close(fpl)

end

println(" Time: ",time()-t0)

max_stars_repo_name: cannin/covid-sicr-test

max_issues_repo_name: cannin/covid-sicr-test

max_forks_repo_name: cannin/covid-sicr-test


text:
@testset "composite: RGB" begin
    @testset "clear: $T" for T in Ts
        @test CompositeClear(RGBA{T}(1, 0.75, 0, 0.6), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0, 0, 0)
    end
    @testset "copy: $T" for T in Ts
        @test CompositeCopy(RGBA{T}(1, 0.75, 0, 0.6), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0.5, 1, 0.6)
    end
    @testset "destination: $T" for T in Ts
        @test CompositeDestination(RGBA{T}(1, 0.75, 0, 0.6), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(1, 0.75, 0, 0.6)
    end
    @testset "source-over: $T" for T in Ts
        @test CompositeSourceOver(RGBA{T}(1, 0.75, 0, 0.6), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(2/7, 4/7, 5/7, 0.84)
        @test CompositeSourceOver(RGBA{T}(1, 0.75, 0, 0), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0.5, 1, 0.6)
        @test CompositeSourceOver(RGBA{T}(1, 0.75, 0, 1), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0.4, 0.6, 0.6, 1)
    end
    @testset "destination-over: $T" for T in Ts
        @test CompositeDestinationOver(RGBA{T}(1, 0.75, 0, 0.6), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(5/7, 19/28, 2/7, 0.84)
        @test CompositeDestinationOver(RGBA{T}(1, 0.75, 0, 0), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0.5, 1, 0.6)
        @test CompositeDestinationOver(RGBA{T}(1, 0.75, 0, 1), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(1, 0.75, 0, 1)
    end
    @testset "source-in: $T" for T in Ts
        @test CompositeSourceIn(RGBA{T}(1, 0.75, 0, 0.6), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0.5, 1, 0.36)
        @test CompositeSourceIn(RGBA{T}(1, 0.75, 0, 0), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0, 0, 0)
        @test CompositeSourceIn(RGBA{T}(1, 0.75, 0, 1), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0.5, 1, 0.6)
    end
    @testset "destination-in: $T" for T in Ts
        @test CompositeDestinationIn(RGBA{T}(1, 0.75, 0, 0.6), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(1, 0.75, 0, 0.36)
        @test CompositeDestinationIn(RGBA{T}(1, 0.75, 0, 0), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0, 0, 0)
        @test CompositeDestinationIn(RGBA{T}(1, 0.75, 0, 1), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(1, 0.75, 0, 0.6)
    end
    @testset "source-out: $T" for T in Ts
        @test CompositeSourceOut(RGBA{T}(1, 0.75, 0, 0.6), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0.5, 1, 0.24)
        @test CompositeSourceOut(RGBA{T}(1, 0.75, 0, 0), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0.5, 1, 0.6)
        @test CompositeSourceOut(RGBA{T}(1, 0.75, 0, 1), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0, 0, 0)
    end
    @testset "destination-out: $T" for T in Ts
        @test CompositeDestinationOut(RGBA{T}(1, 0.75, 0, 0.6), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(1, 0.75, 0, 0.24)
        @test CompositeDestinationOut(RGBA{T}(1, 0.75, 0, 0), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0, 0, 0)
        @test CompositeDestinationOut(RGBA{T}(1, 0.75, 0, 1), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(1, 0.75, 0, 0.4)
    end
    @testset "source-atop: $T" for T in Ts
        @test CompositeSourceAtop(RGBA{T}(1, 0.75, 0, 0.6), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0.4, 0.6, 0.6, 0.6)
        @test CompositeSourceAtop(RGBA{T}(1, 0.75, 0, 0), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0, 0, 0)
        @test CompositeSourceAtop(RGBA{T}(1, 0.75, 0, 1), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0.4, 0.6, 0.6, 1)
    end
    @testset "destination-atop: $T" for T in Ts
        @test CompositeDestinationAtop(RGBA{T}(1, 0.75, 0, 0.6), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0.6, 0.65, 0.4, 0.6)
        @test CompositeDestinationAtop(RGBA{T}(1, 0.75, 0, 0), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0.5, 1, 0.6)
        @test CompositeDestinationAtop(RGBA{T}(1, 0.75, 0, 1), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(1, 0.75, 0, 0.6)
    end
    @testset "xor: $T" for T in Ts
        @test CompositeXor(RGBA{T}(1, 0.75, 0, 0.6), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0.5, 0.625, 0.5, 0.48)
        @test CompositeXor(RGBA{T}(1, 0.75, 0, 0), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0.5, 1, 0.6)
        @test CompositeXor(RGBA{T}(1, 0.75, 0, 1), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(1, 0.75, 0, 0.4)
    end
    @testset "lighter: $T" for T in Ts
        @test CompositeLighter(RGBA{T}(1, 0.75, 0, 0.6), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0.6, 0.75, 0.6, 1)
        @test CompositeLighter(RGBA{T}(1, 0.75, 0, 0), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(0, 0.5, 1, 0.6)
        @test CompositeLighter(RGBA{T}(1, 0.75, 0, 1), RGBA{T}(0, 0.5, 1, 0.6)) ≈ RGBA{T}(1, 1, 0.6, 1)
    end
end

@testset "composite: opaque RGBs" begin
    c1, c2 = RGB{Float32}(1, 0.75, 0), RGB{Float64}(0, 0.5, 1)
    t1, t2 = ARGB(c1), ARGB(c2)
    m = BlendLuminosity
    @testset "$(keyword(op))" for op in (CompositeCopy,
                                         CompositeDestination,
                                         CompositeSourceOver,
                                         CompositeDestinationOver,
                                         CompositeSourceIn,
                                         CompositeDestinationIn,
                                         CompositeSourceAtop,
                                         CompositeDestinationAtop,
                                         CompositeLighter)
        @test ARGB(op(c1, c2, mode=m)) === op(t1, t2, mode=m)
    end
    @testset "$(keyword(op))" for op in (CompositeClear,
                                         CompositeSourceOut,
                                         CompositeDestinationOut,
                                         CompositeXor)
        @test_throws MethodError op(c1, c2, mode=m)
    end
end

@testset "composite: RGB over opaque RGB" begin
    c1, c2 = RGB{Float32}(1, 0.75, 0), RGB{Float64}(0, 0.5, 1)
    t1, t2 = ARGB(c1), ARGB(c2, 0.25)
    m = BlendLuminosity
    @testset "$(keyword(op))" for op in (CompositeDestination,
                                         CompositeSourceOver,
                                         CompositeDestinationOver,
                                         CompositeSourceAtop,
                                         CompositeLighter)
        c = op(c1, c2, opacity=0.25, mode=m)
        @test ARGB(c) === ARGB(op(c1, t2, mode=m)) === op(t1, t2, mode=m)
    end
    @testset "$(keyword(op))" for op in (CompositeClear,
                                         CompositeCopy,
                                         CompositeSourceIn,
                                         CompositeDestinationIn,
                                         CompositeSourceOut,
                                         CompositeDestinationOut,
                                         CompositeDestinationAtop,
                                         CompositeXor)
        @test_throws MethodError op(c1, c2, opacity=0.25, mode=m)
        @test_throws MethodError op(c1, t2, mode=m)
    end
end


text:
module FDFDViz

using FDFD, PyPlot, PyCall, AxisArrays, Printf

export plot_field, plot_device, add_scalebar, add_wavelengthbar

"    plot_field(field::Field; cbar::Bool=false, funcz=real)"
function plot_field(field::Field; cbar::Bool=false, funcz=real, component=nothing)
	fig, ax = subplots(1);
	plot_field(ax, field; cbar=cbar, funcz=funcz, component=component);
	return ax
end

"    plot_field(ax::PyObject, field::Field; cbar::Bool=false, funcz=real)"
function plot_field(ax::PyObject, field::Field; cbar::Bool=false, funcz=real, component=nothing)
	# Set defaults
	isa(field, FieldTM) && component == nothing && (component=:Ez);
	isa(field, FieldTE) && component == nothing && (component=:Hz);
	isa(field, FieldAll) && component == nothing && (component=:Ex);

	Z = funcz.(permutedims(field[:, :, component], (2,1)))

	if funcz == abs
		vmin = 0;
		vmax = +maximum(abs.(Z));
		cmap = "magma";
	elseif funcz == real
		Zmx  = maximum(abs.(Z));
		vmin = -maximum(Zmx);
		vmax = +maximum(Zmx);
		cmap = "RdBu";
	else
		error("Unknown function specified.");
	end

	extents = [ field.grid.bounds[1][1], field.grid.bounds[2][1],
	            field.grid.bounds[1][2], field.grid.bounds[2][2] ];

	mappable = ax[:imshow](Z, cmap=cmap, extent=extents, origin="lower", vmin=vmin, vmax=vmax);
	cbar && colorbar(mappable, ax=ax, label=L"$\vert E \vert$");
	ax[:set_xlabel](L"$x$");
	ax[:set_ylabel](L"$y$");
	ax[:set_title](@sprintf("ω/2π = %.2f THz", real(field.ω/2π/1e12)));
end

"    plot_device(device::AbstractDevice; outline::Bool=false, lc::String=\"k\", lcm::String=\"k\")"
function plot_device(device::AbstractDevice; outline::Bool=false, lc::String="k", lcm::String="k")
	fig, ax = subplots(1);
	plot_device(ax, device; outline=outline, lc=lc, lcm=lcm);
	return ax
end

"    plot_device(ax::PyObject, device::AbstractDevice; outline::Bool=true, lc::String=\"k\", lcm::String=\"k\")"
function plot_device(ax::PyObject, device::AbstractDevice; outline::Bool=true, lc::String="k", lcm::String="k")
	Z = real.(permutedims(device.ϵᵣ, (2,1)));
	Zi = imag.(permutedims(device.ϵᵣ, (2,1)));

	if outline
		ax[:contour](xc(device.grid), yc(device.grid), Z, linewidths=0.25, colors=lc);
		axis("image");
	else
		extents = [ device.grid.bounds[1][1], device.grid.bounds[2][1],
		            device.grid.bounds[1][2], device.grid.bounds[2][2] ];
		ax[:imshow](Z, extent=extents, origin="lower", cmap="YlGnBu");
	end

	ax[:contour](xc(device.grid), yc(device.grid), Zi, colors=lc, linewidths=0.5, linestyles=":");

	if isa(device, ModulatedDevice)
		Z2 = abs.(permutedims(device.Δϵᵣ, (2,1)));
		ax[:contour](xc(device.grid), yc(device.grid), Z2, levels=1, linewidths=0.5, colors=lcm);
	end
	ax[:set_xlabel](L"$x$");
	ax[:set_ylabel](L"$y$");
end

"    add_scalebar(ax::PyObject, pt::Point; fc::String=\"k\", width::Number=1.0, height::Number=0.125, hide::Bool=true)"
function add_scalebar(ax::PyObject, pt::Point; fc::String="k", width::Number=1.0, height::Number=0.125, hide::Bool=true)
    ax[:add_patch](matplotlib[:patches][:Rectangle](pt-[width/2,height/2],width,height,fc=fc));
    ax[:annotate](xy=pt+[0,height/2],s=@sprintf("%d",width)*L" $\mu$m",fontsize="smaller",ha="center",va="bottom",multialignment="center",color=fc);
    if hide
    	ax[:axes][:get_xaxis]()[:set_visible](false);
    	ax[:axes][:get_yaxis]()[:set_visible](false);
    end
end

"    add_wavelengthbar(ax::PyObject, field::Field, pt::Point; fc::String=\"k\", height::Number=0.125)"
function add_wavelengthbar(ax::PyObject, field::Field, pt::Point; fc::String="k", height::Number=0.125)
	λ₀ = 2π*c₀/real(field.ω)/1e-6
    ax[:add_patch](matplotlib[:patches][:Rectangle](pt-[λ₀/2,height/2],λ₀,height,fc=fc));
    ax[:annotate](xy=pt+[0,height/2],s=L"$\lambda_0$",fontsize="smaller",ha="center",va="bottom",multialignment="center",color=fc);
end

end


text:
using Sundials
using Base.Test

examples_path = joinpath(dirname(@__FILE__), "..", "examples")

# run cvode example
println("== start cvode Roberts example (simplified)")
let
include(joinpath(examples_path, "cvode_Roberts_simplified.jl"))
println("result at t=$(t[end]):")
println(res[:,end], "\n")
end

println("== start cvode Roberts example")
let
include(joinpath(examples_path, "cvode_Roberts_dns.jl"))
end

# run ida examples
println("== start ida_Roberts example (simplified)")
let
include(joinpath(examples_path, "ida_Roberts_simplified.jl"))
end

println("== start ida_Roberts example")
let
include(joinpath(examples_path, "ida_Roberts_dns.jl"))
println("result at t=$(t[end]):")
println(yout[:,end], "\n")
end

println("== start ida_Heat2D example")
let
include(joinpath(examples_path, "ida_Heat2D.jl"))
println("result at t=$(t[end]):")
println(yout[:,end], "\n")
end

#= requires Grid package
println("== start ida_Cable example")
let
include(joinpath(examples_path, "ida_Cable.jl"))
end
=#

# run kinsol example
println("== start kinsol example (simplified)")
let
include(joinpath(examples_path, "kinsol_mkin_simplified.jl"))
println("solution:")
println(res)
residual = ones(2)
sysfn(res, residual)
println("residual:")
println(residual, "\n")
end

println("== start kinsol example")
let
include(joinpath(examples_path, "kinsol_mkinTest.jl"))
@test abs(minimum(residual)) < 1e-5
end