plot.py

import numpy as np
import matplotlib.pyplot as plt
import os
import re

N_range = [100,200,300,400,500]

T = 20000
ratio = 0.01
file_dir = './renormalized_sgd_scaling'
plot_dir = './plots/renormalized_sgd_scaling'
os.makedirs(plot_dir, exist_ok=True)
folders = [d for d in os.listdir(file_dir) if f'ratio{ratio}' in d and f'ratio0.015' not in d]

print(folders)

for N in N_range:

    cmap = plt.cm.magma
    color_len = len([d for d in folders if f'ratio{ratio}' in d and f'N{N}' in d])
    colors = [cmap(i) for i in np.linspace(0, 1, color_len)]

    plt.figure(figsize=(15, 10))

    i = 0

    for folder in folders:
        if f'N{N}' in folder:

            match = re.search(r'N(\d+)_B(\d+)_lr([\d.]+)_ratio([\d.]+)', folder)
            N = int(match.group(1))
            B = int(match.group(2))
            lr = float(match.group(3))
            
            f = f'./renormalized_sgd_scaling/{folder}/sgd_empirical.npy'
            if os.path.exists(f):
                data = np.load(f'./renormalized_sgd_scaling/{folder}/sgd_empirical.npy')
                data = data[0][:T]
                print(data.shape)

                # Plotting: Iterations on x-axis
                # mean_sgd = np.mean(data, axis=0)
                # std_sgd = np.std(data, axis=0)
                plt.plot(range(T), data, label=f'N={N} B={B} lr={lr}', color=colors[i])
                i += 1
                # plt.fill_between(range(T), mean_sgd - std_sgd, mean_sgd + std_sgd, color='r', alpha=0.2)

                # TODO: need to get renormalized sgd theoretical risk
                # plt.plot(range(T), sgd_theoretical_risks[B], label='SGD Theoretical', color='r', linestyle='--')
                
                plt.title(f'SGD scaling')
                plt.xlabel('Iteration $t$')
                plt.ylabel('$R(W_t)$')
                plt.yscale('log')
                plt.grid(True)
                plt.legend()

    plt.tight_layout()
    plt.savefig(f'{plot_dir}/iteration_sweep_N_{N}_ratio{ratio}.png')
    plt.show()