Exemple #1
0
    def make_wcs_from_hpx(self, sum_ebins=False, proj='CAR', oversample=2,
                          normalize=True):
        """Make a WCS object and convert HEALPix data into WCS projection

        NOTE: this re-calculates the mapping, if you have already
        calculated the mapping it is much faster to use
        convert_to_cached_wcs() instead

        Parameters
        ----------
        sum_ebins  : bool
           sum energy bins over energy bins before reprojecting

        proj       : str
           WCS-projection

        oversample : int
           Oversampling factor for WCS map

        normalize  : bool
           True -> perserve integral by splitting HEALPix values between bins

        returns (WCS object, np.ndarray() with reprojected data)

        """
        self._wcs_proj = proj
        self._wcs_oversample = oversample
        self._wcs_2d = self.hpx.make_wcs(2, proj=proj, oversample=oversample)
        self._hpx2wcs = HpxToWcsMapping(self.hpx, self._wcs_2d)
        wcs, wcs_data = self.convert_to_cached_wcs(self.counts, sum_ebins,
                                                   normalize)
        return wcs, wcs_data
Exemple #2
0
    def make_wcs_from_hpx(self, sum_ebins=False, proj='CAR', oversample=2,
                          normalize=True):
        """Make a WCS object and convert HEALPix data into WCS projection

        NOTE: this re-calculates the mapping, if you have already
        calculated the mapping it is much faster to use
        convert_to_cached_wcs() instead

        Parameters
        ----------
        sum_ebins  : bool
           sum energy bins over energy bins before reprojecting

        proj       : str
           WCS-projection

        oversample : int
           Oversampling factor for WCS map

        normalize  : bool
           True -> perserve integral by splitting HEALPix values between bins

        returns (WCS object, np.ndarray() with reprojected data)

        """
        self._wcs_proj = proj
        self._wcs_oversample = oversample
        self._wcs_2d = self.hpx.make_wcs(2, proj=proj, oversample=oversample)
        self._hpx2wcs = HpxToWcsMapping(self.hpx, self._wcs_2d)
        wcs, wcs_data = self.convert_to_cached_wcs(self.counts, sum_ebins,
                                                   normalize)
        return wcs, wcs_data
Exemple #3
0
def main():

    import sys
    import argparse

    # Argument defintion
    usage = "usage: %(prog)s [options]"
    description = "Collect all the new source"

    parser = argparse.ArgumentParser(usage, description=__abstract__)

    parser.add_argument("-i",
                        "--input",
                        type=argparse.FileType('r'),
                        required=True,
                        help="Input file")

    parser.add_argument("-e",
                        "--extension",
                        type=str,
                        default="SKYMAP",
                        help="FITS HDU with HEALPix map")

    parser.add_argument("--ebin",
                        type=str,
                        default=None,
                        help="Energy bin, integer or 'ALL'")

    parser.add_argument("--zscale",
                        type=str,
                        default='log',
                        help="Scaling for color scale")

    parser.add_argument("--zmin",
                        type=float,
                        default=None,
                        help="Minimum z-axis value")

    parser.add_argument("--zmax",
                        type=float,
                        default=None,
                        help="Maximum z-axis value")

    parser.add_argument("-o",
                        "--output",
                        type=argparse.FileType('w'),
                        help="Output file.  Leave blank for interactive.")

    # Parse the command line
    args = parser.parse_args(sys.argv[1:])

    # Get the model
    f = pf.open(args.input.name)
    # We need a better check
    maptype = "None"

    model_hdu = f[args.extension]

    hpxmap = HpxMap.create_from_hdulist(f, hdu=args.extension)
    outdata = []

    if args.ebin == "ALL":
        wcsproj = hpxmap.hpx.make_wcs(naxis=2,
                                      proj='MOL',
                                      energies=None,
                                      oversample=2)
        mapping = HpxToWcsMapping(hpxmap.hpx, wcsproj)

        for i, data in enumerate(hpxmap.counts):
            ip = ImagePlotter(data=data, proj=hpxmap.hpx, mapping=mapping)
            fig = plt.figure(i)
            im, ax = ip.plot(zscale=args.zscale,
                             vmin=args.zmin,
                             vmax=args.zmax)
            outdata.append(fig)

    elif args.ebin is None:
        ip = ImagePlotter(data=hpxmap.counts, proj=hpxmap.hpx)
        im, ax = ip.plot(zscale=args.zscale, vmin=args.zmin, vmax=args.zmax)
        outdata.append((im, ax))
    else:
        try:
            ibin = int(args.ebin)
            ip = ImagePlotter(data=hpxmap.counts[ibin], proj=hpxmap.hpx)
            im, ax = ip.plot(zscale=args.zscale,
                             vmin=args.zmin,
                             vmax=args.zmax)
            outdata.append((im, ax))
        except:
            raise ValueError("--ebin argument must be an integer or 'ALL'")

    if args.output is None:
        plt.show()
    else:
        if len(outdata) == 1:
            plt.savefig(args.output.name)
        else:
            base, ext = os.path.splitext(args.output.name)
            for i, fig in enumerate(outdata):
                fig.savefig("%s_%02i%s" % (base, i, ext))
