Improve PSF3D; fix to_energy_dependent_table_psf

gammapy/irf/psf_3d.py

@@ -69,20 +69,17 @@ def info(self):
 
         return ss
 
-    def energy_logcenter(self):
+    def _energy_logcenter(self):
         """Get logcenters of energy bins.
 
         Returns
         -------
         energies : `~astropy.units.Quantity`
             Logcenters of energy bins
         """
-        # TODO: should call helper function here, not re-implement this!
-        return 10 ** ((np.log10(self.energy_hi / Quantity(1, self.energy_hi.unit))
-                       + np.log10(self.energy_lo / Quantity(1, self.energy_lo.unit))) / 2) * Quantity(1,
-                                                                                                      self.energy_lo.unit)
+        return np.sqrt(self.energy_lo * self.energy_hi)
 
-    def rad_center(self):
+    def _rad_center(self):
         """Get centers of rad bins (`~astropy.coordinates.Angle` in deg).
         """
         return ((self.rad_hi + self.rad_lo) / 2).to('deg')
@@ -95,11 +92,10 @@ def read(cls, filename, hdu='PSF_2D_TABLE'):
         ----------
         filename : str
             File name
+        hdu : str
+            HDU name
         """
         filename = str(make_path(filename))
-        # TODO: implement it so that HDUCLASS is used
-        # http://gamma-astro-data-formats.readthedocs.io/en/latest/data_storage/hdu_index/index.html
-
         table = Table.read(filename, hdu=hdu)
         return cls.from_table(table)
 
@@ -189,25 +185,25 @@ def evaluate(self, energy=None, offset=None, rad=None,
         values : `~astropy.units.Quantity`
             Interpolated value
         """
+        from scipy.interpolate import RegularGridInterpolator
         if not interp_kwargs:
             interp_kwargs = dict(bounds_error=False, fill_value=None)
 
-        from scipy.interpolate import RegularGridInterpolator
         if energy is None:
-            energy = self.energy_logcenter()
+            energy = self._energy_logcenter()
         if offset is None:
             offset = self.offset
         if rad is None:
-            rad = self.rad_center()
+            rad = self._rad_center()
 
         energy = Energy(energy).to('TeV')
         offset = Angle(offset).to('deg')
         rad = Angle(rad).to('deg')
 
-        energy_bin = self.energy_logcenter()
+        energy_bin = self._energy_logcenter()
 
         offset_bin = self.offset.to('deg')
-        rad_bin = self.rad_center()
+        rad_bin = self._rad_center()
         points = (rad_bin, offset_bin, energy_bin)
         interpolator = RegularGridInterpolator(points, self.psf_value, **interp_kwargs)
         rr, off, ee = np.meshgrid(rad.value, offset.value, energy.value, indexing='ij')
