torchfits 0.9.2

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

torchfits

torchfits is FITS I/O for PyTorch: a multi-threaded C++ engine with vendored CFITSIO that reads and writes images, headers, HDUs, compressed images, and tables as tensors — without NumPy-to-torch glue. Optional torchfits.data datasets and torchfits.transforms provide header-aware preprocessing for ML training loops.

It is not a full replacement for Astropy, fitsio, or CFITSIO. The supported surface is documented explicitly in docs/parity.md, with source-backed tests for each claimed parity area. WCS, sky-coordinate models, HEALPix, sphere geometry, sky-domain simulation, and photometric physics are out of scope. Transforms cover FITS scale/null/dtype and tensor preprocessing only.

At a Glance

Task	Traditional stack	torchfits equivalent
Read image to GPU	astropy/fitsio → numpy → torch → .to(device)	torchfits.read_tensor("img.fits", device="cuda")
Write tensor to FITS	tensor → numpy → astropy HDU → writeto	torchfits.write("out.fits", tensor)
Filter large table	load all rows → mask in Python	where="MAG < 20" pushdown in C++
Read multi-extension files	manual HDU dispatch	with torchfits.open("mef.fits") as hdul: ...
PyTorch training loop	hand-rolled Dataset + cache tuning	FitsImageDataset + make_loader(..., num_workers=4)
Normalize for model input	ad-hoc scaling in the training script	Compose([BackgroundSubtract(), ZScaleNormalize()])
Verify FITS checksums	comparator-specific helpers	torchfits.verify_checksums(path)

Features

FITS I/O — Multi-threaded C++ core with SIMD-optimized type conversion, memory-mapped image reads, intelligent chunking, and adaptive buffering. Reads and writes images, binary/ASCII tables, compressed images, and multi-extension FITS files with header round-trip coverage.

Table Engine — Arrow-native table API with predicate pushdown (where=), column projection, row slicing, streaming scan(), and in-place mutations (append, insert, update, delete rows and columns). Interop with Pandas, Polars, DuckDB, and PyArrow.

ML Data Layer — torchfits.data ships FitsImageDataset, FitsImageIterableDataset, FitsTableDataset, FitsTableIterableDataset, FitsCutoutDataset, and make_loader with automatic handle-cache warm-up.

Transforms — 28 FITSTransform classes for image stretches, header-aware scaling (FITSHeaderScale, FITSScaleColumns, TNullToNan), spectral/hyperspectral preprocessing, continuum estimators, and outlier rejection. Most ship .inverse() for decoding model outputs back to physical units. See docs/api.md and examples/example_transforms.py.

Compatibility Contract — Parity is tracked by tier: truthful public docs, fitsio core workflow parity, Astropy common workflow parity, selected CFITSIO backend behavior, and explicit non-goals. See docs/parity.md.

What's New in 0.9.2

0.9.2 freezes a leaner Python library surface and ships an early CLI (the operator surface planned for 0.9.3):

• Leaner root API — transforms live under torchfits.transforms only; read_fast / read_image removed from the public surface; torchfits.hdu is a documented namespace.
• CLI — torchfits one-word commands (info, header, verify, diff, stats, table, convert, copy, arith, cutout, compress, decompress, transform, probe, setkey); see docs/cli.md.
• Docs and brand — site logo uses torchfits-logo.png; API docs match the settled exports; README benchmark run IDs track docs/benchmarks.md.
• Deficit honesty — Linux CPU/CUDA strict-gate scores are separate from Mac MPS deficits (see Performance below).

0.9.1 constrained the native ABI to PyTorch 2.10. 0.9.0 added Polars-native table access and handle-cache safety. See docs/changelog.md. Roadmap: docs/roadmap.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_000853 Mac MPS, exhaustive_cpu_20260716_191252 CANFAR CPU, exhaustive_cuda_20260716_191255 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×), Linux CPU and CUDA currently report 0 TorchFits deficits. Mac MPS reports 4 deficits on this refresh (down from an earlier 101-row snapshot); MPS is not the Linux CUDA release gate.

Headline numbers

Case	torchfits	astropy	fitsio	Speedup vs astropy
Large float32 image read (16 MB, CPU)	3.85 ms	16.67 ms	5.89 ms	4.3×
Compressed Rice image (CPU)	9.06 ms	27.77 ms	9.43 ms	3.1×
50× repeated 100×100 cutouts (CPU)	4.68 ms	75.36 ms	4.94 ms	16.7×
Table read (100k rows, 8 cols)	95.3 μs	6.74 ms	59.84 ms	70.6×
Varlen table read (100k rows, 3 cols)	93.9 μs	3.52 ms	288.81 ms	37.5×

By benchmark category

Category	Best speedup vs astropy	Best speedup vs fitsio	Notes
1D images (float32/64, int8–int64)	7.8×	2.3×	All sizes, CPU
2D images (float32/64, int8–int64, uint16/32)	7.7×	2.4×	All sizes, CPU
3D cubes (float32/64, int8–int64)	7.5×	2.1×	Small–medium, CPU
Compressed (gzip, rice, hcompress)	4.3×	1.1×	rice and gzip dominate; hcompress slightly behind fitsio
Scaled (BSCALE/BZERO)	6.1×	1.8×	Automatic integer→float scaling
MEF (multi-extension)	9.9×	2.4×	Small–medium files
Repeated cutouts (50× 100×100)	16.7×	1.1×	Open-once, subset many times
Time series frames	5.1×	1.9×	5 sequential frames
Header reads (all fixtures)	9.5×	1.5×	Sub-100 μs for all backends
Table: read_full	115×	628×	100k rows, 8 cols
Table: projection	147×	91×	Column subset
Table: row_slice	147×	162×	Row range
Table: predicate_filter	57×	25×	WHERE clause
GPU (CUDA) images	78×	3.0×	tiny–large, all dtypes

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.0+.

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)

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