Andreas Specovius / ctools

Commit 77d339936f95ee8260146ed4d98bb48f6776fad0

Authored by Josh Cardenzana
1 parent 329b821a

(#2925) Add 'get_spectrum_file()' and 'get_spectrum_fits()' methods 

Allows extracting SED parameters for plotting
Inline Side-by-side
Showing 1 changed file with 75 additions and 32 deletions
examples/show_spectrum.py
  View file @77d3399
... ... @@ -30,22 +30,38 @@ import gammalib
30 30 import cscripts
31 31  
32 32  
33   -# ============= #
34   -# Plot spectrum #
35   -# ============= #
36   -def plot_spectrum(filename, plotfile):
  33 +def get_spectrum_file(filename):
37 34 """
38   - Plot spectrum
  35 + Extract the spectrum info from a file for plotting
39 36  
40 37 Parameters
41 38 ----------
42 39 filename : str
43 40 Name of spectrum FITS file
44   - plotfile : str
45   - Plot file name
  41 +
  42 + Returns
  43 + -------
  44 + Python dictionary defining spectral plot parameters
46 45 """
47 46 # Read spectrum file
48   - fits = gammalib.GFits(filename)
  47 + fits = gammalib.GFits(filename)
  48 + return get_spectrum_fits(fits)
  49 +
  50 +
  51 +def get_spectrum_fits(fits):
  52 + """
  53 + Extract the spectrum info from a GFits object
  54 +
  55 + Parameters
  56 + ----------
  57 + fits : `~gammalib.GFits`
  58 + Spectral GFits object
  59 +
  60 + Returns
  61 + -------
  62 + Python dictionary defining spectral plot parameters
  63 + """
  64 + # Read spectrum objects
49 65 table = fits.table(1)
50 66 c_energy = table['Energy']
51 67 c_ed = table['ed_Energy']
... ... @@ -56,15 +72,18 @@ def plot_spectrum(filename, plotfile):
56 72 c_upper = table['UpperLimit']
57 73  
58 74 # Initialise arrays to be filled
59   - energies = []
60   - flux = []
61   - ed_engs = []
62   - eu_engs = []
63   - e_flux = []
64   - ul_energies = []
65   - ul_ed_engs = []
66   - ul_eu_engs = []
67   - ul_flux = []
  75 + spec = {
  76 + 'energies' : [],
  77 + 'flux' : [],
  78 + 'ed_engs' : [],
  79 + 'eu_engs' : [],
  80 + 'e_flux' : [],
  81 + 'ul_energies' : [],
  82 + 'ul_ed_engs' : [],
  83 + 'ul_eu_engs' : [],
  84 + 'ul_flux' : [],
  85 + 'yerr' : []
  86 + }
68 87  
69 88 # Loop over rows of the file
70 89 nrows = table.nrows()
... ... @@ -78,33 +97,57 @@ def plot_spectrum(filename, plotfile):
78 97 # If Test Statistic is larger than 9 and flux error is smaller than
79 98 # flux then append flux plots ...
80 99 if ts > 9.0 and e_flx < flx:
81   - energies.append(c_energy.real(row))
82   - flux.append(c_flux.real(row))
83   - ed_engs.append(c_ed.real(row))
84   - eu_engs.append(c_eu.real(row))
85   - e_flux.append(c_eflux.real(row))
  100 + spec['energies'].append(c_energy.real(row))
  101 + spec['flux'].append(c_flux.real(row))
  102 + spec['ed_engs'].append(c_ed.real(row))
  103 + spec['eu_engs'].append(c_eu.real(row))
  104 + spec['e_flux'].append(c_eflux.real(row))
86 105  
87 106 # ... otherwise append upper limit
88 107 else:
89   - ul_energies.append(c_energy.real(row))
90   - ul_flux.append(c_upper.real(row))
91   - ul_ed_engs.append(c_ed.real(row))
92   - ul_eu_engs.append(c_eu.real(row))
  108 + spec['ul_energies'].append(c_energy.real(row))
  109 + spec['ul_flux'].append(c_upper.real(row))
  110 + spec['ul_ed_engs'].append(c_ed.real(row))
  111 + spec['ul_eu_engs'].append(c_eu.real(row))
93 112  
94 113 # Set upper limit errors
95   - yerr = [0.6 * x for x in ul_flux]
  114 + spec['yerr'] = [0.6 * x for x in spec['ul_flux']]
  115 +
  116 + return spec
96 117  
97   - # Plot the spectrum
  118 +
  119 +# ============= #
  120 +# Plot spectrum #
  121 +# ============= #
  122 +def plot_spectrum(filename, plotfile):
  123 + """
  124 + Plot spectrum
  125 +
  126 + Parameters
  127 + ----------
  128 + filename : str
  129 + Name of spectrum FITS file
  130 + plotfile : str
  131 + Plot file name
  132 + """
  133 + spec = get_spectrum_file(filename)
  134 +
  135 + # Create the plot
98 136 plt.figure()
99 137 plt.loglog()
100 138 plt.grid()
101   - plt.errorbar(energies, flux, yerr=e_flux, xerr=[ed_engs, eu_engs],
102   - fmt='ro')
103   - plt.errorbar(ul_energies, ul_flux, xerr=[ul_ed_engs, ul_eu_engs],
104   - yerr=yerr, uplims=True, fmt='ro')
105 139 plt.xlabel('Energy (TeV)')
106 140 plt.ylabel(r'E$^2$ $\times$ dN/dE (erg cm$^{-2}$ s$^{-1}$)')
107 141  
  142 + # Plot the spectrum
  143 + plt.errorbar(spec['energies'], spec['flux'],
  144 + yerr=spec['e_flux'], xerr=[spec['ed_engs'], spec['eu_engs']],
  145 + fmt='ro')
  146 + # Plot upper limits
  147 + plt.errorbar(spec['ul_energies'], spec['ul_flux'],
  148 + yerr=spec['yerr'], xerr=[spec['ul_ed_engs'], spec['ul_eu_engs']],
  149 + uplims=True, fmt='ro')
  150 +
108 151 # Show figure
109 152 if len(plotfile) > 0:
110 153 plt.savefig(plotfile)
... ...