jaxtronomy.ImSim.Numerics package

Submodules

jaxtronomy.ImSim.Numerics.convolution module

class PixelKernelConvolution(kernel=None, convolution_type='fft')[source]

Bases: object

Class to compute convolutions for a given pixelized kernel (fft, grid).

convolution_type=”fft_static” does not result in a speedup in JAXtronomy, thus it is not implemented.

__init__(kernel=None, convolution_type='fft')[source]
Parameters:

  • kernel – 2d array, convolution kernel, can also be supplied at runtime

  • convolution_type – string, ‘fft’, ‘grid’, mode of 2d convolution

pixel_kernel(num_pix=None)[source]

Access pixelated kernel.

Parameters:

num_pix – size of returned kernel (odd number per axis). If None, return the original kernel.

Returns:

pixel kernel centered

copy_transpose()[source]
Returns:

copy of the class with kernel set to the transpose of original one

convolution2d(image, kernel=None)[source]
Parameters:

image – 2d array (image) to be convolved

Returns:

fft convolution

re_size_convolve(image_low_res, image_high_res=None)[source]
Parameters:

  • image_low_res – regular sampled image/model

  • image_high_res – supersampled image/model to be convolved on a regular pixel grid

Returns:

convolved and re-sized image

class SubgridKernelConvolution(kernel_supersampled, supersampling_factor, supersampling_kernel_size=None, convolution_type='fft')[source]

Bases: object

Class to compute the convolution on a supersampled grid with partial convolution computed on the regular grid.

For high resolution convolution, one single convolution is performed using the high resolution kernel and high resolution image.

__init__(kernel_supersampled, supersampling_factor, supersampling_kernel_size=None, convolution_type='fft')[source]
Parameters:

  • kernel_supersampled – kernel in supersampled pixels

  • supersampling_factor – supersampling factor relative to the image pixel grid

  • supersampling_kernel_size – number of pixels (in units of the image pixels) that are convolved with the supersampled kernel

convolution2d(image)[source]
Parameters:

image – 2d array (high resoluton image) to be convolved and re-sized

Returns:

convolved image

re_size_convolve(image_low_res, image_high_res)[source]
Parameters:

image_high_res – supersampled image/model to be convolved on a regular pixel grid

Returns:

convolved and re-sized image

class PartialSubgridKernelConvolution(kernel_supersampled, supersampling_factor, supersampling_kernel_size=None, convolution_type='fft')[source]

Bases: object

Class to compute the convolution on a supersampled grid with partial convolution computed on the regular grid.

For high resolution convolution, rather than performing one single convolution over the entire image, N^2 low-resolution convolutions are performed and summed, where N is the supersampling factor. On CPU, this is a performance boost over performing one high resolution convolution. This class is effectively the same as lenstronomy’s SubgridNumbaConvolution, but with FFT instead of real space convolution.

__init__(kernel_supersampled, supersampling_factor, supersampling_kernel_size=None, convolution_type='fft')[source]
Parameters:

  • kernel_supersampled – kernel in supersampled pixels

  • supersampling_factor – supersampling factor relative to the image pixel grid

  • supersampling_kernel_size – number of pixels (in units of the image pixels) that are convolved with the supersampled kernel

re_size_convolve(image_low_res, image_high_res)[source]
Parameters:

  • image_low_res – 2d array, image at native resolution

  • image_high_res – 2d array, supersampled image

Returns:

convolved and re-sized image

convolution2d(image)[source]
Parameters:

image – 2d array (high resoluton image) to be convolved and re-sized

Returns:

convolved image

class GaussianConvolution(sigma, pixel_scale, supersampling_factor=1, supersampling_convolution=False, truncation=2)[source]

Bases: object

Class to perform a convolution consisting of multiple 2d Gaussians.

Since JAX does not have an ndimage.gaussian_filter function, to perform Gaussian convolutions, we first create Gaussian psf kernels and convolve them using fftconvolve.

__init__(sigma, pixel_scale, supersampling_factor=1, supersampling_convolution=False, truncation=2)[source]
Parameters:

  • sigma – std value of Gaussian kernel

  • pixel_scale – scale of pixel width (to convert sigmas into units of pixels)

  • supersampling_factor – int, ratio of the number of pixels of the high resolution grid to the number of pixels of the original image

  • convolution (supersampling) – bool, determines whether convolution uses supersampled grid or not

  • truncation – float. Truncate the filter at this many standard deviations. Default is 4.0.

