models/dataset.py

"""E3 trajectory generation and HDF5 loading for MP-PDE.

This is an independent implementation from the equations and numerical-method
description in arXiv:2202.03376.  No official repository source is used.
"""

from __future__ import annotations

import argparse
from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait
import json
import multiprocessing
import os
from pathlib import Path
from typing import Any, Dict, Iterable, Mapping, Optional, Tuple

import h5py
import numpy as np
import torch
import yaml
from torch.utils.data import Dataset


SPLITS = ("train", "valid", "test")
PARAMETER_ORDER = ("alpha", "beta", "gamma")
PARALLEL_SCHEMA_VERSION = "mp_pde_e3_parallel_resume_v1"
_WORKER_X: Optional[np.ndarray] = None
_WORKER_TIMES: Optional[np.ndarray] = None
_WORKER_CONFIG: Optional[Dict[str, Any]] = None
_WORKER_STRIDE: Optional[int] = None


def _shift(values: np.ndarray, offset: int) -> np.ndarray:
    """Return values[i + offset] under periodic indexing."""
    return np.roll(values, -offset, axis=-1)


def _weno5_left(values: np.ndarray, epsilon: float) -> np.ndarray:
    """Fifth-order WENO left state at every i+1/2 interface."""
    um2, um1, u0 = _shift(values, -2), _shift(values, -1), values
    up1, up2 = _shift(values, 1), _shift(values, 2)
    p0 = (2.0 * um2 - 7.0 * um1 + 11.0 * u0) / 6.0
    p1 = (-um1 + 5.0 * u0 + 2.0 * up1) / 6.0
    p2 = (2.0 * u0 + 5.0 * up1 - up2) / 6.0
    b0 = (13.0 / 12.0) * (um2 - 2.0 * um1 + u0) ** 2 + 0.25 * (um2 - 4.0 * um1 + 3.0 * u0) ** 2
    b1 = (13.0 / 12.0) * (um1 - 2.0 * u0 + up1) ** 2 + 0.25 * (um1 - up1) ** 2
    b2 = (13.0 / 12.0) * (u0 - 2.0 * up1 + up2) ** 2 + 0.25 * (3.0 * u0 - 4.0 * up1 + up2) ** 2
    alpha = np.stack((0.1 / (epsilon + b0) ** 2, 0.6 / (epsilon + b1) ** 2, 0.3 / (epsilon + b2) ** 2))
    weights = alpha / np.sum(alpha, axis=0, keepdims=True)
    return weights[0] * p0 + weights[1] * p1 + weights[2] * p2


def _weno5_right(values: np.ndarray, epsilon: float) -> np.ndarray:
    """Fifth-order WENO right state at every i+1/2 interface."""
    um1, u0 = _shift(values, -1), values
    up1, up2, up3 = _shift(values, 1), _shift(values, 2), _shift(values, 3)
    p0 = (2.0 * up3 - 7.0 * up2 + 11.0 * up1) / 6.0
    p1 = (-up2 + 5.0 * up1 + 2.0 * u0) / 6.0
    p2 = (2.0 * up1 + 5.0 * u0 - um1) / 6.0
    b0 = (13.0 / 12.0) * (up1 - 2.0 * up2 + up3) ** 2 + 0.25 * (3.0 * up1 - 4.0 * up2 + up3) ** 2
    b1 = (13.0 / 12.0) * (u0 - 2.0 * up1 + up2) ** 2 + 0.25 * (u0 - up2) ** 2
    b2 = (13.0 / 12.0) * (um1 - 2.0 * u0 + up1) ** 2 + 0.25 * (um1 - 4.0 * u0 + 3.0 * up1) ** 2
    alpha = np.stack((0.1 / (epsilon + b0) ** 2, 0.6 / (epsilon + b1) ** 2, 0.3 / (epsilon + b2) ** 2))
    weights = alpha / np.sum(alpha, axis=0, keepdims=True)
    return weights[0] * p0 + weights[1] * p1 + weights[2] * p2


