Merge pull request #1211 from thomasarmstrong/implement_irf_write

Add IRF write methods

examples/example_write_irf.py

@@ -0,0 +1,77 @@
+"""
+Example file how to create fits file using gammapy.irf classes.
+
+"""
+
+import numpy as np
+import astropy.units as u
+from gammapy.irf import EffectiveAreaTable2D, EnergyDispersion2D, Background3D
+from collections import OrderedDict
+from astropy.io import fits
+from gammapy.utils.fits import table_to_fits_table
+
+
+provenance = OrderedDict([('ORIGIN', 'IRAP'),
+                          ('DATE', '2017-09-27T12:02:24'),
+                          ('TELESCOP', 'CTA'),
+                          ('INSTRUME', 'PROD3B'),
+                          ('ETC', 'ETC')])
+
+# Set up some example data
+energy = np.logspace(0, 1, 11) * u.TeV
+energy_lo = energy[:-1]
+energy_hi = energy[1:]
+offset = np.linspace(0, 1, 4) * u.deg
+offset_lo = offset[:-1]
+offset_hi = offset[1:]
+migra = np.linspace(0,3,4)
+migra_lo = migra[:-1]
+migra_hi = migra[1:]
+detx = np.linspace(-6,6,11) * u.deg
+detx_lo = detx[:-1]
+detx_hi = detx[1:]
+dety = np.linspace(-6,6,11) * u.deg
+dety_lo = dety[:-1]
+dety_hi = dety[1:]
+aeff_data = np.ones(shape=(10,3))*u.cm*u.cm
+edisp_data = np.ones(shape=(10, 3, 3))
+bkg_data = np.ones(shape=(10,10,10)) / u.MeV / u.s / u.sr
+
+
+# Create IRF Class objects with data
+aeff = EffectiveAreaTable2D(energy_lo=energy_lo, energy_hi=energy_hi,
+                            offset_lo=offset_lo, offset_hi=offset_hi,
+                            data=aeff_data)
+
+edisp = EnergyDispersion2D(e_true_lo=energy_lo, e_true_hi=energy_hi,
+                           migra_lo=migra_lo, migra_hi=migra_hi,
+                           offset_lo=offset_lo, offset_hi=offset_hi,
+                           data=edisp_data)
+
+bkg = Background3D(energy_lo=energy_lo, energy_hi=energy_hi,
+                   detx_lo=detx_lo, detx_hi=detx_hi,
+                   dety_lo=dety_lo, dety_hi=dety_hi,
+                   data=bkg_data)
+
+# Convert to astropy Table objects
+table_aeff = aeff.to_table()
+table_edisp = edisp.to_table()
+table_bkg = bkg.to_table()
+
+# Add any information that needs to be in the fits header
+table_aeff.meta.update(provenance)
+table_edisp.meta.update(provenance)
+table_bkg.meta.update(provenance)
+
+# Convert to fits HDU objects
+hdu_aeff = table_to_fits_table(table_aeff, name='EFFECTIVE AREA')
+hdu_edisp = table_to_fits_table(table_edisp, name='ENERGY DISPERSION')
+hdu_bkg = table_to_fits_table(table_bkg, name='BACKGROUND')
+prim_hdu = fits.PrimaryHDU()
+
+# Alternatively, HDU can be obtained directly:
+# hdu_aeff = aeff.to_fits(name='EFFECTIVE AREA')
+# ...
+
+
+fits.HDUList([prim_hdu, hdu_aeff, hdu_edisp, hdu_bkg]).writeto('irf_test.fits')

gammapy/irf/background.py

@@ -1,12 +1,14 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 from __future__ import absolute_import, division, print_function, unicode_literals
 from collections import OrderedDict
+from astropy.table import Table
 from astropy.io import fits
 import astropy.units as u
 from ..utils.nddata import NDDataArray, BinnedDataAxis
 from ..utils.scripts import make_path
-from ..utils.fits import fits_table_to_table
+from ..utils.fits import fits_table_to_table, table_to_fits_table
 
+import numpy as np
 __all__ = [
     'Background3D',
 ]
@@ -114,3 +116,20 @@ def read(cls, filename, hdu='BACKGROUND'):
         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['DETX_LO'] = self.data.axis('detx').lo[np.newaxis]
