lenstronomy.ImSim.Numerics.grid — lenstronomy 1.9.1 documentation

# Source code for lenstronomy.ImSim.Numerics.grid

    import numpy as np
    from lenstronomy.Util import util
    from lenstronomy.Util import image_util
    from lenstronomy.Data.coord_transforms import Coordinates1D

    from lenstronomy.Util.package_util import exporter
    export, __all__ = exporter()

    @export
    class AdaptiveGrid(Coordinates1D):
        """
        manages a super-sampled grid on the partial image
        """
        def __init__(self, nx, ny, transform_pix2angle, ra_at_xy_0, dec_at_xy_0, supersampling_indexes,
                     supersampling_factor, flux_evaluate_indexes=None):
            """

            :param nx: number of pixels in x-axis
            :param ny: number of pixels in y-axis
            :param transform_pix2angle: 2x2 matrix, mapping of pixel to coordinate
            :param ra_at_xy_0: ra coordinate at pixel (0,0)
            :param dec_at_xy_0: dec coordinate at pixel (0,0)
            :param supersampling_indexes: bool array of shape nx x ny, corresponding to pixels being super_sampled
            :param supersampling_factor: int, factor (per axis) of super-sampling
            :param 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.
            """
            super(AdaptiveGrid, self).__init__(transform_pix2angle, ra_at_xy_0, dec_at_xy_0)
            self._nx = nx
            self._ny = ny
            self._x_grid, self._y_grid = self.coordinate_grid(nx, ny)
            if flux_evaluate_indexes is None:
                flux_evaluate_indexes = np.ones_like(self._x_grid, dtype=bool)
            supersampled_indexes1d = util.image2array(supersampling_indexes)
            self._high_res_indexes1d = (supersampled_indexes1d) & (flux_evaluate_indexes)
            self._low_res_indexes1d = (np.invert(supersampled_indexes1d)) & (flux_evaluate_indexes)
            self._supersampling_factor = supersampling_factor
            self._num_sub = supersampling_factor * supersampling_factor
            self._x_low_res = self._x_grid[self._low_res_indexes1d]
            self._y_low_res = self._y_grid[self._low_res_indexes1d]
            self._num_low_res = len(self._x_low_res)

        @property
        def coordinates_evaluate(self):
            """

            :return: 1d array of all coordinates being evaluated to perform the image computation
            """
            ra_low, dec_low = self._x_low_res, self._y_low_res
            ra_high, dec_high = self._high_res_coordinates
            ra_joint = np.append(ra_low, ra_high)
            dec_joint = np.append(dec_low, dec_high)
            return ra_joint, dec_joint

        def flux_array2image_low_high(self, flux_array, high_res_return=True):
            """

            :param flux_array: 1d array of low and high resolution flux values corresponding to the coordinates_evaluate order
            :param 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)
            :return: 2d array, 2d array, corresponding to (partial) images in low and high resolution (to be convolved)
            """
            low_res_values = flux_array[0:self._num_low_res]
            high_res_values = flux_array[self._num_low_res:]
            image_low_res = self._merge_low_high_res(low_res_values, high_res_values)
            if high_res_return is True:
                image_high_res_partial = self._high_res_image(high_res_values)
            else:
                image_high_res_partial = None
            return image_low_res, image_high_res_partial

        @property
        def _high_res_coordinates(self):
            """

            :return: 1d arrays of subpixel grid coordinates
            """
            if not hasattr(self, '_x_sub_grid'):
                self._subpixel_coordinates()
            return self._x_high_res, self._y_high_res

        def _merge_low_high_res(self, low_res_values, supersampled_values):
            """
            adds/overwrites the supersampled values on the image

            :param low_res_values: 1d array of image with low resolution
            :param supersampled_values: values of the supersampled sub-pixels
            :return: 2d image
            """

            array = np.zeros_like(self._x_grid)
            array[self._low_res_indexes1d] = low_res_values
            array_low_res_partial = self._average_subgrid(supersampled_values)
            array[self._high_res_indexes1d] = array_low_res_partial
            return util.array2image(array)

        def _high_res_image(self, supersampled_values):
            """

            :param supersampled_values: 1d array of supersampled values corresponding to coordinates
            :return: 2d array of supersampled image (zeros outside supersampled frame)
            """
            high_res = np.zeros((self._nx * self._supersampling_factor, self._ny * self._supersampling_factor))
            count = 0
            for i in range(self._supersampling_factor):
                for j in range(self._supersampling_factor):
                    selected = supersampled_values[count::self._num_sub]
                    high_res[i::self._supersampling_factor, j::self._supersampling_factor] = self._array2image_subset(
                        selected)
                    count += 1
            return high_res

        def _subpixel_coordinates(self):
            """

            :return: 1d arrays of subpixel grid coordinates
            """
            x_grid_select = self._x_grid[self._high_res_indexes1d]
            y_grid_select = self._y_grid[self._high_res_indexes1d]
            x_sub_grid = np.zeros(len(x_grid_select) * self._num_sub)
            y_sub_grid = np.zeros(len(x_grid_select) * self._num_sub)
            count = 0
            for i in range(self._supersampling_factor):
                for j in range(self._supersampling_factor):
                    x_ij = 1. / self._supersampling_factor / 2. + j / float(self._supersampling_factor) - 1. / 2
                    y_ij = 1. / self._supersampling_factor / 2. + i / float(self._supersampling_factor) - 1. / 2
                    delta_ra, delta_dec = self.map_pix2coord(x_ij, y_ij)
                    delta_ra0, delta_dec_0 = self.map_pix2coord(0, 0)
                    x_sub_grid[count::self._num_sub] = x_grid_select + delta_ra - delta_ra0
                    y_sub_grid[count::self._num_sub] = y_grid_select + delta_dec - delta_dec_0
                    count += 1
            self._x_high_res = x_sub_grid
            self._y_high_res = y_sub_grid

        def _average_subgrid(self, subgrid_values):
            """
            averages the values over a pixel

            :param subgrid_values: values (e.g. flux) of subgrid coordinates
            :return: 1d array of size of the supersampled pixels
            """
            values_2d = np.reshape(subgrid_values, (-1, self._num_sub))
            return np.mean(values_2d, axis=1)

        def _array2image_subset(self, array):
            """
            maps a 1d array into a (nx, ny) 2d grid with array populating the idex_mask indices
            :param array: 1d array
            :return: 2d array
            """
            grid1d = np.zeros(self._nx * self._ny)
            grid1d[self._high_res_indexes1d] = array
            grid2d = util.array2image(grid1d, self._nx, self._ny)
            return grid2d

    @export
    class RegularGrid(Coordinates1D):
        """
        manages a super-sampled grid on the partial image
        """
        def __init__(self, nx, ny, transform_pix2angle, ra_at_xy_0, dec_at_xy_0, supersampling_factor=1,
                     flux_evaluate_indexes=None):
            """

            :param nx: number of pixels in x-axis
            :param ny: number of pixels in y-axis
            :param transform_pix2angle: 2x2 matrix, mapping of pixel to coordinate
            :param ra_at_xy_0: ra coordinate at pixel (0,0)
            :param dec_at_xy_0: dec coordinate at pixel (0,0)
            :param supersampling_indexes: bool array of shape nx x ny, corresponding to pixels being super_sampled
            :param supersampling_factor: int, factor (per axis) of super-sampling
            :param 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.
            """
            super(RegularGrid, self).__init__(transform_pix2angle, ra_at_xy_0, dec_at_xy_0)
            self._supersampling_factor = supersampling_factor
            self._nx = nx
            self._ny = ny
            self._x_grid, self._y_grid = self.coordinate_grid(nx, ny)
            if flux_evaluate_indexes is None:
                flux_evaluate_indexes = np.ones_like(self._x_grid, dtype=bool)
            else:
                flux_evaluate_indexes = util.image2array(flux_evaluate_indexes)
            self._compute_indexes = self._subgrid_index(flux_evaluate_indexes, self._supersampling_factor, self._nx, self._ny)

            x_grid_sub, y_grid_sub = util.make_subgrid(self._x_grid, self._y_grid, self._supersampling_factor)
            self._ra_subgrid = x_grid_sub[self._compute_indexes]
            self._dec_subgrid = y_grid_sub[self._compute_indexes]

        @property
        def coordinates_evaluate(self):
            """

            :return: 1d array of all coordinates being evaluated to perform the image computation
            """
            return self._ra_subgrid, self._dec_subgrid

        @property
        def grid_points_spacing(self):
            """
            effective spacing between coordinate points, after supersampling
            :return: sqrt(pixel_area)/supersampling_factor
            """
            return self.pixel_width / self._supersampling_factor

        @property
        def num_grid_points_axes(self):
            """
            effective number of points along each axes, after supersampling
            :return: number of pixels per axis, nx*supersampling_factor ny*supersampling_factor
            """
            return self._nx * self._supersampling_factor, self._ny * self._supersampling_factor

        @property
        def supersampling_factor(self):
            """
            :return: factor (per axis) of super-sampling relative to a pixel
            """
            return self._supersampling_factor

        def flux_array2image_low_high(self, flux_array, **kwargs):
            """

            :param flux_array: 1d array of low and high resolution flux values corresponding to the coordinates_evaluate order
            :return: 2d array, 2d array, corresponding to (partial) images in low and high resolution (to be convolved)
            """
            image = self._array2image(flux_array)
            if self._supersampling_factor > 1:
                image_high_res = image
                image_low_res = image_util.re_size(image, self._supersampling_factor)
            else:
                image_high_res = None
                image_low_res = image
            return image_low_res, image_high_res

        def _subgrid_index(self, idex_mask, subgrid_res, nx, ny):
            """

            :param idex_mask: 1d array of mask of data
            :param subgrid_res: subgrid resolution
            :return: 1d array of equivalent mask in subgrid resolution
            """
            idex_sub = np.repeat(idex_mask, subgrid_res, axis=0)
            idex_sub = util.array2image(idex_sub, nx=nx, ny=ny*subgrid_res)
            idex_sub = np.repeat(idex_sub, subgrid_res, axis=0)
            idex_sub = util.image2array(idex_sub)
            return idex_sub

        def _array2image(self, array):
            """
            maps a 1d array into a (nx, ny) 2d grid with array populating the idex_mask indices

            :param array: 1d array
            :param idex_mask: 1d array of length nx*ny
            :param nx: x-axis of 2d grid
            :param ny: y-axis of 2d grid
            :return:
            """
            nx, ny = self._nx * self._supersampling_factor, self._ny * self._supersampling_factor
            grid1d = np.zeros((nx * ny))
            grid1d[self._compute_indexes] = array
            grid2d = util.array2image(grid1d, nx, ny)
            return grid2d