Update app.py

app.py

@@ -2,81 +2,71 @@
 
 import gradio as gr
 import torch
-import torch.nn.functional as F
 from torch.utils.data import DataLoader
 from huggingface_hub import snapshot_download
 import yaml
 import numpy as np
 from PIL import Image
-import sunpy.visualization.colormaps as sunpy_cm
+import sunpy.map
+import sunpy.net.attrs as a
+from sunpy.net import Fido
+from sunpy.coordinates import Helioprojective
+from astropy.coordinates import SkyCoord
+from astropy.wcs import WCS
+import astropy.units as u
+from reproject import reproject_interp
 import os
-import glob
 import warnings
 import logging
+import datetime
 import matplotlib.pyplot as plt
+import sunpy.visualization.colormaps as sunpy_cm
 
-# --- Use the official Surya modules now that we are in the repo ---
-from surya.datasets.helio import HelioNetCDFDataset, inverse_transform_single_channel
+# --- Use the official Surya modules ---
 from surya.models.helio_spectformer import HelioSpectFormer
-from surya.utils.data import build_scalers, custom_collate_fn
+from surya.utils.data import build_scalers, inverse_transform_single_channel
 
-# Suppress verbose logging and warnings for a cleaner UI
+# --- Configuration ---
 warnings.filterwarnings("ignore")
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# --- Global cache to store expensive-to-load objects ---
-APP_CACHE = {
-    "model": None,
-    "config": None,
-    "scalers": None,
-    "full_results": None, # Will store all prediction results
-    "device": "cuda" if torch.cuda.is_available() else "cpu",
+# Global cache for model, config, etc.
+APP_CACHE = {}
+SDO_CHANNELS_MAP = {
+    "aia94": (a.Wavelength(94, 94, "angstrom"), a.Sample(12 * u.s)),
+    "aia131": (a.Wavelength(131, 131, "angstrom"), a.Sample(12 * u.s)),
+    "aia171": (a.Wavelength(171, 171, "angstrom"), a.Sample(12 * u.s)),
+    "aia193": (a.Wavelength(193, 193, "angstrom"), a.Sample(12 * u.s)),
+    "aia211": (a.Wavelength(211, 211, "angstrom"), a.Sample(12 * u.s)),
+    "aia304": (a.Wavelength(304, 304, "angstrom"), a.Sample(12 * u.s)),
+    "aia335": (a.Wavelength(335, 335, "angstrom"), a.Sample(12 * u.s)),
+    "aia1600": (a.Wavelength(1600, 1600, "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_v": (a.Physobs("los_velocity"), a.Sample(45 * u.s)),
 }
+SDO_CHANNELS = list(SDO_CHANNELS_MAP.keys())
 
-# SDO channels from the test script for the dropdown menu
-SDO_CHANNELS = [
-    "aia94", "aia131", "aia171", "aia193", "aia211", "aia304", "aia335",
-    "aia1600", "hmi_m", "hmi_bx", "hmi_by", "hmi_bz", "hmi_v",
-]
-
-# --- 1. Setup, Download, and Model Loading (adapting fixtures from test_surya.py) ---
-
+# --- 1. Model Loading and Setup ---
 def setup_and_load_model(progress=gr.Progress()):
-    """
-    Handles all initial setup: downloading data, loading configs, and initializing the model.
-    This function will populate the APP_CACHE.
-    """
-    if APP_CACHE["model"] is not None:
-        logger.info("Model and data already loaded. Skipping setup.")
+    if "model" in APP_CACHE:
         return
 
-    # --- Part A: Download data (from download_data fixture) ---
-    progress(0.1, desc="Downloading model weights and config...")
-    snapshot_download(
-        repo_id="nasa-ibm-ai4science/Surya-1.0",
-        local_dir="data/Surya-1.0",
-        allow_patterns=["config.yaml", "scalers.yaml", "surya.366m.v1.pt"],
-    )
-    progress(0.3, desc="Downloading validation data for 2014-01-07...")
-    snapshot_download(
-        repo_id="nasa-ibm-ai4science/Surya-1.0_validation_data",
-        repo_type="dataset",
-        local_dir="data/Surya-1.0_validation_data",
-        allow_patterns="20140107_1[5-9]??.nc",
-    )
+    progress(0.1, desc="Downloading model files (first run only)...")
+    snapshot_download(repo_id="nasa-ibm-ai4science/Surya-1.0", local_dir="data/Surya-1.0",
+                      allow_patterns=["config.yaml", "scalers.yaml", "surya.366m.v1.pt"])
 
-    # --- Part B: Load Config and Scalers (from config & scalers fixtures) ---
-    progress(0.5, desc="Loading configuration and data scalers...")
+    progress(0.5, desc="Loading configuration and scalers...")
     with open("data/Surya-1.0/config.yaml") as fp:
         config = yaml.safe_load(fp)
     APP_CACHE["config"] = config
-    
     scalers_info = yaml.safe_load(open("data/Surya-1.0/scalers.yaml", "r"))
     APP_CACHE["scalers"] = build_scalers(info=scalers_info)
-
-    # --- Part C: Initialize and load model (from model fixture and test function) ---
-    progress(0.7, desc="Initializing model architecture...")
+    
+    progress(0.7, desc="Initializing and loading model...")
     model_config = config["model"]
     model = HelioSpectFormer(
         img_size=model_config["img_size"], patch_size=model_config["patch_size"],
@@ -90,194 +80,214 @@ def setup_and_load_model(progress=gr.Progress()):
         init_weights=False, checkpoint_layers=list(range(model_config["depth"])),
         rpe=model_config["rpe"], ensemble=model_config["ensemble"], finetune=model_config["finetune"],
     )
-
-    progress(0.8, desc=f"Loading model weights to {APP_CACHE['device']}...")
-    path_weights = "data/Surya-1.0/surya.366m.v1.pt"
-    weights = torch.load(path_weights, map_location=torch.device(APP_CACHE["device"]))
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    APP_CACHE["device"] = device
+    weights = torch.load(f"data/Surya-1.0/surya.366m.v1.pt", map_location=torch.device(device))
     model.load_state_dict(weights, strict=True)
-    model.to(APP_CACHE["device"])
+    model.to(device)
     model.eval()
-    
-    n_params = sum(p.numel() for p in model.parameters()) / 1e6
-    logger.info(f"Surya FM: {n_params:.2f}M parameters loaded.")
     APP_CACHE["model"] = model
+    logger.info("Model setup complete.")
 
-# --- 2. Inference Logic (adapting the test loop) ---
-
-def run_full_forecast():
-    """
-    Runs inference on the entire validation dataset and stores results.
-    """
-    if APP_CACHE["full_results"] is not None:
-        return APP_CACHE["full_results"]
-
-    model = APP_CACHE["model"]
+# --- 2. Live Data Fetching and Preprocessing ---
+def fetch_and_process_sdo_data(target_dt, progress):
     config = APP_CACHE["config"]
-    device = APP_CACHE["device"]
-
-    # Create the index file needed by the dataset loader
-    os.makedirs("tests", exist_ok=True)
-    with open("tests/test_surya_index.csv", "w") as f:
-        f.write("path\n")
-        search_path = os.path.join("data/Surya-1.0_validation_data", "**", "*.nc")
-        for nc_file in sorted(glob.glob(search_path, recursive=True)):
-            f.write(f"{nc_file}\n")
+    img_size = config["model"]["img_size"][0]
     
-    # Setup dataset and dataloader (from dataset & dataloader fixtures)
-    dataset = HelioNetCDFDataset(
-        index_path="tests/test_surya_index.csv",
-        time_delta_input_minutes=config["data"]["time_delta_input_minutes"],
-        time_delta_target_minutes=config["data"]["time_delta_target_minutes"],
-        n_input_timestamps=len(config["data"]["time_delta_input_minutes"]),
-        rollout_steps=1, channels=config["data"]["sdo_channels"],
-        scalers=APP_CACHE["scalers"], phase="valid",
-    )
-    dataloader = DataLoader(
-        dataset, shuffle=False, batch_size=1, num_workers=2,
-        pin_memory=True, collate_fn=custom_collate_fn
-    )
-
-    all_results = []
-    with torch.no_grad():
-        for batch_data, batch_metadata in dataloader:
-            input_batch = {k: v.to(device) for k, v in batch_data.items() if k in ["ts", "time_delta_input"]}
+    # Define time windows for input and target (ground truth)
+    input_deltas = config["data"]["time_delta_input_minutes"]
+    target_delta = config["data"]["time_delta_target_minutes"][0]
+    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])))
+
+    # Download data for all required timestamps
+    data_maps = {}
+    total_downloads = len(all_times) * len(SDO_CHANNELS_MAP)
+    downloads_done = 0
+    for t in all_times:
+        data_maps[t] = {}
+        for i, (channel, (physobs, sample)) in enumerate(SDO_CHANNELS_MAP.items()):
+            progress(downloads_done / total_downloads, desc=f"Downloading {channel} for {t.strftime('%H:%M')}...")
             
