""" Information Distribution Analysis Visualization (PMPS vs Sim-CoT) Generates publication-quality figures for information distribution uniformity: Fig 1: Information load per latent token (bar chart comparison) Fig 2: Matching weight heatmap (PMPS, case studies) Fig 3: Prototype utilization histogram Fig 4: Latent-CoT coverage heatmap Usage: python analysis/plot_distribution_analysis.py \ --pmps_dir analysis/results/pmps \ --simcot_dir analysis/results/simcot \ --output_dir analysis/figures/distribution """ import os import json import argparse import numpy as np import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt from matplotlib.colors import LinearSegmentedColormap plt.rcParams.update({ 'font.family': 'serif', 'font.serif': ['Times New Roman', 'DejaVu Serif'], 'font.size': 11, 'axes.titlesize': 13, 'axes.labelsize': 12, 'xtick.labelsize': 10, 'ytick.labelsize': 10, 'legend.fontsize': 10, 'figure.dpi': 300, 'savefig.bbox': 'tight', 'savefig.pad_inches': 0.05, }) COLORS = { 'codi': '#E74C3C', 'pmps': '#2E86C1', 'simcot': '#27AE60', } METHOD_LABELS = { 'codi': 'CODI', 'pmps': 'PMPS (Ours)', 'simcot': 'Sim-CoT', } # ───────────────────────────────────────────────────────────────────────────── # Fig 1: Information Load per Latent Token # ───────────────────────────────────────────────────────────────────────────── def plot_info_load_comparison(pmps_dist_metrics, simcot_info_loads, output_path, num_latent=6): """ Compare information load per latent token between PMPS and Sim-CoT. PMPS: uses matching weight sums from A_c2l Sim-CoT: uses uniform segmentation (N_c / N_l per latent, with remainder going to the last latent) """ fig, axes = plt.subplots(1, 2, figsize=(11, 4.5)) # --- Left: Per-latent info load bar chart --- ax = axes[0] # PMPS info loads pmps_loads = np.array([m['info_load_per_latent'] for m in pmps_dist_metrics]) pmps_mean = pmps_loads.mean(axis=0) pmps_std = pmps_loads.std(axis=0) N_l = len(pmps_mean) # Sim-CoT info loads (if provided, else simulate rigid segmentation) if simcot_info_loads is not None: simcot_loads = np.array(simcot_info_loads) simcot_mean = simcot_loads.mean(axis=0) simcot_std = simcot_loads.std(axis=0) else: simcot_mean = None x = np.arange(N_l) w = 0.35 bars_pmps = ax.bar(x - w/2, pmps_mean, w, yerr=pmps_std, capsize=4, color=COLORS['pmps'], alpha=0.8, label=METHOD_LABELS['pmps']) if simcot_mean is not None: bars_simcot = ax.bar(x + w/2, simcot_mean, w, yerr=simcot_std, capsize=4, color=COLORS['simcot'], alpha=0.8, label=METHOD_LABELS['simcot']) ideal_load = pmps_mean.sum() / N_l ax.axhline(y=ideal_load, color='gray', linestyle='--', alpha=0.6, label=f'Ideal (uniform = {ideal_load:.1f})') ax.set_xlabel('Latent Token Index') ax.set_ylabel('Information Load\n($\\sum_j A_{c2l}(j, i)$)') ax.set_title('(a) Information Load per Latent', fontsize=13) ax.set_xticks(x) ax.set_xticklabels([f'$l_{{{i+1}}}$' for i in range(N_l)]) ax.legend(fontsize=9) # --- Right: CV (Coefficient of Variation) comparison --- ax2 = axes[1] pmps_cvs = [m['info_load_cv'] for m in pmps_dist_metrics] categories = [METHOD_LABELS['pmps']] means = [np.mean(pmps_cvs)] stds = [np.std(pmps_cvs)] bar_colors = [COLORS['pmps']] scatter_data = [pmps_cvs] if simcot_info_loads is not None: simcot_cvs = [np.std(loads) / (np.mean(loads) + 1e-8) for loads in simcot_info_loads] categories.append(METHOD_LABELS['simcot']) means.append(np.mean(simcot_cvs)) stds.append(np.std(simcot_cvs)) bar_colors.append(COLORS['simcot']) scatter_data.append(simcot_cvs) bars = ax2.bar(range(len(categories)), means, yerr=stds, capsize=5, color=bar_colors, alpha=0.7, width=0.5) for i, (d, color) in enumerate(zip(scatter_data, bar_colors)): jitter = np.random.normal(0, 0.03, len(d)) ax2.scatter(i + jitter, d, color=color, alpha=0.15, s=10, zorder=3) ax2.text(i, means[i] + stds[i] + 0.02, f'{means[i]:.3f}', ha='center', fontsize=11, fontweight='bold', color=color) ax2.set_xticks(range(len(categories))) ax2.set_xticklabels(categories, fontsize=12) ax2.set_ylabel('Coefficient of Variation\n(Lower = More Uniform)', fontsize=11) ax2.set_title('(b) Load Uniformity (CV)', fontsize=13) plt.tight_layout() plt.savefig(output_path, dpi=300) plt.savefig(output_path.replace('.png', '.pdf'), dpi=300) plt.close() print(f' Saved: {output_path}') # ───────────────────────────────────────────────────────────────────────────── # Fig 2: Matching Weight Heatmap (Case Studies) # ───────────────────────────────────────────────────────────────────────────── def plot_matching_weight_heatmaps(pmps_internals, output_path, num_cases=3): """ Heatmap of A_c2l for individual samples. Shows which CoT tokens are assigned to which latent tokens. """ n_cases = min(num_cases, len(pmps_internals)) fig, axes = plt.subplots(1, n_cases, figsize=(5 * n_cases, 5)) if n_cases == 1: axes = [axes] cmap = LinearSegmentedColormap.from_list( 'weight', ['#FFFFFF', '#E8F8F5', '#1ABC9C', '#117864'], N=256 ) for idx in range(n_cases): sample = pmps_internals[idx] A_c2l = np.array(sample['A_c2l']) # (N_c, N_l) N_c, N_l = A_c2l.shape ax = axes[idx] im = ax.imshow(A_c2l, cmap=cmap, aspect='auto', vmin=0) ax.set_xlabel('Latent Token Index') ax.set_ylabel('CoT Token Position') ax.set_title(f'Sample #{sample["sample_idx"]}', fontsize=12, fontweight='bold') ax.set_xticks(range(N_l)) ax.set_xticklabels([f'$l_{{{i+1}}}$' for i in range(N_l)]) if N_c > 30: tick_step = max(1, N_c // 10) ax.set_yticks(range(0, N_c, tick_step)) fig.colorbar(im, ax=axes, label='Matching Weight $A_{c2l}$', shrink=0.8, pad=0.02) fig.suptitle('CoT → Latent Matching Weights (PMPS)', fontsize=14, y=1.02) plt.tight_layout() plt.savefig(output_path, dpi=300) plt.savefig(output_path.replace('.png', '.pdf'), dpi=300) plt.close() print(f' Saved: {output_path}') # ───────────────────────────────────────────────────────────────────────────── # Fig 3: Prototype Utilization # ───────────────────────────────────────────────────────────────────────────── def plot_prototype_utilization(pmps_dist_metrics, output_path): """ Prototype utilization histogram + entropy comparison. """ fig, axes = plt.subplots(1, 2, figsize=(11, 4.5)) # Left: Average prototype usage (from latent assignments) ax = axes[0] all_proto_usage = np.array([m['proto_usage_latent'] for m in pmps_dist_metrics]) avg_usage = all_proto_usage.mean(axis=0) std_usage = all_proto_usage.std(axis=0) K = len(avg_usage) uniform_level = 1.0 / K ax.bar(range(K), avg_usage, yerr=std_usage, capsize=2, color=COLORS['pmps'], alpha=0.7, edgecolor='white', linewidth=0.5) ax.axhline(y=uniform_level, color='red', linestyle='--', alpha=0.6, label=f'Uniform = {uniform_level:.4f}') ax.set_xlabel('Prototype Index') ax.set_ylabel('Average Assignment Weight') ax.set_title('(a) Prototype Utilization (Latent Side)', fontsize=13) ax.legend(fontsize=9) if K > 16: ax.set_xticks(range(0, K, max(1, K // 8))) # Right: Entropy comparison ax2 = axes[1] proto_ents = [m['proto_entropy_latent'] for m in pmps_dist_metrics] max_ent = pmps_dist_metrics[0]['max_entropy'] ax2.hist(proto_ents, bins=30, color=COLORS['pmps'], alpha=0.7, edgecolor='white', linewidth=0.5, density=True) ax2.axvline(x=max_ent, color='red', linestyle='--', linewidth=2, alpha=0.7, label=f'Max Entropy = $\\log K$ = {max_ent:.2f}') ax2.axvline(x=np.mean(proto_ents), color=COLORS['pmps'], linestyle='-', linewidth=2, label=f'Mean = {np.mean(proto_ents):.2f}') ax2.set_xlabel('Prototype Assignment Entropy') ax2.set_ylabel('Density') ax2.set_title('(b) Prototype Entropy Distribution', fontsize=13) ax2.legend(fontsize=9) plt.tight_layout() plt.savefig(output_path, dpi=300) plt.savefig(output_path.replace('.png', '.pdf'), dpi=300) plt.close() print(f' Saved: {output_path}') # ───────────────────────────────────────────────────────────────────────────── # Fig 4: Latent-CoT Coverage Comparison # ───────────────────────────────────────────────────────────────────────────── def plot_coverage_comparison(pmps_dist_metrics, simcot_info_loads, output_path): """ Compare how CoT information is distributed across latents. For PMPS: use matching weight entropy per CoT token (higher = distributed more evenly) For Sim-CoT: compute rigid-segment coverage """ fig, ax = plt.subplots(figsize=(7, 4.5)) # PMPS matching entropies pmps_match_ents = [] for m in pmps_dist_metrics: pmps_match_ents.extend(m['match_entropies']) data_list = [pmps_match_ents] labels = [METHOD_LABELS['pmps']] colors = [COLORS['pmps']] # For Sim-CoT: approximate with delta function (each CoT token maps to exactly 1 latent) # entropy = 0 for rigid one-to-one mapping if simcot_info_loads is not None: simcot_ents = [0.0] * len(pmps_match_ents) data_list.append(simcot_ents) labels.append(METHOD_LABELS['simcot']) colors.append(COLORS['simcot']) parts = ax.violinplot(data_list, positions=range(len(labels)), showmedians=True, showextrema=False) for pc, color in zip(parts['bodies'], colors): pc.set_facecolor(color) pc.set_alpha(0.4) parts['cmedians'].set_color('black') parts['cmedians'].set_linewidth(2) for i, (d, color) in enumerate(zip(data_list, colors)): mean_val = np.mean(d) ax.plot(i, mean_val, 'D', color='white', markeredgecolor=color, markersize=8, markeredgewidth=2, zorder=5) ax.annotate(f'{mean_val:.3f}', (i, mean_val), textcoords='offset points', xytext=(18, 5), fontsize=10, color=color, fontweight='bold') ax.set_xticks(range(len(labels))) ax.set_xticklabels(labels, fontsize=12) ax.set_ylabel('Per-CoT-Token Matching Entropy\n(Higher = More Distributed)', fontsize=11) ax.set_title('Matching Weight Distribution\n(How Diversely Each CoT Token Maps to Latents)', fontsize=13) plt.tight_layout() plt.savefig(output_path, dpi=300) plt.savefig(output_path.replace('.png', '.pdf'), dpi=300) plt.close() print(f' Saved: {output_path}') # ───────────────────────────────────────────────────────────────────────────── # Main # ───────────────────────────────────────────────────────────────────────────── def main(): parser = argparse.ArgumentParser() parser.add_argument('--pmps_dir', type=str, required=True) parser.add_argument('--simcot_dir', type=str, default=None) parser.add_argument('--output_dir', type=str, default='analysis/figures/distribution') args = parser.parse_args() os.makedirs(args.output_dir, exist_ok=True) # Load PMPS distribution metrics print('Loading PMPS distribution metrics...') pmps_dist_path = os.path.join(args.pmps_dir, 'distribution_metrics.json') pmps_pmps_path = os.path.join(args.pmps_dir, 'pmps_internals.json') with open(pmps_dist_path, 'r') as f: pmps_dist_metrics = json.load(f) print(f' PMPS: {len(pmps_dist_metrics)} samples with distribution metrics') pmps_internals = None if os.path.exists(pmps_pmps_path): with open(pmps_pmps_path, 'r') as f: pmps_internals = json.load(f) print(f' PMPS: {len(pmps_internals)} samples with PMPS internals') # Load Sim-CoT info loads (if available) simcot_info_loads = None if args.simcot_dir: simcot_dist_path = os.path.join(args.simcot_dir, 'distribution_metrics.json') if os.path.exists(simcot_dist_path): with open(simcot_dist_path, 'r') as f: simcot_dist = json.load(f) simcot_info_loads = [m['info_load_per_latent'] for m in simcot_dist] print(f' Sim-CoT: {len(simcot_info_loads)} samples') else: print(' Sim-CoT: no distribution_metrics.json found, will simulate rigid segmentation') print('\nGenerating figures...') plot_info_load_comparison( pmps_dist_metrics, simcot_info_loads, os.path.join(args.output_dir, 'fig_info_load.png'), ) if pmps_internals: plot_matching_weight_heatmaps( pmps_internals, os.path.join(args.output_dir, 'fig_matching_weights.png'), num_cases=3, ) plot_prototype_utilization( pmps_dist_metrics, os.path.join(args.output_dir, 'fig_prototype_utilization.png'), ) plot_coverage_comparison( pmps_dist_metrics, simcot_info_loads, os.path.join(args.output_dir, 'fig_coverage_comparison.png'), ) print(f'\nAll figures saved to: {args.output_dir}') if __name__ == '__main__': main()