v2: distance-to-transmitter channel — fixes far-field square cutoff

src/dataset.py

@@ -41,9 +41,11 @@ class RadioMapSeerDataset(Dataset):
     }
 
     def __init__(self, data_root=DEFAULT_DATA_ROOT, split="train",
-                 sim="DPM", img_size=256, max_samples=None):
+                 sim="DPM", img_size=256, max_samples=None, add_distance=False):
         self.data_root = Path(data_root)
         self.img_size = img_size
+        self.add_distance = add_distance      # v2: distance-to-Tx channel
+        self._build_cache = {}                # cache shared building maps
 
         self.gain_dir  = self.data_root / "gain" / sim
         self.build_dir = self.data_root / "png" / "buildings_complete"
@@ -79,12 +81,26 @@ class RadioMapSeerDataset(Dataset):
     def __getitem__(self, idx):
         map_id, tx_id = self.samples[idx]
 
-        buildings = _load_gray(self.build_dir / f"{map_id}.png", self.img_size)
-        antenna   = _load_gray(self.ant_dir   / f"{map_id}_{tx_id}.png", self.img_size)
-        gain      = _load_gray(self.gain_dir  / f"{map_id}_{tx_id}.png", self.img_size)
+        # buildings are shared across a map's 80 transmitters -> cache to cut I/O
+        buildings = self._build_cache.get(map_id)
+        if buildings is None:
+            buildings = _load_gray(self.build_dir / f"{map_id}.png", self.img_size)
+            self._build_cache[map_id] = buildings
 
-        x = np.stack([buildings, antenna], axis=0)   # [2, H, W]
-        y = gain[None, :, :]                          # [1, H, W]
+        antenna = _load_gray(self.ant_dir  / f"{map_id}_{tx_id}.png", self.img_size)
+        gain    = _load_gray(self.gain_dir / f"{map_id}_{tx_id}.png", self.img_size)
+
+        channels = [buildings, antenna]
+        if self.add_distance:                              # v2 third channel
+            H = W = self.img_size
+            r, c = np.unravel_index(int(np.argmax(antenna)), antenna.shape)
+            yy, xx = np.mgrid[0:H, 0:W]
+            dist = np.sqrt((yy - r) ** 2 + (xx - c) ** 2).astype(np.float32)
+            dist /= (2 ** 0.5) * H                          # normalize to ~[0, 1]
+            channels.append(dist)
+
+        x = np.stack(channels, axis=0)
+        y = gain[None, :, :]
         return torch.from_numpy(x), torch.from_numpy(y)
 
 

src/train.py

@@ -1,15 +1,18 @@
 """
 train.py — train the U-Net on RadioMapSeer coverage maps.
 
-Examples:
-  # smoke test: tiny subset, 1 epoch — verify the loop and measure VRAM/speed
-  python src\\train.py --max-samples 200 --epochs 1 --batch-size 16
+v1: 2-channel input (buildings, transmitter).
+v2: add --distance for a 3rd channel (distance-to-Tx) that fixes the square cutoff.
 
-  # full run (size it after the smoke test)
-  python src\\train.py --epochs 30 --batch-size 16
+Examples:
+  # v1
+  python src\\train.py --epochs 30 --batch-size 32 --num-workers 8
+  # v2 (separate checkpoint dir so v1 stays intact)
+  python src\\train.py --distance --epochs 30 --batch-size 32 --num-workers 8 --ckpt-dir checkpoints_v2
 """
 
 import argparse
+import csv
 import time
 from pathlib import Path
 
@@ -20,15 +23,17 @@ from torch.utils.data import DataLoader
 from dataset import RadioMapSeerDataset
 from model import UNet
 
+torch.backends.cudnn.benchmark = True   # autotune convs (input size is fixed)
 
-def evaluate(model, loader, criterion, device, amp):
-    model.eval()
+
+def evaluate(net, loader, criterion, device, amp):
+    net.eval()
     total, n = 0.0, 0
     with torch.no_grad():
         for x, y in loader:
             x, y = x.to(device), y.to(device)
             with torch.autocast(device_type=device, dtype=torch.float16, enabled=amp):
-                loss = criterion(model(x), y)
+                loss = criterion(net(x), y)
             total += loss.item() * x.size(0)
             n += x.size(0)
     return total / n
@@ -37,25 +42,28 @@ def evaluate(model, loader, criterion, device, amp):
 def main():
     p = argparse.ArgumentParser()
     p.add_argument("--epochs", type=int, default=30)
-    p.add_argument("--batch-size", type=int, default=16)
+    p.add_argument("--batch-size", type=int, default=32)
     p.add_argument("--lr", type=float, default=1e-3)
     p.add_argument("--img-size", type=int, default=256)
     p.add_argument("--base", type=int, default=64)
