Write instrument, sample

src/sqomega/_build.py

@@ -15,6 +15,7 @@
 
 import numpy as np
 
+from . import _ir as ir
 from ._bytes import Byteorder
 from ._files import open_binary
 from ._low_level_io import LowLevelSqw
@@ -26,9 +27,13 @@
     SqwFileHeader,
     SqwFileType,
     SqwIXExperiment,
+    SqwIXNullInstrument,
+    SqwIXSample,
     SqwMainHeader,
     SqwMultiIXExperiment,
     SqwPixelMetadata,
+    UniqueObjContainer,
+    UniqueRefContainer,
 )
 from ._read_write import write_object_array
 
@@ -68,6 +73,7 @@ def __init__(
         main_header = SqwMainHeader(
             full_filename=self._full_filename,
             title=title,
+            # To be replaced when registering pixel data.
             nfiles=0,
             # To be replaced when writing the file.
             creation_date=datetime(1, 1, 1, tzinfo=timezone.utc),
@@ -77,7 +83,8 @@ def __init__(
         }
 
         self._pix_placeholder: _PixPlaceholder | None = None
-        self._expdata: list[SqwIXExperiment] = []
+        self._instrument: SqwIXNullInstrument | None = None
+        self._sample: SqwIXSample | None = None
 
     @contextmanager
     def create(self) -> Generator[Sqw, None, None]:
@@ -126,9 +133,8 @@ def register_pixel_data(
             raise RuntimeError("SQW builder already has pixel data")
         self._n_dim = n_dims
 
-        self._expdata = experiments
         self._data_blocks[('experiment_info', 'expdata')] = SqwMultiIXExperiment(
-            self._expdata
+            experiments
         )
         self._data_blocks[("pix", "metadata")] = SqwPixelMetadata(
             full_filename=self._full_filename,
@@ -139,12 +145,21 @@ def register_pixel_data(
             n_pixels=n_pixels,
             rows=rows,
         )
+        self._data_blocks[('', 'main_header')].nfiles = len(experiments)
         return self
 
     def add_dnd_metadata(self, block: SqwDndMetadata) -> SqwBuilder:
         self._data_blocks[("data", "metadata")] = block
         return self
 
+    def add_default_instrument(self, instrument: SqwIXNullInstrument) -> SqwBuilder:
+        self._instrument = instrument
+        return self
+
+    def add_default_sample(self, sample: SqwIXSample) -> SqwBuilder:
+        self._sample = sample
+        return self
+
     def _make_file_header(self) -> SqwFileHeader:
         return SqwFileHeader(
             prog_name="horace",
@@ -194,11 +209,29 @@ def _serialize_data_blocks(
 
     def _prepare_data_blocks(self) -> dict[DataBlockName, Any]:
         filepath, filename = self._filepath_and_name
-        return {
+        blocks = {
             key: block.prepare_for_serialization(filepath=filepath, filename=filename)
             for key, block in self._data_blocks.items()
         }
 
+        nfiles = blocks[('', 'main_header')].nfiles
+        if self._instrument is not None:
+            blocks[('experiment_info', 'instruments')] = _broadcast_unique_ref(
+                self._instrument,
+                n=nfiles,
+                baseclass="IX_inst",
+                global_name='GLOBAL_NAME_INSTRUMENTS_CONTAINER',
+            )
+        if self._sample is not None:
+            blocks[('experiment_info', 'samples')] = _broadcast_unique_ref(
+                self._sample,
+                n=nfiles,
+                baseclass="IX_samp",
+                global_name='GLOBAL_NAME_SAMPLES_CONTAINER',
+            )
+
+        return blocks
+
     def _serialize_block_allocation_table(
         self,
         block_descriptors: dict[DataBlockName, SqwDataBlockDescriptor],
@@ -240,7 +273,7 @@ def _serialize_block_allocation_table(
 
     @property
     def _full_filename(self) -> str:
-        return os.fspath(self._stored_path or "")
+        return os.fspath(self._stored_path or "in_memory")
 
     @property
     def _filepath_and_name(self) -> tuple[str, str]:
@@ -269,6 +302,22 @@ def _write_data_block_descriptor(
     return pos
 
 
+def _broadcast_unique_ref(
+    obj: ir.Serializable,
+    n: int,
+    baseclass: str,
+    global_name: str,
+) -> UniqueRefContainer:
+    return UniqueRefContainer(
+        global_name=global_name,
+        objects=UniqueObjContainer(
+            baseclass=baseclass,
+            objects=[obj],
+            indices=[0] * n,
+        ),
+    )
+
+
 @dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
 class _PixPlaceholder:
     n_pixels: int

src/sqomega/_models.py

@@ -235,7 +235,13 @@ class SqwIXSource(ir.Serializable):
     version: ClassVar[float] = 2.0
 
     def _serialize_to_dict(self) -> dict[str, ir.Object]:
-        raise NotImplementedError()
+        return {
+            "serial_name": ir.String(self.serial_name),
+            "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
+        }
 
 
 @dataclass(kw_only=True, slots=True)
@@ -247,7 +253,12 @@ class SqwIXNullInstrument(ir.Serializable):
     version: ClassVar[float] = 2.0
 
     def _serialize_to_dict(self) -> dict[str, ir.Object]:
-        raise NotImplementedError()
+        return {
+            "serial_name": ir.String(self.serial_name),
+            "version": ir.F64(self.version),
+            "name": ir.String(self.name),
+            "source": self.source.serialize_to_ir(),
+        }
 
 
 class EnergyMode(enum.Enum):
@@ -265,7 +276,13 @@ class SqwIXSample(ir.Serializable):
     version: ClassVar[float] = 0.0
 
     def _serialize_to_dict(self) -> dict[str, ir.Object]:
-        raise NotImplementedError()
+        return {
+            "serial_name": ir.String(self.serial_name),
+            "version": ir.F64(self.version),
+            "alatt": _variable_to_float_array(self.lattice_spacing, "1/angstrom"),
+            "angdeg": _variable_to_float_array(self.lattice_angle, "deg"),
+            "name": ir.String(self.name),
+        }
 
 
 # In contrast to SQW files, this model contains the nested
@@ -337,6 +354,51 @@ def _serialize_to_dict(self) -> dict[str, ir.Object]:
         }
 
 
+@dataclass(kw_only=True, slots=True)
+class UniqueRefContainer(ir.Serializable):
+    global_name: str
+    objects: UniqueObjContainer
+
+    serial_name: ClassVar[str] = "unique_references_container"
+    version: ClassVar[float] = 1.0
+
+    def _serialize_to_dict(self) -> dict[str, ir.Object]:
+        return {
+            "serial_name": ir.String(self.serial_name),
+            "version": ir.F64(self.version),
+            "stored_baseclass": ir.String(self.objects.baseclass),
+            "global_name": ir.String(self.global_name),
+            "unique_objects": self.objects.serialize_to_ir().to_object_array(),
+        }
+
+
+@dataclass(kw_only=True, slots=True)
+class UniqueObjContainer(ir.Serializable):
+    baseclass: str
+    objects: list[ir.Serializable]
+    indices: list[int]
+
+    serial_name: ClassVar[str] = "unique_objects_container"
+    version: ClassVar[float] = 1.0
+
+    def _serialize_to_dict(self) -> dict[str, ir.Object]:
+        return {
+            "serial_name": ir.String(self.serial_name),
+            "version": ir.F64(self.version),
+            "baseclass": ir.String(self.baseclass),
+            "unique_objects": ir.CellArray(
+                shape=(len(self.objects),),
+                data=[obj.serialize_to_ir().to_object_array() for obj in self.objects],
+            ),
+            "idx": ir.ObjectArray(
+                shape=(len(self.indices),),
+                # +1 to convert to 1-based indexing
+                data=np.array(self.indices) + 1.0,
+                ty=ir.TypeTag.f64,
+            ),
+        }
+
+
 def _angle_value(x: sc.Variable) -> float:
     return x.to(unit='rad', dtype='float64', copy=False).value
 

src/sqomega/_sqw.py

@@ -472,7 +472,10 @@ def _parse_unique_references_container_1_0(struct: ir.Struct) -> list[Any]:
 def _parse_unique_objects_container_1_0(struct: ir.Struct) -> list[Any]:
     objects = _get_struct_field(struct, "unique_objects").data
     idx = _get_struct_field(struct, "idx").data
+    if not isinstance(idx, np.ndarray):  # it is list[ir.F64]
+        idx = np.array([i.value for i in idx])
     parsed_objects = [_try_parse_block(obj) for obj in objects]
+    # -1 to convert to 0-based index
     return [parsed_objects[int(i) - 1] for i in idx]