reprojax 0.2.0__py3-none-any.whl

reprojax/__init__.py

@@ -0,0 +1,15 @@
+from .adaptive import reproject_adaptive, reproject_adaptive_core
+from .core import PixelTransform
+from .interp import reproject_interp, reproject_interp_batched, reproject_interp_core
+from .transform import build_transform, compose_transform
+
+__all__ = [
+    "PixelTransform",
+    "build_transform",
+    "compose_transform",
+    "reproject_adaptive",
+    "reproject_adaptive_core",
+    "reproject_interp",
+    "reproject_interp_batched",
+    "reproject_interp_core",
+]

reprojax/_astropy_wcs.py

@@ -0,0 +1,74 @@
+"""Host-side extraction of mapprojax WCS parameters and frame rotations from astropy.
+
+Everything in this module runs eagerly in NumPy/astropy, at transform-build
+time — never inside the jit-compiled :class:`~reprojax.core.PixelTransform`
+itself.
+"""
+
+import numpy as np
+from astropy.coordinates import SkyCoord
+from astropy.wcs.utils import wcs_to_celestial_frame
+from mapprojax.sip import Sip
+
+#: FITS projection code -> mapprojax provider key (see transform.py).
+SUPPORTED_PROJECTIONS = ("TAN", "SIN")
+
+
+def projection_code(wcs):
+    """The 3-letter FITS projection code for an astropy WCS's first axis.
+
+    Per the WCS papers, CTYPEia's first 8 characters are always
+    ``"<axis>-<code>"`` with the code in the last 3 characters of that
+    8-character field — true whether or not a further distortion suffix
+    (``-SIP``, ``-TPV``, ...) follows.
+    """
+    ctype = wcs.wcs.ctype[0]
+    code = ctype[:8][-3:].strip()
+    if code not in SUPPORTED_PROJECTIONS:
+        raise NotImplementedError(
+            f"Unsupported projection {code!r} (from CTYPE {ctype!r}); "
+            f"reprojax currently supports: {SUPPORTED_PROJECTIONS}"
+        )
+    return code
+
+
+def wcs_params_from_astropy(wcs):
+    """Extract ``(crpix, cd, crval, sip)`` from an astropy WCS.
+
+    ``crpix`` is converted to mapprojax's 0-based convention (astropy's
+    ``wcs.wcs.crpix`` is 1-based, matching the FITS ``CRPIXi`` keywords).
+    ``cd`` is the CD-matrix-equivalent linear transform regardless of
+    whether the header used ``CDi_j`` or ``PCi_j``/``CDELTi`` keywords.
+    """
+    crpix = (np.asarray(wcs.wcs.crpix, dtype=float) - 1.0).tolist()
+    crval = np.asarray(wcs.wcs.crval, dtype=float).tolist()
+    cd = np.asarray(wcs.pixel_scale_matrix, dtype=float).tolist()
+
+    sip = None
+    if wcs.sip is not None:
+        sip = Sip(a=wcs.sip.a, b=wcs.sip.b, ap=wcs.sip.ap, bp=wcs.sip.bp)
+
+    return crpix, cd, crval, sip
+
+
+def frame_rotation_matrix(wcs_from, wcs_to):
+    """3x3 matrix rotating Cartesian unit vectors from ``wcs_from``'s celestial
+    frame to ``wcs_to``'s.
+
+    Returns the identity when the two WCSes share the same celestial frame
+    (the common case, and the only case that avoids an astropy frame
+    transform).
+    """
+    frame_from = wcs_to_celestial_frame(wcs_from)
+    frame_to = wcs_to_celestial_frame(wcs_to)
+
+    if frame_from.name == frame_to.name:
+        return np.eye(3)
+
+    basis_lon = [0.0, 90.0, 0.0]
+    basis_lat = [0.0, 0.0, 90.0]
+    basis = SkyCoord(basis_lon, basis_lat, unit="deg", frame=frame_from)
+    cart = basis.transform_to(frame_to).cartesian
+
+    # Transformed basis vectors as columns: matrix @ v_from == v_to.
+    return np.stack([cart.x.value, cart.y.value, cart.z.value], axis=0)

reprojax/adaptive.py

