Skip to content

Instantly share code, notes, and snippets.

@anntzer
Created October 19, 2019 17:43
Show Gist options
  • Save anntzer/112a6b0d9567472c83b6c12ccf5fb10c to your computer and use it in GitHub Desktop.
Save anntzer/112a6b0d9567472c83b6c12ccf5fb10c to your computer and use it in GitHub Desktop.
Matplotlib's fourier_demo_wx_sgskip, but for qt.
"""
===============
Qt Fourier Demo
===============
"""
import contextlib
import sys
from matplotlib.backends.backend_qt5agg import FigureCanvas
from matplotlib.backends.qt_compat import (
QtCore as QC, QtGui as QG, QtWidgets as QW)
from matplotlib.figure import Figure
import numpy as np
Qt = QC.Qt
class Layouter:
"""
A helper for hierarchical layouts.
"""
def __init__(self, window, layout):
central_widget = QW.QWidget()
window.setCentralWidget(central_widget)
central_widget.setLayout(layout)
self._layout_stack = [layout]
def __call__(self, method_name, *args):
getattr(self._layout_stack[-1], method_name)(*args)
@contextlib.contextmanager
def sublayout(self, layout):
widget = QW.QWidget()
self("addWidget", widget)
widget.setLayout(layout)
self._layout_stack.append(layout)
try:
yield layout
finally:
self._layout_stack.pop()
class MainWindow(QW.QMainWindow):
def __init__(self, *, parent=None):
super().__init__(parent=parent)
# Set up the figure.
self.figure = Figure(constrained_layout=True)
ax0, ax1 = self.figure.subplots(2)
self.lines = [*ax0.plot([], []), *ax1.plot([], [])]
self.figure.suptitle(
"Click and drag waveforms to change frequency and amplitude")
ax0.set(xlim=[-6, 6], ylim=[0, 1],
xlabel="frequency", ylabel="frequency domain waveform")
ax1.set(xlim=[-2, 2], ylim=[-2, 2],
xlabel="time", ylabel="time domain waveform")
ax0.text(.05, .95,
r"$X(f) = \mathcal{F}\{x(t)\}$",
verticalalignment="top",
transform=ax0.transAxes)
ax1.text(.05, .95,
r"$x(t) = a \cdot \cos(2\pi f_0 t) e^{-\pi t^2}$",
verticalalignment="top",
transform=ax1.transAxes)
# Set up the UI.
layouter = Layouter(self, QW.QHBoxLayout())
layouter("addWidget", FigureCanvas(self.figure))
with layouter.sublayout(QW.QFormLayout()) as sublayout:
self.frequency_widget = QW.QLineEdit(editingFinished=self._update)
self.frequency_widget.setText("2")
self.frequency_widget.setValidator(QG.QDoubleValidator())
layouter("addRow", "frequency", self.frequency_widget)
self.amplitude_widget = QW.QLineEdit(editingFinished=self._update)
self.amplitude_widget.setText("1")
self.amplitude_widget.setValidator(QG.QDoubleValidator())
layouter("addRow", "amplitude", self.amplitude_widget)
self._update()
# Set up the callbacks.
self._button_down_info = (None, None, None, None, None)
self.figure.canvas.callbacks.connect(
"button_press_event", self._on_button_press)
self.figure.canvas.callbacks.connect(
"motion_notify_event", self._on_motion_notify)
self.figure.canvas.callbacks.connect(
"button_release_event", self._on_button_release)
def _on_button_press(self, evt):
if self.lines[0].contains(evt)[0]:
state = "frequency"
elif self.lines[1].contains(evt)[0]:
state = "time"
else:
state = None
self._button_down_info = (
state, evt.xdata, evt.ydata,
max(float(self.frequency_widget.text()), .1),
float(self.amplitude_widget.text()),
)
def _on_motion_notify(self, evt):
state, x0, y0, f0_init, a_init = self._button_down_info
if state is None:
return
x, y = evt.xdata, evt.ydata
if x is None: # outside the axes
return
self.amplitude_widget.setText(str(a_init + (a_init * (y - y0) / y0)))
if state == "frequency":
self.frequency_widget.setText(
str(f0_init + (f0_init * (x - x0) / x0)))
elif state == "time":
if (x - x0) / x0 != -1:
self.frequency_widget.setText(
str(1. / (1. / f0_init + (1. / f0_init * (x - x0) / x0))))
self._update()
def _on_button_release(self, evt):
self._button_down_info = (None, None, None, None, None)
def compute(self, f0, a):
fs = np.arange(-6, 6, 0.02)
ts = np.arange(-2, 2, 0.01)
x = a * np.cos(2 * np.pi * f0 * ts) * np.exp(-np.pi * ts ** 2)
X = a / 2 * \
(np.exp(-np.pi * (fs - f0) ** 2) + np.exp(-np.pi * (fs + f0) ** 2))
return fs, X, ts, x
def _update(self):
f = float(self.frequency_widget.text())
a = float(self.amplitude_widget.text())
x1, y1, x2, y2 = self.compute(f, a)
# update the data of the two waveforms
self.lines[0].set(xdata=x1, ydata=y1)
self.lines[1].set(xdata=x2, ydata=y2)
# make the canvas draw its contents again with the new data
self.figure.canvas.draw()
if __name__ == "__main__":
qapp = QW.QApplication(sys.argv)
app = MainWindow()
app.show()
qapp.exec_()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment