Source code for semi_cr.core.trial.trial

"""The Trial abstraction: one methodology for acquiring data from the OPX1000.

A :class:`Trial` wraps a particular *way* of getting data (e.g. a 1D or 2D gate
sweep): it builds the QUA program, knows which output streams that program
saves, and knows how to turn the raw stream buffers into a physically
meaningful NumPy array.

Consumers only use :meth:`Trial.stream`::

    for frame in trial.stream():        # Jupyter, scripts, ...
        plot(np.abs(frame))

or hand the generator to a background thread for GUI/"on-line" use. Closing
the generator (``.close()``, ``break``, or garbage collection) cancels the
QUA job and closes the quantum machine.

Currently only continuous (infinite-loop) acquisition is supported; a oneshot
mode will be revisited when a concrete use case exists, since oneshot QUA
programs are structured differently.
"""

import threading
import time
from abc import ABC, abstractmethod
from collections.abc import Callable, Generator, Mapping
from typing import TYPE_CHECKING

import numpy as np
from qm import Program, QuantumMachinesManager

if TYPE_CHECKING:
    from qm.type_hinting.config_types import FullQuaConfig

# TODO: Handle groups of gates




class Trial(ABC):
    """One methodology for acquiring data (e.g. a 1D or 2D gate sweep).

    Subclasses provide the QUA program, the machine configuration it runs on,
    and pure post-processing from raw stream buffers to a NumPy array. The
    base class provides the acquisition machinery (:meth:`fetch`,
    :meth:`stream`).

    By convention :meth:`post_process` returns complex I + jQ in volts;
    consumers pick the component or magnitude they care about.
    """

    #: Names of the output streams the QUA program saves. :meth:`fetch` pulls
    #: exactly these; override in subclasses whose programs save different
    #: buffers.
    stream_names: tuple[str, ...] = ("i", "q")

    #: Delay between successive fetch attempts in :meth:`stream`.
    POLL_INTERVAL_S = 0.01

    def __init__(self, qmm: QuantumMachinesManager):
        self._qmm = qmm

    # --- subclass responsibilities -------------------------------------------



    @abstractmethod
    def build_program(self) -> Program:
        """Build the (continuous, infinite-loop) QUA program for this trial."""




    @abstractmethod
    def machine_config(self) -> "FullQuaConfig":
        """The QUA machine configuration the program runs on."""




    @abstractmethod
    def post_process(self, raw: Mapping[str, np.ndarray]) -> np.ndarray:
        """Turn raw stream buffers (keyed by :attr:`stream_names`) into a
        physical-units array, conventionally complex I + jQ volts.

        Must be a pure function so it can be tested with canned arrays.
        """


    @property
    @abstractmethod
    def axes(self) -> tuple[np.ndarray, ...]:
        """Coordinate arrays, one per dimension of the post-processed result,
        so consumers can plot without knowing sweep internals."""

    # --- concrete machinery --------------------------------------------------



    def fetch(self, job) -> dict[str, np.ndarray] | None:
        """Pull one raw frame from the job, or ``None`` if not all streams
        have produced data yet."""
        raw = {name: job.result_handles.get(name).fetch_all() for name in self.stream_names}
        if any(value is None for value in raw.values()):
            return None
        return raw




    def stream(self) -> Generator[np.ndarray, None, None]:
        """Run the trial and yield post-processed frames until closed.

        Owns the quantum-machine/job lifecycle: opens the machine, executes
        the program, and on generator close cancels the job and closes the
        machine. Frames are polled; the same buffer may be yielded more than
        once — deduplication is the consumer's concern.
        """
        qm = self._qmm.open_qm(self.machine_config())
        job = None
        try:
            job = qm.execute(self.build_program())
            while True:
                raw = self.fetch(job)
                if raw is not None:
                    yield self.post_process(raw)
                time.sleep(self.POLL_INTERVAL_S)
        finally:
            if job is not None:
                job.cancel()
            qm.close()






class ThreadedTrialRunner:
    """Consume a trial's stream on a daemon thread for "on-line" (GUI) use.

    Frames are delivered to ``on_frame`` from the background thread; the
    callback is responsible for its own thread safety. ``stop()`` closes the
    stream, which cancels the QUA job and closes the quantum machine.

    Trial parameters changed? Stop this runner and start a new one with a
    freshly built trial.
    """

    def __init__(self, trial: Trial, on_frame: Callable[[np.ndarray], None]):
        self._trial = trial
        self._on_frame = on_frame
        self._thread: threading.Thread | None = None
        self._stop_event = threading.Event()

    @property
    def is_running(self) -> bool:
        return self._thread is not None and self._thread.is_alive()



    def start(self) -> None:
        if self.is_running:
            return
        self._stop_event.clear()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()




    def stop(self) -> None:
        self._stop_event.set()
        if self._thread is not None:
            self._thread.join(timeout=2.0)


    def _run(self) -> None:
        # The for-loop drives Trial.stream(); breaking out closes the
        # generator (via garbage collection at latest), but we close
        # explicitly so hardware is released before stop() returns.
        frames = self._trial.stream()
        try:
            for frame in frames:
                if self._stop_event.is_set():
                    break
                self._on_frame(frame)
        finally:
            frames.close()