Transforms¶

Header-aware preprocessing for FITS images and tables.

When to use¶

• High-dynamic-range visualization (arcsinh, zscale, log/sqrt stretches)
• Model input scaling that you want reusable across a Dataset
• FITS BSCALE / null hygiene (FITSHeaderScale, TNullToNan, column scale)

When not to¶

• You need raw ADU / physical values as stored — use read_tensor / table.read
• One-off arithmetic on a single tensor — plain PyTorch is enough
• Catalog filtering — use table.read(..., where=) (C++ pushdown), not transforms

Wire a pipeline into training with FitsTensorDataset(..., transform=pipeline) or call it on a tensor from read_tensor. See Data module for when to introduce a Dataset / make_loader.

All transforms implement the FITSTransform callable protocol (forward / inverse / __call__). They are not torch.nn.Module subclasses — wrap with :func:as_module (or :class:AsModule) for nn.Sequential. Inverse state is instance-local (_last_* fields); construct one pipeline per worker when num_workers > 0.

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

pipeline = Compose([BackgroundSubtract(), ArcsinhStretch(a=0.1), ZScaleNormalize()])
normalized = pipeline(image)
restored = pipeline.inverse(normalized)

Masks¶

Most transforms accept an optional boolean mask (True = valid). Compose forwards the same mask to every child:

pipeline(image, mask=finite_mask)
pipeline.inverse(normalized, mask=finite_mask)

Invertibility¶

Stateless stretches are the most likely to work under torch.compile; data-dependent normalizers cache Python-side state and may graph-break. There is no certified compile matrix yet.

The implementation lives under torchfits.transforms as a small package (stretch, normalize, fits_meta, clip, rgb) re-exported from torchfits.transforms.

Stretches¶

Stateless, always invertible. Apply non-linear flux scaling for visualization.

ArcsinhStretch(a=1.0)¶

Lupton+ (2004) arcsinh stretch — LSST/SDSS standard.

Inverse: \(x = \frac{\sinh(\text{output} \cdot \operatorname{arcsinh}(a))}{a}\)

LogStretch(a=1000.0, eps=1e-9)¶

Logarithmic stretch for heavy-tailed flux distributions.

Inverse: \(x = \frac{10^{\text{output} \cdot \log_{10}(1 + a)} - 1}{a}\)

SqrtStretch()¶

Square-root stretch — stabilizes Poisson variance.

Inverse: \(x = \text{output}^2\)

Normalizers¶

Data-dependent, invertible. Compute statistics from the image and cache them for inverse transforms.

ZScaleNormalize(contrast=0.25, dim=(-2, -1))¶

IRAF zscale auto-contrast algorithm. Maps display range to [0, 1].

Both clamped to \([x_{\min}, x_{\max}]\).

Inverse: \(x = \text{output} \cdot (z_2 - z_1) + z_1\)

RobustNormalize(dim=(-2, -1))¶

Subtract median, divide by MAD-derived standard deviation.

where \(\text{MAD} = \text{median}(|x - \text{median}(x)|)\).

Inverse: \(x = \text{output} \cdot \text{std\_approx} + \text{median}\)

BackgroundSubtract(dim=(-2, -1))¶

Subtract median background level.

Inverse: \(x = \text{output} + \text{median}\)

PercentileClipNormalize(lower_pct=1, upper_pct=99, dim=(-2, -1))¶

Clip to percentile range, scale to [0, 1].

Inverse: \(x = \text{output} \cdot (\text{upper} - \text{lower}) + \text{lower}\)

MinMaxNormalize(dim=(-2, -1))¶

Min-max normalization to [0, 1] with ULP-safe epsilon.

Inverse: \(x = \text{output} \cdot (v_{\max} - v_{\min}) + v_{\min}\)

GlobalScalarNorm(stat="median", dim=None)¶

Divide by a single scalar statistic. Minimal linear prep.

where scalar is one of: median(x), max(x), mean(x), or \(\sqrt{\text{mean}(x^2)}\) (RMS).

Inverse: \(x = \text{output} \times \text{scalar}\)

Outlier Rejection¶

SigmaClip(n_sigma=3.0, max_iter=5, dim=(-2,-1), fill="mean")¶

Iterative sigma-clipping with mean or median fill.

1. Compute mean \(\mu\) and std \(\sigma\) over dims.
2. Mask pixels where \(|x - \mu| > n_\sigma \cdot \sigma\).
3. Repeat until convergence or max_iter.
4. Replace clipped values with surviving mean or median.

Inverse: None (lossy).

AsymmetricSigmaClip(n_low=3.0, n_high=3.0, dim=(-2,-1))¶

