from __future__ import annotations import numpy as np from pyqtgraph import mkPen from numpy import inf, linspace from numpy.fft import fft, fftfreq, fftshift from ...lib.pg_objects import PlotItem, LogInfiniteLine from nmreval.lib.importer import find_models from nmreval.math import apodization as apodization from nmreval.utils.text import convert from ...Qt import QtCore, QtWidgets from ..._py.apod_dialog import Ui_ApodEdit from ..._py.phase_corr_dialog import Ui_SignalEdit from ...lib.forms import FormWidget class QPreviewDialogs(QtWidgets.QDialog): finished = QtCore.pyqtSignal(str, tuple) def __init__(self, parent=None): super().__init__(parent=parent) self.data = [] self.graphs = [] self.mode = '' def setRange(self, xlim: list, ylim: list, logmode: list[bool]): self.graphicsView.getPlotItem().setLogMode(x=logmode[0], y=logmode[1]) if logmode[0]: xlim = [np.log10(x) for x in xlim] if logmode[1]: ylim = [np.log10(y) for y in ylim] self.graphicsView.setRange(xRange=xlim, yRange=ylim, padding=0, disableAutoRange=True) def add_data(self, x, y): self.data.append((x, y)) real_plt = PlotItem(x=x, y=y.real, pen=mkPen('b'), ) imag_plt = PlotItem(x=x, y=y.imag, pen=mkPen('r')) self.graphs.append((real_plt, imag_plt)) self.graphicsView.addItem(real_plt) self.graphicsView.addItem(imag_plt) def done(self, val): self.cleanup() super().done(val) def close(self): self.cleanup() super().close() def accept(self): self.finished.emit(self.mode, self.get_value()) super().accept() def get_value(self): raise NotImplementedError def cleanup(self): self.blockSignals(True) for line in self.graphs: for g in line: self.graphicsView.removeItem(g) del g self.graphicsView.clear() self.data = [] self.graphs = [] self.blockSignals(False) class QPhasedialog(QPreviewDialogs, Ui_SignalEdit): def __init__(self, parent=None): super().__init__(parent=parent) self.setupUi(self) self.mode = 'ph' self.pvt_line = LogInfiniteLine(pos=0, movable=True) self.graphicsView.addItem(self.pvt_line) self.pvt_line.sigPositionChanged.connect(self.move_line) @QtCore.pyqtSlot(float, name='on_ph1slider_valueChanged') @QtCore.pyqtSlot(float, name='on_ph0slider_valueChanged') def _temp_phase(self, *args): ph0, ph1, pvt = self.get_value() self.pvt_line.setValue(pvt) for i, (x, y) in enumerate(self.data): phasecorr = np.exp(-1j * (ph0 + ph1*(x-pvt)/np.max(x))*np.pi/180.) _y = y * phasecorr self.graphs[i][0].setData(x=x, y=_y.real) self.graphs[i][1].setData(x=x, y=_y.imag) def get_value(self): return float(self.ph0slider.text()), float(self.ph1slider.text()), float(self.pivot_lineedit.text()) def move_line(self, evt): self.pivot_lineedit.setText(f'{evt.value():.5g}') class QApodDialog(QPreviewDialogs, Ui_ApodEdit): def __init__(self, parent=None): super().__init__(parent=parent) self.setupUi(self) self._limits = (-inf, inf), -inf self.apods = [] self.apods = find_models(apodization) self.apodcombobox.blockSignals(True) for ap in self.apods: self.apodcombobox.addItem(ap().name) self.apodcombobox.blockSignals(False) self.apod_graph = PlotItem(x=[], y=[]) self.graphicsView.addItem(self.apod_graph) self.mode = 'ap' self.change_apodization(0) def add_data(self, x, y): real_plt = PlotItem(x=x, y=y.real, pen=mkPen('b')) # imag_plt = (x=x, y=y.imag, pen=pg.mkPen('r')) self.graphicsView.addItem(real_plt) # self.graphicsView.addItem(imag_plt) y_fft = fftshift(fft(y)) x_fft = fftshift(fftfreq(len(x), d=x[1]-x[0])) real_plt_fft = PlotItem(x=x_fft, y=y_fft.real, pen=mkPen('b')) # imag_plt_fft = pg.PlotDataItem(x=x_fft, y=y_fft.imag, pen=pg.mkPen('b')) self.graphicsView_2.addItem(real_plt_fft) # self.graphicsView_2.addItem(imag_plt_fft) self.graphs.append((real_plt, real_plt_fft)) self.data.append((x, y, x_fft)) xlimits = (max(x.min(), self._limits[0][0]), min(x.max(), self._limits[0][1])) ylimit = max(self._limits[1], y.real.max()) self._limits = xlimits, ylimit @QtCore.pyqtSlot(int, name='on_apodcombobox_currentIndexChanged') def change_apodization(self, index): # delete old widgets self.eqn_label.setText(convert(self.apods[index].equation)) while self.widget_layout.count(): item = self.widget_layout.takeAt(0) try: item.widget().deleteLater() except AttributeError: pass # set up parameter widgets for new model for k, v in enumerate(self.apods[index]().params): widgt = FormWidget(name=v) widgt.valueChanged.connect(self._temp_apod) self.widget_layout.addWidget(widgt) self.widget_layout.addStretch() self._temp_apod() def _temp_apod(self): apodmodel = self.apods[self.apodcombobox.currentIndex()] p = self._get_parameter() if self.data: for i, (x, y, x_fft) in enumerate(self.data): y2 = apodmodel.apod(x, *p) _y = y2 * y self.graphs[i][0].setData(x=x, y=_y.real) # self.graphs[i][1].setData(y=_y.imag) y_fft = fftshift(fft(_y)) self.graphs[i][1].setData(x=x_fft, y=y_fft.real) # self.graphs[i][3].setData(y=y_fft.imag) _x_apod = linspace(self._limits[0][0], self._limits[0][1]) try: _y_apod = apodmodel.apod(_x_apod, *p) self.apod_graph.setData(x=_x_apod, y=self._limits[1]*_y_apod) except IndexError: pass def _get_parameter(self): p = [] for i in range(self.widget_layout.count()): item = self.widget_layout.itemAt(i) w = item.widget() try: p.append(w.value) except AttributeError: continue return p def get_value(self): apodmodel = self.apods[self.apodcombobox.currentIndex()] p = self._get_parameter() return p, apodmodel