developed viking functions

src/toolbox_scs/detectors/viking.py

+import numpy as np
+import xarray as xr
+import toolbox_scs as tb
+import matplotlib.pyplot as plt
+
+__all__ = ['Viking']
+
+# -----------------------------------------------------------------------------
+# Viking class
+class Viking:
+    # run
+    PROPOSAL = 2953
+
+    # image range
+    X_RANGE = slice(None, None)
+    Y_RANGE = slice(None, None)
+    
+    # dimension for integration
+    INTEGRATE_DIM = 'newt_y'
+    USE_DARK = False
+    
+    # polynomial degree for background subtraction
+    POLY_DEG = 5
+
+    FIELDS = ['newton']
+    
+    ENERGY_CALIB = [9.8233e-7, 0.0240, 514.4795]
+
+    def set_params(self, **params):
+        for key, value in params.items():
+            setattr(self, key.upper(), value)
+
+    def get_params(self, *params):
+        if not params:
+            params = ('proposal', 'x_range', 'y_range', 'integrate_dim',
+                      'fields',)
+        return {param: getattr(self, param.upper()) for param in params}
+
+    def from_run(self, runNB, proposal=None, add_attrs=True, calibrate=True):
+        if proposal is None:
+            proposal = self.PROPOSAL
+        roi = {'newton': {'newton': {'roi': (self.Y_RANGE, self.X_RANGE),
+                                      'dim': ['newt_y', 'newt_x']}}}
+        run, data = tb.load(proposal, runNB=runNB,
+                            fields=self.FIELDS, rois=roi)
+        data['newton'] = data['newton'].astype(float)
+        if calibrate:
+            data = data.assign_coords(newt_x=np.polyval(self.ENERGY_CALIB,
+                                                        data['newt_x']))
+        if add_attrs:
+            params = self.get_camera_params(run)
+            for k, v in params.items():
+                data.attrs[k] = v
+        return data
+
+    def load_dark(self, runNB=None, proposal=None):
+        if runNB is None:
+            self.USE_DARK = False
+            return
+        if proposal is None:
+            proposal = self.PROPOSAL
+        data = self.from_run(runNB, proposal, add_attrs=False)
+        self.dark_image = data['newton'].mean(dim='trainId')
+        self.dark_image.attrs['runNB'] = runNB
+        self.USE_DARK = True
+
+    def integrate(self, data):
+        imgs = data['newton']
+        if self.USE_DARK:
+            imgs -= self.dark_image
+        data['spectrum'] = imgs.sum(dim=self.INTEGRATE_DIM)
+        return data
+    
+    def get_camera_gain(self, run):
+        """
+        Get the preamp gain of the camera in the Viking spectrometer for
+        a specified run.
+
+        Parameters:
+        ------
+        run: extra_data DataCollection
+            information on the run
+
+        Output:
+        ------
+        gain: int
+        """
+        gain = run.get_run_value('SCS_EXP_NEWTON/CAM/CAMERA',
+                                 'preampGain.value')
+        gain_dict = {0: 1, 1: 2, 2: 4}
+        return gain_dict[gain]
+
+    def get_electrons_per_counts(self, run, gain=None):
+        """
+        Conversion factor from camera digital counts to photoelectrons
+        per count. The values can be found in the camera datasheet
+        but they have been slightly corrected for High Sensitivity
+        mode after analysis of runs 1204, 1207 and 1208, proposal 2937.
+
+        Parameters:
+        ------
+            run: extra_data DataCollection
+                information on the run
+            gain: int
+                the camera preamp gain
+
+        Outputs:
+        ------
+        ret: float
+            photoelectrons per count
+        """
+        if gain is None:
+            gain = self.get_camera_gain(run)
+        hc = run.get_run_value('SCS_EXP_NEWTON/CAM/CAMERA',
+                               'HighCapacity.value')
+        if hc == 0: # High Sensitivity
+            pe_dict = {1: 4., 2: 2.05, 4: 0.97}
+        elif hc == 1: # High Capacity
+            pe_dict = {1: 17.9, 2: 9., 4: 4.5}
+        return pe_dict[gain]
+
+    def get_camera_params(self, run):
+        dic = {'vbin:': 'imageSpecifications.verticalBinning.value',
+               'hbin': 'imageSpecifications.horizontalBinning.value',
+               'startX': 'imageSpecifications.startX.value',
+               'endX': 'imageSpecifications.endX.value',
+               'startY': 'imageSpecifications.startY.value',
+               'endY': 'imageSpecifications.endY.value',
+               'temperature': 'CoolerActual.temperature.value',
+               'high_sensitivity': '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)
+        ret['gain'] = self.get_camera_gain(run)
+        ret['photoelectrons_per_count'] = self.get_electrons_per_counts(run, ret['gain'])
+        return ret
+
+    def removePolyBaseline(self, data, signalRange=[515, 540]):
+        """
+        Removes a polynomial baseline to a signal, assuming a fixed
+        position for the signal.
+        Parameters
+        ----------
+        x: array-like, shape(M,)
+            x-axis
+        spectra: array-like, shape(M,) or (N, M,)
+            the signals to subtract a baseline from. If 2d, the signals
+            are assumed to be stacked on the first axis.
+        deg: int
+            the polynomial degree for fitting a baseline
+        signalRange: list of type(x), length 2
+            the x-interval where to expect the signal. The baseline is fitted to
+            all regions except the one defined by the interval.
+        Output
+        ------
+        spectra_nobl: array-like, shape(M,) or (N, M,)
+            the baseline subtracted spectra
+
+        """
+        if 'spectrum' not in data:
+            return
+        x = data.newt_x
+        spectra = data['spectrum']
+        mask = (x<signalRange[0]) | (x>signalRange[1])
+        if isinstance(x, xr.DataArray):
+            x_bl = x.where(mask, drop=True)
+            bl = spectra.sel(newt_x=x_bl)
+        else:
+            x_bl = x[mask]
+            if len(spectra.shape) == 1:
+                bl = spectra[mask]
+            else:
+                bl = spectra[:, mask]
+        fit = np.polyfit(x_bl, bl.T, self.POLY_DEG)
+        if len(spectra.shape) == 1:
+            return spectra - np.poly1d(fit)(x)
+        final_bl = np.empty(spectra.shape)
+        for t in range(spectra.shape[0]):
+            final_bl[t] = np.poly1d(fit[:, t])(x)
+        data['spectrum_nobg'] = spectra - final_bl
+        return spectra - final_bl
+
+    def xas(self, data, data_ref, thickness=1, plot=False,
+            plot_errors=True, xas_ylim=(-1, 3)):
+        key = 'spectrum_nobg' if 'spectrum_nobg' in data else 'spectrum'
+        if data['newt_x'].equals(data_ref['newt_x']) 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'])
+        
+        ds = xr.Dataset()
+        ds['sample'] = spectrum
+        ds['sample_std'] = std
+        ds['sample_std_err'] = std_err
+        ds['ref'] = spectrum_ref
+        ds['ref_std'] = std
+        ds['ref_std_err'] = std_err
+        ds['absorption'] = spectrum_ref / spectrum
+        ds['absorption_std'] = np.abs(ds['absorption']) * np.sqrt(
+            std_ref**2 / spectrum_ref**2 + std**2 / spectrum**2)
+        ds['absorption_stderr'] = np.abs(ds['absorption']) * np.sqrt(
+            (std_err_ref / spectrum_ref)**2 + (std_err / spectrum)**2)
+        ds['absorptionCoef'] = np.log(ds['absorption']) / thickness
+        ds['absorptionCoef_std'] = ds['absorption_std'] / (thickness * np.abs(ds['absorption']))
+        ds['absorptionCoef_stderr'] = ds['absorption_stderr'] / (thickness * np.abs(ds['absorption']))
+        
+        if plot:
+            plot_viking_xas(ds, plot_errors, xas_ylim)
+        return ds
+
+def plot_viking_xas(xas, plot_errors=True, xas_ylim=(-1, 3)):
+    fig, ax = plt.subplots(3, 1, figsize=(8,8), sharex=True)
+    ax[0].plot(xas.newt_x, xas['ref'])
+    ax[0].grid()
+
+    ax[1].plot(xas.newt_x, xas['sample'])
+    ax[1].grid()
+
+    ax[2].plot(xas.newt_x, xas['absorptionCoef'])
+    ax[2].set_ylim(*xas_ylim)
+    ax[2].grid()
+    
+    if plot_errors:
+        ax[0].fill_between(xas.newt_x,
+                           xas['ref'] - 1.96*xas['ref_std_err'], 
+                           xas['ref'] + 1.96*xas['ref_std_err'],
+                           alpha=0.2)
+        ax[1].fill_between(xas.newt_x,
+                           xas['sample'] - 1.96*xas['sample_std_err'], 
+                           xas['sample'] + 1.96*xas['sample_std_err'],
+                           alpha=0.2)
+        ax[2].fill_between(xas.newt_x,
+                           xas['absorptionCoef'] - 1.96*xas['absorptionCoef_stderr'], 
+                           xas['absorptionCoef'] + 1.96*xas['absorptionCoef_stderr'],
+                           alpha=0.2)