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)

Companions: ivar, mask, and data state¶

Every transform accepts a bare torch.Tensor, the {"flux", "ivar"?, "mask"?} dict that torchfits.data emits when you pass ivar_hdu= / mask_hdu=, or the typed Payload container. Built-ins read and rebuild whichever container you handed them — a dict in gives a dict out, and unknown dict keys (wavelength, meta, …) are carried through untouched.

Inverse variance propagates exactly through affine transforms. If flux becomes a·x + b, then ivar becomes ivar / a² (an additive offset leaves the variance alone). This holds for every normalizer, BackgroundSubtract, MeshBackgroundSubtract, AffineTransform, the header-scaling transforms and SigmaClip-free pipelines.

For the nonlinear stretches — ArcsinhStretch, LogStretch, SqrtStretch — pass propagate_ivar=True to propagate ivar by the delta method (ivar / (df/dx)²). For SqrtStretch on Poisson counts (var(x) = x) this recovers the classic variance-stabilised constant ivar = 4. Where a stretch clamps — LogStretch/SqrtStretch below zero — the map is locally flat and the output no longer constrains the input, so ivar is set to 0 there instead of a spuriously finite value.

Left at the default (False), a companion ivar is passed through unchanged and a warning fires once per instance, because it no longer strictly describes the transformed flux. The clipping transforms (SigmaClip, AsymmetricSigmaClip, PercentileClipNormalize) are projections rather than invertible maps, so they always pass ivar through and warn; the pixels they altered are reported in their _last_mask.

Data states. DataState names the processing stage of a payload so a transform can refuse data it would silently corrupt:

Data read with the default reader is already PHYSICAL, so FITSHeaderScale would scale it a second time. Declare the state in the payload and the transform raises DataStateError instead:

from torchfits.transforms import FITSHeaderScale

FITSHeaderScale(bscale=0.5)({"flux": counts, "state": "physical"})
# DataStateError: FITSHeaderScale expects state [stored], but the payload
# declares 'physical'.  The default reader returns *physical* values already
# (BSCALE/BZERO applied). Read with raw_scale=True for stored codes, or drop
# the header-scaling transform.

Bare tensors carry no declared state and are never rejected, so existing pipelines keep working. calibration_state(header, raw_scale=True) returns the state a reader actually produced, and Payload(flux=..., ivar=..., mask=..., state=...) is the typed equivalent of the dict payload.

The state is maintained: a transform stamps what it produced onto any payload that already declares a state, so FITSHeaderScale on a STORED payload yields a PHYSICAL one and a second FITSHeaderScale in the same pipeline raises instead of scaling twice. inverse() accepts what its own forward() produced and restores the state it consumed, so t.inverse(t(payload)) round-trips. Payloads that never declared a state are left exactly as they were (no state key is added).

Flux-scaling transforms (the normalizers, BackgroundSubtract, MeshBackgroundSubtract, FITSHeaderNormalize) accept STORED, PHYSICAL and NORMALIZED but refuse CONTINUUM_NORMALIZED: dividing an already-continuum-normalized spectrum by another per-spectrum statistic discards the common flux scale the normalization established. Clear the state if you really mean to rescale such data.

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, with analytic inverses. LogStretch / SqrtStretch clamp negative input to zero, so their inverse cannot recover negatives.

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. Negative input is clamped to zero (inverse cannot recover negatives).

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

SqrtStretch()¶

Square-root stretch — stabilizes Poisson variance. Negative input is clamped to zero (inverse cannot recover negatives).

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

Normalizers¶

Data-dependent — compute statistics from the image and cache them for inverse transforms. Require float input; see Invertibility.

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

IRAF zscale auto-contrast algorithm. Maps the display range [z1, z2] to [0, 1] — pixels outside [z1, z2] land outside [0, 1] (no clamp on the output).

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}\)

SigmaNormalize(dim=(-2, -1), stat="mad", zero_preserving=True, eps=1e-12, weighted=False)¶

Rescale so the sky noise is unit variance. This is the normalization the image foundation-model papers actually use (AstroCLIP / AstroPT style per-frame robust scaling), and unlike RobustNormalize it does not shift the data by default:

Zero-preserving mode (default):

Centered mode (zero_preserving=False):

Inverse: \(x = \text{output} \times \sigma + \text{offset}\)

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

Subtract median background level.

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

MeshBackgroundSubtract(mesh=(8, 8), n_sigma=3.0, filter_mesh=True, min_tile_pixels=4, weighted=False)¶

SExtractor-style tile-mesh background: estimate the sky independently in each of mesh tiles (sigma-clipped median + MAD), optionally median-filter the tile grid, then interpolate it back to full resolution and subtract.

