Initial release: code for WCL2026-1544 (context-aware embedding masking via DRL)

update_fig_with_improved.py

@@ -0,0 +1,114 @@
+#!/usr/bin/env python3
+"""
+update_fig_with_improved.py
+Overlay the best-performing DRL improvement variant onto Fig 3(b)
+(per-user CosSim vs SNR). Produces a new PDF
+`fig/wcl_fig_cossim_snr.pdf` with four curves:
+   Joint, MAML, Proposed DRL (paper), Proposed DRL (improved).
+"""
+import csv, os, argparse, sys
+from pathlib import Path
+import numpy as np
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+ROOT = Path(__file__).resolve().parent.parent
+FIG = ROOT / "fig"
+IMP = ROOT / "results_improve"
+SWEEP = ROOT / "results_sweeps"
+RES = ROOT / "results_drl"
+LONG = ROOT / "results_drl_long"
+
+
+def load(path):
+    if not os.path.exists(path):
+        return None
+    with open(path) as f:
+        rows = list(csv.DictReader(f))
+    return rows if rows else None
+
+
+def to_arr(rows):
+    rows = sorted(rows, key=lambda r: float(r["snr_db"]))
+    snrs = np.array([float(r["snr_db"]) for r in rows])
+    cos = np.array([float(r["cos_sim"]) for r in rows])
+    return snrs, cos
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--best", default=None,
+                    help="Variant name under results_improve/. If omitted, "
+                         "picks the best-avg-CosSim automatically.")
+    args = ap.parse_args()
+
+    # Auto-select best if unspecified
+    if args.best is None and IMP.exists():
+        best_name, best_avg = None, -1.0
+        for d in sorted(IMP.iterdir()):
+            rows = load(d / "drl_snr_sweep.csv")
+            if rows is None:
+                continue
+            _, cos = to_arr(rows)
+            if cos.mean() > best_avg:
+                best_avg, best_name = float(cos.mean()), d.name
+        args.best = best_name
+    if args.best is None:
+        print("[FATAL] No improvement variant found in results_improve/",
+              file=sys.stderr)
+        sys.exit(1)
+
+    imp_rows = load(IMP / args.best / "drl_snr_sweep.csv")
+    if imp_rows is None:
+        print(f"[FATAL] results_improve/{args.best}/drl_snr_sweep.csv missing")
+        sys.exit(1)
+
+    joint_rows = load(RES / "joint_snr_sweep.csv")
+    maml_rows = load(SWEEP / "maml_U4_200ep" / "maml_snr_sweep.csv")
+    drl_rows = load(LONG / "drl_snr_sweep.csv") or \
+               load(RES / "drl_snr_sweep.csv")
+    for name, rows in [("Joint", joint_rows), ("MAML", maml_rows),
+                       ("DRL", drl_rows)]:
+        if rows is None:
+            print(f"[FATAL] baseline {name} missing", file=sys.stderr)
+            sys.exit(1)
+
+    snrs, joint_cos = to_arr(joint_rows)
+    _, maml_cos = to_arr(maml_rows)
+    _, drl_cos = to_arr(drl_rows)
+    _, imp_cos = to_arr(imp_rows)
+
+    plt.rcParams.update({
+        "font.size": 9, "axes.labelsize": 9,
+        "legend.fontsize": 7, "xtick.labelsize": 8,
+        "ytick.labelsize": 8, "axes.linewidth": 0.8,
+        "lines.linewidth": 1.3, "figure.dpi": 150,
+    })
+    fig, ax = plt.subplots(figsize=(3.5, 2.6))
+    ax.plot(snrs, joint_cos, "o-", color="#2ca02c",
+            label="Joint (baseline)")
+    ax.plot(snrs, maml_cos, "s--", color="#1f77b4", label="MAML")
+    ax.plot(snrs, drl_cos, "^-", color="#d62728",
+            label="Proposed DRL (paper)")
+    ax.plot(snrs, imp_cos, "D-", color="#9467bd", linewidth=1.6,
+            label=f"Proposed DRL (improved)")
+    ax.set_xlabel("SNR (dB)")
+    ax.set_ylabel(r"Per-user CosSim")
+    lo = min(joint_cos.min(), maml_cos.min(), drl_cos.min(),
+             imp_cos.min()) - 0.02
+    ax.set_ylim(max(0.0, lo), 1.0)
+    ax.grid(True, alpha=0.3)
+    ax.legend(loc="lower right")
+    out = FIG / "wcl_fig_cossim_snr.pdf"
+    fig.savefig(out, bbox_inches="tight")
+    plt.close(fig)
+    print(f"[OK] overlaid best variant '{args.best}' -> {out}")
+    print(f"     avg CosSim: Joint={joint_cos.mean():.4f}, "
+          f"MAML={maml_cos.mean():.4f}, "
+          f"DRL={drl_cos.mean():.4f}, "
+          f"improved={imp_cos.mean():.4f}")
+
+
+if __name__ == "__main__":
+    main()