Neural Network Weight Distribution KDE

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import gaussian_kde

np.random.seed(42)
BG_COLOR = '#0a0a0f'
TEXT_COLOR = 'white'

# Weight distributions at different epochs
epoch_1 = np.random.normal(0, 0.5, 1000)
epoch_50 = np.random.normal(0, 0.3, 1000)
epoch_100 = np.random.normal(0, 0.2, 1000)

fig, ax = plt.subplots(figsize=(10, 6), facecolor=BG_COLOR)
ax.set_facecolor(BG_COLOR)

x_range = np.linspace(-2, 2, 500)

for data, color, label in [(epoch_1, '#F5276C', 'Epoch 1'), 
                            (epoch_50, '#27D3F5', 'Epoch 50'),
                            (epoch_100, '#6CF527', 'Epoch 100')]:
    kde = gaussian_kde(data)
    density = kde(x_range)
    ax.plot(x_range, density, color=color, linewidth=2.5, label=label)
    ax.fill_between(x_range, density, alpha=0.3, color=color)

ax.set_xlabel('Weight Value', fontsize=12, color=TEXT_COLOR, fontweight='500')
ax.set_ylabel('Density', fontsize=12, color=TEXT_COLOR, fontweight='500')
ax.set_title('Neural Network Weight Distribution by Epoch', fontsize=14, color=TEXT_COLOR, fontweight='bold', pad=15)

ax.tick_params(colors='#888888', labelsize=10)
for spine in ax.spines.values():
    spine.set_color('#333333')
ax.legend(facecolor=BG_COLOR, edgecolor='#333333', labelcolor=TEXT_COLOR, fontsize=10)

plt.tight_layout()
plt.show()