AgentA123/project-helios-api

1from contextlib import asynccontextmanager2from datetime import datetime, timedelta, timezone3import os4from pathlib import Path5import time6from typing import Any7 8import httpx9import numpy as np10import torch11from fastapi import FastAPI, HTTPException12from fastapi.middleware.cors import CORSMiddleware13 14from model import SolarTransformer15 16 17NOAA_GOES_XRAY_URL = "https://services.swpc.noaa.gov/json/goes/primary/xrays-1-day.json"18NOAA_GOES_XRAY_7_DAY_URL = (19    "https://services.swpc.noaa.gov/json/goes/primary/xrays-7-day.json"20)21LONG_CHANNEL = "0.1-0.8nm"22MODEL_PATH = Path(__file__).with_name("solar_transformer_model.pt")23APP_NAME = "Solar Flare AI Backend"24VALIDATION_STRIDE_MINUTES = 6025VALIDATION_MAX_WINDOWS = 14426VALIDATION_BATCH_SIZE = 827 28 29class ModelBundle:30    def __init__(self, checkpoint_path: Path):31        checkpoint = self._load_checkpoint(checkpoint_path)32 33        self.feature_cols = list(checkpoint["feature_cols"])34        self.seq_len = int(checkpoint["seq_len"])35        self.step = int(checkpoint.get("step", 1))36        self.threshold = float(checkpoint["threshold"])37        self.forecast_minutes = int(checkpoint["forecast_minutes"])38        self.target_flare_class = str(checkpoint["target_flare_class"])39        self.best_f1 = (40            None if checkpoint.get("best_f1") is None else float(checkpoint["best_f1"])41        )42 43        self.mean = self._as_feature_array(checkpoint["mean"], "mean")44        self.std = self._as_feature_array(checkpoint["std"], "std")45        self.std = np.where(self.std == 0, 1.0, self.std)46 47        state_dict = checkpoint["model_state_dict"]48        max_len = int(state_dict["pos.pe"].shape[1]) if "pos.pe" in state_dict else self.seq_len49        dim_feedforward = (50            int(state_dict["encoder.layers.0.linear1.weight"].shape[0])51            if "encoder.layers.0.linear1.weight" in state_dict52            else 12853        )54        self.model = SolarTransformer(55            input_dim=len(self.feature_cols),56            max_len=max_len,57            dim_feedforward=dim_feedforward,58        )59        self.model.load_state_dict(state_dict)60        self.model.to(torch.device("cpu"))61        self.model.eval()62 63    @staticmethod64    def _load_checkpoint(checkpoint_path: Path) -> dict[str, Any]:65        if not checkpoint_path.exists():66            raise FileNotFoundError(f"Model checkpoint not found: {checkpoint_path}")67 68        try:69            return torch.load(70                checkpoint_path,71                map_location=torch.device("cpu"),72                weights_only=False,73            )74        except TypeError:75            return torch.load(checkpoint_path, map_location=torch.device("cpu"))76 77    def _as_feature_array(self, value: Any, name: str) -> np.ndarray:78        if isinstance(value, dict):79            return np.array([float(value[col]) for col in self.feature_cols], dtype=np.float32)80 81        if hasattr(value, "loc"):82            return np.array(83                [float(value.loc[col]) for col in self.feature_cols],84                dtype=np.float32,85            )86 87        array = np.asarray(value, dtype=np.float32)88        if array.ndim > 1:89            array = array.reshape(-1)90 91        if array.shape[0] != len(self.feature_cols):92            raise ValueError(93                f"Checkpoint {name} has {array.shape[0]} values, "94                f"but feature_cols has {len(self.feature_cols)}."95            )96        return array97 98    def predict_probability(self, features: np.ndarray) -> float:99        normalized = (features - self.mean) / self.std100        tensor = torch.tensor(normalized, dtype=torch.float32).unsqueeze(0)101 102        with torch.no_grad():103            logits = self.model(tensor)104            probability = torch.sigmoid(logits).item()105 106        return float(probability)107 108    def predict_probabilities(self, feature_windows: np.ndarray) -> np.ndarray:109        normalized = (feature_windows - self.mean) / self.std110        tensor = torch.tensor(normalized, dtype=torch.float32)111 112        with torch.no_grad():113            logits = self.model(tensor)114            probabilities = torch.sigmoid(logits).cpu().numpy()115 116        return probabilities.astype(float)117 118    def predict_probabilities_batched(119        self,120        feature_windows: np.ndarray,121        batch_size: int = VALIDATION_BATCH_SIZE,122    ) -> np.ndarray:123        chunks = []124        for start in range(0, len(feature_windows), batch_size):125            batch = feature_windows[start : start + batch_size]126            chunks.append(self.predict_probabilities(batch))127 128        if not chunks:129            return np.asarray([], dtype=float)130 131        return np.concatenate(chunks)132 133 134model_bundle: ModelBundle | None = None135 136 137@asynccontextmanager138async def lifespan(app: FastAPI):139    global model_bundle140    model_bundle = ModelBundle(MODEL_PATH)141    yield142 143 144app = FastAPI(145    title=APP_NAME,146    description="FastAPI backend for real-time GOES X-ray solar flare inference.",147    version="1.0.0",148    lifespan=lifespan,149)150 151app.add_middleware(152    CORSMiddleware,153    allow_origins=["*"],154    allow_credentials=False,155    allow_methods=["*"],156    allow_headers=["*"],157)158 159 160@app.get("/")161async def root() -> dict[str, Any]:162    return {163        "project": APP_NAME,164        "available_endpoints": {165            "root": "/",166            "health": "/health",167            "prediction": "/predict",168            "space_weather": "/space-weather",169            "validation": "/validation",170            "docs": "/docs",171        },172    }173 174 175@app.get("/health")176async def health() -> dict[str, Any]:177    return {178        "status": "ok",179        "model_loaded": model_bundle is not None,180        "data_source": NOAA_GOES_XRAY_URL,181    }182 183 184@app.get("/space-weather")185async def space_weather() -> dict[str, Any]:186    rows = await fetch_goes_long_channel(NOAA_GOES_XRAY_URL)187    if not rows:188        raise HTTPException(189            status_code=503,190            detail="No usable NOAA GOES X-ray rows were available.",191        )192 193    recent_rows = rows[-180:]194    flux_values = [float(row["flux"]) for row in rows]195    return {196        "data_source": NOAA_GOES_XRAY_URL,197        "channel": LONG_CHANNEL,198        "latest_time": rows[-1]["time_tag"],199        "latest_flux": rows[-1]["flux"],200        "data_points_returned": len(recent_rows),201        "data_points_available": len(rows),202        "min_flux": min(flux_values),203        "max_flux": max(flux_values),204        "series": recent_rows,205        "note": "Recent NOAA SWPC GOES primary X-ray flux values in the 0.1-0.8nm channel.",206    }207 208 209@app.get("/predict")210async def predict() -> dict[str, Any]:211    if model_bundle is None:212        raise HTTPException(status_code=503, detail="Model is not loaded.")213 214    rows = await fetch_goes_long_channel(NOAA_GOES_XRAY_URL)215    if len(rows) < model_bundle.seq_len:216        raise HTTPException(217            status_code=503,218            detail=(219                f"Not enough NOAA GOES data points. "220                f"Need {model_bundle.seq_len}, got {len(rows)}."221            ),222        )223 224    recent_rows = rows[-model_bundle.seq_len :]225    features = build_feature_matrix(recent_rows, model_bundle.feature_cols)226    probability = model_bundle.predict_probability(features)227    risk_level = classify_risk(probability)228    input_window_start_time = recent_rows[0]["time_tag"]229    latest_row = recent_rows[-1]230    forecast_start_time = latest_row["time_tag"]231    forecast_end_time = add_minutes_to_time_tag(232        forecast_start_time,233        model_bundle.forecast_minutes,234    )235 236    return {237        "flare_probability": probability,238        "risk_level": risk_level,239        "threshold": model_bundle.threshold,240        "alert": probability >= model_bundle.threshold,241        "latest_time": latest_row["time_tag"],242        "latest_flux": latest_row["flux"],243        "input_window_start_time": input_window_start_time,244        "forecast_start_time": forecast_start_time,245        "forecast_end_time": forecast_end_time,246        "forecast_window_minutes": model_bundle.forecast_minutes,247        "target_flare_class": model_bundle.target_flare_class,248        "model_best_f1": model_bundle.best_f1,249        "data_points_used": len(recent_rows),250        "noaa_rows_available": len(rows),251        "response_generated_at": datetime.now(timezone.utc)252        .isoformat()253        .replace("+00:00", "Z"),254        "interpretation": (255            "The model estimates whether the target flare class is likely "256            "within the forecast window. It does not predict an exact flare "257            "onset time."258        ),259    }260 261 262@app.get("/validation")263async def validation() -> dict[str, Any]:264    if model_bundle is None:265        raise HTTPException(status_code=503, detail="Model is not loaded.")266 267    try:268        rows = await fetch_goes_long_channel(NOAA_GOES_XRAY_7_DAY_URL)269        observed_threshold = flare_flux_threshold(model_bundle.target_flare_class)270        horizon = min(model_bundle.forecast_minutes, max(60, len(rows) // 4))271        required_points = model_bundle.seq_len + horizon272 273        if len(rows) < required_points:274            raise HTTPException(275                status_code=503,276                detail=(277                    "Not enough recent NOAA GOES 7-day data for validation. "278                    f"Need at least {required_points} one-minute points, got {len(rows)}."279                ),280            )281 282        stride = max(VALIDATION_STRIDE_MINUTES, model_bundle.step, 1)283        max_start = len(rows) - required_points284        starts = list(range(0, max_start + 1, stride))285        if len(starts) > VALIDATION_MAX_WINDOWS:286            starts = starts[-VALIDATION_MAX_WINDOWS:]287 288        feature_windows = []289        window_metadata = []290        for start in starts:291            input_rows = rows[start : start + model_bundle.seq_len]292            future_rows = rows[293                start + model_bundle.seq_len : start + model_bundle.seq_len + horizon294            ]295            feature_windows.append(296                build_feature_matrix(input_rows, model_bundle.feature_cols)297            )298            window_metadata.append((input_rows, future_rows))299 300        if not feature_windows:301            raise HTTPException(302                status_code=503,303                detail="No validation windows could be built from recent NOAA data.",304            )305 306        feature_window_array = np.asarray(feature_windows, dtype=np.float32)307        probabilities = model_bundle.predict_probabilities_batched(308            feature_window_array,309            batch_size=VALIDATION_BATCH_SIZE,310        )311 312        results = []313        for probability, (input_rows, future_rows) in zip(314            probabilities,315            window_metadata,316        ):317            probability_value = float(probability)318            predicted_alert = probability_value >= model_bundle.threshold319            future_max_flux = max(float(row["flux"]) for row in future_rows)320            actual_event = future_max_flux >= observed_threshold321 322            results.append(323                {324                    "prediction_time": input_rows[-1]["time_tag"],325                    "observation_window_end": future_rows[-1]["time_tag"],326                    "probability": probability_value,327                    "predicted_alert": bool(predicted_alert),328                    "future_max_flux": future_max_flux,329                    "actual_event": bool(actual_event),330                }331            )332    except HTTPException:333        raise334    except Exception as exc:335        raise HTTPException(336            status_code=500,337            detail=f"Validation failed while processing recent NOAA data: {exc}",338        ) from exc339 340    try:341        stats = confusion_summary(results)342        return {343            **stats,344            "target_flare_class": model_bundle.target_flare_class,345            "observed_proxy_threshold": observed_threshold,346            "observed_flux_threshold_w_m2": observed_threshold,347            "model_threshold": model_bundle.threshold,348            "forecast_window_minutes": horizon,349            "checkpoint_forecast_window_minutes": model_bundle.forecast_minutes,350            "sequence_length_minutes": model_bundle.seq_len,351            "validation_stride_minutes": stride,352            "validation_windows_requested": len(results),353            "validation_max_windows": VALIDATION_MAX_WINDOWS,354            "validation_batch_size": VALIDATION_BATCH_SIZE,355            "data_start_time": rows[0]["time_tag"],356            "data_end_time": rows[-1]["time_tag"],357            "data_source": NOAA_GOES_XRAY_7_DAY_URL,358            "validation_note": (359                "This is an approximate recent backtest using observed GOES "360                "X-ray flux thresholds as a proxy for flare events. It is not "361                "the original DONKI-labeled Kaggle validation set."362            ),363            "method": "Rolling windows over recent 7-day GOES data including quiet and flare-active periods.",364            "sample_windows": results[-10:],365        }366    except Exception as exc:367        raise HTTPException(368            status_code=500,369            detail=f"Validation summary failed: {exc}",370        ) from exc371 372 373async def fetch_goes_long_channel(url: str = NOAA_GOES_XRAY_URL) -> list[dict[str, Any]]:374    try:375        async with httpx.AsyncClient(timeout=20.0) as client:376            response = await client.get(377                url,378                params={"_": str(int(time.time()))},379                headers={"Cache-Control": "no-cache", "Pragma": "no-cache"},380            )381            response.raise_for_status()382            payload = response.json()383    except httpx.HTTPError as exc:384        raise HTTPException(385            status_code=503,386            detail=f"Unable to fetch NOAA GOES X-ray data: {exc}",387        ) from exc388    except ValueError as exc:389        raise HTTPException(390            status_code=503,391            detail="NOAA GOES X-ray response was not valid JSON.",392        ) from exc393 394    if not isinstance(payload, list):395        raise HTTPException(396            status_code=503,397            detail="NOAA GOES X-ray response had an unexpected format.",398        )399 400    filtered: list[dict[str, Any]] = []401    for row in payload:402        if not isinstance(row, dict):403            continue404 405        if row.get("energy") != LONG_CHANNEL:406            continue407 408        time_tag = row.get("time_tag")409        flux = row.get("flux")410        if time_tag is None or flux is None:411            continue412 413        try:414            flux_value = float(flux)415        except (TypeError, ValueError):416            continue417 418        if not np.isfinite(flux_value):419            continue420 421        filtered.append({"time_tag": time_tag, "flux": flux_value})422 423    filtered.sort(key=lambda item: item["time_tag"])424    return filtered425 426 427def build_feature_matrix(428    rows: list[dict[str, Any]],429    feature_cols: list[str],430) -> np.ndarray:431    feature_rows = []432    for row in rows:433        flux = max(float(row["flux"]), 1e-12)434        values = {435            "flux": flux,436            "flux_log": np.log10(flux),437            "sunspot_number": 0.0,438            "radio_flux": 0.0,439        }440 441        try:442            feature_rows.append([float(values[col]) for col in feature_cols])443        except KeyError as exc:444            supported = ", ".join(sorted(values))445            raise HTTPException(446                status_code=500,447                detail=(448                    f"Checkpoint requires unsupported live feature {exc.args[0]!r}. "449                    f"Currently supported live features: {supported}."450                ),451            ) from exc452 453    return np.asarray(feature_rows, dtype=np.float32)454 455 456def classify_risk(probability: float) -> str:457    if probability < 0.30:458        return "LOW"459    if probability < 0.60:460        return "MEDIUM"461    return "HIGH"462 463 464def flare_flux_threshold(target_flare_class: str) -> float:465    flare_class = target_flare_class.strip().upper()[:1]466    thresholds = {467        "C": 1e-6,468        "M": 1e-5,469        "X": 1e-4,470    }471 472    if flare_class not in thresholds:473        raise HTTPException(474            status_code=500,475            detail=(476                "Unsupported target flare class for GOES flux-threshold "477                f"validation: {target_flare_class!r}."478            ),479        )480 481    return thresholds[flare_class]482 483 484def confusion_summary(results: list[dict[str, Any]]) -> dict[str, Any]:485    total = len(results)486    true_positives = sum(487        1 for row in results if row["predicted_alert"] and row["actual_event"]488    )489    false_positives = sum(490        1 for row in results if row["predicted_alert"] and not row["actual_event"]491    )492    false_negatives = sum(493        1 for row in results if not row["predicted_alert"] and row["actual_event"]494    )495    true_negatives = sum(496        1 for row in results if not row["predicted_alert"] and not row["actual_event"]497    )498 499    accuracy = safe_divide(true_positives + true_negatives, total)500    precision = safe_divide(true_positives, true_positives + false_positives)501    recall = safe_divide(true_positives, true_positives + false_negatives)502    if precision is None or recall is None or precision + recall == 0:503        f1_score = None504    else:505        f1_score = 2 * precision * recall / (precision + recall)506 507    return {508        "total_windows_checked": total,509        "predicted_alerts": true_positives + false_positives,510        "actual_flare_level_events": true_positives + false_negatives,511        "true_positives": true_positives,512        "false_positives": false_positives,513        "false_negatives": false_negatives,514        "true_negatives": true_negatives,515        "accuracy": accuracy,516        "precision": precision,517        "recall": recall,518        "f1_score": f1_score,519    }520 521 522def safe_divide(numerator: float, denominator: float) -> float | None:523    if denominator == 0:524        return None525    return numerator / denominator526 527 528def add_minutes_to_time_tag(time_tag: str, minutes: int) -> str:529    normalized = time_tag.replace("Z", "+00:00")530    parsed = datetime.fromisoformat(normalized)531    if parsed.tzinfo is None:532        parsed = parsed.replace(tzinfo=timezone.utc)533 534    return (parsed + timedelta(minutes=minutes)).isoformat().replace("+00:00", "Z")535 536 537if __name__ == "__main__":538    import uvicorn539 540    uvicorn.run(541        "main:app",542        host="0.0.0.0",543        port=int(os.getenv("PORT", "7860")),544    )545