Exemple #4
0
def main():

    import sys
    import argparse

    # Argument defintion
    usage = "usage: %(prog)s [options]" 
    description = "Collect all the new source"

    parser = argparse.ArgumentParser(usage,description=__abstract__)

    parser.add_argument("-i", "--input",type=argparse.FileType('r'),required=True,
                        help="Input file")

    parser.add_argument("-e", "--extension",type=str,default="SKYMAP",
                        help="FITS HDU with HEALPix map")
 
    parser.add_argument("--ebin",type=str,default=None,
                        help="Energy bin, integer or 'ALL'")
    
    parser.add_argument("-o", "--output",type=argparse.FileType('w'),
                        help="Output file.  Leave blank for interactive.")
    
    # Parse the command line
    args = parser.parse_args(sys.argv[1:])

    # Get the model 
    f = pf.open(args.input.name)
    # We need a better check
    maptype = "None"

    model_hdu = f[args.extension]
        
    hpxmap = HpxMap.create_from_hdulist(f,extname=args.extension,ebounds="EBOUNDS")
    outdata = []
     
    if args.ebin == "ALL":
        wcsproj = hpxmap.hpx.make_wcs(naxis=2,proj='AIT',energies=None,oversample=2)
        mapping = HpxToWcsMapping(hpxmap.hpx,wcsproj)
        
        for i,data in enumerate(hpxmap.counts):
            ip =  ImagePlotter(data=data,proj=hpxmap.hpx,mapping=mapping)  
            fig = plt.figure(i)
            im,ax = ip.plot(zscale='log')
            outdata.append(fig)

    elif args.ebin is None:
        ip =  ImagePlotter(data=hpxmap.counts,proj=hpxmap.hpx)  
        im,ax = ip.plot(zscale='log')
        outdata.append((im,ax))        
    else:
        try:
            ibin = int(args.ebin)
            ip =  ImagePlotter(data=hpxmap.counts[ibin],proj=hpxmap.hpx)  
            im,ax = ip.plot(zscale='log')
            outdata.append((im,ax))        
        except:
            print("--ebin argument must be an integer or 'ALL'")

    if args.output is None:
        plt.show()
    else:
        plt.savefig(args.output.name)
