{
    "nbformat": 4,
    "nbformat_minor": 5,
    "metadata": {
        "kernelspec": {
            "display_name": "Python",
            "language": "python",
            "name": "python"
        },
        "language_info": {
            "name": "python"
        }
    },
    "cells": [
        {
            "id": "cb80b3b0",
            "cell_type": "markdown",
            "source": "<div class='alert alert-warning'>\n\nSciPy's interactive examples with Jupyterlite are experimental and may not always work as expected. Execution of cells containing imports may result in large downloads (up to 60MB of content for the first import from SciPy). Load times when importing from SciPy may take roughly 10-20 seconds. If you notice any problems, feel free to open an [issue](https://github.com/scipy/scipy/issues/new/choose).\n\n</div>",
            "metadata": {}
        },
        {
            "id": "ca8e7ad0",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "import numpy as np\nfrom scipy import signal\nimport matplotlib.pyplot as plt\nrng = np.random.default_rng()",
            "outputs": []
        },
        {
            "id": "420ec801",
            "cell_type": "markdown",
            "source": "Generate a test signal, a 2 Vrms sine wave whose frequency is slowly\nmodulated around 3kHz, corrupted by white noise of exponentially\ndecreasing magnitude sampled at 10 kHz.\n",
            "metadata": {}
        },
        {
            "id": "7f2f861c",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "fs = 10e3\nN = 1e5\namp = 2 * np.sqrt(2)\nnoise_power = 0.01 * fs / 2\ntime = np.arange(N) / float(fs)\nmod = 500*np.cos(2*np.pi*0.25*time)\ncarrier = amp * np.sin(2*np.pi*3e3*time + mod)\nnoise = rng.normal(scale=np.sqrt(noise_power),\n                   size=time.shape)\nnoise *= np.exp(-time/5)\nx = carrier + noise",
            "outputs": []
        },
        {
            "id": "c07f3820",
            "cell_type": "markdown",
            "source": "Compute and plot the STFT's magnitude.\n",
            "metadata": {}
        },
        {
            "id": "0ad17c63",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "f, t, Zxx = signal.stft(x, fs, nperseg=1000)\nplt.pcolormesh(t, f, np.abs(Zxx), vmin=0, vmax=amp, shading='gouraud')\nplt.title('STFT Magnitude')\nplt.ylabel('Frequency [Hz]')\nplt.xlabel('Time [sec]')\nplt.show()",
            "outputs": []
        },
        {
            "id": "aed43a90",
            "cell_type": "markdown",
            "source": "Compare the energy of the signal `x` with the energy of its STFT:\n",
            "metadata": {}
        },
        {
            "id": "c482705f",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "E_x = sum(x**2) / fs  # Energy of x\n# Calculate a two-sided STFT with PSD scaling:\nf, t, Zxx = signal.stft(x, fs, nperseg=1000, return_onesided=False,\n                        scaling='psd')\n# Integrate numerically over abs(Zxx)**2:\ndf, dt = f[1] - f[0], t[1] - t[0]\nE_Zxx = sum(np.sum(Zxx.real**2 + Zxx.imag**2, axis=0) * df) * dt\n# The energy is the same, but the numerical errors are quite large:\nnp.isclose(E_x, E_Zxx, rtol=1e-2)",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "True"
                    },
                    "execution_count": null
                }
            ]
        }
    ]
}