ldpc_optical/model/plot_de_results.py

#!/usr/bin/env python3
"""Generate presentation plots from density evolution results."""

import json
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from pathlib import Path

RESULTS_PATH = Path(__file__).parent.parent / 'data' / 'de_results.json'
RESULTS_Z128_PATH = Path(__file__).parent.parent / 'data' / 'de_results_z128.json'
OUT_DIR = Path(__file__).parent.parent / 'data' / 'plots'

# Presentation-friendly style
plt.rcParams.update({
    'font.size': 13,
    'axes.titlesize': 15,
    'axes.labelsize': 14,
    'legend.fontsize': 11,
    'xtick.labelsize': 12,
    'ytick.labelsize': 12,
    'figure.dpi': 150,
    'savefig.bbox': 'tight',
    'savefig.pad_inches': 0.15,
})

COLORS = {
    'original': '#d62728',   # red
    'peg_ring': '#1f77b4',   # blue
    'optimized': '#2ca02c',  # green
}

LABELS = {
    'original': 'Original staircase [7,2,2,2,2,2,2,1]',
    'peg_ring': 'PEG ring [7,3,3,3,2,2,2,2]',
    'optimized': 'DE-optimized [7,4,4,4,4,3,3,3]',
}

MARKERS = {
    'original': 's',
    'peg_ring': '^',
    'optimized': 'o',
}


def load_results():
    with open(RESULTS_PATH) as f:
        return json.load(f)


def plot_fer_comparison(data):
    """FER vs lambda_s for all three matrices."""
    fig, ax = plt.subplots(figsize=(8, 5.5))

    fer_data = data['fer_comparison']
    lam_s = fer_data['lam_s_points']

    for key in ['original', 'peg_ring', 'optimized']:
        fer_vals = [fer_data[key][str(l)]['fer'] for l in lam_s]
        # Replace 0 with small value for log scale
        fer_plot = [max(f, 2e-3) for f in fer_vals]
        ax.semilogy(lam_s, fer_plot,
                     color=COLORS[key], marker=MARKERS[key], markersize=8,
                     linewidth=2, label=LABELS[key],
                     markeredgecolor='white', markeredgewidth=0.8)

    ax.set_xlabel(r'Signal photons/slot ($\lambda_s$)')
    ax.set_ylabel('Frame Error Rate (FER)')
    ax.set_title('FER Comparison: Rate-1/8 QC-LDPC (n=256, k=32)')
    ax.legend(loc='upper right', framealpha=0.9)
    ax.set_xlim(1.5, 10.5)
    ax.set_ylim(1e-3, 1.0)
    ax.grid(True, alpha=0.3, which='both')
    ax.set_xticks([2, 3, 4, 5, 7, 10])
    ax.set_xticklabels(['2', '3', '4', '5', '7', '10'])

    # Annotate the improvement
    ax.annotate('13x fewer\nframe errors',
                xy=(5, 0.01), xytext=(6.5, 0.06),
                fontsize=11, color=COLORS['optimized'],
                arrowprops=dict(arrowstyle='->', color=COLORS['optimized'], lw=1.5),
                ha='center')

    fig.savefig(OUT_DIR / 'fer_comparison.png')
    plt.close(fig)
    print(f'  Saved fer_comparison.png')


def plot_threshold_bars(data):
    """Bar chart of DE thresholds."""
    fig, ax = plt.subplots(figsize=(7, 5))

    thresholds = {
        'Original\nstaircase': data['reference_thresholds']['Original staircase [7,2,2,2,2,2,2,1]'],
        'PEG\nring': data['reference_thresholds']['PEG ring [7,3,3,3,2,2,2,2]'],
        'DE-\noptimized': data['best_threshold'],
    }
    # Shannon limit for rate 1/8 Poisson channel ~ 0.47 photons/slot
    shannon_limit = 0.47

    names = list(thresholds.keys())
    vals = list(thresholds.values())
    colors = [COLORS['original'], COLORS['peg_ring'], COLORS['optimized']]

    bars = ax.bar(names, vals, color=colors, width=0.55, edgecolor='white', linewidth=1.5)

    # Value labels on bars
    for bar, val in zip(bars, vals):
        ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.1,
                f'{val:.1f}', ha='center', va='bottom', fontweight='bold', fontsize=14)

    # Shannon limit line
    ax.axhline(y=shannon_limit, color='black', linestyle='--', linewidth=1.5, alpha=0.6)
    ax.text(2.42, shannon_limit + 0.12, f'Shannon limit ({shannon_limit})',
            fontsize=11, ha='right', va='bottom', style='italic', alpha=0.7)

    ax.set_ylabel(r'DE Threshold ($\lambda_s^*$, photons/slot)')
    ax.set_title('Decoding Threshold Comparison')
    ax.set_ylim(0, 6.5)
    ax.grid(True, alpha=0.2, axis='y')

    # Improvement annotation
    improvement_db = 10 * np.log10(thresholds['Original\nstaircase'] / thresholds['DE-\noptimized'])
    ax.annotate(f'{improvement_db:.1f} dB\nimprovement',
                xy=(2, thresholds['DE-\noptimized']),
                xytext=(2.35, 4.2),
                fontsize=12, fontweight='bold', color=COLORS['optimized'],
                arrowprops=dict(arrowstyle='->', color=COLORS['optimized'], lw=1.5),
                ha='center')

    fig.savefig(OUT_DIR / 'threshold_comparison.png')
    plt.close(fig)
    print(f'  Saved threshold_comparison.png')


