Matplotlib's fourier_demo_wx_sgskip, but for qt.

fourier_demo_qt.py

"""
===============
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_()