lars
1a0324d473
- Implement OLS and Ridge as well as python structure - Start with gradient descent - Create plots for the above mentioned - Create a basic structure for the python source - Install pre commit hooks to check and format the code
134 lines
4.2 KiB
Python
134 lines
4.2 KiB
Python
import matplotlib.pyplot as plt
|
|
from matplotlib.colors import SymLogNorm
|
|
import numpy as np
|
|
|
|
FIG_WIDTH = 6
|
|
|
|
|
|
def get_rc_params():
|
|
colors = ["FF220C", "70D6FF", "8AAA79", "666370", "1C1F33"]
|
|
rcParams = plt.rcParams
|
|
# Use LaTeX for rendering
|
|
# Setup fonts
|
|
rcParams["text.usetex"] = True
|
|
rcParams["font.family"] = "serif"
|
|
rcParams["font.size"] = 12
|
|
rcParams["axes.labelsize"] = 12
|
|
rcParams["axes.titlesize"] = 12
|
|
rcParams["legend.fontsize"] = 10
|
|
rcParams["xtick.labelsize"] = 10
|
|
rcParams["ytick.labelsize"] = 10
|
|
# Figure size and resolution
|
|
rcParams["figure.figsize"] = (FIG_WIDTH, 4)
|
|
rcParams["figure.dpi"] = 300
|
|
# Use colors from the palette
|
|
rcParams["axes.prop_cycle"] = plt.cycler(color=[f"#{color}" for color in colors])
|
|
# Grid
|
|
rcParams["axes.grid"] = True
|
|
rcParams["grid.alpha"] = 0.5
|
|
rcParams["grid.linestyle"] = "--"
|
|
# Point ticks to the inside of the axes
|
|
rcParams["xtick.direction"] = "in"
|
|
rcParams["ytick.direction"] = "in"
|
|
rcParams["xtick.top"] = True
|
|
rcParams["ytick.right"] = True
|
|
return rcParams
|
|
|
|
|
|
def set_rc_params():
|
|
plt.rcParams.update(get_rc_params())
|
|
|
|
|
|
set_rc_params()
|
|
|
|
|
|
def get_figsize(rel_height: float = 2 / 3) -> tuple[float, float]:
|
|
"""Returns a figure size tuple based on the global FIG_WIDTH and a relative height.
|
|
Args:
|
|
rel_height: The relative height of the figure compared to FIG_WIDTH.
|
|
Returns:
|
|
A tuple (width, height) for the figure size.
|
|
"""
|
|
return (FIG_WIDTH, FIG_WIDTH * rel_height)
|
|
|
|
|
|
def mse_r2_plot(
|
|
polynomial_degrees: np.ndarray,
|
|
train_mse_list: list[float],
|
|
mse_list: list[float],
|
|
train_r2_list: list[float],
|
|
r2_list: list[float],
|
|
labels: dict[str, str] | None = None,
|
|
) -> tuple[plt.Figure, tuple[plt.Axes, plt.Axes]]:
|
|
std_labels = {
|
|
"set1": "Train Set",
|
|
"set2": "Test Set",
|
|
"xlabel": "Polynomial Degree",
|
|
"ylabel1": "Mean Squared Error",
|
|
"ylabel2": "$R^2$ Score",
|
|
}
|
|
|
|
if labels is None:
|
|
labels = std_labels
|
|
else:
|
|
for key in std_labels:
|
|
if key not in labels:
|
|
labels[key] = std_labels[key]
|
|
|
|
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=get_figsize(0.5))
|
|
ax1.plot(polynomial_degrees, train_mse_list, marker="o", label=labels["set1"])
|
|
ax1.plot(polynomial_degrees, mse_list, marker="o", label=labels["set2"])
|
|
ax1.legend()
|
|
ax1.set_xlabel(labels["xlabel"])
|
|
ax1.set_ylabel(labels["ylabel1"])
|
|
|
|
ax2.plot(polynomial_degrees, train_r2_list, marker="o", label=labels["set1"])
|
|
ax2.plot(polynomial_degrees, r2_list, marker="o", label=labels["set2"])
|
|
ax2.legend()
|
|
ax2.set_xlabel(labels["xlabel"])
|
|
ax2.set_ylabel(labels["ylabel2"])
|
|
|
|
fig.tight_layout()
|
|
return fig, (ax1, ax2)
|
|
|
|
|
|
def parameter_plot(
|
|
polynomial_degrees: np.ndarray,
|
|
beta_OLS_list: list[np.ndarray],
|
|
labels: dict[str, str] | None = None,
|
|
) -> tuple[plt.Figure, plt.Axes]:
|
|
std_labels = {
|
|
"xlabel": "Polynomial Degree of Fit",
|
|
"ylabel": "Coefficient Index $i$",
|
|
"cbar": "Coefficient Value",
|
|
}
|
|
|
|
if labels is None:
|
|
labels = std_labels
|
|
else:
|
|
for key in std_labels:
|
|
if key not in labels:
|
|
labels[key] = std_labels[key]
|
|
|
|
fig, ax = plt.subplots(figsize=get_figsize(0.5))
|
|
beta_OLS = np.zeros((beta_OLS_list[-1].shape[0], len(polynomial_degrees)))
|
|
beta_OLS[:] = np.nan
|
|
for i in range(len(beta_OLS_list)):
|
|
for j in range(len(beta_OLS_list[i])):
|
|
beta_OLS[j, i] = beta_OLS_list[i][j]
|
|
cmap = plt.get_cmap("coolwarm")
|
|
norm = SymLogNorm(
|
|
linthresh=1e-3, vmin=np.nanmin(beta_OLS), vmax=np.nanmax(beta_OLS)
|
|
)
|
|
im = ax.imshow(beta_OLS, aspect="auto", cmap=cmap, origin="lower", norm=norm)
|
|
cbar = fig.colorbar(im, ax=ax)
|
|
cbar.set_label(labels["cbar"])
|
|
ax.set_xlabel(labels["xlabel"])
|
|
ax.set_ylabel(labels["ylabel"])
|
|
ax.set_xticks(np.arange(0, len(polynomial_degrees), step=3))
|
|
ax.set_xticklabels(polynomial_degrees[::3])
|
|
ax.set_yticks(np.arange(0, beta_OLS.shape[0], 5))
|
|
ax.set_yticklabels(np.arange(1, beta_OLS.shape[0] + 1, step=5))
|
|
fig.tight_layout()
|
|
return fig, ax
|