Fix GTA-UAV evaluation and loss (critical: false negatives + wrong R@K)

PROBLEM: GTA-UAV has overlapping satellite crops (partial IoU).
Standard InfoNCE with diagonal targets treated valid matches as negatives.
R@K checked only diagonal — missed valid matches, artificially low recall.

FIXES:
1. WeightedInfoNCE loss (src/losses/weighted_infonce.py):
   - Per-sample adaptive label smoothing from positive_weights (IoU)
   - Higher weight → sharper target, lower → softer (semi-positive tolerance)
   - Based on Game4Loc reference implementation

2. Multi-match R@K evaluation:
   - Uses dataset.get_all_valid_sat_names() to get ALL valid matches per query
   - R@K counts hit if ANY valid satellite is in top-K (not just diagonal)
   - AP computed as MRR over first valid match

3. Dataset returns positive_weight per sample:
   - Sampled satellite weight passed to loss for adaptive smoothing
   - All valid satellite candidates exposed for evaluation

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

src/losses/weighted_infonce.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+"""Weighted InfoNCE loss for GTA-UAV cross-view geo-localization.
+
+Adapted from Game4Loc (https://github.com/Yux1angJi/GTA-UAV).
+Uses per-sample label smoothing based on positive_weights (IoU/distance)
+to handle partial overlap between drone and satellite crops.
+
+Standard InfoNCE assumes strict 1-to-1 pairs and treats all non-diagonal
+entries as negatives. In GTA-UAV, multiple satellite crops can validly
+match one drone image (partial IoU overlap), causing false negatives.
+WeightedInfoNCE softens this with adaptive label smoothing per sample.
+"""
+
+import math
+
+import gin
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@gin.configurable
+class WeightedInfoNCELoss(nn.Module):
+    """Weighted InfoNCE with adaptive per-sample label smoothing.
+
+    For each sample i, eps_i = 1 - (1 - base_smoothing) / (1 + exp(-k * w_i))
+    where w_i is the positive weight (e.g. IoU with matched satellite crop).
+    Higher weight → lower eps → sharper target (strong positive).
+    Lower weight → higher eps → softer target (weak/semi-positive).
+
+    Args:
+        temperature_init: Initial temperature (or learnable logit_scale).
+        learnable_temperature: If True, temperature is learnable (CLIP-style).
+        label_smoothing: Base label smoothing (used when no weights provided).
+        k: Sigmoid steepness for weight → eps mapping.
+        tau_min: Min clamp for learnable temperature.
+        tau_max: Max clamp for learnable temperature.
+    """
+
+    def __init__(
+        self,
+        temperature_init: float = 0.07,
+        learnable_temperature: bool = True,
+        label_smoothing: float = 0.1,
+        k: float = 5.0,
+        tau_min: float = 0.01,
+        tau_max: float = 0.5,
+    ) -> None:
+        super().__init__()
+        self.label_smoothing = label_smoothing
+        self.k = k
+        self.tau_min = tau_min
+        self.tau_max = tau_max
+        self.learnable_temperature = learnable_temperature
+
+        if learnable_temperature:
+            self.logit_scale = nn.Parameter(
+                torch.tensor(math.log(1.0 / temperature_init))
+            )
+        else:
+            self.logit_scale = None
+            self.temperature = temperature_init
+
+    @property
+    def current_temperature(self) -> float:
+        if self.logit_scale is not None:
+            tau = 1.0 / self.logit_scale.exp().clamp(
+                min=1.0 / self.tau_max, max=1.0 / self.tau_min,
+            ).item()
+            return tau
+        return self.temperature
+
+    def _compute_eps(self, positive_weights: torch.Tensor | None, n: int) -> torch.Tensor | list[float]:
+        """Compute per-sample label smoothing from positive weights."""
+        if positive_weights is not None:
+            # Higher weight → lower eps (sharper, stronger positive).
+            return 1.0 - (1.0 - self.label_smoothing) / (1.0 + torch.exp(-self.k * positive_weights))
+        return [self.label_smoothing] * n
+
+    def _weighted_loss(
+        self,
+        sim_matrix: torch.Tensor,
+        eps_all: torch.Tensor | list[float],
+    ) -> torch.Tensor:
+        """Weighted InfoNCE: per-sample interpolation between hard and uniform targets.
+
+        For each row i:
+          L_i = (1-eps_i) * [-sim[i,i] + logsumexp(sim[i,:])]
+              + eps_i * [-mean(sim[i,:]) + logsumexp(sim[i,:])]
+        """
+        n = sim_matrix.shape[0]
+        total_loss = torch.tensor(0.0, device=sim_matrix.device)
+        for i in range(n):
+            eps = eps_all[i] if isinstance(eps_all, list) else eps_all[i]
+            logsumexp = torch.logsumexp(sim_matrix[i, :], dim=0)
+            total_loss += (1 - eps) * (-sim_matrix[i, i] + logsumexp)
+            total_loss += eps * (-sim_matrix[i, :].mean() + logsumexp)
+        return total_loss / n
+
+    def forward(
+        self,
+        embeddings: dict[str, torch.Tensor],
+        epoch: int = 0,
+        total_epochs: int = 1,
+        positive_weights: torch.Tensor | None = None,
+        queue_negatives: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Compute weighted InfoNCE loss.
+
+        Args:
+            embeddings: Dict with 'query' [B,D], 'gallery' [B,D], 'gate_q', 'gate_g'.
+            positive_weights: Per-sample weight [B] (e.g. IoU with matched sat crop).
+            queue_negatives: Extra negatives [Q,D] from memory bank (not used with weighted loss).
+        """
+        query = embeddings["query"].float()
+        gallery = embeddings["gallery"].float()
+
+        # Temperature.
+        if self.learnable_temperature:
+            clamped = self.logit_scale.float().clamp(
+                min=math.log(1.0 / self.tau_max),
+                max=math.log(1.0 / self.tau_min),
+            )
+            logit_scale = clamped.exp()
+            tau = 1.0 / logit_scale
+        else:
+            logit_scale = 1.0 / self.temperature
+            tau = self.temperature
+
+        sim_q2g = logit_scale * query @ gallery.t()
+        sim_g2q = sim_q2g.t()
+
+        eps = self._compute_eps(positive_weights, query.shape[0])
+
+        loss_q2g = self._weighted_loss(sim_q2g, eps)
+        loss_g2q = self._weighted_loss(sim_g2q, eps)
+        total = (loss_q2g + loss_g2q) / 2
+
+        gate_q = embeddings.get("gate_q", 1.0)
+        gate_g = embeddings.get("gate_g", 1.0)
+
+        return {
+            "total": total,
+            "temperature": tau if isinstance(tau, torch.Tensor) else torch.tensor(tau, device=total.device),
+            "gate_q": torch.tensor(gate_q, device=total.device) if not isinstance(gate_q, torch.Tensor) else gate_q.detach(),
+            "gate_g": torch.tensor(gate_g, device=total.device) if not isinstance(gate_g, torch.Tensor) else gate_g.detach(),
+        }