text:
function pack(X, U)
    z = zeros(num_var - p)
    for t = 1:T
        z[x_idx[t]] = X[t]
        t == T && continue
        z[u_idx[t]] = U[t]
    end
    return z
end

function unpack(z)
    x = [z[x_idx[t]] for t = 1:T]
    u = [z[u_idx[t]] for t = 1:T-1]
    return x, u
end

function primal_bounds(x1, xT)
    zl = -Inf * ones(num_var)
    zu = Inf * ones(num_var)

	zl[x_idx[1]] = x1
	zu[x_idx[1]] = x1

	# zl[x_idx[T]] = xT
	# zu[x_idx[T]] = xT

    return zl, zu
end

function constraint_bounds()
    cl = zeros(num_con)
    cu = zeros(num_con)
    return cl, cu
end

function eval_objective(prob::MOIGeneralProblem, z)
	J = 0.0
	for t = 1:T
		if t < T
			x = z[x_idx[t]]
			u = z[u_idx[t]]
			J += objective(x, u, t)
		else
			x = z[x_idx[t]]
			J += objective(x, nothing, t)
		end
	end
	# parameter regularization
	θ = z[p_idx]
	J += 100.0 * (θ - θ0_policy)' * (θ - θ0_policy)
	return J
end

function eval_objective_gradient!(∇J, z, prob::MOIGeneralProblem)
    ∇J .= 0.0
	for t = 1:T
		if t < T
			x = z[x_idx[t]]
			u = z[u_idx[t]]
			obj_x(a) = objective(a, u, t)
			obj_u(a) = objective(x, a, t)
			∇J[x_idx[t]] = ForwardDiff.gradient(obj_x, x)
			∇J[u_idx[t]] = ForwardDiff.gradient(obj_u, u)
		else
			x = z[x_idx[t]]
			obj_xT(a) = objective(a, nothing, t)
			∇J[x_idx[t]] = ForwardDiff.gradient(obj_xT, x)
		end
	end
	θ = z[p_idx]
	∇J[p_idx] = 2.0 * 100.0 * (θ - θ0_policy)

    return nothing
end

function eval_constraint!(c, z, prob::MOIGeneralProblem)
	# dynamics
    for t = 1:T-1
		x = z[x_idx[t]]
		u = z[u_idx[t]]
		y = z[x_idx[t+1]]
		c[n_idx[t]] = dynamics(y, x, u, t)
	end

	# policy
	θ = z[p_idx]
	for t = 1:T-1
		x = z[x_idx[t]]
		u = z[u_idx[t]]
		c[num_dyn .+ m_idx[t]] = u - policy(x, θ)[1]
	end
    return nothing
end

function eval_constraint_jacobian!(∇c, z, prob::MOIGeneralProblem)
	shift = 0

    for t = 1:T-1
        x = view(z, x_idx[t])
        u = view(z, u_idx[t])
        y = view(z, x_idx[t + 1])

        dyn_x(a) = dynamics(y, a, u, t)
        dyn_u(a) = dynamics(y, x, a, t)
        dyn_y(a) = dynamics(a, x, u, t)

        r_idx = n_idx[t]

        s = n * n
        ∇c[shift .+ (1:s)] = vec(ForwardDiff.jacobian(dyn_x, x))
        shift += s

        s = n * m
        ∇c[shift .+ (1:s)] = vec(ForwardDiff.jacobian(dyn_u, u))
        shift += s

        s = n * n
        ∇c[shift .+ (1:s)] = vec(ForwardDiff.jacobian(dyn_y, y))
        shift += s
    end

	# policy
	θ = z[p_idx]
	for t = 1:T-1
		x = z[x_idx[t]]
		u = z[u_idx[t]]

		px, pθ = policy_gradients(x, θ)

		r_idx = num_dyn .+ m_idx[t]

		s = m * m
		∇c[shift .+ (1:s)] = vec(Diagonal(ones(m)))
		shift += s

		s = m * n
		∇c[shift .+ (1:s)] = -vec(px)
		shift += s

		s = m * p
		∇c[shift .+ (1:s)] = -vec(pθ)
		shift += s
	end

    return nothing
end

function sparsity_jacobian(prob::MOIGeneralProblem)

	row = []
	col = []

	for t = 1:T-1
		r_idx = n_idx[t]
		row_col!(row, col, r_idx, x_idx[t])
		row_col!(row, col, r_idx, u_idx[t])
		row_col!(row, col, r_idx, x_idx[t + 1])
	end

	for t = 1:T-1
		r_idx = num_dyn .+ (m_idx[t])
		row_col!(row, col, r_idx, u_idx[t])
		row_col!(row, col, r_idx, x_idx[t])
		row_col!(row, col, r_idx, p_idx)
	end

	return collect(zip(row, col))
end

## Trajectory Optimization

# time
T = 11
h = 1.0

# dynamics
n = 2
m = 1

A = SMatrix{n,n}([1.0 1.0; 0.0 1.0])
B = SMatrix{n,m}([0.0; 1.0])

function dynamics(y, x, u, t)
	y - (A * x + B * u)
end

# objective
function objective(x, u, t)
	J = 0.0
	if t == T
		Q = Diagonal(100.0 * ones(n))
		J += x' * Q  * x
	else
		Q = Diagonal(ones(n))
		R = Diagonal(0.1 * ones(m))
		J += x' * Q  * x
		J += transpose(u) * R * u
	end
	return J
end

## policy

T_mpc = 5
n = 2
m = 1

x1 = [1.0; 0.0]
xT = [0.0; 0.0]

x_init = x1
_Q = sqrt.(ones(n))
_R = sqrt.(1.0 * ones(m))
_QT = sqrt.(1.0 * ones(n))
x_final = xT

nθ_mpc = n + n + m + n + n + n * n + n * m

num_x_mpc = n * T_mpc
num_u_mpc = m * (T_mpc - 1)
num_y_mpc = n * T_mpc
nz_mpc = num_x_mpc + num_u_mpc + num_y_mpc

x_idx_mpc = [collect((t - 1) * n .+ (1:n)) for t = 1:T_mpc]
u_idx_mpc = [collect(num_x_mpc + (t - 1) * m .+ (1:m)) for t = 1:T_mpc-1]
y_idx_mpc = [collect(num_x_mpc + num_u_mpc + (t - 1) * n .+ (1:n)) for t = 1:T_mpc]

x_mpc = linear_interpolation(x_init, x_final, T_mpc)
u_mpc = [0.01 * randn(m) for t = 1:T_mpc-1]
y_mpc = [zeros(n) for t = 1:T_mpc]

z0 = [vcat(x_mpc...); vcat(u_mpc...); vcat(y_mpc...)]
θ0 = [x_init; _Q; _R; _QT; x_final; vec(A); vec(B)]

function lagrangian(z, θ)
	x_init = θ[1:n]
	_Q = θ[n .+ (1:n)]
	Qt = Diagonal(_Q.^2.0)
	_R = θ[n + n .+ (1:m)]
	Rt = Diagonal(_R.^2.0)
	_QT = θ[n + n + m .+ (1:n)]
	QT = Diagonal(_QT.^2.0)
	x_final = θ[n + n + m + n .+ (1:n)]
	Aθ = reshape(θ[n + n + m + n + n .+ (1:n * n)], n, n)
	Bθ = reshape(θ[n + n + m + n + n + n * n .+ (1:n * m)], n, m)

	L = 0.0

	x1 = view(z, x_idx_mpc[1])
	y1 = view(z, y_idx_mpc[1])

	L += transpose(y1) * (x_init - x1)

	for t = 1:T_mpc-1
		xt = view(z, x_idx_mpc[t])
		ut = view(z, u_idx_mpc[t])
		yt⁺ = view(z, y_idx_mpc[t+1])
		xt⁺ = view(z, x_idx_mpc[t+1])

		L += transpose(xt - x_final) * Qt * (xt - x_final)
		L += (transpose(ut) * Rt * ut)

		L += transpose(yt⁺) *  (xt⁺ - (Aθ * xt + Bθ * ut))
	end

	xT = view(z, x_idx_mpc[T_mpc])
	# yT⁺ = view(z, y_idx_mpc[T+1])

	L += transpose(xT - x_final) * QT * (xT - x_final)
	# L += transpose(yT⁺) * (x_final - xT)

	return L
end

lagrangian(z0, θ0)

@variables z_sym[1:nz_mpc]
@variables θ_sym[1:nθ_mpc]
@variables κ_sym[1:1]

L = lagrangian(z_sym, θ_sym);
L = simplify.(L);

dL = Symbolics.gradient(L, z_sym)
ddL = Symbolics.sparsehessian(L, z_sym)

L_grad = eval(Symbolics.build_function(dL, z_sym, θ_sym, κ_sym)[1])
L_hess = eval(Symbolics.build_function(ddL, z_sym, θ_sym)[1])

L_grad! = eval(Symbolics.build_function(dL, z_sym, θ_sym, κ_sym)[2])
L_hess! = eval(Symbolics.build_function(ddL, z_sym, θ_sym)[2])

L_grad(z0, θ0, 1.0)
L_hess(z0, θ0)

dLθ = Symbolics.jacobian(dL, θ_sym)

dL_jacθ = eval(Symbolics.build_function(dLθ, z_sym, θ_sym)[1])
dL_jacθ! = eval(Symbolics.build_function(dLθ, z_sym, θ_sym)[2])

dL_jacθ(z0, θ0)

# options
opts = InteriorPointOptions(
	κ_init = 1.0,
	κ_tol = 1.0,
	r_tol = 1.0e-8,
	diff_sol = true)

# solver
ip = interior_point(z0, θ0,
	r! = L_grad!, rz! = L_hess!,
	rz = similar(ddL, Float64),
	rθ! = dL_jacθ!,
	rθ = similar(dLθ, Float64),
	opts = opts)

# solve
status = interior_point!(ip)

using Test, Plots
# test
@test status
@test norm(ip.z[x_idx_mpc[1]] - x_init) < 1.0e-8

norm(ip.z[x_idx_mpc[T_mpc]] - x_final)

x_traj_mpc = [ip.z[x_idx_mpc[t]] for t = 1:T_mpc]
u_traj_mpc = [ip.z[u_idx_mpc[t]] for t = 1:T_mpc-1]

