Fetch Illustris TNG halos

fetch-illustris-tng-halos.py

#!/usr/bin/env python3
"""Fetch a halo from the IllustrisTNG API (https://www.tng-project.org/api).

Fetches the metadata of a halo (FoF group) as JSON and, by default, downloads
its full HDF5 cutout: all particle types (gas, dm, stars, bhs) with all fields.
The cutout can optionally be restricted to a subset of particle types/fields.

Usage:
    python fetch_halo.py HALO_ID --api-key KEY [options]

Arguments:
    HALO_ID                 Halo (FoF group) ID in the group catalog
    --api-key KEY           TNG API key (get one at https://www.tng-project.org/users/register/)
    --simulation NAME       Simulation to query (default: TNG100-1), e.g. TNG50-1, TNG300-1
    --snapshot N            Snapshot number (default: 99, i.e. z=0 for TNG runs)
    --redshift Z            Redshift instead of snapshot number; resolved to the
                            simulation's nearest snapshot (mutually exclusive
                            with --snapshot)
    --output PATH           Where to save the metadata JSON (default: halo_<ID>.json)
    --skip-cutout           Only fetch the metadata, skip the HDF5 download
    --subset PARTTYPE=FIELDS
                            Restrict the cutout; repeatable. PARTTYPE is one of
                            gas, dm, stars, bhs; FIELDS is a comma-separated list
                            of snapshot field names. Omit for the entire cutout.
    --cutout-output PATH    Where to save the HDF5 cutout
                            (default: cutout_<sim>_snap<N>_halo<ID>.hdf5)

Examples:
    # Full cutout of halo 450916 in TNG50-1 at z=0
    python fetch_halo.py 450916 --api-key KEY --simulation TNG50-1

    # Same, but selecting the epoch by redshift instead of snapshot number
    python fetch_halo.py 450916 --api-key KEY --simulation TNG50-1 --redshift 2.0

    # Only gas coordinates/density and dm coordinates
    python fetch_halo.py 450916 --api-key KEY --simulation TNG50-1 \\
        --subset gas=Coordinates,Density --subset dm=Coordinates

    # Metadata only
    python fetch_halo.py 450916 --api-key KEY --skip-cutout

Notes:
    - Halo IDs are simulation- and snapshot-specific; the same ID refers to
      different objects in different runs.
    - The data server occasionally answers 503 (overloaded) or 504 (still
      building a large cutout); the script retries with increasing delays
      automatically. If it still fails after all attempts, rerun later.
"""

import argparse
import json
import os
import sys
import time

import requests

BASE_URL = "https://www.tng-project.org/api"
RETRY_DELAYS = [15, 30, 60, 120, 240]  # seconds between attempts on 5xx errors


def get(url, api_key, params=None):
    """Perform an authenticated GET request against the TNG API."""
    headers = {"api-key": api_key}
    response = requests.get(url, params=params, headers=headers, timeout=60)
    response.raise_for_status()
    if response.headers.get("content-type", "").startswith("application/json"):
        return response.json()
    return response.content


def download(url, api_key, output, params=None):
    """Stream a (potentially large) file from the TNG API to disk.

    Retries on 5xx errors: the data server answers 503 when temporarily
    overloaded and 504 while it is still building a large cutout, so waiting
    and retrying usually succeeds.
    """
    headers = {"api-key": api_key}
    for attempt, delay in enumerate([*RETRY_DELAYS, None]):
        try:
            with requests.get(url, params=params, headers=headers, stream=True, timeout=600) as response:
                response.raise_for_status()
                with open(output, "wb") as fh:
                    for chunk in response.iter_content(chunk_size=1024 * 1024):
                        fh.write(chunk)
            return output
        except requests.HTTPError as err:
            if delay is None or err.response is None or err.response.status_code < 500:
                raise
            print(f"Attempt {attempt + 1} failed ({err.response.status_code} "
                  f"{err.response.reason}); retrying in {delay}s...")
            time.sleep(delay)