convolution2d(image)[source]

Convolve the image via FFT convolution.

Parameters:

image – 2d numpy array, image to be convolved

Returns:

convolved image, 2d numpy array

re_size_convolve(image_low_res, image_high_res)[source]
Parameters:

image_high_res – supersampled image/model to be convolved on a regular pixel grid

Returns:

convolved and re-sized image

pixel_kernel(num_pix)[source]

Computes a pixelized kernel from the Gaussian parameters.

Parameters:

num_pix – int, size of kernel (odd number per axis) should be equal to 2 * sigma_scaled * truncation + 1 to be consistent with scipy.ndimage.gaussian_filter

Returns:

pixel kernel centered

jaxtronomy.ImSim.Numerics.grid module

class AdaptiveGrid(nx, ny, transform_pix2angle, ra_at_xy_0, dec_at_xy_0, supersampling_indexes, supersampling_factor, flux_evaluate_indexes=None)[source]

Bases: Coordinates1D

Manages a super-sampled grid on the partial image.

NOTE: The implementation of this class differs from lenstronomy. In lenstronomy, the high resolution image only includes the part of the image given by supersampling_indexes. In JAXtronomy, the high resolution image contains the entire image, with the non-supersampled pixels having their values duplicated along the subpixels.

This difference is due to the fact that JAXtronomy does not implement AdaptiveConvolution, since partial real-space convolution using Numba is not available.

__init__(nx, ny, transform_pix2angle, ra_at_xy_0, dec_at_xy_0, supersampling_indexes, supersampling_factor, flux_evaluate_indexes=None)[source]
Parameters:

  • nx – number of pixels in x-axis

  • ny – number of pixels in y-axis

  • transform_pix2angle – 2x2 matrix, mapping of pixel to coordinate

  • ra_at_xy_0 – ra coordinate at pixel (0,0)

  • dec_at_xy_0 – dec coordinate at pixel (0,0)

  • supersampling_indexes – bool array of shape nx x ny, corresponding to pixels being super_sampled

  • supersampling_factor – int, factor (per axis) of super-sampling

  • flux_evaluate_indexes – bool array of shape nx x ny, corresponding to pixels being evaluated (for both low and high res). Default is None, replaced by setting all pixels to being evaluated.

property coordinates_evaluate
Returns:

1d array of all coordinates being evaluated to perform the image computation

property supersampling_factor
Returns:

factor (per axis) of super-sampling relative to a pixel

flux_array2image_low_high(flux_array, high_res_return=True)[source]

Returns the low res and high res images. The high res implementation differs from lenstronomy in that non-supersampled pixels also have their values set (in lenstronomy they are zero).

Parameters:

  • flux_array – 1d array of low and high resolution flux values corresponding to the coordinates_evaluate order

  • high_res_return – bool, if True also returns the high resolution image (needs more computation and is only needed when convolution is performed on the supersampling level)

Returns:

2d array, 2d array, corresponding to images in low and high resolution (to be convolved)

class RegularGrid(nx, ny, transform_pix2angle, ra_at_xy_0, dec_at_xy_0, supersampling_factor=1, flux_evaluate_indexes=None)[source]

Bases: Coordinates1D

Manages a super-sampled grid on the partial image.

__init__(nx, ny, transform_pix2angle, ra_at_xy_0, dec_at_xy_0, supersampling_factor=1, flux_evaluate_indexes=None)[source]
Parameters:

  • nx – number of pixels in x-axis

  • ny – number of pixels in y-axis

  • transform_pix2angle – 2x2 matrix, mapping of pixel to coordinate

  • ra_at_xy_0 – ra coordinate at pixel (0,0)

  • dec_at_xy_0 – dec coordinate at pixel (0,0)

  • supersampling_factor – int, factor (per axis) of super-sampling

  • flux_evaluate_indexes – bool array of shape nx x ny, corresponding to pixels being evaluated (for both low and high res). Default is None, replaced by setting all pixels to being evaluated. This cannot be a jnp array; must be a np array

property coordinates_evaluate
Returns:

1d array of all coordinates being evaluated to perform the image computation

property grid_points_spacing

Effective spacing between coordinate points, after supersampling.

Returns:

sqrt(pixel_area)/supersampling_factor.

property num_grid_points_axes

Effective number of points along each axes, after supersampling.

Returns:

number of pixels per axis, nx*supersampling_factor ny*supersampling_factor

