add animation

app.py

@@ -1,4 +1,6 @@
 import os
+import tempfile
+from pathlib import Path
 # Set memory optimization environment variables
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
 os.environ['ANEMOI_INFERENCE_NUM_CHUNKS'] = '16'
@@ -14,6 +16,7 @@ from anemoi.inference.runners.simple import SimpleRunner
 from ecmwf.opendata import Client as OpendataClient
 import earthkit.data as ekd
 import earthkit.regrid as ekr
+import matplotlib.animation as animation
 
 # Define parameters (updating to match notebook.py)
 PARAM_SFC = ["10u", "10v", "2d", "2t", "msl", "skt", "sp", "tcw", "lsm", "z", "slor", "sdor"]
@@ -66,6 +69,11 @@ for var in ["t", "u", "v", "w", "q", "z"]:
 # Load the model once at startup
 MODEL = SimpleRunner("aifs-single-mse-1.0.ckpt", device="cuda")  # Default to CUDA
 
+# Create and set custom temp directory
+TEMP_DIR = Path("./gradio_temp")
+TEMP_DIR.mkdir(exist_ok=True)
+os.environ['GRADIO_TEMP_DIR'] = str(TEMP_DIR)
+
 def get_open_data(param, levelist=[]):
     fields = {}
     # Get the data for the current date and the previous date
@@ -90,6 +98,91 @@ def get_open_data(param, levelist=[]):
     
     return fields
 
+def plot_forecast_animation(states, selected_variable):
+    # Setup the figure and axis
+    fig = plt.figure(figsize=(15, 8))
+    ax = plt.axes(projection=ccrs.PlateCarree(central_longitude=0))
+    
+    # Get the first state to setup the plot
+    first_state = states[0]
+    latitudes, longitudes = first_state["latitudes"], first_state["longitudes"]
+    fixed_lons = np.where(longitudes > 180, longitudes - 360, longitudes)
+    triangulation = tri.Triangulation(fixed_lons, latitudes)
+    
+    # Find global min/max for consistent colorbar
+    all_values = [state["fields"][selected_variable] for state in states]
+    vmin, vmax = np.min(all_values), np.max(all_values)
+    
+    # Create a single colorbar that will be reused
+    contour = None
+    cbar_ax = None
+    
+    def update(frame):
+        nonlocal contour, cbar_ax
+        ax.clear()
+        
+        # Set map features
+        ax.set_global()
+        ax.set_extent([-180, 180, -85, 85], crs=ccrs.PlateCarree())
+        ax.coastlines(resolution='50m')
+        ax.add_feature(cfeature.BORDERS, linestyle=":", alpha=0.5)
+        ax.gridlines(draw_labels=True)
+        
+        state = states[frame]
+        values = state["fields"][selected_variable]
+        
+        # Clear the previous colorbar axis if it exists
+        if cbar_ax:
+            cbar_ax.remove()
+        
+        # Create new contour plot
+        contour = ax.tricontourf(triangulation, values,
+                                levels=20, transform=ccrs.PlateCarree(),
+                                cmap='RdBu_r', vmin=vmin, vmax=vmax)
+        
+        # Create new colorbar
+        cbar_ax = fig.add_axes([0.1, 0.05, 0.8, 0.03])  # [left, bottom, width, height]
+        plt.colorbar(contour, cax=cbar_ax, orientation='horizontal')
+        
+        # Format the date string properly
+        forecast_time = state["date"]
+        if isinstance(forecast_time, str):
+            try:
+                forecast_time = datetime.datetime.strptime(forecast_time, "%Y-%m-%d %H:%M:%S")
+            except ValueError:
+                try:
+                    forecast_time = datetime.datetime.strptime(forecast_time, "%Y-%m-%d %H:%M:%S.%f")
+                except ValueError:
+                    forecast_time = DEFAULT_DATE
+        
+        time_str = forecast_time.strftime("%Y-%m-%d %H:%M UTC")
+        
+        # Get variable description from VARIABLE_GROUPS
+        var_desc = None
+        for group in VARIABLE_GROUPS.values():
+            if selected_variable in group:
+                var_desc = group[selected_variable]
+                break
+        var_name = var_desc if var_desc else selected_variable
+        
+        ax.set_title(f"{var_name} - {time_str}")
+    
+    # Create animation
+    anim = animation.FuncAnimation(
+        fig, update,
+        frames=len(states),
+        interval=1000,  # 1 second between frames
+        repeat=True,
+        blit=False  # Must be False to update the colorbar
+    )
+    
+    # Save as MP4
+    temp_file = str(TEMP_DIR / f"forecast_{datetime.datetime.now().timestamp()}.mp4")
+    anim.save(temp_file, writer='ffmpeg', fps=1)
+    plt.close()
+    
+    return temp_file
+
 def run_forecast(date, lead_time, device):
     # Get all required fields
     fields = {}
@@ -122,42 +215,24 @@ def run_forecast(date, lead_time, device):
     if device != MODEL.device:
         MODEL = SimpleRunner("aifs-single-mse-1.0.ckpt", device=device)
     
-    results = []
+    # Collect all states instead of just the last one
+    states = []
     for state in MODEL.run(input_state=input_state, lead_time=lead_time):
-        results.append(state)
-    return results[-1]
+        states.append(state)
+    return states
 