The two trailing dimensions are spatial; leading dimensions are batched, so a [C, H, W] stack or cube works directly. Tiles with fewer than min_tile_pixels valid pixels fall back continuously to the frame-level background.

Inverse: \(x = \text{output} + \text{background}\) (the interpolated map is cached from forward).

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

Clip to percentile range, scale to [0, 1]. inverse() is approximate — values clipped by the clamp cannot be recovered exactly.

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}\)

InterquantileScale(q_low=0.05, q_high=0.95, dim=None, zero_preserving=True, eps=1e-9)¶

Zero-preserving or centered interquantile scale normalization.

In zero-preserving mode (default):

In centered mode (zero_preserving=False):

Inverse: \(x = \text{output} \times s + \text{offset}\)

AffineTransform(scale=1.0, offset=0.0)¶

Explicit flux * scale + offset with exact companion propagation. It is the linear building block every normalizer reduces to, and the right tool when you already know the coefficients — e.g. converting counts to flux from a zeropoint:

from torchfits.transforms import AffineTransform

zp = 25.0
counts_to_flux = AffineTransform(scale=10 ** (-0.4 * zp))
flux = counts_to_flux(payload)  # ivar scaled by scale**-2

Inverse: \(x = (\text{output} - \text{offset}) / \text{scale}\)

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

Masks & robust statistics¶

from torchfits.transforms import apply_mask, combine_masks, mask_from_dq, mask_from_ivar

valid = combine_masks(
    mask_from_dq(payload["dq"], bad_bits=[0, 11]), mask_from_ivar(payload["ivar"])
)
flux = apply_mask(payload["flux"], valid)

Transform statistics are mask-aware and inverse-variance aware. Pass mask= explicitly, or let a dict payload's mask field be picked up automatically (an explicit mask= argument wins). weighted=True on the normalizers, the clippers and estimate_background switches the underlying median / mean / RMS to inverse-variance weights, so low-weight pixels stop dragging the estimate, while the default (weighted=False) stays byte-identical to previous releases.

weighted=True is an exact no-op when no ivar is supplied (there is nothing to weight by), so it can be switched on unconditionally. Note the weighted median uses the inverted-CDF definition, so with uniform weights it returns the lower-middle element rather than the interpolated median; the two agree to within one order-statistic spacing, not bit-for-bit.

from torchfits.transforms import SigmaNormalize

clean = SigmaNormalize(weighted=True)({"flux": flux, "ivar": ivar, "mask": valid})

Helpers¶

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.

rgb(*bands, *, brightness=0.15, saturation=2.0, scene="auto", weights=None, calibrated=False, zeropoints=None)¶

Auto RGB from 1–7 aligned images (or one (C, H, W) cube). Band order is shortest wavelength first: rgb(g, r, i) or rgb(u, g, r, i, z). Returns a [H, W, 3] float tensor in [0, 1] after sRGB encoding.

Uncalibrated (default) subtracts each band's sky median and divides by MAD so ADU scale and sky pedestals do not paint the colour. calibrated=True skips that (bands already on one flux scale). zeropoints= is the AB magnitude of 1 count per band, converted to nanomaggies as counts * 10**(-0.4*(zp - 22.5)).

See examples/example_rgb_sky.py.

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

Lupton asinh RGB from three single-band tensors (same shape). Returns a [H, W, 3] float64 tensor in [0, 1] (channel last; accepts any tensor-convertible inputs). 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):
    view = self.view(x)  # validates the declared data state
    flux = view.flux * self.scale + self.offset
    return view.replace(flux, ivar=self.scale_ivar(view.ivar, self.scale))

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

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 (
    AffineTransform,
    ArcsinhStretch,
    AsModule,
    AsymmetricSigmaClip,
    BackgroundSubtract,
    Compose,
    DataState,
    DataStateError,
    FITSScaleColumns,
    FITSHeaderNormalize,
    FITSHeaderScale,
    GlobalScalarNorm,
    InterquantileScale,
    InterquantileNormalize,
    LogStretch,
    MeshBackgroundSubtract,
    MinMaxNormalize,
    Payload,
    PercentileClipNormalize,
    RobustNormalize,
    SigmaClip,
    SigmaNormalize,
    SqrtStretch,
    TNullToNan,
    ZScaleNormalize,
    apply_mask,
    as_module,
    calibration_state,
    combine_masks,
    estimate_background,
    lupton_rgb,
    mask_from_dq,
    mask_from_ivar,
    mask_from_nan,
    rgb,
    zscale_limits,
)