def resolve_snapshot(simulation, redshift, api_key):
    """Find the simulation snapshot whose redshift is closest to the requested one."""
    try:
        snapshots = get(f"{BASE_URL}/{simulation}/snapshots/", api_key)
    except requests.HTTPError as err:
        sys.exit(f"Could not list snapshots for {simulation}: {err}")
    closest = min(snapshots, key=lambda snap: abs(snap["redshift"] - redshift))
    print(f"Redshift {redshift} -> snapshot {closest['number']} (z={closest['redshift']:.4f})")
    return closest["number"]


def parse_subset(specs):
    """Turn ['gas=Coordinates,Density', 'stars=Masses'] into cutout query params."""
    params = {}
    for spec in specs:
        part_type, _, fields = spec.partition("=")
        if not fields:
            sys.exit(f"Invalid --subset '{spec}': expected PARTTYPE=Field1,Field2,...")
        params[part_type] = fields
    return params


def main():
    parser = argparse.ArgumentParser(description="Fetch a halo from the TNG API.")
    parser.add_argument("halo_id", type=int, help="Halo (FoF group) ID")
    parser.add_argument("--api-key", required=True, help="TNG API key")
    parser.add_argument("--simulation", default="TNG100-1", help="Simulation name (default: TNG100-1)")
    when = parser.add_mutually_exclusive_group()
    when.add_argument("--snapshot", type=int, default=None, help="Snapshot number (default: 99)")
    when.add_argument(
        "--redshift", type=float, default=None,
        help="Redshift; resolved to the nearest snapshot of the simulation "
             "(alternative to --snapshot)",
    )
    parser.add_argument("--output", default=None, help="Path to save the JSON response")
    parser.add_argument(
        "--skip-cutout", action="store_true",
        help="Only fetch the halo metadata, not the HDF5 cutout",
    )
    parser.add_argument(
        "--subset", action="append", default=[], metavar="PARTTYPE=FIELDS",
        help="Restrict the cutout to given fields, e.g. --subset gas=Coordinates,Density "
             "--subset stars=Masses (repeatable; default: entire cutout, all types and fields)",
    )
    parser.add_argument("--cutout-output", default=None, help="Path to save the HDF5 cutout")
    args = parser.parse_args()

    if args.redshift is not None:
        snapshot = resolve_snapshot(args.simulation, args.redshift, args.api_key)
    else:
        snapshot = args.snapshot if args.snapshot is not None else 99
    args.snapshot = snapshot

    url = f"{BASE_URL}/{args.simulation}/snapshots/{args.snapshot}/halos/{args.halo_id}/"
    print(f"Fetching {url}")

    try:
        halo = get(url, args.api_key)
    except requests.HTTPError as err:
        sys.exit(f"Request failed: {err}")

    print(json.dumps(halo, indent=2))

    output = args.output or f"halo_{args.halo_id}.json"
    with open(output, "w") as fh:
        json.dump(halo, fh, indent=2)
    print(f"\nSaved response to {output}")

    if args.skip_cutout:
        return

    cutout_url = f"{url}cutout.hdf5"
    cutout_params = parse_subset(args.subset) or None
    cutout_output = args.cutout_output or f"cutout_{args.simulation}_snap{args.snapshot}_halo{args.halo_id}.hdf5"
    scope = "subset: " + ", ".join(args.subset) if args.subset else "entire cutout"
    print(f"\nDownloading HDF5 cutout ({scope}) from {cutout_url}")

    try:
        download(cutout_url, args.api_key, cutout_output, params=cutout_params)
    except requests.HTTPError as err:
        sys.exit(f"Cutout download failed: {err}")

    size_mb = os.path.getsize(cutout_output) / 1024**2
    print(f"Saved cutout to {cutout_output} ({size_mb:.1f} MB)")


if __name__ == "__main__":
    main()