Data Module¶

PyTorch Dataset / IterableDataset classes over FITS images and tables, plus make_loader (a DataLoader factory with torchfits cache defaults).

For a single file, use Core I/O or Tables. Use Datasets when you need workers, shuffling (map-style via make_loader), or many samples.

Typical stack: read_tensor / table.read → optional transforms → Fits*Dataset → make_loader.

Choosing a Dataset¶

Channels: flux bands / arms / CCDs stack as tensor channels. IVAR and mask are companion tensors (or dict fields) — pass ivar_hdu= / mask_hdu= (or table ivar_column=).

transform=: each dataset calls transform(payload) with one argument (the tensor, or a {"flux", "ivar"?, "mask"?} dict when companions are set). There is no separate mask= kwarg at the Dataset boundary. Custom transforms should branch on dict vs tensor if they need mask-aware wiring — see custom transforms.

Cache: how and when¶

A single-file read_tensor does not need warm-up.

make_loader(ds, optimize_cache=True) (default) calls cache.optimize_for_dataset(ds.files, avg_file_size_mb=...) when the dataset exposes a files list (FitsImageDataset, FitsCubeDataset, FitsSpectrumDataset, FitsTensorIterableDataset, FitsCutoutDataset, and other FitsTensorDataset peers). For remote URLs it also starts prefetch into the disk cache, except FitsCutoutDataset (Range cutouts — no full-file prefetch). Table datasets read one file and do not expose files; optimize_cache=True is then a no-op.

from torchfits import cache

cache.get_cache_stats()
cache.clear_cache()

Warm-up helps when a loader iterates many files (remote prefetch, multi-worker reads, or repeated mosaic cutouts via open_subset_reader).

FitsTensorDataset¶

Map-style general N-D IMAGE reader (any rank). Multi-HDU hdu=[…] stacks flux on dim 0; optional ivar_hdu / mask_hdu return companion tensors.

from torchfits.data import FitsTensorDataset, make_loader

ds = FitsTensorDataset(
    "observations/*.fits",
    hdu=0,
    label_key="CLASS",  # header keyword → int label
    transform=None,  # optional callable
    device="cpu",
    mmap=True,
    add_channel_dim=False,  # Tensor default: leave rank alone
    cache_dir=None,  # optional remote materialization dir
)
loader = make_loader(ds, batch_size=32, num_workers=4)

Returns per item: (payload, label) where payload is a Tensor or {"flux", "ivar"?, "mask"?}.

FitsImageDataset / FitsCubeDataset¶

Subclasses of FitsTensorDataset — same constructor signature and returns, with two defaults changed:

• Image — 2D; add_channel_dim=True by default; multi-band HDUs → [C,H,W].
• Cube — 3D+; add_channel_dim=False by default; adds slice_index.

from torchfits.data import FitsImageDataset, FitsCubeDataset

images = FitsImageDataset("obs/*.fits", hdu=0, label_key="CLASS")
cubes = FitsCubeDataset("cubes/*.fits", hdu=0, slice_index=None)

Returns per item: same as FitsTensorDataset — (payload, label).

FitsImageIterableDataset stays first-class for loaders (same knobs as the tensor iterable, channel dim default on).

FitsSpectrumDataset¶

1D spectra (IMAGE NAXIS=1 or table column= + optional ivar_column=), plus DESI-style 2D [nspec, nwave] via row=. Multi-arm (MOS B/R/Z) uses layout=:

from torchfits.data import FitsSpectrumDataset

ds = FitsSpectrumDataset(
    "spectra/*.fits",
    hdu=["B", "R", "Z"],
    layout="dict",
)

Returns per item: payload per layout (see table above) — no label.

Synthetic DESI-shaped demo: examples/desi_shaped_spectrum.py.

FitsTensorIterableDataset¶

Iterable dataset for multi-worker sharded tensor loading. Each worker processes a deterministic subset — every file is seen exactly once per epoch.

from torchfits.data import FitsTensorIterableDataset

ds = FitsTensorIterableDataset(
    "observations/*.fits",
    hdu=0,
    shuffle=True,
    seed=42,
    add_channel_dim=False,
)

Returns per item: payload (Tensor or flux/ivar/mask dict) — no label.

FitsTableDataset¶

Map-style dataset for row-indexable FITS catalogs. Loads the full filtered table at __init__ — use only when the catalog fits in RAM.

from torchfits.data import FitsTableDataset

ds = FitsTableDataset(
    "catalog.fits",
    hdu=1,
    columns=["RA", "DEC", "MAG_G"],
    where="MAG_G < 20",  # predicate pushdown at load time
    labels=[0, 1, 0, 1, 0, 1],  # optional per-row labels (default 0)
    transform=None,
)

Returns per item: (dict[str, Tensor], torch.Tensor) — row dict and torch.long label.

FitsTableIterableDataset¶

Streams table rows in constant memory. Each __getitem__ yields one dict[str, Tensor] row.

from torchfits.data import FitsTableIterableDataset

ds = FitsTableIterableDataset(
    "survey.fits",
    hdu=1,
    columns=["RA", "DEC"],
    where="DEC > 0",
    batch_size=65536,
    mmap="auto",
)

Returns per item: dict[str, Tensor | Any] — one row.

FitsCutoutDataset¶

Map-style dataset for fixed cutout windows from one or more FITS images.

