Functions for Viking spectrometer analysis

@teichman

src/toolbox_scs/detectors/viking.py

@@ -2,15 +2,73 @@ import numpy as np
 import xarray as xr
 import toolbox_scs as tb
 import matplotlib.pyplot as plt
+from scipy.optimize import curve_fit
+from .hrixs import gauss1d
 
 __all__ = ['Viking']
 
+
+
+def plot_viking_xas(xas, plot_errors=True, xas_ylim=(-1, 3)):
+    fig, ax = plt.subplots(3, 1, figsize=(7,7), sharex=True)
+    ax[0].plot(xas.newt_x, xas['I0'])
+    ax[0].grid()
+    ax[0].set_title('I0 spectra')
+
+    ax[1].plot(xas.newt_x, xas['It'])
+    ax[1].grid()
+    ax[1].set_title('It spectra')
+
+    ax[2].plot(xas.newt_x, xas['absorptionCoef'])
+    ax[2].set_ylim(*xas_ylim)
+    ax[2].grid()
+    ax[2].set_title('XAS')
+
+    if plot_errors:
+        ax[0].fill_between(xas.newt_x,
+                           xas['I0'] - 1.96*xas['I0_stderr'], 
+                           xas['I0'] + 1.96*xas['I0_stderr'],
+                           alpha=0.2)
+        ax[1].fill_between(xas.newt_x,
+                           xas['It'] - 1.96*xas['It_stderr'], 
+                           xas['It'] + 1.96*xas['It_stderr'],
+                           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)
+
+
+def plot_viking_calibration(spectra, x_pos, nrj, energy_calib,
+                            gaussian_fits, runNBs):
+    fig, ax = plt.subplots(2,1, figsize=(7,6), sharex=True)
+    for i in range(len(x_pos)):
+        x = spectra[i].newt_x.values
+        ax[0].plot(spectra[i].newt_x, spectra[i], color=f'C{i}',
+                   label=f'run {runNBs[i]}')
+        ax[0].plot(x, gauss1d(x, *gaussian_fits[i]),
+                 ls='--', color=f'C{i}')
+    ax[0].legend()
+    ax[0].set_ylabel('intensity [arb. un.]')
+    
+    ax[1].scatter(x_pos, nrj, label='data')
+    ax[1].plot(x_pos, np.polyval(energy_calib, x_pos), color='r',
+             label=f"{energy_calib[2]:.4f}+{energy_calib[1]:.4f}"
+                    f"*pixel+{energy_calib[0]:.4e}*pixel^2")
+    ax[1].legend()
+    ax[1].set_xlabel('pixel')
+    ax[1].set_ylabel('energy [eV]')
+    ax[1].grid()
+    fig.suptitle(f'Viking calibration, runs {runNBs}')
+
+
 # -----------------------------------------------------------------------------
 # Viking class
 class Viking:
     def __init__(self, proposalNB):
         self.PROPOSAL = proposalNB
         self.FIELDS = ['newton']
+        # 2 deg polynomial energy calibration
         self.ENERGY_CALIB = [0, 1, 0]
         # dark
         self.USE_DARK = False
@@ -33,7 +91,7 @@ class Viking:
                       'energy_calib', 'use_dark', 'poly_deg', 'fields')
         return {param: getattr(self, param.upper()) for param in params}
 
