Update app.py

app.py

@@ -22,7 +22,6 @@ import sunpy.visualization.colormaps as sunpy_cm
 # --- Use the official Surya modules ---
 from surya.models.helio_spectformer import HelioSpectFormer
 from surya.utils.data import build_scalers
-# *** FIX: Corrected import location for the inverse transform function ***
 from surya.datasets.helio import inverse_transform_single_channel
 
 # --- Configuration ---
@@ -33,15 +32,17 @@ logger = logging.getLogger(__name__)
 
 # Global cache for model, config, etc.
 APP_CACHE = {}
+
+# *** FIX: Corrected the a.Wavelength calls to use astropy units ***
 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)),
+    "aia94": (a.Wavelength(94 * u.angstrom), a.Sample(12 * u.s)),
+    "aia131": (a.Wavelength(131 * u.angstrom), a.Sample(12 * u.s)),
+    "aia171": (a.Wavelength(171 * u.angstrom), a.Sample(12 * u.s)),
+    "aia193": (a.Wavelength(193 * u.angstrom), a.Sample(12 * u.s)),
+    "aia211": (a.Wavelength(211 * u.angstrom), a.Sample(12 * u.s)),
+    "aia304": (a.Wavelength(304 * u.angstrom), a.Sample(12 * u.s)),
+    "aia335": (a.Wavelength(335 * u.angstrom), a.Sample(12 * u.s)),
+    "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
@@ -101,11 +102,12 @@ def fetch_and_process_sdo_data(target_dt, progress):
     all_times = sorted(list(set(input_times + [target_time])))
 
     data_maps = {}
-    total_downloads = len(all_times) * len(SDO_CHANNELS_MAP)
+    total_downloads = len(all_times) * len(SDO_CHANNELS)
     downloads_done = 0
     for t in all_times:
         data_maps[t] = {}
         for i, (channel, (physobs, sample)) in enumerate(SDO_CHANNELS_MAP.items()):
+            downloads_done += 1
             progress(downloads_done / total_downloads, desc=f"Downloading {channel} for {t.strftime('%H:%M')}...")
             
             if channel in ["hmi_by", "hmi_bz"]: 
@@ -113,19 +115,13 @@ def fetch_and_process_sdo_data(target_dt, progress):
                 continue
 
             time_attr = a.Time(t - datetime.timedelta(minutes=10), t + datetime.timedelta(minutes=10))
-            search_query_list = [time_attr, physobs, sample]
-            if "aia" in channel:
-                search_query_list.append(a.Instrument.aia)
-            else:
-                search_query_list.append(a.Instrument.hmi)
-
-            query = Fido.search(*search_query_list)
+            instrument = a.Instrument.hmi if "hmi" in channel else a.Instrument.aia
+            query = Fido.search(time_attr, instrument, physobs, sample)
             
             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
-    
+            
     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
@@ -133,14 +129,10 @@ def fetch_and_process_sdo_data(target_dt, progress):
     output_wcs.wcs.ctype = ['HPLN-TAN', 'HPLT-TAN']
 
     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(processing_done / total_processing, desc=f"Processing {channel} for {t.strftime('%H:%M')}...")
             smap = channel_maps[channel]
-            
             reprojected_data, _ = reproject_interp(smap, output_wcs, shape_out=(img_size, img_size))
             
             exp_time = smap.meta.get('exptime', 1.0)
@@ -150,7 +142,6 @@ def fetch_and_process_sdo_data(target_dt, progress):
             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)
 
     input_tensor_list = [processed_tensors[t] for t in input_times]
@@ -178,15 +169,12 @@ def run_inference(input_tensor):
     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
+    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"][SDO_CHANNELS[c_idx]].get_params()
+    means, stds, epsilons, sl_scale_factors = APP_CACHE["scalers"][channel_name].get_params()
     pred_slice = inverse_transform_single_channel(
-        prediction_tensor[0, c_idx].numpy(),
-        mean=means, std=stds, epsilon=epsilons, sl_scale_factor=sl_scale_factors
+        prediction_tensor[0, c_idx].numpy(), mean=means, std=stds, epsilon=epsilons, sl_scale_factor=sl_scale_factors
     )
     
     vmax = np.quantile(np.nan_to_num(target_map[channel_name].data), 0.995)
@@ -206,8 +194,7 @@ def generate_visualization(last_input_map, prediction_tensor, target_map, channe
 # --- 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.")
+        if not dt_str: raise gr.Error("Please select a date and time.")
         
         progress(0, desc="Initializing...")
         setup_and_load_model(progress)
@@ -216,7 +203,6 @@ def forecast_controller(dt_str, progress=gr.Progress(track_tqdm=True)):
         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)
         
         img_in, img_pred, img_target = generate_visualization(last_input_map, prediction_tensor, target_map, "aia171")
@@ -232,8 +218,7 @@ def forecast_controller(dt_str, progress=gr.Progress(track_tqdm=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
+    if last_input_map is None: return None, None, None
     return generate_visualization(last_input_map, prediction_tensor, target_map, channel_name)
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo: