Merge internal branch into public

Data:
- Add Rayleigh Benard uniform dataset
- Edit information about Shear Flow data

Statistics and Metrics:
- Add RMS statistics
- Add Pearson correlation metrics

Code Refactoring:
- Refine video generation control
- Refactor sample load from HDF5
- Add transformation and augmentation based on resizing and roation
- Allow specifying for dataset split

Author: LTMeyer

.gitignore

@@ -26,6 +26,7 @@ datasets/viscoelastic_instability/data/
 2D/planetswe/data/
 
 datasets/rayleigh_benard/data/
+datasets/rayleigh_benard_uniform/data/
 datasets/acoustic_scattering_inclusions/old_and_problematic_data/
 datasets/shear_flow/data/
 datasets/supernova_explosion_128/data/

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.6.4
+  rev: v0.9.9
   hooks:
     - id: ruff
       args: [--fix]

datasets/rayleigh_benard/README.md

@@ -5,6 +5,7 @@
 **Longer description of the data:**
 Rayleigh-Bénard convection involves fluid dynamics and thermodynamics, seen in a horizontal fluid layer heated from below, forming convective cells due to a temperature gradient. With the lower plate heated and the upper cooled, thermal energy creates density variations, initiating fluid motion. This results in Bénard cells, showcasing warm fluid rising and cool fluid descending. The interplay of buoyancy, conduction, and viscosity leads to complex fluid motion, including vortices and boundary layers.
 
+Note the `rayleigh_benard` and the `rayleigh_benard_uniform` are distinct only by resampling the data to different grids.
 
 **Associated paper**: [Paper 1](https://www.tandfonline.com/doi/pdf/10.1080/14786441608635602), [Paper 2](https://www.cambridge.org/core/journals/journal-of-fluid-mechanics/article/steady-rayleighbenard-convection-between-noslip-boundaries/B4F358EB0AE83BBE9D85968DC5DDD64D
 ).
@@ -21,19 +22,19 @@ $$
 \begin{align*}
 \frac{\partial b}{\partial t} - \kappa \Delta b & = -u\nabla b,
 \\
-\frac{\partial u}{\partial t} - \nu \Delta u + \nabla p - b \vec{e}_z & = -u \nabla u,
+\frac{\partial u}{\partial t} - \nu \Delta u + \nabla p - b \vec{e}_y & = -u \nabla u,
 \end{align*}
 $$
 
-where $\Delta = \nabla \cdot \nabla$ is the spatial Laplacian, $b$ is the buoyancy, $u = (u_x,u_y)$ the (horizontal and vertical) velocity, and $p$ is the pressure, $\vec{e}_z$ is the unit vector in the vertical direction, with the additional constraints $\int p = 0$ (pressure gauge).
+where $\Delta = \nabla \cdot \nabla$ is the spatial Laplacian, $b$ is the buoyancy, $u = (u_x,u_y)$ the (horizontal and vertical) velocity, and $p$ is the pressure, $\vec{e}_y$ is the unit vector in the vertical direction, with the additional constraints $\int p = 0$ (pressure gauge).
 
 The boundary conditions vertically are as follows:
 
 $$
 \begin{align*}
-b(z=0) = Lz ~~~,~~~ b(z=Lz) = 0
+b(y=0) = Ly ~~~,~~~ b(y=Ly) = 0
 \\
-u(z=0) = u(z=Lz) = 0
+u(y=0) = u(y=Ly) = 0
 \end{align*}
 $$
 
@@ -70,9 +71,9 @@ Table: VRMSE metrics on test sets (lower is better). Best results are shown in b
 
 **Size of the ensemble of all simulations:** 342 GB.
 
-**Grid type:** uniform, cartesian coordinates.
+**Grid type:** Cartesian coordinates. Sampled uniformly along the $x$ axis. Sampled at Chebyshev nodes along the $y$ axis.
 
-**Initial conditions:** the buoyancy is composed of a dumped noise added to a linear background  $b(t=0) = (Lz-z)\times\delta b_0 + z(Lz-z) \times\epsilon$ where $\epsilon$ is a Gaussian white noise of scale $10^{-3}$.
+**Initial conditions:** the buoyancy is composed of a dumped noise added to a linear background  $b(t=0) = (Ly-y)\times\delta b_0 + y(Ly-y) \times\epsilon$ where $\epsilon$ is a Gaussian white noise of scale $10^{-3}$.
 The other fields $u$ and $p$ are initialized to $0$.
 
 **Boundary conditions:** periodic on the horizontal direction, Dirichlet conditions on the vertical direction.
@@ -81,7 +82,7 @@ The other fields $u$ and $p$ are initialized to $0$.
 
 **Total time range ($t_{min}$ to $t_{max}$):** $t_{min} = 0$, $t_{max} = 50$.
 