-            with torch.autocast(device_type=device.split(':')[0], dtype=torch.bfloat16):
-                prediction = model(input_batch)
+            # HMI vector fields are not standard products, use LoS as a placeholder for demo
+            instrument = a.Instrument.hmi if "hmi" in channel else a.Instrument.aia
+            if channel in ["hmi_by", "hmi_bz"]: 
+                if data_maps[t].get("hmi_bx"): data_maps[t][channel] = data_maps[t]["hmi_bx"]
+                continue
+
+            time_attr = a.Time(t - datetime.timedelta(minutes=10), t + datetime.timedelta(minutes=10))
+            query = Fido.search(time_attr, a.Instrument.aia, physobs, sample) if "aia" in channel else Fido.search(time_attr, a.Instrument.hmi, physobs, sample)
             
-            # Store the relevant tensors on CPU for later visualization
-            result = {
-                "input": input_batch["ts"].cpu(),
-                "prediction": prediction.cpu(),
-                "target": batch_data["forecast"].cpu(),
-                "input_timestamp": np.datetime_as_string(batch_metadata["timestamps_input"][0][-1], unit='s'),
-                "target_timestamp": np.datetime_as_string(batch_metadata["timestamps_targets"][0][0], unit='s'),
-            }
-            all_results.append(result)
-
-    APP_CACHE["full_results"] = all_results
-    # Cache scalers needed for visualization
-    APP_CACHE["scalers_vis"] = dataset.transformation_inputs()
-    return all_results
-
-# --- 3. Visualization Logic ---
-
-def generate_visualization(results, timestep_index, channel_name):
-    """
-    Generates PIL images for a specific timestep and channel from the results.
-    """
-    if not results:
-        return None, None, None, "No results available. Please run the forecast.", ""
-
-    timestep_data = results[timestep_index]
+            if not query: raise ValueError(f"No data found for {channel} at {t}")
+            files = Fido.fetch(query[0, 0], path="./data/sdo_cache")
+            data_maps[t][channel] = sunpy.map.Map(files[0])
+            downloads_done += 1
+    
+    # Create target WCS for reprojection
+    output_wcs = WCS(naxis=2)
+    output_wcs.wcs.crpix = [(img_size + 1) / 2, (img_size + 1) / 2]
+    output_wcs.wcs.cdelt = np.array([-1.2, 1.2]) * u.arcsec
+    output_wcs.wcs.crval = [0, 0] * u.arcsec
+    output_wcs.wcs.ctype = ['HPLN-TAN', 'HPLT-TAN']
+
+    # Process data
+    processed_tensors = {}
+    for t, channel_maps in data_maps.items():
+        channel_tensors = []
+        for i, channel in enumerate(SDO_CHANNELS):
+            progress(i / len(SDO_CHANNELS), desc=f"Processing {channel} for {t.strftime('%H:%M')}...")
+            smap = channel_maps[channel]
+            
+            # Reproject to common grid
+            reprojected_data, _ = reproject_interp(smap, output_wcs, shape_out=(img_size, img_size))
+            
+            # Normalize by exposure time and apply signed-log transform
+            exp_time = smap.meta.get('exptime', 1.0)
+            if exp_time <= 0: exp_time = 1.0
+            norm_data = reprojected_data / exp_time
+
+            # Apply the same scaling as the training pipeline
+            scaler = APP_CACHE["scalers"][channel]
+            scaled_data = scaler.transform(norm_data)
+            channel_tensors.append(torch.from_numpy(scaled_data.astype(np.float32)))
+        processed_tensors[t] = torch.stack(channel_tensors)
+
+    # Assemble final input and target tensors
+    input_tensor_list = [processed_tensors[t] for t in input_times]
+    input_tensor = torch.stack(input_tensor_list, dim=1).unsqueeze(0) # Add batch dim
+    target_map = data_maps[target_time] # Return raw map for ground truth vis
+    last_input_map = data_maps[input_times[-1]]
+
+    return input_tensor, last_input_map, target_map
+
+# --- 3. Inference and Visualization ---
+def run_inference(input_tensor):
+    logger.info("Running model inference...")
+    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)
+    logger.info("Inference complete.")
+    return prediction.cpu()
+
+def generate_visualization(last_input_map, prediction_tensor, target_map, channel_name):
+    if last_input_map is None:
+        return None, None, None
+    
     c_idx = SDO_CHANNELS.index(channel_name)
