Merge pull request #1357 from registerrier/solid_angle

Add map geom pixel solid angle computation

gammapy/maps/geom.py

@@ -1296,6 +1296,11 @@ def upsample(self, factor):
         """
         pass
 
+    @abc.abstractmethod
+    def solid_angle(self):
+        """Solid angle (`~astropy.units.Quantity` in ``sr``)."""
+        pass
+
     def _fill_header_from_axes(self, header):
 
         for i, ax in enumerate(self.axes):

gammapy/maps/hpx.py

@@ -9,6 +9,7 @@
 from ..extern.six.moves import range
 from astropy.io import fits
 from astropy.coordinates import SkyCoord
+from astropy.units import Quantity
 from .wcs import WcsGeom
 from .geom import MapGeom, MapCoord, pix_tuple_to_idx
 from .geom import coordsys_to_frame, skycoord_to_lonlat
@@ -1712,6 +1713,15 @@ def skydir_to_pix(self, skydir):
 
         return self.pix_to_coord((lat, lon))
 
+    def solid_angle(self):
+        """Solid angle array (`~astropy.units.Quantity` in ``sr``).
+
+        The array has the same dimensionality as ``map.nside``
+        since all pixels have the same solid angle.
+        """
+        import healpy as hp
+        return Quantity(hp.nside2pixarea(self.nside), 'sr')
+
 
 class HpxToWcsMapping(object):
     """Stores the indices need to convert from HEALPIX to WCS.
