Update app.py

app.py

@@ -1,5 +1,3 @@
-# Save this file as in the root of the cloned Surya repository
-
 import gradio as gr
 import torch
 from huggingface_hub import snapshot_download
@@ -20,18 +18,15 @@ import matplotlib.pyplot as plt
 import sunpy.visualization.colormaps as sunpy_cm
 import traceback
 
-# --- Use the official Surya modules ---
 from surya.models.helio_spectformer import HelioSpectFormer
 from surya.utils.data import build_scalers
 from surya.datasets.helio import inverse_transform_single_channel
 
-# --- Configuration ---
 warnings.filterwarnings("ignore", category=UserWarning, module='sunpy')
 warnings.filterwarnings("ignore", category=FutureWarning)
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Global cache for model, config, etc.
 APP_CACHE = {}
 
 SDO_CHANNELS_MAP = {
@@ -45,13 +40,12 @@ SDO_CHANNELS_MAP = {
     "aia1600": (a.Wavelength(1600 * u.angstrom), a.Sample(24 * u.s)),
     "hmi_m": (a.Physobs("intensity"), a.Sample(45 * u.s)),
     "hmi_bx": (a.Physobs("los_magnetic_field"), a.Sample(720 * u.s)),
-    "hmi_by": (a.Physobs("los_magnetic_field"), a.Sample(720 * u.s)), # Placeholder
-    "hmi_bz": (a.Physobs("los_magnetic_field"), a.Sample(720 * u.s)), # Placeholder
+    "hmi_by": (a.Physobs("los_magnetic_field"), a.Sample(720 * u.s)),
+    "hmi_bz": (a.Physobs("los_magnetic_field"), a.Sample(720 * u.s)),
     "hmi_v": (a.Physobs("los_velocity"), a.Sample(45 * u.s)),
 }
 SDO_CHANNELS = list(SDO_CHANNELS_MAP.keys())
 
-# --- 1. Model Loading and Setup ---
 def setup_and_load_model():
     if "model" in APP_CACHE:
         yield "Model already loaded. Skipping setup."
@@ -70,7 +64,18 @@ def setup_and_load_model():
     
     yield "Initializing model architecture..."
     model_config = config["model"]
-    model = HelioSpectFormer(...) # Full model definition
+    model = HelioSpectFormer(
+        img_size=model_config["img_size"], patch_size=model_config["patch_size"],
+        in_chans=len(config["data"]["sdo_channels"]), embed_dim=model_config["embed_dim"],
+        time_embedding={"type": "linear", "time_dim": len(config["data"]["time_delta_input_minutes"])},
+        depth=model_config["depth"], n_spectral_blocks=model_config["n_spectral_blocks"],
+        num_heads=model_config["num_heads"], mlp_ratio=model_config["mlp_ratio"],
+        drop_rate=model_config["drop_rate"], dtype=torch.bfloat16,
+        window_size=model_config["window_size"], dp_rank=model_config["dp_rank"],
+        learned_flow=model_config["learned_flow"], use_latitude_in_learned_flow=model_config["learned_flow"],
+        init_weights=False, checkpoint_layers=list(range(model_config["depth"])),
+        rpe=model_config["rpe"], ensemble=model_config["ensemble"], finetune=model_config["finetune"],
+    )
     device = "cuda" if torch.cuda.is_available() else "cpu"
     APP_CACHE["device"] = device
     
@@ -82,13 +87,12 @@ def setup_and_load_model():
     APP_CACHE["model"] = model
     yield "✅ Model setup complete."
 
-# --- 2. Live Data Fetching and Preprocessing (as a generator) ---
 def fetch_and_process_sdo_data(target_dt, forecast_horizon_minutes):
     config = APP_CACHE["config"]
     img_size = config["model"]["img_size"]
     
     input_deltas = config["data"]["time_delta_input_minutes"]
-    target_delta = forecast_horizon_minutes # Use user-provided horizon
+    target_delta = forecast_horizon_minutes
     input_times = [target_dt + datetime.timedelta(minutes=m) for m in input_deltas]
     target_time = target_dt + datetime.timedelta(minutes=target_delta)
     all_times = sorted(list(set(input_times + [target_time])))
@@ -107,12 +111,11 @@ def fetch_and_process_sdo_data(target_dt, forecast_horizon_minutes):
                 if data_maps[t].get("hmi_bx"): data_maps[t][channel] = data_maps[t]["hmi_bx"]
                 continue
             
-            # *** FIX: Use a.Time.nearest=True for robust fetching instead of a time window ***
             instrument = a.Instrument.hmi if "hmi" in channel else a.Instrument.aia
             query = Fido.search(a.Time(t), instrument, physobs, sample, a.Time.nearest==True)
             
             if not query: raise ValueError(f"No data found for {channel} near {t}")
-            files = Fido.fetch(query, path="./data/sdo_cache") # Fetch the entire result
+            files = Fido.fetch(query, path="./data/sdo_cache")
             data_maps[t][channel] = sunpy.map.Map(files[0])
             
     yield "✅ All files downloaded. Starting preprocessing..."
@@ -146,22 +149,42 @@ def fetch_and_process_sdo_data(target_dt, forecast_horizon_minutes):
 
     yield (input_tensor, last_input_map, target_map)
 
-
-# --- 3. Inference and Visualization ---
 def run_inference(input_tensor):
-    # This function remains the same
-    ...
+    model = APP_CACHE["model"]
+    device = APP_CACHE["device"]
+    time_deltas = APP_CACHE["config"]["data"]["time_delta_input_minutes"]
+    time_delta_tensor = torch.tensor(time_deltas, dtype=torch.float32).unsqueeze(0).to(device)
+    input_batch = {"ts": input_tensor.to(device), "time_delta_input": time_delta_tensor}
+    with torch.no_grad():
+        with torch.autocast(device_type=device.split(':')[0], dtype=torch.bfloat16):
+            prediction = model(input_batch)
+    return prediction.cpu()
 
 def generate_visualization(last_input_map, prediction_tensor, target_map, channel_name):
-    # This function remains the same
-    ...
+    if last_input_map is None: return None, None, None
+    c_idx = SDO_CHANNELS.index(channel_name)
+    scaler = APP_CACHE["scalers"]
+    all_means, all_stds, all_epsilons, all_sl_scale_factors = scaler.get_params()
+    mean, std, epsilon, sl_scale_factor = all_means[c_idx], all_stds[c_idx], all_epsilons[c_idx], all_sl_scale_factors[c_idx]
+    pred_slice = inverse_transform_single_channel(
+        prediction_tensor[0, c_idx].numpy(), mean=mean, std=std, epsilon=epsilon, sl_scale_factor=sl_scale_factor
+    )
+    vmax = np.quantile(np.nan_to_num(target_map[channel_name].data), 0.995)
+    cmap_name = f"sdoaia{channel_name.replace('aia', '')}" if 'aia' in channel_name else 'hmimag'
+    cmap = plt.get_cmap(sunpy_cm.cmlist.get(cmap_name, 'gray'))
+    def to_pil(data, flip=False):
+        data_clipped = np.nan_to_num(data)
+        data_clipped = np.clip(data_clipped, 0, vmax)
+        data_norm = data_clipped / vmax if vmax > 0 else data_clipped
+        colored = (cmap(data_norm)[:, :, :3] * 255).astype(np.uint8)
+        img = Image.fromarray(colored)
+        return img.transpose(Image.Transpose.FLIP_TOP_BOTTOM) if flip else img
+    return to_pil(last_input_map[channel_name].data, flip=True), to_pil(pred_slice, flip=True), to_pil(target_map[channel_name].data, flip=True)
 
-# --- 4. Gradio UI and Controllers ---
 def forecast_controller(date_str, hour, minute, forecast_horizon):
     yield {
         log_box: gr.update(value="Starting forecast...", visible=True),
         run_button: gr.update(interactive=False),
-        # Also disable the other controls
         date_input: gr.update(interactive=False),
         hour_slider: gr.update(interactive=False),
         minute_slider: gr.update(interactive=False),
@@ -175,11 +198,9 @@ def forecast_controller(date_str, hour, minute, forecast_horizon):
         for status in setup_and_load_model():
             yield { log_box: status }
         
-        # Construct datetime from the new UI components
         target_dt = datetime.datetime.fromisoformat(f"{date_str}T{int(hour):02d}:{int(minute):02d}:00")
         
         data_pipeline = fetch_and_process_sdo_data(target_dt, forecast_horizon)
-        # The rest of the generator logic remains the same...
         while True:
             try:
                 status = next(data_pipeline)
@@ -199,14 +220,20 @@ def forecast_controller(date_str, hour, minute, forecast_horizon):
         yield {
             log_box: f"✅ Forecast complete for {target_dt.isoformat()} (+{forecast_horizon} mins).",
             results_group: gr.update(visible=True),
-            # ... update states and images
+            state_last_input: last_input_map,
+            state_prediction: prediction_tensor,
+            state_target: target_map,
+            input_display: img_in,
+            prediction_display: img_pred,
+            target_display: img_target,
         }
 
     except Exception as e:
-        # ... error handling
+        error_str = traceback.format_exc()
+        logger.error(f"An error occurred: {e}\n{error_str}")
+        yield { log_box: f"❌ ERROR: {e}\n\nTraceback:\n{error_str}" }
     
     finally:
-        # Re-enable all controls
         yield {
             run_button: gr.update(interactive=True),
             date_input: gr.update(interactive=True),
@@ -215,14 +242,24 @@ def forecast_controller(date_str, hour, minute, forecast_horizon):
             horizon_slider: gr.update(interactive=True),
         }
 
-# --- 5. Gradio UI Definition ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    # State objects remain the same
-    ...
+    state_last_input = gr.State()
+    state_prediction = gr.State()
+    state_target = gr.State()
 
-    gr.Markdown(...) # Title remains the same
+    gr.Markdown(
+        """
+        <div align='center'>
+        # ☀️ Surya: Live Forecast Demo ☀️
+        ### Generate a real forecast for any recent date using NASA's Heliophysics Model.
+        **Instructions:**
+        1. Pick a date and time (at least 3 hours in the past).
+        2. Click 'Generate Forecast'. **This will be slow (5-15 minutes) as it downloads live data.**
+        3. Once complete, select different channels to explore the multi-spectrum forecast.
+        </div>
+        """
+    )
 
-    # --- NEW: Controls Section ---
     with gr.Accordion("Step 1: Configure Forecast", open=True):
         with gr.Row():
             date_input = gr.Textbox(
@@ -238,20 +275,19 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     
     run_button = gr.Button("🔮 Generate Forecast", variant="primary")
     
-    # --- NEW: Moved log box to its own section ---
     with gr.Accordion("Step 2: View Log", open=False) as log_accordion:
-        log_box = gr.Textbox(label="Log", interactive=False, visible=True, lines=5, max_lines=10)
+        log_box = gr.Textbox(label="Log", interactive=False, visible=False, lines=5, max_lines=10)
 
-    # --- Results section is now Step 3 ---
     with gr.Group(visible=False) as results_group:
         gr.Markdown("### Step 3: Explore Results")
-        channel_selector = gr.Dropdown(...)
+        channel_selector = gr.Dropdown(
+            choices=SDO_CHANNELS, value="aia171", label="🛰️ Select SDO Channel to Visualize"
+        )
         with gr.Row():
-            input_display = gr.Image(...)
-            prediction_display = gr.Image(...)
-            target_display = gr.Image(...)
+            input_display = gr.Image(label="Last Input to Model", height=512, width=512, interactive=False)
+            prediction_display = gr.Image(label="Surya's Forecast", height=512, width=512, interactive=False)
+            target_display = gr.Image(label="Ground Truth", height=512, width=512, interactive=False)
 
-    # --- Event Handlers ---
     run_button.click(
         fn=forecast_controller,
         inputs=[date_input, hour_slider, minute_slider, horizon_slider],
@@ -262,9 +298,12 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         ]
     )
     
-    channel_selector.change(...) # This remains the same
+    channel_selector.change(
+        fn=generate_visualization,
+        inputs=[state_last_input, state_prediction, state_target, channel_selector],
+        outputs=[input_display, prediction_display, target_display]
+    )
 
 if __name__ == "__main__":
-    # Fill in the missing ... from previous versions for the full script
-    # This is a condensed version showing only the key changes
+    os.makedirs("./data/sdo_cache", exist_ok=True)
     demo.launch(debug=True)