cleanup

app.py

@@ -63,10 +63,17 @@ for var in ["t", "u", "v", "w", "q", "z"]:
         var_id = f"{var}_{level}"
         VARIABLE_GROUPS["Pressure Level Variables"][var_id] = f"{var_name} at {level}hPa"
 
+# Load the model once at startup
+MODEL = SimpleRunner("aifs-single-mse-1.0.ckpt", device="cuda")  # Default to CUDA
+
 def get_open_data(param, levelist=[]):
     fields = {}
     # Get the data for the current date and the previous date
+    myiterable = [DEFAULT_DATE - datetime.timedelta(hours=6), DEFAULT_DATE]
+    print(myiterable)
     for date in [DEFAULT_DATE - datetime.timedelta(hours=6), DEFAULT_DATE]:
+        print(f"Fetching data for {date}")
+        # sources can be seen https://earthkit-data.readthedocs.io/en/latest/guide/sources.html#id57
         data = ekd.from_source("ecmwf-open-data", date=date, param=param, levelist=levelist)
         for f in data:
             assert f.to_numpy().shape == (721, 1440)
@@ -108,9 +115,15 @@ def run_forecast(date, lead_time, device):
         fields[f"z_{level}"] = gh * 9.80665
     
     input_state = dict(date=date, fields=fields)
-    runner = SimpleRunner("aifs-single-mse-1.0.ckpt", device=device)
+    
+    # Use the global model instance
+    global MODEL
+    # If device preference changed, move model to new device
+    if device != MODEL.device:
+        MODEL = SimpleRunner("aifs-single-mse-1.0.ckpt", device=device)
+    
     results = []
-    for state in runner.run(input_state=input_state, lead_time=lead_time):
+    for state in MODEL.run(input_state=input_state, lead_time=lead_time):
         results.append(state)
     return results[-1]
 
@@ -139,20 +152,23 @@ for group_name, variables in VARIABLE_GROUPS.items():
     for var_id, desc in sorted(variables.items()):
         DROPDOWN_CHOICES.append((f"{desc} ({var_id})", var_id))
 
-def gradio_interface(date_str, lead_time, device, selected_variable):
-    try:
-        date = datetime.datetime.strptime(date_str, "%Y-%m-%d")
-    except ValueError:
-        raise gr.Error("Please enter a valid date in YYYY-MM-DD format")
-    state = run_forecast(date, lead_time, device)
+def gradio_interface(lead_time, selected_variable):
+    # Use the global latest date
+    global DEFAULT_DATE
+    state = run_forecast(DEFAULT_DATE, lead_time, "cuda")  # Always use CUDA
     return plot_forecast(state, selected_variable)
 
 demo = gr.Interface(
     fn=gradio_interface,
     inputs=[
-        gr.Textbox(value=DEFAULT_DATE.strftime("%Y-%m-%d"), label="Forecast Date (YYYY-MM-DD)"),
-        gr.Slider(minimum=6, maximum=48, step=6, value=12, label="Lead Time (Hours)"),
-        gr.Radio(choices=["cuda", "cpu"], value="cuda", label="Compute Device"),
+        gr.Slider(
+            minimum=6, 
+            maximum=48, 
+            step=6, 
+            value=12, 
+            label="Forecast Hours Ahead",
+            info=f"Latest data available from: {DEFAULT_DATE.strftime('%Y-%m-%d %H:%M UTC')}"
+        ),
         gr.Dropdown(
             choices=DROPDOWN_CHOICES,
             value="2t",  # Default to 2m temperature
@@ -162,7 +178,11 @@ demo = gr.Interface(
     ],
     outputs=gr.Plot(),
     title="AIFS Weather Forecast",
-    description="Interactive visualization of ECMWF AIFS weather forecasts. Select a date, forecast lead time, and meteorological variable to plot."
+    description=f"""
+    Interactive visualization of ECMWF AIFS weather forecasts. 
+    Starting from the latest available data ({DEFAULT_DATE.strftime('%Y-%m-%d %H:%M UTC')}), 
+    select how many hours ahead you want to forecast and which meteorological variable to visualize.
+    """
 )
 
 demo.launch()