plotting_paper.jl

# Used solely to generate the figures used in our paper.
using FileIO, JLD2, LaTeXStrings, Plots, Lux, Measures
x_lqr, x_DMD, n, features, learnt_features = FileIO.load(
    "results/ExampleData.jld2",
    "x_lqr",
    "x_DMD",
    "n",
    "features",
    "learnt_features",
)

width = 320
height = 170
plt1 = plot(
    xlim = (0, 1000),
    framestyle = :box,
    yguidefontsize = 7,
    xguidefontsize = 7,
    xtickfontsize = 6,
    ytickfontsize = 6,
    ylabel = L"$\sum_i^3 x^i_{k\mid k-1}$",
    palette = :seaborn_muted,
    xticks = :none,
    legend=:none,
    foreground_color_legend = nothing,
    legendfontsize=5,
    fontfamily = "Computer Modern",
)
plot!(plt1, x_lqr[1, :] .+ x_lqr[2, :] .+ x_lqr[3, :])
plot!(plt1, x_DMD[1, :] .+ x_DMD[2, :] .+ x_DMD[3, :], xformatter = _ -> "")

function calc_Σ_tr(vec::Matrix{Float64})
    trace = zeros(size(vec)[2])
    for i = 1:size(vec)[2]
        trace[i] = tr(back_from_cholesky_Σ(vec[:, i], dyna)[2])
    end
    return trace
end

plt2 = plot(
    xlim = (0, 1000),
    #ylim=(0,1),
    framestyle = :box,
    yguidefontsize = 7,
    xguidefontsize = 7,
    xtickfontsize = 6,
    ytickfontsize = 6,
    xlabel = L"k",
    palette = :seaborn_muted,
    legend=(0.7,0.35),
    legendfontsize=5,
    ylabel = L"tr($\Sigma_{k \mid k-1}$)",
    fontfamily = "Computer Modern",
)
plot!(plt2, calc_Σ_tr(x_lqr), label = L"using CE-LQR: $K$")
plot!(plt2, calc_Σ_tr(x_DMD), label = L"using SOC-LQR: $K_\Psi$")

plot(plt1, plt2, layout = (2, 1), bottom_margin = [-4mm 0mm], size = (width, height))


savefig("figs/closed_loop_sim.pdf")


plt1 = plot(
    xlim = (0, 100),
    ylim = (-8, 8),
    framestyle = :box,
    yguidefontsize = 7,
    xguidefontsize = 7,
    xtickfontsize = 6,
    ytickfontsize = 6,
    #ylabel = "magnitude",
    xlabel = L"k",
    legend=:outertopright,
    legendfontsize=5,
    palette = :seaborn_muted,
    fontfamily = "Computer Modern",
)
plot!(plt1, x_lqr[1, :] .+ x_lqr[2, :] .+ x_lqr[3, :], color = "red", label=L"CE-LQR: $\sum_i x^i_{k \mid k-1}$")
plot!(plt1, sum(x_true2, dims = 1)[:], color = "red", linestyle = :dot, label=L"CE-LQR: $\sum_i x_k^{i,true}$")

plot!(plt1, x_DMD[1, :] .+ x_DMD[2, :] .+ x_DMD[3, :], color = "black", label=L"SOC-LQR: $\sum_i x^i_{k \mid k-1}$")
plot!(plt1, sum(x_true1, dims = 1)[:], color = "black", linestyle = :dot, label=L"SOC-LQR: $\sum_i x_k^{i,true}$")

plot(plt1, size = (width, height/2))


savefig("figs/compared_to_true.pdf")