property supersampling_factor
Returns:

factor (per axis) of super-sampling relative to a pixel

flux_array2image_low_high(flux_array, **kwargs)[source]
Parameters:

flux_array – 1d array of low and high resolution flux values corresponding to the coordinates_evaluate order

Returns:

2d array, 2d array, corresponding to (partial) images in low and high resolution (to be convolved)

jaxtronomy.ImSim.Numerics.numerics module

class Numerics(pixel_grid, psf, supersampling_factor=1, compute_mode='regular', supersampling_convolution=False, supersampling_kernel_size=5, flux_evaluate_indexes=None, supersampled_indexes=None, compute_indexes=None, point_source_supersampling_factor=1, convolution_kernel_size=None, convolution_type='fft', truncation_conv=None, backend=None)[source]

Bases: PointSourceRendering

This classes manages the numerical options and computations of an image.

The class has two main functions, re_size_convolve() and coordinates_evaluate()

__init__(pixel_grid, psf, supersampling_factor=1, compute_mode='regular', supersampling_convolution=False, supersampling_kernel_size=5, flux_evaluate_indexes=None, supersampled_indexes=None, compute_indexes=None, point_source_supersampling_factor=1, convolution_kernel_size=None, convolution_type='fft', truncation_conv=None, backend=None)[source]
Parameters:

  • pixel_grid – PixelGrid() class instance

  • psf – PSF() class instance

  • compute_mode – options are: ‘regular’, ‘adaptive’. NOTE: adaptive compute mode differs from lenstronomy in that only the ray shooting is done adaptively in JAXtronomy whereas lenstronomy also performs convolution adaptively. Additionally, while lenstronomy’s adaptive compute mode is incompatible with “GAUSSIAN” psf type, in JAXtronomy they are compatible.

  • supersampling_factor – int, factor of higher resolution sub-pixel sampling of surface brightness

  • supersampling_convolution – bool, if True, performs (part of) the convolution on the super-sampled grid/pixels

  • supersampling_kernel_size – int (odd number), size (in regular pixel units) of the super-sampled convolution

  • flux_evaluate_indexes – boolean 2d array of size of image before supersampling (or None, then initiated as gird of True’s). Pixels indicated with True will be used to perform the surface brightness computation (and possible lensing ray-shooting). Pixels marked as False will be assigned a flux value of zero (or ignored in the adaptive convolution)

  • supersampled_indexes – 2d boolean array (only used in mode=’adaptive’) of pixels to be supersampled for ray shooting. All other pixels set to False will not be supersampled for ray shooting

  • compute_indexes – unused in JAXtronomy due to lack of adaptive convolution

  • point_source_supersampling_factor – super-sampling resolution of the point source placing

  • convolution_kernel_size – int, odd number, size of convolution kernel before supersampling. If None, takes size of point_source_kernel Only relevant for psf type PIXEL

  • convolution_type – string, ‘fft’, ‘grid’, ‘fft_static’ mode of 2d convolution

  • truncation_conv – Truncation used for the construction of the convolution kernels (only relevant for Gaussian convolution). By default, the truncation from the psf class will be used. Can be overwritten so that different PSFs are used for convolution and point source rendering.

  • backend – “gpu” or “cpu”. If None, calls jax.default_backend(). If “gpu”, high resolution FFT is done once on the high resolution grid, and if “cpu”, splits it up into many low resolution FFTs

re_size_convolve(flux_array, unconvolved=False)[source]
Parameters:

  • flux_array – 1d array, flux values corresponding to coordinates_evaluate (i.e. flux_array shape must match self._grid with supersample factor and flux_evaluate_indexes)

  • unconvolved – boolean, if True, does not apply a convolution

Returns:

convolved image on regular pixel grid, 2d array

property grid_supersampling_factor
Returns:

supersampling factor set for higher resolution sub-pixel sampling of surface brightness

property coordinates_evaluate
Returns:

1d array of all coordinates being evaluated to perform the image computation

property convolution_class
Returns:

convolution class (can be SubgridKernelConvolution, PixelKernelConvolution)

property grid_class
Returns:

grid class (can be RegularGrid)

jaxtronomy.ImSim.Numerics.numerics_subframe module

class NumericsSubFrame(pixel_grid, psf, supersampling_factor=1, compute_mode='regular', supersampling_convolution=False, supersampling_kernel_size=5, flux_evaluate_indexes=None, supersampled_indexes=None, compute_indexes=None, point_source_supersampling_factor=1, convolution_kernel_size=None, convolution_type='fft', truncation_conv=None)[source]

