niess.nexus¶

def node_name(node) -> str | None:
    """The name of a group, dataset, or named stream node."""
    if not isinstance(node, dict):
        return None
    if 'name' in node:
        return node['name']
    return (node.get('config') or {}).get('name')

def stream(module: str, config: dict, attrs: dict[str, Any] | None = None) -> dict:
    """A filewriter module directive -- ``ev44``, ``da00``, ``f144``, ``link``, ..."""
    node = {'module': module, 'config': config}
    attributes = _attributes(None, attrs)
    if attributes:
        node['attributes'] = attributes
    return node

def load_bindings(path) -> dict[str, Binding]:
    """Read a facility bindings file, such as ECDC's ``bindings.yaml``.

    This is the file ECDC compiles, together with an ``instrument.yaml``, into the NeXus
    template the real instrument's filewriter uses. Each entry says which topic and
    source fill one log. Reading it here lets a niess file name the same streams
    without keeping a second copy of the answers.
    """
    from pathlib import Path
    import yaml
    with Path(path).open('r', encoding='utf-8') as file:
        return parse_bindings(yaml.safe_load(file) or {})

def resolve_stream(translation, default: dict | None = None, name: str = 'data') -> dict | None:
    """The stream group for a component, or ``None`` if it publishes nothing.

    The protocol is never chosen here. Some monitors belong on ``da00`` histograms
    and some on ``ev44`` events; which one is a property of the instrument setup, so
    it is read from the instrument in priority order:

    1. a ``METADATA "nexus_structure_stream_data"`` block on the component -- the
       escape hatch for instruments authored outside niess, emitted verbatim;
    2. a ``nexus_stream`` entry in the component's niess provenance ``extra``,
       which is how a niess component records the choice made when the instrument
       was built;
    3. ``default`` -- the component type's established behaviour, used only when
       the instrument expressed no preference at all.

    A component with no selection and no default gets no stream group rather than a
    guessed one.
    """
    from json import JSONDecodeError, loads

    for metadata in translation.instance.metadata:
        if metadata.mimetype != 'application/json':
            continue
        if metadata.name != 'nexus_structure_stream_data':
            continue
        try:
            return stream_group_from_config(name, loads(metadata.value))
        except JSONDecodeError:
            continue

    selection = None
    if translation.provenance is not None:
        selection = translation.provenance.extra.get('nexus_stream')

    if selection is None:
        selection = default
    if selection is None:
        return None

    return stream_group_from_selection(name, selection)

def component_body(nx_class: str, children=None, attrs=None, name=None,
                   position=None, rotation=None) -> dict:
    """What a translator returns: the class and contents of one component's group.

    ``position`` and ``rotation`` state where the thing sits in the frame it hangs from,
    for the composites that have no placement of their own -- an analyzer is at its
    frame's origin, turned by the Bragg angle; a detector is a distance along its.
    Anything with a position and an orientation of its own is placed by those.
    """
    return {'nx_class': nx_class, 'children': list(children or []),
            'attrs': dict(attrs or {}), 'name': name,
            'position': position, 'rotation': rotation}

def emit(visit: Visit, body: dict) -> None:
    """Put one component's group into the instrument.

    Refuses a name already used directly under NXinstrument. Two groups with one name
    would be one group in the file, silently: the second replaces or merges into the
    first. The accelerator's NXsource is `source`, so a component of that name would
    collide with it.
    """
    context = visit.context
    node = _placed(visit, body)
    name = node_name(node)
    if any(node_name(c) == name for c in context.instrument_group.get('children', [])):
        raise ValueError(f'{name!r} is already directly under {INSTRUMENT_PATH}; '
                         'rename the component')
    context.emitted_names[visit.id] = node_name(node)
    context.pending.append((visit, node))
    add_child(context.instrument_group, node)

