example_nw_estimator.py

import numpy as np
import matplotlib.pyplot as plt
import os


# ------------------------------------------------------------
# Automatically create folder using script name
# ------------------------------------------------------------
script_name = os.path.splitext(os.path.basename(__file__))[0]
save_dir = script_name
os.makedirs(save_dir, exist_ok=True)


# ------------------------------------------------------------
# Data: satellites (keys) and per-satellite inferred positions (values)
# ------------------------------------------------------------
s = np.array([-10.0, 0.0, 8.0])  # satellite positions (keys)

# signed pseudorange-like measurements z_i ≈ s_i - x* + noise
z = np.array([-10.2, 0.1, 8.6])

# per-satellite "candidate" position estimates: x_hat_i = s_i - z_i
x_hat_i = s - z  # [0.2, -0.1, -0.6]

# ------------------------------------------------------------
# NW / Attention parameters
# ------------------------------------------------------------
sigma = 5.0

def kernel(x_query, s_key, sigma):
    """Gaussian kernel k(x, s) = exp(-(x-s)^2 / (2*sigma^2))"""
    return np.exp(-((x_query - s_key) ** 2) / (2.0 * sigma ** 2))

def nw_estimate(x_query):
    """NW regression output: x_hat(x) = sum_i alpha_i(x) * x_hat_i"""
    k = kernel(x_query, s, sigma)              # shape: (3,)
    alpha = k / np.sum(k)                     # normalized weights
    return np.sum(alpha * x_hat_i), alpha

# ------------------------------------------------------------
# Build curve
# ------------------------------------------------------------
x_grid = np.linspace(-15, 15, 600)
y_grid = np.zeros_like(x_grid)

for idx, xq in enumerate(x_grid):
    y_grid[idx], _ = nw_estimate(xq)

# pick a specific query point x0 to annotate
x0 = 0.0
y0, alpha0 = nw_estimate(x0)

# ------------------------------------------------------------
# Enhanced Plot
# ------------------------------------------------------------
plt.figure(figsize=(10, 5.5))

# Plot NW regression curve
plt.plot(x_grid, y_grid, linewidth=3.5, label=r"NW curve $\hat{x}(x)$",
         color='#2E86AB', zorder=2)

# Plot sample points with enhanced styling
plt.scatter(s, x_hat_i, s=150, marker="o",
           label=r"Samples $(s_i,\hat{x}_i)$",
           color='#A23B72', edgecolors='white', linewidth=1.5,
           zorder=5, alpha=0.9)

# Add vertical lines from samples to curve for visual clarity
for i in range(len(s)):
    y_on_curve, _ = nw_estimate(s[i])
    plt.plot([s[i], s[i]], [x_hat_i[i], y_on_curve],
            linestyle=':', linewidth=1, color='gray', alpha=0.4, zorder=1)

# Query point line
plt.axvline(x0, linewidth=2, linestyle="--", color='#F18F01',
           label=rf"Query $x_0={x0:g}$", alpha=0.7, zorder=3)

# NW output point
plt.scatter([x0], [y0], s=200, marker="X", color='#C73E1D',
           edgecolors='white', linewidth=1.5,
           label=rf"NW output $\hat{{x}}(x_0)={y0:.2f}$", zorder=6)

# True position reference line
plt.axhline(0, linewidth=1.5, linestyle='-.', color='#6A994E',
           alpha=0.6, label=r"True position $x^*=0$", zorder=1)

# Annotations for sample points
for i, (si, xi) in enumerate(zip(s, x_hat_i)):
    plt.annotate(f'$s_{i+1}$',
                xy=(si, xi),
                xytext=(si, xi + 0.15),
                fontsize=19,
                ha='center',
                color='#A23B72',
                weight='bold')

# Add weight annotations at x0
# Add weight annotations near each sample point
for i in range(len(s)):
    if alpha0[i] > 0.05:  # Only show significant weights
        plt.annotate(f'$α_{i+1}={alpha0[i]:.3f}$',
                    xy=(s[i], x_hat_i[i] - 0.22),
                    fontsize=14,
                    ha='center',
                    bbox=dict(boxstyle='round,pad=0.3',
                             facecolor='lightyellow',
                             edgecolor='gray',
                             alpha=0.7),
                    zorder=7)

# offset_y = y0 + 0.25
# for i in range(len(s)):
#     if alpha0[i] > 0.05:  # Only show significant weights
#         plt.annotate(f'$α_{i+1}={alpha0[i]:.3f}$',
#                     xy=(s[i], offset_y),
#                     fontsize=9,
#                     ha='center',
#                     bbox=dict(boxstyle='round,pad=0.3',
#                              facecolor='lightyellow',
#                              edgecolor='gray',
#                              alpha=0.7),
#                     zorder=7)

# Labels and title
plt.xlabel(r"Query / prior location $x$", fontsize=15.5, weight='bold')
plt.ylabel(r"Estimated location $\hat{x}(x)$", fontsize=15.5, weight='bold')
plt.title("NW Regression View of Attention-Based LEO Satellite Localization (1D)",
         fontsize=16.5, weight='bold', pad=15)

plt.xticks(fontsize=15.5, fontweight='bold')
plt.yticks(fontsize=15.5, fontweight='bold')

# Legend with better positioning
plt.legend(fontsize=13, loc='lower left', framealpha=0.95,
          edgecolor='gray', fancybox=True)

# Grid styling
plt.grid(True, alpha=0.25, linestyle='-', linewidth=0.5)
plt.grid(True, alpha=0.1, linestyle='-', linewidth=0.5, which='minor')

# Set axis limits for better visualization
plt.xlim(-15, 15)
plt.ylim(-1.2, 0.5)

# Tight layout
plt.tight_layout()

plt.savefig(os.path.join(save_dir, f"{script_name}.pdf"), dpi=300, bbox_inches='tight')
plt.savefig(os.path.join(save_dir, f"{script_name}.png"), dpi=300, bbox_inches='tight')

plt.show()

# ------------------------------------------------------------
# Enhanced console output
# ------------------------------------------------------------
print("="*70)
print("NADARAYA-WATSON LEO LOCALIZATION EXAMPLE")
print("="*70)
print(f"\n{'Satellite Data:':<30}")
print(f"  Positions s_i:              {s}")
print(f"  Measurements z_i:           {z}")
print(f"  Candidate estimates x_hat_i: {x_hat_i}")

print(f"\n{'Query Point Analysis:':<30}")
print(f"  Query location x_0:          {x0}")
print(f"  True position x*:            0")

print(f"\n{'Kernel Values:':<30}")
k_vals = kernel(x0, s, sigma)
for i, (si, ki) in enumerate(zip(s, k_vals)):
    print(f"  k(x_0, s_{i+1}) = k({x0}, {si:>4.0f}):  {ki:.6f}")

print(f"\n{'Normalized Weights:':<30}")
for i, (si, ai) in enumerate(zip(s, alpha0)):
    print(f"  α_{i+1}(x_0, s_{i+1}):            {ai:.6f}  ({ai*100:.2f}%)")

print(f"\n{'Final Estimate:':<30}")
print(f"  x_hat(x_0):                  {y0:.6f}")
print(f"  Estimation error:            {abs(y0 - 0):.6f}")

print(f"\n{'Weighted Contributions:':<30}")
for i in range(len(s)):
    contrib = alpha0[i] * x_hat_i[i]
    print(f"  α_{i+1} × x_hat_{i+1} = {alpha0[i]:.3f} × {x_hat_i[i]:>5.2f} = {contrib:>6.3f}")

print("="*70)