# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# +
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import matplotlib as mpl
import sys
import os
import cmocean
import cmocean.cm as cmo

from matplotlib.lines import Line2D
from matplotlib import pyplot
# -

os.path.exists

# +
basic_wds = [1e-1,1.0,10.,100.]
equiv_wds = [1e-2,1e-3,1e-4,1e-5]
results = []
results = pd.DataFrame()
for basic_idx, basic_wd in enumerate(basic_wds):
    for equiv_idx, equiv_wd in enumerate(equiv_wds):
        fname = "inertia_log_MixedEMLP_basic" + str(basic_wd) + "_equiv" + str(equiv_wd) + ".pkl"
        if os.path.exists("./saved-outputs/" + fname):
            df = pd.read_pickle("./saved-outputs/" + fname)
            if basic_wd == 0.:
                df['basic_wd'] = 1e-7
            else:
                df['basic_wd'] = basic_wd
            if equiv_wd == 0.:
                df['equiv_wd'] = 1e-7
            else:
                df['equiv_wd'] = equiv_wd
            
            results = pd.concat((results, df))

results.columns = ['trial', 'tr_loss', 'te_loss', 'basic_wd', 'equiv_wd']

# +
results['log_basic'] = np.log10(results['basic_wd'])
results['log_equiv'] = np.log10(results['equiv_wd'])
# results = results.replace(to_replace=-7, value=r'-inf')

results["equiv_hue"] = ["$%s$" % x for x in results["log_equiv"]]
results["basic_hue"] = ["$%s$" % x for x in results["log_basic"]]
# -

results[results['te_loss'] > 10.] = None

# +
# results = results[(results['log_basic'] >= 0.) & (results['log_equiv'] <= -2.)]

# +

ncolor = len(results['log_basic'].unique())
ncolor=4

lightcmap = cmocean.tools.crop_by_percent(cmo.matter, ncolor, which='min', N=None)
cpal = sns.color_palette("cmo.matter", n_colors=ncolor)#, desat=0.2)

fs = 16
sns.set_style("whitegrid")

COLOR = '#444444'
mpl.rcParams['text.color'] = COLOR
mpl.rcParams['axes.labelcolor'] = COLOR
mpl.rcParams['xtick.color'] = COLOR
mpl.rcParams['ytick.color'] = COLOR

plt.figure(dpi=150)
# sns.lineplot(x='log_equiv', y='te_loss', hue='basic_hue', data=results[results['log_basic'] < -1.],
#             palette='RdYlBu')

fig, ax = plt.subplots(dpi=150, figsize=(6, 3))
g = sns.lineplot(x='equiv_wd', y='te_loss', data=results,
        hue='basic_hue', markers=['o' for _ in range(ncolor)],
             style='basic_hue', dashes=False, markersize=10,
                ax=ax, palette=cpal, legend=False)
print(g)
ax.legend(loc='lower left', title='Log Basic Precision', bbox_to_anchor=(1., 0.3),
         labels=[r"$10^{-1}$", r"$10^{0}$", r"$10^{1}$", r"$10^{2}$"])
ax.set_xlabel("Equivariant Precision", fontsize=fs)
ax.set_ylabel("Test MSE", fontsize=fs)
ax.set_title("Modified Inertia Test MSE", fontsize=16)
# plt.ylim(0, 10)
# ax.set_xlim(-7, -1)
# ax.set_xticklabels([r"-$\infty$", -6, -5, -4, -3, -2, -1])
ax.tick_params('both', labelsize=fs-2)
ax.axhline(4.791855, linewidth=2.)
ax.axhline(1.053696, linestyle='--', linewidth=2.)
sns.despine()
for axis in ['top','bottom','left','right']:
    ax.spines[axis].set_linewidth(2.)

ax.set_yticks([1., 5.])
ax.set_yticklabels([r"$1x10^{0}$", r"$5x10^{0}$"])
    
handles = []
labels = []
plt.semilogx()
plt.semilogy()
handles.append(Line2D([0], [0], color="#1f77b4", lw=2))
handles.append(Line2D([0], [0], color="#1f77b4", lw=2, ls="--"))
labels.append("MLP")
labels.append("EMLP")

legend1 = fig.legend(handles, labels, bbox_to_anchor=(1.12, 0.35))
pyplot.gca().add_artist(legend1)

ax.set_yticks([1., 5.])
ax.set_yticklabels([r"$10^{0}$", r"$5x10^{0}$"])


plt.savefig("./figs/inertia_equiv_wd.pdf", bbox_inches='tight')
plt.show()
# -

results.log_basic.unique()

# +
equiv_wd = 0.001
basic_wd = 100.0

fname = "./saved-outputs/inertia_log_mlp_basic" + str(basic_wd)
fname = fname + "_equiv" + str(equiv_wd) + ".pkl"

mlp = pd.read_pickle(fname)
# -

mlp

mlp.mean()

# +
equiv_wd = 0.001
basic_wd = 100.0

fname = "./saved-outputs/inertia_log_emlp_basic" + str(basic_wd)
fname = fname + "_equiv" + str(equiv_wd) + ".pkl"

emlp = pd.read_pickle(fname)
# -

emlp.mean()



StarcoderdataJupyterScriptsDedupFiltered