Update app.py

app.py

@@ -70,18 +70,7 @@ def setup_and_load_model():
     
     yield "Initializing model architecture..."
     model_config = config["model"]
-    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"],
-    )
+    model = HelioSpectFormer(...) # Full model definition
     device = "cuda" if torch.cuda.is_available() else "cpu"
     APP_CACHE["device"] = device
     
@@ -94,13 +83,12 @@ def setup_and_load_model():
     yield "✅ Model setup complete."
 
 # --- 2. Live Data Fetching and Preprocessing (as a generator) ---
-def fetch_and_process_sdo_data(target_dt):
+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"]
-    # *** FIX: Access target_delta as an integer, not a list. Removed [0]. ***
-    target_delta = config["data"]["time_delta_target_minutes"]
+    target_delta = forecast_horizon_minutes # Use user-provided horizon
     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])))
@@ -118,13 +106,13 @@ def fetch_and_process_sdo_data(target_dt):
             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))
+            
+            # *** 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(time_attr, instrument, physobs, sample)
+            query = Fido.search(a.Time(t), instrument, physobs, sample, a.Time.nearest==True)
             
-            if not query: raise ValueError(f"No data found for {channel} at {t}")
-            files = Fido.fetch(query[0, 0], path="./data/sdo_cache")
+            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
             data_maps[t][channel] = sunpy.map.Map(files[0])
             
     yield "✅ All files downloaded. Starting preprocessing..."
@@ -161,55 +149,37 @@ def fetch_and_process_sdo_data(target_dt):
 
 # --- 3. Inference and Visualization ---
 def run_inference(input_tensor):
-    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()
+    # This function remains the same
+    ...
 
 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)
-    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)
+    # This function remains the same
+    ...
 
 # --- 4. Gradio UI and Controllers ---
-def forecast_controller(dt_str):
+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),
-        datetime_input: 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),
+        horizon_slider: gr.update(interactive=False),
         results_group: gr.update(visible=False)
     }
     
     try:
-        if not dt_str: raise gr.Error("Please select a date and time.")
+        if not date_str: raise gr.Error("Please select a date.")
         
         for status in setup_and_load_model():
             yield { log_box: status }
-
-        target_dt = datetime.datetime.fromisoformat(dt_str)
         
-        data_pipeline = fetch_and_process_sdo_data(target_dt)
+        # 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)
@@ -227,80 +197,74 @@ def forecast_controller(dt_str):
         img_in, img_pred, img_target = generate_visualization(last_input_map, prediction_tensor, target_map, "aia171")
         
         yield {
-            log_box: f"✅ Forecast complete for {target_dt.isoformat()}.",
+            log_box: f"✅ Forecast complete for {target_dt.isoformat()} (+{forecast_horizon} mins).",
             results_group: gr.update(visible=True),
-            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,
+            # ... update states and images
         }
 
     except Exception as e:
-        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}" }
+        # ... error handling
     
     finally:
+        # Re-enable all controls
         yield {
             run_button: gr.update(interactive=True),
-            datetime_input: gr.update(interactive=True)
+            date_input: gr.update(interactive=True),
+            hour_slider: gr.update(interactive=True),
+            minute_slider: gr.update(interactive=True),
+            horizon_slider: gr.update(interactive=True),
         }
 
 # --- 5. Gradio UI Definition ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    state_last_input = gr.State()
-    state_prediction = gr.State()
-    state_target = gr.State()
+    # State objects remain 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>
-        """
-    )
+    gr.Markdown(...) # Title remains the same
 
-    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=3)).strftime("%Y-%m-%d %H:%M:%S")
+    # --- NEW: Controls Section ---
+    with gr.Accordion("Step 1: Configure Forecast", open=True):
+        with gr.Row():
+            date_input = gr.Textbox(
+                label="Date",
+                value=datetime.date.today().strftime("%Y-%m-%d")
+            )
+            hour_slider = gr.Slider(label="Hour (UTC)", minimum=0, maximum=23, step=1, value=datetime.datetime.utcnow().hour - 3)
+            minute_slider = gr.Slider(label="Minute", minimum=0, maximum=59, step=1, value=datetime.datetime.utcnow().minute)
+        horizon_slider = gr.Slider(
+            label="Forecast Horizon (minutes ahead)",
+            minimum=12, maximum=120, step=12, value=12
         )
-        run_button = gr.Button("🔮 Generate Forecast", variant="primary")
-
-    log_box = gr.Textbox(label="Log", interactive=False, visible=False, lines=5, max_lines=10)
+    
+    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)
 
+    # --- Results section is now Step 3 ---
     with gr.Group(visible=False) as results_group:
-        channel_selector = gr.Dropdown(
-            choices=SDO_CHANNELS, value="aia171", label="🛰️ Select SDO Channel to Visualize"
-        )
+        gr.Markdown("### Step 3: Explore Results")
+        channel_selector = gr.Dropdown(...)
         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)
+            input_display = gr.Image(...)
+            prediction_display = gr.Image(...)
+            target_display = gr.Image(...)
 
+    # --- Event Handlers ---
     run_button.click(
         fn=forecast_controller,
-        inputs=[datetime_input],
+        inputs=[date_input, hour_slider, minute_slider, horizon_slider],
         outputs=[
-            log_box, run_button, datetime_input, results_group,
+            log_box, run_button, date_input, hour_slider, minute_slider, horizon_slider, results_group,
             state_last_input, state_prediction, state_target,
             input_display, prediction_display, target_display
         ]
     )
     
-    channel_selector.change(
-        fn=generate_visualization,
-        inputs=[state_last_input, state_prediction, state_target, channel_selector],
-        outputs=[input_display, prediction_display, target_display]
-    )
+    channel_selector.change(...) # This remains the same
 
 if __name__ == "__main__":
-    os.makedirs("./data/sdo_cache", exist_ok=True)
+    # Fill in the missing ... from previous versions for the full script
+    # This is a condensed version showing only the key changes
     demo.launch(debug=True)