1.1.2.2. autojax.numba module¶

autojax.numba.constant_regularization_matrix_from(coefficient: float, neighbors: ~numpy.ndarray[[<class 'int'>, <class 'int'>], ~numpy.int64], neighbors_sizes: ~numpy.ndarray[[<class 'int'>], ~numpy.int64]) → float64][source]¶

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.numba.curvature_matrix_via_w_compact_sparse_mapping_matrix_direct_from(w_compact: ndarray[tuple[int, int], float64], native_index_for_slim_index: ndarray[tuple[int, int], int64], pix_indexes_for_sub_slim_index: ndarray[tuple[int, int], int64], pix_weights_for_sub_slim_index: ndarray[ndarray[tuple[int, int], float64]], pixels: int) → ndarray[tuple[int, int], float64][source]¶

Calculate the curvature matrix using the compact w_tilde matrix and the sparse mapping matrix.

This calculates it directly without expanding anything in memory. It optimizes for low memory usage but requires more FLOPS.

Memory cost: S^2 <- the output matrix size, i.e. no extra memory is used.

FLOP cost: (4 + 3B^2) M^2, B = pix_size_for_sub_slim_index.mean(), B=3 for Delaunay.

autojax.numba.curvature_matrix_via_w_compact_sparse_mapping_matrix_from(w_compact: ndarray[tuple[int, int], float64], native_index_for_slim_index: ndarray[tuple[int, int], int64], pix_indexes_for_sub_slim_index: ndarray[tuple[int, int], int64], pix_weights_for_sub_slim_index: ndarray[ndarray[tuple[int, int], float64]], pixels: int) → ndarray[tuple[int, int], float64][source]¶

Calculate the curvature matrix using the compact w_tilde matrix and the sparse mapping matrix.

This calculates T^T @ w_tilde @ T as two matrix multiplications, WT = w_tilde @ T and F = T^T @ WT.

This reduces FLOP requirements (by ~(4 + 3B^2)/(2(2 + B)) = 3.1-fold for Delaunay) at the cost of more memory usage by extra amount of MS.

Memory cost: MS + S^2

FLOP cost: 2(2 + B)M^2 + 2MBS, B = pix_size_for_sub_slim_index.mean(), B=3 for Delaunay.

autojax.numba.curvature_matrix_via_w_tilde_curvature_preload_interferometer_from(w_compact: ndarray[tuple[int, int], float64], native_index_for_slim_index: ndarray[tuple[int, int], int64], pix_indexes_for_sub_slim_index: ndarray[tuple[int, int], int64], pix_weights_for_sub_slim_index: ndarray[ndarray[tuple[int, int], float64]], pixels: int) → ndarray[tuple[int, int], float64]¶

Calculate the curvature matrix using the compact w_tilde matrix and the sparse mapping matrix.

This calculates T^T @ w_tilde @ T as two matrix multiplications, WT = w_tilde @ T and F = T^T @ WT.

This reduces FLOP requirements (by ~(4 + 3B^2)/(2(2 + B)) = 3.1-fold for Delaunay) at the cost of more memory usage by extra amount of MS.

Memory cost: MS + S^2

FLOP cost: 2(2 + B)M^2 + 2MBS, B = pix_size_for_sub_slim_index.mean(), B=3 for Delaunay.

autojax.numba.curvature_matrix_via_w_tilde_from(w_tilde: ndarray[tuple[int, int], float64], mapping_matrix: ndarray[tuple[int, int], float64]) → ndarray[tuple[int, int], float64][source]¶

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.numba.curvature_matrix_via_w_wilde_sparse_mapping_matrix_from(w_tilde: ndarray[tuple[int, int], float64], pix_indexes_for_sub_slim_index: ndarray[tuple[int, int], int64], pix_weights_for_sub_slim_index: ndarray[ndarray[tuple[int, int], float64]], pixels: int) → ndarray[tuple[int, int], float64][source]¶

Calculate the curvature matrix using the w_tilde matrix and the sparse mapping matrix.

Memory cost: MS + S^2

FLOP cost: 2BM(M + S), B = pix_size_for_sub_slim_index.mean(), B=3 for Delaunay.

