Add fluence calibration function

src/toolbox_scs/misc/laser_utils.py

 __all__ = [
     'degToRelPower',
     'positionToDelay',
+    'fluenceCalibration',
 ]
 
+import numpy as np
+import matplotlib.pyplot as plt
 
 def positionToDelay(pos, origin=0, invert = False, reflections=1):
     ''' converts a motor position in mm into optical delay in picosecond
@@ -30,3 +33,82 @@ def degToRelPower(x, theta0=0):
             array-like relative power
     '''
     return np.sin(2*(x-theta0)*np.pi/180)**2
+
+def fluenceCalibration(hwp, power_mW, npulses, w0x, w0y=None,
+                       train_rep_rate=10, fit_order=1,
+                       plot=True, xlabel='HWP [%]'):
+    """
+    Given a series of relative powers or half wave plate angles and
+    averaged powers in mW, this routine calculates the corresponding
+    fluence and fits a polynomial to the data.
+    Parameters
+    ----------
+    hwp: array-like (N)
+        angle or relative power from the half wave plate
+    power_mW: array-like (N)
+        measured power in mW by powermeter
+    npulses: int
+        number of pulses per train during power measurement
+    w0x: float
+        radius at 1/e^2 in x-axis
+    w0y: float, optional
+        radius at 1/e^2 in y-axis. If None, w0y=w0x is assumed.
+    train_rep_rate: float
+        repetition rate of the FEL, by default equals to 10 Hz.
+    fit_order: int
+        order of the polynomial fit
+    plot: bool
+        Plot the results if True
+    xlabel: str
+        xlabel for the plot
+
+    Output
+    ------
+    F: ndarray (N)
+        fluence in mJ/cm^2
+    fit: ndarray
+        coefficients of the polynomial fit
+    """
+    power = np.array(power_mW)
+    hwp = np.array(hwp)
+    E = power/(train_rep_rate*npulses)*1e-3 # pulse energy in J
+    if w0y is None:
+        w0y = w0x
+    F = 2*E/(np.pi*w0x*w0y) # fluence in J/m^2
+    
+    fit_E = np.polyfit(hwp, E*1e6, fit_order)
+    fit_F = np.polyfit(hwp, F*1e-1, fit_order)
+    x = np.linspace(hwp.min(), hwp.max(), 100)
+    if plot:
+        plt.figure(figsize=(7,4))
+        ax = plt.subplot(121)
+        plt.title(f'w0x = {w0x*1e6:.0f} $\mu$m, w0y = {w0y*1e6:.0f} $\mu$m')
+        plt.plot(hwp, E*1e6, 'o', label='data')
+        fit_label=''
+        for i in range(len(fit_E)-1, 1, -1):
+            fit_label += f'{fit_E[i]:.2g}x$^{i}$ + '
+            if i%2 == 0:
+                fit_label +='\n'
+        fit_label += f'{fit_E[-2]:.2g}x + {fit_E[-1]:.2g}'
+        plt.plot(x, np.poly1d(fit_E)(x), label=fit_label)
+        plt.ylabel('Pulse energy [$\mu$J]')
+        plt.xlabel('HWP [%]')
+        plt.legend()
+        plt.grid()
+        plt.subplot(122, sharex=ax)
+        plt.title(f'w0x = {w0x*1e6:.0f} $\mu$m, w0y = {w0y*1e6:.0f} $\mu$m')
+        plt.plot(hwp, F*1e-1, 'o', label='data')
+        fit_label=''
+        for i in range(len(fit_F)-1, 1, -1):
+            fit_label += f'{fit_F[i]:.2g}x$^{i}$ + '
+            if i%2 == 0:
+                fit_label +='\n'
+        fit_label += f'{fit_F[-2]:.2g}x + {fit_F[-1]:.2g}'
+        #fit_label = '+'.join([f'{f:.2g}x$^{i}$' for i, f in enumerate(fit_F[:-1])])
+        plt.plot(x, np.poly1d(fit_F)(x), label=fit_label)
+        plt.ylabel('Fluence [mJ/cm$^2$]')
+        plt.xlabel('HWP [%]')
+        plt.legend()
+        plt.grid()
+        plt.tight_layout()
+    return F*1e-1, fit_F, E*1e6, fit_E
\ No newline at end of file