lwm-cde / experiments /source_selection_payoff.py

LWM-CDE: dataset distancing for wireless datasets

7378252 verified 3 days ago

11.1 kB

"""Experiment 1: End-to-end source-selection payoff.
For each held-out target dataset, pick a source under three strategies:
  (a) LWM-CDE distance     : argmin_s D_CDE[s, t]
  (b) Frozen-LWM distance  : argmin_s D_LWM[s, t]
  (c) Random               : random s != t
  (d) Oracle (P-row)       : argmin_s P[s, t]   <-- upper bound
Achieved transfer is P[s_picked, t] (lower = better, since P = 1 - F1).
Report mean +/- std across the held-out targets, plus the gap to oracle and
the gap to random expectation.
Usage:
  python experiments/source_selection_payoff.py \
      --task los-nlos --n_train 400 --n_test 400 --epochs 12
Output: results JSON next to this file plus a summary table to stdout.
"""
from __future__ import annotations

import argparse, contextlib, io, json, os, sys, time
from datetime import datetime

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))   # the repository root

import numpy as np
import torch

from dataset_distancing import (
    DatasetCollection,
    DistanceConfig,
    TransferConfig,
    compute_distances,
    compute_performance_matrix,
)
from dataset_distancing.embeddings import register_lwm_path, load_runner


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--lwm_path",   type=str, default="third_party/lwm-1.1")
    p.add_argument("--cache_dir",  type=str, default="cache/exp1")
    p.add_argument("--gpu",        type=int, default=1)
    p.add_argument("--seed",       type=int, default=42)
    p.add_argument("--n_train",    type=int, default=400)
    p.add_argument("--n_val",      type=int, default=400)
    p.add_argument("--n_test",     type=int, default=400)
    p.add_argument("--task",       type=str, default="los-nlos",
                   choices=["los-nlos", "beam-pred"])
    p.add_argument("--n_ant",      type=int, default=32)
    p.add_argument("--n_sub",      type=int, default=32)
    p.add_argument("--n_beams",    type=int, default=32)
    p.add_argument("--n_cities",   type=int, default=8,
                   help="number of distinct cities (each contributes n_bs base-stations)")
    p.add_argument("--n_bs",       type=int, default=3,
                   help="number of base stations per city (1..3); multi-BS adds within-city variety")
    p.add_argument("--epochs",     type=int, default=12)
    p.add_argument("--lr",         type=float, default=1e-3)
    p.add_argument("--batch_size", type=int, default=128)
    p.add_argument("--cde_epochs", type=int, default=20)
    p.add_argument("--n_random_trials", type=int, default=20,
                   help="random source picks per target (averaged)")
    p.add_argument("--out_json",   type=str, default=None)
    return p.parse_args()


def _select_pool_scenarios(task: str, n_cities: int, n_bs: int):
    """Return a deliberately *diverse* pool of (city, BS) DeepMIMO datasets.
    Mirrors the original paper's recipe of mixing high-LoS US cities with
    low-LoS international cities so the transfer matrix has enough variance to
    distinguish good from bad source choices.
    """
    import deepmimo as dm
    all_scs = dm.search()
    sub6 = [s for s in all_scs if "_3p5" in s and s.startswith("city_")]
    n_low  = n_cities // 2
    n_high = n_cities - n_low
    # First-indexed scenarios = US cities, mostly high LoS
    # Later-indexed scenarios = international cities (Singapore, Dubai, …) low LoS
    cities = sub6[:n_high] + sub6[max(len(sub6)-n_low, n_high):][:n_low]
    pool = []
    for s in cities:
        for tx in range(1, n_bs + 1):
            pool.append({"name": s, "tx_set": tx})
    return pool


