Data Module¶

torchfits.data provides PyTorch Dataset / IterableDataset classes for FITS images and tables, plus make_loader — a thin factory around torch.utils.data.DataLoader with torchfits cache warm-up defaults.

Not training? Stay on Core I/O (read_tensor) or Tables (table.read). Datasets exist for multi-sample epochs, shuffle, and worker parallelism — not for one-off reads.

Layering: read_tensor / table.read → optional transforms → Fits*Dataset → make_loader → training loop. Pass a transform into the dataset when every sample needs the same preprocess; call transforms yourself when exploring a single file.

Choosing a Dataset¶

Channel semantics: flux bands / arms / CCDs stack as tensor channels. IVAR and mask are always companion tensors (or dict fields) — never mixed into flux channels. Pass ivar_hdu= / mask_hdu= (or table ivar_column=).

Dataset transform= signature: each dataset calls transform(payload) with one positional argument — the full payload (a Tensor, or a {"flux", "ivar"?, "mask"?} dict when ivar_hdu= / mask_hdu= are set). There is no separate mask= kwarg at the Dataset boundary, even though FITSTransform.forward(x, mask=None) accepts one when you call a transform directly on a tensor. Write custom transforms to branch on dict vs tensor input if they need mask-aware Dataset wiring — see custom transforms.

Cache: how and when¶

Two independent caches exist: the disk cache (remote downloads, example samples) and the in-process handle/metadata cache (torchfits.cache). Dataset training loops touch both; a single-file read_tensor call touches neither.

Disk cache roots: TORCHFITS_CACHE_DIR / TORCHFITS_REMOTE_CACHE / TORCHFITS_SAMPLE_CACHE — see the User-facing table in Environment variables for defaults.

Datasets accept cache_dir= to override the remote materialization directory per-instance, taking priority over the environment variables for that Dataset only.

Handle/metadata cache: make_loader(ds, optimize_cache=True) (default) calls cache.optimize_for_dataset(ds.files, avg_file_size_mb=...), which sizes the C++ handle cache for the file count. This only fires when the dataset exposes a files list — every Fits*Dataset built on FitsTensorDataset (FitsImageDataset, FitsCubeDataset, FitsSpectrumDataset, FitsTensorIterableDataset, FitsCutoutDataset) has one. FitsTableDataset / FitsTableIterableDataset read one file through CFITSIO's own buffered path and do not expose files; optimize_cache=True is then a documented no-op, not an error.

from torchfits import cache

cache.get_cache_stats()   # hit rate, cpp cache size, config
cache.clear_cache()       # drop Python + C++ I/O + table-handle caches

When not to warm: a single local file read (read_tensor, one-off table.read) opens and closes one handle — there is nothing to warm and optimize_for_dataset / configure_for_environment add overhead without benefit. Warm-up pays off once a loader iterates many files (remote prefetch, multi-worker DataLoader, repeated cutouts from a mosaic 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, ...],  # 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()¶

Thin factory: torch.utils.data.DataLoader(dataset, ...) plus an optional call to cache.optimize_for_dataset before the loader is returned. Any Dataset — torchfits or plain PyTorch — works with either path; the two produce identical batches when optimize_cache=False.

from torchfits.data import make_loader

loader = make_loader(
    ds,
    batch_size=32,
    num_workers=4,
    pin_memory=True,       # faster CPU → GPU transfer
    optimize_cache=True,   # warm handle caches before iteration
    shuffle=True,          # default for map-style; False for iterable
)

Build DataLoader yourself when any of these apply:

• You already have a collate_fn, sampler, or batch_sampler that make_loader's keyword surface doesn't expose directly (pass through **loader_kwargs, but a hand-built call is clearer once you're combining several).
• Your dataset has no files attribute — optimize_cache=True is then a silent no-op (see Cache: how and when); skip the torchfits wrapper and set optimize_cache=False or call DataLoader directly.
• You manage handle-cache warm-up elsewhere (e.g. once per training run rather than per loader construction).

Side-by-side runnable comparison: 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.

Worker Sharding¶

Images (FitsImageIterableDataset)¶

File indices are partitioned by worker_id:

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 per epoch. Shuffling permutes within each worker's shard (epoch-seeded).

Tables (FitsTableIterableDataset)¶

Scan batches are assigned to workers via batch_idx % num_workers == worker_id. Row order within a batch is preserved. Workers do not interleave individual rows.