examples/quickstart.py

"""Quick start: distances between wireless datasets in the frozen LWM space, and four decisions they make.

    python scripts/download_lwm.py        # once: the LWM 1.1 weights from Hugging Face
    python examples/quickstart.py         # from the repository root; a few minutes on a GPU

The script builds LOS/NLOS datasets for eight DeepMIMO base stations (DeepMIMO downloads a scenario the first time it
is used), embeds every channel with the pretrained LWM 1.1, and computes the label-aware sliced Wasserstein distance
between every pair of datasets. It then uses the distances to (i) pick the dataset a new site should learn from,
(ii) rank the datasets to add to the site's few labels, (iii) pick the two datasets that cover the library best for
pretraining, and (iv) check, without labels, whether new data still resembles the site's training data.
"""
import os
import sys

import torch

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

from dataset_distancing import DatasetCollection, DistanceConfig, compute_distances
from dataset_distancing.applications import augmentation_ranking, budgeted_pretraining, source_selection
from dataset_distancing.data import WirelessDataset
from dataset_distancing.embeddings import load_runner

SCENARIOS = [("city_0_newyork_3p5", 1), ("city_0_newyork_3p5", 2), ("city_1_losangeles_3p5", 1), ("city_1_losangeles_3p5", 2),
             ("city_2_chicago_3p5", 1), ("city_3_houston_3p5", 1), ("city_4_phoenix_3p5", 1), ("city_5_philadelphia_3p5", 1)]
SITE = 0                                                                   # the new site: New York, base station 1

torch.manual_seed(0)
library = DatasetCollection.from_scenarios([{"name": s, "tx_set": tx} for s, tx in SCENARIOS], task="los-nlos",
                                           n_ant=32, n_sub=32, max_samples=600, cache_dir="cache/quickstart")
names = library.names()
runner = load_runner()                                                     # LWM 1.1, used as is

# label-aware sliced Wasserstein distance between every pair of datasets
D = compute_distances(library, runner, cfg=DistanceConfig(distance="swd", label_aware=True, n_per_dataset=300))
width = max(len(n) for n in names)
print("\nDistances in the frozen LWM space:")
for i, n in enumerate(names):
    print(f"  {n:<{width}}  " + "  ".join(f"{v:6.4f}" for v in D[i].tolist()))

# (i) transfer learning: the dataset whose model the new site should start from
best = source_selection(D, target_indices=[SITE], top_k=3)[SITE]
print(f"\n(i) {names[SITE]} should learn from {names[best[0]]} (then {', '.join(names[i] for i in best[1:])})")

# (ii) data augmentation: the datasets to add to the site's few labels, closest first
order = augmentation_ranking(D, SITE)
print(f"(ii) data to add to {names[SITE]}, closest first: {', '.join(names[i] for i in order[:3])}")

# (iii) pretraining data: the two datasets that cover the library best (k-medoids on the distances)
picks = budgeted_pretraining(D, k=2)
print(f"(iii) pretrain on {', '.join(names[i] for i in picks)}")

# (iv) monitoring: is new data still like the site's training data? The distance below ignores labels, so a site can
# check new data before anyone labels it. Its normal level is the distance between two disjoint random parts of the
# site's own data, and a threshold three standard deviations above that level flags data that has changed. The parts hold
# 1000 users each: with a few hundred, two parts of one site can sit as far apart as two different sites.
mon = DatasetCollection.from_scenarios([{"name": "city_0_newyork_3p5", "tx_set": 1}, {"name": "city_5_philadelphia_3p5", "tx_set": 1}],
                                       task="los-nlos", n_ant=32, n_sub=32, max_samples=4000, cache_dir="cache/quickstart")
site, other, N = mon[0], mon[1], 1000
free = DistanceConfig(distance="swd", label_aware=False, n_per_dataset=N, sw_projections=512)
def part(ds, seed, k):                                                    # the k-th of disjoint random parts of N users
    idx = torch.randperm(len(ds), generator=torch.Generator().manual_seed(seed))[k * N:(k + 1) * N]
    return WirelessDataset(name=f"{ds.name}_{seed}_{k}", channels=ds.channels[idx])
def dist(ref, new):
    return compute_distances(DatasetCollection(datasets=[ref]), runner, cfg=free, test_collection=DatasetCollection(datasets=[new]))[0, 0].item()
normal = torch.tensor([dist(part(site, s, 0), part(site, s, 1)) for s in range(6)])
threshold = (normal.mean() + 3 * normal.std()).item()
reference = part(site, 99, 0)
print(f"(iv) distance between two parts of {site.name}: {normal.mean():.4f} on average, threshold {threshold:.4f}")
for what, new in [("new data from the same site", part(site, 99, 1)), (f"new data from a different environment ({other.name})", part(other, 99, 0))]:
    d = dist(reference, new)
    print(f"     {what}: distance {d:.4f} -> {'act: retrain, switch or fall back' if d > threshold else 'keep the model'}")