148 lines
5.9 KiB
Python
148 lines
5.9 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.animation as animation
|
|
import time
|
|
|
|
# ==========================================
|
|
# 0. 全局设置
|
|
# ==========================================
|
|
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
|
|
plt.rcParams['axes.unicode_minus'] = False
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"========== 启动 AI 动态追踪生成器 ==========\n使用设备: {device}")
|
|
|
|
N_ant, N_RF, d_lambda = 24, 4, 0.5
|
|
noise_power = 1.0
|
|
INR_linear = 10**(30/10) # 30dB 强杂波
|
|
|
|
def gen_a_batch(theta_deg_tensor):
|
|
theta_rad = torch.deg2rad(theta_deg_tensor)
|
|
n_idx = torch.arange(-(N_ant-1)/2, (N_ant-1)/2 + 0.1, device=device).unsqueeze(0)
|
|
phases = 2 * torch.pi * d_lambda * n_idx * torch.sin(theta_rad.unsqueeze(1))
|
|
return torch.exp(1j * phases).unsqueeze(2)
|
|
|
|
# ==========================================
|
|
# 1. 广域追踪网络 (彻底解决相位死锁)
|
|
# ==========================================
|
|
class DynamicHBFNet(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.net = nn.Sequential(
|
|
nn.Linear(48, 512), nn.GELU(),
|
|
nn.Linear(512, 512), nn.GELU(),
|
|
nn.Linear(512, N_ant * N_RF * 2)
|
|
)
|
|
|
|
def forward(self, w_opt):
|
|
w_flat = w_opt.squeeze(2)
|
|
x = torch.cat([w_flat.real, w_flat.imag], dim=1)
|
|
out = self.net(x).view(-1, N_ant, N_RF, 2)
|
|
complex_out = out[..., 0] + 1j * out[..., 1]
|
|
# 严格恒模约束
|
|
F_RF = complex_out / (torch.abs(complex_out) + 1e-8) / np.sqrt(N_ant)
|
|
return F_RF
|
|
|
|
# ==========================================
|
|
# 2. 广域动态特训
|
|
# ==========================================
|
|
model = DynamicHBFNet().to(device)
|
|
optimizer = optim.Adam(model.parameters(), lr=0.002)
|
|
|
|
Batch_Size, Epochs = 500, 10000
|
|
|
|
print("正在对 AI 进行大范围目标与杂波抗干扰特训...")
|
|
for epoch in range(Epochs):
|
|
model.train()
|
|
optimizer.zero_grad()
|
|
|
|
# 【动态大扫掠特训】：目标 0~40°，杂波 -20~-3°
|
|
theta_t = torch.rand(Batch_Size, device=device) * 40.0
|
|
theta_c = -20.0 + torch.rand(Batch_Size, device=device) * 17.0
|
|
|
|
a_t, a_c = gen_a_batch(theta_t), gen_a_batch(theta_c)
|
|
|
|
R_interf = noise_power * torch.eye(N_ant, device=device).unsqueeze(0) + INR_linear * (a_c @ a_c.mH)
|
|
w_opt = torch.linalg.pinv(R_interf) @ a_t
|
|
w_opt = w_opt / torch.norm(w_opt, dim=1, keepdim=True)
|
|
|
|
F_RF = model(w_opt)
|
|
F_BB = torch.linalg.pinv(F_RF.mH @ F_RF + 1e-3*torch.eye(N_RF, device=device).unsqueeze(0)) @ (F_RF.mH @ w_opt)
|
|
w_ai = F_RF @ F_BB
|
|
w_ai = w_ai / torch.norm(w_ai, dim=1, keepdim=True)
|
|
|
|
sim = torch.abs(torch.sum(w_opt.conj() * w_ai, dim=(1,2)))**2
|
|
loss_sim = torch.mean(1.0 - sim)
|
|
loss_null = torch.mean(torch.abs(torch.sum(a_c.conj() * w_ai, dim=(1,2)))**2)
|
|
|
|
loss = loss_sim + 150.0 * loss_null
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
print("特训完成！正在生成实时追踪 GIF 动图...")
|
|
|
|
# ==========================================
|
|
# 3. 生成精美动态 GIF
|
|
# ==========================================
|
|
model.eval()
|
|
theta_scan = torch.arange(-60, 60.1, 0.5, device=device) # 限制扫描角度让图像更聚焦
|
|
A_scan = gen_a_batch(theta_scan).squeeze(2).T
|
|
|
|
fig, ax = plt.subplots(figsize=(10, 6))
|
|
ax.set_xlim(-60, 60)
|
|
ax.set_ylim(-60, 0)
|
|
ax.set_xlabel('角度 θ (°)', fontsize=12, fontweight='bold')
|
|
ax.set_ylabel('归一化增益 (dB)', fontsize=12, fontweight='bold')
|
|
ax.grid(True, ls='--')
|
|
|
|
line_opt, = ax.plot([], [], 'k-', lw=2, alpha=0.5, label='纯数字最优下限')
|
|
line_ai, = ax.plot([], [], color='#D95319', ls='--', lw=2.5, label='AI 瞬时感知混合波束')
|
|
vline_t = ax.axvline(0, color='g', ls='-.', lw=2, label='机动目标轨迹')
|
|
vline_c = ax.axvline(0, color='r', ls='-.', lw=2, label='强杂波轨迹')
|
|
|
|
title_text = ax.set_title('', fontsize=14, fontweight='bold')
|
|
ax.legend(loc='lower center', fontsize=10)
|
|
|
|
frames_count = 60 # 60帧平滑过渡
|
|
|
|
def update(frame):
|
|
# 目标从 0 匀速扫到 40
|
|
t_ang = 0.0 + frame * (40.0 / (frames_count - 1))
|
|
# 杂波从 -3 匀速扫到 -20
|
|
c_ang = -3.0 - frame * (17.0 / (frames_count - 1))
|
|
|
|
a_t = gen_a_batch(torch.tensor([t_ang], device=device))
|
|
a_c = gen_a_batch(torch.tensor([c_ang], device=device))
|
|
R_interf = noise_power * torch.eye(N_ant, device=device).unsqueeze(0) + INR_linear * (a_c @ a_c.mH)
|
|
|
|
w_opt = torch.linalg.pinv(R_interf) @ a_t
|
|
w_opt /= torch.norm(w_opt, dim=1, keepdim=True)
|
|
|
|
# AI 仅凭一次前向传播，瞬间生成波束！
|
|
with torch.no_grad():
|
|
F_RF = model(w_opt)
|
|
F_BB = torch.linalg.pinv(F_RF.mH @ F_RF + 1e-3*torch.eye(N_RF, device=device).unsqueeze(0)) @ (F_RF.mH @ w_opt)
|
|
w_ai = F_RF @ F_BB
|
|
w_ai /= torch.norm(w_ai, dim=1, keepdim=True)
|
|
|
|
gain_opt = torch.abs(torch.sum(w_opt.conj() * A_scan.unsqueeze(0), dim=1))**2
|
|
bp_opt = 10 * torch.log10(gain_opt).cpu().numpy().flatten()
|
|
bp_opt -= np.max(bp_opt)
|
|
|
|
gain_ai = torch.abs(torch.sum(w_ai.conj() * A_scan.unsqueeze(0), dim=1))**2
|
|
bp_ai = 10 * torch.log10(gain_ai).cpu().numpy().flatten()
|
|
bp_ai -= np.max(bp_ai)
|
|
|
|
line_opt.set_data(theta_scan.cpu().numpy(), bp_opt)
|
|
line_ai.set_data(theta_scan.cpu().numpy(), bp_ai)
|
|
vline_t.set_xdata([t_ang, t_ang])
|
|
vline_c.set_xdata([c_ang, c_ang])
|
|
|
|
title_text.set_text(f'AI 端到端波束实时追踪: 目标 {t_ang:.1f}° | 强杂波 {c_ang:.1f}°')
|
|
return line_opt, line_ai, vline_t, vline_c, title_text
|
|
|
|
# 渲染动图
|
|
anim = animation.FuncAnimation(fig, update, frames=frames_count, interval=100, blit=False)
|
|
anim.save('E:\\Coding\\graduation design\\MIMO_Beamforming_Simulation_Platform\\final\\python\\Deeplearning\\AI_Beam_Tracking.gif', writer='pillow', fps=10) |