+        table['DETX_HI'] = self.data.axis('detx').hi[np.newaxis]
+        table['DETY_LO'] = self.data.axis('dety').lo[np.newaxis]
+        table['DETY_HI'] = self.data.axis('dety').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)

gammapy/irf/effective_area.py

@@ -201,7 +201,7 @@ def to_table(self):
         """
         table = Table()
         table.meta = OrderedDict([
-            ('name', 'SPECRESP'),
+            ('EXTNAME', 'SPECRESP'),
             ('hduclass', 'OGIP'),
             ('hduclas1', 'RESPONSE'),
             ('hduclas2', 'SPECRESP'),
@@ -344,13 +344,14 @@ class EffectiveAreaTable2D(object):
     >>> import numpy as np
     >>> energy = np.logspace(0,1,11) * u.TeV
     >>> offset = np.linspace(0,1,4) * u.deg
-    >>> data = np.ones(shape=(10,4)) * u.cm * u.cm
-    >>> aeff = EffectiveAreaTable2D(energy=energy, offset=offset, data= data)
+    >>> data = np.ones(shape=(10,3)) * u.cm * u.cm
+    >>> aeff = EffectiveAreaTable2D(energy_lo=energy[:-1], energy_hi=energy[1:], offset_lo=offset[:-1], 
+    >>>                             offset_hi=offset[1:], data= data)
     >>> print(aeff)
     Data array summary info
     energy         : size =    11, min =  1.000 TeV, max = 10.000 TeV
     offset         : size =     4, min =  0.000 deg, max =  1.000 deg
-    Data           : size =    40, min =  1.000 cm2, max =  1.000 cm2
+    Data           : size =    30, min =  1.000 cm2, max =  1.000 cm2
     """
     default_interp_kwargs = dict(bounds_error=False, fill_value=None)
     """Default Interpolation kwargs for `~NDDataArray`. Extrapolate."""
@@ -377,14 +378,6 @@ def __str__(self):
         ss += '\n{}'.format(self.data)
         return ss
 
-    @property
-    def energy(self):
-        return self.data.axis('energy')
-
-    @property
-    def offset(self):
-        return self.data.axis('offset')
-
     @property
     def low_threshold(self):
         """Low energy threshold"""
