2D Bond Percolation Simulation with Union-Find

import numpy as np
import matplotlib.pyplot as plt
import time
import jaxfss

class DisjointSet:
    """Union-Find data structure with path compression and union by size."""
    def __init__(self, n: int):
        self.parent = list(range(n))
        self.size = [1] * n

    def find(self, i: int) -> int:
        path = []
        while self.parent[i] != i:
            path.append(i)
            i = self.parent[i]
        for node in path:
            self.parent[node] = i
        return i

    def union(self, i: int, j: int) -> bool:
        root_i = self.find(i)
        root_j = self.find(j)
        if root_i != root_j:
            # Union by size: attach smaller cluster to larger cluster
            if self.size[root_i] < self.size[root_j]:
                root_i, root_j = root_j, root_i
            self.parent[root_j] = root_i
            self.size[root_i] += self.size[root_j]
            return True
        return False

def get_square_lattice_bonds(L: int) -> np.ndarray:
    """Generate all nearest-neighbor pairs (edges) on an L x L open-boundary grid."""
    bonds = []
    for x in range(L):
        for y in range(L):
            idx = x * L + y
            if x + 1 < L:
                bonds.append((idx, (x + 1) * L + y))
            if y + 1 < L:
                bonds.append((idx, x * L + (y + 1)))
    return np.array(bonds, dtype=np.int32)

def get_cluster_grid(L: int, p: float, seed: int = 0):
    rng = np.random.default_rng(seed)
    bonds = get_square_lattice_bonds(L)
    active_bonds = bonds[rng.random(len(bonds)) < p]
    
    N = L * L
    uf = DisjointSet(N)
    for u, v in active_bonds:
        uf.union(u, v)
        
    root_ids = [uf.find(i) for i in range(N)]
    # Map each root to cluster size
    sizes = np.array([uf.size[r] for r in root_ids])
    max_root = root_ids[np.argmax(sizes)]
    
    # Assign unique visual color IDs (reserve -1 for largest cluster)
    unique_roots = list(set(root_ids))
    rng.shuffle(unique_roots)
    root_to_color = {r: (i % 20 + 1) for i, r in enumerate(unique_roots)}
    root_to_color[max_root] = 0  # largest cluster highlighted
    
    grid = np.array([root_to_color[r] for r in root_ids]).reshape(L, L)
    return grid, uf.size[max_root]

L_vis = 32
fig, axes = plt.subplots(1, 3, figsize=(13, 4.5))
probabilities = [0.40, 0.50, 0.60]
titles = ["Subcritical ($p = 0.40$)", "Critical ($p = 0.50$)", "Supercritical ($p = 0.60$)"]

for ax, p_val, title in zip(axes, probabilities, titles):
    grid, s_max = get_cluster_grid(L_vis, p_val, seed=42)
    # Plot clusters with largest cluster colored distinctly
    ax.imshow(grid, cmap="tab20", origin="lower", interpolation="nearest")
    ax.set_title(rf"{title}\n$s_{{\max}} = {s_max} / {L_vis**2}$", fontsize=11)
    ax.axis("off")

plt.tight_layout()
plt.show()

def sample_susceptibility(L: int, p: float, n_samples: int = 100, rng: np.random.Generator = None) -> tuple[float, float]:
    """Compute the truncated mean cluster size chi^f(p, L) averaged over n_samples."""
    if rng is None:
        rng = np.random.default_rng()
    N = L * L
    bonds = get_square_lattice_bonds(L)
    n_bonds = len(bonds)
    
    chis = []
    for _ in range(n_samples):
        active = bonds[rng.random(n_bonds) < p]
        uf = DisjointSet(N)
        for u, v in active:
            uf.union(u, v)
            
        roots = set(uf.find(i) for i in range(N))
        sizes = sorted([uf.size[r] for r in roots], reverse=True)
        # Exclude largest cluster s_max
        finite_sizes = sizes[1:] if len(sizes) > 1 else [0]
        chi_val = sum(s * s for s in finite_sizes) / N
        chis.append(chi_val)
        
    chis = np.array(chis)
    mean = np.mean(chis)
    std_err = np.std(chis, ddof=1) / np.sqrt(n_samples)
    return float(mean), float(std_err)

system_sizes = [16, 32, 64]
p_values = np.linspace(0.35, 0.65, 21)
n_samples = 120
rng = np.random.default_rng(2026)

sim_results = {}
print(f"Running percolation simulation for L in {system_sizes}...")
t_start = time.time()

for L in system_sizes:
    means, errs = [], []
    for p in p_values:
        m, e = sample_susceptibility(L, p, n_samples=n_samples, rng=rng)
        means.append(m)
        errs.append(e)
    sim_results[L] = (np.array(means), np.array(errs))
    print(f"  L = {L:2d} completed in {time.time() - t_start:.2f}s (peak chi = {max(means):.2f})")

print(f"Total simulation time: {time.time() - t_start:.2f}s")

plt.figure(figsize=(7.5, 4.8))
markers = ["o", "s", "^"]

for (L, (means, errs)), m in zip(sim_results.items(), markers):
    plt.errorbar(p_values, means, yerr=errs, fmt=f"{m}-", label=f"$L = {L}$", capsize=3, lw=1.5)

plt.axvline(0.5, color="black", linestyle=":", label="Exact $p_{\mathrm{c}} = 0.5$")
plt.xlabel(r"Occupation Probability $p$", fontsize=12)
plt.ylabel(r"Susceptibility $\chi^{\mathrm{f}}(p, L)$", fontsize=12)
plt.title("2D Bond Percolation: Susceptibility vs. $p$", fontsize=14)
plt.legend(fontsize=11)
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

output_rows = []
for L, (means, errs) in sim_results.items():
    for p, val, err in zip(p_values, means, errs):
        output_rows.append([L, p, val, err])

output_data = np.array(output_rows)
output_path = "percolation_custom.txt"

# Save to file with exact format
np.savetxt(
    output_path,
    output_data,
    fmt=["%d", "%.8f", "%.8f", "%.8f"],
    header="L p chi_f chi_f_err"
)
print(f"Saved {len(output_data)} rows to {output_path}.")

# Verify loading with jaxfss
loaded_dataset = jaxfss.CriticalData.from_file(output_path)
print(f"Successfully loaded into jaxfss: {loaded_dataset}")
print(f"System sizes in dataset: {np.unique(loaded_dataset.Ls)}")