autojax.numba.data_vector_from(mapping_matrix: ndarray[tuple[int, int], float64], dirty_image: ndarray[tuple[int], float64]) → ndarray[tuple[int], float64][source]¶

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.numba.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_compact: ndarray[tuple[int, int], float64], native_index_for_slim_index: ndarray[tuple[int, int], int64], pix_indexes_for_sub_slim_index: ndarray[tuple[int, int], int64], pix_weights_for_sub_slim_index: ndarray[ndarray[tuple[int, int], float64]], neighbors: ndarray[tuple[int, int], int64], neighbors_sizes: ndarray[tuple[int], int64]) → float[source]¶

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
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, B), dtype=int64) – The mappings from a data sub-pixel index to a pixelization pixel index.
pix_weights_for_sub_slim_index (ndarray, shape (M, B), dtype=float64) – The weights of the mappings of every data sub pixel and pixelization 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.numba.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_weights_for_sub_slim_index: ndarray[ndarray[tuple[int, int], float64]], neighbors: ndarray[tuple[int, int], int64], neighbors_sizes: ndarray[tuple[int], int64]) → float[source]¶

Calculates the log likelihood of interferometer data given a model.

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, B), dtype=int64) – The mappings from a data sub-pixel index to a pixelization pixel index.
pix_weights_for_sub_slim_index (ndarray, shape (M, B), dtype=float64) – The weights of the mappings of every data sub pixel and pixelization 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)

autojax.numba.mapping_matrix_from(pix_indexes_for_sub_slim_index: ndarray[tuple[int, int], int64], pix_weights_for_sub_slim_index: ndarray[ndarray[tuple[int, int], float64]], pixels: int) → ndarray[tuple[int, int], float64][source]¶

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, B), dtype=int64) – The mappings from a data sub-pixel index to a pixelization pixel index.
pix_weights_for_sub_slim_index (ndarray, shape (M, B), 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.

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

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.numba.mask_2d_circular_from(shape_native: tuple[int, int], pixel_scales: tuple[float, float], radius: float, centre: tuple[float, float]) → ndarray[source]¶

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.numba.noise_normalization_complex_from(noise_map: ~numpy.ndarray[[<class 'int'>], ~numpy.complex128]) → float[source]¶

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.numba.reconstruction_positive_negative_from(data_vector: ndarray[tuple[int], float64], curvature_reg_matrix: ndarray[tuple[int, int], float64]) → ndarray[tuple[int], float64][source]¶

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.numba.sparse_mapping_matrix_transpose_matmul(matrix: ndarray[tuple[int, int], float64], pix_indexes_for_sub_slim_index: ndarray[tuple[int, int], int64], pix_weights_for_sub_slim_index: ndarray[ndarray[tuple[int, int], float64]], pixels: int) → ndarray[tuple[int, int], float64][source]¶

Calculate T^T @ matrix using the sparse mapping matrix representation.

Assuming matrix is M by S2,

Memory cost: S1 S2 <- the output matrix size (no extra memory is used).

FLOP cost: 2M B S2, B = pix_size_for_sub_slim_index.mean(), B=3 for Delaunay.

autojax.numba.w_compact_curvature_interferometer_from(grid_size: int, noise_map_real: ndarray[tuple[int], float64], uv_wavelengths: ndarray[tuple[int, int], float64], pixel_scale: float) → ndarray[tuple[int, int], float64][source]¶

autojax.numba.w_tilde_curvature_interferometer_from(noise_map_real: ndarray[tuple[int], float64], uv_wavelengths: ndarray[tuple[int, int], float64], grid_radians_slim: ndarray[tuple[int, int], float64]) → ndarray[tuple[int, int], float64][source]¶

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.numba.w_tilde_data_interferometer_from(visibilities_real: ndarray[tuple[int], float64], noise_map_real: ndarray[tuple[int], float64], uv_wavelengths: ndarray[tuple[int, int], float64], grid_radians_slim: ndarray[tuple[int, int], float64]) → ndarray[tuple[int], float64][source]¶

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.numba.w_tilde_via_compact_from(w_compact: ndarray[tuple[int, int], float64], native_index_for_slim_index: ndarray[tuple[int, int], int64]) → ndarray[tuple[int, int], float64][source]¶

autojax.numba.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][source]¶

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