Merge pull request #1312 from JouvinLea/Add_BACKGROUND2D

Add Background 2D class

gammapy/irf/background.py

@@ -7,10 +7,11 @@
 from ..utils.nddata import NDDataArray, BinnedDataAxis
 from ..utils.scripts import make_path
 from ..utils.fits import fits_table_to_table, table_to_fits_table
-
 import numpy as np
+
 __all__ = [
     'Background3D',
+    'Background2D',
 ]
 
 
@@ -132,4 +133,128 @@ def to_table(self):
 
     def to_fits(self, name='BACKGROUND'):
         """Convert to `~astropy.io.fits.BinTable`."""
-        return table_to_fits_table(self.to_table(), name)
+        return table_to_fits_table(self.to_table(), name)
+
+
+class Background2D(object):
+    """Background 2D.
+
+    Data format specification: :ref:`gadf:bkg_32`
+
+    Parameters
+    -----------
+    energy_lo, energy_hi : `~astropy.units.Quantity`
+        Energy binning
+    offset_lo, offset_hi : `~astropy.units.Quantity`
+        FOV coordinate offset-axis binning
+    data : `~astropy.units.Quantity`
+        Background rate (usually: ``s^-1 MeV^-1 sr^-1``)
+    """
+    default_interp_kwargs = dict(bounds_error=False, fill_value=None)
+    """Default Interpolation kwargs for `~NDDataArray`. Extrapolate."""
+
+    def __init__(self, energy_lo, energy_hi,
+                 offset_lo, offset_hi,
+                 data, meta=None, interp_kwargs=None):
+
+        if interp_kwargs is None:
+            interp_kwargs = self.default_interp_kwargs
+        axes = [
+            BinnedDataAxis(
+                energy_lo, energy_hi,
+                interpolation_mode='log', name='energy'),
+            BinnedDataAxis(
+                offset_lo, offset_hi,
+                interpolation_mode='linear', name='offset'),
+        ]
+        self.data = NDDataArray(axes=axes, data=data,
+                                interp_kwargs=interp_kwargs)
+        self.meta = OrderedDict(meta) if meta else OrderedDict()
+
+    def __str__(self):
+        ss = self.__class__.__name__
+        ss += '\n{}'.format(self.data)
+        return ss
+
+    @classmethod
+    def from_table(cls, table):
+        """Read from `~astropy.table.Table`."""
+        # Spec says key should be "BKG", but there are files around
+        # (e.g. CTA 1DC) that use "BGD". For now we support both
+        if 'BKG' in table.colnames:
+            bkg_name = 'BKG'
+        elif 'BGD' in table.colnames:
+            bkg_name = 'BGD'
+        else:
+            raise ValueError('Invalid column names. Need "BKG" or "BGD".')
+
+        # Currently some files (e.g. CTA 1DC) contain unit in the FITS file
+        # '1/s/MeV/sr', which is invalid ( try: astropy.unit.Unit('1/s/MeV/sr')
+        # This should be corrected.
+        # For now, we hard-code the unit here:
+        data_unit = u.Unit('s-1 MeV-1 sr-1')
+        return cls(
+            energy_lo=table['ENERG_LO'].quantity[0],
+            energy_hi=table['ENERG_HI'].quantity[0],
+            offset_lo=table['THETA_LO'].quantity[0],
+            offset_hi=table['THETA_HI'].quantity[0],
+            data=table[bkg_name].data[0] * data_unit,
+            meta=table.meta,
+        )
+
+    @classmethod
+    def from_hdulist(cls, hdulist, hdu='BACKGROUND'):
+        """Create from `~astropy.io.fits.HDUList`."""
+        fits_table = hdulist[hdu]
+        table = fits_table_to_table(fits_table)
+        return cls.from_table(table)
+
+    @classmethod
+    def read(cls, filename, hdu='BACKGROUND'):
+        """Read from file."""
+        filename = make_path(filename)
+        hdulist = fits.open(str(filename))
+        return cls.from_hdulist(hdulist, hdu=hdu)
+
+    def to_table(self):
+        """Convert to `~astropy.table.Table`."""
+        meta = self.meta.copy()
+        table = Table(meta=meta)
+
+        table['THETA_LO'] = self.data.axis('offset').lo[np.newaxis]
+        table['THETA_HI'] = self.data.axis('offset').hi[np.newaxis]
+        table['ENERG_LO'] = self.data.axis('energy').lo[np.newaxis]
+        table['ENERG_HI'] = self.data.axis('energy').hi[np.newaxis]
+        table['BKG'] = self.data.data[np.newaxis]
+        return table
+
+    def to_fits(self, name='BACKGROUND'):
+        """Convert to `~astropy.io.fits.BinTable`."""
+        return table_to_fits_table(self.to_table(), name)
+
+    def evaluate(self, fov_offset, fov_phi=None, energy_reco=None, **kwargs):
+        """
+        Evaluate the `Background2D` at a given offset and energy.
+
+        Parameters
+        ----------
+        fov_offset : `~astropy.coordinates.Angle`
+            Offset in the FOV
+        fov_phi: `~astropy.coordinates.Angle`
+            Azimuth angle in the FOV.
+            Not used for this class since the background model is radially symmetric
+        energy_reco : `~astropy.units.Quantity`
+            Reconstructed energy
+        kwargs : dict
+            option for interpolation for `~scipy.interpolate.RegularGridInterpolator`
+
+        Returns
+        -------
+        array : `~astropy.units.Quantity`
+            Interpolated values, axis order is the same as for the NDData array
+        """
+        if energy_reco is None:
+            energy_reco = self.data.axis('energy').nodes
+
+        array = self.data.evaluate(offset=fov_offset, energy=energy_reco, **kwargs)
+        return array

