Improved some Viking functions

src/toolbox_scs/detectors/viking.py

@@ -138,7 +138,7 @@ class Viking:
                       'use_dark', 'bl_poly_deg', 'bl_signal_range', 'fields')
         return {param: getattr(self, param.upper()) for param in params}
 
-    def from_run(self, runNB, add_attrs=True):
+    def from_run(self, runNB, add_attrs=True, tid_offset=-1):
         """load a run
 
         Load the run `runNB`. A thin wrapper around `toolbox_scs.load`.
@@ -150,6 +150,8 @@ class Viking:
         add_attrs: bool
             if True, adds the camera parameters as attributes to the dataset
             (see get_camera_params())
+        tid_offset: int
+            train Id offset of Newton camera
 
         Output
         ------
@@ -165,11 +167,20 @@ class Viking:
             data2 = v.from_run(155)  # load run 155
             data = xarray.concat([data1, data2], 'trainId')  # combine both
         """
+        #separately load newton image and deal with trainId mismatch
         roi = {'newton': {'newton': {'roi': (self.Y_RANGE, self.X_RANGE),
                                      'dim': ['newt_y', 'newt_x']}}}
-        run, data = tb.load(self.PROPOSAL, runNB,
-                            fields=self.FIELDS, rois=roi)
-        data['newton'] = data['newton'].astype(float)
+        run, newton = tb.load(self.PROPOSAL, runNB, 'newton', rois=roi)
+        newton = newton.shift(trainId=tid_offset).astype(float)
+        
+        #load the rest
+        fields = [f for f in self.FIELDS if f != 'newton']
+        if len(fields) == 0:
+            data = newton
+        else:
+            run, data = tb.load(self.PROPOSAL, runNB,
+                                fields=fields)
+            data = data.merge(newton, join='inner')
         data = data.assign_coords(newt_x=np.polyval(self.ENERGY_CALIB,
                                                     data['newt_x']))
         if add_attrs:
@@ -259,18 +270,24 @@ class Viking:
                'endX': 'imageSpecifications.endX.value',
                'startY': 'imageSpecifications.startY.value',
                'endY': 'imageSpecifications.endY.value',
-               'temperature': 'CoolerActual.temperature.value',
+               'temperature': 'coolerActual.temperature.value',
                'high_capacity': 'HighCapacity.value',
                'exposure_s': 'exposureTime.value'
                }
         ret = {}
-        for k, v in dic.items():
-            ret[k] = run.get_run_value('SCS_EXP_NEWTON/CAM/CAMERA', v)
+        try:
+            for k, v in dic.items():
+                ret[k] = run.get_run_value('SCS_EXP_NEWTON/CAM/CAMERA', v)
+        except Exception as e:
+            print(e)
+            dic['temperature'] = 'CoolerActual.temperature.value'
+            for k, v in dic.items():
+                ret[k] = run.get_run_value('SCS_EXP_NEWTON/CAM/CAMERA', v)
         ret['gain'] = self.get_camera_gain(run)
         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.
@@ -287,10 +304,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,10 +324,10 @@ 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,
+    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),
@@ -320,13 +337,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
@@ -341,15 +360,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
@@ -369,8 +387,8 @@ 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)