<reponame>mikgroup/subtle_data_crimes<filename>subtle_data_crimes/crime_1_zero_padding/Fig5_statistics_knee_data/Fig5_stats.py
'''
This code loads the results for the CS, DictL and DL algorithms and prepares the graphs shown in Fig 5.

(c) <NAME> (UT Berkeley) (2021)
'''

import numpy as np
import sys
import os
sys.path.append("../")  # add folder above
import matplotlib.pyplot as plt

pad_ratio_vec = np.array([1,1.25,1.5,1.75,2])
sampling_type_vec = np.array([1,2])  # 0 = random, 1 = weak var-dens, 2 = strong var-dens


####################################### load CS results ##########################################
CS_data_filename = './CS/CS_test_results.npz'
CS_container = np.load(CS_data_filename,allow_pickle=True)
CS_test_set_NRMSE_arr = CS_container['CS_test_set_NRMSE_arr']
CS_test_set_SSIM_arr  = CS_container['CS_test_set_SSIM_arr']


####################################### load DictL results ##########################################
DictL_data_filename = 'DictL/DictL_test_results.npz'
DictL_container = np.load(DictL_data_filename,allow_pickle=True)
DictL_test_set_NRMSE_arr = DictL_container['DictL_NRMSE_test_set']
DictL_test_set_SSIM_arr  = DictL_container['DictL_SSIM_test_set']


####################################### load DL results ##########################################
DL_data_filename = './DL/DNN_results_NRMSE_and_SSIM_N_examples_122.npz'
DL_container = np.load(DL_data_filename,allow_pickle=True)
DL_test_set_NRMSE_arr = DL_container['NRMSE_test_set']
DL_test_set_SSIM_arr  = DL_container['SSIM_test_set']

#################### prep for plots #################

# prepare x ticks labels for the NRMSE and SSIM graphs
x = pad_ratio_vec
x_ticks_labels = []
for i in range(pad_ratio_vec.shape[0]):
    pad_num = pad_ratio_vec[i]
    if pad_num % 1 ==0.0:
        pad_str = str(int(pad_num))
        x_ticks_labels.append('x{}'.format(pad_str))
    elif pad_num % 1 == 0.5:
        pad_str = str(pad_num)
        x_ticks_labels.append('x{}'.format(pad_str))
    else:
        pad_str = {}
        x_ticks_labels.append('')

markers=['o', 'd', 's', 'h', '8']
colorslist = ['darkcyan','royalblue','k']

figs_path = 'stats_figs'

if not os.path.exists(figs_path):
    os.makedirs(figs_path)

#######################################################

N_methods = 3 # 0 = CS, 1 = DictL, 2 = DNN

