1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
| """
Physics-Guided Multipath-Aware Attention (PGMA) 实现
核心思想：标准 Self-Attention 将输入视为通用序列，
PGMA 将天线几何、波长相关相位关系、多径可行性约束嵌入注意力计算
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class PhysicsGuidedAttention(nn.Module):
"""
物理引导多径感知注意力机制
与标准 Self-Attention 的关键区别：
1. 天线位置编码：基于物理天线间距和波长
2. 多径可行性掩码：基于物理传播约束
3. 相位关系编码：波长相关的相位差
参考：论文 Section III-A, "Physics-Guided Multipath-Aware Attention"
"""
def __init__(self, d_model: int, n_heads: int,
n_antennas: int, antenna_spacing: float,
wavelength: float, dropout: float = 0.1):
super().__init__()
self.d_model = d_model
self.n_heads = n_heads
self.head_dim = d_model // n_heads
self.n_antennas = n_antennas
self.wavelength = wavelength
self.q_proj = nn.Linear(d_model, d_model)
self.k_proj = nn.Linear(d_model, d_model)
self.v_proj = nn.Linear(d_model, d_model)
self.antenna_pos_encoding = self._create_antenna_encoding(
n_antennas, antenna_spacing, wavelength
)
self.multipath_mask = self._create_multipath_mask(
n_antennas, antenna_spacing, wavelength
)
self.out_proj = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)
def _create_antenna_encoding(self, n_antennas: int,
spacing: float,
wavelength: float) -> torch.Tensor:
"""
创建天线位置编码
基于天线间距与波长的比值，生成相位偏移编码
"""
positions = torch.arange(n_antennas).float()
phase_shift = 2 * math.pi * spacing * positions / wavelength
encoding = torch.zeros(n_antennas, self.d_model)
for i in range(self.d_model // 2):
encoding[:, 2*i] = torch.sin(phase_shift * (i+1))
encoding[:, 2*i+1] = torch.cos(phase_shift * (i+1))
return encoding
def _create_multipath_mask(self, n_antennas: int,
spacing: float,
wavelength: float) -> torch.Tensor:
"""
创建多径可行性掩码
基于物理约束：
- 相邻天线的相位差应 < π（避免空间混叠）
- 远距离天线对的多径贡献衰减
"""
mask = torch.ones(n_antennas, n_antennas)
for i in range(n_antennas):
for j in range(n_antennas):
if i != j:
d = abs(i - j) * spacing
phase = 2 * math.pi * d / wavelength
if phase > math.pi:
mask[i, j] = 0.5
mask[i, j] *= 1.0 / (1.0 + d / wavelength)
return mask
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: CSI 特征, shape=(B, N_antennas, d_model)
Returns:
output: 物理增强注意力输出
"""
B, N, D = x.shape
H = self.n_heads
D_h = self.head_dim
x_phys = x + self.antenna_pos_encoding.to(x.device).unsqueeze(0)
Q = self.q_proj(x_phys).view(B, N, H, D_h).transpose(1, 2)
K = self.k_proj(x_phys).view(B, N, H, D_h).transpose(1, 2)
V = self.v_proj(x).view(B, N, H, D_h).transpose(1, 2)
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(D_h)
mask = self.multipath_mask.to(x.device).unsqueeze(0).unsqueeze(0)
scores = scores * mask
attn = F.softmax(scores, dim=-1)
attn = self.dropout(attn)
out = torch.matmul(attn, V)
out = out.transpose(1, 2).contiguous().view(B, N, D)
return self.out_proj(out)
if __name__ == "__main__":
wavelength = 3e8 / 5e9
spacing = wavelength / 2
pga = PhysicsGuidedAttention(
d_model=64, n_heads=4,
n_antennas=4, antenna_spacing=spacing,
wavelength=wavelength
)
x = torch.randn(2, 4, 64)
out = pga(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {out.shape}")
print(f"Antenna encoding shape: {pga.antenna_pos_encoding.shape}")
print(f"Multipath mask:\n{pga.multipath_mask}")
|
