diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 314a151af7..3536f89164 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -1,47 +1,44 @@ // For format details, see https://containers.dev/implementors/json_reference/ { - "name": "Python 3 Developer Container", - "build": { - "dockerfile": "Dockerfile", - "target": "build", - // Only upgrade pip, we will install the project below - "args": { - "PIP_OPTIONS": "--upgrade pip" - }, - }, - "remoteEnv": { - "DISPLAY": "${localEnv:DISPLAY}" - }, - // Set *default* container specific settings.json values on container create. - "settings": { - "python.defaultInterpreterPath": "/venv/bin/python" - }, - "customizations": { - "vscode": { - // Add the IDs of extensions you want installed when the container is created. - "extensions": [ - "ms-python.python", - "tamasfe.even-better-toml", - "redhat.vscode-yaml", - "ryanluker.vscode-coverage-gutters" - ] - } - }, - // Make sure the files we are mapping into the container exist on the host - "initializeCommand": "bash -c 'for i in $HOME/.inputrc; do [ -f $i ] || touch $i; done'", - "runArgs": [ - "--net=host", - "--security-opt=label=type:container_runtime_t" - ], - "mounts": [ - "source=${localEnv:HOME}/.ssh,target=/root/.ssh,type=bind", - "source=${localEnv:HOME}/.inputrc,target=/root/.inputrc,type=bind", - // map in home directory - not strictly necessary but useful - "source=${localEnv:HOME},target=${localEnv:HOME},type=bind,consistency=cached" - ], - // make the workspace folder the same inside and outside of the container - "workspaceMount": "source=${localWorkspaceFolder},target=${localWorkspaceFolder},type=bind", - "workspaceFolder": "${localWorkspaceFolder}", - // After the container is created, install the python project in editable form - "postCreateCommand": "pip install -e .[dev] --config-settings editable_mode=compat" + "name": "Python 3 Developer Container", + "build": { + "dockerfile": "Dockerfile", + "target": "build", + // Only upgrade pip, we will install the project below + "args": { + "PIP_OPTIONS": "--upgrade pip" + } + }, + "remoteEnv": { + "DISPLAY": "${localEnv:DISPLAY}" + }, + "customizations": { + "vscode": { + // Set *default* container specific settings.json values on container create. + // "settings": { + // "python.defaultInterpreterPath": "/venv/bin/python" + // }, + // Add the IDs of extensions you want installed when the container is created. + "extensions": [ + "ms-python.python", + "tamasfe.even-better-toml", + "redhat.vscode-yaml", + "ryanluker.vscode-coverage-gutters" + ] + } + }, + // Make sure the files we are mapping into the container exist on the host + "initializeCommand": "bash -c 'for i in $HOME/.inputrc; do [ -f $i ] || touch $i; done'", + "runArgs": ["--net=host", "--security-opt=label=type:container_runtime_t"], + "mounts": [ + "source=${localEnv:HOME}/.ssh,target=/root/.ssh,type=bind", + "source=${localEnv:HOME}/.inputrc,target=/root/.inputrc,type=bind", + // map in home directory - not strictly necessary but useful + "source=${localEnv:HOME},target=${localEnv:HOME},type=bind,consistency=cached" + ], + // make the workspace folder the same inside and outside of the container + "workspaceMount": "source=${localWorkspaceFolder},target=${localWorkspaceFolder},type=bind", + "workspaceFolder": "${localWorkspaceFolder}", + // After the container is created, install the python project in editable form + "postCreateCommand": "pip install -e .[dev] --config-settings editable_mode=compat" } diff --git a/pyproject.toml b/pyproject.toml index 8f7da05a67..e007a4435c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,8 +52,7 @@ dev = [ "pre-commit", "pydata-sphinx-theme>=0.12", "pyepics>=3.4.2", - "pyside6", - "pyside6-stubs", + "pyside6==6.6.2", "pytest", "pytest-asyncio", "pytest-cov", diff --git a/src/ophyd_async/epics/areadetector/writers/hdf_writer.py b/src/ophyd_async/epics/areadetector/writers/hdf_writer.py index a47638a260..ab1317aaf5 100644 --- a/src/ophyd_async/epics/areadetector/writers/hdf_writer.py +++ b/src/ophyd_async/epics/areadetector/writers/hdf_writer.py @@ -109,12 +109,17 @@ async def collect_stream_docs( await self.hdf.flush_now.set(True) if indices_written: if not self._file: + path = Path(await self.hdf.full_file_name.get_value()) self._file = _HDFFile( self._directory_provider(), # See https://github.com/bluesky/ophyd-async/issues/122 - Path(await self.hdf.full_file_name.get_value()), + path, self._datasets, ) + # stream resource says "here is a dataset", + # stream datum says "here are N frames in that stream resource", + # you get one stream resource and many stream datums per scan + for doc in self._file.stream_resources(): yield "stream_resource", doc for doc in self._file.stream_data(indices_written): diff --git a/src/ophyd_async/sim/__init__.py b/src/ophyd_async/sim/__init__.py new file mode 100644 index 0000000000..9540698461 --- /dev/null +++ b/src/ophyd_async/sim/__init__.py @@ -0,0 +1,11 @@ +from .pattern_generator import PatternGenerator +from .sim_pattern_detector_control import SimPatternDetectorControl +from .sim_pattern_detector_writer import SimPatternDetectorWriter +from .sim_pattern_generator import SimPatternDetector + +__all__ = [ + "PatternGenerator", + "SimPatternDetectorControl", + "SimPatternDetectorWriter", + "SimPatternDetector", +] diff --git a/src/ophyd_async/sim/pattern_generator.py b/src/ophyd_async/sim/pattern_generator.py new file mode 100644 index 0000000000..b6e5434eca --- /dev/null +++ b/src/ophyd_async/sim/pattern_generator.py @@ -0,0 +1,318 @@ +from dataclasses import dataclass +from pathlib import Path +from typing import ( + Any, + AsyncGenerator, + AsyncIterator, + Dict, + Iterator, + List, + Optional, + Sequence, +) + +import h5py +import numpy as np +from bluesky.protocols import Descriptor, StreamAsset +from event_model import ( + ComposeStreamResource, + ComposeStreamResourceBundle, + StreamDatum, + StreamRange, + StreamResource, +) + +from ophyd_async.core import DirectoryInfo, DirectoryProvider +from ophyd_async.core.signal import SignalR, observe_value +from ophyd_async.core.sim_signal_backend import SimSignalBackend +from ophyd_async.core.utils import DEFAULT_TIMEOUT + +# raw data path +DATA_PATH = "/entry/data/data" + +# pixel sum path +SUM_PATH = "/entry/sum" + +MAX_UINT8_VALUE = np.iinfo(np.uint8).max + +SLICE_NAME = "AD_HDF5_SWMR_SLICE" + + +@dataclass +class DatasetConfig: + name: str + shape: Sequence[int] + maxshape: tuple[Any, ...] = (None,) + path: Optional[str] = None + multiplier: Optional[int] = 1 + dtype: Optional[Any] = None + fillvalue: Optional[int] = None + + +def get_full_file_description( + datasets: List[DatasetConfig], outer_shape: tuple[int, ...] +): + full_file_description: Dict[str, Descriptor] = {} + for d in datasets: + source = f"sim://{d.name}" + shape = outer_shape + tuple(d.shape) + dtype = "number" if d.shape == [1] else "array" + descriptor = Descriptor( + source=source, shape=shape, dtype=dtype, external="STREAM:" + ) + key = d.name.replace("/", "_") + full_file_description[key] = descriptor + return full_file_description + + +def generate_gaussian_blob(height: int, width: int) -> np.ndarray: + """Make a Gaussian Blob with float values in range 0..1""" + x, y = np.meshgrid(np.linspace(-1, 1, width), np.linspace(-1, 1, height)) + d = np.sqrt(x * x + y * y) + blob = np.exp(-(d**2)) + return blob + + +def generate_interesting_pattern(x: float, y: float) -> float: + """This function is interesting in x and y in range -10..10, returning + a float value in range 0..1 + """ + z = 0.5 + (np.sin(x) ** 10 + np.cos(10 + y * x) * np.cos(x)) / 2 + return z + + +class HdfStreamProvider: + def __init__( + self, + directory_info: DirectoryInfo, + full_file_name: Path, + datasets: List[DatasetConfig], + ) -> None: + self._last_emitted = 0 + self._bundles: List[ComposeStreamResourceBundle] = self._compose_bundles( + directory_info, full_file_name, datasets + ) + + def _compose_bundles( + self, + directory_info: DirectoryInfo, + full_file_name: Path, + datasets: List[DatasetConfig], + ) -> List[StreamAsset]: + path = str(full_file_name.relative_to(directory_info.root)) + root = str(directory_info.root) + bundler_composer = ComposeStreamResource() + + bundles: List[ComposeStreamResourceBundle] = [] + + bundles = [ + bundler_composer( + spec=SLICE_NAME, + root=root, + resource_path=path, + data_key=d.name.replace("/", "_"), + resource_kwargs={ + "path": d.path, + "multiplier": d.multiplier, + "timestamps": "/entry/instrument/NDAttributes/NDArrayTimeStamp", + }, + ) + for d in datasets + ] + return bundles + + def stream_resources(self) -> Iterator[StreamResource]: + for bundle in self._bundles: + yield bundle.stream_resource_doc + + def stream_data(self, indices_written: int) -> Iterator[StreamDatum]: + # Indices are relative to resource + if indices_written > self._last_emitted: + updated_stream_range = StreamRange( + start=self._last_emitted, + stop=indices_written, + ) + self._last_emitted = indices_written + for bundle in self._bundles: + yield bundle.compose_stream_datum(indices=updated_stream_range) + return None + + def close(self) -> None: + for bundle in self._bundles: + bundle.close() + + +class PatternGenerator: + def __init__( + self, + saturation_exposure_time: float = 1, + detector_width: int = 320, + detector_height: int = 240, + ) -> None: + self.saturation_exposure_time = saturation_exposure_time + self.exposure = saturation_exposure_time + self.x = 0.0 + self.y = 0.0 + self.height = detector_height + self.width = detector_width + self.written_images_counter: int = 0 + + # it automatically initializes to 0 + self.signal_backend = SimSignalBackend(int, "sim://sim_images_counter") + self.sim_signal = SignalR(self.signal_backend) + blob = np.array( + generate_gaussian_blob(width=detector_width, height=detector_height) + * MAX_UINT8_VALUE + ) + self.STARTING_BLOB = blob + self._hdf_stream_provider: Optional[HdfStreamProvider] = None + self._handle_for_h5_file: Optional[h5py.File] = None + self.target_path: Optional[Path] = None + + async def write_image_to_file(self) -> None: + assert self._handle_for_h5_file, "no file has been opened!" + # prepare - resize the fixed hdf5 data structure + # so that the new image can be written + new_layer = self.written_images_counter + 1 + target_dimensions = (new_layer, self.height, self.width) + + # generate the simulated data + intensity: float = generate_interesting_pattern(self.x, self.y) + detector_data: np.uint8 = np.uint8( + self.STARTING_BLOB + * intensity + * self.exposure + / self.saturation_exposure_time + ) + + self._handle_for_h5_file[DATA_PATH].resize(target_dimensions) + + print(f"writing image {new_layer}") + assert self._handle_for_h5_file, "no file has been opened!" + self._handle_for_h5_file[DATA_PATH].resize(target_dimensions) + + self._handle_for_h5_file[SUM_PATH].resize((new_layer,)) + + # write data to disc (intermediate step) + self._handle_for_h5_file[DATA_PATH][self.written_images_counter] = detector_data + self._handle_for_h5_file[SUM_PATH][self.written_images_counter] = np.sum( + detector_data + ) + + # save metadata - so that it's discoverable + self._handle_for_h5_file[DATA_PATH].flush() + self._handle_for_h5_file[SUM_PATH].flush() + + # counter increment is last + # as only at this point the new data is visible from the outside + self.written_images_counter += 1 + await self.signal_backend.put(self.written_images_counter) + + def set_exposure(self, value: float) -> None: + self.exposure = value + + def set_x(self, value: float) -> None: + self.x = value + + def set_y(self, value: float) -> None: + self.y = value + + async def open_file( + self, directory: DirectoryProvider, multiplier: int = 1 + ) -> Dict[str, Descriptor]: + await self.sim_signal.connect() + + self.target_path = self._get_new_path(directory) + + self._handle_for_h5_file = h5py.File(self.target_path, "w", libver="latest") + + assert self._handle_for_h5_file, "not loaded the file right" + + datasets = self._get_datasets() + for d in datasets: + self._handle_for_h5_file.create_dataset( + name=d.name, + shape=d.shape, + dtype=d.dtype, + maxshape=d.maxshape, + ) + + # once datasets written, can switch the model to single writer multiple reader + self._handle_for_h5_file.swmr_mode = True + + outer_shape = (multiplier,) if multiplier > 1 else () + full_file_description = get_full_file_description(datasets, outer_shape) + + # cache state to self + self._datasets = datasets + self.multiplier = multiplier + self._directory_provider = directory + return full_file_description + + def _get_new_path(self, directory: DirectoryProvider) -> Path: + info = directory() + filename = f"{info.prefix}pattern{info.suffix}.h5" + new_path: Path = info.root / info.resource_dir / filename + return new_path + + def _get_datasets(self) -> List[DatasetConfig]: + raw_dataset = DatasetConfig( + # name=data_name, + name=DATA_PATH, + dtype=np.uint8, + shape=(1, self.height, self.width), + maxshape=(None, self.height, self.width), + ) + + sum_dataset = DatasetConfig( + name=SUM_PATH, + dtype=np.float64, + shape=(1,), + maxshape=(None,), + fillvalue=-1, + ) + + datasets: List[DatasetConfig] = [raw_dataset, sum_dataset] + return datasets + + async def collect_stream_docs( + self, indices_written: int + ) -> AsyncIterator[StreamAsset]: + """ + stream resource says "here is a dataset", + stream datum says "here are N frames in that stream resource", + you get one stream resource and many stream datums per scan + """ + if self._handle_for_h5_file: + self._handle_for_h5_file.flush() + # when already something was written to the file + if indices_written: + # if no frames arrived yet, there's no file to speak of + # cannot get the full filename the HDF writer will write + # until the first frame comes in + if not self._hdf_stream_provider: + assert self.target_path, "open file has not been called" + datasets = self._get_datasets() + self._datasets = datasets + self._hdf_stream_provider = HdfStreamProvider( + self._directory_provider(), + self.target_path, + self._datasets, + ) + for doc in self._hdf_stream_provider.stream_resources(): + yield "stream_resource", doc + if self._hdf_stream_provider: + for doc in self._hdf_stream_provider.stream_data(indices_written): + yield "stream_datum", doc + + def close(self) -> None: + if self._handle_for_h5_file: + self._handle_for_h5_file.close() + print("file closed") + self._handle_for_h5_file = None + + async def observe_indices_written( + self, timeout=DEFAULT_TIMEOUT + ) -> AsyncGenerator[int, None]: + async for num_captured in observe_value(self.sim_signal, timeout=timeout): + yield num_captured // self.multiplier diff --git a/src/ophyd_async/sim/sim_pattern_detector_control.py b/src/ophyd_async/sim/sim_pattern_detector_control.py new file mode 100644 index 0000000000..ff5300c4cd --- /dev/null +++ b/src/ophyd_async/sim/sim_pattern_detector_control.py @@ -0,0 +1,55 @@ +import asyncio +from typing import Optional + +from ophyd_async.core import DirectoryProvider +from ophyd_async.core.async_status import AsyncStatus +from ophyd_async.core.detector import DetectorControl, DetectorTrigger +from ophyd_async.sim.pattern_generator import PatternGenerator + + +class SimPatternDetectorControl(DetectorControl): + def __init__( + self, + pattern_generator: PatternGenerator, + directory_provider: DirectoryProvider, + exposure: float = 0.1, + ) -> None: + self.pattern_generator: PatternGenerator = pattern_generator + self.pattern_generator.set_exposure(exposure) + self.directory_provider: DirectoryProvider = directory_provider + self.task: Optional[asyncio.Task] = None + super().__init__() + + async def arm( + self, + num: int, + trigger: DetectorTrigger = DetectorTrigger.internal, + exposure: Optional[float] = 0.01, + ) -> AsyncStatus: + assert exposure is not None + period: float = exposure + self.get_deadtime(exposure) + task = asyncio.create_task( + self._coroutine_for_image_writing(exposure, period, num) + ) + self.task = task + return AsyncStatus(task) + + async def disarm(self): + if self.task: + self.task.cancel() + try: + await self.task + except asyncio.CancelledError: + pass + self.task = None + + def get_deadtime(self, exposure: float) -> float: + return 0.001 + + async def _coroutine_for_image_writing( + self, exposure: float, period: float, frames_number: int + ): + for _ in range(frames_number): + self.pattern_generator.set_exposure(exposure) + await asyncio.sleep(period) + await self.pattern_generator.write_image_to_file() diff --git a/src/ophyd_async/sim/sim_pattern_detector_writer.py b/src/ophyd_async/sim/sim_pattern_detector_writer.py new file mode 100644 index 0000000000..4f83aafa75 --- /dev/null +++ b/src/ophyd_async/sim/sim_pattern_detector_writer.py @@ -0,0 +1,34 @@ +from typing import AsyncGenerator, AsyncIterator, Dict + +from bluesky.protocols import Descriptor + +from ophyd_async.core import DirectoryProvider +from ophyd_async.core.detector import DetectorWriter +from ophyd_async.sim.pattern_generator import PatternGenerator + + +class SimPatternDetectorWriter(DetectorWriter): + pattern_generator: PatternGenerator + + def __init__( + self, pattern_generator: PatternGenerator, directoryProvider: DirectoryProvider + ) -> None: + self.pattern_generator = pattern_generator + self.directory_provider = directoryProvider + + async def open(self, multiplier: int = 1) -> Dict[str, Descriptor]: + return await self.pattern_generator.open_file( + self.directory_provider, multiplier + ) + + async def close(self) -> None: + self.pattern_generator.close() + + def collect_stream_docs(self, indices_written: int) -> AsyncIterator: + return self.pattern_generator.collect_stream_docs(indices_written) + + def observe_indices_written(self, timeout=...) -> AsyncGenerator[int, None]: + return self.pattern_generator.observe_indices_written() + + async def get_indices_written(self) -> int: + return self.pattern_generator.written_images_counter diff --git a/src/ophyd_async/sim/sim_pattern_generator.py b/src/ophyd_async/sim/sim_pattern_generator.py new file mode 100644 index 0000000000..3c2c54d8af --- /dev/null +++ b/src/ophyd_async/sim/sim_pattern_generator.py @@ -0,0 +1,37 @@ +from pathlib import Path +from typing import Sequence + +from ophyd_async.core import DirectoryProvider, StaticDirectoryProvider +from ophyd_async.core.detector import StandardDetector +from ophyd_async.core.signal import SignalR +from ophyd_async.sim.pattern_generator import PatternGenerator + +from .sim_pattern_detector_control import SimPatternDetectorControl +from .sim_pattern_detector_writer import SimPatternDetectorWriter + + +class SimPatternDetector(StandardDetector): + def __init__( + self, + path: Path, + config_sigs: Sequence[SignalR] = [], + name: str = "sim_pattern_detector", + writer_timeout: float = 1, + ) -> None: + self.directory_provider: DirectoryProvider = StaticDirectoryProvider(path) + self.pattern_generator = PatternGenerator() + writer = SimPatternDetectorWriter( + pattern_generator=self.pattern_generator, + directoryProvider=self.directory_provider, + ) + controller = SimPatternDetectorControl( + pattern_generator=self.pattern_generator, + directory_provider=self.directory_provider, + ) + super().__init__( + controller=controller, + writer=writer, + config_sigs=config_sigs, + name=name, + writer_timeout=writer_timeout, + ) diff --git a/tests/core/test_async_status.py b/tests/core/test_async_status.py index ce81254015..78e09abb4a 100644 --- a/tests/core/test_async_status.py +++ b/tests/core/test_async_status.py @@ -63,7 +63,7 @@ async def coroutine_to_wrap(time: float): await asyncio.sleep(time) -async def test_async_status_wrap(): +async def test_async_status_wrap() -> None: wrapped_coroutine = AsyncStatus.wrap(coroutine_to_wrap) status: AsyncStatus = wrapped_coroutine(0.01) diff --git a/tests/epics/areadetector/test_writers.py b/tests/epics/areadetector/test_writers.py index 0c9f4eb092..cad363af4b 100644 --- a/tests/epics/areadetector/test_writers.py +++ b/tests/epics/areadetector/test_writers.py @@ -28,8 +28,8 @@ async def hdf_writer(RE) -> HDFWriter: return HDFWriter( hdf, StaticDirectoryProvider("some_path", "some_prefix"), - lambda: "test", - DummyShapeProvider(), + name_provider=lambda: "test", + shape_provider=DummyShapeProvider(), ) diff --git a/tests/sim/__init__.py b/tests/sim/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/sim/conftest.py b/tests/sim/conftest.py new file mode 100644 index 0000000000..b3ba9301dd --- /dev/null +++ b/tests/sim/conftest.py @@ -0,0 +1,15 @@ +from pathlib import Path + +import pytest + +from ophyd_async.core.device import DeviceCollector +from ophyd_async.sim import SimPatternDetector + + +@pytest.fixture +async def sim_pattern_detector(tmp_path_factory) -> SimPatternDetector: + path: Path = tmp_path_factory.mktemp("tmp") + async with DeviceCollector(sim=True): + sim_pattern_detector = SimPatternDetector(name="PATTERN1", path=path) + + return sim_pattern_detector diff --git a/tests/sim/test_pattern_generator.py b/tests/sim/test_pattern_generator.py new file mode 100644 index 0000000000..7700816a2e --- /dev/null +++ b/tests/sim/test_pattern_generator.py @@ -0,0 +1,72 @@ +import h5py +import numpy as np +import pytest + +from ophyd_async.core import StaticDirectoryProvider +from ophyd_async.sim.pattern_generator import DATA_PATH, SUM_PATH, PatternGenerator + + +@pytest.fixture +async def pattern_generator(): + # path: Path = tmp_path_factory.mktemp("tmp") + pattern_generator = PatternGenerator() + yield pattern_generator + + +async def test_init(pattern_generator: PatternGenerator): + assert pattern_generator.exposure == 1 + assert pattern_generator.height == 240 + assert pattern_generator.width == 320 + assert pattern_generator.written_images_counter == 0 + assert pattern_generator._handle_for_h5_file is None + assert pattern_generator.STARTING_BLOB.shape == (240, 320) + + +def test_initialization(pattern_generator: PatternGenerator): + assert pattern_generator.saturation_exposure_time == 1 + assert pattern_generator.exposure == 1 + assert pattern_generator.x == 0.0 + assert pattern_generator.y == 0.0 + assert pattern_generator.height == 240 + assert pattern_generator.width == 320 + assert pattern_generator.written_images_counter == 0 + assert isinstance(pattern_generator.STARTING_BLOB, np.ndarray) + + +@pytest.mark.asyncio +async def test_open_and_close_file(tmp_path, pattern_generator: PatternGenerator): + dir_provider = StaticDirectoryProvider(str(tmp_path)) + await pattern_generator.open_file(dir_provider) + assert pattern_generator._handle_for_h5_file is not None + assert isinstance(pattern_generator._handle_for_h5_file, h5py.File) + pattern_generator.close() + assert pattern_generator._handle_for_h5_file is None + + +def test_set_exposure(pattern_generator: PatternGenerator): + pattern_generator.set_exposure(0.5) + assert pattern_generator.exposure == 0.5 + + +def test_set_x(pattern_generator: PatternGenerator): + pattern_generator.set_x(5.0) + assert pattern_generator.x == 5.0 + + +def test_set_y(pattern_generator: PatternGenerator): + pattern_generator.set_y(-3.0) + assert pattern_generator.y == -3.0 + + +@pytest.mark.asyncio +async def test_write_image_to_file(tmp_path, pattern_generator: PatternGenerator): + dir_provider = StaticDirectoryProvider(str(tmp_path)) + await pattern_generator.open_file(dir_provider) + + await pattern_generator.write_image_to_file() + assert pattern_generator.written_images_counter == 1 + assert pattern_generator._handle_for_h5_file + assert DATA_PATH in pattern_generator._handle_for_h5_file + assert SUM_PATH in pattern_generator._handle_for_h5_file + + pattern_generator.close() # Clean up diff --git a/tests/sim/test_sim_detector.py b/tests/sim/test_sim_detector.py new file mode 100644 index 0000000000..808dffee90 --- /dev/null +++ b/tests/sim/test_sim_detector.py @@ -0,0 +1,52 @@ +import pytest + +from ophyd_async.core.device import DeviceCollector +from ophyd_async.epics.motion import motor +from ophyd_async.sim.sim_pattern_generator import SimPatternDetector + + +@pytest.fixture +async def sim_motor(): + async with DeviceCollector(sim=True): + sim_motor = motor.Motor("test") + return sim_motor + + +async def test_sim_pattern_detector_initialization( + sim_pattern_detector: SimPatternDetector, +): + assert ( + sim_pattern_detector.pattern_generator + ), "PatternGenerator was not initialized correctly." + + +async def test_detector_creates_controller_and_writer( + sim_pattern_detector: SimPatternDetector, +): + assert sim_pattern_detector.writer + assert sim_pattern_detector.controller + + +async def test_writes_pattern_to_file( + sim_pattern_detector: SimPatternDetector, sim_motor: motor.Motor, tmp_path +): + sim_pattern_detector = SimPatternDetector( + config_sigs=[*sim_motor._read_signals], path=tmp_path + ) + + images_number = 2 + await sim_pattern_detector.controller.arm(num=images_number) + # assert that the file is created and non-empty + assert sim_pattern_detector.writer + + # assert that the file contains data in expected dimensions + + +async def test_set_x_and_y(sim_pattern_detector): + assert sim_pattern_detector.pattern_generator.x == 0 + sim_pattern_detector.pattern_generator.set_x(200) + assert sim_pattern_detector.pattern_generator.x == 200 + + +async def test_initial_blob(sim_pattern_detector: SimPatternDetector): + assert sim_pattern_detector.pattern_generator.STARTING_BLOB.any() diff --git a/tests/sim/test_sim_writer.py b/tests/sim/test_sim_writer.py new file mode 100644 index 0000000000..59ecd5a494 --- /dev/null +++ b/tests/sim/test_sim_writer.py @@ -0,0 +1,45 @@ +from unittest.mock import patch + +import pytest + +from ophyd_async.core import StaticDirectoryProvider +from ophyd_async.core.device import DeviceCollector +from ophyd_async.sim import PatternGenerator +from ophyd_async.sim.sim_pattern_detector_writer import SimPatternDetectorWriter + + +@pytest.fixture +async def writer(tmp_path) -> SimPatternDetectorWriter: + async with DeviceCollector(sim=True): + driver = PatternGenerator() + directory = StaticDirectoryProvider(tmp_path) + + return SimPatternDetectorWriter(driver, directory) + + +async def test_correct_descriptor_doc_after_open(writer: SimPatternDetectorWriter): + with patch("ophyd_async.core.signal.wait_for_value", return_value=None): + descriptor = await writer.open() + + assert descriptor == { + "_entry_data_data": { + "source": "sim:///entry/data/data", + "shape": (1, 240, 320), + "dtype": "array", + "external": "STREAM:", + }, + "_entry_sum": { + "source": "sim:///entry/sum", + "shape": (1,), + "dtype": "array", + "external": "STREAM:", + }, + } + + await writer.close() + + +async def test_collect_stream_docs(writer: SimPatternDetectorWriter): + await writer.open() + [item async for item in writer.collect_stream_docs(1)] + assert writer.pattern_generator._handle_for_h5_file diff --git a/tests/sim/test_streaming_plan.py b/tests/sim/test_streaming_plan.py new file mode 100644 index 0000000000..a077bcc919 --- /dev/null +++ b/tests/sim/test_streaming_plan.py @@ -0,0 +1,51 @@ +from collections import defaultdict +from typing import Dict + +from bluesky import plans as bp +from bluesky.run_engine import RunEngine + +from ophyd_async.sim.sim_pattern_generator import SimPatternDetector + + +def assert_emitted(docs: Dict[str, list], **numbers: int): + assert list(docs) == list(numbers) + assert {name: len(d) for name, d in docs.items()} == numbers + + +# NOTE the async operations with h5py are non-trival +# because of lack of native support for async operations +# see https://github.com/h5py/h5py/issues/837 +async def test_streaming_plan(RE: RunEngine, sim_pattern_detector: SimPatternDetector): + names = [] + docs = [] + + def append_and_print(name, doc): + names.append(name) + docs.append(doc) + + RE.subscribe(append_and_print) + + RE(bp.count([sim_pattern_detector], num=1)) + + print(names) + # NOTE - double resource because double stream + assert names == [ + "start", + "descriptor", + "stream_resource", + "stream_resource", + "stream_datum", + "stream_datum", + "event", + "stop", + ] + await sim_pattern_detector.writer.close() + + +async def test_plan(RE: RunEngine, sim_pattern_detector: SimPatternDetector): + docs = defaultdict(list) + RE(bp.count([sim_pattern_detector]), lambda name, doc: docs[name].append(doc)) + assert_emitted( + docs, start=1, descriptor=1, stream_resource=2, stream_datum=2, event=1, stop=1 + ) + await sim_pattern_detector.writer.close()