code: def test_time_frequency():
      """Test time-frequency transform (PSD and ITC)."""
      # Set parameters
      event_id = 1
      tmin = -0.2
      tmax = 0.498  # Allows exhaustive decimation testing
  
      # Setup for reading the raw data
      raw = read_raw_fif(raw_fname)
      events = read_events(event_fname)
  
      include = []
      exclude = raw.info["bads"] + [
          "MEG 2443",
          "EEG 053",
      ]  # bads + 2 more
  
      # picks MEG gradiometers
      picks = pick_types(
          raw.info,
          meg="grad",
          eeg=False,
          stim=False,
          include=include,
          exclude=exclude,
      )
  
      picks = picks[:2]
      epochs = Epochs(raw, events, event_id, tmin, tmax, picks=picks)
      data = epochs.get_data()
      times = epochs.times
      nave = len(data)
  
      epochs_nopicks = Epochs(raw, events, event_id, tmin, tmax)
  
      freqs = np.arange(6, 20, 5)  # define frequencies of interest
      n_cycles = freqs / 4.0
  
      # Test first with a single epoch
      power, itc = tfr_morlet(
          epochs[0],
          freqs=freqs,
          n_cycles=n_cycles,
          use_fft=True,
          return_itc=True,
      )
  
      # Now compute evoked
      evoked = epochs.average()
      pytest.raises(
          ValueError, tfr_morlet, evoked, freqs, 1.0, return_itc=True
      )
      power, itc = tfr_morlet(
          epochs,
          freqs=freqs,
          n_cycles=n_cycles,
          use_fft=True,
          return_itc=True,
      )
      power_, itc_ = tfr_morlet(
          epochs,
          freqs=freqs,
          n_cycles=n_cycles,
          use_fft=True,
          return_itc=True,
          decim=slice(0, 2),
      )
      # Test picks argument and average parameter
      pytest.raises(
          ValueError,
          tfr_morlet,
          epochs,
          freqs=freqs,
          n_cycles=n_cycles,
          return_itc=True,
          average=False,
      )
  
      power_picks, itc_picks = tfr_morlet(
          epochs_nopicks,
          freqs=freqs,
          n_cycles=n_cycles,
          use_fft=True,
          return_itc=True,
          picks=picks,
          average=True,
      )
  
      epochs_power_picks = tfr_morlet(
          epochs_nopicks,
          freqs=freqs,
          n_cycles=n_cycles,
          use_fft=True,
          return_itc=False,
          picks=picks,
          average=False,
      )
      assert_allclose(
          epochs_power_picks.data[0, 0, 0, 0], 9.130315e-23, rtol=1e-4
      )
      power_picks_avg = epochs_power_picks.average()
      # the actual data arrays here are equivalent, too...
      assert_allclose(power.data, power_picks.data)
      assert_allclose(power.data, power_picks_avg.data)
      assert_allclose(itc.data, itc_picks.data)
  
      # test on evoked
      power_evoked = tfr_morlet(
          evoked, freqs, n_cycles, use_fft=True, return_itc=False
      )
      # one is squared magnitude of the average (evoked) and
      # the other is average of the squared magnitudes (epochs PSD)
      # so values shouldn't match, but shapes should
      assert_array_equal(power.data.shape, power_evoked.data.shape)
      with pytest.raises(
          AssertionError, match="Not equal to tolerance"
      ):
          assert_allclose(power.data, power_evoked.data)
  
      # complex output
      with pytest.raises(
          ValueError, match='must be "power" if average=True'
      ):
          tfr_morlet(
              epochs,
              freqs,
              n_cycles,
              return_itc=False,
              average=True,
              output="complex",
          )
      with pytest.raises(
          ValueError, match="Inter-trial coher.*average=False"
      ):
          tfr_morlet(
              epochs,
              freqs,
              n_cycles,
              return_itc=True,
              average=False,
              output="complex",
          )
      epochs_power_complex = tfr_morlet(
          epochs,
          freqs,
          n_cycles,
          return_itc=False,
          average=False,
          output="complex",
      )
      epochs_amplitude_2 = abs(epochs_power_complex)
      epochs_amplitude_3 = epochs_amplitude_2.copy()
      epochs_amplitude_3.data[
          :
      ] = np.inf  # test that it's actually copied
  
      # test that the power computed via `complex` is equivalent to power
      # computed within the method.
      assert_allclose(
          epochs_amplitude_2.data**2, epochs_power_picks.data
      )
  
      # test that averaging power across tapers when multitaper with
      # output='complex' gives the same as output='power'
      epoch_data = epochs.get_data()
      multitaper_power = tfr_array_multitaper(
          epoch_data,
          epochs.info["sfreq"],
          freqs,
          n_cycles,
          output="power",
      )
      multitaper_complex = tfr_array_multitaper(
          epoch_data,
          epochs.info["sfreq"],
          freqs,
          n_cycles,
          output="complex",
      )
  
      taper_dim = 2
      power_from_complex = (
          multitaper_complex * multitaper_complex.conj()
      ).real.mean(axis=taper_dim)
      assert_allclose(power_from_complex, multitaper_power)
  
      print(itc)  # test repr
      print(itc.ch_names)  # test property
      itc += power  # test add
      itc -= power  # test sub
      ret = itc * 23  # test mult
      itc = ret / 23  # test dic
  
      power = power.apply_baseline(baseline=(-0.1, 0), mode="logratio")
      assert power.baseline == (-0.1, 0)
  
      assert "meg" in power
      assert "grad" in power
      assert "mag" not in power
      assert "eeg" not in power
  
      assert power.nave == nave
      assert itc.nave == nave
      assert power.data.shape == (len(picks), len(freqs), len(times))
      assert power.data.shape == itc.data.shape
      assert power_.data.shape == (len(picks), len(freqs), 2)
      assert power_.data.shape == itc_.data.shape
      assert np.sum(itc.data >= 1) == 0
      assert np.sum(itc.data <= 0) == 0
  
      # grand average
      itc2 = itc.copy()
      itc2.info["bads"] = [itc2.ch_names[0]]  # test channel drop
      gave = grand_average([itc2, itc])
      assert gave.data.shape == (
          itc2.data.shape[0] - 1,
          itc2.data.shape[1],
          itc2.data.shape[2],
      )
      assert itc2.ch_names[1:] == gave.ch_names
      assert gave.nave == 2
      itc2.drop_channels(itc2.info["bads"])
      assert_allclose(gave.data, itc2.data)
      itc2.data = np.ones(itc2.data.shape)
      itc.data = np.zeros(itc.data.shape)
      itc2.nave = 2
      itc.nave = 1
      itc.drop_channels([itc.ch_names[0]])
      combined_itc = combine_tfr([itc2, itc])
      assert_allclose(
          combined_itc.data, np.ones(combined_itc.data.shape) * 2 / 3
      )
  
      # more tests
      power, itc = tfr_morlet(
          epochs,
          freqs=freqs,
          n_cycles=2,
          use_fft=False,
          return_itc=True,
      )
  
      assert power.data.shape == (len(picks), len(freqs), len(times))
      assert power.data.shape == itc.data.shape
      assert np.sum(itc.data >= 1) == 0
      assert np.sum(itc.data <= 0) == 0
  
      tfr = tfr_morlet(
          epochs[0],
          freqs,
          use_fft=True,
          n_cycles=2,
          average=False,
          return_itc=False,
      )
      tfr_data = tfr.data[0]
      assert tfr_data.shape == (len(picks), len(freqs), len(times))
      tfr2 = tfr_morlet(
          epochs[0],
          freqs,
          use_fft=True,
          n_cycles=2,
          decim=slice(0, 2),
          average=False,
          return_itc=False,
      ).data[0]
      assert tfr2.shape == (len(picks), len(freqs), 2)
  
      single_power = tfr_morlet(
          epochs, freqs, 2, average=False, return_itc=False
      ).data
      single_power2 = tfr_morlet(
          epochs,
          freqs,
          2,
          decim=slice(0, 2),
          average=False,
          return_itc=False,
      ).data
      single_power3 = tfr_morlet(
          epochs,
          freqs,
          2,
          decim=slice(1, 3),
          average=False,
          return_itc=False,
      ).data
      single_power4 = tfr_morlet(
          epochs,
          freqs,
          2,
          decim=slice(2, 4),
          average=False,
          return_itc=False,
      ).data
  
      assert_allclose(np.mean(single_power, axis=0), power.data)
      assert_allclose(
          np.mean(single_power2, axis=0), power.data[:, :, :2]
      )
      assert_allclose(
          np.mean(single_power3, axis=0), power.data[:, :, 1:3]
      )
      assert_allclose(
          np.mean(single_power4, axis=0), power.data[:, :, 2:4]
      )
  
      power_pick = power.pick(power.ch_names[:10:2])
      assert_equal(len(power_pick.ch_names), len(power.ch_names[:10:2]))
      assert_equal(power_pick.data.shape[0], len(power.ch_names[:10:2]))
      power_drop = power.drop_channels(power.ch_names[1:10:2])
      assert_equal(power_drop.ch_names, power_pick.ch_names)
      assert_equal(power_pick.data.shape[0], len(power_drop.ch_names))
  
      power_pick, power_drop = mne.equalize_channels(
          [power_pick, power_drop]
      )
      assert_equal(power_pick.ch_names, power_drop.ch_names)
      assert_equal(power_pick.data.shape, power_drop.data.shape)
  
      # Test decimation:
      # 2: multiple of len(times) even
      # 3: multiple odd
      # 8: not multiple, even
      # 9: not multiple, odd
      for decim in [2, 3, 8, 9]:
          for use_fft in [True, False]:
              power, itc = tfr_morlet(
                  epochs,
                  freqs=freqs,
                  n_cycles=2,
                  use_fft=use_fft,
                  return_itc=True,
                  decim=decim,
              )
              assert_equal(
                  power.data.shape[2],
                  np.ceil(float(len(times)) / decim),
              )
      freqs = list(range(50, 55))
      decim = 2
      _, n_chan, n_time = data.shape
      tfr = tfr_morlet(
          epochs[0],
          freqs,
          2.0,
          decim=decim,
          average=False,
          return_itc=False,
      ).data[0]
      assert_equal(tfr.shape, (n_chan, len(freqs), n_time // decim))
  
      # Test cwt modes
      Ws = morlet(512, [10, 20], n_cycles=2)
      pytest.raises(ValueError, cwt, data[0, :, :], Ws, mode="foo")
      for use_fft in [True, False]:
          for mode in ["same", "valid", "full"]:
              cwt(data[0], Ws, use_fft=use_fft, mode=mode)
  
      # Test invalid frequency arguments
      with pytest.raises(
          ValueError, match=" 'freqs' must be greater than 0"
      ):
          tfr_morlet(epochs, freqs=np.arange(0, 3), n_cycles=7)
      with pytest.raises(
          ValueError, match=" 'freqs' must be greater than 0"
      ):
          tfr_morlet(epochs, freqs=np.arange(-4, -1), n_cycles=7)
  
      # Test decim parameter checks
      pytest.raises(
          TypeError,
          tfr_morlet,
          epochs,
          freqs=freqs,
          n_cycles=n_cycles,
          use_fft=True,
          return_itc=True,
          decim="decim",
      )
  
      # When convolving in time, wavelets must not be longer than the data
      pytest.raises(
          ValueError,
          cwt,
          data[0, :, : Ws[0].size - 1],
          Ws,
          use_fft=False,
      )
      with pytest.warns(
          UserWarning, match="one of the wavelets.*is longer"
      ):
          cwt(data[0, :, : Ws[0].size - 1], Ws, use_fft=True)
  
      # Check for off-by-one errors when using wavelets with an even number of
      # samples
      psd = cwt(data[0], [Ws[0][:-1]], use_fft=False, mode="full")
      assert_equal(psd.shape, (2, 1, 420))
  
imports: import datetime
  import re
  from itertools import product
  from pathlib import Path
  import matplotlib.pyplot
  import numpy
  import pytest
  from numpy.testing import assert_allclose, assert_array_equal, assert_equal
  from scipy.signal import morlet2
  import mne
  from mne import Epochs, EpochsArray, Info, Transform, create_info, pick_types, read_events
  from mne.io import read_raw_fif
  from mne.tests.test_epochs import assert_metadata_equal
  from mne.time_frequency import tfr_array_morlet, tfr_array_multitaper
  from mne.time_frequency.tfr import AverageTFR, EpochsTFR, _compute_tfr, _make_dpss, combine_tfr, cwt, fwhm, morlet, read_tfrs, tfr_morlet, tfr_multitaper, write_tfrs
  from mne.utils import catch_logging, grand_average
  from mne.viz.utils import _fake_click, _fake_keypress, _fake_scroll
label_desc: Use image node_id 3 for calling a global function with example usage: test_time_frequency() without return types