1.1.2.3. autojax.original module

autojax.original.constant_regularization_matrix_from(coefficient: float, neighbors: ndarray, neighbors_sizes: ndarray) → ndarray

From the pixel-neighbors array, setup the regularization matrix using the instance regularization scheme.

A complete description of regularizatin and the regularization_matrix can be found in the Regularization class in the module autoarray.inversion.regularization.

Parameters:

• coefficient – The regularization coefficients which controls the degree of smoothing of the inversion reconstruction.

• neighbors (ndarray, shape (S, P), dtype=int64) – An array of length (total_pixels) which provides the index of all neighbors of every pixel in the Voronoi grid (entries of -1 correspond to no neighbor).

• neighbors_sizes (ndarray, shape (S,), dtype=int64) – An array of length (total_pixels) which gives the number of neighbors of every pixel in the Voronoi grid.

Returns:

regularization_matrix – The regularization matrix computed using Regularization where the effective regularization coefficient of every source pixel is the same.

Return type:

ndarray, shape (S, S), dtype=float64

autojax.original.curvature_matrix_via_w_tilde_curvature_preload_interferometer_from(curvature_preload: ndarray[tuple[int, int], float64], pix_indexes_for_sub_slim_index: ndarray[tuple[int, int], int64], pix_size_for_sub_slim_index: ndarray[tuple[int], int64], pix_weights_for_sub_slim_index: ndarray[tuple[int, int], float64], native_index_for_slim_index: ndarray[tuple[int, int], int64], pix_pixels: int) → ndarray[tuple[int, int], float64]

Returns the curvature matrix F (see Warren & Dye 2003) by computing it using w_tilde_preload (see w_tilde_preload_interferometer_from) for an interferometer inversion.

To compute the curvature matrix via w_tilde the following matrix multiplication is normally performed:

curvature_matrix = mapping_matrix.T * w_tilde * mapping matrix

This function speeds this calculation up in two ways:

1. Instead of using w_tilde (dimensions [image_pixels, image_pixels] it uses w_tilde_preload (dimensions [image_pixels, 2]). The massive reduction in the size of this matrix in memory allows for much fast computation.

2. It omits the mapping_matrix and instead uses directly the 1D vector that maps every image pixel to a source pixel native_index_for_slim_index. This exploits the sparsity in the mapping_matrix to directly compute the curvature_matrix (e.g. it condenses the triple matrix multiplication into a double for loop!).

Parameters:

• curvature_preload (ndarray, shape (2N, 2N), dtype=float64) – A matrix that precomputes the values for fast computation of w_tilde, which in this function is used to bypass the creation of w_tilde altogether and go directly to the curvature_matrix.

• pix_indexes_for_sub_slim_index (ndarray, shape (M, 3), dtype=int64) – The mappings from a data sub-pixel index to a pixelization’s mesh pixel index.

• pix_size_for_sub_slim_index (ndarray, shape (M,), dtype=int64) – The number of mappings between each data sub pixel and pixelization pixel.

• pix_weights_for_sub_slim_index (ndarray, shape (M, 3), dtype=float64) – The weights of the mappings of every data sub pixel and pixelization pixel.

• native_index_for_slim_index (ndarray, shape (M, 2), dtype=int64) – An array of shape [total_unmasked_pixels*sub_size] that maps every unmasked sub-pixel to its corresponding native 2D pixel using its (y,x) pixel indexes.

• pix_pixels – The total number of pixels in the pixelization’s mesh that reconstructs the data. pix_pixels = S.

Returns:

curvature_matrix – The curvature matrix F (see Warren & Dye 2003).

Return type:

ndarray, shape (S, S), dtype=float64

autojax.original.curvature_matrix_via_w_tilde_from(w_tilde: ndarray, mapping_matrix: ndarray) → ndarray

Returns the curvature matrix F (see Warren & Dye 2003) from w_tilde.

The dimensions of w_tilde are [image_pixels, image_pixels], meaning that for datasets with many image pixels this matrix can take up 10’s of GB of memory. The calculation of the curvature_matrix via this function will therefore be very slow, and the method curvature_matrix_via_w_tilde_curvature_preload_imaging_from should be used instead.

Parameters:

• w_tilde (ndarray, shape (M, M), dtype=float64) – A matrix of dimensions [image_pixels, image_pixels] that encodes the convolution or NUFFT of every image pixel pair on the noise map.

• mapping_matrix (ndarray, shape (M, S), dtype=float64) – The matrix representing the mappings between sub-grid pixels and pixelization pixels.

Returns:

curvature_matrix – The curvature matrix F (see Warren & Dye 2003).

Return type:

ndarray, shape (S, S), dtype=float64

autojax.original.data_vector_from(mapping_matrix, dirty_image)

Parameters:

• mapping_matrix (ndarray, shape (M, S), dtype=float64) – Matrix representing mappings between sub-grid pixels and pixelization pixels

• dirty_image (ndarray, shape (M,), dtype=float64) – The dirty image used to compute the data vector

Returns:

• data_vector (ndarray, shape (S,), dtype=float64) – The data_vector is a 1D vector whose values are solved for by the simultaneous linear equations constructed by this object.

• The linear algebra is described in the paper https (//arxiv.org/pdf/astro-ph/0302587.pdf), where the)

• data vector is given by equation (4) and the letter D.

• If there are multiple linear objects the data_vectors are concatenated ensuring their values are solved

• for simultaneously.

• The calculation is described in more detail in inversion_util.w_tilde_data_interferometer_from.

autojax.original.log_likelihood_function_via_w_compact_from(dirty_image: ndarray[tuple[int], float64], data: ndarray[tuple[int], complex128], noise_map: ndarray[tuple[int], complex128], w_tilde_preload: ndarray[tuple[int, int], float64], pix_indexes_for_sub_slim_index: ndarray[tuple[int, int], int64], pix_size_for_sub_slim_index: ndarray[ndarray[tuple[int], int64]], pix_weights_for_sub_slim_index: ndarray[ndarray[tuple[int, int], float64]], slim_index_for_sub_slim_index: ndarray[tuple[int], int64], sub_fraction: ndarray[tuple[int], float64], neighbors: ndarray[tuple[int, int], int64], neighbors_sizes: ndarray[tuple[int], int64], native_index_for_slim_index: ndarray[tuple[int, int], int64]) → float

Calculates the log likelihood of interferometer data given a model.

This function combines several steps: 1. Calculates noise normalization from the complex noise map 2. Computes the curvature matrix using w_tilde and mapping matrix 3. Creates a regularization matrix using constant regularization 4. Solves for the reconstruction allowing positive and negative values 5. Combines terms to compute the final log likelihood

Parameters:

• dirty_image (ndarray, shape (M,), dtype=float64) – The dirty image used to compute the data vector

• data (ndarray, shape (K,), dtype=complex128) – The complex interferometer data being fitted

• noise_map (ndarray, shape (K,), dtype=complex128) – The complex noise map of the data

• shape_masked_pixels_2d (tuple(int, int)) – The (y,x) shape corresponding to the extent of unmasked pixels that go vertically and horizontally across the mask. E.g. (N, N), N = 30

• grid_radians_2d (ndarray, shape (N_PRIME, N_PRIME, 2), dtype=float64) – The 2D (y,x) grid of coordinates in radians corresponding to real-space mask within which the image that is Fourier transformed is computed. N_PRIME >= N. E.g. N_PRIME = 100

• uv_wavelengths (ndarray, shape (K, 2), dtype=float64) – The wavelengths of the coordinates in the uv-plane for the interferometer dataset

• pix_indexes_for_sub_slim_index (ndarray, shape (M, 3), dtype=int64) – The mappings from a data sub-pixel index to a pixelization pixel index.

• pix_size_for_sub_slim_index (ndarray, shape (M,), dtype=int64) – The number of mappings between each data sub pixel and pixelization pixel.

• pix_weights_for_sub_slim_index (ndarray, shape (M, 3), dtype=float64) – The weights of the mappings of every data sub pixel and pixelization pixel.

• slim_index_for_sub_slim_index (ndarray, shape (M,), dtype=int64) – The mappings between the data’s sub slimmed indexes and the slimmed indexes on the non sub-sized indexes.

• sub_fraction (ndarray, shape (M,), dtype=float64) – The fractional area each sub-pixel takes up in an pixel.

• neighbors (ndarray, shape (S, P), dtype=int64) – Array providing indices of neighbors for each pixel

• neighbors_sizes (ndarray, shape (S,), dtype=int64) – Array giving number of neighbors for each pixel

• native_index_for_slim_index (ndarray, shape (M, 2), dtype=int64) – An array of shape [total_unmasked_pixels*sub_size] that maps every unmasked sub-pixel to its corresponding native 2D pixel using its (y,x) pixel indexes.

Returns:

The log likelihood value of the model fit to the data

Return type:

float

Notes

The log likelihood calculation follows the formalism described in Warren & Dye 2003 (https://arxiv.org/pdf/astro-ph/0302587.pdf) with additional terms for interferometric data.

Typical sizes: (716 -> 70000 means 716 in the test dataset, but can go up to 70000 in science case)

M = number of image pixels in real_space_mask = 716 -> ~70000 N = The (y,x) shape corresponding to the extent of unmasked pixels that go vertically and horizontally across the mask. E.g. (N, N), N = 30. M ~ N^2 as M only counts unmasked pixels. K = number of visibilitiies = 190 -> ~1e7 (but this is only used to compute w_tilde otuside the likelihood function) P = number of neighbors = 10 -> 3 (for Delaunay) but can go up to 300 for Voronoi (but we can just focus on delaunay for now) S = number of source pixels (e.g. reconstruction.shape) = 716 -> 1000

autojax.original.log_likelihood_function_via_w_tilde_from(dirty_image: ndarray[tuple[int], float64], data: ndarray[tuple[int], complex128], noise_map: ndarray[tuple[int], complex128], w_tilde: ndarray[tuple[int, int], float64], pix_indexes_for_sub_slim_index: ndarray[tuple[int, int], int64], pix_size_for_sub_slim_index: ndarray[ndarray[tuple[int], int64]], pix_weights_for_sub_slim_index: ndarray[ndarray[tuple[int, int], float64]], slim_index_for_sub_slim_index: ndarray[tuple[int], int64], sub_fraction: ndarray[tuple[int], float64], neighbors: ndarray[tuple[int, int], int64], neighbors_sizes: ndarray[tuple[int], int64]) → float

Calculates the log likelihood of interferometer data given a model.

This function combines several steps: 1. Calculates noise normalization from the complex noise map 2. Computes the curvature matrix using w_tilde and mapping matrix 3. Creates a regularization matrix using constant regularization 4. Solves for the reconstruction allowing positive and negative values 5. Combines terms to compute the final log likelihood

Parameters:

• dirty_image (ndarray, shape (M,), dtype=float64) – The dirty image used to compute the data vector

• data (ndarray, shape (K,), dtype=complex128) – The complex interferometer data being fitted

• noise_map (ndarray, shape (K,), dtype=complex128) – The complex noise map of the data

• uv_wavelengths (ndarray, shape (K, 2), dtype=float64) – The wavelengths of the coordinates in the uv-plane for the interferometer dataset

• grid_radians_slim (ndarray, shape (M, 2), dtype=float64) – The 1D (y,x) grid of coordinates in radians corresponding to real-space mask

• pix_indexes_for_sub_slim_index (ndarray, shape (M, 3), dtype=int64) – The mappings from a data sub-pixel index to a pixelization pixel index.

• pix_size_for_sub_slim_index (ndarray, shape (M,), dtype=int64) – The number of mappings between each data sub pixel and pixelization pixel.

• pix_weights_for_sub_slim_index (ndarray, shape (M, 3), dtype=float64) – The weights of the mappings of every data sub pixel and pixelization pixel.

• slim_index_for_sub_slim_index (ndarray, shape (M,), dtype=int64) – The mappings between the data’s sub slimmed indexes and the slimmed indexes on the non sub-sized indexes.

• sub_fraction (ndarray, shape (M,), dtype=float64) – The fractional area each sub-pixel takes up in an pixel.

• neighbors (ndarray, shape (S, P), dtype=int64) – Array providing indices of neighbors for each pixel

• neighbors_sizes (ndarray, shape (S,), dtype=int64) – Array giving number of neighbors for each pixel

Returns:

The log likelihood value of the model fit to the data

Return type:

float

Notes

The log likelihood calculation follows the formalism described in Warren & Dye 2003 (https://arxiv.org/pdf/astro-ph/0302587.pdf) with additional terms for interferometric data.

Typical sizes: (716 -> 70000 means 716 in the test dataset, but can go up to 70000 in science case)

M = number of image pixels in real_space_mask = 716 -> ~70000 K = number of visibilitiies = 190 -> ~1e7 (but this is only used to compute w_tilde otuside the likelihood function) P = number of neighbors = 10 -> 3 (for Delaunay) but can go up to 300 for Voronoi (but we can just focus on delaunay for now) S = number of source pixels (e.g. reconstruction.shape) = 716 -> 1000

autojax.original.mapping_matrix_from(pix_indexes_for_sub_slim_index: ndarray[tuple[int, int], int64], pix_size_for_sub_slim_index: ndarray[ndarray[tuple[int], int64]], pix_weights_for_sub_slim_index: ndarray[ndarray[tuple[int, int], float64]], pixels: int, total_mask_pixels: int, slim_index_for_sub_slim_index: ndarray[tuple[int], int64], sub_fraction: ndarray[tuple[int], float64]) → ndarray[tuple[int, int], float64]

Returns the mapping matrix, which is a matrix representing the mapping between every unmasked sub-pixel of the data and the pixels of a pixelization. Non-zero entries signify a mapping, whereas zeros signify no mapping.

For example, if the data has 5 unmasked pixels (with sub_size=1 so there are not sub-pixels) and the pixelization 3 pixels, with the following mappings:

data pixel 0 -> pixelization pixel 0 data pixel 1 -> pixelization pixel 0 data pixel 2 -> pixelization pixel 1 data pixel 3 -> pixelization pixel 1 data pixel 4 -> pixelization pixel 2

The mapping matrix (which is of dimensions [data_pixels, pixelization_pixels]) would appear as follows:

[1, 0, 0] [0->0] [1, 0, 0] [1->0] [0, 1, 0] [2->1] [0, 1, 0] [3->1] [0, 0, 1] [4->2]

The mapping matrix is actually built using the sub-grid of the grid, whereby each pixel is divided into a grid of sub-pixels which are all paired to pixels in the pixelization. The entries in the mapping matrix now become fractional values dependent on the sub-pixel sizes.

For example, for a 2x2 sub-pixels in each pixel means the fractional value is 1.0/(2.0^2) = 0.25, if we have the following mappings:

data pixel 0 -> data sub pixel 0 -> pixelization pixel 0 data pixel 0 -> data sub pixel 1 -> pixelization pixel 1 data pixel 0 -> data sub pixel 2 -> pixelization pixel 1 data pixel 0 -> data sub pixel 3 -> pixelization pixel 1 data pixel 1 -> data sub pixel 0 -> pixelization pixel 1 data pixel 1 -> data sub pixel 1 -> pixelization pixel 1 data pixel 1 -> data sub pixel 2 -> pixelization pixel 1 data pixel 1 -> data sub pixel 3 -> pixelization pixel 1 data pixel 2 -> data sub pixel 0 -> pixelization pixel 2 data pixel 2 -> data sub pixel 1 -> pixelization pixel 2 data pixel 2 -> data sub pixel 2 -> pixelization pixel 3 data pixel 2 -> data sub pixel 3 -> pixelization pixel 3

The mapping matrix (which is still of dimensions [data_pixels, pixelization_pixels]) appears as follows:

[0.25, 0.75, 0.0, 0.0] [1 sub-pixel maps to pixel 0, 3 map to pixel 1] [ 0.0, 1.0, 0.0, 0.0] [All sub-pixels map to pixel 1] [ 0.0, 0.0, 0.5, 0.5] [2 sub-pixels map to pixel 2, 2 map to pixel 3]

For certain pixelizations each data sub-pixel maps to multiple pixelization pixels in a weighted fashion, for example a Delaunay pixelization where there are 3 mappings per sub-pixel whose weights are determined via a nearest neighbor interpolation scheme.

In this case, each mapping value is multiplied by this interpolation weight (which are in the array pix_weights_for_sub_slim_index) when the mapping matrix is constructed.

Parameters:

• pix_indexes_for_sub_slim_index (ndarray, shape (M, 3), dtype=int64) – The mappings from a data sub-pixel index to a pixelization pixel index.

• pix_size_for_sub_slim_index (ndarray, shape (M,), dtype=int64) – The number of mappings between each data sub pixel and pixelization pixel.

• pix_weights_for_sub_slim_index (ndarray, shape (M, 3), dtype=float64) – The weights of the mappings of every data sub pixel and pixelization pixel.

• pixels – The number of pixels in the pixelization.

• total_mask_pixels – The number of datas pixels in the observed datas and thus on the grid.

• slim_index_for_sub_slim_index (ndarray, shape (M,), dtype=int64) – The mappings between the data’s sub slimmed indexes and the slimmed indexes on the non sub-sized indexes.

• sub_fraction (ndarray, shape (M,), dtype=float64) – The fractional area each sub-pixel takes up in an pixel.

autojax.original.mask_2d_centres_from(shape_native: tuple[int, int], pixel_scales: tuple[float, float], centre: tuple[float, float]) → tuple[float, float]

Returns the (y,x) scaled central coordinates of a mask from its shape, pixel-scales and centre.

The coordinate system is defined such that the positive y axis is up and positive x axis is right.

Parameters:

• shape_native – The (y,x) shape of the 2D array the scaled centre is computed for.

• pixel_scales – The (y,x) scaled units to pixel units conversion factor of the 2D array.

• centre ((float, flloat)) – The (y,x) centre of the 2D mask.

Returns:

The (y,x) scaled central coordinates of the input array.

Return type:

tuple (float, float)

Examples

centres_scaled = centres_from(shape=(5,5), pixel_scales=(0.5, 0.5), centre=(0.0, 0.0))

autojax.original.mask_2d_circular_from(shape_native: tuple[int, int], pixel_scales: tuple[float, float], radius: float, centre: tuple[float, float] = (0.0, 0.0)) → ndarray

Returns a circular mask from the 2D mask array shape and radius of the circle.

This creates a 2D array where all values within the mask radius are unmasked and therefore False.

Parameters:

• shape_native (Tuple[int, int]) – The (y,x) shape of the mask in units of pixels. Usually this is (N_PRIME, N_PRIME).

• pixel_scales – The scaled units to pixel units conversion factor of each pixel, in arcsec.

• radius – The radius (in scaled units) of the circle within which pixels unmasked, in arcsec.

• centre – The centre of the circle used to mask pixels, in arcsec.

Returns:

The 2D mask array whose central pixels are masked as a circle.

Return type:

ndarray

Examples

mask = mask_circular_from(

shape=(10, 10), pixel_scales=0.1, radius=0.5, centre=(0.0, 0.0))

autojax.original.noise_normalization_complex_from(noise_map: ndarray) → float

Returns the noise-map normalization terms of a complex noise-map, summing the noise_map value in every pixel as:

[Noise_Term] = sum(log(2*pi*[Noise]**2.0))

Parameters:

noise_map (ndarray, shape (K,), dtype=complex128) – The masked noise-map of the dataset.

Returns:

noise_normalization

Return type:

float

autojax.original.reconstruction_positive_negative_from(data_vector: ndarray, curvature_reg_matrix: ndarray)

Solve the linear system [F + reg_coeff*H] S = D -> S = [F + reg_coeff*H]^-1 D given by equation (12) of https://arxiv.org/pdf/astro-ph/0302587.pdf

S is the vector of reconstructed inversion values.

This reconstruction uses a linear algebra solver that allows for negative and positives values in the solution. By allowing negative values, the solver is efficient, but there are many inference problems where negative values are nonphysical or undesirable.

This function checks that the solution does not give a linear algebra error (e.g. because the input matrix is not positive-definitive).

It also explicitly checks solutions where all reconstructed values go to the same value, and raises an exception if this occurs. This solution occurs in many scenarios when it is clear not a valid solution, and therefore is checked for and removed.

Parameters:

• data_vector (ndarray, shape (S,), dtype=float64) – The data_vector D which is solved for.

• curvature_reg_matrix (ndarray, shape (S, S), dtype=float64) – The sum of the curvature and regularization matrices.

Returns:

reconstruction

Return type:

ndarray, shape (S,), dtype=float64

autojax.original.w_tilde_curvature_interferometer_from(noise_map_real: ndarray, uv_wavelengths: ndarray, grid_radians_slim: ndarray) → ndarray

The matrix w_tilde is a matrix of dimensions [image_pixels, image_pixels] that encodes the NUFFT of every pair of image pixels given the noise map. This can be used to efficiently compute the curvature matrix via the mappings between image and source pixels, in a way that omits having to perform the NUFFT on every individual source pixel. This provides a significant speed up for inversions of interferometer datasets with large number of visibilities.

The limitation of this matrix is that the dimensions of [image_pixels, image_pixels] can exceed many 10s of GB’s, making it impossible to store in memory and its use in linear algebra calculations extremely. The method w_tilde_preload_interferometer_from describes a compressed representation that overcomes this hurdles. It is advised w_tilde and this method are only used for testing.

Note that the current implementation does not take advantage of the fact that w_tilde is symmetric, due to the use of vectorized operations.

\[\tilde{W}_{ij} = \sum_{k=1}^N \frac{1}{n_k^2} \cos(2\pi[(g_{i1} - g_{j1})u_{k0} + (g_{i0} - g_{j0})u_{k1}])\]

The function is written in a way that the memory use does not depend on size of data K.

Parameters:

• noise_map_real (ndarray, shape (K,), dtype=float64) – The real noise-map values of the interferometer data.

• uv_wavelengths (ndarray, shape (K, 2), dtype=float64) – The wavelengths of the coordinates in the uv-plane for the interferometer dataset that is to be Fourier transformed.

• grid_radians_slim (ndarray, shape (M, 2), dtype=float64) – The 1D (y,x) grid of coordinates in radians corresponding to real-space mask within which the image that is Fourier transformed is computed.

Returns:

curvature_matrix – A matrix that encodes the NUFFT values between the noise map that enables efficient calculation of the curvature matrix.

Return type:

ndarray, shape (M, M), dtype=float64

autojax.original.w_tilde_curvature_preload_interferometer_from(noise_map_real: ndarray[tuple[int], float64], uv_wavelengths: ndarray[tuple[int, int], float64], shape_masked_pixels_2d: tuple[int, int], grid_radians_2d: ndarray[tuple[int, int, int], float64]) → ndarray[tuple[int, int], float64]

Computes a preload matrix for efficient calculation of the NUFFT w_tilde matrix.

The matrix w_tilde is a matrix of dimensions [unmasked_image_pixels, unmasked_image_pixels] that encodes the NUFFT of every pair of image pixels given the noise map. This can be used to efficiently compute the curvature matrix via the mapping matrix, in a way that omits having to perform the NUFFT on every individual source pixel. This provides a significant speed up for inversions of interferometer datasets with large number of visibilities.

The limitation of this matrix is that the dimensions of [image_pixels, image_pixels] can exceed many 10s of GB’s, making it impossible to store in memory and its use in linear algebra calculations extremely expensive. This method creates a preload matrix that can compute the matrix w_tilde via an efficient preloading scheme which exploits the symmetries in the NUFFT.

To compute w_tilde, one first defines a real space mask where every False entry is an unmasked pixel which is used in the calculation, for example:

Mask Example:

x | x | x | x | x | x | x | x | x | x |

x | x | x | x | x | x | x | x | x | x | This is an imaging.Mask2D, where:

x | x | x | x | x | x | x | x | x | x |

x | x | x | x | x | x | x | x | x | x | x = True (Pixel is masked and excluded from lens)

x | x | x | o | o | o | x | x | x | x | o = False (Pixel is not masked and included in lens)

x | x | x | o | o | o | x | x | x | x |

x | x | x | o | o | o | x | x | x | x |

x | x | x | x | x | x | x | x | x | x |

x | x | x | x | x | x | x | x | x | x |

x | x | x | x | x | x | x | x | x | x |

Here, there are 9 unmasked pixels. Indexing of each unmasked pixel goes from the top-left corner right and downwards, therefore:

Indexing Example:

x | x | x | x | x | x | x | x | x | x |

x | x | x | x | x | x | x | x | x | x |

x | x | x | x | x | x | x | x | x | x |

x | x | x | x | x | x | x | x | x | x |

x | x | x | 0 | 1 | 2 | x | x | x | x |

x | x | x | 3 | 4 | 5 | x | x | x | x |

x | x | x | 6 | 7 | 8 | x | x | x | x |

x | x | x | x | x | x | x | x | x | x |

x | x | x | x | x | x | x | x | x | x |

x | x | x | x | x | x | x | x | x | x |

In the standard calculation of w_tilde it is a matrix of dimensions [unmasked_image_pixels, unmasked_pixel_images], therefore for the example mask above it would be dimensions [9, 9]. One performs a double for loop over unmasked_image_pixels, using the (y,x) spatial offset between every possible pair of unmasked image pixels to precompute values that depend on the properties of the NUFFT.

This calculation has a lot of redundancy, because it uses the (y,x) spatial offset between the image pixels. For example, if two image pixels are next to one another by the same spacing the same value will be computed via the NUFFT. For the example mask above:

• The value precomputed for pixel pair [0,1] is the same as pixel pairs [1,2], [3,4], [4,5], [6,7] and [7,8].

• The value precomputed for pixel pair [0,3] is the same as pixel pairs [1,4], [2,5], [3,6], [4,7] and [5,8].

• The values of pixels paired with themselves are also computed repeatedly for the standard calculation (e.g. 9 times using the mask above).

The w_tilde_preload method instead only computes each value once. To do this, it stores the preload values in a matrix of dimensions [shape_masked_pixels_y, shape_masked_pixels_x, 2], where shape_masked_pixels is the (y,x) size of the vertical and horizontal extent of unmasked pixels, e.g. the spatial extent over which the real space grid extends.

Each entry in the matrix w_tilde_preload[:,:,0] provides the precomputed NUFFT value mapping an image pixel to a pixel offset by that much in the y and x directions, for example:

• w_tilde_preload[0,0,0] gives the precomputed values of image pixels that are offset in the y direction by 0 and in the x direction by 0 - the values of pixels paired with themselves.

• w_tilde_preload[1,0,0] gives the precomputed values of image pixels that are offset in the y direction by 1 and in the x direction by 0 - the values of pixel pairs [0,3], [1,4], [2,5], [3,6], [4,7] and [5,8]

• w_tilde_preload[0,1,0] gives the precomputed values of image pixels that are offset in the y direction by 0 and in the x direction by 1 - the values of pixel pairs [0,1], [1,2], [3,4], [4,5], [6,7] and [7,8].

Flipped pairs: The above preloaded values pair all image pixel NUFFT values when a pixel is to the right and/or down of the first image pixel. However, one must also precompute pairs where the paired pixel is to the left of the host pixels. These pairings are stored in w_tilde_preload[:,:,1], and the ordering of these pairings is flipped in the x direction to make it straight forward to use this matrix when computing w_tilde.

Parameters:

• noise_map_real (ndarray) – The real noise-map values of the interferometer data. Shape: (K,) dtype: float64

• uv_wavelengths (ndarray) – The wavelengths of the coordinates in the uv-plane for the interferometer dataset that is to be Fourier transformed. Shape: (K, 2) dtype: float64

• shape_masked_pixels_2d (tuple) – The (y,x) shape corresponding to the extent of unmasked pixels that go vertically and horizontally across the mask. E.g. (N, N), N = 30

• grid_radians_2d (ndarray) – The 2D (y,x) grid of coordinates in radians corresponding to real-space mask within which the image that is Fourier transformed is computed. N_PRIME >= N. E.g. N_PRIME = 100 Shape: (N_PRIME, N_PRIME, 2) dtype: float64

Returns:

A matrix that precomputes the values for fast computation of w_tilde. Shape: (2N, 2N) dtype: float64

Return type:

ndarray

autojax.original.w_tilde_data_interferometer_from(visibilities_real: ndarray, noise_map_real: ndarray, uv_wavelengths: ndarray, grid_radians_slim: ndarray, native_index_for_slim_index) → ndarray

The matrix w_tilde is a matrix of dimensions [image_pixels, image_pixels] that encodes the PSF convolution of every pair of image pixels given the noise map. This can be used to efficiently compute the curvature matrix via the mappings between image and source pixels, in a way that omits having to perform the PSF convolution on every individual source pixel. This provides a significant speed up for inversions of imaging datasets.

When w_tilde is used to perform an inversion, the mapping matrices are not computed, meaning that they cannot be used to compute the data vector. This method creates the vector w_tilde_data which allows for the data vector to be computed efficiently without the mapping matrix.

The matrix w_tilde_data is dimensions [image_pixels] and encodes the PSF convolution with the weight_map, where the weights are the image-pixel values divided by the noise-map values squared:

weight = image / noise**2.0

\[\tilde{w}_{\text{data},i} = \sum_{j=1}^N \left(\frac{N_{r,j}^2}{V_{r,j}}\right)^2 \cos\left(2\pi(g_{i,1}u_{j,0} + g_{i,0}u_{j,1})\right)\]

The function is written in a way that the memory use does not depend on size of data K.

Parameters:

• visibilities_real (ndarray, shape (K,), dtype=float64) – The two dimensional masked image of values which w_tilde_data is computed from.

• noise_map_real (ndarray, shape (K,), dtype=float64) – The two dimensional masked noise-map of values which w_tilde_data is computed from.

• uv_wavelengths (ndarray, shape (K, 2), dtype=float64)

• grid_radians_slim (ndarray, shape (M, 2), dtype=float64)

• native_index_for_slim_index (ndarray, shape (M, 2), dtype=int64) – An array that maps pixels from the slimmed array to the native array.

Returns:

A matrix that encodes the PSF convolution values between the imaging divided by the noise map**2 that enables efficient calculation of the data vector.

Return type:

ndarray, shape (M,), dtype=float64

autojax.original.w_tilde_via_preload_from(w_tilde_preload: ndarray[tuple[int, int], float64], native_index_for_slim_index: ndarray[tuple[int, int], int64]) → ndarray[tuple[int, int], float64]

Use the preloaded w_tilde matrix (see w_tilde_preload_interferometer_from) to compute w_tilde (see w_tilde_interferometer_from) efficiently.

Parameters:

• w_tilde_preload (ndarray, shape (2N, 2N), dtype=float64) – The preloaded values of the NUFFT that enable efficient computation of w_tilde.

• native_index_for_slim_index (ndarray, shape (M, 2), dtype=int64) – An array of shape [total_unmasked_pixels*sub_size] that maps every unmasked sub-pixel to its corresponding native 2D pixel using its (y,x) pixel indexes.

Returns:

ndarray – A matrix that encodes the NUFFT values between the noise map that enables efficient calculation of the curvature matrix.

Return type:

shape (M, M), dtype=float64