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.

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"?} and label is a torch.long scalar. Without labels= or label_key= the label is still there and is always tensor(0) — kept so existing for x, y in loader loops keep working. FitsCutoutDataset, FitsSpectrumDataset and the iterable datasets return the payload alone instead when no labels are configured. Pass return_label=False to get the same payload-only items here:

from torchfits.data import FitsImageDataset

images = FitsImageDataset("obs/*.fits", return_label=False)
x = images[0]  # a Tensor, no (x, tensor(0)) pair

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), with label == tensor(0) when no labels are configured; return_label=False returns the payload alone.

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

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, PHOTZP, ZP, MAGZERO, 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 disjoint, deterministic subset — every file is seen exactly once per epoch. Which order those files arrive in is a separate question, answered by shuffle and seed (see the note below).

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) alone. With labels= or label_key= the item becomes (payload, label), label a torch.long scalar. Each label stays with its own file through rank and worker sharding, shuffle= and shuffle_buffer_size= (duplicate paths keep their own labels):

from torchfits.data import FitsImageIterableDataset, make_loader

files = ["a.fits", "b.fits", "c.fits"]
ds = FitsImageIterableDataset(files, labels=[0, 1, 0], shuffle=True, seed=0)
if __name__ == "__main__":
    for image, label in make_loader(ds, batch_size=2, num_workers=2):
        ...

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, which is tensor(0) for every row when labels is omitted. return_label=False returns the row dict alone.

Rows hold the values as stored, whether or not where= is set: where= only selects rows. Integer columns keep their dtype and a TNULL null comes back as its sentinel value; character columns are uint8 byte tensors (space padded). Convert nulls to NaN with TNullToNan:

import torchfits
from torchfits.data import FitsTableDataset
from torchfits.transforms import TNullToNan

header = torchfits.read_header("catalog.fits", 1)
ds = FitsTableDataset(
    "catalog.fits", where="MAG_G < 20", transform=TNullToNan.from_header(header)
)

FitsTableIterableDataset¶

Streams table rows in constant memory. Each iteration 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.

Iterating the unfiltered dataset yields exactly as many rows as FitsTableDataset reports len for the same file and HDU, including when every column is variable-length. A variable-length column arrives from scan_torch as a Python list rather than a tensor, so the per-chunk row count is read from tensor or list columns.

Rows and chunks follow the same convention as FitsTableDataset: values as stored, with or without where=. A TNULL null is its sentinel value and character columns are uint8 byte tensors; use TNullToNan for NaN nulls. where= accepts the full table.read predicate dialect and is evaluated per chunk on the predicate columns only.

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. The public CutoutSpec type alias names the normalized six-item form: tuple[str, int, int, int, int, int], corresponding to (path, hdu, x1, y1, x2, y2) with half-open upper bounds.

Every spec must describe a window that actually selects pixels. The origin must be non-negative (x1 >= 0, y1 >= 0) and the half-open window must be non-empty (x1 < x2 and y1 < y2); otherwise the constructor raises ValueError. A degenerate window is rejected at construction rather than yielding a 0-pixel tensor, because fits_collate_fn stacks such shapes without complaint and make_loader would hand a model a batch whose numel() is 0. A window whose far edge runs off the image is still clamped to the overlap.

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 with torchfits collation and shuffle defaults.

from torchfits.data import make_loader

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

Use a plain DataLoader when you already own collate_fn or samplers. 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)
• Nested dicts, such as multi-arm spectra {arm: {"flux": Tensor, ...}}, are collated per leaf: batch["R"]["flux"] has shape (batch, n_pixels)

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, first present wins. Rank: RANK (torchrun, DeepSpeed, FSDP), SLURM_PROCID (SLURM), then LOCAL_RANK. World size: WORLD_SIZE, SLURM_NTASKS, then SLURM_NPROCS. 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.

The buffer is re-seeded every epoch, so each epoch mixes differently. It is a reservoir shuffle, not a uniform permutation: an item's chance of landing in a given output slot rises with its position in the stream, which is the standard streaming trade-off (measured here over 20,000 seeds with \(N=8\) and buffer 3: the last input item ended up in the final slot 33.6% of the time, the first 4.4%). The buffer size is therefore the shuffle window — raise it towards the dataset size when you need a stronger mix, at \(O(N)\) memory.