pyglaze 0.1.0__py3-none-any.whl

pyglaze/devtools/mock_device.py

@@ -0,0 +1,367 @@
+from __future__ import annotations
+
+import json
+import struct
+import time
+from abc import ABC, abstractmethod
+from enum import Enum, auto
+from pathlib import Path
+from time import sleep
+
+import numpy as np
+from serial import serialutil
+
+from pyglaze.device.configuration import (
+    DeviceConfiguration,
+    ForceDeviceConfiguration,
+    LeDeviceConfiguration,
+)
+
+
+class MockDevice(ABC):
+    """Base class for Mock devices for testing purposes."""
+
+    @abstractmethod
+    def __init__(
+        self: MockDevice,
+        fail_after: float = np.inf,
+        n_fails: float = np.inf,
+    ) -> None:
+        pass
+
+
+class ForceMockDevice(MockDevice):
+    """Mock device for devices using a FORCE lockin for testing purposes."""
+
+    CONFIG_PATH = Path(__file__).parents[1] / "devtools" / "mockup_config.json"
+    ENCODING: str = "utf-8"
+
+    def __init__(
+        self: ForceMockDevice,
+        fail_after: float = np.inf,
+        n_fails: float = np.inf,
+    ) -> None:
+        self.fail_after = fail_after
+        self.fails_wanted = n_fails
+        self.n_failures = 0
+        self.n_scans = 0
+        self.rng = np.random.default_rng()
+        self.valid_input = True
+        self.experiment_running = False
+
+        self._periods = None
+        self._frequency = None
+        self._sweep_length = None
+        if not self.CONFIG_PATH.exists():
+            conf = {"periods": None, "frequency": None, "sweep_length_ms": None}
+            self.__save(conf)
+
+    @property
+    def periods(self: ForceMockDevice) -> int:
+        """Number of integration periods."""
+        return int(self.__get_val("periods"))
+
+    @periods.setter
+    def periods(self: ForceMockDevice, value: int) -> None:
+        self.__set_val("periods", value)
+
+    @property
+    def frequency(self: ForceMockDevice) -> int:
+        """Modulation frequency in Hz."""
+        return int(self.__get_val("frequency"))
+
+    @frequency.setter
+    def frequency(self: ForceMockDevice, value: int) -> None:
+        self.__set_val("frequency", value)
+
+    @property
+    def sweep_length(self: ForceMockDevice) -> int:
+        """Total sweep length for one pulse."""
+        return int(self.__get_val("sweep_length"))
+
+    @sweep_length.setter
+    def sweep_length(self: ForceMockDevice, value: int) -> None:
+        self.__set_val("sweep_length", value)
+
+    def close(self: ForceMockDevice) -> None:
+        """Mock-close the serial connection."""
+
+    def write(self: ForceMockDevice, input_string: bytes) -> None:
+        """Mock-write to the serial connection."""
+        decoded_input_string = input_string.decode("utf-8")
+        split_input_string = decoded_input_string.split(",")
+        cmd = split_input_string[0]
+        if cmd == "!set timing":
+            self.periods = int(split_input_string[1])
+            self.frequency = int(split_input_string[2])
+        elif cmd == "!set sweep length":
+            self.sweep_length = int(split_input_string[1])
+        elif cmd == "!s":
+            time_pr_point = self.periods / self.frequency
+            self.in_waiting = int(self.sweep_length * 1e-3 / time_pr_point)
+            self.experiment_running = True
+        elif cmd in ["!lut", "!set wave", "!set generator"]:
+            pass
+        elif cmd != "!dat":
+            self.valid_input = False
+
+    def readline(self: ForceMockDevice) -> bytes:
+        """Mock-readline from the serial connection."""
+        if self.valid_input:
+            return self.__run_experiment() if self.experiment_running else b"!A,OK\r"
+
+        return b"A,FAULT\r"
+
+    def read_until(self: ForceMockDevice, expected: bytes = b"\r") -> bytes:  # noqa: ARG002
+        """Mock-read_until from the serial connection."""
+        random_datapoint = self.__create_random_datapoint
+        return random_datapoint.encode(self.ENCODING)
+
+    def __run_experiment(self: ForceMockDevice) -> bytes:
+        return_string = "!A,OK\\r"
+        return_string += (
+            f"!S,ip: {self.periods}, "
+            f"Freq: {self.frequency}, "
+            f"sl: {self.sweep_length}, "
+            f"from: 0.0000, "
+            f"to:1.0000\r"
+        )
+        for _ in range(self.in_waiting):
+            return_string += self.__create_random_datapoint
+        return_string += "!D,DONE\\r"
+        sleep(self.sweep_length * 1e-3)
+        self.n_scans += 1
+        if self.n_scans > self.fail_after and self.n_failures < self.fails_wanted:
+            self.n_failures += 1
+            msg = "MOCK_DEVICE: scan failed"
+            raise serialutil.SerialException(msg)
+
+        return return_string.encode("utf-8")
+
+    @property
+    def __create_random_datapoint(self: ForceMockDevice) -> str:
+        radius = self.rng.random() * 10
+        theta = self.rng.random() * 360
+        return f"!R,{radius},{theta}\r"
+
+    def __save(self: ForceMockDevice, conf: dict) -> None:
+        with self.CONFIG_PATH.open("w") as f:
+            json.dump(conf, f)
+
+    def __get_val(self: ForceMockDevice, key: str) -> int:
+        val: int = getattr(self, f"_{key}")
+        if not val:
+            with self.CONFIG_PATH.open() as f:
+                val = json.load(f)[key]
+        return val
+
+    def __set_val(self: ForceMockDevice, key: str, val: str | float) -> None:
+        with self.CONFIG_PATH.open() as f:
+            config = json.load(f)
+            config[key] = val
+        setattr(self, f"_{key}", val)
+        self.__save(config)
+
+
+class _LeMockState(Enum):
+    IDLE = auto()
+    WAITING_FOR_SETTINGS = auto()
+    WAITING_FOR_LIST = auto()
+    RECEIVED_SETTINGS = auto()
+    RECEIVED_LIST = auto()
+    RECEIVED_STATUS_REQUEST = auto()
+    STARTING_SCAN = auto()
+    SCANNING = auto()
+
+
+class LeMockDevice(MockDevice):
+    """Mock device for devices using a Le lockin for testing purposes."""
+
+    ENCODING = "utf-8"
+    LI_MODULATION_FREQUENCY = 10000
+    TIME_WINDOW = 100e-12
+    DAC_BITWIDTH = 2**12
+
+    def __init__(
+        self: LeMockDevice, fail_after: float = np.inf, n_fails: float = np.inf
+    ) -> None:
+        self.fail_after = fail_after
+        self.fails_wanted = n_fails
+        self.n_failures = 0
+        self.n_scans = 0
+        self.rng = np.random.default_rng()
+        self.state = _LeMockState.IDLE
+        self.is_scanning = False
+        self.n_scanning_points: int | None = None
+        self.integration_periods: int | None = None
+        self.use_ema: bool | None = None
+        self.scanning_list: list[int] | None = None
+        self._scan_start_time: float | None = None
+
+    def write(self: LeMockDevice, input_bytes: bytes) -> None:
+        """Mock-write to the serial connection."""
+        if self.state == _LeMockState.WAITING_FOR_SETTINGS:
+            self._handle_waiting_for_settings(input_bytes)
+            return
+        if self.state == _LeMockState.WAITING_FOR_LIST:
+            self._handle_waiting_for_list(input_bytes)
+            return
+        if self.state == _LeMockState.IDLE:
+            self._handle_idle(input_bytes)
+            return
+        if self.state == _LeMockState.SCANNING:
+            self._handle_scanning(input_bytes)
+            return
+
+        raise NotImplementedError
+
+    def read(self: LeMockDevice, size: int) -> bytes:
+        """Mock-read from the serial connection."""
+        if self.state == _LeMockState.IDLE:
+            return self._create_scan_bytes(n_bytes=size)
+        raise NotImplementedError
+
+    def read_until(self: LeMockDevice, _: bytes = b"\r") -> bytes:  # noqa: PLR0911
+        """Mock-read_until from the serial connection."""
+        if self.state == _LeMockState.WAITING_FOR_SETTINGS:
+            return "ACK: Ready to receive settings.".encode(self.ENCODING)
+        if self.state == _LeMockState.RECEIVED_SETTINGS:
+            self.state = _LeMockState.IDLE
+            return "ACK: Settings received.".encode(self.ENCODING)
+        if self.state == _LeMockState.WAITING_FOR_LIST:
+            return "ACK: Ready to receive list.".encode(self.ENCODING)
+        if self.state == _LeMockState.RECEIVED_LIST:
+            self.state = _LeMockState.IDLE
+            return "ACK: List received.".encode(self.ENCODING)
+        if self.state == _LeMockState.STARTING_SCAN:
+            self.state = _LeMockState.SCANNING
+            self.is_scanning = True
+            return "ACK: Scan started.".encode(self.ENCODING)
+        if self.state == _LeMockState.RECEIVED_STATUS_REQUEST:
+            if self._scan_has_finished():
+                self.state = _LeMockState.IDLE
+                return "ACK: Idle.".encode(self.ENCODING)
+
+            self.state = _LeMockState.SCANNING
+            return "Error: Scan is ongoing.".encode(self.ENCODING)
+
+        msg = f"Unknown state: {self.state}"
+        raise NotImplementedError(msg)
+
+    def close(self: LeMockDevice) -> None:
+        """Mock-close the serial connection."""
+
+    @property
+    def _scanning_time(self: LeMockDevice) -> float:
+        if self.n_scanning_points and self.integration_periods:
+            return (
+                self.n_scanning_points
+                * self.integration_periods
+                / self.LI_MODULATION_FREQUENCY
+            )
+        msg = "Cannot calculate scanning time when n_scanning_points or integration_periods is None"
+        raise ValueError(msg)
+
+    def _handle_idle(self: LeMockDevice, input_bytes: bytes) -> None:
+        msg = input_bytes.decode("utf-8")
+        if msg == "S":
+            self.state = _LeMockState.WAITING_FOR_SETTINGS
+        elif msg == "L":
+            self.state = _LeMockState.WAITING_FOR_LIST
+        elif msg == "G":
+            self.state = _LeMockState.STARTING_SCAN
+            self._scan_start_time = time.time()
+        elif msg == "R":
+            self._scan_has_finished()
+        else:
+            msg = f"Unknown message: {msg}"
+            raise NotImplementedError(msg)
+
+    def _handle_scanning(self: LeMockDevice, input_bytes: bytes) -> None:
+        msg = input_bytes.decode("utf-8")
+        if msg == "H":
+            self.state = _LeMockState.RECEIVED_STATUS_REQUEST
+            return
+        if msg == "R":
+            if self._scan_has_finished():
+                self.state = _LeMockState.IDLE
+            return
+
+        raise NotImplementedError
+
+    def _handle_waiting_for_settings(self: LeMockDevice, input_bytes: bytes) -> None:
+        ints = self._decode_ints(input_bytes)
+        self.n_scanning_points = ints[0]
+        self.integration_periods = ints[1]
+        self.use_ema = bool(ints[2])
+        self.state = _LeMockState.RECEIVED_SETTINGS
+
+    def _handle_waiting_for_list(self: LeMockDevice, input_bytes: bytes) -> None:
+        self.scanning_list = self._decode_ints(input_bytes)
+        self.state = _LeMockState.RECEIVED_LIST
+
+    def _decode_ints(self: LeMockDevice, input_bytes: bytes) -> list[int]:
+        # Convert every two bytes to a 16-bit integer (assuming little-endian format)
+        return [
+            struct.unpack("<H", input_bytes[i : i + 2])[0]
+            for i in range(0, len(input_bytes), 2)
+        ]
+
+    def _scan_has_finished(self: LeMockDevice) -> bool:
+        if not self.is_scanning:
+            return True
+        if self._scan_start_time is None:
+            msg = "Scan start time is None"
+            raise ValueError(msg)
+        scan_finished = time.time() - self._scan_start_time > self._scanning_time
+        if scan_finished:
+            self.is_scanning = False
+            self._scan_start_time = None
+        return scan_finished
+
+    def _create_scan_bytes(self: LeMockDevice, n_bytes: int) -> bytes:  # noqa: ARG002
+        if self.scanning_list is None:
+            msg = "Scanning list is None"
+            raise ValueError(msg)
+
+        self.n_scans += 1
+        if self.n_scans > self.fail_after and self.n_failures < self.fails_wanted:
+            self.n_failures += 1
+            numbers = np.array([])
+        else:
+            numbers = self.rng.random(2 * len(self.scanning_list))
+
+        # Each scanning point will generate an X and a Y value (lockin detection)
+        return struct.pack("<" + "f" * len(numbers), *numbers)
+
+
+def list_mock_devices() -> list[str]:
+    """List all available mock devices."""
+    return [
+        "mock_device",
+        "mock_device_scan_should_fail",
+        "mock_device_fail_first_scan",
+    ]
+
+
+def _mock_device_factory(config: DeviceConfiguration) -> MockDevice:
+    mock_class = _get_mock_class(config)
+    if config.amp_port == "mock_device_scan_should_fail":
+        return mock_class(fail_after=0)
+    if config.amp_port == "mock_device":
+        return mock_class()
+    if config.amp_port == "mock_device_fail_first_scan":
+        return mock_class(fail_after=0, n_fails=1)
+
+    msg = f"Unknown mock device requested: {config.amp_port}. Valid options are: {list_mock_devices()}"
+    raise ValueError(msg)
+
+
+def _get_mock_class(config: DeviceConfiguration) -> type[MockDevice]:
+    if isinstance(config, ForceDeviceConfiguration):
+        return ForceMockDevice
+    if isinstance(config, LeDeviceConfiguration):
+        return LeMockDevice
+
+    msg = f"Unsupported configuration type: {type(config).__name__}"
+    raise ValueError(msg)

