plot_nonlinear_controller.py

import matplotlib.pyplot as plt
from matplotlib import rc
import numpy as np
import os
import seaborn
import random

if __name__ == "__main__":

    seaborn.set()
    seaborn.set_style("ticks")
    seaborn.set_context("paper")
    rc('text', usetex=True)
    rc('pgf', texsystem="pdflatex")
    rc('font', **{'family':'serif', 'serif':['Times New Roman]']})

    # Training plot

    fig, ax = plt.subplots(2, 1, figsize=[3.2, 3.52]) #3.2, 4

    class Colors:
        blue = (0, 160/255, 255/255)
        orange = (255/255, 160/255, 0/255)
        red = (255/255, 80/255, 0/255)
        green = (0/255, 200/255, 0/255)
        grey = (100/255, 100/255, 100/255)
        lightgrey = (150/255, 150/255, 150/255)

    red = Colors.red
    blue = Colors.blue

    files = [file for file in os.listdir('./data/results') if file.endswith('.npz')]
    random.seed(2)
    random.shuffle(files)

    for file in files:

        data = np.load('./data/results/' + file, allow_pickle=True)
        _, system, algorithm, seed, N = file[:-4].split('_')

        if system == 'controller':
            if N == '10':
             c = 'black'
            elif N == '100':
                c = red
            elif N == '1000':
                c = blue

            if algorithm == 'admm':
                #plt.semilogy(data['res_pri'])
                ax[0].semilogy(data['obj'], color=c)
            else:
                ax[1].semilogy(data['obj'], color=c)

    line0, = ax[0].plot(0, 0, 'black')
    line1, = ax[0].plot(0, 0, color=red)
    line2, = ax[0].plot(0, 0, color=blue)
    line0.set_label('N=10')
    line1.set_label('N=100')
    line2.set_label('N=1000')

    ax[0].grid()
    ax[0].set_ylim([1E1, 9E10])
    ax[0].set_xlim([-15, 300])
    ax[0].set_title('ADMM', loc='left')
    ax[0].set_ylabel('Imitaiton Loss')
    ax[0].set_xlabel('Iteration')
    ax[1].set_ylim([1E1, 9E10])
    ax[1].set_xlim([-50, 1000])
    ax[1].grid()
    ax[1].set_title('Projected Gradient Descent', loc='left')
    ax[1].set_ylabel('Imitaiton Loss')
    ax[1].set_xlabel('Iteration')
    fig.legend(handles=[line0, line1, line2], ncol=3, bbox_to_anchor=[1.05, 0.08],
               frameon=False)
    seaborn.despine(trim="true")
    fig.tight_layout(pad=0, h_pad=1)
    fig.subplots_adjust(bottom=0.19)
    plt.savefig('./data/figures/training2.pgf', pad_inches=0)
    plt.show()