@@ -0,0 +1,318 @@
+"""Adaptive anti-aliased reprojection: DeForest (2004), ported from mapprojax.
+
+The core (:func:`reproject_adaptive_core`) consumes a
+:class:`~reprojax.core.PixelTransform` (no mapprojax-class coupling) and
+returns ``(result, footprint)`` matching :mod:`reprojax.interp`'s
+conventions. The wrapper (:func:`reproject_adaptive`) does the astropy-facing
+work, matching ``reproject.reproject_adaptive``'s call/return shape.
+
+This is a from-scratch reimplementation of mapprojax's ``adaptive_reproject``
+against ``reproject_adaptive``'s actual reference algorithm
+(``reproject/adaptive/deforest.pyx``), not a direct port: validating against
+the real algorithm surfaced several places where mapprojax's original
+approximated reproject's behavior rather than matching it, corrected here
+(see the "Deliberate differences from reproject_adaptive" note below).
+
+Deliberate differences from reproject_adaptive
+-----------------------------------------------
+- **Jacobian**: computed exactly via JAX autodiff (``jax.jacfwd`` through the
+  transform), not reproject's finite-difference estimate between
+  neighboring output pixels. Empirically these agree to ~1e-10 on every
+  scenario tested (synthetic WCS pairs at ~1x and ~4x scale, with SIP, and
+  the real SPHEREx pair) — exact autodiff is at least as accurate as a
+  finite-difference approximation of the same quantity, and is the natural
+  choice given reprojax's differentiable transform.
+- **Search window is a fixed maximum size** (``max_window_radius``), not
+  reproject's per-pixel dynamically-sized region: JAX/vmap requires a
+  static shape across every output pixel. Cells outside the true dynamic
+  window (computed per-pixel from the local Jacobian, same as reproject)
+  get zero weight, matching reproject's hard truncation, but
+  ``max_window_radius`` must be set large enough to contain that dynamic
+  window for every pixel in the image — large enough for the most
+  aggressive downsampling present, not just a typical pixel. Too small a
+  value silently truncates the kernel (no error is raised); this is the
+  cost of vmap's uniform per-pixel work requirement.
+- **Only the algorithm's default parameters are supported**:
+  ``kernel_width``, ``sample_region_width``, and ``kernel`` (Gaussian or
+  Hann) are the only tunables. reproject_adaptive's ``boundary_mode`` and
+  ``bad_value_mode`` are fixed at their defaults (``"strict"`` for both —
+  any output pixel whose window extends beyond the input array, or whose
+  window contains a NaN/Inf input value, is set to NaN as a whole) with no
+  other mode supported; ``conserve_flux``, ``despike_jacobian``,
+  ``center_jacobian``, and ``x_cyclic``/``y_cyclic`` are not implemented.
+
+Validation
+----------
+Against ``reproject_adaptive`` with matching default parameters, on
+non-degenerate (non-integer CRPIX/CRVAL/rotation) synthetic WCS pairs:
+identical footprints and value agreement to ~1e-10, at both ~1x and ~4x
+scale, with SIP, and with real masked (NaN) input pixels. See
+tests/test_adaptive_synthetic.py and tests/test_real_spherex.py.
+
+Two things were diagnosed and are *not* papered over with a looser
+tolerance:
+
+1. Exact-integer-tie sensitivity in ``boundary_mode="strict"``'s window
+   check: reproject_adaptive's coordinate pipeline can return a
+   mathematically-exact integer position for special-case inputs (e.g. an
+   output WCS identical to the input WCS), while reprojax's composed
+   transform (via trigonometric round-trip through mapprojax) generally
+   lands ~1e-10 away from that exact integer even when the true answer is
+   exact. Since the window's ceil/floor bounds are then computed from that
+   position, a coincidental floating-point nudge across an exact integer
+   can shift a search window's boundary by one pixel, occasionally
+   disagreeing right at the input array's edge. This is specific to
+   artificially "round" test parameters (integer CRPIX, exactly-matching
+   in/out WCS) rather than real data — it was not observed on the real
+   SPHEREx pair (see tests/test_real_spherex.py), and the synthetic test
+   suite uses deliberately non-round parameters to avoid manufacturing it.
+2. A genuine bug caught during this validation, not present in the shipped
+   code: masking cells outside the per-pixel dynamic search window to zero
+   *weight* is not sufficient on its own — the fixed-size window array can
+   still contain a NaN there (e.g. a masked input pixel that happens to
+   fall within ``max_window_radius`` of an output pixel's mapped position,
+   but outside that pixel's actual, smaller dynamic window), and
+   `0 * nan = nan` under IEEE 754 poisons the weighted sum regardless of
+   the zero weight. Fixed by substituting a safe value outside the dynamic
+   window before the multiply — the same "never multiply an excluded
+   value by NaN" lesson as reprojax.interp's map_coordinates handling.
+"""
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+from .interp import _as_wcs, _output_grid_dtype, _shape_from_header, _to_native_byte_order
+from .transform import build_transform
+
+SUPPORTED_KERNELS = ("gaussian", "hann")
+
+
+@partial(
+    jax.jit,
+    static_argnames=("transform", "kernel", "kernel_width", "sample_region_width", "max_window_radius", "block_size"),
+)
+def reproject_adaptive_core(
+    data,
+    x_out,
+    y_out,
+    transform,
+    kernel="gaussian",
+    kernel_width=1.3,
+    sample_region_width=4.0,
+    max_window_radius=10,
+    block_size=None,
+):
+    """Reproject ``data`` onto the pixel grid given by ``x_out``/``y_out``
+    using the DeForest (2004) adaptive, anti-aliased algorithm.
+
+    Parameters
+    ----------
+    data : array, shape (ny_in, nx_in)
+        Input image.
+    x_out, y_out : arrays, any common broadcastable shape
+        Output-image pixel coordinates to fill (0-based, pixel-center
+        convention — see :mod:`reprojax.core`).
+    transform : PixelTransform
+        Maps ``(x_out, y_out) -> (x_in, y_in)``.
+    kernel : str
+        ``"gaussian"`` (default, matches reproject_adaptive's default) or
+        ``"hann"``.
+    kernel_width : float
+        Gaussian kernel width in pixels (unused for Hann).
+    sample_region_width : float
+        Width, in output pixels, of the sampling region for the Gaussian
+        kernel (unused for Hann, whose extent is determined by the
+        transform's local Jacobian alone).
+    max_window_radius : int
+        Fixed half-width (in input pixels) of the search window; must be
+        large enough to contain the true, dynamically-sized search window
+        for every output pixel (larger for more aggressive downsampling).
+        See the module docstring's note on this being a JAX/vmap-driven
+        departure from reproject_adaptive's per-pixel dynamic sizing.
+    block_size : int, optional
+        Process this many output pixels per `jax.lax.scan` step instead of
+        vectorizing the whole output grid at once, to bound memory use.
+
+    Returns
+    -------
+    result : array, same shape as x_out
+        Reprojected values, dtype matching ``data``. NaN where invalid.
+    footprint : array, same shape as x_out
+        1.0 where the value is valid, 0.0 where not.
+    """
+    if kernel not in SUPPORTED_KERNELS:
+        raise ValueError(f"kernel must be one of {SUPPORTED_KERNELS}, got {kernel!r}")
+
+    ny_in, nx_in = data.shape
+    radius = max_window_radius
+    win_size = 2 * radius + 1
+    padded = jnp.pad(data, radius, mode="constant", constant_values=0.0)
+
+    def process_pixel(x0_out, y0_out):
+        u0, v0 = transform(x0_out, y0_out)
+
+        def t(xy):
+            u, v = transform(xy[0], xy[1])
+            return jnp.stack([u, v])
+
+        J = jax.jacfwd(t)(jnp.array([x0_out, y0_out]))  # d(input)/d(output), (2, 2)
+
+        U, S, Vt = jnp.linalg.svd(J)
+        s_clamped = jnp.maximum(1.0, S)  # never sharpen beyond native resolution
+        j_inv = (Vt.T * (1.0 / s_clamped)) @ U.T  # d(output)/d(input), using the clamped Jacobian
+
+        if kernel == "gaussian":
+            half_w = half_h = sample_region_width * jnp.max(s_clamped) / 2.0
+        else:
+            ji_padded = (U * s_clamped) @ Vt
+            corners = jnp.array([[1.0, 1.0], [1.0, -1.0], [-1.0, 1.0], [-1.0, -1.0]])
+            transformed_corners = corners @ ji_padded.T
+            half_w = jnp.max(jnp.abs(transformed_corners[:, 0]))
+            half_h = jnp.max(jnp.abs(transformed_corners[:, 1]))
+
+        bottom = jnp.ceil(v0 - half_h)
+        top = jnp.floor(v0 + half_h)
+        left = jnp.ceil(u0 - half_w)
+        right = jnp.floor(u0 + half_w)
+
+        u0_int = jnp.rint(u0).astype(int)
+        v0_int = jnp.rint(v0).astype(int)
+        window = jax.lax.dynamic_slice(padded, (v0_int, u0_int), (win_size, win_size))
+
+        win_v_idx, win_u_idx = jnp.meshgrid(jnp.arange(win_size), jnp.arange(win_size), indexing="ij")
+        global_v = v0_int - radius + win_v_idx
+        global_u = u0_int - radius + win_u_idx
+        global_v_f = global_v.astype(v0.dtype)
+        global_u_f = global_u.astype(u0.dtype)
+
+        in_dynamic_window = (global_v_f >= bottom) & (global_v_f <= top) & (global_u_f >= left) & (global_u_f <= right)
+
+        dv = global_v_f - v0
+        du = global_u_f - u0
+        tx = j_inv[0, 0] * du + j_inv[0, 1] * dv
+        ty = j_inv[1, 0] * du + j_inv[1, 1] * dv
+
+        if kernel == "gaussian":
+            raw_weight = jnp.exp(-(tx**2 + ty**2) / kernel_width**2 * 2.0)
+        else:
+            hann_mask = (jnp.abs(tx) < 1.0) & (jnp.abs(ty) < 1.0)
+            raw_weight = jnp.where(hann_mask, (jnp.cos(jnp.pi * tx) + 1.0) * (jnp.cos(jnp.pi * ty) + 1.0), 0.0)
+
+        weight = jnp.where(in_dynamic_window, raw_weight, 0.0)
+
+        # Cells outside the dynamic window carry zero weight, but the
+        # fixed-size window can still contain a NaN there; `0 * nan = nan`
+        # would poison the sum regardless of weight (see module docstring).
+        # A NaN genuinely inside the dynamic window must still poison the
+        # result (strict bad_value_mode) -- only substitute outside it.
+        window_safe = jnp.where(in_dynamic_window, window, 0.0)
+
+        weight_sum = jnp.sum(weight)
+        target = jnp.sum(weight * window_safe)
+
+        in_array_bounds = (global_v >= 0) & (global_v < ny_in) & (global_u >= 0) & (global_u < nx_in)
+        # strict boundary_mode: any dynamic-window cell outside the input
+        # array invalidates the whole output pixel.
+        window_out_of_bounds = jnp.any(in_dynamic_window & ~in_array_bounds)
+        invalid_transform = jnp.isnan(u0) | jnp.isnan(v0) | jnp.any(jnp.isnan(J))
+
+        result = target / weight_sum
+        return jnp.where(window_out_of_bounds | invalid_transform, jnp.nan, result)
+
+    x_flat = x_out.ravel()
+    y_flat = y_out.ravel()
+
+    if block_size is None:
+        flat_result = jax.vmap(process_pixel)(x_flat, y_flat)
+    else:
+        n = x_flat.size
+        pad = (-n) % block_size
+        x_padded = jnp.pad(x_flat, (0, pad))
+        y_padded = jnp.pad(y_flat, (0, pad))
+        x_blocks = x_padded.reshape(-1, block_size)
+        y_blocks = y_padded.reshape(-1, block_size)
+
+        def scan_body(carry, block):
+            bx, by = block
+            return carry, jax.vmap(process_pixel)(bx, by)
+
+        _, blocks_out = jax.lax.scan(scan_body, None, (x_blocks, y_blocks))
+        flat_result = blocks_out.ravel()
+        if pad:
+            flat_result = flat_result[:-pad]
+
+    result = flat_result.reshape(x_out.shape).astype(data.dtype)
+    footprint = (~jnp.isnan(result)).astype(jnp.float32)
+    return result, footprint
+
+
+def reproject_adaptive(
+    input_data,
+    output_projection,
+    shape_out=None,
+    kernel="gaussian",
+    kernel_width=1.3,
+    sample_region_width=4.0,
+    max_window_radius=10,
+    block_size=None,
+    matrix_precision="f64",
+):
+    """Reproject an ``(array, WCS-or-header)`` pair onto ``output_projection``
+    using the DeForest (2004) adaptive, anti-aliased algorithm.
+
+    Matches ``reproject.reproject_adaptive``'s call and return shape for the
+    parameters this port supports — see the module docstring for the
+    (default-matching) scope this covers and what's out of scope.
+
+    Parameters
+    ----------
+    input_data : tuple(array, astropy.wcs.WCS or astropy.io.fits.Header)
+    output_projection : astropy.wcs.WCS or astropy.io.fits.Header
+    shape_out : tuple(int, int), optional
+        Required unless it can be inferred from ``output_projection``
+        (a Header's NAXIS1/NAXIS2, or a WCS with ``array_shape`` set).
+    kernel, kernel_width, sample_region_width, max_window_radius, block_size
+        See :func:`reproject_adaptive_core`.
+    matrix_precision : {"f64", "f32"}
+        See :func:`~reprojax.transform.build_transform`. ``"f32"`` needs
+        no ``jax_enable_x64`` and is faster on GPUs where float64
+        throughput is much lower than float32's, at reduced accuracy
+        (~1e-4 px for typical overlapping-field reprojections; see that
+        function's docstring for the detailed tradeoff).
+
+    Returns
+    -------
+    array, footprint
+    """
+    data, wcs_or_header_in = input_data
+    wcs_in = _as_wcs(wcs_or_header_in)
+    wcs_out = _as_wcs(output_projection)
+
+    if shape_out is None:
+        shape_out = _shape_from_header(output_projection)
+        if shape_out is None:
+            raise ValueError(
+                "shape_out could not be inferred from output_projection; pass it explicitly"
+            )
+    ny, nx = shape_out
+
+    grid_dtype = _output_grid_dtype(matrix_precision)
+    y_out, x_out = jnp.meshgrid(
+        jnp.arange(ny, dtype=grid_dtype), jnp.arange(nx, dtype=grid_dtype), indexing="ij"
+    )
+
+    transform = build_transform(wcs_in, wcs_out, matrix_precision=matrix_precision)
+    data = jnp.asarray(_to_native_byte_order(data))
+    return reproject_adaptive_core(
+        data,
+        x_out,
+        y_out,
+        transform,
+        kernel=kernel,
+        kernel_width=kernel_width,
+        sample_region_width=sample_region_width,
+        max_window_radius=max_window_radius,
+        block_size=block_size,
+    )

