Added resolution estimate.

pes_to_spec/model.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import joblib
 
 import numpy as np
+import lmfit
 from scipy.signal import fftconvolve
 from sklearn.decomposition import IncrementalPCA, PCA
 from sklearn.base import TransformerMixin, BaseEstimator
@@ -32,6 +33,36 @@ def matching_two_ids(a: np.ndarray, b: np.ndarray) -> np.ndarray:
     unique_ids = list(set(a).intersection(b))
     return np.array(unique_ids)
 
+def fwhm(x: np.ndarray, y: np.ndarray) -> float:
+    """Return the full width at half maximum of x."""
+    # half maximum
+    half_max = np.amax(y)*0.5
+    # signum(y - half_max) is zero before and after the half maximum,
+    # and it is 1 in the range above the half maximum
+    # The difference will be +/- 1 only at the transitions
+    d = np.diff(np.sign(y - half_max))
+    left_idx = np.where(d > 0)[0]
+    right_idx = np.where(d < 0)[-1]
+    return x[right_idx] - x[left_idx]
+
+def get_resolution(x: np.ndarray, y: np.ndarray) -> float:
+    """
+    Get resolution from an auto-correlation function y.
+
+    Args:
+      x: Energy axis, assumed centred at 0, symmetric and equally spaced.
+      y: Auto-correlation function.
+    """
+    def gaussian_model(x, amp, cen, wid):
+        """Gaussian."""
+        return amp * np.exp(-0.5 * (x-cen)**2 / (wid**2))
+    gmodel = lmfit.Model(gaussian_model)
+    first = np.where(x > -0.5)[0][0]
+    last = np.where(x > 0.5)[0][0]
+    initial_guess = 0.1 # eV
+    result = gmodel.fit(y[first:last], x=x[first:last], amp=1.0, cen=0.0, wid=initial_guess)
+    return 2.355*result.best_values["wid"]
+
 class PromptNotFoundError(Exception):
     """
     Exception representing the error condition generated by not finding the prompt peak.
@@ -91,7 +122,8 @@ class HighResolutionSmoother(TransformerMixin, BaseEstimator):
         # get the centre value of the energy axis
         mu = energy[:, n_features//2, np.newaxis]
         # generate a gaussian
-        gaussian = np.exp(-0.5*(energy - mu)**2/self.high_res_sigma**2)
+        std = self.high_res_sigma*2.355
+        gaussian = np.exp(-0.5*(energy - mu)**2/std**2)
         gaussian /= np.sum(gaussian, axis=1, keepdims=True)
         # apply it to the data
         high_res_gc = fftconvolve(X, gaussian, mode="same", axes=1)
@@ -873,11 +905,14 @@ class Model(TransformerMixin, BaseEstimator):
         self.wiener_energy_ft = E
         self.transfer_function = H
         h = np.fft.fftshift(np.fft.ifft(H))
-        hmod = np.real(np.absolute(h))
         self.impulse_response = h
-        energy_mu = np.sum(e_axis*hmod)/np.sum(hmod)
-        energy_var = np.sum(((e_axis - energy_mu)**2)*hmod)/np.sum(hmod)
-        self.resolution = np.sqrt(energy_var)
+        self.auto_corr = np.mean(np.fft.fftshift(np.fft.ifft(np.absolute(np.fft.fft(z))**2), axes=(-1,)), axis=0)
+        self.auto_corr = np.real(self.auto_corr)
+        self.auto_corr /= np.amax(self.auto_corr)
+        try:
+            self.resulution = get_resolution(e_axis, self.auto_corr)
+        finally:
+            self.resolution = -1.0
         #print("Resolution:", self.resolution)
 
         # for consistency check per channel