Add make_spectrum.py example script to unit tests

Jürgen Knödlseder
1 parent 4300668e
Showing 2 changed files with 90 additions and 55 deletions
examples/make_spectrum.py
test/test_examples.py
 #! /usr/bin/env python
 # ==========================================================================
-# This script shows how to make a spectrum using obsutils.
+# Makes spectrum and optionally plot it
 #
 # Copyright (C) 2014-2016 Juergen Knoedlseder
 #
@@ -18,7 +18,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
 # ==========================================================================
-#import gammalib
+import sys
 import ctools
 import cscripts
 try:
@@ -33,12 +33,17 @@ except:
 # =============== #
 def make_spectrum():
     """
-    Make a spectrum.
+    Simulate events and generate a source spectrum
+
+    Returns
+    -------
+    spectrum : `~gammalib.GFits`
+        Spectrum FITS file
     """
     # Set script parameters
-    model_name  = "${CTOOLS}/share/models/crab.xml"
-    caldb       = "prod2"
-    irf         = "South_50h"
+    model_name  = 'data/crab.xml'
+    caldb       = 'prod2'
+    irf         = 'South_0.5h'
     ra          =   83.63
     dec         =   22.01
     rad_sim     =    3.0
@@ -50,36 +55,32 @@ def make_spectrum():
     # Simulate events
     sim = ctools.ctobssim()
-    sim["inmodel"] = model_name
-    sim["caldb"]   = caldb
-    sim["irf"]     = irf
-    sim["ra"]      = ra
-    sim["dec"]     = dec
-    sim["rad"]     = rad_sim
-    sim["tmin"]    = tstart
-    sim["tmax"]    = tstop
-    sim["emin"]    = emin
-    sim["emax"]    = emax
+    sim['inmodel'] = model_name
+    sim['caldb']   = caldb
+    sim['irf']     = irf
+    sim['ra']      = ra
+    sim['dec']     = dec
+    sim['rad']     = rad_sim
+    sim['tmin']    = tstart
+    sim['tmax']    = tstop
+    sim['emin']    = emin
+    sim['emax']    = emax
     sim.run()
-    # Generate an energy binning
-    #e_min   = gammalib.GEnergy(emin, "TeV")
-    #e_max   = gammalib.GEnergy(emax, "TeV")
-    #ebounds = gammalib.GEbounds(10, e_min, e_max)
-
-    # Setup csspec run
+    # Create spectrum
     spec = cscripts.csspec(sim.obs())
-    spec["srcname"]  = "Crab"
-    spec["outfile"]  = "spectrum.fits"
-    spec["expcube"]  = "NONE"
-    spec["psfcube"]  = "NONE"
-    spec["bkgcube"]  = "NONE"
-    spec["edisp"]    = False
-    spec["emin"]     = emin
-    spec["emax"]     = emax
-    spec["enumbins"] = enumbins
-    spec["binned"]   = False
+    spec['srcname']  = 'Crab'
+    spec['outfile']  = 'example_spectrum.fits'
+    spec['expcube']  = 'NONE'
+    spec['psfcube']  = 'NONE'
+    spec['bkgcube']  = 'NONE'
+    spec['edisp']    = False
+    spec['emin']     = emin
+    spec['emax']     = emax
+    spec['enumbins'] = enumbins
+    spec['ebinalg']  = 'LOG'
     spec.run()
+    spec.save()
     # Get copy of spectrum
     spectrum = spec.spectrum().copy()
@@ -91,19 +92,26 @@ def make_spectrum():
 # ============= #
 # Plot spectrum #
 # ============= #
-def plot_spectrum(spectrum):
+def plot_spectrum(spectrum, file=False):
     """
     Plot spectrum.
+
+    Parameters
+    ----------
+    spectrum : `~gammalib.GFits`
+        Spectrum FITS file
+    type : str, optional
+        Plot type (blank means display on screen)
     """
-    # Read spectrum file    
+    # Extract columns from spectrum file
     table    = spectrum.table(1)
-    c_energy = table["Energy"]
-    c_ed     = table["ed_Energy"]
-    c_eu     = table["eu_Energy"]
-    c_flux   = table["Flux"]
-    c_eflux  = table["e_Flux"]
-    c_ts     = table["TS"]
-    c_upper  = table["UpperLimit"]
+    c_energy = table['Energy']
+    c_ed     = table['ed_Energy']
+    c_eu     = table['eu_Energy']
+    c_flux   = table['Flux']
+    c_eflux  = table['e_Flux']
+    c_ts     = table['TS']
+    c_upper  = table['UpperLimit']
     # Initialise arrays to be filled
     energies    = []
