from datetime import datetime
from typing import Any
from typing import List
from typing import Tuple
from typing import Union
from silx.io.dictdump import dicttonx
from silx.io.url import DataUrl
from ..types.spectrum import Spectrum
from ..types.xasobject import XASObject
[docs]
def write_xas(
output_file: str,
xas_obj: Union[dict, XASObject],
overwrite: bool = True,
output_group_template: str = "1.{}",
output_group_start_index: int = 1,
) -> List[str]:
"""
Save XAS object in HDF5, one group per spectrum.
:param output_file: HDF5 file name
:param xas_obj: XAS object with data and results to save.
:param overwrite: overwrite existing HDF5 content.
:param output_group_template: HDF5 group name can have a placeholder `"{}"` for the spectrum index.
:param output_group_start_index: Start spectrum index to be used in `output_group_template`.
"""
if not output_file:
raise ValueError("no output filename defined")
if isinstance(xas_obj, dict):
xas_obj = XASObject.from_dict(xas_obj)
if not isinstance(xas_obj, XASObject):
raise TypeError(type(xas_obj))
# Prepare saving
data_paths = []
data_groups = []
last_top_group = None
for i in range(xas_obj.n_spectrum):
out_index = output_group_start_index + i
parts = output_group_template.format(out_index).split("/")
top_group, data_group, data_path = _prepare_nxdict(*parts)
data_paths.append(data_path)
data_groups.append(data_group)
if last_top_group != top_group:
last_top_group = top_group
dicttonx(
top_group,
output_file,
mode="a",
update_mode="add",
)
# The output group template does not have a placeholder
if len(set(data_paths)) != len(data_paths):
data_paths = [
f"{s}_{output_group_start_index+i}" for i, s in enumerate(data_paths, 1)
]
# Save results
urls = []
for spec, data_group, data_path in zip(xas_obj.spectra, data_groups, data_paths):
data_group = dict(data_group)
data_group.update(_build_raw(spec))
data_group.update(_build_results(spec))
dicttonx(
data_group,
output_file,
data_path,
mode="a",
update_mode="replace" if overwrite else "add",
)
url = DataUrl(file_path=output_file, data_path=data_path)
urls.append(url.path())
return urls
def _prepare_nxdict(*parts: str) -> Tuple[dict, dict, str]:
parts = [s for s in parts if s]
if not parts:
parts = ["entry"]
root = {"@NX_class": "NXroot"}
imax = len(parts) - 1
parent = root
for i, name in enumerate(parts):
data_node = {"@NX_class": "NXentry" if i == 0 else "NXcollection"}
if i < imax:
parent[name] = data_node
parent = parent[name]
data_path = "/".join(["", *parts])
return root, data_node, data_path
def _build_raw(spec: Spectrum) -> dict:
data = {"data": _mu_vs_energy(spec)}
data = _prune_none_and_empty_dict(data)
if data:
data["@default"] = "data"
return data
def _build_results(spec: Spectrum) -> dict:
results = {
"statistics": _build_statistics(spec),
"normalized": _normalized_mu_vs_energy(spec),
"raw_enorm": _mu_vs_normalized_energy(spec),
"normalized_enorm": _normalized_mu_vs_normalized_energy(spec),
"exafs": _exafs(spec),
"ft": _ft(spec),
"noise": _noise(spec),
"larch": _nxparameters(**spec.larch_dict.to_nexus()),
"pymca": _nxparameters(**spec.pymca_dict.to_nexus()),
}
results = _prune_none_and_empty_dict(results)
if not results:
return {}
process = {
"@NX_class": "NXprocess",
"program": "est",
"date": datetime.now().astimezone().isoformat(),
**results,
}
return {"results": process}
def _mu_vs_energy(spec: Spectrum) -> dict:
return _nxdata(spec, "energy", "mu", "pre_edge", "post_edge")
def _mu_vs_normalized_energy(spec: Spectrum) -> dict:
return _nxdata(
spec,
"normalized_energy",
"mu",
"pre_edge",
"post_edge",
)
def _normalized_mu_vs_energy(spec: Spectrum) -> dict:
return _nxdata(spec, "energy", "flatten_mu", "normalized_mu")
def _normalized_mu_vs_normalized_energy(spec: Spectrum) -> dict:
return _nxdata(
spec,
"normalized_energy",
"normalized_mu",
"flatten_mu",
)
def _exafs(spec: Spectrum) -> dict:
return _nxdata(spec, "k", "chi", "chi_weighted_k")
def _ft(spec: Spectrum) -> dict:
data = _nxdata(
spec, "ft.radius", "ft.intensity", "ft.real", "ft.imaginary", "ft.phase"
)
if not data:
return {}
data["parameters"] = _nxparameters(**spec.ft.to_nexus("window_weight"))
return data
def _noise(spec: Spectrum) -> dict:
return _nxdata(spec, "energy", "noise_savgol")
def _build_statistics(spec: Spectrum) -> dict:
parameters = {}
parameters.update(spec.to_nexus("e0"))
parameters.update(spec.to_nexus("edge_step"))
parameters.update(spec.to_nexus("noise_e_min"))
parameters.update(spec.to_nexus("noise_e_max"))
parameters.update(spec.to_nexus("raw_noise_savgol"))
parameters.update(spec.to_nexus("norm_noise_savgol"))
if parameters:
return _nxparameters(**parameters)
return {}
def _nxdata(spec: Spectrum, axis: str, signal: str, *auxiliary_signals: str) -> dict:
data = spec.to_nexus(axis, signal, *auxiliary_signals)
data = _prune_none_and_empty_dict(data)
# Extract field name from model path. For example `radius` from `ft.radius`
axis = axis.split(".")[-1]
signal = signal.split(".")[-1]
auxiliary_signals = [s.split(".")[-1] for s in auxiliary_signals]
if signal not in data or axis not in data:
return {}
auxiliary_signals = [s for s in auxiliary_signals if s in data]
data = {"@NX_class": "NXdata", "@signal": signal, "@axes": [axis], **data}
if auxiliary_signals:
data["@auxiliary_signals"] = auxiliary_signals
return data
def _nxparameters(**parameters: Any) -> dict:
parameters = _prune_none_and_empty_dict(parameters)
if not parameters:
return {}
return {"@NX_class": "NXparameters", **parameters}
def _prune_none_and_empty_dict(obj: Any) -> Any:
if isinstance(obj, dict):
cleaned = {}
for k, v in obj.items():
v = _prune_none_and_empty_dict(v)
if v is None:
continue
if isinstance(v, dict) and len(v) == 0:
continue
cleaned[k] = v
return cleaned
return obj