Exemple #5
0
class HpxMap(Map_Base):
    """ Representation of a 2D or 3D counts map using HEALPix. """

    def __init__(self, counts, hpx):
        """ C'tor, fill with a counts vector and a HPX object """
        Map_Base.__init__(self, counts)
        self._hpx = hpx
        self._wcs2d = None
        self._hpx2wcs = None

    @property
    def hpx(self):
        return self._hpx

    @staticmethod
    def create_from_hdu(hdu, ebins):
        """ Creates and returns an HpxMap object from a FITS HDU.

        hdu    : The FITS
        ebins  : Energy bin edges [optional]
        """
        hpx = HPX.create_from_header(hdu.header, ebins)
        colnames = hdu.columns.names
        nebin = 0
        for c in colnames:
            if c.find("CHANNEL") == 0:
                nebin += 1
            pass
        data = np.ndarray((nebin, hpx.npix))
        for i in range(nebin):
            cname = "CHANNEL%i" % (i + 1)
            data[i, 0:] = hdu.data.field(cname)
            pass
        return HpxMap(data, hpx)

    @staticmethod
    def create_from_hdulist(hdulist, extname="SKYMAP", ebounds="EBOUNDS"):
        """ Creates and returns an HpxMap object from a FITS HDUList

        extname : The name of the HDU with the map data
        ebounds : The name of the HDU with the energy bin data
        """
        if ebounds is not None:
            try:
                ebins = utils.read_energy_bounds(hdulist[ebounds])
            except:
                ebins = None
        else:
            ebins = None

        hpxMap = HpxMap.create_from_hdu(hdulist[extname], ebins)
        return hpxMap

    def make_wcs_from_hpx(self, sum_ebins=False, proj='CAR', oversample=2,
                          normalize=True):
        """Make a WCS object and convert HEALPix data into WCS projection

        NOTE: this re-calculates the mapping, if you have already
        calculated the mapping it is much faster to use
        convert_to_cached_wcs() instead

        Parameters
        ----------
        sum_ebins  : bool
           sum energy bins over energy bins before reprojecting

        proj       : str
           WCS-projection

        oversample : int
           Oversampling factor for WCS map

        normalize  : bool
           True -> perserve integral by splitting HEALPix values between bins

        returns (WCS object, np.ndarray() with reprojected data)

        """
        self._wcs_proj = proj
        self._wcs_oversample = oversample
        self._wcs_2d = self.hpx.make_wcs(2, proj=proj, oversample=oversample)
        self._hpx2wcs = HpxToWcsMapping(self.hpx, self._wcs_2d)
        wcs, wcs_data = self.convert_to_cached_wcs(self.counts, sum_ebins,
                                                   normalize)
        return wcs, wcs_data

    def convert_to_cached_wcs(self, hpx_in, sum_ebins=False, normalize=True):
        """ Make a WCS object and convert HEALPix data into WCS projection

        Parameters
        ----------
        hpx_in     : `~numpy.ndarray`
           HEALPix input data
        sum_ebins  : bool
           sum energy bins over energy bins before reprojecting
        normalize  : bool
           True -> perserve integral by splitting HEALPix values between bins

        returns (WCS object, np.ndarray() with reprojected data)
        """
        if self._hpx2wcs is None:
            raise Exception('HpxMap.convert_to_cached_wcs() called '
                            'before make_wcs_from_hpx()')

        if len(hpx_in.shape) == 1:
            wcs_data = np.ndarray(self._hpx2wcs.npix)
            loop_ebins = False
            hpx_data = hpx_in
        elif len(hpx_in.shape) == 2:
            if sum_ebins:
                wcs_data = np.ndarray(self._hpx2wcs.npix)
                hpx_data = hpx_in.sum(0)
                loop_ebins = False
            else:
                wcs_data = np.ndarray((self.counts.shape[0],
                                       self._hpx2wcs.npix[0],
                                       self._hpx2wcs.npix[1]))
                hpx_data = hpx_in
                loop_ebins = True
        else:
            raise Exception('Wrong dimension for HpxMap %i' %
                            len(hpx_in.shape))

        if loop_ebins:
            for i in range(hpx_data.shape[0]):
                self._hpx2wcs.fill_wcs_map_from_hpx_data(
                    hpx_data[i], wcs_data[i], normalize)
                pass
            wcs_data.reshape((self.counts.shape[0], self._hpx2wcs.npix[
                             0], self._hpx2wcs.npix[1]))
            # replace the WCS with a 3D one
            wcs = self.hpx.make_wcs(3, proj=self._wcs_proj,
                                    energies=self.hpx.ebins,
                                    oversample=self._wcs_oversample)
        else:
            self._hpx2wcs.fill_wcs_map_from_hpx_data(
                hpx_data, wcs_data, normalize)
            wcs_data.reshape(self._hpx2wcs.npix)
            wcs = self._wcs_2d

        return wcs, wcs_data

    def get_map_values(self, lons, lats, ibin=None):
        """Return the indices in the flat array corresponding to a set of coordinates

        Parameters
        ----------
        lons  : array-like
           'Longitudes' (RA or GLON)

        lats  : array-like
           'Latitidues' (DEC or GLAT)

        ibin : int or array-like
           Extract data only for a given energy bin.  None -> extract data for all bins

        Returns
        ----------
        vals : numpy.ndarray((n))
           Values of pixels in the flattened map, np.nan used to flag
           coords outside of map
        """
        theta = np.pi / 2. - np.radians(lats)
        phi = np.radians(lons)

        pix = hp.ang2pix(self.hpx.nside, theta, phi, nest=self.hpx.nest)

        if self.data.ndim == 2:
            return self.data[:, pix] if ibin is None else self.data[ibin, pix]
        else:
            return self.data[pix]

    def interpolate(self, lon, lat, egy=None):
        """Interpolate map values.

        """

        if self.data.ndim == 1:
            theta = np.pi / 2. - np.radians(lat)
            phi = np.radians(lon)
            return hp.pixelfunc.get_interp_val(self.counts, theta,
                                               phi, nest=self.hpx.nest)
        else:
            shape = np.broadcast(lon, lat, egy).shape
            lon *= np.ones(shape)
            lat *= np.ones(shape)
            egy *= np.ones(shape)
            theta = np.pi / 2. - np.radians(lat)
            phi = np.radians(lon)
            vals = []
            for i, _ in enumerate(self.hpx.evals):
                v = hp.pixelfunc.get_interp_val(self.counts[i], theta,
                                                phi, nest=self.hpx.nest)
                vals += [np.expand_dims(np.array(v, ndmin=1), -1)]

            vals = np.concatenate(vals, axis=-1)
            xvals = utils.val_to_pix(np.log(self.hpx.evals), np.log(egy))
            return map_coordinates(vals, [np.arange(shape[0]), xvals], order=1)
