Upload uncertainty.py with huggingface_hub

uncertainty.py

@@ -3,20 +3,13 @@ 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
+Methods:
+1. Kinematic Variance - 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.
+3. Spatial Density - data coverage proxy
+4. Flight Phase Entropy - entropy of phase classification in a window
+5. Learned Heteroscedastic - model predicts its own uncertainty (aleatoric)
+6. MC-Dropout - Monte Carlo dropout at inference (epistemic)
 """
 
 import numpy as np
@@ -26,28 +19,9 @@ 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.
-    """
+def uncertainty_kinematic_variance(cog, sog, rot, alt_rate, window=5):
     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)
@@ -55,64 +29,30 @@ def uncertainty_kinematic_variance(
     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
+        cog_var = 1 - R_len
         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.
-    """
+def uncertainty_prediction_residual(east, north, up, timestamps, horizon=3):
     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 +
@@ -120,192 +60,78 @@ def uncertainty_prediction_residual(
         )
         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).
-    """
+def uncertainty_spatial_density(east, north, up, all_trajectories_enu=None, radius_m=5000.0):
     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
-            )
+            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
-    
+            dist = np.sqrt((east[i]-east[i-1])**2 + (north[i]-north[i-1])**2 + (up[i]-up[i-1])**2)
+            scores[i] = dist
+        if N > 1:
+            scores[0] = scores[1]
     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.
-    """
+def uncertainty_flight_phase_entropy(sog, alt_rate, up, window=10):
     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
+            phases[i] = 0
         elif alt_rate[i] > 300:
-            phases[i] = 1  # CLIMB
+            phases[i] = 1
         elif alt_rate[i] < -300:
-            phases[i] = 2  # DESCENT
+            phases[i] = 2
         elif abs(alt_rate[i]) <= 300 and up[i] > 5000:
-            phases[i] = 3  # CRUISE
+            phases[i] = 3
         elif alt_rate[i] < -100 and up[i] < 3000 and sog[i] < 200:
-            phases[i] = 4  # APPROACH
+            phases[i] = 4
         else:
-            phases[i] = 5  # MANEUVERING
+            phases[i] = 5
     
     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)
-    
+        scores[i] = entropy
     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
+    use_temporal_irregularity: bool = False
+    n_bins: int = 16
+    window: int = 5
     
     @property
-    def n_methods(self) -> int:
-        """Number of active uncertainty methods."""
+    def n_methods(self):
         return sum([
             self.use_kinematic_variance,
             self.use_prediction_residual,
@@ -315,56 +141,23 @@ class UncertaintyConfig:
         ])
 
 
-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.
-    """
+def compute_all_uncertainties(east, north, up, timestamps, cog, sog, rot, alt_rate,
+                              config=None, raw_timestamps=None, all_trajectories_enu=None):
+    if config is None:
+        config = UncertaintyConfig()
     results = {}
-    
     if config.use_kinematic_variance:
-        results['kinematic_var'] = uncertainty_kinematic_variance(
-            cog, sog, rot, alt_rate, window=config.window
-        )
-    
+        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
-        )
-    
+        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
-        )
-    
+        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
-        )
-    
+        results['phase_entropy'] = uncertainty_flight_phase_entropy(sog, alt_rate, up, window=config.window * 2)
     return results
 
 
-def discretize_scores(scores: np.ndarray, n_bins: int = 16) -> np.ndarray:
-    """Discretize continuous uncertainty scores into quantile bins."""
+def discretize_scores(scores, n_bins=16):
     if len(np.unique(scores)) < n_bins:
         edges = np.linspace(scores.min(), scores.max() + 1e-10, n_bins + 1)
     else:
@@ -373,213 +166,37 @@ def discretize_scores(scores: np.ndarray, n_bins: int = 16) -> np.ndarray:
     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):
+    def __init__(self, d_model, n_outputs=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, clamped to [-5, 5]
-        """
+    def forward(self, hidden_states):
         return torch.clamp(self.log_var_head(hidden_states), -5.0, 5.0)
 
 
-# ============================================================
-# 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):
+    def __init__(self, d_model, n_methods, n_bins=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),
-        )
+        self.embeds = nn.ModuleList([nn.Embedding(n_bins, d_model) for _ in range(n_methods)])
+        if n_methods > 1:
+            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
-        """
+    def forward(self, uncert_bins):
         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)
-        
+        embeds = [self.embeds[m](uncert_bins[:, :, m]) for m in range(M)]
         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()
+        concat = torch.cat(embeds, dim=-1)
+        weights = self.method_attention(concat)
+        stacked = torch.stack(embeds, dim=-1)
+        return (stacked * weights.unsqueeze(2)).sum(dim=-1)