Merge pull request #651 from scipp/bifrost-sqw

Changes to SQW writer for BIFROST

Author: jl-wynen

src/scippneutron/io/sqw/_build.py

@@ -12,6 +12,7 @@
 from typing import Any, BinaryIO
 
 import numpy as np
+import numpy.typing as npt
 import scipp as sc
 
 from .._files import open_or_pass
@@ -89,7 +90,7 @@ def __init__(
             ("", "main_header"): main_header,
         }
 
-        self._dnd_placeholder: _DndPlaceholder | None = None
+        self._dnd_data: _DndPlaceholder | _DndData | None = None
         self._pix_wrap: _PixWrap | None = None
         self._instrument: SqwIXNullInstrument | None = None
         self._sample: SqwIXSample | None = None
@@ -121,7 +122,7 @@ def create(self, *, chunk_size: int = 8192) -> Path | None:
                         self._pix_wrap.write(sqw_io, chunk_size=chunk_size)  # type: ignore[union-attr]
                     case SqwDataBlockType.dnd:
                         # Type guaranteed by _serialize_data_blocks
-                        self._dnd_placeholder.write(sqw_io)  # type: ignore[union-attr]
+                        self._dnd_data.write(sqw_io)  # type: ignore[union-attr]
                     case _:
                         raise NotImplementedError(
                             f"Unsupported data block type: {descriptor.block_type}"
@@ -131,19 +132,18 @@ def create(self, *, chunk_size: int = 8192) -> Path | None:
 
     def add_pixel_data(
         self,
-        data: sc.DataArray,
+        data: npt.NDArray[np.float32],
         *,
         experiments: list[SqwIXExperiment],
         n_dims: int = 4,
-        rows: tuple[str, ...] = _DEFAULT_PIX_ROWS,
         row_units: tuple[str | None, ...] = _DEFAULT_PIX_ROW_UNITS,
     ) -> SqwBuilder:
         self._n_dims = n_dims
         self._data_blocks[("experiment_info", "expdata")] = SqwMultiIXExperiment(
             experiments
         )
 
-        self._pix_wrap = _split_pix_rows(data, rows, row_units)
+        self._pix_wrap = _PixWrap(row_data=data, row_units=row_units)
         metadata = self._make_pix_metadata(self._pix_wrap)
         self._data_blocks[("pix", "metadata")] = metadata
         self._data_blocks[("", "main_header")].nfiles = len(experiments)
@@ -164,11 +164,22 @@ def _add_dnd_metadata(self, block: SqwDndMetadata) -> SqwBuilder:
         self._data_blocks[("data", "metadata")] = block
         return self
 
-    def add_empty_dnd_data(self, block: SqwDndMetadata) -> SqwBuilder:
+    def add_empty_dnd_data(self, metadata: SqwDndMetadata) -> SqwBuilder:
         # The file must always contain a DND block
-        builder = self._add_dnd_metadata(block)
-        builder._dnd_placeholder = _DndPlaceholder(
-            shape=tuple(map(int, block.axes.n_bins_all_dims))  # type: ignore[call-overload]
+        builder = self._add_dnd_metadata(metadata)
+        builder._dnd_data = _DndPlaceholder(
+            shape=tuple(map(int, metadata.axes.n_bins_all_dims))  # type: ignore[call-overload]
+        )
+        return builder
+
+    def add_dnd_data(
+        self, metadata: SqwDndMetadata, *, data: sc.Variable, counts: sc.Variable
+    ) -> SqwBuilder:
+        builder = self._add_dnd_metadata(metadata)
+        builder._dnd_data = _DndData(
+            shape=tuple(map(int, metadata.axes.n_bins_all_dims)),  # type: ignore[call-overload]
+            data=data,
+            counts=counts,
         )
         return builder
 
@@ -195,12 +206,10 @@ def _make_pix_metadata(self, pix_wrap: _PixWrap) -> SqwPixelMetadata:
             data_range=np.vstack(
                 [
                     (
-                        sc.to_unit(row.min(), unit).value,
-                        sc.to_unit(row.max(), unit).value,
-                    )
-                    for row, unit in zip(
-                        pix_wrap.row_data, pix_wrap.row_units, strict=True
+                        pix_wrap.row_data[:, i].min().astype(np.float64),
+                        pix_wrap.row_data[:, i].max().astype(np.float64),
                     )
+                    for i, unit in enumerate(pix_wrap.row_units)
                 ]
             ),
         )
@@ -232,15 +241,7 @@ def _serialize_data_blocks(
                 locked=False,
             )
 
-        if self._dnd_placeholder is not None:
-            buffers[("data", "nd_data")] = None
-            descriptors[("data", "nd_data")] = SqwDataBlockDescriptor(
-                block_type=SqwDataBlockType.dnd,
-                name=("data", "nd_data"),
-                position=0,
-                size=self._dnd_placeholder.size(),
-                locked=False,
-            )
+        self._serialize_dnd_data(buffers, descriptors)
 
         if self._pix_wrap is not None:
             buffers[("pix", "data_wrap")] = None
@@ -254,6 +255,31 @@ def _serialize_data_blocks(
 
         return buffers, descriptors
 
+    def _serialize_dnd_data(
+        self,
+        buffers: dict[DataBlockName, memoryview | None],
+        descriptors: dict[tuple[str, str], SqwDataBlockDescriptor],
+    ) -> None:
+        match self._dnd_data:
+            case _DndPlaceholder() as placeholder:
+                buffers[("data", "nd_data")] = None
+                descriptors[("data", "nd_data")] = SqwDataBlockDescriptor(
+                    block_type=SqwDataBlockType.dnd,
+                    name=("data", "nd_data"),
+                    position=0,
+                    size=placeholder.size(),
+                    locked=False,
+                )
+            case _DndData() as data:
+                buffers[("data", "nd_data")] = None
+                descriptors[("data", "nd_data")] = SqwDataBlockDescriptor(
+                    block_type=SqwDataBlockType.dnd,
+                    name=("data", "nd_data"),
+                    position=0,
+                    size=data.size(),
+                    locked=False,
+                )
+
     def _prepare_data_blocks(self) -> dict[DataBlockName, Any]:
         filepath, filename = self._filepath_and_name
         blocks = {
@@ -380,33 +406,13 @@ def _to_canonical_block_order(
         ("pix", "data_wrap"),
     )
     blocks = dict(blocks)
-    out = {name: block for name in order if (block := blocks.get(name)) is not None}
+    out = {
+        name: block for name in order if (block := blocks.pop(name, None)) is not None
+    }
     out.update(blocks)  # append remaining blocks if any
     return out
 
 
-def _split_pix_rows(
-    data: sc.DataArray, rows: tuple[str, ...], row_units: tuple[str | None, ...]
-) -> _PixWrap:
-    """Prepare the selected pixel rows for writing."""
-    selected = []
-    for name in rows:
-        if name == 'signal':
-            selected.append(sc.values(data.data))
-        elif name == 'error':
-            selected.append(sc.variances(data.data))
-        else:
-            if data.coords.is_edges(name, data.dim):
-                raise sc.BinEdgeError(
-                    f"Pixel data must not contain bin-edges, got edges for '{name}'."
-                )
-            selected.append(data.coords[name])
-    return _PixWrap(
-        row_data=selected,
-        row_units=row_units,
-    )
-
-
 @dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
 class _DndPlaceholder:
     shape: tuple[int, ...]
@@ -425,34 +431,44 @@ def write(self, sqw_io: LowLevelSqw) -> None:
         sqw_io.write_array(np.zeros(self.shape, dtype="uint64"))
 
 
+@dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
+class _DndData:
+    data: sc.Variable
+    counts: sc.Variable
+    shape: tuple[int, ...]
+
+    def size(self) -> int:
+        n_elem = int(np.prod(self.shape))
+        return 4 + 4 * len(self.shape) + 3 * 8 * n_elem
+
+    def write(self, sqw_io: LowLevelSqw) -> None:
+        sqw_io.write_u32(len(self.shape))
+        for s in self.shape:
+            sqw_io.write_u32(s)
+        sqw_io.write_array(self.data.values.astype("float64", copy=False))
+        sqw_io.write_array(np.sqrt(self.data.variances.astype("float64", copy=False)))
+        # Strictly speaking uint64, but we don't support that in Scipp.
+        # So we have to assume that all numbers are positive anyway.
+        # Then we can avoid unnecessary type casts.
+        sqw_io.write_array(self.counts.values.astype("int64", copy=False))
+
+
 @dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
 class _PixWrap:
-    row_data: list[sc.Variable]
+    row_data: npt.NDArray[np.float32]  # shape: (n_pixels, n_rows)
     row_units: tuple[str | None, ...]
 
     def size(self) -> int:
         # *4 for f32
         return 4 + 8 + self.n_rows() * 4 * self.n_pixels()
 
     def n_rows(self) -> int:
-        return len(self.row_data)
+        return self.row_data.shape[1]
 
     def n_pixels(self) -> int:
-        return len(self.row_data[0])
+        return self.row_data.shape[0]
 
     def write(self, sqw_io: LowLevelSqw, chunk_size: int) -> None:
         sqw_io.write_u32(self.n_rows())
         sqw_io.write_u64(self.n_pixels())
-
-        buffer = np.empty((self.n_pixels(), self.n_rows()), dtype=np.float32)
-        remaining = self.n_pixels()
-        for offset in range(0, self.n_rows(), chunk_size):
-            n = min(chunk_size, remaining)
-            remaining -= n
-            for i_row, (row, unit) in enumerate(
-                zip(self.row_data, self.row_units, strict=True)
-            ):
-                buffer[:n, i_row] = sc.to_unit(
-                    row[offset : offset + chunk_size], unit, copy=False
-                ).values
-            sqw_io.write_array(buffer[:n])
+        sqw_io.write_array_fortran_layout(self.row_data)

src/scippneutron/io/sqw/_ir.py

@@ -150,5 +150,7 @@ def _serialize_field(
             ty=field.ty, shape=(len(field.value),) if field.value else (), data=[field]
         )
     if isinstance(field, Array):
-        return ObjectArray(ty=field.ty, shape=field.value.shape[::-1], data=field.value)
+        return ObjectArray(
+            ty=field.ty, shape=field.value.shape[::-1], data=field.value.T
+        )
     return ObjectArray(ty=field.ty, shape=(1,), data=[field])

src/scippneutron/io/sqw/_low_level_io.py

@@ -155,7 +155,7 @@ def read_array(
             flat = flat.copy()
             self._file.seek(self.position + count * dtype.itemsize)
         else:
-            flat = np.fromfile(self._file, dtype=dtype, count=int(np.prod(shape)))
+            flat = np.fromfile(self._file, dtype=dtype, count=count)
         # Invert the shape because files use column-major layout.
         return flat.reshape(shape[::-1])
 
@@ -198,6 +198,13 @@ def write_chars(self, value: str) -> None:
     @_annotate_write_exception("array")
     def write_array(
         self, array: npt.NDArray[np.float64] | npt.NDArray[np.float32]
+    ) -> None:
+        # Transpose to match the column-major layout of the file.
+        self.write_array_fortran_layout(array.T)
+
+    @_annotate_write_exception("array_fortran_layout")
+    def write_array_fortran_layout(
+        self, array: npt.NDArray[np.float64] | npt.NDArray[np.float32]
     ) -> None:
         out = array.astype(array.dtype.newbyteorder(self.byteorder.value), copy=False)
         if isinstance(self._file, BytesIO):

src/scippneutron/io/sqw/_models.py

@@ -252,7 +252,7 @@ def _serialize_to_dict(
             "version": ir.F64(self.version),
             "name": ir.String(self.name),
             "target_name": ir.String(self.target_name),
-            "frequency": ir.F64(self.frequency.value),  # TODO unit
+            "frequency": ir.F64(self.frequency.to(unit='Hz').value),
         }
 
 
@@ -343,8 +343,8 @@ def _serialize_to_dict(
             "emode": ir.F64(float(self.emode.value)),
             "en": ir.Array(en.values, ty=ir.TypeTag.f64),
             "psi": ir.F64(_angle_value(self.psi)),
-            "u": ir.Array(self.u.values, ty=ir.TypeTag.f64),
-            "v": ir.Array(self.v.values, ty=ir.TypeTag.f64),
+            "u": _variable_to_float_array(self.u, unit="1/angstrom"),
+            "v": _variable_to_float_array(self.v, unit="1/angstrom"),
             "omega": ir.F64(_angle_value(self.omega)),
             "dpsi": ir.F64(_angle_value(self.dpsi)),
             "gl": ir.F64(_angle_value(self.gl)),

src/scippneutron/io/sqw/_sqw.py

@@ -342,7 +342,7 @@ def get_vec(name: str, unit: str) -> sc.Variable | None:
 
     return SqwLineProj(
         lattice_spacing=sc.vector(
-            _get_struct_field(struct, "alatt").data, unit="1/angstrom"
+            _get_struct_field(struct, "alatt").data, unit="angstrom"
         ),
         lattice_angle=sc.vector(_get_struct_field(struct, "angdeg").data, unit="deg"),
         offset=_parse_1d_multi_unit_array(
@@ -394,8 +394,7 @@ def _parse_ix_source_2_0(struct: ir.Struct) -> SqwIXSource:
     return SqwIXSource(
         name=name,
         target_name=target_name,
-        # TODO unit currently unknown
-        frequency=sc.scalar(frequency, unit=None),
+        frequency=sc.scalar(frequency, unit="Hz"),
     )
 
 
@@ -408,7 +407,7 @@ def _parse_ix_null_instrument_1_0(struct: ir.Struct) -> SqwIXNullInstrument:
 def _parse_ix_sample_0_0(struct: ir.Struct) -> SqwIXSample:
     name = _get_scalar_struct_field(struct, "name")
     lattice_spacing = sc.vector(
-        _get_struct_field(struct, "alatt").data, unit="1/angstrom"
+        _get_struct_field(struct, "alatt").data, unit="angstrom"
     )
     lattice_angle = sc.vector(_get_struct_field(struct, "angdeg").data, unit="deg")
     return SqwIXSample(
@@ -436,11 +435,18 @@ def g(n: str) -> Any:
         (e,) = candidate_efix
         efix = sc.scalar(e.value, unit="meV")
 
+    emode = EnergyMode(g("emode"))
+
     raw_en = _get_struct_field(struct, "en").data
-    if isinstance(raw_en, np.ndarray):
-        en = raw_en.squeeze()
+    raw_en = (
+        raw_en
+        if isinstance(raw_en, np.ndarray)
+        else np.array([e.value for e in raw_en])
+    )
+    if emode == EnergyMode.indirect:
+        en = sc.array(dims=["detector", "energy_transfer"], values=raw_en, unit="meV")
     else:
-        en = [e.value for e in raw_en]
+        en = sc.array(dims=["energy_transfer"], values=raw_en.squeeze(), unit="meV")
 
     angle_unit = sc.Unit("deg" if g("angular_is_degree") else "rad")
 
@@ -449,11 +455,11 @@ def g(n: str) -> Any:
         filepath=g("filepath"),
         run_id=int(g("run_id")) - 1,
         efix=efix,
-        emode=EnergyMode(g("emode")),
-        en=sc.array(dims=["energy_transfer"], values=en, unit="meV"),
+        emode=emode,
+        en=en,
         psi=sc.scalar(g("psi"), unit=angle_unit),
-        u=sc.vector(_get_struct_field(struct, "u").data),
-        v=sc.vector(_get_struct_field(struct, "v").data),
+        u=sc.vector(_get_struct_field(struct, "u").data, unit="1/angstrom"),
+        v=sc.vector(_get_struct_field(struct, "v").data, unit="1/angstrom"),
         omega=sc.scalar(g("omega"), unit=angle_unit),
         dpsi=sc.scalar(g("dpsi"), unit=angle_unit),
         gl=sc.scalar(g("gl"), unit=angle_unit),