-**Spatial domain size:** $0 \leq x \leq 4$ horizontally, and $0 \leq z \leq 1$ vertically.
+**Spatial domain size:** $0 \leq x \leq 4$ horizontally, and $0 \leq y \leq 1$ vertically.
 
 **Set of coefficients or non-dimensional parameters evaluated:** $\text{Rayleigh}\in[1e6,1e7,1e8,1e9,1e10], \text{Prandtl}\in[0.1,0.2,0.5,1.0,2.0,5.0,10.0]$. For initial conditions $\delta b_0\in[0.2,0.4,0.6,0.8,1.0]$, the seed used to generate the initial Gaussian white noise are $40,\ldots,49$.
 

datasets/rayleigh_benard_uniform/README.md

@@ -0,0 +1,112 @@
+# Rayleigh-Bénard convection
+
+**One line description of the data:** 2D horizontally-periodic Rayleigh-Benard convection on uniform grid.
+
+**Longer description of the data:**
+Rayleigh-Bénard convection involves fluid dynamics and thermodynamics, seen in a horizontal fluid layer heated from below, forming convective cells due to a temperature gradient. With the lower plate heated and the upper cooled, thermal energy creates density variations, initiating fluid motion. This results in Bénard cells, showcasing warm fluid rising and cool fluid descending. The interplay of buoyancy, conduction, and viscosity leads to complex fluid motion, including vortices and boundary layers.
+
+Note the `rayleigh_benard` and the `rayleigh_benard_uniform` are distinct only by resampling the data to different grids.
+
+**Associated paper**: [Paper 1](https://www.tandfonline.com/doi/pdf/10.1080/14786441608635602), [Paper 2](https://www.cambridge.org/core/journals/journal-of-fluid-mechanics/article/steady-rayleighbenard-convection-between-noslip-boundaries/B4F358EB0AE83BBE9D85968DC5DDD64D
+).
+
+**Data generated by**: [Rudy Morel](https://rudymorel.github.io/), Center for Computational Mathematics, Flatiron Institute.
+
+**Code or software used to generate the data**: [Github repository](https://github.com/RudyMorel/the-well-rbc-sf), based on the software [Dedalus]( https://dedalus-project.readthedocs.io/en/latest/pages/examples/ivp_2d_rayleigh_benard.html).
+
+**Equation**:
+
+While we solve equations in the frequency domain, the original time-domain problem is
+
+$$
+\begin{align*}
+\frac{\partial b}{\partial t} - \kappa \Delta b & = -u\nabla b,
+\\
+\frac{\partial u}{\partial t} - \nu \Delta u + \nabla p - b \vec{e}_y & = -u \nabla u,
+\end{align*}
+$$
+
+where $\Delta = \nabla \cdot \nabla$ is the spatial Laplacian, $b$ is the buoyancy, $u = (u_x,u_y)$ the (horizontal and vertical) velocity, and $p$ is the pressure, $\vec{e}_y$ is the unit vector in the vertical direction, with the additional constraints $\int p = 0$ (pressure gauge).
+
+The boundary conditions vertically are as follows:
+
+$$
+\begin{align*}
+b(y=0) = Ly ~~~,~~~ b(y=Ly) = 0
+\\
+u(y=0) = u(y=Ly) = 0
+\end{align*}
+$$
+
+These PDE are parameterized by the Rayleigh and Prandtl numbers through $\kappa$ and $\nu$.
+
+$$
+\begin{align*}
+\text{(thermal diffusivity)} ~~~~~~~ \kappa & = \big(\text{Rayleigh} * \text{Prandtl}\big)^{-\frac12}
+\\
+\text{(viscosity)} ~~~~~~~ \nu & = \bigg(\frac{\text{Rayleigh}}{\text{Prandtl}}\bigg)^{-\frac12}.
+\end{align*}
+$$
+
+![Gif](https://users.flatironinstitute.org/~polymathic/data/the_well/datasets/rayleigh_benard/gif/buoyancy_good_normalized.gif)
+<!-- <div style="transform: rotate(90deg);">
+  <img src="https://users.flatironinstitute.org/~polymathic/data/the_well/datasets/rayleigh_benard/gif/buoyancy_normalized.gif" alt="Rotated GIF">
+</div> -->
+
+
+| Dataset    | FNO | TFNO  | Unet | CNextU-net
+|:-:|:-:|:-:|:-:|:-:|
+| `rayleigh_benard` | 0.8395  | $\mathbf{0.6566}$ |1.4860|0.6699|
+
+Table: VRMSE metrics on test sets (lower is better). Best results are shown in bold. VRMSE is scaled such that predicting the mean value of the target field results in a score of 1.
+
+## About the data
+
+**Dimension of discretized data:** 200 timesteps of
+512 $\times$ 128 images.
+
+**Fields are available in the data:** buoyancy (scalar field), pressure (scalar field), velocity (vector field).
+
+**Number of simulations:** 1750 (35 PDE parameters $\times$ 50 initial conditions).
+
+**Size of the ensemble of all simulations:** 342 GB.
+
+**Grid type:** uniform, cartesian coordinates.
+
+**Initial conditions:** the buoyancy is composed of a dumped noise added to a linear background  $b(t=0) = (Ly-y)\times\delta b_0 + y(Ly-y) \times\epsilon$ where $\epsilon$ is a Gaussian white noise of scale $10^{-3}$.
+The other fields $u$ and $p$ are initialized to $0$.
+
+**Boundary conditions:** periodic on the horizontal direction, Dirichlet conditions on the vertical direction.
+
+**Simulation time-step:** 0.25.
+
+**Total time range ($t_{min}$ to $t_{max}$):** $t_{min} = 0$, $t_{max} = 50$.
+
+**Spatial domain size:** $0 \leq x \leq 4$ horizontally, and $0 \leq y \leq 1$ vertically.
+
+**Set of coefficients or non-dimensional parameters evaluated:** $\text{Rayleigh}\in[1e6,1e7,1e8,1e9,1e10], \text{Prandtl}\in[0.1,0.2,0.5,1.0,2.0,5.0,10.0]$. For initial conditions $\delta b_0\in[0.2,0.4,0.6,0.8,1.0]$, the seed used to generate the initial Gaussian white noise are $40,\ldots,49$.
+
+**Approximate time to generate the data: per input parameter** from $\sim6\,000s$ to $\sim 50\,000s$ (high Rayleigh numbers take longer), total: $\sim 60$ hours.
+
+**Hardware used to generate the data and precision used for generating the data:** 12 nodes of 64 CPU cores with 8 processes per node, in single precision.
+
+## What is interesting and challenging about the data:
+
+Rayleigh-Bénard convection datasets offer valuable insights into fluid dynamics under thermal gradients, revealing phenomena like turbulent eddies and convection cells.
+The position of such cells are highly sensitive to small variations in the initial conditions.
+Understanding these dynamics is crucial for applications in engineering and environmental science.
+
+Please cite the associated paper if you use this data in your research:
+
+```
+@article{burns2020dedalus,
+  title={Dedalus: A flexible framework for numerical simulations with spectral methods},
+  author={Burns, Keaton J and Vasil, Geoffrey M and Oishi, Jeffrey S and Lecoanet, Daniel and Brown, Benjamin P},
+  journal={Physical Review Research},
+  volume={2},
+  number={2},
+  pages={023068},
+  year={2020},
+  publisher={APS}
+}
+```

datasets/rayleigh_benard_uniform/rayleigh_benard_uniform.yaml

@@ -0,0 +1,98 @@
+dataset_name: rayleigh_benard
+n_spatial_dims: 2
+spatial_resolution:
+- 512
+- 128
+scalar_names: []
+constant_scalar_names:
+- Rayleigh
+- Prandtl
+field_names:
+  0:
+  - buoyancy
+  - pressure
+  1:
+  - velocity_x
+  - velocity_y
+  2: []
+constant_field_names:
+  0: []
+  1: []
+  2: []
+boundary_condition_types:
+- PERIODIC
+- WALL
+n_files: 35
+n_trajectories_per_file:
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+- 5
+n_steps_per_trajectory:
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+- 200
+grid_type: cartesian

datasets/rayleigh_benard_uniform/stats.yaml

@@ -0,0 +1,12 @@
+mean:
+  buoyancy: 0.3554901777056377
+  pressure: 0.00015696872462685817
+  velocity:
+  - -4.087598508971704e-06
+  - 1.503515349262273e-13
+std:
+  buoyancy: 0.19656606235363866
+  pressure: 0.12835234405568197
+  velocity:
+  - 0.20802661326557004
+  - 0.21215223278514359

mkdocs.yml

@@ -36,6 +36,7 @@ nav:
     - planetswe: datasets/planetswe.md
     - post_neutron_star_merger: datasets/post_neutron_star_merger.md
     - rayleigh_benard: datasets/rayleigh_benard.md
+    - rayleigh_benard_uniform: datasets/rayleigh_benard_uniform.md
     - rayleigh_taylor_instability: datasets/rayleigh_taylor_instability.md
     - shear_flow: datasets/shear_flow.md
     - supernova_explosion_64: datasets/supernova_explosion_64.md