Nicer plots

pes_to_spec/test/offline_analysis.py

@@ -15,6 +15,8 @@ matplotlib.use('Agg')
 
 import matplotlib.pyplot as plt
 from matplotlib.gridspec import GridSpec
+from mpl_toolkits.axes_grid.inset_locator import (inset_axes, InsetPosition,
+                                                  mark_inset)
 
 from typing import Dict, Optional
 
@@ -23,7 +25,7 @@ import pandas as pd
 
 SMALL_SIZE = 12
 MEDIUM_SIZE = 18
-BIGGER_SIZE = 24
+BIGGER_SIZE = 22
 
 plt.rc('font', size=BIGGER_SIZE)         # controls default text sizes
 plt.rc('axes', titlesize=BIGGER_SIZE)    # fontsize of the axes title
@@ -53,7 +55,7 @@ def plot_pes(filename: str, pes_raw_int: np.ndarray, first: int, last: int):
     fig.savefig(filename)
     plt.close(fig)
 
-def plot_result(filename: str, spec_pred: Dict[str, np.ndarray], spec_smooth: np.ndarray, spec_raw_pe: np.ndarray, spec_raw_int: Optional[np.ndarray]=None):
+def plot_result(filename: str, spec_pred: Dict[str, np.ndarray], spec_smooth: np.ndarray, spec_raw_pe: np.ndarray, spec_raw_int: Optional[np.ndarray]=None, pes: Optional[np.ndarray]=None, pes_to_show: Optional[str]="", pes_bin: Optional[np.ndarray]=None):
     """
     Plot result with uncertainty band.
 
@@ -63,6 +65,9 @@ def plot_result(filename: str, spec_pred: Dict[str, np.ndarray], spec_smooth: np
       spec_smooth: Smoothened expected result with shape (features,).
       spec_raw_pe: x axis with the photon energy in eV.
       spec_raw_int: Original true expected result with shape (features,).
+      pes: PES spectrum for the inset.
+      pes_to_show: Name of the channel shown.
+      pes_bin: PES bins.
 
     """
     fig = plt.figure(figsize=(12, 8))
@@ -71,19 +76,39 @@ def plot_result(filename: str, spec_pred: Dict[str, np.ndarray], spec_smooth: np
     unc_stat = np.mean(spec_pred["unc"])
     unc_pca = np.mean(spec_pred["pca"])
     unc = np.sqrt(unc_stat**2 + unc_pca**2)
-    ax.plot(spec_raw_pe, spec_smooth, c='b', lw=3, label="High-resolution measurement (smoothened)")
-    ax.plot(spec_raw_pe, spec_pred["expected"], c='r', lw=3, label="High-resolution prediction")
-    ax.fill_between(spec_raw_pe, spec_pred["expected"] - unc, spec_pred["expected"] + unc, facecolor='red', alpha=0.6, label="68% unc.")
+    ax.plot(spec_raw_pe, spec_smooth, c='b', lw=3, label="High-res. measurement (smoothened)")
+    ax.plot(spec_raw_pe, spec_pred["expected"], c='r', ls='--', lw=3, label="High-res. prediction")
+    #ax.fill_between(spec_raw_pe, spec_pred["expected"] - unc, spec_pred["expected"] + unc, facecolor='green', alpha=0.6, label="68% unc.")
+    ax.fill_between(spec_raw_pe, spec_pred["expected"] - unc, spec_pred["expected"] + unc, facecolor='gold', alpha=0.5, label="68% unc.")
     #ax.fill_between(spec_raw_pe, spec_pred["expected"] - unc_stat, spec_pred["expected"] + unc_stat, facecolor='red', alpha=0.6, label="68% unc. (stat.)")
     #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")
     Y = np.amax(spec_smooth)
-    ax.legend(frameon=False, borderaxespad=0)
+    ax.legend(frameon=False, borderaxespad=0, loc='upper left')
     ax.set(title=f"", #avg(stat unc) = {unc_stat}, avg(pca unc) = {unc_pca}",
            xlabel="Photon energy [eV]",
            ylabel="Intensity",
            ylim=(0, 1.2*Y))
+    if pes is not None:
+        ax2 = plt.axes([0,0,1,1])
+        # Manually set the position and relative size of the inset axes within ax1
+        ip = InsetPosition(ax, [0.65,0.6,0.35,0.4])
+        ax2.set_axes_locator(ip)
+        Ypes = np.amax(pes)
+        ax2.plot(pes_bin, pes, c='black', lw=3)
+        ax2.set(title=f"Low-resolution example data",
+                xlabel="Bin",
+                ylabel=f"{pes_to_show}",
+                ylim=(0, 1.2*Ypes),
+                #labelsize=SMALL_SIZE,
+                #xticklabels=dict(fontdict=dict(fontsize=SMALL_SIZE)),
+                #yticklabels=dict(fontdict=dict(fontsize=SMALL_SIZE)),
+                )
+        ax2.title.set_size(SMALL_SIZE)
+        ax2.xaxis.label.set_size(SMALL_SIZE)
+        ax2.yaxis.label.set_size(SMALL_SIZE)
+        ax2.tick_params(axis='both', which='major', labelsize=SMALL_SIZE)
     fig.savefig(filename)
     plt.close(fig)
 
@@ -191,6 +216,9 @@ def main():
     print("Plotting")
     spec_smooth = model.preprocess_high_res(spec_raw_int)
     first, last = model.get_low_resolution_range()
+    first += 10
+    last -= 100
+    pes_to_show = 'channel_1_D'
     # plot
     for tid in test_tids:
         idx = np.where(tid==tids)[0][0]
@@ -200,7 +228,11 @@ def main():
                        for k, item in spec_pred.items()},
                     spec_smooth[idx, :],
                     spec_raw_pe[idx, :],
-                    spec_raw_int[idx, :])
+                    spec_raw_int[idx, :],
+                    pes=-pes_raw[pes_to_show][idx, first:last],
+                    pes_to_show=pes_to_show.replace('_', ' '),
+                    pes_bin=np.arange(first, last)
+                    )
         for ch in channels:
             plot_pes(f"test_pes_{tid}_{ch}.png", pes_raw[ch][idx, first:last], first, last)