gammapy/irf/tests/test_background.py

@@ -2,9 +2,12 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 import pytest
 import astropy.units as u
+import numpy as np
 from numpy.testing import assert_allclose, assert_equal
+from astropy.tests.helper import assert_quantity_allclose
+from astropy.coordinates import Angle
 from ...utils.testing import requires_dependency, requires_data
-from ..background import Background3D
+from ..background import Background3D, Background2D
 from ...utils.fits import table_to_fits_table
 
 
@@ -41,11 +44,68 @@ def test_background_3d_basics(bkg_3d):
 def test_background_3d_evaluate(bkg_3d):
     bkg_rate = bkg_3d.data.evaluate(energy='1 TeV', detx='0.2 deg', dety='0.5 deg')
     assert_allclose(bkg_rate.value, 0.00013352689711418575)
-    assert bkg_rate.unit == u.Unit('s-1 MeV-1 sr-1')
+    assert bkg_rate.unit == 's-1 MeV-1 sr-1'
 
 
 @requires_data('gammapy-extra')
 def test_background_3d_write(bkg_3d):
     hdu = table_to_fits_table(bkg_3d.to_table())
     assert_equal(hdu.data['DETX_LO'][0], bkg_3d.data.axis('detx').lo.value)
     assert hdu.header['TUNIT1'] == bkg_3d.data.axis('detx').lo.unit
+
+
+@pytest.fixture(scope='session')
+def bkg_2d():
+    """A simple Background2D test case"""
+    energy = [1, 10, 100] * u.TeV
+    offset = [0, 1, 2, 3] * u.deg
+    data = np.zeros((2, 3)) * u.Unit('s-1 MeV-1 sr-1')
+    data.value[1, 0] = 2
+    data.value[1, 1] = 4
+    return Background2D(
+        energy_lo=energy[:-1], energy_hi=energy[1:],
+        offset_lo=offset[:-1], offset_hi=offset[1:],
+        data=data,
+    )
+
+
+@requires_dependency('scipy')
+def test_background_2d_evaluate(bkg_2d):
+    # TODO: the test cases here can probably be improved a bit
+    # There's some redundancy, and no case exactly at a node in energy
+
+    # Evaluate at log center between nodes in energy
+    res = bkg_2d.evaluate(fov_offset=[1, 0.5] * u.deg, energy_reco=3.16227766 * u.TeV)
+    assert_allclose(res.value, 0)
+    assert res.shape == (2,)
+    assert res.unit == 's-1 MeV-1 sr-1'
+
+    res = bkg_2d.evaluate(fov_offset=[1, 0.5] * u.deg, energy_reco=31.6227766 * u.TeV)
+    assert_allclose(res.value, [3, 2])
+
+    res = bkg_2d.evaluate(fov_offset=[1, 0.5] * u.deg, energy_reco=[3.16227766, 31.6227766] * u.TeV)
+    assert_allclose(res.value, [[0, 0], [3, 2]])
+    assert res.shape == (2, 2)
+
+    res = bkg_2d.evaluate(fov_offset=1 * u.deg, energy_reco=[3.16227766, 31.6227766] * u.TeV)
+    assert_allclose(res.value, [0, 3])
+    assert res.shape == (2,)
+
+
+def test_background_2d_read_write(tmpdir, bkg_2d):
+    filename = str(tmpdir / "bkg2d.fits")
+    bkg_2d.to_fits().writeto(filename)
+
+    bkg_2d_2 = Background2D.read(filename)
+
+    axis = bkg_2d_2.data.axis('energy')
+    assert axis.nbins == 2
+    assert axis.unit == 'TeV'
+
+    axis = bkg_2d_2.data.axis('offset')
+    assert axis.nbins == 3
+    assert axis.unit == 'deg'
+
+    data = bkg_2d_2.data.data
+    assert data.shape == (2, 3)
+    assert data.unit == 's-1 MeV-1 sr-1'