Oscilloscope_Data_Grabber/main.py

import os
import sys
from rigol_dho_lib import rigol  # Assuming this is the library for the oscilloscope
from rigol_dho_lib.rigol import RigolOsc

os.environ["QT_API"] = "PySide6"

from PySide6.QtWidgets import (
    QMainWindow,
    QApplication,
    QPushButton,
    QVBoxLayout,
    QWidget,
    QLineEdit,
    QLabel,
    QCheckBox,
    QHBoxLayout,
    QMessageBox,
)

from matplotlib.backends.backend_qtagg import FigureCanvas
from matplotlib.figure import Figure


class MplCanvas(FigureCanvas):
    def __init__(self, parent=None, width=5, height=4, dpi=100):
        fig = Figure(figsize=(width, height), dpi=dpi)
        self.axes = fig.add_subplot(111)
        super().__init__(fig)


class OscilloscopeGrabber:
    def __init__(self, osc_addr: str):
        print(f"Connecting to oscilloscope at: {osc_addr}")
        self.rigol = rigol.RigolOsc(osc_addr)


class MainWindow(QMainWindow):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.setWindowTitle("Oscilloscope Waveform Viewer")

        # Default values
        default_osc_address = "TCPIP::192.168.1.3::INSTR"
        default_path = "/home/zychlix/Desktop/pomiary/zrzuty/data.csv"

        # Create main layout
        layout = QVBoxLayout()

        # Oscilloscope address input
        self.osc_addr = QLineEdit(default_osc_address)
        self.osc_addr.setPlaceholderText("Enter oscilloscope address")
        layout.addWidget(QLabel("Oscilloscope Address:"))
        layout.addWidget(self.osc_addr)

        self.path_input = QLineEdit(default_path)
        self.path_input.setPlaceholderText("Enter path")
        layout.addWidget(QLabel("Output File"))
        layout.addWidget(self.path_input)

        save_button = QPushButton("Save Waveform Data")
        save_button.clicked.connect(self.save_waveform_data)
        layout.addWidget(save_button)

        # Connect button
        self.connect_button = QPushButton("Connect to Oscilloscope")
        self.connect_button.clicked.connect(self.connect_to_oscilloscope)
        layout.addWidget(self.connect_button)

        # Channel selector (checkboxes for multiple channels)
        layout.addWidget(QLabel("Select Channels:"))
        self.channel_checkboxes = {}
        channel_layout = QHBoxLayout()
        for i in range(1, 5):  # Assuming 4 channels
            checkbox = QCheckBox(f"Channel {i}")
            self.channel_checkboxes[i] = checkbox
            channel_layout.addWidget(checkbox)
        layout.addLayout(channel_layout)

        # Button to get waveform from oscilloscope
        self.get_waveform_button = QPushButton("Get Waveform")
        self.get_waveform_button.clicked.connect(self.grab_osc_data)
        layout.addWidget(self.get_waveform_button)

        # Matplotlib canvas
        self.canvas = MplCanvas(self, width=5, height=4, dpi=100)
        layout.addWidget(self.canvas)

        # Set central widget
        container = QWidget()
        container.setLayout(layout)
        self.setCentralWidget(container)

        self.show()

    def connect_to_oscilloscope(self):
        try:
            self.og = OscilloscopeGrabber(self.osc_addr.text())
            QMessageBox.information(self, "Success", "Connected to oscilloscope.")
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to connect: {e}")

    def grab_osc_data(self):
        try:
            # Clear the previous plot
            self.canvas.axes.clear()

            # Retrieve and plot waveform data for selected channels
            self.waveform_data = {}  # Initialize the dictionary to store waveform data
            for channel_num, checkbox in self.channel_checkboxes.items():
                if checkbox.isChecked():
                    channel = self.og.rigol.getChannel(channel_num)
                    waveform = channel.getWaveform()
                    time_array = channel.genTimeArray(waveform)

                    # Store the waveform data for saving later
                    self.waveform_data[channel_num] = (time_array, waveform)

                    # Plot the waveform data
                    self.canvas.axes.plot(
                        time_array, waveform, label=f"Channel {channel_num}"
                    )

            # Add labels and legend
            self.canvas.axes.set_title("Oscilloscope Waveform")
            self.canvas.axes.set_xlabel("Time (s)")
            self.canvas.axes.set_ylabel("Amplitude (V)")
            self.canvas.axes.legend()

            # Redraw the canvas
            self.canvas.draw()
        except Exception as e:
            QMessageBox.critical(
                self, "Error", f"Failed to retrieve waveform data: {e}"
            )

    def save_waveform_data(self):
        path = self.path_input.text()
        if not path:
            QMessageBox.warning(self, "Warning", "Please specify a save path.")
            return

        try:
            # Ensure waveform data exists before saving
            if not hasattr(self, "waveform_data") or not self.waveform_data:
                QMessageBox.warning(
                    self,
                    "Warning",
                    "No waveform data to save. Please retrieve data first.",
                )
                return

            # Prepare data for all 4 channels
            all_channels = {}
            max_length = 0
            time_column = None

            for channel_num in range(1, 5):  # Assuming 4 channels
                if channel_num in self.waveform_data:
                    time_array, waveform = self.waveform_data[channel_num]
                    all_channels[channel_num] = waveform
                    if time_column is None:  # Use the first available time array
                        time_column = time_array
                    max_length = max(max_length, len(time_array))
                else:
                    # Fill with empty data if the channel is not selected
                    all_channels[channel_num] = []

            # Ensure all channels have the same length by padding with empty values
            for channel_num in range(1, 5):
                waveform = all_channels[channel_num]
                if len(waveform) < max_length:
                    waveform.extend([""] * (max_length - len(waveform)))
                all_channels[channel_num] = waveform

            # Write to CSV
            with open(path, "w") as f:
                # Write header
                f.write("Time (s),Voltage Ch1,Voltage Ch2,Voltage Ch3,Voltage Ch4\n")

                # Write data row by row
                for i in range(max_length):
                    row = [
                        str(time_column[i])
                        if time_column is not None and len(time_column) > i
                        else ""
                    ]
                    for channel_num in range(1, 5):
                        row.append(
                            str(all_channels[channel_num][i])
                            if i < len(all_channels[channel_num])
                            else ""
                        )
                    f.write(",".join(row) + "\n")

            QMessageBox.information(self, "Success", f"Waveform data saved to {path}.")
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to save waveform data: {e}")


app = QApplication(sys.argv)
w = MainWindow()
app.exec()