@@ -432,7 +425,7 @@ def to_effective_area_table(self, offset, energy=None):
             Energy axis bin edges
         """
         if energy is None:
-            energy = self.energy.bins
+            energy = self.data.axis('energy').bins
 
         energy = EnergyBounds(energy)
         area = self.data.evaluate(offset=offset, energy=energy.log_centers)
@@ -471,15 +464,15 @@ def plot_energy_dependence(self, ax=None, offset=None, energy=None, **kwargs):
             offset = np.linspace(off_min, off_max, 4) * self.data.axis('offset').unit
 
         if energy is None:
-            energy = self.energy.nodes
+            energy = self.data.axis('energy').nodes
 
         for off in offset:
             area = self.data.evaluate(offset=off, energy=energy)
             label = 'offset = {:.1f}'.format(off)
             ax.plot(energy, area.value, label=label, **kwargs)
 
         ax.set_xscale('log')
-        ax.set_xlabel('Energy [{}]'.format(self.energy.unit))
+        ax.set_xlabel('Energy [{}]'.format(self.data.axis('energy').unit))
         ax.set_ylabel('Effective Area [{}]'.format(self.data.data.unit))
         ax.set_xlim(min(energy.value), max(energy.value))
         ax.legend(loc='upper left')
@@ -508,8 +501,8 @@ def plot_offset_dependence(self, ax=None, offset=None, energy=None, **kwargs):
         ax = plt.gca() if ax is None else ax
 
         if energy is None:
-            e_min, e_max = np.log10(self.energy.nodes[[0, -1]].value)
-            energy = np.logspace(e_min, e_max, 4) * self.energy.unit
+            e_min, e_max = np.log10(self.data.axis('energy').nodes[[0, -1]].value)
+            energy = np.logspace(e_min, e_max, 4) * self.data.axis('energy').unit
 
         if offset is None:
             off_lo, off_hi = self.data.axis('offset').nodes[[0, -1]].to('deg').value
@@ -570,3 +563,18 @@ def peek(self, figsize=(15, 5)):
         self.plot_energy_dependence(ax=axes[0])
         self.plot_offset_dependence(ax=axes[1])
         plt.tight_layout()
+
+    def to_table(self):
+        """Convert to `~astropy.table.Table`."""
+        meta = self.meta.copy()
+        table = Table(meta=meta)
+        table['ENERG_LO'] = self.data.axis('energy').lo[np.newaxis]
+        table['ENERG_HI'] = self.data.axis('energy').hi[np.newaxis]
+        table['THETA_LO'] = self.data.axis('offset').lo[np.newaxis]
+        table['THETA_HI'] = self.data.axis('offset').hi[np.newaxis]
+        table['EFFAREA'] = self.data.data.T[np.newaxis]
+        return table
+
+    def to_fits(self, name='EFFECTIVE AREA'):
+        """Convert to `~astropy.io.fits.BinTable`."""
+        return table_to_fits_table(self.to_table(), name)

gammapy/irf/energy_dispersion.py

@@ -10,6 +10,7 @@
 from ..utils.scripts import make_path
 from ..utils.nddata import NDDataArray, BinnedDataAxis
 from ..utils.fits import energy_axis_to_ebounds, fits_table_to_table
+from ..utils.fits import fits_table_to_table, table_to_fits_table
 
 __all__ = [
     'EnergyDispersion',
@@ -592,7 +593,7 @@ class EnergyDispersion2D(object):
     """Default Interpolation kwargs for `~gammapy.utils.nddata.NDDataArray`. Extrapolate."""
 
     def __init__(self, e_true_lo, e_true_hi, migra_lo, migra_hi, offset_lo,
-                 offset_hi, data, interp_kwargs=None):
+                 offset_hi, data, interp_kwargs=None, meta=None,):
         if interp_kwargs is None:
             interp_kwargs = self.default_interp_kwargs
         axes = [
@@ -605,27 +606,13 @@ def __init__(self, e_true_lo, e_true_hi, migra_lo, migra_hi, offset_lo,
         ]
         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
 
-    @property
-    def e_true(self):
-        """True energy axis (`~gammapy.utils.nddata.BinnedDataAxis`)."""
-        return self.data.axis('e_true')
-
-    @property
-    def migra(self):
-        """Energy migration axis (`~gammapy.utils.nddata.BinnedDataAxis`)."""
-        return self.data.axis('migra')
-
-    @property
-    def offset(self):
-        """Field of view offset axis (`~gammapy.utils.nddata.BinnedDataAxis`)."""
-        return self.data.axis('offset')
-
     @classmethod
     def from_gauss(cls, e_true, migra, bias, sigma, offset, pdf_threshold=1e-6):
         """Create Gaussian energy dispersion matrix (`EnergyDispersion2D`).
@@ -680,14 +667,20 @@ def from_gauss(cls, e_true, migra, bias, sigma, offset, pdf_threshold=1e-6):
     @classmethod
     def from_table(cls, table):
         """Create from `~astropy.table.Table`."""
-        e_lo = table['ETRUE_LO'].quantity.squeeze()
-        e_hi = table['ETRUE_HI'].quantity.squeeze()
-        o_lo = table['THETA_LO'].quantity.squeeze()
-        o_hi = table['THETA_HI'].quantity.squeeze()
-        m_lo = table['MIGRA_LO'].quantity.squeeze()
-        m_hi = table['MIGRA_HI'].quantity.squeeze()
-
-        matrix = table['MATRIX'].squeeze().transpose()
+        if 'ENERG_LO' in table.colnames:
+            e_lo = table['ENERG_LO'].quantity[0]
+            e_hi = table['ENERG_HI'].quantity[0]
+        elif 'ETRUE_LO' in table.colnames:
+            e_lo = table['ETRUE_LO'].quantity[0]
+            e_hi = table['ETRUE_HI'].quantity[0]
+        else:
+            raise ValueError('Invalid column names. Need "ENERG_LO/ENERG_HI" or "ETRUE_LO/ETRUE_HI"')
+        o_lo = table['THETA_LO'].quantity[0]
+        o_hi = table['THETA_HI'].quantity[0]
+        m_lo = table['MIGRA_LO'].quantity[0]
+        m_hi = table['MIGRA_HI'].quantity[0]
+
+        matrix = table['MATRIX'].quantity[0].transpose() ## TODO Why does this need to be transposed?
         return cls(e_true_lo=e_lo, e_true_hi=e_hi,
                    offset_lo=o_lo, offset_hi=o_hi,
                    migra_lo=m_lo, migra_hi=m_hi, data=matrix)
