yangshun2005/openscience-cn/hdf5-pde-data-loading

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

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]

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]

from huggingface_hub import hf_hub_download
path = hf_hub_download(
    repo_id="pdebench/Advection",           # or pdebench/Burgers
    filename="1D_Advection_Sols_beta1.0.hdf5",
    repo_type="dataset",
    cache_dir="/tmp/hf_cache",
)

import subprocess
url = f"https://darus.uni-stuttgart.de/api/access/datafile/{file_id}"
subprocess.run([
    "aria2c", "-x", "16", "-s", "16",
    "--max-connection-per-server=16",
    "--min-split-size=10M", "--timeout=600", "--max-tries=5",
    "-d", "/tmp", "-o", filename, url
], check=True, timeout=3600)

import urllib.request
csv_url = "https://raw.githubusercontent.com/pdebench/PDEBench/main/pdebench/data_download/pdebench_data_urls.csv"
data = urllib.request.urlopen(csv_url).read().decode()
for line in data.split('\n'):
    if 'your_dataset' in line:
        print(line)  # filename, URL, path, md5

class PDEDS(Dataset):
    def __init__(self, data, grid, init_step=10):
        self.data = data      # [N, X, T, C]
        self.grid = grid      # [X, 1]
        self.init_step = init_step
    def __len__(self):
        return self.data.shape[0]
    def __getitem__(self, i):
        x = self.data[i, :, :self.init_step, :]  # Input window
        y = self.data[i]                           # Full trajectory (for AR loss)
        return x, y, self.grid
train_loader = DataLoader(
    PDEDS(data[:N_TRAIN], grid, INIT_STEP),
    batch_size=32, shuffle=True, num_workers=2, pin_memory=True
)