diff --git a/pes_to_spec/model.py b/pes_to_spec/model.py index ff38768c59e3e087a38346b66de538088f5e875e..725e1e786e2c0ad6bfec7e6e7256bc063b027bb4 100644 --- a/pes_to_spec/model.py +++ b/pes_to_spec/model.py @@ -677,6 +677,7 @@ class Model(TransformerMixin, BaseEstimator): n_nonlinear_kernel: int=0, poly: bool=False, ): + self.high_res_sigma = high_res_sigma # models self.x_select = SelectRelevantLowResolution(channels, tof_start, delta_tof, poly=poly) x_model_steps = list() @@ -764,7 +765,7 @@ class Model(TransformerMixin, BaseEstimator): Return: weights. """ - self.kde_xgm = KernelDensity() + self.kde_xgm = KernelDensity(bandwidth="scott", kernel="gaussian") self.kde_xgm.fit(intensity.reshape(-1, 1)) self.mu_xgm = np.mean(intensity.reshape(-1), axis=0) self.sigma_xgm = np.std(intensity.reshape(-1), axis=0) @@ -837,18 +838,24 @@ class Model(TransformerMixin, BaseEstimator): V = np.fft.fft(n) de = e[1] - e[0] E = np.fft.fftfreq(len(e), de) + e_axis = np.linspace(-0.5*(e[-1] - e[0]), 0.5*(e[-1] - e[0]), len(e)) H = np.mean(Z/D, axis=0) N = np.mean(np.absolute(V)**2, axis=0) S = np.mean(np.absolute(D)**2, axis=0) # Wiener filter: G = np.conjugate(H) * S / (np.absolute(H)**2 * S + N) + # generate a gaussian + gaussian = np.exp(-0.5*(e_axis)**2/self.high_res_sigma**2) + gaussian /= np.sum(gaussian, axis=1, keepdims=True) + G *= gaussian self.wiener_filter = np.fft.ifft(G) self.wiener_filter_ft = G - self.wiener_energy = E + self.wiener_energy = e_axis + self.wiener_energy_ft = E # get intensity effect intensity = np.sum(z, axis=1) - self.kde_intensity = KernelDensity() + self.kde_intensity = KernelDensity(bandwidth="scott", kernel="gaussian") self.kde_intensity.fit(intensity.reshape(-1, 1)) self.mu_intensity = np.mean(intensity.reshape(-1), axis=0) self.sigma_intensity = np.std(intensity.reshape(-1), axis=0) @@ -965,6 +972,7 @@ class Model(TransformerMixin, BaseEstimator): wiener_filter_ft=self.wiener_filter_ft, wiener_filter=self.wiener_filter, wiener_energy=self.wiener_energy, + wiener_energy_ft=self.wiener_energy_ft, ) ), self.ood, @@ -1009,6 +1017,7 @@ class Model(TransformerMixin, BaseEstimator): obj.sigma_xgm = extra["sigma_xgm"] obj.wiener_filter_ft = extra["wiener_filter_ft"] obj.wiener_filter = extra["wiener_filter"] + obj.wiener_energy_ft = extra["wiener_energy_ft"] obj.wiener_energy = extra["wiener_energy"] return obj diff --git a/pes_to_spec/test/offline_analysis.py b/pes_to_spec/test/offline_analysis.py index 9b9300f194e885841f478687c099f0f334d9ea5f..2c5effd9f17d58a842069c6b894188c603e5022e 100755 --- a/pes_to_spec/test/offline_analysis.py +++ b/pes_to_spec/test/offline_analysis.py @@ -250,7 +250,7 @@ def main(): fig = plt.figure(figsize=(12, 8)) gs = GridSpec(1, 1) ax = fig.add_subplot(gs[0, 0]) - plt.scatter(np.fft.fftshift(model.wiener_energy), np.fft.fftshift(np.absolute(model.wiener_filter_ft))) + 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.]", @@ -262,13 +262,10 @@ def main(): fig = plt.figure(figsize=(12, 8)) gs = GridSpec(1, 1) ax = fig.add_subplot(gs[0, 0]) - de = spec_raw_pe[0, -1] - spec_raw_pe[0,0] - de = np.linspace(-0.5*de, 0.5*de, spec_raw_pe.shape[1]) - plt.scatter(de, np.fft.fftshift(np.absolute(model.wiener_filter))) + 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.]", - yscale='log', ) fig.savefig(os.path.join(args.directory, "wiener.png")) plt.close(fig)