Dev/tests for Zebra Saver IOC

scripts/run-caproto-ioc.sh

@@ -3,7 +3,8 @@
 set -vxeuo pipefail
 
 CAPROTO_IOC="${1:-srx_caproto_iocs.base}"
-
+DEFAULT_PREFIX="BASE:{{Dev:Save1}}:"
+CAPROTO_IOC_PREFIX="${2:-${DEFAULT_PREFIX}}"
 # shellcheck source=/dev/null
 if [ -f "/etc/profile.d/epics.sh" ]; then
     . /etc/profile.d/epics.sh
@@ -12,4 +13,4 @@ fi
 export EPICS_CAS_AUTO_BEACON_ADDR_LIST="no"
 export EPICS_CAS_BEACON_ADDR_LIST="${EPICS_CA_ADDR_LIST:-127.0.0.255}"
 
-python -m "${CAPROTO_IOC}" --prefix="BASE:{{Dev:Save1}}:" --list-pvs
+python -m "${CAPROTO_IOC}" --prefix="${CAPROTO_IOC_PREFIX}" --list-pvs

scripts/run-caproto-zebra-ioc.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+
+set -vxeuo pipefail
+
+SCRIPT_DIR="$(dirname "$0")"
+
+bash "${SCRIPT_DIR}/run-caproto-ioc.sh" srx_caproto_iocs.zebra.caproto_ioc "XF:05IDD-ES:1{{Dev:Zebra2}}:"

src/srx_caproto_iocs/base.py

@@ -16,7 +16,7 @@
 from ophyd import Device, EpicsSignal, EpicsSignalRO
 from ophyd.status import SubscriptionStatus
 
-from .utils import now, save_hdf5
+from .utils import now, save_hdf5_nd
 
 
 class AcqStatuses(Enum):
@@ -104,8 +104,7 @@ async def queue(self, instance, async_lib):
         )
         thread.start()
 