-def plot_forecast(state, selected_variable):
-    latitudes, longitudes = state["latitudes"], state["longitudes"]
-    values = state["fields"][selected_variable]
-    
-    # Create figure with specific projection centered on 0°
-    fig = plt.figure(figsize=(15, 8))
-    ax = plt.axes(projection=ccrs.PlateCarree(central_longitude=0))
-    
-    ax.set_global()
-    ax.set_extent([-180, 180, -90, 90], crs=ccrs.PlateCarree())
-    
-    # Add map features
-    ax.coastlines(resolution='50m')
-    ax.add_feature(cfeature.BORDERS, linestyle=":", alpha=0.5)
-    ax.gridlines(draw_labels=True)
-    
-    # Fix longitudes to be -180 to 180
-    fixed_lons = np.where(longitudes > 180, longitudes - 360, longitudes)
-    
-    # Create triangulation with fixed longitudes
-    triangulation = tri.Triangulation(fixed_lons, latitudes)
-    
-    # Create the contour plot
-    contour = ax.tricontourf(triangulation, values, levels=20, 
-                            transform=ccrs.PlateCarree(), 
-                            cmap='RdBu_r')
-    
-    plt.title(f"{selected_variable} at {state['date']}")
-    plt.colorbar(contour, orientation='horizontal', pad=0.05)
-    
-    return fig
+def update_plot(lead_time, variable):
+    cleanup_old_files()  # Clean up old files before creating new ones
+    states = run_forecast(DEFAULT_DATE, lead_time, "cuda")
+    return plot_forecast_animation(states, variable)
+
+# Add cleanup function for old files
+def cleanup_old_files():
+    # Remove files older than 1 hour
+    current_time = datetime.datetime.now().timestamp()
+    for file in TEMP_DIR.glob("*.mp4"):  # Changed from *.gif to *.mp4
+        if current_time - file.stat().st_mtime > 3600:  # 1 hour in seconds
+            file.unlink(missing_ok=True)
 
 # Create dropdown choices with groups
 DROPDOWN_CHOICES = []
@@ -173,17 +248,31 @@ with gr.Blocks(css="""
         text-align: center;
         margin-bottom: 20px;
     }
+    .subtitle {
+        font-size: 1.2em;
+        line-height: 1.5;
+        margin: 20px 0;
+    }
+    .footer {
+        text-align: center;
+        padding: 20px;
+        margin-top: 20px;
+        border-top: 1px solid #eee;
+    }
 """) as demo:
-    # Centered header section
+    # Header section
     gr.Markdown(f"""
     # AIFS Weather Forecast
     
-    Interactive visualization of ECMWF AIFS weather forecasts. Starting from the latest available data ({DEFAULT_DATE.strftime('%Y-%m-%d %H:%M UTC')}), 
+    <div class="subtitle">
+    Interactive visualization of ECMWF AIFS weather forecasts.<br>
+    Starting from the latest available data ({DEFAULT_DATE.strftime('%Y-%m-%d %H:%M UTC')}),<br>
     select how many hours ahead you want to forecast and which meteorological variable to visualize.
-    """, elem_classes=["centered-header"])
+    </div>
+    """)
     
+    # Main content
     with gr.Row():
-        # Controls column - takes up 1/3 of the width
         with gr.Column(scale=1):
             lead_time = gr.Slider(
                 minimum=6, 
@@ -195,34 +284,52 @@ with gr.Blocks(css="""
             variable = gr.Dropdown(
                 choices=DROPDOWN_CHOICES,
                 value="2t",
-                label="Select Variable to Plot",
-                info="Choose a meteorological variable to visualize"
+                label="Select Variable to Plot"
             )
-            
-            # Add buttons in a row
             with gr.Row():
                 clear_btn = gr.Button("Clear")
                 submit_btn = gr.Button("Submit", variant="primary")
         
-        # Map column - takes up 2/3 of the width
         with gr.Column(scale=2):
-            plot_output = gr.Plot()
+            animation_output = gr.Video()
     
-    # Connect the inputs to the forecast function
-    def update_plot(lead_time, variable):
-        state = run_forecast(DEFAULT_DATE, lead_time, "cuda")
-        return plot_forecast(state, variable)
+    # Footer with fork instructions and model reference
+    gr.Markdown("""
+    <div class="footer">
+    <h3>Want to run this on your own?</h3>
+    You can fork this space and run it yourself:
+    
+    1. Visit <a href="https://huggingface.co/spaces/geobase/aifs-forecast" target="_blank">https://huggingface.co/spaces/geobase/aifs-forecast</a>\n
+    2. Click the "Duplicate this Space" button in the top right\n
+    3. Select your hardware requirements (GPU recommended)\n
+    4. Wait for your copy to deploy
+    
+    <h3>Model Information</h3>
+    This demo uses the <a href="https://huggingface.co/ecmwf/aifs-single-1.0" target="_blank">AIFS Single 1.0</a> model from ECMWF, 
+    which is their first operationally supported Artificial Intelligence Forecasting System. The model produces highly skilled forecasts 
+    for upper-air variables, surface weather parameters, and tropical cyclone tracks.
+    
+    Note: If you encounter any issues with this demo, trying your own fork might work better!
+    </div>
+    """)
     
-    # Clear function to reset to defaults
     def clear():
         return [
-            12,  # Reset slider to default value
-            "2t",  # Reset dropdown to default value
-            None  # Clear the plot
+            12,
+            "2t",
+            None
         ]
     
-    # Connect the buttons
-    submit_btn.click(fn=update_plot, inputs=[lead_time, variable], outputs=plot_output)
-    clear_btn.click(fn=clear, inputs=[], outputs=[lead_time, variable, plot_output])
+    # Connect the inputs to the forecast function
+    submit_btn.click(
+        fn=update_plot, 
+        inputs=[lead_time, variable], 
+        outputs=animation_output
+    )
+    clear_btn.click(
+        fn=clear, 
+        inputs=[], 
+        outputs=[lead_time, variable, animation_output]
+    )
 
 demo.launch()