Restructure package: descriptive study documentation and clean layout

experiments/e3_dynamic_users.py

@@ -0,0 +1,99 @@
+"""E3 — Dynamic user arrivals and departures.
+
+U_max = 8 mask slots, d = 64.  Activity-aware meta-training samples a random
+active subset each batch; at inference the attention softmax is restricted to
+the active set announced by the scheduler (no retraining).
+Compared against: (i) a model trained with all 8 users always active
+(mismatch), and (ii) oracle models retrained for each fixed active count.
+"""
+import numpy as np
+import torch
+
+import lib
+from lib import (SCENARIOS, UWCA, DEVICE, block_masks, channel, gen_embeddings,
+                 mean_cos, save_json, semantic_loss, ser, set_seed)
+
+rng = set_seed(42)
+d, Umax, H = 64, 8, 4
+masks = block_masks(Umax, d)
+scen = SCENARIOS["HIGH8"]
+snrs = [0.0, 4.0, 8.0, 12.0, 16.0, 20.0]
+act_rng = np.random.default_rng(11)
+
+
+def gen(n):
+    return gen_embeddings(n, d, Umax, rng, scen).to(DEVICE)
+
+
+def sample_active(n, k=None):
+    """(n, Umax) bool with k active users (random subset per sample)."""
+    A = np.zeros((n, Umax), dtype=bool)
+    for i in range(n):
+        kk = k if k is not None else int(act_rng.integers(2, Umax + 1))
+        A[i, act_rng.choice(Umax, size=kk, replace=False)] = True
+    return torch.from_numpy(A).to(DEVICE)
+
+
+def train(model, epochs=250, k=None, tag=""):
+    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]
+    opt = torch.optim.Adam([{"params": other, "lr": 1e-3},
+                            {"params": mask_p, "lr": 0.1}])
+    for ep in range(1, epochs + 1):
+        opt.zero_grad()
+        loss = 0.0
+        for s in snrs:
+            E = gen(64)
+            act = sample_active(64, k)
+            Ez = E * act[:, :, None]              # inactive users transmit nothing
+            ch = channel(Ez, snr_db=s)
+            Eh = model(ch["yI"], ch["yQ"], active=act)
+            loss = loss + semantic_loss(Eh, E, 0.1, active=act.float())
+        (loss / len(snrs)).backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
+        opt.step()
+        if ep % 50 == 0:
+            print(f"  [E3-{tag}] ep {ep}/{epochs} loss={float(loss)/len(snrs):.4f}",
+                  flush=True)
+
+
+@torch.no_grad()
+def evaluate(model, k, snr, n_mc=200):
+    s_acc = c_acc = 0.0
+    for _ in range(n_mc):
+        E = gen(64)
+        act = sample_active(64, k)
+        Ez = E * act[:, :, None]
+        ch = channel(Ez, snr_db=snr)
+        Eh = model(ch["yI"], ch["yQ"], active=act)
+        s_acc += ser(Eh, E, active=act.float())
+        c_acc += mean_cos(Eh, E, active=act.float())
+    return s_acc / n_mc, c_acc / n_mc
+
+
+print("[E3] training activity-aware model", flush=True)
+m_act = UWCA(d, Umax, H).to(DEVICE)
+train(m_act, tag="act")
+torch.save(m_act.state_dict(), lib.DATA / "e3_uwca_act.pt")
+
+print("[E3] training fixed-U8 model", flush=True)
+m_fix = UWCA(d, Umax, H).to(DEVICE)
+train(m_fix, k=Umax, tag="fix8")
+
+ks = [2, 3, 4, 5, 6, 7, 8]
+out = {"k": ks, "snr_eval": [10.0, 20.0], "activity": {}, "fixed8": {},
+       "oracle": {}}
+for snr in out["snr_eval"]:
+    out["activity"][str(snr)] = [evaluate(m_act, k, snr) for k in ks]
+    out["fixed8"][str(snr)] = [evaluate(m_fix, k, snr) for k in ks]
+    print(f"[E3] snr={snr} activity={[f'{a[0]:.3f}' for a in out['activity'][str(snr)]]}",
+          flush=True)
+
+for k in [2, 4, 6, 8]:
+    m_o = UWCA(d, Umax, H).to(DEVICE)
+    train(m_o, k=k, epochs=250, tag=f"oracle{k}")
+    out["oracle"][str(k)] = {str(snr): evaluate(m_o, k, snr)
+                             for snr in out["snr_eval"]}
+    print(f"[E3] oracle k={k}: {out['oracle'][str(k)]}", flush=True)
+
+save_json("e3_dynamic_users.json", out)