Create app.py

app.py

@@ -0,0 +1,191 @@
+import torch
+from transformers import AutoModelForImageSegmentation, AutoProcessor
+from sunpy.net import Fido, attrs as a
+import astropy.units as u
+import numpy as np
+from fastapi import FastAPI
+import gradio as gr
+from datetime import datetime, timedelta
+from sklearn.linear_model import LogisticRegression
+import xarray
+import matplotlib.pyplot as plt
+import os
+from pathlib import Path
+
+# FastAPI app
+app = FastAPI()
+
+# Surya model setup
+MODEL_ID = "nasa-ibm-ai4science/ar_segmentation_surya"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = AutoModelForImageSegmentation.from_pretrained(MODEL_ID).to(device)
+processor = AutoProcessor.from_pretrained(MODEL_ID)
+
+# Cache directory for SDO data
+CACHE_DIR = Path("/persistent-storage/sdo_data")
+CACHE_DIR.mkdir(exist_ok=True)
+
+# Simplified NRLMSISE-00 model for neutral density (placeholder for WAM-IPE)
+def nrlmsise_density(altitude_km, flare_factor=1.0):
+    """Estimate neutral density using a simplified NRLMSISE-00 model."""
+    base_density = 1e-12  # kg/m^3 at 550 km (approximate)
+    density = base_density * flare_factor  # Scale by flare/CME impact
+    return density
+
+# Simplified flare/CME probability model
+def train_flare_model():
+    """Train a simple logistic regression model for flare probability."""
+    # Dummy data: active region size (Mm²) vs. flare probability
+    X = np.array([[100], [500], [1000], [2000]])  # Size in Mm²
+    y = np.array([0, 0.5, 0.8, 0.95])  # Flare probability
+    model = LogisticRegression().fit(X, y)
+    return model
+
+flare_model = train_flare_model()
+
+# Drag calculation
+def calculate_drag(density, altitude_km=550, velocity=7.5e3, cd=2.2, area=10):
+    """Calculate drag force and orbital decay for a Starlink satellite."""
+    drag_force = 0.5 * density * velocity**2 * cd * area  # N
+    mass = 260  # kg (Starlink satellite mass, approximate)
+    acceleration = drag_force / mass  # m/s^2
+    decay_rate = acceleration * 86400 / 1000  # km/day
+    return drag_force, decay_rate
+
+# Fetch and preprocess SDO data
+def fetch_sdo_data(start_time, end_time):
+    """Fetch SDO AIA 171Å data for the given time range."""
+    try:
+        query = Fido.search(
+            a.Time(start_time, end_time),
+            a.Instrument("AIA"),
+            a.Wavelength(171 * u.angstrom)
+        )
+        files = Fido.fetch(query, path=str(CACHE_DIR / "aia_{file}"), progress=True)
+        return files
+    except Exception as e:
+        print(f"Error fetching SDO data: {e}")
+        return None
+
+def preprocess_sdo_data(files):
+    """Preprocess SDO data for Surya model."""
+    from sunpy.map import Map
+    if not files:
+        return None
+    sdo_map = Map(files[0])
+    data = sdo_map.data  # 4096x4096 image
+    data = data / np.max(data)  # Normalize
+    data = np.expand_dims(data, axis=(0, 1))  # Shape: (1, 1, H, W)
+    return data
+
+# Run Surya segmentation
+def run_surya_segmentation(sdo_data):
+    """Run Surya AR segmentation on SDO data."""
+    if sdo_data is None:
+        return None
+    inputs = processor(images=sdo_data, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    masks = torch.sigmoid(outputs.logits).cpu().numpy()  # Segmentation masks
+    return masks
+
+# Analyze active regions
+def analyze_active_regions(masks):
+    """Extract active region properties from segmentation masks."""
+    if masks is None:
+        return []
+    mask = masks[0, 0]  # First image, first channel
+    ar_size = np.sum(mask > 0.5) * 0.1  # Approximate size in Mm² (pixel-to-Mm conversion)
+    return [{"id": "AR1", "size_mm2": ar_size}]
+
+# Predict flare/CME and density impact
+def predict_events(active_regions):
+    """Predict flare/CME probabilities and density impact."""
+    if not active_regions:
+        return {"flare_prob": 0.0, "cme_prob": 0.0, "flare_factor": 1.0}
+    ar_size = active_regions[0]["size_mm2"]
+    flare_prob = flare_model.predict_proba([[ar_size]])[0, 1]
+    cme_prob = flare_prob * 0.5  # Simplified: CMEs less likely than flares
+    flare_factor = 1.0 + flare_prob * 2.0  # Scale density by flare intensity
+    return {"flare_prob": flare_prob, "cme_prob": cme_prob, "flare_factor": flare_factor}
+
+# Main prediction function
+def predict_drag(start_time_str):
+    """Predict drag based on SDO data and active region analysis."""
+    try:
+        start_time = datetime.fromisoformat(start_time_str.replace("Z", "+00:00"))
+        end_time = start_time + timedelta(minutes=15)
+        sdo_files = fetch_sdo_data(start_time.isoformat(), end_time.isoformat())
+        sdo_data = preprocess_sdo_data(sdo_files)
+        masks = run_surya_segmentation(sdo_data)
+        active_regions = analyze_active_regions(masks)
+        event_probs = predict_events(active_regions)
+        
+        # Estimate density and drag
+        density = nrlmsise_density(altitude_km=550, flare_factor=event_probs["flare_factor"])
+        drag_force, decay_rate = calculate_drag(density)
+        
+        # Generate recommendations
+        recommendations = []
+        if event_probs["flare_prob"] > 0.7:
+            recommendations.append("Consider raising orbit by 5-10 km.")
+        if decay_rate > 0.5:
+            recommendations.append("Switch satellites to edge-on safe mode.")
+        
+        # Save visualization
+        if masks is not None:
+            plt.imshow(masks[0, 0], cmap="hot")
+            plt.title("Active Region Segmentation")
+            plt.savefig(CACHE_DIR / "ar_mask.png")
+            plt.close()
+        
+        return {
+            "active_regions": active_regions,
+            "flare_probability": event_probs["flare_prob"],
+            "cme_probability": event_probs["cme_prob"],
+            "neutral_density": density,
+            "drag_force": drag_force,
+            "orbital_decay": decay_rate,
+            "recommendations": recommendations,
+            "visualization": str(CACHE_DIR / "ar_mask.png")
+        }
+    except Exception as e:
+        return {"error": str(e)}
+
+# FastAPI endpoint
+@app.post("/predict_drag")
+async def api_predict_drag(start_time: str):
+    """API endpoint for drag prediction."""
+    result = predict_drag(start_time)
+    return result
+
+# Gradio interface
+def gradio_predict(start_time):
+    """Gradio interface for drag prediction."""
+    result = predict_drag(start_time)
+    if "error" in result:
+        return f"Error: {result['error']}"
+    
+    output = f"""
+    **Active Regions**: {result['active_regions']}
+    **Flare Probability**: {result['flare_probability']:.2f}
+    **CME Probability**: {result['cme_probability']:.2f}
+    **Neutral Density**: {result['neutral_density']:.2e} kg/m³
+    **Drag Force**: {result['drag_force']:.2e} N
+    **Orbital Decay**: {result['orbital_decay']:.2f} km/day
+    **Recommendations**: {', '.join(result['recommendations'])}
+    """
+    return output, result["visualization"]
+
+# Gradio app
+iface = gr.Interface(
+    fn=gradio_predict,
+    inputs=gr.Textbox(label="Start Time (ISO format, e.g., 2025-08-21T12:00:00Z)"),
+    outputs=[gr.Textbox(label="Prediction Results"), gr.Image(label="Active Region Visualization")],
+    title="Starlink Drag Prediction App",
+    description="Predict satellite drag based on solar active regions using NASA-IBM Surya model."
+)
+
+# Run Gradio app (Hugging Face Spaces handles this automatically)
+if __name__ == "__main__":
+    iface.launch()