hdf5-pde-data-loading

HDF5 PDE Data Loading

When to Use

• Loading PDE simulation datasets stored in HDF5 format
• PDEBench, PhiFlow, JAX-CFD, or custom simulation outputs
• Multi-variable systems (density, velocity, pressure, etc.)
• Need to downsample spatial/temporal dimensions for memory

HDF5 Layout Detection

PDE datasets come in two common layouts:

Layout 1: Single "tensor" dataset

with h5py.File(path, "r") as f:
    if "tensor" in f:
        ds = f["tensor"]  # shape: [N, T, X, C] or [N, T, X]

Layout 2: Separate variable datasets

with h5py.File(path, "r") as f:
    # Keys like: "density", "Vx", "pressure", or "t", "x", "u"
    rho = f["density"][:]  # [N, T, X]
    vel = f["Vx"][:]
    prs = f["pressure"][:]
    data = np.stack([rho, vel, prs], axis=-1)  # [N, T, X, 3]

Robust loading (handles both):

def load_pde_hdf5(path, res_x=1, res_t=1):
    """Load PDE data from HDF5, handling multiple layouts."""
    with h5py.File(path, "r") as f:
        print(f"Keys: {list(f.keys())}")

        if "tensor" in f:
            ds = f["tensor"]
            raw_shape = ds.shape
            if len(raw_shape) == 4:
                N, T, X, C = raw_shape
            else:
                N, T, X = raw_shape; C = 1

            X_ds = X // res_x
            T_ds = T // res_t if res_t > 1 else T
            data = np.empty((N, X_ds, T_ds, C if len(raw_shape)==4 else 1), dtype=np.float32)

            for s in range(0, N, 500):  # Chunk to avoid OOM
                e = min(s + 500, N)
                chunk = ds[s:e, ::res_t, ::res_x]
                if len(chunk.shape) == 3:
                    data[s:e, :, :, 0] = np.transpose(chunk, (0, 2, 1))
                else:
                    data[s:e] = np.transpose(chunk, (0, 2, 1, 3))
        else:
            # Separate variables — find and stack them
            var_keys = []
            for k in sorted(f.keys()):
                if isinstance(f[k], h5py.Dataset) and len(f[k].shape) >= 3:
                    var_keys.append(k)

            arrays = [f[k][:, ::res_t, ::res_x] for k in var_keys]
            raw = np.stack(arrays, axis=-1)  # [N, T, X, C]
            data = np.transpose(raw, (0, 2, 1, 3)).astype(np.float32)

        # Load grid coordinates
        for grid_key in ["x-coordinate", "x", "X"]:
            if grid_key in f:
                grid = np.array(f[grid_key], dtype=np.float32)[::res_x]
                break
        else:
            grid = np.linspace(0, 1, data.shape[1], dtype=np.float32)

    return data, grid  # data: [N, X, T, C], grid: [X]