autogalaxy.TransformerNUFFT#

class TransformerNUFFT[source]#

Bases: object

JAX-native Non-Uniform FFT for image -> visibilities, backed by nufftax.

This is the default TransformerNUFFT in PyAutoArray. It uses the nufftax library (https://github.com/GragasLab/nufftax), a pure-JAX NUFFT implementation that supports jax.jit, jax.grad, and jax.vmap. Note that nufftax requires JAX; on platforms with no JAX wheels (notably Intel macOS) use TransformerDFT instead.

Convention recipe (matches TransformerDFT to ~1e-13 relative across odd/even/non-square image sizes):

image_flipped = image[::-1, :] x = 2 * pi * u_lambda * pixel_scale_rad y = 2 * pi * v_lambda * pixel_scale_rad offset_x = 0.5 if N_x is even else 0.0 offset_y = 0.5 if N_y is even else 0.0 shift = exp(-i * (offset_x * x + offset_y * y)) visibilities = nufftax.nufft2d2(x, y, image_flipped, eps, -1) * shift

The shift factor is the half-pixel correction between autoarray’s grid centre at index (N - 1) / 2 and nufftax’s mode-0 at index N // 2; nufftax does not apply it internally.

Parameters:

uv_wavelengths (ndarray) – The (u, v) coordinates of the measured visibilities in wavelengths, shape (n_vis, 2).
real_space_mask (Mask2D) – The 2D mask defining the real-space image grid.
eps (float) – Requested NUFFT precision passed to nufftax. Defaults to 1e-12 (effectively machine precision); relax to 1e-9 or 1e-6 for faster execution if marginal accuracy is acceptable.
chunk_size (Optional[int]) – If set to a positive integer, the forward and adjoint NUFFT calls split the visibility axis into chunks of this size and iterate (via jax.lax.scan on the JAX path, a Python loop on the numpy path). This caps the nufftax gather-buffer allocation (~``2 * chunk_size * nspread^2 * dtype_size``) at the cost of per-chunk overhead. Required for visibility counts above ~5M on a 40-80 GB GPU. If None (default), a single one-shot call is used — preserves existing behaviour for small-N callers (sma-class datasets).
xp – Accepted for signature compatibility with the legacy class; not stored. The active backend is selected per-call via the xp argument to visibilities_from / image_from.

grid#: The real-space pixel grid in radians (computed from the mask).

total_visibilities#: Number of measured visibilities.

total_image_pixels#: Number of unmasked pixels in the image grid.

__init__(uv_wavelengths, real_space_mask, eps=1e-12, chunk_size=None, xp=<module 'numpy' from '/home/docs/checkouts/readthedocs.org/user_builds/pyautogalaxy/envs/latest/lib/python3.12/site-packages/numpy/__init__.py'>, **kwargs)[source]#

JAX-native Non-Uniform FFT for image -> visibilities, backed by nufftax.

Convention recipe (matches TransformerDFT to ~1e-13 relative across odd/even/non-square image sizes):

image_flipped = image[::-1, :] x = 2 * pi * u_lambda * pixel_scale_rad y = 2 * pi * v_lambda * pixel_scale_rad offset_x = 0.5 if N_x is even else 0.0 offset_y = 0.5 if N_y is even else 0.0 shift = exp(-i * (offset_x * x + offset_y * y)) visibilities = nufftax.nufft2d2(x, y, image_flipped, eps, -1) * shift

The shift factor is the half-pixel correction between autoarray’s grid centre at index (N - 1) / 2 and nufftax’s mode-0 at index N // 2; nufftax does not apply it internally.

Parameters:

uv_wavelengths (ndarray) – The (u, v) coordinates of the measured visibilities in wavelengths, shape (n_vis, 2).
real_space_mask (Mask2D) – The 2D mask defining the real-space image grid.
eps (float) – Requested NUFFT precision passed to nufftax. Defaults to 1e-12 (effectively machine precision); relax to 1e-9 or 1e-6 for faster execution if marginal accuracy is acceptable.
chunk_size (Optional[int]) – If set to a positive integer, the forward and adjoint NUFFT calls split the visibility axis into chunks of this size and iterate (via jax.lax.scan on the JAX path, a Python loop on the numpy path). This caps the nufftax gather-buffer allocation (~``2 * chunk_size * nspread^2 * dtype_size``) at the cost of per-chunk overhead. Required for visibility counts above ~5M on a 40-80 GB GPU. If None (default), a single one-shot call is used — preserves existing behaviour for small-N callers (sma-class datasets).
xp – Accepted for signature compatibility with the legacy class; not stored. The active backend is selected per-call via the xp argument to visibilities_from / image_from.

grid#: The real-space pixel grid in radians (computed from the mask).

total_visibilities#: Number of measured visibilities.

total_image_pixels#: Number of unmasked pixels in the image grid.

Methods

`__init__`(uv_wavelengths, real_space_mask[, ...])	JAX-native Non-Uniform FFT for image -> visibilities, backed by nufftax.
`image_from`(visibilities[, xp])	Adjoint NUFFT: visibilities -> real-space (dirty) image.
`transform_mapping_matrix`(mapping_matrix[, xp])	Apply the forward NUFFT to each column of a mapping matrix.
`visibilities_from`(image[, xp])	Forward NUFFT: real-space image -> visibilities at the configured uv points.

JAX-native Non-Uniform FFT for image -> visibilities, backed by nufftax.

Convention recipe (matches TransformerDFT to ~1e-13 relative across odd/even/non-square image sizes):

image_flipped = image[::-1, :] x = 2 * pi * u_lambda * pixel_scale_rad y = 2 * pi * v_lambda * pixel_scale_rad offset_x = 0.5 if N_x is even else 0.0 offset_y = 0.5 if N_y is even else 0.0 shift = exp(-i * (offset_x * x + offset_y * y)) visibilities = nufftax.nufft2d2(x, y, image_flipped, eps, -1) * shift

The shift factor is the half-pixel correction between autoarray’s grid centre at index (N - 1) / 2 and nufftax’s mode-0 at index N // 2; nufftax does not apply it internally.

Parameters:

uv_wavelengths (ndarray) – The (u, v) coordinates of the measured visibilities in wavelengths, shape (n_vis, 2).
real_space_mask (Mask2D) – The 2D mask defining the real-space image grid.
eps (float) – Requested NUFFT precision passed to nufftax. Defaults to 1e-12 (effectively machine precision); relax to 1e-9 or 1e-6 for faster execution if marginal accuracy is acceptable.
chunk_size (Optional[int]) – If set to a positive integer, the forward and adjoint NUFFT calls split the visibility axis into chunks of this size and iterate (via jax.lax.scan on the JAX path, a Python loop on the numpy path). This caps the nufftax gather-buffer allocation (~``2 * chunk_size * nspread^2 * dtype_size``) at the cost of per-chunk overhead. Required for visibility counts above ~5M on a 40-80 GB GPU. If None (default), a single one-shot call is used — preserves existing behaviour for small-N callers (sma-class datasets).
xp – Accepted for signature compatibility with the legacy class; not stored. The active backend is selected per-call via the xp argument to visibilities_from / image_from.

grid#: The real-space pixel grid in radians (computed from the mask).

total_visibilities#: Number of measured visibilities.

total_image_pixels#: Number of unmasked pixels in the image grid.

visibilities_from(image, xp=<module 'numpy' from '/home/docs/checkouts/readthedocs.org/user_builds/pyautogalaxy/envs/latest/lib/python3.12/site-packages/numpy/__init__.py'>)[source]#

Forward NUFFT: real-space image -> visibilities at the configured uv points.

For numpy callers (xp=np) the result is materialised back to numpy before being wrapped in Visibilities. For JAX callers (xp=jnp) the result stays as a jax.Array so it can flow through jax.jit / jax.grad / jax.vmap without device round-trips.

image_from(visibilities, xp=<module 'numpy' from '/home/docs/checkouts/readthedocs.org/user_builds/pyautogalaxy/envs/latest/lib/python3.12/site-packages/numpy/__init__.py'>)[source]#

Adjoint NUFFT: visibilities -> real-space (dirty) image.

Implemented as nufftax.nufft2d1 with conj(shift) applied to the visibilities and a final row-flip to return to autoarray’s native orientation. The real part is taken to discard imaginary residue.

Note that this is the mathematical adjoint of visibilities_from, with no kernel deconvolution applied. The values match TransformerDFT.image_from exactly, which is what makes Interferometer.apply_sparse_operator scale-consistent across both transformers.

transform_mapping_matrix(mapping_matrix, xp=<module 'numpy' from '/home/docs/checkouts/readthedocs.org/user_builds/pyautogalaxy/envs/latest/lib/python3.12/site-packages/numpy/__init__.py'>)[source]#

Apply the forward NUFFT to each column of a mapping matrix.

All columns are scattered into a single batched native-shape image of shape (n_src, N_y, N_x) and passed through nufft2d2 in one call (nufft2d2 supports batched f). This avoids the per-column Python loop that, under jax.jit, would unroll into n_src separate NUFFT invocations and blow up the JIT graph for pixelization-heavy fits (notably double-source-plane).