plot_online.py

import os

import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

import matplotlib as mpl
mpl.rcParams['pdf.fonttype'] = 42

np.random.seed(3)


class Plot:
    def __init__(self, res_dir, data_dir='./plot data'):
        self.res_dir = res_dir
        self.data_dir = data_dir

    def load_data(self, test_dir, key, verbose=False):

        # -- load data
        res = {}
        for idx, folder in enumerate([folders for _, folders, _ in os.walk(os.path.join(self.res_dir, test_dir))][0][:20]):

            path = os.path.join(self.res_dir, test_dir, folder)

            if idx == 0:
                res = [np.nan_to_num(np.loadtxt(os.path.join(path, f'{key}.txt')))]
            else:
                res.append(np.nan_to_num(np.loadtxt(os.path.join(path, f'{key}.txt'))))

        return np.array(res)

    def get_quants(self, res, n_boot=500):

        # -- bootstrap sampling
        stat = None
        for i in range(n_boot):
            idx = np.random.choice(range(len(res)), len(res))
            if stat is None:
                np.mean(res[idx], axis=0)
                stat = [np.mean(res[idx], axis=0)]
            else:
                stat.append(np.mean(res[idx], axis=0))

        # -- compute quantiles
        stat = np.array(stat)

        quants = np.quantile(stat, [0.01, 0.5, 0.99], axis=0)
        quants[1] = np.mean(res, axis=0)

        return quants

    def plot_performance(self, test_dirs, fig_name, label=None, color=None, color_bar=None, edgecolor_bar=None, save=True):

        plt.subplots(figsize=(16.7, 4.37))

        # -- plot accuracy
        plt.subplot(1, 2, 1)
        for i, dir_ in enumerate(test_dirs):

            # -- load data
            res = self.load_data(dir_, key='acc')

            # -- get quantiles
            quants = self.get_quants(res)

            # -- plot
            plt.fill_between(range(len(res.transpose())), quants[0, :], quants[2, :], color=color_bar[i], edgecolor=edgecolor_bar[i], lw=0.1)
            plt.plot(range(len(res.transpose())), quants[1, :], color=color[i], label=label[i])

        plt.xlabel('No. training data')
        plt.ylabel('Accuracy')
        plt.yticks([0, .2, .4, .6, .8, 1])
        plt.xlim([-5, 5001])
        plt.ylim([0, 1])
        plt.legend(loc='lower right')

        sns.despine()
        plt.tight_layout()
        plt.savefig(f'./results/{fig_name}', bbox_inches='tight')
        plt.close()

    def plot_multi(self, test_dir, key, fig_name, label_func=None, figsize=(16.7, 4.37), xlabel='No. training data',
                   ylabel=None, xlim=[-5, 5001], ylim=None, legend_loc='lower right'):

        # -- load data
        res = self.load_data(test_dir, key=key)

        # -- get quantiles
        quants = self.get_quants(res)

        # -- plot accuracy
        plt.subplots(figsize=figsize)
        plt.subplot(1, 2, 1)
        colors = sns.color_palette('muted', n_colors=res.shape[2])
        for i in range(res.shape[2]) if key in ['e_norm', 'e_mean', 'e_std'] else range(res.shape[2] - 1):
            plt.plot(range(res.shape[1]), quants[1, :, i], color=colors[i], label=label_func(i))

        # -- plot setting
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)

        plt.xlim(xlim)
        plt.ylim(ylim)
        plt.legend(loc=legend_loc, ncol=int(np.ceil(res.shape[2]/5)), framealpha=0.3)

        sns.despine()
        plt.tight_layout()
        plt.savefig(f'./results/{fig_name}', bbox_inches='tight')
        plt.close()

    def plot_dist(self, test_dir, fig_name, layers=5, label_func=None):
        plt.figure()
        for i in range(layers):
            # -- load data
            res = self.load_data(test_dir, key=f'conf_l{i + 1}').flatten()

            # -- plot distribution
            sns.kdeplot(res, shade=True, alpha=0.2, label=label_func(i), linewidth=3)

        # -- plot settings
        plt.xlabel('Confusion')
        plt.ylabel('Density')
        plt.ylim([-0.01, 12])
        plt.xlim([-1.2, 1.2])
        plt.legend(ncol=int(np.ceil(layers/5)), framealpha=0.3)

        sns.despine()
        plt.tight_layout()
        plt.savefig(f'./results/{fig_name}', bbox_inches='tight')
        plt.close()


