herosdevices.hardware.rigol - Heros Devices 0.5.4 documentation

Skip to content
Back to top

# Source code for herosdevices.hardware.rigol

    """Rigol Devices."""

    import threading

    import numpy as np
    import pyvisa
    from pyvisa.resources.resource import Resource

    from herosdevices.core.templates.oscilloscope import OscilloscopeTemplate
    from herosdevices.helper import log

    class RigolDSA(OscilloscopeTemplate):
        """
        A class to interface Rigol digital spectrum analyzers.

        Args:
            address: Address or IP of the instrument.
            **kwargs: Additional keyword arguments are passed to
                :py:class:`herosdevices.core.templates.oscilloscope.OscilloscopeTemplate`

        Note:
            Currently, this driver only implements retrieving the current trace and getting the marker positions.
        """

        MAX_MARKERS: int = 4

        def __init__(
            self,
            address: str,
            **kwargs,
        ) -> None:
            self.address: str = address
            self.default_config_dict = {}
            super().__init__(config_dict={"default": {}}, **kwargs)
            # initialize pyvisa resource manager
            self._pyvisa_rm: pyvisa.ResourceManager | None = None
            self.open()

        def _start(self) -> bool:
            raise NotImplementedError("Rigol DSA do not support manual trigger")
            return False

        def _stop(self) -> bool:
            raise NotImplementedError("Rigol DSA do not support manual stopping (yet)")
            return False

        def _set_config(self, config: dict) -> bool:  # noqa: ARG002
            """
            Set a config dictionary on the device.

            Note:
                Currently, this driver does not implement configuring the device.
            """
            return False

        def _get_status(self) -> dict:
            return {
                "acquisition_running": self.acquisition_running,
                "idn": self.get_identification(),
            }

        def _open(self) -> Resource:
            log.info(f"Connecting to {self.address}")
            self._pyvisa_rm = pyvisa.ResourceManager("@py")
            return self._pyvisa_rm.open_resource(f"TCPIP0::{self.address}::INSTR")

        def _teardown(self) -> None:
            try:
                self._device.close()
            except pyvisa.errors.Error:
                log.exception("Error while closind DSA.")
            try:
                self._pyvisa_rm.close()
            except pyvisa.errors.Error:
                log.exception("Error while closing pyvisa resource manager.")

        def _acquisition_loop(self) -> None:
            with self.get_scope() as sa:
                # get trace as ASCII CSV
                sa.write(":FORM:DATA ASCII")  # set ASCII mode
                raw = sa.query(":TRAC:DATA? TRACE1")
                # remove any leading/trailing whitespace or newlines
                raw = raw.strip()
                # if the device adds a leading '#' or count
                if raw.startswith("#"):
                    # skip SCPI-style length header
                    raw = raw[2 + int(raw[1]) :]  # '#' + digit-count + digits
                # now parse
                trace = np.fromstring(raw, sep=",")
                start_freq = float(sa.query(":FREQ:STAR?"))
                stop_freq = float(sa.query(":FREQ:STOP?"))
                metadata = {
                    "start_freq": start_freq,
                    "stop_freq": stop_freq,
                    "frame": 0,
                }
            self.acquisition_stopped()
            self.acquisition_data(trace, metadata)

        def _arm(self) -> bool:
            log.debug("Starting acquisition thread")
            self._stop_acquisition_event.clear()
            self._acquisition_thread = threading.Thread(target=self._acquisition_loop)
            self._acquisition_thread.start()
            return True

        def get_identification(self) -> str:
            """
            Query the device identification.

            Returns:
                Device identification string.
            """
            with self.get_scope() as sa:
                return sa.query("*IDN?")

        def get_markers(self) -> dict:
            """
            Query all markers from the Rigol spectrum analyzer.

            Returns:
                Dictionary with :code:`{marker_number: (freq_Hz, amp_dBm)}`.
            """
            markers = {}
            with self.get_scope() as sa:
                for i in range(1, self.MAX_MARKERS + 1):
                    try:
                        state = int(sa.query(f":CALC:MARK{i}:STAT?").strip())
                    except pyvisa.errors.Error:
                        # in case the instrument doesn't support more markers, break
                        break
                    if state == 1:
                        freq = float(sa.query(f":CALC:MARK{i}:X?"))
                        amp = float(sa.query(f":CALC:MARK{i}:Y?"))
                        markers[i] = (freq, amp)
            return markers

        def _observable_data(self) -> dict:
            """
            Implement method :meth:`_observable_data` s.t. this class can be used as a :code:`PolledLocalDatasourceHERO`.

            Returns:
                Dictionary with marker positions and amplitudes.
            """
            out = {}
            for i, (freq, amp) in self.get_markers().items():
                out[f"marker_{i}_frequency"] = (freq, "Hz")
                out[f"marker_{i}_amplitude"] = (amp, "dBm")
            return out