@@ -733,8 +726,8 @@ def to_energy_dispersion(self, offset, e_true=None, e_reco=None):
             Energy dispersion matrix
         """
         offset = Angle(offset)
-        e_true = self.e_true.bins if e_true is None else e_true
-        e_reco = self.e_true.bins if e_reco is None else e_reco
+        e_true = self.data.axis('e_true').bins if e_true is None else e_true
+        e_reco = self.data.axis('e_true').bins if e_reco is None else e_reco
         e_true = EnergyBounds(e_true)
         e_reco = EnergyBounds(e_reco)
 
@@ -779,8 +772,8 @@ def get_response(self, offset, e_true, e_reco=None, migra_step=5e-3):
         # Default: e_reco nodes = migra nodes * e_true nodes
         if e_reco is None:
             e_reco = EnergyBounds.from_lower_and_upper_bounds(
-                self.migra.lo * e_true, self.migra.hi * e_true)
-            migra = self.migra.nodes
+                self.data.axis('migra').lo * e_true, self.data.axis('migra').hi * e_true)
+            migra = self.data.axis('migra').nodes
         # Translate given e_reco binning to migra at bin center
         else:
             e_reco = EnergyBounds(e_reco)
@@ -791,8 +784,8 @@ def get_response(self, offset, e_true, e_reco=None, migra_step=5e-3):
         migra_e_reco = e_reco / e_true
 
         # Define a vector of migration with mig_step step
-        mrec_min = self.migra.lo[0]
-        mrec_max = self.migra.hi[-1]
+        mrec_min = self.data.axis('migra').lo[0]
+        mrec_max = self.data.axis('migra').hi[-1]
         mig_array = np.arange(mrec_min, mrec_max, migra_step)
 
         # Compute energy dispersion probability dP/dm for each element of migration array
@@ -844,7 +837,7 @@ def plot_migration(self, ax=None, offset=None, e_true=None,
             e_true = Energy([0.1, 1, 10], 'TeV')
         else:
             e_true = np.atleast_1d(Energy(e_true))
-        migra = self.migra.nodes if migra is None else migra
+        migra = self.data.axis('migra').nodes if migra is None else migra
 
         for ener in e_true:
             for off in offset:
@@ -888,8 +881,8 @@ def plot_bias(self, ax=None, offset=None, add_cbar=False, **kwargs):
         if offset is None:
             offset = Angle([1], 'deg')
 
-        e_true = self.e_true.bins
-        migra = self.migra.bins
+        e_true = self.data.axis('e_true').bins
+        migra = self.data.axis('migra').bins
 
         x = e_true.value
         y = migra.value
@@ -924,3 +917,20 @@ def peek(self, figsize=(15, 5)):
         edisp.plot_matrix(ax=axes[2])
 
         plt.tight_layout()
+
+    def to_table(self):
+        """Convert to `~astropy.table.Table`."""
+        meta = self.meta.copy()
+        table = Table(meta=meta)
+        table['ENERG_LO'] = self.data.axis('e_true').lo[np.newaxis]
+        table['ENERG_HI'] = self.data.axis('e_true').hi[np.newaxis]
+        table['MIGRA_LO'] = self.data.axis('migra').hi[np.newaxis]
+        table['MIGRA_HI'] = self.data.axis('migra').hi[np.newaxis]
+        table['THETA_LO'] = self.data.axis('offset').lo[np.newaxis]
+        table['THETA_HI'] = self.data.axis('offset').hi[np.newaxis]
+        table['MATRIX'] = self.data.data.T[np.newaxis]
+        return table
+
+    def to_fits(self, name='ENERGY DISPERSION'):
+        """Convert to `~astropy.io.fits.BinTable`."""
+        return table_to_fits_table(self.to_table(), name)