1import logging 2import os 3 4import numpy as np 5from PyQt5 import uic 6from PyQt5.QtCore import QRegExp, Qt, pyqtSignal, pyqtSlot 7from PyQt5.QtGui import QBrush, QColor, QPen, QRegExpValidator 8from PyQt5.QtWidgets import QApplication, QDialog 9 10from urh.plugins.Plugin import SignalEditorPlugin 11from urh.signalprocessing.IQArray import IQArray 12from urh.ui.painting.SceneManager import SceneManager 13from urh.util.Formatter import Formatter 14from urh.util.Logger import logger 15 16 17class InsertSinePlugin(SignalEditorPlugin): 18 insert_sine_wave_clicked = pyqtSignal() 19 20 INSERT_INDICATOR_COLOR = QColor(0, 255, 0, 80) 21 22 def __init__(self): 23 24 self.__dialog_ui = None # type: QDialog 25 self.complex_wave = None 26 27 self.__amplitude = 0.5 28 self.__frequency = 10 29 self.__phase = 0 30 self.__sample_rate = 1e6 31 self.__num_samples = int(1e6) 32 33 self.original_data = None 34 self.draw_data = None 35 self.position = 0 36 37 super().__init__(name="InsertSine") 38 39 @property 40 def dialog_ui(self) -> QDialog: 41 if self.__dialog_ui is None: 42 dir_name = os.path.dirname(os.readlink(__file__)) if os.path.islink(__file__) else os.path.dirname(__file__) 43 44 logging.getLogger().setLevel(logging.WARNING) 45 self.__dialog_ui = uic.loadUi(os.path.realpath(os.path.join(dir_name, "insert_sine_dialog.ui"))) 46 logging.getLogger().setLevel(logger.level) 47 48 self.__dialog_ui.setAttribute(Qt.WA_DeleteOnClose) 49 self.__dialog_ui.setModal(True) 50 self.__dialog_ui.doubleSpinBoxAmplitude.setValue(self.__amplitude) 51 self.__dialog_ui.doubleSpinBoxFrequency.setValue(self.__frequency) 52 self.__dialog_ui.doubleSpinBoxPhase.setValue(self.__phase) 53 self.__dialog_ui.doubleSpinBoxSampleRate.setValue(self.__sample_rate) 54 self.__dialog_ui.doubleSpinBoxNSamples.setValue(self.__num_samples) 55 self.__dialog_ui.lineEditTime.setValidator( 56 QRegExpValidator(QRegExp(r"[0-9]+([nmµ]?|([\.,][0-9]{1,3}[nmµ]?))?$")) 57 ) 58 59 scene_manager = SceneManager(self.dialog_ui.graphicsViewSineWave) 60 self.__dialog_ui.graphicsViewSineWave.scene_manager = scene_manager 61 self.insert_indicator = scene_manager.scene.addRect(0, -2, 0, 4, 62 QPen(QColor(Qt.transparent), 0), 63 QBrush(self.INSERT_INDICATOR_COLOR)) 64 self.insert_indicator.stackBefore(scene_manager.scene.selection_area) 65 66 self.set_time() 67 68 return self.__dialog_ui 69 70 @property 71 def amplitude(self) -> float: 72 return self.__amplitude 73 74 @amplitude.setter 75 def amplitude(self, value: float): 76 if value != self.amplitude: 77 self.__amplitude = value 78 self.draw_sine_wave() 79 80 @property 81 def frequency(self) -> float: 82 return self.__frequency 83 84 @frequency.setter 85 def frequency(self, value: float): 86 if value != self.frequency: 87 self.__frequency = value 88 self.draw_sine_wave() 89 90 @property 91 def phase(self) -> float: 92 return self.__phase 93 94 @phase.setter 95 def phase(self, value: float): 96 if value != self.phase: 97 self.__phase = value 98 self.draw_sine_wave() 99 100 @property 101 def sample_rate(self) -> float: 102 return self.__sample_rate 103 104 @sample_rate.setter 105 def sample_rate(self, value: float): 106 if value != self.sample_rate: 107 self.__sample_rate = value 108 self.set_time() 109 self.draw_sine_wave() 110 111 @property 112 def num_samples(self) -> int: 113 return self.__num_samples 114 115 @num_samples.setter 116 def num_samples(self, value: int): 117 value = int(value) 118 if value != self.num_samples: 119 self.__num_samples = value 120 self.set_time() 121 self.draw_sine_wave() 122 123 def create_connects(self): 124 pass 125 126 def create_dialog_connects(self): 127 self.dialog_ui.doubleSpinBoxAmplitude.editingFinished.connect( 128 self.on_double_spin_box_amplitude_editing_finished) 129 self.dialog_ui.doubleSpinBoxFrequency.editingFinished.connect( 130 self.on_double_spin_box_frequency_editing_finished) 131 self.dialog_ui.doubleSpinBoxPhase.editingFinished.connect(self.on_double_spin_box_phase_editing_finished) 132 self.dialog_ui.doubleSpinBoxSampleRate.editingFinished.connect( 133 self.on_double_spin_box_sample_rate_editing_finished) 134 self.dialog_ui.doubleSpinBoxNSamples.editingFinished.connect(self.on_spin_box_n_samples_editing_finished) 135 self.dialog_ui.lineEditTime.editingFinished.connect(self.on_line_edit_time_editing_finished) 136 self.dialog_ui.buttonBox.accepted.connect(self.on_button_box_accept) 137 self.dialog_ui.buttonBox.rejected.connect(self.on_button_box_reject) 138 self.__dialog_ui.finished.connect(self.on_dialog_finished) 139 140 def get_insert_sine_dialog(self, original_data, position, sample_rate=None, num_samples=None) -> QDialog: 141 if sample_rate is not None: 142 self.__sample_rate = sample_rate 143 self.dialog_ui.doubleSpinBoxSampleRate.setValue(sample_rate) 144 145 if num_samples is not None: 146 self.__num_samples = int(num_samples) 147 self.dialog_ui.doubleSpinBoxNSamples.setValue(num_samples) 148 149 self.original_data = original_data 150 self.position = position 151 152 self.set_time() 153 self.draw_sine_wave() 154 self.create_dialog_connects() 155 156 return self.dialog_ui 157 158 def draw_sine_wave(self): 159 if self.dialog_ui.graphicsViewSineWave.scene_manager: 160 self.dialog_ui.graphicsViewSineWave.scene_manager.clear_path() 161 162 QApplication.instance().setOverrideCursor(Qt.WaitCursor) 163 self.__set_status_of_editable_elements(enabled=False) 164 165 t = np.arange(0, self.num_samples) / self.sample_rate 166 arg = 2 * np.pi * self.frequency * t + self.phase 167 168 self.complex_wave = np.empty(len(arg), dtype=np.complex64) 169 self.complex_wave.real = np.cos(arg) 170 self.complex_wave.imag = np.sin(arg) 171 self.complex_wave = IQArray(self.amplitude * self.complex_wave).convert_to(self.original_data.dtype) 172 173 self.draw_data = np.insert(self.original_data[:, 0], self.position, self.complex_wave[:, 0]) 174 y, h = self.dialog_ui.graphicsViewSineWave.view_rect().y(), self.dialog_ui.graphicsViewSineWave.view_rect().height() 175 self.insert_indicator.setRect(self.position, y - h, self.num_samples, 2 * h + abs(y)) 176 177 self.__set_status_of_editable_elements(enabled=True) 178 QApplication.instance().restoreOverrideCursor() 179 self.dialog_ui.graphicsViewSineWave.plot_data(self.draw_data) 180 self.dialog_ui.graphicsViewSineWave.show_full_scene() 181 182 def __set_status_of_editable_elements(self, enabled: bool): 183 for obj in ("doubleSpinBoxAmplitude", "doubleSpinBoxFrequency", "doubleSpinBoxPhase", 184 "doubleSpinBoxSampleRate", "doubleSpinBoxNSamples", "lineEditTime", "buttonBox"): 185 getattr(self.dialog_ui, obj).setEnabled(enabled) 186 187 def set_time(self): 188 self.dialog_ui.lineEditTime.setText(Formatter.science_time(self.num_samples / self.sample_rate, decimals=3, 189 append_seconds=False, remove_spaces=True)) 190 191 @pyqtSlot() 192 def on_double_spin_box_amplitude_editing_finished(self): 193 self.amplitude = self.dialog_ui.doubleSpinBoxAmplitude.value() 194 195 @pyqtSlot() 196 def on_double_spin_box_frequency_editing_finished(self): 197 self.frequency = self.dialog_ui.doubleSpinBoxFrequency.value() 198 199 @pyqtSlot() 200 def on_double_spin_box_phase_editing_finished(self): 201 self.phase = self.dialog_ui.doubleSpinBoxPhase.value() 202 203 @pyqtSlot() 204 def on_double_spin_box_sample_rate_editing_finished(self): 205 self.sample_rate = self.dialog_ui.doubleSpinBoxSampleRate.value() 206 207 @pyqtSlot() 208 def on_spin_box_n_samples_editing_finished(self): 209 self.num_samples = self.dialog_ui.doubleSpinBoxNSamples.value() 210 211 @pyqtSlot() 212 def on_line_edit_time_editing_finished(self): 213 time_str = self.dialog_ui.lineEditTime.text().replace(",", ".") 214 suffix = "" 215 try: 216 t = float(time_str) 217 except ValueError: 218 suffix = time_str[-1] 219 try: 220 t = float(time_str[:-1]) 221 except ValueError: 222 return 223 224 factor = 10 ** -9 if suffix == "n" else 10 ** -6 if suffix == "µ" else 10 ** -3 if suffix == "m" else 1 225 time_val = t * factor 226 227 if self.sample_rate * time_val >= 1: 228 self.dialog_ui.doubleSpinBoxNSamples.setValue(self.sample_rate * time_val) 229 self.dialog_ui.doubleSpinBoxNSamples.editingFinished.emit() 230 else: 231 self.set_time() 232 233 @pyqtSlot() 234 def on_button_box_reject(self): 235 self.dialog_ui.reject() 236 237 @pyqtSlot() 238 def on_button_box_accept(self): 239 self.insert_sine_wave_clicked.emit() 240 self.dialog_ui.accept() 241 242 @pyqtSlot() 243 def on_dialog_finished(self): 244 self.sender().graphicsViewSineWave.eliminate() 245 self.__dialog_ui = None 246