torchfits 0.5.0b3

High-performance FITS I/O for PyTorch

torchfits

torchfits is a focused FITS I/O library for PyTorch. It reads and writes FITS images, headers, HDUs, compressed images, and FITS tables through a multi-threaded C++ engine with vendored CFITSIO, returning tensor-native data without requiring users to build NumPy-to-torch glue code.

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, and sky-domain simulation workflows are out of scope for torchfits.

At a Glance

Task	Traditional stack	torchfits equivalent
Read image to GPU	astropy/fitsio → numpy → torch → .to(device)	torchfits.read("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: ...
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.

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.

Install

pip install torchfits

Pre-built wheels are available for Linux and macOS (x86_64, arm64). No system CFITSIO needed—it's vendored and compiled automatically.

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

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 benchmark evidence is limited to FITS image I/O and FITS table I/O. Comparators are astropy.io.fits and fitsio; selected CFITSIO behavior is validated through the torchfits native backend and smoke tests.

Methodology, reproducible commands, results, and known deficits: docs/benchmarks.md

Documentation

API Reference	Full public API with signatures and examples
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

GPL-2.0