From ab4dc460cfc5a24fabee67323c54b342b3d5af78 Mon Sep 17 00:00:00 2001 From: DiamondJoseph <53935796+DiamondJoseph@users.noreply.github.com> Date: Wed, 4 Sep 2024 12:22:13 +0100 Subject: [PATCH] Refactor ShapeProvider into DatasetDescriber (#555) * add np_datatype property instead of using final axis --- docs/examples/foo_detector.py | 2 +- src/ophyd_async/core/__init__.py | 4 ++-- src/ophyd_async/core/_providers.py | 8 ++++++-- src/ophyd_async/epics/adaravis/_aravis.py | 2 +- src/ophyd_async/epics/adcore/__init__.py | 4 ++-- src/ophyd_async/epics/adcore/_core_logic.py | 11 +++++++---- src/ophyd_async/epics/adcore/_hdf_writer.py | 20 +++++++------------- src/ophyd_async/epics/adkinetix/_kinetix.py | 2 +- src/ophyd_async/epics/adpilatus/_pilatus.py | 2 +- src/ophyd_async/epics/adsimdetector/_sim.py | 2 +- src/ophyd_async/epics/advimba/_vimba.py | 2 +- tests/epics/adcore/test_scans.py | 2 +- tests/epics/adcore/test_writers.py | 16 ++++++++-------- 13 files changed, 39 insertions(+), 38 deletions(-) diff --git a/docs/examples/foo_detector.py b/docs/examples/foo_detector.py index 66948398b2..bf3e41e6f6 100644 --- a/docs/examples/foo_detector.py +++ b/docs/examples/foo_detector.py @@ -69,7 +69,7 @@ def __init__( self.hdf, path_provider, lambda: self.name, - adcore.ADBaseShapeProvider(self.drv), + adcore.ADBaseDatasetDescriber(self.drv), ), config_sigs=(self.drv.acquire_time,), name=name, diff --git a/src/ophyd_async/core/__init__.py b/src/ophyd_async/core/__init__.py index 71d8ae76d1..6b9c6ccac3 100644 --- a/src/ophyd_async/core/__init__.py +++ b/src/ophyd_async/core/__init__.py @@ -33,11 +33,11 @@ from ._providers import ( AutoIncrementFilenameProvider, AutoIncrementingPathProvider, + DatasetDescriber, FilenameProvider, NameProvider, PathInfo, PathProvider, - ShapeProvider, StaticFilenameProvider, StaticPathProvider, UUIDFilenameProvider, @@ -117,7 +117,7 @@ "NameProvider", "PathInfo", "PathProvider", - "ShapeProvider", + "DatasetDescriber", "StaticFilenameProvider", "StaticPathProvider", "UUIDFilenameProvider", diff --git a/src/ophyd_async/core/_providers.py b/src/ophyd_async/core/_providers.py index 8fe95bf7c4..ffb4565c90 100644 --- a/src/ophyd_async/core/_providers.py +++ b/src/ophyd_async/core/_providers.py @@ -211,7 +211,11 @@ def __call__(self) -> str: """Get the name to be used as a data_key in the descriptor document""" -class ShapeProvider(Protocol): +class DatasetDescriber(Protocol): @abstractmethod - async def __call__(self) -> tuple: + async def np_datatype(self) -> str: + """Represents the numpy datatype""" + + @abstractmethod + async def shape(self) -> tuple[int, ...]: """Get the shape of the data collection""" diff --git a/src/ophyd_async/epics/adaravis/_aravis.py b/src/ophyd_async/epics/adaravis/_aravis.py index db7255f2fd..3a5b9634d7 100644 --- a/src/ophyd_async/epics/adaravis/_aravis.py +++ b/src/ophyd_async/epics/adaravis/_aravis.py @@ -37,7 +37,7 @@ def __init__( self.hdf, path_provider, lambda: self.name, - adcore.ADBaseShapeProvider(self.drv), + adcore.ADBaseDatasetDescriber(self.drv), ), config_sigs=(self.drv.acquire_time,), name=name, diff --git a/src/ophyd_async/epics/adcore/__init__.py b/src/ophyd_async/epics/adcore/__init__.py index b67db6930f..f0878443d8 100644 --- a/src/ophyd_async/epics/adcore/__init__.py +++ b/src/ophyd_async/epics/adcore/__init__.py @@ -7,7 +7,7 @@ ) from ._core_logic import ( DEFAULT_GOOD_STATES, - ADBaseShapeProvider, + ADBaseDatasetDescriber, set_exposure_time_and_acquire_period_if_supplied, start_acquiring_driver_and_ensure_status, ) @@ -31,7 +31,7 @@ "NDFileHDFIO", "NDPluginStatsIO", "DEFAULT_GOOD_STATES", - "ADBaseShapeProvider", + "ADBaseDatasetDescriber", "set_exposure_time_and_acquire_period_if_supplied", "start_acquiring_driver_and_ensure_status", "ADHDFWriter", diff --git a/src/ophyd_async/epics/adcore/_core_logic.py b/src/ophyd_async/epics/adcore/_core_logic.py index 5fa8974d29..21b07406fb 100644 --- a/src/ophyd_async/epics/adcore/_core_logic.py +++ b/src/ophyd_async/epics/adcore/_core_logic.py @@ -4,10 +4,11 @@ from ophyd_async.core import ( DEFAULT_TIMEOUT, AsyncStatus, + DatasetDescriber, DetectorControl, - ShapeProvider, set_and_wait_for_value, ) +from ophyd_async.epics.adcore._utils import convert_ad_dtype_to_np from ._core_io import ADBaseIO, DetectorState @@ -18,15 +19,17 @@ ) -class ADBaseShapeProvider(ShapeProvider): +class ADBaseDatasetDescriber(DatasetDescriber): def __init__(self, driver: ADBaseIO) -> None: self._driver = driver - async def __call__(self) -> tuple: + async def np_datatype(self) -> str: + return convert_ad_dtype_to_np(await self._driver.data_type.get_value()) + + async def shape(self) -> tuple[int, int]: shape = await asyncio.gather( self._driver.array_size_y.get_value(), self._driver.array_size_x.get_value(), - self._driver.data_type.get_value(), ) return shape diff --git a/src/ophyd_async/epics/adcore/_hdf_writer.py b/src/ophyd_async/epics/adcore/_hdf_writer.py index 484f19ab98..2f0e3a766e 100644 --- a/src/ophyd_async/epics/adcore/_hdf_writer.py +++ b/src/ophyd_async/epics/adcore/_hdf_writer.py @@ -8,12 +8,12 @@ from ophyd_async.core import ( DEFAULT_TIMEOUT, AsyncStatus, + DatasetDescriber, DetectorWriter, HDFDataset, HDFFile, NameProvider, PathProvider, - ShapeProvider, observe_value, set_and_wait_for_value, wait_for_value, @@ -22,7 +22,6 @@ from ._core_io import NDArrayBaseIO, NDFileHDFIO from ._utils import ( FileWriteMode, - convert_ad_dtype_to_np, convert_param_dtype_to_np, convert_pv_dtype_to_np, ) @@ -34,13 +33,13 @@ def __init__( hdf: NDFileHDFIO, path_provider: PathProvider, name_provider: NameProvider, - shape_provider: ShapeProvider, + dataset_describer: DatasetDescriber, *plugins: NDArrayBaseIO, ) -> None: self.hdf = hdf self._path_provider = path_provider self._name_provider = name_provider - self._shape_provider = shape_provider + self._dataset_describer = dataset_describer self._plugins = plugins self._capture_status: Optional[AsyncStatus] = None @@ -79,23 +78,18 @@ async def open(self, multiplier: int = 1) -> Dict[str, DataKey]: # Wait for it to start, stashing the status that tells us when it finishes self._capture_status = await set_and_wait_for_value(self.hdf.capture, True) name = self._name_provider() - detector_shape = tuple(await self._shape_provider()) + detector_shape = await self._dataset_describer.shape() + np_dtype = await self._dataset_describer.np_datatype() self._multiplier = multiplier outer_shape = (multiplier,) if multiplier > 1 else () - frame_shape = detector_shape[:-1] if len(detector_shape) > 0 else [] - dtype_numpy = ( - convert_ad_dtype_to_np(detector_shape[-1]) - if len(detector_shape) > 0 - else "" - ) # Add the main data self._datasets = [ HDFDataset( data_key=name, dataset="/entry/data/data", - shape=frame_shape, - dtype_numpy=dtype_numpy, + shape=detector_shape, + dtype_numpy=np_dtype, multiplier=multiplier, ) ] diff --git a/src/ophyd_async/epics/adkinetix/_kinetix.py b/src/ophyd_async/epics/adkinetix/_kinetix.py index 181969077b..1a1f8ae987 100644 --- a/src/ophyd_async/epics/adkinetix/_kinetix.py +++ b/src/ophyd_async/epics/adkinetix/_kinetix.py @@ -33,7 +33,7 @@ def __init__( self.hdf, path_provider, lambda: self.name, - adcore.ADBaseShapeProvider(self.drv), + adcore.ADBaseDatasetDescriber(self.drv), ), config_sigs=(self.drv.acquire_time,), name=name, diff --git a/src/ophyd_async/epics/adpilatus/_pilatus.py b/src/ophyd_async/epics/adpilatus/_pilatus.py index bb42471e3a..113b780fa2 100644 --- a/src/ophyd_async/epics/adpilatus/_pilatus.py +++ b/src/ophyd_async/epics/adpilatus/_pilatus.py @@ -47,7 +47,7 @@ def __init__( self.hdf, path_provider, lambda: self.name, - adcore.ADBaseShapeProvider(self.drv), + adcore.ADBaseDatasetDescriber(self.drv), ), config_sigs=(self.drv.acquire_time,), name=name, diff --git a/src/ophyd_async/epics/adsimdetector/_sim.py b/src/ophyd_async/epics/adsimdetector/_sim.py index 609ff99227..b69937705f 100644 --- a/src/ophyd_async/epics/adsimdetector/_sim.py +++ b/src/ophyd_async/epics/adsimdetector/_sim.py @@ -27,7 +27,7 @@ def __init__( self.hdf, path_provider, lambda: self.name, - adcore.ADBaseShapeProvider(self.drv), + adcore.ADBaseDatasetDescriber(self.drv), ), config_sigs=config_sigs, name=name, diff --git a/src/ophyd_async/epics/advimba/_vimba.py b/src/ophyd_async/epics/advimba/_vimba.py index 40f3f40b66..1dadfc3242 100644 --- a/src/ophyd_async/epics/advimba/_vimba.py +++ b/src/ophyd_async/epics/advimba/_vimba.py @@ -32,7 +32,7 @@ def __init__( self.hdf, path_provider, lambda: self.name, - adcore.ADBaseShapeProvider(self.drv), + adcore.ADBaseDatasetDescriber(self.drv), ), config_sigs=(self.drv.acquire_time,), name=name, diff --git a/tests/epics/adcore/test_scans.py b/tests/epics/adcore/test_scans.py index 87ad710e2e..3936885083 100644 --- a/tests/epics/adcore/test_scans.py +++ b/tests/epics/adcore/test_scans.py @@ -69,7 +69,7 @@ def writer(RE, static_path_provider, tmp_path: Path) -> adcore.ADHDFWriter: hdf, path_provider=static_path_provider, name_provider=lambda: "test", - shape_provider=AsyncMock(), + dataset_describer=AsyncMock(), ) diff --git a/tests/epics/adcore/test_writers.py b/tests/epics/adcore/test_writers.py index ad9ff6b0bc..601504beea 100644 --- a/tests/epics/adcore/test_writers.py +++ b/tests/epics/adcore/test_writers.py @@ -4,9 +4,9 @@ import pytest from ophyd_async.core import ( + DatasetDescriber, DeviceCollector, PathProvider, - ShapeProvider, StandardDetector, StaticPathProvider, ) @@ -16,12 +16,12 @@ from ophyd_async.plan_stubs._nd_attributes import setup_ndattributes, setup_ndstats_sum -class DummyShapeProvider(ShapeProvider): - def __init__(self) -> None: - pass +class DummyDatasetDescriber(DatasetDescriber): + async def np_datatype(self) -> str: + return " tuple: - return (10, 10, adcore.ADBaseDataType.UInt16) + async def shape(self) -> tuple[int, int]: + return (10, 10) @pytest.fixture @@ -35,7 +35,7 @@ async def hdf_writer( hdf, static_path_provider, lambda: "test", - DummyShapeProvider(), + DummyDatasetDescriber(), ) @@ -51,7 +51,7 @@ async def hdf_writer_with_stats( hdf, static_path_provider, lambda: "test", - DummyShapeProvider(), + DummyDatasetDescriber(), stats, )