-    means, stds, epsilons, sl_scale_factors = APP_CACHE["scalers_vis"]
-
-    # Denormalize data for visualization
-    input_slice = inverse_transform_single_channel(
-        timestep_data["input"][0, c_idx, -1].numpy(),
-        mean=means[c_idx], std=stds[c_idx], epsilon=epsilons[c_idx], sl_scale_factor=sl_scale_factors[c_idx]
-    )
+    
+    # Process Prediction
+    means, stds, epsilons, sl_scale_factors = APP_CACHE["scalers"][SDO_CHANNELS[0]].get_params()
     pred_slice = inverse_transform_single_channel(
-        timestep_data["prediction"][0, c_idx].numpy(),
-        mean=means[c_idx], std=stds[c_idx], epsilon=epsilons[c_idx], sl_scale_factor=sl_scale_factors[c_idx]
-    )
-    target_slice = inverse_transform_single_channel(
-        timestep_data["target"][0, c_idx, 0].numpy(),
+        prediction_tensor[0, c_idx].numpy(),
         mean=means[c_idx], std=stds[c_idx], epsilon=epsilons[c_idx], sl_scale_factor=sl_scale_factors[c_idx]
     )
-
-    # Convert to PIL Images using appropriate colormaps
-    vmax = np.quantile(target_slice, 0.995)
+    
+    # Get colormap and normalization
+    vmax = np.quantile(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):
-        data_clipped = np.clip(data, 0, vmax)
-        data_norm = data_clipped / vmax
-        return Image.fromarray((cmap(data_norm)[:, :, :3] * 255).astype(np.uint8)).transpose(Image.Transpose.TRANSPOSE)
-
-    status_text = (f"Displaying Timestep {timestep_index+1}/{len(results)}\n"
-                   f"Input: {timestep_data['input_timestamp']} | Forecast/Target: {timestep_data['target_timestamp']}")
+    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), to_pil(pred_slice, flip=True), to_pil(target_map[channel_name].data)
+
+
+# --- 4. Gradio UI and Controllers ---
+def forecast_controller(dt_str, progress=gr.Progress(track_tqdm=True)):
+    try:
+        if not dt_str:
+            raise gr.Error("Please select a date and time.")
+        
+        progress(0, desc="Initializing...")
+        setup_and_load_model(progress)
+        
+        target_dt = datetime.datetime.fromisoformat(dt_str)
+        logger.info(f"Starting forecast for target time: {target_dt}")
+
+        input_tensor, last_input_map, target_map = fetch_and_process_sdo_data(target_dt, progress)
+        
+        prediction_tensor = run_inference(input_tensor)
+        
+        # Default visualization for aia171
+        img_in, img_pred, img_target = generate_visualization(last_input_map, prediction_tensor, target_map, "aia171")
+        
+        status = f"Forecast complete for {target_dt.isoformat()}. Ready to explore channels."
+        logger.info(status)
+        
+        return (last_input_map, prediction_tensor, target_map, # state
+                img_in, img_pred, img_target, status, gr.update(visible=True))
+                
+    except Exception as e:
+        logger.error(f"An error occurred: {e}", exc_info=True)
+        raise gr.Error(f"Failed to generate forecast. Error: {e}")
+
+def update_visualization_controller(last_input_map, prediction_tensor, target_map, channel_name):
+    if last_input_map is None:
+        return None, None, None
+    return generate_visualization(last_input_map, prediction_tensor, target_map, channel_name)
 