pyglaze/devtools/thz_pulse.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from typing import cast
+
+import numpy as np
+
+from pyglaze.helpers.types import FloatArray
+
+
+def gaussian_derivative_pulse(
+    time: FloatArray,
+    t0: float,
+    sigma: float,
+    signal_to_noise: float | None = None,
+) -> FloatArray:
+    """Simulates a THz pulse as the derivative of a Gaussian.
+
+    Args:
+        time: Times to evaluate pulse at
+        t0: Center position of pulse
+        sigma: Standard deviation of Gaussian
+        signal_to_noise: Ratio between peak of pulse and standard deviation of noise
+
+    Returns:
+        Simulated pulse
+    """
+    signal: np.ndarray = (time - t0) * np.exp(-0.5 * ((time - t0) / sigma) ** 2)
+    noise = (
+        0.0
+        if signal_to_noise is None
+        else np.random.default_rng().normal(
+            scale=1.0 / signal_to_noise, size=len(signal)
+        )
+    )
+    return cast(FloatArray, signal / np.max(signal) + noise)

pyglaze/helpers/types.py

@@ -0,0 +1,20 @@
+from typing import Any, TypeVar, Union
+
+import numpy as np
+from typing_extensions import ParamSpec, TypeAlias
+
+# numpy typing does not work with pipe, hence Union instead
+FloatArray: TypeAlias = np.ndarray[Any, np.dtype[Union[np.float64, np.float32]]]
+ComplexArray: TypeAlias = Union[
+    np.ndarray[Any, np.dtype[Union[np.complex128, np.complex64]]], FloatArray
+]
+F = TypeVar("F", FloatArray, float)
+C = TypeVar("C", ComplexArray, complex)
+
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+JSONConvertible: TypeAlias = Union[
+    list["JSONConvertible"], dict[str, "JSONConvertible"], int, float, str, None
+]

