Adds plotting option for XAS

XAS.py

@@ -10,6 +10,8 @@
 """
 
 import numpy as np
+import matplotlib.gridspec as gridspec
+import matplotlib.pyplot as plt
 
 def absorption(T, Io):
     """ Compute the absorption A = -ln(T/Io)
@@ -97,7 +99,7 @@ def binning(x, data, func, bins=100, bin_length=None):
 
     return bins, res
 
-def xas(nrun, bins=None, Iokey='SCS_SA3', Itkey='MCP3apd', nrjkey='nrj', Iooffset=0):
+def xas(nrun, bins=None, Iokey='SCS_SA3', Itkey='MCP3apd', nrjkey='nrj', Iooffset=0, plot=False):
     """ Compute the XAS spectra from a xarray nrun.
 
         Inputs:
@@ -108,6 +110,7 @@ def xas(nrun, bins=None, Iokey='SCS_SA3', Itkey='MCP3apd', nrjkey='nrj', Iooffse
             Itkey: string for the It fields, typically 'MCP3apd'
             NRJkey: string for the nrj fields, typically 'nrj'
             Iooffset: offset to apply on Io
+            plot: boolean, displays a XAS spectrum if True
 
         Outputs:
             a dictionnary containing:
@@ -154,6 +157,29 @@ def xas(nrun, bins=None, Iokey='SCS_SA3', Itkey='MCP3apd', nrjkey='nrj', Iooffse
 
     bins_c = 0.5*(bins[1:] + bins[:-1])
     
+    if plot:
+        f = plt.figure(figsize=(6.5,6))
+        gs = gridspec.GridSpec(2,1,height_ratios=[4,1])
+        ax1 = plt.subplot(gs[0])
+        ax1.plot(bins_c, muA, color='C1', label=r'$\sigma$')
+        ax1.set_ylabel('XAS')
+        ax1.set_xlabel('Energy (eV)')
+        ax1.legend()
+        ax1_twin = ax1.twinx()
+        ax1_twin.bar(bins_c, nosample['muIo'], width=0.80*(bins_c[1]-bins_c[0]),
+                color='C1', alpha=0.2)
+        ax1_twin.set_ylabel('Io')
+        proposalNB=int(nrun.attrs['runFolder'].split('/')[-4][1:])
+        runNB=int(nrun.attrs['runFolder'].split('/')[-2][1:])
+        ax1.set_title('run {:d} p{:}'.format(runNB, proposalNB))
+        
+        ax2 = plt.subplot(gs[1])
+        ax2.bar(bins_c, nosample['counts'], width=0.80*(bins_c[1]-bins_c[0]),
+                color='C0', alpha=0.2)
+        ax2.set_xlabel('Energy (eV)')
+        ax2.set_ylabel('counts')
+        plt.tight_layout()
+    
     return {'nrj':bins_c, 'muA':muA, 'sterrA':sterrA, 'sigmaA':nosample['sigmaA'],
         'muIo':nosample['muIo'], 'counts':nosample['counts']}