#!/usr/bin/env python3 """Causal Trackformer1.1 intensity and wind-structure inference. The Trackformer1.1 route is responsible for position and the western-Pacific pressure state. This module supplies the previously missing storm-structure outputs without changing that route: maximum sustained wind, central pressure, radius of maximum wind, and the four-quadrant R34/R50/R64 radii. The default weights are a validated residual-anchor spatial ensemble with a secondary structure expert and causal temporal branch. They are loaded only for inference and are conditioned on the same nine-step observed track window plus the current four-channel analysis patch. No positive-lead atmospheric field or official agency forecast is consumed. """ from __future__ import annotations import json import os from pathlib import Path import numpy as np import torch import torch.nn as nn LEADS = 20 NM_TO_KM = 1.852 TARGET_SCALE = np.asarray([100.0, 100.0, 35.0, 20.0, 50.0] + [50.0] * 12, dtype="float32") STRUCTURE_SCALE = TARGET_SCALE[2:] THERMO_ENV_COLS = ( [4, 5, 6, 7] + list(range(8, 20)) + list(range(24, 40)) + [44, 45, 46, 47, 48, 49, 50, 51, 52, 53] ) # The 90th percentile of positive six-hour central-pressure changes in the # training split is 6 hPa. Use that train-only statistic to stop a coarse # forecast-map minimum from disappearing in one step and being replaced by a # different synoptic cell. MAP_PRESSURE_RECOVERY_LIMIT_HPA = 6.0 def sinusoidal(length: int, width: int) -> torch.Tensor: position = torch.arange(length).unsqueeze(1).float() divisor = torch.exp(torch.arange(0, width, 2).float() * (-np.log(10000.0) / width)) result = torch.zeros(length, width) result[:, 0::2] = torch.sin(position * divisor) result[:, 1::2] = torch.cos(position * divisor) return result class StructureSpatialExpert(nn.Module): """Inference copy of the architecture used by the frozen experts.""" def __init__(self, width: int, layers: int, heads: int, structure_residual: bool = False): super().__init__() self.structure_residual = bool(structure_residual) self.track_proj = nn.Linear(len(THERMO_ENV_COLS), width) self.register_buffer("track_time", sinusoidal(9, width).unsqueeze(0)) track_layer = nn.TransformerEncoderLayer( width, heads, width * 4, 0.12, batch_first=True, norm_first=True, activation="gelu", ) self.track_encoder = nn.TransformerEncoder(track_layer, layers) self.field_encoder = nn.Sequential( nn.Conv2d(4, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.SiLU(), nn.Conv2d(64, width, 3, stride=2, padding=1), nn.GroupNorm(8, width), nn.SiLU(), ) self.field_pool = nn.AvgPool2d(kernel_size=2, stride=2) self.field_norm = nn.LayerNorm(width) self.field_pos = nn.Parameter(torch.randn(1, 16, width) * 0.02) decoder_layer = nn.TransformerDecoderLayer( width, heads, width * 4, 0.12, batch_first=True, norm_first=True, activation="gelu", ) self.query = nn.Parameter(torch.randn(1, LEADS, width) * 0.02) self.register_buffer("lead_time", sinusoidal(LEADS, width).unsqueeze(0)) self.decoder = nn.TransformerDecoder(decoder_layer, layers) self.state = nn.Linear(width, 15) self.log_scale = nn.Linear(width, 15) def forward( self, track: torch.Tensor, field: torch.Tensor, current: torch.Tensor, available: torch.Tensor, current_structure: torch.Tensor | None = None, structure_available: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: track_tokens = self.track_encoder(self.track_proj(track[:, :, THERMO_ENV_COLS]) + self.track_time) field_tokens = self.field_pool(self.field_encoder(field)).flatten(2).transpose(1, 2) field_tokens = self.field_norm(field_tokens + self.field_pos) memory = torch.cat([track_tokens, field_tokens], dim=1) query = (self.query + self.lead_time).expand(track.shape[0], -1, -1) hidden = self.decoder(query, memory) state = self.state(hidden) state = state.clone() state[:, :, :2] = state[:, :, :2] + (current * available)[:, None, :] if self.structure_residual: if current_structure is None or structure_available is None: raise ValueError("structure residual mode requires current structure and availability tensors") state[:, :, 2:] = state[:, :, 2:] + current_structure[:, None, :] * structure_available[:, None, :] return state, self.log_scale(hidden) def _device(requested: str | None) -> torch.device: value = requested or os.environ.get("TRACKFORMER_1_1_DEVICE") if value: return torch.device(value) if torch.cuda.is_available(): return torch.device("cuda") if torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") def _calibrated_wind(states: np.ndarray, current_wind: float, calibration: dict) -> np.ndarray: alphas = calibration.get("wind_blend_alpha") if not alphas: return states[:, :, 0] predicted = states[:, :, 0] result = np.empty_like(predicted) for lead, alpha in enumerate(alphas[:LEADS]): result[:, lead] = float(alpha) * predicted[:, lead] + (1.0 - float(alpha)) * current_wind return np.clip(result, 0.0, 190.0) def _calibrated_pressure( states: np.ndarray, current_wind: float, current_pressure: float, previous_wind: float, previous_pressure: float, calibrated_wind: np.ndarray, calibration: dict, ) -> np.ndarray: joint = calibration.get("pressure_joint_calibrations") if not joint: return np.clip(states[:, :, 1], 850.0, 1025.0) predicted_pressure = states[:, :, 1] result = np.empty_like(predicted_pressure) for lead, item in enumerate(joint[:LEADS]): features = np.column_stack([ np.full(len(predicted_pressure), 1.0, dtype="float32"), np.full(len(predicted_pressure), current_pressure, dtype="float32"), np.full(len(predicted_pressure), current_wind, dtype="float32"), calibrated_wind[:, lead], predicted_pressure[:, lead], predicted_pressure[:, lead] - current_pressure, calibrated_wind[:, lead] - current_wind, np.full(len(predicted_pressure), current_pressure - previous_pressure, dtype="float32"), np.full(len(predicted_pressure), current_wind - previous_wind, dtype="float32"), ]) mean = np.asarray(item["mean"], dtype="float32") scale = np.maximum(np.asarray(item["scale"], dtype="float32"), 1e-6) normalized = (features - mean) / scale normalized[:, 0] = 1.0 result[:, lead] = normalized @ np.asarray(item["beta"], dtype="float32") anchor_alpha = calibration.get("pressure_anchor_alpha") if anchor_alpha and np.isfinite(current_pressure) and current_pressure > 0.0: for lead, alpha in enumerate(anchor_alpha[:LEADS]): alpha = float(np.clip(alpha, 0.0, 2.0)) result[:, lead] = alpha * result[:, lead] + (1.0 - alpha) * float(current_pressure) return np.clip(result, 850.0, 1025.0) def _sanitize(states: np.ndarray) -> np.ndarray: """Apply output bounds and preserve the R34 >= R50 >= R64 ordering.""" result = np.asarray(states, dtype="float32").copy() result[:, :, 0] = np.clip(result[:, :, 0], 0.0, 190.0) result[:, :, 1] = np.clip(result[:, :, 1], 850.0, 1025.0) result[:, :, 2] = np.clip(result[:, :, 2], 0.0, 300.0) radii = np.clip(result[:, :, 3:15], 0.0, 1000.0) for offset in range(4): radii[:, :, 4 + offset] = np.minimum(radii[:, :, 4 + offset], radii[:, :, offset]) radii[:, :, 8 + offset] = np.minimum(radii[:, :, 8 + offset], radii[:, :, 4 + offset]) result[:, :, 3:15] = radii return result def _row(mean: np.ndarray, spread: np.ndarray, lead: int) -> dict: # The historical IBTrACS structure labels are nautical miles. Keep the # neural/calibration state in that native unit, but expose all distance # fields in kilometres because the map and route use kilometres. radii = mean[lead, 3:15] * NM_TO_KM radius_spread = spread[lead, 3:15] * NM_TO_KM return { "vmax_kt": round(float(mean[lead, 0]), 2), "vmax_spread_kt": round(float(spread[lead, 0]), 2), "central_pressure_hpa": round(float(mean[lead, 1]), 2), "pressure_spread_hpa": round(float(spread[lead, 1]), 2), "rmw_km": round(float(mean[lead, 2] * NM_TO_KM), 2), "rmw_spread_km": round(float(spread[lead, 2] * NM_TO_KM), 2), "wind_radii_km": [round(float(value), 2) for value in radii], "wind_radii_spread_km": [round(float(value), 2) for value in radius_spread], # Keep the map renderer's generic pressure field alias as well. "pressure_hpa": round(float(mean[lead, 1]), 2), } def _map_grid_feature( pressure: np.ndarray, fields: np.ndarray, latitude: np.ndarray, longitude: np.ndarray, query_latitude: float, query_longitude: float, ) -> dict: """Extract a storm-relative signal from one forecast pressure map. The map can be much coarser than the route grid, especially for the Tip reanalysis. We therefore use a nearby pressure minimum, a broad annulus environment, and quadrant-wise anomaly extents rather than pretending that a single grid cell is an exact storm center. """ pressure = np.asarray(pressure, dtype="float32") fields = np.asarray(fields, dtype="float32") latitude = np.asarray(latitude, dtype="float32").reshape(-1) longitude = np.asarray(longitude, dtype="float32").reshape(-1) lat_order = np.argsort(latitude) lon_order = np.argsort(longitude) latitude = latitude[lat_order] longitude = longitude[lon_order] pressure = pressure[np.ix_(lat_order, lon_order)] fields = fields[:, lat_order, :][:, :, lon_order] lat_grid, lon_grid = np.meshgrid(latitude, longitude, indexing="ij") delta_lon = ((lon_grid - float(query_longitude) + 180.0) % 360.0) - 180.0 delta_lat = lat_grid - float(query_latitude) distance_deg = np.hypot(delta_lat, delta_lon * np.cos(np.deg2rad(float(query_latitude)))) query_row, query_column = np.unravel_index(int(np.nanargmin(distance_deg)), distance_deg.shape) query_pressure = float(pressure[query_row, query_column]) query_wind850 = float(np.hypot(fields[1, query_row, query_column], fields[2, query_row, query_column]) * 1.94384) if fields.shape[0] >= 3 else float("nan") # A minimum farther than this is likely a separate synoptic system rather # than the cyclone represented by the route point. Keep the local query # as a fallback instead of allowing a distant low to control intensity. nearby = distance_deg <= 6.0 finite = np.isfinite(pressure) candidate = np.where(nearby & finite, pressure, np.inf) if not np.isfinite(candidate).any(): row, column = query_row, query_column else: row, column = np.unravel_index(int(np.argmin(candidate)), candidate.shape) center_latitude = float(latitude[row]) center_longitude = float(longitude[column]) center_pressure = float(pressure[row, column]) center_delta_lon = ((lon_grid - center_longitude + 180.0) % 360.0) - 180.0 center_delta_lat = lat_grid - center_latitude center_distance_deg = np.hypot( center_delta_lat, center_delta_lon * np.cos(np.deg2rad(center_latitude)), ) annulus = ( (center_distance_deg >= 2.5) & (center_distance_deg <= 6.0) & np.isfinite(pressure) ) environment = float(np.nanmedian(pressure[annulus])) if annulus.any() else center_pressure + 15.0 anomaly = np.maximum(environment - pressure, 0.0) # Four degrees is a conservative upper bound for a pressure-derived # tropical wind footprint on this coarse map. Without it, a weak broad # gradient would be misreported as an 800-km R34. core = (center_distance_deg <= 4.0) & np.isfinite(anomaly) peak_anomaly = float(np.nanmax(anomaly[core])) if core.any() else 0.0 radii: list[float] = [] # These are pressure-anomaly fractions used as a stable proxy for the # R34/R50/R64 shape. The learned radius output remains the absolute # baseline; only its map-observed expansion/contraction is applied. for fraction in (0.25, 0.50, 0.70): for quadrant in ((1.0, 1.0), (1.0, -1.0), (-1.0, -1.0), (-1.0, 1.0)): north_sign, east_sign = quadrant quadrant_mask = core & (center_delta_lat * north_sign >= 0.0) & (center_delta_lon * east_sign >= 0.0) threshold = peak_anomaly * fraction valid = quadrant_mask & (anomaly >= threshold) if peak_anomaly > 0.0 else np.zeros_like(core) if valid.any(): radii.append(float(np.nanmax(center_distance_deg[valid]) * 111.2)) else: radii.append(float("nan")) wind850 = float(np.hypot(fields[1, row, column], fields[2, row, column]) * 1.94384) if fields.shape[0] >= 3 else float("nan") center_offset_km = float(distance_deg[row, column] * 111.2) center_trusted = bool(center_offset_km <= 333.6 and (environment - center_pressure) >= 4.0) offset_degrees = center_offset_km / 111.2 center_confidence = 1.0 if center_trusted else float(np.clip(1.0 - 0.75 * (offset_degrees - 3.0) / 3.0, 0.25, 0.75)) return { "map_min_pressure_hpa": round(center_pressure, 2), "map_query_pressure_hpa": round(query_pressure, 2) if np.isfinite(query_pressure) else None, "map_environment_pressure_hpa": round(environment, 2), "map_pressure_deficit_hpa": round(max(0.0, environment - center_pressure), 2), "map_center_latitude": round(center_latitude, 3), "map_center_longitude": round(center_longitude % 360.0, 3), "map_center_offset_km": round(center_offset_km, 2), "map_center_trusted": center_trusted, "map_center_confidence": round(center_confidence, 3), "map_wind850_kt": round(wind850, 2) if np.isfinite(wind850) else None, "map_query_wind850_kt": round(query_wind850, 2) if np.isfinite(query_wind850) else None, "map_pressure_radii_km": [round(value, 2) if np.isfinite(value) else None for value in radii], } def couple_forecast_to_pressure_map( structure_rows: list[dict], pressure_states: np.ndarray, field_states: np.ndarray, latitude: np.ndarray, longitude: np.ndarray, base_latitude: float, base_longitude: float, forecast_points: list[dict], current_wind: float, current_pressure: float, ) -> tuple[list[dict], dict]: """Use the causal pressure-map trajectory to correct structure forecasts. The correction is anchored to the observed current wind and pressure, so coarse reanalysis cannot replace a known storm intensity. Future changes in the map minimum drive bounded pressure/wind changes, and changes in the map's quadrant anomaly extent rescale the learned radii. This is the same forecast-state family used by the track route, not a future official field. """ pressure_states = np.asarray(pressure_states, dtype="float32") field_states = np.asarray(field_states, dtype="float32") if pressure_states.ndim != 3 or field_states.ndim != 4 or field_states.shape[1] < 3: raise ValueError(f"invalid map state shapes: {pressure_states.shape}, {field_states.shape}") if len(structure_rows) != len(forecast_points) or len(pressure_states) < len(structure_rows) + 1: raise ValueError("map states, route points, and structure rows have incompatible lengths") if not np.isfinite(current_wind) or not np.isfinite(current_pressure): return structure_rows, { "enabled": False, "reason": "current observed wind and pressure are unavailable", "official_forecasts_used": False, } features = [_map_grid_feature( pressure_states[index], field_states[index], latitude, longitude, base_latitude if index == 0 else float(forecast_points[index - 1]["lat"]), base_longitude if index == 0 else float(forecast_points[index - 1]["lon"]), ) for index in range(len(structure_rows) + 1)] # Use a trusted nearby minimum when available. Otherwise sample the map # at the route point; this keeps another low/typhoon from hijacking the # structure forecast merely because it is the strongest minimum nearby. effective_minimum = np.asarray([ float(item["map_query_pressure_hpa"]) + float(item["map_center_confidence"]) * ( float(item["map_min_pressure_hpa"]) - float(item["map_query_pressure_hpa"]) ) for item in features ], dtype="float32") effective_deficit = np.asarray([ float(item["map_pressure_deficit_hpa"]) * float(item["map_center_confidence"]) for item in features ], dtype="float32") effective_wind850 = np.asarray([ float(item["map_query_wind850_kt"]) + float(item["map_center_confidence"]) * ( float(item["map_wind850_kt"]) - float(item["map_query_wind850_kt"]) ) for item in features ], dtype="float32") # A short causal smoother prevents a one-cell minimum handoff from making # a six-hour intensity jump while retaining the map's lead-time trend. smooth_minimum = effective_minimum.copy() smooth_deficit = effective_deficit.copy() smooth_wind850 = effective_wind850.copy() for index in range(1, len(smooth_minimum)): smooth_minimum[index] = 0.65 * smooth_minimum[index - 1] + 0.35 * effective_minimum[index] smooth_deficit[index] = 0.65 * smooth_deficit[index - 1] + 0.35 * effective_deficit[index] if np.isfinite(effective_wind850[index]) and np.isfinite(smooth_wind850[index - 1]): smooth_wind850[index] = 0.65 * smooth_wind850[index - 1] + 0.35 * effective_wind850[index] if np.isfinite(smooth_minimum[index - 1]) and np.isfinite(smooth_minimum[index]): smooth_minimum[index] = min( smooth_minimum[index], smooth_minimum[index - 1] + MAP_PRESSURE_RECOVERY_LIMIT_HPA, ) reference_radii = np.asarray([ float(value) if value is not None else np.nan for value in features[0]["map_pressure_radii_km"] ], dtype="float32") corrected: list[dict] = [] map_weight = 0.35 for index, base in enumerate(structure_rows, start=1): row = dict(base) # A distant minimum is an environmental candidate, not a reliable # storm-center intensity observation. The previous floor kept more # than one-fifth of the map correction active even when confidence # had fallen to 0.25, which made another synoptic cell rewrite the # central pressure and wind. Let confidence directly control the # intensity correction; the map remains available for diagnostics and # trusted radius adjustments. row_weight = map_weight * float(np.clip(features[index]["map_center_confidence"], 0.0, 1.0)) map_pressure = float(current_pressure + smooth_minimum[index] - smooth_minimum[0]) pressure_signal = float(np.clip( 0.75 * (smooth_minimum[0] - smooth_minimum[index]) + 0.25 * (smooth_deficit[index] - smooth_deficit[0]), -30.0, 30.0, )) wind_signal = pressure_signal if np.isfinite(smooth_wind850[index]) and np.isfinite(smooth_wind850[0]): wind_signal += float(np.clip(0.18 * (smooth_wind850[index] - smooth_wind850[0]), -8.0, 8.0)) map_wind = float(current_wind + 0.70 * wind_signal) row["central_pressure_hpa"] = round(float(np.clip( (1.0 - row_weight) * float(base["central_pressure_hpa"]) + row_weight * map_pressure, 850.0, 1025.0, )), 2) row["pressure_hpa"] = row["central_pressure_hpa"] row["vmax_kt"] = round(float(np.clip( (1.0 - row_weight) * float(base["vmax_kt"]) + row_weight * map_wind, 0.0, 190.0, )), 2) row["pressure_spread_hpa"] = round(float(np.hypot( float(base.get("pressure_spread_hpa", 0.0)), row_weight * abs(map_pressure - float(base["central_pressure_hpa"])), )), 2) row["vmax_spread_kt"] = round(float(np.hypot( float(base.get("vmax_spread_kt", 0.0)), row_weight * abs(map_wind - float(base["vmax_kt"])), )), 2) map_radii = np.asarray([ float(value) if value is not None else np.nan for value in features[index]["map_pressure_radii_km"] ], dtype="float32") base_radii = np.asarray(base["wind_radii_km"], dtype="float32") if not features[0]["map_center_trusted"] or not features[index]["map_center_trusted"]: map_radii[:] = np.nan if np.isfinite(reference_radii).any() and np.isfinite(map_radii).any(): valid = np.isfinite(reference_radii) & np.isfinite(map_radii) & (reference_radii >= 20.0) ratios = np.ones(12, dtype="float32") ratios[valid] = np.clip(map_radii[valid] / reference_radii[valid], 0.60, 1.50) if valid.any(): ratios[~valid] = float(np.clip(np.nanmedian(ratios[valid]), 0.60, 1.50)) radius_factor = 1.0 + 0.25 * (ratios - 1.0) adjusted_radii = base_radii * radius_factor else: adjusted_radii = base_radii adjusted_radii = np.clip(adjusted_radii, 0.0, 1000.0) for quadrant in range(4): adjusted_radii[4 + quadrant] = min(adjusted_radii[4 + quadrant], adjusted_radii[quadrant]) adjusted_radii[8 + quadrant] = min(adjusted_radii[8 + quadrant], adjusted_radii[4 + quadrant]) row["wind_radii_km"] = [round(float(value), 2) for value in adjusted_radii] map_rmw_ratio = 1.0 if features[0]["map_center_trusted"] and features[index]["map_center_trusted"]: ref_r64 = float(np.nanmedian(reference_radii[8:])) if np.isfinite(reference_radii[8:]).any() else np.nan future_r64 = float(np.nanmedian(map_radii[8:])) if np.isfinite(map_radii[8:]).any() else np.nan if np.isfinite(ref_r64) and np.isfinite(future_r64) and ref_r64 >= 20.0: map_rmw_ratio = float(np.clip(future_r64 / ref_r64, 0.75, 1.25)) row["rmw_km"] = round(float(np.clip( float(base["rmw_km"]) * (1.0 + 0.15 * (map_rmw_ratio - 1.0)), 0.0, 300.0, )), 2) row["pressure_map_features"] = features[index] corrected.append(row) metadata = { "enabled": True, "method": "causal forecast pressure-map query plus confidence-weighted local-minimum/anomaly-radius coupling", "map_weight": map_weight, "minimum_confidence": "local minimum contribution decays directly with route-to-minimum confidence; untrusted minima have near-zero central-intensity weight", "pressure_recovery_limit_hpa_per_6h": MAP_PRESSURE_RECOVERY_LIMIT_HPA, "pressure_recovery_limit_source": "training-split 90th percentile of positive six-hour pressure changes", "pressure_signal_weights": { "tracked_map_minimum": 0.75, "map_anomaly_extent": 0.25, }, "wind_pressure_anchor": "observed current wind/pressure; map drives only future changes", "radius_method": "learned radius baseline rescaled by trusted local quadrant pressure-anomaly extent; max four-degree footprint", "map_features": features, "official_forecasts_used": False, "positive_lead_weather_product_used": False, } return corrected, metadata class Trackformer11IntensityEnsemble: """Load frozen Trackformer1.1 experts and emit calibrated structure rows.""" def __init__( self, checkpoint_root: Path, calibration_path: Path | None = None, device: str | None = None, ): self.checkpoint_root = Path(checkpoint_root) self.calibration_path = Path(calibration_path) if calibration_path else None self.device = _device(device) paths = sorted(self.checkpoint_root.glob("trackformer_1_1_intensity_seed*.pt")) if len(paths) < 3: raise FileNotFoundError( f"expected three Trackformer1.1 intensity checkpoints in {self.checkpoint_root}; found {len(paths)}" ) self.models: list[StructureSpatialExpert] = [] for path in paths[:3]: payload = torch.load(path, map_location="cpu", weights_only=False) config = payload["config"] model = StructureSpatialExpert( config["width"], config["layers"], config["heads"], structure_residual=bool(config.get("structure_residual", False)), ) model.load_state_dict(payload["model"]) self.models.append(model.to(self.device).eval()) self.calibration = {} if self.calibration_path and self.calibration_path.exists(): self.calibration = json.loads(self.calibration_path.read_text(encoding="utf-8")) self.structure_models: list[StructureSpatialExpert] = [] structure_root = self.calibration.get("structure_checkpoint_root") if structure_root: structure_root_path = self._resolve_root(structure_root) structure_paths = sorted(structure_root_path.glob("trackformer_1_1_structure_seed*.pt")) for path in structure_paths[:3]: payload = torch.load(path, map_location="cpu", weights_only=False) config = payload["config"] model = StructureSpatialExpert( config["width"], config["layers"], config["heads"], structure_residual=bool(config.get("structure_residual", False)), ) model.load_state_dict(payload["model"]) self.structure_models.append(model.to(self.device).eval()) self.temporal_models = [] temporal_root = self.calibration.get("temporal_checkpoint_root") self.temporal_calibration = self.calibration if temporal_root: from trackformer_1_1_temporal import TemporalStructureSpatial temporal_root_path = self._resolve_root(temporal_root) temporal_paths = sorted(temporal_root_path.glob("trackformer_1_1_temporal_seed*.pt")) for path in temporal_paths[:3]: payload = torch.load(path, map_location="cpu", weights_only=False) config = payload["config"] model = TemporalStructureSpatial( config["width"], config["layers"], config["heads"] ) model.load_state_dict(payload["model"]) self.temporal_models.append(model.to(self.device).eval()) @staticmethod def _resolve_root(value: str | Path) -> Path: path = Path(value) return path if path.is_absolute() else Path(__file__).resolve().parent / path @torch.no_grad() def predict( self, track: np.ndarray, field: np.ndarray, current_wind: float, current_pressure: float, previous_wind: float, previous_pressure: float, current_structure: np.ndarray | None = None, history_field: np.ndarray | None = None, history_available: np.ndarray | None = None, ) -> tuple[list[dict], dict]: track = np.asarray(track, dtype="float32") field = np.asarray(field, dtype="float32") if track.shape != (9, 54): raise ValueError(f"Trackformer1.1 track must have shape (9, 54), got {track.shape}") if field.shape != (4, 17, 17): raise ValueError(f"Trackformer1.1 field must have shape (4, 17, 17), got {field.shape}") if history_field is not None: history_field = np.asarray(history_field, dtype="float32") if history_field.shape != (8, 17, 17): raise ValueError( f"Trackformer1.1 history field must have shape (8, 17, 17), got {history_field.shape}" ) if history_available is None: history_available = np.ones(2, dtype="float32") history_available = np.asarray(history_available, dtype="float32").reshape(-1) if history_available.shape != (2,): raise ValueError("Trackformer1.1 history availability must have shape (2,)") if not np.isfinite(track).all() or not np.isfinite(field).all(): raise ValueError("Trackformer1.1 inputs contain non-finite values") track_clip = self.calibration.get("track_input_clip", {}) lower = np.asarray(track_clip.get("lower", []), dtype="float32") upper = np.asarray(track_clip.get("upper", []), dtype="float32") if lower.shape == (track.shape[1],) and upper.shape == (track.shape[1],): track = np.clip(track, lower[None, :], upper[None, :]) if current_structure is not None: current_structure = np.asarray(current_structure, dtype="float32").reshape(-1) if current_structure.shape != (13,): raise ValueError( "Trackformer1.1 current structure must contain RMW plus twelve radii in native nautical miles" ) structure_available = ( np.isfinite(current_structure).astype("float32") if current_structure is not None else None ) current_available = float(np.isfinite(current_wind) and np.isfinite(current_pressure)) current_values = np.nan_to_num( np.asarray([current_wind, current_pressure], dtype="float32") / TARGET_SCALE[2:4], nan=0.0, ) available = np.asarray([current_available, current_available], dtype="float32") track_tensor = torch.from_numpy(track[None]).to(self.device) field_tensor = torch.from_numpy(field[None]).to(self.device) current_tensor = torch.from_numpy(current_values[None]).to(self.device) available_tensor = torch.from_numpy(available[None]).to(self.device) structure_tensor = None structure_available_tensor = None if current_structure is not None: structure_tensor = torch.from_numpy( np.nan_to_num( current_structure / TARGET_SCALE[4:], nan=0.0, posinf=0.0, neginf=0.0, )[None] ).to(self.device) structure_available_tensor = torch.from_numpy(structure_available[None]).to(self.device) states = np.stack([ model( track_tensor, field_tensor, current_tensor, available_tensor, structure_tensor, structure_available_tensor, )[0][0].detach().cpu().numpy() for model in self.models ]).astype("float32") states *= STRUCTURE_SCALE[None, None, :] if self.structure_models: structure_states = np.stack([ model(track_tensor, field_tensor, current_tensor, available_tensor)[0][0].detach().cpu().numpy() for model in self.structure_models ]).astype("float32") * STRUCTURE_SCALE[None, None, :] expert_alpha = np.asarray( self.calibration.get("structure_expert_alpha", []), dtype="float32" ) if expert_alpha.shape != (LEADS, 13): expert_alpha = np.zeros((LEADS, 13), dtype="float32") for lead in range(LEADS): states[:, lead, 2:] = ( expert_alpha[lead][None, :] * states[:, lead, 2:] + (1.0 - expert_alpha[lead][None, :]) * structure_states[:, lead, 2:] ) temporal_states = None if self.temporal_models and history_field is not None and current_structure is not None: history_tensor = torch.from_numpy(history_field[None]).to(self.device) history_available_tensor = torch.from_numpy(history_available[None]).to(self.device) temporal_states = np.stack([ model( track_tensor, field_tensor, current_tensor, available_tensor, structure_tensor, structure_available_tensor, history_tensor, history_available_tensor, )[0][0].detach().cpu().numpy() for model in self.temporal_models ]).astype("float32") * STRUCTURE_SCALE[None, None, :] if current_available: calibrated_wind = _calibrated_wind(states, float(current_wind), self.calibration) states[:, :, 1] = _calibrated_pressure( states, float(current_wind), float(current_pressure), float(previous_wind) if np.isfinite(previous_wind) else float(current_wind), float(previous_pressure) if np.isfinite(previous_pressure) else float(current_pressure), calibrated_wind, self.calibration, ) states[:, :, 0] = calibrated_wind if temporal_states is not None: temporal_alpha = self.calibration.get("temporal_wind_blend_alpha") temporal_min_wind = float(self.calibration.get("temporal_wind_min_current_kt", 0.0)) if temporal_alpha and float(current_wind) >= temporal_min_wind: temporal_calibration = {"wind_blend_alpha": temporal_alpha} states[:, :, 0] = _calibrated_wind( temporal_states, float(current_wind), temporal_calibration, ) structure_alpha = np.asarray( self.calibration.get("structure_blend_alpha", []), dtype="float32" ) if current_structure is not None and structure_alpha.shape == (LEADS, 13): available_structure = np.isfinite(current_structure) anchor = np.nan_to_num(current_structure, nan=0.0, posinf=0.0, neginf=0.0) for lead in range(LEADS): alpha = np.clip(structure_alpha[lead], 0.0, 1.0) for offset in np.flatnonzero(available_structure): channel = 2 + int(offset) states[:, lead, channel] = ( alpha[int(offset)] * states[:, lead, channel] + (1.0 - alpha[int(offset)]) * anchor[int(offset)] ) states = _sanitize(states) mean = states.mean(axis=0) spread = states.std(axis=0) model_label = ( "Trackformer1.1 residual-anchor spatial ensemble" if self.models and self.models[0].structure_residual else "frozen spatial structure ensemble" ) rows = [_row(mean, spread, lead) for lead in range(LEADS)] metadata = { "model": f"{model_label} used as the Trackformer1.1 intensity head", "checkpoint_count": len(self.models), "calibration": str(self.calibration_path) if self.calibration_path and self.calibration_path.exists() else None, "field_contract": "4x17x17 analysis patch, q/31.75 then clipped to [-4,4]", "outputs": ["vmax_kt", "central_pressure_hpa", "rmw_km", "wind_radii_km"], "native_structure_unit": "nautical_miles", "output_distance_unit": "kilometres", "ensemble_spread": "standard deviation across the three frozen spatial experts", "structure_calibration": "lead- and component-wise validation blend with observed current RMW/radii when available", "structure_expert_blend": bool(self.structure_models), "temporal_wind_branch": bool( self.temporal_models and history_field is not None and current_available and float(current_wind) >= float(self.calibration.get("temporal_wind_min_current_kt", 0.0)) ), "temporal_wind_branch_policy": "same-storm t-12/t-24 analysis patches; wind only; enabled above the validation-selected current-wind gate; base pressure calibration retained", "pressure_anchor_alpha": self.calibration.get("pressure_anchor_alpha"), "pressure_anchor_policy": self.calibration.get("pressure_anchor_policy"), "track_input_policy": self.calibration.get("track_input_clip", {}).get("method"), "official_forecasts_used": False, "positive_lead_weather_used": False, "device": str(self.device), } return rows, metadata