{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n\n# Time-frequency on simulated data (Multitaper vs. Morlet vs. Stockwell vs. Hilbert)\n\nThis example demonstrates the different time-frequency estimation methods\non simulated data. It shows the time-frequency resolution trade-off\nand the problem of estimation variance. In addition it highlights\nalternative functions for generating TFRs without averaging across\ntrials, or by operating on numpy arrays.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Authors: Hari Bharadwaj <hari@nmr.mgh.harvard.edu>\n#          Denis Engemann <denis.engemann@gmail.com>\n#          Chris Holdgraf <choldgraf@berkeley.edu>\n#          Alex Rockhill <aprockhill@mailbox.org>\n#\n# License: BSD-3-Clause\n# Copyright the MNE-Python contributors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nfrom matplotlib import pyplot as plt\n\nfrom mne import Epochs, create_info\nfrom mne.io import RawArray\nfrom mne.time_frequency import AverageTFRArray, EpochsTFRArray, tfr_array_morlet\n\nprint(__doc__)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Simulate data\n\nWe'll simulate data with a known spectro-temporal structure.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "sfreq = 1000.0\nch_names = [\"SIM0001\", \"SIM0002\"]\nch_types = [\"grad\", \"grad\"]\ninfo = create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types)\n\nn_times = 1024  # Just over 1 second epochs\nn_epochs = 40\nseed = 42\nrng = np.random.RandomState(seed)\ndata = rng.randn(len(ch_names), n_times * n_epochs + 200)  # buffer\n\n# Add a 50 Hz sinusoidal burst to the noise and ramp it.\nt = np.arange(n_times, dtype=np.float64) / sfreq\nsignal = np.sin(np.pi * 2.0 * 50.0 * t)  # 50 Hz sinusoid signal\nsignal[np.logical_or(t < 0.45, t > 0.55)] = 0.0  # hard windowing\non_time = np.logical_and(t >= 0.45, t <= 0.55)\nsignal[on_time] *= np.hanning(on_time.sum())  # ramping\ndata[:, 100:-100] += np.tile(signal, n_epochs)  # add signal\n\nraw = RawArray(data, info)\nevents = np.zeros((n_epochs, 3), dtype=int)\nevents[:, 0] = np.arange(n_epochs) * n_times\nepochs = Epochs(\n    raw,\n    events,\n    dict(sin50hz=0),\n    tmin=0,\n    tmax=n_times / sfreq,\n    reject=dict(grad=4000),\n    baseline=None,\n)\n\nepochs.average().plot()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Calculate a time-frequency representation (TFR)\n\nBelow we'll demonstrate the output of several TFR functions in MNE:\n\n* :func:`mne.time_frequency.tfr_multitaper`\n* :func:`mne.time_frequency.tfr_stockwell`\n* :func:`mne.time_frequency.tfr_morlet`\n* :meth:`mne.Epochs.filter` and :meth:`mne.Epochs.apply_hilbert`\n\n### Multitaper transform\nFirst we'll use the multitaper method for calculating the TFR.\nThis creates several orthogonal tapering windows in the TFR estimation,\nwhich reduces variance. We'll also show some of the parameters that can be\ntweaked (e.g., ``time_bandwidth``) that will result in different multitaper\nproperties, and thus a different TFR. You can trade time resolution or\nfrequency resolution or both in order to get a reduction in variance.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "freqs = np.arange(5.0, 100.0, 3.0)\nvmin, vmax = -3.0, 3.0  # Define our color limits.\n\nfig, axs = plt.subplots(1, 3, figsize=(15, 5), sharey=True, layout=\"constrained\")\nfor n_cycles, time_bandwidth, ax, title in zip(\n    [freqs / 2, freqs, freqs / 2],  # number of cycles\n    [2.0, 4.0, 8.0],  # time bandwidth\n    axs,\n    [\n        \"Sim: Least smoothing, most variance\",\n        \"Sim: Less frequency smoothing,\\nmore time smoothing\",\n        \"Sim: Less time smoothing,\\nmore frequency smoothing\",\n    ],\n):\n    power = epochs.compute_tfr(\n        method=\"multitaper\",\n        freqs=freqs,\n        n_cycles=n_cycles,\n        time_bandwidth=time_bandwidth,\n        return_itc=False,\n        average=True,\n    )\n    ax.set_title(title)\n    # Plot results. Baseline correct based on first 100 ms.\n    power.plot(\n        [0],\n        baseline=(0.0, 0.1),\n        mode=\"mean\",\n        vlim=(vmin, vmax),\n        axes=ax,\n        show=False,\n        colorbar=False,\n    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Stockwell (S) transform\n\nStockwell uses a Gaussian window to balance temporal and spectral resolution.\nImportantly, frequency bands are phase-normalized, hence strictly comparable\nwith regard to timing, and, the input signal can be recoverd from the\ntransform in a lossless way if we disregard numerical errors. In this case,\nwe control the spectral / temporal resolution by specifying different widths\nof the gaussian window using the ``width`` parameter.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig, axs = plt.subplots(1, 3, figsize=(15, 5), sharey=True, layout=\"constrained\")\nfmin, fmax = freqs[[0, -1]]\nfor width, ax in zip((0.2, 0.7, 3.0), axs):\n    power = epochs.compute_tfr(method=\"stockwell\", freqs=(fmin, fmax), width=width)\n    power.plot(\n        [0], baseline=(0.0, 0.1), mode=\"mean\", axes=ax, show=False, colorbar=False\n    )\n    ax.set_title(f\"Sim: Using S transform, width = {width:0.1f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Morlet Wavelets\n\nNext, we'll show the TFR using morlet wavelets, which are a sinusoidal wave\nwith a gaussian envelope. We can control the balance between spectral and\ntemporal resolution with the ``n_cycles`` parameter, which defines the\nnumber of cycles to include in the window.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig, axs = plt.subplots(1, 3, figsize=(15, 5), sharey=True, layout=\"constrained\")\nall_n_cycles = [1, 3, freqs / 2.0]\nfor n_cycles, ax in zip(all_n_cycles, axs):\n    power = epochs.compute_tfr(\n        method=\"morlet\", freqs=freqs, n_cycles=n_cycles, return_itc=False, average=True\n    )\n    power.plot(\n        [0],\n        baseline=(0.0, 0.1),\n        mode=\"mean\",\n        vlim=(vmin, vmax),\n        axes=ax,\n        show=False,\n        colorbar=False,\n    )\n    n_cycles = \"scaled by freqs\" if not isinstance(n_cycles, int) else n_cycles\n    ax.set_title(f\"Sim: Using Morlet wavelet, n_cycles = {n_cycles}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Narrow-bandpass Filter and Hilbert Transform\n\nFinally, we'll show a time-frequency representation using a narrow bandpass\nfilter and the Hilbert transform. Choosing the right filter parameters is\nimportant so that you isolate only one oscillation of interest, generally\nthe width of this filter is recommended to be about 2 Hz.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig, axs = plt.subplots(1, 3, figsize=(15, 5), sharey=True, layout=\"constrained\")\nbandwidths = [1.0, 2.0, 4.0]\nfor bandwidth, ax in zip(bandwidths, axs):\n    data = np.zeros(\n        (len(epochs), len(ch_names), freqs.size, epochs.times.size), dtype=complex\n    )\n    for idx, freq in enumerate(freqs):\n        # Filter raw data and re-epoch to avoid the filter being longer than\n        # the epoch data for low frequencies and short epochs, such as here.\n        raw_filter = raw.copy()\n        # NOTE: The bandwidths of the filters are changed from their defaults\n        # to exaggerate differences. With the default transition bandwidths,\n        # these are all very similar because the filters are almost the same.\n        # In practice, using the default is usually a wise choice.\n        raw_filter.filter(\n            l_freq=freq - bandwidth / 2,\n            h_freq=freq + bandwidth / 2,\n            # no negative values for large bandwidth and low freq\n            l_trans_bandwidth=min([4 * bandwidth, freq - bandwidth]),\n            h_trans_bandwidth=4 * bandwidth,\n        )\n        raw_filter.apply_hilbert()\n        epochs_hilb = Epochs(\n            raw_filter, events, tmin=0, tmax=n_times / sfreq, baseline=(0, 0.1)\n        )\n        data[:, :, idx] = epochs_hilb.get_data()\n    power = EpochsTFRArray(epochs.info, data, epochs.times, freqs, method=\"hilbert\")\n    power.average().plot(\n        [0],\n        baseline=(0.0, 0.1),\n        mode=\"mean\",\n        vlim=(0, 0.1),\n        axes=ax,\n        show=False,\n        colorbar=False,\n    )\n    n_cycles = \"scaled by freqs\" if not isinstance(n_cycles, int) else n_cycles\n    ax.set_title(\n        \"Sim: Using narrow bandpass filter Hilbert,\\n\"\n        f\"bandwidth = {bandwidth}, \"\n        f\"transition bandwidth = {4 * bandwidth}\"\n    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Calculating a TFR without averaging over epochs\n\nIt is also possible to calculate a TFR without averaging across trials.\nWe can do this by using ``average=False``. In this case, an instance of\n:class:`mne.time_frequency.EpochsTFR` is returned.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "n_cycles = freqs / 2.0\npower = epochs.compute_tfr(\n    method=\"morlet\", freqs=freqs, n_cycles=n_cycles, return_itc=False, average=False\n)\nprint(type(power))\navgpower = power.average()\navgpower.plot(\n    [0],\n    baseline=(0.0, 0.1),\n    mode=\"mean\",\n    vlim=(vmin, vmax),\n    title=\"Using Morlet wavelets and EpochsTFR\",\n    show=False,\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Operating on arrays\n\nMNE-Python also has functions that operate on :class:`NumPy arrays <numpy.ndarray>`\ninstead of MNE-Python objects. These are :func:`~mne.time_frequency.tfr_array_morlet`\nand :func:`~mne.time_frequency.tfr_array_multitaper`. They expect inputs of the shape\n``(n_epochs, n_channels, n_times)`` and return an array of shape\n``(n_epochs, n_channels, n_freqs, n_times)`` (or optionally, can collapse the epochs\ndimension if you want average power or inter-trial coherence; see ``output`` param).\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "power = tfr_array_morlet(\n    epochs.get_data(),\n    sfreq=epochs.info[\"sfreq\"],\n    freqs=freqs,\n    n_cycles=n_cycles,\n    output=\"avg_power\",\n    zero_mean=False,\n)\n# Put it into a TFR container for easy plotting\ntfr = AverageTFRArray(\n    info=epochs.info, data=power, times=epochs.times, freqs=freqs, nave=len(epochs)\n)\ntfr.plot(\n    baseline=(0.0, 0.1),\n    picks=[0],\n    mode=\"mean\",\n    vlim=(vmin, vmax),\n    title=\"TFR calculated on a NumPy array\",\n)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}