One-pass asymmetric sigma-clip via estimate_background (median + MAD).

Outliers replaced with median.

Inverse: None (lossy).

FITS Metadata¶

FITSHeaderScale(bscale=1.0, bzero=0.0)¶

Apply/remove FITS BSCALE/BZERO linear scaling.

Inverse: \(x = \frac{\text{output} - \text{BZERO}}{\text{BSCALE}}\)

Factory: FITSHeaderScale.from_header(header) — extracts BSCALE/BZERO from a FITS header dict.

FITSScaleColumns(scales)¶

Per-column BSCALE/BZERO for table tensors.

Inverse: \(x[c] = \frac{\text{output}[c] - \text{TZERO}_c}{\text{TSCAL}_c}\)

Factory: FITSScaleColumns.from_header(header).

TNullToNan(nulls)¶

Map FITS TNULL sentinels to NaN in table columns.

Inverse: None (lossy).

Factory: TNullToNan.from_header(header).

FITSHeaderNormalize(header, scale_floats=False)¶

Auto-normalize from BITPIX/BSCALE/BZERO. Integer types mapped to [0, 1].

• Integer types (BITPIX 8/16/32/64): min-max normalize using the native range scaled by BSCALE/BZERO.
• Float types (BITPIX -32/-64): identity unless scale_floats=True.

Inverse: Yes (for normalized types).

Utility¶

Compose(transforms)¶

Chain transforms; inverse() unwinds in reverse order.

pipeline = Compose([
    BackgroundSubtract(),
    ArcsinhStretch(a=0.1),
    ZScaleNormalize(),
])
normalized = pipeline(image)
original = pipeline.inverse(normalized)

FITSTransform¶

Base class for custom transforms. Override forward() and optionally inverse(). __call__ delegates to forward(). Not an nn.Module.

as_module(transform) / AsModule¶

Wrap a FITSTransform as a thin nn.Module for nn.Sequential:

import torch.nn as nn
from torchfits.transforms import ArcsinhStretch, as_module

model = nn.Sequential(as_module(ArcsinhStretch(a=0.1)), nn.Linear(64, 10))

Only the forward pass is exposed; call transform.inverse on the wrapped instance for undo.

lupton_rgb(r, g, b, *, Q=8.0, stretch=0.5)¶

Lupton asinh RGB from three single-band tensors (same shape). Returns a [H, W, 3] float tensor in [0, 1]. Matches Astropy's make_lupton_rgb / LuptonAsinhStretch mapping. See examples/example_lupton_rgb_sdss.py.

Writing a custom transform¶

Subclass FITSTransform, implement forward() (and inverse() if the operation is invertible):

import torch
from torchfits.transforms import FITSTransform


class ScaleOffset(FITSTransform):
    """Affine transform: forward(x) = x * scale + offset."""

    def __init__(self, scale: float, offset: float) -> None:
        self.scale = scale
        self.offset = offset

    def forward(
        self, x: torch.Tensor, mask: torch.Tensor | None = None
    ) -> torch.Tensor:
        return x * self.scale + self.offset

    def inverse(
        self, x: torch.Tensor, mask: torch.Tensor | None = None
    ) -> torch.Tensor:
        return (x - self.offset) / self.scale

Compose it with built-in transforms — Compose calls each child with the same tensor and mask= kwarg, and unwinds inverse() in reverse order:

from torchfits.transforms import BackgroundSubtract, Compose

xf = ScaleOffset(scale=2.0, offset=-10.0)
pipeline = Compose([BackgroundSubtract(), xf])
out = pipeline(image)
restored = pipeline.inverse(out)

Wire it into a Dataset so every sample gets the same preprocessing:

from torchfits.data import FitsImageDataset

dataset = FitsImageDataset("images/*.fits", hdu=0, transform=xf)

def forward(self, x, mask=None):
    if isinstance(x, dict):
        flux = x["flux"] * self.scale + self.offset
        return {**x, "flux": flux}
    return x * self.scale + self.offset

Full runnable version, including the FitsImageDataset wiring above: examples/example_custom_transform.py.

Helpers¶

Importing¶

Import transform classes from torchfits.transforms (namespace-only since 0.9.2 — they are not re-exported at the package root):

from torchfits.transforms import (
    ArcsinhStretch, AsModule, BackgroundSubtract, Compose, ZScaleNormalize,
    RobustNormalize, MinMaxNormalize, PercentileClipNormalize,
    LogStretch, SqrtStretch, GlobalScalarNorm,
    AsymmetricSigmaClip, SigmaClip,
    FITSScaleColumns, TNullToNan, FITSHeaderNormalize,
    as_module, lupton_rgb,
)