Add DemodulatorQoptColoredNoise

src/python_spectrometer/daq/simulator.py

@@ -17,13 +17,17 @@ Add an artificial time delay to mimick finite data acquisition time:
 """
 from __future__ import annotations
 
+import copy
 import dataclasses
 import sys
 import time
-from typing import Callable
+from collections.abc import Callable
+from typing import Literal
 
 import numpy as np
 from qutil.functools import partial
+from qutil.math import cexp
+from qutil.signal_processing import real_space
 
 from .base import DAQ, AcquisitionGenerator
 
@@ -103,6 +107,77 @@ class QoptColoredNoise(DAQ):
         return {'qopt version': qopt.__version__}
 
 
+class DemodulatorQoptColoredNoise(QoptColoredNoise):
+    DTYPE = np.complexfloating
+
+    @staticmethod
+    def demodulate(signal: np.ndarray, IQ: np.ndarray, **settings):
+        # Don't highpass filter
+        settings = copy.deepcopy(settings)
+        settings.pop('f_min', None)
+        return real_space.RC_filter(signal * IQ, **settings)
+
+    def acquire(self, *, n_avg: int, freq: float = 0, filter_order: int = 1,
+                delay: bool | float = False, modulate_signal: bool = False,
+                filter_method: Literal['forward', 'forward-backward'] = 'forward',
+                **settings) -> AcquisitionGenerator[DTYPE]:
+        r"""Simulate demodulated noisy data.
+
+        See Ref. [1]_ for an introduction to Lock-in amplification.
+
+        Parameters
+        ----------
+        n_avg : int
+            Number of outer averages.
+        freq : float, optional
+            Modulation frequency.
+        filter_order : int, optional
+            RC filter order used to filter the demodulated signal.
+        delay : bool | float, optional
+            Simulate a realistic data acquisition duration.
+        modulate_signal : bool, optional
+            Add the simulated noise to the modulation signal to mimic
+            noise picked up by a Lock-In signal travelling through some
+            DUT. Otherwise, mimics the noise at the input of the
+            amplifier.
+
+            In other words, simulate a Lock-In output connected to an
+            input, or just simulate the input.
+
+            Note that if True, noise is assumed to be additive, that
+            is,
+
+            .. math::
+
+                x(t) = s(t) + \delta(t)
+
+            with $s(t)$ the output signal and $\delta(t)$ the noise.
+        filter_method :
+            See :func:`~qutil:qutil.signal_processing.real_space.RC_filter`.
+
+        Yields
+        ------
+        data :
+            Demodulated data in complex IQ-representation.
+
+        References
+        ----------
+        .. [1]: https://www.zhinst.com/europe/en/resources/principles-of-lock-in-detection
+
+        """
+        t = np.arange(0, settings['n_pts'] / settings['fs'], 1 / settings['fs'])
+        # demodulation by √2 exp(-iωt) (ZI convention)
+        IQ = np.sqrt(2) * cexp(-2 * np.pi * freq * t)
+
+        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)
+        )
+
+        return {'qopt version': qopt.__version__}
+
+
 @deprecated("Use QoptColoredNoise instead")
 class qopt_colored_noise(QoptColoredNoise):
     ...