add comments + fix typos

app.py

@@ -1,3 +1,18 @@
+# All the datasets will use the same format: a collection of HDF5 files with data cubes
+# in t0_fields: scalar fields, like density, pressure, energy
+#  the data is of shape (n_trajectories, n_time_steps, x, y)
+# in t1_fields: vector fields, like velocity (size=2 => vx, vy)
+#  the data is of shape (n_trajectories, n_time_steps, x, y, vx/vy)
+# in t2_fields: tensor fields, like ???
+#  the data is of shape (n_trajectories, n_time_steps, x, y, d1, d2), with d1, d2 in [0, 1]
+#   ie, instead of 1 additional dimension for velocity: a (2,2) matrix where each component
+#   (0,0),(1,0),(0,1),(1,1) can be plotted
+# Size:
+# - n_trajectories: 8 to 256
+# - n_time_steps: 101
+# - x: 128 to 512
+# - y: 128 to 512
+# - physical fields: 2 to 8 (density, pressure, energy, velocity…)
 from functools import lru_cache
 
 import gradio as gr
@@ -14,16 +29,18 @@ paths = fs.glob(set_path)
 files = {path: h5py.File(fs.open(path, "rb", cache_type="none"), "r") for path in paths}
 
 def get_scalar_fields(path: str) -> list[str]:
+    # TODO: support t1_fields (vector) and t2_fields (tensor)
     return list(files[path]["t0_fields"].keys())
 
 def get_trajectories(path: str, field: str) -> list[int]:
+    # The first dimension is the trajectory (8 to 256)
     return list(range(len(files[path]["t0_fields"][field])))
 
 @lru_cache(maxsize=4)
 def get_images(path: str, scalar_field: str, trajectory: int) -> list[Image.Image]:
-    #The data is of shape (n_trajectories, n_timesteps, x, y)
+    # The data is of shape (n_trajectories, n_time_steps, x, y)
     out = files[path]["t0_fields"][scalar_field][trajectory]
-    out = np.log(out)
+    out = np.log(out) # not sure why
     out = (out - out.min()) / (out.max() - out.min())
     out = np.uint8(cm.RdBu_r(out) * 255)
     return [Image.fromarray(img) for img in out]
@@ -36,11 +53,11 @@ with gr.Blocks() as demo:
     gr.Markdown(f"# 💠 HDF5 Viewer for the [{repo_id}](https://huggingface.co/datasets/{repo_id}) Dataset 🌊")
     gr.Markdown(f"Showing files at `{set_path}`")
     with gr.Row():
-        files_dropdown = gr.Dropdown(choices=paths, value=paths[0], label="file", scale=4)
-        scalar_fields_dropdown = gr.Dropdown(choices=default_scalar_fields, value=default_scalar_fields[0], label="scalar field")
-        trajectory_dropdown = gr.Dropdown(choices=default_trajectories, value=default_trajectories[0], label="sample")
+        files_dropdown = gr.Dropdown(choices=paths, value=paths[0], label="File", scale=4)
+        scalar_fields_dropdown = gr.Dropdown(choices=default_scalar_fields, value=default_scalar_fields[0], label="Physical field")
+        trajectory_dropdown = gr.Dropdown(choices=default_trajectories, value=default_trajectories[0], label="Trajectory")
     gallery = gr.Gallery(default_images, preview=True, selected_index=len(default_images) // 2)
-    gr.Markdown("_Tip: click on the image to go forward or backards_")
+    gr.Markdown("_Tip: click on the image to go forward or backwards_")
 
     @files_dropdown.select(inputs=[files_dropdown], outputs=[scalar_fields_dropdown, trajectory_dropdown, gallery])
     def _update_file(path: str):