def main():
    args = parse_args()
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    register_lwm_path(args.lwm_path)

    out_dir = os.path.dirname(os.path.abspath(__file__))
    out_json = args.out_json or os.path.join(
        out_dir, f"exp1_{args.task}_c{args.n_cities}b{args.n_bs}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    )

    # ── Pool ──────────────────────────────────────────────────────────────
    scenarios = _select_pool_scenarios(args.task, args.n_cities, args.n_bs)
    print(f"[Exp1] Pool: {len(scenarios)} scenarios for task={args.task} ({args.n_cities} cities × {args.n_bs} BS)")
    for s in scenarios:
        print("  -", s["name"])

    full = DatasetCollection.from_scenarios(
        scenarios,
        task=args.task,
        n_ant=args.n_ant,
        n_sub=args.n_sub,
        n_beams=args.n_beams,
        cache_dir=args.cache_dir,
    )

    # ── Split ─────────────────────────────────────────────────────────────
    train_ds, val_ds, test_ds = [], [], []
    for ds in full:
        tr, va, te = ds.split(
            train=args.n_train, val=args.n_val, test=args.n_test, seed=args.seed,
        )
        train_ds.append(tr); val_ds.append(va); test_ds.append(te)
    train_col = DatasetCollection(datasets=train_ds)
    val_col   = DatasetCollection(datasets=val_ds)
    test_col  = DatasetCollection(datasets=test_ds)

    # ── Runner ────────────────────────────────────────────────────────────
    with contextlib.redirect_stdout(io.StringIO()):
        runner = load_runner(args.lwm_path, weights_file=None)

    # ── P matrix (full pool, symmetric=False to keep the directed truth) ──
    print("[Exp1] Computing transfer matrix P …")
    t0 = time.time()
    P = compute_performance_matrix(
        train_col,
        cfg=TransferConfig(
            task=args.task, epochs=args.epochs, lr=args.lr,
            batch_size=args.batch_size, n_beams=args.n_beams,
            seed=args.seed, symmetrize=False, verbose=False,
        ),
        val_collection=val_col, test_collection=test_col,
    )
    P_np = P.detach().cpu().numpy()
    print(f"[Exp1] P done in {time.time()-t0:.1f}s, shape={P_np.shape}")

    # ── Frozen-LWM distances ──────────────────────────────────────────────
    print("[Exp1] Computing frozen-LWM distances …")
    D_lwm = compute_distances(train_col, runner, cfg=DistanceConfig(
        distance="swd", label_aware=True, n_per_dataset=args.n_train,
        sw_projections=64, asymmetric=False,
    )).detach().cpu().numpy()

    # ── CDE fine-tune on the same pool, then distances ────────────────────
    print("[Exp1] Fine-tuning LWM-CDE on the pool …")
    from dataset_distancing.cde import CDETrainer
    from dataset_distancing import CDEConfig
    cde_cfg = CDEConfig(
        epochs=args.cde_epochs, lr=5e-4,
        n_per_dataset=64, n_per_dataset_eval=128, corr_n_per_dataset=64,
        asymmetric=False,
        distance=DistanceConfig(sw_projections=32, label_aware=True),
    )
    trainer = CDETrainer(runner, cde_cfg)
    trainer.fit(train_col, P, eval_collection=train_col, perf_eval=P, verbose=False)

    print("[Exp1] Computing CDE-refined distances …")
    D_cde = compute_distances(train_col, runner, cfg=DistanceConfig(
        distance="swd", label_aware=True, n_per_dataset=args.n_train,
        sw_projections=64, asymmetric=False,
    )).detach().cpu().numpy()

    # ── Selection payoff ──────────────────────────────────────────────────
    N = P_np.shape[0]
    rng = np.random.default_rng(args.seed)

    def _pick_argmin(M, t):
        col = M[:, t].copy()
        col[t] = np.inf  # forbid self-selection
        return int(np.argmin(col))

    def _pick_random(t):
        choices = [s for s in range(N) if s != t]
        return int(rng.choice(choices))

    def _expected_random(t):
        choices = [s for s in range(N) if s != t]
        return float(np.mean(P_np[choices, t]))

    rows = []
    for t in range(N):
        s_oracle = _pick_argmin(P_np,  t)
        s_lwm    = _pick_argmin(D_lwm, t)
        s_cde    = _pick_argmin(D_cde, t)
        # Random: sample n_random_trials picks and average achieved P
        rand_pcs = [P_np[_pick_random(t), t] for _ in range(args.n_random_trials)]
        rows.append({
            "target":       t,
            "target_name":  test_col[t].name,
            "P_oracle":     float(P_np[s_oracle, t]),
            "P_lwm":        float(P_np[s_lwm,    t]),
            "P_cde":        float(P_np[s_cde,    t]),
            "P_rand_mean":  float(np.mean(rand_pcs)),
            "P_rand_std":   float(np.std(rand_pcs)),
            "P_rand_expected": _expected_random(t),
            "s_oracle":     s_oracle, "s_lwm": s_lwm, "s_cde": s_cde,
        })

    summary = {
        "task":       args.task,
        "n_pool":     N,
        "scenarios":  [s["name"] for s in scenarios],
        "rows":       rows,
        "agg": {
            "oracle_P_mean":  float(np.mean([r["P_oracle"] for r in rows])),
            "oracle_F1_mean": float(1.0 - np.mean([r["P_oracle"] for r in rows])),
            "lwm_P_mean":     float(np.mean([r["P_lwm"]    for r in rows])),
            "lwm_F1_mean":    float(1.0 - np.mean([r["P_lwm"]    for r in rows])),
            "cde_P_mean":     float(np.mean([r["P_cde"]    for r in rows])),
            "cde_F1_mean":    float(1.0 - np.mean([r["P_cde"]    for r in rows])),
            "rand_P_mean":    float(np.mean([r["P_rand_expected"] for r in rows])),
            "rand_F1_mean":   float(1.0 - np.mean([r["P_rand_expected"] for r in rows])),
        },
    }

    print("\n──────── Source-Selection Payoff ────────")
    print(f"task = {args.task}, pool = {N}")
    a = summary["agg"]
    print(f"  Random      : F1 = {a['rand_F1_mean']:.4f}   (P = {a['rand_P_mean']:.4f})")
    print(f"  Frozen LWM  : F1 = {a['lwm_F1_mean']:.4f}    (P = {a['lwm_P_mean']:.4f})")
    print(f"  LWM-CDE     : F1 = {a['cde_F1_mean']:.4f}    (P = {a['cde_P_mean']:.4f})")
    print(f"  Oracle (P)  : F1 = {a['oracle_F1_mean']:.4f} (P = {a['oracle_P_mean']:.4f})")
    if a["rand_F1_mean"] != a["oracle_F1_mean"]:
        gap = a["oracle_F1_mean"] - a["rand_F1_mean"]
        cde_recovered = (a["cde_F1_mean"] - a["rand_F1_mean"]) / gap if gap > 0 else float("nan")
        lwm_recovered = (a["lwm_F1_mean"] - a["rand_F1_mean"]) / gap if gap > 0 else float("nan")
        print(f"  Random→Oracle F1 gap: {gap:+.4f}")
        print(f"  Frozen LWM recovers {100*lwm_recovered:.1f}% of the gap")
        print(f"  LWM-CDE   recovers {100*cde_recovered:.1f}% of the gap")

    with open(out_json, "w") as f:
        json.dump(summary, f, indent=2)
    print(f"\n[Exp1] saved → {out_json}")


if __name__ == "__main__":
    main()

	"""Experiment 1: End-to-end source-selection payoff.

	For each held-out target dataset, pick a source under three strategies:
	(a) LWM-CDE distance : argmin_s D_CDE[s, t]
	(b) Frozen-LWM distance : argmin_s D_LWM[s, t]
	(c) Random : random s != t
	(d) Oracle (P-row) : argmin_s P[s, t] <-- upper bound

	Achieved transfer is P[s_picked, t] (lower = better, since P = 1 - F1).
	Report mean +/- std across the held-out targets, plus the gap to oracle and
	the gap to random expectation.

	Usage:
	python experiments/source_selection_payoff.py \
	--task los-nlos --n_train 400 --n_test 400 --epochs 12

	Output: results JSON next to this file plus a summary table to stdout.
	"""
	from __future__ import annotations

	import argparse, contextlib, io, json, os, sys, time
	from datetime import datetime

	sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) # the repository root

	import numpy as np
	import torch

	from dataset_distancing import (
	DatasetCollection,
	DistanceConfig,
	TransferConfig,
	compute_distances,
	compute_performance_matrix,
	)
	from dataset_distancing.embeddings import register_lwm_path, load_runner


	def parse_args():
	p = argparse.ArgumentParser()
	p.add_argument("--lwm_path", type=str, default="third_party/lwm-1.1")
	p.add_argument("--cache_dir", type=str, default="cache/exp1")
	p.add_argument("--gpu", type=int, default=1)
	p.add_argument("--seed", type=int, default=42)
	p.add_argument("--n_train", type=int, default=400)
	p.add_argument("--n_val", type=int, default=400)
	p.add_argument("--n_test", type=int, default=400)
	p.add_argument("--task", type=str, default="los-nlos",
	choices=["los-nlos", "beam-pred"])
	p.add_argument("--n_ant", type=int, default=32)
	p.add_argument("--n_sub", type=int, default=32)
	p.add_argument("--n_beams", type=int, default=32)
	p.add_argument("--n_cities", type=int, default=8,
	help="number of distinct cities (each contributes n_bs base-stations)")
	p.add_argument("--n_bs", type=int, default=3,
	help="number of base stations per city (1..3); multi-BS adds within-city variety")
	p.add_argument("--epochs", type=int, default=12)
	p.add_argument("--lr", type=float, default=1e-3)
	p.add_argument("--batch_size", type=int, default=128)
	p.add_argument("--cde_epochs", type=int, default=20)
	p.add_argument("--n_random_trials", type=int, default=20,
	help="random source picks per target (averaged)")
	p.add_argument("--out_json", type=str, default=None)
	return p.parse_args()


	def _select_pool_scenarios(task: str, n_cities: int, n_bs: int):
	"""Return a deliberately diverse pool of (city, BS) DeepMIMO datasets.

	Mirrors the original paper's recipe of mixing high-LoS US cities with
	low-LoS international cities so the transfer matrix has enough variance to
	distinguish good from bad source choices.
	"""
	import deepmimo as dm
	all_scs = dm.search()
	sub6 = [s for s in all_scs if "_3p5" in s and s.startswith("city_")]
	n_low = n_cities // 2
	n_high = n_cities - n_low
	# First-indexed scenarios = US cities, mostly high LoS
	# Later-indexed scenarios = international cities (Singapore, Dubai, …) low LoS
	cities = sub6[:n_high] + sub6[max(len(sub6)-n_low, n_high):][:n_low]
	pool = []
	for s in cities:
	for tx in range(1, n_bs + 1):
	pool.append({"name": s, "tx_set": tx})
	return pool


	def main():
	args = parse_args()
	os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)

	torch.manual_seed(args.seed)
	np.random.seed(args.seed)
	register_lwm_path(args.lwm_path)

	out_dir = os.path.dirname(os.path.abspath(__file__))
	out_json = args.out_json or os.path.join(
	out_dir, f"exp1_{args.task}_c{args.n_cities}b{args.n_bs}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
	)

	# ── Pool ──────────────────────────────────────────────────────────────
	scenarios = _select_pool_scenarios(args.task, args.n_cities, args.n_bs)
	print(f"[Exp1] Pool: {len(scenarios)} scenarios for task={args.task} ({args.n_cities} cities × {args.n_bs} BS)")
	for s in scenarios:
	print(" -", s["name"])

	full = DatasetCollection.from_scenarios(
	scenarios,
	task=args.task,
	n_ant=args.n_ant,
	n_sub=args.n_sub,
	n_beams=args.n_beams,
	cache_dir=args.cache_dir,
	)

	# ── Split ─────────────────────────────────────────────────────────────
	train_ds, val_ds, test_ds = [], [], []
	for ds in full:
	tr, va, te = ds.split(
	train=args.n_train, val=args.n_val, test=args.n_test, seed=args.seed,
	)
	train_ds.append(tr); val_ds.append(va); test_ds.append(te)
	train_col = DatasetCollection(datasets=train_ds)
	val_col = DatasetCollection(datasets=val_ds)
	test_col = DatasetCollection(datasets=test_ds)

	# ── Runner ────────────────────────────────────────────────────────────
	with contextlib.redirect_stdout(io.StringIO()):
	runner = load_runner(args.lwm_path, weights_file=None)

	# ── P matrix (full pool, symmetric=False to keep the directed truth) ──
	print("[Exp1] Computing transfer matrix P …")
	t0 = time.time()
	P = compute_performance_matrix(
	train_col,
	cfg=TransferConfig(
	task=args.task, epochs=args.epochs, lr=args.lr,
	batch_size=args.batch_size, n_beams=args.n_beams,
	seed=args.seed, symmetrize=False, verbose=False,
	),
	val_collection=val_col, test_collection=test_col,
	)
	P_np = P.detach().cpu().numpy()
	print(f"[Exp1] P done in {time.time()-t0:.1f}s, shape={P_np.shape}")

	# ── Frozen-LWM distances ──────────────────────────────────────────────
	print("[Exp1] Computing frozen-LWM distances …")
	D_lwm = compute_distances(train_col, runner, cfg=DistanceConfig(
	distance="swd", label_aware=True, n_per_dataset=args.n_train,
	sw_projections=64, asymmetric=False,
	)).detach().cpu().numpy()

	# ── CDE fine-tune on the same pool, then distances ────────────────────
	print("[Exp1] Fine-tuning LWM-CDE on the pool …")
	from dataset_distancing.cde import CDETrainer
	from dataset_distancing import CDEConfig
	cde_cfg = CDEConfig(
	epochs=args.cde_epochs, lr=5e-4,
	n_per_dataset=64, n_per_dataset_eval=128, corr_n_per_dataset=64,
	asymmetric=False,
	distance=DistanceConfig(sw_projections=32, label_aware=True),
	)
	trainer = CDETrainer(runner, cde_cfg)
	trainer.fit(train_col, P, eval_collection=train_col, perf_eval=P, verbose=False)

	print("[Exp1] Computing CDE-refined distances …")
	D_cde = compute_distances(train_col, runner, cfg=DistanceConfig(
	distance="swd", label_aware=True, n_per_dataset=args.n_train,
	sw_projections=64, asymmetric=False,
	)).detach().cpu().numpy()

	# ── Selection payoff ──────────────────────────────────────────────────
	N = P_np.shape[0]
	rng = np.random.default_rng(args.seed)

	def _pick_argmin(M, t):
	col = M[:, t].copy()
	col[t] = np.inf # forbid self-selection
	return int(np.argmin(col))

	def _pick_random(t):
	choices = [s for s in range(N) if s != t]
	return int(rng.choice(choices))

	def _expected_random(t):
	choices = [s for s in range(N) if s != t]
	return float(np.mean(P_np[choices, t]))

	rows = []
	for t in range(N):
	s_oracle = _pick_argmin(P_np, t)
	s_lwm = _pick_argmin(D_lwm, t)
	s_cde = _pick_argmin(D_cde, t)
	# Random: sample n_random_trials picks and average achieved P
	rand_pcs = [P_np[_pick_random(t), t] for _ in range(args.n_random_trials)]
	rows.append({
	"target": t,
	"target_name": test_col[t].name,
	"P_oracle": float(P_np[s_oracle, t]),
	"P_lwm": float(P_np[s_lwm, t]),
	"P_cde": float(P_np[s_cde, t]),
	"P_rand_mean": float(np.mean(rand_pcs)),
	"P_rand_std": float(np.std(rand_pcs)),
	"P_rand_expected": _expected_random(t),
	"s_oracle": s_oracle, "s_lwm": s_lwm, "s_cde": s_cde,
	})

	summary = {
	"task": args.task,
	"n_pool": N,
	"scenarios": [s["name"] for s in scenarios],
	"rows": rows,
	"agg": {
	"oracle_P_mean": float(np.mean([r["P_oracle"] for r in rows])),
	"oracle_F1_mean": float(1.0 - np.mean([r["P_oracle"] for r in rows])),
	"lwm_P_mean": float(np.mean([r["P_lwm"] for r in rows])),
	"lwm_F1_mean": float(1.0 - np.mean([r["P_lwm"] for r in rows])),
	"cde_P_mean": float(np.mean([r["P_cde"] for r in rows])),
	"cde_F1_mean": float(1.0 - np.mean([r["P_cde"] for r in rows])),
	"rand_P_mean": float(np.mean([r["P_rand_expected"] for r in rows])),
	"rand_F1_mean": float(1.0 - np.mean([r["P_rand_expected"] for r in rows])),
	},
	}

	print("\n──────── Source-Selection Payoff ────────")
	print(f"task = {args.task}, pool = {N}")
	a = summary["agg"]
	print(f" Random : F1 = {a['rand_F1_mean']:.4f} (P = {a['rand_P_mean']:.4f})")
	print(f" Frozen LWM : F1 = {a['lwm_F1_mean']:.4f} (P = {a['lwm_P_mean']:.4f})")
	print(f" LWM-CDE : F1 = {a['cde_F1_mean']:.4f} (P = {a['cde_P_mean']:.4f})")
	print(f" Oracle (P) : F1 = {a['oracle_F1_mean']:.4f} (P = {a['oracle_P_mean']:.4f})")
	if a["rand_F1_mean"] != a["oracle_F1_mean"]:
	gap = a["oracle_F1_mean"] - a["rand_F1_mean"]
	cde_recovered = (a["cde_F1_mean"] - a["rand_F1_mean"]) / gap if gap > 0 else float("nan")
	lwm_recovered = (a["lwm_F1_mean"] - a["rand_F1_mean"]) / gap if gap > 0 else float("nan")
	print(f" Random→Oracle F1 gap: {gap:+.4f}")
	print(f" Frozen LWM recovers {100*lwm_recovered:.1f}% of the gap")
	print(f" LWM-CDE recovers {100*cde_recovered:.1f}% of the gap")

	with open(out_json, "w") as f:
	json.dump(summary, f, indent=2)
	print(f"\n[Exp1] saved → {out_json}")


	if __name__ == "__main__":
	main()