def plot_degree_distribution(data):
    """VN degree distribution comparison."""
    fig, axes = plt.subplots(1, 3, figsize=(12, 4), sharey=True)

    distributions = {
        'Original staircase': [7, 2, 2, 2, 2, 2, 2, 1],
        'PEG ring': [7, 3, 3, 3, 2, 2, 2, 2],
        'DE-optimized': data['best_degrees'],
    }
    color_keys = ['original', 'peg_ring', 'optimized']
    col_labels = [f'c{i}' for i in range(8)]

    for ax, (name, degrees), ckey in zip(axes, distributions.items(), color_keys):
        colors_bar = [COLORS[ckey] if i > 0 else '#888888' for i in range(8)]
        ax.bar(col_labels, degrees, color=colors_bar, edgecolor='white', linewidth=1)
        ax.set_title(name, fontsize=12, fontweight='bold')
        ax.set_xlabel('Base matrix column')
        ax.set_ylim(0, 8)
        ax.grid(True, alpha=0.2, axis='y')

        # Annotate average parity degree
        avg_parity = np.mean(degrees[1:])
        ax.text(4, 7.2, f'avg parity dv = {avg_parity:.1f}',
                ha='center', fontsize=10, style='italic')

    axes[0].set_ylabel('Variable node degree')
    fig.suptitle('VN Degree Distributions (col 0 = info, cols 1-7 = parity)',
                 fontsize=13, y=1.02)
    fig.tight_layout()

    fig.savefig(OUT_DIR / 'degree_distributions.png')
    plt.close(fig)
    print(f'  Saved degree_distributions.png')


def plot_base_matrix_heatmap(data):
    """Heatmap of the constructed base matrix."""
    fig, ax = plt.subplots(figsize=(7, 5))

    H = np.array(data['constructed_matrix'])
    m_base, n_base = H.shape

    # Create display matrix: shift values where connected, NaN where not
    display = np.full_like(H, dtype=float, fill_value=np.nan)
    for r in range(m_base):
        for c in range(n_base):
            if H[r, c] >= 0:
                display[r, c] = H[r, c]

    im = ax.imshow(display, cmap='YlGnBu', aspect='auto', vmin=0, vmax=31)

    # Add text annotations
    for r in range(m_base):
        for c in range(n_base):
            if H[r, c] >= 0:
                ax.text(c, r, str(H[r, c]), ha='center', va='center',
                        fontsize=11, fontweight='bold', color='black')
            else:
                ax.text(c, r, '-', ha='center', va='center',
                        fontsize=11, color='#cccccc')

    ax.set_xticks(range(n_base))
    ax.set_xticklabels([f'c{i}' for i in range(n_base)])
    ax.set_yticks(range(m_base))
    ax.set_yticklabels([f'r{i}' for i in range(m_base)])
    ax.set_xlabel('Column (c0=info, c1-c7=parity)')
    ax.set_ylabel('Row (check node group)')
    ax.set_title(f'Optimized Base Matrix (Z=32, girth={data["matrix_checks"]["girth"]})')

    cbar = fig.colorbar(im, ax=ax, shrink=0.8, label='Circulant shift')

    fig.savefig(OUT_DIR / 'base_matrix_heatmap.png')
    plt.close(fig)
    print(f'  Saved base_matrix_heatmap.png')


def load_z128_results():
    """Load Z=128 results if available."""
    if RESULTS_Z128_PATH.exists():
        with open(RESULTS_Z128_PATH) as f:
            return json.load(f)
    return None