Exemple #6
0
class HpxMap(Map_Base):
    """ Representation of a 2D or 3D counts map using HEALPix. """

    def __init__(self, counts, hpx):
        """ C'tor, fill with a counts vector and a HPX object """
        super(HpxMap, self).__init__(counts)
        self._hpx = hpx
        self._wcs2d = None
        self._hpx2wcs = None

    @property
    def hpx(self):
        return self._hpx

    @classmethod
    def create_from_hdu(cls, hdu, ebins):
        """ Creates and returns an HpxMap object from a FITS HDU.

        hdu    : The FITS
        ebins  : Energy bin edges [optional]
        """
        hpx = HPX.create_from_hdu(hdu, ebins)
        colnames = hdu.columns.names
        cnames = []
        if hpx.conv.convname == 'FGST_SRCMAP_SPARSE':
            pixs = hdu.data.field('PIX')
            chans = hdu.data.field('CHANNEL')
            keys = chans * hpx.npix + pixs
            vals = hdu.data.field('VALUE')
            nebin = len(ebins)
            data = np.zeros((nebin, hpx.npix))
            data.flat[keys] = vals
        else:
            for c in colnames:
                if c.find(hpx.conv.colstring) == 0:
                    cnames.append(c)
            nebin = len(cnames)
            data = np.ndarray((nebin, hpx.npix))
            for i, cname in enumerate(cnames):
                data[i, 0:] = hdu.data.field(cname)

        return cls(data, hpx)

    @classmethod
    def create_from_hdulist(cls, hdulist, **kwargs):
        """ Creates and returns an HpxMap object from a FITS HDUList

        extname : The name of the HDU with the map data
        ebounds : The name of the HDU with the energy bin data
        """
        extname = kwargs.get('hdu', hdulist[1].name)
        ebins = fits_utils.find_and_read_ebins(hdulist)
        return cls.create_from_hdu(hdulist[extname], ebins)

    @classmethod
    def create_from_fits(cls, fitsfile, **kwargs):
        hdulist = fits.open(fitsfile)
        return cls.create_from_hdulist(hdulist, **kwargs)

    def create_image_hdu(self, name=None, **kwargs):
        kwargs['extname'] = name
        return self.hpx.make_hdu(self.counts, **kwargs)

    def make_wcs_from_hpx(self, sum_ebins=False, proj='CAR', oversample=2,
                          normalize=True):
        """Make a WCS object and convert HEALPix data into WCS projection

        NOTE: this re-calculates the mapping, if you have already
        calculated the mapping it is much faster to use
        convert_to_cached_wcs() instead

        Parameters
        ----------
        sum_ebins  : bool
           sum energy bins over energy bins before reprojecting

        proj       : str
           WCS-projection

        oversample : int
           Oversampling factor for WCS map

        normalize  : bool
           True -> perserve integral by splitting HEALPix values between bins

        returns (WCS object, np.ndarray() with reprojected data)

        """
        self._wcs_proj = proj
        self._wcs_oversample = oversample
        self._wcs_2d = self.hpx.make_wcs(2, proj=proj, oversample=oversample)
        self._hpx2wcs = HpxToWcsMapping(self.hpx, self._wcs_2d)
        wcs, wcs_data = self.convert_to_cached_wcs(self.counts, sum_ebins,
                                                   normalize)
        return wcs, wcs_data

    def convert_to_cached_wcs(self, hpx_in, sum_ebins=False, normalize=True):
        """ Make a WCS object and convert HEALPix data into WCS projection

        Parameters
        ----------
        hpx_in     : `~numpy.ndarray`
           HEALPix input data
        sum_ebins  : bool
           sum energy bins over energy bins before reprojecting
        normalize  : bool
           True -> perserve integral by splitting HEALPix values between bins

        returns (WCS object, np.ndarray() with reprojected data)
        """
        if self._hpx2wcs is None:
            raise Exception('HpxMap.convert_to_cached_wcs() called '
                            'before make_wcs_from_hpx()')

        if len(hpx_in.shape) == 1:
            wcs_data = np.ndarray(self._hpx2wcs.npix)
            loop_ebins = False
            hpx_data = hpx_in
        elif len(hpx_in.shape) == 2:
            if sum_ebins:
                wcs_data = np.ndarray(self._hpx2wcs.npix)
                hpx_data = hpx_in.sum(0)
                loop_ebins = False
            else:
                wcs_data = np.ndarray((self.counts.shape[0],
                                       self._hpx2wcs.npix[0],
                                       self._hpx2wcs.npix[1]))
                hpx_data = hpx_in
                loop_ebins = True
        else:
            raise Exception('Wrong dimension for HpxMap %i' %
                            len(hpx_in.shape))

        if loop_ebins:
            for i in range(hpx_data.shape[0]):
                self._hpx2wcs.fill_wcs_map_from_hpx_data(
                    hpx_data[i], wcs_data[i], normalize)
                pass
            wcs_data.reshape((self.counts.shape[0], self._hpx2wcs.npix[
                             0], self._hpx2wcs.npix[1]))
            # replace the WCS with a 3D one
            wcs = self.hpx.make_wcs(3, proj=self._wcs_proj,
                                    energies=np.log10(self.hpx.ebins),
                                    oversample=self._wcs_oversample)
        else:
            self._hpx2wcs.fill_wcs_map_from_hpx_data(
                hpx_data, wcs_data, normalize)
            wcs_data.reshape(self._hpx2wcs.npix)
            wcs = self._wcs_2d

        return wcs, wcs_data

    def get_pixel_skydirs(self):
        """Get a list of sky coordinates for the centers of every pixel. """
        sky_coords = self._hpx.get_sky_coords()
        if self.hpx.coordsys == 'GAL':
            return SkyCoord(l=sky_coords.T[0], b=sky_coords.T[1], unit='deg', frame='galactic')
        else:
            return SkyCoord(ra=sky_coords.T[0], dec=sky_coords.T[1], unit='deg', frame='icrs')

    def get_pixel_indices(self, lats, lons):
        """Return the indices in the flat array corresponding to a set of coordinates """
        return self._hpx.get_pixel_indices(lats, lons)

    def sum_over_energy(self):
        """ Reduce a counts cube to a counts map """
        # We sum over axis 0 in the array, and drop the energy binning in the
        # hpx object
        return HpxMap(np.sum(self.counts, axis=0), self.hpx.copy_and_drop_energy())

    def get_map_values(self, lons, lats, ibin=None):
        """Return the indices in the flat array corresponding to a set of coordinates

        Parameters
        ----------
        lons  : array-like
           'Longitudes' (RA or GLON)

        lats  : array-like
           'Latitidues' (DEC or GLAT)

        ibin : int or array-like
           Extract data only for a given energy bin.  None -> extract data for all bins

        Returns
        ----------
        vals : numpy.ndarray((n))
           Values of pixels in the flattened map, np.nan used to flag
           coords outside of map
        """
        theta = np.pi / 2. - np.radians(lats)
        phi = np.radians(lons)

        pix = hp.ang2pix(self.hpx.nside, theta, phi, nest=self.hpx.nest)

        if self.data.ndim == 2:
            return self.data[:, pix] if ibin is None else self.data[ibin, pix]
        else:
            return self.data[pix]

    def interpolate(self, lon, lat, egy=None, interp_log=True):
        """Interpolate map values.

        Parameters
        ----------
        interp_log : bool
            Interpolate the z-coordinate in logspace.

        """

        if self.data.ndim == 1:
            theta = np.pi / 2. - np.radians(lat)
            phi = np.radians(lon)
            return hp.pixelfunc.get_interp_val(self.counts, theta,
                                               phi, nest=self.hpx.nest)
        else:
            return self._interpolate_cube(lon, lat, egy, interp_log)

    def _interpolate_cube(self, lon, lat, egy=None, interp_log=True):
        """Perform interpolation on a healpix cube.  If egy is None
        then interpolation will be performed on the existing energy
        planes.

        """

        shape = np.broadcast(lon, lat, egy).shape
        lon = lon * np.ones(shape)
        lat = lat * np.ones(shape)
        theta = np.pi / 2. - np.radians(lat)
        phi = np.radians(lon)
        vals = []
        for i, _ in enumerate(self.hpx.evals):
            v = hp.pixelfunc.get_interp_val(self.counts[i], theta,
                                            phi, nest=self.hpx.nest)
            vals += [np.expand_dims(np.array(v, ndmin=1), -1)]

        vals = np.concatenate(vals, axis=-1)

        if egy is None:
            return vals.T

        egy = egy * np.ones(shape)

        if interp_log:
            xvals = utils.val_to_pix(np.log(self.hpx.evals), np.log(egy))
        else:
            xvals = utils.val_to_pix(self.hpx.evals, egy)

        vals = vals.reshape((-1, vals.shape[-1]))
        xvals = np.ravel(xvals)
        v = map_coordinates(vals, [np.arange(vals.shape[0]), xvals],
                            order=1)
        return v.reshape(shape)

    def swap_scheme(self):
        """
        """
        hpx_out = self.hpx.make_swapped_hpx()
        if self.hpx.nest:
            if self.data.ndim == 2:
                data_out = np.vstack([hp.pixelfunc.reorder(
                    self.data[i], n2r=True) for i in range(self.data.shape[0])])
            else:
                data_out = hp.pixelfunc.reorder(self.data, n2r=True)
        else:
            if self.data.ndim == 2:
                data_out = np.vstack([hp.pixelfunc.reorder(
                    self.data[i], r2n=True) for i in range(self.data.shape[0])])
            else:
                data_out = hp.pixelfunc.reorder(self.data, r2n=True)
        return HpxMap(data_out, hpx_out)

    def expanded_counts_map(self):
        """ return the full counts map """
        if self.hpx._ipix is None:
            return self.counts

        output = np.zeros(
            (self.counts.shape[0], self.hpx._maxpix), self.counts.dtype)
        for i in range(self.counts.shape[0]):
            output[i][self.hpx._ipix] = self.counts[i]
        return output

    def explicit_counts_map(self, pixels=None):
        """ return a counts map with explicit index scheme

        Parameters
        ----------
        pixels : `np.ndarray` or None
            If set, grab only those pixels.  
            If none, grab only non-zero pixels
        """
        # No pixel index, so build one
        if self.hpx._ipix is None:
            if self.data.ndim == 2:
                summed = self.counts.sum(0)
                if pixels is None:
                    nz = summed.nonzero()[0]
                else:
                    nz = pixels
                data_out = np.vstack(self.data[i].flat[nz]
                                     for i in range(self.data.shape[0]))
            else:
                if pixels is None:
                    nz = self.data.nonzero()[0]
                else:
                    nz = pixels
                data_out = self.data[nz]
            return (nz, data_out)
        else:
            if pixels is None:
                return (self.hpx._ipix, self.data)
        # FIXME, can we catch this
        raise RuntimeError(
            'HPX.explicit_counts_map called with pixels for a map that already has pixels')

    def sparse_counts_map(self):
        """ return a counts map with sparse index scheme
        """
        if self.hpx._ipix is None:
            flatarray = self.data.flattern()
        else:
            flatarray = self.expanded_counts_map()
        nz = flatarray.nonzero()[0]
        data_out = flatarray[nz]
        return (nz, data_out)

    def ud_grade(self, order, preserve_counts=False):
        """
        """
        new_hpx = self.hpx.ud_graded_hpx(order)
        if new_hpx.evals is None:
            nebins = 1
        else:
            nebins = len(new_hpx.evals)
        shape = self.counts.shape

        if preserve_counts:
            power = -2.
        else:
            power = 0

        if len(shape) == 1:
            new_data = hp.pixelfunc.ud_grade(self.counts,
                                             nside_out=new_hpx.nside,
                                             order_in=new_hpx.ordering,
                                             order_out=new_hpx.ordering,
                                             power=power)
        else:
            new_data = np.vstack([hp.pixelfunc.ud_grade(self.counts[i],
                                                        nside_out=new_hpx.nside,
                                                        order_in=new_hpx.ordering,
                                                        order_out=new_hpx.ordering,
                                                        power=power) for i in range(shape[0])])
        return HpxMap(new_data, new_hpx)
