autolens.TransformerNUFFT#

class TransformerNUFFT[source]#

Bases: _PrimaryBeamMixin

Non-Uniform FFT for image -> visibilities, backed by finufft on the NumPy path and nufftax on the JAX path.

This is the default TransformerNUFFT in PyAutoArray. The array namespace of each call picks the backend:

xp=np (eager NumPy) runs finufft (https://github.com/flatironinstitute/finufft), the CPU C++ NUFFT, via autoarray.operators.nufft_cpu. This path never imports JAX, so a NumPy fit (e.g. a numba Nautilus search with forked workers) does not start XLA in its parent process. finufft uses OpenMP and honours OMP_NUM_THREADS.
xp=jnp runs nufftax (https://github.com/GragasLab/nufftax), a pure-JAX port of finufft that supports jax.jit, jax.grad and jax.vmap. nufftax is imported only on this path.

The two agree to ~1e-14 relative. Where finufft has no wheel (e.g. Linux aarch64) the NumPy path falls back to nufftax with a warning; where neither is installed 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 the NUFFT’s mode-0 at index N // 2; neither backend applies 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 the backend. 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 NUFFT 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.
primary_beam – The primary beam P of the antennas, an Array2D on real_space_mask (or a native / slim array). When set, visibilities_from and transform_mapping_matrix multiply their input by P on the slim grid before the scatter to the native grid; image_from stays the un-beamed adjoint. None (default) applies no primary beam.

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/pyautolens/envs/latest/lib/python3.12/site-packages/numpy/__init__.py'>, primary_beam=None, **kwargs)[source]#

Non-Uniform FFT for image -> visibilities, backed by finufft on the NumPy path and nufftax on the JAX path.

This is the default TransformerNUFFT in PyAutoArray. The array namespace of each call picks the backend:

xp=np (eager NumPy) runs finufft (https://github.com/flatironinstitute/finufft), the CPU C++ NUFFT, via autoarray.operators.nufft_cpu. This path never imports JAX, so a NumPy fit (e.g. a numba Nautilus search with forked workers) does not start XLA in its parent process. finufft uses OpenMP and honours OMP_NUM_THREADS.
xp=jnp runs nufftax (https://github.com/GragasLab/nufftax), a pure-JAX port of finufft that supports jax.jit, jax.grad and jax.vmap. nufftax is imported only on this path.

The two agree to ~1e-14 relative. Where finufft has no wheel (e.g. Linux aarch64) the NumPy path falls back to nufftax with a warning; where neither is installed 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 the NUFFT’s mode-0 at index N // 2; neither backend applies 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 the backend. 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 NUFFT 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.
primary_beam – The primary beam P of the antennas, an Array2D on real_space_mask (or a native / slim array). When set, visibilities_from and transform_mapping_matrix multiply their input by P on the slim grid before the scatter to the native grid; image_from stays the un-beamed adjoint. None (default) applies no primary beam.

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[, ...])	Non-Uniform FFT for image -> visibilities, backed by finufft on the NumPy path and nufftax on the JAX path.
`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.
`with_primary_beam`(primary_beam)	Returns a shallow copy of this transformer whose forward transforms apply the primary beam primary_beam (an Array2D on the real-space mask, or a native / slim array), or none if it is None.

Attributes

primary_beam_slim

Non-Uniform FFT for image -> visibilities, backed by finufft on the NumPy path and nufftax on the JAX path.

This is the default TransformerNUFFT in PyAutoArray. The array namespace of each call picks the backend:

xp=np (eager NumPy) runs finufft (https://github.com/flatironinstitute/finufft), the CPU C++ NUFFT, via autoarray.operators.nufft_cpu. This path never imports JAX, so a NumPy fit (e.g. a numba Nautilus search with forked workers) does not start XLA in its parent process. finufft uses OpenMP and honours OMP_NUM_THREADS.
xp=jnp runs nufftax (https://github.com/GragasLab/nufftax), a pure-JAX port of finufft that supports jax.jit, jax.grad and jax.vmap. nufftax is imported only on this path.

The two agree to ~1e-14 relative. Where finufft has no wheel (e.g. Linux aarch64) the NumPy path falls back to nufftax with a warning; where neither is installed 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 the NUFFT’s mode-0 at index N // 2; neither backend applies 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 the backend. 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 NUFFT 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.
primary_beam – The primary beam P of the antennas, an Array2D on real_space_mask (or a native / slim array). When set, visibilities_from and transform_mapping_matrix multiply their input by P on the slim grid before the scatter to the native grid; image_from stays the un-beamed adjoint. None (default) applies no primary beam.

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/pyautolens/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.

If the transformer has a primary_beam_slim the image is multiplied by it on the slim grid before the scatter to the native grid, so the visibilities are those of P * image.

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

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

Implemented as a type-1 nufft2d1 (finufft via nufft_cpu for xp=np, nufftax for xp=jnp) 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.

The primary beam is deliberately not applied: this is the adjoint of the un-beamed transform, so the dirty image, dirty beam and the w-tilde quantities built from it do not depend on primary_beam_slim.

transform_mapping_matrix(mapping_matrix, xp=<module 'numpy' from '/home/docs/checkouts/readthedocs.org/user_builds/pyautolens/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).

If the transformer has a primary_beam_slim every column is multiplied by it on the slim grid before the scatter, giving F P M.

The numpy path runs finufft (nufft_cpu), the JAX path nufftax.