-    @stage.putter
-    async def stage(self, instance, value):
+    async def _stage(self, instance, value):
         """The stage method to perform preparation of a dataset to save the data."""
         if (
             instance.value in [True, StageStates.STAGED.value]
@@ -150,7 +149,11 @@ async def stage(self, instance, value):
 
         return False
 
-    def _get_current_dataset(self, frame):
+    @stage.putter
+    async def stage(self, *args, **kwargs):
+        return await self._stage(*args, **kwargs)
+
+    async def _get_current_dataset(self, frame):
         """The method to return a desired dataset.
 
         See https://scikit-image.org/docs/stable/auto_examples/data/plot_3d.html
@@ -184,7 +187,7 @@ async def acquire(self, instance, value):
         # Delegate saving the resulting data to a blocking callback in a thread.
         payload = {
             "filename": self.full_file_path.value,
-            "data": self._get_current_dataset(frame=self.frame_num.value),
+            "data": await self._get_current_dataset(frame=self.frame_num.value),
             "uid": str(uuid.uuid4()),
             "timestamp": ttime.time(),
             "frame_number": self.frame_num.value,
@@ -210,7 +213,7 @@ def saver(request_queue, response_queue):
             data = received["data"]
             frame_number = received["frame_number"]
             try:
-                save_hdf5(fname=filename, data=data, mode="a")
+                save_hdf5_nd(fname=filename, data=data, mode="x", group_path="enc1")
                 print(
                     f"{now()}: saved {frame_number=} {data.shape} data into:\n  {filename}"
                 )

src/srx_caproto_iocs/utils.py

@@ -15,15 +15,41 @@ def now(as_object=False):
     return _now.isoformat()
 
 
-def save_hdf5(
+def save_hdf5_zebra(
+    fname,
+    data,
+    dtype="float32",
+    mode="x",
+):
+    """The function to export the 1-D data to an HDF5 file.
+
+    Check https://docs.h5py.org/en/stable/high/file.html#opening-creating-files for modes:
+
+        r           Readonly, file must exist (default)
+        r+          Read/write, file must exist
+        w           Create file, truncate if exists
+        w- or x     Create file, fail if exists
+        a           Read/write if exists, create otherwise
+    """
+    with h5py.File(fname, mode, libver="latest") as h5file_desc:
+        for pvname, value in data.items():
+            dataset = h5file_desc.create_dataset(
+                pvname,
+                data=value,
+                dtype=dtype,
+            )
+            dataset.flush()
+
+
+def save_hdf5_nd(
     fname,
     data,
     group_name="/entry",
     group_path="data/data",
     dtype="float32",
     mode="x",
 ):
-    """The function to export the data to an HDF5 file.
+    """The function to export the N-D data to an HDF5 file (N>1).
 
     Check https://docs.h5py.org/en/stable/high/file.html#opening-creating-files for modes:
 
@@ -54,5 +80,5 @@ def save_hdf5(
         h5file_desc.swmr_mode = True
 
         dataset.resize((frame_num + 1, *frame_shape))
-        dataset[frame_num, :, :] = data
+        dataset[frame_num, ...] = data
         dataset.flush()

src/srx_caproto_iocs/zebra/caproto_ioc.py

@@ -1,20 +1,227 @@
 from __future__ import annotations
 
 import textwrap
+from enum import Enum
 
-from caproto.server import run, template_arg_parser
+from caproto import ChannelType
+from caproto.server import pvproperty, run, template_arg_parser
 
 from ..base import CaprotoSaveIOC, check_args
+from ..utils import now, save_hdf5_zebra
+
+# def export_nano_zebra_data(zebra, filepath, fastaxis):
+#     j = 0
+#     while zebra.pc.data_in_progress.get() == 1:
+#         print("Waiting for zebra...")
+#         ttime.sleep(0.1)
+#         j += 1
+#         if j > 10:
+#             print("THE ZEBRA IS BEHAVING BADLY CARRYING ON")
+#             break
+
+#     time_d = zebra.pc.data.time.get()
+#     enc1_d = zebra.pc.data.enc1.get()
+#     enc2_d = zebra.pc.data.enc2.get()
+#     enc3_d = zebra.pc.data.enc3.get()
+
+#     px = zebra.pc.pulse_step.get()
+#     if fastaxis == 'NANOHOR':
+#         # Add half pixelsize to correct encoder
+#         enc1_d = enc1_d + (px / 2)
+#     elif fastaxis == 'NANOVER':
+#         # Add half pixelsize to correct encoder
+#         enc2_d = enc2_d + (px / 2)
+#     elif fastaxis == 'NANOZ':
+#         # Add half pixelsize to correct encoder
+#         enc3_d = enc3_d + (px / 2)
+
+#     size = (len(time_d),)
+#     with h5py.File(filepath, "w") as f:
+#         dset0 = f.create_dataset("zebra_time", size, dtype="f")
+#         dset0[...] = np.array(time_d)
+#         dset1 = f.create_dataset("enc1", size, dtype="f")
+#         dset1[...] = np.array(enc1_d)
+#         dset2 = f.create_dataset("enc2", size, dtype="f")
+#         dset2[...] = np.array(enc2_d)
+#         dset3 = f.create_dataset("enc3", size, dtype="f")
+#         dset3[...] = np.array(enc3_d)
+
+
+# class ZebraPositionCaptureData(Device):
+#     """
+#     Data arrays for the Zebra position capture function and their metadata.
+#     """
+#     # Data arrays
+#     ...
+#     enc1 = Cpt(EpicsSignal, "PC_ENC1")  # XF:05IDD-ES:1{Dev:Zebra2}:PC_ENC1
+#     enc2 = Cpt(EpicsSignal, "PC_ENC2")  # XF:05IDD-ES:1{Dev:Zebra2}:PC_ENC2
+#     enc3 = Cpt(EpicsSignal, "PC_ENC3")  # XF:05IDD-ES:1{Dev:Zebra2}:PC_ENC3
+#     time = Cpt(EpicsSignal, "PC_TIME")  # XF:05IDD-ES:1{Dev:Zebra2}:PC_TIME
+#     ...
+
+# class ZebraPositionCapture(Device):
+#     """
+#     Signals for the position capture function of the Zebra
+#     """
+
+#     # Configuration settings and status PVs
+#     ...
+#     pulse_step = Cpt(EpicsSignalWithRBV, "PC_PULSE_STEP")  # XF:05IDD-ES:1{Dev:Zebra2}:PC_PULSE_STEP
+#     ...
+#     data_in_progress = Cpt(EpicsSignalRO, "ARRAY_ACQ")  # XF:05IDD-ES:1{Dev:Zebra2}:ARRAY_ACQ
+#     ...
+#     data = Cpt(ZebraPositionCaptureData, "")
+
+# nanoZebra = SRXZebra(
+#    "XF:05IDD-ES:1{Dev:Zebra2}:", name="nanoZebra",
+#    read_attrs=["pc.data.enc1", "pc.data.enc2", "pc.data.enc3", "pc.data.time"],
+# )
+
+DEFAULT_MAX_LENGTH = 100_000
+
+
+class DevTypes(Enum):
+    """Enum class for devices."""
+
+    ZEBRA = "zebra"
+    SCALER = "scaler"
 
 
 class ZebraSaveIOC(CaprotoSaveIOC):
     """Zebra caproto save IOC."""
 
+    dev_type = pvproperty(
+        value=DevTypes.ZEBRA.value,
+        enum_strings=[x.value for x in DevTypes],
+        dtype=ChannelType.ENUM,
+        doc="Pick device type",
+    )
+
+    enc1 = pvproperty(
+        value=0,
+        doc="enc1 data",
+        max_length=DEFAULT_MAX_LENGTH,
+    )
+
+    enc2 = pvproperty(
+        value=0,
+        doc="enc2 data",
+        max_length=DEFAULT_MAX_LENGTH,
+    )
+
+    enc3 = pvproperty(
+        value=0,
+        doc="enc3 data",
+        max_length=DEFAULT_MAX_LENGTH,
+    )
+
+    zebra_time = pvproperty(
+        value=0,
+        doc="zebra time",
+        max_length=DEFAULT_MAX_LENGTH,
+    )
+
+    i0 = pvproperty(
+        value=0,
+        doc="i0 data",
+        max_length=DEFAULT_MAX_LENGTH,
+    )
+
+    im = pvproperty(
+        value=0,
+        doc="im data",
+        max_length=DEFAULT_MAX_LENGTH,
+    )
+
+    it = pvproperty(
+        value=0,
+        doc="it data",
+        max_length=DEFAULT_MAX_LENGTH,
+    )
+
+    sis_time = pvproperty(
+        value=0,
+        doc="sis time",
+        max_length=DEFAULT_MAX_LENGTH,
+    )
+
+    # def __init__(self, *args, external_pvs=None, **kwargs):
+    #     """Init method.
+
+    #     external_pvs : dict
+    #         a dictionary of external PVs with keys as human-readable names.
+    #     """
+    #     super().__init__(*args, **kwargs)
+    #     self._external_pvs = external_pvs
+
+    async def _get_current_dataset(
+        self, *args, **kwargs
+    ):  # , frame, external_pv="enc1"):
+        # client_context = Context()
+        # (pvobject,) = await client_context.get_pvs(self._external_pvs[external_pv])
+        # print(f"{pvobject = }")
+        # # pvobject = pvobjects[0]
+        # ret = await pvobject.read()
+
+        if self.dev_type == DevTypes.ZEBRA:
+            pvnames = ["enc1", "enc2", "enc3", "zebra_time"]
+        else:
+            pvnames = ["i0", "im", "it", "sis_time"]
+
+        dataset = {}
+        for pvname in pvnames:
+            dataset[pvname] = getattr(self, pvname).value
+
+        print(f"{now()}:\n{dataset}")
+
+        return dataset
+
+    @staticmethod
+    def saver(request_queue, response_queue):
+        """The saver callback for threading-based queueing."""
+        while True:
+            received = request_queue.get()
+            filename = received["filename"]
+            data = received["data"]
+            # 'frame_number' is not used for this exporter.
+            frame_number = received["frame_number"]  # noqa: F841
+            try:
+                save_hdf5_zebra(fname=filename, data=data, mode="x")
+                print(f"{now()}: saved data into:\n  {filename}")
+
+                success = True
+                error_message = ""
+            except Exception as exc:  # pylint: disable=broad-exception-caught
+                success = False
+                error_message = exc
+                print(
+                    f"Cannot save file {filename!r} due to the following exception:\n{exc}"
+                )
+
+            response = {"success": success, "error_message": error_message}
+            response_queue.put(response)
+
 
 if __name__ == "__main__":
     parser, split_args = template_arg_parser(
         default_prefix="", desc=textwrap.dedent(ZebraSaveIOC.__doc__)
     )
     ioc_options, run_options = check_args(parser, split_args)
-    ioc = ZebraSaveIOC(**ioc_options)  # TODO: pass libca IOC PVs of interest
+
+    # external_pv_prefix = (
+    #     ioc_options["prefix"].replace("{{", "{").replace("}}", "}")
+    # )  # "XF:05IDD-ES:1{Dev:Zebra2}:"
+
+    # external_pvs = {
+    #     "pulse_step": external_pv_prefix + "PC_PULSE_STEP",
+    #     "data_in_progress": external_pv_prefix + "ARRAY_ACQ",
+    #     "enc1": external_pv_prefix + "PC_ENC1",
+    #     "enc2": external_pv_prefix + "PC_ENC2",
+    #     "enc3": external_pv_prefix + "PC_ENC3",
+    #     "enc4": external_pv_prefix + "PC_ENC4",
+    #     "time": external_pv_prefix + "PC_TIME",
+    # }
+
+    # ioc = ZebraSaveIOC(external_pvs=external_pvs, **ioc_options)
+    ioc = ZebraSaveIOC(**ioc_options)
     run(ioc.pvdb, **run_options)