Exemple #7
0
class HpxMap(Map_Base):
    """ Representation of a 2D or 3D counts map using HEALPix. """

    def __init__(self, counts, hpx):
        """ C'tor, fill with a counts vector and a HPX object """
        Map_Base.__init__(self, counts)
        self._hpx = hpx
        self._wcs2d = None
        self._hpx2wcs = None

    @property
    def hpx(self):
        return self._hpx

    @classmethod
    def create_from_hdu(cls, hdu, ebins):
        """ Creates and returns an HpxMap object from a FITS HDU.

        hdu    : The FITS
        ebins  : Energy bin edges [optional]
        """
        hpx = HPX.create_from_header(hdu.header, ebins)
        colnames = hdu.columns.names
        cnames = []
        if hpx.conv.convname == 'FGST_SRCMAP_SPARSE':
            keys = hdu.data.field('KEY')
            vals = hdu.data.field('VALUE')
            nebin = len(ebins)
            data = np.zeros((nebin, hpx.npix))
            data.flat[keys] = vals
        else:
            for c in colnames:
                if c.find(hpx.conv.colstring) == 0:
                    cnames.append(c)
            nebin = len(cnames)
            data = np.ndarray((nebin, hpx.npix))
            for i, cname in enumerate(cnames):
                data[i, 0:] = hdu.data.field(cname)

        return cls(data, hpx)

    @classmethod
    def create_from_hdulist(cls, hdulist, **kwargs):
        """ Creates and returns an HpxMap object from a FITS HDUList

        extname : The name of the HDU with the map data
        ebounds : The name of the HDU with the energy bin data
        """
        extname = kwargs.get('hdu', 'SKYMAP')
        ebins = fits_utils.find_and_read_ebins(hdulist)
        return cls.create_from_hdu(hdulist[extname], ebins)

    def create_image_hdu(self, name=None, **kwargs):
        kwargs['extname'] = name
        return self.hpx.make_hdu(self.counts, **kwargs)

    @classmethod
    def create_from_fits(cls, fitsfile, **kwargs):
        hdulist = fits.open(fitsfile)
        return cls.create_from_hdulist(hdulist, **kwargs)

    def create_image_hdu(self, name=None, **kwargs):
        kwargs['extname'] = name
        return self.hpx.make_hdu(self.counts, **kwargs)

    def make_wcs_from_hpx(self, sum_ebins=False, proj='CAR', oversample=2,
                          normalize=True):
        """Make a WCS object and convert HEALPix data into WCS projection

        NOTE: this re-calculates the mapping, if you have already
        calculated the mapping it is much faster to use
        convert_to_cached_wcs() instead

        Parameters
        ----------
        sum_ebins  : bool
           sum energy bins over energy bins before reprojecting

        proj       : str
           WCS-projection

        oversample : int
           Oversampling factor for WCS map

        normalize  : bool
           True -> perserve integral by splitting HEALPix values between bins

        returns (WCS object, np.ndarray() with reprojected data)

        """
        self._wcs_proj = proj
        self._wcs_oversample = oversample
        self._wcs_2d = self.hpx.make_wcs(2, proj=proj, oversample=oversample)
        self._hpx2wcs = HpxToWcsMapping(self.hpx, self._wcs_2d)
        wcs, wcs_data = self.convert_to_cached_wcs(self.counts, sum_ebins,
                                                   normalize)
        return wcs, wcs_data

    def convert_to_cached_wcs(self, hpx_in, sum_ebins=False, normalize=True):
        """ Make a WCS object and convert HEALPix data into WCS projection

        Parameters
        ----------
        hpx_in     : `~numpy.ndarray`
           HEALPix input data
        sum_ebins  : bool
           sum energy bins over energy bins before reprojecting
        normalize  : bool
           True -> perserve integral by splitting HEALPix values between bins

        returns (WCS object, np.ndarray() with reprojected data)
        """
        if self._hpx2wcs is None:
            raise Exception('HpxMap.convert_to_cached_wcs() called '
                            'before make_wcs_from_hpx()')

        if len(hpx_in.shape) == 1:
            wcs_data = np.ndarray(self._hpx2wcs.npix)
            loop_ebins = False
            hpx_data = hpx_in
        elif len(hpx_in.shape) == 2:
            if sum_ebins:
                wcs_data = np.ndarray(self._hpx2wcs.npix)
                hpx_data = hpx_in.sum(0)
                loop_ebins = False
            else:
                wcs_data = np.ndarray((self.counts.shape[0],
                                       self._hpx2wcs.npix[0],
                                       self._hpx2wcs.npix[1]))
                hpx_data = hpx_in
                loop_ebins = True
        else:
            raise Exception('Wrong dimension for HpxMap %i' %
                            len(hpx_in.shape))

        if loop_ebins:
            for i in range(hpx_data.shape[0]):
                self._hpx2wcs.fill_wcs_map_from_hpx_data(
                    hpx_data[i], wcs_data[i], normalize)
                pass
            wcs_data.reshape((self.counts.shape[0], self._hpx2wcs.npix[
                             0], self._hpx2wcs.npix[1]))
            # replace the WCS with a 3D one
            wcs = self.hpx.make_wcs(3, proj=self._wcs_proj,
                                    energies=self.hpx.ebins,
                                    oversample=self._wcs_oversample)
        else:
            self._hpx2wcs.fill_wcs_map_from_hpx_data(
                hpx_data, wcs_data, normalize)
            wcs_data.reshape(self._hpx2wcs.npix)
            wcs = self._wcs_2d

        return wcs, wcs_data

    def get_pixel_skydirs(self):
        """Get a list of sky coordinates for the centers of every pixel. """
        sky_coords = self._hpx.get_sky_coords()
        if self.hpx.coordsys == 'GAL':
            frame = Galactic
        else:
            frame = ICRS
        return SkyCoord(sky_coords[0], sky_coords[1], frame=frame, unit='deg')

    def get_pixel_indices(self, lats, lons):
        """Return the indices in the flat array corresponding to a set of coordinates """
        return self._hpx.get_pixel_indices(lats, lons)

    def sum_over_energy(self):
        """ Reduce a counts cube to a counts map """
        # We sum over axis 0 in the array, and drop the energy binning in the
        # hpx object
        return HpxMap(np.sum(self.counts, axis=0), self.hpx.copy_and_drop_energy())

    def get_map_values(self, lons, lats, ibin=None):
        """Return the indices in the flat array corresponding to a set of coordinates

        Parameters
        ----------
        lons  : array-like
           'Longitudes' (RA or GLON)

        lats  : array-like
           'Latitidues' (DEC or GLAT)

        ibin : int or array-like
           Extract data only for a given energy bin.  None -> extract data for all bins

        Returns
        ----------
        vals : numpy.ndarray((n))
           Values of pixels in the flattened map, np.nan used to flag
           coords outside of map
        """
        theta = np.pi / 2. - np.radians(lats)
        phi = np.radians(lons)

        pix = hp.ang2pix(self.hpx.nside, theta, phi, nest=self.hpx.nest)

        if self.data.ndim == 2:
            return self.data[:, pix] if ibin is None else self.data[ibin, pix]
        else:
            return self.data[pix]

    def interpolate(self, lon, lat, egy=None, interp_log=True):
        """Interpolate map values.

        Parameters
        ----------
        interp_log : bool
            Interpolate the z-coordinate in logspace.

        """

        if self.data.ndim == 1:
            theta = np.pi / 2. - np.radians(lat)
            phi = np.radians(lon)
            return hp.pixelfunc.get_interp_val(self.counts, theta,
                                               phi, nest=self.hpx.nest)
        else:
            return self._interpolate_cube(lon, lat, egy, interp_log)

    def _interpolate_cube(self, lon, lat, egy=None, interp_log=True):
        """Perform interpolation on a healpix cube.  If egy is None
        then interpolation will be performed on the existing energy
        planes.

        """

        shape = np.broadcast(lon, lat, egy).shape
        lon = lon * np.ones(shape)
        lat = lat * np.ones(shape)
        theta = np.pi / 2. - np.radians(lat)
        phi = np.radians(lon)
        vals = []
        for i, _ in enumerate(self.hpx.evals):
            v = hp.pixelfunc.get_interp_val(self.counts[i], theta,
                                            phi, nest=self.hpx.nest)
            vals += [np.expand_dims(np.array(v, ndmin=1), -1)]

        vals = np.concatenate(vals, axis=-1)

        if egy is None:
            return vals.T

        egy = egy * np.ones(shape)

        if interp_log:
            xvals = utils.val_to_pix(np.log(self.hpx.evals), np.log(egy))
        else:
            xvals = utils.val_to_pix(self.hpx.evals, egy)

        vals = vals.reshape((-1, vals.shape[-1]))
        xvals = np.ravel(xvals)
        v = map_coordinates(vals, [np.arange(vals.shape[0]), xvals],
                            order=1)
        return v.reshape(shape)

    def swap_scheme(self):
        """
        """
        hpx_out = self.hpx.make_swapped_hpx()
        if self.hpx.nest:
            if self.data.ndim == 2:
                data_out = np.vstack([hp.pixelfunc.reorder(
                    self.data[i], n2r=True) for i in range(self.data.shape[0])])
            else:
                data_out = hp.pixelfunc.reorder(self.data, n2r=True)
        else:
            if self.data.ndim == 2:
                data_out = np.vstack([hp.pixelfunc.reorder(
                    self.data[i], r2n=True) for i in range(self.data.shape[0])])
            else:
                data_out = hp.pixelfunc.reorder(self.data, r2n=True)
        return HpxMap(data_out, hpx_out)

    def ud_grade(self, order, preserve_counts=False):
        """
        """
        new_hpx = self.hpx.ud_graded_hpx(order)
        nebins = len(new_hpx.evals)
        shape = self.counts.shape

        if preserve_counts:
            power = -2.
        else:
            power = 0

        if len(shape) == 1:
            new_data = hp.pixelfunc.ud_grade(self.counts,
                                             nside_out=new_hpx.nside,
                                             order_in=new_hpx.ordering,
                                             order_out=new_hpx.ordering,
                                             power=power)
        else:
            new_data = np.vstack([hp.pixelfunc.ud_grade(self.counts[i],
                                                        nside_out=new_hpx.nside,
                                                        order_in=new_hpx.ordering,
                                                        order_out=new_hpx.ordering,
                                                        power=power) for i in range(shape[0])])
        return HpxMap(new_data, new_hpx)