def to_nexus_structure(instrument, registry=None, nxlog_root: str | None = None,
                       streams=None, efu_monitors: bool = False) -> dict:
    """Convert ``instrument`` to ESS NeXus Structure JSON.

    ``streams`` says where a driven axis's numbers come from: ``SIMULATED`` (the
    default) writes the simulation's own parameter names, ``REAL`` writes the EPICS
    positioners the components declare. One tree, two files; see
    :mod:`niess.nexus.bindings`.

    ``efu_monitors`` describes every monitor that has a cbm channel by the stream its
    EFU publishes -- for a simulation whose monitor rays are collected and replayed to
    that EFU. Otherwise a monitor describes the histogram McStas produces.
    """
    from .bindings import as_streams
    context = NexusContext(
        instrument=instrument,
        nxlog_root=DEFAULT_NXLOG_ROOT if nxlog_root is None else nxlog_root,
        streams=as_streams(streams), efu_monitors=efu_monitors)
    walk(instrument, NEXUS_REGISTRY if registry is None else registry, context=context)
    # once, at the end: what feeds what is not known until everything has a name
    for visit, node in context.pending:
        for direction, names in zip(('inputs', 'outputs'), context.neighbours(visit)):
            if names:
                # Datasets, not attributes. NXcomponent declares `inputs` and `outputs`
                # as fields, so every class extending it inherits them as fields, and a
                # validator reading the NXDL rejects them written any other way.
                # One name is written as a string rather than a list of one, which is
                # what the standard and every reader of these files expect.
                add_child(node, dataset(direction,
                                        names[0] if len(names) == 1 else names))
    if getattr(instrument, 'origin', None):
        _rebase_on_origin(context.instrument_group, instrument.origin)
    entry = group('entry', nx_class='NXentry',
                  children=[context.instrument_group])
    return {'children': [entry]}

def translator(*classes):
    """Register a function returning a component body for one or more niess classes."""
    def decorate(func):
        def run(visit: Visit):
            body = func(visit)
            if body is not None:
                emit(visit, body)
            return body

        holder = type(func.__name__, (), {'leaf': staticmethod(run),
                                          '__doc__': func.__doc__})
        for klass in classes:
            NEXUS_REGISTRY.register(klass)(holder)
        return func

    return decorate

def bifrost_detector_source(arc, triplet) -> str:
    """Which triplet an ev44 message came from, as the source string ICD 01 uses."""
    return f'arc={arc};triplet={triplet}'

def icd_pixel(resolution, arc, triplet, tube, position):
    """Pixel id per ICD 01 v6; ``position`` runs from 0 to ``resolution - 1``."""
    return 27 * resolution * arc + 9 * resolution * tube + resolution * triplet + position + 1

def arc_and_triplet(visit: Visit) -> tuple[int, int]:
    """Which arc and which triplet, from where the thing sits in the instrument.

    `niess.nexus` matches a regex against a generated ``WHEN`` clause for this, because
    an emitted McStas instrument is all it has. The tree knows: a triplet is the arm it
    belongs to, in the channel that arm belongs to.
    """
    from ..bifrost.arm import Arm
    from ..bifrost.channel import Channel
    arm, channel = visit.ancestor(Arm), visit.ancestor(Channel)
    return (0 if arm is None else arm.index, 0 if channel is None else channel.index)