@@ -1862,8 +1872,8 @@ def fill_wcs_map_from_hpx_data(self, hpx_data, wcs_data, normalize=True,
             wcs_data[wcs_slice] = hpx_data[hpx_slice]
 
         if fill_nan:
-            valid = np.swapaxes(self._valid.reshape(shape),-1,-2)
-            valid = valid*np.ones_like(wcs_data, dtype=bool)
+            valid = np.swapaxes(self._valid.reshape(shape), -1, -2)
+            valid = valid * np.ones_like(wcs_data, dtype=bool)
             wcs_data[~valid] = np.nan
         return wcs_data
 

gammapy/maps/tests/test_hpx.py

@@ -69,7 +69,7 @@ def test_ravel_hpx_index():
 
 
 def make_test_nside(nside, nside0, nside1):
-    npix = 12 * nside**2
+    npix = 12 * nside ** 2
     nside_test = np.concatenate((nside0 * np.ones(npix // 2, dtype=int),
                                  nside1 * np.ones(npix // 2, dtype=int)))
     return nside_test
@@ -79,10 +79,9 @@ def make_test_nside(nside, nside0, nside1):
                          [(4, 2, True), (8, 2, True), (8, make_test_nside(8, 4, 2), True),
                           (4, 2, False), (8, 2, False), (8, make_test_nside(8, 4, 2), False), ])
 def test_get_superpixels(nside_subpix, nside_superpix, nest):
-
     import healpy as hp
 
-    npix = 12 * nside_subpix**2
+    npix = 12 * nside_subpix ** 2
     subpix = np.arange(npix)
     ang_subpix = hp.pix2ang(nside_subpix, subpix, nest=nest)
     superpix = get_superpixels(subpix, nside_subpix, nside_superpix, nest=nest)
@@ -107,7 +106,7 @@ def test_get_superpixels(nside_subpix, nside_superpix, nest):
 def test_get_subpixels(nside_superpix, nside_subpix, nest):
     import healpy as hp
 
-    npix = 12 * nside_superpix**2
+    npix = 12 * nside_superpix ** 2
     superpix = np.arange(npix)
     subpix = get_subpixels(superpix, nside_superpix, nside_subpix, nest=nest)
     ang1 = hp.pix2ang(nside_subpix, subpix, nest=nest)
@@ -532,3 +531,14 @@ def test_hpxgeom_downsample(nside, nested, coordsys, region, axes):
     assert_allclose(geom.nside, 2 * geom_down.nside)
     coords = geom.get_coord(flat=True)
     assert np.all(geom_down.contains(coords))
+
+
+def test_hpxgeom_solid_angle():
+    geom = HpxGeom.create(nside=8, coordsys='GAL',
+                          axes=[MapAxis.from_edges([0, 2, 3])])
+
+    solid_angle = geom.solid_angle()
+    print(float(solid_angle.value))
+
+    assert solid_angle.shape == (1,)
+    assert_allclose(solid_angle.value, 0.016362461737446838)

gammapy/maps/tests/test_wcs.py

@@ -5,6 +5,7 @@
 from numpy.testing import assert_allclose
 from astropy.io import fits
 from astropy.coordinates import SkyCoord
+import astropy.units as u
 from ..wcs import WcsGeom
 from ..geom import MapAxis
 
@@ -74,7 +75,7 @@ def test_wcsgeom_test_coord_to_idx(npix, binsz, coordsys, proj, skydir, axes):
 
     if not geom.is_allsky:
         coords = geom.center_coord[:2] + \
-            tuple([ax.center[0] for ax in geom.axes])
+                 tuple([ax.center[0] for ax in geom.axes])
         coords[0][...] += 2.0 * np.max(geom.width[0])
         idx = geom.coord_to_idx(coords)
         assert_allclose(np.full_like(coords[0], -1, dtype=int), idx[0])
@@ -121,3 +122,37 @@ def test_wcsgeom_contains(npix, binsz, coordsys, proj, skydir, axes):
     if not geom.is_allsky:
         coords = [0.0, 0.0] + [ax.center[0] for ax in geom.axes]
         assert_allclose(geom.contains(coords), np.zeros((1,), dtype=bool))
+
+
+def test_wcsgeom_solid_angle():
+    # Test using a CAR projection map with an extra axis
+    binsz = 1.0 * u.deg
+    npix = 10
+    geom = WcsGeom.create(skydir=(0, 0), npix=(npix, npix),
+                          binsz=binsz, coordsys='GAL', proj='CAR',
+                          axes=[MapAxis.from_edges([0, 2, 3])])
+
+    solid_angle = geom.solid_angle()
+
+    # Check array size
+    assert solid_angle.shape == (2, npix, npix)
+
+    # Note: test is valid for small enough bin sizes since
+    # WcsGeom.solid_angle() approximates the true solid angle value
+
+    # Test at b = 0 deg
+    solid_lat0 = binsz.to('rad').value * (np.sin(binsz)) * u.sr
+    assert_allclose(solid_angle[0, 5, 5], solid_lat0, rtol=1e-3)
+
+    # Test at b = 5 deg
+    solid_lat5 = binsz.to('rad').value * (np.sin(5 * binsz) - np.sin(4 * binsz.to('rad').value)) * u.sr
+    assert_allclose(solid_angle[0, 9, 5], solid_lat5, rtol=1e-3)
+
+
+def test_wcsgeom_solid_angle_ait():
+    # Pixels that don't correspond to locations on ths sky
+    # should have solid angles set to NaN
+    ait_geom = WcsGeom.create(skydir=(0, 0), npix=(10, 4), binsz=50,
+                              coordsys='GAL', proj='AIT')
+    solid_angle = ait_geom.solid_angle()
+    assert np.isnan(solid_angle[0, 0])

gammapy/maps/wcs.py

@@ -5,6 +5,8 @@
 from astropy.wcs import WCS
 from astropy.io import fits
 from astropy.coordinates import SkyCoord
+from astropy.coordinates.angle_utilities import angular_separation
+import astropy.units as u
 from ..image.utils import make_header
 from ..utils.wcs import get_resampled_wcs
 from .geom import MapGeom, MapCoord, pix_tuple_to_idx, skycoord_to_lonlat
@@ -400,8 +402,8 @@ def _get_pix_coords(self, idx=None, mode='center'):
         npix = copy.deepcopy(self.npix)
 
         if mode == 'edge':
-            npix[0] += 1
-            npix[1] += 1
+            for pix_num in npix[:2]:
+                pix_num += 1
 
         if self.axes and not self.is_regular:
 
@@ -427,7 +429,7 @@ def _get_pix_coords(self, idx=None, mode='center'):
                     s_img = (slice(0, npix0), slice(0, npix1),) + idx_img
                 else:
                     s_img = (slice(0, npix0), slice(0, npix1),) + \
-                        (0,) * len(self.axes)
+                            (0,) * len(self.axes)
 
                 pix2[0][s_img] = pix_img[0]
                 pix2[1][s_img] = pix_img[1]
@@ -455,10 +457,10 @@ def _get_pix_coords(self, idx=None, mode='center'):
             pix[i][~m] = np.nan
         return pix
 
-#        shape = np.broadcast(*coords).shape
-#        m = [np.isfinite(c) for c in coords]
-#        m = np.broadcast_to(np.prod(m),shape)
-#        return tuple([np.ravel(np.broadcast_to(t,shape)[m]) for t in pix])
+    #        shape = np.broadcast(*coords).shape
+    #        m = [np.isfinite(c) for c in coords]
+    #        m = np.broadcast_to(np.prod(m),shape)
+    #        return tuple([np.ravel(np.broadcast_to(t,shape)[m]) for t in pix])
 
     def get_coord(self, idx=None, flat=False):
         pix = self._get_pix_coords(idx=idx)
@@ -531,7 +533,7 @@ def pix_to_idx(self, pix, clip=False):
                     np.putmask(idxs[i], (idx < 0) | (idx >= npix[i]), -1)
                 else:
                     np.putmask(idxs[i], (idx < 0) | (
-                        idx >= self.axes[i - 2].nbin), -1)
+                            idx >= self.axes[i - 2].nbin), -1)
 
         return idxs
 
@@ -594,6 +596,33 @@ def upsample(self, factor):
     def to_slice(self, slices):
         raise NotImplementedError
 
+    def solid_angle(self):
+        """Solid angle array (`~astropy.units.Quantity` in ``sr``).
+
+        The array has the same dimension as the WcsGeom object.
+        To return solid angles for the spatial dimensions only use: WcsGeom.to_image().solid_angle()
+        """
+        # TODO: Improve by exposing a mode 'edge' for get_coord
+        # Note that edge is applied only to spatial coordinates in the following call
+        # Note also that pix_to_coord is already called in _get_pix_coords.
+        # This should be made more efficient.
+        pix = self._get_pix_coords(mode='edge')
+        coord = self.pix_to_coord(pix)
+        lon = coord[0] * np.pi / 180.
+        lat = coord[1] * np.pi / 180.
+
+        # Compute solid angle using the approximation that it's
+        # the product between angular separation of pixel corners.
+        # First index is "y", second index is "x"
+        ylo_xlo = lon[..., :-1, :-1], lat[..., :-1, :-1]
+        ylo_xhi = lon[..., :-1, 1:], lat[..., :-1, 1:]
+        yhi_xlo = lon[..., 1:, :-1], lat[..., 1:, :-1]
+
+        dx = angular_separation(*(ylo_xlo + ylo_xhi))
+        dy = angular_separation(*(ylo_xlo + yhi_xlo))
+
+        return u.Quantity(dx * dy, 'sr')
+
 
 def create_wcs(skydir, coordsys='CEL', projection='AIT',
                cdelt=1.0, crpix=1., axes=None):