torchfits 1.0b1

FITS I/O for PyTorch with native tensor reads, datasets, and transforms

torchfits

torchfits reads and writes FITS files as PyTorch tensors. A multi-threaded C++ engine (vendored CFITSIO) handles images, tables, headers, compression, and MEF files. Optional datasets, transforms, and a torchfits CLI sit on top.

pip install torchfits

Requires Python 3.10+ and PyTorch 2.10. Docs: astroai.github.io/torchfits.

At a Glance

Task	torchfits
Image → GPU tensor	torchfits.read_tensor("img.fits", device="cuda")
Write a tensor	torchfits.write("out.fits", tensor)
Filter a catalog in C++	table.read(..., where="MAG < 20")
Open a MEF	with torchfits.open("mef.fits") as hdul: …
Train	FitsImageDataset + make_loader(..., num_workers=4)
Shell	torchfits info / header / convert / …

Features

• Fast FITS I/O — mmap image reads, compressed images, MEF, checksums
• Tables — Arrow-native with where= pushdown, scan/stream, Parquet/CSV/TSV/Arrow IPC export
• ML — torchfits.data datasets + make_loader
• Transforms — stretches, FITS scale/null handling, spectral prep (torchfits.transforms)
• CLI — MEF-aware inspect/convert tools (docs/cli.md)

Supported feature matrix: docs/parity.md.

What's New in 0.9.2 / 0.9.3

• CLI — torchfits for info, header, verify, stats, table, cutout, convert, … (docs/cli.md)
• Leaner imports — transforms from torchfits.transforms; use read / read_tensor (not read_fast / read_image)
• Convert — tables → Parquet, CSV, TSV, or Arrow IPC; images → Lupton PNG
• Scorecard — CUDA 0 strict deficits; Linux CPU 1; Mac MPS 16 (docs/benchmarks.md)

Full notes: docs/changelog.md.

Transforms

from torchfits.transforms import ArcsinhStretch, BackgroundSubtract, Compose, ZScaleNormalize

pipeline = Compose([BackgroundSubtract(), ArcsinhStretch(a=0.1), ZScaleNormalize()])
normalized = pipeline(image)              # forward → model input
restored = pipeline.inverse(normalized)   # inverse → physical flux

Representative classes (full catalog in docs/api.md):

Category	Examples	Inverse
Image stretches	ArcsinhStretch, LogStretch, SqrtStretch, ZScaleNormalize, RobustNormalize, MinMaxNormalize, PercentileClipNormalize	✓
Background / normalization	BackgroundSubtract, GlobalScalarNorm, FITSHeaderNormalize	✓
Header / table	FITSHeaderScale, FITSScaleColumns, TNullToNan	✓
Spectral	ContinuumNormalize, ContinuumRemoval, DopplerShift, SpectralBinning	✓ (except BandMath)
Continuum estimators	AsymmetricLeastSquares, AlphaShapeContinuum, WaveletDecompose, SavitzkyGolayFilter, RunningPercentile, UpperEnvelopeContinuum	✓
Outlier / time	SigmaClip, AsymmetricSigmaClip, PhaseFold	✗ (lossy or many-to-one)

Runnable demos: examples/example_transforms.py (image pipeline), examples/example_hyperspectral.py (spectral cube), examples/example_time_series.py (light curves).

Performance

Lab multi-host exhaustive scorecard (exhaustive_mps_20260717_040150 Mac MPS, exhaustive_cpu_20260717_040146 CANFAR CPU, exhaustive_cuda_20260717_042840 CANFAR CUDA); see docs/benchmarks.md for methodology, full exhaustive table, category summaries, RSS columns, and deficit transparency.

Under the strict gate (images: any lag; Arrow tables: ≤1.05×), CANFAR CUDA reports 0 TorchFits deficits; Linux CPU 1 (narrow 1M-row predicate); Mac MPS 16. MPS is not the Linux CUDA release gate.

Headline numbers

Case	torchfits	astropy	fitsio	Speedup vs astropy
Large float32 image read (16 MB, CPU)	6.51 ms	13.95 ms	8.83 ms	2.1×
Compressed Rice image (CPU)	15.17 ms	75.41 ms	18.45 ms	5.2×
50× repeated 100×100 cutouts (CPU)	21.75 ms	335.26 ms	21.03 ms	18.3×
Table read (100k rows, 8 cols)	6.97 ms	95.60 ms	30.20 ms	13.7×
Varlen table read (100k rows, 3 cols)	258.11 ms	1.624 s	337.40 ms	6.4×

By benchmark category

Category ranges and the full exhaustive table live in docs/benchmarks.md (CANFAR CUDA exhaustive_cuda_20260717_042840). Use the headline table above for this release’s absolute timings.

Current deficits

Scorecard policy (same-mmap peers):