from torchfits.data import FitsCutoutDataset

cutouts = [
    ("mosaic.fits", 0, 100, 200, 164),  # (path, hdu, x, y, size)
    ("mosaic.fits", 0, 300, 400, 164),
]
ds = FitsCutoutDataset(cutouts, transform=None, device="cpu")

Accepts (path, hdu, x, y, size) or (path, hdu, x1, y1, x2, y2) tuples.

Returns per item: Tensor — raw cutout tensor without a label. This dataset is unsupervised by design (it has no labels parameter). For supervised patch classification, subclass FitsCutoutDataset and override __getitem__ to return (tensor, label), or create a custom Dataset that pairs read_subset calls with labels.

make_loader()¶

Builds a torch.utils.data.DataLoader and, when optimize_cache=True, calls cache.optimize_for_dataset if the dataset exposes files.

from torchfits.data import make_loader

loader = make_loader(
    ds,
    batch_size=32,
    num_workers=4,
    pin_memory=True,
    optimize_cache=True,
    shuffle=True,  # default for map-style; False for iterable
)

Use a plain DataLoader when you already own collate_fn / samplers, or when the dataset has no files list (optimize_cache would no-op). See examples/example_make_loader_vs_dataloader.py.

Returns: torch.utils.data.DataLoader

fits_collate_fn¶

Default collation function for torchfits datasets:

• list[Tensor] → stacked Tensor (all must have the same shape)
• list[(Tensor, Tensor)] → (stacked_images, stacked_labels)
• list[dict[str, Tensor]] → dict[str, stacked_Tensor]
• list[(dict[str, Tensor], Tensor)] → (stacked_dict, stacked_labels)

All tensors in a batch must have identical shapes — torch.stack is used under the hood. Raises ValueError on non-tensor columns (strings, VLA lists). For variable-size images, either pad externally before collation or supply a custom collate_fn.

FitsCubeIterableDataset¶

Streaming peer for 3D+ datacubes (e.g. IFUs, velocity cubes, radio cubes). Optionally slices along the leading spectral/velocity axis on the fly via slice_index=.

from torchfits.data import FitsCubeIterableDataset

ds = FitsCubeIterableDataset(
    "cubes/*.fits",
    hdu=0,
    slice_index=15,  # Extracts channel 15 on the fly
    shuffle=True,
    shuffle_buffer_size=100,  # In-flight reservoir shuffle
)

FitsSpectrumIterableDataset¶

Streaming 1D multi-arm spectra from binary table columns (e.g. DESI, BOSS) or multi-HDU image extensions.

from torchfits.data import FitsSpectrumIterableDataset

ds = FitsSpectrumIterableDataset(
    "spectra/*.fits",
    hdu=["B_FLUX", "R_FLUX", "Z_FLUX"],
    ivar_hdu=["B_IVAR", "R_IVAR", "Z_IVAR"],
    layout="stack",  # 'dict', 'stack' [C, nwave], or 'concat'
)

FitsStagedCutoutIterableDataset¶

Asynchronous staged cutout extraction for massive astronomical survey mosaics (CFHT MegaCam, HSC, Rubin LSST, MegaPipe). Downloads remote mosaics to fast ephemeral scratch (e.g. $SLURM_TMPDIR), samples \(K\) random stamps rapidly using open_subset_reader, and automatically deletes the staged file upon completion.

from torchfits.data import FitsStagedCutoutIterableDataset

ds = FitsStagedCutoutIterableDataset(
    paths=["https://archive.org/mosaic1.fits", "https://archive.org/mosaic2.fits"],
    cutouts_per_file=500,  # 500 cutouts sampled per staged mosaic
    cutout_size=(128, 128),
    staging_dir=None,  # Auto-detects $SLURM_TMPDIR, $TMPDIR, or scratch
    cleanup=True,  # Delete staged mosaic when finished sampling
    shuffle_buffer_size=1000,  # Mix stamps across mosaics in-flight
)

Distributed & Worker Sharding¶

torchfits iterable datasets natively partition across both multi-process DataLoader workers and multi-GPU / multi-node distributed training (rank and world_size).

Automatic Coordinate Resolution¶

When rank and world_size are omitted, torchfits automatically checks: 1. Cluster environment variables: RANK & WORLD_SIZE (torchrun, DeepSpeed, FSDP) or SLURM_PROCID & SLURM_NTASKS (SLURM). 2. torch.distributed process group (if initialized). 3. Single-worker fallback (rank=0, world_size=1).

2-Level Strided Sharding¶

1. Rank Sharding: Datasets first partition files across distributed ranks: files[rank::world_size]. Each GPU or node processes a disjoint subset of files.
2. Worker Sharding: Within each rank, the rank's subset is partitioned across DataLoader workers:

   per_worker = total // num_workers
   remainder  = total %  num_workers
   start      = worker_id * per_worker + min(worker_id, remainder)
   size       = per_worker + (1 if worker_id < remainder else 0)
   

Every file is seen exactly once across the entire cluster without duplication.

Reservoir Shuffle Buffer¶

Because streaming datasets do not load all samples into memory at once, setting shuffle_buffer_size=N enables an \(O(N)\) streaming reservoir shuffle. Incoming items from the dataset stream are mixed with past items in a rolling buffer, yielding pseudo-random batches across files and mosaics.