pyglaze/helpers/utilities.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from functools import wraps
+from typing import TYPE_CHECKING, Callable, cast
+
+import serial
+import serial.tools.list_ports
+
+from pyglaze.helpers.types import P, T
+
+if TYPE_CHECKING:
+    import logging
+
+APP_NAME = "Glaze"
+LOGGER_NAME = "glaze-logger"
+
+
+def list_serial_ports() -> list[str]:
+    """Lists available serial ports for device conneciton.
+
+    Returns:
+        list[str]: Paths to available ports.
+    """
+    skip_ports_substrings = ["Bluetooth", "debug"]
+
+    ports = []
+    for port in serial.tools.list_ports.comports():
+        if any(substring in port.device for substring in skip_ports_substrings):
+            continue
+        ports.append(port.device)
+    return ports
+
+
+@dataclass
+class _BackoffRetry:
+    """Decorator for retrying a function, using exponential backoff, if it fails.
+
+    Args:
+        max_tries: The maximum number of times the function should be tried.
+        max_backoff: The maximum backoff time in seconds.
+        backoff_base: The base of the exponential backoff.
+        logger: A Logger class to use for logging. If None, messages are printed.
+
+    Returns:
+        The function that is decorated.
+    """
+
+    max_tries: int = 5
+    max_backoff: float = 5
+    backoff_base: float = 0.01
+    logger: logging.Logger | None = None
+
+    def __call__(self: _BackoffRetry, func: Callable[P, T]) -> Callable[P, T]:
+        """Try the function `max_tries` times, with exponential backoff if it fails."""
+
+        @wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+            for tries in range(self.max_tries - 1):
+                try:
+                    return cast(T, func(*args, **kwargs))
+                except (KeyboardInterrupt, SystemExit):
+                    raise
+                except Exception as e:  # noqa: BLE001
+                    self._log(
+                        f"{func.__name__} failed {tries+1} time(s) with: '{e}'. Trying again"
+                    )
+                backoff = min(self.backoff_base * 2**tries, self.max_backoff)
+                time.sleep(backoff)
+            self._log(f"{func.__name__}: Last try ({tries+2}).")
+            return cast(T, func(*args, **kwargs))
+
+        return wrapper
+
+    def _log(self: _BackoffRetry, msg: str) -> None:
+        if self.logger:
+            self.logger.warning(msg)
+        else:
+            pass