-    return to_pil(input_slice), to_pil(pred_slice), to_pil(target_slice), status_text
 
-# --- 4. Gradio Controller Functions ---
-
-def forecast_controller(progress=gr.Progress(track_tqdm=True)):
-    """Main function for the 'Generate Forecast' button."""
-    progress(0, desc="Starting setup...")
-    setup_and_load_model(progress)
-    
-    logger.info("Running forecast on all validation timesteps...")
-    progress(0.9, desc="Running inference on validation data...")
-    results = run_full_forecast()
-    logger.info(f"Forecast complete. {len(results)} timesteps processed.")
-    
-    # Generate the first visualization
-    img_in, img_pred, img_target, status = generate_visualization(results, 0, SDO_CHANNELS[2]) # Default to aia171
-    
-    # Update the slider to be interactive and have the correct number of steps
-    slider_update = gr.Slider(minimum=1, maximum=len(results), step=1, value=1, interactive=True,
-                              label="Forecast Timestep")
-                              
-    return results, img_in, img_pred, img_target, status, slider_update
-
-def update_visualization_controller(results, timestep, channel):
-    """Called when a slider or dropdown is changed."""
-    return generate_visualization(results, timestep - 1, channel)
-
-
-# --- 5. Gradio UI Layout ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    # State object to hold the results of the full inference run
-    state_results = gr.State()
+    # State objects to hold the data after a forecast is run
+    state_last_input = gr.State()
+    state_prediction = gr.State()
+    state_target = gr.State()
 
     gr.Markdown(
         """
         <div align='center'>
-        # ☀️ Surya: Live Model Demo ☀️
-        ### An Interactive Interface for NASA's Heliophysics Foundation Model
-        This demo runs the **actual** Surya model on its official validation data for **2014-01-07**.
-        <br>
-        **Instructions:** 1. Click 'Generate Forecast'. 2. Use the controls to explore the results.
+        # ☀️ 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 1 hour 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>
         """
     )
 
     with gr.Row():
