Include chunk shape as a parameter in stream resource for HDF dataset (…

…#544)

* Adding record for num frames in chunk along with chunk_size field in desc

* Attributes are saved all in a single chunk

* Update tests to account for chunk_size datakey parameter

* Chunk size should be in sres not desc

* Move chunk size to sres parameters

* Refactor tests to reflect changes

* chunk size can be int or none

Co-authored-by: Eugene <ymatviych@bnl.gov>

* Update chunk size signal to non-zero in one of the AD test sets

* Use correct chunk size for PandA, make sure we use chunk size auto

* Add comment on chunk size

* Make chunk_size a tuple that explicitly describes all chunk dims

* Make sure chunk size is tuple even with one dim, update tests, simplify ad standard det tests

* Make chunk_size always tuple of int, default to empty tuple

* Use readback value to avoid disconnect between actual value and signal get

* Follow import convention for tests

* Make use of slicing for detector name in ad_standard_det_factory clearer

* Rename chunk size to chunk shape

* Add space for linting

* Fix test

* Fix merge conflict

* Simplifying ad standard det factory fixture

* Fix unawaited task issue

* kinetix fixture doesn't need to be async

---------

Co-authored-by: Eugene <ymatviych@bnl.gov>

src/ophyd_async/core/_hdf_dataset.py

@@ -20,6 +20,8 @@ class HDFDataset:
  dtype_numpy: str = ""
  multiplier: int = 1
  swmr: bool = False
+ # Represents explicit chunk size written to disk.
+ chunk_shape: tuple[int, ...] = ()
 
 
 SLICE_NAME = "AD_HDF5_SWMR_SLICE"
@@ -66,6 +68,7 @@ def __init__(
  "dataset": ds.dataset,
  "swmr": ds.swmr,
  "multiplier": ds.multiplier,
+ "chunk_shape": ds.chunk_shape,
  },
  uid=None,
  validate=True,

src/ophyd_async/epics/adcore/_core_io.py

@@ -135,4 +135,6 @@ def __init__(self, prefix: str, name="") -> None:
  self.array_size0 = epics_signal_r(int, prefix + "ArraySize0")
  self.array_size1 = epics_signal_r(int, prefix + "ArraySize1")
  self.create_directory = epics_signal_rw(int, prefix + "CreateDirectory")
+ self.num_frames_chunks = epics_signal_r(int, prefix + "NumFramesChunks_RBV")
+ self.chunk_size_auto = epics_signal_rw_rbv(bool, prefix + "ChunkSizeAuto")
  super().__init__(prefix, name)

src/ophyd_async/epics/adcore/_hdf_writer.py

@@ -56,6 +56,9 @@ async def open(self, multiplier: int = 1) -> dict[str, DataKey]:
  # when directory path PV is processed.
  await self.hdf.create_directory.set(info.create_dir_depth)
 
+ # Make sure we are using chunk auto-sizing
+ await asyncio.gather(self.hdf.chunk_size_auto.set(True))
+
  await asyncio.gather(
  self.hdf.num_extra_dims.set(0),
  self.hdf.lazy_open.set(True),
@@ -84,6 +87,9 @@ async def open(self, multiplier: int = 1) -> dict[str, DataKey]:
  self._multiplier = multiplier
  outer_shape = (multiplier,) if multiplier > 1 else ()
 
+ # Determine number of frames that will be saved per HDF chunk
+ frames_per_chunk = await self.hdf.num_frames_chunks.get_value()
+
  # Add the main data
  self._datasets = [
  HDFDataset(
@@ -92,6 +98,7 @@ async def open(self, multiplier: int = 1) -> dict[str, DataKey]:
  shape=detector_shape,
  dtype_numpy=np_dtype,
  multiplier=multiplier,
+ chunk_shape=(frames_per_chunk, *detector_shape),
  )
  ]
  # And all the scalar datasets
@@ -118,6 +125,9 @@ async def open(self, multiplier: int = 1) -> dict[str, DataKey]:
  (),
  np_datatype,
  multiplier,
+ # NDAttributes appear to always be configured with
+ # this chunk size
+ chunk_shape=(16384,),
  )
  )
 

src/ophyd_async/fastcs/panda/_writer.py

@@ -105,7 +105,11 @@ async def _update_datasets(self) -> None:
 
  capture_table = await self.panda_data_block.datasets.get_value()
  self._datasets = [
- HDFDataset(dataset_name, "/" + dataset_name, [1], multiplier=1)
+ # TODO: Update chunk size to read signal once available in IOC
+ # Currently PandA IOC sets chunk size to 1024 points per chunk
+ HDFDataset(
+ dataset_name, "/" + dataset_name, [1], multiplier=1, chunk_shape=(1024,)
+ )
  for dataset_name in capture_table["name"]
  ]
 

