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).

Cubes and IFU cubes¶

The leading axis of a cube is the spectral (or velocity / channel) axis. Most IFU pipelines want a range, not a single channel:

from torchfits.data import FitsCubeDataset

# A rest-wavelength window instead of the whole cube.
ifs = FitsCubeDataset("cubes/*.fits", spectral_slice=(1200, 2600), label_key="OBJID")

The same window is applied to companion ivar / mask extensions, so they stay aligned with flux. FitsCubeIterableDataset(..., spectral_slice=(a, b)) does the same while streaming.

Discovering bands¶

discover_bands(path) lists every 2D+ image extension with its name, role and photometric metadata, so datasets can be built from band names instead of hand-counted HDU indices:

from torchfits.data import FitsImageDataset, discover_bands

for band in discover_bands("mosaic.fits"):
    print(band.name, band.role, band.shape, band.zeropoint)
# G flux (4096, 4096) 25.3
# R flux (4096, 4096) 25.1
# Z flux (4096, 4096) 24.8
# G_IVAR ivar (4096, 4096) None

ds = FitsImageDataset.from_bands("mosaic.fits", bands=["G", "R", "Z"])

BandInfo.role is "flux", "ivar", "mask" or "wavelength", inferred from the name suffix, so from_bands(bands=None) selects the flux extensions and automatically attaches matching G_IVAR / G_DQ companions when every selected band has one. A *_DQ companion is a bitfield, so from_bands decodes it into a boolean validity mask by default (mask_is_dq=False to override, bad_bits= to name the fatal bits); mixing *_DQ and *_MASK companions across bands is rejected as ambiguous. Zeropoints come from PHOTZEROPOINT, ZP, MAGZERO, PHOTZP, MAGZP, ABMAGZERO or ZEROPOINT (first present wins; ZP and friends are read as HIERARCH where needed):

zp = discover_bands("mosaic.fits")[0].zeropoint
counts_to_flux = 10 ** (-0.4 * zp)

# or, straight from the dataset:
ds.band_zeropoints()  # {'G': 25.3, 'R': 25.1, 'Z': 24.8}

BandInfo.flux_scale(exptime_normalized=False) returns that factor divided by EXPTIME for flux per second. Feed it to AffineTransform to apply it inside a pipeline.

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). When labels= or label_key= is given the item becomes (payload, label); otherwise the payload is returned alone, so existing call sites are unchanged.

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, or one chunk of batch_size rows when as_batches=True.

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.