# -*- coding: utf-8 -*-
import os
import numpy as np
import pint
from ewokscore import Task
from ewoksfluo.tasks.hdf5_utils import link_bliss_scan
# from PIL import Image
from silx.io import h5py_utils
from silx.io import url
ureg = pint.UnitRegistry()
DEFAULTS = {
"Nsample": 1000,
"normalization": None,
"ratio_edge": None,
"fw": None,
}
[docs]
class VibrationsSpectrum(
Task,
input_names=[
"bliss_scan_uri",
"output_root_uri",
"counters",
"output_name",
"output_group",
],
optional_input_names=[
"Nsample",
"normalization",
"ratio_edge",
"fw",
],
output_names=[
"bliss_scan_uri",
"output_root_uri",
"ratio_edge",
],
):
[docs]
def run(self):
inputs = {**DEFAULTS, **self.get_input_values()}
bliss_scan_uri = url.DataUrl(inputs["bliss_scan_uri"])
# epath = ESRFPath(bliss_root_uri.file_path())
output_root_uri = url.DataUrl(inputs["output_root_uri"])
norm = inputs["normalization"]
N = inputs["Nsample"]
data = None
with h5py_utils.open_item(
bliss_scan_uri.file_path(), bliss_scan_uri.data_path()
) as scan:
for cnt, f in inputs["counters"].items():
_data = scan["measurement/{}".format(cnt)][()]
acq_time = float(scan["instrument/fscan_parameters/acq_time"][()])
if norm is not None:
ndata = scan.parent[norm]["measurement/{}".format(cnt)][()]
self.outputs.ratio_edge = np.mean(_data) / np.mean(ndata)
_data = _data / np.mean(ndata)
else:
self.outputs.ratio_edge = None
if data is None:
data = f * _data
else:
data += f * _data
if data.shape[0] < N:
raise RuntimeError("No enough points!")
fw = inputs["fw"]
ratio_edge = inputs["ratio_edge"]
if fw is not None and ratio_edge is not None: # Simulation
p3 = fw / 2.0 / np.arccosh(1.0 / (0.5) ** 0.5)
x0 = p3 * np.arctanh(2.0 * ratio_edge - 1.0)
data = data - np.mean(data)
data_sim = (1.0 - np.tanh((data - x0) / p3)) / 2.0
print(f"{p3=} {x0=} {data=} {data_sim=}")
data = data_sim
nsub = data.shape[0] // N
data = data[: nsub * N]
data.shape = nsub, N
data = data - np.mean(data, 1)[:, None]
spectrum = np.abs(np.fft.rfft(data, axis=1))
os.makedirs(os.path.dirname(output_root_uri.file_path()), exist_ok=True)
cnt = inputs["output_name"]
with h5py_utils.open_item(output_root_uri.file_path(), "/", mode="a") as fd:
grp = fd.require_group(output_root_uri.data_path())
link_bliss_scan(grp, bliss_scan_uri)
grp = grp.require_group("vibrations")
grp.attrs["NX_class"] = "NXprocessing"
grp = grp.require_group(inputs["output_group"])
grp.attrs["default"] = "spectrum"
grp.attrs["NX_class"] = "NXcollection"
grp_data = grp.require_group("timefreq")
grp_data.attrs["NX_class"] = "NXdata"
grp_data.attrs["axes"] = [
".",
"freq",
]
grp_data.attrs["interpretation"] = "image"
if "freq" not in grp_data:
grp_data.create_dataset("freq", data=np.fft.rfftfreq(N, acq_time))
if cnt in grp_data:
del grp_data[cnt]
else:
if "signal" not in grp_data.attrs:
grp_data.attrs["signal"] = cnt
elif "auxiliary_signals" not in grp_data.attrs:
grp_data.attrs["auxiliary_signals"] = [
cnt,
]
else:
grp_data.attrs["auxiliary_signals"] = [
*grp_data.attrs["auxiliary_signals"],
cnt,
]
grp_data.create_dataset(cnt, data=spectrum)
grp_data = grp.require_group("spectrum")
grp_data.attrs["NX_class"] = "NXdata"
grp_data.attrs["axes"] = [
"freq",
]
grp_data.attrs["interpretation"] = "spectrum"
if "freq" not in grp_data:
grp_data.create_dataset("freq", data=np.fft.rfftfreq(N, acq_time))
if cnt in grp_data:
del grp_data[cnt]
else:
if "signal" not in grp_data.attrs:
grp_data.attrs["signal"] = cnt
elif "auxiliary_signals" not in grp_data.attrs:
grp_data.attrs["auxiliary_signals"] = [
cnt,
]
else:
grp_data.attrs["auxiliary_signals"] = [
*grp_data.attrs["auxiliary_signals"],
cnt,
]
grp_data.create_dataset(cnt, data=np.mean(spectrum, 0))
self.outputs.bliss_scan_uri = self.inputs.bliss_scan_uri
self.outputs.output_root_uri = self.inputs.output_root_uri