a simple waveform generator with magicgui

generator.py

"""
See: https://docs.scipy.org/doc/scipy/reference/signal.html#waveforms

"""
from dataclasses import dataclass, field
from functools import partial

import matplotlib.pyplot as plt
import numpy as np
from scipy import signal


@dataclass
class Signal:
    """Constructs a 1D signal"""

    func: callable
    t: float = field(init=False)
    duration: float = 1.0
    size: int = 100

    def __post_init__(self):
        self.t = np.linspace(0, self.duration, self.size)
        self.data = self.func(self.t)

    def plot(self, ax=None, **kwargs):
        if ax is None:
            fig, ax = plt.subplots()
        else:
            fig = ax.get_figure()
        ax.plot(self.t, self.data, **kwargs)
        return fig, ax


def sine(
    duration: float = 1.0, size: int = 100, freq: float = 0.5, phase: float = 0.0
) -> Signal:
    """Returns a 1D sine wave

    Parameters
    ----------
    duration: float
       the duration of the signal in seconds
    freq: float
       the frequency of the signal in Hz
    phase: float
       the phase of the signal (in degrees)
    """

    print("******", duration)
    sig = Signal(
        duration=duration,
        size=size,
        func=lambda t: np.sin(t * (2 * np.pi * freq) + phase * np.pi / 180),
    )
    return sig


def chirp(
    duration: float = 1.0,
    size: int = 100,
    f0: float = 1.0,
    t1: float = 0.5,
    f1: float = 2.0,
    phase: float = 0.0,
) -> Signal:
    """Frequency-swept cosine generator

    See scipy.signal.chirp
    """
    sig = Signal(
        duration=duration,
        size=size,
        func=partial(signal.chirp, f0=f0, t1=t1, f1=f1, phi=phase),
    )
    return sig


def sawtooth(
    duration: float = 1.0, size: int = 100, freq: float = 1.0, width: float = 1.0
) -> Signal:
    """Return a periodic sawtooth or triangle waveform.

    See scipy.signal.sawtooth
    """
    sig = Signal(
        duration=duration,
        size=size,
        func=lambda t: signal.sawtooth(2 * np.pi * freq * t, width=width),
    )
    return sig


def square(
    duration: float = 1.0, size: int = 100, freq: float = 1.0, duty: float = 0.5
) -> Signal:
    """Return a periodic sawtooth or triangle waveform.

    See scipy.signal.square
    """
    sig = Signal(
        duration=duration,
        size=size,
        func=lambda t: signal.square(2 * np.pi * freq * t, duty=duty),
    )
    return sig


def on_off(
    duration: float = 1.0, size: int = 100, t_on: float = 0, t_off: float = -1
) -> Signal:

    data = np.ones(size)
    data[: int(size * t_on / duration)] = -1
    if t_off > 0:
        data[int(size * t_off / duration) :] = -1
    sig = Signal(duration=duration, size=size, func=lambda t: data)
    return sig

gui.py

from enum import Enum

import matplotlib.pyplot as plt

# from openephys_event import EventListener
from generator import Signal, chirp, on_off, sawtooth, sine, square
from magicgui import magicgui, type_map, widgets
from matplotlib.backends.backend_qt5agg import FigureCanvas

WAVEFORMS = {
    "sine": sine,
    "chirp": chirp,
    "sawtooth": sawtooth,
    "square": square,
    "on_off": on_off,
}


class Select(Enum):
    Sine = "sine"
    Chirp = "chirp"
    Sawtooth = "sawtooth"
    Square = "square"
    OnOff = "on_off"


class WaveForm(widgets.Container):
    def __init__(self):
        super().__init__()
        self.fig, self.ax = plt.subplots()
        self.native.layout().addWidget(FigureCanvas(self.fig))
        self.waveform = sine
        self.controls = None
        self.append(self.signal_widget)
        self.update_controls()

    @magicgui(auto_call=True)
    def signal_widget(self, select: Select = Select.Sine) -> widgets.Container:
        self.waveform = WAVEFORMS[select.value]
        self.update_controls()

    def update_controls(self):
        if self.controls is not None:
            self.remove(self.controls)
        self.controls = magicgui(auto_call=True)(self.waveform)
        self.append(self.controls)
        self.controls.called.connect(self.update_graph)

    def update_graph(self, sig: Signal):
        print(f"Called with {self.controls}")
        self.ax.cla()
        print(self.ax)
        sig.plot(ax=self.ax)
        self.fig.canvas.draw()

if __name__ == "__main__":
    wv = WaveForm()
    wv.show()