Bases: PointSourceRendering

This class finds the optimal rectangular sub-frame of a data to be modelled that contains all the flux_evaluate_indexes and performs the numerical calculations only in this frame and then patches zeros around it to match the full data size.

__init__(pixel_grid, psf, supersampling_factor=1, compute_mode='regular', supersampling_convolution=False, supersampling_kernel_size=5, flux_evaluate_indexes=None, supersampled_indexes=None, compute_indexes=None, point_source_supersampling_factor=1, convolution_kernel_size=None, convolution_type='fft', truncation_conv=None)[source]
Parameters:

  • pixel_grid – PixelGrid() class instance

  • psf – PSF() class instance

  • compute_mode – options are: ‘regular’, ‘adaptive’

  • supersampling_factor – int, factor of higher resolution sub-pixel sampling of surface brightness

  • supersampling_convolution – bool, if True, performs (part of) the convolution on the super-sampled grid/pixels

  • supersampling_kernel_size – int (odd number), size (in regular pixel units) of the super-sampled convolution

  • flux_evaluate_indexes – boolean 2d array of size of image before supersampling (or None, then initiated as grid of True’s). Pixels indicated with True will be used to perform the surface brightness computation (and possible lensing ray-shooting). Pixels marked as False will be assigned a flux value of zero (or ignored in the adaptive convolution)

  • supersampled_indexes – 2d boolean array (only used in mode=’adaptive’) of pixels to be supersampled (in surface brightness and if supersampling_convolution=True also in convolution)

  • compute_indexes – 2d boolean array (only used in mode=’adaptive’), marks pixel that the resonse after convolution is computed (all others =0). This can be set to likelihood_mask in the Likelihood module for consistency.

  • point_source_supersampling_factor – super-sampling resolution of the point source placing

  • convolution_kernel_size – int, odd number, size of convolution kernel. If None, takes size of point_source_kernel

  • convolution_type – string, ‘fft’, ‘grid’, ‘fft_static’ mode of 2d convolution

  • truncation_conv – Truncation used for the construction of the convolution kernels (only relevant for Gaussian convolution). By default, the truncation from the psf class will be used. Can be overwritten so that different PSFs are used for convolution and point source rendering.

re_size_convolve(flux_array, unconvolved=False)[source]
Parameters:

flux_array – 1d array, flux values corresponding to coordinates_evaluate

Returns:

convolved image on regular pixel grid, 2d array

property grid_supersampling_factor
Returns:

supersampling factor set for higher resolution sub-pixel sampling of surface brightness

property coordinates_evaluate
Returns:

1d array of all coordinates being evaluated to perform the image computation

property convolution_class
Returns:

convolution class (can be SubgridKernelConvolution, PixelKernelConvolution, MultiGaussianConvolution, …)

property grid_class
Returns:

grid class (can be RegularGrid, AdaptiveGrid)

jaxtronomy.ImSim.Numerics.point_source_rendering module

class PointSourceRendering(pixel_grid, supersampling_factor, psf)[source]

Bases: object

Numerics to compute the point source response on an image.

__init__(pixel_grid, supersampling_factor, psf)[source]
Parameters:

  • pixel_grid – PixelGrid() instance

  • supersampling_factor – int, factor of supersampling of point source if None, then uses the supersampling factor of the original PSF

  • psf – PSF() instance

point_source_rendering(ra_pos, dec_pos, amp, unconvolved=False)[source]
Parameters:

  • ra_pos – list of RA positions of point source(s)

  • dec_pos – list of DEC positions of point source(s)

  • amp – list of amplitudes of point source(s)

  • unconvolved – bool, if True, renders point source on a single pixel instead of proper PSF

Returns:

2d numpy array of size of the image with the point source(s) rendered

psf_variance_map(ra_pos, dec_pos, amp, data, fix_psf_variance_map=False)[source]

Variance of PSF error.

Parameters:

  • ra_pos – image positions of point sources

  • dec_pos – image positions of point sources

  • amp – amplitude of modeled point sources

  • data – 2d numpy array of the data

  • fix_psf_variance_map – bool, if True, estimates the error based on the input (modeled) amplitude, else uses the data to do so.

Returns:

2d array of size of the image with error terms (sigma**2) expected from inaccuracies in the PSF modeling

Module contents