• Images / cubes / spectra / cutouts: any lag above float-timer ε is a deficit (rice/hcompress included — no percent floor).
• Arrow tables: allow up to 1.05×.

Prior “0 deficit” claims used a 25% lag floor and are retracted. Re-score after the SIMD endian + thin device + WHERE⇒mmap-scan fixes; see docs/benchmarks.md.

GPU integer reads: Default read(..., device="cuda") applies BSCALE/BZERO on device and returns float32 for generic scaled pixels — good for ML. For native integer dtypes (int8, uint16) matching fitsio, use read_tensor(..., raw_scale=True) or rely on the automatic signed-byte / unsigned-integer fast paths (see benchmarks doc). Tables remain CPU-resident in all backends; GPU rows measure host decode + H2D copy, not disk→GPU bypass.

ML DataLoader (local diagnostic, not in lab CSV): 30×512² float32, CPU, 2 epochs — torchfits 1.12× vs fitsio on Rice-compressed files; uncompressed within ~4%. make_loader(..., optimize_cache=True) warms handle caches automatically when the dataset exposes a files attribute.

Install

pip install torchfits

Pre-built wheels are available for Linux x86_64 and macOS arm64. No system CFITSIO is needed—it is vendored and compiled automatically. Other architectures install from source when a compatible compiler and PyTorch are available.

From source:

git clone https://github.com/astroai/torchfits.git
cd torchfits
pip install -e .

Requires Python 3.10+, a C++17 compiler, CMake 3.21+, and PyTorch 2.10.

Quick Start

Read an image to GPU

import torchfits

data, header = torchfits.read("science.fits", device="cuda", return_header=True)
# data: torch.Tensor on CUDA, shape e.g. (4096, 4096), dtype torch.float32

tensor = torchfits.read_tensor("science.fits", hdu=0, device="cuda")

PyTorch DataLoader

from torchfits.data import FitsImageDataset, make_loader

ds = FitsImageDataset("observations/*.fits", label_key="CLASS")
loader = make_loader(ds, batch_size=32, num_workers=4)

for images, labels in loader:
    ...  # images: [B, 1, H, W] when add_channel_dim=True (default)

Filter and stream a catalog

# Predicate pushdown — only matching rows leave C++
table = torchfits.table.read(
    "catalog.fits",
    columns=["RA", "DEC", "MAG_G"],
    where="MAG_G < 20.0 AND CLASS_STAR > 0.9",
)
# table: pyarrow.Table

# Stream 100M rows in constant memory
for batch in torchfits.table.scan("survey.fits", batch_size=50_000):
    process(batch)  # batch: pyarrow.RecordBatch

Multi-HDU access

with torchfits.open("multi_ext.fits") as hdul:
    print(hdul)            # pretty-printed summary
    img = hdul[0].data     # image tensor
    tbl = hdul[1].data     # dict-like table accessor
    tbl_filtered = hdul[1].filter("FLUX > 100 AND FLAG = 0")

Write back

torchfits.write("output.fits", data, header=header, overwrite=True)
# table_dict is a dict of column names to 1D arrays/tensors
torchfits.table.write("catalog_out.fits", table_dict, header=header, overwrite=True)

Shell (CLI)

torchfits info science.fits
torchfits header science.fits --keyword OBJECT --json
torchfits verify science.fits
torchfits stats science.fits --hdu 0
torchfits convert catalog.fits out.csv --to csv --hdu 1

Benchmarks

torchfits is benchmarked across FITS image I/O (1D/2D/3D, all integer and float dtypes, compressed, scaled, MEF, cutouts, time series) and FITS table I/O (read, projection, row slicing, predicate filtering, streaming). GPU (CUDA) transport rows are included for image reads.

Comparators are astropy.io.fits and fitsio; selected CFITSIO behavior is validated through the torchfits native backend and smoke tests.

Methodology, full exhaustive table, category summaries, and known deficits: docs/benchmarks.md

Documentation

Published site: astroai.github.io/torchfits

Documentation site	Browse all docs on GitHub Pages
API Reference	Full public API with signatures and examples
CLI	torchfits command-line tools
Migration from Astropy	Side-by-side workflow translation
Migration from fitsio	Side-by-side workflow translation
Dataset migration	Removed FITSDataset → torchfits.data
Roadmap	FITS I/O roadmap and parity tiers
Parity Matrix	Supported, partial, unsupported, and out-of-scope features
Examples	Runnable scripts for every major workflow
Installation	Build from source, GPU setup, troubleshooting
Benchmarks	Methodology, commands, and latest numbers
Changelog	Version history and migration notes
Release Checklist	Maintainer guide for cutting releases

Contributing

git clone https://github.com/astroai/torchfits.git
cd torchfits
pixi install
pixi run test

The project uses pixi for environment management, ruff for linting, and pytest for testing.

License

MIT