diff --git a/src/instamatic-tem-emulator/simulation/camera.py b/src/instamatic-tem-emulator/simulation/camera.py index e56bbea..5c088ef 100644 --- a/src/instamatic-tem-emulator/simulation/camera.py +++ b/src/instamatic-tem-emulator/simulation/camera.py @@ -1,6 +1,7 @@ from __future__ import annotations -from typing import Tuple +import time +from typing import Tuple, Optional, Generator import numpy as np @@ -78,16 +79,33 @@ def get_image(self, exposure: float = None, binsize: int = None, **kwargs) -> np shape = (shape_x // binsize, shape_y // binsize) if mode == 'diff': - return self.stage.get_diffraction_pattern( - shape=shape, x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max - ) + image_factory = self.stage.get_diffraction_pattern else: - img = self.stage.get_image( - shape=shape, x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max - ) - # Add some noise - img += np.random.random_integers(1, 10, shape).astype(img.dtype) - return img + image_factory = self.stage.get_image + + t_start = time.perf_counter() + img = image_factory(shape=shape, x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max) + img += np.random.randint(low=1, high=10, size=shape, dtype=img.dtype) + + t_deadline = t_start + exposure + t_remaining = t_deadline - time.perf_counter() + if t_remaining > 0: + if t_remaining > 0.002: + time.sleep(t_remaining - 0.001) + while time.perf_counter() < t_deadline: + pass + + return img + + def get_movie( + self, + n_frames: int, + exposure: Optional[float] = None, + binsize: Optional[int] = None, + **kwargs, + ) -> Generator[np.ndarray, None, None]: + for _ in range(n_frames): + yield self.get_image(exposure=exposure, binsize=binsize, **kwargs) def _mag_to_ranges(self, mag: float) -> Tuple[float, float, float, float]: # assume 50x = 2mm full size diff --git a/src/instamatic-tem-emulator/start_server.py b/src/instamatic-tem-emulator/start_server.py index 78a731f..af9c9a3 100644 --- a/src/instamatic-tem-emulator/start_server.py +++ b/src/instamatic-tem-emulator/start_server.py @@ -1,9 +1,11 @@ import datetime +import inspect import logging import socket import threading import time import traceback +import uuid from argparse import ArgumentParser from dataclasses import dataclass, field from functools import partial @@ -20,6 +22,7 @@ from simulation.camera import CameraEmulator +_generators = {} stop_program_event = threading.Event() TEM_PORT = config.settings.tem_server_port @@ -121,6 +124,10 @@ def run(self) -> None: try: ret = self.evaluate(func_name, args, kwargs) status = 200 + if inspect.isgenerator(ret): + gen_id = uuid.uuid4().hex + _generators[gen_id] = ret + ret = {'__generator__': gen_id} except Exception as e: traceback.print_exc() self._device_kind.log.exception(e) @@ -133,16 +140,30 @@ def run(self) -> None: self._device_kind.log.info('Terminating ' + thread_desc) def evaluate(self, func_name: str, args: list, kwargs: dict) -> Any: - """Eval and call `self._device.func_name` with `args` and `kwargs`.""" + """Eval function `func_name` on `self.device` with `args` & `kwargs`.""" self._device_kind.log.debug(f'eval {func_name}, {args}, {kwargs}') + + if func_name == '__gen_next__': + gen = _generators[kwargs['id']] + try: + ret = next(gen) + SharedImageProxy.push(image=ret) + return {'name': NAME, 'shape': ret.shape, 'dtype': str(ret.dtype)} + except StopIteration: + del _generators[kwargs['id']] + return + + if func_name == "__gen_close__": + _generators.pop(kwargs['id'], None) + return + f = getattr(self.device, func_name) - try: - ret = f(*args, **kwargs) - except TypeError: # TypeError: 'attribute class' object is not callable - ret = f - if func_name in {'get_image', 'get_movie'}: + ret = f(*args, **kwargs) if callable(f) else f + + if func_name in {'get_image', }: SharedImageProxy.push(image=ret) ret = {'name': NAME, 'shape': ret.shape, 'dtype': str(ret.dtype)} + return ret