| """ |
| AirTrackLM - Uncertainty Methods |
| ================================= |
| Multiple uncertainty quantification approaches for trajectory prediction. |
| |
| 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 - 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 |
| import torch |
| import torch.nn as nn |
| from typing import Dict, List, Optional, Tuple |
| from dataclasses import dataclass |
|
|
|
|
| def uncertainty_kinematic_variance(cog, sog, rot, alt_rate, window=5): |
| N = len(cog) |
| scores = np.zeros(N) |
| 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) |
| 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 |
| 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 |
|
|
|
|
| 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): |
| 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] |
| 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 = np.sqrt( |
| (pred_e - east[i + horizon])**2 + |
| (pred_n - north[i + horizon])**2 + |
| (pred_u - up[i + horizon])**2 |
| ) |
| scores[i] = residual |
| |
| if N > horizon + 1: |
| scores[0] = scores[1] |
| scores[N-horizon:] = scores[N-horizon-1] |
| return scores |
|
|
|
|
| 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: |
| 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]) |
| 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) |
| scores[i] = 1.0 / max(count, 1) |
| else: |
| 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 |
| if N > 1: |
| scores[0] = scores[1] |
| return scores |
|
|
|
|
| def uncertainty_flight_phase_entropy(sog, alt_rate, up, window=10): |
| N = len(sog) |
| phases = np.zeros(N, dtype=np.int64) |
| for i in range(N): |
| if sog[i] < 30 and up[i] < 500: |
| phases[i] = 0 |
| elif alt_rate[i] > 300: |
| phases[i] = 1 |
| elif alt_rate[i] < -300: |
| phases[i] = 2 |
| elif abs(alt_rate[i]) <= 300 and up[i] > 5000: |
| phases[i] = 3 |
| elif alt_rate[i] < -100 and up[i] < 3000 and sog[i] < 200: |
| phases[i] = 4 |
| else: |
| 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] |
| 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 |
| return scores |
|
|
|
|
| @dataclass |
| class UncertaintyConfig: |
| 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 |
| n_bins: int = 16 |
| window: int = 5 |
| |
| @property |
| def n_methods(self): |
| 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, 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) |
| 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) |
| return results |
|
|
|
|
| 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: |
| 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) |
|
|
|
|
| class HeteroscedasticHead(nn.Module): |
| def __init__(self, d_model, n_outputs=6): |
| super().__init__() |
| self.log_var_head = nn.Sequential( |
| nn.Linear(d_model, d_model // 2), |
| nn.GELU(), |
| nn.Linear(d_model // 2, n_outputs), |
| ) |
| |
| def forward(self, hidden_states): |
| return torch.clamp(self.log_var_head(hidden_states), -5.0, 5.0) |
|
|
|
|
| class MultiUncertaintyEmbedding(nn.Module): |
| def __init__(self, d_model, n_methods, n_bins=16): |
| super().__init__() |
| self.n_methods = n_methods |
| self.n_bins = n_bins |
| 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): |
| B, L, M = uncert_bins.shape |
| embeds = [self.embeds[m](uncert_bins[:, :, m]) for m in range(M)] |
| if M == 1: |
| return embeds[0] |
| 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) |
|
|