def godunov_quadratic_flux(left: np.ndarray, right: np.ndarray) -> np.ndarray:
    """Godunov flux for the convex scalar flux f(u)=u**2."""
    left_flux, right_flux = left * left, right * right
    rarefaction = left <= right
    rare_flux = np.where((left <= 0.0) & (right >= 0.0), 0.0, np.minimum(left_flux, right_flux))
    shock_flux = np.maximum(left_flux, right_flux)
    return np.where(rarefaction, rare_flux, shock_flux)


def weno5_flux_derivative(values: np.ndarray, dx: float, epsilon: float = 1.0e-6) -> np.ndarray:
    """Conservative derivative d_x(u**2) on a periodic uniform grid."""
    interface_flux = godunov_quadratic_flux(_weno5_left(values, epsilon), _weno5_right(values, epsilon))
    return (interface_flux - np.roll(interface_flux, 1, axis=-1)) / dx


def fourth_order_second_derivative(values: np.ndarray, dx: float) -> np.ndarray:
    return (-_shift(values, 2) + 16.0 * _shift(values, 1) - 30.0 * values + 16.0 * _shift(values, -1) - _shift(values, -2)) / (12.0 * dx**2)


def fourth_order_third_derivative(values: np.ndarray, dx: float) -> np.ndarray:
    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)


def sample_e3_parameters(rng: np.random.Generator, cfg: Mapping[str, Any]) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
    equation, forcing = cfg["equation"], cfg["forcing"]
    params = np.array(
        [rng.uniform(*equation["alpha_range"]), rng.uniform(*equation["beta_range"]), rng.uniform(*equation["gamma_range"])],
        dtype=np.float64,
    )
    policy = cfg["ambiguity_policy"]
    if policy == "paper_strict":
        omega_range = forcing["paper_omega_range"]
    elif policy == "official_consistency":
        omega_range = forcing["official_consistency_omega_range"]
    else:
        raise ValueError(f"Unknown ambiguity_policy={policy!r}")
    terms = int(forcing["terms"])
    provenance = {
        "amplitude": rng.uniform(*forcing["amplitude_range"], size=terms),
        "omega": rng.uniform(*omega_range, size=terms),
        "mode": rng.choice(np.asarray(forcing["modes"], dtype=np.int64), size=terms),
        "phase": rng.uniform(*forcing["phase_range"], size=terms),
    }
    return params, provenance


def evaluate_forcing(time: float, x: np.ndarray, forcing: Mapping[str, np.ndarray], domain_length: float) -> np.ndarray:
    phase = (
        forcing["omega"][:, None] * time
        + 2.0 * np.pi * forcing["mode"][:, None] * x[None, :] / domain_length
        + forcing["phase"][:, None]
    )
    return np.sum(forcing["amplitude"][:, None] * np.sin(phase), axis=0)


def e3_rhs(
    time: float,
    state: np.ndarray,
    x: np.ndarray,
    params: np.ndarray,
    forcing: Mapping[str, np.ndarray],
    domain_length: float,
    weno_epsilon: float,
) -> np.ndarray:
    alpha, beta, gamma = params
    dx = domain_length / state.shape[-1]
    return (
        evaluate_forcing(time, x, forcing, domain_length)
        - alpha * weno5_flux_derivative(state, dx, weno_epsilon)
        + beta * fourth_order_second_derivative(state, dx)
        - gamma * fourth_order_third_derivative(state, dx)
    )


def _stable_step(state: np.ndarray, params: np.ndarray, dx: float, cfl: float) -> float:
    alpha, beta, gamma = np.abs(params)
    limits = []
    wave_speed = 2.0 * alpha * float(np.max(np.abs(state)))
    if wave_speed > 1.0e-14:
        limits.append(dx / wave_speed)
    if beta > 1.0e-14:
        limits.append(dx**2 / (2.0 * beta))
    if gamma > 1.0e-14:
        limits.append(dx**3 / (6.0 * gamma))
    return cfl * min(limits) if limits else np.inf


