<commit_before>"""Perform weak-lensing analysis of simulated sources.
"""

import numpy as np

import astropy.table

import galsim

import descwl.model

class OverlapResults(object):
    """Results of analyzing effects of overlapping sources on weak lensing.

    Results are normally obtained by calling :meth:`OverlapAnalyzer.finalize` after
    simulating an image or using :meth:`descwl.output.Reader` to load saved
    simulation results.

    Indices returned by the `find` methods below can be used to lookup analysis
    results via `table[index]` and the corresponding simulated datacube using
    `stamps[index]` and `bounds[index]`.

    Args:
        survey(descwl.survey.Survey): Simulated survey that results are based on.
        table(astropy.table.Table): Table of analysis results with one row per galaxy.
        stamps(array): Array of :class:`numpy.ndarray` postage-stamp datacubes.
        bounds(array): Array of galsim.BoundsI objects giving pixel coordinates of
            each datacube in the full simulated image.

    Raises:
        RuntimeError: Image datacubes have unexpected number of slices.
    """
    def __init__(self,survey,table,stamps,bounds):
        self.survey = survey
        self.table = table
        self.stamps = stamps
        self.bounds = bounds
        self.num_objects = len(self.table)
        self.num_slices = self.stamps[0].shape[0]
        self.slice_labels = ['dflux','dx','dy','dscale','dg1','dg2']
        if self.num_slices not in (1,len(self.slice_labels)):
            raise RuntimeError('Image datacubes have unexpected number of slices (%d).'
                % self.num_slices)
        self.locals = { name: self.table[name] for name in self.table.colnames }

    def select(self,selector):
        """Select objects.

        This function is implemented using :func:`eval` so should only be used when the
        source of the selector string is trusted.

        Args:
            selector(str): A selection string consisting of a boolean expression of
                column table names, e.g. 'snr_iso > 5'.  The special strings 'ALL' and
                'NONE' do what you would expect.

        Returns:
            :class:`numpy.ndarray`: Boolean array containing a selection mask with one
                entry per object whose True/False value indicates if the object has
                been selected.

        Raises:
            RuntimeError: Selector uses an undefined variable name.
        """
        if selector == 'ALL':
            return np.ones(self.num_objects,dtype=bool)
        elif selector == 'NONE':
            return np.zeros(self.num_objects,dtype=bool)
        else:
            try:
                return eval(selector,self.locals)
            except NameError,e:
                raise RuntimeError('%s in selector %r.' % (e.message,selector))

    def get_stamp(self,index,datacube_index=0):
        """Return the simulated postage stamp for a single galaxy.

        Args:
            index(int): Index of the requested galaxy in our results. Use
                meth:`find_galaxy` to get the index for an identifier.
            datacube_index(int): Which slice of the datacube to return.

        Returns:
            galsim.ImageView: A GalSim image for the requested galaxy.

        Raises:
            RuntimeError: No such galaxy in the results.
        """
        if index < 0 or index >= len(self.stamps):
            raise RuntimeError('No such galaxy with index=%d.' % index)
        bounds = self.bounds[index]
        image = galsim.Image(self.stamps[index][datacube_index],
            xmin = bounds.xmin,ymin = bounds.ymin,
            scale = self.survey.pixel_scale,make_const = True)
        return image

    def get_subimage(self,indices,datacube_index=0):
        """Return simulated subimage of a set of objects.

        Args:
            indices(iterable): Indices of the objects to include in the subimage.

        Returns:
            galsim.Image: Image of the selected objects or None if indices is empty.

        Raises:
            RuntimeError: An index is out of range.
        """
        if len(indices) == 0:
            return None
        subimage_bounds = galsim.BoundsI()
        try:
            for index in indices:
                subimage_bounds += self.bounds[index]
        except IndexError:
            raise RuntimeError('Index out of range %d' % index)
        subimage = galsim.Image(bounds = subimage_bounds, scale = self.survey.pixel_scale)
        for index in indices:
            stamp = self.get_stamp(index,datacube_index)
            overlap = subimage_bounds & stamp.bounds
            subimage[overlap] += stamp[overlap]
        return subimage

    def get_fisher_images(self,indices):
        """Return Fisher-matrix images for a set of objects.

        Args:
            indices(iterable): Indices of the objects to include in the subimage.

        Returns:
            numpy.ndarray: Array with shape (npar,npar,height,width) where npar = 6*len(indices)
                is the number of floating parameters used for the calculation, and (height,width)
                are the dimensions of the subimage containing the requested objects.
                The returned array is symmetric in its first two indices. Returns None if
                indices is empty.

        Raises:
            RuntimeError: An index is out of range.
        """
        nselected = len(indices)
        npartials = self.num_slices
        npar = nselected*npartials
        background = self.get_subimage(indices)
        if nselected == 0 or background is None:
            return None
        height,width = background.array.shape
        sky_level = self.survey.mean_sky_level

        fisher_images = np.empty((npar,npar,height,width))
        stamp1 = background.copy()
        stamp2 = background.copy()

        for index1 in range(npar):
            galaxy1 = indices[index1//npartials]
            slice1 = index1%npartials
            stamp1.array[:] = 0.
            stamp1[self.bounds[galaxy1]] = self.get_stamp(galaxy1,slice1)
            if slice1 == 0:
                # Normalize to give partial with respect to added flux in electrons.
                stamp1 /= self.table['flux'][galaxy1]
            for index2 in range(index1+1):
                galaxy2 = indices[index2//npartials]
                slice2 = index2%npartials
                stamp2.array[:] = 0.
                stamp2[self.bounds[galaxy2]] = self.get_stamp(galaxy2,slice2)
                if slice2 == 0:
                    # Normalize to give partial with respect to added flux in electrons.
                    stamp2 /= self.table['flux'][galaxy2]
                fisher_images[index1,index2] = stamp1.array*stamp2.array/(
                    background.array + sky_level)
                if index2 < index1:
                    fisher_images[index2,index1] = fisher_images[index1,index2]
        return fisher_images

class OverlapAnalyzer(object):
    """Analyze impact of overlapping sources on weak lensing.

    Args:
        survey(descwl.survey.Survey): Simulated survey to describe with FITS header keywords.
    """
    def __init__(self,survey):
        self.survey = survey
        self.models = [ ]
        self.stamps = [ ]
        self.bounds = [ ]

    def add_galaxy(self,model,stamps,bounds):
        """Add one galaxy to be analyzed.

        Args:
            model(:class:`descwl.model.Galaxy`): The galaxy model used for rendering.
            stamps(:class:`numpy.ndarray`): Array of shape (nstamp,width,height) containing
                pixel values for nstamp stamps of dimensions (width,height).<commit_msg>Add get_matrices method<commit_after>"""Perform weak-lensing analysis of simulated sources.
"""

import numpy as np

import astropy.table

import galsim

import descwl.model

class OverlapResults(object):
    """Results of analyzing effects of overlapping sources on weak lensing.

    Results are normally obtained by calling :meth:`OverlapAnalyzer.finalize` after
    simulating an image or using :meth:`descwl.output.Reader` to load saved
    simulation results.

    Indices returned by the `find` methods below can be used to lookup analysis
    results via `table[index]` and the corresponding simulated datacube using
    `stamps[index]` and `bounds[index]`.

    Args:
        survey(descwl.survey.Survey): Simulated survey that results are based on.
        table(astropy.table.Table): Table of analysis results with one row per galaxy.
        stamps(array): Array of :class:`numpy.ndarray` postage-stamp datacubes.
        bounds(array): Array of galsim.BoundsI objects giving pixel coordinates of
            each datacube in the full simulated image.

    Raises:
        RuntimeError: Image datacubes have unexpected number of slices.
    """
    def __init__(self,survey,table,stamps,bounds):
        self.survey = survey
        self.table = table
        self.stamps = stamps
        self.bounds = bounds
        self.num_objects = len(self.table)
        self.num_slices = self.stamps[0].shape[0]
        self.slice_labels = ['dflux','dx','dy','dscale','dg1','dg2']
        if self.num_slices not in (1,len(self.slice_labels)):
            raise RuntimeError('Image datacubes have unexpected number of slices (%d).'
                % self.num_slices)
        self.locals = { name: self.table[name] for name in self.table.colnames }

    def select(self,selector):
        """Select objects.

        This function is implemented using :func:`eval` so should only be used when the
        source of the selector string is trusted.

        Args:
            selector(str): A selection string consisting of a boolean expression of
                column table names, e.g. 'snr_iso > 5'.  The special strings 'ALL' and
                'NONE' do what you would expect.

        Returns:
            :class:`numpy.ndarray`: Boolean array containing a selection mask with one
                entry per object whose True/False value indicates if the object has
                been selected.

        Raises:
            RuntimeError: Selector uses an undefined variable name.
        """
        if selector == 'ALL':
            return np.ones(self.num_objects,dtype=bool)
        elif selector == 'NONE':
            return np.zeros(self.num_objects,dtype=bool)
        else:
            try:
                return eval(selector,self.locals)
            except NameError,e:
                raise RuntimeError('%s in selector %r.' % (e.message,selector))

    def get_stamp(self,index,datacube_index=0):
        """Return the simulated postage stamp for a single galaxy.

        Args:
            index(int): Index of the requested galaxy in our results. Use
                meth:`find_galaxy` to get the index for an identifier.
            datacube_index(int): Which slice of the datacube to return.

        Returns:
            galsim.ImageView: A GalSim image for the requested galaxy.

        Raises:
            RuntimeError: No such galaxy in the results.
        """
        if index < 0 or index >= len(self.stamps):
            raise RuntimeError('No such galaxy with index=%d.' % index)
        bounds = self.bounds[index]
        image = galsim.Image(self.stamps[index][datacube_index],
            xmin = bounds.xmin,ymin = bounds.ymin,
            scale = self.survey.pixel_scale,make_const = True)
        return image

    def get_subimage(self,indices,datacube_index=0):
        """Return simulated subimage of a set of objects.

        Args:
            indices(iterable): Indices of the objects to include in the subimage.

        Returns:
            galsim.Image: Image of the selected objects or None if indices is empty.

        Raises:
            RuntimeError: An index is out of range.
        """
        if len(indices) == 0:
            return None
        subimage_bounds = galsim.BoundsI()
        try:
            for index in indices:
                subimage_bounds += self.bounds[index]
        except IndexError:
            raise RuntimeError('Index out of range %d' % index)
        subimage = galsim.Image(bounds = subimage_bounds, scale = self.survey.pixel_scale)
        for index in indices:
            stamp = self.get_stamp(index,datacube_index)
            overlap = subimage_bounds & stamp.bounds
            subimage[overlap] += stamp[overlap]
        return subimage

    def get_fisher_images(self,indices):
        """Return Fisher-matrix images for a set of objects.

        Fisher matrix images are derived from the partial derivatives with respect to
        six parameters: total flux in detected electrons, centroid positions in x and y
        in arcseconds, flux-preserving radial scale (dimensionless), and shear g1,g2
        with \|g\| = (a-b)/(a+b).

        Args:
            indices(iterable): Indices of the objects to include in the subimage.

        Returns:
            numpy.ndarray: Array with shape (npar,npar,height,width) where npar = 6*len(indices)
                is the number of floating parameters used for the calculation, and (height,width)
                are the dimensions of the subimage containing the requested objects.
                The returned array is symmetric in its first two indices. Returns None if
                indices is empty.

        Raises:
            RuntimeError: An index is out of range.
        """
        nselected = len(indices)
        npartials = self.num_slices
        npar = nselected*npartials
        background = self.get_subimage(indices)
        if nselected == 0 or background is None:
            return None
        height,width = background.array.shape
        sky_level = self.survey.mean_sky_level

        fisher_images = np.empty((npar,npar,height,width))
        stamp1 = background.copy()
        stamp2 = background.copy()

        for index1 in range(npar):
            galaxy1 = indices[index1//npartials]
            slice1 = index1%npartials
            stamp1.array[:] = 0.
            stamp1[self.bounds[galaxy1]] = self.get_stamp(galaxy1,slice1)
            if slice1 == 0:
                # Normalize to give partial with respect to added flux in electrons.
                stamp1 /= self.table['flux'][galaxy1]
            for index2 in range(index1+1):
                galaxy2 = indices[index2//npartials]
                slice2 = index2%npartials
                stamp2.array[:] = 0.
                stamp2[self.bounds[galaxy2]] = self.get_stamp(galaxy2,slice2)
                if slice2 == 0:
                    # Normalize to give partial with respect to added flux in electrons.
                    stamp2 /= self.table['flux'][galaxy2]
                fisher_images[index1,index2] = stamp1.array*stamp2.array/(
                    background.array + sky_level)
                if index2 < index1:
                    fisher_images[index2,index1] = fisher_images[index1,index2]
        return fisher_images

    def get_matrices(self,fisher_images):
        """Return matrices derived from Fisher-matrix images.

        If the Fisher matrix is not invertible, `covariance`, `variance` and `correlation`
        will be returned as None.  If any variances are <= 0, `correlation` will be
        returned as None.

        Args:
            fisher_images(numpy.ndarray): Fisher-matrix images returned by
                :meth:`get_fisher_images`. These are assumed to be symmetric in the
                first two indices, but this is not checked.

        Returns:
            tuple: Tuple `(fisher,covariance,variance,correlation)` of :class:`numpy.ndarray`
                where `variance` has shape (npar,) and all other arrays are symmetric with
                shape (npar,npar).  If any array cannot be calculated, it will be returned
                as None.

        Raises:
            RuntimeError: `fisher_images` has unexpected shape.
        """
        shape = fisher_images.shape
        if len(shape) != 4 or shape[0] != shape[1]:
            raise RuntimeError('Fisher images have unexpected shape %r' % shape)

        fisher = np.sum(fisher_images,axis=(2,3))
        covariance,variance,correlation = None,None,None
        try:
            covariance = np.linalg.inv(fisher)
            variance = np.diag(covariance)
            if np.min(variance) > 0:
                correlation = covariance/np.sqrt(np.outer(variance,variance))
        except np.linalg.LinAlgError:
            pass

        return fisher,covariance,variance,correlation

class OverlapAnalyzer(object):
    """Analyze impact of overlapping sources on weak lensing.

    Args:
        survey(descwl.survey.Survey): Simulated survey to describe with FITS header keywords.
    """
    def __init__(self,survey):
        self.survey = survey
        self.models = [ ]
        self.stamps = [ ]
        self.bounds = [ ]

    def add_galaxy(self,model,stamps,bounds):
        """Add one galaxy to be analyzed.

        Args:
            model(:class:`descwl.model.Galaxy`): The galaxy model used for rendering.
            stamps(:class:`numpy.ndarray`): Array of shape (nstamp,width,height) containing
                pixel values for nstamp stamps of dimensions (width,height).
{'fraction_non_alphanumeric': 0.08327526132404181, 'fraction_numerical': 0.0070267131242741, 'mean_word_length': 2.1708709261646106, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '721360', 'n_tokens_mistral': 4773, 'n_tokens_neox': 4602, 'n_words': 1480}