Merge pull request #1317 from registerrier/mapaxis_type

Add function to compute counts maps

gammapy/cube/__init__.py

@@ -8,3 +8,4 @@
 from .utils import *
 from .models import *
 from .cube_pipe import *
+from .basic_cube import *

gammapy/cube/basic_cube.py

@@ -0,0 +1,44 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+"""Functions to perform basic functions for map and cube analysis.
+"""
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+__all__ = [
+    'fill_map_counts',
+]
+
+
+def fill_map_counts(count_map, event_list):
+    """Fill events into a counts map.
+
+    The energy of the events is used for a non-spatial axis homogeneous to energy.
+    The other non-spatial axis names should have an entry in the column names of the ``EventList``
+
+    Parameters
+    ----------
+    count_map : `~gammapy.maps.Map`
+        Map object, will be filled by this function.
+    event_list : `~gammapy.data.EventList`
+        Event list
+    """
+    geom = count_map.geom
+
+    # Make a coordinate dictionary; skycoord is always added
+    coord_dict = dict(skycoord=event_list.radec)
+
+    # Now add one coordinate for each extra map axis
+    for axis in geom.axes:
+        if axis.type == 'energy':
+            # This axis is the energy. We treat it differently because axis.name could be e.g. 'energy_reco'
+            coord_dict[axis.name] = event_list.energy.to(axis.unit)
+        # TODO: add proper extraction for time
+        else:
+            # We look for other axes name in the table column names (case insensitive)
+            colnames = [_.upper() for _ in event_list.table.colnames]
+            if axis.name.upper() in colnames:
+                column_name = event_list.table.colnames[colnames.index(axis.name.upper())]
+                coord_dict.update({axis.name: event_list.table[column_name].to(axis.unit)})
+            else:
+                raise ValueError("Cannot find MapGeom axis {!r} in EventList".format(axis.name))
+
+    count_map.fill_by_coord(coord_dict)

gammapy/cube/tests/test_basic_cube.py

@@ -0,0 +1,87 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+import pytest
+import numpy as np
+from astropy.coordinates import SkyCoord
+import astropy.units as u
+from ...utils.testing import requires_dependency, requires_data
+from ...maps import MapAxis, WcsGeom, HpxGeom, Map
+from ...data import DataStore, EventList
+from ...cube.basic_cube import fill_map_counts
+
+
+pytest.importorskip('scipy')
+
+# TODO: move these two cases to different test functions?
+# Would allow to change the asserts in `test_fill_map_counts` to something much more specific / useful, no?
+axis_energy_reco = MapAxis(np.logspace(-1., 1.5, 10), interp='log', node_type='edge',
+                           name='energy_reco', unit='TeV')
+
+# initial time of run 110380 from cta 1DC
+times = np.linspace(664504199, 664504199 + 1900., 10)
+axis_time = MapAxis(times, node_type='edge', name='time', unit='s')
+
+pos_cta = SkyCoord(0.0, 0.0, frame='galactic', unit='deg')
+
+geom_cta = {'binsz': 0.02, 'coordsys': 'GAL', 'width': 15 * u.deg,
+            'skydir': pos_cta, 'axes': [axis_energy_reco]}
+geom_cta_time = {'binsz': 0.02, 'coordsys': 'GAL', 'width': 15 * u.deg,
+                 'skydir': pos_cta, 'axes': [axis_energy_reco, axis_time]}
+
+
+# TODO: change the test event list to something that's created from scratch,
+# using values so that it's possible to make simple assert statements on the
+# map data in the tests below, i.e. have pixels that should receive 0, 1 or 2 counts
+def make_test_event_list():
+    ds = DataStore.from_dir('$GAMMAPY_EXTRA/datasets/cta-1dc/index/gps/')
+    return ds.obs(110380).events
+
+
+@pytest.fixture(scope='session')
+def evt_2fhl():
+    return EventList.read('$GAMMAPY_EXTRA/datasets/fermi_2fhl/2fhl_events.fits.gz')
+
+
+@requires_data('gammapy-extra')
+@pytest.mark.parametrize('geom_opts', [geom_cta, geom_cta_time])
+def test_fill_map_counts(geom_opts):
+    events = make_test_event_list()
+    geom = WcsGeom.create(**geom_opts)
+    cntmap = Map.from_geom(geom)
+
+    fill_map_counts(cntmap, events)
+
+    # Number of entries in the map
+    nmap = cntmap.data.sum()
+    # number of events
+    valid = np.ones_like(events.energy.value)
+    for axis in geom.axes:
+        if axis.name == 'energy_reco':
+            valid = np.logical_and(events.energy.value >= axis.edges[0], valid)
+            valid = np.logical_and(events.energy.value <= axis.edges[-1], valid)
+        else:
+            valid = np.logical_and(events.table[axis.name.upper()] >= axis.edges[0], valid)
+            valid = np.logical_and(events.table[axis.name.upper()] <= axis.edges[-1], valid)
+
+    nevt = valid.sum()
+    assert nmap == nevt
+
+
+@requires_data('gammapy-extra')
+@requires_dependency('healpy')
+def test_fill_map_counts_hpx(evt_2fhl):
+    # This tests healpix maps fill with non standard non spatial axis
+
+    axis_zen = MapAxis([0, 45, 180], node_type='edge', name='zenith_angle', unit=u.deg)
+    geom = HpxGeom(256, coordsys='GAL', axes=[axis_zen])
+    m = Map.from_geom(geom)
+
+    fill_map_counts(m, evt_2fhl)
+
+    nmap_l = np.sum(m.data[0])
+    nmap_h = np.sum(m.data[1])
+
+    nevt_l = np.sum(evt_2fhl.table['ZENITH_ANGLE'] < 45)
+    nevt_h = np.sum(evt_2fhl.table['ZENITH_ANGLE'] > 45)
+
+    assert nmap_l == nevt_l
+    assert nmap_h == nevt_h

gammapy/maps/geom.py

@@ -355,6 +355,14 @@ def __init__(self, nodes, interp='lin', name='',
         self._node_type = node_type
         self._unit = u.Unit('' if unit is None else unit)
 
+        # Set axis type from its unit
+        if self._unit.is_equivalent("eV"):
+            self._type = 'energy'
+        elif self._unit.is_equivalent("s"):
+            self._type = 'time'
+        else:
+            self._type = 'any'
+
         # Set pixel coordinate of first node
         if node_type == 'edge':
             self._pix_offset = -0.5
@@ -417,6 +425,11 @@ def unit(self):
         """Return coordinate axis unit."""
         return self._unit
 
+    @property
+    def type(self):
+        """Return coordinate axis type."""
+        return self._type
+
     @classmethod
     def from_bounds(cls, lo_bnd, hi_bnd, nbin, **kwargs):
         """Generate an axis object from a lower/upper bound and number of bins.