for method_i in range(N_methods):
    if method_i==0:
        method_str = 'CS'
        NRMSE_arr = CS_test_set_NRMSE_arr
        SSIM_arr = CS_test_set_SSIM_arr
    elif method_i==1:
        method_str = 'DictL'
        NRMSE_arr = DictL_test_set_NRMSE_arr
        SSIM_arr = DictL_test_set_SSIM_arr
    elif method_i == 2:
        method_str = 'DL'
        NRMSE_arr = DL_test_set_NRMSE_arr
        SSIM_arr = DL_test_set_SSIM_arr

    print(f'--------- {method_str} --------')


    ############ NRMSE figure ###############


    # NRMSE figure
    fig = plt.figure()
    for j in range(sampling_type_vec.shape[0]):

        NRMSE_vec = NRMSE_arr[:, :, j].squeeze()
        NRMSE_av_per_pad_N = NRMSE_vec.mean(axis=0)
        NRMSE_std_per_pad_N = NRMSE_vec.std(axis=0)

        NRMSE_diff = NRMSE_av_per_pad_N[-1] - NRMSE_av_per_pad_N[0]
        NRMSE_drop = NRMSE_diff / NRMSE_av_per_pad_N[0] * 100

        NRMSE_diff_vec = NRMSE_av_per_pad_N - NRMSE_av_per_pad_N[0]
        NRMSE_change_vec = NRMSE_diff_vec/ NRMSE_av_per_pad_N[0] * 100

        # create labels & plot lines showing the NRMSE decline
        if sampling_type_vec[j] == 1:
            label_str = 'weak VD'
            # plot horizontal and vertical straight lines:
            plt.plot(np.array([1, 2]), np.array([NRMSE_av_per_pad_N[0], NRMSE_av_per_pad_N[0]]), color='dimgray',
                     linestyle='--')
            plt.plot(np.array([2, 2]), np.array([NRMSE_av_per_pad_N[0], NRMSE_av_per_pad_N[-1]]), color='dimgray',
                     linestyle='solid')

        elif sampling_type_vec[j] == 2:
            label_str = 'strong VD'

        plt.plot(pad_ratio_vec, NRMSE_av_per_pad_N, linestyle='solid', label=label_str,color=colorslist[j],
                     marker=markers[j])

        plt.fill_between(pad_ratio_vec, (NRMSE_av_per_pad_N - NRMSE_std_per_pad_N/2), (NRMSE_av_per_pad_N + NRMSE_std_per_pad_N/2), color=colorslist[j],alpha=0.1)

        # print results
        print(f'{method_str} {label_str}; NRMSE init {NRMSE_av_per_pad_N[0]:.4f}; NRMSE end {NRMSE_av_per_pad_N[-1]:.4f}; NRMSE drop {NRMSE_drop:.1f}%')


    plt.ylabel('NRMSE', fontsize=20)
    ax = plt.gca()
    ax.set_xticks(x)
    ax.set_xticklabels(x_ticks_labels, fontsize=18)
    plt.xlim(0.95,2.05)
    plt.ylim(0.006, 0.02)
    ax.set_yticks((0.01, 0.015, 0.02, 0.02))
    plt.yticks(fontsize=20)
    plt.xticks(fontsize=20)
    if method_i==2:
        ax.legend(fontsize=20,loc='upper right')

    plt.tight_layout()
    plt.show()
    figname_NRMSE = figs_path + f'/{method_str}_NRMSE_stats'
    fig.savefig(fname=figname_NRMSE, dpi=1200)


    ############################################################################
    # SSIM figure
    fig = plt.figure()
    for j in range(sampling_type_vec.shape[0]):

        SSIM_vec = SSIM_arr[:, :, j].squeeze()
        SSIM_av_per_pad_N = SSIM_vec.mean(axis=0)
        SSIM_std_per_pad_N = SSIM_vec.std(axis=0)

        SSIM_diff = SSIM_av_per_pad_N[-1] - SSIM_av_per_pad_N[0]
        SSIM_drop = SSIM_diff / SSIM_av_per_pad_N[0] * 100

        SSIM_diff_vec = SSIM_av_per_pad_N - SSIM_av_per_pad_N[0]
        SSIM_change_vec = SSIM_diff_vec/ SSIM_av_per_pad_N[0] * 100

        # create labels
        if sampling_type_vec[j] == 1:
            label_str = 'weak VD'
        elif sampling_type_vec[j] == 2:
            label_str = 'strong VD'

        plt.plot(pad_ratio_vec, SSIM_av_per_pad_N, linestyle='solid', label=label_str,color=colorslist[j],
                     marker=markers[j])

        plt.fill_between(pad_ratio_vec, (SSIM_av_per_pad_N - SSIM_std_per_pad_N/2), (SSIM_av_per_pad_N + SSIM_std_per_pad_N/2), color=colorslist[j],alpha=0.1)

        # print results
        #print(f'SSIM av {label_str}: {SSIM_av_per_pad_N}')
        print(f'{method_str} {label_str}; SSIM init {SSIM_av_per_pad_N[0]:.2f} ; SSIM end  {SSIM_av_per_pad_N[-1]:.2f} ; SSIM rise {SSIM_drop:.1f}%')


    plt.xlabel('Zero padding', fontsize=18)
    plt.ylabel('SSIM', fontsize=20)
    ax = plt.gca()
    ax.set_xticks(x)
    ax.set_xticklabels(x_ticks_labels, fontsize=18)
    plt.xlim(0.95, 2.05)
    plt.locator_params(axis='y', nbins=4)
    plt.yticks(fontsize=20)
    plt.xticks(fontsize=20)
    plt.ylim(0.847, 0.97)
    ax.set_yticks((0.85, 0.9, 0.95))
    if method_i == 2:
        ax.legend(fontsize=20, loc='lower right')

    plt.tight_layout()
    plt.show()

    figname_SSIM = figs_path + f'/{method_str}_SSIM_stats'
    fig.savefig(fname=figname_SSIM, dpi=1200)