plot(hcat(x_traj_mpc...)')
plot(hcat(u_traj_mpc...)')

function policy(x, θ)
	x_mpc = linear_interpolation(x, x_final, T_mpc)
	u_mpc = [0.01 * randn(m) for t = 1:T_mpc-1]
	y_mpc = [zeros(n) for t = 1:T_mpc]

	# initialized solver primals
	ip.z .= copy([vcat(x_mpc...); vcat(u_mpc...); vcat(y_mpc...)])

	# initialized solver parameters
	_Q = θ[1:n]
	_R = θ[n .+ (1:m)]
	_QT = θ[n + m .+ (1:n)]
	_A = θ[n + m + n .+ (1:n * n)]
	_B = θ[n + m + n + n * n .+ (1:n * m)]

	ip.θ .= copy([x; _Q; _R; _QT; x_final; _A; _B])

	status = interior_point!(ip)

	!status && (@warn "solver failure")

	return ip.z[u_idx_mpc[1]], ip
	# return reshape(θ, m, n) * x, nothing
end

function policy_gradients(x, θ)
	_, ip = policy(x, θ)

	ux = ip.δz[u_idx_mpc[1], 1:n]
	uθ = ip.δz[u_idx_mpc[1], n .+ (1:(n + m + n + n * n + n * m))]

	return ux, uθ

	# px(a) = policy(a, θ)[1]
	# pθ(a) = policy(x, a)[1]
	# return ForwardDiff.jacobian(px, x), ForwardDiff.jacobian(pθ, θ)
end

# policy(x1, 0.01 * randn(p))
# policy_gradients(x1, 0.01 * randn(p))
θ0_policy = [_Q; _R; _QT; vec(A); vec(B)]
policy(x1, θ0_policy)[1]
policy_gradients(x1, θ0_policy)

# p = m * n
p = n + m + n + n * n + n * m

# trajectory indices
x_idx = [(t - 1) * (n + m) .+ (1:n) for t = 1:T]
u_idx = [(t - 1) * (n + m) + n .+ (1:m) for t = 1:T-1]
n_idx = [(t - 1) * n .+ (1:n) for t = 1:T-1]
m_idx = [(t - 1) * m .+ (1:m) for t = 1:T-1]
p_idx = n * T + m * (T - 1) .+ (1:p)

# MOI setup
num_x = n * T
num_u = m * (T - 1)
num_var = num_x + num_u + p

num_dyn = n * (T - 1)
num_policy = m * (T - 1)
num_con = num_dyn + num_policy

# initialization
x1 = [1.0; 0.0]
xT = [0.0; 0.0]

# U = [0.1 * randn(m) for t = 1:T-1]

# rollout
x_hist = [copy(x1)]
u_hist = []
Js = 0.0
for t = 1:T-1
	push!(u_hist, policy(x_hist[end], θ0_policy)[1])
	# push!(u_hist, policy(x_hist[end], θ_sol)[1])

	Js += objective(x_hist[end], u_hist[end], t)
	push!(x_hist, A * x_hist[end] + B * u_hist[end])
end
Js += objective(x_hist[end], nothing, T)
@show Js
@show norm(x_hist[end] - xT)

z = [pack(x_hist, u_hist); θ0_policy]
# z = [pack(X, U); 0.01 * randn(p)]# _Q; _R; _QT]

prob = MOIGeneralProblem(num_var, num_con, primal_bounds(x1, xT), constraint_bounds())
prob_moi = moi_problem(prob)

z_sol, info = solve(prob_moi, z,
		tol = 1.0e-2,
		c_tol = 1.0e-2,
		max_iter = 1000,
		nlp = :ipopt)

x_sol, u_sol = unpack(z_sol)
θ_sol = z_sol[p_idx]

norm(θ_sol - θ0_policy)
# M = reshape(θ_sol, m, n)

using Plots

plot(hcat(x_sol...)')
plot!(hcat(x_hist...)')
plot(hcat(u_sol..., u_sol[end])', linetype = :steppost)
plot!(hcat(u_hist..., u_hist[end])', linetype = :steppost)

## simulate policy
x1_sim = copy(x1) + 0.0 * randn(n)

x_hist = [x1_sim]
u_hist = []
Js = 0.0
for t = 1:T-1
	# push!(u_hist, policy(x_hist[end], θ0_policy)[1])
	push!(u_hist, policy(x_hist[end], θ_sol)[1])

	Js += objective(x_hist[end], u_hist[end], t)
	push!(x_hist, A * x_hist[end] + B * u_hist[end])
end
Js += objective(x_hist[end], nothing, T)
@show Js
plot(hcat(x_hist...)')
plot(hcat(u_hist..., u_hist[end])', linetype = :steppost)

@show norm(x_hist[end] - xT)
Q_sol = Diagonal(θ_sol[1:n].^2.0)
R_sol = Diagonal(θ_sol[n .+ (1:m)].^2.0)
QT_sol = Diagonal(θ_sol[n + m .+ (1:n)].^2.0)
reshape(θ_sol[n + m + n .+ (1:n * n)], n, n)
reshape(θ_sol[n + m + n + n * n .+ (1:n * m)], n, m)

eigen(reshape(θ_sol[n + m + n .+ (1:n * n)], n, n))

max_issues_repo_name: thowell/DirectMotionPlanning


text:
using Printf, Test, Plots, LinearAlgebra

include("../../hybmc/mathutils/Helpers.jl")

include("../../hybmc/termstructures/YieldCurve.jl")
include("../../hybmc/models/HullWhiteModel.jl")
include("../../hybmc/models/AssetModel.jl")
include("../../hybmc/models/DeterministicModel.jl")
include("../../hybmc/models/QuasiGaussianModel.jl")
include("../../hybmc/models/HybridModel.jl")

include("../../hybmc/simulations/McSimulation.jl")
include("../../hybmc/simulations/Payoffs.jl")

using LinearAlgebra

function HwModel(rate=0.01,vol=0.0050,mean=0.03)
    curve = YieldCurve(rate)
    times = [ 10.0 ]
    vols  = [ vol  ]
    return HullWhiteModel(curve, mean, times, vols)
end

function fwd(mcSim::McSimulation,p::Payoff)
    samples = discountedAt(p,mcSim)
    fwd = mean(samples) /
        discount(mcSim.model.domRatesModel.yieldCurve,obsTime(p))
    err = std(samples) / sqrt(size(samples)[1]) /
        discount(mcSim.model.domRatesModel.yieldCurve,obsTime(p))
    return (fwd, err)
end

function setUp()
    ### full smile/skew model
    # domestic rates
    domAlias = "EUR"
    eurCurve = YieldCurve(0.03)
    d     = 2
    times = [  10.0     ]
    sigma = [ 0.0060 ;
              0.0040 ]
    slope = [ 0.10  ;
              0.15  ]
    curve = [ 0.05  ;
              0.10  ]
    delta = [ 1.0, 10.0  ]
    chi   = [ 0.01, 0.15 ]
    Gamma = [ 1.0   0.6  ;
              0.6   1.0  ]
    eurRatesModel = QuasiGaussianModel(eurCurve,d,times,sigma,slope,curve,delta,chi,Gamma)
    #
    forAliases = [ "USD", "GBP" ]
    spotS0     = [   1.0,   2.0 ]
    spotVol    = [   0.3,   0.2 ]
    forAssetModels = [
        AssetModel(S0, vol) for (S0, vol) in zip(spotS0,spotVol) ]
    # USD rates
    usdCurve = YieldCurve(0.02)
    d     = 3
    times = [  10.0   ]
    sigma = [  0.0060 ;
               0.0050 ;
               0.0040 ]
    slope = [  0.10   ;
               0.20   ;
               0.30   ]
    curve = [  0.05   ;
               0.10   ;
               0.20   ]
    delta = [ 1.0, 5.0, 20.0 ]
    chi   = [ 0.01, 0.05, 0.15 ]
    Gamma = [ 1.0  0.8  0.6 ;
              0.8  1.0  0.8 ;
              0.6  0.8  1.0 ]
    usdRatesModel = QuasiGaussianModel(usdCurve,d,times,sigma,slope,curve,delta,chi,Gamma)
    #
    gbpRatesModel = HwModel()
    #
           # 'EUR_x_0', 'EUR_x_1', 'USD_logS', 'USD_x_0', 'USD_x_1', 'USD_x_2', 'GBP_logS', 'GBP_x'
    corr = [  1.0        0.0         0.5       -0.5        0.0        0.0        -0.5       -0.5   ;   # EUR_x_0
              0.0        1.0         0.0        0.0       -0.5        0.0        -0.5        0.0   ;   # EUR_x_1
              0.5        0.0         1.0       -0.5       -0.5       -0.5         0.0        0.0   ;   # USD_logS
             -0.5        0.0        -0.5        1.0        0.0        0.0         0.0        0.0   ;   # USD_x_0
              0.0       -0.5        -0.5        0.0        1.0        0.0         0.0        0.0   ;   # USD_x_1
              0.0        0.0        -0.5        0.0        0.0        1.0         0.0        0.0   ;   # USD_x_2
             -0.5       -0.5         0.0        0.0        0.0        0.0         1.0        0.5   ;   # GBP_logS
             -0.5        0.0         0.0        0.0        0.0        0.0         0.5        1.0   ]   # GBP_x
    model = HybridModel(domAlias,eurRatesModel,forAliases,forAssetModels,[usdRatesModel,gbpRatesModel],corr)
    return model
end

@testset "test_ModelSetup" begin
    model = setUp()
    @test stateSize(model) == 24
    @test factors(model) == 8
end

@testset "test_HybridSimulation" begin
    model = setUp()
    times = append!(Array(range(0,stop=10.0,length=11)), [10.5])
    nPaths = 2^13
    seed = 314159265359
    # risk-neutral simulation
    mcSim = McSimulation(model,times,nPaths,seed,false)
    #
    T = 10.0
    P = Pay(Fixed(1.0),T)
    (fw, err) = fwd(mcSim,P)
    # domestic numeraire
    @printf("1.0   @ %4.1lfy %8.6lf - mc_err = %8.6lf\n", T,fw,err)
    @test abs(fw - 1.0) < 1.5*err
    # foreign assets
    for (k, alias) in enumerate(model.forAliases)
        p = Asset(T,alias)
        xT = model.forAssetModels[k].X0 *
            discount(model.forRatesModels[k].yieldCurve,T) /
            discount(model.domRatesModel.yieldCurve,T)
        (fw, err) = fwd(mcSim,p)
        @printf("%s   @ %4.1lfy %8.6lf vs %8.6lf (curve) - mc_err = %8.6lf\n", alias,T,fw,xT,err)
        @test abs(fw - xT) < 2.0*err
    end
    # domestic Libor rate
    Tstart = 10.0
    Tend = 10.5
    L = Pay(LiborRate(T,Tstart,Tend,alias="EUR"),Tend)
    (fw, err) = fwd(mcSim,L)
    Lref = (discount(mcSim.model.domRatesModel.yieldCurve,Tstart) /
            discount(mcSim.model.domRatesModel.yieldCurve, Tend) - 1) /
            (Tend - Tstart)
    @printf("L_EUR @ %4.1lfy %8.6lf vs %8.6lf (curve) - mc_err = %8.6lf\n", T,fw,Lref,err)
    @test abs(fw - Lref) < 2.0*err
    # foreign Lbor rates
    for (k, alias) in enumerate(model.forAliases)
        L = Pay(LiborRate(T,Tstart,Tend,alias=alias)*Asset(Tend,alias),Tend)
        (fw, err) = fwd(mcSim,L)
        fw *= discount(mcSim.model.domRatesModel.yieldCurve, Tend) /
              discount(mcSim.model.forRatesModels[k].yieldCurve, Tend) /
              mcSim.model.forAssetModels[k].X0
        err *= discount(mcSim.model.domRatesModel.yieldCurve, Tend) /
               discount(mcSim.model.forRatesModels[k].yieldCurve, Tend) /
               mcSim.model.forAssetModels[k].X0
        Lref = (discount(mcSim.model.forRatesModels[k].yieldCurve,Tstart) /
                discount(mcSim.model.forRatesModels[k].yieldCurve, Tend) - 1) /
                (Tend - Tstart)
        @printf("L_%s @ %4.1lfy %8.6lf vs %8.6lf (curve) - mc_err = %8.6lf\n", alias,T,fw,Lref,err)
        @test abs(fw - Lref) < 2.0*err
    end
end

@testset "test_HybridVolAdjusterCalculation" begin
    model = setUp()
    hybAdjTimes = Array(range(0.0, stop=20.0, length=21))
    hybVolAdj = hybridVolAdjuster(model,hybAdjTimes)
    model = HybridModel(model,hybAdjTimes,hybVolAdj)
    # plot(hybAdjTimes, transpose(hybVolAdj))
    # savefig("hybVolAdj.png")
    times = Array(range(0.0, stop=10.0, length=11))
    nPaths = 2^13
    seed = 314159265359
    # risk-neutral simulation
    mcSim = McSimulation(model,times,nPaths,seed,false)
    #
    T = 10.0
    for (k, alias) in enumerate(model.forAliases)
        # ATM forward
        xT = model.forAssetModels[k].X0 *
             discount(model.forRatesModels[k].yieldCurve, T) /
             discount(model.domRatesModel.yieldCurve, T)
        K = Fixed(xT)
        Z = Fixed(0.0)
        C = Pay(Max(Asset(T,alias)-K,Z),T)
        (fw, err) = fwd(mcSim,C)
        vol = BlackImpliedVol(fw,xT,xT,T,1.0)
        vega = BlackVega(xT,xT,vol,T)
        err /= vega
        volRef = model.forAssetModels[k].sigma
        @printf("C_%s @ %4.1lfy %8.6lf vs %8.6lf (curve) - mc_err = %8.6lf\n", alias,T,vol,volRef,err)
        @test abs(vol - volRef) < 5.0*err
        P = Pay(Max(K-Asset(T,alias),Z),T)
        (fw, err) = fwd(mcSim,P)
        vol = BlackImpliedVol(fw,xT,xT,T,-1.0)
        vega = BlackVega(xT,xT,vol,T)
        err /= vega
        volRef = model.forAssetModels[k].sigma
        @printf("P_%s @ %4.1lfy %8.6lf vs %8.6lf (curve) - mc_err = %8.6lf\n", alias,T,vol,volRef,err)
        @test abs(vol - volRef) < 5.0*err
    end
end

ext: jl

max_stars_repo_name: sschlenkrich/HybridMonteCarlo

max_issues_repo_name: sschlenkrich/HybridMonteCarlo

max_forks_repo_name: sschlenkrich/HybridMonteCarlo


text:
#using Pkg
#Pkg.add("DifferentialEquations")
using Plots
using DifferentialEquations

a = 0.25
eps = 0.002
Iapp = 0.0
gam = 1.1
p = [a,eps,gam,Iapp]
tspan = (0.0,1000.0)

function FN(du,u,t,p)
    du[1] = u[1]*(a-u[1])*(u[1]-1)-u[2]+Iapp
    du[2] = eps*(u[1]-gam*u[2])
end

function apply(u0)
    prob = ODEProblem(FN,u0,tspan)
    sol = solve(prob)
end

sol1 = apply([0.1,0.0])
sol2 = apply([0.2,0.0])
sol3 = apply([0.3,0.0])
sol4 = apply([0.4,0.0])

p1 = plot(sol1,linecolor = :red)
p2 = plot!(sol2, linecolor = :blue)
p3 = plot!(sol3, linecolor = :green)
p4 = plot!(sol4, linecolor = :yellow)

# Phase-space
p1 = plot(sol1,vars = (1,2),linecolor = :red,xlims = (-0.25,1),ylims = (0,0.125))
p2 = plot!(sol2,vars = (1,2), linecolor = :blue,xlims = (-0.25,1),ylims = (0,0.125))
p3 = plot!(sol3,vars = (1,2), linecolor = :green,xlims = (-0.25,1),ylims = (0,0.125))
p4 = plot!(sol4,vars = (1,2), linecolor = :yellow,xlims = (-0.25,1),ylims = (0,0.125), xlabel="v",ylabel="w")

# Nullclines
# dv = -v^3 + (a-1)*v^2 -a*v -w = 0
#-> w = -v^3 + (a-1)*v^2 -a*v
# dw = eps*(v-gam*w) = 0
# -> w = v/gamma

function FNnc(u,p)
    ncv = u.*(a.-u).*(u.-1).+Iapp
    ncw = u./gam
    nc = [ncv,ncw]
end

u = collect(-0.25:0.1:1.5)
nc = FNnc(u,p)
plot!(u,nc)

# One sees that v grows for v0>0.25 (yellow&green) until
# they intersect with the nullcline (turkis) and then decreases
# while w grows rapidly.
# They all follow the trajectory until they reach the fixed point (0,0)

## Testing influence of Iapp
u0 = [0.3;0.0]
Iapp = 0.02
sol5 = apply(u0)
nc5 = FNnc(u,[a,eps,gam,Iapp])
Iapp = 0.1
sol6 = apply(u0)
nc6 = FNnc(u,[a,eps,gam,Iapp])
Iapp = 0.5
sol7 = apply(u0)
nc7 = FNnc(u,[a,eps,gam,Iapp])
Iapp = 0.7
sol8 = apply(u0)
nc8 = FNnc(u,[a,eps,gam,Iapp])

p5 = plot(sol5,linecolor = :red)
p6 = plot!(sol6,linecolor = :blue)
p7 = plot!(sol7,linecolor = :green)
p8 = plot!(sol8,linecolor = :yellow)

p5 = plot(sol5,vars = (1,2),linecolor = :red, label="Iapp=0.02")
p5nc = plot!(u,nc5,linecolor = :red, linestyle = :dash, label="Nullcline, Iapp=0.02",ylim=(0,0.8))
p6 = plot!(sol6,vars = (1,2),linecolor = :blue, label="Iapp=0.1")
p6nc = plot!(u,nc6,linecolor = :blue, linestyle = :dash, label="Nullcline, Iapp=0.1")
p7 = plot!(sol7,vars = (1,2),linecolor = :green, label="Iapp=0.5")
p7nc = plot!(u,nc7,linecolor = :green, linestyle = :dash, label="Nullcline, Iapp=0.5")
p8 = plot!(sol8,vars = (1,2),linecolor = :yellow, label="Iapp=0.7",xlim=(-0.25,1.5))
p8nc = plot!(u,nc8,linecolor = :yellow, linestyle = :dash, label="Nullcline, Iapp=0.7")

# Iapp gives higher/lower threshold for firing


text:
"""
# Arguments
- `pomdp::CovidPOMDP` - Simulation parameters
- `state::CovidState` - Current Sim State
- `action::CovidAction` - Current Sim Action
- `T::Int` - Simulation duration (days)
"""
function POMDPs.simulate(
    pomdp::POMDP{S},
    state::S = S(pomdp),
    action::CovidAction = CovidAction(0.0);
    T::Int = 50
    ) where S <: CovidState

    stateHist = Vector{S}(undef, T)
    testHist = zeros(Int,T)
    actionHist = fill(action, T)
    rewardHist = zeros(Float64, T)
    testHist = Vector{Int}(undef, T)

    for day in 1:T
        stateHist[day] = state

        sp, new_infections, pos_tests = sim_step(pomdp, state, action)
        r = reward(pomdp, state, action,sp)

        testHist[day] = sum(pos_tests)
        rewardHist[day] = r
        state = sp
    end

    return SimHist(
        stateHist,
        pomdp.N,
        T,
        testHist,
        actionHist,
        rewardHist,
        Matrix{S}(undef,0,0)
    )
end

function POMDPs.simulate(
    pomdp::POMDP{S},
    b,
    planner::Policy;
    s = rand(b),
    T::Int=50,
    upd = BootstrapFilter,
    n_p::Int = 10_000,
    progress::Bool=false) where S <: CovidState

    stateHist = Vector{S}(undef, T)
    actionHist = zeros(CovidAction, T)
    rewardHist = zeros(Float64, T)
    testHist = Vector{Int}(undef, T)
    beliefHist = Matrix{S}(undef, n_p, T)

    single_step_pomdp = unity_test_period(pomdp)
    upd = upd(single_step_pomdp, n_p)

    prog = Progress(T; enabled=progress)
    for day in 1:T

        if (day-1)%pomdp.test_period == 0
            a = POMDPs.action(planner, b)
        else
            a = actionHist[day-1]
        end
        isone(day) && (b = initialize_belief(upd, b))

        stateHist[day] = s
        actionHist[day] = a
        beliefHist[:,day] .= b.particles

        s, o, r = POMDPs.gen(single_step_pomdp, s, a)
        b = update(upd, b, a, o)

        rewardHist[day] = r
        testHist[day] = o

        next!(prog)
    end
    return SimHist(stateHist, pomdp.N, T, testHist, actionHist, rewardHist, beliefHist)
end

@inline function shift_inf!(I::Vector)
    @inbounds for i in length(I):-1:2
        I[i] = I[i-1]
    end
    I[1] = 0
    I
end

"""
Progress testing state forward
People k days from receiving test back are now k-1 days from receiving test
Tested individuals with infection age t move to infection age t + 1
"""
@inline function shift_test!(T::Matrix)
    s1,s2 = size(T)
    @inbounds for i in 1:(s1-1)
        @views T[i,:] .= T[i+1,:]
    end
    @inbounds for j in s2:-1:2
        @views T[:,j] .= T[:,j-1]
    end
    T[:,1] .= 0
    T[end,:] .= 0
    return T
end


text:
@testset "▶ types/ax_elem               " begin
    # TITLE
    t = GPlot.Title(text="title")
    @test t.text == "title"
    @test isnothing(t.textstyle.font)
    @test isnothing(t.textstyle.hei)
    @test isnothing(t.textstyle.color)
    @test isnothing(t.dist)

    t = GPlot.Title(text="title", textstyle=GPlot.TextStyle(font="psh"), dist=0.3)
    @test t.textstyle.font == "psh"
    @test t.dist == 0.3

    # TICKS AND TICKSLABELS
    t = GPlot.Ticks()
    @test t.labels.names == String[]
    @test t.labels.off == false
    @test isnothing(t.labels.angle)
    @test isnothing(t.labels.format)
    @test isnothing(t.labels.shift)
    @test isnothing(t.labels.dist)
    @test isnothing(t.linestyle.color)
    @test t.places == Float64[]
    @test t.off == false
    @test isnothing(t.length)
    @test t.symticks == false
    @test t.grid == false
end

@testset "▶ /ax_elem                    " begin
    f = Figure()
    # titles and co
    title("blah"); erase!(f)
    xtitle("blah"); erase!(f)
    x2title("blah"); erase!(f)
    ytitle("blah"); erase!(f)
    y2title("blah")
    @test f.axes[1].y2axis.title.text  == "blah"
    x2title("blih", font="psh")
    @test f.axes[1].x2axis.title.textstyle.font == "psh"
    x2title("blah")
    @test isnothing(f.axes[1].x2axis.title.textstyle.font)
    @test f.axes[1].x2axis.title.text == "blah"
    y2title("hello")
    @test f.axes[1].y2axis.title.text == "hello"
    erase!(f)
    xtitle("hello")
    erase!(f)
    title("blah", font="psh"); title("blih")
    @test f.axes[1].title.text == "blih"

    # ticks and ticklabels
    erase!(f)
    xticks([1, 2], ["A", "B"]); x2ticks([3, 4])
    yticks([1, 2], ["A", "B"]); y2ticks([3, 4])
    @test f.axes[1].xaxis.ticks.places == [1., 2.]
    @test f.axes[1].yaxis.ticks.places == [1., 2.]
    @test f.axes[1].x2axis.ticks.places == [3., 4.]
    @test f.axes[1].y2axis.ticks.places == [3., 4.]
    @test f.axes[1].xaxis.ticks.labels.names == ["A", "B"]
    @test f.axes[1].yaxis.ticks.labels.names == ["A", "B"]
    @test isempty(f.axes[1].x2axis.ticks.labels.names)
    @test isempty(f.axes[1].y2axis.ticks.labels.names)
    xticks([1.0, 2.3])
    @test f.axes[1].xaxis.ticks.places == [1.0, 2.3]
    @test isempty(f.axes[1].xaxis.ticks.labels.names)

    @test_throws ArgumentError xticks([1, 2], ["A", "B", "C"])

    x2ticks([3, 5])
    @test f.axes[1].x2axis.ticks.places == [3., 5.]
    yticks([3, 5])
    @test f.axes[1].yaxis.ticks.places == [3., 5.]
    y2ticks([3, 5])
    @test f.axes[1].y2axis.ticks.places == [3., 5.]

    # -- empty
    cla()
    xticks(color="blue")
    @test gca().xaxis.ticks.places == Float64[]
    @test gca().xaxis.ticks.labels.names == String[]
    @test gca().xaxis.ticks.labels.textstyle.color == colorant"blue"

    @test_throws ArgumentError xticks(Float64[], ["blah", "blah"])

    # grid
    cla()
    grid(axis=["x", "y"], lstyle="--", color="lightgray")
    @test gca().xaxis.ticks.linestyle.lstyle == 9
    @test gca().yaxis.ticks.linestyle.lstyle == 9
    @test gca().xaxis.ticks.linestyle.color == colorant"lightgray"
    @test gca().yaxis.ticks.linestyle.color == colorant"lightgray"

    x2ticks("off")
    @test gca().x2axis.ticks.off == true
    @test gca().x2axis.ticks.labels.off == true
    x2ticks("on")
    @test gca().x2axis.ticks.labels.off == false
    @test_throws ArgumentError xticks("not_reco")

    # more ax elem
    cla()
    xaxis("off")
    @test gca().xaxis.off == true
    xaxis("on")
    @test gca().xaxis.off == false
end

@testset "▶ set_prop/ax_elem            " begin
    f = Figure()
    # Title
    title("blah", dist=1); @test f.axes[1].title.dist == 1.0
    @test_throws GPlot.OptionValueError title("b", dist=-1)
    @test_throws GPlot.UnknownOptionError  title("b", blah=5)
    # Ticks
    y2ticks([1, 2], off=true); @test f.axes[1].y2axis.ticks.off
    @test_throws GPlot.NotImplementedError y2ticks([1, 2], symticks=true)
    erase!(f)
    x2ticks([1, 2], ["a", "b"], tickscolor="blue")
    @test f.axes[1].x2axis.ticks.linestyle.color == colorant"blue"
    y2ticks([1, 2], ["a", "b"], angle=45)
    @test  f.axes[1].y2axis.ticks.labels.angle == 45
    @test_throws GPlot.NotImplementedError yticks([1, 2], format="something")
    x2ticks([1, 2], shift=0.1)
    @test f.axes[1].x2axis.ticks.labels.shift == 0.1
    yticks([1, 2], off=true)
    @test f.axes[1].yaxis.ticks.off == true
    yticks([1, 2], ["a", "b"], hidelabels=true)
    @test f.axes[1].yaxis.ticks.labels.off == true
end

@testset "▶ apply_gle/ax_elem           " begin
    g = G.GLE()
    f = G.Figure();
    G.add_axes2d!()

    # title (see also apply_textstyle)
    erase!(f)
    title("title")
    xlabel("xlabel", dist=0.5)
    G.apply_axes!(g, f.axes[1], f.id, 1); s = String(take!(g))
    isin(s, "title \"title\"")
    isin(s, "xtitle \"xlabel\" dist 0.5")

    # ticks (see also apply_linestyle, apply_tickslabels)
    erase!(f)
    y2ticks([1, 2])
    xticks([1, 2], off=true)
    G.apply_axes!(g, f.axes[1], f.id, 1); s = String(take!(g))
    isin(s, "y2places 1.0 2.0")
    isin(s, "xticks off")
    # XXX test symticks, length

    # tickslabels (see also apply_textstyle)
    erase!(f)
    x2ticks([1, 2], ["a", "b"], dist=0.5, angle=45, shift=1)
    G.apply_axes!(g, f.axes[1], f.id, 1); s = String(take!(g))
    isin(s, "x2places 1.0 2.0")
    isin(s, "x2names \"a\" \"b\"")
    isin(s, "x2labels  dist 0.5")
    isin(s, "x2axis angle 45.0 shift 1.0")
end


text:
#=

  Safe cracking problem in Julia ConstraintSolver.jl

  From the Oz Primer:
  http://www.comp.nus.edu.sg/~henz/projects/puzzles/digits/index.html
  """
  The code of Professor Smart's safe is a sequence of 9 distinct
  nonzero digits C1 .. C9 such that the following equations and
  inequations are satisfied:

        C4 - C6   =   C7
   C1 * C2 * C3   =   C8 + C9
   C2 + C3 + C6   <   C8
             C9   <   C8

   and

   C1 <> 1, C2 <> 2, ..., C9 <> 9

  can you find the correct combination?
  """

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

=#

using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function safe_cracking(print_solutions=true,all_solutions=true)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>6,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    n = 9
    @variable(model, 1 <= x[1:n] <= n, Int)
    c1,c2,c3,c4,c5,c6,c7,c8,c9 = x

    @constraint(model, x in CS.AllDifferentSet())

    for i in 1:n
        @constraint(model,x[i] != i)
    end

    @constraint(model,c4 - c6 == c7)

    # This is a non-linear constraint which ConstraintSolver.jl
    # cannot handle. Hence we have to rely on table instead...
    # @constraint(model,c1 * c2 * c3 == c8 + c9) # non-linear
    @variable(model,1 <= c123 <= 999, Int)
    c123_table = mult_table3(1,9)
    @constraint(model,[c1,c2,c3,c123] in CS.TableSet(c123_table))
    @constraint(model,c123 == c8 + c9)

    @constraint(model,c2 + c3 + c6 <= c8 - 1)
    @constraint(model,c9 <= c8 -1)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                println("x:$x_val")

            end
        end
    else
        println("status:$status")
    end

    return status
end

@time safe_cracking(true,true)

max_stars_repo_name: Wikunia/hakank

max_issues_repo_name: Wikunia/hakank

max_forks_repo_name: Wikunia/hakank


text:
using DataStructures, LinearAlgebra

"""
    CornerSpaceRenorm.cgs(X)

Apply the classical Gram-Schmidt orthogonalisation procedure to a set of vectors.
Vectors are passed as columns of a dense matrix `X`.
"""
function cgs(X::Matrix{T}) where T <: Number
    m, n = size(X);
    V = copy(X);
    Q = zeros(T, size(X));

    @inbounds @views for j = 1:n
        for k = 1:j-1
            axpy!(-BLAS.dotc(m, Q[:,k], 1, V[:,j], 1), Q[:,k], V[:,j])
        end
        Q[:,j] .= V[:,j] ./ sqrt(BLAS.dotc(n, V[:,j], 1, V[:,j], 1));
    end

    return Q
end

"""
    CornerSpaceRenorm.cgs!(X)

In-place classical Gram-Schmidt orthogonalisation procedure.
See also: [`CornerSpaceRenorm.cgs`](@ref)
"""
function cgs!(X::Matrix{T}) where T <: Number
    m, n = size(X);
    V = copy(X);

    @inbounds @views for j = 1:n
        for k = 1:j-1
            axpy!(-BLAS.dotc(m, X[:,k], 1, V[:,j], 1), X[:,k], V[:,j])
        end
        X[:,j] .= V[:,j] ./ sqrt(BLAS.dotc(n, V[:,j], 1, V[:,j], 1));
    end

    return X
end

"""
    CornerSpaceRenorm.mgs(X)

Apply the modified Gram-Schmidt orthogonalisation procedure to a set of vectors.
Vectors are passed as columns of a dense matrix `X`.
"""
function mgs(X::Matrix{T}) where T <: Number
    m, n = size(X);
    V = copy(X);
    Q = copy(X);

    for j = 1:n
        @inbounds @views Q[:,j] .= V[:,j]./norm(V[:,j]);
        @inbounds @views for k = j+1:n
            axpy!(-BLAS.dotc(m, Q[:,j], 1, V[:,k], 1), Q[:,j], V[:,k])
        end
    end

    return Q
end

"""
    CornerSpaceRenorm.mgs!(X)

In-place modified Gram-Schmidt orthogonalisation procedure.
See also: [`CornerSpaceRenorm.mgs`](@ref)
"""
function mgs!(X::Matrix{T}) where T <: Number
    m, n = size(X);
    V = copy(X);

    for j = 1:n
        @inbounds @views X[:,j] .= V[:,j]./norm(V[:,j]);
        @inbounds @views for k = j+1:n
            axpy!(-BLAS.dotc(m, X[:,j], 1, V[:,k], 1), X[:,j], V[:,k])
        end
    end

    return X
end

"""
    CornerSpaceRenorm.maxk(a, k)

Find `k` largest elements of an array and return their indices.
# Arguments
* `a`: some array.
* `k`: number of largest elements to be found.
"""
function maxk(a::AbstractVector{T}, k::Int) where T<:Real
    ix = partialsortperm(a, 1:k, rev=true)
    @inbounds @views collect(ix), collect(a[ix[1:k]])
end

"""
    CornerSpaceRenorm.max_prod_pairs(a, b, k)

Find `k` pairs `(a_i,b_i)` with largest products from two arrays `a` and `b`.
# Arguments
* `a`: some array.
* `b`: some array.
* `k`: number of largest products to be found.
"""
function max_prod_pairs(a::AbstractVector{T1}, b::AbstractVector{T2}, k::Int) where {T1<:Real, T2<:Real}
    ix_a, _a = maxk(a, min(k,length(a)))
    ix_b, _b = maxk(b, min(k,length(b)))
    heap_idcs = [CartesianIndex(i,1) for i in 1:length(_a)]
    h = MutableBinaryMaxHeap([_a[idx[1]] * _b[idx[2]] for idx in heap_idcs])
    handles = Array{Tuple{Int64,Int64},1}(undef, k)
    prod_pairs = zeros(eltype(a), k)
    i = 1
    while i <= k && length(h) > 0
        prod_pairs[i], heap_idx = top_with_handle(h);
        handles[i] = (ix_a[heap_idcs[heap_idx][1]], ix_b[heap_idcs[heap_idx][2]])
        pop!(h);
        if heap_idcs[heap_idx][2] < length(_b)
            push!(heap_idcs, CartesianIndex(heap_idcs[heap_idx][1], heap_idcs[heap_idx][2]+1))
            push!(h, _a[heap_idcs[heap_idx][1]] * _b[heap_idcs[heap_idx][2]+1])
        end
        i += 1
    end
    return handles, prod_pairs
end

"""
    CornerSpaceRenorm.corner_subspace(ρA, ρB, M)

Find the `SubspaceBasis` associated with the `M`-dimensional corner defined by the
pair of states `(ρA, ρB)`, the pairs of eigenkets of A and B defining the corner
basis and the eigenkets of A and B.
# Arguments
* `ρA`: density matrix of the subsystem A.
* `ρB`: density matrix of the subsystem B.
* `M`: corner dimension.
"""
function corner_subspace(ρA::DenseMatrix, bA::Basis, ρB::DenseMatrix, bB::Basis, M::Int)
    ps_A, αs_A = eigen(ρA); # αs_A[:,i] : i-th eigenvector
    ps_B, αs_B = eigen(ρB);
    mgs!(αs_A);
    mgs!(αs_B);
    @inbounds ϕs_A = [Ket(bA, αs_A[:,i]) for i in 1:length(ps_A)]
    @inbounds ϕs_B = [Ket(bB, αs_B[:,i]) for i in 1:length(ps_B)]
    handles, prod_pairs = max_prod_pairs(real.(ps_A), real.(ps_B), M)
    bC = typeof(bA) <: CompositeBasis ? CompositeBasis(bA.bases...,bB.bases...) : CompositeBasis(bA,bB)
    return CornerBasis(bA, bB, M), handles, ϕs_A, ϕs_B
end
function corner_subspace(ρA::AbstractOperator{B1,B1}, ρB::AbstractOperator{B2,B2}, M::Int) where {B1<:Basis,B2<:Basis}
    if isdense(ρA) && isdense(ρB)
        return corner_subspace(ρA.data, ρA.basis_l, ρB.data, ρB.basis_l, M)
    else
        return corner_subspace(Matrix(ρA.data), ρA.basis_l, Matrix(ρB.data), ρB.basis_l, M)
    end
end

"""
    hermitianize(x)

Extract the Hermitian part of `x`.
# Arguments
* `x`: some operator.
"""
function hermitianize(x::AbstractOperator{B,B}) where {B<:Basis}
    ishermitian(x) && return x
    n = LinearAlgebra.checksquare(x.data);
    y = copy(x);
    @inbounds @views for j = 1:n, i = 1:j
        y.data[i,j] = (y.data[i,j] + conj(y.data[j,i])) / 2.
        y.data[j,i] = conj(y.data[i,j]);
    end
    return y
end

"""
    hermitianize!(x)

Extract the Hermitian part of `x`.
# Arguments
* `x`: some operator.
"""
function hermitianize!(x::AbstractOperator{B,B}) where {B<:Basis}
    ishermitian(x) && return x
    n = LinearAlgebra.checksquare(x.data);
    @inbounds @views for j = 1:n, i = 1:j
        x.data[i,j] = (x.data[i,j] + conj(x.data[j,i])) / 2.
        x.data[j,i] = conj(x.data[i,j]);
    end
    return x
end

"""
    merge(s1, s2, d, ρ1, ρ2, M)

Merge two `NdSystem`s along the `d`-th direction with corner compression.
# Arguments
* `s1`: `NdSystem`.
* `s2`: `NdSystem`.
* `d`: merging direction.
* `ρ1`: state of the system `s1`.
* `ρ2`: state of the system `s2`.
* `M`: corner dimension.
"""
function Base.merge(s1::NdSystem{N},s2::NdSystem{N},d::Integer,
                    ρ1::Operator{B1,B1,T1},ρ2::Operator{B2,B2,T2},M::Int;
                    return_dm=false) where {N,B1<:Basis,B2<:Basis,T1<:DenseMatrix,T2<:DenseMatrix}
    # TO DO: tests
    @assert 0 < M <= length(s1.gbasis)*length(s2.gbasis) "The corner dimension must not exceed that of the total Hilbert space."
    length(s1.lattice.Vext[d]) == length(s2.lattice.Vint[d]) || throw(DimensionMismatch("Lattices cannot be merged in this direction."))
    @assert s1.lattice.pbc == s2.lattice.pbc "Cannot merge systems with periodic and open open boundary conditions."
    lattice = union(s1.lattice,s2.lattice,d)

    all(s1.trate .≈ s2.trate) || @info "The two input systems do not have equal tunnelling rates along all directions."

    bC, handles, ϕs_1, ϕs_2 = corner_subspace(ρ1,ρ2,M)
    vt_1 = map(x->transpose(x.data), ϕs_1)
    vc_1 = map(x->conj(x.data), ϕs_1)
    vt_2 = map(x->transpose(x.data), ϕs_2)
    vc_2 = map(x->conj(x.data), ϕs_2)
    cache1 = [similar(first(vc_1)) for i in 1:Threads.nthreads()]
    cache2 = [similar(first(vc_2)) for i in 1:Threads.nthreads()]

    function 𝒫1(op)
        # TO DO: take advantage of orthogonormality to get rid of the scalar product on subspace 2
        opC = DenseOperator(bC)
        Threads.@threads for j in 1:length(handles)
            hj = handles[j]
            mul!(cache1[Threads.threadid()], op.data, vc_1[hj[1]])
            for i in 1:length(handles)
                hi = handles[i]
                opC.data[i,j] = vt_1[hi[1]] * cache1[Threads.threadid()] * (vt_2[hi[2]]*vc_2[hj[2]])
            end
        end
        return opC
    end
    function 𝒫1_bis(op)
        # TO DO: take advantage of orthogonormality to get rid of the scalar product on subspace 2
        opC = DenseOperator(bC)
        for j in 1:length(handles)
            hj = handles[j]
            mul!(cache1[Threads.threadid()], op.data, vc_1[hj[1]])
            for i in 1:length(handles)
                hi = handles[i]
                opC.data[i,j] = vt_1[hi[1]] * cache1[Threads.threadid()] * (vt_2[hi[2]]*vc_2[hj[2]])
            end
        end
        return opC
    end
    function 𝒫2(op)
        # TO DO: take advantage of orthogonormality to get rid of the scalar product on subspace 1
        opC = DenseOperator(bC)
        Threads.@threads for j in 1:length(handles)
            hj = handles[j]
            mul!(cache2[Threads.threadid()], op.data, vc_2[hj[2]])
            for i in 1:length(handles)
                hi = handles[i]
                opC.data[i,j] = vt_2[hi[2]] * cache2[Threads.threadid()] * (vt_1[hi[1]]*vc_1[hj[1]])
            end
        end
        return opC
    end
    function 𝒫(op1,op2)
        opC = DenseOperator(bC)
        Threads.@threads for j in 1:length(handles)
            hj = handles[j]
            mul!(cache1[Threads.threadid()], op1.data, vc_1[hj[1]])
            mul!(cache2[Threads.threadid()], op2.data, vc_2[hj[2]])
            for i in 1:length(handles)
                hi = handles[i]
                opC.data[i,j] = (vt_1[hi[1]] * cache1[Threads.threadid()]) * (vt_2[hi[2]] * cache2[Threads.threadid()])
            end
        end
        return opC
    end

    # TO DO: exploit Hermicianity of H to compute half of the matrix elements in the corner
    H = DenseOperator(bC)
    if lattice.pbc
        Ht1 = zeros(ComplexF64,size(s1.H.data))
        @inbounds for i in 1:length(s1.Htext[d])
            Ht1 .+= s1.trate[d] * (s1.Htext[d][i] * dagger(s1.Htint[d][i])).data;
        end
        Ht2 = zeros(ComplexF64,size(s2.H.data))
        @inbounds for i in 1:length(s2.Htext[d])
                Ht2 .+= s2.trate[d] * (s2.Htext[d][i] * dagger(s2.Htint[d][i])).data;
            end
        H.data .= 𝒫1(DenseOperator(s1.gbasis, s1.H.data .- (Ht1 .+ Ht1'))).data .+ 𝒫2(DenseOperator(s2.gbasis, s2.H.data .- (Ht2 .+ Ht2'))).data;
    else
        H.data .= 𝒫1(s1.H).data .+ 𝒫2(s2.H).data;
    end
    gbasis = H.basis_l;
    Ht = zeros(ComplexF64,size(H.data));
    @inbounds for i in 1:length(s1.Htext[d])
        Ht .= s1.trate[d] * 𝒫(s1.Htext[d][i],dagger(s2.Htint[d][i])).data;
        if lattice.pbc
            Ht .+= s1.trate[d] * 𝒫(dagger(s1.Htint[d][i]), s2.Htext[d][i]).data;
        end
        H.data .+= Ht .+ Ht';
    end
    hermitianize!(H);

    J = [[𝒫1(s1.J[i]) for i in 1:length(s1.J)]; [𝒫2(s2.J[i]) for i in 1:length(s2.J)]]

    #return H
    d⊥ = [_d for _d in 1:N if _d!= d]
    T = typeof(H)
    Htint::Array{Array{T,1},1} = [Array{T,1}(undef,0) for i in 1:N]
    Htext::Array{Array{T,1},1} = [Array{T,1}(undef,0) for i in 1:N]

    Htint[d] = 𝒫1.(s1.Htint[d])
    Htext[d] = 𝒫2.(s2.Htext[d])
    for _d in d⊥
        Htint[_d] = [𝒫1.(s1.Htint[_d]); 𝒫2.(s2.Htint[_d])]
        Htext[_d] = [𝒫1.(s1.Htext[_d]); 𝒫2.(s2.Htext[_d])]
    end

    # TO DO: Test that observable types are compatible
    obs = [[Dict{String,typeof(H)}() for i in 1:nv(s1.lattice)]; [Dict{String,typeof(H)}() for i in 1:nv(s2.lattice)]]

    for i in 1:length(s1.observables)
        for k in keys(s1.observables[i])
            obs[i][k] = 𝒫1(s1.observables[i][k])
        end
    end
    for i in 1:length(s2.observables)
        for k in keys(s2.observables[i])
            obs[length(s1.observables)+i][k] = 𝒫2(s2.observables[i][k])
        end
    end

    if return_dm
        return NdSystem{N,typeof(lattice),typeof(gbasis),typeof(H),eltype(first(Htint)),eltype(J),typeof(H)}(lattice,gbasis,H,s1.trate,Tuple(Htint),Tuple(Htext),J,obs), normalize(𝒫(ρ1, ρ2))
    else
        return NdSystem{N,typeof(lattice),typeof(gbasis),typeof(H),eltype(first(Htint)),eltype(J),typeof(H)}(lattice,gbasis,H,s1.trate,Tuple(Htint),Tuple(Htext),J,obs)
    end
end


text:
mutable struct IdealRobot <: AbstractRobot
    pose_::Vector{Float64}
    agent_::AbstractAgent
    radius_::Float64
    color_::String
    poses_::Vector{Vector{Float64}}
    sensor_::Union{AbstractSensor,Nothing}
    function IdealRobot(
        pose::Vector{Float64},
        agent::AbstractAgent,
        sensor::Union{AbstractSensor,Nothing},
        radius = 0.05,
        color = "blue",
    )
        if typeof(sensor) == Nothing
            new([pose[1], pose[2], pose[3]], agent, radius, color, [copy(pose)], nothing)
        else
            new([pose[1], pose[2], pose[3]], agent, radius, color, [copy(pose)], sensor)
        end
    end
end

function draw(robot::IdealRobot, p::Plot{T}) where {T}
    # robot
    p = scatter!(
        [robot.pose_[1]],
        [robot.pose_[2]],
        markersize = robot.radius_ * 100,
        color = robot.color_,
        markeralpha = 0.5,
        legend = nothing,
        aspect_ratio = :equal,
    )
    θ = robot.pose_[3]
    # pose
    p = quiver!(
        [robot.pose_[1]],
        [robot.pose_[2]],
        quiver = ([robot.radius_ * cos(θ) * 5], [robot.radius_ * sin(θ) * 5]),
        color = "black",
    )
    # traj
    x_his = [pose[1] for pose in robot.poses_]
    y_his = [pose[2] for pose in robot.poses_]
    p = plot!(x_his, y_his, color = "black", lw = 0.5)
    # camera
    draw(robot.sensor_, robot.pose_, p)
end

function state_transition(cur_pose::Vector{Float64}, v::Float64, ω::Float64, dt::Float64)
    θ = cur_pose[3]
    new_pose = [0.0, 0.0, 0.0]
    if abs(ω) < 1e-10
        dpose = [v * cos(θ), v * sin(θ), ω] * dt
        new_pose = cur_pose .+ dpose
    else
        dpose = [
            v / ω * (sin(θ + ω * dt) - sin(θ)),
            v / ω * (-cos(θ + ω * dt) + cos(θ)),
            ω * dt,
        ]
        new_pose = cur_pose .+ dpose
    end
    return new_pose
end

function state_transition(robot::IdealRobot, v::Float64, ω::Float64, dt::Float64)
    θ = robot.pose_[3]
    new_pose = [0.0, 0.0, 0.0]
    if abs(ω) < 1e-10
        dpose = [v * cos(θ), v * sin(θ), ω] * dt
        new_pose = robot.pose_ .+ dpose
    else
        dpose = [
            v / ω * (sin(θ + ω * dt) - sin(θ)),
            v / ω * (-cos(θ + ω * dt) + cos(θ)),
            ω * dt,
        ]
        new_pose = robot.pose_ .+ dpose
    end
    push!(robot.poses_, copy(robot.pose_))
    robot.pose_ = copy(new_pose)
    return new_pose
end

mutable struct RealRobot <: AbstractRobot
    pose_::Vector{Float64}
    agent_::AbstractAgent
    radius_::Float64
    color_::String
    poses_::Vector{Vector{Float64}}
    sensor_::Union{AbstractSensor,Nothing}
    # motion uncertainty
    ## movement noise
    noise_::Exponential{Float64}
    theta_noise_::Normal{Float64}
    distance_until_noise_::Float64
    ## velocity bias noise
    bias_rate_v_::Float64
    bias_rate_ω_::Float64
    ## stuck noise
    stuck_noise_::Exponential{Float64}
    escape_noise_::Exponential{Float64}
    time_until_stuck_::Float64
    time_until_escape_::Float64
    is_stuck_::Bool
    ## kidnap
    kidnap_noise_::Exponential{Float64}
    time_until_kidnap_::Float64
    kidnap_distrib_::PoseUniform

    function RealRobot(
        pose::Vector{Float64},
        agent::AbstractAgent,
        sensor::Union{AbstractSensor,Nothing};
        radius = 0.05,
        color = "blue",
        noise_per_meter = 5.0,
        noise_std = pi / 60,
        bias_rate_stds = (0.1, 0.1),
        expected_stuck_time = 1e100,
        expected_escape_time = 1e-100,
        expected_kidnap_time = 1e100,
        kidnap_range_x = [-5.0, 5.0],
        kidnap_range_y = [-5.0, 5.0],
    )
        noise = Exponential(1.0 / (1e-100 + noise_per_meter))
        stuck_noise = Exponential(expected_stuck_time)
        escape_noise = Exponential(expected_escape_time)
        kidnap_noise = Exponential(expected_kidnap_time)

        new(
            [pose[1], pose[2], pose[3]],
            agent,
            radius,
            color,
            [copy(pose)],
            sensor,
            noise, ## movement noise
            Normal(0.0, noise_std),
            rand(noise),
            rand(Normal(1.0, bias_rate_stds[1])), ## velocity noise
            rand(Normal(1.0, bias_rate_stds[2])),
            stuck_noise, ## stuck noise
            escape_noise,
            rand(stuck_noise),
            rand(escape_noise),
            false,
            kidnap_noise, ## kidnap
            rand(kidnap_noise),
            PoseUniform(kidnap_range_x, kidnap_range_y),
        )
    end
end

function apply_bias_error(robot::RealRobot, v::Float64, ω::Float64)
    return v * robot.bias_rate_v_, ω * robot.bias_rate_ω_
end

function apply_stuck_error(robot::RealRobot, v::Float64, ω::Float64, dt::Float64)
    if robot.is_stuck_
        robot.time_until_escape_ -= dt
        if robot.time_until_escape_ <= 0.0
            robot.time_until_escape_ += rand(robot.escape_noise_)
            robot.is_stuck_ = false
        end
    else
        robot.time_until_stuck_ -= dt
        if robot.time_until_stuck_ <= 0.0
            robot.time_until_stuck_ += rand(robot.stuck_noise_)
            robot.is_stuck_ = true
        end
    end

    return v * convert(Float64, robot.is_stuck_), ω * convert(Float64, robot.is_stuck_)
end

function apply_kidnap(robot::RealRobot, dt::Float64)
    robot.time_until_kidnap_ -= dt
    if robot.time_until_kidnap_ <= 0.0
        robot.time_until_kidnap_ += rand(robot.kidnap_noise_)
        robot.pose_ = uniform(robot.kidnap_distrib_)
    end
end

function state_transition(
    robot::RealRobot,
    v_::Float64,
    ω_::Float64,
    dt::Float64;
    move_noise = false,
    vel_bias_noise = false,
    stuck_noise = false,
    kidnap = false,
)
    v = v_
    ω = ω_

    ## velocity bias  noise
    if vel_bias_noise
        v, ω = apply_bias_error(robot, v, ω)
    end

    ## stuck noise
    if stuck_noise
        v, ω = apply_stuck_error(robot, v, ω, dt)
    end

    θ = robot.pose_[3]
    new_pose = [0.0, 0.0, 0.0]
    if abs(ω) < 1e-10
        dpose = [v * cos(θ), v * sin(θ), ω] * dt
        new_pose = robot.pose_ .+ dpose
    else
        dpose = [
            v / ω * (sin(θ + ω * dt) - sin(θ)),
            v / ω * (-cos(θ + ω * dt) + cos(θ)),
            ω * dt,
        ]
        new_pose = robot.pose_ .+ dpose
    end

    ## movement noise
    if move_noise
        robot.distance_until_noise_ -= (abs(v) * dt + robot.radius_ * abs(ω) * dt)
        if robot.distance_until_noise_ <= 0.0
            robot.distance_until_noise_ += rand(robot.noise_)
            new_pose[3] += rand(robot.theta_noise_)
        end
    end

    # save trajectory
    push!(robot.poses_, copy(robot.pose_))

    robot.pose_ = copy(new_pose)

    ## kidnap
    if kidnap
        apply_kidnap(robot, dt)
    end
end

function draw(robot::RealRobot, p::Plot{T}) where {T}
    # robot
    p = scatter!(
        [robot.pose_[1]],
        [robot.pose_[2]],
        markersize = robot.radius_ * 100,
        color = robot.color_,
        markeralpha = 0.5,
        legend = nothing,
        aspect_ratio = :equal,
    )
    θ = robot.pose_[3]
    # pose
    p = quiver!(
        [robot.pose_[1]],
        [robot.pose_[2]],
        quiver = ([robot.radius_ * cos(θ) * 5], [robot.radius_ * sin(θ) * 5]),
        color = "black",
    )
    # traj
    x_his = [pose[1] for pose in robot.poses_]
    y_his = [pose[2] for pose in robot.poses_]
    p = plot!(x_his, y_his, color = "black", lw = 0.5)
    # camera
    draw(robot.sensor_, robot.pose_, p)
    # agent
    draw(robot.agent_, p)
end

max_stars_repo_name: soblin/JuliaProbo

max_issues_repo_name: soblin/JuliaProbo

max_forks_repo_name: soblin/JuliaProbo


text:
# test saving and loading of trees

using Test
using IncrementalInference

@testset "Test loading and saving of Bayes (Junction) tree" begin

    fg = generateCanonicalFG_Kaess(graphinit=false)
    tree = buildTreeReset!(fg)

    # save and load tree
    saveTree(tree)
    tree2 = loadTree()

    # perform a few spot checks to see that the trees are similar
    @test length(tree.cliques) == length(tree2.cliques)
    @test getEliminationOrder(tree) == getEliminationOrder(tree2)

    for (clid,cl) in tree.cliques
      fsyms = getFrontals(cl)
      cl2 = getClique(tree2, fsyms[1])
      fsyms2 = getFrontals(cl2)
      @test fsyms == fsyms2
      @test getCliqSeparatorVarIds(cl) == getCliqSeparatorVarIds(cl2)
      @test typeof(cl) == typeof(cl2)
    end

end

@testset "Test loading and saving of Bayes (Junction) tree" begin

    fg = generateCanonicalFG_Kaess(graphinit=false)
    tree = buildTreeReset!(fg)

    # save and load tree as array
    filepath = saveTree([tree;deepcopy(tree)])
    trees = loadTree(filepath)

    # perform a few spot checks to see that the trees are similar
    @test length(tree.cliques) == length(trees[1].cliques)
    @test getEliminationOrder(tree) == getEliminationOrder(trees[1])

    for (clid,cl) in tree.cliques
      fsyms = getFrontals(cl)
      cl2 = getClique(trees[1], fsyms[1])
      fsyms2 = getFrontals(cl2)
      @test fsyms == fsyms2
      @test getCliqSeparatorVarIds(cl) == getCliqSeparatorVarIds(cl2)
      @test typeof(cl) == typeof(cl2)
    end

end


text:
module ParametricTypeAlias

# This definition will be stripped by creload - June'17
const MyVector{T} = Array{T, 1}

end


text:
# Script to replicate table 6 and 7 (regression from equation 17 and 18)

using DataFrames, FixedEffectModels, RegressionTables, StatsBase, Statistics, LinearAlgebra, CSV

# -------------- Previous scripts -----------------------------------------------------------------------------------------------------------------------

include("src/data_wrangling.jl") # formats the raw input data obtained from WIOD

# -------------- Country-sector level I/O-Table ---------------------------------------------------------------------------------------------------------

function country_sector_level_composition_indicators(df::DataFrame, year::Int64, output::String)

    ID, FD, GO, VA, IV, other = obtain_matrices(df, year) # N*S×N*S, N*S×N, N*S×1, N*S×1, N*S×N, N*S×1

    # net inventory correction
    iv_correction = [GO[i]/(GO[i]-IV[i, j]) for i in 1:n_ctrys*n_sectors, j in 1:n_ctrys] # N*S×N, to match country wise inventory changes
    # notice difference in column size, need spread inventory correction of country column to country-sector columns
    # i.e. ID[1,1:35]/iv_correction[1,1], however problem with indexing => ceil(Int, 34/35)=1, ceil(Int, 36/35)=2
    ID_corrected = [ID[i, j]*iv_correction[i, ceil(Int, j/n_sectors)] for i in 1:n_ctrys*n_sectors, j in 1:n_ctrys*n_sectors] # N*S×N*S
    FD_corrected = [FD[i, j]*iv_correction[i, j] for i in 1:n_ctrys*n_sectors, j in 1:n_ctrys] # N*S×N

    ID = replace!(ID_corrected, NaN => 0.0) # sometimes zeros are interpreted as NaN (particular for sector 35, private households with employed persons)
    FD = replace!(FD_corrected, NaN => 0.0)

    FD_ctry = [sum(FD[:, j]) for j in 1:n_ctrys] # 1×N
    FD_ctry = transpose(FD_ctry)

    sector_iterator = 0:n_sectors:n_ctrys*n_sectors-1
    # sum the same sector over all destination countries and divide by total final demand of destination country
    α = [sum(FD[sector_iterator .+ i, j]/FD_ctry[j]) for i in 1:n_sectors, j in 1:n_ctrys] # S×N, each column sums to 1.0

    # sum the same sector over all destination country-industry pairs and divide by corresponding output of country-industry
    γ = [sum(ID[sector_iterator .+ i, j])/GO[j] for i in 1:n_sectors, j in 1:n_ctrys*n_sectors] # S×N*S

    # in DataFrame format
    # add ifelse since dimensions are different, i.e. S×N vs S×N*S
    if (output == "α")
        table_α = hcat(fill(year, n_sectors), 1:35, α)
        namestable = ["year", "row_sector"]
        namestable = [namestable; ctrys]
        table_α = DataFrame(table_α, namestable)
        table_α_long = stack(table_α, Not([:year, :row_sector]))
        rename!(table_α_long, :variable => :country)

        result = table_α_long

    else
        table_γ = hcat(fill(year, n_sectors), 1:35, γ)
        namestable = ["year", "row_sector"]
        namestable = [namestable; repeat(ctrys, inner=n_sectors) .* "_" .* string.(repeat(1:35, outer=n_ctrys))]
        table_γ = DataFrame(table_γ, namestable)
        table_γ_long = stack(table_γ, Not([:year, :row_sector]))
        table_γ_long = hcat(table_γ_long, DataFrame(reduce(vcat, permutedims.(split.(table_γ_long.variable, "_"))), [:country, :col_sector]))
        transform!(table_γ_long, :col_sector => ByRow(x -> parse(Int64,x)) => :col_sector, renamecols=false) # change type to Int

        result = table_γ_long[:, Not(:variable)]
    end

    transform!(result, :country => ByRow(string) => :country, renamecols=false) # change type to String
    transform!(result, [:year, :row_sector] .=> ByRow(Int64) .=> [:year, :row_sector], renamecols=false) # change type to Int
    transform!(result, :value => ByRow(Float64) => :value, renamecols=false) # change type to Float64

    return result
end

t_tab6 = DataFrame(year=Int[], country=String[], row_sector=Int64[], value=Float64[]) # initialize empty DataFrames
t_tab7 = DataFrame(year=Int[], country=String[], row_sector=Int64[], col_sector=Int64[], value=Float64[])

for i in years
    append!(t_tab6, country_sector_level_composition_indicators(df, i, "α"))
    append!(t_tab7, country_sector_level_composition_indicators(df, i, "γ"))
end

CSV.write("clean/t_tab6.csv", t_tab6)
CSV.write("clean/t_tab7.csv", t_tab7)

# -------------- Table 6 -----------------------------------------------------------------------------------------------------------------------------

t_tab6 = CSV.read("clean/t_tab6.csv", DataFrame) # also imports types correctly!

function reg_tab_6(df::DataFrame)

    #df.value .= ifelse.(df.value .<= 0.0, 1e-18, df.value) # since log cannot deal with 0
    subset!(df, :value => ByRow(x -> x > 0.0)) # same results as with treatment above

    df[:, :country_industry] .= df[:, :country] .* "_" .* string.(df[:, :row_sector]) # for fixed effects
    transform!(df, :row_sector => ByRow(x -> x in goods_sectors ? 1 : 0) => :row_sector_dummy) # dummy for *reporting* sector, "goods" = 1

    df_goods = subset(df, :row_sector_dummy => ByRow(x -> x==1)) # only goods sectors
    df_services = subset(df, :row_sector_dummy => ByRow(x -> x!=1)) # only service sectors

    regression = Dict{String,Any}() # initializing dictionary to store regression results (Any could be replaced by type FixedEffectModel)
    df_names = ["all", "goods", "services"] # names of dictionary "keys"
    df_array = [df, df_goods, df_services] # names of DataFrames used

    for i in 1:length(df_names)
        push!(regression, "reg_1_" * df_names[i] =>
        FixedEffectModels.reg(df_array[i], @formula(log(value) ~ year + fe(country_industry)), Vcov.cluster(:country, :row_sector, :year)))

        push!(regression, "reg_2_" * df_names[i] =>
        FixedEffectModels.reg(df_array[i], @formula(log(value) ~ year + fe(country_industry)), Vcov.cluster(:country, :row_sector, :year),
            contrasts = Dict(:year => DummyCoding(base=1995)))) # treat years as dummies
    end

    reg = RegressionTables.regtable(regression["reg_1_all"], regression["reg_2_all"], regression["reg_1_goods"], regression["reg_2_goods"],
    regression["reg_1_services"], regression["reg_2_services"];
    renderSettings = asciiOutput(), regression_statistics=[:nobs, :r2], print_fe_section=true, estimformat="%0.4f") # output on the console

    return reg
end

reg_tab_6(t_tab6)

# -------------- Table 7 -----------------------------------------------------------------------------------------------------------------------------

t_tab7 = CSV.read("clean/t_tab7.csv", DataFrame)

function reg_tab_7(df::DataFrame)

    df.value .= ifelse.(isinf.(df.value), 0.0, df.value) # treat inf as 0.0
    df.value .= ifelse.(df.value .<= 0.0, 1e-18, df.value) # since log cannot deal with 0
    #subset!(df, :value => ByRow(x -> x > 0.0)) # same results as with treatment above

    df[:, :dest_country_industry] .= df[:, :country] .* "_" .* string.(df[:, :col_sector]) # for SE clustering and fixed effects (purchasing country-industry)

    transform!(df, :row_sector => ByRow(x -> x in goods_sectors ? 1 : 0) => :row_sector_dummy) # dummy for *reporting* sector, "goods" = 1
    df_goods = subset(df, :row_sector_dummy => ByRow(x -> x==1)) # only goods sectors (as origin)
    df_services = subset(df, :row_sector_dummy => ByRow(x -> x!=1)) # only service sectors (as origin)

    regression = Dict{String,Any}() # initializing dictionary to store regression results (Any could be replaced by type FixedEffectModel)
    df_names = ["all", "goods", "services"] # names of dictionary "keys"
    df_array = [df, df_goods, df_services] # names of DataFrames used

    for i in 1:length(df_names)
        push!(regression, "reg_1_" * df_names[i] =>
        FixedEffectModels.reg(df_array[i], @formula(log(value) ~ year + fe(row_sector)&fe(dest_country_industry)), Vcov.cluster(:row_sector, :dest_country_industry, :year)))

        push!(regression, "reg_2_" * df_names[i] =>
        FixedEffectModels.reg(df_array[i], @formula(log(value) ~ year + fe(row_sector)&fe(dest_country_industry)), Vcov.cluster(:row_sector, :dest_country_industry, :year),
            contrasts = Dict(:year => DummyCoding(base=1995)))) # treat years as dummies
    end

    reg = RegressionTables.regtable(regression["reg_1_all"], regression["reg_2_all"], regression["reg_1_goods"], regression["reg_2_goods"],
    regression["reg_1_services"], regression["reg_2_services"];
    renderSettings = asciiOutput(), regression_statistics=[:nobs, :r2], print_fe_section=true, estimformat="%0.4f") # output on the console

    return reg
end

reg_tab_7(t_tab7)


text:
"""
GenX: An Configurable Capacity Expansion Model
Copyright (C) 2021,  Massachusetts Institute of Technology
This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.
A complete copy of the GNU General Public License v2 (GPLv2) is available
in LICENSE.txt.  Users uncompressing this from an archive may not have
received this license file.  If not, see <http://www.gnu.org/licenses/>.
"""

@doc raw"""
	write_charge(path::AbstractString, inputs::Dict, setup::Dict, EP::Model)

Function for writing the charging energy values of the different storage technologies.
"""
function write_charge(path::AbstractString, inputs::Dict, setup::Dict, EP::Model)
	dfGen = inputs["dfGen"]
	G = inputs["G"]     # Number of resources (generators, storage, DR, and DERs)
	T = inputs["T"]     # Number of time steps (hours)
	# Power withdrawn to charge each resource in each time step
	dfCharge = DataFrame(Resource = inputs["RESOURCES"], Zone = dfGen[!,:Zone], AnnualSum = Array{Union{Missing,Float32}}(undef, G))
	charge = zeros(G,T)
	for i in 1:G
		if setup["ParameterScale"] ==1
			if i in inputs["STOR_ALL"]
				charge[i,:] = value.(EP[:vCHARGE])[i,:] * ModelScalingFactor
			elseif i in inputs["FLEX"]
				charge[i,:] = value.(EP[:vCHARGE_FLEX])[i,:] * ModelScalingFactor
			end
		else
			if i in inputs["STOR_ALL"]
				charge[i,:] = value.(EP[:vCHARGE])[i,:]
			elseif i in inputs["FLEX"]
				charge[i,:] = value.(EP[:vCHARGE_FLEX])[i,:]
			end
		end
		dfCharge[!,:AnnualSum][i] = sum(inputs["omega"].* charge[i,:])
	end
	dfCharge = hcat(dfCharge, DataFrame(charge, :auto))
	auxNew_Names=[Symbol("Resource");Symbol("Zone");Symbol("AnnualSum");[Symbol("t$t") for t in 1:T]]
	rename!(dfCharge,auxNew_Names)
	total = DataFrame(["Total" 0 sum(dfCharge[!,:AnnualSum]) fill(0.0, (1,T))], :auto)
	for t in 1:T
		total[:,t+3] .= sum(dfCharge[:,Symbol("t$t")][union(inputs["STOR_ALL"],inputs["FLEX"])])
	end
	rename!(total,auxNew_Names)
	dfCharge = vcat(dfCharge, total)
	CSV.write(joinpath(path, "charge.csv"), dftranspose(dfCharge, false), writeheader=false)
	return dfCharge
end

max_stars_repo_name: Betristor/LTESOM

max_issues_repo_name: Betristor/LTESOM

max_forks_repo_name: Betristor/LTESOM


text:
function add_variable_cost!(
    container::OptimizationContainer,
    ::U,
    devices::IS.FlattenIteratorWrapper{T},
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    for d in devices
        op_cost_data = PSY.get_operation_cost(d)
        _add_variable_cost_to_objective!(container, U(), d, op_cost_data, V())
    end
    return
end

function add_variable_cost!(
    container::OptimizationContainer,
    ::U,
    service::T,
    ::V,
) where {T <: PSY.ReserveDemandCurve, U <: VariableType, V <: StepwiseCostReserve}
    _add_variable_cost_to_objective!(container, U(), service, V())
    return
end

function add_shut_down_cost!(
    container::OptimizationContainer,
    ::U,
    devices::IS.FlattenIteratorWrapper{T},
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    multiplier = objective_function_multiplier(U(), V())
    for d in devices
        op_cost_data = PSY.get_operation_cost(d)
        cost_term = shut_down_cost(op_cost_data, d, V())
        iszero(cost_term) && continue
        for t in get_time_steps(container)
            _add_proportional_term!(container, U(), d, cost_term * multiplier, t)
        end
    end
    return
end

function add_proportional_cost!(
    container::OptimizationContainer,
    ::U,
    devices::IS.FlattenIteratorWrapper{T},
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    multiplier = objective_function_multiplier(U(), V())
    for d in devices
        op_cost_data = PSY.get_operation_cost(d)
        cost_term = proportional_cost(op_cost_data, U(), d, V())
        iszero(cost_term) && continue
        for t in get_time_steps(container)
            _add_proportional_term!(container, U(), d, cost_term * multiplier, t)
        end
    end
    return
end

function add_proportional_cost!(
    container::OptimizationContainer,
    ::U,
    devices::IS.FlattenIteratorWrapper{T},
    ::V,
) where {T <: PSY.ThermalGen, U <: OnVariable, V <: AbstractCompactUnitCommitment}
    multiplier = objective_function_multiplier(U(), V())
    for d in devices
        op_cost_data = PSY.get_operation_cost(d)
        cost_term = proportional_cost(op_cost_data, U(), d, V())
        iszero(cost_term) && continue
        for t in get_time_steps(container)
            exp = _add_proportional_term!(container, U(), d, cost_term * multiplier, t)
            add_to_expression!(container, ProductionCostExpression, exp, d, t)
        end
    end
    return
end

function add_proportional_cost!(
    container::OptimizationContainer,
    ::U,
    devices::IS.FlattenIteratorWrapper{T},
    ::V,
) where {
    T <: PSY.Storage,
    U <: Union{ActivePowerInVariable, ActivePowerOutVariable},
    V <: AbstractDeviceFormulation,
}
    multiplier = objective_function_multiplier(U(), V())
    for d in devices
        for t in get_time_steps(container)
            _add_proportional_term!(container, U(), d, COST_EPSILON * multiplier, t)
        end
    end
    return
end

function add_proportional_cost!(
    container::OptimizationContainer,
    ::U,
    devices::IS.FlattenIteratorWrapper{T},
    ::V,
) where {
    T <: PSY.Component,
    U <: Union{EnergySurplusVariable, EnergyShortageVariable},
    V <: AbstractDeviceFormulation,
}
    base_p = get_base_power(container)
    multiplier = objective_function_multiplier(U(), V())
    for d in devices
        op_cost_data = PSY.get_operation_cost(d)
        cost_term = proportional_cost(op_cost_data, U(), d, V())
        iszero(cost_term) && continue
        for t in get_time_steps(container)
            _add_proportional_term!(container, U(), d, cost_term * multiplier * base_p, t)
        end
    end
    return
end

function add_proportional_cost!(
    container::OptimizationContainer,
    ::U,
    service::T,
    ::V,
) where {
    T <: Union{PSY.Reserve, PSY.ReserveNonSpinning},
    U <: ActivePowerReserveVariable,
    V <: AbstractReservesFormulation,
}
    base_p = get_base_power(container)
    reserve_variable = get_variable(container, U(), T, PSY.get_name(service))
    for index in Iterators.product(axes(reserve_variable)...)
        add_to_objective_invariant_expression!(
            container,
            DEFAULT_RESERVE_COST / base_p * reserve_variable[index...],
        )
    end
    return
end

function add_proportional_cost!(
    container::OptimizationContainer,
    ::U,
    areas::IS.FlattenIteratorWrapper{T},
    ::PIDSmoothACE,
) where {T <: PSY.Area, U <: LiftVariable}
    lift_variable = get_variable(container, U(), T)
    for index in Iterators.product(axes(lift_variable)...)
        add_to_objective_invariant_expression!(
            container,
            SERVICES_SLACK_COST * lift_variable[index...],
        )
    end
    return
end

function _add_variable_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Component,
    op_cost::PSY.OperationalCost,
    ::U,
) where {T <: VariableType, U <: AbstractDeviceFormulation}
    variable_cost_data = variable_cost(op_cost, T(), component, U())
    _add_variable_cost_to_objective!(container, T(), component, variable_cost_data, U())
    return
end

function _add_variable_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Component,
    op_cost::PSY.MarketBidCost,
    ::U,
) where {T <: VariableType, U <: AbstractDeviceFormulation}
    component_name = PSY.get_name(component)
    @debug "Market Bid" _group = LOG_GROUP_COST_FUNCTIONS component_name
    time_steps = get_time_steps(container)
    initial_time = get_initial_time(container)
    variable_cost_forecast = PSY.get_variable_cost(
        component,
        op_cost;
        start_time=initial_time,
        len=length(time_steps),
    )
    variable_cost_forecast_values = TimeSeries.values(variable_cost_forecast)
    parameter_container = _get_cost_function_parameter_container(
        container,
        CostFunctionParameter(),
        component,
        T(),
        U(),
        eltype(variable_cost_forecast_values),
    )
    pwl_cost_expressions =
        _add_pwl_term!(container, component, variable_cost_forecast_values, T(), U())
    jump_model = get_jump_model(container)
    for t in time_steps
        set_parameter!(
            parameter_container,
            jump_model,
            PSY.get_cost(variable_cost_forecast_values[t]),
            # Using 1.0 here since we want to reuse the existing code that adds the mulitpler
            #  of base power times the time delta.
            1.0,
            component_name,
            t,
        )
        add_to_expression!(
            container,
            ProductionCostExpression,
            pwl_cost_expressions[t],
            component,
            t,
        )
        add_to_objective_variant_expression!(container, pwl_cost_expressions[t])
    end

    # Service Cost Bid
    ancillary_services = PSY.get_ancillary_services(op_cost)
    for service in ancillary_services
        _add_service_bid_cost!(container, component, service)
    end
    return
end

function _add_variable_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Reserve,
    ::U,
) where {T <: VariableType, U <: StepwiseCostReserve}
    component_name = PSY.get_name(component)
    @debug "PWL Variable Cost" _group = LOG_GROUP_COST_FUNCTIONS component_name
    # If array is full of tuples with zeros return 0.0
    time_steps = get_time_steps(container)
    initial_time = get_initial_time(container)
    variable_cost_forecast = get_time_series(container, component, "variable_cost")
    variable_cost_forecast_values = TimeSeries.values(variable_cost_forecast)
    variable_cost_forecast_values = map(PSY.VariableCost, variable_cost_forecast_values)
    parameter_container = _get_cost_function_parameter_container(
        container,
        CostFunctionParameter(),
        component,
        T(),
        U(),
        eltype(variable_cost_forecast_values),
    )
    pwl_cost_expressions =
        _add_pwl_term!(container, component, variable_cost_forecast_values, T(), U())
    jump_model = get_jump_model(container)
    for t in time_steps
        set_parameter!(
            parameter_container,
            jump_model,
            PSY.get_cost(variable_cost_forecast_values[t]),
            # Using 1.0 here since we want to reuse the existing code that adds the mulitpler
            #  of base power times the time delta.
            1.0,
            component_name,
            t,
        )
        add_to_objective_variant_expression!(container, pwl_cost_expressions[t])
    end
    return
end

function add_start_up_cost!(
    container::OptimizationContainer,
    ::U,
    devices::IS.FlattenIteratorWrapper{T},
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    for d in devices
        op_cost_data = PSY.get_operation_cost(d)
        _add_start_up_cost_to_objective!(container, U(), d, op_cost_data, V())
    end
    return
end

function _add_start_up_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Component,
    op_cost::PSY.OperationalCost,
    ::U,
) where {T <: VariableType, U <: AbstractDeviceFormulation}
    cost_term = start_up_cost(op_cost, component, U())
    iszero(cost_term) && return
    multiplier = objective_function_multiplier(T(), U())
    for t in get_time_steps(container)
        _add_proportional_term!(container, T(), component, cost_term * multiplier, t)
    end
    return
end

const MULTI_START_COST_MAP = Dict{DataType, Int}(
    HotStartVariable => 1,
    WarmStartVariable => 2,
    ColdStartVariable => 3,
)

function _add_start_up_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Component,
    op_cost::Union{PSY.MultiStartCost, PSY.MarketBidCost},
    ::U,
) where {T <: VariableType, U <: ThermalMultiStartUnitCommitment}
    cost_terms = start_up_cost(op_cost, component, U())
    cost_term = cost_terms[MULTI_START_COST_MAP[T]]
    iszero(cost_term) && return
    multiplier = objective_function_multiplier(T(), U())
    for t in get_time_steps(container)
        _add_proportional_term!(container, T(), component, cost_term * multiplier, t)
    end
    return
end

_get_cost_function_data_type(::Type{PSY.VariableCost{T}}) where {T} = T

function _get_cost_function_parameter_container(
    container::OptimizationContainer,
    ::S,
    component::T,
    ::U,
    ::V,
    cost_type::DataType,
) where {
    S <: ObjectiveFunctionParameter,
    T <: PSY.Component,
    U <: VariableType,
    V <: Union{AbstractDeviceFormulation, AbstractServiceFormulation},
}
    if has_container_key(container, S, T)
        return get_parameter(container, S(), T)
    else
        container_axes = axes(get_variable(container, U(), T))
        if has_container_key(container, OnStatusParameter, T)
            sos_val = SOSStatusVariable.PARAMETER
        else
            sos_val = sos_status(component, V())
        end
        return add_param_container!(
            container,
            S(),
            T,
            sos_val,
            U,
            uses_compact_power(component, V()),
            _get_cost_function_data_type(cost_type),
            container_axes...,
        )
    end
end

function _add_service_bid_cost!(
    container::OptimizationContainer,
    component::PSY.Component,
    service::PSY.Reserve{T},
) where {T <: PSY.ReserveDirection}
    time_steps = get_time_steps(container)
    initial_time = get_initial_time(container)
    base_power = get_base_power(container)
    forecast_data = PSY.get_services_bid(
        component,
        PSY.get_operation_cost(component),
        service;
        start_time=initial_time,
        len=length(time_steps),
    )
    forecast_data_values = PSY.get_cost.(TimeSeries.values(forecast_data)) .* base_power
    reserve_variable = get_variable(container, U(), T, PSY.get_name(service))
    component_name = PSY.get_name(component)
    for t in time_steps
        add_to_objective_invariant_expression!(
            container,
            forecast_data_values[t] * reserve_variable[component_name, t],
        )
    end
end

function _add_service_bid_cost!(::OptimizationContainer, ::PSY.Component, ::PSY.Service) end

function _add_service_bid_cost!(
    ::OptimizationContainer,
    ::PSY.Component,
    service::PSY.ReserveDemandCurve{T},
) where {T <: PSY.ReserveDirection}
    error(
        "The Current version doesn't supports cost bid for ReserveDemandCurve services, \\
        please change the forecast data for $(PSY.get_name(service)) \\
        and open a feature request",
    )
    return
end

"""
Adds to the cost function cost terms for sum of variables with common factor to be used for cost expression for optimization_container model.

# Arguments

  - container::OptimizationContainer : the optimization_container model built in PowerSimulations
  - var_key::VariableKey: The variable name
  - component_name::String: The component_name of the variable container
  - cost_component::PSY.VariableCost{Float64} : container for cost to be associated with variable
"""
function _add_variable_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Component,
    cost_component::PSY.VariableCost{Float64},
    ::U,
) where {T <: VariableType, U <: AbstractDeviceFormulation}
    multiplier = objective_function_multiplier(T(), U())
    base_power = get_base_power(container)
    cost_data = PSY.get_cost(cost_component)
    resolution = get_resolution(container)
    dt = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    for time_period in get_time_steps(container)
        linear_cost = _add_proportional_term!(
            container,
            T(),
            component,
            cost_data * multiplier * base_power * dt,
            time_period,
        )
        add_to_expression!(
            container,
            ProductionCostExpression,
            linear_cost,
            component,
            time_period,
        )
    end
    return
end

@doc raw"""
Adds to the cost function cost terms for sum of variables with common factor to be used for cost expression for optimization_container model.

# Equation

``` gen_cost = dt*sign*(sum(variable.^2)*cost_data[1] + sum(variable)*cost_data[2]) ```

# LaTeX

`` cost = dt\times sign (sum_{i\in I} c_1 v_i^2 + sum_{i\in I} c_2 v_i ) ``

for quadratic factor large enough. If the first term of the quadratic objective is 0.0, adds a
linear cost term `sum(variable)*cost_data[2]`

# Arguments

* container::OptimizationContainer : the optimization_container model built in PowerSimulations
* var_key::VariableKey: The variable name
* component_name::String: The component_name of the variable container
* cost_component::PSY.VariableCost{NTuple{2, Float64}} : container for quadratic and linear factors
"""
function _add_variable_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Component,
    cost_component::PSY.VariableCost{NTuple{2, Float64}},
    ::U,
) where {T <: VariableType, U <: AbstractDeviceFormulation}
    multiplier = objective_function_multiplier(T(), U())
    base_power = get_base_power(container)
    cost_data = PSY.get_cost(cost_component)
    resolution = get_resolution(container)
    dt = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    for time_period in get_time_steps(container)
        if cost_data[1] >= eps()
            cost_term = _add_quadratic_term!(
                container,
                T(),
                component,
                cost_data,
                base_power,
                multiplier * dt,
                time_period,
            )
        else
            cost_term = _add_proportional_term!(
                container,
                T(),
                component,
                cost_data[2] * multiplier * base_power * dt,
                time_period,
            )
        end
        add_to_expression!(
            container,
            ProductionCostExpression,
            cost_term,
            component,
            time_period,
        )
    end
    return
end

"""
Creates piecewise linear cost function using a sum of variables and expression with sign and time step included.

# Arguments

  - container::OptimizationContainer : the optimization_container model built in PowerSimulations
  - var_key::VariableKey: The variable name
  - component_name::String: The component_name of the variable container
  - cost_component::PSY.VariableCost{Vector{NTuple{2, Float64}}}
"""
function _add_variable_cost_to_objective!(
    container::OptimizationContainer,
    ::T,
    component::PSY.Component,
    cost_component::PSY.VariableCost{Vector{NTuple{2, Float64}}},
    ::U,
) where {T <: VariableType, U <: AbstractDeviceFormulation}
    component_name = PSY.get_name(component)
    @debug "PWL Variable Cost" _group = LOG_GROUP_COST_FUNCTIONS component_name
    # If array is full of tuples with zeros return 0.0
    cost_data = PSY.get_cost(cost_component)
    if all(iszero.(last.(cost_data)))
        @debug "All cost terms for component $(component_name) are 0.0" _group =
            LOG_GROUP_COST_FUNCTIONS
        return
    end
    pwl_cost_expressions = _add_pwl_term!(container, component, cost_data, T(), U())
    for t in get_time_steps(container)
        add_to_expression!(
            container,
            ProductionCostExpression,
            pwl_cost_expressions[t],
            component,
            t,
        )
        add_to_objective_invariant_expression!(container, pwl_cost_expressions[t])
    end
    return
end

"""
Returns True/False depending on compatibility of the cost data with the convex implementation method
"""
function _slope_convexity_check(slopes::Vector{Float64})
    flag = true
    for ix in 1:(length(slopes) - 1)
        if slopes[ix] > slopes[ix + 1]
            @debug slopes _group = LOG_GROUP_COST_FUNCTIONS
            return flag = false
        end
    end
    return flag
end

function _check_pwl_compact_data(
    min::Float64,
    max::Float64,
    data::Vector{Tuple{Float64, Float64}},
    base_power::Float64,
)
    return isapprox(max - min, data[end][2] / base_power) && iszero(data[1][2])
end

function _check_pwl_compact_data(
    d::PSY.Component,
    data::Vector{Tuple{Float64, Float64}},
    base_power::Float64,
)
    min = PSY.get_active_power_limits(d).min
    max = PSY.get_active_power_limits(d).max
    return _check_pwl_compact_data(min, max, data, base_power)
end

function _get_sos_value(
    container::OptimizationContainer,
    ::Type{V},
    component::T,
) where {T <: PSY.Component, V <: AbstractDeviceFormulation}
    if has_container_key(container, OnStatusParameter, T)
        sos_val = SOSStatusVariable.PARAMETER
    else
        sos_val = sos_status(component, V())
    end
    return sos_val
end

function _get_sos_value(
    container::OptimizationContainer,
    ::Type{V},
    component::T,
) where {T <: PSY.Component, V <: AbstractServiceFormulation}
    return SOSStatusVariable.NO_VARIABLE
end

function _add_pwl_term!(
    container::OptimizationContainer,
    component::T,
    cost_data::Vector{PSY.VariableCost{Float64}},
    ::U,
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    multiplier = objective_function_multiplier(U(), V())
    resolution = get_resolution(container)
    dt = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    base_power = get_base_power(container)
    # Re-scale breakpoints by Basepower
    time_steps = get_time_steps(container)
    cost_expressions = Vector{JuMP.AffExpr}(undef, time_steps[end])
    for t in time_steps
        proportial_value = PSY.get_cost(cost_data[t]) * multiplier * base_power * dt
        cost_expressions[t] =
            _add_proportional_term!(container, U(), component, proportial_value, t)
    end
    return cost_expressions
end

"""
Add PWL cost terms for data coming from the MarketBidCost
"""
function _add_pwl_term!(
    container::OptimizationContainer,
    component::T,
    cost_data::Vector{PSY.VariableCost{Vector{Tuple{Float64, Float64}}}},
    ::U,
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    multiplier = objective_function_multiplier(U(), V())
    resolution = get_resolution(container)
    dt = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    base_power = get_base_power(container)
    # Re-scale breakpoints by Basepower
    name = PSY.get_name(component)
    time_steps = get_time_steps(container)
    pwl_cost_expressions = Vector{JuMP.AffExpr}(undef, time_steps[end])
    sos_val = _get_sos_value(container, V, component)
    for t in time_steps
        data = PSY.get_cost(cost_data[t])
        is_power_data_compact = _check_pwl_compact_data(component, data, base_power)
        if !uses_compact_power(component, V()) && is_power_data_compact
            error(
                "The data provided is not compatible with formulation $V. Use a formulation compatible with Compact Cost Functions",
            )
            # data = _convert_to_full_variable_cost(data, component)
        elseif uses_compact_power(component, V()) && !is_power_data_compact
            data = _convert_to_compact_variable_cost(data)
        else
            @debug uses_compact_power(component, V()) name T V
            @debug is_power_data_compact name T V
        end
        slopes = PSY.get_slopes(data)
        # First element of the return is the average cost at P_min.
        # Shouldn't be passed for convexity check
        is_convex = _slope_convexity_check(slopes[2:end])
        break_points = map(x -> last(x), data) ./ base_power
        _add_pwl_variables!(container, T, name, t, data)
        _add_pwl_constraint!(container, component, U(), break_points, sos_val, t)
        if !is_convex
            _add_pwl_sos_constraint!(container, component, U(), break_points, sos_val, t)
        end
        pwl_cost = _get_pwl_cost_expression(container, component, t, data, multiplier * dt)
        pwl_cost_expressions[t] = pwl_cost
    end
    return pwl_cost_expressions
end

function _add_pwl_term!(
    container::OptimizationContainer,
    component::T,
    cost_data::Vector{PSY.VariableCost{Vector{Tuple{Float64, Float64}}}},
    ::U,
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractServiceFormulation}
    multiplier = objective_function_multiplier(U(), V())
    resolution = get_resolution(container)
    dt = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    base_power = get_base_power(container)
    # Re-scale breakpoints by Basepower
    name = PSY.get_name(component)
    time_steps = get_time_steps(container)
    pwl_cost_expressions = Vector{JuMP.AffExpr}(undef, time_steps[end])
    sos_val = _get_sos_value(container, V, component)
    for t in time_steps
        data = PSY.get_cost(cost_data[t])
        slopes = PSY.get_slopes(data)
        # Shouldn't be passed for convexity check
        is_convex = false
        break_points = map(x -> last(x), data) ./ base_power
        _add_pwl_variables!(container, T, name, t, data)
        _add_pwl_constraint!(container, component, U(), break_points, sos_val, t)
        if !is_convex
            _add_pwl_sos_constraint!(container, component, U(), break_points, sos_val, t)
        end
        pwl_cost = _get_pwl_cost_expression(container, component, t, data, multiplier * dt)
        pwl_cost_expressions[t] = pwl_cost
    end
    return pwl_cost_expressions
end

"""
Add PWL cost terms for data coming from a constant PWL cost function
"""
function _add_pwl_term!(
    container::OptimizationContainer,
    component::T,
    data::Vector{NTuple{2, Float64}},
    ::U,
    ::V,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    multiplier = objective_function_multiplier(U(), V())
    resolution = get_resolution(container)
    dt = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    base_power = get_base_power(container)
    # Re-scale breakpoints by Basepower
    name = PSY.get_name(component)

    is_power_data_compact = _check_pwl_compact_data(component, data, base_power)

    if !uses_compact_power(component, V()) && is_power_data_compact
        error(
            "The data provided is not compatible with formulation $V. Use a formulation compatible with Compact Cost Functions",
        )
        # data = _convert_to_full_variable_cost(data, component)
    elseif uses_compact_power(component, V()) && !is_power_data_compact
        data = _convert_to_compact_variable_cost(data)
    else
        @debug uses_compact_power(component, V()) name T V
        @debug is_power_data_compact name T V
    end

    slopes = PSY.get_slopes(data)
    # First element of the return is the average cost at P_min.
    # Shouldn't be passed for convexity check
    is_convex = _slope_convexity_check(slopes[2:end])
    time_steps = get_time_steps(container)
    pwl_cost_expressions = Vector{JuMP.AffExpr}(undef, time_steps[end])
    break_points = map(x -> last(x), data) ./ base_power
    sos_val = _get_sos_value(container, V, component)
    for t in time_steps
        _add_pwl_variables!(container, T, name, t, data)
        _add_pwl_constraint!(container, component, U(), break_points, sos_val, t)
        if !is_convex
            _add_pwl_sos_constraint!(container, component, U(), break_points, sos_val, t)
        end
        pwl_cost = _get_pwl_cost_expression(container, component, t, data, multiplier * dt)
        pwl_cost_expressions[t] = pwl_cost
    end
    return pwl_cost_expressions
end

function _add_pwl_term!(
    container::OptimizationContainer,
    component::T,
    data::Vector{NTuple{2, Float64}},
    ::U,
    ::V,
) where {T <: PSY.ThermalGen, U <: VariableType, V <: ThermalDispatchNoMin}
    multiplier = objective_function_multiplier(U(), V())
    resolution = get_resolution(container)
    dt = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    component_name = PSY.get_name(component)
    @debug "PWL cost function detected for device $(component_name) using $V"
    base_power = get_base_power(container)
    slopes = PSY.get_slopes(data)
    if any(slopes .< 0) || !_slope_convexity_check(slopes[2:end])
        throw(
            IS.InvalidValue(
                "The PWL cost data provided for generator $(component_name) is not compatible with $U.",
            ),
        )
    end

    if _check_pwl_compact_data(component, data, base_power)
        error("The data provided is not compatible with formulation $V. \\
              Use a formulation compatible with Compact Cost Functions")
    end

    if slopes[1] != 0.0
        @debug "PWL has no 0.0 intercept for generator $(component_name)"
        # adds a first intercept a x = 0.0 and Y below the intercept of the first tuple to make convex equivalent
        first_pair = data[1]
        intercept_point = (0.0, first_pair[2] - COST_EPSILON)
        data = vcat(intercept_point, data)
        @assert _slope_convexity_check(slopes)
    end

    time_steps = get_time_steps(container)
    pwl_cost_expressions = Vector{JuMP.AffExpr}(undef, time_steps[end])
    break_points = map(x -> last(x), data) ./ base_power
    sos_val = _get_sos_value(container, V, component)
    for t in time_steps
        _add_pwl_variables!(container, T, component_name, t, data)
        _add_pwl_constraint!(container, component, U(), break_points, sos_val, t)
        pwl_cost = _get_pwl_cost_expression(container, component, t, data, multiplier * dt)
        pwl_cost_expressions[t] = pwl_cost
    end
    return pwl_cost_expressions
end

function _add_pwl_variables!(
    container::OptimizationContainer,
    ::Type{T},
    component_name::String,
    time_period::Int,
    cost_data::Vector{NTuple{2, Float64}},
) where {T <: PSY.Component}
    var_container = lazy_container_addition!(container, PieceWiseLinearCostVariable(), T)
    pwlvars = Array{JuMP.VariableRef}(undef, length(cost_data))
    for i in 1:length(cost_data)
        pwlvars[i] =
            var_container[(component_name, i, time_period)] = JuMP.@variable(
                get_jump_model(container),
                base_name = "PieceWiseLinearCostVariable_$(component_name)_{pwl_$(i), $time_period}",
                lower_bound = 0.0,
                upper_bound = 1.0
            )
    end
    return pwlvars
end

function _add_pwl_constraint!(
    container::OptimizationContainer,
    component::T,
    ::U,
    break_points::Vector{Float64},
    sos_status::SOSStatusVariable,
    period::Int,
) where {T <: PSY.Component, U <: VariableType}
    variables = get_variable(container, U(), T)
    const_container = lazy_container_addition!(
        container,
        PieceWiseLinearCostConstraint(),
        T,
        axes(variables)...,
    )
    len_cost_data = length(break_points)
    jump_model = get_jump_model(container)
    pwl_vars = get_variable(container, PieceWiseLinearCostVariable(), T)
    name = PSY.get_name(component)
    const_container[name, period] = JuMP.@constraint(
        jump_model,
        variables[name, period] ==
        sum(pwl_vars[name, ix, period] * break_points[ix] for ix in 1:len_cost_data)
    )

    if sos_status == SOSStatusVariable.NO_VARIABLE
        bin = 1.0
        @debug "Using Piecewise Linear cost function but no variable/parameter ref for ON status is passed. Default status will be set to online (1.0)" _group =
            LOG_GROUP_COST_FUNCTIONS

    elseif sos_status == SOSStatusVariable.PARAMETER
        bin = get_parameter(container, OnStatusParameter(), T).parameter_array[name, period]
        @debug "Using Piecewise Linear cost function with parameter OnStatusParameter, $T" _group =
            LOG_GROUP_COST_FUNCTIONS
    elseif sos_status == SOSStatusVariable.VARIABLE
        bin = get_variable(container, OnVariable(), T)[name, period]
        @debug "Using Piecewise Linear cost function with variable OnVariable $T" _group =
            LOG_GROUP_COST_FUNCTIONS
    else
        @assert false
    end

    JuMP.@constraint(
        jump_model,
        sum(pwl_vars[name, i, period] for i in 1:len_cost_data) == bin
    )
    return
end

function _add_pwl_sos_constraint!(
    container::OptimizationContainer,
    component::T,
    ::U,
    break_points::Vector{Float64},
    sos_status::SOSStatusVariable,
    period::Int,
) where {T <: PSY.Component, U <: VariableType}
    name = PSY.get_name(component)
    @warn(
        "The cost function provided for $(name) is not compatible with a linear PWL cost function.
  An SOS-2 formulation will be added to the model. This will result in additional binary variables."
    )

    jump_model = get_jump_model(container)
    pwl_vars = get_variable(container, PieceWiseLinearCostVariable(), T)
    bp_count = length(break_points)
    pwl_vars_subset = [pwl_vars[name, i, period] for i in 1:bp_count]
    JuMP.@constraint(jump_model, pwl_vars_subset in MOI.SOS2(collect(1:bp_count)))
    return
end

function _get_pwl_cost_expression(
    container::OptimizationContainer,
    component::T,
    time_period::Int,
    cost_data::Vector{NTuple{2, Float64}},
    multiplier::Float64,
) where {T <: PSY.Component}
    name = PSY.get_name(component)
    pwl_var_container = get_variable(container, PieceWiseLinearCostVariable(), T)
    gen_cost = JuMP.AffExpr(0.0)
    slopes = PSY.get_slopes(cost_data)
    upb = PSY.get_breakpoint_upperbounds(cost_data)
    for i in 1:length(cost_data)
        JuMP.add_to_expression!(
            gen_cost,
            cost_data[i][1] * multiplier * pwl_var_container[(name, i, time_period)],
        )
    end
    return gen_cost
end

function _get_no_load_cost(
    component::T,
    ::V,
    ::U,
) where {T <: PSY.Component, U <: VariableType, V <: AbstractDeviceFormulation}
    return no_load_cost(PSY.get_operation_cost(component), U(), component, V())
end

function _convert_to_compact_variable_cost(
    var_cost::Vector{NTuple{2, Float64}},
    no_load_cost::Float64,
    p_min::Float64,
)
    return [(c - no_load_cost, pp - p_min) for (c, pp) in var_cost]
end

function _convert_to_compact_variable_cost(var_cost::Vector{NTuple{2, Float64}})
    no_load_cost, p_min = var_cost[1]
    return _convert_to_compact_variable_cost(var_cost, no_load_cost, p_min)
end

function _add_proportional_term!(
    container::OptimizationContainer,
    ::T,
    component::U,
    linear_term::Float64,
    time_period::Int,
) where {T <: VariableType, U <: PSY.Component}
    component_name = PSY.get_name(component)
    @debug "Linear Variable Cost" _group = LOG_GROUP_COST_FUNCTIONS component_name
    variable = get_variable(container, T(), U)[component_name, time_period]
    lin_cost = variable * linear_term
    add_to_objective_invariant_expression!(container, lin_cost)
    return lin_cost
end

function _add_quadratic_term!(
    container::OptimizationContainer,
    ::T,
    component::U,
    q_terms::NTuple{2, Float64},
    var_multiplier::Float64,
    expression_multiplier::Float64,
    time_period::Int,
) where {T <: VariableType, U <: PSY.Component}
    component_name = PSY.get_name(component)
    @debug "$component_name Quadratic Variable Cost" _group = LOG_GROUP_COST_FUNCTIONS component_name
    var = get_variable(container, T(), U)[component_name, time_period]
    q_cost_ = (var * var_multiplier) .^ 2 * q_terms[1] + var * var_multiplier * q_terms[2]
    q_cost = q_cost_ * expression_multiplier
    add_to_objective_invariant_expression!(container, q_cost)
    return q_cost
end

function _add_quadratic_term!(
    container::OptimizationContainer,
    ::T,
    component::U,
    q_terms::NTuple{2, Float64},
    var_multiplier::Float64,
    expression_multiplier::Float64,
    time_period::Int,
) where {T <: PowerAboveMinimumVariable, U <: PSY.ThermalGen}
    component_name = PSY.get_name(component)
    p_min = PSY.get_active_power_limits(component).min
    @debug "$component_name Quadratic Variable Cost" _group = LOG_GROUP_COST_FUNCTIONS component_name
    var = get_variable(container, T(), U)[component_name, time_period]
    q_cost_ =
        (var * var_multiplier) .^ 2 * q_terms[1] +
        var * var_multiplier * (q_terms[2] + 2 * q_terms[1] * p_min)
    q_cost = q_cost_ * expression_multiplier
    add_to_objective_invariant_expression!(container, q_cost)
    return q_cost
end