491 lines
19 KiB
Python
Executable File
491 lines
19 KiB
Python
Executable File
"""
|
|
Shared library for TWC revision-2 experiments (new submission).
|
|
Single-signal uplink model matching the manuscript:
|
|
y = sum_v g_v (e_v ⊙ m_v) + n, g_v = |h_v| e^{jΔφ_v}
|
|
All experiments import from here. Seed fixed = 42.
|
|
"""
|
|
import copy
|
|
import json
|
|
import math
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
DATA = Path(__file__).resolve().parent / "data"
|
|
DATA.mkdir(exist_ok=True)
|
|
|
|
SCENARIOS = {
|
|
"HIGH": {"beta_u": [0.65, 0.65, 0.60, 0.60], "scenes": [0, 0, 0, 0]},
|
|
"LOW": {"beta_u": [0.65, 0.05, 0.05, 0.05], "scenes": [0, 1, 2, 3]},
|
|
"MIX": {"beta_u": [0.65, 0.65, 0.05, 0.05], "scenes": [0, 0, 1, 2]},
|
|
# 8-slot high-correlation scenario for dynamic-user experiment
|
|
"HIGH8": {"beta_u": [0.60] * 8, "scenes": [0] * 8},
|
|
}
|
|
|
|
SNR_GRID = np.arange(0.0, 20.0 + 1e-6, 2.0)
|
|
|
|
|
|
def beta_matrix(scen):
|
|
"""True relevance matrix beta_{u,v} = beta_u beta_v [same scene], else 0."""
|
|
b = np.asarray(scen["beta_u"], dtype=np.float64)
|
|
sc = np.asarray(scen["scenes"])
|
|
U = len(b)
|
|
B = np.zeros((U, U))
|
|
for u in range(U):
|
|
for v in range(U):
|
|
B[u, v] = 1.0 if u == v else (b[u] * b[v] if sc[u] == sc[v] else 0.0)
|
|
return B
|
|
|
|
|
|
def gen_embeddings(n, d, U, rng, scen):
|
|
"""Unit-norm ground-truth embeddings (n, U, d): e_u = sqrt(1-b^2) p + b s."""
|
|
blend, scenes = scen["beta_u"], scen["scenes"]
|
|
svecs = {}
|
|
for sc in sorted(set(scenes)):
|
|
v = rng.standard_normal(d)
|
|
svecs[sc] = v / (np.linalg.norm(v) + 1e-8)
|
|
embs = []
|
|
for u in range(U):
|
|
b = blend[u]
|
|
p = rng.standard_normal((n, d))
|
|
p /= np.linalg.norm(p, axis=-1, keepdims=True) + 1e-8
|
|
e = np.sqrt(max(1 - b * b, 0.0)) * p + b * svecs[scenes[u]][None, :]
|
|
e /= np.linalg.norm(e, axis=-1, keepdims=True) + 1e-8
|
|
embs.append(e)
|
|
return torch.from_numpy(np.stack(embs, 1)).float()
|
|
|
|
|
|
class ViewNets(nn.Module):
|
|
"""Fixed random per-user nonlinear view functions g_u([kappa*s; p_u])."""
|
|
|
|
def __init__(self, d, U, seed=7):
|
|
super().__init__()
|
|
g = torch.Generator().manual_seed(seed)
|
|
self.nets = nn.ModuleList()
|
|
for _ in range(U):
|
|
l1 = nn.Linear(2 * d, 2 * d)
|
|
l2 = nn.Linear(2 * d, d)
|
|
for l in (l1, l2):
|
|
nn.init.normal_(l.weight, std=(2.0 / l.in_features) ** 0.5, generator=g)
|
|
nn.init.zeros_(l.bias)
|
|
self.nets.append(nn.Sequential(l1, nn.Tanh(), l2))
|
|
for p in self.parameters():
|
|
p.requires_grad_(False)
|
|
|
|
@torch.no_grad()
|
|
def gen(self, n, d, U, rng, kappa, shared_scene=True):
|
|
if shared_scene:
|
|
s = rng.standard_normal((n, d)) / math.sqrt(d)
|
|
s = np.repeat(s[:, None, :], U, axis=1)
|
|
else:
|
|
s = rng.standard_normal((n, U, d)) / math.sqrt(d)
|
|
p = rng.standard_normal((n, U, d)) / math.sqrt(d)
|
|
s = torch.from_numpy(s).float().to(DEVICE)
|
|
p = torch.from_numpy(p).float().to(DEVICE)
|
|
outs = []
|
|
for u in range(U):
|
|
x = torch.cat([kappa * s[:, u], p[:, u]], -1)
|
|
outs.append(F.normalize(self.nets[u](x), dim=-1))
|
|
return torch.stack(outs, 1) # (n, U, d)
|
|
|
|
|
|
def block_masks(U, d, device=DEVICE):
|
|
dpu = d // U
|
|
m = torch.zeros(U, d, device=device)
|
|
for u in range(U):
|
|
m[u, u * dpu:(u + 1) * dpu] = 1.0
|
|
return m
|
|
|
|
|
|
def channel(E, snr_db, *, phase_sigma_deg=0.0, fading="rayleigh", rician_K_dB=None,
|
|
nakagami_m=None, offsets=None, h_err_sigma=0.0, masks=None, use_masks=True):
|
|
"""Single-signal uplink. Returns dict with yI, yQ, h (true magnitude), h_hat.
|
|
|
|
E : (n, U, d) ground-truth embeddings on DEVICE.
|
|
offsets : (n, U) integer per-user timing offsets (symbols), or None.
|
|
Noise convention: per-rail noise std = sqrt(mean|y_tx|^2 / snr_lin); with
|
|
phase_sigma = 0 the model reduces exactly to the real-valued model.
|
|
"""
|
|
n, U, d = E.shape
|
|
if masks is None:
|
|
masks = block_masks(U, d, E.device)
|
|
X = E * masks[None] if use_masks else E.clone()
|
|
|
|
if offsets is not None:
|
|
Xs = torch.zeros_like(X)
|
|
offs = offsets
|
|
for u in range(U):
|
|
for o in offs[:, u].unique():
|
|
o = int(o.item())
|
|
idx = offs[:, u] == o
|
|
if o == 0:
|
|
Xs[idx, u] = X[idx, u]
|
|
else:
|
|
Xs[idx, u, o:] = X[idx, u, :d - o]
|
|
X = Xs
|
|
|
|
if fading == "rayleigh":
|
|
hI = torch.randn(n, U, device=E.device) * (0.5 ** 0.5)
|
|
hQ = torch.randn(n, U, device=E.device) * (0.5 ** 0.5)
|
|
hmag = (hI ** 2 + hQ ** 2).sqrt()
|
|
elif fading == "rician":
|
|
K = 10 ** (rician_K_dB / 10.0)
|
|
mu = math.sqrt(K / (K + 1))
|
|
sig = math.sqrt(1.0 / (2 * (K + 1)))
|
|
hI = mu + torch.randn(n, U, device=E.device) * sig
|
|
hQ = torch.randn(n, U, device=E.device) * sig
|
|
hmag = (hI ** 2 + hQ ** 2).sqrt()
|
|
elif fading == "nakagami":
|
|
m = nakagami_m
|
|
gam = torch.distributions.Gamma(m, m).sample((n, U)).to(E.device)
|
|
hmag = gam.sqrt()
|
|
else:
|
|
raise ValueError(fading)
|
|
|
|
dphi = torch.randn(n, U, device=E.device) * math.radians(phase_sigma_deg)
|
|
gI = hmag * torch.cos(dphi)
|
|
gQ = hmag * torch.sin(dphi)
|
|
|
|
yI = (gI[:, :, None] * X).sum(1)
|
|
yQ = (gQ[:, :, None] * X).sum(1)
|
|
P = (yI ** 2 + yQ ** 2).mean()
|
|
nstd = (P / (10 ** (snr_db / 10.0))).sqrt()
|
|
yI = yI + torch.randn_like(yI) * nstd
|
|
yQ = yQ + torch.randn_like(yQ) * nstd
|
|
h_hat = hmag * (1 + torch.randn_like(hmag) * h_err_sigma) if h_err_sigma > 0 else hmag
|
|
return {"yI": yI, "yQ": yQ, "h": hmag, "h_hat": h_hat, "nvar": float(nstd ** 2)}
|
|
|
|
|
|
# ---------------------------------------------------------------- decoders --
|
|
class UWCA(nn.Module):
|
|
"""User-wise cross-attention decoder on the single superimposed signal.
|
|
|
|
iq=True : keys/values read the stacked [I; Q] rails (2d input).
|
|
forward(yI, yQ, active) with active (n, U) bool or None.
|
|
"""
|
|
|
|
def __init__(self, d, U, H=4, iq=False):
|
|
super().__init__()
|
|
assert d % H == 0
|
|
self.d, self.U, self.H, self.dk = d, U, H, d // H
|
|
self.iq = iq
|
|
din = 2 * d if iq else d
|
|
self.q_vectors = nn.Parameter(torch.randn(U, d) * d ** -0.5)
|
|
self.eta = nn.Parameter(torch.ones(1))
|
|
self.W_K = nn.Linear(din, d, bias=False)
|
|
self.W_V = nn.Linear(din, d, bias=False)
|
|
self.W_O = nn.Linear(d, d, bias=False)
|
|
self.norm = nn.LayerNorm(d)
|
|
init_logits = torch.full((U, d), -3.0)
|
|
dpu = d // U
|
|
for u in range(U):
|
|
init_logits[u, u * dpu:(u + 1) * dpu] = 3.0
|
|
self.mask_logits = nn.Parameter(init_logits)
|
|
|
|
def soft_masks(self):
|
|
return torch.sigmoid(self.mask_logits)
|
|
|
|
def forward(self, yI, yQ=None, active=None, topk_mask=None, return_alpha=False):
|
|
# yI: (n, d) shared signal, or (n, U, d) per-candidate aligned copies
|
|
U, H, dk = self.U, self.H, self.dk
|
|
m = self.soft_masks() # (U, d)
|
|
if yI.dim() == 2:
|
|
n, d = yI.shape
|
|
yIc = yI[:, None, :].expand(-1, U, -1)
|
|
else:
|
|
n, _, d = yI.shape
|
|
yIc = yI
|
|
R = yIc * m[None] # (n, U, d)
|
|
if self.iq:
|
|
yQc = yQ[:, None, :].expand(-1, U, -1) if yQ.dim() == 2 else yQ
|
|
RQ = yQc * m[None]
|
|
Rin = torch.cat([R, RQ], -1) # (n, U, 2d)
|
|
else:
|
|
Rin = R
|
|
K = self.W_K(Rin).view(n, U, H, dk) # (n, Uk, H, dk)
|
|
V = self.W_V(Rin).view(n, U, H, dk)
|
|
Q = self.q_vectors.view(U, H, dk) # (Uq, H, dk)
|
|
scores = torch.einsum("qhk,nihk->nqhi", Q, K) * self.eta / dk ** 0.5
|
|
if active is not None: # (n, U) bool
|
|
scores = scores.masked_fill(~active[:, None, None, :], -1e9)
|
|
if topk_mask is not None: # (U, U) bool keep
|
|
scores = scores.masked_fill(~topk_mask[None, :, None, :], -1e9)
|
|
alpha = F.softmax(scores, dim=-1) # (n, Uq, H, Uk)
|
|
ctx = torch.einsum("nqhi,nihk->nqhk", alpha, V).reshape(n, U, d)
|
|
own = yIc * m[None]
|
|
out = F.normalize(self.norm(self.W_O(ctx) + own), dim=-1)
|
|
if return_alpha:
|
|
return out, alpha.mean(dim=(0, 2))
|
|
return out
|
|
|
|
|
|
class Encoder(nn.Module):
|
|
"""Trainable semantic encoder for the end-to-end experiment."""
|
|
|
|
def __init__(self, d):
|
|
super().__init__()
|
|
self.net = nn.Sequential(nn.Linear(d, 2 * d), nn.LayerNorm(2 * d),
|
|
nn.GELU(), nn.Linear(2 * d, d))
|
|
|
|
def forward(self, x):
|
|
return F.normalize(self.net(x), dim=-1)
|
|
|
|
|
|
def ofdma_decode(yI, masks):
|
|
return F.normalize(yI[:, None, :] * masks[None], dim=-1)
|
|
|
|
|
|
def tdma_proj_decode(E, snr_db, rng_t):
|
|
"""Orthogonal scheme with an arbitrary (random orthonormal) d/U-dim
|
|
projection per user instead of coordinate masks: z_u = h_u P_u e_u + n."""
|
|
n, U, d = E.shape
|
|
dpu = d // U
|
|
Q, _ = torch.linalg.qr(torch.randn(d, d, generator=rng_t).to(E.device))
|
|
outs = []
|
|
hmag = (torch.randn(n, U, device=E.device) ** 2 +
|
|
torch.randn(n, U, device=E.device) ** 2).sqrt() * 0.5 ** 0.5
|
|
snr_lin = 10 ** (snr_db / 10.0)
|
|
for u in range(U):
|
|
P = Q[u * dpu:(u + 1) * dpu] # (dpu, d)
|
|
z = hmag[:, u:u + 1] * (E[:, u] @ P.T) # (n, dpu)
|
|
nstd = (z.pow(2).mean() / snr_lin).sqrt()
|
|
z = z + torch.randn_like(z) * nstd
|
|
outs.append(F.normalize(z @ P, dim=-1))
|
|
return torch.stack(outs, 1)
|
|
|
|
|
|
def lmmse_decode(yI, h, nvar, B, masks, genie=True):
|
|
"""Closed-form linear MMSE on the block model. C_vv = I/d, C_uv = B_uv I/d.
|
|
Block-diagonal C_y => per-block Wiener weights w_uv = h_v B_uv/d / (h_v^2/d + nvar).
|
|
genie=False zeroes the cross terms (correlation-blind)."""
|
|
n, d = yI.shape
|
|
U = masks.shape[0]
|
|
Bm = torch.as_tensor(B, dtype=torch.float32, device=yI.device)
|
|
if not genie:
|
|
Bm = torch.eye(U, device=yI.device)
|
|
yb = yI[:, None, :] * masks[None] # (n, Uv, d) block pieces
|
|
w = (h[:, None, :] * Bm[None] / d) / (h[:, None, :] ** 2 / d + nvar) # (n,Uu,Uv)
|
|
est = torch.einsum("nuv,nvd->nud", w, yb)
|
|
return F.normalize(est, dim=-1)
|
|
|
|
|
|
def noma_sic_decode(E, snr_db):
|
|
"""Full-band power-domain NOMA with SIC (no masks)."""
|
|
n, U, d = E.shape
|
|
pa = torch.tensor([0.40, 0.30, 0.20, 0.10], device=E.device)[:U]
|
|
pa = pa / pa.sum()
|
|
h = (torch.randn(n, U, 1, device=E.device) ** 2 +
|
|
torch.randn(n, U, 1, device=E.device) ** 2).sqrt() * 0.5 ** 0.5
|
|
y = (E * pa.sqrt()[None, :, None] * h).sum(1)
|
|
nstd = (y.pow(2).mean() / 10 ** (snr_db / 10.0)).sqrt()
|
|
y = y + torch.randn(n, d, device=E.device) * nstd
|
|
order = torch.argsort(pa, descending=True)
|
|
res = y.clone()
|
|
out = torch.zeros_like(E)
|
|
for ui in order:
|
|
u = int(ui.item())
|
|
eh = F.normalize(res / (h[:, u] + 1e-8), dim=-1)
|
|
out[:, u] = eh
|
|
res = res - h[:, u] * pa[u].sqrt() * eh
|
|
return out
|
|
|
|
|
|
# ------------------------------------------------------------------ losses --
|
|
def semantic_loss(Ehat, E, lam=0.1, active=None):
|
|
cos = (Ehat * E).sum(-1)
|
|
if active is not None:
|
|
distortion = ((1 - cos) * active).sum() / active.sum()
|
|
else:
|
|
distortion = (1 - cos).mean()
|
|
U = Ehat.shape[1]
|
|
emb = Ehat.mean(0)
|
|
ec = emb - emb.mean(1, keepdim=True)
|
|
en = F.normalize(ec, dim=1)
|
|
C = en @ en.T
|
|
off = C[~torch.eye(U, dtype=torch.bool, device=Ehat.device)].abs().mean()
|
|
return distortion + lam * off
|
|
|
|
|
|
def ser(Ehat, E, tau=0.45, active=None):
|
|
bad = ((Ehat * E).sum(-1) < tau).float()
|
|
if active is not None:
|
|
return float((bad * active).sum() / active.sum())
|
|
return float(bad.mean())
|
|
|
|
|
|
def mean_cos(Ehat, E, active=None):
|
|
c = (Ehat * E).sum(-1)
|
|
if active is not None:
|
|
return float((c * active).sum() / active.sum())
|
|
return float(c.mean())
|
|
|
|
|
|
def sample_corr(A, Bt):
|
|
"""Mean per-sample Pearson correlation across the d dims of two (n,d) tensors."""
|
|
Ac = A - A.mean(-1, keepdim=True)
|
|
Bc = Bt - Bt.mean(-1, keepdim=True)
|
|
num = (Ac * Bc).sum(-1)
|
|
den = Ac.norm(dim=-1) * Bc.norm(dim=-1) + 1e-9
|
|
return float((num / den).mean())
|
|
|
|
|
|
# ----------------------------------------------------------------- training --
|
|
def train_multitask(model, gen_fn, tasks, epochs=300, batch=64, lam=0.1,
|
|
outer_lr=1e-3, mask_lr_mult=100.0, log_every=50, tag="",
|
|
encoder=None, extra_loss=None, log_state=None):
|
|
"""Multi-SNR/-condition aggregated training (the manuscript's outer objective
|
|
without inner adaptation). gen_fn(batch) -> E ground truth on DEVICE.
|
|
tasks: list of dicts of channel kwargs incl. 'snr_db'."""
|
|
params = []
|
|
mask_p = [p for nm, p in model.named_parameters() if "mask_logits" in nm]
|
|
other = [p for nm, p in model.named_parameters() if "mask_logits" not in nm]
|
|
params = [{"params": other, "lr": outer_lr},
|
|
{"params": mask_p, "lr": outer_lr * mask_lr_mult}]
|
|
if encoder is not None:
|
|
params.append({"params": encoder.parameters(), "lr": outer_lr})
|
|
opt = torch.optim.Adam(params)
|
|
hist = []
|
|
for ep in range(1, epochs + 1):
|
|
loss_acc = 0.0
|
|
opt.zero_grad()
|
|
for t in tasks:
|
|
E = gen_fn(batch)
|
|
if encoder is not None:
|
|
n, U, d = E.shape
|
|
E = encoder(E.reshape(-1, d)).reshape(n, U, d)
|
|
ch = channel(E, **t)
|
|
Eh = model(ch["yI"], ch["yQ"])
|
|
L = semantic_loss(Eh, E, lam)
|
|
if extra_loss is not None:
|
|
L = L + extra_loss(E)
|
|
loss_acc += L
|
|
(loss_acc / len(tasks)).backward()
|
|
nn.utils.clip_grad_norm_(model.parameters(), 5.0)
|
|
opt.step()
|
|
hist.append(float(loss_acc) / len(tasks))
|
|
if log_state is not None:
|
|
gn = sum(float(p.grad.norm()) ** 2 for p in model.parameters()
|
|
if p.grad is not None) ** 0.5
|
|
log_state.append({"ep": ep, "eta": float(model.eta), "gnorm": gn})
|
|
if ep % log_every == 0:
|
|
print(f" [{tag}] ep {ep}/{epochs} loss={hist[-1]:.4f}", flush=True)
|
|
return hist
|
|
|
|
|
|
def fomaml_train(model, gen_fn, tasks, epochs=200, batch=64, lam=0.1,
|
|
inner_lr=0.01, inner_steps=5, outer_lr=1e-3,
|
|
tasks_per_step=8, log_every=25, tag="fomaml", log_state=None,
|
|
rng=None):
|
|
"""Proper first-order MAML: inner SGD on support, outer update from query
|
|
gradients evaluated at the adapted parameters."""
|
|
opt = torch.optim.Adam(model.parameters(), lr=outer_lr)
|
|
rng = rng or np.random.default_rng(0)
|
|
hist = []
|
|
names = [nm for nm, _ in model.named_parameters()]
|
|
for ep in range(1, epochs + 1):
|
|
idx = rng.choice(len(tasks), size=min(tasks_per_step, len(tasks)),
|
|
replace=False)
|
|
grads = {nm: torch.zeros_like(p) for nm, p in model.named_parameters()}
|
|
qloss_acc = 0.0
|
|
for ti in idx:
|
|
t = tasks[ti]
|
|
adapted = copy.deepcopy(model)
|
|
iopt = torch.optim.SGD(adapted.parameters(), lr=inner_lr)
|
|
for _ in range(inner_steps):
|
|
E = gen_fn(batch)
|
|
ch = channel(E, **t)
|
|
L = semantic_loss(adapted(ch["yI"], ch["yQ"]), E, lam)
|
|
iopt.zero_grad(); L.backward(); iopt.step()
|
|
E = gen_fn(batch)
|
|
ch = channel(E, **t)
|
|
qL = semantic_loss(adapted(ch["yI"], ch["yQ"]), E, lam)
|
|
adapted.zero_grad(); qL.backward()
|
|
for nm, p in adapted.named_parameters():
|
|
if p.grad is not None:
|
|
grads[nm] += p.grad
|
|
qloss_acc += float(qL)
|
|
opt.zero_grad()
|
|
for nm, p in model.named_parameters():
|
|
p.grad = grads[nm] / len(idx)
|
|
nn.utils.clip_grad_norm_(model.parameters(), 5.0)
|
|
opt.step()
|
|
hist.append(qloss_acc / len(idx))
|
|
if log_state is not None:
|
|
gn = sum(float(g.norm()) ** 2 for g in grads.values()) ** 0.5 / len(idx)
|
|
log_state.append({"ep": ep, "eta": float(model.eta), "gnorm": gn})
|
|
if ep % log_every == 0:
|
|
print(f" [{tag}] ep {ep}/{epochs} qloss={hist[-1]:.4f} "
|
|
f"eta={float(model.eta):.3f}", flush=True)
|
|
return hist
|
|
|
|
|
|
def adapt(model, gen_fn, task, steps=5, inner_lr=0.01, batch=64, lam=0.1):
|
|
adapted = copy.deepcopy(model)
|
|
iopt = torch.optim.SGD(adapted.parameters(), lr=inner_lr)
|
|
for _ in range(steps):
|
|
E = gen_fn(batch)
|
|
ch = channel(E, **task)
|
|
L = semantic_loss(adapted(ch["yI"], ch["yQ"]), E, lam)
|
|
iopt.zero_grad(); L.backward(); iopt.step()
|
|
return adapted
|
|
|
|
|
|
# --------------------------------------------------------------- evaluation --
|
|
@torch.no_grad()
|
|
def eval_scheme(scheme, gen_fn, task, n_mc=200, batch=64, tau=0.45, model=None,
|
|
B=None, masks=None, encoder=None, rng_t=None, active_fn=None,
|
|
topk_mask=None):
|
|
"""Returns (ser, cos) for one task/channel config."""
|
|
s_acc = c_acc = 0.0
|
|
for _ in range(n_mc):
|
|
E = gen_fn(batch)
|
|
if encoder is not None:
|
|
n, U, d = E.shape
|
|
E = encoder(E.reshape(-1, d)).reshape(n, U, d)
|
|
act = active_fn(E.shape[0]) if active_fn is not None else None
|
|
if scheme == "uwca":
|
|
ch = channel(E, **task)
|
|
Eh = model(ch["yI"], ch["yQ"], active=act, topk_mask=topk_mask)
|
|
elif scheme == "ofdma":
|
|
ch = channel(E, **task)
|
|
m = masks if masks is not None else block_masks(E.shape[1], E.shape[2],
|
|
E.device)
|
|
Eh = ofdma_decode(ch["yI"], m)
|
|
elif scheme == "sfdma":
|
|
ch = channel(E, **task)
|
|
m = masks
|
|
Eh = torch.stack([F.normalize((ch["yI"] * m[u]) /
|
|
(ch["h_hat"][:, u:u + 1] + 1e-8), dim=-1)
|
|
for u in range(m.shape[0])], 1)
|
|
elif scheme == "noma":
|
|
Eh = noma_sic_decode(E, task["snr_db"])
|
|
elif scheme in ("lmmse_genie", "lmmse_blind"):
|
|
ch = channel(E, **task)
|
|
Eh = lmmse_decode(ch["yI"], ch["h_hat"], ch["nvar"], B, masks,
|
|
genie=(scheme == "lmmse_genie"))
|
|
elif scheme == "tdma_proj":
|
|
Eh = tdma_proj_decode(E, task["snr_db"], rng_t)
|
|
else:
|
|
raise ValueError(scheme)
|
|
s_acc += ser(Eh, E, tau, act)
|
|
c_acc += mean_cos(Eh, E, act)
|
|
return s_acc / n_mc, c_acc / n_mc
|
|
|
|
|
|
def save_json(name, obj):
|
|
p = DATA / name
|
|
with open(p, "w") as f:
|
|
json.dump(obj, f, indent=1)
|
|
print(f"saved -> {p}", flush=True)
|
|
|
|
|
|
def set_seed(seed=42):
|
|
torch.manual_seed(seed)
|
|
np.random.seed(seed)
|
|
return np.random.default_rng(seed)
|