reprojax/core.py

@@ -0,0 +1,55 @@
+"""The ``PixelTransform`` contract.
+
+This is reprojax's central seam. Reprojection algorithms are written against
+this contract, not against any particular WCS library — mapprojax is the
+default provider (see :mod:`reprojax.transform`), and other providers (e.g.
+an astropy-WCS ``pure_callback`` fallback for projections mapprojax doesn't
+cover) can be added later behind the same interface.
+
+Contract
+--------
+A ``PixelTransform`` is a callable::
+
+    x_in, y_in = transform(x_out, y_out)
+
+mapping output-image pixel coordinates to input-image pixel coordinates,
+with the following conventions:
+
+- **0-based, pixel-center convention**: pixel ``0`` is the center of the
+  first pixel (the ``scipy.ndimage.map_coordinates`` / FITS-``origin=0``
+  convention), not FITS's 1-based convention.
+- **Shape**: ``x_out``, ``y_out`` may be arrays of any broadcastable shape;
+  ``x_in``, ``y_in`` come back the same shape. Batching over many images is
+  done by ``vmap``-ing the transform (and whatever consumes it), not by the
+  transform itself handling a batch dimension.
+- **Invalid points are NaN**: pixels with no valid corresponding input
+  position (off the projection's visible hemisphere, off the celestial
+  sphere) come back as NaN in both ``x_in`` and ``y_in``.
+- **Units**: everything at this boundary is in pixels. Providers built on
+  spherical projections will generally work in radians internally, but that
+  is an implementation detail, not part of the contract.
+- **Jit/grad-safe**: a conforming transform is built from pure array
+  operations (no host callbacks, no data-dependent Python control flow) so
+  that it can be used under ``jax.jit`` and differentiated with
+  ``jax.grad``.
+
+Providers may optionally expose ``pixel_map(x_out, y_out, *,
+bounds_shape=(ny, nx))``. It returns finite mapped coordinates and a boolean
+validity mask; interpolation uses this capability when available and falls
+back to the two-result callable contract otherwise.
+"""
+
+from typing import Any, Protocol, Tuple
+
+
+class PixelTransform(Protocol):
+    """Callable mapping output pixel coordinates to input pixel coordinates.
+
+    See the module docstring for the full contract. This is a
+    :class:`typing.Protocol`: any callable with this signature satisfies it
+    structurally, with no need to import or subclass anything from
+    reprojax.
+    """
+
+    def __call__(self, x_out: Any, y_out: Any) -> Tuple[Any, Any]:
+        ...

