Add uncertainty.py

uncertainty.py

@@ -0,0 +1,585 @@
+"""
+AirTrackLM - Uncertainty Methods
+=================================
+Multiple uncertainty quantification approaches for trajectory prediction.
+
+Methods implemented:
+1. Kinematic Variance (original) - sliding window variance of COG/SOG/ROT/alt_rate
+2. Prediction Residual - deviation from constant-velocity prediction model
+3. Spatial Density - how many other trajectory points exist nearby (data coverage proxy)
+4. Flight Phase Entropy - uncertainty from ambiguity in flight phase classification
+5. Temporal Irregularity - variance in time gaps between raw ADS-B messages
+6. Learned Heteroscedastic - model predicts its own uncertainty (aleatoric)
+7. MC-Dropout - Monte Carlo dropout at inference (epistemic, model-level)
+
+Methods 1-5 produce per-timestep scalar scores at preprocessing time.
+Methods 6-7 are model-architecture modifications used during training/inference.
+
+All preprocessing methods are discretized into N bins and embedded as learned embeddings.
+The model selects which uncertainty method(s) to use via config.
+"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass
+
+
+# ============================================================
+# 1. Kinematic Variance (original method)
+# ============================================================
+
+def uncertainty_kinematic_variance(
+    cog: np.ndarray,
+    sog: np.ndarray,
+    rot: np.ndarray,
+    alt_rate: np.ndarray,
+    window: int = 5,
+) -> np.ndarray:
+    """
+    Sliding-window variance of kinematic features.
+    High variance = maneuvering = high uncertainty.
+    
+    This is a measure of how unpredictable the trajectory has been
+    in the recent past.
+    """
+    N = len(cog)
+    scores = np.zeros(N)
+    
+    # Pre-compute global variances for normalization
+    global_sog_var = max(np.var(sog), 1e-10)
+    global_rot_var = max(np.var(rot), 1e-10)
+    global_alt_var = max(np.var(alt_rate), 1e-10)
+    
+    for i in range(N):
+        start = max(0, i - window + 1)
+        w = slice(start, i + 1)
+        
+        # Circular variance for COG
+        cog_rad = np.radians(cog[w])
+        R_len = np.sqrt(np.mean(np.cos(cog_rad))**2 + np.mean(np.sin(cog_rad))**2)
+        cog_var = 1 - R_len  # [0, 1]
+        
+        # Normalized variances for others
+        sog_var = np.var(sog[w]) / global_sog_var if len(sog[w]) > 1 else 0
+        rot_var = np.var(rot[w]) / global_rot_var if len(rot[w]) > 1 else 0
+        alt_var = np.var(alt_rate[w]) / global_alt_var if len(alt_rate[w]) > 1 else 0
+        
+        scores[i] = cog_var + sog_var + rot_var + alt_var
+    
+    return scores
+
+
+# ============================================================
+# 2. Prediction Residual Uncertainty
+# ============================================================
+
+def uncertainty_prediction_residual(
+    east: np.ndarray,
+    north: np.ndarray,
+    up: np.ndarray,
+    timestamps: np.ndarray,
+    horizon: int = 3,
+) -> np.ndarray:
+    """
+    Constant-velocity prediction residual.
+    
+    At each time step t, we use the velocity at t to predict where
+    the aircraft will be at t+horizon. The error between this prediction
+    and the actual position at t+horizon measures how well a simple
+    kinematic model captures the motion.
+    
+    High residual = the aircraft is doing something unexpected = high uncertainty.
+    
+    This captures maneuver onset, turbulence, and ATC-induced deviations.
+    """
+    N = len(east)
+    scores = np.zeros(N)
+    
+    dt = np.diff(timestamps)
+    dt = np.maximum(dt, 1e-6)
+    
+    for i in range(1, N - horizon):
+        # Velocity at time i (backward difference)
+        vx = (east[i] - east[i-1]) / dt[i-1]
+        vy = (north[i] - north[i-1]) / dt[i-1]
+        vz = (up[i] - up[i-1]) / dt[i-1]
+        
+        # Predict position at i+horizon using constant velocity
+        dt_pred = timestamps[i + horizon] - timestamps[i]
+        pred_e = east[i] + vx * dt_pred
+        pred_n = north[i] + vy * dt_pred
+        pred_u = up[i] + vz * dt_pred
+        
+        # Residual (3D Euclidean distance in meters)
+        residual = np.sqrt(
+            (pred_e - east[i + horizon])**2 +
+            (pred_n - north[i + horizon])**2 +
+            (pred_u - up[i + horizon])**2
+        )
+        scores[i] = residual
+    
+    # Fill edges with nearest computed value
+    if N > horizon + 1:
+        scores[0] = scores[1]
+        scores[N-horizon:] = scores[N-horizon-1]
+    
+    return scores
+
+
+# ============================================================
+# 3. Spatial Density Uncertainty
+# ============================================================
+
+def uncertainty_spatial_density(
+    east: np.ndarray,
+    north: np.ndarray,
+    up: np.ndarray,
+    all_trajectories_enu: Optional[List[Tuple[np.ndarray, np.ndarray, np.ndarray]]] = None,
+    radius_m: float = 5000.0,
+) -> np.ndarray:
+    """
+    Spatial data density proxy.
+    
+    For each point, count how many training trajectory points exist
+    within `radius_m` meters. Low density = underrepresented region = 
+    high uncertainty (the model has seen few examples of this area).
+    
+    This is a data-centric uncertainty measure — it tells you where the
+    model is likely to be poorly calibrated due to lack of training data.
+    
+    If all_trajectories_enu is None, uses self-density (how spread out
+    this trajectory's own points are, inverse of local point density).
+    """
+    N = len(east)
+    scores = np.zeros(N)
+    
+    if all_trajectories_enu is not None:
+        # Build a flat array of all training points
+        all_e = np.concatenate([t[0] for t in all_trajectories_enu])
+        all_n = np.concatenate([t[1] for t in all_trajectories_enu])
+        all_u = np.concatenate([t[2] for t in all_trajectories_enu])
+        
+        # For efficiency, subsample if too many points
+        if len(all_e) > 50000:
+            idx = np.random.choice(len(all_e), 50000, replace=False)
+            all_e, all_n, all_u = all_e[idx], all_n[idx], all_u[idx]
+        
+        for i in range(N):
+            dists = np.sqrt(
+                (all_e - east[i])**2 +
+                (all_n - north[i])**2 +
+                (all_u - up[i])**2
+            )
+            count = np.sum(dists < radius_m)
+            # Inverse density: fewer points nearby = higher uncertainty
+            scores[i] = 1.0 / max(count, 1)
+    else:
+        # Self-density: use spacing between consecutive points
+        # More spread out = higher velocity = potentially higher uncertainty
+        for i in range(1, N):
+            dist = np.sqrt(
+                (east[i] - east[i-1])**2 +
+                (north[i] - north[i-1])**2 +
+                (up[i] - up[i-1])**2
+            )
+            scores[i] = dist  # larger steps = less "dense" sampling = more uncertain
+        scores[0] = scores[1] if N > 1 else 0
+    
+    return scores
+
+
+# ============================================================
+# 4. Flight Phase Entropy
+# ============================================================
+
+def uncertainty_flight_phase_entropy(
+    sog: np.ndarray,
+    alt_rate: np.ndarray,
+    up: np.ndarray,
+    window: int = 10,
+) -> np.ndarray:
+    """
+    Entropy of flight phase classification within a window.
+    
+    We classify each point into a flight phase based on simple heuristics:
+      - CLIMB: alt_rate > 300 ft/min
+      - DESCENT: alt_rate < -300 ft/min
+      - CRUISE: alt_rate in [-300, 300] and altitude > 5000m
+      - APPROACH: alt_rate < -100 and altitude < 3000m and SOG < 200 kts
+      - GROUND: SOG < 30 kts and altitude < 500m
+      - MANEUVERING: everything else
+    
+    If the window has a mix of phases, the entropy is high → uncertain
+    about what the aircraft is doing → hard to predict next state.
+    """
+    N = len(sog)
+    
+    # Classify each point
+    phases = np.zeros(N, dtype=np.int64)
+    for i in range(N):
+        if sog[i] < 30 and up[i] < 500:
+            phases[i] = 0  # GROUND
+        elif alt_rate[i] > 300:
+            phases[i] = 1  # CLIMB
+        elif alt_rate[i] < -300:
+            phases[i] = 2  # DESCENT
+        elif abs(alt_rate[i]) <= 300 and up[i] > 5000:
+            phases[i] = 3  # CRUISE
+        elif alt_rate[i] < -100 and up[i] < 3000 and sog[i] < 200:
+            phases[i] = 4  # APPROACH
+        else:
+            phases[i] = 5  # MANEUVERING
+    
+    n_phases = 6
+    scores = np.zeros(N)
+    
+    for i in range(N):
+        start = max(0, i - window + 1)
+        w_phases = phases[start:i+1]
+        
+        # Compute entropy of phase distribution in window
+        counts = np.bincount(w_phases, minlength=n_phases).astype(float)
+        probs = counts / counts.sum()
+        probs = probs[probs > 0]
+        entropy = -np.sum(probs * np.log2(probs))
+        
+        scores[i] = entropy  # max entropy = log2(6) ≈ 2.58
+    
+    return scores
+
+
+# ============================================================
+# 5. Temporal Irregularity
+# ============================================================
+
+def uncertainty_temporal_irregularity(
+    raw_timestamps: np.ndarray,
+    resampled_timestamps: np.ndarray,
+    window: int = 10,
+) -> np.ndarray:
+    """
+    Variance of original (pre-resampling) time gaps mapped to resampled points.
+    
+    ADS-B messages arrive irregularly. When messages are sparse or bursty,
+    the resampled positions rely heavily on interpolation → higher uncertainty.
+    
+    We measure this by looking at how many raw messages fall near each
+    resampled point. Fewer raw messages = more interpolation = more uncertain.
+    """
+    N = len(resampled_timestamps)
+    scores = np.zeros(N)
+    
+    # For each resampled point, find nearest raw timestamp
+    for i in range(N):
+        t = resampled_timestamps[i]
+        # Find raw messages within a window around this resampled time
+        dt_half = (resampled_timestamps[min(i+1, N-1)] - resampled_timestamps[max(i-1, 0)]) / 2
+        dt_half = max(dt_half, 1.0)
+        
+        nearby = np.abs(raw_timestamps - t) < dt_half
+        count = nearby.sum()
+        
+        # Fewer raw messages nearby = higher uncertainty
+        scores[i] = 1.0 / max(count, 1)
+    
+    return scores
+
+
+# ============================================================
+# 6. Multi-method Uncertainty Combiner (Preprocessing)
+# ============================================================
+
+@dataclass
+class UncertaintyConfig:
+    """Configuration for which uncertainty methods to use."""
+    use_kinematic_variance: bool = True
+    use_prediction_residual: bool = True
+    use_spatial_density: bool = True
+    use_flight_phase_entropy: bool = True
+    use_temporal_irregularity: bool = False  # requires raw timestamps
+    
+    n_bins: int = 16  # bins per method
+    window: int = 5   # sliding window size
+    
+    @property
+    def n_methods(self) -> int:
+        """Number of active uncertainty methods."""
+        return sum([
+            self.use_kinematic_variance,
+            self.use_prediction_residual,
+            self.use_spatial_density,
+            self.use_flight_phase_entropy,
+            self.use_temporal_irregularity,
+        ])
+
+
+def compute_all_uncertainties(
+    east: np.ndarray,
+    north: np.ndarray,
+    up: np.ndarray,
+    timestamps: np.ndarray,
+    cog: np.ndarray,
+    sog: np.ndarray,
+    rot: np.ndarray,
+    alt_rate: np.ndarray,
+    config: UncertaintyConfig = UncertaintyConfig(),
+    raw_timestamps: Optional[np.ndarray] = None,
+    all_trajectories_enu: Optional[List] = None,
+) -> Dict[str, np.ndarray]:
+    """
+    Compute all configured uncertainty methods.
+    
+    Returns dict of method_name → raw scores (N,) arrays.
+    """
+    results = {}
+    
+    if config.use_kinematic_variance:
+        results['kinematic_var'] = uncertainty_kinematic_variance(
+            cog, sog, rot, alt_rate, window=config.window
+        )
+    
+    if config.use_prediction_residual:
+        results['pred_residual'] = uncertainty_prediction_residual(
+            east, north, up, timestamps, horizon=3
+        )
+    
+    if config.use_spatial_density:
+        results['spatial_density'] = uncertainty_spatial_density(
+            east, north, up, all_trajectories_enu
+        )
+    
+    if config.use_flight_phase_entropy:
+        results['phase_entropy'] = uncertainty_flight_phase_entropy(
+            sog, alt_rate, up, window=config.window * 2
+        )
+    
+    if config.use_temporal_irregularity and raw_timestamps is not None:
+        results['temporal_irreg'] = uncertainty_temporal_irregularity(
+            raw_timestamps, timestamps, window=config.window
+        )
+    
+    return results
+
+
+def discretize_scores(scores: np.ndarray, n_bins: int = 16) -> np.ndarray:
+    """Discretize continuous uncertainty scores into quantile bins."""
+    if len(np.unique(scores)) < n_bins:
+        edges = np.linspace(scores.min(), scores.max() + 1e-10, n_bins + 1)
+    else:
+        edges = np.quantile(scores, np.linspace(0, 1, n_bins + 1))
+        edges[-1] += 1e-10
+    return np.clip(np.digitize(scores, edges) - 1, 0, n_bins - 1)
+
+
+# ============================================================
+# 7. Learned Heteroscedastic Uncertainty (Model Component)
+# ============================================================
+
+class HeteroscedasticHead(nn.Module):
+    """
+    Predicts per-timestep aleatoric uncertainty (log-variance) alongside
+    the main predictions. Used in the loss function to weight samples
+    by their predicted uncertainty.
+    
+    Based on Kendall & Gal (2017) "What Uncertainties Do We Need in 
+    Bayesian Deep Learning for Computer Vision?"
+    
+    The model learns to predict log(σ²) for each output, and the loss
+    becomes: L = (1/2σ²) * ||y - ŷ||² + (1/2) * log(σ²)
+    
+    This lets the model say "I'm unsure about this prediction" and
+    reduce the gradient signal from noisy/ambiguous samples.
+    """
+    
+    def __init__(self, d_model: int, n_outputs: int = 6):
+        super().__init__()
+        # Predict log-variance for each output head
+        self.log_var_head = nn.Sequential(
+            nn.Linear(d_model, d_model // 2),
+            nn.GELU(),
+            nn.Linear(d_model // 2, n_outputs),
+        )
+        # n_outputs: [geohash, cog, sog, rot, alt_rate, continuous]
+    
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: (B, L, d_model)
+        Returns:
+            log_var: (B, L, n_outputs) — predicted log-variance per head
+        """
+        return self.log_var_head(hidden_states)
+
+
+# ============================================================
+# 8. MC-Dropout Module
+# ============================================================
+
+class MCDropoutWrapper(nn.Module):
+    """
+    Wrapper that enables dropout at inference time for Monte Carlo
+    uncertainty estimation.
+    
+    Usage:
+        model = MCDropoutWrapper(base_model, n_samples=10)
+        predictions, uncertainty = model.predict_with_uncertainty(batch)
+    
+    The uncertainty is the variance across multiple stochastic forward passes.
+    """
+    
+    def __init__(self, model: nn.Module, n_samples: int = 10):
+        super().__init__()
+        self.model = model
+        self.n_samples = n_samples
+    
+    def enable_dropout(self):
+        """Enable dropout layers during inference."""
+        for module in self.model.modules():
+            if isinstance(module, nn.Dropout):
+                module.train()
+    
+    @torch.no_grad()
+    def predict_with_uncertainty(
+        self, batch: Dict[str, torch.Tensor]
+    ) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        """
+        Run multiple forward passes with dropout enabled.
+        
+        Returns:
+            mean_predictions: dict of mean logits per head
+            uncertainty: dict of variance per head (epistemic uncertainty)
+        """
+        self.model.eval()
+        self.enable_dropout()
+        
+        all_predictions = []
+        for _ in range(self.n_samples):
+            pred = self.model(batch)
+            all_predictions.append(pred)
+        
+        # Compute mean and variance across samples
+        mean_predictions = {}
+        uncertainty = {}
+        
+        for key in all_predictions[0].keys():
+            stacked = torch.stack([p[key] for p in all_predictions], dim=0)  # (n_samples, B, L, ...)
+            mean_predictions[key] = stacked.mean(dim=0)
+            uncertainty[key] = stacked.var(dim=0)  # epistemic uncertainty
+        
+        return mean_predictions, uncertainty
+
+
+# ============================================================
+# 9. Multi-Method Uncertainty Embedding (Model Component)
+# ============================================================
+
+class MultiUncertaintyEmbedding(nn.Module):
+    """
+    Embeds multiple uncertainty signals and combines them.
+    
+    Each preprocessing-based uncertainty method gets its own embedding table.
+    The embeddings are summed (like the main feature embeddings).
+    
+    Optionally includes the heteroscedastic head for learned uncertainty.
+    """
+    
+    def __init__(self, d_model: int, n_methods: int, n_bins: int = 16):
+        super().__init__()
+        self.n_methods = n_methods
+        self.n_bins = n_bins
+        
+        # One embedding table per uncertainty method
+        self.embeds = nn.ModuleList([
+            nn.Embedding(n_bins, d_model) for _ in range(n_methods)
+        ])
+        
+        # Learned combination weights (attention over methods)
+        self.method_attention = nn.Sequential(
+            nn.Linear(d_model * n_methods, n_methods),
+            nn.Softmax(dim=-1),
+        )
+    
+    def forward(self, uncert_bins: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            uncert_bins: (B, L, n_methods) long — bin indices per method
+        Returns:
+            (B, L, d_model) — combined uncertainty embedding
+        """
+        B, L, M = uncert_bins.shape
+        assert M == self.n_methods, f"Expected {self.n_methods} methods, got {M}"
+        
+        # Embed each method
+        embeds = []
+        for m in range(M):
+            embeds.append(self.embeds[m](uncert_bins[:, :, m]))  # (B, L, d_model)
+        
+        if M == 1:
+            return embeds[0]
+        
+        # Learned attention weighting
+        concat = torch.cat(embeds, dim=-1)  # (B, L, d_model * M)
+        weights = self.method_attention(concat)  # (B, L, M)
+        
+        # Weighted sum
+        stacked = torch.stack(embeds, dim=-1)  # (B, L, d_model, M)
+        weighted = (stacked * weights.unsqueeze(2)).sum(dim=-1)  # (B, L, d_model)
+        
+        return weighted
+
+
+# ============================================================
+# 10. Heteroscedastic Loss
+# ============================================================
+
+class HeteroscedasticLoss(nn.Module):
+    """
+    Attenuated loss that uses predicted log-variance to weight samples.
+    
+    For regression: L = (1/2) * exp(-s) * ||y - ŷ||² + (1/2) * s
+    For classification: L = exp(-s) * CE(y, ŷ) + (1/2) * s
+    
+    where s = log(σ²) is the predicted log-variance.
+    
+    This encourages the model to predict high uncertainty for difficult
+    samples and low uncertainty for easy ones.
+    """
+    
+    def __init__(self):
+        super().__init__()
+        self.ce = nn.CrossEntropyLoss(reduction='none')
+        self.bce = nn.BCEWithLogitsLoss(reduction='none')
+    
+    def classification_loss(
+        self, logits: torch.Tensor, targets: torch.Tensor, log_var: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Args:
+            logits: (B*L, n_classes)
+            targets: (B*L,) long
+            log_var: (B*L,) — predicted log-variance
+        Returns:
+            scalar loss
+        """
+        ce = self.ce(logits, targets)  # (B*L,)
+        precision = torch.exp(-log_var)
+        loss = precision * ce + 0.5 * log_var
+        return loss.mean()
+    
+    def regression_loss(
+        self, pred: torch.Tensor, target: torch.Tensor, log_var: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Args:
+            pred: (B, L, D)
+            target: (B, L, D)
+            log_var: (B, L) — predicted log-variance
+        Returns:
+            scalar loss
+        """
+        mse = ((pred - target) ** 2).sum(dim=-1)  # (B, L)
+        precision = torch.exp(-log_var)
+        loss = 0.5 * precision * mse + 0.5 * log_var
+        return loss.mean()