Update XAS function and remove printouts

src/toolbox_scs/detectors/viking.py

@@ -290,7 +290,6 @@ class Viking:
         """
         if key not in data:
             return
-        print('hello')
         x = data.newt_x
         spectra = data[key]
         mask = xr.ones_like(x, dtype=bool)
@@ -302,7 +301,6 @@ class Viking:
             for xrange in ranges:
                 mask = mask & ((x < xrange[0]) | (x > xrange[1]))
         x_bl = x.where(mask, drop=True).astype(int)
-        print(x_bl)
         bl = spectra.sel(newt_x=x_bl)
         fit = np.polyfit(x_bl, bl.T, self.BL_POLY_DEG)
         if len(spectra.shape) == 1:
@@ -312,8 +310,9 @@ class Viking:
             final_bl[t] = np.polyval(fit[:, t], x)
         data[key+'_nobl'] = spectra - final_bl
         return data
-
-    def xas(self, data, data_ref, thickness=1, plot=False,
+    
+    
+    def xas(self, sam, ref, thickness=1, dim='newt_x', plot=False,
             plot_errors=True, xas_ylim=(-1, 3)):
         """
         Given two independent datasets (one with sample and one reference),
@@ -323,13 +322,15 @@ class Viking:
 
         Parameters
         ----------
-        data: xarray Dataset
-            the dataset containing the spectra with sample
-        data_ref: xarray Dataset
-            the dataset containing the spectra without sample
+        sam: xarray DataArray
+            the data array containing the spectra with sample
+        ref: xarray DataArray
+            the data array containing the spectra without sample
         thickness: float
             the thickness used for the calculation of the absorption
             coefficient
+        dim: string
+            the name of the dimension along the dispersion axis
         plot: bool
             If True, plot the resulting average spectra.
         plot_errors: bool
@@ -344,15 +345,14 @@ class Viking:
             I0, It, absorptionCoef and their associated errors.
         """
 
-        key = 'spectrum_nobl' if 'spectrum_nobl' in data else 'spectrum'
-        if data['newt_x'].equals(data_ref['newt_x']) is False:
+        if sam[dim].equals(ref[dim]) is False:
             return
-        spectrum = data[key].mean(dim='trainId')
-        std = data[key].std(dim='trainId')
-        std_err = std / np.sqrt(data.sizes['trainId'])
-        spectrum_ref = data_ref[key].mean(dim='trainId')
-        std_ref = data_ref[key].std(dim='trainId')
-        std_err_ref = std_ref / np.sqrt(data_ref.sizes['trainId'])
+        spectrum = sam.mean(dim='trainId')
+        std = sam.std(dim='trainId')
+        std_err = std / np.sqrt(sam.sizes['trainId'])
+        spectrum_ref = ref.mean(dim='trainId')
+        std_ref = ref.std(dim='trainId')
+        std_err_ref = std_ref / np.sqrt(ref.sizes['trainId'])
 
         ds = xr.Dataset()
         ds['It'] = spectrum
@@ -372,13 +372,14 @@ class Viking:
                                                      np.abs(absorption))
         ds['absorptionCoef_stderr'] = absorption_stderr / (thickness *
                                                            np.abs(absorption))
-        ds.attrs['n_It'] = data[key].sizes['trainId']
-        ds.attrs['n_I0'] = data_ref[key].sizes['trainId']
+        ds.attrs['n_It'] = sam.sizes['trainId']
+        ds.attrs['n_I0'] = ref.sizes['trainId']
 
         if plot:
             plot_viking_xas(ds, plot_errors, xas_ylim)
         return ds
 
+
     def calibrate(self, runList, plot=True):
         """
         This routine determines the calibration coefficients to translate the