def plot_z128_comparison(data_z32, data_z128):
    """
    FER vs lambda_s comparison between Z=32 and Z=128 for the
    DE-optimized degree distribution [7,4,4,4,4,3,3,3].

    Shows both lifting factors on the same axes, with the optimized
    matrix in each case plus the PEG ring reference.
    """
    fig, ax = plt.subplots(figsize=(9, 6))

    # Colors and styles for Z=32 vs Z=128
    z32_color = '#2ca02c'    # green (matches optimized)
    z128_color = '#9467bd'   # purple
    z32_peg_color = '#1f77b4'  # blue (matches peg_ring)
    z128_peg_color = '#ff7f0e' # orange

    # Z=32 optimized data (from de_results.json)
    fer_z32 = data_z32['fer_comparison']
    lam_s_z32 = fer_z32['lam_s_points']
    fer_z32_opt = [fer_z32['optimized'][str(l)]['fer'] for l in lam_s_z32]
    fer_z32_peg = [fer_z32['peg_ring'][str(l)]['fer'] for l in lam_s_z32]

    # Z=128 data
    fer_z128 = data_z128['fer_results']
    lam_s_z128 = fer_z128['lam_s_points']
    fer_z128_opt = [fer_z128['optimized'][str(l)]['fer'] for l in lam_s_z128]
    fer_z128_peg = [fer_z128['peg_ring'][str(l)]['fer'] for l in lam_s_z128]

    # Replace 0 with small value for log scale
    floor = 5e-3

    # Plot Z=32 curves (solid lines)
    ax.semilogy(lam_s_z32, [max(f, floor) for f in fer_z32_opt],
                color=z32_color, marker='o', markersize=8, linewidth=2,
                label='Optimized Z=32 (n=256, offset MS)',
                markeredgecolor='white', markeredgewidth=0.8)
    ax.semilogy(lam_s_z32, [max(f, floor) for f in fer_z32_peg],
                color=z32_peg_color, marker='^', markersize=8, linewidth=2,
                label='PEG ring Z=32 (n=256, offset MS)',
                markeredgecolor='white', markeredgewidth=0.8,
                linestyle='--', alpha=0.7)

    # Plot Z=128 curves (dashed lines, larger markers)
    ax.semilogy(lam_s_z128, [max(f, floor) for f in fer_z128_opt],
                color=z128_color, marker='D', markersize=9, linewidth=2.5,
                label=f'Optimized Z=128 (n=1024, norm MS, a={data_z128["alpha"]})',
                markeredgecolor='white', markeredgewidth=0.8)
    ax.semilogy(lam_s_z128, [max(f, floor) for f in fer_z128_peg],
                color=z128_peg_color, marker='v', markersize=9, linewidth=2.5,
                label=f'PEG ring Z=128 (n=1024, norm MS, a={data_z128["alpha"]})',
                markeredgecolor='white', markeredgewidth=0.8,
                linestyle='--', alpha=0.7)

    ax.set_xlabel(r'Signal photons/slot ($\lambda_s$)')
    ax.set_ylabel('Frame Error Rate (FER)')
    ax.set_title('FER Comparison: Z=32 vs Z=128\nDE-optimized [7,4,4,4,4,3,3,3]')
    ax.legend(loc='upper right', framealpha=0.9, fontsize=10)
    ax.set_xlim(1.0, 11.0)
    ax.set_ylim(floor / 2, 1.1)
    ax.grid(True, alpha=0.3, which='both')
    ax.set_xticks([1.5, 2, 3, 4, 5, 7, 10])
    ax.set_xticklabels(['1.5', '2', '3', '4', '5', '7', '10'])

    # Add Z=128 matrix info annotation
    girth_128 = data_z128.get('girth', '?')
    ax.text(0.02, 0.02,
            f'Z=128: girth={girth_128}, n=1024, k=128\n'
            f'Z=32: girth={data_z32["matrix_checks"]["girth"]}, n=256, k=32',
            transform=ax.transAxes, fontsize=9, verticalalignment='bottom',
            bbox=dict(boxstyle='round,pad=0.3', facecolor='wheat', alpha=0.5))

    fig.savefig(OUT_DIR / 'z128_fer_comparison.png')
    plt.close(fig)
    print(f'  Saved z128_fer_comparison.png')


def main():
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    data = load_results()

    print('Generating plots...')
    plot_fer_comparison(data)
    plot_threshold_bars(data)
    plot_degree_distribution(data)
    plot_base_matrix_heatmap(data)

    # Z=128 comparison plot
    data_z128 = load_z128_results()
    if data_z128 is not None:
        plot_z128_comparison(data, data_z128)
    else:
        print('  Skipping Z=128 comparison (data/de_results_z128.json not found)')

    print(f'\nAll plots saved to {OUT_DIR}/')


if __name__ == '__main__':
    main()