Merge pull request #115 from scipp/json-generator

Add "JSON" generator to facilitate testing of live workflows in downstream packages

src/ess/reduce/nexus/json_generator.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024 Scipp contributors (https://github.com/scipp)
+"""Generators for "JSON" data from a NeXus file, for the purpose of testing."""
+
+from collections.abc import Generator
+from typing import Any
+
+import scipp as sc
+
+
+def _as_str_attr(value: Any, name: str) -> dict:
+    val = str(value)
+    return {"string_size": len(val), "type": "string", "name": name, "values": val}
+
+
+def _variable_to_json(var: sc.Variable, name: str):
+    attrs = []
+    if var.dtype == sc.DType.datetime64:
+        offset = var.min()
+        attrs.append(_as_str_attr(offset.value, "offset"))
+        var = var - offset
+    if var.unit is not None:
+        attrs.append(_as_str_attr(var.unit, "units"))
+    return {
+        "module": "dataset",
+        "config": {
+            "name": name,
+            "values": var.values,
+            "size": list(var.shape),
+            "type": str(var.dtype),
+        },
+        "attributes": attrs,
+    }
+
+
+_event_index_0 = sc.array(dims=('dummy',), values=[0], unit=None)
+
+
+def _event_data_pulse_to_json(pulse: sc.DataArray) -> dict:
+    content = pulse.value.coords
+    event_time_zero = sc.concat([pulse.coords['event_time_zero']], 'dummy')
+    event_time_offset = content['event_time_offset']
+    event_id = content.get('event_id')
+    # I think we always have a pixel_id in the flatbuffer, so monitors just get ones?
+    if event_id is None:
+        event_id = sc.ones(sizes=event_time_offset.sizes, dtype='int32', unit=None)
+    children = [
+        _variable_to_json(event_time_zero, name='event_time_zero'),
+        _variable_to_json(event_time_offset, name='event_time_offset'),
+        _variable_to_json(event_id, name='event_id'),
+        _variable_to_json(_event_index_0, name='event_index'),
+    ]
+    group = {
+        "type": "group",
+        "name": "events_0",
+        "children": children,
+        "attributes": [_as_str_attr("NXevent_data", name="NX_class")],
+    }
+    return group
+
+
+def event_data_generator(data: sc.DataArray) -> Generator[dict, None, None]:
+    """
+    Generate JSON data for event data from a NeXus file.
+
+    Parameters
+    ----------
+    data:
+        A data array with event data, equivalent to what ScippNexus would load from an
+        NXevent_data group in a NeXus file.
+
+    Yields
+    ------
+    :
+        A dict of data for a single event data pulse that can be wrapped in a
+        :py:class:`ess.reduce.nexus.json_nexus.JSONGroup`.
+    """
+    for pulse in data:
+        yield _event_data_pulse_to_json(pulse)

tests/nexus/json_generator_test.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024 Scipp contributors (https://github.com/scipp)
+import pytest
+import scippnexus as snx
+from scipp.testing import assert_identical
+
+from ess.reduce import data
+from ess.reduce.nexus.json_generator import event_data_generator
+from ess.reduce.nexus.json_nexus import json_nexus_group
+
+
+def test_event_data_generator_monitor_events_round_trip() -> None:
+    filename = data.loki_tutorial_sample_run_60250()
+    monitor = snx.load(filename, root='entry/instrument/monitor_1/monitor_1_events')
+    generator = event_data_generator(monitor)
+    for i in range(len(monitor)):
+        group = json_nexus_group(next(generator))
+        assert_identical(group[()], monitor[i : i + 1])
+    with pytest.raises(StopIteration):
+        next(generator)
+
+
+def test_event_data_generator_detector_events_round_trip() -> None:
+    filename = data.loki_tutorial_sample_run_60250()
+    detector = snx.load(
+        filename, root='entry/instrument/larmor_detector/larmor_detector_events'
+    )
+    generator = event_data_generator(detector)
+    for i in range(100):
+        group = json_nexus_group(next(generator))
+        assert_identical(group[()], detector[i : i + 1])
+
+
+def test_event_data_generator_without_event_id_yields_ones() -> None:
+    filename = data.loki_tutorial_sample_run_60250()
+    base = snx.load(filename, root='entry/instrument/monitor_1/monitor_1_events')
+    monitor = base.bins.drop_coords('event_id')
+    generator = event_data_generator(monitor)
+    for i in range(100):
+        group = json_nexus_group(next(generator))
+        assert_identical(group[()], base[i : i + 1])

tests/nexus/json_nexus_test.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024 Scipp contributors (https://github.com/scipp)
 import json
 from copy import deepcopy
 from functools import lru_cache