langevin_sampling_example.py

import numpy as np
import matplotlib.pyplot as plt

# Set seed for reproducibility
np.random.seed(123)

# Parameters
T = 10000       # number of time steps
s = 0.01        # step size
dim = 2         # 2D space

# Initialize Langevin dynamics
z = np.zeros((T, dim))
z[0] = np.array([4.0, -3.0])  # Initial position

# Langevin update: z_{t+1} = z_t - s * grad U(z_t) + sqrt(2s) * N(0, I)
# where grad U(z) = z for standard Gaussian
for t in range(T - 1):
    grad_U = z[t]
    noise = np.sqrt(2 * s) * np.random.randn(dim)
    z[t + 1] = z[t] - s * grad_U + noise

# Plot trajectory and samples
fig, ax = plt.subplots(figsize=(7, 7))

# Plot trajectory path
ax.plot(z[:, 0], z[:, 1], color='gray', alpha=0.3, linewidth=1, label='Trajectory')

# Use second half of samples for visualization (post burn-in)
samples = z[int(T/2):]
ax.scatter(samples[:, 0], samples[:, 1], s=2, alpha=0.4, label='Samples')

# Mark the starting and final points
ax.plot(z[0, 0], z[0, 1], 'ro', label='Start')
ax.plot(z[-1, 0], z[-1, 1], 'go', label='End')

# Styling
ax.set_title("2D Langevin Dynamics Trajectory and Samples")
ax.set_xlabel("z₁")
ax.set_ylabel("z₂")
ax.set_aspect('equal')
ax.grid(True)
ax.legend()
plt.tight_layout()
plt.show()