reprojax/interp.py

@@ -0,0 +1,278 @@
+"""Nearest/bilinear reprojection: the pure-JAX core and the reproject-shaped wrapper.
+
+The core (:func:`reproject_interp_core`) consumes a :class:`~reprojax.core.PixelTransform`
+and is a pure function of arrays — jit/vmap/grad-safe. The wrapper
+(:func:`reproject_interp`) does the astropy-facing work (header/WCS parsing,
+transform composition) and matches ``reproject``'s call/return shape.
+"""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from astropy.io.fits import Header
+from astropy.wcs import WCS
+from jax.scipy.ndimage import map_coordinates
+
+from .transform import build_transform, compose_transform
+
+SUPPORTED_ORDERS = (0, 1)
+
+
+def _output_grid_dtype(matrix_precision):
+    """The dtype for the output pixel grid built by the wrappers.
+
+    The grid must match the transform's coordinate dtype, not a fixed
+    float64: under ``matrix_precision="f32"`` every per-pixel call runs in
+    float32, so a float64 grid only costs accuracy and, with
+    ``jax_enable_x64`` off, emits a truncation warning on every call.
+    """
+    return jnp.float32 if matrix_precision == "f32" else jnp.float64
+
+# Below this coverage, at least one interpolation neighbor was outside the
+# input array and contributed through `cval`. The threshold sits far below
+# any real partial-coverage value (>= 1/(order+2)) and far above float32
+# rounding noise, so it only ever separates "fully covered" from "not".
+_COVERAGE_VALID_THRESHOLD = 1e-5
+
+
+def _clip_border_pixels(coord, size):
+    """Snap coordinates in the outer half-pixel border band to the border
+    pixel's center index.
+
+    ``map_coordinates`` (like plain scipy) treats pixel values as defined
+    only at integer pixel centers, so a coordinate in the outer half of the
+    border pixel is technically "outside" the defined domain even though it
+    is still within that pixel. ``reproject`` corrects for this with the
+    same clip (see ``reproject._array_utils._clip_coords``); we match it so
+    our footprints agree with reproject's at the image edges.
+    """
+    lower = (coord < 0) & (coord >= -0.5)
+    upper = (coord < size - 0.5) & (coord >= size - 1)
+    coord = jnp.where(lower, 0.0, coord)
+    coord = jnp.where(upper, size - 1.0, coord)
+    return coord
+
+
+def reproject_interp_core(data, x_out, y_out, transform, order=1):
+    """Reproject ``data`` onto the pixel grid given by ``x_out``/``y_out``.
+
+    Parameters
+    ----------
+    data : array, shape (ny_in, nx_in)
+        Input image.
+    x_out, y_out : arrays, any common broadcastable shape
+        Output-image pixel coordinates to fill (0-based, pixel-center
+        convention — see :mod:`reprojax.core`).
+    transform : PixelTransform
+        Maps ``(x_out, y_out) -> (x_in, y_in)``.
+    order : int
+        0 (nearest) or 1 (bilinear).
+
+    Returns
+    -------
+    result : array, same shape as x_out
+        Reprojected values, dtype matching ``data``. NaN where invalid.
+    footprint : array, same shape as x_out
+        1.0 where the value is valid, 0.0 where not.
+    """
+    if order not in SUPPORTED_ORDERS:
+        raise ValueError(f"order must be one of {SUPPORTED_ORDERS} (bicubic and higher are not supported), got {order!r}")
+
+    ny, nx = data.shape
+    pixel_map = getattr(transform, "pixel_map", None)
+    has_pixel_map = callable(pixel_map)
+    if has_pixel_map:
+        x_in, y_in, valid = pixel_map(x_out, y_out, bounds_shape=(ny, nx))
+        coord_is_nan = ~valid
+    else:
+        x_in, y_in = transform(x_out, y_out)
+        coord_is_nan = jnp.isnan(x_in) | jnp.isnan(y_in)
+        valid = ~coord_is_nan
+    # Generic providers may return NaN coordinates, which must never reach
+    # map_coordinates: NaN interpolation weights can corrupt data gradients
+    # even when the final result is masked invalid. pixel_map providers
+    # already guarantee finite zero stand-ins, so they need no second
+    # sanitization pass.
+    if has_pixel_map:
+        x_safe, y_safe = x_in, y_in
+    else:
+        x_safe = jnp.where(coord_is_nan, 0.0, x_in)
+        y_safe = jnp.where(coord_is_nan, 0.0, y_in)
+    y_clipped = _clip_border_pixels(y_safe, ny)
+    x_clipped = _clip_border_pixels(x_safe, nx)
+    coords = [y_clipped, x_clipped]
+
+    result = map_coordinates(data, coords, order=order, mode="constant", cval=0.0)
+
+    if has_pixel_map:
+        valid_coord = valid
+    else:
+        # For generic providers, retain the existing coverage interpolation:
+        # its value is the fraction of interpolation weight that landed on
+        # in-bounds neighbors. A float32 fraction is exact enough for the
+        # threshold below, and matches the float32 path under x64-off.
+        coverage_input = jnp.ones((ny, nx), dtype=jnp.float32)
+        coverage = map_coordinates(
+            coverage_input, coords, order=order, mode="constant", cval=0.0
+        )
+        valid_coord = (
+            valid
+            & (coverage >= 1.0 - _COVERAGE_VALID_THRESHOLD)
+            & ~coord_is_nan
+        )
+
+    result = jnp.where(valid_coord, result, jnp.nan).astype(data.dtype)
+    # Footprint tracks the final result, not just coordinate validity: `data`
+    # itself can contain NaN (real masked/bad pixels), which propagates
+    # through interpolation even when the coordinate is fully in-bounds.
+    # Matches reproject's own footprint definition (~isnan(output)).
+    footprint = (~jnp.isnan(result)).astype(jnp.float32)
+    return result, footprint
+
+
+def reproject_interp_batched(
+    data_batch,
+    mwcs_in_batch,
+    mwcs_out,
+    rotation_batch,
+    x_out,
+    y_out,
+    order=1,
+):
+    """Reproject a batch of exposures onto one shared output grid.
+
+    Batching is `jax.vmap` over the whole per-exposure pipeline (build the
+    transform, then reproject) — per the design, batching happens by
+    vmapping the core, not by threading a batch dimension through a single
+    transform. This is a thin, documented recipe: build ``mwcs_in_batch``
+    (and ``rotation_batch``) yourself with mapprojax's ``Tan``/``Sin`` (which are
+    JAX pytrees), each leaf carrying a leading batch dimension, and everything
+    else follows the same path as the single-exposure core.
+
+    Parameters
+    ----------
+    data_batch : array, shape (n, ny_in, nx_in)
+        Stacked input images. Individual exposures may differ in size only
+        if you pad them to a common shape first — vmap requires uniform
+        shapes across the batch.
+    mwcs_in_batch : a mapprojax WCS (e.g. ``mapprojax.wcs.Tan``), batched
+        Construct with each parameter (crpix, cd, crval, and SIP
+        coefficients if used) carrying a leading dimension of size ``n`` —
+        e.g. ``Tan(crpix_batch, cd_batch, crval_batch, sip=sip_batch,
+        precision="f64")`` (``precision=`` is required).
+        Parameters that don't vary across the batch (e.g. a shared crpix for
+        exposures from the same detector) still need that leading dimension;
+        broadcast them yourself (``jnp.broadcast_to``) before constructing.
+        Every exposure must share the same SIP-or-not structure (vmap
+        requires a uniform pytree).
+    mwcs_out : a mapprojax WCS
+        The single, shared output projection (not batched).
+    rotation_batch : array, shape (n, 3, 3)
+        Per-exposure celestial frame rotation (see :func:`build_transform`).
+        Often identical across the batch (same input/output frame pair for
+        every exposure); broadcast to shape ``(n, 3, 3)`` in that case.
+    x_out, y_out : arrays
+        The shared output pixel grid (not batched).
+    order : int
+        0 (nearest) or 1 (bilinear).
+
+    Returns
+    -------
+    results, footprints : arrays, shape (n, *x_out.shape)
+        Element-wise identical to calling :func:`reproject_interp_core` once
+        per exposure with the same ``x_out``/``y_out``.
+    """
+
+    def _single(data, mwcs_in, rotation):
+        transform = compose_transform(mwcs_in, mwcs_out, rotation)
+        return reproject_interp_core(data, x_out, y_out, transform, order=order)
+
+    return jax.vmap(_single, in_axes=(0, 0, 0))(data_batch, mwcs_in_batch, rotation_batch)
+
+
+def _as_wcs(wcs_or_header):
+    if isinstance(wcs_or_header, WCS):
+        return wcs_or_header
+    if isinstance(wcs_or_header, Header):
+        return WCS(wcs_or_header)
+    raise TypeError(f"expected an astropy WCS or Header, got {type(wcs_or_header)!r}")
+
+
+def _to_native_byte_order(data):
+    """FITS data is always big-endian; JAX only accepts native-byte-order
+    arrays. A plain ``jnp.asarray`` on data read straight from
+    ``astropy.io.fits`` fails outright, so the wrapper normalizes here
+    before anything touches JAX.
+    """
+    data = np.asarray(data)
+    if not data.dtype.isnative:
+        data = data.astype(data.dtype.newbyteorder("="))
+    return data
+
+
+def _shape_from_header(wcs_or_header):
+    if isinstance(wcs_or_header, Header):
+        header = wcs_or_header
+    elif getattr(wcs_or_header, "array_shape", None) is not None:
+        return tuple(wcs_or_header.array_shape)
+    else:
+        return None
+    if "NAXIS2" in header and "NAXIS1" in header:
+        return (header["NAXIS2"], header["NAXIS1"])
+    return None
+
+
+def reproject_interp(input_data, output_projection, shape_out=None, order=1, matrix_precision="f64"):
+    """Reproject an ``(array, WCS-or-header)`` pair onto ``output_projection``.
+
+    Matches ``reproject.reproject_interp``'s call and return shape for the
+    subset of input forms reprojax currently supports.
+
+    Parameters
+    ----------
+    input_data : tuple(array, astropy.wcs.WCS or astropy.io.fits.Header)
+    output_projection : astropy.wcs.WCS or astropy.io.fits.Header
+    shape_out : tuple(int, int), optional
+        Required unless it can be inferred from ``output_projection``
+        (a Header's NAXIS1/NAXIS2, or a WCS with ``array_shape`` set).
+    order : int
+        0 (nearest) or 1 (bilinear).
+    matrix_precision : {"f64", "f32"}
+        See :func:`~reprojax.transform.build_transform`. ``"f32"`` needs
+        no ``jax_enable_x64`` and is faster on GPUs where float64
+        throughput is much lower than float32's, at reduced accuracy
+        (~1e-4 px for typical overlapping-field reprojections; see that
+        function's docstring for the detailed tradeoff).
+
+    Returns
+    -------
+    array, footprint
+
+    Note
+    ----
+    Like :func:`~reprojax.adaptive.reproject_adaptive`, this calls
+    :func:`~reprojax.transform.build_transform` internally, so repeated
+    calls with an equivalent ``wcs_in``/``wcs_out`` pair reuse the same
+    transform object rather than rebuilding (and retracing) it — see that
+    function's docstring for the caching details.
+    """
+    data, wcs_or_header_in = input_data
+    wcs_in = _as_wcs(wcs_or_header_in)
+    wcs_out = _as_wcs(output_projection)
+
+    if shape_out is None:
+        shape_out = _shape_from_header(output_projection)
+        if shape_out is None:
+            raise ValueError(
+                "shape_out could not be inferred from output_projection; pass it explicitly"
+            )
+    ny, nx = shape_out
+
+    grid_dtype = _output_grid_dtype(matrix_precision)
+    y_out, x_out = jnp.meshgrid(
+        jnp.arange(ny, dtype=grid_dtype), jnp.arange(nx, dtype=grid_dtype), indexing="ij"
+    )
+
+    transform = build_transform(wcs_in, wcs_out, matrix_precision=matrix_precision)
+    data = jnp.asarray(_to_native_byte_order(data))
+    return reproject_interp_core(data, x_out, y_out, transform, order=order)