@@ -216,14 +212,14 @@ def evaluate(self, energy=None, offset=None, rad=None,
         data_interp = interpolator(pix_coords)
         return Quantity(data_interp.reshape(shape), self.psf_value.unit)
 
-    def to_energy_dependent_table_psf(self, theta=None, exposure=None):
+    def to_energy_dependent_table_psf(self, theta='0 deg', exposure=None):
         """
         Convert PSF3D in EnergyDependentTablePSF.
 
         Parameters
         ----------
         theta : `~astropy.coordinates.Angle`
-            Offset in the field of view. Default theta = 0 deg
+            Offset in the field of view
         exposure : `~astropy.units.Quantity`
             Energy dependent exposure. Should be in units equivalent to 'cm^2 s'.
             Default exposure = 1.
@@ -233,17 +229,18 @@ def to_energy_dependent_table_psf(self, theta=None, exposure=None):
         table_psf : `~gammapy.irf.EnergyDependentTablePSF`
             Energy-dependent PSF
         """
-        theta = theta or Angle(0, 'deg')
-        energies = self.energy_logcenter()
-        rad = self.rad_center()
-        psf_value = self.evaluate(offset=theta).quantity[0].T
+        theta = Angle(theta)
+        energies = self._energy_logcenter()
+        rad = self._rad_center()
+        psf_value = self.evaluate(offset=theta)
+        psf_value = psf_value[:, 0, :].transpose()
 
         return EnergyDependentTablePSF(
             energy=energies, rad=rad,
             exposure=exposure, psf_value=psf_value,
         )
 
-    def to_table_psf(self, energy, theta=None, interp_kwargs=None, **kwargs):
+    def to_table_psf(self, energy, theta='0 deg', interp_kwargs=None, **kwargs):
         """Evaluate the `EnergyOffsetArray` at one given energy.
 
         Parameters
@@ -260,14 +257,15 @@ def to_table_psf(self, energy, theta=None, interp_kwargs=None, **kwargs):
         table : `~astropy.table.Table`
             Table with two columns: offset, value
         """
-        theta = theta or Angle(0, 'deg')
-
+        energy = Quantity(energy)
+        theta = Angle(theta)
         psf_value = self.evaluate(energy, theta, interp_kwargs=interp_kwargs).squeeze()
-        table_psf = TablePSF(self.rad_center(), psf_value, **kwargs)
+        rad = self._rad_center()
+        table_psf = TablePSF(rad, psf_value, **kwargs)
 
         return table_psf
 
-    def containment_radius(self, energy, theta=None, fraction=0.68, interp_kwargs=None):
+    def containment_radius(self, energy, theta='0 deg', fraction=0.68, interp_kwargs=None):
         """Containment radius.
 
         Parameters
@@ -284,10 +282,11 @@ def containment_radius(self, energy, theta=None, fraction=0.68, interp_kwargs=No
         radius : `~astropy.units.Quantity`
             Containment radius in deg
         """
-        if theta is None:
-            theta = Angle(0, 'deg')
+        energy = Quantity(energy)
         if energy.ndim == 0:
             energy = Quantity([energy.value], energy.unit)
+
+        theta = Angle(theta)
         if theta.ndim == 0:
             theta = Quantity([theta.value], theta.unit)
 
@@ -331,7 +330,7 @@ def plot_containment_vs_energy(self, fractions=[0.68, 0.95],
         ax.set_xlabel('Energy (TeV)')
         ax.set_ylabel('Containment radius (deg)')
 
-    def plot_psf_vs_rad(self, theta=Angle(0, 'deg'), energy=Quantity(1, 'TeV')):
+    def plot_psf_vs_rad(self, theta='0 deg', energy=Quantity(1, 'TeV')):
         """Plot PSF vs rad.
 
         Parameters
@@ -341,6 +340,7 @@ def plot_psf_vs_rad(self, theta=Angle(0, 'deg'), energy=Quantity(1, 'TeV')):
         theta : `~astropy.coordinates.Angle`
             Offset in the field of view. Default theta = 0 deg
         """
+        theta = Angle(theta)
         table = self.to_table_psf(energy=energy, theta=theta)
         return table.plot_psf_vs_rad()
 
@@ -360,7 +360,7 @@ def plot_containment(self, fraction=0.68, ax=None, show_safe_energy=False,
 
         ax = plt.gca() if ax is None else ax
 
-        energy = self.energy_logcenter()
+        energy = self._energy_logcenter()
         offset = self.offset
 
         # Set up and compute data

gammapy/irf/tests/test_psf_3d.py

@@ -1,52 +1,90 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 from __future__ import absolute_import, division, print_function, unicode_literals
-from astropy.tests.helper import assert_quantity_allclose
-from astropy.table import Table
+import pytest
+from numpy.testing import assert_allclose
+from astropy import units as u
 from ...utils.testing import requires_dependency, requires_data, mpl_savefig_check
-from ...utils.scripts import make_path
 from ...irf import PSF3D
 
+pytest.importorskip('scipy')
+
+
+@pytest.fixture(scope='session')
+def psf_3d():
+    filename = '$GAMMAPY_EXTRA/test_datasets/psf_table_023523.fits.gz'
+    return PSF3D.read(filename)
+
+
+@requires_data('gammapy-extra')
+def test_psf_3d_basics(psf_3d):
+    assert_allclose(psf_3d.rad_lo[-1].value, 0.6704476475715637)
+    assert psf_3d.rad_lo.shape == (900,)
+    assert psf_3d.rad_lo.unit == 'deg'
+
+    assert_allclose(psf_3d.energy_lo[0].value, 0.02)
+    assert psf_3d.energy_lo.shape == (18,)
+    assert psf_3d.energy_lo.unit == 'TeV'
+
+    assert psf_3d.psf_value.shape == (900, 6, 18)
+    assert psf_3d.psf_value.unit == 'sr-1'
+
+    assert_allclose(psf_3d.energy_thresh_lo.value, 0.1)
+    assert psf_3d.energy_lo.unit == 'TeV'
+
+    assert 'PSF3D' in psf_3d.info()
+
+
+@requires_data('gammapy-extra')
+def test_psf_3d_evaluate(psf_3d):
+    q = psf_3d.evaluate(energy='1 TeV', offset='0.3 deg', rad='0.1 deg')
+    assert_allclose(q.value, 21417.213824)
+    # TODO: is this the shape we want here?
+    assert q.shape == (1, 1, 1)
+    assert q.unit == 'sr-1'
+
 
 @requires_data('gammapy-extra')
-@requires_dependency('scipy')
-class TestPSF3D:
-    def setup(self):
-        filename = '$GAMMAPY_EXTRA/test_datasets/psf_table_023523.fits.gz'
-        self.psf = PSF3D.read(filename)
-        filename = str(make_path(filename))
-        self.table = Table.read(filename, 'PSF_2D_TABLE')
-
-    def test_read(self):
-        table = self.table
-        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]
-        rad_lo = table["RAD_LO"].quantity[0]
-        rad_hi = table["RAD_HI"].quantity[0]
-        psf_value = table["RPSF"].quantity[0]
-
-        assert_quantity_allclose(self.psf.energy_lo, energy_lo)
-        assert_quantity_allclose(self.psf.energy_hi, energy_hi)
-        assert_quantity_allclose(self.psf.offset, offset_lo)
-        assert_quantity_allclose(self.psf.offset, offset_hi)
-        assert_quantity_allclose(self.psf.rad_lo, rad_lo)
-        assert_quantity_allclose(self.psf.rad_hi, rad_hi)
-        assert_quantity_allclose(self.psf.psf_value, psf_value)
-
-    def test_write(self, tmpdir):
-        filename = str(tmpdir / 'psf.fits')
-        self.psf.write(filename)
-        psf2 = PSF3D.read(filename)
-        assert_quantity_allclose(self.psf.energy_lo, psf2.energy_lo)
-        assert_quantity_allclose(self.psf.energy_hi, psf2.energy_hi)
-        assert_quantity_allclose(self.psf.offset, psf2.offset)
-        assert_quantity_allclose(self.psf.rad_lo, psf2.rad_lo)
-        assert_quantity_allclose(self.psf.rad_hi, psf2.rad_hi)
-        assert_quantity_allclose(self.psf.psf_value, psf2.psf_value)
-
-    @requires_dependency('matplotlib')
-    def test_peek(self):
-        self.psf.peek()
-        mpl_savefig_check()
+def test_to_energy_dependent_table_psf(psf_3d):
+    psf = psf_3d.to_energy_dependent_table_psf()
+    assert psf.psf_value.shape == (18, 900)
+    radius = psf.table_psf_at_energy('1 TeV').containment_radius(0.68).deg
+    assert_allclose(radius, 0.171445, atol=1e-4)
+
+
+@requires_data('gammapy-extra')
+def test_psf_3d_containment_radius(psf_3d):
+    q = psf_3d.containment_radius(energy='1 TeV')
+    assert_allclose(q.value, 0.171445, rtol=1e-3)
+    assert q.isscalar
+    assert q.unit == 'deg'
+
+    q = psf_3d.containment_radius(energy=[1, 3] * u.TeV)
+    assert_allclose(q.value, [0.171445, 0.157455], rtol=1e-3)
+    assert q.shape == (2,)
+
+
+@requires_data('gammapy-extra')
+def test_psf_3d_write(psf_3d, tmpdir):
+    filename = str(tmpdir / 'temp.fits')
+    psf_3d.write(filename)
+    psf_3d = PSF3D.read(filename)
+
+    assert_allclose(psf_3d.energy_lo[0].value, 0.02)
+
+
+@requires_dependency('matplotlib')
+def test_psf_3d_plot_vs_rad(psf_3d):
+    psf_3d.plot_psf_vs_rad()
+    mpl_savefig_check()
+
+
+@requires_dependency('matplotlib')
+def test_psf_3d_plot_containment(psf_3d):
+    psf_3d.plot_containment(show_safe_energy=True)
+    mpl_savefig_check()
+
 
+@requires_dependency('matplotlib')
+def test_psf_3d_peek(psf_3d):
+    psf_3d.peek()
+    mpl_savefig_check()