{
    "nbformat": 4,
    "nbformat_minor": 5,
    "metadata": {
        "kernelspec": {
            "display_name": "Python",
            "language": "python",
            "name": "python"
        },
        "language_info": {
            "name": "python"
        }
    },
    "cells": [
        {
            "id": "53c56699",
            "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": "015b1f15",
            "cell_type": "markdown",
            "source": "The following example shows the magnitude of the STFT of a sine with\nvarying frequency $f_i(t)$ (marked by a red dashed line in the plot):\n",
            "metadata": {}
        },
        {
            "id": "7b4affa5",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "import numpy as np\nimport matplotlib.pyplot as plt\nfrom scipy.signal import ShortTimeFFT\nfrom scipy.signal.windows import gaussian\n\nT_x, N = 1 / 20, 1000  # 20 Hz sampling rate for 50 s signal\nt_x = np.arange(N) * T_x  # time indexes for signal\nf_i = 1 * np.arctan((t_x - t_x[N // 2]) / 2) + 5  # varying frequency\nx = np.sin(2*np.pi*np.cumsum(f_i)*T_x) # the signal",
            "outputs": []
        },
        {
            "id": "1e47cf7f",
            "cell_type": "markdown",
            "source": "The utilized Gaussian window is 50 samples or 2.5 s long. The parameter\n``mfft=200`` in `ShortTimeFFT` causes the spectrum to be oversampled\nby a factor of 4:\n",
            "metadata": {}
        },
        {
            "id": "03921acb",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "g_std = 8  # standard deviation for Gaussian window in samples\nw = gaussian(50, std=g_std, sym=True)  # symmetric Gaussian window\nSFT = ShortTimeFFT(w, hop=10, fs=1/T_x, mfft=200, scale_to='magnitude')\nSx = SFT.stft(x)  # perform the STFT",
            "outputs": []
        },
        {
            "id": "8c509ac5",
            "cell_type": "markdown",
            "source": "In the plot, the time extent of the signal `x` is marked by vertical dashed\nlines. Note that the SFT produces values outside the time range of `x`. The\nshaded areas on the left and the right indicate border effects caused\nby the window slices in that area not fully being inside time range of\n`x`:\n",
            "metadata": {}
        },
        {
            "id": "80cdb290",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "fig1, ax1 = plt.subplots(figsize=(6., 4.))  # enlarge plot a bit\nt_lo, t_hi = SFT.extent(N)[:2]  # time range of plot\nax1.set_title(rf\"STFT ({SFT.m_num*SFT.T:g}$\\,s$ Gaussian window, \" +\n              rf\"$\\sigma_t={g_std*SFT.T}\\,$s)\")\nax1.set(xlabel=f\"Time $t$ in seconds ({SFT.p_num(N)} slices, \" +\n               rf\"$\\Delta t = {SFT.delta_t:g}\\,$s)\",\n        ylabel=f\"Freq. $f$ in Hz ({SFT.f_pts} bins, \" +\n               rf\"$\\Delta f = {SFT.delta_f:g}\\,$Hz)\",\n        xlim=(t_lo, t_hi))\n\nim1 = ax1.imshow(abs(Sx), origin='lower', aspect='auto',\n                 extent=SFT.extent(N), cmap='viridis')\nax1.plot(t_x, f_i, 'r--', alpha=.5, label='$f_i(t)$')\nfig1.colorbar(im1, label=\"Magnitude $|S_x(t, f)|$\")\n\n# Shade areas where window slices stick out to the side:\nfor t0_, t1_ in [(t_lo, SFT.lower_border_end[0] * SFT.T),\n                 (SFT.upper_border_begin(N)[0] * SFT.T, t_hi)]:\n    ax1.axvspan(t0_, t1_, color='w', linewidth=0, alpha=.2)\nfor t_ in [0, N * SFT.T]:  # mark signal borders with vertical line:\n    ax1.axvline(t_, color='y', linestyle='--', alpha=0.5)\nax1.legend()\nfig1.tight_layout()\nplt.show()",
            "outputs": []
        },
        {
            "id": "cb12e6b2",
            "cell_type": "markdown",
            "source": "Reconstructing the signal with the `~ShortTimeFFT.istft` is\nstraightforward, but note that the length of `x1` should be specified,\nsince the STFT length increases in `hop` steps:\n",
            "metadata": {}
        },
        {
            "id": "fe2fec73",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "SFT.invertible  # check if invertible",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "True"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "b610da40",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "x1 = SFT.istft(Sx, k1=N)\nnp.allclose(x, x1)",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "True"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "37805592",
            "cell_type": "markdown",
            "source": "It is possible to calculate the STFT of signal parts:\n",
            "metadata": {}
        },
        {
            "id": "1b1901c3",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "N2 = SFT.nearest_k_p(N // 2)\nSx0 = SFT.stft(x[:N2])\nSx1 = SFT.stft(x[N2:])",
            "outputs": []
        },
        {
            "id": "fede2632",
            "cell_type": "markdown",
            "source": "When assembling sequential STFT parts together, the overlap needs to be\nconsidered:\n",
            "metadata": {}
        },
        {
            "id": "ea780c5a",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "p0_ub = SFT.upper_border_begin(N2)[1] - SFT.p_min\np1_le = SFT.lower_border_end[1] - SFT.p_min\nSx01 = np.hstack((Sx0[:, :p0_ub],\n                  Sx0[:, p0_ub:] + Sx1[:, :p1_le],\n                  Sx1[:, p1_le:]))\nnp.allclose(Sx01, Sx)  # Compare with SFT of complete signal",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "True"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "26751cf2",
            "cell_type": "markdown",
            "source": "It is also possible to calculate the `itsft` for signal parts:\n",
            "metadata": {}
        },
        {
            "id": "e945d618",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "y_p = SFT.istft(Sx, N//3, N//2)\nnp.allclose(y_p, x[N//3:N//2])",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "True"
                    },
                    "execution_count": null
                }
            ]
        }
    ]
}