reprojax/transform.py

@@ -0,0 +1,231 @@
+"""Build a :class:`~reprojax.core.PixelTransform` from a pair of astropy WCS objects.
+
+mapprojax provides the projection math (TAN, SIN, with optional SIP
+distortion); astropy provides celestial frame determination and rotation.
+Everything astropy-facing runs once, eagerly, at build time in
+:func:`build_transform`; the returned callable is pure JAX and safe under
+``jit``/``vmap``/``grad``.
+"""
+
+from collections import OrderedDict
+
+import jax.numpy as jnp
+from astropy.wcs.utils import wcs_to_celestial_frame
+from mapprojax.wcs import Sin, Tan
+from mapprojax.reproject import Reprojector
+from mapprojax.utils import apply_rotation
+
+from ._astropy_wcs import frame_rotation_matrix, projection_code, wcs_params_from_astropy
+
+_PROVIDERS = {"TAN": Tan, "SIN": Sin}
+MATRIX_PRECISIONS = ("f64", "f32")
+
+# build_transform's cache: keyed on the numeric content of wcs_in/wcs_out
+# (not object identity), so repeated calls with re-parsed-but-identical WCS
+# (e.g. the same header read from disk each time, or the same pointing's
+# multiple spectral channels) hit the cache too, not just calls reusing the
+# same Python WCS object. Bounded and LRU-evicted since each entry pins a
+# compiled transform closure (and jit will hold a compilation cache entry
+# keyed on its identity) in memory for as long as the cache keeps it.
+_TRANSFORM_CACHE_MAXSIZE = 32
+_TRANSFORM_CACHE = OrderedDict()
+
+
+def _mapprojax_wcs(astropy_wcs, precision="f64"):
+    code = projection_code(astropy_wcs)
+    crpix, cd, crval, sip = wcs_params_from_astropy(astropy_wcs)
+    return _PROVIDERS[code](crpix, cd, crval, sip=sip, precision=precision)
+
+
+def _sip_fingerprint(sip):
+    if sip is None:
+        return None
+    return tuple(None if arr is None else (arr.shape, arr.tobytes()) for arr in (sip.a, sip.b, sip.ap, sip.bp))
+
+
+def _wcs_fingerprint(astropy_wcs):
+    code = projection_code(astropy_wcs)
+    crpix, cd, crval, sip = wcs_params_from_astropy(astropy_wcs)
+    frame = repr(wcs_to_celestial_frame(astropy_wcs))
+    return (code, tuple(crpix), tuple(map(tuple, cd)), tuple(crval), _sip_fingerprint(sip), frame)
+
+
+def compose_transform(mwcs_in, mwcs_out, rotation):
+    """Build a ``PixelTransform`` from already-constructed mapprojax WCS
+    objects and a fixed frame-rotation matrix — the pure-JAX half of
+    :func:`build_transform`, with no astropy involved.
+
+    ``mwcs_in``/``mwcs_out`` are JAX pytrees (mapprojax registers
+    ``Tan``/``Sin``/``Sip`` for exactly this), so unlike
+    ``build_transform`` (which takes non-traceable ``astropy.wcs.WCS``
+    objects), this function's WCS parameters — CRVAL, CD, SIP coefficients —
+    are differentiable: construct ``mwcs_in``/``mwcs_out`` inside a function
+    you pass to ``jax.grad``, and gradients flow through here into whatever
+    consumes the returned transform.
+
+    The relative rotation ``mwcs_in.r_matrix @ rotation @ mwcs_out.r_matrix.T``
+    is formed in JAX (so it stays differentiable with respect to CRVAL) and
+    handed to :meth:`mapprojax.reproject.Reprojector.from_matrix` with
+    ``matrix_precision="f64"``; the returned transform is that
+    ``Reprojector``'s ``pix_to_pix``.
+
+    Parameters
+    ----------
+    mwcs_in, mwcs_out : mapprojax WCS
+        ``mapprojax.wcs.Tan`` or ``mapprojax.wcs.Sin``, constructed with
+        ``precision="f64"`` (the ``precision=`` argument is required), and
+        optionally carrying a ``mapprojax.sip.Sip``.
+    rotation : array, shape (3, 3)
+        Fixed celestial-frame rotation from ``mwcs_out``'s frame to
+        ``mwcs_in``'s (not a differentiable parameter — it comes from
+        astropy's frame machinery in ``build_transform``, but a constant
+        array is equally valid to pass directly).
+    Returns
+    -------
+    PixelTransform
+    """
+
+    m_rel = mwcs_in.r_matrix @ rotation @ jnp.swapaxes(mwcs_out.r_matrix, -1, -2)
+    reprojector = Reprojector.from_matrix(
+        mwcs_out, mwcs_in, m_rel, matrix_precision="f64"
+    )
+
+    def transform(x_out, y_out):
+        return reprojector.pix_to_pix(x_out, y_out)
+
+    def pixel_map(x_out, y_out, *, bounds_shape=None):
+        return reprojector.pixel_map(x_out, y_out, bounds_shape=bounds_shape)
+
+    transform.pixel_map = pixel_map
+    return transform
+
+
+def compose_transform_fp32(reprojector):
+    """Build a ``PixelTransform`` from a :class:`mapprojax.reproject.Reprojector`
+    -- the fp64-free counterpart of :func:`compose_transform`.
+
+    The celestial<->native rotation runs once, in true float64, at
+    ``Reprojector`` construction time (not per pixel -- see its
+    docstring); every per-pixel call after that (``pix_to_native -> apply
+    the composed rotation -> native_to_pix``) runs in float32,
+    independent of ``jax_enable_x64``.
+
+    Note
+    ----
+    Unlike :func:`compose_transform`, this is *not* differentiable with
+    respect to either WCS's ``crval``: ``Reprojector``'s rotation is
+    precomputed outside JAX's trace (see its docstring). It remains
+    differentiable with respect to the output pixel coordinates
+    (``x_out``, ``y_out``) -- all ``reproject_adaptive``'s
+    ``jax.jacfwd`` needs -- and, through the target leg, with respect to
+    ``crpix``/``cd``/SIP coefficients.
+
+    Parameters
+    ----------
+    reprojector : mapprojax.reproject.Reprojector
+        Maps output pixels (its ``wcs_src``) to input pixels (its
+        ``wcs_tgt``) -- see :func:`build_transform`.
+    Returns
+    -------
+    PixelTransform
+    """
+    def transform(x_out, y_out):
+        return reprojector.pix_to_pix(x_out, y_out)
+
+    def pixel_map(x_out, y_out, *, bounds_shape=None):
+        return reprojector.pixel_map(x_out, y_out, bounds_shape=bounds_shape)
+
+    transform.pixel_map = pixel_map
+    return transform
+
+
+def build_transform(wcs_in, wcs_out, matrix_precision="f64"):
+    """Build a ``PixelTransform`` mapping output pixels to input pixels.
+
+    The composed pipeline, per output pixel: output pixel -> output-frame
+    unit vector (``wcs_out.unproj_xyz``) -> input-frame unit vector
+    (celestial frame rotation) -> input pixel (``wcs_in.proj_xyz``).
+
+    Parameters
+    ----------
+    wcs_in, wcs_out : astropy.wcs.WCS
+        TAN or SIN projections, optionally with SIP distortion.
+    matrix_precision : {"f64", "f32"}
+        - ``"f64"`` (default): builds ``mwcs_in``/``mwcs_out`` as mapprojax
+          ``Tan``/``Sin`` with ``precision="f64"`` and composes them via
+          :func:`compose_transform` (``Reprojector.from_matrix`` with
+          ``matrix_precision="f64"``), entirely in float64. Requires
+          ``jax_enable_x64`` for correctness.
+        - ``"f32"``: builds ``precision="f32"`` WCS objects and a
+          :class:`mapprojax.reproject.Reprojector` with
+          ``matrix_precision="f32"``, which composes the two WCSes'
+          rotations with the frame rotation into one matrix, once, in true
+          float64 on the host, independent of ``jax_enable_x64``; every
+          per-pixel call then runs in float32. No ``jax_enable_x64``
+          needed, and faster on GPUs where float64 throughput is much
+          lower than float32's. Less accurate than ``"f64"`` -- a few
+          e-4 px on the real SPHEREx ICRS<->Galactic pair in
+          ``tests/test_real_spherex.py``, growing with field-center
+          separation and with finer plate scales (see mapprojax's
+          ``Reprojector`` docstring); not differentiable with respect to
+          either WCS's ``crval`` (see :func:`compose_transform_fp32`).
+
+    Returns
+    -------
+    PixelTransform
+
+    Note
+    ----
+    This function itself is not traceable (it constructs and reads
+    ``astropy.wcs.WCS`` objects, which are not JAX pytrees) — the returned
+    *transform* is pure JAX, but you cannot ``jax.grad`` through
+    ``build_transform`` with respect to ``wcs_in``/``wcs_out``. To
+    differentiate with respect to WCS parameters (CRVAL, CD, SIP
+    coefficients), construct mapprojax WCS objects directly and use
+    :func:`compose_transform`.
+
+    Note
+    ----
+    Results are cached on the numeric content of ``wcs_in``/``wcs_out`` (CD,
+    CRPIX, CRVAL, SIP coefficients, projection code, and celestial frame)
+    — calling this again with an equivalent
+    WCS pair, even freshly reconstructed from the same header, returns the
+    *same* transform object rather than rebuilding it. Reusing that object
+    also keeps the compiled interpolation/adaptive call paths on their
+    established cache key; a changed WCS fingerprint or static configuration
+    can require retracing, especially for the autodiff-heavy adaptive core.
+    Reusing the same WCS pair across many calls (e.g. per-band exposures
+    sharing one pointing) pays the construction and compilation costs once.
+    """
+    if matrix_precision not in MATRIX_PRECISIONS:
+        raise ValueError(f"matrix_precision must be one of {MATRIX_PRECISIONS}, got {matrix_precision!r}")
+
+    cache_key = (_wcs_fingerprint(wcs_in), _wcs_fingerprint(wcs_out), matrix_precision)
+    cached = _TRANSFORM_CACHE.get(cache_key)
+    if cached is not None:
+        _TRANSFORM_CACHE.move_to_end(cache_key)
+        return cached
+
+    if matrix_precision == "f32":
+        wcs_in_np = _mapprojax_wcs(wcs_in, precision="f32")
+        wcs_out_np = _mapprojax_wcs(wcs_out, precision="f32")
+        frame_rotation = frame_rotation_matrix(wcs_out, wcs_in)
+        reprojector = Reprojector(wcs_out_np, wcs_in_np, frame_rotation=frame_rotation, matrix_precision="f32")
+        transform = compose_transform_fp32(reprojector)
+    else:
+        mwcs_in = _mapprojax_wcs(wcs_in)
+        mwcs_out = _mapprojax_wcs(wcs_out)
+        rotation = jnp.asarray(frame_rotation_matrix(wcs_out, wcs_in))
+        transform = compose_transform(mwcs_in, mwcs_out, rotation)
+
+    _TRANSFORM_CACHE[cache_key] = transform
+    if len(_TRANSFORM_CACHE) > _TRANSFORM_CACHE_MAXSIZE:
+        _TRANSFORM_CACHE.popitem(last=False)
+    return transform
+
+
+def _clear_transform_cache():
+    _TRANSFORM_CACHE.clear()
+
+
+build_transform.cache_clear = _clear_transform_cache