tests/epics/adaravis/test_aravis.py

@@ -1,11 +1,9 @@
 import re
 
 import pytest
-from bluesky.run_engine import RunEngine
 
 from ophyd_async.core import (
  DetectorTrigger,
- DeviceCollector,
  PathProvider,
  TriggerInfo,
  set_mock_value,
@@ -14,14 +12,8 @@
 
 
 @pytest.fixture
-async def test_adaravis(
- RE: RunEngine,
- static_path_provider: PathProvider,
-) -> adaravis.AravisDetector:
- async with DeviceCollector(mock=True):
- test_adaravis = adaravis.AravisDetector("ADARAVIS:", static_path_provider)
-
- return test_adaravis
+def test_adaravis(ad_standard_det_factory) -> adaravis.AravisDetector:
+ return ad_standard_det_factory(adaravis.AravisDetector)
 
 
 @pytest.mark.parametrize("exposure_time", [0.0, 0.1, 1.0, 10.0, 100.0])
@@ -80,7 +72,7 @@ def test_gpio_pin_limited(test_adaravis: adaravis.AravisDetector):
 
 
 async def test_hints_from_hdf_writer(test_adaravis: adaravis.AravisDetector):
- assert test_adaravis.hints == {"fields": ["test_adaravis"]}
+ assert test_adaravis.hints == {"fields": ["test_adaravis1"]}
 
 
 async def test_can_read(test_adaravis: adaravis.AravisDetector):
@@ -98,9 +90,9 @@ async def test_decribe_describes_writer_dataset(
  await test_adaravis.stage()
  await test_adaravis.prepare(one_shot_trigger_info)
  assert await test_adaravis.describe() == {
- "test_adaravis": {
- "source": "mock+ca://ADARAVIS:HDF1:FullFileName_RBV",
- "shape": (0, 0),
+ "test_adaravis1": {
+ "source": "mock+ca://ARAVIS1:HDF1:FullFileName_RBV",
+ "shape": (10, 10),
  "dtype": "array",
  "dtype_numpy": "|i1",
  "external": "STREAM:",
@@ -125,12 +117,13 @@ async def test_can_collect(
  assert docs[0][0] == "stream_resource"
  stream_resource = docs[0][1]
  sr_uid = stream_resource["uid"]
- assert stream_resource["data_key"] == "test_adaravis"
+ assert stream_resource["data_key"] == "test_adaravis1"
  assert stream_resource["uri"] == "file://localhost" + str(full_file_name)
  assert stream_resource["parameters"] == {
  "dataset": "/entry/data/data",
  "swmr": False,
  "multiplier": 1,
+ "chunk_shape": (1, 10, 10),
  }
  assert docs[1][0] == "stream_datum"
  stream_datum = docs[1][1]
@@ -148,9 +141,9 @@ async def test_can_decribe_collect(
  await test_adaravis.stage()
  await test_adaravis.prepare(one_shot_trigger_info)
  assert (await test_adaravis.describe_collect()) == {
- "test_adaravis": {
- "source": "mock+ca://ADARAVIS:HDF1:FullFileName_RBV",
- "shape": (0, 0),
+ "test_adaravis1": {
+ "source": "mock+ca://ARAVIS1:HDF1:FullFileName_RBV",
+ "shape": (10, 10),
  "dtype": "array",
  "dtype_numpy": "|i1",
  "external": "STREAM:",

tests/epics/adcore/test_writers.py

@@ -46,6 +46,9 @@ async def hdf_writer_with_stats(
  hdf = adcore.NDFileHDFIO("HDF:")
  stats = adcore.NDPluginStatsIO("FOO:")
 
+ # Set number of frames per chunk to something reasonable
+ set_mock_value(hdf.num_frames_chunks, 2)
+
  return adcore.ADHDFWriter(
  hdf,
  static_path_provider,

tests/epics/adkinetix/test_kinetix.py

@@ -1,25 +1,18 @@
 import pytest
-from bluesky.run_engine import RunEngine
 
 from ophyd_async.core import (
  DetectorTrigger,
- DeviceCollector,
  StaticPathProvider,
  set_mock_value,
 )
 from ophyd_async.core._detector import TriggerInfo
 from ophyd_async.epics import adkinetix
+from ophyd_async.epics.adkinetix._kinetix_io import KinetixTriggerMode
 
 
 @pytest.fixture
-async def test_adkinetix(
- RE: RunEngine,
- static_path_provider: StaticPathProvider,
-) -> adkinetix.KinetixDetector:
- async with DeviceCollector(mock=True):
- test_adkinetix = adkinetix.KinetixDetector("KINETIX:", static_path_provider)
-
- return test_adkinetix
+def test_adkinetix(ad_standard_det_factory):
+ return ad_standard_det_factory(adkinetix.KinetixDetector)
 
 
 async def test_get_deadtime(
@@ -30,7 +23,7 @@ async def test_get_deadtime(
 
 
 async def test_trigger_modes(test_adkinetix: adkinetix.KinetixDetector):
- set_mock_value(test_adkinetix.drv.trigger_mode, "Internal")
+ set_mock_value(test_adkinetix.drv.trigger_mode, KinetixTriggerMode.internal)
 
  async def setup_trigger_mode(trig_mode: DetectorTrigger):
  await test_adkinetix.controller.prepare(
@@ -58,7 +51,7 @@ async def setup_trigger_mode(trig_mode: DetectorTrigger):
 
 
 async def test_hints_from_hdf_writer(test_adkinetix: adkinetix.KinetixDetector):
- assert test_adkinetix.hints == {"fields": ["test_adkinetix"]}
+ assert test_adkinetix.hints == {"fields": ["test_adkinetix1"]}
 
 
 async def test_can_read(test_adkinetix: adkinetix.KinetixDetector):
@@ -76,9 +69,9 @@ async def test_decribe_describes_writer_dataset(
  await test_adkinetix.stage()
  await test_adkinetix.prepare(one_shot_trigger_info)
  assert await test_adkinetix.describe() == {
- "test_adkinetix": {
- "source": "mock+ca://KINETIX:HDF1:FullFileName_RBV",
- "shape": (0, 0),
+ "test_adkinetix1": {
+ "source": "mock+ca://KINETIX1:HDF1:FullFileName_RBV",
+ "shape": (10, 10),
  "dtype": "array",
  "dtype_numpy": "|i1",
  "external": "STREAM:",
@@ -103,12 +96,13 @@ async def test_can_collect(
  assert docs[0][0] == "stream_resource"
  stream_resource = docs[0][1]
  sr_uid = stream_resource["uid"]
- assert stream_resource["data_key"] == "test_adkinetix"
+ assert stream_resource["data_key"] == "test_adkinetix1"
  assert stream_resource["uri"] == "file://localhost" + str(full_file_name)
  assert stream_resource["parameters"] == {
  "dataset": "/entry/data/data",
  "swmr": False,
  "multiplier": 1,
+ "chunk_shape": (1, 10, 10),
  }
  assert docs[1][0] == "stream_datum"
  stream_datum = docs[1][1]
@@ -126,9 +120,9 @@ async def test_can_decribe_collect(
  await test_adkinetix.stage()
  await test_adkinetix.prepare(one_shot_trigger_info)
  assert (await test_adkinetix.describe_collect()) == {
- "test_adkinetix": {
- "source": "mock+ca://KINETIX:HDF1:FullFileName_RBV",
- "shape": (0, 0),
+ "test_adkinetix1": {
+ "source": "mock+ca://KINETIX1:HDF1:FullFileName_RBV",
+ "shape": (10, 10),
  "dtype": "array",
  "dtype_numpy": "|i1",
  "external": "STREAM:",

tests/epics/advimba/test_vimba.py

@@ -1,25 +1,22 @@
 import pytest
-from bluesky.run_engine import RunEngine
 
 from ophyd_async.core import (
  DetectorTrigger,
- DeviceCollector,
  PathProvider,
  set_mock_value,
 )
 from ophyd_async.core._detector import TriggerInfo
 from ophyd_async.epics import advimba
+from ophyd_async.epics.advimba._vimba_io import (
+ VimbaExposeOutMode,
+ VimbaOnOff,
+ VimbaTriggerSource,
+)
 
 
 @pytest.fixture
-async def test_advimba(
- RE: RunEngine,
- static_path_provider: PathProvider,
-) -> advimba.VimbaDetector:
- async with DeviceCollector(mock=True):
- test_advimba = advimba.VimbaDetector("VIMBA:", static_path_provider)
-
- return test_advimba
+def test_advimba(ad_standard_det_factory) -> advimba.VimbaDetector:
+ return ad_standard_det_factory(advimba.VimbaDetector)
 
 
 async def test_get_deadtime(
@@ -30,9 +27,9 @@ async def test_get_deadtime(
 
 
 async def test_arming_trig_modes(test_advimba: advimba.VimbaDetector):
- set_mock_value(test_advimba.drv.trigger_source, "Freerun")
- set_mock_value(test_advimba.drv.trigger_mode, "Off")
- set_mock_value(test_advimba.drv.exposure_mode, "Timed")
+ set_mock_value(test_advimba.drv.trigger_source, VimbaTriggerSource.freerun)
+ set_mock_value(test_advimba.drv.trigger_mode, VimbaOnOff.off)
+ set_mock_value(test_advimba.drv.exposure_mode, VimbaExposeOutMode.timed)
 
  async def setup_trigger_mode(trig_mode: DetectorTrigger):
  await test_advimba.controller.prepare(TriggerInfo(number=1, trigger=trig_mode))
@@ -68,7 +65,7 @@ async def setup_trigger_mode(trig_mode: DetectorTrigger):
 
 
 async def test_hints_from_hdf_writer(test_advimba: advimba.VimbaDetector):
- assert test_advimba.hints == {"fields": ["test_advimba"]}
+ assert test_advimba.hints == {"fields": ["test_advimba1"]}
 
 
 async def test_can_read(test_advimba: advimba.VimbaDetector):
@@ -86,9 +83,9 @@ async def test_decribe_describes_writer_dataset(
  await test_advimba.stage()
  await test_advimba.prepare(one_shot_trigger_info)
  assert await test_advimba.describe() == {
- "test_advimba": {
- "source": "mock+ca://VIMBA:HDF1:FullFileName_RBV",
- "shape": (0, 0),
+ "test_advimba1": {
+ "source": "mock+ca://VIMBA1:HDF1:FullFileName_RBV",
+ "shape": (10, 10),
  "dtype": "array",
  "dtype_numpy": "|i1",
  "external": "STREAM:",
@@ -113,12 +110,13 @@ async def test_can_collect(
  assert docs[0][0] == "stream_resource"
  stream_resource = docs[0][1]
  sr_uid = stream_resource["uid"]
- assert stream_resource["data_key"] == "test_advimba"
+ assert stream_resource["data_key"] == "test_advimba1"
  assert stream_resource["uri"] == "file://localhost" + str(full_file_name)
  assert stream_resource["parameters"] == {
  "dataset": "/entry/data/data",
  "swmr": False,
  "multiplier": 1,
+ "chunk_shape": (1, 10, 10),
  }
  assert docs[1][0] == "stream_datum"
  stream_datum = docs[1][1]
@@ -136,9 +134,9 @@ async def test_can_decribe_collect(
  await test_advimba.stage()
  await test_advimba.prepare(one_shot_trigger_info)
  assert (await test_advimba.describe_collect()) == {
- "test_advimba": {
- "source": "mock+ca://VIMBA:HDF1:FullFileName_RBV",
- "shape": (0, 0),
+ "test_advimba1": {
+ "source": "mock+ca://VIMBA1:HDF1:FullFileName_RBV",
+ "shape": (10, 10),
  "dtype": "array",
  "dtype_numpy": "|i1",
  "external": "STREAM:",

tests/epics/conftest.py

@@ -0,0 +1,38 @@
+from collections.abc import Callable
+
+import pytest
+from bluesky.run_engine import RunEngine
+
+from ophyd_async.core._detector import StandardDetector
+from ophyd_async.core._device import DeviceCollector
+from ophyd_async.core._mock_signal_utils import set_mock_value
+
+
+@pytest.fixture
+def ad_standard_det_factory(
+ RE: RunEngine,
+ static_path_provider,
+) -> Callable:
+ def generate_ad_standard_det(
+ ad_standard_detector_class, number=1
+ ) -> StandardDetector:
+ # Dynamically generate a name based on the class of detector
+ detector_name = ad_standard_detector_class.__name__
+ if detector_name.endswith("Detector"):
+ detector_name = detector_name[: -len("Detector")]
+
+ with DeviceCollector(mock=True):
+ test_adstandard_det = ad_standard_detector_class(
+ f"{detector_name.upper()}{number}:",
+ static_path_provider,
+ name=f"test_ad{detector_name.lower()}{number}",
+ )
+
+ # Set number of frames per chunk and frame dimensions to something reasonable
+ set_mock_value(test_adstandard_det.hdf.num_frames_chunks, 1)
+ set_mock_value(test_adstandard_det.drv.array_size_x, 10)
+ set_mock_value(test_adstandard_det.drv.array_size_y, 10)
+
+ return test_adstandard_det
+
+ return generate_ad_standard_det