def generate_trajectory(
    x: np.ndarray,
    save_times: np.ndarray,
    params: np.ndarray,
    forcing: Mapping[str, np.ndarray],
    cfg: Mapping[str, Any],
) -> np.ndarray:
    """Integrate one trajectory using RK4 and stability-limited substeps."""
    domain_length = float(cfg["domain_length"])
    generation = cfg["generation"]
    dx = domain_length / x.size
    state = evaluate_forcing(float(save_times[0]), x, forcing, domain_length).astype(np.float64)
    trajectory = np.empty((save_times.size, x.size), dtype=np.float32)
    trajectory[0] = state
    current_time = float(save_times[0])
    for output_index, target_time in enumerate(save_times[1:], start=1):
        substeps = 0
        while current_time < float(target_time) - 1.0e-14:
            stable = _stable_step(state, params, dx, float(generation["cfl"]))
            remaining = float(target_time) - current_time
            step = min(stable, remaining)
            if step < float(generation["min_dt"]) and remaining > float(generation["min_dt"]):
                raise RuntimeError(f"Stable RK4 step {step:.3e} fell below min_dt at t={current_time:.6g}")
            step = remaining if remaining <= float(generation["min_dt"]) else step
            rhs_args = (x, params, forcing, domain_length, float(generation["weno_epsilon"]))
            k1 = e3_rhs(current_time, state, *rhs_args)
            k2 = e3_rhs(current_time + 0.5 * step, state + 0.5 * step * k1, *rhs_args)
            k3 = e3_rhs(current_time + 0.5 * step, state + 0.5 * step * k2, *rhs_args)
            k4 = e3_rhs(current_time + step, state + step * k3, *rhs_args)
            state = state + (step / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
            current_time += step
            substeps += 1
            if substeps > int(generation["max_substeps"]):
                raise RuntimeError(f"max_substeps exceeded while advancing to t={target_time:.6g}")
            if not np.all(np.isfinite(state)):
                raise FloatingPointError(f"Non-finite E3 state at t={current_time:.6g}")
        current_time = float(target_time)
        trajectory[output_index] = state
    return trajectory


def _generation_config(config: Mapping[str, Any]) -> Dict[str, Any]:
    data = config["data"]
    return {
        "ambiguity_policy": config["experiment"]["ambiguity_policy"],
        "domain_length": data["domain_length"],
        "equation": data["equation"],
        "forcing": data["forcing"],
        "generation": data["generation"],
    }


def _initialize_generation_worker(
    x_high: np.ndarray, times: np.ndarray, generation_cfg: Mapping[str, Any], stride: int
) -> None:
    """Initialize immutable state used by one trajectory worker process."""
    global _WORKER_X, _WORKER_TIMES, _WORKER_CONFIG, _WORKER_STRIDE
    _WORKER_X = x_high
    _WORKER_TIMES = times
    _WORKER_CONFIG = dict(generation_cfg)
    _WORKER_STRIDE = int(stride)


def _generate_sample_worker(payload: Tuple[int, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]) -> Tuple[int, np.ndarray]:
    """Generate one independent trajectory; HDF5 remains owned by the parent."""
    if _WORKER_X is None or _WORKER_TIMES is None or _WORKER_CONFIG is None or _WORKER_STRIDE is None:
        raise RuntimeError("Parallel E3 worker was not initialized")
    sample_index, params, amplitude, omega, mode, phase = payload
    forcing = {"amplitude": amplitude, "omega": omega, "mode": mode, "phase": phase}
    trajectory = generate_trajectory(_WORKER_X, _WORKER_TIMES, params, forcing, _WORKER_CONFIG)
    return sample_index, trajectory[:, ::_WORKER_STRIDE]


def _sample_payloads(group: h5py.Group, indices: Iterable[int]) -> Iterable[Tuple[int, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]]:
    for sample_index in indices:
        yield (
            sample_index,
            np.asarray(group["params"][sample_index], dtype=np.float64),
            np.asarray(group["forcing_amplitude"][sample_index], dtype=np.float64),
            np.asarray(group["forcing_omega"][sample_index], dtype=np.float64),
            np.asarray(group["forcing_mode"][sample_index], dtype=np.int64),
            np.asarray(group["forcing_phase"][sample_index], dtype=np.float64),
        )


def _parallel_results(
    payloads: Iterable[Tuple[int, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]],
    workers: int,
    max_in_flight: int,
    x_high: np.ndarray,
    times: np.ndarray,
    generation_cfg: Mapping[str, Any],
    stride: int,
) -> Iterable[Tuple[int, np.ndarray]]:
    """Yield completed trajectories while keeping the process queue bounded."""
    if workers == 1:
        _initialize_generation_worker(x_high, times, generation_cfg, stride)
        for payload in payloads:
            yield _generate_sample_worker(payload)
        return

    thread_variables = ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_NUM_THREADS")
    previous_environment = {name: os.environ.get(name) for name in thread_variables}
    for name in thread_variables:
        os.environ[name] = "1"
    executor = ProcessPoolExecutor(
        max_workers=workers,
        mp_context=multiprocessing.get_context("spawn"),
        initializer=_initialize_generation_worker,
        initargs=(x_high, times, dict(generation_cfg), stride),
    )
    iterator = iter(payloads)
    futures = set()
    try:
        for _ in range(max_in_flight):
            try:
                futures.add(executor.submit(_generate_sample_worker, next(iterator)))
            except StopIteration:
                break
        while futures:
            completed, futures = wait(futures, return_when=FIRST_COMPLETED)
            for future in completed:
                yield future.result()
                try:
                    futures.add(executor.submit(_generate_sample_worker, next(iterator)))
                except StopIteration:
                    pass
    finally:
        executor.shutdown(wait=True, cancel_futures=True)
        for name, value in previous_environment.items():
            if value is None:
                os.environ.pop(name, None)
            else:
                os.environ[name] = value


def _generation_signature(
    config: Mapping[str, Any], counts: Mapping[str, int], nt: int, high_nx: int, target_nx: int
) -> str:
    payload = {
        "schema_version": PARALLEL_SCHEMA_VERSION,
        "paper": str(config["experiment"]["paper"]),
        "ambiguity_policy": str(config["experiment"]["ambiguity_policy"]),
        "generation_config": _generation_config(config),
        "counts": {name: int(counts[name]) for name in SPLITS},
        "nt": int(nt),
        "high_resolution_nx": int(high_nx),
        "resolution": int(target_nx),
        "seed": int(config["data"]["seed"]),
    }
    return json.dumps(payload, sort_keys=True, separators=(",", ":"))


def _initialize_partial_file(
    partial_path: Path,
    config: Mapping[str, Any],
    counts: Mapping[str, int],
    nt: int,
    high_nx: int,
    target_nx: int,
    x: np.ndarray,
    times: np.ndarray,
    generation_cfg: Mapping[str, Any],
    compression: Optional[str],
    compression_opts: Optional[int],
    signature: str,
) -> None:
    data = config["data"]
    base_seed = int(data["seed"])
    with h5py.File(partial_path, "w") as handle:
        handle.attrs["schema_version"] = PARALLEL_SCHEMA_VERSION
        handle.attrs["generation_signature"] = signature
        handle.attrs["experiment"] = "MP-PDE E3"
        handle.attrs["paper"] = str(config["experiment"]["paper"])
        handle.attrs["ambiguity_policy"] = str(config["experiment"]["ambiguity_policy"])
        handle.attrs["parameter_order"] = json.dumps(PARAMETER_ORDER)
        handle.attrs["generation_config"] = json.dumps(generation_cfg, sort_keys=True)
        handle.attrs["seed"] = base_seed
        handle.attrs["high_resolution_nx"] = high_nx
        handle.attrs["resolution"] = target_nx
        for split_index, split in enumerate(SPLITS):
            count = int(counts[split])
            rng = np.random.default_rng(np.random.SeedSequence([base_seed, split_index]))
            group = handle.create_group(split)
            group.attrs["seed_derivation"] = json.dumps([base_seed, split_index])
            group.create_dataset("x", data=x)
            group.create_dataset("t", data=times.astype(np.float32))
            group.create_dataset(
                "u", shape=(count, nt, target_nx), dtype="f4", chunks=(1, min(nt, 32), target_nx),
                compression=compression, compression_opts=compression_opts,
            )
            params_ds = group.create_dataset("params", shape=(count, 3), dtype="f4")
            terms = int(data["forcing"]["terms"])
            amplitude_ds = group.create_dataset("forcing_amplitude", shape=(count, terms), dtype="f4")
            omega_ds = group.create_dataset("forcing_omega", shape=(count, terms), dtype="f4")
            mode_ds = group.create_dataset("forcing_mode", shape=(count, terms), dtype="i8")
            phase_ds = group.create_dataset("forcing_phase", shape=(count, terms), dtype="f4")
            group.create_dataset("completed", shape=(count,), dtype="bool", data=np.zeros(count, dtype=bool))
            for sample_index in range(count):
                params, forcing = sample_e3_parameters(rng, generation_cfg)
                params_ds[sample_index] = params
                amplitude_ds[sample_index] = forcing["amplitude"]
                omega_ds[sample_index] = forcing["omega"]
                mode_ds[sample_index] = forcing["mode"]
                phase_ds[sample_index] = forcing["phase"]
        handle.flush()


def generate_e3_hdf5(
    config: Mapping[str, Any],
    output_path: Path | str,
    *,
    sample_counts: Optional[Mapping[str, int]] = None,
    nt: Optional[int] = None,
    high_resolution_nx: Optional[int] = None,
    resolution: Optional[int] = None,
    workers: Optional[int] = None,
    max_in_flight: Optional[int] = None,
    flush_every: Optional[int] = None,
    resume_partial: Optional[bool] = None,
    overwrite: bool = False,
) -> Path:
    """Generate E3 splits with process workers and a single resumable HDF5 writer."""
    output_path = Path(output_path)
    if output_path.exists() and not overwrite:
        raise FileExistsError(f"Refusing to overwrite existing dataset: {output_path}")
    output_path.parent.mkdir(parents=True, exist_ok=True)
    partial_path = output_path.with_suffix(output_path.suffix + ".partial")

    data = config["data"]
    parallel = data.get("parallel_generation", {})
    workers = int(workers if workers is not None else parallel.get("workers", 1))
    max_in_flight = int(max_in_flight if max_in_flight is not None else parallel.get("max_in_flight", 2 * workers))
    flush_every = int(flush_every if flush_every is not None else parallel.get("flush_every", workers))
    resume_partial = bool(resume_partial if resume_partial is not None else parallel.get("resume_partial", True))
    if workers < 1 or max_in_flight < workers or flush_every < 1:
        raise ValueError("workers>=1, max_in_flight>=workers, and flush_every>=1 are required")
    nt = int(nt or data["num_time_points"])
    high_nx = int(high_resolution_nx or data["high_resolution_nx"])
    target_nx = int(resolution or data["resolution"])
    if high_nx % target_nx != 0:
        raise ValueError(f"high_resolution_nx={high_nx} must be divisible by resolution={target_nx}")
    if high_nx < 7 or nt < 2:
        raise ValueError("Generation requires high_resolution_nx>=7 and nt>=2")
    counts = dict(sample_counts or {name: int(data[f"{name}_samples"]) for name in SPLITS})
    if set(counts) != set(SPLITS) or any(int(counts[name]) <= 0 for name in SPLITS):
        raise ValueError(f"sample_counts must provide positive counts for {SPLITS}")

    domain_length, final_time = float(data["domain_length"]), float(data["final_time"])
    x_high = np.linspace(0.0, domain_length, high_nx, endpoint=False, dtype=np.float64)
    stride = high_nx // target_nx
    x = x_high[::stride].astype(np.float32)
    times = np.linspace(0.0, final_time, nt, dtype=np.float64)
    generation_cfg = _generation_config(config)
    compression = data["generation"].get("compression")
    compression_opts = int(data["generation"].get("compression_level", 4)) if compression == "gzip" else None
    signature = _generation_signature(config, counts, nt, high_nx, target_nx)
    if partial_path.exists() and overwrite:
        partial_path.unlink()
    if partial_path.exists() and not resume_partial:
        raise FileExistsError(f"Partial dataset exists; enable resume_partial or use --overwrite: {partial_path}")
    if not partial_path.exists():
        _initialize_partial_file(
            partial_path, config, counts, nt, high_nx, target_nx, x, times, generation_cfg,
            compression, compression_opts, signature,
        )

    try:
        with h5py.File(partial_path, "r+") as handle:
            stored_signature = str(handle.attrs.get("generation_signature", ""))
            if stored_signature != signature:
                raise ValueError("Partial dataset configuration does not match this run; archive it or use --overwrite")
            print(
                f"[generate] workers={workers} max_in_flight={max_in_flight} flush_every={flush_every} "
                f"resume_partial={resume_partial}", flush=True,
            )
            for split in SPLITS:
                count = int(counts[split])
                group = handle[split]
                completed_ds = group["completed"]
                pending_indices = np.flatnonzero(~np.asarray(completed_ds[:], dtype=bool)).tolist()
                completed_count = count - len(pending_indices)
                if pending_indices:
                    print(f"[generate] split={split} resume_completed={completed_count}/{count}", flush=True)
                payloads = _sample_payloads(group, pending_indices)
                for sample_index, trajectory in _parallel_results(
                    payloads, workers, max_in_flight, x_high, times, generation_cfg, stride
                ):
                    group["u"][sample_index] = trajectory
                    completed_ds[sample_index] = True
                    completed_count += 1
                    if completed_count % flush_every == 0 or completed_count == count:
                        handle.flush()
                    print(
                        f"[generate] split={split} completed={completed_count}/{count} "
                        f"sample_index={sample_index} workers={workers}", flush=True,
                    )
                if not np.all(np.asarray(completed_ds[:], dtype=bool)):
                    raise RuntimeError(f"Split {split} is incomplete after generation")
                handle.flush()
        os.replace(partial_path, output_path)
    except Exception:
        print(f"Generation failed; partial file retained at {partial_path}", flush=True)
        raise
    return output_path


class E3Dataset(Dataset):
    """Lazy, process-safe reader for one E3 HDF5 split."""

    def __init__(self, path: Path | str, split: str, expected_nt: Optional[int] = None, expected_nx: Optional[int] = None):
        self.path = Path(path)
        self.split = split
        self._handle: Optional[h5py.File] = None
        if split not in SPLITS:
            raise ValueError(f"split must be one of {SPLITS}, got {split!r}")
        if not self.path.is_file():
            raise FileNotFoundError(f"E3 dataset not found: {self.path}")
        with h5py.File(self.path, "r") as handle:
            if split not in handle:
                raise KeyError(f"Missing HDF5 group {split!r}")
            group = handle[split]
            required = {"u", "x", "t", "params", "forcing_amplitude", "forcing_omega", "forcing_mode", "forcing_phase"}
            missing = required.difference(group.keys())
            if missing:
                raise KeyError(f"Missing HDF5 fields in {split}: {sorted(missing)}")
            shape = group["u"].shape
            if len(shape) != 3 or group["params"].shape != (shape[0], 3) or group["x"].shape != (shape[2],) or group["t"].shape != (shape[1],):
                raise ValueError(f"Inconsistent E3 schema in split={split}: u={shape}")
            if expected_nt is not None and shape[1] != expected_nt:
                raise ValueError(f"Expected nt={expected_nt}, found {shape[1]}")
            if expected_nx is not None and shape[2] != expected_nx:
                raise ValueError(f"Expected nx={expected_nx}, found {shape[2]}")
            self.length, self.nt, self.nx = shape
            x, t = group["x"][:], group["t"][:]
            if not np.all(np.isfinite(x)) or not np.all(np.isfinite(t)) or np.any(np.diff(t) <= 0.0):
                raise ValueError("Grid/time metadata are non-finite or non-monotone")
            if x.size > 1 and not np.allclose(np.diff(x), np.diff(x)[0], rtol=1e-5, atol=1e-7):
                raise ValueError("E3 x grid must be uniform")
            if t.size > 1 and not np.allclose(np.diff(t), np.diff(t)[0], rtol=1e-5, atol=1e-7):
                raise ValueError("E3 saved times must be uniform")

    def _group(self) -> h5py.Group:
        if self._handle is None:
            self._handle = h5py.File(self.path, "r")
        return self._handle[self.split]

    def __len__(self) -> int:
        return self.length

    def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:
        group = self._group()
        trajectory = np.asarray(group["u"][index], dtype=np.float32)
        params = np.asarray(group["params"][index], dtype=np.float32)
        if not np.all(np.isfinite(trajectory)) or not np.all(np.isfinite(params)):
            raise FloatingPointError(f"Non-finite data at split={self.split}, sample={index}")
        return {
            "u": torch.from_numpy(trajectory),
            "x": torch.from_numpy(np.asarray(group["x"][:], dtype=np.float32)),
            "t": torch.from_numpy(np.asarray(group["t"][:], dtype=np.float32)),
            "params": torch.from_numpy(params),
            "index": torch.tensor(index, dtype=torch.long),
        }

    def close(self) -> None:
        if self._handle is not None:
            self._handle.close()
            self._handle = None

    def __del__(self) -> None:
        self.close()


def _load_yaml(path: Path) -> Dict[str, Any]:
    with path.open("r", encoding="utf-8") as stream:
        config = yaml.safe_load(stream)
    if not isinstance(config, dict):
        raise ValueError(f"Config must contain a mapping: {path}")
    return config


def main() -> None:
    project_root = Path(__file__).resolve().parents[1]
    parser = argparse.ArgumentParser(description="Generate the MP-PDE E3 HDF5 dataset")
    parser.add_argument("--config", type=Path, default=project_root / "config/config.yaml")
    parser.add_argument("--output", type=Path, default=None)
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--train-samples", type=int)
    parser.add_argument("--valid-samples", type=int)
    parser.add_argument("--test-samples", type=int)
    parser.add_argument("--nt", type=int)
    parser.add_argument("--high-resolution-nx", type=int)
    parser.add_argument("--resolution", type=int)
    parser.add_argument("--workers", type=int)
    parser.add_argument("--max-in-flight", type=int)
    parser.add_argument("--flush-every", type=int)
    parser.add_argument("--no-resume-partial", action="store_true")
    args = parser.parse_args()
    config = _load_yaml(args.config.resolve())
    configured = Path(config["paths"]["data"])
    output = args.output or (configured if configured.is_absolute() else project_root / configured)
    default_counts = {name: int(config["data"][f"{name}_samples"]) for name in SPLITS}
    counts = {
        "train": args.train_samples or default_counts["train"],
        "valid": args.valid_samples or default_counts["valid"],
        "test": args.test_samples or default_counts["test"],
    }
    generated = generate_e3_hdf5(
        config, output, sample_counts=counts, nt=args.nt, high_resolution_nx=args.high_resolution_nx,
        resolution=args.resolution, workers=args.workers, max_in_flight=args.max_in_flight,
        flush_every=args.flush_every, resume_partial=False if args.no_resume_partial else None,
        overwrite=args.overwrite,
    )
    print(f"Generated E3 dataset: {generated}", flush=True)


if __name__ == "__main__":
    main()