ssqueezepy oscilattions

plot_oscillations.py

import numpy as np
import matplotlib.pyplot as plt

from ssqueezepy import ssq_cwt, cwt


SAMPLE_RATE = 44100


def sine(freq, num_samples=SAMPLE_RATE // 4, amp=1.0):
    ts = np.arange(num_samples) / SAMPLE_RATE
    wave = amp * np.sin(2.0 * np.pi * freq * ts)
    return wave


def mix(*waves):
    return np.array(waves).sum(axis=0) / len(waves)


def plot(Tx, ax):
    xs = np.arange(Tx.shape[1] + 1)
    ys = np.arange(Tx.shape[0] + 1)
    img = ax.pcolormesh(xs, ys, np.abs(Tx))


def plot_ssq(wave, ax):
    kw = dict(nv=32, wavelet=('morlet', {'mu': 4.0}))
    Tx, *_ = ssq_cwt(wave, fs=SAMPLE_RATE, **kw)
    plot(Tx, ax)


if __name__ == "__main__":
    fig, axes = plt.subplots(1, 3, figsize=(18, 6))

    plot_ssq(sine(440.0), axes[0])
    plot_ssq(sine(880.0), axes[1])
    plot_ssq(mix(sine(440.0), sine(880.0)), axes[2])

    fig.tight_layout()
    plt.savefig("/tmp/osci.png")
    plt.show()