-    p.add_argument("--max-samples", type=int, default=None,
-                   help="Limit samples for quick smoke tests")
+    p.add_argument("--max-samples", type=int, default=None)
     p.add_argument("--num-workers", type=int, default=4)
     p.add_argument("--ckpt-dir", default="checkpoints")
+    p.add_argument("--distance", action="store_true",
+                   help="v2: add a distance-to-transmitter input channel")
+    p.add_argument("--compile", action="store_true",
+                   help="wrap model in torch.compile (optional; flaky on Windows)")
     args = p.parse_args()
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
     amp = (device == "cuda")
-    print(f"Device: {device} | AMP: {amp}")
+    print(f"Device: {device} | AMP: {amp} | distance: {args.distance} | compile: {args.compile}")
 
     train_ds = RadioMapSeerDataset(split="train", img_size=args.img_size,
-                                   max_samples=args.max_samples)
+                                   max_samples=args.max_samples, add_distance=args.distance)
     val_max = None if args.max_samples is None else max(20, args.max_samples // 5)
     val_ds = RadioMapSeerDataset(split="val", img_size=args.img_size,
-                                 max_samples=val_max)
+                                 max_samples=val_max, add_distance=args.distance)
     print(f"Train: {len(train_ds)} | Val: {len(val_ds)}")
 
     dl_kwargs = dict(batch_size=args.batch_size, num_workers=args.num_workers,
@@ -63,7 +71,9 @@ def main():
     train_loader = DataLoader(train_ds, shuffle=True, **dl_kwargs)
     val_loader = DataLoader(val_ds, shuffle=False, **dl_kwargs)
 
-    model = UNet(in_channels=2, out_channels=1, base=args.base).to(device)
+    in_ch = 3 if args.distance else 2
+    model = UNet(in_channels=in_ch, out_channels=1, base=args.base).to(device)
+    net = torch.compile(model) if args.compile else model
     criterion = nn.MSELoss()
     optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
     scaler = torch.amp.GradScaler(enabled=amp)
@@ -72,20 +82,19 @@ def main():
     ckpt_dir.mkdir(exist_ok=True)
     best_val = float("inf")
 
-    import csv
     history_path = ckpt_dir / "history.csv"
     with open(history_path, "w", newline="") as f:
         csv.writer(f).writerow(["epoch", "train_mse", "val_mse"])
 
     for epoch in range(1, args.epochs + 1):
-        model.train()
+        net.train()
         t0 = time.time()
         running, n = 0.0, 0
         for i, (x, y) in enumerate(train_loader, 1):
             x, y = x.to(device), y.to(device)
             optimizer.zero_grad(set_to_none=True)
             with torch.autocast(device_type=device, dtype=torch.float16, enabled=amp):
-                loss = criterion(model(x), y)
+                loss = criterion(net(x), y)
             scaler.scale(loss).backward()
             scaler.step(optimizer)
             scaler.update()
@@ -95,7 +104,7 @@ def main():
                 print(f"  epoch {epoch} step {i}/{len(train_loader)} | loss {running/n:.5f}")
 
         train_loss = running / n
-        val_loss = evaluate(model, val_loader, criterion, device, amp)
+        val_loss = evaluate(net, val_loader, criterion, device, amp)
         dt = time.time() - t0
         print(f"[epoch {epoch}/{args.epochs}] train MSE {train_loss:.5f} | "
               f"val MSE {val_loss:.5f} | val RMSE {val_loss**0.5:.5f} | {dt:.0f}s")

src/visualize.py

@@ -31,14 +31,17 @@ def main():
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
     ckpt = torch.load(args.ckpt, map_location=device, weights_only=False)
-    base = ckpt.get("args", {}).get("base", 64)
-    model = UNet(in_channels=2, out_channels=1, base=base).to(device)
+    cargs = ckpt.get("args", {})
+    base = cargs.get("base", 64)
+    add_dist = bool(cargs.get("distance", False))
+    in_ch = 3 if add_dist else 2
+    model = UNet(in_channels=in_ch, out_channels=1, base=base).to(device)
     model.load_state_dict(ckpt["model_state"])
     model.eval()
-    print(f"Loaded {args.ckpt} (epoch {ckpt.get('epoch', '?')}, "
-          f"val MSE {ckpt.get('val_loss', float('nan')):.5f})")
+    print(f"Loaded {args.ckpt} | in_channels={in_ch} | distance={add_dist} | "
+          f"val MSE {ckpt.get('val_loss', float('nan')):.5f}")
 
-    ds = RadioMapSeerDataset(split="test", img_size=args.img_size)  # unseen cities
+    ds = RadioMapSeerDataset(split="test", img_size=args.img_size, add_distance=add_dist)
     rng = np.random.default_rng(args.seed)
     idxs = rng.choice(len(ds), size=args.n, replace=False)