pyglaze/interpolation/__init__.py

@@ -0,0 +1,3 @@
+from .interpolation import ws_interpolate
+
+__all__ = ["ws_interpolate"]

pyglaze/interpolation/interpolation.py

@@ -0,0 +1,24 @@
+import numpy as np
+
+from pyglaze.helpers.types import FloatArray
+
+
+def ws_interpolate(
+    times: FloatArray, pulse: FloatArray, interp_times: FloatArray
+) -> FloatArray:
+    """Performs Whittaker-Shannon interpolation at the supplied times given a pulse.
+
+    Args:
+        times: Sampling times
+        pulse: A sampled pulse satisfying the Nyquist criterion
+        interp_times: Array of times at which to interpolate
+
+    Returns:
+        FloatArray: Interpolated values
+    """
+    dt = times[1] - times[0]
+    _range = np.arange(len(pulse))
+    # times must be zero-centered for formula to work
+    sinc = np.sinc((interp_times[:, np.newaxis] - times[0] - dt * _range) / dt)
+
+    return np.asarray(np.sum(pulse * sinc, axis=1))

pyglaze/scanning/__init__.py

@@ -0,0 +1,4 @@
+from .client import GlazeClient
+from .scanner import Scanner
+
+__all__ = ["GlazeClient", "Scanner"]

