Added Wiener deconvolution.

pes_to_spec/test/offline_analysis.py

@@ -18,8 +18,10 @@ matplotlib.use('Agg')
 
 import matplotlib.pyplot as plt
 from matplotlib.gridspec import GridSpec
-from mpl_toolkits.axes_grid.inset_locator import InsetPosition
+from mpl_toolkits.axes_grid1.inset_locator import InsetPosition
 import seaborn as sns
+import scipy
+from scipy.signal import fftconvolve
 
 from typing import Dict, Optional
 
@@ -65,7 +67,8 @@ def plot_result(filename: str,
                 spec_raw_int: Optional[np.ndarray]=None,
                 pes: Optional[np.ndarray]=None,
                 pes_to_show: Optional[str]="",
-                pes_bin: Optional[np.ndarray]=None):
+                pes_bin: Optional[np.ndarray]=None,
+                wiener: Optional[np.ndarray]=None):
     """
     Plot result with uncertainty band.
 
@@ -78,6 +81,7 @@ def plot_result(filename: str,
       pes: PES spectrum for the inset.
       pes_to_show: Name of the channel shown.
       pes_bin: PES bins.
+      wiener: A Wiener filter to use to improve the filter estimate.
 
     """
     fig = plt.figure(figsize=(12, 8))
@@ -94,6 +98,9 @@ def plot_result(filename: str,
     #ax.fill_between(spec_raw_pe, spec_pred["expected"] - unc_pca, spec_pred["expected"] + unc_pca, facecolor='magenta', alpha=0.6, label="68% unc. (syst., PCA)")
     #if spec_raw_int is not None:
     #    ax.plot(spec_raw_pe, spec_raw_int, c='b', lw=1, ls='--', label="High-resolution measurement")
+    if wiener is not None:
+        deconvolved = fftconvolve(spec_pred["expected"], wiener, mode="same")
+        ax.plot(spec_raw_pe, deconvolved, c='g', ls='-.', lw=3, label="Wiener filter result")
     Y = np.amax(spec_smooth)
     ax.legend(frameon=False, borderaxespad=0, loc='upper left')
     ax.set(title=f"", #avg(stat unc) = {unc_stat}, avg(pca unc) = {unc_pca}",
@@ -239,6 +246,30 @@ def main():
     t += [time_ns() - start]
     t_names += ["Load"]
 
+    # plot Wiener filter
+    fig = plt.figure(figsize=(12, 8))
+    gs = GridSpec(1, 1)
+    ax = fig.add_subplot(gs[0, 0])
+    plt.plot(np.fft.fftshift(model.wiener_energy_ft), np.fft.fftshift(np.absolute(model.wiener_filter_ft)))
+    ax.set(title=f"",
+           xlabel=r"Reciprocal energy [1/eV]",
+           ylabel="Filter intensity [a.u.]",
+           yscale='log',
+           )
+    fig.savefig(os.path.join(args.directory, "wiener_ft.png"))
+    plt.close(fig)
+
+    fig = plt.figure(figsize=(12, 8))
+    gs = GridSpec(1, 1)
+    ax = fig.add_subplot(gs[0, 0])
+    plt.plot(model.wiener_energy, np.fft.fftshift(np.absolute(model.wiener_filter)))
+    ax.set(title=f"",
+           xlabel=r"Energy [eV]",
+           ylabel="Filter value [a.u.]",
+           )
+    fig.savefig(os.path.join(args.directory, "wiener.png"))
+    plt.close(fig)
+
     print("Check consistency")
     start = time_ns()
     Z = model.check_compatibility(pes_raw_t)
@@ -297,7 +328,7 @@ def main():
         ax2.tick_params(axis='both', which='major', labelsize=SMALL_SIZE)
         fig.savefig(os.path.join(args.directory, "intensity_vs_chi2.png"))
         plt.close(fig)
-        
+
         fig = plt.figure(figsize=(12, 8))
         gs = GridSpec(1, 1)
         ax = fig.add_subplot(gs[0, 0])
@@ -317,7 +348,7 @@ def main():
                 color='black', fontsize=15)
         fig.savefig(os.path.join(args.directory, "chi2.png"))
         plt.close(fig)
-        
+
         fig = plt.figure(figsize=(12, 8))
         gs = GridSpec(1, 1)
         ax = fig.add_subplot(gs[0, 0])
@@ -336,7 +367,7 @@ def main():
                 color='black', fontsize=15)
         fig.savefig(os.path.join(args.directory, "intensity.png"))
         plt.close(fig)
-        
+
         # rmse
         rmse = np.sqrt(np.mean((spec_smooth - spec_pred["expected"])**2, axis=1))
         fig = plt.figure(figsize=(12, 8))
@@ -350,7 +381,7 @@ def main():
                )
         fig.savefig(os.path.join(args.directory, "intensity_vs_rmse.png"))
         plt.close(fig)
-        
+
         fig = plt.figure(figsize=(12, 8))
         gs = GridSpec(1, 1)
         ax = fig.add_subplot(gs[0, 0])
@@ -370,7 +401,7 @@ def main():
                 color='black', fontsize=15)
         fig.savefig(os.path.join(args.directory, "rmse.png"))
         plt.close(fig)
-        
+
         # SPEC integral w.r.t XGM intensity
         fig = plt.figure(figsize=(12, 8))
         gs = GridSpec(1, 1)
@@ -382,7 +413,7 @@ def main():
                )
         fig.savefig(os.path.join(args.directory, "xgm_vs_intensity.png"))
         plt.close(fig)
-        
+
         # SPEC integral w.r.t XGM intensity
         fig = plt.figure(figsize=(12, 8))
         gs = GridSpec(1, 1)
@@ -394,7 +425,7 @@ def main():
                )
         fig.savefig(os.path.join(args.directory, "expected_vs_intensity.png"))
         plt.close(fig)
-        
+
         fig = plt.figure(figsize=(12, 8))
         gs = GridSpec(1, 1)
         ax = fig.add_subplot(gs[0, 0])
@@ -422,7 +453,8 @@ def main():
                     spec_raw_int[idx, :] if showSpec else None,
                     #pes=-pes_raw[pes_to_show][idx, first:last],
                     #pes_to_show=pes_to_show.replace('_', ' '),
-                    #pes_bin=np.arange(first, last)
+                    #pes_bin=np.arange(first, last),
+                    wiener=model.wiener_filter
                     )
         for ch in channels:
             plot_pes(os.path.join(args.directory, f"test_pes_{tid}_{ch}.png"),