-    def from_run(self, runNB, proposal=None, add_attrs=True, calibrate=True):
+    def from_run(self, runNB, proposal=None, add_attrs=True):
         if proposal is None:
             proposal = self.PROPOSAL
         roi = {'newton': {'newton': {'roi': (self.Y_RANGE, self.X_RANGE),
@@ -41,9 +99,8 @@ class Viking:
         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']))
+        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():
@@ -63,7 +120,8 @@ class Viking:
         imgs = data['newton']
         if self.USE_DARK:
             imgs = imgs - self.dark_image
-        data['spectrum'] = imgs.mean(dim=self.INTEGRATE_DIM)
+        spectrum = imgs.mean(dim=self.INTEGRATE_DIM)
+        data['spectrum'] = data.attrs['photoelectrons_per_count']*spectrum
         return data
     
     def get_camera_gain(self, run):
@@ -180,6 +238,35 @@ class Viking:
 
     def xas(self, data, data_ref, thickness=1, plot=False,
             plot_errors=True, xas_ylim=(-1, 3)):
+        """
+        Given two independent datasets (one with sample and one reference),
+        this calculates the average XAS spectrum (absorption coefficient),
+        associated standard deviation and standard error. The absorption
+        coefficient is defined as -log(It/I0)/thickness.
+        
+        Parameters
+        ----------
+        data: xarray Dataset
+            the dataset containing the spectra with sample
+        data_ref: xarray Dataset
+            the dataset containing the spectra without sample
+        thickness: float
+            the thickness used for the calculation of the absorption
+            coefficient
+        plot: bool
+            If True, plot the resulting average spectra.
+        plot_errors: bool
+            If True, adds the 95% confidence interval on the spectra.
+        xas_ylim: tuple or list of float
+            the y limits for the XAS plot.
+        
+        Output
+        ------
+        xas: xarray Dataset
+            the dataset containing the computed XAS quantities:
+            I0, It, absorptionCoef and their associated errors.
+        """
+
         key = 'spectrum_nobg' if 'spectrum_nobg' in data else 'spectrum'
         if data['newt_x'].equals(data_ref['newt_x']) is False:
             return 
@@ -213,32 +300,40 @@ class Viking:
             plot_viking_xas(ds, plot_errors, xas_ylim)
         return ds
 
+    def calibrate(self, runList, plot=True):
+        self.ENERGY_CALIB = [0, 1, 0]
+        x_pos = []
+        nrj = []
+        spectra = []
+        gaussian_fits = []
+        runNBs = []
+        for i, runNB in enumerate(runList):
+            if plot:
+                print(runNB)
+            ds = self.from_run(runNB)
+            self.integrate(ds)
+            avg_spectrum = ds['spectrum'].mean(dim='trainId')
+            p0 = [np.max(avg_spectrum)-np.min(avg_spectrum),
+                  avg_spectrum.argmax().values, 10, np.min(avg_spectrum)]
+            try:
+                popt, pcov = curve_fit(gauss1d, avg_spectrum['newt_x'].values,
+                                       avg_spectrum.values, p0=p0)
+                x_pos.append(popt[1])
+                gaussian_fits.append(popt)
+                spectra.append(avg_spectrum)
+                runNBs.append(runNB)
+                nrj.append(tb.load_run_values(self.PROPOSAL, runNB)['nrj'])
+            except Exception as e:
+                print(f'error with run {runNB}:', e)
+        x_pos = np.array(x_pos)
+        nrj = np.array(nrj)
+        idx = np.argsort(x_pos)
+        x_pos = x_pos[idx]
+        nrj = nrj[idx]
+        energy_calib = np.polyfit(x_pos, nrj, 2)
+        if plot:
+            plot_viking_calibration(spectra, x_pos, nrj, energy_calib, 
+                                    gaussian_fits, runNBs)
+        self.ENERGY_CALIB = energy_calib
+        return energy_calib
 
-def plot_viking_xas(xas, plot_errors=True, xas_ylim=(-1, 3)):
-    fig, ax = plt.subplots(3, 1, figsize=(7,7), sharex=True)
-    ax[0].plot(xas.newt_x, xas['I0'])
-    ax[0].grid()
-    ax[0].set_title('I0 spectra')
-
-    ax[1].plot(xas.newt_x, xas['It'])
-    ax[1].grid()
-    ax[1].set_title('It spectra')
-
-    ax[2].plot(xas.newt_x, xas['absorptionCoef'])
-    ax[2].set_ylim(*xas_ylim)
-    ax[2].grid()
-    ax[2].set_title('XAS')
-
-    if plot_errors:
-        ax[0].fill_between(xas.newt_x,
-                           xas['I0'] - 1.96*xas['I0_stderr'], 
-                           xas['I0'] + 1.96*xas['I0_stderr'],
-                           alpha=0.2)
-        ax[1].fill_between(xas.newt_x,
-                           xas['It'] - 1.96*xas['It_stderr'], 
-                           xas['It'] + 1.96*xas['It_stderr'],
-                           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)