From 52d8fac62288a9d3dd95fa2f9246470f51129e78 Mon Sep 17 00:00:00 2001
From: Martin Teichmann <martin.teichmann@xfel.eu>
Date: Wed, 9 Mar 2022 17:46:30 +0100
Subject: [PATCH] use xarrays for keeping data
do not store the data in our own class, employ xarray for that.
---
src/toolbox_scs/detectors/hrixs.py | 173 +++++++++++++++--------------
1 file changed, 88 insertions(+), 85 deletions(-)
diff --git a/src/toolbox_scs/detectors/hrixs.py b/src/toolbox_scs/detectors/hrixs.py
index 59d51e2..8a6b9e2 100644
--- a/src/toolbox_scs/detectors/hrixs.py
+++ b/src/toolbox_scs/detectors/hrixs.py
@@ -1,6 +1,8 @@
from functools import lru_cache
import numpy as np
+import matplotlib.pyplot as plt
+from scipy.optimize import leastsq
from scipy.optimize import curve_fit
from scipy.signal import fftconvolve
@@ -56,6 +58,26 @@ def find_curvature(image, frangex=None, frangey=None,
return curv[:-1][::-1]
+def find_curvature(img, args, plot=False):
+ def parabola(x, a, b, c, s=0, h=0, o=0):
+ return (a*x + b)*x + c
+ def gauss(y, x, a, b, c, s, h, o=0):
+ return h * np.exp(-((y - parabola(x, a, b, c)) / (2 * s))**2) + o
+ x = np.arange(img.shape[1])[None, :]
+ y = np.arange(img.shape[0])[:, None]
+
+ if plot:
+ plt.figure(figsize=(10,10))
+ plt.imshow(img, cmap='gray', aspect='auto', vmin=26000, vmax=27000, interpolation='nearest')
+ plt.plot(x[0, :], parabola(x[0, :], *args))
+
+ args, _ = leastsq(lambda args: (gauss(y, x, *args) - img).ravel(), args)
+
+ if plot:
+ plt.plot(x[0, :], parabola(x[0, :], *args))
+ return args
+
+
def correct_curvature(image, factor=None, axis=1):
if factor is None:
return
@@ -257,7 +279,7 @@ class hRIXS:
PROPOSAL = 2769
# image range
- X_RANGE = np.s_[1300:-100]
+ X_RANGE = np.s_[:]
Y_RANGE = np.s_[:]
# centroid
@@ -272,100 +294,81 @@ class hRIXS:
METHOD = 'centroid' # ['centroid', 'integral']
- @classmethod
- def set_params(cls, **params):
- for key, value in params.items():
- setattr(cls, key.upper(), value)
+ FIELDS = ['hRIXS_det', 'Delay', 'hRIXS_index',
+ 'hRIXS_delay', 'hRIXS_norm']
- def __set_params(self, **params):
- self.__class__.set_params(**params)
- self.refresh()
+ def set_params(self, **params):
+ for key, value in params.items():
+ setattr(self, key.upper(), value)
- @classmethod
- def get_params(cls, *params):
+ def get_params(self, *params):
if not params:
params = ('proposal', 'x_range', 'y_range',
'threshold', 'curve_a', 'curve_b',
'factor', 'range', 'bins',
- 'method')
- return {param: getattr(cls, param.upper()) for param in params}
-
- def refresh(self):
- cls = self.__class__
- for cached in ['_centroid', '_integral']:
- getattr(cls, cached).fget.cache_clear()
+ 'method', 'fields')
+ return {param: getattr(self, param.upper()) for param in params}
- def __init__(self, images, norm=None):
- self.images = images
- self.norm = norm
-
- # class/instance method compatibility
- self.set_params = self.__set_params
-
- @classmethod
- def from_run(cls, runNB, proposal=None, first_wrong=False):
+ def from_run(self, runNB, proposal=None, extra_fields=()):
if proposal is None:
- proposal = cls.PROPOSAL
-
- run, data = tb.load(proposal, runNB=runNB, fields=['hRIXS_det'])
-
- # Get slow train data
- mnemo = tb.mnemonics_for_run(run)['SCS_slowTrain']
- slow_train = run[mnemo['source'], mnemo['key']].ndarray().sum()
-
- return cls(images=data['hRIXS_det'][1 if first_wrong else 0:].data,
- norm=slow_train)
-
- @property
- @lru_cache()
- def _centroid(self):
- return sum((centroid(image[self.Y_RANGE, self.X_RANGE].T,
- threshold=self.THRESHOLD,
- curvature=(self.CURVE_A, self.CURVE_B), )
- for image in self.images), [])
-
- def _centroid_spectrum(self, bins=None, range=None, normalize=True):
+ proposal = self.PROPOSAL
+ run, data = tb.load(proposal, runNB=runNB,
+ fields=self.FIELDS + list(extra_fields))
+
+ return data
+
+ def find_curvature(self, runNB, proposal=None, plot=True):
+ data = self.from_run(runNB, proposal)
+
+ image = data['hRIXS_det'].sum(dim='trainId') \
+ .to_numpy()[self.X_RANGE, self.Y_RANGE].T
+ spec = (image - image[:10, :].mean()).mean(axis=1)
+ mean = np.average(np.arange(len(spec)), weights=spec)
+ args = find_curvature(
+ image, (-2e-7, 0.02, mean - 0.02 * image.shape[1] / 2, 3,
+ spec.max(), image.mean()), plot=plot)
+ self.CURVE_B, self.CURVE_A, *_ = args
+ return self.CURVE_A, self.CURVE_B
+
+ def centroid(self, data, bins=None, range=None):
if bins is None:
bins = self.BINS
if range is None:
range = self.RANGE
- r = np.array(self._centroid)
- hy, hx = np.histogram(r[:, 0], bins=bins, range=range)
- if normalize and self.norm is not None:
- hy = hy / self.norm
-
- return (hx[:-1] + hx[1:]) / 2, hy
-
- @property
- @lru_cache()
- def _integral(self):
- return sum((integrate(image[self.Y_RANGE, self.X_RANGE].T,
- factor=self.FACTOR,
- range=self.RANGE,
- curvature=(self.CURVE_A, self.CURVE_B))
- for image in self.images))
-
- def _integral_spectrum(self, normalize=True):
- values = self._integral
- if normalize and self.norm is not None:
- values = values / self.norm
- return np.arange(values.size), values
-
- @property
- def corrected(self):
- return decentroid(self._centroid)
-
- def spectrum(self, normalize=True):
- spec_func = (self._centroid_spectrum if self.METHOD.lower() == 'centroid'
- else self._integral_spectrum)
- return spec_func(normalize=normalize)
-
- def __sub__(self, other):
- px, py = self.spectrum()
- mx, my = other.spectrum()
- return (px + mx) / 2, py - my
-
- def __add__(self, other):
- return self.__class__(images=list(self.images) + list(other.images),
- norm=self.norm + other.norm)
+ ret = np.zeros((len(data["hRIXS_det"]), bins))
+ for image, r in zip(data["hRIXS_det"], ret):
+ c = centroid(
+ image.to_numpy()[self.X_RANGE, self.Y_RANGE].T,
+ threshold=self.THRESHOLD,
+ curvature=(self.CURVE_A, self.CURVE_B))
+ if not len(c):
+ continue
+ r = np.array(c)
+ hy, hx = np.histogram(r[:, 0], bins=bins, range=range)
+ ret[:] = hy
+
+ data = data.assign_coords(energy=np.linspace(*range, bins))
+ return data.assign(spectrum=(("trainId", "energy"), ret))
+
+ aggregators = dict(
+ hRIXS_det=lambda x, dim: x.sum(dim=dim),
+ Delay=lambda x, dim: x.mean(dim=dim),
+ hRIXS_delay=lambda x, dim: x.mean(dim=dim),
+ hRIXS_norm=lambda x, dim: x.sum(dim=dim),
+ spectrum=lambda x, dim: x.sum(dim=dim),
+ )
+
+ def aggregator(self, da, dim):
+ agg = self.aggregators.get(da.name)
+ if agg is None:
+ return None
+ return agg(da, dim=dim)
+
+ def aggregate(self, ds, dim="trainId"):
+ ret = ds.map(self.aggregator, dim=dim)
+ ret = ret.drop_vars([n for n in ret if n not in self.aggregators])
+ return ret
+
+ def normalize(self, data):
+ return data.assign(normalized=data["spectrum"] / data["hRIXS_norm"])
--
GitLab