example_attention_leo_localization.py

from __future__ import annotations

import os
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D


def softmax(x: np.ndarray) -> np.ndarray:
    """Numerically stable softmax computation."""
    x = np.asarray(x, dtype=float)
    x = x - np.max(x)
    ex = np.exp(x)
    return ex / np.sum(ex)


def save_figure(fig: plt.Figure, out_base: Path) -> None:
    """
    Save figure in both PNG (300 dpi) and PDF formats.

    Args:
        fig: Matplotlib figure object
        out_base: Output path without extension
    """
    png_path = out_base.with_suffix(".png")
    pdf_path = out_base.with_suffix(".pdf")
    fig.savefig(png_path, dpi=300, bbox_inches="tight")
    fig.savefig(pdf_path, bbox_inches="tight")
    plt.close(fig)


def main() -> None:
    # ================================================================
    # Configuration and Output Setup
    # ================================================================
    try:
        script_path = Path(__file__).resolve()
        out_dir = script_path.parent
        prefix = script_path.stem
    except NameError:
        out_dir = Path(os.getcwd()).resolve()
        prefix = "attention_leo_example"

    # Professional color scheme
    colors = {
        'uniform': '#7FA6C9',  # Muted blue
        'attention': '#E87D3E',  # Vibrant orange
        'candidate': '#5C8A6F',  # Forest green
        'prior': '#9B59B6',  # Purple
        'grid': '#E0E0E0'  # Light gray
    }

    # ================================================================
    # LEO Satellite Localization Example
    # ================================================================
    # Satellite positions and pseudorange measurements
    s = np.array([100.0, 140.0, 180.0])
    rho = np.array([30.0, 72.0, 110.0])

    # Candidate position values
    v = s - rho  # [70, 68, 70]

    # Prior estimate
    x0 = 70.0

    # Quality indicators
    r = np.abs(s - x0 - rho)  # Residuals: [0, 2, 0]
    snr = np.array([0.9, 0.4, 0.8])  # Signal-to-noise ratios

    # Key construction: [SNR, -residual]
    K = np.column_stack([snr, -r])

    # Query: prefer high SNR and small residuals
    q = np.array([1.0, 1.0])

    # Attention computation
    scores = q @ K.T
    alpha = softmax(scores)

    # Fusion results
    x_attn = float(np.sum(alpha * v))
    alpha_uniform = np.ones_like(alpha) / alpha.size
    x_uniform = float(np.mean(v))

    # ================================================================
    # Display Results
    # ================================================================
    np.set_printoptions(precision=6, suppress=True)
    print("=" * 60)
    print("LEO Satellite Localization: Attention-Based Fusion")
    print("=" * 60)
    print(f"\nSatellite positions s:        {s}")
    print(f"Pseudorange measurements ρ:   {rho}")
    print(f"Candidate positions v = s-ρ:  {v}")
    print(f"\nPrior estimate x₀:            {x0}")
    print(f"Residuals |r|:                {r}")
    print(f"Signal-to-noise ratios:       {snr}")
    print(f"\nKey matrix K [SNR, -|r|]:")
    print(K)
    print(f"\nQuery q:                      {q}")
    print(f"Similarity scores qKᵀ:        {scores}")
    print(f"\nAttention weights α:          {alpha}")
    print(f"Uniform weights:              {alpha_uniform}")
    print(f"\nFused estimate (attention):   {x_attn:.4f}")
    print(f"Fused estimate (uniform):     {x_uniform:.4f}")
    print("=" * 60)

    # ================================================================
    # Figure 1: Weight Comparison with Enhanced Visualization
    # ================================================================
    sat_indices = np.arange(1, 4)

    # Increased figure height to accommodate labels
    fig1, ax1 = plt.subplots(figsize=(6.5, 4.5))

    bar_width = 0.35
    x_pos = sat_indices

    # Create bars with enhanced styling
    bars1 = ax1.bar(x_pos - bar_width / 2, alpha_uniform, width=bar_width,
                    label='Uniform pooling', color=colors['uniform'],
                    edgecolor='white', linewidth=1.5, alpha=0.85)

    bars2 = ax1.bar(x_pos + bar_width / 2, alpha, width=bar_width,
                    label='Attention weights', color=colors['attention'],
                    edgecolor='white', linewidth=1.5, alpha=0.95)

    # Add value labels on bars
    for bars in [bars1, bars2]:
        for bar in bars:
            height = bar.get_height()
            ax1.text(bar.get_x() + bar.get_width() / 2., height,
                     f'{height:.3f}',
                     ha='center', va='bottom', fontsize=9, fontweight='bold')

    # Styling
    ax1.set_xticks(sat_indices)
    ax1.set_xticklabels([f'Sat {i}' for i in sat_indices], fontsize=12)
    ax1.set_xlabel('Satellite Index', fontsize=13, fontweight='bold', labelpad=35)
    ax1.set_ylabel('Weight', fontsize=13, fontweight='bold')
    ax1.set_title('Uniform Pooling vs. Attention-Based Weighting',
                  fontsize=14, fontweight='bold', pad=15)
    ax1.set_ylim(0.0, max(0.65, float(np.max(alpha)) * 1.2))

    # Enhanced grid
    ax1.grid(axis='y', alpha=0.3, linestyle='--', linewidth=0.8, color=colors['grid'])
    ax1.set_axisbelow(True)

    # Legend with better positioning
    ax1.legend(loc='upper right', frameon=True, shadow=True,
               fontsize=11, framealpha=0.95, edgecolor='gray')

    # Add quality annotations below x-axis (moved down to avoid overlap)
    quality_labels = [
        f"SNR={snr[i]:.1f}, |r|={r[i]:.0f}"
        for i in range(len(sat_indices))
    ]
    for i, (x, label) in enumerate(zip(sat_indices, quality_labels)):
        # Use figure coordinates to place text below axis
        ax1.text(x, -0.08, label, ha='center', va='top',
                 fontsize=12, fontweight='bold', style='italic', color='#555555',
                 transform=ax1.transData)

    fig1.tight_layout()
    save_figure(fig1, out_dir / f"{prefix}_weights")

    # ================================================================
    # Figure 2: Enhanced Position Visualization
    # ================================================================
    fig2, ax2 = plt.subplots(figsize=(9.5, 4.5))

    # Font size configuration for position figure
    FS_VALUE = 14  # v_i labels
    FS_QUALITY = 14  # SNR and alpha
    FS_LEGEND = 13  # legend
    FS_TICKS = 12

    # Vertical offset for different elements
    y_offset = 0.20

    # Plot candidates - use unique positions for visualization
    # For overlapping points at x=70, plot them with slight x-offset for clarity
    v_plot = v.copy()
    v_plot[0] = v[0] - 0.15  # v1 slightly left
    v_plot[2] = v[2] + 0.15  # v3 slightly right

    # Plot each candidate individually for better control
    for i in range(len(v)):
        ax2.scatter(v_plot[i], 0.0, s=200,
                    c=colors['candidate'], marker='o',
                    edgecolors='white', linewidths=2,
                    zorder=5, alpha=0.9)

    # Add a note for the overlapping candidates
    ax2.plot([v_plot[0], v_plot[2]], [0.0, 0.0],
             color=colors['candidate'], linewidth=2, alpha=0.3, zorder=3)

    # Smart annotation positioning to avoid overlap
    # Stagger labels vertically for overlapping positions
    annotation_configs = [
        {'offset': (-35, 45), 'ha': 'right', 'arrow_rad': -0.3},  # v1: upper left
        {'offset': (0, 35), 'ha': 'center', 'arrow_rad': 0},  # v2: centered
        {'offset': (35, 45), 'ha': 'left', 'arrow_rad': 0.3}  # v3: upper right
    ]

    for i, vi in enumerate(v):
        config = annotation_configs[i]

        # Value label with background box
        idx = i + 1
        label_text = f'$v_{idx}={vi:.0f}$'
        ax2.annotate(label_text,
                     xy=(v_plot[i], 0.0), xytext=config['offset'],
                     textcoords='offset points', ha=config['ha'],
                     fontsize=FS_VALUE, fontweight='bold',
                     bbox=dict(boxstyle='round,pad=0.5',
                               facecolor=colors['candidate'],
                               edgecolor='white', alpha=0.35, linewidth=1.5),
                     arrowprops=dict(arrowstyle='->',
                                     connectionstyle=f"arc3,rad={config['arrow_rad']}",
                                     color=colors['candidate'], lw=1.2, alpha=0.6))

    # Quality info below - with better spacing for overlapping positions
    quality_offsets = [
        (-35, -40),  # v1: left
        (0, -40),  # v2: center
        (35, -40)  # v3: right
    ]

    for i in range(len(v)):
        offset_x, offset_y = quality_offsets[i]
        idx = i + 1
        #quality_text = f'SNR={snr[i]:.1f}\n$\\alpha_{idx}$={alpha[i]:.3f}'
        quality_text = (
            f'SNR={snr[i]:.1f}\n'
            f'$\\alpha_{{{idx}}}=\\mathbf{{{alpha[i]:.3f}}}$'
        )
        ax2.annotate(quality_text,
                     xy=(v_plot[i], 0.0), xytext=(offset_x, offset_y),
                     textcoords='offset points', ha='center',
                     fontsize=FS_QUALITY, color='#555555', style='italic')

    # Plot fusion results with distinct markers
    ax2.scatter([x_uniform], [-y_offset], marker='s', s=220,
                c=colors['uniform'], edgecolors='white', linewidths=2,
                zorder=6, alpha=0.9)

    ax2.scatter([x_attn], [y_offset], marker='*', s=320,
                c=colors['attention'], edgecolors='white', linewidths=2,
                zorder=6, alpha=0.95)

    # Add vertical lines to show position
    ax2.axvline(x_uniform, color=colors['uniform'], linestyle='--',
                alpha=0.4, linewidth=1.5, zorder=1)
    ax2.axvline(x_attn, color=colors['attention'], linestyle='--',
                alpha=0.4, linewidth=1.5, zorder=1)

    # Enhanced baseline
    xmin = float(min(np.min(v), x_uniform, x_attn)) - 2.0
    xmax = float(max(np.max(v), x_uniform, x_attn)) + 2.0
    ax2.hlines(0.0, xmin, xmax, linewidth=2.0, color='black', alpha=0.3, zorder=2)

    # Styling
    ax2.set_xlim(xmin, xmax)
    ax2.set_ylim(-0.45, 0.45)
    ax2.set_yticks([])
    ax2.set_xlabel('Position', fontsize=14, fontweight='bold', labelpad=10)
    ax2.set_title('Candidate Positions and Fusion Estimates',
                  fontsize=17, fontweight='bold', pad=15)
    ax2.tick_params(axis='both', labelsize=FS_TICKS)

    # Enhanced legend - create custom legend
    legend_elements = [
        Line2D([0], [0], marker='o', color='w',
               markerfacecolor=colors['candidate'], markersize=13,
               markeredgecolor='white', markeredgewidth=1.5,
               label='Satellite candidates'),
        Line2D([0], [0], marker='s', color='w',
               markerfacecolor=colors['uniform'], markersize=13,
               markeredgecolor='white', markeredgewidth=1.5,
               label=f'Uniform: $\\bar{{x}}={x_uniform:.2f}$'),
        Line2D([0], [0], marker='*', color='w',
               markerfacecolor=colors['attention'], markersize=13,
               markeredgecolor='white', markeredgewidth=1.5,
               label=f'Attention: $\\hat{{x}}={x_attn:.2f}$')
    ]
    ax2.legend(
        handles=legend_elements,
        loc='lower left',
        ncol=1,
        frameon=True,
        fontsize=FS_LEGEND,
        borderpad=0.6,
        labelspacing=0.4,
        handletextpad=0.5,
        borderaxespad=0.6,
        framealpha=0.90,
        edgecolor='gray'
    )

    # Add grid
    ax2.grid(axis='x', alpha=0.2, linestyle=':', linewidth=0.8,
             color=colors['grid'])
    ax2.set_axisbelow(True)

    fig2.tight_layout()
    save_figure(fig2, out_dir / f"{prefix}_positions")

    # ================================================================
    # Summary
    # ================================================================
    print(f"\n{'=' * 60}")
    print(f"Figures saved to: {out_dir}")
    print(f"  • {prefix}_weights.png / .pdf")
    print(f"  • {prefix}_positions.png / .pdf")
    print(f"{'=' * 60}\n")


if __name__ == "__main__":
    main()