OneScience-Group/MP_PDE

1"""E3 trajectory generation and HDF5 loading for MP-PDE.2 3This is an independent implementation from the equations and numerical-method4description in arXiv:2202.03376.  No official repository source is used.5"""6 7from __future__ import annotations8 9import argparse10from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait11import json12import multiprocessing13import os14from pathlib import Path15from typing import Any, Dict, Iterable, Mapping, Optional, Tuple16 17import h5py18import numpy as np19import torch20import yaml21from torch.utils.data import Dataset22 23 24SPLITS = ("train", "valid", "test")25PARAMETER_ORDER = ("alpha", "beta", "gamma")26PARALLEL_SCHEMA_VERSION = "mp_pde_e3_parallel_resume_v1"27_WORKER_X: Optional[np.ndarray] = None28_WORKER_TIMES: Optional[np.ndarray] = None29_WORKER_CONFIG: Optional[Dict[str, Any]] = None30_WORKER_STRIDE: Optional[int] = None31 32 33def _shift(values: np.ndarray, offset: int) -> np.ndarray:34    """Return values[i + offset] under periodic indexing."""35    return np.roll(values, -offset, axis=-1)36 37 38def _weno5_left(values: np.ndarray, epsilon: float) -> np.ndarray:39    """Fifth-order WENO left state at every i+1/2 interface."""40    um2, um1, u0 = _shift(values, -2), _shift(values, -1), values41    up1, up2 = _shift(values, 1), _shift(values, 2)42    p0 = (2.0 * um2 - 7.0 * um1 + 11.0 * u0) / 6.043    p1 = (-um1 + 5.0 * u0 + 2.0 * up1) / 6.044    p2 = (2.0 * u0 + 5.0 * up1 - up2) / 6.045    b0 = (13.0 / 12.0) * (um2 - 2.0 * um1 + u0) ** 2 + 0.25 * (um2 - 4.0 * um1 + 3.0 * u0) ** 246    b1 = (13.0 / 12.0) * (um1 - 2.0 * u0 + up1) ** 2 + 0.25 * (um1 - up1) ** 247    b2 = (13.0 / 12.0) * (u0 - 2.0 * up1 + up2) ** 2 + 0.25 * (3.0 * u0 - 4.0 * up1 + up2) ** 248    alpha = np.stack((0.1 / (epsilon + b0) ** 2, 0.6 / (epsilon + b1) ** 2, 0.3 / (epsilon + b2) ** 2))49    weights = alpha / np.sum(alpha, axis=0, keepdims=True)50    return weights[0] * p0 + weights[1] * p1 + weights[2] * p251 52 53def _weno5_right(values: np.ndarray, epsilon: float) -> np.ndarray:54    """Fifth-order WENO right state at every i+1/2 interface."""55    um1, u0 = _shift(values, -1), values56    up1, up2, up3 = _shift(values, 1), _shift(values, 2), _shift(values, 3)57    p0 = (2.0 * up3 - 7.0 * up2 + 11.0 * up1) / 6.058    p1 = (-up2 + 5.0 * up1 + 2.0 * u0) / 6.059    p2 = (2.0 * up1 + 5.0 * u0 - um1) / 6.060    b0 = (13.0 / 12.0) * (up1 - 2.0 * up2 + up3) ** 2 + 0.25 * (3.0 * up1 - 4.0 * up2 + up3) ** 261    b1 = (13.0 / 12.0) * (u0 - 2.0 * up1 + up2) ** 2 + 0.25 * (u0 - up2) ** 262    b2 = (13.0 / 12.0) * (um1 - 2.0 * u0 + up1) ** 2 + 0.25 * (um1 - 4.0 * u0 + 3.0 * up1) ** 263    alpha = np.stack((0.1 / (epsilon + b0) ** 2, 0.6 / (epsilon + b1) ** 2, 0.3 / (epsilon + b2) ** 2))64    weights = alpha / np.sum(alpha, axis=0, keepdims=True)65    return weights[0] * p0 + weights[1] * p1 + weights[2] * p266 67 68def godunov_quadratic_flux(left: np.ndarray, right: np.ndarray) -> np.ndarray:69    """Godunov flux for the convex scalar flux f(u)=u**2."""70    left_flux, right_flux = left * left, right * right71    rarefaction = left <= right72    rare_flux = np.where((left <= 0.0) & (right >= 0.0), 0.0, np.minimum(left_flux, right_flux))73    shock_flux = np.maximum(left_flux, right_flux)74    return np.where(rarefaction, rare_flux, shock_flux)75 76 77def weno5_flux_derivative(values: np.ndarray, dx: float, epsilon: float = 1.0e-6) -> np.ndarray:78    """Conservative derivative d_x(u**2) on a periodic uniform grid."""79    interface_flux = godunov_quadratic_flux(_weno5_left(values, epsilon), _weno5_right(values, epsilon))80    return (interface_flux - np.roll(interface_flux, 1, axis=-1)) / dx81 82 83def fourth_order_second_derivative(values: np.ndarray, dx: float) -> np.ndarray:84    return (-_shift(values, 2) + 16.0 * _shift(values, 1) - 30.0 * values + 16.0 * _shift(values, -1) - _shift(values, -2)) / (12.0 * dx**2)85 86 87def fourth_order_third_derivative(values: np.ndarray, dx: float) -> np.ndarray:88    return (_shift(values, -3) - 8.0 * _shift(values, -2) + 13.0 * _shift(values, -1) - 13.0 * _shift(values, 1) + 8.0 * _shift(values, 2) - _shift(values, 3)) / (8.0 * dx**3)89 90 91def sample_e3_parameters(rng: np.random.Generator, cfg: Mapping[str, Any]) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:92    equation, forcing = cfg["equation"], cfg["forcing"]93    params = np.array(94        [rng.uniform(*equation["alpha_range"]), rng.uniform(*equation["beta_range"]), rng.uniform(*equation["gamma_range"])],95        dtype=np.float64,96    )97    policy = cfg["ambiguity_policy"]98    if policy == "paper_strict":99        omega_range = forcing["paper_omega_range"]100    elif policy == "official_consistency":101        omega_range = forcing["official_consistency_omega_range"]102    else:103        raise ValueError(f"Unknown ambiguity_policy={policy!r}")104    terms = int(forcing["terms"])105    provenance = {106        "amplitude": rng.uniform(*forcing["amplitude_range"], size=terms),107        "omega": rng.uniform(*omega_range, size=terms),108        "mode": rng.choice(np.asarray(forcing["modes"], dtype=np.int64), size=terms),109        "phase": rng.uniform(*forcing["phase_range"], size=terms),110    }111    return params, provenance112 113 114def evaluate_forcing(time: float, x: np.ndarray, forcing: Mapping[str, np.ndarray], domain_length: float) -> np.ndarray:115    phase = (116        forcing["omega"][:, None] * time117        + 2.0 * np.pi * forcing["mode"][:, None] * x[None, :] / domain_length118        + forcing["phase"][:, None]119    )120    return np.sum(forcing["amplitude"][:, None] * np.sin(phase), axis=0)121 122 123def e3_rhs(124    time: float,125    state: np.ndarray,126    x: np.ndarray,127    params: np.ndarray,128    forcing: Mapping[str, np.ndarray],129    domain_length: float,130    weno_epsilon: float,131) -> np.ndarray:132    alpha, beta, gamma = params133    dx = domain_length / state.shape[-1]134    return (135        evaluate_forcing(time, x, forcing, domain_length)136        - alpha * weno5_flux_derivative(state, dx, weno_epsilon)137        + beta * fourth_order_second_derivative(state, dx)138        - gamma * fourth_order_third_derivative(state, dx)139    )140 141 142def _stable_step(state: np.ndarray, params: np.ndarray, dx: float, cfl: float) -> float:143    alpha, beta, gamma = np.abs(params)144    limits = []145    wave_speed = 2.0 * alpha * float(np.max(np.abs(state)))146    if wave_speed > 1.0e-14:147        limits.append(dx / wave_speed)148    if beta > 1.0e-14:149        limits.append(dx**2 / (2.0 * beta))150    if gamma > 1.0e-14:151        limits.append(dx**3 / (6.0 * gamma))152    return cfl * min(limits) if limits else np.inf153 154 155def generate_trajectory(156    x: np.ndarray,157    save_times: np.ndarray,158    params: np.ndarray,159    forcing: Mapping[str, np.ndarray],160    cfg: Mapping[str, Any],161) -> np.ndarray:162    """Integrate one trajectory using RK4 and stability-limited substeps."""163    domain_length = float(cfg["domain_length"])164    generation = cfg["generation"]165    dx = domain_length / x.size166    state = evaluate_forcing(float(save_times[0]), x, forcing, domain_length).astype(np.float64)167    trajectory = np.empty((save_times.size, x.size), dtype=np.float32)168    trajectory[0] = state169    current_time = float(save_times[0])170    for output_index, target_time in enumerate(save_times[1:], start=1):171        substeps = 0172        while current_time < float(target_time) - 1.0e-14:173            stable = _stable_step(state, params, dx, float(generation["cfl"]))174            remaining = float(target_time) - current_time175            step = min(stable, remaining)176            if step < float(generation["min_dt"]) and remaining > float(generation["min_dt"]):177                raise RuntimeError(f"Stable RK4 step {step:.3e} fell below min_dt at t={current_time:.6g}")178            step = remaining if remaining <= float(generation["min_dt"]) else step179            rhs_args = (x, params, forcing, domain_length, float(generation["weno_epsilon"]))180            k1 = e3_rhs(current_time, state, *rhs_args)181            k2 = e3_rhs(current_time + 0.5 * step, state + 0.5 * step * k1, *rhs_args)182            k3 = e3_rhs(current_time + 0.5 * step, state + 0.5 * step * k2, *rhs_args)183            k4 = e3_rhs(current_time + step, state + step * k3, *rhs_args)184            state = state + (step / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)185            current_time += step186            substeps += 1187            if substeps > int(generation["max_substeps"]):188                raise RuntimeError(f"max_substeps exceeded while advancing to t={target_time:.6g}")189            if not np.all(np.isfinite(state)):190                raise FloatingPointError(f"Non-finite E3 state at t={current_time:.6g}")191        current_time = float(target_time)192        trajectory[output_index] = state193    return trajectory194 195 196def _generation_config(config: Mapping[str, Any]) -> Dict[str, Any]:197    data = config["data"]198    return {199        "ambiguity_policy": config["experiment"]["ambiguity_policy"],200        "domain_length": data["domain_length"],201        "equation": data["equation"],202        "forcing": data["forcing"],203        "generation": data["generation"],204    }205 206 207def _initialize_generation_worker(208    x_high: np.ndarray, times: np.ndarray, generation_cfg: Mapping[str, Any], stride: int209) -> None:210    """Initialize immutable state used by one trajectory worker process."""211    global _WORKER_X, _WORKER_TIMES, _WORKER_CONFIG, _WORKER_STRIDE212    _WORKER_X = x_high213    _WORKER_TIMES = times214    _WORKER_CONFIG = dict(generation_cfg)215    _WORKER_STRIDE = int(stride)216 217 218def _generate_sample_worker(payload: Tuple[int, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]) -> Tuple[int, np.ndarray]:219    """Generate one independent trajectory; HDF5 remains owned by the parent."""220    if _WORKER_X is None or _WORKER_TIMES is None or _WORKER_CONFIG is None or _WORKER_STRIDE is None:221        raise RuntimeError("Parallel E3 worker was not initialized")222    sample_index, params, amplitude, omega, mode, phase = payload223    forcing = {"amplitude": amplitude, "omega": omega, "mode": mode, "phase": phase}224    trajectory = generate_trajectory(_WORKER_X, _WORKER_TIMES, params, forcing, _WORKER_CONFIG)225    return sample_index, trajectory[:, ::_WORKER_STRIDE]226 227 228def _sample_payloads(group: h5py.Group, indices: Iterable[int]) -> Iterable[Tuple[int, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]]:229    for sample_index in indices:230        yield (231            sample_index,232            np.asarray(group["params"][sample_index], dtype=np.float64),233            np.asarray(group["forcing_amplitude"][sample_index], dtype=np.float64),234            np.asarray(group["forcing_omega"][sample_index], dtype=np.float64),235            np.asarray(group["forcing_mode"][sample_index], dtype=np.int64),236            np.asarray(group["forcing_phase"][sample_index], dtype=np.float64),237        )238 239 240def _parallel_results(241    payloads: Iterable[Tuple[int, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]],242    workers: int,243    max_in_flight: int,244    x_high: np.ndarray,245    times: np.ndarray,246    generation_cfg: Mapping[str, Any],247    stride: int,248) -> Iterable[Tuple[int, np.ndarray]]:249    """Yield completed trajectories while keeping the process queue bounded."""250    if workers == 1:251        _initialize_generation_worker(x_high, times, generation_cfg, stride)252        for payload in payloads:253            yield _generate_sample_worker(payload)254        return255 256    thread_variables = ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_NUM_THREADS")257    previous_environment = {name: os.environ.get(name) for name in thread_variables}258    for name in thread_variables:259        os.environ[name] = "1"260    executor = ProcessPoolExecutor(261        max_workers=workers,262        mp_context=multiprocessing.get_context("spawn"),263        initializer=_initialize_generation_worker,264        initargs=(x_high, times, dict(generation_cfg), stride),265    )266    iterator = iter(payloads)267    futures = set()268    try:269        for _ in range(max_in_flight):270            try:271                futures.add(executor.submit(_generate_sample_worker, next(iterator)))272            except StopIteration:273                break274        while futures:275            completed, futures = wait(futures, return_when=FIRST_COMPLETED)276            for future in completed:277                yield future.result()278                try:279                    futures.add(executor.submit(_generate_sample_worker, next(iterator)))280                except StopIteration:281                    pass282    finally:283        executor.shutdown(wait=True, cancel_futures=True)284        for name, value in previous_environment.items():285            if value is None:286                os.environ.pop(name, None)287            else:288                os.environ[name] = value289 290 291def _generation_signature(292    config: Mapping[str, Any], counts: Mapping[str, int], nt: int, high_nx: int, target_nx: int293) -> str:294    payload = {295        "schema_version": PARALLEL_SCHEMA_VERSION,296        "paper": str(config["experiment"]["paper"]),297        "ambiguity_policy": str(config["experiment"]["ambiguity_policy"]),298        "generation_config": _generation_config(config),299        "counts": {name: int(counts[name]) for name in SPLITS},300        "nt": int(nt),301        "high_resolution_nx": int(high_nx),302        "resolution": int(target_nx),303        "seed": int(config["data"]["seed"]),304    }305    return json.dumps(payload, sort_keys=True, separators=(",", ":"))306 307 308def _initialize_partial_file(309    partial_path: Path,310    config: Mapping[str, Any],311    counts: Mapping[str, int],312    nt: int,313    high_nx: int,314    target_nx: int,315    x: np.ndarray,316    times: np.ndarray,317    generation_cfg: Mapping[str, Any],318    compression: Optional[str],319    compression_opts: Optional[int],320    signature: str,321) -> None:322    data = config["data"]323    base_seed = int(data["seed"])324    with h5py.File(partial_path, "w") as handle:325        handle.attrs["schema_version"] = PARALLEL_SCHEMA_VERSION326        handle.attrs["generation_signature"] = signature327        handle.attrs["experiment"] = "MP-PDE E3"328        handle.attrs["paper"] = str(config["experiment"]["paper"])329        handle.attrs["ambiguity_policy"] = str(config["experiment"]["ambiguity_policy"])330        handle.attrs["parameter_order"] = json.dumps(PARAMETER_ORDER)331        handle.attrs["generation_config"] = json.dumps(generation_cfg, sort_keys=True)332        handle.attrs["seed"] = base_seed333        handle.attrs["high_resolution_nx"] = high_nx334        handle.attrs["resolution"] = target_nx335        for split_index, split in enumerate(SPLITS):336            count = int(counts[split])337            rng = np.random.default_rng(np.random.SeedSequence([base_seed, split_index]))338            group = handle.create_group(split)339            group.attrs["seed_derivation"] = json.dumps([base_seed, split_index])340            group.create_dataset("x", data=x)341            group.create_dataset("t", data=times.astype(np.float32))342            group.create_dataset(343                "u", shape=(count, nt, target_nx), dtype="f4", chunks=(1, min(nt, 32), target_nx),344                compression=compression, compression_opts=compression_opts,345            )346            params_ds = group.create_dataset("params", shape=(count, 3), dtype="f4")347            terms = int(data["forcing"]["terms"])348            amplitude_ds = group.create_dataset("forcing_amplitude", shape=(count, terms), dtype="f4")349            omega_ds = group.create_dataset("forcing_omega", shape=(count, terms), dtype="f4")350            mode_ds = group.create_dataset("forcing_mode", shape=(count, terms), dtype="i8")351            phase_ds = group.create_dataset("forcing_phase", shape=(count, terms), dtype="f4")352            group.create_dataset("completed", shape=(count,), dtype="bool", data=np.zeros(count, dtype=bool))353            for sample_index in range(count):354                params, forcing = sample_e3_parameters(rng, generation_cfg)355                params_ds[sample_index] = params356                amplitude_ds[sample_index] = forcing["amplitude"]357                omega_ds[sample_index] = forcing["omega"]358                mode_ds[sample_index] = forcing["mode"]359                phase_ds[sample_index] = forcing["phase"]360        handle.flush()361 362 363def generate_e3_hdf5(364    config: Mapping[str, Any],365    output_path: Path | str,366    *,367    sample_counts: Optional[Mapping[str, int]] = None,368    nt: Optional[int] = None,369    high_resolution_nx: Optional[int] = None,370    resolution: Optional[int] = None,371    workers: Optional[int] = None,372    max_in_flight: Optional[int] = None,373    flush_every: Optional[int] = None,374    resume_partial: Optional[bool] = None,375    overwrite: bool = False,376) -> Path:377    """Generate E3 splits with process workers and a single resumable HDF5 writer."""378    output_path = Path(output_path)379    if output_path.exists() and not overwrite:380        raise FileExistsError(f"Refusing to overwrite existing dataset: {output_path}")381    output_path.parent.mkdir(parents=True, exist_ok=True)382    partial_path = output_path.with_suffix(output_path.suffix + ".partial")383 384    data = config["data"]385    parallel = data.get("parallel_generation", {})386    workers = int(workers if workers is not None else parallel.get("workers", 1))387    max_in_flight = int(max_in_flight if max_in_flight is not None else parallel.get("max_in_flight", 2 * workers))388    flush_every = int(flush_every if flush_every is not None else parallel.get("flush_every", workers))389    resume_partial = bool(resume_partial if resume_partial is not None else parallel.get("resume_partial", True))390    if workers < 1 or max_in_flight < workers or flush_every < 1:391        raise ValueError("workers>=1, max_in_flight>=workers, and flush_every>=1 are required")392    nt = int(nt or data["num_time_points"])393    high_nx = int(high_resolution_nx or data["high_resolution_nx"])394    target_nx = int(resolution or data["resolution"])395    if high_nx % target_nx != 0:396        raise ValueError(f"high_resolution_nx={high_nx} must be divisible by resolution={target_nx}")397    if high_nx < 7 or nt < 2:398        raise ValueError("Generation requires high_resolution_nx>=7 and nt>=2")399    counts = dict(sample_counts or {name: int(data[f"{name}_samples"]) for name in SPLITS})400    if set(counts) != set(SPLITS) or any(int(counts[name]) <= 0 for name in SPLITS):401        raise ValueError(f"sample_counts must provide positive counts for {SPLITS}")402 403    domain_length, final_time = float(data["domain_length"]), float(data["final_time"])404    x_high = np.linspace(0.0, domain_length, high_nx, endpoint=False, dtype=np.float64)405    stride = high_nx // target_nx406    x = x_high[::stride].astype(np.float32)407    times = np.linspace(0.0, final_time, nt, dtype=np.float64)408    generation_cfg = _generation_config(config)409    compression = data["generation"].get("compression")410    compression_opts = int(data["generation"].get("compression_level", 4)) if compression == "gzip" else None411    signature = _generation_signature(config, counts, nt, high_nx, target_nx)412    if partial_path.exists() and overwrite:413        partial_path.unlink()414    if partial_path.exists() and not resume_partial:415        raise FileExistsError(f"Partial dataset exists; enable resume_partial or use --overwrite: {partial_path}")416    if not partial_path.exists():417        _initialize_partial_file(418            partial_path, config, counts, nt, high_nx, target_nx, x, times, generation_cfg,419            compression, compression_opts, signature,420        )421 422    try:423        with h5py.File(partial_path, "r+") as handle:424            stored_signature = str(handle.attrs.get("generation_signature", ""))425            if stored_signature != signature:426                raise ValueError("Partial dataset configuration does not match this run; archive it or use --overwrite")427            print(428                f"[generate] workers={workers} max_in_flight={max_in_flight} flush_every={flush_every} "429                f"resume_partial={resume_partial}", flush=True,430            )431            for split in SPLITS:432                count = int(counts[split])433                group = handle[split]434                completed_ds = group["completed"]435                pending_indices = np.flatnonzero(~np.asarray(completed_ds[:], dtype=bool)).tolist()436                completed_count = count - len(pending_indices)437                if pending_indices:438                    print(f"[generate] split={split} resume_completed={completed_count}/{count}", flush=True)439                payloads = _sample_payloads(group, pending_indices)440                for sample_index, trajectory in _parallel_results(441                    payloads, workers, max_in_flight, x_high, times, generation_cfg, stride442                ):443                    group["u"][sample_index] = trajectory444                    completed_ds[sample_index] = True445                    completed_count += 1446                    if completed_count % flush_every == 0 or completed_count == count:447                        handle.flush()448                    print(449                        f"[generate] split={split} completed={completed_count}/{count} "450                        f"sample_index={sample_index} workers={workers}", flush=True,451                    )452                if not np.all(np.asarray(completed_ds[:], dtype=bool)):453                    raise RuntimeError(f"Split {split} is incomplete after generation")454                handle.flush()455        os.replace(partial_path, output_path)456    except Exception:457        print(f"Generation failed; partial file retained at {partial_path}", flush=True)458        raise459    return output_path460 461 462class E3Dataset(Dataset):463    """Lazy, process-safe reader for one E3 HDF5 split."""464 465    def __init__(self, path: Path | str, split: str, expected_nt: Optional[int] = None, expected_nx: Optional[int] = None):466        self.path = Path(path)467        self.split = split468        self._handle: Optional[h5py.File] = None469        if split not in SPLITS:470            raise ValueError(f"split must be one of {SPLITS}, got {split!r}")471        if not self.path.is_file():472            raise FileNotFoundError(f"E3 dataset not found: {self.path}")473        with h5py.File(self.path, "r") as handle:474            if split not in handle:475                raise KeyError(f"Missing HDF5 group {split!r}")476            group = handle[split]477            required = {"u", "x", "t", "params", "forcing_amplitude", "forcing_omega", "forcing_mode", "forcing_phase"}478            missing = required.difference(group.keys())479            if missing:480                raise KeyError(f"Missing HDF5 fields in {split}: {sorted(missing)}")481            shape = group["u"].shape482            if len(shape) != 3 or group["params"].shape != (shape[0], 3) or group["x"].shape != (shape[2],) or group["t"].shape != (shape[1],):483                raise ValueError(f"Inconsistent E3 schema in split={split}: u={shape}")484            if expected_nt is not None and shape[1] != expected_nt:485                raise ValueError(f"Expected nt={expected_nt}, found {shape[1]}")486            if expected_nx is not None and shape[2] != expected_nx:487                raise ValueError(f"Expected nx={expected_nx}, found {shape[2]}")488            self.length, self.nt, self.nx = shape489            x, t = group["x"][:], group["t"][:]490            if not np.all(np.isfinite(x)) or not np.all(np.isfinite(t)) or np.any(np.diff(t) <= 0.0):491                raise ValueError("Grid/time metadata are non-finite or non-monotone")492            if x.size > 1 and not np.allclose(np.diff(x), np.diff(x)[0], rtol=1e-5, atol=1e-7):493                raise ValueError("E3 x grid must be uniform")494            if t.size > 1 and not np.allclose(np.diff(t), np.diff(t)[0], rtol=1e-5, atol=1e-7):495                raise ValueError("E3 saved times must be uniform")496 497    def _group(self) -> h5py.Group:498        if self._handle is None:499            self._handle = h5py.File(self.path, "r")500        return self._handle[self.split]501 502    def __len__(self) -> int:503        return self.length504 505    def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:506        group = self._group()507        trajectory = np.asarray(group["u"][index], dtype=np.float32)508        params = np.asarray(group["params"][index], dtype=np.float32)509        if not np.all(np.isfinite(trajectory)) or not np.all(np.isfinite(params)):510            raise FloatingPointError(f"Non-finite data at split={self.split}, sample={index}")511        return {512            "u": torch.from_numpy(trajectory),513            "x": torch.from_numpy(np.asarray(group["x"][:], dtype=np.float32)),514            "t": torch.from_numpy(np.asarray(group["t"][:], dtype=np.float32)),515            "params": torch.from_numpy(params),516            "index": torch.tensor(index, dtype=torch.long),517        }518 519    def close(self) -> None:520        if self._handle is not None:521            self._handle.close()522            self._handle = None523 524    def __del__(self) -> None:525        self.close()526 527 528def _load_yaml(path: Path) -> Dict[str, Any]:529    with path.open("r", encoding="utf-8") as stream:530        config = yaml.safe_load(stream)531    if not isinstance(config, dict):532        raise ValueError(f"Config must contain a mapping: {path}")533    return config534 535 536def main() -> None:537    project_root = Path(__file__).resolve().parents[1]538    parser = argparse.ArgumentParser(description="Generate the MP-PDE E3 HDF5 dataset")539    parser.add_argument("--config", type=Path, default=project_root / "config/config.yaml")540    parser.add_argument("--output", type=Path, default=None)541    parser.add_argument("--overwrite", action="store_true")542    parser.add_argument("--train-samples", type=int)543    parser.add_argument("--valid-samples", type=int)544    parser.add_argument("--test-samples", type=int)545    parser.add_argument("--nt", type=int)546    parser.add_argument("--high-resolution-nx", type=int)547    parser.add_argument("--resolution", type=int)548    parser.add_argument("--workers", type=int)549    parser.add_argument("--max-in-flight", type=int)550    parser.add_argument("--flush-every", type=int)551    parser.add_argument("--no-resume-partial", action="store_true")552    args = parser.parse_args()553    config = _load_yaml(args.config.resolve())554    configured = Path(config["paths"]["data"])555    output = args.output or (configured if configured.is_absolute() else project_root / configured)556    default_counts = {name: int(config["data"][f"{name}_samples"]) for name in SPLITS}557    counts = {558        "train": args.train_samples or default_counts["train"],559        "valid": args.valid_samples or default_counts["valid"],560        "test": args.test_samples or default_counts["test"],561    }562    generated = generate_e3_hdf5(563        config, output, sample_counts=counts, nt=args.nt, high_resolution_nx=args.high_resolution_nx,564        resolution=args.resolution, workers=args.workers, max_in_flight=args.max_in_flight,565        flush_every=args.flush_every, resume_partial=False if args.no_resume_partial else None,566        overwrite=args.overwrite,567    )568    print(f"Generated E3 dataset: {generated}", flush=True)569 570 571if __name__ == "__main__":572    main()573