pyglaze/scanning/_asyncscanner.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from multiprocessing import Event, Pipe, Process, Queue, synchronize
+from queue import Empty, Full
+from typing import TYPE_CHECKING
+
+from serial import serialutil
+
+from pyglaze.datamodels.waveform import UnprocessedWaveform, _TimestampedWaveform
+from pyglaze.scanning.scanner import Scanner
+
+if TYPE_CHECKING:
+    import logging
+    from multiprocessing.connection import Connection
+
+    from pyglaze.device.configuration import DeviceConfiguration
+
+
+@dataclass
+class _ScannerHealth:
+    is_alive: bool
+    is_healthy: bool
+    error: Exception | None
+
+
+@dataclass
+class _AsyncScanner:
+    """Used by GlazeClient to starts a scanner in a new process and read scans from shared memory."""
+
+    queue_maxsize: int = 10
+    startup_timeout: float = 30.0
+    logger: logging.Logger | None = None
+    is_scanning: bool = False
+    _child_process: Process = field(init=False)
+    _shared_mem: Queue[_TimestampedWaveform] = field(init=False)
+    _SCAN_TIMEOUT: float = field(init=False)
+    _stop_signal: synchronize.Event = field(init=False)
+    _scanner_conn: Connection = field(init=False)
+
+    def start_scan(self: _AsyncScanner, config: DeviceConfiguration) -> None:
+        """Starts continuously scanning in new process.
+
+        Args:
+            config: Device configurtaion
+        """
+        self._SCAN_TIMEOUT = config._sweep_length_ms * 2e-3 + 1  # noqa: SLF001, access to private attribute for backwards compatibility
+        self._shared_mem = Queue(maxsize=self.queue_maxsize)
+        self._stop_signal = Event()
+        self._scanner_conn, child_conn = Pipe()
+        self._child_process = Process(
+            target=_AsyncScanner._run_scanner,
+            args=[config, self._shared_mem, self._stop_signal, child_conn],
+        )
+        self._child_process.start()
+
+        # Wait for scanner to start
+        if not self._scanner_conn.poll(timeout=self.startup_timeout):
+            self.stop_scan()
+            err_msg = "Scanner timed out"
+            raise TimeoutError(err_msg)
+
+        msg: _ScannerHealth = self._scanner_conn.recv()
+        if msg.is_healthy and msg.is_alive:
+            self.is_scanning = True
+        else:
+            self.stop_scan()
+
+        if msg.error:
+            if self.logger:
+                self.logger.error(str(msg.error))
+            raise msg.error
+
+    def stop_scan(self: _AsyncScanner) -> None:
+        self._stop_signal.set()
+        self._child_process.join()
+        self._child_process.close()
+        self.is_scanning = False
+
+    def get_scans(self: _AsyncScanner, n_pulses: int) -> list[UnprocessedWaveform]:
+        call_time = datetime.now()  # noqa: DTZ005
+        stamped_pulse = self._get_scan()
+
+        while stamped_pulse.timestamp < call_time:
+            stamped_pulse = self._get_scan()
+
+        return [self._get_scan().waveform for _ in range(n_pulses)]
+
+    def get_next(self: _AsyncScanner, averaged_over_n: int = 1) -> UnprocessedWaveform:
+        return UnprocessedWaveform.average(
+            [self._get_scan().waveform for _ in range(averaged_over_n)]
+        )
+
+    def _get_scan(self: _AsyncScanner) -> _TimestampedWaveform:
+        try:
+            return self._shared_mem.get(timeout=self._SCAN_TIMEOUT)
+        except Empty as err:
+            if self._scanner_conn.poll(timeout=self.startup_timeout):
+                msg: _ScannerHealth = self._scanner_conn.recv()
+                if not msg.is_alive:
+                    self.is_scanning = False
+                if msg.error:
+                    raise msg.error from err
+            raise
+
+    @staticmethod
+    def _run_scanner(
+        config: DeviceConfiguration,
+        shared_mem: Queue[_TimestampedWaveform],
+        stop_signal: synchronize.Event,
+        parent_conn: Connection,
+    ) -> None:
+        try:
+            scanner = Scanner(config=config)
+            parent_conn.send(_ScannerHealth(is_alive=True, is_healthy=True, error=None))
+        except (serialutil.SerialException, TimeoutError) as e:
+            parent_conn.send(_ScannerHealth(is_alive=False, is_healthy=False, error=e))
+            return
+
+        while not stop_signal.is_set():
+            try:
+                waveform = _TimestampedWaveform(datetime.now(), scanner.scan())  # noqa: DTZ005
+            except (serialutil.SerialException, TimeoutError) as e:
+                parent_conn.send(
+                    _ScannerHealth(is_alive=False, is_healthy=False, error=e)
+                )
+                break
+
+            try:
+                shared_mem.put_nowait(waveform)
+            except Full:
+                # when full, remove the oldest scan from the list before putting a new
+                shared_mem.get_nowait()
+                shared_mem.put_nowait(waveform)
+
+        # Empty queue before shutting down
+        try:
+            while 1:
+                shared_mem.get_nowait()
+        except Empty:
+            # this call required - see https://docs.python.org/3.9/library/multiprocessing.html#programming-guidelines
+            shared_mem.cancel_join_thread()
+            parent_conn.close()
+            shared_mem.cancel_join_thread()
+            parent_conn.close()