-        with gr.Column(scale=1):
-            run_button = gr.Button("🔮 1. Generate Full Forecast", variant="primary")
-            status_box = gr.Textbox(label="Status", interactive=False, value="Ready.", lines=2)
-            channel_selector = gr.Dropdown(
-                choices=SDO_CHANNELS, value="aia171", label="🛰️ 2. Select SDO Channel"
-            )
-            timestep_slider = gr.Slider(
-                minimum=1, maximum=8, step=1, value=1, interactive=False, label="Forecast Timestep"
-            )
-        with gr.Column(scale=3):
-             with gr.Row():
-                input_display = gr.Image(label="Last Input", height=512, width=512, interactive=False)
-                prediction_display = gr.Image(label="Model Forecast", height=512, width=512, interactive=False)
-                target_display = gr.Image(label="Ground Truth", height=512, width=512, interactive=False)
+        datetime_input = gr.Textbox(label="Enter Forecast Start Time (YYYY-MM-DD HH:MM:SS)",
+                                    value=(datetime.datetime.now() - datetime.timedelta(hours=2)).strftime("%Y-%m-%d %H:%M:%S"))
+        run_button = gr.Button("🔮 Generate Forecast", variant="primary")
+
+    with gr.Group(visible=False) as results_group:
+        status_box = gr.Textbox(label="Status", interactive=False)
+        channel_selector = gr.Dropdown(choices=SDO_CHANNELS, value="aia171", label="🛰️ Select SDO Channel")
+        with gr.Row():
+            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,
-        outputs=[state_results, input_display, prediction_display, target_display, status_box, timestep_slider]
+        inputs=[datetime_input],
+        outputs=[state_last_input, state_prediction, state_target,
+                 input_display, prediction_display, target_display, status_box, results_group]
     )
     
-    # When the user changes the channel or timestep, call the visualization update function
     channel_selector.change(
         fn=update_visualization_controller,
-        inputs=[state_results, timestep_slider, channel_selector],
-        outputs=[input_display, prediction_display, target_display, status_box]
-    )
-    timestep_slider.change(
-        fn=update_visualization_controller,
-        inputs=[state_results, timestep_slider, channel_selector],
-        outputs=[input_display, prediction_display, target_display, status_box]
+        inputs=[state_last_input, state_prediction, state_target, channel_selector],
+        outputs=[input_display, prediction_display, target_display]
     )
 
 if __name__ == "__main__":
+    # Create cache directory if it doesn't exist
+    os.makedirs("./data/sdo_cache", exist_ok=True)
     demo.launch(debug=True)