labels

app.py

@@ -23,6 +23,43 @@ LEVELS = [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 150, 100, 50]
 SOIL_LEVELS = [1, 2]
 DEFAULT_DATE = OpendataClient().latest()
 
+# First define the variable descriptions
+VARIABLE_DESCRIPTIONS = {
+    # Surface variables (10m)
+    "10u": "10m U Wind Component",
+    "10v": "10m V Wind Component",
+    "2d": "2m Dewpoint Temperature",
+    "2t": "2m Temperature",
+    "msl": "Mean Sea Level Pressure",
+    "skt": "Skin Temperature",
+    "sp": "Surface Pressure",
+    "tcw": "Total Column Water",
+    "lsm": "Land-Sea Mask",
+    "z": "Surface Geopotential",
+    "slor": "Slope of Sub-gridscale Orography",
+    "sdor": "Standard Deviation of Orography",
+    
+    # Soil variables
+    "stl1": "Soil Temperature Level 1",
+    "stl2": "Soil Temperature Level 2",
+    "swvl1": "Soil Water Volume Level 1",
+    "swvl2": "Soil Water Volume Level 2",
+}
+
+# Add pressure level variable descriptions dynamically
+for var in ["t", "u", "v", "w", "q", "z"]:
+    var_name = {
+        "t": "Temperature",
+        "u": "U Wind Component",
+        "v": "V Wind Component",
+        "w": "Vertical Velocity",
+        "q": "Specific Humidity",
+        "z": "Geopotential"
+    }[var]
+    
+    for level in LEVELS:
+        VARIABLE_DESCRIPTIONS[f"{var}_{level}"] = f"{var_name} at {level}hPa"
+
 def get_open_data(param, levelist=[]):
     fields = {}
     # Get the data for the current date and the previous date
@@ -74,36 +111,59 @@ def run_forecast(date, lead_time, device):
         results.append(state)
     return results[-1]
 
-def plot_forecast(state):
+def plot_forecast(state, selected_variable):
     latitudes, longitudes = state["latitudes"], state["longitudes"]
-    values = state["fields"]["100u"]
+    values = state["fields"][selected_variable]
     fig, ax = plt.subplots(figsize=(11, 6), subplot_kw={"projection": ccrs.PlateCarree()})
     ax.coastlines()
     ax.add_feature(cfeature.BORDERS, linestyle=":")
     triangulation = tri.Triangulation(longitudes, latitudes)
     contour = ax.tricontourf(triangulation, values, levels=20, transform=ccrs.PlateCarree(), cmap="RdBu")
-    plt.title(f"100m winds at {state['date']}")
+    plt.title(f"{selected_variable} at {state['date']}")
     plt.colorbar(contour)
     return fig
 
-def gradio_interface(date_str, lead_time, device):
+# Then create the available variables list
+AVAILABLE_VARIABLES = (
+    # Surface variables
+    ["10u", "10v", "2d", "2t", "msl", "skt", "sp", "tcw", "lsm", "z", "slor", "sdor"] +
+    # Soil variables
+    ["stl1", "stl2", "swvl1", "swvl2"] +
+    # Pressure level variables (adding level suffix)
+    [f"{var}_{level}" for var in ["t", "u", "v", "w", "q", "z"] 
+     for level in LEVELS]
+)
+
+# Finally create the dropdown choices
+DROPDOWN_CHOICES = [
+    (f"{VARIABLE_DESCRIPTIONS[var_id]} ({var_id})", var_id)
+    for var_id in sorted(AVAILABLE_VARIABLES)
+]
+
+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)
-    return plot_forecast(state)
+    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.Radio(choices=["cuda", "cpu"], value="cuda", label="Compute Device"),
+        gr.Dropdown(
+            choices=DROPDOWN_CHOICES,
+            value="t_850",  # This should be the variable ID
+            label="Select Variable to Plot",
+            info="Choose a meteorological variable to visualize"
+        )
     ],
     outputs=gr.Plot(),
     title="AIFS Weather Forecast",
-    description="Run ECMWF AIFS forecasts based on selected parameters."
+    description="Interactive visualization of ECMWF AIFS weather forecasts. Select a date, forecast lead time, and meteorological variable to plot."
 )
 
 demo.launch()