Add and improve "edges" mode for map pix coordinates

gammapy/cube/tests/test_models.py

@@ -241,32 +241,32 @@ def test_compute_dnde(evaluator):
         out = evaluator.compute_dnde()
         assert out.shape == (2, 4, 5)
         assert out.unit == 'cm-2 s-1 TeV-1 deg-2'
-        assert_allclose(out.value.mean(), 7.460919e-14)
+        assert_allclose(out.value.mean(), 7.460919e-14, rtol=1e-5)
 
     @staticmethod
     def test_compute_flux(evaluator):
         out = evaluator.compute_flux()
         assert out.shape == (2, 4, 5)
         assert out.unit == 'cm-2 s-1'
-        assert_allclose(out.value.mean(), 1.828206748668197e-14)
+        assert_allclose(out.value.mean(), 1.828206748668197e-14, rtol=1e-5)
 
     @staticmethod
     def test_apply_psf(evaluator):
         flux = evaluator.compute_flux()
         npred = evaluator.apply_exposure(flux)
         out = evaluator.apply_psf(npred)
         assert out.data.shape == (2, 4, 5)
-        assert_allclose(out.data.mean(), 1.2574065e-08)
+        assert_allclose(out.data.mean(), 1.2574065e-08, rtol=1e-5)
 
     @staticmethod
     def test_apply_edisp(evaluator):
         flux = evaluator.compute_flux()
         out = evaluator.apply_edisp(flux.value)
         assert out.shape == (2, 4, 5)
-        assert_allclose(out.mean(), 1.828206748668197e-14)
+        assert_allclose(out.mean(), 1.828206748668197e-14, rtol=1e-5)
 
     @staticmethod
     def test_compute_npred(evaluator):
         out = evaluator.compute_npred()
         assert out.shape == (2, 4, 5)
-        assert_allclose(out.sum(), 45.02963e-07)
+        assert_allclose(out.sum(), 45.02963e-07, rtol=1e-5)

gammapy/cube/tests/test_new.py

@@ -83,8 +83,8 @@ def test_make_map_fov_background(bkg_3d, counts_cube):
     )
 
     assert m.data.shape == (15, 120, 200)
-    assert_allclose(m.data[0, 0, 0], 0.013959366925415072)
-    assert_allclose(m.data.sum(), 1356.2551841113177)
+    assert_allclose(m.data[0, 0, 0], 0.013959, rtol=1e-4)
+    assert_allclose(m.data.sum(), 1356.2551, rtol=1e-5)
 
     # TODO: Check that `offset_max` is working properly
     # pos = SkyCoord(85.6, 23, unit='deg')

gammapy/maps/tests/test_wcs.py

@@ -4,6 +4,7 @@
 import numpy as np
 from numpy.testing import assert_allclose
 from astropy.io import fits
+from astropy.tests.helper import assert_quantity_allclose
 from astropy.coordinates import SkyCoord
 import astropy.units as u
 from ..wcs import WcsGeom
@@ -137,16 +138,11 @@ def test_wcsgeom_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)
+    assert_quantity_allclose(solid_angle[0, 5, 5], 0.0003046 * u.sr, 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)
+    assert_quantity_allclose(solid_angle[0, 9, 5], 0.0003038  * u.sr, rtol=1e-3)
 
 
 def test_wcsgeom_solid_angle_ait():
@@ -158,6 +154,30 @@ def test_wcsgeom_solid_angle_ait():
     assert np.isnan(solid_angle[0, 0])
 
 
+def test_wcsgeom_get_coord():
+    geom = WcsGeom.create(skydir=(0, 0), npix=(4, 3), binsz=1,
+                          coordsys='GAL', proj='CAR')
+    coord = geom.get_coord(mode='edges')
+    assert_allclose(coord.lon[0, 0], 2)
+    assert_allclose(coord.lat[0, 0], -1.5)
+
+
+def test_wcsgeom_get_pix_coords():
+    geom = WcsGeom.create(skydir=(0, 0), npix=(4, 3), binsz=1,
+                          coordsys='GAL', proj='CAR', axes=axes1)
+    idx_center = geom.get_pix(mode='center')
+
+    for idx in idx_center:
+        assert idx.shape == (2, 3, 4)
+        assert_allclose(idx[0, 0, 0], 0)
+
+    idx_edge = geom.get_pix(mode='edges')
+    for idx, desired in zip(idx_edge, [-0.5, -0.5, 0]):
+        assert idx.shape == (2, 4, 5)
+        assert_allclose(idx[0, 0, 0], desired)
+
+
+
 def test_geom_repr():
     geom = WcsGeom.create(skydir=(0, 0), npix=(10, 4), binsz=50,
                           coordsys='GAL', proj='AIT')