def register_bifrost() -> None:
    """BIFROST's analyzers and detectors, read off the objects."""
    from ..bifrost.analyzer import Analyzer
    from ..bifrost.triplet import Triplet
    from ..components.filter import RadialFilterCollimator
    from ..components.geometry import Off, Cylinder

    def register(klass, func):
        def run(visit):
            body = func(visit)
            if body is not None:
                emit(visit, body)
            return SKIP

        BIFROST_REGISTRY.register(klass)(
            type(func.__name__, (), {'enter': staticmethod(run),
                                     'leaf': staticmethod(run),
                                     '__doc__': func.__doc__}))

    def analyzer(visit):
        """A Rowland-geometry analyzer, as a segmented NXcrystal.

        The blade count, shape and mosaic are the analyzer's own fields. The other way
        round they are McStas component parameters -- NH, zwidth, yheight, mosaic --
        read back out of a component call. Each blade is one face of the OFF geometry,
        at its own place on the Rowland circle and turned to focus, as the arm says.
        """
        import numpy as np
        obj = visit.obj
        blade = obj.central_blade
        mosaic = float(blade.mosaic.to(unit='arcminute').value)
        centres, normals, perp_q, perp_plane, _ = visit.ancestor(_arm()).obj.local_blades()
        count = len(centres)

        up = np.array([0., 1., 0.])      # the Rowland cylinder's axis
        half_width, half_height = perp_q / 2, perp_plane / 2
        vertices, faces = [], []
        for i, (centre, normal) in enumerate(zip(centres, normals)):
            along = np.cross(normal, up)  # across the blade, in the Rowland plane
            for h, w in ((-1, -1), (-1, 1), (1, 1), (1, -1)):
                # to a nanometre: the digits beyond are the arithmetic's, and differ
                # between numpy builds -- enough to change a written structure
                corner = centre + h * half_height * up + w * half_width * along
                vertices.append([round(float(v), 9) for v in corner])
            faces.append([4 * i, 4 * i + 1, 4 * i + 2, 4 * i + 3])

        return component_body('NXcrystal', [
            dataset('usage', 'Bragg'),
            dataset('d_spacing', float(obj.central_blade.plane_spacing.to(unit='angstrom').value),
                    attrs={'units': 'angstrom'}),
            dataset('segment_width', float(perp_q), attrs={'units': 'm'}),
            dataset('segment_height', float(perp_plane), attrs={'units': 'm'}),
            dataset('segment_columns', count),
            dataset('segment_rows', 1),
            dataset('mosaic_horizontal', mosaic, attrs={'units': 'arcminutes'}),
            dataset('mosaic_vertical', mosaic, attrs={'units': 'arcminutes'}),
            Off(vertices, faces).to_nexus('geometry'),
        ], name=visit.emit_name('monochromator'),
           rotation=(0.0, float(visit.ancestor(_arm()).obj.analyzer_theta.value), 0.0))

    def detector(visit):
        """A triplet of He3 tubes: one shared cylinder, repositioned per pixel."""
        import numpy as np
        from scipp import vector

        obj = visit.obj
        arc, triplet = arc_and_triplet(visit)
        tubes = obj.tubes
        nj = int(tubes[0].elements)
        offsets, height, radius = obj.local_tubes()
        ni = len(offsets)
        di = np.asarray(offsets)
        half_pixel = height / nj / 2
        dj = -np.linspace(-height / 2 + half_pixel, height / 2 - half_pixel, nj)
        grid_j, grid_i = np.meshgrid(dj, di)

        numbers = [[icd_pixel(nj, arc, triplet, tube, position)
                    for position in range(nj)] for tube in range(ni)]

        geometry = Cylinder.single({
            'length': vector([0, 2*half_pixel, 0], unit='m'),
            'radius': vector([radius, 0, 0], unit='m'),
        }).to_nexus('geometry')

        # Event streaming is what these tubes do: without an NXevent_data group the
        # filewriter has nothing to fill the detector with. Which topic and which source
        # is the facility's business rather than the tubes', so the triplet may say, and
        # then a bindings file may; unset, these have always been BIFROST's own topic
        # and an arc/triplet source.
        name = visit.emit_name('triplet')
        selection = obj.stream
        if selection is None and hasattr(visit.context.streams, 'stream'):
            selection = visit.context.streams.stream(name)
        if selection is None:
            selection = {'module': 'ev44',
                         'source': bifrost_detector_source(arc, triplet),
                         'topic': BIFROST_DETECTOR_TOPIC}
        stream = visit.context.stream_group(selection)

        return component_body('NXdetector', [
            dataset('detector_number',
                    np.array(numbers).astype('int32').tolist(), dtype='int32'),
            # straight after detector_number, which is where the other route and
            # moreniius before it put it
            stream,
            dataset('x_pixel_offset', grid_i.tolist(), dtype='double',
                    attrs={'units': 'm'}),
            dataset('y_pixel_offset', grid_j.tolist(), dtype='double',
                    attrs={'units': 'm'}),
            dataset('x_pixel_size', 2 * radius, attrs={'units': 'm'}),
            dataset('y_pixel_size', height / nj, attrs={'units': 'm'}),
            dataset('diameter', 2 * radius, attrs={'units': 'm'}),
            dataset('type', f'{ni} He3 tubes in series'),
            geometry,
        ], name=name,
           position=vector([0., 0., 1.]) * visit.ancestor(_arm()).obj.analyzer_detector_distance)

    def collimator(visit):
        obj = visit.obj
        return component_body('NXcollimator', [
            dataset('type', 'radial'),
            dataset('divergence_x',
                    float(obj.collimation_angle.to(unit='degree').value),
                    attrs={'units': 'degrees'}),
        ])

    register(Analyzer, analyzer)
    register(Triplet, detector)
    BIFROST_REGISTRY.register(RadialFilterCollimator)(
        type('collimator', (), {'leaf': staticmethod(
            lambda visit: emit(visit, collimator(visit)))}))