Reproducibility package: UWCA semantic multiple access (TWC submission)

rev2/e1_fair_baselines.py

@@ -0,0 +1,54 @@
+"""E1 — Degrees-of-freedom fairness (R1.10, R3.7).
+
+Adds full-dimensional receivers on the SAME received signal:
+  - lmmse_blind : optimal linear receiver with cross-user correlation set to 0
+                  (proves the d/U ceiling is fundamental to correlation-blind
+                  processing, not an artifact of the OFDMA baseline)
+  - lmmse_genie : optimal linear receiver given the true relevance matrix
+                  (genie-aided upper reference; UWCA should approach it)
+  - tdma_proj   : orthogonal scheme with an arbitrary orthonormal projection
+                  (proves any orthogonal partition is statistically identical
+                  to coordinate masking for isotropic embeddings)
+Also trains the UWCA decoder per scenario under the single-signal model and
+saves checkpoints for reuse (E6).
+"""
+import numpy as np
+import torch
+
+import lib
+from lib import (SCENARIOS, SNR_GRID, UWCA, DEVICE, beta_matrix, block_masks,
+                 eval_scheme, gen_embeddings, save_json, set_seed, train_multitask)
+
+rng = set_seed(42)
+d, U, H = 64, 4, 4
+masks = block_masks(U, d)
+tasks = [{"snr_db": float(s)} for s in np.arange(0, 21, 4)]
+
+out = {"snr": SNR_GRID.tolist(), "scenarios": {}}
+for scen_name in ["HIGH", "LOW", "MIX"]:
+    scen = SCENARIOS[scen_name]
+    B = beta_matrix(scen)
+
+    def gen(n, scen=scen):
+        return gen_embeddings(n, d, U, rng, scen).to(DEVICE)
+
+    model = UWCA(d, U, H).to(DEVICE)
+    train_multitask(model, gen, tasks, epochs=300, tag=f"E1-{scen_name}")
+    torch.save(model.state_dict(), lib.DATA / f"e1_uwca_{scen_name}.pt")
+
+    res = {}
+    rng_t = torch.Generator().manual_seed(1)
+    for scheme in ["uwca", "ofdma", "sfdma", "noma", "lmmse_blind",
+                   "lmmse_genie", "tdma_proj"]:
+        sers, coss = [], []
+        for snr in SNR_GRID:
+            t = {"snr_db": float(snr)}
+            s, c = eval_scheme(scheme, gen, t, n_mc=200, model=model, B=B,
+                               masks=masks, rng_t=rng_t)
+            sers.append(s); coss.append(c)
+        res[scheme] = {"ser": sers, "cos": coss}
+        print(f"[E1-{scen_name}] {scheme}: SER@10dB={sers[5]:.3f} "
+              f"cos@10dB={coss[5]:.3f}", flush=True)
+    out["scenarios"][scen_name] = res
+
+save_json("e1_fair_baselines.json", out)