@@ -115,7 +123,7 @@ def plot_spectrum(spectrum):
     ul_ed_engs  = []
     ul_eu_engs  = []
     ul_flux     = []
-
+ 
     # Loop over rows of the file
     nrows = table.nrows()
     for row in range(nrows):
@@ -146,31 +154,43 @@ def plot_spectrum(spectrum):
     # Create figure
     plt.figure()
-    plt.title("Crab spectrum")
+    plt.title('Crab spectrum')
     # Plot the spectrum 
     plt.loglog()
     plt.grid()
-    plt.errorbar(energies, flux, yerr=e_flux, xerr=[ed_engs, eu_engs], fmt='ro')
-    plt.errorbar(ul_energies, ul_flux, xerr=[ul_ed_engs, ul_eu_engs], yerr=1.0e-11, uplims=True, fmt='ro')
-    plt.xlabel("Energy (TeV)")
-    plt.ylabel(r"E$^{2}$dN/dE (erg cm$^{-2}$ s$^{-1}$)")    
-
-    # Show spectrum
-    plt.show()
+    plt.errorbar(energies, flux, yerr=e_flux, xerr=[ed_engs, eu_engs],
+                 fmt='ro')
+    plt.errorbar(ul_energies, ul_flux, xerr=[ul_ed_engs, ul_eu_engs],
+                 yerr=1.0e-11, uplims=True, fmt='ro')
+    plt.xlabel('Energy (TeV)')
+    plt.ylabel(r'E$^{2}$dN/dE (erg cm$^{-2}$ s$^{-1}$)')
+
+    # Show spectrum or save it into file
+    if file:
+        plt.savefig('example_spectrum.eps')
+    else:
+        plt.show()
     # Return
     return
-#==========================#
+# ======================== #
 # Main routine entry point #
-#==========================#
+# ======================== #
 if __name__ == '__main__':
+    # Initialise option
+    file = False
+
+    # Get any command line arguments
+    if (len(sys.argv) > 1):
+        file = True
+
     # Generate spectrum
     spectrum = make_spectrum()
-    # Plot spectrum
+    # If matplotlib is installed then plot spectrum
     if has_matplotlib:
-        plot_spectrum(spectrum)
+        plot_spectrum(spectrum, file=file)
@@ -71,6 +71,7 @@ class Test(gammalib.GPythonTestSuite):
         # Execute Python script, make sure we catch any exception
         try:
             rc = os.system(cmd+' > example_'+logname+'.log 2>&1')
+            #rc = os.system(cmd)
         except:
             pass
@@ -91,6 +92,7 @@ class Test(gammalib.GPythonTestSuite):
         # Append tests
         self.append(self.test_generate_prod3_irfs, 'Test generate_prod3_irfs')
         self.append(self.test_make_pointings, 'Test make_pointings')
+        self.append(self.test_make_spectrum, 'Test make_spectrum')
         self.append(self.test_pipeline_binned_disk, 'Test pipeline_binned_disk')
         self.append(self.test_pipeline_binned_mem, 'Test pipeline_binned_mem')
         self.append(self.test_pipeline_stacked_disk, 'Test pipeline_stacked_disk')
@@ -110,7 +112,7 @@ class Test(gammalib.GPythonTestSuite):
         script = 'make_pointings'
         # Execute script
-        self._execute_python(script, args='-h', log='make_pointings_help')
+        self._execute_python(script, args='-h', log='make_pointings_h')
         self._execute_python(script, args='gps', log='make_pointings_gps')
         self._execute_python(script, args='gps3', log='make_pointings_gps3')
         self._execute_python(script, args='extgal', log='make_pointings_extgal')
@@ -122,6 +124,19 @@ class Test(gammalib.GPythonTestSuite):
         # Return
         return
+    # Test make_spectrum
+    def test_make_spectrum(self):
+        """
+        Test make_spectrum
+        """
+        # Execute script
+        self._execute_python('make_spectrum', args='eps')
+
+        #TODO: Do any testing
+
+        # Return
+        return
+
     # Test generate_prod3_irfs
     def test_generate_prod3_irfs(self):
         """