gammapy/maps/wcs.py

@@ -134,6 +134,8 @@ class WcsGeom(MapGeom):
         Serialization format convention.  This sets the default format
         that will be used when writing this geometry to a file.
     """
+    _slice_spatial_axes = slice(0, 2)
+    _slice_non_spatial_axes = slice(2, -1)
 
     def __init__(self, wcs, npix, cdelt=None, crpix=None, axes=None, conv='gadf'):
         self._wcs = wcs
@@ -469,14 +471,25 @@ def get_image_shape(self, idx):
     def get_image_wcs(self, idx):
         raise NotImplementedError
 
-    def get_idx(self, idx=None, local=False, flat=False):
-        pix = self._get_pix_coords(idx=idx, mode='center')
+    def get_idx(self, idx=None, flat=False):
+        pix = self.get_pix(idx=idx, mode='center')
         if flat:
             pix = tuple([p[np.isfinite(p)] for p in pix])
         return pix_tuple_to_idx(pix)
 
-    def _get_pix_coords(self, idx=None, mode='center'):
+    def get_pix(self, idx=None, mode='center'):
+        """Get map pix coordinates from the geometry.
 
+        Paramters
+        ---------
+        mode : {'center', 'edges'}
+            Get center or edge pix coordinates for the spatial axes.
+
+        Returns
+        -------
+        coord : tuple
+            Map pix coordinate tuple.
+        """
         # FIXME: Figure out if there is some way to employ open/sparse
         # vectors
 
@@ -486,7 +499,7 @@ def _get_pix_coords(self, idx=None, mode='center'):
 
         npix = copy.deepcopy(self.npix)
 
-        if mode == 'edge':
+        if mode == 'edges':
             for pix_num in npix[:2]:
                 pix_num += 1
 
@@ -533,32 +546,36 @@ def _get_pix_coords(self, idx=None, mode='center'):
             pix = np.meshgrid(*pix[::-1], indexing='ij', sparse=False)[::-1]
 
         if mode == 'edges':
-            for i in range(len(pix)):
-                pix[i] -= 0.5
+            for pix_array in pix[self._slice_spatial_axes]:
+                pix_array -= 0.5
 
         coords = self.pix_to_coord(pix)
         m = np.isfinite(coords[0])
         for i in range(len(pix)):
             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])
+    def get_coord(self, idx=None, flat=False, mode='center'):
+        """Get map coordinates from the geometry.
+
+        Paramters
+        ---------
+        mode : {'center', 'edges'}
+            Get center or edge coordinates for the spatial axes.
 
-    def get_coord(self, idx=None, flat=False):
-        pix = self._get_pix_coords(idx=idx)
+        Returns
+        -------
+        coord : `~MapCoord`
+            Map coordinate object.
+        """
+        pix = self.get_pix(idx=idx, mode=mode)
         coords = self.pix_to_coord(pix)
+
         if flat:
             coords = tuple([c[np.isfinite(c)] for c in coords])
 
-        cdict = OrderedDict([
-            ('lon', coords[0]),
-            ('lat', coords[1]),
-        ])
-        for i, axis in enumerate(self.axes):
-            cdict[axis.name] = coords[i + 2]
+        axes_names = ['lon', 'lat'] + [ax.name for ax in self.axes]
+        cdict = OrderedDict(zip(axes_names, coords))
 
         return MapCoord.create(cdict, coordsys=self.coordsys)
 
@@ -686,23 +703,15 @@ def upsample(self, factor):
         return self.__class__(wcs, npix, cdelt=cdelt,
                               axes=copy.deepcopy(self.axes))
 
-    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.
+        coord = self.get_coord(mode='edges')
+        lon = coord.lon * np.pi / 180.
+        lat = coord.lat * np.pi / 180.
 
         # Compute solid angle using the approximation that it's
         # the product between angular separation of pixel corners.