diff --git a/pes_to_spec/bnn.py b/pes_to_spec/bnn.py
index 887512e093e555bb65bef8442585eb0c84801146..5f8832ee8448f975f21851ef514c934b581aae62 100644
--- a/pes_to_spec/bnn.py
+++ b/pes_to_spec/bnn.py
@@ -2,13 +2,259 @@ from sklearn.base import BaseEstimator, RegressorMixin
from typing import Any, Dict, Optional, Union, Tuple
import numpy as np
+import math
from scipy.special import gamma
import torch
import torch.nn as nn
-import torchbnn as bnn
+import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader
+class BayesLinearEmpiricalPrior(nn.Module):
+ """
+ Applies Bayesian Linear
+
+ Args:
+ prior_mu (Float): mean of prior normal distribution.
+ prior_sigma (Float): sigma of prior normal distribution.
+
+ """
+ __constants__ = ['prior_mu', 'prior_sigma', 'bias', 'in_features', 'out_features']
+
+ def __init__(self, prior_mu, prior_sigma, in_features, out_features, bias=True):
+ super(BayesLinearEmpiricalPrior, self).__init__()
+ self.in_features = in_features
+ self.out_features = out_features
+
+ self.prior_mu = prior_mu
+ self.prior_sigma = prior_sigma
+
+ self.prior_log_sigma_w = nn.Parameter(torch.ones((out_features, in_features))*np.log(prior_sigma))
+ self.prior_log_sigma_b = nn.Parameter(torch.ones((out_features,))*np.log(prior_sigma))
+
+ self.weight_mu = nn.Parameter(torch.Tensor(out_features, in_features))
+ self.weight_log_sigma = nn.Parameter(torch.Tensor(out_features, in_features))
+ self.register_buffer('weight_eps', None)
+
+ if bias is None or bias is False :
+ self.bias = False
+ else :
+ self.bias = True
+
+ if self.bias:
+ self.bias_mu = nn.Parameter(torch.Tensor(out_features))
+ self.bias_log_sigma = nn.Parameter(torch.Tensor(out_features))
+ self.register_buffer('bias_eps', None)
+ else:
+ self.register_parameter('bias_mu', None)
+ self.register_parameter('bias_log_sigma', None)
+ self.register_buffer('bias_eps', None)
+
+ self.reset_parameters()
+
+ def reset_parameters(self):
+ stdv = 1. / np.sqrt(self.weight_mu.size(1))
+ self.weight_mu.data.uniform_(-stdv, stdv)
+ self.weight_log_sigma.data.fill_(np.log(self.prior_sigma))
+ if self.bias :
+ self.bias_mu.data.uniform_(-stdv, stdv)
+ self.bias_log_sigma.data.fill_(np.log(self.prior_sigma))
+
+ def freeze(self) :
+ self.weight_eps = torch.randn_like(self.weight_log_sigma)
+ if self.bias :
+ self.bias_eps = torch.randn_like(self.bias_log_sigma)
+
+ def unfreeze(self) :
+ self.weight_eps = None
+ if self.bias :
+ self.bias_eps = None
+
+ def forward(self, input):
+ r"""
+ Overriden.
+ """
+ if self.weight_eps is None :
+ weight = self.weight_mu + torch.exp(self.weight_log_sigma) * torch.randn_like(self.weight_log_sigma)
+ else :
+ weight = self.weight_mu + torch.exp(self.weight_log_sigma) * self.weight_eps
+
+ if self.bias:
+ if self.bias_eps is None :
+ bias = self.bias_mu + torch.exp(self.bias_log_sigma) * torch.randn_like(self.bias_log_sigma)
+ else :
+ bias = self.bias_mu + torch.exp(self.bias_log_sigma) * self.bias_eps
+ else :
+ bias = None
+
+ return F.linear(input, weight, bias)
+
+ def extra_repr(self):
+ r"""
+ Overriden.
+ """
+ return 'prior_mu={}, prior_sigma={}, in_features={}, out_features={}, bias={}'.format(self.prior_mu, self.prior_sigma, self.in_features, self.out_features, self.bias is not None)
+
+class BayesLinear(nn.Module):
+ r"""
+ Applies Bayesian Linear
+
+ Arguments:
+ prior_mu (Float): mean of prior normal distribution.
+ prior_sigma (Float): sigma of prior normal distribution.
+
+ """
+ __constants__ = ['prior_mu', 'bias', 'in_features', 'out_features']
+
+ def __init__(self, prior_mu, prior_sigma, in_features, out_features, bias=True):
+ super(BayesLinear, self).__init__()
+ self.in_features = in_features
+ self.out_features = out_features
+
+ self.prior_mu = prior_mu
+ self.prior_log_sigma = nn.Parameter(torch.ones(1)*np.log(prior_sigma), requires_grad=True)
+
+ self.weight_mu = nn.Parameter(torch.Tensor(out_features, in_features))
+ self.weight_log_sigma = nn.Parameter(torch.Tensor(out_features, in_features))
+ self.register_buffer('weight_eps', None)
+
+ if bias is None or bias is False:
+ self.bias = False
+ else :
+ self.bias = True
+
+ if self.bias:
+ self.bias_mu = nn.Parameter(torch.Tensor(out_features))
+ self.bias_log_sigma = nn.Parameter(torch.Tensor(out_features))
+ self.register_buffer('bias_eps', None)
+ else:
+ self.register_parameter('bias_mu', None)
+ self.register_parameter('bias_log_sigma', None)
+ self.register_buffer('bias_eps', None)
+
+ self.reset_parameters()
+
+ def reset_parameters(self):
+ stdv = 1. / np.sqrt(self.weight_mu.size(1))
+ self.weight_mu.data.uniform_(-stdv, stdv)
+ self.weight_log_sigma.data.fill_(self.prior_log_sigma.detach()[0])
+ if self.bias :
+ self.bias_mu.data.uniform_(-stdv, stdv)
+ self.bias_log_sigma.data.fill_(self.prior_log_sigma.detach()[0])
+
+ def freeze(self) :
+ self.weight_eps = torch.randn_like(self.weight_log_sigma)
+ if self.bias :
+ self.bias_eps = torch.randn_like(self.bias_log_sigma)
+
+ def unfreeze(self) :
+ self.weight_eps = None
+ if self.bias:
+ self.bias_eps = None
+
+ def forward(self, input):
+ r"""
+ Overriden.
+ """
+ if self.weight_eps is None :
+ weight = self.weight_mu + torch.exp(self.weight_log_sigma) * torch.randn_like(self.weight_log_sigma)
+ else :
+ weight = self.weight_mu + torch.exp(self.weight_log_sigma) * self.weight_eps
+
+ if self.bias:
+ if self.bias_eps is None :
+ bias = self.bias_mu + torch.exp(self.bias_log_sigma) * torch.randn_like(self.bias_log_sigma)
+ else :
+ bias = self.bias_mu + torch.exp(self.bias_log_sigma) * self.bias_eps
+ else :
+ bias = None
+
+ return F.linear(input, weight, bias)
+
+ def extra_repr(self):
+ r"""
+ Overriden.
+ """
+ return 'prior_mu={}, prior_sigma={}, in_features={}, out_features={}, bias={}'.format(self.prior_mu, self.prior_sigma, self.in_features, self.out_features, self.bias is not None)
+
+def _kl_loss(mu_0, log_sigma_0, mu_1, log_sigma_1) :
+ """
+ An method for calculating KL divergence between two Normal distribtuion.
+
+ Arguments:
+ mu_0 (Float) : mean of normal distribution.
+ log_sigma_0 (Float): log(standard deviation of normal distribution).
+ mu_1 (Float): mean of normal distribution.
+ log_sigma_1 (Float): log(standard deviation of normal distribution).
+
+ """
+ if isinstance(log_sigma_1, float):
+ sigma_1 = np.exp(log_sigma_1)
+ else:
+ sigma_1 = torch.exp(log_sigma_1)
+ kl = log_sigma_1 - log_sigma_0 + \
+ (torch.exp(log_sigma_0)**2 + (mu_0-mu_1)**2)/(2*sigma_1**2) - 0.5
+ return kl.sum()
+
+class BKLLoss(nn.Module):
+ """
+ Loss for calculating KL divergence of baysian neural network model.
+
+ Args:
+ reduction (string, optional): Specifies the reduction to apply to the output:
+ ``'mean'``: the sum of the output will be divided by the number of
+ elements of the output.
+ ``'sum'``: the output will be summed.
+ last_layer_only (Bool): True for return only the last layer's KL divergence.
+ """
+ __constants__ = ['reduction']
+
+ def __init__(self, reduction='mean', last_layer_only=False):
+ super(BKLLoss, self).__init__()
+ self.last_layer_only = last_layer_only
+ self.reduction = reduction
+
+ def forward(self, model):
+ """
+ Args:
+ model (nn.Module): a model to be calculated for KL-divergence.
+ """
+ #return bayesian_kl_loss(model, reduction=self.reduction, last_layer_only=self.last_layer_only)
+ device = torch.device("cuda" if next(model.parameters()).is_cuda else "cpu")
+ kl = torch.Tensor([0]).to(device)
+ kl_sum = torch.Tensor([0]).to(device)
+ n = torch.Tensor([0]).to(device)
+
+ for m in model.modules() :
+ if isinstance(m, (BayesLinearEmpiricalPrior)):
+ kl = _kl_loss(m.weight_mu, m.weight_log_sigma, m.prior_mu, m.prior_log_sigma_w)
+ kl_sum += kl
+ n += len(m.weight_mu.view(-1))
+
+ if m.bias :
+ kl = _kl_loss(m.bias_mu, m.bias_log_sigma, m.prior_mu, m.prior_log_sigma_b)
+ kl_sum += kl
+ n += len(m.bias_mu.view(-1))
+ if isinstance(m, (BayesLinear)):
+ kl = _kl_loss(m.weight_mu, m.weight_log_sigma, m.prior_mu, m.prior_log_sigma)
+ kl_sum += kl
+ n += len(m.weight_mu.view(-1))
+
+ if m.bias :
+ kl = _kl_loss(m.bias_mu, m.bias_log_sigma, m.prior_mu, m.prior_log_sigma)
+ kl_sum += kl
+ n += len(m.bias_mu.view(-1))
+
+ if self.last_layer_only or n == 0 :
+ return kl
+
+ if self.reduction == 'mean':
+ return kl_sum/n
+ elif self.reduction == 'sum':
+ return kl_sum
+ else:
+ raise ValueError(f"{self.reduction} is not valid")
+
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self, name, fmt=':f'):
@@ -49,6 +295,7 @@ class ProgressMeter(object):
return '[' + fmt + '/' + fmt.format(num_batches) + ']'
+
class BNN(nn.Module):
"""
A model Bayesian Neural network.
@@ -61,13 +308,13 @@ class BNN(nn.Module):
between the prediction and the true value. The standard deviation of the Gaussian is left as a
parameter to be fit: sigma.
"""
- def __init__(self, input_dimension: int=1, output_dimension: int=1):
+ def __init__(self, input_dimension: int=1, output_dimension: int=1, rvm: bool=False):
super(BNN, self).__init__()
hidden_dimension = 50
# controls the aleatoric uncertainty
self.log_isigma2 = nn.Parameter(-torch.ones(1, output_dimension)*np.log(0.1**2), requires_grad=True)
# controls the weight hyperprior
- self.log_ilambda2 = nn.Parameter(-torch.ones(1)*np.log(0.1**2), requires_grad=True)
+ self.log_ilambda2 = -np.log(0.1**2)
# inverse Gamma hyper prior alpha and beta
#
@@ -82,8 +329,8 @@ class BNN(nn.Module):
# and the only regularization is to prevent the weights from becoming > 18 + 3 sqrt(var) ~= 50, making this a very loose regularization.
# An alternative would be to set the (alpha, beta) both to very low values, whichmakes the hyper prior become closer to the non-informative Jeffrey's prior.
# Using this alternative (ie: (0.1, 0.1) for the weights' hyper prior) leads to very large lambda and numerical issues with the fit.
- self.alpha_lambda = 0.001
- self.beta_lambda = 0.001
+ self.alpha_lambda = 0.0001
+ self.beta_lambda = 0.0001
# Hyperprior choice on the likelihood noise level:
# The likelihood noise level is controlled by sigma in the likelihood and it should be allowed to be very broad, but different
@@ -92,20 +339,45 @@ class BNN(nn.Module):
# Making both alpha and beta small makes the gamma distribution closer to the Jeffey's prior, which makes it non-informative
# This seems to lead to a larger training time, though.
# Since, after standardization, we know to expect the variance to be of order (1), we can select also alpha and beta leading to high variance in this range
- self.alpha_sigma = 0.001
- self.beta_sigma = 0.001
-
- self.model = nn.Sequential(
- bnn.BayesLinear(prior_mu=0.0,
- prior_sigma=torch.exp(-0.5*self.log_ilambda2),
+ self.alpha_sigma = 0.0001
+ self.beta_sigma = 0.0001
+
+ if rvm:
+ self.model = nn.Sequential(
+ BayesLinearEmpiricalPrior(prior_mu=0.0,
+ prior_sigma=np.exp(-0.5*self.log_ilambda2),
+ in_features=input_dimension,
+ out_features=hidden_dimension),
+ nn.ReLU(),
+ BayesLinearEmpiricalPrior(prior_mu=0.0,
+ prior_sigma=np.exp(-0.5*self.log_ilambda2),
+ in_features=hidden_dimension,
+ out_features=output_dimension)
+ )
+ else:
+ self.model = nn.Sequential(
+ BayesLinear(prior_mu=0.0,
+ prior_sigma=np.exp(-0.5*self.log_ilambda2),
in_features=input_dimension,
out_features=hidden_dimension),
- nn.ReLU(),
- bnn.BayesLinear(prior_mu=0.0,
- prior_sigma=torch.exp(-0.5*self.log_ilambda2),
+ nn.ReLU(),
+ BayesLinear(prior_mu=0.0,
+ prior_sigma=np.exp(-0.5*self.log_ilambda2),
in_features=hidden_dimension,
out_features=output_dimension)
- )
+ )
+ self.rvm = rvm
+
+ def prune(self):
+ """Prune weights."""
+ with torch.no_grad():
+ for layer in self.model.modules():
+ if isinstance(layer, BayesLinearEmpiricalPrior):
+ log_isigma2 = -2.0*layer.prior_log_sigma_w
+ isigma2 = torch.exp(log_isigma2)
+ keep = isigma2 < 1e4
+ layer.weight_mu[~keep] *= 0.0
+ layer.weight_log_sigma[~keep] = -12.0
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
@@ -138,10 +410,21 @@ class BNN(nn.Module):
# with a standardized input, this hyperprior forces sigma to be
# on avg. 1 and it is broad enough to allow for different sigma
isigma2 = torch.exp(self.log_isigma2)
- neg_log_hyperprior_noise = self.neg_log_gamma(self.log_isigma2, isigma2, self.alpha_sigma, self.beta_sigma)
- ilambda2 = torch.exp(self.log_ilambda2)
- neg_log_hyperprior_weights = self.neg_log_gamma(self.log_ilambda2, ilambda2, self.alpha_lambda, self.beta_lambda)
- return neg_log_hyperprior_noise.sum() + neg_log_hyperprior_weights.sum()
+ neg_log_hyperprior_noise = self.neg_log_gamma(self.log_isigma2, isigma2, self.alpha_sigma, self.beta_sigma).sum()
+ if self.rvm:
+ log_ilambda2 = [-2.0*self.model[0].prior_log_sigma_w,
+ -2.0*self.model[2].prior_log_sigma_w,
+ -2.0*self.model[0].prior_log_sigma_b,
+ -2.0*self.model[2].prior_log_sigma_b
+ ]
+ else:
+ log_ilambda2 = [-2.0*self.model[0].prior_log_sigma,
+ -2.0*self.model[2].prior_log_sigma,
+ ]
+ ilambda2 = [torch.exp(k) for k in log_ilambda2]
+ neg_log_hyperprior_weights = sum(self.neg_log_gamma(log_k, k, self.alpha_lambda, self.beta_lambda).sum()
+ for log_k, k in zip(log_ilambda2, ilambda2))
+ return neg_log_hyperprior_noise + neg_log_hyperprior_weights
def aleatoric_uncertainty(self) -> torch.Tensor:
"""
@@ -154,7 +437,18 @@ class BNN(nn.Module):
"""
Get the weights precision.
"""
- return torch.exp(self.log_ilambda2[0])
+ if self.rvm:
+ log_ilambda2 = [-2.0*self.model[0].prior_log_sigma_w,
+ -2.0*self.model[2].prior_log_sigma_w,
+ -2.0*self.model[0].prior_log_sigma_b,
+ -2.0*self.model[2].prior_log_sigma_b
+ ]
+ else:
+ log_ilambda2 = [-2.0*self.model[0].prior_log_sigma,
+ -2.0*self.model[2].prior_log_sigma,
+ ]
+ ilambda2 = [torch.exp(k) for k in log_ilambda2]
+ return sum(k.mean() for k in ilambda2)/len(ilambda2)
class BNNModel(RegressorMixin, BaseEstimator):
"""
@@ -162,14 +456,16 @@ class BNNModel(RegressorMixin, BaseEstimator):
Args:
"""
- def __init__(self, state_dict=None):
+ def __init__(self, state_dict=None, rvm: bool=False, n_epochs: int=250):
if state_dict is not None:
Nx = state_dict["model.0.weight_mu"].shape[1]
Ny = state_dict["model.2.weight_mu"].shape[0]
- self.model = BNN(Nx, Ny)
+ self.model = BNN(Nx, Ny, rvm=rvm)
self.model.load_state_dict(state_dict)
else:
- self.model = BNN()
+ self.model = BNN(rvm=rvm)
+ self.rvm = rvm
+ self.n_epochs = n_epochs
self.model.eval()
def state_dict(self) -> Dict[str, Any]:
@@ -197,7 +493,7 @@ class BNNModel(RegressorMixin, BaseEstimator):
torch.from_numpy(weights))
# create model
- self.model = BNN(X.shape[1], y.shape[1])
+ self.model = BNN(X.shape[1], y.shape[1], rvm=self.rvm)
# prepare data loader
B = 50
@@ -218,12 +514,11 @@ class BNNModel(RegressorMixin, BaseEstimator):
weight_prior /= float(B)
# KL loss
- kl_loss = bnn.BKLLoss(reduction='sum', last_layer_only=False)
+ kl_loss = BKLLoss(reduction='sum', last_layer_only=False)
# train
self.model.train()
- epochs = 500
- for epoch in range(epochs):
+ for epoch in range(self.n_epochs):
meter = {k: AverageMeter(k, ':6.3f')
for k in ('loss', '-log(lkl)', '-log(prior)', '-log(hyper)', 'sigma', 'w.prec.')}
progress = ProgressMeter(
@@ -251,6 +546,9 @@ class BNNModel(RegressorMixin, BaseEstimator):
meter['w.prec.'].update(self.model.w_precision().detach().cpu().item(), B)
progress.display(len(loader))
+ if self.rvm:
+ self.model.prune()
+
self.model.eval()
return self
diff --git a/pes_to_spec/model.py b/pes_to_spec/model.py
index 5751510b4a5e98187ffd1081849c4dba0223d2fa..4d93ec82f883e304963e96c4c698b2cd100c2a1d 100644
--- a/pes_to_spec/model.py
+++ b/pes_to_spec/model.py
@@ -260,12 +260,14 @@ class SelectRelevantLowResolution(TransformerMixin, BaseEstimator):
raise NotImplementedError("The low-resolution data cannot be transformed before the prompt has been identified. Call the fit function first.")
if pulse_spacing is None:
pulse_spacing = {ch: [0] for ch in X.keys()}
- y = X
+ y = {channel: item for channel, item in X.items()
+ if channel in self.channels}
if self.delta_tof is not None:
first = max(0, self.tof_start - self.delta_tof)
last = min(X[self.channels[0]].shape[1], self.tof_start + self.delta_tof)
y = {channel: np.stack([item[:, (first + delta):(last + delta)] for delta in pulse_spacing[channel]], axis=1)
- for channel, item in X.items()}
+ for channel, item in X.items()
+ if channel in self.channels}
if not keep_dictionary_structure:
selected = list(y.values())
if pulse_energy is not None:
@@ -319,9 +321,9 @@ class SelectRelevantLowResolution(TransformerMixin, BaseEstimator):
self.tof_start = self.estimate_prompt_peak(X)
X_tr = self.transform(X, keep_dictionary_structure=True)
self.mean = {ch: np.mean(X_tr[ch], axis=0, keepdims=True)
- for ch in X.keys()}
+ for ch in X_tr.keys()}
self.std = {ch: np.std(X_tr[ch], axis=0, keepdims=True)
- for ch in X.keys()}
+ for ch in X_tr.keys()}
return self
def debug_peak_finding(self, X: Dict[str, np.ndarray], filename: str):
@@ -574,7 +576,9 @@ class Model(TransformerMixin, BaseEstimator):
Set to None to perform no selection.
validation_size: Fraction (number between 0 and 1) of the data to take for
validation and systematic uncertainty estimate.
- model_type: Which model to use. "bnn" for a BNN, "ridge" for Ridge and "ard" for ARD.
+ model_type: Which model to use. "bnn" for a BNN, "bnn_rvm" for a BNN with RVM, "ridge" for Ridge and "ard" for ARD.
+ n_peaks: Minimum numbr of peaks in the grating spectrometer.
+ n_bnn_epochs: Number of BNN epochs for training.
"""
def __init__(self,
@@ -586,11 +590,13 @@ class Model(TransformerMixin, BaseEstimator):
tof_start: Optional[int]=None,
delta_tof: Optional[int]=300,
validation_size: float=0.05,
- model_type: Literal["bnn", "ridge", "ard"]="ard",
+ model_type: Literal["bnn", "bnn_rvm", "ridge", "ard"]="ard",
+ n_peaks: int=0,
+ n_bnn_epochs: int=500,
):
self.high_res_sigma = high_res_sigma
# models
- self.x_select = SelectRelevantLowResolution(channels, tof_start, delta_tof, poly=(model_type not in ["bnn"]))
+ self.x_select = SelectRelevantLowResolution(channels, tof_start, delta_tof, poly=(model_type not in ["bnn", "bnn_rvm"]))
x_model_steps = list()
x_model_steps += [
('pca', PCA(n_pca_lr, whiten=True)),
@@ -605,7 +611,9 @@ class Model(TransformerMixin, BaseEstimator):
self.ood = {ch: UncorrelatedDeviation(sigma=5)
for ch in channels+['full']}
if model_type == "bnn":
- self.fit_model = BNNModel()
+ self.fit_model = BNNModel(n_epochs=n_bnn_epochs)
+ elif model_type == "bnn_rvm":
+ self.fit_model = BNNModel(n_epochs=n_bnn_epochs, rvm=True)
elif model_type == "ridge":
self.fit_model = MultiOutputRidgeWithStd(BayesianRidge(n_iter=300, tol=1e-8, verbose=True), n_jobs=8)
elif model_type == "ard":
@@ -625,9 +633,12 @@ class Model(TransformerMixin, BaseEstimator):
# size of the test subset
self.validation_size = validation_size
+ # minimum number of peaks
+ self.n_peaks = n_peaks
+
def n_pars(self) -> float:
"""Get number of parameters."""
- if self.model_type == "bnn":
+ if self.model_type in ("bnn", "bnn_rvm"):
return sum(p.numel() for p in self.fit_model.model.parameters())
return sum(len(estimator.coef_) + 1 for estimator in self.fit_model.estimators_)
@@ -735,7 +746,7 @@ class Model(TransformerMixin, BaseEstimator):
"""
print("Checking data quality in high-resolution data.")
peaks = self.count_peaks(high_res_data, high_res_photon_energy)
- filter_hr = (peaks > 3)
+ filter_hr = (peaks >= self.n_peaks)
print("Fitting PCA on low-resolution data.")
self.x_select.fit(low_res_data)
@@ -985,6 +996,7 @@ class Model(TransformerMixin, BaseEstimator):
unc=unc.reshape((B, P, -1)),
pca=pca_unc,
total_unc=total_unc.reshape((B, P, -1)),
+ nopca_unc=unc.reshape((B, P, -1)),
expected_pca=high_pca.reshape((B, P, -1)),
expected_pca_unc=high_pca_unc.reshape((B, P, -1)),
)
@@ -1015,7 +1027,7 @@ class Model(TransformerMixin, BaseEstimator):
joblib.dump([self.x_select,
self.x_model,
self.y_model,
- self.fit_model.state_dict() if self.model_type == "bnn" else self.fit_model,
+ self.fit_model.state_dict() if self.model_type in ("bnn", "bnn_rvm") else self.fit_model,
self.channel_pca,
#self.channel_fit_model
DataHolder(dict(
@@ -1073,7 +1085,9 @@ class Model(TransformerMixin, BaseEstimator):
obj.x_model = x_model
obj.y_model = y_model
if obj.model_type == "bnn":
- obj.fit_model = BNNModel(state_dict=fit_model)
+ obj.fit_model = BNNModel(state_dict=fit_model, rvm=False)
+ elif obj.model_type == "bnn_rvm":
+ obj.fit_model = BNNModel(state_dict=fit_model, rvm=True)
else:
obj.fit_model = fit_model
obj.channel_pca = channel_pca
diff --git a/pes_to_spec/test/channel_sensitivity.py b/pes_to_spec/test/channel_sensitivity.py
new file mode 100755
index 0000000000000000000000000000000000000000..2ff437238966b8a8a2f2665245bb2a956841a16d
--- /dev/null
+++ b/pes_to_spec/test/channel_sensitivity.py
@@ -0,0 +1,300 @@
+#!/usr/bin/env python
+
+import sys
+sys.path.append('.')
+sys.path.append('..')
+
+import os
+import argparse
+
+import numpy as np
+from extra_data import open_run, by_id, RunDirectory
+from pes_to_spec.model import Model, matching_ids
+
+# for helper plots
+from pes_to_spec.model import SelectRelevantLowResolution
+from sklearn.decomposition import PCA
+
+from itertools import product
+
+import pandas as pd
+from copy import deepcopy
+import scipy
+from scipy.signal import fftconvolve
+
+from typing import Dict, Optional
+
+from time import time_ns
+import pandas as pd
+
+def get_gas(run, tids):
+ gas_sources = [
+ "SA3_XTD10_PES/DCTRL/V30300S_NITROGEN",
+ "SA3_XTD10_PES/DCTRL/V30310S_NEON",
+ "SA3_XTD10_PES/DCTRL/V30320S_KRYPTON",
+ "SA3_XTD10_PES/DCTRL/V30330S_XENON",
+ ]
+ gas_active = list()
+ for gas in gas_sources:
+ # check if this gas source is interlocked
+ if gas in run.all_sources and run[gas, "interlock.AActionState.value"].ndarray().sum() == 0:
+ # it is not, so this gas was used
+ gas_active += [gas.split("/")[-1].split("_")[-1]]
+ gas = "_".join(gas_active)
+ return gas
+
+
+def save_result(filename: str,
+ spec_pred: Dict[str, np.ndarray],
+ spec_smooth: np.ndarray,
+ spec_raw_pe: np.ndarray,
+ intensity: float,
+ #spec_raw_int: Optional[np.ndarray]=None,
+ pes: Optional[np.ndarray]=None,
+ pes_to_show: Optional[str]="",
+ first: Optional[int]=None,
+ last: Optional[int]=None,
+ ):
+ """
+ Plot result with uncertainty band.
+
+ Args:
+ filename: Output file name.
+ spec_pred: Predicted result with uncertainty bands in a dictionary.
+ 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.
+ intensity: The XGM intensity in uJ.
+
+ """
+ unc_stat = spec_pred["unc"]
+ unc_pca = spec_pred["pca"]
+ unc = np.sqrt(unc_stat**2 + unc_pca**2)
+ df = pd.DataFrame(dict(energy=spec_raw_pe,
+ spec=spec_smooth,
+ prediction=spec_pred["expected"],
+ unc=unc,
+ unc_pca=unc_pca,
+ unc_stat=unc_stat,
+ beam_intensity=intensity*1e-3*np.ones_like(spec_raw_pe),
+ deconvolved=spec_pred["deconvolved"]
+ ))
+ df.to_csv(filename)
+ if pes is not None:
+ pes_data = deepcopy(pes)
+ pes_data['bin'] = np.arange(len(pes['channel_1_D']))
+ pes_data['first'] = first*np.ones_like(pes_data['bin'])
+ pes_data['last'] = last*np.ones_like(pes_data['bin'])
+ df = pd.DataFrame(pes_data)
+ df.to_csv(filename.replace('.pdf', '_pes.csv'))
+
+def save_pes_result(filename: str,
+ pes: Optional[np.ndarray]=None,
+ first: Optional[int]=None,
+ last: Optional[int]=None,
+ ):
+ """
+ Plot result with uncertainty band.
+
+ Args:
+ filename: Output file name.
+ spec_pred: Predicted result with uncertainty bands in a dictionary.
+ 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.
+ intensity: The XGM intensity in uJ.
+
+ """
+ pes_data = deepcopy(pes)
+ pes_data['bin'] = np.arange(len(pes['channel_1_D']))
+ pes_data['first'] = first*np.ones_like(pes_data['bin'])
+ pes_data['last'] = last*np.ones_like(pes_data['bin'])
+ df = pd.DataFrame(pes_data)
+ df.to_csv(filename)
+
+def main():
+ """
+ Main entry point. Reads some data, trains and predicts.
+ """
+ parser = argparse.ArgumentParser(prog="offline_analysis", description="Test pes2spec doing an offline analysis of the data.")
+ parser.add_argument('-p', '--proposal', type=int, metavar='INT', help='Proposal number', default=2828)
+ parser.add_argument('-r', '--run', type=str, metavar='INT,INT,...', help='Run numbers, comma-separated.', default=206)
+ parser.add_argument('-t', '--test-run', type=int, metavar='INT', help='Run to test', default=None)
+ parser.add_argument('-d', '--directory', type=str, metavar='DIRECTORY', default=".", help='Where to save the results.')
+ parser.add_argument('-S', '--spec', type=str, metavar='NAME', default="SA3_XTD10_SPECT/MDL/SPECTROMETER_SCS_NAVITAR:output", help='SPEC name')
+ parser.add_argument('-P', '--pes', type=str, metavar='NAME', default="SA3_XTD10_PES/ADC/1:network", help='PES name')
+ parser.add_argument('-X', '--xgm', type=str, metavar='NAME', default="SA3_XTD10_XGM/XGM/DOOCS:output", help='XGM name')
+ parser.add_argument('-o', '--offset', type=int, metavar='INT', default=0, help='Train ID offset')
+ parser.add_argument('-c', '--xgm_cut', type=float, metavar='INTENSITY', default=0, help='XGM intensity threshold in uJ.')
+ parser.add_argument('-T', '--model-type', type=str, metavar='TYPE', default="ard", choices=["bnn", "bnn_rvm", "ridge", "ard"], help='Which model type to use.')
+ parser.add_argument('-w', '--weight', action="store_true", default=False, help='Whether to reweight data as a function of the pulse energy to make it invariant to that.')
+
+ args = parser.parse_args()
+
+ print("Opening run ...")
+ runs = args.run.split(',')
+ runs = [int(r) for r in runs]
+ # get run
+ run_list = [open_run(proposal=args.proposal, run=r) for r in runs]
+ run = run_list[0]
+ if len(run_list) > 1:
+ run = run.union(*run_list[1:])
+
+ run_test = run
+ other_run_test = False
+ if "test_run" in args and args.test_run is not None:
+ other_run_test = True
+ run_test = open_run(proposal=args.proposal, run=args.test_run)
+
+ spec_offset = args.offset
+ spec_name = args.spec
+ pes_name = args.pes
+ xgm_name = args.xgm
+
+ pes_tid = run[pes_name, "digitizers.trainId"].ndarray()
+ xgm_tid = run[xgm_name, "data.trainId"].ndarray()
+
+ spec_tid = spec_offset + run[spec_name, "data.trainId"].ndarray()
+ # these are the train ID intersection
+ # this could have been done by a select call in the RunDirectory, but it would not correct for the spec_offset
+ tids = matching_ids(spec_tid, pes_tid, xgm_tid)
+
+ # read the spec photon energy and intensity
+ spec_raw_pe = run[spec_name, "data.photonEnergy"].select_trains(by_id[tids - spec_offset]).ndarray()
+ spec_raw_int = run[spec_name, "data.intensityDistribution"].select_trains(by_id[tids - spec_offset]).ndarray()
+
+
+ # reserve part of it for the test stage
+ train_tids = tids[:-10]
+ if other_run_test:
+ pes_tidt = run_test[pes_name, "digitizers.trainId"].ndarray()
+ xgm_tidt = run_test[xgm_name, "data.trainId"].ndarray()
+ spec_tidt = run_test[spec_name, "data.trainId"].ndarray()
+ test_tids = matching_ids(spec_tidt, pes_tidt, xgm_tidt)
+ else:
+ test_tids = tids
+ print(f"Number of train IDs: {len(train_tids)}")
+ print(f"Number of test IDs: {len(test_tids)}")
+
+ # read the PES data for each channel
+ channels = [f"channel_{i}_{l}"
+ for i, l in product([1,3,4], ["A", "B", "C", "D"])]
+ pes_raw = {ch: run[pes_name, f"digitizers.{ch}.raw.samples"].select_trains(by_id[tids]).ndarray()
+ for ch in channels}
+ pes_raw_t = {ch: run_test[pes_name, f"digitizers.{ch}.raw.samples"].select_trains(by_id[test_tids]).ndarray()
+ for ch in channels}
+
+ # select test SPEC data
+ spec_raw_pe_t = run_test[spec_name, "data.photonEnergy"].select_trains(by_id[test_tids - spec_offset]).ndarray()
+ spec_raw_int_t = run_test[spec_name, "data.intensityDistribution"].select_trains(by_id[test_tids - spec_offset]).ndarray()
+
+ print("Data in memory.")
+
+ # read the XGM information
+ #xgm_pressure = run['SA3_XTD10_XGM/XGM/DOOCS', "pressure.pressureFiltered.value"].select_trains(by_id[tids]).ndarray()
+ #xgm_pe = run['SA3_XTD10_XGM/XGM/DOOCS:output', "data.intensitySa3TD"].select_trains(by_id[tids]).ndarray()
+ #retvol_raw = run["SA3_XTD10_PES/MDL/DAQ_MPOD", "u212.value"].select_trains(by_id[tids]).ndarray()
+ #retvol_raw_timestamp = run["SA3_XTD10_PES/MDL/DAQ_MPOD", "u212.timestamp"].select_trains(by_id[tids]).ndarray()
+
+ xgm_flux = run['SA3_XTD10_XGM/XGM/DOOCS:output', "data.intensitySa3TD"].select_trains(by_id[tids]).ndarray()[:, 0][:, np.newaxis]
+ xgm_flux_t = run_test['SA3_XTD10_XGM/XGM/DOOCS:output', "data.intensitySa3TD"].select_trains(by_id[test_tids]).ndarray()[:, 0][:, np.newaxis]
+
+ print(f"Intensity in training: {np.mean(xgm_flux):.2e} +/- {np.std(xgm_flux):.2e}")
+ print(f"Intensity in testing: {np.mean(xgm_flux_t):.2e} +/- {np.std(xgm_flux_t):.2e}")
+
+ pressure = run["SA3_XTD10_PES/GAUGE/G30310F", "value"].select_trains(by_id[tids]).ndarray()
+ pressure_t = run_test["SA3_XTD10_PES/GAUGE/G30310F", "value"].select_trains(by_id[test_tids]).ndarray()
+ print(f"Pressure in training: {np.mean(pressure):.2e} +/- {np.std(pressure):.2e}")
+ print(f"Pressure in testing: {np.mean(pressure_t):.2e} +/- {np.std(pressure_t):.2e}")
+
+ voltage = run["SA3_XTD10_PES/MDL/DAQ_MPOD", "u212.value"].select_trains(by_id[tids]).ndarray()
+ voltage_t = run_test["SA3_XTD10_PES/MDL/DAQ_MPOD", "u212.value"].select_trains(by_id[test_tids]).ndarray()
+ print(f"Voltage in training: {np.mean(voltage):.2f} +/- {np.std(voltage):.2f}")
+ print(f"Voltage in testing: {np.mean(voltage_t):.2f} +/- {np.std(voltage_t):.2f}")
+
+ gas = get_gas(run, tids)
+ gas_t = get_gas(run_test, test_tids)
+ print(f"Gas in training: {gas}")
+ print(f"Gas in testing: {gas_t}")
+
+ t = list()
+ t_names = list()
+ t_nch = list()
+
+ train_idx = np.isin(tids, train_tids) & (xgm_flux[:,0] > args.xgm_cut)
+
+ # we just need this for training and we need to avoid copying it, which blows up the memoray usage
+ for k in pes_raw.keys():
+ pes_raw[k] = pes_raw[k][train_idx]
+
+
+ nch_axis = np.arange(1, len(channels)+1)
+ resolution = list()
+ rmse = list()
+ delta_rmse = list()
+ chi2_prepca = list()
+ unc = list()
+ delta_unc = list()
+ for nch in nch_axis:
+ model = Model(channels=channels[:nch], model_type=args.model_type)
+
+ print(f"Fitting using {nch} channels")
+ start = time_ns()
+ model.uniformize(xgm_flux[train_idx])
+ model.fit(pes_raw,
+ spec_raw_int[train_idx],
+ spec_raw_pe[train_idx],
+ pulse_energy=xgm_flux[train_idx],
+ )
+ t += [time_ns() - start]
+ t_names += ["Fit"]
+ t_nch += [nch]
+
+ resolution += [model.resolution]
+
+ # transfer function
+ print(f"Resolution: {model.resolution:.2f} eV")
+
+ # test
+ print("Predict")
+ start = time_ns()
+ spec_pred = model.predict(pes_raw_t, pulse_energy=xgm_flux_t)
+ spec_smooth = model.preprocess_high_res(spec_raw_int_t)
+ t += [time_ns() - start]
+ t_names += ["Predict"]
+ t_nch += [nch]
+
+ spec_smooth_pca = model.y_model['pca'].transform(spec_smooth)
+ unc2 = spec_pred["expected_pca_unc"]**2
+ pca_var = (spec_pred["expected_pca"].std(axis=0, keepdims=True)**2).reshape(1, 1, -1)
+ ndof_prepca = float(spec_smooth_pca.shape[-1])
+ print("Expected pca std:", pca_var)
+ chi2_prepca += [np.mean(np.sum((spec_smooth_pca[:, np.newaxis, :] - spec_pred["expected_pca"])**2/unc2, axis=(-1, -2)))/ndof_prepca]
+ rmse += [np.mean(np.sqrt(np.mean((spec_smooth[:, np.newaxis, :] - spec_pred["expected"])**2, axis=(-1, -2))))]
+ delta_rmse += [np.std(np.sqrt(np.mean((spec_smooth[:, np.newaxis, :] - spec_pred["expected"])**2, axis=(-1, -2))))]
+ unc += [np.mean(np.mean(spec_pred["total_unc"], axis=(-1, -2)))]
+ delta_unc += [np.std(np.mean(spec_pred["total_unc"], axis=(-1, -2)))]
+
+ df = pd.DataFrame(dict(number_channels=nch_axis,
+ resolution=resolution,
+ rmse=rmse,
+ delta_rmse=delta_rmse,
+ chi2_prepca=chi2_prepca,
+ unc=unc,
+ delta_unc=delta_unc,
+ ))
+ df.to_csv(os.path.join(args.directory, "number_channel_effect.csv"))
+
+
+ print("Time taken in ms")
+ df_time = pd.DataFrame(data=dict(time=t, name=t_names, nch=t_nch))
+ df_time.time *= 1e-6
+ df_time.to_csv(os.path.join(args.directory, "number_channel_time.csv"))
+
+if __name__ == '__main__':
+ main()
+
diff --git a/pes_to_spec/test/offline_analysis.py b/pes_to_spec/test/offline_analysis.py
index a9ba79450dffcacbcc0b93b19eec85337c741174..3911d0f630ed8f4239d528be1ff1bc288ea7ceae 100755
--- a/pes_to_spec/test/offline_analysis.py
+++ b/pes_to_spec/test/offline_analysis.py
@@ -11,17 +11,14 @@ import numpy as np
from extra_data import open_run, by_id, RunDirectory
from pes_to_spec.model import Model, matching_ids
-from itertools import product
+# for helper plots
+from pes_to_spec.model import SelectRelevantLowResolution
+from sklearn.decomposition import PCA
-import matplotlib
-matplotlib.use('Agg')
+from itertools import product
import pandas as pd
from copy import deepcopy
-import matplotlib.pyplot as plt
-from matplotlib.gridspec import GridSpec
-from mpl_toolkits.axes_grid1.inset_locator import InsetPosition
-import seaborn as sns
import scipy
from scipy.signal import fftconvolve
@@ -30,18 +27,6 @@ from typing import Dict, Optional
from time import time_ns
import pandas as pd
-SMALL_SIZE = 12
-MEDIUM_SIZE = 18
-BIGGER_SIZE = 22
-
-plt.rc('font', size=BIGGER_SIZE) # controls default text sizes
-plt.rc('axes', titlesize=BIGGER_SIZE) # fontsize of the axes title
-plt.rc('axes', labelsize=BIGGER_SIZE) # fontsize of the x and y labels
-plt.rc('xtick', labelsize=BIGGER_SIZE) # fontsize of the tick labels
-plt.rc('ytick', labelsize=BIGGER_SIZE) # fontsize of the tick labels
-plt.rc('legend', fontsize=MEDIUM_SIZE) # legend fontsize
-plt.rc('figure', titlesize=BIGGER_SIZE) # fontsize of the figure title
-
def get_gas(run, tids):
gas_sources = [
"SA3_XTD10_PES/DCTRL/V30300S_NITROGEN",
@@ -59,27 +44,7 @@ def get_gas(run, tids):
return gas
-def plot_pes(filename: str, pes_raw_int: np.ndarray, first: int, last: int):
- """
- Plot low-resolution spectrum.
-
- Args:
- filename: Output file name.
- pes_raw_int: Low-resolution spectrum.
-
- """
- fig = plt.figure(figsize=(16, 8))
- gs = GridSpec(1, 1)
- ax = fig.add_subplot(gs[0, 0])
- ax.plot(np.arange(first, last), pes_raw_int, c='b', lw=3, label="Low-resolution measurement")
- ax.legend()
- ax.set(title=f"",
- xlabel="ToF index",
- ylabel="Intensity")
- fig.savefig(filename)
- plt.close(fig)
-
-def plot_result(filename: str,
+def save_result(filename: str,
spec_pred: Dict[str, np.ndarray],
spec_smooth: np.ndarray,
spec_raw_pe: np.ndarray,
@@ -87,7 +52,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,
+ first: Optional[int]=None,
+ last: Optional[int]=None,
):
"""
Plot result with uncertainty band.
@@ -100,7 +66,6 @@ def plot_result(filename: str,
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.
intensity: The XGM intensity in uJ.
"""
@@ -111,64 +76,45 @@ def plot_result(filename: str,
spec=spec_smooth,
prediction=spec_pred["expected"],
unc=unc,
- beam_intensity=intensity*1e-3*np.ones_like(spec_raw_pe)
+ unc_pca=unc_pca,
+ unc_stat=unc_stat,
+ beam_intensity=intensity*1e-3*np.ones_like(spec_raw_pe),
+ deconvolved=spec_pred["deconvolved"]
))
- df.to_csv(filename.replace('.png', '.csv'))
+ df.to_csv(filename)
+ if pes is not None:
+ pes_data = deepcopy(pes)
+ pes_data['bin'] = np.arange(len(pes['channel_1_D']))
+ pes_data['first'] = first*np.ones_like(pes_data['bin'])
+ pes_data['last'] = last*np.ones_like(pes_data['bin'])
+ df = pd.DataFrame(pes_data)
+ df.to_csv(filename.replace('.pdf', '_pes.csv'))
+
+def save_pes_result(filename: str,
+ pes: Optional[np.ndarray]=None,
+ first: Optional[int]=None,
+ last: Optional[int]=None,
+ ):
+ """
+ Plot result with uncertainty band.
+
+ Args:
+ filename: Output file name.
+ spec_pred: Predicted result with uncertainty bands in a dictionary.
+ 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.
+ intensity: The XGM intensity in uJ.
+
+ """
pes_data = deepcopy(pes)
pes_data['bin'] = np.arange(len(pes['channel_1_D']))
+ pes_data['first'] = first*np.ones_like(pes_data['bin'])
+ pes_data['last'] = last*np.ones_like(pes_data['bin'])
df = pd.DataFrame(pes_data)
- df.to_csv(filename.replace('.png', '_pes.csv'))
-
- fig = plt.figure(figsize=(12, 8))
- gs = GridSpec(1, 1)
- ax = fig.add_subplot(gs[0, 0])
- 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")
- #if wiener is not None:
- # deconvolved = fftconvolve(spec_pred["expected"], wiener, mode="same")
- #ax.plot(spec_raw_pe, spec_pred["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"Beam intensity: {intensity*1e-3:.1f} mJ", loc="left")
- ax.spines['top'].set_visible(False)
- ax.spines['right'].set_visible(False)
- ax.set(
- xlabel="Photon energy [eV]",
- ylabel="Intensity",
- ylim=(0, 1.3*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])
- ip = InsetPosition(ax, [0.72,0.7,0.35,0.4])
- ax2.set_axes_locator(ip)
- if pes_to_show == "sum":
- pes_plot = sum([pes[k][pes_bin] for k in pes.keys()])
- pes_label = r"$\sum$ PES channels"
- else:
- pes_plot = pes[pes_to_show][pes_bin]
- pes_label = pes_to_show
- ax2.plot(pes_bin, pes_plot, c='black', lw=3)
- ax2.set(title=f"Low-resolution example data",
- xlabel="Bin",
- ylabel=pes_label,
- ylim=(0, None),
- #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)
+ df.to_csv(filename)
def main():
"""
@@ -180,12 +126,12 @@ def main():
parser.add_argument('-t', '--test-run', type=int, metavar='INT', help='Run to test', default=None)
parser.add_argument('-m', '--model', type=str, metavar='FILENAME', default="", help='Model to load. If given, do not train a model and just do inference with this one.')
parser.add_argument('-d', '--directory', type=str, metavar='DIRECTORY', default=".", help='Where to save the results.')
- parser.add_argument('-S', '--spec', type=str, metavar='NAME', default="SA3_XTD10_SPECT/MDL/SPECTROMETER_SQS_NAVITAR:output", help='SPEC name')
+ parser.add_argument('-S', '--spec', type=str, metavar='NAME', default="SA3_XTD10_SPECT/MDL/SPECTROMETER_SCS_NAVITAR:output", help='SPEC name')
parser.add_argument('-P', '--pes', type=str, metavar='NAME', default="SA3_XTD10_PES/ADC/1:network", help='PES name')
parser.add_argument('-X', '--xgm', type=str, metavar='NAME', default="SA3_XTD10_XGM/XGM/DOOCS:output", help='XGM name')
parser.add_argument('-o', '--offset', type=int, metavar='INT', default=0, help='Train ID offset')
parser.add_argument('-c', '--xgm_cut', type=float, metavar='INTENSITY', default=0, help='XGM intensity threshold in uJ.')
- parser.add_argument('-T', '--model-type', type=str, metavar='TYPE', default="ard", choices=["bnn", "ridge", "ard"], help='Which model type to use.')
+ parser.add_argument('-T', '--model-type', type=str, metavar='TYPE', default="ard", choices=["bnn", "bnn_rvm", "ridge", "ard"], help='Which model type to use.')
parser.add_argument('-w', '--weight', action="store_true", default=False, help='Whether to reweight data as a function of the pulse energy to make it invariant to that.')
args = parser.parse_args()
@@ -253,8 +199,10 @@ def main():
print(f"Number of test IDs: {len(test_tids)}")
# read the PES data for each channel
+ #channels = [f"channel_{i}_{l}"
+ # for i, l in product(range(1, 5), ["A", "B", "C", "D"])]
channels = [f"channel_{i}_{l}"
- for i, l in product(range(1, 5), ["A", "B", "C", "D"])]
+ for i, l in product([1,3,4], ["A", "B", "C", "D"])]
pes_raw = {ch: run[pes_name, f"digitizers.{ch}.raw.samples"].select_trains(by_id[tids]).ndarray()
for ch in channels}
pes_raw_t = {ch: run_test[pes_name, f"digitizers.{ch}.raw.samples"].select_trains(by_id[test_tids]).ndarray()
@@ -296,14 +244,14 @@ def main():
t = list()
t_names = list()
- model = Model(model_type=args.model_type)
+ model = Model(channels=channels, model_type=args.model_type)
train_idx = np.isin(tids, train_tids) & (xgm_flux[:,0] > args.xgm_cut)
# we just need this for training and we need to avoid copying it, which blows up the memoray usage
for k in pes_raw.keys():
pes_raw[k] = pes_raw[k][train_idx]
- model.debug_peak_finding(pes_raw, os.path.join(args.directory, "test_peak_finding.png"))
+ model.debug_peak_finding(pes_raw, os.path.join(args.directory, "test_peak_finding.pdf"))
if len(args.model) == 0:
print("Fitting")
start = time_ns()
@@ -334,44 +282,36 @@ def main():
t += [time_ns() - start]
t_names += ["Load"]
+ # save PCA information
+ pes_raw_select = model.x_select.transform(pes_raw, pulse_energy=xgm_flux[train_idx])
+ ch = channels[0]
+ idx = 0
+ first = model.x_select.tof_start - model.x_select.delta_tof
+ last = model.x_select.tof_start + model.x_select.delta_tof
+ B, P, _ = pes_raw_select.shape
+ pes_raw_select = pes_raw_select.reshape((B*P, -1))
+ pca = PCA(None, whiten=True)
+ pca.fit(pes_raw_select)
+ df = pd.DataFrame(dict(variance_ratio=pca.explained_variance_ratio_,
+ n_comp=600*np.ones_like(pca.explained_variance_ratio_),
+ ))
+ df.to_csv(os.path.join(args.directory, "pca_pes.csv"))
+
+ pca_spec = PCA(None, whiten=True)
+ pca_spec.fit(spec_raw_int[train_idx])
+ df = pd.DataFrame(dict(variance_ratio=pca_spec.explained_variance_ratio_,
+ n_comp=20*np.ones_like(pca_spec.explained_variance_ratio_),
+ ))
+ df.to_csv(os.path.join(args.directory, "pca_spec.csv"))
+
# transfer function
- fig = plt.figure(figsize=(12, 8))
- gs = GridSpec(1, 1)
- ax = fig.add_subplot(gs[0, 0])
- plt.plot(model.wiener_energy, np.absolute(model.impulse_response))
- ax.set(title=f"",
- xlabel=r"Energy [eV]",
- ylabel="Response [a.u.]",
- yscale='log',
- )
- fig.savefig(os.path.join(args.directory, "impulse.png"))
- plt.close(fig)
print(f"Resolution: {model.resolution:.2f} eV")
- # 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.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)
+ df = pd.DataFrame(dict(wiener_energy=model.wiener_energy,
+ wiener_filter=model.wiener_filter,
+ impulse=model.impulse_response,
+ resolution=model.resolution*np.ones_like(model.wiener_energy)))
+ df.to_csv(os.path.join(args.directory, "model.csv"))
print("Check consistency")
start = time_ns()
@@ -410,206 +350,52 @@ def main():
chi2 = np.sum((spec_smooth[:, np.newaxis, :] - spec_pred["expected"])**2/(spec_pred["total_unc"]**2), axis=(-1, -2))
ndof = spec_smooth.shape[1]
print(f"Chi2 after PCA: {np.mean(chi2):.2f}, ndof: {ndof}, chi2/ndof: {np.mean(chi2/ndof):.2f}")
- fig = plt.figure(figsize=(12, 8))
- gs = GridSpec(1, 1)
- ax = fig.add_subplot(gs[0, 0])
- ax.scatter(chi2/ndof, xgm_flux_t[:,0], c='r', s=20)
- ax.set(title=f"", #avg(stat unc) = {unc_stat}, avg(pca unc) = {unc_pca}",
- xlabel=r"$\chi^2/$ndof",
- ylabel="XGM intensity [uJ]",
- xlim=(0, 5),
- )
- 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)
- ax2.scatter(chi2/ndof, xgm_flux_t[:,0], c='r', s=30)
- #ax2.scatter(chi2/ndof, np.sum(spec_pred["expected"], axis=1)*de, c='b', s=30)
- #ax2.scatter(chi2/ndof, np.sum(spec_raw_int, axis=1)*de, c='g', s=30)
- ax2.set(title="",
- xlabel=r"$\chi^2/$ndof",
- ylabel=f"XGM intensity [uJ]",
- )
- 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(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])
- sns.histplot(x=chi2/ndof, kde=True, linewidth=3, ax=ax)
- ax.set(title=f"",
- xlabel=r"$\chi^2/$ndof",
- ylabel="Counts [a.u.]",
- xlim=(0, 5),
- )
- #ax.text(0.90, 0.95, fr"$\mu = ${np.mean(chi2/ndof):.2f}",
- # verticalalignment='top', horizontalalignment='right',
- # transform=ax.transAxes,
- # color='black', fontsize=15)
- #ax.text(0.90, 0.90, fr"$\sigma = ${np.std(chi2/ndof):.2f}",
- # verticalalignment='top', horizontalalignment='right',
- # transform=ax.transAxes,
- # color='black', fontsize=15)
- fig.savefig(os.path.join(args.directory, "chi2.png"))
- plt.close(fig)
spec_smooth_pca = model.y_model['pca'].transform(spec_smooth)
- chi2_prepca = np.sum((spec_smooth_pca[:, np.newaxis, :] - spec_pred["expected_pca"])**2/(spec_pred["expected_pca_unc"]**2), axis=(-1, -2))
+ unc2 = spec_pred["expected_pca_unc"]**2
+ pca_var = (spec_pred["expected_pca"].std(axis=0, keepdims=True)**2).reshape(1, 1, -1)
+ print("Expected pca std:", pca_var)
+ chi2_prepca = np.sum((spec_smooth_pca[:, np.newaxis, :] - spec_pred["expected_pca"])**2/unc2, axis=(-1, -2))
+
ndof_prepca = float(spec_smooth_pca.shape[-1])
print(f"Chi2 before PCA: {np.mean(chi2_prepca):.2f}, ndof: {ndof_prepca}, chi2/ndof: {np.mean(chi2_prepca/ndof_prepca):.2f} +/- {np.std(chi2_prepca/ndof_prepca):.2f}")
- fig = plt.figure(figsize=(12, 8))
- gs = GridSpec(1, 1)
- ax = fig.add_subplot(gs[0, 0])
- sns.histplot(x=chi2_prepca/ndof_prepca, kde=True, linewidth=3, ax=ax)
- ax.set(title=f"",
- xlabel=r"$\chi^2/$ndof before undoing PCA",
- ylabel="Counts [a.u.]",
- xlim=(0, 5),
- )
- #ax.text(0.90, 0.95, fr"$\mu = ${np.mean(chi2/ndof):.2f}",
- # verticalalignment='top', horizontalalignment='right',
- # transform=ax.transAxes,
- # color='black', fontsize=15)
- #ax.text(0.90, 0.90, fr"$\sigma = ${np.std(chi2/ndof):.2f}",
- # verticalalignment='top', horizontalalignment='right',
- # transform=ax.transAxes,
- # color='black', fontsize=15)
- fig.savefig(os.path.join(args.directory, "chi2_prepca.png"))
- plt.close(fig)
-
- fig = plt.figure(figsize=(12, 8))
- gs = GridSpec(1, 1)
- ax = fig.add_subplot(gs[0, 0])
- ax.scatter(chi2_prepca/ndof_prepca, xgm_flux_t[:,0], c='r', s=20)
- ax.set(title=f"",
- xlabel=r"$\chi^2/$ndof before undoing PCA",
- ylabel="XGM intensity [uJ]",
- xlim=(0, 5),
- ylim=(0, np.mean(xgm_flux_t) + 3*np.std(xgm_flux_t))
- )
- fig.savefig(os.path.join(args.directory, "intensity_vs_chi2_prepca.png"))
- plt.close(fig)
res_prepca = np.sum((spec_smooth_pca[:, np.newaxis, :] - spec_pred["expected_pca"])/spec_pred["expected_pca_unc"], axis=1)
- n_plots = res_prepca.shape[1]//10
- fig = plt.figure(figsize=(8*n_plots, 8))
- gs = GridSpec(1, n_plots)
- for i_plot in range(n_plots):
- ax = fig.add_subplot(gs[0, i_plot])
- sns.kdeplot(data={f"Dim. {k+1}": res_prepca[:, k] for k in range(i_plot*10, i_plot*10 + 10)},
- linewidth=3, ax=ax)
- ax.set(title=f"",
- xlabel=r"residue/uncertainty [a.u.]",
- ylabel="Counts [a.u.]",
- xlim=(-3, 3),
- )
- ax.legend(frameon=False)
- fig.savefig(os.path.join(args.directory, "res_prepca.png"))
- plt.close(fig)
-
- fig = plt.figure(figsize=(12, 8))
- gs = GridSpec(1, 1)
- ax = fig.add_subplot(gs[0, 0])
- sns.histplot(x=xgm_flux_t[:,0], kde=True, linewidth=3, ax=ax)
- ax.set(title=f"",
- xlabel="XGM intensity [uJ]",
- ylabel="Counts [a.u.]",
- )
- #ax.text(0.90, 0.95, fr"$\mu = ${np.mean(xgm_flux_t[:,0]):.2f}",
- # verticalalignment='top', horizontalalignment='right',
- # transform=ax.transAxes,
- # color='black', fontsize=15)
- #ax.text(0.90, 0.90, fr"$\sigma = ${np.std(xgm_flux_t[:,0]):.2f}",
- # verticalalignment='top', horizontalalignment='right',
- # transform=ax.transAxes,
- # color='black', fontsize=15)
- plt.tight_layout()
- fig.savefig(os.path.join(args.directory, "intensity.png"))
- plt.close(fig)
# rmse
rmse = np.sqrt(np.mean((spec_smooth[:, np.newaxis, :] - spec_pred["expected"])**2, axis=(-1, -2)))
- fig = plt.figure(figsize=(12, 8))
- gs = GridSpec(1, 1)
- ax = fig.add_subplot(gs[0, 0])
- ax.scatter(rmse, xgm_flux_t[:,0], c='r', s=30)
- ax = plt.gca()
- ax.set(title=f"",
- xlabel=r"Root-mean-squared error",
- ylabel="XGM intensity [uJ]",
- )
- 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])
- sns.histplot(x=rmse, kde=True, linewidth=3, ax=ax)
- ax.set(title=f"",
- xlabel="Root-mean-squared error",
- ylabel="Counts [a.u.]",
- )
- #ax.text(0.90, 0.95, fr"$\mu = ${np.mean(rmse):.2f}",
- # verticalalignment='top', horizontalalignment='right',
- # transform=ax.transAxes,
- # color='black', fontsize=15)
- #ax.text(0.90, 0.90, fr"$\sigma = ${np.std(rmse):.2f}",
- # verticalalignment='top', horizontalalignment='right',
- # transform=ax.transAxes,
- # 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)
- #ax = fig.add_subplot(gs[0, 0])
- #sns.regplot(x=np.sum(spec_raw_int_t, axis=1)*de, y=xgm_flux_t[:,0], color='r', robust=True, ax=ax)
- #ax.set(title=f"",
- # xlabel="SPEC (raw) integral",
- # ylabel="XGM Intensity [uJ]",
- # )
- #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)
- #ax = fig.add_subplot(gs[0, 0])
- #sns.regplot(x=np.sum(spec_raw_int_t, axis=-1)*de, y=np.sum(spec_pred["expected"], axis=(-1, -2))*de, color='r', robust=True, ax=ax)
- #ax.set(title=f"",
- # xlabel="SPEC (raw) integral",
- # ylabel="Predicted integral",
- # )
- #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])
- #sns.regplot(x=np.sum(spec_pred["expected"], axis=(-1, -2))*de, y=xgm_flux_t[:,0], color='r', robust=True, ax=ax)
- #ax.set(title=f"",
- # xlabel="Predicted integral",
- # ylabel="XGM intensity [uJ]",
- # )
- #fig.savefig(os.path.join(args.directory, "xgm_vs_expected.png"))
- #plt.close(fig)
+ nopca_unc = np.sqrt(np.mean(spec_pred["nopca_unc"]**2, axis=(-1, -2)))
+ total_unc = np.sqrt(np.mean(spec_pred["total_unc"]**2, axis=(-1, -2)))
+ median_unc = np.median(spec_pred["total_unc"], axis=(-1, -2))
+
+ q = dict(chi2_prepca=chi2_prepca,
+ ndof=spec_smooth_pca.shape[-1]*np.ones_like(chi2_prepca),
+ xgm_flux_t=xgm_flux_t[:,0],
+ rmse=rmse,
+ nopca_unc=nopca_unc,
+ total_unc=total_unc,
+ median_unc=median_unc,
+ root_mean_squared_pca_unc=np.sqrt((spec_pred["expected_pca_unc"][:, 0, :]**2).mean(axis=-1))
+ )
+ q.update({f'res_prepca_{k}': res_prepca[:, k]
+ for k in range(res_prepca.shape[1])
+ }
+ )
+ q.update({f'unc_prepca_{k}': spec_pred["expected_pca_unc"][:, 0, k]
+ for k in range(spec_pred["expected_pca_unc"].shape[-1])
+ }
+ )
+ df = pd.DataFrame(q)
+ df.to_csv(os.path.join(args.directory, "quality.csv"))
first, last = model.get_low_resolution_range()
- first = max(0, first+250)
- last = min(last, pes_raw_t["channel_1_D"].shape[1]-1)
- pes_to_show = 'sum'
# plot
high_int_idx = np.argsort(xgm_flux_t[:,0])
for q in [10, 25, 50, 75, 100]:
qi = int(len(high_int_idx)*(q/100.0))
for idx in high_int_idx[qi-10:qi]:
tid = test_tids[idx]
- plot_result(os.path.join(args.directory, f"test_q{q}_{tid}.png"),
+ save_result(os.path.join(args.directory, f"test_q{q}_{tid}.csv"),
{k: item[idx, 0, ...] if k != "pca"
else item[0, ...]
for k, item in spec_pred.items()},
@@ -617,14 +403,13 @@ def main():
spec_raw_pe_t[idx, :] if showSpec else None,
#spec_raw_int_t[idx, :] if showSpec else None,
intensity=xgm_flux_t[idx,0],
+ )
+ save_pes_result(os.path.join(args.directory, f"test_q{q}_{tid}_pes.csv"),
pes={k: -item[idx, :]
for k, item in pes_raw_t.items()},
- pes_to_show=pes_to_show,
- pes_bin=np.arange(first, last),
+ first=first,
+ last=last,
)
- #for ch in channels:
- # plot_pes(os.path.join(args.directory, f"test_pes_{tid}_{ch}.png"),
- # pes_raw_t[ch][idx, first:last], first, last)
if __name__ == '__main__':
main()
diff --git a/pes_to_spec/test/plot_channel_sensitivity.py b/pes_to_spec/test/plot_channel_sensitivity.py
new file mode 100755
index 0000000000000000000000000000000000000000..813b4ecb5f4785e47b79a5b710ad286bbb28e22e
--- /dev/null
+++ b/pes_to_spec/test/plot_channel_sensitivity.py
@@ -0,0 +1,93 @@
+#!/usr/bin/env python
+
+import os
+import re
+
+from typing import Optional, Tuple, Dict
+
+import matplotlib
+matplotlib.use('Agg')
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.gridspec import GridSpec
+import seaborn as sns
+
+SMALL_SIZE = 12
+MEDIUM_SIZE = 22
+BIGGER_SIZE = 26
+
+plt.rc('font', size=BIGGER_SIZE) # controls default text sizes
+plt.rc('axes', titlesize=BIGGER_SIZE) # fontsize of the axes title
+plt.rc('axes', labelsize=BIGGER_SIZE) # fontsize of the x and y labels
+plt.rc('xtick', labelsize=BIGGER_SIZE) # fontsize of the tick labels
+plt.rc('ytick', labelsize=BIGGER_SIZE) # fontsize of the tick labels
+plt.rc('legend', fontsize=MEDIUM_SIZE) # legend fontsize
+plt.rc('figure', titlesize=BIGGER_SIZE) # fontsize of the figure title
+
+def plot_resolution(df: pd.DataFrame, filename: str):
+ fig = plt.figure(figsize=(12, 8))
+ gs = GridSpec(1, 1)
+ ax = fig.add_subplot(gs[0, 0])
+ ax.plot(df.number_channels, df.resolution, c='b', lw=3)
+ ax.spines['top'].set_visible(False)
+ ax.spines['right'].set_visible(False)
+ ax.set(
+ xlabel="Number of channels",
+ ylabel="Average resolution [eV]",
+ )
+ plt.tight_layout()
+ fig.savefig(filename)
+ plt.close(fig)
+
+def plot_unc(df: pd.DataFrame, filename: str):
+ fig = plt.figure(figsize=(12, 8))
+ gs = GridSpec(1, 1)
+ ax = fig.add_subplot(gs[0, 0])
+ #ax.plot(df.number_channels, 2*df.unc, c='tab:blue', lw=3, alpha=0.7, label="Avg. 95% CL uncertainty band")
+ #ax.fill_between(df.number_channels, 2*df.unc - 2*df.delta_unc, 2*df.unc + 2*df.delta_unc, color='tab:blue', alpha=0.2)
+ ax.errorbar(df.number_channels, 2*df.unc, yerr=2*df.delta_unc, color='tab:blue', alpha=0.5, marker='o', markersize=20, lw=3, linestyle='none', label="95% CL uncertainty band")
+ ax.spines['top'].set_visible(False)
+ ax.spines['right'].set_visible(False)
+ ax.set(
+ xlabel="Number of channels",
+ ylabel="Grating spectrometer intensity [a.u.]",
+ )
+ #rax = ax.twinx()
+ rax = ax
+ #rax.plot(df.number_channels, df.rmse, c='tab:red', lw=3, alpha=0.7, label="Avg. root-mean-squared error")
+ #rax.fill_between(df.number_channels, df.rmse - df.delta_rmse, df.rmse + df.delta_rmse, color='tab:red', alpha=0.2)
+ rax.errorbar(df.number_channels, df.rmse, yerr=df.delta_rmse, color='tab:red', alpha=0.5, marker='^', markersize=20, lw=3, linestyle='none', label="Root-mean-squared error")
+ #rax.spines['right'].set_color('tab:red')
+ #rax.spines['top'].set_visible(False)
+ #rax.tick_params(axis='y', colors='tab:red')
+ #rax.set_ylabel("Root-mean-squared error [a.u.]")
+ #rax.yaxis.label.set_color("tab:red")
+ ax.legend(frameon=False)
+ plt.tight_layout()
+ fig.savefig(filename)
+ plt.close(fig)
+
+def plot_rmse(df: pd.DataFrame, filename: str):
+ fig = plt.figure(figsize=(12, 8))
+ gs = GridSpec(1, 1)
+ ax = fig.add_subplot(gs[0, 0])
+ ax.plot(df.number_channels, df.rmse, c='b', lw=3)
+ ax.spines['top'].set_visible(False)
+ ax.spines['right'].set_visible(False)
+ ax.set(
+ xlabel="Number of channels",
+ ylabel="Root-mean-squared error [a.u.]",
+ )
+ plt.tight_layout()
+ fig.savefig(filename)
+ plt.close(fig)
+
+if __name__ == '__main__':
+ indir = 'p900331r69t70'
+ df = pd.read_csv(f'{indir}/number_channel_effect.csv')
+
+ plot_rmse(df, "rmse.pdf")
+ plot_unc(df, "unc.pdf")
+ plot_resolution(df, "resolution.pdf")
+
diff --git a/pes_to_spec/test/prepare_plots.py b/pes_to_spec/test/prepare_plots.py
new file mode 100755
index 0000000000000000000000000000000000000000..3f6ecf64f0b7199a1b69bc22b0293cb649e016a7
--- /dev/null
+++ b/pes_to_spec/test/prepare_plots.py
@@ -0,0 +1,442 @@
+#!/usr/bin/env python
+
+import os
+import re
+
+from typing import Optional, Tuple, Dict
+
+import matplotlib
+matplotlib.use('Agg')
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.gridspec import GridSpec
+import seaborn as sns
+
+SMALL_SIZE = 12
+MEDIUM_SIZE = 22
+BIGGER_SIZE = 26
+
+plt.rc('font', size=BIGGER_SIZE) # controls default text sizes
+plt.rc('axes', titlesize=BIGGER_SIZE) # fontsize of the axes title
+plt.rc('axes', labelsize=BIGGER_SIZE) # fontsize of the x and y labels
+plt.rc('xtick', labelsize=BIGGER_SIZE) # fontsize of the tick labels
+plt.rc('ytick', labelsize=BIGGER_SIZE) # fontsize of the tick labels
+plt.rc('legend', fontsize=MEDIUM_SIZE) # legend fontsize
+plt.rc('figure', titlesize=BIGGER_SIZE) # fontsize of the figure title
+
+def plot_final(df: pd.DataFrame, filename: str):
+ fig = plt.figure(figsize=(12, 8))
+ gs = GridSpec(1, 1)
+ ax = fig.add_subplot(gs[0, 0])
+ ax.plot(df.energy, df.spec, c='b', lw=3, label="Grating spectrometer")
+ ax.plot(df.energy, df.prediction, c='r', ls='--', lw=3, label="Prediction")
+ ax.fill_between(df.energy, df.prediction - 2*df.unc, df.prediction + 2*df.unc, facecolor='gold', alpha=0.5, label="95% unc. (total)")
+ ax.fill_between(df.energy, df.prediction - 2*df.unc_pca, df.prediction + 2*df.unc_pca, facecolor='magenta', alpha=0.5, label="95% unc. (PCA only)")
+ Y = np.amax(df.spec)
+ ax.legend(frameon=False, borderaxespad=0, loc='upper left')
+ ax.set_title(f"Beam intensity: {df.beam_intensity.iloc[0]:.1f} mJ", loc="left")
+ ax.spines['top'].set_visible(False)
+ ax.spines['right'].set_visible(False)
+ ax.set(
+ xlabel="Photon energy [eV]",
+ ylabel="Intensity [a.u.]",
+ ylim=(0, 1.3*Y))
+ plt.tight_layout()
+ fig.savefig(filename)
+ plt.close(fig)
+
+def plot_chi2(df: pd.DataFrame, filename: str):
+ fig = plt.figure(figsize=(12, 8))
+ gs = GridSpec(1, 1)
+ ax = fig.add_subplot(gs[0, 0])
+ sns.histplot(x=df.chi2_prepca/df.ndof.iloc[0], kde=True, linewidth=3, ax=ax)
+ ax.set(title=f"",
+ xlabel=r"$\chi^2/$ndof",
+ ylabel="Counts [a.u.]",
+ xlim=(0, 5),
+ )
+ fig.savefig(filename)
+ plt.close(fig)
+
+def plot_rmse(df: pd.DataFrame, filename: str):
+ fig = plt.figure(figsize=(12, 8))
+ gs = GridSpec(1, 1)
+ ax = fig.add_subplot(gs[0, 0])
+ sns.histplot(x=df.rmse, kde=True, linewidth=3, ax=ax)
+ ax.set(title=f"",
+ xlabel=r"Root-mean-square-error",
+ ylabel="Counts [a.u.]",
+ )
+ fig.savefig(filename)
+ plt.close(fig)
+
+def plot_residue(df: pd.DataFrame, filename: str):
+ cols = [k for k in df.columns if "res_prepca" in k]
+ df_res = df.loc[:, cols]
+ n_plots = len(df_res.columns)//10
+ fig = plt.figure(figsize=(8*n_plots, 8))
+ gs = GridSpec(1, n_plots)
+ for i_plot in range(n_plots):
+ ax = fig.add_subplot(gs[0, i_plot])
+ sns.kdeplot(data={f"Dim. {k+1}": df_res.loc[:, cols[k]] for k in range(i_plot*10, i_plot*10 + 10)},
+ linewidth=3, ax=ax)
+ ax.set(title=f"",
+ xlabel=r"residue/uncertainty [a.u.]",
+ ylabel="Counts [a.u.]",
+ xlim=(-3, 3),
+ )
+ ax.legend(frameon=False)
+ fig.savefig(filename)
+ plt.close(fig)
+
+def plot_residue_corr(df: pd.DataFrame, filename: str):
+ cols = [k for k in df.columns if "res_prepca" in k]
+ df_res = df.loc[:, cols]
+ df_res.columns = [re.match(r"res_prepca_([0-9]*)", k).groups()[0] for k in df_res.columns]
+ fig = plt.figure(figsize=(8, 8))
+ corr = df_res.corr()
+ mask = np.triu(np.ones_like(corr, dtype=bool))
+ cmap = sns.diverging_palette(230, 20, as_cmap=True)
+ sns.heatmap(corr, mask=mask, cmap=cmap, center=0,
+ square=True, linewidths=0.5, vmin=-1, vmax=1)
+ fig.savefig(filename)
+ plt.close(fig)
+
+def plot_chi2_intensity(df: pd.DataFrame, filename: str):
+ fig = plt.figure(figsize=(12, 8))
+ gs = GridSpec(1, 1)
+ ax = fig.add_subplot(gs[0, 0])
+ sns.kdeplot(x=df.chi2_prepca/df.ndof.iloc[0], y=df.xgm_flux_t*1e-3,
+ fill=True,
+ ax=ax)
+ sns.scatterplot(x=df.chi2_prepca/df.ndof.iloc[0], y=df.xgm_flux_t*1e-3,
+ s=5,
+ alpha=0.4,
+ c="tab:red",
+ #size=df.root_mean_squared_pca_unc,
+ #sizes=(20, 200),
+ ax=ax)
+ ax = plt.gca()
+ ax.set(title=f"",
+ xlabel=r"$\chi^2/$ndof",
+ ylabel="Beam intensity [mJ]",
+ xlim=(0, 5),
+ ylim=(0, df.xgm_flux_t.mean()*1e-3 + 3*df.xgm_flux_t.std()*1e-3)
+ )
+ ax.spines['top'].set_visible(False)
+ ax.spines['right'].set_visible(False)
+ plt.tight_layout()
+ fig.savefig(filename)
+ plt.close(fig)
+
+def plot_rmse_intensity(df: pd.DataFrame, filename: str):
+ fig = plt.figure(figsize=(12, 8))
+ gs = GridSpec(1, 1)
+ ax = fig.add_subplot(gs[0, 0])
+ sns.kdeplot(x=df.rmse, y=df.xgm_flux_t*1e-3,
+ fill=True,
+ ax=ax)
+ sns.scatterplot(x=df.rmse, y=df.xgm_flux_t*1e-3,
+ s=5,
+ alpha=0.4,
+ c="tab:red",
+ #size=df.root_mean_squared_pca_unc,
+ #sizes=(20, 200),
+ ax=ax)
+ ax = plt.gca()
+ ax.set(title=f"",
+ xlabel=r"Root-mean-squared error [a.u.]",
+ ylabel="Beam intensity [mJ]",
+ ylim=(0, df.xgm_flux_t.mean()*1e-3 + 3*df.xgm_flux_t.std()*1e-3)
+ )
+ ax.spines['top'].set_visible(False)
+ ax.spines['right'].set_visible(False)
+ plt.tight_layout()
+ fig.savefig(filename)
+ plt.close(fig)
+
+def plot_unc_intensity(df: pd.DataFrame, filename: str):
+ fig = plt.figure(figsize=(12, 8))
+ gs = GridSpec(1, 1)
+ ax = fig.add_subplot(gs[0, 0])
+ sns.kdeplot(x=df.total_unc, y=df.xgm_flux_t*1e-3,
+ fill=True,
+ ax=ax)
+ sns.scatterplot(x=df.total_unc, y=df.xgm_flux_t*1e-3,
+ s=5,
+ alpha=0.4,
+ c="tab:red",
+ #size=df.root_mean_squared_pca_unc,
+ #sizes=(20, 200),
+ ax=ax)
+ ax = plt.gca()
+ ax.set(title=f"",
+ xlabel=r"Uncertainty [a.u.]",
+ ylabel="Beam intensity [mJ]",
+ ylim=(0, df.xgm_flux_t.mean()*1e-3 + 3*df.xgm_flux_t.std()*1e-3)
+ )
+ ax.spines['top'].set_visible(False)
+ ax.spines['right'].set_visible(False)
+ plt.tight_layout()
+ fig.savefig(filename)
+ plt.close(fig)
+
+def plot_unc_rmse(df: pd.DataFrame, filename: str):
+ fig = plt.figure(figsize=(12, 8))
+ gs = GridSpec(1, 1)
+ ax = fig.add_subplot(gs[0, 0])
+ sns.kdeplot(x=2*df.total_unc, y=df.rmse,
+ fill=True,
+ ax=ax)
+ sns.scatterplot(x=2*df.total_unc, y=df.rmse,
+ s=5,
+ alpha=0.4,
+ c="tab:red",
+ #size=df.root_mean_squared_pca_unc,
+ #sizes=(20, 200),
+ ax=ax)
+ ax = plt.gca()
+ ax.set(title=f"",
+ xlabel=r"Root-mean-squared unc. [a.u.]",
+ ylabel="Root-mean-squared error [a.u.]",
+ )
+ ax.spines['top'].set_visible(False)
+ ax.spines['right'].set_visible(False)
+ plt.tight_layout()
+ fig.savefig(filename)
+ plt.close(fig)
+
+def pca_variance_plot(df: pd.DataFrame, filename: str, max_comp_frac: float=0.99):
+ """
+ Plot variance contribution.
+
+ Args:
+ filename: Output file name.
+ variance_ratio: Contribution of each component's variance.
+
+ """
+ fig = plt.figure(figsize=(12, 8))
+ gs = GridSpec(1, 1)
+ ax = fig.add_subplot(gs[0, 0])
+ c = np.cumsum(df.variance_ratio)
+ n_comp = int(df.n_comp.iloc[0])
+ ax.bar(1+np.arange(len(df.variance_ratio)), df.variance_ratio*100, color='tab:red', alpha=0.3, label="Per component")
+ ax.plot(1+np.arange(len(df.variance_ratio)), c*100, c='tab:blue', lw=5, label="Cumulative")
+ ax.plot([n_comp, n_comp], [0, c[n_comp]*100], lw=3, ls='--', c='m', label="Components kept")
+ ax.plot([0, n_comp], [c[n_comp]*100, c[n_comp]*100], lw=3, ls='--', c='m')
+ ax.legend(frameon=False)
+ print(f"PCA plot: total n. components: {len(df.variance_ratio)}")
+ x_max = np.where(c > max_comp_frac)[0][0]
+ print(f"Fraction of variance: {c[n_comp]}")
+ ax.set_yscale('log')
+ ax.set(title=f"",
+ xlabel="Component",
+ ylabel="Variance contribution [%]",
+ xlim=(1, x_max),
+ ylim=(0.01, 100))
+ ax.spines['top'].set_visible(False)
+ ax.spines['right'].set_visible(False)
+ plt.tight_layout()
+ fig.savefig(filename)
+ plt.close(fig)
+
+def moving_average(a, n=3):
+ ret = np.cumsum(a)
+ ret[n:] = ret[n:] - ret[:-n]
+ return ret[n - 1:] / n
+
+def plot_impulse(df: pd.DataFrame, filename: str):
+ """
+ Plot variance contribution.
+
+ Args:
+ filename: Output file name.
+ variance_ratio: Contribution of each component's variance.
+
+ """
+ fig = plt.figure(figsize=(12, 8))
+ gs = GridSpec(1, 1)
+ ax = fig.add_subplot(gs[0, 0])
+ x = df.wiener_energy.to_numpy()
+ y = np.absolute(df.impulse.to_numpy())
+ #x_new = np.linspace(-6, 6, 601)
+ #spl = make_interp_spline(x, np.log10(y), k=3)
+ #y_new = np.power(10, spl(x_new))
+ x_new = moving_average(x, n=5)
+ y_new = moving_average(y, n=5)
+ sel = (x_new >= -5.1) & (x_new <= 5.1)
+ ax.plot(x_new[sel], y_new[sel], c='tab:blue', lw=3)
+ ax.set_yscale('log')
+ ax.set(title=f"",
+ xlabel="Energy [eV]",
+ ylim=(1e-4, 0.4),
+ ylabel="Response [a.u.]",
+ )
+ ax.spines['top'].set_visible(False)
+ ax.spines['right'].set_visible(False)
+ plt.tight_layout()
+ fig.savefig(filename)
+ plt.close(fig)
+
+def plot_wiener(df: pd.DataFrame, filename: str):
+ """
+ Plot variance contribution.
+
+ Args:
+ filename: Output file name.
+ variance_ratio: Contribution of each component's variance.
+
+ """
+ fig = plt.figure(figsize=(12, 8))
+ gs = GridSpec(1, 1)
+ ax = fig.add_subplot(gs[0, 0])
+ ax.plot(df.wiener_energy, np.absolute(df.wiener_filter), c='tab:blue', lw=3)
+ ax.set_yscale('log')
+ ax.set(title=f"",
+ xlabel="Energy [eV]",
+ ylim=(1e-3, 1),
+ ylabel="Response [a.u.]",
+ )
+ ax.spines['top'].set_visible(False)
+ ax.spines['right'].set_visible(False)
+ plt.tight_layout()
+ fig.savefig(filename)
+ plt.close(fig)
+
+def plot_pes(df: pd.DataFrame, channel: Dict[str, int], filename: str, fast_range: Optional[Tuple[int, int]]=None, Ne1s: Optional[Tuple[int, int]]=None, label: Optional[Dict[str, str]]=None, refs: Optional[Dict[str, Dict[int, float]]]=None, counts_to_mv: Optional[float]=None):
+ """
+ Plot low-resolution spectrum.
+
+ Args:
+ filename: Output file name.
+ pes_raw_int: Low-resolution spectrum.
+
+ """
+ fig = plt.figure(figsize=(12, 8))
+ gs = GridSpec(1, 1)
+ ax = fig.add_subplot(gs[0, 0])
+ first, last = df.loc[:, 'first'].iloc[0], df.loc[:, 'last'].iloc[0]
+ first = first+220
+ last = last-270
+ print("Range:", first, last)
+ sel = (df.bin >= first) & (df.bin < last)
+ x = df.loc[sel, "bin"].to_numpy()
+ col = dict()
+ colors = ["tab:red", "tab:blue"]
+ p = list()
+ # plot each channel
+ for ich, ch in enumerate(channel.keys()):
+ if label is None:
+ sch = ch.replace('_', '')[-2:]
+ else:
+ sch = label[ch]
+ y = df.loc[sel, ch].to_numpy().astype(np.float32)
+ if counts_to_mv is not None:
+ y *= counts_to_mv
+ c = colors[ich]
+ col[ch] = c
+ p += [ax.plot(x, y, lw=2, c=c, label=sch)]
+ ax.set(title=f"",
+ ylim=(0, None),
+ #xlabel="Time-of-flight index",
+ xlabel="Samples",
+ ylabel="Counts [a.u.]" if counts_to_mv is None else "Digitizer reading [mV]")
+ ax.spines['top'].set_visible(False)
+ ax.spines['right'].set_visible(False)
+ minY, maxY = ax.get_ylim()
+ # show reference energy lines
+ if refs is not None:
+ for ich, ch in enumerate(channel.keys()):
+ for tof, energy in refs[ch].items():
+ ax.axvline(tof, 0, 0.5 + ich*0.17, ls='-.', lw=1, c=col[ch])
+ ax.text(tof-1, (0.51 + ich*0.18)*maxY, f"{energy} eV", fontsize=14, rotation="vertical", color=col[ch])
+ # show prompt line
+ for ch, prompt in channel.items():
+ ax.axvline(x=prompt, ls='--', lw=1, c=col[ch])
+ ax.text(prompt-3, 0.5*maxY, "Prompt", fontsize=16, rotation="vertical", color=col[ch])
+ # show the fast electrons range
+ if fast_range is not None:
+ x1, x2 = fast_range
+ xtext = int(x1 + (x2 - x1)*0.3)
+ ytext = 0.9*maxY
+ ax.fill_between([x1, x2], minY, maxY, alpha=0.2, facecolor="tab:olive")
+ ax.text(xtext, ytext, "Valence", fontsize=18, fontweight='bold')
+ ax.text(xtext, ytext-0.05*maxY, "Auger", fontsize=18, fontweight='bold')
+ # show the Ne 1s range
+ if Ne1s is not None:
+ x1, x2 = Ne1s
+ xtext = int(x1 + (x2 - x1)*0.3)
+ ytext = 0.9*maxY
+ ax.fill_between([x1, x2], minY, maxY, alpha=0.2, facecolor="tab:cyan")
+ ax.text(xtext, ytext, "Ne 1s", fontsize=22, fontweight='bold')
+ ns_per_sample = 0.5
+ cax = dict()
+ def f_(ch):
+ return (lambda kk: (np.array(kk) - int(channel[ch]))*ns_per_sample)
+ def i_(ch):
+ return (lambda kk: np.array(kk)/ns_per_sample + int(channel[ch]))
+ forward_ = {ch: f_(ch) for ch in channel}
+ inverse_ = {ch: i_(ch) for ch in channel}
+ for ich, (ch, prompt) in enumerate(channel.items()):
+ cax[ch] = ax.secondary_xaxis(1.0+0.07*ich, functions=(forward_[ch], inverse_[ch]))
+ #cax[ch].spines['left'].set_visible(False)
+ cax[ch].spines['top'].set_position(('outward', 10))
+ cax[ch].spines['top'].set_color(col[ch])
+ cax[ch].tick_params(axis='x', colors=col[ch], labelsize=16)
+ if ich == len(channel)-1:
+ cax[ch].set_xlabel('Time-of-flight [ns]', fontsize=16)
+ #cax[ch].xaxis.label.set_color(col[ch])
+ #cax[ch].title.set_color(col[ch])
+ ax.legend(frameon=False, loc='center')
+ plt.tight_layout()
+ fig.savefig(filename)
+ plt.close(fig)
+
+if __name__ == '__main__':
+ indir = 'p900331r69t70'
+ channel = {'channel_4_A': 2639,
+ 'channel_3_B': 2646,
+ }
+ label = {'channel_4_A': r'22.5$^\circ$',
+ 'channel_3_B': r'225$^\circ$',
+ }
+ Ne1s = (2710, 2742)
+ fast_range = (2650, 2670)
+ refs={'channel_4_A': {2716:1002.5, 2722:997.5},
+ 'channel_3_B': {2723:1002.5, 2729:997.5}
+ }
+ counts_to_mv = 40.0/100.0
+ #channel = 'sum'
+ #for fname in os.listdir(indir):
+ # if re.match(r'test_q100_[0-9]*\.csv', fname):
+ # fname = fname[:-4]
+ # print(f"Plotting {fname}")
+ # plot_final(pd.read_csv(f'{indir}/{fname}.csv'), f'{fname}.pdf')
+ # plot_pes(pd.read_csv(f'{indir}/{fname}_pes.csv'), channel, f'{fname}_pes.pdf')
+
+ for fname in ('test_q100_1724098413', 'test_q100_1724098596', 'test_q50_1724099445'):
+ plot_final(pd.read_csv(f'{indir}/{fname}.csv'), f'{fname}.pdf')
+ plot_pes(pd.read_csv(f'{indir}/{fname}_pes.csv'), channel, f'{fname}_pes.pdf',
+ fast_range=fast_range, Ne1s=Ne1s, label=label, refs=refs,
+ counts_to_mv=counts_to_mv)
+
+ plot_chi2(pd.read_csv(f'{indir}/quality.csv'), f'chi2_prepca.pdf')
+ plot_chi2_intensity(pd.read_csv(f'{indir}/quality.csv'), f'intensity_vs_chi2_prepca.pdf')
+ plot_unc_intensity(pd.read_csv(f'{indir}/quality.csv'), f'intensity_vs_unc.pdf')
+ plot_unc_rmse(pd.read_csv(f'{indir}/quality.csv'), f'rmse_vs_unc.pdf')
+ plot_rmse(pd.read_csv(f'{indir}/quality.csv'), f'rmse.pdf')
+ plot_rmse_intensity(pd.read_csv(f'{indir}/quality.csv'), f'intensity_vs_rmse.pdf')
+
+ plot_residue(pd.read_csv(f'{indir}/quality.csv'), f'residue.pdf')
+ plot_residue_corr(pd.read_csv(f'{indir}/quality.csv'), f'residue_corr.pdf')
+
+ df_model = pd.read_csv(f'{indir}/model.csv')
+ df_model.impulse = df_model.impulse.str.replace('i','j').apply(lambda x: np.complex(x))
+ df_model.wiener_filter = df_model.wiener_filter.str.replace('i','j').apply(lambda x: np.complex(x))
+ plot_impulse(df_model, f'impulse.pdf')
+ plot_wiener(df_model, f'wiener.pdf')
+
+ pca_variance_plot(pd.read_csv(f'{indir}/pca_spec.csv'), f'pca_spec.pdf', max_comp_frac=0.99)
+ pca_variance_plot(pd.read_csv(f'{indir}/pca_pes.csv'), f'pca_pes.pdf', max_comp_frac=0.95)
+
diff --git a/pyproject.toml b/pyproject.toml
index 86ea8d577381b8e30672742abe3d63e0d5ff09b6..9ea807b4f8cce6209cbb409d43aec2cc6e905656 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -29,7 +29,6 @@ dependencies = [
"scipy>=1.6",
"scikit-learn>=1.2.0",
"torch",
- "torchbnn",
]
[project.optional-dependencies]
diff --git a/run_tests.sh b/run_tests.sh
new file mode 100644
index 0000000000000000000000000000000000000000..41fb02678541b9814f347b9666bc11c9658bf56a
--- /dev/null
+++ b/run_tests.sh
@@ -0,0 +1,71 @@
+#!/bin/bash
+#SBATCH --partition=exfel
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --time=8:00:00
+#SBATCH --job-name=pes2spec
+#SBATCH -o slurm.%x.err.txt
+#SBATCH -e slurm.%x.err.txt
+#SBATCH --reservation=exfel_ml
+
+optstring="d:"
+
+DIR=results
+while getopts ${optstring} arg; do
+ case ${arg} in
+ d)
+ DIR=${OPTARG}
+ S=$OPTIND
+ ;;
+ *)
+ S=$((OPTIND))
+ break
+ ;;
+ esac
+done
+OPTS="${@:$S}"
+echo "Options: $OPTS"
+
+source /usr/share/Modules/init/sh
+
+module load exfel exfel_anaconda3
+
+cd $HOME/scratch/karabo/devices/pes_to_spec
+source env/bin/activate
+pwd
+export PYTHONPATH=$PYTHONPATH:$PWD
+
+mkdir $DIR
+
+do_it() {
+ p=$1
+ r=$2
+ rt=$3
+ output=$DIR/p${p}r${r}t${rt}
+ mkdir -p $output
+ echo "Proposal $p, run $r, test at run $rt"
+ CMD=(./pes_to_spec/test/offline_analysis.py -p $p -r $r -t $rt -d $output ${@:4})
+ echo "${CMD[*]}"
+ ${CMD[*]} 2>&1 | tee $output/log.txt
+}
+
+do_it 900331 69 70 $OPTS
+
+# train in run 2 and test in run 3
+#do_it 3384 2 3 $OPTS
+
+# new runs:
+#for run in 2 4
+#do
+# do_it 3384 $run $run $OPTS
+#done
+
+# train in run 4 and test in run 3
+#do_it 3384 4 3 $OPTS
+
+# old run
+#do_it 2828 206 206 $OPTS
+#do_it 2828 206 207 $OPTS
+
+#do_it 2828 207 207 $OPTS
+
diff --git a/run_tests_number_channels.sh b/run_tests_number_channels.sh
new file mode 100644
index 0000000000000000000000000000000000000000..3609bfb528e88460d60f043fb0bbcc7802a6ac84
--- /dev/null
+++ b/run_tests_number_channels.sh
@@ -0,0 +1,53 @@
+#!/bin/bash
+#SBATCH --partition=exfel
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --time=8:00:00
+#SBATCH --job-name=nch_pes2spec
+#SBATCH -o slurm.%x.err.txt
+#SBATCH -e slurm.%x.err.txt
+#SBATCH --reservation=exfel_ml
+
+optstring="d:"
+
+DIR=results
+while getopts ${optstring} arg; do
+ case ${arg} in
+ d)
+ DIR=${OPTARG}
+ S=$OPTIND
+ ;;
+ *)
+ S=$((OPTIND))
+ break
+ ;;
+ esac
+done
+OPTS="${@:$S}"
+echo "Options: $OPTS"
+
+source /usr/share/Modules/init/sh
+
+module load exfel exfel_anaconda3
+
+cd $HOME/scratch/karabo/devices/pes_to_spec
+source env/bin/activate
+pwd
+export PYTHONPATH=$PYTHONPATH:$PWD
+
+mkdir $DIR
+
+do_it() {
+ p=$1
+ r=$2
+ rt=$3
+ output=$DIR/p${p}r${r}t${rt}
+ mkdir -p $output
+ echo "Proposal $p, run $r, test at run $rt"
+ CMD=(./pes_to_spec/test/channel_sensitivity.py -p $p -r $r -t $rt -d $output ${@:4})
+ echo "${CMD[*]}"
+ ${CMD[*]} 2>&1 | tee $output/log.txt
+}
+
+do_it 900331 69 70 $OPTS
+