back to it! first real try

app.py

@@ -1,7 +1,142 @@
 import gradio as gr
 
-def greet(name):
-    return "Hello " + name + "!!"
+import os
+import matplotlib.pyplot as plt
+import netCDF4 as nc
+from huggingface_hub import hf_hub_download
+from utils import fig2img, create_gif_from_frames
+
+# Global dictionary to store datasets in memory
+datasets = {}
+dataset_metadata = {"CE-RPUI": {"key":"data", "features":["density", "horz_velocity", "vertical_velocity", "pressure", "energy"]}, "CE-RM": {"key":"solution", "features":["density", "horz_velocity", "vertical_velocity", "pressure", "energy"]}, "NS-PwC": {"key":"velocity", "features":["horz_velocity", "vertical_velocity", "passive_tracer"]}}
+feature_label_map = {"density": "Density", "horz_velocity": "Horizontal Velocity", "vertical_velocity": "Vertical Velocity", "pressure": "Pressure", "energy": "Energy", "passive_tracer": "Passive Tracer"}
+
+def load_dataset_into_memory(dataset_name):
+    if dataset_name not in datasets:
+        print(f"Loading dataset: {dataset_name}...")
+
+        dataset_path = hf_hub_download(
+            repo_id=f"camlab-ethz/{dataset_name}",
+            filename=dataset_metadata[dataset_name]['key'] + "_0.nc",
+            repo_type="dataset",
+            token=os.getenv("HF_TOKEN")
+        )
+
+        dataset = nc.Dataset(dataset_path)
+        datasets[dataset_name] = dataset.variables[dataset_metadata[dataset_name]['key']]  # Store only the data variable
+
+    return datasets[dataset_name]
+
+# Function to generate GIFs dynamically
+def generate_gifs(dataset_name, index):
+    data = load_dataset_into_memory(dataset_name)  # Load dataset from memory
+    first_traj = data[index]  # Extract trajectory at given index
+
+    # Extract features
+    if dataset_name == "NS-PwC":
+        features = {
+            "horz_velocity": first_traj[:, 0, :, :],
+            "vertical_velocity": first_traj[:, 1, :, :],
+            "passive_tracer": first_traj[:, 2, :, :],
+        }
+    else:
+        features = {
+            "density": first_traj[:, 0, :, :],
+            "horz_velocity": first_traj[:, 1, :, :],
+            "vertical_velocity": first_traj[:, 2, :, :],
+            "pressure": first_traj[:, 3, :, :],
+            "energy": first_traj[:, 4, :, :]
+        }
+
+    gif_paths = []
+    output_dir = "episode_gifs"
+    os.makedirs(output_dir, exist_ok=True)  # Ensure output directory exists
+
+    for feature_name, feature in features.items():
+        gif_filename = f"{output_dir}/{dataset_name.lower()}_{feature_name}_{index}.gif"
+
+        # Only generate if the GIF doesn't already exist
+        if not os.path.exists(gif_filename):
+            print(f"Generating GIF: {gif_filename}")
+            frames = []
+            for i in range(len(feature)):
+                fig = plt.figure(figsize=(1.28, 1.28), dpi=100, frameon=False)
+                ax = plt.Axes(fig, [0., 0., 1., 1.])
+                ax.set_axis_off()
+                fig.add_axes(ax)
+                
+                ax.imshow(feature[i, :, :])
+                frames.append(fig2img(fig))
+                plt.close(fig)
+
+            create_gif_from_frames(frames, gif_filename)
+        else:
+            print(f"Using cached GIF: {gif_filename}")
+
+        gif_paths.append(gif_filename)
+
+    return gif_paths  # Return paths to update the Gradio UI
+
+
+gif_css = """
+#gif-label {
+    text-align: center;
+    width: 100%;
+    display: flex;
+    justify-content: center;
+}
+"""
+
+css = """
+#label {
+    text-align: center;
+    width: 100%;
+}
+"""
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown(
+        """
+        # Poseidon: Efficient Foundation Models for PDEs
+
+        Partial Differential Equations (PDEs) are fundamental in describing various physical phenomena, such as heat conduction, fluid dynamics, and electromagnetism. Poseidon leverages advanced modeling techniques to provide efficient solutions for complex PDEs, enabling faster computations and deeper insights into these phenomena.
+        """
+    )
+
+    with gr.Row():
+        dataset_selector = gr.Dropdown(choices=["CE-RPUI", "CE-RM", "NS-PwC"], value="CE-RPUI", label="Select Dataset")
+        index_slider = gr.Slider(minimum=0, maximum=100, value=0, step=1, label="Select Index")
+
+    gif_outputs = []
+    label_outputs = []
+    with gr.Row(equal_height=True) as gif_container:
+        for _ in range(5):
+            with gr.Column(scale=1, min_width=0):
+                output = gr.Image(type="filepath", show_label=False, container=False)
+                gif_outputs.append(output)
+
+                with gr.Column():  # Wrap label inside a column
+                    label = gr.Markdown(elem_id="label")
+                    label_outputs.append(label)
+
+    def update_layout(dataset_name, index):
+        features = dataset_metadata[dataset_name]["features"]
+        gif_paths = generate_gifs(dataset_name, index)
+
+        output_values = []
+        label_values = []
+        for i in range(5):
+            if i < len(features):
+                output_values.append(gr.update(value=gif_paths[i], visible=True))
+                label_values.append(gr.update(value=f"**{feature_label_map[features[i]]}**", visible=True))
+            else:
+                output_values.append(gr.update(visible=False))
+                label_values.append(gr.update(visible=False))
+
+        return output_values + label_values
+
+    dataset_selector.change(update_layout, inputs=[dataset_selector, index_slider], outputs=gif_outputs + label_outputs)
+    index_slider.change(update_layout, inputs=[dataset_selector, index_slider], outputs=gif_outputs + label_outputs)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()
+

requirements.txt

@@ -0,0 +1,5 @@
+gradio
+matplotlib
+netCDF4
+huggingface_hub
+pillow

utils.py

@@ -0,0 +1,28 @@
+import io
+import os
+from PIL import Image
+
+def fig2img(fig):
+    """Convert a Matplotlib figure to a PIL Image and return it"""
+    # Save figure to a temporary buffer.
+    buf = io.BytesIO()
+    fig.savefig(buf)
+    buf.seek(0)
+    
+    # Create PIL image from buffer
+    img = Image.open(buf)
+    return img
+
+def create_gif_from_frames(frames, filename):
+    """Create a GIF from a list of PIL Image frames"""
+    # Create output directory if it doesn't exist
+    os.makedirs('episode_gifs', exist_ok=True)
+    
+    # Save the frames as GIF
+    frames[0].save(
+        f'{filename}',
+        save_all=True,
+        append_images=frames[1:],
+        duration=200,  # Duration for each frame in milliseconds
+        loop=0
+    )