Add key argument to viking removePolyBaseline()

src/toolbox_scs/detectors/viking.py

@@ -270,7 +270,7 @@ class Viking:
         ret['photoelectrons_per_count'] = self.e_per_counts(run, ret['gain'])
         return ret
 
-    def removePolyBaseline(self, data):
+    def removePolyBaseline(self, data, key='spectrum'):
         """
         Removes a polynomial baseline to a spectrum, assuming a fixed
         position for the signal.
@@ -279,6 +279,7 @@ class Viking:
         ----------
         data: xarray Dataset
             The Viking data containing the variable "spectrum"
+            or given by argument key
             
         Output
         ------
@@ -287,10 +288,10 @@ class Viking:
             containing the baseline subtracted spectra.
 
         """
-        if 'spectrum' not in data:
+        if key not in data:
             return
         x = data.newt_x
-        spectra = data['spectrum']
+        spectra = data[key]
         mask = xr.ones_like(x, dtype=bool)
         if len(self.BL_SIGNAL_RANGE) > 0:
             if not hasattr(self.BL_SIGNAL_RANGE[0], '__len__'):
@@ -307,7 +308,7 @@ class Viking:
         final_bl = np.empty(spectra.shape)
         for t in range(spectra.shape[0]):
             final_bl[t] = np.polyval(fit[:, t], x)
-        data['spectrum_nobl'] = spectra - final_bl
+        data[key+'_nobl'] = spectra - final_bl
         return data
 
     def xas(self, data, data_ref, thickness=1, plot=False,