def main():
    # -- plot setting
    sns.set(context='paper', style='white', font_scale=1.85, rc={"lines.linewidth": 2.0})
    sns.set_palette('muted')

    # -- tests: 'Baseline', 'Weight_transport', 'Vanishing_gradients', 'Naive_initialization'
    dir_name = 'Weight_transport'
    test_name = dir_name
    my_plot = Plot('./results')

    if dir_name in ['Baseline', 'Weight_transport', 'Naive_initialization']:
        L = 5
    elif dir_name == 'Vanishing_gradients':
        L = 10

    # -- label functions
    y_label_func = lambda x: r'$y_{}$'.format(x)
    w_label_func = lambda x: r'$w_{{{},{}}}$'.format(x, x+1)
    e_label_func = lambda x: rf'$e_{{{x}}}$'
    None_label_func = lambda x: None

    # -- stats to be plotted
    stats = {
        # 'stat1': ['Stat 1 (norm)', w_label_func, [-10, 15000]],
        # 'stat2': ['Stat 2 (norm)', w_label_func, [-10, 9000]],
        # 'stat3': ['Stat 3 (norm)', w_label_func, [-10, 10000]]
        # 'stat4': ['Stat 4 (I-yyT)', w_label_func, [0, 60]],
        # 'stat6': ['Stat 6 (Oja)', w_label_func, [-10, 3000]],
        # 'stat7': ['Stat 7 (Error, compressive)', w_label_func, [-5, 200]],
        # 'stat8': ['Stat 8 (Error, expanding)', w_label_func, [-2, 60]],
        # 'stat9': ['Stat 9 (tr(cov)/k)', w_label_func, [0, 1.]],
        # 'stat10': ['Stat 10 (|cov|)', w_label_func, None],
        # 'stat11': [r'$\lambda_1(cov)$', w_label_func, [0, 12]],
        # 'stat12': [r'$\|I - cov\|$', w_label_func, [0, 12]],
        # 'stat13': [r'$\lambda_2(cov)$', w_label_func, [0, 12]],
        # 'stat14': [r'$\lambda_3(cov)$', w_label_func, [0, 12]],
        # 'stat15': [r'$\lambda_4(cov)$', w_label_func, [0, 12]],
        # 'stat16': [r'$\Sigma\lambda(cov)$', w_label_func, [0, 40]],
        # 'stat5': ['Stat 5 (norm)', w_label_func, None]
        # 'y_mean': ['Activation mean', y_label_func, [0., 1.5]],  # activation mean
        # 'y_std': ['Activation std', y_label_func, [0., 2.]],  # activation std
        # 'oja_err': ["Oja's error", w_label_func, [0., 300]],  # Oja's error
        # 'orth_err': ['Orthogonality error', w_label_func, [0., 15]],  # orthogonality error
        # 'conf_avg': ['Confusion average', w_label_func, [-1., 1.]],  # confusion average
        # 'conf_min': ['Confusion minimum', w_label_func, [-1., 1.]],  # confusion minimum
        # 'e_mean': ['Activation gradient (error) mean', e_label_func, [-0.006, 0.012]],  # activation gradient (error) mean
        # 'e_std': ['Activation gradient (error) std', e_label_func, [-0.01, 0.12]],  # activation gradient (error) std
        # 'e_norm': ['Activation gradient (error) norm', e_label_func, [-0.01, 0.9]],  # activation gradient (error) norm
        # 'wgrad_mean': ['Weight gradient mean', w_label_func, [-0.0001, 0.0001]],  # weight gradient mean
        # 'wgrad_std': ['Weight gradient std', w_label_func, [-0.000001, 0.0015]],  # weight gradient std
        # 'wgrad_norm': ['Weight gradient norm', w_label_func, [-0.001, 0.2]],  # weight gradient norm
    }
    for i in range(9):
        stats[f'VEs_{i+1}'] = [f'VE_{i+1}', None_label_func, [-0.05, 1.]]
    #         f'VE_s_{i}: ['Stat 9 (tr(cov)/k)', w_label_func, [0, 1.]],

    # -- plot accuracy
    my_plot.plot_performance([f'{dir_name}/Test_1', f'{dir_name}/Test_2'], label=['BP', 'BP+Oja'], color=['k', 'b'],
                             color_bar=[(0, 0, 0, 0.2), (0, 0, 1, 0.2)], edgecolor_bar=[(0, 0, 0, 0.2), (0, 0, 1, 0.2)],
                             fig_name=f'{test_name}_MNIST_acc.png')
    my_plot.plot_performance([f'{dir_name}/Test_3', f'{dir_name}/Test_4'], label=['BP', 'BP+Oja'], color=['k', 'b'],
                             color_bar=[(0, 0, 0, 0.2), (0, 0, 1, 0.2)], edgecolor_bar=[(0, 0, 0, 0.2), (0, 0, 1, 0.2)],
                             fig_name=f'{test_name}_FashionMNIST_acc.png')

    for i, test_desc in enumerate(['MNIST_BP', 'MNIST_BPOja', 'FashionMNIST_BP', 'FashionMNIST_BPOja']):

        # -- plot stats
        for key, (title, label_func, ylim) in stats.items():
            my_plot.plot_multi(test_dir=f'{dir_name}/Test_{i+1}', key=key, fig_name=f'{test_name}_{test_desc}_{key}.png',
                               label_func=label_func, ylabel=title, ylim=ylim)

        # -- plot distributions
        my_plot.plot_dist(test_dir=f'{dir_name}/Test_{i+1}', fig_name=f'{test_name}_{test_desc}_conf_dist.png', layers=L,
                          label_func=w_label_func)


if __name__ == '__main__':
    main()