Make emulated camera image and movie consider exposure

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

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