show_spectrum.py now adapts to energy scale of data and shows filename in title

Jürgen Knödlseder
1 parent 1cbb5663
Showing 1 changed file with 39 additions and 17 deletions
examples/show_spectrum.py
@@ -2,7 +2,7 @@
 # ==========================================================================
 # Display spectrum generated by csspec
 #
-# Copyright (C) 2015-2020 Juergen Knoedlseder
+# Copyright (C) 2015-2021 Juergen Knoedlseder
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by
@@ -18,6 +18,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
 # ==========================================================================
+import os
 import sys
 try:
     import matplotlib.pyplot as plt
@@ -70,7 +71,7 @@ def get_spectrum_fits(fits):
     ----------
     fits : `~gammalib.GFits`
         Spectral GFits object
-    
+
     Returns
     -------
     spec : dict
@@ -168,7 +169,7 @@ def get_spectrum_fits(fits):
 # ========================= #
 # Plot delta log-likelihood #
 # ========================= #
-def plot_dloglike(spec):
+def plot_dloglike(ax, spec):
     """
     Plot delta log-likelihood
@@ -187,7 +188,7 @@ def plot_dloglike(spec):
     # Compute the logarithmic center of each bin
     fluxpnts = 10 ** ((y_bounds[1:]+y_bounds[:-1]) / 2.0)
-    
+
     # Scale the boundaries on the grid
     y_bounds = 10 ** y_bounds
@@ -201,12 +202,12 @@ def plot_dloglike(spec):
     dll_hist = [[row[i] for row in dll_hist] for i in range(steps)]
     # Plot the likelihood profiles
-    plt.pcolormesh(np.array(spec['engs_scan']), y_bounds,
-                   np.array(dll_hist), vmin=-10.0, vmax=0.0,
-                   cmap='Reds', alpha=0.9)
-    cbar = plt.colorbar()
+    c = ax.pcolormesh(np.array(spec['engs_scan']), y_bounds,
+                      np.array(dll_hist), vmin=-10.0, vmax=0.0,
+                      cmap='Reds', alpha=0.9)
+    cbar = plt.colorbar(c)
     cbar.set_label(r'$\Delta$ log-likelihood')
-    plt.grid()
+    ax.grid()
     # Return
     return
@@ -229,27 +230,48 @@ def plot_spectrum(filename, plotfile):
     # Get spectrum parameters
     spec = get_spectrum_file(filename)
+    # Set energy scale of spectrum
+    if max(spec['energies']) < 1.0e-3:
+        xlabel = 'Energy (MeV)'
+        conv   = 1.0e6
+    elif max(spec['energies']) < 1.0:
+        xlabel = 'Energy (GeV)'
+        conv   = 1.0e3
+    else:
+        xlabel = 'Energy (TeV)'
+        conv   = 1.0
+    if conv != 1.0:
+        spec['energies']    = [e * 1.0e6 for e in spec['energies']]
+        spec['ed_engs']     = [e * 1.0e6 for e in spec['ed_engs']]
+        spec['eu_engs']     = [e * 1.0e6 for e in spec['eu_engs']]
+        spec['ul_energies'] = [e * 1.0e6 for e in spec['ul_energies']]
+        spec['ul_ed_engs']  = [e * 1.0e6 for e in spec['ul_ed_engs']]
+        spec['ul_eu_engs']  = [e * 1.0e6 for e in spec['ul_eu_engs']]
+
     # Create the plot
-    plt.figure()
-    plt.loglog()
-    plt.grid()
-    plt.xlabel('Energy (TeV)')
-    plt.ylabel(r'E$^2$ $\times$ dN/dE (erg cm$^{-2}$ s$^{-1}$)')
+    fig = plt.figure()
+    fig.subplots_adjust(left=0.17, right=0.98, top=0.92, bottom=0.10)
+    ax = fig.add_subplot(111)
+    ax.loglog()
+    ax.grid()
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel(r'E$^2$ $\times$ dN/dE (erg cm$^{-2}$ s$^{-1}$)')
+    ax.set_title(os.path.basename(filename))
     # Plot the spectrum
-    plt.errorbar(spec['energies'], spec['flux'], 
+    ax.errorbar(spec['energies'], spec['flux'],
                  yerr=spec['e_flux'], xerr=[spec['ed_engs'], spec['eu_engs']],
                  fmt='ro')
     # Plot upper limits
     if len(spec['ul_energies'])	> 0:
-        plt.errorbar(spec['ul_energies'], spec['ul_flux'], yerr=spec['yerr'],
+        ax.errorbar(spec['ul_energies'], spec['ul_flux'], yerr=spec['yerr'],
                      xerr=[spec['ul_ed_engs'], spec['ul_eu_engs']],
                      uplims=True, fmt='ro')
     # Plot log-likelihood profiles
     if len(spec['dll_scan']) > 0:
-        plot_dloglike(spec)
+        plot_dloglike(ax, spec)
     # Show figure
     if len(plotfile) > 0: