added plot function, evaluate function, chages to load function

gammapy/datasets/load.py

@@ -34,6 +34,7 @@
            'load_electron_spectrum',
            'load_arf_fits_table',
            'load_aeff2D_fits_table',
+           'load_edisp2D_fits_table',
            'load_psf_fits_table',
            ]
 
@@ -140,11 +141,12 @@ def load_edisp2D_fits_table():
 
     Returns
     -------
-    hdu_list : `~astropy.io.fits.HDUList`
-        aeff2D file contents.
+    hdu: `~astropy.io.fits.BinTableHDU`
+        edisp2D file contents.
     """
     filename = get_path('irfs/edisp2D.fits')
-    return fits.open(filename)
+    hdulist = fits.open(filename)
+    return hdulist[1]
 
 def load_poisson_stats_image(extra_info=False, return_filenames=False):
     """Load Poisson statistics counts image of a Gaussian source on flat background.

gammapy/irf/energy_dispersion.py

@@ -614,10 +614,71 @@ def to_energy_dispersion(self, offset, energy):
         """
         pass
 
+    def evaluate(self, offset=None, energy=None, migra=None):
+        """One line docu string
+        """
+
+        if offset is None:
+            offset = self.offset
+        if energy is None:
+            energy = self.energy
+        if migra is None:
+            migra = self.migra
+
+        if not isinstance(energy, Quantity):
+            raise ValueError("Energy must be a Quantity object.")
+        if not isinstance(offset, Angle):
+            raise ValueError("Offset must be an Angle object.")
+
+        offset = offset.to('degree')
+        energy = energy.to('TeV')
+
+        val = self._eval(offset=offset, energy=energy, migra=migra)
+
+        return val
+
     def _eval(self, offset=None, energy=None, migra=None):
+
+        x = np.asarray(offset.value)
+        y = np.asarray(migra)
+        z = np.asarray(np.log10(energy.value))
+        ax = [x,y,z]
+
+        in_shape = (ax[0].size, ax[1].size, ax[2].size)
+
+        for i,s in enumerate(ax):
+            if ax[i].shape == ():
+                ax[i] = ax[i].reshape(1)
+
+        pts = [[xx, yy, zz] for xx in ax[0] for yy in ax[1] for zz in ax[2]]
+
+        val_array = self._linear(pts)
+        return val_array.reshape(in_shape).squeeze()
+
+    def plot_migra_hist(self, ax=None, offset=None, energy=None, migra=None, **kwargs):
+        """Plot migration histogram for given offset and true energy.
+        """
+        import matplotlib.pyplot as plt
+
+        ax = plt.gca() if ax is None else ax
+
+        if offset is None:
+            val = self.offset
+            offset = Angle(0.5, 'deg')
+
+        if energy is None:
+            energy = Quantity(10, 'TeV')
+
+        if migra is None:
+            migra = self.migra
+
+        disp = self.evaluate(offset=offset, energy=energy, migra=migra)
+        plt.plot(migra, disp, **kwargs)
+
+        plt.xlabel('RecoEnergy')
+        plt.ylabel('Probability')
 
-        val = self._linear(offset.value, migra, np.log10(energy.value))
-        return Quantity(val, self.eff_area.unit)
+        return ax
 
     def _prepare_linear_interpolator(self):
         """Linear interpolation in N dimensions

gammapy/irf/tests/test_energy_dispersion.py

@@ -1,33 +1,41 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 from __future__ import (absolute_import, division, print_function,
                         unicode_literals)
-from numpy.testing import assert_allclose
+from numpy.testing import assert_allclose, assert_equal
 from astropy.tests.helper import pytest
-from ...irf import EnergyDispersion
-
+from ...irf import EnergyDispersion, EnergyDispersion2D
+from ...datasets import load_edisp2D_fits_table
 
 @pytest.mark.xfail
 def test_EnergyDispersion():
     edisp = EnergyDispersion()
     pdf = edisp(3, 4)
     assert_allclose(pdf, 42)
 
-
-
 def test_EnergyDispersion2D():
 
     # Read test effective area file
-    effarea = EffectiveAreaTable2D.from_fits(
-        load_aeff2D_fits_table())
-
-    effarea.interpolation_method = method
+    edisp = EnergyDispersion2D.from_fits(
+        load_edisp2D_fits_table())
 
     # Check that nodes are evaluated correctly
-    e_node = 42
+    e_node = 12
     off_node = 3
-    offset = effarea.offset[off_node]
-    energy = effarea.energy[e_node]
-    actual = effarea.evaluate(offset, energy)
-    desired = effarea.eff_area[off_node, e_node]
-    assert_allclose(actual, desired)
+    m_node = 5
+    offset = edisp.offset[off_node]
+    energy = edisp.energy.log_centers[e_node]
+    migra = edisp.migra[m_node]
+    actual = edisp.evaluate(offset, energy, migra)
+    desired = edisp.dispersion[off_node, m_node, e_node]
+    assert_allclose(actual, desired, rtol=1e-06)
+
+
+    # Check that values between node make sense
+    #THINK ABOUT WHAT MAKES SENSE
+    energy2 = edisp.energy.log_centers[e_node + 1]
+    upper = edisp.evaluate(offset, energy, migra)
+    lower = edisp.evaluate(offset, energy2, migra)
+    e_val = (energy + energy2) / 2
+    actual = edisp.evaluate(offset, e_val, migra)
+    #assert_equal(lower > actual and actual > upper, True)