Add more comments and examples README

examples/README.md

@@ -0,0 +1,3 @@
+# Examples
+
+The scripts in this folder reproduce the plots appearing in the blog post [Transformers from Spin Models: Approximate Free Energy Minimization](https://mcbal.github.io/post/transformers-from-spin-models-approximate-free-energy-minimization/).

examples/model_beta_sweep.py

@@ -1,5 +1,5 @@
 # Sweep across inverse temperature `beta` for fixed inputs to a `VectorSpinModel` module
-# and plot `phi` and its derivatives.
+# and return plots of `phi` and its derivatives as an animated GIF.
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -13,10 +13,16 @@
 beta_min, beta_max, beta_num_steps = -0.6, 1.0, 100  # log10-space
 plot_values_cutoff = 500
 
+# Fixed inputs and initial auxiliary variable.
 x = (torch.randn(1, num_spins, dim) / np.sqrt(dim)).requires_grad_()
 t0 = 0.5*torch.ones(1)
 
 
+##############
+# PLOT STUFF #
+##############
+
+
 def filter_array(a, threshold=plot_values_cutoff):
     idx = np.where(np.abs(a) > threshold)
     a[idx] = np.nan
@@ -27,17 +33,20 @@ def simple_update(frame, fig, axs):
 
     print(f'{frame:.4f} / {10**beta_max}')
 
+    # Setup vector-spin model.
     model = VectorSpinModel(
         num_spins=num_spins,
         dim=dim,
         beta=frame,
     )
-    out = model(x, t0=t0, return_afe=True, return_magnetizations=True)
+    # Run forward pass.
+    out = model(x, t0=t0, return_afe=True)
 
     t_star = out.t_star[0][0].detach().numpy()
     afe_star = out.afe[0][0].detach().numpy()
     t_min, t_max, t_step = 0.0, 3.0, 0.001
     t_range = torch.arange(t_min, t_max, t_step)[:, None]
+
     phis = np.array(model._phi(t_range, x[:1, :, :].repeat(t_range.numel(), 1, 1)).detach().numpy())
     grad_phis = np.array(model._jac_phi(
         t_range, x[:1, :, :].repeat(t_range.numel(), 1, 1)).detach().numpy())

examples/model_fwd_bwd.py

@@ -29,9 +29,9 @@
 # Run backward on sum of free energies across batch dimension.
 out.afe.sum().backward()
 
-###############################################################
-# Plot internally-used function `phi(t)` and its derivatives. #
-###############################################################
+##############
+# PLOT STUFF #
+##############
 
 if x.size(0) == 1 and t0.numel() == 1:
 

examples/model_two_spins.py

@@ -27,6 +27,10 @@
     f'✨ t_star {out.t_star.detach()}\n✨ magnetizations: {out.magnetizations}\n✨ approximate free energy: {out.afe.detach()}'
 )
 
+##############
+# PLOT STUFF #
+##############
+
 
 def filter_array(a, threshold=50):
     idx = np.where(np.abs(a) > threshold)