diff --git a/pes_to_spec/__init__.py b/pes_to_spec/__init__.py
index ca6f3e9f1fcf657fec6836e8f13c3c178112bc3b..439057e6c3c9215cc40f5848237329de2dbbc43d 100644
--- a/pes_to_spec/__init__.py
+++ b/pes_to_spec/__init__.py
@@ -2,4 +2,4 @@
Estimate high-resolution photon spectrometer data from low-resolution non-invasive measurements.
"""
-VERSION = "0.0.9"
+VERSION = "0.1.0"
diff --git a/pes_to_spec/model.py b/pes_to_spec/model.py
index f54f7d43bd0d533256edbdbab8117e1996d68785..ff38768c59e3e087a38346b66de538088f5e875e 100644
--- a/pes_to_spec/model.py
+++ b/pes_to_spec/model.py
@@ -3,23 +3,18 @@ from __future__ import annotations
import joblib
import numpy as np
-import scipy.stats
from scipy.signal import fftconvolve
#from scipy.signal import find_peaks_cwt
from scipy.optimize import fmin_l_bfgs_b
from sklearn.decomposition import PCA
-from sklearn.pipeline import Pipeline, FeatureUnion
-from sklearn.preprocessing import FunctionTransformer
from sklearn.base import TransformerMixin, BaseEstimator
from sklearn.base import RegressorMixin
+from sklearn.pipeline import Pipeline
from sklearn.kernel_approximation import Nystroem
-from sklearn.preprocessing import PolynomialFeatures
-from sklearn.linear_model import ARDRegression
+#from sklearn.linear_model import ARDRegression
from sklearn.linear_model import BayesianRidge
from sklearn.neighbors import KernelDensity
from sklearn.ensemble import IsolationForest
-#from sklearn.neighbors import LocalOutlierFactor
-#from sklearn.covariance import EllipticEnvelope
from functools import reduce
from itertools import product
@@ -166,6 +161,7 @@ class FitModel(RegressorMixin, BaseEstimator):
self.pars = FitModel.get_pars(sc_op[0], self.structure)
self.nll_train = sc_op[1]
self.nll_train_expected = np.mean(self.nll(X, pars=self.pars), axis=0, keepdims=True)
+ return self
def predict(self, X: np.ndarray, return_std: bool=False) -> Union[Tuple[np.ndarray, np.ndarray], np.ndarray]:
"""
@@ -770,8 +766,8 @@ class Model(TransformerMixin, BaseEstimator):
"""
self.kde_xgm = KernelDensity()
self.kde_xgm.fit(intensity.reshape(-1, 1))
- self.mu_xgm = np.mean(intensity, axis=0)
- self.sigma_xgm = np.mean(intensity, axis=0)
+ self.mu_xgm = np.mean(intensity.reshape(-1), axis=0)
+ self.sigma_xgm = np.std(intensity.reshape(-1), axis=0)
q = np.quantile(intensity, [0.10, 0.90])
l, h = q[0], q[1]
x = intensity*((intensity > l) & (intensity < h)) + l*(intensity <= l) + h*(intensity >= h)
@@ -807,11 +803,7 @@ class Model(TransformerMixin, BaseEstimator):
x_t = self.x_model.fit_transform(low_res_select)
print("Fitting PCA on high-resolution data.")
y_t = self.y_model.fit_transform(high_res_data, smoothen__energy=high_res_photon_energy)
- intensity = np.sum(y_t, axis=1)
- self.kde_intensity = KernelDensity()
- self.kde_intensity.fit(intensity.reshape(-1, 1))
- self.mu_intensity = np.mean(intensity, axis=0)
- self.sigma_intensity = np.std(intensity, axis=0)
+
#self.fit_model.set_params(fex__gamma=1.0/float(x_t.shape[0]))
print("Fitting outlier detection")
self.ood['full'].fit(x_t)
@@ -827,6 +819,40 @@ class Model(TransformerMixin, BaseEstimator):
high_pca_unc = np.sqrt(np.mean((high_res - high_pca_rec)**2, axis=0, keepdims=True))
self.y_model['unc'].set_uncertainty(high_pca_unc)
+ print("Calculate transfer function")
+ # Model: z(e) = conv(h(e), d(e)) + n(e)
+ # d: true high-resolution data
+ # h: the effect of the model
+ # z: estimated high-resolution data
+ # n: noise (uncertainty)
+ # e: energy
+ # true signal (as far as we can get -- it is smoothened, but this is the model target)
+ d = high_res[inliers]
+ y_pred, n = self.fit_model.predict(x_t[inliers], return_std=True)
+ z = self.y_model['pca'].inverse_transform(y_pred)
+ n = np.fabs(self.y_model['pca'].inverse_transform(y_pred + n) - z)
+ e = high_res_photon_energy[0,:] if len(high_res_photon_energy.shape) == 2 else high_res_photon_energy
+ D = np.fft.fft(d)
+ Z = np.fft.fft(z)
+ V = np.fft.fft(n)
+ de = e[1] - e[0]
+ E = np.fft.fftfreq(len(e), de)
+ 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)
+ self.wiener_filter = np.fft.ifft(G)
+ self.wiener_filter_ft = G
+ self.wiener_energy = E
+
+ # get intensity effect
+ intensity = np.sum(z, axis=1)
+ self.kde_intensity = KernelDensity()
+ 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)
+
# for consistency check per channel
selection_model = self.x_select
low_res_selected = selection_model.transform(low_res_data, keep_dictionary_structure=True)
@@ -932,10 +958,15 @@ class Model(TransformerMixin, BaseEstimator):
self.fit_model,
self.channel_pca,
#self.channel_fit_model
- DataHolder(self.mu_intensity),
- DataHolder(self.sigma_intensity),
- DataHolder(self.mu_xgm),
- DataHolder(self.sigma_xgm),
+ DataHolder(dict(mu_intensity=self.mu_intensity,
+ sigma_intensity=self.sigma_intensity,
+ mu_xgm=self.mu_xgm,
+ sigma_xgm=self.sigma_xgm,
+ wiener_filter_ft=self.wiener_filter_ft,
+ wiener_filter=self.wiener_filter,
+ wiener_energy=self.wiener_energy,
+ )
+ ),
self.ood,
self.kde_xgm,
self.kde_intensity,
@@ -955,10 +986,7 @@ class Model(TransformerMixin, BaseEstimator):
x_model, y_model, fit_model,
channel_pca,
#channel_fit_model
- mu_intensity,
- sigma_intensity,
- mu_xgm,
- sigma_xgm,
+ extra,
ood,
kde_xgm,
kde_intensity,
@@ -971,11 +999,16 @@ class Model(TransformerMixin, BaseEstimator):
obj.channel_pca = channel_pca
#obj.channel_fit_model = channel_fit_model
obj.ood = ood
- obj.mu_intensity = mu_intensity.get_data()
- obj.sigma_intensity = sigma_intensity.get_data()
- obj.mu_xgm = mu_xgm.get_data()
- obj.sigma_xgm = sigma_xgm.get_data()
obj.kde_xgm = kde_xgm
obj.kde_intensity = kde_intensity
+
+ extra = extra.get_data()
+ obj.mu_intensity = extra["mu_intensity"]
+ obj.sigma_intensity = extra["sigma_intensity"]
+ obj.mu_xgm = extra["mu_xgm"]
+ obj.sigma_xgm = extra["sigma_xgm"]
+ obj.wiener_filter_ft = extra["wiener_filter_ft"]
+ obj.wiener_filter = extra["wiener_filter"]
+ 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 790a5b368054825f9672b165fe753b4ddbfbe993..9b9300f194e885841f478687c099f0f334d9ea5f 100755
--- a/pes_to_spec/test/offline_analysis.py
+++ b/pes_to_spec/test/offline_analysis.py
@@ -20,6 +20,8 @@ import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from mpl_toolkits.axes_grid.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,33 @@ 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.scatter(np.fft.fftshift(model.wiener_energy), 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])
+ 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)))
+ 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)
+
print("Check consistency")
start = time_ns()
Z = model.check_compatibility(pes_raw_t)
@@ -297,7 +331,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 +351,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 +370,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 +384,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 +404,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 +416,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 +428,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 +456,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"),