Add with_delay decorator

src/python_spectrometer/daq/simulator.py

@@ -25,7 +25,7 @@ from collections.abc import Callable
 from typing import Literal
 
 import numpy as np
-from qutil.functools import partial
+from qutil.functools import partial, wraps
 from qutil.math import cexp
 from qutil.signal_processing import real_space
 
@@ -47,6 +47,33 @@ except ImportError as e:
                       "'pip install qopt.'") from e
 
 
+def with_delay(meth):
+    """Wraps an acquisition generator to accept the *delay* kwarg."""
+
+    @wraps(meth)
+    def wrapped(self, *, delay: bool | float = False, **settings):
+        skip_wait = getattr(wrapped, '__with_delay', False)
+
+        if delay is True:
+            delay = 1 / settings['df']
+
+        it = meth(self, **settings)
+        while True:
+            tic = time.perf_counter()
+            try:
+                data = next(it)
+            except StopIteration as stop:
+                return stop.value
+            else:
+                if delay and not skip_wait:
+                    time.sleep(delay - (time.perf_counter() - tic))
+
+                yield data
+
+    wrapped.__with_delay = True
+    return wrapped
+
+
 class MonochromaticNoise(DAQ):
     """Generate monochromatic sinusoidal noise with random phase.
 
@@ -56,22 +83,16 @@ class MonochromaticNoise(DAQ):
     Inherits from the base DAQ class and implements the required acquire method.
     """
 
-    def acquire(self, *, n_avg: int, fs: float, n_pts: int, delay: bool | float = False,
-                A: float = 1, f_0: float = 50, **settings) -> AcquisitionGenerator[DAQ.DTYPE]:
+    @with_delay
+    def acquire(self, *, n_avg: int, fs: float, n_pts: int, A: float = 1, f_0: float = 50,
+                **settings) -> AcquisitionGenerator[DAQ.DTYPE]:
         """Generate sinusoidal data with random phase."""
 
-        if delay is True:
-            delay = 1 / settings['df']
-
         t = np.arange(0, n_pts / fs, 1 / fs)
         rng = np.random.default_rng()
 
         for _ in range(n_avg):
-            tic = time.perf_counter()
-            data = np.sin(2 * np.pi * (t * f_0 + rng.random()))
-            if delay:
-                time.sleep(delay - (time.perf_counter() - tic))
-            yield data
+            yield np.sin(2 * np.pi * (t * f_0 + rng.random()))
 
 
 @dataclasses.dataclass
@@ -110,7 +131,8 @@ class QoptColoredNoise(DAQ):
         """White noise power spectral density with amplitude S_0."""
         return np.full_like(f, S_0)
 
-    def acquire(self, *, n_avg: int, fs: float, n_pts: int, delay: bool | float = False,
+    @with_delay
+    def acquire(self, *, n_avg: int, fs: float, n_pts: int,
                 **settings) -> AcquisitionGenerator[DAQ.DTYPE]:
         """Executes a measurement and yields the resulting timetrace.
 
@@ -119,21 +141,14 @@ class QoptColoredNoise(DAQ):
         actually acquire data with the given settings; if float delay
         by the given amount.
         """
-        if delay is True:
-            delay = 1 / settings['df']
-
         for _ in range(n_avg):
-            tic = time.perf_counter()
-            data = qopt.noise.fast_colored_noise(
+            yield qopt.noise.fast_colored_noise(
                 partial(
                     settings.get('spectral_density', self.spectral_density),
                     **settings
                 ),
                 dt=1/fs, n_samples=n_pts, output_shape=()
             )
-            if delay:
-                time.sleep(delay - (time.perf_counter() - tic))
-            yield data
 
         # This is the place to return metadata (possibly obtained from the instrument)
         return {'qopt version': qopt.__version__}
@@ -175,8 +190,9 @@ class DemodulatorQoptColoredNoise(QoptColoredNoise):
         settings.pop('f_min', None)
         return real_space.RC_filter(signal * IQ, **settings)
 
+    @with_delay
     def acquire(self, *, n_avg: int, freq: float = 0, filter_order: int = 1,
-                delay: bool | float = False, modulate_signal: bool = False,
+                modulate_signal: bool = False,
                 filter_method: Literal['forward', 'forward-backward'] = 'forward',
                 **settings) -> AcquisitionGenerator[DTYPE]:
         r"""Simulate demodulated noisy data.
@@ -234,7 +250,7 @@ class DemodulatorQoptColoredNoise(QoptColoredNoise):
         yield from (
             self.demodulate(IQ.real + data if modulate_signal else data, IQ,
                             order=filter_order, method=filter_method, **settings)
-            for data in super().acquire(n_avg=n_avg, delay=delay, **settings)
+            for data in super().acquire(n_avg=n_avg, **settings)
         )
 
         return {'qopt version': qopt.__version__}