403 lines
15 KiB
Python
403 lines
15 KiB
Python
|
||
import os
|
||
import math
|
||
import time
|
||
import numpy as np
|
||
import torch
|
||
import matplotlib.pyplot as plt
|
||
|
||
# ============================================================
|
||
# 0. 全局设置
|
||
# ============================================================
|
||
plt.rcParams['font.sans-serif'] = [
|
||
'Microsoft YaHei', 'SimHei', 'PingFang SC',
|
||
'Noto Sans CJK SC', 'Arial Unicode MS', 'DejaVu Sans'
|
||
]
|
||
plt.rcParams['axes.unicode_minus'] = False
|
||
plt.rcParams['mathtext.fontset'] = 'stix'
|
||
|
||
SAVE_DIR = "robust_tsc_omp_cuda_results"
|
||
os.makedirs(SAVE_DIR, exist_ok=True)
|
||
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
torch.manual_seed(2024)
|
||
np.random.seed(2024)
|
||
|
||
# ============================================================
|
||
# 1. 参数区
|
||
# 说明：
|
||
# - 本脚本默认使用 CUDA（若可用）
|
||
# - 为避免“卡死”，默认采用较稳妥的 Monte Carlo 与扫描步长
|
||
# - 所有图都会生成并在脚本结束时统一以窗口形式显示
|
||
# ============================================================
|
||
ARRAY_TYPE = 'NULA' # 'ULA' 或 'NULA'
|
||
N_ant = 24
|
||
N_RF = 4
|
||
d_lambda = 0.5
|
||
|
||
target_angle = 5.0
|
||
clutter_angle = -15.0
|
||
noise_power = 1.0
|
||
INR_dB = 30.0
|
||
INR_linear = 10 ** (INR_dB / 10)
|
||
|
||
L_snapshots = 100
|
||
Monte_Carlo = 300
|
||
SNR_dB_range = np.arange(-10, 22, 2)
|
||
|
||
theta_prior = 5.0
|
||
window_half_width = 5.0
|
||
sector_min = theta_prior - window_half_width
|
||
sector_max = theta_prior + window_half_width
|
||
|
||
# TSC-OMP 字典与测角搜索栅格
|
||
tsc_grid_np = np.arange(sector_min, sector_max + 0.1, 0.1)
|
||
scan_grid_np = np.arange(sector_min, sector_max + 0.001, 0.05) # 比 0.01 更快
|
||
|
||
# 对角加载系数
|
||
alpha_list = [0.00, 0.01, 0.03, 0.05, 0.10]
|
||
|
||
# 是否保存图片
|
||
SAVE_FIG = True
|
||
|
||
# ============================================================
|
||
# 2. 阵列位置
|
||
# ============================================================
|
||
def get_array_positions(array_type='ULA', N=24, d=0.5):
|
||
if array_type.upper() == 'ULA':
|
||
pos = np.arange(-(N - 1) / 2, (N - 1) / 2 + 1) * d
|
||
return pos.astype(np.float32)
|
||
elif array_type.upper() == 'NULA':
|
||
# 与论文一致的示意性 NULA
|
||
base = 0.7
|
||
jump = 1.03
|
||
pos = []
|
||
x = 0.0
|
||
for _ in range(12):
|
||
pos.append(x)
|
||
x += base
|
||
x += (jump - base)
|
||
for _ in range(12, 24):
|
||
pos.append(x)
|
||
x += base
|
||
pos = np.array(pos, dtype=np.float32)
|
||
pos = pos - np.mean(pos)
|
||
return pos
|
||
else:
|
||
raise ValueError("ARRAY_TYPE 只能是 'ULA' 或 'NULA'")
|
||
|
||
array_pos_np = get_array_positions(ARRAY_TYPE, N_ant, d_lambda)
|
||
array_pos_t = torch.tensor(array_pos_np, dtype=torch.float32, device=device).view(-1, 1)
|
||
|
||
# ============================================================
|
||
# 3. 基础函数（CUDA）
|
||
# ============================================================
|
||
def sync():
|
||
if device.type == 'cuda':
|
||
torch.cuda.synchronize()
|
||
|
||
def steering_vector(theta_deg):
|
||
theta = torch.as_tensor(theta_deg, dtype=torch.float32, device=device)
|
||
theta_rad = torch.deg2rad(theta)
|
||
phase = 2 * torch.pi * array_pos_t * torch.sin(theta_rad)
|
||
return torch.exp(1j * phase)
|
||
|
||
def steering_matrix(theta_deg_array):
|
||
theta = torch.as_tensor(theta_deg_array, dtype=torch.float32, device=device).view(1, -1)
|
||
theta_rad = torch.deg2rad(theta)
|
||
phase = 2 * torch.pi * array_pos_t @ torch.sin(theta_rad)
|
||
return torch.exp(1j * phase)
|
||
|
||
def build_dictionary(angle_grid_np, normalize=True):
|
||
A = steering_matrix(angle_grid_np)
|
||
if normalize:
|
||
A = A / math.sqrt(A.shape[0])
|
||
return A
|
||
|
||
def build_prior_broad(sector_min, sector_max, step=0.5):
|
||
angles = np.arange(sector_min, sector_max + 1e-12, step, dtype=np.float32)
|
||
A = steering_matrix(angles)
|
||
a = torch.sum(A, dim=1, keepdim=True)
|
||
a = a / torch.linalg.norm(a)
|
||
return a
|
||
|
||
def extract_omp(w_opt, A_dic, angle_grid_np, N_RF):
|
||
"""
|
||
CUDA 版 OMP / TSC-OMP
|
||
"""
|
||
res = w_opt.clone()
|
||
F_RF = None
|
||
chosen_angles = []
|
||
|
||
for _ in range(N_RF):
|
||
proj = torch.abs(A_dic.mH @ res).flatten()
|
||
idx = int(torch.argmax(proj).item())
|
||
atom = A_dic[:, idx:idx+1]
|
||
F_RF = atom if F_RF is None else torch.cat((F_RF, atom), dim=1)
|
||
chosen_angles.append(float(angle_grid_np[idx]))
|
||
res = w_opt - F_RF @ (torch.linalg.pinv(F_RF) @ w_opt)
|
||
|
||
return F_RF, chosen_angles
|
||
|
||
def build_prior_weight(target_angle, clutter_angle, sector_min, sector_max):
|
||
a_clutter = steering_vector(clutter_angle)
|
||
R_interf = noise_power * torch.eye(N_ant, dtype=torch.complex64, device=device) \
|
||
+ INR_linear * (a_clutter @ a_clutter.mH)
|
||
a_prior = build_prior_broad(sector_min, sector_max, step=0.5)
|
||
w_prior = torch.linalg.solve(R_interf, a_prior)
|
||
w_prior = w_prior / torch.linalg.norm(w_prior)
|
||
return w_prior
|
||
|
||
def calc_hybrid_weight_loaded(F_RF, R_in, a_target, alpha_dl=0.05):
|
||
a_eff = F_RF.mH @ a_target
|
||
R_eff = F_RF.mH @ R_in @ F_RF
|
||
gamma = alpha_dl * torch.real(torch.trace(R_eff)) / R_eff.shape[0]
|
||
R_eff_loaded = R_eff + gamma * torch.eye(R_eff.shape[0], dtype=torch.complex64, device=device)
|
||
inv_R = torch.linalg.inv(R_eff_loaded)
|
||
F_BB = (inv_R @ a_eff) / (a_eff.mH @ inv_R @ a_eff)
|
||
w = F_RF @ F_BB
|
||
w = w / torch.linalg.norm(w)
|
||
cond_number = torch.linalg.cond(R_eff_loaded).real.item()
|
||
return w, R_eff_loaded, cond_number
|
||
|
||
def robust_capon_spectrum_batch(F_RF, X_batch, A_scan, alpha_dl=0.05):
|
||
"""
|
||
X_batch: [MC, N_ant, L]
|
||
A_scan: [N_ant, Ns]
|
||
返回：
|
||
P: [MC, Ns]
|
||
mean_cond: 平均条件数
|
||
"""
|
||
Y = torch.matmul(F_RF.mH.unsqueeze(0), X_batch) # [MC, N_RF, L]
|
||
R_bb = torch.matmul(Y, Y.mH) / X_batch.shape[-1] # [MC, N_RF, N_RF]
|
||
|
||
trace_R = torch.real(torch.diagonal(R_bb, dim1=-2, dim2=-1).sum(-1))
|
||
gamma = alpha_dl * trace_R / F_RF.shape[1]
|
||
eye = torch.eye(F_RF.shape[1], dtype=torch.complex64, device=device).unsqueeze(0)
|
||
R_bb_loaded = R_bb + gamma[:, None, None] * eye
|
||
|
||
inv_R = torch.linalg.inv(R_bb_loaded) # [MC, N_RF, N_RF]
|
||
A_eff = F_RF.mH @ A_scan # [N_RF, Ns]
|
||
|
||
# P = 1 / |a^H invR a|
|
||
denom = torch.einsum('rn,brs,sn->bn', A_eff.conj(), inv_R, A_eff)
|
||
P = 1.0 / torch.clamp(torch.abs(denom), min=1e-12)
|
||
|
||
cond_vals = torch.linalg.cond(R_bb_loaded).real
|
||
mean_cond = torch.mean(cond_vals).item()
|
||
return P, mean_cond
|
||
|
||
def calc_sinr(w, a_target, a_clutter, sig_p):
|
||
R_interf = noise_power * torch.eye(N_ant, dtype=torch.complex64, device=device) \
|
||
+ INR_linear * (a_clutter @ a_clutter.mH)
|
||
num = sig_p * (torch.abs(w.mH @ a_target).item() ** 2)
|
||
den = torch.real(w.mH @ R_interf @ w).item()
|
||
return 10 * np.log10(max(num / max(den, 1e-12), 1e-12))
|
||
|
||
# ============================================================
|
||
# 4. 固定 TSC-OMP 模拟矩阵
|
||
# ============================================================
|
||
print(f"当前设备: {device}")
|
||
sync()
|
||
tic = time.perf_counter()
|
||
|
||
w_prior = build_prior_weight(
|
||
target_angle=target_angle,
|
||
clutter_angle=clutter_angle,
|
||
sector_min=sector_min,
|
||
sector_max=sector_max
|
||
)
|
||
A_dic_tsc = build_dictionary(tsc_grid_np, normalize=True)
|
||
F_RF_tsc, chosen_angles = extract_omp(w_prior, A_dic_tsc, tsc_grid_np, N_RF)
|
||
|
||
a_target = steering_vector(target_angle)
|
||
a_clutter = steering_vector(clutter_angle)
|
||
|
||
sync()
|
||
print(f"TSC-OMP 选中的模拟原子角度: {chosen_angles}")
|
||
print(f"固定模拟矩阵构建耗时: {time.perf_counter() - tic:.3f} s")
|
||
|
||
# ============================================================
|
||
# 5. 静态方向图：不同加载系数
|
||
# ============================================================
|
||
sig_power_static = 10 ** (0 / 10.0)
|
||
s_t = np.sqrt(sig_power_static / 2) * (
|
||
torch.randn(1, L_snapshots, dtype=torch.float32, device=device)
|
||
+ 1j * torch.randn(1, L_snapshots, dtype=torch.float32, device=device)
|
||
)
|
||
s_c = np.sqrt(INR_linear / 2) * (
|
||
torch.randn(1, L_snapshots, dtype=torch.float32, device=device)
|
||
+ 1j * torch.randn(1, L_snapshots, dtype=torch.float32, device=device)
|
||
)
|
||
n_t = np.sqrt(noise_power / 2) * (
|
||
torch.randn(N_ant, L_snapshots, dtype=torch.float32, device=device)
|
||
+ 1j * torch.randn(N_ant, L_snapshots, dtype=torch.float32, device=device)
|
||
)
|
||
X_static = a_target @ s_t + a_clutter @ s_c + n_t
|
||
R_in_static = (X_static @ X_static.mH) / L_snapshots
|
||
|
||
theta_scan_np = np.arange(-60, 60.1, 0.1)
|
||
A_bp = steering_matrix(theta_scan_np)
|
||
|
||
plt.figure(figsize=(10, 6))
|
||
for alpha_dl in alpha_list:
|
||
w_hyb, _, _ = calc_hybrid_weight_loaded(F_RF_tsc, R_in_static, a_target, alpha_dl=alpha_dl)
|
||
BP = torch.abs(w_hyb.mH @ A_bp).flatten() ** 2
|
||
BP = 10 * torch.log10(BP / torch.max(BP) + 1e-12)
|
||
plt.plot(theta_scan_np, BP.detach().cpu().numpy(), linewidth=2, label=f'α={alpha_dl:.2f}')
|
||
|
||
plt.axvline(target_angle, color='g', linestyle='--', label=f'目标 {target_angle:.2f}°')
|
||
plt.axvline(clutter_angle, color='r', linestyle='--', label=f'杂波 {clutter_angle:.2f}°')
|
||
plt.xlabel('角度 θ (°)')
|
||
plt.ylabel('归一化增益 (dB)')
|
||
plt.title(f'不同对角加载系数下 TSC-OMP 波束方向图 ({ARRAY_TYPE})')
|
||
plt.grid(True, linestyle='--', alpha=0.5)
|
||
plt.legend()
|
||
plt.tight_layout()
|
||
if SAVE_FIG:
|
||
plt.savefig(os.path.join(SAVE_DIR, f'BP_loaded_compare_{ARRAY_TYPE}.png'), dpi=300)
|
||
|
||
# ============================================================
|
||
# 6. 输出 SINR / RMSE / 条件数 对比
|
||
# ============================================================
|
||
SINR_loaded = {alpha: [] for alpha in alpha_list}
|
||
RMSE_loaded = {alpha: [] for alpha in alpha_list}
|
||
COND_loaded = {alpha: [] for alpha in alpha_list}
|
||
|
||
A_scan = steering_matrix(scan_grid_np)
|
||
|
||
for snr_db in SNR_dB_range:
|
||
sync()
|
||
tic = time.perf_counter()
|
||
|
||
sig_p = 10 ** (snr_db / 10.0)
|
||
|
||
# 一次性批量生成 Monte Carlo 数据
|
||
s_t = np.sqrt(sig_p / 2) * (
|
||
torch.randn(Monte_Carlo, 1, L_snapshots, dtype=torch.float32, device=device)
|
||
+ 1j * torch.randn(Monte_Carlo, 1, L_snapshots, dtype=torch.float32, device=device)
|
||
)
|
||
s_c = np.sqrt(INR_linear / 2) * (
|
||
torch.randn(Monte_Carlo, 1, L_snapshots, dtype=torch.float32, device=device)
|
||
+ 1j * torch.randn(Monte_Carlo, 1, L_snapshots, dtype=torch.float32, device=device)
|
||
)
|
||
n_t = np.sqrt(noise_power / 2) * (
|
||
torch.randn(Monte_Carlo, N_ant, L_snapshots, dtype=torch.float32, device=device)
|
||
+ 1j * torch.randn(Monte_Carlo, N_ant, L_snapshots, dtype=torch.float32, device=device)
|
||
)
|
||
X_batch = a_target.unsqueeze(0) * s_t + a_clutter.unsqueeze(0) * s_c + n_t
|
||
R_in_batch = torch.matmul(X_batch, X_batch.mH) / L_snapshots
|
||
|
||
for alpha_dl in alpha_list:
|
||
# --- SINR：使用每个 Monte Carlo 样本求权，最后取平均 ---
|
||
sinr_vals = []
|
||
cond_eff_vals = []
|
||
|
||
for mc in range(Monte_Carlo):
|
||
w_hyb, _, cond_eff = calc_hybrid_weight_loaded(
|
||
F_RF_tsc, R_in_batch[mc], a_target, alpha_dl=alpha_dl
|
||
)
|
||
sinr_vals.append(calc_sinr(w_hyb, a_target, a_clutter, sig_p))
|
||
cond_eff_vals.append(cond_eff)
|
||
|
||
SINR_loaded[alpha_dl].append(float(np.mean(sinr_vals)))
|
||
|
||
# --- RMSE 与 R_bb 条件数：批量 Capon 谱 ---
|
||
P, mean_cond_bb = robust_capon_spectrum_batch(
|
||
F_RF_tsc, X_batch, A_scan, alpha_dl=alpha_dl
|
||
)
|
||
est_idx = torch.argmax(P, dim=1)
|
||
est_angles = torch.as_tensor(scan_grid_np, dtype=torch.float32, device=device)[est_idx]
|
||
rmse = torch.sqrt(torch.mean((est_angles - target_angle) ** 2)).item()
|
||
|
||
RMSE_loaded[alpha_dl].append(rmse)
|
||
COND_loaded[alpha_dl].append(mean_cond_bb)
|
||
|
||
sync()
|
||
print(f"SNR={snr_db:>3} dB 完成，耗时 {time.perf_counter() - tic:.2f} s")
|
||
|
||
# ============================================================
|
||
# 7. 输出 SINR 图
|
||
# ============================================================
|
||
plt.figure(figsize=(10, 6))
|
||
for alpha_dl in alpha_list:
|
||
plt.plot(SNR_dB_range, SINR_loaded[alpha_dl], marker='o', linewidth=2, label=f'α={alpha_dl:.2f}')
|
||
plt.xlabel('输入 SNR (dB)')
|
||
plt.ylabel('输出 SINR (dB)')
|
||
plt.title(f'不同对角加载系数下 TSC-OMP 输出 SINR 对比 ({ARRAY_TYPE})')
|
||
plt.grid(True, linestyle='--', alpha=0.5)
|
||
plt.legend()
|
||
plt.tight_layout()
|
||
if SAVE_FIG:
|
||
plt.savefig(os.path.join(SAVE_DIR, f'SINR_loaded_compare_{ARRAY_TYPE}.png'), dpi=300)
|
||
|
||
# ============================================================
|
||
# 8. RMSE 图
|
||
# ============================================================
|
||
plt.figure(figsize=(10, 6))
|
||
for alpha_dl in alpha_list:
|
||
plt.semilogy(SNR_dB_range, RMSE_loaded[alpha_dl], marker='s', linewidth=2, label=f'α={alpha_dl:.2f}')
|
||
plt.xlabel('输入 SNR (dB)')
|
||
plt.ylabel('测角 RMSE (°)')
|
||
plt.title(f'不同对角加载系数下 TSC-OMP 测角 RMSE 对比 ({ARRAY_TYPE})')
|
||
plt.grid(True, which='both', linestyle='--', alpha=0.5)
|
||
plt.legend()
|
||
plt.tight_layout()
|
||
if SAVE_FIG:
|
||
plt.savefig(os.path.join(SAVE_DIR, f'RMSE_loaded_compare_{ARRAY_TYPE}.png'), dpi=300)
|
||
|
||
# ============================================================
|
||
# 9. 条件数图
|
||
# ============================================================
|
||
plt.figure(figsize=(10, 6))
|
||
for alpha_dl in alpha_list:
|
||
plt.semilogy(SNR_dB_range, COND_loaded[alpha_dl], marker='d', linewidth=2, label=f'α={alpha_dl:.2f}')
|
||
plt.xlabel('输入 SNR (dB)')
|
||
plt.ylabel('低维协方差矩阵条件数')
|
||
plt.title(f'不同对角加载系数下低维协方差矩阵条件数 ({ARRAY_TYPE})')
|
||
plt.grid(True, which='both', linestyle='--', alpha=0.5)
|
||
plt.legend()
|
||
plt.tight_layout()
|
||
if SAVE_FIG:
|
||
plt.savefig(os.path.join(SAVE_DIR, f'COND_loaded_compare_{ARRAY_TYPE}.png'), dpi=300)
|
||
|
||
# ============================================================
|
||
# 10. 低SNR点 α 扫描汇总图
|
||
# ============================================================
|
||
low_snr_pick = 0
|
||
alpha_arr = np.array(alpha_list, dtype=float)
|
||
rmse_pick = np.array([RMSE_loaded[a][np.where(SNR_dB_range == low_snr_pick)[0][0]] for a in alpha_list], dtype=float)
|
||
cond_pick = np.array([COND_loaded[a][np.where(SNR_dB_range == low_snr_pick)[0][0]] for a in alpha_list], dtype=float)
|
||
|
||
fig = plt.figure(figsize=(10, 6))
|
||
ax1 = fig.add_subplot(111)
|
||
ax1.plot(alpha_arr, rmse_pick, marker='o', linewidth=2, color='#0072BD', label='RMSE')
|
||
ax1.set_xlabel('对角加载系数 α')
|
||
ax1.set_ylabel('RMSE (°)', color='#0072BD')
|
||
ax1.tick_params(axis='y', labelcolor='#0072BD')
|
||
ax1.grid(True, linestyle='--', alpha=0.5)
|
||
|
||
ax2 = ax1.twinx()
|
||
ax2.semilogy(alpha_arr, cond_pick, marker='s', linewidth=2, color='#D95319', label='条件数')
|
||
ax2.set_ylabel('条件数', color='#D95319')
|
||
ax2.tick_params(axis='y', labelcolor='#D95319')
|
||
|
||
plt.title(f'低SNR={low_snr_pick} dB 下加载系数对 RMSE 与条件数的影响 ({ARRAY_TYPE})')
|
||
fig.tight_layout()
|
||
if SAVE_FIG:
|
||
plt.savefig(os.path.join(SAVE_DIR, f'alpha_sweep_lowSNR_{ARRAY_TYPE}.png'), dpi=300)
|
||
|
||
# ============================================================
|
||
# 11. 打印总结
|
||
# ============================================================
|
||
print("\n===== 鲁棒性提升实验完成 =====")
|
||
for alpha_dl in alpha_list:
|
||
print(f"α={alpha_dl:.2f} : "
|
||
f"平均SINR={np.mean(SINR_loaded[alpha_dl]):.3f} dB, "
|
||
f"平均RMSE={np.mean(RMSE_loaded[alpha_dl]):.4f}°, "
|
||
f"平均条件数={np.mean(COND_loaded[alpha_dl]):.3e}")
|
||
|
||
print(f"\n结果图片已保存到: {SAVE_DIR}")
|
||
print("即将以窗口形式显示所有图片……")
|
||
plt.show()
|