Source code for ewoksfluo.io.nexus

import datetime
import logging
import os
from contextlib import ExitStack
from contextlib import contextmanager
from typing import Iterator
from typing import Mapping
from typing import Optional

import h5py
from ewoksdata.data.hdf5.dataset_writer import DatasetWriter
from silx.io import h5py_utils

from . import hdf5
from .types import ScanData

logger = logging.getLogger(__name__)




def save_as_bliss_scan(
    filename: str,
    entry_name: str,
    scan_data: Iterator[ScanData],
    positioners: Optional[Mapping[str, float]] = None,
    title: Optional[str] = None,
    **openoptions,
):
    openoptions.setdefault("mode", "a")
    os.makedirs(os.path.abspath(os.path.dirname(filename)), exist_ok=True)
    with h5py_utils.File(filename, **openoptions) as f:
        if entry_name in f:
            logger.warning("%s::/%s already exists", filename, entry_name)
            return

        with ExitStack() as stack:
            entry = stack.enter_context(_writer_context(f, entry_name))
            instrument = entry.create_group("instrument")
            instrument.attrs["NX_class"] = "NXinstrument"
            measurement = entry.create_group("measurement")
            measurement.attrs["NX_class"] = "NXcollection"

            if title:
                entry["title"] = title

            if positioners:
                collection = instrument.create_group("positioners_start")
                collection.attrs["NX_class"] = "NXcollection"
                for k, v in positioners.items():
                    collection[k] = v
                positioners = dict(positioners)

            writers = dict()
            for data in scan_data:
                detector = instrument.require_group(data.group)
                if data.detector_type == "positioner":
                    detector.attrs["NX_class"] = "NXpositioner"
                    if positioners:
                        positioners[data.group] = data.data
                else:
                    detector.attrs["NX_class"] = "NXdetector"
                    if "type" not in detector and data.detector_type:
                        detector["type"] = data.detector_type

                key = data.group, data.name
                writer = writers.get(key)
                if writer is None:
                    writer = stack.enter_context(DatasetWriter(detector, data.name))
                    writers[key] = writer
                    if data.local_alias:
                        detector[data.local_alias] = h5py.SoftLink(writer.dataset_name)
                    if data.global_alias and data.global_alias not in measurement:
                        measurement[data.global_alias] = h5py.SoftLink(
                            writer.dataset_name
                        )

                writer.add_points(data.data)

        if positioners:
            collection = instrument.create_group("positioners")
            collection.attrs["NX_class"] = "NXcollection"
            for k, v in positioners.items():
                collection[k] = v



@contextmanager
def _writer_context(root: hdf5.FileType, entry_name: str) -> Iterator[h5py.Group]:
    root.attrs["NX_class"] = "NXroot"
    root.attrs["creator"] = "ewoksfluo"
    entry = root.create_group(entry_name)
    try:
        entry.attrs["NX_class"] = "NXentry"
        entry.attrs["start_time"] = _timestamp()
        writer = entry.create_group("writer")
        writer.attrs["NX_class"] = "NXnote"
        writer["status"] = "STARTING"
        yield entry
    except Exception:
        writer["status"][()] = "FAILED"
        raise
    else:
        writer["status"][()] = "SUCCEEDED"
    finally:
        entry["end_time"] = _timestamp()


def _timestamp() -> str:
    return datetime.datetime.now().astimezone().isoformat()