Dataset Downloading API¤

The Downloading API was co-authored with Gemini

The neojax.data module provides robust dataset downloading utilities to fetch standardized Partial Differential Equation (PDE) datasets from various public hosting services (like Hugging Face Hub, Zenodo, DaRUS Stuttgart Dataverse, or direct HTTP/HTTPS URLs).

Downloader Architecture¤

Scientific PDE datasets are typically large (ranging from hundreds of megabytes to tens of gigabytes). To download these files reliably over standard internet connections, neojax implements a resilient download pipeline:

Connection Retries: Automatically retries downloads on network timeouts or transient errors using exponential backoff.
Range-Based Resume: Checks the size of existing partial downloads (marked with a .part extension) and sends HTTP Range requests to resume downloading from the offset where the connection was cut, preventing wasted bandwidth.
Progress Tracking: Uses tqdm for interactive progress bars when available, falling back to a clean terminal-printing progress bar otherwise.
Checksum Verification: Verifies finished files against expected MD5 or SHA256 hashes, deleting and re-downloading if a mismatch is found.

Registered Datasets¤

The registry contains 37 standardized PDE datasets (including 20 PDEgym datasets from Hugging Face, 2 verification datasets from Zenodo, and 15 complete PDEBench datasets from Stuttgart's DaRUS repository).

PDEgym

Poseidon: Efficient Foundation Models for PDEs

Huggingface collection

@inproceedings{herde2024_poseidon,
    author = {Herde, Maximilian and Raoni\'{c}, Bogdan and Rohner, Tobias and K\"{a}ppeli, Roger and Molinaro, Roberto and de B\'{e}zenac, Emmanuel and Mishra, Siddhartha},
    booktitle = {Advances in Neural Information Processing Systems},
    doi = {10.52202/079017-2311},
    editor = {A. Globerson and L. Mackey and D. Belgrave and A. Fan and U. Paquet and J. Tomczak and C. Zhang},
    pages = {72525--72624},
    publisher = {Curran Associates, Inc.},
    title = {Poseidon: Efficient Foundation Models for PDEs},
    url = {https://proceedings.neurips.cc/paper_files/paper/2024/file/84e1b1ec17bb11c57234e96433022a9a-Paper-Conference.pdf},
    volume = {37},
    year = {2024}
}

PDEBench

PDEBENCH: An Extensive Benchmark for Scientific Machine Learning

@article{takamoto2022pdebench,
    title={Pdebench: An extensive benchmark for scientific machine learning},
    author={Takamoto, Makoto and Praditia, Timothy and Leiteritz, Raphael and MacKinlay, Daniel and Alesiani, Francesco and Pfl{\"u}ger, Dirk and Niepert, Mathias},
    journal={Advances in neural information processing systems},
    volume={35},
    pages={1596--1611},
    year={2022}
}

Hugging Face (`camlab-ethz/pdegym`)¤

The PDEgym collection represents a diverse range of physical systems used for pretraining and benchmarking operator foundation models:

Key	Description
`ace`	Allen-Cahn equation
`ce_crp`	Compressible Euler equations (Centrally Rarefaction/Pressure problem)
`ce_gauss`	Compressible Euler equations with Gaussian initial conditions
`ce_kh`	Compressible Euler equations (Kelvin-Helmholtz instability)
`ce_rm`	Compressible Euler equations (Richtmyer-Meshkov problem)
`ce_rp`	Compressible Euler equations (Riemann problem)
`ce_rpui`	Compressible Euler equations (Riemann problem variant)
`fns_kf`	Forced incompressible Navier-Stokes equations with Kolmogorov forcing
`gce_rt`	Gravity-driven compressible Euler (Rayleigh-Taylor instability)
`helmholtz`	Solutions to the Helmholtz equation
`ns_bb`	Navier-Stokes equations (Backwards-facing step variant)
`ns_gauss`	Navier-Stokes equations with Gaussian initial conditions
`ns_pwc`	Navier-Stokes equations (Piecewise constant)
`ns_sines`	Navier-Stokes equations with sine-based initial conditions
`ns_sl`	Navier-Stokes equations (Shear Layer)
`ns_svs`	Navier-Stokes equations (Small-scale vortex structures)
`poisson_gauss`	Poisson equation with Gaussian sources
`se_af`	Compressible Euler equations (Steady Airfoil)
`wave_gauss`	Wave equation with Gaussian initial conditions and wave speeds
`wave_layer`	Wave equation with layered wave speeds

Zenodo (PDEBench Verification)¤

Quick reference datasets:

Key	Description
`burgers_1d`	1D Burgers' equation Salsa T10 verification dataset
`darcy_flow_2d`	2D Darcy Flow beta 1.0 verification dataset

DaRUS Stuttgart Dataverse (PDEBench)¤

The complete set of PDEBench problems resolved dynamically using Dataverse file ID queries:

Key	Description
`pdebench_advection_1d`	1D Advection dataset (8 files)
`pdebench_burgers_1d`	1D Burgers' equation dataset (12 files)
`pdebench_cfd_1d_random`	1D Compressible Flow with random initial conditions (4 files)
`pdebench_cfd_1d_shock`	1D Compressible Flow with shock tube boundaries (1 file)
`pdebench_diffusion_sorption_1d`	1D Diffusion-Sorption equation dataset (1 file)
`pdebench_cfd_2d_random`	2D Compressible Flow with random initial conditions (6 files)
`pdebench_cfd_2d_turbulent`	2D Compressible Flow with turbulent dynamics (2 files)
`pdebench_darcy_2d`	2D Darcy Flow dataset with variable betas (5 files)
`pdebench_diffusion_reaction_2d`	2D Diffusion-Reaction equation dataset (1 file)
`pdebench_reaction_diffusion_2d`	2D Reaction-Diffusion FitzHugh-Nagumo dataset (16 files)
`pdebench_brusselator_2d`	2D Reaction-Diffusion Brusselator dataset (1 file)
`pdebench_cfd_3d_random`	3D Compressible Flow with random initial conditions (2 files)
`pdebench_cfd_3d_turbulent`	3D Compressible Flow with turbulent dynamics (1 file)
`pdebench_sod_1d`	1D Sod Shock Tube dataset (1 file)
`pdebench_incompressible_2d`	2D Incompressible Inhomogeneous Navier-Stokes dataset (275 files)

API Reference¤

`neojax.data.download.dataset.download_dataset(name: str, target_dir: str | pathlib.Path, force: bool = False) -> list[str]` ¤

Downloads a dataset from the registry by name.

Parameters:

Name	Type	Description	Default
`name`	`str`	The name of the registered dataset (case-insensitive).	required
`target_dir`	`str \| Path`	Directory where the files should be saved.	required
`force`	`bool`	If True, forces re-downloading even if files exist and pass validation.	`False`

Returns:

Type	Description
`list[str]`	A list of absolute file paths to the downloaded files.

Raises:

Type	Description
`ValueError`	If the dataset name is not registered, or if the downloader is invalid.

`neojax.data.download.http_downloader.HTTPDownloader` ¤

Downloader that fetches files from HTTP/HTTPS URLs.

Resilience features include connection retries with exponential backoff, range-based resuming of interrupted downloads, and checksum verification.

Parameters:

Name	Type	Description	Default
`urls`	`dict[str, str]`	A dictionary mapping target filenames (relative to target_dir) to their download URLs.	required
`checksums`	`dict[str, str] \| None`	Optional dictionary mapping filenames to their expected checksum strings (e.g. 'md5:xxx' or 'sha256:xxx').	`None`
`chunk_size`	`int`	Chunk size in bytes for reading the stream. Default is 1MB.	`1048576`
`max_retries`	`int`	Maximum number of download retries on failure. Default is 5.	`5`
`backoff_factor`	`float`	Exponential backoff factor for retries. Default is 1.5.	`1.5`

Internal Attributes

These fields store the configuration of the downloader.

urls (dict[str, str]): Target filenames mapped to download URLs.
checksums (dict[str, str]): Target filenames mapped to expected checksums.
chunk_size (int): Chunk size in bytes.
max_retries (int): Max download retries on failure.
backoff_factor (float): Backoff factor for exponential backoff.

`download(target_dir: str | pathlib.Path, force: bool = False) -> list[pathlib.Path]` ¤

Download all configured URLs.

Parameters:

Name	Type	Description	Default
`target_dir`	`str \| Path`	Directory where the files should be saved.	required
`force`	`bool`	If True, forces re-downloading even if files exist and pass validation.	`False`

Returns:

Type	Description
`list[Path]`	A list of absolute file paths to the downloaded files.

Raises:

Type	Description
`RequestException`	If the HTTP request/connection fails.
`ValueError`	If a downloaded file fails checksum verification.
`OSError`	If a file system error occurs during reading or writing.
`RuntimeError`	If the download loop terminates unexpectedly.

`neojax.data.download.huggingface_downloader.HuggingFaceDownloader` ¤

Downloader that fetches files from Hugging Face Hub repositories.

Queries the Hugging Face Dataset API to list repository siblings, filters out repository configuration/metadata files, and downloads the dataset files using direct LFS resolve links.

Parameters:

Name	Type	Description	Default
`repo_id`	`str`	Hugging Face dataset repository ID (e.g. 'camlab-ethz/ACE').	required
`filenames`	`list[str] \| None`	Optional list of specific files to download. If None, downloads all non-metadata files.	`None`
`revision`	`str`	Git branch or commit hash to download from. Defaults to 'main'.	`'main'`
`max_retries`	`int`	Retries for the HTTP downloader. Default is 5.	`5`
`backoff_factor`	`float`	Backoff factor for retries. Default is 1.5.	`1.5`

Internal Attributes

These fields store the configuration of the downloader.

repo_id (str): Hugging Face repository ID.
filenames (list[str] | None): Optional list of specific files to download.
revision (str): Git branch/tag/commit hash.
max_retries (int): Max download retries.
backoff_factor (float): Backoff factor.

`download(target_dir: str | pathlib.Path, force: bool = False) -> list[str]` ¤

Query Hugging Face API to list files and download them.

Parameters:

Name	Type	Description	Default
`target_dir`	`str \| Path`	Directory where the files should be saved.	required
`force`	`bool`	If True, forces re-downloading even if files exist and pass validation.	`False`

Returns:

Type	Description
`list[str]`	A list of absolute file paths to the downloaded files.

Raises:

Type	Description
`HTTPError`	If the HTTP request to the Hugging Face API fails.
`ValueError`	If no files are found in the Hugging Face repository, or if no matching files are found to download.

`neojax.data.download.zenodo_downloader.ZenodoDownloader` ¤

Downloader that fetches files from Zenodo record IDs.

Queries the Zenodo Record API to discover download URLs and checksums, then delegates the download to HTTPDownloader.

Parameters:

Name	Type	Description	Default
`record_id`	`str`	The Zenodo Record ID.	required
`filenames`	`list[str] \| None`	Optional list of specific filenames to download. If None, downloads all files in the record.	`None`
`max_retries`	`int`	Retries for the HTTP downloader. Default is 5.	`5`
`backoff_factor`	`float`	Backoff factor for retries. Default is 1.5.	`1.5`

Internal Attributes

These fields store the configuration of the downloader.

record_id (str): The Zenodo Record ID.
filenames (list[str] | None): Optional list of specific filenames to download.
max_retries (int): Max download retries.
backoff_factor (float): Backoff factor.

`download(target_dir: str | pathlib.Path, force: bool = False) -> list[str]` ¤

Query Zenodo API to list and download files.

Parameters:

Name	Type	Description	Default
`target_dir`	`str \| Path`	Directory where the files should be saved.	required
`force`	`bool`	If True, forces re-downloading even if files exist and pass validation.	`False`

Returns:

Type	Description
`list[str]`	A list of absolute file paths to the downloaded files.

Raises:

Type	Description
`HTTPError`	If the HTTP request to the Zenodo API fails.
`ValueError`	If no files are found in the Zenodo record, or if none of the specified filenames are present in the record.

`neojax.data.download.dataverse_downloader.DataverseDownloader` ¤

Downloader that fetches files from Dataverse datasets using DOIs.

Queries the Dataverse API to resolve file IDs and checksums (MD5), then delegates the download to HTTPDownloader.

Parameters:

Name	Type	Description	Default
`doi`	`str`	The persistent ID (DOI) of the dataset (e.g. '10.18419/darus-2986').	required
`filenames`	`list[str] \| None`	Optional list of specific filenames to download. If None, downloads all files in the dataset.	`None`
`server_url`	`str`	Base URL of the Dataverse server instance. Defaults to "https://darus.uni-stuttgart.de".	`'https://darus.uni-stuttgart.de'`
`max_retries`	`int`	Retries for the HTTP downloader. Default is 5.	`5`
`backoff_factor`	`float`	Backoff factor for retries. Default is 1.5.	`1.5`

Internal Attributes

These fields store the configuration of the downloader.

doi (str): Dataset DOI.
filenames (list[str] | None): Optional list of specific files to download.
server_url (str): Base URL of the Dataverse instance.
max_retries (int): Max download retries.
backoff_factor (float): Backoff factor.

`download(target_dir: str | pathlib.Path, force: bool = False) -> list[str]` ¤

Query Dataverse API to list and download files.

Parameters:

Name	Type	Description	Default
`target_dir`	`str \| Path`	Directory where the files should be saved.	required
`force`	`bool`	If True, forces re-downloading even if files exist and pass validation.	`False`

Returns:

Type	Description
`list[str]`	A list of absolute file paths to the downloaded files.

Raises:

Type	Description
`HTTPError`	If the HTTP request to the Dataverse API fails.
`ValueError`	If no files are found in the dataset, or if none of the specified filenames are present in the dataset.

Dataset Downloading API¤

Downloader Architecture¤

Registered Datasets¤

Hugging Face (camlab-ethz/pdegym)¤

Zenodo (PDEBench Verification)¤

DaRUS Stuttgart Dataverse (PDEBench)¤

API Reference¤

neojax.data.download.dataset.download_dataset(name: str, target_dir: str | pathlib.Path, force: bool = False) -> list[str] ¤

neojax.data.download.http_downloader.HTTPDownloader ¤

download(target_dir: str | pathlib.Path, force: bool = False) -> list[pathlib.Path] ¤

neojax.data.download.huggingface_downloader.HuggingFaceDownloader ¤

download(target_dir: str | pathlib.Path, force: bool = False) -> list[str] ¤

neojax.data.download.zenodo_downloader.ZenodoDownloader ¤

download(target_dir: str | pathlib.Path, force: bool = False) -> list[str] ¤

neojax.data.download.dataverse_downloader.DataverseDownloader ¤

download(target_dir: str | pathlib.Path, force: bool = False) -> list[str] ¤

Hugging Face (`camlab-ethz/pdegym`)¤

`neojax.data.download.dataset.download_dataset(name: str, target_dir: str | pathlib.Path, force: bool = False) -> list[str]` ¤

`neojax.data.download.http_downloader.HTTPDownloader` ¤

`download(target_dir: str | pathlib.Path, force: bool = False) -> list[pathlib.Path]` ¤

`neojax.data.download.huggingface_downloader.HuggingFaceDownloader` ¤

`download(target_dir: str | pathlib.Path, force: bool = False) -> list[str]` ¤

`neojax.data.download.zenodo_downloader.ZenodoDownloader` ¤

`download(target_dir: str | pathlib.Path, force: bool = False) -> list[str]` ¤

`neojax.data.download.dataverse_downloader.DataverseDownloader` ¤

`download(target_dir: str | pathlib.Path, force: bool = False) -> list[str]` ¤