Update app.py

app.py

@@ -2,7 +2,6 @@
 
 import gradio as gr
 import torch
-from torch.utils.data import DataLoader
 from huggingface_hub import snapshot_download
 import yaml
 import numpy as np
@@ -10,11 +9,9 @@ from PIL import Image
 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
+from reproject import reproject_interp # Correct import statement
 import os
 import warnings
 import logging
@@ -27,7 +24,8 @@ from surya.models.helio_spectformer import HelioSpectFormer
 from surya.utils.data import build_scalers, inverse_transform_single_channel
 
 # --- Configuration ---
-warnings.filterwarnings("ignore")
+warnings.filterwarnings("ignore", category=UserWarning, module='sunpy')
+warnings.filterwarnings("ignore", category=FutureWarning)
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -94,14 +92,12 @@ def fetch_and_process_sdo_data(target_dt, progress):
     config = APP_CACHE["config"]
     img_size = config["model"]["img_size"][0]
     
-    # 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
@@ -110,53 +106,56 @@ def fetch_and_process_sdo_data(target_dt, progress):
         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')}...")
             
-            # 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)
+            search_query = [time_attr, physobs, sample]
+            # AIA and HMI queries are slightly different
+            if "aia" in channel:
+                search_query.append(a.Instrument.aia)
+            else:
+                search_query.append(a.Instrument.hmi)
+
+            query = Fido.search(*search_query)
             
             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 = {}
+    total_processing = len(all_times) * len(SDO_CHANNELS)
+    processing_done = 0
     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')}...")
+            progress(processing_done / total_processing, 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
+            if exp_time is None or 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)))
+            processing_done += 1
         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
+    input_tensor = torch.stack(input_tensor_list, dim=1).unsqueeze(0)
+    target_map = data_maps[target_time]
     last_input_map = data_maps[input_times[-1]]
 
     return input_tensor, last_input_map, target_map
@@ -191,8 +190,7 @@ def generate_visualization(last_input_map, prediction_tensor, target_map, channe
         mean=means[c_idx], std=stds[c_idx], epsilon=epsilons[c_idx], sl_scale_factor=sl_scale_factors[c_idx]
     )
     
-    # Get colormap and normalization
-    vmax = np.quantile(target_map[channel_name].data, 0.995)
+    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'))
     
@@ -206,7 +204,6 @@ def generate_visualization(last_input_map, prediction_tensor, target_map, channe
 
     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:
@@ -223,13 +220,12 @@ def forecast_controller(dt_str, progress=gr.Progress(track_tqdm=True)):
         
         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
+        return (last_input_map, prediction_tensor, target_map,
                 img_in, img_pred, img_target, status, gr.update(visible=True))
                 
     except Exception as e:
@@ -243,7 +239,6 @@ def update_visualization_controller(last_input_map, prediction_tensor, target_ma
 
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    # State objects to hold the data after a forecast is run
     state_last_input = gr.State()
     state_prediction = gr.State()
     state_target = gr.State()
@@ -263,7 +258,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     with gr.Row():
         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"))
+                                    value=(datetime.datetime.now() - datetime.timedelta(hours=3)).strftime("%Y-%m-%d %H:%M:%S"))
         run_button = gr.Button("🔮 Generate Forecast", variant="primary")
 
     with gr.Group(visible=False) as results_group:
@@ -288,6 +283,5 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     )
 
 if __name__ == "__main__":
-    # Create cache directory if it doesn't exist
     os.makedirs("./data/sdo_cache", exist_ok=True)
     demo.launch(debug=True)