{ "cells": [ { "cell_type": "markdown", "id": "e0b5341b", "metadata": {}, "source": [ "# Dataset Distancing Tutorial\n", "\n", "End-to-end walkthrough: generate wireless channel datasets from DeepMIMO scenarios,\n", "compute pairwise distances in the LWM embedding space, compute transfer performance,\n", "correlate them, and apply distances for source selection, data augmentation, and\n", "budgeted pretraining." ] }, { "cell_type": "markdown", "id": "3ba2ce0e", "metadata": {}, "source": [ "## Setup\n", "\n", "Add the LWM package to `sys.path` so that `from_scenarios()` and `load_runner()` can find it." ] }, { "cell_type": "code", "execution_count": null, "id": "b2b026f9", "metadata": {}, "outputs": [], "source": [ "import sys, os\n", "\n", "# Use GPU 1 (GPU 0 may be occupied by other users)\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n", "\n", "# Add LWM package to sys.path (adjust this path to your LWM installation)\n", "LWM_PATH = \"third_party/lwm-1.1\"\n", "if LWM_PATH not in sys.path:\n", " sys.path.insert(0, LWM_PATH)\n", "\n", "# Cache directory — datasets, trained heads, and scores are persisted here.\n", "# Adding a new scenario only triggers computation for that scenario.\n", "CACHE_DIR = \"cache/tutorial_los_nlos\"\n", "\n", "%matplotlib inline\n", "import matplotlib.pyplot as plt\n", "\n", "import torch\n", "print(f\"PyTorch {torch.__version__}, CUDA available: {torch.cuda.is_available()}\")\n", "if torch.cuda.is_available():\n", " print(f\"Using: {torch.cuda.get_device_name(0)}, \"\n", " f\"Free: {torch.cuda.mem_get_info()[0]/1e9:.1f} GB\")" ] }, { "cell_type": "markdown", "id": "840849ba", "metadata": {}, "source": [ "## Step 1: Generate Datasets\n", "\n", "Datasets are generated on the fly from [DeepMIMO](https://deepmimo.net/) ray-tracing scenarios. Each scenario becomes one `WirelessDataset` with complex channel tensors and task labels. Scenarios are auto-downloaded on first use.\n", "\n", "**Note:** Using geographically diverse scenarios (US + international cities) is important for meaningful distance–performance correlation." ] }, { "cell_type": "code", "execution_count": null, "id": "279146f3", "metadata": {}, "outputs": [], "source": [ "from dataset_distancing import DatasetCollection\n", "import deepmimo as dm\n", "from lwm.data.scenarios import ScenarioConfig\n", "\n", "# Discover 3.5 GHz city scenarios via DeepMIMO search API\n", "all_scenarios = dm.search({\"bands\": [\"sub6\"], \"environment\": \"outdoor\"})\n", "city_3p5 = [s for s in all_scenarios if \"city_\" in s and \"3p5\" in s and \"_s\" not in s]\n", "\n", "# Split into multi-TX (US, high LoS) and single-TX (international, low LoS)\n", "multi_tx = set(dm.search({\"bands\": [\"sub6\"], \"numTx\": {\"min\": 2}}))\n", "us_cities = [s for s in city_3p5 if s in multi_tx][:2]\n", "intl_cities = [s for s in city_3p5 if s not in multi_tx][:2]\n", "\n", "# Build ScenarioConfigs: US cities get two TX sets each for LoS diversity;\n", "# international cities use TX set 1 (only one available).\n", "configs = []\n", "for name in us_cities:\n", " configs.append(ScenarioConfig(name, tx_set=1, n_ant=32, n_sub=32)) # high LoS\n", " configs.append(ScenarioConfig(name, tx_set=2, n_ant=32, n_sub=32)) # different LoS\n", "for name in intl_cities:\n", " configs.append(ScenarioConfig(name, tx_set=1, n_ant=32, n_sub=32)) # low LoS\n", "\n", "collection = DatasetCollection.from_scenarios(\n", " configs,\n", " task=\"los-nlos\",\n", " n_ant=32, n_sub=32,\n", " max_samples=3000,\n", " balance_classes=False,\n", " cache_dir=CACHE_DIR,\n", ")\n", "\n", "print(f\"Generated {len(collection)} datasets:\")\n", "for ds in collection:\n", " n0 = (ds.labels == 0).sum().item()\n", " n1 = (ds.labels == 1).sum().item()\n", " print(f\" {ds.name}: {len(ds)} samples (LoS={n0}, NLoS={n1}, LoS%={100*n0/len(ds):.0f}%)\")\n" ] }, { "cell_type": "markdown", "id": "b1759490", "metadata": {}, "source": [ "## Step 2: Compute Pairwise Distances\n", "\n", "Load the LWM encoder, embed each dataset, and compute **label-aware** Sliced Wasserstein Distances (SWD).\n", "\n", "Label-aware distances compute per-class embedding distances and average them, isolating channel similarity from label-ratio overlap. This requires datasets with **diverse class ratios** — mixing high-LoS and low-LoS scenarios. With the pretrained LWM encoder, expect **Pearson ~0.5–0.7**.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "7d7ddf69", "metadata": {}, "outputs": [], "source": [ "from dataset_distancing import compute_distances, DistanceConfig\n", "from dataset_distancing.embeddings import load_runner\n", "import torch, numpy as np\n", "\n", "torch.manual_seed(42)\n", "np.random.seed(42)\n", "\n", "# Pretrained LWM encoder\n", "runner = load_runner(LWM_PATH)\n", "\n", "D = compute_distances(collection, runner, cfg=DistanceConfig(\n", " distance=\"swd\",\n", " label_aware=True, # per-class distances — critical for correlation\n", " n_per_dataset=500,\n", " sw_projections=64,\n", "))\n", "\n", "print(f\"Distance matrix ({D.shape[0]}x{D.shape[1]}):\")\n", "print(D)\n" ] }, { "cell_type": "markdown", "id": "5c91dd6e", "metadata": {}, "source": [ "## Step 3: Compute Transfer Performance Matrix\n", "\n", "For every pair (i, j), train a lightweight LoS/NLoS classifier on dataset i and evaluate weighted F1 on dataset j. The matrix entry is `P[i,j] = 1 - F1` (lower = better transfer). The matrix is automatically symmetrized." ] }, { "cell_type": "code", "execution_count": null, "id": "644b91ce", "metadata": {}, "outputs": [], "source": [ "from dataset_distancing import compute_performance_matrix, TransferConfig\n", "\n", "P = compute_performance_matrix(\n", " collection,\n", " cfg=TransferConfig(\n", " task=\"los-nlos\",\n", " epochs=25,\n", " batch_size=32,\n", " verbose=True,\n", " ),\n", ")\n", "\n", "print(f\"\\nPerformance matrix (1 - F1):\")\n", "print(P)\n" ] }, { "cell_type": "markdown", "id": "e68328f0", "metadata": {}, "source": [ "## Step 4: Correlate Distances with Transfer Performance\n", "\n", "Compute Pearson and Spearman correlation between the upper-triangular entries of the distance matrix D and the performance matrix P. High positive correlation means larger distance predicts worse transfer (higher 1-F1).\n", "\n", "With label-aware SWD and the pretrained LWM encoder on datasets with diverse class ratios, expect **Pearson ~0.6–0.7**. CDE fine-tuning (below) can push this to **~0.8**.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "1ceba9fa", "metadata": {}, "outputs": [], "source": [ "from dataset_distancing import correlate\n", "\n", "stats = correlate(D, P)\n", "\n", "print(f\"Pearson: r = {stats['pearson']:.4f} (p = {stats['pearson_p']:.2e})\")\n", "print(f\"Spearman: rho = {stats['spearman']:.4f} (p = {stats['spearman_p']:.2e})\")" ] }, { "cell_type": "code", "execution_count": null, "id": "5975b355", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "\n", "short_names = [n.replace(\"_3p5_tx1\", \"\").replace(\"city_\", \"\") for n in collection.names()]\n", "\n", "fig, axes = plt.subplots(1, 3, figsize=(18, 5))\n", "\n", "# Distance matrix heatmap\n", "im0 = axes[0].imshow(D.numpy(), cmap=\"viridis\")\n", "axes[0].set_xticks(range(len(short_names)))\n", "axes[0].set_yticks(range(len(short_names)))\n", "axes[0].set_xticklabels(short_names, rotation=45, ha=\"right\", fontsize=7)\n", "axes[0].set_yticklabels(short_names, fontsize=7)\n", "axes[0].set_title(\"Distance Matrix D (SWD)\")\n", "fig.colorbar(im0, ax=axes[0], shrink=0.8)\n", "\n", "# Performance matrix heatmap\n", "im1 = axes[1].imshow(P.numpy(), cmap=\"magma\")\n", "axes[1].set_xticks(range(len(short_names)))\n", "axes[1].set_yticks(range(len(short_names)))\n", "axes[1].set_xticklabels(short_names, rotation=45, ha=\"right\", fontsize=7)\n", "axes[1].set_yticklabels(short_names, fontsize=7)\n", "axes[1].set_title(\"Performance Matrix P (1 - F1)\")\n", "fig.colorbar(im1, ax=axes[1], shrink=0.8)\n", "\n", "# Scatter: distance vs performance (upper triangle)\n", "tri = np.triu_indices(len(collection), k=1)\n", "d_flat = D.numpy()[tri]\n", "p_flat = P.numpy()[tri]\n", "axes[2].scatter(d_flat, p_flat, alpha=0.6, edgecolors=\"k\", linewidths=0.5)\n", "# Fit line\n", "m, b = np.polyfit(d_flat, p_flat, 1)\n", "x_line = np.linspace(d_flat.min(), d_flat.max(), 100)\n", "axes[2].plot(x_line, m * x_line + b, \"r--\", linewidth=1.5)\n", "axes[2].set_xlabel(\"Dataset Distance (SWD)\")\n", "axes[2].set_ylabel(\"Transfer Gap (1 - F1)\")\n", "axes[2].set_title(f\"Correlation: r={stats['pearson']:.3f}, ρ={stats['spearman']:.3f}\")\n", "axes[2].grid(True, alpha=0.3)\n", "\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Optional: CDE Fine-Tuning on These Scenarios\n", "\n", "Train LWM-CDE directly on the current scenario set to maximize distance–performance correlation. This fine-tunes the top LWM layers and a small projection head so that embedding distances better predict transfer performance.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from dataset_distancing.cde import CDETrainer\n", "from dataset_distancing import CDEConfig\n", "\n", "# Reload runner (fine-tuning modifies the model in-place)\n", "runner_ft = load_runner(LWM_PATH)\n", "\n", "trainer = CDETrainer(runner_ft, CDEConfig(\n", " epochs=60,\n", " lr=5e-4,\n", " finetune_start_layer=6, # unfreeze top 6 layers\n", " n_per_dataset=64,\n", " n_per_dataset_eval=128,\n", " w_listwise=1.0,\n", " w_correlation=0.5,\n", " w_consistency=0.2,\n", " w_distill=0.2,\n", " eval_every=5,\n", " early_stop_patience=6,\n", " distance=DistanceConfig(sw_projections=16, label_aware=True),\n", "))\n", "\n", "# Fine-tune using the performance matrix as supervision\n", "stats = trainer.fit(collection, P, eval_collection=collection, perf_eval=P, verbose=True)\n", "print(f\"\\nAfter CDE fine-tuning: Pearson={stats['pearson']:.4f}, Spearman={stats['spearman']:.4f}\")\n", "\n", "# Recompute distances with the fine-tuned model + head\n", "D_cde = trainer.compute_distances(collection, DistanceConfig(\n", " distance=\"swd\",\n", " label_aware=True,\n", " n_per_dataset=512,\n", " sw_projections=256,\n", "))\n", "\n", "stats_cde = correlate(D_cde, P)\n", "print(f\"\\nFine-tuned distances: Pearson={stats_cde['pearson']:.4f}, Spearman={stats_cde['spearman']:.4f}\")\n", "print(f\"Baseline pretrained: Pearson={correlate(D, P)['pearson']:.4f}\")\n", "\n", "# Save the fine-tuned model (backbone + head)\n", "trainer.save(f\"{CACHE_DIR}/cde_finetuned\")\n", "print(f\"\\nCDE model saved to {CACHE_DIR}/cde_finetuned/\")\n" ] }, { "cell_type": "markdown", "id": "1fed0c5b", "metadata": {}, "source": [ "## Step 5a: Source Selection\n", "\n", "For each target dataset, find the nearest source datasets in the distance space. The prediction: training on the nearest source yields the best transfer. We verify against the actual performance matrix." ] }, { "cell_type": "code", "execution_count": null, "id": "14647897", "metadata": {}, "outputs": [], "source": [ "from dataset_distancing.applications import source_selection\n", "\n", "picks = source_selection(D, target_indices=[7, 9], top_k=2)\n", "\n", "print(\"Source selection (top-2 nearest by distance):\")\n", "for t, srcs in picks.items():\n", " print(f\" {collection[t].name} <- {[collection[s].name for s in srcs]}\")\n", "\n", "# Verify against performance matrix (ground truth)\n", "print(\"\\nGround truth (lowest 1-F1 = best actual transfer):\")\n", "for t in [7, 9]:\n", " row = P[t].clone()\n", " row[t] = float(\"inf\")\n", " best = row.argsort()[:2].tolist()\n", " print(f\" {collection[t].name} <- {[collection[s].name for s in best]}\")" ] }, { "cell_type": "markdown", "id": "0c58bd5e", "metadata": {}, "source": [ "## Step 5b: Data Augmentation Ranking\n", "\n", "Rank all candidate datasets by distance to a label-scarce target. The nearest dataset is predicted to provide the most useful augmentation data." ] }, { "cell_type": "code", "execution_count": null, "id": "37fb4c5e", "metadata": {}, "outputs": [], "source": [ "from dataset_distancing.applications import augmentation_ranking\n", "\n", "target_idx = 0\n", "ranking = augmentation_ranking(D, target_idx=target_idx)\n", "\n", "print(f\"Augmentation ranking for {collection[target_idx].name} (nearest first):\")\n", "for rank, idx in enumerate(ranking, 1):\n", " print(f\" {rank}. {collection[idx].name} \"\n", " f\"(distance = {D[target_idx, idx]:.4f}, \"\n", " f\"1-F1 = {P[target_idx, idx]:.4f})\")" ] }, { "cell_type": "markdown", "id": "cbe65884", "metadata": {}, "source": [ "## Step 5c: Budgeted Pretraining\n", "\n", "Select k diverse representative datasets that cover the dataset manifold. Four selection algorithms are available." ] }, { "cell_type": "code", "execution_count": null, "id": "93183ebf", "metadata": {}, "outputs": [], "source": [ "from dataset_distancing.applications import budgeted_pretraining\n", "from dataset_distancing import select_datasets\n", "\n", "# k-medoids (default) -- minimizes total assignment cost\n", "subset_kmed = budgeted_pretraining(D, k=3, method=\"kmedoids\")\n", "print(f\"k-medoids (k=3): {[collection[i].name for i in subset_kmed]}\")\n", "\n", "# k-center (farthest-first traversal) -- minimizes worst-case assignment\n", "subset_kc = select_datasets(D, k=3, method=\"kcenter\")\n", "print(f\"k-center (k=3): {[collection[i].name for i in subset_kc]}\")\n", "\n", "# Facility location (greedy submodular) -- maximizes similarity coverage\n", "subset_fl = select_datasets(D, k=3, method=\"fl\")\n", "print(f\"Facility-loc (k=3): {[collection[i].name for i in subset_fl]}\")\n", "\n", "# Max-sum diversity -- maximizes total pairwise distance\n", "subset_msd = select_datasets(D, k=3, method=\"msd\")\n", "print(f\"Max-sum div (k=3): {[collection[i].name for i in subset_msd]}\")" ] }, { "cell_type": "markdown", "id": "4ef4d8e5", "metadata": {}, "source": [ "## Summary" ] }, { "cell_type": "code", "execution_count": null, "id": "3b377da7", "metadata": {}, "outputs": [], "source": [ "names = collection.names()\n", "\n", "print(\"=\" * 60)\n", "print(\"DATASET DISTANCING RESULTS\")\n", "print(\"=\" * 60)\n", "\n", "print(f\"\\nDatasets ({len(names)}):\")\n", "for i, name in enumerate(names):\n", " print(f\" [{i}] {name} ({len(collection[i])} samples)\")\n", "\n", "print(f\"\\nDistance matrix D (SWD):\")\n", "print(D.numpy().round(4))\n", "\n", "print(f\"\\nPerformance matrix P (1 - F1):\")\n", "print(P.numpy().round(4))\n", "\n", "print(f\"\\nCorrelation (D vs P):\")\n", "print(f\" Pearson r = {stats['pearson']:.4f} (p = {stats['pearson_p']:.2e})\")\n", "print(f\" Spearman rho = {stats['spearman']:.4f} (p = {stats['spearman_p']:.2e})\")\n", "\n", "print(f\"\\nSource selection (top-2):\")\n", "for t, srcs in picks.items():\n", " print(f\" {names[t]} <- {[names[s] for s in srcs]}\")\n", "\n", "print(f\"\\nAugmentation ranking for {names[target_idx]}:\")\n", "for rank, idx in enumerate(ranking, 1):\n", " print(f\" {rank}. {names[idx]}\")\n", "\n", "print(f\"\\nBudgeted pretraining (k=3, k-medoids):\")\n", "print(f\" {[names[i] for i in subset_kmed]}\")\n", "\n", "print(\"\\n\" + \"=\" * 60)" ] } ], "metadata": { "kernelspec": { "display_name": "lwmComp_env", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.0" } }, "nbformat": 4, "nbformat_minor": 5 }