reprojax-0.2.0.dist-info/METADATA

@@ -0,0 +1,198 @@
+Metadata-Version: 2.4
+Name: reprojax
+Version: 0.2.0
+Summary: JAX-native, differentiable reprojection of astronomical images (a port of astropy's reproject)
+License-Expression: BSD-3-Clause
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=2.4.6
+Requires-Dist: jax>=0.10.2
+Requires-Dist: astropy>=7.2.0
+Requires-Dist: mapprojax<0.4,>=0.3.0
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: reproject>=0.21; extra == "test"
+Dynamic: license-file
+
+# reprojax
+
+A JAX-native, differentiable port of astropy's [`reproject`](https://github.com/astropy/reproject) package.
+
+reprojax reprojects astronomical images between WCS frames on top of
+[mapprojax](https://github.com/leejjoon/mapprojax), a JAX port of CDS's
+`mapproj` crate providing fast, differentiable pixel↔sky transforms. The goal
+is GPU-accelerated, batchable, gradient-friendly reprojection, validated
+against the original `reproject` package as the accuracy reference.
+
+reprojax is a JAX port of `reproject`'s algorithms and API shape; see
+`LICENSE` for `reproject`'s BSD-3-Clause license, carried alongside
+reprojax's own.
+
+## Status
+
+0.2.0, first public release. The validated scope is deliberate:
+
+- **Projections**: TAN and SIN, with SIP distortion. Anything else raises
+  `NotImplementedError`.
+- **Algorithms**: bilinear/nearest interpolation (`reproject_interp`) and the
+  DeForest (2004) adaptive anti-aliased algorithm (`reproject_adaptive`).
+- **Reference**: `reproject` 0.21.0 is the golden accuracy reference. It is a
+  **test-time** dependency only; nothing in `reprojax` imports it at runtime.
+- **Not included**: cubic interpolation, flux-conserving resampling, mosaic
+  coadds, filename/HDU/`NDData` input forms. See the
+  [issue tracker](https://github.com/leejjoon/reprojax/issues) for the plan.
+
+## Precision: read this before your first call
+
+`matrix_precision` selects the transform's coordinate precision, and the
+**default is `"f64"`, which requires JAX x64**. Without it your first call
+raises:
+
+```
+ValueError: precision='f64' requires float64. Call
+jax.config.update('jax_enable_x64', True) before constructing it.
+```
+
+Two valid choices:
+
+| Mode | Requirement | Use it for |
+| --- | --- | --- |
+| `"f64"` (default) | `jax.config.update("jax_enable_x64", True)` before importing/creating anything JAX | Full-precision work, astrometry, gradients with respect to WCS parameters |
+| `"f32"` | none — works with x64 off | GPU throughput; the celestial rotation is still composed in float64 on the host, only the per-pixel path is float32 |
+
+`"f64"` is the accurate default. `"f32"` is the fast path and is documented
+in terms of its own measured caps, not as an equivalent to `"f64"` (below).
+
+## Quick start
+
+No data files needed. This runs exactly as written, from an installed package:
+
+```python
+import numpy as np
+import jax
+jax.config.update("jax_enable_x64", True)  # required by the default matrix_precision="f64"
+
+from astropy.io import fits
+from astropy.wcs import WCS
+from reprojax import reproject_interp
+
+ny_in = nx_in = 64
+hdr_in = fits.Header()
+hdr_in["NAXIS"] = 2
+hdr_in["NAXIS1"] = hdr_in["NAXIS2"] = ny_in
+hdr_in["CTYPE1"], hdr_in["CTYPE2"] = "RA---TAN", "DEC--TAN"
+hdr_in["CRPIX1"] = hdr_in["CRPIX2"] = ny_in / 2 + 0.5
+hdr_in["CRVAL1"], hdr_in["CRVAL2"] = 210.0, 54.0
+hdr_in["CDELT1"], hdr_in["CDELT2"] = -1e-3, 1e-3
+
+y, x = np.mgrid[:ny_in, :nx_in]
+data = np.exp(-(((x - 32) ** 2 + (y - 32) ** 2) / (2 * 6.0 ** 2))).astype(np.float32)
+
+hdr_out = hdr_in.copy()
+hdr_out["CRVAL1"] = 210.02  # shift the output grid
+
+array, footprint = reproject_interp((data, WCS(hdr_in)), hdr_out, shape_out=(40, 40))
+
+assert array.shape == footprint.shape == (40, 40)
+assert np.isfinite(array).any()
+print(float(np.nanmax(array)))
+```
+
+`array` holds the resampled image (NaN where no valid input pixel was
+covered) and `footprint` is 1.0 where a valid input contributed, 0.0
+otherwise. Adding `matrix_precision="f32"` runs the same call without x64.
+
+## Architecture
+
+- `reprojax.core` — the `PixelTransform` contract: a callable
+  `transform(x_out, y_out) -> (x_in, y_in)` mapping output-image pixel
+  coordinates to input-image pixel coordinates. This is the seam the rest of
+  the package is built around.
+- `reprojax.transform` — builds a `PixelTransform` from a pair of astropy
+  WCS objects, using mapprojax for the projection math and astropy for
+  celestial frame rotation.
+- Mapprojax-backed transforms also expose an optional
+  `transform.pixel_map(x_out, y_out, bounds_shape=(ny, nx))` capability. It
+  supplies finite coordinates and geometric validity to interpolation while
+  the two-result callable contract remains available for generic providers.
+
+## Accuracy
+
+Two kinds of numbers appear here, and they are labelled: an **observation**
+is what the two implementations measured on a case, an **enforced tolerance**
+is what the test suite requires.
+
+Real SPHEREx pair (2040² ICRS TAN+SIP exposure onto the 1201² Galactic TAN
+template, compared against `reproject` 0.21.0):
+
+| Algorithm | Observation | Enforced |
+| --- | --- | --- |
+| `reproject_interp` | 1,345,444 of 1,345,460 jointly-valid pixels agree to `rtol <= 1e-8` | `rtol=1e-6, atol=0.01` over all jointly-valid pixels, plus exact footprint equality |
+| `reproject_adaptive` | median disagreement ~`4e-8`; 1,334,193 of 1,334,200 jointly-valid pixels form the bulk | `rtol=1e-6, atol=0.005`, plus exact footprint equality |
+
+In both algorithms the pixels outside the stated bulk are 16 (interp) and 7
+(adaptive) pixels sitting at one real, un-flagged data spike: this fixture
+drops the `FLAGS` extension that would normally mask it, so a sub-micropixel
+position difference next to a ~10⁴-per-pixel local gradient is amplified into
+a value difference of up to ~0.01 (interp) / ~0.003 (adaptive). The fixtures
+are described in [`tests/data/README.md`](tests/data/README.md).
+
+Synthetic TAN / TAN-SIP pairs: interpolated images agree within enforced
+`rtol=1e-8, atol=1e-9`; adaptive within enforced `rtol=1e-6, atol=1e-8`.
+
+`matrix_precision="f32"` has **enforced coordinate caps** of `6e-4` px
+(forward) and `5e-4` px (reverse) per direction on the real pair. These bound
+the coordinate map only — they are not image-agreement bounds, and they are
+not a claim of float64 equivalence.
+
+On a Strix Halo host (ROCm device, JAX 0.11.1 + rocm7 plugins 0.11.1),
+batched and `jax.jit`-compiled reprojax reprojects 8 SPHEREx-scale
+(2040² → 1201²) exposures in 0.026 s. The looped CPU `reproject_interp`
+baseline for the same work measured 4.4–7.9 s across runs on this host, i.e.
+roughly 170–300×. The reprojax figure is stable; the baseline is what varies.
+This is a cross-backend batched-vs-looped comparison, not a like-for-like
+backend comparison, and the CPU skips the batched GPU-only tests.
+
+## Development
+
+```bash
+python -m pip install -e ".[test]" -c ci/latest.txt
+python -m pytest tests/
+```
+
+Using a `ci/` constraints file matters: the golden reference is pinned to
+`reproject==0.21.0` there, and the suite asserts the installed version. Use
+`ci/minimum.txt` for the oldest supported stack (Python 3.11). If a
+constraint conflicts with your environment, install `.[test]` without it and
+expect the version guard to tell you which pin is off.
+
+The test suite needs `tests/data/`, which is 16 MB of real SPHEREx
+exposures and is therefore not shipped in the sdist. Install from a checkout
+when you want to run the tests.
+
+## License
+
+BSD-3-Clause. The SIP and reprojection algorithms are ported from
+[`reproject`](https://github.com/astropy/reproject); its license is carried in
+[`LICENSE`](LICENSE) alongside reprojax's own.
+
+## Test data
+
+The integration tests use a real SPHEREx exposure and a Galactic template
+header. Only the exposure is mission data: NASA's Science Mission
+Directorate releases mission data into the public domain (SPD-41), while the
+template header is project metadata under reprojax's own licence. If you
+publish work that uses the exposure, include the acknowledgement IRSA
+requests:
+
+> This publication makes use of data products from the Spectro-Photometer for
+the History of the Universe, Epoch of Reionization and Ices Explorer
+> (SPHEREx), which is a joint project of the Jet Propulsion Laboratory and the
+> California Institute of Technology, and is funded by the National
+> Aeronautics and Space Administration.
+
+Provenance, checksums and the citation DOIs are in
+[`tests/data/README.md`](tests/data/README.md).

reprojax-0.2.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+reprojax/__init__.py,sha256=YYsuLMCsd7olVt0avymZIHvAP8W-w-FR05h1b_9Vmkk,470
+reprojax/_astropy_wcs.py,sha256=thtdOSwdV5cJSY6HwB5G9zDzVK_Uan-QMC7P8a0WtpA,2733
+reprojax/adaptive.py,sha256=AggxrpAmgVFpUPCqvbbYtzl6ouRR8jJIPvpA2PgWIeY,14038
+reprojax/core.py,sha256=Ad6y0cC9ceTiWAD0A3YtCgbYgRgehvZkcoaPJqdE3H8,2427
+reprojax/interp.py,sha256=YeDMbE6fL05qgHI9WuzeuScOYQh7BbMwtGBIGDeuB1E,11419
+reprojax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+reprojax/transform.py,sha256=8F_1-BGhMo7FCuZY9J5B-zEz6rS7KT1GS71rYlKOMiM,9992
+reprojax-0.2.0.dist-info/licenses/LICENSE,sha256=i43S9vaJvTB0wJJqbreVhKW33qADGXP2ZkgY6A9XuPI,3130
+reprojax-0.2.0.dist-info/METADATA,sha256=_Du-WoIdixNt2UmFAnPiZc-XOzsm82NeVyAgGe_gxU4,8604
+reprojax-0.2.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+reprojax-0.2.0.dist-info/top_level.txt,sha256=EwQoY_6Xd2nriNTaiQU7SBqK7IKHjclqGM_bDkpvank,9
+reprojax-0.2.0.dist-info/RECORD,,

reprojax-0.2.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

reprojax-0.2.0.dist-info/licenses/LICENSE

@@ -0,0 +1,63 @@
+Copyright (c) 2026, Jae-Joon Lee
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the reprojax project nor the names of its contributors
+  may be used to endorse or promote products derived from this software
+  without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+POSSIBILITY OF SUCH DAMAGE.
+
+---
+
+reprojax is a JAX port of astropy's `reproject` package
+(https://github.com/astropy/reproject), which carries the following
+license:
+
+Copyright (c) 2011-16, Thomas P. Robitaille
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the Astropy project nor the names of its contributors
+  may be used to endorse or promote products derived from this software
+  without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+POSSIBILITY OF SUCH DAMAGE.

reprojax-0.2.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+reprojax