1 files changed, 41 insertions, 0 deletions
diff --git a/viz/gen_fig1.py b/viz/gen_fig1.py
new file mode 100644
index 0000000..a9e23f7
--- /dev/null
+++ b/viz/gen_fig1.py
@@ -0,0 +1,41 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+# Set aesthetic styling
+plt.style.use('seaborn-v0_8-darkgrid' if 'seaborn-v0_8-darkgrid' in plt.style.available else 'default')
+fig, ax = plt.subplots(figsize=(9, 5), dpi=300)
+
+layers = [f'L{i}' for i in range(6)]
+proxy_r = [-0.661, -0.673, -0.653, -0.648, -0.562, -0.665]
+faithful_r = [-0.643, -0.668, -0.672, -0.685, -0.637, -0.606]
+spearman_r = [-0.844, -0.851, -0.825, -0.769, -0.716, -0.808]
+
+x = np.arange(len(layers))
+width = 0.25
+
+rects1 = ax.bar(x - width, faithful_r, width, label='Faithful Lens (Pearson r)', color='#6366f1', alpha=0.9)
+rects2 = ax.bar(x, proxy_r, width, label='Proxy Lens (Pearson r)', color='#a855f7', alpha=0.7)
+rects3 = ax.bar(x + width, spearman_r, width, label='Faithful Lens (Spearman ρ)', color='#06b6d4', alpha=0.8)
+
+ax.set_ylabel('Correlation with Token Frequency', fontsize=12, fontweight='bold', labelpad=10)
+ax.set_title('Token Frequency Anti-Correlation Across Layers (nanoGPT 10.65M)', fontsize=14, fontweight='bold', pad=15)
+ax.set_xticks(x)
+ax.set_xticklabels(layers, fontsize=11)
+ax.axhline(0, color='black', linewidth=0.8, linestyle='--')
+ax.set_ylim(-1.0, 0.1)
+ax.legend(frameon=True, facecolor='#ffffff', edgecolor='none', fontsize=10, loc='lower left')
+
+# Add values above/below bars
+for rect in rects1:
+    h = rect.get_height()
+    ax.annotate(f'{h:.2f}', xy=(rect.get_x() + rect.get_width() / 2, h),
+                xytext=(0, 6), textcoords="offset points", ha='center', va='bottom', fontsize=8, color='white', fontweight='bold')
+
+for rect in rects3:
+    h = rect.get_height()
+    ax.annotate(f'{h:.2f}', xy=(rect.get_x() + rect.get_width() / 2, h),
+                xytext=(0, 6), textcoords="offset points", ha='center', va='bottom', fontsize=8, color='white', fontweight='bold')
+
+plt.tight_layout()
+plt.savefig('docs/assets/fig1_layer_correlation.png', dpi=300)
+print("Saved Fig 1")