from __future__ import annotations
import numpy as np
from astropy.wcs import WCS
from astropy.coordinates import SkyCoord, SkyOffsetFrame
import astropy.units as u
def _sky_residuals_arcsec(
wcs: WCS, xy: np.ndarray, frame, off: SkyOffsetFrame, sky_off: SkyCoord, coslat: np.ndarray
) -> np.ndarray:
"""
Residuals in offset-frame arcsec around the field center.
The reference positions are precomputed in the offset frame (`sky_off`,
with `coslat` = cos of their offset latitudes) since they do not change
during the fit. Longitude offsets are scaled by cos(lat) so that both
components are proper angular distances even far from the center.
Returns concatenated [d_lon_arcsec, d_lat_arcsec] per point.
"""
ra, dec = wcs.all_pix2world(xy[:, 0], xy[:, 1], 0)
model = SkyCoord(ra * u.deg, dec * u.deg, frame=frame)
model_off = model.transform_to(off)
dlon = (model_off.lon - sky_off.lon).to_value(u.arcsec) * coslat
dlat = (model_off.lat - sky_off.lat).to_value(u.arcsec)
return np.concatenate([dlon, dlat])
def _spherical_mean(sky: SkyCoord) -> SkyCoord:
"""Direction of the mean unit vector of *sky* (robust to the lon=0/360 wrap)."""
xyz = sky.cartesian.xyz.value
mean = np.mean(np.atleast_2d(xyz), axis=1)
norm = np.linalg.norm(mean)
if not np.isfinite(norm) or norm <= 0:
raise ValueError("Cannot determine field center from reference positions")
mean = mean / norm
return SkyCoord(
np.degrees(np.arctan2(mean[1], mean[0])) * u.deg,
np.degrees(np.arcsin(np.clip(mean[2], -1.0, 1.0))) * u.deg,
frame=sky.frame,
)
def _pack_params(w: WCS, pv_deg: int) -> np.ndarray:
# CRPIX, CRVAL, CD(2x2), PV2_0..PV2_pv_deg
p = []
p += [float(w.wcs.crpix[0]), float(w.wcs.crpix[1])]
p += [float(w.wcs.crval[0]), float(w.wcs.crval[1])]
cd = np.array(w.wcs.cd, dtype=float)
p += [cd[0, 0], cd[0, 1], cd[1, 0], cd[1, 1]]
pv = np.zeros(pv_deg + 1, dtype=float)
# astropy stores PV in w.wcs.get_pv() / w.wcs.set_pv(); get_pv returns list of (i, m, value)
for i, m, val in w.wcs.get_pv():
if i == 2 and 0 <= m <= pv_deg:
pv[m] = float(val)
p += pv.tolist()
return np.array(p, dtype=float)
def _unpack_params_to_wcs(base: WCS, p: np.ndarray, pv_deg: int) -> WCS:
w = base.deepcopy()
w.wcs.crpix = [p[0], p[1]]
w.wcs.crval = [p[2], p[3]]
w.wcs.cd = np.array([[p[4], p[5]], [p[6], p[7]]], dtype=float)
pv_vals = p[8 : 8 + (pv_deg + 1)]
# Preserve any PV keywords for other axes (e.g. PV1_*)
pv_list = [(i, m, val) for (i, m, val) in base.wcs.get_pv() if i != 2]
pv_list += [(2, m, float(pv_vals[m])) for m in range(pv_deg + 1)]
w.wcs.set_pv(pv_list)
return w
def _normalize_zpn_pv1(w: WCS) -> WCS:
"""Normalize a ZPN WCS so that PV2_1 = 1, absorbing the scale into CD.
The ZPN projection R = Σ PV2_m·θ^m has a degeneracy: scaling all PV
by ``c`` and CD by ``c`` gives an identical pixel→sky mapping. This
function normalizes PV2_1 to 1 (the FITS standard convention), moving
any absorbed plate-scale factor into the CD matrix.
Parameters
----------
w : WCS
ZPN WCS (modified in place).
Returns
-------
WCS
The same object, normalized.
"""
pv_dict = {}
for i, m, val in w.wcs.get_pv():
if i == 2:
pv_dict[m] = val
pv1 = pv_dict.get(1, 1.0)
if not np.isfinite(pv1) or abs(pv1) < 1e-15 or abs(pv1 - 1.0) < 1e-12:
return w # already normalized or can't normalize
scale = 1.0 / pv1
# CD' = CD/pv1 so that intermediate coords scale together with PV.
cd = np.array(w.wcs.cd, dtype=float)
w.wcs.cd = cd / pv1
# Scale all PV2 coefficients
pv_list = [(i, m, val) for (i, m, val) in w.wcs.get_pv() if i != 2]
for m, val in pv_dict.items():
pv_list.append((2, m, float(val * scale)))
w.wcs.set_pv(pv_list)
return w
[docs]
def fit_zpn_wcs_from_points(
xy: np.ndarray,
sky: SkyCoord,
wcs_init: WCS,
pv_deg: int = 5,
fit_crpix: bool = True,
fit_crval: bool = True,
fit_cd: bool = True,
fit_pv: bool = True,
robust_loss: str = "soft_l1",
f_scale_arcsec: float = 2.0,
max_nfev: int = 200,
verbose: bool = False,
):
"""
Fit a ZPN WCS by optimizing WCS parameters against matched (x,y) <-> (ra,dec).
Parameters
----------
xy : (N,2) array
Pixel coordinates (0-based as in astropy WCS, i.e. origin=0).
sky : SkyCoord (N)
Reference sky positions.
wcs_init : astropy.wcs.WCS
Initial WCS; MUST already be ZPN (RA---ZPN/DEC--ZPN) or at least usable as base.
pv_deg : int
Degree for PV2_m coefficients to fit (PV2_0..PV2_pv_deg).
fit_* : bool
Toggle which parameter blocks to optimize.
robust_loss : str
Passed to scipy.optimize.least_squares(loss=...).
Good options: 'linear', 'soft_l1', 'huber', 'cauchy'.
f_scale_arcsec : float
Robust loss scale in arcsec.
max_nfev : int
Optimization iterations (SciPy).
verbose : bool or callable
Whether to show verbose messages during the run of the function or not.
May also be a print-like function.
Returns
-------
wcs_best : astropy.wcs.WCS
result : scipy OptimizeResult (or None if nothing was fitted)
Notes
-----
For stability, the solver runs in two stages when *fit_pv* is True:
it first fits CRPIX/CRVAL/CD with PV fixed, then fits all free
parameters (including PV) with conservative bounds to prevent
invalid projections.
"""
try:
from scipy.optimize import least_squares
except ImportError as e:
raise RuntimeError(
"fit_zpn_wcs_from_points requires SciPy (scipy.optimize.least_squares)"
) from e
# Simple wrapper around print for logging in verbose mode only
log = (verbose if callable(verbose) else print) if verbose else lambda *args, **kwargs: None
xy = np.asarray(xy, dtype=float)
if xy.ndim != 2 or xy.shape[1] != 2:
raise ValueError("xy must be (N,2)")
if len(sky) != xy.shape[0]:
raise ValueError("sky and xy must have the same length")
# Parameter packing assumes the CD convention; fold PC/CDELT into CD if needed
if not wcs_init.wcs.has_cd():
cd = np.array(wcs_init.pixel_scale_matrix, dtype=float)
wcs_init = wcs_init.deepcopy()
wcs_init.wcs.cdelt = np.ones(2)
if wcs_init.wcs.has_pc():
wcs_init.wcs.__delattr__("pc")
wcs_init.wcs.cd = cd
if fit_pv:
pv1_init = 0.0
for i, m, val in wcs_init.wcs.get_pv():
if i == 2 and m == 1:
pv1_init = float(val)
if not np.isfinite(pv1_init) or pv1_init == 0:
raise ValueError(
"Initial ZPN WCS has PV2_1 = 0 or undefined; the linear radial term "
"must be non-zero to fit PV coefficients "
"(e.g. initialize the WCS with tan_wcs_to_zpn())"
)
# Stable center for residual evaluation: spherical mean of the reference
# positions (robust to the lon=0/360 wrap, unlike a coordinate-wise median)
center = _spherical_mean(sky)
# The offset frame and reference positions in it are constant during the
# fit, so compute them once here rather than per residual evaluation
off = SkyOffsetFrame(origin=center)
sky_off = sky.transform_to(off)
coslat = np.cos(sky_off.lat.to_value(u.rad))
def _estimate_theta_max_deg(w: WCS) -> float | None:
if w.pixel_shape is None:
return None
nx, ny = w.pixel_shape
if nx is None or ny is None:
return None
# Prefer a pixel-scale estimate (robust to projection issues)
pixel_scale_deg = None
try:
pscales = w.proj_plane_pixel_scales()
if hasattr(pscales[0], "to_value"):
pixel_scale_deg = float(np.mean([s.to_value(u.deg) for s in pscales]))
else:
pixel_scale_deg = float(np.mean(pscales))
except Exception:
pixel_scale_deg = None
if pixel_scale_deg is not None and np.isfinite(pixel_scale_deg) and pixel_scale_deg > 0:
r_pix = 0.5 * np.hypot(nx, ny)
return max(float(pixel_scale_deg * r_pix), 0.01)
# Fallback: estimate from footprint if possible
try:
crpix = np.array(w.wcs.crpix, dtype=float)
corners = np.array(
[[0.0, 0.0], [nx - 1.0, 0.0], [0.0, ny - 1.0], [nx - 1.0, ny - 1.0]],
dtype=float,
)
ra_c, dec_c = w.all_pix2world(crpix[0], crpix[1], 0)
ra_k, dec_k = w.all_pix2world(corners[:, 0], corners[:, 1], 0)
if (
np.isfinite(ra_c)
and np.isfinite(dec_c)
and np.all(np.isfinite(ra_k))
and np.all(np.isfinite(dec_k))
):
center_c = SkyCoord(ra_c * u.deg, dec_c * u.deg, frame="icrs")
sky_k = SkyCoord(ra_k * u.deg, dec_k * u.deg, frame="icrs")
theta = float(np.max(center_c.separation(sky_k).to_value(u.deg)))
return max(theta, 0.01)
except Exception:
pass
return None
def _make_bounds(p0: np.ndarray, mask: np.ndarray, base: WCS, allow_pv: bool) -> tuple:
# Build bounds over the full parameter vector, then slice by mask.
# CRPIX/CRVAL/CD unbounded by default
lb = np.full_like(p0, -np.inf, dtype=float)
ub = np.full_like(p0, +np.inf, dtype=float)
if allow_pv and pv_deg >= 1:
# PV bounds: keep solution in a physically plausible neighborhood
pv0 = float(p0[8 + 0])
pv1 = float(p0[8 + 1]) # non-zero, guaranteed by the check above
theta_max_deg = _estimate_theta_max_deg(base)
pv1_abs = abs(pv1)
# PV2_0 ~ 0 (allow small drift; keep init inside bounds)
pv0_abs = max(1e-3, abs(pv0) * 2.0)
lb[8 + 0] = pv0 - pv0_abs
ub[8 + 0] = pv0 + pv0_abs
# PV2_1 (linear): allow moderate range, keep the initial sign
lo = min(max(0.1, pv1_abs * 0.2), pv1_abs)
hi = pv1_abs * 5.0
if pv1 > 0:
lb[8 + 1], ub[8 + 1] = lo, hi
else:
lb[8 + 1], ub[8 + 1] = -hi, -lo
# Higher-order PV terms: limit contribution at field edge
pv_frac = 0.3 # allow ~30% of linear term at the edge
for m in range(2, pv_deg + 1):
pv_m = float(p0[8 + m])
if theta_max_deg is not None and np.isfinite(theta_max_deg) and theta_max_deg > 0:
abs_bound = pv_frac * pv1_abs / (theta_max_deg ** (m - 1))
else:
abs_bound = max(abs(pv_m) * 5.0, 1e-6)
# Keep initial value inside bounds
abs_bound = max(abs_bound, abs(pv_m) * 2.0)
lb[8 + m] = pv_m - abs_bound
ub[8 + m] = pv_m + abs_bound
return lb[mask], ub[mask]
def _fit_with_mask(base: WCS, allow_pv: bool):
p0 = _pack_params(base, pv_deg=pv_deg)
# Build a mask over parameters to optionally freeze blocks
mask = np.ones_like(p0, dtype=bool)
# indices: 0..1 CRPIX, 2..3 CRVAL, 4..7 CD, 8.. PV
if not fit_crpix:
mask[0:2] = False
if not fit_crval:
mask[2:4] = False
if not fit_cd:
mask[4:8] = False
if not allow_pv:
mask[8:] = False
elif pv_deg >= 1:
# PV2_1 (linear scale) is degenerate with CD — let it float.
# Caller normalizes after fit via _normalize_zpn_pv1().
pass
if not np.any(mask):
return base, None
free0 = p0[mask]
lb, ub = _make_bounds(p0, mask, base, allow_pv=allow_pv)
def make_wcs_from_free(free: np.ndarray) -> WCS:
p = p0.copy()
p[mask] = free
return _unpack_params_to_wcs(base, p, pv_deg=pv_deg)
def fun(free: np.ndarray) -> np.ndarray:
w = make_wcs_from_free(free)
try:
res = _sky_residuals_arcsec(w, xy, sky.frame, off, sky_off, coslat)
except ValueError:
# wcslib may reject some in-bounds PV combinations as an
# invalid projection; report a very poor fit instead of
# crashing so the optimizer backs off
return np.full(2 * xy.shape[0], 1e6)
if not np.all(np.isfinite(res)):
res = np.nan_to_num(res, nan=1e6, posinf=1e6, neginf=-1e6)
return res
res = least_squares(
fun,
free0,
bounds=(lb, ub),
loss=robust_loss,
f_scale=float(f_scale_arcsec),
max_nfev=int(max_nfev),
x_scale="jac",
verbose=0,
)
w_best = make_wcs_from_free(res.x)
return w_best, res
def _log_residuals(label: str, res) -> None:
if res is None:
return
n = xy.shape[0]
r = np.hypot(res.fun[:n], res.fun[n:])
log(
f"ZPN fit {label}: median residual {np.median(r):.3g} arcsec "
f"over {n} points, {res.nfev} function evaluations"
)
# Two-stage fit: first solve CRPIX/CRVAL/CD with PV fixed,
# then allow PV with conservative bounds.
w_curr = wcs_init
res_last = None
if fit_pv and (fit_crpix or fit_crval or fit_cd):
w_curr, res_last = _fit_with_mask(w_curr, allow_pv=False)
_log_residuals("linear stage (PV fixed)", res_last)
w_curr, res_last = _fit_with_mask(w_curr, allow_pv=fit_pv)
_log_residuals("final stage", res_last)
# Normalize so PV2_1 = 1 (breaks the CD/PV degeneracy)
if fit_pv:
w_curr = _normalize_zpn_pv1(w_curr)
return w_curr, res_last
def _fit_zpn_sip(
wcs_zpn,
xy,
sky,
sip_degree=2,
pv_deg=5,
n_iter=15,
robust_loss="soft_l1",
f_scale_arcsec=2.0,
max_nfev=200,
verbose=False,
):
"""Add SIP distortion corrections on top of a ZPN WCS.
The SIP polynomials A(u,v), B(u,v) correct for non-radial distortions
(coma, astigmatism, etc.) that the ZPN radial model cannot capture.
Iterates between fitting SIP coefficients (pixel-space residuals) and
re-optimizing PV parameters with SIP applied.
Parameters
----------
wcs_zpn : WCS
ZPN WCS with PV parameters already fitted (no SIP).
xy : (N, 2) array
Pixel coordinates of matched sources (0-based).
sky : SkyCoord
Reference sky positions.
sip_degree : int
SIP polynomial order (default 2).
pv_deg : int
ZPN PV polynomial degree for re-fitting.
n_iter : int
Maximal number of PV+SIP alternation iterations.
robust_loss, f_scale_arcsec, max_nfev :
Passed to ``fit_zpn_wcs_from_points`` for PV re-fitting.
verbose : bool or callable
Whether to show verbose messages during the run of the function or not.
May also be a print-like function.
Returns
-------
wcs_sip : WCS
ZPN-SIP WCS with both PV and SIP coefficients.
"""
from astropy.wcs.wcs import Sip
# Simple wrapper around print for logging in verbose mode only
log = (verbose if callable(verbose) else print) if verbose else lambda *args, **kwargs: None
xy = np.asarray(xy, dtype=float)
w = wcs_zpn.deepcopy()
# SIP term indices: (p, q) with 2 <= p+q <= sip_degree
pq = []
for total in range(2, sip_degree + 1):
for q in range(total + 1):
pq.append((total - q, q))
crpix = np.array(w.wcs.crpix, dtype=float) # 1-based
prev_q90 = np.inf
for iteration in range(n_iter):
# Strip SIP to get the base ZPN-only WCS
w_nosip = w.deepcopy()
w_nosip.sip = None
if '-SIP' in w_nosip.wcs.ctype[0]:
w_nosip.wcs.ctype = [c.replace('-SIP', '') for c in w_nosip.wcs.ctype]
# Undistorted pixel coords: where ZPN projection places catalog stars.
# quiet=True prevents NoConvergence for points near the projection edge;
# non-finite results are excluded below.
x_cat, y_cat = w_nosip.all_world2pix(sky.ra.deg, sky.dec.deg, 0, quiet=True)
u_cat = x_cat - (crpix[0] - 1)
v_cat = y_cat - (crpix[1] - 1)
# Observed (distorted) pixel coords
u_obs = xy[:, 0] - (crpix[0] - 1)
v_obs = xy[:, 1] - (crpix[1] - 1)
# SIP convention: u_distorted = u_undistorted + A(u, v)
# u_obs = undistorted (raw pixel); u_cat = distorted (from ZPN inverse)
# So: A(u_obs, v_obs) = u_cat - u_obs
du = u_cat - u_obs
dv = v_cat - v_obs
# Filter outliers (use median absolute deviation)
dist = np.sqrt(du**2 + dv**2)
finite = np.isfinite(dist)
if np.sum(finite) < len(pq) + 5:
log(f"SIP iteration {iteration}: too few finite points ({np.sum(finite)}), stopping")
break
med_dist = np.median(dist[finite])
mad = np.median(np.abs(dist[finite] - med_dist))
good = finite & (dist < med_dist + 5 * max(mad, 0.5)) # generous 5-sigma clip
if np.sum(good) < len(pq) + 5:
log(f"SIP iteration {iteration}: too few points after clipping ({np.sum(good)}), stopping")
break
# SIP basis uses undistorted (raw pixel) coordinates
u_fit = u_obs[good]
v_fit = v_obs[good]
du_fit = du[good]
dv_fit = dv[good]
# Normalize coordinates to prevent ill-conditioning for high
# SIP degrees. Without this, u^5 ~ 2000^5 ~ 3e16 creates a
# design matrix with condition number > 1e14.
u_scale = max(np.abs(u_fit).max(), 1.0)
v_scale = max(np.abs(v_fit).max(), 1.0)
u_norm = u_fit / u_scale
v_norm = v_fit / v_scale
# Build SIP basis matrix in normalized coordinates
basis = np.column_stack([u_norm**p * v_norm**q for p, q in pq])
# Fit A and B coefficients via least squares (normalized)
ca_norm, _, _, _ = np.linalg.lstsq(basis, du_fit, rcond=None)
cb_norm, _, _, _ = np.linalg.lstsq(basis, dv_fit, rcond=None)
# Convert back to original (unnormalized) SIP coefficients
a_vals = np.zeros((sip_degree + 1, sip_degree + 1))
b_vals = np.zeros((sip_degree + 1, sip_degree + 1))
for k, (p, q) in enumerate(pq):
scale_pq = u_scale**p * v_scale**q
a_vals[p, q] = ca_norm[k] / scale_pq
b_vals[p, q] = cb_norm[k] / scale_pq
# Apply SIP to the WCS
if '-SIP' not in w.wcs.ctype[0]:
w.wcs.ctype = [c + '-SIP' for c in w.wcs.ctype]
w.sip = Sip(
a_vals,
b_vals,
np.zeros((sip_degree + 1, sip_degree + 1)),
np.zeros((sip_degree + 1, sip_degree + 1)),
crpix,
)
# Check convergence: stop when SIP corrections stabilize.
# Use 90th percentile (sensitive to outer-field + center) rather
# than median which converges before the center is corrected.
curr_q90 = np.percentile(dist[good], 90)
log(f"SIP iteration {iteration}: q90 distortion {curr_q90:.4g} pix over {np.sum(good)} points")
if iteration >= 2 and prev_q90 < np.inf:
rel_change = abs(curr_q90 - prev_q90) / max(prev_q90, 1e-6)
if rel_change < 0.005:
log(f"SIP converged after {iteration + 1} iterations")
break
prev_q90 = curr_q90
# Re-fit PV with SIP now applied (last iteration skip PV refit)
if iteration < n_iter - 1:
w, _ = fit_zpn_wcs_from_points(
xy,
sky,
w,
pv_deg=pv_deg,
robust_loss=robust_loss,
f_scale_arcsec=f_scale_arcsec,
max_nfev=max_nfev,
verbose=verbose,
)
crpix = np.array(w.wcs.crpix, dtype=float)
return w
[docs]
def tan_wcs_to_zpn(
w_tan: WCS,
pv_deg: int = 5,
n_samples: int = 256,
theta_max_deg: float | None = None,
) -> WCS:
"""
Convert a celestial TAN WCS into a ZPN WCS with PV2_m initialized to approximate TAN.
Notes
-----
TAN (gnomonic) radial law: r = tan(theta) [in radians]
In "degrees" units (common in FITS WCS plane coordinates), that's:
r_deg = tan(theta_rad) * (180/pi)
ZPN radial law: r_deg ≈ sum_{m=0..M} PV2_m * theta_deg^m
We set PV2_0 = 0 and fit PV2_1..PV2_M to approximate the TAN law
over theta in [0, theta_max_deg].
SIP and other pixel-space distortions of the input WCS are dropped
(they are not representable in the radial ZPN model); refit them on
top of the result if needed, e.g. with :func:`_fit_zpn_sip`.
Parameters
----------
w_tan : astropy.wcs.WCS
Input TAN WCS (2D celestial).
pv_deg : int
Highest PV degree to initialize (PV2_0..PV2_pv_deg).
5–9 is usually plenty; higher can get wiggly.
n_samples : int
Samples used for the polynomial fit.
theta_max_deg : float or None
Max angular radius (deg) over which to match TAN. If None,
estimated from image footprint corners using pixel_shape.
Returns
-------
w_zpn : astropy.wcs.WCS
A ZPN WCS with same CRVAL/CRPIX/CD and PV2_m initialized.
"""
ctype1, ctype2 = w_tan.wcs.ctype
if len(ctype1) < 8 or len(ctype2) < 8:
raise ValueError("Expected CTYPE like 'RA---TAN'/'DEC--TAN'.")
# Clean linear WCS (CRPIX/CRVAL/CD + metadata) with ZPN projection;
# strips SIP and PV parameters
w = _wcs_to_linear(w_tan, "ZPN")
# Estimate theta_max from footprint if not provided
if theta_max_deg is None:
if w.pixel_shape is None:
raise ValueError(
"w_tan.pixel_shape is None; provide theta_max_deg explicitly "
"or set w_tan.pixel_shape = (nx, ny)."
)
nx, ny = w.pixel_shape # (NAXIS1, NAXIS2)
crpix = np.array(w.wcs.crpix, dtype=float)
# Corners in pixel coordinates (origin=0 convention for astropy WCS)
corners = np.array(
[
[0.0, 0.0],
[nx - 1.0, 0.0],
[0.0, ny - 1.0],
[nx - 1.0, ny - 1.0],
],
dtype=float,
)
# Sky positions of corners and center under TAN WCS
ra_c, dec_c = w_tan.all_pix2world(crpix[0], crpix[1], 0)
center = SkyCoord(ra_c * u.deg, dec_c * u.deg, frame="icrs")
ra_k, dec_k = w_tan.all_pix2world(corners[:, 0], corners[:, 1], 0)
sky_k = SkyCoord(ra_k * u.deg, dec_k * u.deg, frame="icrs")
theta_max_deg = float(np.max(center.separation(sky_k).to_value(u.deg)))
# Safety floor
theta_max_deg = max(theta_max_deg, 0.01)
# Fit PV2_m so that ZPN radial r(theta) ~ TAN radial r(theta)
# Use degrees for theta and degrees for r on plane.
theta = np.linspace(0.0, theta_max_deg, n_samples, dtype=float)
theta_rad = np.deg2rad(theta)
r_tan_deg = np.tan(theta_rad) * (180.0 / np.pi)
# Build design matrix for m=1..pv_deg (PV2_0 fixed to 0)
# r ≈ sum c[m-1] * theta^m
A = np.vstack([theta**m for m in range(1, pv_deg + 1)]).T
# Mild weighting emphasizing the central region: stabilizes the fit
# against the steep TAN growth near theta_max; only used for initialization
wgt = 1.0 / (1.0 + (theta / (0.6 * theta_max_deg)) ** 2)
Aw = A * wgt[:, None]
bw = r_tan_deg * wgt
coeffs, *_ = np.linalg.lstsq(Aw, bw, rcond=None)
pv_list = [(2, 0, 0.0)] + [(2, m, float(coeffs[m - 1])) for m in range(1, pv_deg + 1)]
w.wcs.set_pv(pv_list)
return w
[docs]
def convert_wcs_projection(
wcs_input: WCS,
target_projection: str,
pv_deg: int = 5,
) -> WCS:
"""Convert a celestial WCS to a different projection type.
Parameters
----------
wcs_input : WCS
Input WCS (any celestial projection).
target_projection : str
Target projection code, e.g. ``'TAN'``, ``'ZPN'``, ``'STG'``,
``'ARC'``, ``'ZEA'``, ``'SIN'``, etc.
pv_deg : int, optional
ZPN PV polynomial degree (only used when *target_projection* is
``'ZPN'``). Default 5.
Returns
-------
WCS
New WCS with the target projection. For ZPN the PV coefficients
are initialised from the TAN radial law — an exact match for TAN
input, and a generic starting point (meant to be refined by a
subsequent fit) for other input projections. For other targets
CRPIX/CRVAL/CD are copied and CTYPE is replaced, so the mapping
away from the reference point changes; refit afterwards.
"""
target = target_projection.upper().strip()
# Detect current projection code (last 3 chars of CTYPE after the dash)
try:
cur_proj = wcs_input.wcs.ctype[0].split('-')[-1]
# Strip trailing SIP suffix if present
if cur_proj == 'SIP':
parts = wcs_input.wcs.ctype[0].replace('-SIP', '').split('-')
cur_proj = parts[-1]
except Exception:
cur_proj = ''
if cur_proj == target:
return wcs_input.deepcopy()
# ZPN needs special initialisation (PV coefficients)
if target == 'ZPN':
# tan_wcs_to_zpn works from TAN (with or without SIP)
# For non-TAN input, we first need a TAN-like WCS base
if 'TAN' in cur_proj:
return tan_wcs_to_zpn(wcs_input, pv_deg=pv_deg)
else:
# Build a minimal TAN WCS from the input's linear terms,
# then convert to ZPN
w_tan = _wcs_to_linear(wcs_input, 'TAN')
return tan_wcs_to_zpn(w_tan, pv_deg=pv_deg)
# For all other projections: copy linear WCS with new CTYPE
return _wcs_to_linear(wcs_input, target)
def _wcs_to_linear(wcs_input: WCS, proj_code: str) -> WCS:
"""Create a clean linear WCS with the same pointing but a new projection.
Copies CRPIX, CRVAL, CD (or derives CD from PC+CDELT), metadata, and
pixel_shape. Strips SIP distortion and PV parameters.
"""
w = WCS(naxis=2)
# CD matrix; pixel_scale_matrix correctly folds CDELT_i * PC_ij
# (row scaling) when the input uses the PC/CDELT convention
cd = np.array(wcs_input.pixel_scale_matrix, dtype=float)
w.wcs.crpix = np.array(wcs_input.wcs.crpix, dtype=float)
w.wcs.crval = np.array(wcs_input.wcs.crval, dtype=float)
w.wcs.cd = cd
# Build CTYPE: preserve axis names (e.g. 'RA---' / 'DEC--')
ctype1, ctype2 = wcs_input.wcs.ctype
# Strip any existing projection + SIP suffix
base1 = ctype1.replace('-SIP', '')[:5]
base2 = ctype2.replace('-SIP', '')[:5]
w.wcs.ctype = (base1 + proj_code, base2 + proj_code)
# Copy metadata
for attr in ('cunit', 'radesys', 'equinox'):
try:
setattr(w.wcs, attr, getattr(wcs_input.wcs, attr))
except Exception:
pass
for attr in ('lonpole', 'latpole'):
try:
val = getattr(wcs_input.wcs, attr)
if np.isfinite(val):
setattr(w.wcs, attr, float(val))
except Exception:
pass
if wcs_input.pixel_shape is not None:
w.pixel_shape = wcs_input.pixel_shape
return w
def _fit_tan_sip_robust(
xy,
world_coords,
proj_point="center",
projection=None,
sip_degree=2,
robust_loss="soft_l1",
f_scale=None,
verbose=False,
):
"""Fit TAN+SIP WCS with robust loss function.
Replicates astropy's ``fit_wcs_from_points`` two-stage fitting
(linear WCS, then joint CD + CRPIX + SIP) but uses a robust loss
in the linear stage and a two-pass approach for the SIP stage:
first L2 to capture the distortion pattern, then robust re-fit
with data-driven ``f_scale`` to suppress outliers.
Parameters
----------
xy : tuple of arrays ``(x, y)``
Pixel coordinates.
world_coords : `~astropy.coordinates.SkyCoord`
Reference sky positions.
proj_point : str or SkyCoord
Projection center ('center' or explicit).
projection : WCS or None
Template WCS; used as initial guess for CD matrix.
sip_degree : int
SIP polynomial degree.
robust_loss : str
Loss function for `scipy.optimize.least_squares`.
f_scale : float or None
Soft margin for robust loss in the SIP stage. If None
(default), estimated adaptively from the L2 SIP residuals.
The linear stage always uses a data-driven scale.
verbose : bool or callable
Whether to show verbose messages during the run of the function or not.
May also be a print-like function.
Returns
-------
wcs : `~astropy.wcs.WCS`
"""
from scipy.optimize import least_squares
from astropy.wcs.utils import celestial_frame_to_wcs
from astropy.wcs.wcs import Sip
try:
# Private astropy helpers (stable since astropy 3.x, but not public API)
from astropy.wcs.utils import _linear_wcs_fit, _sip_fit
except ImportError as e:
raise RuntimeError(
"astropy internals (_linear_wcs_fit/_sip_fit) are not available in "
"this astropy version; _fit_tan_sip_robust needs updating"
) from e
# Simple wrapper around print for logging in verbose mode only
log = (verbose if callable(verbose) else print) if verbose else lambda *args, **kwargs: None
xp, yp = xy
try:
lon, lat = world_coords.data.lon.deg, world_coords.data.lat.deg
except AttributeError:
unit_sph = world_coords.unit_spherical
lon, lat = unit_sph.lon.deg, unit_sph.lat.deg
use_center_as_proj_point = str(proj_point) == "center"
# Build WCS template
if isinstance(projection, str) or projection is None:
proj_code = projection if isinstance(projection, str) else "TAN"
wcs = celestial_frame_to_wcs(frame=world_coords.frame, projection=proj_code)
else:
wcs = projection.deepcopy()
wcs.sip = None
if wcs.wcs.has_pc():
# pixel_scale_matrix folds CDELT_i * PC_ij with correct row scaling
wcs.wcs.cd = wcs.pixel_scale_matrix
wcs.wcs.cdelt = (1.0, 1.0)
wcs.wcs.__delattr__("pc")
xpmin, xpmax = xp.min(), xp.max()
ypmin, ypmax = yp.min(), yp.max()
wcs.pixel_shape = (
1 if xpmax <= 0.0 else int(np.ceil(xpmax)),
1 if ypmax <= 0.0 else int(np.ceil(ypmax)),
)
if use_center_as_proj_point:
sc1 = SkyCoord(lon.min() * u.deg, lat.max() * u.deg)
sc2 = SkyCoord(lon.max() * u.deg, lat.min() * u.deg)
pa = sc1.position_angle(sc2)
sep = sc1.separation(sc2)
midpoint_sc = sc1.directional_offset_by(pa, sep / 2)
wcs.wcs.crval = (midpoint_sc.data.lon.deg, midpoint_sc.data.lat.deg)
wcs.wcs.crpix = ((xpmax + xpmin) / 2.0, (ypmax + ypmin) / 2.0)
else:
proj_point = proj_point.transform_to(world_coords.frame)
wcs.wcs.crval = (proj_point.data.lon.deg, proj_point.data.lat.deg)
close = lambda l, p: p[np.argmin(np.abs(l))]
wcs.wcs.crpix = (
close(lon - wcs.wcs.crval[0], xp + 1),
close(lat - wcs.wcs.crval[1], yp + 1),
)
if xpmin == xpmax:
xpmin, xpmax = xpmin - 0.5, xpmax + 0.5
if ypmin == ypmax:
ypmin, ypmax = ypmin - 0.5, ypmax + 0.5
# --- Stage 1: linear WCS (CD + CRPIX) ---
# Use robust loss here: linear residuals include distortion as systematic
# pattern that looks like outliers; robust loss prevents them from biasing
# the linear fit.
p0_lin = np.concatenate([wcs.wcs.cd.flatten(), wcs.wcs.crpix.flatten()])
lin_resids = _linear_wcs_fit(p0_lin, lon, lat, xp, yp, wcs)
lin_f_scale = max(float(np.median(np.abs(lin_resids)) * 3), 1.0 / 3600)
fit = least_squares(
_linear_wcs_fit,
p0_lin,
args=(lon, lat, xp, yp, wcs),
bounds=[
[-np.inf, -np.inf, -np.inf, -np.inf, xpmin + 1, ypmin + 1],
[np.inf, np.inf, np.inf, np.inf, xpmax + 1, ypmax + 1],
],
loss=robust_loss,
f_scale=lin_f_scale,
method="trf",
)
wcs.wcs.crpix = np.array(fit.x[4:6])
wcs.wcs.cd = np.array(fit.x[0:4].reshape((2, 2)))
log(
f"TAN-SIP linear stage: median residual "
f"{np.median(np.abs(fit.fun)) * 3600:.3g} arcsec"
)
# --- Stage 2: joint CD + CRPIX + SIP ---
if "-SIP" not in wcs.wcs.ctype[0]:
wcs.wcs.ctype = [x + "-SIP" for x in wcs.wcs.ctype]
coef_names = [
f"{i}_{j}"
for i in range(sip_degree + 1)
for j in range(sip_degree + 1)
if (i + j) < (sip_degree + 1) and (i + j) > 1
]
sip_bounds = (
[xpmin + 1, ypmin + 1] + [-np.inf] * (4 + 2 * len(coef_names)),
[xpmax + 1, ypmax + 1] + [np.inf] * (4 + 2 * len(coef_names)),
)
# Parameter layout of _sip_fit: [CRPIX(2), CD(4), A_coeffs, B_coeffs]
sip_args = (lon, lat, xp, yp, wcs, sip_degree, coef_names)
# Pass 1: L2 fit to capture the full distortion pattern (SIP starts at 0)
p0_sip = np.concatenate(
(
np.array(wcs.wcs.crpix),
wcs.wcs.cd.flatten(),
np.zeros(2 * len(coef_names)),
)
)
# Compute parameter scales for the optimizer.
# SIP coefficients A_{p,q} multiply (u-crpix)^p * (v-crpix)^q where
# pixel offsets can be ~U pixels. Expected coefficient magnitude is
# ~1/U^(p+q), which spans many orders of magnitude for high SIP orders.
# Without proper scaling the optimizer cannot find meaningful SIP≥4
# coefficients on wide-field images.
U = max(
abs(xpmax - wcs.wcs.crpix[0]),
abs(xpmin - wcs.wcs.crpix[0]),
abs(ypmax - wcs.wcs.crpix[1]),
abs(ypmin - wcs.wcs.crpix[1]),
1.0,
)
sip_x_scale = np.ones(len(p0_sip))
# CRPIX: O(pixels) → scale 1
# CD: O(deg/pixel) → use current magnitude
cd_scale = max(np.abs(wcs.wcs.cd).max(), 1e-10)
sip_x_scale[2:6] = cd_scale
# SIP coefficients: O(1/U^(p+q))
for k, coef_name in enumerate(coef_names):
p_deg, q_deg = int(coef_name[0]), int(coef_name[2])
order = p_deg + q_deg
scale = 1.0 / U**order
sip_x_scale[6 + k] = scale
sip_x_scale[6 + len(coef_names) + k] = scale
fit = least_squares(
_sip_fit,
p0_sip,
args=sip_args,
bounds=sip_bounds,
x_scale=sip_x_scale,
)
# Pass 2: robust re-fit starting from L2 solution
# f_scale estimated from L2 residuals so distortion signal is inlier
sip_resids = fit.fun
if f_scale is None:
sip_f_scale = max(float(np.median(np.abs(sip_resids)) * 3), 1e-7)
else:
sip_f_scale = float(f_scale)
fit = least_squares(
_sip_fit,
fit.x, # warm-start from L2 solution
args=sip_args,
bounds=sip_bounds,
loss=robust_loss,
f_scale=sip_f_scale,
method="trf",
x_scale=sip_x_scale,
)
log(
f"TAN-SIP robust stage: median residual "
f"{np.median(np.abs(fit.fun)) * 3600:.3g} arcsec, f_scale {sip_f_scale * 3600:.3g} arcsec"
)
coef_fit = (
list(fit.x[6 : 6 + len(coef_names)]),
list(fit.x[6 + len(coef_names) :]),
)
wcs.wcs.cd = fit.x[2:6].reshape((2, 2))
wcs.wcs.crpix = fit.x[0:2]
a_vals = np.zeros((sip_degree + 1, sip_degree + 1))
b_vals = np.zeros((sip_degree + 1, sip_degree + 1))
for coef_name in coef_names:
a_vals[int(coef_name[0])][int(coef_name[2])] = coef_fit[0].pop(0)
b_vals[int(coef_name[0])][int(coef_name[2])] = coef_fit[1].pop(0)
wcs.sip = Sip(
a_vals,
b_vals,
np.zeros((sip_degree + 1, sip_degree + 1)),
np.zeros((sip_degree + 1, sip_degree + 1)),
wcs.wcs.crpix,
)
return wcs
[docs]
def fit_wcs_from_points(
xy,
world_coords,
proj_point="center",
projection=None,
sip_degree=None,
pv_deg=5,
verbose=False,
):
"""Drop-in wrapper around :func:`astropy.wcs.utils.fit_wcs_from_points`
that also handles **ZPN** projection (which astropy does not natively fit)
and **ZPN-SIP** (ZPN radial distortion plus SIP polynomial corrections).
Parameters
----------
xy : tuple of arrays ``(x, y)`` or ``(2, N)`` array
Pixel coordinates (same convention as the astropy function).
world_coords : `~astropy.coordinates.SkyCoord`
Reference sky positions.
proj_point : str, optional
Passed through to astropy for non-ZPN projections. Ignored for ZPN,
where the projection center comes from the template WCS.
projection : `~astropy.wcs.WCS` or other, optional
Projection template. If this is a WCS with ``RA---ZPN / DEC--ZPN``
CTYPEs, the ZPN fitter is used instead of astropy's.
sip_degree : int or None, optional
SIP polynomial degree. For TAN projections, controls SIP distortion
order. For ZPN projections, if > 0, SIP corrections are fitted on
top of ZPN PV parameters to capture non-radial distortions. For
other projections it is ignored (SIP is not standard there).
pv_deg : int, optional
ZPN PV polynomial degree (``PV2_0 … PV2_pv_deg``). Default 5.
verbose : bool or callable, optional
Whether to show verbose messages during the run of the function or not.
May also be a print-like function.
Returns
-------
wcs : `~astropy.wcs.WCS`
Fitted WCS (same return type as the astropy function).
"""
from astropy.wcs.utils import fit_wcs_from_points as _astropy_fit
# Simple wrapper around print for logging in verbose mode only
log = (verbose if callable(verbose) else print) if verbose else lambda *args, **kwargs: None
# Detect ZPN projection
is_zpn = False
if isinstance(projection, WCS):
try:
is_zpn = "ZPN" in projection.wcs.ctype[0]
except Exception:
pass
if is_zpn:
# Convert xy to (N, 2) array expected by fit_zpn_wcs_from_points
if isinstance(xy, (list, tuple)) and len(xy) == 2:
# (x_array, y_array) form
xy_arr = np.column_stack(
[np.asarray(xy[0], dtype=float), np.asarray(xy[1], dtype=float)]
)
else:
xy_arr = np.asarray(xy, dtype=float)
if xy_arr.ndim == 2 and xy_arr.shape[0] == 2 and xy_arr.shape[1] != 2:
# (2, N) -> (N, 2)
xy_arr = xy_arr.T
zpn_deg = int(pv_deg)
# First fit ZPN PV parameters (radial distortion)
log(f"Fitting ZPN WCS with pv_deg={zpn_deg} using {len(xy_arr)} points")
wcs_best, _result = fit_zpn_wcs_from_points(
xy_arr, world_coords, wcs_init=projection, pv_deg=zpn_deg, verbose=verbose
)
# Then fit SIP corrections for non-radial distortions
if sip_degree is not None and int(sip_degree) > 0:
log(f"Fitting SIP degree {int(sip_degree)} corrections on top of ZPN")
wcs_best = _fit_zpn_sip(
wcs_best,
xy_arr,
world_coords,
sip_degree=int(sip_degree),
pv_deg=zpn_deg,
verbose=verbose,
)
return wcs_best
# ---------- Non-ZPN ----------
# SIP only makes sense for TAN-based projections
effective_sip = sip_degree
if isinstance(projection, WCS):
try:
if "TAN" not in projection.wcs.ctype[0]:
effective_sip = None
except Exception:
pass
if effective_sip is not None and effective_sip > 0:
log(f"Fitting TAN-SIP WCS with sip_degree={int(effective_sip)}")
return _fit_tan_sip_robust(
xy,
world_coords,
proj_point=proj_point,
projection=projection,
sip_degree=int(effective_sip),
verbose=verbose,
)
if effective_sip != sip_degree:
log("SIP is not supported for this projection, ignoring sip_degree")
log("Fitting WCS using astropy fit_wcs_from_points")
return _astropy_fit(
xy,
world_coords,
proj_point=proj_point,
projection=projection,
)